diff --git a/src/gateway.ts b/src/gateway.ts index d38e234..dd9fa19 100644 --- a/src/gateway.ts +++ b/src/gateway.ts @@ -28,6 +28,10 @@ export interface GatewayOptions { agents?: EnabledAgent[]; /** Override the bearer token (otherwise read from ~/.a2a/credentials.json). */ token?: string | null; + /** Test/advanced override for one-shot upstream JSON-RPC calls. */ + upstreamRequestTimeoutMs?: number; + /** Test/advanced override for idle SSE upstream tools/call streams. */ + upstreamStreamIdleTimeoutMs?: number; /** Optional override for the underlying MCP Server (tests). */ server?: Server; } @@ -56,6 +60,8 @@ export async function buildGateway(opts: GatewayOptions = {}): Promise(method: string, params: unknown): Promise { const id = _nextId++; + const timeoutMs = this.cfg.requestTimeoutMs ?? DEFAULT_RPC_TIMEOUT_MS; + const controller = new AbortController(); + let timedOut = false; + const timeout = setTimeout(() => { + timedOut = true; + controller.abort(); + }, timeoutMs); const headers: Record = { "content-type": "application/json", accept: "application/json", }; if (this.cfg.token) headers["authorization"] = `Bearer ${this.cfg.token}`; - const resp = await fetch(`${this.cfg.url}/mcp`, { - method: "POST", - headers, - body: JSON.stringify({ jsonrpc: JSON_RPC, id, method, params }), - }); + try { + const resp = await fetch(`${this.cfg.url}/mcp`, { + method: "POST", + headers, + body: JSON.stringify({ jsonrpc: JSON_RPC, id, method, params }), + signal: controller.signal, + }); - if (!resp.ok) { - const text = await resp.text().catch(() => ""); - throw new UpstreamError(this.cfg.name, resp.status, text || resp.statusText); - } + if (!resp.ok) { + const text = await resp.text().catch(() => ""); + throw new UpstreamError(this.cfg.name, resp.status, text || resp.statusText); + } - const env: any = await resp.json(); - if (env?.error) { - throw new UpstreamError( - this.cfg.name, - 500, - `${env.error.code}: ${env.error.message}`, - ); + const env: any = await resp.json(); + if (env?.error) { + throw new UpstreamError( + this.cfg.name, + 500, + `${env.error.code}: ${env.error.message}`, + ); + } + return env?.result as T; + } catch (err) { + if (timedOut) { + throw new UpstreamError( + this.cfg.name, + 504, + `request timed out after ${timeoutMs}ms`, + ); + } + throw err; + } finally { + clearTimeout(timeout); } - return env?.result as T; } async listTools(): Promise { @@ -136,22 +161,51 @@ export class UpstreamAgent { }) => void | Promise, ): Promise { const id = _nextId++; + const idleTimeoutMs = + this.cfg.streamIdleTimeoutMs ?? DEFAULT_STREAM_IDLE_TIMEOUT_MS; + const controller = new AbortController(); + let timedOut = false; + let idleTimer: ReturnType | undefined; + const resetIdleTimer = () => { + if (idleTimer) clearTimeout(idleTimer); + idleTimer = setTimeout(() => { + timedOut = true; + controller.abort(); + }, idleTimeoutMs); + }; + resetIdleTimer(); const headers: Record = { "content-type": "application/json", accept: "text/event-stream, application/json", }; if (this.cfg.token) headers["authorization"] = `Bearer ${this.cfg.token}`; - const resp = await fetch(`${this.cfg.url}/mcp`, { - method: "POST", - headers, - body: JSON.stringify({ - jsonrpc: JSON_RPC, id, method: "tools/call", params, - }), - }); + let resp: Response; + try { + resp = await fetch(`${this.cfg.url}/mcp`, { + method: "POST", + headers, + body: JSON.stringify({ + jsonrpc: JSON_RPC, id, method: "tools/call", params, + }), + signal: controller.signal, + }); + } catch (err) { + clearTimeout(idleTimer); + if (timedOut) { + throw new UpstreamError( + this.cfg.name, + 504, + `stream idle timeout after ${idleTimeoutMs}ms`, + ); + } + throw err; + } + resetIdleTimer(); if (!resp.ok) { const text = await resp.text().catch(() => ""); + clearTimeout(idleTimer); throw new UpstreamError(this.cfg.name, resp.status, text || resp.statusText); } @@ -160,15 +214,18 @@ export class UpstreamAgent { // Server fell back to JSON (older a2a-pack or non-SSE upstream). const env: any = await resp.json(); if (env?.error) { + clearTimeout(idleTimer); throw new UpstreamError( this.cfg.name, 500, `${env.error.code}: ${env.error.message}`, ); } + clearTimeout(idleTimer); return env?.result as UpstreamCallResult; } if (!resp.body) { + clearTimeout(idleTimer); throw new UpstreamError(this.cfg.name, 502, "empty stream body"); } @@ -177,8 +234,24 @@ export class UpstreamAgent { let buf = ""; try { while (true) { - const { done, value } = await reader.read(); - if (value) buf += decoder.decode(value, { stream: true }); + let done: boolean; + let value: Uint8Array | undefined; + try { + ({ done, value } = await reader.read()); + } catch (err) { + if (timedOut) { + throw new UpstreamError( + this.cfg.name, + 504, + `stream idle timeout after ${idleTimeoutMs}ms`, + ); + } + throw err; + } + if (value) { + resetIdleTimer(); + buf += decoder.decode(value, { stream: true }); + } let sep = buf.indexOf("\n\n"); while (sep !== -1) { const frame = buf.slice(0, sep); @@ -189,6 +262,10 @@ export class UpstreamAgent { } if (dataLines.length > 0) { const raw = dataLines.join("\n"); + if (raw === "[DONE]") { + sep = buf.indexOf("\n\n"); + continue; + } let msg: any; try { msg = JSON.parse(raw); } catch { msg = null; } if (msg && msg.id === id) { @@ -221,6 +298,7 @@ export class UpstreamAgent { if (done) break; } } finally { + clearTimeout(idleTimer); try { reader.releaseLock(); } catch { /* ignore */ } } throw new UpstreamError(this.cfg.name, 502, "stream ended without result"); diff --git a/tests/gateway.test.ts b/tests/gateway.test.ts index 9765690..01a60b4 100644 --- a/tests/gateway.test.ts +++ b/tests/gateway.test.ts @@ -9,7 +9,7 @@ import { ListToolsRequestSchema, } from "@modelcontextprotocol/sdk/types.js"; -import { buildGateway, SEP } from "../src/gateway.js"; +import { buildGateway, SEP, type GatewayOptions } from "../src/gateway.js"; import type { EnabledAgent } from "../src/config.js"; /** @@ -70,25 +70,65 @@ function startFakeAgent(opts: { }); } +function startSseAgent(opts: { + onCall: (body: any, res: http.ServerResponse) => void; +}): Promise<{ url: string; close: () => Promise; requests: any[] }> { + const requests: any[] = []; + return new Promise((resolve) => { + const server = http.createServer((req, res) => { + if (req.method !== "POST" || req.url !== "/mcp") { + res.statusCode = 404; + res.end(); + return; + } + const chunks: Buffer[] = []; + req.on("data", (c) => chunks.push(c)); + req.on("end", () => { + const body = JSON.parse(Buffer.concat(chunks).toString("utf8")); + requests.push({ body }); + if (body.method !== "tools/call") { + res.setHeader("content-type", "application/json"); + res.end(JSON.stringify({ jsonrpc: "2.0", id: body.id, result: {} })); + return; + } + res.writeHead(200, { "content-type": "text/event-stream" }); + opts.onCall(body, res); + }); + }); + server.listen(0, "127.0.0.1", () => { + const { port } = server.address() as AddressInfo; + resolve({ + url: `http://127.0.0.1:${port}`, + requests, + close: () => new Promise((r) => server.close(() => r())), + }); + }); + }); +} + /** Drive an in-process gateway Server via the SDK's request handlers. */ async function callGateway( agents: EnabledAgent[], token: string | null, method: "tools/list" | "tools/call", params: any, + opts: { + gateway?: Partial>; + extra?: any; + } = {}, ) { const mcpServer = new Server( { name: "test", version: "0.0.0" }, { capabilities: { tools: { listChanged: false } } }, ); - await buildGateway({ agents, token, server: mcpServer }); + await buildGateway({ agents, token, server: mcpServer, ...opts.gateway }); // Reach into the server's registered handlers — the SDK exposes them via // a private map, so we route through the public schema-keyed lookup we // know is wired up. const schema = method === "tools/list" ? ListToolsRequestSchema : CallToolRequestSchema; const handler = (mcpServer as any)._requestHandlers.get(schema.shape.method.value); if (!handler) throw new Error("handler not registered"); - return handler({ method, params }, {} as any); + return handler({ method, params }, opts.extra ?? ({} as any)); } test("tools/list aggregates upstream tools with agent prefix", async () => { @@ -209,3 +249,74 @@ test("tools/call surfaces upstream JSON-RPC error as soft error", async () => { await upstream.close(); } }); + +test("tools/call forwards upstream SSE progress notifications", async () => { + const upstream = await startSseAgent({ + onCall: (body, res) => { + res.write( + `data: ${JSON.stringify({ + jsonrpc: "2.0", + method: "notifications/progress", + params: { progressToken: "p1", message: "halfway" }, + })}\n\n`, + ); + res.end( + `data: ${JSON.stringify({ + jsonrpc: "2.0", + id: body.id, + result: { + content: [{ type: "text", text: "done" }], + isError: false, + }, + })}\n\n`, + ); + }, + }); + const notifications: any[] = []; + try { + const result = await callGateway( + [{ name: "streamer", url: upstream.url, addedAt: "now" }], + null, + "tools/call", + { name: `streamer${SEP}build`, arguments: {} }, + { + extra: { + requestId: "req-1", + sendNotification: async (notification: any) => { + notifications.push(notification); + }, + }, + }, + ); + assert.equal(result.isError, false); + assert.deepEqual(notifications, [ + { + method: "notifications/progress", + params: { progressToken: "p1", message: "halfway" }, + }, + ]); + } finally { + await upstream.close(); + } +}); + +test("tools/call reports upstream SSE idle timeout", async () => { + const upstream = await startSseAgent({ + onCall: (_body, res) => { + res.write(": connected\n\n"); + }, + }); + try { + const result = await callGateway( + [{ name: "slow", url: upstream.url, addedAt: "now" }], + null, + "tools/call", + { name: `slow${SEP}build`, arguments: {} }, + { gateway: { upstreamStreamIdleTimeoutMs: 25 } }, + ); + assert.equal(result.isError, true); + assert.match(result.content[0].text, /stream idle timeout/); + } finally { + await upstream.close(); + } +});