codexis/extension/codexis.ts

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