my-pi/extensions/codegraph.ts

import type {
	AgentToolResult,
	ExecOptions,
	ExecResult,
	ExtensionAPI,
	ExtensionContext,
} from "@earendil-works/pi-coding-agent";
import { stat } from "node:fs/promises";
import { dirname, join, resolve } from "node:path";

import { CodegraphRpcClient } from "./codegraph-client.ts";

const INIT_TIMEOUT_MS = 600_000;
const IDLE_CLOSE_MS = 10 * 60_000;

export interface ExploreParams {
	query: string;
	maxFiles?: number;
	projectPath?: string;
}

export type CodegraphExec = (
	command: string,
	args: string[],
	options?: ExecOptions,
) => Promise<ExecResult>;

/** The slice of CodegraphRpcClient the handler depends on (test seam). */
export interface CodegraphExploreClient {
	readonly dead: boolean;
	start(): Promise<void>;
	callExplore(
		args: Record<string, unknown>,
		signal?: AbortSignal,
	): Promise<Record<string, unknown>>;
	close(): Promise<void> | void;
	schemaPropertyNames(): readonly string[];
}

export interface CodegraphExploreDeps {
	/** One-shot exec — used only for `codegraph init` after user approval. */
	exec: CodegraphExec;
	/** One client per index root; production wires `new CodegraphRpcClient({ root })`. */
	createClient: (root: string) => CodegraphExploreClient;
	now?: () => number;
}

/** Walk up from startDir looking for a .codegraph/codegraph.db index file. */
export async function findIndexedRoot(
	startDir: string,
): Promise<string | undefined> {
	let dir = resolve(startDir);
	for (;;) {
		try {
			if ((await stat(join(dir, ".codegraph", "codegraph.db"))).isFile())
				return dir;
		} catch {
			// Missing or unreadable index — keep walking up.
		}
		const parent = dirname(dir);
		if (parent === dir) return undefined;
		dir = parent;
	}
}

function toResult(text: string): AgentToolResult {
	return { content: [{ type: "text", text }] };
}

function failureText(result: ExecResult): string {
	const detail = result.stdout.trim() || result.stderr.trim();
	return detail
		? `（退出码 ${result.code}）${detail.slice(0, 400)}`
		: `（退出码 ${result.code}）`;
}

function errorMessage(error: unknown): string {
	return error instanceof Error ? error.message : String(error);
}

function notIndexedText(startDir: string): string {
	return `CodeGraph 未在 ${startDir} 建立索引（用户未授权或未安装）。本会话对该目录请改用 find/grep/read/LSP；如需启用 CodeGraph，请用户安装 codegraph CLI 并运行 codegraph init。`;
}

function serviceFailureText(error: unknown): string {
	const detail = errorMessage(error);
	if (/ENOENT/u.test(detail)) {
		return `未找到 codegraph CLI（${detail}）。请用户运行 install.sh 安装；本次请改用 find/grep/read/LSP。`;
	}
	return `CodeGraph 服务调用失败（${detail}）。请改用 find/grep/read/LSP 完成本次探索。`;
}

type IndexDecision = { root: string } | { bounce: string };

/**
 * Native codegraph_explore tool backed by `codegraph serve --mcp`.
 *
 * Each index root gets one long-lived child (lazily started on first use,
 * closed after 10 idle minutes or on session shutdown) spoken to through the
 * bundle's own stdio JSON-RPC client — no MCP adapter. Unlike the one-shot
 * `codegraph explore` CLI, the serve process runs a file watcher that keeps
 * the index synced, so there is no result cache and no manual sync step.
 *
 * When the target has no index, the tool asks the user once whether to run
 * `codegraph init`; a decline is remembered for the rest of the session. All
 * failures bounce the model to find/grep/read/LSP — the tool never throws.
 */
export function createCodegraphExploreHandler(deps: CodegraphExploreDeps) {
	const clients = new Map<string, CodegraphExploreClient>();
	const declined = new Set<string>();
	const pendingDecisions = new Map<string, Promise<IndexDecision>>();
	const idleTimers = new Map<string, ReturnType<typeof setTimeout>>();
	const registeredProps = ["maxFiles", "projectPath", "query"];
	let driftWarned = false;

	async function decideIndex(
		startDir: string,
		ctx: ExtensionContext,
		signal?: AbortSignal,
	): Promise<IndexDecision> {
		const indexed = await findIndexedRoot(startDir);
		if (indexed !== undefined) return { root: indexed };
		if (!ctx.hasUI) return { bounce: notIndexedText(startDir) };

		const confirmed = await ctx.ui.confirm(
			"启用 CodeGraph",
			`项目 ${startDir} 尚未建立 CodeGraph 索引。\n初始化会写入 .codegraph/ 并全量解析源码（大项目需要一些时间）。是否现在初始化？`,
		);
		if (!confirmed) {
			declined.add(startDir);
			return { bounce: notIndexedText(startDir) };
		}

		const init = await deps.exec("codegraph", ["init", startDir], {
			cwd: startDir,
			timeout: INIT_TIMEOUT_MS,
			signal,
		});
		if (init.code !== 0) {
			return {
				bounce: `codegraph init 失败${failureText(init)}。请改用 find/grep/read/LSP 完成本次探索。`,
			};
		}
		return { root: (await findIndexedRoot(startDir)) ?? startDir };
	}

	function decideIndexOnce(
		startDir: string,
		ctx: ExtensionContext,
		signal?: AbortSignal,
	) {
		let pending = pendingDecisions.get(startDir);
		if (!pending) {
			pending = decideIndex(startDir, ctx, signal).finally(() =>
				pendingDecisions.delete(startDir),
			);
			pendingDecisions.set(startDir, pending);
		}
		return pending;
	}

	function armIdleTimer(root: string, client: CodegraphExploreClient): void {
		const existing = idleTimers.get(root);
		if (existing) clearTimeout(existing);
		const timer = setTimeout(() => {
			idleTimers.delete(root);
			if (clients.get(root) === client) clients.delete(root);
			void Promise.resolve(client.close()).catch(() => {});
		}, IDLE_CLOSE_MS);
		timer.unref();
		idleTimers.set(root, timer);
	}

	function dropClient(root: string, client: CodegraphExploreClient): void {
		if (clients.get(root) === client) clients.delete(root);
		const timer = idleTimers.get(root);
		if (timer) {
			clearTimeout(timer);
			idleTimers.delete(root);
		}
		void Promise.resolve(client.close()).catch(() => {});
	}

	function checkSchemaDrift(
		client: CodegraphExploreClient,
		ctx: ExtensionContext,
	): void {
		if (driftWarned || !ctx.hasUI) return;
		const upstream = [...client.schemaPropertyNames()].sort();
		if (JSON.stringify(upstream) === JSON.stringify(registeredProps)) return;
		driftWarned = true;
		ctx.ui.notify(
			`CodeGraph 上游 codegraph_explore 参数已变化：上游 ${upstream.join(", ") || "(无)"}，扩展注册 ${registeredProps.join(", ")}。`,
			"warning",
		);
	}

	async function handleExplore(
		params: ExploreParams,
		ctx: ExtensionContext,
		signal?: AbortSignal,
	): Promise<AgentToolResult> {
		const startDir = resolve(params.projectPath ?? ctx.cwd);
		if (declined.has(startDir)) return toResult(notIndexedText(startDir));

		let root: string;
		try {
			const decision = await decideIndexOnce(startDir, ctx, signal);
			if ("bounce" in decision) return toResult(decision.bounce);
			root = decision.root;
		} catch (error) {
			return toResult(serviceFailureText(error));
		}

		let client: CodegraphExploreClient | undefined;
		const attempt = async (): Promise<AgentToolResult> => {
			const existing = clients.get(root);
			const active =
				existing && !existing.dead ? existing : deps.createClient(root);
			client = active;
			clients.set(root, active);
			armIdleTimer(root, active);
			try {
				await active.start();
				checkSchemaDrift(active, ctx);
				const result = await active.callExplore(
					params as unknown as Record<string, unknown>,
					signal,
				);
				const content = result.content;
				return {
					content: Array.isArray(content)
						? (content as AgentToolResult["content"])
						: [{ type: "text", text: "codegraph_explore 返回空输出。" }],
				};
			} catch (error) {
				// A user abort says nothing about the child's health — keep it warm.
				if (!signal?.aborted) dropClient(root, active);
				throw error;
			}
		};

		try {
			return await attempt();
		} catch (first) {
			// Retry once only when the child died — a live client failing a call
			// (RPC error, isError) is not fixed by respawning.
			if (!client?.dead) return toResult(serviceFailureText(first));
		}

		try {
			return await attempt();
		} catch (second) {
			return toResult(serviceFailureText(second));
		}
	}

	async function close(): Promise<void> {
		for (const timer of idleTimers.values()) clearTimeout(timer);
		idleTimers.clear();
		const all = [...clients.values()];
		clients.clear();
		await Promise.all(
			all.map((client) =>
				Promise.resolve(client.close()).catch(() => {}),
			),
		);
	}

	return { handle: handleExplore, close };
}

export default function codegraphExtension(pi: ExtensionAPI): void {
	const handler = createCodegraphExploreHandler({
		exec: (command, args, options) => pi.exec(command, args, options),
		createClient: (root) => new CodegraphRpcClient({ root }),
	});
	pi.on("session_shutdown", () => handler.close());

	pi.registerTool({
		name: "codegraph_explore",
		label: "codegraph_explore",
		description:
			"PRIMARY TOOL — call FIRST for almost any question OR before an edit: how does X work, architecture, a bug, where/what is X, surveying an area, or the symbols you are about to change. Returns the verbatim source of the relevant symbols grouped by file in ONE capped call (Read-equivalent — treat the shown source as already Read; do NOT re-open those files), plus the call path among them. Query can be a natural-language question OR a bag of symbol/file names. Usually the ONLY call you need — more accurate context, in far fewer tokens and round-trips than a search/Read/Grep loop.",
		parameters: {
			type: "object",
			properties: {
				query: {
					type: "string",
					description:
						'Symbol names, file names, or short code terms to explore (e.g., "AuthService loginUser session-manager", "GraphTraverser BFS impact traversal.ts"). For a flow question, name the symbols spanning the flow (e.g. "mutateElement renderScene"). A natural-language question works too.',
				},
				maxFiles: {
					type: "number",
					description:
						"Maximum number of files to include source code from (default: 12)",
				},
				projectPath: {
					type: "string",
					description:
						"Absolute path to the project to query (or any directory inside it) — codegraph uses the nearest .codegraph/ index at or above that path. Omit to use this session's default project. Pass it to query a second codebase, or when the session root has no index of its own (e.g. a monorepo where only sub-projects are indexed).",
				},
			},
			required: ["query"],
			additionalProperties: false,
		},
		async execute(_id, params, signal, _onUpdate, ctx) {
			return handler.handle(params as ExploreParams, ctx, signal);
		},
	});
}