@bachi/pi-coder 1.0.0 → 1.1.0

package/extensions/mcp/client.test.ts

@@ -0,0 +1,518 @@
+/**
+ * Tests for client.ts — 三种传输的真实端到端用例。
+ *
+ * Run with:  node --test clients/pi/extensions/mcp/client.test.ts
+ *
+ * stdio 用的是真正 spawn 出来的 fixture 子进程（fixtures/fake-mcp-server.mjs），HTTP/SSE 用的是
+ * 测试里现起的 node:http 服务 —— 不 mock 传输层，因为这一层的坑几乎全在「字节怎么流」上：
+ * 半个 JSON、SSE 响应体、session id 往返、进程中途退出、超时与取消。
+ *
+ * 真实 wechat-local-mcp 的手感由 `npm run mcp:probe`（scripts/mcp-probe.mjs）验证，
+ * 不放进单测：那需要本机微信数据。
+ */
+
+import assert from "node:assert/strict";
+import { mkdtempSync, readFileSync, rmSync } from "node:fs";
+import { createServer, type IncomingMessage, type Server, type ServerResponse } from "node:http";
+import type { AddressInfo } from "node:net";
+import { tmpdir } from "node:os";
+import { join } from "node:path";
+import { fileURLToPath } from "node:url";
+import { after, describe, it } from "node:test";
+
+import { McpClient } from "./client.ts";
+import { McpConnectionError, McpError, MCP_PROTOCOL_VERSION } from "./protocol.ts";
+import type { McpRemoteServer, McpStdioServer } from "./config.ts";
+
+const FIXTURE = fileURLToPath(new URL("./fixtures/fake-mcp-server.mjs", import.meta.url));
+const TOKEN_HELPER = fileURLToPath(new URL("./fixtures/token-helper.mjs", import.meta.url));
+
+function stdioConfig(overrides: Partial<McpStdioServer> = {}): McpStdioServer {
+	return {
+		name: "fake",
+		transport: "stdio",
+		command: process.execPath,
+		args: [FIXTURE],
+		env: {},
+		timeoutMs: 5000,
+		enabled: true,
+		source: "test",
+		...overrides,
+	};
+}
+
+const clients: McpClient[] = [];
+const servers: Server[] = [];
+
+async function connectStdio(overrides: Partial<McpStdioServer> = {}, diagnostics: string[] = []): Promise<McpClient> {
+	const client = await McpClient.connect(stdioConfig(overrides), {
+		handshakeTimeoutMs: 5000,
+		onDiagnostic: (line) => diagnostics.push(line),
+	});
+	clients.push(client);
+	return client;
+}
+
+after(async () => {
+	await Promise.all(clients.map((client) => client.close().catch(() => {})));
+	await Promise.all(servers.map((server) => new Promise<void>((resolve) => server.close(() => resolve()))));
+});
+
+describe("stdio transport", () => {
+	it("完成握手并读出 serverInfo / 能力 / 工具表", async () => {
+		const client = await connectStdio();
+		assert.equal(client.serverInfo.name, "fake-mcp-server");
+		assert.equal(client.serverInfo.version, "9.9.9");
+		assert.equal(client.protocolVersion, "2025-06-18");
+		assert.deepEqual(Object.keys(client.capabilities), ["tools"]);
+
+		const tools = await client.listTools();
+		assert.equal(tools.length, 9);
+		assert.equal(tools[0]?.name, "echo");
+		assert.equal(tools[0]?.annotations?.readOnlyHint, true);
+		// 没有 inputSchema 的工具也要能列出来（index.ts 那边会给它兜一个空对象 schema）。
+		assert.equal(tools.find((tool) => tool.name === "no_schema")?.inputSchema, undefined);
+	});
+
+	it("tools/call 把参数原样送到服务端并把文本取回来", async () => {
+		const client = await connectStdio();
+		const result = await client.callTool("echo", { text: "你好", nested: { a: 1 } });
+		assert.equal(result.isError, false);
+		assert.deepEqual(JSON.parse((result.content[0] as { text: string }).text), {
+			text: "你好",
+			nested: { a: 1 },
+		});
+	});
+
+	it("isError: true 的结果照原样返回（由上层决定怎么报给模型）", async () => {
+		const client = await connectStdio();
+		const result = await client.callTool("fail", {});
+		assert.equal(result.isError, true);
+		assert.equal((result.content[0] as { text: string }).text, "工具内部失败了");
+	});
+
+	it("服务端返回 JSON-RPC 错误时抛 McpError 并保留错误码", async () => {
+		const client = await connectStdio();
+		await assert.rejects(
+			() => client.callTool("nope", {}),
+			(error: unknown) => error instanceof McpError && error.code === -32602 && /Unknown tool/.test(error.message),
+		);
+	});
+
+	it("超过 timeout 的调用抛超时错误，并发出 cancelled 通知", async () => {
+		const diagnostics: string[] = [];
+		const client = await connectStdio({ timeoutMs: 250 }, diagnostics);
+		await assert.rejects(
+			() => client.callTool("delay", { ms: 3000 }),
+			(error: unknown) => error instanceof McpError && /超时/.test(error.message),
+		);
+		// fixture 收到 notifications/cancelled 会往 stderr 写一行，经 onDiagnostic 回到这里。
+		await waitFor(() => diagnostics.some((line) => line.startsWith("cancelled request")));
+	});
+
+	it("AbortSignal 取消在途调用", async () => {
+		const diagnostics: string[] = [];
+		const client = await connectStdio({ timeoutMs: 10_000 }, diagnostics);
+		const controller = new AbortController();
+		const pending = client.callTool("delay", { ms: 3000 }, { signal: controller.signal });
+		setTimeout(() => controller.abort(), 50);
+		await assert.rejects(pending, (error: unknown) => error instanceof McpConnectionError && /已取消/.test(error.message));
+		await waitFor(() => diagnostics.some((line) => line.startsWith("cancelled request")));
+	});
+
+	it("子进程退出时在途请求被拒（不是永远挂着）", async () => {
+		const client = await connectStdio();
+		// exit 工具会先回响应再退出，所以这里等的是「连接变成关闭」这件事。
+		await client.callTool("exit", {}).catch(() => undefined);
+		await waitFor(() => client.isClosed);
+		assert.equal(client.isClosed, true);
+		await assert.rejects(() => client.callTool("echo", { text: "x" }), McpConnectionError);
+	});
+
+	it("采集子进程 stderr 作为诊断信息", async () => {
+		const diagnostics: string[] = [];
+		const client = await connectStdio({}, diagnostics);
+		await client.callTool("stderr", {});
+		await waitFor(() => diagnostics.includes("这是一行诊断输出"));
+	});
+
+	it("回复服务端的反向请求（回「未实现」而不是傻等）", async () => {
+		const diagnostics: string[] = [];
+		const client = await connectStdio({}, diagnostics);
+		await client.callTool("server_request", {});
+		await waitFor(() => diagnostics.some((line) => /answered: error -32601/.test(line)));
+	});
+
+	it("close() 之后再次调用直接失败", async () => {
+		const client = await connectStdio();
+		await client.close();
+		assert.equal(client.isClosed, true);
+		await assert.rejects(() => client.listTools(), McpConnectionError);
+	});
+
+	it("command 不存在时给出可读的连接错误", async () => {
+		await assert.rejects(
+			() =>
+				McpClient.connect(stdioConfig({ command: "/nonexistent/mcp-binary-xyz" }), {
+					handshakeTimeoutMs: 3000,
+				}),
+			(error: unknown) =>
+				error instanceof McpConnectionError && /子进程启动失败|连接失败/.test(error.message),
+		);
+	});
+});
+
+describe("streamable HTTP transport", () => {
+	it("JSON 响应：握手 + 工具调用，并带上 session id 与协议版本头", async () => {
+		const seen: Array<{ headers: IncomingMessage["headers"]; body: unknown }> = [];
+		const http = await startHttpServer((req, res, body) => {
+			seen.push({ headers: req.headers, body });
+			const message = body as { id?: number; method?: string };
+			if (message.method === "initialize") {
+				res.setHeader("mcp-session-id", "sess-42");
+				respondJson(res, 200, {
+					jsonrpc: "2.0",
+					id: message.id,
+					result: { protocolVersion: "2025-06-18", capabilities: { tools: {} }, serverInfo: { name: "http-fake", version: "1.0" } },
+				});
+				return;
+			}
+			if (message.method === "notifications/initialized") {
+				res.writeHead(202).end();
+				return;
+			}
+			if (message.method === "tools/list") {
+				respondJson(res, 200, {
+					jsonrpc: "2.0",
+					id: message.id,
+					result: { tools: [{ name: "ping", description: "ping", inputSchema: { type: "object", properties: {} } }] },
+				});
+				return;
+			}
+			respondJson(res, 200, { jsonrpc: "2.0", id: message.id, result: { content: [{ type: "text", text: "pong" }] } });
+		});
+		const client = await McpClient.connect({ ...http.config, name: "http-fake" });
+		clients.push(client);
+
+		const tools = await client.listTools();
+		assert.equal(tools[0]?.name, "ping");
+		const result = await client.callTool("ping", {});
+		assert.equal((result.content[0] as { text: string }).text, "pong");
+
+		const notInitialize = seen.filter((entry) => (entry.body as { method?: string }).method !== "initialize");
+		assert.ok(notInitialize.length > 0);
+		for (const entry of notInitialize) {
+			assert.equal(entry.headers["mcp-session-id"], "sess-42");
+			assert.equal(entry.headers["mcp-protocol-version"], MCP_PROTOCOL_VERSION);
+		}
+	});
+
+	it("SSE 响应体（且流不主动结束）也能拿到响应", async () => {
+		const http = await startHttpServer((req, res, body) => {
+			const message = body as { id?: number; method?: string };
+			if (message.method === "initialize") {
+				respondJson(res, 200, {
+					jsonrpc: "2.0",
+					id: message.id,
+					result: { protocolVersion: "2025-06-18", capabilities: {}, serverInfo: { name: "sse-body", version: "1.0" } },
+				});
+				return;
+			}
+			if (message.method === "notifications/initialized") {
+				res.writeHead(202).end();
+				return;
+			}
+			// 响应走 SSE，之后**不关流**：客户端必须自己收工。
+			res.writeHead(200, { "content-type": "text/event-stream" });
+			res.write(`data: ${JSON.stringify({ jsonrpc: "2.0", id: message.id, result: { tools: [] } })}\n\n`);
+		});
+		const client = await McpClient.connect({ ...http.config, name: "sse-body" });
+		clients.push(client);
+		assert.deepEqual(await client.listTools(), []);
+		// 上一个请求的 SSE 流还挂着（服务端没关），下一个请求照样发得出去。
+		const second = await client.callTool("anything", {});
+		assert.deepEqual(second.content, []);
+	});
+
+	it("HTTP 错误状态给出带状态码与 body 的错误", async () => {
+		const http = await startHttpServer((_req, res) => {
+			respondJson(res, 500, { error: "backend exploded" });
+		});
+		await assert.rejects(
+			() => McpClient.connect({ ...http.config, name: "broken" }, { handshakeTimeoutMs: 2000 }),
+			(error: unknown) => error instanceof McpConnectionError && /HTTP 500/.test(error.message) && /backend exploded/.test(error.message),
+		);
+	});
+
+	it("空响应体被视为协议错误", async () => {
+		const http = await startHttpServer((_req, res) => {
+			res.writeHead(200, { "content-type": "application/json" }).end();
+		});
+		await assert.rejects(
+			() => McpClient.connect({ ...http.config, name: "empty" }, { handshakeTimeoutMs: 2000 }),
+			(error: unknown) => error instanceof McpConnectionError && /为空/.test(error.message),
+		);
+	});
+});
+
+describe("legacy SSE transport", () => {
+	it("GET 流里的 endpoint 事件 + POST 到该端点 + 响应从流里回来", async () => {
+		const streams = new Set<ServerResponse>();
+		const http = await startHttpServer((req, res, body) => {
+			if (req.method === "GET") {
+				res.writeHead(200, { "content-type": "text/event-stream" });
+				res.write("event: endpoint\ndata: /messages?session=abc\n\n");
+				streams.add(res);
+				res.on("close", () => streams.delete(res));
+				return;
+			}
+			const message = body as { id?: number; method?: string };
+			res.writeHead(202).end();
+			if (message.method === "notifications/initialized") return;
+			const result =
+				message.method === "initialize"
+					? { protocolVersion: "2024-11-05", capabilities: {}, serverInfo: { name: "legacy-sse", version: "1.0" } }
+					: { tools: [{ name: "old_tool", description: "老协议工具", inputSchema: { type: "object", properties: {} } }] };
+			const payload = `data: ${JSON.stringify({ jsonrpc: "2.0", id: message.id, result })}\n\n`;
+			for (const stream of streams) stream.write(payload);
+		});
+
+		const client = await McpClient.connect({ ...http.config, transport: "sse", name: "legacy" });
+		clients.push(client);
+		assert.equal(client.serverInfo.name, "legacy-sse");
+		assert.equal(client.protocolVersion, "2024-11-05");
+		const tools = await client.listTools();
+		assert.equal(tools[0]?.name, "old_tool");
+	});
+});
+
+describe("dynamic headers (headersCommand)", () => {
+	const tempDirs: string[] = [];
+
+	function counterPath(): string {
+		const dir = mkdtempSync(join(tmpdir(), "mcp-header-test-"));
+		tempDirs.push(dir);
+		return join(dir, "token-count");
+	}
+
+	after(() => {
+		for (const dir of tempDirs) rmSync(dir, { recursive: true, force: true });
+	});
+
+	/** 只放行指定的 Authorization 值，其余一律 401 —— 这正是头命令要解决的场景。 */
+	async function startTokenGatedServer(expected: string): Promise<{
+		config: McpRemoteServer;
+		received: string[];
+		requests: () => number;
+	}> {
+		const received: string[] = [];
+		let requests = 0;
+		const http = await startHttpServer((req, res, body) => {
+			requests += 1;
+			const auth = req.headers["authorization"] ?? "";
+			received.push(auth);
+			if (auth !== expected) {
+				res.writeHead(401, { "content-type": "application/json" }).end('{"error":"unauthorized"}');
+				return;
+			}
+			const message = body as { id?: number; method?: string };
+			if (message.method === "initialize") {
+				respondJson(res, 200, {
+					jsonrpc: "2.0",
+					id: message.id,
+					result: { protocolVersion: "2025-06-18", capabilities: {}, serverInfo: { name: "gated", version: "1.0" } },
+				});
+				return;
+			}
+			if (message.method === "notifications/initialized") {
+				res.writeHead(202).end();
+				return;
+			}
+			respondJson(res, 200, {
+				jsonrpc: "2.0",
+				id: message.id,
+				result: { tools: [{ name: "ping", description: "ping", inputSchema: { type: "object", properties: {} } }] },
+			});
+		});
+		return { config: { ...http.config, name: "gated", headersCommandTimeoutMs: 10_000 }, received, requests: () => requests };
+	}
+
+	it("命令取来的头在首次请求前就已生效，且每次连接只跑一次", async () => {
+		const counter = counterPath();
+		const gated = await startTokenGatedServer("Bearer token-1");
+		const client = await McpClient.connect({
+			...gated.config,
+			headersCommand: `${process.execPath} ${TOKEN_HELPER} --file ${counter}`,
+		});
+		clients.push(client);
+		const tools = await client.listTools();
+		assert.equal(tools[0]?.name, "ping");
+		// initialize / notifications / tools-list 三个请求带的是同一个 token：命令只在连接时跑了一次。
+		assert.ok(gated.received.every((value) => value === "Bearer token-1"));
+		assert.equal(readFileSync(counter, "utf8").trim(), "1");
+	});
+
+	it("401 后重跑命令，头变了就用新头重试一次", async () => {
+		const counter = counterPath();
+		const gated = await startTokenGatedServer("Bearer token-2");
+		const client = await McpClient.connect({
+			...gated.config,
+			headersCommand: `${process.execPath} ${TOKEN_HELPER} --file ${counter}`,
+		});
+		clients.push(client);
+		assert.equal(client.serverInfo.name, "gated");
+		assert.deepEqual(gated.received.slice(0, 2), ["Bearer token-1", "Bearer token-2"]);
+		assert.equal(readFileSync(counter, "utf8").trim(), "2");
+	});
+
+	it("命令每次返回同一个头时不重试", async () => {
+		const gated = await startTokenGatedServer("Bearer never-matches");
+		await assert.rejects(
+			() =>
+				McpClient.connect({
+					...gated.config,
+					headersCommand: `${process.execPath} ${TOKEN_HELPER} --fixed token-1`,
+				}),
+			(error: unknown) => error instanceof McpConnectionError && /HTTP 401/.test(error.message),
+		);
+		assert.equal(gated.requests(), 1, "头没变化就不该重试");
+	});
+
+	it("头命令失败：记诊断、退回静态 headers 继续连", async () => {
+		const diagnostics: string[] = [];
+		const gated = await startTokenGatedServer("Bearer static-token");
+		const client = await McpClient.connect(
+			{
+				...gated.config,
+				headers: { Authorization: "Bearer static-token" },
+				headersCommand: `${process.execPath} ${TOKEN_HELPER} --fail`,
+			},
+			{ onDiagnostic: (line) => diagnostics.push(line) },
+		);
+		clients.push(client);
+		assert.equal(client.serverInfo.name, "gated");
+		assert.ok(
+			// fixture 往 stderr 写的是 "token 服务连不上"，runCommand 把它拼进错误信息。
+			diagnostics.some((line) => line.includes("头命令失败") && line.includes("token 服务连不上")),
+			`诊断里应记录头命令失败：${diagnostics.join(" | ")}`,
+		);
+	});
+
+	it("被 401 拒绝时，错误信息里包含头命令的失败原因", async () => {
+		const gated = await startTokenGatedServer("Bearer static-token");
+		await assert.rejects(
+			() =>
+				McpClient.connect({
+					...gated.config,
+					headers: { Authorization: "Bearer wrong-static" },
+					headersCommand: `${process.execPath} ${TOKEN_HELPER} --fail`,
+				}),
+			(error: unknown) => {
+				const message = error instanceof Error ? error.message : String(error);
+				assert.match(message, /HTTP 401/);
+				assert.match(message, /头命令失败/);
+				return true;
+			},
+		);
+	});
+
+	it("`Key: Value` 行格式的命令输出也能用", async () => {
+		const gated = await startTokenGatedServer("Bearer token-fixed");
+		const client = await McpClient.connect({
+			...gated.config,
+			headersCommand: `${process.execPath} ${TOKEN_HELPER} --lines`,
+		});
+		clients.push(client);
+		assert.equal(client.serverInfo.name, "gated");
+	});
+
+	it("旧版 SSE：GET 长连接与 POST 都带上动态头", async () => {
+		const streams = new Set<ServerResponse>();
+		const http = await startHttpServer((req, res, body) => {
+			if ((req.headers["authorization"] ?? "") !== "Bearer token-1") {
+				res.writeHead(401).end();
+				return;
+			}
+			if (req.method === "GET") {
+				res.writeHead(200, { "content-type": "text/event-stream" });
+				res.write("event: endpoint\ndata: /messages\n\n");
+				streams.add(res);
+				res.on("close", () => streams.delete(res));
+				return;
+			}
+			const message = body as { id?: number; method?: string };
+			res.writeHead(202).end();
+			if (message.method === "notifications/initialized") return;
+			const result =
+				message.method === "initialize"
+					? { protocolVersion: "2024-11-05", capabilities: {}, serverInfo: { name: "legacy-gated", version: "1.0" } }
+					: { tools: [] };
+			for (const stream of streams) {
+				stream.write(`data: ${JSON.stringify({ jsonrpc: "2.0", id: message.id, result })}\n\n`);
+			}
+		});
+		const client = await McpClient.connect({
+			...http.config,
+			transport: "sse",
+			name: "legacy-gated",
+			headersCommand: `${process.execPath} ${TOKEN_HELPER} --fixed token-1`,
+		});
+		clients.push(client);
+		assert.equal(client.serverInfo.name, "legacy-gated");
+		assert.deepEqual(await client.listTools(), []);
+	});
+});
+
+interface HttpFixture {
+	config: McpRemoteServer;
+	close: () => Promise<void>;
+}
+
+type HttpHandler = (req: IncomingMessage, res: ServerResponse, body: unknown) => void;
+
+async function startHttpServer(handler: HttpHandler): Promise<HttpFixture> {
+	const server = createServer((req, res) => {
+		const chunks: Buffer[] = [];
+		req.on("data", (chunk: Buffer) => chunks.push(chunk));
+		req.on("end", () => {
+			const raw = Buffer.concat(chunks).toString("utf8");
+			let body: unknown;
+			try {
+				body = raw ? JSON.parse(raw) : undefined;
+			} catch {
+				body = raw;
+			}
+			handler(req, res, body);
+		});
+	});
+	servers.push(server);
+	await new Promise<void>((resolve) => server.listen(0, "127.0.0.1", resolve));
+	const { port } = server.address() as AddressInfo;
+	const base = `http://127.0.0.1:${port}`;
+	return {
+		config: {
+			name: "http",
+			transport: "http",
+			url: `${base}/mcp`,
+			headers: {},
+			timeoutMs: 5000,
+			enabled: true,
+			source: "test",
+		},
+		close: () => new Promise<void>((resolve) => server.close(() => resolve())),
+	};
+}
+
+function respondJson(res: ServerResponse, status: number, payload: unknown): void {
+	res.writeHead(status, { "content-type": "application/json" }).end(JSON.stringify(payload));
+}
+
+/** 轮询等待条件成立，避免给超时/退出这类异步动作硬编码 sleep。 */
+async function waitFor(predicate: () => boolean, timeoutMs = 3000): Promise<void> {
+	const deadline = Date.now() + timeoutMs;
+	while (Date.now() < deadline) {
+		if (predicate()) return;
+		await new Promise((resolve) => setTimeout(resolve, 10));
+	}
+	throw new Error("waitFor 超时");
+}