pi-graphnav 0.0.0-stage → 0.1.0

package/extensions/core/mcp.ts

@@ -0,0 +1,183 @@
+/**
+ * A long-lived `codegraph serve --mcp` process per project, spoken to over stdio (JSON-RPC 2.0,
+ * one message per line).
+ *
+ * Long-lived on purpose: the server's file watcher keeps the index in sync while pi runs,
+ * and no call pays for a process start. A call that times out kills the process; the next
+ * call starts a fresh one.
+ */
+import { type ChildProcessWithoutNullStreams, spawn as nodeSpawn } from "node:child_process";
+import { pathToFileURL } from "node:url";
+
+/** The part of a child process the client uses (so tests can pass a fake). */
+export interface ServerProcess {
+	stdin: { write(chunk: string): unknown };
+	stdout: { on(event: "data", fn: (chunk: Buffer | string) => void): unknown };
+	stderr: { on(event: "data", fn: (chunk: Buffer | string) => void): unknown };
+	on(event: "exit", fn: (code: number | null) => void): unknown;
+	on(event: "error", fn: (err: Error) => void): unknown;
+	kill(): unknown;
+}
+
+export type Spawner = (args: string[], cwd: string) => ServerProcess;
+
+export const defaultSpawner: Spawner = (args, cwd) =>
+	nodeSpawn("codegraph", args, {
+		cwd,
+		env: process.env,
+		stdio: ["pipe", "pipe", "pipe"],
+		windowsHide: true,
+		// npm installs a .cmd shim on Windows, which only a shell resolves.
+		shell: process.platform === "win32",
+	}) as ChildProcessWithoutNullStreams;
+
+export interface McpResult {
+	content?: { type: string; text?: string }[];
+	isError?: boolean;
+}
+
+export const MAX_STDERR = 2000;
+
+/** Remove ANSI colours and anything that looks like a secret from server diagnostics. */
+export function sanitize(text: string): string {
+	return text
+		.replace(/\u001b\[[0-9;]*m/g, "")
+		.replace(/\b([A-Z0-9_]*(?:TOKEN|SECRET|PASSWORD|API_KEY|APIKEY)[A-Z0-9_]*=)\S+/gi, "$1[redacted]")
+		.replace(/\bBearer\s+\S+/gi, "Bearer [redacted]")
+		.trim()
+		.slice(-MAX_STDERR);
+}
+
+interface Pending {
+	resolve: (v: any) => void;
+	reject: (e: Error) => void;
+}
+
+export class CodeGraphServer {
+	private proc?: ServerProcess;
+	private ready?: Promise<void>;
+	private pending = new Map<number, Pending>();
+	private nextId = 1;
+	private buffer = "";
+	private stderr = "";
+
+	constructor(
+		readonly root: string,
+		private readonly opts: { watch: boolean; timeoutMs: number; spawn?: Spawner },
+	) {}
+
+	get running(): boolean {
+		return this.proc !== undefined;
+	}
+
+	private start(): Promise<void> {
+		const args = ["serve", "--mcp", "--path", this.root, ...(this.opts.watch ? [] : ["--no-watch"])];
+		const proc = (this.opts.spawn ?? defaultSpawner)(args, this.root);
+		this.proc = proc;
+		this.buffer = "";
+		this.stderr = "";
+		proc.stdout.on("data", (chunk) => this.onData(String(chunk)));
+		proc.stderr.on("data", (chunk) => {
+			this.stderr = (this.stderr + String(chunk)).slice(-MAX_STDERR);
+		});
+		proc.on("error", (err) => this.fail(proc, err));
+		proc.on("exit", (code) => this.fail(proc, new Error(sanitize(this.stderr) || `codegraph serve exited with code ${code}`)));
+		const rootUri = pathToFileURL(this.root).href;
+		return this.request("initialize", {
+			protocolVersion: "2024-11-05",
+			rootUri,
+			workspaceFolders: [{ uri: rootUri, name: this.root.split(/[\\/]/).pop() || this.root }],
+			capabilities: {},
+			clientInfo: { name: "pi-graphnav", version: "1" },
+		}).then(() => {
+			this.send({ jsonrpc: "2.0", method: "notifications/initialized", params: {} });
+		});
+	}
+
+	/** The process died or could not start: fail what waits on it and forget it. */
+	private fail(proc: ServerProcess, err: Error) {
+		if (this.proc !== proc) return;
+		this.proc = undefined;
+		this.ready = undefined;
+		for (const p of this.pending.values()) p.reject(err);
+		this.pending.clear();
+	}
+
+	private onData(chunk: string) {
+		this.buffer += chunk;
+		let nl: number;
+		while ((nl = this.buffer.indexOf("\n")) !== -1) {
+			const line = this.buffer.slice(0, nl).trim();
+			this.buffer = this.buffer.slice(nl + 1);
+			if (!line) continue;
+			let msg: any;
+			try {
+				msg = JSON.parse(line);
+			} catch {
+				continue; // a log line on stdout
+			}
+			const p = typeof msg?.id === "number" ? this.pending.get(msg.id) : undefined;
+			if (!p) continue;
+			this.pending.delete(msg.id);
+			if (msg.error) p.reject(new Error(String(msg.error.message ?? JSON.stringify(msg.error))));
+			else p.resolve(msg.result);
+		}
+	}
+
+	private send(payload: unknown) {
+		this.proc?.stdin.write(`${JSON.stringify(payload)}\n`);
+	}
+
+	private request(method: string, params: unknown): Promise<any> {
+		const id = this.nextId++;
+		return new Promise((resolve, reject) => {
+			this.pending.set(id, { resolve, reject });
+			this.send({ jsonrpc: "2.0", id, method, params });
+		});
+	}
+
+	/** Call a CodeGraph tool. Times out after `timeoutMs` (the process is then restarted on the next call). */
+	async call(name: string, args: Record<string, unknown>, signal?: AbortSignal): Promise<McpResult> {
+		if (signal?.aborted) throw new Error("Aborted");
+		this.ready ??= this.start().catch((e: unknown) => {
+			this.close();
+			throw e;
+		});
+		let timer: ReturnType<typeof setTimeout> | undefined;
+		let onAbort: (() => void) | undefined;
+		let id: number | undefined;
+		const guard = new Promise<never>((_, reject) => {
+			timer = setTimeout(() => {
+				this.close();
+				reject(new Error(`CodeGraph did not answer within ${Math.round(this.opts.timeoutMs / 1000)}s. If this repeats, run "codegraph unlock" in ${this.root}.`));
+			}, this.opts.timeoutMs);
+			onAbort = () => reject(new Error("Aborted"));
+			signal?.addEventListener("abort", onAbort, { once: true });
+		});
+		try {
+			const answer = this.ready.then(() => {
+				if (signal?.aborted) throw new Error("Aborted");
+				id = this.nextId;
+				return this.request("tools/call", { name, arguments: args });
+			});
+			return await Promise.race([answer, guard]);
+		} finally {
+			// An abandoned request (timeout, abort) must not wait in `pending` for an answer nobody reads.
+			if (id !== undefined) this.pending.delete(id);
+			clearTimeout(timer);
+			if (onAbort) signal?.removeEventListener("abort", onAbort);
+		}
+	}
+
+	/** Stop the process. Idempotent. */
+	close() {
+		const proc = this.proc;
+		if (!proc) return;
+		this.fail(proc, new Error("CodeGraph server closed"));
+		try {
+			proc.kill();
+		} catch {
+			// already gone
+		}
+	}
+}

package/extensions/core/project.ts

@@ -0,0 +1,126 @@
+/**
+ * Which project CodeGraph serves, and what keeping its index current takes.
+ *
+ * Pure, except where a function takes the file system or a process runner as a parameter.
+ */
+import { existsSync } from "node:fs";
+import { homedir } from "node:os";
+import { dirname, join, parse, resolve } from "node:path";
+
+export interface Project {
+	root: string;
+	/** A `.codegraph/` directory exists at the root. */
+	hasIndexDir: boolean;
+	/** The root is a git work tree (so `init` there is expected, not a surprise). */
+	isRepo: boolean;
+}
+
+/**
+ * The project for `cwd`: the nearest directory with `.codegraph/`, else the nearest git
+ * work tree. Never the home directory or a file-system root: indexing those is never wanted.
+ */
+export function findProject(cwd: string, exists: (p: string) => boolean = existsSync, home: string = homedir()): Project | undefined {
+	let gitRoot: string | undefined;
+	let dir = resolve(cwd);
+	for (;;) {
+		if (dir === home || dir === parse(dir).root) break;
+		if (exists(join(dir, ".codegraph"))) return { root: dir, hasIndexDir: true, isRepo: exists(join(dir, ".git")) };
+		if (gitRoot === undefined && exists(join(dir, ".git"))) gitRoot = dir;
+		const up = dirname(dir);
+		if (up === dir) break;
+		dir = up;
+	}
+	return gitRoot === undefined ? undefined : { root: gitRoot, hasIndexDir: false, isRepo: true };
+}
+
+/** What `codegraph status -j` says, reduced to what matters here. */
+export interface IndexStatus {
+	initialized: boolean;
+	/** `complete`, `indexing`, … (`unknown` when the field is missing). */
+	state: string;
+	files: number;
+	pending: number;
+	reindexRecommended: boolean;
+	/** The index was built by another CodeGraph version than the one installed. */
+	versionMismatch: boolean;
+}
+
+const num = (v: unknown): number => (typeof v === "number" && Number.isFinite(v) ? v : 0);
+
+/** Parse `codegraph status -j`. Garbage (an old CodeGraph, a crash) reads as "not initialized". */
+export function parseStatus(json: string): IndexStatus {
+	let raw: any;
+	try {
+		raw = JSON.parse(json);
+	} catch {
+		raw = undefined;
+	}
+	const r = typeof raw === "object" && raw !== null ? raw : {};
+	const index = typeof r.index === "object" && r.index !== null ? r.index : {};
+	const pending = typeof r.pendingChanges === "object" && r.pendingChanges !== null ? r.pendingChanges : {};
+	return {
+		initialized: r.initialized === true,
+		state: typeof index.state === "string" ? index.state : "unknown",
+		files: num(r.fileCount),
+		pending: num(pending.added) + num(pending.modified) + num(pending.removed),
+		reindexRecommended: index.reindexRecommended === true,
+		versionMismatch:
+			typeof index.builtWithVersion === "string" && typeof r.version === "string" && index.builtWithVersion !== r.version,
+	};
+}
+
+export type IndexAction = "init" | "index" | "sync" | "none" | "skip";
+
+/**
+ * What brings the index up to date: `init` when there is none (only in a git work tree, and
+ * only when allowed), a full `index` when it is stale or broken, `sync` for pending changes.
+ */
+export function planIndex(status: IndexStatus, project: Project, opts: { autoInit: boolean }): IndexAction {
+	if (!status.initialized) return opts.autoInit && project.isRepo ? "init" : "skip";
+	if (status.reindexRecommended || status.versionMismatch || (status.state !== "complete" && status.state !== "indexing")) return "index";
+	if (status.pending > 0) return "sync";
+	return "none";
+}
+
+/** The CLI arguments for an action. */
+export function actionArgs(action: Exclude<IndexAction, "none" | "skip">): string[] {
+	switch (action) {
+		case "init":
+			return ["init", "-i", "-y"];
+		case "index":
+			return ["index", "-q"];
+		case "sync":
+			return ["sync", "-q"];
+	}
+}
+
+/** Runs `codegraph <args>` in `cwd`; resolves with the exit code and output, never rejects. */
+export type Runner = (args: string[], cwd: string) => Promise<{ code: number; stdout: string; stderr: string }>;
+
+export interface EnsureResult {
+	action: IndexAction;
+	status: IndexStatus;
+	/** Set when CodeGraph is missing or a step failed. */
+	error?: string;
+}
+
+/** Bring the index of `project` up to date: status, then init/index/sync as planned, then status again. */
+export async function ensureIndex(project: Project, run: Runner, opts: { autoInit: boolean }): Promise<EnsureResult> {
+	const first = await run(["status", "-j"], project.root);
+	if (first.code === 127) return { action: "skip", status: parseStatus(""), error: "the codegraph CLI is not installed (npm i -g @colbymchenry/codegraph)" };
+	const status = parseStatus(first.stdout);
+	const action = planIndex(status, project, opts);
+	if (action === "none" || action === "skip") return { action, status };
+	const res = await run(actionArgs(action), project.root);
+	const after = parseStatus((await run(["status", "-j"], project.root)).stdout);
+	if (res.code !== 0) return { action, status: after, error: `codegraph ${action} failed (exit ${res.code}): ${lastLine(res.stderr || res.stdout)}` };
+	return { action, status: after };
+}
+
+const lastLine = (s: string): string =>
+	s
+		.replace(/\u001b\[[0-9;]*m/g, "")
+		.trim()
+		.split("\n")
+		.pop()!
+		.slice(0, 300);

package/extensions/core/tools.ts

@@ -0,0 +1,142 @@
+/**
+ * The `codegraph_*` tools: names, parameters and the guidance the model sees.
+ *
+ * The names match the CodeGraph MCP server's tools, so calls pass straight through.
+ */
+import { Type } from "@earendil-works/pi-ai";
+
+const projectPath = Type.Optional(
+	Type.String({ description: "Absolute path of another indexed project. Defaults to the current project." }),
+);
+
+const kind = Type.Optional(
+	Type.Union(
+		["function", "method", "class", "interface", "type", "variable", "route", "component"].map((k) => Type.Literal(k)),
+		{ description: "Only symbols of this kind." },
+	),
+);
+
+export const TOOLS = [
+	{
+		name: "codegraph_explore",
+		label: "CodeGraph Explore",
+		description:
+			"Start here for any question about how code works, where something lives or what to change: returns the source of the relevant symbols, grouped by file, with the call paths between them. Accepts symbol names, file names or plain code terms.",
+		snippet: "Explore an area of the codebase: relevant symbols' source and call paths in one call",
+		parameters: Type.Object({
+			query: Type.String({ description: "Symbols, files or code terms, e.g. \"parseConfig loadSettings\"." }),
+			maxFiles: Type.Optional(Type.Number({ description: "Most files to return (default 12)." })),
+			projectPath,
+		}),
+	},
+	{
+		name: "codegraph_search",
+		label: "CodeGraph Search",
+		description: "Find symbols by name (or part of one): kind, file and line of each definition. Use instead of grep for a symbol.",
+		snippet: "Find where a symbol is defined (instead of rg/grep)",
+		parameters: Type.Object({
+			query: Type.String({ description: "Symbol name or part of one." }),
+			kind,
+			limit: Type.Optional(Type.Number({ description: "Most results (default 10)." })),
+			projectPath,
+		}),
+	},
+	{
+		name: "codegraph_node",
+		label: "CodeGraph Node",
+		description: "One symbol in detail: signature, location, callers and callees; with includeCode its source.",
+		snippet: "One symbol's details, callers and callees",
+		parameters: Type.Object({
+			symbol: Type.String({ description: "Symbol name." }),
+			includeCode: Type.Optional(Type.Boolean({ description: "Include the source (default false)." })),
+			projectPath,
+		}),
+	},
+	{
+		name: "codegraph_callers",
+		label: "CodeGraph Callers",
+		description: "Every function or method that calls a symbol: its usages, without grepping.",
+		snippet: "Who calls a symbol (its usages)",
+		parameters: Type.Object({
+			symbol: Type.String({ description: "Symbol name." }),
+			limit: Type.Optional(Type.Number({ description: "Most results (default 20)." })),
+			projectPath,
+		}),
+	},
+	{
+		name: "codegraph_callees",
+		label: "CodeGraph Callees",
+		description: "Every function or method a symbol calls.",
+		snippet: "What a symbol calls",
+		parameters: Type.Object({
+			symbol: Type.String({ description: "Symbol name." }),
+			limit: Type.Optional(Type.Number({ description: "Most results (default 20)." })),
+			projectPath,
+		}),
+	},
+	{
+		name: "codegraph_impact",
+		label: "CodeGraph Impact",
+		description: "What a change to a symbol affects, transitively. Call it before changing a signature or behaviour.",
+		snippet: "What breaks if a symbol changes",
+		parameters: Type.Object({
+			symbol: Type.String({ description: "Symbol name." }),
+			depth: Type.Optional(Type.Number({ description: "How many levels to follow (default 2)." })),
+			projectPath,
+		}),
+	},
+	{
+		name: "codegraph_files",
+		label: "CodeGraph Files",
+		description: "The project's indexed files as a tree, flat list or grouped by language, with symbol counts.",
+		snippet: "Project file structure from the index",
+		parameters: Type.Object({
+			path: Type.Optional(Type.String({ description: "Only files under this root-relative prefix, e.g. \"src/core\"." })),
+			pattern: Type.Optional(Type.String({ description: "Glob on file paths, e.g. \"*.test.ts\"." })),
+			format: Type.Optional(Type.Union([Type.Literal("tree"), Type.Literal("flat"), Type.Literal("grouped")])),
+			includeMetadata: Type.Optional(Type.Boolean()),
+			maxDepth: Type.Optional(Type.Number()),
+			projectPath,
+		}),
+	},
+	{
+		name: "codegraph_status",
+		label: "CodeGraph Status",
+		description: "Index statistics for the project: files, symbols, last update.",
+		snippet: "Index status",
+		parameters: Type.Object({ projectPath }),
+	},
+] as const;
+
+export type ToolName = (typeof TOOLS)[number]["name"];
+export const TOOL_NAMES: readonly string[] = TOOLS.map((t) => t.name);
+export const isCodeGraphTool = (name: string): boolean => TOOL_NAMES.includes(name);
+
+/** Guidelines added to the system prompt while the tools are active (static: they never break the prompt cache). */
+export const GUIDELINES = [
+	"Understand code through CodeGraph first: codegraph_explore for how something works or where to change it, codegraph_search to find a definition, codegraph_callers/codegraph_callees for relations, codegraph_impact before changing a signature, codegraph_files for structure.",
+	"Do not grep (rg, grep, git grep) for symbol names: such searches are redirected to CodeGraph. Text search is for string literals, comments, config and docs, or after CodeGraph came up empty.",
+	"Read a file once CodeGraph has told you where to look, with offset/limit around the reported lines.",
+];
+
+/** Largest result handed to the model; the rest is cut with a note. */
+export const MAX_RESULT_CHARS = 60_000;
+
+export function resultText(result: { content?: { type: string; text?: string }[] } | undefined): string {
+	const text = (result?.content ?? [])
+		.filter((p) => p?.type === "text" && typeof p.text === "string")
+		.map((p) => p.text)
+		.join("\n");
+	if (text.length <= MAX_RESULT_CHARS) return text || "(no result)";
+	return `${text.slice(0, MAX_RESULT_CHARS)}\n\n[cut at ${MAX_RESULT_CHARS} characters: narrow the query, lower maxFiles/limit or pass a path]`;
+}
+
+/** `codegraph_files` wants a root-relative POSIX prefix; accept an absolute path inside the project too. */
+export function normalizeFilesPath(path: unknown, root: string): string | undefined {
+	if (typeof path !== "string" || path.trim() === "") return undefined;
+	let p = path.trim().replace(/\\/g, "/");
+	const r = root.replace(/\\/g, "/").replace(/\/$/, "");
+	if (p === r) return undefined;
+	if (p.startsWith(`${r}/`)) p = p.slice(r.length + 1).replace(/^(?:\.?\/)+/, "");
+	return p.replace(/^\.\/(?:\.?\/)*/, "").replace(/\/$/, "") || undefined;
+}