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