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.
 *
 * - 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";

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