304 lines
11 KiB
TypeScript
304 lines
11 KiB
TypeScript
/**
|
|
* 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<string, unknown>[];
|
|
|
|
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<string, string>();
|
|
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<string, unknown>;
|
|
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<string, Array<Record<string, unknown>>>();
|
|
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<string, unknown>;
|
|
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<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();
|
|
}
|
|
},
|
|
});
|
|
}
|