fix(lbug): stream CodeEmbedding cache under the connection lock

Spill vectors once the in-memory limit is crossed and fail the load
instead of adopting an empty snapshot, so incremental analyze cannot
OOM or quietly drop the restore cache (#3306).

Co-authored-by: Cursor <cursoragent@cursor.com>

gitnexus/src/core/lbug/lbug-adapter.ts

@@ -34,7 +34,16 @@ import type { GraphEmitManifest } from './graph-emit-sink.js';
 import type { PdgEmitManifest } from './pdg-emit-sink.js';
 import { PDG_EDGE_TYPES } from './pdg-emit-sink.js';
 import { getNodeLabel as deriveNodeLabel, type WriteStreamFactory } from './rel-pair-routing.js';
-import { EMBEDDABLE_LABELS, type CachedEmbedding } from '../embeddings/types.js';
+import { EMBEDDABLE_LABELS } from '../embeddings/types.js';
+import {
+  abortCachedEmbeddingsBuilder,
+  createCachedEmbeddingsBuilder,
+  emptyCachedEmbeddingsSnapshot,
+  finalizeCachedEmbeddingsSnapshot,
+  ingestCachedEmbeddingRow,
+  type CachedEmbeddingsSnapshot,
+  type LoadCachedEmbeddingsOptions,
+} from '../embeddings/embedding-restore-spill.js';
 import {
   extensionManager,
   getFtsCapability,
@@ -2069,28 +2078,31 @@ export const getLbugStats = async (): Promise<{
 
 /**
  * Load cached embeddings from LadybugDB before a rebuild.
- * Returns all embedding vectors so they can be re-inserted after the graph is reloaded,
- * avoiding expensive re-embedding of unchanged nodes.
+ *
+ * Streams `CodeEmbedding` rows with `hasNext`/`getNext` under `withConnLock`
+ * (#2264, #3306). Vectors are spilled to a temp Float32 file once the table
+ * exceeds the in-memory row limit so incremental analyze cannot OOM the V8
+ * heap by materializing every `number[]` up front. Small tables still return
+ * in-RAM `embeddings` for existing callers/tests.
  *
  * Detects old schema (no chunkIndex column) and returns empty cache to trigger rebuild.
  */
-export const loadCachedEmbeddings = async (): Promise<{
-  embeddingNodeIds: Set<string>;
-  embeddings: CachedEmbedding[];
-}> => {
+export const loadCachedEmbeddings = async (
+  options?: LoadCachedEmbeddingsOptions,
+): Promise<CachedEmbeddingsSnapshot> => {
   const c = conn;
   if (!c) {
-    return { embeddingNodeIds: new Set(), embeddings: [] };
+    return emptyCachedEmbeddingsSnapshot();
   }
 
   // The whole read runs inside the connection lock (#2264 review P2). It's safe
   // today only by call-ordering (loadCachedEmbeddings runs before the WAL driver
   // starts), but the lock makes it robust to future reordering — a concurrent
   // CHECKPOINT on the singleton connection is the documented corruption trigger.
-  // Leaf read: no nested withConnLock-wrapped helpers inside.
+  // Leaf read: no nested withConnLock-wrapped helpers inside. Do NOT call
+  // `streamQuery` here — that path is unlocked and would race a CHECKPOINT.
   return withConnLock(async () => {
-    const embeddingNodeIds = new Set<string>();
-    const embeddings: CachedEmbedding[] = [];
+    const builder = createCachedEmbeddingsBuilder(options);
     try {
       // Schema migration detection: query with new columns to verify schema version.
       // Old schema only had (nodeId, embedding); new schema adds (id, chunkIndex, startLine, endLine, contentHash).
@@ -2104,51 +2116,46 @@ export const loadCachedEmbeddings = async (): Promise<{
         );
         await readQueryRows(check);
       } catch {
-        return { embeddingNodeIds: new Set(), embeddings: [] };
+        abortCachedEmbeddingsBuilder(builder);
+        return emptyCachedEmbeddingsSnapshot();
       }
 
-      // Try to read contentHash alongside chunk columns
-      let rows: any;
+      let queryResult: lbug.QueryResult | lbug.QueryResult[] | undefined;
       let hasContentHash = true;
       try {
-        rows = await c.query(
-          `MATCH (e:${EMBEDDING_TABLE_NAME}) RETURN e.nodeId AS nodeId, e.chunkIndex AS chunkIndex, e.startLine AS startLine, e.endLine AS endLine, e.embedding AS embedding, e.contentHash AS contentHash`,
-        );
-      } catch (err: any) {
-        // Fallback for legacy DBs without contentHash column
-        const msg = err?.message ?? '';
-        if (isMissingColumnOrTableError(msg)) {
-          hasContentHash = false;
-          rows = await c.query(
-            `MATCH (e:${EMBEDDING_TABLE_NAME}) RETURN e.nodeId AS nodeId, e.chunkIndex AS chunkIndex, e.startLine AS startLine, e.endLine AS endLine, e.embedding AS embedding`,
+        try {
+          queryResult = await c.query(
+            `MATCH (e:${EMBEDDING_TABLE_NAME}) RETURN e.nodeId AS nodeId, e.chunkIndex AS chunkIndex, e.startLine AS startLine, e.endLine AS endLine, e.embedding AS embedding, e.contentHash AS contentHash`,
           );
-        } else {
-          throw err;
+        } catch (err: any) {
+          // Fallback for legacy DBs without contentHash column
+          const msg = err?.message ?? '';
+          if (isMissingColumnOrTableError(msg)) {
+            hasContentHash = false;
+            queryResult = await c.query(
+              `MATCH (e:${EMBEDDING_TABLE_NAME}) RETURN e.nodeId AS nodeId, e.chunkIndex AS chunkIndex, e.startLine AS startLine, e.endLine AS endLine, e.embedding AS embedding`,
+            );
+          } else {
+            throw err;
+          }
         }
-      }
-      for (const row of await readQueryRows(rows)) {
-        const nodeId = String(row.nodeId ?? row[0] ?? '');
-        if (!nodeId) continue;
-        embeddingNodeIds.add(nodeId);
-        const embedding = row.embedding ?? row[4];
-        if (embedding) {
-          embeddings.push({
-            nodeId,
-            chunkIndex: Number(row.chunkIndex ?? row[1] ?? 0),
-            startLine: Number(row.startLine ?? row[2] ?? 0),
-            endLine: Number(row.endLine ?? row[3] ?? 0),
-            embedding: Array.isArray(embedding)
-              ? embedding.map(Number)
-              : Array.from(embedding as any).map(Number),
-            contentHash: hasContentHash ? (row.contentHash ?? row[5] ?? undefined) : undefined,
-          });
+        const results = Array.isArray(queryResult) ? queryResult : [queryResult];
+        const result = results[0];
+        while (await result.hasNext()) {
+          const row = await result.getNext();
+          ingestCachedEmbeddingRow(builder, row, hasContentHash);
         }
+        return finalizeCachedEmbeddingsSnapshot(builder);
+      } catch (err) {
+        abortCachedEmbeddingsBuilder(builder);
+        throw err;
+      } finally {
+        if (queryResult) await closeQueryResults(queryResult);
       }
-    } catch {
-      /* embedding table may not exist */
+    } catch (err) {
+      abortCachedEmbeddingsBuilder(builder);
+      throw err;
     }
-
-    return { embeddingNodeIds, embeddings };
   });
 };
 

gitnexus/test/integration/load-cached-embeddings-spill.test.ts

@@ -0,0 +1,106 @@
+/**
+ * Real-DB coverage for #3306: loadCachedEmbeddings must stream CodeEmbedding
+ * rows instead of getAll()+map(Number) of the whole table, and must be able
+ * to spill vectors so incremental analyze does not keep every embedding in
+ * the V8 heap.
+ */
+import fs from 'node:fs';
+import path from 'node:path';
+import { afterEach, describe, expect, it } from 'vitest';
+import { createTempDir, type TestDBHandle } from '../helpers/test-db.js';
+import { EMBEDDING_DIMS } from '../../src/core/lbug/schema.js';
+import { batchInsertEmbeddings } from '../../src/core/embeddings/embedding-pipeline.js';
+import {
+  disposeEmbeddingSpill,
+  materializeCachedEmbeddings,
+} from '../../src/core/embeddings/embedding-restore-spill.js';
+
+describe('loadCachedEmbeddings streaming (#3306)', () => {
+  let tmp: TestDBHandle | undefined;
+
+  afterEach(async () => {
+    const adapter = await import('../../src/core/lbug/lbug-adapter.js');
+    try {
+      await adapter.closeLbug();
+    } catch {
+      /* already closed */
+    }
+    await tmp?.cleanup();
+    tmp = undefined;
+  });
+
+  async function seedDb(rowCount: number) {
+    tmp = await createTempDir('gitnexus-lbug-');
+    const dbPath = path.join(tmp.dbPath, 'lbug');
+    const adapter = await import('../../src/core/lbug/lbug-adapter.js');
+    await adapter.initLbug(dbPath);
+    const rows = Array.from({ length: rowCount }, (_, i) => ({
+      nodeId: `Function:src/f${i}.ts:fn${i}:1`,
+      chunkIndex: 0,
+      startLine: 1,
+      endLine: 3,
+      embedding: Array.from({ length: EMBEDDING_DIMS }, (__, d) => (d === 0 ? i + 1 : 0)),
+      contentHash: `hash-${i}`,
+    }));
+    await batchInsertEmbeddings(adapter.executeWithReusedStatement, rows);
+    return { adapter, rows };
+  }
+
+  it('materializes a small table in RAM (skip-fts / mock-compatible shape)', async () => {
+    const { adapter, rows } = await seedDb(3);
+    const cached = await adapter.loadCachedEmbeddings();
+    expect(cached.spill).toBeUndefined();
+    expect(cached.embeddings).toHaveLength(3);
+    expect(cached.rows).toHaveLength(3);
+    expect(cached.embeddingNodeIds.size).toBe(3);
+    expect(cached.embeddings.map((e) => e.nodeId).sort()).toEqual(rows.map((r) => r.nodeId).sort());
+    expect(cached.embeddings.find((e) => e.nodeId === rows[1]!.nodeId)?.embedding[0]).toBe(2);
+  });
+
+  it('streams into a spill file when the in-memory limit is 0 and restores a subset', async () => {
+    const { adapter, rows } = await seedDb(12);
+    const cached = await adapter.loadCachedEmbeddings({ inMemoryRowLimit: 0 });
+    try {
+      expect(cached.embeddings).toEqual([]);
+      expect(cached.spill?.rowCount).toBe(12);
+      expect(cached.rows).toHaveLength(12);
+      const subset = materializeCachedEmbeddings(
+        cached,
+        cached.rows.filter((_, i) => i % 5 === 0),
+      );
+      expect(subset).toHaveLength(3);
+      expect(subset[0]?.embedding[0]).toBe(1);
+      expect(subset[1]?.embedding[0]).toBe(6);
+      expect(subset[2]?.embedding[0]).toBe(11);
+      expect(subset[0]?.contentHash).toBe(rows[0]!.contentHash);
+    } finally {
+      disposeEmbeddingSpill(cached.spill);
+    }
+  });
+
+  it('flips from RAM to spill once a non-zero in-memory limit is crossed', async () => {
+    const { adapter } = await seedDb(8);
+    const cached = await adapter.loadCachedEmbeddings({ inMemoryRowLimit: 4 });
+    try {
+      expect(cached.embeddings).toEqual([]);
+      expect(cached.spill?.rowCount).toBe(8);
+      expect(cached.rows).toHaveLength(8);
+      expect(fs.statSync(cached.spill!.path).size).toBe(12 + 8 * EMBEDDING_DIMS * 4);
+      const subset = materializeCachedEmbeddings(cached, cached.rows.slice(2, 4));
+      expect(subset).toHaveLength(2);
+      expect(subset[0]?.embedding[0]).toBe(3);
+      expect(subset[1]?.embedding[0]).toBe(4);
+    } finally {
+      disposeEmbeddingSpill(cached.spill);
+    }
+  });
+
+  it('surfaces a spill write failure instead of adopting an empty snapshot', async () => {
+    const { adapter } = await seedDb(3);
+    const spillDir = path.join(tmp!.dbPath, 'not-a-directory');
+    fs.writeFileSync(spillDir, 'x');
+    await expect(adapter.loadCachedEmbeddings({ inMemoryRowLimit: 0, spillDir })).rejects.toThrow(
+      /ENOTDIR|not a directory|ENOSPC|EACCES/i,
+    );
+  });
+});

gitnexus/vitest.config.ts

@@ -100,6 +100,7 @@ export default defineConfig({
             'test/integration/analyze-wal-checkpoint-failure.test.ts',
             'test/integration/lbug-non-ascii-path.test.ts',
             'test/integration/lbug-conn-serialization.test.ts',
+            'test/integration/load-cached-embeddings-spill.test.ts',
             'test/integration/group/manifest-resolve-symbol-2325.test.ts',
             'test/integration/group/manifest-synthetic-impact-lbug.test.ts',
             'test/integration/group/http-route-resolve-symbol.test.ts',
@@ -179,6 +180,7 @@ export default defineConfig({
             'test/integration/analyze-wal-checkpoint-failure.test.ts',
             'test/integration/lbug-non-ascii-path.test.ts',
             'test/integration/lbug-conn-serialization.test.ts',
+            'test/integration/load-cached-embeddings-spill.test.ts',
             'test/integration/group/manifest-resolve-symbol-2325.test.ts',
             'test/integration/group/manifest-synthetic-impact-lbug.test.ts',
             'test/integration/group/http-route-resolve-symbol.test.ts',