diff --git a/packages/agent/src/agent-loop.ts b/packages/agent/src/agent-loop.ts index e44086909..abe5cf64d 100644 --- a/packages/agent/src/agent-loop.ts +++ b/packages/agent/src/agent-loop.ts @@ -517,12 +517,19 @@ async function executeToolCalls( const hasSequentialToolCall = toolCalls.some( (tc) => currentContext.tools?.find((t) => t.name === tc.name)?.executionMode === "sequential", ); + const nestedQueue: NestedCallQueue = { tail: Promise.resolve() }; if (config.toolExecution === "sequential" || hasSequentialToolCall) { - return executeToolCallsSequential(currentContext, assistantMessage, toolCalls, config, signal, emit); + return executeToolCallsSequential(currentContext, assistantMessage, toolCalls, config, signal, emit, nestedQueue); } - return executeToolCallsParallel(currentContext, assistantMessage, toolCalls, config, signal, emit); + return executeToolCallsParallel(currentContext, assistantMessage, toolCalls, config, signal, emit, nestedQueue); } +/** + * Serializes nested calls that must not run concurrently (sequential tools, or all nested calls when + * the loop runs sequentially). Shared by all tool calls of one assistant message. + */ +type NestedCallQueue = { tail: Promise }; + type ExecutedToolCallBatch = { messages: ToolResultMessage[]; terminate: boolean; @@ -535,6 +542,7 @@ async function executeToolCallsSequential( config: AgentLoopConfig, signal: AbortSignal | undefined, emit: AgentEventSink, + nestedQueue: NestedCallQueue, ): Promise { const finalizedCalls: FinalizedToolCallOutcome[] = []; const messages: ToolResultMessage[] = []; @@ -563,6 +571,7 @@ async function executeToolCallsSequential( preparation, signal, emitToolExecutionUpdate(preparation.toolCall, emit), + nestedQueue, ); finalized = await finalizeExecutedToolCall( currentContext, @@ -598,6 +607,7 @@ async function executeToolCallsParallel( config: AgentLoopConfig, signal: AbortSignal | undefined, emit: AgentEventSink, + nestedQueue: NestedCallQueue, ): Promise { const finalizedCalls: FinalizedToolCallEntry[] = []; @@ -641,6 +651,7 @@ async function executeToolCallsParallel( preparation, signal, emitToolExecutionUpdate(preparation.toolCall, emit), + nestedQueue, ); const finalized = await finalizeExecutedToolCall( currentContext, @@ -802,6 +813,8 @@ function emitToolExecutionUpdate(toolCall: AgentToolCall, emit: AgentEventSink): /** * Services handed to a tool while it runs. Nested calls get ids `/` and run through * the same prepare, execute, and finalize steps as model-issued calls, without emitting events. + * `nestedQueue` is undefined inside a call that already holds the queue, so its own nested calls + * do not wait on themselves. */ function createToolContext( currentContext: AgentContext, @@ -809,6 +822,7 @@ function createToolContext( config: AgentLoopConfig, toolCall: AgentToolCall, signal: AbortSignal | undefined, + nestedQueue: NestedCallQueue | undefined, ): AgentToolContext { let nextNestedId = 1; return { @@ -828,6 +842,7 @@ function createToolContext( }, options?.signal ?? signal, options?.onUpdate, + nestedQueue, ), }; } @@ -840,35 +855,53 @@ async function executeNestedToolCall( toolCall: AgentToolCall, signal: AbortSignal | undefined, onUpdate: AgentToolUpdateCallback | undefined, + nestedQueue: NestedCallQueue | undefined, ): Promise { - const preparation = await prepareToolCall( - currentContext, - assistantMessage, - toolCall, - config, - signal, - parentToolCall, - ); - if (preparation.kind === "immediate") { - return { toolCall, result: preparation.result, isError: preparation.isError }; + const exclusive = + nestedQueue && + (config.toolExecution === "sequential" || + currentContext.tools?.find((t) => t.name === toolCall.name)?.executionMode === "sequential"); + let release: (() => void) | undefined; + if (exclusive) { + const previous = nestedQueue.tail; + nestedQueue.tail = new Promise((resolve) => { + release = resolve; + }); + await previous; + } + try { + const preparation = await prepareToolCall( + currentContext, + assistantMessage, + toolCall, + config, + signal, + parentToolCall, + ); + if (preparation.kind === "immediate") { + return { toolCall, result: preparation.result, isError: preparation.isError }; + } + const executed = await executePreparedToolCall( + currentContext, + assistantMessage, + config, + preparation, + signal, + (partialResult) => onUpdate?.(partialResult), + exclusive ? undefined : nestedQueue, + ); + return await finalizeExecutedToolCall( + currentContext, + assistantMessage, + preparation, + executed, + config, + signal, + parentToolCall, + ); + } finally { + release?.(); } - const executed = await executePreparedToolCall( - currentContext, - assistantMessage, - config, - preparation, - signal, - (partialResult) => onUpdate?.(partialResult), - ); - return finalizeExecutedToolCall( - currentContext, - assistantMessage, - preparation, - executed, - config, - signal, - parentToolCall, - ); } async function executePreparedToolCall( @@ -878,6 +911,7 @@ async function executePreparedToolCall( prepared: PreparedToolCall, signal: AbortSignal | undefined, onUpdate: ToolUpdateSink, + nestedQueue: NestedCallQueue | undefined, ): Promise { const updateEvents: Promise[] = []; let acceptingUpdates = true; @@ -891,7 +925,7 @@ async function executePreparedToolCall( if (!acceptingUpdates) return; updateEvents.push(Promise.resolve(onUpdate(partialResult))); }, - createToolContext(currentContext, assistantMessage, config, prepared.toolCall, signal), + createToolContext(currentContext, assistantMessage, config, prepared.toolCall, signal, nestedQueue), ); acceptingUpdates = false; await Promise.all(updateEvents); diff --git a/packages/agent/test/agent-loop.test.ts b/packages/agent/test/agent-loop.test.ts index bbf9d89a6..6efc03040 100644 --- a/packages/agent/test/agent-loop.test.ts +++ b/packages/agent/test/agent-loop.test.ts @@ -2242,6 +2242,46 @@ describe("nested tool calls", () => { await stream.result(); expect(seen).toEqual([undefined, { secret: 2 }]); }); + + it("serializes concurrent nested calls to sequential tools", async () => { + const schema = Type.Object({}); + let active = 0; + let maxActive = { sequential: 0, parallel: 0 }; + const makeTool = (name: "sequential" | "parallel"): AgentTool => ({ + name, + label: name, + description: name, + parameters: schema, + executionMode: name === "sequential" ? "sequential" : undefined, + async execute() { + active++; + maxActive = { ...maxActive, [name]: Math.max(maxActive[name], active) }; + await new Promise((resolve) => setTimeout(resolve, 5)); + active--; + return { content: [], details: {} }; + }, + }); + const runner: AgentTool = { + name: "runner", + label: "Runner", + description: "Runs other tools", + parameters: schema, + async execute(_id, _params, _signal, _onUpdate, context) { + await Promise.all([1, 2, 3].map(() => context!.executeTool("sequential", {}))); + await Promise.all([1, 2, 3].map(() => context!.executeTool("parallel", {}))); + return { content: [], details: {} }; + }, + }; + const stream = agentLoop( + [createUserMessage("go")], + { messages: [], tools: [makeTool("sequential"), makeTool("parallel"), runner] }, + { model: createModel(), convertToLlm: identityConverter }, + undefined, + scriptedStream({ type: "toolCall", id: "call-1", name: "runner", arguments: {} }), + ); + await stream.result(); + expect(maxActive).toEqual({ sequential: 1, parallel: 3 }); + }); }); describe("nestedOnly tools", () => { diff --git a/packages/coding-agent/src/core/tools/codemode-execute.ts b/packages/coding-agent/src/core/tools/codemode-execute.ts index 88536a659..293353cd7 100644 --- a/packages/coding-agent/src/core/tools/codemode-execute.ts +++ b/packages/coding-agent/src/core/tools/codemode-execute.ts @@ -275,7 +275,9 @@ export async function executeCodemode( let result: CodemodeResult; try { - result = await sandbox.execute(code, { signal, store: readCodemodeStore(ctx.sessionManager.getBranch()) }); + // Without a session (plain Agent or direct call) the store starts empty and writes are dropped. + const store = ctx.sessionManager ? readCodemodeStore(ctx.sessionManager.getBranch()) : {}; + result = await sandbox.execute(code, { signal, store }); } finally { await sandbox.close(); } diff --git a/packages/coding-agent/src/extensions/mcp/oauth.ts b/packages/coding-agent/src/extensions/mcp/oauth.ts index d236ccca0..7a5bb309f 100644 --- a/packages/coding-agent/src/extensions/mcp/oauth.ts +++ b/packages/coding-agent/src/extensions/mcp/oauth.ts @@ -120,20 +120,23 @@ function createProvider( * Auth provider for MCP connections: sends the stored access token and refreshes it after a 401. * Throws `McpOAuthAuthorizationRequiredError` when the user has to sign in. `onChallenge` receives * the server's `WWW-Authenticate` challenge so sign-in can use its resource metadata URL and scope. + * `settings` is only called when a refresh is needed, so a secret that fails to resolve fails the + * refresh instead of the whole connection setup. */ export function createMcpAuthProvider(options: { serverUrl: string; store: McpOAuthStateStore; - settings: McpOAuthSettings; + settings: () => McpOAuthSettings; onChallenge: (challenge: OAuthChallenge) => void; }): AuthProvider { - const { serverUrl, store, settings } = options; + const { serverUrl, store } = options; return { token: async () => (await store.load())?.tokens?.access_token, onUnauthorized: async (context) => { options.onChallenge(parseWwwAuthenticate(context.response.headers.get("www-authenticate"))); const state = await store.load(); if (!state?.tokens?.refresh_token) throw new McpOAuthAuthorizationRequiredError(); + const settings = options.settings(); const redirectUrl = configuredRedirectUrl(settings) ?? registeredRedirectUrls(state.clientInformation)[0] ?? diff --git a/packages/coding-agent/src/extensions/mcp/runtime.ts b/packages/coding-agent/src/extensions/mcp/runtime.ts index a8a736911..ee31acf96 100644 --- a/packages/coding-agent/src/extensions/mcp/runtime.ts +++ b/packages/coding-agent/src/extensions/mcp/runtime.ts @@ -11,6 +11,7 @@ import { type CallToolResult, McpClient, type McpRequestOptions, + McpSessionExpiredError, type Tool as McpTool, type McpTransport, StdioTransport, @@ -109,7 +110,7 @@ export class McpServerConnection implements McpToolCaller { ? createMcpAuthProvider({ serverUrl: url, store: options.credentials.forServer(url), - settings: this.oauthSettings(), + settings: () => this.oauthSettings(), onChallenge: (challenge) => { this.challenge = challenge; }, @@ -149,14 +150,23 @@ export class McpServerConnection implements McpToolCaller { } async callTool(name: string, args: Record, options: McpRequestOptions): Promise { - const client = await this.getClient(); - try { - return await client.callTool(name, args, options); - } catch (error) { - if (!(error instanceof McpOAuthAuthorizationRequiredError)) throw error; - await this.dropClient(client); - this.markNeedsAuth(); - throw new Error(signInRequiredMessage(this.entry.name)); + for (let attempt = 1; ; attempt++) { + const client = await this.getClient(); + try { + return await client.callTool(name, args, options); + } catch (error) { + if (error instanceof McpSessionExpiredError && attempt === 1) { + // The server no longer knows the session (restart, deploy), so it did not run this call. + // Retry once on a new session. The old client is detached but not closed: closing would + // fail its other in-flight calls, which instead get the same 404 and retry the same way. + if (this.client === client) this.client = undefined; + continue; + } + if (!(error instanceof McpOAuthAuthorizationRequiredError)) throw error; + await this.dropClient(client); + this.markNeedsAuth(); + throw new Error(signInRequiredMessage(this.entry.name)); + } } } diff --git a/packages/coding-agent/test/mcp-extension.test.ts b/packages/coding-agent/test/mcp-extension.test.ts index bb48ca23c..c6b76e43a 100644 --- a/packages/coding-agent/test/mcp-extension.test.ts +++ b/packages/coding-agent/test/mcp-extension.test.ts @@ -1,8 +1,12 @@ import { mkdirSync, mkdtempSync, rmSync, writeFileSync } from "node:fs"; import { tmpdir } from "node:os"; import { join } from "node:path"; +import { type JsonRpcMessage, LATEST_PROTOCOL_VERSION, McpSessionExpiredError } from "@earendil-works/pi-mcp"; +import { createInMemoryTransportPair } from "@earendil-works/pi-mcp/testing"; import { afterEach, describe, expect, it } from "vitest"; -import { loadMcpConfig } from "../src/extensions/mcp/config.ts"; +import { InMemoryAuthStorageBackend } from "../src/core/auth-storage.ts"; +import { loadMcpConfig, type McpServerEntry } from "../src/extensions/mcp/config.ts"; +import { McpOAuthCredentialStore, McpServerConnection } from "../src/extensions/mcp/runtime.ts"; import { convertMcpResult, createMcpToolName } from "../src/extensions/mcp/tools.ts"; // Config values are resolved at connect time, so the literal reference must survive loading. @@ -93,3 +97,74 @@ describe("MCP tools", () => { ); }); }); + +describe("MCP connections", () => { + /** In-memory server that answers initialize, tools/list, and tools/call with "ok". */ + function createTransport(options: { expireFirstCall?: boolean } = {}) { + const pair = createInMemoryTransportPair(); + pair.server.onMessage((message) => { + if (!("id" in message) || !("method" in message)) return; + const result = + message.method === "initialize" + ? { + protocolVersion: LATEST_PROTOCOL_VERSION, + capabilities: { tools: {} }, + serverInfo: { name: "fake", version: "1.0.0" }, + } + : message.method === "tools/list" + ? { tools: [] } + : { content: [{ type: "text", text: "ok" }] }; + queueMicrotask(() => void pair.server.send({ jsonrpc: "2.0", id: message.id, result })); + }); + void pair.server.start(); + if (options.expireFirstCall) { + const send = pair.client.send.bind(pair.client); + // Simulates the HTTP transport's 404 for a session the server no longer knows. + pair.client.send = async (message: JsonRpcMessage) => { + if ("method" in message && message.method === "tools/call") throw new McpSessionExpiredError("gone"); + return send(message); + }; + } + return pair.client; + } + + function connect(entry: McpServerEntry, transports: (() => ReturnType)[]) { + let opened = 0; + const connection = new McpServerConnection({ + entry, + cwd: process.cwd(), + createTransport: () => transports[opened++](), + credentials: new McpOAuthCredentialStore(new InMemoryAuthStorageBackend()), + onTools: () => {}, + }); + return { connection, opened: () => opened }; + } + + it("starts a new session and retries once when the session expired", async () => { + const { connection, opened } = connect({ name: "fake", config: { command: "unused" }, source: "test" }, [ + () => createTransport({ expireFirstCall: true }), + () => createTransport(), + ]); + const results = await Promise.all([connection.callTool("echo", {}, {}), connection.callTool("echo", {}, {})]); + expect(results).toEqual([ + { content: [{ type: "text", text: "ok" }] }, + { content: [{ type: "text", text: "ok" }] }, + ]); + expect(opened()).toBe(2); + await connection.close(); + }); + + it("resolves the OAuth client secret lazily", async () => { + const { connection } = connect( + { + name: "fake", + config: { url: "http://unused.invalid", oauth: { clientSecret: "!exit 1" } }, + source: "test", + }, + [() => createTransport()], + ); + expect(await connection.callTool("echo", {}, {})).toEqual({ content: [{ type: "text", text: "ok" }] }); + expect(() => connection.oauthSettings()).toThrow("oauth.clientSecret"); + await connection.close(); + }); +}); diff --git a/packages/coding-agent/test/suite/agent-session-codemode.test.ts b/packages/coding-agent/test/suite/agent-session-codemode.test.ts index e470ada75..b1eac7f36 100644 --- a/packages/coding-agent/test/suite/agent-session-codemode.test.ts +++ b/packages/coding-agent/test/suite/agent-session-codemode.test.ts @@ -377,6 +377,19 @@ describe("codemode options and store", () => { expect(storeEntries(harness)).toEqual([]); }); + it("runs without a session, starting from an empty store", async () => { + const tool = createCodemodeTool(); + const toolCall = { type: "toolCall" as const, id: "direct", name: tool.name, arguments: {} }; + const result = await tool.execute("direct", { code: increment }, undefined, undefined, { + toolCall, + tools: [], + executeTool: async () => { + throw new Error("unexpected nested call"); + }, + }); + expect(result.content).toEqual([{ type: "text", text: "1" }]); + }); + it("loads the values written on the current branch", async () => { const harness = await setup(); await run(harness, increment); diff --git a/packages/mcp/src/client.ts b/packages/mcp/src/client.ts index c9bb0972d..d13b89474 100644 --- a/packages/mcp/src/client.ts +++ b/packages/mcp/src/client.ts @@ -154,7 +154,8 @@ export class McpClient { this.transport = transport; this.disposers = [ transport.onMessage((message) => this.handleMessage(message)), - transport.onError((error) => this.handleTransportError(error)), + // Transport errors are reported only. Pending requests fail when the transport closes. + transport.onError((error) => this.emitError(error)), transport.onClose(() => this.handleTransportClose()), ]; @@ -454,11 +455,6 @@ export class McpClient { } } - private handleTransportError(error: Error): void { - this.rejectPending(error); - this.emitError(error); - } - private handleTransportClose(): void { this.markClosed(new McpConnectionClosedError()); } diff --git a/packages/mcp/src/transports/streamable-http.ts b/packages/mcp/src/transports/streamable-http.ts index abb427455..00b7e52a9 100644 --- a/packages/mcp/src/transports/streamable-http.ts +++ b/packages/mcp/src/transports/streamable-http.ts @@ -1,5 +1,13 @@ import type { AuthProvider, McpFetch } from "../auth-provider.ts"; -import { type JsonRpcMessage, McpConnectionClosedError, parseJsonRpcMessage } from "../protocol/jsonrpc.ts"; +import { + isJsonRpcRequest, + isJsonRpcResponse, + JSON_RPC_ERROR_CODES, + type JsonRpcMessage, + McpConnectionClosedError, + parseJsonRpcMessage, + toError, +} from "../protocol/jsonrpc.ts"; import { DEFAULT_MAX_MESSAGE_BYTES, type McpTransport, TransportEvents } from "./transport.ts"; const MAX_ERROR_BODY_BYTES = 8 * 1024; @@ -201,7 +209,23 @@ export class StreamableHttpTransport extends TransportEvents implements McpTrans return; } if (type === "text/event-stream" && response.body) { - void this.consumeSse(response.body).catch((error) => this.emitError(error)); + const requestId = isJsonRpcRequest(message) ? message.id : undefined; + let answered = false; + void this.consumeSse(response.body, (received) => { + if (isJsonRpcResponse(received) && received.id === requestId) answered = true; + }).catch((error) => { + this.emitError(error); + if (requestId === undefined || answered || this.closed) return; + // Fail only the request this stream belongs to. The connection itself is still usable. + this.emitMessage({ + jsonrpc: "2.0", + id: requestId, + error: { + code: JSON_RPC_ERROR_CODES.internalError, + message: `MCP response stream failed: ${toError(error).message}`, + }, + }); + }); return; } if (response.headers.get("content-length") === "0") return; @@ -232,12 +256,17 @@ export class StreamableHttpTransport extends TransportEvents implements McpTrans throw new McpHttpError(response.status, `MCP HTTP request failed with status ${response.status}`, body); } - private async consumeSse(stream: ReadableStream): Promise { + private async consumeSse( + stream: ReadableStream, + onMessage?: (message: JsonRpcMessage) => void, + ): Promise { await consumeSseStream(stream, { maxEventBytes: this.options.maxMessageBytes ?? DEFAULT_MAX_MESSAGE_BYTES, onEvent: (event) => { if (event.id) this.lastEventId = event.id; - this.emitMessage(parseJsonRpcMessage(JSON.parse(event.data))); + const message = parseJsonRpcMessage(JSON.parse(event.data)); + onMessage?.(message); + this.emitMessage(message); }, }); } diff --git a/packages/mcp/test/client.test.ts b/packages/mcp/test/client.test.ts index 5eea8a98a..59eaff628 100644 --- a/packages/mcp/test/client.test.ts +++ b/packages/mcp/test/client.test.ts @@ -189,6 +189,29 @@ describe("McpClient", () => { await client.close(); }); + it("reports transport errors without failing pending requests", async () => { + const { clientTransport, server } = await createServer(); + const client = new McpClient({ name: "test-client", version: "1.0.0" }); + await client.connect(clientTransport); + const errors: Error[] = []; + client.onError((error) => errors.push(error)); + let respond = () => {}; + server.setHandler( + "tools/call", + () => + new Promise((resolve) => { + respond = () => resolve({ content: [] }); + }), + ); + const call = client.callTool("wait"); + await new Promise((resolve) => setTimeout(resolve, 0)); + clientTransport.emitError(new Error("stray log line")); + respond(); + expect(await call).toEqual({ content: [] }); + expect(errors.map((error) => error.message)).toEqual(["stray log line"]); + await client.close(); + }); + it("answers roots/list and dispatches notifications", async () => { const { clientTransport, server } = await createServer(); const client = new McpClient({ diff --git a/packages/mcp/test/streamable-http.test.ts b/packages/mcp/test/streamable-http.test.ts index 40d15513a..e82760a10 100644 --- a/packages/mcp/test/streamable-http.test.ts +++ b/packages/mcp/test/streamable-http.test.ts @@ -28,6 +28,7 @@ async function protocolHandler( request: IncomingMessage, response: ServerResponse, requests: RecordedRequest[], + body?: Record, ): Promise { if (request.method === "GET") { requests.push({ method: "GET", headers: request.headers }); @@ -41,7 +42,7 @@ async function protocolHandler( response.end(); return; } - const message = JSON.parse(await readBody(request)) as Record; + const message = body ?? (JSON.parse(await readBody(request)) as Record); requests.push({ method: request.method ?? "", headers: request.headers, message }); if (!("id" in message)) { response.statusCode = 202; @@ -137,6 +138,36 @@ describe("StreamableHttpTransport", () => { } satisfies Partial); }); + it("fails only the request whose SSE stream breaks", async () => { + let releaseSlow = () => {}; + const slowGate = new Promise((resolve) => { + releaseSlow = resolve; + }); + const { url } = await startServer(async (request, response, requests) => { + if (request.method !== "POST") return protocolHandler(request, response, requests); + const message = JSON.parse(await readBody(request)) as Record; + const name = (message.params as { name?: string } | undefined)?.name; + if (name === "broken") { + response.writeHead(200, { "content-type": "text/event-stream" }); + response.end("data: not json\n\n"); + return; + } + if (name === "slow") await slowGate; + await protocolHandler(request, response, requests, message); + }); + const client = new McpClient({ name: "http-test", version: "1.0.0" }); + const errors: Error[] = []; + client.onError((error) => errors.push(error)); + await client.connect(new StreamableHttpTransport({ url, openGetStream: false })); + + const slow = client.callTool("slow"); + await expect(client.callTool("broken")).rejects.toThrow("MCP response stream failed"); + releaseSlow(); + expect(await slow).toEqual({ content: [{ type: "text", text: "hello" }] }); + expect(errors).toHaveLength(1); + await client.close(); + }); + it("classifies an expired established session", async () => { let posts = 0; const { url } = await startServer(async (request, response, requests) => {