From fe64274cc707a255749596e5efa3b7f5684040fb Mon Sep 17 00:00:00 2001 From: JonathanLab Date: Wed, 22 Jul 2026 17:20:25 +0200 Subject: [PATCH 01/28] 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/28] 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/28] 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/28] 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/28] 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/28] 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/28] 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/28] 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/28] 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 462ee21388eb963bcac2d66a5a65e84b56af48ba Mon Sep 17 00:00:00 2001 From: JonathanLab Date: Tue, 28 Jul 2026 12:18:03 +0200 Subject: [PATCH 19/28] 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 20/28] 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 21/28] 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 22/28] 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 23/28] 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 24/28] 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 25/28] 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 26/28] 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 894534d97720a9b2fafe13bbad4c28e1fb692d61 Mon Sep 17 00:00:00 2001 From: JonathanLab Date: Wed, 29 Jul 2026 12:43:08 +0200 Subject: [PATCH 27/28] 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 28/28] 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 {