From 9318048f514e4a2f6bf61470b3a003f52a1c2790 Mon Sep 17 00:00:00 2001 From: Mohamed Boudra Date: Wed, 24 Dec 2025 09:05:24 +0000 Subject: [PATCH] ensure mcp waits surface permissions --- .../server/src/server/agent/agent-manager.ts | 17 +++- .../src/server/agent/agent-mcp.e2e.test.ts | 79 +++++++++++++------ .../server/src/server/agent/mcp-server.ts | 12 ++- 3 files changed, 76 insertions(+), 32 deletions(-) diff --git a/packages/server/src/server/agent/agent-manager.ts b/packages/server/src/server/agent/agent-manager.ts index bc8275d7c..be59baed1 100644 --- a/packages/server/src/server/agent/agent-manager.ts +++ b/packages/server/src/server/agent/agent-manager.ts @@ -60,6 +60,7 @@ export type AgentManagerOptions = { export type WaitForAgentOptions = { signal?: AbortSignal; + waitForActive?: boolean; }; export type WaitForAgentResult = { @@ -611,9 +612,12 @@ export class AgentManager { }; } - if (!isAgentBusy(snapshot.lifecycle)) { + const initialStatus = snapshot.lifecycle; + const initialBusy = isAgentBusy(initialStatus); + const waitForActive = options?.waitForActive ?? false; + if (!waitForActive && !initialBusy) { return { - status: snapshot.lifecycle, + status: initialStatus, permission: null, lastMessage: this.getLastAssistantMessage(agentId) }; @@ -632,7 +636,8 @@ export class AgentManager { return; } - let currentStatus: AgentLifecycleStatus = snapshot.lifecycle; + let currentStatus: AgentLifecycleStatus = initialStatus; + let hasStarted = initialBusy; // Bug #3 Fix: Declare unsubscribe and abortHandler upfront so cleanup can reference them let unsubscribe: (() => void) | null = null; @@ -691,7 +696,11 @@ export class AgentManager { finish(pending); return; } - if (!isAgentBusy(event.agent.lifecycle)) { + if (isAgentBusy(event.agent.lifecycle)) { + hasStarted = true; + return; + } + if (!waitForActive || hasStarted) { finish(null); } } diff --git a/packages/server/src/server/agent/agent-mcp.e2e.test.ts b/packages/server/src/server/agent/agent-mcp.e2e.test.ts index d12d7e44c..440e492c4 100644 --- a/packages/server/src/server/agent/agent-mcp.e2e.test.ts +++ b/packages/server/src/server/agent/agent-mcp.e2e.test.ts @@ -160,6 +160,12 @@ describe("agent MCP end-to-end", () => { let agentId: string | null = null; try { + const initialPrompt = [ + "You must call the tool named shell.", + "Run this command exactly: [\"bash\", \"-lc\", \"echo ok > mcp-smoke.txt\"].", + "After the tool runs, reply with done and stop.", + ].join("\n"); + const result = (await client.callTool({ name: "create_agent", args: { @@ -167,6 +173,7 @@ describe("agent MCP end-to-end", () => { title: "MCP e2e smoke", agentType: "claude", initialMode: "default", + initialPrompt, background: false, }, })) as McpToolResult; @@ -175,35 +182,13 @@ describe("agent MCP end-to-end", () => { expect(payload).toBeTruthy(); agentId = payload?.agentId as string | null; expect(agentId).toBeTruthy(); - - const prompt = [ - "You must call the tool named shell.", - "Run this command exactly: [\"bash\", \"-lc\", \"echo ok > mcp-smoke.txt\"].", - "After the tool runs, reply with done and stop.", - ].join("\n"); - - await client.callTool({ - name: "send_agent_prompt", - args: { - agentId, - prompt, - sessionMode: "default", - background: true, - }, - }); - - const waitResult = (await client.callTool({ - name: "wait_for_agent", - args: { agentId }, - })) as McpToolResult; - const waitPayload = getStructuredContent(waitResult); - const permission = waitPayload?.permission as PermissionPayload | null; - expect(permission?.id).toBeTruthy(); + const createPermission = payload?.permission as PermissionPayload | null; + expect(createPermission?.id).toBeTruthy(); await client.callTool({ name: "respond_to_permission", args: { agentId, - requestId: permission!.id, + requestId: createPermission!.id, response: { behavior: "allow" }, }, }); @@ -226,6 +211,50 @@ describe("agent MCP end-to-end", () => { throw new Error(`${(error as Error).message}\n${details}`); } expect(contents.trim()).toBe("ok"); + + const prompt = [ + "You must call the tool named shell.", + "Run this command exactly: [\"bash\", \"-lc\", \"echo ok-2 > mcp-smoke-2.txt\"].", + "After the tool runs, reply with done and stop.", + ].join("\n"); + + const promptResult = (await client.callTool({ + name: "send_agent_prompt", + args: { + agentId, + prompt, + sessionMode: "default", + background: false, + }, + })) as McpToolResult; + + const promptPayload = getStructuredContent(promptResult); + const promptPermission = promptPayload?.permission as PermissionPayload | null; + expect(promptPermission?.id).toBeTruthy(); + + const waitPermissionResult = (await client.callTool({ + name: "wait_for_agent", + args: { agentId }, + })) as McpToolResult; + const waitPermissionPayload = getStructuredContent(waitPermissionResult); + const waitPermission = + waitPermissionPayload?.permission as PermissionPayload | null; + expect(waitPermission?.id).toBe(promptPermission?.id); + + await client.callTool({ + name: "respond_to_permission", + args: { + agentId, + requestId: promptPermission!.id, + response: { behavior: "allow" }, + }, + }); + + await waitForAgentCompletion(client, agentId); + + const secondFilePath = path.join(agentCwd, "mcp-smoke-2.txt"); + const secondContents = await waitForFile(secondFilePath); + expect(secondContents.trim()).toBe("ok-2"); } finally { if (agentId) { await client.callTool({ name: "kill_agent", args: { agentId } }); diff --git a/packages/server/src/server/agent/mcp-server.ts b/packages/server/src/server/agent/mcp-server.ts index 1c8c2e4bb..716b27520 100644 --- a/packages/server/src/server/agent/mcp-server.ts +++ b/packages/server/src/server/agent/mcp-server.ts @@ -62,7 +62,8 @@ function expandPath(path: string): string { async function waitForAgentWithTimeout( agentManager: AgentManager, agentId: string, - existingSignal?: AbortSignal + existingSignal?: AbortSignal, + options?: { waitForActive?: boolean } ): Promise { const timeoutSignal = AbortSignal.timeout(AGENT_WAIT_TIMEOUT_MS); const abortController = new AbortController(); @@ -99,6 +100,7 @@ async function waitForAgentWithTimeout( try { const result = await agentManager.waitForAgentEvent(agentId, { signal: abortController.signal, + waitForActive: options?.waitForActive, }); return result; } catch (error) { @@ -297,7 +299,9 @@ export async function createAgentMcpServer( // If not running in background, wait for completion if (!background) { - const result = await waitForAgentWithTimeout(agentManager, snapshot.id); + const result = await waitForAgentWithTimeout(agentManager, snapshot.id, undefined, { + waitForActive: true, + }); const responseData = { agentId: snapshot.id, @@ -509,7 +513,9 @@ export async function createAgentMcpServer( // If not running in background, wait for completion if (!background) { - const result = await waitForAgentWithTimeout(agentManager, agentId); + const result = await waitForAgentWithTimeout(agentManager, agentId, undefined, { + waitForActive: true, + }); const responseData = { success: true,