From 4ff51157c0e0fa9d7df48b041ab2c35c72b5142b Mon Sep 17 00:00:00 2001 From: Robert Date: Sat, 23 May 2026 13:32:25 -0400 Subject: [PATCH] Relay upstream MCP elicitation --- src/gateway.ts | 14 +++++++ src/upstream.ts | 97 ++++++++++++++++++++++++++++++++++++------- tests/gateway.test.ts | 94 +++++++++++++++++++++++++++++++++++++++++ 3 files changed, 190 insertions(+), 15 deletions(-) diff --git a/src/gateway.ts b/src/gateway.ts index dd9fa19..6ff7b8b 100644 --- a/src/gateway.ts +++ b/src/gateway.ts @@ -11,7 +11,9 @@ import { Server } from "@modelcontextprotocol/sdk/server/index.js"; import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js"; import { CallToolRequestSchema, + ElicitResultSchema, ListToolsRequestSchema, + ResultSchema, } from "@modelcontextprotocol/sdk/types.js"; import { loadConfig, type EnabledAgent } from "./config.js"; @@ -136,6 +138,18 @@ export async function buildGateway(opts: GatewayOptions = {}): Promise { + const schema = upstreamReq.method === "elicitation/create" + ? ElicitResultSchema + : ResultSchema; + return extra.sendRequest( + { + method: upstreamReq.method, + params: upstreamReq.params, + } as any, + schema as any, + ); + }, }, ); // Pass through the upstream's CallToolResult as-is. Upstreams already diff --git a/src/upstream.ts b/src/upstream.ts index b88e73f..a977f21 100644 --- a/src/upstream.ts +++ b/src/upstream.ts @@ -43,6 +43,12 @@ export interface UpstreamCallResult { isError?: boolean; } +export interface UpstreamServerRequest { + id: string | number; + method: string; + params?: Record; +} + export class UpstreamError extends Error { constructor( public agent: string, @@ -129,6 +135,7 @@ export class UpstreamAgent { method: string; params?: Record; }) => void | Promise; + onRequest?: (req: UpstreamServerRequest) => unknown | Promise; progressToken?: string | number; } = {}, ): Promise { @@ -139,7 +146,10 @@ export class UpstreamAgent { if (opts.progressToken !== undefined) meta.progressToken = opts.progressToken; const params: Record = { name, arguments: arguments_ }; if (Object.keys(meta).length > 0) params._meta = meta; - return this.streamingCallTool(params, opts.onNotification); + return this.streamingCallTool(params, { + onNotification: opts.onNotification, + onRequest: opts.onRequest, + }); } /** @@ -148,17 +158,18 @@ export class UpstreamAgent { * a deploy + a wait-for-live-card poll), and a plain JSON POST gets * killed by ingress idle timeouts (~60s default on traefik). The * upstream's SSE response interleaves any elicitation/create requests - * and a final tools/call result; here we drain the stream and return - * the final result. (Elicit relaying to the MCP client is not wired - * yet — any elicit that fires will time out on the upstream side and - * the skill will surface a soft error.) + * and a final tools/call result; here we relay server-to-client + * requests through the MCP client and return the final result. */ private async streamingCallTool( params: Record, - onNotification?: (notif: { - method: string; - params?: Record; - }) => void | Promise, + handlers: { + onNotification?: (notif: { + method: string; + params?: Record; + }) => void | Promise; + onRequest?: (req: UpstreamServerRequest) => unknown | Promise; + } = {}, ): Promise { const id = _nextId++; const idleTimeoutMs = @@ -224,6 +235,8 @@ export class UpstreamAgent { return env?.result as UpstreamCallResult; } + const upstreamSessionId = resp.headers.get("mcp-session-id"); + if (!resp.body) { clearTimeout(idleTimer); throw new UpstreamError(this.cfg.name, 502, "empty stream body"); @@ -278,19 +291,24 @@ export class UpstreamAgent { return msg.result as UpstreamCallResult; } // Otherwise it's a server→client message. Notifications - // (no id) → forward to the MCP client so long-running - // skills keep the client's tool-call timer alive. Requests - // (has id, has method) — e.g. elicitation/create — are - // not relayed yet, so they will time out upstream. - if (msg && typeof msg.method === "string" && !("id" in msg) && onNotification) { + // (no id) are best-effort. Requests (has id, has method), such + // as elicitation/create, must be relayed and answered upstream. + if (msg && typeof msg.method === "string" && !("id" in msg) && handlers.onNotification) { try { - await onNotification({ + await handlers.onNotification({ method: msg.method, params: msg.params as Record | undefined, }); } catch { // Best-effort; never break the stream on a notify error. } + } else if (msg && typeof msg.method === "string" && "id" in msg) { + await this.relayServerRequest(upstreamSessionId, { + id: msg.id as string | number, + method: msg.method, + params: msg.params as Record | undefined, + }, handlers.onRequest); + resetIdleTimer(); } } sep = buf.indexOf("\n\n"); @@ -303,4 +321,53 @@ export class UpstreamAgent { } throw new UpstreamError(this.cfg.name, 502, "stream ended without result"); } + + private async relayServerRequest( + upstreamSessionId: string | null, + req: UpstreamServerRequest, + onRequest?: (req: UpstreamServerRequest) => unknown | Promise, + ): Promise { + if (!upstreamSessionId) { + throw new UpstreamError( + this.cfg.name, + 502, + `upstream request ${req.method} missing Mcp-Session-Id`, + ); + } + let envelope: Record; + try { + if (!onRequest) throw new Error(`no client request relay for ${req.method}`); + const result = await onRequest(req); + envelope = { jsonrpc: JSON_RPC, id: req.id, result }; + } catch (err) { + envelope = { + jsonrpc: JSON_RPC, + id: req.id, + error: { + code: -32603, + message: err instanceof Error ? err.message : String(err), + }, + }; + } + + const headers: Record = { + "content-type": "application/json", + accept: "application/json", + "mcp-session-id": upstreamSessionId, + }; + if (this.cfg.token) headers.authorization = `Bearer ${this.cfg.token}`; + const resp = await fetch(`${this.cfg.url}/mcp`, { + method: "POST", + headers, + body: JSON.stringify(envelope), + }); + if (!resp.ok) { + const text = await resp.text().catch(() => ""); + throw new UpstreamError( + this.cfg.name, + resp.status, + text || `failed to deliver response for ${req.method}`, + ); + } + } } diff --git a/tests/gateway.test.ts b/tests/gateway.test.ts index 01a60b4..aff81f1 100644 --- a/tests/gateway.test.ts +++ b/tests/gateway.test.ts @@ -320,3 +320,97 @@ test("tools/call reports upstream SSE idle timeout", async () => { await upstream.close(); } }); + +test("tools/call relays upstream elicitation requests", async () => { + let resolveResponse: (() => void) | undefined; + const responseSeen = new Promise((resolve) => { resolveResponse = resolve; }); + let upstreamResponse: any = null; + const requests: any[] = []; + 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({ headers: req.headers, body }); + if (body.method === "tools/call") { + res.writeHead(200, { + "content-type": "text/event-stream", + "Mcp-Session-Id": "sess-1", + }); + res.write( + `data: ${JSON.stringify({ + jsonrpc: "2.0", + id: "ask-1", + method: "elicitation/create", + params: { message: "Name?", requestedSchema: { type: "object" } }, + })}\n\n`, + ); + void responseSeen.then(() => { + res.end( + `data: ${JSON.stringify({ + jsonrpc: "2.0", + id: body.id, + result: { + content: [{ type: "text", text: upstreamResponse.result.content.answer }], + isError: false, + }, + })}\n\n`, + ); + }); + return; + } + if (body.id === "ask-1" && req.headers["mcp-session-id"] === "sess-1") { + upstreamResponse = body; + res.statusCode = 202; + res.end(); + resolveResponse?.(); + return; + } + res.setHeader("content-type", "application/json"); + res.end(JSON.stringify({ jsonrpc: "2.0", id: body.id, result: {} })); + }); + }); + const upstream = await new Promise<{ url: string; close: () => Promise }>((resolve) => { + server.listen(0, "127.0.0.1", () => { + const { port } = server.address() as AddressInfo; + resolve({ + url: `http://127.0.0.1:${port}`, + close: () => new Promise((r) => server.close(() => r())), + }); + }); + }); + const clientRequests: any[] = []; + try { + const result = await callGateway( + [{ name: "human", url: upstream.url, addedAt: "now" }], + null, + "tools/call", + { name: `human${SEP}run_demo`, arguments: {} }, + { + extra: { + requestId: "req-elicitation", + sendNotification: async () => {}, + sendRequest: async (request: any) => { + clientRequests.push(request); + return { action: "accept", content: { answer: "Alice" } }; + }, + }, + }, + ); + assert.equal(result.isError, false); + assert.equal(result.content[0].text, "Alice"); + assert.equal(clientRequests[0].method, "elicitation/create"); + assert.deepEqual(upstreamResponse.result, { + action: "accept", + content: { answer: "Alice" }, + }); + assert.ok(requests.some((r) => r.body.id === "ask-1")); + } finally { + await upstream.close(); + } +});