my-pi/extensions/codegraph-client.ts

/**
 * Minimal stdio JSON-RPC 2.0 client for `codegraph serve --mcp -p <root>`.
 *
 * Same approach as pi-web-search/clients/exa.ts (plain JSON-RPC, no MCP
 * adapter or SDK): `codegraph serve --mcp` speaks one JSON object per line on
 * stdio. Unlike the one-shot `codegraph explore` CLI, the serve process runs a
 * file watcher that keeps the index synced, so newly written symbols become
 * searchable without a manual `codegraph sync`.
 *
 * One client exists per index root; `extensions/codegraph.ts` owns the map,
 * idle closing, and the one-shot respawn retry.
 */

import {
	spawn as nodeSpawn,
	type ChildProcess,
	type SpawnOptions,
} from "node:child_process";

const PROTOCOL_VERSION = "2025-06-18";
const DEFAULT_REQUEST_TIMEOUT_MS = 120_000;
const STDERR_TAIL_LINES = 20;
const CLOSE_GRACE_MS = 1_000;

export type CodegraphSpawn = (
	command: string,
	args: string[],
	options: SpawnOptions,
) => ChildProcess;

export interface CodegraphRpcClientOptions {
	/** Index root — passed to `codegraph serve --mcp -p <root>` and used as cwd. */
	root: string;
	/** Test seam: replace the child_process spawn. */
	spawn?: CodegraphSpawn;
	/** Default "codegraph". */
	command?: string;
	/** Default ["serve", "--mcp", "-p", root]. */
	args?: string[];
	/** Per-request timeout; default 120 s. */
	requestTimeoutMs?: number;
}

interface PendingRequest {
	resolve: (result: unknown) => void;
	reject: (error: Error) => void;
}

function isRecord(value: unknown): value is Record<string, unknown> {
	return typeof value === "object" && value !== null && !Array.isArray(value);
}

function contentText(result: Record<string, unknown>): string {
	const content = Array.isArray(result.content)
		? (result.content as Record<string, unknown>[])
		: [];
	return content
		.filter((part) => part?.type === "text" && typeof part.text === "string")
		.map((part) => part.text as string)
		.join("\n")
		.trim();
}

export class CodegraphRpcClient {
	private readonly spawnFn: CodegraphSpawn;
	private readonly command: string;
	private readonly args: string[];
	private readonly timeoutMs: number;

	private child: ChildProcess | undefined;
	private stdoutBuffer = "";
	private stderrLines: string[] = [];
	private pending = new Map<number, PendingRequest>();
	private nextId = 1;
	private startPromise: Promise<void> | undefined;
	/** null = upstream tool declares no inputSchema.properties → forward all args. */
	private schemaProps: ReadonlySet<string> | null = null;
	private deadFlag = false;
	private closedFlag = false;
	private deathFinalized = false;
	private exitDetail = "";
	private deathReason = "codegraph serve 子进程已退出";

	private readonly options: CodegraphRpcClientOptions;

	constructor(options: CodegraphRpcClientOptions) {
		this.options = options;
		this.spawnFn = options.spawn ?? nodeSpawn;
		this.command = options.command ?? "codegraph";
		this.args =
			options.args ?? ["serve", "--mcp", "-p", options.root];
		this.timeoutMs = options.requestTimeoutMs ?? DEFAULT_REQUEST_TIMEOUT_MS;
	}

	/** True once the child has exited or errored — the handler respawns instead. */
	get dead(): boolean {
		return this.deadFlag;
	}

	/** Property names of the upstream codegraph_explore inputSchema. */
	schemaPropertyNames(): readonly string[] {
		return this.schemaProps === null ? [] : [...this.schemaProps];
	}

	/** Spawn + initialize + tools/list. Memoized: concurrent callers share it. */
	start(): Promise<void> {
		if (!this.startPromise) {
			this.startPromise = this.doStart();
			// A dead child during startup rejects this promise; attach a no-op
			// branch so it never surfaces as an unhandled rejection when a
			// racing close() kills the child with nobody awaiting start().
			this.startPromise.catch(() => {});
		}
		return this.startPromise;
	}

	private async doStart(): Promise<void> {
		let child: ChildProcess;
		try {
			child = this.spawnFn(this.command, this.args, {
				cwd: this.options.root,
				stdio: ["pipe", "pipe", "pipe"],
			});
		} catch (error) {
			this.markDead(
				`无法启动 codegraph serve：${error instanceof Error ? error.message : String(error)}`,
			);
			throw this.deathError();
		}
		this.child = child;
		this.attach(child);
		await this.request("initialize", {
			protocolVersion: PROTOCOL_VERSION,
			capabilities: {},
			clientInfo: { name: "my-pi-codegraph", version: "1.0.0" },
		});
		this.notify("notifications/initialized");
		const list = (await this.request("tools/list", {})) as Record<
			string,
			unknown
		>;
		const tools = Array.isArray(list.tools) ? list.tools : [];
		const tool = tools.find(
			(entry) => isRecord(entry) && entry.name === "codegraph_explore",
		);
		if (!tool) {
			throw new Error("codegraph serve 未提供 codegraph_explore 工具");
		}
		const properties =
			isRecord(tool) && isRecord(tool.inputSchema)
				? tool.inputSchema.properties
				: undefined;
		this.schemaProps = isRecord(properties)
			? new Set(Object.keys(properties))
			: null;
	}

	/**
	 * `tools/call codegraph_explore`. Forwards only the argument properties the
	 * upstream inputSchema declares (e.g. `projectPath` is forwarded only when
	 * upstream lists it). Returns the raw result object; the handler passes
	 * `result.content` through verbatim.
	 */
	async callExplore(
		args: Record<string, unknown>,
		signal?: AbortSignal,
	): Promise<Record<string, unknown>> {
		await this.start();
		const forwarded: Record<string, unknown> = {};
		for (const [key, value] of Object.entries(args)) {
			if (this.schemaProps === null || this.schemaProps.has(key)) {
				forwarded[key] = value;
			}
		}
		const result = await this.request(
			"tools/call",
			{ name: "codegraph_explore", arguments: forwarded },
			signal,
		);
		if (isRecord(result) && result.isError) {
			throw new Error(
				`codegraph_explore 返回错误：${contentText(result) || "unknown"}`,
			);
		}
		return isRecord(result) ? result : {};
	}

	/** End stdin, then SIGTERM after a 1 s grace. Idempotent. */
	async close(): Promise<void> {
		if (this.closedFlag) return;
		this.closedFlag = true;
		const child = this.child;
		this.failAll(new Error("codegraph serve 客户端已关闭"));
		if (!child) return;
		try {
			child.stdin?.end();
		} catch {
			// stdin already closed.
		}
		if (child.exitCode !== null || child.signalCode !== null) return;
		await new Promise<void>((resolve) => {
			const timer = setTimeout(() => {
				try {
					child.kill("SIGTERM");
				} catch {
					// Already gone.
				}
				resolve();
			}, CLOSE_GRACE_MS);
			timer.unref();
			child.once("exit", () => {
				clearTimeout(timer);
				resolve();
			});
			child.once("error", () => {
				clearTimeout(timer);
				resolve();
			});
		});
	}

	// ── internals ────────────────────────────────────────────────────────────

	private attach(child: ChildProcess): void {
		// EPIPE on a write after the child died is reported via exit/close.
		child.stdin?.on("error", () => {});
		child.stdout?.setEncoding("utf8");
		child.stdout?.on("data", (chunk: string) => {
			this.stdoutBuffer += chunk;
			for (;;) {
				const newline = this.stdoutBuffer.indexOf("\n");
				if (newline < 0) break;
				const line = this.stdoutBuffer.slice(0, newline).trim();
				this.stdoutBuffer = this.stdoutBuffer.slice(newline + 1);
				if (line) this.handleLine(line);
			}
		});
		child.stderr?.setEncoding("utf8");
		child.stderr?.on("data", (chunk: string) => {
			for (const line of chunk.split(/\r?\n/u)) {
				if (!line.trim()) continue;
				this.stderrLines.push(line);
				if (this.stderrLines.length > STDERR_TAIL_LINES) {
					this.stderrLines.shift();
				}
			}
		});
		child.once("error", (error) => {
			// Spawn failure (e.g. ENOENT) — the message carries the detail.
			this.markDead(`无法启动 codegraph serve：${error.message}`);
		});
		child.once("exit", (code, signal) => {
			this.deadFlag = true;
			this.exitDetail =
				signal !== null ? `signal ${signal}` : `code ${code}`;
		});
		// Reject pending on "close" (stdio drained), not "exit" — otherwise the
		// final stderr lines may not have been delivered yet.
		child.once("close", () => {
			this.markDead(`codegraph serve 子进程退出（${this.exitDetail}）`);
		});
	}

	private handleLine(line: string): void {
		let message: unknown;
		try {
			message = JSON.parse(line);
		} catch {
			return; // Not JSON — server chatter; ignore.
		}
		if (!isRecord(message)) return;

		// Server → client request: we implement no methods; reply -32601.
		if (typeof message.method === "string" && message.id !== undefined) {
			this.sendRaw({
				jsonrpc: "2.0",
				id: message.id,
				error: { code: -32601, message: "Method not found" },
			});
			return;
		}
		// Response to one of our requests.
		if (message.id !== undefined) {
			const entry = this.pending.get(Number(message.id));
			if (!entry) return;
			if (message.error !== undefined) {
				const detail = isRecord(message.error)
					? (message.error.message ?? JSON.stringify(message.error))
					: String(message.error);
				entry.reject(new Error(`codegraph serve ${String(detail)}`));
			} else {
				entry.resolve(message.result);
			}
		}
		// Anything else (server notifications) is ignored.
	}

	private request(
		method: string,
		params: Record<string, unknown>,
		signal?: AbortSignal,
	): Promise<unknown> {
		if (this.deadFlag || this.closedFlag) {
			return Promise.reject(this.deathError());
		}
		const stdin = this.child?.stdin;
		if (!this.child || !stdin || stdin.destroyed) {
			return Promise.reject(new Error("codegraph serve 子进程不可用"));
		}
		if (signal?.aborted) {
			return Promise.reject(new Error("codegraph serve 请求已中止"));
		}
		const id = this.nextId++;
		return new Promise<unknown>((resolve, reject) => {
			let onAbort: (() => void) | undefined;
			const finish = (action: () => void): void => {
				clearTimeout(timer);
				if (onAbort) signal?.removeEventListener("abort", onAbort);
				this.pending.delete(id);
				action();
			};
			const timer = setTimeout(() => {
				finish(() =>
					reject(new Error(`codegraph serve 请求 ${method} 超时`)),
				);
			}, this.timeoutMs);
			timer.unref();
			onAbort = () => {
				this.notify("notifications/cancelled", {
					requestId: id,
					reason: "aborted",
				});
				finish(() => reject(new Error("codegraph serve 请求已中止")));
			};
			if (signal) {
				signal.addEventListener("abort", onAbort, { once: true });
			}
			this.pending.set(id, {
				resolve: (result) => finish(() => resolve(result)),
				reject: (error) => finish(() => reject(error)),
			});
			this.sendRaw({ jsonrpc: "2.0", id, method, params });
		});
	}

	private notify(method: string, params?: Record<string, unknown>): void {
		this.sendRaw({ jsonrpc: "2.0", method, params });
	}

	private sendRaw(message: Record<string, unknown>): void {
		try {
			this.child?.stdin?.write(`${JSON.stringify(message)}\n`);
		} catch {
			// stdin already closed — the exit handler will fail pending calls.
		}
	}

	private markDead(reason: string): void {
		if (this.deathFinalized) return;
		this.deathFinalized = true;
		this.deadFlag = true;
		const tail = this.stderrLines.join("\n").trim();
		this.deathReason = tail ? `${reason}；stderr: ${tail}` : reason;
		this.failAll(this.deathError());
	}

	private deathError(): Error {
		return new Error(
			this.closedFlag && !this.deadFlag
				? "codegraph serve 客户端已关闭"
				: this.deathReason,
		);
	}

	private failAll(error: Error): void {
		const pending = [...this.pending.values()];
		this.pending.clear();
		for (const entry of pending) entry.reject(error);
	}
}