From 5a5ec6c67653241f9fe08b5b9d714003ccbc8605 Mon Sep 17 00:00:00 2001 From: Chris Volzer Date: Mon, 29 Jun 2026 10:26:45 -0400 Subject: [PATCH] fix(codex): approve MCP tools with PreToolUse hooks --- apps/code/scripts/download-binaries.mjs | 19 +- apps/code/scripts/download-binaries.test.mjs | 2 + packages/agent/package.json | 8 + packages/agent/src/acp-extensions.ts | 6 +- .../claude/claude-agent.refresh.test.ts | 23 +- .../agent/src/adapters/claude/claude-agent.ts | 22 +- .../agent/src/adapters/claude/hooks.test.ts | 4 +- packages/agent/src/adapters/claude/hooks.ts | 5 +- .../permissions/permission-handlers.test.ts | 23 +- .../claude/permissions/permission-handlers.ts | 2 +- .../codex/codex-agent.refresh.test.ts | 19 + .../agent/src/adapters/codex/codex-agent.ts | 32 +- .../src/adapters/codex/codex-client.test.ts | 10 +- .../agent/src/adapters/codex/codex-client.ts | 7 +- .../adapters/codex/mcp-approval-hook.test.ts | 120 +++ .../src/adapters/codex/mcp-approval-hook.ts | 351 +++++++++ .../agent/src/adapters/codex/spawn.test.ts | 24 + packages/agent/src/adapters/codex/spawn.ts | 15 + packages/agent/src/agent.ts | 1 + packages/agent/src/index.ts | 6 +- .../agent/src/server/agent-server.test.ts | 118 +++ packages/agent/src/server/agent-server.ts | 734 +++++++++++++----- .../agent/src/server/question-relay.test.ts | 136 +++- packages/agent/src/server/schemas.test.ts | 12 +- packages/agent/src/server/schemas.ts | 13 +- packages/agent/src/types.ts | 2 + packages/agent/tsup.config.ts | 2 + .../host-router/src/routers/agent.router.ts | 9 + .../hooks/useMcpInstallationTools.ts | 26 +- .../src/services/agent/agent.test.ts | 433 +++++++++++ .../src/services/agent/agent.ts | 437 +++++++++-- .../src/services/agent/codex-home.test.ts | 33 + .../src/services/agent/codex-home.ts | 9 + .../src/services/agent/schemas.ts | 12 + 34 files changed, 2383 insertions(+), 292 deletions(-) create mode 100644 packages/agent/src/adapters/codex/mcp-approval-hook.test.ts create mode 100644 packages/agent/src/adapters/codex/mcp-approval-hook.ts diff --git a/apps/code/scripts/download-binaries.mjs b/apps/code/scripts/download-binaries.mjs index 26a3e10f6a..af33b2b482 100644 --- a/apps/code/scripts/download-binaries.mjs +++ b/apps/code/scripts/download-binaries.mjs @@ -6,8 +6,10 @@ import { createWriteStream, existsSync, mkdirSync, + readFileSync, realpathSync, rmSync, + writeFileSync, } from "node:fs"; import { dirname, join } from "node:path"; import { pipeline } from "node:stream/promises"; @@ -21,7 +23,7 @@ const DEST_DIR = join(__dirname, "..", "resources", "codex-acp"); const BINARIES = [ { name: "codex-acp", - version: "0.14.0", + version: "0.16.0", getUrl: (version, target) => { const ext = target.includes("windows") ? "zip" : "tar.gz"; return `https://github.com/zed-industries/codex-acp/releases/download/v${version}/codex-acp-${version}-${target}.${ext}`; @@ -140,12 +142,22 @@ async function downloadBinary(binary) { const binaryName = process.platform === "win32" ? `${binary.name}.exe` : binary.name; const binaryPath = join(DEST_DIR, binaryName); + const versionPath = join(DEST_DIR, `${binary.name}.version`); console.log(`\n[${binary.name}] v${binary.version}`); if (existsSync(binaryPath)) { - console.log(` Already exists: ${binaryPath}`); - return; + const currentVersion = existsSync(versionPath) + ? readFileSync(versionPath, "utf8").trim() + : null; + if (currentVersion === binary.version) { + console.log(` Already exists: ${binaryPath}`); + return; + } + console.log( + ` Replacing existing ${binary.name}${currentVersion ? ` v${currentVersion}` : ""} with v${binary.version}`, + ); + rmSync(binaryPath, { force: true }); } const target = binary.getTarget(); @@ -170,6 +182,7 @@ async function downloadBinary(binary) { if (process.platform === "darwin") { signForMacOS(binaryPath); } + writeFileSync(versionPath, `${binary.version}\n`); console.log(` Ready: ${binaryPath}`); } diff --git a/apps/code/scripts/download-binaries.test.mjs b/apps/code/scripts/download-binaries.test.mjs index 85e01b8085..3c000c6960 100644 --- a/apps/code/scripts/download-binaries.test.mjs +++ b/apps/code/scripts/download-binaries.test.mjs @@ -16,8 +16,10 @@ vi.mock("node:fs", () => { createWriteStream: vi.fn(() => ({})), existsSync: vi.fn(() => true), mkdirSync: vi.fn(), + readFileSync: vi.fn(() => ""), realpathSync: vi.fn(() => "/not/the/entrypoint"), rmSync: vi.fn(), + writeFileSync: vi.fn(), }; return { ...fns, default: fns }; }); diff --git a/packages/agent/package.json b/packages/agent/package.json index b6f4326c64..b1d801bb15 100644 --- a/packages/agent/package.json +++ b/packages/agent/package.json @@ -20,6 +20,10 @@ "types": "./dist/gateway-models.d.ts", "import": "./dist/gateway-models.js" }, + "./mcp-store/tool-keys": { + "types": "./dist/mcp-store/tool-keys.d.ts", + "import": "./dist/mcp-store/tool-keys.js" + }, "./posthog-api": { "types": "./dist/posthog-api.d.ts", "import": "./dist/posthog-api.js" @@ -68,6 +72,10 @@ "types": "./dist/adapters/claude/mcp/tool-metadata.d.ts", "import": "./dist/adapters/claude/mcp/tool-metadata.js" }, + "./adapters/codex/mcp-approval-hook": { + "types": "./dist/adapters/codex/mcp-approval-hook.d.ts", + "import": "./dist/adapters/codex/mcp-approval-hook.js" + }, "./execution-mode": { "types": "./dist/execution-mode.d.ts", "import": "./dist/execution-mode.js" diff --git a/packages/agent/src/acp-extensions.ts b/packages/agent/src/acp-extensions.ts index 332972e79e..870886aee2 100644 --- a/packages/agent/src/acp-extensions.ts +++ b/packages/agent/src/acp-extensions.ts @@ -89,9 +89,9 @@ export const POSTHOG_NOTIFICATIONS = { export const POSTHOG_METHODS = { /** * Client requests a session refresh between turns. Payload may include - * `mcpServers` to trigger a resume-with-new-options reinit; future fields - * can extend this without adding new methods. Returns once the refresh has - * completed so the caller can safely send the next prompt. + * `mcpServers` to trigger a resume-with-new-options reinit or + * `mcpToolApprovals` to refresh the in-memory approval gate. Returns once + * the refresh has completed so the caller can safely send the next prompt. */ REFRESH_SESSION: "_posthog/refresh_session", } as const; diff --git a/packages/agent/src/adapters/claude/claude-agent.refresh.test.ts b/packages/agent/src/adapters/claude/claude-agent.refresh.test.ts index 762569e627..0863b4b646 100644 --- a/packages/agent/src/adapters/claude/claude-agent.refresh.test.ts +++ b/packages/agent/src/adapters/claude/claude-agent.refresh.test.ts @@ -59,13 +59,14 @@ vi.mock("@anthropic-ai/claude-agent-sdk", () => ({ const fetchMcpToolMetadataMock = vi.fn().mockResolvedValue(undefined); const clearMcpToolMetadataCacheMock = vi.fn(); const clearMcpToolApprovalCacheMock = vi.fn(); +const setMcpToolApprovalStatesMock = vi.fn(); vi.mock("./mcp/tool-metadata", () => ({ fetchMcpToolMetadata: fetchMcpToolMetadataMock, getConnectedMcpServerNames: vi.fn().mockReturnValue([]), getCachedMcpTools: vi.fn().mockReturnValue([]), clearMcpToolMetadataCache: clearMcpToolMetadataCacheMock, clearMcpToolApprovalCache: clearMcpToolApprovalCacheMock, - setMcpToolApprovalStates: vi.fn(), + setMcpToolApprovalStates: setMcpToolApprovalStatesMock, })); // Import after the mocks so ClaudeAcpAgent resolves the mocked SDK @@ -173,6 +174,7 @@ describe("ClaudeAcpAgent.extMethod refresh_session", () => { fetchMcpToolMetadataMock.mockClear(); clearMcpToolMetadataCacheMock.mockClear(); clearMcpToolApprovalCacheMock.mockClear(); + setMcpToolApprovalStatesMock.mockClear(); }); it("returns methodNotFound for unknown extension methods", async () => { @@ -191,6 +193,25 @@ describe("ClaudeAcpAgent.extMethod refresh_session", () => { ).rejects.toThrow(/requires at least one refreshable field/); }); + it("updates MCP approval state without rebuilding when only approvals changed", async () => { + const agent = makeAgent(); + const { oldQuery, endSpy } = installFakeSession(agent, "s-approval"); + + const result = await agent.extMethod(POSTHOG_METHODS.REFRESH_SESSION, { + mcpToolApprovals: { + mcp__Linear__search: "needs_approval", + }, + }); + + expect(result).toEqual({ refreshed: true }); + expect(setMcpToolApprovalStatesMock).toHaveBeenCalledWith({ + mcp__Linear__search: "needs_approval", + }); + expect(oldQuery.interrupt).not.toHaveBeenCalled(); + expect(endSpy).not.toHaveBeenCalled(); + expect(createdQueries).toHaveLength(0); + }); + it("rejects when mcpServers is not an array", async () => { const agent = makeAgent(); installFakeSession(agent, "s-malformed"); diff --git a/packages/agent/src/adapters/claude/claude-agent.ts b/packages/agent/src/adapters/claude/claude-agent.ts index 68fef52689..b565ebd6c1 100644 --- a/packages/agent/src/adapters/claude/claude-agent.ts +++ b/packages/agent/src/adapters/claude/claude-agent.ts @@ -60,6 +60,7 @@ import { POSTHOG_PRODUCTS, type PostHogProductId, } from "../../posthog-products"; +import { mcpToolApprovalsSchema } from "../../server/schemas"; import type { PostHogAPIConfig } from "../../types"; import { isCloudRun, @@ -1231,10 +1232,25 @@ export class ClaudeAcpAgent extends BaseAcpAgent { throw RequestError.methodNotFound(method); } - // Trust boundary: refresh is only safe when the caller is trusted infra - // (e.g. the sandbox agent-server). Do not route this method from - // untrusted clients — parseMcpServers does no URL/command validation. + const approvalResult = + params.mcpToolApprovals === undefined + ? undefined + : mcpToolApprovalsSchema.safeParse(params.mcpToolApprovals); + if (approvalResult && !approvalResult.success) { + throw new RequestError( + -32602, + "refresh_session: mcpToolApprovals must be an object of approval states", + ); + } + if (approvalResult) { + setMcpToolApprovalStates(approvalResult.data); + } + + // Trust boundary: server refresh is only safe when the caller is trusted + // infra. Do not route mcpServers from untrusted clients — parseMcpServers + // does no URL/command validation. if (params.mcpServers === undefined) { + if (approvalResult) return { refreshed: true }; throw new RequestError( -32602, "refresh_session requires at least one refreshable field (e.g. mcpServers)", diff --git a/packages/agent/src/adapters/claude/hooks.test.ts b/packages/agent/src/adapters/claude/hooks.test.ts index d71790e632..e1114a0f0b 100644 --- a/packages/agent/src/adapters/claude/hooks.test.ts +++ b/packages/agent/src/adapters/claude/hooks.test.ts @@ -369,7 +369,7 @@ describe("createPreToolUseHook", () => { clearMcpToolApprovalCache(); }); - test("respects an explicit allow rule for a needs_approval MCP Store tool", async () => { + test("routes needs_approval MCP Store tools to ask even with an old allow rule", async () => { clearMcpToolMetadataCache(); clearMcpToolApprovalCache(); setMcpToolApprovalStates({ @@ -388,7 +388,7 @@ describe("createPreToolUseHook", () => { ); expect(result).toMatchObject({ - hookSpecificOutput: { permissionDecision: "allow" }, + hookSpecificOutput: { permissionDecision: "ask" }, }); clearMcpToolMetadataCache(); clearMcpToolApprovalCache(); diff --git a/packages/agent/src/adapters/claude/hooks.ts b/packages/agent/src/adapters/claude/hooks.ts index cb49001679..465354c63a 100644 --- a/packages/agent/src/adapters/claude/hooks.ts +++ b/packages/agent/src/adapters/claude/hooks.ts @@ -416,8 +416,8 @@ export const createPreToolUseHook = // gate that runs for these tools. Force a decision here: deny blocked // tools outright, and route needs_approval tools back through canUseTool // (via "ask") so handleMcpApprovalFlow relays the approval dialog to the - // desktop and persists the result to the backend. An explicit allow rule - // (e.g. a prior "always allow") takes precedence so the prompt sticks. + // desktop and persists the result to the backend. A later Settings change + // back to needs_approval must override any old local allow rule. const mcpApprovalState = getMcpToolApprovalState(toolName); if (mcpApprovalState === "do_not_use") { logger.info(`[PreToolUseHook] Blocking do_not_use MCP tool: ${toolName}`); @@ -433,7 +433,6 @@ export const createPreToolUseHook = } if ( mcpApprovalState === "needs_approval" && - permissionCheck.decision !== "allow" && permissionCheck.decision !== "deny" ) { logger.info( diff --git a/packages/agent/src/adapters/claude/permissions/permission-handlers.test.ts b/packages/agent/src/adapters/claude/permissions/permission-handlers.test.ts index 395e0e42c5..2ff453684a 100644 --- a/packages/agent/src/adapters/claude/permissions/permission-handlers.test.ts +++ b/packages/agent/src/adapters/claude/permissions/permission-handlers.test.ts @@ -126,7 +126,7 @@ describe("canUseTool MCP approval enforcement", () => { ); }); - it("updates the approval cache after the user approves a tool", async () => { + it("keeps needs_approval after the user allows a tool once", async () => { const toolName = "mcp__Linear__search"; setMcpToolApprovalStates({ [toolName]: "needs_approval", @@ -134,6 +134,27 @@ describe("canUseTool MCP approval enforcement", () => { const result = await canUseTool(createContext(toolName)); + expect(result.behavior).toBe("allow"); + expect(getMcpToolApprovalState(toolName)).toBe("needs_approval"); + }); + + it("updates the approval cache after the user always allows a tool", async () => { + const toolName = "mcp__Linear__search"; + setMcpToolApprovalStates({ + [toolName]: "needs_approval", + }); + + const result = await canUseTool( + createContext(toolName, { + client: { + sessionUpdate: vi.fn().mockResolvedValue(undefined), + requestPermission: vi.fn().mockResolvedValue({ + outcome: { outcome: "selected", optionId: "allow_always" }, + }), + }, + }), + ); + expect(result.behavior).toBe("allow"); expect(getMcpToolApprovalState(toolName)).toBe("approved"); }); diff --git a/packages/agent/src/adapters/claude/permissions/permission-handlers.ts b/packages/agent/src/adapters/claude/permissions/permission-handlers.ts index 31f78502ec..26f177a0e1 100644 --- a/packages/agent/src/adapters/claude/permissions/permission-handlers.ts +++ b/packages/agent/src/adapters/claude/permissions/permission-handlers.ts @@ -509,8 +509,8 @@ async function handleMcpApprovalFlow( (response.outcome.optionId === "allow" || response.outcome.optionId === "allow_always") ) { - setMcpToolApprovalStates({ [approvalToolName]: "approved" }); if (response.outcome.optionId === "allow_always") { + setMcpToolApprovalStates({ [approvalToolName]: "approved" }); return { behavior: "allow", updatedInput: toolInput as Record, diff --git a/packages/agent/src/adapters/codex/codex-agent.refresh.test.ts b/packages/agent/src/adapters/codex/codex-agent.refresh.test.ts index 6f780851bb..cd9cb8d599 100644 --- a/packages/agent/src/adapters/codex/codex-agent.refresh.test.ts +++ b/packages/agent/src/adapters/codex/codex-agent.refresh.test.ts @@ -121,6 +121,7 @@ type PrivateAgent = { }; configOptions: unknown[]; taskRunId?: string; + mcpToolApprovals?: Record; }; codexProcess: SpawnHandle; codexConnection: MockCodexConnection; @@ -188,6 +189,24 @@ describe("CodexAcpAgent.extMethod refresh_session", () => { ).rejects.toThrow(/at least one refreshable field/); }); + it("updates MCP approval state without respawning when only approvals changed", async () => { + const agent = makeAgent(); + const { priv } = primeSession(agent, "s-approval"); + + const result = await agent.extMethod(POSTHOG_METHODS.REFRESH_SESSION, { + mcpToolApprovals: { + mcp__Linear__search: "needs_approval", + }, + }); + + expect(result).toEqual({ refreshed: true }); + expect(priv.sessionState.mcpToolApprovals).toEqual({ + mcp__Linear__search: "needs_approval", + }); + expect(spawnedProcesses).toHaveLength(1); + expect(createdConnections).toHaveLength(1); + }); + it("rejects when mcpServers is not an array", async () => { const agent = makeAgent(); await expect( diff --git a/packages/agent/src/adapters/codex/codex-agent.ts b/packages/agent/src/adapters/codex/codex-agent.ts index 66f04d6600..fd4e1f419a 100644 --- a/packages/agent/src/adapters/codex/codex-agent.ts +++ b/packages/agent/src/adapters/codex/codex-agent.ts @@ -57,9 +57,10 @@ import { isCodexNativeMode, type PermissionMode, } from "../../execution-mode"; -import type { - McpToolApprovalsConfig, - McpToolInstallationsConfig, +import { + type McpToolApprovalsConfig, + type McpToolInstallationsConfig, + mcpToolApprovalsSchema, } from "../../server/schemas"; import type { PostHogAPIConfig, ProcessSpawnedCallback } from "../../types"; import { isCloudRun } from "../../utils/common"; @@ -881,11 +882,28 @@ export class CodexAcpAgent extends BaseAcpAgent { throw RequestError.methodNotFound(method); } - // Trust boundary: refresh is only safe when the caller is trusted infra - // (e.g. the sandbox agent-server). Do not route this method from - // untrusted clients — mcpServers contents are forwarded verbatim to - // codex-acp with no URL/command validation. + const approvalResult = + params.mcpToolApprovals === undefined + ? undefined + : mcpToolApprovalsSchema.safeParse(params.mcpToolApprovals); + if (approvalResult && !approvalResult.success) { + throw new RequestError( + -32602, + "refresh_session: mcpToolApprovals must be an object of approval states", + ); + } + if (approvalResult) { + this.sessionState.mcpToolApprovals = { + ...this.sessionState.mcpToolApprovals, + ...approvalResult.data, + }; + } + + // Trust boundary: server refresh is only safe when the caller is trusted + // infra. Do not route mcpServers from untrusted clients — contents are + // forwarded verbatim to codex-acp with no URL/command validation. if (params.mcpServers === undefined) { + if (approvalResult) return { refreshed: true }; throw new RequestError( -32602, "refresh_session requires at least one refreshable field (e.g. mcpServers)", diff --git a/packages/agent/src/adapters/codex/codex-client.test.ts b/packages/agent/src/adapters/codex/codex-client.test.ts index 010915c78b..3bed763e82 100644 --- a/packages/agent/src/adapters/codex/codex-client.test.ts +++ b/packages/agent/src/adapters/codex/codex-client.test.ts @@ -141,7 +141,7 @@ describe("createCodexClient MCP Store permissions", () => { }; } - test("relays needs_approval MCP tools by bare title even in full-access mode", async () => { + test("relays needs_approval MCP tools by bare title even in full-access mode without promoting allow once", async () => { const toolName = "mcp__Granola__query_granola_meetings"; const upstream = makePermissionUpstream(); const sessionState = createSessionState("sess", "/tmp", { @@ -178,12 +178,14 @@ describe("createCodexClient MCP Store permissions", () => { }), }), ); - expect(sessionState.mcpToolApprovals?.[toolName]).toBe("approved"); + expect(sessionState.mcpToolApprovals?.[toolName]).toBe("needs_approval"); }); - test("relays needs_approval MCP tools with unsanitized MCP server names", async () => { + test("promotes needs_approval MCP tools only when the user chooses always allow", async () => { const toolName = "mcp__Granola_Meetings__query_granola_meetings"; - const upstream = makePermissionUpstream(); + const upstream = makePermissionUpstream({ + outcome: { outcome: "selected" as const, optionId: "allow_always" }, + }); const sessionState = createSessionState("sess", "/tmp", { permissionMode: "full-access", mcpToolApprovals: { [toolName]: "needs_approval" }, diff --git a/packages/agent/src/adapters/codex/codex-client.ts b/packages/agent/src/adapters/codex/codex-client.ts index 9fcb2156fc..102d3b9097 100644 --- a/packages/agent/src/adapters/codex/codex-client.ts +++ b/packages/agent/src/adapters/codex/codex-client.ts @@ -121,11 +121,10 @@ function withMcpToolName( }; } -function isAllowResponse(response: RequestPermissionResponse): boolean { +function isAlwaysAllowResponse(response: RequestPermissionResponse): boolean { return ( response.outcome?.outcome === "selected" && - (response.outcome.optionId === "allow" || - response.outcome.optionId === "allow_always") + response.outcome.optionId === "allow_always" ); } @@ -217,7 +216,7 @@ export function createCodexClient( const response = await upstreamClient.requestPermission( withMcpToolName(params, mcpToolName), ); - if (isAllowResponse(response)) { + if (isAlwaysAllowResponse(response)) { sessionState.mcpToolApprovals = { ...sessionState.mcpToolApprovals, [mcpToolName]: "approved", diff --git a/packages/agent/src/adapters/codex/mcp-approval-hook.test.ts b/packages/agent/src/adapters/codex/mcp-approval-hook.test.ts new file mode 100644 index 0000000000..295b039899 --- /dev/null +++ b/packages/agent/src/adapters/codex/mcp-approval-hook.test.ts @@ -0,0 +1,120 @@ +import { afterEach, describe, expect, it, vi } from "vitest"; +import { CodexMcpApprovalHookBridge } from "./mcp-approval-hook"; + +const noopLogger = { + debug() {}, + info() {}, + warn() {}, + error() {}, +}; + +describe("CodexMcpApprovalHookBridge", () => { + let bridge: CodexMcpApprovalHookBridge | undefined; + + afterEach(async () => { + await bridge?.stop(); + bridge = undefined; + }); + + it("forwards PreToolUse hook input to the approval handler", async () => { + const preToolUse = vi.fn().mockResolvedValue({ + action: "deny" as const, + message: "blocked", + }); + bridge = new CodexMcpApprovalHookBridge( + { + preToolUse, + postToolUse: vi.fn(), + }, + noopLogger, + ); + const env = await bridge.start(); + + const response = await fetch(`${env.bridgeUrl}/pre-tool-use`, { + method: "POST", + headers: { + authorization: `Bearer ${env.bridgeToken}`, + "content-type": "application/json", + }, + body: JSON.stringify({ + hook_event_name: "PreToolUse", + tool_name: "mcp__Granola__query_granola_meetings", + tool_use_id: "tool-1", + tool_input: { limit: 10 }, + }), + }); + + expect(response.status).toBe(200); + expect(await response.json()).toEqual({ + action: "deny", + message: "blocked", + }); + expect(preToolUse).toHaveBeenCalledWith( + expect.objectContaining({ + hookEventName: "PreToolUse", + toolName: "mcp__Granola__query_granola_meetings", + toolUseId: "tool-1", + toolInput: { limit: 10 }, + }), + ); + }); + + it("forwards PostToolUse hook input to the cleanup handler", async () => { + const postToolUse = vi.fn().mockResolvedValue(undefined); + bridge = new CodexMcpApprovalHookBridge( + { + preToolUse: vi.fn(), + postToolUse, + }, + noopLogger, + ); + const env = await bridge.start(); + + const response = await fetch(`${env.bridgeUrl}/post-tool-use`, { + method: "POST", + headers: { + authorization: `Bearer ${env.bridgeToken}`, + "content-type": "application/json", + }, + body: JSON.stringify({ + hook_event_name: "PostToolUse", + tool_name: "mcp__Granola__query_granola_meetings", + tool_use_id: "tool-1", + tool_response: { ok: true }, + }), + }); + + expect(response.status).toBe(200); + expect(await response.json()).toEqual({ action: "allow" }); + expect(postToolUse).toHaveBeenCalledWith( + expect.objectContaining({ + hookEventName: "PostToolUse", + toolName: "mcp__Granola__query_granola_meetings", + toolUseId: "tool-1", + toolResponse: { ok: true }, + }), + ); + }); + + it("rejects requests without the bridge token", async () => { + bridge = new CodexMcpApprovalHookBridge( + { + preToolUse: vi.fn(), + postToolUse: vi.fn(), + }, + noopLogger, + ); + const env = await bridge.start(); + + const response = await fetch(`${env.bridgeUrl}/pre-tool-use`, { + method: "POST", + headers: { "content-type": "application/json" }, + body: JSON.stringify({ + hook_event_name: "PreToolUse", + tool_name: "mcp__Granola__query_granola_meetings", + }), + }); + + expect(response.status).toBe(401); + }); +}); diff --git a/packages/agent/src/adapters/codex/mcp-approval-hook.ts b/packages/agent/src/adapters/codex/mcp-approval-hook.ts new file mode 100644 index 0000000000..4d4dda3c9a --- /dev/null +++ b/packages/agent/src/adapters/codex/mcp-approval-hook.ts @@ -0,0 +1,351 @@ +import http from "node:http"; +import { randomBytes } from "node:crypto"; +import { chmod, mkdir, writeFile } from "node:fs/promises"; +import { join } from "node:path"; + +const BRIDGE_TOKEN_BYTES = 32; +const DEFAULT_HOOK_TIMEOUT_SECONDS = 600; + +export const CODEX_MCP_APPROVAL_HOOK_ENV = { + bridgeUrl: "POSTHOG_CODEX_MCP_APPROVAL_BRIDGE_URL", + bridgeToken: "POSTHOG_CODEX_MCP_APPROVAL_BRIDGE_TOKEN", +} as const; + +export interface CodexMcpApprovalHookEnv { + bridgeUrl: string; + bridgeToken: string; +} + +export interface CodexMcpApprovalHookInput { + hookEventName: "PreToolUse" | "PostToolUse"; + toolName: string; + toolUseId?: string; + toolInput?: unknown; + toolResponse?: unknown; + turnId?: string; + sessionId?: string; + raw: Record; +} + +export type CodexMcpApprovalHookDecision = + | { action: "allow" } + | { action: "deny"; message: string }; + +export interface CodexMcpApprovalHookHandler { + preToolUse( + input: CodexMcpApprovalHookInput, + ): Promise; + postToolUse(input: CodexMcpApprovalHookInput): Promise; +} + +interface HookLogger { + info(message: string, data?: unknown): void; + warn(message: string, data?: unknown): void; + error(message: string, data?: unknown): void; + debug(message: string, data?: unknown): void; +} + +function toHookInput(raw: unknown): CodexMcpApprovalHookInput | null { + if (!raw || typeof raw !== "object" || Array.isArray(raw)) { + return null; + } + + const record = raw as Record; + const hookEventName = record.hook_event_name; + const toolName = record.tool_name; + if ( + (hookEventName !== "PreToolUse" && hookEventName !== "PostToolUse") || + typeof toolName !== "string" + ) { + return null; + } + + return { + hookEventName, + toolName, + toolUseId: + typeof record.tool_use_id === "string" ? record.tool_use_id : undefined, + toolInput: record.tool_input, + toolResponse: record.tool_response, + turnId: typeof record.turn_id === "string" ? record.turn_id : undefined, + sessionId: + typeof record.session_id === "string" ? record.session_id : undefined, + raw: record, + }; +} + +function readRequestBody(req: http.IncomingMessage): Promise { + return new Promise((resolve, reject) => { + const chunks: Buffer[] = []; + req.on("data", (chunk: Buffer) => chunks.push(chunk)); + req.on("end", () => resolve(Buffer.concat(chunks).toString("utf8"))); + req.on("error", reject); + }); +} + +function writeJson( + res: http.ServerResponse, + status: number, + body: Record, +): void { + res.writeHead(status, { "content-type": "application/json" }); + res.end(JSON.stringify(body)); +} + +function createBridgeToken(): string { + return randomBytes(BRIDGE_TOKEN_BYTES).toString("base64url"); +} + +export class CodexMcpApprovalHookBridge { + private server: http.Server | null = null; + private port: number | null = null; + private readonly token = createBridgeToken(); + + constructor( + private readonly handler: CodexMcpApprovalHookHandler, + private readonly logger: HookLogger, + ) {} + + async start(): Promise { + if (this.server && this.port) { + return { + bridgeUrl: `http://127.0.0.1:${this.port}`, + bridgeToken: this.token, + }; + } + + const server = http.createServer((req, res) => { + this.handleRequest(req, res).catch((error) => { + this.logger.error("Codex MCP approval hook bridge error", { + error: error instanceof Error ? error.message : String(error), + }); + if (!res.headersSent) { + writeJson(res, 500, { + action: "deny", + message: "MCP approval hook bridge failed.", + }); + } else { + res.end(); + } + }); + }); + + await new Promise((resolve, reject) => { + server.listen(0, "127.0.0.1", () => { + const address = server.address(); + if (!address || typeof address === "string") { + reject(new Error("Failed to start Codex MCP approval hook bridge")); + return; + } + this.server = server; + this.port = address.port; + this.logger.info("Codex MCP approval hook bridge started", { + port: this.port, + }); + resolve(); + }); + server.on("error", reject); + }); + + return { + bridgeUrl: `http://127.0.0.1:${this.port}`, + bridgeToken: this.token, + }; + } + + async stop(): Promise { + if (!this.server) return; + const server = this.server; + await new Promise((resolve) => { + server.close(() => resolve()); + }); + this.server = null; + this.port = null; + } + + private async handleRequest( + req: http.IncomingMessage, + res: http.ServerResponse, + ): Promise { + if (req.method !== "POST") { + writeJson(res, 405, { error: "method_not_allowed" }); + return; + } + + const authorization = req.headers.authorization ?? ""; + if (authorization !== `Bearer ${this.token}`) { + writeJson(res, 401, { error: "unauthorized" }); + return; + } + + const url = new URL(req.url ?? "/", "http://127.0.0.1"); + const rawBody = await readRequestBody(req); + let parsed: unknown; + try { + parsed = rawBody ? JSON.parse(rawBody) : null; + } catch { + writeJson(res, 400, { + action: "deny", + message: "Invalid Codex hook JSON.", + }); + return; + } + const input = toHookInput(parsed); + if (!input) { + writeJson(res, 400, { + action: "deny", + message: "Invalid Codex hook input.", + }); + return; + } + + if (url.pathname === "/pre-tool-use") { + const decision = await this.handler.preToolUse(input); + writeJson(res, 200, decision); + return; + } + + if (url.pathname === "/post-tool-use") { + await this.handler.postToolUse(input); + writeJson(res, 200, { action: "allow" }); + return; + } + + writeJson(res, 404, { error: "not_found" }); + } +} + +function quoteCommandArg(value: string): string { + if (process.platform === "win32") { + return `"${value.replace(/"/g, '\\"')}"`; + } + return `'${value.replace(/'/g, "'\\''")}'`; +} + +function buildHookScript(): string { + return `#!/usr/bin/env node +const BRIDGE_URL_ENV = ${JSON.stringify(CODEX_MCP_APPROVAL_HOOK_ENV.bridgeUrl)}; +const BRIDGE_TOKEN_ENV = ${JSON.stringify(CODEX_MCP_APPROVAL_HOOK_ENV.bridgeToken)}; + +function readStdin() { + return new Promise((resolve, reject) => { + let data = ""; + process.stdin.setEncoding("utf8"); + process.stdin.on("data", (chunk) => { + data += chunk; + }); + process.stdin.on("end", () => resolve(data)); + process.stdin.on("error", reject); + }); +} + +function denyPreToolUse(message) { + process.stdout.write(JSON.stringify({ + hookSpecificOutput: { + hookEventName: "PreToolUse", + permissionDecision: "deny", + permissionDecisionReason: message, + }, + })); +} + +async function callBridge(path, input) { + const bridgeUrl = process.env[BRIDGE_URL_ENV]; + const bridgeToken = process.env[BRIDGE_TOKEN_ENV]; + if (!bridgeUrl || !bridgeToken) { + throw new Error("Codex MCP approval hook bridge is not configured."); + } + + const response = await fetch(new URL(path, bridgeUrl), { + method: "POST", + headers: { + "authorization": \`Bearer \${bridgeToken}\`, + "content-type": "application/json", + }, + body: JSON.stringify(input), + }); + const text = await response.text(); + const body = text ? JSON.parse(text) : {}; + if (!response.ok) { + throw new Error(typeof body.message === "string" ? body.message : "Codex MCP approval hook bridge rejected the request."); + } + return body; +} + +async function main() { + const raw = await readStdin(); + const input = raw ? JSON.parse(raw) : {}; + const event = input.hook_event_name; + if (event !== "PreToolUse" && event !== "PostToolUse") { + return; + } + + try { + const body = await callBridge( + event === "PostToolUse" ? "/post-tool-use" : "/pre-tool-use", + input, + ); + if (event === "PreToolUse" && body.action === "deny") { + denyPreToolUse(typeof body.message === "string" ? body.message : "MCP tool call was denied."); + } + } catch (error) { + if (event === "PreToolUse") { + denyPreToolUse(error instanceof Error ? error.message : String(error)); + } + } +} + +main().catch((error) => { + denyPreToolUse(error instanceof Error ? error.message : String(error)); +}); +`; +} + +export async function installCodexMcpApprovalHook(options: { + codexHome: string; + runtimeCommand: string; + timeoutSeconds?: number; +}): Promise { + const hooksDir = join(options.codexHome, "hooks"); + await mkdir(hooksDir, { recursive: true }); + + const scriptPath = join(hooksDir, "posthog-mcp-approval-hook.js"); + await writeFile(scriptPath, buildHookScript(), { mode: 0o700 }); + await chmod(scriptPath, 0o700); + + const command = `${quoteCommandArg(options.runtimeCommand)} ${quoteCommandArg(scriptPath)}`; + const hookConfig = { + hooks: { + PreToolUse: [ + { + matcher: "^mcp__.*", + hooks: [ + { + type: "command", + command, + timeout: options.timeoutSeconds ?? DEFAULT_HOOK_TIMEOUT_SECONDS, + statusMessage: "Checking MCP tool approval", + }, + ], + }, + ], + PostToolUse: [ + { + matcher: "^mcp__.*", + hooks: [ + { + type: "command", + command, + timeout: options.timeoutSeconds ?? DEFAULT_HOOK_TIMEOUT_SECONDS, + }, + ], + }, + ], + }, + }; + + await writeFile( + join(options.codexHome, "hooks.json"), + `${JSON.stringify(hookConfig, null, 2)}\n`, + ); +} diff --git a/packages/agent/src/adapters/codex/spawn.test.ts b/packages/agent/src/adapters/codex/spawn.test.ts index 44b77e56b1..c82d757f55 100644 --- a/packages/agent/src/adapters/codex/spawn.test.ts +++ b/packages/agent/src/adapters/codex/spawn.test.ts @@ -100,3 +100,27 @@ describe("spawnCodexProcess codex home", () => { expect(env.CODEX_HOME).toBe(process.env.CODEX_HOME); }); }); + +describe("spawnCodexProcess MCP approval hook", () => { + it("enables hooks and passes the hook bridge environment when configured", () => { + spawnMock.mockClear(); + spawnMock.mockReturnValue(makeFakeChild()); + + spawnCodexProcess({ + logger: new Logger({ debug: false }), + codexMcpApprovalHook: { + bridgeUrl: "http://127.0.0.1:4567", + bridgeToken: "secret-token", + }, + }); + + const args: string[] = spawnMock.mock.calls[0][1]; + expect(args).toContain("features.hooks=true"); + + const env = spawnMock.mock.calls[0][2].env; + expect(env.POSTHOG_CODEX_MCP_APPROVAL_BRIDGE_URL).toBe( + "http://127.0.0.1:4567", + ); + expect(env.POSTHOG_CODEX_MCP_APPROVAL_BRIDGE_TOKEN).toBe("secret-token"); + }); +}); diff --git a/packages/agent/src/adapters/codex/spawn.ts b/packages/agent/src/adapters/codex/spawn.ts index 9e14e1a8cd..2778864c9b 100644 --- a/packages/agent/src/adapters/codex/spawn.ts +++ b/packages/agent/src/adapters/codex/spawn.ts @@ -4,6 +4,10 @@ import { delimiter, dirname } from "node:path"; import type { Readable, Writable } from "node:stream"; import type { ProcessSpawnedCallback } from "../../types"; import { Logger } from "../../utils/logger"; +import { + CODEX_MCP_APPROVAL_HOOK_ENV, + type CodexMcpApprovalHookEnv, +} from "./mcp-approval-hook"; import type { CodexSettings } from "./settings"; export interface CodexProcessOptions { @@ -20,6 +24,7 @@ export interface CodexProcessOptions { developerInstructions?: string; binaryPath?: string; codexHome?: string; + codexMcpApprovalHook?: CodexMcpApprovalHookEnv; logger?: Logger; processCallbacks?: ProcessSpawnedCallback; settings?: CodexSettings; @@ -38,6 +43,9 @@ function buildConfigArgs(options: CodexProcessOptions): string[] { const args: string[] = []; args.push("-c", `features.remote_models=false`); + if (options.codexMcpApprovalHook) { + args.push("-c", `features.hooks=true`); + } // Disable the user's local MCPs one-by-one so Codex only uses the MCPs we // provide via ACP. We can't use `-c mcp_servers={}` because that makes Codex @@ -127,6 +135,13 @@ export function spawnCodexProcess(options: CodexProcessOptions): CodexProcess { env.CODEX_HOME = options.codexHome; } + if (options.codexMcpApprovalHook) { + env[CODEX_MCP_APPROVAL_HOOK_ENV.bridgeUrl] = + options.codexMcpApprovalHook.bridgeUrl; + env[CODEX_MCP_APPROVAL_HOOK_ENV.bridgeToken] = + options.codexMcpApprovalHook.bridgeToken; + } + const { command, args } = findCodexBinary(options); if (options.binaryPath && existsSync(options.binaryPath)) { diff --git a/packages/agent/src/agent.ts b/packages/agent/src/agent.ts index 28d26d627c..8b95b4509e 100644 --- a/packages/agent/src/agent.ts +++ b/packages/agent/src/agent.ts @@ -141,6 +141,7 @@ export class Agent { apiKey: gatewayConfig.apiKey, binaryPath: options.codexBinaryPath, codexHome: options.codexHome, + codexMcpApprovalHook: options.codexMcpApprovalHook, model: sanitizedModel, reasoningEffort: options.reasoningEffort, developerInstructions: options.developerInstructions, diff --git a/packages/agent/src/index.ts b/packages/agent/src/index.ts index c73615f419..3c520e485a 100644 --- a/packages/agent/src/index.ts +++ b/packages/agent/src/index.ts @@ -1,4 +1,8 @@ -export { isNotification, POSTHOG_NOTIFICATIONS } from "./acp-extensions"; +export { + isNotification, + POSTHOG_METHODS, + POSTHOG_NOTIFICATIONS, +} from "./acp-extensions"; export { getMcpToolMetadata, isMcpToolReadOnly, diff --git a/packages/agent/src/server/agent-server.test.ts b/packages/agent/src/server/agent-server.test.ts index cac2989962..fa60b8fe4d 100644 --- a/packages/agent/src/server/agent-server.test.ts +++ b/packages/agent/src/server/agent-server.test.ts @@ -235,6 +235,33 @@ interface TestableServer { permissionMode: PermissionMode; runtimeAdapter: "claude" | "codex"; }): Record; + approveCodexMcpToolUse(input: { + hookEventName: "PreToolUse"; + toolName: string; + toolUseId?: string; + raw: Record; + }): Promise<{ action: "allow" | "deny"; message?: string }>; + completeCodexMcpToolUse(input: { + hookEventName: "PostToolUse"; + toolName: string; + toolUseId?: string; + raw: Record; + }): Promise; + resolvePermission(requestId: string, optionId: string): boolean; + app: { + request: (url: string, init?: RequestInit) => Promise; + }; + config: { + mcpToolApprovals?: Record; + mcpToolInstallations?: Record< + string, + { installationId: string; toolName: string } + >; + }; + posthogAPI: { + updateMcpToolApproval: ReturnType; + }; + session: unknown; } interface NativeResumeTestServer { @@ -438,6 +465,91 @@ describe("AgentServer HTTP Mode", () => { }, 30000); }); + describe("cloud Codex MCP approval hook", () => { + it("asks for approval before Codex MCP Store tool calls and restores allow-once approvals", async () => { + const testServer = createServer({ + repositoryPath: undefined, + runtimeAdapter: "codex", + mcpToolApprovals: { + mcp__granola__query_granola_meetings: "needs_approval", + }, + mcpToolInstallations: { + mcp__granola__query_granola_meetings: { + installationId: "inst-1", + toolName: "query_granola_meetings", + }, + }, + }) as unknown as TestableServer; + const appendRawLine = vi.fn(); + const approvalSpy = vi.fn(async () => {}); + testServer.posthogAPI = { updateMcpToolApproval: approvalSpy }; + testServer.session = { + payload: { + run_id: "test-run-id", + task_id: "test-task-id", + team_id: 1, + user_id: 1, + distinct_id: "test-distinct-id", + mode: "interactive", + }, + acpSessionId: "codex-session-id", + sseController: null, + hasDesktopConnected: false, + permissionMode: "auto", + logWriter: { appendRawLine, flush: vi.fn(async () => {}) }, + acpConnection: { cleanup: vi.fn(async () => {}) }, + }; + + const decisionPromise = testServer.approveCodexMcpToolUse({ + hookEventName: "PreToolUse", + toolName: "mcp__granola__query_granola_meetings", + toolUseId: "tool-call-1", + raw: {}, + }); + + await vi.waitFor(() => expect(appendRawLine).toHaveBeenCalled()); + + const request = appendRawLine.mock.calls + .map(([, line]) => JSON.parse(line)) + .find((n) => n?.method === "_posthog/permission_request"); + expect(request).toBeTruthy(); + expect(request.params.toolCall.rawInput).toEqual({ + toolName: "mcp__granola__query_granola_meetings", + name: "query_granola_meetings", + }); + + expect( + testServer.resolvePermission(request.params.requestId, "allow"), + ).toBe(true); + + await expect(decisionPromise).resolves.toEqual({ action: "allow" }); + expect(approvalSpy).toHaveBeenNthCalledWith( + 1, + "inst-1", + "query_granola_meetings", + "approved", + ); + + await testServer.completeCodexMcpToolUse({ + hookEventName: "PostToolUse", + toolName: "mcp__granola__query_granola_meetings", + toolUseId: "tool-call-1", + raw: {}, + }); + + expect(approvalSpy).toHaveBeenNthCalledWith( + 2, + "inst-1", + "query_granola_meetings", + "needs_approval", + ); + expect( + testServer.config.mcpToolApprovals + ?.mcp__granola__query_granola_meetings, + ).toBe("needs_approval"); + }); + }); + describe("turn completion", () => { function stubSessionCleanup(testServer: unknown): { cleanupSession: (options?: { @@ -1282,6 +1394,12 @@ describe("AgentServer HTTP Mode", () => { expect( (s as unknown as TestableServer).buildCodexInstructions(sessionPrompt), ).toContain("Cloud Task Execution"); + expect( + (s as unknown as TestableServer).buildCodexInstructions(sessionPrompt), + ).toContain("MCP Tool Approvals"); + expect( + (s as unknown as TestableServer).buildCodexInstructions(sessionPrompt), + ).not.toContain("query_granola_meetings"); }); }); diff --git a/packages/agent/src/server/agent-server.ts b/packages/agent/src/server/agent-server.ts index 1dc32247ae..7b00b8e229 100644 --- a/packages/agent/src/server/agent-server.ts +++ b/packages/agent/src/server/agent-server.ts @@ -1,4 +1,5 @@ -import { access, mkdir, writeFile } from "node:fs/promises"; +import { access, mkdir, mkdtemp, rm, writeFile } from "node:fs/promises"; +import { tmpdir } from "node:os"; import { basename, join } from "node:path"; import { pathToFileURL } from "node:url"; import type { @@ -27,6 +28,13 @@ import { getSessionJsonlPath, hydrateSessionJsonl, } from "../adapters/claude/session/jsonl-hydration"; +import { + CodexMcpApprovalHookBridge, + type CodexMcpApprovalHookDecision, + type CodexMcpApprovalHookEnv, + type CodexMcpApprovalHookInput, + installCodexMcpApprovalHook, +} from "../adapters/codex/mcp-approval-hook"; import type { GatewayEnv } from "../adapters/claude/session/options"; import { type AgentErrorClassification, @@ -77,6 +85,7 @@ import { type JwtPayload, JwtValidationError, validateJwt } from "./jwt"; import { handoffLocalGitStateSchema, jsonRpcRequestSchema, + type RemoteMcpServer, validateCommandParams, } from "./schemas"; import type { AgentServerConfig } from "./types"; @@ -108,6 +117,19 @@ type MessageCallback = (message: unknown) => void; export const SSE_KEEPALIVE_INTERVAL_MS = 25_000; +const CODEX_MCP_APPROVAL_INSTRUCTIONS = `## MCP Tool Approvals + +Some MCP tools may pause for explicit user approval. When an MCP tool call is waiting for approval, wait for the permission result before continuing. If permission is denied or unavailable, explain that the tool cannot be used without approval.`; + +function getSelectedPermissionOption( + options: RequestPermissionRequest["options"], + response: RequestPermissionResponse, +): RequestPermissionRequest["options"][number] | undefined { + if (response.outcome?.outcome !== "selected") return undefined; + const { optionId } = response.outcome; + return options.find((option) => option.optionId === optionId); +} + interface CloudSessionMeta extends Record { sessionId: string; taskRunId: string; @@ -124,6 +146,13 @@ interface CloudSessionMeta extends Record { claudeCode?: { options: Record }; } +interface TemporaryMcpStoreToolApproval { + toolKey: string; + installationId: string; + toolName: string; + restoreApprovalState: "needs_approval"; +} + class NdJsonTap { private decoder = new TextDecoder(); private buffer = ""; @@ -250,6 +279,8 @@ interface ActiveSession { permissionMode: PermissionMode; /** Whether a desktop client has ever connected via SSE during this session */ hasDesktopConnected: boolean; + codexMcpApprovalHookBridge?: CodexMcpApprovalHookBridge; + codexHome?: string; pendingHandoffGitState?: HandoffLocalGitState; } @@ -292,6 +323,10 @@ export class AgentServer { // causing a second session to be created and duplicate Slack messages to be sent. private initializationPromise: Promise | null = null; private pendingEvents: Record[] = []; + private temporaryMcpStoreToolApprovals = new Map< + string, + TemporaryMcpStoreToolApproval + >(); private pendingPermissions = new Map< string, { @@ -456,35 +491,42 @@ export class AgentServer { ); } - private async approveMcpStoreToolIfNeeded( + private async applyMcpStoreToolApprovalResponse( params: RequestPermissionRequest, response: RequestPermissionResponse, - ): Promise { - const approved = - response.outcome?.outcome === "selected" && - (response.outcome.optionId === "allow" || - response.outcome.optionId === "allow_always"); - if (!approved) { - return; - } - + ): Promise<"none" | "temporary" | "persisted" | "failed"> { const toolName = this.getMcpToolNameFromPermissionRequest(params); if ( !toolName || this.config.mcpToolApprovals?.[toolName] !== "needs_approval" ) { - return; + return "none"; + } + + const selectedOption = getSelectedPermissionOption( + params.options, + response, + ); + const selectedOptionId = + response.outcome?.outcome === "selected" + ? response.outcome.optionId + : undefined; + const shouldPersistApproval = + selectedOption?.kind === "allow_always" || + selectedOptionId === "allow_always"; + const shouldAllowOnce = + selectedOption?.kind === "allow_once" || selectedOptionId === "allow"; + + if (!shouldPersistApproval && !shouldAllowOnce) return "none"; + + const toolCallId = params.toolCall?.toolCallId; + if (shouldAllowOnce && !toolCallId) { + return "failed"; } const installation = this.config.mcpToolInstallations?.[toolName]; if (!installation) { - this.logger.warn( - "Cannot persist MCP Store tool approval: missing installation ref", - { - toolName, - }, - ); - return; + return "failed"; } try { @@ -493,17 +535,205 @@ export class AgentServer { installation.toolName, "approved", ); + if (shouldAllowOnce && toolCallId) { + this.temporaryMcpStoreToolApprovals.set(toolCallId, { + toolKey: toolName, + installationId: installation.installationId, + toolName: installation.toolName, + restoreApprovalState: "needs_approval", + }); + return "temporary"; + } + this.config.mcpToolApprovals = { ...this.config.mcpToolApprovals, [toolName]: "approved", }; - this.logger.info("Persisted MCP Store tool approval", { toolName }); - } catch (error) { - this.logger.warn("Failed to persist MCP Store tool approval", { - toolName, - error: error instanceof Error ? error.message : String(error), - }); + if (toolCallId) { + this.temporaryMcpStoreToolApprovals.delete(toolCallId); + } + return "persisted"; + } catch { + return "failed"; + } + } + + private hasOtherTemporaryMcpStoreApproval( + toolCallId: string, + approval: TemporaryMcpStoreToolApproval, + ): boolean { + for (const [otherToolCallId, otherApproval] of this + .temporaryMcpStoreToolApprovals) { + if (otherToolCallId === toolCallId) continue; + if ( + otherApproval.installationId === approval.installationId && + otherApproval.toolName === approval.toolName + ) { + return true; + } + } + return false; + } + + private async restoreTemporaryMcpStoreToolApproval( + toolCallId: string, + _reason: string, + ): Promise { + const approval = this.temporaryMcpStoreToolApprovals.get(toolCallId); + if (!approval) return; + + this.temporaryMcpStoreToolApprovals.delete(toolCallId); + + if (this.hasOtherTemporaryMcpStoreApproval(toolCallId, approval)) { + return; + } + + if ( + this.config.mcpToolApprovals?.[approval.toolKey] !== + approval.restoreApprovalState + ) { + return; + } + + try { + await this.posthogAPI.updateMcpToolApproval( + approval.installationId, + approval.toolName, + approval.restoreApprovalState, + ); + } catch {} + } + + private async restoreTemporaryMcpStoreToolApprovals( + reason: string, + ): Promise { + const toolCallIds = Array.from(this.temporaryMcpStoreToolApprovals.keys()); + for (const toolCallId of toolCallIds) { + await this.restoreTemporaryMcpStoreToolApproval(toolCallId, reason); + } + } + + private isAllowPermissionResponse( + options: RequestPermissionRequest["options"], + response: RequestPermissionResponse, + ): boolean { + const selectedOption = getSelectedPermissionOption(options, response); + const selectedOptionId = + response.outcome?.outcome === "selected" + ? response.outcome.optionId + : undefined; + return ( + selectedOption?.kind === "allow_once" || + selectedOption?.kind === "allow_always" || + selectedOptionId === "allow" || + selectedOptionId === "allow_always" + ); + } + + private resolveCodexHookMcpToolKey( + input: CodexMcpApprovalHookInput, + ): string | null { + return resolveMcpStoreToolKey(input.toolName, { + approvals: this.config.mcpToolApprovals, + installations: this.config.mcpToolInstallations, + }); + } + + private async approveCodexMcpToolUse( + input: CodexMcpApprovalHookInput, + ): Promise { + const toolKey = this.resolveCodexHookMcpToolKey(input); + if (!toolKey) { + return { action: "allow" }; + } + + const installation = this.config.mcpToolInstallations?.[toolKey]; + const displayToolName = installation?.toolName ?? input.toolName; + const approvalState = this.config.mcpToolApprovals?.[toolKey]; + + if (approvalState === "approved") { + return { action: "allow" }; + } + if (approvalState === "do_not_use") { + return { + action: "deny", + message: `Tool '${displayToolName}' has been disabled by the user`, + }; + } + if (approvalState !== "needs_approval") { + return { action: "allow" }; + } + if (!this.session) { + return { + action: "deny", + message: "MCP tool call cannot proceed because the session is gone.", + }; + } + if (!input.toolUseId) { + return { + action: "deny", + message: + "MCP tool call cannot proceed because Codex did not provide a tool call id.", + }; } + + const toolCallId = input.toolUseId; + const options: RequestPermissionRequest["options"] = [ + { kind: "allow_once", optionId: "allow", name: "Allow" }, + { + kind: "allow_always", + optionId: "allow_always", + name: "Always allow", + }, + { kind: "reject_once", optionId: "reject", name: "Reject" }, + ]; + const permissionRequest: RequestPermissionRequest = { + sessionId: this.session.acpSessionId, + options, + toolCall: { + toolCallId, + title: displayToolName, + kind: "other", + rawInput: { toolName: toolKey, name: displayToolName }, + }, + }; + + const response = await this.relayPermissionToClient(permissionRequest); + if (!this.isAllowPermissionResponse(options, response)) { + const feedback = (response._meta as Record | undefined) + ?.customInput; + return { + action: "deny", + message: + typeof feedback === "string" && feedback.trim() + ? `User refused permission to run tool with feedback: ${feedback.trim()}` + : "User refused permission to run tool", + }; + } + + const application = await this.applyMcpStoreToolApprovalResponse( + permissionRequest, + response, + ); + if (application === "failed") { + return { + action: "deny", + message: + "MCP tool approval could not be applied before calling the tool.", + }; + } + + return { action: "allow" }; + } + + private async completeCodexMcpToolUse( + input: CodexMcpApprovalHookInput, + ): Promise { + if (!input.toolUseId) return; + await this.restoreTemporaryMcpStoreToolApproval( + input.toolUseId, + "codex_post_tool_use", + ); } private createApp(): Hono { @@ -1051,6 +1281,12 @@ export class AgentServer { const mcpServers = Array.isArray(params.mcpServers) ? params.mcpServers : []; + const mcpToolApprovals = + params.mcpToolApprovals && + typeof params.mcpToolApprovals === "object" && + !Array.isArray(params.mcpToolApprovals) + ? (params.mcpToolApprovals as AgentServerConfig["mcpToolApprovals"]) + : undefined; const refreshedCredentials = Array.isArray(params.refreshedCredentials) ? (params.refreshedCredentials as string[]) : []; @@ -1064,7 +1300,20 @@ export class AgentServer { ); } + if (mcpToolApprovals) { + this.config.mcpToolApprovals = { + ...this.config.mcpToolApprovals, + ...mcpToolApprovals, + }; + } + if (mcpServers.length === 0) { + if (mcpToolApprovals) { + return await this.session.clientConnection.extMethod( + POSTHOG_METHODS.REFRESH_SESSION, + { mcpToolApprovals }, + ); + } return { refreshed: true }; } @@ -1074,7 +1323,10 @@ export class AgentServer { return await this.session.clientConnection.extMethod( POSTHOG_METHODS.REFRESH_SESSION, - { mcpServers }, + { + mcpServers: mcpServers as RemoteMcpServer[], + ...(mcpToolApprovals && { mcpToolApprovals }), + }, ); } @@ -1237,197 +1489,245 @@ export class AgentServer { logger: new Logger({ debug: true, prefix: "[SessionLogWriter]" }), }); - const acpConnection = createAcpConnection({ - adapter: runtimeAdapter, - taskRunId: payload.run_id, - taskId: payload.task_id, - deviceType: deviceInfo.type, - logWriter, - logger: this.logger, - claudeGatewayEnv: runtimeAdapter !== "codex" ? gatewayEnv : undefined, - codexOptions: - runtimeAdapter === "codex" - ? { - cwd: this.config.repositoryPath ?? "/tmp/workspace", - apiBaseUrl: gatewayEnv.openaiBaseUrl, - apiKey: this.config.apiKey, - model: this.config.model ?? DEFAULT_CODEX_MODEL, - reasoningEffort: this.config.reasoningEffort, - developerInstructions: codexInstructions, - } - : undefined, - onStructuredOutput: async (output) => { - await this.posthogAPI.setTaskRunOutput( - payload.task_id, - payload.run_id, + let codexHome: string | undefined; + let codexMcpApprovalHookBridge: CodexMcpApprovalHookBridge | undefined; + let codexMcpApprovalHook: CodexMcpApprovalHookEnv | undefined; + let sessionCreated = false; + try { + if (runtimeAdapter === "codex") { + codexMcpApprovalHookBridge = new CodexMcpApprovalHookBridge( { - output, + preToolUse: (input) => this.approveCodexMcpToolUse(input), + postToolUse: (input) => this.completeCodexMcpToolUse(input), }, + this.logger, ); - }, - }); - - // Tap both streams to broadcast all ACP messages via SSE (mimics local transport) - this.adapterEmittedTurnComplete = false; - const onAcpMessage = (message: unknown) => { - if (isTurnCompleteNotification(message)) { - this.adapterEmittedTurnComplete = true; + codexMcpApprovalHook = await codexMcpApprovalHookBridge.start(); + codexHome = await mkdtemp(join(tmpdir(), "posthog-codex-home-")); + await installCodexMcpApprovalHook({ + codexHome, + runtimeCommand: process.execPath, + }); } - this.broadcastEvent({ - type: "notification", - timestamp: new Date().toISOString(), - notification: message, - }); - }; - - const tappedReadable = createTappedReadableStream( - acpConnection.clientStreams.readable as ReadableStream, - onAcpMessage, - this.logger, - ); - const tappedWritable = createTappedWritableStream( - acpConnection.clientStreams.writable as WritableStream, - onAcpMessage, - this.logger, - ); + const acpConnection = createAcpConnection({ + adapter: runtimeAdapter, + taskRunId: payload.run_id, + taskId: payload.task_id, + deviceType: deviceInfo.type, + logWriter, + logger: this.logger, + claudeGatewayEnv: runtimeAdapter !== "codex" ? gatewayEnv : undefined, + codexOptions: + runtimeAdapter === "codex" + ? { + cwd: this.config.repositoryPath ?? "/tmp/workspace", + apiBaseUrl: gatewayEnv.openaiBaseUrl, + apiKey: this.config.apiKey, + codexHome, + codexMcpApprovalHook, + model: this.config.model ?? DEFAULT_CODEX_MODEL, + reasoningEffort: this.config.reasoningEffort, + developerInstructions: codexInstructions, + } + : undefined, + onStructuredOutput: async (output) => { + await this.posthogAPI.setTaskRunOutput( + payload.task_id, + payload.run_id, + { + output, + }, + ); + }, + }); - const clientStream = ndJsonStream(tappedWritable, tappedReadable); + // Tap both streams to broadcast all ACP messages via SSE (mimics local transport) + this.adapterEmittedTurnComplete = false; + const onAcpMessage = (message: unknown) => { + if (isTurnCompleteNotification(message)) { + this.adapterEmittedTurnComplete = true; + } + this.broadcastEvent({ + type: "notification", + timestamp: new Date().toISOString(), + notification: message, + }); + }; - const clientConnection = new ClientSideConnection( - () => this.createCloudClient(payload), - clientStream, - ); + const tappedReadable = createTappedReadableStream( + acpConnection.clientStreams.readable as ReadableStream, + onAcpMessage, + this.logger, + ); - await clientConnection.initialize({ - protocolVersion: PROTOCOL_VERSION, - clientCapabilities: {}, - }); + const tappedWritable = createTappedWritableStream( + acpConnection.clientStreams.writable as WritableStream, + onAcpMessage, + this.logger, + ); - const runState = preTaskRun?.state as Record | undefined; - // Preserve native Codex modes for cloud runs so they behave the same as - // local sessions. Claude keeps the historical auto-approved default when - // PostHog Code has not explicitly selected a mode. - const initialPermissionMode: PermissionMode = - typeof runState?.initial_permission_mode === "string" - ? (runState.initial_permission_mode as PermissionMode) - : runtimeAdapter === "codex" - ? "auto" - : "bypassPermissions"; - const sessionCwd = this.config.repositoryPath ?? "/tmp/workspace"; - const sessionMeta = this.buildCloudSessionMeta({ - payload, - sessionSystemPrompt, - preTask, - permissionMode: initialPermissionMode, - runtimeAdapter, - }); + const clientStream = ndJsonStream(tappedWritable, tappedReadable); - const nativeResume = await this.prepareNativeResume( - payload, - posthogAPI, - preTaskRun, - runtimeAdapter, - sessionCwd, - initialPermissionMode, - ); + const clientConnection = new ClientSideConnection( + () => this.createCloudClient(payload), + clientStream, + ); - let acpSessionId: string; - if (nativeResume) { - await clientConnection.resumeSession({ - sessionId: nativeResume.sessionId, - cwd: sessionCwd, - mcpServers: this.config.mcpServers ?? [], - _meta: { ...sessionMeta, sessionId: nativeResume.sessionId }, + await clientConnection.initialize({ + protocolVersion: PROTOCOL_VERSION, + clientCapabilities: {}, }); - acpSessionId = nativeResume.sessionId; - this.nativeResume = nativeResume; - this.logger.debug("ACP session resumed", { - acpSessionId, - runId: payload.run_id, - warm: nativeResume.warm, - }); - } else { - const sessionResponse = await clientConnection.newSession({ - cwd: sessionCwd, - mcpServers: this.config.mcpServers ?? [], - _meta: sessionMeta, - }); - acpSessionId = sessionResponse.sessionId; - this.logger.debug("ACP session created", { - acpSessionId, - runId: payload.run_id, + + const runState = preTaskRun?.state as Record | undefined; + // Preserve native Codex modes for cloud runs so they behave the same as + // local sessions. Claude keeps the historical auto-approved default when + // PostHog Code has not explicitly selected a mode. + const initialPermissionMode: PermissionMode = + typeof runState?.initial_permission_mode === "string" + ? (runState.initial_permission_mode as PermissionMode) + : runtimeAdapter === "codex" + ? "auto" + : "bypassPermissions"; + const sessionCwd = this.config.repositoryPath ?? "/tmp/workspace"; + const sessionMeta = this.buildCloudSessionMeta({ + payload, + sessionSystemPrompt, + preTask, + permissionMode: initialPermissionMode, + runtimeAdapter, }); - } - this.evaluatedPrUrls.clear(); - this.prAttributionChain = Promise.resolve(); + const nativeResume = await this.prepareNativeResume( + payload, + posthogAPI, + preTaskRun, + runtimeAdapter, + sessionCwd, + initialPermissionMode, + ); - this.session = { - payload, - acpSessionId, - acpConnection, - clientConnection, - sseController, - deviceInfo, - logWriter, - permissionMode: initialPermissionMode, - hasDesktopConnected: sseController !== null, - pendingHandoffGitState: undefined, - }; + let acpSessionId: string; + const sessionMcpServers = this.config.mcpServers ?? []; + if (nativeResume) { + await clientConnection.resumeSession({ + sessionId: nativeResume.sessionId, + cwd: sessionCwd, + mcpServers: sessionMcpServers, + _meta: { ...sessionMeta, sessionId: nativeResume.sessionId }, + }); + acpSessionId = nativeResume.sessionId; + this.nativeResume = nativeResume; + this.logger.debug("ACP session resumed", { + acpSessionId, + runId: payload.run_id, + warm: nativeResume.warm, + }); + } else { + const sessionResponse = await clientConnection.newSession({ + cwd: sessionCwd, + mcpServers: sessionMcpServers, + _meta: sessionMeta, + }); + acpSessionId = sessionResponse.sessionId; + this.logger.debug("ACP session created", { + acpSessionId, + runId: payload.run_id, + }); + } - this.logger = new Logger({ - debug: true, - prefix: "[AgentServer]", - onLog: (level, scope, message, data) => { - this.emitConsoleLog(level, scope, message, data); - }, - }); + this.evaluatedPrUrls.clear(); + this.prAttributionChain = Promise.resolve(); - this.sessionReadyBootMs = Math.round(process.uptime() * 1000); - this.logger.debug("Session initialized successfully", { - bootMs: this.sessionReadyBootMs, - }); - this.logger.debug( - `Agent version: ${this.config.version ?? packageJson.version}`, - ); - await logAgentshRuntimeInfo(this.logger); - this.logger.debug(`Initial permission mode: ${initialPermissionMode}`); - - // Lifecycle handshake: clients gate "agent is ready to accept user - // messages" on this notification. Persisted to the session log so - // warm reconnects (sandbox restart with snapshot resume) replay it - // and see the agent come online again. - const runStartedNotification = { - jsonrpc: "2.0" as const, - method: POSTHOG_NOTIFICATIONS.RUN_STARTED, - params: { - sessionId: acpSessionId, - runId: payload.run_id, - taskId: payload.task_id, - agentVersion: this.config.version ?? packageJson.version, - }, - }; - this.broadcastEvent({ - type: "notification", - timestamp: new Date().toISOString(), - notification: runStartedNotification, - }); - this.session.logWriter.appendRawLine( - payload.run_id, - JSON.stringify(runStartedNotification), - ); + this.session = { + payload, + acpSessionId, + acpConnection, + clientConnection, + sseController, + deviceInfo, + logWriter, + permissionMode: initialPermissionMode, + hasDesktopConnected: sseController !== null, + codexMcpApprovalHookBridge, + codexHome, + pendingHandoffGitState: undefined, + }; + sessionCreated = true; - // Signal in_progress so the UI can start polling for updates - this.posthogAPI - .updateTaskRun(payload.task_id, payload.run_id, { - status: "in_progress", - }) - .catch((err) => - this.logger.debug("Failed to set task run to in_progress", err), + this.logger = new Logger({ + debug: true, + prefix: "[AgentServer]", + onLog: (level, scope, message, data) => { + this.emitConsoleLog(level, scope, message, data); + }, + }); + + this.sessionReadyBootMs = Math.round(process.uptime() * 1000); + this.logger.debug("Session initialized successfully", { + bootMs: this.sessionReadyBootMs, + }); + this.logger.debug( + `Agent version: ${this.config.version ?? packageJson.version}`, + ); + await logAgentshRuntimeInfo(this.logger); + this.logger.debug(`Initial permission mode: ${initialPermissionMode}`); + + // Lifecycle handshake: clients gate "agent is ready to accept user + // messages" on this notification. Persisted to the session log so + // warm reconnects (sandbox restart with snapshot resume) replay it + // and see the agent come online again. + const runStartedNotification = { + jsonrpc: "2.0" as const, + method: POSTHOG_NOTIFICATIONS.RUN_STARTED, + params: { + sessionId: acpSessionId, + runId: payload.run_id, + taskId: payload.task_id, + agentVersion: this.config.version ?? packageJson.version, + }, + }; + this.broadcastEvent({ + type: "notification", + timestamp: new Date().toISOString(), + notification: runStartedNotification, + }); + this.session.logWriter.appendRawLine( + payload.run_id, + JSON.stringify(runStartedNotification), ); - await this.sendInitialTaskMessage(payload, preTaskRun); + // Signal in_progress so the UI can start polling for updates + this.posthogAPI + .updateTaskRun(payload.task_id, payload.run_id, { + status: "in_progress", + }) + .catch((err) => + this.logger.debug("Failed to set task run to in_progress", err), + ); + + await this.sendInitialTaskMessage(payload, preTaskRun); + } catch (error) { + if (!sessionCreated) { + await codexMcpApprovalHookBridge?.stop().catch((stopError) => { + this.logger.debug( + "Failed to stop Codex MCP approval hook bridge after init failure", + { error: stopError }, + ); + }); + if (codexHome) { + await rm(codexHome, { recursive: true, force: true }).catch( + (rmError) => { + this.logger.debug( + "Failed to remove Codex home after init failure", + { error: rmError }, + ); + }, + ); + } + } + throw error; + } } private extractErrorClassification(error: unknown): { @@ -2050,9 +2350,11 @@ export class AgentServer { private buildCodexInstructions( systemPrompt: string | { append: string }, ): string { - return typeof systemPrompt === "string" - ? systemPrompt - : systemPrompt.append; + const instructions = + typeof systemPrompt === "string" ? systemPrompt : systemPrompt.append; + return [instructions, CODEX_MCP_APPROVAL_INSTRUCTIONS] + .filter(Boolean) + .join("\n\n"); } /** @@ -2640,7 +2942,7 @@ ${signedCommitInstructions} sessionPermissionMode, }); const response = await this.relayPermissionToClient(params); - await this.approveMcpStoreToolIfNeeded(params, response); + await this.applyMcpStoreToolApprovalResponse(params, response); return response; } } @@ -2694,10 +2996,26 @@ ${signedCommitInstructions} const meta = (params.update?._meta as Record) ?.claudeCode as Record | undefined; const toolName = meta?.toolName as string | undefined; + const toolCallId = + typeof params.update?.toolCallId === "string" + ? params.update.toolCallId + : undefined; + const status = + typeof params.update?.status === "string" + ? params.update.status + : undefined; const toolResponse = meta?.toolResponse as | Record | undefined; + if ( + toolName?.startsWith("mcp__") && + toolCallId && + (status === "completed" || status === "failed") + ) { + await this.restoreTemporaryMcpStoreToolApproval(toolCallId, status); + } + if ( (toolName === "Write" || toolName === "Edit" || @@ -2940,6 +3258,20 @@ ${signedCommitInstructions} } this.pendingPermissions.clear(); + await this.restoreTemporaryMcpStoreToolApprovals("session_cleanup"); + await this.session.codexMcpApprovalHookBridge?.stop().catch((error) => { + this.logger.debug("Failed to stop Codex MCP approval hook bridge", { + error, + }); + }); + if (this.session.codexHome) { + await rm(this.session.codexHome, { recursive: true, force: true }).catch( + (error) => { + this.logger.debug("Failed to remove Codex home", { error }); + }, + ); + } + try { await this.session.acpConnection.cleanup(); } catch (error) { diff --git a/packages/agent/src/server/question-relay.test.ts b/packages/agent/src/server/question-relay.test.ts index 147e454ae5..7303197144 100644 --- a/packages/agent/src/server/question-relay.test.ts +++ b/packages/agent/src/server/question-relay.test.ts @@ -20,6 +20,10 @@ interface TestableAgentServer { outcome: { outcome: string; optionId?: string }; _meta?: { message?: string }; }>; + sessionUpdate: (params: { + sessionId: string; + update?: Record; + }) => Promise; }; questionRelayedToSlack: boolean; session: unknown; @@ -409,11 +413,11 @@ describe("Question relay", () => { expect(request.params.toolCallId).toBe("tool-1"); expect( - server.resolvePermission(request.params.requestId, "allow"), + server.resolvePermission(request.params.requestId, "allow_always"), ).toBe(true); await expect(permissionPromise).resolves.toMatchObject({ - outcome: { outcome: "selected", optionId: "allow" }, + outcome: { outcome: "selected", optionId: "allow_always" }, }); expect(approvalSpy).toHaveBeenCalledWith( "inst-1", @@ -425,6 +429,126 @@ describe("Question relay", () => { ); }); + it("temporarily approves MCP Store tools for allow once and restores after the tool call", async () => { + const appendRawLine = vi.fn(); + const approvalSpy = vi + .spyOn(server.posthogAPI, "updateMcpToolApproval") + .mockResolvedValue(undefined); + + server.config.mcpToolApprovals = { + mcp__Granola__query_granola_meetings: "needs_approval", + }; + server.config.mcpToolInstallations = { + mcp__Granola__query_granola_meetings: { + installationId: "inst-1", + toolName: "query_granola_meetings", + }, + }; + server.session = { + payload: TEST_PAYLOAD, + sseController: null, + hasDesktopConnected: false, + permissionMode: "default", + logWriter: { appendRawLine }, + }; + + const client = server.createCloudClient(TEST_PAYLOAD); + const permissionPromise = client.requestPermission({ + options: ALLOW_OPTIONS, + toolCall: { + toolCallId: "tool-1", + title: "query_granola_meetings", + kind: "other", + rawInput: { toolName: "mcp__Granola__query_granola_meetings" }, + }, + }); + + const request = appendRawLine.mock.calls + .map(([, line]) => JSON.parse(line)) + .find((n) => n?.method === "_posthog/permission_request"); + expect(request).toBeTruthy(); + expect(request.params.toolCallId).toBe("tool-1"); + + expect( + server.resolvePermission(request.params.requestId, "allow"), + ).toBe(true); + + await expect(permissionPromise).resolves.toMatchObject({ + outcome: { outcome: "selected", optionId: "allow" }, + }); + expect(approvalSpy).toHaveBeenCalledWith( + "inst-1", + "query_granola_meetings", + "approved", + ); + expect( + server.config.mcpToolApprovals.mcp__Granola__query_granola_meetings, + ).toBe("needs_approval"); + + await client.sessionUpdate({ + sessionId: "test-session-id", + update: { + sessionUpdate: "tool_call_update", + toolCallId: "tool-1", + status: "completed", + _meta: { + claudeCode: { + toolName: "mcp__Granola__query_granola_meetings", + }, + }, + }, + }); + + expect(approvalSpy).toHaveBeenCalledWith( + "inst-1", + "query_granola_meetings", + "needs_approval", + ); + }); + + it("does not relay a permission request after a backend MCP approval failure", async () => { + const appendRawLine = vi.fn(); + const approvalSpy = vi + .spyOn(server.posthogAPI, "updateMcpToolApproval") + .mockResolvedValue(undefined); + + server.config.mcpToolApprovals = { + mcp__Granola__query_granola_meetings: "needs_approval", + }; + server.config.mcpToolInstallations = { + mcp__Granola__query_granola_meetings: { + installationId: "inst-1", + toolName: "query_granola_meetings", + }, + }; + server.session = { + payload: TEST_PAYLOAD, + acpSessionId: "codex-session-id", + sseController: null, + hasDesktopConnected: false, + permissionMode: "auto", + logWriter: { appendRawLine }, + }; + + const client = server.createCloudClient(TEST_PAYLOAD); + await client.sessionUpdate({ + sessionId: "codex-session-id", + update: { + sessionUpdate: "tool_call_update", + toolCallId: "call-granola", + status: "failed", + rawOutput: + "tool call error: tool call failed for `granola/query_granola_meetings`\n\nCaused by:\n Mcp error: -32001: Tool 'query_granola_meetings' requires approval before it can be called", + }, + }); + + expect(appendRawLine).not.toHaveBeenCalled(); + expect(approvalSpy).not.toHaveBeenCalled(); + expect( + server.config.mcpToolApprovals.mcp__Granola__query_granola_meetings, + ).toBe("needs_approval"); + }); + it("recognizes MCP Store approval requests by bare upstream tool name", async () => { const appendRawLine = vi.fn(); const approvalSpy = vi @@ -466,11 +590,11 @@ describe("Question relay", () => { expect(request.params.toolCallId).toBe("tool-1"); expect( - server.resolvePermission(request.params.requestId, "allow"), + server.resolvePermission(request.params.requestId, "allow_always"), ).toBe(true); await expect(permissionPromise).resolves.toMatchObject({ - outcome: { outcome: "selected", optionId: "allow" }, + outcome: { outcome: "selected", optionId: "allow_always" }, }); expect(approvalSpy).toHaveBeenCalledWith( "inst-1", @@ -525,11 +649,11 @@ describe("Question relay", () => { expect(request).toBeTruthy(); expect( - server.resolvePermission(request.params.requestId, "allow"), + server.resolvePermission(request.params.requestId, "allow_always"), ).toBe(true); await expect(permissionPromise).resolves.toMatchObject({ - outcome: { outcome: "selected", optionId: "allow" }, + outcome: { outcome: "selected", optionId: "allow_always" }, }); expect(approvalSpy).toHaveBeenCalledWith( "inst-1", diff --git a/packages/agent/src/server/schemas.test.ts b/packages/agent/src/server/schemas.test.ts index b7f228e025..a8df130c87 100644 --- a/packages/agent/src/server/schemas.test.ts +++ b/packages/agent/src/server/schemas.test.ts @@ -260,7 +260,17 @@ describe("validateCommandParams", () => { expect(result.success).toBe(true); }); - it("rejects refresh_session without mcpServers", () => { + it("accepts refresh_session with only MCP approval updates", () => { + const result = validateCommandParams("_posthog/refresh_session", { + mcpToolApprovals: { + mcp__Linear__search: "needs_approval", + }, + }); + + expect(result.success).toBe(true); + }); + + it("rejects refresh_session without refreshable fields", () => { const result = validateCommandParams("_posthog/refresh_session", {}); expect(result.success).toBe(false); diff --git a/packages/agent/src/server/schemas.ts b/packages/agent/src/server/schemas.ts index c1bd1b7b4c..0ae12e3fda 100644 --- a/packages/agent/src/server/schemas.ts +++ b/packages/agent/src/server/schemas.ts @@ -116,9 +116,16 @@ export const setConfigOptionParamsSchema = z.object({ value: z.string().min(1, "value is required"), }); -export const refreshSessionParamsSchema = z.object({ - mcpServers: mcpServersSchema, -}); +export const refreshSessionParamsSchema = z + .object({ + mcpServers: mcpServersSchema.optional(), + mcpToolApprovals: mcpToolApprovalsSchema.optional(), + }) + .refine( + (params) => + params.mcpServers !== undefined || params.mcpToolApprovals !== undefined, + { error: "refresh_session requires at least one refreshable field" }, + ); export const closeParamsSchema = z .object({ diff --git a/packages/agent/src/types.ts b/packages/agent/src/types.ts index 0056590678..3dbfe696ac 100644 --- a/packages/agent/src/types.ts +++ b/packages/agent/src/types.ts @@ -3,6 +3,7 @@ import type { HandoffLocalGitState as GitHandoffLocalGitState, PostHogAPIConfig, } from "@posthog/shared"; +import type { CodexMcpApprovalHookEnv } from "./adapters/codex/mcp-approval-hook"; import type { EffortLevel } from "@posthog/shared/domain-types"; export type { @@ -54,6 +55,7 @@ export interface TaskExecutionOptions { gatewayUrl?: string; codexBinaryPath?: string; codexHome?: string; + codexMcpApprovalHook?: CodexMcpApprovalHookEnv; reasoningEffort?: EffortLevel; /** * Codex-only. Appended on top of the model's base prompt via the Codex diff --git a/packages/agent/tsup.config.ts b/packages/agent/tsup.config.ts index 19269bb8d4..20ce4afa9e 100644 --- a/packages/agent/tsup.config.ts +++ b/packages/agent/tsup.config.ts @@ -109,12 +109,14 @@ export default defineConfig([ "src/pr-url-detector.ts", "src/resume.ts", "src/types.ts", + "src/mcp-store/tool-keys.ts", "src/adapters/claude/questions/utils.ts", "src/adapters/claude/permissions/permission-options.ts", "src/adapters/claude/tools.ts", "src/adapters/claude/conversion/tool-use-to-acp.ts", "src/adapters/claude/session/jsonl-hydration.ts", "src/adapters/claude/session/models.ts", + "src/adapters/codex/mcp-approval-hook.ts", "src/adapters/codex/models.ts", "src/adapters/claude/mcp/tool-metadata.ts", "src/adapters/codex/structured-output-mcp-server.ts", diff --git a/packages/host-router/src/routers/agent.router.ts b/packages/host-router/src/routers/agent.router.ts index e62ea18cf9..bfc6619848 100644 --- a/packages/host-router/src/routers/agent.router.ts +++ b/packages/host-router/src/routers/agent.router.ts @@ -25,6 +25,7 @@ import { setConfigOptionInput, startSessionInput, subscribeSessionInput, + updateMcpToolApprovalForActiveSessionsInput, } from "@posthog/workspace-server/services/agent/schemas"; import { PROCESS_TRACKING_SERVICE } from "@posthog/workspace-server/services/process-tracking/identifiers"; import type { ProcessTrackingService } from "@posthog/workspace-server/services/process-tracking/process-tracking"; @@ -81,6 +82,14 @@ export const agentRouter = router({ .setSessionConfigOption(input.sessionId, input.configId, input.value), ), + updateMcpToolApprovalForActiveSessions: publicProcedure + .input(updateMcpToolApprovalForActiveSessionsInput) + .mutation(({ ctx, input }) => + ctx.container + .get(AGENT_SERVICE) + .updateMcpToolApprovalForActiveSessions(input), + ), + onSessionEvent: publicProcedure .input(subscribeSessionInput) .subscription(async function* (opts) { diff --git a/packages/ui/src/features/mcp-servers/hooks/useMcpInstallationTools.ts b/packages/ui/src/features/mcp-servers/hooks/useMcpInstallationTools.ts index 4caa73d3b7..b98ce5ff09 100644 --- a/packages/ui/src/features/mcp-servers/hooks/useMcpInstallationTools.ts +++ b/packages/ui/src/features/mcp-servers/hooks/useMcpInstallationTools.ts @@ -4,7 +4,7 @@ import type { } from "@posthog/api-client/posthog-client"; import { dispatchBulkApproval } from "@posthog/core/mcp-servers/toolBulk"; import { shouldAutoRefreshTools } from "@posthog/core/mcp-servers/toolRefresh"; -import { useHostTRPC } from "@posthog/host-router/react"; +import { useHostTRPC, useHostTRPCClient } from "@posthog/host-router/react"; import { useAuthenticatedMutation } from "@posthog/ui/hooks/useAuthenticatedMutation"; import { useAuthenticatedQuery } from "@posthog/ui/hooks/useAuthenticatedQuery"; import { useQueryClient } from "@tanstack/react-query"; @@ -28,6 +28,7 @@ export function useMcpInstallationTools( options: UseMcpInstallationToolsOptions = {}, ) { const trpc = useHostTRPC(); + const trpcClient = useHostTRPCClient(); const queryClient = useQueryClient(); const queryKey = [ @@ -68,8 +69,16 @@ export function useMcpInstallationTools( ); }, { - onSuccess: () => { + onSuccess: (_data, vars) => { invalidate(); + if (!installationId) return; + void trpcClient.agent.updateMcpToolApprovalForActiveSessions + .mutate({ + installationId, + toolName: vars.toolName, + approvalState: vars.approval_state, + }) + .catch(() => undefined); }, onError: (error: Error) => { toast.error(error.message || "Failed to update tool approval"); @@ -96,8 +105,19 @@ export function useMcpInstallationTools( ); }, { - onSuccess: () => { + onSuccess: (_data, vars) => { invalidate(); + if (!installationId) return; + const changedTools = vars.targetTools ?? tools ?? []; + for (const tool of changedTools) { + void trpcClient.agent.updateMcpToolApprovalForActiveSessions + .mutate({ + installationId, + toolName: tool.tool_name, + approvalState: vars.approval_state, + }) + .catch(() => undefined); + } }, onError: (error: Error) => { toast.error(error.message || "Failed to update tool approvals"); diff --git a/packages/workspace-server/src/services/agent/agent.test.ts b/packages/workspace-server/src/services/agent/agent.test.ts index adeffc2afb..0590c8d01f 100644 --- a/packages/workspace-server/src/services/agent/agent.test.ts +++ b/packages/workspace-server/src/services/agent/agent.test.ts @@ -1,3 +1,8 @@ +import type { + RequestPermissionRequest, + RequestPermissionResponse, + SessionNotification, +} from "@agentclientprotocol/sdk"; import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; // --- Hoisted mocks --- @@ -22,6 +27,7 @@ const mockClientSideConnection = vi.hoisted(() => this.newSession = mockNewSession; this.loadSession = vi.fn().mockResolvedValue({ configOptions: [] }); this.resumeSession = vi.fn().mockResolvedValue({ configOptions: [] }); + this.extMethod = vi.fn().mockResolvedValue({ refreshed: true }); }), ); @@ -45,6 +51,20 @@ const mockAgentConstructor = vi.hoisted(() => }), ); +const mockCodexHome = vi.hoisted(() => ({ + prepareCodexHome: vi.fn().mockResolvedValue("/mock/codex-home"), + cleanupCodexHome: vi.fn().mockResolvedValue(undefined), +})); + +const mockCodexHookBridge = vi.hoisted(() => ({ + start: vi.fn().mockResolvedValue({ + bridgeUrl: "http://127.0.0.1:4567", + bridgeToken: "hook-token", + }), + stop: vi.fn().mockResolvedValue(undefined), + constructor: vi.fn(), +})); + // --- Module mocks --- vi.mock("electron", () => ({ @@ -63,6 +83,14 @@ vi.mock("@agentclientprotocol/sdk", () => ({ vi.mock("@posthog/agent", () => ({ isMcpToolReadOnly: vi.fn(() => false), + isNotification: vi.fn(() => false), + POSTHOG_METHODS: { + REFRESH_SESSION: "_posthog/refresh_session", + }, + POSTHOG_NOTIFICATIONS: { + SDK_SESSION: "_posthog/sdk_session", + USAGE_UPDATE: "_posthog/usage_update", + }, })); vi.mock("@posthog/agent/posthog-api", () => ({ @@ -82,6 +110,22 @@ vi.mock("@posthog/agent/adapters/claude/session/jsonl-hydration", () => ({ hydrateSessionJsonl: vi.fn().mockResolvedValue(undefined), })); +vi.mock("@posthog/agent/adapters/codex/mcp-approval-hook", () => ({ + CodexMcpApprovalHookBridge: vi.fn().mockImplementation(function ( + this: Record, + handler, + ) { + mockCodexHookBridge.constructor(handler); + this.start = mockCodexHookBridge.start; + this.stop = mockCodexHookBridge.stop; + }), +})); + +vi.mock("./codex-home", () => ({ + prepareCodexHome: mockCodexHome.prepareCodexHome, + cleanupCodexHome: mockCodexHome.cleanupCodexHome, +})); + vi.mock("node:fs", async (importOriginal) => { const original = await importOriginal(); return { @@ -101,6 +145,7 @@ vi.mock("node:fs", async (importOriginal) => { // --- Import after mocks --- import type { RegisteredFolder } from "../folders/schemas"; import { AgentService, buildAutoApproveOutcome } from "./agent"; +import { AgentServiceEvent } from "./schemas"; // --- Test helpers --- @@ -145,6 +190,7 @@ function createMockDependencies() { toolApprovals: {}, toolInstallations: {}, }), + updateMcpToolApproval: vi.fn().mockResolvedValue(undefined), }, mcpAppsService: { setServerConfigs: vi.fn(), @@ -199,12 +245,62 @@ const baseSessionParams = { projectId: 1, }; +type TestAcpClient = { + requestPermission( + params: RequestPermissionRequest, + ): Promise; + sessionUpdate(params: SessionNotification): Promise; +}; + +function getCreatedAcpClient(): TestAcpClient { + const latestCall = + mockClientSideConnection.mock.calls[ + mockClientSideConnection.mock.calls.length - 1 + ]; + const createClient = latestCall?.[0] as + | ((agent: unknown) => TestAcpClient) + | undefined; + if (!createClient) { + throw new Error("No ACP client connection was created"); + } + return createClient({}); +} + describe("AgentService", () => { let service: AgentService; let deps: ReturnType; beforeEach(() => { vi.clearAllMocks(); + mockNewSession.mockResolvedValue({ + sessionId: "test-session-id", + configOptions: [], + }); + mockClientSideConnection.mockImplementation(function ( + this: Record, + ) { + this.initialize = vi.fn().mockResolvedValue({}); + this.newSession = mockNewSession; + this.loadSession = vi.fn().mockResolvedValue({ configOptions: [] }); + this.resumeSession = vi.fn().mockResolvedValue({ configOptions: [] }); + this.extMethod = vi.fn().mockResolvedValue({ refreshed: true }); + }); + mockAgentRun.mockImplementation(() => + Promise.resolve({ + clientStreams: { + readable: new ReadableStream(), + writable: new WritableStream(), + }, + }), + ); + mockAgentConstructor.mockImplementation(function ( + this: Record, + ) { + this.run = mockAgentRun; + this.cleanup = vi.fn().mockResolvedValue(undefined); + this.getPosthogAPI = vi.fn(); + this.flushAllLogs = vi.fn().mockResolvedValue(undefined); + }); // The Codex MCP reachability probe hits the network; default it to "reachable" // so unrelated session tests stay deterministic and offline-safe. @@ -344,6 +440,342 @@ describe("AgentService", () => { }), ); }); + + it("pushes MCP approval changes into active sessions", async () => { + deps.agentAuthAdapter.buildMcpServers.mockResolvedValueOnce({ + servers: [ + { + name: "posthog", + type: "http", + url: "https://mcp.posthog.com/mcp", + headers: [], + }, + ], + toolApprovals: { + mcp__Linear__search: "approved", + }, + toolInstallations: { + mcp__Linear__search: { + installationId: "inst-1", + toolName: "search", + }, + }, + }); + + await service.startSession({ + ...baseSessionParams, + adapter: "codex", + }); + + await service.updateMcpToolApprovalForActiveSessions({ + installationId: "inst-1", + toolName: "search", + approvalState: "needs_approval", + }); + + const connection = mockClientSideConnection.mock.instances[0] as { + extMethod: ReturnType; + }; + expect(connection.extMethod).toHaveBeenCalledWith( + "_posthog/refresh_session", + { + mcpToolApprovals: { + mcp__Linear__search: "needs_approval", + }, + }, + ); + }); + + it("temporarily approves an MCP Store tool for allow once and restores it after the tool call", async () => { + const toolKey = "mcp__Granola__query_granola_meetings"; + deps.agentAuthAdapter.buildMcpServers.mockResolvedValueOnce({ + servers: [ + { + name: "granola", + type: "http", + url: "https://mcp.posthog.com/granola", + headers: [], + }, + ], + toolApprovals: { + [toolKey]: "needs_approval", + }, + toolInstallations: { + [toolKey]: { + installationId: "inst-granola", + toolName: "query_granola_meetings", + }, + }, + }); + + await service.startSession({ + ...baseSessionParams, + adapter: "codex", + }); + + const client = getCreatedAcpClient(); + const permissionPromise = client.requestPermission({ + sessionId: "test-session-id", + options: [ + { kind: "allow_once", optionId: "allow", name: "Allow" }, + { + kind: "allow_always", + optionId: "allow_always", + name: "Always allow", + }, + ], + toolCall: { + toolCallId: "tc-granola", + title: "query_granola_meetings", + kind: "read", + rawInput: { toolName: toolKey }, + }, + }); + + service.respondToPermission("run-1", "tc-granola", "allow"); + + await expect(permissionPromise).resolves.toMatchObject({ + outcome: { outcome: "selected", optionId: "allow" }, + }); + expect(deps.agentAuthAdapter.updateMcpToolApproval).toHaveBeenCalledWith( + expect.objectContaining({ + apiHost: "https://app.posthog.com", + projectId: 1, + }), + "inst-granola", + "query_granola_meetings", + "approved", + ); + + await client.sessionUpdate({ + sessionId: "test-session-id", + update: { + sessionUpdate: "tool_call_update", + toolCallId: "tc-granola", + status: "completed", + rawOutput: { ok: true }, + _meta: { + claudeCode: { + toolName: toolKey, + }, + }, + }, + }); + + expect(deps.agentAuthAdapter.updateMcpToolApproval).toHaveBeenCalledWith( + expect.objectContaining({ + apiHost: "https://app.posthog.com", + projectId: 1, + }), + "inst-granola", + "query_granola_meetings", + "needs_approval", + ); + expect(deps.mcpAppsService.notifyToolResult).toHaveBeenCalledWith( + toolKey, + "tc-granola", + { ok: true }, + false, + ); + }); + + it("persists MCP Store approval for allow always without restoring it", async () => { + const toolKey = "mcp__Granola__query_granola_meetings"; + deps.agentAuthAdapter.buildMcpServers.mockResolvedValueOnce({ + servers: [ + { + name: "granola", + type: "http", + url: "https://mcp.posthog.com/granola", + headers: [], + }, + ], + toolApprovals: { + [toolKey]: "needs_approval", + }, + toolInstallations: { + [toolKey]: { + installationId: "inst-granola", + toolName: "query_granola_meetings", + }, + }, + }); + + await service.startSession({ + ...baseSessionParams, + adapter: "codex", + }); + + const client = getCreatedAcpClient(); + const permissionPromise = client.requestPermission({ + sessionId: "test-session-id", + options: [ + { kind: "allow_once", optionId: "allow", name: "Allow" }, + { + kind: "allow_always", + optionId: "allow_always", + name: "Always allow", + }, + ], + toolCall: { + toolCallId: "tc-granola", + title: "query_granola_meetings", + kind: "read", + rawInput: { toolName: toolKey }, + }, + }); + + service.respondToPermission("run-1", "tc-granola", "allow_always"); + + await expect(permissionPromise).resolves.toMatchObject({ + outcome: { outcome: "selected", optionId: "allow_always" }, + }); + + await client.sessionUpdate({ + sessionId: "test-session-id", + update: { + sessionUpdate: "tool_call_update", + toolCallId: "tc-granola", + status: "completed", + rawOutput: { ok: true }, + _meta: { + claudeCode: { + toolName: toolKey, + }, + }, + }, + }); + + expect(deps.agentAuthAdapter.updateMcpToolApproval).toHaveBeenCalledTimes( + 1, + ); + expect(deps.agentAuthAdapter.updateMcpToolApproval).toHaveBeenCalledWith( + expect.objectContaining({ + apiHost: "https://app.posthog.com", + projectId: 1, + }), + "inst-granola", + "query_granola_meetings", + "approved", + ); + }); + + it("prompts and temporarily approves Codex MCP Store calls through the PreToolUse hook", async () => { + const toolKey = "mcp__Granola__query_granola_meetings"; + deps.agentAuthAdapter.buildMcpServers.mockResolvedValueOnce({ + servers: [ + { + name: "granola", + type: "http", + url: "https://mcp.posthog.com/granola", + headers: [], + }, + ], + toolApprovals: { + [toolKey]: "needs_approval", + }, + toolInstallations: { + [toolKey]: { + installationId: "inst-granola", + toolName: "query_granola_meetings", + }, + }, + }); + + await service.startSession({ + ...baseSessionParams, + adapter: "codex", + }); + + expect(mockAgentRun).toHaveBeenCalledWith( + "task-1", + "run-1", + expect.objectContaining({ + codexHome: "/mock/codex-home", + codexMcpApprovalHook: { + bridgeUrl: "http://127.0.0.1:4567", + bridgeToken: "hook-token", + }, + }), + ); + + const hookHandler = mockCodexHookBridge.constructor.mock.calls[0]?.[0] as + | { + preToolUse(input: { + hookEventName: "PreToolUse"; + toolName: string; + toolUseId: string; + raw: Record; + }): Promise<{ action: "allow" | "deny"; message?: string }>; + postToolUse(input: { + hookEventName: "PostToolUse"; + toolName: string; + toolUseId: string; + raw: Record; + }): Promise; + } + | undefined; + expect(hookHandler).toBeDefined(); + + const decisionPromise = hookHandler?.preToolUse({ + hookEventName: "PreToolUse", + toolName: toolKey, + toolUseId: "tc-granola", + raw: {}, + }); + await Promise.resolve(); + + const permissionRequest = ( + service.emit as ReturnType + ).mock.calls.find( + ([eventName]) => eventName === AgentServiceEvent.PermissionRequest, + )?.[1] as + | { + taskRunId: string; + toolCall?: { toolCallId?: string; rawInput?: unknown }; + } + | undefined; + expect(permissionRequest?.taskRunId).toBe("run-1"); + expect(permissionRequest?.toolCall?.rawInput).toEqual({ + toolName: toolKey, + name: "query_granola_meetings", + }); + + service.respondToPermission( + "run-1", + permissionRequest?.toolCall?.toolCallId ?? "", + "allow", + ); + + const decision = await decisionPromise; + expect(decision?.action).toBe("allow"); + expect(deps.agentAuthAdapter.updateMcpToolApproval).toHaveBeenCalledWith( + expect.objectContaining({ + apiHost: "https://app.posthog.com", + projectId: 1, + }), + "inst-granola", + "query_granola_meetings", + "approved", + ); + + await hookHandler?.postToolUse({ + hookEventName: "PostToolUse", + toolName: toolKey, + toolUseId: "tc-granola", + raw: {}, + }); + + expect(deps.agentAuthAdapter.updateMcpToolApproval).toHaveBeenCalledWith( + expect.objectContaining({ + apiHost: "https://app.posthog.com", + projectId: 1, + }), + "inst-granola", + "query_granola_meetings", + "needs_approval", + ); + }); }); describe("idle timeout", () => { @@ -366,6 +798,7 @@ describe("AgentService", () => { config: {}, promptPending: false, inFlightMcpToolCalls: new Map(), + temporaryMcpToolApprovals: new Map(), mcpToolApprovals: {}, toolInstallations: {}, ...overrides, diff --git a/packages/workspace-server/src/services/agent/agent.ts b/packages/workspace-server/src/services/agent/agent.ts index 7d28a3d71b..aef1412a19 100644 --- a/packages/workspace-server/src/services/agent/agent.ts +++ b/packages/workspace-server/src/services/agent/agent.ts @@ -15,9 +15,19 @@ import { import { isMcpToolReadOnly, isNotification, + POSTHOG_METHODS, POSTHOG_NOTIFICATIONS, } from "@posthog/agent"; -import type { McpToolApprovals } from "@posthog/agent/adapters/claude/mcp/tool-metadata"; +import { + CodexMcpApprovalHookBridge, + type CodexMcpApprovalHookDecision, + type CodexMcpApprovalHookEnv, + type CodexMcpApprovalHookInput, +} from "@posthog/agent/adapters/codex/mcp-approval-hook"; +import type { + McpToolApprovalState, + McpToolApprovals, +} from "@posthog/agent/adapters/claude/mcp/tool-metadata"; import { hydrateSessionJsonl } from "@posthog/agent/adapters/claude/session/jsonl-hydration"; import { getReasoningEffortOptions } from "@posthog/agent/adapters/reasoning-effort"; import { Agent } from "@posthog/agent/agent"; @@ -35,6 +45,7 @@ import { isAnthropicModel, isOpenAIModel, } from "@posthog/agent/gateway-models"; +import { resolveMcpStoreToolKey } from "@posthog/agent/mcp-store/tool-keys"; import { getLlmGatewayUrl } from "@posthog/agent/posthog-api"; import { findPrUrl, wasCreatedRecently } from "@posthog/agent/pr-url-detector"; import type * as AgentTypes from "@posthog/agent/types"; @@ -298,15 +309,27 @@ interface ManagedSession { configOptions?: SessionConfigOption[]; /** Tracks in-flight MCP tool calls (toolCallId → toolKey) for cancellation */ inFlightMcpToolCalls: Map; + /** One-time backend approvals to restore after their tool call finishes */ + temporaryMcpToolApprovals: Map; /** MCP tool approval states fetched at session start */ mcpToolApprovals: McpToolApprovals; /** Maps tool keys to their installation for backend approval updates */ toolInstallations: McpToolInstallations; + codexMcpApprovalHookBridge?: CodexMcpApprovalHookBridge; // Reset per session. `evaluatedPrUrls` dedupes the GitHub lookup per URL. prAttributed: boolean; evaluatedPrUrls: Set; } +interface TemporaryMcpToolApproval { + toolKey: string; + installationId: string; + toolName: string; + restoreApprovalState: McpToolApprovalState; +} + +type McpToolApprovalApplication = "none" | "persisted" | "temporary" | "failed"; + /** Get the agent session ID from a managed session, throwing if not set. */ function getAgentSessionId(session: ManagedSession): string { const { sessionId } = session.config; @@ -329,6 +352,25 @@ export function buildAutoApproveOutcome( return { outcome: "selected", optionId }; } +function getSelectedPermissionOption( + options: RequestPermissionRequest["options"], + response: RequestPermissionResponse, +): RequestPermissionRequest["options"][number] | undefined { + if (response.outcome?.outcome !== "selected") return undefined; + const { optionId } = response.outcome; + return options.find((option) => option.optionId === optionId); +} + +function isSelectedPermissionOption( + response: RequestPermissionResponse, + optionId: string, +): boolean { + return ( + response.outcome?.outcome === "selected" && + response.outcome.optionId === optionId + ); +} + interface PendingPermission { resolve: (response: RequestPermissionResponse) => void; reject: (error: Error) => void; @@ -454,6 +496,307 @@ export class AgentService extends TypedEventEmitter { this.recordActivity(taskRunId); } + public async updateMcpToolApprovalForActiveSessions(input: { + installationId: string; + toolName: string; + approvalState: McpToolApprovalState; + }): Promise { + const refreshes: Promise[] = []; + + for (const session of this.sessions.values()) { + const patch: McpToolApprovals = {}; + for (const [toolKey, ref] of Object.entries(session.toolInstallations)) { + if ( + ref.installationId === input.installationId && + ref.toolName === input.toolName + ) { + patch[toolKey] = input.approvalState; + } + } + + if (Object.keys(patch).length === 0) continue; + + session.mcpToolApprovals = { + ...session.mcpToolApprovals, + ...patch, + }; + + refreshes.push( + session.clientSideConnection + .extMethod(POSTHOG_METHODS.REFRESH_SESSION, { + mcpToolApprovals: patch, + }) + .then(() => undefined) + .catch(() => undefined), + ); + } + + await Promise.all(refreshes); + } + + private async applyMcpToolApprovalResponse( + session: ManagedSession | undefined, + toolCallId: string, + toolKey: string, + options: RequestPermissionRequest["options"], + response: RequestPermissionResponse, + ): Promise { + if ( + !session || + !toolCallId || + session.mcpToolApprovals?.[toolKey] !== "needs_approval" + ) { + return "none"; + } + + const selectedOption = getSelectedPermissionOption(options, response); + const selectedOptionId = + response.outcome?.outcome === "selected" + ? response.outcome.optionId + : undefined; + const shouldPersistApproval = + selectedOption?.kind === "allow_always" || + selectedOptionId === "allow_always"; + const shouldAllowOnce = + selectedOption?.kind === "allow_once" || selectedOptionId === "allow"; + + if (!shouldPersistApproval && !shouldAllowOnce) return "none"; + + const installation = session.toolInstallations[toolKey]; + if (!installation) { + return "failed"; + } + + try { + await this.agentAuthAdapter.updateMcpToolApproval( + session.config.credentials, + installation.installationId, + installation.toolName, + "approved", + ); + } catch { + return "failed"; + } + + if (shouldPersistApproval) { + session.mcpToolApprovals[toolKey] = "approved"; + session.temporaryMcpToolApprovals.delete(toolCallId); + return "persisted"; + } + + session.temporaryMcpToolApprovals.set(toolCallId, { + toolKey, + installationId: installation.installationId, + toolName: installation.toolName, + restoreApprovalState: "needs_approval", + }); + return "temporary"; + } + + private hasOtherTemporaryMcpApproval( + session: ManagedSession, + toolCallId: string, + approval: TemporaryMcpToolApproval, + ): boolean { + for (const [ + otherToolCallId, + otherApproval, + ] of session.temporaryMcpToolApprovals) { + if (otherToolCallId === toolCallId) continue; + if ( + otherApproval.installationId === approval.installationId && + otherApproval.toolName === approval.toolName + ) { + return true; + } + } + return false; + } + + private async restoreTemporaryMcpToolApproval( + session: ManagedSession, + toolCallId: string, + _reason: string, + ): Promise { + const approval = session.temporaryMcpToolApprovals.get(toolCallId); + if (!approval) return; + + session.temporaryMcpToolApprovals.delete(toolCallId); + + if (this.hasOtherTemporaryMcpApproval(session, toolCallId, approval)) { + return; + } + + if ( + session.mcpToolApprovals?.[approval.toolKey] !== + approval.restoreApprovalState + ) { + return; + } + + try { + await this.agentAuthAdapter.updateMcpToolApproval( + session.config.credentials, + approval.installationId, + approval.toolName, + approval.restoreApprovalState, + ); + } catch {} + } + + private async restoreTemporaryMcpToolApprovals( + session: ManagedSession, + reason: string, + ): Promise { + const toolCallIds = Array.from(session.temporaryMcpToolApprovals.keys()); + for (const toolCallId of toolCallIds) { + await this.restoreTemporaryMcpToolApproval(session, toolCallId, reason); + } + } + + private resolveCodexHookMcpToolKey( + session: ManagedSession, + input: CodexMcpApprovalHookInput, + ): string | null { + return resolveMcpStoreToolKey(input.toolName, { + approvals: session.mcpToolApprovals, + installations: session.toolInstallations, + }); + } + + private async approveCodexMcpToolUse( + taskRunId: string, + input: CodexMcpApprovalHookInput, + ): Promise { + const session = this.sessions.get(taskRunId); + if (!session) { + return { + action: "deny", + message: "MCP tool call cannot proceed because the session is gone.", + }; + } + + const toolKey = this.resolveCodexHookMcpToolKey(session, input); + if (!toolKey) { + return { action: "allow" }; + } + + const installation = session.toolInstallations[toolKey]; + const displayToolName = installation?.toolName ?? input.toolName; + const approvalState = session.mcpToolApprovals?.[toolKey]; + if (approvalState === "approved") { + return { action: "allow" }; + } + if (approvalState === "do_not_use") { + return { + action: "deny", + message: `Tool '${displayToolName}' has been disabled by the user`, + }; + } + if (approvalState !== "needs_approval") { + return { action: "allow" }; + } + + if (!input.toolUseId) { + return { + action: "deny", + message: + "MCP tool call cannot proceed because Codex did not provide a tool call id.", + }; + } + + const toolCallId = input.toolUseId; + const options: RequestPermissionRequest["options"] = [ + { kind: "allow_once", optionId: "allow", name: "Allow" }, + { + kind: "allow_always", + optionId: "allow_always", + name: "Always allow", + }, + { kind: "reject_once", optionId: "reject", name: "Reject" }, + ]; + + this.sleepService.release(taskRunId); + try { + const response = await new Promise( + (resolve, reject) => { + const key = `${taskRunId}:${toolCallId}`; + this.pendingPermissions.set(key, { + resolve, + reject, + taskRunId, + toolCallId, + }); + + this.emit(AgentServiceEvent.PermissionRequest, { + taskRunId, + options, + toolCall: { + toolCallId, + title: displayToolName, + kind: "other", + rawInput: { toolName: toolKey, name: displayToolName }, + }, + }); + }, + ); + + const selectedOption = getSelectedPermissionOption(options, response); + const shouldAllow = + selectedOption?.kind === "allow_once" || + selectedOption?.kind === "allow_always" || + isSelectedPermissionOption(response, "allow") || + isSelectedPermissionOption(response, "allow_always"); + + if (!shouldAllow) { + const feedback = (response._meta as Record | undefined) + ?.customInput; + return { + action: "deny", + message: + typeof feedback === "string" && feedback.trim() + ? `User refused permission to run tool with feedback: ${feedback.trim()}` + : "User refused permission to run tool", + }; + } + + const application = await this.applyMcpToolApprovalResponse( + session, + toolCallId, + toolKey, + options, + response, + ); + + if (application === "failed") { + return { + action: "deny", + message: + "MCP tool approval could not be applied before calling the tool.", + }; + } + + return { action: "allow" }; + } finally { + if (this.sessions.has(taskRunId)) { + this.sleepService.acquire(taskRunId); + } + } + } + + private async completeCodexMcpToolUse( + taskRunId: string, + input: CodexMcpApprovalHookInput, + ): Promise { + const session = this.sessions.get(taskRunId); + if (!session || !input.toolUseId) return; + await this.restoreTemporaryMcpToolApproval( + session, + input.toolUseId, + "codex_post_tool_use", + ); + } + /** * Cancel a pending permission request. * This resolves the promise with a "cancelled" outcome per ACP spec. @@ -748,6 +1091,7 @@ If a repository IS genuinely required, attach one in this priority order: debug: isDevBuild(), onLog: this.onAgentLog, }); + let codexMcpApprovalHookBridge: CodexMcpApprovalHookBridge | undefined; try { const systemPrompt = this.buildSystemPrompt( @@ -766,21 +1110,25 @@ If a repository IS genuinely required, attach one in this priority order: ); let codexHome: string | undefined; + let codexMcpApprovalHook: CodexMcpApprovalHookEnv | undefined; if (adapter === "codex") { - try { - codexHome = await prepareCodexHome({ - appDataPath: this.storagePaths.appDataPath, - taskRunId, - bundledSkillsDir, - log: this.log, - }); - } catch (err) { - // A skills-prep failure must not kill the session; Codex falls back - // to its default home and the user's own ~/.agents/skills. - this.log.warn("Failed to prepare codex home", { - error: err instanceof Error ? err.message : String(err), - }); - } + codexMcpApprovalHookBridge = new CodexMcpApprovalHookBridge( + { + preToolUse: (input) => + this.approveCodexMcpToolUse(taskRunId, input), + postToolUse: (input) => + this.completeCodexMcpToolUse(taskRunId, input), + }, + this.log, + ); + codexMcpApprovalHook = await codexMcpApprovalHookBridge.start(); + codexHome = await prepareCodexHome({ + appDataPath: this.storagePaths.appDataPath, + taskRunId, + bundledSkillsDir, + hookRuntimeCommand: process.execPath, + log: this.log, + }); } const acpConnection = await agent.run(taskId, taskRunId, { @@ -789,6 +1137,7 @@ If a repository IS genuinely required, attach one in this priority order: codexBinaryPath: adapter === "codex" ? this.getCodexBinaryPath() : undefined, codexHome, + codexMcpApprovalHook, model, reasoningEffort: adapter === "codex" ? effort : undefined, developerInstructions: @@ -1053,8 +1402,10 @@ If a repository IS genuinely required, attach one in this priority order: promptPending: false, configOptions, inFlightMcpToolCalls: new Map(), + temporaryMcpToolApprovals: new Map(), mcpToolApprovals: toolApprovals, toolInstallations, + codexMcpApprovalHookBridge, prAttributed: false, evaluatedPrUrls: new Set(), }; @@ -1067,6 +1418,11 @@ If a repository IS genuinely required, attach one in this priority order: } return session; } catch (err) { + await codexMcpApprovalHookBridge?.stop().catch(() => { + this.log.debug("Codex MCP approval hook bridge cleanup failed", { + taskRunId, + }); + }); try { await agent.cleanup(); } catch { @@ -1556,7 +1912,13 @@ For git operations while detached: }); } this.cancelInFlightMcpToolCalls(session); + await this.restoreTemporaryMcpToolApprovals(session, "session_cleanup"); this.sleepService.release(taskRunId); + await session.codexMcpApprovalHookBridge?.stop().catch(() => { + this.log.debug("Codex MCP approval hook bridge cleanup failed", { + taskRunId, + }); + }); try { await session.agent.cleanup(); } catch { @@ -1681,36 +2043,14 @@ For git operations while detached: }, ); - const approved = - response.outcome?.outcome === "selected" && - (response.outcome.optionId === "allow" || - response.outcome.optionId === "allow_always"); - if (approved && toolName) { - const session = service.sessions.get(taskRunId); - if ( - session?.mcpToolApprovals?.[toolName] === "needs_approval" && - session.toolInstallations[toolName] - ) { - const { installationId, toolName: rawToolName } = - session.toolInstallations[toolName]; - try { - await service.agentAuthAdapter.updateMcpToolApproval( - session.config.credentials, - installationId, - rawToolName, - "approved", - ); - session.mcpToolApprovals[toolName] = "approved"; - } catch (err) { - service.log.warn( - "Failed to update tool approval on backend", - { - toolName, - error: err instanceof Error ? err.message : String(err), - }, - ); - } - } + if (toolName) { + await service.applyMcpToolApprovalResponse( + service.sessions.get(taskRunId), + toolCallId, + toolName, + params.options, + response, + ); } return response; @@ -1793,6 +2133,13 @@ For git operations while detached: update.status === "failed" ) { session?.inFlightMcpToolCalls.delete(update.toolCallId); + if (session) { + await service.restoreTemporaryMcpToolApproval( + session, + update.toolCallId, + update.status, + ); + } service.mcpAppsService.notifyToolResult( toolName, update.toolCallId, diff --git a/packages/workspace-server/src/services/agent/codex-home.test.ts b/packages/workspace-server/src/services/agent/codex-home.test.ts index 985c64a401..d758facaa8 100644 --- a/packages/workspace-server/src/services/agent/codex-home.test.ts +++ b/packages/workspace-server/src/services/agent/codex-home.test.ts @@ -102,6 +102,39 @@ describe("prepareCodexHome", () => { expect(await readlink(link)).toBe(realpathSync(configPath)); }); + it("installs the Codex MCP approval hooks when a runtime command is provided", async () => { + const codexHome = await prepareCodexHome({ + appDataPath, + taskRunId, + bundledSkillsDir, + hookRuntimeCommand: "/usr/local/bin/node", + log: noopLog, + }); + + const hooksJson = JSON.parse( + readFileSync(path.join(codexHome, "hooks.json"), "utf-8"), + ) as { + hooks: { + PreToolUse: Array<{ + matcher: string; + hooks: Array<{ command: string }>; + }>; + PostToolUse: Array<{ + matcher: string; + hooks: Array<{ command: string }>; + }>; + }; + }; + expect(hooksJson.hooks.PreToolUse[0].matcher).toBe("^mcp__.*"); + expect(hooksJson.hooks.PostToolUse[0].matcher).toBe("^mcp__.*"); + expect(hooksJson.hooks.PreToolUse[0].hooks[0].command).toContain( + "/usr/local/bin/node", + ); + expect( + existsSync(path.join(codexHome, "hooks", "posthog-mcp-approval-hook.js")), + ).toBe(true); + }); + it("rebuilds the skills dir, dropping stale links", async () => { await createSkill(bundledSkillsDir, "first"); await prepareCodexHome({ diff --git a/packages/workspace-server/src/services/agent/codex-home.ts b/packages/workspace-server/src/services/agent/codex-home.ts index 723f367f73..bc068e99b8 100644 --- a/packages/workspace-server/src/services/agent/codex-home.ts +++ b/packages/workspace-server/src/services/agent/codex-home.ts @@ -1,6 +1,7 @@ import * as fs from "node:fs"; import * as os from "node:os"; import * as path from "node:path"; +import { installCodexMcpApprovalHook } from "@posthog/agent/adapters/codex/mcp-approval-hook"; import { findSkillDirs, getUserSkillsDir, @@ -55,6 +56,7 @@ export async function prepareCodexHome(options: { appDataPath: string; taskRunId: string; bundledSkillsDir: string; + hookRuntimeCommand?: string; log: AgentScopedLogger; }): Promise { const codexHome = getCodexHomeDir(options.appDataPath, options.taskRunId); @@ -92,5 +94,12 @@ export async function prepareCodexHome(options: { } } + if (options.hookRuntimeCommand) { + await installCodexMcpApprovalHook({ + codexHome, + runtimeCommand: options.hookRuntimeCommand, + }); + } + return codexHome; } diff --git a/packages/workspace-server/src/services/agent/schemas.ts b/packages/workspace-server/src/services/agent/schemas.ts index 493e79943e..456a81a257 100644 --- a/packages/workspace-server/src/services/agent/schemas.ts +++ b/packages/workspace-server/src/services/agent/schemas.ts @@ -277,6 +277,18 @@ export const respondToPermissionInput = z.object({ export type RespondToPermissionInput = z.infer; +const mcpToolApprovalStateSchema = z.enum([ + "approved", + "needs_approval", + "do_not_use", +]); + +export const updateMcpToolApprovalForActiveSessionsInput = z.object({ + installationId: z.string().min(1), + toolName: z.string().min(1), + approvalState: mcpToolApprovalStateSchema, +}); + // Permission cancellation input for tRPC export const cancelPermissionInput = z.object({ taskRunId: z.string(),