From 79acf694ea4b442bf63733eebc1ad943fc6f43d0 Mon Sep 17 00:00:00 2001 From: Evan Reichard Date: Wed, 15 Apr 2026 15:44:33 -0400 Subject: [PATCH] feat: call graph --- extension/codexis.ts | 176 ------------------------ extension/index.ts | 303 +++++++++++++++++++++++++++++++++++++++++ extension/package.json | 19 +++ 3 files changed, 322 insertions(+), 176 deletions(-) delete mode 100644 extension/codexis.ts create mode 100644 extension/index.ts create mode 100644 extension/package.json diff --git a/extension/codexis.ts b/extension/codexis.ts deleted file mode 100644 index f1e587f..0000000 --- a/extension/codexis.ts +++ /dev/null @@ -1,176 +0,0 @@ -/** - * Codexis - Code index query tool for pi - * - * Provides a single tool that queries the .codexis/index.db SQLite database - * containing symbols, files, and line numbers for the codebase. - * - * - Only registers if `codexis` binary is in PATH - * - Auto-indexes on first tool call if DB is missing - * - Re-indexes (incremental) on each call to keep index fresh - */ - -import { Type } from "@sinclair/typebox"; -import type { ExtensionAPI } from "@mariozechner/pi-coding-agent"; -import { existsSync } from "node:fs"; -import { join } from "node:path"; -import Database from "better-sqlite3"; - -// @ts-expect-error - generated file -import schema from "./schema.sql?raw"; - -const SCHEMA = schema; - -const DESCRIPTION = `Query the code index database (.codexis/index.db). Run read-only SQL to find symbols, files, and line numbers across the codebase. - -${SCHEMA} - -Example queries: - -- Find where a function is defined - SELECT f.path, s.line FROM symbols s JOIN files f ON s.file_id=f.id WHERE s.name='HandleRequest' - - -- Public API of a package - SELECT s.name, s.kind, s.line, f.path FROM symbols s JOIN files f ON s.file_id=f.id WHERE f.package='server' AND s.exported=1 - - -- All types in a directory - SELECT s.name, s.line, f.path FROM symbols s JOIN files f ON s.file_id=f.id WHERE f.path LIKE 'backend/api/%' AND s.kind='type' - - -- Methods on a class/type (via parent_id) - SELECT c.name as parent, s.name, s.kind, s.line FROM symbols s JOIN symbols c ON s.parent_id=c.id WHERE c.name='AuthService' - - -- Overview: symbols per area - SELECT CASE WHEN f.path LIKE 'backend/%' THEN 'backend' WHEN f.path LIKE 'frontend/%' THEN 'frontend' ELSE 'other' END as area, COUNT(*) FROM symbols s JOIN files f ON s.file_id=f.id GROUP BY area - - -- Full-text search for content across all files - SELECT f.path, snippet(file_contents, 1, '>>>', '<<<', '...', 20) as match FROM file_contents fc JOIN files f ON f.id=fc.file_id WHERE file_contents MATCH 'handleRequest' ORDER BY rank LIMIT 10 - - -- FTS search scoped to a directory - SELECT f.path, snippet(file_contents, 1, '>>>', '<<<', '...', 20) as match FROM file_contents fc JOIN files f ON f.id=fc.file_id WHERE file_contents MATCH 'database migration' AND f.path LIKE 'backend/%' ORDER BY rank LIMIT 10`; - -function findGitRoot(cwd: string): string | null { - try { - const { execSync } = require("node:child_process"); - return execSync("git rev-parse --show-toplevel", { - cwd, - encoding: "utf-8", - stdio: ["pipe", "pipe", "pipe"], - }).trim(); - } catch { - return null; - } -} - -function codexisAvailable(): boolean { - try { - const { execSync } = require("node:child_process"); - execSync("which codexis", { stdio: ["pipe", "pipe", "pipe"] }); - return true; - } catch { - return false; - } -} - -export default function (pi: ExtensionAPI) { - if (!codexisAvailable()) { - return; - } - - pi.registerTool({ - name: "Codexis", - label: "Codexis", - description: DESCRIPTION, - parameters: Type.Object({ - sql: Type.String({ - description: "SQL query to run against the code index database", - }), - }), - - async execute(_toolCallId, params, signal, _onUpdate, ctx) { - const gitRoot = findGitRoot(ctx.cwd); - if (!gitRoot) { - throw new Error("Not in a git repository"); - } - - const dbPath = join(gitRoot, ".codexis", "index.db"); - - // Run incremental index to keep DB fresh (fast if nothing changed) - const indexResult = await pi.exec("codexis", [gitRoot], { - signal, - timeout: 120_000, - }); - - if (indexResult.code !== 0) { - throw new Error( - `codexis indexing failed (exit ${indexResult.code}): ${indexResult.stderr}`, - ); - } - - if (!existsSync(dbPath)) { - throw new Error( - "codexis ran but no index.db was produced. Check codexis output.", - ); - } - - const db = new Database(dbPath, { readonly: true }); - try { - // Block writes - const normalized = params.sql.trim().toUpperCase(); - if ( - !normalized.startsWith("SELECT") && - !normalized.startsWith("WITH") && - !normalized.startsWith("EXPLAIN") && - !normalized.startsWith("PRAGMA") - ) { - throw new Error( - "Only SELECT, WITH, EXPLAIN, and PRAGMA queries are allowed", - ); - } - - const stmt = db.prepare(params.sql); - const rows = stmt.all(); - - if (rows.length === 0) { - return { - content: [{ type: "text", text: "No results." }], - details: { rowCount: 0 }, - }; - } - - // Format as aligned text table - const columns = Object.keys(rows[0] as Record); - const data = rows.map((row) => { - const r = row as Record; - return columns.map((col) => String(r[col] ?? "NULL")); - }); - - const widths = columns.map((col, i) => - Math.max(col.length, ...data.map((row) => row[i].length)), - ); - - const header = columns - .map((col, i) => col.padEnd(widths[i])) - .join(" "); - const separator = widths.map((w) => "-".repeat(w)).join(" "); - const body = data - .map((row) => row.map((val, i) => val.padEnd(widths[i])).join(" ")) - .join("\n"); - - const result = `${header}\n${separator}\n${body}`; - - // Truncate if huge - const maxLen = 48000; - const truncated = - result.length > maxLen - ? result.slice(0, maxLen) + - `\n\n[Truncated: ${rows.length} rows total, showing partial results. Narrow your query.]` - : result; - - return { - content: [{ type: "text", text: truncated }], - details: { rowCount: rows.length }, - }; - } finally { - db.close(); - } - }, - }); -} diff --git a/extension/index.ts b/extension/index.ts new file mode 100644 index 0000000..f257ec1 --- /dev/null +++ b/extension/index.ts @@ -0,0 +1,303 @@ +/** + * Codexis - Code index query tool for pi + * + * Provides a single tool that queries the .codexis/index.db SQLite database + * containing symbols, files, and line numbers for the codebase. + * + * - Only registers if `codexis` binary is in PATH + * - Auto-indexes on first tool call if DB is missing + * - Re-indexes (incremental) on each call to keep index fresh + */ + +import { readFileSync } from "node:fs"; +import { dirname, join } from "node:path"; +import { fileURLToPath } from "node:url"; +import { Type } from "@sinclair/typebox"; +import type { ExtensionAPI } from "@mariozechner/pi-coding-agent"; +import { existsSync } from "node:fs"; +import Database from "better-sqlite3"; + +const baseDir = dirname(fileURLToPath(import.meta.url)); + +const SCHEMA = readFileSync(join(baseDir, "schema.sql"), "utf-8"); + +const DESCRIPTION = `Query the code index database (.codexis/index.db). Run read-only SQL to find symbols, files, and line numbers across the codebase. + +${SCHEMA} + +Example queries: + -- Find where a function is defined + SELECT f.path, s.line FROM symbols s JOIN files f ON s.file_id=f.id WHERE s.name='HandleRequest' + + -- Public API of a package + SELECT s.name, s.kind, s.line, f.path FROM symbols s JOIN files f ON s.file_id=f.id WHERE f.package='server' AND s.exported=1 + + -- All types in a directory + SELECT s.name, s.line, f.path FROM symbols s JOIN files f ON s.file_id=f.id WHERE f.path LIKE 'backend/api/%' AND s.kind='type' + + -- Methods on a class/type (via parent_id) + SELECT c.name as parent, s.name, s.kind, s.line FROM symbols s JOIN symbols c ON s.parent_id=c.id WHERE c.name='AuthService' + + -- Overview: symbols per area + SELECT CASE WHEN f.path LIKE 'backend/%' THEN 'backend' WHEN f.path LIKE 'frontend/%' THEN 'frontend' ELSE 'other' END as area, COUNT(*) FROM symbols s JOIN files f ON s.file_id=f.id GROUP BY area + + -- Full-text search for content across all files + SELECT f.path, snippet(file_contents, 1, '>>>', '<<<', '...', 20) as match FROM file_contents fc JOIN files f ON f.id=fc.file_id WHERE file_contents MATCH 'handleRequest' ORDER BY rank LIMIT 10 + + -- FTS search scoped to a directory + SELECT f.path, snippet(file_contents, 1, '>>>', '<<<', '...', 20) as match FROM file_contents fc JOIN files f ON f.id=fc.file_id WHERE file_contents MATCH 'database migration' AND f.path LIKE 'backend/%' ORDER BY rank LIMIT 10`; + +function findGitRoot(cwd: string): string | null { + try { + const { execSync } = require("node:child_process"); + return execSync("git rev-parse --show-toplevel", { + cwd, + encoding: "utf-8", + stdio: ["pipe", "pipe", "pipe"], + }).trim(); + } catch { + return null; + } +} + +function codexisAvailable(): boolean { + try { + const { execSync } = require("node:child_process"); + execSync("which codexis", { stdio: ["pipe", "pipe", "pipe"] }); + return true; + } catch { + return false; + } +} + +export default function (pi: ExtensionAPI) { + if (!codexisAvailable()) { + return; + } + + pi.registerTool({ + name: "Codexis", + label: "Codexis", + description: DESCRIPTION, + parameters: Type.Object({ + type: Type.Enum( + { + sql: "sql", + call_graph: "call_graph", + }, + { + description: + "Query type: 'sql' for custom SQL queries, 'call_graph' for function call graph", + }, + ), + sql: Type.Optional( + Type.String({ + description: + "SQL query to run (required when type='sql', ignored when type='call_graph')", + }), + ), + }), + + async execute(_toolCallId, params, signal, _onUpdate, ctx) { + const gitRoot = findGitRoot(ctx.cwd); + if (!gitRoot) { + throw new Error("Not in a git repository"); + } + + const dbPath = join(gitRoot, ".codexis", "index.db"); + + // Run incremental index to keep DB fresh (fast if nothing changed) + const indexResult = await pi.exec("codexis", [gitRoot], { + signal, + timeout: 120_000, + }); + + if (indexResult.code !== 0) { + throw new Error( + `codexis indexing failed (exit ${indexResult.code}): ${indexResult.stderr}`, + ); + } + + if (!existsSync(dbPath)) { + throw new Error( + "codexis ran but no index.db was produced. Check codexis output.", + ); + } + + const db = new Database(dbPath, { readonly: true }); + try { + let rows: Record[]; + + if (params.type === "call_graph") { + const callGraphSql = ` + WITH defs AS ( + SELECT s.id, s.name, s.line as start_line, COALESCE(s.line_end, s.line + 1) as end_line, + f.path, f.id as file_id + FROM symbols s JOIN files f ON s.file_id = f.id + WHERE s.kind IN ('function', 'method') AND s.parent_id IS NULL + ), + func_counts AS ( + SELECT name, COUNT(*) as n FROM defs GROUP BY name + ), + calls AS ( + SELECT s.name, s.line as call_line, f.id as caller_file_id + FROM symbols s JOIN files f ON s.file_id = f.id + WHERE s.kind = 'reference' + ), + call_sites AS ( + SELECT DISTINCT + caller.name as caller, + caller.path as caller_file, + caller.start_line || '-' || caller.end_line as caller_lines, + c.call_line, + c.name as callee_name, + fc.n as defs_count + FROM calls c + JOIN defs caller ON c.caller_file_id = caller.file_id + AND c.call_line >= caller.start_line AND c.call_line <= caller.end_line + JOIN func_counts fc ON c.name = fc.name + ) + SELECT + caller, + caller_file, + caller_lines, + call_line, + callee_name as callee, + CASE WHEN defs_count > 1 THEN '' ELSE callee.path END as callee_file, + CASE WHEN defs_count > 1 THEN '' ELSE callee.start_line || '-' || callee.end_line END as callee_lines + FROM call_sites + LEFT JOIN defs callee ON call_sites.callee_name = callee.name AND defs_count <= 1 + ORDER BY caller_file, call_line; + `; + rows = db.prepare(callGraphSql).all(); + } else { + // type === 'sql' mode + if (!params.sql) { + throw new Error("sql parameter is required when type='sql'"); + } + const normalized = params.sql.trim().toUpperCase(); + if ( + !normalized.startsWith("SELECT") && + !normalized.startsWith("WITH") && + !normalized.startsWith("EXPLAIN") && + !normalized.startsWith("PRAGMA") + ) { + throw new Error( + "Only SELECT, WITH, EXPLAIN, and PRAGMA queries are allowed", + ); + } + rows = db.prepare(params.sql).all(); + } + + if (rows.length === 0) { + return { + content: [{ type: "text", text: "No results." }], + details: { rowCount: 0 }, + }; + } + + if (params.type === "call_graph") { + // Format call graph with file aliases + const fileMap = new Map(); + let aliasCounter = 1; + const fileToAlias = (path: string) => { + if (!fileMap.has(path)) { + fileMap.set(path, String.fromCharCode(64 + aliasCounter++)); // A, B, C... + } + return fileMap.get(path)!; + }; + + // First pass: collect all files + for (const row of rows) { + const r = row as Record; + fileToAlias(String(r.caller_file)); + if (r.callee_file) { + fileToAlias(String(r.callee_file)); + } + } + + // Build file legend + const fileLegend = Array.from(fileMap.entries()) + .map(([path, alias]) => ` ${alias} = ${path}`) + .join("\n"); + + // Group calls by caller + const grouped = new Map>>(); + for (const row of rows) { + const key = `${row.caller} (${row.caller_file}:${row.caller_lines})`; + if (!grouped.has(key)) { + grouped.set(key, []); + } + grouped.get(key)!.push(row); + } + + // Build call graph output + const callGraphLines: string[] = []; + callGraphLines.push("Files:"); + callGraphLines.push(fileLegend); + callGraphLines.push(""); + + for (const [callerKey, calls] of grouped) { + callGraphLines.push(callerKey); + for (const call of calls) { + const r = call as Record; + const callLine = String(r.call_line); + const callee = String(r.callee); + const calleeFile = String(r.callee_file); + const calleeLines = String(r.callee_lines); + const calleeAlias = calleeFile ? fileToAlias(calleeFile) : "??"; + const target = calleeFile + ? `${calleeAlias}:${calleeLines}` + : "?? (ambiguous)"; + callGraphLines.push( + ` ${callLine.padEnd(4)} -> ${callee.padEnd(20)} ${target}`, + ); + } + } + + const result = callGraphLines.join("\n"); + return { + content: [{ type: "text", text: result }], + details: { rowCount: rows.length }, + }; + } + + // Default: Format as aligned text table + const columns = Object.keys(rows[0] as Record); + const data = rows.map((row) => { + const r = row as Record; + return columns.map((col) => String(r[col] ?? "NULL")); + }); + + const widths = columns.map((col, i) => + Math.max(col.length, ...data.map((row) => row[i].length)), + ); + + const header = columns + .map((col, i) => col.padEnd(widths[i])) + .join(" "); + const separator = widths.map((w) => "-".repeat(w)).join(" "); + const body = data + .map((row) => row.map((val, i) => val.padEnd(widths[i])).join(" ")) + .join("\n"); + + const result = `${header}\n${separator}\n${body}`; + + // Truncate if huge + const maxLen = 48000; + const truncated = + result.length > maxLen + ? result.slice(0, maxLen) + + `\n\n[Truncated: ${rows.length} rows total, showing partial results. Narrow your query.]` + : result; + + return { + content: [{ type: "text", text: truncated }], + details: { rowCount: rows.length }, + }; + } finally { + db.close(); + } + }, + }); +} diff --git a/extension/package.json b/extension/package.json new file mode 100644 index 0000000..0dd22b4 --- /dev/null +++ b/extension/package.json @@ -0,0 +1,19 @@ +{ + "name": "codexis", + "private": true, + "version": "1.0.0", + "type": "module", + "scripts": { + "clean": "echo 'nothing to clean'", + "build": "echo 'nothing to build'", + "check": "echo 'nothing to check'" + }, + "pi": { + "extensions": [ + "./index.ts" + ] + }, + "dependencies": { + "better-sqlite3": "^12.9.0" + } +}