From ecb0f31fb4172e29ef05aadafdbb433d40c50d39 Mon Sep 17 00:00:00 2001 From: JonathanLab Date: Mon, 13 Jul 2026 14:38:20 +0200 Subject: [PATCH 1/3] feat(harness): expose native Pi runtime --- packages/agent/tsup.config.ts | 1 + 1 file changed, 1 insertion(+) diff --git a/packages/agent/tsup.config.ts b/packages/agent/tsup.config.ts index 93a1e2ce88..cc9513da78 100644 --- a/packages/agent/tsup.config.ts +++ b/packages/agent/tsup.config.ts @@ -114,6 +114,7 @@ export default defineConfig([ "src/posthog-products.ts", "src/pr-url-detector.ts", "src/pi/rpc-client.ts", + "src/pi/rpc-host.ts", "src/resume.ts", "src/types.ts", "src/adapters/claude/questions/utils.ts", From 1c87a93f0ee97f756b08705c3e5d03fedd4cb749 Mon Sep 17 00:00:00 2001 From: JonathanLab Date: Wed, 15 Jul 2026 13:17:58 +0200 Subject: [PATCH 2/3] feat(harness): integrate pi runtime into desktop app --- packages/agent/tsup.config.ts | 1 - 1 file changed, 1 deletion(-) diff --git a/packages/agent/tsup.config.ts b/packages/agent/tsup.config.ts index cc9513da78..93a1e2ce88 100644 --- a/packages/agent/tsup.config.ts +++ b/packages/agent/tsup.config.ts @@ -114,7 +114,6 @@ export default defineConfig([ "src/posthog-products.ts", "src/pr-url-detector.ts", "src/pi/rpc-client.ts", - "src/pi/rpc-host.ts", "src/resume.ts", "src/types.ts", "src/adapters/claude/questions/utils.ts", From 4287e6937dd4d0bf0c54f184be9cbfddd8a0122a Mon Sep 17 00:00:00 2001 From: JonathanLab Date: Wed, 15 Jul 2026 15:53:59 +0200 Subject: [PATCH 3/3] feat: hook up pi runtime to ChatThread --- apps/code/src/main/di/bindings.ts | 4 + apps/code/src/main/di/container.ts | 50 +- .../desktop-pi-runtime-factory.test.ts | 20 + .../desktop-pi-runtime-factory.ts | 24 + apps/code/src/renderer/di/bindings.ts | 5 + apps/code/src/renderer/di/container.ts | 5 + .../platform-adapters/trpc-pi-runner.ts | 17 +- apps/web/src/web-container.ts | 9 + packages/agent/package.json | 17 +- .../agent/src/pi/conversation/toolKind.ts | 14 + .../src/pi/conversation/toolTranslator.ts | 22 + .../conversation/tools/bashTranslator.test.ts | 83 ++++ .../pi/conversation/tools/bashTranslator.ts | 22 + .../conversation/tools/editTranslator.test.ts | 102 ++++ .../pi/conversation/tools/editTranslator.ts | 55 +++ .../conversation/tools/findTranslator.test.ts | 51 ++ .../pi/conversation/tools/findTranslator.ts | 26 ++ .../conversation/tools/grepTranslator.test.ts | 59 +++ .../pi/conversation/tools/grepTranslator.ts | 55 +++ .../conversation/tools/lsTranslator.test.ts | 57 +++ .../src/pi/conversation/tools/lsTranslator.ts | 33 ++ .../conversation/tools/readTranslator.test.ts | 94 ++++ .../pi/conversation/tools/readTranslator.ts | 35 ++ .../tools/writeTranslator.test.ts | 40 ++ .../pi/conversation/tools/writeTranslator.ts | 18 + .../translatePiConversation.test.ts | 434 ++++++++++++++++++ .../conversation/translatePiConversation.ts | 348 ++++++++++++++ .../conversation/translatePiMessage.test.ts | 113 +++++ .../src/pi/conversation/translatePiMessage.ts | 299 ++++++++++++ packages/agent/src/pi/rpc-client.test.ts | 35 +- packages/agent/src/pi/rpc-client.ts | 19 + packages/agent/src/pi/runtime.test.ts | 76 +++ packages/agent/src/pi/runtime.ts | 83 ++++ packages/agent/src/pi/types.ts | 45 ++ packages/agent/tsconfig.build.json | 8 +- packages/agent/tsup.config.ts | 5 +- packages/core/package.json | 1 + packages/core/src/pi-runtime/identifiers.ts | 1 + .../core/src/pi-runtime/pi-runtime.module.ts | 7 + .../pi-runtime/piSessionController.test.ts | 197 ++++++++ .../src/pi-runtime/piSessionController.ts | 352 ++++++++++++++ .../core/src/pi-runtime/piSessionStore.ts | 37 ++ packages/host-router/src/pi-session-client.ts | 96 ++++ .../src/routers/pi-session.router.ts | 272 ++++++++++- packages/shared/src/agent-conversation.ts | 158 +++++++ packages/shared/src/index.ts | 23 + packages/shared/src/pi-session.ts | 7 + packages/ui/package.json | 1 - .../pi-sessions/PiSessionControls.tsx | 259 +++++++++++ .../features/pi-sessions/PiSessionView.tsx | 367 +++++++++------ .../pi-sessions/piSessionFeed.test.ts | 143 ------ .../src/features/pi-sessions/piSessionFeed.ts | 99 ---- .../pi-sessions/useEnsurePiSession.ts | 30 -- .../sessions/components/ThreadView.tsx | 13 +- .../buildAgentConversationItems.test.ts | 207 +++++++++ .../components/buildConversationItems.ts | 237 ++++++++-- .../chat-thread/ChatMarkdown.test.tsx | 15 + .../components/chat-thread/ChatMarkdown.tsx | 5 + .../components/chat-thread/ChatThread.tsx | 74 ++- .../new-thread/buildThreadGroups.ts | 2 +- .../session-update/SessionUpdateView.tsx | 12 +- .../StatusNotificationView.test.ts | 18 + .../session-update/StatusNotificationView.tsx | 83 +++- .../hooks/useAgentConversationItems.ts | 16 + packages/ui/src/features/sessions/types.ts | 68 ++- .../services/auth-proxy/auth-proxy.test.ts | 20 + .../src/services/auth-proxy/auth-proxy.ts | 6 +- .../src/services/mcp-proxy/mcp-proxy.test.ts | 6 + .../src/services/mcp-proxy/mcp-proxy.ts | 6 +- .../src/services/pi-session/identifiers.ts | 13 + .../services/pi-session/pi-session.test.ts | 37 ++ .../src/services/pi-session/pi-session.ts | 412 ++++++++++++++++- .../src/services/pi-session/schemas.ts | 304 +++++++++++- pnpm-lock.yaml | 91 ++-- pnpm-workspace.yaml | 2 +- 75 files changed, 5442 insertions(+), 637 deletions(-) create mode 100644 apps/code/src/main/platform-adapters/desktop-pi-runtime-factory.test.ts create mode 100644 apps/code/src/main/platform-adapters/desktop-pi-runtime-factory.ts create mode 100644 packages/agent/src/pi/conversation/toolKind.ts create mode 100644 packages/agent/src/pi/conversation/toolTranslator.ts create mode 100644 packages/agent/src/pi/conversation/tools/bashTranslator.test.ts create mode 100644 packages/agent/src/pi/conversation/tools/bashTranslator.ts create mode 100644 packages/agent/src/pi/conversation/tools/editTranslator.test.ts create mode 100644 packages/agent/src/pi/conversation/tools/editTranslator.ts create mode 100644 packages/agent/src/pi/conversation/tools/findTranslator.test.ts create mode 100644 packages/agent/src/pi/conversation/tools/findTranslator.ts create mode 100644 packages/agent/src/pi/conversation/tools/grepTranslator.test.ts create mode 100644 packages/agent/src/pi/conversation/tools/grepTranslator.ts create mode 100644 packages/agent/src/pi/conversation/tools/lsTranslator.test.ts create mode 100644 packages/agent/src/pi/conversation/tools/lsTranslator.ts create mode 100644 packages/agent/src/pi/conversation/tools/readTranslator.test.ts create mode 100644 packages/agent/src/pi/conversation/tools/readTranslator.ts create mode 100644 packages/agent/src/pi/conversation/tools/writeTranslator.test.ts create mode 100644 packages/agent/src/pi/conversation/tools/writeTranslator.ts create mode 100644 packages/agent/src/pi/conversation/translatePiConversation.test.ts create mode 100644 packages/agent/src/pi/conversation/translatePiConversation.ts create mode 100644 packages/agent/src/pi/conversation/translatePiMessage.test.ts create mode 100644 packages/agent/src/pi/conversation/translatePiMessage.ts create mode 100644 packages/agent/src/pi/runtime.test.ts create mode 100644 packages/agent/src/pi/runtime.ts create mode 100644 packages/agent/src/pi/types.ts create mode 100644 packages/core/src/pi-runtime/pi-runtime.module.ts create mode 100644 packages/core/src/pi-runtime/piSessionController.test.ts create mode 100644 packages/core/src/pi-runtime/piSessionController.ts create mode 100644 packages/core/src/pi-runtime/piSessionStore.ts create mode 100644 packages/host-router/src/pi-session-client.ts create mode 100644 packages/shared/src/agent-conversation.ts create mode 100644 packages/shared/src/pi-session.ts create mode 100644 packages/ui/src/features/pi-sessions/PiSessionControls.tsx delete mode 100644 packages/ui/src/features/pi-sessions/piSessionFeed.test.ts delete mode 100644 packages/ui/src/features/pi-sessions/piSessionFeed.ts delete mode 100644 packages/ui/src/features/pi-sessions/useEnsurePiSession.ts create mode 100644 packages/ui/src/features/sessions/components/buildAgentConversationItems.test.ts create mode 100644 packages/ui/src/features/sessions/components/chat-thread/ChatMarkdown.test.tsx create mode 100644 packages/ui/src/features/sessions/components/session-update/StatusNotificationView.test.ts create mode 100644 packages/ui/src/features/sessions/hooks/useAgentConversationItems.ts create mode 100644 packages/workspace-server/src/services/pi-session/pi-session.test.ts diff --git a/apps/code/src/main/di/bindings.ts b/apps/code/src/main/di/bindings.ts index 7bd793ee9a..6cc67ecbf5 100644 --- a/apps/code/src/main/di/bindings.ts +++ b/apps/code/src/main/di/bindings.ts @@ -188,7 +188,9 @@ import type { } from "@posthog/workspace-server/services/mcp-relay/identifiers"; import type { PI_RPC_CLIENT_FACTORY, + PI_RUNTIME_FACTORY, PiRpcClientFactory, + PiRuntimeFactory, } from "@posthog/workspace-server/services/pi-session/identifiers"; import type { PosthogPluginService } from "@posthog/workspace-server/services/posthog-plugin/posthog-plugin"; import type { ProcessTrackingService } from "@posthog/workspace-server/services/process-tracking/process-tracking"; @@ -356,6 +358,8 @@ export interface MainBindings { [AGENT_LOGGER]: RootLogger; [PI_RPC_CLIENT_FACTORY]: PiRpcClientFactory; + [PI_RUNTIME_FACTORY]: PiRuntimeFactory; + // Logger [ROOT_LOGGER]: RootLogger; diff --git a/apps/code/src/main/di/container.ts b/apps/code/src/main/di/container.ts index 00922bc392..e9f84c43ad 100644 --- a/apps/code/src/main/di/container.ts +++ b/apps/code/src/main/di/container.ts @@ -189,6 +189,7 @@ import { onboardingImportModule } from "@posthog/workspace-server/services/onboa import { osModule } from "@posthog/workspace-server/services/os/os.module"; import { PI_RPC_CLIENT_FACTORY, + PI_RUNTIME_FACTORY, PI_SESSION_SERVICE, } from "@posthog/workspace-server/services/pi-session/identifiers"; import type { PiSessionService } from "@posthog/workspace-server/services/pi-session/pi-session"; @@ -230,6 +231,7 @@ import { workspaceMetadataModule } from "@posthog/workspace-server/services/work import ExternalAppsStoreImpl from "electron-store"; import type { FileWatcherBridge } from "../index"; import { DesktopPiRpcClientFactory } from "../platform-adapters/desktop-pi-rpc-client-factory"; +import { DesktopPiRuntimeFactory } from "../platform-adapters/desktop-pi-runtime-factory"; import { ElectronAppLifecycle } from "../platform-adapters/electron-app-lifecycle"; import { ElectronAppMeta } from "../platform-adapters/electron-app-meta"; import { ElectronAppMetrics } from "../platform-adapters/electron-app-metrics"; @@ -325,6 +327,17 @@ import { WORKTREE_REPOSITORY as MAIN_WORKTREE_REPOSITORY, } from "./tokens"; +async function cancelTaskSessions( + agentService: AgentService, + piSessionService: PiSessionService, + taskId: string, +): Promise { + await Promise.all([ + agentService.cancelSessionsByTaskId(taskId), + piSessionService.stop(taskId), + ]); +} + export const container = new TypedContainer({ defaultScope: "Singleton", }); @@ -366,6 +379,7 @@ container .bind(MAIN_DEFAULT_ADDITIONAL_DIRECTORY_REPOSITORY) .toService(DEFAULT_ADDITIONAL_DIRECTORY_REPOSITORY); container.load(agentModule); +container.bind(PI_RUNTIME_FACTORY).to(DesktopPiRuntimeFactory); container.load(piSessionModule); container.bind(AGENT_SLEEP_COORDINATOR).toService(MAIN_SLEEP_SERVICE); container.bind(AGENT_MCP_APPS).toService(MCP_APPS_SERVICE); @@ -406,12 +420,12 @@ container.bind(MCP_PROXY_AUTH).toDynamicValue((ctx) => { }); container.load(archiveModule); container.bind(ARCHIVE_SESSION_CANCELLER).toDynamicValue((ctx) => ({ - cancelSessionsByTaskId: async (taskId: string) => { - await Promise.all([ - ctx.get(AGENT_SERVICE).cancelSessionsByTaskId(taskId), - ctx.get(PI_SESSION_SERVICE).stop(taskId), - ]); - }, + cancelSessionsByTaskId: (taskId: string) => + cancelTaskSessions( + ctx.get(AGENT_SERVICE), + ctx.get(PI_SESSION_SERVICE), + taskId, + ), })); container.bind(ARCHIVE_FILE_WATCHER).toDynamicValue((ctx) => ({ stopWatching: async (worktreePath: string) => { @@ -422,12 +436,12 @@ container.bind(ARCHIVE_FILE_WATCHER).toDynamicValue((ctx) => ({ })); container.load(suspensionModule); container.bind(SUSPENSION_SESSION_CANCELLER).toDynamicValue((ctx) => ({ - cancelSessionsByTaskId: async (taskId: string) => { - await Promise.all([ - ctx.get(AGENT_SERVICE).cancelSessionsByTaskId(taskId), - ctx.get(PI_SESSION_SERVICE).stop(taskId), - ]); - }, + cancelSessionsByTaskId: (taskId: string) => + cancelTaskSessions( + ctx.get(AGENT_SERVICE), + ctx.get(PI_SESSION_SERVICE), + taskId, + ), })); container.bind(SUSPENSION_FILE_WATCHER).toDynamicValue((ctx) => ({ stopWatching: async (worktreePath: string) => { @@ -705,12 +719,12 @@ container.load(workspaceModule); container.bind(WORKSPACE_AGENT).toDynamicValue((ctx): WorkspaceAgent => { const agent = ctx.get(AGENT_SERVICE); return { - cancelSessionsByTaskId: async (taskId) => { - await Promise.all([ - agent.cancelSessionsByTaskId(taskId), - ctx.get(PI_SESSION_SERVICE).stop(taskId), - ]); - }, + cancelSessionsByTaskId: (taskId) => + cancelTaskSessions( + agent, + ctx.get(PI_SESSION_SERVICE), + taskId, + ), onAgentFileActivity: (handler) => agent.on(AgentServiceEvent.AgentFileActivity, handler), }; diff --git a/apps/code/src/main/platform-adapters/desktop-pi-runtime-factory.test.ts b/apps/code/src/main/platform-adapters/desktop-pi-runtime-factory.test.ts new file mode 100644 index 0000000000..d9b7a19f3f --- /dev/null +++ b/apps/code/src/main/platform-adapters/desktop-pi-runtime-factory.test.ts @@ -0,0 +1,20 @@ +import { PiRuntime } from "@posthog/agent/pi/runtime"; +import type { PiRpcClientFactory } from "@posthog/workspace-server/services/pi-session/identifiers"; +import { describe, expect, it, vi } from "vitest"; +import { DesktopPiRuntimeFactory } from "./desktop-pi-runtime-factory"; + +describe("DesktopPiRuntimeFactory", () => { + it("wraps the host-authenticated RPC client", async () => { + const client = { onEvent: vi.fn() }; + const clientFactory = { + create: vi.fn(async () => client), + } as unknown as PiRpcClientFactory; + const factory = new DesktopPiRuntimeFactory(clientFactory); + + const runtime = await factory.create({ cwd: "/workspace" }); + + expect(runtime).toBeInstanceOf(PiRuntime); + expect(runtime.client).toBe(client); + expect(clientFactory.create).toHaveBeenCalledWith({ cwd: "/workspace" }); + }); +}); diff --git a/apps/code/src/main/platform-adapters/desktop-pi-runtime-factory.ts b/apps/code/src/main/platform-adapters/desktop-pi-runtime-factory.ts new file mode 100644 index 0000000000..b085050ff5 --- /dev/null +++ b/apps/code/src/main/platform-adapters/desktop-pi-runtime-factory.ts @@ -0,0 +1,24 @@ +import { PiRuntime } from "@posthog/agent/pi/runtime"; +import { + PI_RPC_CLIENT_FACTORY, + type PiRpcClientFactory, + type PiRuntimeFactory, +} from "@posthog/workspace-server/services/pi-session/identifiers"; +import { inject, injectable } from "inversify"; + +@injectable() +export class DesktopPiRuntimeFactory implements PiRuntimeFactory { + constructor( + @inject(PI_RPC_CLIENT_FACTORY) + private readonly clientFactory: PiRpcClientFactory, + ) {} + + async create(input: { + cwd: string; + model?: string; + sessionFile?: string; + }): Promise { + const client = await this.clientFactory.create(input); + return new PiRuntime(client); + } +} diff --git a/apps/code/src/renderer/di/bindings.ts b/apps/code/src/renderer/di/bindings.ts index 3d25557e3f..f8556c970f 100644 --- a/apps/code/src/renderer/di/bindings.ts +++ b/apps/code/src/renderer/di/bindings.ts @@ -75,6 +75,10 @@ import { } from "@posthog/core/onboarding/identifiers"; import { PI_RUNNER } from "@posthog/core/pi-runtime/identifiers"; import type { PiRunner } from "@posthog/core/pi-runtime/piRunner"; +import { + PI_SESSION_CLIENT, + type PiSessionClient, +} from "@posthog/core/pi-runtime/piSessionController"; import { type BundleLocalSkill, CLOUD_ARTIFACT_BUNDLE_LOCAL_SKILL, @@ -291,6 +295,7 @@ export interface RendererBindings { [ANALYTICS_TRACKER]: AnalyticsTracker; [TASK_CREATION_HOST]: ITaskCreationHost; [PI_RUNNER]: PiRunner; + [PI_SESSION_CLIENT]: PiSessionClient; [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 9006191ad1..5d7d27ddf4 100644 --- a/apps/code/src/renderer/di/container.ts +++ b/apps/code/src/renderer/di/container.ts @@ -33,7 +33,9 @@ import type { LlmMessage } from "@posthog/core/llm-gateway/schemas"; import { LOCAL_MCP_WORKSPACE_CLIENT } from "@posthog/core/local-mcp/identifiers"; import type { LocalMcpWorkspaceClient } from "@posthog/core/local-mcp/localMcpImport"; 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 { CLOUD_ARTIFACT_BUNDLE_LOCAL_SKILL, CLOUD_ARTIFACT_READ_FILE_AS_BASE64, @@ -87,6 +89,7 @@ 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 { BROWSER_TABS_CLIENT, type BrowserTabsClient, @@ -294,6 +297,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.load(piRuntimeModule); container.bind(TASK_CREATION_EFFECTS).toConstantValue(taskCreationEffects); container.bind(RENDERER_TASK_SERVICE).to(TaskService); container.bind(TASK_SERVICE).toService(RENDERER_TASK_SERVICE); diff --git a/apps/code/src/renderer/platform-adapters/trpc-pi-runner.ts b/apps/code/src/renderer/platform-adapters/trpc-pi-runner.ts index 3f4311f7fb..438f0b485a 100644 --- a/apps/code/src/renderer/platform-adapters/trpc-pi-runner.ts +++ b/apps/code/src/renderer/platform-adapters/trpc-pi-runner.ts @@ -3,26 +3,27 @@ import type { PiRunInput, PiRunner, } from "@posthog/core/pi-runtime/piRunner"; -import { resolveService } from "@posthog/di/container"; import { HOST_TRPC_CLIENT, type HostTrpcClient, } from "@posthog/host-router/client"; +import { inject, injectable } from "inversify"; -function hostClient(): HostTrpcClient { - return resolveService(HOST_TRPC_CLIENT); -} - +@injectable() export class TrpcPiRunner implements PiRunner { + constructor( + @inject(HOST_TRPC_CLIENT) private readonly hostClient: HostTrpcClient, + ) {} + async create(input: PiRunInput): Promise { - await hostClient().piSession.start.mutate(input); + await this.hostClient.piSession.start.mutate(input); } resume(input: PiResumeInput): Promise { - return hostClient().piSession.resume.mutate(input); + return this.hostClient.piSession.resume.mutate(input); } stop(taskId: string): Promise { - return hostClient().piSession.stop.mutate({ taskId }); + return this.hostClient.piSession.stop.mutate({ taskId }); } } diff --git a/apps/web/src/web-container.ts b/apps/web/src/web-container.ts index 6d15ba6884..9fdfa4e7d9 100644 --- a/apps/web/src/web-container.ts +++ b/apps/web/src/web-container.ts @@ -1,12 +1,18 @@ import "reflect-metadata"; import { TypedContainer } from "@inversifyjs/strongly-typed"; import { taskThreadCoreModule } from "@posthog/core/canvas/taskThread.module"; +import { piRuntimeModule } from "@posthog/core/pi-runtime/pi-runtime.module"; +import { + PI_SESSION_CLIENT, + type PiSessionClient, +} from "@posthog/core/pi-runtime/piSessionController"; import { setRootContainer } from "@posthog/di/container"; import { ROOT_LOGGER, type RootLogger } from "@posthog/di/logger"; import { HOST_TRPC_CLIENT, type HostTrpcClient, } from "@posthog/host-router/client"; +import { TrpcPiSessionClient } from "@posthog/host-router/pi-session-client"; import { sandboxProxyHtml } from "@posthog/shared/mcp-sandbox-proxy"; import { AUTH_SIDE_EFFECTS, @@ -37,6 +43,7 @@ import { hostTrpcClient } from "./web-trpc"; interface WebBindings { [HOST_TRPC_CLIENT]: HostTrpcClient; + [PI_SESSION_CLIENT]: PiSessionClient; [ROOT_LOGGER]: RootLogger; [FEATURE_FLAGS]: FeatureFlags; [ANALYTICS_TRACKER]: AnalyticsTracker; @@ -54,6 +61,8 @@ export const container = new TypedContainer({ // Keystone: the same typed host client the renderer binds, over HTTP not IPC. container.bind(HOST_TRPC_CLIENT).toConstantValue(hostTrpcClient); +container.bind(PI_SESSION_CLIENT).to(TrpcPiSessionClient); +container.load(piRuntimeModule); // Logger: web uses console; electron uses electron-log. Same RootLogger shape. const scoped = (name?: string): RootLogger => ({ diff --git a/packages/agent/package.json b/packages/agent/package.json index 7f06e60866..1f253ca628 100644 --- a/packages/agent/package.json +++ b/packages/agent/package.json @@ -32,6 +32,18 @@ "types": "./dist/pi/rpc-client.d.ts", "import": "./dist/pi/rpc-client.js" }, + "./pi/conversation": { + "types": "./dist/pi/conversation/translatePiConversation.d.ts", + "import": "./dist/pi/conversation/translatePiConversation.js" + }, + "./pi/runtime": { + "types": "./dist/pi/runtime.d.ts", + "import": "./dist/pi/runtime.js" + }, + "./pi/types": { + "types": "./dist/pi/types.d.ts", + "import": "./dist/pi/types.js" + }, "./pr-url-detector": { "types": "./dist/pr-url-detector.d.ts", "import": "./dist/pr-url-detector.js" @@ -116,7 +128,8 @@ "author": "PostHog", "license": "MIT", "scripts": { - "build": "node ../../scripts/rimraf.mjs dist && tsup && node build/verify-local-tools-mcp-server.mjs", + "build": "node ../../scripts/rimraf.mjs dist && tsup && pnpm build:types && node build/verify-local-tools-mcp-server.mjs", + "build:types": "tsc -p tsconfig.build.json", "dev": "tsup --watch", "test": "vitest run", "test:watch": "vitest", @@ -144,6 +157,8 @@ "@agentclientprotocol/sdk": "1.1.0", "@anthropic-ai/claude-agent-sdk": "0.3.197", "@anthropic-ai/sdk": "0.109.0", + "@earendil-works/pi-agent-core": "catalog:", + "@earendil-works/pi-ai": "catalog:", "@earendil-works/pi-coding-agent": "catalog:", "@hono/node-server": "^1.19.9", "@openai/codex": "0.144.0", diff --git a/packages/agent/src/pi/conversation/toolKind.ts b/packages/agent/src/pi/conversation/toolKind.ts new file mode 100644 index 0000000000..8ba14de912 --- /dev/null +++ b/packages/agent/src/pi/conversation/toolKind.ts @@ -0,0 +1,14 @@ +import type { ToolsOptions } from "@earendil-works/pi-coding-agent"; +import type { AgentToolKind } from "@posthog/shared"; + +export type PiToolName = keyof ToolsOptions; + +export const TOOL_KIND_BY_NAME: Record = { + read: "read", + edit: "edit", + write: "edit", + bash: "execute", + grep: "search", + find: "search", + ls: "read", +}; diff --git a/packages/agent/src/pi/conversation/toolTranslator.ts b/packages/agent/src/pi/conversation/toolTranslator.ts new file mode 100644 index 0000000000..be9017193f --- /dev/null +++ b/packages/agent/src/pi/conversation/toolTranslator.ts @@ -0,0 +1,22 @@ +import type { ImageContent, TextContent } from "@earendil-works/pi-ai"; +import type { + AgentToolCallContent, + AgentToolCallLocation, +} from "@posthog/shared"; + +export interface PiToolTranslatorInput { + toolCallId: string; + arguments: unknown; + resultContent?: (TextContent | ImageContent)[]; + details?: unknown; + isError?: boolean; +} + +export interface PiToolTranslatorOutput { + locations?: AgentToolCallLocation[]; + content?: AgentToolCallContent[]; +} + +export type PiToolTranslator = ( + input: PiToolTranslatorInput, +) => PiToolTranslatorOutput; diff --git a/packages/agent/src/pi/conversation/tools/bashTranslator.test.ts b/packages/agent/src/pi/conversation/tools/bashTranslator.test.ts new file mode 100644 index 0000000000..8b58be6381 --- /dev/null +++ b/packages/agent/src/pi/conversation/tools/bashTranslator.test.ts @@ -0,0 +1,83 @@ +import { describe, expect, it } from "vitest"; +import { bashTranslator } from "./bashTranslator"; + +describe("bashTranslator", () => { + it("surfaces command output as text content on success", () => { + const result = bashTranslator({ + toolCallId: "tool-1", + arguments: { command: "echo hi" }, + resultContent: [{ type: "text", text: "hi\n" }], + details: { truncation: undefined, fullOutputPath: undefined }, + isError: false, + }); + + expect(result).toEqual({ + content: [ + { + type: "content", + content: { type: "text", text: "hi\n" }, + }, + ], + }); + }); + + it("surfaces stderr-style output for a failed command", () => { + const result = bashTranslator({ + toolCallId: "tool-2", + arguments: { command: "false" }, + resultContent: [ + { type: "text", text: "command failed with exit code 1" }, + ], + isError: true, + }); + + expect(result).toEqual({ + content: [ + { + type: "content", + content: { type: "text", text: "command failed with exit code 1" }, + }, + ], + }); + }); + + it("returns no content when the result has no text blocks", () => { + const result = bashTranslator({ + toolCallId: "tool-3", + arguments: { command: "echo hi" }, + resultContent: [], + }); + + expect(result).toEqual({}); + }); + + it("joins multiple text blocks and ignores image blocks", () => { + const result = bashTranslator({ + toolCallId: "tool-4", + arguments: { command: "cat file.txt" }, + resultContent: [ + { type: "text", text: "line one\n" }, + { type: "image", data: "abc", mimeType: "image/png" }, + { type: "text", text: "line two\n" }, + ], + }); + + expect(result).toEqual({ + content: [ + { + type: "content", + content: { type: "text", text: "line one\nline two\n" }, + }, + ], + }); + }); + + it("returns no content when resultContent is missing", () => { + const result = bashTranslator({ + toolCallId: "tool-5", + arguments: { command: "echo hi" }, + }); + + expect(result).toEqual({}); + }); +}); diff --git a/packages/agent/src/pi/conversation/tools/bashTranslator.ts b/packages/agent/src/pi/conversation/tools/bashTranslator.ts new file mode 100644 index 0000000000..694dc0de16 --- /dev/null +++ b/packages/agent/src/pi/conversation/tools/bashTranslator.ts @@ -0,0 +1,22 @@ +import type { AgentToolCallContent } from "@posthog/shared"; +import type { PiToolTranslator } from "../toolTranslator"; + +export const bashTranslator: PiToolTranslator = ({ resultContent }) => { + const outputText = resultContent + ?.filter((block) => block.type === "text") + .map((block) => block.text) + .join(""); + + if (!outputText) { + return {}; + } + + const content: AgentToolCallContent[] = [ + { + type: "content", + content: { type: "text", text: outputText }, + }, + ]; + + return { content }; +}; diff --git a/packages/agent/src/pi/conversation/tools/editTranslator.test.ts b/packages/agent/src/pi/conversation/tools/editTranslator.test.ts new file mode 100644 index 0000000000..75f0e4a38f --- /dev/null +++ b/packages/agent/src/pi/conversation/tools/editTranslator.test.ts @@ -0,0 +1,102 @@ +import { describe, expect, it } from "vitest"; +import { editTranslator } from "./editTranslator"; + +describe("editTranslator", () => { + it("produces a diff from the edit arguments and a location for the path", () => { + const output = editTranslator({ + toolCallId: "1", + arguments: { + path: "src/foo.ts", + edits: [{ oldText: "const a = 1;", newText: "const a = 2;" }], + }, + details: { + diff: "- const a = 1;\n+ const a = 2;", + patch: "@@ -1 +1 @@\n-const a = 1;\n+const a = 2;", + firstChangedLine: 3, + }, + }); + + expect(output.locations).toEqual([{ path: "src/foo.ts", line: 3 }]); + expect(output.content).toEqual([ + { + type: "diff", + path: "src/foo.ts", + oldText: "const a = 1;", + newText: "const a = 2;", + }, + ]); + }); + + it("joins multiple edits into a single diff", () => { + const output = editTranslator({ + toolCallId: "1", + arguments: { + path: "src/foo.ts", + edits: [ + { oldText: "const a = 1;", newText: "const a = 2;" }, + { oldText: "const b = 1;", newText: "const b = 2;" }, + ], + }, + details: { + diff: "irrelevant", + patch: "irrelevant", + }, + }); + + expect(output.content).toEqual([ + { + type: "diff", + path: "src/foo.ts", + oldText: "const a = 1;\nconst b = 1;", + newText: "const a = 2;\nconst b = 2;", + }, + ]); + }); + + it("falls back to the details diff string when edits are missing", () => { + const output = editTranslator({ + toolCallId: "1", + arguments: { path: "src/foo.ts" }, + details: { + diff: "- const a = 1;\n+ const a = 2;", + patch: "@@ -1 +1 @@", + }, + }); + + expect(output.locations).toEqual([{ path: "src/foo.ts", line: undefined }]); + expect(output.content).toEqual([ + { + type: "diff", + path: "src/foo.ts", + newText: "- const a = 1;\n+ const a = 2;", + }, + ]); + }); + + it("falls back to result text content on error with no details", () => { + const output = editTranslator({ + toolCallId: "1", + arguments: { path: "src/foo.ts" }, + resultContent: [{ type: "text", text: "permission denied" }], + isError: true, + }); + + expect(output.locations).toEqual([{ path: "src/foo.ts", line: undefined }]); + expect(output.content).toEqual([ + { + type: "content", + content: { type: "text", text: "permission denied" }, + }, + ]); + }); + + it("returns no locations or content when arguments are missing entirely", () => { + const output = editTranslator({ + toolCallId: "1", + arguments: undefined, + }); + + expect(output.locations).toBeUndefined(); + expect(output.content).toBeUndefined(); + }); +}); diff --git a/packages/agent/src/pi/conversation/tools/editTranslator.ts b/packages/agent/src/pi/conversation/tools/editTranslator.ts new file mode 100644 index 0000000000..1b717624d0 --- /dev/null +++ b/packages/agent/src/pi/conversation/tools/editTranslator.ts @@ -0,0 +1,55 @@ +import type { + EditToolDetails, + EditToolInput, +} from "@earendil-works/pi-coding-agent"; +import type { + AgentToolCallContent, + AgentToolCallLocation, +} from "@posthog/shared"; +import type { PiToolTranslator } from "../toolTranslator"; + +export const editTranslator: PiToolTranslator = ({ + arguments: rawArguments, + resultContent, + details: rawDetails, +}) => { + const args = rawArguments as EditToolInput | undefined; + const details = rawDetails as EditToolDetails | undefined; + + const locations: AgentToolCallLocation[] | undefined = args?.path + ? [{ path: args.path, line: details?.firstChangedLine }] + : undefined; + + if (args?.path && args.edits && args.edits.length > 0) { + const diff: AgentToolCallContent = { + type: "diff", + path: args.path, + oldText: args.edits.map((edit) => edit.oldText).join("\n"), + newText: args.edits.map((edit) => edit.newText).join("\n"), + }; + + return { locations, content: [diff] }; + } + + if (args?.path && details?.diff) { + const diff: AgentToolCallContent = { + type: "diff", + path: args.path, + newText: details.diff, + }; + + return { locations, content: [diff] }; + } + + const textContent = resultContent + ?.filter((block) => block.type === "text") + .map((block) => ({ + type: "content" as const, + content: { type: "text" as const, text: block.text }, + })); + + return { + locations, + content: textContent && textContent.length > 0 ? textContent : undefined, + }; +}; diff --git a/packages/agent/src/pi/conversation/tools/findTranslator.test.ts b/packages/agent/src/pi/conversation/tools/findTranslator.test.ts new file mode 100644 index 0000000000..f862ae05c2 --- /dev/null +++ b/packages/agent/src/pi/conversation/tools/findTranslator.test.ts @@ -0,0 +1,51 @@ +import { describe, expect, it } from "vitest"; +import { findTranslator } from "./findTranslator"; + +describe("findTranslator", () => { + it("emits a location from the search path and text content from results", () => { + const output = findTranslator({ + toolCallId: "call-1", + arguments: { pattern: "*.ts", path: "packages/core/src" }, + resultContent: [ + { + type: "text", + text: "packages/core/src/foo.ts\npackages/core/src/bar.ts", + }, + ], + }); + + expect(output.locations).toEqual([{ path: "packages/core/src" }]); + expect(output.content).toEqual([ + { + type: "content", + content: { + type: "text", + text: "packages/core/src/foo.ts\npackages/core/src/bar.ts", + }, + }, + ]); + }); + + it("omits locations and content when path and resultContent are absent", () => { + const output = findTranslator({ + toolCallId: "call-2", + arguments: { pattern: "*.ts" }, + }); + + expect(output.locations).toBeUndefined(); + expect(output.content).toBeUndefined(); + }); + + it("returns no content when result blocks are all images", () => { + const output = findTranslator({ + toolCallId: "call-3", + arguments: { pattern: "*.png", path: "assets" }, + resultContent: [ + { type: "image", data: "base64data", mimeType: "image/png" }, + ], + }); + + expect(output.locations).toEqual([{ path: "assets" }]); + expect(output.content).toBeUndefined(); + }); +}); diff --git a/packages/agent/src/pi/conversation/tools/findTranslator.ts b/packages/agent/src/pi/conversation/tools/findTranslator.ts new file mode 100644 index 0000000000..cfa526b69e --- /dev/null +++ b/packages/agent/src/pi/conversation/tools/findTranslator.ts @@ -0,0 +1,26 @@ +import type { FindToolInput } from "@earendil-works/pi-coding-agent"; +import type { AgentToolCallLocation } from "@posthog/shared"; +import type { PiToolTranslator } from "../toolTranslator"; + +export const findTranslator: PiToolTranslator = ({ + arguments: rawArguments, + resultContent, +}) => { + const args = rawArguments as FindToolInput | undefined; + + const locations: AgentToolCallLocation[] | undefined = args?.path + ? [{ path: args.path }] + : undefined; + + const content = resultContent + ?.filter((block) => block.type === "text") + .map((block) => ({ + type: "content" as const, + content: { type: "text" as const, text: block.text }, + })); + + return { + locations, + content: content && content.length > 0 ? content : undefined, + }; +}; diff --git a/packages/agent/src/pi/conversation/tools/grepTranslator.test.ts b/packages/agent/src/pi/conversation/tools/grepTranslator.test.ts new file mode 100644 index 0000000000..820f9bcae7 --- /dev/null +++ b/packages/agent/src/pi/conversation/tools/grepTranslator.test.ts @@ -0,0 +1,59 @@ +import { describe, expect, it } from "vitest"; +import { grepTranslator } from "./grepTranslator"; + +describe("grepTranslator", () => { + it("returns locations and text content on success", () => { + const result = grepTranslator({ + toolCallId: "call-1", + arguments: { pattern: "foo", path: "src" }, + resultContent: [ + { type: "text", text: "src/a.ts:1:foo\nsrc/b.ts:2:foo bar" }, + ], + details: undefined, + isError: false, + }); + + expect(result.locations).toEqual([{ path: "src" }]); + expect(result.content).toEqual([ + { + type: "content", + content: { type: "text", text: "src/a.ts:1:foo\nsrc/b.ts:2:foo bar" }, + }, + ]); + }); + + it("appends truncation notes and omits locations when path is missing", () => { + const result = grepTranslator({ + toolCallId: "call-2", + arguments: { pattern: "foo" }, + resultContent: [{ type: "text", text: "match" }], + details: { matchLimitReached: 100, linesTruncated: true }, + isError: false, + }); + + expect(result.locations).toBeUndefined(); + expect(result.content).toEqual([ + { type: "content", content: { type: "text", text: "match" } }, + { + type: "content", + content: { + type: "text", + text: "Match limit reached at 100 matches. Some lines were truncated.", + }, + }, + ]); + }); + + it("returns undefined content and locations when there is nothing to report", () => { + const result = grepTranslator({ + toolCallId: "call-3", + arguments: { pattern: "foo" }, + resultContent: [], + details: undefined, + isError: true, + }); + + expect(result.locations).toBeUndefined(); + expect(result.content).toBeUndefined(); + }); +}); diff --git a/packages/agent/src/pi/conversation/tools/grepTranslator.ts b/packages/agent/src/pi/conversation/tools/grepTranslator.ts new file mode 100644 index 0000000000..0686d7bfb4 --- /dev/null +++ b/packages/agent/src/pi/conversation/tools/grepTranslator.ts @@ -0,0 +1,55 @@ +import type { + GrepToolDetails, + GrepToolInput, +} from "@earendil-works/pi-coding-agent"; +import type { + AgentToolCallContentBlock, + AgentToolCallLocation, +} from "@posthog/shared"; +import type { PiToolTranslator } from "../toolTranslator"; + +export const grepTranslator: PiToolTranslator = ({ + arguments: rawArguments, + resultContent, + details: rawDetails, +}) => { + const args = rawArguments as GrepToolInput; + const details = rawDetails as GrepToolDetails | undefined; + + const locations: AgentToolCallLocation[] = []; + if (args?.path) { + locations.push({ path: args.path }); + } + + const content: AgentToolCallContentBlock[] = []; + const resultText = (resultContent ?? []) + .filter((block) => block.type === "text") + .map((block) => block.text) + .join("\n"); + + if (resultText) { + content.push({ + type: "content", + content: { type: "text", text: resultText }, + }); + } + + const notes: string[] = []; + if (details?.matchLimitReached !== undefined) { + notes.push(`Match limit reached at ${details.matchLimitReached} matches.`); + } + if (details?.linesTruncated) { + notes.push("Some lines were truncated."); + } + if (notes.length > 0) { + content.push({ + type: "content", + content: { type: "text", text: notes.join(" ") }, + }); + } + + return { + locations: locations.length > 0 ? locations : undefined, + content: content.length > 0 ? content : undefined, + }; +}; diff --git a/packages/agent/src/pi/conversation/tools/lsTranslator.test.ts b/packages/agent/src/pi/conversation/tools/lsTranslator.test.ts new file mode 100644 index 0000000000..bfe71d3a21 --- /dev/null +++ b/packages/agent/src/pi/conversation/tools/lsTranslator.test.ts @@ -0,0 +1,57 @@ +import { describe, expect, it } from "vitest"; +import { lsTranslator } from "./lsTranslator"; + +describe("lsTranslator", () => { + it("returns a location for the listed path and text content from the result", () => { + const output = lsTranslator({ + toolCallId: "call-1", + arguments: { path: "src/features" }, + resultContent: [{ type: "text", text: "sessions/\ntools/" }], + }); + + expect(output.locations).toEqual([{ path: "src/features" }]); + expect(output.content).toEqual([ + { type: "content", content: { type: "text", text: "sessions/\ntools/" } }, + ]); + }); + + it("omits locations when no path is given and content when result is empty", () => { + const output = lsTranslator({ + toolCallId: "call-2", + arguments: {}, + resultContent: [], + }); + + expect(output.locations).toBeUndefined(); + expect(output.content).toBeUndefined(); + }); + + it("still surfaces content on an error result", () => { + const output = lsTranslator({ + toolCallId: "call-3", + arguments: { path: "missing-dir" }, + resultContent: [{ type: "text", text: "ENOENT: no such directory" }], + isError: true, + }); + + expect(output.locations).toEqual([{ path: "missing-dir" }]); + expect(output.content).toEqual([ + { + type: "content", + content: { type: "text", text: "ENOENT: no such directory" }, + }, + ]); + }); + + it("ignores image content blocks", () => { + const output = lsTranslator({ + toolCallId: "call-4", + arguments: { path: "src" }, + resultContent: [ + { type: "image", data: "base64data", mimeType: "image/png" }, + ], + }); + + expect(output.content).toBeUndefined(); + }); +}); diff --git a/packages/agent/src/pi/conversation/tools/lsTranslator.ts b/packages/agent/src/pi/conversation/tools/lsTranslator.ts new file mode 100644 index 0000000000..d50893580e --- /dev/null +++ b/packages/agent/src/pi/conversation/tools/lsTranslator.ts @@ -0,0 +1,33 @@ +import type { LsToolInput } from "@earendil-works/pi-coding-agent"; +import type { + AgentToolCallContent, + AgentToolCallLocation, +} from "@posthog/shared"; +import type { PiToolTranslator } from "../toolTranslator"; + +export const lsTranslator: PiToolTranslator = ({ + arguments: args, + resultContent, +}) => { + const input = args as LsToolInput | undefined; + + const locations: AgentToolCallLocation[] = []; + if (input?.path) { + locations.push({ path: input.path }); + } + + const content: AgentToolCallContent[] = []; + for (const block of resultContent ?? []) { + if (block.type === "text") { + content.push({ + type: "content", + content: { type: "text", text: block.text }, + }); + } + } + + return { + locations: locations.length > 0 ? locations : undefined, + content: content.length > 0 ? content : undefined, + }; +}; diff --git a/packages/agent/src/pi/conversation/tools/readTranslator.test.ts b/packages/agent/src/pi/conversation/tools/readTranslator.test.ts new file mode 100644 index 0000000000..37e7398540 --- /dev/null +++ b/packages/agent/src/pi/conversation/tools/readTranslator.test.ts @@ -0,0 +1,94 @@ +import type { + ReadToolDetails, + ReadToolInput, +} from "@earendil-works/pi-coding-agent"; +import { describe, expect, it } from "vitest"; +import { readTranslator } from "./readTranslator"; + +describe("readTranslator", () => { + it("returns the read path as a location and the file text as content", () => { + const args: ReadToolInput = { path: "src/index.ts" }; + + const output = readTranslator({ + toolCallId: "1", + arguments: args, + resultContent: [{ type: "text", text: "export const x = 1;" }], + isError: false, + }); + + expect(output).toEqual({ + locations: [{ path: "src/index.ts" }], + content: [ + { + type: "content", + content: { type: "text", text: "export const x = 1;" }, + }, + ], + }); + }); + + it("appends truncation info when the read result was truncated", () => { + const args: ReadToolInput = { path: "src/big.ts", offset: 0, limit: 100 }; + const details: ReadToolDetails = { + truncation: { + content: "line1\nline2", + truncated: true, + truncatedBy: "lines", + totalLines: 5000, + totalBytes: 100000, + outputLines: 100, + outputBytes: 2000, + lastLinePartial: false, + firstLineExceedsLimit: false, + maxLines: 2000, + maxBytes: 50000, + }, + }; + + const output = readTranslator({ + toolCallId: "2", + arguments: args, + resultContent: [{ type: "text", text: "line1\nline2" }], + details, + isError: false, + }); + + expect(output.locations).toEqual([{ path: "src/big.ts" }]); + expect(output.content).toHaveLength(1); + const [content] = output.content ?? []; + expect(content?.type).toBe("content"); + if (content?.type === "content" && content.content.type === "text") { + expect(content.content.text).toContain("line1\nline2"); + expect(content.content.text).toContain( + "truncated: showing 100 of 5000 lines", + ); + } + }); + + it("omits content when the result has no text block, e.g. an error result", () => { + const args: ReadToolInput = { path: "src/missing.ts" }; + + const output = readTranslator({ + toolCallId: "3", + arguments: args, + resultContent: undefined, + isError: true, + }); + + expect(output).toEqual({ + locations: [{ path: "src/missing.ts" }], + content: undefined, + }); + }); + + it("omits locations when arguments are missing", () => { + const output = readTranslator({ + toolCallId: "4", + arguments: undefined, + resultContent: [{ type: "text", text: "content" }], + isError: false, + }); + + expect(output.locations).toBeUndefined(); + }); +}); diff --git a/packages/agent/src/pi/conversation/tools/readTranslator.ts b/packages/agent/src/pi/conversation/tools/readTranslator.ts new file mode 100644 index 0000000000..4086394c94 --- /dev/null +++ b/packages/agent/src/pi/conversation/tools/readTranslator.ts @@ -0,0 +1,35 @@ +import type { + ReadToolDetails, + ReadToolInput, +} from "@earendil-works/pi-coding-agent"; +import type { AgentToolCallContent } from "@posthog/shared"; +import type { PiToolTranslator } from "../toolTranslator"; + +export const readTranslator: PiToolTranslator = ({ + arguments: rawArguments, + resultContent, + details: rawDetails, +}) => { + const args = rawArguments as ReadToolInput | undefined; + const details = rawDetails as ReadToolDetails | undefined; + + const locations = args?.path ? [{ path: args.path }] : undefined; + + const textBlock = resultContent?.find((block) => block.type === "text"); + const content: AgentToolCallContent[] = []; + + if (textBlock && textBlock.type === "text") { + let text = textBlock.text; + + if (details?.truncation?.truncated) { + text = `${text}\n[truncated: showing ${details.truncation.outputLines} of ${details.truncation.totalLines} lines]`; + } + + content.push({ type: "content", content: { type: "text", text } }); + } + + return { + locations, + content: content.length > 0 ? content : undefined, + }; +}; diff --git a/packages/agent/src/pi/conversation/tools/writeTranslator.test.ts b/packages/agent/src/pi/conversation/tools/writeTranslator.test.ts new file mode 100644 index 0000000000..caffa7403e --- /dev/null +++ b/packages/agent/src/pi/conversation/tools/writeTranslator.test.ts @@ -0,0 +1,40 @@ +import { describe, expect, it } from "vitest"; +import { writeTranslator } from "./writeTranslator"; + +describe("writeTranslator", () => { + it("produces a diff content block and location for the written file", () => { + const result = writeTranslator({ + toolCallId: "call-1", + arguments: { path: "src/foo.ts", content: "export const foo = 1;\n" }, + }); + + expect(result.locations).toEqual([{ path: "src/foo.ts" }]); + expect(result.content).toEqual([ + { + type: "diff", + path: "src/foo.ts", + oldText: null, + newText: "export const foo = 1;\n", + }, + ]); + }); + + it("still produces a diff when the result is an error", () => { + const result = writeTranslator({ + toolCallId: "call-2", + arguments: { path: "src/bar.ts", content: "" }, + isError: true, + resultContent: [{ type: "text", text: "permission denied" }], + }); + + expect(result.locations).toEqual([{ path: "src/bar.ts" }]); + expect(result.content).toEqual([ + { + type: "diff", + path: "src/bar.ts", + oldText: null, + newText: "", + }, + ]); + }); +}); diff --git a/packages/agent/src/pi/conversation/tools/writeTranslator.ts b/packages/agent/src/pi/conversation/tools/writeTranslator.ts new file mode 100644 index 0000000000..7e2e4b4bf8 --- /dev/null +++ b/packages/agent/src/pi/conversation/tools/writeTranslator.ts @@ -0,0 +1,18 @@ +import type { WriteToolInput } from "@earendil-works/pi-coding-agent"; +import type { PiToolTranslator } from "../toolTranslator"; + +export const writeTranslator: PiToolTranslator = ({ arguments: args }) => { + const input = args as WriteToolInput; + + return { + locations: [{ path: input.path }], + content: [ + { + type: "diff", + path: input.path, + oldText: null, + newText: input.content, + }, + ], + }; +}; diff --git a/packages/agent/src/pi/conversation/translatePiConversation.test.ts b/packages/agent/src/pi/conversation/translatePiConversation.test.ts new file mode 100644 index 0000000000..a55458e0f8 --- /dev/null +++ b/packages/agent/src/pi/conversation/translatePiConversation.test.ts @@ -0,0 +1,434 @@ +import type { AssistantMessage } from "@earendil-works/pi-ai"; +import { describe, expect, it } from "vitest"; +import { createPiConversationTranslator } from "./translatePiConversation"; + +function assistant( + content: AssistantMessage["content"], + stopReason: AssistantMessage["stopReason"] = "stop", + timestamp = 10, +): AssistantMessage { + return { + role: "assistant", + content, + api: "anthropic-messages" as AssistantMessage["api"], + provider: "anthropic" as AssistantMessage["provider"], + model: "test-model", + usage: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 0, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + stopReason, + timestamp, + }; +} + +describe("createPiConversationTranslator", () => { + it("keeps complete assistant content when translating history", () => { + const translator = createPiConversationTranslator(); + + expect( + translator.translateHistoryMessage( + assistant([{ type: "text", text: "complete" }]), + ), + ).toContainEqual({ + type: "assistant_message_chunk", + timestamp: 10, + content: { type: "text", text: "complete" }, + }); + }); + + it("uses message_update deltas without repeating cumulative text at message_end", () => { + const translator = createPiConversationTranslator(); + const message = assistant([{ type: "text", text: "complete" }]); + + const streamed = translator.translateEvent({ + type: "message_update", + message, + assistantMessageEvent: { + type: "text_delta", + contentIndex: 0, + delta: "complete", + partial: message, + }, + }); + const ended = translator.translateEvent({ type: "message_end", message }); + + expect(streamed).toEqual([ + { + type: "assistant_message_chunk", + timestamp: 10, + content: { type: "text", text: "complete" }, + }, + ]); + expect(ended).toEqual([]); + }); + + it("does not repeat streamed content when assistant timestamps collide", () => { + const translator = createPiConversationTranslator(); + const first = assistant([{ type: "text", text: "first" }]); + const second = assistant([{ type: "text", text: "second" }]); + + translator.translateEvent({ + type: "message_update", + message: first, + assistantMessageEvent: { + type: "text_delta", + contentIndex: 0, + delta: "first", + partial: first, + }, + }); + translator.translateEvent({ type: "message_end", message: first }); + translator.translateEvent({ + type: "message_update", + message: second, + assistantMessageEvent: { + type: "text_delta", + contentIndex: 0, + delta: "second", + partial: second, + }, + }); + + expect( + translator.translateEvent({ type: "message_end", message: second }), + ).toEqual([]); + }); + + it("completes a turn using the latest runtime timestamp", () => { + const translator = createPiConversationTranslator(); + const laterMessage = assistant( + [{ type: "text", text: "later" }], + "stop", + 20, + ); + const earlierMessage = assistant( + [{ type: "text", text: "earlier" }], + "stop", + 10, + ); + + translator.translateEvent({ type: "message_end", message: laterMessage }); + translator.translateEvent({ type: "message_end", message: earlierMessage }); + + expect(translator.translateEvent({ type: "agent_settled" })).toEqual([ + { type: "turn_completed", timestamp: 20 }, + ]); + }); + + it("translates retry lifecycle without rendering transient runtime errors", () => { + const translator = createPiConversationTranslator(); + const failedMessage = { + ...assistant([], "error"), + errorMessage: "Rate limited", + }; + + expect( + translator.translateEvent({ + type: "message_end", + message: failedMessage, + }), + ).toEqual([]); + expect( + translator.translateEvent({ + type: "agent_end", + messages: [failedMessage], + willRetry: true, + }), + ).toEqual([]); + expect( + translator.translateEvent({ + type: "auto_retry_start", + attempt: 1, + maxAttempts: 3, + delayMs: 1000, + errorMessage: "Rate limited", + }), + ).toEqual([ + { + type: "runtime_status", + timestamp: 10, + status: "retrying", + message: "Rate limited", + attempt: 1, + maxAttempts: 3, + delayMs: 1000, + }, + ]); + expect( + translator.translateEvent({ + type: "auto_retry_end", + success: true, + attempt: 1, + }), + ).toEqual([ + { + type: "runtime_status", + timestamp: 10, + status: "retrying", + isComplete: true, + }, + ]); + }); + + it("renders terminal Pi runtime errors inline", () => { + const translator = createPiConversationTranslator(); + const failedMessage = { + ...assistant([], "error"), + errorMessage: "Authentication failed", + }; + + translator.translateEvent({ type: "message_end", message: failedMessage }); + + expect( + translator.translateEvent({ + type: "agent_end", + messages: [failedMessage], + willRetry: false, + }), + ).toEqual([ + { + type: "runtime_error", + timestamp: 10, + errorType: "pi_runtime", + message: "Authentication failed", + }, + ]); + }); + + it("translates compaction lifecycle into generic runtime statuses", () => { + const translator = createPiConversationTranslator(); + translator.translateHistoryMessage( + assistant([{ type: "text", text: "complete" }]), + ); + + expect( + translator.translateEvent({ + type: "compaction_start", + reason: "manual", + }), + ).toEqual([ + { + type: "runtime_status", + timestamp: 10, + status: "compacting", + }, + ]); + + expect( + translator.translateEvent({ + type: "compaction_end", + reason: "manual", + result: undefined, + aborted: false, + willRetry: false, + }), + ).toEqual([ + { + type: "runtime_status", + timestamp: 10, + status: "compacting", + isComplete: true, + }, + ]); + }); + + it("translates compaction failures with their error", () => { + const translator = createPiConversationTranslator(); + + expect( + translator.translateEvent({ + type: "compaction_end", + reason: "manual", + result: undefined, + aborted: false, + willRetry: false, + errorMessage: "Not enough messages", + }), + ).toEqual([ + { + type: "runtime_status", + timestamp: 0, + status: "compacting_failed", + error: "Not enough messages", + }, + ]); + }); + + it("translates direct bash history into the generic execute tool UI", () => { + const translator = createPiConversationTranslator(); + + expect( + translator.translateHistoryMessage({ + role: "bashExecution", + command: "pwd", + output: "/tmp/project", + exitCode: 0, + cancelled: false, + truncated: false, + timestamp: 20, + }), + ).toEqual([ + { + type: "tool_call_started", + timestamp: 20, + toolCall: { + id: "pi-bash-20", + title: "pwd", + kind: "execute", + status: "in_progress", + rawInput: { command: "pwd" }, + }, + }, + { + type: "tool_call_updated", + timestamp: 20, + toolCall: { + id: "pi-bash-20", + status: "completed", + rawOutput: "/tmp/project", + content: [ + { + type: "content", + content: { type: "text", text: "/tmp/project" }, + }, + ], + }, + }, + ]); + }); + + it("streams tool execution start, output updates, and completion", () => { + const translator = createPiConversationTranslator(); + const message = assistant( + [ + { + type: "toolCall", + id: "tool-1", + name: "bash", + arguments: { command: "printf hello" }, + }, + ], + "toolUse", + ); + + translator.translateEvent({ type: "message_end", message }); + + expect( + translator.translateEvent({ + type: "tool_execution_start", + toolCallId: "tool-1", + toolName: "bash", + args: { command: "printf hello" }, + }), + ).toEqual([ + { + type: "tool_call_updated", + timestamp: 10, + toolCall: { id: "tool-1", status: "in_progress" }, + }, + ]); + + expect( + translator.translateEvent({ + type: "tool_execution_update", + toolCallId: "tool-1", + toolName: "bash", + args: { command: "printf hello" }, + partialResult: { + content: [{ type: "text", text: "hel" }], + details: undefined, + }, + }), + ).toEqual([ + { + type: "tool_call_updated", + timestamp: 10, + toolCall: { + id: "tool-1", + status: "in_progress", + rawOutput: [{ type: "text", text: "hel" }], + content: [ + { + type: "content", + content: { type: "text", text: "hel" }, + }, + ], + }, + }, + ]); + + expect( + translator.translateEvent({ + type: "tool_execution_end", + toolCallId: "tool-1", + toolName: "bash", + result: { + content: [{ type: "text", text: "hello" }], + details: undefined, + }, + isError: false, + }), + ).toEqual([ + { + type: "tool_call_updated", + timestamp: 10, + toolCall: { + id: "tool-1", + status: "completed", + rawOutput: [{ type: "text", text: "hello" }], + content: [ + { + type: "content", + content: { type: "text", text: "hello" }, + }, + ], + }, + }, + ]); + }); + + it("preserves tool calls when filtering streamed assistant content", () => { + const translator = createPiConversationTranslator(); + const message = assistant([ + { type: "text", text: "running" }, + { + type: "toolCall", + id: "tool-1", + name: "bash", + arguments: { command: "pwd" }, + }, + ]); + + translator.translateEvent({ + type: "message_update", + message, + assistantMessageEvent: { + type: "text_delta", + contentIndex: 0, + delta: "running", + partial: message, + }, + }); + + expect(translator.translateEvent({ type: "message_end", message })).toEqual( + [ + { + type: "tool_call_started", + timestamp: 10, + toolCall: { + id: "tool-1", + title: "bash", + kind: "execute", + status: "pending", + rawInput: { command: "pwd" }, + }, + }, + ], + ); + }); +}); diff --git a/packages/agent/src/pi/conversation/translatePiConversation.ts b/packages/agent/src/pi/conversation/translatePiConversation.ts new file mode 100644 index 0000000000..b02cb82c58 --- /dev/null +++ b/packages/agent/src/pi/conversation/translatePiConversation.ts @@ -0,0 +1,348 @@ +import type { AssistantMessage, Message } from "@earendil-works/pi-ai"; +import type { AgentSessionEvent } from "@earendil-works/pi-coding-agent"; +import type { AgentConversationEvent } from "@posthog/shared"; +import { createPiMessageTranslator } from "./translatePiMessage"; + +type AgentMessage = Extract< + AgentSessionEvent, + { type: "message_end" } +>["message"]; + +function isMessage(message: AgentMessage): message is Message { + return ( + message.role === "user" || + message.role === "assistant" || + message.role === "toolResult" + ); +} + +function customMessageEvents(message: AgentMessage): AgentConversationEvent[] { + if (message.role === "bashExecution") { + const id = `pi-bash-${message.timestamp}`; + const failed = message.cancelled || (message.exitCode ?? 0) !== 0; + + return [ + { + type: "tool_call_started", + timestamp: message.timestamp, + toolCall: { + id, + title: message.command, + kind: "execute", + status: "in_progress", + rawInput: { command: message.command }, + }, + }, + { + type: "tool_call_updated", + timestamp: message.timestamp, + toolCall: { + id, + status: failed ? "failed" : "completed", + rawOutput: message.output, + content: message.output + ? [ + { + type: "content", + content: { type: "text", text: message.output }, + }, + ] + : [], + }, + }, + ]; + } + + let text: string | undefined; + + if ( + message.role === "branchSummary" || + message.role === "compactionSummary" + ) { + text = message.summary; + } else if (message.role === "custom" && message.display) { + text = + typeof message.content === "string" + ? message.content + : message.content + .flatMap((content) => + content.type === "text" ? [content.text] : [], + ) + .join("\n"); + } + + if (!text) { + return []; + } + + return [ + { + type: "assistant_message_chunk", + timestamp: message.timestamp, + content: { type: "text", text }, + }, + ]; +} + +function isAssistantMessage( + message: AgentMessage, +): message is AssistantMessage { + return message.role === "assistant"; +} + +export interface PiConversationTranslator { + translateHistoryMessage(message: AgentMessage): AgentConversationEvent[]; + translateEvent(event: AgentSessionEvent): AgentConversationEvent[]; +} + +export function createPiConversationTranslator(): PiConversationTranslator { + const messageTranslator = createPiMessageTranslator(); + const streamedAssistantTimestamps = new Set(); + let historyTurnActive = false; + let latestRuntimeTimestamp = 0; + let latestConversationTimestamp = 0; + let pendingRuntimeError: AgentConversationEvent | undefined; + + function translateHistoryMessage( + message: AgentMessage, + ): AgentConversationEvent[] { + const events: AgentConversationEvent[] = []; + latestConversationTimestamp = Math.max( + latestConversationTimestamp, + message.timestamp, + ); + + if (message.role === "user" && historyTurnActive) { + events.push({ + type: "turn_completed", + timestamp: message.timestamp, + }); + historyTurnActive = false; + } + + if (isMessage(message)) { + events.push(...messageTranslator.translate(message)); + } else { + events.push(...customMessageEvents(message)); + } + + if (message.role === "user") { + historyTurnActive = true; + } + + if ( + isAssistantMessage(message) && + message.stopReason !== "toolUse" && + historyTurnActive + ) { + events.push({ + type: "turn_completed", + timestamp: message.timestamp, + stopReason: message.stopReason, + }); + historyTurnActive = false; + } + + return events; + } + + function translateEvent(event: AgentSessionEvent): AgentConversationEvent[] { + if (event.type === "message_update") { + const update = event.assistantMessageEvent; + latestRuntimeTimestamp = Math.max( + latestRuntimeTimestamp, + event.message.timestamp, + ); + latestConversationTimestamp = Math.max( + latestConversationTimestamp, + event.message.timestamp, + ); + + if (update.type === "text_delta" && update.delta) { + streamedAssistantTimestamps.add(event.message.timestamp); + return [ + { + type: "assistant_message_chunk", + timestamp: event.message.timestamp, + content: { type: "text", text: update.delta }, + }, + ]; + } + + if (update.type === "thinking_delta" && update.delta) { + streamedAssistantTimestamps.add(event.message.timestamp); + return [ + { + type: "assistant_thought_chunk", + timestamp: event.message.timestamp, + content: { type: "text", text: update.delta }, + }, + ]; + } + + return []; + } + + if (event.type === "tool_execution_start") { + return messageTranslator.translateToolExecutionStart( + event.toolCallId, + event.toolName, + event.args, + latestRuntimeTimestamp, + ); + } + + if (event.type === "tool_execution_update") { + return messageTranslator.translateToolExecutionUpdate( + event.toolCallId, + event.toolName, + event.args, + event.partialResult, + latestRuntimeTimestamp, + ); + } + + if (event.type === "tool_execution_end") { + return messageTranslator.translateToolExecutionEnd( + event.toolCallId, + event.toolName, + event.result, + event.isError, + latestRuntimeTimestamp, + ); + } + + if (event.type === "message_end") { + latestRuntimeTimestamp = Math.max( + latestRuntimeTimestamp, + event.message.timestamp, + ); + latestConversationTimestamp = Math.max( + latestConversationTimestamp, + event.message.timestamp, + ); + + if (!isMessage(event.message)) { + return customMessageEvents(event.message); + } + + const events = messageTranslator.translate(event.message); + const runtimeError = events.find( + (translated) => translated.type === "runtime_error", + ); + if (runtimeError) { + pendingRuntimeError = runtimeError; + } + + const visibleEvents = events.filter( + (translated) => translated.type !== "runtime_error", + ); + if ( + event.message.role !== "assistant" || + !streamedAssistantTimestamps.has(event.message.timestamp) + ) { + return visibleEvents; + } + + return visibleEvents.filter( + (translated) => + translated.type !== "assistant_message_chunk" && + translated.type !== "assistant_thought_chunk", + ); + } + + if (event.type === "agent_end") { + const runtimeError = pendingRuntimeError; + pendingRuntimeError = undefined; + + if (!event.willRetry && runtimeError) { + return [runtimeError]; + } + + return []; + } + + if (event.type === "auto_retry_start") { + return [ + { + type: "runtime_status", + timestamp: latestConversationTimestamp, + status: "retrying", + message: event.errorMessage, + attempt: event.attempt, + maxAttempts: event.maxAttempts, + delayMs: event.delayMs, + }, + ]; + } + + if (event.type === "auto_retry_end") { + const events: AgentConversationEvent[] = [ + { + type: "runtime_status", + timestamp: latestConversationTimestamp, + status: "retrying", + isComplete: true, + }, + ]; + + if (!event.success && event.finalError) { + events.push({ + type: "runtime_error", + timestamp: latestConversationTimestamp, + errorType: "pi_runtime", + message: event.finalError, + }); + } + + return events; + } + + if (event.type === "compaction_start") { + return [ + { + type: "runtime_status", + timestamp: latestConversationTimestamp, + status: "compacting", + }, + ]; + } + + if (event.type === "compaction_end") { + if (event.aborted || event.errorMessage) { + return [ + { + type: "runtime_status", + timestamp: latestConversationTimestamp, + status: "compacting_failed", + error: + event.errorMessage ?? + (event.aborted ? "Compaction cancelled" : undefined), + }, + ]; + } + + return [ + { + type: "runtime_status", + timestamp: latestConversationTimestamp, + status: "compacting", + isComplete: true, + }, + ]; + } + + if (event.type === "agent_settled") { + streamedAssistantTimestamps.clear(); + + const timestamp = latestRuntimeTimestamp; + latestRuntimeTimestamp = 0; + + return timestamp > 0 ? [{ type: "turn_completed", timestamp }] : []; + } + + return []; + } + + return { translateHistoryMessage, translateEvent }; +} diff --git a/packages/agent/src/pi/conversation/translatePiMessage.test.ts b/packages/agent/src/pi/conversation/translatePiMessage.test.ts new file mode 100644 index 0000000000..0393331230 --- /dev/null +++ b/packages/agent/src/pi/conversation/translatePiMessage.test.ts @@ -0,0 +1,113 @@ +import type { AssistantMessage, UserMessage } from "@earendil-works/pi-ai"; +import { describe, expect, it } from "vitest"; +import { createPiMessageTranslator } from "./translatePiMessage"; + +function makeAssistant(content: AssistantMessage["content"]): AssistantMessage { + return { + role: "assistant", + content, + api: "anthropic-messages" as AssistantMessage["api"], + provider: "anthropic" as AssistantMessage["provider"], + model: "test-model", + usage: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 0, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + stopReason: "stop", + timestamp: 0, + }; +} + +describe("createPiMessageTranslator", () => { + it("translates a string user message into a user message chunk", () => { + const translator = createPiMessageTranslator(); + + const message: UserMessage = { + role: "user", + content: "hello there", + timestamp: 0, + }; + + expect(translator.translate(message)).toEqual([ + { + type: "user_message", + id: "pi-user-0-1", + timestamp: 0, + content: [{ type: "text", text: "hello there" }], + }, + ]); + }); + + it("translates user message content blocks into user message chunks", () => { + const translator = createPiMessageTranslator(); + + const message: UserMessage = { + role: "user", + content: [ + { type: "text", text: "first" }, + { type: "image", data: "abc", mimeType: "image/png" }, + ], + timestamp: 0, + }; + + expect(translator.translate(message)).toEqual([ + { + type: "user_message", + id: "pi-user-0-1", + timestamp: 0, + content: [ + { type: "text", text: "first" }, + { type: "image", data: "abc", mimeType: "image/png" }, + ], + }, + ]); + }); + + it("translates a plain text assistant message into an agent message chunk", () => { + const translator = createPiMessageTranslator(); + + const message = makeAssistant([{ type: "text", text: "working on it" }]); + + expect(translator.translate(message)).toEqual([ + { + type: "assistant_message_chunk", + timestamp: 0, + content: { type: "text", text: "working on it" }, + }, + ]); + }); + + it("translates assistant runtime errors", () => { + const translator = createPiMessageTranslator(); + const message = makeAssistant([]); + message.stopReason = "error"; + message.errorMessage = "Provider unavailable"; + + expect(translator.translate(message)).toEqual([ + { + type: "runtime_error", + timestamp: 0, + errorType: "pi_runtime", + message: "Provider unavailable", + }, + ]); + }); + + it("translates assistant thinking into an agent thought chunk", () => { + const translator = createPiMessageTranslator(); + + const message = makeAssistant([{ type: "thinking", thinking: "hmm" }]); + + expect(translator.translate(message)).toEqual([ + { + type: "assistant_thought_chunk", + timestamp: 0, + content: { type: "text", text: "hmm" }, + }, + ]); + }); +}); diff --git a/packages/agent/src/pi/conversation/translatePiMessage.ts b/packages/agent/src/pi/conversation/translatePiMessage.ts new file mode 100644 index 0000000000..874dc06242 --- /dev/null +++ b/packages/agent/src/pi/conversation/translatePiMessage.ts @@ -0,0 +1,299 @@ +import type { + AssistantMessage, + Message, + ToolResultMessage, + UserMessage, +} from "@earendil-works/pi-ai"; +import type { + AgentContent, + AgentConversationEvent, + AgentToolCallStatus, +} from "@posthog/shared"; +import { type PiToolName, TOOL_KIND_BY_NAME } from "./toolKind"; +import { bashTranslator } from "./tools/bashTranslator"; +import { editTranslator } from "./tools/editTranslator"; +import { findTranslator } from "./tools/findTranslator"; +import { grepTranslator } from "./tools/grepTranslator"; +import { lsTranslator } from "./tools/lsTranslator"; +import { readTranslator } from "./tools/readTranslator"; +import { writeTranslator } from "./tools/writeTranslator"; +import type { PiToolTranslator } from "./toolTranslator"; + +const TRANSLATOR_BY_NAME: Record = { + read: readTranslator, + bash: bashTranslator, + edit: editTranslator, + write: writeTranslator, + grep: grepTranslator, + find: findTranslator, + ls: lsTranslator, +}; + +interface PendingToolCall { + name: string; + arguments: unknown; +} + +interface PiToolExecutionResult { + content: ToolResultMessage["content"]; + details?: unknown; +} + +function isPiToolName(name: string): name is PiToolName { + return name in TOOL_KIND_BY_NAME; +} + +function toContent(block: { + type: string; + text?: string; + data?: string; + mimeType?: string; +}): AgentContent | undefined { + if (block.type === "text" && typeof block.text === "string") { + return { type: "text", text: block.text }; + } + + if ( + block.type === "image" && + typeof block.data === "string" && + typeof block.mimeType === "string" + ) { + return { type: "image", data: block.data, mimeType: block.mimeType }; + } + + return undefined; +} + +export interface PiMessageTranslator { + translate(message: Message): AgentConversationEvent[]; + translateToolExecutionStart( + toolCallId: string, + toolName: string, + args: unknown, + timestamp: number, + ): AgentConversationEvent[]; + translateToolExecutionUpdate( + toolCallId: string, + toolName: string, + args: unknown, + result: PiToolExecutionResult, + timestamp: number, + ): AgentConversationEvent[]; + translateToolExecutionEnd( + toolCallId: string, + toolName: string, + result: PiToolExecutionResult, + isError: boolean, + timestamp: number, + ): AgentConversationEvent[]; +} + +export function createPiMessageTranslator(): PiMessageTranslator { + const pendingToolCalls = new Map(); + let userMessageId = 0; + + function translateUser(message: UserMessage): AgentConversationEvent[] { + const content = + typeof message.content === "string" + ? [{ type: "text" as const, text: message.content }] + : message.content.flatMap((block) => { + const translated = toContent(block); + return translated ? [translated] : []; + }); + + if (content.length === 0) { + return []; + } + + userMessageId += 1; + + return [ + { + type: "user_message", + id: `pi-user-${message.timestamp}-${userMessageId}`, + timestamp: message.timestamp, + content, + }, + ]; + } + + function translateAssistant( + message: AssistantMessage, + ): AgentConversationEvent[] { + const events: AgentConversationEvent[] = []; + + for (const block of message.content) { + if (block.type === "text") { + events.push({ + type: "assistant_message_chunk", + timestamp: message.timestamp, + content: { type: "text", text: block.text }, + }); + continue; + } + + if (block.type === "thinking") { + events.push({ + type: "assistant_thought_chunk", + timestamp: message.timestamp, + content: { type: "text", text: block.thinking }, + }); + continue; + } + + if (block.type === "toolCall") { + pendingToolCalls.set(block.id, { + name: block.name, + arguments: block.arguments, + }); + + const kind = isPiToolName(block.name) + ? TOOL_KIND_BY_NAME[block.name] + : null; + + events.push({ + type: "tool_call_started", + timestamp: message.timestamp, + toolCall: { + id: block.id, + title: block.name, + kind, + status: "pending", + rawInput: block.arguments, + }, + }); + } + } + + if (message.stopReason === "error") { + events.push({ + type: "runtime_error", + timestamp: message.timestamp, + errorType: "pi_runtime", + message: message.errorMessage ?? "Pi runtime failed", + }); + } + + return events; + } + + function translateToolExecution( + toolCallId: string, + toolName: string, + args: unknown, + result: PiToolExecutionResult, + status: AgentToolCallStatus, + timestamp: number, + ): AgentConversationEvent[] { + const toolCall: Extract< + AgentConversationEvent, + { type: "tool_call_updated" } + >["toolCall"] = { + id: toolCallId, + status, + rawOutput: result.content, + }; + + const translator = isPiToolName(toolName) + ? TRANSLATOR_BY_NAME[toolName] + : undefined; + + if (translator) { + const output = translator({ + toolCallId, + arguments: args, + resultContent: result.content, + details: result.details, + isError: status === "failed", + }); + + if (output.content) { + toolCall.content = output.content; + } + + if (output.locations) { + toolCall.locations = output.locations; + } + } + + return [{ type: "tool_call_updated", timestamp, toolCall }]; + } + + function translateToolResult( + message: ToolResultMessage, + ): AgentConversationEvent[] { + const pending = pendingToolCalls.get(message.toolCallId); + pendingToolCalls.delete(message.toolCallId); + + return translateToolExecution( + message.toolCallId, + message.toolName, + pending?.arguments, + { content: message.content, details: message.details }, + message.isError ? "failed" : "completed", + message.timestamp, + ); + } + + return { + translate(message: Message): AgentConversationEvent[] { + if (message.role === "user") { + return translateUser(message); + } + + if (message.role === "assistant") { + return translateAssistant(message); + } + + return translateToolResult(message); + }, + + translateToolExecutionStart(toolCallId, toolName, args, timestamp) { + pendingToolCalls.set(toolCallId, { name: toolName, arguments: args }); + + return [ + { + type: "tool_call_updated", + timestamp, + toolCall: { id: toolCallId, status: "in_progress" }, + }, + ]; + }, + + translateToolExecutionUpdate( + toolCallId, + toolName, + args, + result, + timestamp, + ) { + return translateToolExecution( + toolCallId, + toolName, + args, + result, + "in_progress", + timestamp, + ); + }, + + translateToolExecutionEnd( + toolCallId, + toolName, + result, + isError, + timestamp, + ) { + const pending = pendingToolCalls.get(toolCallId); + + return translateToolExecution( + toolCallId, + toolName, + pending?.arguments, + result, + isError ? "failed" : "completed", + timestamp, + ); + }, + }; +} diff --git a/packages/agent/src/pi/rpc-client.test.ts b/packages/agent/src/pi/rpc-client.test.ts index 4c8239d698..ff4dcc3c23 100644 --- a/packages/agent/src/pi/rpc-client.test.ts +++ b/packages/agent/src/pi/rpc-client.test.ts @@ -1,6 +1,39 @@ import { RpcClient } from "@earendil-works/pi-coding-agent"; import { describe, expect, it } from "vitest"; -import { createPiRpcClient } from "./rpc-client"; +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"], + }), + ], + ); + }); +}); 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 85896bf6f4..5806b77d1c 100644 --- a/packages/agent/src/pi/rpc-client.ts +++ b/packages/agent/src/pi/rpc-client.ts @@ -2,15 +2,34 @@ 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 type { PosthogProviderOptions } from "@posthog/harness/extensions/posthog-provider/provider"; import { safePiEnvironment } from "./rpc-environment"; +import type { PiModelOption, PiThinkingLevel } from "./types"; export type PiRpcClient = RpcClient; +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; }; diff --git a/packages/agent/src/pi/runtime.test.ts b/packages/agent/src/pi/runtime.test.ts new file mode 100644 index 0000000000..b9cf4656de --- /dev/null +++ b/packages/agent/src/pi/runtime.test.ts @@ -0,0 +1,76 @@ +import type { AssistantMessage } from "@earendil-works/pi-ai"; +import type { + AgentSessionEvent, + RpcClient, +} from "@earendil-works/pi-coding-agent"; +import { describe, expect, it, vi } from "vitest"; +import { PiRuntime } from "./runtime"; + +function assistant(text: string): AssistantMessage { + return { + role: "assistant", + content: [{ type: "text", text }], + api: "anthropic-messages" as AssistantMessage["api"], + provider: "anthropic" as AssistantMessage["provider"], + model: "test-model", + usage: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 0, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + stopReason: "stop", + timestamp: 1, + }; +} + +function createClient(messages: AssistantMessage[] = []) { + 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) }; +} + +describe("PiRuntime", () => { + it("normalizes live Pi events before forwarding them", () => { + const { client, emit } = createClient(); + const runtime = new PiRuntime(client); + const conversationListener = vi.fn(); + runtime.onConversationEvent(conversationListener); + + emit({ type: "message_end", message: assistant("hello") }); + + expect(conversationListener).toHaveBeenCalledWith({ + type: "assistant_message_chunk", + timestamp: 1, + 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 new file mode 100644 index 0000000000..102c517fcb --- /dev/null +++ b/packages/agent/src/pi/runtime.ts @@ -0,0 +1,83 @@ +import type { AgentSessionEvent } from "@earendil-works/pi-coding-agent"; +import type { AgentConversationEvent } from "@posthog/shared"; +import { + createPiConversationTranslator, + type PiConversationTranslator, +} from "./conversation/translatePiConversation"; +import { + createPiRpcClient, + getAvailableModelsWithThinkingLevels, + getPiRpcClientProcess, + type PiRpcClient, + type PiRpcClientOptions, +} from "./rpc-client"; +import type { PiModelOption } from "./types"; + +export class PiRuntime { + readonly client: PiRpcClient; + + private readonly translator: PiConversationTranslator; + private readonly runtimeListeners = new Set< + (event: AgentSessionEvent) => void + >(); + private readonly conversationListeners = new Set< + (event: AgentConversationEvent) => void + >(); + + constructor(client: PiRpcClient) { + this.client = client; + this.translator = createPiConversationTranslator(); + client.onEvent((event) => this.handleEvent(event)); + } + + get process() { + return getPiRpcClientProcess(this.client); + } + + onRuntimeEvent(listener: (event: AgentSessionEvent) => void): () => void { + this.runtimeListeners.add(listener); + return () => this.runtimeListeners.delete(listener); + } + + onConversationEvent( + listener: (event: AgentConversationEvent) => void, + ): () => void { + this.conversationListeners.add(listener); + 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); + } + + const conversationEvents = this.translator.translateEvent(event); + for (const conversationEvent of conversationEvents) { + for (const listener of this.conversationListeners) { + listener(conversationEvent); + } + } + } +} + +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 new file mode 100644 index 0000000000..30443c7596 --- /dev/null +++ b/packages/agent/src/pi/types.ts @@ -0,0 +1,45 @@ +import type { QueueMode, ThinkingLevel } from "@earendil-works/pi-agent-core"; +import type { + RpcClient, + RpcSessionState, +} from "@earendil-works/pi-coding-agent"; + +function exhaustiveValues() { + return ( + values: Values & ([T] extends [Values[number]] ? unknown : never), + ): Values => values; +} + +export type PiThinkingLevel = ThinkingLevel; +export type PiQueueMode = QueueMode; + +export const PI_THINKING_LEVELS = exhaustiveValues()([ + "off", + "minimal", + "low", + "medium", + "high", + "xhigh", + "max", +]); + +export const PI_QUEUE_MODES = exhaustiveValues()([ + "all", + "one-at-a-time", +]); + +export type PiNativeModelInfo = Awaited< + ReturnType +>[number]; + +export type PiModelOption = PiNativeModelInfo & { + thinkingLevels: PiThinkingLevel[]; +}; + +export type PiCommand = Awaited>[number]; + +export type PiSessionStatus = Omit & { + model?: Pick, "provider" | "id">; +}; + +export type PiSessionStats = Awaited>; diff --git a/packages/agent/tsconfig.build.json b/packages/agent/tsconfig.build.json index 8db109166a..2659241485 100644 --- a/packages/agent/tsconfig.build.json +++ b/packages/agent/tsconfig.build.json @@ -2,13 +2,15 @@ "extends": "./tsconfig.json", "compilerOptions": { "noEmit": false, + "rootDir": "src", "outDir": "dist", "declaration": true, - "declarationMap": true, - "emitDeclarationOnly": false, - "sourceMap": true, + "declarationMap": false, + "emitDeclarationOnly": true, + "sourceMap": false, "removeComments": false, "allowImportingTsExtensions": false }, + "include": ["src/**/*.ts"], "exclude": ["dist", "node_modules", "**/*.test.ts", "**/*.spec.ts"] } diff --git a/packages/agent/tsup.config.ts b/packages/agent/tsup.config.ts index 93a1e2ce88..57ebcd9f6a 100644 --- a/packages/agent/tsup.config.ts +++ b/packages/agent/tsup.config.ts @@ -114,6 +114,9 @@ export default defineConfig([ "src/posthog-products.ts", "src/pr-url-detector.ts", "src/pi/rpc-client.ts", + "src/pi/runtime.ts", + "src/pi/types.ts", + "src/pi/conversation/translatePiConversation.ts", "src/resume.ts", "src/types.ts", "src/adapters/claude/questions/utils.ts", @@ -132,7 +135,7 @@ export default defineConfig([ "src/server/agent-server.ts", ], format: ["esm"], - dts: true, + dts: false, clean: false, // noExternal inlines CJS deps (e.g. simple-git via @posthog/git) whose // dynamic `require(...)` calls throw in ESM output unless a real require diff --git a/packages/core/package.json b/packages/core/package.json index 32d069cbae..5a5a1aa81e 100644 --- a/packages/core/package.json +++ b/packages/core/package.json @@ -20,6 +20,7 @@ "@modelcontextprotocol/ext-apps": "^1.1.2", "@modelcontextprotocol/sdk": "^1.12.1", "@pierre/diffs": "^1.2.10", + "@posthog/agent": "workspace:*", "@posthog/api-client": "workspace:*", "@posthog/di": "workspace:*", "@posthog/platform": "workspace:*", diff --git a/packages/core/src/pi-runtime/identifiers.ts b/packages/core/src/pi-runtime/identifiers.ts index 42e0c14ab3..a91ad780f4 100644 --- a/packages/core/src/pi-runtime/identifiers.ts +++ b/packages/core/src/pi-runtime/identifiers.ts @@ -1 +1,2 @@ export const PI_RUNNER = Symbol.for("posthog.pi.runner"); +export const PI_SESSION_CONTROLLER = Symbol.for("posthog.pi.sessionController"); diff --git a/packages/core/src/pi-runtime/pi-runtime.module.ts b/packages/core/src/pi-runtime/pi-runtime.module.ts new file mode 100644 index 0000000000..4fd1d3a7f1 --- /dev/null +++ b/packages/core/src/pi-runtime/pi-runtime.module.ts @@ -0,0 +1,7 @@ +import { ContainerModule } from "inversify"; +import { PI_SESSION_CONTROLLER } from "./identifiers"; +import { PiSessionController } from "./piSessionController"; + +export const piRuntimeModule = new ContainerModule(({ bind }) => { + 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 new file mode 100644 index 0000000000..1c771e0d46 --- /dev/null +++ b/packages/core/src/pi-runtime/piSessionController.test.ts @@ -0,0 +1,197 @@ +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, + PiSessionController, +} from "./piSessionController"; + +function createController( + client = createClient(), + taskService = { + openTask: vi.fn(async () => ({ success: true })), + } as unknown as TaskService, +): PiSessionController { + return new PiSessionController(client, taskService); +} + +function createClient(): PiSessionClient { + return { + health: vi.fn(async () => ({ state: "idle" as const })), + conversation: vi.fn(async () => []), + status: vi.fn(async () => ({ + thinkingLevel: "off" as const, + isStreaming: false, + isCompacting: false, + steeringMode: "all" as const, + followUpMode: "all" as const, + sessionId: "session-1", + autoCompactionEnabled: true, + messageCount: 0, + pendingMessageCount: 0, + })), + availableModels: vi.fn(async () => []), + commands: vi.fn(async () => []), + subscribe: vi.fn(() => () => {}), + 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 })), + 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 () => {}), + }; +} + +describe("PiSessionController", () => { + it.each([ + { + text: "hello", + streaming: false, + mode: "steer" as const, + action: "prompt", + }, + { text: "hello", streaming: true, mode: "steer" as const, action: "steer" }, + { + text: "hello", + streaming: true, + mode: "queue" as const, + action: "followUp", + }, + { + text: "/compact keep details", + streaming: false, + mode: "steer" as const, + action: "compact", + }, + ])("classifies $action submissions", ({ text, streaming, mode, action }) => { + const controller = createController(); + + expect(controller.getSubmitAction(text, streaming, mode)).toBe(action); + }); + + it.each([ + { + text: "hello", + streaming: false, + mode: "steer" as const, + method: "prompt" as const, + expectedArgs: ["task-1", "hello"], + }, + { + text: "hello", + streaming: true, + mode: "steer" as const, + method: "steer" as const, + expectedArgs: ["task-1", "hello"], + }, + { + text: "hello", + streaming: true, + mode: "queue" as const, + method: "followUp" as const, + expectedArgs: ["task-1", "hello"], + }, + { + text: "/compact keep details", + streaming: false, + mode: "steer" as const, + method: "compact" as const, + expectedArgs: ["task-1", "keep details"], + }, + ])("routes submissions through $method", async (input) => { + const client = createClient(); + const controller = createController(client); + + await controller.submit("task-1", input.text, input.streaming, input.mode); + + expect(client[input.method]).toHaveBeenCalledWith(...input.expectedArgs); + }); + + it("opens cold tasks before connecting", async () => { + const client = createClient(); + 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"); + + expect(openTask).toHaveBeenCalledWith("task-1"); + expect(controller.store.getState().sessions["task-1"]).toMatchObject({ + connectionState: "connected", + }); + }); + + it("makes the transcript available before model discovery finishes", async () => { + let resolveModels: (models: []) => void = () => {}; + const models = new Promise<[]>((resolve) => { + resolveModels = resolve; + }); + const initialEvent: AgentConversationEvent = { + type: "assistant_thought_chunk", + timestamp: 1, + content: { type: "text", text: "working" }, + }; + const client = createClient(); + vi.mocked(client.conversation).mockResolvedValue([initialEvent]); + vi.mocked(client.status).mockResolvedValue({ + thinkingLevel: "high", + isStreaming: true, + isCompacting: false, + steeringMode: "all", + followUpMode: "all", + sessionId: "session-1", + autoCompactionEnabled: true, + messageCount: 1, + pendingMessageCount: 0, + }); + vi.mocked(client.availableModels).mockReturnValue(models); + const controller = createController(client); + + const connection = controller.connect("task-1"); + + await vi.waitFor(() => { + expect(controller.store.getState().sessions["task-1"]).toMatchObject({ + events: [initialEvent], + status: { isStreaming: true }, + }); + }); + + resolveModels([]); + await connection; + }); + + it("loads session state and appends normalized runtime events", async () => { + const initialEvent: AgentConversationEvent = { + type: "assistant_message_chunk", + timestamp: 1, + content: { type: "text", text: "hello" }, + }; + const liveEvent: AgentConversationEvent = { + type: "runtime_status", + timestamp: 2, + status: "compacting", + }; + let onEvent: (event: AgentConversationEvent) => void = () => {}; + const client = createClient(); + vi.mocked(client.conversation).mockResolvedValue([initialEvent]); + vi.mocked(client.subscribe).mockImplementation((_taskId, handler) => { + onEvent = handler; + return () => {}; + }); + const controller = createController(client); + + await controller.connect("task-1"); + onEvent(liveEvent); + + expect(controller.store.getState().sessions["task-1"]).toMatchObject({ + events: [initialEvent, liveEvent], + status: { isCompacting: true }, + }); + }); +}); diff --git a/packages/core/src/pi-runtime/piSessionController.ts b/packages/core/src/pi-runtime/piSessionController.ts new file mode 100644 index 0000000000..6ad5ef873e --- /dev/null +++ b/packages/core/src/pi-runtime/piSessionController.ts @@ -0,0 +1,352 @@ +import type { + PiCommand, + PiModelOption, + PiQueueMode, + PiSessionStatus, + PiThinkingLevel, +} from "@posthog/agent/pi/types"; +import type { + AgentConversationEvent, + PiMessagingMode, + PiRuntimeHealth, +} from "@posthog/shared"; +import { inject, injectable } from "inversify"; +import { parseCommandLine } from "../message-editor/commands"; +import { TASK_SERVICE, type TaskService } from "../task-detail/taskService"; +import { + createEmptyPiControllerSession, + createPiSessionStore, + type PiControllerSessionState, + type PiSessionStore, +} from "./piSessionStore"; + +export type { + PiModelOption, + PiQueueMode, + PiThinkingLevel, +} from "@posthog/agent/pi/types"; + +export const PI_SESSION_CLIENT = Symbol.for("posthog.pi.sessionClient"); + +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, + 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 type PiSubmitResult = "prompt" | "steer" | "followUp" | "compact"; + +@injectable() +export class PiSessionController { + readonly store: PiSessionStore = createPiSessionStore(); + + private readonly subscriptions = new Map void>(); + private readonly liveEvents = new Map(); + private readonly connections = new Map>(); + private readonly readiness = new Map>(); + + constructor( + @inject(PI_SESSION_CLIENT) private readonly client: PiSessionClient, + @inject(TASK_SERVICE) private readonly taskService: TaskService, + ) {} + + ensureConnected(taskId: string): Promise { + this.ensureSubscription(taskId); + + const existing = this.readiness.get(taskId); + if (existing) { + return existing; + } + + this.updateSession(taskId, { + connectionState: "connecting", + error: undefined, + }); + const readiness = this.ensureConnectedInternal(taskId) + .then(() => { + this.updateSession(taskId, { connectionState: "connected" }); + }) + .catch((error) => { + this.updateSession(taskId, { + connectionState: "failed", + error: error instanceof Error ? error.message : String(error), + }); + throw error; + }) + .finally(() => { + this.readiness.delete(taskId); + }); + this.readiness.set(taskId, readiness); + return readiness; + } + + connect(taskId: string): Promise { + this.ensureSubscription(taskId); + + const existing = this.connections.get(taskId); + if (existing) { + return existing; + } + + this.updateSession(taskId, { error: undefined }); + + const connection = this.loadSession(taskId).finally(() => { + this.connections.delete(taskId); + }); + this.connections.set(taskId, connection); + return connection; + } + + disconnect(taskId: string): void { + this.subscriptions.get(taskId)?.(); + this.subscriptions.delete(taskId); + } + + getSubmitAction( + text: string, + isStreaming: boolean, + messagingMode: PiMessagingMode, + ): PiSubmitResult { + const command = parseCommandLine(text.trim()); + if (command?.name === "compact") { + return "compact"; + } + + if (!isStreaming) { + return "prompt"; + } + + return messagingMode === "steer" ? "steer" : "followUp"; + } + + async submit( + taskId: string, + text: string, + isStreaming: boolean, + messagingMode: PiMessagingMode, + ): Promise { + 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; + } + } + + async setModel(taskId: string, model: PiModelOption): Promise { + await this.client.setModel(taskId, model.provider, model.id); + await this.refreshStatus(taskId); + } + + async setThinkingLevel( + taskId: string, + level: PiThinkingLevel, + ): Promise { + await this.client.setThinkingLevel(taskId, level); + await this.refreshStatus(taskId); + } + + async setQueueMode( + taskId: string, + messagingMode: PiMessagingMode, + queueMode: PiQueueMode, + ): Promise { + if (messagingMode === "steer") { + await this.client.setSteeringMode(taskId, queueMode); + } else { + await this.client.setFollowUpMode(taskId, queueMode); + } + await this.refreshStatus(taskId); + } + + async bash(taskId: string, command: string): Promise { + this.updateSession(taskId, { isBashRunning: true }); + try { + await this.client.bash(taskId, command); + await this.refreshConversation(taskId); + } finally { + this.updateSession(taskId, { isBashRunning: false }); + } + } + + async abort(taskId: string): Promise { + await this.client.abort(taskId); + await this.refreshStatus(taskId); + } + + async abortBash(taskId: string): Promise { + await this.client.abortBash(taskId); + this.updateSession(taskId, { isBashRunning: false }); + } + + private async ensureConnectedInternal(taskId: string): Promise { + const health = await this.client.health(taskId); + if (health.state === "cold") { + const result = await this.taskService.openTask(taskId); + if (!result.success) { + throw new Error(result.error); + } + } + + await this.connect(taskId); + } + + private ensureSubscription(taskId: string): void { + if (this.subscriptions.has(taskId)) { + return; + } + + const unsubscribe = this.client.subscribe( + taskId, + (event) => this.handleEvent(taskId, event), + (error) => { + this.updateSession(taskId, { + error: error instanceof Error ? error.message : String(error), + }); + }, + ); + this.subscriptions.set(taskId, unsubscribe); + } + + private async loadSession(taskId: string): Promise { + try { + const [events, status] = await Promise.all([ + this.client.conversation(taskId), + this.client.status(taskId), + ]); + const liveEvents = status.isStreaming + ? (this.liveEvents.get(taskId) ?? []) + : []; + const currentSession = this.getSession(taskId); + this.liveEvents.set(taskId, liveEvents); + this.setSession(taskId, { + connectionState: "connected", + events: [...events, ...liveEvents], + status, + models: currentSession.models, + commands: currentSession.commands, + isBashRunning: false, + error: undefined, + }); + + const [models, commands] = await Promise.all([ + this.client.availableModels(taskId), + this.client.commands(taskId), + ]); + this.updateSession(taskId, { models, commands }); + } catch (error) { + this.updateSession(taskId, { + error: error instanceof Error ? error.message : String(error), + }); + throw error; + } + } + + private handleEvent(taskId: string, event: AgentConversationEvent): void { + 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") { + status = { ...status, isCompacting: !event.isComplete }; + } else if (event.status === "compacting_failed") { + status = { ...status, isCompacting: false }; + } + } + if (status && event.type === "turn_completed") { + status = { ...status, isStreaming: false }; + } + + this.updateSession(taskId, { + events: [...session.events, event], + status, + }); + + if (event.type === "turn_completed") { + const capturedCount = liveEvents.length; + void this.refreshConversation(taskId, capturedCount); + } + } + + private async refreshConversation( + taskId: string, + capturedLiveCount?: number, + ): Promise { + const events = await this.client.conversation(taskId); + const liveEvents = this.liveEvents.get(taskId) ?? []; + const remainingEvents = + capturedLiveCount === undefined + ? [] + : liveEvents.slice(capturedLiveCount); + this.liveEvents.set(taskId, remainingEvents); + this.updateSession(taskId, { + events: [...events, ...remainingEvents], + }); + } + + private async refreshStatus(taskId: string): Promise { + const status = await this.client.status(taskId); + this.updateSession(taskId, { status }); + } + + private getSession(taskId: string): PiControllerSessionState { + return ( + this.store.getState().sessions[taskId] ?? createEmptyPiControllerSession() + ); + } + + private setSession(taskId: string, session: PiControllerSessionState): void { + this.store.setState((state) => ({ + sessions: { ...state.sessions, [taskId]: session }, + })); + } + + private updateSession( + taskId: string, + update: Partial, + ): void { + this.setSession(taskId, { ...this.getSession(taskId), ...update }); + } +} diff --git a/packages/core/src/pi-runtime/piSessionStore.ts b/packages/core/src/pi-runtime/piSessionStore.ts new file mode 100644 index 0000000000..1b4519a59f --- /dev/null +++ b/packages/core/src/pi-runtime/piSessionStore.ts @@ -0,0 +1,37 @@ +import type { + PiCommand, + PiModelOption, + PiSessionStatus, +} from "@posthog/agent/pi/types"; +import type { AgentConversationEvent } from "@posthog/shared"; +import { createStore, type StoreApi } from "zustand/vanilla"; + +export interface PiControllerSessionState { + connectionState: "connecting" | "connected" | "failed"; + events: AgentConversationEvent[]; + models: PiModelOption[]; + commands: PiCommand[]; + status?: PiSessionStatus; + error?: string; + isBashRunning: boolean; +} + +export interface PiSessionState { + sessions: Record; +} + +export type PiSessionStore = StoreApi; + +export function createPiSessionStore(): PiSessionStore { + return createStore(() => ({ sessions: {} })); +} + +export function createEmptyPiControllerSession(): PiControllerSessionState { + return { + connectionState: "connecting", + events: [], + models: [], + commands: [], + isBashRunning: false, + }; +} diff --git a/packages/host-router/src/pi-session-client.ts b/packages/host-router/src/pi-session-client.ts new file mode 100644 index 0000000000..d168a0c914 --- /dev/null +++ b/packages/host-router/src/pi-session-client.ts @@ -0,0 +1,96 @@ +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/routers/pi-session.router.ts b/packages/host-router/src/routers/pi-session.router.ts index 67eba086d1..625fc0e009 100644 --- a/packages/host-router/src/routers/pi-session.router.ts +++ b/packages/host-router/src/routers/pi-session.router.ts @@ -2,10 +2,40 @@ 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, + piSessionHealthOutput, + piSessionLastAssistantTextOutput, + piSessionMessageInput, + piSessionModelInput, + piSessionModelOutput, + piSessionNameInput, + piSessionNewInput, + piSessionPathInput, + piSessionPromptAndWaitInput, piSessionPromptInput, + piSessionQueueModeInput, piSessionStartOutput, + piSessionStatusOutput, + piSessionStderrOutput, + piSessionThinkingCycleOutput, + piSessionThinkingLevelInput, + piSessionTimeoutInput, piSessionTranscriptInput, + piSessionUnknownOutput, resumePiSessionInput, startPiSessionInput, } from "@posthog/workspace-server/services/pi-session/schemas"; @@ -26,7 +56,31 @@ export const piSessionRouter = router({ prompt: publicProcedure .input(piSessionPromptInput) .mutation(({ ctx, input }) => - getService(ctx.container).prompt(input.taskId, input.prompt), + getService(ctx.container).prompt( + input.taskId, + input.prompt, + input.images, + ), + ), + + 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 @@ -35,20 +89,232 @@ export const piSessionRouter = router({ getService(ctx.container).abort(input.taskId), ), - stop: publicProcedure + 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) - .mutation(({ ctx, input }) => getService(ctx.container).stop(input.taskId)), + .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) + .mutation(({ ctx, input }) => getService(ctx.container).stop(input.taskId)), + + health: publicProcedure + .input(piSessionTranscriptInput) + .output(piSessionHealthOutput) + .query(({ ctx, input }) => getService(ctx.container).health(input.taskId)), + onEvent: publicProcedure .input(piSessionTranscriptInput) .subscription(async function* (opts) { diff --git a/packages/shared/src/agent-conversation.ts b/packages/shared/src/agent-conversation.ts new file mode 100644 index 0000000000..9f03b345f6 --- /dev/null +++ b/packages/shared/src/agent-conversation.ts @@ -0,0 +1,158 @@ +export type AgentToolKind = + | "read" + | "edit" + | "delete" + | "move" + | "search" + | "execute" + | "think" + | "fetch" + | "switch_mode" + | "question" + | "other"; + +export type AgentToolCallStatus = + | "pending" + | "in_progress" + | "completed" + | "failed"; + +export interface AgentTextContent { + type: "text"; + text: string; +} + +export interface AgentImageContent { + type: "image"; + data: string; + mimeType: string; +} + +export interface AgentAudioContent { + type: "audio"; + data: string; + mimeType: string; +} + +export interface AgentResourceLinkContent { + type: "resource_link"; + uri: string; + name: string; + description?: string | null; + mimeType?: string | null; + size?: number | null; + title?: string | null; +} + +export interface AgentTextResource { + uri: string; + mimeType?: string | null; + text: string; +} + +export interface AgentBlobResource { + uri: string; + mimeType?: string | null; + blob: string; +} + +export interface AgentEmbeddedResourceContent { + type: "resource"; + resource: AgentTextResource | AgentBlobResource; +} + +export type AgentContent = + | AgentTextContent + | AgentImageContent + | AgentAudioContent + | AgentResourceLinkContent + | AgentEmbeddedResourceContent; + +export interface AgentToolCallContentBlock { + type: "content"; + content: AgentContent; +} + +export interface AgentToolCallDiff { + type: "diff"; + path: string; + oldText?: string | null; + newText: string; +} + +export interface AgentToolCallTerminal { + type: "terminal"; + terminalId: string; +} + +export type AgentToolCallContent = + | AgentToolCallContentBlock + | AgentToolCallDiff + | AgentToolCallTerminal; + +export interface AgentToolCallLocation { + path: string; + line?: number | null; +} + +export interface AgentToolCall { + id: string; + title: string; + kind?: AgentToolKind | null; + status?: AgentToolCallStatus | null; + content?: AgentToolCallContent[]; + locations?: AgentToolCallLocation[]; + rawInput?: unknown; + rawOutput?: unknown; + parentId?: string; +} + +export type AgentConversationEvent = + | { + type: "user_message"; + id: string; + timestamp: number; + content: AgentContent[]; + } + | { + type: "assistant_message_chunk"; + timestamp: number; + content: AgentContent; + } + | { + type: "assistant_thought_chunk"; + timestamp: number; + content: AgentContent; + } + | { + type: "tool_call_started"; + timestamp: number; + toolCall: AgentToolCall; + } + | { + type: "tool_call_updated"; + timestamp: number; + toolCall: Pick & Partial>; + } + | { + type: "runtime_status"; + timestamp: number; + status: string; + isComplete?: boolean; + error?: string; + message?: string; + attempt?: number; + maxAttempts?: number; + delayMs?: number; + } + | { + type: "runtime_error"; + timestamp: number; + errorType: string; + message: string; + } + | { + type: "turn_completed"; + timestamp: number; + stopReason?: string; + }; diff --git a/packages/shared/src/index.ts b/packages/shared/src/index.ts index ae104cb828..2158b3eae3 100644 --- a/packages/shared/src/index.ts +++ b/packages/shared/src/index.ts @@ -1,4 +1,23 @@ export * from "./adapter"; +export type { + AgentAudioContent, + AgentBlobResource, + AgentContent, + AgentConversationEvent, + AgentEmbeddedResourceContent, + AgentImageContent, + AgentResourceLinkContent, + AgentTextContent, + AgentTextResource, + AgentToolCall, + AgentToolCallContent, + AgentToolCallContentBlock, + AgentToolCallDiff, + AgentToolCallLocation, + AgentToolCallStatus, + AgentToolCallTerminal, + AgentToolKind, +} from "./agent-conversation"; export * from "./agent-runtime"; export * from "./analytics-events"; export { type ArchivedTask, archivedTaskSchema } from "./archive-domain"; @@ -166,6 +185,10 @@ export { pathToFileUri, toRelativePath, } from "./path"; +export type { + PiMessagingMode, + PiRuntimeHealth, +} from "./pi-session"; export { buildPrOutput, mergePrUrls, diff --git a/packages/shared/src/pi-session.ts b/packages/shared/src/pi-session.ts new file mode 100644 index 0000000000..c17268c2d9 --- /dev/null +++ b/packages/shared/src/pi-session.ts @@ -0,0 +1,7 @@ +export type PiMessagingMode = "steer" | "queue"; + +export interface PiRuntimeHealth { + state: "cold" | "starting" | "idle" | "streaming"; + pid?: number; + lastUsedAt?: number; +} diff --git a/packages/ui/package.json b/packages/ui/package.json index ef0cfd7514..a93ec5c3a5 100644 --- a/packages/ui/package.json +++ b/packages/ui/package.json @@ -18,7 +18,6 @@ }, "dependencies": { "@agentclientprotocol/sdk": "0.22.1", - "@earendil-works/pi-coding-agent": "catalog:", "@base-ui/react": "^1.3.0", "@codemirror/lang-angular": "^0.1.4", "@codemirror/lang-cpp": "^6.0.3", diff --git a/packages/ui/src/features/pi-sessions/PiSessionControls.tsx b/packages/ui/src/features/pi-sessions/PiSessionControls.tsx new file mode 100644 index 0000000000..48bd2fb75e --- /dev/null +++ b/packages/ui/src/features/pi-sessions/PiSessionControls.tsx @@ -0,0 +1,259 @@ +import { Brain, CaretDown, Lightning, Stack } from "@phosphor-icons/react"; +import type { + PiModelOption, + PiQueueMode, + PiThinkingLevel, +} from "@posthog/core/pi-runtime/piSessionController"; +import { + Button, + DropdownMenu, + DropdownMenuContent, + DropdownMenuRadioGroup, + DropdownMenuRadioItem, + DropdownMenuSeparator, + DropdownMenuTrigger, + MenuLabel, +} from "@posthog/quill"; +import type { MessagingMode } from "@posthog/ui/features/sessions/messagingModeStore"; +import { Fragment } from "react"; + +interface PiModelSelectorProps { + models: PiModelOption[]; + currentModel?: Pick; + disabled?: boolean; + onChange: (model: PiModelOption) => void; +} + +function modelKey(model: Pick): string { + return JSON.stringify([model.provider, model.id]); +} + +export function PiModelSelector({ + models, + currentModel, + disabled, + onChange, +}: PiModelSelectorProps) { + if (models.length === 0) { + return null; + } + + const modelsByProvider = new Map(); + for (const model of models) { + const providerModels = modelsByProvider.get(model.provider) ?? []; + providerModels.push(model); + modelsByProvider.set(model.provider, providerModels); + } + + const currentValue = currentModel ? modelKey(currentModel) : ""; + const currentLabel = currentModel?.id ?? "Model"; + + return ( + + + {currentLabel} + + + } + /> + + { + const model = models.find( + (candidate) => modelKey(candidate) === value, + ); + if (model) { + onChange(model); + } + }} + > + {[...modelsByProvider.entries()].map( + ([provider, providerModels], index) => ( + + {index > 0 && } + {provider} + {providerModels.map((model) => ( + + {model.id} + + ))} + + ), + )} + + + + ); +} + +const thinkingLevelLabels: Record = { + off: "Off", + minimal: "Minimal", + low: "Low", + medium: "Medium", + high: "High", + xhigh: "Extra high", + max: "Max", +}; + +interface PiThinkingLevelSelectorProps { + level: PiThinkingLevel; + levels: PiThinkingLevel[]; + disabled?: boolean; + onChange: (level: PiThinkingLevel) => void; +} + +export function PiThinkingLevelSelector({ + level, + levels, + disabled, + onChange, +}: PiThinkingLevelSelectorProps) { + const activeLabel = thinkingLevelLabels[level] ?? level; + + return ( + + + + {activeLabel} + + + } + /> + + Thinking + onChange(value as PiThinkingLevel)} + > + {levels.map((value) => ( + + {thinkingLevelLabels[value] ?? value} + + ))} + + + + ); +} + +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") { + label = "Steer"; + } else if (queuedCount > 0) { + label = `Queue (${queuedCount})`; + } + + const colorClass = mode === "steer" ? "text-purple-11" : "text-gray-11"; + + return ( + + + + {mode === "steer" ? ( + + ) : ( + + )} + + {label} + + + } + /> + + While Pi is generating + onModeChange(value as MessagingMode)} + > + + Steer at the next tool boundary + + + 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 f97cdaa3c4..f0c5894c7e 100644 --- a/packages/ui/src/features/pi-sessions/PiSessionView.tsx +++ b/packages/ui/src/features/pi-sessions/PiSessionView.tsx @@ -1,196 +1,263 @@ -import { useHostTRPC, useHostTRPCClient } from "@posthog/host-router/react"; +import { PI_SESSION_CONTROLLER } from "@posthog/core/pi-runtime/identifiers"; +import type { + PiModelOption, + PiQueueMode, + PiSessionController, + PiThinkingLevel, +} from "@posthog/core/pi-runtime/piSessionController"; +import { useService } from "@posthog/di/react"; import { - Button, Empty, EmptyDescription, EmptyHeader, EmptyTitle, } from "@posthog/quill"; -import { useQuery } from "@tanstack/react-query"; -import { useSubscription } from "@trpc/tanstack-react-query"; -import { useEffect, useMemo, useState } from "react"; +import { PromptInput } from "@posthog/ui/features/message-editor/components/PromptInput"; +import { useDraftStore } from "@posthog/ui/features/message-editor/draftStore"; +import { ChatThread } from "@posthog/ui/features/sessions/components/chat-thread/ChatThread"; +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 { toast } from "@posthog/ui/primitives/toast"; +import { TaskDetailSkeleton } from "@posthog/ui/router/routeSkeletons"; +import { Box, Flex } from "@radix-ui/themes"; +import { useCallback, useEffect } from "react"; +import { useStore } from "zustand"; import { - applyPiEvent, - emptyLiveFeed, - type PiEntries, - PiEntriesSyncer, - type PiEvent, - type PiLiveFeed, - type PiMessage, -} from "./piSessionFeed"; -import { useEnsurePiSession } from "./useEnsurePiSession"; + PiMessagingModeSelector, + PiModelSelector, + PiThinkingLevelSelector, +} from "./PiSessionControls"; interface PiSessionViewProps { taskId: string; } -type PiMessageWithContent = Extract; +export function PiSessionView({ taskId }: PiSessionViewProps) { + const piSessionController = useService( + PI_SESSION_CONTROLLER, + ); + const session = useStore( + piSessionController.store, + (state) => state.sessions[taskId], + ); + const draftActions = useDraftStore((state) => state.actions); + const workspace = useWorkspace(taskId); + const repoPath = workspace?.worktreePath ?? workspace?.folderPath; + const messagingMode = useMessagingMode(taskId); + const setMessagingMode = useMessagingModeStore((state) => state.setMode); -function PiMessageView({ message }: { message: PiMessage }) { - if (message.role === "bashExecution") { - return <>{message.output}; - } - if ( - message.role === "branchSummary" || - message.role === "compactionSummary" - ) { - return <>{message.summary}; - } - if ("content" in message) { - return <>{messageContentText(message)}; - } - return null; -} + useEffect(() => { + void piSessionController.ensureConnected(taskId); + return () => piSessionController.disconnect(taskId); + }, [piSessionController, taskId]); -function messageContentText(message: PiMessageWithContent): string { - if (typeof message.content === "string") { - return message.content; - } + const sessionAvailable = session?.connectionState === "connected"; + const status = session?.status; + const isStreaming = status?.isStreaming ?? false; + const isCompacting = status?.isCompacting ?? false; + const isBashRunning = session?.isBashRunning ?? false; - return message.content - .flatMap((part) => (part.type === "text" ? [part.text] : [])) - .join("\n"); -} + useEffect(() => { + draftActions.setContext(taskId, { + taskId, + repoPath, + disabled: !sessionAvailable || isCompacting, + isLoading: isStreaming || isBashRunning, + }); + }, [ + draftActions, + isBashRunning, + isCompacting, + isStreaming, + repoPath, + sessionAvailable, + taskId, + ]); -function messageBubbleClass(role: string): string { - if (role === "user") { - return "mb-3 ml-auto max-w-[80%] rounded-lg bg-accent-3 p-3 text-sm"; - } - return "mb-3 max-w-[80%] whitespace-pre-wrap rounded-lg bg-gray-3 p-3 text-sm"; -} + useEffect(() => { + if (!session?.commands) { + return; + } -export function PiSessionView({ taskId }: PiSessionViewProps) { - const trpc = useHostTRPC(); - const client = useHostTRPCClient(); - const { error: ensureError, isSuccess: sessionReady } = - useEnsurePiSession(taskId); - - const [prompt, setPrompt] = useState(""); - const [liveFeed, setLiveFeed] = useState(emptyLiveFeed); - const [syncedEntries, setSyncedEntries] = useState( - undefined, + const piCommands = session.commands + .filter((command) => command.name !== "compact") + .map((command) => ({ + name: command.name, + description: command.description ?? "", + })); + + draftActions.setCommands(taskId, [ + { + name: "compact", + description: "Compact the current Pi session context", + input: { hint: "optional instructions" }, + }, + ...piCommands, + ]); + }, [draftActions, session?.commands, taskId]); + + const sendPrompt = useCallback( + (text: string) => { + const message = text.trim(); + if (!message) { + return; + } + + const action = piSessionController.getSubmitAction( + message, + isStreaming, + messagingMode, + ); + void piSessionController + .submit(taskId, message, isStreaming, messagingMode) + .then(() => { + if (action === "compact") { + toast.success("Pi context compacted"); + } + }) + .catch(() => { + const failureMessage = + action === "compact" + ? "Failed to compact Pi context" + : "Failed to send message to Pi"; + toast.error(failureMessage); + }); + }, + [isStreaming, messagingMode, piSessionController, taskId], ); - const { data: fetchedEntries } = useQuery({ - ...trpc.piSession.entries.queryOptions({ taskId }), - enabled: sessionReady, - }); - const { error: statusError } = useQuery({ - ...trpc.piSession.status.queryOptions({ taskId }), - enabled: sessionReady, - }); - - const syncer = useMemo( - () => - new PiEntriesSyncer( - (since) => client.piSession.entries.query({ taskId, since }), - setSyncedEntries, - ), - [client, taskId], + const setModel = useCallback( + (model: PiModelOption) => { + void piSessionController + .setModel(taskId, model) + .catch(() => toast.error("Failed to change Pi model")); + }, + [piSessionController, taskId], ); - useEffect(() => { - syncer.seed(fetchedEntries); - }, [syncer, fetchedEntries]); + const setThinkingLevel = useCallback( + (level: PiThinkingLevel) => { + void piSessionController + .setThinkingLevel(taskId, level) + .catch(() => toast.error("Failed to change Pi thinking level")); + }, + [piSessionController, taskId], + ); - const history = syncedEntries ?? fetchedEntries; + const setQueueMode = useCallback( + (mode: PiQueueMode) => { + void piSessionController + .setQueueMode(taskId, messagingMode, mode) + .catch(() => toast.error("Failed to change Pi queue behavior")); + }, + [messagingMode, piSessionController, taskId], + ); - useSubscription( - trpc.piSession.onEvent.subscriptionOptions( - { taskId }, - { - enabled: sessionReady, - onData: (event: PiEvent) => { - setLiveFeed((feed) => applyPiEvent(feed, event)); + const toggleMessagingMode = useCallback(() => { + const nextMode = messagingMode === "steer" ? "queue" : "steer"; + setMessagingMode(taskId, nextMode); + }, [messagingMode, setMessagingMode, taskId]); - if (event.type === "agent_settled") { - void syncer.sync().then(() => setLiveFeed(emptyLiveFeed)); - } - }, - }, - ), - ); + const runBashCommand = (command: string) => { + void piSessionController + .bash(taskId, command) + .catch(() => toast.error("Failed to run Pi bash command")); + }; - const send = async () => { - const text = prompt.trim(); - if (!text) { + const cancelPrompt = () => { + if (isBashRunning) { + void piSessionController.abortBash(taskId); return; } - await client.piSession.prompt.mutate({ taskId, prompt: text }); - setPrompt(""); + + void piSessionController.abort(taskId); }; - const sessionError = ensureError ?? statusError; + const sessionError = session?.error; if (sessionError) { return ( Pi session failed to start - {sessionError.message} + {sessionError} ); } - if (!sessionReady) { - return ( - - - Starting Pi session… - - - ); + if (!session || !status) { + return ; } + const pending = isStreaming || isBashRunning; + const currentModel = session.models.find( + (model) => + model.provider === status.model?.provider && model.id === status.model.id, + ); + const thinkingLevels = currentModel?.thinkingLevels ?? []; + const supportsThinking = thinkingLevels.some((level) => level !== "off"); + const queueMode = + messagingMode === "steer" ? status.steeringMode : status.followUpMode; + return ( -
-
- {history?.entries.map((entry) => { - if (entry.type !== "message") { - return null; + + + + + + } - - return ( -
- -
- ); - })} - {liveFeed.liveMessages.map((message) => ( -
- -
- ))} - {liveFeed.streamingMessage ? ( -
- -
- ) : null} -
-
-