popagent

Latest change 39b4a69ecda2154310822d948569825fcb333b1c - Size code-index embedding batches for the local CPU embedder by AkurAI Build

import { Client } from "@modelcontextprotocol/sdk/client/index.js";
import {
  getDefaultEnvironment,
  StdioClientTransport,
} from "@modelcontextprotocol/sdk/client/stdio.js";
import { createHash } from "node:crypto";
import { mkdir, readFile, writeFile } from "node:fs/promises";
import { dirname, isAbsolute, join, posix, resolve } from "node:path";
import { fileURLToPath } from "node:url";
import { agentWorkspaces, type AgentWorkspaceStore } from "./agent-workspaces";

const MAX_CONTEXT_RESULTS = 20;
const MAX_IMPACT_PATHS = 8;
const DEFAULT_TOKEN_BUDGET = 4_000;
const INDEX_REQUEST_TIMEOUT_MS = 240_000;
// A cold full index of a repository on the local CPU embedder takes many
// minutes, far beyond one agent tool call. Warm-ups run in the background
// with their own generous budget; agent queries only ever wait for the fast
// incremental refresh of an already-warm index.
const INDEX_WARM_TIMEOUT_MS = 45 * 60_000;
const INDEX_EXCLUDE = "/.codebase-index/";

export type CodeIndexCall = {
  name: string;
  arguments: Record<string, unknown>;
};

export type CodeIndexResult = {
  text: string;
  isError: boolean;
};

export interface CodeIndexSession {
  call(call: CodeIndexCall, signal?: AbortSignal, timeoutMs?: number): Promise<CodeIndexResult>;
  close(): Promise<void>;
}

export type CodeIndexWarmth =
  | { state: "ready" }
  | { state: "warming"; startedAt: number }
  | { state: "cold" };

export interface CodeIndexBackend {
  execute(repositoryPath: string, calls: CodeIndexCall[], signal?: AbortSignal): Promise<CodeIndexResult[]>;
  /** Build the repository index in the background; resolves when it is ready. */
  warm?(repositoryPath: string): Promise<void>;
  warmth?(repositoryPath: string): CodeIndexWarmth;
}
type McpTextContent = { type: "text"; text: string };

type EmbeddingConfig = {
  baseUrl: string;
  model: string;
  dimensions: number;
};

type BackendOptions = {
  dataRoot?: string;
  embedding?: EmbeddingConfig;
  openSession?: (input: {
    repositoryPath: string;
    configPath: string;
  }) => Promise<CodeIndexSession>;
};

type WorkspaceResolver = Pick<AgentWorkspaceStore, "resolveRepository">;
type ChangedFiles = (repositoryPath: string) => Promise<string[]>;

function messageOf(error: unknown): string {
  return error instanceof Error ? error.message : String(error);
}

function boundedInteger(value: number, minimum: number, maximum: number, name: string): number {
  if (!Number.isInteger(value) || value < minimum || value > maximum) {
    throw new Error(`${name} must be an integer from ${minimum} to ${maximum}`);
  }
  return value;
}

function normalizedRelativePath(value: string): string {
  const normalized = posix.normalize(value.trim().replaceAll("\\", "/"));
  if (!normalized || normalized === "." || isAbsolute(normalized) || normalized === ".." || normalized.startsWith("../")) {
    throw new Error("Code intelligence paths must be repository-relative");
  }
  return normalized;
}

async function git(repositoryPath: string, args: string[]): Promise<string> {
  const child = Bun.spawn(["git", "-C", repositoryPath, ...args], {
    stdout: "pipe",
    stderr: "pipe",
  });
  if (await child.exited !== 0) {
    throw new Error((await new Response(child.stderr).text()).trim() || "Git inspection failed");
  }
  return new Response(child.stdout).text();
}

async function changedFiles(repositoryPath: string): Promise<string[]> {
  const status = await git(repositoryPath, ["status", "--porcelain=v1"]);
  return status.split("\n")
    .filter(Boolean)
    .map((line) => line.slice(3).replace(/^.* -> /, ""))
    .map(normalizedRelativePath)
    .slice(0, MAX_IMPACT_PATHS);
}

function configuredEmbedding(): EmbeddingConfig {
  const baseUrl = (process.env.POPAGENT_EMBEDDING_URL ?? process.env.NINEROUTER_URL)?.replace(/\/$/, "");
  if (!baseUrl) throw new Error("POPAGENT_EMBEDDING_URL or NINEROUTER_URL is required for code intelligence");
  const dimensions = Number(process.env.POPAGENT_EMBEDDING_DIMENSION ?? 1_536);
  if (!Number.isInteger(dimensions) || dimensions < 1) {
    throw new Error("POPAGENT_EMBEDDING_DIMENSION must be a positive integer");
  }
  return {
    baseUrl,
    model: process.env.POPAGENT_EMBEDDING_MODEL ?? "text-embedding-3-small",
    dimensions,
  };
}

function indexConfig(embedding: EmbeddingConfig): string {
  return `${JSON.stringify({
    embeddingProvider: "custom",
    customProvider: {
      baseUrl: embedding.baseUrl,
      model: embedding.model,
      dimensions: embedding.dimensions,
      // The local llama.cpp embedder is CPU-bound and serves one request at a
      // time: 8k-token chunks at concurrency 3 exceeded the 120 s request
      // budget and every batch failed, so the index never filled. Smaller
      // chunks, one request in flight, and a generous timeout let it finish.
      maxTokens: 2_048,
      timeoutMs: 600_000,
      concurrency: 1,
      requestIntervalMs: 0,
    },
    scope: "project",
    indexing: {
      autoIndex: false,
      watchFiles: false,
      requireProjectMarker: true,
      maxDepth: -1,
      maxFilesPerDirectory: 10_000,
      fallbackToTextOnMaxChunks: true,
    },
    additionalInclude: [
      "**/*.sql",
      "**/*.css",
      "**/*.html",
      "**/*.sh",
      "**/*.md",
      "**/*.json",
      "**/*.yml",
      "**/*.yaml",
      "**/*.toml",
    ],
    search: {
      maxResults: MAX_CONTEXT_RESULTS,
      fusionStrategy: "rrf",
      rerankTopN: 20,
    },
  }, null, 2)}\n`;
}

async function ensureFile(path: string, content: string, mode: number): Promise<void> {
  const current = await readFile(path, "utf8").catch(() => undefined);
  if (current === content) return;
  await mkdir(dirname(path), { recursive: true, mode: 0o700 });
  await writeFile(path, content, { mode });
}

async function excludeIndexFromGit(repositoryPath: string): Promise<void> {
  const fallback = join(repositoryPath, ".git", "info", "exclude");
  const gitPath = await git(repositoryPath, ["rev-parse", "--path-format=absolute", "--git-path", "info/exclude"])
    .then((path) => path.trim())
    .catch(() => fallback);
  const existing = await readFile(gitPath, "utf8").catch(() => "");
  if (existing.split("\n").includes(INDEX_EXCLUDE)) return;
  const prefix = existing && !existing.endsWith("\n") ? "\n" : "";
  await mkdir(dirname(gitPath), { recursive: true, mode: 0o700 });
  await writeFile(gitPath, `${existing}${prefix}# Popagent derived code intelligence\n${INDEX_EXCLUDE}\n`, { mode: 0o600 });
}

function codebaseIndexCli(): string {
  const packageEntry = fileURLToPath(import.meta.resolve("opencode-codebase-index"));
  return join(dirname(packageEntry), "cli.js");
}

async function openMcpSession(input: {
  repositoryPath: string;
  configPath: string;
}): Promise<CodeIndexSession> {
  const transport = new StdioClientTransport({
    command: process.execPath,
    args: [
      codebaseIndexCli(),
      "--project",
      input.repositoryPath,
      "--host",
      "codex",
      "--config",
      input.configPath,
    ],
    cwd: input.repositoryPath,
    env: getDefaultEnvironment(),
    stderr: "pipe",
  });
  transport.stderr?.on("data", () => undefined);
  const client = new Client({ name: "popagent-code-intelligence", version: "1.0.0" });
  await client.connect(transport);
  return {
    call: async ({ name, arguments: arguments_ }, signal, timeoutMs = INDEX_REQUEST_TIMEOUT_MS) => {
      const result = await client.callTool(
        { name, arguments: arguments_ },
        undefined,
        {
          signal,
          timeout: timeoutMs,
          resetTimeoutOnProgress: true,
          maxTotalTimeout: timeoutMs,
        },
      );
      if (!("content" in result) || !Array.isArray(result.content)) {
        return { text: String(result.toolResult ?? ""), isError: false };
      }
      const text = result.content
        .filter((part): part is McpTextContent =>
          typeof part === "object" && part !== null && "type" in part && part.type === "text" && "text" in part && typeof part.text === "string")
        .map((part) => part.text)
        .join("\n")
        .trim();
      return { text, isError: Boolean(result.isError) };
    },
    close: () => client.close(),
  };
}

export class OpenCodebaseIndexBackend implements CodeIndexBackend {
  private readonly dataRoot: string;
  private readonly embedding: EmbeddingConfig;
  private readonly openSession: NonNullable<BackendOptions["openSession"]>;

  constructor(options: BackendOptions = {}) {
    this.dataRoot = resolve(options.dataRoot ?? join(process.env.POPAGENT_DATA_DIR ?? join(process.cwd(), "data"), "code-intelligence"));
    this.embedding = options.embedding ?? configuredEmbedding();
    this.openSession = options.openSession ?? openMcpSession;
  }

  private readonly ready = new Set<string>();
  private readonly warming = new Map<string, { startedAt: number; promise: Promise<void> }>();
  // The local embedder serves one batch at a time, so parallel warm-ups only
  // starve each other; run them one after another.
  private warmQueue: Promise<void> = Promise.resolve();

  private keyFor(repositoryPath: string): string {
    return createHash("sha256").update(resolve(repositoryPath)).digest("hex");
  }

  private readyMarker(key: string): string {
    return join(this.dataRoot, "ready", `${key}`);
  }

  private async isReady(repositoryPath: string): Promise<boolean> {
    const key = this.keyFor(repositoryPath);
    if (this.ready.has(key)) return true;
    try {
      const marker = await readFile(this.readyMarker(key), "utf8");
      // The marker records the repository path so a stale hash collision or a
      // moved checkout never reads as warm.
      if (marker.trim() === resolve(repositoryPath)) {
        this.ready.add(key);
        return true;
      }
    } catch {
      // No marker: cold.
    }
    return false;
  }

  warmth(repositoryPath: string): CodeIndexWarmth {
    const key = this.keyFor(repositoryPath);
    if (this.ready.has(key)) return { state: "ready" };
    const active = this.warming.get(key);
    return active ? { state: "warming", startedAt: active.startedAt } : { state: "cold" };
  }

  async warm(repositoryPath: string): Promise<void> {
    const key = this.keyFor(repositoryPath);
    if (this.ready.has(key)) return;
    const active = this.warming.get(key);
    if (active) return active.promise;
    const startedAt = Date.now();
    const promise = (this.warmQueue = this.warmQueue.catch(() => undefined).then(async () => {
      if (await this.isReady(repositoryPath)) return;
      const session = await this.session(repositoryPath);
      try {
        const refreshed = await session.call({
          name: "index_codebase",
          arguments: { force: false, estimateOnly: false, verbose: false },
        }, undefined, INDEX_WARM_TIMEOUT_MS);
        if (refreshed.isError) throw new Error(refreshed.text || "Code intelligence index warm-up failed");
      } finally {
        await session.close();
      }
      await ensureFile(this.readyMarker(key), resolve(repositoryPath), 0o600);
      this.ready.add(key);
    })).finally(() => {
      this.warming.delete(key);
    });
    this.warming.set(key, { startedAt, promise });
    return promise;
  }

  private async session(repositoryPath: string): Promise<CodeIndexSession> {
    await excludeIndexFromGit(repositoryPath);
    const configPath = join(this.dataRoot, "configs", `${this.keyFor(repositoryPath)}.json`);
    await ensureFile(configPath, indexConfig(this.embedding), 0o600);
    return this.openSession({ repositoryPath, configPath });
  }

  async execute(repositoryPath: string, calls: CodeIndexCall[], signal?: AbortSignal): Promise<CodeIndexResult[]> {
    signal?.throwIfAborted();
    if (!(await this.isReady(repositoryPath))) {
      // Never make an agent tool call wait for a cold full index. Start (or
      // keep) the background warm-up and answer immediately so the agent can
      // fall back to grep/read and retry semantic lookups later.
      const warmth = this.warmth(repositoryPath);
      void this.warm(repositoryPath).catch(() => undefined);
      const elapsed = warmth.state === "warming" ? Math.round((Date.now() - warmth.startedAt) / 1000) : 0;
      const text = warmth.state === "warming"
        ? `Code intelligence index for this repository is still warming (started ${elapsed}s ago). Use workspace grep/read for now and retry this lookup later in the task.`
        : "Code intelligence index for this repository is cold; a background warm-up has started. Use workspace grep/read for now and retry this lookup later in the task.";
      return calls.map(() => ({ text, isError: true }));
    }
    const session = await this.session(repositoryPath);
    try {
      const refreshed = await session.call({
        name: "index_codebase",
        arguments: { force: false, estimateOnly: false, verbose: false },
      }, signal);
      if (refreshed.isError) throw new Error(refreshed.text || "Code intelligence index refresh failed");
      const results: CodeIndexResult[] = [];
      for (const call of calls) {
        signal?.throwIfAborted();
        results.push(await session.call(call, signal));
      }
      return results;
    } finally {
      await session.close();
    }
  }
}

export class CodeIntelligenceService {
  constructor(
    private readonly workspaces: WorkspaceResolver = agentWorkspaces,
    private readonly backend: CodeIndexBackend = new OpenCodebaseIndexBackend(),
    private readonly listChangedFiles: ChangedFiles = changedFiles,
  ) {}

  private async repository(workspaceId: string, repositoryPath?: string): Promise<string> {
    if (repositoryPath) return resolve(repositoryPath);
    return (await this.workspaces.resolveRepository(workspaceId)).path;
  }

  /**
   * Build the semantic index for a repository ahead of agent work so research
   * can use natural-language lookups immediately instead of re-reading files.
   * Resolves when the index is ready; safe to call repeatedly.
   */
  async warm(repositoryPath: string): Promise<void> {
    await this.backend.warm?.(resolve(repositoryPath));
  }

  warmth(repositoryPath: string): CodeIndexWarmth {
    return this.backend.warmth?.(resolve(repositoryPath)) ?? { state: "ready" };
  }

  async codeContext(
    workspaceId: string,
    query: string,
    limit = 8,
    repositoryPath?: string,
    signal?: AbortSignal,
  ): Promise<{ evidence: string }> {
    const request = query.trim();
    if (!request || request.length > 2_000) throw new Error("Code context query must be 1–2000 characters");
    boundedInteger(limit, 1, MAX_CONTEXT_RESULTS, "Code context limit");
    const [result] = await this.backend.execute(await this.repository(workspaceId, repositoryPath), [{
      name: "codebase_context",
      arguments: { query: request, limit, tokenBudget: DEFAULT_TOKEN_BUDGET },
    }], signal);
    if (!result || result.isError) throw new Error(result?.text || "Code context lookup failed");
    return { evidence: result.text };
  }

  async symbolContext(
    workspaceId: string,
    symbol: string,
    filePath?: string,
    repositoryPath?: string,
    signal?: AbortSignal,
  ): Promise<{ definition: string; callers: string; callees: string }> {
    const name = symbol.trim();
    if (!name || name.length > 512) throw new Error("Symbol must be 1–512 characters");
    const normalizedFile = filePath ? normalizedRelativePath(filePath) : undefined;
    const directory = normalizedFile?.includes("/") ? posix.dirname(normalizedFile) : undefined;
    const results = await this.backend.execute(await this.repository(workspaceId, repositoryPath), [
      {
        name: "implementation_lookup",
        arguments: { query: name, limit: 5, ...(directory ? { directory } : {}) },
      },
      {
        name: "call_graph",
        arguments: { name, direction: "callers", ...(normalizedFile ? { filePath: normalizedFile } : {}) },
      },
      {
        name: "call_graph",
        arguments: { name, direction: "callees", ...(normalizedFile ? { filePath: normalizedFile } : {}) },
      },
    ], signal);
    const [definition, callers, callees] = results;
    if (!definition || !callers || !callees) throw new Error("Symbol context lookup returned incomplete evidence");
    return {
      definition: definition.text,
      callers: callers.text,
      callees: callees.text,
    };
  }

  async changeImpact(
    workspaceId: string,
    paths: string[] = [],
    repositoryPath?: string,
    signal?: AbortSignal,
  ): Promise<{
    paths: string[];
    branchImpact: string;
    files: Array<{ path: string; evidence: string }>;
    unresolved: string[];
  }> {
    const repository = await this.repository(workspaceId, repositoryPath);
    const selected = [...new Set((paths.length ? paths : await this.listChangedFiles(repository)).map(normalizedRelativePath))]
      .slice(0, MAX_IMPACT_PATHS);
    const calls: CodeIndexCall[] = [{
      name: "pr_impact",
      arguments: { maxDepth: 5, direction: "both", checkConflicts: false },
    }, ...selected.map((path) => ({
      name: "codebase_edit_context",
      arguments: {
        query: `Assess the blast radius and affected tests for changes in ${path}.`,
        filePath: path,
        callerLimit: 8,
        calleeLimit: 8,
        tokenBudget: DEFAULT_TOKEN_BUDGET,
      },
    }))];
    const [branch, ...files] = await this.backend.execute(repository, calls, signal);
    const unresolved: string[] = [];
    if (!selected.length) unresolved.push("No working-tree paths were supplied or changed.");
    files.forEach((result, index) => {
      if (result?.isError) unresolved.push(`${selected[index]}: ${result.text || "impact lookup failed"}`);
    });
    return {
      paths: selected,
      branchImpact: branch?.text ?? "Branch impact unavailable.",
      files: files.map((result, index) => ({ path: selected[index]!, evidence: result?.text ?? "" })),
      unresolved,
    };
  }
}

export const codeIntelligence = new CodeIntelligenceService();

export function codeIntelligenceError(error: unknown): string {
  return messageOf(error);
}