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
This commit is contained in:
@@ -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<string, unknown>);
|
||||
const data = rows.map((row) => {
|
||||
const r = row as Record<string, unknown>;
|
||||
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<string, unknown>);
|
||||
const data = rows.map((row) => {
|
||||
const r = row as Record<string, unknown>;
|
||||
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();
|
||||
}
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user