diff --git a/services/agent-harness/src/tools/mcp-test-fixture.ts b/services/agent-harness/src/tools/mcp-test-fixture.ts new file mode 100644 index 0000000000..58db1b0a54 --- /dev/null +++ b/services/agent-harness/src/tools/mcp-test-fixture.ts @@ -0,0 +1,79 @@ +import type { Tool } from '@modelcontextprotocol/sdk/types.js'; + +export function mcpTestFixture() { + const key = crypto.randomUUID(); + const shape = { + type: 'object' as const, + properties: { [key]: { type: 'integer', minimum: 1 } }, + required: [key], + additionalProperties: false, + }; + const tool: Tool = { + name: 'remote', + inputSchema: shape, + outputSchema: shape, + annotations: { readOnlyHint: true }, + }; + const connection = { + serverId: 'configured', + configurationVersion: '1', + url: 'https://gateway.example/mcp-connect/user/owner/configured/route', + authorization: 'Bearer derived-secret', + }; + const state = { + tool, + result: { content: [], structuredContent: { [key]: 2 } } as unknown, + effects: [] as unknown[], + status: 200, + lose: false, + overflow: false, + stall: false, + response: undefined as Response | ((message: object) => Response) | undefined, + requests: [] as { url: string; init: RequestInit }[], + }; + const fetchImpl: typeof fetch = async (target, init) => { + if (init?.redirect !== 'manual') throw new Error('Redirect controls are missing'); + const url = + typeof target === 'string' ? target : target instanceof URL ? target.href : target.url; + state.requests.push({ url, init }); + if (init.method === 'GET') return new Response(null, { status: 405 }); + if (typeof init.body !== 'string') throw new Error('Expected a JSON request body'); + const message = JSON.parse(init.body) as { + id?: string | number; + method: string; + params: { arguments: unknown }; + }; + if (!('id' in message)) return new Response(null, { status: 202 }); + if (state.status !== 200) return new Response('provider-secret', { status: state.status }); + let result: unknown = { + protocolVersion: '2025-11-25', + capabilities: { tools: {} }, + serverInfo: { name: 'remote', version: '1' }, + }; + if (message.method === 'tools/list') result = { tools: [state.tool] }; + if (message.method === 'tools/call') { + state.effects.push(message.params.arguments); + if (state.lose) throw new Error('provider-secret'); + if (state.stall) + return new Promise((_resolve, reject) => + init.signal?.addEventListener('abort', () => reject(new Error('timeout')), { once: true }) + ); + if (state.overflow) + return new Response( + new ReadableStream({ + start(controller) { + controller.enqueue(new Uint8Array(4096)); + controller.close(); + }, + }), + { headers: { 'Content-Type': 'application/json' } } + ); + result = state.result; + } + const response = { jsonrpc: '2.0', id: message.id, result }; + return typeof state.response === 'function' + ? state.response(response) + : (state.response ?? Response.json(response)); + }; + return { key, state, connection, fetchImpl }; +} diff --git a/services/agent-harness/src/tools/mcp-transport.test.ts b/services/agent-harness/src/tools/mcp-transport.test.ts new file mode 100644 index 0000000000..3f6d55c500 --- /dev/null +++ b/services/agent-harness/src/tools/mcp-transport.test.ts @@ -0,0 +1,224 @@ +import { expect, it, onTestFinished, vi } from 'vitest'; +import { Client } from '@modelcontextprotocol/sdk/client/index.js'; +import { ErrorCode, McpError } from '@modelcontextprotocol/sdk/types.js'; +import { CfWorkerJsonSchemaValidator } from '@modelcontextprotocol/sdk/validation/cfworker-provider.js'; +import { createMcpTransportFactory } from './mcp-transport'; +import { mcpTestFixture } from './mcp-test-fixture'; + +const closed = new McpError(ErrorCode.ConnectionClosed, 'Connection closed'); +const list = { jsonrpc: '2.0', id: 1, method: 'tools/list' } as const; +const call = (args: object) => ({ + ...list, + method: 'tools/call', + params: { name: 'remote', arguments: args }, +}); +const reply = (body: BodyInit, type = 'application/json') => + new Response(body, { headers: { 'Content-Type': type } }); +const streamReply = (source: UnderlyingSource, type = 'application/json') => + reply(new ReadableStream(source), type); +const sseData = (message: object) => `data: ${JSON.stringify(message)}\n\n`; +function fixture(httpResponseBytes = 4096, signal?: AbortSignal) { + const provider = mcpTestFixture(); + const abort = new AbortController(); + const factory = createMcpTransportFactory( + signal ? AbortSignal.any([abort.signal, signal]) : abort.signal, + httpResponseBytes, + reason => { + throw new Error(reason); + }, + provider.fetchImpl + ); + return { + ...provider, + abort, + transport(connection = provider.connection) { + const transport = factory(connection); + onTestFinished(() => transport.close()); + return transport; + }, + async client(transport = this.transport()) { + const client = new Client( + { name: 'test', version: '1' }, + { jsonSchemaValidator: new CfWorkerJsonSchemaValidator() } + ); + await client.connect(transport); + return client; + }, + }; +} + +it.each([ + [401, 'reauthorization_required'], + [403, 'reauthorization_required'], + [503, 'unavailable_server'], + [302, 'unsafe_destination'], + [307, 'unsafe_destination'], + [308, 'unsafe_destination'], + [400, 'unsafe_destination'], + ['redirected', 'unsafe_destination'], + ['url', 'unsafe_destination'], + ['lose', 'unavailable_server'], + ['body', 'unavailable_server'], + ['json', 'unavailable_server'], +] as const)('sanitizes %s without retrying a mutation', async (mode, reason) => { + const f = fixture(); + if (typeof mode === 'number') f.state.status = mode; + f.state.lose = mode === 'lose'; + if (mode === 'redirected' || mode === 'url') { + f.state.response = reply('{}'); + Object.defineProperty(f.state.response, mode, { + value: mode === 'url' ? 'https://other.example' : true, + }); + } + if (mode === 'body') + f.state.response = streamReply({ + start(controller) { + controller.error(new Error('provider-secret')); + }, + }); + if (mode === 'json') f.state.response = reply('provider-secret'); + const transport = f.transport(); + const errors: string[] = []; + transport.onerror = error => errors.push(error.message); + await transport.start(); + await expect(transport.send(call({}))).rejects.toThrow(reason); + expect(errors).toEqual([reason]); + expect(f.state.requests).toHaveLength(1); + expect(f.state.effects).toEqual(typeof mode === 'number' ? [] : [{}]); +}); + +it.each([ + 'not a URL', + 'http://gateway.example', + 'https://user:secret@gateway.example', + 'https://gateway.example?secret', + 'https://gateway.example#secret', +])('rejects unsafe URL %s', url => { + const f = fixture(); + expect(() => f.transport({ ...f.connection, url })).toThrow('unsafe_destination'); + expect(f.state.requests).toEqual([]); +}); + +it('uses SDK initialization, discovery, and calls without changing the authorized destination', async () => { + const f = fixture(); + const transport = f.transport(); + const destination = f.connection.url; + f.connection.url = 'https://other.example'; + f.connection.authorization = 'Bearer replacement'; + const client = await f.client(transport); + expect(await client.listTools()).toEqual({ tools: [f.state.tool] }); + expect(await client.callTool(call({ [f.key]: 2 }).params)).toEqual(f.state.result); + f.state.response = message => + reply(sseData({ ...message, result: { tools: [] } }), 'text/event-stream'); + expect(await client.listTools()).toEqual({ tools: [] }); + expect(f.state.effects).toEqual([{ [f.key]: 2 }]); + for (const { url, init } of f.state.requests) { + expect(url).toBe(destination); + expect(new Headers(init.headers).get('authorization')).toBe('Bearer derived-secret'); + expect(new Headers(init.headers).get('accept')).toContain('text/event-stream'); + } + expect(f.state.requests.some(({ init }) => init.method === 'GET')).toBe(true); +}); + +it.each(['cancel', 'fetch', 'disconnected', 'body', 'json', 'envelope'])( + 'rejects %s requests and blocks the failed operation', + async mode => { + const f = fixture(); + const secondConnection = { ...f.connection, url: `${f.connection.url}/second` }; + const transports = [f.transport(), f.transport(secondConnection)]; + const clients = await Promise.all(transports.map(transport => f.client(transport))); + const streams: [ReadableStreamDefaultController, object][] = []; + f.state.stall = mode === 'fetch'; + f.state.response = message => + streamReply( + { + start(controller) { + streams.push([controller, message]); + }, + }, + 'text/event-stream' + ); + const args = { [f.key]: 2 }; + const pending = Promise.allSettled(clients.map(client => client.callTool(call(args).params))); + await vi.waitFor(() => expect(f.state.effects).toEqual([args, args])); + const requests = f.state.requests.length; + const secret = 'provider-secret https://user:secret@gateway.example Bearer derived-secret'; + if (mode === 'cancel' || mode === 'fetch') f.abort.abort(new Error(secret)); + else { + const stream = streams[0]; + if (!stream) throw new Error('Missing SSE request'); + const [controller, message] = stream; + if (mode === 'body') controller.error(new Error(secret)); + else { + const data = + mode === 'disconnected' + ? 'id: resume\nevent: ping\ndata: {}\nretry: 1\n\n' + : `data: ${mode === 'json' ? secret : '{}'}\n\n${sseData(message)}`; + controller.enqueue(new TextEncoder().encode(data)); + controller.close(); + } + } + for (const outcome of await pending) + expect(outcome).toEqual({ status: 'rejected', reason: closed }); + for (const transport of transports) + await expect(transport.send(call(args))).rejects.toThrow('unavailable_server'); + expect(() => f.transport()).toThrow('unavailable_server'); + expect(f.state.effects).toEqual([args, args]); + expect(f.state.requests).toHaveLength(requests); + }, + 1000 +); + +it('shares streamed UTF-8 bytes across connections and cancels overflow before parsing', async () => { + const f = fixture(1100); + const first = f.transport(); + const second = f.transport({ ...f.connection, url: `${f.connection.url}/second` }); + const messages: unknown[] = []; + first.onmessage = second.onmessage = message => messages.push(message); + await first.start(); + await second.start(); + await first.send(list); + expect(messages).toEqual([{ jsonrpc: '2.0', id: 1, result: { tools: [f.state.tool] } }]); + let cancelled = false, + chunks = 0; + f.state.response = streamReply({ + pull(controller) { + if (++chunks <= 5) controller.enqueue(new TextEncoder().encode('é'.repeat(100))); + else controller.close(); + }, + cancel() { + cancelled = true; + }, + }); + f.state.response.headers.set('Content-Length', '1'); + await expect(second.send(list)).rejects.toThrow('limit_exceeded'); + await vi.waitFor(() => expect(cancelled).toBe(true)); + expect(messages).toHaveLength(1); + expect(() => f.transport()).toThrow('limit_exceeded'); +}); + +it.each(['fetch', 'body', 'sse', 'initialize', 'aborted'])( + 'honors the caller deadline during %s', + async phase => { + const f = fixture(4096, phase === 'aborted' ? AbortSignal.abort() : AbortSignal.timeout(50)); + const initializing = phase === 'initialize' || phase === 'aborted'; + const client = initializing ? undefined : await f.client(); + f.state.stall = phase === 'fetch'; + const streamed = phase === 'body' || phase === 'sse' || phase === 'initialize'; + let cancelled = false; + if (streamed) + f.state.response = streamReply( + { + cancel() { + cancelled = true; + }, + }, + phase === 'body' ? 'application/json' : 'text/event-stream' + ); + const pending = client ? client.callTool(call({}).params) : f.client(); + await expect(pending).rejects.toEqual(closed); + expect(f.state.effects).toEqual(initializing ? [] : [{}]); + if (streamed) expect(cancelled).toBe(true); + }, + 1000 +); diff --git a/services/agent-harness/src/tools/mcp-transport.ts b/services/agent-harness/src/tools/mcp-transport.ts new file mode 100644 index 0000000000..be6b16f898 --- /dev/null +++ b/services/agent-harness/src/tools/mcp-transport.ts @@ -0,0 +1,150 @@ +import { StreamableHTTPClientTransport } from '@modelcontextprotocol/sdk/client/streamableHttp.js'; +import type { Transport } from '@modelcontextprotocol/sdk/shared/transport.js'; + +export type McpConnection = { + serverId: string; + configurationVersion: string; + url: string; + authorization: string; +}; +type McpTransportFailure = + | 'unsafe_destination' + | 'reauthorization_required' + | 'unavailable_server' + | 'limit_exceeded'; + +/** Create one factory per operation. The caller supplies its deadline/cancellation signal. */ +export function createMcpTransportFactory( + signal: AbortSignal, + httpResponseBytes: number, + onFailure: (reason: McpTransportFailure) => never, + fetchImpl: typeof fetch = fetch +) { + let receivedBytes = 0; + let failure: McpTransportFailure | undefined; + const connections = new Set(); + const stop = (reason: McpTransportFailure): never => { + failure ??= reason; + for (const sdk of connections) void sdk.close().catch(() => undefined); + return onFailure(failure); + }; + const fail = () => { + if (failure) return; + try { + stop('unavailable_server'); + } catch { + // onFailure throws; closing the SDK settles pending Client requests. + } + }; + signal.addEventListener('abort', fail, { once: true }); + return (connection: McpConnection): Transport => { + if (failure) stop(failure); + let url: URL; + try { + url = new URL(connection.url); + } catch { + return stop('unsafe_destination'); + } + if (url.protocol !== 'https:' || url.username || url.password || url.search || url.hash) + stop('unsafe_destination'); + // Authorization supplies this scoped gateway route; the gateway retains upstream DNS checks. + const destination = url.href; + const sdk = new StreamableHTTPClientTransport(url, { + requestInit: { headers: { Authorization: connection.authorization } }, + reconnectionOptions: { + maxRetries: 0, + initialReconnectionDelay: 0, + maxReconnectionDelay: 0, + reconnectionDelayGrowFactor: 1, + }, + fetch: async (target, init) => { + if (failure) stop(failure); + if (String(target) !== destination) stop('unsafe_destination'); + const requestSignal = AbortSignal.any([signal, ...(init?.signal ? [init.signal] : [])]); + requestSignal.throwIfAborted(); + const response = await fetchImpl(destination, { + ...init, + redirect: 'manual', + signal: requestSignal, + }); + if ( + (response.status >= 300 && response.status < 400) || + response.redirected || + (response.url && response.url !== destination) || + response.status === 400 + ) { + void response.body?.cancel().catch(() => undefined); + stop('unsafe_destination'); + } + if (!response.ok && !(init?.method === 'GET' && response.status === 405)) { + void response.body?.cancel().catch(() => undefined); + stop( + response.status === 401 || response.status === 403 + ? 'reauthorization_required' + : 'unavailable_server' + ); + } + const reader = response.body?.getReader(); + return new Response( + reader + ? new ReadableStream({ + start(controller) { + const cancel = () => { + controller.error(new Error(failure ?? 'unavailable_server')); + void reader.cancel().catch(() => undefined); + }; + requestSignal.addEventListener('abort', cancel, { once: true }); + void reader.closed + .finally(() => requestSignal.removeEventListener('abort', cancel)) + .catch(() => undefined); + if (requestSignal.aborted) cancel(); + }, + async pull(controller) { + const { done, value } = await reader.read(); + if (done) { + controller.close(); + return; + } + receivedBytes += value.byteLength; + if (receivedBytes > httpResponseBytes) stop('limit_exceeded'); + controller.enqueue(value); + }, + cancel: () => reader.cancel().catch(() => undefined), + }) + : null, + { status: response.status, headers: response.headers } + ); + }, + }); + // Keep SDK parsing and session handling, but never expose its provider-bearing diagnostics. + const transport: Transport = { + start: () => sdk.start(), + close: () => sdk.close(), + async send(message, options) { + if (failure) stop(failure); + try { + await sdk.send(message, options); + } catch { + stop(failure ?? 'unavailable_server'); + } + }, + get sessionId() { + return sdk.sessionId; + }, + setProtocolVersion: version => sdk.setProtocolVersion(version), + }; + sdk.onmessage = message => { + if (!failure) transport.onmessage?.(message); + }; + sdk.onclose = () => { + if (!connections.delete(sdk)) return; + transport.onclose?.(); + if (failure) transport.onerror?.(new Error(failure)); + }; + sdk.onerror = () => { + if (connections.has(sdk)) fail(); + }; + connections.add(sdk); + return transport; + }; +}