pi-mega-compact 0.4.24 → 0.4.26

package/src/store/sqlite.ts

@@ -872,6 +872,19 @@ export function hasCheckpoint(sessionId: string, checkpointId: string, stateDir:
   return row !== undefined;
 }
 
+/** Fetch a single checkpoint by (session, id), or undefined if absent. */
+export function getCheckpoint(
+  sessionId: string,
+  checkpointId: string,
+  stateDir: string = getStateDir(),
+): StoredCheckpoint | undefined {
+  const db = openStore(stateDir);
+  const row = db
+    .prepare("SELECT * FROM context_chunks WHERE session_id = ? AND id = ? LIMIT 1")
+    .get(normalizeSessionId(sessionId), checkpointId) as any;
+  return row ? rowToCheckpoint(row) : undefined;
+}
+
 /** Mark a checkpoint's dedup_status (e.g. 'removed' by SemDeDup). */
 export function setDedupStatus(
   checkpointId: string,

package/src/store/vectorIndex.test.ts

@@ -0,0 +1,116 @@
+/**
+ * vectorIndex.test.ts — Slice 2 async PGlite/HNSW vector index.
+ *
+ * Proves: cross-repo nearest-neighbor recall, repoId scoping, the dimension
+ * guard (non-512 vectors skipped, never corrupt the index), and graceful
+ * degradation when the index is disabled (kill-switch) — all without touching
+ * the synchronous node:sqlite store.
+ *
+ * The index is a WASM Postgres (PGlite) — fully local, zero network
+ * (PREVENT-PI-004). Each test isolates state via MEGACOMPACT_VECTOR_INDEX_DIR.
+ */
+
+import { test } from "node:test";
+import assert from "node:assert/strict";
+import { mkdtempSync, rmSync } from "node:fs";
+import { tmpdir } from "node:os";
+import { join } from "node:path";
+import {
+  EMBEDDING_DIM,
+  initVectorIndex,
+  upsertEmbedding,
+  searchAsync,
+  closeVectorIndex,
+  isVectorIndexDisabled,
+} from "./vectorIndex.js";
+
+/** A 512-dim unit-ish vector with a single spike at `idx` (deterministic NN). */
+function spikeVec(idx: number, magnitude = 1): number[] {
+  const v = new Array<number>(EMBEDDING_DIM).fill(0);
+  v[idx % EMBEDDING_DIM] = magnitude;
+  return v;
+}
+
+function isolateIndexDir(): string {
+  const dir = mkdtempSync(join(tmpdir(), "mc-vecidx-"));
+  process.env.MEGACOMPACT_VECTOR_INDEX_DIR = dir;
+  return dir;
+}
+
+test("cross-repo HNSW nearest-neighbor recall across repos + repoId scoping", async () => {
+  delete process.env.MEGACOMPACT_PGLITE_DISABLED;
+  const dir = isolateIndexDir();
+  try {
+    await closeVectorIndex(); // ensure a fresh singleton for this dir
+    const pg = await initVectorIndex();
+    assert.ok(pg, "index should initialize (PGlite WASM available)");
+
+    // repoA: two checkpoints; repoB: one checkpoint. Distinct spike directions.
+    await upsertEmbedding("/repoA/.pi/mega-compact", "sessA", "chkpt_001", spikeVec(0));
+    await upsertEmbedding("/repoA/.pi/mega-compact", "sessA", "chkpt_002", spikeVec(5));
+    await upsertEmbedding("/repoB/.pi/mega-compact", "sessB", "chkpt_001", spikeVec(0));
+
+    // Cross-repo query near spike(0): nearest are the two spike(0) rows, one per repo.
+    const cross = await searchAsync(spikeVec(0), { k: 2 });
+    assert.equal(cross.length, 2, "cross-repo returns two nearest");
+    const repos = new Set(cross.map((h) => h.repoId));
+    assert.ok(repos.has("/repoA/.pi/mega-compact"), "hit from repoA");
+    assert.ok(repos.has("/repoB/.pi/mega-compact"), "hit from repoB");
+    assert.ok(cross[0].score > 0.99, "top hit is near-identical (cosine ~1)");
+
+    // Scoped to repoA only: excludes repoB even though repoB has an identical vec.
+    const scoped = await searchAsync(spikeVec(0), { k: 5, repoId: "/repoA/.pi/mega-compact" });
+    assert.ok(scoped.length >= 1, "scoped returns repoA hits");
+    assert.ok(
+      scoped.every((h) => h.repoId === "/repoA/.pi/mega-compact"),
+      "repoId filter excludes other repos",
+    );
+  } finally {
+    await closeVectorIndex();
+    rmSync(dir, { recursive: true, force: true });
+    delete process.env.MEGACOMPACT_VECTOR_INDEX_DIR;
+  }
+});
+
+test("dimension guard: non-512 vectors are skipped, never corrupt the index", async () => {
+  delete process.env.MEGACOMPACT_PGLITE_DISABLED;
+  const dir = isolateIndexDir();
+  try {
+    await closeVectorIndex();
+    await initVectorIndex();
+    // Wrong-dimension vector (BYO embedder mismatch) must be silently skipped.
+    await upsertEmbedding("/repoC/.pi/mega-compact", "sessC", "chkpt_001", [1, 2, 3]);
+    const hits = await searchAsync(spikeVec(0), { k: 5 });
+    assert.equal(hits.length, 0, "no rows stored for a mismatched-dim vector");
+
+    // A correct-dim vector still stores fine afterward (index not corrupted).
+    await upsertEmbedding("/repoC/.pi/mega-compact", "sessC", "chkpt_002", spikeVec(3));
+    const ok = await searchAsync(spikeVec(3), { k: 1 });
+    assert.equal(ok.length, 1, "valid vector stored after a skipped one");
+    assert.equal(ok[0].checkpointId, "chkpt_002");
+  } finally {
+    await closeVectorIndex();
+    rmSync(dir, { recursive: true, force: true });
+    delete process.env.MEGACOMPACT_VECTOR_INDEX_DIR;
+  }
+});
+
+test("kill-switch: MEGACOMPACT_PGLITE_DISABLED disables the index gracefully", async () => {
+  const dir = isolateIndexDir();
+  process.env.MEGACOMPACT_PGLITE_DISABLED = "true";
+  try {
+    await closeVectorIndex();
+    assert.equal(isVectorIndexDisabled(), true, "kill-switch reported disabled");
+    const pg = await initVectorIndex();
+    assert.equal(pg, undefined, "init returns undefined when disabled");
+    // Upsert + search are no-ops that never throw and return empty.
+    await upsertEmbedding("/repoD/.pi/mega-compact", "sessD", "chkpt_001", spikeVec(0));
+    const hits = await searchAsync(spikeVec(0), { k: 3 });
+    assert.deepEqual(hits, [], "search returns [] when disabled");
+  } finally {
+    delete process.env.MEGACOMPACT_PGLITE_DISABLED;
+    await closeVectorIndex();
+    rmSync(dir, { recursive: true, force: true });
+    delete process.env.MEGACOMPACT_VECTOR_INDEX_DIR;
+  }
+});

package/src/store/vectorIndex.ts

@@ -0,0 +1,273 @@
+/**
+ * vectorIndex.ts — Slice 2 async vector index (PGlite/pgvector HNSW).
+ *
+ * A REDUNDANT, additive, ASYNC index layered over the synchronous node:sqlite
+ * store (which remains the authoritative source of truth). The sync linear
+ * cosine scan over `embedding_blob` stays the DEFAULT recall path; this index
+ * exists only to provide real cross-repo / cross-session HNSW nearest-neighbor
+ * recall. It is best-effort and non-fatal: any init/write failure degrades to
+ * the sync scan and must NEVER break add(), compaction, or extension load.
+ *
+ * PREVENT-PI-004: PGlite is WASM Postgres — fully local, zero network.
+ *
+ * Index topology (decision 2026-07-15): ONE global PGlite DB, `repo_id` is a
+ * first-class column. `searchAsync(q, k, {repoId?})` → omit repoId for cross-repo
+ * NN, pass repoId to scope to a single repo. The sync store is per-repo (state
+ * dir); this global index is the thing that makes cross-repo recall possible.
+ */
+
+import { homedir } from "node:os";
+import { join } from "node:path";
+import { mkdirSync, rmSync, existsSync } from "node:fs";
+
+// PGlite + pgvector are script-free WASM (no native build) → survive pi's
+// install-script block. Imported lazily so a missing/broken package degrades
+// gracefully instead of crashing module load.
+import { PGlite, type PGlite as PGliteInstance } from "@electric-sql/pglite";
+import { vector } from "@electric-sql/pglite-pgvector";
+
+/** Vector dimension produced by the default TrigramEmbedder (src/embedder.ts). */
+export const EMBEDDING_DIM = 512;
+
+/** A single recall hit returned by the async index. */
+export interface VectorIndexHit {
+  repoId: string;
+  sessionId: string;
+  checkpointId: string;
+  /** Cosine similarity in [0,1] (1 = identical). */
+  score: number;
+}
+
+let db: PGliteInstance | undefined;
+let initPromise: Promise<PGliteInstance | undefined> | undefined;
+let disabled = false;
+let warned = false;
+
+function indexDir(): string {
+  const override = process.env.MEGACOMPACT_VECTOR_INDEX_DIR;
+  if (override && override.trim() !== "") return override;
+  try {
+    return join(homedir(), ".pi", "mega-compact-vector");
+  } catch {
+    return join("/tmp", ".mega-compact-vector");
+  }
+}
+
+function logWarn(msg: string): void {
+  // Never throw — degradation is the whole point. One warning per process.
+  if (warned) return;
+  warned = true;
+  try {
+    console.warn(`[mega-compact:vectorIndex] ${msg} (falling back to sync scan)`);
+  } catch {
+    /* ignore */
+  }
+}
+
+/** Honor the emergency kill-switch. When set, the index is fully disabled. */
+export function isVectorIndexDisabled(): boolean {
+  return (
+    disabled ||
+    process.env.MEGACOMPACT_PGLITE_DISABLED === "true" ||
+    process.env.MEGACOMPACT_PGLITE_DISABLED === "1"
+  );
+}
+
+/**
+ * Lazily open + schema-init the global PGlite DB. Idempotent and safe to call
+ * from many places. Returns undefined when disabled/unavailable so callers can
+ * fall back to the synchronous scan. Never throws.
+ */
+export function initVectorIndex(): Promise<PGliteInstance | undefined> {
+  if (isVectorIndexDisabled()) return Promise.resolve(undefined);
+  if (db) return Promise.resolve(db);
+  if (initPromise) return initPromise;
+  initPromise = openPgLite(/* retryOnCorrupt */ true);
+  return initPromise;
+}
+
+/**
+ * Open + schema-init PGlite. When `retryOnCorrupt` is true, a WASM-level
+ * abort (typically from a corrupted/torn data dir) triggers a delete + one
+ * retry — the dir is rebuilt from scratch by PGlite's initdb.
+ */
+async function openPgLite(
+  retryOnCorrupt: boolean,
+): Promise<PGliteInstance | undefined> {
+  try {
+    const dir = indexDir();
+    mkdirSync(dir, { recursive: true });
+    const pg = await new PGlite({
+      dataDir: dir,
+      extensions: { vector },
+    });
+    await pg.exec("CREATE EXTENSION IF NOT EXISTS vector;");
+    await pg.exec(`
+      CREATE TABLE IF NOT EXISTS vector_index (
+        repo_id       TEXT NOT NULL,
+        session_id    TEXT NOT NULL,
+        checkpoint_id TEXT NOT NULL,
+        embedding     vector(${EMBEDDING_DIM}) NOT NULL,
+        PRIMARY KEY (repo_id, session_id, checkpoint_id)
+      );
+    `);
+    // HNSW index over cosine distance for fast NN. Created idempotently.
+    await pg.exec(
+      "CREATE INDEX IF NOT EXISTS vector_index_hnsw ON vector_index USING hnsw (embedding vector_cosine_ops);",
+    );
+    db = pg;
+    return pg;
+  } catch (err) {
+    const msg = err instanceof Error ? err.message : String(err);
+    // Self-heal: a WASM Aborted() typically means the data dir is corrupted
+    // (torn WAL from concurrent access). Delete it and retry once.
+    if (
+      retryOnCorrupt &&
+      (msg.includes("Aborted") || msg.includes("RuntimeError"))
+    ) {
+      try {
+        const dir = indexDir();
+        if (existsSync(dir)) {
+          rmSync(dir, { recursive: true, force: true });
+        }
+        // Clear singleton state so the retry starts fresh.
+        initPromise = undefined;
+        return openPgLite(/* retryOnCorrupt */ false);
+      } catch {
+        // Self-heal failed — fall through to disable.
+      }
+    }
+    disabled = true;
+    logWarn(`init failed: ${msg}`);
+    return undefined;
+  }
+}
+
+function toVectorLiteral(v: number[]): string {
+  // pgvector text form: [a,b,c]. Guard against NaN/Inf for a clean literal.
+  const parts = v.map((x) => (Number.isFinite(x) ? x : 0));
+  return `[${parts.join(",")}]`;
+}
+
+/**
+ * Best-effort upsert of one checkpoint embedding into the global index.
+ * Dimension-mismatched vectors (e.g. a BYO embedder with dim ≠ 512) are skipped
+ * rather than corrupting the index. Fire-and-forget: resolved promise only;
+ * callers must NOT await this on the sync path. Never throws.
+ */
+export async function upsertEmbedding(
+  repoId: string,
+  sessionId: string,
+  checkpointId: string,
+  embedding: number[],
+): Promise<void> {
+  if (isVectorIndexDisabled()) return;
+  if (!embedding || embedding.length !== EMBEDDING_DIM) {
+    // Dimension guard: skip without corrupting the fixed-dim index.
+    return;
+  }
+  try {
+    const pg = await initVectorIndex();
+    if (!pg) return;
+    const lit = toVectorLiteral(embedding);
+    await pg.query(
+      `INSERT INTO vector_index (repo_id, session_id, checkpoint_id, embedding)
+       VALUES ($1, $2, $3, $4::vector)
+       ON CONFLICT (repo_id, session_id, checkpoint_id)
+       DO UPDATE SET embedding = EXCLUDED.embedding;`,
+      [repoId, sessionId, checkpointId, lit],
+    );
+  } catch (err) {
+    disabled = true;
+    logWarn(`upsert failed: ${err instanceof Error ? err.message : String(err)}`);
+  }
+}
+
+export interface SearchAsyncOpts {
+  /** When provided, scope the NN search to a single repo; omit for cross-repo. */
+  repoId?: string;
+  /** Max hits (default 3). */
+  k?: number;
+}
+
+/**
+ * Cross-repo (or single-repo) HNSW nearest-neighbor search. Returns hits sorted
+ * by descending similarity. Never throws — on any failure returns [].
+ */
+export async function searchAsync(
+  query: number[],
+  opts: SearchAsyncOpts = {},
+): Promise<VectorIndexHit[]> {
+  if (isVectorIndexDisabled() || !query || query.length !== EMBEDDING_DIM) return [];
+  const k = opts.k ?? 3;
+  const repoId = opts.repoId;
+  try {
+    const pg = await initVectorIndex();
+    if (!pg) return [];
+    const lit = toVectorLiteral(query);
+    const params: unknown[] = [lit, k];
+    let sql =
+      "SELECT repo_id, session_id, checkpoint_id, 1 - (embedding <=> $1::vector) AS score " +
+      "FROM vector_index";
+    if (repoId) {
+      sql += " WHERE repo_id = $3";
+      params.push(repoId);
+    }
+    sql += " ORDER BY embedding <=> $1::vector LIMIT $2";
+    const res = await pg.query(sql, params);
+    return res.rows.map((r: any) => ({
+      repoId: r.repo_id as string,
+      sessionId: r.session_id as string,
+      checkpointId: r.checkpoint_id as string,
+      score: r.score as number,
+    }));
+  } catch (err) {
+    disabled = true;
+    logWarn(`search failed: ${err instanceof Error ? err.message : String(err)}`);
+    return [];
+  }
+}
+
+/** Close the index (test teardown / shutdown). Safe to call when unopened. */
+export async function closeVectorIndex(): Promise<void> {
+  if (db) {
+    try {
+      await db.close();
+    } catch {
+      /* ignore */
+    }
+  }
+  db = undefined;
+  initPromise = undefined;
+  disabled = false;
+  warned = false;
+}
+
+/**
+ * Rebuild the entire index from the authoritative node:sqlite store. Used for
+ * backfill + DR. `enumerateRepoStateDirs` yields each repo's state dir; we read
+ * its checkpoint embeddings and bulk upsert. Best-effort: counts successes and
+ * skips failures. Returns {upserted, errors}.
+ */
+export async function rebuildFromSqlite(
+  enumerateRepoStateDirs: () => Iterable<{ repoId: string; stateDir: string }>,
+  readCheckpoints: (
+    stateDir: string,
+  ) => Iterable<{ sessionId: string; checkpointId: string; embedding: number[] }>,
+): Promise<{ upserted: number; errors: number }> {
+  let upserted = 0;
+  let errors = 0;
+  const pg = await initVectorIndex();
+  if (!pg) return { upserted, errors: 1 };
+  for (const repo of enumerateRepoStateDirs()) {
+    for (const cp of readCheckpoints(repo.stateDir)) {
+      try {
+        await upsertEmbedding(repo.repoId, cp.sessionId, cp.checkpointId, cp.embedding);
+        upserted++;
+      } catch {
+        errors++;
+      }
+    }
+  }
+  return { upserted, errors };
+}

package/src/vectorStore.ts

@@ -26,6 +26,7 @@ import {
   listCheckpoints,
   nextCheckpointId,
   upsertCheckpoint,
+  getCheckpoint,
   loadSessionState,
   saveSessionState,
   upsertMinhashSignature,
@@ -38,6 +39,11 @@ import {
   repoStats as repoStatsFromStore,
   dataInvariantStats,
 } from "./store/sqlite.js";
+import {
+  initVectorIndex,
+  searchAsync as vectorIndexSearch,
+  type VectorIndexHit,
+} from "./store/vectorIndex.js";
 import { rehydrateRaptorTree } from "./dedup/raptor/index.js";
 import { stagedExpansion } from "./dedup/raptor/retrieval.js";
 import { migrateJsonToSqlite } from "./store/migrate.js";
@@ -98,6 +104,13 @@ export class VectorStore {
   private readonly cfg: DedupConfigShape;
   /** Optional monitoring target (Sprint 14). Undefined → no monitoring. */
   private readonly eventsPath?: string;
+  /**
+   * Repo key for the async PGlite vector index (Slice 2). We use the stateDir
+   * itself as the repo id — it is already unique per repo and available here
+   * without crossing into the pi-runtime layer (src/ stays pi-agnostic). The
+   * global index keys on repoId so recall can span repos.
+   */
+  private readonly repoId: string;
 
   constructor(
     opts: {
@@ -110,10 +123,13 @@ export class VectorStore {
       config?: DedupConfigShape;
       /** Optional events.log path for decision monitoring (Sprint 14). */
       eventsPath?: string;
+      /** Repo id for the async cross-repo vector index. Defaults to stateDir. */
+      repoId?: string;
     } = {},
   ) {
     this.embedder = opts.embedder ?? defaultEmbedder();
     this.stateDir = opts.stateDir ?? getStateDir();
+    this.repoId = opts.repoId ?? this.stateDir;
     // Sprint 14: all tier flags/thresholds flow from the single config source
     // (DedupConfig). The legacy opts.dedupSim / opts.l2Enabled remain accepted
     // for backward-compat callers but flags are authoritative via `cfg`.
@@ -472,6 +488,60 @@ export class VectorStore {
     return ranked;
   }
 
+  /**
+   * Slice 2: async cross-repo (or single-repo) recall via the PGlite/HNSW index.
+   *
+   * This is the ONLY async recall surface and is a BONUS path — the synchronous
+   * `search()` above remains the default. `opts.repoId` scopes to one repo; omit
+   * it for cross-repo nearest-neighbor recall (the headline capability the sync
+   * per-session scan cannot provide).
+   *
+   * Best-effort: if the index is disabled/empty/failing, we fall back to the
+   * synchronous per-session `search()` for THIS repo so callers always get a
+   * sensible result. Hydrates each hit's StoredCheckpoint from the authoritative
+   * node:sqlite store (the hit's repoId doubles as that repo's stateDir), then
+   * MMR-dedupes the merged set.
+   */
+  async searchAsync(
+    sessionId: string,
+    query: string,
+    k = 3,
+    opts: { repoId?: string; crossRepo?: boolean } = {},
+  ): Promise<SearchHit[]> {
+    const sid = normalizeSessionId(sessionId);
+    const qv = this.embedder.embed(query);
+    // repoId filter: explicit opts.repoId wins; else this repo unless crossRepo.
+    const repoId = opts.repoId ?? (opts.crossRepo ? undefined : this.repoId);
+    let indexHits: VectorIndexHit[] = [];
+    try {
+      await initVectorIndex();
+      indexHits = await vectorIndexSearch(qv, { k: Math.max(k * 2, k), repoId });
+    } catch {
+      indexHits = [];
+    }
+    if (indexHits.length === 0) {
+      // Index empty/unavailable → synchronous per-session fallback (this repo).
+      return this.search(sid, query, k);
+    }
+    // Hydrate each index hit from the authoritative node:sqlite store. repoId is
+    // that repo's stateDir, so cross-repo hits resolve against their own store.
+    const hydrated: SearchHit[] = [];
+    for (const h of indexHits) {
+      const cp = getCheckpoint(h.sessionId, h.checkpointId, h.repoId);
+      if (cp && cp.dedupStatus !== "removed") {
+        hydrated.push({ checkpoint: cp, score: h.score });
+      }
+    }
+    if (hydrated.length === 0) return this.search(sid, query, k);
+    // MMR-dedupe the merged candidate set for diversity (mirrors sync search).
+    const mmrItems: MmrItem<SearchHit>[] = hydrated.map((h) => ({
+      item: h,
+      vector: h.checkpoint.embedding,
+      relevance: h.score,
+    }));
+    return mmrRerank(mmrItems, k, this.cfg.MMR_LAMBDA);
+  }
+
   /**
    * Serve the RAPTOR tree for a query (Fix D): rehydrate the persisted tree and
    * return its staged-expansion leaf hits as SearchHits. Returns [] when no tree