From 0167bd60438845bbf5a44bce44122a3d3e330bc2 Mon Sep 17 00:00:00 2001 From: Aiden Cline Date: Fri, 30 Oct 2026 20:02:54 -0500 Subject: [PATCH] perf(core): skip the MCP SDK tool cache and validate only the called tool --- packages/core/src/mcp/client.ts | 50 ++++++++++++----- packages/core/test/mcp.test.ts | 95 ++++++++++++++++++++++++++++++++- 2 files changed, 129 insertions(+), 16 deletions(-) diff --git a/packages/core/src/mcp/client.ts b/packages/core/src/mcp/client.ts index 0f51129ac05..99d269f87cd 100644 --- a/packages/core/src/mcp/client.ts +++ b/packages/core/src/mcp/client.ts @@ -4,6 +4,7 @@ import path from "node:path" import { pathToFileURL } from "node:url" import { Client, + ProtocolError, SdkError, SdkErrorCode, SdkHttpError, @@ -37,6 +38,8 @@ const DEFAULT_CATALOG_TIMEOUT = 30_000 const DEFAULT_EXECUTION_TIMEOUT = 12 * 60 * 60 * 1_000 // 12 hours const ATTACHED_EXECUTION_TIMEOUT = 10 * 60 * 1_000 // 10 minutes const TERMINATE_TIMEOUT = 1_000 +// SEP-2243 HeaderMismatch; the SDK does not export the constant. +const HEADER_MISMATCH = -32020 const toError = (error: unknown) => (error instanceof Error ? error : new Error(String(error))) // HTTP statuses and network error codes live on error properties and are lost once flattened to a message. @@ -233,7 +236,7 @@ export const connect = Effect.fnUntraced(function* ( // Our SDK patch leaves tools/list items unchecked, so one malformed tool no longer fails every tool // on the server; they are checked here instead. The SDK runs a handler set on the transport before - // connect ahead of its own, on the same message, so it caches only the tools kept here. + // connect ahead of its own, on the same message, so it returns only the tools kept here. const onmessage = transport.onmessage transport.onmessage = (message, extra) => { if (isJSONRPCResultResponse(message) && Array.isArray(message.result.tools)) @@ -353,7 +356,13 @@ export const connect = Effect.fnUntraced(function* ( yield* Effect.tryPromise(() => client.close()).pipe(Effect.ignore) }), ) - const catalog = { timeout: timeout?.catalog ?? DEFAULT_CATALOG_TIMEOUT } + // Lists skip the SDK response cache: it re-stringifies every list, and on the first tools/call + // re-parses it and compiles an output validator for every listed tool. Each call passes its own + // listed definition instead, so only that tool's outputSchema is compiled and modern + // Mcp-Param-* header mirroring still sees the schema. Lists are refetched only on list_changed or + // reconnect, so modern list TTLs are not honored. + const catalog = { timeout: timeout?.catalog ?? DEFAULT_CATALOG_TIMEOUT, cacheMode: "bypass" as const } + const listed = new Map() const execution = timeout?.execution ?? DEFAULT_EXECUTION_TIMEOUT const request = (what: string, run: (signal: AbortSignal) => Promise) => Effect.tryPromise({ try: run, catch: failure }).pipe( @@ -379,7 +388,11 @@ export const connect = Effect.fnUntraced(function* ( tools: () => advertised("tools") ? list("list MCP tools", (signal) => client.listTools(undefined, { ...catalog, signal })).pipe( - Effect.map((r) => r.tools), + Effect.map((r) => { + listed.clear() + for (const tool of r.tools) listed.set(tool.name, tool) + return r.tools + }), ) : Effect.succeed([]), prompts: () => @@ -411,17 +424,26 @@ export const connect = Effect.fnUntraced(function* ( client.getPrompt({ name: input.name, arguments: input.args ?? {} }, { signal, timeout: execution }), ), callTool: (input) => - request("call MCP tool", (signal) => - client.callTool( - { - name: input.name, - arguments: input.args ?? {}, - ...(input.sessionID === undefined ? {} : { _meta: { "ai.opencode/sessionID": input.sessionID } }), - }, - // Requesting progress keeps long calls alive under the SDK's timeout; execution is the hard wall. - { signal, timeout: execution, onprogress: () => {} }, - ), - ).pipe(Effect.map(toCallToolResult)), + request("call MCP tool", (signal) => { + const params = { + name: input.name, + arguments: input.args ?? {}, + ...(input.sessionID === undefined ? {} : { _meta: { "ai.opencode/sessionID": input.sessionID } }), + } + // Requesting progress keeps long calls alive under the SDK's timeout; execution is the hard wall. + const options = { signal, timeout: execution, onprogress: () => {} } + const toolDefinition = listed.get(input.name) + if (toolDefinition === undefined) return client.callTool(params, options) + // A modern server rejects headers mirrored from a stale definition before running the tool. + // Without a definition the SDK refetches the list itself and retries. + return client + .callTool(params, { ...options, toolDefinition }) + .catch((error) => + error instanceof ProtocolError && error.code === HEADER_MISMATCH + ? client.callTool(params, options) + : Promise.reject(error), + ) + }).pipe(Effect.map(toCallToolResult)), onClose: (callback) => { client.onclose = () => callback(lastError ? `Connection closed: ${lastError}` : "Connection closed") }, diff --git a/packages/core/test/mcp.test.ts b/packages/core/test/mcp.test.ts index 45cf2bb4443..e9d99877ccd 100644 --- a/packages/core/test/mcp.test.ts +++ b/packages/core/test/mcp.test.ts @@ -2,7 +2,12 @@ import path from "node:path" import fs from "node:fs/promises" import os from "node:os" import { describe, expect, test } from "bun:test" -import { Client, InMemoryTransport, StreamableHTTPClientTransport } from "@modelcontextprotocol/client" +import { + Client, + InMemoryTransport, + StreamableHTTPClientTransport, + isJSONRPCRequest, +} from "@modelcontextprotocol/client" import { createMcpHandler, inputRequired, @@ -103,6 +108,8 @@ function resourceServer( * malformed description. */ malformedTool?: boolean + /** Listed after the default tools. A call's `structured` argument comes back as structured content. */ + tools?: McpClient.Tool[] /** Legacy servers give up on the URL elicitation after this many milliseconds. */ elicitationTimeout?: number respond?: (request: Request) => Response | undefined | Promise @@ -169,6 +176,7 @@ function resourceServer( : input.urlElicitation ? [{ name: "url-elicitation", inputSchema: { type: "object" as const, properties: {} } }] : [{ name: "echo", inputSchema: { type: "object" as const, properties: {} } }]), + ...(input.tools ?? []), ], }) }) @@ -218,7 +226,10 @@ function resourceServer( sessionID: request.params._meta?.["ai.opencode/sessionID"], progressToken: request.params._meta?.progressToken, }) - return Promise.resolve({ content: [] }) + const structured = request.params.arguments?.structured + return Promise.resolve( + structured === undefined ? { content: [] } : { content: [], structuredContent: structured }, + ) }) } protocol.setRequestHandler("prompts/list", () => Promise.resolve({ prompts: [{ name: "greet" }] })) @@ -1227,6 +1238,86 @@ describe.each([ ), ) }) + + test("validates structured content against the called tool's output schema", async () => { + await Effect.runPromise( + Effect.scoped( + Effect.gen(function* () { + const server = yield* resourceServer({ + modern: era === "modern", + resources: false, + tools: [ + { + name: "typed", + inputSchema: { type: "object" }, + outputSchema: { type: "object", properties: { value: { type: "string" } }, required: ["value"] }, + }, + ], + }) + const connection = yield* connect( + "typed", + new ConfigMCP.Remote({ type: "remote", url: server.url, oauth: false, protocol }), + import.meta.dir, + ) + yield* connection.tools() + + expect(yield* connection.callTool({ name: "typed", args: { structured: { value: "ok" } } })).toMatchObject({ + structured: { value: "ok" }, + }) + expect( + (yield* connection.callTool({ name: "typed", args: { structured: { value: 42 } } }).pipe(Effect.flip)) + .message, + ).toContain("Structured content does not match the tool's output schema") + expect((yield* connection.callTool({ name: "typed" }).pipe(Effect.flip)).message).toContain( + "has an output schema but did not return structured content", + ) + expect(yield* connection.callTool({ name: "echo", args: { structured: { value: 42 } } })).toMatchObject({ + structured: { value: 42 }, + }) + }), + ), + ) + }) +}) + +test("retries a modern MCP tool call whose listed header declaration went stale", async () => { + await Effect.runPromise( + Effect.scoped( + Effect.gen(function* () { + const region = { type: "string" as const, "x-mcp-header": "Region" } + const rejected: Array = [] + const server = yield* resourceServer({ + modern: true, + resources: false, + tools: [{ name: "regional", inputSchema: { type: "object", properties: { region } } }], + // Stands in for the server-side check: the region must arrive in the currently declared header. + respond: async (request) => { + const body: unknown = request.method === "POST" ? await request.clone().json() : undefined + if (!isJSONRPCRequest(body) || body.method !== "tools/call") return undefined + if (request.headers.get(`mcp-param-${region["x-mcp-header"]}`) === "eu") return undefined + rejected.push(request.headers.get("mcp-param-region")) + return Response.json( + { jsonrpc: "2.0", id: body.id, error: { code: -32020, message: "Header mismatch" } }, + { status: 400 }, + ) + }, + }) + const connection = yield* connect( + "regional", + new ConfigMCP.Remote({ type: "remote", url: server.url, oauth: false, protocol: "2026-07-28" }), + import.meta.dir, + ) + yield* connection.tools() + region["x-mcp-header"] = "Zone" + + expect((yield* connection.callTool({ name: "regional", args: { region: "eu" } })).isError).toBe(false) + expect(server.state.toolCalls.map((call) => call.arguments)).toEqual([{ region: "eu" }]) + // First the stale Region header, then no header; the SDK then relists and sends Zone. + expect(rejected).toEqual(["eu", null]) + expect(server.state.toolLists).toBe(2) + }), + ), + ) }) test("lists, reads, and reports MCP resource changes", async () => {