From dfd61f899a115ae93f0556dd3bb72e603e0aff47 Mon Sep 17 00:00:00 2001 From: Evan Reichard Date: Fri, 10 Apr 2026 15:56:47 -0400 Subject: [PATCH] refactor(extension): auto-index on tool call and skip registration when binary missing - Only register tool if `codexis` binary is in PATH - Run incremental index via `pi.exec` on each tool call to keep DB fresh - Remove `findDatabase` helper; derive DB path from git root directly - Replace `defineTool` with inline `pi.registerTool` call - Update imports (`@sinclair/typebox`, lazy `execSync`) - Fix output flag help text in main.go --- extension/codexis.ts | 210 ++++++++++++++++++++++++------------------- main.go | 2 +- 2 files changed, 119 insertions(+), 93 deletions(-) diff --git a/extension/codexis.ts b/extension/codexis.ts index b123365..169a508 100644 --- a/extension/codexis.ts +++ b/extension/codexis.ts @@ -3,11 +3,14 @@ * * 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 "@mariozechner/pi-ai"; -import { defineTool, type ExtensionAPI } from "@mariozechner/pi-coding-agent"; -import { execSync } from "node:child_process"; +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"; @@ -56,6 +59,7 @@ Example queries: function findGitRoot(cwd: string): string | null { try { + const { execSync } = require("node:child_process"); return execSync("git rev-parse --show-toplevel", { cwd, encoding: "utf-8", @@ -66,96 +70,118 @@ function findGitRoot(cwd: string): string | 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; +function codexisAvailable(): boolean { + try { + const { execSync } = require("node:child_process"); + execSync("which codexis", { stdio: ["pipe", "pipe", "pipe"] }); + return true; + } catch { + return false; + } } -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); + 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/main.go b/main.go index 4bc1220..e4c5887 100644 --- a/main.go +++ b/main.go @@ -20,7 +20,7 @@ const dbFileName = "index.db" func main() { force := flag.Bool("force", false, "Force full re-index (ignore file hashes)") - output := flag.String("o", "", "Output database path (default: /.codexis.db)") + output := flag.String("o", "", "Output database path (default: /.codexis/index.db)") flag.Parse() root := "."