From e24236da2cc9c7512a1eafb4f21115c3623e027e Mon Sep 17 00:00:00 2001 From: JonathanLab Date: Thu, 30 Jul 2026 13:55:12 +0200 Subject: [PATCH 1/3] 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); From 1b9e3023b50f2587e0a7b79b80c738e6d2305aec Mon Sep 17 00:00:00 2001 From: JonathanLab Date: Thu, 30 Jul 2026 14:18:17 +0200 Subject: [PATCH 2/3] fix(pi): preserve runtime MCP tools --- apps/code/electron-builder.ts | 1 + .../src/extensions/mcp/extension.test.ts | 24 +++++++++++++++++++ .../harness/src/extensions/mcp/extension.ts | 6 ++++- 3 files changed, 30 insertions(+), 1 deletion(-) diff --git a/apps/code/electron-builder.ts b/apps/code/electron-builder.ts index 70bd5102ad..39ceacd959 100644 --- a/apps/code/electron-builder.ts +++ b/apps/code/electron-builder.ts @@ -51,6 +51,7 @@ const config: Configuration = { ".vite/build/grammars/**", ".vite/build/rpc-host.js", ".vite/build/rpc-host.js.map", + ".vite/build/mcp-server.js", ...asarUnpackGlobs, ], diff --git a/packages/harness/src/extensions/mcp/extension.test.ts b/packages/harness/src/extensions/mcp/extension.test.ts index a9818b48ee..6c5dc84d07 100644 --- a/packages/harness/src/extensions/mcp/extension.test.ts +++ b/packages/harness/src/extensions/mcp/extension.test.ts @@ -320,6 +320,30 @@ describe("createMcpExtension", () => { ); }); + it("starts runtime servers when loading file config fails", async () => { + const mock = createMockMcpServer([ECHO_TOOL]); + const { pi, emit, getActive } = fakePi(); + createMcpExtension({ + configLoader: vi.fn().mockRejectedValue(new Error("bad config")), + runtimeServers: { + runtime: { command: "unused", lifecycle: "eager", directTools: true }, + }, + transportFactory: mock.transportFactory, + })(pi); + const { ctx, notify } = fakeCtx(); + + await emit("session_start", { reason: "startup" }, ctx); + + expect(notify).toHaveBeenCalledWith( + expect.stringContaining("bad config"), + "error", + ); + expect(getActive()).toContain("mcp_runtime_echo"); + + await emit("session_shutdown", { reason: "quit" }, ctx); + await mock.close(); + }); + it("notifies when an eager server fails to start", async () => { const mock = createMockMcpServer([]); const { pi, emit } = fakePi(); diff --git a/packages/harness/src/extensions/mcp/extension.ts b/packages/harness/src/extensions/mcp/extension.ts index c34423f91d..650d6e1815 100644 --- a/packages/harness/src/extensions/mcp/extension.ts +++ b/packages/harness/src/extensions/mcp/extension.ts @@ -240,7 +240,11 @@ export function createMcpExtension( if (ctx.hasUI) { ctx.ui.notify(`mcp: config error — ${describeError(err)}`, "error"); } - return; + const runtimeServers = options.runtimeServers ?? {}; + if (Object.keys(runtimeServers).length === 0) { + return; + } + nextConfig = mergeRuntimeServers(emptyConfig(), runtimeServers); } currentCwd = ctx.cwd; From 0573a01a7876057a8d82d2382b860f762e743429 Mon Sep 17 00:00:00 2001 From: JonathanLab Date: Thu, 30 Jul 2026 14:35:15 +0200 Subject: [PATCH 3/3] fix(agent): keep GitHub tokens out of MCP env --- .../adapters/local-tools/mcp-server-config.test.ts | 13 ++++++------- .../src/adapters/local-tools/mcp-server-config.ts | 12 ------------ .../agent/src/adapters/local-tools/mcp-server.ts | 5 ++--- 3 files changed, 8 insertions(+), 22 deletions(-) diff --git a/packages/agent/src/adapters/local-tools/mcp-server-config.test.ts b/packages/agent/src/adapters/local-tools/mcp-server-config.test.ts index e89170e39f..5d524044a4 100644 --- a/packages/agent/src/adapters/local-tools/mcp-server-config.test.ts +++ b/packages/agent/src/adapters/local-tools/mcp-server-config.test.ts @@ -17,8 +17,7 @@ describe("buildLocalToolsServer", () => { }; beforeEach(() => { - // The signed-git gate reads IS_SANDBOX and the token vars; clear them so each - // case controls the cloud signal (meta.environment) and token explicitly. + // Clear credential state so each case controls the cloud signal explicitly. delete process.env.IS_SANDBOX; delete process.env.GH_TOKEN; delete process.env.GITHUB_TOKEN; @@ -60,9 +59,8 @@ describe("buildLocalToolsServer", () => { const envNames = server?.env.map((e) => e.name) ?? []; expect(envNames).toContain("POSTHOG_LOCAL_TOOLS_CTX"); expect(envNames).toContain("POSTHOG_LOCAL_TOOLS_ENABLED"); - // Token is forwarded to the child so its own git remote ops authenticate. - expect(envNames).toContain("GH_TOKEN"); - expect(envNames).toContain("GITHUB_TOKEN"); + expect(envNames).not.toContain("GH_TOKEN"); + expect(envNames).not.toContain("GITHUB_TOKEN"); // Codex strips ELECTRON_RUN_AS_NODE from its own env, and process.execPath // is the app binary in packaged installs; without this the server boots // the full desktop app instead of running the script. @@ -78,13 +76,14 @@ describe("buildLocalToolsServer", () => { Buffer.from(ctxEntry?.value ?? "", "base64").toString("utf-8"), ); expect(ctx.cwd).toBe("/repo"); - expect(ctx.token).toBe("ghs_test"); + expect(ctx.token).toBeUndefined(); expect(ctx.taskId).toBe("task-1"); expect(ctx.taskRunId).toBe("run-1"); expect(ctx.baseBranch).toBe("master"); }); - it("returns a server but omits token env vars when no token is present", () => { + it("does not serialize a parent GitHub token into the child config", () => { + process.env.GH_TOKEN = "ghs_test"; const server = buildLocalToolsServer( { cwd: "/repo" }, { environment: "cloud" }, diff --git a/packages/agent/src/adapters/local-tools/mcp-server-config.ts b/packages/agent/src/adapters/local-tools/mcp-server-config.ts index c06ba32ce3..92f19e08d2 100644 --- a/packages/agent/src/adapters/local-tools/mcp-server-config.ts +++ b/packages/agent/src/adapters/local-tools/mcp-server-config.ts @@ -4,8 +4,6 @@ */ 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 { @@ -43,15 +41,6 @@ function toMcpServerStdio( // binary, which boots the full app without this var. Inert on real node. { name: "ELECTRON_RUN_AS_NODE", value: "1" }, ]; - if (ctx.token) { - // Token also on the child env so its own git remote ops authenticate. - env.push( - ...Object.entries(ghTokenEnv(ctx.token)).map(([name, value]) => ({ - name, - value, - })), - ); - } return { name: LOCAL_TOOLS_MCP_NAME, command: process.execPath, @@ -75,7 +64,6 @@ export function buildLocalToolsServer( } const toolCtx: LocalToolCtx = { cwd, - token: resolveGithubToken(), taskId: resolveTaskId(meta), taskRunId: meta?.taskRunId, baseBranch: meta?.baseBranch, diff --git a/packages/agent/src/adapters/local-tools/mcp-server.ts b/packages/agent/src/adapters/local-tools/mcp-server.ts index c8553f4660..fc871f7b74 100644 --- a/packages/agent/src/adapters/local-tools/mcp-server.ts +++ b/packages/agent/src/adapters/local-tools/mcp-server.ts @@ -11,7 +11,7 @@ import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js"; import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js"; -import { readGithubTokenFromEnv } from "@posthog/git/signed-commit"; +import { resolveGithubToken } from "../../utils/github-token"; import { LOCAL_TOOLS, LOCAL_TOOLS_MCP_NAME, type LocalToolCtx } from "./index"; function die(message: string): never { @@ -28,7 +28,6 @@ let parsed: { cwd: string; taskId?: string; taskRunId?: string; - token?: string; baseBranch?: string; }; try { @@ -43,7 +42,7 @@ if (!parsed.cwd) { const ctx: LocalToolCtx = { cwd: parsed.cwd, - token: parsed.token ?? readGithubTokenFromEnv(), + token: resolveGithubToken(), taskId: parsed.taskId, taskRunId: parsed.taskRunId, baseBranch: parsed.baseBranch,