Relay upstream MCP elicitation
All checks were successful
build / test (push) Successful in 31s

This commit is contained in:
2026-05-23 13:32:25 -04:00
parent 2fd97d8d09
commit 4ff51157c0
3 changed files with 190 additions and 15 deletions

View File

@@ -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<GatewayHa
// best effort
}
},
onRequest: async (upstreamReq) => {
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

View File

@@ -43,6 +43,12 @@ export interface UpstreamCallResult {
isError?: boolean;
}
export interface UpstreamServerRequest {
id: string | number;
method: string;
params?: Record<string, unknown>;
}
export class UpstreamError extends Error {
constructor(
public agent: string,
@@ -129,6 +135,7 @@ export class UpstreamAgent {
method: string;
params?: Record<string, unknown>;
}) => void | Promise<void>;
onRequest?: (req: UpstreamServerRequest) => unknown | Promise<unknown>;
progressToken?: string | number;
} = {},
): Promise<UpstreamCallResult> {
@@ -139,7 +146,10 @@ export class UpstreamAgent {
if (opts.progressToken !== undefined) meta.progressToken = opts.progressToken;
const params: Record<string, unknown> = { 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<string, unknown>,
handlers: {
onNotification?: (notif: {
method: string;
params?: Record<string, unknown>;
}) => void | Promise<void>,
}) => void | Promise<void>;
onRequest?: (req: UpstreamServerRequest) => unknown | Promise<unknown>;
} = {},
): Promise<UpstreamCallResult> {
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<string, unknown> | 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<string, unknown> | 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<unknown>,
): Promise<void> {
if (!upstreamSessionId) {
throw new UpstreamError(
this.cfg.name,
502,
`upstream request ${req.method} missing Mcp-Session-Id`,
);
}
let envelope: Record<string, unknown>;
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<string, string> = {
"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}`,
);
}
}
}

View File

@@ -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<void>((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<void> }>((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();
}
});