diff --git a/src/agent/tools.ts b/src/agent/tools.ts index 22a98521d..9b1b10097 100644 --- a/src/agent/tools.ts +++ b/src/agent/tools.ts @@ -15,6 +15,7 @@ import { import { manageTasksDefinition } from "./tasks.js"; import { validateView } from "../tui/view/validate.js"; import { SETTINGS_DIR_NAME } from "../branding.js"; +import { startBoundMcpHandshake } from "../exec/mcp-handshake.js"; import { advertiseShellGuardTimeout, type ShellTimeoutConfig, @@ -312,6 +313,10 @@ export interface MCPConnectCallbacks { // Headless hosts must not advertise an auth callback they cannot complete; // its presence marks the OAuth flow as interactive. interactiveAuth: boolean; + // The TUI sets this so a dial abort (the 15s handshake cap) is dropped once + // the browser auth URL is showing. The browser-auth budget and a real + // cancel still bound that wait. Headless leaves it unset. + browserAuthOutlivesDialAbort?: boolean; // Fired whenever a server's connection state changes. onStatus: (state: MCPServerState) => void; // Fired after a server connects and its tools register, with the full @@ -1269,6 +1274,7 @@ export async function createAgentToolset( return; } } + startBoundMcpHandshake(signal); if (staleOrDisabled()) return; callbacks.onStatus({ name: config.name, state: "connecting" }); let result: MCPConnectResult; @@ -1279,6 +1285,12 @@ export async function createAgentToolset( ...(callbacks.interactiveAuth ? { onAuthURL: (name: string, url: string) => { + // The parent signal is the dial cap plus exit. Once the + // operator has an auth URL, that cap must not close the + // redirect or burn the single browser attempt. Dispose and + // per-server cancel stay on connectionSignal. + if (callbacks.browserAuthOutlivesDialAbort === true) + forwarded?.disarm(); if (!disposed && !staleOrDisabled()) { callbacks.onStatus({ name, state: "needs-auth", url }); } diff --git a/src/exec/mcp-handshake.ts b/src/exec/mcp-handshake.ts index 7f561de43..1dd730c40 100644 --- a/src/exec/mcp-handshake.ts +++ b/src/exec/mcp-handshake.ts @@ -25,26 +25,56 @@ export async function awaitExecMcpConnect( } } +const deferredHandshakeStarts = new WeakMap void>(); + +export function bindDeferredMcpHandshake( + signal: AbortSignal, + start: () => void, +): void { + deferredHandshakeStarts.set(signal, start); +} + +// First post-trust dial starts a deferred cap. Later servers in the same +// batch share that window. Unbound signals (headless) are already armed. +export function startBoundMcpHandshake(signal: AbortSignal | undefined): void { + if (signal === undefined) return; + const start = deferredHandshakeStarts.get(signal); + if (start === undefined) return; + deferredHandshakeStarts.delete(signal); + start(); +} + // Connect-only abort. The MCP client ties `signal` to the transport lifecycle, // so AbortSignal.timeout would kill a handshake that already succeeded. The // toolset forwards this signal per server and detaches on settle, so abort // only reaches handshakes still in flight. Disarm only when the batch // fulfills — a rejected Promise.all still leaves sibling forwards armed. -export function armExecMcpHandshakeAbort(timeoutMs: number): { +// `when` "later" does not start the timer until `start` (TUI trust gate). +export function armExecMcpHandshakeAbort( + timeoutMs: number, + when: "now" | "later" = "now", +): { signal: AbortSignal; disarm: () => void; + start: () => void; } { const controller = new AbortController(); - const timer = setTimeout(() => { - controller.abort(); - }, timeoutMs); + let timer: ReturnType | undefined; let disarmed = false; + const start = (): void => { + if (disarmed || timer !== undefined) return; + timer = setTimeout(() => { + controller.abort(); + }, timeoutMs); + }; + if (when === "now") start(); return { signal: controller.signal, + start, disarm: () => { if (disarmed) return; disarmed = true; - clearTimeout(timer); + if (timer !== undefined) clearTimeout(timer); }, }; } diff --git a/src/tui/runner/mcp-connect.test.ts b/src/tui/runner/mcp-connect.test.ts new file mode 100644 index 000000000..b5af79b2f --- /dev/null +++ b/src/tui/runner/mcp-connect.test.ts @@ -0,0 +1,192 @@ +import { afterEach, beforeEach, describe, expect, test } from "bun:test"; +import { mkdtempSync, rmSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { + installMcpConnectMock, + mcpTestPermissionGate, +} from "../../../testkit/mcp-connect-mock.js"; +import type { MCPServerState } from "../../agent/tools.js"; +import { runTuiMcpDial } from "./mcp-connect.js"; + +const dirs: string[] = []; + +function tempCwd(): string { + const dir = mkdtempSync(join(tmpdir(), "corbits-tui-mcp-dial-")); + dirs.push(dir); + return dir; +} + +const mock = await installMcpConnectMock( + import.meta.resolve("../../mcp/client.js"), + { + initialTools: [{ name: "list", description: "List", inputSchema: {} }], + abortableDeferred: true, + teardownOnAbort: true, + }, +); + +const { createAgentToolset } = await import("../../agent/tools.js"); + +const acme = { + name: "acme", + type: "http" as const, + url: "https://mcp.acme.test/mcp", +}; +const lin = { + name: "lin", + type: "http" as const, + url: "https://mcp.lin.test/mcp", +}; +const stuck = { + name: "stuck", + type: "stdio" as const, + command: "stuck-mcp", +}; + +function callbacks( + states: MCPServerState[], + browserAuthOutlivesDialAbort = false, +) { + return { + interactiveAuth: true, + ...(browserAuthOutlivesDialAbort + ? { browserAuthOutlivesDialAbort: true as const } + : {}), + onStatus: (state: MCPServerState) => { + states.push(state); + }, + onToolsChanged: () => undefined, + }; +} + +afterEach(() => { + for (const dir of dirs.splice(0)) { + rmSync(dir, { recursive: true, force: true }); + } +}); + +beforeEach(() => { + mock.reset(); +}); + +describe("TUI MCP dial timeout", () => { + test("a hung dial fails, the live sibling stays connected, and resume runs", async () => { + mock.hangNames.add("lin"); + const toolset = await createAgentToolset({ + cwd: tempCwd(), + permissionGate: mcpTestPermissionGate(), + onOperatorGate: async () => ({ kind: "cancel" }), + mcpServers: [acme, lin], + }); + const states: MCPServerState[] = []; + const exit = new AbortController(); + let resumed = false; + try { + const started = Date.now(); + await runTuiMcpDial( + (signal) => toolset.connectMCP(callbacks(states), signal), + exit.signal, + 40, + ).then(async () => { + resumed = true; + }); + expect(Date.now() - started).toBeLessThan(500); + expect(resumed).toBe(true); + expect(states.some((s) => s.name === "lin" && s.state === "failed")).toBe( + true, + ); + expect( + states.some((s) => s.name === "acme" && s.state === "connected"), + ).toBe(true); + expect(toolset.hasMCPServer("acme")).toBe(true); + expect(toolset.hasMCPServer("lin")).toBe(false); + expect(mock.closedClients).not.toContain("acme"); + expect( + toolset.dynamicRunner.currentDefinitions().map((d) => d.name), + ).toContain("mcp__acme__list"); + } finally { + await toolset.dispose(); + } + }); + + test("unfinished browser auth survives the handshake cap and the live sibling stays mounted", async () => { + mock.browserAuthNames.add("lin"); + mock.hangNames.add("stuck"); + const toolset = await createAgentToolset({ + cwd: tempCwd(), + permissionGate: mcpTestPermissionGate(), + onOperatorGate: async () => ({ kind: "cancel" }), + mcpServers: [acme, lin, stuck], + }); + const states: MCPServerState[] = []; + const exit = new AbortController(); + let resumed = false; + const dial = runTuiMcpDial( + (signal) => + toolset.connectMCP(callbacks(states, true), signal).then(async () => { + resumed = true; + }), + exit.signal, + 40, + ); + try { + await new Promise((resolve) => setTimeout(resolve, 80)); + expect(resumed).toBe(false); + expect(states.some((s) => s.name === "lin" && s.state === "failed")).toBe( + false, + ); + expect( + states.some((s) => s.name === "lin" && s.state === "needs-auth"), + ).toBe(true); + expect( + states.some((s) => s.name === "stuck" && s.state === "failed"), + ).toBe(true); + expect( + states.some((s) => s.name === "acme" && s.state === "connected"), + ).toBe(true); + expect(toolset.hasMCPServer("acme")).toBe(true); + expect(toolset.hasMCPServer("stuck")).toBe(false); + expect(mock.closedClients).not.toContain("acme"); + expect( + toolset.dynamicRunner.currentDefinitions().map((d) => d.name), + ).toContain("mcp__acme__list"); + expect(mock.authWaitAborts).toBe(0); + } finally { + await toolset.dispose(); + await dial; + } + }); + + test("a trust prompt longer than the dial cap still connects on a live signal", async () => { + const toolset = await createAgentToolset({ + cwd: tempCwd(), + permissionGate: mcpTestPermissionGate(), + onOperatorGate: async () => ({ kind: "cancel" }), + mcpServers: [acme], + mcpServersSource: "local", + projectTrust: { + trustedPluginPaths: [], + trustedMcpFingerprints: [], + trustedGrantFingerprints: [], + }, + requestMcpTrust: async () => { + await new Promise((resolve) => setTimeout(resolve, 80)); + return true; + }, + }); + const exit = new AbortController(); + try { + await runTuiMcpDial( + (signal) => toolset.connectMCP(callbacks([]), signal), + exit.signal, + 40, + ); + expect(mock.connectOptions).toHaveLength(1); + expect(mock.connectOptions[0]?.signal?.aborted).toBe(false); + expect(toolset.hasMCPServer("acme")).toBe(true); + } finally { + await toolset.dispose(); + } + }); +}); diff --git a/src/tui/runner/mcp-connect.ts b/src/tui/runner/mcp-connect.ts new file mode 100644 index 000000000..1c9c8f906 --- /dev/null +++ b/src/tui/runner/mcp-connect.ts @@ -0,0 +1,26 @@ +import { + armExecMcpHandshakeAbort, + bindDeferredMcpHandshake, + EXEC_MCP_HANDSHAKE_TIMEOUT_MS, + followExecMcpHandshake, +} from "../../exec/mcp-handshake.js"; + +/** + * Bound one TUI MCP dial (or the startup batch) the same way headless does. + * The timer starts on the first dial after trust filtering, then aborts + * handshakes still in flight. A server that reaches interactive browser auth + * drops this abort and keeps the browser-auth budget; a dial that never gets + * that far still fails closed. The toolset detaches each server's forward on + * settle, and this disarms the timer when the dial fulfills, so a live sibling + * is not torn down. + */ +export function runTuiMcpDial( + dial: (signal: AbortSignal) => Promise, + exitSignal: AbortSignal, + timeoutMs: number = EXEC_MCP_HANDSHAKE_TIMEOUT_MS, +): Promise { + const handshake = armExecMcpHandshakeAbort(timeoutMs, "later"); + const signal = AbortSignal.any([exitSignal, handshake.signal]); + bindDeferredMcpHandshake(signal, handshake.start); + return followExecMcpHandshake(dial(signal), handshake); +} diff --git a/src/tui/runner/mcp.ts b/src/tui/runner/mcp.ts index 326433b52..06370d508 100644 --- a/src/tui/runner/mcp.ts +++ b/src/tui/runner/mcp.ts @@ -29,6 +29,7 @@ import { import { mergeMcpSurfaceEntries, isBuiltinRow } from "../mcp-list.js"; import { nextMcpCatalog } from "../mcp-catalog.js"; import type { MCPConnectCallbacks } from "../../agent/tools.js"; +import { runTuiMcpDial } from "./mcp-connect.js"; import { type RunnerServices, type RunnerState } from "./state.js"; import { SETTINGS_DIR_NAME, LOG_NAMESPACE_ROOT } from "../../branding.js"; @@ -46,6 +47,9 @@ export function wireMcp( ): McpWiring { const mcpConnectCallbacks: MCPConnectCallbacks = { interactiveAuth: true, + // Browser auth has its own wait. The startup dial cap must not fail that + // row or close the redirect once the auth URL is up. + browserAuthOutlivesDialAbort: true, onStatus: (status) => { services.mcpStates.set(status.name, status); services.emitter.emit("mcp.status", status); @@ -66,19 +70,17 @@ export function wireMcp( }; const connectLateMCPServer = (server: MCPServerConfig): void => { - void services.toolset - .connectMCPServer( - server, - mcpConnectCallbacks, - services.mcpConnectController.signal, - ) - .catch((err: unknown) => { - if (err instanceof Error && err.name === "AbortError") return; - tuiLogger.warn("Late MCP connect failed for {server}: {error}", { - server: server.name, - error: err instanceof Error ? err.message : String(err), - }); + void runTuiMcpDial( + (signal) => + services.toolset.connectMCPServer(server, mcpConnectCallbacks, signal), + services.mcpConnectController.signal, + ).catch((err: unknown) => { + if (err instanceof Error && err.name === "AbortError") return; + tuiLogger.warn("Late MCP connect failed for {server}: {error}", { + server: server.name, + error: err instanceof Error ? err.message : String(err), }); + }); }; state.connectLateMCPServer = connectLateMCPServer; @@ -239,18 +241,16 @@ function createMcpSurface( }; } // Manual retry: cancel any backoff loop first so exactly one dial runs. - void services.toolset - .retryMCPServer( - server, - mcpConnectCallbacks, - services.mcpConnectController.signal, - ) - .catch((err: unknown) => { - tuiLogger.warn("MCP retry failed for {server}: {error}", { - server: server.name, - error: err instanceof Error ? err.message : String(err), - }); + void runTuiMcpDial( + (signal) => + services.toolset.retryMCPServer(server, mcpConnectCallbacks, signal), + services.mcpConnectController.signal, + ).catch((err: unknown) => { + tuiLogger.warn("MCP retry failed for {server}: {error}", { + server: server.name, + error: err instanceof Error ? err.message : String(err), }); + }); return { ok: true, message: `Retrying ${server.name}; connecting now.` }; }, setEnabled: async (name: string, enabled: boolean) => { diff --git a/src/tui/runner/wiring.ts b/src/tui/runner/wiring.ts index e818a6335..440f84088 100644 --- a/src/tui/runner/wiring.ts +++ b/src/tui/runner/wiring.ts @@ -72,6 +72,7 @@ import { createRuntimeShutdown } from "./shutdown.js"; import { ASK_DEADLINE_MS } from "../../subagent/session-store.js"; import { resumeTranscriptLoadErrorBlock } from "./exit.js"; import { userInboundMessage } from "./submit.js"; +import { runTuiMcpDial } from "./mcp-connect.js"; import { hostOf, liveAgent, @@ -578,10 +579,14 @@ export function wirePostStartup( // Connect MCP after the TUI is up so auth surfaces as a copyable link, not // a browser pop; each server that mounts tools requests an agent rebuild - // through its callbacks. On settle, resume any persisted workflow. Aborted - // on exit. - void services.toolset - .connectMCP(mcpConnectCallbacks, services.mcpConnectController.signal) + // through its callbacks. A hung dial that never reaches browser auth aborts + // on the same cap as headless and reports failed. A dial already waiting on + // the browser keeps that wait, so resume is not stuck behind a sibling that + // never settles and auth is not failed by the cap. + void runTuiMcpDial( + (signal) => services.toolset.connectMCP(mcpConnectCallbacks, signal), + services.mcpConnectController.signal, + ) .then(async () => { // Capability map reflects connected servers; restore any persisted workflow. await services.workflowHost.resume(); diff --git a/testkit/mcp-connect-mock.ts b/testkit/mcp-connect-mock.ts index 46ba62080..36416cdd8 100644 --- a/testkit/mcp-connect-mock.ts +++ b/testkit/mcp-connect-mock.ts @@ -47,6 +47,11 @@ export interface McpConnectMock { authResourceCloses: number; /** Server names whose connect never settles until the signal aborts. */ hangNames: Set; + /** + * Server names that offer an auth URL and then wait until the connect + * signal aborts. Other servers in the same batch are unaffected. + */ + browserAuthNames: Set; connectedTools: MCPTool[]; toolCallResult: string; calls: McpToolCall[]; @@ -127,6 +132,7 @@ export async function installMcpConnectMock( authWaitAborts: 0, authResourceCloses: 0, hangNames: new Set(), + browserAuthNames: new Set(), connectedTools: initialTools, toolCallResult, calls: [], @@ -146,6 +152,7 @@ export async function installMcpConnectMock( mock.authWaitAborts = 0; mock.authResourceCloses = 0; mock.hangNames.clear(); + mock.browserAuthNames.clear(); mock.connectedTools = initialTools; mock.toolCallResult = toolCallResult; mock.calls.length = 0; @@ -165,6 +172,29 @@ export async function installMcpConnectMock( const generation = ++mock.connectGeneration; const mode = mock.mode; + if (mock.browserAuthNames.has(config.name)) { + connectOptions.onAuthURL?.( + config.name, + mock.authURL ?? "https://auth.test/authorize", + ); + await new Promise((resolve) => { + const onAbort = (): void => { + mock.authWaitAborts += 1; + mock.authResourceCloses += 1; + resolve(); + }; + if (connectOptions.signal?.aborted === true) onAbort(); + else + connectOptions.signal?.addEventListener("abort", onAbort, { + once: true, + }); + }); + return { + ok: false as const, + serverName: config.name, + error: "authorization aborted", + }; + } if (mode === "auth" || mock.blockOnAuth || mock.authURL !== null) { connectOptions.onAuthURL?.( config.name,