diff --git a/packages/agent/src/adapters/claude/hooks.test.ts b/packages/agent/src/adapters/claude/hooks.test.ts index 0c6415b4ce..d71790e632 100644 --- a/packages/agent/src/adapters/claude/hooks.test.ts +++ b/packages/agent/src/adapters/claude/hooks.test.ts @@ -16,6 +16,11 @@ import { createTaskHook, type EnrichedReadCache, } from "./hooks"; +import { + clearMcpToolApprovalCache, + clearMcpToolMetadataCache, + setMcpToolApprovalStates, +} from "./mcp/tool-metadata"; import type { PermissionCheckResult, SettingsManager, @@ -313,6 +318,81 @@ describe("createPreToolUseHook", () => { hookSpecificOutput: { permissionDecision: "ask" }, }); }); + + test("routes needs_approval MCP Store tool to canUseTool via ask", async () => { + clearMcpToolMetadataCache(); + clearMcpToolApprovalCache(); + setMcpToolApprovalStates({ + mcp__granola__query_granola_meetings: "needs_approval", + }); + const settingsManager = buildSettingsManagerStub({ decision: "ask" }); + const hook = createPreToolUseHook(settingsManager, logger); + const result = await hook( + buildPreToolUseHookInput("mcp__granola__query_granola_meetings", {}), + undefined, + { signal: new AbortController().signal }, + ); + + expect(result).toMatchObject({ + continue: true, + hookSpecificOutput: { + hookEventName: "PreToolUse", + permissionDecision: "ask", + }, + }); + clearMcpToolMetadataCache(); + clearMcpToolApprovalCache(); + }); + + test("denies do_not_use MCP Store tool", async () => { + clearMcpToolMetadataCache(); + clearMcpToolApprovalCache(); + setMcpToolApprovalStates({ + mcp__granola__query_granola_meetings: "do_not_use", + }); + const settingsManager = buildSettingsManagerStub({ decision: "ask" }); + const hook = createPreToolUseHook(settingsManager, logger); + const result = await hook( + buildPreToolUseHookInput("mcp__granola__query_granola_meetings", {}), + undefined, + { signal: new AbortController().signal }, + ); + + expect(result).toMatchObject({ + continue: true, + hookSpecificOutput: { + hookEventName: "PreToolUse", + permissionDecision: "deny", + }, + }); + clearMcpToolMetadataCache(); + clearMcpToolApprovalCache(); + }); + + test("respects an explicit allow rule for a needs_approval MCP Store tool", async () => { + clearMcpToolMetadataCache(); + clearMcpToolApprovalCache(); + setMcpToolApprovalStates({ + mcp__granola__query_granola_meetings: "needs_approval", + }); + const settingsManager = buildSettingsManagerStub({ + decision: "allow", + rule: "mcp__granola__query_granola_meetings", + source: "allow", + }); + const hook = createPreToolUseHook(settingsManager, logger); + const result = await hook( + buildPreToolUseHookInput("mcp__granola__query_granola_meetings", {}), + undefined, + { signal: new AbortController().signal }, + ); + + expect(result).toMatchObject({ + hookSpecificOutput: { permissionDecision: "allow" }, + }); + clearMcpToolMetadataCache(); + clearMcpToolApprovalCache(); + }); }); describe("createSignedCommitGuardHook", () => { diff --git a/packages/agent/src/adapters/claude/hooks.ts b/packages/agent/src/adapters/claude/hooks.ts index c3d21423a3..cb49001679 100644 --- a/packages/agent/src/adapters/claude/hooks.ts +++ b/packages/agent/src/adapters/claude/hooks.ts @@ -7,6 +7,7 @@ import type { Logger } from "../../utils/logger"; import { SIGNED_COMMIT_QUALIFIED_TOOL_NAME } from "../signed-commit-shared"; import { stripCatLineNumbers } from "./conversion/sdk-to-acp"; import type { TaskState } from "./conversion/task-state"; +import { getMcpToolApprovalState } from "./mcp/tool-metadata"; import { extractPostHogSubTool, isPostHogDestructiveSubTool, @@ -410,6 +411,45 @@ export const createPreToolUseHook = ); } + // MCP Store per-tool approval. Cloud sessions run in bypassPermissions, + // which skips canUseTool entirely — so this PreToolUse hook is the only + // 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. + const mcpApprovalState = getMcpToolApprovalState(toolName); + if (mcpApprovalState === "do_not_use") { + logger.info(`[PreToolUseHook] Blocking do_not_use MCP tool: ${toolName}`); + return { + continue: true, + hookSpecificOutput: { + hookEventName: "PreToolUse" as const, + permissionDecision: "deny" as const, + permissionDecisionReason: + "This tool has been blocked. To re-enable it, go to Settings > MCP Servers in PostHog Code.", + }, + }; + } + if ( + mcpApprovalState === "needs_approval" && + permissionCheck.decision !== "allow" && + permissionCheck.decision !== "deny" + ) { + logger.info( + `[PreToolUseHook] Routing needs_approval MCP tool to canUseTool: ${toolName}`, + ); + return { + continue: true, + hookSpecificOutput: { + hookEventName: "PreToolUse" as const, + permissionDecision: "ask" as const, + permissionDecisionReason: + "This MCP tool requires approval before it can be used.", + }, + }; + } + // Defer destructive PostHog exec sub-tools to canUseTool so the // sub-tool gate can re-prompt. Returning `{ continue: true }` is // not enough — the SDK then falls back to its default permission diff --git a/packages/agent/src/adapters/claude/mcp/tool-metadata.test.ts b/packages/agent/src/adapters/claude/mcp/tool-metadata.test.ts index d896520450..ff4361f3b2 100644 --- a/packages/agent/src/adapters/claude/mcp/tool-metadata.test.ts +++ b/packages/agent/src/adapters/claude/mcp/tool-metadata.test.ts @@ -1,8 +1,10 @@ import { beforeEach, describe, expect, it } from "vitest"; import { + clearMcpToolApprovalCache, clearMcpToolMetadataCache, getMcpToolApprovalState, getMcpToolMetadata, + getMcpToolMetadataKey, isMcpToolReadOnly, sanitizeMcpServerName, setMcpToolApprovalStates, @@ -11,6 +13,7 @@ import { describe("tool-metadata approval states", () => { beforeEach(() => { clearMcpToolMetadataCache(); + clearMcpToolApprovalCache(); }); describe("setMcpToolApprovalStates", () => { @@ -56,12 +59,50 @@ describe("tool-metadata approval states", () => { expect(getMcpToolApprovalState("mcp__server__unknown")).toBeUndefined(); }); + it("normalizes MCP server names before looking up approval state", () => { + setMcpToolApprovalStates({ + mcp__Granola_Meetings__query_granola_meetings: "needs_approval", + }); + + expect( + getMcpToolApprovalState( + "mcp__Granola Meetings__query_granola_meetings", + ), + ).toBe("needs_approval"); + }); + + it("resolves a bare upstream tool name when it uniquely maps to an MCP Store tool", () => { + setMcpToolApprovalStates({ + mcp__Granola_Meetings__query_granola_meetings: "needs_approval", + }); + + expect(getMcpToolApprovalState("query_granola_meetings")).toBe( + "needs_approval", + ); + expect(getMcpToolMetadataKey("query_granola_meetings")).toBe( + "mcp__Granola_Meetings__query_granola_meetings", + ); + }); + it("returns the correct state", () => { setMcpToolApprovalStates({ mcp__s__t: "needs_approval", }); expect(getMcpToolApprovalState("mcp__s__t")).toBe("needs_approval"); }); + + it("survives a metadata cache clear (MCP server refresh/reconnect)", () => { + setMcpToolApprovalStates({ + mcp__granola__query_granola_meetings: "needs_approval", + }); + + // A server reconnect wipes the metadata cache mid-session. + clearMcpToolMetadataCache(); + + expect( + getMcpToolApprovalState("mcp__granola__query_granola_meetings"), + ).toBe("needs_approval"); + }); }); describe("isMcpToolReadOnly with approval states", () => { diff --git a/packages/agent/src/adapters/claude/mcp/tool-metadata.ts b/packages/agent/src/adapters/claude/mcp/tool-metadata.ts index 0f13d589ac..91bddc6e4a 100644 --- a/packages/agent/src/adapters/claude/mcp/tool-metadata.ts +++ b/packages/agent/src/adapters/claude/mcp/tool-metadata.ts @@ -1,4 +1,11 @@ import type { McpServerStatus, Query } from "@anthropic-ai/claude-agent-sdk"; +import { + buildMcpToolKey, + normalizeMcpToolKey, + parseMcpToolKey, + resolveMcpStoreToolKey, + sanitizeMcpServerName, +} from "../../../mcp-store/tool-keys"; import { Logger } from "../../../utils/logger"; export type McpToolApprovalState = "approved" | "needs_approval" | "do_not_use"; @@ -9,22 +16,25 @@ export type McpToolApprovals = Record; export interface McpToolMetadata { readOnly: boolean; name: string; + serverName?: string; description?: string; approvalState?: McpToolApprovalState; } const mcpToolMetadataCache: Map = new Map(); +// Per-tool approval state lives in its own store, keyed by normalized tool key. +// It must outlive the metadata cache: `clearMcpToolMetadataCache()` runs on +// every MCP server refresh/reconnect, and approval state is session config +// (set once from the start payload), not volatile per-connection metadata. +// Keeping them separate means a reconnect can't silently drop approvals and +// let a needs_approval tool slip through the gate. +const mcpToolApprovalCache: Map = new Map(); + const PENDING_RETRY_INTERVAL_MS = 1_000; const PENDING_MAX_RETRIES = 10; -export function sanitizeMcpServerName(name: string): string { - return name.replace(/[^a-zA-Z0-9_-]/g, "_"); -} - -function buildToolKey(serverName: string, toolName: string): string { - return `mcp__${serverName}__${toolName}`; -} +export { sanitizeMcpServerName }; function delay(ms: number): Promise { return new Promise((resolve) => setTimeout(resolve, ms)); @@ -56,15 +66,17 @@ export async function fetchMcpToolMetadata( let readOnlyCount = 0; for (const tool of server.tools) { - const toolKey = buildToolKey(server.name, tool.name); + const toolKey = buildMcpToolKey(server.name, tool.name); const readOnly = tool.annotations?.readOnly === true; const existing = mcpToolMetadataCache.get(toolKey); mcpToolMetadataCache.set(toolKey, { readOnly, name: tool.name, + serverName: server.name, description: tool.description, - approvalState: existing?.approvalState, + approvalState: + mcpToolApprovalCache.get(toolKey) ?? existing?.approvalState, }); if (readOnly) readOnlyCount++; } @@ -99,11 +111,24 @@ export async function fetchMcpToolMetadata( export function getMcpToolMetadata( toolName: string, ): McpToolMetadata | undefined { - return mcpToolMetadataCache.get(toolName); + const key = getMcpToolMetadataKey(toolName); + return key ? mcpToolMetadataCache.get(key) : undefined; +} + +export function getMcpToolMetadataKey(toolName: string): string | undefined { + const resolvedKey = resolveMcpStoreToolKey(toolName, { + approvals: Object.fromEntries( + [...mcpToolMetadataCache.keys()].map((key) => [key, true]), + ), + }); + const normalizedToolName = normalizeMcpToolKey(toolName); + if (mcpToolMetadataCache.has(toolName)) return toolName; + if (mcpToolMetadataCache.has(normalizedToolName)) return normalizedToolName; + return resolvedKey ?? undefined; } export function isMcpToolReadOnly(toolName: string): boolean { - const metadata = mcpToolMetadataCache.get(toolName); + const metadata = getMcpToolMetadata(toolName); return metadata?.readOnly === true; } @@ -125,18 +150,41 @@ export function getCachedMcpTools(): McpToolMetadata[] { export function getMcpToolApprovalState( toolName: string, ): McpToolApprovalState | undefined { - return mcpToolMetadataCache.get(toolName)?.approvalState; + if (mcpToolApprovalCache.size === 0) { + return undefined; + } + // Direct/normalized hit first, then resolve server-name or title variants + // against the known approval keys. Reads the approval store (not the + // metadata cache) so it survives MCP metadata refreshes. + const direct = + mcpToolApprovalCache.get(toolName) ?? + mcpToolApprovalCache.get(normalizeMcpToolKey(toolName)); + if (direct) { + return direct; + } + const resolvedKey = resolveMcpStoreToolKey(toolName, { + approvals: Object.fromEntries( + [...mcpToolApprovalCache.keys()].map((key) => [key, true]), + ), + }); + return resolvedKey ? mcpToolApprovalCache.get(resolvedKey) : undefined; } export function setMcpToolApprovalStates(approvals: McpToolApprovals): void { for (const [toolKey, approvalState] of Object.entries(approvals)) { - const existing = mcpToolMetadataCache.get(toolKey); + const normalizedToolKey = normalizeMcpToolKey(toolKey); + mcpToolApprovalCache.set(normalizedToolKey, approvalState); + // Mirror onto the metadata cache entry when present so UI/debug reads stay + // consistent; the approval store above remains the source of truth. + const existing = mcpToolMetadataCache.get(normalizedToolKey); if (existing) { existing.approvalState = approvalState; } else { - mcpToolMetadataCache.set(toolKey, { + const parsed = parseMcpToolKey(normalizedToolKey); + mcpToolMetadataCache.set(normalizedToolKey, { readOnly: false, - name: toolKey, + name: parsed?.toolName ?? normalizedToolKey, + serverName: parsed?.serverName, approvalState, }); } @@ -146,3 +194,9 @@ export function setMcpToolApprovalStates(approvals: McpToolApprovals): void { export function clearMcpToolMetadataCache(): void { mcpToolMetadataCache.clear(); } + +/** Reset per-tool approval state. Approvals survive metadata refreshes, so + * this is only for session teardown and tests — not the per-refresh clear. */ +export function clearMcpToolApprovalCache(): void { + mcpToolApprovalCache.clear(); +} 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 528a219e18..345a8bb9fa 100644 --- a/packages/agent/src/adapters/claude/permissions/permission-handlers.test.ts +++ b/packages/agent/src/adapters/claude/permissions/permission-handlers.test.ts @@ -79,6 +79,62 @@ describe("canUseTool MCP approval enforcement", () => { ); }); + it("routes needs_approval MCP tools when the runtime tool name has an unsanitized server name", async () => { + setMcpToolApprovalStates({ + mcp__Granola_Meetings__query_granola_meetings: "needs_approval", + }); + + const context = createContext( + "mcp__Granola Meetings__query_granola_meetings", + ); + const result = await canUseTool(context); + + expect(result.behavior).toBe("allow"); + expect(context.client.requestPermission).toHaveBeenCalledWith( + expect.objectContaining({ + toolCall: expect.objectContaining({ + rawInput: expect.objectContaining({ + toolName: "mcp__Granola_Meetings__query_granola_meetings", + }), + }), + }), + ); + }); + + it("routes needs_approval MCP tools when Claude reports the bare upstream tool name", async () => { + setMcpToolApprovalStates({ + mcp__Granola_Meetings__query_granola_meetings: "needs_approval", + }); + + const context = createContext("query_granola_meetings"); + const result = await canUseTool(context); + + expect(result.behavior).toBe("allow"); + expect(context.client.requestPermission).toHaveBeenCalledWith( + expect.objectContaining({ + toolCall: expect.objectContaining({ + title: + "The agent wants to call query_granola_meetings (Granola_Meetings)", + rawInput: expect.objectContaining({ + toolName: "mcp__Granola_Meetings__query_granola_meetings", + }), + }), + }), + ); + }); + + it("does not treat built-in Claude tools as bare MCP Store matches", async () => { + setMcpToolApprovalStates({ + mcp__server__Read: "needs_approval", + }); + + const context = createContext("Read"); + const result = await canUseTool(context); + + expect(result.behavior).toBe("allow"); + expect(context.client.requestPermission).not.toHaveBeenCalled(); + }); + it("allows approved MCP tools through normal flow", async () => { setMcpToolApprovalStates({ mcp__server__approved_tool: "approved", diff --git a/packages/agent/src/adapters/claude/permissions/permission-handlers.ts b/packages/agent/src/adapters/claude/permissions/permission-handlers.ts index 485c266c80..925fb2bd83 100644 --- a/packages/agent/src/adapters/claude/permissions/permission-handlers.ts +++ b/packages/agent/src/adapters/claude/permissions/permission-handlers.ts @@ -12,6 +12,7 @@ import { toolInfoFromToolUse } from "../conversion/tool-use-to-acp"; import { getMcpToolApprovalState, getMcpToolMetadata, + getMcpToolMetadataKey, } from "../mcp/tool-metadata"; import { getClaudePlansDir, @@ -25,7 +26,15 @@ import { OPTION_PREFIX, type QuestionItem, } from "../questions/utils"; -import { isToolAllowedForMode, WRITE_TOOLS } from "../tools"; +import { + AGENT_TOOLS, + BASH_TOOLS, + isToolAllowedForMode, + READ_TOOLS, + SEARCH_TOOLS, + WEB_TOOLS, + WRITE_TOOLS, +} from "../tools"; import type { Session } from "../types"; import { buildExitPlanModePermissionOptions, @@ -429,13 +438,33 @@ function parseMcpToolName(toolName: string): { }; } +const CLAUDE_BUILTIN_TOOLS = new Set([ + ...READ_TOOLS, + ...WRITE_TOOLS, + ...BASH_TOOLS, + ...SEARCH_TOOLS, + ...WEB_TOOLS, + ...AGENT_TOOLS, + "EnterPlanMode", + "ExitPlanMode", + "AskUserQuestion", +]); + +function shouldCheckMcpStoreApproval(toolName: string): boolean { + return toolName.startsWith("mcp__") || !CLAUDE_BUILTIN_TOOLS.has(toolName); +} + async function handleMcpApprovalFlow( context: ToolHandlerContext, ): Promise { const { toolName, toolInput, toolUseID, client, sessionId } = context; - const { serverName, tool: displayTool } = parseMcpToolName(toolName); + const approvalToolName = getMcpToolMetadataKey(toolName) ?? toolName; const metadata = getMcpToolMetadata(toolName); + const { serverName: parsedServerName, tool: parsedToolName } = + parseMcpToolName(approvalToolName); + const serverName = metadata?.serverName ?? parsedServerName; + const displayTool = metadata?.name ?? parsedToolName; const description = metadata?.description ? `\n\n${metadata.description}` : ""; @@ -463,7 +492,10 @@ async function handleMcpApprovalFlow( content: description ? [{ type: "content" as const, content: text(description) }] : [], - rawInput: { ...(toolInput as Record), toolName }, + rawInput: { + ...(toolInput as Record), + toolName: approvalToolName, + }, }, }); @@ -669,7 +701,7 @@ export async function canUseTool( } } - if (toolName.startsWith("mcp__")) { + if (shouldCheckMcpStoreApproval(toolName)) { const approvalState = getMcpToolApprovalState(toolName); if (approvalState === "do_not_use") { @@ -682,7 +714,9 @@ export async function canUseTool( if (approvalState === "needs_approval") { return handleMcpApprovalFlow(context); } + } + if (toolName.startsWith("mcp__")) { if (isPostHogExecTool(toolName)) { const subTool = extractPostHogSubTool(toolInput); if (subTool && isPostHogDestructiveSubTool(subTool)) { diff --git a/packages/agent/src/adapters/codex/codex-agent.ts b/packages/agent/src/adapters/codex/codex-agent.ts index f1f053fdd9..66f04d6600 100644 --- a/packages/agent/src/adapters/codex/codex-agent.ts +++ b/packages/agent/src/adapters/codex/codex-agent.ts @@ -57,6 +57,10 @@ import { isCodexNativeMode, type PermissionMode, } from "../../execution-mode"; +import type { + McpToolApprovalsConfig, + McpToolInstallationsConfig, +} from "../../server/schemas"; import type { PostHogAPIConfig, ProcessSpawnedCallback } from "../../types"; import { isCloudRun } from "../../utils/common"; import { resolveGithubToken } from "../../utils/github-token"; @@ -122,6 +126,8 @@ interface NewSessionMeta { disableBuiltInTools?: boolean; allowedDomains?: string[]; jsonSchema?: Record | null; + mcpToolApprovals?: McpToolApprovalsConfig; + mcpToolInstallations?: McpToolInstallationsConfig; } export interface CodexAcpAgentOptions { @@ -462,6 +468,8 @@ export class CodexAcpAgent extends BaseAcpAgent { modeId: response.modes?.currentModeId ?? "auto", modelId: modelIdFromConfigOptions(response.configOptions), permissionMode: requestedPermissionMode, + mcpToolApprovals: meta?.mcpToolApprovals, + mcpToolInstallations: meta?.mcpToolInstallations, }); this.sessionId = response.sessionId; this.sessionState.configOptions = response.configOptions ?? []; @@ -518,6 +526,8 @@ export class CodexAcpAgent extends BaseAcpAgent { taskId: resolveTaskId(meta), modeId: response.modes?.currentModeId ?? "auto", permissionMode: currentPermissionMode, + mcpToolApprovals: meta?.mcpToolApprovals, + mcpToolInstallations: meta?.mcpToolInstallations, }); this.sessionId = params.sessionId; this.sessionState.configOptions = response.configOptions ?? []; @@ -565,6 +575,8 @@ export class CodexAcpAgent extends BaseAcpAgent { taskId: resolveTaskId(meta), modeId: loadResponse.modes?.currentModeId ?? "auto", permissionMode: currentPermissionMode, + mcpToolApprovals: meta?.mcpToolApprovals, + mcpToolInstallations: meta?.mcpToolInstallations, }); this.sessionId = params.sessionId; this.sessionState.configOptions = loadResponse.configOptions ?? []; @@ -612,6 +624,8 @@ export class CodexAcpAgent extends BaseAcpAgent { taskId: resolveTaskId(meta), modeId: newResponse.modes?.currentModeId ?? "auto", permissionMode: requestedPermissionMode, + mcpToolApprovals: meta?.mcpToolApprovals, + mcpToolInstallations: meta?.mcpToolInstallations, }); this.sessionId = newResponse.sessionId; this.sessionState.configOptions = newResponse.configOptions ?? []; diff --git a/packages/agent/src/adapters/codex/codex-client.test.ts b/packages/agent/src/adapters/codex/codex-client.test.ts index 56a8d8b1d0..010915c78b 100644 --- a/packages/agent/src/adapters/codex/codex-client.test.ts +++ b/packages/agent/src/adapters/codex/codex-client.test.ts @@ -112,6 +112,138 @@ describe("createCodexClient readTextFile", () => { }); }); +describe("createCodexClient MCP Store permissions", () => { + const logger = new Logger({ debug: false, prefix: "[test]" }); + const ALLOW_OPTIONS = [ + { kind: "allow_once" as const, optionId: "allow", name: "Allow" }, + { kind: "reject_once" as const, optionId: "reject", name: "Reject" }, + ]; + + function makePermissionUpstream( + response = { + outcome: { outcome: "selected" as const, optionId: "allow" }, + }, + ): AgentSideConnection & { requestPermission: ReturnType } { + return { + sessionUpdate: vi.fn(async () => {}), + requestPermission: vi.fn(async () => response), + readTextFile: vi.fn(), + writeTextFile: vi.fn(), + createTerminal: vi.fn(), + terminalOutput: vi.fn(), + releaseTerminal: vi.fn(), + waitForTerminalExit: vi.fn(), + killTerminal: vi.fn(), + extMethod: vi.fn(), + extNotification: vi.fn(), + } as unknown as AgentSideConnection & { + requestPermission: ReturnType; + }; + } + + test("relays needs_approval MCP tools by bare title even in full-access mode", async () => { + const toolName = "mcp__Granola__query_granola_meetings"; + const upstream = makePermissionUpstream(); + const sessionState = createSessionState("sess", "/tmp", { + permissionMode: "full-access", + mcpToolApprovals: { [toolName]: "needs_approval" }, + mcpToolInstallations: { + [toolName]: { + installationId: "inst-1", + toolName: "query_granola_meetings", + }, + }, + }); + const client = createCodexClient(upstream, logger, sessionState); + + const result = await client.requestPermission?.({ + options: ALLOW_OPTIONS, + toolCall: { + toolCallId: "tc-1", + title: "query_granola_meetings", + kind: "other", + rawInput: { search: "today" }, + }, + } as never); + + expect(result?.outcome.outcome).toBe("selected"); + expect(upstream.requestPermission).toHaveBeenCalledTimes(1); + expect(upstream.requestPermission).toHaveBeenCalledWith( + expect.objectContaining({ + toolCall: expect.objectContaining({ + rawInput: expect.objectContaining({ + search: "today", + toolName, + }), + }), + }), + ); + expect(sessionState.mcpToolApprovals?.[toolName]).toBe("approved"); + }); + + test("relays needs_approval MCP tools with unsanitized MCP server names", async () => { + const toolName = "mcp__Granola_Meetings__query_granola_meetings"; + const upstream = makePermissionUpstream(); + const sessionState = createSessionState("sess", "/tmp", { + permissionMode: "full-access", + mcpToolApprovals: { [toolName]: "needs_approval" }, + mcpToolInstallations: { + [toolName]: { + installationId: "inst-1", + toolName: "query_granola_meetings", + }, + }, + }); + const client = createCodexClient(upstream, logger, sessionState); + + await client.requestPermission?.({ + options: ALLOW_OPTIONS, + toolCall: { + toolCallId: "tc-1", + title: "mcp__Granola Meetings__query_granola_meetings", + kind: "other", + }, + } as never); + + expect(upstream.requestPermission).toHaveBeenCalledWith( + expect.objectContaining({ + toolCall: expect.objectContaining({ + rawInput: expect.objectContaining({ toolName }), + }), + }), + ); + expect(sessionState.mcpToolApprovals?.[toolName]).toBe("approved"); + }); + + test("cancels blocked MCP tools before they reach the upstream permission flow", async () => { + const toolName = "mcp__Granola__query_granola_meetings"; + const upstream = makePermissionUpstream(); + const sessionState = createSessionState("sess", "/tmp", { + permissionMode: "full-access", + mcpToolApprovals: { [toolName]: "do_not_use" }, + mcpToolInstallations: { + [toolName]: { + installationId: "inst-1", + toolName: "query_granola_meetings", + }, + }, + }); + const client = createCodexClient(upstream, logger, sessionState); + + const result = await client.requestPermission?.({ + options: ALLOW_OPTIONS, + toolCall: { + toolCallId: "tc-1", + title: "query_granola_meetings", + kind: "other", + }, + } as never); + + expect(result?.outcome.outcome).toBe("cancelled"); + expect(upstream.requestPermission).not.toHaveBeenCalled(); + }); +}); + describe("createCodexClient onStructuredOutput", () => { const logger = new Logger({ debug: false, prefix: "[test]" }); const sessionState = createSessionState("sess", "/tmp"); diff --git a/packages/agent/src/adapters/codex/codex-client.ts b/packages/agent/src/adapters/codex/codex-client.ts index ebe05a4fe2..9fcb2156fc 100644 --- a/packages/agent/src/adapters/codex/codex-client.ts +++ b/packages/agent/src/adapters/codex/codex-client.ts @@ -34,6 +34,7 @@ import { type FileEnrichmentDeps, } from "../../enrichment/file-enricher"; import type { PermissionMode } from "../../execution-mode"; +import { resolveMcpStoreToolKey } from "../../mcp-store/tool-keys"; import type { Logger } from "../../utils/logger"; import type { CodexSessionState } from "./session-state"; import { @@ -77,6 +78,57 @@ function toRecord(value: unknown): Record | null { return null; } +function getMcpToolNameFromPermissionRequest( + sessionState: CodexSessionState, + params: RequestPermissionRequest, +): string | null { + const rawInput = toRecord(params.toolCall?.rawInput); + const meta = toRecord(params.toolCall?._meta); + const claudeCode = toRecord(meta?.claudeCode); + const candidates = [ + typeof rawInput?.toolName === "string" ? rawInput.toolName : null, + typeof rawInput?.name === "string" ? rawInput.name : null, + typeof meta?.toolName === "string" ? meta.toolName : null, + typeof claudeCode?.toolName === "string" ? claudeCode.toolName : null, + typeof params.toolCall?.title === "string" ? params.toolCall.title : null, + ]; + + for (const candidate of candidates) { + const toolName = resolveMcpStoreToolKey(candidate, { + approvals: sessionState.mcpToolApprovals, + installations: sessionState.mcpToolInstallations, + }); + if (toolName) return toolName; + } + + return null; +} + +function withMcpToolName( + params: RequestPermissionRequest, + toolName: string, +): RequestPermissionRequest { + if (!params.toolCall) return params; + return { + ...params, + toolCall: { + ...params.toolCall, + rawInput: { + ...(toRecord(params.toolCall.rawInput) ?? {}), + toolName, + }, + }, + }; +} + +function isAllowResponse(response: RequestPermissionResponse): boolean { + return ( + response.outcome?.outcome === "selected" && + (response.outcome.optionId === "allow" || + response.outcome.optionId === "allow_always") + ); +} + const AUTO_APPROVED_KINDS: Record> = { default: new Set(["read", "search", "fetch", "think"]), acceptEdits: new Set(["read", "edit", "search", "fetch", "think"]), @@ -143,6 +195,36 @@ export function createCodexClient( params: RequestPermissionRequest, ): Promise { const kind = params.toolCall?.kind as ToolKind | null | undefined; + const mcpToolName = getMcpToolNameFromPermissionRequest( + sessionState, + params, + ); + const mcpApprovalState = mcpToolName + ? sessionState.mcpToolApprovals?.[mcpToolName] + : undefined; + + if (mcpToolName && mcpApprovalState === "do_not_use") { + return { + outcome: { outcome: "cancelled" }, + _meta: { + message: + "This tool has been blocked. To re-enable it, go to Settings > MCP Servers in PostHog Code.", + }, + }; + } + + if (mcpToolName && mcpApprovalState === "needs_approval") { + const response = await upstreamClient.requestPermission( + withMcpToolName(params, mcpToolName), + ); + if (isAllowResponse(response)) { + sessionState.mcpToolApprovals = { + ...sessionState.mcpToolApprovals, + [mcpToolName]: "approved", + }; + } + return response; + } if (shouldAutoApprove(sessionState.permissionMode, kind)) { logger.debug("Auto-approving permission", { diff --git a/packages/agent/src/adapters/codex/session-state.ts b/packages/agent/src/adapters/codex/session-state.ts index 9aa8694a85..fa3ba43bc6 100644 --- a/packages/agent/src/adapters/codex/session-state.ts +++ b/packages/agent/src/adapters/codex/session-state.ts @@ -1,5 +1,9 @@ import type { SessionConfigOption } from "@agentclientprotocol/sdk"; import type { PermissionMode } from "../../execution-mode"; +import type { + McpToolApprovalsConfig, + McpToolInstallationsConfig, +} from "../../server/schemas"; import type { ContextBreakdownBaseline } from "../claude/context-breakdown"; export interface CodexUsage { @@ -22,6 +26,8 @@ export interface CodexSessionState { permissionMode: PermissionMode; taskRunId?: string; taskId?: string; + mcpToolApprovals?: McpToolApprovalsConfig; + mcpToolInstallations?: McpToolInstallationsConfig; } export function createSessionState( @@ -33,6 +39,8 @@ export function createSessionState( modeId?: string; modelId?: string; permissionMode?: PermissionMode; + mcpToolApprovals?: McpToolApprovalsConfig; + mcpToolInstallations?: McpToolInstallationsConfig; }, ): CodexSessionState { return { @@ -50,6 +58,8 @@ export function createSessionState( permissionMode: opts?.permissionMode ?? "auto", taskRunId: opts?.taskRunId, taskId: opts?.taskId, + mcpToolApprovals: opts?.mcpToolApprovals, + mcpToolInstallations: opts?.mcpToolInstallations, }; } @@ -66,6 +76,8 @@ export function resetSessionState( modeId?: string; modelId?: string; permissionMode?: PermissionMode; + mcpToolApprovals?: McpToolApprovalsConfig; + mcpToolInstallations?: McpToolInstallationsConfig; }, ): void { state.sessionId = sessionId; @@ -85,6 +97,8 @@ export function resetSessionState( state.permissionMode = opts?.permissionMode ?? "auto"; state.taskRunId = opts?.taskRunId; state.taskId = opts?.taskId; + state.mcpToolApprovals = opts?.mcpToolApprovals; + state.mcpToolInstallations = opts?.mcpToolInstallations; } export function resetUsage(state: CodexSessionState): void { diff --git a/packages/agent/src/mcp-store/tool-keys.ts b/packages/agent/src/mcp-store/tool-keys.ts new file mode 100644 index 0000000000..605be8a1b3 --- /dev/null +++ b/packages/agent/src/mcp-store/tool-keys.ts @@ -0,0 +1,160 @@ +export type McpToolApprovalMap = Record; + +export interface McpToolInstallationRef { + toolName: string; +} + +export type McpToolInstallationMap = Record; + +interface ParsedMcpToolKey { + serverName: string; + toolName: string; +} + +export function sanitizeMcpServerName(name: string): string { + return name.replace(/[^a-zA-Z0-9_-]/g, "_"); +} + +function serverIdentity(name: string): string { + return sanitizeMcpServerName(name) + .replace(/_+/g, "_") + .replace(/^_+|_+$/g, "") + .toLowerCase(); +} + +export function buildMcpToolKey(serverName: string, toolName: string): string { + return `mcp__${sanitizeMcpServerName(serverName)}__${toolName}`; +} + +export function parseMcpToolKey(toolName: string): ParsedMcpToolKey | null { + if (!toolName.startsWith("mcp__")) { + return null; + } + + const parts = toolName.split("__"); + if (parts.length < 3) { + return null; + } + + return { + serverName: parts[1], + toolName: parts.slice(2).join("__"), + }; +} + +export function normalizeMcpToolKey(toolName: string): string { + const parsed = parseMcpToolKey(toolName); + return parsed + ? buildMcpToolKey(parsed.serverName, parsed.toolName) + : toolName; +} + +function hasApproval(approvals: McpToolApprovalMap | undefined, key: string) { + return Object.hasOwn(approvals ?? {}, key); +} + +function approvalKeys(approvals: McpToolApprovalMap | undefined): string[] { + return Object.keys(approvals ?? {}); +} + +function uniqueMatch(matches: string[]): string | null { + const unique = [...new Set(matches)]; + return unique.length === 1 ? unique[0] : null; +} + +export function resolveMcpStoreToolKeyFromParts(params: { + serverName?: string | null; + toolName: string; + approvals?: McpToolApprovalMap; + installations?: McpToolInstallationMap; +}): string | null { + const { serverName, toolName, approvals, installations } = params; + + if (serverName) { + const normalized = buildMcpToolKey(serverName, toolName); + if (hasApproval(approvals, normalized)) { + return normalized; + } + + const wantedServer = serverIdentity(serverName); + const matches = approvalKeys(approvals).filter((key) => { + const parsed = parseMcpToolKey(key); + return ( + parsed?.toolName === toolName && + serverIdentity(parsed.serverName) === wantedServer + ); + }); + const match = uniqueMatch(matches); + if (match) return match; + } + + const installationMatches = Object.entries(installations ?? {}) + .filter( + ([key, installation]) => + installation.toolName === toolName && hasApproval(approvals, key), + ) + .map(([key]) => key); + const installationMatch = uniqueMatch(installationMatches); + if (installationMatch) return installationMatch; + + const approvalMatches = approvalKeys(approvals).filter( + (key) => parseMcpToolKey(key)?.toolName === toolName, + ); + return uniqueMatch(approvalMatches); +} + +function parseApprovalTitle(title: string): { + serverName: string; + toolName: string; +} | null { + const match = title.match(/^The agent wants to call\s+(.+?)\s+\((.+)\)$/); + if (!match) { + return null; + } + return { toolName: match[1], serverName: match[2] }; +} + +export function resolveMcpStoreToolKey( + candidate: string | null | undefined, + params: { + approvals?: McpToolApprovalMap; + installations?: McpToolInstallationMap; + }, +): string | null { + if (!candidate) return null; + + const trimmed = candidate.trim(); + if (!trimmed) return null; + + if (hasApproval(params.approvals, trimmed)) { + return trimmed; + } + + const normalized = normalizeMcpToolKey(trimmed); + if (normalized !== trimmed && hasApproval(params.approvals, normalized)) { + return normalized; + } + + const parsed = parseMcpToolKey(trimmed); + if (parsed) { + return resolveMcpStoreToolKeyFromParts({ + ...params, + serverName: parsed.serverName, + toolName: parsed.toolName, + }); + } + + const titleParts = parseApprovalTitle(trimmed); + if (titleParts) { + return resolveMcpStoreToolKeyFromParts({ + ...params, + serverName: titleParts.serverName, + toolName: titleParts.toolName, + }); + } + + return resolveMcpStoreToolKeyFromParts({ + ...params, + toolName: trimmed, + }); +} diff --git a/packages/agent/src/posthog-api.ts b/packages/agent/src/posthog-api.ts index fc8b9ccd5a..515ee5a9a9 100644 --- a/packages/agent/src/posthog-api.ts +++ b/packages/agent/src/posthog-api.ts @@ -190,6 +190,20 @@ export class PostHogAPIClient { ); } + async updateMcpToolApproval( + installationId: string, + toolName: string, + approvalState: "approved" | "needs_approval" | "do_not_use", + ): Promise { + await this.apiRequest( + `/api/environments/${this.getTeamId()}/mcp_server_installations/${installationId}/tools/${encodeURIComponent(toolName)}/`, + { + method: "PATCH", + body: JSON.stringify({ approval_state: approvalState }), + }, + ); + } + async appendTaskRunLog( taskId: string, runId: string, diff --git a/packages/agent/src/server/agent-server.ts b/packages/agent/src/server/agent-server.ts index 22bf913294..963a5db86d 100644 --- a/packages/agent/src/server/agent-server.ts +++ b/packages/agent/src/server/agent-server.ts @@ -40,6 +40,7 @@ import { import type { PermissionMode } from "../execution-mode"; import { DEFAULT_CODEX_MODEL } from "../gateway-models"; import { HandoffCheckpointTracker } from "../handoff-checkpoint"; +import { resolveMcpStoreToolKey } from "../mcp-store/tool-keys"; import { PostHogAPIClient } from "../posthog-api"; import { findPrUrl, wasCreatedRecently } from "../pr-url-detector"; import { @@ -364,6 +365,131 @@ export class AgentServer { return mode === "default" || mode === "auto" || mode === "read-only"; } + private resolveMcpToolName(candidate: string | null): string | null { + return resolveMcpStoreToolKey(candidate, { + approvals: this.config.mcpToolApprovals, + installations: this.config.mcpToolInstallations, + }); + } + + private getMcpToolNameFromPermissionRequest( + params: RequestPermissionRequest, + ): string | null { + const rawInput = + params.toolCall?.rawInput && + typeof params.toolCall.rawInput === "object" && + !Array.isArray(params.toolCall.rawInput) + ? (params.toolCall.rawInput as Record) + : null; + const rawInputToolName = + typeof rawInput?.toolName === "string" ? rawInput.toolName : null; + const toolNameFromRawInputToolName = + this.resolveMcpToolName(rawInputToolName); + if (toolNameFromRawInputToolName) { + return toolNameFromRawInputToolName; + } + + const rawInputName = + typeof rawInput?.name === "string" ? rawInput.name : null; + const toolNameFromRawInputName = this.resolveMcpToolName(rawInputName); + if (toolNameFromRawInputName) { + return toolNameFromRawInputName; + } + + const meta = + params.toolCall?._meta && + typeof params.toolCall._meta === "object" && + !Array.isArray(params.toolCall._meta) + ? (params.toolCall._meta as Record) + : null; + const claudeCode = + meta?.claudeCode && + typeof meta.claudeCode === "object" && + !Array.isArray(meta.claudeCode) + ? (meta.claudeCode as Record) + : null; + const topLevelMetaToolName = + typeof meta?.toolName === "string" ? meta.toolName : null; + const toolNameFromTopLevelMeta = + this.resolveMcpToolName(topLevelMetaToolName); + if (toolNameFromTopLevelMeta) { + return toolNameFromTopLevelMeta; + } + + const metaToolName = + typeof claudeCode?.toolName === "string" ? claudeCode.toolName : null; + const toolNameFromMeta = this.resolveMcpToolName(metaToolName); + if (toolNameFromMeta) { + return toolNameFromMeta; + } + + const title = + typeof params.toolCall?.title === "string" ? params.toolCall.title : null; + const toolNameFromTitle = this.resolveMcpToolName(title); + if (toolNameFromTitle) { + return toolNameFromTitle; + } + + return null; + } + + private isMcpStoreApprovalRequest(params: RequestPermissionRequest): boolean { + const toolName = this.getMcpToolNameFromPermissionRequest(params); + return Boolean( + toolName && this.config.mcpToolApprovals?.[toolName] === "needs_approval", + ); + } + + private async approveMcpStoreToolIfNeeded( + params: RequestPermissionRequest, + response: RequestPermissionResponse, + ): Promise { + const approved = + response.outcome?.outcome === "selected" && + (response.outcome.optionId === "allow" || + response.outcome.optionId === "allow_always"); + if (!approved) { + return; + } + + const toolName = this.getMcpToolNameFromPermissionRequest(params); + if ( + !toolName || + this.config.mcpToolApprovals?.[toolName] !== "needs_approval" + ) { + return; + } + + const installation = this.config.mcpToolInstallations?.[toolName]; + if (!installation) { + this.logger.warn( + "Cannot persist MCP Store tool approval: missing installation ref", + { + toolName, + }, + ); + return; + } + + try { + await this.posthogAPI.updateMcpToolApproval( + installation.installationId, + installation.toolName, + "approved", + ); + 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), + }); + } + } + private createApp(): Hono { const app = new Hono(); @@ -1183,6 +1309,9 @@ export class AgentServer { allowedDomains: this.config.allowedDomains, jsonSchema: preTask?.json_schema ?? null, permissionMode: initialPermissionMode, + ...(this.config.mcpToolApprovals && { + mcpToolApprovals: this.config.mcpToolApprovals, + }), ...(this.config.baseBranch && { baseBranch: this.config.baseBranch }), ...this.buildClaudeCodeSessionMeta(runtimeAdapter), }; @@ -2458,21 +2587,26 @@ ${signedCommitInstructions} { const isQuestion = codeToolKind === "question"; const sessionPermissionMode = this.getSessionPermissionMode(); + const needsMcpStoreApproval = this.isMcpStoreApprovalRequest(params); const needsDesktopApproval = isQuestion || this.shouldRelayPermissionToClient(sessionPermissionMode); if ( isPlanApproval || + needsMcpStoreApproval || (needsDesktopApproval && this.session?.hasDesktopConnected) ) { this.logger.debug("Relaying permission request", { kind: params.toolCall?.kind, isQuestion, + needsMcpStoreApproval, hasDesktopConnected: this.session?.hasDesktopConnected ?? false, sessionPermissionMode, }); - return this.relayPermissionToClient(params); + const response = await this.relayPermissionToClient(params); + await this.approveMcpStoreToolIfNeeded(params, response); + return response; } } diff --git a/packages/agent/src/server/bin.ts b/packages/agent/src/server/bin.ts index c70368ae9a..9b14e1013f 100644 --- a/packages/agent/src/server/bin.ts +++ b/packages/agent/src/server/bin.ts @@ -3,7 +3,12 @@ import { Command } from "commander"; import { z } from "zod/v4"; import { isSupportedReasoningEffort } from "../adapters/reasoning-effort"; import { AgentServer } from "./agent-server"; -import { claudeCodeConfigSchema, mcpServersSchema } from "./schemas"; +import { + claudeCodeConfigSchema, + mcpServersSchema, + mcpToolApprovalsSchema, + mcpToolInstallationsSchema, +} from "./schemas"; const envSchema = z.object({ JWT_PUBLIC_KEY: z @@ -95,6 +100,14 @@ program "--mcpServers ", "MCP servers config as JSON array (ACP McpServer[] format)", ) + .option( + "--mcpToolApprovals ", + "MCP tool approval states as JSON object keyed by mcp__server__tool", + ) + .option( + "--mcpToolInstallations ", + "MCP tool installation refs as JSON object keyed by mcp__server__tool", + ) .option("--createPr ", "Whether this run may publish changes") .option("--baseBranch ", "Base branch for PR creation") .option( @@ -126,6 +139,16 @@ program mcpServersSchema, "--mcpServers", ); + const mcpToolApprovals = parseJsonOption( + options.mcpToolApprovals, + mcpToolApprovalsSchema, + "--mcpToolApprovals", + ); + const mcpToolInstallations = parseJsonOption( + options.mcpToolInstallations, + mcpToolInstallationsSchema, + "--mcpToolInstallations", + ); const claudeCode = parseJsonOption( options.claudeCodeConfig, claudeCodeConfigSchema, @@ -169,6 +192,8 @@ program runId: options.runId, createPr, mcpServers, + mcpToolApprovals, + mcpToolInstallations, baseBranch: options.baseBranch, claudeCode, allowedDomains, diff --git a/packages/agent/src/server/question-relay.test.ts b/packages/agent/src/server/question-relay.test.ts index e00643d0a3..147e454ae5 100644 --- a/packages/agent/src/server/question-relay.test.ts +++ b/packages/agent/src/server/question-relay.test.ts @@ -17,7 +17,7 @@ interface TestableAgentServer { options: unknown[]; toolCall: unknown; }) => Promise<{ - outcome: { outcome: string }; + outcome: { outcome: string; optionId?: string }; _meta?: { message?: string }; }>; }; @@ -25,6 +25,17 @@ interface TestableAgentServer { session: unknown; relayAgentResponse: (payload: Record) => Promise; sendInitialTaskMessage: (payload: Record) => Promise; + resolvePermission: (requestId: string, optionId: string) => boolean; + config: { + mcpToolApprovals?: Record< + string, + "approved" | "needs_approval" | "do_not_use" + >; + mcpToolInstallations?: Record< + string, + { installationId: string; toolName: string } + >; + }; } const TEST_PAYLOAD = { @@ -356,6 +367,180 @@ describe("Question relay", () => { expect(brokenSseController.send).toHaveBeenCalledTimes(1); expect(appendRawLine).toHaveBeenCalledTimes(2); }); + + it("relays MCP Store approval requests without a desktop connection and persists approval", async () => { + const appendRawLine = vi.fn(); + const approvalSpy = vi + .spyOn(server.posthogAPI, "updateMcpToolApproval") + .mockResolvedValue(undefined); + + server.config.mcpToolApprovals = { + mcp__Linear__search: "needs_approval", + }; + server.config.mcpToolInstallations = { + mcp__Linear__search: { + installationId: "inst-1", + toolName: "search", + }, + }; + 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: "The agent wants to call search (Linear)", + kind: "other", + rawInput: { toolName: "mcp__Linear__search" }, + }, + }); + + 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", + "search", + "approved", + ); + expect(server.config.mcpToolApprovals.mcp__Linear__search).toBe( + "approved", + ); + }); + + it("recognizes MCP Store approval requests by bare upstream tool name", 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: "bypassPermissions", + 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: { query: "today" }, + }, + }); + + 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("approved"); + }); + + it("recognizes MCP Store approval requests with unsanitized MCP server names", async () => { + const appendRawLine = vi.fn(); + const approvalSpy = vi + .spyOn(server.posthogAPI, "updateMcpToolApproval") + .mockResolvedValue(undefined); + + server.config.mcpToolApprovals = { + mcp__Granola_Meetings__query_granola_meetings: "needs_approval", + }; + server.config.mcpToolInstallations = { + mcp__Granola_Meetings__query_granola_meetings: { + installationId: "inst-1", + toolName: "query_granola_meetings", + }, + }; + server.session = { + payload: TEST_PAYLOAD, + sseController: null, + hasDesktopConnected: false, + permissionMode: "bypassPermissions", + logWriter: { appendRawLine }, + }; + + const client = server.createCloudClient(TEST_PAYLOAD); + const permissionPromise = client.requestPermission({ + options: ALLOW_OPTIONS, + toolCall: { + toolCallId: "tool-1", + title: + "The agent wants to call query_granola_meetings (Granola Meetings)", + kind: "other", + rawInput: { + toolName: "mcp__Granola Meetings__query_granola_meetings", + }, + }, + }); + + const request = appendRawLine.mock.calls + .map(([, line]) => JSON.parse(line)) + .find((n) => n?.method === "_posthog/permission_request"); + expect(request).toBeTruthy(); + + 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_Meetings__query_granola_meetings, + ).toBe("approved"); + }); }); describe("with createPr disabled", () => { diff --git a/packages/agent/src/server/schemas.test.ts b/packages/agent/src/server/schemas.test.ts index 898f1a270e..b7f228e025 100644 --- a/packages/agent/src/server/schemas.test.ts +++ b/packages/agent/src/server/schemas.test.ts @@ -1,5 +1,10 @@ import { describe, expect, it } from "vitest"; -import { mcpServersSchema, validateCommandParams } from "./schemas"; +import { + mcpServersSchema, + mcpToolApprovalsSchema, + mcpToolInstallationsSchema, + validateCommandParams, +} from "./schemas"; describe("mcpServersSchema", () => { it("accepts a valid HTTP server", () => { @@ -116,6 +121,40 @@ describe("mcpServersSchema", () => { }); }); +describe("MCP tool approval config schemas", () => { + it("accepts approval state and installation maps", () => { + expect( + mcpToolApprovalsSchema.safeParse({ + mcp__Linear__search: "needs_approval", + mcp__Linear__create_ticket: "approved", + }).success, + ).toBe(true); + + expect( + mcpToolInstallationsSchema.safeParse({ + mcp__Linear__search: { + installationId: "inst-1", + toolName: "search", + }, + }).success, + ).toBe(true); + }); + + it("rejects invalid approval state and incomplete installation refs", () => { + expect( + mcpToolApprovalsSchema.safeParse({ + mcp__Linear__search: "ask_me_later", + }).success, + ).toBe(false); + + expect( + mcpToolInstallationsSchema.safeParse({ + mcp__Linear__search: { installationId: "inst-1" }, + }).success, + ).toBe(false); + }); +}); + describe("validateCommandParams", () => { it("accepts structured user_message content arrays", () => { const result = validateCommandParams("user_message", { diff --git a/packages/agent/src/server/schemas.ts b/packages/agent/src/server/schemas.ts index 2dfa791a9f..c1bd1b7b4c 100644 --- a/packages/agent/src/server/schemas.ts +++ b/packages/agent/src/server/schemas.ts @@ -29,6 +29,30 @@ export const mcpServersSchema = z.array(remoteMcpServerSchema); export type RemoteMcpServer = z.infer; +const mcpToolApprovalStateSchema = z.enum([ + "approved", + "needs_approval", + "do_not_use", +]); + +export const mcpToolApprovalsSchema = z.record( + z.string(), + mcpToolApprovalStateSchema, +); + +export const mcpToolInstallationsSchema = z.record( + z.string(), + z.object({ + installationId: z.string().min(1), + toolName: z.string().min(1), + }), +); + +export type McpToolApprovalsConfig = z.infer; +export type McpToolInstallationsConfig = z.infer< + typeof mcpToolInstallationsSchema +>; + export const claudeCodeConfigSchema = z.object({ systemPrompt: z .union([ diff --git a/packages/agent/src/server/types.ts b/packages/agent/src/server/types.ts index d11cb7748c..8a6c376486 100644 --- a/packages/agent/src/server/types.ts +++ b/packages/agent/src/server/types.ts @@ -1,5 +1,9 @@ import type { AgentMode } from "../types"; -import type { RemoteMcpServer } from "./schemas"; +import type { + McpToolApprovalsConfig, + McpToolInstallationsConfig, + RemoteMcpServer, +} from "./schemas"; export interface ClaudeCodeConfig { systemPrompt?: @@ -23,6 +27,8 @@ export interface AgentServerConfig { createPr?: boolean; version?: string; mcpServers?: RemoteMcpServer[]; + mcpToolApprovals?: McpToolApprovalsConfig; + mcpToolInstallations?: McpToolInstallationsConfig; baseBranch?: string; claudeCode?: ClaudeCodeConfig; allowedDomains?: string[];