/** * 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. */ import { Type } from "@mariozechner/pi-ai"; import { defineTool, type ExtensionAPI } from "@mariozechner/pi-coding-agent"; import { execSync } from "node:child_process"; import { existsSync } from "node:fs"; import { join } from "node:path"; import Database from "better-sqlite3"; const SCHEMA = `-- .codexis/index.db schema: -- -- files: indexed source files -- id INTEGER PRIMARY KEY -- path TEXT NOT NULL UNIQUE -- relative to repo root -- language TEXT NOT NULL -- e.g. 'go', 'typescript', 'python', 'tsx', 'proto' -- package TEXT -- package/module name (from AST or directory) -- hash TEXT NOT NULL -- sha256, for incremental indexing -- indexed_at DATETIME -- -- symbols: definitions extracted via tree-sitter -- id INTEGER PRIMARY KEY -- file_id INTEGER NOT NULL REFERENCES files(id) -- name TEXT NOT NULL -- kind TEXT NOT NULL -- one of: 'function','method','class','type','interface','constant','variable','constructor' -- line INTEGER NOT NULL -- 1-indexed -- line_end INTEGER -- end of definition body -- col INTEGER -- col_end INTEGER -- exported BOOLEAN -- language-specific visibility -- parent_id INTEGER REFERENCES symbols(id) -- e.g. method→class, field→struct`; 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`; function findGitRoot(cwd: string): string | null { try { return execSync("git rev-parse --show-toplevel", { cwd, encoding: "utf-8", stdio: ["pipe", "pipe", "pipe"], }).trim(); } catch { return null; } } function findDatabase(cwd: string): string | null { const gitRoot = findGitRoot(cwd); if (!gitRoot) return null; const dbPath = join(gitRoot, ".codexis", "index.db"); if (!existsSync(dbPath)) return null; return dbPath; } const codexisTool = defineTool({ 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 dbPath = findDatabase(ctx.cwd); if (!dbPath) { throw new Error( "No code index found. Run `codexis` in the repo root to generate .codexis/index.db" ); } 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(); } }, }); export default function (pi: ExtensionAPI) { pi.registerTool(codexisTool); }