From 8e757a19a80dbe3309636b06e74f9b7927a8a5db Mon Sep 17 00:00:00 2001 From: "posthog[bot]" <206114724+posthog[bot]@users.noreply.github.com> Date: Tue, 28 Jul 2026 08:23:36 +0000 Subject: [PATCH 1/2] feat(agent): return artifact download URL from upload_artifact The upload_artifact local tool now surfaces the presigned download URL that the finalize-upload endpoint returns for each artifact, so agents can link to the file directly in their final response. The URL is read defensively (it rides in on TaskRunArtifact once the api-client is regenerated against the updated OpenAPI spec) and omitted when absent, so this is backward-compatible with older backends. Depends on the posthog backend change that adds `url` to the finalize response. Generated-By: PostHog Code Task-Id: d3757c3d-15f3-485a-b4f3-d0b760f78963 --- docs/cloud-task-artifacts.md | 2 ++ .../local-tools/tools/upload-artifact.test.ts | 4 ++++ .../local-tools/tools/upload-artifact.ts | 21 ++++++++++++------- 3 files changed, 20 insertions(+), 7 deletions(-) diff --git a/docs/cloud-task-artifacts.md b/docs/cloud-task-artifacts.md index fa5a17315d..9f103b1cd0 100644 --- a/docs/cloud-task-artifacts.md +++ b/docs/cloud-task-artifacts.md @@ -4,4 +4,6 @@ Cloud agents run with `/tmp/workspace` as the sandbox workspace. A repository ta Files left in the sandbox filesystem are temporary and cannot be downloaded after the sandbox is released. Agents should call the `upload_artifact` tool for every non-code deliverable they create. The tool accepts files inside the session working directory, uploads them directly to task object storage, and registers them as `output` artifacts on the task run. Registered artifacts are available through the existing task artifact download endpoint. +On success the tool returns a presigned download URL for the uploaded file (minted by the finalize-upload endpoint), so an agent can link to the artifact directly in its final response. The URL is time-limited and not persisted on the run manifest; re-fetch the artifact through the download endpoint for a durable link. + Repository changes should continue to be delivered through git rather than duplicated as task artifacts. A single uploaded artifact is limited to 30 MB. diff --git a/packages/agent/src/adapters/local-tools/tools/upload-artifact.test.ts b/packages/agent/src/adapters/local-tools/tools/upload-artifact.test.ts index bb49893f72..a3b156e856 100644 --- a/packages/agent/src/adapters/local-tools/tools/upload-artifact.test.ts +++ b/packages/agent/src/adapters/local-tools/tools/upload-artifact.test.ts @@ -47,6 +47,7 @@ describe("uploadArtifactTool", () => { size: 7, storage_path: "tasks/artifacts/report.csv", uploaded_at: "2026-01-01T00:00:00Z", + url: "https://storage.example/download/report.csv", }, ]); }); @@ -83,6 +84,9 @@ describe("uploadArtifactTool", () => { }), ], ); + expect(result.content[0]?.text).toContain( + "https://storage.example/download/report.csv", + ); }); it("rejects files outside the session workspace", async () => { diff --git a/packages/agent/src/adapters/local-tools/tools/upload-artifact.ts b/packages/agent/src/adapters/local-tools/tools/upload-artifact.ts index 12fd7a242a..78fbb500f5 100644 --- a/packages/agent/src/adapters/local-tools/tools/upload-artifact.ts +++ b/packages/agent/src/adapters/local-tools/tools/upload-artifact.ts @@ -11,7 +11,8 @@ export const uploadArtifactTool = defineLocalTool({ description: "Deliver a file you created to the user as a downloadable task artifact. " + "Call this for every non-code deliverable (reports, images, archives, data files, and similar output) " + - "before your final response. The file must be inside the session workspace. Repository changes belong in git and should not be uploaded.", + "before your final response. The file must be inside the session workspace. Repository changes belong in git and should not be uploaded. " + + "On success the result includes a download URL for the uploaded file, which you can reference in your final response.", schema: { path: z .string() @@ -115,19 +116,25 @@ export const uploadArtifactTool = defineLocalTool({ }, ], ); - if ( - !finalized.some( - (artifact) => artifact.storage_path === upload.storage_path, - ) - ) { + const finalizedEntry = finalized.find( + (artifact) => artifact.storage_path === upload.storage_path, + ); + if (!finalizedEntry) { return errorResult("PostHog did not confirm the artifact upload."); } + // The finalize endpoint returns a presigned download URL on each artifact. + // Read it defensively: the field rides in on TaskRunArtifact once the + // api-client is regenerated against the updated OpenAPI spec, and older + // backends simply omit it. + const downloadUrl = (finalizedEntry as { url?: string }).url; + const linkText = downloadUrl ? ` Download URL: ${downloadUrl}` : ""; + return { content: [ { type: "text", - text: `Uploaded ${name} as a downloadable task artifact. Mention it in your final response.`, + text: `Uploaded ${name} as a downloadable task artifact.${linkText} Mention it in your final response.`, }, ], }; From e24236da2cc9c7512a1eafb4f21115c3623e027e Mon Sep 17 00:00:00 2001 From: JonathanLab Date: Thu, 30 Jul 2026 13:55:12 +0200 Subject: [PATCH 2/2] feat(pi): inject local MCP tools --- apps/code/src/main/di/bindings.ts | 4 - apps/code/src/main/di/container.ts | 3 - .../desktop-pi-rpc-client-factory.test.ts | 1 + .../desktop-pi-rpc-client-factory.ts | 1 + .../desktop-pi-runtime-factory.test.ts | 20 --- .../desktop-pi-runtime-factory.ts | 24 ---- apps/code/vite-main-plugins.mts | 41 ++++-- .../build/verify-local-tools-mcp-server.mjs | 6 +- .../codex-app-server-agent.ts | 5 +- .../mcp-server-config.test.ts} | 4 +- .../mcp-server-config.ts} | 12 +- .../mcp-server.integration.test.ts | 58 ++++++++ .../mcp-server.ts} | 17 +-- packages/agent/src/pi/rpc-client.test.ts | 55 +++++++ packages/agent/src/pi/rpc-client.ts | 47 +++++- packages/agent/src/pi/rpc-host.ts | 14 ++ .../agent/src/server/pi-agent-server.test.ts | 46 ++++++ packages/agent/src/server/pi-agent-server.ts | 7 + .../agent/src/utils/resolve-bundled-script.ts | 7 +- packages/agent/tsup.config.ts | 4 +- .../harness/src/extensions/mcp/config.test.ts | 39 ++++- packages/harness/src/extensions/mcp/config.ts | 15 ++ .../harness/src/extensions/mcp/extension.ts | 11 +- .../mcp/runtime-servers.integration.test.ts | 134 ++++++++++++++++++ packages/harness/src/extensions/registry.ts | 13 +- .../src/services/pi-session/identifiers.ts | 13 -- .../services/pi-session/pi-session.test.ts | 54 +++---- .../src/services/pi-session/pi-session.ts | 21 ++- 28 files changed, 518 insertions(+), 158 deletions(-) delete mode 100644 apps/code/src/main/platform-adapters/desktop-pi-runtime-factory.test.ts delete mode 100644 apps/code/src/main/platform-adapters/desktop-pi-runtime-factory.ts rename packages/agent/src/adapters/{codex-app-server/local-tools-mcp.test.ts => local-tools/mcp-server-config.test.ts} (97%) rename packages/agent/src/adapters/{codex-app-server/local-tools-mcp.ts => local-tools/mcp-server-config.ts} (88%) create mode 100644 packages/agent/src/adapters/local-tools/mcp-server.integration.test.ts rename packages/agent/src/adapters/{codex-app-server/local-tools-mcp-server.ts => local-tools/mcp-server.ts} (77%) create mode 100644 packages/harness/src/extensions/mcp/runtime-servers.integration.test.ts diff --git a/apps/code/src/main/di/bindings.ts b/apps/code/src/main/di/bindings.ts index 6cc67ecbf5..7bd793ee9a 100644 --- a/apps/code/src/main/di/bindings.ts +++ b/apps/code/src/main/di/bindings.ts @@ -188,9 +188,7 @@ 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"; @@ -358,8 +356,6 @@ 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 04c245598f..5874e536ed 100644 --- a/apps/code/src/main/di/container.ts +++ b/apps/code/src/main/di/container.ts @@ -188,7 +188,6 @@ 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"; @@ -231,7 +230,6 @@ 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"; @@ -379,7 +377,6 @@ 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); diff --git a/apps/code/src/main/platform-adapters/desktop-pi-rpc-client-factory.test.ts b/apps/code/src/main/platform-adapters/desktop-pi-rpc-client-factory.test.ts index 323237e3ef..d52ec55bfb 100644 --- a/apps/code/src/main/platform-adapters/desktop-pi-rpc-client-factory.test.ts +++ b/apps/code/src/main/platform-adapters/desktop-pi-rpc-client-factory.test.ts @@ -33,6 +33,7 @@ describe("DesktopPiRpcClientFactory", () => { ); expect(createPiRpcClient).toHaveBeenCalledWith({ cwd: "/workspace", + capabilities: { environment: "local" }, providerOptions: { region: "eu", baseUrl: "http://127.0.0.1:1234", diff --git a/apps/code/src/main/platform-adapters/desktop-pi-rpc-client-factory.ts b/apps/code/src/main/platform-adapters/desktop-pi-rpc-client-factory.ts index 37d36f54dc..f9006dc27a 100644 --- a/apps/code/src/main/platform-adapters/desktop-pi-rpc-client-factory.ts +++ b/apps/code/src/main/platform-adapters/desktop-pi-rpc-client-factory.ts @@ -38,6 +38,7 @@ export class DesktopPiRpcClientFactory implements PiRpcClientFactory { return createPiRpcClient({ ...input, + capabilities: { environment: "local" }, providerOptions: { region: credentials.region, baseUrl, 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 deleted file mode 100644 index d9b7a19f3f..0000000000 --- a/apps/code/src/main/platform-adapters/desktop-pi-runtime-factory.test.ts +++ /dev/null @@ -1,20 +0,0 @@ -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 deleted file mode 100644 index b085050ff5..0000000000 --- a/apps/code/src/main/platform-adapters/desktop-pi-runtime-factory.ts +++ /dev/null @@ -1,24 +0,0 @@ -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/vite-main-plugins.mts b/apps/code/vite-main-plugins.mts index 1fe97279c7..180097ef8f 100644 --- a/apps/code/vite-main-plugins.mts +++ b/apps/code/vite-main-plugins.mts @@ -167,21 +167,36 @@ export function copyPiRpcHost(): Plugin { return { name: "copy-pi-rpc-host", writeBundle() { - const candidates = [ - join(__dirname, "node_modules/@posthog/agent/dist/pi/rpc-host.js"), - join( - __dirname, - "../../node_modules/@posthog/agent/dist/pi/rpc-host.js", - ), - join(__dirname, "../../packages/agent/dist/pi/rpc-host.js"), + const assets = [ + { + name: "Pi RPC host", + source: "pi/rpc-host.js", + destination: "rpc-host.js", + }, + { + name: "Pi local-tools MCP server", + source: "adapters/local-tools/mcp-server.js", + destination: "mcp-server.js", + }, ]; - const source = candidates.find((candidate) => existsSync(candidate)); - if (!source) { - throw new Error( - `[copy-pi-rpc-host] Unable to find Pi RPC host, required at runtime by createPiRpcClient. Build @posthog/agent first. Checked:\n ${candidates.join("\n ")}`, - ); + + for (const asset of assets) { + const candidates = [ + join(__dirname, `node_modules/@posthog/agent/dist/${asset.source}`), + join( + __dirname, + `../../node_modules/@posthog/agent/dist/${asset.source}`, + ), + join(__dirname, `../../packages/agent/dist/${asset.source}`), + ]; + const source = candidates.find((candidate) => existsSync(candidate)); + if (!source) { + throw new Error( + `[copy-pi-rpc-host] Unable to find ${asset.name}. Build @posthog/agent first. Checked:\n ${candidates.join("\n ")}`, + ); + } + copyFileSync(source, join(__dirname, ".vite/build", asset.destination)); } - copyFileSync(source, join(__dirname, ".vite/build/rpc-host.js")); }, }; } diff --git a/packages/agent/build/verify-local-tools-mcp-server.mjs b/packages/agent/build/verify-local-tools-mcp-server.mjs index 89fafa1e0e..fb760e7050 100644 --- a/packages/agent/build/verify-local-tools-mcp-server.mjs +++ b/packages/agent/build/verify-local-tools-mcp-server.mjs @@ -1,5 +1,5 @@ // Post-build smoke check: boot the bundled local-tools MCP server exactly the -// way Codex spawns it in a cloud sandbox (a bare `node dist/...js` process) and +// way an agent runtime spawns it in a cloud sandbox (a bare `node dist/...js` process) and // assert it answers `tools/list`. Catches bundling regressions — e.g. inlined // CJS deps whose dynamic `require()` throws in ESM output — that unit tests // running from src can never see, and that otherwise fail silently in cloud @@ -10,7 +10,7 @@ import { fileURLToPath } from "node:url"; const script = resolve( fileURLToPath(new URL(".", import.meta.url)), - "../dist/adapters/codex-app-server/local-tools-mcp-server.js", + "../dist/adapters/local-tools/mcp-server.js", ); const ctx = Buffer.from( @@ -20,7 +20,7 @@ const ctx = Buffer.from( const child = spawn(process.execPath, [script], { env: { ...process.env, - // Mirror the production Codex spawn: if this ever runs from an + // Mirror the production child spawn: if this ever runs from an // Electron-hosted process, execPath is the app binary, not node. ELECTRON_RUN_AS_NODE: "1", POSTHOG_LOCAL_TOOLS_CTX: ctx, diff --git a/packages/agent/src/adapters/codex-app-server/codex-app-server-agent.ts b/packages/agent/src/adapters/codex-app-server/codex-app-server-agent.ts index 25e54abe5b..9508a74fbd 100644 --- a/packages/agent/src/adapters/codex-app-server/codex-app-server-agent.ts +++ b/packages/agent/src/adapters/codex-app-server/codex-app-server-agent.ts @@ -53,6 +53,10 @@ import { estimateTokens, } from "../claude/context-breakdown"; import { isLocalSkillCommandChunk } from "../local-skill"; +import { + buildLocalToolsServer, + type LocalToolsMeta, +} from "../local-tools/mcp-server-config"; import { resolveSpokenNarration } from "../session-meta"; import { AppServerClient, @@ -67,7 +71,6 @@ import { buildUsageBreakdownParams, } from "./ext-notifications"; import { type CodexUserInput, toCodexInput } from "./input"; -import { buildLocalToolsServer, type LocalToolsMeta } from "./local-tools-mcp"; import { type AppServerItem, changePaths, diff --git a/packages/agent/src/adapters/codex-app-server/local-tools-mcp.test.ts b/packages/agent/src/adapters/local-tools/mcp-server-config.test.ts similarity index 97% rename from packages/agent/src/adapters/codex-app-server/local-tools-mcp.test.ts rename to packages/agent/src/adapters/local-tools/mcp-server-config.test.ts index 0b53b4ceac..e89170e39f 100644 --- a/packages/agent/src/adapters/codex-app-server/local-tools-mcp.test.ts +++ b/packages/agent/src/adapters/local-tools/mcp-server-config.test.ts @@ -1,6 +1,6 @@ import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { LOCAL_TOOLS_MCP_NAME } from "../local-tools"; -import { buildLocalToolsServer } from "./local-tools-mcp"; +import { buildLocalToolsServer } from "./mcp-server-config"; // The dist asset isn't on the walk-up path in unit tests, so make existsSync // succeed; nothing spawns the script — we only inspect the path. @@ -55,7 +55,7 @@ describe("buildLocalToolsServer", () => { expect(server?.name).toBe(LOCAL_TOOLS_MCP_NAME); expect(server?.command).toBe(process.execPath); expect(server?.args).toHaveLength(1); - expect(server?.args[0]).toMatch(/local-tools-mcp-server\.js$/); + expect(server?.args[0]).toMatch(/mcp-server\.js$/); const envNames = server?.env.map((e) => e.name) ?? []; expect(envNames).toContain("POSTHOG_LOCAL_TOOLS_CTX"); diff --git a/packages/agent/src/adapters/codex-app-server/local-tools-mcp.ts b/packages/agent/src/adapters/local-tools/mcp-server-config.ts similarity index 88% rename from packages/agent/src/adapters/codex-app-server/local-tools-mcp.ts rename to packages/agent/src/adapters/local-tools/mcp-server-config.ts index b46d5102c2..c06ba32ce3 100644 --- a/packages/agent/src/adapters/codex-app-server/local-tools-mcp.ts +++ b/packages/agent/src/adapters/local-tools/mcp-server-config.ts @@ -1,21 +1,19 @@ /** - * Builds the stdio local-tools MCP server config to inject into a Codex - * app-server thread's `config.mcp_servers`. - * Returns the ACP `McpServerStdio` shape so the existing translation layer stays - * the single owner of the ACP→Codex map. + * Builds the stdio local-tools MCP server config for agent runtimes. + * Returns the ACP `McpServerStdio` shape used by the Codex adapter and Pi. */ import type { McpServerStdio } from "@agentclientprotocol/sdk"; import { ghTokenEnv } from "@posthog/git/signed-commit"; import { resolveGithubToken } from "../../utils/github-token"; import { resolveBundledMcpScript } from "../../utils/resolve-bundled-script"; +import { resolveTaskId } from "../session-meta"; import { enabledLocalTools, LOCAL_TOOLS_MCP_NAME, type LocalToolCtx, type LocalToolGateMeta, -} from "../local-tools"; -import { resolveTaskId } from "../session-meta"; +} from "./index"; /** * Gate inputs the local-tools server needs beyond `LocalToolGateMeta`: the task id @@ -34,7 +32,7 @@ function toMcpServerStdio( enabledNames: string[], ): McpServerStdio { const scriptPath = resolveBundledMcpScript( - "adapters/codex-app-server/local-tools-mcp-server.js", + "adapters/local-tools/mcp-server.js", ); const ctxBase64 = Buffer.from(JSON.stringify(ctx)).toString("base64"); const env = [ diff --git a/packages/agent/src/adapters/local-tools/mcp-server.integration.test.ts b/packages/agent/src/adapters/local-tools/mcp-server.integration.test.ts new file mode 100644 index 0000000000..6cb2583e6e --- /dev/null +++ b/packages/agent/src/adapters/local-tools/mcp-server.integration.test.ts @@ -0,0 +1,58 @@ +import { mkdtemp, rm } from "node:fs/promises"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { Client } from "@modelcontextprotocol/sdk/client/index.js"; +import { StdioClientTransport } from "@modelcontextprotocol/sdk/client/stdio.js"; +import { afterEach, describe, expect, it } from "vitest"; + +const directories: string[] = []; + +async function temporaryDirectory(): Promise { + const directory = await mkdtemp(join(tmpdir(), "local-tools-mcp-server-")); + directories.push(directory); + return directory; +} + +afterEach(async () => { + await Promise.all( + directories + .splice(0) + .map((directory) => rm(directory, { force: true, recursive: true })), + ); +}); + +describe("local-tools MCP server", () => { + it("starts the bundled server and executes an enabled tool", async () => { + const cwd = await temporaryDirectory(); + const client = new Client({ name: "local-tools-test", version: "1.0.0" }); + const transport = new StdioClientTransport({ + command: process.execPath, + args: [ + join(process.cwd(), "../../node_modules/tsx/dist/cli.mjs"), + join(process.cwd(), "src/adapters/local-tools/mcp-server.ts"), + ], + cwd, + env: { + PATH: process.env.PATH ?? "", + POSTHOG_LOCAL_TOOLS_CTX: Buffer.from( + JSON.stringify({ cwd, taskId: "task-1", taskRunId: "run-1" }), + ).toString("base64"), + POSTHOG_LOCAL_TOOLS_ENABLED: "speak", + }, + }); + + try { + await client.connect(transport); + const tools = await client.listTools(); + const result = await client.callTool({ + name: "speak", + arguments: { text: "Completed the task", kind: "done" }, + }); + + expect(tools.tools.map((tool) => tool.name)).toEqual(["speak"]); + expect(result.content).toEqual([{ type: "text", text: "ok" }]); + } finally { + await client.close(); + } + }); +}); diff --git a/packages/agent/src/adapters/codex-app-server/local-tools-mcp-server.ts b/packages/agent/src/adapters/local-tools/mcp-server.ts similarity index 77% rename from packages/agent/src/adapters/codex-app-server/local-tools-mcp-server.ts rename to packages/agent/src/adapters/local-tools/mcp-server.ts index 92faafd74b..c8553f4660 100644 --- a/packages/agent/src/adapters/codex-app-server/local-tools-mcp-server.ts +++ b/packages/agent/src/adapters/local-tools/mcp-server.ts @@ -1,25 +1,18 @@ /** - * Standalone stdio MCP server exposing the general local tools to the Codex - * app-server adapter, which spawns it as an MCP server process. Reads its context - * (cwd, taskId, token) from POSTHOG_LOCAL_TOOLS_CTX and the set of tools to - * register from POSTHOG_LOCAL_TOOLS_ENABLED (both set by the parent, which has - * already evaluated each tool's gate) — then registers those registry tools, - * the same ones the Claude adapter exposes in-process. + * Standalone stdio MCP server exposing the general local tools to agent runtimes. + * It reads its context from POSTHOG_LOCAL_TOOLS_CTX and its enabled tool set from + * POSTHOG_LOCAL_TOOLS_ENABLED, then registers the corresponding registry tools. * * Usage: * POSTHOG_LOCAL_TOOLS_CTX= \ * POSTHOG_LOCAL_TOOLS_ENABLED=git_signed_commit \ - * node local-tools-mcp-server.js + * node mcp-server.js */ import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js"; import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js"; import { readGithubTokenFromEnv } from "@posthog/git/signed-commit"; -import { - LOCAL_TOOLS, - LOCAL_TOOLS_MCP_NAME, - type LocalToolCtx, -} from "../local-tools"; +import { LOCAL_TOOLS, LOCAL_TOOLS_MCP_NAME, type LocalToolCtx } from "./index"; function die(message: string): never { process.stderr.write(`[local-tools-mcp-server] ${message}\n`); diff --git a/packages/agent/src/pi/rpc-client.test.ts b/packages/agent/src/pi/rpc-client.test.ts index eb0abae742..54492ba621 100644 --- a/packages/agent/src/pi/rpc-client.test.ts +++ b/packages/agent/src/pi/rpc-client.test.ts @@ -3,6 +3,13 @@ import { tmpdir } from "node:os"; import { join } from "node:path"; import { RpcClient } from "@earendil-works/pi-coding-agent"; import { describe, expect, it, vi } from "vitest"; + +const buildLocalToolsServer = vi.hoisted(() => vi.fn()); + +vi.mock("../adapters/local-tools/mcp-server-config", () => ({ + buildLocalToolsServer, +})); + import { createPiRpcClient } from "./rpc-client"; describe("createPiRpcClient", () => { @@ -10,6 +17,7 @@ describe("createPiRpcClient", () => { const client = createPiRpcClient({ cwd: "/workspace", model: "claude-opus-4-8", + capabilities: { environment: "local" }, providerOptions: { region: "us", baseUrl: "http://127.0.0.1:1234", @@ -48,6 +56,7 @@ process.stdin.resume(); const client = createPiRpcClient({ cliPath: hostPath, cwd: directory, + capabilities: { environment: "local" }, providerOptions: { apiKey: "proxy-key" }, }); @@ -62,6 +71,51 @@ process.stdin.resume(); } }); + it("passes local MCP configuration through the private bootstrap pipe", async () => { + const directory = await mkdtemp(join(tmpdir(), "pi-local-mcp-bootstrap-")); + const hostPath = join(directory, "host.mjs"); + const capturePath = join(directory, "capture.json"); + await writeFile( + hostPath, + ` +import { readFileSync, writeFileSync } from "node:fs"; + +const bootstrap = JSON.parse(readFileSync(3, "utf8")); +writeFileSync(${JSON.stringify(capturePath)}, JSON.stringify(bootstrap)); +process.stdin.resume(); +`, + ); + buildLocalToolsServer.mockReturnValueOnce({ + name: "posthog-code-tools", + command: "node", + args: ["local-tools.js"], + env: [{ name: "POSTHOG_LOCAL_TOOLS_ENABLED", value: "upload_artifact" }], + }); + const client = createPiRpcClient({ + cliPath: hostPath, + cwd: directory, + capabilities: { + environment: "cloud", + taskId: "task-1", + taskRunId: "run-1", + }, + providerOptions: { apiKey: "proxy-key" }, + }); + + try { + await client.start(); + await vi.waitFor(async () => { + await expect(readFile(capturePath, "utf8")).resolves.not.toBe(""); + }); + await expect(readFile(capturePath, "utf8")).resolves.toContain( + '"name":"posthog-code-tools"', + ); + } finally { + await client.stop(); + await rm(directory, { recursive: true }); + } + }); + it("uses the private host channel without changing Pi RPC", async () => { const directory = await mkdtemp(join(tmpdir(), "pi-host-channel-")); const hostPath = join(directory, "host.mjs"); @@ -83,6 +137,7 @@ process.on("message", (request) => { const client = createPiRpcClient({ cliPath: hostPath, cwd: directory, + capabilities: { environment: "local" }, providerOptions: { apiKey: "proxy-key" }, }); diff --git a/packages/agent/src/pi/rpc-client.ts b/packages/agent/src/pi/rpc-client.ts index 757ac8fca9..5d11a84507 100644 --- a/packages/agent/src/pi/rpc-client.ts +++ b/packages/agent/src/pi/rpc-client.ts @@ -7,6 +7,10 @@ import { RpcClient, type RpcClientOptions, } from "@earendil-works/pi-coding-agent"; +import { + buildLocalToolsServer, + type LocalToolsMeta, +} from "../adapters/local-tools/mcp-server-config"; import { safePiEnvironment } from "./rpc-environment"; import type { PiQueueSnapshot } from "./types"; @@ -21,6 +25,35 @@ export interface PiRpcProviderOptions { baseUrl?: string; } +export interface PiLocalMcpServer { + name: string; + command: string; + args: string[]; + env: Record; +} + +export interface PiRuntimeCapabilities extends LocalToolsMeta { + environment: "local" | "cloud"; +} + +function createPiLocalMcpServer( + cwd: string, + capabilities: PiRuntimeCapabilities, +): PiLocalMcpServer | undefined { + const server = buildLocalToolsServer({ cwd }, capabilities); + if (!server) { + return undefined; + } + return { + name: server.name, + command: server.command, + args: server.args, + env: Object.fromEntries( + (server.env ?? []).map(({ name, value }) => [name, value]), + ), + }; +} + type RpcClientProcessAccess = { process?: ChildProcess; }; @@ -84,6 +117,7 @@ class SecurePiRpcClient extends RpcClient { constructor( private readonly secureOptions: RpcClientOptions, private readonly providerOptions: PiRpcProviderOptions, + private readonly localMcpServer: PiLocalMcpServer | undefined, ) { super(secureOptions); } @@ -162,7 +196,10 @@ class SecurePiRpcClient extends RpcClient { const bootstrapPipe = child.stdio[3] as Writable | null; bootstrapPipe?.on("error", () => {}); bootstrapPipe?.end( - JSON.stringify({ providerOptions: this.providerOptions }), + JSON.stringify({ + providerOptions: this.providerOptions, + localMcpServer: this.localMcpServer, + }), ); await new Promise((resolve) => setTimeout(resolve, 100)); @@ -265,10 +302,15 @@ export type PiRpcClientOptions = Pick< > & { sessionFile?: string; providerOptions: PiRpcProviderOptions; + capabilities: PiRuntimeCapabilities; }; export function createPiRpcClient(options: PiRpcClientOptions): PiRpcClient { - const { sessionFile, providerOptions, ...rpcOptions } = options; + const { sessionFile, providerOptions, capabilities, ...rpcOptions } = options; + const localMcpServer = createPiLocalMcpServer( + rpcOptions.cwd ?? process.cwd(), + capabilities, + ); const args = sessionFile ? ["--session-file", sessionFile] : []; const cliPath = rpcOptions.cliPath ?? @@ -281,5 +323,6 @@ export function createPiRpcClient(options: PiRpcClientOptions): PiRpcClient { provider: "posthog", }, providerOptions, + localMcpServer, ); } diff --git a/packages/agent/src/pi/rpc-host.ts b/packages/agent/src/pi/rpc-host.ts index 9516177b4e..a1c0befe57 100644 --- a/packages/agent/src/pi/rpc-host.ts +++ b/packages/agent/src/pi/rpc-host.ts @@ -6,10 +6,12 @@ import { POSTHOG_PI_QUEUE_ENTRY_TYPE, readPersistedPiQueue, } from "./queue-persistence"; +import type { PiLocalMcpServer } from "./rpc-client"; import { sanitizePiHostEnvironment } from "./rpc-environment"; interface PiRpcBootstrap { providerOptions?: PosthogProviderOptions; + localMcpServer?: PiLocalMcpServer; } interface PiHostRequest { @@ -25,6 +27,7 @@ function argumentValue(name: string): string | undefined { const bootstrap = JSON.parse(readFileSync(3, "utf8")) as PiRpcBootstrap; const providerOptions = bootstrap.providerOptions; +const localMcpServer = bootstrap.localMcpServer; if (!providerOptions?.apiKey) { throw new Error("Pi RPC host requires PostHog provider credentials"); } @@ -38,6 +41,17 @@ const sessionManager = sessionFile const runtime = await createHarnessRuntime({ cwd, sessionManager, + runtimeMcpServers: localMcpServer + ? { + [localMcpServer.name]: { + command: localMcpServer.command, + args: localMcpServer.args, + env: localMcpServer.env, + lifecycle: "eager", + directTools: true, + }, + } + : undefined, ...providerOptions, }); diff --git a/packages/agent/src/server/pi-agent-server.test.ts b/packages/agent/src/server/pi-agent-server.test.ts index 21c2463223..26438218c9 100644 --- a/packages/agent/src/server/pi-agent-server.test.ts +++ b/packages/agent/src/server/pi-agent-server.test.ts @@ -2,6 +2,14 @@ import { mkdtemp, readFile, rm, writeFile } from "node:fs/promises"; import { tmpdir } from "node:os"; import { join } from "node:path"; import { describe, expect, it, vi } from "vitest"; + +const createPiRpcClient = vi.hoisted(() => vi.fn()); + +vi.mock("../pi/rpc-client", async (importOriginal) => ({ + ...(await importOriginal()), + createPiRpcClient, +})); + import { PiAgentServer } from "./pi-agent-server"; import type { AgentServerConfig } from "./types"; @@ -22,6 +30,44 @@ function config(overrides: Partial = {}): AgentServerConfig { } describe("PiAgentServer", () => { + it("passes enabled local tools into the Pi runtime", async () => { + const client = { + onEvent: vi.fn(), + start: vi.fn(async () => {}), + getState: vi.fn(async () => ({ sessionFile: "/tmp/session.jsonl" })), + }; + createPiRpcClient.mockReturnValue(client); + const server = new PiAgentServer( + config({ baseBranch: "main" }), + ) as unknown as { + posthogAPI: Record; + createSession(payload: Record): Promise; + }; + server.posthogAPI = { + getTaskSession: vi.fn(async () => ({ + id: "session-storage-1", + content_sha256: null, + })), + downloadTaskSession: vi.fn(async () => ""), + updateTaskRun: vi.fn(async () => ({})), + appendTaskRunLog: vi.fn(async () => ({})), + }; + + await server.createSession({ task_id: "task-1", run_id: "run-1" }); + + expect(createPiRpcClient).toHaveBeenCalledWith( + expect.objectContaining({ + capabilities: { + environment: "cloud", + taskId: "task-1", + taskRunId: "run-1", + baseBranch: "main", + background: false, + }, + }), + ); + }); + it("logs session initialization diagnostics when setup fails", async () => { const payload = { task_id: "task-1", run_id: "run-1" }; const server = new PiAgentServer( diff --git a/packages/agent/src/server/pi-agent-server.ts b/packages/agent/src/server/pi-agent-server.ts index 0fb220e900..994299008c 100644 --- a/packages/agent/src/server/pi-agent-server.ts +++ b/packages/agent/src/server/pi-agent-server.ts @@ -493,6 +493,13 @@ export class PiAgentServer { cwd, model: this.config.model, sessionFile: restoredSessionFile, + capabilities: { + environment: "cloud", + taskId: payload.task_id, + taskRunId: payload.run_id, + baseBranch: this.config.baseBranch, + background: this.config.mode === "background", + }, providerOptions: { apiKey: this.config.apiKey, baseUrl: resolveLlmGatewayUrl( diff --git a/packages/agent/src/utils/resolve-bundled-script.ts b/packages/agent/src/utils/resolve-bundled-script.ts index c1c4fa36a8..eb40be01b3 100644 --- a/packages/agent/src/utils/resolve-bundled-script.ts +++ b/packages/agent/src/utils/resolve-bundled-script.ts @@ -1,5 +1,5 @@ import { existsSync } from "node:fs"; -import { resolve as resolvePath } from "node:path"; +import { basename, resolve as resolvePath } from "node:path"; /** * Resolve a shared dist asset relative to the compiled adapter location. When @@ -11,8 +11,9 @@ import { resolve as resolvePath } from "node:path"; export function resolveBundledMcpScript(rel: string): string { let dir = import.meta.dirname ?? __dirname; for (let i = 0; i < 5; i++) { - const candidate = resolvePath(dir, rel); - if (existsSync(candidate)) return candidate; + const candidates = [resolvePath(dir, rel), resolvePath(dir, basename(rel))]; + const candidate = candidates.find((path) => existsSync(path)); + if (candidate) return candidate; dir = resolvePath(dir, ".."); } throw new Error( diff --git a/packages/agent/tsup.config.ts b/packages/agent/tsup.config.ts index 55bfd03fea..04c6af4997 100644 --- a/packages/agent/tsup.config.ts +++ b/packages/agent/tsup.config.ts @@ -144,7 +144,7 @@ export default defineConfig([ "src/adapters/claude/session/mcp-config.ts", "src/adapters/claude/session/models.ts", "src/adapters/codex-app-server/models.ts", - "src/adapters/codex-app-server/local-tools-mcp-server.ts", + "src/adapters/local-tools/mcp-server.ts", "src/adapters/claude/mcp/tool-metadata.ts", "src/adapters/reasoning-effort.ts", "src/execution-mode.ts", @@ -157,7 +157,7 @@ export default defineConfig([ 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 - // exists. Entries spawned directly by node (local-tools-mcp-server.js) + // exists. Entries spawned directly by node (mcp-server.js) // crash at import time without this shim. banner: { js: 'import { createRequire as __createRequire } from "node:module"; const require = __createRequire(import.meta.url);', diff --git a/packages/harness/src/extensions/mcp/config.test.ts b/packages/harness/src/extensions/mcp/config.test.ts index 80889712dd..1a36dea746 100644 --- a/packages/harness/src/extensions/mcp/config.test.ts +++ b/packages/harness/src/extensions/mcp/config.test.ts @@ -2,7 +2,12 @@ import { mkdtemp, rm, writeFile } from "node:fs/promises"; import { tmpdir } from "node:os"; import { join } from "node:path"; import { afterEach, beforeEach, describe, expect, it } from "vitest"; -import { loadConfig, mergeRawConfigs, parseConfig } from "./config"; +import { + loadConfig, + mergeRawConfigs, + mergeRuntimeServers, + parseConfig, +} from "./config"; import { McpError } from "./errors"; describe("parseConfig", () => { @@ -106,6 +111,38 @@ describe("mergeRawConfigs", () => { }); }); +describe("mergeRuntimeServers", () => { + it("adds first-party runtime servers after file-configured servers", () => { + const config = parseConfig( + { + mcpServers: { + "posthog-code-tools": { command: "untrusted-command" }, + fileOnly: { command: "node" }, + }, + }, + "test", + ); + + const merged = mergeRuntimeServers(config, { + "posthog-code-tools": { + command: "node", + args: ["local-tools.js"], + env: { POSTHOG_LOCAL_TOOLS_ENABLED: "upload_artifact" }, + lifecycle: "eager", + directTools: true, + }, + }); + + expect(merged.mcpServers["posthog-code-tools"]).toMatchObject({ + command: "node", + args: ["local-tools.js"], + lifecycle: "eager", + directTools: true, + }); + expect(merged.mcpServers.fileOnly?.command).toBe("node"); + }); +}); + describe("loadConfig", () => { let dir: string; let globalPath: string; diff --git a/packages/harness/src/extensions/mcp/config.ts b/packages/harness/src/extensions/mcp/config.ts index 93102b3ace..a96cbb553c 100644 --- a/packages/harness/src/extensions/mcp/config.ts +++ b/packages/harness/src/extensions/mcp/config.ts @@ -139,6 +139,8 @@ const mcpConfigSchema = z.object({ }); export type McpAuthConfig = z.output; +export type McpServerInput = z.input; +export type RuntimeMcpServers = Record; export type McpServerConfig = z.output; export type McpSettings = z.output; export type McpConfig = z.output; @@ -147,6 +149,19 @@ export function emptyConfig(): McpConfig { return mcpConfigSchema.parse({}); } +export function mergeRuntimeServers( + config: McpConfig, + servers: RuntimeMcpServers, +): McpConfig { + return parseConfig( + { + settings: config.settings, + mcpServers: { ...config.mcpServers, ...servers }, + }, + "runtime MCP servers", + ); +} + async function readJsonFile(path: string): Promise { let text: string; try { diff --git a/packages/harness/src/extensions/mcp/extension.ts b/packages/harness/src/extensions/mcp/extension.ts index 70f149590c..c34423f91d 100644 --- a/packages/harness/src/extensions/mcp/extension.ts +++ b/packages/harness/src/extensions/mcp/extension.ts @@ -29,8 +29,8 @@ import type { OAuthFlowRunner } from "./auth-flow"; import { runOAuthFlow } from "./auth-flow"; import { McpAuthStorage } from "./auth-storage"; import { CallbackServer } from "./callback-server"; -import type { ConfigLoader, McpConfig } from "./config"; -import { emptyConfig, loadConfig } from "./config"; +import type { ConfigLoader, McpConfig, RuntimeMcpServers } from "./config"; +import { emptyConfig, loadConfig, mergeRuntimeServers } from "./config"; import { describeError } from "./errors"; import { createMcpProxyTool } from "./proxy-tool"; import type { TransportFactory } from "./server-manager"; @@ -49,6 +49,7 @@ export interface McpExtensionOptions { toolCache?: McpToolCache; /** Override the interactive OAuth flow (tests). */ oauthFlow?: OAuthFlowRunner; + runtimeServers?: RuntimeMcpServers; } /** Open a URL in the user's default browser (best effort). Exported for tests. */ @@ -228,9 +229,13 @@ export function createMcpExtension( pi.on("session_start", async (_event, ctx) => { let nextConfig: McpConfig; try { - nextConfig = await configLoader(ctx.cwd, { + const fileConfig = await configLoader(ctx.cwd, { includeProject: ctx.isProjectTrusted(), }); + nextConfig = mergeRuntimeServers( + fileConfig, + options.runtimeServers ?? {}, + ); } catch (err) { if (ctx.hasUI) { ctx.ui.notify(`mcp: config error — ${describeError(err)}`, "error"); diff --git a/packages/harness/src/extensions/mcp/runtime-servers.integration.test.ts b/packages/harness/src/extensions/mcp/runtime-servers.integration.test.ts new file mode 100644 index 0000000000..31b8bf8379 --- /dev/null +++ b/packages/harness/src/extensions/mcp/runtime-servers.integration.test.ts @@ -0,0 +1,134 @@ +import { mkdtemp, rm } from "node:fs/promises"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { InMemoryCredentialStore } from "@earendil-works/pi-ai"; +import { + fauxAssistantMessage, + fauxToolCall, + registerFauxProvider, +} from "@earendil-works/pi-ai/compat"; +import { + createAgentSessionFromServices, + createAgentSessionServices, + type ExtensionAPI, + ModelRuntime, + SessionManager, +} from "@earendil-works/pi-coding-agent"; +import { afterEach, describe, expect, it, vi } from "vitest"; +import { parseConfig } from "./config"; +import { createMcpExtension } from "./extension"; +import { createMockMcpServer } from "./test-support"; + +const directories: string[] = []; + +async function temporaryDirectory(): Promise { + const directory = await mkdtemp(join(tmpdir(), "mcp-runtime-server-")); + directories.push(directory); + return directory; +} + +afterEach(async () => { + await Promise.all( + directories + .splice(0) + .map((directory) => rm(directory, { force: true, recursive: true })), + ); +}); + +describe("runtime MCP servers", () => { + it("lets a Pi session call a first-party runtime tool", async () => { + const cwd = await temporaryDirectory(); + const agentDir = await temporaryDirectory(); + const called = vi.fn(); + const mcpServer = createMockMcpServer([ + { + name: "echo", + description: "Echo text back", + inputSchema: { + type: "object", + properties: { text: { type: "string" } }, + required: ["text"], + }, + handler: (args) => { + called(args); + return { content: [{ type: "text", text: `echo: ${args.text}` }] }; + }, + }, + ]); + const faux = registerFauxProvider(); + const model = faux.getModel(); + faux.setResponses([ + fauxAssistantMessage( + [ + fauxToolCall("mcp_posthog_code_tools_echo", { + text: "hello", + }), + ], + { stopReason: "toolUse" }, + ), + fauxAssistantMessage("finished"), + ]); + + const credentials = new InMemoryCredentialStore(); + await credentials.modify(model.provider, async () => ({ + type: "api_key", + key: "faux-key", + })); + const modelRuntime = await ModelRuntime.create({ credentials }); + const services = await createAgentSessionServices({ + agentDir, + modelRuntime, + cwd, + resourceLoaderOptions: { + extensionFactories: [ + { + name: "faux-provider", + factory(pi: ExtensionAPI) { + pi.registerProvider(model.provider, { + baseUrl: model.baseUrl, + apiKey: "faux-key", + api: faux.api, + models: faux.models, + }); + }, + }, + { + name: "mcp", + factory: createMcpExtension({ + configLoader: async () => parseConfig({}, "test"), + runtimeServers: { + "posthog-code-tools": { + command: "node", + args: ["local-tools.js"], + lifecycle: "eager", + directTools: true, + }, + }, + transportFactory: mcpServer.transportFactory, + }), + }, + ], + noPromptTemplates: true, + noSkills: true, + noThemes: true, + }, + }); + const { session } = await createAgentSessionFromServices({ + services, + sessionManager: SessionManager.inMemory(cwd), + model, + }); + + try { + await session.bindExtensions({}); + await session.prompt("Call the first-party echo tool"); + + expect(called).toHaveBeenCalledWith({ text: "hello" }); + expect(faux.state.callCount).toBe(2); + } finally { + session.dispose(); + faux.unregister(); + await mcpServer.close(); + } + }); +}); diff --git a/packages/harness/src/extensions/registry.ts b/packages/harness/src/extensions/registry.ts index ac813b00c4..a544c540a2 100644 --- a/packages/harness/src/extensions/registry.ts +++ b/packages/harness/src/extensions/registry.ts @@ -5,13 +5,16 @@ import type { } from "@earendil-works/pi-coding-agent"; import type { HogBrandingOptions } from "./hog-branding/extension"; import { createHogBrandingExtension } from "./hog-branding/extension"; +import type { RuntimeMcpServers } from "./mcp/config"; import { createMcpExtension } from "./mcp/extension"; import { createPosthogProviderExtension } from "./posthog-provider/extension"; import type { PosthogProviderOptions } from "./posthog-provider/provider"; import { createWebAccessExtension } from "./web-access/extension"; export type HarnessExtensionOptions = PosthogProviderOptions & - HogBrandingOptions; + HogBrandingOptions & { + runtimeMcpServers?: RuntimeMcpServers; + }; interface HarnessExtension { name: string; @@ -22,9 +25,11 @@ const EXTENSIONS: HarnessExtension[] = [ { name: "hog-branding", create: createHogBrandingExtension }, { name: "posthog-provider", create: createPosthogProviderExtension }, { name: "web-access", create: createWebAccessExtension }, - // createMcpExtension's options are test seams (config loader, transport - // factory), not HarnessExtensionOptions, so drop the registry options. - { name: "mcp", create: () => createMcpExtension() }, + { + name: "mcp", + create: (options) => + createMcpExtension({ runtimeServers: options.runtimeMcpServers }), + }, ]; export const HARNESS_EXTENSION_NAMES: readonly string[] = EXTENSIONS.map( diff --git a/packages/workspace-server/src/services/pi-session/identifiers.ts b/packages/workspace-server/src/services/pi-session/identifiers.ts index b9a978124c..b1061315c4 100644 --- a/packages/workspace-server/src/services/pi-session/identifiers.ts +++ b/packages/workspace-server/src/services/pi-session/identifiers.ts @@ -2,7 +2,6 @@ import type { PiRpcClient, PiRpcClientOptions, } from "@posthog/agent/pi/rpc-client"; -import type { PiRuntime } from "@posthog/agent/pi/runtime"; export interface PiRpcClientFactory { create( @@ -14,18 +13,6 @@ export const PI_RPC_CLIENT_FACTORY = Symbol.for( "posthog.workspace.piRpcClientFactory", ); -export interface PiRuntimeFactory { - create(input: { - cwd: string; - model?: string; - sessionFile?: string; - }): Promise; -} - -export const PI_RUNTIME_FACTORY = Symbol.for( - "posthog.workspace.piRuntimeFactory", -); - export const PI_SESSION_SERVICE = Symbol.for( "posthog.workspace.piSessionService", ); diff --git a/packages/workspace-server/src/services/pi-session/pi-session.test.ts b/packages/workspace-server/src/services/pi-session/pi-session.test.ts index 9195c8e53d..6ee077f2c6 100644 --- a/packages/workspace-server/src/services/pi-session/pi-session.test.ts +++ b/packages/workspace-server/src/services/pi-session/pi-session.test.ts @@ -1,11 +1,10 @@ import type { PiRpcClient } from "@posthog/agent/pi/rpc-client"; -import type { RpcCommand, RpcResponse } from "@posthog/agent/pi/rpc-transport"; -import type { PiRuntime } from "@posthog/agent/pi/runtime"; +import type { RpcResponse } from "@posthog/agent/pi/rpc-transport"; import type { RootLogger } from "@posthog/di/logger"; import { afterEach, describe, expect, it, vi } from "vitest"; import type { ITaskMetadataRepository } from "../../db/repositories/task-metadata-repository"; import type { ProcessTrackingService } from "../process-tracking/process-tracking"; -import type { PiRuntimeFactory } from "./identifiers"; +import type { PiRpcClientFactory } from "./identifiers"; import { PiSessionService, selectPiPoolEvictionCandidate } from "./pi-session"; const scopedLogger = { @@ -158,7 +157,7 @@ describe("PiSessionService task session config", () => { vi.fn(async () => new Response(content)), ); const service = new PiSessionService( - {} as PiRuntimeFactory, + {} as PiRpcClientFactory, {} as ITaskMetadataRepository, {} as ProcessTrackingService, rootLogger, @@ -178,6 +177,7 @@ describe("PiSessionService start", () => { const setThinkingLevel = vi.fn().mockResolvedValue(undefined); const prompt = vi.fn().mockResolvedValue(undefined); const client = { + onEvent: vi.fn(() => () => {}), start: vi.fn().mockResolvedValue(undefined), stop: vi.fn().mockResolvedValue(undefined), getState: vi.fn().mockResolvedValue({ @@ -188,14 +188,9 @@ describe("PiSessionService start", () => { setThinkingLevel, prompt, } as unknown as PiRpcClient; - const runtimeFactory = { - create: vi.fn(async () => ({ - client, - process: undefined, - onRuntimeEvent: vi.fn(), - onConversationEvent: vi.fn(), - })), - } as unknown as PiRuntimeFactory; + const clientFactory = { + create: vi.fn(async () => client), + } as unknown as PiRpcClientFactory; const taskMetadataRepository = { upsert: vi.fn(), } as unknown as ITaskMetadataRepository; @@ -204,7 +199,7 @@ describe("PiSessionService start", () => { unregister: vi.fn(), } as unknown as ProcessTrackingService; const service = new PiSessionService( - runtimeFactory, + clientFactory, taskMetadataRepository, processTracking, rootLogger, @@ -236,6 +231,7 @@ describe("PiSessionService RPC request pinning", () => { followUp: string[]; }) => void = () => {}; const firstClient = { + onEvent: vi.fn(() => () => {}), start: vi.fn().mockResolvedValue(undefined), stop: vi.fn().mockResolvedValue(undefined), getState: vi.fn().mockResolvedValue({ @@ -256,6 +252,7 @@ describe("PiSessionService RPC request pinning", () => { ), } as unknown as PiRpcClient; const secondClient = { + onEvent: vi.fn(() => () => {}), start: vi.fn().mockResolvedValue(undefined), stop: vi.fn().mockResolvedValue(undefined), getState: vi.fn().mockResolvedValue({ @@ -265,6 +262,7 @@ describe("PiSessionService RPC request pinning", () => { send: vi.fn(), } as unknown as PiRpcClient; const thirdClient = { + onEvent: vi.fn(() => () => {}), start: vi.fn().mockResolvedValue(undefined), stop: vi.fn().mockResolvedValue(undefined), getState: vi.fn().mockResolvedValue({ @@ -274,24 +272,9 @@ describe("PiSessionService RPC request pinning", () => { send: vi.fn(), } as unknown as PiRpcClient; const clients = [firstClient, secondClient, thirdClient]; - const runtimeFactory = { - create: vi.fn(async () => { - const client = clients.shift() as PiRpcClient; - return { - client, - process: undefined, - sendCommand: vi.fn((command) => - ( - client as unknown as { - send(command: RpcCommand): Promise; - } - ).send(command), - ), - onRuntimeEvent: vi.fn(), - onConversationEvent: vi.fn(), - } as unknown as PiRuntime; - }), - } as PiRuntimeFactory; + const clientFactory = { + create: vi.fn(async () => clients.shift() as PiRpcClient), + } as PiRpcClientFactory; const taskMetadataRepository = { findByTaskId: vi.fn((taskId: string) => ({ piSessionFile: `/tmp/${taskId}.jsonl`, @@ -303,7 +286,7 @@ describe("PiSessionService RPC request pinning", () => { unregister: vi.fn(), } as unknown as ProcessTrackingService; const service = new PiSessionService( - runtimeFactory, + clientFactory, taskMetadataRepository, processTracking, rootLogger, @@ -311,15 +294,16 @@ describe("PiSessionService RPC request pinning", () => { await service.resume({ taskId: "first", cwd: "/tmp" }); const bashRequest = service.request("first", { - type: "bash", - command: "sleep 1", + id: "request-1", + type: "prompt", + message: "hello", }); const queueRequest = service.getQueue("first"); await service.resume({ taskId: "second", cwd: "/tmp" }); expect(firstClient.stop).not.toHaveBeenCalled(); - requestResolvers[0](successfulResponse("bash")); + requestResolvers[0](successfulResponse("prompt")); await bashRequest; await vi.waitFor(() => expect(secondClient.stop).toHaveBeenCalledOnce()); expect(firstClient.stop).not.toHaveBeenCalled(); diff --git a/packages/workspace-server/src/services/pi-session/pi-session.ts b/packages/workspace-server/src/services/pi-session/pi-session.ts index cdc1f7af15..8b1d5d4742 100644 --- a/packages/workspace-server/src/services/pi-session/pi-session.ts +++ b/packages/workspace-server/src/services/pi-session/pi-session.ts @@ -7,7 +7,7 @@ import { } from "@earendil-works/pi-coding-agent"; import type { PiRpcClient } from "@posthog/agent/pi/rpc-client"; import type { RpcCommand, RpcResponse } from "@posthog/agent/pi/rpc-transport"; -import type { PiRuntime } from "@posthog/agent/pi/runtime"; +import { PiRuntime } from "@posthog/agent/pi/runtime"; import { PI_THINKING_LEVELS, type PiPersistedSessionConfig, @@ -24,7 +24,7 @@ import { TASK_METADATA_REPOSITORY } from "../../db/identifiers"; import type { ITaskMetadataRepository } from "../../db/repositories/task-metadata-repository"; import { PROCESS_TRACKING_SERVICE } from "../process-tracking/identifiers"; import type { ProcessTrackingService } from "../process-tracking/process-tracking"; -import { PI_RUNTIME_FACTORY, type PiRuntimeFactory } from "./identifiers"; +import { PI_RPC_CLIENT_FACTORY, type PiRpcClientFactory } from "./identifiers"; import type { StartPiSessionInput } from "./schemas"; type PiPoolSessionState = "starting" | "idle" | "streaming"; @@ -91,8 +91,8 @@ export class PiSessionService extends TypedEventEmitter { private readonly log: ReturnType; constructor( - @inject(PI_RUNTIME_FACTORY) - private readonly runtimeFactory: PiRuntimeFactory, + @inject(PI_RPC_CLIENT_FACTORY) + private readonly clientFactory: PiRpcClientFactory, @inject(TASK_METADATA_REPOSITORY) private readonly taskMetadataRepository: ITaskMetadataRepository, @inject(PROCESS_TRACKING_SERVICE) @@ -114,7 +114,7 @@ export class PiSessionService extends TypedEventEmitter { ): Promise<{ sessionFile: string | null; sessionId: string }> { await this.stopLocked(input.taskId); - const runtime = await this.runtimeFactory.create({ + const runtime = await this.createRuntime({ cwd: input.cwd, model: input.model, }); @@ -171,7 +171,7 @@ export class PiSessionService extends TypedEventEmitter { await this.stopLocked(input.taskId); - const runtime = await this.runtimeFactory.create({ + const runtime = await this.createRuntime({ cwd: input.cwd, sessionFile, }); @@ -251,6 +251,15 @@ export class PiSessionService extends TypedEventEmitter { await this.runExclusive(taskId, () => this.stopLocked(taskId)); } + private async createRuntime(input: { + cwd: string; + model?: string; + sessionFile?: string; + }): Promise { + const client = await this.clientFactory.create(input); + return new PiRuntime(client); + } + private async stopLocked(taskId: string): Promise { const session = this.sessions.get(taskId);