my-pi/tests/codegraph.test.ts

import assert from "node:assert/strict";
import { EventEmitter } from "node:events";
import {
	execFileSync,
	spawn,
	type ChildProcess,
} from "node:child_process";
import {
	mkdir,
	mkdtemp,
	readFile,
	rm,
	writeFile,
} from "node:fs/promises";
import { tmpdir } from "node:os";
import { join } from "node:path";
import { PassThrough } from "node:stream";
import { fileURLToPath } from "node:url";
import test from "node:test";

import {
	CodegraphRpcClient,
	type CodegraphRpcClientOptions,
} from "../extensions/codegraph-client.ts";
import {
	createCodegraphExploreHandler,
	findIndexedRoot,
	type CodegraphExec,
	type CodegraphExploreClient,
	type ExploreParams,
} from "../extensions/codegraph.ts";

const FIXTURE = fileURLToPath(
	new URL("./fixtures/fake-codegraph-mcp.mjs", import.meta.url),
);

const sleep = (ms: number) => new Promise((r) => setTimeout(r, ms));

// ── fixtures ────────────────────────────────────────────────────────────────

async function makeIndexedProject(): Promise<{
	root: string;
	cleanup: () => Promise<void>;
}> {
	const root = await mkdtemp(join(tmpdir(), "my-pi-codegraph-"));
	await mkdir(join(root, ".codegraph"));
	await writeFile(join(root, ".codegraph", "codegraph.db"), "");
	return { root, cleanup: () => rm(root, { recursive: true, force: true }) };
}

/** Spawn the .mjs fixture under process.execPath with the given env. */
function fixtureSpawn(extraEnv: Record<string, string>, children?: ChildProcess[]) {
	return (_command: string, _args: string[], options: object) => {
		const child = spawn(process.execPath, [FIXTURE], {
			...(options as object),
			env: { ...process.env, ...extraEnv },
		});
		children?.push(child);
		return child;
	};
}

function fixtureClient(
	root: string,
	extraEnv: Record<string, string> = {},
	overrides: Partial<CodegraphRpcClientOptions> = {},
	children?: ChildProcess[],
): CodegraphRpcClient {
	return new CodegraphRpcClient({
		root,
		requestTimeoutMs: 5_000,
		spawn: fixtureSpawn(
			{ FAKE_CG_MODE: "ok", ...extraEnv },
			children,
		),
		...overrides,
	});
}

function mockCtx(
	cwd: string,
	ui?: { confirm?: () => Promise<boolean>; notify?: (m: string, t?: string) => void },
) {
	return {
		cwd,
		hasUI: ui !== undefined,
		ui: {
			confirm: ui?.confirm ?? (async () => false),
			notify: ui?.notify ?? (() => {}),
		},
	} as never;
}

function mockExec(
	handler: (command: string, args: string[]) => {
		code?: number;
		stdout?: string;
		stderr?: string;
	},
) {
	const calls: Array<{ command: string; args: string[] }> = [];
	const exec: CodegraphExec = async (command, args) => {
		calls.push({ command, args });
		const result = handler(command, args);
		return {
			code: result.code ?? 0,
			stdout: result.stdout ?? "",
			stderr: result.stderr ?? "",
			killed: false,
		};
	};
	return { exec, calls };
}

function makeHandler(
	exec: CodegraphExec,
	createClient: (root: string) => CodegraphExploreClient,
) {
	return createCodegraphExploreHandler({ exec, createClient });
}

// ── findIndexedRoot (unchanged behavior) ────────────────────────────────────

test("findIndexedRoot locates the nearest ancestor index", async () => {
	const { root, cleanup } = await makeIndexedProject();
	try {
		const nested = join(root, "src", "deep");
		await mkdir(nested, { recursive: true });
		assert.equal(await findIndexedRoot(nested), root);
		assert.equal(await findIndexedRoot(root), root);
	} finally {
		await cleanup();
	}
});

test("findIndexedRoot returns undefined when no index exists", async () => {
	const root = await mkdtemp(join(tmpdir(), "my-pi-codegraph-"));
	try {
		assert.equal(await findIndexedRoot(root), undefined);
	} finally {
		await rm(root, { recursive: true, force: true });
	}
});

// ── CodegraphRpcClient against the fixture ──────────────────────────────────

test("client handshake + explore returns upstream content verbatim", async () => {
	const { root, cleanup } = await makeIndexedProject();
	const client = fixtureClient(root);
	try {
		await client.start();
		const result = await client.callExplore({
			query: "alphaOne",
			maxFiles: 5,
			projectPath: "/some/dir",
		});
		const content = result.content as Array<{ type: string; text: string }>;
		const args = JSON.parse(content[0].text.slice("explore:".length));
		assert.deepEqual(args, {
			query: "alphaOne",
			maxFiles: 5,
			projectPath: "/some/dir",
		});
	} finally {
		await client.close();
		await cleanup();
	}
});

test("client forwards only properties the upstream schema declares", async () => {
	const { root, cleanup } = await makeIndexedProject();
	const client = fixtureClient(root, {
		FAKE_CG_SCHEMA: '["query","maxFiles"]',
	});
	try {
		await client.start();
		const result = await client.callExplore({
			query: "q",
			maxFiles: 3,
			projectPath: "/should/not/forward",
		});
		const content = result.content as Array<{ text: string }>;
		const args = JSON.parse(content[0].text.slice("explore:".length));
		assert.deepEqual(args, { query: "q", maxFiles: 3 });
	} finally {
		await client.close();
		await cleanup();
	}
});

test("client start fails when the server lacks codegraph_explore", async () => {
	const { root, cleanup } = await makeIndexedProject();
	const client = fixtureClient(root, { FAKE_CG_MODE: "no-tool" });
	try {
		await assert.rejects(client.start(), /未提供 codegraph_explore 工具/);
	} finally {
		await client.close();
		await cleanup();
	}
});

test("client abort sends notifications/cancelled and rejects", async () => {
	const dir = await mkdtemp(join(tmpdir(), "my-pi-codegraph-"));
	const logFile = join(dir, "rpc.log");
	const client = fixtureClient(dir, {
		FAKE_CG_MODE: "hang",
		FAKE_CG_LOG: logFile,
	});
	try {
		await client.start();
		const controller = new AbortController();
		const call = client.callExplore({ query: "q" }, controller.signal);
		setTimeout(() => controller.abort(), 50);
		await assert.rejects(call);
		await client.close();
		const lines = (await readFile(logFile, "utf8"))
			.trim()
			.split("\n")
			.map((l) => JSON.parse(l));
		const cancelled = lines.find(
			(m) => m.method === "notifications/cancelled",
		);
		assert.ok(cancelled, "server must see notifications/cancelled");
		assert.equal(cancelled.params.requestId, 3);
	} finally {
		await rm(dir, { recursive: true, force: true });
	}
});

test("client request timeout rejects a hung call", async () => {
	const { root, cleanup } = await makeIndexedProject();
	const client = fixtureClient(
		root,
		{ FAKE_CG_MODE: "hang" },
		{ requestTimeoutMs: 300 },
	);
	try {
		await client.start();
		await assert.rejects(client.callExplore({ query: "q" }), /超时/);
	} finally {
		await client.close();
		await cleanup();
	}
});

test("client exit rejects the pending call with the stderr tail", async () => {
	const dir = await mkdtemp(join(tmpdir(), "my-pi-codegraph-"));
	const client = fixtureClient(dir, {
		FAKE_CG_MODE: "crash-first",
		FAKE_CG_MARKER: join(dir, "crash.marker"),
	});
	try {
		await client.start();
		await assert.rejects(
			client.callExplore({ query: "q" }),
			/index corrupted/,
		);
		assert.equal(client.dead, true);
	} finally {
		await client.close();
		await rm(dir, { recursive: true, force: true });
	}
});

test("client close() terminates the child process", async () => {
	const { root, cleanup } = await makeIndexedProject();
	const children: ChildProcess[] = [];
	const client = fixtureClient(root, {}, {}, children);
	try {
		await client.start();
		await client.close();
		assert.equal(children.length, 1);
		assert.ok(
			children[0].exitCode !== null || children[0].signalCode !== null,
			"child must be gone after close()",
		);
	} finally {
		await client.close();
		await cleanup();
	}
});

test("a stdin EPIPE after child death rejects pending requests without throwing", async () => {
	const dir = await mkdtemp(join(tmpdir(), "my-pi-codegraph-"));
	const child = new EventEmitter() as ChildProcess;
	child.stdin = new PassThrough();
	child.stdout = new PassThrough();
	child.stderr = new PassThrough();
	child.exitCode = null;
	child.signalCode = null;
	const client = new CodegraphRpcClient({
		root: dir,
		spawn: () => child,
		requestTimeoutMs: 5_000,
	});
	try {
		const start = client.start();
		await sleep(20);
		// Without the stdin "error" listener this emit throws uncaught.
		child.stdin.emit(
			"error",
			Object.assign(new Error("write EPIPE"), { code: "EPIPE" }),
		);
		child.exitCode = 1;
		child.emit("exit", 1, null);
		child.emit("close", 1, null);
		await assert.rejects(start, /子进程退出/);
		assert.equal(client.dead, true);
	} finally {
		await client.close();
		await rm(dir, { recursive: true, force: true });
	}
});

// ── handler ─────────────────────────────────────────────────────────────────

test("indexed project explores through the serve child", async () => {
	const { root, cleanup } = await makeIndexedProject();
	const { exec, calls } = mockExec(() => ({ stdout: "unreachable" }));
	const handler = makeHandler(exec, (r) => fixtureClient(r));
	try {
		const result = await handler.handle({ query: "auth" }, mockCtx(root));
		const content = result.content as Array<{ text: string }>;
		assert.match(content[0].text, /^explore:\{/);
		assert.equal(calls.length, 0, "indexed project never runs codegraph init");
	} finally {
		await handler.close();
		await cleanup();
	}
});

test("declined init is remembered for the rest of the session", async () => {
	const root = await mkdtemp(join(tmpdir(), "my-pi-codegraph-"));
	let confirmCalls = 0;
	const { exec, calls } = mockExec(() => ({ stdout: "unreachable" }));
	const handler = makeHandler(exec, (r) => fixtureClient(r));
	const ctx = mockCtx(root, {
		confirm: async () => {
			confirmCalls += 1;
			return false;
		},
	});
	try {
		const first = await handler.handle({ query: "a" }, ctx);
		const second = await handler.handle({ query: "b" }, ctx);
		assert.match(first.content[0].text ?? "", /find\/grep\/read\/LSP/);
		assert.match(second.content[0].text ?? "", /find\/grep\/read\/LSP/);
		assert.equal(confirmCalls, 1, "decline is remembered — no second prompt");
		assert.equal(calls.length, 0, "declined init must not spawn the CLI");
	} finally {
		await handler.close();
		await rm(root, { recursive: true, force: true });
	}
});

test("no UI bounces to built-in tools without prompting", async () => {
	const root = await mkdtemp(join(tmpdir(), "my-pi-codegraph-"));
	const { exec, calls } = mockExec(() => ({ stdout: "unreachable" }));
	const handler = makeHandler(exec, (r) => fixtureClient(r));
	try {
		const result = await handler.handle({ query: "auth" }, mockCtx(root));
		assert.match(result.content[0].text ?? "", /find\/grep\/read\/LSP/);
		assert.equal(calls.length, 0);
	} finally {
		await handler.close();
		await rm(root, { recursive: true, force: true });
	}
});

test("uninitialized project inits on approval then explores", async () => {
	const root = await mkdtemp(join(tmpdir(), "my-pi-codegraph-"));
	const { exec, calls } = mockExec((_command, args) => {
		if (args[0] === "init") return { stdout: "indexed" };
		return { stdout: "unreachable" };
	});
	// Emulate codegraph init materializing the index.
	const execWithSideEffect: CodegraphExec = async (command, args, options) => {
		const result = await exec(command, args, options);
		if (args[0] === "init") {
			await mkdir(join(root, ".codegraph"));
			await writeFile(join(root, ".codegraph", "codegraph.db"), "");
		}
		return result;
	};
	const handler = makeHandler(execWithSideEffect, (r) => fixtureClient(r));
	try {
		const result = await handler.handle(
			{ query: "auth" },
			mockCtx(root, { confirm: async () => true }),
		);
		const content = result.content as Array<{ text: string }>;
		assert.match(content[0].text, /^explore:\{/);
		assert.deepEqual(
			calls.map(({ args }) => args[0]),
			["init"],
		);
		assert.deepEqual(calls[0].args, ["init", root]);
	} finally {
		await handler.close();
		await rm(root, { recursive: true, force: true });
	}
});

test("concurrent calls on an uninitialized project share one init decision", async () => {
	const root = await mkdtemp(join(tmpdir(), "my-pi-codegraph-"));
	let confirmCalls = 0;
	const { exec, calls } = mockExec(() => ({ stdout: "ok" }));
	const execWithSideEffect: CodegraphExec = async (command, args, options) => {
		const result = await exec(command, args, options);
		if (args[0] === "init") {
			await mkdir(join(root, ".codegraph"));
			await writeFile(join(root, ".codegraph", "codegraph.db"), "");
		}
		return result;
	};
	const handler = makeHandler(execWithSideEffect, (r) => fixtureClient(r));
	try {
		const ctx = mockCtx(root, {
			confirm: async () => {
				confirmCalls += 1;
				return true;
			},
		});
		const [a, b] = await Promise.all([
			handler.handle({ query: "one" }, ctx),
			handler.handle({ query: "two" }, ctx),
		]);
		assert.equal(a.content[0].type, "text");
		assert.equal(b.content[0].type, "text");
		assert.equal(confirmCalls, 1, "parallel callers must not double-prompt");
		assert.equal(
			calls.filter(({ args }) => args[0] === "init").length,
			1,
			"parallel callers must share a single init",
		);
	} finally {
		await handler.close();
		await rm(root, { recursive: true, force: true });
	}
});

test("init failure bounces with exit code instead of exploring", async () => {
	const root = await mkdtemp(join(tmpdir(), "my-pi-codegraph-"));
	const { exec } = mockExec((_command, args) =>
		args[0] === "init" ? { code: 1, stderr: "boom" } : { stdout: "unreachable" },
	);
	const handler = makeHandler(exec, (r) => fixtureClient(r));
	try {
		const result = await handler.handle(
			{ query: "auth" },
			mockCtx(root, { confirm: async () => true }),
		);
		assert.match(result.content[0].text ?? "", /codegraph init 失败（退出码 1）/);
	} finally {
		await handler.close();
		await rm(root, { recursive: true, force: true });
	}
});

test("a dead serve child is transparently replaced once", async () => {
	const dir = await mkdtemp(join(tmpdir(), "my-pi-codegraph-"));
	await mkdir(join(dir, ".codegraph"));
	await writeFile(join(dir, ".codegraph", "codegraph.db"), "");
	const env = {
		FAKE_CG_MODE: "crash-first",
		FAKE_CG_MARKER: join(dir, "crash.marker"),
	};
	let created = 0;
	const { exec } = mockExec(() => ({ stdout: "unreachable" }));
	const handler = makeHandler(exec, (r) => {
		created += 1;
		return fixtureClient(r, env);
	});
	try {
		const result = await handler.handle({ query: "auth" }, mockCtx(dir));
		const content = result.content as Array<{ text: string }>;
		assert.match(content[0].text, /^explore:\{/);
		assert.equal(created, 2, "the crashed client is replaced exactly once");
	} finally {
		await handler.close();
		await rm(dir, { recursive: true, force: true });
	}
});

test("a persistently crashing child bounces after one retry", async () => {
	const { root, cleanup } = await makeIndexedProject();
	let created = 0;
	const { exec } = mockExec(() => ({ stdout: "unreachable" }));
	const handler = makeHandler(exec, (r) => {
		created += 1;
		return fixtureClient(r, { FAKE_CG_MODE: "crash" });
	});
	try {
		const result = await handler.handle({ query: "auth" }, mockCtx(root));
		assert.match(result.content[0].text ?? "", /CodeGraph 服务调用失败/);
		assert.equal(created, 2, "one retry, then bounce");
	} finally {
		await handler.close();
		await cleanup();
	}
});

test("missing CLI surfaces an install hint instead of throwing", async () => {
	const { root, cleanup } = await makeIndexedProject();
	const { exec } = mockExec(() => ({ stdout: "unreachable" }));
	const handler = makeHandler(exec, () => ({
		dead: false,
		start: async () => {
			throw new Error("spawn codegraph ENOENT");
		},
		callExplore: async () => ({}),
		close: async () => {},
		schemaPropertyNames: () => [],
	}));
	try {
		const result = await handler.handle({ query: "auth" }, mockCtx(root));
		assert.match(result.content[0].text ?? "", /未找到 codegraph CLI/);
		assert.match(result.content[0].text ?? "", /install\.sh/);
	} finally {
		await handler.close();
		await cleanup();
	}
});

test("projectPath targets a different indexed root with its own client", async () => {
	const main = await makeIndexedProject();
	const other = await makeIndexedProject();
	const roots: string[] = [];
	const { exec } = mockExec(() => ({ stdout: "unreachable" }));
	const handler = makeHandler(exec, (r) => {
		roots.push(r);
		return fixtureClient(r);
	});
	try {
		await handler.handle({ query: "a" }, mockCtx(main.root));
		await handler.handle(
			{ query: "b", projectPath: other.root } satisfies ExploreParams,
			mockCtx(main.root),
		);
		assert.deepEqual(roots, [main.root, other.root]);
	} finally {
		await handler.close();
		await main.cleanup();
		await other.cleanup();
	}
});

test("schema drift notifies once per process", async () => {
	const { root, cleanup } = await makeIndexedProject();
	const notices: string[] = [];
	const { exec } = mockExec(() => ({ stdout: "unreachable" }));
	const handler = makeHandler(exec, (r) =>
		fixtureClient(r, { FAKE_CG_SCHEMA: '["query","maxFiles"]' }),
	);
	const ctx = mockCtx(root, {
		confirm: async () => true,
		notify: (message: string) => {
			notices.push(message);
		},
	});
	try {
		await handler.handle({ query: "a" }, ctx);
		await handler.handle({ query: "b" }, ctx);
		assert.equal(notices.length, 1);
		assert.match(notices[0], /参数已变化/);
	} finally {
		await handler.close();
		await cleanup();
	}
});

// ── real CLI integration ────────────────────────────────────────────────────

const codegraphAvailable = (() => {
	try {
		execFileSync("codegraph", ["--version"], { stdio: "ignore" });
		return true;
	} catch {
		return false;
	}
})();

test(
	"real codegraph serve --mcp keeps the index synced via its watcher",
	{ skip: !codegraphAvailable },
	async () => {
		const dir = await mkdtemp(join(tmpdir(), "my-pi-codegraph-real-"));
		const client = new CodegraphRpcClient({ root: dir });
		try {
			execFileSync("git", ["init", "-q"], { cwd: dir, stdio: "ignore" });
			await writeFile(
				join(dir, "a.ts"),
				"export function alphaOne() { return 1; }\n",
			);
			execFileSync("codegraph", ["init", dir], {
				cwd: dir,
				stdio: "ignore",
				timeout: 120_000,
			});

			const first = await client.callExplore({ query: "alphaOne" });
			const firstText = (first.content as Array<{ text?: string }>)
				.map((c) => c.text ?? "")
				.join("\n");
			assert.match(firstText, /alphaOne/);

			await writeFile(
				join(dir, "a.ts"),
				"export function alphaOne() { return 1; }\nexport function betaTwo() { return 2; }\n",
			);
			const deadline = Date.now() + 15_000;
			let found = false;
			while (Date.now() < deadline && !found) {
				const res = await client.callExplore({ query: "betaTwo" });
				found = (res.content as Array<{ text?: string }>).some((c) =>
					(c.text ?? "").includes("betaTwo"),
				);
				if (!found) await sleep(500);
			}
			assert.ok(found, "watcher should make betaTwo searchable within 15 s");
		} finally {
			await client.close();
			await rm(dir, { recursive: true, force: true });
		}
	},
);