@bachi/pi-coder 1.0.0 → 1.1.0

package/extensions/mcp/index.ts

@@ -0,0 +1,295 @@
+/**
+ * mcp — 让 pi 用上本机 / 远程的 MCP 服务器（stdio、streamable HTTP、旧版 SSE）。
+ *
+ * 设计取舍（对齐用户「最简单够用」的要求）：
+ *   - **每个 MCP 工具直接注册成一个 pi 工具**，名字 `mcp__<server>__<tool>`（Claude Code 同款），
+ *     不走「一个 mcp 代理工具 + 参数指定 server/tool」那套（pi-mcp-adapter 的做法）——
+ *     工具少时直连对模型友好得多，代价只是 system prompt 变长。
+ *   - **配置沿用 Claude Code 的 `.mcp.json` 形状**，全局 `~/.pi/agent/mcp.json` + 项目
+ *     `.mcp.json`（项目同名覆盖全局）。用户已有的 `.mcp.json` 抄一份即可。
+ *   - **会话开始时连接、会话结束断开**。工具表必须先 `tools/list` 才能注册，所以不能等到
+ *     首次调用才连；多个 server 并行握手，单个 server 失败只影响它自己。
+ *   - 诊断信息（子进程 stderr、协议异常）**只进内存环形缓冲**，不写 stdout/stderr ——
+ *     interactive pi 里往 stderr 写会糊在输入框上（见仓库里 subagent-log-guard 的存在理由）。
+ *     要看就用 `/mcp` 或 `/mcp <server>`。
+ *
+ * 命令：
+ *   /mcp            当前 server / 工具 / 配置来源一览
+ *   /mcp reload     重新读配置、重连、重注册工具（改完 mcp.json 不用重启 pi）
+ *   /mcp <server>   单个 server 的详情与最近诊断输出
+ */
+
+import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
+import { Type } from "typebox";
+
+import { McpClient, type McpToolInfo } from "./client.ts";
+import { loadMcpConfig, type McpConfigIssue, type McpServerConfig } from "./config.ts";
+import { McpConnectionError, McpError } from "./protocol.ts";
+import {
+	mcpContentToPiContent,
+	normalizeInputSchema,
+	piToolName,
+	toolDescription,
+	toolPromptSnippet,
+	truncateText,
+	type PiToolContent,
+} from "./tools.ts";
+
+/** 每个 server 保留多少行诊断信息（stderr / 协议异常）。 */
+const DIAGNOSTIC_KEEP_LINES = 20;
+
+type ServerStatus = "ready" | "disabled" | "error";
+
+interface ServerState {
+	config: McpServerConfig;
+	status: ServerStatus;
+	client?: McpClient;
+	tools: McpToolInfo[];
+	error?: string;
+	diagnostics: string[];
+}
+
+interface McpRuntime {
+	cwd: string;
+	sources: string[];
+	issues: McpConfigIssue[];
+	servers: ServerState[];
+}
+
+let runtime: McpRuntime | undefined;
+
+export default function (pi: ExtensionAPI) {
+	pi.on("session_start", async (_event, ctx) => {
+		await stopRuntime();
+		runtime = await startRuntime(pi, ctx);
+		const failures = runtime.servers.filter((server) => server.status === "error");
+		const problems = failures.length + runtime.issues.length;
+		if (problems > 0) {
+			const ready = runtime.servers.filter((server) => server.status === "ready").length;
+			const detail = [
+				`MCP：${runtime.servers.length} 个 server，${ready} 个就绪`,
+				...failures.map((server) => `${server.config.name}: ${server.error ?? "连接失败"}`),
+				...runtime.issues.map((issue) => `${issue.server ?? issue.source}: ${issue.message}`),
+				"用 /mcp 查看详情",
+			].join("\n");
+			ctx.ui.notify(detail, "warning");
+		}
+	});
+
+	pi.on("session_shutdown", async () => {
+		await stopRuntime();
+	});
+
+	pi.registerCommand("mcp", {
+		description: "MCP 服务器状态；/mcp reload 重连；/mcp <server> 看单个 server 详情",
+		handler: async (args, ctx) => {
+			const argument = args.trim();
+			if (argument === "reload") {
+				if (runtime) ctx.ui.notify(`正在重连 ${runtime.servers.length} 个 MCP server…`, "info");
+				await stopRuntime();
+				runtime = await startRuntime(pi, ctx);
+				ctx.ui.notify(formatStatus(runtime, ctx), runtimeHasProblems(runtime) ? "warning" : "info");
+				return;
+			}
+			if (!runtime) {
+				ctx.ui.notify("MCP 尚未初始化（本会话启动时没有可用配置）", "warning");
+				return;
+			}
+			if (argument) {
+				const server = runtime.servers.find((candidate) => candidate.config.name === argument);
+				if (!server) {
+					ctx.ui.notify(`没有名为 "${argument}" 的 MCP server。用 /mcp 查看已配置的 server。`, "warning");
+					return;
+				}
+				ctx.ui.notify(formatServerDetail(server), server.status === "error" ? "error" : "info");
+				return;
+			}
+			ctx.ui.notify(formatStatus(runtime, ctx), runtimeHasProblems(runtime) ? "warning" : "info");
+		},
+	});
+}
+
+async function startRuntime(pi: ExtensionAPI, ctx: ExtensionContext): Promise<McpRuntime> {
+	const loaded = loadMcpConfig({ cwd: ctx.cwd });
+	const servers: ServerState[] = loaded.servers.map((config) => ({
+		config,
+		status: config.enabled ? "error" : "disabled",
+		tools: [],
+		error: config.enabled ? undefined : "配置里已禁用（enabled: false）",
+		diagnostics: [],
+	}));
+
+	const enabled = servers.filter((server) => server.config.enabled);
+	// connectServer 只写状态、不抛异常：一个 server 挂了不能影响其它 server 和 pi 本体。
+	await Promise.all(enabled.map((server) => connectServer(pi, server)));
+
+	return { cwd: ctx.cwd, sources: loaded.sources, issues: loaded.issues, servers };
+}
+
+async function connectServer(pi: ExtensionAPI, state: ServerState): Promise<void> {
+	const { config } = state;
+	try {
+		const client = await McpClient.connect(config, {
+			onDiagnostic: (line) => {
+				state.diagnostics.push(line);
+				if (state.diagnostics.length > DIAGNOSTIC_KEEP_LINES) state.diagnostics.shift();
+			},
+		});
+		state.client = client;
+		state.tools = await client.listTools();
+		state.status = "ready";
+		state.error = undefined;
+		registerTools(pi, state);
+	} catch (error) {
+		state.status = "error";
+		state.error = describeError(error);
+	}
+}
+
+/**
+ * 把一个 MCP server 的工具全部注册成 pi 工具。
+ *
+ * 同名覆盖是 pi 的行为（`extension.tools` 是 Map），所以 `/mcp reload` 重复注册不会留下旧定义。
+ */
+function registerTools(pi: ExtensionAPI, state: ServerState): void {
+	const { config } = state;
+	for (const tool of state.tools) {
+		const name = piToolName(config.name, tool.name);
+		pi.registerTool({
+			name,
+			label: `MCP: ${tool.name}`,
+			description: toolDescription(config.name, tool),
+			promptSnippet: toolPromptSnippet(tool.description),
+			// MCP 给的是标准 JSON Schema；Type.Unsafe 让 pi 原样用它做参数校验。
+			parameters: Type.Unsafe(normalizeInputSchema(tool.inputSchema)),
+			async execute(_toolCallId, params, signal) {
+				const client = state.client;
+				if (!client || client.isClosed) {
+					throw new Error(
+						`MCP server "${config.name}" 未连接${state.error ? `（${state.error}）` : ""}。运行 /mcp reload 重连。`,
+					);
+				}
+				const started = Date.now();
+				const args = (params ?? {}) as Record<string, unknown>;
+				const result = await client.callTool(tool.name, args, { signal });
+				const mapped = mcpContentToPiContent(result.content);
+				const truncated = truncateText(mapped.text);
+				const durationMs = Date.now() - started;
+				const content = buildToolContent(mapped.content, truncated.text, truncated.truncated);
+
+				if (result.isError) {
+					throw new Error(
+						`MCP 工具 ${config.name}/${tool.name} 返回错误：\n${truncated.text || "(无输出)"}`,
+					);
+				}
+
+				return {
+					content,
+					details: {
+						server: config.name,
+						tool: tool.name,
+						durationMs,
+						truncated: truncated.truncated,
+						totalBytes: truncated.totalBytes,
+						totalLines: truncated.totalLines,
+						notes: mapped.notes,
+					},
+				};
+			},
+		});
+	}
+}
+
+/**
+ * 截断只作用于文本：图片块必须原样留下。
+ * 未截断时保持 MCP 给的块顺序（文本与图片的相对位置有意义），截断时退化成「一段文本 + 图片」。
+ */
+function buildToolContent(blocks: PiToolContent[], truncatedText: string, truncated: boolean): PiToolContent[] {
+	if (blocks.length === 0) return [{ type: "text", text: "(MCP 工具返回空内容)" }];
+	if (!truncated) return blocks;
+	return [{ type: "text", text: truncatedText }, ...blocks.filter((block) => block.type === "image")];
+}
+
+async function stopRuntime(): Promise<void> {
+	const current = runtime;
+	runtime = undefined;
+	if (!current) return;
+	await Promise.all(current.servers.map((server) => server.client?.close().catch(() => {})));
+}
+
+function runtimeHasProblems(rt: McpRuntime): boolean {
+	return rt.issues.length > 0 || rt.servers.some((server) => server.status === "error");
+}
+
+function formatStatus(rt: McpRuntime, ctx: ExtensionContext): string {
+	const ready = rt.servers.filter((server) => server.status === "ready");
+	const toolCount = ready.reduce((total, server) => total + server.tools.length, 0);
+	const lines: string[] = [`MCP：${ready.length}/${rt.servers.length} 个 server 就绪，共 ${toolCount} 个工具`];
+	if (rt.servers.length === 0) {
+		lines.push(`(未发现配置。写一个 ~/.pi/agent/mcp.json 或项目根的 .mcp.json，然后 /mcp reload)`);
+	}
+	for (const server of rt.servers) {
+		const mark = server.status === "ready" ? "●" : server.status === "disabled" ? "○" : "✗";
+		const parts: string[] = [];
+		if (server.client) {
+			parts.push(shortTransport(server.config));
+			parts.push(`v${server.client.serverInfo.version ?? "?"}`);
+		}
+		if (server.status === "ready") parts.push(`${server.tools.length} tools`);
+		if (server.config.transport !== "stdio" && server.config.headersCommand) parts.push("headersCommand");
+		if (server.status === "error" && server.error) parts.push(firstLine(server.error));
+		lines.push(`${mark} ${server.config.name}${parts.length > 0 ? ` · ${parts.join(" · ")}` : ""}`);
+	}
+	for (const issue of rt.issues) {
+		lines.push(`⚠ ${issue.server ? `${issue.server}: ` : ""}${issue.message}${issue.server ? "" : ` (${issue.source})`}`);
+	}
+	if (rt.sources.length > 0) lines.push(`配置：${rt.sources.join("  ")}`);
+	lines.push(`工作目录：${rt.cwd}`);
+	return lines.join("\n");
+}
+
+function formatServerDetail(server: ServerState): string {
+	const lines: string[] = [];
+	lines.push(`${server.config.name}（${server.config.transport}，${server.status}）`);
+	const client = server.client;
+	if (client) {
+		lines.push(`server: ${client.serverInfo.name ?? "?"} ${client.serverInfo.version ?? ""} · 协议 ${client.protocolVersion}`);
+		lines.push(`连接: ${client.transportLabel}`);
+		lines.push(`能力: ${Object.keys(client.capabilities).join(", ") || "(无)"}`);
+	}
+	if (server.error) lines.push(`错误: ${server.error}`);
+	if (server.config.transport !== "stdio" && server.config.headersCommand) {
+		// 只展示命令本身（用户自己写的配置），头的**值**任何情况下都不打印。
+		const command = server.config.headersCommand;
+		lines.push(`头命令: ${command.length > 100 ? `${command.slice(0, 100)}…` : command}`);
+	}
+	if (server.tools.length > 0) {
+		lines.push(`工具（${server.tools.length}）:`);
+		for (const tool of server.tools) {
+			const short = tool.description?.replace(/\s+/g, " ").slice(0, 60) ?? "";
+			lines.push(`  mcp__${server.config.name}__${tool.name}${short ? ` — ${short}` : ""}`);
+		}
+	}
+	if (server.diagnostics.length > 0) {
+		lines.push(`最近诊断（${server.diagnostics.length} 行）:`);
+		for (const line of server.diagnostics.slice(-8)) lines.push(`  ${firstLine(line).slice(0, 160)}`);
+	}
+	return lines.join("\n");
+}
+
+function describeError(error: unknown): string {
+	if (error instanceof McpConnectionError) return error.message;
+	if (error instanceof McpError) return `${error.message}（code ${error.code}）`;
+	if (error instanceof Error) return error.message;
+	return String(error);
+}
+
+function firstLine(text: string): string {
+	const index = text.indexOf("\n");
+	return index === -1 ? text : text.slice(0, index);
+}
+
+/** 状态行里的短传输标识：stdio 只写 "stdio"，完整命令留给 `/mcp <server>` 与报错信息。 */
+function shortTransport(config: McpServerConfig): string {
+	return config.transport === "stdio" ? "stdio" : `${config.transport} ${config.url}`;
+}

package/extensions/mcp/protocol.test.ts

@@ -0,0 +1,179 @@
+/**
+ * Tests for protocol.ts — JSON-RPC 编解码与 SSE 分帧。
+ *
+ * Run with:  node --test clients/pi/extensions/mcp/protocol.test.ts
+ *
+ * 这里只测「字节流 → 消息」的边界：半个 JSON、半个 SSE 事件、跨 chunk 的字段、\r\n。
+ * 真实服务端的握手成功与否由 client.test.ts 的 stdio 端到端用例负责。
+ */
+
+import assert from "node:assert/strict";
+import { describe, it } from "node:test";
+
+import {
+	createCancelledNotification,
+	createLineDecoder,
+	createSseDecoder,
+	encodeStdioFrame,
+	isJsonRpcRequestOrNotification,
+	isJsonRpcResponse,
+	JSONRPC_PARSE_ERROR,
+	McpError,
+	parseJsonRpcMessage,
+	parseSseBlock,
+	toErrorFromResponse,
+} from "./protocol.ts";
+
+describe("createLineDecoder", () => {
+	it("把一个 chunk 里的多行都切出来", () => {
+		const decode = createLineDecoder();
+		assert.deepEqual(decode("a\nb\nc\n"), ["a", "b", "c"]);
+	});
+
+	it("半个 JSON 留在缓冲里，等下一块拼上再吐", () => {
+		const decode = createLineDecoder();
+		assert.deepEqual(decode('{"jsonrpc":"2.0",'), []);
+		assert.deepEqual(decode('"id":1}\n'), ['{"jsonrpc":"2.0","id":1}']);
+	});
+
+	it("兼容 \\r\\n", () => {
+		const decode = createLineDecoder();
+		assert.deepEqual(decode("a\r\nb\r\n"), ["a", "b"]);
+	});
+
+	it("空行原样返回（调用方负责跳过）", () => {
+		const decode = createLineDecoder();
+		assert.deepEqual(decode("a\n\nb\n"), ["a", "", "b"]);
+	});
+
+	it("一行跨三个 chunk 也能拼出来", () => {
+		const decode = createLineDecoder();
+		decode("hel");
+		decode("lo ");
+		assert.deepEqual(decode("world\n"), ["hello world"]);
+	});
+});
+
+describe("createSseDecoder", () => {
+	it("按空行切事件，data 多行用 \\n 拼接", () => {
+		const decode = createSseDecoder();
+		const events = decode("event: message\ndata: a\ndata: b\n\n");
+		assert.equal(events.length, 1);
+		assert.equal(events[0]?.event, "message");
+		assert.equal(events[0]?.data, "a\nb");
+	});
+
+	it("没有 event 字段的事件也算（MCP 的响应就是纯 data）", () => {
+		const decode = createSseDecoder();
+		const events = decode('data: {"jsonrpc":"2.0","id":1,"result":{}}\n\n');
+		assert.equal(events.length, 1);
+		assert.equal(events[0]?.event, undefined);
+	});
+
+	it("半个事件留在缓冲里", () => {
+		const decode = createSseDecoder();
+		assert.deepEqual(decode('data: {"a":'), []);
+		assert.deepEqual(decode('1}\n\n'), [{ event: undefined, data: '{"a":1}', id: undefined }]);
+	});
+
+	it("忽略注释行（心跳）", () => {
+		const decode = createSseDecoder();
+		assert.deepEqual(decode(": ping\n\n"), []);
+	});
+
+	it("\\r\\n 分隔的事件也能切", () => {
+		const decode = createSseDecoder();
+		const events = decode("data: x\r\n\r\n");
+		assert.equal(events[0]?.data, "x");
+	});
+
+	it("一个 chunk 里的多个事件都吐出来", () => {
+		const decode = createSseDecoder();
+		const events = decode("data: 1\n\ndata: 2\n\n");
+		assert.deepEqual(events.map((event) => event.data), ["1", "2"]);
+	});
+
+	it("id 字段被保留", () => {
+		const decode = createSseDecoder();
+		assert.equal(decode("id: 42\ndata: x\n\n")[0]?.id, "42");
+	});
+});
+
+describe("parseSseBlock", () => {
+	it("只有 data 的块直接可用", () => {
+		assert.deepEqual(parseSseBlock("data: hello"), { event: undefined, data: "hello", id: undefined });
+	});
+
+	it("去掉冒号后的单个空格，保留后续空格", () => {
+		assert.equal(parseSseBlock("data:  two spaces")?.data, " two spaces");
+	});
+
+	it("没有 data 的块返回 undefined", () => {
+		assert.equal(parseSseBlock("event: ping"), undefined);
+	});
+});
+
+describe("encodeStdioFrame", () => {
+	it("一行 JSON 加换行，消息里不会出现裸换行", () => {
+		const frame = encodeStdioFrame({ jsonrpc: "2.0", id: 1, method: "x", params: { text: "a\nb" } });
+		assert.ok(frame.endsWith("\n"));
+		assert.equal(frame.trimEnd().split("\n").length, 1);
+	});
+});
+
+describe("parseJsonRpcMessage", () => {
+	it("解析成功", () => {
+		const message = parseJsonRpcMessage('{"jsonrpc":"2.0","id":1,"result":{"ok":true}}');
+		assert.ok(isJsonRpcResponse(message));
+	});
+
+	it("坏 JSON 抛带 parse error 码的 McpError", () => {
+		assert.throws(
+			() => parseJsonRpcMessage("{oops"),
+			(error: unknown) => error instanceof McpError && error.code === JSONRPC_PARSE_ERROR,
+		);
+	});
+
+	it("数组/字面量不是合法 JSON-RPC 消息", () => {
+		assert.throws(() => parseJsonRpcMessage("[1,2]"), McpError);
+		assert.throws(() => parseJsonRpcMessage("42"), McpError);
+	});
+});
+
+describe("isJsonRpcResponse / isJsonRpcRequestOrNotification", () => {
+	it("按 method 字段区分响应与请求", () => {
+		assert.ok(isJsonRpcResponse({ jsonrpc: "2.0", id: 1, result: {} } as never));
+		assert.ok(isJsonRpcResponse({ jsonrpc: "2.0", id: 1, error: { code: 1, message: "x" } } as never));
+		assert.ok(!isJsonRpcResponse({ jsonrpc: "2.0", id: 1, method: "tools/list" } as never));
+		assert.ok(isJsonRpcRequestOrNotification({ jsonrpc: "2.0", method: "notifications/initialized" } as never));
+	});
+});
+
+describe("toErrorFromResponse", () => {
+	it("把 data 拼进 message", () => {
+		const error = toErrorFromResponse({ code: -32602, message: "bad params", data: { field: "x" } });
+		assert.equal(error.code, -32602);
+		assert.match(error.message, /bad params/);
+		assert.match(error.message, /field/);
+	});
+
+	it("data 过长时截断，别让错误信息本身撑爆日志", () => {
+		const error = toErrorFromResponse({ code: 1, message: "x", data: "y".repeat(2000) });
+		assert.ok(error.message.length < 600);
+		assert.match(error.message, /…/);
+	});
+
+	it("循环引用的 data 不会让拼消息本身抛异常", () => {
+		const data: Record<string, unknown> = {};
+		data.self = data;
+		assert.doesNotThrow(() => toErrorFromResponse({ code: 1, message: "x", data }));
+	});
+});
+
+describe("createCancelledNotification", () => {
+	it("形状符合 notifications/cancelled", () => {
+		const notification = createCancelledNotification(7, "timeout");
+		assert.equal(notification.method, "notifications/cancelled");
+		assert.deepEqual(notification.params, { requestId: 7, reason: "timeout" });
+	});
+});