This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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>,
|
||||
onNotification?: (notif: {
|
||||
method: string;
|
||||
params?: Record<string, unknown>;
|
||||
}) => void | Promise<void>,
|
||||
handlers: {
|
||||
onNotification?: (notif: {
|
||||
method: string;
|
||||
params?: Record<string, unknown>;
|
||||
}) => 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}`,
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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();
|
||||
}
|
||||
});
|
||||
|
||||
Reference in New Issue
Block a user