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";

const NOT_INDEXED_MARKER = "no .codegraph/ index exists";
const EXPLORE_TIMEOUT_MS = 120_000;
const INIT_TIMEOUT_MS = 600_000;
const CACHE_MAX_ENTRIES = 32;

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

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

/** 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;
	}
}

export function buildExploreArgs(params: ExploreParams, root: string): string[] {
	const args = ["explore", "-p", root];
	if (params.maxFiles !== undefined) args.push("--max-files", String(params.maxFiles));
	args.push(params.query);
	return args;
}

export function isNotIndexedOutput(text: string): boolean {
	return text.includes(NOT_INDEXED_MARKER);
}

async function databaseMtime(root: string): Promise<number | undefined> {
	try {
		return (await stat(join(root, ".codegraph", "codegraph.db"))).mtimeMs;
	} catch {
		return undefined;
	}
}

function cacheKey(root: string, params: ExploreParams): string {
	return `${root}${params.query}${params.maxFiles ?? ""}`;
}

function resultText(result: ExecResult): string {
	return (result.stdout.trim() || result.stderr.trim());
}

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

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

/**
 * Native codegraph_explore tool backed by the machine-level `codegraph` CLI.
 *
 * Each call is a short-lived `codegraph explore` process — no MCP server, no
 * daemon requirement, nothing resident between calls. Results are memoized per
 * (root, query, maxFiles) keyed on the index database mtime: any reindex or
 * sync bumps the mtime and invalidates stale entries automatically.
 *
 * When the target has no index, the tool asks the user once (per concurrent
 * decision) whether to run `codegraph init`; on approval it indexes and then
 * answers the original call, on decline the model is bounced to built-in tools.
 */
export function createCodegraphExploreHandler(
	exec: CodegraphExec,
	cache: Map<string, { mtime: number; text: string }>,
	pendingDecisions: Map<string, Promise<string | undefined>>,
) {
	async function decideIndex(startDir: string, ctx: ExtensionContext, signal?: AbortSignal) {
		const indexed = await findIndexedRoot(startDir);
		if (indexed !== undefined) return indexed;
		if (!ctx.hasUI) return undefined;

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

		const init = await exec("codegraph", ["init", startDir], {
			cwd: startDir,
			timeout: INIT_TIMEOUT_MS,
			signal,
		});
		if (init.code !== 0) {
			throw new Error(`codegraph init 失败${failureText(init)}`);
		}
		return (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 remember(key: string, mtime: number | undefined, text: string) {
		if (mtime === undefined) return;
		if (!cache.has(key) && cache.size >= CACHE_MAX_ENTRIES) {
			const oldest = cache.keys().next().value;
			if (oldest !== undefined) cache.delete(oldest);
		}
		cache.set(key, { mtime, text });
	}

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

		let root: string;
		try {
			const decided = await decideIndexOnce(startDir, ctx, signal);
			if (decided === undefined) {
				return toResult(
					`CodeGraph 未在 ${startDir} 建立索引（用户未授权或未安装）。本会话对该目录请改用 find/grep/read/LSP；如需启用 CodeGraph，请用户安装 codegraph CLI 并运行 codegraph init。`,
				);
			}
			root = decided;
		} catch (error) {
			const detail = error instanceof Error ? error.message : String(error);
			return toResult(`${detail}。请改用 find/grep/read/LSP 完成本次探索。`);
		}

		const key = cacheKey(root, params);
		const mtime = await databaseMtime(root);
		const hit = cache.get(key);
		if (hit && hit.mtime === mtime) return toResult(hit.text);

		let explored: ExecResult;
		try {
			explored = await exec("codegraph", buildExploreArgs(params, root), {
				cwd: root,
				timeout: EXPLORE_TIMEOUT_MS,
				signal,
			});
		} catch (error) {
			const detail = error instanceof Error ? error.message : String(error);
			return toResult(
				`无法执行 codegraph CLI（${detail}）。若未安装，请用户运行 install.sh 安装；本次请改用 find/grep/read/LSP。`,
			);
		}
		const output = resultText(explored);
		if (explored.code !== 0) {
			return toResult(`codegraph explore 失败${failureText(explored)}。请改用 find/grep/read/LSP。`);
		}
		if (isNotIndexedOutput(output)) {
			// Index vanished between our check and the CLI run — surface the
			// upstream guidance verbatim; the next call re-enters the init path.
			return toResult(output);
		}
		const text = output || "codegraph explore 返回空输出。";
		remember(key, mtime, text);
		return toResult(text);
	};
}

export default function codegraphExtension(pi: ExtensionAPI): void {
	const cache = new Map<string, { mtime: number; text: string }>();
	const pendingDecisions = new Map<string, Promise<string | undefined>>();
	const handleExplore = createCodegraphExploreHandler(
		(command, args, options) => pi.exec(command, args, options),
		cache,
		pendingDecisions,
	);

	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 handleExplore(params as ExploreParams, ctx, signal);
		},
	});
}