import { execFile, execFileSync } from "node:child_process"; import { appendFileSync, existsSync, lstatSync, mkdirSync, readFileSync, realpathSync, statSync } from "node:fs"; import { homedir, tmpdir } from "node:os"; import { basename, dirname, join, resolve } from "node:path"; import { promisify } from "node:util"; import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent"; const CODEGRAPH_OPERATION = { INIT: "init", QUERY: "query", EXPLORE: "explore", } as const; const CODEGRAPH_STATUS = { UNAVAILABLE: "unavailable", FAILED: "failed", } as const; type CodeGraphOperation = (typeof CODEGRAPH_OPERATION)[keyof typeof CODEGRAPH_OPERATION]; type CodeGraphStatus = (typeof CODEGRAPH_STATUS)[keyof typeof CODEGRAPH_STATUS]; export interface CodeGraphToolParameters { operation: CodeGraphOperation; query?: string; limit?: number; } export interface CodeGraphCommandResult { stdout: string; stderr: string; } export interface CodeGraphRunOptions { cwd: string; signal?: AbortSignal; maxBuffer: number; } interface CodeGraphLocalExclude { status: "added" | "present" | "failed"; path?: string; reason?: string; } interface CodeGraphFallbackDetails { status: CodeGraphStatus; operation: CodeGraphOperation; cwd: string; fallback: string; localExclude?: CodeGraphLocalExclude; } export type CodeGraphRunner = ( args: readonly string[], options: CodeGraphRunOptions, ) => Promise; const CODEGRAPH_TOOL_PARAMETERS = { type: "object", additionalProperties: false, required: ["operation"], properties: { operation: { type: "string", enum: Object.values(CODEGRAPH_OPERATION) }, query: { type: "string", minLength: 1, maxLength: 2_000 }, limit: { type: "integer", minimum: 1, maximum: 20 }, }, } as const; const DEFAULT_LIMIT = 10; const MAX_LIMIT = 20; const MAX_OUTPUT_CHARS = 100_000; const PROCESS_MAX_BUFFER = MAX_OUTPUT_CHARS * 2; const FALLBACK_INSTRUCTIONS = "Use read, grep, and find for this exploration."; // query and explore read an existing index; on a slow filesystem (for example // WSL /mnt mounts) they could otherwise block the turn indefinitely (#1901). // init is the legitimately slow first index, so only the user's signal bounds it. const CODEGRAPH_QUERY_TIMEOUT_MS = 180_000; const execFileAsync = promisify(execFile); function resolveWorkspaceCwd(cwd: string): string { const resolved = realpathSync(cwd); if (!lstatSync(resolved).isDirectory()) { throw new Error("CodeGraph can run only in the current workspace directory."); } if (resolved === realpathSync(homedir()) || resolved === realpathSync(tmpdir())) { throw new Error("CodeGraph requires a real Git project root equal to the current workspace, not HOME or a temporary directory."); } try { const root = realpathSync(execFileSync("git", ["rev-parse", "--show-toplevel"], { cwd: resolved, encoding: "utf8", stdio: ["ignore", "pipe", "ignore"], windowsHide: true, }).trim()); if (root !== resolved) { throw new Error("CodeGraph requires a real Git project root equal to the current workspace."); } return resolved; } catch (error) { if (error instanceof Error && /real Git project root/.test(error.message)) throw error; throw new Error("CodeGraph requires a real Git project root equal to the current workspace."); } } function assertSafeIndexDirectory(cwd: string): void { try { const index = lstatSync(join(cwd, ".codegraph")); if (index.isSymbolicLink() || !index.isDirectory()) { throw new Error("CodeGraph .codegraph must be a real directory when it already exists."); } } catch (error: unknown) { if (typeof error === "object" && error !== null && "code" in error && error.code === "ENOENT") return; throw error; } } function indexExists(cwd: string): boolean { try { return lstatSync(join(cwd, ".codegraph")).isDirectory(); } catch { return false; } } function gitOutput(cwd: string, args: readonly string[]): string { return execFileSync("git", [...args], { cwd, encoding: "utf8", stdio: ["ignore", "pipe", "ignore"], windowsHide: true, }).trim(); } // CodeGraph writes .codegraph/.gitignore with `*` and `!.gitignore`, so that file // stays untracked and enters every diff. Keep the whole index out of git through // the worktree's local exclude file; tracked files and .gitignore are never edited. const LOCAL_EXCLUDE_PATTERN = ".codegraph/"; const EQUIVALENT_EXCLUDE_PATTERNS = new Set([".codegraph", ".codegraph/", "/.codegraph", "/.codegraph/"]); function excludeIndexLocally(cwd: string): CodeGraphLocalExclude | undefined { let path: string | undefined; try { if (gitOutput(cwd, ["rev-parse", "--is-inside-work-tree"]) !== "true") return undefined; // --git-path resolves the exclude file a linked worktree actually reads. path = resolve(cwd, gitOutput(cwd, ["rev-parse", "--git-path", "info/exclude"])); const current = existsSync(path) ? readFileSync(path, "utf8") : ""; mkdirSync(dirname(path), { recursive: true }); path = join(realpathSync(dirname(path)), basename(path)); if (current.split(/\r?\n/).some((line) => EQUIVALENT_EXCLUDE_PATTERNS.has(line.trim()))) { return { status: "present", path }; } const separator = current.length === 0 || current.endsWith("\n") ? "" : "\n"; appendFileSync(path, `${separator}${LOCAL_EXCLUDE_PATTERN}\n`); return { status: "added", path }; } catch (error: unknown) { const reason = error instanceof Error ? error.message.split("\n")[0] : String(error); return { status: "failed", ...(path ? { path } : {}), reason }; } } function withLocalExcludeNote(text: string, localExclude: CodeGraphLocalExclude | undefined): string { if (localExclude?.status !== "failed") return text; return `${text}\n\n.codegraph/ is not excluded from git locally (${localExclude.reason}). Keep .codegraph/ out of the change.`; } function resolveLimit(limit: number | undefined): number { if (limit === undefined) return DEFAULT_LIMIT; if (!Number.isInteger(limit) || limit < 1 || limit > MAX_LIMIT) { throw new Error(`CodeGraph limit must be an integer between 1 and ${MAX_LIMIT}.`); } return limit; } function requireQuery(query: string | undefined): string { if (typeof query !== "string" || query.trim().length === 0) { throw new Error("CodeGraph query is required for query and explore operations."); } if (query.length > 2_000) { throw new Error("CodeGraph query must not exceed 2000 characters."); } return query; } function commandArguments(parameters: CodeGraphToolParameters, cwd: string): string[] { switch (parameters.operation) { case CODEGRAPH_OPERATION.INIT: return [CODEGRAPH_OPERATION.INIT, cwd]; case CODEGRAPH_OPERATION.QUERY: { const query = requireQuery(parameters.query); return [ CODEGRAPH_OPERATION.QUERY, "--path", cwd, "--limit", String(resolveLimit(parameters.limit)), "--", query, ]; } case CODEGRAPH_OPERATION.EXPLORE: { const query = requireQuery(parameters.query); return [ CODEGRAPH_OPERATION.EXPLORE, "--path", cwd, "--max-files", String(resolveLimit(parameters.limit)), "--", query, ]; } } } function truncateOutput(output: string): string { return output.length <= MAX_OUTPUT_CHARS ? output : `${output.slice(0, MAX_OUTPUT_CHARS)}\n\n[CodeGraph output truncated]`; } function codeGraphFailureDetails( error: unknown, operation: CodeGraphOperation, cwd: string, ): CodeGraphFallbackDetails { const status = typeof error === "object" && error !== null && "code" in error && error.code === "ENOENT" ? CODEGRAPH_STATUS.UNAVAILABLE : CODEGRAPH_STATUS.FAILED; return { status, operation, cwd, fallback: FALLBACK_INSTRUCTIONS }; } function codeGraphFailureMessage(status: CodeGraphStatus): string { return status === CODEGRAPH_STATUS.UNAVAILABLE ? `CodeGraph is unavailable because the codegraph binary was not found. ${FALLBACK_INSTRUCTIONS}` : `CodeGraph failed to run. ${FALLBACK_INSTRUCTIONS}`; } function isEnoent(error: unknown): boolean { return ( typeof error === "object" && error !== null && "code" in error && (error as { code?: unknown }).code === "ENOENT" ); } /** Find `codegraph.cmd` shims in Windows PATH order. */ function* codeGraphCmdPathsOnPath(): Iterable { const pathEnv = process.env.Path ?? process.env.PATH ?? ""; for (const entry of pathEnv.split(";")) { const dir = entry.length >= 2 && entry.startsWith('"') && entry.endsWith('"') ? entry.slice(1, -1) : entry; if (!dir) continue; const candidate = join(dir, "codegraph.cmd"); if (existsSync(candidate)) yield candidate; } } /** * Resolve the actual JS entry of the `@colbymchenry/codegraph` package from the * npm global bin dir (sibling of the `codegraph.cmd` shim) and run it with the * current Node executable. No shell involvement, so args are never interpreted. */ export function codeGraphNodeScript(cmdPath: string): string | undefined { const binDir = dirname(cmdPath); const pkgRoot = join(binDir, "node_modules", "@colbymchenry", "codegraph"); const pkgJsonPath = join(pkgRoot, "package.json"); if (!existsSync(pkgJsonPath)) return undefined; try { const pkg = JSON.parse(readFileSync(pkgJsonPath, "utf8")) as { bin?: unknown }; const bin = pkg.bin; const target = typeof bin === "string" ? bin : bin && typeof bin === "object" && "codegraph" in bin && typeof (bin as Record).codegraph === "string" ? (bin as Record).codegraph : undefined; if (typeof target !== "string" || !target) return undefined; const script = join(pkgRoot, target); return statSync(script).isFile() ? script : undefined; } catch { return undefined; } } function* codeGraphNodeScriptsOnPath(): Generator { for (const cmdPath of codeGraphCmdPathsOnPath()) { const script = codeGraphNodeScript(cmdPath); if (script) yield script; } } /** Resolve the first valid CodeGraph npm entry point in Windows PATH order. */ export function findCodeGraphNodeScriptOnPath(): string | undefined { return codeGraphNodeScriptsOnPath().next().value; } const runCodeGraphCommand: CodeGraphRunner = async (args, options) => { const runOptions = { cwd: options.cwd, signal: options.signal, maxBuffer: options.maxBuffer, windowsHide: true, }; let unavailableError: unknown; try { const result = await execFileAsync("codegraph", [...args], runOptions); return { stdout: result.stdout, stderr: result.stderr }; } catch (error) { // Windows npm installs only shims (codegraph.cmd/.ps1 and a shell script) // with no codegraph.exe, so plain execFile (CreateProcess, no shell) // always fails with ENOENT even when the shim is on PATH. if (process.platform !== "win32" || !isEnoent(error)) throw error; unavailableError = error; } // Resolve the real package script from each npm global bin directory in PATH // order and run it through the current Node executable. This remains // shell-free, so argument boundaries are never interpreted by a shell. for (const script of codeGraphNodeScriptsOnPath()) { try { const result = await execFileAsync(process.execPath, [script, ...args], runOptions); return { stdout: result.stdout, stderr: result.stderr }; } catch (innerError) { if (!isEnoent(innerError)) throw innerError; } } throw unavailableError; }; export function createCodeGraphTool( runner: CodeGraphRunner = runCodeGraphCommand, { queryTimeoutMs = CODEGRAPH_QUERY_TIMEOUT_MS }: { queryTimeoutMs?: number } = {}, ) { return { name: "codegraph", renderShell: "self" as const, label: "CodeGraph", description: "Initialize, search, or explore the CodeGraph index for the current Pi workspace only. This tool never accepts a project path or shell command.", promptSnippet: "Initialize and query CodeGraph for the current workspace without shell access", promptGuidelines: [ "Use operation init before querying when the current workspace has no .codegraph index.", "Use query for symbol search and explore for source plus call paths. Do not use this tool to run arbitrary commands or target another directory.", ], parameters: CODEGRAPH_TOOL_PARAMETERS, executionMode: "sequential" as const, async execute( _toolCallId: string, parameters: CodeGraphToolParameters, signal: AbortSignal | undefined, _onUpdate: undefined, ctx: ExtensionContext, ): Promise<{ content: { type: "text"; text: string }[]; details: { operation: CodeGraphOperation; cwd: string; args: string[]; localExclude?: CodeGraphLocalExclude } | CodeGraphFallbackDetails }> { const cwd = resolveWorkspaceCwd(ctx.cwd); assertSafeIndexDirectory(cwd); const args = commandArguments(parameters, cwd); const timeout = parameters.operation === "init" ? undefined : AbortSignal.timeout(queryTimeoutMs); const runSignal = timeout && signal ? AbortSignal.any([signal, timeout]) : (timeout ?? signal); try { const result = await runner(args, { cwd, signal: runSignal, maxBuffer: PROCESS_MAX_BUFFER }); const output = truncateOutput([result.stdout, result.stderr].filter(Boolean).join("\n")); const localExclude = parameters.operation === CODEGRAPH_OPERATION.INIT || indexExists(cwd) ? excludeIndexLocally(cwd) : undefined; return { content: [{ type: "text" as const, text: withLocalExcludeNote(output || "CodeGraph completed without output.", localExclude) }], details: { operation: parameters.operation, cwd, args, ...(localExclude ? { localExclude } : {}) }, }; } catch (error: unknown) { const details = codeGraphFailureDetails(error, parameters.operation, cwd); // A failed init can still leave a partial index behind. const localExclude = indexExists(cwd) ? excludeIndexLocally(cwd) : undefined; if (localExclude) details.localExclude = localExclude; const timedOut = timeout?.aborted === true && signal?.aborted !== true; const text = timedOut ? `CodeGraph timed out after ${queryTimeoutMs} ms (a slow filesystem such as a WSL /mnt mount, or a very large workspace). ${FALLBACK_INSTRUCTIONS}` : codeGraphFailureMessage(details.status); return { content: [{ type: "text" as const, text: withLocalExcludeNote(text, localExclude) }], details, }; } }, }; } export function registerCodeGraphTool(pi: ExtensionAPI): void { pi.registerTool(createCodeGraphTool()); } export default function codeGraphTools(pi: ExtensionAPI): void { registerCodeGraphTool(pi); }