558 lines
18 KiB
TypeScript
558 lines
18 KiB
TypeScript
import { test } from "node:test";
|
|
import assert from "node:assert/strict";
|
|
import http from "node:http";
|
|
import { AddressInfo } from "node:net";
|
|
import os from "node:os";
|
|
import path from "node:path";
|
|
import { promises as fs } from "node:fs";
|
|
|
|
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 { loadConfig, 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",
|
|
{},
|
|
);
|
|
const upstreamTools = out.tools.filter((tool: any) => tool.name.includes(SEP));
|
|
assert.equal(upstreamTools.length, 1);
|
|
assert.equal(upstreamTools[0].name, `greeter${SEP}greet`);
|
|
assert.match(upstreamTools[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)
|
|
.filter((name: string) => name.includes(SEP));
|
|
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();
|
|
}
|
|
});
|
|
|
|
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();
|
|
}
|
|
});
|
|
|
|
function startFakeControlPlane(): Promise<{ url: string; close: () => Promise<void>; requests: any[] }> {
|
|
const requests: any[] = [];
|
|
return new Promise((resolve) => {
|
|
const server = http.createServer((req, res) => {
|
|
requests.push({ method: req.method, url: req.url, auth: req.headers.authorization });
|
|
if (req.method === "GET" && req.url?.startsWith("/v1/agents/search")) {
|
|
res.setHeader("content-type", "application/json");
|
|
res.end(
|
|
JSON.stringify([
|
|
{
|
|
name: "soc2-evidence",
|
|
description: "Collect SOC2 evidence",
|
|
status: "running",
|
|
public: true,
|
|
url: "https://soc2.example.test",
|
|
score: 0.91,
|
|
match_source: "lexical",
|
|
setup_required: false,
|
|
skills: [{ name: "collect", description: "", tags: ["security"], input_fields: [] }],
|
|
},
|
|
]),
|
|
);
|
|
return;
|
|
}
|
|
if (req.method === "GET" && req.url === "/v1/agents/soc2-evidence") {
|
|
res.setHeader("content-type", "application/json");
|
|
res.end(
|
|
JSON.stringify({
|
|
name: "soc2-evidence",
|
|
version: "1.0.0",
|
|
status: "running",
|
|
url: "https://soc2.example.test",
|
|
description: "Collect SOC2 evidence",
|
|
card: {
|
|
capabilities: {
|
|
a2a_cloud_import: {
|
|
mcp: { mode: "generated", path: "/mcp" },
|
|
},
|
|
},
|
|
},
|
|
}),
|
|
);
|
|
return;
|
|
}
|
|
res.statusCode = 404;
|
|
res.end("not found");
|
|
});
|
|
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())),
|
|
});
|
|
});
|
|
});
|
|
}
|
|
|
|
async function withTempHome<T>(fn: () => Promise<T>): Promise<T> {
|
|
const originalHome = process.env.HOME;
|
|
const tmpHome = await fs.mkdtemp(path.join(os.tmpdir(), "a2a-mcp-gateway-"));
|
|
process.env.HOME = tmpHome;
|
|
try {
|
|
return await fn();
|
|
} finally {
|
|
if (originalHome !== undefined) process.env.HOME = originalHome;
|
|
else delete process.env.HOME;
|
|
await fs.rm(tmpHome, { recursive: true, force: true });
|
|
}
|
|
}
|
|
|
|
test("tools/list includes a2a management tools", async () => {
|
|
const out = await callGateway([], "token", "tools/list", {});
|
|
const names = out.tools.map((tool: any) => tool.name);
|
|
assert.ok(names.includes("a2a_search_agents"));
|
|
assert.ok(names.includes("a2a_add_agent"));
|
|
assert.ok(names.includes("a2a_list_enabled_agents"));
|
|
});
|
|
|
|
test("a2a_search_agents returns visible control-plane results", async () => {
|
|
const cp = await startFakeControlPlane();
|
|
try {
|
|
const result = await callGateway(
|
|
[],
|
|
"access-token",
|
|
"tools/call",
|
|
{ name: "a2a_search_agents", arguments: { query: "soc2", tags: ["security"] } },
|
|
{ gateway: { apiUrl: cp.url } },
|
|
);
|
|
assert.equal(result.isError, undefined);
|
|
assert.equal(result.structuredContent.agents[0].name, "soc2-evidence");
|
|
assert.equal(result.structuredContent.agents[0].enabled, false);
|
|
assert.equal(cp.requests[0].auth, "Bearer access-token");
|
|
assert.match(cp.requests[0].url, /q=soc2/);
|
|
assert.match(cp.requests[0].url, /tag=security/);
|
|
} finally {
|
|
await cp.close();
|
|
}
|
|
});
|
|
|
|
test("a2a_add_agent persists control-plane agent and reports list change", async () => {
|
|
await withTempHome(async () => {
|
|
const cp = await startFakeControlPlane();
|
|
const notifications: any[] = [];
|
|
try {
|
|
const result = await callGateway(
|
|
[],
|
|
"access-token",
|
|
"tools/call",
|
|
{ name: "a2a_add_agent", arguments: { name: "soc2-evidence" } },
|
|
{
|
|
gateway: { apiUrl: cp.url },
|
|
extra: {
|
|
requestId: "req-add",
|
|
sendNotification: async (notification: any) => {
|
|
notifications.push(notification);
|
|
},
|
|
sendRequest: async () => ({}),
|
|
},
|
|
},
|
|
);
|
|
assert.equal(result.isError, undefined);
|
|
assert.match(result.content[0].text, /Added soc2-evidence/);
|
|
assert.equal(notifications[0].method, "notifications/tools/list_changed");
|
|
|
|
const cfg = await loadConfig();
|
|
assert.equal(cfg.agents[0].name, "soc2-evidence");
|
|
assert.equal(cfg.agents[0].url, `${cp.url}/v1/agents/soc2-evidence`);
|
|
assert.equal(cfg.agents[0].mcpPath, "/mcp");
|
|
} finally {
|
|
await cp.close();
|
|
}
|
|
});
|
|
});
|