323 lines
9.8 KiB
TypeScript
323 lines
9.8 KiB
TypeScript
import { test } from "node:test";
|
|
import assert from "node:assert/strict";
|
|
import http from "node:http";
|
|
import { AddressInfo } from "node:net";
|
|
|
|
import { Server } from "@modelcontextprotocol/sdk/server/index.js";
|
|
import {
|
|
CallToolRequestSchema,
|
|
ListToolsRequestSchema,
|
|
} from "@modelcontextprotocol/sdk/types.js";
|
|
|
|
import { buildGateway, SEP, type GatewayOptions } from "../src/gateway.js";
|
|
import type { EnabledAgent } from "../src/config.js";
|
|
|
|
/**
|
|
* Fake upstream A2A agent: speaks the JSON-RPC subset the gateway uses
|
|
* (initialize, tools/list, tools/call). One handler per server instance.
|
|
*/
|
|
function startFakeAgent(opts: {
|
|
tools: Array<{ name: string; description: string; inputSchema: any; outputSchema?: any }>;
|
|
onCall: (name: string, args: any) => any;
|
|
authHeader?: string;
|
|
}): Promise<{ url: string; close: () => Promise<void>; 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 auth = req.headers["authorization"];
|
|
const chunks: Buffer[] = [];
|
|
req.on("data", (c) => chunks.push(c));
|
|
req.on("end", () => {
|
|
const body = JSON.parse(Buffer.concat(chunks).toString("utf8"));
|
|
requests.push({ auth, body });
|
|
let result: any;
|
|
if (body.method === "tools/list") {
|
|
result = { tools: opts.tools };
|
|
} else if (body.method === "tools/call") {
|
|
try {
|
|
result = opts.onCall(body.params.name, body.params.arguments);
|
|
} catch (err) {
|
|
res.setHeader("content-type", "application/json");
|
|
res.end(
|
|
JSON.stringify({
|
|
jsonrpc: "2.0",
|
|
id: body.id,
|
|
error: { code: -32603, message: (err as Error).message },
|
|
}),
|
|
);
|
|
return;
|
|
}
|
|
} else {
|
|
result = {};
|
|
}
|
|
res.setHeader("content-type", "application/json");
|
|
res.end(JSON.stringify({ jsonrpc: "2.0", id: body.id, result }));
|
|
});
|
|
});
|
|
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())),
|
|
});
|
|
});
|
|
});
|
|
}
|
|
|
|
function startSseAgent(opts: {
|
|
onCall: (body: any, res: http.ServerResponse) => void;
|
|
}): Promise<{ url: string; close: () => Promise<void>; 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<Omit<GatewayOptions, "agents" | "token" | "server">>;
|
|
extra?: any;
|
|
} = {},
|
|
) {
|
|
const mcpServer = new Server(
|
|
{ name: "test", version: "0.0.0" },
|
|
{ capabilities: { tools: { listChanged: false } } },
|
|
);
|
|
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 }, opts.extra ?? ({} as any));
|
|
}
|
|
|
|
test("tools/list aggregates upstream tools with agent prefix", async () => {
|
|
const upstream = await startFakeAgent({
|
|
tools: [
|
|
{
|
|
name: "greet",
|
|
description: "Greet someone",
|
|
inputSchema: { type: "object", properties: { who: { type: "string" } } },
|
|
},
|
|
],
|
|
onCall: () => ({ content: [{ type: "text", text: "hi" }], isError: false }),
|
|
});
|
|
try {
|
|
const out = await callGateway(
|
|
[{ name: "greeter", url: upstream.url, addedAt: "now" }],
|
|
null,
|
|
"tools/list",
|
|
{},
|
|
);
|
|
assert.equal(out.tools.length, 1);
|
|
assert.equal(out.tools[0].name, `greeter${SEP}greet`);
|
|
assert.match(out.tools[0].description, /\[greeter\]/);
|
|
} finally {
|
|
await upstream.close();
|
|
}
|
|
});
|
|
|
|
test("tools/list survives a single upstream failure", async () => {
|
|
const good = await startFakeAgent({
|
|
tools: [
|
|
{
|
|
name: "ping",
|
|
description: "p",
|
|
inputSchema: { type: "object", properties: {} },
|
|
},
|
|
],
|
|
onCall: () => ({ content: [], isError: false }),
|
|
});
|
|
try {
|
|
const out = await callGateway(
|
|
[
|
|
{ name: "broken", url: "http://127.0.0.1:1", addedAt: "now" },
|
|
{ name: "good", url: good.url, addedAt: "now" },
|
|
],
|
|
null,
|
|
"tools/list",
|
|
{},
|
|
);
|
|
const names = out.tools.map((t: any) => t.name);
|
|
assert.deepEqual(names, [`good${SEP}ping`]);
|
|
} finally {
|
|
await good.close();
|
|
}
|
|
});
|
|
|
|
test("tools/call routes by agent prefix and forwards the bearer token", async () => {
|
|
const upstream = await startFakeAgent({
|
|
tools: [],
|
|
onCall: (name, args) => ({
|
|
content: [{ type: "text", text: `called ${name} with ${args.who}` }],
|
|
structuredContent: { result: `hello ${args.who}` },
|
|
isError: false,
|
|
}),
|
|
});
|
|
try {
|
|
const result = await callGateway(
|
|
[{ name: "greeter", url: upstream.url, addedAt: "now" }],
|
|
"tkn-123",
|
|
"tools/call",
|
|
{ name: `greeter${SEP}greet`, arguments: { who: "world" } },
|
|
);
|
|
assert.equal(result.isError, false);
|
|
assert.equal(result.structuredContent.result, "hello world");
|
|
const req = upstream.requests.find((r) => r.body.method === "tools/call");
|
|
assert.equal(req.auth, "Bearer tkn-123");
|
|
assert.equal(req.body.params.name, "greet"); // prefix stripped
|
|
} finally {
|
|
await upstream.close();
|
|
}
|
|
});
|
|
|
|
test("tools/call with unknown agent returns soft error", async () => {
|
|
const result = await callGateway([], null, "tools/call", {
|
|
name: `ghost${SEP}skill`,
|
|
arguments: {},
|
|
});
|
|
assert.equal(result.isError, true);
|
|
assert.match(result.content[0].text, /unknown agent: ghost/);
|
|
});
|
|
|
|
test("tools/call without prefix returns soft error", async () => {
|
|
const result = await callGateway([], null, "tools/call", {
|
|
name: "no_prefix",
|
|
arguments: {},
|
|
});
|
|
assert.equal(result.isError, true);
|
|
assert.match(result.content[0].text, /missing agent prefix/);
|
|
});
|
|
|
|
test("tools/call surfaces upstream JSON-RPC error as soft error", async () => {
|
|
const upstream = await startFakeAgent({
|
|
tools: [],
|
|
onCall: () => {
|
|
throw new Error("kaboom");
|
|
},
|
|
});
|
|
try {
|
|
const result = await callGateway(
|
|
[{ name: "boom", url: upstream.url, addedAt: "now" }],
|
|
null,
|
|
"tools/call",
|
|
{ name: `boom${SEP}explode`, arguments: {} },
|
|
);
|
|
assert.equal(result.isError, true);
|
|
assert.match(result.content[0].text, /kaboom/);
|
|
} finally {
|
|
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();
|
|
}
|
|
});
|