From fe64274cc707a255749596e5efa3b7f5684040fb Mon Sep 17 00:00:00 2001 From: JonathanLab Date: Wed, 22 Jul 2026 17:20:25 +0200 Subject: [PATCH 01/45] feat: add durable cloud Pi runtime sessions --- apps/code/src/renderer/di/bindings.ts | 14 +- apps/code/src/renderer/di/container.ts | 14 +- apps/web/src/web-container.ts | 25 +- packages/agent/package.json | 10 +- .../conversation/translatePiMessage.test.ts | 82 ++- .../src/pi/conversation/translatePiMessage.ts | 22 + packages/agent/src/pi/remote-rpc-client.ts | 164 +++++ packages/agent/src/pi/rpc-client.test.ts | 35 +- packages/agent/src/pi/rpc-client.ts | 39 +- packages/agent/src/pi/rpc-host.ts | 8 +- packages/agent/src/pi/rpc-transport.test.ts | 62 ++ packages/agent/src/pi/rpc-transport.ts | 61 ++ packages/agent/src/pi/runtime.test.ts | 22 +- packages/agent/src/pi/runtime.ts | 31 +- packages/agent/src/pi/types.ts | 4 +- packages/agent/src/posthog-api.test.ts | 158 +++++ packages/agent/src/posthog-api.ts | 120 +++- packages/agent/src/server/bin.ts | 20 +- .../agent/src/server/pi-agent-server.test.ts | 376 ++++++++++ packages/agent/src/server/pi-agent-server.ts | 651 ++++++++++++++++++ packages/agent/src/server/types.ts | 3 + packages/agent/tsup.config.ts | 20 +- .../core/src/cloud-task/cloud-task-types.ts | 69 -- packages/core/src/cloud-task/cloud-task.ts | 16 +- .../core/src/cloud-task/cloudTaskClient.ts | 23 + packages/core/src/cloud-task/schemas.ts | 9 +- .../pi-runtime/cloudPiSessionClient.test.ts | 351 ++++++++++ .../src/pi-runtime/cloudPiSessionClient.ts | 371 ++++++++++ .../core/src/pi-runtime/pi-runtime.module.ts | 7 +- .../pi-runtime/piSessionController.test.ts | 358 +++++++++- .../src/pi-runtime/piSessionController.ts | 411 ++++++++--- .../src/pi-runtime/piSessionProvider.test.ts | 154 +++++ .../core/src/pi-runtime/piSessionProvider.ts | 85 +++ .../core/src/pi-runtime/piSessionStore.ts | 3 + .../core/src/task-detail/piTaskCreator.ts | 197 ------ .../src/task-detail/taskCreationApiClient.ts | 1 + .../src/task-detail/taskCreationSaga.test.ts | 144 +++- .../core/src/task-detail/taskCreationSaga.ts | 57 +- .../core/src/task-detail/taskService.test.ts | 51 ++ packages/core/src/task-detail/taskService.ts | 82 ++- .../extensions/posthog-provider/provider.ts | 19 +- packages/harness/src/runtime.test.ts | 48 +- packages/harness/src/runtime.ts | 70 +- packages/host-router/package.json | 1 + packages/host-router/src/cloud-task-client.ts | 47 ++ .../host-router/src/pi-runner.ts | 0 packages/host-router/src/pi-session-client.ts | 96 --- .../host-router/src/pi-session-factory.ts | 61 ++ .../src/routers/cloud-task.router.ts | 7 + .../src/routers/pi-session.router.ts | 295 +------- packages/shared/src/session-events.ts | 3 + .../features/pi-sessions/PiSessionView.tsx | 16 +- .../buildAgentConversationItems.test.ts | 54 ++ .../task-detail/components/TaskDetail.tsx | 15 +- .../task-detail/components/TaskInput.tsx | 3 +- .../services/pi-session/pi-session.test.ts | 205 +++++- .../src/services/pi-session/pi-session.ts | 286 ++------ .../src/services/pi-session/schemas.ts | 311 +-------- pnpm-lock.yaml | 87 +-- pnpm-workspace.yaml | 8 +- 60 files changed, 4328 insertions(+), 1634 deletions(-) create mode 100644 packages/agent/src/pi/remote-rpc-client.ts create mode 100644 packages/agent/src/pi/rpc-transport.test.ts create mode 100644 packages/agent/src/pi/rpc-transport.ts create mode 100644 packages/agent/src/server/pi-agent-server.test.ts create mode 100644 packages/agent/src/server/pi-agent-server.ts delete mode 100644 packages/core/src/cloud-task/cloud-task-types.ts create mode 100644 packages/core/src/cloud-task/cloudTaskClient.ts create mode 100644 packages/core/src/pi-runtime/cloudPiSessionClient.test.ts create mode 100644 packages/core/src/pi-runtime/cloudPiSessionClient.ts create mode 100644 packages/core/src/pi-runtime/piSessionProvider.test.ts create mode 100644 packages/core/src/pi-runtime/piSessionProvider.ts delete mode 100644 packages/core/src/task-detail/piTaskCreator.ts create mode 100644 packages/host-router/src/cloud-task-client.ts rename apps/code/src/renderer/platform-adapters/trpc-pi-runner.ts => packages/host-router/src/pi-runner.ts (100%) delete mode 100644 packages/host-router/src/pi-session-client.ts create mode 100644 packages/host-router/src/pi-session-factory.ts diff --git a/apps/code/src/renderer/di/bindings.ts b/apps/code/src/renderer/di/bindings.ts index f8556c970f..5bd2b7626d 100644 --- a/apps/code/src/renderer/di/bindings.ts +++ b/apps/code/src/renderer/di/bindings.ts @@ -13,6 +13,10 @@ import { type AutoresearchSessionClient, type AutoresearchStorageClient, } from "@posthog/core/autoresearch/identifiers"; +import { + CLOUD_TASK_CLIENT, + type CloudTaskClient, +} from "@posthog/core/cloud-task/cloudTaskClient"; import { CODE_REVIEW_WORKSPACE_CLIENT, REVERT_HUNK_SERVICE, @@ -76,8 +80,10 @@ import { import { PI_RUNNER } from "@posthog/core/pi-runtime/identifiers"; import type { PiRunner } from "@posthog/core/pi-runtime/piRunner"; import { - PI_SESSION_CLIENT, - type PiSessionClient, + LOCAL_PI_SESSION_FACTORY, + PI_SESSION_PROVIDER, + type PiSessionFactory, + type PiSessionProvider, } from "@posthog/core/pi-runtime/piSessionController"; import { type BundleLocalSkill, @@ -295,7 +301,9 @@ export interface RendererBindings { [ANALYTICS_TRACKER]: AnalyticsTracker; [TASK_CREATION_HOST]: ITaskCreationHost; [PI_RUNNER]: PiRunner; - [PI_SESSION_CLIENT]: PiSessionClient; + [PI_SESSION_PROVIDER]: PiSessionProvider; + [LOCAL_PI_SESSION_FACTORY]: PiSessionFactory; + [CLOUD_TASK_CLIENT]: CloudTaskClient; [TASK_CREATION_EFFECTS]: TaskCreationEffects; [RENDERER_TASK_SERVICE]: TaskService; [TASK_SERVICE]: TaskService; diff --git a/apps/code/src/renderer/di/container.ts b/apps/code/src/renderer/di/container.ts index 5d7d27ddf4..dba9224fd8 100644 --- a/apps/code/src/renderer/di/container.ts +++ b/apps/code/src/renderer/di/container.ts @@ -2,6 +2,10 @@ import "reflect-metadata"; import { useDevFlagsStore } from "@features/dev-toolbar/devFlagsStore"; import { TypedContainer } from "@inversifyjs/strongly-typed"; import type { TrpcRouter } from "@main/trpc/router"; +import { + CLOUD_TASK_CLIENT, + type CloudTaskClient, +} from "@posthog/core/cloud-task/cloudTaskClient"; import { CODE_REVIEW_WORKSPACE_CLIENT, REVERT_HUNK_SERVICE, @@ -35,7 +39,7 @@ import type { LocalMcpWorkspaceClient } from "@posthog/core/local-mcp/localMcpIm import { PI_RUNNER } from "@posthog/core/pi-runtime/identifiers"; import { piRuntimeModule } from "@posthog/core/pi-runtime/pi-runtime.module"; import type { PiRunner } from "@posthog/core/pi-runtime/piRunner"; -import { PI_SESSION_CLIENT } from "@posthog/core/pi-runtime/piSessionController"; +import { LOCAL_PI_SESSION_FACTORY } from "@posthog/core/pi-runtime/piSessionController"; import { CLOUD_ARTIFACT_BUNDLE_LOCAL_SKILL, CLOUD_ARTIFACT_READ_FILE_AS_BASE64, @@ -89,7 +93,9 @@ import { import { WorkspaceSetupService } from "@posthog/core/workspace/WorkspaceSetupService"; import { setRootContainer } from "@posthog/di/container"; import { HOST_TRPC_CLIENT } from "@posthog/host-router/client"; -import { TrpcPiSessionClient } from "@posthog/host-router/pi-session-client"; +import { TrpcCloudTaskClient } from "@posthog/host-router/cloud-task-client"; +import { TrpcPiRunner } from "@posthog/host-router/pi-runner"; +import { TrpcPiSessionFactory } from "@posthog/host-router/pi-session-factory"; import { BROWSER_TABS_CLIENT, type BrowserTabsClient, @@ -156,7 +162,6 @@ import { trpcClient } from "@renderer/trpc"; import { hostTrpcClient } from "@renderer/trpc/client"; import type { TRPCClient } from "@trpc/client"; import { hostLog, logger } from "@utils/logger"; -import { TrpcPiRunner } from "../platform-adapters/trpc-pi-runner"; import type { RendererBindings } from "./bindings"; import { TASK_SERVICE as RENDERER_TASK_SERVICE, TRPC_CLIENT } from "./tokens"; @@ -297,7 +302,8 @@ container // Bind services container.bind(TASK_CREATION_HOST).to(TrpcTaskCreationHost); container.bind(PI_RUNNER).to(TrpcPiRunner); -container.bind(PI_SESSION_CLIENT).to(TrpcPiSessionClient); +container.bind(LOCAL_PI_SESSION_FACTORY).to(TrpcPiSessionFactory); +container.bind(CLOUD_TASK_CLIENT).to(TrpcCloudTaskClient); container.load(piRuntimeModule); container.bind(TASK_CREATION_EFFECTS).toConstantValue(taskCreationEffects); container.bind(RENDERER_TASK_SERVICE).to(TaskService); diff --git a/apps/web/src/web-container.ts b/apps/web/src/web-container.ts index ebeacc79b1..f090f1a60b 100644 --- a/apps/web/src/web-container.ts +++ b/apps/web/src/web-container.ts @@ -29,6 +29,10 @@ import { canvasCoreModule } from "@posthog/core/canvas/canvas.module"; import { taskThreadCoreModule } from "@posthog/core/canvas/taskThread.module"; import type { CloudTaskService } from "@posthog/core/cloud-task/cloud-task"; import { cloudTaskModule } from "@posthog/core/cloud-task/cloud-task.module"; +import { + CLOUD_TASK_CLIENT, + type CloudTaskClient, +} from "@posthog/core/cloud-task/cloudTaskClient"; import { CLOUD_TASK_AUTH, CLOUD_TASK_SERVICE, @@ -86,10 +90,14 @@ import { type GithubConnectClient as OnboardingGithubConnectContract, } from "@posthog/core/onboarding/identifiers"; import { onboardingModule } from "@posthog/core/onboarding/onboarding.module"; +import { PI_RUNNER } from "@posthog/core/pi-runtime/identifiers"; import { piRuntimeModule } from "@posthog/core/pi-runtime/pi-runtime.module"; +import type { PiRunner } from "@posthog/core/pi-runtime/piRunner"; import { - PI_SESSION_CLIENT, - type PiSessionClient, + LOCAL_PI_SESSION_FACTORY, + PI_SESSION_PROVIDER, + type PiSessionFactory, + type PiSessionProvider, } from "@posthog/core/pi-runtime/piSessionController"; import { type BundleLocalSkill, @@ -161,7 +169,9 @@ import { HOST_TRPC_CLIENT, type HostTrpcClient, } from "@posthog/host-router/client"; -import { TrpcPiSessionClient } from "@posthog/host-router/pi-session-client"; +import { TrpcCloudTaskClient } from "@posthog/host-router/cloud-task-client"; +import { TrpcPiRunner } from "@posthog/host-router/pi-runner"; +import { TrpcPiSessionFactory } from "@posthog/host-router/pi-session-factory"; import { ANALYTICS_SERVICE, type IAnalytics, @@ -313,7 +323,10 @@ import { hostTrpcClient } from "./web-trpc"; interface WebBindings { [HOST_TRPC_CLIENT]: HostTrpcClient; - [PI_SESSION_CLIENT]: PiSessionClient; + [PI_SESSION_PROVIDER]: PiSessionProvider; + [LOCAL_PI_SESSION_FACTORY]: PiSessionFactory; + [CLOUD_TASK_CLIENT]: CloudTaskClient; + [PI_RUNNER]: PiRunner; [ROOT_LOGGER]: RootLogger; [HOST_LOGGER]: HostLogger; [FEATURE_FLAGS]: FeatureFlags; @@ -394,7 +407,9 @@ export const container = new TypedContainer({ // Keystone: the same typed host client the renderer binds — served in-process // here (web-trpc.ts) instead of over Electron IPC. container.bind(HOST_TRPC_CLIENT).toConstantValue(hostTrpcClient); -container.bind(PI_SESSION_CLIENT).to(TrpcPiSessionClient); +container.bind(LOCAL_PI_SESSION_FACTORY).to(TrpcPiSessionFactory); +container.bind(CLOUD_TASK_CLIENT).to(TrpcCloudTaskClient); +container.bind(PI_RUNNER).to(TrpcPiRunner); container.load(piRuntimeModule); // Logger: web uses console; electron uses electron-log. Same RootLogger shape. diff --git a/packages/agent/package.json b/packages/agent/package.json index cd94555a5c..7a13079bcb 100644 --- a/packages/agent/package.json +++ b/packages/agent/package.json @@ -32,6 +32,14 @@ "types": "./dist/pi/rpc-client.d.ts", "import": "./dist/pi/rpc-client.js" }, + "./pi/rpc-transport": { + "types": "./dist/pi/rpc-transport.d.ts", + "import": "./dist/pi/rpc-transport.js" + }, + "./pi/remote-rpc-client": { + "types": "./dist/pi/remote-rpc-client.d.ts", + "import": "./dist/pi/remote-rpc-client.js" + }, "./pi/conversation": { "types": "./dist/pi/conversation/translatePiConversation.d.ts", "import": "./dist/pi/conversation/translatePiConversation.js" @@ -114,7 +122,7 @@ } }, "bin": { - "agent-server": "./dist/server/bin.cjs" + "agent-server": "./dist/server/bin.js" }, "type": "module", "keywords": [ diff --git a/packages/agent/src/pi/conversation/translatePiMessage.test.ts b/packages/agent/src/pi/conversation/translatePiMessage.test.ts index 0393331230..b0743919ce 100644 --- a/packages/agent/src/pi/conversation/translatePiMessage.test.ts +++ b/packages/agent/src/pi/conversation/translatePiMessage.test.ts @@ -1,4 +1,8 @@ -import type { AssistantMessage, UserMessage } from "@earendil-works/pi-ai"; +import type { + AssistantMessage, + ToolResultMessage, + UserMessage, +} from "@earendil-works/pi-ai"; import { describe, expect, it } from "vitest"; import { createPiMessageTranslator } from "./translatePiMessage"; @@ -110,4 +114,80 @@ describe("createPiMessageTranslator", () => { }, ]); }); + + it("provides generic rendered content for extension tool results", () => { + const translator = createPiMessageTranslator(); + const content: ToolResultMessage["content"] = [ + { type: "text", text: "Found " }, + { type: "text", text: "three matches" }, + ]; + const message: ToolResultMessage = { + role: "toolResult", + toolCallId: "extension-1", + toolName: "web_search", + content, + details: { resultCount: 3 }, + isError: false, + timestamp: 12, + }; + + expect(translator.translate(message)).toEqual([ + { + type: "tool_call_updated", + timestamp: 12, + toolCall: { + id: "extension-1", + status: "completed", + rawOutput: content, + content: [ + { + type: "content", + content: { type: "text", text: "Found three matches" }, + }, + ], + }, + }, + ]); + }); + + it("keeps built-in tool translation and raw output", () => { + const translator = createPiMessageTranslator(); + const content: ToolResultMessage["content"] = [ + { type: "text", text: "file contents" }, + ]; + + translator.translateToolExecutionStart( + "read-1", + "read", + { path: "src/file.ts" }, + 1, + ); + + expect( + translator.translateToolExecutionEnd( + "read-1", + "read", + { content }, + false, + 2, + ), + ).toEqual([ + { + type: "tool_call_updated", + timestamp: 2, + toolCall: { + id: "read-1", + status: "completed", + rawOutput: content, + locations: [{ path: "src/file.ts" }], + content: [ + { + type: "content", + content: { type: "text", text: "file contents" }, + }, + ], + }, + }, + ]); + }); }); diff --git a/packages/agent/src/pi/conversation/translatePiMessage.ts b/packages/agent/src/pi/conversation/translatePiMessage.ts index 874dc06242..dbccce3a30 100644 --- a/packages/agent/src/pi/conversation/translatePiMessage.ts +++ b/packages/agent/src/pi/conversation/translatePiMessage.ts @@ -7,6 +7,7 @@ import type { import type { AgentContent, AgentConversationEvent, + AgentToolCallContent, AgentToolCallStatus, } from "@posthog/shared"; import { type PiToolName, TOOL_KIND_BY_NAME } from "./toolKind"; @@ -43,6 +44,21 @@ function isPiToolName(name: string): name is PiToolName { return name in TOOL_KIND_BY_NAME; } +function toGenericToolContent( + resultContent: ToolResultMessage["content"], +): AgentToolCallContent[] | undefined { + const text = resultContent + .filter((block) => block.type === "text") + .map((block) => block.text) + .join(""); + + if (!text) { + return undefined; + } + + return [{ type: "content", content: { type: "text", text } }]; +} + function toContent(block: { type: string; text?: string; @@ -214,6 +230,12 @@ export function createPiMessageTranslator(): PiMessageTranslator { if (output.locations) { toolCall.locations = output.locations; } + } else { + const content = toGenericToolContent(result.content); + + if (content) { + toolCall.content = content; + } } return [{ type: "tool_call_updated", timestamp, toolCall }]; diff --git a/packages/agent/src/pi/remote-rpc-client.ts b/packages/agent/src/pi/remote-rpc-client.ts new file mode 100644 index 0000000000..fd3debda1e --- /dev/null +++ b/packages/agent/src/pi/remote-rpc-client.ts @@ -0,0 +1,164 @@ +import type { + RpcClient, + RpcCommand, + RpcResponse, +} from "@earendil-works/pi-coding-agent"; +import type { AgentConversationEvent } from "@posthog/shared"; +import { createPiConversationTranslator } from "./conversation/translatePiConversation"; +import { type PiRpcTransport, parsePiRpcResponse } from "./rpc-transport"; + +export type PiRemoteRpcClient = Pick< + RpcClient, + | "prompt" + | "steer" + | "followUp" + | "abort" + | "getState" + | "setModel" + | "getAvailableModels" + | "getAvailableThinkingLevels" + | "setThinkingLevel" + | "setSteeringMode" + | "setFollowUpMode" + | "compact" + | "bash" + | "abortBash" + | "getEntries" + | "getCommands" +>; + +export async function getRemotePiConversation( + client: Pick, +): Promise { + const entries = await client.getEntries(); + const translator = createPiConversationTranslator(); + const events: AgentConversationEvent[] = []; + + for (const entry of entries.entries) { + if (entry.type === "message") { + events.push(...translator.translateHistoryMessage(entry.message)); + } + } + + return events; +} + +export class RemotePiRpcClient implements PiRemoteRpcClient { + constructor(private readonly transport: PiRpcTransport) {} + + async prompt( + message: string, + images?: Parameters[1], + ): Promise { + await this.request({ type: "prompt", message, images }); + } + + async steer( + message: string, + images?: Parameters[1], + ): Promise { + await this.request({ type: "steer", message, images }); + } + + async followUp( + message: string, + images?: Parameters[1], + ): Promise { + await this.request({ type: "follow_up", message, images }); + } + + async abort(): Promise { + await this.request({ type: "abort" }); + } + + async getState(): ReturnType { + return this.data(await this.request({ type: "get_state" })); + } + + async setModel( + provider: string, + modelId: string, + ): ReturnType { + return this.data( + await this.request({ type: "set_model", provider, modelId }), + ); + } + + async getAvailableModels(): ReturnType< + PiRemoteRpcClient["getAvailableModels"] + > { + const data = this.data<{ + models: Awaited>; + }>(await this.request({ type: "get_available_models" })); + return data.models; + } + + async getAvailableThinkingLevels(): ReturnType< + PiRemoteRpcClient["getAvailableThinkingLevels"] + > { + const data = this.data<{ + levels: Awaited< + ReturnType + >; + }>(await this.request({ type: "get_available_thinking_levels" })); + return data.levels; + } + + async setThinkingLevel( + level: Parameters[0], + ): Promise { + await this.request({ type: "set_thinking_level", level }); + } + + async setSteeringMode( + mode: Parameters[0], + ): Promise { + await this.request({ type: "set_steering_mode", mode }); + } + + async setFollowUpMode( + mode: Parameters[0], + ): Promise { + await this.request({ type: "set_follow_up_mode", mode }); + } + + async compact( + customInstructions?: string, + ): ReturnType { + return this.data( + await this.request({ type: "compact", customInstructions }), + ); + } + + async bash(command: string): ReturnType { + return this.data(await this.request({ type: "bash", command })); + } + + async abortBash(): Promise { + await this.request({ type: "abort_bash" }); + } + + async getEntries( + since?: string, + ): ReturnType { + return this.data(await this.request({ type: "get_entries", since })); + } + + async getCommands(): ReturnType { + const data = this.data<{ + commands: Awaited>; + }>(await this.request({ type: "get_commands" })); + return data.commands; + } + + private async request(command: RpcCommand): Promise { + return parsePiRpcResponse(await this.transport.request(command)); + } + + private data(response: RpcResponse): T { + if (!response.success) { + throw new Error(response.error); + } + return (response as unknown as { data: T }).data; + } +} diff --git a/packages/agent/src/pi/rpc-client.test.ts b/packages/agent/src/pi/rpc-client.test.ts index ff4dcc3c23..4c8239d698 100644 --- a/packages/agent/src/pi/rpc-client.test.ts +++ b/packages/agent/src/pi/rpc-client.test.ts @@ -1,39 +1,6 @@ import { RpcClient } from "@earendil-works/pi-coding-agent"; import { describe, expect, it } from "vitest"; -import { - createPiRpcClient, - getAvailableModelsWithThinkingLevels, - type PiRpcClient, -} from "./rpc-client"; - -describe("getAvailableModelsWithThinkingLevels", () => { - it("uses Pi's per-model capability map", async () => { - const client = { - getAvailableModels: async () => [ - { - provider: "openai", - id: "gpt-5.6", - contextWindow: 200000, - reasoning: true, - thinkingLevelMap: { - off: "none", - minimal: null, - xhigh: "xhigh", - max: "max", - }, - }, - ], - } as unknown as PiRpcClient; - - await expect(getAvailableModelsWithThinkingLevels(client)).resolves.toEqual( - [ - expect.objectContaining({ - thinkingLevels: ["off", "low", "medium", "high", "xhigh", "max"], - }), - ], - ); - }); -}); +import { createPiRpcClient } from "./rpc-client"; describe("createPiRpcClient", () => { it("does not put provider credentials in the child environment", () => { diff --git a/packages/agent/src/pi/rpc-client.ts b/packages/agent/src/pi/rpc-client.ts index 1479a4d6af..b0e6b3ae0e 100644 --- a/packages/agent/src/pi/rpc-client.ts +++ b/packages/agent/src/pi/rpc-client.ts @@ -2,17 +2,11 @@ import { type ChildProcess, spawn } from "node:child_process"; import type { Writable } from "node:stream"; import { StringDecoder } from "node:string_decoder"; import { fileURLToPath } from "node:url"; -import { - type Api, - getSupportedThinkingLevels, - type Model, -} from "@earendil-works/pi-ai"; import { RpcClient, type RpcClientOptions, } from "@earendil-works/pi-coding-agent"; import { safePiEnvironment } from "./rpc-environment"; -import type { PiModelOption, PiThinkingLevel } from "./types"; export type PiRpcClient = RpcClient; @@ -22,19 +16,6 @@ export interface PiRpcProviderOptions { baseUrl?: string; } -export async function getAvailableModelsWithThinkingLevels( - client: PiRpcClient, -): Promise { - const models = await client.getAvailableModels(); - - return models.map((model) => ({ - ...model, - thinkingLevels: getSupportedThinkingLevels( - model as unknown as Model, - ) as PiThinkingLevel[], - })); -} - type RpcClientProcessAccess = { process?: ChildProcess; }; @@ -76,6 +57,7 @@ class SecurePiRpcClient extends RpcClient { constructor( private readonly secureOptions: RpcClientOptions, private readonly providerOptions: PiRpcProviderOptions, + private readonly sessionDir?: string, ) { super(secureOptions); } @@ -148,7 +130,10 @@ class SecurePiRpcClient extends RpcClient { const bootstrapPipe = child.stdio[3] as Writable | null; bootstrapPipe?.end( - JSON.stringify({ providerOptions: this.providerOptions }), + JSON.stringify({ + providerOptions: this.providerOptions, + sessionDir: this.sessionDir, + }), ); await new Promise((resolve) => setTimeout(resolve, 100)); @@ -167,21 +152,29 @@ export function getPiRpcClientProcess( return (client as unknown as RpcClientProcessAccess).process ?? null; } -export type PiRpcClientOptions = Pick & { +export type PiRpcClientOptions = Pick< + RpcClientOptions, + "cliPath" | "cwd" | "model" +> & { + sessionDir?: string; sessionFile?: string; providerOptions: PiRpcProviderOptions; }; export function createPiRpcClient(options: PiRpcClientOptions): PiRpcClient { - const { sessionFile, providerOptions, ...rpcOptions } = options; + const { sessionDir, sessionFile, providerOptions, ...rpcOptions } = options; const args = sessionFile ? ["--session-file", sessionFile] : []; + const cliPath = + rpcOptions.cliPath ?? + fileURLToPath(new URL("./rpc-host.js", import.meta.url)); return new SecurePiRpcClient( { ...rpcOptions, args, - cliPath: fileURLToPath(new URL("./rpc-host.js", import.meta.url)), + cliPath, provider: "posthog", }, providerOptions, + sessionDir, ); } diff --git a/packages/agent/src/pi/rpc-host.ts b/packages/agent/src/pi/rpc-host.ts index 0ce15c5e68..f07fe90e1b 100644 --- a/packages/agent/src/pi/rpc-host.ts +++ b/packages/agent/src/pi/rpc-host.ts @@ -6,6 +6,7 @@ import { sanitizePiHostEnvironment } from "./rpc-environment"; interface PiRpcBootstrap { providerOptions?: PosthogProviderOptions; + sessionDir?: string; } function argumentValue(name: string): string | undefined { @@ -24,7 +25,7 @@ const cwd = process.cwd(); const sessionFile = argumentValue("--session-file"); const sessionManager = sessionFile ? SessionManager.open(sessionFile, undefined, cwd) - : undefined; + : SessionManager.create(cwd, bootstrap.sessionDir); const runtime = await createHarnessRuntime({ cwd, sessionManager, @@ -33,7 +34,10 @@ const runtime = await createHarnessRuntime({ const requestedModel = argumentValue("--model")?.replace(/^posthog\//, ""); if (requestedModel) { - const model = runtime.services.modelRegistry.find("posthog", requestedModel); + const model = runtime.services.modelRuntime.getModel( + "posthog", + requestedModel, + ); if (!model) { throw new Error(`PostHog model not found: ${requestedModel}`); } diff --git a/packages/agent/src/pi/rpc-transport.test.ts b/packages/agent/src/pi/rpc-transport.test.ts new file mode 100644 index 0000000000..4c03d09e0b --- /dev/null +++ b/packages/agent/src/pi/rpc-transport.test.ts @@ -0,0 +1,62 @@ +import type { RpcCommand, RpcResponse } from "@earendil-works/pi-coding-agent"; +import { describe, expect, it, vi } from "vitest"; +import { RemotePiRpcClient } from "./remote-rpc-client"; +import { piRpcCommandSchema } from "./rpc-transport"; + +function response(command: RpcCommand, data?: unknown): RpcResponse { + return { + type: "response" as const, + command: command.type, + success: true as const, + ...(data === undefined ? {} : { data }), + } as RpcResponse; +} + +describe("RemotePiRpcClient", () => { + it("uses Pi's native methods to encode commands over an injected transport", async () => { + const request = vi.fn(async (command: RpcCommand) => { + if (command.type === "compact") { + return response(command, { + summary: "summary", + firstKeptEntryId: "entry-1", + tokensBefore: 100, + }); + } + if (command.type === "get_available_thinking_levels") { + return response(command, { levels: ["off", "high", "xhigh"] }); + } + return response(command); + }); + const client = new RemotePiRpcClient({ request }); + + await client.setFollowUpMode("one-at-a-time"); + const compaction = await client.compact("retain decisions"); + const thinkingLevels = await client.getAvailableThinkingLevels(); + + expect(request).toHaveBeenNthCalledWith(1, { + type: "set_follow_up_mode", + mode: "one-at-a-time", + }); + expect(request).toHaveBeenNthCalledWith(2, { + type: "compact", + customInstructions: "retain decisions", + }); + expect(request).toHaveBeenNthCalledWith(3, { + type: "get_available_thinking_levels", + }); + expect(compaction.summary).toBe("summary"); + expect(thinkingLevels).toEqual(["off", "high", "xhigh"]); + }); + + it("rejects malformed responses from every transport", async () => { + const client = new RemotePiRpcClient({ + request: vi.fn(async () => ({ type: "not-a-response" })), + }); + + await expect(client.getState()).rejects.toThrow(); + }); + + it("requires a native command type at the transport boundary", () => { + expect(() => piRpcCommandSchema.parse({ mode: "invalid" })).toThrow(); + }); +}); diff --git a/packages/agent/src/pi/rpc-transport.ts b/packages/agent/src/pi/rpc-transport.ts new file mode 100644 index 0000000000..37b2f86710 --- /dev/null +++ b/packages/agent/src/pi/rpc-transport.ts @@ -0,0 +1,61 @@ +import type { + AgentSessionEvent, + RpcClient, + RpcCommand, + RpcResponse, +} from "@earendil-works/pi-coding-agent"; +import { z } from "zod/v4"; + +export type { RpcCommand, RpcResponse } from "@earendil-works/pi-coding-agent"; + +export const piRpcCommandSchema = z + .object({ + id: z.string().optional(), + type: z.string().min(1), + }) + .loose() + .transform((command) => command as RpcCommand); + +export const piRpcResponseSchema = z.discriminatedUnion("success", [ + z + .object({ + id: z.string().optional(), + type: z.literal("response"), + command: z.string(), + success: z.literal(true), + data: z.unknown().optional(), + }) + .loose(), + z + .object({ + id: z.string().optional(), + type: z.literal("response"), + command: z.string(), + success: z.literal(false), + error: z.string(), + }) + .loose(), +]); + +export function parsePiRpcResponse(value: unknown): RpcResponse { + return piRpcResponseSchema.parse(value) as RpcResponse; +} + +export interface PiRpcTransport { + request(command: RpcCommand): Promise; + onEvent?(listener: (event: AgentSessionEvent) => void): () => void; + start?(): Promise; + stop?(): Promise; +} + +interface RpcClientInternals { + send(command: RpcCommand): Promise; +} + +export function sendPiRpcCommand( + client: RpcClient, + command: RpcCommand, +): Promise { + const internals = client as unknown as RpcClientInternals; + return internals.send(command); +} diff --git a/packages/agent/src/pi/runtime.test.ts b/packages/agent/src/pi/runtime.test.ts index b9cf4656de..34d3677937 100644 --- a/packages/agent/src/pi/runtime.test.ts +++ b/packages/agent/src/pi/runtime.test.ts @@ -26,22 +26,13 @@ function assistant(text: string): AssistantMessage { }; } -function createClient(messages: AssistantMessage[] = []) { +function createClient() { let listener: (event: AgentSessionEvent) => void = () => {}; const client = { onEvent: vi.fn((nextListener) => { listener = nextListener; return () => {}; }), - getEntries: vi.fn(async () => ({ - entries: messages.map((message, index) => ({ - type: "message" as const, - id: `entry-${index}`, - parentId: null, - timestamp: new Date().toISOString(), - message, - })), - })), } as unknown as RpcClient; return { client, emit: (event: AgentSessionEvent) => listener(event) }; @@ -62,15 +53,4 @@ describe("PiRuntime", () => { content: { type: "text", text: "hello" }, }); }); - - it("normalizes persisted conversation history", async () => { - const { client } = createClient([assistant("history")]); - const runtime = new PiRuntime(client); - - await expect(runtime.conversation()).resolves.toContainEqual({ - type: "assistant_message_chunk", - timestamp: 1, - content: { type: "text", text: "history" }, - }); - }); }); diff --git a/packages/agent/src/pi/runtime.ts b/packages/agent/src/pi/runtime.ts index 102c517fcb..652e6fa902 100644 --- a/packages/agent/src/pi/runtime.ts +++ b/packages/agent/src/pi/runtime.ts @@ -4,14 +4,7 @@ import { createPiConversationTranslator, type PiConversationTranslator, } from "./conversation/translatePiConversation"; -import { - createPiRpcClient, - getAvailableModelsWithThinkingLevels, - getPiRpcClientProcess, - type PiRpcClient, - type PiRpcClientOptions, -} from "./rpc-client"; -import type { PiModelOption } from "./types"; +import { getPiRpcClientProcess, type PiRpcClient } from "./rpc-client"; export class PiRuntime { readonly client: PiRpcClient; @@ -46,24 +39,6 @@ export class PiRuntime { return () => this.conversationListeners.delete(listener); } - availableModels(): Promise { - return getAvailableModelsWithThinkingLevels(this.client); - } - - async conversation(): Promise { - const entries = await this.client.getEntries(); - const translator = createPiConversationTranslator(); - const events: AgentConversationEvent[] = []; - - for (const entry of entries.entries) { - if (entry.type === "message") { - events.push(...translator.translateHistoryMessage(entry.message)); - } - } - - return events; - } - private handleEvent(event: AgentSessionEvent): void { for (const listener of this.runtimeListeners) { listener(event); @@ -77,7 +52,3 @@ export class PiRuntime { } } } - -export function createPiRuntime(options: PiRpcClientOptions): PiRuntime { - return new PiRuntime(createPiRpcClient(options)); -} diff --git a/packages/agent/src/pi/types.ts b/packages/agent/src/pi/types.ts index 30443c7596..9ad9c747af 100644 --- a/packages/agent/src/pi/types.ts +++ b/packages/agent/src/pi/types.ts @@ -32,9 +32,7 @@ export type PiNativeModelInfo = Awaited< ReturnType >[number]; -export type PiModelOption = PiNativeModelInfo & { - thinkingLevels: PiThinkingLevel[]; -}; +export type PiModelOption = PiNativeModelInfo; export type PiCommand = Awaited>[number]; diff --git a/packages/agent/src/posthog-api.test.ts b/packages/agent/src/posthog-api.test.ts index 651f539b3c..992af534b6 100644 --- a/packages/agent/src/posthog-api.test.ts +++ b/packages/agent/src/posthog-api.test.ts @@ -127,6 +127,164 @@ describe("PostHogAPIClient", () => { }, ); + it("loads and syncs the durable task session", async () => { + const client = new PostHogAPIClient({ + apiUrl: "https://app.posthog.com", + getApiKey: vi.fn().mockResolvedValue("token"), + projectId: 7, + }); + const content = '{"type":"session"}\n'; + const access = { + id: "session-1", + download_url: "https://storage.example/session.jsonl", + revision: 3, + }; + const prepared = { + id: "session-1", + sync_id: "sync-1", + upload: { + url: "https://storage.example/upload", + fields: { key: "task-sessions/session-1/uploads/4.jsonl" }, + }, + }; + mockFetch + .mockResolvedValueOnce({ + ok: true, + json: vi.fn().mockResolvedValue(access), + }) + .mockResolvedValueOnce({ + ok: true, + text: vi.fn().mockResolvedValue(content), + }) + .mockResolvedValueOnce({ + ok: true, + json: vi.fn().mockResolvedValue(prepared), + }) + .mockResolvedValueOnce({ ok: true }) + .mockResolvedValueOnce({ + ok: true, + json: vi.fn().mockResolvedValue({ id: "session-1", revision: 4 }), + }); + + const storage = await client.getTaskSession("task-1", "run-1"); + await expect(client.downloadTaskSession(storage)).resolves.toBe(content); + await expect( + client.syncTaskSession("task-1", "run-1", "sandbox-1", 3, content), + ).resolves.toBe(4); + + expect(mockFetch).toHaveBeenNthCalledWith( + 4, + "https://storage.example/upload", + expect.objectContaining({ method: "POST", body: expect.any(FormData) }), + ); + expect(mockFetch).toHaveBeenLastCalledWith( + "https://app.posthog.com/api/projects/7/tasks/task-1/runs/run-1/task_session_sync/", + expect.objectContaining({ + method: "POST", + body: JSON.stringify({ + sandbox_id: "sandbox-1", + sync_id: "sync-1", + expected_revision: 3, + }), + }), + ); + }); + + it("recovers an ambiguous finalize only when its prepared object was promoted", async () => { + const client = new PostHogAPIClient({ + apiUrl: "https://app.posthog.com", + getApiKey: vi.fn().mockResolvedValue("token"), + projectId: 7, + }); + const prepared = { + id: "session-1", + sync_id: "sync-1", + upload: { + url: "https://storage.example/upload", + fields: { + key: "task-sessions/org/task/session/uploads/4-sync-1.jsonl", + }, + }, + }; + mockFetch + .mockResolvedValueOnce({ + ok: true, + json: vi.fn().mockResolvedValue(prepared), + }) + .mockResolvedValueOnce({ ok: true }) + .mockResolvedValueOnce({ + ok: false, + status: 504, + json: vi.fn().mockResolvedValue({ error: "Gateway timeout" }), + }) + .mockResolvedValueOnce({ + ok: true, + json: vi.fn().mockResolvedValue({ + id: "session-1", + revision: 4, + download_url: + "https://storage.example/task-sessions/org/task/session/revisions/4-sync-1.jsonl?signature=abc", + }), + }); + + await expect( + client.syncTaskSession( + "task-1", + "run-1", + "sandbox-1", + 3, + '{"type":"session"}\n', + ), + ).resolves.toBe(4); + }); + + it("rejects an ambiguous finalize when a competing revision was promoted", async () => { + const client = new PostHogAPIClient({ + apiUrl: "https://app.posthog.com", + getApiKey: vi.fn().mockResolvedValue("token"), + projectId: 7, + }); + mockFetch + .mockResolvedValueOnce({ + ok: true, + json: vi.fn().mockResolvedValue({ + id: "session-1", + sync_id: "sync-1", + upload: { + url: "https://storage.example/upload", + fields: { + key: "task-sessions/org/task/session/uploads/4-sync-1.jsonl", + }, + }, + }), + }) + .mockResolvedValueOnce({ ok: true }) + .mockResolvedValueOnce({ + ok: false, + status: 409, + json: vi.fn().mockResolvedValue({ error: "Stale revision" }), + }) + .mockResolvedValueOnce({ + ok: true, + json: vi.fn().mockResolvedValue({ + id: "session-1", + revision: 4, + download_url: + "https://storage.example/task-sessions/org/task/session/revisions/4-sync-2.jsonl?signature=abc", + }), + }); + + await expect( + client.syncTaskSession( + "task-1", + "run-1", + "sandbox-1", + 3, + '{"type":"session"}\n', + ), + ).rejects.toThrow("Stale revision"); + }); + it("returns only the artifacts created by the current upload request", async () => { const client = new PostHogAPIClient({ apiUrl: "https://app.posthog.com", diff --git a/packages/agent/src/posthog-api.ts b/packages/agent/src/posthog-api.ts index b3596fe424..09623c686d 100644 --- a/packages/agent/src/posthog-api.ts +++ b/packages/agent/src/posthog-api.ts @@ -1,3 +1,4 @@ +import type { StoredLogEntry } from "@posthog/shared"; import packageJson from "../package.json" with { type: "json" }; import type { ArtifactType, @@ -40,6 +41,18 @@ export interface PreparedTaskArtifactUpload { presigned_post: { url: string; fields: Record }; } +export interface TaskSessionStorageAccess { + id: string; + download_url: string; + revision: number; +} + +export interface TaskSessionSyncUpload { + id: string; + sync_id: string; + upload: { url: string; fields: Record }; +} + export interface TaskArtifactFinalizeUploadPayload { id: string; name: string; @@ -210,10 +223,115 @@ export class PostHogAPIClient { ); } + async getTaskSession( + taskId: string, + runId: string, + ): Promise { + const teamId = this.getTeamId(); + return this.apiRequest( + `/api/projects/${teamId}/tasks/${taskId}/runs/${runId}/task_session/`, + ); + } + + async downloadTaskSession(access: TaskSessionStorageAccess): Promise { + const response = await fetch(access.download_url, { + signal: AbortSignal.timeout(30_000), + }); + if (response.status === 404) { + return ""; + } + if (!response.ok) { + throw new Error( + `Failed to download task session: [${response.status}] ${response.statusText}`, + ); + } + return response.text(); + } + + async syncTaskSession( + taskId: string, + runId: string, + sandboxId: string, + expectedRevision: number, + content: string, + ): Promise { + const teamId = this.getTeamId(); + const prepared = await this.apiRequest( + `/api/projects/${teamId}/tasks/${taskId}/runs/${runId}/task_session_sync_prepare/`, + { + method: "POST", + body: JSON.stringify({ + sandbox_id: sandboxId, + expected_revision: expectedRevision, + }), + signal: AbortSignal.timeout(30_000), + }, + ); + const form = new FormData(); + for (const [key, value] of Object.entries(prepared.upload.fields)) { + form.append(key, value); + } + form.append("file", new Blob([content]), "session.jsonl"); + const uploadResponse = await fetch(prepared.upload.url, { + method: "POST", + body: form, + signal: AbortSignal.timeout(30_000), + }); + if (!uploadResponse.ok) { + throw new Error( + `Failed to upload task session: [${uploadResponse.status}] ${uploadResponse.statusText}`, + ); + } + + try { + const result = await this.apiRequest<{ id: string; revision: number }>( + `/api/projects/${teamId}/tasks/${taskId}/runs/${runId}/task_session_sync/`, + { + method: "POST", + body: JSON.stringify({ + sandbox_id: sandboxId, + sync_id: prepared.sync_id, + expected_revision: expectedRevision, + }), + signal: AbortSignal.timeout(30_000), + }, + ); + return result.revision; + } catch (error) { + const current = await this.getTaskSession(taskId, runId); + const uploadStoragePath = prepared.upload.fields.key; + const promotedStoragePath = uploadStoragePath?.replace( + "/uploads/", + "/revisions/", + ); + if ( + current.id === prepared.id && + current.revision === expectedRevision + 1 && + promotedStoragePath && + this.isTaskSessionStoragePath(current.download_url, promotedStoragePath) + ) { + return current.revision; + } + throw error; + } + } + + private isTaskSessionStoragePath( + downloadUrl: string, + storagePath: string, + ): boolean { + try { + const pathname = decodeURIComponent(new URL(downloadUrl).pathname); + return pathname.endsWith(`/${storagePath}`); + } catch { + return false; + } + } + async appendTaskRunLog( taskId: string, runId: string, - entries: StoredEntry[], + entries: (StoredEntry | StoredLogEntry)[], ): Promise { const teamId = this.getTeamId(); return this.apiRequest( diff --git a/packages/agent/src/server/bin.ts b/packages/agent/src/server/bin.ts index 1bdfef445e..14e0626c6b 100644 --- a/packages/agent/src/server/bin.ts +++ b/packages/agent/src/server/bin.ts @@ -1,15 +1,19 @@ #!/usr/bin/env node +import { realpathSync } from "node:fs"; +import { dirname, resolve } from "node:path"; import { Command } from "commander"; import { z } from "zod/v4"; import { isSupportedReasoningEffort } from "../adapters/reasoning-effort"; import { DEFAULT_POSTHOG_EXEC_PERMISSION_REGEX_SOURCE } from "../posthog-exec-permission"; import { AgentServer } from "./agent-server"; +import { PiAgentServer } from "./pi-agent-server"; import { claudeCodeConfigSchema, mcpServersSchema, posthogExecPermissionRegexSchema, relayMcpServerNamesSchema, } from "./schemas"; +import type { AgentServerConfig } from "./types"; const envSchema = z.object({ JWT_PUBLIC_KEY: z @@ -33,6 +37,8 @@ const envSchema = z.object({ }) .regex(/^\d+$/, "POSTHOG_PROJECT_ID must be a numeric string") .transform((val) => parseInt(val, 10)), + POSTHOG_AGENT_PROTOCOL: z.enum(["acp", "pi"]).optional(), + POSTHOG_SANDBOX_ID: z.string().min(1).optional(), POSTHOG_CODE_RUNTIME_ADAPTER: z.enum(["claude", "codex"]).optional(), POSTHOG_CODE_MODEL: z.string().optional(), POSTHOG_CODE_REASONING_EFFORT: z @@ -215,7 +221,7 @@ program ); } - const server = new AgentServer({ + const serverConfig: AgentServerConfig = { port: parseInt(options.port, 10), agentStateDir: env.POSTHOG_AGENT_STATE_DIR, jwtPublicKey: env.JWT_PUBLIC_KEY, @@ -233,6 +239,7 @@ program mode, taskId: options.taskId, runId: options.runId, + sandboxId: env.POSTHOG_SANDBOX_ID, createPr, autoPublish, mcpServers, @@ -241,10 +248,19 @@ program baseBranch: options.baseBranch, claudeCode, allowedDomains, + protocol: env.POSTHOG_AGENT_PROTOCOL, + piRpcHostPath: resolve( + dirname(realpathSync(process.argv[1])), + "../pi/rpc-host.js", + ), runtimeAdapter: env.POSTHOG_CODE_RUNTIME_ADAPTER, model: env.POSTHOG_CODE_MODEL, reasoningEffort: env.POSTHOG_CODE_REASONING_EFFORT, - }); + }; + const server = + env.POSTHOG_AGENT_PROTOCOL === "pi" + ? new PiAgentServer(serverConfig) + : new AgentServer(serverConfig); process.on("SIGINT", async () => { await server.stop(); diff --git a/packages/agent/src/server/pi-agent-server.test.ts b/packages/agent/src/server/pi-agent-server.test.ts new file mode 100644 index 0000000000..2fba7922b2 --- /dev/null +++ b/packages/agent/src/server/pi-agent-server.test.ts @@ -0,0 +1,376 @@ +import { mkdtemp, rm, writeFile } from "node:fs/promises"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { describe, expect, it, vi } from "vitest"; +import { PiAgentServer } from "./pi-agent-server"; +import type { AgentServerConfig } from "./types"; + +function config(): AgentServerConfig { + return { + port: 0, + jwtPublicKey: "public-key", + apiUrl: "https://us.posthog.com", + apiKey: "token", + projectId: 1, + mode: "interactive", + taskId: "task-1", + runId: "run-1", + sandboxId: "sandbox-1", + }; +} + +describe("PiAgentServer", () => { + it("persists translated Pi events at the turn boundary", async () => { + const appendTaskRunLog = vi.fn(async () => ({})); + const server = new PiAgentServer(config()) as unknown as { + posthogAPI: { appendTaskRunLog: typeof appendTaskRunLog }; + handleEvent(event: Record): void; + logFlushQueue: Promise; + }; + server.posthogAPI.appendTaskRunLog = appendTaskRunLog; + + server.handleEvent({ + type: "user_message", + timestamp: 1, + content: [{ type: "text", text: "hello" }], + }); + server.handleEvent({ type: "turn_completed", timestamp: 2 }); + await server.logFlushQueue; + + expect(appendTaskRunLog).toHaveBeenCalledWith("task-1", "run-1", [ + { + type: "pi_event", + timestamp: expect.any(String), + event: { + type: "user_message", + timestamp: 1, + content: [{ type: "text", text: "hello" }], + }, + }, + { + type: "pi_event", + timestamp: expect.any(String), + event: { type: "turn_completed", timestamp: 2 }, + }, + ]); + }); + + it("uses native Pi prompt for an idle cloud user message", async () => { + const prompt = vi.fn(async () => {}); + const followUp = vi.fn(async () => {}); + const server = new PiAgentServer(config()) as unknown as { + session: unknown; + executeCommand( + method: string, + params: Record, + ): Promise; + }; + server.session = { + runtime: { + client: { + getState: vi.fn(async () => ({ isStreaming: false })), + prompt, + followUp, + }, + }, + }; + + await server.executeCommand("user_message", { content: "hello" }); + + expect(prompt).toHaveBeenCalledWith("hello"); + expect(followUp).not.toHaveBeenCalled(); + }); + + it("deduplicates completed user messages by their stable messageId", async () => { + const prompt = vi.fn(async () => {}); + const server = new PiAgentServer(config()) as unknown as { + app: { + request(path: string, init: RequestInit): Promise; + }; + authenticate(): { task_id: string; run_id: string }; + session: unknown; + }; + server.authenticate = () => ({ task_id: "task-1", run_id: "run-1" }); + server.session = { + payload: { run_id: "run-1" }, + runtime: { + client: { + getState: vi.fn(async () => ({ isStreaming: false })), + prompt, + }, + }, + }; + const command = (id: number) => + server.app.request("/command", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ + jsonrpc: "2.0", + id, + method: "user_message", + params: { content: "hello", messageId: "message-1" }, + }), + }); + + const firstResponse = await command(1); + const retryResponse = await command(2); + + expect(firstResponse.status).toBe(200); + expect(retryResponse.status).toBe(200); + expect(prompt).toHaveBeenCalledOnce(); + expect(prompt).toHaveBeenCalledWith("hello"); + }); + + it("deduplicates concurrent in-flight user-message deliveries", async () => { + let resolvePrompt: (() => void) | undefined; + const prompt = vi.fn( + () => + new Promise((resolve) => { + resolvePrompt = resolve; + }), + ); + const server = new PiAgentServer(config()) as unknown as { + session: unknown; + executeCommand( + method: string, + params: Record, + ): Promise; + }; + server.session = { + runtime: { + client: { + getState: vi.fn(async () => ({ isStreaming: false })), + prompt, + }, + }, + }; + const params = { content: "hello", messageId: "message-1" }; + + const firstDelivery = server.executeCommand("user_message", params); + const concurrentDelivery = server.executeCommand("user_message", params); + await vi.waitFor(() => expect(prompt).toHaveBeenCalledOnce()); + resolvePrompt?.(); + await Promise.all([firstDelivery, concurrentDelivery]); + + expect(prompt).toHaveBeenCalledOnce(); + }); + + it("allows a failed user-message delivery to be retried", async () => { + const prompt = vi + .fn() + .mockRejectedValueOnce(new Error("delivery failed")) + .mockResolvedValueOnce(undefined); + const server = new PiAgentServer(config()) as unknown as { + session: unknown; + executeCommand( + method: string, + params: Record, + ): Promise; + }; + server.session = { + runtime: { + client: { + getState: vi.fn(async () => ({ isStreaming: false })), + prompt, + }, + }, + }; + const params = { content: "hello", messageId: "message-1" }; + + await expect(server.executeCommand("user_message", params)).rejects.toThrow( + "delivery failed", + ); + await expect( + server.executeCommand("user_message", params), + ).resolves.toBeUndefined(); + + expect(prompt).toHaveBeenCalledTimes(2); + }); + + it("bounds completed user-message deliveries without evicting recent IDs", async () => { + const prompt = vi.fn(async () => {}); + const server = new PiAgentServer(config()) as unknown as { + session: unknown; + completedUserMessageDeliveries: Map; + executeCommand( + method: string, + params: Record, + ): Promise; + }; + server.session = { + runtime: { + client: { + getState: vi.fn(async () => ({ isStreaming: false })), + prompt, + }, + }, + }; + + for (let index = 0; index <= 500; index++) { + await server.executeCommand("user_message", { + content: `message ${index}`, + messageId: `message-${index}`, + }); + } + + expect(server.completedUserMessageDeliveries.size).toBe(500); + expect(server.completedUserMessageDeliveries.has("message-0")).toBe(false); + expect(server.completedUserMessageDeliveries.has("message-1")).toBe(true); + + await server.executeCommand("user_message", { + content: "recent retry", + messageId: "message-1", + }); + + expect(prompt).toHaveBeenCalledTimes(501); + }); + + it("does not install an SSE controller canceled during initialization", async () => { + let finishInitialization: (() => void) | undefined; + const initializationGate = new Promise((resolve) => { + finishInitialization = resolve; + }); + const controller = { send: vi.fn(), close: vi.fn() }; + const payload = { task_id: "task-1", run_id: "run-1" }; + type TestController = typeof controller; + type TestPayload = typeof payload; + const server = new PiAgentServer(config()) as unknown as { + session: { + payload: TestPayload; + sseController: TestController | null; + } | null; + createSession(sessionPayload: TestPayload): Promise; + initializeSession( + sessionPayload: TestPayload, + sseController: TestController, + ): Promise; + cancelSseController(sseController: TestController): void; + }; + server.createSession = vi.fn(async (sessionPayload) => { + await initializationGate; + server.session = { payload: sessionPayload, sseController: null }; + }); + + const initialization = server.initializeSession(payload, controller); + server.cancelSseController(controller); + finishInitialization?.(); + await initialization; + + expect(server.session?.sseController).toBeNull(); + expect(controller.send).not.toHaveBeenCalled(); + }); + + it("preserves a replacement SSE controller when the old stream cancels", () => { + const oldController = { send: vi.fn(), close: vi.fn() }; + const replacementController = { send: vi.fn(), close: vi.fn() }; + const server = new PiAgentServer(config()) as unknown as { + session: { sseController: typeof replacementController } | null; + cancelSseController(controller: typeof oldController): void; + }; + server.session = { sseController: replacementController }; + + server.cancelSseController(oldController); + + expect(server.session?.sseController).toBe(replacementController); + + server.cancelSseController(replacementController); + + expect(server.session?.sseController).toBeNull(); + }); + + it("forwards native Pi RPC commands without redefining operations", async () => { + const send = vi.fn(async () => ({ + type: "response", + command: "set_follow_up_mode", + success: true, + })); + const server = new PiAgentServer(config()) as unknown as { + session: unknown; + executeCommand( + method: string, + params: Record, + ): Promise; + }; + server.session = { runtime: { client: { send } } }; + const command = { + type: "set_follow_up_mode", + mode: "one-at-a-time", + }; + + const response = await server.executeCommand("pi/rpc", { command }); + + expect(send).toHaveBeenCalledWith(command); + expect(response).toEqual({ + type: "response", + command: "set_follow_up_mode", + success: true, + }); + }); + + it("waits for Pi to create the native session file before syncing", async () => { + const directory = await mkdtemp(join(tmpdir(), "pi-session-sync-")); + const syncTaskSession = vi.fn(async () => 1); + const server = new PiAgentServer(config()) as unknown as { + sessionFile: string; + posthogAPI: { syncTaskSession: typeof syncTaskSession }; + syncTaskSession(): Promise; + }; + server.sessionFile = join(directory, "not-created.jsonl"); + server.posthogAPI = { syncTaskSession }; + + await server.syncTaskSession(); + + expect(syncTaskSession).not.toHaveBeenCalled(); + await rm(directory, { recursive: true }); + }); + + it("syncs changed native session JSONL to durable task storage", async () => { + const directory = await mkdtemp(join(tmpdir(), "pi-session-sync-")); + const sessionFile = join(directory, "session.jsonl"); + const content = '{"type":"session"}\n'; + await writeFile(sessionFile, content); + const syncTaskSession = vi.fn(async () => 1); + const server = new PiAgentServer(config()) as unknown as { + sessionFile: string; + posthogAPI: { syncTaskSession: typeof syncTaskSession }; + syncTaskSession(): Promise; + }; + server.sessionFile = sessionFile; + server.posthogAPI = { syncTaskSession }; + + await server.syncTaskSession(); + await server.syncTaskSession(); + + expect(syncTaskSession).toHaveBeenCalledOnce(); + expect(syncTaskSession).toHaveBeenCalledWith( + "task-1", + "run-1", + "sandbox-1", + 0, + content, + ); + await rm(directory, { recursive: true }); + }); + + it("publishes runtime-neutral Pi conversation events", () => { + const send = vi.fn(); + const server = new PiAgentServer(config()) as unknown as { + session: unknown; + handleEvent(event: unknown): void; + }; + server.session = { sseController: { send } }; + + server.handleEvent({ + type: "assistant_message_chunk", + timestamp: 1, + content: { type: "text", text: "hello" }, + }); + + expect(send).toHaveBeenCalledWith( + expect.objectContaining({ + type: "pi_event", + event: expect.objectContaining({ type: "assistant_message_chunk" }), + }), + ); + }); +}); diff --git a/packages/agent/src/server/pi-agent-server.ts b/packages/agent/src/server/pi-agent-server.ts new file mode 100644 index 0000000000..fe6750dc42 --- /dev/null +++ b/packages/agent/src/server/pi-agent-server.ts @@ -0,0 +1,651 @@ +import { access, mkdir, readFile, writeFile } from "node:fs/promises"; +import { join } from "node:path"; +import type { ServerType } from "@hono/node-server"; +import { serve } from "@hono/node-server"; +import type { AgentConversationEvent, StoredLogEntry } from "@posthog/shared"; +import { Hono } from "hono"; +import { z } from "zod/v4"; +import { createPiRpcClient, type PiRpcClient } from "../pi/rpc-client"; +import { + piRpcCommandSchema, + type RpcCommand, + sendPiRpcCommand, +} from "../pi/rpc-transport"; +import { PiRuntime } from "../pi/runtime"; +import { PostHogAPIClient } from "../posthog-api"; +import { Logger } from "../utils/logger"; +import { TaskRunEventStreamSender } from "./event-stream-sender"; +import { type JwtPayload, JwtValidationError, validateJwt } from "./jwt"; +import { jsonRpcRequestSchema } from "./schemas"; +import type { AgentServerConfig } from "./types"; + +interface SseController { + send(data: unknown): void; + close(): void; +} + +interface PiCloudSession { + payload: JwtPayload; + runtime: PiRuntime; + sseController: SseController | null; + unsubscribe: () => void; +} + +const SESSION_SYNC_INTERVAL_MS = 5_000; +const COMPLETED_USER_MESSAGE_DELIVERY_LIMIT = 500; +const emptySchema = z.object({}); + +const userMessageCommandSchema = z.object({ + content: z.string().min(1), + messageId: z.string().min(1).optional(), +}); + +const commandSchemas = { + user_message: userMessageCommandSchema, + cancel: emptySchema, + "pi/rpc": z.object({ command: piRpcCommandSchema }), +} as const; + +type PiCommandMethod = keyof typeof commandSchemas; + +export class PiAgentServer { + private readonly app: Hono; + private readonly logger = new Logger({ + debug: true, + prefix: "[PiAgentServer]", + }); + private readonly posthogAPI: PostHogAPIClient; + private readonly eventStreamSender: TaskRunEventStreamSender | null; + private server: ServerType | null = null; + private session: PiCloudSession | null = null; + private initializationPromise: Promise | null = null; + private pendingEvents: Record[] = []; + private sessionReadyBootMs?: number; + private sessionInitMs?: number; + private sessionFile: string | null = null; + private lastSyncedSessionContent = ""; + private sessionRevision = 0; + private sessionSyncInterval: ReturnType | null = null; + private sessionSyncQueue: Promise = Promise.resolve(); + private pendingLogEntries: StoredLogEntry[] = []; + private logFlushQueue: Promise = Promise.resolve(); + private readonly canceledSseControllers = new WeakSet(); + private readonly userMessageDeliveries = new Map>(); + private readonly completedUserMessageDeliveries = new Map(); + + constructor(private readonly config: AgentServerConfig) { + this.posthogAPI = new PostHogAPIClient({ + apiUrl: config.apiUrl, + projectId: config.projectId, + getApiKey: () => config.apiKey, + userAgent: `posthog/pi-cloud`, + }); + this.eventStreamSender = config.eventIngestToken + ? new TaskRunEventStreamSender({ + apiUrl: config.apiUrl, + eventIngestBaseUrl: config.eventIngestBaseUrl, + keepProxyStreamOpen: config.eventIngestKeepStreamOpen, + projectId: config.projectId, + taskId: config.taskId, + runId: config.runId, + token: config.eventIngestToken, + logger: this.logger.child("EventIngest"), + streamWindowMs: config.eventIngestStreamWindowMs, + }) + : null; + this.app = this.createApp(); + } + + async start(): Promise { + await new Promise((resolve) => { + this.server = serve( + { fetch: this.app.fetch, port: this.config.port }, + () => resolve(), + ); + }); + + const payload: JwtPayload = { + task_id: this.config.taskId, + run_id: this.config.runId, + team_id: this.config.projectId, + user_id: 0, + distinct_id: "pi-agent-server", + mode: this.config.mode, + }; + await this.initializeSession(payload, null); + } + + async stop(): Promise { + const session = this.session; + if (this.sessionSyncInterval) { + clearInterval(this.sessionSyncInterval); + this.sessionSyncInterval = null; + } + if (session) { + await session.runtime.client.abort().catch(() => undefined); + await session.runtime.client.waitForIdle(5_000).catch(() => undefined); + await this.syncTaskSession().catch((error) => + this.logger.error("Failed to sync Pi session during shutdown", error), + ); + session.unsubscribe(); + await session.runtime.client.stop(); + } + this.session = null; + await this.flushConversationLog().catch((error) => + this.logger.error("Failed to persist Pi events during shutdown", error), + ); + await this.eventStreamSender?.stop(); + this.server?.close(); + this.server = null; + } + + async reportFatalError(error: unknown): Promise { + const message = error instanceof Error ? error.message : String(error); + this.broadcast({ + type: "pi_event", + timestamp: new Date().toISOString(), + event: { + type: "runtime_error", + timestamp: Date.now(), + errorType: "agent_server_crash", + message, + } satisfies AgentConversationEvent, + }); + await Promise.all([ + this.syncTaskSession(), + this.flushConversationLog(), + ]).catch((syncError) => + this.logger.error("Failed to persist crashed Pi session", syncError), + ); + await this.posthogAPI + .updateTaskRun(this.config.taskId, this.config.runId, { + status: "failed", + error_message: `Pi agent server crashed: ${message}`, + }) + .catch(() => undefined); + await this.eventStreamSender?.stop(); + } + + private createApp(): Hono { + const app = new Hono(); + + app.get("/health", (context) => + context.json({ + status: "ok", + hasSession: this.session !== null, + bootMs: this.sessionReadyBootMs, + sessionInitMs: this.sessionInitMs, + }), + ); + + app.get("/events", async (context) => { + let payload: JwtPayload; + try { + payload = this.authenticate(context.req.header.bind(context.req)); + } catch (error) { + return context.json( + { error: error instanceof Error ? error.message : "Invalid token" }, + 401, + ); + } + + const encoder = new TextEncoder(); + let keepalive: ReturnType | null = null; + let sseController: SseController | null = null; + const stream = new ReadableStream({ + start: async (controller) => { + sseController = { + send: (data) => + controller.enqueue( + encoder.encode(`data: ${JSON.stringify(data)}\n\n`), + ), + close: () => controller.close(), + }; + keepalive = setInterval(() => { + controller.enqueue(encoder.encode(": keepalive\n\n")); + }, 25_000); + await this.initializeSession(payload, sseController); + if (this.session?.sseController !== sseController) { + return; + } + this.replayPendingEvents(); + sseController.send({ type: "connected", run_id: payload.run_id }); + }, + cancel: () => { + if (keepalive) { + clearInterval(keepalive); + } + this.cancelSseController(sseController); + }, + }); + + return new Response(stream, { + headers: { + "Content-Type": "text/event-stream", + "Cache-Control": "no-cache", + Connection: "keep-alive", + }, + }); + }); + + app.post("/command", async (context) => { + let payload: JwtPayload; + try { + payload = this.authenticate(context.req.header.bind(context.req)); + } catch (error) { + return context.json( + { error: error instanceof Error ? error.message : "Invalid token" }, + 401, + ); + } + if (!this.session || this.session.payload.run_id !== payload.run_id) { + return context.json({ error: "No active session for this run" }, 400); + } + + const request = jsonRpcRequestSchema.safeParse( + await context.req.json().catch(() => null), + ); + if (!request.success) { + return context.json({ error: "Invalid JSON-RPC request" }, 400); + } + + const method = request.data.method as PiCommandMethod; + const schema = commandSchemas[method]; + if (!schema) { + return context.json({ + jsonrpc: "2.0", + id: request.data.id, + error: { + code: -32601, + message: `Unknown method: ${request.data.method}`, + }, + }); + } + const params = schema.safeParse(request.data.params ?? {}); + if (!params.success) { + return context.json({ + jsonrpc: "2.0", + id: request.data.id, + error: { code: -32602, message: params.error.message }, + }); + } + + try { + const result = await this.executeCommand( + method, + params.data as Record, + ); + return context.json({ jsonrpc: "2.0", id: request.data.id, result }); + } catch (error) { + return context.json({ + jsonrpc: "2.0", + id: request.data.id, + error: { + code: -32000, + message: error instanceof Error ? error.message : "Unknown error", + }, + }); + } + }); + + return app; + } + + private async initializeSession( + payload: JwtPayload, + sseController: SseController | null, + ): Promise { + if (this.session?.payload.run_id === payload.run_id) { + this.installSseController(sseController); + return; + } + if (this.initializationPromise) { + await this.initializationPromise; + this.installSseController(sseController); + return; + } + + const initializationPromise = this.createSession(payload); + this.initializationPromise = initializationPromise; + try { + await initializationPromise; + } finally { + if (this.initializationPromise === initializationPromise) { + this.initializationPromise = null; + } + } + this.installSseController(sseController); + } + + private async createSession(payload: JwtPayload): Promise { + const startedAt = Date.now(); + await this.waitForRepoReady(); + const taskRun = await this.posthogAPI.getTaskRun( + payload.task_id, + payload.run_id, + ); + const task = await this.posthogAPI.getTask(payload.task_id); + const state = (taskRun.state ?? {}) as Record; + const cwd = this.config.repositoryPath ?? "/tmp/workspace"; + if (!this.config.sandboxId) { + throw new Error("Pi task session persistence requires a sandbox ID"); + } + const sessionStorage = await this.posthogAPI.getTaskSession( + payload.task_id, + payload.run_id, + ); + const persistedSessionContent = + await this.posthogAPI.downloadTaskSession(sessionStorage); + const sessionDir = join("/tmp", "posthog-pi-sessions", sessionStorage.id); + await mkdir(sessionDir, { recursive: true }); + const restoredSessionFile = persistedSessionContent + ? join(sessionDir, "session.jsonl") + : undefined; + if (restoredSessionFile) { + await writeFile(restoredSessionFile, persistedSessionContent, "utf8"); + } + this.lastSyncedSessionContent = persistedSessionContent; + this.sessionRevision = sessionStorage.revision; + + const client = createPiRpcClient({ + cliPath: this.config.piRpcHostPath, + cwd, + sessionDir, + model: this.config.model, + sessionFile: restoredSessionFile, + providerOptions: { + apiKey: this.config.apiKey, + baseUrl: this.posthogAPI.getLlmGatewayUrl(), + }, + }); + const runtime = new PiRuntime(client); + const unsubscribeConversation = runtime.onConversationEvent((event) => + this.handleEvent(event), + ); + const unsubscribeRuntime = runtime.onRuntimeEvent((event) => { + if (event.type === "agent_settled") { + void Promise.all([ + this.syncTaskSession(), + this.flushConversationLog(), + ]).catch((error) => + this.logger.error("Failed to persist settled Pi turn", error), + ); + } + }); + await client.start(); + const runtimeState = await client.getState(); + this.sessionFile = runtimeState.sessionFile ?? restoredSessionFile ?? null; + const unsubscribe = () => { + unsubscribeConversation(); + unsubscribeRuntime(); + }; + + this.session = { payload, runtime, sseController: null, unsubscribe }; + await this.syncTaskSession(); + this.sessionSyncInterval = setInterval(() => { + void this.syncTaskSession().catch((error) => + this.logger.error("Failed to sync active Pi session", error), + ); + }, SESSION_SYNC_INTERVAL_MS); + this.sessionReadyBootMs = Math.round(process.uptime() * 1000); + this.sessionInitMs = Date.now() - startedAt; + await this.posthogAPI.updateTaskRun(payload.task_id, payload.run_id, { + status: "in_progress", + }); + this.broadcast({ + type: "pi_run_started", + timestamp: new Date().toISOString(), + taskId: payload.task_id, + runId: payload.run_id, + }); + + const pendingMessage = + typeof state.pending_user_message === "string" + ? state.pending_user_message + : null; + const prompt = pendingMessage?.trim() || task.description?.trim(); + const prewarmed = state.prewarmed === true; + if (prompt && (!prewarmed || pendingMessage)) { + await client.prompt(prompt); + } + } + + private handleEvent(event: AgentConversationEvent): void { + this.broadcast({ + type: "pi_event", + timestamp: new Date().toISOString(), + event, + }); + } + + private async executeCommand( + method: PiCommandMethod, + params: Record, + ): Promise { + const runtime = this.session?.runtime; + if (!runtime) { + throw new Error("No active Pi runtime"); + } + const client = runtime.client; + switch (method) { + case "user_message": + return this.deliverUserMessage(client, params); + case "cancel": + return client.abort(); + case "pi/rpc": + return sendPiRpcCommand(client, params.command as RpcCommand); + } + } + + private async deliverUserMessage( + client: PiRpcClient, + params: Record, + ): Promise { + const messageId = + typeof params.messageId === "string" ? params.messageId : null; + if (!messageId) { + return this.dispatchUserMessage(client, String(params.content)); + } + + if (this.completedUserMessageDeliveries.has(messageId)) { + return this.completedUserMessageDeliveries.get(messageId); + } + + const existingDelivery = this.userMessageDeliveries.get(messageId); + if (existingDelivery) { + return existingDelivery; + } + + const delivery = this.dispatchUserMessage(client, String(params.content)); + this.userMessageDeliveries.set(messageId, delivery); + try { + const result = await delivery; + if (this.userMessageDeliveries.get(messageId) === delivery) { + this.userMessageDeliveries.delete(messageId); + this.completedUserMessageDeliveries.set(messageId, result); + this.evictCompletedUserMessageDeliveries(); + } + return result; + } catch (error) { + if (this.userMessageDeliveries.get(messageId) === delivery) { + this.userMessageDeliveries.delete(messageId); + } + throw error; + } + } + + private async dispatchUserMessage( + client: PiRpcClient, + content: string, + ): Promise { + const state = await client.getState(); + if (state.isStreaming) { + return client.followUp(content); + } + return client.prompt(content); + } + + private evictCompletedUserMessageDeliveries(): void { + while ( + this.completedUserMessageDeliveries.size > + COMPLETED_USER_MESSAGE_DELIVERY_LIMIT + ) { + const oldestMessageId = this.completedUserMessageDeliveries + .keys() + .next().value; + if (oldestMessageId === undefined) { + return; + } + this.completedUserMessageDeliveries.delete(oldestMessageId); + } + } + + private installSseController(sseController: SseController | null): void { + if (sseController && !this.canceledSseControllers.has(sseController)) { + if (this.session) { + this.session.sseController = sseController; + } + } + } + + private cancelSseController(sseController: SseController | null): void { + if (!sseController) { + return; + } + this.canceledSseControllers.add(sseController); + if (this.session?.sseController === sseController) { + this.session.sseController = null; + } + } + + private syncTaskSession(): Promise { + const sync = this.sessionSyncQueue.then(async () => { + if (!this.sessionFile) { + return; + } + + let content: string; + try { + content = await readFile(this.sessionFile, "utf8"); + } catch (error) { + if ((error as NodeJS.ErrnoException).code === "ENOENT") { + return; + } + throw error; + } + if (content === this.lastSyncedSessionContent) { + return; + } + + if (!this.config.sandboxId) { + throw new Error("Pi task session persistence requires a sandbox ID"); + } + this.sessionRevision = await this.posthogAPI.syncTaskSession( + this.config.taskId, + this.config.runId, + this.config.sandboxId, + this.sessionRevision, + content, + ); + this.lastSyncedSessionContent = content; + }); + this.sessionSyncQueue = sync.catch(() => undefined); + return sync; + } + + private broadcast(event: Record): void { + if (event.type === "pi_event" || event.type === "pi_run_started") { + this.pendingLogEntries.push({ + type: event.type, + timestamp: + typeof event.timestamp === "string" ? event.timestamp : undefined, + event: + event.type === "pi_event" + ? (event.event as AgentConversationEvent) + : undefined, + }); + if ( + event.type === "pi_run_started" || + (event.event as { type?: string } | undefined)?.type === + "turn_completed" + ) { + void this.flushConversationLog().catch((error) => + this.logger.error("Failed to persist Pi conversation events", error), + ); + } + } + + this.eventStreamSender?.enqueue(event); + if (this.session?.sseController) { + this.session.sseController.send(event); + } else { + this.pendingEvents.push(event); + } + } + + private flushConversationLog(): Promise { + if (this.pendingLogEntries.length === 0) { + return this.logFlushQueue; + } + + const entries = this.pendingLogEntries; + this.pendingLogEntries = []; + const flush = this.logFlushQueue + .then(() => + this.posthogAPI.appendTaskRunLog( + this.config.taskId, + this.config.runId, + entries, + ), + ) + .then(() => undefined) + .catch((error) => { + this.pendingLogEntries = [...entries, ...this.pendingLogEntries]; + throw error; + }); + this.logFlushQueue = flush.catch(() => undefined); + return flush; + } + + private replayPendingEvents(): void { + const controller = this.session?.sseController; + if (!controller) { + return; + } + const events = this.pendingEvents; + this.pendingEvents = []; + for (const event of events) { + controller.send(event); + } + } + + private authenticate( + getHeader: (name: string) => string | undefined, + ): JwtPayload { + const authHeader = getHeader("authorization"); + if (!authHeader?.startsWith("Bearer ")) { + throw new JwtValidationError( + "Missing authorization header", + "invalid_token", + ); + } + return validateJwt(authHeader.slice(7), this.config.jwtPublicKey); + } + + private async waitForRepoReady(): Promise { + const path = this.config.repoReadyFile; + if (!path) { + return; + } + const deadline = Date.now() + 10 * 60_000; + while (Date.now() < deadline) { + try { + await access(path); + return; + } catch { + await new Promise((resolve) => setTimeout(resolve, 250)); + } + } + throw new Error(`Repository readiness file was not created: ${path}`); + } +} diff --git a/packages/agent/src/server/types.ts b/packages/agent/src/server/types.ts index 34b94a45e1..5dfe5114f2 100644 --- a/packages/agent/src/server/types.ts +++ b/packages/agent/src/server/types.ts @@ -27,6 +27,7 @@ export interface AgentServerConfig { mode: AgentMode; taskId: string; runId: string; + sandboxId?: string; createPr?: boolean; // User-opted auto-publish: push and open a draft PR on completion even for // manual (non-automated-origin) cloud runs. createPr=false still wins. @@ -47,6 +48,8 @@ export interface AgentServerConfig { baseBranch?: string; claudeCode?: ClaudeCodeConfig; allowedDomains?: string[]; + protocol?: "acp" | "pi"; + piRpcHostPath?: string; runtimeAdapter?: Adapter; model?: string; reasoningEffort?: "low" | "medium" | "high" | "xhigh" | "max"; diff --git a/packages/agent/tsup.config.ts b/packages/agent/tsup.config.ts index 57ebcd9f6a..f84607296e 100644 --- a/packages/agent/tsup.config.ts +++ b/packages/agent/tsup.config.ts @@ -103,6 +103,21 @@ const sharedOptions = { }; export default defineConfig([ + { + entry: { + "pi/rpc-transport": "src/pi/rpc-transport.ts", + "pi/remote-rpc-client": "src/pi/remote-rpc-client.ts", + }, + format: ["esm"], + dts: false, + clean: false, + sourcemap: true, + splitting: false, + outDir: "dist", + target: "es2022", + platform: "browser", + external: ["@earendil-works/pi-ai", "@posthog/shared", "zod"], + }, { entry: [ "src/index.ts", @@ -167,9 +182,12 @@ export default defineConfig([ }, { entry: { "server/bin": "src/server/bin.ts" }, - format: ["cjs"], + format: ["esm"], dts: false, clean: false, + banner: { + js: 'import { createRequire as __createRequire } from "node:module"; const require = __createRequire(import.meta.url);', + }, ...sharedOptions, }, { diff --git a/packages/core/src/cloud-task/cloud-task-types.ts b/packages/core/src/cloud-task/cloud-task-types.ts deleted file mode 100644 index a2cd0c377e..0000000000 --- a/packages/core/src/cloud-task/cloud-task-types.ts +++ /dev/null @@ -1,69 +0,0 @@ -import type { StoredLogEntry, TaskRunStatus } from "@posthog/shared"; - -interface CloudTaskUpdateBase { - taskId: string; - runId: string; -} - -export interface CloudTaskLogsUpdate extends CloudTaskUpdateBase { - kind: "logs"; - newEntries: StoredLogEntry[]; - totalEntryCount: number; -} - -export interface CloudTaskStatusUpdate extends CloudTaskUpdateBase { - kind: "status"; - status?: TaskRunStatus; - stage?: string | null; - output?: Record | null; - errorMessage?: string | null; - branch?: string | null; - sandboxAlive?: boolean | null; -} - -export interface CloudTaskSnapshotUpdate extends CloudTaskUpdateBase { - kind: "snapshot"; - newEntries: StoredLogEntry[]; - totalEntryCount: number; - status?: TaskRunStatus; - stage?: string | null; - output?: Record | null; - errorMessage?: string | null; - branch?: string | null; - sandboxAlive?: boolean | null; -} - -export interface CloudTaskErrorUpdate extends CloudTaskUpdateBase { - kind: "error"; - errorTitle: string; - errorMessage: string; - retryable: boolean; -} - -export interface CloudPermissionOption { - kind: string; - optionId: string; - name: string; - _meta?: Record; -} - -export interface CloudTaskPermissionRequestUpdate extends CloudTaskUpdateBase { - kind: "permission_request"; - requestId: string; - toolCall: { - toolCallId: string; - title: string; - kind: string; - content?: unknown[]; - rawInput?: Record; - _meta?: Record; - }; - options: CloudPermissionOption[]; -} - -export type CloudTaskUpdatePayload = - | CloudTaskLogsUpdate - | CloudTaskStatusUpdate - | CloudTaskSnapshotUpdate - | CloudTaskErrorUpdate - | CloudTaskPermissionRequestUpdate; diff --git a/packages/core/src/cloud-task/cloud-task.ts b/packages/core/src/cloud-task/cloud-task.ts index 22b5a1ddab..19ff6d0215 100644 --- a/packages/core/src/cloud-task/cloud-task.ts +++ b/packages/core/src/cloud-task/cloud-task.ts @@ -15,8 +15,8 @@ import { TypedEventEmitter, } from "@posthog/shared"; import { ANALYTICS_EVENTS } from "@posthog/shared/analytics-events"; +import type { CloudTaskPermissionRequestUpdate } from "@posthog/shared/domain-types"; import { inject, injectable, optional, preDestroy } from "inversify"; -import type { CloudTaskPermissionRequestUpdate } from "./cloud-task-types"; import { CLOUD_TASK_AUTH, type ICloudTaskAuth, @@ -706,6 +706,10 @@ export class CloudTaskService extends TypedEventEmitter { } } + getCloudContext(): Promise<{ apiHost: string; teamId: number } | null> { + return this.auth.getCloudContext(); + } + watch(input: WatchInput): void { const key = watcherKey(input.taskId, input.runId); @@ -861,7 +865,13 @@ export class CloudTaskService extends TypedEventEmitter { status: response.status, error: errorMessage, }); - return { success: false, error: errorMessage }; + const retryable = [400, 502, 503, 504].includes(response.status); + return { + success: false, + error: errorMessage, + status: response.status, + retryable, + }; } const data = (await response.json()) as { @@ -896,7 +906,7 @@ export class CloudTaskService extends TypedEventEmitter { method: input.method, error: errorMessage, }); - return { success: false, error: errorMessage }; + return { success: false, error: errorMessage, retryable: true }; } } diff --git a/packages/core/src/cloud-task/cloudTaskClient.ts b/packages/core/src/cloud-task/cloudTaskClient.ts new file mode 100644 index 0000000000..2a86b34f04 --- /dev/null +++ b/packages/core/src/cloud-task/cloudTaskClient.ts @@ -0,0 +1,23 @@ +import type { CloudTaskUpdatePayload } from "@posthog/shared/domain-types"; +import type { SendCommandInput, SendCommandOutput } from "./schemas"; + +export const CLOUD_TASK_CLIENT = Symbol.for("posthog.cloudTask.client"); + +export interface CloudTaskClient { + getContext(): Promise<{ apiHost: string; teamId: number } | null>; + watch(input: { + taskId: string; + runId: string; + apiHost: string; + teamId: number; + }): Promise; + unwatch(taskId: string, runId: string): Promise; + subscribe( + taskId: string, + runId: string, + onUpdate: (update: CloudTaskUpdatePayload) => void, + onError: (error: unknown) => void, + onStarted: () => void, + ): () => void; + sendCommand(input: SendCommandInput): Promise; +} diff --git a/packages/core/src/cloud-task/schemas.ts b/packages/core/src/cloud-task/schemas.ts index d694e52141..7436d0a891 100644 --- a/packages/core/src/cloud-task/schemas.ts +++ b/packages/core/src/cloud-task/schemas.ts @@ -1,9 +1,13 @@ import type { TaskRunStatus } from "@posthog/shared"; +import type { CloudTaskUpdatePayload } from "@posthog/shared/domain-types"; import { z } from "zod"; -import type { CloudTaskUpdatePayload } from "./cloud-task-types"; export type { CloudTaskUpdatePayload, TaskRunStatus }; +export const cloudContextOutput = z + .object({ apiHost: z.string(), teamId: z.number() }) + .nullable(); + export const TERMINAL_STATUSES = ["completed", "failed", "cancelled"] as const; export function isTerminalStatus( @@ -65,6 +69,7 @@ export const sendCommandInput = z.object({ "permission_response", "set_config_option", "mcp_response", + "pi/rpc", ]), params: z.record(z.string(), z.unknown()).optional(), }); @@ -84,6 +89,8 @@ export const sendCommandOutput = z.object({ success: z.boolean(), result: z.unknown().optional(), error: z.string().optional(), + status: z.number().optional(), + retryable: z.boolean().optional(), }); export type SendCommandOutput = z.infer; diff --git a/packages/core/src/pi-runtime/cloudPiSessionClient.test.ts b/packages/core/src/pi-runtime/cloudPiSessionClient.test.ts new file mode 100644 index 0000000000..9a822a033a --- /dev/null +++ b/packages/core/src/pi-runtime/cloudPiSessionClient.test.ts @@ -0,0 +1,351 @@ +import type { TaskService } from "@posthog/core/task-detail/taskService"; +import type { AgentConversationEvent } from "@posthog/shared"; +import type { CloudTaskUpdatePayload } from "@posthog/shared/domain-types"; +import { describe, expect, it, vi } from "vitest"; +import type { CloudTaskClient } from "../cloud-task/cloudTaskClient"; +import { CloudPiSessionClient } from "./cloudPiSessionClient"; +import { + PiSessionController, + type PiSessionProvider, +} from "./piSessionController"; + +function createCloudTaskClient(autoStart = true) { + let onUpdate: (update: CloudTaskUpdatePayload) => void = () => {}; + let onError: (error: unknown) => void = () => {}; + let onStarted: () => void = () => {}; + const unsubscribe = vi.fn(); + const client: CloudTaskClient = { + getContext: vi.fn(async () => null), + watch: vi.fn(async () => {}), + unwatch: vi.fn(async () => {}), + subscribe: vi.fn((_taskId, _runId, handler, errorHandler, started) => { + onUpdate = handler; + onError = errorHandler; + onStarted = started; + if (autoStart) { + onStarted(); + } + return unsubscribe; + }), + sendCommand: vi.fn(async () => ({ success: false })), + }; + + return { + client, + startSubscription: () => onStarted(), + sendUpdate: (update: CloudTaskUpdatePayload) => onUpdate(update), + sendError: (error: unknown) => onError(error), + unsubscribe, + }; +} + +function context(status: "in_progress" | "completed") { + return { + taskId: "task-1", + runId: "run-1", + runStatus: status, + apiHost: "https://us.posthog.com", + teamId: 1, + }; +} + +const snapshotEvent: AgentConversationEvent = { + type: "assistant_message_chunk", + timestamp: 1, + content: { type: "text", text: "durable response" }, +}; + +describe("CloudPiSessionClient", () => { + it("waits for the native Pi readiness event before startup RPC commands", async () => { + const cloud = createCloudTaskClient(); + vi.mocked(cloud.client.sendCommand).mockResolvedValue({ + success: true, + result: { + type: "response", + command: "get_state", + success: true, + data: { isStreaming: true }, + }, + }); + const session = new CloudPiSessionClient( + cloud.client, + context("in_progress"), + ); + session.onConversationEvent(vi.fn(), vi.fn()); + + const state = session.client.getState(); + expect(cloud.client.sendCommand).not.toHaveBeenCalled(); + + cloud.sendUpdate({ + taskId: "task-1", + runId: "run-1", + kind: "logs", + newEntries: [{ type: "pi_run_started" }], + totalEntryCount: 1, + }); + + await expect(state).resolves.toMatchObject({ isStreaming: true }); + expect(cloud.client.sendCommand).toHaveBeenCalledOnce(); + }); + + it("falls back to task state readiness for older Pi servers", async () => { + vi.useFakeTimers(); + try { + const cloud = createCloudTaskClient(); + const waitUntilReady = vi.fn(async () => "in_progress" as const); + vi.mocked(cloud.client.sendCommand).mockResolvedValue({ + success: true, + result: { + type: "response", + command: "get_state", + success: true, + data: { isStreaming: true }, + }, + }); + const session = new CloudPiSessionClient(cloud.client, { + ...context("in_progress"), + waitUntilReady, + }); + + const state = session.client.getState(); + await vi.advanceTimersByTimeAsync(40_000); + + await expect(state).resolves.toMatchObject({ isStreaming: true }); + expect(waitUntilReady).toHaveBeenCalledOnce(); + expect(cloud.client.sendCommand).toHaveBeenCalledOnce(); + } finally { + vi.useRealTimers(); + } + }); + + it("waits for subscription readiness before watching and only unsubscribes on cleanup", async () => { + const cloud = createCloudTaskClient(false); + vi.mocked(cloud.client.watch).mockImplementation(async () => { + cloud.sendUpdate({ + taskId: "task-1", + runId: "run-1", + kind: "snapshot", + status: "completed", + newEntries: [{ type: "pi_event", event: snapshotEvent }], + totalEntryCount: 1, + }); + }); + const session = new CloudPiSessionClient( + cloud.client, + context("completed"), + ); + + const cleanup = session.onConversationEvent(vi.fn(), vi.fn()); + const conversation = session.getConversation(); + expect(cloud.client.watch).not.toHaveBeenCalled(); + + cloud.startSubscription(); + + await expect(conversation).resolves.toEqual([snapshotEvent]); + expect(cloud.client.watch).toHaveBeenCalledTimes(1); + + cleanup(); + expect(cloud.unsubscribe).toHaveBeenCalledTimes(1); + expect(cloud.client.unwatch).not.toHaveBeenCalled(); + }); + + it("rejects terminal history when the update subscription fails", async () => { + const cloud = createCloudTaskClient(); + const session = new CloudPiSessionClient( + cloud.client, + context("completed"), + ); + const onError = vi.fn(); + session.onConversationEvent(vi.fn(), onError); + + const conversation = session.getConversation(); + const error = new Error("subscription failed"); + cloud.sendError(error); + + await expect(conversation).rejects.toThrow("subscription failed"); + expect(onError).toHaveBeenCalledWith(error); + }); + + it("loads terminal history from the cloud snapshot without sandbox RPC", async () => { + const cloud = createCloudTaskClient(); + const session = new CloudPiSessionClient( + cloud.client, + context("completed"), + ); + const events: AgentConversationEvent[] = []; + session.onConversationEvent((event) => events.push(event), vi.fn()); + + const conversation = session.getConversation(); + cloud.sendUpdate({ + taskId: "task-1", + runId: "run-1", + kind: "snapshot", + status: "completed", + newEntries: [{ type: "pi_event", event: snapshotEvent }], + totalEntryCount: 1, + }); + + await expect(conversation).resolves.toEqual([snapshotEvent]); + await expect(session.health()).resolves.toEqual({ state: "cold" }); + await expect(session.client.getState()).resolves.toMatchObject({ + isStreaming: false, + }); + await expect(session.client.getAvailableModels()).resolves.toEqual([]); + await expect(session.client.getCommands()).resolves.toEqual([]); + expect(events).toEqual([ + snapshotEvent, + expect.objectContaining({ type: "turn_completed" }), + ]); + expect(cloud.client.sendCommand).not.toHaveBeenCalled(); + }); + + it("does not install streaming state after a terminal snapshot arrives during controller load", async () => { + const cloud = createCloudTaskClient(); + let resolveEntries: (result: { success: false; retryable: true }) => void = + () => {}; + const entries = new Promise<{ success: false; retryable: true }>( + (resolve) => { + resolveEntries = resolve; + }, + ); + vi.mocked(cloud.client.sendCommand).mockImplementation(async (input) => { + const command = input.params?.command as { type: string }; + if (command.type === "get_entries") { + return entries; + } + if (command.type === "get_state") { + return { + success: true, + result: { + type: "response", + command: "get_state", + success: true, + data: { + thinkingLevel: "off", + isStreaming: true, + isCompacting: false, + steeringMode: "all", + followUpMode: "all", + sessionId: "run-1", + autoCompactionEnabled: true, + messageCount: 1, + pendingMessageCount: 0, + }, + }, + }; + } + + return { success: false }; + }); + const session = new CloudPiSessionClient( + cloud.client, + context("in_progress"), + ); + const provider: PiSessionProvider = { + get: vi.fn(async () => session), + }; + const controller = new PiSessionController(provider, {} as TaskService); + + const connection = controller.connect("task-1"); + await vi.waitFor(() => { + expect(cloud.client.subscribe).toHaveBeenCalledTimes(1); + }); + cloud.sendUpdate({ + taskId: "task-1", + runId: "run-1", + kind: "logs", + newEntries: [{ type: "pi_run_started" }], + totalEntryCount: 1, + }); + await vi.waitFor(() => { + expect(cloud.client.sendCommand).toHaveBeenCalledTimes(1); + }); + cloud.sendUpdate({ + taskId: "task-1", + runId: "run-1", + kind: "snapshot", + status: "completed", + newEntries: [{ type: "pi_event", event: snapshotEvent }], + totalEntryCount: 1, + }); + resolveEntries({ success: false, retryable: true }); + + await connection; + + const controllerSession = controller.store.getState().sessions["task-1"]; + expect(controllerSession.events).toContain(snapshotEvent); + expect(controllerSession.status).toMatchObject({ isStreaming: false }); + }); + + it("switches readiness retries to terminal responses when the run finishes", async () => { + const cloud = createCloudTaskClient(); + const session = new CloudPiSessionClient( + cloud.client, + context("in_progress"), + ); + session.onConversationEvent(vi.fn(), vi.fn()); + vi.mocked(cloud.client.sendCommand).mockImplementation(async () => { + cloud.sendUpdate({ + taskId: "task-1", + runId: "run-1", + kind: "snapshot", + status: "completed", + newEntries: [{ type: "pi_event", event: snapshotEvent }], + totalEntryCount: 1, + }); + return { success: false, retryable: true }; + }); + cloud.sendUpdate({ + taskId: "task-1", + runId: "run-1", + kind: "logs", + newEntries: [{ type: "pi_run_started" }], + totalEntryCount: 1, + }); + + await expect(session.getConversation()).resolves.toEqual([snapshotEvent]); + expect(cloud.client.sendCommand).toHaveBeenCalledTimes(1); + }); + + it("processes reconnect snapshots and clears streaming on terminal status", async () => { + const cloud = createCloudTaskClient(); + const session = new CloudPiSessionClient( + cloud.client, + context("in_progress"), + ); + const events: AgentConversationEvent[] = []; + session.onConversationEvent((event) => events.push(event), vi.fn()); + + cloud.sendUpdate({ + taskId: "task-1", + runId: "run-1", + kind: "snapshot", + status: "in_progress", + newEntries: [{ type: "pi_event", event: snapshotEvent }], + totalEntryCount: 1, + }); + cloud.sendUpdate({ + taskId: "task-1", + runId: "run-1", + kind: "snapshot", + status: "in_progress", + newEntries: [{ type: "pi_event", event: snapshotEvent }], + totalEntryCount: 1, + }); + cloud.sendUpdate({ + taskId: "task-1", + runId: "run-1", + kind: "status", + status: "failed", + }); + + expect(events).toEqual([ + snapshotEvent, + expect.objectContaining({ type: "turn_completed" }), + ]); + await expect(session.client.abort()).rejects.toThrow( + "Cloud task run run-1 is failed", + ); + expect(cloud.client.sendCommand).not.toHaveBeenCalled(); + }); +}); diff --git a/packages/core/src/pi-runtime/cloudPiSessionClient.ts b/packages/core/src/pi-runtime/cloudPiSessionClient.ts new file mode 100644 index 0000000000..0c2e6594c4 --- /dev/null +++ b/packages/core/src/pi-runtime/cloudPiSessionClient.ts @@ -0,0 +1,371 @@ +import { + getRemotePiConversation, + type PiRemoteRpcClient, + RemotePiRpcClient, +} from "@posthog/agent/pi/remote-rpc-client"; +import type { RpcCommand } from "@posthog/agent/pi/rpc-transport"; +import type { + AgentConversationEvent, + PiRuntimeHealth, + StoredLogEntry, + TaskRunStatus, +} from "@posthog/shared"; +import type { CloudTaskUpdatePayload } from "@posthog/shared/domain-types"; +import type { CloudTaskClient } from "../cloud-task/cloudTaskClient"; +import { isTerminalStatus } from "../cloud-task/schemas"; +import type { PiSession } from "./piSessionController"; + +const readinessCommands = new Set([ + "get_state", + "get_entries", + "get_available_models", + "get_available_thinking_levels", + "get_commands", +]); + +export interface CloudPiSessionContext { + taskId: string; + runId: string; + runStatus: TaskRunStatus; + apiHost: string; + teamId: number; + waitUntilReady?: () => Promise; +} + +export class CloudPiSessionClient implements PiSession { + readonly client: PiRemoteRpcClient; + + private runStatus: TaskRunStatus; + private snapshotEvents: AgentConversationEvent[] = []; + private hasSnapshot = false; + private resolveSnapshot: () => void = () => {}; + private rejectSnapshot: (error: unknown) => void = () => {}; + private readonly snapshotReceived = new Promise((resolve, reject) => { + this.resolveSnapshot = resolve; + this.rejectSnapshot = reject; + }); + private runtimeReady = false; + private resolveRuntimeReady: () => void = () => {}; + private readonly runtimeReadyReceived = new Promise((resolve) => { + this.resolveRuntimeReady = resolve; + }); + private terminalEventSent = false; + private resolveTerminalStatus: () => void = () => {}; + private readonly terminalStatusReceived = new Promise((resolve) => { + this.resolveTerminalStatus = resolve; + }); + + constructor( + private readonly cloudTaskClient: CloudTaskClient, + private readonly context: CloudPiSessionContext, + ) { + this.runStatus = context.runStatus; + if (isTerminalStatus(this.runStatus)) { + this.resolveTerminalStatus(); + } + void this.snapshotReceived.catch(() => {}); + this.client = new RemotePiRpcClient({ + request: (command) => this.request(command), + }); + } + + health(): Promise { + if (this.runStatus === "in_progress") { + return Promise.resolve({ state: "streaming" }); + } + if (isTerminalStatus(this.runStatus)) { + return Promise.resolve({ state: "cold" }); + } + return Promise.resolve({ state: "starting" }); + } + + async getConversation(): Promise { + if (!isTerminalStatus(this.runStatus)) { + const conversation = await getRemotePiConversation(this.client); + if (!isTerminalStatus(this.runStatus)) { + return conversation; + } + } + + await this.snapshotReceived; + return this.snapshotEvents; + } + + onConversationEvent( + onEvent: (event: AgentConversationEvent) => void, + onError: (error: unknown) => void, + ): () => void { + let active = true; + const unsubscribe = this.cloudTaskClient.subscribe( + this.context.taskId, + this.context.runId, + (update) => this.handleUpdate(update, onEvent, onError), + (error) => { + if (isTerminalStatus(this.runStatus)) { + this.rejectSnapshot(error); + } + onError(error); + }, + () => { + if (!active) { + return; + } + + void this.cloudTaskClient + .watch({ + taskId: this.context.taskId, + runId: this.context.runId, + apiHost: this.context.apiHost, + teamId: this.context.teamId, + }) + .catch((error) => { + if (isTerminalStatus(this.runStatus)) { + this.rejectSnapshot(error); + } + onError(error); + }); + }, + ); + + return () => { + active = false; + unsubscribe(); + }; + } + + private handleUpdate( + update: CloudTaskUpdatePayload, + onEvent: (event: AgentConversationEvent) => void, + onError: (error: unknown) => void, + ): void { + if ( + (update.kind === "snapshot" || update.kind === "logs") && + update.newEntries.some((entry) => entry.type === "pi_run_started") + ) { + this.markRuntimeReady(); + } + + if (update.kind === "error") { + const error = new Error(update.errorMessage); + if (isTerminalStatus(this.runStatus)) { + this.rejectSnapshot(error); + } + onError(error); + return; + } + + if (update.kind === "snapshot") { + const events = this.getPiEvents(update.newEntries); + let unchangedEventCount = 0; + while ( + unchangedEventCount < events.length && + unchangedEventCount < this.snapshotEvents.length && + this.eventsEqual( + events[unchangedEventCount], + this.snapshotEvents[unchangedEventCount], + ) + ) { + unchangedEventCount += 1; + } + + this.snapshotEvents = events; + this.hasSnapshot = true; + this.resolveSnapshot(); + for (const event of events.slice(unchangedEventCount)) { + onEvent(event); + } + } else if (update.kind === "logs") { + const events = this.getPiEvents(update.newEntries); + this.snapshotEvents = [...this.snapshotEvents, ...events]; + for (const event of events) { + onEvent(event); + } + } + + if ( + (update.kind === "snapshot" || update.kind === "status") && + update.status + ) { + this.runStatus = update.status; + } + + if (isTerminalStatus(this.runStatus)) { + this.resolveTerminalStatus(); + if (!this.terminalEventSent) { + this.terminalEventSent = true; + onEvent({ type: "turn_completed", timestamp: Date.now() }); + } + } + } + + private eventsEqual( + left: AgentConversationEvent, + right: AgentConversationEvent, + ): boolean { + return JSON.stringify(left) === JSON.stringify(right); + } + + private getPiEvents(entries: StoredLogEntry[]): AgentConversationEvent[] { + const events: AgentConversationEvent[] = []; + for (const entry of entries) { + if (entry.type === "pi_event" && entry.event) { + events.push(entry.event); + } + } + return events; + } + + private async request(command: RpcCommand): Promise { + if (isTerminalStatus(this.runStatus)) { + return this.terminalResponseWhenReady(command); + } + + if (readinessCommands.has(command.type)) { + await this.waitForRuntimeReady(); + if (isTerminalStatus(this.runStatus)) { + return this.terminalResponseWhenReady(command); + } + } + + const input = { + taskId: this.context.taskId, + runId: this.context.runId, + apiHost: this.context.apiHost, + teamId: this.context.teamId, + method: "pi/rpc" as const, + params: { command }, + }; + const maxAttempts = readinessCommands.has(command.type) ? 3 : 1; + + for (let attempt = 1; attempt <= maxAttempts; attempt += 1) { + if (isTerminalStatus(this.runStatus)) { + return this.terminalResponseWhenReady(command); + } + + const result = await this.cloudTaskClient.sendCommand(input); + if (result.success) { + return result.result; + } + if (isTerminalStatus(this.runStatus)) { + return this.terminalResponseWhenReady(command); + } + + const error = result.error ?? `Pi RPC command failed: ${command.type}`; + if (attempt === maxAttempts || !result.retryable) { + throw new Error(error); + } + + await Promise.race([ + new Promise((resolve) => setTimeout(resolve, 1_000)), + this.terminalStatusReceived, + ]); + } + + throw new Error(`Pi RPC command failed: ${command.type}`); + } + + private markRuntimeReady(): void { + if (this.runtimeReady) { + return; + } + this.runtimeReady = true; + this.resolveRuntimeReady(); + } + + private async waitForRuntimeReady(): Promise { + if (this.runtimeReady || isTerminalStatus(this.runStatus)) { + return; + } + + const readiness = await new Promise<"ready" | "terminal" | "fallback">( + (resolve) => { + let settled = false; + const settle = (value: "ready" | "terminal" | "fallback") => { + if (settled) { + return; + } + settled = true; + clearTimeout(timeout); + resolve(value); + }; + const timeout = setTimeout(() => settle("fallback"), 10_000); + void this.runtimeReadyReceived.then(() => settle("ready")); + void this.terminalStatusReceived.then(() => settle("terminal")); + }, + ); + if (readiness !== "fallback" || !this.context.waitUntilReady) { + return; + } + + this.runStatus = await this.context.waitUntilReady(); + if (isTerminalStatus(this.runStatus) || this.runtimeReady) { + return; + } + + const nativeReadiness = await new Promise<"ready" | "terminal" | "legacy">( + (resolve) => { + let settled = false; + const settle = (value: "ready" | "terminal" | "legacy") => { + if (settled) { + return; + } + settled = true; + clearTimeout(timeout); + resolve(value); + }; + const timeout = setTimeout(() => settle("legacy"), 30_000); + void this.runtimeReadyReceived.then(() => settle("ready")); + void this.terminalStatusReceived.then(() => settle("terminal")); + }, + ); + if (nativeReadiness === "legacy") { + this.markRuntimeReady(); + } + } + + private async terminalResponseWhenReady( + command: RpcCommand, + ): Promise { + if (command.type === "get_entries" && !this.hasSnapshot) { + await this.snapshotReceived; + } + + return this.terminalResponse(command); + } + + private terminalResponse(command: RpcCommand): unknown { + let data: unknown; + if (command.type === "get_state") { + data = { + isStreaming: false, + isCompacting: false, + thinkingLevel: "off", + steeringMode: "all", + followUpMode: "all", + sessionId: this.context.runId, + autoCompactionEnabled: true, + messageCount: 0, + pendingMessageCount: 0, + }; + } else if (command.type === "get_available_models") { + data = { models: [] }; + } else if (command.type === "get_available_thinking_levels") { + data = { levels: [] }; + } else if (command.type === "get_commands") { + data = { commands: [] }; + } else if (command.type === "get_entries" && this.hasSnapshot) { + data = { entries: [] }; + } else { + throw new Error( + `Cloud task run ${this.context.runId} is ${this.runStatus}`, + ); + } + + return { + type: "response", + command: command.type, + success: true, + data, + }; + } +} diff --git a/packages/core/src/pi-runtime/pi-runtime.module.ts b/packages/core/src/pi-runtime/pi-runtime.module.ts index 4fd1d3a7f1..10c46b4906 100644 --- a/packages/core/src/pi-runtime/pi-runtime.module.ts +++ b/packages/core/src/pi-runtime/pi-runtime.module.ts @@ -1,7 +1,12 @@ import { ContainerModule } from "inversify"; import { PI_SESSION_CONTROLLER } from "./identifiers"; -import { PiSessionController } from "./piSessionController"; +import { + PI_SESSION_PROVIDER, + PiSessionController, +} from "./piSessionController"; +import { RoutingPiSessionProvider } from "./piSessionProvider"; export const piRuntimeModule = new ContainerModule(({ bind }) => { + bind(PI_SESSION_PROVIDER).to(RoutingPiSessionProvider).inSingletonScope(); bind(PI_SESSION_CONTROLLER).to(PiSessionController).inSingletonScope(); }); diff --git a/packages/core/src/pi-runtime/piSessionController.test.ts b/packages/core/src/pi-runtime/piSessionController.test.ts index 1c771e0d46..3dfa35ef75 100644 --- a/packages/core/src/pi-runtime/piSessionController.test.ts +++ b/packages/core/src/pi-runtime/piSessionController.test.ts @@ -1,25 +1,28 @@ +import type { PiRemoteRpcClient } from "@posthog/agent/pi/remote-rpc-client"; import type { TaskService } from "@posthog/core/task-detail/taskService"; import type { AgentConversationEvent } from "@posthog/shared"; import { describe, expect, it, vi } from "vitest"; import { - type PiSessionClient, + type PiSession, PiSessionController, + type PiSessionProvider, } from "./piSessionController"; function createController( - client = createClient(), + session = createSession(), taskService = { openTask: vi.fn(async () => ({ success: true })), } as unknown as TaskService, ): PiSessionController { - return new PiSessionController(client, taskService); + const provider: PiSessionProvider = { + get: vi.fn(async () => session), + }; + return new PiSessionController(provider, taskService); } -function createClient(): PiSessionClient { - return { - health: vi.fn(async () => ({ state: "idle" as const })), - conversation: vi.fn(async () => []), - status: vi.fn(async () => ({ +function createSession(): PiSession { + const client = { + getState: vi.fn(async () => ({ thinkingLevel: "off" as const, isStreaming: false, isCompacting: false, @@ -30,20 +33,27 @@ function createClient(): PiSessionClient { messageCount: 0, pendingMessageCount: 0, })), - availableModels: vi.fn(async () => []), - commands: vi.fn(async () => []), - subscribe: vi.fn(() => () => {}), + getAvailableModels: vi.fn(async () => []), + getAvailableThinkingLevels: vi.fn(async () => ["off" as const]), + getCommands: vi.fn(async () => []), prompt: vi.fn(async () => {}), steer: vi.fn(async () => {}), followUp: vi.fn(async () => {}), compact: vi.fn(async () => undefined), - setModel: vi.fn(async (_taskId, provider, id) => ({ provider, id })), + setModel: vi.fn(async (provider, id) => ({ provider, id })), setThinkingLevel: vi.fn(async () => {}), setSteeringMode: vi.fn(async () => {}), setFollowUpMode: vi.fn(async () => {}), bash: vi.fn(async () => undefined), abort: vi.fn(async () => {}), abortBash: vi.fn(async () => {}), + } as unknown as PiRemoteRpcClient; + + return { + client, + health: vi.fn(async () => ({ state: "idle" as const })), + getConversation: vi.fn(async () => []), + onConversationEvent: vi.fn(() => () => {}), }; } @@ -80,53 +90,170 @@ describe("PiSessionController", () => { streaming: false, mode: "steer" as const, method: "prompt" as const, - expectedArgs: ["task-1", "hello"], + expectedArgs: ["hello"], }, { text: "hello", streaming: true, mode: "steer" as const, method: "steer" as const, - expectedArgs: ["task-1", "hello"], + expectedArgs: ["hello"], }, { text: "hello", streaming: true, mode: "queue" as const, method: "followUp" as const, - expectedArgs: ["task-1", "hello"], + expectedArgs: ["hello"], }, { text: "/compact keep details", streaming: false, mode: "steer" as const, method: "compact" as const, - expectedArgs: ["task-1", "keep details"], + expectedArgs: ["keep details"], }, ])("routes submissions through $method", async (input) => { - const client = createClient(); + const client = createSession(); const controller = createController(client); await controller.submit("task-1", input.text, input.streaming, input.mode); - expect(client[input.method]).toHaveBeenCalledWith(...input.expectedArgs); + expect(client.client[input.method]).toHaveBeenCalledWith( + ...input.expectedArgs, + ); + }); + + it("keeps a connected transcript usable when a command fails", async () => { + const initialEvent: AgentConversationEvent = { + type: "user_message", + id: "message-1", + timestamp: 1, + content: [{ type: "text", text: "hello" }], + }; + const session = createSession(); + vi.mocked(session.getConversation).mockResolvedValue([initialEvent]); + vi.mocked(session.client.prompt).mockRejectedValue( + new Error("temporary command failure"), + ); + const controller = createController(session); + + await controller.connect("task-1"); + await expect( + controller.submit("task-1", "retry me", false, "steer"), + ).rejects.toThrow("temporary command failure"); + + expect(controller.store.getState().sessions["task-1"]).toMatchObject({ + connectionState: "connected", + events: [initialEvent], + error: undefined, + }); + }); + + it("owns and releases the bound session lifetime", async () => { + const session = createSession(); + const provider: PiSessionProvider = { + get: vi.fn(async () => session), + }; + const controller = new PiSessionController(provider, {} as TaskService); + + await controller.ensureConnected("task-1"); + await controller.setThinkingLevel("task-1", "high"); + + expect(provider.get).toHaveBeenCalledOnce(); + + controller.disconnect("task-1"); + await controller.ensureConnected("task-1"); + + expect(provider.get).toHaveBeenCalledTimes(2); }); it("opens cold tasks before connecting", async () => { - const client = createClient(); + const client = createSession(); vi.mocked(client.health).mockResolvedValue({ state: "cold" }); const openTask = vi.fn(async () => ({ success: true })); const taskService = { openTask } as unknown as TaskService; const controller = createController(client, taskService); - await controller.ensureConnected("task-1"); + await controller.ensureConnected("task-1", "run-1"); - expect(openTask).toHaveBeenCalledWith("task-1"); + expect(openTask).toHaveBeenCalledWith("task-1", "run-1"); expect(controller.store.getState().sessions["task-1"]).toMatchObject({ connectionState: "connected", }); }); + it("refreshes native thinking levels after changing models", async () => { + const session = createSession(); + const client = session.client; + vi.mocked(client.getState).mockResolvedValue({ + thinkingLevel: "high", + isStreaming: false, + isCompacting: false, + steeringMode: "all", + followUpMode: "all", + sessionId: "session-1", + autoCompactionEnabled: true, + messageCount: 0, + pendingMessageCount: 0, + model: { + provider: "posthog", + id: "model-2", + name: "Model 2", + api: "anthropic-messages", + baseUrl: "https://example.com", + reasoning: true, + input: ["text"], + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, + contextWindow: 200_000, + maxTokens: 8_000, + }, + }); + vi.mocked(client.getAvailableModels).mockResolvedValue([ + { + provider: "posthog", + id: "model-1", + contextWindow: 100_000, + reasoning: true, + }, + { + provider: "posthog", + id: "model-2", + contextWindow: 200_000, + reasoning: true, + }, + ]); + vi.mocked(client.getAvailableThinkingLevels).mockResolvedValue([ + "off", + "low", + "medium", + "high", + "xhigh", + ]); + const controller = createController(session); + await controller.ensureConnected("task-1"); + + await controller.setModel("task-1", { + provider: "posthog", + id: "model-2", + contextWindow: 200_000, + reasoning: true, + }); + + const state = controller.store.getState().sessions["task-1"]; + expect(state?.models).toEqual([ + expect.objectContaining({ id: "model-1" }), + expect.objectContaining({ id: "model-2" }), + ]); + expect(state?.thinkingLevels).toEqual([ + "off", + "low", + "medium", + "high", + "xhigh", + ]); + }); + it("makes the transcript available before model discovery finishes", async () => { let resolveModels: (models: []) => void = () => {}; const models = new Promise<[]>((resolve) => { @@ -137,9 +264,9 @@ describe("PiSessionController", () => { timestamp: 1, content: { type: "text", text: "working" }, }; - const client = createClient(); - vi.mocked(client.conversation).mockResolvedValue([initialEvent]); - vi.mocked(client.status).mockResolvedValue({ + const client = createSession(); + vi.mocked(client.getConversation).mockResolvedValue([initialEvent]); + vi.mocked(client.client.getState).mockResolvedValue({ thinkingLevel: "high", isStreaming: true, isCompacting: false, @@ -150,7 +277,7 @@ describe("PiSessionController", () => { messageCount: 1, pendingMessageCount: 0, }); - vi.mocked(client.availableModels).mockReturnValue(models); + vi.mocked(client.client.getAvailableModels).mockReturnValue(models); const controller = createController(client); const connection = controller.connect("task-1"); @@ -166,6 +293,181 @@ describe("PiSessionController", () => { await connection; }); + it("reconciles structurally equal live events included in native history", async () => { + const nativeEvent: AgentConversationEvent = { + type: "user_message", + id: "native-message-id", + timestamp: 1, + content: [{ type: "text", text: "hello" }], + }; + const liveEvent: AgentConversationEvent = { + ...nativeEvent, + id: "live-message-id", + content: [{ type: "text", text: "hello" }], + }; + let resolveConversation: (events: AgentConversationEvent[]) => void = + () => {}; + const conversation = new Promise((resolve) => { + resolveConversation = resolve; + }); + let onEvent: (event: AgentConversationEvent) => void = () => {}; + let subscribed = false; + const session = createSession(); + vi.mocked(session.getConversation).mockReturnValue(conversation); + vi.mocked(session.onConversationEvent).mockImplementation((handler) => { + onEvent = handler; + subscribed = true; + return () => {}; + }); + const controller = createController(session); + + const connection = controller.connect("task-1"); + await vi.waitFor(() => expect(subscribed).toBe(true)); + onEvent(liveEvent); + resolveConversation([nativeEvent]); + await connection; + + expect(controller.store.getState().sessions["task-1"].events).toEqual([ + nativeEvent, + ]); + }); + + it("does not append streamed assistant text already present in native history", async () => { + const nativeEvent: AgentConversationEvent = { + type: "assistant_message_chunk", + timestamp: 1, + content: { type: "text", text: "hello world" }, + }; + const liveEvent: AgentConversationEvent = { + type: "assistant_message_chunk", + timestamp: 1, + content: { type: "text", text: "world" }, + }; + let resolveConversation: (events: AgentConversationEvent[]) => void = + () => {}; + const conversation = new Promise((resolve) => { + resolveConversation = resolve; + }); + let onEvent: (event: AgentConversationEvent) => void = () => {}; + let subscribed = false; + const session = createSession(); + vi.mocked(session.getConversation).mockReturnValue(conversation); + vi.mocked(session.onConversationEvent).mockImplementation((handler) => { + onEvent = handler; + subscribed = true; + return () => {}; + }); + const controller = createController(session); + + const connection = controller.connect("task-1"); + await vi.waitFor(() => expect(subscribed).toBe(true)); + onEvent(liveEvent); + resolveConversation([nativeEvent]); + await connection; + + expect(controller.store.getState().sessions["task-1"].events).toEqual([ + nativeEvent, + ]); + }); + + it("does not let an older load overwrite a newer live refresh", async () => { + const nativeEvent: AgentConversationEvent = { + type: "user_message", + id: "message-1", + timestamp: 1, + content: [{ type: "text", text: "newer history" }], + }; + const turnCompleted: AgentConversationEvent = { + type: "turn_completed", + timestamp: 2, + }; + let resolveInitialConversation: (events: AgentConversationEvent[]) => void = + () => {}; + const initialConversation = new Promise( + (resolve) => { + resolveInitialConversation = resolve; + }, + ); + let onEvent: (event: AgentConversationEvent) => void = () => {}; + let subscribed = false; + const session = createSession(); + vi.mocked(session.getConversation) + .mockReturnValueOnce(initialConversation) + .mockResolvedValueOnce([nativeEvent, turnCompleted]); + vi.mocked(session.onConversationEvent).mockImplementation((handler) => { + onEvent = handler; + subscribed = true; + return () => {}; + }); + const controller = createController(session); + + const connection = controller.connect("task-1"); + await vi.waitFor(() => expect(subscribed).toBe(true)); + onEvent(turnCompleted); + await vi.waitFor(() => + expect(controller.store.getState().sessions["task-1"].events).toEqual([ + nativeEvent, + turnCompleted, + ]), + ); + resolveInitialConversation([]); + await connection; + + expect(controller.store.getState().sessions["task-1"].events).toEqual([ + nativeEvent, + turnCompleted, + ]); + }); + + it("drops retained live events when reconnecting after disconnect", async () => { + const liveEvent: AgentConversationEvent = { + type: "assistant_message_chunk", + timestamp: 1, + content: { type: "text", text: "stale" }, + }; + let onEvent: (event: AgentConversationEvent) => void = () => {}; + const session = createSession(); + vi.mocked(session.onConversationEvent).mockImplementation((handler) => { + onEvent = handler; + return () => {}; + }); + const controller = createController(session); + + await controller.connect("task-1"); + onEvent(liveEvent); + controller.disconnect("task-1"); + await controller.connect("task-1"); + + expect(controller.store.getState().sessions["task-1"].events).toEqual([]); + }); + + it("catches conversation refresh failures triggered by live events", async () => { + const turnCompleted: AgentConversationEvent = { + type: "turn_completed", + timestamp: 1, + }; + let onEvent: (event: AgentConversationEvent) => void = () => {}; + const session = createSession(); + vi.mocked(session.onConversationEvent).mockImplementation((handler) => { + onEvent = handler; + return () => {}; + }); + const controller = createController(session); + + await controller.connect("task-1"); + vi.mocked(session.getConversation).mockRejectedValueOnce( + new Error("refresh failed"), + ); + onEvent(turnCompleted); + await vi.waitFor(() => + expect(session.getConversation).toHaveBeenCalledTimes(2), + ); + + expect(controller.store.getState().sessions["task-1"].events).toEqual([ + turnCompleted, + ]); + }); + it("loads session state and appends normalized runtime events", async () => { const initialEvent: AgentConversationEvent = { type: "assistant_message_chunk", @@ -178,9 +480,9 @@ describe("PiSessionController", () => { status: "compacting", }; let onEvent: (event: AgentConversationEvent) => void = () => {}; - const client = createClient(); - vi.mocked(client.conversation).mockResolvedValue([initialEvent]); - vi.mocked(client.subscribe).mockImplementation((_taskId, handler) => { + const client = createSession(); + vi.mocked(client.getConversation).mockResolvedValue([initialEvent]); + vi.mocked(client.onConversationEvent).mockImplementation((handler) => { onEvent = handler; return () => {}; }); diff --git a/packages/core/src/pi-runtime/piSessionController.ts b/packages/core/src/pi-runtime/piSessionController.ts index 6ad5ef873e..edb17bba12 100644 --- a/packages/core/src/pi-runtime/piSessionController.ts +++ b/packages/core/src/pi-runtime/piSessionController.ts @@ -1,8 +1,7 @@ +import type { PiRemoteRpcClient } from "@posthog/agent/pi/remote-rpc-client"; import type { - PiCommand, PiModelOption, PiQueueMode, - PiSessionStatus, PiThinkingLevel, } from "@posthog/agent/pi/types"; import type { @@ -26,53 +25,50 @@ export type { PiThinkingLevel, } from "@posthog/agent/pi/types"; -export const PI_SESSION_CLIENT = Symbol.for("posthog.pi.sessionClient"); +export const PI_SESSION_PROVIDER = Symbol.for("posthog.pi.sessionProvider"); +export const LOCAL_PI_SESSION_FACTORY = Symbol.for( + "posthog.pi.localSessionFactory", +); -export interface PiSessionClient { - health(taskId: string): Promise; - conversation(taskId: string): Promise; - status(taskId: string): Promise; - availableModels(taskId: string): Promise; - commands(taskId: string): Promise; - subscribe( - taskId: string, +export interface PiSession { + client: PiRemoteRpcClient; + health(): Promise; + getConversation(): Promise; + onConversationEvent( onEvent: (event: AgentConversationEvent) => void, onError: (error: unknown) => void, ): () => void; - prompt(taskId: string, prompt: string): Promise; - steer(taskId: string, message: string): Promise; - followUp(taskId: string, message: string): Promise; - compact(taskId: string, customInstructions?: string): Promise; - setModel( - taskId: string, - provider: string, - modelId: string, - ): Promise<{ provider: string; id: string }>; - setThinkingLevel(taskId: string, level: PiThinkingLevel): Promise; - setSteeringMode(taskId: string, mode: PiQueueMode): Promise; - setFollowUpMode(taskId: string, mode: PiQueueMode): Promise; - bash(taskId: string, command: string): Promise; - abort(taskId: string): Promise; - abortBash(taskId: string): Promise; } +export interface PiSessionFactory { + get(taskId: string, taskRunId?: string): Promise; +} + +export type PiSessionProvider = PiSessionFactory; + export type PiSubmitResult = "prompt" | "steer" | "followUp" | "compact"; @injectable() export class PiSessionController { readonly store: PiSessionStore = createPiSessionStore(); + private readonly sessions = new Map>(); private readonly subscriptions = new Map void>(); private readonly liveEvents = new Map(); private readonly connections = new Map>(); private readonly readiness = new Map>(); + private readonly conversationRequestVersions = new Map(); + private readonly conversationAppliedVersions = new Map(); + private readonly sessionVersions = new Map(); + private readonly taskRunIds = new Map(); constructor( - @inject(PI_SESSION_CLIENT) private readonly client: PiSessionClient, + @inject(PI_SESSION_PROVIDER) private readonly provider: PiSessionProvider, @inject(TASK_SERVICE) private readonly taskService: TaskService, ) {} - ensureConnected(taskId: string): Promise { + ensureConnected(taskId: string, taskRunId?: string): Promise { + this.bindTaskRun(taskId, taskRunId); this.ensureSubscription(taskId); const existing = this.readiness.get(taskId); @@ -84,25 +80,33 @@ export class PiSessionController { connectionState: "connecting", error: undefined, }); + const connectedSessionVersion = this.getSessionVersion(taskId); const readiness = this.ensureConnectedInternal(taskId) .then(() => { - this.updateSession(taskId, { connectionState: "connected" }); + if (this.getSessionVersion(taskId) === connectedSessionVersion) { + this.updateSession(taskId, { connectionState: "connected" }); + } }) .catch((error) => { - this.updateSession(taskId, { - connectionState: "failed", - error: error instanceof Error ? error.message : String(error), - }); + if (this.getSessionVersion(taskId) === connectedSessionVersion) { + this.updateSession(taskId, { + connectionState: "failed", + error: error instanceof Error ? error.message : String(error), + }); + } throw error; }) .finally(() => { - this.readiness.delete(taskId); + if (this.readiness.get(taskId) === readiness) { + this.readiness.delete(taskId); + } }); this.readiness.set(taskId, readiness); return readiness; } - connect(taskId: string): Promise { + connect(taskId: string, taskRunId?: string): Promise { + this.bindTaskRun(taskId, taskRunId); this.ensureSubscription(taskId); const existing = this.connections.get(taskId); @@ -113,15 +117,25 @@ export class PiSessionController { this.updateSession(taskId, { error: undefined }); const connection = this.loadSession(taskId).finally(() => { - this.connections.delete(taskId); + if (this.connections.get(taskId) === connection) { + this.connections.delete(taskId); + } }); this.connections.set(taskId, connection); return connection; } disconnect(taskId: string): void { + this.advanceSessionVersion(taskId); this.subscriptions.get(taskId)?.(); this.subscriptions.delete(taskId); + this.sessions.delete(taskId); + this.taskRunIds.delete(taskId); + this.liveEvents.delete(taskId); + this.connections.delete(taskId); + this.readiness.delete(taskId); + this.conversationRequestVersions.delete(taskId); + this.conversationAppliedVersions.delete(taskId); } getSubmitAction( @@ -150,40 +164,38 @@ export class PiSessionController { const message = text.trim(); const action = this.getSubmitAction(message, isStreaming, messagingMode); - try { - if (action === "compact") { - const command = parseCommandLine(message); - const customInstructions = command?.args?.trim() || undefined; - await this.client.compact(taskId, customInstructions); - await this.refreshConversation(taskId); - } else if (action === "prompt") { - await this.client.prompt(taskId, message); - } else if (action === "steer") { - await this.client.steer(taskId, message); - } else { - await this.client.followUp(taskId, message); - } - - await this.refreshStatus(taskId); - return action; - } catch (error) { - this.updateSession(taskId, { - error: error instanceof Error ? error.message : String(error), - }); - throw error; + const session = await this.getPiSession(taskId); + if (action === "compact") { + const command = parseCommandLine(message); + const customInstructions = command?.args?.trim() || undefined; + await session.client.compact(customInstructions); + await this.refreshConversation(taskId); + } else if (action === "prompt") { + await session.client.prompt(message); + } else if (action === "steer") { + await session.client.steer(message); + } else { + await session.client.followUp(message); } + + await this.refreshStatus(taskId); + return action; } async setModel(taskId: string, model: PiModelOption): Promise { - await this.client.setModel(taskId, model.provider, model.id); + const session = await this.getPiSession(taskId); + await session.client.setModel(model.provider, model.id); await this.refreshStatus(taskId); + const thinkingLevels = await session.client.getAvailableThinkingLevels(); + this.updateSession(taskId, { thinkingLevels }); } async setThinkingLevel( taskId: string, level: PiThinkingLevel, ): Promise { - await this.client.setThinkingLevel(taskId, level); + const session = await this.getPiSession(taskId); + await session.client.setThinkingLevel(level); await this.refreshStatus(taskId); } @@ -192,10 +204,11 @@ export class PiSessionController { messagingMode: PiMessagingMode, queueMode: PiQueueMode, ): Promise { + const session = await this.getPiSession(taskId); if (messagingMode === "steer") { - await this.client.setSteeringMode(taskId, queueMode); + await session.client.setSteeringMode(queueMode); } else { - await this.client.setFollowUpMode(taskId, queueMode); + await session.client.setFollowUpMode(queueMode); } await this.refreshStatus(taskId); } @@ -203,7 +216,8 @@ export class PiSessionController { async bash(taskId: string, command: string): Promise { this.updateSession(taskId, { isBashRunning: true }); try { - await this.client.bash(taskId, command); + const session = await this.getPiSession(taskId); + await session.client.bash(command); await this.refreshConversation(taskId); } finally { this.updateSession(taskId, { isBashRunning: false }); @@ -211,22 +225,34 @@ export class PiSessionController { } async abort(taskId: string): Promise { - await this.client.abort(taskId); + const session = await this.getPiSession(taskId); + await session.client.abort(); await this.refreshStatus(taskId); } async abortBash(taskId: string): Promise { - await this.client.abortBash(taskId); + const session = await this.getPiSession(taskId); + await session.client.abortBash(); this.updateSession(taskId, { isBashRunning: false }); } private async ensureConnectedInternal(taskId: string): Promise { - const health = await this.client.health(taskId); + const session = await this.getPiSession(taskId); + const health = await session.health(); if (health.state === "cold") { - const result = await this.taskService.openTask(taskId); + const taskRunId = this.taskRunIds.get(taskId); + const result = taskRunId + ? await this.taskService.openTask(taskId, taskRunId) + : await this.taskService.openTask(taskId); if (!result.success) { throw new Error(result.error); } + + this.subscriptions.get(taskId)?.(); + this.subscriptions.delete(taskId); + this.sessions.delete(taskId); + this.connections.delete(taskId); + this.ensureSubscription(taskId); } await this.connect(taskId); @@ -237,48 +263,78 @@ export class PiSessionController { return; } - const unsubscribe = this.client.subscribe( - taskId, - (event) => this.handleEvent(taskId, event), - (error) => { + let disposed = false; + let unsubscribe: (() => void) | undefined; + void this.getPiSession(taskId) + .then((session) => { + if (disposed) { + return; + } + unsubscribe = session.onConversationEvent( + (event) => this.handleEvent(taskId, event), + (error) => { + this.updateSession(taskId, { + error: error instanceof Error ? error.message : String(error), + }); + }, + ); + }) + .catch((error) => { this.updateSession(taskId, { error: error instanceof Error ? error.message : String(error), }); - }, - ); - this.subscriptions.set(taskId, unsubscribe); + }); + this.subscriptions.set(taskId, () => { + disposed = true; + unsubscribe?.(); + }); } private async loadSession(taskId: string): Promise { + const connectedSessionVersion = this.getSessionVersion(taskId); + const conversationVersion = this.nextConversationVersion(taskId); try { - const [events, status] = await Promise.all([ - this.client.conversation(taskId), - this.client.status(taskId), - ]); - const liveEvents = status.isStreaming - ? (this.liveEvents.get(taskId) ?? []) - : []; + const session = await this.getPiSession(taskId); + const events = await session.getConversation(); + const status = await session.client.getState(); + if (this.getSessionVersion(taskId) !== connectedSessionVersion) { + return; + } + const currentSession = this.getSession(taskId); - this.liveEvents.set(taskId, liveEvents); + let reconciledEvents = currentSession.events; + if (this.shouldApplyConversation(taskId, conversationVersion)) { + const liveEvents = this.liveEvents.get(taskId) ?? []; + const newLiveEvents = this.reconcileLiveEvents(events, liveEvents); + this.liveEvents.set(taskId, newLiveEvents); + reconciledEvents = [...events, ...newLiveEvents]; + } + this.setSession(taskId, { connectionState: "connected", - events: [...events, ...liveEvents], + events: reconciledEvents, status, models: currentSession.models, + thinkingLevels: currentSession.thinkingLevels, commands: currentSession.commands, isBashRunning: false, error: undefined, }); - const [models, commands] = await Promise.all([ - this.client.availableModels(taskId), - this.client.commands(taskId), + const [models, thinkingLevels, commands] = await Promise.all([ + session.client.getAvailableModels(), + session.client.getAvailableThinkingLevels(), + session.client.getCommands(), ]); - this.updateSession(taskId, { models, commands }); + if (this.getSessionVersion(taskId) === connectedSessionVersion) { + this.updateSession(taskId, { models, thinkingLevels, commands }); + } } catch (error) { - this.updateSession(taskId, { - error: error instanceof Error ? error.message : String(error), - }); + if (this.getSessionVersion(taskId) === connectedSessionVersion) { + this.updateSession(taskId, { + error: error instanceof Error ? error.message : String(error), + }); + } throw error; } } @@ -305,32 +361,185 @@ export class PiSessionController { }); if (event.type === "turn_completed") { - const capturedCount = liveEvents.length; - void this.refreshConversation(taskId, capturedCount); + void this.refreshConversation(taskId).catch(() => {}); } } - private async refreshConversation( - taskId: string, - capturedLiveCount?: number, - ): Promise { - const events = await this.client.conversation(taskId); + private async refreshConversation(taskId: string): Promise { + const conversationVersion = this.nextConversationVersion(taskId); + const sessionVersion = this.getSessionVersion(taskId); + const session = await this.getPiSession(taskId); + const events = await session.getConversation(); + if ( + this.getSessionVersion(taskId) !== sessionVersion || + !this.shouldApplyConversation(taskId, conversationVersion) + ) { + return; + } + const liveEvents = this.liveEvents.get(taskId) ?? []; - const remainingEvents = - capturedLiveCount === undefined - ? [] - : liveEvents.slice(capturedLiveCount); + const remainingEvents = this.reconcileLiveEvents(events, liveEvents); this.liveEvents.set(taskId, remainingEvents); this.updateSession(taskId, { events: [...events, ...remainingEvents], }); } + private nextConversationVersion(taskId: string): number { + const version = (this.conversationRequestVersions.get(taskId) ?? 0) + 1; + this.conversationRequestVersions.set(taskId, version); + return version; + } + + private shouldApplyConversation(taskId: string, version: number): boolean { + const appliedVersion = this.conversationAppliedVersions.get(taskId) ?? 0; + if (version < appliedVersion) { + return false; + } + + this.conversationAppliedVersions.set(taskId, version); + return true; + } + + private reconcileLiveEvents( + nativeEvents: AgentConversationEvent[], + liveEvents: AgentConversationEvent[], + ): AgentConversationEvent[] { + const nativeEventCounts = new Map(); + const nativeTextByMessage = new Map(); + + for (const event of nativeEvents) { + const textMessageKey = this.getTextMessageKey(event); + const text = this.getTextContent(event); + if (textMessageKey && text !== undefined) { + const nativeText = nativeTextByMessage.get(textMessageKey) ?? ""; + nativeTextByMessage.set(textMessageKey, nativeText + text); + continue; + } + + const key = this.getEventKey(event); + nativeEventCounts.set(key, (nativeEventCounts.get(key) ?? 0) + 1); + } + + const nativeTextOffsets = new Map(); + return liveEvents.filter((event) => { + const textMessageKey = this.getTextMessageKey(event); + const text = this.getTextContent(event); + if (textMessageKey && text !== undefined) { + const nativeText = nativeTextByMessage.get(textMessageKey); + if (nativeText === undefined) { + return true; + } + + const offset = nativeTextOffsets.get(textMessageKey) ?? 0; + const matchIndex = nativeText.indexOf(text, offset); + if (matchIndex === -1) { + return true; + } + + nativeTextOffsets.set(textMessageKey, matchIndex + text.length); + return false; + } + + const key = this.getEventKey(event); + const nativeCount = nativeEventCounts.get(key) ?? 0; + if (nativeCount === 0) { + return true; + } + + nativeEventCounts.set(key, nativeCount - 1); + return false; + }); + } + + private getTextMessageKey(event: AgentConversationEvent): string | undefined { + if ( + event.type !== "assistant_message_chunk" && + event.type !== "assistant_thought_chunk" + ) { + return undefined; + } + + if (event.content.type !== "text") { + return undefined; + } + + return `${event.type}:${event.timestamp}`; + } + + private getTextContent(event: AgentConversationEvent): string | undefined { + if ( + event.type !== "assistant_message_chunk" && + event.type !== "assistant_thought_chunk" + ) { + return undefined; + } + + return event.content.type === "text" ? event.content.text : undefined; + } + + private getEventKey(event: AgentConversationEvent): string { + if (event.type === "user_message") { + return JSON.stringify({ + type: event.type, + timestamp: event.timestamp, + content: event.content, + }); + } + + return JSON.stringify(event); + } + private async refreshStatus(taskId: string): Promise { - const status = await this.client.status(taskId); + const session = await this.getPiSession(taskId); + const status = await session.client.getState(); this.updateSession(taskId, { status }); } + private bindTaskRun(taskId: string, taskRunId?: string): void { + const currentTaskRunId = this.taskRunIds.get(taskId); + if (!taskRunId || currentTaskRunId === taskRunId) { + return; + } + + if (currentTaskRunId) { + this.advanceSessionVersion(taskId); + this.subscriptions.get(taskId)?.(); + this.subscriptions.delete(taskId); + this.sessions.delete(taskId); + this.liveEvents.delete(taskId); + this.connections.delete(taskId); + this.readiness.delete(taskId); + this.conversationRequestVersions.delete(taskId); + this.conversationAppliedVersions.delete(taskId); + } + this.taskRunIds.set(taskId, taskRunId); + } + + private getPiSession(taskId: string): Promise { + const existing = this.sessions.get(taskId); + if (existing) { + return existing; + } + + const session = this.provider.get(taskId, this.taskRunIds.get(taskId)); + this.sessions.set(taskId, session); + void session.catch(() => { + if (this.sessions.get(taskId) === session) { + this.sessions.delete(taskId); + } + }); + return session; + } + + private getSessionVersion(taskId: string): number { + return this.sessionVersions.get(taskId) ?? 0; + } + + private advanceSessionVersion(taskId: string): void { + this.sessionVersions.set(taskId, this.getSessionVersion(taskId) + 1); + } + private getSession(taskId: string): PiControllerSessionState { return ( this.store.getState().sessions[taskId] ?? createEmptyPiControllerSession() diff --git a/packages/core/src/pi-runtime/piSessionProvider.test.ts b/packages/core/src/pi-runtime/piSessionProvider.test.ts new file mode 100644 index 0000000000..e0f10aa73b --- /dev/null +++ b/packages/core/src/pi-runtime/piSessionProvider.test.ts @@ -0,0 +1,154 @@ +import type { PiRemoteRpcClient } from "@posthog/agent/pi/remote-rpc-client"; +import { describe, expect, it, vi } from "vitest"; +import type { CloudTaskClient } from "../cloud-task/cloudTaskClient"; +import type { TaskService } from "../task-detail/taskService"; +import type { PiSession, PiSessionFactory } from "./piSessionController"; +import { RoutingPiSessionProvider } from "./piSessionProvider"; + +function localSession(): PiSession { + const client = { + getState: vi.fn(async () => ({ isStreaming: false })), + getAvailableModels: vi.fn(async () => []), + getCommands: vi.fn(async () => []), + prompt: vi.fn(async () => {}), + steer: vi.fn(async () => {}), + followUp: vi.fn(async () => {}), + compact: vi.fn(async () => undefined), + setModel: vi.fn(async () => ({ provider: "posthog", id: "model" })), + setThinkingLevel: vi.fn(async () => {}), + setSteeringMode: vi.fn(async () => {}), + setFollowUpMode: vi.fn(async () => {}), + bash: vi.fn(async () => undefined), + abort: vi.fn(async () => {}), + abortBash: vi.fn(async () => {}), + } as unknown as PiRemoteRpcClient; + + return { + client, + health: vi.fn(async () => ({ state: "idle" as const })), + getConversation: vi.fn(async () => []), + onConversationEvent: vi.fn(() => () => {}), + }; +} + +function localFactory(session: PiSession): PiSessionFactory { + return { + get: vi.fn(async () => session), + }; +} + +function cloudTaskClient(): CloudTaskClient { + return { + getContext: vi.fn(async () => ({ + apiHost: "https://us.posthog.com", + teamId: 1, + })), + watch: vi.fn(async () => {}), + unwatch: vi.fn(async () => {}), + subscribe: vi.fn(() => () => {}), + sendCommand: vi.fn(async (input) => ({ + success: true, + result: { + type: "response", + command: input.params?.command + ? (input.params.command as { type: string }).type + : "unknown", + success: true, + }, + })), + }; +} + +function taskService(environment: "local" | "cloud"): TaskService { + return { + getTask: vi.fn(async () => ({ + id: "task-1", + runtime: "pi", + latest_run: + environment === "cloud" + ? { id: "run-1", environment: "cloud", status: "in_progress" } + : null, + })), + } as unknown as TaskService; +} + +describe("RoutingPiSessionProvider", () => { + it("returns a cloud session bound to its task run", async () => { + const local = localSession(); + const cloudTasks = cloudTaskClient(); + const provider = new RoutingPiSessionProvider( + localFactory(local), + cloudTasks, + taskService("cloud"), + ); + + const session = await provider.get("task-1"); + await session.client.steer("change direction"); + + expect(cloudTasks.sendCommand).toHaveBeenCalledWith({ + taskId: "task-1", + runId: "run-1", + apiHost: "https://us.posthog.com", + teamId: 1, + method: "pi/rpc", + params: { + command: { type: "steer", message: "change direction" }, + }, + }); + expect(local.client.steer).not.toHaveBeenCalled(); + }); + + it("binds explicit historical runs and invalidates changed run context", async () => { + const local = localSession(); + const cloudTasks = cloudTaskClient(); + const getTask = vi.fn(async (_taskId: string, taskRunId?: string) => ({ + id: "task-1", + runtime: "pi" as const, + latest_run: { + id: taskRunId ?? "run-latest", + environment: "cloud" as const, + status: "in_progress" as const, + }, + })); + const provider = new RoutingPiSessionProvider( + localFactory(local), + cloudTasks, + { getTask } as unknown as TaskService, + ); + + const historical = await provider.get("task-1", "run-old"); + await historical.client.abort(); + const replacement = await provider.get("task-1", "run-new"); + await replacement.client.abort(); + + expect(getTask).toHaveBeenNthCalledWith(1, "task-1", "run-old"); + expect(getTask).toHaveBeenNthCalledWith(2, "task-1", "run-new"); + expect(cloudTasks.sendCommand).toHaveBeenNthCalledWith( + 1, + expect.objectContaining({ runId: "run-old" }), + ); + expect(cloudTasks.sendCommand).toHaveBeenNthCalledWith( + 2, + expect.objectContaining({ runId: "run-new" }), + ); + }); + + it("delegates local session lifetime to the controller", async () => { + const local = localSession(); + const localSessions = localFactory(local); + const cloudTasks = cloudTaskClient(); + const tasks = taskService("local"); + const provider = new RoutingPiSessionProvider( + localSessions, + cloudTasks, + tasks, + ); + + const first = await provider.get("task-1"); + const second = await provider.get("task-1"); + + expect(first).toBe(local); + expect(second).toBe(local); + expect(localSessions.get).toHaveBeenCalledTimes(2); + }); +}); diff --git a/packages/core/src/pi-runtime/piSessionProvider.ts b/packages/core/src/pi-runtime/piSessionProvider.ts new file mode 100644 index 0000000000..5dd4be1bf1 --- /dev/null +++ b/packages/core/src/pi-runtime/piSessionProvider.ts @@ -0,0 +1,85 @@ +import type { TaskRunStatus } from "@posthog/shared"; +import { inject, injectable } from "inversify"; +import { + CLOUD_TASK_CLIENT, + type CloudTaskClient, +} from "../cloud-task/cloudTaskClient"; +import { isTerminalStatus } from "../cloud-task/schemas"; +import { TASK_SERVICE, type TaskService } from "../task-detail/taskService"; +import { + CloudPiSessionClient, + type CloudPiSessionContext, +} from "./cloudPiSessionClient"; +import { + LOCAL_PI_SESSION_FACTORY, + type PiSession, + type PiSessionFactory, + type PiSessionProvider, +} from "./piSessionController"; + +@injectable() +export class RoutingPiSessionProvider implements PiSessionProvider { + constructor( + @inject(LOCAL_PI_SESSION_FACTORY) + private readonly localFactory: PiSessionFactory, + @inject(CLOUD_TASK_CLIENT) + private readonly cloudTaskClient: CloudTaskClient, + @inject(TASK_SERVICE) + private readonly taskService: TaskService, + ) {} + + async get(taskId: string, taskRunId?: string): Promise { + const cloudContext = await this.resolveCloudContext(taskId, taskRunId); + if (cloudContext) { + return new CloudPiSessionClient(this.cloudTaskClient, cloudContext); + } + return this.localFactory.get(taskId); + } + + private async resolveCloudContext( + taskId: string, + taskRunId?: string, + ): Promise { + const [task, context] = await Promise.all([ + this.taskService.getTask(taskId, taskRunId), + this.cloudTaskClient.getContext(), + ]); + const run = task.latest_run; + if (!context || !run || run.environment !== "cloud") { + return null; + } + + return { + taskId, + runId: run.id, + runStatus: run.status, + waitUntilReady: () => this.waitUntilCloudRuntimeReady(taskId, run.id), + ...context, + }; + } + + private async waitUntilCloudRuntimeReady( + taskId: string, + runId: string, + ): Promise { + let delayMs = 1_000; + for (let attempt = 0; attempt < 60; attempt += 1) { + const task = await this.taskService.getTask(taskId, runId); + const run = task.latest_run; + if (!run || run.id !== runId) { + throw new Error(`Cloud task run ${runId} is unavailable`); + } + if ( + isTerminalStatus(run.status) || + typeof run.state?.sandbox_url === "string" + ) { + return run.status; + } + + await new Promise((resolve) => setTimeout(resolve, delayMs)); + delayMs = Math.min(delayMs * 2, 5_000); + } + + throw new Error(`Cloud task run ${runId} did not become ready`); + } +} diff --git a/packages/core/src/pi-runtime/piSessionStore.ts b/packages/core/src/pi-runtime/piSessionStore.ts index 1b4519a59f..4a15b18a8e 100644 --- a/packages/core/src/pi-runtime/piSessionStore.ts +++ b/packages/core/src/pi-runtime/piSessionStore.ts @@ -2,6 +2,7 @@ import type { PiCommand, PiModelOption, PiSessionStatus, + PiThinkingLevel, } from "@posthog/agent/pi/types"; import type { AgentConversationEvent } from "@posthog/shared"; import { createStore, type StoreApi } from "zustand/vanilla"; @@ -10,6 +11,7 @@ export interface PiControllerSessionState { connectionState: "connecting" | "connected" | "failed"; events: AgentConversationEvent[]; models: PiModelOption[]; + thinkingLevels: PiThinkingLevel[]; commands: PiCommand[]; status?: PiSessionStatus; error?: string; @@ -31,6 +33,7 @@ export function createEmptyPiControllerSession(): PiControllerSessionState { connectionState: "connecting", events: [], models: [], + thinkingLevels: [], commands: [], isBashRunning: false, }; diff --git a/packages/core/src/task-detail/piTaskCreator.ts b/packages/core/src/task-detail/piTaskCreator.ts deleted file mode 100644 index b453b6f5aa..0000000000 --- a/packages/core/src/task-detail/piTaskCreator.ts +++ /dev/null @@ -1,197 +0,0 @@ -import { - Saga, - type SagaLogger, - type TaskCreationInput, - type TaskCreationOutput, - type Workspace, -} from "@posthog/shared"; -import type { Task } from "@posthog/shared/domain-types"; -import type { PiRunner } from "../pi-runtime/piRunner"; -import type { TaskCreationApiClient } from "./taskCreationApiClient"; -import type { ITaskCreationHost } from "./taskCreationHost"; -import { resolveTaskRepository } from "./taskRepository"; - -export interface PiTaskCreatorDeps { - posthogClient: TaskCreationApiClient; - host: ITaskCreationHost; - piRunner: PiRunner; - onTaskReady?: (output: TaskCreationOutput) => void; -} - -export class PiTaskCreator extends Saga { - readonly sagaName = "PiTaskCreator"; - - constructor( - private readonly deps: PiTaskCreatorDeps, - logger?: SagaLogger, - ) { - super(logger); - } - - protected async execute( - input: TaskCreationInput, - ): Promise { - if (input.workspaceMode === "cloud") { - throw new Error("Pi tasks are only supported in local workspaces"); - } - - const task = await this.createTask(input); - const repoPath = input.repoPath; - let workspace: Workspace | null = null; - - if (repoPath) { - workspace = await this.createWorkspace(task, repoPath, input); - } else if (input.allowNoRepo) { - workspace = await this.createScratchWorkspace(task); - } - - if (!workspace) { - throw new Error("Pi tasks require a workspace or scratch directory"); - } - - const cwd = workspace.worktreePath ?? workspace.folderPath; - const additionalDirectories = (input.additionalDirectories ?? []).filter( - (path) => path && path !== input.repoPath, - ); - if (additionalDirectories.length > 0) { - await this.step({ - name: "additional_directories", - execute: async () => { - await Promise.all( - additionalDirectories.map((path) => - this.deps.host.addAdditionalDirectory({ taskId: task.id, path }), - ), - ); - return { taskId: task.id, paths: additionalDirectories }; - }, - rollback: async ({ taskId, paths }) => { - await Promise.all( - paths.map((path) => - this.deps.host.removeAdditionalDirectory({ taskId, path }), - ), - ); - }, - }); - } - - await this.step({ - name: "pi_session", - execute: async () => { - await this.deps.piRunner.create({ - taskId: task.id, - cwd, - prompt: input.content ?? "", - model: input.model, - }); - return { taskId: task.id }; - }, - rollback: async ({ taskId }) => this.deps.piRunner.stop(taskId), - }); - - this.deps.onTaskReady?.({ task, workspace }); - return { task, workspace }; - } - - private async createTask(input: TaskCreationInput): Promise { - const repository = await resolveTaskRepository( - input, - this.deps.host, - this.log, - ); - - return this.step({ - name: "task_creation", - execute: async () => - (await this.deps.posthogClient.createTask({ - description: input.content ?? "", - repository: repository ?? undefined, - origin_product: input.signalReportId - ? "signal_report" - : "user_created", - signal_report: input.signalReportId ?? undefined, - channel: input.channelId ?? undefined, - runtime: "pi", - })) as unknown as Task, - rollback: async (task) => this.deps.posthogClient.deleteTask(task.id), - }); - } - - private async createWorkspace( - task: Task, - repoPath: string, - input: TaskCreationInput, - ): Promise { - const folder = await this.deps.host.getFolders().then(async (folders) => { - const existing = folders.find((candidate) => candidate.path === repoPath); - return existing ?? this.deps.host.addFolder({ folderPath: repoPath }); - }); - const workspaceInfo = await this.step({ - name: "workspace_creation", - execute: () => - this.deps.host.createWorkspace({ - taskId: task.id, - mainRepoPath: repoPath, - folderId: folder.id, - folderPath: repoPath, - mode: input.workspaceMode ?? "local", - branch: input.branch ?? undefined, - allowRemoteBranchCheckout: input.allowRemoteBranchCheckout, - reuseExistingWorktree: input.reuseExistingWorktree, - }), - rollback: () => - this.deps.host.deleteWorkspace({ - taskId: task.id, - mainRepoPath: repoPath, - }), - }); - - const workspaceMode = input.workspaceMode ?? "local"; - const worktree = workspaceInfo.worktree; - if (workspaceMode === "worktree" && !worktree) { - throw new Error("Pi worktree creation did not return a worktree"); - } - if (worktree) { - return { - taskId: task.id, - folderId: folder.id, - folderPath: repoPath, - mode: workspaceMode, - worktreePath: worktree.worktreePath, - worktreeName: worktree.worktreeName, - branchName: worktree.branchName, - baseBranch: worktree.baseBranch, - linkedBranch: workspaceInfo.linkedBranch, - createdAt: worktree.createdAt, - }; - } - - return { - taskId: task.id, - folderId: folder.id, - folderPath: repoPath, - mode: "local", - worktreePath: null, - worktreeName: null, - branchName: workspaceInfo.branchName, - baseBranch: input.branch ?? null, - linkedBranch: workspaceInfo.linkedBranch, - createdAt: new Date().toISOString(), - }; - } - - private async createScratchWorkspace(task: Task): Promise { - const folderPath = await this.deps.host.ensureScratchDir(task.id); - return { - taskId: task.id, - folderId: "", - folderPath, - mode: "local", - worktreePath: null, - worktreeName: null, - branchName: null, - baseBranch: null, - linkedBranch: null, - createdAt: new Date().toISOString(), - }; - } -} diff --git a/packages/core/src/task-detail/taskCreationApiClient.ts b/packages/core/src/task-detail/taskCreationApiClient.ts index 9a24ccd2e8..6a39cd08c6 100644 --- a/packages/core/src/task-detail/taskCreationApiClient.ts +++ b/packages/core/src/task-detail/taskCreationApiClient.ts @@ -45,4 +45,5 @@ export interface TaskCreationApiClient { runId: string, options?: StartTaskRunClientOptions, ): Promise; + resumeRunInCloud(taskId: string, runId: string): Promise; } diff --git a/packages/core/src/task-detail/taskCreationSaga.test.ts b/packages/core/src/task-detail/taskCreationSaga.test.ts index 5558ab6580..b18eb8ac9f 100644 --- a/packages/core/src/task-detail/taskCreationSaga.test.ts +++ b/packages/core/src/task-detail/taskCreationSaga.test.ts @@ -10,8 +10,6 @@ const mockHost = vi.hoisted(() => ({ getAuthenticatedClient: vi.fn(), getTaskDirectory: vi.fn(), ensureScratchDir: vi.fn(), - startPiSession: vi.fn(), - stopPiSession: vi.fn(), getWorkspace: vi.fn(), createWorkspace: vi.fn(), deleteWorkspace: vi.fn(), @@ -37,12 +35,17 @@ const mockHost = vi.hoisted(() => ({ linkTaskBranch: vi.fn(), })); -import { PiTaskCreator } from "./piTaskCreator"; import { TaskCreationSaga } from "./taskCreationSaga"; import { buildWorktreeAdoptionInput } from "./taskInput"; const host = mockHost as unknown as ITaskCreationHost; +const piRunner = { + create: vi.fn(async () => {}), + resume: vi.fn(async () => {}), + stop: vi.fn(async () => {}), +}; + const sessionService = { connectToTask: vi.fn(), disconnectFromTask: vi.fn(), @@ -97,6 +100,7 @@ function makeSaga( } as never, host, sessionService, + piRunner, track: vi.fn(), ...extra, }); @@ -253,6 +257,7 @@ describe("TaskCreationSaga", () => { } as never, host, sessionService, + piRunner, track: vi.fn(), }); @@ -303,6 +308,7 @@ describe("TaskCreationSaga", () => { } as never, host, sessionService, + piRunner, track: vi.fn(), }); @@ -345,17 +351,7 @@ describe("TaskCreationSaga", () => { it("starts a Pi session without creating an ACP session", async () => { const createdTask = createTask({ repository: undefined }); const createTaskRequest = vi.fn().mockResolvedValue(createdTask); - const saga = new PiTaskCreator({ - posthogClient: { - createTask: createTaskRequest, - deleteTask: vi.fn(), - } as never, - host, - piRunner: { - create: mockHost.startPiSession, - stop: mockHost.stopPiSession, - } as never, - }); + const saga = makeSaga({ createTask: createTaskRequest }); const result = await saga.run({ content: "Draft a launch email", @@ -369,13 +365,131 @@ describe("TaskCreationSaga", () => { expect(createTaskRequest).toHaveBeenCalledWith( expect.objectContaining({ runtime: "pi" }), ); - expect(mockHost.startPiSession).toHaveBeenCalledWith({ + expect(piRunner.create).toHaveBeenCalledWith({ taskId: "task-123", cwd: "/tmp/scratch/task-123", prompt: "Draft a launch email", model: "claude-sonnet", }); expect(sessionService.connectToTask).not.toHaveBeenCalled(); + expect(sessionService.markTaskCreationInFlight).not.toHaveBeenCalled(); + }); + + it("starts a cloud Pi run without creating a local runtime", async () => { + const createdTask = createTask({ repository: "posthog/posthog" }); + const startedTask = createTask({ latest_run: createRun(), runtime: "pi" }); + const createTaskRun = vi.fn().mockResolvedValue(createRun()); + const startTaskRun = vi.fn().mockResolvedValue(startedTask); + const saga = makeSaga({ + createTask: vi.fn().mockResolvedValue(createdTask), + createTaskRun, + startTaskRun, + }); + + const result = await saga.run({ + content: "Fix the cloud build", + repository: "posthog/posthog", + workspaceMode: "cloud", + runtime: "pi", + branch: "main", + adapter: "codex", + model: "gpt-5.4", + reasoningLevel: "high", + }); + + expect(result.success).toBe(true); + expect(createTaskRun).toHaveBeenCalledWith("task-123", { + environment: "cloud", + mode: "interactive", + branch: "main", + adapter: "codex", + model: "gpt-5.4", + reasoningLevel: "high", + sandboxEnvironmentId: undefined, + customImageId: undefined, + prAuthorshipMode: "user", + autoPublish: undefined, + rtkEnabled: undefined, + runSource: "manual", + signalReportId: undefined, + homeQuickAction: undefined, + importedMcpServers: undefined, + relayedMcpServers: undefined, + initialPermissionMode: "auto", + }); + expect(startTaskRun).toHaveBeenCalledWith("task-123", "run-123", { + pendingUserMessage: "Fix the cloud build", + pendingUserArtifactIds: undefined, + }); + expect(piRunner.create).not.toHaveBeenCalled(); + }); + + it("uploads initial Pi files and local skill bundles before starting the cloud run", async () => { + const createdTask = createTask({ repository: "posthog/posthog" }); + const startedTask = createTask({ latest_run: createRun(), runtime: "pi" }); + const createTaskRun = vi.fn().mockResolvedValue(createRun()); + const startTaskRun = vi.fn().mockResolvedValue(startedTask); + const skillTag = + ' do it '; + const messageText = + ' do it'; + const skillBundles = [ + { name: "my-skill", source: "user" as const, path: "/skills/my-skill" }, + ]; + + mockHost.resolveLocalSkillCommandPrompt.mockResolvedValue(skillTag); + mockHost.getCloudPromptTransport.mockReturnValue({ + filePaths: ["/tmp/input.txt"], + skillBundles, + messageText, + promptText: "do it\n\nAttached files: input.txt", + }); + mockHost.uploadRunAttachments.mockResolvedValue([ + "file-artifact", + "skill-artifact", + ]); + + const saga = makeSaga({ + createTask: vi.fn().mockResolvedValue(createdTask), + createTaskRun, + startTaskRun, + }); + + const result = await saga.run({ + content: '/my-skill do it ', + filePaths: ["/tmp/input.txt"], + repository: "posthog/posthog", + workspaceMode: "cloud", + runtime: "pi", + adapter: "codex", + model: "gpt-5.4", + reasoningLevel: "high", + }); + + expect(result.success).toBe(true); + expect(mockHost.resolveLocalSkillCommandPrompt).toHaveBeenCalledWith( + '/my-skill do it ', + ); + expect(mockHost.getCloudPromptTransport).toHaveBeenCalledWith(skillTag, [ + "/tmp/input.txt", + ]); + expect(mockHost.uploadRunAttachments).toHaveBeenCalledWith( + expect.anything(), + "task-123", + "run-123", + ["/tmp/input.txt"], + skillBundles, + ); + expect(startTaskRun).toHaveBeenCalledWith("task-123", "run-123", { + pendingUserMessage: messageText, + pendingUserArtifactIds: ["file-artifact", "skill-artifact"], + }); + expect(createTaskRun.mock.invocationCallOrder[0]).toBeLessThan( + mockHost.uploadRunAttachments.mock.invocationCallOrder[0], + ); + expect( + mockHost.uploadRunAttachments.mock.invocationCallOrder[0], + ).toBeLessThan(startTaskRun.mock.invocationCallOrder[0]); }); it("uploads initial cloud attachments before starting the run", async () => { diff --git a/packages/core/src/task-detail/taskCreationSaga.ts b/packages/core/src/task-detail/taskCreationSaga.ts index 2e3d0c5bad..d85de2b2c0 100644 --- a/packages/core/src/task-detail/taskCreationSaga.ts +++ b/packages/core/src/task-detail/taskCreationSaga.ts @@ -18,6 +18,7 @@ import { } from "@posthog/shared"; import { ANALYTICS_EVENTS } from "@posthog/shared/analytics-events"; import type { Task } from "@posthog/shared/domain-types"; +import type { PiRunner } from "../pi-runtime/piRunner"; import type { TaskCreationApiClient } from "./taskCreationApiClient"; import type { CloudPromptTransport, @@ -30,6 +31,7 @@ export interface TaskCreationDeps { posthogClient: TaskCreationApiClient; host: ITaskCreationHost; sessionService: SessionService; + piRunner: PiRunner; onTaskReady?: (output: TaskCreationOutput) => void; track: (event: string, props?: Record) => void; } @@ -87,15 +89,18 @@ export class TaskCreationSaga extends Saga< input: TaskCreationInput, ): Promise { const taskId = input.taskId; + const isPiRuntime = input.runtime === "pi"; const folderPromise = !taskId && input.repoPath ? this.resolveFolder(input.repoPath) : undefined; - const importedClaude = await this.importClaudeSession(input); + const importedClaude = isPiRuntime + ? undefined + : await this.importClaudeSession(input); const warmPayload = - !taskId && input.workspaceMode === "cloud" + !isPiRuntime && !taskId && input.workspaceMode === "cloud" ? await this.prepareWarmActivation(input) : null; @@ -105,9 +110,9 @@ export class TaskCreationSaga extends Saga< ) : await this.createTask(input, warmPayload); - // Session reconcile auto-recovers run-less local tasks; mark this one as - // mid-creation so the recovery doesn't race the agent_session step below. - this.deps.sessionService.markTaskCreationInFlight(task.id); + if (!isPiRuntime) { + this.deps.sessionService.markTaskCreationInFlight(task.id); + } if (importedClaude && input.repoPath) { await this.recordClaudeImport(input, importedClaude, task.id); @@ -384,7 +389,7 @@ export class TaskCreationSaga extends Saga< // the optimistic placeholder would show the bare task description with // no CONTEXT.md / personalization chip. Hand the augmented message to // the session service so it seeds the placeholder right away. - if (augmented && pendingUserMessage) { + if (!isPiRuntime && augmented && pendingUserMessage) { this.deps.sessionService.rememberInitialCloudPrompt( task.id, pendingUserMessage, @@ -418,7 +423,7 @@ export class TaskCreationSaga extends Saga< throw new Error("Failed to create cloud run"); } - if (input.relayedMcpServers?.length) { + if (!isPiRuntime && input.relayedMcpServers?.length) { // Best-effort: relay designation failing must not fail creation — // the run still works, minus desktop-relayed servers. await this.deps.sessionService @@ -485,7 +490,7 @@ export class TaskCreationSaga extends Saga< if (shouldConnect) { const initialPrompt = - !input.taskId && input.content + !isPiRuntime && !input.taskId && input.content ? await this.readOnlyStep("build_prompt_blocks", () => buildPromptBlocks( input.content ?? "", @@ -510,6 +515,16 @@ export class TaskCreationSaga extends Saga< await this.step({ name: "agent_session", execute: async () => { + if (isPiRuntime) { + await this.deps.piRunner.create({ + taskId: task.id, + cwd: agentCwd ?? "", + prompt: input.content ?? "", + model: input.model, + }); + return { taskId: task.id }; + } + const connectParams: ConnectParams = { task, repoPath: agentCwd ?? "", @@ -530,6 +545,10 @@ export class TaskCreationSaga extends Saga< return { taskId: task.id }; }, rollback: async ({ taskId }) => { + if (isPiRuntime) { + await this.deps.piRunner.stop(taskId); + return; + } this.log.info("Rolling back: disconnecting agent session", { taskId, }); @@ -749,6 +768,8 @@ export class TaskCreationSaga extends Saga< name: "task_creation", execute: async () => { const description = input.taskDescription ?? input.content ?? ""; + const canActivateWarmRun = + input.runtime !== "pi" && !warmPayload?.suppressWarmReuse; const result = await this.deps.posthogClient.createTask({ description, repository: repository ?? undefined, @@ -768,36 +789,40 @@ export class TaskCreationSaga extends Saga< // The server associates the task with the report and records the implementation // task_run artefact — no relationship label is sent (associations are unlabelled). branch: - input.workspaceMode === "cloud" && !warmPayload?.suppressWarmReuse + input.workspaceMode === "cloud" && canActivateWarmRun ? (input.branch ?? null) : undefined, runtime_adapter: - input.workspaceMode === "cloud" + input.workspaceMode === "cloud" && canActivateWarmRun ? (input.adapter ?? null) : undefined, model: - input.workspaceMode === "cloud" ? (input.model ?? null) : undefined, + input.workspaceMode === "cloud" && canActivateWarmRun + ? (input.model ?? null) + : undefined, reasoning_effort: - input.workspaceMode === "cloud" + input.workspaceMode === "cloud" && canActivateWarmRun ? (input.reasoningLevel ?? null) : undefined, sandbox_environment_id: - input.workspaceMode === "cloud" && !warmPayload?.suppressWarmReuse + input.workspaceMode === "cloud" && canActivateWarmRun ? input.sandboxEnvironmentId : undefined, custom_image_id: - input.workspaceMode === "cloud" && !warmPayload?.suppressWarmReuse + input.workspaceMode === "cloud" && canActivateWarmRun ? input.customImageId : undefined, signal_report: input.signalReportId ?? undefined, channel: input.channelId ?? undefined, - runtime: "acp", + runtime: input.runtime ?? "acp", pending_user_message: warmPayload?.pendingUserMessage, pending_user_artifact_ids: warmPayload?.pendingUserArtifactIds, // If creation activates a pre-warmed run, this is the only request // that can carry the choice — the saga skips run creation entirely. auto_publish: - input.workspaceMode === "cloud" && input.cloudAutoPublish + input.workspaceMode === "cloud" && + canActivateWarmRun && + input.cloudAutoPublish ? true : undefined, }); diff --git a/packages/core/src/task-detail/taskService.test.ts b/packages/core/src/task-detail/taskService.test.ts index 7883f3f4b5..9fa97e251f 100644 --- a/packages/core/src/task-detail/taskService.test.ts +++ b/packages/core/src/task-detail/taskService.test.ts @@ -54,6 +54,57 @@ function makeService(): TaskService { return new TaskService(host, sessionService, effects, piRunner, rootLogger); } +describe("TaskService.openTask", () => { + it("resumes a completed cloud Pi run without starting a local runtime", async () => { + const resumedRun = { + id: "run-1", + environment: "cloud", + status: "queued", + }; + const completedRun = { + id: "run-1", + environment: "cloud", + status: "completed", + }; + const api = { + getTask: vi.fn(async () => ({ + id: "task-1", + runtime: "pi", + latest_run: completedRun, + })), + getTaskRun: vi.fn(async () => completedRun), + resumeRunInCloud: vi.fn(async () => resumedRun), + }; + const workspace = { folderPath: "/repo" }; + const host = { + getAuthenticatedClient: vi.fn(async () => api), + getWorkspace: vi.fn(async () => workspace), + } as unknown as ITaskCreationHost; + const piRunner = { + create: vi.fn(), + resume: vi.fn(), + stop: vi.fn(), + } as unknown as PiRunner; + const service = new TaskService( + host, + {} as SessionService, + {} as TaskCreationEffects, + piRunner, + rootLogger, + ); + + const result = await service.openTask("task-1", "run-1"); + + expect(result.success).toBe(true); + expect(api.resumeRunInCloud).toHaveBeenCalledWith("task-1", "run-1"); + expect(piRunner.resume).not.toHaveBeenCalled(); + if (result.success) { + expect(result.data.task.latest_run).toBe(resumedRun); + expect(result.data.workspace).toBe(workspace); + } + }); +}); + describe("TaskService.createTask validation", () => { it("rejects an input with neither content nor a taskDescription", async () => { const result = await makeService().createTask({ diff --git a/packages/core/src/task-detail/taskService.ts b/packages/core/src/task-detail/taskService.ts index 73ea84990f..3d7eec16f0 100644 --- a/packages/core/src/task-detail/taskService.ts +++ b/packages/core/src/task-detail/taskService.ts @@ -14,7 +14,6 @@ import { inject, injectable } from "inversify"; import { PI_RUNNER } from "../pi-runtime/identifiers"; import type { PiRunner } from "../pi-runtime/piRunner"; import { TASK_CREATION_EFFECTS, TASK_CREATION_HOST } from "./identifiers"; -import { PiTaskCreator } from "./piTaskCreator"; import type { TaskCreationEffects } from "./taskCreationEffects"; import type { ITaskCreationHost } from "./taskCreationHost"; import { TaskCreationSaga } from "./taskCreationSaga"; @@ -108,31 +107,18 @@ export class TaskService { } } - let result: CreateTaskResult; - if (input.runtime === "pi") { - const creator = new PiTaskCreator( - { - posthogClient, - host: this.host, - piRunner: this.piRunner, - onTaskReady, - }, - this.log, - ); - result = await creator.run(input); - } else { - const creator = new TaskCreationSaga( - { - posthogClient, - host: this.host, - sessionService: this.sessionService, - track: (event, props) => this.host.track(event, props), - onTaskReady, - }, - this.log, - ); - result = await creator.run(input); - } + const creator = new TaskCreationSaga( + { + posthogClient, + host: this.host, + sessionService: this.sessionService, + piRunner: this.piRunner, + track: (event, props) => this.host.track(event, props), + onTaskReady, + }, + this.log, + ); + const result = await creator.run(input); if (result.success) { this.effects.onWorkspaceCreated(result.data); @@ -142,6 +128,20 @@ export class TaskService { return result; } + public async getTask(taskId: string, taskRunId?: string): Promise { + const posthogClient = await this.host.getAuthenticatedClient(); + if (!posthogClient) { + throw new Error("Not authenticated"); + } + + const task = await posthogClient.getTask(taskId); + if (taskRunId) { + task.latest_run = await posthogClient.getTaskRun(taskId, taskRunId); + } + + return task; + } + public async openTask( taskId: string, taskRunId?: string, @@ -174,6 +174,35 @@ export class TaskService { const runtime = task.runtime === "pi" ? "pi" : "acp"; const existingWorkspace = await this.host.getWorkspace(taskId); + if (runtime === "pi" && task.latest_run?.environment === "cloud") { + try { + if ( + task.latest_run.status === "completed" || + task.latest_run.status === "failed" || + task.latest_run.status === "cancelled" + ) { + task.latest_run = await posthogClient.resumeRunInCloud( + taskId, + task.latest_run.id, + ); + } + + return { + success: true, + data: { task, workspace: existingWorkspace }, + }; + } catch (error) { + return { + success: false, + error: + error instanceof Error + ? error.message + : "Failed to resume cloud Pi session", + failedStep: "pi_session", + }; + } + } + if (existingWorkspace) { this.log.info("Workspace already exists, fetching task only", { taskId }); try { @@ -223,6 +252,7 @@ export class TaskService { posthogClient, host: this.host, sessionService: this.sessionService, + piRunner: this.piRunner, track: (event, props) => this.host.track(event, props), }, this.log, diff --git a/packages/harness/src/extensions/posthog-provider/provider.ts b/packages/harness/src/extensions/posthog-provider/provider.ts index 07980c777b..dc8deaf166 100644 --- a/packages/harness/src/extensions/posthog-provider/provider.ts +++ b/packages/harness/src/extensions/posthog-provider/provider.ts @@ -1,6 +1,10 @@ -import type { Api, Model, OAuthCredentials } from "@earendil-works/pi-ai"; import type { - AuthStorage, + Api, + CredentialStore, + Model, + OAuthCredentials, +} from "@earendil-works/pi-ai"; +import type { ProviderConfig, ProviderModelConfig, } from "@earendil-works/pi-coding-agent"; @@ -36,11 +40,14 @@ export function parsePosthogOAuthCredentials( : null; } -export function setPosthogOAuthCredentials( - storage: AuthStorage, +export async function setPosthogOAuthCredentials( + storage: CredentialStore, credentials: PosthogOAuthCredentials, -): void { - storage.set(POSTHOG_PROVIDER_NAME, { type: "oauth", ...credentials }); +): Promise { + await storage.modify(POSTHOG_PROVIDER_NAME, async () => ({ + type: "oauth", + ...credentials, + })); } /** diff --git a/packages/harness/src/runtime.test.ts b/packages/harness/src/runtime.test.ts index 335a50b543..375d3f7f0a 100644 --- a/packages/harness/src/runtime.test.ts +++ b/packages/harness/src/runtime.test.ts @@ -2,6 +2,7 @@ import { existsSync } from "node:fs"; import { mkdtemp, readFile, rm, writeFile } from "node:fs/promises"; import { tmpdir } from "node:os"; import { join } from "node:path"; +import { InMemoryCredentialStore } from "@earendil-works/pi-ai"; import { afterEach, describe, expect, it, vi } from "vitest"; import { createHarnessRuntime } from "./runtime"; @@ -34,7 +35,7 @@ describe("createHarnessRuntime", () => { const runtime = await createHarnessRuntime({ agentDir, - authStorage: pi.AuthStorage.inMemory(), + credentialStore: new InMemoryCredentialStore(), cwd, sessionManager: pi.SessionManager.inMemory(cwd), }); @@ -42,6 +43,7 @@ describe("createHarnessRuntime", () => { try { expect(runtime).toBeInstanceOf(pi.AgentSessionRuntime); expect(runtime.session.model?.provider).toBe("posthog"); + expect(runtime.session.getAvailableThinkingLevels()).toContain("off"); expect(runtime.services.settingsManager.isProjectTrusted()).toBe(false); expect( runtime.services.resourceLoader @@ -82,11 +84,9 @@ describe("createHarnessRuntime", () => { }); try { - expect(runtime.services.authStorage.get("posthog")).toMatchObject({ - type: "oauth", - access: "access-token", - refresh: "refresh-token", - }); + await expect( + runtime.services.modelRuntime.getAuth("posthog"), + ).resolves.toMatchObject({ auth: { apiKey: "access-token" } }); expect(existsSync(join(agentDir, "auth.json"))).toBe(false); } finally { await runtime.dispose(); @@ -123,14 +123,12 @@ describe("createHarnessRuntime", () => { }); try { - expect(runtime.services.authStorage.get("anthropic")).toMatchObject({ - type: "api_key", - key: "anthropic-key", - }); - expect(runtime.services.authStorage.get("posthog")).toMatchObject({ - access: "access-token", - refresh: "refresh-token", - }); + await expect( + runtime.services.modelRuntime.getAuth("anthropic"), + ).resolves.toMatchObject({ auth: { apiKey: "anthropic-key" } }); + await expect( + runtime.services.modelRuntime.getAuth("posthog"), + ).resolves.toMatchObject({ auth: { apiKey: "access-token" } }); expect(JSON.parse(await readFile(authPath, "utf8"))).toEqual( storedCredentials, ); @@ -166,18 +164,23 @@ describe("createHarnessRuntime", () => { try { await expect( - runtime.services.modelRegistry.getApiKeyForProvider("posthog"), - ).resolves.toBe("proxy-key"); + runtime.services.modelRuntime.getAuth("posthog"), + ).resolves.toMatchObject({ auth: { apiKey: "proxy-key" } }); } finally { await runtime.dispose(); } }); - it("uses file-backed auth storage when no desktop credentials are provided", async () => { + it("uses file-backed credentials when desktop credentials are absent", async () => { vi.stubEnv("PI_OFFLINE", "1"); const pi = await import("@earendil-works/pi-coding-agent"); const cwd = await temporaryDirectory(); const agentDir = await temporaryDirectory(); + const authPath = join(agentDir, "auth.json"); + await writeFile( + authPath, + JSON.stringify({ posthog: { type: "api_key", key: "stored-key" } }), + ); const runtime = await createHarnessRuntime({ agentDir, @@ -186,14 +189,9 @@ describe("createHarnessRuntime", () => { }); try { - runtime.services.authStorage.set("posthog", { - type: "oauth", - access: "access-token", - refresh: "refresh-token", - expires: Date.now() + 60_000, - }); - - expect(existsSync(join(agentDir, "auth.json"))).toBe(true); + await expect( + runtime.services.modelRuntime.listCredentials(), + ).resolves.toContainEqual({ providerId: "posthog", type: "api_key" }); } finally { await runtime.dispose(); } diff --git a/packages/harness/src/runtime.ts b/packages/harness/src/runtime.ts index 91b691c23d..8692a42b3f 100644 --- a/packages/harness/src/runtime.ts +++ b/packages/harness/src/runtime.ts @@ -1,8 +1,12 @@ import { readFileSync } from "node:fs"; import { join } from "node:path"; +import { + type Credential, + type CredentialStore, + InMemoryCredentialStore, +} from "@earendil-works/pi-ai"; import type { AgentSessionRuntime, - AuthStorage, CreateAgentSessionFromServicesOptions, CreateAgentSessionRuntimeFactory, CreateAgentSessionServicesOptions, @@ -16,19 +20,28 @@ import { import type { HarnessExtensionOptions } from "./extensions/registry"; type PiRuntimeTarget = Parameters[0]; -type AuthStorageSnapshot = Parameters[0]; +type CredentialSnapshot = Record; -function loadAuthStorageSnapshot( - authPath: string, -): AuthStorageSnapshot | undefined { +function loadCredentialSnapshot(authPath: string): CredentialSnapshot { try { - return JSON.parse(readFileSync(authPath, "utf8")) as AuthStorageSnapshot; + return JSON.parse(readFileSync(authPath, "utf8")) as CredentialSnapshot; } catch { - return undefined; + return {}; } } +async function createCredentialStore( + snapshot: CredentialSnapshot, +): Promise { + const store = new InMemoryCredentialStore(); + for (const [providerId, credential] of Object.entries(snapshot)) { + await store.modify(providerId, async () => credential); + } + return store; +} + export type HarnessRuntimeOptions = HarnessExtensionOptions & { + credentialStore?: CredentialStore; posthogOAuthCredentials?: PosthogOAuthCredentials; } & Partial< Pick< @@ -54,7 +67,8 @@ export type HarnessRuntimeOptions = HarnessExtensionOptions & { export async function createHarnessRuntime( options: HarnessRuntimeOptions = {}, ): Promise { - const { posthogOAuthCredentials, ...runtimeOptions } = options; + const { credentialStore, posthogOAuthCredentials, ...runtimeOptions } = + options; // Pi reads its application branding when the SDK is first evaluated. Keep // every runtime import below dynamic so this always happens first. installHogBrandEnv(); @@ -75,23 +89,26 @@ export async function createHarnessRuntime( sessionStartEvent, }) => { const authPath = join(runtimeAgentDir, "auth.json"); - const authStorage = - runtimeOptions.authStorage ?? + const credentials = + credentialStore ?? (posthogOAuthCredentials - ? pi.AuthStorage.inMemory(loadAuthStorageSnapshot(authPath)) - : pi.AuthStorage.create(authPath)); - if (posthogOAuthCredentials) { - setPosthogOAuthCredentials(authStorage, posthogOAuthCredentials); - } - if (options.apiKey) { - authStorage.setRuntimeApiKey(POSTHOG_PROVIDER_NAME, options.apiKey); + ? await createCredentialStore(loadCredentialSnapshot(authPath)) + : undefined); + if (credentials && posthogOAuthCredentials) { + await setPosthogOAuthCredentials(credentials, posthogOAuthCredentials); } + const modelRuntime = + runtimeOptions.modelRuntime ?? + (await pi.ModelRuntime.create({ + authPath, + credentials, + })); const services = await pi.createAgentSessionServices({ ...runtimeOptions, cwd: runtimeCwd, agentDir: runtimeAgentDir, - authStorage, + modelRuntime, settingsManager: options.settingsManager ?? pi.SettingsManager.create(runtimeCwd, runtimeAgentDir, { @@ -106,13 +123,20 @@ export async function createHarnessRuntime( }, }); - const preferredModel = services.modelRegistry.find( - "posthog", + if (options.apiKey) { + await services.modelRuntime.setRuntimeApiKey( + POSTHOG_PROVIDER_NAME, + options.apiKey, + ); + } + + const preferredModel = services.modelRuntime.getModel( + POSTHOG_PROVIDER_NAME, DEFAULT_MODEL, ); - const fallbackModel = services.modelRegistry - .getAll() - .find((model) => model.provider === "posthog"); + const fallbackModel = services.modelRuntime + .getModels(POSTHOG_PROVIDER_NAME) + .at(0); const created = await pi.createAgentSessionFromServices({ ...runtimeOptions, diff --git a/packages/host-router/package.json b/packages/host-router/package.json index c4df528c7c..40c2d7cf4c 100644 --- a/packages/host-router/package.json +++ b/packages/host-router/package.json @@ -17,6 +17,7 @@ "dependencies": { "@agentclientprotocol/sdk": "0.22.1", "@json-render/core": "^0.19.0", + "@posthog/agent": "workspace:*", "@posthog/core": "workspace:*", "@posthog/di": "workspace:*", "@posthog/host-trpc": "workspace:*", diff --git a/packages/host-router/src/cloud-task-client.ts b/packages/host-router/src/cloud-task-client.ts new file mode 100644 index 0000000000..6fc9534b91 --- /dev/null +++ b/packages/host-router/src/cloud-task-client.ts @@ -0,0 +1,47 @@ +import type { CloudTaskClient } from "@posthog/core/cloud-task/cloudTaskClient"; +import type { SendCommandInput } from "@posthog/core/cloud-task/schemas"; +import type { CloudTaskUpdatePayload } from "@posthog/shared/domain-types"; +import { inject, injectable } from "inversify"; +import { HOST_TRPC_CLIENT, type HostTrpcClient } from "./client"; + +@injectable() +export class TrpcCloudTaskClient implements CloudTaskClient { + constructor( + @inject(HOST_TRPC_CLIENT) private readonly client: HostTrpcClient, + ) {} + + getContext(): Promise<{ apiHost: string; teamId: number } | null> { + return this.client.cloudTask.context.query(); + } + + async watch(input: { + taskId: string; + runId: string; + apiHost: string; + teamId: number; + }): Promise { + await this.client.cloudTask.watch.mutate(input); + } + + async unwatch(taskId: string, runId: string): Promise { + await this.client.cloudTask.unwatch.mutate({ taskId, runId }); + } + + subscribe( + taskId: string, + runId: string, + onUpdate: (update: CloudTaskUpdatePayload) => void, + onError: (error: unknown) => void, + onStarted: () => void, + ): () => void { + const subscription = this.client.cloudTask.onUpdate.subscribe( + { taskId, runId }, + { onData: onUpdate, onError, onStarted }, + ); + return () => subscription.unsubscribe(); + } + + sendCommand(input: SendCommandInput) { + return this.client.cloudTask.sendCommand.mutate(input); + } +} diff --git a/apps/code/src/renderer/platform-adapters/trpc-pi-runner.ts b/packages/host-router/src/pi-runner.ts similarity index 100% rename from apps/code/src/renderer/platform-adapters/trpc-pi-runner.ts rename to packages/host-router/src/pi-runner.ts diff --git a/packages/host-router/src/pi-session-client.ts b/packages/host-router/src/pi-session-client.ts deleted file mode 100644 index d168a0c914..0000000000 --- a/packages/host-router/src/pi-session-client.ts +++ /dev/null @@ -1,96 +0,0 @@ -import type { PiSessionClient } from "@posthog/core/pi-runtime/piSessionController"; -import { inject, injectable } from "inversify"; -import { HOST_TRPC_CLIENT, type HostTrpcClient } from "./client"; - -@injectable() -export class TrpcPiSessionClient implements PiSessionClient { - constructor( - @inject(HOST_TRPC_CLIENT) private readonly client: HostTrpcClient, - ) {} - - health(taskId: string) { - return this.client.piSession.health.query({ taskId }); - } - - conversation(taskId: string) { - return this.client.piSession.conversation.query({ taskId }); - } - - status(taskId: string) { - return this.client.piSession.status.query({ taskId }); - } - - availableModels(taskId: string) { - return this.client.piSession.availableModels.query({ taskId }); - } - - commands(taskId: string) { - return this.client.piSession.commands.query({ taskId }); - } - - subscribe( - taskId: string, - onEvent: Parameters[1], - onError: Parameters[2], - ): () => void { - const subscription = this.client.piSession.onEvent.subscribe( - { taskId }, - { onData: onEvent, onError }, - ); - - return () => subscription.unsubscribe(); - } - - prompt(taskId: string, prompt: string) { - return this.client.piSession.prompt.mutate({ taskId, prompt }); - } - - steer(taskId: string, message: string) { - return this.client.piSession.steer.mutate({ taskId, message }); - } - - followUp(taskId: string, message: string) { - return this.client.piSession.followUp.mutate({ taskId, message }); - } - - compact(taskId: string, customInstructions?: string) { - return this.client.piSession.compact.mutate({ taskId, customInstructions }); - } - - setModel(taskId: string, provider: string, modelId: string) { - return this.client.piSession.setModel.mutate({ taskId, provider, modelId }); - } - - setThinkingLevel( - taskId: string, - level: Parameters[1], - ) { - return this.client.piSession.setThinkingLevel.mutate({ taskId, level }); - } - - setSteeringMode( - taskId: string, - mode: Parameters[1], - ) { - return this.client.piSession.setSteeringMode.mutate({ taskId, mode }); - } - - setFollowUpMode( - taskId: string, - mode: Parameters[1], - ) { - return this.client.piSession.setFollowUpMode.mutate({ taskId, mode }); - } - - bash(taskId: string, command: string) { - return this.client.piSession.bash.mutate({ taskId, command }); - } - - abort(taskId: string) { - return this.client.piSession.abort.mutate({ taskId }); - } - - abortBash(taskId: string) { - return this.client.piSession.abortBash.mutate({ taskId }); - } -} diff --git a/packages/host-router/src/pi-session-factory.ts b/packages/host-router/src/pi-session-factory.ts new file mode 100644 index 0000000000..53548e8a93 --- /dev/null +++ b/packages/host-router/src/pi-session-factory.ts @@ -0,0 +1,61 @@ +import { + getRemotePiConversation, + RemotePiRpcClient, +} from "@posthog/agent/pi/remote-rpc-client"; +import type { + PiSession, + PiSessionFactory, +} from "@posthog/core/pi-runtime/piSessionController"; +import { inject, injectable } from "inversify"; +import type { HostTrpcClient } from "./client"; +import { HOST_TRPC_CLIENT } from "./client"; + +class TrpcPiSession implements PiSession { + readonly client: RemotePiRpcClient; + + constructor( + private readonly hostClient: HostTrpcClient, + private readonly taskId: string, + ) { + this.client = new RemotePiRpcClient({ + request: async (command) => { + const response = await this.hostClient.piSession.rpc.mutate({ + taskId: this.taskId, + command, + }); + return response; + }, + }); + } + + health() { + return this.hostClient.piSession.health.query({ taskId: this.taskId }); + } + + getConversation() { + return getRemotePiConversation(this.client); + } + + onConversationEvent( + onEvent: Parameters[0], + onError: Parameters[1], + ): () => void { + const subscription = this.hostClient.piSession.onEvent.subscribe( + { taskId: this.taskId }, + { onData: onEvent, onError }, + ); + + return () => subscription.unsubscribe(); + } +} + +@injectable() +export class TrpcPiSessionFactory implements PiSessionFactory { + constructor( + @inject(HOST_TRPC_CLIENT) private readonly client: HostTrpcClient, + ) {} + + get(taskId: string): Promise { + return Promise.resolve(new TrpcPiSession(this.client, taskId)); + } +} diff --git a/packages/host-router/src/routers/cloud-task.router.ts b/packages/host-router/src/routers/cloud-task.router.ts index 15d577ce59..57061998d3 100644 --- a/packages/host-router/src/routers/cloud-task.router.ts +++ b/packages/host-router/src/routers/cloud-task.router.ts @@ -2,6 +2,7 @@ import type { CloudTaskService } from "@posthog/core/cloud-task/cloud-task"; import { CLOUD_TASK_SERVICE } from "@posthog/core/cloud-task/identifiers"; import { CloudTaskEvent, + cloudContextOutput, designateRelayedMcpServersInput, onUpdateInput, retryInput, @@ -15,6 +16,12 @@ import { import { publicProcedure, router } from "@posthog/host-trpc/trpc"; export const cloudTaskRouter = router({ + context: publicProcedure + .output(cloudContextOutput) + .query(({ ctx }) => + ctx.container.get(CLOUD_TASK_SERVICE).getCloudContext(), + ), + watch: publicProcedure .input(watchInput) .mutation(({ ctx, input }) => diff --git a/packages/host-router/src/routers/pi-session.router.ts b/packages/host-router/src/routers/pi-session.router.ts index 625fc0e009..067948189b 100644 --- a/packages/host-router/src/routers/pi-session.router.ts +++ b/packages/host-router/src/routers/pi-session.router.ts @@ -2,40 +2,11 @@ import { publicProcedure, router } from "@posthog/host-trpc/trpc"; import { PI_SESSION_SERVICE } from "@posthog/workspace-server/services/pi-session/identifiers"; import type { PiSessionService } from "@posthog/workspace-server/services/pi-session/pi-session"; import { - piConversationOutput, - piSessionAvailableModelsOutput, - piSessionBashInput, - piSessionBashOutput, - piSessionCancelledOutput, - piSessionCommandsOutput, - piSessionCompactInput, - piSessionCycleModelOutput, - piSessionEnabledInput, - piSessionEntriesInput, - piSessionEntryInput, - piSessionExportInput, - piSessionExportOutput, - piSessionForkMessagesOutput, - piSessionForkOutput, + piRpcResponseSchema, piSessionHealthOutput, - piSessionLastAssistantTextOutput, - piSessionMessageInput, - piSessionModelInput, - piSessionModelOutput, - piSessionNameInput, - piSessionNewInput, - piSessionPathInput, - piSessionPromptAndWaitInput, - piSessionPromptInput, - piSessionQueueModeInput, + piSessionRpcInput, piSessionStartOutput, - piSessionStatusOutput, - piSessionStderrOutput, - piSessionThinkingCycleOutput, - piSessionThinkingLevelInput, - piSessionTimeoutInput, - piSessionTranscriptInput, - piSessionUnknownOutput, + piSessionTaskInput, resumePiSessionInput, startPiSessionInput, } from "@posthog/workspace-server/services/pi-session/schemas"; @@ -53,270 +24,24 @@ export const piSessionRouter = router({ .input(resumePiSessionInput) .mutation(({ ctx, input }) => getService(ctx.container).resume(input)), - prompt: publicProcedure - .input(piSessionPromptInput) + rpc: publicProcedure + .input(piSessionRpcInput) + .output(piRpcResponseSchema) .mutation(({ ctx, input }) => - getService(ctx.container).prompt( - input.taskId, - input.prompt, - input.images, - ), + getService(ctx.container).request(input.taskId, input.command), ), - steer: publicProcedure - .input(piSessionMessageInput) - .mutation(({ ctx, input }) => - getService(ctx.container).steer( - input.taskId, - input.message, - input.images, - ), - ), - - followUp: publicProcedure - .input(piSessionMessageInput) - .mutation(({ ctx, input }) => - getService(ctx.container).followUp( - input.taskId, - input.message, - input.images, - ), - ), - - abort: publicProcedure - .input(piSessionTranscriptInput) - .mutation(({ ctx, input }) => - getService(ctx.container).abort(input.taskId), - ), - - newSession: publicProcedure - .input(piSessionNewInput) - .output(piSessionCancelledOutput) - .mutation(({ ctx, input }) => - getService(ctx.container).newSession(input.taskId, input.parentSession), - ), - - setModel: publicProcedure - .input(piSessionModelInput) - .output(piSessionModelOutput) - .mutation(({ ctx, input }) => - getService(ctx.container).setModel( - input.taskId, - input.provider, - input.modelId, - ), - ), - - cycleModel: publicProcedure - .input(piSessionTranscriptInput) - .output(piSessionCycleModelOutput) - .mutation(({ ctx, input }) => - getService(ctx.container).cycleModel(input.taskId), - ), - - availableModels: publicProcedure - .input(piSessionTranscriptInput) - .output(piSessionAvailableModelsOutput) - .query(({ ctx, input }) => - getService(ctx.container).availableModels(input.taskId), - ), - - setThinkingLevel: publicProcedure - .input(piSessionThinkingLevelInput) - .mutation(({ ctx, input }) => - getService(ctx.container).setThinkingLevel(input.taskId, input.level), - ), - - cycleThinkingLevel: publicProcedure - .input(piSessionTranscriptInput) - .output(piSessionThinkingCycleOutput) - .mutation(({ ctx, input }) => - getService(ctx.container).cycleThinkingLevel(input.taskId), - ), - - setSteeringMode: publicProcedure - .input(piSessionQueueModeInput) - .mutation(({ ctx, input }) => - getService(ctx.container).setSteeringMode(input.taskId, input.mode), - ), - - setFollowUpMode: publicProcedure - .input(piSessionQueueModeInput) - .mutation(({ ctx, input }) => - getService(ctx.container).setFollowUpMode(input.taskId, input.mode), - ), - - compact: publicProcedure - .input(piSessionCompactInput) - .output(piSessionUnknownOutput) - .mutation(({ ctx, input }) => - getService(ctx.container).compact(input.taskId, input.customInstructions), - ), - - setAutoCompaction: publicProcedure - .input(piSessionEnabledInput) - .mutation(({ ctx, input }) => - getService(ctx.container).setAutoCompaction(input.taskId, input.enabled), - ), - - setAutoRetry: publicProcedure - .input(piSessionEnabledInput) - .mutation(({ ctx, input }) => - getService(ctx.container).setAutoRetry(input.taskId, input.enabled), - ), - - abortRetry: publicProcedure - .input(piSessionTranscriptInput) - .mutation(({ ctx, input }) => - getService(ctx.container).abortRetry(input.taskId), - ), - - bash: publicProcedure - .input(piSessionBashInput) - .output(piSessionBashOutput) - .mutation(({ ctx, input }) => - getService(ctx.container).bash(input.taskId, input.command), - ), - - abortBash: publicProcedure - .input(piSessionTranscriptInput) - .mutation(({ ctx, input }) => - getService(ctx.container).abortBash(input.taskId), - ), - - sessionStats: publicProcedure - .input(piSessionTranscriptInput) - .output(piSessionUnknownOutput) - .query(({ ctx, input }) => - getService(ctx.container).sessionStats(input.taskId), - ), - - exportHtml: publicProcedure - .input(piSessionExportInput) - .output(piSessionExportOutput) - .mutation(({ ctx, input }) => - getService(ctx.container).exportHtml(input.taskId, input.outputPath), - ), - - switchSession: publicProcedure - .input(piSessionPathInput) - .output(piSessionCancelledOutput) - .mutation(({ ctx, input }) => - getService(ctx.container).switchSession(input.taskId, input.sessionPath), - ), - - fork: publicProcedure - .input(piSessionEntryInput) - .output(piSessionForkOutput) - .mutation(({ ctx, input }) => - getService(ctx.container).fork(input.taskId, input.entryId), - ), - - clone: publicProcedure - .input(piSessionTranscriptInput) - .output(piSessionCancelledOutput) - .mutation(({ ctx, input }) => - getService(ctx.container).clone(input.taskId), - ), - - forkMessages: publicProcedure - .input(piSessionTranscriptInput) - .output(piSessionForkMessagesOutput) - .query(({ ctx, input }) => - getService(ctx.container).forkMessages(input.taskId), - ), - - setSessionName: publicProcedure - .input(piSessionNameInput) - .mutation(({ ctx, input }) => - getService(ctx.container).setSessionName(input.taskId, input.name), - ), - - status: publicProcedure - .input(piSessionTranscriptInput) - .output(piSessionStatusOutput) - .query(({ ctx, input }) => getService(ctx.container).status(input.taskId)), - - conversation: publicProcedure - .input(piSessionTranscriptInput) - .output(piConversationOutput) - .query(({ ctx, input }) => - getService(ctx.container).conversation(input.taskId), - ), - - entries: publicProcedure - .input(piSessionEntriesInput) - .query(({ ctx, input }) => - getService(ctx.container).entries(input.taskId, input.since), - ), - - tree: publicProcedure - .input(piSessionTranscriptInput) - .output(piSessionUnknownOutput) - .query(({ ctx, input }) => getService(ctx.container).tree(input.taskId)), - - lastAssistantText: publicProcedure - .input(piSessionTranscriptInput) - .output(piSessionLastAssistantTextOutput) - .query(({ ctx, input }) => - getService(ctx.container).lastAssistantText(input.taskId), - ), - - messages: publicProcedure - .input(piSessionTranscriptInput) - .output(piSessionUnknownOutput) - .query(({ ctx, input }) => - getService(ctx.container).messages(input.taskId), - ), - - commands: publicProcedure - .input(piSessionTranscriptInput) - .output(piSessionCommandsOutput) - .query(({ ctx, input }) => - getService(ctx.container).commands(input.taskId), - ), - - waitForIdle: publicProcedure - .input(piSessionTimeoutInput) - .mutation(({ ctx, input }) => - getService(ctx.container).waitForIdle(input.taskId, input.timeout), - ), - - collectEvents: publicProcedure - .input(piSessionTimeoutInput) - .output(piSessionUnknownOutput) - .mutation(({ ctx, input }) => - getService(ctx.container).collectEvents(input.taskId, input.timeout), - ), - - promptAndWait: publicProcedure - .input(piSessionPromptAndWaitInput) - .output(piSessionUnknownOutput) - .mutation(({ ctx, input }) => - getService(ctx.container).promptAndWait( - input.taskId, - input.prompt, - input.images, - input.timeout, - ), - ), - - stderr: publicProcedure - .input(piSessionTranscriptInput) - .output(piSessionStderrOutput) - .query(({ ctx, input }) => getService(ctx.container).stderr(input.taskId)), - stop: publicProcedure - .input(piSessionTranscriptInput) + .input(piSessionTaskInput) .mutation(({ ctx, input }) => getService(ctx.container).stop(input.taskId)), health: publicProcedure - .input(piSessionTranscriptInput) + .input(piSessionTaskInput) .output(piSessionHealthOutput) .query(({ ctx, input }) => getService(ctx.container).health(input.taskId)), onEvent: publicProcedure - .input(piSessionTranscriptInput) + .input(piSessionTaskInput) .subscription(async function* (opts) { const service = getService(opts.ctx.container); const iterable = service.toIterable("event", { signal: opts.signal }); diff --git a/packages/shared/src/session-events.ts b/packages/shared/src/session-events.ts index c2a4a375b3..6bf934e2ee 100644 --- a/packages/shared/src/session-events.ts +++ b/packages/shared/src/session-events.ts @@ -1,3 +1,5 @@ +import type { AgentConversationEvent } from "./agent-conversation"; + /** * JSON-RPC message types for ACP protocol communication. * These types are used in both main process (session-manager.ts) @@ -74,6 +76,7 @@ export const IMPORTED_USER_PROMPT_META_KEY = "importedUserPrompt"; export interface StoredLogEntry { type: string; timestamp?: string; + event?: AgentConversationEvent; notification?: { id?: number; method?: string; diff --git a/packages/ui/src/features/pi-sessions/PiSessionView.tsx b/packages/ui/src/features/pi-sessions/PiSessionView.tsx index f0c5894c7e..d094094042 100644 --- a/packages/ui/src/features/pi-sessions/PiSessionView.tsx +++ b/packages/ui/src/features/pi-sessions/PiSessionView.tsx @@ -32,9 +32,10 @@ import { interface PiSessionViewProps { taskId: string; + taskRunId?: string; } -export function PiSessionView({ taskId }: PiSessionViewProps) { +export function PiSessionView({ taskId, taskRunId }: PiSessionViewProps) { const piSessionController = useService( PI_SESSION_CONTROLLER, ); @@ -49,9 +50,9 @@ export function PiSessionView({ taskId }: PiSessionViewProps) { const setMessagingMode = useMessagingModeStore((state) => state.setMode); useEffect(() => { - void piSessionController.ensureConnected(taskId); + void piSessionController.ensureConnected(taskId, taskRunId).catch(() => {}); return () => piSessionController.disconnect(taskId); - }, [piSessionController, taskId]); + }, [piSessionController, taskId, taskRunId]); const sessionAvailable = session?.connectionState === "connected"; const status = session?.status; @@ -192,12 +193,9 @@ export function PiSessionView({ taskId }: PiSessionViewProps) { } const pending = isStreaming || isBashRunning; - const currentModel = session.models.find( - (model) => - model.provider === status.model?.provider && model.id === status.model.id, + const supportsThinking = session.thinkingLevels.some( + (level) => level !== "off", ); - const thinkingLevels = currentModel?.thinkingLevels ?? []; - const supportsThinking = thinkingLevels.some((level) => level !== "off"); const queueMode = messagingMode === "steer" ? status.steeringMode : status.followUpMode; @@ -236,7 +234,7 @@ export function PiSessionView({ taskId }: PiSessionViewProps) { supportsThinking ? ( diff --git a/packages/ui/src/features/sessions/components/buildAgentConversationItems.test.ts b/packages/ui/src/features/sessions/components/buildAgentConversationItems.test.ts index 29d869124b..a767a8aae5 100644 --- a/packages/ui/src/features/sessions/components/buildAgentConversationItems.test.ts +++ b/packages/ui/src/features/sessions/components/buildAgentConversationItems.test.ts @@ -86,6 +86,60 @@ describe("buildAgentConversationItems", () => { }); }); + it("keeps generic extension tool result content for rendering", () => { + const rawOutput = [{ type: "text", text: "Workflow finished" }]; + const result = buildAgentConversationItems( + [ + { + type: "tool_call_started", + timestamp: 1, + toolCall: { + id: "workflow-1", + title: "workflow", + kind: null, + status: "pending", + rawInput: { name: "release" }, + }, + }, + { + type: "tool_call_updated", + timestamp: 2, + toolCall: { + id: "workflow-1", + status: "completed", + rawOutput, + content: [ + { + type: "content", + content: { type: "text", text: "Workflow finished" }, + }, + ], + }, + }, + ], + false, + ); + + expect(result.items).toContainEqual( + expect.objectContaining({ + type: "session_update", + update: expect.objectContaining({ + sessionUpdate: "tool_call", + toolCallId: "workflow-1", + title: "workflow", + status: "completed", + rawOutput, + content: [ + { + type: "content", + content: { type: "text", text: "Workflow finished" }, + }, + ], + }), + }), + ); + }); + it("builds and completes a generic compaction status", () => { const result = buildAgentConversationItems( [ diff --git a/packages/ui/src/features/task-detail/components/TaskDetail.tsx b/packages/ui/src/features/task-detail/components/TaskDetail.tsx index 865d8438c6..18653f9e22 100644 --- a/packages/ui/src/features/task-detail/components/TaskDetail.tsx +++ b/packages/ui/src/features/task-detail/components/TaskDetail.tsx @@ -51,6 +51,17 @@ export function TaskDetail({ channelId, }: TaskDetailProps) { const taskId = initialTask.id; + const selectedTaskRunRef = useRef({ + taskId, + taskRunId: initialTask.latest_run?.id, + }); + if (selectedTaskRunRef.current.taskId !== taskId) { + selectedTaskRunRef.current = { + taskId, + taskRunId: initialTask.latest_run?.id, + }; + } + const selectedTaskRunId = selectedTaskRunRef.current.taskRunId; const { task } = useTaskData({ taskId, initialTask }); const runtime = task.runtime === "pi" ? "pi" : "acp"; @@ -273,7 +284,9 @@ export function TaskDetail({ - {runtime === "pi" && } + {runtime === "pi" && ( + + )} {runtime === "acp" && } diff --git a/packages/ui/src/features/task-detail/components/TaskInput.tsx b/packages/ui/src/features/task-detail/components/TaskInput.tsx index 0c0d51fca6..9e8511f85b 100644 --- a/packages/ui/src/features/task-detail/components/TaskInput.tsx +++ b/packages/ui/src/features/task-detail/components/TaskInput.tsx @@ -392,7 +392,6 @@ export function TaskInput({ const setWorkspaceMode = (mode: WorkspaceMode) => { didResolveWorkspaceModeRef.current = true; - if (mode === "cloud") setRuntime("acp"); setWorkspaceModeState(mode); setLastUsedWorkspaceMode(mode); if (mode !== "cloud") { @@ -1106,7 +1105,7 @@ export function TaskInput({ align="center" className="absolute bottom-full left-0 mb-2 min-w-0" > - {piHarnessEnabled && workspaceMode !== "cloud" && ( + {piHarnessEnabled && ( scopedLogger, +}; + +function successfulResponse(command: string): RpcResponse { + return { + type: "response", + command, + success: true, + } as RpcResponse; +} + +afterEach(() => { + vi.unstubAllEnvs(); + vi.restoreAllMocks(); +}); describe("selectPiPoolEvictionCandidate", () => { it("selects the least recently used idle session", () => { expect( selectPiPoolEvictionCandidate([ - { taskId: "recent", state: "idle", lastUsedAt: 30 }, - { taskId: "oldest", state: "idle", lastUsedAt: 10 }, - { taskId: "middle", state: "idle", lastUsedAt: 20 }, + { + taskId: "recent", + state: "idle", + lastUsedAt: 30, + activeRequestCount: 0, + }, + { + taskId: "oldest", + state: "idle", + lastUsedAt: 10, + activeRequestCount: 0, + }, + { + taskId: "middle", + state: "idle", + lastUsedAt: 20, + activeRequestCount: 0, + }, ]), ).toBe("oldest"); }); - it("pins streaming, starting, and protected sessions", () => { + it("pins streaming, starting, protected, and requested sessions", () => { expect( selectPiPoolEvictionCandidate( [ - { taskId: "streaming", state: "streaming", lastUsedAt: 1 }, - { taskId: "starting", state: "starting", lastUsedAt: 2 }, - { taskId: "protected", state: "idle", lastUsedAt: 3 }, - { taskId: "evictable", state: "idle", lastUsedAt: 4 }, + { + taskId: "streaming", + state: "streaming", + lastUsedAt: 1, + activeRequestCount: 0, + }, + { + taskId: "starting", + state: "starting", + lastUsedAt: 2, + activeRequestCount: 0, + }, + { + taskId: "requested", + state: "idle", + lastUsedAt: 3, + activeRequestCount: 1, + }, + { + taskId: "protected", + state: "idle", + lastUsedAt: 4, + activeRequestCount: 0, + }, + { + taskId: "evictable", + state: "idle", + lastUsedAt: 5, + activeRequestCount: 0, + }, ], "protected", ), @@ -29,9 +101,118 @@ describe("selectPiPoolEvictionCandidate", () => { it("returns null when every session is pinned", () => { expect( selectPiPoolEvictionCandidate([ - { taskId: "streaming", state: "streaming", lastUsedAt: 1 }, - { taskId: "starting", state: "starting", lastUsedAt: 2 }, + { + taskId: "streaming", + state: "streaming", + lastUsedAt: 1, + activeRequestCount: 0, + }, + { + taskId: "starting", + state: "starting", + lastUsedAt: 2, + activeRequestCount: 0, + }, + { + taskId: "requested", + state: "idle", + lastUsedAt: 3, + activeRequestCount: 2, + }, ]), ).toBeNull(); }); }); + +describe("PiSessionService RPC request pinning", () => { + it("keeps a session pinned until every concurrent generic command settles", async () => { + vi.stubEnv("POSTHOG_CODE_PI_HOT_POOL_SIZE", "1"); + let timestamp = 0; + vi.spyOn(Date, "now").mockImplementation(() => timestamp++); + + const requestResolvers: Array<(response: RpcResponse) => void> = []; + const firstClient = { + start: vi.fn().mockResolvedValue(undefined), + stop: vi.fn().mockResolvedValue(undefined), + getState: vi.fn().mockResolvedValue({ + isStreaming: false, + sessionFile: "/tmp/first.jsonl", + }), + send: vi.fn( + () => + new Promise((resolve) => { + requestResolvers.push(resolve); + }), + ), + } as unknown as PiRpcClient; + const secondClient = { + start: vi.fn().mockResolvedValue(undefined), + stop: vi.fn().mockResolvedValue(undefined), + getState: vi.fn().mockResolvedValue({ + isStreaming: false, + sessionFile: "/tmp/second.jsonl", + }), + send: vi.fn(), + } as unknown as PiRpcClient; + const thirdClient = { + start: vi.fn().mockResolvedValue(undefined), + stop: vi.fn().mockResolvedValue(undefined), + getState: vi.fn().mockResolvedValue({ + isStreaming: false, + sessionFile: "/tmp/third.jsonl", + }), + send: vi.fn(), + } as unknown as PiRpcClient; + const clients = [firstClient, secondClient, thirdClient]; + const runtimeFactory = { + create: vi.fn(async () => { + const client = clients.shift() as PiRpcClient; + return { + client, + process: undefined, + onRuntimeEvent: vi.fn(), + onConversationEvent: vi.fn(), + } as unknown as PiRuntime; + }), + } as PiRuntimeFactory; + const taskMetadataRepository = { + findByTaskId: vi.fn((taskId: string) => ({ + piSessionFile: `/tmp/${taskId}.jsonl`, + })), + upsert: vi.fn(), + } as unknown as ITaskMetadataRepository; + const processTracking = { + register: vi.fn(), + unregister: vi.fn(), + } as unknown as ProcessTrackingService; + const service = new PiSessionService( + runtimeFactory, + taskMetadataRepository, + processTracking, + rootLogger, + ); + + await service.resume({ taskId: "first", cwd: "/tmp" }); + const bashRequest = service.request("first", { + type: "bash", + command: "sleep 1", + }); + const compactRequest = service.request("first", { type: "compact" }); + + await service.resume({ taskId: "second", cwd: "/tmp" }); + expect(firstClient.stop).not.toHaveBeenCalled(); + + requestResolvers[0](successfulResponse("bash")); + await bashRequest; + await vi.waitFor(() => expect(secondClient.stop).toHaveBeenCalledOnce()); + expect(firstClient.stop).not.toHaveBeenCalled(); + + requestResolvers[1](successfulResponse("compact")); + await compactRequest; + expect(firstClient.stop).not.toHaveBeenCalled(); + + await service.resume({ taskId: "third", cwd: "/tmp" }); + expect(firstClient.stop).toHaveBeenCalledOnce(); + expect(thirdClient.stop).not.toHaveBeenCalled(); + }); +}); diff --git a/packages/workspace-server/src/services/pi-session/pi-session.ts b/packages/workspace-server/src/services/pi-session/pi-session.ts index 837dee3823..d421049991 100644 --- a/packages/workspace-server/src/services/pi-session/pi-session.ts +++ b/packages/workspace-server/src/services/pi-session/pi-session.ts @@ -1,6 +1,10 @@ import type { PiRpcClient } from "@posthog/agent/pi/rpc-client"; +import { + type RpcCommand, + type RpcResponse, + sendPiRpcCommand, +} from "@posthog/agent/pi/rpc-transport"; import type { PiRuntime } from "@posthog/agent/pi/runtime"; -import type { PiModelOption } from "@posthog/agent/pi/types"; import { ROOT_LOGGER, type RootLogger } from "@posthog/di/logger"; import { type AgentConversationEvent, @@ -21,6 +25,7 @@ interface PiPoolEntry { taskId: string; state: PiPoolSessionState; lastUsedAt: number; + activeRequestCount: number; } export function selectPiPoolEvictionCandidate( @@ -29,7 +34,10 @@ export function selectPiPoolEvictionCandidate( ): string | null { const candidate = entries .filter( - (entry) => entry.taskId !== protectedTaskId && entry.state === "idle", + (entry) => + entry.taskId !== protectedTaskId && + entry.state === "idle" && + entry.activeRequestCount === 0, ) .sort((left, right) => left.lastUsedAt - right.lastUsedAt)[0]; @@ -47,6 +55,7 @@ interface ManagedPiSession { runtime: PiRuntime; state: PiPoolSessionState; lastUsedAt: number; + activeRequestCount: number; pid?: number; } @@ -161,227 +170,27 @@ export class PiSessionService extends TypedEventEmitter { await this.startSession(input.taskId, client, session, async () => {}); } - async prompt( - taskId: string, - prompt: string, - images?: Parameters[1], - ): Promise { - await this.requireSession(taskId).client.prompt(prompt, images); - } - - async steer( - taskId: string, - message: string, - images?: Parameters[1], - ): Promise { - await this.requireSession(taskId).client.steer(message, images); - } - - async followUp( - taskId: string, - message: string, - images?: Parameters[1], - ): Promise { - await this.requireSession(taskId).client.followUp(message, images); - } - - async abort(taskId: string): Promise { - await this.requireSession(taskId).client.abort(); - } - - async newSession( - taskId: string, - parentSession?: string, - ): ReturnType { - const result = - await this.requireSession(taskId).client.newSession(parentSession); - - if (!result.cancelled) { - await this.persistSessionState(taskId); - } - - return result; - } - - setModel( - taskId: string, - provider: string, - modelId: string, - ): ReturnType { - return this.requireSession(taskId).client.setModel(provider, modelId); - } - - cycleModel(taskId: string): ReturnType { - return this.requireSession(taskId).client.cycleModel(); - } - - availableModels(taskId: string): Promise { - return this.requireSession(taskId).runtime.availableModels(); - } - - setThinkingLevel( - taskId: string, - level: Parameters[0], - ): ReturnType { - return this.requireSession(taskId).client.setThinkingLevel(level); - } - - cycleThinkingLevel( - taskId: string, - ): ReturnType { - return this.requireSession(taskId).client.cycleThinkingLevel(); - } - - setSteeringMode( - taskId: string, - mode: Parameters[0], - ): ReturnType { - return this.requireSession(taskId).client.setSteeringMode(mode); - } - - setFollowUpMode( - taskId: string, - mode: Parameters[0], - ): ReturnType { - return this.requireSession(taskId).client.setFollowUpMode(mode); - } - - compact( - taskId: string, - customInstructions?: string, - ): ReturnType { - return this.requireSession(taskId).client.compact(customInstructions); - } - - setAutoCompaction( - taskId: string, - enabled: boolean, - ): ReturnType { - return this.requireSession(taskId).client.setAutoCompaction(enabled); - } - - setAutoRetry( - taskId: string, - enabled: boolean, - ): ReturnType { - return this.requireSession(taskId).client.setAutoRetry(enabled); - } - - abortRetry(taskId: string): ReturnType { - return this.requireSession(taskId).client.abortRetry(); - } - - bash(taskId: string, command: string): ReturnType { - return this.requireSession(taskId).client.bash(command); - } - - abortBash(taskId: string): ReturnType { - return this.requireSession(taskId).client.abortBash(); - } - - sessionStats(taskId: string): ReturnType { - return this.requireSession(taskId).client.getSessionStats(); - } - - exportHtml( - taskId: string, - outputPath?: string, - ): ReturnType { - return this.requireSession(taskId).client.exportHtml(outputPath); - } - - async switchSession( - taskId: string, - sessionPath: string, - ): ReturnType { - const result = - await this.requireSession(taskId).client.switchSession(sessionPath); - - if (!result.cancelled) { - await this.persistSessionState(taskId); - } + async request(taskId: string, command: RpcCommand): Promise { + const session = this.requireSession(taskId); + session.activeRequestCount += 1; - return result; - } - - async fork(taskId: string, entryId: string): ReturnType { - const result = await this.requireSession(taskId).client.fork(entryId); - - if (!result.cancelled) { - await this.persistSessionState(taskId); - } - - return result; - } - - async clone(taskId: string): ReturnType { - const result = await this.requireSession(taskId).client.clone(); + try { + const response = await sendPiRpcCommand(session.client, command); + + if ( + response.success && + ["new_session", "switch_session", "fork", "clone"].includes( + command.type, + ) + ) { + await this.persistSessionState(taskId); + } - if (!result.cancelled) { - await this.persistSessionState(taskId); + return response; + } finally { + session.activeRequestCount -= 1; + void this.enforceHotPoolLimit(); } - - return result; - } - - forkMessages(taskId: string): ReturnType { - return this.requireSession(taskId).client.getForkMessages(); - } - - tree(taskId: string): ReturnType { - return this.requireSession(taskId).client.getTree(); - } - - lastAssistantText( - taskId: string, - ): ReturnType { - return this.requireSession(taskId).client.getLastAssistantText(); - } - - setSessionName( - taskId: string, - name: string, - ): ReturnType { - return this.requireSession(taskId).client.setSessionName(name); - } - - messages(taskId: string): ReturnType { - return this.requireSession(taskId).client.getMessages(); - } - - commands(taskId: string): ReturnType { - return this.requireSession(taskId).client.getCommands(); - } - - waitForIdle( - taskId: string, - timeout?: number, - ): ReturnType { - return this.requireSession(taskId).client.waitForIdle(timeout); - } - - collectEvents( - taskId: string, - timeout?: number, - ): ReturnType { - return this.requireSession(taskId).client.collectEvents(timeout); - } - - promptAndWait( - taskId: string, - prompt: string, - images?: Parameters[1], - timeout?: number, - ): ReturnType { - return this.requireSession(taskId).client.promptAndWait( - prompt, - images, - timeout, - ); - } - - stderr(taskId: string): string { - return this.requireSession(taskId).client.getStderr(); } async stop(taskId: string): Promise { @@ -420,21 +229,6 @@ export class PiSessionService extends TypedEventEmitter { }; } - status(taskId: string): ReturnType { - return this.requireSession(taskId).client.getState(); - } - - conversation(taskId: string): Promise { - return this.requireSession(taskId).runtime.conversation(); - } - - entries( - taskId: string, - since?: string, - ): ReturnType { - return this.requireSession(taskId).client.getEntries(since); - } - async cleanup(): Promise { await Promise.all( [...this.sessions.keys()].map((taskId) => this.stop(taskId)), @@ -513,6 +307,7 @@ export class PiSessionService extends TypedEventEmitter { runtime, state: "starting", lastUsedAt: Date.now(), + activeRequestCount: 0, }; this.sessions.set(taskId, session); @@ -614,6 +409,7 @@ export class PiSessionService extends TypedEventEmitter { taskId, state: session.state, lastUsedAt: session.lastUsedAt, + activeRequestCount: session.activeRequestCount, })), protectedTaskId, ); @@ -621,12 +417,22 @@ export class PiSessionService extends TypedEventEmitter { if (!taskId) { return; } - this.log.info("Evicting least recently used Pi session", { - taskId, - maxHotSessions: this.maxHotSessions, - }); try { - await this.stop(taskId); + await this.runExclusive(taskId, async () => { + const session = this.sessions.get(taskId); + const isEvictable = + session?.state === "idle" && session.activeRequestCount === 0; + + if (!isEvictable || taskId === protectedTaskId) { + return; + } + + this.log.info("Evicting least recently used Pi session", { + taskId, + maxHotSessions: this.maxHotSessions, + }); + await this.stopLocked(taskId); + }); } catch (error) { this.log.warn("Failed to evict Pi session", { taskId, error }); return; diff --git a/packages/workspace-server/src/services/pi-session/schemas.ts b/packages/workspace-server/src/services/pi-session/schemas.ts index 03dcb93c2f..544411d20c 100644 --- a/packages/workspace-server/src/services/pi-session/schemas.ts +++ b/packages/workspace-server/src/services/pi-session/schemas.ts @@ -1,153 +1,10 @@ import { - PI_QUEUE_MODES, - PI_THINKING_LEVELS, - type PiCommand, - type PiModelOption, - type PiSessionStatus, -} from "@posthog/agent/pi/types"; + piRpcCommandSchema, + piRpcResponseSchema, +} from "@posthog/agent/pi/rpc-transport"; import { z } from "zod"; -const agentContent = z.discriminatedUnion("type", [ - z.object({ type: z.literal("text"), text: z.string() }), - z.object({ - type: z.literal("image"), - data: z.string(), - mimeType: z.string(), - }), - z.object({ - type: z.literal("audio"), - data: z.string(), - mimeType: z.string(), - }), - z.object({ - type: z.literal("resource_link"), - uri: z.string(), - name: z.string(), - description: z.string().nullable().optional(), - mimeType: z.string().nullable().optional(), - size: z.number().nullable().optional(), - title: z.string().nullable().optional(), - }), - z.object({ - type: z.literal("resource"), - resource: z.union([ - z.object({ - uri: z.string(), - mimeType: z.string().nullable().optional(), - text: z.string(), - }), - z.object({ - uri: z.string(), - mimeType: z.string().nullable().optional(), - blob: z.string(), - }), - ]), - }), -]); - -const agentToolContent = z.discriminatedUnion("type", [ - z.object({ type: z.literal("content"), content: agentContent }), - z.object({ - type: z.literal("diff"), - path: z.string(), - oldText: z.string().nullable().optional(), - newText: z.string(), - }), - z.object({ type: z.literal("terminal"), terminalId: z.string() }), -]); - -const agentToolCall = z.object({ - id: z.string(), - title: z.string(), - kind: z - .enum([ - "read", - "edit", - "delete", - "move", - "search", - "execute", - "think", - "fetch", - "switch_mode", - "question", - "other", - ]) - .nullable() - .optional(), - status: z - .enum(["pending", "in_progress", "completed", "failed"]) - .nullable() - .optional(), - content: z.array(agentToolContent).optional(), - locations: z - .array( - z.object({ path: z.string(), line: z.number().nullable().optional() }), - ) - .optional(), - rawInput: z.unknown().optional(), - rawOutput: z.unknown().optional(), - parentId: z.string().optional(), -}); - -export const piConversationEvent = z.discriminatedUnion("type", [ - z.object({ - type: z.literal("user_message"), - id: z.string(), - timestamp: z.number(), - content: z.array(agentContent), - }), - z.object({ - type: z.literal("assistant_message_chunk"), - timestamp: z.number(), - content: agentContent, - }), - z.object({ - type: z.literal("assistant_thought_chunk"), - timestamp: z.number(), - content: agentContent, - }), - z.object({ - type: z.literal("tool_call_started"), - timestamp: z.number(), - toolCall: agentToolCall, - }), - z.object({ - type: z.literal("tool_call_updated"), - timestamp: z.number(), - toolCall: agentToolCall.partial().required({ id: true }), - }), - z.object({ - type: z.literal("runtime_status"), - timestamp: z.number(), - status: z.string(), - isComplete: z.boolean().optional(), - error: z.string().optional(), - message: z.string().optional(), - attempt: z.number().optional(), - maxAttempts: z.number().optional(), - delayMs: z.number().optional(), - }), - z.object({ - type: z.literal("runtime_error"), - timestamp: z.number(), - errorType: z.string(), - message: z.string(), - }), - z.object({ - type: z.literal("turn_completed"), - timestamp: z.number(), - stopReason: z.string().optional(), - }), -]); - -export const piConversationOutput = z.array(piConversationEvent); - -export const piImageContent = z.object({ - type: z.literal("image"), - data: z.string(), - mimeType: z.string(), -}); +export { piRpcResponseSchema }; export const startPiSessionInput = z.object({ taskId: z.string(), @@ -156,6 +13,8 @@ export const startPiSessionInput = z.object({ model: z.string().optional(), }); +export type StartPiSessionInput = z.infer; + export const piSessionStartOutput = z.object({ sessionFile: z.string().nullable(), sessionId: z.string(), @@ -167,164 +26,14 @@ export const piSessionHealthOutput = z.object({ lastUsedAt: z.number().optional(), }); -export const piSessionVoidOutput = z.void(); - -export const piSessionCancelledOutput = z.object({ cancelled: z.boolean() }); - -export const piSessionModelOutput = z.object({ - provider: z.string(), - id: z.string(), -}); - -export const piThinkingLevel = z.enum(PI_THINKING_LEVELS); -export const piQueueMode = z.enum(PI_QUEUE_MODES); - -export const piSessionCycleModelOutput = z - .object({ - model: piSessionModelOutput, - thinkingLevel: piThinkingLevel, - isScoped: z.boolean(), - }) - .nullable(); - -export const piSessionAvailableModelsOutput = z.array( - piSessionModelOutput.extend({ - contextWindow: z.number(), - reasoning: z.boolean(), - thinkingLevels: z.array(piThinkingLevel), - }), -) satisfies z.ZodType; - -export const piSessionThinkingCycleOutput = z - .object({ level: piThinkingLevel }) - .nullable(); - -export const piSessionStatusOutput = z.object({ - model: piSessionModelOutput.optional(), - thinkingLevel: piThinkingLevel, - isStreaming: z.boolean(), - isCompacting: z.boolean(), - steeringMode: piQueueMode, - followUpMode: piQueueMode, - sessionFile: z.string().optional(), - sessionId: z.string(), - sessionName: z.string().optional(), - autoCompactionEnabled: z.boolean(), - messageCount: z.number(), - pendingMessageCount: z.number(), -}) satisfies z.ZodType; - -export const piSessionBashOutput = z.object({ - output: z.string(), - exitCode: z.number().optional(), - cancelled: z.boolean(), - truncated: z.boolean(), - fullOutputPath: z.string().optional(), -}); - -export const piSessionExportOutput = z.object({ path: z.string() }); - -export const piSessionForkOutput = z.object({ - text: z.string(), - cancelled: z.boolean(), -}); - -export const piSessionForkMessagesOutput = z.array( - z.object({ entryId: z.string(), text: z.string() }), -); - -export const piSessionCommandsOutput = z.array( - z.object({ - name: z.string(), - description: z.string().optional(), - source: z.enum(["extension", "prompt", "skill"]), - sourceInfo: z.object({ - path: z.string(), - source: z.string(), - scope: z.enum(["user", "project", "temporary"]), - origin: z.enum(["package", "top-level"]), - baseDir: z.string().optional(), - }), - }), -) satisfies z.ZodType; - -export const piSessionLastAssistantTextOutput = z.string().nullable(); -export const piSessionStderrOutput = z.string(); -export const piSessionUnknownOutput = z.unknown(); - export const resumePiSessionInput = z.object({ taskId: z.string(), cwd: z.string(), }); -export const piSessionTranscriptInput = z.object({ taskId: z.string() }); - -export const piSessionPromptInput = piSessionTranscriptInput.extend({ - prompt: z.string().min(1), - images: z.array(piImageContent).optional(), -}); - -export const piSessionMessageInput = piSessionTranscriptInput.extend({ - message: z.string().min(1), - images: z.array(piImageContent).optional(), -}); - -export const piSessionBashInput = piSessionTranscriptInput.extend({ - command: z.string().min(1), -}); - -export const piSessionModelInput = piSessionTranscriptInput.extend({ - provider: z.string().min(1), - modelId: z.string().min(1), -}); - -export const piSessionThinkingLevelInput = piSessionTranscriptInput.extend({ - level: piThinkingLevel, -}); - -export const piSessionQueueModeInput = piSessionTranscriptInput.extend({ - mode: piQueueMode, -}); - -export const piSessionCompactInput = piSessionTranscriptInput.extend({ - customInstructions: z.string().optional(), -}); - -export const piSessionEnabledInput = piSessionTranscriptInput.extend({ - enabled: z.boolean(), -}); - -export const piSessionNewInput = piSessionTranscriptInput.extend({ - parentSession: z.string().optional(), -}); +export const piSessionTaskInput = z.object({ taskId: z.string() }); -export const piSessionPathInput = piSessionTranscriptInput.extend({ - sessionPath: z.string().min(1), -}); - -export const piSessionEntryInput = piSessionTranscriptInput.extend({ - entryId: z.string().min(1), -}); - -export const piSessionNameInput = piSessionTranscriptInput.extend({ - name: z.string().min(1), -}); - -export const piSessionExportInput = piSessionTranscriptInput.extend({ - outputPath: z.string().optional(), -}); - -export const piSessionTimeoutInput = piSessionTranscriptInput.extend({ - timeout: z.number().int().positive().optional(), -}); - -export const piSessionPromptAndWaitInput = piSessionPromptInput.extend({ - timeout: z.number().int().positive().optional(), -}); - -export const piSessionEntriesInput = piSessionTranscriptInput.extend({ - since: z.string().optional(), +export const piSessionRpcInput = z.object({ + taskId: z.string(), + command: piRpcCommandSchema, }); - -export type StartPiSessionInput = z.infer; -export type PiSessionPromptInput = z.infer; diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 4a8c6ced07..4df2224b92 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -7,17 +7,17 @@ settings: catalogs: default: '@earendil-works/pi-agent-core': - specifier: 0.80.6 - version: 0.80.6 + specifier: 0.81.0 + version: 0.81.0 '@earendil-works/pi-ai': - specifier: 0.80.6 - version: 0.80.6 + specifier: 0.81.0 + version: 0.81.0 '@earendil-works/pi-coding-agent': - specifier: 0.80.6 - version: 0.80.6 + specifier: 0.81.0 + version: 0.81.0 '@earendil-works/pi-tui': - specifier: 0.80.6 - version: 0.80.6 + specifier: 0.81.0 + version: 0.81.0 '@hono/node-server': specifier: ^1.13.7 version: 1.19.9 @@ -762,13 +762,13 @@ importers: version: 0.109.0(zod@4.4.3) '@earendil-works/pi-agent-core': specifier: 'catalog:' - version: 0.80.6(@modelcontextprotocol/sdk@1.29.0(zod@4.4.3))(ws@8.19.0)(zod@4.4.3) + version: 0.81.0(@modelcontextprotocol/sdk@1.29.0(zod@4.4.3))(ws@8.19.0)(zod@4.4.3) '@earendil-works/pi-ai': specifier: 'catalog:' - version: 0.80.6(@modelcontextprotocol/sdk@1.29.0(zod@4.4.3))(ws@8.19.0)(zod@4.4.3) + version: 0.81.0(@modelcontextprotocol/sdk@1.29.0(zod@4.4.3))(ws@8.19.0)(zod@4.4.3) '@earendil-works/pi-coding-agent': specifier: 'catalog:' - version: 0.80.6(@modelcontextprotocol/sdk@1.29.0(zod@4.4.3))(ws@8.19.0)(zod@4.4.3) + version: 0.81.0(@modelcontextprotocol/sdk@1.29.0(zod@4.4.3))(ws@8.19.0)(zod@4.4.3) '@hono/node-server': specifier: ^1.19.9 version: 1.19.9(hono@4.11.7) @@ -1041,13 +1041,13 @@ importers: dependencies: '@earendil-works/pi-ai': specifier: 'catalog:' - version: 0.80.6(@modelcontextprotocol/sdk@1.29.0(zod@4.4.3))(ws@8.19.0)(zod@4.4.3) + version: 0.81.0(@modelcontextprotocol/sdk@1.29.0(zod@4.4.3))(ws@8.19.0)(zod@4.4.3) '@earendil-works/pi-coding-agent': specifier: 'catalog:' - version: 0.80.6(@modelcontextprotocol/sdk@1.29.0(zod@4.4.3))(ws@8.19.0)(zod@4.4.3) + version: 0.81.0(@modelcontextprotocol/sdk@1.29.0(zod@4.4.3))(ws@8.19.0)(zod@4.4.3) '@earendil-works/pi-tui': specifier: 'catalog:' - version: 0.80.6 + version: 0.81.0 '@modelcontextprotocol/sdk': specifier: ^1.29.0 version: 1.29.0(zod@4.4.3) @@ -1088,6 +1088,9 @@ importers: '@json-render/core': specifier: ^0.19.0 version: 0.19.0(zod@4.4.3) + '@posthog/agent': + specifier: workspace:* + version: link:../agent '@posthog/core': specifier: workspace:* version: link:../core @@ -2734,22 +2737,22 @@ packages: '@drizzle-team/brocli@0.10.2': resolution: {integrity: sha512-z33Il7l5dKjUgGULTqBsQBQwckHh5AbIuxhdsIxDDiZAzBOrZO6q9ogcWC65kU382AfynTfgNumVcNIjuIua6w==} - '@earendil-works/pi-agent-core@0.80.6': - resolution: {integrity: sha512-Lvn89ko42h5ETUb6Z0Ku6ldskEqXaTdQBYvSa0+7bdG9V6rUEpXptv5e0OVZ1HDcvi8s6/2lGCQWsxKX+DFHNw==} + '@earendil-works/pi-agent-core@0.81.0': + resolution: {integrity: sha512-Fe6JYbW0CGVvmD4dwuFreynwZxlTElIMFRqYDl4llRiwfYZjjtAqbV1hY8MrdGqP15+FySBv5RGodOsHqlhvZQ==} engines: {node: '>=22.19.0'} - '@earendil-works/pi-ai@0.80.6': - resolution: {integrity: sha512-7xfLk8sANBp+bpPEbjoOZTbPxsa+++b1JXAoSJsNa3vbs9AHHEclmvg54XLQcxH+fuwaeti/g2jeIfJ+mVYLpA==} + '@earendil-works/pi-ai@0.81.0': + resolution: {integrity: sha512-n3lDV1Px/2BOp86rUJkoHcXuQ6uf7711VEtSjE5gTwrnPeMAEjzV3LnbQ3wTF9bUxl253Igi3U35U1CMF5POng==} engines: {node: '>=22.19.0'} hasBin: true - '@earendil-works/pi-coding-agent@0.80.6': - resolution: {integrity: sha512-vcfD6tOk402isLl3Cm/qbn2O10TvgroMp1+/fEGM24ZdvETFCdOYv5VZ7m59EI5fPsjfSJh+CpQ5bhBrhfOg7g==} + '@earendil-works/pi-coding-agent@0.81.0': + resolution: {integrity: sha512-2p0Dnx+3fkPLga8M82eg14ZYNLcFLhqxxKyVVfqUSIio9Xx4p7UjvJtopx/6PTeJKmdJl1/xOm/c02AFcJ+l/g==} engines: {node: '>=22.19.0'} hasBin: true - '@earendil-works/pi-tui@0.80.6': - resolution: {integrity: sha512-bSuzS4EVSqEPj/Qr/p9eqCESfKsGuDNbl77EGci8Iaqqt/C/XCBZL1MjXaxSWW1NsT5afjp/Cb0NTPzOLv/aPA==} + '@earendil-works/pi-tui@0.81.0': + resolution: {integrity: sha512-vdrrV//CldG3Dt20vj+iCRcHOlPwxlOBJwwLg1KE/92IO+TlEwbf4uwKXgFZHGVh9lEJXtYsKzIRb/nHw5YuAA==} engines: {node: '>=22.19.0'} '@ecies/ciphers@0.2.6': @@ -9554,8 +9557,8 @@ packages: resolution: {integrity: sha512-0je+qPKHEMohvfRTCEo3CrPG6cAzAYgmzKyxRiYSSDkS6eGJdyVJm7WaYA5ECaAD9wLB2T4EEeymA5aFVcYXCA==} engines: {node: '>=6'} - deslop-js@0.7.7: - resolution: {integrity: sha512-mUr2B3FZM+vJZK4w6ZRwuHm6cXVWM5kS4xGRjRLC0yyen/zD+/ZI7QtH/UqUN0LaEIQZL2+P8PMcsDOIB5xgPQ==} + deslop-js@0.7.8: + resolution: {integrity: sha512-QMmb3Z/ARvYZmZneudb8cnY/4mVvZTdhUyA9TC2skwOcm7KvY9zyOdn0TApQc4rL0VM2TffFkmo3ky/lJZX7qw==} destroy@1.2.0: resolution: {integrity: sha512-2sJGJTaXIIaR1w4iJSNoN0hnMY7Gpc/n8D4qSCJw8QqFWXf7cuAgnEHxBpweaVcPevC2l3KpjYCx3NypQQgaJg==} @@ -12924,8 +12927,8 @@ packages: engines: {node: ^20.19.0 || >=22.12.0} hasBin: true - oxlint-plugin-react-doctor@0.7.7: - resolution: {integrity: sha512-OMl/8k3rcxTSyMJWUL/pdq/ZZKFtv1F4kyM+AWYsj/YYaL/nftL/idrJOxyd2CuFa5Nrbz2DlK3QFRej1uwQ/Q==} + oxlint-plugin-react-doctor@0.7.8: + resolution: {integrity: sha512-3f9/jFLIC/KRLPYqxiXSk20cq47luGy9Oz5Ru7nK7w0EI9B9zuMvAU92bK9jFiwFE7Lc9PA8ER6Y+naYaGFQGw==} engines: {node: ^20.19.0 || >=22.13.0} oxlint@1.66.0: @@ -13571,8 +13574,8 @@ packages: resolution: {integrity: sha512-+NRMYs2DyTP4/tqWz371Oo50JqmWltR1h2gcdgUMAWZJIAvrd0/SqlCfx7tpzpl/s36rzw6qH2MjoNrxtRNYhA==} engines: {node: ^20.9.0 || >=22} - react-doctor@0.7.7: - resolution: {integrity: sha512-kHlzPDZNRKlQbXsXgnFh1I4mL0ybt0xJca3/uRMZafaXR4AG5zf7iFMqtdRUiFWHeefE4vGf4YQmkYkXzE1rlg==} + react-doctor@0.7.8: + resolution: {integrity: sha512-G3spmtZJE/gWWPRJ3rpgUWTPRDJpEmdRja7iNZ7RAXlfpEO+NWVzPTca/cPI9hLwPo2Aq5/BZggo5JDBrwGrlA==} engines: {node: ^20.19.0 || >=22.13.0} hasBin: true @@ -17152,9 +17155,9 @@ snapshots: '@drizzle-team/brocli@0.10.2': {} - '@earendil-works/pi-agent-core@0.80.6(@modelcontextprotocol/sdk@1.29.0(zod@4.4.3))(ws@8.19.0)(zod@4.4.3)': + '@earendil-works/pi-agent-core@0.81.0(@modelcontextprotocol/sdk@1.29.0(zod@4.4.3))(ws@8.19.0)(zod@4.4.3)': dependencies: - '@earendil-works/pi-ai': 0.80.6(@modelcontextprotocol/sdk@1.29.0(zod@4.4.3))(ws@8.19.0)(zod@4.4.3) + '@earendil-works/pi-ai': 0.81.0(@modelcontextprotocol/sdk@1.29.0(zod@4.4.3))(ws@8.19.0)(zod@4.4.3) ignore: 7.0.5 typebox: 1.1.38 yaml: 2.9.0 @@ -17166,7 +17169,7 @@ snapshots: - ws - zod - '@earendil-works/pi-ai@0.80.6(@modelcontextprotocol/sdk@1.29.0(zod@4.4.3))(ws@8.19.0)(zod@4.4.3)': + '@earendil-works/pi-ai@0.81.0(@modelcontextprotocol/sdk@1.29.0(zod@4.4.3))(ws@8.19.0)(zod@4.4.3)': dependencies: '@anthropic-ai/sdk': 0.91.1(zod@4.4.3) '@aws-sdk/client-bedrock-runtime': 3.1048.0 @@ -17187,11 +17190,11 @@ snapshots: - ws - zod - '@earendil-works/pi-coding-agent@0.80.6(@modelcontextprotocol/sdk@1.29.0(zod@4.4.3))(ws@8.19.0)(zod@4.4.3)': + '@earendil-works/pi-coding-agent@0.81.0(@modelcontextprotocol/sdk@1.29.0(zod@4.4.3))(ws@8.19.0)(zod@4.4.3)': dependencies: - '@earendil-works/pi-agent-core': 0.80.6(@modelcontextprotocol/sdk@1.29.0(zod@4.4.3))(ws@8.19.0)(zod@4.4.3) - '@earendil-works/pi-ai': 0.80.6(@modelcontextprotocol/sdk@1.29.0(zod@4.4.3))(ws@8.19.0)(zod@4.4.3) - '@earendil-works/pi-tui': 0.80.6 + '@earendil-works/pi-agent-core': 0.81.0(@modelcontextprotocol/sdk@1.29.0(zod@4.4.3))(ws@8.19.0)(zod@4.4.3) + '@earendil-works/pi-ai': 0.81.0(@modelcontextprotocol/sdk@1.29.0(zod@4.4.3))(ws@8.19.0)(zod@4.4.3) + '@earendil-works/pi-tui': 0.81.0 '@silvia-odwyer/photon-node': 0.3.4 chalk: 5.6.2 cross-spawn: 7.0.6 @@ -17217,7 +17220,7 @@ snapshots: - ws - zod - '@earendil-works/pi-tui@0.80.6': + '@earendil-works/pi-tui@0.81.0': dependencies: get-east-asian-width: 1.6.0 marked: 18.0.5 @@ -24193,7 +24196,7 @@ snapshots: dequal@2.0.3: {} - deslop-js@0.7.7: + deslop-js@0.7.8: dependencies: '@oxc-project/types': 0.138.0 fast-glob: 3.3.3 @@ -28570,7 +28573,7 @@ snapshots: '@oxfmt/binding-win32-ia32-msvc': 0.45.0 '@oxfmt/binding-win32-x64-msvc': 0.45.0 - oxlint-plugin-react-doctor@0.7.7: + oxlint-plugin-react-doctor@0.7.8: dependencies: '@typescript-eslint/types': 8.62.0 eslint-scope: 9.1.2 @@ -29323,19 +29326,19 @@ snapshots: transitivePeerDependencies: - supports-color - react-doctor@0.7.7(@opentelemetry/core@2.8.0(@opentelemetry/api@1.9.1))(eslint@10.5.0(jiti@2.7.0)): + react-doctor@0.7.8(@opentelemetry/core@2.8.0(@opentelemetry/api@1.9.1))(eslint@10.5.0(jiti@2.7.0)): dependencies: '@babel/code-frame': 7.29.0 '@sentry/node': 10.61.0(@opentelemetry/core@2.8.0(@opentelemetry/api@1.9.1)) agent-install: 0.0.5 conf: 15.1.0 confbox: 0.2.4 - deslop-js: 0.7.7 + deslop-js: 0.7.8 eslint-plugin-react-hooks: 7.1.1(eslint@10.5.0(jiti@2.7.0)) jiti: 2.7.0 magicast: 0.5.3 oxlint: 1.66.0 - oxlint-plugin-react-doctor: 0.7.7 + oxlint-plugin-react-doctor: 0.7.8 prompts: 2.4.2 typescript: 5.9.3 vscode-languageserver: 9.0.1 @@ -29605,7 +29608,7 @@ snapshots: preact: 10.29.2 prompts: 2.4.2 react: 19.2.6 - react-doctor: 0.7.7(@opentelemetry/core@2.8.0(@opentelemetry/api@1.9.1))(eslint@10.5.0(jiti@2.7.0)) + react-doctor: 0.7.8(@opentelemetry/core@2.8.0(@opentelemetry/api@1.9.1))(eslint@10.5.0(jiti@2.7.0)) react-dom: 19.2.6(react@19.2.6) react-grab: 0.1.48(react@19.2.6) optionalDependencies: diff --git a/pnpm-workspace.yaml b/pnpm-workspace.yaml index 4b7843a910..52d080ec0c 100644 --- a/pnpm-workspace.yaml +++ b/pnpm-workspace.yaml @@ -6,10 +6,10 @@ packages: catalog: '@hono/node-server': ^1.13.7 '@hono/trpc-server': ^0.3.4 - '@earendil-works/pi-agent-core': 0.80.6 - '@earendil-works/pi-ai': 0.80.6 - '@earendil-works/pi-coding-agent': 0.80.6 - '@earendil-works/pi-tui': 0.80.6 + '@earendil-works/pi-agent-core': 0.81.0 + '@earendil-works/pi-ai': 0.81.0 + '@earendil-works/pi-coding-agent': 0.81.0 + '@earendil-works/pi-tui': 0.81.0 '@parcel/watcher': ^2.5.6 '@phosphor-icons/react': ^2.1.10 '@posthog/quill': 0.3.0-beta.24 From 69695d8471ec66b2ac3e8d5fd0ca8a2a32446ddd Mon Sep 17 00:00:00 2001 From: JonathanLab Date: Thu, 23 Jul 2026 13:44:49 +0200 Subject: [PATCH 02/45] fix: harden cloud Pi session lifecycle --- packages/agent/src/server/pi-agent-server.ts | 17 +++++------------ .../task-detail/components/TaskDetail.tsx | 13 +------------ 2 files changed, 6 insertions(+), 24 deletions(-) diff --git a/packages/agent/src/server/pi-agent-server.ts b/packages/agent/src/server/pi-agent-server.ts index fe6750dc42..f30e788fba 100644 --- a/packages/agent/src/server/pi-agent-server.ts +++ b/packages/agent/src/server/pi-agent-server.ts @@ -13,6 +13,7 @@ import { } from "../pi/rpc-transport"; import { PiRuntime } from "../pi/runtime"; import { PostHogAPIClient } from "../posthog-api"; +import { resolveLlmGatewayUrl } from "../utils/gateway"; import { Logger } from "../utils/logger"; import { TaskRunEventStreamSender } from "./event-stream-sender"; import { type JwtPayload, JwtValidationError, validateJwt } from "./jwt"; @@ -31,7 +32,6 @@ interface PiCloudSession { unsubscribe: () => void; } -const SESSION_SYNC_INTERVAL_MS = 5_000; const COMPLETED_USER_MESSAGE_DELIVERY_LIMIT = 500; const emptySchema = z.object({}); @@ -65,7 +65,6 @@ export class PiAgentServer { private sessionFile: string | null = null; private lastSyncedSessionContent = ""; private sessionRevision = 0; - private sessionSyncInterval: ReturnType | null = null; private sessionSyncQueue: Promise = Promise.resolve(); private pendingLogEntries: StoredLogEntry[] = []; private logFlushQueue: Promise = Promise.resolve(); @@ -117,10 +116,6 @@ export class PiAgentServer { async stop(): Promise { const session = this.session; - if (this.sessionSyncInterval) { - clearInterval(this.sessionSyncInterval); - this.sessionSyncInterval = null; - } if (session) { await session.runtime.client.abort().catch(() => undefined); await session.runtime.client.waitForIdle(5_000).catch(() => undefined); @@ -355,7 +350,10 @@ export class PiAgentServer { sessionFile: restoredSessionFile, providerOptions: { apiKey: this.config.apiKey, - baseUrl: this.posthogAPI.getLlmGatewayUrl(), + baseUrl: resolveLlmGatewayUrl( + process.env.LLM_GATEWAY_URL, + this.config.apiUrl, + ), }, }); const runtime = new PiRuntime(client); @@ -382,11 +380,6 @@ export class PiAgentServer { this.session = { payload, runtime, sseController: null, unsubscribe }; await this.syncTaskSession(); - this.sessionSyncInterval = setInterval(() => { - void this.syncTaskSession().catch((error) => - this.logger.error("Failed to sync active Pi session", error), - ); - }, SESSION_SYNC_INTERVAL_MS); this.sessionReadyBootMs = Math.round(process.uptime() * 1000); this.sessionInitMs = Date.now() - startedAt; await this.posthogAPI.updateTaskRun(payload.task_id, payload.run_id, { diff --git a/packages/ui/src/features/task-detail/components/TaskDetail.tsx b/packages/ui/src/features/task-detail/components/TaskDetail.tsx index 18653f9e22..458b6d954a 100644 --- a/packages/ui/src/features/task-detail/components/TaskDetail.tsx +++ b/packages/ui/src/features/task-detail/components/TaskDetail.tsx @@ -51,20 +51,9 @@ export function TaskDetail({ channelId, }: TaskDetailProps) { const taskId = initialTask.id; - const selectedTaskRunRef = useRef({ - taskId, - taskRunId: initialTask.latest_run?.id, - }); - if (selectedTaskRunRef.current.taskId !== taskId) { - selectedTaskRunRef.current = { - taskId, - taskRunId: initialTask.latest_run?.id, - }; - } - const selectedTaskRunId = selectedTaskRunRef.current.taskRunId; - const { task } = useTaskData({ taskId, initialTask }); const runtime = task.runtime === "pi" ? "pi" : "acp"; + const selectedTaskRunId = task.latest_run?.id; const effectiveRepoPath = useCwd(taskId); From e1814d69002272f4a5dd671eb4928a99836ba80f Mon Sep 17 00:00:00 2001 From: JonathanLab Date: Thu, 23 Jul 2026 13:46:31 +0200 Subject: [PATCH 03/45] fix(agent): bind Pi credentials to the configured run --- .../agent/src/server/pi-agent-server.test.ts | 19 +++++++++++++++++++ packages/agent/src/server/pi-agent-server.ts | 17 ++++++++++++++++- 2 files changed, 35 insertions(+), 1 deletion(-) diff --git a/packages/agent/src/server/pi-agent-server.test.ts b/packages/agent/src/server/pi-agent-server.test.ts index 2fba7922b2..215c47aa21 100644 --- a/packages/agent/src/server/pi-agent-server.test.ts +++ b/packages/agent/src/server/pi-agent-server.test.ts @@ -20,6 +20,25 @@ function config(): AgentServerConfig { } describe("PiAgentServer", () => { + it.each([ + ["task", { task_id: "task-2", run_id: "run-1", team_id: 1 }], + ["run", { task_id: "task-1", run_id: "run-2", team_id: 1 }], + ["team", { task_id: "task-1", run_id: "run-1", team_id: 2 }], + ])("rejects a token for a different %s", (_field, identity) => { + const server = new PiAgentServer(config()) as unknown as { + assertConfiguredRun(payload: Record): void; + }; + + expect(() => + server.assertConfiguredRun({ + ...identity, + user_id: 1, + distinct_id: "user-1", + mode: "interactive", + }), + ).toThrow("Token does not match the configured task run"); + }); + it("persists translated Pi events at the turn boundary", async () => { const appendTaskRunLog = vi.fn(async () => ({})); const server = new PiAgentServer(config()) as unknown as { diff --git a/packages/agent/src/server/pi-agent-server.ts b/packages/agent/src/server/pi-agent-server.ts index f30e788fba..cebc51152d 100644 --- a/packages/agent/src/server/pi-agent-server.ts +++ b/packages/agent/src/server/pi-agent-server.ts @@ -622,7 +622,22 @@ export class PiAgentServer { "invalid_token", ); } - return validateJwt(authHeader.slice(7), this.config.jwtPublicKey); + const payload = validateJwt(authHeader.slice(7), this.config.jwtPublicKey); + this.assertConfiguredRun(payload); + return payload; + } + + private assertConfiguredRun(payload: JwtPayload): void { + if ( + payload.task_id !== this.config.taskId || + payload.run_id !== this.config.runId || + payload.team_id !== this.config.projectId + ) { + throw new JwtValidationError( + "Token does not match the configured task run", + "invalid_token", + ); + } } private async waitForRepoReady(): Promise { From 65e04c1970e33d958153f3acd7bc9b7a5d76035f Mon Sep 17 00:00:00 2001 From: JonathanLab Date: Thu, 23 Jul 2026 14:40:15 +0200 Subject: [PATCH 04/45] feat: show cloud Pi provisioning progress --- packages/core/src/cloud-task/schemas.ts | 8 +++ .../pi-runtime/cloudPiSessionClient.test.ts | 44 +++++++++++++++++ .../src/pi-runtime/cloudPiSessionClient.ts | 44 +++++++++++++++-- packages/shared/src/agent-conversation.ts | 11 +++++ .../features/pi-sessions/PiSessionView.tsx | 17 ++++++- .../buildAgentConversationItems.test.ts | 49 +++++++++++++++++++ .../components/buildConversationItems.ts | 5 ++ .../session-update/ProgressGroupView.tsx | 12 +++-- 8 files changed, 180 insertions(+), 10 deletions(-) diff --git a/packages/core/src/cloud-task/schemas.ts b/packages/core/src/cloud-task/schemas.ts index 7436d0a891..8ac35b0c11 100644 --- a/packages/core/src/cloud-task/schemas.ts +++ b/packages/core/src/cloud-task/schemas.ts @@ -22,6 +22,14 @@ export function isTerminalStatus( // --- Events --- +export const progressNotificationParams = z.object({ + step: z.string().min(1), + status: z.enum(["in_progress", "completed", "failed"]), + label: z.string().min(1), + group: z.string().min(1), + detail: z.string().optional(), +}); + export const CloudTaskEvent = { Update: "cloud-task-update", } as const; diff --git a/packages/core/src/pi-runtime/cloudPiSessionClient.test.ts b/packages/core/src/pi-runtime/cloudPiSessionClient.test.ts index 9a822a033a..a1a8fbbf27 100644 --- a/packages/core/src/pi-runtime/cloudPiSessionClient.test.ts +++ b/packages/core/src/pi-runtime/cloudPiSessionClient.test.ts @@ -166,6 +166,50 @@ describe("CloudPiSessionClient", () => { expect(onError).toHaveBeenCalledWith(error); }); + it("streams provisioning progress before the Pi runtime is ready", () => { + const cloud = createCloudTaskClient(); + const session = new CloudPiSessionClient( + cloud.client, + context("in_progress"), + ); + const events: AgentConversationEvent[] = []; + session.onConversationEvent((event) => events.push(event), vi.fn()); + + cloud.sendUpdate({ + taskId: "task-1", + runId: "run-1", + kind: "logs", + newEntries: [ + { + type: "notification", + timestamp: "2026-07-23T12:00:00.000Z", + notification: { + method: "_posthog/progress", + params: { + step: "sandbox", + status: "in_progress", + label: "Setting up sandbox", + group: "setup:run-1", + }, + }, + }, + ], + totalEntryCount: 1, + }); + + expect(events).toEqual([ + { + type: "progress", + timestamp: Date.parse("2026-07-23T12:00:00.000Z"), + step: "sandbox", + status: "in_progress", + label: "Setting up sandbox", + group: "setup:run-1", + }, + ]); + expect(cloud.client.sendCommand).not.toHaveBeenCalled(); + }); + it("loads terminal history from the cloud snapshot without sandbox RPC", async () => { const cloud = createCloudTaskClient(); const session = new CloudPiSessionClient( diff --git a/packages/core/src/pi-runtime/cloudPiSessionClient.ts b/packages/core/src/pi-runtime/cloudPiSessionClient.ts index 0c2e6594c4..5ec78b7bd7 100644 --- a/packages/core/src/pi-runtime/cloudPiSessionClient.ts +++ b/packages/core/src/pi-runtime/cloudPiSessionClient.ts @@ -12,7 +12,10 @@ import type { } from "@posthog/shared"; import type { CloudTaskUpdatePayload } from "@posthog/shared/domain-types"; import type { CloudTaskClient } from "../cloud-task/cloudTaskClient"; -import { isTerminalStatus } from "../cloud-task/schemas"; +import { + isTerminalStatus, + progressNotificationParams, +} from "../cloud-task/schemas"; import type { PiSession } from "./piSessionController"; const readinessCommands = new Set([ @@ -155,7 +158,7 @@ export class CloudPiSessionClient implements PiSession { } if (update.kind === "snapshot") { - const events = this.getPiEvents(update.newEntries); + const events = this.getConversationEvents(update.newEntries); let unchangedEventCount = 0; while ( unchangedEventCount < events.length && @@ -175,7 +178,7 @@ export class CloudPiSessionClient implements PiSession { onEvent(event); } } else if (update.kind === "logs") { - const events = this.getPiEvents(update.newEntries); + const events = this.getConversationEvents(update.newEntries); this.snapshotEvents = [...this.snapshotEvents, ...events]; for (const event of events) { onEvent(event); @@ -205,16 +208,49 @@ export class CloudPiSessionClient implements PiSession { return JSON.stringify(left) === JSON.stringify(right); } - private getPiEvents(entries: StoredLogEntry[]): AgentConversationEvent[] { + private getConversationEvents( + entries: StoredLogEntry[], + ): AgentConversationEvent[] { const events: AgentConversationEvent[] = []; for (const entry of entries) { if (entry.type === "pi_event" && entry.event) { events.push(entry.event); + continue; + } + + const progress = this.getProgressEvent(entry); + if (progress) { + events.push(progress); } } return events; } + private getProgressEvent( + entry: StoredLogEntry, + ): AgentConversationEvent | null { + if ( + entry.notification?.method !== "_posthog/progress" && + entry.notification?.method !== "__posthog/progress" + ) { + return null; + } + + const params = progressNotificationParams.safeParse( + entry.notification.params, + ); + const timestamp = Date.parse(entry.timestamp ?? ""); + if (!params.success || Number.isNaN(timestamp)) { + return null; + } + + return { + type: "progress", + timestamp, + ...params.data, + }; + } + private async request(command: RpcCommand): Promise { if (isTerminalStatus(this.runStatus)) { return this.terminalResponseWhenReady(command); diff --git a/packages/shared/src/agent-conversation.ts b/packages/shared/src/agent-conversation.ts index 9f03b345f6..8644f1eb8b 100644 --- a/packages/shared/src/agent-conversation.ts +++ b/packages/shared/src/agent-conversation.ts @@ -17,6 +17,8 @@ export type AgentToolCallStatus = | "completed" | "failed"; +export type AgentProgressStatus = "in_progress" | "completed" | "failed"; + export interface AgentTextContent { type: "text"; text: string; @@ -134,6 +136,15 @@ export type AgentConversationEvent = timestamp: number; toolCall: Pick & Partial>; } + | { + type: "progress"; + timestamp: number; + step: string; + status: AgentProgressStatus; + label: string; + group: string; + detail?: string; + } | { type: "runtime_status"; timestamp: number; diff --git a/packages/ui/src/features/pi-sessions/PiSessionView.tsx b/packages/ui/src/features/pi-sessions/PiSessionView.tsx index d094094042..7e7d944c2d 100644 --- a/packages/ui/src/features/pi-sessions/PiSessionView.tsx +++ b/packages/ui/src/features/pi-sessions/PiSessionView.tsx @@ -188,10 +188,25 @@ export function PiSessionView({ taskId, taskRunId }: PiSessionViewProps) { ); } - if (!session || !status) { + if (!session || (!status && session.events.length === 0)) { return ; } + if (!status) { + return ( + + + + + + ); + } + const pending = isStreaming || isBashRunning; const supportsThinking = session.thinkingLevels.some( (level) => level !== "off", diff --git a/packages/ui/src/features/sessions/components/buildAgentConversationItems.test.ts b/packages/ui/src/features/sessions/components/buildAgentConversationItems.test.ts index a767a8aae5..b4207e7613 100644 --- a/packages/ui/src/features/sessions/components/buildAgentConversationItems.test.ts +++ b/packages/ui/src/features/sessions/components/buildAgentConversationItems.test.ts @@ -140,6 +140,55 @@ describe("buildAgentConversationItems", () => { ); }); + it("groups runtime-neutral provisioning progress", () => { + const result = buildAgentConversationItems( + [ + { + type: "progress", + timestamp: 1, + step: "sandbox", + status: "completed", + label: "Set up sandbox", + group: "setup:run-1", + }, + { + type: "progress", + timestamp: 2, + step: "clone", + status: "in_progress", + label: "Cloning repository", + group: "setup:run-1", + detail: "posthog/code", + }, + ], + true, + ); + + expect(result.items).toContainEqual( + expect.objectContaining({ + type: "session_update", + update: { + sessionUpdate: "progress_group", + isActive: true, + steps: [ + { + key: "sandbox", + status: "completed", + label: "Set up sandbox", + detail: undefined, + }, + { + key: "clone", + status: "in_progress", + label: "Cloning repository", + detail: "posthog/code", + }, + ], + }, + }), + ); + }); + it("builds and completes a generic compaction status", () => { const result = buildAgentConversationItems( [ diff --git a/packages/ui/src/features/sessions/components/buildConversationItems.ts b/packages/ui/src/features/sessions/components/buildConversationItems.ts index 83d33c41c1..a6609ede9d 100644 --- a/packages/ui/src/features/sessions/components/buildConversationItems.ts +++ b/packages/ui/src/features/sessions/components/buildConversationItems.ts @@ -372,6 +372,11 @@ export function processAgentConversationEvent( return; } + if (event.type === "progress") { + handleProgress(b, event, event.timestamp); + return; + } + if (event.type === "runtime_status") { handleRuntimeStatus(b, event, event.timestamp); return; diff --git a/packages/ui/src/features/sessions/components/session-update/ProgressGroupView.tsx b/packages/ui/src/features/sessions/components/session-update/ProgressGroupView.tsx index 9ca9b06d8f..e5b3efcd76 100644 --- a/packages/ui/src/features/sessions/components/session-update/ProgressGroupView.tsx +++ b/packages/ui/src/features/sessions/components/session-update/ProgressGroupView.tsx @@ -48,10 +48,12 @@ export function ProgressGroupView({ // trigger is disabled and forced open, so the user sees progress stream in without a flicker between // consecutive step transitions. Once the turn completes, the header auto-collapses (default: open) // and becomes interactive. Single-step groups have no header — the one step row IS the whole view. + const isSettled = turnComplete && !isActive; + if (!chatChrome) { const isOpen = !hasHeader ? true - : !turnComplete + : !isSettled ? true : (userToggledOpen ?? true); const summaryLabel = resolveHeaderLabel(steps) ?? ""; @@ -61,11 +63,11 @@ export function ProgressGroupView({ { - if (hasHeader && turnComplete) setUserToggledOpen(next); + if (hasHeader && isSettled) setUserToggledOpen(next); }} > {hasHeader && ( - + + )} + + ); } - if (!session || (!status && session.events.length === 0)) { + if (!status && !hasTranscript) { return ; } @@ -241,6 +296,17 @@ export function PiSessionView({ taskId, taskRunId }: PiSessionViewProps) { return ( + {isConnecting && hasTranscript && ( + + )} + {session.errorMessage && hasTranscript && ( + + )} { @@ -70,10 +76,10 @@ export function CloudInitializingView({ - {heading} + {visibleHeading} - {subtitle} + {visibleSubtitle} diff --git a/packages/ui/src/features/sessions/components/CloudSessionLifecycle.tsx b/packages/ui/src/features/sessions/components/CloudSessionLifecycle.tsx new file mode 100644 index 0000000000..6b279c88c8 --- /dev/null +++ b/packages/ui/src/features/sessions/components/CloudSessionLifecycle.tsx @@ -0,0 +1,85 @@ +import { Spinner, Warning } from "@phosphor-icons/react"; +import { Button, Flex, Text } from "@radix-ui/themes"; + +interface CloudConnectionBannerProps { + message: string; +} + +export function CloudConnectionBanner({ message }: CloudConnectionBannerProps) { + return ( + + + + {message} + + + ); +} + +interface CloudStreamDisconnectedBannerProps { + errorTitle?: string; + errorMessage?: string; + onRetry?: () => void; + onRestart?: () => void; +} + +export function CloudStreamDisconnectedBanner({ + errorTitle, + errorMessage, + onRetry, + onRestart, +}: CloudStreamDisconnectedBannerProps) { + return ( + + + + {errorTitle && ( + + {errorTitle} + + )} + {errorMessage && ( + + {errorMessage} + + )} + + + {onRetry && ( + + )} + {onRestart && ( + + )} + + + ); +} + +export function ConnectingToAgent() { + return ( + <> + + + Connecting to agent... + + + ); +} diff --git a/packages/ui/src/features/sessions/components/SessionView.tsx b/packages/ui/src/features/sessions/components/SessionView.tsx index 163f19fbc4..b538f82e8e 100644 --- a/packages/ui/src/features/sessions/components/SessionView.tsx +++ b/packages/ui/src/features/sessions/components/SessionView.tsx @@ -17,6 +17,10 @@ import { useAutoFocusOnTyping } from "@posthog/ui/features/message-editor/useAut import { resolveAndAttachDroppedFiles } from "@posthog/ui/features/message-editor/utils/persistFile"; import { PermissionSelector } from "@posthog/ui/features/permissions/PermissionSelector"; import { CloudInitializingView } from "@posthog/ui/features/sessions/components/CloudInitializingView"; +import { + CloudStreamDisconnectedBanner, + ConnectingToAgent, +} from "@posthog/ui/features/sessions/components/CloudSessionLifecycle"; import type { PromptRecallHandler } from "@posthog/ui/features/sessions/components/chat-thread/composerPromptRecall"; import { copyFromContextMenu, @@ -95,17 +99,6 @@ interface SessionViewProps { const DEFAULT_ERROR_MESSAGE = "Failed to resume this session. The working directory may have been deleted. Please start a new session."; -function ConnectingToAgent() { - return ( - <> - - - Connecting to agent... - - - ); -} - /** Centers composer-slot content at the chat width (or compact padding). */ function ComposerWidth({ compact, @@ -149,48 +142,6 @@ function ComposerSlot({ ); } -interface CloudStreamDisconnectedBannerProps { - errorTitle?: string; - errorMessage?: string; - onRetry?: () => void; -} - -function CloudStreamDisconnectedBanner({ - errorTitle, - errorMessage, - onRetry, -}: CloudStreamDisconnectedBannerProps) { - return ( - - - - {errorTitle && ( - - {errorTitle} - - )} - {errorMessage && ( - - {errorMessage} - - )} - - {onRetry && ( - - )} - - ); -} - export function SessionView({ events, taskId, diff --git a/packages/ui/src/shell/GlobalEventHandlers.tsx b/packages/ui/src/shell/GlobalEventHandlers.tsx index 4ba6fa0000..ca74f0f1e8 100644 --- a/packages/ui/src/shell/GlobalEventHandlers.tsx +++ b/packages/ui/src/shell/GlobalEventHandlers.tsx @@ -1,8 +1,10 @@ +import { PI_SESSION_CONTROLLER } from "@posthog/core/pi-runtime/identifiers"; +import type { PiSessionController } from "@posthog/core/pi-runtime/piSessionController"; import { SESSION_SERVICE, type SessionService, } from "@posthog/core/sessions/sessionService"; -import { useService } from "@posthog/di/react"; +import { useService, useServiceOptional } from "@posthog/di/react"; import { useHostTRPC } from "@posthog/host-router/react"; import { PROJECT_BLUEBIRD_FLAG } from "@posthog/shared"; import type { Task } from "@posthog/shared/domain-types"; @@ -46,6 +48,9 @@ export function GlobalEventHandlers({ }: GlobalEventHandlersProps) { const trpcReact = useHostTRPC(); const sessionService = useService(SESSION_SERVICE); + const piSessionController = useServiceOptional( + PI_SESSION_CONTROLLER, + ); const commandMenuOpen = useCommandMenuStore((s) => s.isOpen); const openSettingsDialog = openSettings; const view = useAppView(); @@ -282,10 +287,11 @@ export function GlobalEventHandlers({ const handleFocus = () => { loadFolders(); sessionService.retryUnhealthyCloudSessions(); + piSessionController?.retryUnhealthyCloudSessions(); }; window.addEventListener("focus", handleFocus); return () => window.removeEventListener("focus", handleFocus); - }, [loadFolders, sessionService]); + }, [loadFolders, piSessionController, sessionService]); // Freeze perpetual CSS animations while the window is backgrounded (see the // `.ph-window-blurred` rule in globals.css). Driven by the shared focus store From d59247c73262d62bda564530cb6e9e68cc3276c3 Mon Sep 17 00:00:00 2001 From: JonathanLab Date: Mon, 27 Jul 2026 16:01:18 +0200 Subject: [PATCH 14/45] feat(pi): smooth cloud token streaming --- .../agent/src/server/event-stream-sender.ts | 7 +++++++ .../features/pi-sessions/PiSessionView.tsx | 20 +++++++++++++------ 2 files changed, 21 insertions(+), 6 deletions(-) diff --git a/packages/agent/src/server/event-stream-sender.ts b/packages/agent/src/server/event-stream-sender.ts index 05b17fbedd..a4f30b4cf0 100644 --- a/packages/agent/src/server/event-stream-sender.ts +++ b/packages/agent/src/server/event-stream-sender.ts @@ -103,6 +103,7 @@ export class TaskRunEventStreamSender { config.logger.info("Event ingest target resolved", { ingestUrl: this.ingestUrl, routedToProxy: usingProxy, + persistentUpload: !usingProxy || config.keepProxyStreamOpen === true, }); this.maxBufferedEvents = config.maxBufferedEvents ?? DEFAULT_MAX_BUFFERED_EVENTS; @@ -478,6 +479,12 @@ export class TaskRunEventStreamSender { await this.applyIngestResponse(response, "Event ingest stream"); this.sequenceSynced = true; + this.config.logger.debug("Task run event ingest stream delivered", { + durationMs: Date.now() - stream.startedAtMs, + sentBytes: stream.sentBytes, + sentEvents: stream.sentEvents, + sentThroughSeq: stream.sentThroughSeq, + }); } private async abortActiveStream(): Promise { diff --git a/packages/ui/src/features/pi-sessions/PiSessionView.tsx b/packages/ui/src/features/pi-sessions/PiSessionView.tsx index 70ae95b64b..abede68273 100644 --- a/packages/ui/src/features/pi-sessions/PiSessionView.tsx +++ b/packages/ui/src/features/pi-sessions/PiSessionView.tsx @@ -24,6 +24,7 @@ import { CloudStreamDisconnectedBanner, } from "@posthog/ui/features/sessions/components/CloudSessionLifecycle"; import { ChatThread } from "@posthog/ui/features/sessions/components/chat-thread/ChatThread"; +import type { PromptRecallHandler } from "@posthog/ui/features/sessions/components/chat-thread/composerPromptRecall"; import { CHAT_CONTENT_MAX_WIDTH } from "@posthog/ui/features/sessions/constants"; import { useMessagingMode } from "@posthog/ui/features/sessions/hooks/useMessagingMode"; import { useMessagingModeStore } from "@posthog/ui/features/sessions/messagingModeStore"; @@ -32,7 +33,7 @@ import { useConnectivity } from "@posthog/ui/hooks/useConnectivity"; import { toast } from "@posthog/ui/primitives/toast"; import { TaskDetailSkeleton } from "@posthog/ui/router/routeSkeletons"; import { Box, Flex } from "@radix-ui/themes"; -import { type ReactElement, useCallback, useEffect } from "react"; +import { type ReactElement, useCallback, useEffect, useRef } from "react"; import { useStore } from "zustand"; import { PiMessagingModeSelector, @@ -59,6 +60,11 @@ export function PiSessionView({ taskId, taskRunId }: PiSessionViewProps) { const messagingMode = useMessagingMode(taskId); const setMessagingMode = useMessagingModeStore((state) => state.setMode); const { isOnline } = useConnectivity(); + const promptRecallRef = useRef(null); + const handlePromptRecall = useCallback( + (direction) => promptRecallRef.current?.(direction) ?? null, + [], + ); useEffect(() => { void piSessionController.ensureConnected(taskId, taskRunId).catch(() => {}); @@ -247,7 +253,7 @@ export function PiSessionView({ taskId, taskRunId }: PiSessionViewProps) { return ; } - const pending = status ? isStreaming || isBashRunning : false; + const controlsPending = status ? isStreaming || isBashRunning : false; let modelSelector: ReactElement = ; let reasoningSelector: ReactElement | null = ( @@ -259,7 +265,7 @@ export function PiSessionView({ taskId, taskRunId }: PiSessionViewProps) { ); @@ -273,7 +279,7 @@ export function PiSessionView({ taskId, taskRunId }: PiSessionViewProps) { ) : null; @@ -310,9 +316,10 @@ export function PiSessionView({ taskId, taskRunId }: PiSessionViewProps) { Date: Mon, 27 Jul 2026 17:27:50 +0200 Subject: [PATCH 15/45] feat(pi): add durable single-message queue --- .../conversation/translatePiConversation.ts | 11 + .../agent/src/pi/queue-persistence.test.ts | 48 ++++ packages/agent/src/pi/queue-persistence.ts | 30 +++ packages/agent/src/pi/remote-rpc-client.ts | 14 -- packages/agent/src/pi/rpc-client.test.ts | 43 +++- packages/agent/src/pi/rpc-client.ts | 111 ++++++++- packages/agent/src/pi/rpc-host.ts | 63 ++++++ packages/agent/src/pi/rpc-transport.test.ts | 8 +- packages/agent/src/pi/runtime.test.ts | 30 ++- packages/agent/src/pi/types.ts | 13 +- .../agent/src/server/pi-agent-server.test.ts | 28 +++ packages/agent/src/server/pi-agent-server.ts | 11 + packages/core/src/cloud-task/schemas.ts | 2 + .../pi-runtime/cloudPiSessionClient.test.ts | 33 ++- .../src/pi-runtime/cloudPiSessionClient.ts | 44 +++- .../pi-runtime/piSessionController.test.ts | 166 +++++++++++++- .../src/pi-runtime/piSessionController.ts | 214 +++++++++++++++--- .../src/pi-runtime/piSessionProvider.test.ts | 4 +- .../core/src/pi-runtime/piSessionStore.ts | 3 + .../host-router/src/pi-session-factory.ts | 8 + .../src/routers/pi-session.router.ts | 15 ++ packages/shared/src/agent-conversation.ts | 6 + .../pi-sessions/PiQueuedMessagesDock.test.tsx | 67 ++++++ .../pi-sessions/PiQueuedMessagesDock.tsx | 34 +++ .../pi-sessions/PiSessionControls.tsx | 16 -- .../features/pi-sessions/PiSessionView.tsx | 67 ++++-- .../components/buildConversationItems.ts | 4 + .../session-update/QueuedMessageView.tsx | 20 +- .../features/settings/settingsStore.test.ts | 22 ++ .../ui/src/features/settings/settingsStore.ts | 13 +- .../task-detail/components/TaskInput.tsx | 17 +- .../src/services/pi-session/pi-session.ts | 9 + .../src/services/pi-session/schemas.ts | 5 + 33 files changed, 1065 insertions(+), 114 deletions(-) create mode 100644 packages/agent/src/pi/queue-persistence.test.ts create mode 100644 packages/agent/src/pi/queue-persistence.ts create mode 100644 packages/ui/src/features/pi-sessions/PiQueuedMessagesDock.test.tsx create mode 100644 packages/ui/src/features/pi-sessions/PiQueuedMessagesDock.tsx diff --git a/packages/agent/src/pi/conversation/translatePiConversation.ts b/packages/agent/src/pi/conversation/translatePiConversation.ts index fb78985c80..26e353014c 100644 --- a/packages/agent/src/pi/conversation/translatePiConversation.ts +++ b/packages/agent/src/pi/conversation/translatePiConversation.ts @@ -317,6 +317,17 @@ export function createPiConversationTranslator(): PiConversationTranslator { ]; } + if (event.type === "queue_update") { + return [ + { + type: "queue_update", + timestamp: Date.now(), + steering: [...event.steering], + followUp: [...event.followUp], + }, + ]; + } + if (event.type === "message_end") { latestRuntimeTimestamp = Math.max( latestRuntimeTimestamp, diff --git a/packages/agent/src/pi/queue-persistence.test.ts b/packages/agent/src/pi/queue-persistence.test.ts new file mode 100644 index 0000000000..656db83158 --- /dev/null +++ b/packages/agent/src/pi/queue-persistence.test.ts @@ -0,0 +1,48 @@ +import type { SessionEntry } from "@earendil-works/pi-coding-agent"; +import { describe, expect, it } from "vitest"; +import { + POSTHOG_PI_QUEUE_ENTRY_TYPE, + readPersistedPiQueue, +} from "./queue-persistence"; + +describe("readPersistedPiQueue", () => { + it("uses the latest valid persisted queue snapshot", () => { + const entries = [ + { + type: "custom", + customType: POSTHOG_PI_QUEUE_ENTRY_TYPE, + data: { steering: ["old"], followUp: [] }, + }, + { + type: "custom", + customType: POSTHOG_PI_QUEUE_ENTRY_TYPE, + data: { steering: ["new"], followUp: ["later"] }, + }, + ] as SessionEntry[]; + + expect(readPersistedPiQueue(entries)).toEqual({ + steering: ["new"], + followUp: ["later"], + }); + }); + + it("ignores malformed and unrelated custom entries", () => { + const entries = [ + { + type: "custom", + customType: POSTHOG_PI_QUEUE_ENTRY_TYPE, + data: { steering: [1], followUp: [] }, + }, + { + type: "custom", + customType: "other", + data: { steering: ["other"], followUp: [] }, + }, + ] as SessionEntry[]; + + expect(readPersistedPiQueue(entries)).toEqual({ + steering: [], + followUp: [], + }); + }); +}); diff --git a/packages/agent/src/pi/queue-persistence.ts b/packages/agent/src/pi/queue-persistence.ts new file mode 100644 index 0000000000..abf2854181 --- /dev/null +++ b/packages/agent/src/pi/queue-persistence.ts @@ -0,0 +1,30 @@ +import type { SessionEntry } from "@earendil-works/pi-coding-agent"; +import type { PiQueueSnapshot } from "./types"; + +export const POSTHOG_PI_QUEUE_ENTRY_TYPE = "posthog.pi.queue"; + +export function readPersistedPiQueue(entries: SessionEntry[]): PiQueueSnapshot { + for (let index = entries.length - 1; index >= 0; index -= 1) { + const entry = entries[index]; + if ( + entry?.type !== "custom" || + entry.customType !== POSTHOG_PI_QUEUE_ENTRY_TYPE + ) { + continue; + } + const data = entry.data as Partial | undefined; + if ( + Array.isArray(data?.steering) && + data.steering.every((message) => typeof message === "string") && + Array.isArray(data.followUp) && + data.followUp.every((message) => typeof message === "string") + ) { + return { + steering: [...data.steering], + followUp: [...data.followUp], + }; + } + } + + return { steering: [], followUp: [] }; +} diff --git a/packages/agent/src/pi/remote-rpc-client.ts b/packages/agent/src/pi/remote-rpc-client.ts index 66820ee642..1a284759ff 100644 --- a/packages/agent/src/pi/remote-rpc-client.ts +++ b/packages/agent/src/pi/remote-rpc-client.ts @@ -18,8 +18,6 @@ export type PiRemoteRpcClient = Pick< | "getAvailableModels" | "getAvailableThinkingLevels" | "setThinkingLevel" - | "setSteeringMode" - | "setFollowUpMode" | "compact" | "bash" | "abortBash" @@ -116,18 +114,6 @@ export class RemotePiRpcClient implements PiRemoteRpcClient { await this.request({ type: "set_thinking_level", level }); } - async setSteeringMode( - mode: Parameters[0], - ): Promise { - await this.request({ type: "set_steering_mode", mode }); - } - - async setFollowUpMode( - mode: Parameters[0], - ): Promise { - await this.request({ type: "set_follow_up_mode", mode }); - } - async compact( customInstructions?: string, ): ReturnType { diff --git a/packages/agent/src/pi/rpc-client.test.ts b/packages/agent/src/pi/rpc-client.test.ts index 4c8239d698..3da0303468 100644 --- a/packages/agent/src/pi/rpc-client.test.ts +++ b/packages/agent/src/pi/rpc-client.test.ts @@ -1,6 +1,9 @@ +import { mkdtemp, rm, writeFile } from "node:fs/promises"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; import { RpcClient } from "@earendil-works/pi-coding-agent"; import { describe, expect, it } from "vitest"; -import { createPiRpcClient } from "./rpc-client"; +import { createPiRpcClient, getPiRpcClientProcess } from "./rpc-client"; describe("createPiRpcClient", () => { it("does not put provider credentials in the child environment", () => { @@ -27,4 +30,42 @@ describe("createPiRpcClient", () => { .options.env, ).toBeUndefined(); }); + + it("uses the private host channel without changing Pi RPC", async () => { + const directory = await mkdtemp(join(tmpdir(), "pi-host-channel-")); + const hostPath = join(directory, "host.mjs"); + await writeFile( + hostPath, + ` +process.stdin.resume(); +process.on("message", (request) => { + const data = request.method === "clear_queue" + ? { steering: ["cleared"], followUp: [] } + : { steering: ["queued"], followUp: ["later"] }; + process.send({ type: "posthog_pi_host_response", id: request.id, data }); +}); +`, + ); + const client = createPiRpcClient({ + cliPath: hostPath, + cwd: directory, + providerOptions: { apiKey: "proxy-key" }, + }); + + try { + await client.start(); + + await expect(client.getQueue()).resolves.toEqual({ + steering: ["queued"], + followUp: ["later"], + }); + await expect(client.clearQueue()).resolves.toEqual({ + steering: ["cleared"], + followUp: [], + }); + } finally { + getPiRpcClientProcess(client)?.kill(); + await rm(directory, { recursive: true }); + } + }); }); diff --git a/packages/agent/src/pi/rpc-client.ts b/packages/agent/src/pi/rpc-client.ts index e565e360e0..3f9654338a 100644 --- a/packages/agent/src/pi/rpc-client.ts +++ b/packages/agent/src/pi/rpc-client.ts @@ -1,4 +1,5 @@ import { type ChildProcess, spawn } from "node:child_process"; +import { randomUUID } from "node:crypto"; import type { Writable } from "node:stream"; import { StringDecoder } from "node:string_decoder"; import { fileURLToPath } from "node:url"; @@ -7,8 +8,12 @@ import { type RpcClientOptions, } from "@earendil-works/pi-coding-agent"; import { safePiEnvironment } from "./rpc-environment"; +import type { PiQueueSnapshot } from "./types"; -export type PiRpcClient = RpcClient; +export type PiRpcClient = RpcClient & { + getQueue(): Promise; + clearQueue(): Promise; +}; export interface PiRpcProviderOptions { region?: "us" | "eu" | "dev"; @@ -33,6 +38,19 @@ interface RpcClientInternals { rejectPendingRequests(error: Error): void; } +interface PiHostRequest { + type: "posthog_pi_host_request"; + id: string; + method: "get_queue" | "clear_queue"; +} + +interface PiHostResponse { + type: "posthog_pi_host_response"; + id: string; + data?: PiQueueSnapshot; + error?: string; +} + function attachJsonlReader( stream: NodeJS.ReadableStream, onLine: (line: string) => void, @@ -54,6 +72,15 @@ function attachJsonlReader( } class SecurePiRpcClient extends RpcClient { + private readonly hostRequests = new Map< + string, + { + resolve: (snapshot: PiQueueSnapshot) => void; + reject: (error: Error) => void; + timeout: ReturnType; + } + >(); + constructor( private readonly secureOptions: RpcClientOptions, private readonly providerOptions: PiRpcProviderOptions, @@ -85,7 +112,7 @@ class SecurePiRpcClient extends RpcClient { { cwd: this.secureOptions.cwd, env: safePiEnvironment(process.env), - stdio: ["pipe", "pipe", "pipe", "pipe"], + stdio: ["pipe", "pipe", "pipe", "pipe", "ipc"], }, ); internals.process = child; @@ -101,7 +128,9 @@ class SecurePiRpcClient extends RpcClient { const error = internals.createProcessExitError(code, signal); internals.exitError = error; internals.rejectPendingRequests(error); + this.rejectHostRequests(error); }); + child.on("message", (message: unknown) => this.handleHostResponse(message)); child.once("error", (error) => { if (internals.process !== child) { return; @@ -140,6 +169,84 @@ class SecurePiRpcClient extends RpcClient { ); } } + + getQueue(): Promise { + return this.sendHostRequest("get_queue"); + } + + clearQueue(): Promise { + return this.sendHostRequest("clear_queue"); + } + + private sendHostRequest( + method: PiHostRequest["method"], + ): Promise { + const process = (this as unknown as RpcClientInternals).process; + if (!process?.connected) { + return Promise.reject(new Error("Pi RPC host is not connected")); + } + + const id = randomUUID(); + const request: PiHostRequest = { + type: "posthog_pi_host_request", + id, + method, + }; + + return new Promise((resolve, reject) => { + const timeout = setTimeout(() => { + this.hostRequests.delete(id); + reject(new Error(`Pi RPC host request timed out: ${method}`)); + }, 10_000); + this.hostRequests.set(id, { resolve, reject, timeout }); + process.send?.(request, (error) => { + if (!error) { + return; + } + const pending = this.hostRequests.get(id); + if (pending) { + clearTimeout(pending.timeout); + this.hostRequests.delete(id); + } + reject(error); + }); + }); + } + + private handleHostResponse(message: unknown): void { + const response = message as Partial; + if ( + response.type !== "posthog_pi_host_response" || + typeof response.id !== "string" + ) { + return; + } + + const request = this.hostRequests.get(response.id); + if (!request) { + return; + } + this.hostRequests.delete(response.id); + clearTimeout(request.timeout); + + if (typeof response.error === "string") { + request.reject(new Error(response.error)); + return; + } + if (!response.data) { + request.reject(new Error("Pi RPC host returned an empty queue response")); + return; + } + request.resolve(response.data); + } + + private rejectHostRequests(error: Error): void { + for (const request of this.hostRequests.values()) { + clearTimeout(request.timeout); + request.reject(error); + } + this.hostRequests.clear(); + } } export function getPiRpcClientProcess( diff --git a/packages/agent/src/pi/rpc-host.ts b/packages/agent/src/pi/rpc-host.ts index eb13b730bb..9516177b4e 100644 --- a/packages/agent/src/pi/rpc-host.ts +++ b/packages/agent/src/pi/rpc-host.ts @@ -2,12 +2,22 @@ import { readFileSync } from "node:fs"; import { SessionManager } from "@earendil-works/pi-coding-agent"; import { createHarnessRuntime, runRpcMode } from "@posthog/harness"; import type { PosthogProviderOptions } from "@posthog/harness/extensions/posthog-provider/provider"; +import { + POSTHOG_PI_QUEUE_ENTRY_TYPE, + readPersistedPiQueue, +} from "./queue-persistence"; import { sanitizePiHostEnvironment } from "./rpc-environment"; interface PiRpcBootstrap { providerOptions?: PosthogProviderOptions; } +interface PiHostRequest { + type: "posthog_pi_host_request"; + id: string; + method: "get_queue" | "clear_queue"; +} + function argumentValue(name: string): string | undefined { const index = process.argv.indexOf(name); return index === -1 ? undefined : process.argv[index + 1]; @@ -31,6 +41,26 @@ const runtime = await createHarnessRuntime({ ...providerOptions, }); +const persistedQueue = readPersistedPiQueue(sessionManager.getEntries()); +for (const message of persistedQueue.steering) { + await runtime.session.steer(message); +} +for (const message of persistedQueue.followUp) { + await runtime.session.followUp(message); +} +runtime.session.subscribe((event) => { + if (event.type !== "queue_update") { + return; + } + runtime.session.sessionManager.appendCustomEntry( + POSTHOG_PI_QUEUE_ENTRY_TYPE, + { + steering: [...event.steering], + followUp: [...event.followUp], + }, + ); +}); + const requestedModel = argumentValue("--model")?.replace(/^posthog\//, ""); if (requestedModel) { const model = runtime.services.modelRuntime.getModel( @@ -43,4 +73,37 @@ if (requestedModel) { await runtime.session.setModel(model); } +process.on("message", (message: unknown) => { + const request = message as Partial; + if ( + request.type !== "posthog_pi_host_request" || + typeof request.id !== "string" || + (request.method !== "get_queue" && request.method !== "clear_queue") + ) { + return; + } + + try { + const session = runtime.session; + const data = + request.method === "clear_queue" + ? session.clearQueue() + : { + steering: [...session.getSteeringMessages()], + followUp: [...session.getFollowUpMessages()], + }; + process.send?.({ + type: "posthog_pi_host_response", + id: request.id, + data, + }); + } catch (error) { + process.send?.({ + type: "posthog_pi_host_response", + id: request.id, + error: error instanceof Error ? error.message : String(error), + }); + } +}); + await runRpcMode(runtime); diff --git a/packages/agent/src/pi/rpc-transport.test.ts b/packages/agent/src/pi/rpc-transport.test.ts index ae56c91148..ec28950723 100644 --- a/packages/agent/src/pi/rpc-transport.test.ts +++ b/packages/agent/src/pi/rpc-transport.test.ts @@ -29,21 +29,15 @@ describe("RemotePiRpcClient", () => { }); const client = new RemotePiRpcClient({ request }); - await client.setFollowUpMode("one-at-a-time"); const compaction = await client.compact("retain decisions"); const thinkingLevels = await client.getAvailableThinkingLevels(); expect(request).toHaveBeenNthCalledWith(1, { - id: expect.any(String), - type: "set_follow_up_mode", - mode: "one-at-a-time", - }); - expect(request).toHaveBeenNthCalledWith(2, { id: expect.any(String), type: "compact", customInstructions: "retain decisions", }); - expect(request).toHaveBeenNthCalledWith(3, { + expect(request).toHaveBeenNthCalledWith(2, { id: expect.any(String), type: "get_available_thinking_levels", }); diff --git a/packages/agent/src/pi/runtime.test.ts b/packages/agent/src/pi/runtime.test.ts index 3dafdee015..201d42d9a5 100644 --- a/packages/agent/src/pi/runtime.test.ts +++ b/packages/agent/src/pi/runtime.test.ts @@ -1,9 +1,7 @@ import type { AssistantMessage, UserMessage } from "@earendil-works/pi-ai"; -import type { - AgentSessionEvent, - RpcClient, -} from "@earendil-works/pi-coding-agent"; +import type { AgentSessionEvent } from "@earendil-works/pi-coding-agent"; import { describe, expect, it, vi } from "vitest"; +import type { PiRpcClient } from "./rpc-client"; import { PiRuntime } from "./runtime"; function assistant(text: string): AssistantMessage { @@ -35,7 +33,9 @@ function createClient() { return () => {}; }), send, - } as unknown as RpcClient; + getQueue: vi.fn(async () => ({ steering: [], followUp: [] })), + clearQueue: vi.fn(async () => ({ steering: [], followUp: [] })), + } as unknown as PiRpcClient; return { client, @@ -165,6 +165,26 @@ describe("PiRuntime", () => { ); }); + it("forwards native queue snapshots", () => { + const { client, emit } = createClient(); + const runtime = new PiRuntime(client); + const conversationListener = vi.fn(); + runtime.onConversationEvent(conversationListener); + + emit({ + type: "queue_update", + steering: ["fix this"], + followUp: ["then summarize"], + }); + + expect(conversationListener).toHaveBeenCalledWith({ + type: "queue_update", + timestamp: expect.any(Number), + steering: ["fix this"], + followUp: ["then summarize"], + }); + }); + it("normalizes live Pi events before forwarding them", () => { const { client, emit } = createClient(); const runtime = new PiRuntime(client); diff --git a/packages/agent/src/pi/types.ts b/packages/agent/src/pi/types.ts index d6deb53b33..3fe75e7f17 100644 --- a/packages/agent/src/pi/types.ts +++ b/packages/agent/src/pi/types.ts @@ -1,4 +1,4 @@ -import type { QueueMode, ThinkingLevel } from "@earendil-works/pi-agent-core"; +import type { ThinkingLevel } from "@earendil-works/pi-agent-core"; import type { RpcClient, RpcSessionState, @@ -11,7 +11,6 @@ function exhaustiveValues() { } export type PiThinkingLevel = ThinkingLevel; -export type PiQueueMode = QueueMode; export const PI_THINKING_LEVELS = exhaustiveValues()([ "off", @@ -23,11 +22,6 @@ export const PI_THINKING_LEVELS = exhaustiveValues()([ "max", ]); -export const PI_QUEUE_MODES = exhaustiveValues()([ - "all", - "one-at-a-time", -]); - export type PiNativeModelInfo = Awaited< ReturnType >[number]; @@ -38,4 +32,9 @@ export type PiSessionStatus = Omit & { model?: Pick, "provider" | "id">; }; +export interface PiQueueSnapshot { + steering: string[]; + followUp: string[]; +} + export type PiSessionStats = Awaited>; diff --git a/packages/agent/src/server/pi-agent-server.test.ts b/packages/agent/src/server/pi-agent-server.test.ts index 76149cdedd..22fccc7736 100644 --- a/packages/agent/src/server/pi-agent-server.test.ts +++ b/packages/agent/src/server/pi-agent-server.test.ts @@ -298,6 +298,34 @@ describe("PiAgentServer", () => { }); }); + it.each([ + ["queue_get", "getQueue"], + ["queue_clear", "clearQueue"], + ] as const)( + "forwards %s through the private Pi host API", + async (method, operation) => { + const queue = { + steering: ["fix this"], + followUp: ["then summarize"], + }; + const client = { + getQueue: vi.fn(async () => queue), + clearQueue: vi.fn(async () => queue), + }; + const server = new PiAgentServer(config()) as unknown as { + session: unknown; + executeCommand( + method: string, + params: Record, + ): Promise; + }; + server.session = { runtime: { client } }; + + await expect(server.executeCommand(method, {})).resolves.toEqual(queue); + expect(client[operation]).toHaveBeenCalledOnce(); + }, + ); + it("waits for Pi to create the native session file before syncing", async () => { const directory = await mkdtemp(join(tmpdir(), "pi-session-sync-")); const syncTaskSession = vi.fn(async () => "content-hash"); diff --git a/packages/agent/src/server/pi-agent-server.ts b/packages/agent/src/server/pi-agent-server.ts index 6e99506434..0315adcfa9 100644 --- a/packages/agent/src/server/pi-agent-server.ts +++ b/packages/agent/src/server/pi-agent-server.ts @@ -50,6 +50,8 @@ const userMessageCommandSchema = z const commandSchemas = { user_message: userMessageCommandSchema, cancel: emptySchema, + queue_get: emptySchema, + queue_clear: emptySchema, "pi/rpc": z.object({ command: piRpcCommandSchema }), } as const; @@ -428,6 +430,11 @@ export class PiAgentServer { timestamp: new Date().toISOString(), event: { ...event, sourceId: id }, }); + if (event.type === "queue_update") { + void this.syncTaskSession().catch((error) => + this.logger.error("Failed to persist Pi queue state", error), + ); + } } private async executeCommand( @@ -444,6 +451,10 @@ export class PiAgentServer { return this.deliverUserMessage(runtime, params); case "cancel": return client.abort(); + case "queue_get": + return client.getQueue(); + case "queue_clear": + return client.clearQueue(); case "pi/rpc": return runtime.sendCommand(params.command as RpcCommand); } diff --git a/packages/core/src/cloud-task/schemas.ts b/packages/core/src/cloud-task/schemas.ts index 76383d98fd..ca2680e433 100644 --- a/packages/core/src/cloud-task/schemas.ts +++ b/packages/core/src/cloud-task/schemas.ts @@ -79,6 +79,8 @@ export const sendCommandInput = z.object({ "set_config_option", "mcp_response", "pi/rpc", + "queue_get", + "queue_clear", ]), params: z.record(z.string(), z.unknown()).optional(), }); diff --git a/packages/core/src/pi-runtime/cloudPiSessionClient.test.ts b/packages/core/src/pi-runtime/cloudPiSessionClient.test.ts index 0769929d3b..28b2ff8241 100644 --- a/packages/core/src/pi-runtime/cloudPiSessionClient.test.ts +++ b/packages/core/src/pi-runtime/cloudPiSessionClient.test.ts @@ -297,6 +297,12 @@ describe("CloudPiSessionClient", () => { resolveState = resolve; }); vi.mocked(cloud.client.sendCommand).mockImplementation(async (input) => { + if (input.method === "queue_get") { + return { + success: true, + result: { steering: [], followUp: [] }, + }; + } const command = input.params?.command as { type: string }; if (command.type === "get_state") { return state; @@ -324,7 +330,7 @@ describe("CloudPiSessionClient", () => { totalEntryCount: 1, }); await vi.waitFor(() => { - expect(cloud.client.sendCommand).toHaveBeenCalledTimes(1); + expect(cloud.client.sendCommand).toHaveBeenCalledTimes(2); }); cloud.sendUpdate({ taskId: "task-1", @@ -398,6 +404,31 @@ describe("CloudPiSessionClient", () => { expect(cloud.client.sendCommand).toHaveBeenCalledTimes(1); }); + it("keeps sessions compatible with queue-unaware cloud runtimes", async () => { + const cloud = createCloudTaskClient(); + vi.mocked(cloud.client.sendCommand).mockResolvedValue({ + success: false, + error: "Unknown method: queue_get", + }); + const session = new CloudPiSessionClient( + cloud.client, + context("in_progress"), + ); + session.onConversationEvent(vi.fn(), vi.fn()); + cloud.sendUpdate({ + taskId: "task-1", + runId: "run-1", + kind: "logs", + newEntries: [{ type: "pi_run_started" }], + totalEntryCount: 1, + }); + + await expect(session.getQueue()).resolves.toEqual({ + steering: [], + followUp: [], + }); + }); + it("preserves structured backend failure details", () => { const cloud = createCloudTaskClient(); const session = new CloudPiSessionClient( diff --git a/packages/core/src/pi-runtime/cloudPiSessionClient.ts b/packages/core/src/pi-runtime/cloudPiSessionClient.ts index 4de9aa6b9d..0466ddfe0f 100644 --- a/packages/core/src/pi-runtime/cloudPiSessionClient.ts +++ b/packages/core/src/pi-runtime/cloudPiSessionClient.ts @@ -3,6 +3,7 @@ import { RemotePiRpcClient, } from "@posthog/agent/pi/remote-rpc-client"; import type { RpcCommand } from "@posthog/agent/pi/rpc-transport"; +import type { PiQueueSnapshot } from "@posthog/agent/pi/types"; import type { AgentConversationEvent, PiRuntimeHealth, @@ -45,8 +46,6 @@ function createTerminalPiRpcClient( getAvailableModels: async () => [], getAvailableThinkingLevels: async () => [], setThinkingLevel: rejectCommand, - setSteeringMode: rejectCommand, - setFollowUpMode: rejectCommand, compact: rejectCommand, bash: rejectCommand, abortBash: rejectCommand, @@ -127,6 +126,25 @@ export class CloudPiSessionClient implements PiSession { await this.cloudTaskClient.retry(this.context.taskId, this.context.runId); } + async getQueue(): Promise { + try { + return await this.requestQueue("queue_get"); + } catch (error) { + const message = error instanceof Error ? error.message : String(error); + if ( + message.includes("Unknown method: queue_get") || + message.includes("queue_get is not supported") + ) { + return { steering: [], followUp: [] }; + } + throw error; + } + } + + clearQueue(): Promise { + return this.requestQueue("queue_clear"); + } + async sendUserMessage( type: "prompt" | "steer" | "follow_up", message: string, @@ -352,6 +370,28 @@ export class CloudPiSessionClient implements PiSession { }; } + private async requestQueue( + method: "queue_get" | "queue_clear", + ): Promise { + await this.waitForRuntimeReady(); + if (isTerminalStatus(this.runStatus)) { + return { steering: [], followUp: [] }; + } + const result = await this.cloudTaskClient.sendCommand({ + taskId: this.context.taskId, + runId: this.context.runId, + apiHost: this.context.apiHost, + teamId: this.context.teamId, + id: globalThis.crypto.randomUUID(), + method, + params: {}, + }); + if (!result.success) { + throw new Error(result.error ?? `Pi queue command failed: ${method}`); + } + return result.result as PiQueueSnapshot; + } + private async request(command: RpcCommand): Promise { await this.waitForRuntimeReady(); if (isTerminalStatus(this.runStatus)) { diff --git a/packages/core/src/pi-runtime/piSessionController.test.ts b/packages/core/src/pi-runtime/piSessionController.test.ts index 562f2ced64..395c775198 100644 --- a/packages/core/src/pi-runtime/piSessionController.test.ts +++ b/packages/core/src/pi-runtime/piSessionController.test.ts @@ -42,8 +42,6 @@ function createSession(): PiSession { compact: vi.fn(async () => undefined), setModel: vi.fn(async (provider, id) => ({ provider, id })), setThinkingLevel: vi.fn(async () => {}), - setSteeringMode: vi.fn(async () => {}), - setFollowUpMode: vi.fn(async () => {}), bash: vi.fn(async () => undefined), abort: vi.fn(async () => {}), abortBash: vi.fn(async () => {}), @@ -53,6 +51,8 @@ function createSession(): PiSession { client, health: vi.fn(async () => ({ state: "idle" as const })), getConversation: vi.fn(async () => []), + getQueue: vi.fn(async () => ({ steering: [], followUp: [] })), + clearQueue: vi.fn(async () => ({ steering: [], followUp: [] })), onConversationEvent: vi.fn(() => () => {}), }; } @@ -163,13 +163,63 @@ describe("PiSessionController", () => { ); }); + it("allows only one queued message and keeps it out of the transcript", async () => { + const session = createSession(); + session.sendUserMessage = vi.fn(async () => {}); + vi.mocked(session.getQueue) + .mockResolvedValueOnce({ steering: [], followUp: [] }) + .mockResolvedValue({ steering: [], followUp: ["first"] }); + const controller = createController(session, { + prepareCloudPiMessage: vi.fn(async (_taskId, _runId, content) => ({ + content, + artifactIds: [], + })), + } as unknown as TaskService); + + await controller.connect("task-1", "run-1"); + await controller.submit("task-1", "first", true, "queue"); + + expect(controller.store.getState().sessions["task-1"]).toMatchObject({ + events: [], + queue: { steering: [], followUp: ["first"] }, + }); + await expect( + controller.submit("task-1", "second", true, "queue"), + ).rejects.toThrow("Pi already has a queued message"); + expect(session.sendUserMessage).toHaveBeenCalledOnce(); + }); + + it("clears the optimistic queue when Pi accepted the message as a prompt", async () => { + const session = createSession(); + session.sendUserMessage = vi.fn(async () => {}); + const controller = createController(session, { + prepareCloudPiMessage: vi.fn(async () => ({ + content: "continue", + artifactIds: [], + })), + } as unknown as TaskService); + + await controller.connect("task-1", "run-1"); + await controller.submit("task-1", "continue", true, "queue"); + + expect(controller.store.getState().sessions["task-1"].queue).toEqual({ + steering: [], + followUp: [], + }); + }); + it("marks a submitted turn as streaming while the command starts", async () => { let resolveSend: () => void = () => {}; const sending = new Promise((resolve) => { resolveSend = resolve; }); + let onEvent: (event: AgentConversationEvent) => void = () => {}; const session = createSession(); session.sendUserMessage = vi.fn(() => sending); + vi.mocked(session.onConversationEvent).mockImplementation((handler) => { + onEvent = handler; + return () => {}; + }); const controller = createController(session, { prepareCloudPiMessage: vi.fn(async () => ({ content: "hello", @@ -188,6 +238,70 @@ describe("PiSessionController", () => { resolveSend(); await submission; + expect(controller.store.getState().sessions["task-1"].status).toMatchObject( + { isStreaming: true }, + ); + + onEvent({ type: "turn_completed", timestamp: 2 }); + + expect(controller.store.getState().sessions["task-1"].status).toMatchObject( + { isStreaming: false }, + ); + }); + + it("restores a native queue after retry replaces the runtime", async () => { + let onEvent: (event: AgentConversationEvent) => void = () => {}; + const session = createSession(); + session.retry = vi.fn(async () => {}); + vi.mocked(session.onConversationEvent).mockImplementation((handler) => { + onEvent = handler; + return () => {}; + }); + const controller = createController(session); + + await controller.connect("task-1", "run-1"); + onEvent({ + type: "queue_update", + timestamp: 1, + steering: ["fix this"], + followUp: ["then summarize"], + }); + + await controller.retry("task-1"); + + expect(session.client.prompt).toHaveBeenCalledWith("fix this"); + expect(session.client.followUp).toHaveBeenCalledWith("then summarize"); + }); + + it("does not restore a captured queue after the task disconnects", async () => { + let resolveRetry: () => void = () => {}; + const retrying = new Promise((resolve) => { + resolveRetry = resolve; + }); + let onEvent: (event: AgentConversationEvent) => void = () => {}; + const session = createSession(); + session.retry = vi.fn(() => retrying); + vi.mocked(session.onConversationEvent).mockImplementation((handler) => { + onEvent = handler; + return () => {}; + }); + const controller = createController(session); + + await controller.connect("task-1", "run-1"); + onEvent({ + type: "queue_update", + timestamp: 1, + steering: ["already handled"], + followUp: [], + }); + const retry = controller.retry("task-1"); + await vi.waitFor(() => expect(session.retry).toHaveBeenCalledOnce()); + + controller.disconnect("task-1"); + resolveRetry(); + await retry; + + expect(session.client.prompt).not.toHaveBeenCalled(); }); it("retries a cloud session without discarding its transcript", async () => { @@ -552,6 +666,54 @@ describe("PiSessionController", () => { expect(controller.store.getState().sessions["task-1"].events).toEqual([]); }); + it("tracks native queue updates without adding them to the transcript", async () => { + let onEvent: (event: AgentConversationEvent) => void = () => {}; + const session = createSession(); + vi.mocked(session.onConversationEvent).mockImplementation((handler) => { + onEvent = handler; + return () => {}; + }); + const controller = createController(session); + + await controller.connect("task-1"); + onEvent({ + type: "queue_update", + timestamp: 1, + steering: ["fix this"], + followUp: ["then summarize"], + }); + + expect(controller.store.getState().sessions["task-1"]).toMatchObject({ + events: [], + queue: { + steering: ["fix this"], + followUp: ["then summarize"], + }, + status: { pendingMessageCount: 2 }, + }); + }); + + it("clears the native queue and returns its contents for editing", async () => { + const session = createSession(); + vi.mocked(session.clearQueue).mockResolvedValue({ + steering: ["fix this"], + followUp: ["then summarize"], + }); + const controller = createController(session); + + await controller.connect("task-1"); + const queue = await controller.clearQueue("task-1"); + + expect(queue).toEqual({ + steering: ["fix this"], + followUp: ["then summarize"], + }); + expect(controller.store.getState().sessions["task-1"].queue).toEqual({ + steering: [], + followUp: [], + }); + }); + it("uses live turn completion without reloading native history", async () => { const turnCompleted: AgentConversationEvent = { type: "turn_completed", diff --git a/packages/core/src/pi-runtime/piSessionController.ts b/packages/core/src/pi-runtime/piSessionController.ts index db860eea65..e5885798ac 100644 --- a/packages/core/src/pi-runtime/piSessionController.ts +++ b/packages/core/src/pi-runtime/piSessionController.ts @@ -1,7 +1,7 @@ import type { PiRemoteRpcClient } from "@posthog/agent/pi/remote-rpc-client"; import type { PiNativeModelInfo, - PiQueueMode, + PiQueueSnapshot, PiThinkingLevel, } from "@posthog/agent/pi/types"; import type { @@ -22,7 +22,7 @@ import { export type { PiNativeModelInfo, - PiQueueMode, + PiQueueSnapshot, PiThinkingLevel, } from "@posthog/agent/pi/types"; @@ -38,6 +38,8 @@ export interface PiSession { readonly resumeRequired?: boolean; readonly cloudStatus?: TaskRunStatus; retry?(): Promise; + getQueue(): Promise; + clearQueue(): Promise; sendUserMessage?( type: "prompt" | "steer" | "follow_up", message: string, @@ -88,6 +90,8 @@ export class PiSessionController { private readonly connections = new Map>(); private readonly readiness = new Map>(); private readonly sessionVersions = new Map(); + private readonly queueRevisions = new Map(); + private readonly queuesToRestore = new Map(); private readonly taskRunIds = new Map(); constructor( @@ -165,10 +169,13 @@ export class PiSessionController { this.resetTransport(taskId); this.taskRunIds.delete(taskId); this.liveEvents.delete(taskId); + this.queueRevisions.delete(taskId); + this.queuesToRestore.delete(taskId); } async retry(taskId: string): Promise { const taskRunId = this.taskRunIds.get(taskId); + this.captureQueueForRestore(taskId); const session = await this.getPiSession(taskId); this.updateSession(taskId, { connectionState: "connecting", @@ -186,8 +193,17 @@ export class PiSessionController { } } + async clearQueue(taskId: string): Promise { + const session = await this.getPiSession(taskId); + const queue = await session.clearQueue(); + this.queuesToRestore.delete(taskId); + this.applyQueue(taskId, { steering: [], followUp: [] }); + return queue; + } + async restart(taskId: string): Promise { const taskRunId = this.taskRunIds.get(taskId); + this.captureQueueForRestore(taskId); if (!taskRunId) { await this.retry(taskId); return; @@ -254,7 +270,18 @@ export class PiSessionController { const action = this.getSubmitAction(message, isStreaming, messagingMode); const currentSession = await this.getPiSession(taskId); - const wasStreaming = this.getSession(taskId).status?.isStreaming ?? false; + const controllerSession = this.getSession(taskId); + const wasStreaming = controllerSession.status?.isStreaming ?? false; + const queuesMessage = action === "steer" || action === "followUp"; + const queuedMessageCount = + controllerSession.queue.steering.length + + controllerSession.queue.followUp.length; + if (queuesMessage && queuedMessageCount > 0) { + throw new Error("Pi already has a queued message"); + } + const refreshAfterSubmit = + action === "compact" || + this.isExtensionCommand(controllerSession, message); if (action === "compact") { const session = await this.getWritablePiSession(taskId); const command = parseCommandLine(message); @@ -265,9 +292,18 @@ export class PiSessionController { const messageId = currentSession.sendUserMessage ? globalThis.crypto.randomUUID() : undefined; - if (messageId) { + const hasOptimisticTranscriptMessage = Boolean( + messageId && action === "prompt", + ); + if (messageId && hasOptimisticTranscriptMessage) { this.appendOptimisticUserMessage(taskId, messageId, message); } + if (queuesMessage) { + this.applyQueue(taskId, { + steering: action === "steer" ? [message] : [], + followUp: action === "followUp" ? [message] : [], + }); + } this.markTurnPending(taskId); if (currentSession.resumeRequired) { this.updateSession(taskId, { connectionState: "connecting" }); @@ -298,16 +334,24 @@ export class PiSessionController { } else { await session.client.followUp(message); } + if (queuesMessage) { + await this.refreshQueue(taskId, session); + } } catch (error) { - if (messageId) { + if (messageId && hasOptimisticTranscriptMessage) { this.removeUserMessage(taskId, messageId); } + if (queuesMessage) { + this.applyQueue(taskId, controllerSession.queue); + } this.setTurnStreaming(taskId, wasStreaming); throw error; } } - await this.refreshStatus(taskId); + if (refreshAfterSubmit) { + await this.refreshStatus(taskId); + } return action; } @@ -331,20 +375,6 @@ export class PiSessionController { await this.refreshStatus(taskId); } - async setQueueMode( - taskId: string, - messagingMode: PiMessagingMode, - queueMode: PiQueueMode, - ): Promise { - const session = await this.getPiSession(taskId); - if (messagingMode === "steer") { - await session.client.setSteeringMode(queueMode); - } else { - await session.client.setFollowUpMode(queueMode); - } - await this.refreshStatus(taskId); - } - async bash(taskId: string, command: string): Promise { this.updateSession(taskId, { isBashRunning: true }); try { @@ -419,18 +449,28 @@ export class PiSessionController { const connectedSessionVersion = this.getSessionVersion(taskId); try { const session = await this.getPiSession(taskId); - const events = await session.getConversation(); - const status = await session.client.getState(); + const queueRevision = this.queueRevisions.get(taskId) ?? 0; + const [events, status, queue] = await Promise.all([ + session.getConversation(), + session.client.getState(), + session.getQueue(), + ]); if (this.getSessionVersion(taskId) !== connectedSessionVersion) { return; } const currentSession = this.getSession(taskId); + const conversationEvents = events.filter( + (event) => event.type !== "queue_update", + ); const liveEvents = this.liveEvents.get(taskId) ?? []; - const newLiveEvents = this.reconcileLiveEvents(events, liveEvents); + const newLiveEvents = this.reconcileLiveEvents( + conversationEvents, + liveEvents, + ); this.liveEvents.set(taskId, newLiveEvents); const historyUserMessageIds = new Set( - events.flatMap((event) => + conversationEvents.flatMap((event) => event.type === "user_message" ? [event.id] : [], ), ); @@ -441,26 +481,38 @@ export class PiSessionController { !historyUserMessageIds.has(event.id)), ); const reconciledEvents = [ - ...events, + ...conversationEvents, ...newLiveEvents, ...optimisticEvents, ]; + const resolvedQueue = + (this.queueRevisions.get(taskId) ?? 0) === queueRevision + ? queue + : currentSession.queue; + const resolvedStatus = { + ...status, + pendingMessageCount: + resolvedQueue.steering.length + resolvedQueue.followUp.length, + }; this.setSession(taskId, { connectionState: "connected", events: reconciledEvents, - status, + status: resolvedStatus, models: currentSession.models, modelsLoaded: currentSession.modelsLoaded, thinkingLevels: currentSession.thinkingLevels, thinkingLevelsLoaded: currentSession.thinkingLevelsLoaded, commands: currentSession.commands, + queue: resolvedQueue, isBashRunning: false, errorTitle: undefined, errorMessage: undefined, errorRetryable: undefined, }); + await this.restoreQueueIfNeeded(taskId, session, resolvedStatus); + await Promise.all([ session.client.getAvailableModels().then((models) => { if (this.getSessionVersion(taskId) === connectedSessionVersion) { @@ -490,6 +542,15 @@ export class PiSessionController { } private handleEvent(taskId: string, event: AgentConversationEvent): void { + if (event.type === "queue_update") { + const queue = { + steering: event.steering, + followUp: event.followUp, + }; + this.applyQueue(taskId, queue); + return; + } + const liveEvents = [...(this.liveEvents.get(taskId) ?? []), event]; this.liveEvents.set(taskId, liveEvents); const session = this.getSession(taskId); @@ -552,6 +613,107 @@ export class PiSessionController { ); } + private captureQueueForRestore(taskId: string): void { + const queue = this.getSession(taskId).queue; + if (queue.steering.length === 0 && queue.followUp.length === 0) { + return; + } + this.queuesToRestore.set(taskId, { + steering: [...queue.steering], + followUp: [...queue.followUp], + }); + } + + private async restoreQueueIfNeeded( + taskId: string, + session: PiSession, + status: NonNullable, + ): Promise { + const queue = this.getSession(taskId).queue; + const queueToRestore = this.queuesToRestore.get(taskId); + if (!queueToRestore) { + return; + } + if (queue.steering.length > 0 || queue.followUp.length > 0) { + this.queuesToRestore.delete(taskId); + return; + } + + const messages = [ + ...queueToRestore.steering.map((content) => ({ + content, + mode: "steer" as const, + })), + ...queueToRestore.followUp.map((content) => ({ + content, + mode: "follow_up" as const, + })), + ]; + if (!status.isStreaming) { + const first = messages.shift(); + if (first) { + await session.client.prompt(first.content); + } + } + + for (const message of messages) { + if (message.mode === "steer") { + await session.client.steer(message.content); + } else { + await session.client.followUp(message.content); + } + } + this.queuesToRestore.delete(taskId); + } + + private async refreshQueue( + taskId: string, + session: PiSession, + ): Promise { + try { + const queue = await session.getQueue(); + this.applyQueue(taskId, queue); + } catch { + return; + } + } + + private applyQueue(taskId: string, queue: PiQueueSnapshot): void { + this.queueRevisions.set(taskId, (this.queueRevisions.get(taskId) ?? 0) + 1); + this.updateSession(taskId, { + queue, + status: this.withPendingMessageCount(taskId, queue), + }); + } + + private withPendingMessageCount( + taskId: string, + queue: PiQueueSnapshot, + ): PiControllerSessionState["status"] { + const status = this.getSession(taskId).status; + if (!status) { + return undefined; + } + return { + ...status, + pendingMessageCount: queue.steering.length + queue.followUp.length, + }; + } + + private isExtensionCommand( + session: PiControllerSessionState, + message: string, + ): boolean { + const command = parseCommandLine(message); + if (!command) { + return false; + } + return session.commands.some( + (available) => + available.name === command.name && available.source === "extension", + ); + } + private markTurnPending(taskId: string): void { this.setTurnStreaming(taskId, true); } diff --git a/packages/core/src/pi-runtime/piSessionProvider.test.ts b/packages/core/src/pi-runtime/piSessionProvider.test.ts index f73be39cb2..36ce1d9250 100644 --- a/packages/core/src/pi-runtime/piSessionProvider.test.ts +++ b/packages/core/src/pi-runtime/piSessionProvider.test.ts @@ -16,8 +16,6 @@ function localSession(): PiSession { compact: vi.fn(async () => undefined), setModel: vi.fn(async () => ({ provider: "posthog", id: "model" })), setThinkingLevel: vi.fn(async () => {}), - setSteeringMode: vi.fn(async () => {}), - setFollowUpMode: vi.fn(async () => {}), bash: vi.fn(async () => undefined), abort: vi.fn(async () => {}), abortBash: vi.fn(async () => {}), @@ -27,6 +25,8 @@ function localSession(): PiSession { client, health: vi.fn(async () => ({ state: "idle" as const })), getConversation: vi.fn(async () => []), + getQueue: vi.fn(async () => ({ steering: [], followUp: [] })), + clearQueue: vi.fn(async () => ({ steering: [], followUp: [] })), onConversationEvent: vi.fn(() => () => {}), }; } diff --git a/packages/core/src/pi-runtime/piSessionStore.ts b/packages/core/src/pi-runtime/piSessionStore.ts index 33428060cd..bad7e09a7c 100644 --- a/packages/core/src/pi-runtime/piSessionStore.ts +++ b/packages/core/src/pi-runtime/piSessionStore.ts @@ -1,6 +1,7 @@ import type { PiCommand, PiNativeModelInfo, + PiQueueSnapshot, PiSessionStatus, PiThinkingLevel, } from "@posthog/agent/pi/types"; @@ -19,6 +20,7 @@ export interface PiControllerSessionState { thinkingLevels: PiThinkingLevel[]; thinkingLevelsLoaded: boolean; commands: PiCommand[]; + queue: PiQueueSnapshot; status?: PiSessionStatus; cloudStatus?: TaskRunStatus; errorTitle?: string; @@ -46,6 +48,7 @@ export function createEmptyPiControllerSession(): PiControllerSessionState { thinkingLevels: [], thinkingLevelsLoaded: false, commands: [], + queue: { steering: [], followUp: [] }, isBashRunning: false, }; } diff --git a/packages/host-router/src/pi-session-factory.ts b/packages/host-router/src/pi-session-factory.ts index 53548e8a93..9ff1c0eaad 100644 --- a/packages/host-router/src/pi-session-factory.ts +++ b/packages/host-router/src/pi-session-factory.ts @@ -36,6 +36,14 @@ class TrpcPiSession implements PiSession { return getRemotePiConversation(this.client); } + getQueue() { + return this.hostClient.piSession.getQueue.query({ taskId: this.taskId }); + } + + clearQueue() { + return this.hostClient.piSession.clearQueue.mutate({ taskId: this.taskId }); + } + onConversationEvent( onEvent: Parameters[0], onError: Parameters[1], diff --git a/packages/host-router/src/routers/pi-session.router.ts b/packages/host-router/src/routers/pi-session.router.ts index 067948189b..1ab86122b2 100644 --- a/packages/host-router/src/routers/pi-session.router.ts +++ b/packages/host-router/src/routers/pi-session.router.ts @@ -2,6 +2,7 @@ import { publicProcedure, router } from "@posthog/host-trpc/trpc"; import { PI_SESSION_SERVICE } from "@posthog/workspace-server/services/pi-session/identifiers"; import type { PiSessionService } from "@posthog/workspace-server/services/pi-session/pi-session"; import { + piQueueSnapshotOutput, piRpcResponseSchema, piSessionHealthOutput, piSessionRpcInput, @@ -40,6 +41,20 @@ export const piSessionRouter = router({ .output(piSessionHealthOutput) .query(({ ctx, input }) => getService(ctx.container).health(input.taskId)), + getQueue: publicProcedure + .input(piSessionTaskInput) + .output(piQueueSnapshotOutput) + .query(({ ctx, input }) => + getService(ctx.container).getQueue(input.taskId), + ), + + clearQueue: publicProcedure + .input(piSessionTaskInput) + .output(piQueueSnapshotOutput) + .mutation(({ ctx, input }) => + getService(ctx.container).clearQueue(input.taskId), + ), + onEvent: publicProcedure .input(piSessionTaskInput) .subscription(async function* (opts) { diff --git a/packages/shared/src/agent-conversation.ts b/packages/shared/src/agent-conversation.ts index cffe371dcc..545b182e06 100644 --- a/packages/shared/src/agent-conversation.ts +++ b/packages/shared/src/agent-conversation.ts @@ -150,6 +150,12 @@ export type AgentConversationEvent = ( group: string; detail?: string; } + | { + type: "queue_update"; + timestamp: number; + steering: string[]; + followUp: string[]; + } | { type: "runtime_status"; timestamp: number; diff --git a/packages/ui/src/features/pi-sessions/PiQueuedMessagesDock.test.tsx b/packages/ui/src/features/pi-sessions/PiQueuedMessagesDock.test.tsx new file mode 100644 index 0000000000..d245710e2b --- /dev/null +++ b/packages/ui/src/features/pi-sessions/PiQueuedMessagesDock.test.tsx @@ -0,0 +1,67 @@ +import { Theme } from "@radix-ui/themes"; +import { fireEvent, render, screen } from "@testing-library/react"; +import { describe, expect, it, vi } from "vitest"; +import { PiQueuedMessagesDock } from "./PiQueuedMessagesDock"; + +describe("PiQueuedMessagesDock", () => { + it("renders the single native queued message with ACP-style actions", () => { + const onEdit = vi.fn(); + const onRemove = vi.fn(); + + render( + + + , + ); + + expect(screen.getByText("then summarize")).toBeInTheDocument(); + + fireEvent.click( + screen.getByRole("button", { name: "Edit queued message" }), + ); + expect(onEdit).toHaveBeenCalledOnce(); + + fireEvent.click( + screen.getByRole("button", { name: "Discard queued message" }), + ); + expect(onRemove).toHaveBeenCalledOnce(); + }); + + it("preserves legacy multi-message queues behind one edit action", () => { + render( + + + , + ); + + expect(screen.getByText("first")).toBeInTheDocument(); + expect(screen.getByText("second")).toBeInTheDocument(); + expect(screen.getByText("third")).toBeInTheDocument(); + expect( + screen.queryByRole("button", { name: "Discard queued message" }), + ).not.toBeInTheDocument(); + }); + + it("does not render an empty queue", () => { + const { container } = render( + , + ); + + expect(container).toBeEmptyDOMElement(); + }); +}); diff --git a/packages/ui/src/features/pi-sessions/PiQueuedMessagesDock.tsx b/packages/ui/src/features/pi-sessions/PiQueuedMessagesDock.tsx new file mode 100644 index 0000000000..8b5ff30bdb --- /dev/null +++ b/packages/ui/src/features/pi-sessions/PiQueuedMessagesDock.tsx @@ -0,0 +1,34 @@ +import type { PiQueueSnapshot } from "@posthog/core/pi-runtime/piSessionController"; +import { QueuedMessageView } from "@posthog/ui/features/sessions/components/session-update/QueuedMessageView"; + +interface PiQueuedMessagesDockProps { + queue: PiQueueSnapshot; + onEdit(): void; + onRemove(): void; +} + +export function PiQueuedMessagesDock({ + queue, + onEdit, + onRemove, +}: PiQueuedMessagesDockProps) { + const messages = [...queue.steering, ...queue.followUp]; + const content = messages.join("\n\n"); + if (!content) { + return null; + } + + return ( +
+ +
+ ); +} diff --git a/packages/ui/src/features/pi-sessions/PiSessionControls.tsx b/packages/ui/src/features/pi-sessions/PiSessionControls.tsx index 0fe0e6d45a..a295fa5ea4 100644 --- a/packages/ui/src/features/pi-sessions/PiSessionControls.tsx +++ b/packages/ui/src/features/pi-sessions/PiSessionControls.tsx @@ -1,7 +1,6 @@ import { Brain, CaretDown, Lightning, Stack } from "@phosphor-icons/react"; import type { PiModelSelection, - PiQueueMode, PiThinkingLevel, } from "@posthog/core/pi-runtime/piSessionController"; import { @@ -177,20 +176,16 @@ export function PiThinkingLevelSelector({ interface PiMessagingModeSelectorProps { mode: MessagingMode; - queueMode: PiQueueMode; queuedCount: number; disabled?: boolean; onModeChange: (mode: MessagingMode) => void; - onQueueModeChange: (mode: PiQueueMode) => void; } export function PiMessagingModeSelector({ mode, - queueMode, queuedCount, disabled, onModeChange, - onQueueModeChange, }: PiMessagingModeSelectorProps) { let label = "Queue"; if (mode === "steer") { @@ -242,17 +237,6 @@ export function PiMessagingModeSelector({ Queue for the next turn - - Process queued messages - onQueueModeChange(value as PiQueueMode)} - > - - One per turn - - All at once - ); diff --git a/packages/ui/src/features/pi-sessions/PiSessionView.tsx b/packages/ui/src/features/pi-sessions/PiSessionView.tsx index abede68273..0b333bc7a7 100644 --- a/packages/ui/src/features/pi-sessions/PiSessionView.tsx +++ b/packages/ui/src/features/pi-sessions/PiSessionView.tsx @@ -1,7 +1,10 @@ +import { + contentToXml, + xmlToContent, +} from "@posthog/core/message-editor/content"; import { PI_SESSION_CONTROLLER } from "@posthog/core/pi-runtime/identifiers"; import type { PiModelSelection, - PiQueueMode, PiSessionController, PiThinkingLevel, } from "@posthog/core/pi-runtime/piSessionController"; @@ -26,7 +29,6 @@ import { import { ChatThread } from "@posthog/ui/features/sessions/components/chat-thread/ChatThread"; import type { PromptRecallHandler } from "@posthog/ui/features/sessions/components/chat-thread/composerPromptRecall"; import { CHAT_CONTENT_MAX_WIDTH } from "@posthog/ui/features/sessions/constants"; -import { useMessagingMode } from "@posthog/ui/features/sessions/hooks/useMessagingMode"; import { useMessagingModeStore } from "@posthog/ui/features/sessions/messagingModeStore"; import { useWorkspace } from "@posthog/ui/features/workspace/useWorkspace"; import { useConnectivity } from "@posthog/ui/hooks/useConnectivity"; @@ -35,6 +37,7 @@ import { TaskDetailSkeleton } from "@posthog/ui/router/routeSkeletons"; import { Box, Flex } from "@radix-ui/themes"; import { type ReactElement, useCallback, useEffect, useRef } from "react"; import { useStore } from "zustand"; +import { PiQueuedMessagesDock } from "./PiQueuedMessagesDock"; import { PiMessagingModeSelector, PiModelSelector, @@ -57,7 +60,9 @@ export function PiSessionView({ taskId, taskRunId }: PiSessionViewProps) { const draftActions = useDraftStore((state) => state.actions); const workspace = useWorkspace(taskId); const repoPath = workspace?.worktreePath ?? workspace?.folderPath; - const messagingMode = useMessagingMode(taskId); + const messagingMode = useMessagingModeStore( + (state) => state.modesByTaskId[taskId] ?? "steer", + ); const setMessagingMode = useMessagingModeStore((state) => state.setMode); const { isOnline } = useConnectivity(); const promptRecallRef = useRef(null); @@ -163,15 +168,6 @@ export function PiSessionView({ taskId, taskRunId }: PiSessionViewProps) { [piSessionController, taskId], ); - const setQueueMode = useCallback( - (mode: PiQueueMode) => { - void piSessionController - .setQueueMode(taskId, messagingMode, mode) - .catch(() => toast.error("Failed to change Pi queue behavior")); - }, - [messagingMode, piSessionController, taskId], - ); - const toggleMessagingMode = useCallback(() => { const nextMode = messagingMode === "steer" ? "queue" : "steer"; setMessagingMode(taskId, nextMode); @@ -204,6 +200,32 @@ export function PiSessionView({ taskId, taskRunId }: PiSessionViewProps) { .catch(() => toast.error("Failed to restart Pi")); }, [piSessionController, taskId]); + const editQueuedMessage = useCallback(() => { + void piSessionController + .clearQueue(taskId) + .then((queue) => { + const queuedText = [...queue.steering, ...queue.followUp].join("\n\n"); + if (!queuedText) { + return; + } + const draft = draftActions.getDraft(taskId); + const draftText = + typeof draft === "string" ? draft : draft ? contentToXml(draft) : ""; + const content = [queuedText, draftText] + .filter((value) => value.trim()) + .join("\n\n"); + draftActions.setPendingContent(taskId, xmlToContent(content)); + draftActions.requestFocus(taskId); + }) + .catch(() => toast.error("Failed to edit queued Pi message")); + }, [draftActions, piSessionController, taskId]); + + const removeQueuedMessage = useCallback(() => { + void piSessionController + .clearQueue(taskId) + .catch(() => toast.error("Failed to discard queued Pi message")); + }, [piSessionController, taskId]); + if (!session) { return ; } @@ -254,6 +276,8 @@ export function PiSessionView({ taskId, taskRunId }: PiSessionViewProps) { } const controlsPending = status ? isStreaming || isBashRunning : false; + const hasQueuedMessage = + session.queue.steering.length + session.queue.followUp.length > 0; let modelSelector: ReactElement = ; let reasoningSelector: ReactElement | null = ( @@ -286,16 +310,12 @@ export function PiSessionView({ taskId, taskRunId }: PiSessionViewProps) { } if (status) { - const queueMode = - messagingMode === "steer" ? status.steeringMode : status.followUpMode; messagingModeToggle = ( setMessagingMode(taskId, mode)} - onQueueModeChange={setQueueMode} /> ); } @@ -326,6 +346,11 @@ export function PiSessionView({ taskId, taskRunId }: PiSessionViewProps) { className="mx-auto w-full px-2 pb-3" style={{ maxWidth: CHAT_CONTENT_MAX_WIDTH }} > + - + {dragHandleRef && ( + + )} { }); }); + it("persists and rehydrates the last used agent runtime", async () => { + useSettingsStore.getState().setLastUsedAgentRuntime("pi"); + + await waitForPersistedWrite(); + + const lastCall = setItem.mock.calls[setItem.mock.calls.length - 1]; + const persisted = JSON.parse(lastCall[1]); + expect(persisted.state.lastUsedAgentRuntime).toBe("pi"); + + getItem.mockResolvedValue( + JSON.stringify({ + state: { lastUsedAgentRuntime: "pi" }, + version: 0, + }), + ); + useSettingsStore.setState({ lastUsedAgentRuntime: "acp" }); + + await useSettingsStore.persist.rehydrate(); + + expect(useSettingsStore.getState().lastUsedAgentRuntime).toBe("pi"); + }); + it("persists the last used cloud repository", async () => { useSettingsStore.getState().setLastUsedCloudRepository("posthog/posthog"); diff --git a/packages/ui/src/features/settings/settingsStore.ts b/packages/ui/src/features/settings/settingsStore.ts index 2d71ae4c17..197fe99281 100644 --- a/packages/ui/src/features/settings/settingsStore.ts +++ b/packages/ui/src/features/settings/settingsStore.ts @@ -1,5 +1,10 @@ import type { UserRepositoryIntegrationRef } from "@posthog/core/integrations/repositories"; -import type { Adapter, ExecutionMode, WorkspaceMode } from "@posthog/shared"; +import type { + Adapter, + AgentRuntime, + ExecutionMode, + WorkspaceMode, +} from "@posthog/shared"; import { COLLAPSE_MODE_DEFAULT, type CollapseMode, @@ -103,6 +108,7 @@ interface SettingsStore { lastUsedRunMode: "local" | "cloud"; lastUsedLocalWorkspaceMode: LocalWorkspaceMode; lastUsedWorkspaceMode: WorkspaceMode; + lastUsedAgentRuntime: AgentRuntime; lastUsedAdapter: AgentAdapter; lastUsedModel: string | null; lastUsedReasoningEffort: string | null; @@ -124,6 +130,7 @@ interface SettingsStore { setLastUsedRunMode: (mode: "local" | "cloud") => void; setLastUsedLocalWorkspaceMode: (mode: LocalWorkspaceMode) => void; setLastUsedWorkspaceMode: (mode: WorkspaceMode) => void; + setLastUsedAgentRuntime: (runtime: AgentRuntime) => void; setLastUsedAdapter: (adapter: AgentAdapter) => void; setLastUsedModel: (model: string) => void; setLastUsedReasoningEffort: (effort: string) => void; @@ -296,6 +303,7 @@ export const useSettingsStore = create()( lastUsedRunMode: "local", lastUsedLocalWorkspaceMode: "local", lastUsedWorkspaceMode: DEFAULT_WORKSPACE_MODE, + lastUsedAgentRuntime: "acp", lastUsedAdapter: "claude", lastUsedModel: null, lastUsedReasoningEffort: null, @@ -313,6 +321,8 @@ export const useSettingsStore = create()( setLastUsedLocalWorkspaceMode: (mode) => set({ lastUsedLocalWorkspaceMode: mode }), setLastUsedWorkspaceMode: (mode) => set({ lastUsedWorkspaceMode: mode }), + setLastUsedAgentRuntime: (runtime) => + set({ lastUsedAgentRuntime: runtime }), setLastUsedAdapter: (adapter) => set({ lastUsedAdapter: adapter }), setLastUsedModel: (model) => set({ lastUsedModel: model }), setLastUsedReasoningEffort: (effort) => @@ -521,6 +531,7 @@ export const useSettingsStore = create()( lastUsedRunMode: state.lastUsedRunMode, lastUsedLocalWorkspaceMode: state.lastUsedLocalWorkspaceMode, lastUsedWorkspaceMode: state.lastUsedWorkspaceMode, + lastUsedAgentRuntime: state.lastUsedAgentRuntime, lastUsedAdapter: state.lastUsedAdapter, lastUsedModel: state.lastUsedModel, lastUsedReasoningEffort: state.lastUsedReasoningEffort, diff --git a/packages/ui/src/features/task-detail/components/TaskInput.tsx b/packages/ui/src/features/task-detail/components/TaskInput.tsx index 2e2a1d798d..06f270c347 100644 --- a/packages/ui/src/features/task-detail/components/TaskInput.tsx +++ b/packages/ui/src/features/task-detail/components/TaskInput.tsx @@ -206,6 +206,8 @@ export function TaskInput({ lastUsedLocalWorkspaceMode, lastUsedWorkspaceMode, setLastUsedWorkspaceMode, + lastUsedAgentRuntime, + setLastUsedAgentRuntime, lastUsedAdapter, setLastUsedAdapter, lastUsedCloudRepository, @@ -250,6 +252,7 @@ export function TaskInput({ const [isCreatingBranch, setIsCreatingBranch] = useState(false); const [selectedBranch, setSelectedBranch] = useState(null); const [runtime, setRuntime] = useState("acp"); + const didResolveRuntimeRef = useRef(false); const [selectedPiModelId, setSelectedPiModelId] = useState( null, ); @@ -370,6 +373,16 @@ export function TaskInput({ hasGithubIntegration, }); + useEffect(() => { + if (didResolveRuntimeRef.current || !settingsHydrated || !flagsLoaded) { + return; + } + didResolveRuntimeRef.current = true; + setRuntime( + piHarnessEnabled && lastUsedAgentRuntime === "pi" ? "pi" : "acp", + ); + }, [flagsLoaded, lastUsedAgentRuntime, piHarnessEnabled, settingsHydrated]); + const [workspaceMode, setWorkspaceModeState] = useState(() => { if (initialCloudRepository) return "cloud"; if (!localWorkspaces) return "cloud"; @@ -1023,12 +1036,14 @@ export function TaskInput({ const handleRuntimeChange = useCallback( (nextRuntime: AgentRuntime) => { + didResolveRuntimeRef.current = true; setRuntime(nextRuntime); + setLastUsedAgentRuntime(nextRuntime); if (nextRuntime === "pi") { useAutoresearchDraftStore.getState().clearDraft(sessionId); } }, - [sessionId], + [sessionId, setLastUsedAgentRuntime], ); const handlePiModelChange = useCallback( diff --git a/packages/workspace-server/src/services/pi-session/pi-session.ts b/packages/workspace-server/src/services/pi-session/pi-session.ts index 3eb3d11481..b7aa4634a1 100644 --- a/packages/workspace-server/src/services/pi-session/pi-session.ts +++ b/packages/workspace-server/src/services/pi-session/pi-session.ts @@ -1,6 +1,7 @@ import type { PiRpcClient } from "@posthog/agent/pi/rpc-client"; import type { RpcCommand, RpcResponse } from "@posthog/agent/pi/rpc-transport"; import type { PiRuntime } from "@posthog/agent/pi/runtime"; +import type { PiQueueSnapshot } from "@posthog/agent/pi/types"; import { ROOT_LOGGER, type RootLogger } from "@posthog/di/logger"; import { type AgentConversationEvent, @@ -192,6 +193,14 @@ export class PiSessionService extends TypedEventEmitter { } } + getQueue(taskId: string): Promise { + return this.requireSession(taskId).client.getQueue(); + } + + clearQueue(taskId: string): Promise { + return this.requireSession(taskId).client.clearQueue(); + } + async stop(taskId: string): Promise { await this.runExclusive(taskId, () => this.stopLocked(taskId)); } diff --git a/packages/workspace-server/src/services/pi-session/schemas.ts b/packages/workspace-server/src/services/pi-session/schemas.ts index 67c19db561..ff8a401cff 100644 --- a/packages/workspace-server/src/services/pi-session/schemas.ts +++ b/packages/workspace-server/src/services/pi-session/schemas.ts @@ -36,6 +36,11 @@ export const resumePiSessionInput = z.object({ export const piSessionTaskInput = z.object({ taskId: z.string() }); +export const piQueueSnapshotOutput = z.object({ + steering: z.array(z.string()), + followUp: z.array(z.string()), +}); + export const piSessionRpcInput = z.object({ taskId: z.string(), command: piRpcCommandSchema, From 103fa6e86bb84fc66b9f59e1f0d03b32c0dc271a Mon Sep 17 00:00:00 2001 From: JonathanLab Date: Mon, 27 Jul 2026 18:36:54 +0200 Subject: [PATCH 16/45] feat(pi): improve cloud error recovery --- .../pi-runtime/piSessionController.test.ts | 248 +++++++++++- .../src/pi-runtime/piSessionController.ts | 361 ++++++++++++++---- .../core/src/pi-runtime/piSessionStore.ts | 19 +- packages/shared/src/errors.test.ts | 16 + packages/shared/src/errors.ts | 64 ++++ packages/shared/src/index.ts | 3 + .../features/pi-sessions/PiSessionView.tsx | 145 +++++-- 7 files changed, 748 insertions(+), 108 deletions(-) diff --git a/packages/core/src/pi-runtime/piSessionController.test.ts b/packages/core/src/pi-runtime/piSessionController.test.ts index 395c775198..807dc59d08 100644 --- a/packages/core/src/pi-runtime/piSessionController.test.ts +++ b/packages/core/src/pi-runtime/piSessionController.test.ts @@ -1,8 +1,10 @@ import type { PiRemoteRpcClient } from "@posthog/agent/pi/remote-rpc-client"; +import type { AuthService } from "@posthog/core/auth/auth"; import type { TaskService } from "@posthog/core/task-detail/taskService"; import type { AgentConversationEvent } from "@posthog/shared"; import { describe, expect, it, vi } from "vitest"; import { + PiOperationError, type PiSession, PiSessionController, type PiSessionProvider, @@ -13,11 +15,12 @@ function createController( taskService = { openTask: vi.fn(async () => ({ success: true })), } as unknown as TaskService, + authService?: AuthService, ): PiSessionController { const provider: PiSessionProvider = { get: vi.fn(async () => session), }; - return new PiSessionController(provider, taskService); + return new PiSessionController(provider, taskService, authService); } function createSession(): PiSession { @@ -163,6 +166,230 @@ describe("PiSessionController", () => { ); }); + it("waits for cloud authentication restoration before sending", async () => { + let authStatus: "restoring" | "authenticated" = "restoring"; + let onStateChange: (state: { status: "authenticated" }) => void = () => {}; + const authService = { + getState: vi.fn(() => ({ status: authStatus })), + on: vi.fn((_event, handler) => { + onStateChange = handler; + }), + off: vi.fn(), + } as unknown as AuthService; + const session = createSession(); + session.sendUserMessage = vi.fn(async () => {}); + const controller = createController( + session, + { + prepareCloudPiMessage: vi.fn(async () => ({ + content: "hello", + artifactIds: [], + })), + } as unknown as TaskService, + authService, + ); + + await controller.connect("task-1", "run-1"); + const submission = controller.submit("task-1", "hello", false, "steer"); + await vi.waitFor(() => { + expect(controller.store.getState().sessions["task-1"].authRestoring).toBe( + true, + ); + }); + expect(session.sendUserMessage).not.toHaveBeenCalled(); + await expect( + controller.submit("task-1", "second", false, "steer"), + ).rejects.toMatchObject({ + failure: { + kind: "authentication", + recoveryPrompt: "second", + }, + }); + + authStatus = "authenticated"; + onStateChange({ status: "authenticated" }); + await submission; + + expect(session.sendUserMessage).toHaveBeenCalledOnce(); + expect(controller.store.getState().sessions["task-1"].authRestoring).toBe( + false, + ); + }); + + it("cancels auth-held submissions on disconnect and preserves the prompt", async () => { + const authService = { + getState: vi.fn(() => ({ status: "restoring" })), + on: vi.fn(), + off: vi.fn(), + } as unknown as AuthService; + const session = createSession(); + session.sendUserMessage = vi.fn(async () => {}); + const controller = createController( + session, + {} as TaskService, + authService, + ); + + await controller.connect("task-1", "run-1"); + const submission = controller.submit( + "task-1", + "do not lose this", + false, + "steer", + ); + await vi.waitFor(() => { + expect(controller.store.getState().sessions["task-1"].authRestoring).toBe( + true, + ); + }); + + controller.disconnect("task-1"); + + await expect(submission).rejects.toBeInstanceOf(PiOperationError); + expect(session.sendUserMessage).not.toHaveBeenCalled(); + expect(controller.store.getState().sessions["task-1"].error).toMatchObject({ + scope: "operation", + kind: "authentication", + recoveryPrompt: "do not lose this", + }); + }); + + it("classifies usage limits without failing the session", async () => { + const session = createSession(); + session.sendUserMessage = vi.fn(async () => { + throw new Error("Rate limit exceeded: User burst rate limit exceeded"); + }); + const controller = createController(session, { + prepareCloudPiMessage: vi.fn(async () => ({ + content: "hello", + artifactIds: [], + })), + } as unknown as TaskService); + + await controller.connect("task-1", "run-1"); + await expect( + controller.submit("task-1", "hello", false, "steer"), + ).rejects.toBeInstanceOf(PiOperationError); + + expect(controller.store.getState().sessions["task-1"]).toMatchObject({ + connectionState: "connected", + error: { + scope: "operation", + kind: "usage_limit", + title: "Usage limit reached", + limitCause: "org_limit", + }, + }); + }); + + it("classifies streamed transient provider errors as retryable", async () => { + let onEvent: (event: AgentConversationEvent) => void = () => {}; + const session = createSession(); + vi.mocked(session.onConversationEvent).mockImplementation((handler) => { + onEvent = handler; + return () => {}; + }); + const controller = createController(session); + + await controller.connect("task-1"); + onEvent({ + type: "runtime_error", + timestamp: 1, + errorType: "upstream_timeout", + message: "API Error: request timed out", + }); + + expect(controller.store.getState().sessions["task-1"].error).toMatchObject({ + scope: "operation", + kind: "transient", + title: "Provider temporarily unavailable", + retryable: true, + }); + expect(controller.store.getState().sessions["task-1"].connectionState).toBe( + "connected", + ); + }); + + it("keeps fatal runtime errors in a retryable disconnected state", async () => { + let onEvent: (event: AgentConversationEvent) => void = () => {}; + const session = createSession(); + vi.mocked(session.onConversationEvent).mockImplementation((handler) => { + onEvent = handler; + return () => {}; + }); + const controller = createController(session); + + await controller.connect("task-1"); + onEvent({ + type: "runtime_error", + timestamp: 1, + errorType: "agent_error", + message: "process exited unexpectedly", + }); + + expect(controller.store.getState().sessions["task-1"]).toMatchObject({ + connectionState: "disconnected", + error: { + scope: "connection", + kind: "fatal_session", + title: "Failed to send message", + retryable: true, + }, + }); + }); + + it("uses action-specific model errors", async () => { + const session = createSession(); + vi.mocked(session.client.setModel).mockRejectedValue( + new Error("Model is unavailable"), + ); + const controller = createController(session); + + await controller.connect("task-1"); + await expect( + controller.setModel("task-1", { provider: "posthog", id: "missing" }), + ).rejects.toBeInstanceOf(PiOperationError); + + expect(controller.store.getState().sessions["task-1"].error).toMatchObject({ + scope: "operation", + kind: "unknown", + title: "Failed to change Pi model", + message: "Model is unavailable", + }); + }); + + it("surfaces compaction failure details and resets compacting state", async () => { + let onEvent: (event: AgentConversationEvent) => void = () => {}; + const session = createSession(); + vi.mocked(session.onConversationEvent).mockImplementation((handler) => { + onEvent = handler; + return () => {}; + }); + const controller = createController(session); + + await controller.connect("task-1"); + onEvent({ + type: "runtime_status", + timestamp: 1, + status: "compacting", + }); + onEvent({ + type: "runtime_status", + timestamp: 2, + status: "compacting_failed", + error: "Summary request timed out", + }); + + expect(controller.store.getState().sessions["task-1"]).toMatchObject({ + status: { isCompacting: false }, + error: { + scope: "operation", + title: "Failed to compact Pi context", + message: "Summary request timed out", + }, + }); + }); + it("allows only one queued message and keeps it out of the transcript", async () => { const session = createSession(); session.sendUserMessage = vi.fn(async () => {}); @@ -322,8 +549,15 @@ describe("PiSessionController", () => { "task-1": { ...state.sessions["task-1"], connectionState: "disconnected", - errorMessage: "stream dropped", - errorRetryable: true, + error: { + id: "connection-error", + scope: "connection", + kind: "unknown", + title: "Connection failed", + message: "stream dropped", + retryable: true, + limitCause: null, + }, }, }, })); @@ -334,7 +568,7 @@ describe("PiSessionController", () => { expect(controller.store.getState().sessions["task-1"]).toMatchObject({ connectionState: "connected", events: [initialEvent], - errorMessage: undefined, + error: undefined, }); }); @@ -396,7 +630,11 @@ describe("PiSessionController", () => { expect(controller.store.getState().sessions["task-1"]).toMatchObject({ connectionState: "connected", events: [initialEvent], - errorMessage: undefined, + error: { + scope: "operation", + title: "Failed to send message", + message: "temporary command failure", + }, }); }); diff --git a/packages/core/src/pi-runtime/piSessionController.ts b/packages/core/src/pi-runtime/piSessionController.ts index e5885798ac..d253a6bbb4 100644 --- a/packages/core/src/pi-runtime/piSessionController.ts +++ b/packages/core/src/pi-runtime/piSessionController.ts @@ -4,19 +4,25 @@ import type { PiQueueSnapshot, PiThinkingLevel, } from "@posthog/agent/pi/types"; -import type { - AgentConversationEvent, - PiMessagingMode, - PiRuntimeHealth, - TaskRunStatus, +import { + type AgentConversationEvent, + classifyPromptFailure, + type PiMessagingMode, + type PiRuntimeHealth, + type PromptFailure, + type TaskRunStatus, } from "@posthog/shared"; -import { inject, injectable } from "inversify"; +import { inject, injectable, optional } from "inversify"; +import type { AuthService } from "../auth/auth"; +import { AUTH_SERVICE } from "../auth/auth.module"; +import { AuthServiceEvent } from "../auth/schemas"; import { parseCommandLine } from "../message-editor/commands"; import { TASK_SERVICE, type TaskService } from "../task-detail/taskService"; import { createEmptyPiControllerSession, createPiSessionStore, type PiControllerSessionState, + type PiSessionError, type PiSessionStore, } from "./piSessionStore"; @@ -63,6 +69,24 @@ export type PiSessionProvider = PiSessionFactory; export type PiSubmitResult = "prompt" | "steer" | "followUp" | "compact"; +type PiOperation = + | "prompt" + | "compact" + | "model" + | "thinking" + | "bash" + | "cancel" + | "queue" + | "retry" + | "restart"; + +export class PiOperationError extends Error { + constructor(readonly failure: PiSessionError) { + super(failure.message); + this.name = "PiOperationError"; + } +} + function normalizeSessionError(error: unknown): { title: string; message: string; @@ -92,11 +116,15 @@ export class PiSessionController { private readonly sessionVersions = new Map(); private readonly queueRevisions = new Map(); private readonly queuesToRestore = new Map(); + private readonly cancelAuthRestoration = new Map void>(); private readonly taskRunIds = new Map(); constructor( @inject(PI_SESSION_PROVIDER) private readonly provider: PiSessionProvider, @inject(TASK_SERVICE) private readonly taskService: TaskService, + @inject(AUTH_SERVICE) + @optional() + private readonly authService?: AuthService, ) {} ensureConnected(taskId: string, taskRunId?: string): Promise { @@ -110,9 +138,7 @@ export class PiSessionController { this.updateSession(taskId, { connectionState: "connecting", - errorTitle: undefined, - errorMessage: undefined, - errorRetryable: undefined, + error: undefined, }); const connectedSessionVersion = this.getSessionVersion(taskId); const readiness = this.ensureConnectedInternal(taskId) @@ -120,9 +146,7 @@ export class PiSessionController { if (this.getSessionVersion(taskId) === connectedSessionVersion) { this.updateSession(taskId, { connectionState: "connected", - errorTitle: undefined, - errorMessage: undefined, - errorRetryable: undefined, + error: undefined, }); } }) @@ -150,11 +174,7 @@ export class PiSessionController { return existing; } - this.updateSession(taskId, { - errorTitle: undefined, - errorMessage: undefined, - errorRetryable: undefined, - }); + this.updateSession(taskId, { error: undefined }); const connection = this.loadSession(taskId).finally(() => { if (this.connections.get(taskId) === connection) { @@ -166,6 +186,7 @@ export class PiSessionController { } disconnect(taskId: string): void { + this.cancelAuthRestoration.get(taskId)?.(); this.resetTransport(taskId); this.taskRunIds.delete(taskId); this.liveEvents.delete(taskId); @@ -179,26 +200,27 @@ export class PiSessionController { const session = await this.getPiSession(taskId); this.updateSession(taskId, { connectionState: "connecting", - errorTitle: undefined, - errorMessage: undefined, - errorRetryable: undefined, + error: undefined, }); try { await session.retry?.(); this.resetTransport(taskId); await this.ensureConnected(taskId, taskRunId); } catch (error) { - this.applySessionError(taskId, error); - throw error; + throw this.recordOperationFailure(taskId, "retry", error); } } async clearQueue(taskId: string): Promise { - const session = await this.getPiSession(taskId); - const queue = await session.clearQueue(); - this.queuesToRestore.delete(taskId); - this.applyQueue(taskId, { steering: [], followUp: [] }); - return queue; + try { + const session = await this.getPiSession(taskId); + const queue = await session.clearQueue(); + this.queuesToRestore.delete(taskId); + this.applyQueue(taskId, { steering: [], followUp: [] }); + return queue; + } catch (error) { + throw this.recordOperationFailure(taskId, "queue", error); + } } async restart(taskId: string): Promise { @@ -211,9 +233,7 @@ export class PiSessionController { this.updateSession(taskId, { connectionState: "connecting", - errorTitle: undefined, - errorMessage: undefined, - errorRetryable: undefined, + error: undefined, }); try { const resumedRun = await this.taskService.resumeCloudPiRun( @@ -223,8 +243,7 @@ export class PiSessionController { this.resetTransport(taskId); await this.ensureConnected(taskId, resumedRun.id); } catch (error) { - this.applySessionError(taskId, error); - throw error; + throw this.recordOperationFailure(taskId, "restart", error); } } @@ -234,7 +253,7 @@ export class PiSessionController { )) { if ( session.cloudStatus !== undefined && - session.errorRetryable && + session.error?.retryable && (session.connectionState === "disconnected" || session.connectionState === "error") ) { @@ -270,7 +289,38 @@ export class PiSessionController { const action = this.getSubmitAction(message, isStreaming, messagingMode); const currentSession = await this.getPiSession(taskId); + const submissionSessionVersion = this.getSessionVersion(taskId); + if (this.getSession(taskId).authRestoring) { + throw this.recordOperationFailure( + taskId, + "prompt", + new Error("Authentication required while the session restores"), + undefined, + message, + ); + } + if (currentSession.sendUserMessage) { + try { + await this.waitForAuthRestoration(taskId); + if (this.getSessionVersion(taskId) !== submissionSessionVersion) { + throw new Error( + "Authentication required; submission cancelled after session changed", + ); + } + } catch (error) { + throw this.recordOperationFailure( + taskId, + "prompt", + error, + undefined, + message, + ); + } + } const controllerSession = this.getSession(taskId); + if (controllerSession.error?.scope === "operation") { + this.updateSession(taskId, { error: undefined }); + } const wasStreaming = controllerSession.status?.isStreaming ?? false; const queuesMessage = action === "steer" || action === "followUp"; const queuedMessageCount = @@ -283,10 +333,14 @@ export class PiSessionController { action === "compact" || this.isExtensionCommand(controllerSession, message); if (action === "compact") { - const session = await this.getWritablePiSession(taskId); - const command = parseCommandLine(message); - const customInstructions = command?.args?.trim() || undefined; - await session.client.compact(customInstructions); + try { + const session = await this.getWritablePiSession(taskId); + const command = parseCommandLine(message); + const customInstructions = command?.args?.trim() || undefined; + await session.client.compact(customInstructions); + } catch (error) { + throw this.recordOperationFailure(taskId, "compact", error); + } } else { const commandType = action === "followUp" ? "follow_up" : action; const messageId = currentSession.sendUserMessage @@ -345,7 +399,8 @@ export class PiSessionController { this.applyQueue(taskId, controllerSession.queue); } this.setTurnStreaming(taskId, wasStreaming); - throw error; + const operation = queuesMessage ? "queue" : "prompt"; + throw this.recordOperationFailure(taskId, operation, error); } } @@ -356,23 +411,31 @@ export class PiSessionController { } async setModel(taskId: string, model: PiModelSelection): Promise { - const session = await this.getPiSession(taskId); - await session.client.setModel(model.provider, model.id); - await this.refreshStatus(taskId); - const thinkingLevels = await session.client.getAvailableThinkingLevels(); - this.updateSession(taskId, { - thinkingLevels, - thinkingLevelsLoaded: true, - }); + try { + const session = await this.getPiSession(taskId); + await session.client.setModel(model.provider, model.id); + await this.refreshStatus(taskId); + const thinkingLevels = await session.client.getAvailableThinkingLevels(); + this.updateSession(taskId, { + thinkingLevels, + thinkingLevelsLoaded: true, + }); + } catch (error) { + throw this.recordOperationFailure(taskId, "model", error); + } } async setThinkingLevel( taskId: string, level: PiThinkingLevel, ): Promise { - const session = await this.getPiSession(taskId); - await session.client.setThinkingLevel(level); - await this.refreshStatus(taskId); + try { + const session = await this.getPiSession(taskId); + await session.client.setThinkingLevel(level); + await this.refreshStatus(taskId); + } catch (error) { + throw this.recordOperationFailure(taskId, "thinking", error); + } } async bash(taskId: string, command: string): Promise { @@ -380,21 +443,31 @@ export class PiSessionController { try { const session = await this.getPiSession(taskId); await session.client.bash(command); + } catch (error) { + throw this.recordOperationFailure(taskId, "bash", error); } finally { this.updateSession(taskId, { isBashRunning: false }); } } async abort(taskId: string): Promise { - const session = await this.getPiSession(taskId); - await session.client.abort(); - await this.refreshStatus(taskId); + try { + const session = await this.getPiSession(taskId); + await session.client.abort(); + await this.refreshStatus(taskId); + } catch (error) { + throw this.recordOperationFailure(taskId, "cancel", error); + } } async abortBash(taskId: string): Promise { - const session = await this.getPiSession(taskId); - await session.client.abortBash(); - this.updateSession(taskId, { isBashRunning: false }); + try { + const session = await this.getPiSession(taskId); + await session.client.abortBash(); + this.updateSession(taskId, { isBashRunning: false }); + } catch (error) { + throw this.recordOperationFailure(taskId, "cancel", error); + } } private async ensureConnectedInternal(taskId: string): Promise { @@ -505,10 +578,12 @@ export class PiSessionController { thinkingLevelsLoaded: currentSession.thinkingLevelsLoaded, commands: currentSession.commands, queue: resolvedQueue, + error: + currentSession.error?.scope === "operation" + ? currentSession.error + : undefined, + authRestoring: currentSession.authRestoring, isBashRunning: false, - errorTitle: undefined, - errorMessage: undefined, - errorRetryable: undefined, }); await this.restoreQueueIfNeeded(taskId, session, resolvedStatus); @@ -551,6 +626,15 @@ export class PiSessionController { return; } + if (event.type === "runtime_error") { + this.recordOperationFailure( + taskId, + "prompt", + new Error(event.message), + event.errorType, + ); + } + const liveEvents = [...(this.liveEvents.get(taskId) ?? []), event]; this.liveEvents.set(taskId, liveEvents); const session = this.getSession(taskId); @@ -560,6 +644,11 @@ export class PiSessionController { status = { ...status, isCompacting: !event.isComplete }; } else if (event.status === "compacting_failed") { status = { ...status, isCompacting: false }; + this.recordOperationFailure( + taskId, + "compact", + new Error(event.error ?? event.message ?? "Compaction failed"), + ); } } const hasTurnActivity = @@ -588,14 +677,22 @@ export class PiSessionController { events.push(event); } + const latestSession = this.getSession(taskId); + const preserveConnectionError = + event.type === "runtime_error" && + latestSession.error?.scope === "connection"; + const preserveOperationError = latestSession.error?.scope === "operation"; this.updateSession(taskId, { connectionState: - event.type === "progress" ? session.connectionState : "connected", + event.type === "progress" || preserveConnectionError + ? latestSession.connectionState + : "connected", events, status, - errorTitle: undefined, - errorMessage: undefined, - errorRetryable: undefined, + error: + preserveConnectionError || preserveOperationError + ? latestSession.error + : undefined, }); } @@ -613,6 +710,133 @@ export class PiSessionController { ); } + acknowledgeOperationFailure(taskId: string, failureId: string): void { + const session = this.getSession(taskId); + if ( + session.error?.scope === "operation" && + session.error.id === failureId + ) { + this.updateSession(taskId, { error: undefined }); + } + } + + private async waitForAuthRestoration(taskId: string): Promise { + if ( + !this.authService || + this.authService.getState().status !== "restoring" + ) { + return; + } + + this.updateSession(taskId, { authRestoring: true }); + try { + await new Promise((resolve, reject) => { + const cleanup = () => { + this.authService?.off( + AuthServiceEvent.StateChanged, + handleStateChange, + ); + this.cancelAuthRestoration.delete(taskId); + }; + const handleStateChange = ( + state: ReturnType, + ) => { + if (state.status === "restoring") { + return; + } + cleanup(); + if (state.status === "authenticated") { + resolve(); + } else { + reject(new Error("Authentication required for cloud commands")); + } + }; + this.cancelAuthRestoration.set(taskId, () => { + cleanup(); + reject( + new Error( + "Authentication required; submission cancelled while restoring", + ), + ); + }); + this.authService?.on(AuthServiceEvent.StateChanged, handleStateChange); + if (this.authService) { + handleStateChange(this.authService.getState()); + } + }); + } finally { + this.cancelAuthRestoration.delete(taskId); + this.updateSession(taskId, { authRestoring: false }); + } + } + + private recordOperationFailure( + taskId: string, + operation: PiOperation, + error: unknown, + errorType?: string, + recoveryPrompt?: string, + ): PiOperationError { + const details = (error as { data?: { details?: string } })?.data?.details; + const classified = classifyPromptFailure(error, details, errorType); + const retryable = + classified.retryable || + ((operation === "retry" || operation === "restart") && + classified.kind === "unknown"); + const scope = + classified.kind === "fatal_session" || + operation === "retry" || + operation === "restart" + ? "connection" + : "operation"; + const failure: PiSessionError = { + id: globalThis.crypto.randomUUID(), + scope, + kind: classified.kind, + title: this.errorTitleForOperation(operation, classified), + message: classified.message, + retryable, + limitCause: classified.limitCause, + recoveryPrompt, + }; + this.updateSession(taskId, { + error: failure, + ...(scope === "connection" + ? { + connectionState: retryable ? "disconnected" : "error", + } + : {}), + }); + return new PiOperationError(failure); + } + + private errorTitleForOperation( + operation: PiOperation, + failure: PromptFailure, + ): string { + if (failure.kind === "usage_limit") { + return "Usage limit reached"; + } + if (failure.kind === "transient") { + return "Provider temporarily unavailable"; + } + if (failure.kind === "authentication") { + return "Authentication required"; + } + const titles: Record = { + prompt: "Failed to send message", + compact: "Failed to compact Pi context", + model: "Failed to change Pi model", + thinking: "Failed to change Pi thinking level", + bash: "Failed to run Pi bash command", + cancel: "Failed to stop Pi", + queue: "Failed to update queued message", + retry: "Failed to reconnect to Pi", + restart: "Failed to restart Pi", + }; + return titles[operation]; + } + private captureQueueForRestore(taskId: string): void { const queue = this.getSession(taskId).queue; if (queue.steering.length === 0 && queue.followUp.length === 0) { @@ -840,11 +1064,18 @@ export class PiSessionController { private applySessionError(taskId: string, error: unknown): void { const failure = normalizeSessionError(error); + const classified = classifyPromptFailure(error); this.updateSession(taskId, { connectionState: failure.retryable ? "disconnected" : "error", - errorTitle: failure.title, - errorMessage: failure.message, - errorRetryable: failure.retryable, + error: { + id: globalThis.crypto.randomUUID(), + scope: "connection", + kind: classified.kind, + title: failure.title, + message: failure.message, + retryable: failure.retryable, + limitCause: classified.limitCause, + }, }); } diff --git a/packages/core/src/pi-runtime/piSessionStore.ts b/packages/core/src/pi-runtime/piSessionStore.ts index bad7e09a7c..4453366a05 100644 --- a/packages/core/src/pi-runtime/piSessionStore.ts +++ b/packages/core/src/pi-runtime/piSessionStore.ts @@ -7,11 +7,24 @@ import type { } from "@posthog/agent/pi/types"; import type { AgentConversationEvent, + GatewayLimitCause, + PromptFailureKind, SessionStatus, TaskRunStatus, } from "@posthog/shared"; import { createStore, type StoreApi } from "zustand/vanilla"; +export interface PiSessionError { + id: string; + scope: "connection" | "operation"; + kind: PromptFailureKind; + title: string; + message: string; + retryable: boolean; + limitCause: GatewayLimitCause | null; + recoveryPrompt?: string; +} + export interface PiControllerSessionState { connectionState: SessionStatus; events: AgentConversationEvent[]; @@ -23,9 +36,8 @@ export interface PiControllerSessionState { queue: PiQueueSnapshot; status?: PiSessionStatus; cloudStatus?: TaskRunStatus; - errorTitle?: string; - errorMessage?: string; - errorRetryable?: boolean; + error?: PiSessionError; + authRestoring: boolean; isBashRunning: boolean; } @@ -49,6 +61,7 @@ export function createEmptyPiControllerSession(): PiControllerSessionState { thinkingLevelsLoaded: false, commands: [], queue: { steering: [], followUp: [] }, + authRestoring: false, isBashRunning: false, }; } diff --git a/packages/shared/src/errors.test.ts b/packages/shared/src/errors.test.ts index 05b39ca57f..4b5b59c8e1 100644 --- a/packages/shared/src/errors.test.ts +++ b/packages/shared/src/errors.test.ts @@ -1,6 +1,7 @@ import { describe, expect, it } from "vitest"; import { classifyGatewayLimitError, + classifyPromptFailure, getErrorMessage, isAuthError, isFatalSessionError, @@ -141,6 +142,21 @@ describe("classifyGatewayLimitError", () => { }); }); +describe("classifyPromptFailure", () => { + it.each([ + ["Rate limit exceeded", undefined, "usage_limit", false], + ["API Error: 529 overloaded", undefined, "transient", true], + ["boom", "upstream_timeout", "transient", true], + ["Authentication required", undefined, "authentication", true], + ["process exited", undefined, "fatal_session", true], + ["invalid model", undefined, "unknown", false], + ] as const)("classifies %j as %s", (message, errorType, kind, retryable) => { + expect( + classifyPromptFailure(new Error(message), undefined, errorType), + ).toMatchObject({ kind, retryable }); + }); +}); + describe("isFatalSessionError", () => { it.each([ "internal error", diff --git a/packages/shared/src/errors.ts b/packages/shared/src/errors.ts index 8a0508a832..3e98b7c60a 100644 --- a/packages/shared/src/errors.ts +++ b/packages/shared/src/errors.ts @@ -155,6 +155,70 @@ export function isTransientUpstreamError( ); } +export type PromptFailureKind = + | "usage_limit" + | "transient" + | "authentication" + | "fatal_session" + | "unknown"; + +export interface PromptFailure { + kind: PromptFailureKind; + message: string; + retryable: boolean; + limitCause: GatewayLimitCause | null; +} + +export function classifyPromptFailure( + error: unknown, + errorDetails?: string, + errorType?: string, +): PromptFailure { + const message = getErrorMessage(error) || String(error); + const limitCause = classifyGatewayLimitError(message, errorDetails); + if (limitCause !== null || isRateLimitError(message, errorDetails)) { + return { + kind: "usage_limit", + message, + retryable: false, + limitCause, + }; + } + if ( + errorType?.startsWith("upstream_") || + isTransientUpstreamError(message, errorDetails) + ) { + return { + kind: "transient", + message, + retryable: true, + limitCause: null, + }; + } + if (isNotAuthenticatedError(error) || isAuthError(error)) { + return { + kind: "authentication", + message, + retryable: true, + limitCause: null, + }; + } + if (isFatalSessionError(message, errorDetails)) { + return { + kind: "fatal_session", + message, + retryable: true, + limitCause: null, + }; + } + return { + kind: "unknown", + message, + retryable: false, + limitCause: null, + }; +} + export function isFatalSessionError( errorMessage: string, errorDetails?: string, diff --git a/packages/shared/src/index.ts b/packages/shared/src/index.ts index ff03602e7e..a837d48b97 100644 --- a/packages/shared/src/index.ts +++ b/packages/shared/src/index.ts @@ -93,6 +93,7 @@ export type { SignalReportPriority, Task } from "./domain-types"; export * from "./enrichment"; export { classifyGatewayLimitError, + classifyPromptFailure, type GatewayLimitCause, getErrorMessage, isAuthError, @@ -101,6 +102,8 @@ export { isRateLimitError, isTransientUpstreamError, NotAuthenticatedError, + type PromptFailure, + type PromptFailureKind, type SerializedError, serializeError, } from "./errors"; diff --git a/packages/ui/src/features/pi-sessions/PiSessionView.tsx b/packages/ui/src/features/pi-sessions/PiSessionView.tsx index 0b333bc7a7..f85e9be648 100644 --- a/packages/ui/src/features/pi-sessions/PiSessionView.tsx +++ b/packages/ui/src/features/pi-sessions/PiSessionView.tsx @@ -3,10 +3,11 @@ import { xmlToContent, } from "@posthog/core/message-editor/content"; import { PI_SESSION_CONTROLLER } from "@posthog/core/pi-runtime/identifiers"; -import type { - PiModelSelection, - PiSessionController, - PiThinkingLevel, +import { + type PiModelSelection, + PiOperationError, + type PiSessionController, + type PiThinkingLevel, } from "@posthog/core/pi-runtime/piSessionController"; import { useService } from "@posthog/di/react"; import { @@ -19,6 +20,7 @@ import { Skeleton, } from "@posthog/quill"; import type { AgentConversationEvent } from "@posthog/shared"; +import { useUsageLimitStore } from "@posthog/ui/features/billing/usageLimitStore"; import { PromptInput } from "@posthog/ui/features/message-editor/components/PromptInput"; import { useDraftStore } from "@posthog/ui/features/message-editor/draftStore"; import { CloudInitializingView } from "@posthog/ui/features/sessions/components/CloudInitializingView"; @@ -65,6 +67,7 @@ export function PiSessionView({ taskId, taskRunId }: PiSessionViewProps) { ); const setMessagingMode = useMessagingModeStore((state) => state.setMode); const { isOnline } = useConnectivity(); + const showUsageLimit = useUsageLimitStore((state) => state.show); const promptRecallRef = useRef(null); const handlePromptRecall = useCallback( (direction) => promptRecallRef.current?.(direction) ?? null, @@ -120,6 +123,16 @@ export function PiSessionView({ taskId, taskRunId }: PiSessionViewProps) { ]); }, [draftActions, session?.commands, taskId]); + const handleControllerError = useCallback( + (error: unknown, fallback: string) => { + if (error instanceof PiOperationError) { + return; + } + toast.error(fallback); + }, + [], + ); + const sendPrompt = useCallback( (text: string) => { const message = text.trim(); @@ -139,33 +152,43 @@ export function PiSessionView({ taskId, taskRunId }: PiSessionViewProps) { toast.success("Pi context compacted"); } }) - .catch(() => { + .catch((error) => { const failureMessage = action === "compact" ? "Failed to compact Pi context" : "Failed to send message to Pi"; - toast.error(failureMessage); + handleControllerError(error, failureMessage); }); }, - [isStreaming, messagingMode, piSessionController, taskId], + [ + handleControllerError, + isStreaming, + messagingMode, + piSessionController, + taskId, + ], ); const setModel = useCallback( (model: PiModelSelection) => { void piSessionController .setModel(taskId, model) - .catch(() => toast.error("Failed to change Pi model")); + .catch((error) => + handleControllerError(error, "Failed to change Pi model"), + ); }, - [piSessionController, taskId], + [handleControllerError, piSessionController, taskId], ); const setThinkingLevel = useCallback( (level: PiThinkingLevel) => { void piSessionController .setThinkingLevel(taskId, level) - .catch(() => toast.error("Failed to change Pi thinking level")); + .catch((error) => + handleControllerError(error, "Failed to change Pi thinking level"), + ); }, - [piSessionController, taskId], + [handleControllerError, piSessionController, taskId], ); const toggleMessagingMode = useCallback(() => { @@ -176,29 +199,37 @@ export function PiSessionView({ taskId, taskRunId }: PiSessionViewProps) { const runBashCommand = (command: string) => { void piSessionController .bash(taskId, command) - .catch(() => toast.error("Failed to run Pi bash command")); + .catch((error) => + handleControllerError(error, "Failed to run Pi bash command"), + ); }; const cancelPrompt = () => { if (isBashRunning) { - void piSessionController.abortBash(taskId); + void piSessionController + .abortBash(taskId) + .catch((error) => handleControllerError(error, "Failed to stop bash")); return; } - void piSessionController.abort(taskId); + void piSessionController + .abort(taskId) + .catch((error) => handleControllerError(error, "Failed to stop Pi")); }; const retry = useCallback(() => { void piSessionController .retry(taskId) - .catch(() => toast.error("Failed to reconnect to Pi")); - }, [piSessionController, taskId]); + .catch((error) => + handleControllerError(error, "Failed to reconnect to Pi"), + ); + }, [handleControllerError, piSessionController, taskId]); const restart = useCallback(() => { void piSessionController .restart(taskId) - .catch(() => toast.error("Failed to restart Pi")); - }, [piSessionController, taskId]); + .catch((error) => handleControllerError(error, "Failed to restart Pi")); + }, [handleControllerError, piSessionController, taskId]); const editQueuedMessage = useCallback(() => { void piSessionController @@ -217,14 +248,46 @@ export function PiSessionView({ taskId, taskRunId }: PiSessionViewProps) { draftActions.setPendingContent(taskId, xmlToContent(content)); draftActions.requestFocus(taskId); }) - .catch(() => toast.error("Failed to edit queued Pi message")); - }, [draftActions, piSessionController, taskId]); + .catch((error) => + handleControllerError(error, "Failed to edit queued Pi message"), + ); + }, [draftActions, handleControllerError, piSessionController, taskId]); const removeQueuedMessage = useCallback(() => { void piSessionController .clearQueue(taskId) - .catch(() => toast.error("Failed to discard queued Pi message")); - }, [piSessionController, taskId]); + .catch((error) => + handleControllerError(error, "Failed to discard queued Pi message"), + ); + }, [handleControllerError, piSessionController, taskId]); + + useEffect(() => { + const failure = session?.error; + if (!failure || failure.scope !== "operation") { + return; + } + if (failure.recoveryPrompt) { + draftActions.setPendingContent( + taskId, + xmlToContent(failure.recoveryPrompt), + ); + draftActions.requestFocus(taskId); + } + if (failure.kind === "usage_limit") { + showUsageLimit( + failure.limitCause ? { cause: failure.limitCause } : undefined, + ); + } else { + toast.error(failure.title, { description: failure.message }); + } + piSessionController.acknowledgeOperationFailure(taskId, failure.id); + }, [ + draftActions, + piSessionController, + session?.error, + showUsageLimit, + taskId, + ]); if (!session) { return ; @@ -235,6 +298,9 @@ export function PiSessionView({ taskId, taskRunId }: PiSessionViewProps) { event.type === "progress" && event.status === "in_progress", ); const isConnecting = session.connectionState === "connecting"; + const isAuthRestoring = session.authRestoring; + const connectionError = + session.error?.scope === "connection" ? session.error : undefined; const hasTranscript = session.events.some( (event) => event.type !== "progress", ); @@ -250,15 +316,15 @@ export function PiSessionView({ taskId, taskRunId }: PiSessionViewProps) { ); } - if (session.errorMessage && !hasTranscript) { + if (connectionError && !hasTranscript) { return ( - {session.errorTitle ?? "Pi session failed"} - {session.errorMessage} + {connectionError.title} + {connectionError.message} - {session.errorRetryable && ( + {connectionError.retryable && ( @@ -322,14 +388,17 @@ export function PiSessionView({ taskId, taskRunId }: PiSessionViewProps) { return ( - {isConnecting && hasTranscript && ( + {isAuthRestoring && ( + + )} + {isConnecting && hasTranscript && !isAuthRestoring && ( )} - {session.errorMessage && hasTranscript && ( + {connectionError && hasTranscript && ( )} @@ -359,14 +428,20 @@ export function PiSessionView({ taskId, taskRunId }: PiSessionViewProps) { disabled={isCompacting} isLoading={controlsPending} submitDisabledExternal={ - !sessionAvailable || !status || !isOnline || hasQueuedMessage + !sessionAvailable || + !status || + !isOnline || + hasQueuedMessage || + isAuthRestoring } submitTooltipOverride={ !isOnline ? "No internet connection" - : hasQueuedMessage - ? "A message is already queued" - : undefined + : isAuthRestoring + ? "Restoring authentication" + : hasQueuedMessage + ? "A message is already queued" + : undefined } enableBashMode enableCommands From d18f3b25ae38def9c579e7aeb8df0c5c5433300e Mon Sep 17 00:00:00 2001 From: JonathanLab Date: Tue, 28 Jul 2026 10:02:15 +0200 Subject: [PATCH 17/45] feat(pi): show native session usage --- packages/agent/src/pi/remote-rpc-client.ts | 5 ++ packages/agent/src/pi/rpc-transport.test.ts | 14 +++++ .../pi-runtime/cloudPiSessionClient.test.ts | 2 +- .../src/pi-runtime/cloudPiSessionClient.ts | 1 + .../pi-runtime/piSessionController.test.ts | 48 ++++++++++++++++ .../src/pi-runtime/piSessionController.ts | 23 +++++++- .../core/src/pi-runtime/piSessionStore.ts | 2 + .../src/pi-runtime/piSessionUsage.test.ts | 55 +++++++++++++++++++ .../core/src/pi-runtime/piSessionUsage.ts | 25 +++++++++ packages/core/src/sessions/contextUsage.ts | 1 + packages/harness/src/runtime.test.ts | 30 ++++++++++ packages/harness/src/runtime.ts | 8 ++- .../features/pi-sessions/PiSessionView.tsx | 3 + .../ContextBreakdownPopover.test.tsx | 11 ++++ .../components/ContextBreakdownPopover.tsx | 6 +- .../components/chat-thread/ChatThread.tsx | 4 ++ .../chat-thread/ChatThreadFooter.tsx | 6 +- 17 files changed, 238 insertions(+), 6 deletions(-) create mode 100644 packages/core/src/pi-runtime/piSessionUsage.test.ts create mode 100644 packages/core/src/pi-runtime/piSessionUsage.ts diff --git a/packages/agent/src/pi/remote-rpc-client.ts b/packages/agent/src/pi/remote-rpc-client.ts index 1a284759ff..3a576b292d 100644 --- a/packages/agent/src/pi/remote-rpc-client.ts +++ b/packages/agent/src/pi/remote-rpc-client.ts @@ -14,6 +14,7 @@ export type PiRemoteRpcClient = Pick< | "followUp" | "abort" | "getState" + | "getSessionStats" | "setModel" | "getAvailableModels" | "getAvailableThinkingLevels" @@ -79,6 +80,10 @@ export class RemotePiRpcClient implements PiRemoteRpcClient { return this.data(await this.request({ type: "get_state" })); } + async getSessionStats(): ReturnType { + return this.data(await this.request({ type: "get_session_stats" })); + } + async setModel( provider: string, modelId: string, diff --git a/packages/agent/src/pi/rpc-transport.test.ts b/packages/agent/src/pi/rpc-transport.test.ts index ec28950723..5ea9b03c03 100644 --- a/packages/agent/src/pi/rpc-transport.test.ts +++ b/packages/agent/src/pi/rpc-transport.test.ts @@ -25,12 +25,21 @@ describe("RemotePiRpcClient", () => { if (command.type === "get_available_thinking_levels") { return response(command, { levels: ["off", "high", "xhigh"] }); } + if (command.type === "get_session_stats") { + return response(command, { + sessionId: "session-1", + totalMessages: 2, + tokens: { total: 120 }, + cost: 0.01, + }); + } return response(command); }); const client = new RemotePiRpcClient({ request }); const compaction = await client.compact("retain decisions"); const thinkingLevels = await client.getAvailableThinkingLevels(); + const stats = await client.getSessionStats(); expect(request).toHaveBeenNthCalledWith(1, { id: expect.any(String), @@ -41,8 +50,13 @@ describe("RemotePiRpcClient", () => { id: expect.any(String), type: "get_available_thinking_levels", }); + expect(request).toHaveBeenNthCalledWith(3, { + id: expect.any(String), + type: "get_session_stats", + }); expect(compaction.summary).toBe("summary"); expect(thinkingLevels).toEqual(["off", "high", "xhigh"]); + expect(stats).toMatchObject({ tokens: { total: 120 }, cost: 0.01 }); }); it("rejects malformed responses from every transport", async () => { diff --git a/packages/core/src/pi-runtime/cloudPiSessionClient.test.ts b/packages/core/src/pi-runtime/cloudPiSessionClient.test.ts index 28b2ff8241..819390a0f6 100644 --- a/packages/core/src/pi-runtime/cloudPiSessionClient.test.ts +++ b/packages/core/src/pi-runtime/cloudPiSessionClient.test.ts @@ -330,7 +330,7 @@ describe("CloudPiSessionClient", () => { totalEntryCount: 1, }); await vi.waitFor(() => { - expect(cloud.client.sendCommand).toHaveBeenCalledTimes(2); + expect(cloud.client.sendCommand).toHaveBeenCalledTimes(3); }); cloud.sendUpdate({ taskId: "task-1", diff --git a/packages/core/src/pi-runtime/cloudPiSessionClient.ts b/packages/core/src/pi-runtime/cloudPiSessionClient.ts index 0466ddfe0f..e1a473a79b 100644 --- a/packages/core/src/pi-runtime/cloudPiSessionClient.ts +++ b/packages/core/src/pi-runtime/cloudPiSessionClient.ts @@ -42,6 +42,7 @@ function createTerminalPiRpcClient( messageCount: 0, pendingMessageCount: 0, }), + getSessionStats: rejectCommand, setModel: rejectCommand, getAvailableModels: async () => [], getAvailableThinkingLevels: async () => [], diff --git a/packages/core/src/pi-runtime/piSessionController.test.ts b/packages/core/src/pi-runtime/piSessionController.test.ts index 807dc59d08..8a42ab2868 100644 --- a/packages/core/src/pi-runtime/piSessionController.test.ts +++ b/packages/core/src/pi-runtime/piSessionController.test.ts @@ -36,6 +36,24 @@ function createSession(): PiSession { messageCount: 0, pendingMessageCount: 0, })), + getSessionStats: vi.fn(async () => ({ + sessionFile: undefined, + sessionId: "session-1", + userMessages: 0, + assistantMessages: 0, + toolCalls: 0, + toolResults: 0, + totalMessages: 0, + tokens: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + total: 0, + }, + cost: 0, + contextUsage: undefined, + })), getAvailableModels: vi.fn(async () => []), getAvailableThinkingLevels: vi.fn(async () => ["off" as const]), getCommands: vi.fn(async () => []), @@ -966,8 +984,38 @@ describe("PiSessionController", () => { const controller = createController(session); await controller.connect("task-1"); + vi.mocked(session.client.getSessionStats).mockResolvedValueOnce({ + sessionFile: undefined, + sessionId: "session-1", + userMessages: 1, + assistantMessages: 1, + toolCalls: 0, + toolResults: 0, + totalMessages: 2, + tokens: { + input: 1_000, + output: 500, + cacheRead: 0, + cacheWrite: 0, + total: 1_500, + }, + cost: 0.03, + contextUsage: { + tokens: 12_000, + contextWindow: 100_000, + percent: 12, + }, + }); onEvent(turnCompleted); + await vi.waitFor(() => { + expect( + controller.store.getState().sessions["task-1"].stats, + ).toMatchObject({ + cost: 0.03, + contextUsage: { tokens: 12_000, contextWindow: 100_000 }, + }); + }); expect(session.getConversation).toHaveBeenCalledOnce(); expect(controller.store.getState().sessions["task-1"].events).toEqual([ turnCompleted, diff --git a/packages/core/src/pi-runtime/piSessionController.ts b/packages/core/src/pi-runtime/piSessionController.ts index d253a6bbb4..b63465c5d0 100644 --- a/packages/core/src/pi-runtime/piSessionController.ts +++ b/packages/core/src/pi-runtime/piSessionController.ts @@ -415,6 +415,7 @@ export class PiSessionController { const session = await this.getPiSession(taskId); await session.client.setModel(model.provider, model.id); await this.refreshStatus(taskId); + await this.refreshStats(taskId); const thinkingLevels = await session.client.getAvailableThinkingLevels(); this.updateSession(taskId, { thinkingLevels, @@ -523,10 +524,12 @@ export class PiSessionController { try { const session = await this.getPiSession(taskId); const queueRevision = this.queueRevisions.get(taskId) ?? 0; - const [events, status, queue] = await Promise.all([ + const retainedStats = this.getSession(taskId).stats; + const [events, status, queue, stats] = await Promise.all([ session.getConversation(), session.client.getState(), session.getQueue(), + session.client.getSessionStats().catch(() => retainedStats), ]); if (this.getSessionVersion(taskId) !== connectedSessionVersion) { return; @@ -572,6 +575,7 @@ export class PiSessionController { connectionState: "connected", events: reconciledEvents, status: resolvedStatus, + stats, models: currentSession.models, modelsLoaded: currentSession.modelsLoaded, thinkingLevels: currentSession.thinkingLevels, @@ -694,6 +698,23 @@ export class PiSessionController { ? latestSession.error : undefined, }); + + if (event.type === "turn_completed") { + void this.refreshStats(taskId); + } + } + + private async refreshStats(taskId: string): Promise { + const sessionVersion = this.getSessionVersion(taskId); + try { + const session = await this.getPiSession(taskId); + const stats = await session.client.getSessionStats(); + if (this.getSessionVersion(taskId) === sessionVersion) { + this.updateSession(taskId, { stats }); + } + } catch { + return; + } } private reconcileLiveEvents( diff --git a/packages/core/src/pi-runtime/piSessionStore.ts b/packages/core/src/pi-runtime/piSessionStore.ts index 4453366a05..73a9073801 100644 --- a/packages/core/src/pi-runtime/piSessionStore.ts +++ b/packages/core/src/pi-runtime/piSessionStore.ts @@ -2,6 +2,7 @@ import type { PiCommand, PiNativeModelInfo, PiQueueSnapshot, + PiSessionStats, PiSessionStatus, PiThinkingLevel, } from "@posthog/agent/pi/types"; @@ -35,6 +36,7 @@ export interface PiControllerSessionState { commands: PiCommand[]; queue: PiQueueSnapshot; status?: PiSessionStatus; + stats?: PiSessionStats; cloudStatus?: TaskRunStatus; error?: PiSessionError; authRestoring: boolean; diff --git a/packages/core/src/pi-runtime/piSessionUsage.test.ts b/packages/core/src/pi-runtime/piSessionUsage.test.ts new file mode 100644 index 0000000000..d25f30bc4f --- /dev/null +++ b/packages/core/src/pi-runtime/piSessionUsage.test.ts @@ -0,0 +1,55 @@ +import type { PiSessionStats } from "@posthog/agent/pi/types"; +import { describe, expect, it } from "vitest"; +import { toPiContextUsage } from "./piSessionUsage"; + +function stats( + contextUsage: PiSessionStats["contextUsage"], + cost = 0, +): PiSessionStats { + return { + sessionFile: undefined, + sessionId: "session-1", + userMessages: 1, + assistantMessages: 1, + toolCalls: 0, + toolResults: 0, + totalMessages: 2, + tokens: { + input: 1_000, + output: 500, + cacheRead: 100, + cacheWrite: 50, + total: 1_650, + }, + cost, + contextUsage, + }; +} + +describe("toPiContextUsage", () => { + it("maps native Pi context and cost statistics to shared context usage", () => { + expect( + toPiContextUsage( + stats( + { tokens: 38_323, contextWindow: 1_000_000, percent: 3.8323 }, + 0.42, + ), + ), + ).toEqual({ + used: 38_323, + size: 1_000_000, + percentage: 4, + cost: { amount: 0.42, currency: "USD" }, + breakdown: null, + breakdownAvailable: false, + }); + }); + + it("hides context usage while Pi cannot estimate it", () => { + expect( + toPiContextUsage( + stats({ tokens: null, contextWindow: 100_000, percent: null }), + ), + ).toBeNull(); + }); +}); diff --git a/packages/core/src/pi-runtime/piSessionUsage.ts b/packages/core/src/pi-runtime/piSessionUsage.ts new file mode 100644 index 0000000000..87036a9bbf --- /dev/null +++ b/packages/core/src/pi-runtime/piSessionUsage.ts @@ -0,0 +1,25 @@ +import type { PiSessionStats } from "@posthog/agent/pi/types"; +import type { ContextUsage } from "../sessions/contextUsage"; + +export function toPiContextUsage( + stats: PiSessionStats | undefined, +): ContextUsage | null { + const usage = stats?.contextUsage; + if (!usage || usage.tokens === null) { + return null; + } + + return { + used: usage.tokens, + size: usage.contextWindow, + percentage: Math.round( + usage.percent ?? + (usage.contextWindow > 0 + ? (usage.tokens / usage.contextWindow) * 100 + : 0), + ), + cost: stats.cost > 0 ? { amount: stats.cost, currency: "USD" } : null, + breakdown: null, + breakdownAvailable: false, + }; +} diff --git a/packages/core/src/sessions/contextUsage.ts b/packages/core/src/sessions/contextUsage.ts index df3f4001e3..97d5d9ed11 100644 --- a/packages/core/src/sessions/contextUsage.ts +++ b/packages/core/src/sessions/contextUsage.ts @@ -18,6 +18,7 @@ export interface ContextUsage { /** Cumulative estimated session cost, summed across turns; `null` if none reported (e.g. codex). */ cost: { amount: number; currency: string } | null; breakdown: ContextBreakdown | null; + breakdownAvailable?: boolean; } type ContextUsageAggregate = Omit; diff --git a/packages/harness/src/runtime.test.ts b/packages/harness/src/runtime.test.ts index 375d3f7f0a..c7ab40634f 100644 --- a/packages/harness/src/runtime.test.ts +++ b/packages/harness/src/runtime.test.ts @@ -65,6 +65,36 @@ describe("createHarnessRuntime", () => { }, ); + it("restores the session model before calculating context usage", async () => { + vi.stubEnv("PI_OFFLINE", "1"); + const pi = await import("@earendil-works/pi-coding-agent"); + const cwd = await temporaryDirectory(); + const agentDir = await temporaryDirectory(); + const sessionManager = pi.SessionManager.inMemory(cwd); + sessionManager.appendModelChange("posthog", "claude-haiku-4-5"); + sessionManager.appendMessage({ + role: "user", + content: "continue", + timestamp: Date.now(), + }); + + const runtime = await createHarnessRuntime({ + agentDir, + apiKey: "proxy-key", + cwd, + sessionManager, + }); + + try { + expect(runtime.session.model?.id).toBe("claude-haiku-4-5"); + expect(runtime.session.getSessionStats().contextUsage).toMatchObject({ + contextWindow: 200_000, + }); + } finally { + await runtime.dispose(); + } + }); + it("keeps desktop-provided OAuth credentials in memory without touching auth.json", async () => { vi.stubEnv("PI_OFFLINE", "1"); const pi = await import("@earendil-works/pi-coding-agent"); diff --git a/packages/harness/src/runtime.ts b/packages/harness/src/runtime.ts index 8692a42b3f..906b53b616 100644 --- a/packages/harness/src/runtime.ts +++ b/packages/harness/src/runtime.ts @@ -137,13 +137,19 @@ export async function createHarnessRuntime( const fallbackModel = services.modelRuntime .getModels(POSTHOG_PROVIDER_NAME) .at(0); + const existingSession = sessionManager.buildSessionContext(); + const hasRestorableModel = + existingSession.messages.length > 0 && existingSession.model !== null; + const defaultModel = hasRestorableModel + ? undefined + : (preferredModel ?? fallbackModel); const created = await pi.createAgentSessionFromServices({ ...runtimeOptions, services, sessionManager, sessionStartEvent, - model: runtimeOptions.model ?? preferredModel ?? fallbackModel, + model: runtimeOptions.model ?? defaultModel, }); return { diff --git a/packages/ui/src/features/pi-sessions/PiSessionView.tsx b/packages/ui/src/features/pi-sessions/PiSessionView.tsx index f85e9be648..e2b42638a2 100644 --- a/packages/ui/src/features/pi-sessions/PiSessionView.tsx +++ b/packages/ui/src/features/pi-sessions/PiSessionView.tsx @@ -9,6 +9,7 @@ import { type PiSessionController, type PiThinkingLevel, } from "@posthog/core/pi-runtime/piSessionController"; +import { toPiContextUsage } from "@posthog/core/pi-runtime/piSessionUsage"; import { useService } from "@posthog/di/react"; import { Button, @@ -301,6 +302,7 @@ export function PiSessionView({ taskId, taskRunId }: PiSessionViewProps) { const isAuthRestoring = session.authRestoring; const connectionError = session.error?.scope === "connection" ? session.error : undefined; + const contextUsage = toPiContextUsage(session.stats); const hasTranscript = session.events.some( (event) => event.type !== "progress", ); @@ -408,6 +410,7 @@ export function PiSessionView({ taskId, taskRunId }: PiSessionViewProps) { isPromptPending={isStreaming} taskId={taskId} repoPath={repoPath} + usage={contextUsage} promptRecallRef={promptRecallRef} /> diff --git a/packages/ui/src/features/sessions/components/ContextBreakdownPopover.test.tsx b/packages/ui/src/features/sessions/components/ContextBreakdownPopover.test.tsx index 2cc33beb79..cf30e422f2 100644 --- a/packages/ui/src/features/sessions/components/ContextBreakdownPopover.test.tsx +++ b/packages/ui/src/features/sessions/components/ContextBreakdownPopover.test.tsx @@ -54,6 +54,17 @@ describe("ContextBreakdownPopover", () => { ).toBeInTheDocument(); }); + it("does not mention a breakdown when the runtime cannot provide one", () => { + render( + + + , + ); + expect(screen.queryByText(/breakdown/i)).not.toBeInTheDocument(); + }); + it("renders one row per non-zero category", () => { render( diff --git a/packages/ui/src/features/sessions/components/ContextBreakdownPopover.tsx b/packages/ui/src/features/sessions/components/ContextBreakdownPopover.tsx index 4269e1504f..2379254317 100644 --- a/packages/ui/src/features/sessions/components/ContextBreakdownPopover.tsx +++ b/packages/ui/src/features/sessions/components/ContextBreakdownPopover.tsx @@ -47,7 +47,7 @@ export function ContextBreakdownPopover({ )} - {breakdown ? ( + {breakdown && ( {CONTEXT_CATEGORIES.filter((c) => breakdown[c.key] > 0).map((cat) => ( ))} - ) : ( + )} + + {!breakdown && usage.breakdownAvailable !== false && ( Detailed breakdown available after the first response. diff --git a/packages/ui/src/features/sessions/components/chat-thread/ChatThread.tsx b/packages/ui/src/features/sessions/components/chat-thread/ChatThread.tsx index 1365e714ff..637756b780 100644 --- a/packages/ui/src/features/sessions/components/chat-thread/ChatThread.tsx +++ b/packages/ui/src/features/sessions/components/chat-thread/ChatThread.tsx @@ -7,6 +7,7 @@ import { Scroll, } from "@phosphor-icons/react"; import { WorkerPoolContextProvider } from "@pierre/diffs/react"; +import type { ContextUsage } from "@posthog/core/sessions/contextUsage"; import { useService } from "@posthog/di/react"; import { Button, @@ -1099,6 +1100,7 @@ interface SharedChatThreadProps { repoPath?: string | null; task?: Task; taskId?: string; + usage?: ContextUsage | null; } export interface ChatThreadProps extends SharedChatThreadProps { @@ -1149,6 +1151,7 @@ function ChatThreadRenderer({ repoPath, task, taskId, + usage, promptRecallRef, }: ChatThreadRendererProps) { const diffWorkerFactory = useService(DIFF_WORKER_FACTORY); @@ -1275,6 +1278,7 @@ function ChatThreadRenderer({ promptStartedAt={promptStartedAt} task={task} taskId={taskId} + usage={usage} /> } /> diff --git a/packages/ui/src/features/sessions/components/chat-thread/ChatThreadFooter.tsx b/packages/ui/src/features/sessions/components/chat-thread/ChatThreadFooter.tsx index cabb5777e3..07143d07d3 100644 --- a/packages/ui/src/features/sessions/components/chat-thread/ChatThreadFooter.tsx +++ b/packages/ui/src/features/sessions/components/chat-thread/ChatThreadFooter.tsx @@ -1,3 +1,4 @@ +import type { ContextUsage } from "@posthog/core/sessions/contextUsage"; import type { AcpMessage } from "@posthog/shared"; import type { Task } from "@posthog/shared/domain-types"; import { SessionFooter } from "@posthog/ui/features/sessions/components/SessionFooter"; @@ -16,6 +17,7 @@ interface ChatThreadFooterProps { promptStartedAt?: number | null; task?: Task; taskId?: string; + usage?: ContextUsage | null; } /** @@ -34,9 +36,11 @@ export function ChatThreadFooter({ promptStartedAt, task, taskId, + usage, }: ChatThreadFooterProps) { const showDebugLogs = useSettingsStore((s) => s.debugLogsCloudRuns); - const contextUsage = useContextUsage(events); + const eventContextUsage = useContextUsage(events); + const contextUsage = usage === undefined ? eventContextUsage : usage; const { lastTurnInfo, isCompacting, completedToolCallCount } = useConversationItems(events, isPromptPending, { showDebugLogs }); const pendingPermissions = usePendingPermissionsForTask(taskId ?? ""); From 7a3e0239c6b1bfce72e72cd5bfb3c6a7f202e807 Mon Sep 17 00:00:00 2001 From: JonathanLab Date: Tue, 28 Jul 2026 10:31:13 +0200 Subject: [PATCH 18/45] fix(pi): disable cloud shell RPC --- .../src/pi-runtime/piSessionController.test.ts | 14 ++++++++++++++ .../core/src/pi-runtime/piSessionController.ts | 1 + .../ui/src/features/pi-sessions/PiSessionView.tsx | 2 +- 3 files changed, 16 insertions(+), 1 deletion(-) diff --git a/packages/core/src/pi-runtime/piSessionController.test.ts b/packages/core/src/pi-runtime/piSessionController.test.ts index 8a42ab2868..93b58b276c 100644 --- a/packages/core/src/pi-runtime/piSessionController.test.ts +++ b/packages/core/src/pi-runtime/piSessionController.test.ts @@ -1022,6 +1022,20 @@ describe("PiSessionController", () => { ]); }); + it("retains cloud status when loading a cloud session", async () => { + const session = { + ...createSession(), + cloudStatus: "in_progress" as const, + }; + const controller = createController(session); + + await controller.connect("task-1", "run-1"); + + expect(controller.store.getState().sessions["task-1"].cloudStatus).toBe( + "in_progress", + ); + }); + it("loads session state and appends normalized runtime events", async () => { const initialEvent: AgentConversationEvent = { type: "assistant_message_chunk", diff --git a/packages/core/src/pi-runtime/piSessionController.ts b/packages/core/src/pi-runtime/piSessionController.ts index b63465c5d0..95663fddaa 100644 --- a/packages/core/src/pi-runtime/piSessionController.ts +++ b/packages/core/src/pi-runtime/piSessionController.ts @@ -576,6 +576,7 @@ export class PiSessionController { events: reconciledEvents, status: resolvedStatus, stats, + cloudStatus: session.cloudStatus, models: currentSession.models, modelsLoaded: currentSession.modelsLoaded, thinkingLevels: currentSession.thinkingLevels, diff --git a/packages/ui/src/features/pi-sessions/PiSessionView.tsx b/packages/ui/src/features/pi-sessions/PiSessionView.tsx index e2b42638a2..e8550446fa 100644 --- a/packages/ui/src/features/pi-sessions/PiSessionView.tsx +++ b/packages/ui/src/features/pi-sessions/PiSessionView.tsx @@ -446,7 +446,7 @@ export function PiSessionView({ taskId, taskRunId }: PiSessionViewProps) { ? "A message is already queued" : undefined } - enableBashMode + enableBashMode={session.cloudStatus === undefined} enableCommands modelSelector={modelSelector} reasoningSelector={reasoningSelector} From 484a6c7b88e6578dbe5dbdd74b6cc1753900a985 Mon Sep 17 00:00:00 2001 From: Richard Solomou Date: Fri, 24 Jul 2026 19:17:22 +0300 Subject: [PATCH 19/45] refactor(shared): extract task contracts and model policy Generated-By: PostHog Code Task-Id: c1bbe3cf-742b-4b24-bf96-d11a18b4cf22 --- .../agent/src/adapters/reasoning-effort.ts | 45 +- packages/agent/src/gateway-models.test.ts | 252 +----------- packages/agent/src/gateway-models.ts | 274 ++----------- packages/agent/src/utils/gateway.ts | 23 +- packages/api-client/src/posthog-client.ts | 2 +- packages/shared/src/cloud-task-models.test.ts | 200 +++++++++ packages/shared/src/cloud-task-models.ts | 385 ++++++++++++++++++ packages/shared/src/domain-types.test.ts | 19 +- packages/shared/src/domain-types.ts | 34 +- packages/shared/src/index.ts | 74 +++- packages/shared/src/reasoning-effort.test.ts | 20 + packages/shared/src/reasoning-effort.ts | 77 ++++ packages/shared/src/sessions.ts | 2 +- packages/shared/src/task-automation.test.ts | 66 +++ packages/shared/src/task-automation.ts | 58 +++ packages/shared/src/task.test.ts | 53 +++ packages/shared/src/task.ts | 103 +---- .../src/services/agent/agent.ts | 117 +----- 18 files changed, 1035 insertions(+), 769 deletions(-) create mode 100644 packages/shared/src/cloud-task-models.test.ts create mode 100644 packages/shared/src/cloud-task-models.ts create mode 100644 packages/shared/src/reasoning-effort.test.ts create mode 100644 packages/shared/src/reasoning-effort.ts create mode 100644 packages/shared/src/task-automation.test.ts create mode 100644 packages/shared/src/task-automation.ts create mode 100644 packages/shared/src/task.test.ts diff --git a/packages/agent/src/adapters/reasoning-effort.ts b/packages/agent/src/adapters/reasoning-effort.ts index 590bb3be86..2cc1da50f8 100644 --- a/packages/agent/src/adapters/reasoning-effort.ts +++ b/packages/agent/src/adapters/reasoning-effort.ts @@ -1,39 +1,6 @@ -import type { Adapter } from "@posthog/shared"; -import { getEffortOptions as getClaudeEffortOptions } from "./claude/session/models"; -import { getReasoningEffortOptions as getCodexReasoningEffortOptions } from "./codex-app-server/models"; - -export type SupportedReasoningEffort = - | "low" - | "medium" - | "high" - | "xhigh" - | "max"; - -export interface ReasoningEffortOption { - value: SupportedReasoningEffort; - name: string; -} - -export function getReasoningEffortOptions( - adapter: Adapter, - modelId: string, -): ReasoningEffortOption[] | null { - const options = - adapter === "codex" - ? getCodexReasoningEffortOptions(modelId) - : getClaudeEffortOptions(modelId); - - return options as ReasoningEffortOption[] | null; -} - -export function isSupportedReasoningEffort( - adapter: Adapter, - modelId: string, - value: string, -): value is SupportedReasoningEffort { - return ( - getReasoningEffortOptions(adapter, modelId)?.some( - (option) => option.value === value, - ) ?? false - ); -} +export { + getReasoningEffortOptions, + isSupportedReasoningEffort, + type ReasoningEffortOption, + type SupportedReasoningEffort, +} from "@posthog/shared"; diff --git a/packages/agent/src/gateway-models.test.ts b/packages/agent/src/gateway-models.test.ts index 5aa7438a8d..8c988ec1e9 100644 --- a/packages/agent/src/gateway-models.test.ts +++ b/packages/agent/src/gateway-models.test.ts @@ -1,163 +1,5 @@ import { afterEach, describe, expect, it, vi } from "vitest"; -import { - compareModelsForPicker, - fetchGatewayModels, - fetchModelsList, - formatGatewayModelName, - type GatewayModel, - getClaudeModelRecency, - isAnthropicModel, - isBlockedModelId, - isCloudflareModel, - pickAllowedModel, -} from "./gateway-models"; - -const model = (id: string, owned_by = ""): GatewayModel => ({ - id, - owned_by, - context_window: 128000, - supports_streaming: true, - supports_vision: false, - allowed: true, -}); - -describe("formatGatewayModelName", () => { - it("keeps Claude models in friendly title case", () => { - expect( - formatGatewayModelName({ - id: "claude-opus-4-8", - owned_by: "anthropic", - context_window: 200000, - supports_streaming: true, - supports_vision: true, - allowed: true, - }), - ).toBe("Claude Opus 4.8"); - }); - - it("uppercases the GPT acronym in OpenAI model ids", () => { - expect( - formatGatewayModelName({ - id: "GPT-5.5", - owned_by: "openai", - context_window: 200000, - supports_streaming: true, - supports_vision: true, - allowed: true, - }), - ).toBe("GPT-5.5"); - }); - - it("strips the openai/ prefix, uppercases GPT, and title-cases the suffix", () => { - expect( - formatGatewayModelName({ - id: "openai/gpt-5.6-sol", - owned_by: "openai", - context_window: 200000, - supports_streaming: true, - supports_vision: true, - allowed: true, - }), - ).toBe("GPT-5.6 Sol"); - }); - - it("formats Cloudflare models as the final path segment with GLM uppercased", () => { - expect( - formatGatewayModelName({ - id: "@cf/zai-org/glm-5.2", - owned_by: "cloudflare", - context_window: 128000, - supports_streaming: true, - supports_vision: false, - allowed: true, - }), - ).toBe("GLM-5.2"); - }); - - it("leaves non-acronym Cloudflare models lowercase", () => { - expect( - formatGatewayModelName({ - id: "@cf/meta/llama-3.1-8b-instruct", - owned_by: "cloudflare", - context_window: 128000, - supports_streaming: true, - supports_vision: false, - allowed: true, - }), - ).toBe("llama-3.1-8b-instruct"); - }); - - it("blocks deprecated Claude gateway models", () => { - expect(isBlockedModelId("claude-opus-4-5")).toBe(true); - expect(isBlockedModelId("claude-opus-4-6")).toBe(true); - expect(isBlockedModelId("claude-sonnet-4-5")).toBe(true); - expect(isBlockedModelId("claude-haiku-4-5")).toBe(true); - expect(isBlockedModelId("ANTHROPIC/CLAUDE-HAIKU-4-5")).toBe(true); - }); - - it("blocks deprecated Codex gateway models", () => { - expect(isBlockedModelId("gpt-5.2")).toBe(true); - expect(isBlockedModelId("gpt-5.3")).toBe(true); - expect(isBlockedModelId("gpt-5.3-codex")).toBe(true); - expect(isBlockedModelId("openai/gpt-5.2")).toBe(true); - expect(isBlockedModelId("OPENAI/GPT-5.3")).toBe(true); - expect(isBlockedModelId("OPENAI/GPT-5.3-CODEX")).toBe(true); - }); -}); - -describe("getClaudeModelRecency", () => { - it.each([ - ["claude-haiku-4-5", 4005], - ["claude-sonnet-4-6", 4006], - ["claude-opus-4-7", 4007], - ["claude-opus-4-8", 4008], - ["claude-opus-5", 5000], - ["claude-sonnet-5", 5000], - ["claude-fable-5", 5000], - ])("ranks %s by its embedded version (%i)", (modelId, rank) => { - expect(getClaudeModelRecency(modelId)).toBe(rank); - }); - - it("ignores a trailing date suffix when reading the version", () => { - expect(getClaudeModelRecency("claude-haiku-4-5-20251001")).toBe(4005); - }); - - it("ranks a model with no recognisable version as newest", () => { - expect(getClaudeModelRecency("claude-mystery")).toBe( - Number.MAX_SAFE_INTEGER, - ); - expect(getClaudeModelRecency("claude-mystery")).toBeGreaterThan( - getClaudeModelRecency("claude-fable-5"), - ); - }); -}); - -describe("compareModelsForPicker", () => { - it("groups models by family least capable first, newest version first", () => { - // The picker opens upward, so least-capable-first DOM order puts the most - // capable family (Fable) nearest the trigger — the visual top of the menu. - // Models as the gateway might return them — arbitrary order. - const gatewayOrder = [ - "claude-fable-5", - "claude-opus-4-7", - "claude-mystery", - "claude-sonnet-5", - "claude-haiku-4-5", - "claude-sonnet-4-6", - "claude-opus-4-8", - ]; - const displayed = [...gatewayOrder].sort(compareModelsForPicker); - expect(displayed).toEqual([ - "claude-haiku-4-5", - "claude-sonnet-5", - "claude-sonnet-4-6", - "claude-opus-4-8", - "claude-opus-4-7", - "claude-fable-5", - "claude-mystery", - ]); - }); -}); +import { fetchGatewayModels, fetchModelsList } from "./gateway-models"; describe("gateway model fetch timeout", () => { afterEach(() => { @@ -242,96 +84,4 @@ describe("gateway models cache", () => { expect(fetchMock).toHaveBeenCalledTimes(1); expect(cached[0]?.allowed).toBe(false); }); - - it("corrects stale GLM 5.2 context-window metadata", async () => { - vi.spyOn(globalThis, "fetch").mockResolvedValue( - new Response( - JSON.stringify({ - object: "list", - data: [ - { - id: "@cf/zai-org/glm-5.2", - owned_by: "cloudflare", - context_window: 128_000, - supports_streaming: true, - supports_vision: false, - }, - ], - }), - { status: 200, headers: { "Content-Type": "application/json" } }, - ), - ); - - const models = await fetchGatewayModels({ - gatewayUrl: "https://gateway.glm-context-test", - }); - - expect(models[0]?.context_window).toBe(1_000_000); - }); -}); - -describe("isCloudflareModel", () => { - it.each([ - { id: "@cf/zai-org/glm-5.2", owned_by: "cloudflare", expected: true }, - { id: "claude-opus-4-8", owned_by: "anthropic", expected: false }, - { id: "@cf/zai-org/glm-5.2", owned_by: "", expected: true }, - { id: "gpt-5.5", owned_by: "", expected: false }, - // A Cloudflare-served model can report an upstream owner; the `@cf/` prefix still wins. - { id: "@cf/openai/gpt-oss", owned_by: "openai", expected: true }, - ])( - "isCloudflareModel($id, owned_by=$owned_by) → $expected", - ({ id, owned_by, expected }) => { - expect(isCloudflareModel(model(id, owned_by))).toBe(expected); - }, - ); - - it("does not classify Cloudflare models as Anthropic", () => { - // The Claude adapter accepts both, but they must stay distinguishable. - const glm = model("@cf/zai-org/glm-5.2", "cloudflare"); - expect(isCloudflareModel(glm)).toBe(true); - expect(isAnthropicModel(glm)).toBe(false); - }); -}); - -describe("pickAllowedModel", () => { - const entry = (id: string, allowed: boolean) => ({ id, allowed }); - - it.each([ - [ - "keeps an allowed preferred model", - [entry("claude-opus-4-8", true)], - "claude-opus-4-8", - "claude-opus-4-8", - ], - [ - "keeps a preferred model absent from the list", - [entry("claude-opus-4-8", true)], - "claude-sonnet-5", - "claude-sonnet-5", - ], - [ - "moves a restricted preferred model to the newest allowed one", - [ - entry("claude-opus-4-8", false), - entry("claude-sonnet-4-6", true), - entry("@cf/zai-org/glm-5.2", true), - ], - "claude-opus-4-8", - "@cf/zai-org/glm-5.2", - ], - [ - "keeps the preferred model when everything is restricted", - [entry("claude-opus-4-8", false)], - "claude-opus-4-8", - "claude-opus-4-8", - ], - [ - "keeps the preferred model when the list is empty", - [], - "claude-opus-4-8", - "claude-opus-4-8", - ], - ] as const)("%s", (_name, models, preferred, expected) => { - expect(pickAllowedModel(models, preferred)).toBe(expected); - }); }); diff --git a/packages/agent/src/gateway-models.ts b/packages/agent/src/gateway-models.ts index 116d9bd7d1..35ae10df5a 100644 --- a/packages/agent/src/gateway-models.ts +++ b/packages/agent/src/gateway-models.ts @@ -1,20 +1,28 @@ -export interface GatewayModel { - id: string; - owned_by: string; - context_window: number; - supports_streaming: boolean; - supports_vision: boolean; - // Free-tier model gate: authenticated fetches mark models outside the - // caller's plan `allowed: false`. Anonymous fetches and older gateways - // don't mark, so absence means allowed. - allowed: boolean; - restriction_reason?: string | null; -} - -interface GatewayModelsResponse { - object: "list"; - data: Array & { allowed?: boolean }>; -} +import { + type GatewayModel, + normalizeGatewayModelsResponse, +} from "@posthog/shared"; + +export { + BLOCKED_GATEWAY_MODEL_IDS, + buildCloudTaskConfigOptions, + type CloudTaskConfigOption, + type CloudTaskConfigSelectOption, + compareModelsForPicker, + DEFAULT_CODEX_MODEL, + DEFAULT_GATEWAY_MODEL, + formatGatewayModelName, + formatModelId, + type GatewayModel, + getClaudeModelRecency, + getProviderName, + isAnthropicModel, + isBlockedModelId, + isCloudflareModel, + isCloudflareModelId, + isOpenAIModel, + pickAllowedModel, +} from "@posthog/shared"; export interface FetchGatewayModelsOptions { gatewayUrl: string; @@ -22,47 +30,6 @@ export interface FetchGatewayModelsOptions { authToken?: string; } -export const DEFAULT_GATEWAY_MODEL = "claude-opus-4-8"; - -export const DEFAULT_CODEX_MODEL = "gpt-5.5"; - -const BLOCKED_MODELS = new Set([ - "gpt-5-mini", - "openai/gpt-5-mini", - "gpt-5.2", - "openai/gpt-5.2", - "gpt-5.3", - "openai/gpt-5.3", - "gpt-5.3-codex", - "openai/gpt-5.3-codex", - "claude-opus-4-5", - "anthropic/claude-opus-4-5", - "claude-opus-4-6", - "anthropic/claude-opus-4-6", - "claude-sonnet-4-5", - "anthropic/claude-sonnet-4-5", - "claude-haiku-4-5", - "anthropic/claude-haiku-4-5", -]); - -export function isBlockedModelId(modelId: string): boolean { - return BLOCKED_MODELS.has(modelId.toLowerCase()); -} - -interface ModelsListEntry { - id?: string; - owned_by?: string; - allowed?: boolean; - restriction_reason?: string | null; -} - -type ModelsListResponse = - | { - data?: ModelsListEntry[]; - models?: ModelsListEntry[]; - } - | ModelsListEntry[]; - const CACHE_TTL = 10 * 60 * 1000; // 10 minutes // Bound the gateway /v1/models request so a stalled connection cannot hold up @@ -71,10 +38,6 @@ const CACHE_TTL = 10 * 60 * 1000; // 10 minutes // the callers fall through to `return []`. const GATEWAY_FETCH_TIMEOUT_MS = 10_000; -const MODEL_CONTEXT_WINDOW_OVERRIDES: Readonly> = { - "@cf/zai-org/glm-5.2": 1_000_000, -}; - // Restriction marks are identity-scoped (free-tier marks are authed-only and // differ per org), so cache entries are keyed on the exact token — an org // switch in the same process must never be served the old org's marks. A @@ -125,17 +88,7 @@ export async function fetchGatewayModels( return []; } - const data = (await response.json()) as GatewayModelsResponse; - const models = (data.data ?? []) - .filter((m) => !isBlockedModelId(m.id)) - .map((m) => ({ - ...m, - context_window: Math.max( - m.context_window, - MODEL_CONTEXT_WINDOW_OVERRIDES[m.id] ?? 0, - ), - allowed: m.allowed !== false, - })); + const models = normalizeGatewayModelsResponse(await response.json()); gatewayModelsCache = { models, expiry: Date.now() + CACHE_TTL, @@ -148,36 +101,6 @@ export async function fetchGatewayModels( } } -export function isAnthropicModel(model: GatewayModel): boolean { - if (model.owned_by) { - return model.owned_by === "anthropic"; - } - return model.id.startsWith("claude-") || model.id.startsWith("anthropic/"); -} - -export function isOpenAIModel(model: GatewayModel): boolean { - if (model.owned_by) { - return model.owned_by === "openai"; - } - return model.id.startsWith("gpt-") || model.id.startsWith("openai/"); -} - -// Cloudflare Workers AI model ids carry the `@cf/` path prefix (e.g. `@cf/zai-org/glm-5.2`). Kept as -// a standalone id-only check so callers that only have a model id (not a full GatewayModel) — like the -// Claude adapter's desync guard — share one source of truth with `isCloudflareModel`. -export function isCloudflareModelId(modelId: string): boolean { - return modelId.startsWith("@cf/"); -} - -// Cloudflare Workers AI models (e.g. `@cf/zai-org/glm-5.2`). The gateway serves these over both its -// OpenAI and Anthropic-Messages surfaces (it translates the `@cf/` path), so the Claude adapter can -// drive them just like an Anthropic model. The `@cf/` path prefix is the structural, always-present -// signal, so honour it regardless of `owned_by` — a Cloudflare-served model can report an upstream -// owner (e.g. `@cf/openai/...` with `owned_by: "openai"`) and must still classify as Cloudflare. -export function isCloudflareModel(model: GatewayModel): boolean { - return isCloudflareModelId(model.id) || model.owned_by === "cloudflare"; -} - export interface ModelInfo { id: string; owned_by?: string; @@ -208,22 +131,14 @@ export async function fetchModelsList( if (!response.ok) { return []; } - const data = (await response.json()) as ModelsListResponse; - const models = Array.isArray(data) - ? data - : (data.data ?? data.models ?? []); - const results: ModelInfo[] = []; - for (const model of models) { - const id = model?.id ? String(model.id) : ""; - if (!id) continue; - if (isBlockedModelId(id)) continue; - results.push({ - id, - owned_by: model?.owned_by, - allowed: model?.allowed !== false, - restriction_reason: model?.restriction_reason ?? null, - }); - } + const results = normalizeGatewayModelsResponse(await response.json()).map( + (model) => ({ + id: model.id, + owned_by: model.owned_by || undefined, + allowed: model.allowed, + restriction_reason: model.restriction_reason, + }), + ); modelsListCache = { models: results, expiry: Date.now() + CACHE_TTL, @@ -235,124 +150,3 @@ export async function fetchModelsList( return []; } } - -/** - * The model a session should start on: the preferred id when present and - * allowed, else the newest allowed model — a free-tier org must not default - * onto a model that 403s its first message. Falls back to the preferred id - * when the list is empty (fetch failed) or nothing is allowed (all locked — - * the picker gate communicates that state better than a silent swap). - */ -export function pickAllowedModel( - models: ReadonlyArray>, - preferred: string, -): string { - if (models.length === 0) return preferred; - const preferredEntry = models.find((m) => m.id === preferred); - if (!preferredEntry || preferredEntry.allowed) return preferred; - const allowed = models.filter((m) => m.allowed); - if (allowed.length === 0) return preferred; - return allowed.reduce((best, candidate) => - getClaudeModelRecency(candidate.id) >= getClaudeModelRecency(best.id) - ? candidate - : best, - ).id; -} - -const PROVIDER_NAMES: Record = { - anthropic: "Anthropic", - openai: "OpenAI", - "google-vertex": "Gemini", -}; - -export function getProviderName(ownedBy: string): string { - return PROVIDER_NAMES[ownedBy] ?? ownedBy; -} - -// Version embedded in the model id, e.g. "claude-opus-4-8" -> 4008. Ids with no -// recognisable version rank newest. A trailing date suffix is ignored. -export function getClaudeModelRecency(modelId: string): number { - const match = modelId.toLowerCase().match(/-(\d+)(?:[-.](\d+))?/); - if (!match) return Number.MAX_SAFE_INTEGER; - const major = Number(match[1]); - const minor = match[2] ? Number(match[2]) : 0; - return major * 1000 + minor; -} - -// Families ordered least-capable first. The picker opens upward (side="top") -// from the composer, so items later in this list render nearer the trigger and -// read as the top of the menu — this puts the most-capable family (Fable) on -// top. Unknown families sort after all known ones. -const MODEL_FAMILY_ORDER = ["haiku", "sonnet", "opus", "fable"]; - -function getModelFamilyRank(modelId: string): number { - const id = modelId.toLowerCase(); - const index = MODEL_FAMILY_ORDER.findIndex((family) => id.includes(family)); - return index === -1 ? MODEL_FAMILY_ORDER.length : index; -} - -// Group by family, then newest version first within each family. -export function compareModelsForPicker(a: string, b: string): number { - const familyDiff = getModelFamilyRank(a) - getModelFamilyRank(b); - if (familyDiff !== 0) return familyDiff; - return getClaudeModelRecency(b) - getClaudeModelRecency(a); -} - -const PROVIDER_PREFIXES = ["anthropic/", "openai/", "google-vertex/"]; - -const KNOWN_ACRONYMS = new Set(["gpt", "glm"]); - -// For a known acronym, uppercase it, keep the version attached, and title-case -// any suffix: "gpt-5.6-sol" -> "GPT-5.6 Sol", "glm-5.2" -> "GLM-5.2". Other ids -// stay lowercase to avoid mangling ordinary names (e.g. "llama-3.1-8b"). -function formatProviderModelName(modelId: string): string { - const [acronym, version, ...suffix] = modelId.split("-"); - if (!KNOWN_ACRONYMS.has(acronym.toLowerCase())) return modelId.toLowerCase(); - const head = version - ? `${acronym.toUpperCase()}-${version}` - : acronym.toUpperCase(); - const tail = suffix.map( - (word) => word.charAt(0).toUpperCase() + word.slice(1).toLowerCase(), - ); - return [head, ...tail].join(" "); -} - -export function formatGatewayModelName(model: GatewayModel): string { - if (isCloudflareModel(model)) { - return formatProviderModelName(model.id.split("/").pop() ?? model.id); - } - - if (isOpenAIModel(model)) { - return formatProviderModelName(stripProviderPrefix(model.id)); - } - - return formatModelId(model.id); -} - -function stripProviderPrefix(modelId: string): string { - for (const prefix of PROVIDER_PREFIXES) { - if (modelId.startsWith(prefix)) { - return modelId.slice(prefix.length); - } - } - return modelId; -} - -export function formatModelId(modelId: string): string { - let cleanId = modelId; - for (const prefix of PROVIDER_PREFIXES) { - if (cleanId.startsWith(prefix)) { - cleanId = cleanId.slice(prefix.length); - break; - } - } - - cleanId = cleanId.replace(/(\d)-(\d)/g, "$1.$2"); - - const words = cleanId.split(/[-_]/).map((word) => { - if (word.match(/^[0-9.]+$/)) return word; - return word.charAt(0).toUpperCase() + word.slice(1).toLowerCase(); - }); - - return words.join(" "); -} diff --git a/packages/agent/src/utils/gateway.ts b/packages/agent/src/utils/gateway.ts index b258086070..5e97841742 100644 --- a/packages/agent/src/utils/gateway.ts +++ b/packages/agent/src/utils/gateway.ts @@ -1,3 +1,5 @@ +import { getCloudTaskGatewayUrl } from "@posthog/shared"; + export type GatewayProduct = | "posthog_code" | "background_agents" @@ -37,26 +39,7 @@ export { } from "@posthog/shared/posthog-property-headers"; function getGatewayBaseUrl(posthogHost: string): string { - const url = new URL(posthogHost); - const hostname = url.hostname; - - if (hostname === "localhost" || hostname === "127.0.0.1") { - return `${url.protocol}//localhost:3308`; - } - - if (hostname === "host.docker.internal") { - return `${url.protocol}//host.docker.internal:3308`; - } - - // The hosted dev environment runs its own LLM gateway with its own auth DB, - // so a dev-minted `pha_` token can't be routed to the US gateway — that's - // a different DB and returns 401 Authentication required. - if (hostname === "app.dev.posthog.dev") { - return "https://gateway.dev.posthog.dev"; - } - - const region = hostname.match(/^(us|eu)\.posthog\.com$/)?.[1] ?? "us"; - return `https://gateway.${region}.posthog.com`; + return getCloudTaskGatewayUrl(posthogHost).replace(/\/posthog_code$/, ""); } export function getLlmGatewayUrl( diff --git a/packages/api-client/src/posthog-client.ts b/packages/api-client/src/posthog-client.ts index 00bbbc7700..906e77f290 100644 --- a/packages/api-client/src/posthog-client.ts +++ b/packages/api-client/src/posthog-client.ts @@ -1,5 +1,4 @@ import "./generated.augment"; -import { isSupportedReasoningEffort } from "@posthog/agent/adapters/reasoning-effort"; import type { Adapter, CloudMcpServerImport, @@ -15,6 +14,7 @@ import type { import { DISMISSAL_REASON_OPTIONS, type DismissalReasonOptionValue, + isSupportedReasoningEffort, resolveCloudInitialPermissionMode, } from "@posthog/shared"; import type { diff --git a/packages/shared/src/cloud-task-models.test.ts b/packages/shared/src/cloud-task-models.test.ts new file mode 100644 index 0000000000..31418a7775 --- /dev/null +++ b/packages/shared/src/cloud-task-models.test.ts @@ -0,0 +1,200 @@ +import { describe, expect, it } from "vitest"; +import { + buildCloudTaskConfigOptions, + compareModelsForPicker, + formatGatewayModelName, + type GatewayModel, + getClaudeModelRecency, + isAnthropicModel, + isBlockedModelId, + isCloudflareModel, + normalizeGatewayModelsResponse, + pickAllowedModel, +} from "./cloud-task-models"; + +const model = ( + id: string, + owned_by = "anthropic", + allowed = true, +): GatewayModel => ({ + id, + owned_by, + context_window: 128000, + supports_streaming: true, + supports_vision: false, + allowed, +}); + +describe("formatGatewayModelName", () => { + it.each([ + [model("claude-opus-4-8"), "Claude Opus 4.8"], + [model("GPT-5.5", "openai"), "GPT-5.5"], + [model("openai/gpt-5.6-sol", "openai"), "GPT-5.6 Sol"], + [model("@cf/zai-org/glm-5.2", "cloudflare"), "GLM-5.2"], + [ + model("@cf/meta/llama-3.1-8b-instruct", "cloudflare"), + "llama-3.1-8b-instruct", + ], + ])("formats $id", (gatewayModel, expected) => { + expect(formatGatewayModelName(gatewayModel)).toBe(expected); + }); +}); + +describe("normalizeGatewayModelsResponse", () => { + it("corrects stale GLM 5.2 context-window metadata", () => { + const models = normalizeGatewayModelsResponse([ + model("@cf/zai-org/glm-5.2", "cloudflare"), + ]); + + expect(models[0]?.context_window).toBe(1_000_000); + }); +}); + +describe("isBlockedModelId", () => { + it.each([ + "claude-opus-4-5", + "claude-opus-4-6", + "claude-sonnet-4-5", + "ANTHROPIC/CLAUDE-HAIKU-4-5", + "gpt-5.2", + "gpt-5.3", + "gpt-5.3-codex", + "OPENAI/GPT-5.3-CODEX", + ])("blocks %s", (modelId) => { + expect(isBlockedModelId(modelId)).toBe(true); + }); +}); + +describe("getClaudeModelRecency", () => { + it.each([ + ["claude-haiku-4-5", 4005], + ["claude-sonnet-4-6", 4006], + ["claude-opus-4-7", 4007], + ["claude-opus-4-8", 4008], + ["claude-sonnet-5", 5000], + ])("ranks %s", (modelId, expected) => { + expect(getClaudeModelRecency(modelId)).toBe(expected); + }); + + it("ignores trailing dates and ranks unknown versions newest", () => { + expect(getClaudeModelRecency("claude-haiku-4-5-20251001")).toBe(4005); + expect(getClaudeModelRecency("claude-mystery")).toBe( + Number.MAX_SAFE_INTEGER, + ); + }); +}); + +describe("compareModelsForPicker", () => { + it("groups by capability and sorts newest first", () => { + const displayed = [ + "claude-fable-5", + "claude-opus-4-7", + "claude-mystery", + "claude-sonnet-5", + "claude-haiku-4-5", + "claude-sonnet-4-6", + "claude-opus-4-8", + ].sort(compareModelsForPicker); + + expect(displayed).toEqual([ + "claude-fable-5", + "claude-opus-4-8", + "claude-opus-4-7", + "claude-sonnet-5", + "claude-sonnet-4-6", + "claude-haiku-4-5", + "claude-mystery", + ]); + }); +}); + +describe("model classification", () => { + it("keeps Cloudflare models distinct from Anthropic", () => { + const gatewayModel = model("@cf/openai/gpt-oss", "openai"); + expect(isCloudflareModel(gatewayModel)).toBe(true); + expect(isAnthropicModel(gatewayModel)).toBe(false); + }); +}); + +describe("pickAllowedModel", () => { + const entry = (id: string, allowed: boolean) => ({ id, allowed }); + + it.each([ + [[entry("claude-opus-4-8", true)], "claude-opus-4-8", "claude-opus-4-8"], + [[entry("claude-opus-4-8", true)], "claude-sonnet-5", "claude-sonnet-5"], + [ + [ + entry("claude-opus-4-8", false), + entry("claude-sonnet-4-6", true), + entry("@cf/zai-org/glm-5.2", true), + ], + "claude-opus-4-8", + "@cf/zai-org/glm-5.2", + ], + [[entry("claude-opus-4-8", false)], "claude-opus-4-8", "claude-opus-4-8"], + [[], "claude-opus-4-8", "claude-opus-4-8"], + ] as const)("selects an allowed default", (models, preferred, expected) => { + expect(pickAllowedModel(models, preferred)).toBe(expected); + }); +}); + +describe("buildCloudTaskConfigOptions", () => { + it("builds Claude options with restrictions and reasoning policy", () => { + const options = buildCloudTaskConfigOptions( + [ + model("gpt-5.5", "openai"), + model("claude-opus-4-7", "anthropic"), + model("claude-opus-4-8", "anthropic", false), + model("@cf/zai-org/glm-5.2", "cloudflare"), + ], + "claude", + ); + + expect(options).toMatchObject([ + { id: "mode", currentValue: "plan" }, + { + id: "model", + currentValue: "@cf/zai-org/glm-5.2", + options: [ + { value: "claude-opus-4-7" }, + { + value: "claude-opus-4-8", + _meta: { "posthog.code/restrictedModel": true }, + }, + { value: "@cf/zai-org/glm-5.2" }, + ], + }, + ]); + expect(options.map((option) => option.id)).toEqual(["mode", "model"]); + }); + + it("builds Codex options with the shared default and reasoning levels", () => { + const options = buildCloudTaskConfigOptions( + [ + model("claude-opus-4-8"), + model("gpt-5.6", "openai"), + model("gpt-5.5", "openai"), + ], + "codex", + ); + + expect(options).toMatchObject([ + { id: "mode", currentValue: "auto" }, + { + id: "model", + currentValue: "gpt-5.5", + options: [{ value: "gpt-5.6" }, { value: "gpt-5.5" }], + }, + { + id: "reasoning_effort", + currentValue: "high", + options: [ + { value: "low" }, + { value: "medium" }, + { value: "high" }, + { value: "xhigh" }, + ], + }, + ]); + }); +}); diff --git a/packages/shared/src/cloud-task-models.ts b/packages/shared/src/cloud-task-models.ts new file mode 100644 index 0000000000..a0fc5f4477 --- /dev/null +++ b/packages/shared/src/cloud-task-models.ts @@ -0,0 +1,385 @@ +import type { Adapter } from "./adapter"; +import { CODEX_MODE_PRESETS } from "./execution-modes"; +import { restrictedModelMeta } from "./models"; +import { getReasoningEffortOptions } from "./reasoning-effort"; + +export interface GatewayModel { + id: string; + owned_by: string; + context_window: number; + supports_streaming: boolean; + supports_vision: boolean; + allowed: boolean; + restriction_reason?: string | null; +} + +interface GatewayModelsResponse { + data?: unknown[]; + models?: unknown[]; +} + +export interface CloudTaskConfigSelectOption { + value: string; + name: string; + description?: string; + _meta?: Record; +} + +export interface CloudTaskConfigOption { + id: string; + name: string; + type: "select"; + currentValue: string; + options: CloudTaskConfigSelectOption[]; + category: "mode" | "model" | "thought_level"; + description: string; +} + +export interface CloudTaskModePreset { + id: string; + name: string; + description: string; +} + +export const DEFAULT_GATEWAY_MODEL = "claude-opus-4-8"; + +export const DEFAULT_CODEX_MODEL = "gpt-5.5"; + +export const BLOCKED_GATEWAY_MODEL_IDS = [ + "gpt-5-mini", + "openai/gpt-5-mini", + "gpt-5.2", + "openai/gpt-5.2", + "gpt-5.3", + "openai/gpt-5.3", + "gpt-5.3-codex", + "openai/gpt-5.3-codex", + "claude-opus-4-5", + "anthropic/claude-opus-4-5", + "claude-opus-4-6", + "anthropic/claude-opus-4-6", + "claude-sonnet-4-5", + "anthropic/claude-sonnet-4-5", + "claude-haiku-4-5", + "anthropic/claude-haiku-4-5", +] as const; + +const BLOCKED_GATEWAY_MODELS = new Set(BLOCKED_GATEWAY_MODEL_IDS); + +const CLAUDE_MODE_PRESETS: readonly CloudTaskModePreset[] = [ + { + id: "default", + name: "Default", + description: "Standard behavior, prompts for dangerous operations", + }, + { + id: "acceptEdits", + name: "Accept Edits", + description: "Auto-accept file edit operations", + }, + { + id: "plan", + name: "Plan Mode", + description: "Planning mode, no actual tool execution", + }, + { + id: "bypassPermissions", + name: "Bypass Permissions", + description: "Auto-accept all permission requests", + }, + { + id: "auto", + name: "Auto Mode", + description: "Auto-approve file edits and shell commands", + }, +]; + +const PROVIDER_NAMES: Record = { + anthropic: "Anthropic", + openai: "OpenAI", + "google-vertex": "Gemini", +}; + +const MODEL_FAMILY_ORDER = ["fable", "opus", "sonnet", "haiku"]; +const PROVIDER_PREFIXES = ["anthropic/", "openai/", "google-vertex/"]; +const KNOWN_ACRONYMS = new Set(["gpt", "glm"]); +const MODEL_CONTEXT_WINDOW_OVERRIDES: Readonly> = { + "@cf/zai-org/glm-5.2": 1_000_000, +}; + +export function getCloudTaskGatewayUrl(posthogHost: string): string { + const url = new URL(posthogHost); + let gatewayBaseUrl: string; + + if (url.hostname === "localhost" || url.hostname === "127.0.0.1") { + gatewayBaseUrl = `${url.protocol}//localhost:3308`; + } else if (url.hostname === "host.docker.internal") { + gatewayBaseUrl = `${url.protocol}//host.docker.internal:3308`; + } else if (url.hostname === "app.dev.posthog.dev") { + gatewayBaseUrl = "https://gateway.dev.posthog.dev"; + } else { + const region = url.hostname.match(/^(us|eu)\.posthog\.com$/)?.[1] ?? "us"; + gatewayBaseUrl = `https://gateway.${region}.posthog.com`; + } + + return `${gatewayBaseUrl}/posthog_code`; +} + +function isGatewayModel(value: unknown): value is Partial & { + id: string; +} { + return ( + typeof value === "object" && + value !== null && + typeof (value as { id?: unknown }).id === "string" + ); +} + +export function normalizeGatewayModelsResponse(value: unknown): GatewayModel[] { + const response = value as GatewayModelsResponse; + const entries = Array.isArray(value) + ? value + : Array.isArray(response?.data) + ? response.data + : Array.isArray(response?.models) + ? response.models + : []; + + return entries + .filter(isGatewayModel) + .filter((model) => !isBlockedModelId(model.id)) + .map((model) => ({ + id: model.id, + owned_by: model.owned_by ?? "", + context_window: Math.max( + model.context_window ?? 0, + MODEL_CONTEXT_WINDOW_OVERRIDES[model.id] ?? 0, + ), + supports_streaming: model.supports_streaming ?? false, + supports_vision: model.supports_vision ?? false, + allowed: model.allowed !== false, + restriction_reason: model.restriction_reason ?? null, + })); +} + +export function isBlockedModelId(modelId: string): boolean { + return BLOCKED_GATEWAY_MODELS.has(modelId.toLowerCase()); +} + +export function isAnthropicModel(model: GatewayModel): boolean { + if (model.owned_by) { + return model.owned_by === "anthropic"; + } + return model.id.startsWith("claude-") || model.id.startsWith("anthropic/"); +} + +export function isOpenAIModel(model: GatewayModel): boolean { + if (model.owned_by) { + return model.owned_by === "openai"; + } + return model.id.startsWith("gpt-") || model.id.startsWith("openai/"); +} + +export function isCloudflareModelId(modelId: string): boolean { + return modelId.startsWith("@cf/"); +} + +export function isGlmModelId(modelId: string): boolean { + return modelId.toLowerCase().includes("glm"); +} + +export function isCloudflareModel(model: GatewayModel): boolean { + return isCloudflareModelId(model.id) || model.owned_by === "cloudflare"; +} + +export function pickAllowedModel( + models: ReadonlyArray>, + preferred: string, +): string { + if (models.length === 0) return preferred; + const preferredEntry = models.find((model) => model.id === preferred); + if (!preferredEntry || preferredEntry.allowed) return preferred; + const allowed = models.filter((model) => model.allowed); + if (allowed.length === 0) return preferred; + return allowed.reduce((best, candidate) => + getClaudeModelRecency(candidate.id) >= getClaudeModelRecency(best.id) + ? candidate + : best, + ).id; +} + +export function getProviderName(ownedBy: string): string { + return PROVIDER_NAMES[ownedBy] ?? ownedBy; +} + +export function getClaudeModelRecency(modelId: string): number { + const match = modelId.toLowerCase().match(/-(\d+)(?:[-.](\d+))?/); + if (!match) return Number.MAX_SAFE_INTEGER; + const major = Number(match[1]); + const minor = match[2] ? Number(match[2]) : 0; + return major * 1000 + minor; +} + +function getModelFamilyRank(modelId: string): number { + const normalizedModelId = modelId.toLowerCase(); + const index = MODEL_FAMILY_ORDER.findIndex((family) => + normalizedModelId.includes(family), + ); + return index === -1 ? MODEL_FAMILY_ORDER.length : index; +} + +export function compareModelsForPicker(a: string, b: string): number { + const familyDiff = getModelFamilyRank(a) - getModelFamilyRank(b); + if (familyDiff !== 0) return familyDiff; + return getClaudeModelRecency(b) - getClaudeModelRecency(a); +} + +function stripProviderPrefix(modelId: string): string { + for (const prefix of PROVIDER_PREFIXES) { + if (modelId.startsWith(prefix)) { + return modelId.slice(prefix.length); + } + } + return modelId; +} + +function formatProviderModelName(modelId: string): string { + const [acronym, version, ...suffix] = modelId.split("-"); + if (!KNOWN_ACRONYMS.has(acronym.toLowerCase())) return modelId.toLowerCase(); + const head = version + ? `${acronym.toUpperCase()}-${version}` + : acronym.toUpperCase(); + const tail = suffix.map( + (word) => word.charAt(0).toUpperCase() + word.slice(1).toLowerCase(), + ); + return [head, ...tail].join(" "); +} + +export function formatGatewayModelName(model: GatewayModel): string { + if (isCloudflareModel(model)) { + return formatProviderModelName(model.id.split("/").pop() ?? model.id); + } + if (isOpenAIModel(model)) { + return formatProviderModelName(stripProviderPrefix(model.id)); + } + return formatModelId(model.id); +} + +export function formatModelId(modelId: string): string { + const cleanId = stripProviderPrefix(modelId).replace(/(\d)-(\d)/g, "$1.$2"); + return cleanId + .split(/[-_]/) + .map((word) => { + if (/^[0-9.]+$/.test(word)) return word; + return word.charAt(0).toUpperCase() + word.slice(1).toLowerCase(); + }) + .join(" "); +} + +function getAdapterModels( + models: readonly GatewayModel[], + adapter: Adapter, +): GatewayModel[] { + return models.filter((model) => + adapter === "codex" + ? isOpenAIModel(model) + : isAnthropicModel(model) || isCloudflareModel(model), + ); +} + +function getModeOptions( + adapter: Adapter, + modePresets?: readonly CloudTaskModePreset[], +): CloudTaskConfigSelectOption[] { + const modes = + modePresets ?? + (adapter === "codex" ? CODEX_MODE_PRESETS : CLAUDE_MODE_PRESETS); + return modes.map((mode) => ({ + value: mode.id, + name: mode.name, + description: mode.description, + })); +} + +export function buildCloudTaskConfigOptions( + models: readonly GatewayModel[], + adapter: Adapter, + modePresets?: readonly CloudTaskModePreset[], +): CloudTaskConfigOption[] { + const adapterModels = getAdapterModels(models, adapter); + const modelOptions: CloudTaskConfigSelectOption[] = adapterModels.map( + (model) => ({ + value: model.id, + name: formatGatewayModelName(model), + description: `Context: ${model.context_window.toLocaleString()} tokens`, + ...(model.allowed ? {} : { _meta: restrictedModelMeta() }), + }), + ); + + if (adapter === "claude") { + modelOptions.sort( + (a, b) => getClaudeModelRecency(a.value) - getClaudeModelRecency(b.value), + ); + } + + const defaultModel = + adapter === "codex" + ? (modelOptions.find((option) => option.value === DEFAULT_CODEX_MODEL) + ?.value ?? + modelOptions[0]?.value ?? + "") + : DEFAULT_GATEWAY_MODEL; + const preferredModelId = modelOptions.some( + (option) => option.value === defaultModel, + ) + ? defaultModel + : (modelOptions[0]?.value ?? defaultModel); + const resolvedModelId = pickAllowedModel(adapterModels, preferredModelId); + + if (!modelOptions.some((option) => option.value === resolvedModelId)) { + modelOptions.unshift({ + value: resolvedModelId, + name: resolvedModelId, + description: "Custom model", + }); + } + + const configOptions: CloudTaskConfigOption[] = [ + { + id: "mode", + name: "Approval Preset", + type: "select", + currentValue: adapter === "codex" ? "auto" : "plan", + options: getModeOptions(adapter, modePresets), + category: "mode", + description: "Choose an approval and sandboxing preset for your session", + }, + { + id: "model", + name: "Model", + type: "select", + currentValue: resolvedModelId, + options: modelOptions, + category: "model", + description: "Choose which model the agent should use", + }, + ]; + + const reasoningOptions = getReasoningEffortOptions(adapter, resolvedModelId); + if (reasoningOptions) { + configOptions.push({ + id: adapter === "codex" ? "reasoning_effort" : "effort", + name: adapter === "codex" ? "Reasoning Level" : "Effort", + type: "select", + currentValue: "high", + options: reasoningOptions, + category: "thought_level", + description: + adapter === "codex" + ? "Controls how much reasoning effort the model uses" + : "Controls how much effort Claude puts into its response", + }); + } + + return configOptions; +} diff --git a/packages/shared/src/domain-types.test.ts b/packages/shared/src/domain-types.test.ts index f9b060cbed..1264ae0462 100644 --- a/packages/shared/src/domain-types.test.ts +++ b/packages/shared/src/domain-types.test.ts @@ -1,5 +1,22 @@ import { describe, expect, it } from "vitest"; -import { isContentlessTask } from "./domain-types"; +import { + isContentlessTask, + isTerminalStatus, + TERMINAL_STATUSES, +} from "./domain-types"; + +describe("task run statuses", () => { + it.each(TERMINAL_STATUSES)("identifies %s as terminal", (status) => { + expect(isTerminalStatus(status)).toBe(true); + }); + + it.each(["not_started", "queued", "in_progress", "unknown", null, undefined])( + "identifies %s as non-terminal", + (status) => { + expect(isTerminalStatus(status)).toBe(false); + }, + ); +}); describe("isContentlessTask", () => { it.each([ diff --git a/packages/shared/src/domain-types.ts b/packages/shared/src/domain-types.ts index 4aa90fd2c7..3600423dae 100644 --- a/packages/shared/src/domain-types.ts +++ b/packages/shared/src/domain-types.ts @@ -4,6 +4,7 @@ import type { AgentRuntime } from "./agent-runtime"; import type { DismissalReasonOptionValue } from "./dismissal-reasons"; import type { StoredLogEntry } from "./session-events"; import type { TaskRunArtifact } from "./task"; +import type { UploadableSkillSource } from "./skills"; // Execution mode schema and type - shared between main and renderer export const executionModeSchema = z.enum([ @@ -192,6 +193,37 @@ export type TaskRunStatus = | "failed" | "cancelled"; +export type TaskRunEnvironment = "local" | "cloud"; + +export type ArtifactType = + | "plan" + | "context" + | "reference" + | "output" + | "artifact" + | "user_attachment" + | "skill_bundle"; + +export interface TaskRunArtifactMetadata { + skill_name: string; + skill_source: UploadableSkillSource; + content_sha256: string; + bundle_format: "zip"; + schema_version: number; +} + +export interface TaskRunArtifact { + id?: string; + name: string; + type: ArtifactType; + source?: string; + size?: number; + content_type?: string; + metadata?: TaskRunArtifactMetadata; + storage_path?: string; + uploaded_at?: string; +} + export const TERMINAL_STATUSES = ["completed", "failed", "cancelled"] as const; export function isTerminalStatus( @@ -220,7 +252,7 @@ export interface TaskRun { model?: string | null; reasoning_effort?: "low" | "medium" | "high" | "xhigh" | "max" | null; stage?: string | null; // Current stage (e.g., 'research', 'plan', 'build') - environment?: "local" | "cloud"; + environment?: TaskRunEnvironment; status: TaskRunStatus; log_url: string; error_message: string | null; diff --git a/packages/shared/src/index.ts b/packages/shared/src/index.ts index ff03602e7e..5809830aba 100644 --- a/packages/shared/src/index.ts +++ b/packages/shared/src/index.ts @@ -71,6 +71,30 @@ export { promptBlocksToText, serializeCloudPrompt, } from "./cloud-prompt"; +export { + BLOCKED_GATEWAY_MODEL_IDS, + buildCloudTaskConfigOptions, + type CloudTaskConfigOption, + type CloudTaskConfigSelectOption, + type CloudTaskModePreset, + compareModelsForPicker, + DEFAULT_CODEX_MODEL, + DEFAULT_GATEWAY_MODEL, + formatGatewayModelName, + formatModelId, + type GatewayModel, + getClaudeModelRecency, + getCloudTaskGatewayUrl, + getProviderName, + isAnthropicModel, + isBlockedModelId, + isCloudflareModel, + isCloudflareModelId, + isGlmModelId, + isOpenAIModel, + normalizeGatewayModelsResponse, + pickAllowedModel, +} from "./cloud-task-models"; export { buildInboxDeeplink, buildScoutDeeplink, @@ -87,9 +111,28 @@ export { export { DISMISSAL_REASON_OPTIONS, type DismissalReasonOptionValue, + dismissalReasonLabel, isDismissalReasonSnooze, } from "./dismissal-reasons"; -export type { SignalReportPriority, Task } from "./domain-types"; +export { + type ArtifactType, + type CloudPermissionOption, + type CloudTaskErrorUpdate, + type CloudTaskLogsUpdate, + type CloudTaskPermissionRequestUpdate, + type CloudTaskSnapshotUpdate, + type CloudTaskStatusUpdate, + type CloudTaskUpdatePayload, + isTerminalStatus, + type SignalReportPriority, + type Task, + type TaskRun, + type TaskRunArtifact, + type TaskRunArtifactMetadata, + type TaskRunEnvironment, + type TaskRunStatus, + TERMINAL_STATUSES, +} from "./domain-types"; export * from "./enrichment"; export { classifyGatewayLimitError, @@ -213,6 +256,13 @@ export { isPrivateIpv4Octets, isPrivateIpv6Literal, } from "./private-network"; +export { + DEFAULT_REASONING_EFFORT, + getReasoningEffortOptions, + isSupportedReasoningEffort, + type ReasoningEffortOption, + type SupportedReasoningEffort, +} from "./reasoning-effort"; export { type CloudRegion, formatRegionBadge, @@ -276,15 +326,19 @@ export { serializeSkillMarkdown, stripFrontmatter, } from "./skills"; -export type { - ArtifactType, - PostHogAPIConfig, - TaskRun, - TaskRunArtifact, - TaskRunArtifactMetadata, - TaskRunEnvironment, - TaskRunStatus, -} from "./task"; +export type { PostHogAPIConfig } from "./task"; +export { + type CreateTaskAutomationOptions, + createTaskAutomationSchema, + type TaskAutomation, + type TaskAutomationList, + type TaskAutomationValidationErrorDetails, + taskAutomationListSchema, + taskAutomationSchema, + taskAutomationValidationErrorSchema, + type UpdateTaskAutomationOptions, + updateTaskAutomationSchema, +} from "./task-automation"; export type { TaskCreationInput, TaskCreationOutput, diff --git a/packages/shared/src/reasoning-effort.test.ts b/packages/shared/src/reasoning-effort.test.ts new file mode 100644 index 0000000000..5b1e301dd3 --- /dev/null +++ b/packages/shared/src/reasoning-effort.test.ts @@ -0,0 +1,20 @@ +import { describe, expect, it } from "vitest"; +import { isSupportedReasoningEffort } from "./reasoning-effort"; + +describe("isSupportedReasoningEffort", () => { + it.each([ + ["codex", "gpt-5.5", "xhigh", true], + ["codex", "gpt-5.6-sol", "max", true], + ["codex", "gpt-5.4", "max", false], + ["claude", "claude-opus-4-8", "xhigh", true], + ["claude", "claude-sonnet-4-6", "xhigh", false], + ["claude", "claude-opus-4-8", "minimal", false], + ] as const)( + "validates %s %s effort %s", + (adapter, modelId, effort, expected) => { + expect(isSupportedReasoningEffort(adapter, modelId, effort)).toBe( + expected, + ); + }, + ); +}); diff --git a/packages/shared/src/reasoning-effort.ts b/packages/shared/src/reasoning-effort.ts new file mode 100644 index 0000000000..15f534a612 --- /dev/null +++ b/packages/shared/src/reasoning-effort.ts @@ -0,0 +1,77 @@ +import type { Adapter } from "./adapter"; + +export type SupportedReasoningEffort = + | "low" + | "medium" + | "high" + | "xhigh" + | "max"; + +export const DEFAULT_REASONING_EFFORT: SupportedReasoningEffort = "high"; + +export interface ReasoningEffortOption { + value: SupportedReasoningEffort; + name: string; +} + +const BASE_OPTIONS: ReasoningEffortOption[] = [ + { value: "low", name: "Low" }, + { value: "medium", name: "Medium" }, + { value: "high", name: "High" }, +]; + +const CLAUDE_MODELS_WITH_EFFORT = new Set([ + "claude-opus-4-7", + "claude-opus-4-8", + "claude-sonnet-4-6", + "claude-sonnet-5", + "claude-fable-5", +]); + +const CLAUDE_MODELS_WITH_XHIGH_EFFORT = new Set([ + "claude-opus-4-7", + "claude-opus-4-8", + "claude-sonnet-5", + "claude-fable-5", +]); + +export function getReasoningEffortOptions( + adapter: Adapter, + modelId: string, +): ReasoningEffortOption[] | null { + if (adapter === "claude" && !CLAUDE_MODELS_WITH_EFFORT.has(modelId)) { + return null; + } + + const options = [...BASE_OPTIONS]; + const normalizedModelId = modelId.toLowerCase(); + const supportsXhigh = + adapter === "claude" + ? CLAUDE_MODELS_WITH_XHIGH_EFFORT.has(modelId) + : normalizedModelId.includes("gpt-5.5") || + normalizedModelId.includes("gpt-5.6"); + + if (supportsXhigh) { + options.push({ value: "xhigh", name: "Extra High" }); + } + if ( + (adapter === "claude" && supportsXhigh) || + (adapter === "codex" && normalizedModelId.includes("gpt-5.6")) + ) { + options.push({ value: "max", name: "Max" }); + } + + return options; +} + +export function isSupportedReasoningEffort( + adapter: Adapter, + modelId: string, + value: string, +): value is SupportedReasoningEffort { + return ( + getReasoningEffortOptions(adapter, modelId)?.some( + (option) => option.value === value, + ) ?? false + ); +} diff --git a/packages/shared/src/sessions.ts b/packages/shared/src/sessions.ts index 76a49cd736..181ff36c20 100644 --- a/packages/shared/src/sessions.ts +++ b/packages/shared/src/sessions.ts @@ -8,9 +8,9 @@ import type { } from "@agentclientprotocol/sdk"; import type { Adapter } from "./adapter"; import type { SkillButtonId } from "./analytics-events"; +import type { TaskRunArtifact, TaskRunStatus } from "./domain-types"; import type { ExecutionMode } from "./exec-types"; import type { AcpMessage } from "./session-events"; -import type { TaskRunArtifact, TaskRunStatus } from "./task"; export type { Adapter }; diff --git a/packages/shared/src/task-automation.test.ts b/packages/shared/src/task-automation.test.ts new file mode 100644 index 0000000000..57e6fcaa97 --- /dev/null +++ b/packages/shared/src/task-automation.test.ts @@ -0,0 +1,66 @@ +import { describe, expect, expectTypeOf, it } from "vitest"; +import { + type CreateTaskAutomationOptions, + createTaskAutomationSchema, + taskAutomationSchema, + taskAutomationValidationErrorSchema, + type UpdateTaskAutomationOptions, + updateTaskAutomationSchema, +} from "./task-automation"; + +describe("task automation contracts", () => { + it("normalizes optional automation response fields", () => { + expect( + taskAutomationSchema.parse({ + id: "automation-1", + name: "Daily PRs", + prompt: "Check PRs", + repository: "posthog/posthog", + cron_expression: "0 9 * * *", + last_run_at: null, + last_run_status: null, + last_task_id: null, + last_task_run_id: null, + last_error: null, + created_at: "2026-07-21T00:00:00Z", + updated_at: "2026-07-21T00:00:00Z", + }), + ).toMatchObject({ + github_integration: null, + timezone: null, + template_id: null, + enabled: true, + }); + }); + + it("keeps create fields required and update fields partial", () => { + const create = createTaskAutomationSchema.parse({ + name: "Daily PRs", + prompt: "Check PRs", + repository: "posthog/posthog", + cron_expression: "0 9 * * *", + timezone: "Europe/London", + }); + const update = updateTaskAutomationSchema.parse({ enabled: false }); + + expect(create.timezone).toBe("Europe/London"); + expect(update).toEqual({ enabled: false }); + expectTypeOf(create).toEqualTypeOf(); + expectTypeOf(update).toEqualTypeOf(); + }); + + it("preserves backend validation field attribution", () => { + expect( + taskAutomationValidationErrorSchema.parse({ + type: "validation_error", + detail: "Enter a valid cron expression.", + attr: "cron_expression", + }), + ).toEqual({ + type: "validation_error", + code: "invalid_input", + detail: "Enter a valid cron expression.", + attr: "cron_expression", + }); + }); +}); diff --git a/packages/shared/src/task-automation.ts b/packages/shared/src/task-automation.ts new file mode 100644 index 0000000000..0f3e4a4bde --- /dev/null +++ b/packages/shared/src/task-automation.ts @@ -0,0 +1,58 @@ +import { z } from "zod"; + +export const taskAutomationSchema = z.object({ + id: z.string(), + name: z.string(), + prompt: z.string(), + repository: z.string(), + github_integration: z.number().nullable().default(null), + cron_expression: z.string(), + timezone: z.string().nullable().default(null), + template_id: z.string().nullable().default(null), + enabled: z.boolean().default(true), + last_run_at: z.string().nullable(), + last_run_status: z.string().nullable(), + last_task_id: z.string().nullable(), + last_task_run_id: z.string().nullable(), + last_error: z.string().nullable(), + created_at: z.string(), + updated_at: z.string(), +}); +export type TaskAutomation = z.infer; + +export const taskAutomationListSchema = z.object({ + count: z.number(), + next: z.string().nullable().optional(), + previous: z.string().nullable().optional(), + results: z.array(taskAutomationSchema), +}); +export type TaskAutomationList = z.infer; + +export const createTaskAutomationSchema = z.object({ + name: z.string(), + prompt: z.string(), + repository: z.string(), + github_integration: z.number().nullable().optional(), + cron_expression: z.string(), + timezone: z.string(), + template_id: z.string().nullable().optional(), + enabled: z.boolean().optional(), +}); +export type CreateTaskAutomationOptions = z.infer< + typeof createTaskAutomationSchema +>; + +export const updateTaskAutomationSchema = createTaskAutomationSchema.partial(); +export type UpdateTaskAutomationOptions = z.infer< + typeof updateTaskAutomationSchema +>; + +export const taskAutomationValidationErrorSchema = z.object({ + type: z.string().optional(), + code: z.string().default("invalid_input"), + detail: z.string(), + attr: z.string().nullable().default(null), +}); +export type TaskAutomationValidationErrorDetails = z.infer< + typeof taskAutomationValidationErrorSchema +>; diff --git a/packages/shared/src/task.test.ts b/packages/shared/src/task.test.ts new file mode 100644 index 0000000000..8d666f2b11 --- /dev/null +++ b/packages/shared/src/task.test.ts @@ -0,0 +1,53 @@ +import { describe, expect, expectTypeOf, it } from "vitest"; +import type { + Task, + TaskRun, + TaskRunArtifact, + TaskRunStatus, +} from "./domain-types"; +import { + type CloudPermissionOption, + type CloudTaskUpdatePayload, + isTerminalStatus, + type Task as RootTask, + type TaskRun as RootTaskRun, + type TaskRunArtifact as RootTaskRunArtifact, + type TaskRunStatus as RootTaskRunStatus, + TERMINAL_STATUSES, +} from "./index"; +import type { + Task as LegacyTask, + TaskRun as LegacyTaskRun, + TaskRunArtifact as LegacyTaskRunArtifact, + TaskRunStatus as LegacyTaskRunStatus, +} from "./task"; + +describe("cloud task contract exports", () => { + it("keeps legacy and root task exports canonical", () => { + expectTypeOf().toEqualTypeOf(); + expectTypeOf().toEqualTypeOf(); + expectTypeOf().toEqualTypeOf(); + expectTypeOf().toEqualTypeOf(); + expectTypeOf().toEqualTypeOf(); + expectTypeOf().toEqualTypeOf(); + expectTypeOf().toEqualTypeOf(); + expectTypeOf().toEqualTypeOf(); + }); + + it("exports cloud permission and update payload contracts from the root", () => { + expectTypeOf().toMatchTypeOf<{ + kind: string; + optionId: string; + name: string; + }>(); + expectTypeOf().toEqualTypeOf< + "logs" | "status" | "snapshot" | "error" | "permission_request" + >(); + }); + + it("exports terminal status helpers from the root", () => { + expect(TERMINAL_STATUSES).toEqual(["completed", "failed", "cancelled"]); + expect(isTerminalStatus("completed")).toBe(true); + expect(isTerminalStatus("in_progress")).toBe(false); + }); +}); diff --git a/packages/shared/src/task.ts b/packages/shared/src/task.ts index b93c6e96d7..f6593d21bf 100644 --- a/packages/shared/src/task.ts +++ b/packages/shared/src/task.ts @@ -1,97 +1,12 @@ -// PostHog Task model (matches the desktop task API's OpenAPI schema) -import type { AgentRuntime } from "./agent-runtime"; -import type { UploadableSkillSource } from "./skills"; - -export interface Task { - id: string; - task_number?: number; - slug?: string; - title: string; - description: string; - origin_product: - | "error_tracking" - | "eval_clusters" - | "user_created" - | "support_queue" - | "session_summaries" - | "signal_report" - | "signals_scout" - | "slack"; - signal_report?: string | null; // Inbox report UUID when origin_product is "signal_report" - github_integration?: number | null; - repository: string; // Format: "organization/repository" (e.g., "posthog/posthog-js") - json_schema?: Record | null; // JSON schema for task output validation - internal?: boolean; - runtime?: AgentRuntime; - created_at: string; - updated_at: string; - created_by?: { - id: number; - uuid: string; - distinct_id: string; - first_name: string; - email: string; - }; - latest_run?: TaskRun; -} - -export type ArtifactType = - | "plan" - | "context" - | "reference" - | "output" - | "artifact" - | "user_attachment" - | "skill_bundle"; - -export interface TaskRunArtifactMetadata { - skill_name: string; - skill_source: UploadableSkillSource; - content_sha256: string; - bundle_format: "zip"; - schema_version: number; -} - -export interface TaskRunArtifact { - id?: string; - name: string; - type: ArtifactType; - source?: string; - size?: number; - content_type?: string; - metadata?: TaskRunArtifactMetadata; - storage_path?: string; - uploaded_at?: string; -} - -export type TaskRunStatus = - | "not_started" - | "queued" - | "in_progress" - | "completed" - | "failed" - | "cancelled"; - -export type TaskRunEnvironment = "local" | "cloud"; - -// TaskRun model - represents individual execution runs of tasks -export interface TaskRun { - id: string; - task: string; // Task ID - team: number; - branch: string | null; - stage: string | null; // Current stage (e.g., 'research', 'plan', 'build') - environment: TaskRunEnvironment; - status: TaskRunStatus; - log_url: string; - error_message: string | null; - output: Record | null; // Structured output (PR URL, commit SHA, etc.) - state: Record; // Intermediate run state (defaults to {}, never null) - artifacts?: TaskRunArtifact[]; - created_at: string; - updated_at: string; - completed_at: string | null; -} +export type { + ArtifactType, + Task, + TaskRun, + TaskRunArtifact, + TaskRunArtifactMetadata, + TaskRunEnvironment, + TaskRunStatus, +} from "./domain-types"; export interface PostHogAPIConfig { apiUrl: string; diff --git a/packages/workspace-server/src/services/agent/agent.ts b/packages/workspace-server/src/services/agent/agent.ts index 9654cb30b7..bd0821e254 100644 --- a/packages/workspace-server/src/services/agent/agent.ts +++ b/packages/workspace-server/src/services/agent/agent.ts @@ -20,24 +20,16 @@ import { } from "@posthog/agent"; import type { McpToolApprovals } from "@posthog/agent/adapters/claude/mcp/tool-metadata"; import { hydrateSessionJsonl } from "@posthog/agent/adapters/claude/session/jsonl-hydration"; -import { getReasoningEffortOptions } from "@posthog/agent/adapters/reasoning-effort"; import { Agent } from "@posthog/agent/agent"; import { getAvailableCodexModes, getAvailableModes, } from "@posthog/agent/execution-mode"; import { - DEFAULT_CODEX_MODEL, - DEFAULT_GATEWAY_MODEL, fetchGatewayModels, formatGatewayModelName, - type GatewayModel, getClaudeModelRecency, getProviderName, - isAnthropicModel, - isCloudflareModel, - isOpenAIModel, - pickAllowedModel, } from "@posthog/agent/gateway-models"; import { getLlmGatewayUrl } from "@posthog/agent/posthog-api"; import { @@ -72,10 +64,10 @@ import { import { type AcpMessage, type Adapter, + buildCloudTaskConfigOptions, type ExecutionMode, isAuthError, resolveCloudInitialPermissionMode, - restrictedModelMeta, serializeError, TypedEventEmitter, } from "@posthog/shared"; @@ -2395,111 +2387,14 @@ For git operations while detached: adapter: Adapter = "claude", ): Promise { const gatewayUrl = getLlmGatewayUrl(apiHost); - // Authenticated so the gateway can mark plan-restricted models; falls - // back to an anonymous fetch (everything allowed) without auth. const gatewayModels = await fetchGatewayModels({ gatewayUrl, authToken: (await this.agentAuthAdapter.gatewayAuthToken()) ?? undefined, }); - - // The Claude adapter can also drive Cloudflare `@cf/` models the gateway serves over its - // Anthropic-Messages surface, so the preview/default-model path must offer them too — otherwise an - // advertised `@cf/*` model is dropped here and the pre-session run falls back to Opus. - const modelFilter = - adapter === "codex" - ? isOpenAIModel - : (model: GatewayModel) => - isAnthropicModel(model) || isCloudflareModel(model); - - const adapterModels = gatewayModels.filter((model) => modelFilter(model)); - const modelOptions = adapterModels.map((model) => ({ - value: model.id, - name: formatGatewayModelName(model), - description: `Context: ${model.context_window.toLocaleString()} tokens`, - // Locked models stay listed so the picker can gate them instead of - // silently dropping them. - ...(model.allowed ? {} : { _meta: restrictedModelMeta() }), - })); - - // The gateway returns models in an arbitrary order. Sort Claude models - // oldest-to-newest so the picker is deterministic and the newest model - // lands at the end of the list, closest to the trigger. - if (adapter === "claude") { - modelOptions.sort( - (a, b) => - getClaudeModelRecency(a.value) - getClaudeModelRecency(b.value), - ); - } - - const defaultModel = - adapter === "codex" - ? (modelOptions.find((o) => o.value === DEFAULT_CODEX_MODEL)?.value ?? - modelOptions[0]?.value ?? - "") - : DEFAULT_GATEWAY_MODEL; - - const preferredModelId = modelOptions.some((o) => o.value === defaultModel) - ? defaultModel - : (modelOptions[0]?.value ?? defaultModel); - // Never preselect a model the org's plan can't use — it would 403 on the - // first message. - const resolvedModelId = pickAllowedModel(adapterModels, preferredModelId); - - if (!modelOptions.some((o) => o.value === resolvedModelId)) { - modelOptions.unshift({ - value: resolvedModelId, - name: resolvedModelId, - description: "Custom model", - }); - } - - const modes = - adapter === "codex" ? getAvailableCodexModes() : getAvailableModes(); - const modeOptions = modes.map((mode) => ({ - value: mode.id, - name: mode.name, - description: mode.description ?? undefined, - })); - const defaultMode = adapter === "codex" ? "auto" : "plan"; - - const configOptions: SessionConfigOption[] = [ - { - id: "mode", - name: "Approval Preset", - type: "select", - currentValue: defaultMode, - options: modeOptions, - category: "mode", - description: - "Choose an approval and sandboxing preset for your session", - }, - { - id: "model", - name: "Model", - type: "select", - currentValue: resolvedModelId, - options: modelOptions, - category: "model", - description: "Choose which model Claude should use", - }, - ]; - - const effortOpts = getReasoningEffortOptions(adapter, resolvedModelId); - if (effortOpts) { - configOptions.push({ - id: adapter === "codex" ? "reasoning_effort" : "effort", - name: adapter === "codex" ? "Reasoning Level" : "Effort", - type: "select", - currentValue: "high", - options: effortOpts, - category: "thought_level", - description: - adapter === "codex" - ? "Controls how much reasoning effort the model uses" - : "Controls how much effort Claude puts into its response", - }); - } - - return configOptions; + return buildCloudTaskConfigOptions( + gatewayModels, + adapter, + adapter === "codex" ? getAvailableCodexModes() : getAvailableModes(), + ) as SessionConfigOption[]; } } From ec5fb80c4e293a8a3f29a5bbfadb1559320cfb6c Mon Sep 17 00:00:00 2001 From: Richard Solomou Date: Fri, 24 Jul 2026 19:45:36 +0300 Subject: [PATCH 20/45] fix(models): preserve GLM effort policy Generated-By: PostHog Code Task-Id: c1bbe3cf-742b-4b24-bf96-d11a18b4cf22 --- packages/shared/src/cloud-task-models.test.ts | 11 ++++- packages/shared/src/reasoning-effort.test.ts | 3 ++ packages/shared/src/reasoning-effort.ts | 48 ++++++++++--------- 3 files changed, 38 insertions(+), 24 deletions(-) diff --git a/packages/shared/src/cloud-task-models.test.ts b/packages/shared/src/cloud-task-models.test.ts index 31418a7775..afbad989f8 100644 --- a/packages/shared/src/cloud-task-models.test.ts +++ b/packages/shared/src/cloud-task-models.test.ts @@ -164,8 +164,17 @@ describe("buildCloudTaskConfigOptions", () => { { value: "@cf/zai-org/glm-5.2" }, ], }, + { + id: "effort", + currentValue: "high", + options: [{ value: "high" }, { value: "max" }], + }, + ]); + expect(options.map((option) => option.id)).toEqual([ + "mode", + "model", + "effort", ]); - expect(options.map((option) => option.id)).toEqual(["mode", "model"]); }); it("builds Codex options with the shared default and reasoning levels", () => { diff --git a/packages/shared/src/reasoning-effort.test.ts b/packages/shared/src/reasoning-effort.test.ts index 5b1e301dd3..2ac4a9bbff 100644 --- a/packages/shared/src/reasoning-effort.test.ts +++ b/packages/shared/src/reasoning-effort.test.ts @@ -8,6 +8,9 @@ describe("isSupportedReasoningEffort", () => { ["codex", "gpt-5.4", "max", false], ["claude", "claude-opus-4-8", "xhigh", true], ["claude", "claude-sonnet-4-6", "xhigh", false], + ["claude", "@cf/zai-org/glm-5.2", "high", true], + ["claude", "@cf/zai-org/glm-5.2", "max", true], + ["claude", "@cf/zai-org/glm-5.2", "medium", false], ["claude", "claude-opus-4-8", "minimal", false], ] as const)( "validates %s %s effort %s", diff --git a/packages/shared/src/reasoning-effort.ts b/packages/shared/src/reasoning-effort.ts index 15f534a612..2fe12b7232 100644 --- a/packages/shared/src/reasoning-effort.ts +++ b/packages/shared/src/reasoning-effort.ts @@ -20,44 +20,46 @@ const BASE_OPTIONS: ReasoningEffortOption[] = [ { value: "high", name: "High" }, ]; -const CLAUDE_MODELS_WITH_EFFORT = new Set([ - "claude-opus-4-7", - "claude-opus-4-8", - "claude-sonnet-4-6", - "claude-sonnet-5", - "claude-fable-5", -]); +const CLAUDE_MODEL_EFFORTS: Readonly< + Record +> = { + "claude-opus-4-7": ["low", "medium", "high", "xhigh", "max"], + "claude-opus-4-8": ["low", "medium", "high", "xhigh", "max"], + "claude-sonnet-4-6": ["low", "medium", "high"], + "claude-sonnet-5": ["low", "medium", "high", "xhigh", "max"], + "claude-fable-5": ["low", "medium", "high", "xhigh", "max"], + "@cf/zai-org/glm-5.2": ["high", "max"], +}; -const CLAUDE_MODELS_WITH_XHIGH_EFFORT = new Set([ - "claude-opus-4-7", - "claude-opus-4-8", - "claude-sonnet-5", - "claude-fable-5", -]); +const EFFORT_NAMES: Record = { + low: "Low", + medium: "Medium", + high: "High", + xhigh: "Extra High", + max: "Max", +}; export function getReasoningEffortOptions( adapter: Adapter, modelId: string, ): ReasoningEffortOption[] | null { - if (adapter === "claude" && !CLAUDE_MODELS_WITH_EFFORT.has(modelId)) { - return null; + if (adapter === "claude") { + const efforts = CLAUDE_MODEL_EFFORTS[modelId]; + return ( + efforts?.map((value) => ({ value, name: EFFORT_NAMES[value] })) ?? null + ); } const options = [...BASE_OPTIONS]; const normalizedModelId = modelId.toLowerCase(); const supportsXhigh = - adapter === "claude" - ? CLAUDE_MODELS_WITH_XHIGH_EFFORT.has(modelId) - : normalizedModelId.includes("gpt-5.5") || - normalizedModelId.includes("gpt-5.6"); + normalizedModelId.includes("gpt-5.5") || + normalizedModelId.includes("gpt-5.6"); if (supportsXhigh) { options.push({ value: "xhigh", name: "Extra High" }); } - if ( - (adapter === "claude" && supportsXhigh) || - (adapter === "codex" && normalizedModelId.includes("gpt-5.6")) - ) { + if (adapter === "codex" && normalizedModelId.includes("gpt-5.6")) { options.push({ value: "max", name: "Max" }); } From fa2dda48da93a111d7ddca4081ce3763433da716 Mon Sep 17 00:00:00 2001 From: Richard Solomou Date: Sat, 25 Jul 2026 01:55:53 +0300 Subject: [PATCH 21/45] test(ui): wait for async content Generated-By: PostHog Code Task-Id: c1bbe3cf-742b-4b24-bf96-d11a18b4cf22 --- .../components/session-update/PlanApprovalView.test.tsx | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/packages/ui/src/features/sessions/components/session-update/PlanApprovalView.test.tsx b/packages/ui/src/features/sessions/components/session-update/PlanApprovalView.test.tsx index ef705ba77c..dd83ad5f6b 100644 --- a/packages/ui/src/features/sessions/components/session-update/PlanApprovalView.test.tsx +++ b/packages/ui/src/features/sessions/components/session-update/PlanApprovalView.test.tsx @@ -110,7 +110,7 @@ describe("PlanApprovalView", () => { ).toBeInTheDocument(); }); - it("uses updated content instead of stale raw input while streaming", () => { + it("uses updated content instead of stale raw input while streaming", async () => { renderView({ toolCall: makeToolCall({ status: "in_progress", @@ -124,7 +124,7 @@ describe("PlanApprovalView", () => { }), }); - expect(screen.getByText("Updated plan")).toBeInTheDocument(); + expect(await screen.findByText("Updated plan")).toBeInTheDocument(); expect(screen.queryByText("Initial plan")).not.toBeInTheDocument(); }); From 0960bc378d53c5e5e3cdb261eb619953bb69f18a Mon Sep 17 00:00:00 2001 From: Richard Solomou Date: Sat, 25 Jul 2026 02:06:34 +0300 Subject: [PATCH 22/45] test(ui): isolate plan approval presentation Generated-By: PostHog Code Task-Id: c1bbe3cf-742b-4b24-bf96-d11a18b4cf22 --- .../components/session-update/PlanApprovalView.test.tsx | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/packages/ui/src/features/sessions/components/session-update/PlanApprovalView.test.tsx b/packages/ui/src/features/sessions/components/session-update/PlanApprovalView.test.tsx index dd83ad5f6b..3bd1add8e8 100644 --- a/packages/ui/src/features/sessions/components/session-update/PlanApprovalView.test.tsx +++ b/packages/ui/src/features/sessions/components/session-update/PlanApprovalView.test.tsx @@ -2,9 +2,13 @@ import type { ToolCall } from "@posthog/ui/features/sessions/types"; import { Theme } from "@radix-ui/themes"; import { render, screen } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; -import { describe, expect, it } from "vitest"; +import { describe, expect, it, vi } from "vitest"; import { PlanApprovalView } from "./PlanApprovalView"; +vi.mock("../../../permissions/PlanContent", () => ({ + PlanContent: ({ plan }: { plan: string }) =>
{plan}
, +})); + const PLAN_MARKER = "Sentinel plan body for testing"; function makeToolCall(overrides: Partial = {}): ToolCall { @@ -124,7 +128,7 @@ describe("PlanApprovalView", () => { }), }); - expect(await screen.findByText("Updated plan")).toBeInTheDocument(); + expect(screen.getByText("Updated plan")).toBeInTheDocument(); expect(screen.queryByText("Initial plan")).not.toBeInTheDocument(); }); From 3f608f21ce8b2a1cf0a4277ed6fd77832371ee8f Mon Sep 17 00:00:00 2001 From: Richard Solomou Date: Tue, 28 Jul 2026 12:05:07 +0300 Subject: [PATCH 23/45] fix(shared): Preserve canonical task artifacts Keep the extracted domain declaration as the single source after rebasing onto current main. Generated-By: PostHog Code Task-Id: c1bbe3cf-742b-4b24-bf96-d11a18b4cf22 --- packages/shared/src/domain-types.ts | 1 - 1 file changed, 1 deletion(-) diff --git a/packages/shared/src/domain-types.ts b/packages/shared/src/domain-types.ts index 3600423dae..8f1e8dc56d 100644 --- a/packages/shared/src/domain-types.ts +++ b/packages/shared/src/domain-types.ts @@ -3,7 +3,6 @@ import type { Adapter } from "./adapter"; import type { AgentRuntime } from "./agent-runtime"; import type { DismissalReasonOptionValue } from "./dismissal-reasons"; import type { StoredLogEntry } from "./session-events"; -import type { TaskRunArtifact } from "./task"; import type { UploadableSkillSource } from "./skills"; // Execution mode schema and type - shared between main and renderer From 9c4235c3f9f0b9ca26c0eb2aac4f547288b57029 Mon Sep 17 00:00:00 2001 From: Richard Solomou Date: Tue, 28 Jul 2026 12:27:36 +0300 Subject: [PATCH 24/45] test(agent): match canonical model picker order Generated-By: PostHog Code Task-Id: c1bbe3cf-742b-4b24-bf96-d11a18b4cf22 --- packages/agent/src/adapters/claude/session/model-config.test.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/agent/src/adapters/claude/session/model-config.test.ts b/packages/agent/src/adapters/claude/session/model-config.test.ts index 22fd0bcf69..5561454baa 100644 --- a/packages/agent/src/adapters/claude/session/model-config.test.ts +++ b/packages/agent/src/adapters/claude/session/model-config.test.ts @@ -26,8 +26,8 @@ describe("applyAvailableModelsAllowlist", () => { "claude-opus-4-8", ]).options, ).toEqual([ - { value: "claude-sonnet-4-6", name: "Claude Sonnet 4.6" }, { value: "claude-opus-4-8", name: "Claude Opus 4.8" }, + { value: "claude-sonnet-4-6", name: "Claude Sonnet 4.6" }, ]); }); From 462ee21388eb963bcac2d66a5a65e84b56af48ba Mon Sep 17 00:00:00 2001 From: JonathanLab Date: Tue, 28 Jul 2026 12:18:03 +0200 Subject: [PATCH 25/45] Revert "fix(pi): disable cloud shell RPC" This reverts commit 7a3e0239c6b1bfce72e72cd5bfb3c6a7f202e807. --- .../src/pi-runtime/piSessionController.test.ts | 14 -------------- .../core/src/pi-runtime/piSessionController.ts | 1 - .../ui/src/features/pi-sessions/PiSessionView.tsx | 2 +- 3 files changed, 1 insertion(+), 16 deletions(-) diff --git a/packages/core/src/pi-runtime/piSessionController.test.ts b/packages/core/src/pi-runtime/piSessionController.test.ts index 93b58b276c..8a42ab2868 100644 --- a/packages/core/src/pi-runtime/piSessionController.test.ts +++ b/packages/core/src/pi-runtime/piSessionController.test.ts @@ -1022,20 +1022,6 @@ describe("PiSessionController", () => { ]); }); - it("retains cloud status when loading a cloud session", async () => { - const session = { - ...createSession(), - cloudStatus: "in_progress" as const, - }; - const controller = createController(session); - - await controller.connect("task-1", "run-1"); - - expect(controller.store.getState().sessions["task-1"].cloudStatus).toBe( - "in_progress", - ); - }); - it("loads session state and appends normalized runtime events", async () => { const initialEvent: AgentConversationEvent = { type: "assistant_message_chunk", diff --git a/packages/core/src/pi-runtime/piSessionController.ts b/packages/core/src/pi-runtime/piSessionController.ts index 95663fddaa..b63465c5d0 100644 --- a/packages/core/src/pi-runtime/piSessionController.ts +++ b/packages/core/src/pi-runtime/piSessionController.ts @@ -576,7 +576,6 @@ export class PiSessionController { events: reconciledEvents, status: resolvedStatus, stats, - cloudStatus: session.cloudStatus, models: currentSession.models, modelsLoaded: currentSession.modelsLoaded, thinkingLevels: currentSession.thinkingLevels, diff --git a/packages/ui/src/features/pi-sessions/PiSessionView.tsx b/packages/ui/src/features/pi-sessions/PiSessionView.tsx index e8550446fa..e2b42638a2 100644 --- a/packages/ui/src/features/pi-sessions/PiSessionView.tsx +++ b/packages/ui/src/features/pi-sessions/PiSessionView.tsx @@ -446,7 +446,7 @@ export function PiSessionView({ taskId, taskRunId }: PiSessionViewProps) { ? "A message is already queued" : undefined } - enableBashMode={session.cloudStatus === undefined} + enableBashMode enableCommands modelSelector={modelSelector} reasoningSelector={reasoningSelector} From 3f07f999a5120ef6918e9358765432a69a1130cb Mon Sep 17 00:00:00 2001 From: JonathanLab Date: Tue, 28 Jul 2026 12:44:34 +0200 Subject: [PATCH 26/45] fix(agent): await Pi RPC shutdown --- packages/agent/src/pi/rpc-client.test.ts | 4 +- pnpm-lock.yaml | 88 ++++++++++++------------ pnpm-workspace.yaml | 8 +-- 3 files changed, 50 insertions(+), 50 deletions(-) diff --git a/packages/agent/src/pi/rpc-client.test.ts b/packages/agent/src/pi/rpc-client.test.ts index 3da0303468..6579fca723 100644 --- a/packages/agent/src/pi/rpc-client.test.ts +++ b/packages/agent/src/pi/rpc-client.test.ts @@ -3,7 +3,7 @@ import { tmpdir } from "node:os"; import { join } from "node:path"; import { RpcClient } from "@earendil-works/pi-coding-agent"; import { describe, expect, it } from "vitest"; -import { createPiRpcClient, getPiRpcClientProcess } from "./rpc-client"; +import { createPiRpcClient } from "./rpc-client"; describe("createPiRpcClient", () => { it("does not put provider credentials in the child environment", () => { @@ -64,7 +64,7 @@ process.on("message", (request) => { followUp: [], }); } finally { - getPiRpcClientProcess(client)?.kill(); + await client.stop(); await rm(directory, { recursive: true }); } }); diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 17af034a07..62db42dcd7 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -7,17 +7,17 @@ settings: catalogs: default: '@earendil-works/pi-agent-core': - specifier: 0.82.0 - version: 0.82.0 + specifier: 0.82.1 + version: 0.82.1 '@earendil-works/pi-ai': - specifier: 0.82.0 - version: 0.82.0 + specifier: 0.82.1 + version: 0.82.1 '@earendil-works/pi-coding-agent': - specifier: 0.82.0 - version: 0.82.0 + specifier: 0.82.1 + version: 0.82.1 '@earendil-works/pi-tui': - specifier: 0.82.0 - version: 0.82.0 + specifier: 0.82.1 + version: 0.82.1 '@hono/node-server': specifier: ^1.13.7 version: 1.19.9 @@ -762,13 +762,13 @@ importers: version: 0.109.0(zod@4.4.3) '@earendil-works/pi-agent-core': specifier: 'catalog:' - version: 0.82.0(@modelcontextprotocol/sdk@1.29.0(zod@4.4.3))(ws@8.19.0)(zod@4.4.3) + version: 0.82.1(@modelcontextprotocol/sdk@1.29.0(zod@4.4.3))(ws@8.19.0)(zod@4.4.3) '@earendil-works/pi-ai': specifier: 'catalog:' - version: 0.82.0(@modelcontextprotocol/sdk@1.29.0(zod@4.4.3))(ws@8.19.0)(zod@4.4.3) + version: 0.82.1(@modelcontextprotocol/sdk@1.29.0(zod@4.4.3))(ws@8.19.0)(zod@4.4.3) '@earendil-works/pi-coding-agent': specifier: 'catalog:' - version: 0.82.0(@modelcontextprotocol/sdk@1.29.0(zod@4.4.3))(ws@8.19.0)(zod@4.4.3) + version: 0.82.1(@modelcontextprotocol/sdk@1.29.0(zod@4.4.3))(ws@8.19.0)(zod@4.4.3) '@hono/node-server': specifier: ^1.19.9 version: 1.19.9(hono@4.11.7) @@ -1050,13 +1050,13 @@ importers: dependencies: '@earendil-works/pi-ai': specifier: 'catalog:' - version: 0.82.0(@modelcontextprotocol/sdk@1.29.0(zod@4.4.3))(ws@8.19.0)(zod@4.4.3) + version: 0.82.1(@modelcontextprotocol/sdk@1.29.0(zod@4.4.3))(ws@8.19.0)(zod@4.4.3) '@earendil-works/pi-coding-agent': specifier: 'catalog:' - version: 0.82.0(@modelcontextprotocol/sdk@1.29.0(zod@4.4.3))(ws@8.19.0)(zod@4.4.3) + version: 0.82.1(@modelcontextprotocol/sdk@1.29.0(zod@4.4.3))(ws@8.19.0)(zod@4.4.3) '@earendil-works/pi-tui': specifier: 'catalog:' - version: 0.82.0 + version: 0.82.1 '@modelcontextprotocol/sdk': specifier: ^1.29.0 version: 1.29.0(zod@4.4.3) @@ -2746,22 +2746,22 @@ packages: '@drizzle-team/brocli@0.10.2': resolution: {integrity: sha512-z33Il7l5dKjUgGULTqBsQBQwckHh5AbIuxhdsIxDDiZAzBOrZO6q9ogcWC65kU382AfynTfgNumVcNIjuIua6w==} - '@earendil-works/pi-agent-core@0.82.0': - resolution: {integrity: sha512-bnS9DpOKK5T/F/gQkaOnYdMsuuciWiScfAHHWC+k5OQ0HxjSqMFQvp8keurULLoT4+v8NHv4V14pNvd4hsfC0Q==} + '@earendil-works/pi-agent-core@0.82.1': + resolution: {integrity: sha512-Z3kloziJIE2dmrisRckZX8zDca/gIv9/YdFAzeoqpHiLV2wsni6bL4hInNSjVKLbqT+4kqLIkph2JQLKvSepjg==} engines: {node: '>=22.19.0'} - '@earendil-works/pi-ai@0.82.0': - resolution: {integrity: sha512-8MvW9+zno13sXDuT2kFMnWeTNUufUhPeZDRVO+igGoBRCDWgn7Xh2FkRQI1mRuet6QhF4ENQuLYdIAOyG6BhNw==} + '@earendil-works/pi-ai@0.82.1': + resolution: {integrity: sha512-3WFYRhEp3lQB3444EhPMBcM7zSaEUE3eJgHOR7s4081NLqbw/FsWilIKWXSua0Gv3sRr7m9xMidR3pPDE7jI/A==} engines: {node: '>=22.19.0'} hasBin: true - '@earendil-works/pi-coding-agent@0.82.0': - resolution: {integrity: sha512-Qnqgn9zhJFQ2HZ8R4iNuGhyCk93XX6+eUw9i+TjTuo47amzCy93ft3bB6yaUCleCrNO58dJDHYSGNHv/GAPWKg==} + '@earendil-works/pi-coding-agent@0.82.1': + resolution: {integrity: sha512-zbkAhoIuDPMF3pKuja0ajZabrMWU29FUMV9A/XMXT/XC1yXs5xt6t6t13GogQFsDrDqbFP4DkZQO1w8rWRAzYA==} engines: {node: '>=22.19.0'} hasBin: true - '@earendil-works/pi-tui@0.82.0': - resolution: {integrity: sha512-9IDjQOXne7t9l2s2YcjnIBxsVNVPE7qScVSB3YmFlXsBW4pfo2gOElTxggV84KrRiGqABnlFPBWbf0k54hszHQ==} + '@earendil-works/pi-tui@0.82.1': + resolution: {integrity: sha512-9yN8hALfKaxZq7n54EMxqhFCWnMi6LHkraMJ/1YjHiATq75XrI6XDMVppn9EDtiK7Fks8hUe1SDXUTrIvwRWfQ==} engines: {node: '>=22.19.0'} '@ecies/ciphers@0.2.6': @@ -2855,11 +2855,11 @@ packages: '@esbuild-kit/core-utils@3.3.2': resolution: {integrity: sha512-sPRAnw9CdSsRmEtnsl2WXWdyquogVpB3yZ3dgwJfe8zrOzTsV7cJvmwrKVa+0ma5BoiGJ+BoqkMvawbayKUsqQ==} - deprecated: 'Merged into tsx: https://tsx.hirok.io' + deprecated: 'Merged into tsx: https://tsx.is' '@esbuild-kit/esm-loader@2.6.5': resolution: {integrity: sha512-FxEMIkJKnodyA1OaCUoEvbYRkoZlLZ4d/eXFu9Fh8CbBBgP5EmZxrfTRyN0qpXZ4vOvqnE5YdRdcrmUUXuU+dA==} - deprecated: 'Merged into tsx: https://tsx.hirok.io' + deprecated: 'Merged into tsx: https://tsx.is' '@esbuild/aix-ppc64@0.25.12': resolution: {integrity: sha512-Hhmwd6CInZ3dwpuGTF8fJG6yoWmsToE+vYgD4nytZVxcu1ulHpUQRAB1UJ8+N1Am3Mz4+xOByoQoSZf4D+CpkA==} @@ -9572,8 +9572,8 @@ packages: resolution: {integrity: sha512-0je+qPKHEMohvfRTCEo3CrPG6cAzAYgmzKyxRiYSSDkS6eGJdyVJm7WaYA5ECaAD9wLB2T4EEeymA5aFVcYXCA==} engines: {node: '>=6'} - deslop-js@0.7.8: - resolution: {integrity: sha512-QMmb3Z/ARvYZmZneudb8cnY/4mVvZTdhUyA9TC2skwOcm7KvY9zyOdn0TApQc4rL0VM2TffFkmo3ky/lJZX7qw==} + deslop-js@0.8.3: + resolution: {integrity: sha512-axNV/iX3Zq9xt0MYesmbBGxneeeY/HrYgXTsaM4+GOrdxXP9JyCfTdC4j8zx09bBML94oSIpFNyn9U1f0oEPqQ==} destroy@1.2.0: resolution: {integrity: sha512-2sJGJTaXIIaR1w4iJSNoN0hnMY7Gpc/n8D4qSCJw8QqFWXf7cuAgnEHxBpweaVcPevC2l3KpjYCx3NypQQgaJg==} @@ -12942,8 +12942,8 @@ packages: engines: {node: ^20.19.0 || >=22.12.0} hasBin: true - oxlint-plugin-react-doctor@0.7.8: - resolution: {integrity: sha512-3f9/jFLIC/KRLPYqxiXSk20cq47luGy9Oz5Ru7nK7w0EI9B9zuMvAU92bK9jFiwFE7Lc9PA8ER6Y+naYaGFQGw==} + oxlint-plugin-react-doctor@0.8.3: + resolution: {integrity: sha512-S1Gq1H9+BpziApWcZ/sWPNYYy1FMkFmtSEGiqIJ4/Aac76LfkeutPFtSRlkJF1JlcrbYFf7LXcfcxZCJgmVBBQ==} engines: {node: ^20.19.0 || >=22.13.0} oxlint@1.66.0: @@ -13589,8 +13589,8 @@ packages: resolution: {integrity: sha512-+NRMYs2DyTP4/tqWz371Oo50JqmWltR1h2gcdgUMAWZJIAvrd0/SqlCfx7tpzpl/s36rzw6qH2MjoNrxtRNYhA==} engines: {node: ^20.9.0 || >=22} - react-doctor@0.7.8: - resolution: {integrity: sha512-G3spmtZJE/gWWPRJ3rpgUWTPRDJpEmdRja7iNZ7RAXlfpEO+NWVzPTca/cPI9hLwPo2Aq5/BZggo5JDBrwGrlA==} + react-doctor@0.8.3: + resolution: {integrity: sha512-FfG7YQKb1yv1UNk2gknZzLAPx6L3RMOhYsYbslr4MSpVMNk5xBHlu9VxjtS0br0DEBWjkZovrf/C1MrchiIzAw==} engines: {node: ^20.19.0 || >=22.13.0} hasBin: true @@ -17170,9 +17170,9 @@ snapshots: '@drizzle-team/brocli@0.10.2': {} - '@earendil-works/pi-agent-core@0.82.0(@modelcontextprotocol/sdk@1.29.0(zod@4.4.3))(ws@8.19.0)(zod@4.4.3)': + '@earendil-works/pi-agent-core@0.82.1(@modelcontextprotocol/sdk@1.29.0(zod@4.4.3))(ws@8.19.0)(zod@4.4.3)': dependencies: - '@earendil-works/pi-ai': 0.82.0(@modelcontextprotocol/sdk@1.29.0(zod@4.4.3))(ws@8.19.0)(zod@4.4.3) + '@earendil-works/pi-ai': 0.82.1(@modelcontextprotocol/sdk@1.29.0(zod@4.4.3))(ws@8.19.0)(zod@4.4.3) diff: 8.0.4 ignore: 7.0.5 typebox: 1.1.38 @@ -17185,7 +17185,7 @@ snapshots: - ws - zod - '@earendil-works/pi-ai@0.82.0(@modelcontextprotocol/sdk@1.29.0(zod@4.4.3))(ws@8.19.0)(zod@4.4.3)': + '@earendil-works/pi-ai@0.82.1(@modelcontextprotocol/sdk@1.29.0(zod@4.4.3))(ws@8.19.0)(zod@4.4.3)': dependencies: '@anthropic-ai/sdk': 0.91.1(zod@4.4.3) '@aws-sdk/client-bedrock-runtime': 3.1048.0 @@ -17206,11 +17206,11 @@ snapshots: - ws - zod - '@earendil-works/pi-coding-agent@0.82.0(@modelcontextprotocol/sdk@1.29.0(zod@4.4.3))(ws@8.19.0)(zod@4.4.3)': + '@earendil-works/pi-coding-agent@0.82.1(@modelcontextprotocol/sdk@1.29.0(zod@4.4.3))(ws@8.19.0)(zod@4.4.3)': dependencies: - '@earendil-works/pi-agent-core': 0.82.0(@modelcontextprotocol/sdk@1.29.0(zod@4.4.3))(ws@8.19.0)(zod@4.4.3) - '@earendil-works/pi-ai': 0.82.0(@modelcontextprotocol/sdk@1.29.0(zod@4.4.3))(ws@8.19.0)(zod@4.4.3) - '@earendil-works/pi-tui': 0.82.0 + '@earendil-works/pi-agent-core': 0.82.1(@modelcontextprotocol/sdk@1.29.0(zod@4.4.3))(ws@8.19.0)(zod@4.4.3) + '@earendil-works/pi-ai': 0.82.1(@modelcontextprotocol/sdk@1.29.0(zod@4.4.3))(ws@8.19.0)(zod@4.4.3) + '@earendil-works/pi-tui': 0.82.1 '@silvia-odwyer/photon-node': 0.3.4 chalk: 5.6.2 cross-spawn: 7.0.6 @@ -17236,7 +17236,7 @@ snapshots: - ws - zod - '@earendil-works/pi-tui@0.82.0': + '@earendil-works/pi-tui@0.82.1': dependencies: get-east-asian-width: 1.6.0 marked: 18.0.5 @@ -24229,7 +24229,7 @@ snapshots: dequal@2.0.3: {} - deslop-js@0.7.8: + deslop-js@0.8.3: dependencies: '@oxc-project/types': 0.138.0 fast-glob: 3.3.3 @@ -28606,7 +28606,7 @@ snapshots: '@oxfmt/binding-win32-ia32-msvc': 0.45.0 '@oxfmt/binding-win32-x64-msvc': 0.45.0 - oxlint-plugin-react-doctor@0.7.8: + oxlint-plugin-react-doctor@0.8.3: dependencies: '@typescript-eslint/types': 8.62.0 eslint-scope: 9.1.2 @@ -29359,19 +29359,19 @@ snapshots: transitivePeerDependencies: - supports-color - react-doctor@0.7.8(@opentelemetry/core@2.8.0(@opentelemetry/api@1.9.1))(@opentelemetry/exporter-trace-otlp-http@0.208.0(@opentelemetry/api@1.9.1))(eslint@10.5.0(jiti@2.7.0)): + react-doctor@0.8.3(@opentelemetry/core@2.8.0(@opentelemetry/api@1.9.1))(@opentelemetry/exporter-trace-otlp-http@0.208.0(@opentelemetry/api@1.9.1))(eslint@10.5.0(jiti@2.7.0)): dependencies: '@babel/code-frame': 7.29.0 '@sentry/node': 10.61.0(@opentelemetry/core@2.8.0(@opentelemetry/api@1.9.1))(@opentelemetry/exporter-trace-otlp-http@0.208.0(@opentelemetry/api@1.9.1)) agent-install: 0.0.5 conf: 15.1.0 confbox: 0.2.4 - deslop-js: 0.7.8 + deslop-js: 0.8.3 eslint-plugin-react-hooks: 7.1.1(eslint@10.5.0(jiti@2.7.0)) jiti: 2.7.0 magicast: 0.5.3 oxlint: 1.66.0 - oxlint-plugin-react-doctor: 0.7.8 + oxlint-plugin-react-doctor: 0.8.3 prompts: 2.4.2 typescript: 5.9.3 vscode-languageserver: 9.0.1 @@ -29641,7 +29641,7 @@ snapshots: preact: 10.29.2 prompts: 2.4.2 react: 19.2.6 - react-doctor: 0.7.8(@opentelemetry/core@2.8.0(@opentelemetry/api@1.9.1))(@opentelemetry/exporter-trace-otlp-http@0.208.0(@opentelemetry/api@1.9.1))(eslint@10.5.0(jiti@2.7.0)) + react-doctor: 0.8.3(@opentelemetry/core@2.8.0(@opentelemetry/api@1.9.1))(@opentelemetry/exporter-trace-otlp-http@0.208.0(@opentelemetry/api@1.9.1))(eslint@10.5.0(jiti@2.7.0)) react-dom: 19.2.6(react@19.2.6) react-grab: 0.1.48(react@19.2.6) optionalDependencies: diff --git a/pnpm-workspace.yaml b/pnpm-workspace.yaml index 11a53d7625..d0aa2e24e7 100644 --- a/pnpm-workspace.yaml +++ b/pnpm-workspace.yaml @@ -6,10 +6,10 @@ packages: catalog: '@hono/node-server': ^1.13.7 '@hono/trpc-server': ^0.3.4 - '@earendil-works/pi-agent-core': 0.82.0 - '@earendil-works/pi-ai': 0.82.0 - '@earendil-works/pi-coding-agent': 0.82.0 - '@earendil-works/pi-tui': 0.82.0 + '@earendil-works/pi-agent-core': 0.82.1 + '@earendil-works/pi-ai': 0.82.1 + '@earendil-works/pi-coding-agent': 0.82.1 + '@earendil-works/pi-tui': 0.82.1 '@parcel/watcher': ^2.5.6 '@phosphor-icons/react': ^2.1.10 '@posthog/quill': 0.3.0-beta.24 From f5669368f29af902f96dcad43e0099545663879e Mon Sep 17 00:00:00 2001 From: JonathanLab Date: Tue, 28 Jul 2026 12:56:24 +0200 Subject: [PATCH 27/45] fix(agent): handle Pi bootstrap pipe closure --- packages/agent/src/pi/rpc-client.test.ts | 3 +++ packages/agent/src/pi/rpc-client.ts | 1 + 2 files changed, 4 insertions(+) diff --git a/packages/agent/src/pi/rpc-client.test.ts b/packages/agent/src/pi/rpc-client.test.ts index 6579fca723..25db00d0a4 100644 --- a/packages/agent/src/pi/rpc-client.test.ts +++ b/packages/agent/src/pi/rpc-client.test.ts @@ -37,6 +37,9 @@ describe("createPiRpcClient", () => { await writeFile( hostPath, ` +import { closeSync } from "node:fs"; + +closeSync(3); process.stdin.resume(); process.on("message", (request) => { const data = request.method === "clear_queue" diff --git a/packages/agent/src/pi/rpc-client.ts b/packages/agent/src/pi/rpc-client.ts index 3f9654338a..068153bdff 100644 --- a/packages/agent/src/pi/rpc-client.ts +++ b/packages/agent/src/pi/rpc-client.ts @@ -157,6 +157,7 @@ class SecurePiRpcClient extends RpcClient { } const bootstrapPipe = child.stdio[3] as Writable | null; + bootstrapPipe?.on("error", () => {}); bootstrapPipe?.end( JSON.stringify({ providerOptions: this.providerOptions }), ); From a77aaae140c340979cb0e6957d078337c9e57576 Mon Sep 17 00:00:00 2001 From: "posthog[bot]" <206114724+posthog[bot]@users.noreply.github.com> Date: Tue, 28 Jul 2026 11:14:34 +0000 Subject: [PATCH 28/45] chore(visual): update storybook baselines 12 updated, 12 removed Run: 9e2a836d-06fa-4eb0-80ab-9cafd1241ac9 Co-authored-by: jonathanlab <32547391+jonathanlab@users.noreply.github.com> --- apps/code/snapshots.yml | 48 ++++++++++++++++++++--------------------- 1 file changed, 24 insertions(+), 24 deletions(-) diff --git a/apps/code/snapshots.yml b/apps/code/snapshots.yml index 3cbfb34a08..20fbaa732f 100644 --- a/apps/code/snapshots.yml +++ b/apps/code/snapshots.yml @@ -72,30 +72,6 @@ snapshots: hash: v1.k4693efd2.1d199b2c4bba8034cb18fb5b49866b64a6eb3311591add5f58ec288710dce2e9.FCc9egSaO1Onih417yciSYKNROe8zqDZyyS6xuXCZGw billing-usagemeter--zero-spend-limit--light: hash: v1.k4693efd2.b50c18736bbde45fc89f4c0ac7fd616286e5c1ecc2a56910fffca13eaa13d8a2.31z4F0JmvMeR4yqk5pjs_OQ3MypnM2miZI9al4wuc8c - channels-taskfeedrow--agent-origin--dark: - hash: v1.k4693efd2.82f8c70a399c9ea768201933e202fcc2fde74332c5153aeabf6835ace79beee6.-SVboDRIZ_-nh4PtMliCv7e7iCtdUm-1m2xQ9Xzcdfw - channels-taskfeedrow--agent-origin--light: - hash: v1.k4693efd2.4c1db3470d5ec9b4db9872d0e048314d0078979e8fbdb602b8c648ba70d7209f._KPUvj46p33u7A4Xndz-git-4rp7VKca1MpeTztlB6w - channels-taskfeedrow--human-email-only--dark: - hash: v1.k4693efd2.0f4dece4643b7375b77746347658f4fe65b37817cdfa9ad1ba697067e619f428.f0jKZ4OYqsmeJjXffjt-amK9cwE81JF0lgSMqtn37PA - channels-taskfeedrow--human-email-only--light: - hash: v1.k4693efd2.14db0510f7ea5c57985cf0cb60ab08f5e146c7243eeb1b0b963f74e85c7d60f4.ijHQLK0lZjjUlOQuzYzisImU9DWsNmZM7i5GYfWU4E0 - channels-taskfeedrow--human-started--dark: - hash: v1.k4693efd2.e0abdae2e8ac2ef29793ea7afc2a2580c6610e07a3c640ad782af7fd3c12592a.YUgF_BTPzZoqnWXvyLBRs3w9C_P3Yfh4-cP4mUqs2NQ - channels-taskfeedrow--human-started--light: - hash: v1.k4693efd2.84c26deb1a587fe061238b3982b555167575893bacc9cd2667d4d3f74646261f.abqc0Voe6FIQFflzDJTv1KewcZ4XxcDz1a0mFlJu7oM - channels-taskfeedrow--long-prompt--dark: - hash: v1.k4693efd2.876d34660bc267af79d39a971a681ba279a870061dba10b905412112c2a5e4dd.rSw1X068udMs8Arglu_WgX5RL8WZlpAc8_kVONPQhF4 - channels-taskfeedrow--long-prompt--light: - hash: v1.k4693efd2.d0d6e4b6bfa257c3f46d991777f72c345437b0be2ee16a182fa925d3ece7dc9e.6PGeOlMxVauJSQia0FAIirzaYio79-I7H4x-9LvEUbw - channels-taskfeedrow--no-prompt--dark: - hash: v1.k4693efd2.9fa967f1a9acdeba0c50a9e45ae649f118ee26938037dbefd1bb0577067b03d4.va1lGsscqLW86-5yCKc3H6pQ8Az7_HBItkgGTnTQiYU - channels-taskfeedrow--no-prompt--light: - hash: v1.k4693efd2.a02339aecdb6fc327fdf6586e28ceeebf8490442fe217eb8bb46f1ce66bff78a.2l7XdvdW3D0GlLmPwDdTFgjZWdav68bZX_LYResVKwo - channels-taskfeedrow--no-starter--dark: - hash: v1.k4693efd2.66fff212f14afa9dd3bc532698a042383014a0cabd4331d30eff878148e1773d.RiectZxFTqk1husLYI-UV5ko1uZ4lpbjWstE7Hy3MJE - channels-taskfeedrow--no-starter--light: - hash: v1.k4693efd2.9ee4ad64ed7d2c5c62d0cd682d903b09da643d462d49aae0858afa500e92256b.OhGj64l-o7iWtxzJo1qt0A_CU9ME0ZRr_5e455RGUJQ components-permissions-permissionselector--create-new-file--dark: hash: v1.k4693efd2.c54203a4e636b83b3d24d7ed9c4ace8659db87cd8f231d9dc2ecc03320e31646.epDm7LebiLzlp0uuZBrE-Obt_anAn0xsE8bHFnm5vos components-permissions-permissionselector--create-new-file--light: @@ -708,6 +684,30 @@ snapshots: hash: v1.k4693efd2.cfff4c6bbb0acef1c9941b21d2f69c24b234c2c206282acf6fcada0cb2bd2386.xILjVPp8ecQmGbPjGDdhbRyi9KFPpA_1nfYF3bfjwuc skill-buttons-skillbuttonsmenu--default--light: hash: v1.k4693efd2.eb9bd49b9700641f6f3c90653447c9767c52e997afec9843e8649b90752faad3.VdxEtBwtkU3ioy2evcNbrEKxPjQ1GDUpZAQzFe-Uho4 + spaces-taskfeedrow--agent-origin--dark: + hash: v1.k4693efd2.82f8c70a399c9ea768201933e202fcc2fde74332c5153aeabf6835ace79beee6.yEZq_qnKxItv5u0Ydcqa1B4VCwv8lwXUvcaFScE2XO0 + spaces-taskfeedrow--agent-origin--light: + hash: v1.k4693efd2.4c1db3470d5ec9b4db9872d0e048314d0078979e8fbdb602b8c648ba70d7209f.wt_Ufkw30jHFnxEMFXE0wQG8ignuCURuUUYIvNMnJ6U + spaces-taskfeedrow--human-email-only--dark: + hash: v1.k4693efd2.611794921d5c6270a211f4a6bc3d91fd63a74a58f0c905a79271cbc7d9c73eb8.balV69yNKs9biJzweJYwpk5e1UjZ11fFersE9UMhvEU + spaces-taskfeedrow--human-email-only--light: + hash: v1.k4693efd2.961288ce410b6376bab92428ad148b881c768e9edea4821db4818d13682f50c4.79pzkpuXwARCmUS4HtZRvb2n7dM8HUZZs-ErIXX4wwg + spaces-taskfeedrow--human-started--dark: + hash: v1.k4693efd2.6e1cf0720234b1c64ba04f6f3af965d32438551771bc3680c4391bb8463b8eee.4O-iWXxbDv-L9MsuPt6KoDvK7C95R0Jj6LTYjDwwEhk + spaces-taskfeedrow--human-started--light: + hash: v1.k4693efd2.7156c1e5b2359b5f792529aa9041ad9277842d51c3423bda4f0c029c58dcef28.PF-JYrCTfHxEzFG9dnabJxDYS47wQXnxurx8Ev_LfkY + spaces-taskfeedrow--long-prompt--dark: + hash: v1.k4693efd2.00acb2487749ee5a1db2c95da8ae0e2fb2de0815d653973e60c7ceb3ff808135.pdV864kTAUM1IgcuME846VwInafp3f-WFFxxxE4H_oY + spaces-taskfeedrow--long-prompt--light: + hash: v1.k4693efd2.8a6b372abdf9c60a558c59a703f8e7c8f65fcb55fd99e6c97c6eca23e6a76781.HbZ7kfuE60F5SWucdbYLhy98POhk8MPPX2TLA_r2hz8 + spaces-taskfeedrow--no-prompt--dark: + hash: v1.k4693efd2.ff77770c9b5c656b213c139270322defa38306d714abeee6b5fe4e17e2056f88.KnYgUDb_jYEfH5aSoJbt1cwyJXAfDo_eKmIfn0fxjjc + spaces-taskfeedrow--no-prompt--light: + hash: v1.k4693efd2.d3d755ee7ea6eccbed93ea18e6dcbccbd9f77f019f5c153cb0278131b08c324f.l_KwHWHFTzODefftPGtlWsXLFi6tbH78ZRQeeY7X4Z0 + spaces-taskfeedrow--no-starter--dark: + hash: v1.k4693efd2.66fff212f14afa9dd3bc532698a042383014a0cabd4331d30eff878148e1773d.-FvbS5RpndI7QXNTpQhbQShDktPrOV_jdcWDY4MUYF4 + spaces-taskfeedrow--no-starter--light: + hash: v1.k4693efd2.9ee4ad64ed7d2c5c62d0cd682d903b09da643d462d49aae0858afa500e92256b.lCc1Q5KBHSLnYp7N9k6HEBxDPjf0SblH9zDLHQ7dYo8 task-detail-continueclisessions--importing--dark: hash: v1.k4693efd2.3196e21b9daa2f183c79eaf2c70a5775365dbee3f1d920c64927e20a0fadcbf9.FjHhjEX3urGbycQiDehA9JQ7H2zfyl1id9rDSVjrtoI task-detail-continueclisessions--importing--light: From 3890211965c586b214549755752e599e9e9951d7 Mon Sep 17 00:00:00 2001 From: JonathanLab Date: Tue, 28 Jul 2026 18:41:51 +0200 Subject: [PATCH 29/45] fix(pi): address cloud runtime review findings --- .../translatePiConversation.test.ts | 25 ++++++ .../conversation/translatePiConversation.ts | 26 +++++- .../conversation/translatePiMessage.test.ts | 33 +++++++ .../src/pi/conversation/translatePiMessage.ts | 29 ++++-- packages/agent/src/pi/model-catalog.test.ts | 26 +++++- packages/agent/src/pi/model-catalog.ts | 5 +- packages/agent/src/pi/runtime.test.ts | 59 ++++++++++++ packages/agent/src/pi/runtime.ts | 17 ++++ packages/agent/src/server/bin.ts | 1 + .../agent/src/server/pi-agent-server.test.ts | 46 +++++++++- packages/agent/src/server/pi-agent-server.ts | 90 +++++++++++++++---- .../pi-runtime/cloudPiSessionClient.test.ts | 4 +- .../src/pi-runtime/cloudPiSessionClient.ts | 4 +- .../pi-runtime/piSessionController.test.ts | 90 +++++++++++++++++++ .../src/pi-runtime/piSessionController.ts | 21 +++-- .../features/pi-sessions/PiSessionView.tsx | 10 ++- .../components/chat-thread/ChatThread.tsx | 14 ++- .../chat-thread/ChatThreadFooter.tsx | 15 +++- .../features/settings/settingsStore.test.ts | 19 ++++ .../ui/src/features/settings/settingsStore.ts | 5 ++ .../task-detail/components/TaskInput.tsx | 9 +- .../services/pi-session/pi-session.test.ts | 18 +++- .../src/services/pi-session/pi-session.ts | 37 +++++--- 23 files changed, 538 insertions(+), 65 deletions(-) diff --git a/packages/agent/src/pi/conversation/translatePiConversation.test.ts b/packages/agent/src/pi/conversation/translatePiConversation.test.ts index 1f9980d4d8..44453dd9de 100644 --- a/packages/agent/src/pi/conversation/translatePiConversation.test.ts +++ b/packages/agent/src/pi/conversation/translatePiConversation.test.ts @@ -411,6 +411,31 @@ describe("createPiConversationTranslator", () => { ]); }); + it("preserves streamed direct bash output when the command fails", () => { + const translator = createPiConversationTranslator(); + const [started] = translator.beginDirectBash("failing-command"); + if (started?.type !== "tool_call_started") { + throw new Error("Expected a direct bash tool call"); + } + + translator.translateEvent({ + type: "bash_execution_update", + id: "req_1", + delta: "partial output", + }); + + expect(translator.failDirectBash("transport failed")).toMatchObject([ + { + type: "tool_call_updated", + toolCall: { + id: started.toolCall.id, + status: "failed", + rawOutput: "partial output\n\ntransport failed", + }, + }, + ]); + }); + it("streams tool execution start, output updates, and completion", () => { const translator = createPiConversationTranslator(); const message = assistant( diff --git a/packages/agent/src/pi/conversation/translatePiConversation.ts b/packages/agent/src/pi/conversation/translatePiConversation.ts index 26e353014c..ce0f906259 100644 --- a/packages/agent/src/pi/conversation/translatePiConversation.ts +++ b/packages/agent/src/pi/conversation/translatePiConversation.ts @@ -114,6 +114,7 @@ export function createPiConversationTranslator(): PiConversationTranslator { let directBashSequence = 0; let activeDirectBash: | { + nextOutputSize: number; output: string; startedAt: number; toolCallId: string; @@ -123,7 +124,12 @@ export function createPiConversationTranslator(): PiConversationTranslator { function beginDirectBash(command: string): AgentConversationEvent[] { const startedAt = Date.now(); const toolCallId = `pi-bash-live-${startedAt}-${++directBashSequence}`; - activeDirectBash = { output: "", startedAt, toolCallId }; + activeDirectBash = { + nextOutputSize: 4_096, + output: "", + startedAt, + toolCallId, + }; return [ { @@ -179,7 +185,10 @@ export function createPiConversationTranslator(): PiConversationTranslator { } function failDirectBash(message: string): AgentConversationEvent[] { - return finishDirectBash("failed", message); + const output = [activeDirectBash?.output, message] + .filter(Boolean) + .join("\n\n"); + return finishDirectBash("failed", output); } function translateHistoryMessage( @@ -298,6 +307,15 @@ export function createPiConversationTranslator(): PiConversationTranslator { } directBash.output += event.delta; + if (directBash.output.length >= 4_096) { + if (directBash.output.length < directBash.nextOutputSize) { + return []; + } + while (directBash.nextOutputSize <= directBash.output.length) { + directBash.nextOutputSize *= 2; + } + } + return [ { type: "tool_call_updated", @@ -441,6 +459,10 @@ export function createPiConversationTranslator(): PiConversationTranslator { const timestamp = event.result?.summary ? Math.max(Date.now(), latestConversationTimestamp + 1) : latestConversationTimestamp; + latestConversationTimestamp = Math.max( + latestConversationTimestamp, + timestamp, + ); const events: AgentConversationEvent[] = [ { type: "runtime_status", diff --git a/packages/agent/src/pi/conversation/translatePiMessage.test.ts b/packages/agent/src/pi/conversation/translatePiMessage.test.ts index 40df9d453f..cad1a6b6cb 100644 --- a/packages/agent/src/pi/conversation/translatePiMessage.test.ts +++ b/packages/agent/src/pi/conversation/translatePiMessage.test.ts @@ -160,6 +160,39 @@ describe("createPiMessageTranslator", () => { ]); }); + it("preserves images in generic extension tool results", () => { + const translator = createPiMessageTranslator(); + const content: ToolResultMessage["content"] = [ + { type: "image", data: "aW1hZ2U=", mimeType: "image/png" }, + ]; + const message: ToolResultMessage = { + role: "toolResult", + toolCallId: "extension-image", + toolName: "screenshot", + content, + isError: false, + timestamp: 12, + }; + + expect(translator.translate(message)).toMatchObject([ + { + type: "tool_call_updated", + toolCall: { + content: [ + { + type: "content", + content: { + type: "image", + data: "aW1hZ2U=", + mimeType: "image/png", + }, + }, + ], + }, + }, + ]); + }); + it("keeps built-in tool translation and raw output", () => { const translator = createPiMessageTranslator(); const content: ToolResultMessage["content"] = [ diff --git a/packages/agent/src/pi/conversation/translatePiMessage.ts b/packages/agent/src/pi/conversation/translatePiMessage.ts index 983e3d2dff..17f6ce1b07 100644 --- a/packages/agent/src/pi/conversation/translatePiMessage.ts +++ b/packages/agent/src/pi/conversation/translatePiMessage.ts @@ -47,16 +47,31 @@ function isPiToolName(name: string): name is PiToolName { function toGenericToolContent( resultContent: ToolResultMessage["content"], ): AgentToolCallContent[] | undefined { - const text = resultContent - .filter((block) => block.type === "text") - .map((block) => block.text) - .join(""); + const content: AgentToolCallContent[] = []; + let text = ""; - if (!text) { - return undefined; + const appendText = () => { + if (!text) { + return; + } + content.push({ type: "content", content: { type: "text", text } }); + text = ""; + }; + + for (const block of resultContent) { + if (block.type === "text") { + text += block.text; + continue; + } + const translated = toContent(block); + if (translated) { + appendText(); + content.push({ type: "content", content: translated }); + } } + appendText(); - return [{ type: "content", content: { type: "text", text } }]; + return content.length > 0 ? content : undefined; } function toContent(block: { diff --git a/packages/agent/src/pi/model-catalog.test.ts b/packages/agent/src/pi/model-catalog.test.ts index 1b9bc27e56..aad2ffed61 100644 --- a/packages/agent/src/pi/model-catalog.test.ts +++ b/packages/agent/src/pi/model-catalog.test.ts @@ -1,7 +1,15 @@ -import { describe, expect, it } from "vitest"; -import { resolvePosthogPiModelCatalog } from "./model-catalog"; +import { afterEach, describe, expect, it, vi } from "vitest"; +import { + fetchPosthogPiModelCatalog, + resolvePosthogPiModelCatalog, +} from "./model-catalog"; describe("resolvePosthogPiModelCatalog", () => { + afterEach(() => { + delete process.env.PI_OFFLINE; + vi.unstubAllGlobals(); + }); + it("uses the PostHog provider model configuration for gateway models", () => { const models = resolvePosthogPiModelCatalog( [ @@ -42,4 +50,18 @@ describe("resolvePosthogPiModelCatalog", () => { }), ]); }); + + it("uses fallback models without fetching while offline", async () => { + process.env.PI_OFFLINE = "1"; + const fetch = vi.fn(); + vi.stubGlobal("fetch", fetch); + + const models = await fetchPosthogPiModelCatalog( + "https://gateway.example.com", + "us", + ); + + expect(fetch).not.toHaveBeenCalled(); + expect(models.length).toBeGreaterThan(0); + }); }); diff --git a/packages/agent/src/pi/model-catalog.ts b/packages/agent/src/pi/model-catalog.ts index aef4b2f7b7..96b6bb45d5 100644 --- a/packages/agent/src/pi/model-catalog.ts +++ b/packages/agent/src/pi/model-catalog.ts @@ -42,6 +42,9 @@ export async function fetchPosthogPiModelCatalog( region: CloudRegion, apiKey?: string, ): Promise { - const models = await fetchPosthogGatewayModels(gatewayUrl, apiKey); + const models = + process.env.PI_OFFLINE || process.env.HARNESS_STATIC_MODELS + ? [] + : await fetchPosthogGatewayModels(gatewayUrl, apiKey); return resolvePosthogPiModelCatalog(models, region); } diff --git a/packages/agent/src/pi/runtime.test.ts b/packages/agent/src/pi/runtime.test.ts index 201d42d9a5..e1c82972fa 100644 --- a/packages/agent/src/pi/runtime.test.ts +++ b/packages/agent/src/pi/runtime.test.ts @@ -165,6 +165,65 @@ describe("PiRuntime", () => { ); }); + it("drops cleared queued message ids before matching later messages", async () => { + const { client, emit, send } = createClient(); + const runtime = new PiRuntime(client); + const conversationListener = vi.fn(); + runtime.onConversationEvent(conversationListener); + send.mockResolvedValue({ + type: "response", + command: "steer", + success: true, + }); + + await runtime.sendCommand({ + id: "cleared-id", + type: "steer", + message: "continue", + }); + runtime.clearPendingQueuedUserMessages(); + send.mockImplementationOnce(async () => { + emit({ + type: "message_end", + message: { role: "user", content: "continue", timestamp: 1 }, + }); + return { type: "response", command: "prompt", success: true }; + }); + await runtime.sendCommand({ + id: "current-id", + type: "prompt", + message: "continue", + }); + + expect(conversationListener).toHaveBeenCalledWith( + expect.objectContaining({ type: "user_message", id: "current-id" }), + ); + }); + + it("rejects concurrent direct bash commands", async () => { + const { client, send } = createClient(); + const runtime = new PiRuntime(client); + let resolveBash: (value: unknown) => void = () => {}; + send.mockImplementationOnce( + () => + new Promise((resolve) => { + resolveBash = resolve; + }), + ); + + const first = runtime.sendCommand({ type: "bash", command: "sleep 1" }); + await expect( + runtime.sendCommand({ type: "bash", command: "pwd" }), + ).rejects.toThrow("already running"); + resolveBash({ + type: "response", + command: "bash", + success: true, + data: { output: "", exitCode: 0, cancelled: false }, + }); + await first; + }); + it("forwards native queue snapshots", () => { const { client, emit } = createClient(); const runtime = new PiRuntime(client); diff --git a/packages/agent/src/pi/runtime.ts b/packages/agent/src/pi/runtime.ts index f89d95790d..31a2a776a3 100644 --- a/packages/agent/src/pi/runtime.ts +++ b/packages/agent/src/pi/runtime.ts @@ -25,7 +25,9 @@ export class PiRuntime { private readonly pendingUserMessages: Array<{ id: string; message: string; + type: "prompt" | "steer" | "follow_up"; }> = []; + private directBashActive = false; constructor(client: PiRpcClient) { this.client = client; @@ -58,6 +60,7 @@ export class PiRuntime { this.pendingUserMessages.push({ id: command.id, message: command.message, + type: command.type, }); } if (command.type !== "bash") { @@ -75,6 +78,10 @@ export class PiRuntime { } } + if (this.directBashActive) { + throw new Error("A Pi bash command is already running"); + } + this.directBashActive = true; this.emitConversationEvents( this.translator.beginDirectBash(command.command), ); @@ -93,6 +100,16 @@ export class PiRuntime { const message = error instanceof Error ? error.message : String(error); this.emitConversationEvents(this.translator.failDirectBash(message)); throw error; + } finally { + this.directBashActive = false; + } + } + + clearPendingQueuedUserMessages(): void { + for (let index = this.pendingUserMessages.length - 1; index >= 0; index--) { + if (this.pendingUserMessages[index]?.type !== "prompt") { + this.pendingUserMessages.splice(index, 1); + } } } diff --git a/packages/agent/src/server/bin.ts b/packages/agent/src/server/bin.ts index 3b5826c0bd..2c7e22d16b 100644 --- a/packages/agent/src/server/bin.ts +++ b/packages/agent/src/server/bin.ts @@ -182,6 +182,7 @@ program // bodies can't leak it. Defense in depth, not a boundary: same-UID // processes can still read the container's initial env via /proc. delete process.env.POSTHOG_AGENT_OTEL_LOGS_TOKEN; + delete process.env.POSTHOG_TASK_RUN_SESSION_TOKEN; const mode = options.mode === "background" ? "background" : "interactive"; const createPr = parseBooleanOption(options.createPr, "--createPr"); diff --git a/packages/agent/src/server/pi-agent-server.test.ts b/packages/agent/src/server/pi-agent-server.test.ts index 22fccc7736..c42b815578 100644 --- a/packages/agent/src/server/pi-agent-server.test.ts +++ b/packages/agent/src/server/pi-agent-server.test.ts @@ -83,6 +83,44 @@ describe("PiAgentServer", () => { ]); }); + it("bounds events retained while no SSE client is connected", () => { + const server = new PiAgentServer(config()) as unknown as { + broadcast(event: Record): void; + pendingEvents: Record[]; + }; + + for (let index = 0; index < 1_100; index++) { + server.broadcast({ type: "test", index }); + } + + expect(server.pendingEvents).toHaveLength(1_000); + expect(server.pendingEvents[0]).toEqual({ type: "test", index: 100 }); + }); + + it("flushes long-running conversation logs in bounded batches", async () => { + const appendTaskRunLog = vi.fn( + async (_taskId: string, _runId: string, _entries: unknown[]) => ({}), + ); + const server = new PiAgentServer(config()) as unknown as { + posthogAPI: { appendTaskRunLog: typeof appendTaskRunLog }; + handleEvent(event: Record): void; + logFlushQueue: Promise; + }; + server.posthogAPI.appendTaskRunLog = appendTaskRunLog; + + for (let index = 0; index < 100; index++) { + server.handleEvent({ + type: "assistant_message_chunk", + timestamp: index, + content: { type: "text", text: String(index) }, + }); + } + await server.logFlushQueue; + + expect(appendTaskRunLog).toHaveBeenCalledOnce(); + expect(appendTaskRunLog.mock.calls[0]?.[2]).toHaveLength(100); + }); + it("uses the durable message id for an idle native Pi prompt", async () => { const sendCommand = vi.fn( async (_command: Record) => ({}), @@ -312,6 +350,7 @@ describe("PiAgentServer", () => { getQueue: vi.fn(async () => queue), clearQueue: vi.fn(async () => queue), }; + const clearPendingQueuedUserMessages = vi.fn(); const server = new PiAgentServer(config()) as unknown as { session: unknown; executeCommand( @@ -319,10 +358,15 @@ describe("PiAgentServer", () => { params: Record, ): Promise; }; - server.session = { runtime: { client } }; + server.session = { + runtime: { client, clearPendingQueuedUserMessages }, + }; await expect(server.executeCommand(method, {})).resolves.toEqual(queue); expect(client[operation]).toHaveBeenCalledOnce(); + expect(clearPendingQueuedUserMessages).toHaveBeenCalledTimes( + method === "queue_clear" ? 1 : 0, + ); }, ); diff --git a/packages/agent/src/server/pi-agent-server.ts b/packages/agent/src/server/pi-agent-server.ts index 0315adcfa9..853aeb23b0 100644 --- a/packages/agent/src/server/pi-agent-server.ts +++ b/packages/agent/src/server/pi-agent-server.ts @@ -34,6 +34,9 @@ interface PiCloudSession { } const emptySchema = z.object({}); +const MAX_PENDING_EVENTS = 1_000; +const MAX_PENDING_LOG_ENTRIES = 10_000; +const LOG_FLUSH_ENTRY_COUNT = 100; const userMessageCommandSchema = z .object({ @@ -78,6 +81,8 @@ export class PiAgentServer { private settledPersistenceQueue: Promise = Promise.resolve(); private pendingLogEntries: StoredLogEntry[] = []; private logFlushQueue: Promise = Promise.resolve(); + private logFlushActive = false; + private logFlushRequested = false; private readonly canceledSseControllers = new WeakSet(); constructor(private readonly config: AgentServerConfig) { @@ -148,7 +153,11 @@ export class PiAgentServer { this.logger.error("Failed to sync Pi session during shutdown", error), ); session.unsubscribe(); - await session.runtime.client.stop(); + await session.runtime.client + .stop() + .catch((error) => + this.logger.error("Failed to stop Pi client during shutdown", error), + ); } this.session = null; await this.flushConversationLog().catch((error) => @@ -171,6 +180,15 @@ export class PiAgentServer { message, } satisfies AgentConversationEvent, }); + await this.settledPersistenceQueue.catch((syncError) => + this.logger.error( + "Failed to persist settled Pi turn after crash", + syncError, + ), + ); + await this.syncTaskSession().catch((syncError) => + this.logger.error("Failed to sync crashed Pi session", syncError), + ); await this.flushConversationLog().catch((syncError) => this.logger.error("Failed to persist crashed Pi events", syncError), ); @@ -378,6 +396,7 @@ export class PiAgentServer { const client = createPiRpcClient({ cliPath: this.config.piRpcHostPath, cwd, + model: this.config.model, sessionFile: restoredSessionFile, providerOptions: { apiKey: this.config.apiKey, @@ -453,8 +472,11 @@ export class PiAgentServer { return client.abort(); case "queue_get": return client.getQueue(); - case "queue_clear": - return client.clearQueue(); + case "queue_clear": { + const queue = await client.clearQueue(); + runtime.clearPendingQueuedUserMessages(); + return queue; + } case "pi/rpc": return runtime.sendCommand(params.command as RpcCommand); } @@ -639,8 +661,15 @@ export class PiAgentServer { ? (event.event as AgentConversationEvent) : undefined, }); + if (this.pendingLogEntries.length > MAX_PENDING_LOG_ENTRIES) { + this.pendingLogEntries.splice( + 0, + this.pendingLogEntries.length - MAX_PENDING_LOG_ENTRIES, + ); + } if ( event.type === "pi_run_started" || + this.pendingLogEntries.length >= LOG_FLUSH_ENTRY_COUNT || (event.event as { type?: string } | undefined)?.type === "turn_completed" ) { @@ -655,29 +684,54 @@ export class PiAgentServer { this.session.sseController.send(event); } else { this.pendingEvents.push(event); + if (this.pendingEvents.length > MAX_PENDING_EVENTS) { + this.pendingEvents.splice( + 0, + this.pendingEvents.length - MAX_PENDING_EVENTS, + ); + } } } private flushConversationLog(): Promise { + if (this.logFlushActive) { + this.logFlushRequested = true; + return this.logFlushQueue; + } if (this.pendingLogEntries.length === 0) { return this.logFlushQueue; } - const entries = this.pendingLogEntries; - this.pendingLogEntries = []; - const flush = this.logFlushQueue - .then(() => - this.posthogAPI.appendTaskRunLog( - this.config.taskId, - this.config.runId, - entries, - ), - ) - .then(() => undefined) - .catch((error) => { - this.pendingLogEntries = [...entries, ...this.pendingLogEntries]; - throw error; - }); + this.logFlushActive = true; + const flush = (async () => { + do { + this.logFlushRequested = false; + const entries = this.pendingLogEntries; + this.pendingLogEntries = []; + if (entries.length === 0) { + return; + } + try { + await this.posthogAPI.appendTaskRunLog( + this.config.taskId, + this.config.runId, + entries, + ); + } catch (error) { + this.pendingLogEntries = [ + ...entries, + ...this.pendingLogEntries, + ].slice(-MAX_PENDING_LOG_ENTRIES); + throw error; + } + } while ( + this.logFlushRequested || + this.pendingLogEntries.length >= LOG_FLUSH_ENTRY_COUNT + ); + })().finally(() => { + this.logFlushActive = false; + }); + this.logFlushQueue = flush.catch(() => undefined); return flush; } diff --git a/packages/core/src/pi-runtime/cloudPiSessionClient.test.ts b/packages/core/src/pi-runtime/cloudPiSessionClient.test.ts index 819390a0f6..8a5fa8a457 100644 --- a/packages/core/src/pi-runtime/cloudPiSessionClient.test.ts +++ b/packages/core/src/pi-runtime/cloudPiSessionClient.test.ts @@ -126,7 +126,7 @@ describe("CloudPiSessionClient", () => { } }); - it("accepts an in-progress reconnect snapshot as runtime readiness", async () => { + it("accepts a reconnect snapshot with a start entry as runtime readiness", async () => { const cloud = createCloudTaskClient(); vi.mocked(cloud.client.sendCommand).mockResolvedValue({ success: true, @@ -145,7 +145,7 @@ describe("CloudPiSessionClient", () => { taskId: "task-1", runId: "run-1", kind: "snapshot", - status: "in_progress", + status: "queued", newEntries: [{ type: "pi_run_started" }], totalEntryCount: 1, }); diff --git a/packages/core/src/pi-runtime/cloudPiSessionClient.ts b/packages/core/src/pi-runtime/cloudPiSessionClient.ts index e1a473a79b..f0ac8739f4 100644 --- a/packages/core/src/pi-runtime/cloudPiSessionClient.ts +++ b/packages/core/src/pi-runtime/cloudPiSessionClient.ts @@ -237,10 +237,8 @@ export class CloudPiSessionClient implements PiSession { onError: (error: unknown) => void, onCloudStatus?: (status: TaskRunStatus) => void, ): void { - const snapshotCanProveReadiness = - update.kind === "snapshot" && update.status === "in_progress"; const hasCurrentReadinessEvent = - (update.kind === "logs" || snapshotCanProveReadiness) && + (update.kind === "logs" || update.kind === "snapshot") && update.newEntries.some((entry) => entry.type === "pi_run_started"); if (hasCurrentReadinessEvent) { this.markRuntimeReady(); diff --git a/packages/core/src/pi-runtime/piSessionController.test.ts b/packages/core/src/pi-runtime/piSessionController.test.ts index 8a42ab2868..f52d947a12 100644 --- a/packages/core/src/pi-runtime/piSessionController.test.ts +++ b/packages/core/src/pi-runtime/piSessionController.test.ts @@ -518,6 +518,35 @@ describe("PiSessionController", () => { expect(session.client.followUp).toHaveBeenCalledWith("then summarize"); }); + it("does not replay an already restored prompt after a later queue failure", async () => { + let onEvent: (event: AgentConversationEvent) => void = () => {}; + const session = createSession(); + session.retry = vi.fn(async () => {}); + vi.mocked(session.onConversationEvent).mockImplementation((handler) => { + onEvent = handler; + return () => {}; + }); + vi.mocked(session.client.followUp).mockRejectedValueOnce( + new Error("queue unavailable"), + ); + const controller = createController(session); + + await controller.connect("task-1", "run-1"); + onEvent({ + type: "queue_update", + timestamp: 1, + steering: ["fix this"], + followUp: ["then summarize"], + }); + + await expect(controller.retry("task-1")).rejects.toThrow( + "queue unavailable", + ); + await controller.retry("task-1"); + + expect(session.client.prompt).toHaveBeenCalledTimes(1); + }); + it("does not restore a captured queue after the task disconnects", async () => { let resolveRetry: () => void = () => {}; const retrying = new Promise((resolve) => { @@ -590,6 +619,67 @@ describe("PiSessionController", () => { }); }); + it("does not retry disconnected cloud sessions after their view unmounts", async () => { + const session = createSession(); + session.retry = vi.fn(async () => {}); + const controller = createController(session); + + await controller.connect("task-1", "run-1"); + controller.store.setState((state) => ({ + sessions: { + ...state.sessions, + "task-1": { + ...state.sessions["task-1"], + connectionState: "disconnected", + cloudStatus: "in_progress", + error: { + id: "connection-error", + scope: "connection", + kind: "unknown", + title: "Connection failed", + message: "stream dropped", + retryable: true, + limitCause: null, + }, + }, + }, + })); + controller.disconnect("task-1"); + + controller.retryUnhealthyCloudSessions(); + + expect(session.retry).not.toHaveBeenCalled(); + }); + + it("deduplicates concurrent retry requests", async () => { + let resolveRetry: () => void = () => {}; + const session = createSession(); + session.retry = vi.fn( + () => + new Promise((resolve) => { + resolveRetry = resolve; + }), + ); + const controller = createController(session); + + await controller.connect("task-1", "run-1"); + controller.store.setState((state) => ({ + sessions: { + ...state.sessions, + "task-1": { + ...state.sessions["task-1"], + connectionState: "disconnected", + }, + }, + })); + + const first = controller.retry("task-1"); + const second = controller.retry("task-1"); + await vi.waitFor(() => expect(session.retry).toHaveBeenCalledOnce()); + resolveRetry(); + await Promise.all([first, second]); + }); + it("uses the live bash operation without reloading native history", async () => { const session = createSession(); const controller = createController(session); diff --git a/packages/core/src/pi-runtime/piSessionController.ts b/packages/core/src/pi-runtime/piSessionController.ts index b63465c5d0..b163811299 100644 --- a/packages/core/src/pi-runtime/piSessionController.ts +++ b/packages/core/src/pi-runtime/piSessionController.ts @@ -118,6 +118,7 @@ export class PiSessionController { private readonly queuesToRestore = new Map(); private readonly cancelAuthRestoration = new Map void>(); private readonly taskRunIds = new Map(); + private readonly activeTaskIds = new Set(); constructor( @inject(PI_SESSION_PROVIDER) private readonly provider: PiSessionProvider, @@ -128,6 +129,7 @@ export class PiSessionController { ) {} ensureConnected(taskId: string, taskRunId?: string): Promise { + this.activeTaskIds.add(taskId); this.bindTaskRun(taskId, taskRunId); this.ensureSubscription(taskId); @@ -166,6 +168,7 @@ export class PiSessionController { } connect(taskId: string, taskRunId?: string): Promise { + this.activeTaskIds.add(taskId); this.bindTaskRun(taskId, taskRunId); this.ensureSubscription(taskId); @@ -192,17 +195,21 @@ export class PiSessionController { this.liveEvents.delete(taskId); this.queueRevisions.delete(taskId); this.queuesToRestore.delete(taskId); + this.activeTaskIds.delete(taskId); } async retry(taskId: string): Promise { - const taskRunId = this.taskRunIds.get(taskId); - this.captureQueueForRestore(taskId); - const session = await this.getPiSession(taskId); + if (this.getSession(taskId).connectionState === "connecting") { + return; + } this.updateSession(taskId, { connectionState: "connecting", error: undefined, }); + const taskRunId = this.taskRunIds.get(taskId); + this.captureQueueForRestore(taskId); try { + const session = await this.getPiSession(taskId); await session.retry?.(); this.resetTransport(taskId); await this.ensureConnected(taskId, taskRunId); @@ -224,8 +231,10 @@ export class PiSessionController { } async restart(taskId: string): Promise { + if (this.getSession(taskId).connectionState === "connecting") { + return; + } const taskRunId = this.taskRunIds.get(taskId); - this.captureQueueForRestore(taskId); if (!taskRunId) { await this.retry(taskId); return; @@ -235,6 +244,7 @@ export class PiSessionController { connectionState: "connecting", error: undefined, }); + this.captureQueueForRestore(taskId); try { const resumedRun = await this.taskService.resumeCloudPiRun( taskId, @@ -252,6 +262,7 @@ export class PiSessionController { this.store.getState().sessions, )) { if ( + this.activeTaskIds.has(taskId) && session.cloudStatus !== undefined && session.error?.retryable && (session.connectionState === "disconnected" || @@ -894,6 +905,7 @@ export class PiSessionController { mode: "follow_up" as const, })), ]; + this.queuesToRestore.delete(taskId); if (!status.isStreaming) { const first = messages.shift(); if (first) { @@ -908,7 +920,6 @@ export class PiSessionController { await session.client.followUp(message.content); } } - this.queuesToRestore.delete(taskId); } private async refreshQueue( diff --git a/packages/ui/src/features/pi-sessions/PiSessionView.tsx b/packages/ui/src/features/pi-sessions/PiSessionView.tsx index e2b42638a2..702f66919f 100644 --- a/packages/ui/src/features/pi-sessions/PiSessionView.tsx +++ b/packages/ui/src/features/pi-sessions/PiSessionView.tsx @@ -1,5 +1,6 @@ import { contentToXml, + isContentEmpty, xmlToContent, } from "@posthog/core/message-editor/content"; import { PI_SESSION_CONTROLLER } from "@posthog/core/pi-runtime/identifiers"; @@ -37,6 +38,7 @@ import { useWorkspace } from "@posthog/ui/features/workspace/useWorkspace"; import { useConnectivity } from "@posthog/ui/hooks/useConnectivity"; import { toast } from "@posthog/ui/primitives/toast"; import { TaskDetailSkeleton } from "@posthog/ui/router/routeSkeletons"; +import { logger } from "@posthog/ui/shell/logger"; import { Box, Flex } from "@radix-ui/themes"; import { type ReactElement, useCallback, useEffect, useRef } from "react"; import { useStore } from "zustand"; @@ -47,6 +49,8 @@ import { PiThinkingLevelSelector, } from "./PiSessionControls"; +const log = logger.scope("pi-session-view"); + interface PiSessionViewProps { taskId: string; taskRunId?: string; @@ -126,6 +130,7 @@ export function PiSessionView({ taskId, taskRunId }: PiSessionViewProps) { const handleControllerError = useCallback( (error: unknown, fallback: string) => { + log.error(fallback, error); if (error instanceof PiOperationError) { return; } @@ -267,7 +272,10 @@ export function PiSessionView({ taskId, taskRunId }: PiSessionViewProps) { if (!failure || failure.scope !== "operation") { return; } - if (failure.recoveryPrompt) { + if ( + failure.recoveryPrompt && + isContentEmpty(useDraftStore.getState().drafts[taskId] ?? null) + ) { draftActions.setPendingContent( taskId, xmlToContent(failure.recoveryPrompt), diff --git a/packages/ui/src/features/sessions/components/chat-thread/ChatThread.tsx b/packages/ui/src/features/sessions/components/chat-thread/ChatThread.tsx index 1652d2b84b..1a67e5eec3 100644 --- a/packages/ui/src/features/sessions/components/chat-thread/ChatThread.tsx +++ b/packages/ui/src/features/sessions/components/chat-thread/ChatThread.tsx @@ -35,7 +35,10 @@ import { SHORTCUTS } from "@posthog/ui/features/command/keyboard-shortcuts"; import { useSmoothedText } from "@posthog/ui/features/editor/components/useSmoothedText"; import { useFeatureFlag } from "@posthog/ui/features/feature-flags/useFeatureFlag"; import { usePanelLayoutStore } from "@posthog/ui/features/panels/panelLayoutStore"; -import type { ConversationItem } from "@posthog/ui/features/sessions/components/buildConversationItems"; +import type { + BuildResult, + ConversationItem, +} from "@posthog/ui/features/sessions/components/buildConversationItems"; import { CloudArtifactDownloads } from "@posthog/ui/features/sessions/components/CloudArtifactDownloads"; import { ChatMarkdown, @@ -961,6 +964,7 @@ interface SharedChatThreadProps { task?: Task; taskId?: string; usage?: ContextUsage | null; + footerState?: Omit; } export interface ChatThreadProps extends SharedChatThreadProps { @@ -972,7 +976,10 @@ export interface AcpChatThreadProps extends SharedChatThreadProps { } export function ChatThread({ events, ...props }: ChatThreadProps) { - const { items } = useAgentConversationItems(events, props.isPromptPending); + const { items, ...footerState } = useAgentConversationItems( + events, + props.isPromptPending, + ); return ( ); } @@ -1014,6 +1022,7 @@ function ChatThreadRenderer({ task, taskId, usage, + footerState, promptRecallRef, }: ChatThreadRendererProps) { const diffWorkerFactory = useService(DIFF_WORKER_FACTORY); @@ -1140,6 +1149,7 @@ function ChatThreadRenderer({ task={task} taskId={taskId} usage={usage} + footerState={footerState} /> ); diff --git a/packages/ui/src/features/sessions/components/chat-thread/ChatThreadFooter.tsx b/packages/ui/src/features/sessions/components/chat-thread/ChatThreadFooter.tsx index 07143d07d3..80af68e991 100644 --- a/packages/ui/src/features/sessions/components/chat-thread/ChatThreadFooter.tsx +++ b/packages/ui/src/features/sessions/components/chat-thread/ChatThreadFooter.tsx @@ -1,6 +1,7 @@ import type { ContextUsage } from "@posthog/core/sessions/contextUsage"; import type { AcpMessage } from "@posthog/shared"; import type { Task } from "@posthog/shared/domain-types"; +import type { BuildResult } from "@posthog/ui/features/sessions/components/buildConversationItems"; import { SessionFooter } from "@posthog/ui/features/sessions/components/SessionFooter"; import { useContextUsage } from "@posthog/ui/features/sessions/hooks/useContextUsage"; import { useConversationItems } from "@posthog/ui/features/sessions/hooks/useConversationItems"; @@ -18,6 +19,7 @@ interface ChatThreadFooterProps { task?: Task; taskId?: string; usage?: ContextUsage | null; + footerState?: Omit; } /** @@ -37,12 +39,21 @@ export function ChatThreadFooter({ task, taskId, usage, + footerState, }: ChatThreadFooterProps) { const showDebugLogs = useSettingsStore((s) => s.debugLogsCloudRuns); const eventContextUsage = useContextUsage(events); const contextUsage = usage === undefined ? eventContextUsage : usage; - const { lastTurnInfo, isCompacting, completedToolCallCount } = - useConversationItems(events, isPromptPending, { showDebugLogs }); + const eventFooterState = useConversationItems(events, isPromptPending, { + showDebugLogs, + }); + const lastTurnInfo = + footerState?.lastTurnInfo ?? eventFooterState.lastTurnInfo; + const isCompacting = + footerState?.isCompacting ?? eventFooterState.isCompacting; + const completedToolCallCount = + footerState?.completedToolCallCount ?? + eventFooterState.completedToolCallCount; const pendingPermissions = usePendingPermissionsForTask(taskId ?? ""); const queuedCount = useQueuedMessagesForTask(taskId).length; const session = useSessionForTask(taskId); diff --git a/packages/ui/src/features/settings/settingsStore.test.ts b/packages/ui/src/features/settings/settingsStore.test.ts index dc7c3ecb76..0cedfa6639 100644 --- a/packages/ui/src/features/settings/settingsStore.test.ts +++ b/packages/ui/src/features/settings/settingsStore.test.ts @@ -84,6 +84,25 @@ describe("feature settingsStore cloud selections", () => { expect(useSettingsStore.getState().lastUsedAgentRuntime).toBe("pi"); }); + it("persists Pi and ACP model selections independently", async () => { + const settings = useSettingsStore.getState(); + settings.setLastUsedModel("claude-sonnet-4-5"); + settings.setLastUsedPiModel("claude-opus-4-8"); + + expect(useSettingsStore.getState()).toMatchObject({ + lastUsedModel: "claude-sonnet-4-5", + lastUsedPiModel: "claude-opus-4-8", + }); + await waitForPersistedWrite(); + + const lastCall = setItem.mock.calls[setItem.mock.calls.length - 1]; + const persisted = JSON.parse(lastCall[1]); + expect(persisted.state).toMatchObject({ + lastUsedModel: "claude-sonnet-4-5", + lastUsedPiModel: "claude-opus-4-8", + }); + }); + it("persists the last used cloud repository", async () => { useSettingsStore.getState().setLastUsedCloudRepository("posthog/posthog"); diff --git a/packages/ui/src/features/settings/settingsStore.ts b/packages/ui/src/features/settings/settingsStore.ts index 197fe99281..da38be2ecd 100644 --- a/packages/ui/src/features/settings/settingsStore.ts +++ b/packages/ui/src/features/settings/settingsStore.ts @@ -111,6 +111,7 @@ interface SettingsStore { lastUsedAgentRuntime: AgentRuntime; lastUsedAdapter: AgentAdapter; lastUsedModel: string | null; + lastUsedPiModel: string | null; lastUsedReasoningEffort: string | null; lastUsedCloudRepository: string | null; cachedCloudRepositoryMap: Record; @@ -133,6 +134,7 @@ interface SettingsStore { setLastUsedAgentRuntime: (runtime: AgentRuntime) => void; setLastUsedAdapter: (adapter: AgentAdapter) => void; setLastUsedModel: (model: string) => void; + setLastUsedPiModel: (model: string) => void; setLastUsedReasoningEffort: (effort: string) => void; setLastUsedCloudRepository: (repo: string | null) => void; setCachedCloudRepositoryMap: ( @@ -306,6 +308,7 @@ export const useSettingsStore = create()( lastUsedAgentRuntime: "acp", lastUsedAdapter: "claude", lastUsedModel: null, + lastUsedPiModel: null, lastUsedReasoningEffort: null, lastUsedCloudRepository: null, cachedCloudRepositoryMap: {}, @@ -325,6 +328,7 @@ export const useSettingsStore = create()( set({ lastUsedAgentRuntime: runtime }), setLastUsedAdapter: (adapter) => set({ lastUsedAdapter: adapter }), setLastUsedModel: (model) => set({ lastUsedModel: model }), + setLastUsedPiModel: (model) => set({ lastUsedPiModel: model }), setLastUsedReasoningEffort: (effort) => set({ lastUsedReasoningEffort: effort }), setLastUsedCloudRepository: (repo) => @@ -534,6 +538,7 @@ export const useSettingsStore = create()( lastUsedAgentRuntime: state.lastUsedAgentRuntime, lastUsedAdapter: state.lastUsedAdapter, lastUsedModel: state.lastUsedModel, + lastUsedPiModel: state.lastUsedPiModel, lastUsedReasoningEffort: state.lastUsedReasoningEffort, lastUsedCloudRepository: state.lastUsedCloudRepository, cachedCloudRepositoryMap: state.cachedCloudRepositoryMap, diff --git a/packages/ui/src/features/task-detail/components/TaskInput.tsx b/packages/ui/src/features/task-detail/components/TaskInput.tsx index 06f270c347..78bbd72511 100644 --- a/packages/ui/src/features/task-detail/components/TaskInput.tsx +++ b/packages/ui/src/features/task-detail/components/TaskInput.tsx @@ -220,8 +220,9 @@ export function TaskInput({ defaultInitialTaskMode, lastUsedInitialTaskMode, setLastUsedReasoningEffort, - lastUsedModel, setLastUsedModel, + lastUsedPiModel, + setLastUsedPiModel, _hasHydrated: settingsHydrated, } = useSettingsStore(); const { data: skills } = useSkills(); @@ -757,7 +758,7 @@ export function TaskInput({ thoughtOption?.type === "select" ? thoughtOption.currentValue : undefined; const currentPiModel = piModelCatalog.find((model) => model.id === selectedPiModelId) ?? - piModelCatalog.find((model) => model.id === lastUsedModel) ?? + piModelCatalog.find((model) => model.id === lastUsedPiModel) ?? piModelCatalog[0]; const piThinkingLevels = currentPiModel?.thinkingLevels ?? []; const currentPiThinkingLevel = piThinkingLevels.includes( @@ -1049,9 +1050,9 @@ export function TaskInput({ const handlePiModelChange = useCallback( (model: PiModelSelection) => { setSelectedPiModelId(model.id); - setLastUsedModel(model.id); + setLastUsedPiModel(model.id); }, - [setLastUsedModel], + [setLastUsedPiModel], ); const handlePiThinkingLevelChange = useCallback((level: PiThinkingLevel) => { diff --git a/packages/workspace-server/src/services/pi-session/pi-session.test.ts b/packages/workspace-server/src/services/pi-session/pi-session.test.ts index 7432e16c2a..b15fb58dd3 100644 --- a/packages/workspace-server/src/services/pi-session/pi-session.test.ts +++ b/packages/workspace-server/src/services/pi-session/pi-session.test.ts @@ -176,12 +176,16 @@ describe("PiSessionService start", () => { }); describe("PiSessionService RPC request pinning", () => { - it("keeps a session pinned until every concurrent generic command settles", async () => { + it("keeps a session pinned until command and queue requests settle", async () => { vi.stubEnv("POSTHOG_CODE_PI_HOT_POOL_SIZE", "1"); let timestamp = 0; vi.spyOn(Date, "now").mockImplementation(() => timestamp++); const requestResolvers: Array<(response: RpcResponse) => void> = []; + let resolveQueue: (queue: { + steering: string[]; + followUp: string[]; + }) => void = () => {}; const firstClient = { start: vi.fn().mockResolvedValue(undefined), stop: vi.fn().mockResolvedValue(undefined), @@ -195,6 +199,12 @@ describe("PiSessionService RPC request pinning", () => { requestResolvers.push(resolve); }), ), + getQueue: vi.fn( + () => + new Promise<{ steering: string[]; followUp: string[] }>((resolve) => { + resolveQueue = resolve; + }), + ), } as unknown as PiRpcClient; const secondClient = { start: vi.fn().mockResolvedValue(undefined), @@ -255,7 +265,7 @@ describe("PiSessionService RPC request pinning", () => { type: "bash", command: "sleep 1", }); - const compactRequest = service.request("first", { type: "compact" }); + const queueRequest = service.getQueue("first"); await service.resume({ taskId: "second", cwd: "/tmp" }); expect(firstClient.stop).not.toHaveBeenCalled(); @@ -265,8 +275,8 @@ describe("PiSessionService RPC request pinning", () => { await vi.waitFor(() => expect(secondClient.stop).toHaveBeenCalledOnce()); expect(firstClient.stop).not.toHaveBeenCalled(); - requestResolvers[1](successfulResponse("compact")); - await compactRequest; + resolveQueue({ steering: [], followUp: [] }); + await queueRequest; expect(firstClient.stop).not.toHaveBeenCalled(); await service.resume({ taskId: "third", cwd: "/tmp" }); diff --git a/packages/workspace-server/src/services/pi-session/pi-session.ts b/packages/workspace-server/src/services/pi-session/pi-session.ts index b7aa4634a1..bef202c952 100644 --- a/packages/workspace-server/src/services/pi-session/pi-session.ts +++ b/packages/workspace-server/src/services/pi-session/pi-session.ts @@ -170,11 +170,8 @@ export class PiSessionService extends TypedEventEmitter { await this.startSession(input.taskId, client, session, async () => {}); } - async request(taskId: string, command: RpcCommand): Promise { - const session = this.requireSession(taskId); - session.activeRequestCount += 1; - - try { + request(taskId: string, command: RpcCommand): Promise { + return this.withActiveRequest(taskId, async (session) => { const response = await session.runtime.sendCommand(command); if ( @@ -187,18 +184,21 @@ export class PiSessionService extends TypedEventEmitter { } return response; - } finally { - session.activeRequestCount -= 1; - void this.enforceHotPoolLimit(); - } + }); } getQueue(taskId: string): Promise { - return this.requireSession(taskId).client.getQueue(); + return this.withActiveRequest(taskId, (session) => + session.client.getQueue(), + ); } clearQueue(taskId: string): Promise { - return this.requireSession(taskId).client.clearQueue(); + return this.withActiveRequest(taskId, async (session) => { + const queue = await session.client.clearQueue(); + session.runtime.clearPendingQueuedUserMessages(); + return queue; + }); } async stop(taskId: string): Promise { @@ -365,6 +365,21 @@ export class PiSessionService extends TypedEventEmitter { }); } + private async withActiveRequest( + taskId: string, + operation: (session: ManagedPiSession) => Promise, + ): Promise { + const session = this.requireSession(taskId); + session.activeRequestCount += 1; + + try { + return await operation(session); + } finally { + session.activeRequestCount -= 1; + void this.enforceHotPoolLimit(); + } + } + private requireSession(taskId: string): ManagedPiSession { const session = this.sessions.get(taskId); From ad8223961544ba95161708ff1f99670bbb6f99e2 Mon Sep 17 00:00:00 2001 From: JonathanLab Date: Tue, 28 Jul 2026 19:52:35 +0200 Subject: [PATCH 30/45] fix(pi): recover stale cloud sessions --- packages/agent/src/posthog-api.test.ts | 21 ++++ packages/agent/src/posthog-api.ts | 3 + packages/agent/src/server/pi-agent-server.ts | 3 + .../pi-runtime/cloudPiSessionClient.test.ts | 12 +- .../src/pi-runtime/cloudPiSessionClient.ts | 8 +- .../pi-runtime/piSessionController.test.ts | 118 +++++++++++++++++- .../src/pi-runtime/piSessionController.ts | 79 ++++++++++-- .../features/pi-sessions/PiSessionView.tsx | 16 +-- 8 files changed, 242 insertions(+), 18 deletions(-) diff --git a/packages/agent/src/posthog-api.test.ts b/packages/agent/src/posthog-api.test.ts index 200a5d90a5..cbc600ebf4 100644 --- a/packages/agent/src/posthog-api.test.ts +++ b/packages/agent/src/posthog-api.test.ts @@ -198,6 +198,27 @@ describe("PostHogAPIClient", () => { expect(mockFetch).not.toHaveBeenCalled(); }); + it("treats a missing stored task session object as empty", async () => { + const client = new PostHogAPIClient({ + apiUrl: "https://app.posthog.com", + getApiKey: vi.fn().mockResolvedValue("token"), + projectId: 7, + }); + mockFetch.mockResolvedValueOnce({ + ok: false, + status: 404, + statusText: "Not Found", + }); + + await expect( + client.downloadTaskSession({ + id: "session-1", + download_url: "https://storage.example/missing.jsonl", + content_sha256: "old-hash", + }), + ).resolves.toBe(""); + }); + it("surfaces an uncertain task session replacement without retrying", async () => { const client = new PostHogAPIClient({ apiUrl: "https://app.posthog.com", diff --git a/packages/agent/src/posthog-api.ts b/packages/agent/src/posthog-api.ts index 126467ce42..99b2574ff0 100644 --- a/packages/agent/src/posthog-api.ts +++ b/packages/agent/src/posthog-api.ts @@ -236,6 +236,9 @@ export class PostHogAPIClient { const response = await fetch(access.download_url, { signal: AbortSignal.timeout(30_000), }); + if (response.status === 404) { + return ""; + } if (!response.ok) { throw new Error( `Failed to download task session: [${response.status}] ${response.statusText}`, diff --git a/packages/agent/src/server/pi-agent-server.ts b/packages/agent/src/server/pi-agent-server.ts index 853aeb23b0..66d3185eaf 100644 --- a/packages/agent/src/server/pi-agent-server.ts +++ b/packages/agent/src/server/pi-agent-server.ts @@ -420,6 +420,9 @@ export class PiAgentServer { } }); await client.start(); + if (this.config.reasoningEffort) { + await client.setThinkingLevel(this.config.reasoningEffort); + } const runtimeState = await client.getState(); this.sessionFile = runtimeState.sessionFile ?? restoredSessionFile ?? null; const unsubscribe = () => { diff --git a/packages/core/src/pi-runtime/cloudPiSessionClient.test.ts b/packages/core/src/pi-runtime/cloudPiSessionClient.test.ts index 8a5fa8a457..d34e6d5dd9 100644 --- a/packages/core/src/pi-runtime/cloudPiSessionClient.test.ts +++ b/packages/core/src/pi-runtime/cloudPiSessionClient.test.ts @@ -126,7 +126,7 @@ describe("CloudPiSessionClient", () => { } }); - it("accepts a reconnect snapshot with a start entry as runtime readiness", async () => { + it("waits for a fresh start when a queued resume snapshot contains an old start", async () => { const cloud = createCloudTaskClient(); vi.mocked(cloud.client.sendCommand).mockResolvedValue({ success: true, @@ -150,6 +150,16 @@ describe("CloudPiSessionClient", () => { totalEntryCount: 1, }); + expect(cloud.client.sendCommand).not.toHaveBeenCalled(); + + cloud.sendUpdate({ + taskId: "task-1", + runId: "run-1", + kind: "logs", + newEntries: [{ type: "pi_run_started" }], + totalEntryCount: 2, + }); + await expect(state).resolves.toMatchObject({ isStreaming: false }); expect(cloud.client.sendCommand).toHaveBeenCalledOnce(); }); diff --git a/packages/core/src/pi-runtime/cloudPiSessionClient.ts b/packages/core/src/pi-runtime/cloudPiSessionClient.ts index f0ac8739f4..e38777f775 100644 --- a/packages/core/src/pi-runtime/cloudPiSessionClient.ts +++ b/packages/core/src/pi-runtime/cloudPiSessionClient.ts @@ -119,6 +119,10 @@ export class CloudPiSessionClient implements PiSession { return isTerminalStatus(this.runStatus); } + get taskRunId(): string { + return this.context.runId; + } + get cloudStatus(): TaskRunStatus { return this.runStatus; } @@ -237,8 +241,10 @@ export class CloudPiSessionClient implements PiSession { onError: (error: unknown) => void, onCloudStatus?: (status: TaskRunStatus) => void, ): void { + const snapshotCanProveReadiness = + update.kind === "snapshot" && update.status === "in_progress"; const hasCurrentReadinessEvent = - (update.kind === "logs" || update.kind === "snapshot") && + (update.kind === "logs" || snapshotCanProveReadiness) && update.newEntries.some((entry) => entry.type === "pi_run_started"); if (hasCurrentReadinessEvent) { this.markRuntimeReady(); diff --git a/packages/core/src/pi-runtime/piSessionController.test.ts b/packages/core/src/pi-runtime/piSessionController.test.ts index f52d947a12..3256d5098c 100644 --- a/packages/core/src/pi-runtime/piSessionController.test.ts +++ b/packages/core/src/pi-runtime/piSessionController.test.ts @@ -690,10 +690,46 @@ describe("PiSessionController", () => { expect(session.getConversation).not.toHaveBeenCalled(); }); + it("does not mark direct bash events as assistant streaming", async () => { + let onEvent: (event: AgentConversationEvent) => void = () => {}; + const session = createSession(); + vi.mocked(session.onConversationEvent).mockImplementation((handler) => { + onEvent = handler; + return () => {}; + }); + const controller = createController(session); + + await controller.connect("task-1", "run-1"); + onEvent({ + type: "tool_call_started", + timestamp: 1, + toolCall: { + id: "pi-bash-live-1-1", + title: "printf hello", + kind: "execute", + status: "in_progress", + rawInput: { command: "printf hello" }, + }, + }); + onEvent({ + type: "tool_call_updated", + timestamp: 2, + toolCall: { + id: "pi-bash-live-1-1", + status: "completed", + }, + }); + + expect( + controller.store.getState().sessions["task-1"].status?.isStreaming, + ).toBe(false); + }); + it("resumes a terminal cloud run only when a message is submitted", async () => { const terminalSession = { ...createSession(), resumeRequired: true, + taskRunId: "run-1", }; const resumedSession = createSession(); const provider = { @@ -706,7 +742,7 @@ describe("PiSessionController", () => { const taskService = { resumeCloudPiRun } as unknown as TaskService; const controller = new PiSessionController(provider, taskService); - await controller.connect("task-1", "run-1"); + await controller.connect("task-1"); expect(resumeCloudPiRun).not.toHaveBeenCalled(); @@ -716,6 +752,46 @@ describe("PiSessionController", () => { expect(resumedSession.client.prompt).toHaveBeenCalledWith("continue"); }); + it("resumes and retries a message when the prior sandbox is gone", async () => { + const staleSession = { + ...createSession(), + taskRunId: "run-1", + sendUserMessage: vi.fn(async () => { + throw new Error("No active sandbox for this task run"); + }), + }; + const resumedSession = { + ...createSession(), + sendUserMessage: vi.fn(async () => {}), + }; + const provider = { + get: vi + .fn() + .mockResolvedValueOnce(staleSession) + .mockResolvedValue(resumedSession), + } as PiSessionProvider; + const resumeCloudPiRun = vi.fn(async () => ({ id: "run-2" })); + const taskService = { + prepareCloudPiMessage: vi.fn(async () => ({ + content: "continue", + artifactIds: [], + })), + resumeCloudPiRun, + } as unknown as TaskService; + const controller = new PiSessionController(provider, taskService); + + await controller.connect("task-1"); + await controller.submit("task-1", "continue", false, "steer"); + + expect(resumeCloudPiRun).toHaveBeenCalledWith("task-1", "run-1"); + expect(resumedSession.sendUserMessage).toHaveBeenCalledWith( + "prompt", + "continue", + [], + expect.any(String), + ); + }); + it("keeps a connected transcript usable when a command fails", async () => { const initialEvent: AgentConversationEvent = { type: "user_message", @@ -950,6 +1026,46 @@ describe("PiSessionController", () => { ]); }); + it("does not briefly duplicate retained events during reconnect snapshots", async () => { + const retainedEvent: AgentConversationEvent = { + type: "assistant_message_chunk", + timestamp: 1, + content: { type: "text", text: "retained" }, + sourceId: "pi-entry-1:0", + }; + let onEvent: (event: AgentConversationEvent) => void = () => {}; + const session = createSession(); + vi.mocked(session.onConversationEvent).mockImplementation((handler) => { + onEvent = handler; + return () => {}; + }); + const controller = createController(session); + + await controller.connect("task-1", "run-1"); + onEvent(retainedEvent); + controller.disconnect("task-1"); + + let resolveConversation: (events: AgentConversationEvent[]) => void = + () => {}; + vi.mocked(session.getConversation).mockReturnValue( + new Promise((resolve) => { + resolveConversation = resolve; + }), + ); + const reconnect = controller.connect("task-1", "run-1"); + await vi.waitFor(() => + expect(session.onConversationEvent).toHaveBeenCalledTimes(2), + ); + onEvent(retainedEvent); + + expect(controller.store.getState().sessions["task-1"].events).toEqual([ + retainedEvent, + ]); + + resolveConversation([retainedEvent]); + await reconnect; + }); + it("does not append streamed assistant text already present in native history", async () => { const nativeEvent: AgentConversationEvent = { type: "assistant_message_chunk", diff --git a/packages/core/src/pi-runtime/piSessionController.ts b/packages/core/src/pi-runtime/piSessionController.ts index b163811299..5b2a66e5f4 100644 --- a/packages/core/src/pi-runtime/piSessionController.ts +++ b/packages/core/src/pi-runtime/piSessionController.ts @@ -43,6 +43,7 @@ export interface PiSession { client: PiRemoteRpcClient; readonly resumeRequired?: boolean; readonly cloudStatus?: TaskRunStatus; + readonly taskRunId?: string; retry?(): Promise; getQueue(): Promise; clearQueue(): Promise; @@ -386,7 +387,9 @@ export class PiSessionController { message, ) : { content: message, artifactIds: [] }; - await session.sendUserMessage( + await this.sendCloudUserMessage( + taskId, + session, commandType, prepared.content, prepared.artifactIds, @@ -641,6 +644,14 @@ export class PiSessionController { return; } + const session = this.getSession(taskId); + if ( + event.sourceId && + session.events.some((existing) => existing.sourceId === event.sourceId) + ) { + return; + } + if (event.type === "runtime_error") { this.recordOperationFailure( taskId, @@ -652,7 +663,6 @@ export class PiSessionController { const liveEvents = [...(this.liveEvents.get(taskId) ?? []), event]; this.liveEvents.set(taskId, liveEvents); - const session = this.getSession(taskId); let status = session.status; if (status && event.type === "runtime_status") { if (event.status === "compacting") { @@ -666,11 +676,16 @@ export class PiSessionController { ); } } + const isDirectBashEvent = + (event.type === "tool_call_started" || + event.type === "tool_call_updated") && + event.toolCall.id.startsWith("pi-bash-"); const hasTurnActivity = - event.type === "assistant_message_chunk" || - event.type === "assistant_thought_chunk" || - event.type === "tool_call_started" || - event.type === "tool_call_updated"; + !isDirectBashEvent && + (event.type === "assistant_message_chunk" || + event.type === "assistant_thought_chunk" || + event.type === "tool_call_started" || + event.type === "tool_call_updated"); if (status && hasTurnActivity) { status = { ...status, isStreaming: true }; } @@ -1020,9 +1035,50 @@ export class PiSessionController { this.updateSession(taskId, { status }); } + private async sendCloudUserMessage( + taskId: string, + session: PiSession, + type: "prompt" | "steer" | "follow_up", + content: string, + artifactIds: string[], + messageId: string, + ): Promise { + if (!session.sendUserMessage) { + throw new Error("Cloud Pi session cannot send messages"); + } + + try { + await session.sendUserMessage(type, content, artifactIds, messageId); + return; + } catch (error) { + const message = error instanceof Error ? error.message : String(error); + const taskRunId = this.taskRunIds.get(taskId) ?? session.taskRunId; + if (!taskRunId || !message.includes("No active sandbox")) { + throw error; + } + + const resumedRun = await this.taskService.resumeCloudPiRun( + taskId, + taskRunId, + ); + this.resetTransport(taskId); + await this.ensureConnected(taskId, resumedRun.id); + const resumedSession = await this.getPiSession(taskId); + if (!resumedSession.sendUserMessage) { + throw new Error("Resumed cloud Pi session cannot send messages"); + } + await resumedSession.sendUserMessage( + type, + content, + artifactIds, + messageId, + ); + } + } + private async getWritablePiSession(taskId: string): Promise { const session = await this.getPiSession(taskId); - const taskRunId = this.taskRunIds.get(taskId); + const taskRunId = this.taskRunIds.get(taskId) ?? session.taskRunId; if (!session.resumeRequired || !taskRunId) { return session; } @@ -1064,7 +1120,14 @@ export class PiSessionController { return existing; } - const session = this.provider.get(taskId, this.taskRunIds.get(taskId)); + const session = this.provider + .get(taskId, this.taskRunIds.get(taskId)) + .then((resolved) => { + if (resolved.taskRunId && !this.taskRunIds.has(taskId)) { + this.taskRunIds.set(taskId, resolved.taskRunId); + } + return resolved; + }); this.sessions.set(taskId, session); void session.catch(() => { if (this.sessions.get(taskId) === session) { diff --git a/packages/ui/src/features/pi-sessions/PiSessionView.tsx b/packages/ui/src/features/pi-sessions/PiSessionView.tsx index 702f66919f..9cd84670ed 100644 --- a/packages/ui/src/features/pi-sessions/PiSessionView.tsx +++ b/packages/ui/src/features/pi-sessions/PiSessionView.tsx @@ -84,7 +84,6 @@ export function PiSessionView({ taskId, taskRunId }: PiSessionViewProps) { return () => piSessionController.disconnect(taskId); }, [piSessionController, taskId, taskRunId]); - const sessionAvailable = session?.connectionState === "connected"; const status = session?.status; const isStreaming = status?.isStreaming ?? false; const isCompacting = status?.isCompacting ?? false; @@ -314,6 +313,8 @@ export function PiSessionView({ taskId, taskRunId }: PiSessionViewProps) { const hasTranscript = session.events.some( (event) => event.type !== "progress", ); + const sessionAvailable = + session.connectionState === "connected" || hasTranscript; if (isConnecting && !hasTranscript) { return ( @@ -354,11 +355,15 @@ export function PiSessionView({ taskId, taskRunId }: PiSessionViewProps) { const controlsPending = status ? isStreaming || isBashRunning : false; const hasQueuedMessage = session.queue.steering.length + session.queue.followUp.length > 0; - let modelSelector: ReactElement = ; + let modelSelector: ReactElement = ( + + ); let reasoningSelector: ReactElement | null = ( - + + ); + let messagingModeToggle: ReactElement = ( + ); - let messagingModeToggle: ReactElement = ; if (status && session.modelsLoaded) { modelSelector = ( @@ -401,9 +406,6 @@ export function PiSessionView({ taskId, taskRunId }: PiSessionViewProps) { {isAuthRestoring && ( )} - {isConnecting && hasTranscript && !isAuthRestoring && ( - - )} {connectionError && hasTranscript && ( Date: Tue, 28 Jul 2026 20:15:49 +0200 Subject: [PATCH 31/45] refactor(pi): remove frontend runtime heuristics --- .../translatePiConversation.test.ts | 31 ++++++++- .../conversation/translatePiConversation.ts | 26 ++++--- .../agent/src/server/pi-agent-server.test.ts | 36 ++++++++++ packages/agent/src/server/pi-agent-server.ts | 69 ++++++++++++++++++- .../pi-runtime/cloudPiSessionClient.test.ts | 35 ++++++++++ .../src/pi-runtime/cloudPiSessionClient.ts | 31 ++++++++- .../pi-runtime/piSessionController.test.ts | 2 + .../src/pi-runtime/piSessionController.ts | 2 +- .../src/task-detail/taskCreationEffects.ts | 2 + .../core/src/task-detail/taskService.test.ts | 31 +++++++++ packages/core/src/task-detail/taskService.ts | 4 +- packages/shared/src/agent-conversation.ts | 1 + .../task-detail/taskCreationEffectsImpl.ts | 15 ++++ 13 files changed, 268 insertions(+), 17 deletions(-) diff --git a/packages/agent/src/pi/conversation/translatePiConversation.test.ts b/packages/agent/src/pi/conversation/translatePiConversation.test.ts index 44453dd9de..4db19cb58e 100644 --- a/packages/agent/src/pi/conversation/translatePiConversation.test.ts +++ b/packages/agent/src/pi/conversation/translatePiConversation.test.ts @@ -1,5 +1,5 @@ import type { AssistantMessage } from "@earendil-works/pi-ai"; -import { describe, expect, it } from "vitest"; +import { describe, expect, it, vi } from "vitest"; import { createPiConversationTranslator } from "./translatePiConversation"; function assistant( @@ -99,7 +99,9 @@ describe("createPiConversationTranslator", () => { ).toEqual([]); }); - it("completes a turn using the latest runtime timestamp", () => { + it("completes a turn using the settlement time", () => { + vi.useFakeTimers(); + vi.setSystemTime(30); const translator = createPiConversationTranslator(); const laterMessage = assistant( [{ type: "text", text: "later" }], @@ -116,8 +118,9 @@ describe("createPiConversationTranslator", () => { translator.translateEvent({ type: "message_end", message: earlierMessage }); expect(translator.translateEvent({ type: "agent_settled" })).toEqual([ - { type: "turn_completed", timestamp: 20 }, + { type: "turn_completed", timestamp: 30 }, ]); + vi.useRealTimers(); }); it("translates retry lifecycle without rendering transient runtime errors", () => { @@ -310,6 +313,7 @@ describe("createPiConversationTranslator", () => { kind: "execute", status: "in_progress", rawInput: { command: "pwd" }, + origin: "user_shell", }, }, { @@ -319,6 +323,7 @@ describe("createPiConversationTranslator", () => { id: "pi-bash-20", status: "completed", rawOutput: "/tmp/project", + origin: "user_shell", content: [ { type: "content", @@ -411,6 +416,26 @@ describe("createPiConversationTranslator", () => { ]); }); + it("throttles direct bash output by encoded byte size", () => { + const translator = createPiConversationTranslator(); + translator.beginDirectBash("unicode-output"); + + expect( + translator.translateEvent({ + type: "bash_execution_update", + id: "req_1", + delta: "🙂".repeat(1_024), + }), + ).toHaveLength(1); + expect( + translator.translateEvent({ + type: "bash_execution_update", + id: "req_1", + delta: "x", + }), + ).toEqual([]); + }); + it("preserves streamed direct bash output when the command fails", () => { const translator = createPiConversationTranslator(); const [started] = translator.beginDirectBash("failing-command"); diff --git a/packages/agent/src/pi/conversation/translatePiConversation.ts b/packages/agent/src/pi/conversation/translatePiConversation.ts index ce0f906259..5bc20647ab 100644 --- a/packages/agent/src/pi/conversation/translatePiConversation.ts +++ b/packages/agent/src/pi/conversation/translatePiConversation.ts @@ -1,3 +1,4 @@ +import { Buffer } from "node:buffer"; import type { AssistantMessage, Message } from "@earendil-works/pi-ai"; import type { AgentSessionEvent } from "@earendil-works/pi-coding-agent"; import type { AgentConversationEvent } from "@posthog/shared"; @@ -31,6 +32,7 @@ function customMessageEvents(message: AgentMessage): AgentConversationEvent[] { kind: "execute", status: "in_progress", rawInput: { command: message.command }, + origin: "user_shell", }, }, { @@ -40,6 +42,7 @@ function customMessageEvents(message: AgentMessage): AgentConversationEvent[] { id, status: failed ? "failed" : "completed", rawOutput: message.output, + origin: "user_shell", content: message.output ? [ { @@ -114,8 +117,9 @@ export function createPiConversationTranslator(): PiConversationTranslator { let directBashSequence = 0; let activeDirectBash: | { - nextOutputSize: number; + nextOutputBytes: number; output: string; + outputBytes: number; startedAt: number; toolCallId: string; } @@ -125,8 +129,9 @@ export function createPiConversationTranslator(): PiConversationTranslator { const startedAt = Date.now(); const toolCallId = `pi-bash-live-${startedAt}-${++directBashSequence}`; activeDirectBash = { - nextOutputSize: 4_096, + nextOutputBytes: 4_096, output: "", + outputBytes: 0, startedAt, toolCallId, }; @@ -141,6 +146,7 @@ export function createPiConversationTranslator(): PiConversationTranslator { kind: "execute", status: "in_progress", rawInput: { command }, + origin: "user_shell", }, }, ]; @@ -164,6 +170,7 @@ export function createPiConversationTranslator(): PiConversationTranslator { id: directBash.toolCallId, status, rawOutput: output, + origin: "user_shell", content: output ? [ { @@ -307,12 +314,13 @@ export function createPiConversationTranslator(): PiConversationTranslator { } directBash.output += event.delta; - if (directBash.output.length >= 4_096) { - if (directBash.output.length < directBash.nextOutputSize) { + directBash.outputBytes += Buffer.byteLength(event.delta, "utf8"); + if (directBash.outputBytes >= 4_096) { + if (directBash.outputBytes < directBash.nextOutputBytes) { return []; } - while (directBash.nextOutputSize <= directBash.output.length) { - directBash.nextOutputSize *= 2; + while (directBash.nextOutputBytes <= directBash.outputBytes) { + directBash.nextOutputBytes *= 2; } } @@ -322,6 +330,7 @@ export function createPiConversationTranslator(): PiConversationTranslator { timestamp: directBash.startedAt, toolCall: { id: directBash.toolCallId, + origin: "user_shell", content: directBash.output ? [ { @@ -485,10 +494,11 @@ export function createPiConversationTranslator(): PiConversationTranslator { if (event.type === "agent_settled") { streamedAssistantTimestamps.clear(); - const timestamp = latestRuntimeTimestamp; + const timestamp = Math.max(Date.now(), latestRuntimeTimestamp); + const hadRuntimeActivity = latestRuntimeTimestamp > 0; latestRuntimeTimestamp = 0; - return timestamp > 0 ? [{ type: "turn_completed", timestamp }] : []; + return hadRuntimeActivity ? [{ type: "turn_completed", timestamp }] : []; } return []; diff --git a/packages/agent/src/server/pi-agent-server.test.ts b/packages/agent/src/server/pi-agent-server.test.ts index c42b815578..c185666512 100644 --- a/packages/agent/src/server/pi-agent-server.test.ts +++ b/packages/agent/src/server/pi-agent-server.test.ts @@ -97,6 +97,42 @@ describe("PiAgentServer", () => { expect(server.pendingEvents[0]).toEqual({ type: "test", index: 100 }); }); + it("coalesces replay and log buffers for repeated tool updates", () => { + const server = new PiAgentServer(config()) as unknown as { + broadcast(event: Record): void; + pendingEvents: Record[]; + pendingLogEntries: Array<{ event?: Record }>; + }; + + server.broadcast({ + type: "pi_event", + event: { + type: "tool_call_updated", + timestamp: 1, + toolCall: { id: "tool-1", content: [{ type: "content" }] }, + }, + }); + server.broadcast({ + type: "pi_event", + event: { + type: "tool_call_updated", + timestamp: 2, + toolCall: { id: "tool-1", status: "completed" }, + }, + }); + + expect(server.pendingEvents).toHaveLength(1); + expect(server.pendingLogEntries).toHaveLength(1); + expect(server.pendingLogEntries[0]?.event).toMatchObject({ + timestamp: 2, + toolCall: { + id: "tool-1", + status: "completed", + content: [{ type: "content" }], + }, + }); + }); + it("flushes long-running conversation logs in bounded batches", async () => { const appendTaskRunLog = vi.fn( async (_taskId: string, _runId: string, _entries: unknown[]) => ({}), diff --git a/packages/agent/src/server/pi-agent-server.ts b/packages/agent/src/server/pi-agent-server.ts index 66d3185eaf..a218fe5d7c 100644 --- a/packages/agent/src/server/pi-agent-server.ts +++ b/packages/agent/src/server/pi-agent-server.ts @@ -60,6 +60,29 @@ const commandSchemas = { type PiCommandMethod = keyof typeof commandSchemas; +function updatedToolCallId( + event: AgentConversationEvent | undefined, +): string | undefined { + return event?.type === "tool_call_updated" ? event.toolCall.id : undefined; +} + +function mergeToolCallUpdate( + previous: AgentConversationEvent | undefined, + next: AgentConversationEvent | undefined, +): AgentConversationEvent | undefined { + if ( + previous?.type !== "tool_call_updated" || + next?.type !== "tool_call_updated" || + previous.toolCall.id !== next.toolCall.id + ) { + return next; + } + return { + ...next, + toolCall: { ...previous.toolCall, ...next.toolCall }, + }; +} + export class PiAgentServer { private readonly app: Hono; private readonly logger = new Logger({ @@ -654,7 +677,7 @@ export class PiAgentServer { private broadcast(event: Record): void { if (event.type === "pi_event" || event.type === "pi_run_started") { - this.pendingLogEntries.push({ + const logEntry: StoredLogEntry = { id: typeof event.id === "string" ? event.id : undefined, type: event.type, timestamp: @@ -663,7 +686,22 @@ export class PiAgentServer { event.type === "pi_event" ? (event.event as AgentConversationEvent) : undefined, - }); + }; + const toolCallId = updatedToolCallId(logEntry.event); + const pendingLogIndex = toolCallId + ? this.pendingLogEntries.findLastIndex( + (entry) => updatedToolCallId(entry.event) === toolCallId, + ) + : -1; + if (pendingLogIndex >= 0) { + const previous = this.pendingLogEntries[pendingLogIndex]; + this.pendingLogEntries[pendingLogIndex] = { + ...logEntry, + event: mergeToolCallUpdate(previous?.event, logEntry.event), + }; + } else { + this.pendingLogEntries.push(logEntry); + } if (this.pendingLogEntries.length > MAX_PENDING_LOG_ENTRIES) { this.pendingLogEntries.splice( 0, @@ -686,7 +724,32 @@ export class PiAgentServer { if (this.session?.sseController) { this.session.sseController.send(event); } else { - this.pendingEvents.push(event); + const toolCallId = updatedToolCallId( + event.type === "pi_event" + ? (event.event as AgentConversationEvent) + : undefined, + ); + const pendingEventIndex = toolCallId + ? this.pendingEvents.findLastIndex( + (pending) => + pending.type === "pi_event" && + updatedToolCallId( + pending.event as AgentConversationEvent | undefined, + ) === toolCallId, + ) + : -1; + if (pendingEventIndex >= 0) { + const previous = this.pendingEvents[pendingEventIndex]; + this.pendingEvents[pendingEventIndex] = { + ...event, + event: mergeToolCallUpdate( + previous?.event as AgentConversationEvent | undefined, + event.event as AgentConversationEvent | undefined, + ), + }; + } else { + this.pendingEvents.push(event); + } if (this.pendingEvents.length > MAX_PENDING_EVENTS) { this.pendingEvents.splice( 0, diff --git a/packages/core/src/pi-runtime/cloudPiSessionClient.test.ts b/packages/core/src/pi-runtime/cloudPiSessionClient.test.ts index d34e6d5dd9..476d2c2259 100644 --- a/packages/core/src/pi-runtime/cloudPiSessionClient.test.ts +++ b/packages/core/src/pi-runtime/cloudPiSessionClient.test.ts @@ -258,6 +258,41 @@ describe("CloudPiSessionClient", () => { expect(cloud.client.sendCommand).not.toHaveBeenCalled(); }); + it("normalizes legacy direct bash events at the cloud boundary", async () => { + const cloud = createCloudTaskClient(); + const session = new CloudPiSessionClient( + cloud.client, + context("completed"), + ); + session.onConversationEvent(vi.fn(), vi.fn()); + + const conversation = session.getConversation(); + cloud.sendUpdate({ + taskId: "task-1", + runId: "run-1", + kind: "snapshot", + status: "completed", + newEntries: [ + { + type: "pi_event", + event: { + type: "tool_call_updated", + timestamp: 1, + toolCall: { id: "pi-bash-1", status: "completed" }, + }, + }, + ], + totalEntryCount: 1, + }); + + await expect(conversation).resolves.toEqual([ + expect.objectContaining({ + type: "tool_call_updated", + toolCall: expect.objectContaining({ origin: "user_shell" }), + }), + ]); + }); + it("loads terminal history from the cloud snapshot without sandbox RPC", async () => { const cloud = createCloudTaskClient(); const session = new CloudPiSessionClient( diff --git a/packages/core/src/pi-runtime/cloudPiSessionClient.ts b/packages/core/src/pi-runtime/cloudPiSessionClient.ts index e38777f775..1f29951487 100644 --- a/packages/core/src/pi-runtime/cloudPiSessionClient.ts +++ b/packages/core/src/pi-runtime/cloudPiSessionClient.ts @@ -338,7 +338,10 @@ export class CloudPiSessionClient implements PiSession { const sourceId = entry.id ?? `${this.context.runId}:log:${firstEntryIndex + index}`; if (entry.type === "pi_event" && entry.event) { - events.push({ ...entry.event, sourceId }); + events.push({ + ...this.normalizeLegacyEvent(entry.event), + sourceId, + }); continue; } @@ -350,6 +353,32 @@ export class CloudPiSessionClient implements PiSession { return events; } + private normalizeLegacyEvent( + event: AgentConversationEvent, + ): AgentConversationEvent { + if ( + event.type === "tool_call_started" && + event.toolCall.origin === undefined && + event.toolCall.id.startsWith("pi-bash-") + ) { + return { + ...event, + toolCall: { ...event.toolCall, origin: "user_shell" }, + }; + } + if ( + event.type === "tool_call_updated" && + event.toolCall.origin === undefined && + event.toolCall.id.startsWith("pi-bash-") + ) { + return { + ...event, + toolCall: { ...event.toolCall, origin: "user_shell" }, + }; + } + return event; + } + private getProgressEvent( entry: StoredLogEntry, ): AgentConversationEvent | null { diff --git a/packages/core/src/pi-runtime/piSessionController.test.ts b/packages/core/src/pi-runtime/piSessionController.test.ts index 3256d5098c..842359ea8e 100644 --- a/packages/core/src/pi-runtime/piSessionController.test.ts +++ b/packages/core/src/pi-runtime/piSessionController.test.ts @@ -709,6 +709,7 @@ describe("PiSessionController", () => { kind: "execute", status: "in_progress", rawInput: { command: "printf hello" }, + origin: "user_shell", }, }); onEvent({ @@ -717,6 +718,7 @@ describe("PiSessionController", () => { toolCall: { id: "pi-bash-live-1-1", status: "completed", + origin: "user_shell", }, }); diff --git a/packages/core/src/pi-runtime/piSessionController.ts b/packages/core/src/pi-runtime/piSessionController.ts index 5b2a66e5f4..b48c79dae5 100644 --- a/packages/core/src/pi-runtime/piSessionController.ts +++ b/packages/core/src/pi-runtime/piSessionController.ts @@ -679,7 +679,7 @@ export class PiSessionController { const isDirectBashEvent = (event.type === "tool_call_started" || event.type === "tool_call_updated") && - event.toolCall.id.startsWith("pi-bash-"); + event.toolCall.origin === "user_shell"; const hasTurnActivity = !isDirectBashEvent && (event.type === "assistant_message_chunk" || diff --git a/packages/core/src/task-detail/taskCreationEffects.ts b/packages/core/src/task-detail/taskCreationEffects.ts index 60981a8d4b..cce2d882ce 100644 --- a/packages/core/src/task-detail/taskCreationEffects.ts +++ b/packages/core/src/task-detail/taskCreationEffects.ts @@ -1,4 +1,5 @@ import type { TaskCreationInput, TaskCreationOutput } from "@posthog/shared"; +import type { TaskRun } from "@posthog/shared/domain-types"; /** * Host-side reactions to a successful task-creation: optimistic workspace @@ -9,4 +10,5 @@ import type { TaskCreationInput, TaskCreationOutput } from "@posthog/shared"; export interface TaskCreationEffects { onWorkspaceCreated(output: TaskCreationOutput): void; onCreateSuccess(output: TaskCreationOutput, input?: TaskCreationInput): void; + onRunResumed(taskId: string, run: TaskRun): void; } diff --git a/packages/core/src/task-detail/taskService.test.ts b/packages/core/src/task-detail/taskService.test.ts index 7b71417c16..9b147bd059 100644 --- a/packages/core/src/task-detail/taskService.test.ts +++ b/packages/core/src/task-detail/taskService.test.ts @@ -100,6 +100,37 @@ describe("TaskService.openTask", () => { }); }); +describe("TaskService.resumeCloudPiRun", () => { + it("publishes the resumed run to host state", async () => { + const run = { + id: "run-2", + task_id: "task-1", + environment: "cloud", + status: "queued", + }; + const api = { + resumeRunInCloud: vi.fn(async () => run), + }; + const effects = { + onRunResumed: vi.fn(), + } as unknown as TaskCreationEffects; + const service = new TaskService( + { + getAuthenticatedClient: vi.fn(async () => api), + } as unknown as ITaskCreationHost, + {} as SessionService, + effects, + {} as PiRunner, + rootLogger, + ); + + await expect(service.resumeCloudPiRun("task-1", "run-1")).resolves.toBe( + run, + ); + expect(effects.onRunResumed).toHaveBeenCalledWith("task-1", run); + }); +}); + describe("TaskService.createTask validation", () => { it("rejects an input with neither content nor a taskDescription", async () => { const result = await makeService().createTask({ diff --git a/packages/core/src/task-detail/taskService.ts b/packages/core/src/task-detail/taskService.ts index 35af7c32df..f9601d463f 100644 --- a/packages/core/src/task-detail/taskService.ts +++ b/packages/core/src/task-detail/taskService.ts @@ -191,7 +191,9 @@ export class TaskService { throw new Error("Not authenticated"); } - return posthogClient.resumeRunInCloud(taskId, taskRunId); + const run = await posthogClient.resumeRunInCloud(taskId, taskRunId); + this.effects.onRunResumed(taskId, run); + return run; } public async openTask( diff --git a/packages/shared/src/agent-conversation.ts b/packages/shared/src/agent-conversation.ts index 545b182e06..e1ec62cd2b 100644 --- a/packages/shared/src/agent-conversation.ts +++ b/packages/shared/src/agent-conversation.ts @@ -108,6 +108,7 @@ export interface AgentToolCall { rawInput?: unknown; rawOutput?: unknown; parentId?: string; + origin?: "agent" | "user_shell"; } interface AgentConversationEventIdentity { diff --git a/packages/ui/src/features/task-detail/taskCreationEffectsImpl.ts b/packages/ui/src/features/task-detail/taskCreationEffectsImpl.ts index de24f5e04f..5340024b44 100644 --- a/packages/ui/src/features/task-detail/taskCreationEffectsImpl.ts +++ b/packages/ui/src/features/task-detail/taskCreationEffectsImpl.ts @@ -5,12 +5,14 @@ import type { TaskCreationOutput, Workspace, } from "@posthog/shared"; +import type { Task, TaskRun } from "@posthog/shared/domain-types"; import { IMPERATIVE_QUERY_CLIENT, type ImperativeQueryClient, } from "../../shell/queryClient"; import { useDraftStore } from "../message-editor/draftStore"; import { useSettingsStore } from "../settings/settingsStore"; +import { taskKeys } from "../tasks/taskKeys"; import { WORKSPACE_QUERY_KEY } from "../workspace/identifiers"; function queryClient(): ImperativeQueryClient { @@ -29,6 +31,19 @@ export const taskCreationEffects: TaskCreationEffects = { void client.invalidateQueries({ queryKey: WORKSPACE_QUERY_KEY }); }, + onRunResumed(taskId: string, run: TaskRun): void { + const client = queryClient(); + client.setQueryData(taskKeys.detail(taskId), (task) => + task ? { ...task, latest_run: run } : task, + ); + client.setQueriesData({ queryKey: taskKeys.lists() }, (tasks) => + tasks?.map((task) => + task.id === taskId ? { ...task, latest_run: run } : task, + ), + ); + void client.invalidateQueries({ queryKey: taskKeys.allSummaries() }); + }, + onCreateSuccess(output: TaskCreationOutput, input?: TaskCreationInput): void { if (!input) return; From c47d64d28517aee6ea6851e4aae16f7283b94bb4 Mon Sep 17 00:00:00 2001 From: JonathanLab Date: Tue, 28 Jul 2026 20:21:50 +0200 Subject: [PATCH 32/45] fix(agent): keep Pi event translation browser-safe --- .../agent/src/pi/conversation/translatePiConversation.ts | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/packages/agent/src/pi/conversation/translatePiConversation.ts b/packages/agent/src/pi/conversation/translatePiConversation.ts index 5bc20647ab..745145623d 100644 --- a/packages/agent/src/pi/conversation/translatePiConversation.ts +++ b/packages/agent/src/pi/conversation/translatePiConversation.ts @@ -1,4 +1,3 @@ -import { Buffer } from "node:buffer"; import type { AssistantMessage, Message } from "@earendil-works/pi-ai"; import type { AgentSessionEvent } from "@earendil-works/pi-coding-agent"; import type { AgentConversationEvent } from "@posthog/shared"; @@ -9,6 +8,8 @@ type AgentMessage = Extract< { type: "message_end" } >["message"]; +const utf8Encoder = new TextEncoder(); + function isMessage(message: AgentMessage): message is Message { return ( message.role === "user" || @@ -314,7 +315,7 @@ export function createPiConversationTranslator(): PiConversationTranslator { } directBash.output += event.delta; - directBash.outputBytes += Buffer.byteLength(event.delta, "utf8"); + directBash.outputBytes += utf8Encoder.encode(event.delta).byteLength; if (directBash.outputBytes >= 4_096) { if (directBash.outputBytes < directBash.nextOutputBytes) { return []; From b69d17d464d63c87e1b054625fbd1fd4b16cc9f6 Mon Sep 17 00:00:00 2001 From: Richard Solomou Date: Fri, 24 Jul 2026 19:18:07 +0300 Subject: [PATCH 33/45] refactor(api-client): extract cloud task transport Generated-By: PostHog Code Task-Id: c1bbe3cf-742b-4b24-bf96-d11a18b4cf22 --- packages/api-client/package.json | 1 - packages/api-client/src/fetcher.test.ts | 39 ++ packages/api-client/src/fetcher.ts | 25 +- packages/api-client/src/index.ts | 6 +- .../src/posthog-client.automations.test.ts | 192 +++++++++ .../api-client/src/posthog-client.test.ts | 349 ++++++++++++++++- packages/api-client/src/posthog-client.ts | 368 +++++++++++++++--- .../api-client/src/task-normalization.test.ts | 115 ++++++ packages/api-client/src/task-normalization.ts | 203 ++++++++++ pnpm-lock.yaml | 12 +- 10 files changed, 1244 insertions(+), 66 deletions(-) create mode 100644 packages/api-client/src/posthog-client.automations.test.ts create mode 100644 packages/api-client/src/task-normalization.test.ts create mode 100644 packages/api-client/src/task-normalization.ts diff --git a/packages/api-client/package.json b/packages/api-client/package.json index bee0ebc547..3a592020a4 100644 --- a/packages/api-client/package.json +++ b/packages/api-client/package.json @@ -25,7 +25,6 @@ "src/**/*" ], "dependencies": { - "@posthog/agent": "workspace:*", "@posthog/shared": "workspace:*" } } diff --git a/packages/api-client/src/fetcher.test.ts b/packages/api-client/src/fetcher.test.ts index c205949418..6e1e17a351 100644 --- a/packages/api-client/src/fetcher.test.ts +++ b/packages/api-client/src/fetcher.test.ts @@ -53,6 +53,45 @@ describe("buildApiFetcher", () => { expect(mockFetch.mock.calls[0][1].headers.get("Authorization")).toBe( "Bearer my-token", ); + expect(mockFetch.mock.calls[0][1].headers.get("User-Agent")).toBe( + "posthog/desktop.hog.dev; version: test", + ); + }); + + it("uses an injected fetch implementation and custom user agent", async () => { + const injectedFetch = vi.fn().mockResolvedValueOnce(ok()); + const fetcher = buildApiFetcher({ + getAccessToken: vi.fn().mockResolvedValue("token"), + refreshAccessToken: vi.fn().mockResolvedValue("new-token"), + appVersion: "1.2.3", + fetch: injectedFetch, + userAgent: "posthog/mobile; version: 1.2.3", + }); + + await fetcher.fetch(mockInput); + + expect(injectedFetch).toHaveBeenCalledTimes(1); + expect(mockFetch).not.toHaveBeenCalled(); + expect(injectedFetch.mock.calls[0][1].headers.get("User-Agent")).toBe( + "posthog/mobile; version: 1.2.3", + ); + }); + + it("omits the user agent when explicitly disabled", async () => { + const injectedFetch = vi.fn().mockResolvedValueOnce(ok()); + const fetcher = buildApiFetcher({ + getAccessToken: vi.fn().mockResolvedValue("token"), + refreshAccessToken: vi.fn().mockResolvedValue("new-token"), + appVersion: "1.2.3", + fetch: injectedFetch, + userAgent: null, + }); + + await fetcher.fetch(mockInput); + + expect(injectedFetch.mock.calls[0][1].headers.has("User-Agent")).toBe( + false, + ); }); it("retries once with a freshly fetched token on 401", async () => { diff --git a/packages/api-client/src/fetcher.ts b/packages/api-client/src/fetcher.ts index 6bf59aa9f8..bc061a3030 100644 --- a/packages/api-client/src/fetcher.ts +++ b/packages/api-client/src/fetcher.ts @@ -1,9 +1,16 @@ import type { createApiClient } from "./generated"; +export type FetchImplementation = ( + input: string | URL | Request, + init?: RequestInit, +) => Promise; + export type ApiFetcherConfig = { getAccessToken: () => Promise; refreshAccessToken: () => Promise; appVersion: string; + fetch?: FetchImplementation; + userAgent?: string | null; }; /** @@ -13,11 +20,13 @@ export type ApiFetcherConfig = { */ export class ApiRequestError extends Error { readonly status: number; + readonly body: unknown; - constructor(status: number, serializedBody: string) { + constructor(status: number, serializedBody: string, body?: unknown) { super(`Failed request: [${status}] ${serializedBody}`); this.name = "ApiRequestError"; this.status = status; + this.body = body; } } @@ -29,7 +38,11 @@ export function requestErrorStatus(error: unknown): number | undefined { export const buildApiFetcher: ( config: ApiFetcherConfig, ) => Parameters[0] = (config) => { - const userAgent = `posthog/desktop.hog.dev; version: ${config.appVersion}`; + const fetchImpl = config.fetch ?? globalThis.fetch; + const userAgent = + config.userAgent === undefined + ? `posthog/desktop.hog.dev; version: ${config.appVersion}` + : config.userAgent; const makeRequest = async ( input: Parameters[0]["fetch"]>[0], @@ -38,7 +51,9 @@ export const buildApiFetcher: ( const headers = new Headers(); headers.set("Authorization", `Bearer ${token}`); headers.set("Content-Type", "application/json"); - headers.set("User-Agent", userAgent); + if (userAgent) { + headers.set("User-Agent", userAgent); + } if (input.urlSearchParams) { input.url.search = input.urlSearchParams.toString(); @@ -59,7 +74,7 @@ export const buildApiFetcher: ( } try { - const response = await fetch(input.url, { + const response = await fetchImpl(input.url, { method: input.method.toUpperCase(), ...(body && { body }), headers, @@ -114,6 +129,7 @@ export const buildApiFetcher: ( throw new ApiRequestError( response.status, JSON.stringify(errorResponse), + errorResponse, ); } } @@ -128,6 +144,7 @@ export const buildApiFetcher: ( throw new ApiRequestError( response.status, JSON.stringify(errorResponse), + errorResponse, ); } diff --git a/packages/api-client/src/index.ts b/packages/api-client/src/index.ts index e6b7c3c639..a18d6ff4b0 100644 --- a/packages/api-client/src/index.ts +++ b/packages/api-client/src/index.ts @@ -1,6 +1,10 @@ import "./generated.augment"; -export { type ApiFetcherConfig, buildApiFetcher } from "./fetcher"; +export { + type ApiFetcherConfig, + buildApiFetcher, + type FetchImplementation, +} from "./fetcher"; export { createApiClient, type Schemas } from "./generated"; export { createLoop, diff --git a/packages/api-client/src/posthog-client.automations.test.ts b/packages/api-client/src/posthog-client.automations.test.ts new file mode 100644 index 0000000000..3d3386ef3a --- /dev/null +++ b/packages/api-client/src/posthog-client.automations.test.ts @@ -0,0 +1,192 @@ +import { beforeEach, describe, expect, it, vi } from "vitest"; +import { + PostHogAPIClient, + TaskAutomationValidationError, +} from "./posthog-client"; + +const automationPayload = { + id: "automation-1", + name: "Daily PRs", + prompt: "Check PRs", + repository: "posthog/posthog", + github_integration: 7, + cron_expression: "0 9 * * *", + timezone: "Europe/London", + template_id: "llm-skill:daily-prs", + enabled: true, + last_run_at: null, + last_run_status: null, + last_task_id: null, + last_task_run_id: null, + last_error: null, + created_at: "2026-07-21T00:00:00Z", + updated_at: "2026-07-21T00:00:00Z", +}; + +function jsonResponse(body: unknown, status = 200): Response { + return new Response(JSON.stringify(body), { + status, + headers: { "Content-Type": "application/json" }, + }); +} + +describe("PostHogAPIClient task automations", () => { + const fetch = vi.fn(); + const client = new PostHogAPIClient( + "https://app.posthog.test", + async () => "access-token", + async () => "refreshed-token", + 42, + { appVersion: "test", fetch }, + ); + + beforeEach(() => { + fetch.mockReset(); + }); + + it("lists automations and normalizes optional response fields", async () => { + const minimalPayload = { + ...automationPayload, + github_integration: undefined, + timezone: undefined, + template_id: undefined, + enabled: undefined, + }; + fetch.mockResolvedValueOnce( + jsonResponse({ + count: 1, + next: null, + previous: null, + results: [minimalPayload], + }), + ); + + await expect(client.listTaskAutomations()).resolves.toEqual([ + expect.objectContaining({ + id: "automation-1", + github_integration: null, + timezone: null, + template_id: null, + enabled: true, + }), + ]); + expect(fetch).toHaveBeenCalledWith( + new URL( + "https://app.posthog.test/api/projects/42/task_automations/?limit=500", + ), + expect.objectContaining({ method: "GET" }), + ); + }); + + it("gets and creates automations through generated endpoints", async () => { + fetch + .mockResolvedValueOnce(jsonResponse(automationPayload)) + .mockResolvedValueOnce(jsonResponse(automationPayload, 201)); + + await expect(client.getTaskAutomation("automation-1")).resolves.toEqual( + automationPayload, + ); + await expect( + client.createTaskAutomation({ + name: "Daily PRs", + prompt: "Check PRs", + repository: "posthog/posthog", + github_integration: 7, + cron_expression: "0 9 * * *", + timezone: "Europe/London", + template_id: "llm-skill:daily-prs", + enabled: true, + }), + ).resolves.toEqual(automationPayload); + + expect(fetch).toHaveBeenNthCalledWith( + 2, + new URL("https://app.posthog.test/api/projects/42/task_automations/"), + expect.objectContaining({ + method: "POST", + body: JSON.stringify({ + name: "Daily PRs", + prompt: "Check PRs", + repository: "posthog/posthog", + github_integration: 7, + cron_expression: "0 9 * * *", + timezone: "Europe/London", + template_id: "llm-skill:daily-prs", + enabled: true, + }), + }), + ); + }); + + it("updates, deletes, and runs automations", async () => { + fetch + .mockResolvedValueOnce( + jsonResponse({ ...automationPayload, enabled: false }), + ) + .mockResolvedValueOnce(new Response(null, { status: 204 })) + .mockResolvedValueOnce(jsonResponse(automationPayload)); + + await expect( + client.updateTaskAutomation("automation-1", { enabled: false }), + ).resolves.toMatchObject({ enabled: false }); + await expect( + client.deleteTaskAutomation("automation-1"), + ).resolves.toBeUndefined(); + await expect(client.runTaskAutomation("automation-1")).resolves.toEqual( + automationPayload, + ); + + expect(fetch).toHaveBeenNthCalledWith( + 1, + new URL( + "https://app.posthog.test/api/projects/42/task_automations/automation-1/", + ), + expect.objectContaining({ + method: "PATCH", + body: JSON.stringify({ enabled: false }), + }), + ); + expect(fetch).toHaveBeenNthCalledWith( + 3, + new URL( + "https://app.posthog.test/api/projects/42/task_automations/automation-1/run/", + ), + expect.objectContaining({ method: "POST" }), + ); + expect(fetch.mock.calls[2]?.[1]?.body).toBeUndefined(); + }); + + it("preserves validation detail, code, and field attribution", async () => { + fetch.mockResolvedValueOnce( + new Response( + JSON.stringify({ + type: "validation_error", + code: "invalid_input", + detail: "Enter a valid cron expression.", + attr: "cron_expression", + }), + { + status: 400, + statusText: "Bad Request", + headers: { "Content-Type": "application/json" }, + }, + ), + ); + + const request = client.createTaskAutomation({ + name: "Daily PRs", + prompt: "Check PRs", + repository: "posthog/posthog", + cron_expression: "not a cron", + timezone: "Europe/London", + }); + + await expect(request).rejects.toBeInstanceOf(TaskAutomationValidationError); + await expect(request).rejects.toMatchObject({ + status: 400, + code: "invalid_input", + attr: "cron_expression", + message: "Enter a valid cron expression.", + }); + }); +}); diff --git a/packages/api-client/src/posthog-client.test.ts b/packages/api-client/src/posthog-client.test.ts index 42a36681cb..da93bc7960 100644 --- a/packages/api-client/src/posthog-client.test.ts +++ b/packages/api-client/src/posthog-client.test.ts @@ -1,8 +1,337 @@ import { describe, expect, it, vi } from "vitest"; import { ApiRequestError } from "./fetcher"; -import { PostHogAPIClient } from "./posthog-client"; +import { CloudCommandError, PostHogAPIClient } from "./posthog-client"; describe("PostHogAPIClient", () => { + it.each([ + "user_message", + "permission_response", + "set_config_option", + "cancel", + ] as const)("sends the %s cloud run command", async (method) => { + const fetch = vi.fn().mockResolvedValue( + new Response(JSON.stringify({ result: { accepted: true } }), { + status: 200, + }), + ); + const client = new PostHogAPIClient( + "https://app.posthog.test", + async () => "token", + async () => "token", + 42, + { fetch }, + ); + + await expect( + client.sendCloudRunCommand("task-1", "run-1", method, { + value: "payload", + }), + ).resolves.toEqual({ accepted: true }); + + expect(fetch).toHaveBeenCalledWith( + new URL( + "https://app.posthog.test/api/projects/42/tasks/task-1/runs/run-1/command/", + ), + expect.objectContaining({ + method: "POST", + body: expect.any(String), + }), + ); + const request = fetch.mock.calls[0][1] as RequestInit; + expect(JSON.parse(request.body as string)).toMatchObject({ + jsonrpc: "2.0", + method, + params: { value: "payload" }, + }); + }); + + it("throws structured cloud command errors for HTTP failures", async () => { + const fetch = vi.fn().mockResolvedValue( + new Response( + JSON.stringify({ error: "No active sandbox for this run" }), + { + status: 409, + statusText: "Conflict", + }, + ), + ); + const client = new PostHogAPIClient( + "https://app.posthog.test", + async () => "token", + async () => "token", + 42, + { fetch }, + ); + + const error = await client + .sendCloudRunCommand("task-1", "run-1", "user_message") + .catch((caught: unknown) => caught); + + expect(error).toMatchObject({ + name: "CloudCommandError", + method: "user_message", + status: 409, + backendError: "No active sandbox for this run", + }); + expect(error).toBeInstanceOf(CloudCommandError); + expect((error as CloudCommandError).isSandboxInactive()).toBe(true); + }); + + it("throws structured cloud command errors for JSON-RPC failures", async () => { + const fetch = vi + .fn() + .mockResolvedValue( + new Response( + JSON.stringify({ error: { message: "Permission request expired" } }), + { status: 200 }, + ), + ); + const client = new PostHogAPIClient( + "https://app.posthog.test", + async () => "token", + async () => "token", + 42, + { fetch }, + ); + + await expect( + client.sendCloudRunCommand("task-1", "run-1", "permission_response"), + ).rejects.toMatchObject({ + method: "permission_response", + status: 200, + backendError: "Permission request expired", + }); + }); + + it("preserves the legacy sendRunCommand result contract", async () => { + const fetch = vi.fn().mockResolvedValue( + new Response(JSON.stringify({ error: "Run is unavailable" }), { + status: 503, + }), + ); + const client = new PostHogAPIClient( + "https://app.posthog.test", + async () => "token", + async () => "token", + 42, + { fetch }, + ); + + await expect( + client.sendRunCommand("task-1", "run-1", "set_config_option"), + ).resolves.toEqual({ + success: false, + error: "Cloud command 'set_config_option' failed: 503 Run is unavailable", + }); + }); + + it("cancels a cloud task run with an optional reason", async () => { + const fetch = vi + .fn() + .mockResolvedValue( + new Response(JSON.stringify({ status: "cancelled" }), { status: 200 }), + ); + const client = new PostHogAPIClient( + "https://app.posthog.test", + async () => "token", + async () => "token", + 42, + { fetch }, + ); + + await expect( + client.cancelTaskRun("task-1", "run-1", "user requested"), + ).resolves.toEqual({ status: "cancelled" }); + + expect(fetch).toHaveBeenCalledWith( + new URL( + "https://app.posthog.test/api/projects/42/tasks/task-1/runs/run-1/cancel/", + ), + expect.objectContaining({ + method: "POST", + body: JSON.stringify({ reason: "user requested" }), + }), + ); + }); + + it("cancels a cloud task run with an empty body by default", async () => { + const fetch = vi + .fn() + .mockResolvedValue(new Response(null, { status: 204 })); + const client = new PostHogAPIClient( + "https://app.posthog.test", + async () => "token", + async () => "token", + 42, + { fetch }, + ); + + await expect(client.cancelTaskRun("task-1", "run-1")).resolves.toEqual({}); + + const request = fetch.mock.calls[0][1] as RequestInit; + expect(request.body).toBe(JSON.stringify({})); + }); + + it("builds cloud task config from the authenticated gateway catalog", async () => { + const fetch = vi.fn().mockResolvedValue( + new Response( + JSON.stringify({ + data: [ + { + id: "claude-opus-4-8", + owned_by: "anthropic", + context_window: 200000, + supports_streaming: true, + supports_vision: true, + allowed: true, + }, + { + id: "claude-fable-5", + owned_by: "anthropic", + context_window: 200000, + supports_streaming: true, + supports_vision: true, + allowed: false, + }, + ], + }), + { status: 200, headers: { "Content-Type": "application/json" } }, + ), + ); + const client = new PostHogAPIClient( + "https://eu.posthog.com", + async () => "token", + async () => "token", + 123, + { fetch }, + ); + + const options = await client.getCloudTaskConfigOptions("claude"); + + expect(fetch).toHaveBeenCalledWith( + new URL("https://gateway.eu.posthog.com/posthog_code/v1/models"), + expect.objectContaining({ method: "GET" }), + ); + expect(options.find((option) => option.category === "model")).toMatchObject( + { + currentValue: "claude-opus-4-8", + options: [ + expect.objectContaining({ value: "claude-opus-4-8" }), + expect.objectContaining({ + value: "claude-fable-5", + _meta: expect.any(Object), + }), + ], + }, + ); + }); + + it("uses the configured fetch implementation for task log URLs", async () => { + const fetch = vi + .fn() + .mockResolvedValue( + new Response( + '{"type":"notification","timestamp":"2026-07-21T00:00:00Z"}\n', + { status: 200 }, + ), + ); + const client = new PostHogAPIClient( + "http://localhost:8000", + async () => "token", + async () => "token", + 123, + { fetch }, + ); + vi.spyOn(client, "getTask").mockResolvedValue({ + id: "task-1", + task_number: 1, + slug: "task-1", + title: "Task", + description: "Task", + created_at: "2026-07-21T00:00:00Z", + updated_at: "2026-07-21T00:00:00Z", + origin_product: "user_created", + latest_run: { + id: "run-1", + task: "task-1", + team: 123, + branch: null, + status: "in_progress", + log_url: "https://logs.posthog.test/run-1.jsonl", + error_message: null, + output: null, + state: {}, + created_at: "2026-07-21T00:00:00Z", + updated_at: "2026-07-21T00:00:00Z", + completed_at: null, + }, + }); + + await expect(client.getTaskLogs("task-1")).resolves.toHaveLength(1); + expect(fetch).toHaveBeenCalledWith("https://logs.posthog.test/run-1.jsonl"); + }); + + it.each([ + { + label: "desktop default", + options: undefined, + expectedConnectFrom: "posthog_code", + expectedUserAgent: "posthog/desktop.hog.dev; version: unknown", + }, + { + label: "mobile configuration", + options: { + appVersion: "1.2.3", + userAgent: "posthog/mobile; version: 1.2.3", + githubConnectFrom: "posthog_mobile", + }, + expectedConnectFrom: "posthog_mobile", + expectedUserAgent: "posthog/mobile; version: 1.2.3", + }, + ])( + "uses $label identity for GitHub connections", + async ({ options, expectedConnectFrom, expectedUserAgent }) => { + const fetch = vi.fn().mockResolvedValue( + new Response( + JSON.stringify({ install_url: "https://github.com/login/oauth" }), + { + status: 200, + headers: { "Content-Type": "application/json" }, + }, + ), + ); + const client = new PostHogAPIClient( + "http://localhost:8000", + async () => "token", + async () => "token", + 123, + { ...options, fetch }, + ); + + await expect(client.startGithubUserIntegrationConnect()).resolves.toEqual( + { + install_url: "https://github.com/login/oauth", + }, + ); + + expect(fetch).toHaveBeenCalledWith( + new URL( + "http://localhost:8000/api/users/@me/integrations/github/start/", + ), + expect.objectContaining({ + method: "POST", + body: JSON.stringify({ + team_id: 123, + connect_from: expectedConnectFrom, + }), + }), + ); + expect(fetch.mock.calls[0][1].headers.get("User-Agent")).toBe( + expectedUserAgent, + ); + }, + ); + it("sends supported reasoning effort for cloud Codex runs", async () => { const client = new PostHogAPIClient( "http://localhost:8000", @@ -257,7 +586,13 @@ describe("PostHogAPIClient", () => { reasoningLevel: "high", initialPermissionMode: "auto", }), - ).resolves.toEqual({ id: "run-123", environment: "cloud" }); + ).resolves.toMatchObject({ + id: "run-123", + task: "task-123", + team: 123, + environment: "cloud", + status: "not_started", + }); expect(fetch).toHaveBeenCalledWith( expect.objectContaining({ @@ -435,7 +770,15 @@ describe("PostHogAPIClient", () => { pendingUserMessage: "Read the attached file first", pendingUserArtifactIds: ["artifact-1"], }), - ).resolves.toEqual({ id: "task-123", latest_run: { id: "run-123" } }); + ).resolves.toMatchObject({ + id: "task-123", + latest_run: { + id: "run-123", + task: "task-123", + team: 123, + status: "not_started", + }, + }); expect(fetch).toHaveBeenCalledWith( expect.objectContaining({ diff --git a/packages/api-client/src/posthog-client.ts b/packages/api-client/src/posthog-client.ts index 906e77f290..b2ef9fa8a8 100644 --- a/packages/api-client/src/posthog-client.ts +++ b/packages/api-client/src/posthog-client.ts @@ -4,18 +4,30 @@ import type { CloudMcpServerImport, CloudMcpServerRelayDesignation, CloudRunSource, + CreateTaskAutomationOptions, ExecutionMode, PrAuthorshipMode, SourceProduct, SourceType, StoredLogEntry, + TaskAutomation, TaskRunArtifactMetadata, + UpdateTaskAutomationOptions, } from "@posthog/shared"; import { + buildCloudTaskConfigOptions, + type CloudTaskConfigOption, + createTaskAutomationSchema, DISMISSAL_REASON_OPTIONS, type DismissalReasonOptionValue, + getCloudTaskGatewayUrl, isSupportedReasoningEffort, + normalizeGatewayModelsResponse, resolveCloudInitialPermissionMode, + taskAutomationListSchema, + taskAutomationSchema, + taskAutomationValidationErrorSchema, + updateTaskAutomationSchema, } from "@posthog/shared"; import type { AgentAnalyticsData, @@ -101,9 +113,18 @@ import { type HogQLGrid, shapeAgentAnalytics, } from "./agent-analytics"; -import { buildApiFetcher, requestErrorStatus } from "./fetcher"; +import { + ApiRequestError, + buildApiFetcher, + type FetchImplementation, + requestErrorStatus, +} from "./fetcher"; import { createApiClient, type Schemas } from "./generated"; import type { SpendAnalysisResponse } from "./spend-analysis"; +import { + normalizeTaskResponse, + normalizeTaskRunResponse, +} from "./task-normalization"; export interface ApiClientLogger { warn(...args: unknown[]): void; } @@ -122,6 +143,13 @@ export function setPosthogApiClientAppVersion(version: string): void { clientAppVersion = version; } +export interface PostHogAPIClientOptions { + fetch?: FetchImplementation; + appVersion?: string; + userAgent?: string | null; + githubConnectFrom?: string; +} + export function getPosthogApiClientAppVersion(): string { return clientAppVersion; } @@ -163,6 +191,36 @@ export class CloudUsageLimitError extends Error { } } +export class TaskAutomationValidationError extends Error { + readonly status = 400; + readonly code: string; + readonly attr: string | null; + + constructor(details: { + detail: string; + code: string; + attr: string | null; + }) { + super(details.detail); + this.name = "TaskAutomationValidationError"; + this.code = details.code; + this.attr = details.attr; + } +} + +function rethrowTaskAutomationError(error: unknown): never { + if (error instanceof ApiRequestError && error.status === 400) { + const validationError = taskAutomationValidationErrorSchema.safeParse( + error.body, + ); + if (validationError.success) { + throw new TaskAutomationValidationError(validationError.data); + } + } + + throw error; +} + export const MCP_CATEGORIES = [ { id: "all", label: "All" }, { id: "business", label: "Business Operations" }, @@ -581,7 +639,7 @@ export interface FinalizedTaskArtifactUpload { uploaded_at?: string; } -interface CloudRunOptions { +export interface CloudRunOptions { adapter?: Adapter; model?: string; reasoningLevel?: string; @@ -602,6 +660,56 @@ interface CloudRunOptions { relayedMcpServers?: CloudMcpServerRelayDesignation[]; } +export type CloudRunCommandMethod = + | "user_message" + | "permission_response" + | "set_config_option" + | "cancel" + | "close"; + +export class CloudCommandError extends Error { + readonly status: number; + readonly backendError: string | null; + readonly method: CloudRunCommandMethod; + + constructor( + method: CloudRunCommandMethod, + status: number, + backendError: string | null, + message: string, + ) { + super(message); + this.name = "CloudCommandError"; + this.method = method; + this.status = status; + this.backendError = backendError; + } + + isSandboxInactive(): boolean { + const backendError = this.backendError?.toLowerCase(); + return ( + this.status === 404 || + backendError?.includes("no active sandbox") === true || + backendError?.includes("returned 404") === true + ); + } +} + +function cloudCommandBackendError(payload: unknown): string | null { + if (typeof payload === "string") return payload || null; + if (!payload || typeof payload !== "object") return null; + + const error = "error" in payload ? payload.error : null; + if (typeof error === "string") return error || null; + if (error && typeof error === "object" && "message" in error) { + return typeof error.message === "string" ? error.message : null; + } + if ("message" in payload && typeof payload.message === "string") { + return payload.message; + } + return null; +} + interface CreateTaskRunOptions extends CloudRunOptions { environment?: "local" | "cloud"; mode?: "interactive" | "background"; @@ -1342,19 +1450,28 @@ function previewTokenHeader( export class PostHogAPIClient { private api: ReturnType; private _teamId: number | null = null; + private githubConnectFrom: string; + private readonly fetch: FetchImplementation; + private readonly apiHost: string; constructor( apiHost: string, getAccessToken: () => Promise, refreshAccessToken: () => Promise, teamId?: number, + options: PostHogAPIClientOptions = {}, ) { const baseUrl = apiHost.endsWith("/") ? apiHost.slice(0, -1) : apiHost; + this.apiHost = baseUrl; + this.githubConnectFrom = options.githubConnectFrom ?? "posthog_code"; + this.fetch = options.fetch ?? globalThis.fetch.bind(globalThis); this.api = createApiClient( buildApiFetcher({ getAccessToken, refreshAccessToken, - appVersion: clientAppVersion, + appVersion: options.appVersion ?? clientAppVersion, + fetch: options.fetch, + userAgent: options.userAgent, }), baseUrl, ); @@ -1391,6 +1508,21 @@ export class PostHogAPIClient { return data; } + async getCloudTaskConfigOptions( + adapter: Adapter = "claude", + ): Promise { + const url = new URL(`${getCloudTaskGatewayUrl(this.apiHost)}/v1/models`); + const response = await this.api.fetcher.fetch({ + method: "get", + url, + path: url.pathname, + }); + return buildCloudTaskConfigOptions( + normalizeGatewayModelsResponse(await response.json()), + adapter, + ); + } + // Desktop file system — the backend surface that backs canvas channels // (top-level folders) and dashboards. These routes aren't in the generated // OpenAPI client, so we use the raw fetcher. @@ -1755,7 +1887,10 @@ export class PostHogAPIClient { url, path: urlPath, overrides: { - body: JSON.stringify({ team_id: id, connect_from: "posthog_code" }), + body: JSON.stringify({ + team_id: id, + connect_from: this.githubConnectFrom, + }), }, }); if (!response.ok) { @@ -2257,7 +2392,7 @@ export class PostHogAPIClient { originProduct?: string; internal?: boolean; channel?: string; - }) { + }): Promise { const teamId = await this.getTeamId(); const params: Record = { limit: 500, @@ -2288,7 +2423,9 @@ export class PostHogAPIClient { query: params, }); - return data.results ?? []; + return (data.results ?? []).map((task) => + normalizeTaskResponse(task, { teamId }), + ); } async getTaskSummaries(ids: string[]) { @@ -2330,7 +2467,102 @@ export class PostHogAPIClient { const data = await this.api.get(`/api/projects/{project_id}/tasks/{id}/`, { path: { project_id: teamId.toString(), id: taskId }, }); - return data as unknown as Task; + return normalizeTaskResponse(data, { teamId }); + } + + async listTaskAutomations(options?: { + limit?: number; + offset?: number; + }): Promise { + const teamId = await this.getTeamId(); + const data = await this.api.get( + `/api/projects/{project_id}/task_automations/`, + { + path: { project_id: teamId.toString() }, + query: { + limit: options?.limit ?? 500, + ...(options?.offset === undefined ? {} : { offset: options.offset }), + }, + }, + ); + + return taskAutomationListSchema.parse(data).results; + } + + async getTaskAutomation(automationId: string): Promise { + const teamId = await this.getTeamId(); + const data = await this.api.get( + `/api/projects/{project_id}/task_automations/{id}/`, + { + path: { project_id: teamId.toString(), id: automationId }, + }, + ); + + return taskAutomationSchema.parse(data); + } + + async createTaskAutomation( + options: CreateTaskAutomationOptions, + ): Promise { + const teamId = await this.getTeamId(); + const body = createTaskAutomationSchema.parse(options); + + try { + const data = await this.api.post( + `/api/projects/{project_id}/task_automations/`, + { + path: { project_id: teamId.toString() }, + body: body as Schemas.TaskAutomation, + }, + ); + return taskAutomationSchema.parse(data); + } catch (error) { + rethrowTaskAutomationError(error); + } + } + + async updateTaskAutomation( + automationId: string, + updates: UpdateTaskAutomationOptions, + ): Promise { + const teamId = await this.getTeamId(); + const body = updateTaskAutomationSchema.parse(updates); + + try { + const data = await this.api.patch( + `/api/projects/{project_id}/task_automations/{id}/`, + { + path: { project_id: teamId.toString(), id: automationId }, + body, + }, + ); + return taskAutomationSchema.parse(data); + } catch (error) { + rethrowTaskAutomationError(error); + } + } + + async deleteTaskAutomation(automationId: string): Promise { + const teamId = await this.getTeamId(); + await this.api.delete(`/api/projects/{project_id}/task_automations/{id}/`, { + path: { project_id: teamId.toString(), id: automationId }, + }); + } + + async runTaskAutomation(automationId: string): Promise { + const teamId = await this.getTeamId(); + const path = `/api/projects/${teamId}/task_automations/${automationId}/run/`; + + try { + const response = await this.api.fetcher.fetch({ + method: "post", + path, + url: new URL(`${this.api.baseUrl}${path}`), + }); + return taskAutomationSchema.parse(await response.json()); + } catch (error) { + rethrowTaskAutomationError(error); + } } async createTask( @@ -2357,7 +2589,7 @@ export class PostHogAPIClient { pending_user_artifact_ids?: string[]; auto_publish?: boolean; }, - ) { + ): Promise { const teamId = await this.getTeamId(); const { origin_product: originProduct, ...taskOptions } = options; @@ -2369,10 +2601,13 @@ export class PostHogAPIClient { } as unknown as Schemas.Task, }); - return data; + return normalizeTaskResponse(data, { teamId }); } - async updateTask(taskId: string, updates: Partial) { + async updateTask( + taskId: string, + updates: Partial, + ): Promise { const teamId = await this.getTeamId(); const data = await this.api.patch( `/api/projects/{project_id}/tasks/{id}/`, @@ -2382,7 +2617,7 @@ export class PostHogAPIClient { }, ); - return data; + return normalizeTaskResponse(data, { teamId }); } async deleteTask(taskId: string) { @@ -2681,9 +2916,28 @@ export class PostHogAPIClient { async sendRunCommand( taskId: string, runId: string, - method: "user_message" | "cancel" | "close", + method: CloudRunCommandMethod, params?: Record, ): Promise<{ success: boolean; result?: unknown; error?: string }> { + try { + return { + success: true, + result: await this.sendCloudRunCommand(taskId, runId, method, params), + }; + } catch (error) { + return { + success: false, + error: error instanceof Error ? error.message : "Unknown error", + }; + } + } + + async sendCloudRunCommand( + taskId: string, + runId: string, + method: CloudRunCommandMethod, + params: Record = {}, + ): Promise { const teamId = await this.getTeamId(); const url = new URL( `${this.api.baseUrl}/api/projects/${teamId}/tasks/${taskId}/runs/${runId}/command/`, @@ -2691,7 +2945,7 @@ export class PostHogAPIClient { const body = { jsonrpc: "2.0", method, - params: params ?? {}, + params, id: `posthog-code-${Date.now()}`, }; @@ -2705,39 +2959,54 @@ export class PostHogAPIClient { }, }); - if (!response.ok) { - const errorText = await response.text().catch(() => ""); - let errorMessage = `Command failed: ${response.statusText}`; - try { - const errorJson = JSON.parse(errorText); - errorMessage = - errorJson.error?.message ?? errorJson.error ?? errorMessage; - } catch { - if (errorText) errorMessage = errorText; - } - return { success: false, error: errorMessage }; - } - const data = (await response.json()) as { - error?: { message?: string }; + error?: unknown; result?: unknown; }; if (data.error) { - return { - success: false, - error: data.error.message ?? JSON.stringify(data.error), - }; + const backendError = cloudCommandBackendError(data); + throw new CloudCommandError( + method, + response.status, + backendError, + `Cloud command '${method}' error: ${backendError ?? JSON.stringify(data.error)}`, + ); } - return { success: true, result: data.result }; + return data.result; } catch (error) { - return { - success: false, - error: error instanceof Error ? error.message : "Unknown error", - }; + if (error instanceof CloudCommandError) throw error; + if (error instanceof ApiRequestError) { + const backendError = cloudCommandBackendError(error.body); + throw new CloudCommandError( + method, + error.status, + backendError, + `Cloud command '${method}' failed: ${error.status}${backendError ? ` ${backendError}` : ""}`, + ); + } + throw error; } } + async cancelTaskRun( + taskId: string, + runId: string, + reason?: string, + ): Promise<{ status?: string }> { + const teamId = await this.getTeamId(); + const path = `/api/projects/${teamId}/tasks/${taskId}/runs/${runId}/cancel/`; + const response = await this.api.fetcher.fetch({ + method: "post", + url: new URL(`${this.api.baseUrl}${path}`), + path, + overrides: { + body: JSON.stringify(reason ? { reason } : {}), + }, + }); + return (await response.json().catch(() => ({}))) as { status?: string }; + } + async runTaskInCloud( taskId: string, branch?: string | null, @@ -2761,7 +3030,7 @@ export class PostHogAPIClient { }), ); - return data as unknown as Task; + return normalizeTaskResponse(data, { teamId }); } async warmTask(options: { @@ -3001,7 +3270,8 @@ export class PostHogAPIClient { throw new Error(`Failed to resume run in cloud: ${response.statusText}`); } - return (await response.json()) as TaskRun; + const data = (await response.json()) as Schemas.TaskRunDetail; + return normalizeTaskRunResponse(data, { teamId, taskId }); } async listTaskRuns(taskId: string): Promise { @@ -3019,8 +3289,11 @@ export class PostHogAPIClient { throw new Error(`Failed to fetch task runs: ${response.statusText}`); } - const data = (await response.json()) as { results?: TaskRun[] }; - return data.results ?? []; + const data = + (await response.json()) as Partial; + return (data.results ?? []).map((run) => + normalizeTaskRunResponse(run, { teamId, taskId }), + ); } async getTaskRun(taskId: string, runId: string): Promise { @@ -3038,7 +3311,8 @@ export class PostHogAPIClient { throw new Error(`Failed to fetch task run: ${response.statusText}`); } - return (await response.json()) as TaskRun; + const data = (await response.json()) as Schemas.TaskRunDetail; + return normalizeTaskRunResponse(data, { teamId, taskId }); } async createTaskRun( @@ -3070,7 +3344,8 @@ export class PostHogAPIClient { throw new Error(`Failed to create task run: ${response.statusText}`); } - return (await response.json()) as TaskRun; + const data = (await response.json()) as Schemas.TaskRunDetail; + return normalizeTaskRunResponse(data, { teamId, taskId }); } async startTaskRun( @@ -3100,7 +3375,8 @@ export class PostHogAPIClient { throw new Error(`Failed to start task run: ${response.statusText}`); } - return (await response.json()) as Task; + const data = (await response.json()) as Schemas.Task; + return normalizeTaskResponse(data, { teamId }); } async updateTaskRun( @@ -3125,7 +3401,7 @@ export class PostHogAPIClient { body: updates as Record, }, ); - return data as unknown as TaskRun; + return normalizeTaskRunResponse(data, { teamId, taskId }); } /** @@ -3215,14 +3491,14 @@ export class PostHogAPIClient { async getTaskLogs(taskId: string): Promise { try { - const task = (await this.getTask(taskId)) as unknown as Task; + const task = await this.getTask(taskId); const logUrl = task?.latest_run?.log_url; if (!logUrl) { return []; } - const response = await fetch(logUrl); + const response = await this.fetch(logUrl); if (!response.ok) { log.warn( diff --git a/packages/api-client/src/task-normalization.test.ts b/packages/api-client/src/task-normalization.test.ts new file mode 100644 index 0000000000..9c9739a51f --- /dev/null +++ b/packages/api-client/src/task-normalization.test.ts @@ -0,0 +1,115 @@ +import { describe, expect, it } from "vitest"; +import { + normalizeTaskResponse, + normalizeTaskRunResponse, +} from "./task-normalization"; + +describe("task response normalization", () => { + it("normalizes legacy started runs and nullable generated fields", () => { + expect( + normalizeTaskRunResponse( + { + id: "run-1", + task: "task-1", + status: "started", + branch: null, + stage: null, + runtime_adapter: null, + model: null, + reasoning_effort: null, + log_url: null, + error_message: null, + output: null, + state: null, + artifacts: [ + { + id: "artifact-1", + name: "result.txt", + type: "legacy_type", + storage_path: "tasks/result.txt", + uploaded_at: "2026-07-21T00:00:00Z", + }, + ], + created_at: "2026-07-21T00:00:00Z", + updated_at: "2026-07-21T00:01:00Z", + completed_at: null, + }, + { teamId: 123 }, + ), + ).toEqual({ + id: "run-1", + task: "task-1", + team: 123, + branch: null, + stage: null, + runtime_adapter: null, + model: null, + reasoning_effort: null, + status: "in_progress", + log_url: "", + error_message: null, + output: null, + state: {}, + artifacts: [ + { + id: "artifact-1", + name: "result.txt", + type: "artifact", + storage_path: "tasks/result.txt", + uploaded_at: "2026-07-21T00:00:00Z", + }, + ], + created_at: "2026-07-21T00:00:00Z", + updated_at: "2026-07-21T00:01:00Z", + completed_at: null, + }); + }); + + it("normalizes task responses and their generated latest-run records", () => { + expect( + normalizeTaskResponse( + { + id: "task-1", + task_number: null, + slug: "task-1", + repository: null, + github_integration: null, + github_user_integration: null, + json_schema: null, + signal_report: null, + channel: null, + latest_run: { + id: "run-1", + status: "started", + log_url: null, + }, + created_at: "2026-07-21T00:00:00Z", + updated_at: "2026-07-21T00:01:00Z", + }, + { teamId: 123 }, + ), + ).toMatchObject({ + id: "task-1", + task_number: null, + slug: "task-1", + title: "", + description: "", + origin_product: "", + repository: null, + github_integration: null, + github_user_integration: null, + json_schema: null, + signal_report: null, + channel: null, + latest_run: { + id: "run-1", + task: "task-1", + team: 123, + status: "in_progress", + log_url: "", + output: null, + state: {}, + }, + }); + }); +}); diff --git a/packages/api-client/src/task-normalization.ts b/packages/api-client/src/task-normalization.ts new file mode 100644 index 0000000000..9571a5b27e --- /dev/null +++ b/packages/api-client/src/task-normalization.ts @@ -0,0 +1,203 @@ +import type { + ArtifactType, + Task, + TaskRun, + TaskRunArtifact, + TaskRunArtifactMetadata, + TaskRunStatus, +} from "@posthog/shared/domain-types"; +import type { Schemas } from "./generated"; + +type TaskRunResponseDTO = Partial< + Omit +> & { + id: string; + artifacts?: Array< + Schemas.TaskRunArtifactResponse & { metadata?: unknown } + > | null; + status?: Schemas.StatusA35Enum | "started" | null; + team?: number | null; +}; + +type TaskResponseDTO = Partial< + Omit +> & { + id: string; + channel?: string | null; + created_by?: Schemas.UserBasic | null; + github_user_integration?: string | null; + json_schema?: unknown | null; + latest_run?: Record | null; + runtime?: unknown; +}; + +function isRecord(value: unknown): value is Record { + return typeof value === "object" && value !== null && !Array.isArray(value); +} + +function isTaskRunResponseDTO(value: unknown): value is TaskRunResponseDTO { + return isRecord(value) && typeof value.id === "string"; +} + +function normalizeTaskRunStatus(status: unknown): TaskRunStatus { + switch (status) { + case "started": + return "in_progress"; + case "not_started": + case "queued": + case "in_progress": + case "completed": + case "failed": + case "cancelled": + return status; + default: + return "not_started"; + } +} + +function normalizeArtifactType(type: string): ArtifactType { + switch (type) { + case "plan": + case "context": + case "reference": + case "output": + case "artifact": + case "user_attachment": + case "skill_bundle": + return type; + default: + return "artifact"; + } +} + +function normalizeArtifactMetadata( + value: unknown, +): TaskRunArtifactMetadata | undefined { + if ( + !isRecord(value) || + typeof value.skill_name !== "string" || + (value.skill_source !== "user" && + value.skill_source !== "repo" && + value.skill_source !== "marketplace" && + value.skill_source !== "codex") + ) { + return undefined; + } + if ( + typeof value.content_sha256 !== "string" || + value.bundle_format !== "zip" || + typeof value.schema_version !== "number" + ) { + return undefined; + } + + return { + skill_name: value.skill_name, + skill_source: value.skill_source, + content_sha256: value.content_sha256, + bundle_format: value.bundle_format, + schema_version: value.schema_version, + }; +} + +function normalizeTaskRunArtifact( + artifact: NonNullable[number], +): TaskRunArtifact { + const metadata = normalizeArtifactMetadata(artifact.metadata); + + return { + ...(artifact.id === undefined ? {} : { id: artifact.id }), + name: artifact.name, + type: normalizeArtifactType(artifact.type), + ...(artifact.source === undefined ? {} : { source: artifact.source }), + ...(artifact.size === undefined ? {} : { size: artifact.size }), + ...(artifact.content_type === undefined + ? {} + : { content_type: artifact.content_type }), + ...(metadata === undefined ? {} : { metadata }), + ...(artifact.storage_path === undefined + ? {} + : { storage_path: artifact.storage_path }), + ...(artifact.uploaded_at === undefined + ? {} + : { uploaded_at: artifact.uploaded_at }), + }; +} + +export function normalizeTaskRunResponse( + dto: TaskRunResponseDTO, + context: { teamId: number; taskId?: string }, +): TaskRun { + return { + id: dto.id, + task: dto.task ?? context.taskId ?? "", + team: dto.team ?? context.teamId, + branch: dto.branch ?? null, + ...(dto.runtime_adapter === undefined + ? {} + : { runtime_adapter: dto.runtime_adapter }), + ...(dto.model === undefined ? {} : { model: dto.model }), + ...(dto.reasoning_effort === undefined + ? {} + : { reasoning_effort: dto.reasoning_effort }), + ...(dto.stage === undefined ? {} : { stage: dto.stage }), + ...(dto.environment === undefined ? {} : { environment: dto.environment }), + status: normalizeTaskRunStatus(dto.status), + log_url: dto.log_url ?? "", + error_message: dto.error_message ?? null, + output: isRecord(dto.output) ? dto.output : null, + state: isRecord(dto.state) ? dto.state : {}, + ...(dto.artifacts == null + ? {} + : { artifacts: dto.artifacts.map(normalizeTaskRunArtifact) }), + created_at: dto.created_at ?? "", + updated_at: dto.updated_at ?? "", + completed_at: dto.completed_at ?? null, + }; +} + +export function normalizeTaskResponse( + dto: TaskResponseDTO, + context: { teamId: number }, +): Task { + const jsonSchema = isRecord(dto.json_schema) ? dto.json_schema : null; + const runtime = + dto.runtime === "acp" || dto.runtime === "pi" ? dto.runtime : undefined; + + const latestRun = isTaskRunResponseDTO(dto.latest_run) + ? normalizeTaskRunResponse(dto.latest_run, { + teamId: context.teamId, + taskId: dto.id, + }) + : undefined; + + return { + id: dto.id, + task_number: dto.task_number ?? null, + slug: dto.slug ?? "", + title: dto.title ?? "", + ...(dto.title_manually_set === undefined + ? {} + : { title_manually_set: dto.title_manually_set }), + description: dto.description ?? "", + created_at: dto.created_at ?? "", + updated_at: dto.updated_at ?? "", + ...(dto.created_by === undefined ? {} : { created_by: dto.created_by }), + origin_product: dto.origin_product ?? "", + ...(dto.repository === undefined ? {} : { repository: dto.repository }), + ...(dto.github_integration === undefined + ? {} + : { github_integration: dto.github_integration }), + ...(dto.github_user_integration === undefined + ? {} + : { github_user_integration: dto.github_user_integration }), + ...(dto.json_schema === undefined ? {} : { json_schema: jsonSchema }), + ...(dto.signal_report === undefined + ? {} + : { signal_report: dto.signal_report }), + ...(dto.internal === undefined ? {} : { internal: dto.internal }), + ...(runtime === undefined ? {} : { runtime }), + ...(dto.channel === undefined ? {} : { channel: dto.channel }), + ...(latestRun === undefined ? {} : { latest_run: latestRun }), + }; +} diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 65ef2c8cbf..5cdfefdf93 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -869,9 +869,6 @@ importers: packages/api-client: dependencies: - '@posthog/agent': - specifier: workspace:* - version: link:../agent '@posthog/shared': specifier: workspace:* version: link:../shared @@ -19215,13 +19212,6 @@ snapshots: '@tybys/wasm-util': 0.10.2 optional: true - '@napi-rs/wasm-runtime@1.1.6(@emnapi/core@1.10.0)(@emnapi/runtime@1.10.0)': - dependencies: - '@emnapi/core': 1.10.0 - '@emnapi/runtime': 1.10.0 - '@tybys/wasm-util': 0.10.3 - optional: true - '@napi-rs/wasm-runtime@1.1.6(@emnapi/core@1.11.1)(@emnapi/runtime@1.11.1)': dependencies: '@emnapi/core': 1.11.1 @@ -22659,7 +22649,7 @@ snapshots: dependencies: '@emnapi/core': 1.10.0 '@emnapi/runtime': 1.10.0 - '@napi-rs/wasm-runtime': 1.1.6(@emnapi/core@1.10.0)(@emnapi/runtime@1.10.0) + '@napi-rs/wasm-runtime': 1.1.5(@emnapi/core@1.10.0)(@emnapi/runtime@1.10.0) optional: true '@unrs/resolver-binding-win32-arm64-msvc@1.12.2': From 01aec6d6b84e9751dec936413b67bc88700bb125 Mon Sep 17 00:00:00 2001 From: Richard Solomou Date: Fri, 24 Jul 2026 19:18:39 +0300 Subject: [PATCH 34/45] refactor(core): extract cloud task policies Generated-By: PostHog Code Task-Id: c1bbe3cf-742b-4b24-bf96-d11a18b4cf22 --- .../automations/automationSchedule.test.ts | 223 ++++++++++++++++++ .../src/automations/automationSchedule.ts | 216 +++++++++++++++++ .../src/automations/automationStatus.test.ts | 80 +++++++ .../core/src/automations/automationStatus.ts | 83 +++++++ packages/core/src/inbox/artefacts.test.ts | 159 ++++++++++++- packages/core/src/inbox/artefacts.ts | 86 +++++++ .../core/src/inbox/reportFiltering.test.ts | 10 + packages/core/src/inbox/reportFiltering.ts | 12 +- .../core/src/sessions/cloudSessionConfig.ts | 9 +- packages/core/src/sessions/executionModes.ts | 4 +- .../sessions/portableSessionEvents.test.ts | 71 ++++++ .../src/sessions/portableSessionEvents.ts | 98 ++++++++ .../core/src/sessions/sessionActivity.test.ts | 152 ++++++++++++ packages/core/src/sessions/sessionActivity.ts | 129 ++++++++++ packages/core/src/tasks/taskActivity.test.ts | 132 +++++++++++ packages/core/src/tasks/taskActivity.ts | 45 ++++ packages/core/src/tasks/taskArchive.test.ts | 44 ++++ packages/core/src/tasks/taskArchive.ts | 6 + .../src/tasks/taskStatusPresentation.test.ts | 69 ++++++ .../core/src/tasks/taskStatusPresentation.ts | 37 +++ 20 files changed, 1659 insertions(+), 6 deletions(-) create mode 100644 packages/core/src/automations/automationSchedule.test.ts create mode 100644 packages/core/src/automations/automationSchedule.ts create mode 100644 packages/core/src/automations/automationStatus.test.ts create mode 100644 packages/core/src/automations/automationStatus.ts create mode 100644 packages/core/src/sessions/portableSessionEvents.test.ts create mode 100644 packages/core/src/sessions/portableSessionEvents.ts create mode 100644 packages/core/src/sessions/sessionActivity.test.ts create mode 100644 packages/core/src/sessions/sessionActivity.ts create mode 100644 packages/core/src/tasks/taskActivity.test.ts create mode 100644 packages/core/src/tasks/taskActivity.ts create mode 100644 packages/core/src/tasks/taskArchive.test.ts create mode 100644 packages/core/src/tasks/taskArchive.ts create mode 100644 packages/core/src/tasks/taskStatusPresentation.test.ts create mode 100644 packages/core/src/tasks/taskStatusPresentation.ts diff --git a/packages/core/src/automations/automationSchedule.test.ts b/packages/core/src/automations/automationSchedule.test.ts new file mode 100644 index 0000000000..5deb2487a0 --- /dev/null +++ b/packages/core/src/automations/automationSchedule.test.ts @@ -0,0 +1,223 @@ +import { describe, expect, it } from "vitest"; +import { + type AutomationScheduleDraft, + buildCronExpression, + createDefaultScheduleDraft, + deriveAutomationName, + formatAutomationScheduleSummary, + formatScheduleSummary, + parseCronExpression, + sanitizeHour, + sanitizeMinute, + WEEKDAY_OPTIONS, +} from "./automationSchedule"; + +describe("automationSchedule", () => { + it("creates the default daily schedule draft", () => { + expect(createDefaultScheduleDraft()).toEqual({ + mode: "daily", + hour: "09", + minute: "00", + weekday: "1", + rawCron: "0 9 * * *", + }); + }); + + it("provides cron weekday values in display order", () => { + expect(WEEKDAY_OPTIONS).toEqual([ + { value: "1", label: "Mon" }, + { value: "2", label: "Tue" }, + { value: "3", label: "Wed" }, + { value: "4", label: "Thu" }, + { value: "5", label: "Fri" }, + { value: "6", label: "Sat" }, + { value: "0", label: "Sun" }, + ]); + }); + + it.each([ + ["", ""], + ["a", ""], + ["7", "07"], + ["09", "09"], + ["2x3", "23"], + ["24", "23"], + ["999", "23"], + ])("sanitizes hour input %j to %j", (input, expected) => { + expect(sanitizeHour(input)).toBe(expected); + }); + + it.each([ + ["", ""], + ["a", ""], + ["7", "07"], + ["09", "09"], + ["5x9", "59"], + ["60", "59"], + ["999", "59"], + ])("sanitizes minute input %j to %j", (input, expected) => { + expect(sanitizeMinute(input)).toBe(expected); + }); + + it.each<{ + name: string; + changes: Partial; + expected: string; + }>([ + { + name: "hourly", + changes: { mode: "hourly", minute: "15" }, + expected: "15 * * * *", + }, + { + name: "daily", + changes: { mode: "daily", hour: "09", minute: "15" }, + expected: "15 9 * * *", + }, + { + name: "weekdays", + changes: { mode: "weekdays", hour: "10", minute: "00" }, + expected: "0 10 * * 1-5", + }, + { + name: "weekly", + changes: { + mode: "weekly", + hour: "11", + minute: "30", + weekday: "4", + }, + expected: "30 11 * * 4", + }, + { + name: "weekly with a missing weekday", + changes: { mode: "weekly", weekday: "" }, + expected: "0 9 * * 1", + }, + { + name: "preset with missing time values", + changes: { mode: "daily", hour: "", minute: "" }, + expected: "0 9 * * *", + }, + { + name: "custom", + changes: { mode: "custom", rawCron: " */15 * * * * " }, + expected: "*/15 * * * *", + }, + ])("builds the $name cron expression", ({ changes, expected }) => { + expect( + buildCronExpression({ ...createDefaultScheduleDraft(), ...changes }), + ).toBe(expected); + }); + + it.each([ + [ + "15 * * * *", + { + mode: "hourly", + hour: "09", + minute: "15", + weekday: "*", + rawCron: "15 * * * *", + }, + ], + [ + "0 9 * * *", + { + mode: "daily", + hour: "09", + minute: "00", + weekday: "*", + rawCron: "0 9 * * *", + }, + ], + [ + "0 9 * * 1-5", + { + mode: "weekdays", + hour: "09", + minute: "00", + weekday: "1", + rawCron: "0 9 * * 1-5", + }, + ], + [ + "30 14 * * 2", + { + mode: "weekly", + hour: "14", + minute: "30", + weekday: "2", + rawCron: "30 14 * * 2", + }, + ], + ] as const)("parses %s into a schedule draft", (cron, expected) => { + expect(parseCronExpression(cron)).toEqual(expected); + }); + + it.each(["*/15 * * * *", "0 9 1 * *", "0 9 * 1 *", "0 9 * * 1,3"])( + "keeps unsupported cron expression %s in custom mode", + (cron) => { + expect(parseCronExpression(cron)).toMatchObject({ + mode: "custom", + rawCron: cron, + }); + }, + ); + + it("normalizes surrounding and repeated cron whitespace", () => { + expect(parseCronExpression(" 5 8 * * * ")).toEqual({ + mode: "daily", + hour: "08", + minute: "05", + weekday: "*", + rawCron: "5 8 * * *", + }); + }); + + it("uses default draft fields for a cron expression with the wrong arity", () => { + expect(parseCronExpression("0 9 * *")).toEqual({ + mode: "custom", + hour: "09", + minute: "00", + weekday: "1", + rawCron: "0 9 * *", + }); + }); + + it("derives a compact name from the first non-empty prompt line", () => { + expect( + deriveAutomationName( + "\n Review every open PostHog PR for stale comments \nIgnore this line", + ), + ).toBe("Review every open PostHog PR for stale comments"); + }); + + it("returns an empty name for a blank prompt", () => { + expect(deriveAutomationName(" \n\t\n ")).toBe(""); + }); + + it("limits derived names to 80 characters", () => { + expect(deriveAutomationName("a".repeat(100))).toBe("a".repeat(80)); + }); + + it.each([ + ["15 * * * *", "Europe/London", "Every hour at :15 · Europe/London"], + ["0 9 * * *", null, "Daily at 09:00"], + ["0 9 * * 1-5", "UTC", "Weekdays at 09:00 · UTC"], + ["30 14 * * 2", undefined, "Tue at 14:30"], + ["30 14 * * 7", "UTC", "Weekly at 14:30 · UTC"], + ["*/15 * * * *", "UTC", "Custom schedule · UTC"], + ])("formats %s with timezone %j", (cronExpression, timezone, expected) => { + expect(formatScheduleSummary(cronExpression, timezone)).toBe(expected); + }); + + it("formats a schedule from an automation-shaped input", () => { + expect( + formatAutomationScheduleSummary({ + cron_expression: "0 18 * * 0", + timezone: "America/New_York", + }), + ).toBe("Sun at 18:00 · America/New_York"); + }); +}); diff --git a/packages/core/src/automations/automationSchedule.ts b/packages/core/src/automations/automationSchedule.ts new file mode 100644 index 0000000000..527e8da497 --- /dev/null +++ b/packages/core/src/automations/automationSchedule.ts @@ -0,0 +1,216 @@ +export type AutomationScheduleMode = + | "hourly" + | "daily" + | "weekdays" + | "weekly" + | "custom"; + +export interface AutomationScheduleDraft { + mode: AutomationScheduleMode; + hour: string; + minute: string; + weekday: string; + rawCron: string; +} + +export interface AutomationScheduleSummaryInput { + cron_expression: string; + timezone?: string | null; +} + +export const WEEKDAY_OPTIONS = [ + { value: "1", label: "Mon" }, + { value: "2", label: "Tue" }, + { value: "3", label: "Wed" }, + { value: "4", label: "Thu" }, + { value: "5", label: "Fri" }, + { value: "6", label: "Sat" }, + { value: "0", label: "Sun" }, +] as const; + +export function createDefaultScheduleDraft(): AutomationScheduleDraft { + return { + mode: "daily", + hour: "09", + minute: "00", + weekday: "1", + rawCron: "0 9 * * *", + }; +} + +function padTimePart(value: string): string { + return value.padStart(2, "0"); +} + +export function sanitizeHour(value: string): string { + const digitsOnly = value.replace(/\D/g, "").slice(0, 2); + if (!digitsOnly) { + return ""; + } + + return String(Math.min(23, Number(digitsOnly))).padStart(2, "0"); +} + +export function sanitizeMinute(value: string): string { + const digitsOnly = value.replace(/\D/g, "").slice(0, 2); + if (!digitsOnly) { + return ""; + } + + return String(Math.min(59, Number(digitsOnly))).padStart(2, "0"); +} + +export function buildCronExpression(draft: AutomationScheduleDraft): string { + if (draft.mode === "custom") { + return draft.rawCron.trim(); + } + + const minute = draft.minute ? String(Number(draft.minute)) : "0"; + const hour = draft.hour ? String(Number(draft.hour)) : "9"; + + switch (draft.mode) { + case "hourly": + return `${minute} * * * *`; + case "weekdays": + return `${minute} ${hour} * * 1-5`; + case "weekly": + return `${minute} ${hour} * * ${draft.weekday || "1"}`; + default: + return `${minute} ${hour} * * *`; + } +} + +export function parseCronExpression( + cronExpression: string, +): AutomationScheduleDraft { + const normalized = cronExpression.trim(); + const parts = normalized.split(/\s+/); + + if (parts.length !== 5) { + return { + ...createDefaultScheduleDraft(), + mode: "custom", + rawCron: normalized, + }; + } + + const [minute, hour, dayOfMonth, month, dayOfWeek] = parts; + const isNumericMinute = /^\d{1,2}$/.test(minute); + const isNumericHour = /^\d{1,2}$/.test(hour); + const draftBase = { + hour: padTimePart(hour), + minute: padTimePart(minute), + weekday: dayOfWeek, + rawCron: normalized, + }; + + if ( + isNumericMinute && + hour === "*" && + dayOfMonth === "*" && + month === "*" && + dayOfWeek === "*" + ) { + return { + ...draftBase, + mode: "hourly", + hour: "09", + }; + } + + if ( + isNumericMinute && + isNumericHour && + dayOfMonth === "*" && + month === "*" && + dayOfWeek === "*" + ) { + return { + ...draftBase, + mode: "daily", + }; + } + + if ( + isNumericMinute && + isNumericHour && + dayOfMonth === "*" && + month === "*" && + dayOfWeek === "1-5" + ) { + return { + ...draftBase, + mode: "weekdays", + weekday: "1", + }; + } + + if ( + isNumericMinute && + isNumericHour && + dayOfMonth === "*" && + month === "*" && + /^\d$/.test(dayOfWeek) + ) { + return { + ...draftBase, + mode: "weekly", + }; + } + + return { + ...draftBase, + mode: "custom", + }; +} + +export function deriveAutomationName(prompt: string): string { + const normalized = prompt + .split("\n") + .map((line) => line.trim()) + .find(Boolean); + + if (!normalized) { + return ""; + } + + return normalized.replace(/\s+/g, " ").slice(0, 80); +} + +function formatTime(hour: string, minute: string): string { + return `${padTimePart(hour)}:${padTimePart(minute)}`; +} + +export function formatScheduleSummary( + cronExpression: string, + timezone: string | null | undefined, +): string { + const draft = parseCronExpression(cronExpression); + const suffix = timezone ? ` · ${timezone}` : ""; + + switch (draft.mode) { + case "hourly": + return `Every hour at :${padTimePart(draft.minute)}${suffix}`; + case "weekdays": + return `Weekdays at ${formatTime(draft.hour, draft.minute)}${suffix}`; + case "weekly": { + const label = + WEEKDAY_OPTIONS.find((option) => option.value === draft.weekday) + ?.label ?? "Weekly"; + return `${label} at ${formatTime(draft.hour, draft.minute)}${suffix}`; + } + case "custom": + return `Custom schedule${suffix}`; + default: + return `Daily at ${formatTime(draft.hour, draft.minute)}${suffix}`; + } +} + +export function formatAutomationScheduleSummary( + automation: AutomationScheduleSummaryInput, +): string { + return formatScheduleSummary( + automation.cron_expression, + automation.timezone ?? null, + ); +} diff --git a/packages/core/src/automations/automationStatus.test.ts b/packages/core/src/automations/automationStatus.test.ts new file mode 100644 index 0000000000..862fc3c0b5 --- /dev/null +++ b/packages/core/src/automations/automationStatus.test.ts @@ -0,0 +1,80 @@ +import { describe, expect, it } from "vitest"; +import { + type AutomationStatusPresentation, + type AutomationTaskRunStatus, + getAutomationStatusPresentation, +} from "./automationStatus"; + +describe("automationStatus", () => { + it.each<{ + status: AutomationTaskRunStatus; + expected: AutomationStatusPresentation | null; + }>([ + { + status: "not_started", + expected: { label: "Queued", tone: "warning", iconKind: "queued" }, + }, + { + status: "queued", + expected: { label: "Queued", tone: "warning", iconKind: "queued" }, + }, + { status: "started", expected: null }, + { status: "in_progress", expected: null }, + { + status: "completed", + expected: { label: "Success", tone: "success", iconKind: "success" }, + }, + { + status: "failed", + expected: { label: "Failed", tone: "error", iconKind: "failed" }, + }, + { + status: "cancelled", + expected: { label: "Failed", tone: "error", iconKind: "failed" }, + }, + ])( + "maps task-run status $status to renderer-neutral presentation data", + ({ status, expected }) => { + expect( + getAutomationStatusPresentation({ + lastRunStatus: "success", + lastTaskRunStatus: status, + }), + ).toEqual(expected); + }, + ); + + it.each([ + ["running", null], + ["success", { label: "Success", tone: "success", iconKind: "success" }], + ["failed", { label: "Failed", tone: "error", iconKind: "failed" }], + [null, { label: "Never run", tone: "neutral", iconKind: "never-run" }], + ["unknown", { label: "Never run", tone: "neutral", iconKind: "never-run" }], + ] as const)( + "falls back from automation status %j to semantic presentation data", + (lastRunStatus, expected) => { + expect(getAutomationStatusPresentation({ lastRunStatus })).toEqual( + expected, + ); + }, + ); + + it("prioritizes linked task-run detail over the automation-level status", () => { + expect( + getAutomationStatusPresentation({ + lastRunStatus: "failed", + lastTaskRunStatus: "completed", + }), + ).toEqual({ + label: "Success", + tone: "success", + iconKind: "success", + }); + }); + + it("does not expose renderer-specific class names", () => { + expect( + getAutomationStatusPresentation({ lastRunStatus: "success" }), + ).not.toHaveProperty("className"); + }); +}); diff --git a/packages/core/src/automations/automationStatus.ts b/packages/core/src/automations/automationStatus.ts new file mode 100644 index 0000000000..4bcc13735e --- /dev/null +++ b/packages/core/src/automations/automationStatus.ts @@ -0,0 +1,83 @@ +export type AutomationTaskRunStatus = + | "not_started" + | "queued" + | "started" + | "in_progress" + | "completed" + | "failed" + | "cancelled"; + +export interface AutomationStatusInput { + lastRunStatus: string | null; + lastTaskRunStatus?: AutomationTaskRunStatus | null; +} + +export type AutomationStatusTone = "neutral" | "warning" | "success" | "error"; + +export type AutomationStatusIconKind = + | "queued" + | "success" + | "failed" + | "never-run"; + +export interface AutomationStatusPresentation { + label: string; + tone: AutomationStatusTone; + iconKind: AutomationStatusIconKind; +} + +export function getAutomationStatusPresentation({ + lastRunStatus, + lastTaskRunStatus, +}: AutomationStatusInput): AutomationStatusPresentation | null { + switch (lastTaskRunStatus) { + case "not_started": + case "queued": + return { + label: "Queued", + tone: "warning", + iconKind: "queued", + }; + case "started": + case "in_progress": + return null; + case "completed": + return { + label: "Success", + tone: "success", + iconKind: "success", + }; + case "failed": + case "cancelled": + return { + label: "Failed", + tone: "error", + iconKind: "failed", + }; + default: + break; + } + + switch (lastRunStatus) { + case "running": + return null; + case "success": + return { + label: "Success", + tone: "success", + iconKind: "success", + }; + case "failed": + return { + label: "Failed", + tone: "error", + iconKind: "failed", + }; + default: + return { + label: "Never run", + tone: "neutral", + iconKind: "never-run", + }; + } +} diff --git a/packages/core/src/inbox/artefacts.test.ts b/packages/core/src/inbox/artefacts.test.ts index a4e972613e..3d733ffbd3 100644 --- a/packages/core/src/inbox/artefacts.test.ts +++ b/packages/core/src/inbox/artefacts.test.ts @@ -1,11 +1,43 @@ -import type { SuggestedReviewer } from "@posthog/shared/types"; +import type { + AvailableSuggestedReviewer, + SuggestedReviewer, +} from "@posthog/shared/types"; import { describe, expect, it } from "vitest"; import { + buildReviewerOptions, extractSuggestedReviewers, + orderSuggestedReviewers, reviewerInitials, + reviewerMatchesAvailable, + reviewerOptionLabel, suggestedReviewerDisplayName, + toSuggestedReviewerWriteContent, } from "./artefacts"; +function makeReviewer( + partial: Partial = {}, +): SuggestedReviewer { + return { + github_login: "octocat", + github_name: "The Octocat", + relevant_commits: [], + user: null, + ...partial, + }; +} + +function makeAvailableReviewer( + partial: Partial = {}, +): AvailableSuggestedReviewer { + return { + uuid: "uuid-1", + name: "Ada Lovelace", + email: "ada@example.com", + github_login: "ada", + ...partial, + }; +} + describe("artefacts", () => { it("extracts suggested reviewers from artefacts", () => { const reviewers: SuggestedReviewer[] = [ @@ -46,4 +78,129 @@ describe("artefacts", () => { expect(reviewerInitials("Ben W.", null)).toBe("BW"); expect(reviewerInitials("", "ben@posthog.com")).toBe("BE"); }); + + it("moves the current user to the front", () => { + const reviewers = [ + makeReviewer({ + github_login: "a", + user: { + id: 1, + uuid: "uuid-a", + email: "a@posthog.com", + first_name: "a", + last_name: "", + }, + }), + makeReviewer({ + github_login: "me", + user: { + id: 2, + uuid: "uuid-me", + email: "me@posthog.com", + first_name: "me", + last_name: "", + }, + }), + ]; + + expect( + orderSuggestedReviewers(reviewers, "uuid-me").map( + (reviewer) => reviewer.github_login, + ), + ).toEqual(["me", "a"]); + }); + + it("deduplicates reviewer options and pins the current user first", () => { + const options = buildReviewerOptions( + [ + makeAvailableReviewer({ uuid: "b", name: "Bob" }), + makeAvailableReviewer({ uuid: "a", name: "Ada" }), + makeAvailableReviewer({ uuid: "a", name: "Ada duplicate" }), + ], + "b", + ); + + expect(options.map((option) => option.uuid)).toEqual(["b", "a"]); + }); + + it("labels the current reviewer", () => { + expect( + reviewerOptionLabel({ + uuid: "uuid-me", + name: "Ada", + email: "ada@example.com", + github_login: "ada", + isMe: true, + }), + ).toBe("Ada (Me)"); + }); + + it.each([ + { + name: "user uuid", + reviewer: makeReviewer({ + github_login: "", + user: { + id: 1, + uuid: "uuid-1", + email: "", + first_name: "", + last_name: "", + }, + }), + expected: true, + }, + { + name: "case-insensitive GitHub login", + reviewer: makeReviewer({ github_login: "ADA" }), + expected: true, + }, + { + name: "different reviewer", + reviewer: makeReviewer(), + expected: false, + }, + ])("matches an available reviewer by $name", ({ reviewer, expected }) => { + expect(reviewerMatchesAvailable(reviewer, makeAvailableReviewer())).toBe( + expected, + ); + }); + + it.each([ + { + name: "GitHub login", + reviewer: makeReviewer({ + github_login: "ada", + user: { + id: 1, + uuid: "uuid-1", + email: "", + first_name: "", + last_name: "", + }, + }), + expected: [{ github_login: "ada" }], + }, + { + name: "user uuid fallback", + reviewer: makeReviewer({ + github_login: "", + user: { + id: 1, + uuid: "uuid-1", + email: "", + first_name: "", + last_name: "", + }, + }), + expected: [{ user_uuid: "uuid-1" }], + }, + { + name: "unresolved reviewer", + reviewer: makeReviewer({ github_login: "" }), + expected: [], + }, + ])("builds write content from the $name", ({ reviewer, expected }) => { + expect(toSuggestedReviewerWriteContent([reviewer])).toEqual(expected); + }); }); diff --git a/packages/core/src/inbox/artefacts.ts b/packages/core/src/inbox/artefacts.ts index 1978f710bc..dc5b2eddb2 100644 --- a/packages/core/src/inbox/artefacts.ts +++ b/packages/core/src/inbox/artefacts.ts @@ -1,8 +1,18 @@ import type { + AvailableSuggestedReviewer, RepoSelectionArtefact, SuggestedReviewer, + SuggestedReviewerWriteEntry, } from "@posthog/shared/types"; +export interface ReviewerOption { + uuid: string; + name: string; + email: string; + github_login: string; + isMe: boolean; +} + function hasRepositoryContent( content: unknown, ): content is RepoSelectionArtefact["content"] { @@ -48,6 +58,82 @@ export function extractSuggestedReviewers( return artefact?.content ?? []; } +export function orderSuggestedReviewers( + reviewers: SuggestedReviewer[], + currentUserUuid: string | null | undefined, +): SuggestedReviewer[] { + if (!currentUserUuid) return reviewers; + const currentUserIndex = reviewers.findIndex( + (reviewer) => reviewer.user?.uuid === currentUserUuid, + ); + if (currentUserIndex <= 0) return reviewers; + return [ + reviewers[currentUserIndex], + ...reviewers.filter((_, index) => index !== currentUserIndex), + ]; +} + +export function buildReviewerOptions( + reviewers: AvailableSuggestedReviewer[], + currentUserUuid: string | undefined, +): ReviewerOption[] { + const seen = new Set(); + const options: ReviewerOption[] = []; + + for (const reviewer of reviewers) { + if (!reviewer.uuid || seen.has(reviewer.uuid)) continue; + seen.add(reviewer.uuid); + options.push({ + uuid: reviewer.uuid, + name: reviewer.name?.trim() || "", + email: reviewer.email?.trim() || "", + github_login: reviewer.github_login?.trim() || "", + isMe: reviewer.uuid === currentUserUuid, + }); + } + + options.sort((first, second) => { + if (first.isMe && !second.isMe) return -1; + if (!first.isMe && second.isMe) return 1; + return (first.name || first.email).localeCompare( + second.name || second.email, + ); + }); + + return options; +} + +export function reviewerOptionLabel(reviewer: ReviewerOption): string { + const base = reviewer.name || reviewer.email || "Unknown user"; + return reviewer.isMe ? `${base} (Me)` : base; +} + +export function reviewerMatchesAvailable( + reviewer: SuggestedReviewer, + available: AvailableSuggestedReviewer, +): boolean { + if (reviewer.user?.uuid && reviewer.user.uuid === available.uuid) { + return true; + } + return ( + !!reviewer.github_login && + !!available.github_login && + reviewer.github_login.toLowerCase() === available.github_login.toLowerCase() + ); +} + +export function toSuggestedReviewerWriteContent( + reviewers: SuggestedReviewer[], +): SuggestedReviewerWriteEntry[] { + return reviewers + .map((reviewer): SuggestedReviewerWriteEntry | null => { + if (reviewer.github_login) return { github_login: reviewer.github_login }; + if (reviewer.user?.uuid) return { user_uuid: reviewer.user.uuid }; + return null; + }) + .filter((entry): entry is SuggestedReviewerWriteEntry => entry !== null); +} + const AVATAR_PALETTE = [ "bg-(--orange-9) text-white", "bg-(--blue-9) text-white", diff --git a/packages/core/src/inbox/reportFiltering.test.ts b/packages/core/src/inbox/reportFiltering.test.ts index f2be318fac..3362ca5cde 100644 --- a/packages/core/src/inbox/reportFiltering.test.ts +++ b/packages/core/src/inbox/reportFiltering.test.ts @@ -6,8 +6,18 @@ import { buildSignalReportListOrdering, buildSuggestedReviewerFilterParam, filterReportsBySearch, + INBOX_PIPELINE_STATUS_FILTER, + INBOX_PIPELINE_STATUSES, } from "./reportFiltering"; +describe("inbox pipeline statuses", () => { + it("derives the API filter from the typed status list", () => { + expect(INBOX_PIPELINE_STATUS_FILTER).toBe( + INBOX_PIPELINE_STATUSES.join(","), + ); + }); +}); + function makeReport(overrides: Partial = {}): SignalReport { return { id: "1", diff --git a/packages/core/src/inbox/reportFiltering.ts b/packages/core/src/inbox/reportFiltering.ts index 06d36038b5..2a8271b2cb 100644 --- a/packages/core/src/inbox/reportFiltering.ts +++ b/packages/core/src/inbox/reportFiltering.ts @@ -5,12 +5,20 @@ import type { SignalReportStatus, } from "@posthog/shared/types"; +export const INBOX_PIPELINE_STATUSES = [ + "ready", + "pending_input", + "in_progress", + "failed", + "candidate", + "potential", +] as const satisfies readonly SignalReportStatus[]; + /** * Comma-separated statuses for the inbox query. We pull `failed` so the Runs * tab can surface failed runs in its Recently finished section. */ -export const INBOX_PIPELINE_STATUS_FILTER = - "potential,candidate,in_progress,ready,pending_input,failed"; +export const INBOX_PIPELINE_STATUS_FILTER = INBOX_PIPELINE_STATUSES.join(","); /** * Status filter for the Archive tab — the two terminal, not-in-inbox states: diff --git a/packages/core/src/sessions/cloudSessionConfig.ts b/packages/core/src/sessions/cloudSessionConfig.ts index ecb989b446..b6a9e4f8f4 100644 --- a/packages/core/src/sessions/cloudSessionConfig.ts +++ b/packages/core/src/sessions/cloudSessionConfig.ts @@ -1,6 +1,10 @@ import type { SessionConfigOption } from "@agentclientprotocol/sdk"; import type { Adapter, StoredLogEntry } from "@posthog/shared"; -import { getAvailableCodexModes, getAvailableModes } from "./executionModes"; +import { + DEFAULT_CLAUDE_EXECUTION_MODE, + getAvailableCodexModes, + getAvailableModes, +} from "./executionModes"; /** * Pure derivations of cloud session config options. No store or host access — @@ -54,7 +58,8 @@ export function buildCloudDefaultConfigOptions( ): SessionConfigOption[] { const modes = adapter === "codex" ? getAvailableCodexModes() : getAvailableModes(); - const fallbackMode = adapter === "codex" ? "auto" : "plan"; + const fallbackMode = + adapter === "codex" ? "auto" : DEFAULT_CLAUDE_EXECUTION_MODE; const currentMode = typeof initialMode === "string" && modes.some((mode) => mode.id === initialMode) diff --git a/packages/core/src/sessions/executionModes.ts b/packages/core/src/sessions/executionModes.ts index 4a32413c12..2ccbadf354 100644 --- a/packages/core/src/sessions/executionModes.ts +++ b/packages/core/src/sessions/executionModes.ts @@ -1,4 +1,4 @@ -import { CODEX_MODE_PRESETS } from "@posthog/shared"; +import { CODEX_MODE_PRESETS, type ExecutionMode } from "@posthog/shared"; export interface ModeInfo { id: string; @@ -6,6 +6,8 @@ export interface ModeInfo { description: string; } +export const DEFAULT_CLAUDE_EXECUTION_MODE: ExecutionMode = "plan"; + const availableModes: ModeInfo[] = [ { id: "default", diff --git a/packages/core/src/sessions/portableSessionEvents.test.ts b/packages/core/src/sessions/portableSessionEvents.test.ts new file mode 100644 index 0000000000..d777a0ec8b --- /dev/null +++ b/packages/core/src/sessions/portableSessionEvents.test.ts @@ -0,0 +1,71 @@ +import type { StoredLogEntry } from "@posthog/shared"; +import { describe, expect, it, vi } from "vitest"; +import { + convertStoredEntriesToPortableSessionEvents, + inferStoredLogEntryDirection, +} from "./portableSessionEvents"; + +describe("inferStoredLogEntryDirection", () => { + it.each([ + [ + "client requests", + { notification: { id: 1, method: "session/prompt" } }, + "client", + ], + ["agent responses", { notification: { id: 1, result: {} } }, "agent"], + [ + "agent notifications", + { notification: { method: "session/update" } }, + "agent", + ], + ["missing messages", {}, "agent"], + ] as const)("classifies %s", (_name, entry, expected) => { + expect(inferStoredLogEntryDirection(entry as StoredLogEntry)).toBe( + expected, + ); + }); +}); + +describe("convertStoredEntriesToPortableSessionEvents", () => { + it("projects session updates alongside their raw ACP message", () => { + const notification = { + update: { + sessionUpdate: "user_message_chunk", + content: { type: "text", text: "hello" }, + }, + }; + const events = convertStoredEntriesToPortableSessionEvents([ + { + type: "notification", + timestamp: "2026-07-21T12:00:00.000Z", + notification: { method: "session/update", params: notification }, + }, + ]); + + expect(events).toEqual([ + { + type: "acp_message", + direction: "agent", + ts: 1_784_635_200_000, + message: { method: "session/update", params: notification }, + }, + { + type: "session_update", + ts: 1_784_635_200_000, + notification, + }, + ]); + }); + + it("uses the current time when an entry has no timestamp", () => { + vi.useFakeTimers(); + vi.setSystemTime(new Date("2026-07-21T12:00:00.000Z")); + + const events = convertStoredEntriesToPortableSessionEvents([ + { type: "response", notification: { id: 1, result: {} } }, + ]); + + expect(events[0]?.ts).toBe(1_784_635_200_000); + vi.useRealTimers(); + }); +}); diff --git a/packages/core/src/sessions/portableSessionEvents.ts b/packages/core/src/sessions/portableSessionEvents.ts new file mode 100644 index 0000000000..5ee36c7642 --- /dev/null +++ b/packages/core/src/sessions/portableSessionEvents.ts @@ -0,0 +1,98 @@ +import type { JsonRpcMessage, StoredLogEntry } from "@posthog/shared"; + +export type PortableSessionToolCallStatus = + | "pending" + | "in_progress" + | "completed" + | "failed" + | null; + +export interface PortableSessionUpdate { + sessionUpdate?: string; + content?: { type: string; text: string }; + attachments?: Array<{ + kind: "image" | "document"; + uri: string; + fileName: string; + mimeType?: string; + }>; + title?: string; + toolCallId?: string; + status?: PortableSessionToolCallStatus; + rawInput?: Record; + rawOutput?: unknown; + entries?: Array<{ + content: string; + status: "pending" | "in_progress" | "completed" | "failed"; + priority: string; + }>; + _meta?: { + claudeCode?: { + toolName?: string; + parentToolCallId?: string; + }; + }; +} + +export interface PortableSessionNotification { + update?: PortableSessionUpdate; +} + +export interface PortableSessionAcpMessage { + type: "acp_message"; + direction: "client" | "agent"; + ts: number; + message: JsonRpcMessage; +} + +export interface PortableSessionUpdateEvent { + type: "session_update"; + ts: number; + notification: PortableSessionNotification; +} + +export type PortableSessionEvent = + | PortableSessionAcpMessage + | PortableSessionUpdateEvent; + +export function inferStoredLogEntryDirection( + entry: StoredLogEntry, +): "client" | "agent" { + const message = entry.notification; + if (!message) return "agent"; + if (message.id !== undefined && message.method !== undefined) return "client"; + return "agent"; +} + +export function convertStoredEntriesToPortableSessionEvents( + entries: readonly StoredLogEntry[], +): PortableSessionEvent[] { + const events: PortableSessionEvent[] = []; + + for (const entry of entries) { + const ts = entry.timestamp + ? new Date(entry.timestamp).getTime() + : Date.now(); + + events.push({ + type: "acp_message", + direction: inferStoredLogEntryDirection(entry), + ts, + message: (entry.notification ?? {}) as JsonRpcMessage, + }); + + if ( + entry.type === "notification" && + entry.notification?.method === "session/update" && + entry.notification.params + ) { + events.push({ + type: "session_update", + ts, + notification: entry.notification.params as PortableSessionNotification, + }); + } + } + + return events; +} diff --git a/packages/core/src/sessions/sessionActivity.test.ts b/packages/core/src/sessions/sessionActivity.test.ts new file mode 100644 index 0000000000..db795f2118 --- /dev/null +++ b/packages/core/src/sessions/sessionActivity.test.ts @@ -0,0 +1,152 @@ +import { describe, expect, it } from "vitest"; +import type { + PortableSessionEvent, + PortableSessionToolCallStatus, +} from "./portableSessionEvents"; +import { + countUserMessages, + getSessionActivityPhase, + isSessionAwaitingUserInput, +} from "./sessionActivity"; + +function userMessage(ts = 1): PortableSessionEvent { + return { + type: "session_update", + ts, + notification: { + update: { + sessionUpdate: "user_message_chunk", + content: { type: "text", text: "Yes" }, + }, + }, + }; +} + +function questionToolCall( + status: PortableSessionToolCallStatus, + sessionUpdate = "tool_call", +): PortableSessionEvent { + return { + type: "session_update", + ts: 1, + notification: { + update: { + sessionUpdate, + toolCallId: "question-1", + status, + rawInput: { questions: [{ question: "Proceed?", options: [] }] }, + _meta: { claudeCode: { toolName: "AskUserQuestion" } }, + }, + }, + }; +} + +function acpNotification(method: string): PortableSessionEvent { + return { + type: "acp_message", + direction: "agent", + ts: 1, + message: { method }, + }; +} + +describe("isSessionAwaitingUserInput", () => { + it("tracks question tools until a metadata-free completion update", () => { + const completion: PortableSessionEvent = { + type: "session_update", + ts: 2, + notification: { + update: { + sessionUpdate: "tool_call_update", + toolCallId: "question-1", + status: "completed", + }, + }, + }; + + expect( + isSessionAwaitingUserInput([questionToolCall("pending"), completion]), + ).toBe(false); + }); + + it("clears questions when the user responds", () => { + expect( + isSessionAwaitingUserInput([questionToolCall("pending"), userMessage(2)]), + ).toBe(false); + }); + + it("honors explicit waiting and terminal backend markers", () => { + expect( + isSessionAwaitingUserInput([ + acpNotification("_posthog/awaiting_user_input"), + ]), + ).toBe(true); + expect( + isSessionAwaitingUserInput([ + acpNotification("_posthog/awaiting_user_input"), + acpNotification("_posthog/turn_complete"), + ]), + ).toBe(false); + }); +}); + +describe("countUserMessages", () => { + it("counts only projected user message updates", () => { + expect( + countUserMessages([ + userMessage(), + questionToolCall("pending"), + userMessage(2), + ]), + ).toBe(2); + }); +}); + +describe("getSessionActivityPhase", () => { + it.each([ + ["retrying", true, undefined, "connecting"], + [ + "awaiting agent output", + false, + { isPromptPending: true, awaitingAgentOutput: true }, + "connecting", + ], + [ + "working", + false, + { isPromptPending: true, awaitingAgentOutput: false }, + "working", + ], + [ + "not pending", + false, + { isPromptPending: false, awaitingAgentOutput: false }, + "idle", + ], + [ + "terminal", + false, + { + isPromptPending: true, + awaitingAgentOutput: false, + terminalStatus: "completed" as const, + }, + "idle", + ], + [ + "waiting for user", + false, + { + isPromptPending: true, + awaitingAgentOutput: false, + events: [questionToolCall("pending")], + }, + "idle", + ], + ] as const)( + "returns the expected phase while %s", + (_name, retrying, session, expected) => { + expect(getSessionActivityPhase({ retrying, session })).toBe(expected); + }, + ); +}); diff --git a/packages/core/src/sessions/sessionActivity.ts b/packages/core/src/sessions/sessionActivity.ts new file mode 100644 index 0000000000..6e2804ff47 --- /dev/null +++ b/packages/core/src/sessions/sessionActivity.ts @@ -0,0 +1,129 @@ +import { isNotification, POSTHOG_NOTIFICATIONS } from "./acpNotifications"; +import type { + PortableSessionEvent, + PortableSessionNotification, + PortableSessionToolCallStatus, +} from "./portableSessionEvents"; + +export type SessionActivityPhase = "idle" | "connecting" | "working"; + +export interface SessionActivityState { + isPromptPending?: boolean; + awaitingAgentOutput?: boolean; + terminalStatus?: "failed" | "completed"; + events?: readonly PortableSessionEvent[]; +} + +function isQuestionNotification( + notification: PortableSessionNotification, +): boolean { + const update = notification.update; + if (!update) return false; + + const rawToolName = update._meta?.claudeCode?.toolName; + if (typeof rawToolName === "string" && /question/i.test(rawToolName)) { + return true; + } + + const rawInput = update.rawInput; + if (!rawInput) return false; + if (Array.isArray(rawInput.questions)) return true; + + const nestedInput = rawInput.input; + return ( + typeof nestedInput === "object" && + nestedInput !== null && + Array.isArray((nestedInput as { questions?: unknown }).questions) + ); +} + +function isPendingQuestionStatus( + status: PortableSessionToolCallStatus | undefined, +): boolean { + return status === null || status === "pending" || status === "in_progress"; +} + +export function isSessionAwaitingUserInput( + events: readonly PortableSessionEvent[] = [], +): boolean { + let awaitingUserInput = false; + const questionStatuses = new Map< + string, + PortableSessionToolCallStatus | undefined + >(); + + for (const event of events) { + if (event.type === "session_update") { + const update = event.notification.update; + const sessionUpdate = update?.sessionUpdate; + + if (sessionUpdate === "user_message_chunk") { + awaitingUserInput = false; + questionStatuses.clear(); + continue; + } + + if ( + sessionUpdate === "tool_call" || + sessionUpdate === "tool_call_update" + ) { + const toolCallId = update?.toolCallId; + const isKnownQuestion = toolCallId + ? questionStatuses.has(toolCallId) + : false; + if (!isKnownQuestion && !isQuestionNotification(event.notification)) { + continue; + } + + questionStatuses.set( + toolCallId ?? `question-${event.ts}`, + update?.status, + ); + awaitingUserInput = [...questionStatuses.values()].some( + isPendingQuestionStatus, + ); + } + + continue; + } + + const method = "method" in event.message ? event.message.method : undefined; + if (method === "_posthog/awaiting_user_input") { + awaitingUserInput = true; + continue; + } + + if ( + isNotification(method, POSTHOG_NOTIFICATIONS.TURN_COMPLETE) || + isNotification(method, POSTHOG_NOTIFICATIONS.TASK_COMPLETE) || + isNotification(method, POSTHOG_NOTIFICATIONS.ERROR) + ) { + awaitingUserInput = false; + questionStatuses.clear(); + } + } + + return awaitingUserInput; +} + +export function countUserMessages( + events: readonly PortableSessionEvent[] = [], +): number { + return events.filter( + (event) => + event.type === "session_update" && + event.notification.update?.sessionUpdate === "user_message_chunk", + ).length; +} + +export function getSessionActivityPhase(args: { + retrying: boolean; + session?: SessionActivityState | null; +}): SessionActivityPhase { + const { retrying, session } = args; + + if (retrying) return "connecting"; + if (!session?.isPromptPending || session.terminalStatus) return "idle"; + if (isSessionAwaitingUserInput(session.events)) return "idle"; + return session.awaitingAgentOutput ? "connecting" : "working"; +} diff --git a/packages/core/src/tasks/taskActivity.test.ts b/packages/core/src/tasks/taskActivity.test.ts new file mode 100644 index 0000000000..35d1eadaff --- /dev/null +++ b/packages/core/src/tasks/taskActivity.test.ts @@ -0,0 +1,132 @@ +import type { Task } from "@posthog/shared/domain-types"; +import { describe, expect, it } from "vitest"; +import { filterAndSortTasks, taskActivityTimestamp } from "./taskActivity"; + +function makeTask(overrides: Partial = {}): Task { + return { + id: "task-1", + task_number: 1, + slug: "task-1", + title: "A real task", + description: "Do the thing", + created_at: "2026-01-01T00:00:00Z", + updated_at: "2026-01-02T00:00:00Z", + origin_product: "tasks", + ...overrides, + }; +} + +describe("taskActivityTimestamp", () => { + it("uses creation time in created mode", () => { + const task = makeTask({ + created_at: "2026-01-01T00:00:00Z", + updated_at: "2026-01-03T00:00:00Z", + }); + + expect(taskActivityTimestamp(task, "created")).toBe( + new Date("2026-01-01T00:00:00Z").getTime(), + ); + }); + + it("uses the latest task or run update in updated mode", () => { + const task = makeTask({ + updated_at: "2026-01-02T00:00:00Z", + latest_run: { + id: "run-1", + task: "task-1", + team: 1, + branch: null, + status: "completed", + log_url: "", + error_message: null, + output: null, + state: {}, + created_at: "2026-01-01T00:00:00Z", + updated_at: "2026-01-04T00:00:00Z", + completed_at: "2026-01-04T00:00:00Z", + }, + }); + + expect(taskActivityTimestamp(task, "updated")).toBe( + new Date("2026-01-04T00:00:00Z").getTime(), + ); + }); +}); + +describe("filterAndSortTasks", () => { + it.each([ + { title: "", description: "" }, + { title: " ", description: "\n\t" }, + ])("hides contentless placeholder tasks", ({ title, description }) => { + const placeholder = makeTask({ id: "placeholder", title, description }); + const realTask = makeTask({ id: "real" }); + + expect( + filterAndSortTasks([placeholder, realTask], "updated", false, "").map( + (task) => task.id, + ), + ).toEqual(["real"]); + }); + + it("selects internal or external tasks", () => { + const externalTask = makeTask({ id: "external", internal: false }); + const internalTask = makeTask({ id: "internal", internal: true }); + + expect( + filterAndSortTasks( + [externalTask, internalTask], + "updated", + false, + "", + ).map((task) => task.id), + ).toEqual(["external"]); + expect( + filterAndSortTasks([externalTask, internalTask], "updated", true, "").map( + (task) => task.id, + ), + ).toEqual(["internal"]); + }); + + it.each([ + ["title", { title: "Fix Login" }], + ["slug", { slug: "fix-login" }], + ["description", { description: "Fix Login" }], + ] as const)("matches a case-insensitive %s filter", (_field, overrides) => { + const matchingTask = makeTask({ id: "matching", ...overrides }); + const otherTask = makeTask({ id: "other", title: "Unrelated" }); + + expect( + filterAndSortTasks( + [otherTask, matchingTask], + "updated", + false, + "LOGIN", + ).map((task) => task.id), + ).toEqual(["matching"]); + }); + + it("sorts by the selected activity timestamp without mutating input", () => { + const olderCreated = makeTask({ + id: "older-created", + created_at: "2026-01-01T00:00:00Z", + updated_at: "2026-01-04T00:00:00Z", + }); + const newerCreated = makeTask({ + id: "newer-created", + created_at: "2026-01-02T00:00:00Z", + updated_at: "2026-01-03T00:00:00Z", + }); + const tasks = [olderCreated, newerCreated]; + + expect( + filterAndSortTasks(tasks, "created", false, "").map((task) => task.id), + ).toEqual(["newer-created", "older-created"]); + expect( + filterAndSortTasks(tasks, "updated", false, "").map((task) => task.id), + ).toEqual(["older-created", "newer-created"]); + expect(tasks.map((task) => task.id)).toEqual([ + "older-created", + "newer-created", + ]); + }); +}); diff --git a/packages/core/src/tasks/taskActivity.ts b/packages/core/src/tasks/taskActivity.ts new file mode 100644 index 0000000000..6a52168e63 --- /dev/null +++ b/packages/core/src/tasks/taskActivity.ts @@ -0,0 +1,45 @@ +import { isContentlessTask, type Task } from "@posthog/shared/domain-types"; + +export type TaskActivitySortMode = "created" | "updated"; + +export function taskActivityTimestamp( + task: Pick, + sortMode: TaskActivitySortMode, +): number { + if (sortMode === "created") { + return new Date(task.created_at).getTime(); + } + + const runUpdatedAt = task.latest_run?.updated_at; + return Math.max( + runUpdatedAt ? new Date(runUpdatedAt).getTime() : 0, + new Date(task.updated_at ?? task.created_at).getTime(), + ); +} + +export function filterAndSortTasks( + tasks: readonly Task[], + sortMode: TaskActivitySortMode, + showInternal: boolean, + filter: string, +): Task[] { + const normalizedFilter = filter.toLowerCase(); + + return tasks + .filter((task) => !isContentlessTask(task)) + .filter((task) => + showInternal ? task.internal === true : task.internal !== true, + ) + .filter( + (task) => + !normalizedFilter || + task.title.toLowerCase().includes(normalizedFilter) || + task.slug.toLowerCase().includes(normalizedFilter) || + task.description?.toLowerCase().includes(normalizedFilter), + ) + .sort( + (firstTask, secondTask) => + taskActivityTimestamp(secondTask, sortMode) - + taskActivityTimestamp(firstTask, sortMode), + ); +} diff --git a/packages/core/src/tasks/taskArchive.test.ts b/packages/core/src/tasks/taskArchive.test.ts new file mode 100644 index 0000000000..ab22bbd459 --- /dev/null +++ b/packages/core/src/tasks/taskArchive.test.ts @@ -0,0 +1,44 @@ +import type { Task, TaskRunStatus } from "@posthog/shared/domain-types"; +import { describe, expect, it } from "vitest"; +import { isTaskRunning } from "./taskArchive"; + +function makeTask(status?: TaskRunStatus): Pick { + return { + latest_run: status + ? { + id: "run-1", + task: "task-1", + team: 1, + branch: null, + status, + log_url: "", + error_message: null, + output: null, + state: {}, + created_at: "2026-01-01T00:00:00Z", + updated_at: "2026-01-01T00:00:00Z", + completed_at: null, + } + : undefined, + }; +} + +describe("isTaskRunning", () => { + it("returns false when a task has no run", () => { + expect(isTaskRunning(makeTask())).toBe(false); + }); + + it.each(["not_started", "queued", "in_progress"] as const)( + "returns true for %s", + (status) => { + expect(isTaskRunning(makeTask(status))).toBe(true); + }, + ); + + it.each(["completed", "failed", "cancelled"] as const)( + "returns false for %s", + (status) => { + expect(isTaskRunning(makeTask(status))).toBe(false); + }, + ); +}); diff --git a/packages/core/src/tasks/taskArchive.ts b/packages/core/src/tasks/taskArchive.ts new file mode 100644 index 0000000000..db52bb6310 --- /dev/null +++ b/packages/core/src/tasks/taskArchive.ts @@ -0,0 +1,6 @@ +import { isTerminalStatus, type Task } from "@posthog/shared/domain-types"; + +export function isTaskRunning(task: Pick): boolean { + const status = task.latest_run?.status; + return status !== undefined && !isTerminalStatus(status); +} diff --git a/packages/core/src/tasks/taskStatusPresentation.test.ts b/packages/core/src/tasks/taskStatusPresentation.test.ts new file mode 100644 index 0000000000..a7895cf59e --- /dev/null +++ b/packages/core/src/tasks/taskStatusPresentation.test.ts @@ -0,0 +1,69 @@ +import type { Task, TaskRun } from "@posthog/shared/domain-types"; +import { describe, expect, it } from "vitest"; +import { getTaskStatusPresentationKind } from "./taskStatusPresentation"; + +function makeTask(latestRun?: Partial): Pick { + return { + latest_run: latestRun + ? { + id: "run-1", + task: "task-1", + team: 1, + branch: null, + status: "not_started", + log_url: "", + error_message: null, + output: null, + state: {}, + created_at: "2026-01-01T00:00:00Z", + updated_at: "2026-01-01T00:00:00Z", + completed_at: null, + ...latestRun, + } + : undefined, + }; +} + +describe("getTaskStatusPresentationKind", () => { + it("prioritizes a pull request over cloud presentation", () => { + expect( + getTaskStatusPresentationKind( + makeTask({ + environment: "cloud", + status: "in_progress", + output: { pr_url: "https://github.com/PostHog/code/pull/123" }, + }), + ), + ).toBe("pr"); + }); + + it.each([ + "not_started", + "queued", + "in_progress", + "completed", + "failed", + "cancelled", + ] as const)("uses chat presentation for cloud status %s", (status) => { + expect( + getTaskStatusPresentationKind(makeTask({ environment: "cloud", status })), + ).toBe("chat"); + }); + + it.each([ + ["completed", "completed"], + ["failed", "failed"], + ["in_progress", "running"], + ["queued", "started"], + ["not_started", "chat"], + ["cancelled", "chat"], + ] as const)("maps local status %s to %s", (status, expected) => { + expect( + getTaskStatusPresentationKind(makeTask({ environment: "local", status })), + ).toBe(expected); + }); + + it("falls back to chat when a task has no run", () => { + expect(getTaskStatusPresentationKind(makeTask())).toBe("chat"); + }); +}); diff --git a/packages/core/src/tasks/taskStatusPresentation.ts b/packages/core/src/tasks/taskStatusPresentation.ts new file mode 100644 index 0000000000..968a17bb76 --- /dev/null +++ b/packages/core/src/tasks/taskStatusPresentation.ts @@ -0,0 +1,37 @@ +import { readPrUrls } from "@posthog/shared"; +import type { Task } from "@posthog/shared/domain-types"; + +export type TaskStatusPresentationKind = + | "pr" + | "completed" + | "failed" + | "running" + | "started" + | "chat"; + +export function getTaskStatusPresentationKind( + task: Pick, +): TaskStatusPresentationKind { + const latestRun = task.latest_run; + + if (readPrUrls(latestRun?.output)[0]) { + return "pr"; + } + + if (latestRun?.environment === "cloud") { + return "chat"; + } + + switch (latestRun?.status) { + case "completed": + return "completed"; + case "failed": + return "failed"; + case "in_progress": + return "running"; + case "queued": + return "started"; + default: + return "chat"; + } +} From 53e538364092cc252cd383a9c8595a5c72268b8a Mon Sep 17 00:00:00 2001 From: Richard Solomou Date: Fri, 24 Jul 2026 19:19:34 +0300 Subject: [PATCH 35/45] refactor(core): rename cloud task service as engine Generated-By: PostHog Code Task-Id: c1bbe3cf-742b-4b24-bf96-d11a18b4cf22 --- apps/web/src/web-container.ts | 2 +- .../core/src/cloud-task/{cloud-task.ts => cloud-task-engine.ts} | 0 packages/core/src/cloud-task/cloud-task.module.ts | 2 +- packages/core/src/cloud-task/cloud-task.test.ts | 2 +- packages/core/src/handoff/handoff.ts | 2 +- packages/host-router/src/routers/cloud-task.router.ts | 2 +- 6 files changed, 5 insertions(+), 5 deletions(-) rename packages/core/src/cloud-task/{cloud-task.ts => cloud-task-engine.ts} (100%) diff --git a/apps/web/src/web-container.ts b/apps/web/src/web-container.ts index 3075d139ce..cf17cd493e 100644 --- a/apps/web/src/web-container.ts +++ b/apps/web/src/web-container.ts @@ -27,8 +27,8 @@ import { } from "@posthog/core/auth/identifiers"; import { canvasCoreModule } from "@posthog/core/canvas/canvas.module"; import { taskThreadCoreModule } from "@posthog/core/canvas/taskThread.module"; -import type { CloudTaskService } from "@posthog/core/cloud-task/cloud-task"; import { cloudTaskModule } from "@posthog/core/cloud-task/cloud-task.module"; +import type { CloudTaskService } from "@posthog/core/cloud-task/cloud-task-engine"; import { CLOUD_TASK_AUTH, CLOUD_TASK_SERVICE, diff --git a/packages/core/src/cloud-task/cloud-task.ts b/packages/core/src/cloud-task/cloud-task-engine.ts similarity index 100% rename from packages/core/src/cloud-task/cloud-task.ts rename to packages/core/src/cloud-task/cloud-task-engine.ts diff --git a/packages/core/src/cloud-task/cloud-task.module.ts b/packages/core/src/cloud-task/cloud-task.module.ts index 02f0cc91db..464011d14f 100644 --- a/packages/core/src/cloud-task/cloud-task.module.ts +++ b/packages/core/src/cloud-task/cloud-task.module.ts @@ -1,5 +1,5 @@ import { ContainerModule } from "inversify"; -import { CloudTaskService } from "./cloud-task"; +import { CloudTaskService } from "./cloud-task-engine"; import { CLOUD_TASK_SERVICE } from "./identifiers"; export const cloudTaskModule = new ContainerModule(({ bind }) => { diff --git a/packages/core/src/cloud-task/cloud-task.test.ts b/packages/core/src/cloud-task/cloud-task.test.ts index a7fbf37c66..9002f948ff 100644 --- a/packages/core/src/cloud-task/cloud-task.test.ts +++ b/packages/core/src/cloud-task/cloud-task.test.ts @@ -22,7 +22,7 @@ const fetchRouter = vi.hoisted(() => }), ); -import { CloudTaskService } from "./cloud-task"; +import { CloudTaskService } from "./cloud-task-engine"; const mockAuthService = { authenticatedFetch: vi.fn(), diff --git a/packages/core/src/handoff/handoff.ts b/packages/core/src/handoff/handoff.ts index c85c2fcfa8..2b8d477099 100644 --- a/packages/core/src/handoff/handoff.ts +++ b/packages/core/src/handoff/handoff.ts @@ -5,7 +5,7 @@ import { TypedEventEmitter, } from "@posthog/shared"; import { inject, injectable } from "inversify"; -import type { CloudTaskService } from "../cloud-task/cloud-task"; +import type { CloudTaskService } from "../cloud-task/cloud-task-engine"; import { CLOUD_TASK_SERVICE } from "../cloud-task/identifiers"; import { HandoffSaga, type HandoffSagaDeps } from "./handoff-saga"; import { diff --git a/packages/host-router/src/routers/cloud-task.router.ts b/packages/host-router/src/routers/cloud-task.router.ts index 15d577ce59..4546995404 100644 --- a/packages/host-router/src/routers/cloud-task.router.ts +++ b/packages/host-router/src/routers/cloud-task.router.ts @@ -1,4 +1,4 @@ -import type { CloudTaskService } from "@posthog/core/cloud-task/cloud-task"; +import type { CloudTaskService } from "@posthog/core/cloud-task/cloud-task-engine"; import { CLOUD_TASK_SERVICE } from "@posthog/core/cloud-task/identifiers"; import { CloudTaskEvent, From 2245a2e61032c11dc77773d4fcd64afc9a121e2a Mon Sep 17 00:00:00 2001 From: Richard Solomou Date: Fri, 24 Jul 2026 19:20:17 +0300 Subject: [PATCH 36/45] refactor(core): extract portable cloud task engine Generated-By: PostHog Code Task-Id: c1bbe3cf-742b-4b24-bf96-d11a18b4cf22 --- apps/web/src/web-container.ts | 2 +- .../core/src/cloud-task/cloud-task-engine.ts | 96 ++++++++++++------- .../src/cloud-task/cloud-task-service.test.ts | 27 ++++++ .../core/src/cloud-task/cloud-task.module.ts | 2 +- .../core/src/cloud-task/cloud-task.test.ts | 47 +++++---- packages/core/src/cloud-task/cloud-task.ts | 35 +++++++ packages/core/src/cloud-task/schemas.ts | 22 ++--- packages/core/src/handoff/handoff.ts | 2 +- packages/core/vitest.config.ts | 23 +++++ .../src/routers/cloud-task.router.ts | 2 +- 10 files changed, 183 insertions(+), 75 deletions(-) create mode 100644 packages/core/src/cloud-task/cloud-task-service.test.ts create mode 100644 packages/core/src/cloud-task/cloud-task.ts create mode 100644 packages/core/vitest.config.ts diff --git a/apps/web/src/web-container.ts b/apps/web/src/web-container.ts index cf17cd493e..3075d139ce 100644 --- a/apps/web/src/web-container.ts +++ b/apps/web/src/web-container.ts @@ -27,8 +27,8 @@ import { } from "@posthog/core/auth/identifiers"; import { canvasCoreModule } from "@posthog/core/canvas/canvas.module"; import { taskThreadCoreModule } from "@posthog/core/canvas/taskThread.module"; +import type { CloudTaskService } from "@posthog/core/cloud-task/cloud-task"; import { cloudTaskModule } from "@posthog/core/cloud-task/cloud-task.module"; -import type { CloudTaskService } from "@posthog/core/cloud-task/cloud-task-engine"; import { CLOUD_TASK_AUTH, CLOUD_TASK_SERVICE, diff --git a/packages/core/src/cloud-task/cloud-task-engine.ts b/packages/core/src/cloud-task/cloud-task-engine.ts index 22b5a1ddab..006b0ab6d2 100644 --- a/packages/core/src/cloud-task/cloud-task-engine.ts +++ b/packages/core/src/cloud-task/cloud-task-engine.ts @@ -1,37 +1,24 @@ +import type { RootLogger, ScopedLogger } from "@posthog/di/logger"; +import type { IAnalytics } from "@posthog/platform/analytics"; import { - ROOT_LOGGER, - type RootLogger, - type ScopedLogger, -} from "@posthog/di/logger"; -import { - ANALYTICS_SERVICE, - type IAnalytics, -} from "@posthog/platform/analytics"; -import type { StoredLogEntry } from "@posthog/shared"; -import { + type CloudTaskPermissionRequestUpdate, + isTerminalStatus, mcpToolKey, posthogToolMeta, + type StoredLogEntry, serializeError, + type TaskRunStatus, TypedEventEmitter, } from "@posthog/shared"; import { ANALYTICS_EVENTS } from "@posthog/shared/analytics-events"; -import { inject, injectable, optional, preDestroy } from "inversify"; -import type { CloudTaskPermissionRequestUpdate } from "./cloud-task-types"; -import { - CLOUD_TASK_AUTH, - type ICloudTaskAuth, - MCP_RELAY_EXECUTOR, - type McpRelayExecutor, -} from "./identifiers"; +import type { ICloudTaskAuth, McpRelayExecutor } from "./identifiers"; import { CloudTaskEvent, type CloudTaskEvents, - isTerminalStatus, type SendCommandInput, type SendCommandOutput, type StopInput, type StopOutput, - type TaskRunStatus, type WatchInput, } from "./schemas"; import { type SseEvent, SseEventParser } from "./sse-parser"; @@ -435,23 +422,45 @@ function sandboxAlivePayload(watcher: { lastSandboxAlive: boolean | null }): { : { sandboxAlive: watcher.lastSandboxAlive }; } -@injectable() -export class CloudTaskService extends TypedEventEmitter { +export interface CloudTaskEngineDependencies { + auth: ICloudTaskAuth; + analytics: IAnalytics; + logger: RootLogger; + mcpRelayExecutor?: McpRelayExecutor | null; + streamFetch?: CloudTaskFetch; +} + +export type CloudTaskFetch = ( + input: string | URL | Request, + init?: RequestInit, +) => Promise; + +export function createCloudTaskEngine( + dependencies: CloudTaskEngineDependencies, +): CloudTaskEngine { + return new CloudTaskEngine(dependencies); +} + +export class CloudTaskEngine extends TypedEventEmitter { private watchers = new Map(); private readonly log: ScopedLogger; - - constructor( - @inject(CLOUD_TASK_AUTH) - private readonly auth: ICloudTaskAuth, - @inject(ANALYTICS_SERVICE) - private readonly analytics: IAnalytics, - @inject(ROOT_LOGGER) - logger: RootLogger, - @inject(MCP_RELAY_EXECUTOR) - @optional() - private readonly mcpRelayExecutor: McpRelayExecutor | null = null, - ) { + private readonly auth: ICloudTaskAuth; + private readonly analytics: IAnalytics; + private readonly mcpRelayExecutor: McpRelayExecutor | null; + private readonly streamFetch: CloudTaskFetch; + + constructor({ + auth, + analytics, + logger, + mcpRelayExecutor = null, + streamFetch = globalThis.fetch.bind(globalThis), + }: CloudTaskEngineDependencies) { super(); + this.auth = auth; + this.analytics = analytics; + this.mcpRelayExecutor = mcpRelayExecutor; + this.streamFetch = streamFetch; this.log = logger.scope("cloud-task"); } @@ -770,6 +779,22 @@ export class CloudTaskService extends TypedEventEmitter { void this.bootstrapWatcher(key); } + reconnectIfDisconnected(taskId: string, runId: string): void { + const key = watcherKey(taskId, runId); + const watcher = this.watchers.get(key); + if ( + !watcher || + watcher.sseAbortController || + watcher.reconnectTimeoutId || + watcher.isBootstrapping || + isTerminalStatus(watcher.lastStatus) + ) { + return; + } + + void this.connectSse(key); + } + // Resets a watcher to its pre-bootstrap state so bootstrapWatcher can rebuild it from server truth. private resetWatcherForRebootstrap(watcher: WatcherState): void { watcher.reconnectAttempts = 0; @@ -959,7 +984,6 @@ export class CloudTaskService extends TypedEventEmitter { } } - @preDestroy() unwatchAll(): void { for (const key of [...this.watchers.keys()]) { this.stopWatcher(key); @@ -1306,7 +1330,7 @@ export class CloudTaskService extends TypedEventEmitter { try { // The proxy authenticates with the run-scoped Bearer token; the Django leg uses the session. const response = usingProxy - ? await fetch(url.toString(), { + ? await this.streamFetch(url.toString(), { method: "GET", headers, signal: controller.signal, diff --git a/packages/core/src/cloud-task/cloud-task-service.test.ts b/packages/core/src/cloud-task/cloud-task-service.test.ts new file mode 100644 index 0000000000..644f9dd7cb --- /dev/null +++ b/packages/core/src/cloud-task/cloud-task-service.test.ts @@ -0,0 +1,27 @@ +import { describe, expect, it, vi } from "vitest"; +import { CloudTaskService } from "./cloud-task"; +import { CloudTaskEngine } from "./cloud-task-engine"; + +describe("CloudTaskService", () => { + it("preserves the injectable service API as a thin engine wrapper", () => { + const scopedLog = { + debug: vi.fn(), + info: vi.fn(), + warn: vi.fn(), + error: vi.fn(), + }; + const service = new CloudTaskService( + { + authenticatedFetch: vi.fn(), + getCloudContext: vi.fn(), + }, + { track: vi.fn() } as never, + { ...scopedLog, scope: vi.fn(() => scopedLog) }, + ); + + expect(service).toBeInstanceOf(CloudTaskEngine); + expect(service.watch).toBeTypeOf("function"); + expect(service.retry).toBeTypeOf("function"); + expect(service.unwatchAll).toBeTypeOf("function"); + }); +}); diff --git a/packages/core/src/cloud-task/cloud-task.module.ts b/packages/core/src/cloud-task/cloud-task.module.ts index 464011d14f..02f0cc91db 100644 --- a/packages/core/src/cloud-task/cloud-task.module.ts +++ b/packages/core/src/cloud-task/cloud-task.module.ts @@ -1,5 +1,5 @@ import { ContainerModule } from "inversify"; -import { CloudTaskService } from "./cloud-task-engine"; +import { CloudTaskService } from "./cloud-task"; import { CLOUD_TASK_SERVICE } from "./identifiers"; export const cloudTaskModule = new ContainerModule(({ bind }) => { diff --git a/packages/core/src/cloud-task/cloud-task.test.ts b/packages/core/src/cloud-task/cloud-task.test.ts index 9002f948ff..4178bd3d78 100644 --- a/packages/core/src/cloud-task/cloud-task.test.ts +++ b/packages/core/src/cloud-task/cloud-task.test.ts @@ -5,14 +5,17 @@ const mockNetFetch = vi.hoisted(() => vi.fn()); const mockStreamFetch = vi.hoisted(() => vi.fn()); const mockStreamTokenFetch = vi.hoisted(() => vi.fn()); -// The service now uses global fetch for BOTH authenticated API calls (JSON) -// and SSE streaming. The two used to be distinct (net.fetch vs global fetch). // Route by URL: /stream_token/ → token mock (read-leg resolution), the stream leg // (Django /stream/ or proxy /v1/runs/:run/stream) → stream mock, everything else → API mock. // The token mock has a Django-path default so existing fixtures (which never set it) are untouched. const fetchRouter = vi.hoisted(() => - vi.fn((input: string | Request, init?: RequestInit) => { - const url = typeof input === "string" ? input : input.url; + vi.fn((input: string | URL | Request, init?: RequestInit) => { + const url = + typeof input === "string" + ? input + : input instanceof URL + ? input.toString() + : input.url; const impl = url.includes("/stream_token/") ? mockStreamTokenFetch : /\/stream(\/|\?|$)/.test(url) @@ -22,7 +25,10 @@ const fetchRouter = vi.hoisted(() => }), ); -import { CloudTaskService } from "./cloud-task-engine"; +import { + type CloudTaskEngine, + createCloudTaskEngine, +} from "./cloud-task-engine"; const mockAuthService = { authenticatedFetch: vi.fn(), @@ -86,8 +92,8 @@ async function waitFor( } } -describe("CloudTaskService", () => { - let service: CloudTaskService; +describe("CloudTaskEngine", () => { + let service: CloudTaskEngine; beforeEach(() => { const scopedLog = { @@ -98,11 +104,12 @@ describe("CloudTaskService", () => { }; const loggerMock = { ...scopedLog, scope: vi.fn(() => scopedLog) }; const analyticsMock = { track: vi.fn() }; - service = new CloudTaskService( - mockAuthService as never, - analyticsMock as never, - loggerMock, - ); + service = createCloudTaskEngine({ + auth: mockAuthService as never, + analytics: analyticsMock as never, + logger: loggerMock, + streamFetch: fetchRouter, + }); mockNetFetch.mockReset(); mockStreamFetch.mockReset(); mockStreamTokenFetch.mockReset(); @@ -3077,8 +3084,8 @@ describe("CloudTaskService", () => { }); }); -describe("CloudTaskService MCP relay", () => { - let relayService: CloudTaskService; +describe("CloudTaskEngine MCP relay", () => { + let relayService: CloudTaskEngine; let mcpRelayExecutor: { execute: ReturnType; closeRun: ReturnType; @@ -3099,12 +3106,12 @@ describe("CloudTaskService MCP relay", () => { })), closeRun: vi.fn(async () => {}), }; - relayService = new CloudTaskService( - mockAuthService as never, - analyticsMock as never, - loggerMock, - mcpRelayExecutor as never, - ); + relayService = createCloudTaskEngine({ + auth: mockAuthService as never, + analytics: analyticsMock as never, + logger: loggerMock, + mcpRelayExecutor: mcpRelayExecutor as never, + }); mockNetFetch.mockReset(); mockStreamFetch.mockReset(); diff --git a/packages/core/src/cloud-task/cloud-task.ts b/packages/core/src/cloud-task/cloud-task.ts new file mode 100644 index 0000000000..1e2003d85a --- /dev/null +++ b/packages/core/src/cloud-task/cloud-task.ts @@ -0,0 +1,35 @@ +import { ROOT_LOGGER, type RootLogger } from "@posthog/di/logger"; +import { + ANALYTICS_SERVICE, + type IAnalytics, +} from "@posthog/platform/analytics"; +import { inject, injectable, optional, preDestroy } from "inversify"; +import { CloudTaskEngine } from "./cloud-task-engine"; +import { + CLOUD_TASK_AUTH, + type ICloudTaskAuth, + MCP_RELAY_EXECUTOR, + type McpRelayExecutor, +} from "./identifiers"; + +@injectable() +export class CloudTaskService extends CloudTaskEngine { + constructor( + @inject(CLOUD_TASK_AUTH) + auth: ICloudTaskAuth, + @inject(ANALYTICS_SERVICE) + analytics: IAnalytics, + @inject(ROOT_LOGGER) + logger: RootLogger, + @inject(MCP_RELAY_EXECUTOR) + @optional() + mcpRelayExecutor: McpRelayExecutor | null = null, + ) { + super({ auth, analytics, logger, mcpRelayExecutor }); + } + + @preDestroy() + override unwatchAll(): void { + super.unwatchAll(); + } +} diff --git a/packages/core/src/cloud-task/schemas.ts b/packages/core/src/cloud-task/schemas.ts index d694e52141..b8c03eb202 100644 --- a/packages/core/src/cloud-task/schemas.ts +++ b/packages/core/src/cloud-task/schemas.ts @@ -1,20 +1,12 @@ -import type { TaskRunStatus } from "@posthog/shared"; +import type { CloudTaskUpdatePayload } from "@posthog/shared"; import { z } from "zod"; -import type { CloudTaskUpdatePayload } from "./cloud-task-types"; -export type { CloudTaskUpdatePayload, TaskRunStatus }; - -export const TERMINAL_STATUSES = ["completed", "failed", "cancelled"] as const; - -export function isTerminalStatus( - status: TaskRunStatus | string | null | undefined, -): boolean { - return ( - status !== null && - status !== undefined && - TERMINAL_STATUSES.includes(status as (typeof TERMINAL_STATUSES)[number]) - ); -} +export { + type CloudTaskUpdatePayload, + isTerminalStatus, + type TaskRunStatus, + TERMINAL_STATUSES, +} from "@posthog/shared"; // --- Events --- diff --git a/packages/core/src/handoff/handoff.ts b/packages/core/src/handoff/handoff.ts index 2b8d477099..c85c2fcfa8 100644 --- a/packages/core/src/handoff/handoff.ts +++ b/packages/core/src/handoff/handoff.ts @@ -5,7 +5,7 @@ import { TypedEventEmitter, } from "@posthog/shared"; import { inject, injectable } from "inversify"; -import type { CloudTaskService } from "../cloud-task/cloud-task-engine"; +import type { CloudTaskService } from "../cloud-task/cloud-task"; import { CLOUD_TASK_SERVICE } from "../cloud-task/identifiers"; import { HandoffSaga, type HandoffSagaDeps } from "./handoff-saga"; import { diff --git a/packages/core/vitest.config.ts b/packages/core/vitest.config.ts new file mode 100644 index 0000000000..bed14b24e0 --- /dev/null +++ b/packages/core/vitest.config.ts @@ -0,0 +1,23 @@ +import { defineConfig } from "vitest/config"; +import { trunkTestOptions } from "../../vitest.config.base"; + +export default defineConfig({ + oxc: false, + esbuild: { + tsconfigRaw: { + compilerOptions: { + experimentalDecorators: true, + target: "ES2022", + useDefineForClassFields: false, + verbatimModuleSyntax: true, + }, + }, + }, + test: { + globals: true, + ...trunkTestOptions, + environment: "node", + include: ["src/**/*.test.ts", "src/**/*.test.tsx"], + exclude: ["**/node_modules/**", "**/dist/**"], + }, +}); diff --git a/packages/host-router/src/routers/cloud-task.router.ts b/packages/host-router/src/routers/cloud-task.router.ts index 4546995404..15d577ce59 100644 --- a/packages/host-router/src/routers/cloud-task.router.ts +++ b/packages/host-router/src/routers/cloud-task.router.ts @@ -1,4 +1,4 @@ -import type { CloudTaskService } from "@posthog/core/cloud-task/cloud-task-engine"; +import type { CloudTaskService } from "@posthog/core/cloud-task/cloud-task"; import { CLOUD_TASK_SERVICE } from "@posthog/core/cloud-task/identifiers"; import { CloudTaskEvent, From 36f9b398e6afb14a921e49148574f39f104a66c5 Mon Sep 17 00:00:00 2001 From: Richard Solomou Date: Fri, 24 Jul 2026 19:21:03 +0300 Subject: [PATCH 37/45] refactor(core): extract repository integration semantics Generated-By: PostHog Code Task-Id: c1bbe3cf-742b-4b24-bf96-d11a18b4cf22 --- .../src/integrations/repositories.test.ts | 66 ++++++++++ .../core/src/integrations/repositories.ts | 124 ++++++++++++++++++ 2 files changed, 190 insertions(+) diff --git a/packages/core/src/integrations/repositories.test.ts b/packages/core/src/integrations/repositories.test.ts index ce6af19540..80eec8ad36 100644 --- a/packages/core/src/integrations/repositories.test.ts +++ b/packages/core/src/integrations/repositories.test.ts @@ -1,5 +1,7 @@ import { describe, expect, it } from "vitest"; import { + buildTeamRepositoryOptions, + buildUserRepositoryOptions, combineGithubRepositories, combineRepositoryPicker, combineUserGithubRepositories, @@ -7,8 +9,11 @@ import { isEmptyRepositoryMap, isRepoInIntegration, normalizeRepoKey, + normalizeRepositoryNames, type RepositoryCacheAction, type RepositoryQueryResult, + repositoryLoadWarning, + repositoryOptionsEqual, resolveEffectiveUserRepositoryMap, resolveUserRepositoryCacheAction, sameUserRepositoryMap, @@ -18,6 +23,67 @@ import { type UserRepositoryIntegrationRef, } from "./repositories"; +describe("repository options", () => { + it("normalizes repository names", () => { + expect(normalizeRepositoryNames(["PostHog/Code", ""])).toEqual([ + "posthog/code", + ]); + }); + + it("builds sorted team options with integration labels", () => { + expect( + buildTeamRepositoryOptions( + [ + { id: 2, display_name: "Work" }, + { id: 1, config: { account: { login: "personal" } } }, + ], + { 1: ["z/repo"], 2: ["a/repo"] }, + ), + ).toEqual([ + { integrationId: 2, integrationLabel: "Work", repository: "a/repo" }, + { + integrationId: 1, + integrationLabel: "personal", + repository: "z/repo", + }, + ]); + }); + + it("builds user options with the same shape", () => { + expect( + buildUserRepositoryOptions( + [{ id: "user-1", installation_id: "42", account: { name: "Me" } }], + { 42: ["posthog/code"] }, + ), + ).toEqual([ + { integrationId: 42, integrationLabel: "Me", repository: "posthog/code" }, + ]); + }); + + it.each([ + [0, 2, null], + [1, 2, "Some GitHub repositories could not be loaded. Pull to retry."], + [2, 2, "Could not load GitHub repositories. Pull to retry."], + ])( + "describes %i of %i failed repository loads", + (failed, total, expected) => { + expect(repositoryLoadWarning(failed, total)).toBe(expected); + }, + ); + + it("compares option lists by content", () => { + const options = [ + { integrationId: 1, integrationLabel: "Me", repository: "a/repo" }, + ]; + expect( + repositoryOptionsEqual( + options, + options.map((option) => ({ ...option })), + ), + ).toBe(true); + }); +}); + function result( data: T | undefined, flags: Partial, "data">> = {}, diff --git a/packages/core/src/integrations/repositories.ts b/packages/core/src/integrations/repositories.ts index 43a7fbf74e..ea9bb6d202 100644 --- a/packages/core/src/integrations/repositories.ts +++ b/packages/core/src/integrations/repositories.ts @@ -5,6 +5,130 @@ export interface RepositoryQueryResult { isRefetching: boolean; } +export interface RepositoryOption { + integrationId: number; + integrationLabel: string; + repository: string; +} + +export interface RepositorySelection { + integrationId: number | null; + repository: string | null; +} + +export interface TeamRepositoryIntegration { + id: number; + display_name?: string; + config?: { account?: { login?: string } }; +} + +export interface UserRepositoryIntegration { + id: string; + installation_id: string; + account?: { name?: string | null } | null; +} + +export function normalizeRepositoryNames( + repositories: ReadonlyArray, +): string[] { + return repositories + .map((repository) => repository.toLowerCase()) + .filter((repository) => repository.length > 0); +} + +export function repositoryLoadWarning( + failedCount: number, + totalCount: number, +): string | null { + if (failedCount === 0) return null; + return failedCount === totalCount + ? "Could not load GitHub repositories. Pull to retry." + : "Some GitHub repositories could not be loaded. Pull to retry."; +} + +export function buildTeamRepositoryOptions( + integrations: ReadonlyArray, + repositoriesByIntegration: Readonly>, +): RepositoryOption[] { + return integrations + .flatMap((integration) => + (repositoriesByIntegration[integration.id] ?? []).map((repository) => ({ + integrationId: integration.id, + integrationLabel: + integration.display_name ?? + integration.config?.account?.login ?? + `GitHub ${integration.id}`, + repository, + })), + ) + .sort((left, right) => left.repository.localeCompare(right.repository)); +} + +export function buildUserRepositoryOptions( + integrations: ReadonlyArray, + repositoriesByInstallation: Readonly>, +): RepositoryOption[] { + return integrations + .flatMap((integration) => + (repositoriesByInstallation[integration.installation_id] ?? []).map( + (repository) => ({ + integrationId: Number(integration.installation_id), + integrationLabel: + integration.account?.name ?? + `GitHub ${integration.installation_id}`, + repository, + }), + ), + ) + .sort((left, right) => left.repository.localeCompare(right.repository)); +} + +export function repositoryOptionsEqual( + left: ReadonlyArray, + right: ReadonlyArray, +): boolean { + return ( + left.length === right.length && + left.every((option, index) => { + const other = right[index]; + return ( + other?.integrationId === option.integrationId && + other.integrationLabel === option.integrationLabel && + other.repository === option.repository + ); + }) + ); +} + +export function findRepositoryOption( + options: ReadonlyArray, + selection: RepositorySelection, +): RepositoryOption | null { + if (!selection.integrationId || !selection.repository) return null; + return ( + options.find( + (option) => + option.integrationId === selection.integrationId && + option.repository === selection.repository, + ) ?? null + ); +} + +export function toRepositorySelection( + option: RepositoryOption | null, +): RepositorySelection { + return { + integrationId: option?.integrationId ?? null, + repository: option?.repository ?? null, + }; +} + +export function isRepositorySelectionComplete( + selection: RepositorySelection, +): boolean { + return !!selection.integrationId && !!selection.repository; +} + export interface TeamRepositoriesResult { integrationId: number; repos?: string[] | null; From 3003336e7f5c74b0181f0f79a4ec92524c372071 Mon Sep 17 00:00:00 2001 From: Richard Solomou Date: Fri, 24 Jul 2026 19:21:33 +0300 Subject: [PATCH 38/45] refactor(core): extract pending prompt recovery Generated-By: PostHog Code Task-Id: c1bbe3cf-742b-4b24-bf96-d11a18b4cf22 --- .../core/src/tasks/pendingPrompts.test.ts | 37 +++++++++++++++ packages/core/src/tasks/pendingPrompts.ts | 47 +++++++++++++++++++ .../task-detail/hooks/useTaskCreation.ts | 7 ++- .../ui/src/shell/pendingTaskPromptStore.ts | 45 +++++++----------- 4 files changed, 105 insertions(+), 31 deletions(-) create mode 100644 packages/core/src/tasks/pendingPrompts.test.ts create mode 100644 packages/core/src/tasks/pendingPrompts.ts diff --git a/packages/core/src/tasks/pendingPrompts.test.ts b/packages/core/src/tasks/pendingPrompts.test.ts new file mode 100644 index 0000000000..8a733ad113 --- /dev/null +++ b/packages/core/src/tasks/pendingPrompts.test.ts @@ -0,0 +1,37 @@ +import { describe, expect, it } from "vitest"; +import { + buildPendingPromptKey, + capPendingPrompts, + listPendingPromptsNewestFirst, + selectNewestPendingPrompt, +} from "./pendingPrompts"; + +describe("pending prompts", () => { + it("keeps the newest prompts up to the limit", () => { + expect( + capPendingPrompts( + { + old: { createdAt: 1 }, + middle: { createdAt: 2 }, + newest: { createdAt: 3 }, + }, + 2, + ), + ).toEqual({ middle: { createdAt: 2 }, newest: { createdAt: 3 } }); + }); + + it("orders prompts newest first and selects the newest", () => { + const prompts = { old: { createdAt: 1 }, new: { createdAt: 2 } }; + expect( + listPendingPromptsNewestFirst(prompts).map(({ key }) => key), + ).toEqual(["new", "old"]); + expect(selectNewestPendingPrompt(prompts)?.key).toBe("new"); + }); + + it.each([ + ["uuid", 1, "abc", "uuid"], + [null, 123, "abc", "pending-123-abc"], + ])("builds a portable pending key", (uuid, timestamp, entropy, expected) => { + expect(buildPendingPromptKey(uuid, timestamp, entropy)).toBe(expected); + }); +}); diff --git a/packages/core/src/tasks/pendingPrompts.ts b/packages/core/src/tasks/pendingPrompts.ts new file mode 100644 index 0000000000..e4f77eafa2 --- /dev/null +++ b/packages/core/src/tasks/pendingPrompts.ts @@ -0,0 +1,47 @@ +export const MAX_RECOVERABLE_PROMPTS = 20; + +export interface TimestampedPendingPrompt { + createdAt: number; +} + +export interface RecoverablePendingPrompt< + TPrompt extends TimestampedPendingPrompt, +> { + key: string; + prompt: TPrompt; +} + +export function capPendingPrompts( + byKey: Record, + limit: number = MAX_RECOVERABLE_PROMPTS, +): Record { + const keys = Object.keys(byKey); + if (keys.length <= limit) return byKey; + + const kept = keys + .sort((left, right) => byKey[right].createdAt - byKey[left].createdAt) + .slice(0, limit); + return Object.fromEntries(kept.map((key) => [key, byKey[key]])); +} + +export function listPendingPromptsNewestFirst< + TPrompt extends TimestampedPendingPrompt, +>(byKey: Record): RecoverablePendingPrompt[] { + return Object.entries(byKey) + .map(([key, prompt]) => ({ key, prompt })) + .sort((left, right) => right.prompt.createdAt - left.prompt.createdAt); +} + +export function selectNewestPendingPrompt< + TPrompt extends TimestampedPendingPrompt, +>(byKey: Record): RecoverablePendingPrompt | null { + return listPendingPromptsNewestFirst(byKey)[0] ?? null; +} + +export function buildPendingPromptKey( + randomUuid: string | null, + timestamp: number, + entropy: string, +): string { + return randomUuid ?? `pending-${timestamp}-${entropy}`; +} diff --git a/packages/ui/src/features/task-detail/hooks/useTaskCreation.ts b/packages/ui/src/features/task-detail/hooks/useTaskCreation.ts index d13a1c445f..9b3abdaae0 100644 --- a/packages/ui/src/features/task-detail/hooks/useTaskCreation.ts +++ b/packages/ui/src/features/task-detail/hooks/useTaskCreation.ts @@ -31,7 +31,10 @@ import { useConnectivity } from "../../../hooks/useConnectivity"; import { toast } from "../../../primitives/toast"; import { track } from "../../../shell/analytics"; import { logger } from "../../../shell/logger"; -import { pendingTaskPromptStoreApi } from "../../../shell/pendingTaskPromptStore"; +import { + generatePendingTaskKey, + pendingTaskPromptStoreApi, +} from "../../../shell/pendingTaskPromptStore"; import { titleAttachmentStoreApi } from "../../../shell/titleAttachmentStore"; import { useAuthStateValue } from "../../auth/store"; import { assertCloudUsageAvailable } from "../../billing/preflightCloudUsage"; @@ -319,7 +322,7 @@ export function useTaskCreation({ const shouldShowPendingView = !onTaskCreated && !!plainPromptText; const pendingTaskKey = shouldShowPendingView - ? (globalThis.crypto?.randomUUID?.() ?? `pending-${Date.now()}`) + ? generatePendingTaskKey() : null; if (pendingTaskKey) { diff --git a/packages/ui/src/shell/pendingTaskPromptStore.ts b/packages/ui/src/shell/pendingTaskPromptStore.ts index b91fefd3f4..cd227ef24a 100644 --- a/packages/ui/src/shell/pendingTaskPromptStore.ts +++ b/packages/ui/src/shell/pendingTaskPromptStore.ts @@ -1,13 +1,8 @@ import type { UserMessageAttachment } from "@posthog/ui/features/sessions/userMessageTypes"; -import { logger } from "@posthog/ui/shell/logger"; import { electronStorage } from "@posthog/ui/shell/rendererStorage"; import { create } from "zustand"; import { persist } from "zustand/middleware"; -const log = logger.scope("pending-task-prompts"); - -const MAX_PENDING_PROMPTS = 20; - export interface PendingTaskPrompt { promptText: string; attachments: UserMessageAttachment[]; @@ -16,26 +11,6 @@ export interface PendingTaskPrompt { export type PendingTaskPromptInput = Omit; -function capToNewest( - byKey: Record, -): Record { - const keys = Object.keys(byKey); - if (keys.length <= MAX_PENDING_PROMPTS) { - return byKey; - } - const keptKeys = keys - .sort((a, b) => byKey[b].createdAt - byKey[a].createdAt) - .slice(0, MAX_PENDING_PROMPTS); - log.warn("Dropping oldest unrecovered prompts beyond cap", { - dropped: keys.length - keptKeys.length, - }); - const kept: Record = {}; - for (const key of keptKeys) { - kept[key] = byKey[key]; - } - return kept; -} - interface PendingTaskPromptStore { byKey: Record; _hasHydrated: boolean; @@ -54,7 +29,7 @@ export const usePendingTaskPromptStore = create()( setHasHydrated: (hydrated) => set({ _hasHydrated: hydrated }), set: (key, prompt) => set((state) => ({ - byKey: capToNewest({ + byKey: capPendingPrompts({ ...state.byKey, [key]: { ...prompt, createdAt: Date.now() }, }), @@ -110,9 +85,7 @@ export const pendingTaskPromptStoreApi = { usePendingTaskPromptStore.getState().move(fromKey, toKey), clear: (key: string) => usePendingTaskPromptStore.getState().clear(key), getAllNewestFirst: (): RecoverablePendingPrompt[] => - Object.entries(usePendingTaskPromptStore.getState().byKey) - .map(([key, prompt]) => ({ key, prompt })) - .sort((a, b) => b.prompt.createdAt - a.prompt.createdAt), + listPendingPromptsNewestFirst(usePendingTaskPromptStore.getState().byKey), whenHydrated: (): Promise => { if (usePendingTaskPromptStore.getState()._hasHydrated) { return Promise.resolve(); @@ -128,6 +101,14 @@ export const pendingTaskPromptStoreApi = { }, }; +export function generatePendingTaskKey(): string { + return buildPendingPromptKey( + globalThis.crypto?.randomUUID?.() ?? null, + Date.now(), + Math.random().toString(36).slice(2, 10), + ); +} + export function usePendingTaskPrompt( key: string | undefined, ): PendingTaskPrompt | undefined { @@ -135,3 +116,9 @@ export function usePendingTaskPrompt( key ? state.byKey[key] : undefined, ); } + +import { + buildPendingPromptKey, + capPendingPrompts, + listPendingPromptsNewestFirst, +} from "@posthog/core/tasks/pendingPrompts"; From b5204c8a1322c09f3b1f9a6432529a65479d787f Mon Sep 17 00:00:00 2001 From: Richard Solomou Date: Fri, 24 Jul 2026 19:22:05 +0300 Subject: [PATCH 39/45] refactor(core): extract plan approval presentation Generated-By: PostHog Code Task-Id: c1bbe3cf-742b-4b24-bf96-d11a18b4cf22 --- .../sessions/planApprovalPresentation.test.ts | 29 +++++++++++++++++++ .../src/sessions/planApprovalPresentation.ts | 23 +++++++++++++++ .../session-update/PlanApprovalView.tsx | 19 ++++-------- 3 files changed, 57 insertions(+), 14 deletions(-) create mode 100644 packages/core/src/sessions/planApprovalPresentation.test.ts create mode 100644 packages/core/src/sessions/planApprovalPresentation.ts diff --git a/packages/core/src/sessions/planApprovalPresentation.test.ts b/packages/core/src/sessions/planApprovalPresentation.test.ts new file mode 100644 index 0000000000..27f9636189 --- /dev/null +++ b/packages/core/src/sessions/planApprovalPresentation.test.ts @@ -0,0 +1,29 @@ +import { describe, expect, it } from "vitest"; +import { extractPlanText } from "./planApprovalPresentation"; + +describe("extractPlanText", () => { + it.each([ + [{ rawInput: { plan: "Raw plan" } }, "Raw plan"], + [{ content: [{ text: "Direct content" }] }, "Direct content"], + [ + { + content: [ + { type: "content", content: { type: "text", text: "Nested" } }, + ], + }, + "Nested", + ], + [{ rawInput: {}, content: [] }, null], + ])("extracts plan presentation from %o", (toolCall, expected) => { + expect(extractPlanText(toolCall)).toBe(expected); + }); + + it("prefers the canonical raw plan over rendered content", () => { + expect( + extractPlanText({ + rawInput: { plan: "Canonical" }, + content: [{ text: "Rendered" }], + }), + ).toBe("Canonical"); + }); +}); diff --git a/packages/core/src/sessions/planApprovalPresentation.ts b/packages/core/src/sessions/planApprovalPresentation.ts new file mode 100644 index 0000000000..2dd43553b6 --- /dev/null +++ b/packages/core/src/sessions/planApprovalPresentation.ts @@ -0,0 +1,23 @@ +function extractTextContent(item: unknown): string | null { + if (!item || typeof item !== "object") return null; + const record = item as Record; + if (typeof record.text === "string") return record.text; + + if (!record.content || typeof record.content !== "object") return null; + const content = record.content as Record; + return typeof content.text === "string" ? content.text : null; +} + +export function extractPlanText(toolCall: { + rawInput?: { plan?: unknown } | null; + content?: readonly unknown[] | null; +}): string | null { + const rawPlan = toolCall.rawInput?.plan; + if (typeof rawPlan === "string" && rawPlan.trim()) return rawPlan; + + for (const item of toolCall.content ?? []) { + const text = extractTextContent(item); + if (text?.trim()) return text; + } + return null; +} diff --git a/packages/ui/src/features/sessions/components/session-update/PlanApprovalView.tsx b/packages/ui/src/features/sessions/components/session-update/PlanApprovalView.tsx index b70c281953..bf609d89d2 100644 --- a/packages/ui/src/features/sessions/components/session-update/PlanApprovalView.tsx +++ b/packages/ui/src/features/sessions/components/session-update/PlanApprovalView.tsx @@ -1,4 +1,5 @@ import { CaretDown, CaretRight, CheckCircle } from "@phosphor-icons/react"; +import { extractPlanText } from "@posthog/core/sessions/planApprovalPresentation"; import { Box, Flex, Text } from "@radix-ui/themes"; import { useMemo, useState } from "react"; import { PlanContent } from "../../../permissions/PlanContent"; @@ -32,20 +33,10 @@ export function PlanApprovalView({ | undefined; const isHistoricalPlan = rawInput?.historical === true; - const planText = useMemo(() => { - if (content?.length) { - const textContent = content.find((c) => c.type === "content"); - if (textContent && "content" in textContent) { - const inner = textContent.content as - | { type?: string; text?: string } - | undefined; - if (inner?.type === "text" && inner.text) { - return inner.text; - } - } - } - return rawInput?.plan ?? null; - }, [content, rawInput?.plan]); + const planText = useMemo( + () => extractPlanText({ rawInput, content }), + [content, rawInput], + ); const wasNotApproved = isFailed || wasCancelled; const showResult = isHistoricalPlan || isComplete || wasNotApproved; From 7ccd2592d6cc8f4117fe8287442ee751aeeade10 Mon Sep 17 00:00:00 2001 From: Richard Solomou Date: Fri, 24 Jul 2026 19:22:25 +0300 Subject: [PATCH 40/45] refactor(core): extract permission option presentation Generated-By: PostHog Code Task-Id: c1bbe3cf-742b-4b24-bf96-d11a18b4cf22 --- .../src/sessions/permissionResponse.test.ts | 52 ++++++++++++++ .../core/src/sessions/permissionResponse.ts | 67 +++++++++++++++++++ .../permissions/PlanApprovalSelector.tsx | 42 +++--------- packages/ui/src/features/permissions/types.ts | 9 ++- 4 files changed, 131 insertions(+), 39 deletions(-) diff --git a/packages/core/src/sessions/permissionResponse.test.ts b/packages/core/src/sessions/permissionResponse.test.ts index 228a28142d..d467ca6154 100644 --- a/packages/core/src/sessions/permissionResponse.test.ts +++ b/packages/core/src/sessions/permissionResponse.test.ts @@ -2,10 +2,62 @@ import type { PermissionRequest } from "@posthog/shared"; import { describe, expect, it } from "vitest"; import { formatPermissionAnswerPrompt, + getPermissionOptionMeta, isOtherPermissionOption, + isPermissionApproval, + isPermissionRejection, + permissionOptionUsesCustomInput, planPermissionResponse, + resolveInitialPlanApprovalOption, + selectPlanPermissionOptions, } from "./permissionResponse"; +describe("permission option presentation", () => { + const approveOnce = { + optionId: "default", + name: "Approve", + kind: "allow_once" as const, + }; + const approveAuto = { + optionId: "auto", + name: "Approve automatically", + kind: "allow_always" as const, + }; + const reject = { + optionId: "reject_with_feedback", + name: "Reject", + kind: "reject_once" as const, + _meta: { customInput: true, description: "Explain why" }, + }; + + it("classifies approval, rejection, and custom-input options", () => { + expect(isPermissionApproval(approveOnce)).toBe(true); + expect(isPermissionRejection(reject)).toBe(true); + expect(permissionOptionUsesCustomInput(reject)).toBe(true); + expect(getPermissionOptionMeta(reject)).toEqual({ + customInput: true, + description: "Explain why", + }); + }); + + it("selects plan options and prefers a feedback rejection", () => { + expect(selectPlanPermissionOptions([approveOnce, reject])).toEqual({ + approvals: [approveOnce], + rejection: reject, + }); + }); + + it.each([ + ["default", "default"], + [null, "auto"], + ["missing", "auto"], + ])("resolves preferred approval %s", (preferred, expected) => { + expect( + resolveInitialPlanApprovalOption([approveOnce, approveAuto], preferred), + ).toBe(expected); + }); +}); + function makePermission( options: Array<{ optionId: string; diff --git a/packages/core/src/sessions/permissionResponse.ts b/packages/core/src/sessions/permissionResponse.ts index eef81a7f10..5cc4a596b5 100644 --- a/packages/core/src/sessions/permissionResponse.ts +++ b/packages/core/src/sessions/permissionResponse.ts @@ -1,5 +1,72 @@ import type { PermissionRequest } from "@posthog/shared"; +export type PermissionOption = PermissionRequest["options"][number]; + +export function getPermissionOptionMeta(option: PermissionOption): { + customInput: boolean; + description?: string; +} { + const meta = option._meta as + | { customInput?: boolean; description?: string } + | null + | undefined; + return { + customInput: meta?.customInput === true, + ...(meta?.description ? { description: meta.description } : {}), + }; +} + +export function isPermissionApproval(option: PermissionOption): boolean { + return option.kind === "allow_once" || option.kind === "allow_always"; +} + +export function isPermissionRejection(option: PermissionOption): boolean { + return ( + option.kind === "reject_once" || + option.kind === "reject_always" || + option.optionId.includes("reject") + ); +} + +export function permissionOptionUsesCustomInput( + option: PermissionOption, +): boolean { + return ( + isOtherPermissionOption(option.optionId) || + getPermissionOptionMeta(option).customInput + ); +} + +export function selectPlanPermissionOptions(options: PermissionOption[]): { + approvals: PermissionOption[]; + rejection: PermissionOption | null; +} { + const approvals = options.filter(isPermissionApproval); + const rejections = options.filter(isPermissionRejection); + return { + approvals, + rejection: + rejections.find(permissionOptionUsesCustomInput) ?? rejections[0] ?? null, + }; +} + +export function resolveInitialPlanApprovalOption( + approvals: PermissionOption[], + preferredOptionId?: string | null, +): string | undefined { + const has = (optionId: string): boolean => + approvals.some((option) => option.optionId === optionId); + return ( + (preferredOptionId && has(preferredOptionId) + ? preferredOptionId + : undefined) ?? + (has("auto") ? "auto" : undefined) ?? + approvals.find((option) => option.optionId === "default")?.optionId ?? + approvals.find((option) => option.kind === "allow_once")?.optionId ?? + approvals[0]?.optionId + ); +} + const OTHER_OPTION_ID = "_other"; const OTHER_OPTION_ID_ALT = "other"; diff --git a/packages/ui/src/features/permissions/PlanApprovalSelector.tsx b/packages/ui/src/features/permissions/PlanApprovalSelector.tsx index b22a77951f..145add84aa 100644 --- a/packages/ui/src/features/permissions/PlanApprovalSelector.tsx +++ b/packages/ui/src/features/permissions/PlanApprovalSelector.tsx @@ -1,7 +1,8 @@ -import type { - PermissionOption, - SessionConfigOption, -} from "@agentclientprotocol/sdk"; +import type { SessionConfigOption } from "@agentclientprotocol/sdk"; +import { + resolveInitialPlanApprovalOption, + selectPlanPermissionOptions, +} from "@posthog/core/sessions/permissionResponse"; import type { ExecutionMode } from "@posthog/shared"; import { ModeSelector } from "@posthog/ui/features/message-editor/components/ModeSelector"; import { MODE_LABELS } from "@posthog/ui/features/sessions/modeStyles"; @@ -17,21 +18,6 @@ import { type BasePermissionProps, toSelectorOptions } from "./types"; const TITLE = "Implementation Plan"; const QUESTION = "Approve this plan to proceed?"; -function isApprove(option: PermissionOption): boolean { - return option.kind === "allow_once" || option.kind === "allow_always"; -} - -function isReject(option: PermissionOption): boolean { - return option.kind === "reject_once" || option.kind === "reject_always"; -} - -function hasCustomInput(option: PermissionOption): boolean { - return ( - (option._meta as { customInput?: boolean } | null | undefined) - ?.customInput === true - ); -} - // Don't steal focus from an interactive element in a different grid cell // (multi-task view). Mirrors the guard in useActionSelectorState. function isInteractiveElementInDifferentCell( @@ -66,11 +52,8 @@ export function PlanApprovalSelector({ onSelect, onCancel, }: BasePermissionProps) { - const approveOptions = useMemo(() => options.filter(isApprove), [options]); - const rejectOption = useMemo( - () => - options.find((o) => isReject(o) && hasCustomInput(o)) ?? - options.find(isReject), + const { approvals: approveOptions, rejection: rejectOption } = useMemo( + () => selectPlanPermissionOptions(options), [options], ); @@ -87,16 +70,7 @@ export function PlanApprovalSelector({ // via `useMemo` (rather than seeding a `useState` once) means it stays // correct once the store finishes hydrating. const initialMode = useMemo(() => { - const has = (id: string) => approveOptions.some((o) => o.optionId === id); - return ( - (lastApprovalMode && has(lastApprovalMode) - ? lastApprovalMode - : undefined) ?? - (has("auto") ? "auto" : undefined) ?? - approveOptions.find((o) => o.optionId === "default")?.optionId ?? - approveOptions.find((o) => o.kind === "allow_once")?.optionId ?? - approveOptions[0]?.optionId - ); + return resolveInitialPlanApprovalOption(approveOptions, lastApprovalMode); }, [approveOptions, lastApprovalMode]); // Only the user's own pick lives in state; everything else derives from diff --git a/packages/ui/src/features/permissions/types.ts b/packages/ui/src/features/permissions/types.ts index 1505da5062..1794965cda 100644 --- a/packages/ui/src/features/permissions/types.ts +++ b/packages/ui/src/features/permissions/types.ts @@ -3,6 +3,7 @@ import type { RequestPermissionRequest, ToolCallContent, } from "@agentclientprotocol/sdk"; +import { getPermissionOptionMeta } from "@posthog/core/sessions/permissionResponse"; import type { CodeToolKind } from "@posthog/ui/features/sessions/types"; import type { SelectorOption } from "@posthog/ui/primitives/ActionSelector"; @@ -26,14 +27,12 @@ export function toSelectorOptions( options: PermissionOption[], ): SelectorOption[] { return options.map((opt) => { - const meta = opt._meta as - | { description?: string; customInput?: boolean } - | undefined; + const meta = getPermissionOptionMeta(opt); return { id: opt.optionId, label: opt.name, - description: meta?.description, - customInput: meta?.customInput, + description: meta.description, + customInput: meta.customInput, }; }); } From 046422e63a97b1c2125b878cc17b7aa734cb3edd Mon Sep 17 00:00:00 2001 From: Richard Solomou Date: Fri, 24 Jul 2026 19:22:47 +0300 Subject: [PATCH 41/45] refactor(core): extract composer controls Generated-By: PostHog Code Task-Id: c1bbe3cf-742b-4b24-bf96-d11a18b4cf22 --- .../src/task-detail/composerControls.test.ts | 27 +++++ .../core/src/task-detail/composerControls.ts | 99 +++++++++++++++++++ .../components/UnifiedModelSelector.tsx | 4 +- 3 files changed, 128 insertions(+), 2 deletions(-) create mode 100644 packages/core/src/task-detail/composerControls.test.ts create mode 100644 packages/core/src/task-detail/composerControls.ts diff --git a/packages/core/src/task-detail/composerControls.test.ts b/packages/core/src/task-detail/composerControls.test.ts new file mode 100644 index 0000000000..0ba02d02cc --- /dev/null +++ b/packages/core/src/task-detail/composerControls.test.ts @@ -0,0 +1,27 @@ +import { describe, expect, it } from "vitest"; +import { resolveComposerPrimaryAction } from "./composerControls"; + +describe("resolveComposerPrimaryAction", () => { + it.each([ + [{ hasContent: true }, "send"], + [{ canStop: true }, "stop"], + [{ canStop: true, hasContent: true }, "send"], + [{ canStop: true, hasContent: true, allowSendWhileRunning: false }, "stop"], + [{ isRecording: true }, "mic-stop"], + [{}, "mic"], + [{ disabled: true, hasContent: true }, "disabled"], + [{ isTranscribing: true }, "disabled"], + ])("derives %s", (overrides, expected) => { + expect( + resolveComposerPrimaryAction({ + hasContent: false, + disabled: false, + isRecording: false, + isTranscribing: false, + canStop: false, + allowSendWhileRunning: true, + ...overrides, + }), + ).toBe(expected); + }); +}); diff --git a/packages/core/src/task-detail/composerControls.ts b/packages/core/src/task-detail/composerControls.ts new file mode 100644 index 0000000000..e9694cccef --- /dev/null +++ b/packages/core/src/task-detail/composerControls.ts @@ -0,0 +1,99 @@ +import { + type Adapter, + type CloudTaskConfigOption, + DEFAULT_REASONING_EFFORT, + isRestrictedModelOption, + isSupportedReasoningEffort, + type SupportedReasoningEffort, +} from "@posthog/shared"; + +export interface ComposerModelOption { + value: string; + label: string; + description?: string; + disabled: boolean; +} + +export function getModelConfigOption( + configOptions: readonly CloudTaskConfigOption[], +): CloudTaskConfigOption { + const option = configOptions.find((item) => item.category === "model"); + if (!option) throw new Error("Cloud task model configuration is unavailable"); + return option; +} + +export function getComposerModelOptions( + modelOption: CloudTaskConfigOption, +): ComposerModelOption[] { + return modelOption.options.map((option) => ({ + value: option.value, + label: option.name, + description: option.description, + disabled: isRestrictedModelOption(option._meta), + })); +} + +export function getConfigOptionLabel( + options: ReadonlyArray<{ value: string; name: string }>, + value: string | undefined, +): string | undefined { + return options.find((option) => option.value === value)?.name ?? value; +} + +export function resolveAvailableModel( + modelOption: CloudTaskConfigOption, + value: string, +): string { + const selected = modelOption.options.find((option) => option.value === value); + return selected && !isRestrictedModelOption(selected._meta) + ? value + : modelOption.currentValue; +} + +export function resolveComposerModelChange({ + adapter, + modelOption, + requestedModel, + reasoning, +}: { + adapter: Adapter; + modelOption: CloudTaskConfigOption; + requestedModel: string; + reasoning: SupportedReasoningEffort; +}): { model: string; reasoning: SupportedReasoningEffort } { + const model = resolveAvailableModel(modelOption, requestedModel); + return { + model, + reasoning: isSupportedReasoningEffort(adapter, model, reasoning) + ? reasoning + : DEFAULT_REASONING_EFFORT, + }; +} + +export type ComposerPrimaryAction = + | "send" + | "stop" + | "mic" + | "mic-stop" + | "disabled"; + +export function resolveComposerPrimaryAction({ + hasContent, + disabled, + isRecording, + isTranscribing, + canStop, + allowSendWhileRunning, +}: { + hasContent: boolean; + disabled: boolean; + isRecording: boolean; + isTranscribing: boolean; + canStop: boolean; + allowSendWhileRunning: boolean; +}): ComposerPrimaryAction { + if (disabled || isTranscribing) return "disabled"; + if (canStop && (!allowSendWhileRunning || !hasContent)) return "stop"; + if (hasContent && !isRecording) return "send"; + return isRecording ? "mic-stop" : "mic"; +} diff --git a/packages/ui/src/features/sessions/components/UnifiedModelSelector.tsx b/packages/ui/src/features/sessions/components/UnifiedModelSelector.tsx index 06b3956f41..d34dbd9135 100644 --- a/packages/ui/src/features/sessions/components/UnifiedModelSelector.tsx +++ b/packages/ui/src/features/sessions/components/UnifiedModelSelector.tsx @@ -9,6 +9,7 @@ import { Robot, Spinner, } from "@phosphor-icons/react"; +import { getConfigOptionLabel } from "@posthog/core/task-detail/composerControls"; import { Button, DropdownMenu, @@ -78,8 +79,7 @@ export function UnifiedModelSelector({ }, [selectOption]); const currentValue = selectOption?.currentValue; - const currentLabel = - options.find((opt) => opt.value === currentValue)?.name ?? currentValue; + const currentLabel = getConfigOptionLabel(options, currentValue); const otherAdapter = getOtherAdapter(adapter); From 37f1c0ba9d4125f2ff41da71e43e60a2c8c71bf3 Mon Sep 17 00:00:00 2001 From: Richard Solomou Date: Fri, 24 Jul 2026 19:23:30 +0300 Subject: [PATCH 42/45] refactor(core): extract composer model policy Generated-By: PostHog Code Task-Id: c1bbe3cf-742b-4b24-bf96-d11a18b4cf22 --- .../task-detail/composerModelPolicy.test.ts | 42 +++++++++++++++++++ .../src/task-detail/composerModelPolicy.ts | 35 ++++++++++++++++ 2 files changed, 77 insertions(+) create mode 100644 packages/core/src/task-detail/composerModelPolicy.test.ts create mode 100644 packages/core/src/task-detail/composerModelPolicy.ts diff --git a/packages/core/src/task-detail/composerModelPolicy.test.ts b/packages/core/src/task-detail/composerModelPolicy.test.ts new file mode 100644 index 0000000000..235c6199f6 --- /dev/null +++ b/packages/core/src/task-detail/composerModelPolicy.test.ts @@ -0,0 +1,42 @@ +import { + type Adapter, + type CloudTaskConfigOption, + DEFAULT_GATEWAY_MODEL, + restrictedModelMeta, + type SupportedReasoningEffort, +} from "@posthog/shared"; +import { expect, it } from "vitest"; +import { resolveCloudComposerModelChange } from "./composerModelPolicy"; + +const modelOption: CloudTaskConfigOption = { + id: "model", + name: "Model", + type: "select", + currentValue: DEFAULT_GATEWAY_MODEL, + options: [ + { value: DEFAULT_GATEWAY_MODEL, name: "Claude" }, + { value: "restricted", name: "Restricted", _meta: restrictedModelMeta() }, + { value: "gpt-5.3-codex", name: "Codex" }, + ], + category: "model", + description: "Choose a model", +}; + +it.each([ + ["claude", DEFAULT_GATEWAY_MODEL, "high", DEFAULT_GATEWAY_MODEL, "high"], + ["claude", "restricted", "high", DEFAULT_GATEWAY_MODEL, "high"], + ["claude", "missing", "high", DEFAULT_GATEWAY_MODEL, "high"], + ["codex", "gpt-5.3-codex", "xhigh", "gpt-5.3-codex", "high"], +] as const)( + "resolves %s model %s with %s reasoning", + (adapter, requestedModel, reasoning, expectedModel, expectedReasoning) => { + expect( + resolveCloudComposerModelChange({ + adapter: adapter as Adapter, + modelOption, + requestedModel, + reasoning: reasoning as SupportedReasoningEffort, + }), + ).toEqual({ model: expectedModel, reasoning: expectedReasoning }); + }, +); diff --git a/packages/core/src/task-detail/composerModelPolicy.ts b/packages/core/src/task-detail/composerModelPolicy.ts new file mode 100644 index 0000000000..d9783335f2 --- /dev/null +++ b/packages/core/src/task-detail/composerModelPolicy.ts @@ -0,0 +1,35 @@ +import { + type Adapter, + type CloudTaskConfigOption, + DEFAULT_REASONING_EFFORT, + isRestrictedModelOption, + isSupportedReasoningEffort, + type SupportedReasoningEffort, +} from "@posthog/shared"; + +export function resolveCloudComposerModelChange({ + adapter, + modelOption, + requestedModel, + reasoning, +}: { + adapter: Adapter; + modelOption: CloudTaskConfigOption; + requestedModel: string; + reasoning: SupportedReasoningEffort; +}): { model: string; reasoning: SupportedReasoningEffort } { + const selected = modelOption.options.find( + (option) => option.value === requestedModel, + ); + const model = + selected && !isRestrictedModelOption(selected._meta) + ? requestedModel + : modelOption.currentValue; + + return { + model, + reasoning: isSupportedReasoningEffort(adapter, model, reasoning) + ? reasoning + : DEFAULT_REASONING_EFFORT, + }; +} From 0a1c884c84bc23155b1d1e0420da7c00828bad89 Mon Sep 17 00:00:00 2001 From: Richard Solomou Date: Sat, 25 Jul 2026 02:14:00 +0300 Subject: [PATCH 43/45] fix(core): prefer streamed plan content Generated-By: PostHog Code Task-Id: c1bbe3cf-742b-4b24-bf96-d11a18b4cf22 --- packages/core/src/sessions/planApprovalPresentation.test.ts | 4 ++-- packages/core/src/sessions/planApprovalPresentation.ts | 6 +++--- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/packages/core/src/sessions/planApprovalPresentation.test.ts b/packages/core/src/sessions/planApprovalPresentation.test.ts index 27f9636189..8519c8a41f 100644 --- a/packages/core/src/sessions/planApprovalPresentation.test.ts +++ b/packages/core/src/sessions/planApprovalPresentation.test.ts @@ -18,12 +18,12 @@ describe("extractPlanText", () => { expect(extractPlanText(toolCall)).toBe(expected); }); - it("prefers the canonical raw plan over rendered content", () => { + it("prefers streamed content over stale raw input", () => { expect( extractPlanText({ rawInput: { plan: "Canonical" }, content: [{ text: "Rendered" }], }), - ).toBe("Canonical"); + ).toBe("Rendered"); }); }); diff --git a/packages/core/src/sessions/planApprovalPresentation.ts b/packages/core/src/sessions/planApprovalPresentation.ts index 2dd43553b6..0817576da7 100644 --- a/packages/core/src/sessions/planApprovalPresentation.ts +++ b/packages/core/src/sessions/planApprovalPresentation.ts @@ -12,12 +12,12 @@ export function extractPlanText(toolCall: { rawInput?: { plan?: unknown } | null; content?: readonly unknown[] | null; }): string | null { - const rawPlan = toolCall.rawInput?.plan; - if (typeof rawPlan === "string" && rawPlan.trim()) return rawPlan; - for (const item of toolCall.content ?? []) { const text = extractTextContent(item); if (text?.trim()) return text; } + + const rawPlan = toolCall.rawInput?.plan; + if (typeof rawPlan === "string" && rawPlan.trim()) return rawPlan; return null; } From 894534d97720a9b2fafe13bbad4c28e1fb692d61 Mon Sep 17 00:00:00 2001 From: JonathanLab Date: Wed, 29 Jul 2026 12:43:08 +0200 Subject: [PATCH 44/45] fix(pi): restore cold session controls from TaskSession --- packages/agent/src/pi/types.ts | 5 ++ .../api-client/src/posthog-client.test.ts | 36 ++++++++++ packages/api-client/src/posthog-client.ts | 29 +++++++++ .../pi-runtime/cloudPiSessionClient.test.ts | 19 ++++++ .../src/pi-runtime/cloudPiSessionClient.ts | 16 ++++- .../pi-runtime/piSessionController.test.ts | 65 +++++++++++++++++++ .../src/pi-runtime/piSessionController.ts | 63 +++++++++++++++++- .../src/pi-runtime/piSessionProvider.test.ts | 39 ++++++++++- .../core/src/pi-runtime/piSessionProvider.ts | 11 ++++ .../core/src/pi-runtime/piSessionStore.ts | 2 +- .../src/task-detail/taskCreationApiClient.ts | 5 ++ packages/core/src/task-detail/taskService.ts | 17 ++++- .../host-router/src/pi-session-factory.ts | 4 ++ .../src/routers/pi-session.router.ts | 9 +++ .../features/pi-sessions/PiSessionView.tsx | 11 +++- .../services/pi-session/pi-session.test.ts | 49 ++++++++++++++ .../src/services/pi-session/pi-session.ts | 48 +++++++++++++- .../src/services/pi-session/schemas.ts | 22 +++++++ 18 files changed, 439 insertions(+), 11 deletions(-) diff --git a/packages/agent/src/pi/types.ts b/packages/agent/src/pi/types.ts index 3fe75e7f17..1771097372 100644 --- a/packages/agent/src/pi/types.ts +++ b/packages/agent/src/pi/types.ts @@ -26,6 +26,11 @@ export type PiNativeModelInfo = Awaited< ReturnType >[number]; +export interface PiPersistedSessionConfig { + model: { provider: string; id: string } | null; + thinkingLevel: PiThinkingLevel; +} + export type PiCommand = Awaited>[number]; export type PiSessionStatus = Omit & { diff --git a/packages/api-client/src/posthog-client.test.ts b/packages/api-client/src/posthog-client.test.ts index 42a36681cb..70f5b22ab4 100644 --- a/packages/api-client/src/posthog-client.test.ts +++ b/packages/api-client/src/posthog-client.test.ts @@ -278,6 +278,42 @@ describe("PostHogAPIClient", () => { ); }); + it("loads native task session storage access", async () => { + const storage = { + id: "session-1", + download_url: "https://storage.example/session.jsonl", + content_sha256: "hash", + }; + const fetch = vi.fn().mockResolvedValue({ + ok: true, + json: async () => storage, + }); + const client = new PostHogAPIClient( + "http://localhost:8000", + async () => "token", + async () => "token", + 123, + ); + ( + client as unknown as { + api: { baseUrl: string; fetcher: { fetch: typeof fetch } }; + } + ).api = { + baseUrl: "http://localhost:8000", + fetcher: { fetch }, + }; + + await expect( + client.getTaskSessionStorageAccess("task-123", "run-123"), + ).resolves.toEqual(storage); + expect(fetch).toHaveBeenCalledWith( + expect.objectContaining({ + method: "get", + path: "/api/projects/123/tasks/task-123/runs/run-123/task_session/", + }), + ); + }); + it("maps the permission mode per adapter when creating task runs", async () => { const fetch = vi.fn().mockResolvedValue({ ok: true, diff --git a/packages/api-client/src/posthog-client.ts b/packages/api-client/src/posthog-client.ts index bcabc27b59..a8c0117b96 100644 --- a/packages/api-client/src/posthog-client.ts +++ b/packages/api-client/src/posthog-client.ts @@ -145,6 +145,12 @@ export interface TaskRunSessionLogsResult { complete: boolean; } +export interface TaskSessionStorageAccess { + id: string; + download_url: string | null; + content_sha256: string | null; +} + /** Thrown when the backend rejects a cloud run with a 429 usage-limit error. */ export class CloudUsageLimitError extends Error { limitType: UsageLimitType; @@ -2987,6 +2993,29 @@ export class PostHogAPIClient { return data.url; } + async getTaskSessionStorageAccess( + taskId: string, + runId: string, + ): Promise { + const teamId = await this.getTeamId(); + const url = new URL( + `${this.api.baseUrl}/api/projects/${teamId}/tasks/${taskId}/runs/${runId}/task_session/`, + ); + const response = await this.api.fetcher.fetch({ + method: "get", + url, + path: `/api/projects/${teamId}/tasks/${taskId}/runs/${runId}/task_session/`, + }); + if (response.status === 404) { + return null; + } + if (!response.ok) { + throw new Error(`Failed to load task session: ${response.statusText}`); + } + + return (await response.json()) as TaskSessionStorageAccess; + } + async resumeRunInCloud(taskId: string, runId: string): Promise { const teamId = await this.getTeamId(); const url = new URL( diff --git a/packages/core/src/pi-runtime/cloudPiSessionClient.test.ts b/packages/core/src/pi-runtime/cloudPiSessionClient.test.ts index 476d2c2259..14ff92b77e 100644 --- a/packages/core/src/pi-runtime/cloudPiSessionClient.test.ts +++ b/packages/core/src/pi-runtime/cloudPiSessionClient.test.ts @@ -293,6 +293,25 @@ describe("CloudPiSessionClient", () => { ]); }); + it("serves persisted native config while the cloud runtime is cold", async () => { + const cloud = createCloudTaskClient(); + const session = new CloudPiSessionClient(cloud.client, { + ...context("completed"), + persistedConfig: { + model: { provider: "posthog", id: "claude-opus-4-8" }, + thinkingLevel: "high", + }, + }); + + await expect(session.client.getState()).resolves.toMatchObject({ + thinkingLevel: "high", + }); + expect(session.persistedConfig).toEqual({ + model: { provider: "posthog", id: "claude-opus-4-8" }, + thinkingLevel: "high", + }); + }); + it("loads terminal history from the cloud snapshot without sandbox RPC", async () => { const cloud = createCloudTaskClient(); const session = new CloudPiSessionClient( diff --git a/packages/core/src/pi-runtime/cloudPiSessionClient.ts b/packages/core/src/pi-runtime/cloudPiSessionClient.ts index 1f29951487..238d508101 100644 --- a/packages/core/src/pi-runtime/cloudPiSessionClient.ts +++ b/packages/core/src/pi-runtime/cloudPiSessionClient.ts @@ -3,7 +3,10 @@ import { RemotePiRpcClient, } from "@posthog/agent/pi/remote-rpc-client"; import type { RpcCommand } from "@posthog/agent/pi/rpc-transport"; -import type { PiQueueSnapshot } from "@posthog/agent/pi/types"; +import type { + PiPersistedSessionConfig, + PiQueueSnapshot, +} from "@posthog/agent/pi/types"; import type { AgentConversationEvent, PiRuntimeHealth, @@ -21,6 +24,7 @@ import type { PiSession } from "./piSessionController"; function createTerminalPiRpcClient( runId: string, getRunStatus: () => TaskRunStatus, + persistedConfig?: PiPersistedSessionConfig | null, ): PiRemoteRpcClient { const rejectCommand = async (): Promise => { throw new Error(`Cloud task run ${runId} is ${getRunStatus()}`); @@ -34,7 +38,7 @@ function createTerminalPiRpcClient( getState: async () => ({ isStreaming: false, isCompacting: false, - thinkingLevel: "off", + thinkingLevel: persistedConfig?.thinkingLevel ?? "off", steeringMode: "all", followUpMode: "all", sessionId: runId, @@ -61,6 +65,7 @@ export interface CloudPiSessionContext { runStatus: TaskRunStatus; apiHost: string; teamId: number; + persistedConfig?: PiPersistedSessionConfig | null; } export class CloudPiSessionClient implements PiSession { @@ -106,6 +111,7 @@ export class CloudPiSessionClient implements PiSession { this.terminalClient = createTerminalPiRpcClient( context.runId, () => this.runStatus, + context.persistedConfig, ); } @@ -123,6 +129,10 @@ export class CloudPiSessionClient implements PiSession { return this.context.runId; } + get persistedConfig(): PiPersistedSessionConfig | null | undefined { + return this.context.persistedConfig; + } + get cloudStatus(): TaskRunStatus { return this.runStatus; } @@ -470,7 +480,7 @@ export class CloudPiSessionClient implements PiSession { data: { isStreaming: false, isCompacting: false, - thinkingLevel: "off", + thinkingLevel: this.context.persistedConfig?.thinkingLevel ?? "off", steeringMode: "all", followUpMode: "all", sessionId: this.context.runId, diff --git a/packages/core/src/pi-runtime/piSessionController.test.ts b/packages/core/src/pi-runtime/piSessionController.test.ts index 842359ea8e..897e1799ee 100644 --- a/packages/core/src/pi-runtime/piSessionController.test.ts +++ b/packages/core/src/pi-runtime/piSessionController.test.ts @@ -727,6 +727,71 @@ describe("PiSessionController", () => { ).toBe(false); }); + it("hydrates cold model controls from persisted native config", async () => { + const session = { + ...createSession(), + persistedConfig: { + model: { provider: "posthog", id: "claude-opus-4-8" }, + thinkingLevel: "high" as const, + }, + }; + vi.mocked(session.client.getState).mockResolvedValue({ + thinkingLevel: "high", + isStreaming: false, + isCompacting: false, + steeringMode: "all", + followUpMode: "all", + sessionId: "session-1", + autoCompactionEnabled: true, + messageCount: 0, + pendingMessageCount: 0, + }); + vi.mocked(session.client.getAvailableModels).mockResolvedValue([]); + vi.mocked(session.client.getAvailableThinkingLevels).mockResolvedValue([]); + const controller = createController(session); + + await controller.connect("task-1", "run-1"); + + expect(controller.store.getState().sessions["task-1"]).toMatchObject({ + status: { + model: { provider: "posthog", id: "claude-opus-4-8" }, + thinkingLevel: "high", + }, + models: [{ provider: "posthog", id: "claude-opus-4-8" }], + thinkingLevels: ["high"], + modelsLoaded: true, + thinkingLevelsLoaded: true, + }); + }); + + it("keeps persisted controls when the old sandbox is unavailable", async () => { + const session = { + ...createSession(), + taskRunId: "run-1", + persistedConfig: { + model: { provider: "posthog", id: "claude-opus-4-8" }, + thinkingLevel: "high" as const, + }, + }; + vi.mocked(session.client.getState).mockRejectedValue( + new Error("No active sandbox for this task run"), + ); + const controller = createController(session); + + await expect(controller.connect("task-1", "run-1")).rejects.toThrow( + "No active sandbox", + ); + + expect(controller.store.getState().sessions["task-1"]).toMatchObject({ + status: { + model: { provider: "posthog", id: "claude-opus-4-8" }, + thinkingLevel: "high", + }, + models: [{ provider: "posthog", id: "claude-opus-4-8" }], + thinkingLevels: ["high"], + }); + }); + it("resumes a terminal cloud run only when a message is submitted", async () => { const terminalSession = { ...createSession(), diff --git a/packages/core/src/pi-runtime/piSessionController.ts b/packages/core/src/pi-runtime/piSessionController.ts index b48c79dae5..f2b202350d 100644 --- a/packages/core/src/pi-runtime/piSessionController.ts +++ b/packages/core/src/pi-runtime/piSessionController.ts @@ -1,6 +1,7 @@ import type { PiRemoteRpcClient } from "@posthog/agent/pi/remote-rpc-client"; import type { PiNativeModelInfo, + PiPersistedSessionConfig, PiQueueSnapshot, PiThinkingLevel, } from "@posthog/agent/pi/types"; @@ -44,6 +45,7 @@ export interface PiSession { readonly resumeRequired?: boolean; readonly cloudStatus?: TaskRunStatus; readonly taskRunId?: string; + readonly persistedConfig?: PiPersistedSessionConfig | null; retry?(): Promise; getQueue(): Promise; clearQueue(): Promise; @@ -64,6 +66,9 @@ export interface PiSession { export interface PiSessionFactory { get(taskId: string, taskRunId?: string): Promise; + readSessionConfig?( + downloadUrl: string, + ): Promise; } export type PiSessionProvider = PiSessionFactory; @@ -519,6 +524,7 @@ export class PiSessionController { if (disposed) { return; } + this.applyPersistedConfig(taskId, session); this.updateSession(taskId, { cloudStatus: session.cloudStatus }); unsubscribe = session.onConversationEvent( (event) => this.handleEvent(taskId, event), @@ -533,6 +539,40 @@ export class PiSessionController { }); } + private applyPersistedConfig(taskId: string, session: PiSession): void { + const config = session.persistedConfig; + if (!config) { + return; + } + + const current = this.getSession(taskId); + const status = current.status + ? { + ...current.status, + model: config.model ?? undefined, + thinkingLevel: config.thinkingLevel, + } + : { + isStreaming: false, + isCompacting: false, + thinkingLevel: config.thinkingLevel, + model: config.model ?? undefined, + steeringMode: "all" as const, + followUpMode: "all" as const, + sessionId: session.taskRunId ?? taskId, + autoCompactionEnabled: true, + messageCount: current.events.length, + pendingMessageCount: 0, + }; + this.updateSession(taskId, { + status, + models: config.model ? [config.model] : [], + modelsLoaded: true, + thinkingLevels: [config.thinkingLevel], + thinkingLevelsLoaded: true, + }); + } + private async loadSession(taskId: string): Promise { const connectedSessionVersion = this.getSessionVersion(taskId); try { @@ -581,6 +621,9 @@ export class PiSessionController { : currentSession.queue; const resolvedStatus = { ...status, + model: status.model + ? { provider: status.model.provider, id: status.model.id } + : (session.persistedConfig?.model ?? undefined), pendingMessageCount: resolvedQueue.steering.length + resolvedQueue.followUp.length, }; @@ -609,13 +652,29 @@ export class PiSessionController { await Promise.all([ session.client.getAvailableModels().then((models) => { if (this.getSessionVersion(taskId) === connectedSessionVersion) { - this.updateSession(taskId, { models, modelsLoaded: true }); + const persistedModel = session.persistedConfig?.model; + this.updateSession(taskId, { + models: + models.length > 0 + ? models + : persistedModel + ? [persistedModel] + : [], + modelsLoaded: true, + }); } }), session.client.getAvailableThinkingLevels().then((thinkingLevels) => { if (this.getSessionVersion(taskId) === connectedSessionVersion) { + const persistedThinkingLevel = + session.persistedConfig?.thinkingLevel; this.updateSession(taskId, { - thinkingLevels, + thinkingLevels: + thinkingLevels.length > 0 + ? thinkingLevels + : persistedThinkingLevel + ? [persistedThinkingLevel] + : [], thinkingLevelsLoaded: true, }); } diff --git a/packages/core/src/pi-runtime/piSessionProvider.test.ts b/packages/core/src/pi-runtime/piSessionProvider.test.ts index 36ce1d9250..eca455d0e8 100644 --- a/packages/core/src/pi-runtime/piSessionProvider.test.ts +++ b/packages/core/src/pi-runtime/piSessionProvider.test.ts @@ -71,6 +71,7 @@ function cloudTaskClient(): CloudTaskClient { function taskService(environment: "local" | "cloud"): TaskService { return { + getCloudPiTaskSessionStorage: vi.fn(async () => null), getTask: vi.fn(async () => ({ id: "task-1", runtime: "pi", @@ -112,6 +113,39 @@ describe("RoutingPiSessionProvider", () => { expect(local.client.steer).not.toHaveBeenCalled(); }); + it("reads persisted configuration through Pi's session reader", async () => { + const local = localSession(); + const localSessions = { + ...localFactory(local), + readSessionConfig: vi.fn(async () => ({ + model: { provider: "posthog", id: "claude-opus-4-8" }, + thinkingLevel: "high" as const, + })), + }; + const provider = new RoutingPiSessionProvider( + localSessions, + cloudTaskClient(), + { + ...taskService("cloud"), + getCloudPiTaskSessionStorage: vi.fn(async () => ({ + id: "session-1", + download_url: "https://storage.example/session.jsonl", + content_sha256: "hash", + })), + } as unknown as TaskService, + ); + + const session = await provider.get("task-1"); + + expect(localSessions.readSessionConfig).toHaveBeenCalledWith( + "https://storage.example/session.jsonl", + ); + expect(session.persistedConfig).toEqual({ + model: { provider: "posthog", id: "claude-opus-4-8" }, + thinkingLevel: "high", + }); + }); + it("binds explicit historical runs and invalidates changed run context", async () => { const local = localSession(); const cloudTasks = cloudTaskClient(); @@ -127,7 +161,10 @@ describe("RoutingPiSessionProvider", () => { const provider = new RoutingPiSessionProvider( localFactory(local), cloudTasks, - { getTask } as unknown as TaskService, + { + getTask, + getCloudPiTaskSessionStorage: vi.fn(async () => null), + } as unknown as TaskService, ); const historical = await provider.get("task-1", "run-old"); diff --git a/packages/core/src/pi-runtime/piSessionProvider.ts b/packages/core/src/pi-runtime/piSessionProvider.ts index 30528f0614..b5dc0c9ae5 100644 --- a/packages/core/src/pi-runtime/piSessionProvider.ts +++ b/packages/core/src/pi-runtime/piSessionProvider.ts @@ -47,10 +47,21 @@ export class RoutingPiSessionProvider implements PiSessionProvider { return null; } + const storage = await this.taskService + .getCloudPiTaskSessionStorage(taskId, run.id) + .catch(() => null); + const persistedConfig = + storage?.download_url && this.localFactory.readSessionConfig + ? await this.localFactory + .readSessionConfig(storage.download_url) + .catch(() => null) + : null; + return { taskId, runId: run.id, runStatus: run.status, + persistedConfig, ...context, }; } diff --git a/packages/core/src/pi-runtime/piSessionStore.ts b/packages/core/src/pi-runtime/piSessionStore.ts index 73a9073801..a1dd61b0e1 100644 --- a/packages/core/src/pi-runtime/piSessionStore.ts +++ b/packages/core/src/pi-runtime/piSessionStore.ts @@ -29,7 +29,7 @@ export interface PiSessionError { export interface PiControllerSessionState { connectionState: SessionStatus; events: AgentConversationEvent[]; - models: PiNativeModelInfo[]; + models: Array>; modelsLoaded: boolean; thinkingLevels: PiThinkingLevel[]; thinkingLevelsLoaded: boolean; diff --git a/packages/core/src/task-detail/taskCreationApiClient.ts b/packages/core/src/task-detail/taskCreationApiClient.ts index 0677f4352b..428189f9b0 100644 --- a/packages/core/src/task-detail/taskCreationApiClient.ts +++ b/packages/core/src/task-detail/taskCreationApiClient.ts @@ -1,3 +1,4 @@ +import type { TaskSessionStorageAccess } from "@posthog/api-client/posthog-client"; import type { Adapter, CloudMcpServerImport, @@ -46,5 +47,9 @@ export interface TaskCreationApiClient { runId: string, options?: StartTaskRunClientOptions, ): Promise; + getTaskSessionStorageAccess( + taskId: string, + runId: string, + ): Promise; resumeRunInCloud(taskId: string, runId: string): Promise; } diff --git a/packages/core/src/task-detail/taskService.ts b/packages/core/src/task-detail/taskService.ts index f9601d463f..52c76216a6 100644 --- a/packages/core/src/task-detail/taskService.ts +++ b/packages/core/src/task-detail/taskService.ts @@ -1,4 +1,7 @@ -import { CLOUD_USAGE_LIMIT_ERROR_MESSAGE } from "@posthog/api-client/posthog-client"; +import { + CLOUD_USAGE_LIMIT_ERROR_MESSAGE, + type TaskSessionStorageAccess, +} from "@posthog/api-client/posthog-client"; import { SESSION_SERVICE, type SessionService, @@ -182,6 +185,18 @@ export class TaskService { return task; } + public async getCloudPiTaskSessionStorage( + taskId: string, + taskRunId: string, + ): Promise { + const posthogClient = await this.host.getAuthenticatedClient(); + if (!posthogClient) { + throw new Error("Not authenticated"); + } + + return posthogClient.getTaskSessionStorageAccess(taskId, taskRunId); + } + public async resumeCloudPiRun( taskId: string, taskRunId: string, diff --git a/packages/host-router/src/pi-session-factory.ts b/packages/host-router/src/pi-session-factory.ts index 9ff1c0eaad..79f133ef66 100644 --- a/packages/host-router/src/pi-session-factory.ts +++ b/packages/host-router/src/pi-session-factory.ts @@ -66,4 +66,8 @@ export class TrpcPiSessionFactory implements PiSessionFactory { get(taskId: string): Promise { return Promise.resolve(new TrpcPiSession(this.client, taskId)); } + + readSessionConfig(downloadUrl: string) { + return this.client.piSession.readSessionConfig.query({ downloadUrl }); + } } diff --git a/packages/host-router/src/routers/pi-session.router.ts b/packages/host-router/src/routers/pi-session.router.ts index 1ab86122b2..23a6705c68 100644 --- a/packages/host-router/src/routers/pi-session.router.ts +++ b/packages/host-router/src/routers/pi-session.router.ts @@ -4,6 +4,8 @@ import type { PiSessionService } from "@posthog/workspace-server/services/pi-ses import { piQueueSnapshotOutput, piRpcResponseSchema, + piSessionConfigInput, + piSessionConfigOutput, piSessionHealthOutput, piSessionRpcInput, piSessionStartOutput, @@ -41,6 +43,13 @@ export const piSessionRouter = router({ .output(piSessionHealthOutput) .query(({ ctx, input }) => getService(ctx.container).health(input.taskId)), + readSessionConfig: publicProcedure + .input(piSessionConfigInput) + .output(piSessionConfigOutput) + .query(({ ctx, input }) => + getService(ctx.container).readSessionConfig(input.downloadUrl), + ), + getQueue: publicProcedure .input(piSessionTaskInput) .output(piQueueSnapshotOutput) diff --git a/packages/ui/src/features/pi-sessions/PiSessionView.tsx b/packages/ui/src/features/pi-sessions/PiSessionView.tsx index 9cd84670ed..d3e2c4317e 100644 --- a/packages/ui/src/features/pi-sessions/PiSessionView.tsx +++ b/packages/ui/src/features/pi-sessions/PiSessionView.tsx @@ -22,6 +22,7 @@ import { Skeleton, } from "@posthog/quill"; import type { AgentConversationEvent } from "@posthog/shared"; +import { isTerminalStatus } from "@posthog/shared/domain-types"; import { useUsageLimitStore } from "@posthog/ui/features/billing/usageLimitStore"; import { PromptInput } from "@posthog/ui/features/message-editor/components/PromptInput"; import { useDraftStore } from "@posthog/ui/features/message-editor/draftStore"; @@ -353,6 +354,12 @@ export function PiSessionView({ taskId, taskRunId }: PiSessionViewProps) { } const controlsPending = status ? isStreaming || isBashRunning : false; + const controlsDisabled = + controlsPending || + isCompacting || + session.connectionState !== "connected" || + (session.cloudStatus !== undefined && + isTerminalStatus(session.cloudStatus)); const hasQueuedMessage = session.queue.steering.length + session.queue.followUp.length > 0; let modelSelector: ReactElement = ( @@ -370,7 +377,7 @@ export function PiSessionView({ taskId, taskRunId }: PiSessionViewProps) { ); @@ -384,7 +391,7 @@ export function PiSessionView({ taskId, taskRunId }: PiSessionViewProps) { ) : null; diff --git a/packages/workspace-server/src/services/pi-session/pi-session.test.ts b/packages/workspace-server/src/services/pi-session/pi-session.test.ts index b15fb58dd3..9195c8e53d 100644 --- a/packages/workspace-server/src/services/pi-session/pi-session.test.ts +++ b/packages/workspace-server/src/services/pi-session/pi-session.test.ts @@ -29,6 +29,7 @@ function successfulResponse(command: string): RpcResponse { afterEach(() => { vi.unstubAllEnvs(); + vi.unstubAllGlobals(); vi.restoreAllMocks(); }); @@ -124,6 +125,54 @@ describe("selectPiPoolEvictionCandidate", () => { }); }); +describe("PiSessionService task session config", () => { + it("uses Pi session context resolution for model and thinking", async () => { + const content = [ + { + type: "session", + version: 3, + id: "session-1", + timestamp: "2026-01-01T00:00:00.000Z", + cwd: "/repo", + }, + { + type: "model_change", + id: "model-1", + parentId: null, + timestamp: "2026-01-01T00:00:01.000Z", + provider: "posthog", + modelId: "claude-opus-4-8", + }, + { + type: "thinking_level_change", + id: "thinking-1", + parentId: "model-1", + timestamp: "2026-01-01T00:00:02.000Z", + thinkingLevel: "high", + }, + ] + .map((entry) => JSON.stringify(entry)) + .join("\n"); + vi.stubGlobal( + "fetch", + vi.fn(async () => new Response(content)), + ); + const service = new PiSessionService( + {} as PiRuntimeFactory, + {} as ITaskMetadataRepository, + {} as ProcessTrackingService, + rootLogger, + ); + + await expect( + service.readSessionConfig("https://storage.example/session.jsonl"), + ).resolves.toEqual({ + model: { provider: "posthog", id: "claude-opus-4-8" }, + thinkingLevel: "high", + }); + }); +}); + describe("PiSessionService start", () => { it("sets the selected thinking level before the initial prompt", async () => { const setThinkingLevel = vi.fn().mockResolvedValue(undefined); diff --git a/packages/workspace-server/src/services/pi-session/pi-session.ts b/packages/workspace-server/src/services/pi-session/pi-session.ts index bef202c952..cdc1f7af15 100644 --- a/packages/workspace-server/src/services/pi-session/pi-session.ts +++ b/packages/workspace-server/src/services/pi-session/pi-session.ts @@ -1,7 +1,18 @@ +import { + buildSessionContext, + type FileEntry, + migrateSessionEntries, + parseSessionEntries, + type SessionEntry, +} from "@earendil-works/pi-coding-agent"; import type { PiRpcClient } from "@posthog/agent/pi/rpc-client"; import type { RpcCommand, RpcResponse } from "@posthog/agent/pi/rpc-transport"; import type { PiRuntime } from "@posthog/agent/pi/runtime"; -import type { PiQueueSnapshot } from "@posthog/agent/pi/types"; +import { + PI_THINKING_LEVELS, + type PiPersistedSessionConfig, + type PiQueueSnapshot, +} from "@posthog/agent/pi/types"; import { ROOT_LOGGER, type RootLogger } from "@posthog/di/logger"; import { type AgentConversationEvent, @@ -201,6 +212,41 @@ export class PiSessionService extends TypedEventEmitter { }); } + async readSessionConfig( + downloadUrl: string, + ): Promise { + const response = await fetch(downloadUrl, { + signal: AbortSignal.timeout(30_000), + }); + if (response.status === 404) { + return null; + } + if (!response.ok) { + throw new Error( + `Failed to download Pi task session: ${response.statusText}`, + ); + } + + const fileEntries = parseSessionEntries( + await response.text(), + ) as FileEntry[]; + migrateSessionEntries(fileEntries); + const entries = fileEntries.filter( + (entry): entry is SessionEntry => entry.type !== "session", + ); + const context = buildSessionContext(entries); + const thinkingLevel = PI_THINKING_LEVELS.find( + (level) => level === context.thinkingLevel, + ); + + return { + model: context.model + ? { provider: context.model.provider, id: context.model.modelId } + : null, + thinkingLevel: thinkingLevel ?? "off", + }; + } + async stop(taskId: string): Promise { await this.runExclusive(taskId, () => this.stopLocked(taskId)); } diff --git a/packages/workspace-server/src/services/pi-session/schemas.ts b/packages/workspace-server/src/services/pi-session/schemas.ts index ff8a401cff..344b6738a9 100644 --- a/packages/workspace-server/src/services/pi-session/schemas.ts +++ b/packages/workspace-server/src/services/pi-session/schemas.ts @@ -36,6 +36,28 @@ export const resumePiSessionInput = z.object({ export const piSessionTaskInput = z.object({ taskId: z.string() }); +export const piSessionConfigInput = z.object({ downloadUrl: z.url() }); + +export const piSessionConfigOutput = z + .object({ + model: z + .object({ + provider: z.string(), + id: z.string(), + }) + .nullable(), + thinkingLevel: z.enum([ + "off", + "minimal", + "low", + "medium", + "high", + "xhigh", + "max", + ]), + }) + .nullable(); + export const piQueueSnapshotOutput = z.object({ steering: z.array(z.string()), followUp: z.array(z.string()), From 97cd41cb5f55e4164593e163fadcd739080cb3ad Mon Sep 17 00:00:00 2001 From: JonathanLab Date: Wed, 29 Jul 2026 13:06:03 +0200 Subject: [PATCH 45/45] fix(pi): disable unsupported harness extensions --- packages/harness/src/extensions/registry.ts | 6 ------ packages/harness/src/runtime.test.ts | 16 ++++++++++------ 2 files changed, 10 insertions(+), 12 deletions(-) diff --git a/packages/harness/src/extensions/registry.ts b/packages/harness/src/extensions/registry.ts index 1e19739ca7..ac813b00c4 100644 --- a/packages/harness/src/extensions/registry.ts +++ b/packages/harness/src/extensions/registry.ts @@ -3,15 +3,12 @@ import type { ExtensionFactory, InlineExtension, } from "@earendil-works/pi-coding-agent"; -import { createBackgroundJobsExtension } from "./background-jobs/extension"; import type { HogBrandingOptions } from "./hog-branding/extension"; import { createHogBrandingExtension } from "./hog-branding/extension"; import { createMcpExtension } from "./mcp/extension"; import { createPosthogProviderExtension } from "./posthog-provider/extension"; import type { PosthogProviderOptions } from "./posthog-provider/provider"; -import { createSubagentExtension } from "./subagent/extension"; import { createWebAccessExtension } from "./web-access/extension"; -import { createWorkflowExtension } from "./workflow/extension"; export type HarnessExtensionOptions = PosthogProviderOptions & HogBrandingOptions; @@ -25,9 +22,6 @@ const EXTENSIONS: HarnessExtension[] = [ { name: "hog-branding", create: createHogBrandingExtension }, { name: "posthog-provider", create: createPosthogProviderExtension }, { name: "web-access", create: createWebAccessExtension }, - { name: "background-jobs", create: () => createBackgroundJobsExtension() }, - { name: "subagent", create: createSubagentExtension }, - { name: "workflow", create: createWorkflowExtension }, // createMcpExtension's options are test seams (config loader, transport // factory), not HarnessExtensionOptions, so drop the registry options. { name: "mcp", create: () => createMcpExtension() }, diff --git a/packages/harness/src/runtime.test.ts b/packages/harness/src/runtime.test.ts index c7ab40634f..6314e26b35 100644 --- a/packages/harness/src/runtime.test.ts +++ b/packages/harness/src/runtime.test.ts @@ -45,18 +45,22 @@ describe("createHarnessRuntime", () => { expect(runtime.session.model?.provider).toBe("posthog"); expect(runtime.session.getAvailableThinkingLevels()).toContain("off"); expect(runtime.services.settingsManager.isProjectTrusted()).toBe(false); - expect( - runtime.services.resourceLoader - .getExtensions() - .extensions.map((extension) => extension.path), - ).toEqual( + const extensionPaths = runtime.services.resourceLoader + .getExtensions() + .extensions.map((extension) => extension.path); + expect(extensionPaths).toEqual( expect.arrayContaining([ "", "", "", + "", + ]), + ); + expect(extensionPaths).not.toEqual( + expect.arrayContaining([ + "", "", "", - "", ]), ); } finally {