diff options
Diffstat (limited to 'packages/mcp/src')
| -rw-r--r-- | packages/mcp/src/client.test.ts | 456 | ||||
| -rw-r--r-- | packages/mcp/src/client.ts | 242 | ||||
| -rw-r--r-- | packages/mcp/src/config.test.ts | 218 | ||||
| -rw-r--r-- | packages/mcp/src/config.ts | 114 | ||||
| -rw-r--r-- | packages/mcp/src/extension.test.ts | 697 | ||||
| -rw-r--r-- | packages/mcp/src/extension.ts | 367 | ||||
| -rw-r--r-- | packages/mcp/src/framing.test.ts | 241 | ||||
| -rw-r--r-- | packages/mcp/src/framing.ts | 238 | ||||
| -rw-r--r-- | packages/mcp/src/index.ts | 46 | ||||
| -rw-r--r-- | packages/mcp/src/manager.test.ts | 417 | ||||
| -rw-r--r-- | packages/mcp/src/manager.ts | 367 | ||||
| -rw-r--r-- | packages/mcp/src/registry.test.ts | 416 | ||||
| -rw-r--r-- | packages/mcp/src/registry.ts | 104 | ||||
| -rw-r--r-- | packages/mcp/src/rpc.test.ts | 169 | ||||
| -rw-r--r-- | packages/mcp/src/rpc.ts | 164 | ||||
| -rw-r--r-- | packages/mcp/src/timeout.test.ts | 105 | ||||
| -rw-r--r-- | packages/mcp/src/timeout.ts | 114 | ||||
| -rw-r--r-- | packages/mcp/src/transport.test.ts | 301 | ||||
| -rw-r--r-- | packages/mcp/src/transport.ts | 146 | ||||
| -rw-r--r-- | packages/mcp/src/types.ts | 76 |
20 files changed, 2845 insertions, 2153 deletions
diff --git a/packages/mcp/src/client.test.ts b/packages/mcp/src/client.test.ts index 695bdcc..4542a36 100644 --- a/packages/mcp/src/client.test.ts +++ b/packages/mcp/src/client.test.ts @@ -3,204 +3,270 @@ import { McpClient } from "./client.js"; import type { Connection } from "./transport.js"; function makeMockConnection(): Connection & { - responses: Map<string, unknown>; - feedResponse: (method: string, result: unknown) => void; - notifications: Array<{ method: string; params: unknown }>; + responses: Map<string, unknown>; + feedResponse: (method: string, result: unknown) => void; + notifications: Array<{ method: string; params: unknown }>; } { - const responses = new Map<string, unknown>(); - const pendingRequests = new Map<number, { method: string; resolve: (v: unknown) => void }>(); - let nextId = 1; - const notifications: Array<{ method: string; params: unknown }> = []; - const notificationHandlers = new Map<string, (params: unknown) => void>(); - - return { - responses, - notifications, - feedResponse: (_method: string, _result: unknown) => {}, - send: (method: string, _params?: unknown) => { - const id = nextId++; - return new Promise((resolve) => { - pendingRequests.set(id, { method, resolve }); - // Auto-respond for initialize - if (method === "initialize") { - resolve({ - protocolVersion: "2025-11-25", - capabilities: { tools: { listChanged: true } }, - serverInfo: { name: "test-server", version: "1.0.0" }, - }); - } else if (method === "tools/list") { - resolve({ - tools: [ - { - name: "test_tool", - description: "A test tool", - inputSchema: { type: "object", properties: { input: { type: "string" } } }, - }, - ], - }); - } else if (method === "tools/call") { - resolve({ - content: [{ type: "text", text: "result from tool" }], - isError: false, - }); - } - }); - }, - notify: (method: string, params?: unknown) => { - notifications.push({ method, params }); - }, - onNotification: (method: string, handler: (params: unknown) => void) => { - notificationHandlers.set(method, handler); - }, - close: () => {}, - pid: 999, - }; + const responses = new Map<string, unknown>(); + const pendingRequests = new Map<number, { method: string; resolve: (v: unknown) => void }>(); + let nextId = 1; + const notifications: Array<{ method: string; params: unknown }> = []; + const notificationHandlers = new Map<string, (params: unknown) => void>(); + + return { + responses, + notifications, + feedResponse: (_method: string, _result: unknown) => {}, + send: (method: string, _params?: unknown) => { + const id = nextId++; + return new Promise((resolve) => { + pendingRequests.set(id, { method, resolve }); + // Auto-respond for initialize + if (method === "initialize") { + resolve({ + protocolVersion: "2025-11-25", + capabilities: { tools: { listChanged: true } }, + serverInfo: { name: "test-server", version: "1.0.0" }, + }); + } else if (method === "tools/list") { + resolve({ + tools: [ + { + name: "test_tool", + description: "A test tool", + inputSchema: { type: "object", properties: { input: { type: "string" } } }, + }, + ], + }); + } else if (method === "tools/call") { + resolve({ + content: [{ type: "text", text: "result from tool" }], + isError: false, + }); + } + }); + }, + notify: (method: string, params?: unknown) => { + notifications.push({ method, params }); + }, + onNotification: (method: string, handler: (params: unknown) => void) => { + notificationHandlers.set(method, handler); + }, + close: () => {}, + pid: 999, + }; } describe("McpClient", () => { - it("initialize sends correct protocolVersion + capabilities", async () => { - const conn = makeMockConnection(); - const client = new McpClient({ connection: conn }); - - const result = await client.initialize(); - - expect(result.protocolVersion).toBe("2025-11-25"); - expect(result.capabilities.tools?.listChanged).toBe(true); - expect(result.serverInfo.name).toBe("test-server"); - expect(client.getState()).toBe("connected"); + it("initialize sends correct protocolVersion + capabilities", async () => { + const conn = makeMockConnection(); + const client = new McpClient({ connection: conn }); + + const result = await client.initialize(); + + expect(result.protocolVersion).toBe("2025-11-25"); + expect(result.capabilities.tools?.listChanged).toBe(true); + expect(result.serverInfo.name).toBe("test-server"); + expect(client.getState()).toBe("connected"); - // Should have sent notifications/initialized - expect(conn.notifications.length).toBe(1); - expect(conn.notifications[0].method).toBe("notifications/initialized"); - }); - - it("listTools returns parsed tools", async () => { - const conn = makeMockConnection(); - const client = new McpClient({ connection: conn }); - - await client.initialize(); - const tools = await client.listTools(); - - expect(tools.length).toBe(1); - expect(tools[0].name).toBe("test_tool"); - expect(tools[0].description).toBe("A test tool"); - }); - - it("callTool sends name + arguments", async () => { - const conn = makeMockConnection(); - let callParams: unknown = null; - const origSend = conn.send.bind(conn); - conn.send = (method: string, params?: unknown) => { - if (method === "tools/call") callParams = params; - return origSend(method, params); - }; - - const client = new McpClient({ connection: conn }); - - await client.initialize(); - const result = await client.callTool("test_tool", { input: "hello" }); - - expect(callParams).toEqual({ name: "test_tool", arguments: { input: "hello" } }); - expect(result.content).toEqual([{ type: "text", text: "result from tool" }]); - expect(result.isError).toBe(false); - }); - - it("list_changed triggers re-list", async () => { - const conn = makeMockConnection(); - const notificationHandlers = new Map<string, (params: unknown) => void>(); - conn.onNotification = (method: string, handler: (params: unknown) => void) => { - notificationHandlers.set(method, handler); - }; - - const client = new McpClient({ connection: conn }); - - let toolsChangedFired = false; - client.onToolsChanged(() => { - toolsChangedFired = true; - }); - - await client.initialize(); - - // Simulate list_changed notification - const handler = notificationHandlers.get("notifications/tools/list_changed"); - expect(handler).toBeDefined(); - handler?.(undefined); - - expect(toolsChangedFired).toBe(true); - }); - - it("handles server error on initialize", async () => { - const conn = makeMockConnection(); - conn.send = (method: string) => { - if (method === "initialize") { - return Promise.reject(new Error("Server startup failed")); - } - return Promise.resolve({}); - }; - - const client = new McpClient({ connection: conn }); - - await expect(client.initialize()).rejects.toThrow("Server startup failed"); - expect(client.getState()).toBe("error"); - }); - - it("callTool rejects when not connected", async () => { - const conn = makeMockConnection(); - const client = new McpClient({ connection: conn }); - - await expect(client.callTool("test", {})).rejects.toThrow("Client not connected"); - }); - - it("listTools rejects when not connected", async () => { - const conn = makeMockConnection(); - const client = new McpClient({ connection: conn }); - - await expect(client.listTools()).rejects.toThrow("Client not connected"); - }); - - it("close sets state to disconnected", async () => { - const conn = makeMockConnection(); - const client = new McpClient({ connection: conn }); - - await client.initialize(); - expect(client.getState()).toBe("connected"); - - client.close(); - expect(client.getState()).toBe("disconnected"); - }); - - it("callTool with abort signal", async () => { - const conn = makeMockConnection(); - let resolveRequest: ((v: unknown) => void) | null = null; - conn.send = (method: string) => { - if (method === "tools/call") { - return new Promise((resolve) => { - resolveRequest = resolve; - }); - } - if (method === "initialize") { - return Promise.resolve({ - protocolVersion: "2025-11-25", - capabilities: {}, - serverInfo: { name: "test", version: "1.0.0" }, - }); - } - return Promise.resolve({}); - }; - - const client = new McpClient({ connection: conn }); - await client.initialize(); - - const controller = new AbortController(); - const callPromise = client.callTool("test", {}, controller.signal); - - controller.abort(); - - await expect(callPromise).rejects.toThrow("Aborted"); - - // Clean up - resolveRequest?.({ - content: [{ type: "text", text: "too late" }], - }); - }); + // Should have sent notifications/initialized + expect(conn.notifications.length).toBe(1); + expect(conn.notifications[0].method).toBe("notifications/initialized"); + }); + + it("listTools returns parsed tools", async () => { + const conn = makeMockConnection(); + const client = new McpClient({ connection: conn }); + + await client.initialize(); + const tools = await client.listTools(); + + expect(tools.length).toBe(1); + expect(tools[0].name).toBe("test_tool"); + expect(tools[0].description).toBe("A test tool"); + }); + + it("callTool sends name + arguments", async () => { + const conn = makeMockConnection(); + let callParams: unknown = null; + const origSend = conn.send.bind(conn); + conn.send = (method: string, params?: unknown) => { + if (method === "tools/call") callParams = params; + return origSend(method, params); + }; + + const client = new McpClient({ connection: conn }); + + await client.initialize(); + const result = await client.callTool("test_tool", { input: "hello" }); + + expect(callParams).toEqual({ name: "test_tool", arguments: { input: "hello" } }); + expect(result.content).toEqual([{ type: "text", text: "result from tool" }]); + expect(result.isError).toBe(false); + }); + + it("list_changed triggers re-list", async () => { + const conn = makeMockConnection(); + const notificationHandlers = new Map<string, (params: unknown) => void>(); + conn.onNotification = (method: string, handler: (params: unknown) => void) => { + notificationHandlers.set(method, handler); + }; + + const client = new McpClient({ connection: conn }); + + let toolsChangedFired = false; + client.onToolsChanged(() => { + toolsChangedFired = true; + }); + + await client.initialize(); + + // Simulate list_changed notification + const handler = notificationHandlers.get("notifications/tools/list_changed"); + expect(handler).toBeDefined(); + handler?.(undefined); + + expect(toolsChangedFired).toBe(true); + }); + + it("handles server error on initialize", async () => { + const conn = makeMockConnection(); + conn.send = (method: string) => { + if (method === "initialize") { + return Promise.reject(new Error("Server startup failed")); + } + return Promise.resolve({}); + }; + + const client = new McpClient({ connection: conn }); + + await expect(client.initialize()).rejects.toThrow("Server startup failed"); + expect(client.getState()).toBe("error"); + }); + + it("callTool rejects when not connected", async () => { + const conn = makeMockConnection(); + const client = new McpClient({ connection: conn }); + + await expect(client.callTool("test", {})).rejects.toThrow("Client not connected"); + }); + + it("listTools rejects when not connected", async () => { + const conn = makeMockConnection(); + const client = new McpClient({ connection: conn }); + + await expect(client.listTools()).rejects.toThrow("Client not connected"); + }); + + it("close sets state to disconnected", async () => { + const conn = makeMockConnection(); + const client = new McpClient({ connection: conn }); + + await client.initialize(); + expect(client.getState()).toBe("connected"); + + client.close(); + expect(client.getState()).toBe("disconnected"); + }); + + it("callTool with abort signal", async () => { + const conn = makeMockConnection(); + let resolveRequest: ((v: unknown) => void) | null = null; + conn.send = (method: string) => { + if (method === "tools/call") { + return new Promise((resolve) => { + resolveRequest = resolve; + }); + } + if (method === "initialize") { + return Promise.resolve({ + protocolVersion: "2025-11-25", + capabilities: {}, + serverInfo: { name: "test", version: "1.0.0" }, + }); + } + return Promise.resolve({}); + }; + + const client = new McpClient({ connection: conn }); + await client.initialize(); + + const controller = new AbortController(); + const callPromise = client.callTool("test", {}, controller.signal); + + controller.abort(); + + await expect(callPromise).rejects.toThrow("Aborted"); + + // Clean up + resolveRequest?.({ + content: [{ type: "text", text: "too late" }], + }); + }); + + /** A connection whose initialize never responds (simulates a framing- + * incompatible server like chrome-devtools-mcp under Content-Length framing): + * the pending JSON-RPC request would hang forever without a timeout/abort. */ + function makeHangingConnection(): Connection { + const never = new Promise<unknown>(() => {}); + return { + send: () => never, + notify: () => {}, + onNotification: () => {}, + close: () => {}, + pid: 1, + }; + } + + it("initialize raises McpTimeoutError when the server never responds", async () => { + const { McpTimeoutError } = await import("./timeout.js"); + const client = new McpClient({ connection: makeHangingConnection() }); + await expect(client.initialize(undefined, 20)).rejects.toBeInstanceOf(McpTimeoutError); + expect(client.getState()).toBe("error"); + }); + + it("initialize is abortable: an aborting signal rejects immediately", async () => { + const client = new McpClient({ connection: makeHangingConnection() }); + const controller = new AbortController(); + const p = client.initialize(controller.signal, 50_000); + controller.abort(); + await expect(p).rejects.toThrow("Aborted"); + expect(client.getState()).toBe("error"); + }); + + it("initialize rejects immediately when the signal is already aborted", async () => { + const client = new McpClient({ connection: makeHangingConnection() }); + const controller = new AbortController(); + controller.abort(); + await expect(client.initialize(controller.signal, 50_000)).rejects.toThrow("Aborted"); + }); + + it("listTools raises McpTimeoutError when the server never responds", async () => { + const { McpTimeoutError } = await import("./timeout.js"); + // Reach "connected" with a fast (auto-responding) connection, then swap in + // a hanging connection for the tools/list call. + const fast = makeMockConnection(); + const client = new McpClient({ connection: fast }); + await client.initialize(); + + const hanging = makeHangingConnection(); + // Swap the connection so tools/list hangs. + (client as unknown as { connection: Connection }).connection = hanging; + + await expect(client.listTools(undefined, 20)).rejects.toBeInstanceOf(McpTimeoutError); + }); + + it("listTools is abortable", async () => { + const fast = makeMockConnection(); + const client = new McpClient({ connection: fast }); + await client.initialize(); + + const hanging = makeHangingConnection(); + (client as unknown as { connection: Connection }).connection = hanging; + + const controller = new AbortController(); + const p = client.listTools(controller.signal, 50_000); + controller.abort(); + await expect(p).rejects.toThrow("Aborted"); + }); }); diff --git a/packages/mcp/src/client.ts b/packages/mcp/src/client.ts index 17463d9..6a69c00 100644 --- a/packages/mcp/src/client.ts +++ b/packages/mcp/src/client.ts @@ -3,121 +3,153 @@ * * Manages a single MCP server connection: initialize handshake, * tool discovery, tool invocation, and list_changed notifications. + * + * Every awaited handshake/list operation is bounded by `withTimeout` (a default + * timeout) and an optional `AbortSignal`, so a misbehaving or framing- + * incompatible server can never hang the caller (the per-turn tools filter) + * indefinitely. `callTool` was already abort-aware; `initialize`/`listTools` + * now are too. */ +import { MCP_DEFAULT_TIMEOUT_MS, withTimeout } from "./timeout.js"; import type { Connection } from "./transport.js"; import type { - McpCallResult, - McpInitializeResult, - McpListToolsResult, - McpServerCapabilities, - McpToolInfo, + McpCallResult, + McpInitializeResult, + McpListToolsResult, + McpServerCapabilities, + McpToolInfo, } from "./types.js"; export type McpClientState = "disconnected" | "connecting" | "connected" | "error"; export interface McpClientDeps { - readonly connection: Connection; + readonly connection: Connection; } export class McpClient { - private state: McpClientState = "disconnected"; - private capabilities: McpServerCapabilities = {}; - private tools: readonly McpToolInfo[] = []; - private connection: Connection; - private toolsChangedHandler: (() => void) | null = null; - - constructor(deps: McpClientDeps) { - this.connection = deps.connection; - } - - getState(): McpClientState { - return this.state; - } - - getCapabilities(): McpServerCapabilities { - return this.capabilities; - } - - getTools(): readonly McpToolInfo[] { - return this.tools; - } - - onToolsChanged(handler: () => void): void { - this.toolsChangedHandler = handler; - } - - async initialize(): Promise<McpInitializeResult> { - this.state = "connecting"; - try { - const result = (await this.connection.send("initialize", { - protocolVersion: "2025-11-25", - capabilities: {}, - clientInfo: { name: "dispatch", version: "0.0.0" }, - })) as McpInitializeResult; - - this.capabilities = result.capabilities; - this.connection.notify("notifications/initialized", {}); - - this.connection.onNotification("notifications/tools/list_changed", () => { - if (this.toolsChangedHandler) { - this.toolsChangedHandler(); - } - }); - - this.state = "connected"; - return result; - } catch (err: unknown) { - this.state = "error"; - throw err; - } - } - - async listTools(): Promise<readonly McpToolInfo[]> { - if (this.state !== "connected") { - throw new Error("Client not connected"); - } - const result = (await this.connection.send("tools/list")) as McpListToolsResult; - this.tools = result.tools; - return this.tools; - } - - async callTool(name: string, args: unknown, signal?: AbortSignal): Promise<McpCallResult> { - if (this.state !== "connected") { - throw new Error("Client not connected"); - } - - if (signal?.aborted) { - throw new Error("Aborted"); - } - - const resultPromise = this.connection.send("tools/call", { - name, - arguments: args, - }) as Promise<McpCallResult>; - - if (!signal) { - return resultPromise; - } - - return new Promise<McpCallResult>((resolve, reject) => { - const onAbort = () => reject(new Error("Aborted")); - signal.addEventListener("abort", onAbort, { once: true }); - resultPromise.then( - (result) => { - signal.removeEventListener("abort", onAbort); - resolve(result); - }, - (err) => { - signal.removeEventListener("abort", onAbort); - reject(err); - }, - ); - }); - } - - close(): void { - this.state = "disconnected"; - this.connection.close(); - } + private state: McpClientState = "disconnected"; + private capabilities: McpServerCapabilities = {}; + private tools: readonly McpToolInfo[] = []; + private connection: Connection; + private toolsChangedHandler: (() => void) | null = null; + + constructor(deps: McpClientDeps) { + this.connection = deps.connection; + } + + getState(): McpClientState { + return this.state; + } + + getCapabilities(): McpServerCapabilities { + return this.capabilities; + } + + getTools(): readonly McpToolInfo[] { + return this.tools; + } + + onToolsChanged(handler: () => void): void { + this.toolsChangedHandler = handler; + } + + /** + * Perform the MCP `initialize` handshake. Bounded by `timeoutMs` (default + * {@link MCP_DEFAULT_TIMEOUT_MS}) and the optional `signal` (the turn's abort + * signal) so a server that never responds cannot hang the caller forever. + */ + async initialize( + signal?: AbortSignal, + timeoutMs: number = MCP_DEFAULT_TIMEOUT_MS, + ): Promise<McpInitializeResult> { + this.state = "connecting"; + try { + const result = (await withTimeout( + this.connection.send("initialize", { + protocolVersion: "2025-11-25", + capabilities: {}, + clientInfo: { name: "dispatch", version: "0.0.0" }, + }), + "initialize", + timeoutMs, + signal, + )) as McpInitializeResult; + + this.capabilities = result.capabilities; + this.connection.notify("notifications/initialized", {}); + + this.connection.onNotification("notifications/tools/list_changed", () => { + if (this.toolsChangedHandler) { + this.toolsChangedHandler(); + } + }); + + this.state = "connected"; + return result; + } catch (err: unknown) { + this.state = "error"; + throw err; + } + } + + /** + * List the server's tools. Bounded by `timeoutMs` (default + * {@link MCP_DEFAULT_TIMEOUT_MS}) and the optional `signal`. + */ + async listTools( + signal?: AbortSignal, + timeoutMs: number = MCP_DEFAULT_TIMEOUT_MS, + ): Promise<readonly McpToolInfo[]> { + if (this.state !== "connected") { + throw new Error("Client not connected"); + } + const result = (await withTimeout( + this.connection.send("tools/list"), + "tools/list", + timeoutMs, + signal, + )) as McpListToolsResult; + this.tools = result.tools; + return this.tools; + } + + async callTool(name: string, args: unknown, signal?: AbortSignal): Promise<McpCallResult> { + if (this.state !== "connected") { + throw new Error("Client not connected"); + } + + if (signal?.aborted) { + throw new Error("Aborted"); + } + + const resultPromise = this.connection.send("tools/call", { + name, + arguments: args, + }) as Promise<McpCallResult>; + + if (!signal) { + return resultPromise; + } + + return new Promise<McpCallResult>((resolve, reject) => { + const onAbort = () => reject(new Error("Aborted")); + signal.addEventListener("abort", onAbort, { once: true }); + resultPromise.then( + (result) => { + signal.removeEventListener("abort", onAbort); + resolve(result); + }, + (err) => { + signal.removeEventListener("abort", onAbort); + reject(err); + }, + ); + }); + } + + close(): void { + this.state = "disconnected"; + this.connection.close(); + } } diff --git a/packages/mcp/src/config.test.ts b/packages/mcp/src/config.test.ts index 38c9b50..c83061f 100644 --- a/packages/mcp/src/config.test.ts +++ b/packages/mcp/src/config.test.ts @@ -2,113 +2,113 @@ import { describe, expect, it } from "vitest"; import { resolveServers } from "./config.js"; describe("resolveServers", () => { - it("resolves from .dispatch/mcp.json", () => { - const dispatchConfig = JSON.stringify({ - servers: { - freecad: { command: "uvx", args: ["freecad-mcp"], env: { KEY: "val" } }, - }, - }); - - const result = resolveServers({ dispatchMcpJson: dispatchConfig, opencodeJson: null }); - - expect(result.servers.length).toBe(1); - expect(result.servers[0].id).toBe("freecad"); - expect(result.servers[0].command).toEqual(["uvx", "freecad-mcp"]); - expect(result.servers[0].env).toEqual({ KEY: "val" }); - expect(result.servers[0].configSource).toBe(".dispatch/mcp.json"); - expect(result.shadowed).toBe(false); - }); - - it("falls back to opencode.json mcp key", () => { - const opencodeConfig = JSON.stringify({ - mcp: { - chrome: { command: "npx", args: ["chrome-devtools-mcp@latest"] }, - }, - }); - - const result = resolveServers({ dispatchMcpJson: null, opencodeJson: opencodeConfig }); - - expect(result.servers.length).toBe(1); - expect(result.servers[0].id).toBe("chrome"); - expect(result.servers[0].command).toEqual(["npx", "chrome-devtools-mcp@latest"]); - expect(result.servers[0].configSource).toBe("opencode.json"); - expect(result.shadowed).toBe(false); - }); - - it("shadow warning when both present", () => { - const dispatchConfig = JSON.stringify({ - servers: { - freecad: { command: "uvx", args: ["freecad-mcp"] }, - }, - }); - const opencodeConfig = JSON.stringify({ - mcp: { - chrome: { command: "npx", args: ["chrome-devtools-mcp@latest"] }, - }, - }); - - const result = resolveServers({ - dispatchMcpJson: dispatchConfig, - opencodeJson: opencodeConfig, - }); - - expect(result.servers.length).toBe(1); - expect(result.servers[0].id).toBe("freecad"); - expect(result.shadowed).toBe(true); - }); - - it("empty when neither present", () => { - const result = resolveServers({ dispatchMcpJson: null, opencodeJson: null }); - - expect(result.servers.length).toBe(0); - expect(result.shadowed).toBe(false); - }); - - it("empty when dispatch has no servers key", () => { - const result = resolveServers({ - dispatchMcpJson: JSON.stringify({}), - opencodeJson: null, - }); - - expect(result.servers.length).toBe(0); - expect(result.shadowed).toBe(false); - }); - - it("handles malformed JSON gracefully", () => { - const result = resolveServers({ - dispatchMcpJson: "not valid json", - opencodeJson: "{ also bad", - }); - - expect(result.servers.length).toBe(0); - expect(result.shadowed).toBe(false); - }); - - it("server without args", () => { - const dispatchConfig = JSON.stringify({ - servers: { - simple: { command: "my-server" }, - }, - }); - - const result = resolveServers({ dispatchMcpJson: dispatchConfig, opencodeJson: null }); - - expect(result.servers.length).toBe(1); - expect(result.servers[0].command).toEqual(["my-server"]); - }); - - it("multiple servers from dispatch", () => { - const dispatchConfig = JSON.stringify({ - servers: { - a: { command: "server-a" }, - b: { command: "server-b", args: ["--port", "3000"] }, - }, - }); - - const result = resolveServers({ dispatchMcpJson: dispatchConfig, opencodeJson: null }); - - expect(result.servers.length).toBe(2); - const ids = result.servers.map((s) => s.id).sort(); - expect(ids).toEqual(["a", "b"]); - }); + it("resolves from .dispatch/mcp.json", () => { + const dispatchConfig = JSON.stringify({ + servers: { + freecad: { command: "uvx", args: ["freecad-mcp"], env: { KEY: "val" } }, + }, + }); + + const result = resolveServers({ dispatchMcpJson: dispatchConfig, opencodeJson: null }); + + expect(result.servers.length).toBe(1); + expect(result.servers[0].id).toBe("freecad"); + expect(result.servers[0].command).toEqual(["uvx", "freecad-mcp"]); + expect(result.servers[0].env).toEqual({ KEY: "val" }); + expect(result.servers[0].configSource).toBe(".dispatch/mcp.json"); + expect(result.shadowed).toBe(false); + }); + + it("falls back to opencode.json mcp key", () => { + const opencodeConfig = JSON.stringify({ + mcp: { + chrome: { command: "npx", args: ["chrome-devtools-mcp@latest"] }, + }, + }); + + const result = resolveServers({ dispatchMcpJson: null, opencodeJson: opencodeConfig }); + + expect(result.servers.length).toBe(1); + expect(result.servers[0].id).toBe("chrome"); + expect(result.servers[0].command).toEqual(["npx", "chrome-devtools-mcp@latest"]); + expect(result.servers[0].configSource).toBe("opencode.json"); + expect(result.shadowed).toBe(false); + }); + + it("shadow warning when both present", () => { + const dispatchConfig = JSON.stringify({ + servers: { + freecad: { command: "uvx", args: ["freecad-mcp"] }, + }, + }); + const opencodeConfig = JSON.stringify({ + mcp: { + chrome: { command: "npx", args: ["chrome-devtools-mcp@latest"] }, + }, + }); + + const result = resolveServers({ + dispatchMcpJson: dispatchConfig, + opencodeJson: opencodeConfig, + }); + + expect(result.servers.length).toBe(1); + expect(result.servers[0].id).toBe("freecad"); + expect(result.shadowed).toBe(true); + }); + + it("empty when neither present", () => { + const result = resolveServers({ dispatchMcpJson: null, opencodeJson: null }); + + expect(result.servers.length).toBe(0); + expect(result.shadowed).toBe(false); + }); + + it("empty when dispatch has no servers key", () => { + const result = resolveServers({ + dispatchMcpJson: JSON.stringify({}), + opencodeJson: null, + }); + + expect(result.servers.length).toBe(0); + expect(result.shadowed).toBe(false); + }); + + it("handles malformed JSON gracefully", () => { + const result = resolveServers({ + dispatchMcpJson: "not valid json", + opencodeJson: "{ also bad", + }); + + expect(result.servers.length).toBe(0); + expect(result.shadowed).toBe(false); + }); + + it("server without args", () => { + const dispatchConfig = JSON.stringify({ + servers: { + simple: { command: "my-server" }, + }, + }); + + const result = resolveServers({ dispatchMcpJson: dispatchConfig, opencodeJson: null }); + + expect(result.servers.length).toBe(1); + expect(result.servers[0].command).toEqual(["my-server"]); + }); + + it("multiple servers from dispatch", () => { + const dispatchConfig = JSON.stringify({ + servers: { + a: { command: "server-a" }, + b: { command: "server-b", args: ["--port", "3000"] }, + }, + }); + + const result = resolveServers({ dispatchMcpJson: dispatchConfig, opencodeJson: null }); + + expect(result.servers.length).toBe(2); + const ids = result.servers.map((s) => s.id).sort(); + expect(ids).toEqual(["a", "b"]); + }); }); diff --git a/packages/mcp/src/config.ts b/packages/mcp/src/config.ts index 20ed749..599f414 100644 --- a/packages/mcp/src/config.ts +++ b/packages/mcp/src/config.ts @@ -11,79 +11,79 @@ import type { McpServerConfig, ResolvedMcpServer, ResolveResult } from "./types.js"; export interface ResolveServersDeps { - readonly dispatchMcpJson: string | null; - readonly opencodeJson: string | null; + readonly dispatchMcpJson: string | null; + readonly opencodeJson: string | null; } export interface DispatchMcpConfig { - readonly servers?: Readonly<Record<string, McpServerConfig>>; + readonly servers?: Readonly<Record<string, McpServerConfig>>; } export interface OpencodeJsonConfig { - readonly mcp?: Readonly<Record<string, McpServerConfig>>; + readonly mcp?: Readonly<Record<string, McpServerConfig>>; } export function resolveServers(deps: ResolveServersDeps): ResolveResult { - const result = new Map<string, ResolvedMcpServer>(); + const result = new Map<string, ResolvedMcpServer>(); - // Parse opencode.json once — used both as the fallback source and to detect - // whether a present `.dispatch/mcp.json` silently shadows its `mcp` key. - let opencodeConfig: OpencodeJsonConfig | null = null; - if (deps.opencodeJson) { - try { - opencodeConfig = JSON.parse(deps.opencodeJson) as OpencodeJsonConfig; - } catch { - // ignore parse errors - } - } - const opencodeHasMcp = !!opencodeConfig?.mcp && Object.keys(opencodeConfig.mcp).length > 0; + // Parse opencode.json once — used both as the fallback source and to detect + // whether a present `.dispatch/mcp.json` silently shadows its `mcp` key. + let opencodeConfig: OpencodeJsonConfig | null = null; + if (deps.opencodeJson) { + try { + opencodeConfig = JSON.parse(deps.opencodeJson) as OpencodeJsonConfig; + } catch { + // ignore parse errors + } + } + const opencodeHasMcp = !!opencodeConfig?.mcp && Object.keys(opencodeConfig.mcp).length > 0; - // 1. cwd/.dispatch/mcp.json (highest precedence) - let dispatchHadServers = false; - if (deps.dispatchMcpJson) { - try { - const config = JSON.parse(deps.dispatchMcpJson) as DispatchMcpConfig; - if (config.servers) { - for (const [key, server] of Object.entries(config.servers)) { - const resolved = resolveServer(key, server, ".dispatch/mcp.json"); - result.set(resolved.id, resolved); - } - dispatchHadServers = result.size > 0; - } - } catch { - // ignore parse errors - } - } + // 1. cwd/.dispatch/mcp.json (highest precedence) + let dispatchHadServers = false; + if (deps.dispatchMcpJson) { + try { + const config = JSON.parse(deps.dispatchMcpJson) as DispatchMcpConfig; + if (config.servers) { + for (const [key, server] of Object.entries(config.servers)) { + const resolved = resolveServer(key, server, ".dispatch/mcp.json"); + result.set(resolved.id, resolved); + } + dispatchHadServers = result.size > 0; + } + } catch { + // ignore parse errors + } + } - // 2. fallback cwd/opencode.json mcp key (only when dispatch yielded nothing) - if (result.size === 0 && opencodeConfig?.mcp) { - for (const [key, server] of Object.entries(opencodeConfig.mcp)) { - const resolved = resolveServer(key, server, "opencode.json"); - result.set(resolved.id, resolved); - } - } + // 2. fallback cwd/opencode.json mcp key (only when dispatch yielded nothing) + if (result.size === 0 && opencodeConfig?.mcp) { + for (const [key, server] of Object.entries(opencodeConfig.mcp)) { + const resolved = resolveServer(key, server, "opencode.json"); + result.set(resolved.id, resolved); + } + } - // No built-in servers — MCP has no built-in registry. + // No built-in servers — MCP has no built-in registry. - // `.dispatch/mcp.json` silently shadows `opencode.json`'s mcp key when both - // declare servers — the opencode entry is skipped with no warning otherwise. - const shadowed = dispatchHadServers && opencodeHasMcp; - return { servers: [...result.values()], shadowed }; + // `.dispatch/mcp.json` silently shadows `opencode.json`'s mcp key when both + // declare servers — the opencode entry is skipped with no warning otherwise. + const shadowed = dispatchHadServers && opencodeHasMcp; + return { servers: [...result.values()], shadowed }; } function resolveServer( - key: string, - config: McpServerConfig, - configSource: ".dispatch/mcp.json" | "opencode.json", + key: string, + config: McpServerConfig, + configSource: ".dispatch/mcp.json" | "opencode.json", ): ResolvedMcpServer { - const command = [config.command, ...(config.args ?? [])]; - const result: ResolvedMcpServer = { - id: key, - command, - configSource, - }; - if (config.env) { - (result as { env?: Readonly<Record<string, string>> }).env = config.env; - } - return result; + const command = [config.command, ...(config.args ?? [])]; + const result: ResolvedMcpServer = { + id: key, + command, + configSource, + }; + if (config.env) { + (result as { env?: Readonly<Record<string, string>> }).env = config.env; + } + return result; } diff --git a/packages/mcp/src/extension.test.ts b/packages/mcp/src/extension.test.ts index 75515fb..9e029d2 100644 --- a/packages/mcp/src/extension.test.ts +++ b/packages/mcp/src/extension.test.ts @@ -11,67 +11,67 @@ import type { McpToolInfo } from "./types.js"; // --------------------------------------------------------------------------- const stubTool = (name: string): ToolContract => ({ - name, - description: "", - parameters: { type: "object" }, - execute: async () => ({ content: "" }), + name, + description: "", + parameters: { type: "object" }, + execute: async () => ({ content: "" }), }); describe("filterMcpTools (pure)", () => { - it("keeps non-MCP tools and connected-server tools, removes disconnected-server tools", () => { - const toolToServer = new Map<string, string>([ - ["a__x", "a"], - ["b__y", "b"], - ]); - const connected = new Set<string>(["a"]); - - const result = filterMcpTools( - { - tools: [stubTool("a__x"), stubTool("b__y"), stubTool("other")], - cwd: "/p", - conversationId: "c", - }, - toolToServer, - connected, - ); - - expect(result.tools.map((t) => t.name).sort()).toEqual(["a__x", "other"]); - expect(result.cwd).toBe("/p"); - expect(result.conversationId).toBe("c"); - }); - - it("removes all MCP tools when no server is connected", () => { - const result = filterMcpTools( - { tools: [stubTool("a__x")], conversationId: "c" }, - new Map<string, string>([["a__x", "a"]]), - new Set<string>(), - ); - expect(result.tools).toHaveLength(0); - expect(result.conversationId).toBe("c"); - expect(result.cwd).toBeUndefined(); - expect(result.computerId).toBeUndefined(); - }); - - it("preserves computerId when set (mirrors cwd/conversationId preservation)", () => { - const toolToServer = new Map<string, string>([["a__x", "a"]]); - const connected = new Set<string>(["a"]); - - const result = filterMcpTools( - { - tools: [stubTool("a__x"), stubTool("other")], - cwd: "/p", - computerId: "ssh-host", - conversationId: "c", - }, - toolToServer, - connected, - ); - - expect(result.tools.map((t) => t.name).sort()).toEqual(["a__x", "other"]); - expect(result.computerId).toBe("ssh-host"); - expect(result.cwd).toBe("/p"); - expect(result.conversationId).toBe("c"); - }); + it("keeps non-MCP tools and connected-server tools, removes disconnected-server tools", () => { + const toolToServer = new Map<string, string>([ + ["a__x", "a"], + ["b__y", "b"], + ]); + const connected = new Set<string>(["a"]); + + const result = filterMcpTools( + { + tools: [stubTool("a__x"), stubTool("b__y"), stubTool("other")], + cwd: "/p", + conversationId: "c", + }, + toolToServer, + connected, + ); + + expect(result.tools.map((t) => t.name).sort()).toEqual(["a__x", "other"]); + expect(result.cwd).toBe("/p"); + expect(result.conversationId).toBe("c"); + }); + + it("removes all MCP tools when no server is connected", () => { + const result = filterMcpTools( + { tools: [stubTool("a__x")], conversationId: "c" }, + new Map<string, string>([["a__x", "a"]]), + new Set<string>(), + ); + expect(result.tools).toHaveLength(0); + expect(result.conversationId).toBe("c"); + expect(result.cwd).toBeUndefined(); + expect(result.computerId).toBeUndefined(); + }); + + it("preserves computerId when set (mirrors cwd/conversationId preservation)", () => { + const toolToServer = new Map<string, string>([["a__x", "a"]]); + const connected = new Set<string>(["a"]); + + const result = filterMcpTools( + { + tools: [stubTool("a__x"), stubTool("other")], + cwd: "/p", + computerId: "ssh-host", + conversationId: "c", + }, + toolToServer, + connected, + ); + + expect(result.tools.map((t) => t.name).sort()).toEqual(["a__x", "other"]); + expect(result.computerId).toBe("ssh-host"); + expect(result.cwd).toBe("/p"); + expect(result.conversationId).toBe("c"); + }); }); // --------------------------------------------------------------------------- @@ -79,9 +79,12 @@ describe("filterMcpTools (pure)", () => { // --------------------------------------------------------------------------- interface FakeServer { - tools: McpToolInfo[]; - failInitialize: boolean; - emitListChanged: () => void; + tools: McpToolInfo[]; + failInitialize: boolean; + /** When true, the spawn never responds to `initialize` (a hanging / + * framing-incompatible server) — used to exercise timeout/abort paths. */ + hangInitialize: boolean; + emitListChanged: () => void; } /** @@ -90,87 +93,91 @@ interface FakeServer { * transport → framing → rpc → client → manager end to end. */ function makeFakeSpawn(server: FakeServer): SpawnProcess { - const decoder = new FrameDecoder(); - let dataListeners: Array<(data: Uint8Array) => void> = []; - - const emit = (frame: Uint8Array) => { - for (const cb of dataListeners) cb(frame); - }; - - const spawn: SpawnProcess = (_command, _opts) => { - // Each spawn is a fresh process; reset listeners so a reconnect (after - // shutdown) doesn't feed closed rpc instances. - dataListeners = []; - const process: SpawnedProcess = { - stdin: { - write: (bytes: Uint8Array) => { - for (const msg of decoder.decode(bytes)) { - const parsed = JSON.parse(msg) as { - id?: number; - method?: string; - params?: unknown; - }; - const id = parsed.id ?? 0; - const method = parsed.method; - if (method === "initialize") { - if (server.failInitialize) { - emit( - encode( - JSON.stringify({ - jsonrpc: "2.0", - id, - error: { code: -32603, message: "initialize failed" }, - }), - ), - ); - } else { - emit( - encode( - JSON.stringify({ - jsonrpc: "2.0", - id, - result: { - protocolVersion: "2025-11-25", - capabilities: { tools: { listChanged: true } }, - serverInfo: { name: "fake", version: "0.0.0" }, - }, - }), - ), - ); - } - } else if (method === "tools/list") { - emit(encode(JSON.stringify({ jsonrpc: "2.0", id, result: { tools: server.tools } }))); - } else if (method === "tools/call") { - emit( - encode( - JSON.stringify({ - jsonrpc: "2.0", - id, - result: { content: [{ type: "text", text: "ok" }], isError: false }, - }), - ), - ); - } - // notifications (notifications/initialized): no response. - } - }, - }, - stdout: { - on: (event: string, cb: (data: Uint8Array) => void) => { - if (event === "data") dataListeners.push(cb); - }, - }, - pid: 7000, - kill: () => {}, - }; - return process; - }; - - server.emitListChanged = () => { - emit(encode(JSON.stringify({ jsonrpc: "2.0", method: "notifications/tools/list_changed" }))); - }; - - return spawn; + const decoder = new FrameDecoder(); + let dataListeners: Array<(data: Uint8Array) => void> = []; + + const emit = (frame: Uint8Array) => { + for (const cb of dataListeners) cb(frame); + }; + + const spawn: SpawnProcess = (_command, _opts) => { + // Each spawn is a fresh process; reset listeners so a reconnect (after + // shutdown) doesn't feed closed rpc instances. + dataListeners = []; + const process: SpawnedProcess = { + stdin: { + write: (bytes: Uint8Array) => { + for (const msg of decoder.decode(bytes)) { + const parsed = JSON.parse(msg) as { + id?: number; + method?: string; + params?: unknown; + }; + const id = parsed.id ?? 0; + const method = parsed.method; + if (method === "initialize") { + if (server.hangInitialize) { + // Never respond — simulates a framing-incompatible server + // (e.g. chrome-devtools-mcp under the old Content-Length + // framing). The connect must be bounded by timeout/abort. + } else if (server.failInitialize) { + emit( + encode( + JSON.stringify({ + jsonrpc: "2.0", + id, + error: { code: -32603, message: "initialize failed" }, + }), + ), + ); + } else { + emit( + encode( + JSON.stringify({ + jsonrpc: "2.0", + id, + result: { + protocolVersion: "2025-11-25", + capabilities: { tools: { listChanged: true } }, + serverInfo: { name: "fake", version: "0.0.0" }, + }, + }), + ), + ); + } + } else if (method === "tools/list") { + emit(encode(JSON.stringify({ jsonrpc: "2.0", id, result: { tools: server.tools } }))); + } else if (method === "tools/call") { + emit( + encode( + JSON.stringify({ + jsonrpc: "2.0", + id, + result: { content: [{ type: "text", text: "ok" }], isError: false }, + }), + ), + ); + } + // notifications (notifications/initialized): no response. + } + }, + }, + stdout: { + on: (event: string, cb: (data: Uint8Array) => void) => { + if (event === "data") dataListeners.push(cb); + }, + }, + pid: 7000, + kill: () => {}, + }; + return process; + }; + + server.emitListChanged = () => { + emit(encode(JSON.stringify({ jsonrpc: "2.0", method: "notifications/tools/list_changed" }))); + }; + + return spawn; } // --------------------------------------------------------------------------- @@ -178,51 +185,51 @@ function makeFakeSpawn(server: FakeServer): SpawnProcess { // --------------------------------------------------------------------------- function makeFakeHost(): { - host: HostAPI; - tools: Map<string, ToolContract>; - getFilter: () => ((a: ToolAssembly) => Promise<ToolAssembly>) | null; - getService: () => unknown; + host: HostAPI; + tools: Map<string, ToolContract>; + getFilter: () => ((a: ToolAssembly) => Promise<ToolAssembly>) | null; + getService: () => unknown; } { - const tools = new Map<string, ToolContract>(); - let filterFn: ((a: ToolAssembly) => Promise<ToolAssembly>) | null = null; - let service: unknown = null; - - const noopSpan = { - id: "s", - log: {} as Logger, - setAttributes: () => {}, - addLink: () => {}, - child: () => noopSpan, - end: () => {}, - }; - const noopLogger: Logger = { - info: () => {}, - warn: () => {}, - error: () => {}, - debug: () => {}, - child: () => noopLogger, - span: () => noopSpan, - }; - - const host = { - defineTool: (t: ToolContract) => { - tools.set(t.name, t); - }, - addFilter: (_hook: typeof toolsFilter, fn: (a: ToolAssembly) => Promise<ToolAssembly>) => { - filterFn = fn; - return () => { - filterFn = null; - }; - }, - provideService: (_handle: unknown, impl: unknown) => { - service = impl; - }, - getService: () => service, - getTools: () => tools, - logger: noopLogger, - } as unknown as HostAPI; - - return { host, tools, getFilter: () => filterFn, getService: () => service }; + const tools = new Map<string, ToolContract>(); + let filterFn: ((a: ToolAssembly) => Promise<ToolAssembly>) | null = null; + let service: unknown = null; + + const noopSpan = { + id: "s", + log: {} as Logger, + setAttributes: () => {}, + addLink: () => {}, + child: () => noopSpan, + end: () => {}, + }; + const noopLogger: Logger = { + info: () => {}, + warn: () => {}, + error: () => {}, + debug: () => {}, + child: () => noopLogger, + span: () => noopSpan, + }; + + const host = { + defineTool: (t: ToolContract) => { + tools.set(t.name, t); + }, + addFilter: (_hook: typeof toolsFilter, fn: (a: ToolAssembly) => Promise<ToolAssembly>) => { + filterFn = fn; + return () => { + filterFn = null; + }; + }, + provideService: (_handle: unknown, impl: unknown) => { + service = impl; + }, + getService: () => service, + getTools: () => tools, + logger: noopLogger, + } as unknown as HostAPI; + + return { host, tools, getFilter: () => filterFn, getService: () => service }; } // --------------------------------------------------------------------------- @@ -232,132 +239,228 @@ function makeFakeHost(): { const dispatchConfig = (servers: Record<string, unknown>): string => JSON.stringify({ servers }); const tool = (name: string, description = name): McpToolInfo => ({ - name, - description, - inputSchema: { type: "object" }, + name, + description, + inputSchema: { type: "object" }, }); const assembly = (tools: ToolContract[], cwd = "/proj"): ToolAssembly => ({ - tools, - cwd, - conversationId: "conv-1", + tools, + cwd, + conversationId: "conv-1", }); const flush = () => new Promise((r) => setTimeout(r, 0)); function makeServer(initialTools: McpToolInfo[]): FakeServer { - return { tools: [...initialTools], failInitialize: false, emitListChanged: () => {} }; + return { + tools: [...initialTools], + failInitialize: false, + hangInitialize: false, + emitListChanged: () => {}, + }; } function makeExt(server: FakeServer, configJson: string): Extension { - return makeMcpExtension({ - spawn: makeFakeSpawn(server), - readFile: async (path) => (path.endsWith(".dispatch/mcp.json") ? configJson : null), - getCwd: () => "/proj", - }); + return makeMcpExtension({ + spawn: makeFakeSpawn(server), + readFile: async (path) => (path.endsWith(".dispatch/mcp.json") ? configJson : null), + getCwd: () => "/proj", + }); } // --------------------------------------------------------------------------- // Lifecycle tests // --------------------------------------------------------------------------- describe("mcp extension lifecycle", () => { - /** Get the registered filter, throwing if activation did not register one. */ - function requireFilter(getFilter: () => ((a: ToolAssembly) => Promise<ToolAssembly>) | null) { - const filter = getFilter(); - if (!filter) throw new Error("toolsFilter was not registered"); - return filter; - } - - /** Look up a registered tool, throwing if absent. */ - function requireTool(tools: Map<string, ToolContract>, name: string) { - const t = tools.get(name); - if (!t) throw new Error(`tool ${name} not registered`); - return t; - } - - it("registers tools on connect", async () => { - const server = makeServer([tool("create_object", "Create an object")]); - const ext = makeExt(server, dispatchConfig({ freecad: { command: "fake" } })); - const { host, tools, getFilter } = makeFakeHost(); - - ext.activate(host); - const filter = requireFilter(getFilter); - - // Running the filter triggers lazy connect + register. - await filter(assembly([])); - - expect(tools.has("freecad__create_object")).toBe(true); - const t = requireTool(tools, "freecad__create_object"); - expect(t.description).toBe("[freecad] Create an object"); - expect(t.concurrencySafe).toBe(false); - ext.deactivate?.(); - }); - - it("toolsFilter keeps connected-server tools and removes disconnected-server tools", async () => { - const server = makeServer([tool("create_object")]); - const ext = makeExt(server, dispatchConfig({ freecad: { command: "fake" } })); - const { host, tools, getFilter } = makeFakeHost(); - ext.activate(host); - const filter = requireFilter(getFilter); - - // Connect + register the tool. - await filter(assembly([])); - const registered = requireTool(tools, "freecad__create_object"); - - // Connected server → tool passes through the filter. - const kept = await filter(assembly([registered])); - expect(kept.tools.map((t) => t.name)).toContain("freecad__create_object"); - - // Disconnect: deactivate shuts down the client (clearing it), then make - // the server fail to reconnect. toolToServer still maps the tool, so the - // filter drops it because the server is no longer connected. - ext.deactivate?.(); - server.failInitialize = true; - - const removed = await filter(assembly([registered])); - expect(removed.tools.map((t) => t.name)).not.toContain("freecad__create_object"); - }); - - it("re-registers tools on list_changed", async () => { - const server = makeServer([tool("first_tool")]); - const ext = makeExt(server, dispatchConfig({ freecad: { command: "fake" } })); - const { host, tools, getFilter } = makeFakeHost(); - ext.activate(host); - const filter = requireFilter(getFilter); - - await filter(assembly([])); - expect(tools.has("freecad__first_tool")).toBe(true); - - // Server changes its tool set, then announces list_changed. - server.tools = [tool("first_tool"), tool("second_tool", "The second")]; - server.emitListChanged(); - - // Let the async onToolsChanged handler (re-list + re-register) flush. - await flush(); - - expect(tools.has("freecad__second_tool")).toBe(true); - expect(requireTool(tools, "freecad__second_tool").description).toBe("[freecad] The second"); - ext.deactivate?.(); - }); - - it("deactivate shuts down all clients", async () => { - const server = makeServer([tool("create_object")]); - const ext = makeExt(server, dispatchConfig({ freecad: { command: "fake" } })); - const { host, getFilter, getService } = makeFakeHost(); - ext.activate(host); - const filter = requireFilter(getFilter); - - await filter(assembly([])); - - const service = getService() as { - status: (cwd: string) => Promise<readonly { state: string }[]>; - }; - const before = await service.status("/proj"); - expect(before[0].state).toBe("connected"); - - ext.deactivate?.(); - - const after = await service.status("/proj"); - expect(after[0].state).toBe("disconnected"); - }); + /** Get the registered filter, throwing if activation did not register one. */ + function requireFilter(getFilter: () => ((a: ToolAssembly) => Promise<ToolAssembly>) | null) { + const filter = getFilter(); + if (!filter) throw new Error("toolsFilter was not registered"); + return filter; + } + + /** Look up a registered tool, throwing if absent. */ + function requireTool(tools: Map<string, ToolContract>, name: string) { + const t = tools.get(name); + if (!t) throw new Error(`tool ${name} not registered`); + return t; + } + + it("registers tools on connect", async () => { + const server = makeServer([tool("create_object", "Create an object")]); + const ext = makeExt(server, dispatchConfig({ freecad: { command: "fake" } })); + const { host, tools, getFilter } = makeFakeHost(); + + ext.activate(host); + const filter = requireFilter(getFilter); + + // Running the filter triggers lazy connect + register. + await filter(assembly([])); + + expect(tools.has("freecad__create_object")).toBe(true); + const t = requireTool(tools, "freecad__create_object"); + expect(t.description).toBe("[freecad] Create an object"); + expect(t.concurrencySafe).toBe(false); + ext.deactivate?.(); + }); + + it("toolsFilter keeps connected-server tools and removes disconnected-server tools", async () => { + const server = makeServer([tool("create_object")]); + const ext = makeExt(server, dispatchConfig({ freecad: { command: "fake" } })); + const { host, tools, getFilter } = makeFakeHost(); + ext.activate(host); + const filter = requireFilter(getFilter); + + // Connect + register the tool. + await filter(assembly([])); + const registered = requireTool(tools, "freecad__create_object"); + + // Connected server → tool passes through the filter. + const kept = await filter(assembly([registered])); + expect(kept.tools.map((t) => t.name)).toContain("freecad__create_object"); + + // Disconnect: deactivate shuts down the client (clearing it), then make + // the server fail to reconnect. toolToServer still maps the tool, so the + // filter drops it because the server is no longer connected. + ext.deactivate?.(); + server.failInitialize = true; + + const removed = await filter(assembly([registered])); + expect(removed.tools.map((t) => t.name)).not.toContain("freecad__create_object"); + }); + + it("re-registers tools on list_changed", async () => { + const server = makeServer([tool("first_tool")]); + const ext = makeExt(server, dispatchConfig({ freecad: { command: "fake" } })); + const { host, tools, getFilter } = makeFakeHost(); + ext.activate(host); + const filter = requireFilter(getFilter); + + await filter(assembly([])); + expect(tools.has("freecad__first_tool")).toBe(true); + + // Server changes its tool set, then announces list_changed. + server.tools = [tool("first_tool"), tool("second_tool", "The second")]; + server.emitListChanged(); + + // Let the async onToolsChanged handler (re-list + re-register) flush. + await flush(); + + expect(tools.has("freecad__second_tool")).toBe(true); + expect(requireTool(tools, "freecad__second_tool").description).toBe("[freecad] The second"); + ext.deactivate?.(); + }); + + it("deactivate shuts down all clients", async () => { + const server = makeServer([tool("create_object")]); + const ext = makeExt(server, dispatchConfig({ freecad: { command: "fake" } })); + const { host, getFilter, getService } = makeFakeHost(); + ext.activate(host); + const filter = requireFilter(getFilter); + + await filter(assembly([])); + + const service = getService() as { + status: (cwd: string) => Promise<readonly { state: string }[]>; + }; + const before = await service.status("/proj"); + expect(before[0].state).toBe("connected"); + + ext.deactivate?.(); + + const after = await service.status("/proj"); + expect(after[0].state).toBe("disconnected"); + }); + + // ------------------------------------------------------------------------- + // Bug 2 + Bug 3: a misbehaving/hanging server must not hang a turn, and the + // turn's AbortSignal (assembly.signal) must interrupt a stuck connect. + // ------------------------------------------------------------------------- + + it("degrades gracefully (no MCP tools) when the turn's signal is already aborted", async () => { + const server = makeServer([tool("create_object")]); + const ext = makeExt(server, dispatchConfig({ freecad: { command: "fake" } })); + const { host, getFilter } = makeFakeHost(); + ext.activate(host); + const filter = requireFilter(getFilter); + + const controller = new AbortController(); + controller.abort(); + + const base = assembly([]); + // The filter must NOT hang on the (never-needed) connect: an aborted turn + // signal propagates to initialize, which rejects immediately. + const result = await filter({ + tools: base.tools, + cwd: base.cwd, + conversationId: base.conversationId, + signal: controller.signal, + }); + + expect(result.tools).toEqual([]); + ext.deactivate?.(); + }); + + it("the turn's signal aborts a hanging server connect (POST /stop interrupts)", async () => { + // hangInitialize: the spawn never responds to initialize (a framing- + // incompatible / misbehaving server). Without abort propagation this + // would hang the filter until MCP_CONNECT_TIMEOUT_MS; with propagation + // the abort breaks it immediately. + const server = makeServer([tool("create_object")]); + server.hangInitialize = true; + const ext = makeExt(server, dispatchConfig({ chrome: { command: "fake" } })); + const { host, getFilter } = makeFakeHost(); + ext.activate(host); + const filter = requireFilter(getFilter); + + const controller = new AbortController(); + const base = assembly([]); + const resultPromise = filter({ + tools: base.tools, + cwd: base.cwd, + conversationId: base.conversationId, + signal: controller.signal, + }); + + // Let the filter progress into the hanging initialize (withTimeout has + // its abort listener armed), THEN abort — a true mid-flight cancel + // simulating POST /conversations/:id/stop. Without signal propagation + // this would hang ~30s (the connect backstop) and time out the test. + await flush(); + controller.abort(); + + const result = await resultPromise; + // Degraded: no MCP tools surfaced, and the filter resolved (did not hang). + expect(result.tools).toEqual([]); + ext.deactivate?.(); + }); + + it("non-MCP tools pass through unchanged when an MCP connect fails", async () => { + // failInitialize: the server rejects initialize (a fast failure, not a + // hang) so the connect degrades promptly without waiting on a backstop. + const server = makeServer([tool("create_object")]); + server.failInitialize = true; + const ext = makeExt(server, dispatchConfig({ chrome: { command: "fake" } })); + const { host, getFilter } = makeFakeHost(); + ext.activate(host); + const filter = requireFilter(getFilter); + + const stubNonMcp: ToolContract = { + name: "run_shell", + description: "kept", + parameters: { type: "object" }, + execute: async () => ({ content: "" }), + }; + + const result = await filter({ + tools: [stubNonMcp], + cwd: "/proj", + conversationId: "c", + }); + + // Non-MCP tool survives; the failed MCP server contributed no tools. + expect(result.tools.map((t) => t.name)).toEqual(["run_shell"]); + ext.deactivate?.(); + }); }); diff --git a/packages/mcp/src/extension.ts b/packages/mcp/src/extension.ts index e1c4d52..bb951ee 100644 --- a/packages/mcp/src/extension.ts +++ b/packages/mcp/src/extension.ts @@ -18,6 +18,7 @@ import { resolveServers } from "./config.js"; import type { Logger } from "./manager.js"; import { McpManager } from "./manager.js"; import { adaptTool, namespace } from "./registry.js"; +import { MCP_CONNECT_TIMEOUT_MS } from "./timeout.js"; import type { SpawnedProcess, SpawnProcess } from "./transport.js"; import { createStdioTransport } from "./transport.js"; import type { McpServerStatus, McpService, ResolvedMcpServer } from "./types.js"; @@ -26,9 +27,9 @@ export const mcpServiceHandle: ServiceHandle<McpService> = defineService<McpServ /** Filesystem + process adapters injected into the extension for testability. */ export interface McpExtensionDeps { - readonly spawn: SpawnProcess; - readonly readFile: (path: string) => Promise<string | null>; - readonly getCwd: () => string; + readonly spawn: SpawnProcess; + readonly readFile: (path: string) => Promise<string | null>; + readonly getCwd: () => string; } /** @@ -37,192 +38,224 @@ export interface McpExtensionDeps { * Extracted from the filter handler so it is unit-testable without I/O. */ export function filterMcpTools( - assembly: ToolAssembly, - toolToServer: ReadonlyMap<string, string>, - connectedServerIds: ReadonlySet<string>, + assembly: ToolAssembly, + toolToServer: ReadonlyMap<string, string>, + connectedServerIds: ReadonlySet<string>, ): ToolAssembly { - const filtered = assembly.tools.filter((tool) => { - const serverId = toolToServer.get(tool.name); - if (serverId === undefined) return true; - return connectedServerIds.has(serverId); - }); - return { - tools: filtered, - ...(assembly.cwd !== undefined && { cwd: assembly.cwd }), - ...(assembly.computerId !== undefined && { computerId: assembly.computerId }), - conversationId: assembly.conversationId, - }; + const filtered = assembly.tools.filter((tool) => { + const serverId = toolToServer.get(tool.name); + if (serverId === undefined) return true; + return connectedServerIds.has(serverId); + }); + return { + tools: filtered, + ...(assembly.cwd !== undefined && { cwd: assembly.cwd }), + ...(assembly.computerId !== undefined && { computerId: assembly.computerId }), + conversationId: assembly.conversationId, + }; } /** Map a host Logger to the manager's narrower Logger surface. */ function wrapLogger(logger: HostAPI["logger"]): Logger { - return { - info: (msg, attrs) => logger.info(msg, attrs), - warn: (msg, attrs) => logger.warn(msg, attrs), - error: (msg, attrs) => logger.error(msg, attrs), - }; + return { + info: (msg, attrs) => logger.info(msg, attrs), + warn: (msg, attrs) => logger.warn(msg, attrs), + error: (msg, attrs) => logger.error(msg, attrs), + }; } export function makeMcpExtension(deps: McpExtensionDeps): Extension { - // Module-scoped store so deactivate can reach the manager. Lives in the - // factory closure so each built extension has its own. - const store: { manager: McpManager | null } = { manager: null }; - - return { - manifest: { - id: "mcp", - name: "Model Context Protocol", - version: "0.0.0", - apiVersion: "^0.1.0", - trust: "bundled", - activation: "eager", - dependsOn: ["session-orchestrator"], - capabilities: { spawn: true }, - contributes: { tools: [], services: ["mcp"] }, - }, - activate(host: HostAPI) { - const logger = host.logger; - - const connectionFactory = (server: ResolvedMcpServer, cwd: string) => { - return createStdioTransport( - { - spawn: deps.spawn, - command: server.command, - ...(server.env !== undefined && { env: server.env }), - }, - cwd, - ); - }; - - const manager = new McpManager( - { spawn: deps.spawn, logger: wrapLogger(logger) }, - connectionFactory, - ); - - // Track which tool names belong to which server for the filter. - const toolToServer = new Map<string, string>(); - - function registerToolsFromClient(serverId: string, client: McpClient): void { - const tools = client.getTools(); - for (const mcpTool of tools) { - const name = namespace(serverId, mcpTool.name); - toolToServer.set(name, serverId); - host.defineTool(adaptTool(serverId, mcpTool, client)); - } - } - - async function connectAndRegister(server: ResolvedMcpServer, cwd: string): Promise<void> { - const client = await manager.ensureConnected(server, cwd); - registerToolsFromClient(server.id, client); - - // Wire list_changed → re-list → re-register. onToolsChanged replaces - // the handler; ensureConnected returns the same cached client so this - // is idempotent across turns. - client.onToolsChanged(async () => { - try { - await client.listTools(); - registerToolsFromClient(server.id, client); - } catch (err: unknown) { - logger.error("MCP tools re-list failed", { - serverId: server.id, - error: err instanceof Error ? err.message : String(err), - }); - } - }); - } - - // Resolve config + ensure servers connected, then drop tools whose - // server is not connected. Lazy-spawn happens here (first turn). - host.addFilter(toolsFilter, async (assembly: ToolAssembly): Promise<ToolAssembly> => { - const cwd = assembly.cwd ?? deps.getCwd(); - const dispatchMcpJson = await deps.readFile(joinPath(cwd, ".dispatch", "mcp.json")); - const opencodeJson = await deps.readFile(joinPath(cwd, "opencode.json")); - const { servers } = resolveServers({ dispatchMcpJson, opencodeJson }); - - for (const server of servers) { - try { - await connectAndRegister(server, cwd); - } catch { - // Connection failure — the manager tracks broken state. - } - } - - const statuses = manager.status(servers); - const connectedIds = new Set( - statuses.filter((s) => s.state === "connected").map((s) => s.id), - ); - - return filterMcpTools(assembly, toolToServer, connectedIds); - }); - - // Provide the MCP service (status introspection). - const service: McpService = { - async status(cwd: string): Promise<readonly McpServerStatus[]> { - const dispatchMcpJson = await deps.readFile(joinPath(cwd, ".dispatch", "mcp.json")); - const opencodeJson = await deps.readFile(joinPath(cwd, "opencode.json")); - const { servers } = resolveServers({ dispatchMcpJson, opencodeJson }); - return manager.status(servers); - }, - }; - host.provideService(mcpServiceHandle, service); - - store.manager = manager; - logger.info("MCP extension activated"); - }, - deactivate() { - store.manager?.shutdownAll(); - store.manager = null; - }, - }; + // Module-scoped store so deactivate can reach the manager. Lives in the + // factory closure so each built extension has its own. + const store: { manager: McpManager | null } = { manager: null }; + + return { + manifest: { + id: "mcp", + name: "Model Context Protocol", + version: "0.0.0", + apiVersion: "^0.1.0", + trust: "bundled", + activation: "eager", + dependsOn: ["session-orchestrator"], + capabilities: { spawn: true }, + contributes: { tools: [], services: ["mcp"] }, + }, + activate(host: HostAPI) { + const logger = host.logger; + + const connectionFactory = (server: ResolvedMcpServer, cwd: string) => { + return createStdioTransport( + { + spawn: deps.spawn, + command: server.command, + ...(server.env !== undefined && { env: server.env }), + }, + cwd, + ); + }; + + const manager = new McpManager( + { spawn: deps.spawn, logger: wrapLogger(logger) }, + connectionFactory, + ); + + // Track which tool names belong to which server for the filter. + const toolToServer = new Map<string, string>(); + + function registerToolsFromClient(serverId: string, client: McpClient): void { + const tools = client.getTools(); + for (const mcpTool of tools) { + const name = namespace(serverId, mcpTool.name); + toolToServer.set(name, serverId); + host.defineTool(adaptTool(serverId, mcpTool, client)); + } + } + + async function connectAndRegister( + server: ResolvedMcpServer, + cwd: string, + signal?: AbortSignal, + ): Promise<void> { + const client = await manager.ensureConnected(server, cwd, signal); + registerToolsFromClient(server.id, client); + + // Wire list_changed → re-list → re-register. onToolsChanged replaces + // the handler; ensureConnected returns the same cached client so this + // is idempotent across turns. The async re-list runs LATER (not during + // this filter), so it is NOT bound to the filter's signal (already + // done) — it relies on listTools()'s own default timeout instead. + client.onToolsChanged(async () => { + try { + await client.listTools(); + registerToolsFromClient(server.id, client); + } catch (err: unknown) { + logger.error("MCP tools re-list failed", { + serverId: server.id, + error: err instanceof Error ? err.message : String(err), + }); + } + }); + } + + // Resolve config + ensure servers connected, then drop tools whose + // server is not connected. Lazy-spawn happens here (first turn). + // + // The whole connect phase is wrapped in a per-filter AbortController + // that fires on EITHER (a) the turn's signal (`assembly.signal`, so + // POST /conversations/:id/stop interrupts a stuck connect immediately) + // OR (b) a timeout (`MCP_CONNECT_TIMEOUT_MS`, so a misbehaving / + // framing-incompatible server cannot hang the turn forever). On abort + // we degrade gracefully: skip MCP tools for this turn rather than block. + host.addFilter(toolsFilter, async (assembly: ToolAssembly): Promise<ToolAssembly> => { + const cwd = assembly.cwd ?? deps.getCwd(); + const dispatchMcpJson = await deps.readFile(joinPath(cwd, ".dispatch", "mcp.json")); + const opencodeJson = await deps.readFile(joinPath(cwd, "opencode.json")); + const { servers } = resolveServers({ dispatchMcpJson, opencodeJson }); + + const controller = new AbortController(); + const parentSignal = assembly.signal; + const onParentAbort = (): void => controller.abort(); + if (parentSignal !== undefined) { + if (parentSignal.aborted) { + controller.abort(); + } else { + parentSignal.addEventListener("abort", onParentAbort, { once: true }); + } + } + const timer = setTimeout(() => controller.abort(), MCP_CONNECT_TIMEOUT_MS); + + try { + for (const server of servers) { + try { + await connectAndRegister(server, cwd, controller.signal); + } catch { + // Connection failure / timeout / aborted — the manager tracks + // broken state; we keep going (or abort cascades) below. + } + if (controller.signal.aborted) break; + } + } finally { + clearTimeout(timer); + parentSignal?.removeEventListener("abort", onParentAbort); + } + + const statuses = manager.status(servers); + const connectedIds = new Set( + statuses.filter((s) => s.state === "connected").map((s) => s.id), + ); + + return filterMcpTools(assembly, toolToServer, connectedIds); + }); + + // Provide the MCP service (status introspection). + const service: McpService = { + async status(cwd: string): Promise<readonly McpServerStatus[]> { + const dispatchMcpJson = await deps.readFile(joinPath(cwd, ".dispatch", "mcp.json")); + const opencodeJson = await deps.readFile(joinPath(cwd, "opencode.json")); + const { servers } = resolveServers({ dispatchMcpJson, opencodeJson }); + return manager.status(servers); + }, + }; + host.provideService(mcpServiceHandle, service); + + store.manager = manager; + logger.info("MCP extension activated"); + }, + deactivate() { + store.manager?.shutdownAll(); + store.manager = null; + }, + }; } // --- real Bun-backed adapters (production wiring) --- function realSpawn( - command: readonly string[], - opts: { readonly cwd: string; readonly env?: Readonly<Record<string, string>> | undefined }, + command: readonly string[], + opts: { readonly cwd: string; readonly env?: Readonly<Record<string, string>> | undefined }, ): SpawnedProcess { - const env: Record<string, string | undefined> = { ...process.env }; - if (opts.env) { - for (const [key, value] of Object.entries(opts.env)) { - env[key] = value; - } - } - const proc = Bun.spawn(command as string[], { - cwd: opts.cwd, - env: env as Record<string, string>, - stdin: "pipe", - stdout: "pipe", - stderr: "pipe", - }); - return { - stdin: proc.stdin, - stdout: proc.stdout, - stderr: proc.stderr, - pid: proc.pid, - kill: () => proc.kill(), - }; + const env: Record<string, string | undefined> = { ...process.env }; + if (opts.env) { + for (const [key, value] of Object.entries(opts.env)) { + env[key] = value; + } + } + const proc = Bun.spawn(command as string[], { + cwd: opts.cwd, + env: env as Record<string, string>, + stdin: "pipe", + stdout: "pipe", + stderr: "pipe", + }); + return { + stdin: proc.stdin, + stdout: proc.stdout, + stderr: proc.stderr, + pid: proc.pid, + kill: () => proc.kill(), + }; } async function realReadFile(path: string): Promise<string | null> { - try { - const file = Bun.file(path); - if (await file.exists()) { - return file.text(); - } - return null; - } catch { - return null; - } + try { + const file = Bun.file(path); + if (await file.exists()) { + return file.text(); + } + return null; + } catch { + return null; + } } function joinPath(...parts: readonly string[]): string { - return parts.join("/"); + return parts.join("/"); } /** Production extension: real Bun spawn + filesystem reads. */ export const extension: Extension = makeMcpExtension({ - spawn: realSpawn, - readFile: realReadFile, - getCwd: () => process.cwd(), + spawn: realSpawn, + readFile: realReadFile, + getCwd: () => process.cwd(), }); diff --git a/packages/mcp/src/framing.test.ts b/packages/mcp/src/framing.test.ts index be8cb8e..9832c74 100644 --- a/packages/mcp/src/framing.test.ts +++ b/packages/mcp/src/framing.test.ts @@ -1,92 +1,163 @@ import { describe, expect, it } from "vitest"; import { encode, FrameDecoder } from "./framing.js"; +/** Build a legacy Content-Length frame (for exercising the decoder's CL path). */ +function contentLengthFrame(body: string): Uint8Array { + const bodyBytes = new TextEncoder().encode(body); + return new TextEncoder().encode(`Content-Length: ${bodyBytes.length}\r\n\r\n${body}`); +} + describe("encode", () => { - it("produces correct Content-Length header", () => { - const msg = '{"jsonrpc":"2.0","id":1,"method":"initialize"}'; - const encoded = encode(msg); - const text = new TextDecoder().decode(encoded); - expect(text).toBe(`Content-Length: ${new TextEncoder().encode(msg).length}\r\n\r\n${msg}`); - }); - - it("handles empty message", () => { - const encoded = encode(""); - const text = new TextDecoder().decode(encoded); - expect(text).toBe("Content-Length: 0\r\n\r\n"); - }); + it("produces newline-delimited JSON (current MCP spec framing)", () => { + const msg = '{"jsonrpc":"2.0","id":1,"method":"initialize"}'; + const encoded = encode(msg); + const text = new TextDecoder().decode(encoded); + expect(text).toBe(`${msg}\n`); + }); + + it("appends a trailing newline to an empty message", () => { + const encoded = encode(""); + const text = new TextDecoder().decode(encoded); + expect(text).toBe("\n"); + }); }); -describe("FrameDecoder", () => { - it("reassembles a complete message from one chunk", () => { - const msg = '{"jsonrpc":"2.0","id":1}'; - const encoded = encode(msg); - const decoder = new FrameDecoder(); - const messages = decoder.decode(encoded); - expect(messages).toEqual([msg]); - }); - - it("handles split across chunks", () => { - const msg = '{"jsonrpc":"2.0","id":1,"method":"initialize"}'; - const encoded = encode(msg); - const mid = Math.floor(encoded.length / 2); - const chunk1 = encoded.slice(0, mid); - const chunk2 = encoded.slice(mid); - - const decoder = new FrameDecoder(); - const result1 = decoder.decode(chunk1); - expect(result1).toEqual([]); - - const result2 = decoder.decode(chunk2); - expect(result2).toEqual([msg]); - }); - - it("handles two messages in one chunk", () => { - const msg1 = '{"jsonrpc":"2.0","id":1}'; - const msg2 = '{"jsonrpc":"2.0","id":2}'; - const encoded1 = encode(msg1); - const encoded2 = encode(msg2); - const combined = new Uint8Array(encoded1.length + encoded2.length); - combined.set(encoded1); - combined.set(encoded2, encoded1.length); - - const decoder = new FrameDecoder(); - const messages = decoder.decode(combined); - expect(messages).toEqual([msg1, msg2]); - }); - - it("rejects negative Content-Length by skipping header", () => { - const header = "Content-Length: -5\r\n\r\n"; - const encoded = new TextEncoder().encode(`${header}extra`); - const decoder = new FrameDecoder(); - const messages = decoder.decode(encoded); - // Negative length does not match the digit capture, so the header is skipped. - expect(messages).toEqual([]); - }); - - it("rejects zero Content-Length", () => { - const encoded = encode(""); - const decoder = new FrameDecoder(); - const messages = decoder.decode(encoded); - expect(messages).toEqual([""]); - }); - - it("reassembles multi-byte UTF-8 content (byte-length, not char-length)", () => { - // "héllo" — é is two UTF-8 bytes; Content-Length counts bytes. - const msg = '{"text":"héllo 🚀"}'; - const encoded = encode(msg); - expect(new TextEncoder().encode(msg).length).toBeGreaterThan(msg.length); - - const decoder = new FrameDecoder(); - const messages = decoder.decode(encoded); - expect(messages).toEqual([msg]); - }); - - it("reassembles multi-byte content split across a chunk boundary", () => { - const msg = '{"text":"日本語のテスト"}'; - const encoded = encode(msg); - const mid = Math.floor(encoded.length / 2); - const decoder = new FrameDecoder(); - expect(decoder.decode(encoded.slice(0, mid))).toEqual([]); - expect(decoder.decode(encoded.slice(mid))).toEqual([msg]); - }); +describe("FrameDecoder — newline-delimited JSON", () => { + it("decodes a single newline-delimited message", () => { + const msg = '{"jsonrpc":"2.0","id":1}'; + const decoder = new FrameDecoder(); + expect(decoder.decode(encode(msg))).toEqual([msg]); + }); + + it("decodes a CRLF-terminated message (trailing \\r tolerated)", () => { + const msg = '{"jsonrpc":"2.0","id":1}'; + const decoder = new FrameDecoder(); + const framed = new TextEncoder().encode(`${msg}\r\n`); + expect(decoder.decode(framed)).toEqual([msg]); + }); + + it("handles a split across chunks", () => { + const msg = '{"jsonrpc":"2.0","id":1,"method":"initialize"}'; + const encoded = encode(msg); + const mid = Math.floor(encoded.length / 2); + + const decoder = new FrameDecoder(); + expect(decoder.decode(encoded.slice(0, mid))).toEqual([]); + expect(decoder.decode(encoded.slice(mid))).toEqual([msg]); + }); + + it("handles two messages in one chunk", () => { + const msg1 = '{"jsonrpc":"2.0","id":1}'; + const msg2 = '{"jsonrpc":"2.0","id":2}'; + const combined = new Uint8Array(encode(msg1).length + encode(msg2).length); + combined.set(encode(msg1)); + combined.set(encode(msg2), encode(msg1).length); + + const decoder = new FrameDecoder(); + expect(decoder.decode(combined)).toEqual([msg1, msg2]); + }); + + it("skips blank lines between messages", () => { + const msg = '{"jsonrpc":"2.0","id":1}'; + const framed = new TextEncoder().encode(`\n\n${msg}\n\n`); + const decoder = new FrameDecoder(); + expect(decoder.decode(framed)).toEqual([msg]); + }); + + it("reassembles multi-byte UTF-8 content (byte-aware, not char-aware)", () => { + const msg = '{"text":"héllo 🚀"}'; + expect(new TextEncoder().encode(msg).length).toBeGreaterThan(msg.length); + const decoder = new FrameDecoder(); + expect(decoder.decode(encode(msg))).toEqual([msg]); + }); + + it("reassembles multi-byte content split across a chunk boundary", () => { + const msg = '{"text":"日本語のテスト"}'; + const encoded = encode(msg); + const mid = Math.floor(encoded.length / 2); + const decoder = new FrameDecoder(); + expect(decoder.decode(encoded.slice(0, mid))).toEqual([]); + expect(decoder.decode(encoded.slice(mid))).toEqual([msg]); + }); + + it("does not split a JSON string containing an escaped \\n (no raw newline)", () => { + // JSON escapes newlines inside strings as the two chars `\` + `n`; a raw + // 0x0a only ever appears as a message separator. So a JSON body carrying + // an embedded newline literal survives intact. + const msg = '{"text":"line1\\nline2"}'; + const decoder = new FrameDecoder(); + expect(decoder.decode(encode(msg))).toEqual([msg]); + }); +}); + +describe("FrameDecoder — legacy Content-Length framing (auto-detected)", () => { + it("decodes a Content-Length-framed message", () => { + const msg = '{"jsonrpc":"2.0","id":1}'; + const decoder = new FrameDecoder(); + expect(decoder.decode(contentLengthFrame(msg))).toEqual([msg]); + }); + + it("reassembles a Content-Length frame split across chunks", () => { + const msg = '{"jsonrpc":"2.0","id":1,"method":"initialize"}'; + const encoded = contentLengthFrame(msg); + const mid = Math.floor(encoded.length / 2); + const decoder = new FrameDecoder(); + expect(decoder.decode(encoded.slice(0, mid))).toEqual([]); + expect(decoder.decode(encoded.slice(mid))).toEqual([msg]); + }); + + it("decodes two Content-Length frames in one chunk", () => { + const msg1 = '{"jsonrpc":"2.0","id":1}'; + const msg2 = '{"jsonrpc":"2.0","id":2}'; + const combined = new Uint8Array( + contentLengthFrame(msg1).length + contentLengthFrame(msg2).length, + ); + combined.set(contentLengthFrame(msg1)); + combined.set(contentLengthFrame(msg2), contentLengthFrame(msg1).length); + + const decoder = new FrameDecoder(); + expect(decoder.decode(combined)).toEqual([msg1, msg2]); + }); + + it("rejects negative Content-Length by skipping header", () => { + const encoded = new TextEncoder().encode("Content-Length: -5\r\n\r\nextra"); + const decoder = new FrameDecoder(); + expect(decoder.decode(encoded)).toEqual([]); + }); + + it("accepts zero Content-Length as an empty message", () => { + const decoder = new FrameDecoder(); + expect(decoder.decode(contentLengthFrame(""))).toEqual([""]); + }); + + it("reassembles multi-byte UTF-8 via Content-Length (byte count)", () => { + const msg = '{"text":"héllo 🚀"}'; + const decoder = new FrameDecoder(); + expect(decoder.decode(contentLengthFrame(msg))).toEqual([msg]); + }); + + it("does not mis-read a partial 'Content-Length' prefix as newline-delimited", () => { + // A buffer that is a partial prefix of "Content-Length:" must WAIT for more + // bytes rather than being split on a (nonexistent) newline. + const decoder = new FrameDecoder(); + const partial = new TextEncoder().encode("Content-Len"); + expect(decoder.decode(partial)).toEqual([]); + const rest = new TextEncoder().encode("gth: 3\r\n\r\nabc"); + expect(decoder.decode(rest)).toEqual(["abc"]); + }); +}); + +describe("FrameDecoder — mixed framings", () => { + it("decodes a Content-Length frame followed by a newline-delimited message", () => { + const clMsg = '{"jsonrpc":"2.0","id":1}'; + const nlMsg = '{"jsonrpc":"2.0","id":2}'; + const cl = contentLengthFrame(clMsg); + const nl = encode(nlMsg); + const combined = new Uint8Array(cl.length + nl.length); + combined.set(cl); + combined.set(nl, cl.length); + + const decoder = new FrameDecoder(); + expect(decoder.decode(combined)).toEqual([clMsg, nlMsg]); + }); }); diff --git a/packages/mcp/src/framing.ts b/packages/mcp/src/framing.ts index a6e8b99..d30b871 100644 --- a/packages/mcp/src/framing.ts +++ b/packages/mcp/src/framing.ts @@ -1,100 +1,192 @@ /** - * Content-Length framing for MCP stdio transport. + * MCP stdio framing. * - * Each JSON-RPC message is framed as: - * Content-Length: <byte-length>\r\n\r\n<JSON bytes> + * The MCP spec (revision 2025-03-26 and later, including 2025-11-25) frames + * JSON-RPC messages over stdio as **newline-delimited JSON**: each message is + * a single JSON document followed by a `\n` (or `\r\n`). This replaced the + * older LSP-style `Content-Length: N\r\n\r\n<JSON>` framing that the protocol + * inherited originally. Modern servers (e.g. chrome-devtools-mcp) speak ONLY + * newline-delimited JSON — they emit `<JSON>\n` on stdout and do not respond to + * `Content-Length`-framed input at all. * - * Same framing as LSP — the MCP spec inherited this from the LSP base protocol. + * Outgoing messages are therefore encoded as newline-delimited JSON (the + * current spec default), so we can talk to modern servers. The decoder + * **auto-detects** the incoming framing per message — it accepts BOTH + * newline-delimited JSON and legacy `Content-Length` frames — so we still + * interoperate with servers that respond with the old framing. * * PURE: no I/O. Operates on bytes (Uint8Array) so multi-byte UTF-8 content is * handled correctly — `Content-Length` is a *byte* count, not a character count. */ +const HEADER_PREFIX = "Content-Length:"; const HEADER_SEP = "\r\n\r\n"; const CONTENT_LENGTH_RE = /Content-Length:\s*(\d+)/i; -const SEP_BYTES = new TextEncoder().encode(HEADER_SEP); +const encoder = new TextEncoder(); +const SEP_BYTES = encoder.encode(HEADER_SEP); + +const CR = 0x0d; +const LF = 0x0a; /** - * Encode a JSON string into a single Content-Length-framed message. - * Returns the full frame (header + blank line + body) as bytes. + * Encode a JSON string as a single newline-delimited frame (the current MCP + * spec stdio framing): `<JSON>\n`. This is what we send to MCP servers. */ export function encode(msg: string): Uint8Array { - const body = new TextEncoder().encode(msg); - const header = `Content-Length: ${body.length}\r\n\r\n`; - const frame = new TextEncoder().encode(header); - const result = new Uint8Array(frame.length + body.length); - result.set(frame); - result.set(body, frame.length); - return result; + return encoder.encode(`${msg}\n`); } /** Find the first occurrence of `needle` in `haystack` at or after `from`. -1 if absent. */ function indexOfBytes(haystack: Uint8Array, needle: Uint8Array, from: number): number { - if (needle.length === 0) return from; - const max = haystack.length - needle.length; - for (let i = from; i <= max; i++) { - let match = true; - for (let j = 0; j < needle.length; j++) { - if (haystack[i + j] !== needle[j]) { - match = false; - break; - } - } - if (match) return i; - } - return -1; + if (needle.length === 0) return from; + const max = haystack.length - needle.length; + for (let i = from; i <= max; i++) { + let match = true; + for (let j = 0; j < needle.length; j++) { + if (haystack[i + j] !== needle[j]) { + match = false; + break; + } + } + if (match) return i; + } + return -1; +} + +/** Find the first occurrence of a single byte at or after `from`. -1 if absent. */ +function indexOfByte(haystack: Uint8Array, needle: number, from: number): number { + for (let i = from; i < haystack.length; i++) { + if (haystack[i] === needle) return i; + } + return -1; +} + +/** Lowercase an ASCII byte (A-Z → a-z); leave everything else unchanged. */ +function toLowerByte(b: number): number { + return b >= 0x41 && b <= 0x5a ? b + 0x20 : b; +} + +function isLineTerminator(b: number | undefined): boolean { + return b === CR || b === LF; +} + +/** + * How does `buf` relate to the `Content-Length:` header prefix (case-insensitive)? + * - `HEADER_PREFIX.length` → `buf` starts with the full `Content-Length:` prefix. + * - a positive number `< HEADER_PREFIX.length` → `buf` is a (possibly partial) + * prefix of `Content-Length:` — ambiguous, the caller must wait for more bytes. + * - `-1` → `buf` definitively does NOT start with `Content-Length:` (not CL framing). + */ +function contentLengthPrefixLength(buf: Uint8Array): number { + const n = Math.min(buf.length, HEADER_PREFIX.length); + for (let i = 0; i < n; i++) { + const a = buf[i]; + if (a === undefined) return -1; // unreachable: i < n <= buf.length + if (toLowerByte(a) !== toLowerByte(HEADER_PREFIX.charCodeAt(i))) return -1; + } + return n; } /** * Feed raw bytes into the decoder. Returns all complete JSON messages that can * be extracted from the accumulated buffer. Buffers partial frames across calls. + * + * Auto-detects framing per message: a `Content-Length:`-prefixed buffer is parsed + * as a Content-Length frame (legacy/LSP-style); anything else is parsed as + * newline-delimited JSON (current MCP spec). Both framings may be mixed in a + * single stream. */ export class FrameDecoder { - private buf: Uint8Array = new Uint8Array(0); - private readonly decoder = new TextDecoder(); - - decode(chunk: Uint8Array): string[] { - // Append the incoming chunk to the internal byte buffer. - const next = new Uint8Array(this.buf.length + chunk.length); - next.set(this.buf); - next.set(chunk, this.buf.length); - this.buf = next; - - const messages: string[] = []; - - while (true) { - const sepIdx = indexOfBytes(this.buf, SEP_BYTES, 0); - if (sepIdx === -1) break; - - // The header block is everything before the separator; parse - // Content-Length from it (ASCII, so decoding the slice is safe). - const headerText = this.decoder.decode(this.buf.subarray(0, sepIdx)); - const match = CONTENT_LENGTH_RE.exec(headerText); - const bodyStart = sepIdx + SEP_BYTES.length; - - if (!match?.[1]) { - // No usable Content-Length — drop this header and continue scanning. - this.buf = this.buf.subarray(bodyStart); - continue; - } - - const length = Number.parseInt(match[1], 10); - if (length < 0) { - this.buf = this.buf.subarray(bodyStart); - continue; - } - - if (this.buf.length - bodyStart < length) { - // Body not fully received yet; wait for more bytes. - break; - } - - // Decode exactly `length` body bytes (preserves multi-byte UTF-8). - messages.push(this.decoder.decode(this.buf.subarray(bodyStart, bodyStart + length))); - this.buf = this.buf.subarray(bodyStart + length); - } - - return messages; - } + private buf: Uint8Array = new Uint8Array(0); + private readonly decoder = new TextDecoder(); + + decode(chunk: Uint8Array): string[] { + // Append the incoming chunk to the internal byte buffer. + if (chunk.length > 0) { + const next = new Uint8Array(this.buf.length + chunk.length); + next.set(this.buf); + next.set(chunk, this.buf.length); + this.buf = next; + } + + const messages: string[] = []; + + while (true) { + // 1. Skip leading CR/LF whitespace between frames. + let i = 0; + while (i < this.buf.length && isLineTerminator(this.buf[i])) i++; + if (i > 0) this.buf = this.buf.subarray(i); + if (this.buf.length === 0) break; + + // 2. Detect framing. + const prefix = contentLengthPrefixLength(this.buf); + if (prefix >= 0 && prefix < HEADER_PREFIX.length) { + // Buffer is a (possibly partial) prefix of "Content-Length:" — ambiguous; + // wait for more bytes before deciding this is (or isn't) a CL frame. + break; + } + if (prefix === HEADER_PREFIX.length) { + // Content-Length framing. + const result = this.tryParseContentLength(); + if (result === "incomplete") break; + if (result !== "skip") messages.push(result); + continue; + } + + // 3. Newline-delimited JSON: one message per line, terminated by `\n` + // (a trailing `\r` before the `\n` is tolerated). A raw newline byte + // can only ever appear as a message separator — JSON escapes newlines + // inside strings as the two characters `\n`, never a literal 0x0a — so + // splitting on `\n` bytes is safe for valid JSON. + const nl = indexOfByte(this.buf, LF, 0); + if (nl === -1) break; // incomplete line — wait for more bytes + + let end = nl; + if (end > 0 && this.buf[end - 1] === CR) end--; + const text = this.decoder.decode(this.buf.subarray(0, end)); + this.buf = this.buf.subarray(nl + 1); + if (text.length > 0) messages.push(text); + } + + return messages; + } + + /** + * Parse one Content-Length frame from the front of `this.buf`. Returns: + * - the decoded body string (possibly `""` for a zero-length body), + * - `"incomplete"` if the header or body hasn't fully arrived (caller waits), + * - `"skip"` if the header was consumed but carried no usable Content-Length + * (caller continues scanning without emitting a message). + */ + private tryParseContentLength(): string | "incomplete" | "skip" { + const sepIdx = indexOfBytes(this.buf, SEP_BYTES, 0); + if (sepIdx === -1) return "incomplete"; // header not fully received yet + + const headerText = this.decoder.decode(this.buf.subarray(0, sepIdx)); + const match = CONTENT_LENGTH_RE.exec(headerText); + const bodyStart = sepIdx + SEP_BYTES.length; + + if (!match?.[1]) { + // No usable Content-Length — drop this header and continue scanning. + this.buf = this.buf.subarray(bodyStart); + return "skip"; + } + + const length = Number.parseInt(match[1], 10); + if (length < 0) { + this.buf = this.buf.subarray(bodyStart); + return "skip"; + } + + if (this.buf.length - bodyStart < length) { + // Body not fully received yet; wait for more bytes. + return "incomplete"; + } + + // Decode exactly `length` body bytes (preserves multi-byte UTF-8). + const body = this.decoder.decode(this.buf.subarray(bodyStart, bodyStart + length)); + this.buf = this.buf.subarray(bodyStart + length); + return body; + } } diff --git a/packages/mcp/src/index.ts b/packages/mcp/src/index.ts index 1f2667d..6d186d7 100644 --- a/packages/mcp/src/index.ts +++ b/packages/mcp/src/index.ts @@ -1,31 +1,37 @@ export { McpClient, type McpClientState } from "./client.js"; export { type ResolveServersDeps, resolveServers } from "./config.js"; export { - extension, - filterMcpTools, - type McpExtensionDeps, - makeMcpExtension, - mcpServiceHandle, + extension, + filterMcpTools, + type McpExtensionDeps, + makeMcpExtension, + mcpServiceHandle, } from "./extension.js"; export { encode, FrameDecoder } from "./framing.js"; export { type Logger, McpManager, type McpManagerDeps } from "./manager.js"; export { adaptTool, flattenContent, namespace } from "./registry.js"; export { - type Connection, - createStdioTransport, - type SpawnedProcess, - type SpawnProcess, + MCP_CONNECT_TIMEOUT_MS, + MCP_DEFAULT_TIMEOUT_MS, + McpTimeoutError, + withTimeout, +} from "./timeout.js"; +export { + type Connection, + createStdioTransport, + type SpawnedProcess, + type SpawnProcess, } from "./transport.js"; export type { - McpCallResult, - McpContentItem, - McpServerCapabilities, - McpServerConfig, - McpServerState, - McpServerStatus, - McpService, - McpToolCaller, - McpToolInfo, - ResolvedMcpServer, - ResolveResult, + McpCallResult, + McpContentItem, + McpServerCapabilities, + McpServerConfig, + McpServerState, + McpServerStatus, + McpService, + McpToolCaller, + McpToolInfo, + ResolvedMcpServer, + ResolveResult, } from "./types.js"; diff --git a/packages/mcp/src/manager.test.ts b/packages/mcp/src/manager.test.ts index 275b078..d6179b7 100644 --- a/packages/mcp/src/manager.test.ts +++ b/packages/mcp/src/manager.test.ts @@ -4,213 +4,238 @@ import type { Connection } from "./transport.js"; import type { ResolvedMcpServer } from "./types.js"; function makeMockConnection(): Connection { - return { - send: async (method: string) => { - if (method === "initialize") { - return { - protocolVersion: "2025-11-25", - capabilities: { tools: { listChanged: false } }, - serverInfo: { name: "test", version: "1.0.0" }, - }; - } - if (method === "tools/list") { - return { - tools: [ - { - name: "tool_a", - description: "Tool A", - inputSchema: { type: "object" }, - }, - ], - }; - } - return {}; - }, - notify: () => {}, - onNotification: () => {}, - close: () => {}, - pid: 100, - }; + return { + send: async (method: string) => { + if (method === "initialize") { + return { + protocolVersion: "2025-11-25", + capabilities: { tools: { listChanged: false } }, + serverInfo: { name: "test", version: "1.0.0" }, + }; + } + if (method === "tools/list") { + return { + tools: [ + { + name: "tool_a", + description: "Tool A", + inputSchema: { type: "object" }, + }, + ], + }; + } + return {}; + }, + notify: () => {}, + onNotification: () => {}, + close: () => {}, + pid: 100, + }; } function makeBrokenConnection(): Connection { - return { - send: async () => { - throw new Error("Connection refused"); - }, - notify: () => {}, - onNotification: () => {}, - close: () => {}, - pid: 101, - }; + return { + send: async () => { + throw new Error("Connection refused"); + }, + notify: () => {}, + onNotification: () => {}, + close: () => {}, + pid: 101, + }; } function makeManager( - connectionFactory: (_server: ResolvedMcpServer) => { - connection: Connection; - promise: Promise<void>; - }, + connectionFactory: (_server: ResolvedMcpServer) => { + connection: Connection; + promise: Promise<void>; + }, ): McpManager { - const currentTime = 1000; - const deps: McpManagerDeps = { - spawn: () => ({ - stdin: { write: () => {} }, - stdout: { on: () => {} }, - pid: 1, - kill: () => {}, - }), - now: () => currentTime, - }; - - const factory = (_server: ResolvedMcpServer, _cwd: string) => connectionFactory(_server); - - const manager = new McpManager(deps, factory); - return manager; + const currentTime = 1000; + const deps: McpManagerDeps = { + spawn: () => ({ + stdin: { write: () => {} }, + stdout: { on: () => {} }, + pid: 1, + kill: () => {}, + }), + now: () => currentTime, + }; + + const factory = (_server: ResolvedMcpServer, _cwd: string) => connectionFactory(_server); + + const manager = new McpManager(deps, factory); + return manager; } const testServer: ResolvedMcpServer = { - id: "test-server", - command: ["test-cmd"], - configSource: ".dispatch/mcp.json", + id: "test-server", + command: ["test-cmd"], + configSource: ".dispatch/mcp.json", }; describe("McpManager", () => { - it("lazy-spawn on first access", async () => { - const conn = makeMockConnection(); - let spawnCount = 0; - const manager = makeManager((_server) => { - spawnCount++; - return { connection: conn, promise: Promise.resolve() }; - }); - - const client = await manager.ensureConnected(testServer, "/tmp"); - expect(client).toBeDefined(); - expect(spawnCount).toBe(1); - }); - - it("reuses existing client on second access", async () => { - const conn = makeMockConnection(); - let spawnCount = 0; - const manager = makeManager(() => { - spawnCount++; - return { connection: conn, promise: Promise.resolve() }; - }); - - const client1 = await manager.ensureConnected(testServer, "/tmp"); - const client2 = await manager.ensureConnected(testServer, "/tmp"); - expect(client1).toBe(client2); - expect(spawnCount).toBe(1); - }); - - it("status returns server states", async () => { - const conn = makeMockConnection(); - const manager = makeManager(() => { - return { connection: conn, promise: Promise.resolve() }; - }); - - // Before connecting - let statuses = manager.status([testServer]); - expect(statuses.length).toBe(1); - expect(statuses[0].state).toBe("disconnected"); - - // After connecting - await manager.ensureConnected(testServer, "/tmp"); - statuses = manager.status([testServer]); - expect(statuses.length).toBe(1); - expect(statuses[0].state).toBe("connected"); - expect(statuses[0].toolCount).toBe(1); - }); - - it("shutdownAll kills all clients", async () => { - let closed = false; - const conn: Connection = { - send: async (method: string) => { - if (method === "initialize") { - return { - protocolVersion: "2025-11-25", - capabilities: {}, - serverInfo: { name: "test", version: "1.0.0" }, - }; - } - if (method === "tools/list") return { tools: [] }; - return {}; - }, - notify: () => {}, - onNotification: () => {}, - close: () => { - closed = true; - }, - pid: 200, - }; - - const manager = makeManager(() => { - return { connection: conn, promise: Promise.resolve() }; - }); - - await manager.ensureConnected(testServer, "/tmp"); - manager.shutdownAll(); - - expect(closed).toBe(true); - const statuses = manager.status([testServer]); - expect(statuses[0].state).toBe("disconnected"); - }); - - it("broken server reports error state", async () => { - const manager = makeManager(() => { - return { connection: makeBrokenConnection(), promise: Promise.resolve() }; - }); - - await expect(manager.ensureConnected(testServer, "/tmp")).rejects.toThrow(); - - const statuses = manager.status([testServer]); - expect(statuses[0].state).toBe("error"); - expect(statuses[0].error).toContain("test-server"); - }); - - it("broken server retries after backoff", async () => { - let currentTime = 1000; - const brokenConn = makeBrokenConnection(); - const goodConn = makeMockConnection(); - let useGood = false; - - const deps: McpManagerDeps = { - spawn: () => ({ - stdin: { write: () => {} }, - stdout: { on: () => {} }, - pid: 1, - kill: () => {}, - }), - now: () => currentTime, - }; - - const factory = (_server: ResolvedMcpServer, _cwd: string) => { - const conn = useGood ? goodConn : brokenConn; - return { connection: conn, promise: Promise.resolve() }; - }; - - const manager = new McpManager(deps, factory); - - // First attempt fails - await expect(manager.ensureConnected(testServer, "/tmp")).rejects.toThrow(); - expect(manager.status([testServer])[0].state).toBe("error"); - - // Not enough time passed — still broken - currentTime += 29_000; - expect(manager.status([testServer])[0].state).toBe("error"); - - // After backoff — should allow retry - useGood = true; - currentTime += 2_000; - const statuses = manager.status([testServer]); - // After backoff, status() clears the broken entry - expect(statuses[0].state).toBe("disconnected"); - }); - - it("getClient returns undefined for unknown server", () => { - const manager = makeManager(() => { - return { connection: makeMockConnection(), promise: Promise.resolve() }; - }); - - expect(manager.getClient("nonexistent")).toBeUndefined(); - }); + it("lazy-spawn on first access", async () => { + const conn = makeMockConnection(); + let spawnCount = 0; + const manager = makeManager((_server) => { + spawnCount++; + return { connection: conn, promise: Promise.resolve() }; + }); + + const client = await manager.ensureConnected(testServer, "/tmp"); + expect(client).toBeDefined(); + expect(spawnCount).toBe(1); + }); + + it("reuses existing client on second access", async () => { + const conn = makeMockConnection(); + let spawnCount = 0; + const manager = makeManager(() => { + spawnCount++; + return { connection: conn, promise: Promise.resolve() }; + }); + + const client1 = await manager.ensureConnected(testServer, "/tmp"); + const client2 = await manager.ensureConnected(testServer, "/tmp"); + expect(client1).toBe(client2); + expect(spawnCount).toBe(1); + }); + + it("status returns server states", async () => { + const conn = makeMockConnection(); + const manager = makeManager(() => { + return { connection: conn, promise: Promise.resolve() }; + }); + + // Before connecting + let statuses = manager.status([testServer]); + expect(statuses.length).toBe(1); + expect(statuses[0].state).toBe("disconnected"); + + // After connecting + await manager.ensureConnected(testServer, "/tmp"); + statuses = manager.status([testServer]); + expect(statuses.length).toBe(1); + expect(statuses[0].state).toBe("connected"); + expect(statuses[0].toolCount).toBe(1); + }); + + it("shutdownAll kills all clients", async () => { + let closed = false; + const conn: Connection = { + send: async (method: string) => { + if (method === "initialize") { + return { + protocolVersion: "2025-11-25", + capabilities: {}, + serverInfo: { name: "test", version: "1.0.0" }, + }; + } + if (method === "tools/list") return { tools: [] }; + return {}; + }, + notify: () => {}, + onNotification: () => {}, + close: () => { + closed = true; + }, + pid: 200, + }; + + const manager = makeManager(() => { + return { connection: conn, promise: Promise.resolve() }; + }); + + await manager.ensureConnected(testServer, "/tmp"); + manager.shutdownAll(); + + expect(closed).toBe(true); + const statuses = manager.status([testServer]); + expect(statuses[0].state).toBe("disconnected"); + }); + + it("broken server reports error state", async () => { + const manager = makeManager(() => { + return { connection: makeBrokenConnection(), promise: Promise.resolve() }; + }); + + await expect(manager.ensureConnected(testServer, "/tmp")).rejects.toThrow(); + + const statuses = manager.status([testServer]); + expect(statuses[0].state).toBe("error"); + expect(statuses[0].error).toContain("test-server"); + }); + + it("broken server retries after backoff", async () => { + let currentTime = 1000; + const brokenConn = makeBrokenConnection(); + const goodConn = makeMockConnection(); + let useGood = false; + + const deps: McpManagerDeps = { + spawn: () => ({ + stdin: { write: () => {} }, + stdout: { on: () => {} }, + pid: 1, + kill: () => {}, + }), + now: () => currentTime, + }; + + const factory = (_server: ResolvedMcpServer, _cwd: string) => { + const conn = useGood ? goodConn : brokenConn; + return { connection: conn, promise: Promise.resolve() }; + }; + + const manager = new McpManager(deps, factory); + + // First attempt fails + await expect(manager.ensureConnected(testServer, "/tmp")).rejects.toThrow(); + expect(manager.status([testServer])[0].state).toBe("error"); + + // Not enough time passed — still broken + currentTime += 29_000; + expect(manager.status([testServer])[0].state).toBe("error"); + + // After backoff — should allow retry + useGood = true; + currentTime += 2_000; + const statuses = manager.status([testServer]); + // After backoff, status() clears the broken entry + expect(statuses[0].state).toBe("disconnected"); + }); + + it("getClient returns undefined for unknown server", () => { + const manager = makeManager(() => { + return { connection: makeMockConnection(), promise: Promise.resolve() }; + }); + + expect(manager.getClient("nonexistent")).toBeUndefined(); + }); + + it("forwards the abort signal: a hanging initialize is interrupted", async () => { + // A connection whose initialize never resolves (a misbehaving / + // framing-incompatible server). Without a signal this hangs forever. + const hangingConn: Connection = { + send: () => new Promise(() => {}), + notify: () => {}, + onNotification: () => {}, + close: () => {}, + pid: 300, + }; + const manager = makeManager(() => { + return { connection: hangingConn, promise: Promise.resolve() }; + }); + + const controller = new AbortController(); + const connectPromise = manager.ensureConnected(testServer, "/tmp", controller.signal); + + // Abort mid-connect — the signal must reach initialize() and break it. + controller.abort(); + + await expect(connectPromise).rejects.toThrow(); + // The server is recorded as broken (not silently wedged). + expect(manager.status([testServer])[0].state).toBe("error"); + }); }); diff --git a/packages/mcp/src/manager.ts b/packages/mcp/src/manager.ts index a9c06de..1d719c6 100644 --- a/packages/mcp/src/manager.ts +++ b/packages/mcp/src/manager.ts @@ -12,194 +12,215 @@ import type { Connection, SpawnProcess } from "./transport.js"; import type { McpServerState, McpServerStatus, ResolvedMcpServer } from "./types.js"; export interface Logger { - readonly info: (msg: string, attrs?: Record<string, string | number | boolean | null>) => void; - readonly warn: (msg: string, attrs?: Record<string, string | number | boolean | null>) => void; - readonly error: (msg: string, attrs?: Record<string, unknown>) => void; + readonly info: (msg: string, attrs?: Record<string, string | number | boolean | null>) => void; + readonly warn: (msg: string, attrs?: Record<string, string | number | boolean | null>) => void; + readonly error: (msg: string, attrs?: Record<string, unknown>) => void; } export interface McpManagerDeps { - readonly spawn: SpawnProcess; - readonly logger?: Logger; - readonly now?: () => number; + readonly spawn: SpawnProcess; + readonly logger?: Logger; + readonly now?: () => number; } export type ConnectionFactory = ( - server: ResolvedMcpServer, - cwd: string, + server: ResolvedMcpServer, + cwd: string, ) => { connection: Connection; promise: Promise<void> }; type ClientEntry = { - readonly client: McpClient; - readonly server: ResolvedMcpServer; - readonly promise: Promise<void>; + readonly client: McpClient; + readonly server: ResolvedMcpServer; + readonly promise: Promise<void>; }; type BrokenEntry = { - readonly brokenAt: number; - readonly error: string; + readonly brokenAt: number; + readonly error: string; }; const BACKOFF_MS = 30_000; export class McpManager { - private clients = new Map<string, ClientEntry>(); - private broken = new Map<string, BrokenEntry>(); - private spawning = new Map<string, Promise<void>>(); - private readonly deps: McpManagerDeps; - private readonly connectionFactory: ConnectionFactory; - private readonly now: () => number; - - constructor(deps: McpManagerDeps, connectionFactory: ConnectionFactory) { - this.deps = deps; - this.connectionFactory = connectionFactory; - this.now = deps.now ?? Date.now; - } - - getClient(serverId: string): McpClient | undefined { - return this.clients.get(serverId)?.client; - } - - getServerState(serverId: string): McpServerState { - const brokenEntry = this.broken.get(serverId); - if (brokenEntry) { - const backoffElapsed = this.now() - brokenEntry.brokenAt >= BACKOFF_MS; - if (backoffElapsed) { - this.broken.delete(serverId); - } else { - return "error"; - } - } - - const entry = this.clients.get(serverId); - if (!entry) return "disconnected"; - - const state = entry.client.getState(); - if (state === "error") return "error"; - if (state === "connecting") return "connecting"; - if (state === "connected") return "connected"; - return "disconnected"; - } - - status(servers: readonly ResolvedMcpServer[]): McpServerStatus[] { - const results: McpServerStatus[] = []; - - for (const server of servers) { - const state = this.getServerState(server.id); - const entry = this.clients.get(server.id); - const brokenEntry = this.broken.get(server.id); - - const status: McpServerStatus = { - id: server.id, - state, - toolCount: entry?.client.getTools().length ?? 0, - }; - if (state === "error" && brokenEntry) { - (status as { error?: string }).error = brokenEntry.error; - } else if (state === "error" && entry?.client.getState() === "error") { - (status as { error?: string }).error = brokenEntry?.error ?? `${server.id}: client error`; - } - results.push(status); - } - - return results; - } - - async ensureConnected(server: ResolvedMcpServer, cwd: string): Promise<McpClient> { - const existing = this.clients.get(server.id); - if (existing && existing.client.getState() === "connected") { - return existing.client; - } - - const brokenEntry = this.broken.get(server.id); - if (brokenEntry) { - const backoffElapsed = this.now() - brokenEntry.brokenAt >= BACKOFF_MS; - if (!backoffElapsed) { - throw new Error(brokenEntry.error); - } - this.broken.delete(server.id); - } - - await this.spawnClient(server, cwd); - const entry = this.clients.get(server.id); - if (!entry) { - throw new Error(`Failed to spawn MCP client for ${server.id}`); - } - if (entry.client.getState() === "error") { - const brokenNow = this.broken.get(server.id); - throw new Error(brokenNow?.error ?? `${server.id}: client error`); - } - return entry.client; - } - - private async spawnClient(server: ResolvedMcpServer, cwd: string): Promise<void> { - const existingSpawn = this.spawning.get(server.id); - if (existingSpawn) return existingSpawn; - - const spawnPromise = this.doSpawn(server, cwd); - this.spawning.set(server.id, spawnPromise); - - try { - await spawnPromise; - } finally { - this.spawning.delete(server.id); - } - } - - private async doSpawn(server: ResolvedMcpServer, cwd: string): Promise<void> { - const { connection, promise } = this.connectionFactory(server, cwd); - - const client = new McpClient({ connection }); - - const entry: ClientEntry = { - client, - server, - promise: this.initClient(client, server, promise), - }; - - this.clients.set(server.id, entry); - await entry.promise; - - // If initialization failed, the client is in an error state and broken[] - // is already populated. Drop the half-created client (and reap its child - // process) so a later retry spawns fresh instead of returning a dead entry. - if (client.getState() === "error") { - this.clients.delete(server.id); - client.close(); - } - } - - private async initClient( - client: McpClient, - server: ResolvedMcpServer, - _transportPromise: Promise<void>, - ): Promise<void> { - try { - await client.initialize(); - await client.listTools(); - this.deps.logger?.info("MCP server connected", { - serverId: server.id, - toolCount: String(client.getTools().length), - }); - } catch (err: unknown) { - const message = err instanceof Error ? err.message : String(err); - this.broken.set(server.id, { - brokenAt: this.now(), - error: `${server.id}: ${message}`, - }); - this.deps.logger?.warn("MCP server failed to connect", { - serverId: server.id, - error: message, - }); - } - } - - shutdownAll(): void { - for (const [, entry] of this.clients) { - entry.client.close(); - } - this.clients.clear(); - this.broken.clear(); - this.spawning.clear(); - this.deps.logger?.info("All MCP servers shut down"); - } + private clients = new Map<string, ClientEntry>(); + private broken = new Map<string, BrokenEntry>(); + private spawning = new Map<string, Promise<void>>(); + private readonly deps: McpManagerDeps; + private readonly connectionFactory: ConnectionFactory; + private readonly now: () => number; + + constructor(deps: McpManagerDeps, connectionFactory: ConnectionFactory) { + this.deps = deps; + this.connectionFactory = connectionFactory; + this.now = deps.now ?? Date.now; + } + + getClient(serverId: string): McpClient | undefined { + return this.clients.get(serverId)?.client; + } + + getServerState(serverId: string): McpServerState { + const brokenEntry = this.broken.get(serverId); + if (brokenEntry) { + const backoffElapsed = this.now() - brokenEntry.brokenAt >= BACKOFF_MS; + if (backoffElapsed) { + this.broken.delete(serverId); + } else { + return "error"; + } + } + + const entry = this.clients.get(serverId); + if (!entry) return "disconnected"; + + const state = entry.client.getState(); + if (state === "error") return "error"; + if (state === "connecting") return "connecting"; + if (state === "connected") return "connected"; + return "disconnected"; + } + + status(servers: readonly ResolvedMcpServer[]): McpServerStatus[] { + const results: McpServerStatus[] = []; + + for (const server of servers) { + const state = this.getServerState(server.id); + const entry = this.clients.get(server.id); + const brokenEntry = this.broken.get(server.id); + + const status: McpServerStatus = { + id: server.id, + state, + toolCount: entry?.client.getTools().length ?? 0, + }; + if (state === "error" && brokenEntry) { + (status as { error?: string }).error = brokenEntry.error; + } else if (state === "error" && entry?.client.getState() === "error") { + (status as { error?: string }).error = brokenEntry?.error ?? `${server.id}: client error`; + } + results.push(status); + } + + return results; + } + + /** + * Ensure a client for `server` is connected, lazily spawning + handshaking on + * first access. The optional `signal` (the turn's abort signal) is forwarded + * into the `initialize`/`listTools` handshake so `POST /conversations/:id/stop` + * can interrupt a stuck connect; the operations are independently bounded by + * their own default timeout, so a misbehaving server cannot hang a turn even + * when no signal is supplied. + */ + async ensureConnected( + server: ResolvedMcpServer, + cwd: string, + signal?: AbortSignal, + ): Promise<McpClient> { + const existing = this.clients.get(server.id); + if (existing && existing.client.getState() === "connected") { + return existing.client; + } + + const brokenEntry = this.broken.get(server.id); + if (brokenEntry) { + const backoffElapsed = this.now() - brokenEntry.brokenAt >= BACKOFF_MS; + if (!backoffElapsed) { + throw new Error(brokenEntry.error); + } + this.broken.delete(server.id); + } + + await this.spawnClient(server, cwd, signal); + const entry = this.clients.get(server.id); + if (!entry) { + throw new Error(`Failed to spawn MCP client for ${server.id}`); + } + if (entry.client.getState() === "error") { + const brokenNow = this.broken.get(server.id); + throw new Error(brokenNow?.error ?? `${server.id}: client error`); + } + return entry.client; + } + + private async spawnClient( + server: ResolvedMcpServer, + cwd: string, + signal?: AbortSignal, + ): Promise<void> { + const existingSpawn = this.spawning.get(server.id); + if (existingSpawn) return existingSpawn; + + const spawnPromise = this.doSpawn(server, cwd, signal); + this.spawning.set(server.id, spawnPromise); + + try { + await spawnPromise; + } finally { + this.spawning.delete(server.id); + } + } + + private async doSpawn( + server: ResolvedMcpServer, + cwd: string, + signal?: AbortSignal, + ): Promise<void> { + const { connection, promise } = this.connectionFactory(server, cwd); + + const client = new McpClient({ connection }); + + const entry: ClientEntry = { + client, + server, + promise: this.initClient(client, server, promise, signal), + }; + + this.clients.set(server.id, entry); + await entry.promise; + + // If initialization failed, the client is in an error state and broken[] + // is already populated. Drop the half-created client (and reap its child + // process) so a later retry spawns fresh instead of returning a dead entry. + if (client.getState() === "error") { + this.clients.delete(server.id); + client.close(); + } + } + + private async initClient( + client: McpClient, + server: ResolvedMcpServer, + _transportPromise: Promise<void>, + signal?: AbortSignal, + ): Promise<void> { + try { + await client.initialize(signal); + await client.listTools(signal); + this.deps.logger?.info("MCP server connected", { + serverId: server.id, + toolCount: String(client.getTools().length), + }); + } catch (err: unknown) { + const message = err instanceof Error ? err.message : String(err); + this.broken.set(server.id, { + brokenAt: this.now(), + error: `${server.id}: ${message}`, + }); + this.deps.logger?.warn("MCP server failed to connect", { + serverId: server.id, + error: message, + }); + } + } + + shutdownAll(): void { + for (const [, entry] of this.clients) { + entry.client.close(); + } + this.clients.clear(); + this.broken.clear(); + this.spawning.clear(); + this.deps.logger?.info("All MCP servers shut down"); + } } diff --git a/packages/mcp/src/registry.test.ts b/packages/mcp/src/registry.test.ts index 4c88828..37cfea5 100644 --- a/packages/mcp/src/registry.test.ts +++ b/packages/mcp/src/registry.test.ts @@ -4,21 +4,21 @@ import { adaptTool, flattenContent, namespace } from "./registry.js"; import type { McpCallResult, McpContentItem, McpToolCaller, McpToolInfo } from "./types.js"; const mockSpan = { - id: "span-1", - log: {} as Logger, - setAttributes: () => {}, - addLink: () => {}, - child: () => mockSpan, - end: () => {}, + id: "span-1", + log: {} as Logger, + setAttributes: () => {}, + addLink: () => {}, + child: () => mockSpan, + end: () => {}, }; const mockLogger: Logger = { - info: () => {}, - warn: () => {}, - error: () => {}, - debug: () => {}, - child: () => mockLogger, - span: () => mockSpan, + info: () => {}, + warn: () => {}, + error: () => {}, + debug: () => {}, + child: () => mockLogger, + span: () => mockSpan, }; /** @@ -26,212 +26,212 @@ const mockLogger: Logger = { * It records calls and returns a configurable result. */ function makeCaller( - result: McpCallResult, + result: McpCallResult, ): McpToolCaller & { calls: Array<{ name: string; args: unknown }> } { - const calls: Array<{ name: string; args: unknown }> = []; - return { - calls, - callTool: async (name: string, args: unknown) => { - calls.push({ name, args }); - return result; - }, - }; + const calls: Array<{ name: string; args: unknown }> = []; + return { + calls, + callTool: async (name: string, args: unknown) => { + calls.push({ name, args }); + return result; + }, + }; } describe("namespace", () => { - it("produces <serverId>__<toolName>", () => { - expect(namespace("freecad", "create_object")).toBe("freecad__create_object"); - }); + it("produces <serverId>__<toolName>", () => { + expect(namespace("freecad", "create_object")).toBe("freecad__create_object"); + }); - it("handles serverId with special chars", () => { - expect(namespace("chrome-devtools", "navigate")).toBe("chrome-devtools__navigate"); - }); + it("handles serverId with special chars", () => { + expect(namespace("chrome-devtools", "navigate")).toBe("chrome-devtools__navigate"); + }); - it("handles empty toolName", () => { - expect(namespace("server", "")).toBe("server__"); - }); + it("handles empty toolName", () => { + expect(namespace("server", "")).toBe("server__"); + }); }); describe("flattenContent", () => { - it("flattens text content", () => { - const content: McpContentItem[] = [{ type: "text", text: "hello world" }]; - expect(flattenContent(content)).toBe("hello world"); - }); - - it("flattens image content", () => { - const content: McpContentItem[] = [ - { type: "image", data: "base64data", mimeType: "image/png" }, - ]; - expect(flattenContent(content)).toBe("[image: image/png]"); - }); - - it("flattens resource content with text", () => { - const content: McpContentItem[] = [ - { type: "resource", resource: { uri: "file:///test", text: "resource text" } }, - ]; - expect(flattenContent(content)).toBe("resource text"); - }); - - it("flattens resource content without text", () => { - const content: McpContentItem[] = [{ type: "resource", resource: { uri: "file:///test" } }]; - expect(flattenContent(content)).toBe("[resource: file:///test]"); - }); - - it("joins multiple items with newline", () => { - const content: McpContentItem[] = [ - { type: "text", text: "first" }, - { type: "text", text: "second" }, - ]; - expect(flattenContent(content)).toBe("first\nsecond"); - }); - - it("returns empty string for empty content", () => { - expect(flattenContent([])).toBe(""); - }); - - it("handles mixed content types", () => { - const content: McpContentItem[] = [ - { type: "text", text: "here is an image:" }, - { type: "image", data: "data", mimeType: "image/jpeg" }, - { type: "resource", resource: { uri: "file:///x", text: "some data" } }, - ]; - expect(flattenContent(content)).toBe("here is an image:\n[image: image/jpeg]\nsome data"); - }); + it("flattens text content", () => { + const content: McpContentItem[] = [{ type: "text", text: "hello world" }]; + expect(flattenContent(content)).toBe("hello world"); + }); + + it("flattens image content", () => { + const content: McpContentItem[] = [ + { type: "image", data: "base64data", mimeType: "image/png" }, + ]; + expect(flattenContent(content)).toBe("[image: image/png]"); + }); + + it("flattens resource content with text", () => { + const content: McpContentItem[] = [ + { type: "resource", resource: { uri: "file:///test", text: "resource text" } }, + ]; + expect(flattenContent(content)).toBe("resource text"); + }); + + it("flattens resource content without text", () => { + const content: McpContentItem[] = [{ type: "resource", resource: { uri: "file:///test" } }]; + expect(flattenContent(content)).toBe("[resource: file:///test]"); + }); + + it("joins multiple items with newline", () => { + const content: McpContentItem[] = [ + { type: "text", text: "first" }, + { type: "text", text: "second" }, + ]; + expect(flattenContent(content)).toBe("first\nsecond"); + }); + + it("returns empty string for empty content", () => { + expect(flattenContent([])).toBe(""); + }); + + it("handles mixed content types", () => { + const content: McpContentItem[] = [ + { type: "text", text: "here is an image:" }, + { type: "image", data: "data", mimeType: "image/jpeg" }, + { type: "resource", resource: { uri: "file:///x", text: "some data" } }, + ]; + expect(flattenContent(content)).toBe("here is an image:\n[image: image/jpeg]\nsome data"); + }); }); describe("adaptTool", () => { - it("maps inputSchema to ToolParameterSchema", () => { - const mcpTool: McpToolInfo = { - name: "create_obj", - description: "Create an object", - inputSchema: { - type: "object", - properties: { name: { type: "string", description: "Object name" } }, - required: ["name"], - }, - }; - const caller = makeCaller({ content: [] }); - const adapted = adaptTool("freecad", mcpTool, caller); - - expect(adapted.name).toBe("freecad__create_obj"); - expect(adapted.description).toBe("[freecad] Create an object"); - expect(adapted.parameters.type).toBe("object"); - expect(adapted.parameters.properties).toEqual({ - name: { type: "string", description: "Object name" }, - }); - expect(adapted.parameters.required).toEqual(["name"]); - expect(adapted.concurrencySafe).toBe(false); - }); - - it("execute proxies to callTool", async () => { - const mcpTool: McpToolInfo = { - name: "test_tool", - description: "Test", - inputSchema: { type: "object" }, - }; - const caller = makeCaller({ - content: [{ type: "text", text: "called" }], - isError: false, - }); - const adapted = adaptTool("server", mcpTool, caller); - - const result = await adapted.execute( - { input: "value" }, - { - toolCallId: "call-1", - onOutput: () => {}, - signal: new AbortController().signal, - log: mockLogger, - }, - ); - - expect(result.content).toBe("called"); - expect(result.isError).toBe(false); - expect(caller.calls).toEqual([{ name: "test_tool", args: { input: "value" } }]); - }); - - it("execute propagates isError", async () => { - const caller = makeCaller({ - content: [{ type: "text", text: "error occurred" }], - isError: true, - }); - - const mcpTool: McpToolInfo = { - name: "fail_tool", - description: "Fails", - inputSchema: { type: "object" }, - }; - const adapted = adaptTool("server", mcpTool, caller); - - const result = await adapted.execute( - {}, - { - toolCallId: "call-2", - onOutput: () => {}, - signal: new AbortController().signal, - log: mockLogger, - }, - ); - - expect(result.isError).toBe(true); - expect(result.content).toBe("error occurred"); - }); - - it("handles inputSchema without optional fields", () => { - const mcpTool: McpToolInfo = { - name: "simple", - description: "Simple tool", - inputSchema: { type: "object" }, - }; - const caller = makeCaller({ content: [] }); - const adapted = adaptTool("server", mcpTool, caller); - - expect(adapted.parameters.type).toBe("object"); - expect(adapted.parameters.properties).toBeUndefined(); - expect(adapted.parameters.required).toBeUndefined(); - expect(adapted.parameters.additionalProperties).toBeUndefined(); - }); - - it("preserves additionalProperties", () => { - const mcpTool: McpToolInfo = { - name: "flex", - description: "Flexible", - inputSchema: { - type: "object", - additionalProperties: true, - }, - }; - const caller = makeCaller({ content: [] }); - const adapted = adaptTool("server", mcpTool, caller); - - expect(adapted.parameters.additionalProperties).toBe(true); - }); - - it("execute flattens multi-item content", async () => { - const caller = makeCaller({ - content: [ - { type: "text", text: "summary" }, - { type: "image", data: "d", mimeType: "image/png" }, - ], - isError: false, - }); - const mcpTool: McpToolInfo = { - name: "multi", - description: "Multi", - inputSchema: { type: "object" }, - }; - const adapted = adaptTool("srv", mcpTool, caller); - - const result = await adapted.execute( - {}, - { - toolCallId: "c", - onOutput: () => {}, - signal: new AbortController().signal, - log: mockLogger, - }, - ); - - expect(result.content).toBe("summary\n[image: image/png]"); - }); + it("maps inputSchema to ToolParameterSchema", () => { + const mcpTool: McpToolInfo = { + name: "create_obj", + description: "Create an object", + inputSchema: { + type: "object", + properties: { name: { type: "string", description: "Object name" } }, + required: ["name"], + }, + }; + const caller = makeCaller({ content: [] }); + const adapted = adaptTool("freecad", mcpTool, caller); + + expect(adapted.name).toBe("freecad__create_obj"); + expect(adapted.description).toBe("[freecad] Create an object"); + expect(adapted.parameters.type).toBe("object"); + expect(adapted.parameters.properties).toEqual({ + name: { type: "string", description: "Object name" }, + }); + expect(adapted.parameters.required).toEqual(["name"]); + expect(adapted.concurrencySafe).toBe(false); + }); + + it("execute proxies to callTool", async () => { + const mcpTool: McpToolInfo = { + name: "test_tool", + description: "Test", + inputSchema: { type: "object" }, + }; + const caller = makeCaller({ + content: [{ type: "text", text: "called" }], + isError: false, + }); + const adapted = adaptTool("server", mcpTool, caller); + + const result = await adapted.execute( + { input: "value" }, + { + toolCallId: "call-1", + onOutput: () => {}, + signal: new AbortController().signal, + log: mockLogger, + }, + ); + + expect(result.content).toBe("called"); + expect(result.isError).toBe(false); + expect(caller.calls).toEqual([{ name: "test_tool", args: { input: "value" } }]); + }); + + it("execute propagates isError", async () => { + const caller = makeCaller({ + content: [{ type: "text", text: "error occurred" }], + isError: true, + }); + + const mcpTool: McpToolInfo = { + name: "fail_tool", + description: "Fails", + inputSchema: { type: "object" }, + }; + const adapted = adaptTool("server", mcpTool, caller); + + const result = await adapted.execute( + {}, + { + toolCallId: "call-2", + onOutput: () => {}, + signal: new AbortController().signal, + log: mockLogger, + }, + ); + + expect(result.isError).toBe(true); + expect(result.content).toBe("error occurred"); + }); + + it("handles inputSchema without optional fields", () => { + const mcpTool: McpToolInfo = { + name: "simple", + description: "Simple tool", + inputSchema: { type: "object" }, + }; + const caller = makeCaller({ content: [] }); + const adapted = adaptTool("server", mcpTool, caller); + + expect(adapted.parameters.type).toBe("object"); + expect(adapted.parameters.properties).toBeUndefined(); + expect(adapted.parameters.required).toBeUndefined(); + expect(adapted.parameters.additionalProperties).toBeUndefined(); + }); + + it("preserves additionalProperties", () => { + const mcpTool: McpToolInfo = { + name: "flex", + description: "Flexible", + inputSchema: { + type: "object", + additionalProperties: true, + }, + }; + const caller = makeCaller({ content: [] }); + const adapted = adaptTool("server", mcpTool, caller); + + expect(adapted.parameters.additionalProperties).toBe(true); + }); + + it("execute flattens multi-item content", async () => { + const caller = makeCaller({ + content: [ + { type: "text", text: "summary" }, + { type: "image", data: "d", mimeType: "image/png" }, + ], + isError: false, + }); + const mcpTool: McpToolInfo = { + name: "multi", + description: "Multi", + inputSchema: { type: "object" }, + }; + const adapted = adaptTool("srv", mcpTool, caller); + + const result = await adapted.execute( + {}, + { + toolCallId: "c", + onOutput: () => {}, + signal: new AbortController().signal, + log: mockLogger, + }, + ); + + expect(result.content).toBe("summary\n[image: image/png]"); + }); }); diff --git a/packages/mcp/src/registry.ts b/packages/mcp/src/registry.ts index 7c31c88..e1c95e9 100644 --- a/packages/mcp/src/registry.ts +++ b/packages/mcp/src/registry.ts @@ -9,71 +9,71 @@ */ import type { - ToolContract, - ToolExecuteContext, - ToolParameterSchema, - ToolResult, + ToolContract, + ToolExecuteContext, + ToolParameterSchema, + ToolResult, } from "@dispatch/kernel"; import type { McpContentItem, McpToolCaller, McpToolInfo } from "./types.js"; const NAMESPACE_SEP = "__"; export function namespace(serverId: string, toolName: string): string { - return `${serverId}${NAMESPACE_SEP}${toolName}`; + return `${serverId}${NAMESPACE_SEP}${toolName}`; } export function adaptTool( - serverId: string, - mcpTool: McpToolInfo, - caller: McpToolCaller, + serverId: string, + mcpTool: McpToolInfo, + caller: McpToolCaller, ): ToolContract { - const parameters: ToolParameterSchema = { - type: "object", - ...(mcpTool.inputSchema.properties !== undefined && { - properties: mcpTool.inputSchema.properties, - }), - ...(mcpTool.inputSchema.required !== undefined && { - required: mcpTool.inputSchema.required, - }), - ...(mcpTool.inputSchema.additionalProperties !== undefined && { - additionalProperties: mcpTool.inputSchema.additionalProperties, - }), - }; + const parameters: ToolParameterSchema = { + type: "object", + ...(mcpTool.inputSchema.properties !== undefined && { + properties: mcpTool.inputSchema.properties, + }), + ...(mcpTool.inputSchema.required !== undefined && { + required: mcpTool.inputSchema.required, + }), + ...(mcpTool.inputSchema.additionalProperties !== undefined && { + additionalProperties: mcpTool.inputSchema.additionalProperties, + }), + }; - return { - name: namespace(serverId, mcpTool.name), - description: `[${serverId}] ${mcpTool.description}`, - parameters, - concurrencySafe: false, - execute: async (args: unknown, ctx: ToolExecuteContext): Promise<ToolResult> => { - const result = await caller.callTool(mcpTool.name, args, ctx.signal); - const toolResult: ToolResult = { - content: flattenContent(result.content), - }; - if (result.isError !== undefined) { - (toolResult as { isError?: boolean }).isError = result.isError; - } - return toolResult; - }, - }; + return { + name: namespace(serverId, mcpTool.name), + description: `[${serverId}] ${mcpTool.description}`, + parameters, + concurrencySafe: false, + execute: async (args: unknown, ctx: ToolExecuteContext): Promise<ToolResult> => { + const result = await caller.callTool(mcpTool.name, args, ctx.signal); + const toolResult: ToolResult = { + content: flattenContent(result.content), + }; + if (result.isError !== undefined) { + (toolResult as { isError?: boolean }).isError = result.isError; + } + return toolResult; + }, + }; } export function flattenContent(content: readonly McpContentItem[]): string { - if (content.length === 0) return ""; + if (content.length === 0) return ""; - const parts: string[] = []; - for (const item of content) { - if (item.type === "text" && item.text !== undefined) { - parts.push(item.text); - } else if (item.type === "image" && item.mimeType !== undefined) { - parts.push(`[image: ${item.mimeType}]`); - } else if (item.type === "resource" && item.resource !== undefined) { - if (item.resource.text !== undefined) { - parts.push(item.resource.text); - } else { - parts.push(`[resource: ${item.resource.uri}]`); - } - } - } - return parts.join("\n"); + const parts: string[] = []; + for (const item of content) { + if (item.type === "text" && item.text !== undefined) { + parts.push(item.text); + } else if (item.type === "image" && item.mimeType !== undefined) { + parts.push(`[image: ${item.mimeType}]`); + } else if (item.type === "resource" && item.resource !== undefined) { + if (item.resource.text !== undefined) { + parts.push(item.resource.text); + } else { + parts.push(`[resource: ${item.resource.uri}]`); + } + } + } + return parts.join("\n"); } diff --git a/packages/mcp/src/rpc.test.ts b/packages/mcp/src/rpc.test.ts index 9ffb121..34b6af4 100644 --- a/packages/mcp/src/rpc.test.ts +++ b/packages/mcp/src/rpc.test.ts @@ -2,116 +2,115 @@ import { describe, expect, it } from "vitest"; import { JsonRpcClient } from "./rpc.js"; function makeClient(): { - client: JsonRpcClient; - written: Uint8Array[]; - feedMessage: (msg: unknown) => void; + client: JsonRpcClient; + written: Uint8Array[]; + feedMessage: (msg: unknown) => void; } { - const written: Uint8Array[] = []; - const client = new JsonRpcClient((bytes) => { - written.push(bytes); - }); - return { - client, - written, - feedMessage: (msg: unknown) => { - client.handleMessage(JSON.stringify(msg)); - }, - }; + const written: Uint8Array[] = []; + const client = new JsonRpcClient((bytes) => { + written.push(bytes); + }); + return { + client, + written, + feedMessage: (msg: unknown) => { + client.handleMessage(JSON.stringify(msg)); + }, + }; } describe("JsonRpcClient", () => { - it("request returns result", async () => { - const { client, feedMessage } = makeClient(); + it("request returns result", async () => { + const { client, feedMessage } = makeClient(); - const resultPromise = client.request("initialize", { protocolVersion: "2025-11-25" }); + const resultPromise = client.request("initialize", { protocolVersion: "2025-11-25" }); - feedMessage({ jsonrpc: "2.0", id: 1, result: { protocolVersion: "2025-11-25" } }); + feedMessage({ jsonrpc: "2.0", id: 1, result: { protocolVersion: "2025-11-25" } }); - const result = await resultPromise; - expect(result).toEqual({ protocolVersion: "2025-11-25" }); - }); + const result = await resultPromise; + expect(result).toEqual({ protocolVersion: "2025-11-25" }); + }); - it("request rejects on error response", async () => { - const { client, feedMessage } = makeClient(); + it("request rejects on error response", async () => { + const { client, feedMessage } = makeClient(); - const resultPromise = client.request("bad-method"); + const resultPromise = client.request("bad-method"); - feedMessage({ - jsonrpc: "2.0", - id: 1, - error: { code: -32601, message: "Method not found" }, - }); + feedMessage({ + jsonrpc: "2.0", + id: 1, + error: { code: -32601, message: "Method not found" }, + }); - await expect(resultPromise).rejects.toThrow("Method not found"); - }); + await expect(resultPromise).rejects.toThrow("Method not found"); + }); - it("notify sends without expecting response", () => { - const { client, written } = makeClient(); + it("notify sends without expecting response", () => { + const { client, written } = makeClient(); - client.notify("notifications/initialized", {}); + client.notify("notifications/initialized", {}); - expect(written.length).toBe(1); - const sent = new TextDecoder().decode(written[0]); - expect(sent).toContain('"method":"notifications/initialized"'); - expect(sent).not.toContain('"id"'); - }); + expect(written.length).toBe(1); + const sent = new TextDecoder().decode(written[0]); + expect(sent).toContain('"method":"notifications/initialized"'); + expect(sent).not.toContain('"id"'); + }); - it("onNotification fires for matching method", () => { - const { client, feedMessage } = makeClient(); + it("onNotification fires for matching method", () => { + const { client, feedMessage } = makeClient(); - let received: unknown = "unset"; - client.onNotification("notifications/tools/list_changed", (params) => { - received = params; - }); + let received: unknown = "unset"; + client.onNotification("notifications/tools/list_changed", (params) => { + received = params; + }); - feedMessage({ jsonrpc: "2.0", method: "notifications/tools/list_changed", params: { a: 1 } }); + feedMessage({ jsonrpc: "2.0", method: "notifications/tools/list_changed", params: { a: 1 } }); - expect(received).toEqual({ a: 1 }); - }); + expect(received).toEqual({ a: 1 }); + }); - it("pending request rejected on close", async () => { - const { client } = makeClient(); + it("pending request rejected on close", async () => { + const { client } = makeClient(); - const resultPromise = client.request("slow-method"); + const resultPromise = client.request("slow-method"); - client.close(); + client.close(); - await expect(resultPromise).rejects.toThrow("Connection closed"); - }); + await expect(resultPromise).rejects.toThrow("Connection closed"); + }); - it("incremental request ids", () => { - const { client, written } = makeClient(); + it("incremental request ids", () => { + const { client, written } = makeClient(); - client.request("a"); - client.request("b"); - client.request("c"); + client.request("a"); + client.request("b"); + client.request("c"); - expect(written.length).toBe(3); - // Extract JSON body from Content-Length framed messages - const parse = (bytes: Uint8Array): { id: number } => { - const text = new TextDecoder().decode(bytes); - const bodyStart = text.indexOf("\r\n\r\n") + 4; - return JSON.parse(text.slice(bodyStart)) as { id: number }; - }; - const msg1 = parse(written[0]); - const msg2 = parse(written[1]); - const msg3 = parse(written[2]); - expect(msg1.id).toBe(1); - expect(msg2.id).toBe(2); - expect(msg3.id).toBe(3); - }); + expect(written.length).toBe(3); + // Outgoing messages are newline-delimited JSON (current MCP spec framing). + const parse = (bytes: Uint8Array): { id: number } => { + const text = new TextDecoder().decode(bytes); + return JSON.parse(text.replace(/\r?\n$/, "")) as { id: number }; + }; + const msg1 = parse(written[0]); + const msg2 = parse(written[1]); + const msg3 = parse(written[2]); + expect(msg1.id).toBe(1); + expect(msg2.id).toBe(2); + expect(msg3.id).toBe(3); + }); - it("notify after close is silently dropped", () => { - const { client, written } = makeClient(); - client.close(); - const count = written.length; - client.notify("test"); - expect(written.length).toBe(count); - }); + it("notify after close is silently dropped", () => { + const { client, written } = makeClient(); + client.close(); + const count = written.length; + client.notify("test"); + expect(written.length).toBe(count); + }); - it("request after close rejects immediately", async () => { - const { client } = makeClient(); - client.close(); - await expect(client.request("test")).rejects.toThrow("Connection closed"); - }); + it("request after close rejects immediately", async () => { + const { client } = makeClient(); + client.close(); + await expect(client.request("test")).rejects.toThrow("Connection closed"); + }); }); diff --git a/packages/mcp/src/rpc.ts b/packages/mcp/src/rpc.ts index eef8d73..aa89159 100644 --- a/packages/mcp/src/rpc.ts +++ b/packages/mcp/src/rpc.ts @@ -10,94 +10,94 @@ import { encode } from "./framing.js"; export type WriteFn = (bytes: Uint8Array) => void; export interface PendingRequest { - readonly resolve: (value: unknown) => void; - readonly reject: (reason: unknown) => void; + readonly resolve: (value: unknown) => void; + readonly reject: (reason: unknown) => void; } export type NotificationHandler = (params: unknown) => void; export interface JsonRpcMessage { - readonly jsonrpc: "2.0"; - readonly id?: number | string | undefined; - readonly method?: string | undefined; - readonly params?: unknown; - readonly result?: unknown; - readonly error?: - | { readonly code: number; readonly message: string; readonly data?: unknown } - | undefined; + readonly jsonrpc: "2.0"; + readonly id?: number | string | undefined; + readonly method?: string | undefined; + readonly params?: unknown; + readonly result?: unknown; + readonly error?: + | { readonly code: number; readonly message: string; readonly data?: unknown } + | undefined; } export class JsonRpcClient { - private nextId = 1; - private pending = new Map<number | string, PendingRequest>(); - private notificationHandlers = new Map<string, NotificationHandler>(); - private write: WriteFn; - private closed = false; - - constructor(write: WriteFn) { - this.write = write; - } - - request(method: string, params?: unknown): Promise<unknown> { - if (this.closed) { - return Promise.reject(new Error("Connection closed")); - } - const id = this.nextId++; - const msg: JsonRpcMessage = { jsonrpc: "2.0", id, method, params }; - return new Promise((resolve, reject) => { - this.pending.set(id, { resolve, reject }); - this.sendMessage(msg); - }); - } - - notify(method: string, params?: unknown): void { - if (this.closed) return; - const msg: JsonRpcMessage = { jsonrpc: "2.0", method, params }; - this.sendMessage(msg); - } - - onNotification(method: string, handler: NotificationHandler): void { - this.notificationHandlers.set(method, handler); - } - - handleMessage(json: string): void { - const msg = JSON.parse(json) as JsonRpcMessage; - const { id, method } = msg; - - if (method !== undefined && id === undefined) { - this.handleIncomingNotification(method, msg.params); - } else if (id !== undefined) { - this.handleResponse(id, msg); - } - } - - private sendMessage(msg: JsonRpcMessage): void { - this.write(encode(JSON.stringify(msg))); - } - - private handleResponse(id: number | string, msg: JsonRpcMessage): void { - const entry = this.pending.get(id); - if (!entry) return; - this.pending.delete(id); - if (msg.error) { - entry.reject(new Error(msg.error.message)); - } else { - entry.resolve(msg.result); - } - } - - private handleIncomingNotification(method: string, params: unknown): void { - const handler = this.notificationHandlers.get(method); - if (handler) { - handler(params); - } - } - - close(): void { - this.closed = true; - for (const entry of this.pending.values()) { - entry.reject(new Error("Connection closed")); - } - this.pending.clear(); - } + private nextId = 1; + private pending = new Map<number | string, PendingRequest>(); + private notificationHandlers = new Map<string, NotificationHandler>(); + private write: WriteFn; + private closed = false; + + constructor(write: WriteFn) { + this.write = write; + } + + request(method: string, params?: unknown): Promise<unknown> { + if (this.closed) { + return Promise.reject(new Error("Connection closed")); + } + const id = this.nextId++; + const msg: JsonRpcMessage = { jsonrpc: "2.0", id, method, params }; + return new Promise((resolve, reject) => { + this.pending.set(id, { resolve, reject }); + this.sendMessage(msg); + }); + } + + notify(method: string, params?: unknown): void { + if (this.closed) return; + const msg: JsonRpcMessage = { jsonrpc: "2.0", method, params }; + this.sendMessage(msg); + } + + onNotification(method: string, handler: NotificationHandler): void { + this.notificationHandlers.set(method, handler); + } + + handleMessage(json: string): void { + const msg = JSON.parse(json) as JsonRpcMessage; + const { id, method } = msg; + + if (method !== undefined && id === undefined) { + this.handleIncomingNotification(method, msg.params); + } else if (id !== undefined) { + this.handleResponse(id, msg); + } + } + + private sendMessage(msg: JsonRpcMessage): void { + this.write(encode(JSON.stringify(msg))); + } + + private handleResponse(id: number | string, msg: JsonRpcMessage): void { + const entry = this.pending.get(id); + if (!entry) return; + this.pending.delete(id); + if (msg.error) { + entry.reject(new Error(msg.error.message)); + } else { + entry.resolve(msg.result); + } + } + + private handleIncomingNotification(method: string, params: unknown): void { + const handler = this.notificationHandlers.get(method); + if (handler) { + handler(params); + } + } + + close(): void { + this.closed = true; + for (const entry of this.pending.values()) { + entry.reject(new Error("Connection closed")); + } + this.pending.clear(); + } } diff --git a/packages/mcp/src/timeout.test.ts b/packages/mcp/src/timeout.test.ts new file mode 100644 index 0000000..38961d0 --- /dev/null +++ b/packages/mcp/src/timeout.test.ts @@ -0,0 +1,105 @@ +import { describe, expect, it } from "vitest"; +import { McpTimeoutError, withTimeout } from "./timeout.js"; + +/** + * `withTimeout` uses a single `setTimeout` (the only edge). Tests drive it + * deterministically via `AbortController` for the abort path, and use real + * (tiny) timers for the timeout path — vitest's async scheduler is fast enough + * at single-digit ms that this is stable. + */ + +describe("withTimeout", () => { + it("resolves with the underlying value when the promise settles first", async () => { + const result = await withTimeout(Promise.resolve("ok"), "initialize", 1000); + expect(result).toBe("ok"); + }); + + it("rejects with the underlying error when the promise rejects first", async () => { + await expect( + withTimeout(Promise.reject(new Error("boom")), "initialize", 1000), + ).rejects.toThrow("boom"); + }); + + it("rejects with McpTimeoutError when the timeout fires first", async () => { + const never = new Promise<string>(() => {}); // never settles + const p = withTimeout(never, "initialize", 20); + await expect(p).rejects.toBeInstanceOf(McpTimeoutError); + await expect(p).rejects.toMatchObject({ method: "initialize", timeoutMs: 20 }); + }); + + it("cleans up the timer after the promise settles (no leak)", async () => { + // If the timer were not cleared, vitest would report an unhandled timer; + // resolving quickly should leave nothing pending. + const result = await withTimeout(Promise.resolve(1), "tools/list", 50_000); + expect(result).toBe(1); + }); + + it("rejects with Error('Aborted') when the signal is already aborted", async () => { + const controller = new AbortController(); + controller.abort(); + await expect( + withTimeout(new Promise(() => {}), "initialize", 50_000, controller.signal), + ).rejects.toThrow("Aborted"); + }); + + it("rejects with Error('Aborted') when the signal aborts mid-flight", async () => { + const controller = new AbortController(); + const never = new Promise<string>(() => {}); + const p = withTimeout(never, "initialize", 50_000, controller.signal); + controller.abort(); + await expect(p).rejects.toThrow("Aborted"); + }); + + it("abort beats timeout (immediate cancellation)", async () => { + const controller = new AbortController(); + const never = new Promise<string>(() => {}); + const p = withTimeout(never, "initialize", 50_000, controller.signal); + controller.abort(); + await expect(p).rejects.toThrow("Aborted"); + }); + + it("passes through when timeout is disabled (0) and no signal", async () => { + const result = await withTimeout(Promise.resolve("passthrough"), "initialize", 0); + expect(result).toBe("passthrough"); + }); + + it("passes through when timeout is Infinity and no signal", async () => { + const result = await withTimeout(Promise.resolve(42), "tools/list", Number.POSITIVE_INFINITY); + expect(result).toBe(42); + }); + + it("still honors the signal when timeout is disabled", async () => { + const controller = new AbortController(); + const p = withTimeout(new Promise<string>(() => {}), "initialize", 0, controller.signal); + controller.abort(); + await expect(p).rejects.toThrow("Aborted"); + }); + + it("removes the abort listener after the promise settles", async () => { + const controller = new AbortController(); + const addSpy = controller.signal.addEventListener.bind(controller.signal); + let added = 0; + const spiedAdd = (...args: unknown[]) => { + added++; + return (addSpy as (...a: unknown[]) => void)(...args); + }; + controller.signal.addEventListener = spiedAdd as typeof controller.signal.addEventListener; + + await withTimeout(Promise.resolve("ok"), "initialize", 50_000, controller.signal); + // The resolve path cleans up: the listener was added then removed. The + // important assertion is that resolution works and no abort fires after. + expect(added).toBe(1); + // Aborting after settlement must NOT reject the already-settled promise. + controller.abort(); + }); +}); + +describe("McpTimeoutError", () => { + it("carries method + timeoutMs and a descriptive message", () => { + const err = new McpTimeoutError("initialize", 5000); + expect(err.method).toBe("initialize"); + expect(err.timeoutMs).toBe(5000); + expect(err.message).toBe("MCP initialize timed out after 5000ms"); + expect(err.name).toBe("McpTimeoutError"); + }); +}); diff --git a/packages/mcp/src/timeout.ts b/packages/mcp/src/timeout.ts new file mode 100644 index 0000000..294bfc5 --- /dev/null +++ b/packages/mcp/src/timeout.ts @@ -0,0 +1,114 @@ +/** + * Timeout + abort helper for MCP operations. + * + * A single misbehaving or framing-incompatible MCP server must never be able to + * hang an agent turn indefinitely: the JSON-RPC `initialize` / `tools/list` + * requests are awaited in-band during the per-turn tools filter, so a server + * that never responds would block the whole turn forever. `withTimeout` bounds + * any such awaited operation by BOTH a timeout (always) and an optional + * `AbortSignal` (so the turn's stop can interrupt an in-flight connect). + * + * Edge effect: uses `setTimeout` (the clock is the only I/O). The timer + the + * signal listener are always cleaned up on settlement, so a resolved operation + * never leaks a pending timer. The underlying promise ALWAYS has a handler + * attached (even when abort/timeout wins first), so it can never surface as an + * unhandled rejection. Mocking the OUTERMOST edge (real clock) is fine; tests + * drive this via `AbortController` (deterministic) rather than the timer. + */ + +/** Default per-operation timeout for MCP handshake/tool-list requests (ms). */ +export const MCP_DEFAULT_TIMEOUT_MS = 30_000; + +/** + * Backstop timeout bounding the ENTIRE per-turn MCP connect phase (spawn + + * initialize + listTools across all configured servers), applied by the tools + * filter. A misbehaving or framing-incompatible server cannot hang a turn + * longer than this; on expiry the filter degrades gracefully (skips MCP tools + * for that turn) instead of blocking the turn. Generous enough to absorb a + * legitimate slow server startup (e.g. a browser-launching MCP server). + */ +export const MCP_CONNECT_TIMEOUT_MS = 30_000; + +/** + * Raised when an MCP operation does not settle within its timeout. Distinct + * from a plain `Error` so callers (the manager's broken-state tracking, tests) + * can tell a timeout/incompatibility apart from a server-reported RPC error. + */ +export class McpTimeoutError extends Error { + readonly method: string; + readonly timeoutMs: number; + constructor(method: string, timeoutMs: number) { + super(`MCP ${method} timed out after ${timeoutMs}ms`); + this.name = "McpTimeoutError"; + this.method = method; + this.timeoutMs = timeoutMs; + } +} + +/** + * Race `promise` against a timeout (always) and an optional `AbortSignal`. + * Resolves/rejects with `promise`'s outcome if it settles first; rejects with + * `McpTimeoutError` on timeout, or `Error("Aborted")` if `signal` aborts first. + * + * @param method JSON-RPC method name (for the timeout message). + * @param timeoutMs Milliseconds before a timeout is raised. Pass `0` or + * `Infinity` to disable the timeout (only the `signal` then bounds the call). + * @param signal Optional abort signal — typically the turn's signal, so + * `POST /conversations/:id/stop` can interrupt a stuck connect immediately. + */ +export function withTimeout<T>( + promise: Promise<T>, + method: string, + timeoutMs: number, + signal?: AbortSignal, +): Promise<T> { + // No timeout and no signal → pass straight through (nothing to race). + const hasTimeout = timeoutMs > 0 && Number.isFinite(timeoutMs); + if (!hasTimeout && signal === undefined) { + return promise; + } + + return new Promise<T>((resolve, reject) => { + let settled = false; + + const finish = (action: () => void): void => { + if (settled) return; + settled = true; + cleanup(); + action(); + }; + + let timer: ReturnType<typeof setTimeout> | undefined; + const onAbort = (): void => { + finish(() => reject(new Error("Aborted"))); + }; + + const cleanup = (): void => { + if (timer !== undefined) clearTimeout(timer); + if (signal !== undefined) signal.removeEventListener("abort", onAbort); + }; + + if (hasTimeout) { + timer = setTimeout( + () => finish(() => reject(new McpTimeoutError(method, timeoutMs))), + timeoutMs, + ); + } + if (signal !== undefined) { + if (signal.aborted) { + // Already aborted: abort wins immediately. The `.then` below still + // attaches a handler so the underlying promise never rejects unhandled. + finish(() => reject(new Error("Aborted"))); + } else { + signal.addEventListener("abort", onAbort, { once: true }); + } + } + + // Always attach handlers so the underlying promise is never unhandled — + // even when abort/timeout already won (settled), this is a no-op. + promise.then( + (value) => finish(() => resolve(value)), + (err: unknown) => finish(() => reject(err)), + ); + }); +} diff --git a/packages/mcp/src/transport.test.ts b/packages/mcp/src/transport.test.ts index b369e74..8c69ad2 100644 --- a/packages/mcp/src/transport.test.ts +++ b/packages/mcp/src/transport.test.ts @@ -9,146 +9,171 @@ import { createStdioTransport } from "./transport.js"; * what we wrote to the child's stdin (our outgoing framed messages). */ function makePipe(): { - process: SpawnedProcess; - emitStdout: (data: Uint8Array) => void; - emitEnd: () => void; - writtenToStdin: () => Uint8Array[]; - killed: () => boolean; + process: SpawnedProcess; + emitStdout: (data: Uint8Array) => void; + emitEnd: () => void; + writtenToStdin: () => Uint8Array[]; + killed: () => boolean; } { - const dataListeners: Array<(data: Uint8Array) => void> = []; - const endListeners: Array<() => void> = []; - const stdinWrites: Uint8Array[] = []; - let killed = false; - - const process: SpawnedProcess = { - stdin: { - write: (bytes: Uint8Array) => { - stdinWrites.push(bytes); - }, - }, - stdout: { - on: (event: string, cb: (data: Uint8Array) => void) => { - if (event === "data") dataListeners.push(cb); - else if (event === "end") endListeners.push(cb as unknown as () => void); - }, - }, - pid: 12345, - kill: () => { - killed = true; - }, - }; - - return { - process, - emitStdout: (data: Uint8Array) => { - for (const cb of dataListeners) cb(data); - }, - emitEnd: () => { - for (const cb of endListeners) cb(); - }, - writtenToStdin: () => stdinWrites, - killed: () => killed, - }; + const dataListeners: Array<(data: Uint8Array) => void> = []; + const endListeners: Array<() => void> = []; + const stdinWrites: Uint8Array[] = []; + let killed = false; + + const process: SpawnedProcess = { + stdin: { + write: (bytes: Uint8Array) => { + stdinWrites.push(bytes); + }, + }, + stdout: { + on: (event: string, cb: (data: Uint8Array) => void) => { + if (event === "data") dataListeners.push(cb); + else if (event === "end") endListeners.push(cb as unknown as () => void); + }, + }, + pid: 12345, + kill: () => { + killed = true; + }, + }; + + return { + process, + emitStdout: (data: Uint8Array) => { + for (const cb of dataListeners) cb(data); + }, + emitEnd: () => { + for (const cb of endListeners) cb(); + }, + writtenToStdin: () => stdinWrites, + killed: () => killed, + }; } describe("createStdioTransport", () => { - it("creates connection with correct pid", () => { - const pair = makePipe(); - const spawn: SpawnProcess = () => pair.process; - - const { connection } = createStdioTransport({ spawn, command: ["test-server"] }, "/tmp"); - - expect(connection.pid).toBe(12345); - connection.close(); - }); - - it("connection sends framed messages via stdin", () => { - const pair = makePipe(); - const spawn: SpawnProcess = () => pair.process; - - const { connection } = createStdioTransport({ spawn, command: ["test"] }, "/tmp"); - - connection.notify("test/method", { key: "value" }); - - const writes = pair.writtenToStdin(); - expect(writes.length).toBe(1); - const text = new TextDecoder().decode(writes[0]); - expect(text).toContain("Content-Length:"); - expect(text).toContain('"method":"test/method"'); - connection.close(); - }); - - it("close kills the child process", () => { - const pair = makePipe(); - const spawn: SpawnProcess = () => pair.process; - - const { connection } = createStdioTransport({ spawn, command: ["test"] }, "/tmp"); - - connection.close(); - expect(pair.killed()).toBe(true); - }); - - it("pipes stdout through framing: a notification triggers onNotification", async () => { - const pair = makePipe(); - const spawn: SpawnProcess = () => pair.process; - - const { connection } = createStdioTransport({ spawn, command: ["test"] }, "/tmp"); - - let received: unknown = null; - connection.onNotification("notifications/tools/list_changed", (params) => { - received = params; - }); - - // Simulate the server writing a framed notification to stdout. - const notification = JSON.stringify({ - jsonrpc: "2.0", - method: "notifications/tools/list_changed", - params: { reason: "tools added" }, - }); - pair.emitStdout(encode(notification)); - - // onNotification is invoked synchronously inside the data handler. - expect(received).toEqual({ reason: "tools added" }); - connection.close(); - }); - - it("pipes stdout through framing: a response resolves a request", async () => { - const pair = makePipe(); - const spawn: SpawnProcess = () => pair.process; - - const { connection } = createStdioTransport({ spawn, command: ["test"] }, "/tmp"); - - const resultPromise = connection.send("tools/list"); - - // The request was framed and written to stdin; respond via stdout. - const response = JSON.stringify({ - jsonrpc: "2.0", - id: 1, - result: { tools: [{ name: "t", description: "d", inputSchema: { type: "object" } }] }, - }); - pair.emitStdout(encode(response)); - - const result = await resultPromise; - expect(result).toEqual({ - tools: [{ name: "t", description: "d", inputSchema: { type: "object" } }], - }); - connection.close(); - }); - - it("handles a frame split across two stdout chunks", async () => { - const pair = makePipe(); - const spawn: SpawnProcess = () => pair.process; - - const { connection } = createStdioTransport({ spawn, command: ["test"] }, "/tmp"); - - const resultPromise = connection.send("ping"); - - const response = encode(JSON.stringify({ jsonrpc: "2.0", id: 1, result: { ok: true } })); - const mid = Math.floor(response.length / 2); - pair.emitStdout(response.slice(0, mid)); - pair.emitStdout(response.slice(mid)); - - await expect(resultPromise).resolves.toEqual({ ok: true }); - connection.close(); - }); + it("creates connection with correct pid", () => { + const pair = makePipe(); + const spawn: SpawnProcess = () => pair.process; + + const { connection } = createStdioTransport({ spawn, command: ["test-server"] }, "/tmp"); + + expect(connection.pid).toBe(12345); + connection.close(); + }); + + it("connection sends newline-delimited messages via stdin (current MCP spec)", () => { + const pair = makePipe(); + const spawn: SpawnProcess = () => pair.process; + + const { connection } = createStdioTransport({ spawn, command: ["test"] }, "/tmp"); + + connection.notify("test/method", { key: "value" }); + + const writes = pair.writtenToStdin(); + expect(writes.length).toBe(1); + const text = new TextDecoder().decode(writes[0]); + // Outgoing framing is newline-delimited JSON (not Content-Length). + expect(text).not.toContain("Content-Length:"); + expect(text).toContain('"method":"test/method"'); + expect(text.endsWith("\n")).toBe(true); + connection.close(); + }); + + it("decodes a newline-delimited server response (auto-detect)", async () => { + const pair = makePipe(); + const spawn: SpawnProcess = () => pair.process; + + const { connection } = createStdioTransport({ spawn, command: ["test"] }, "/tmp"); + + const resultPromise = connection.send("tools/list"); + + // Server responds with newline-delimited JSON (e.g. chrome-devtools-mcp). + const response = `${JSON.stringify({ + jsonrpc: "2.0", + id: 1, + result: { tools: [{ name: "t", description: "d", inputSchema: { type: "object" } }] }, + })}\n`; + pair.emitStdout(new TextEncoder().encode(response)); + + const result = await resultPromise; + expect(result).toEqual({ + tools: [{ name: "t", description: "d", inputSchema: { type: "object" } }], + }); + connection.close(); + }); + + it("close kills the child process", () => { + const pair = makePipe(); + const spawn: SpawnProcess = () => pair.process; + + const { connection } = createStdioTransport({ spawn, command: ["test"] }, "/tmp"); + + connection.close(); + expect(pair.killed()).toBe(true); + }); + + it("pipes stdout through framing: a notification triggers onNotification", async () => { + const pair = makePipe(); + const spawn: SpawnProcess = () => pair.process; + + const { connection } = createStdioTransport({ spawn, command: ["test"] }, "/tmp"); + + let received: unknown = null; + connection.onNotification("notifications/tools/list_changed", (params) => { + received = params; + }); + + // Simulate the server writing a framed notification to stdout. + const notification = JSON.stringify({ + jsonrpc: "2.0", + method: "notifications/tools/list_changed", + params: { reason: "tools added" }, + }); + pair.emitStdout(encode(notification)); + + // onNotification is invoked synchronously inside the data handler. + expect(received).toEqual({ reason: "tools added" }); + connection.close(); + }); + + it("pipes stdout through framing: a response resolves a request", async () => { + const pair = makePipe(); + const spawn: SpawnProcess = () => pair.process; + + const { connection } = createStdioTransport({ spawn, command: ["test"] }, "/tmp"); + + const resultPromise = connection.send("tools/list"); + + // The request was framed and written to stdin; respond via stdout. + const response = JSON.stringify({ + jsonrpc: "2.0", + id: 1, + result: { tools: [{ name: "t", description: "d", inputSchema: { type: "object" } }] }, + }); + pair.emitStdout(encode(response)); + + const result = await resultPromise; + expect(result).toEqual({ + tools: [{ name: "t", description: "d", inputSchema: { type: "object" } }], + }); + connection.close(); + }); + + it("handles a frame split across two stdout chunks", async () => { + const pair = makePipe(); + const spawn: SpawnProcess = () => pair.process; + + const { connection } = createStdioTransport({ spawn, command: ["test"] }, "/tmp"); + + const resultPromise = connection.send("ping"); + + const response = encode(JSON.stringify({ jsonrpc: "2.0", id: 1, result: { ok: true } })); + const mid = Math.floor(response.length / 2); + pair.emitStdout(response.slice(0, mid)); + pair.emitStdout(response.slice(mid)); + + await expect(resultPromise).resolves.toEqual({ ok: true }); + connection.close(); + }); }); diff --git a/packages/mcp/src/transport.ts b/packages/mcp/src/transport.ts index b1e3ec5..492879f 100644 --- a/packages/mcp/src/transport.ts +++ b/packages/mcp/src/transport.ts @@ -7,95 +7,95 @@ import { FrameDecoder } from "./framing.js"; import { JsonRpcClient, type WriteFn } from "./rpc.js"; export interface SpawnedProcess { - readonly stdin: { readonly write: (bytes: Uint8Array) => void }; - readonly stdout: - | AsyncIterable<Uint8Array> - | { readonly on: (event: string, cb: (data: Uint8Array) => void) => void }; - readonly stderr?: - | AsyncIterable<Uint8Array> - | { readonly on: (event: string, cb: (data: Uint8Array) => void) => void } - | undefined; - readonly pid: number | undefined; - readonly kill: () => void; + readonly stdin: { readonly write: (bytes: Uint8Array) => void }; + readonly stdout: + | AsyncIterable<Uint8Array> + | { readonly on: (event: string, cb: (data: Uint8Array) => void) => void }; + readonly stderr?: + | AsyncIterable<Uint8Array> + | { readonly on: (event: string, cb: (data: Uint8Array) => void) => void } + | undefined; + readonly pid: number | undefined; + readonly kill: () => void; } export type SpawnProcess = ( - command: readonly string[], - opts: { readonly cwd: string; readonly env?: Readonly<Record<string, string>> | undefined }, + command: readonly string[], + opts: { readonly cwd: string; readonly env?: Readonly<Record<string, string>> | undefined }, ) => SpawnedProcess; export interface Connection { - readonly send: (method: string, params?: unknown) => Promise<unknown>; - readonly notify: (method: string, params?: unknown) => void; - readonly onNotification: (method: string, handler: (params: unknown) => void) => void; - readonly close: () => void; - readonly pid: number | undefined; + readonly send: (method: string, params?: unknown) => Promise<unknown>; + readonly notify: (method: string, params?: unknown) => void; + readonly onNotification: (method: string, handler: (params: unknown) => void) => void; + readonly close: () => void; + readonly pid: number | undefined; } export interface StdioTransportDeps { - readonly spawn: SpawnProcess; - readonly command: readonly string[]; - readonly env?: Readonly<Record<string, string>>; + readonly spawn: SpawnProcess; + readonly command: readonly string[]; + readonly env?: Readonly<Record<string, string>>; } export function createStdioTransport( - deps: StdioTransportDeps, - cwd: string, + deps: StdioTransportDeps, + cwd: string, ): { connection: Connection; promise: Promise<void> } { - const spawnOpts: { readonly cwd: string; readonly env?: Readonly<Record<string, string>> } = { - cwd, - }; - if (deps.env) { - (spawnOpts as { env?: Readonly<Record<string, string>> }).env = deps.env; - } + const spawnOpts: { readonly cwd: string; readonly env?: Readonly<Record<string, string>> } = { + cwd, + }; + if (deps.env) { + (spawnOpts as { env?: Readonly<Record<string, string>> }).env = deps.env; + } - const proc = deps.spawn(deps.command, spawnOpts); - const decoder = new FrameDecoder(); + const proc = deps.spawn(deps.command, spawnOpts); + const decoder = new FrameDecoder(); - const writeFn: WriteFn = (bytes) => proc.stdin.write(bytes); - const rpc = new JsonRpcClient(writeFn); + const writeFn: WriteFn = (bytes) => proc.stdin.write(bytes); + const rpc = new JsonRpcClient(writeFn); - const stdoutSource = proc.stdout; - const promise = new Promise<void>((resolve, reject) => { - if (Symbol.asyncIterator in stdoutSource) { - (async () => { - try { - for await (const chunk of stdoutSource as AsyncIterable<Uint8Array>) { - const messages = decoder.decode(chunk); - for (const msg of messages) { - rpc.handleMessage(msg); - } - } - resolve(); - } catch (err: unknown) { - reject(err); - } - })(); - } else { - const source = stdoutSource as { - readonly on: (event: string, cb: (data: Uint8Array) => void) => void; - }; - source.on("data", (data: Uint8Array) => { - const messages = decoder.decode(data); - for (const msg of messages) { - rpc.handleMessage(msg); - } - }); - source.on("end", () => resolve()); - source.on("error", (err: unknown) => reject(err)); - } - }); + const stdoutSource = proc.stdout; + const promise = new Promise<void>((resolve, reject) => { + if (Symbol.asyncIterator in stdoutSource) { + (async () => { + try { + for await (const chunk of stdoutSource as AsyncIterable<Uint8Array>) { + const messages = decoder.decode(chunk); + for (const msg of messages) { + rpc.handleMessage(msg); + } + } + resolve(); + } catch (err: unknown) { + reject(err); + } + })(); + } else { + const source = stdoutSource as { + readonly on: (event: string, cb: (data: Uint8Array) => void) => void; + }; + source.on("data", (data: Uint8Array) => { + const messages = decoder.decode(data); + for (const msg of messages) { + rpc.handleMessage(msg); + } + }); + source.on("end", () => resolve()); + source.on("error", (err: unknown) => reject(err)); + } + }); - const connection: Connection = { - send: (method, params) => rpc.request(method, params), - notify: (method, params) => rpc.notify(method, params), - onNotification: (method, handler) => rpc.onNotification(method, handler), - close: () => { - rpc.close(); - proc.kill(); - }, - pid: proc.pid, - }; + const connection: Connection = { + send: (method, params) => rpc.request(method, params), + notify: (method, params) => rpc.notify(method, params), + onNotification: (method, handler) => rpc.onNotification(method, handler), + close: () => { + rpc.close(); + proc.kill(); + }, + pid: proc.pid, + }; - return { connection, promise }; + return { connection, promise }; } diff --git a/packages/mcp/src/types.ts b/packages/mcp/src/types.ts index 522511b..fb851fd 100644 --- a/packages/mcp/src/types.ts +++ b/packages/mcp/src/types.ts @@ -5,75 +5,75 @@ import type { ToolParameterSchema } from "@dispatch/kernel"; export interface McpServerConfig { - readonly command: string; - readonly args?: readonly string[]; - readonly env?: Readonly<Record<string, string>>; + readonly command: string; + readonly args?: readonly string[]; + readonly env?: Readonly<Record<string, string>>; } export interface ResolvedMcpServer { - readonly id: string; - readonly command: readonly string[]; - readonly env?: Readonly<Record<string, string>>; - readonly configSource: ".dispatch/mcp.json" | "opencode.json"; + readonly id: string; + readonly command: readonly string[]; + readonly env?: Readonly<Record<string, string>>; + readonly configSource: ".dispatch/mcp.json" | "opencode.json"; } export interface ResolveResult { - readonly servers: readonly ResolvedMcpServer[]; - readonly shadowed: boolean; + readonly servers: readonly ResolvedMcpServer[]; + readonly shadowed: boolean; } export type McpServerState = "connecting" | "connected" | "error" | "disconnected"; export interface McpServerStatus { - readonly id: string; - readonly state: McpServerState; - readonly error?: string; - readonly toolCount: number; + readonly id: string; + readonly state: McpServerState; + readonly error?: string; + readonly toolCount: number; } export interface McpToolInfo { - readonly name: string; - readonly description: string; - readonly inputSchema: ToolParameterSchema; + readonly name: string; + readonly description: string; + readonly inputSchema: ToolParameterSchema; } export interface McpContentItem { - readonly type: string; - readonly text?: string; - readonly data?: string; - readonly mimeType?: string; - readonly resource?: { - readonly uri: string; - readonly text?: string; - }; + readonly type: string; + readonly text?: string; + readonly data?: string; + readonly mimeType?: string; + readonly resource?: { + readonly uri: string; + readonly text?: string; + }; } export interface McpCallResult { - readonly content: readonly McpContentItem[]; - readonly isError?: boolean; + readonly content: readonly McpContentItem[]; + readonly isError?: boolean; } export interface McpServerCapabilities { - readonly tools?: { - readonly listChanged?: boolean; - }; + readonly tools?: { + readonly listChanged?: boolean; + }; } export interface McpInitializeResult { - readonly protocolVersion: string; - readonly capabilities: McpServerCapabilities; - readonly serverInfo: { - readonly name: string; - readonly version: string; - }; + readonly protocolVersion: string; + readonly capabilities: McpServerCapabilities; + readonly serverInfo: { + readonly name: string; + readonly version: string; + }; } export interface McpService { - readonly status: (cwd: string) => Promise<readonly McpServerStatus[]>; + readonly status: (cwd: string) => Promise<readonly McpServerStatus[]>; } export interface McpListToolsResult { - readonly tools: readonly McpToolInfo[]; + readonly tools: readonly McpToolInfo[]; } /** @@ -83,5 +83,5 @@ export interface McpListToolsResult { * `McpClient` satisfies this structurally. */ export interface McpToolCaller { - readonly callTool: (name: string, args: unknown, signal?: AbortSignal) => Promise<McpCallResult>; + readonly callTool: (name: string, args: unknown, signal?: AbortSignal) => Promise<McpCallResult>; } |
