/** * embed-partial-chunks.test.ts — a document that embedded chunk 0 and lost its * later chunks must be re-selected on the next pass (i-xeekgx6h). * * Before the fix, "does this document need embedding?" was asked as "is there a * vector at seq 0?", while chunk inserts are per-chunk best-effort. Compose the * two and a document whose first chunk landed during an upstream outage is * complete FOREVER: the 30-minute embed cron never selects it again, and its * missing chunks are unreachable. Nothing errors — semantic search just answers * from a document with holes in it. * * The negative control matters as much as the fix here: a fully embedded corpus * must NOT be re-selected, or the repair costs a full re-embed every pass. * * In-memory SQLite + stub provider — no node-llama-cpp. */ import { describe, test, expect, beforeEach, afterEach } from "vitest"; import { mkdtempSync, rmSync } from "node:fs"; import { tmpdir } from "node:os"; import { join } from "node:path"; import { createStore, generateEmbeddings, getHashesNeedingEmbedding, type Store, } from "../src/store.js"; import type { EmbeddingProvider, ProviderEmbedding, ProviderHealth, } from "../src/embedding/provider.js"; /** * Stub provider that can drop chunks the way an unreachable upstream does. * * `failBatchTextsFrom` is an index into the texts handed to `embedBatch` across * the whole run: everything at or after it comes back `null`, which is exactly * the contract `generateEmbeddings` treats as "this chunk has no embedding" and * skips inserting. `embed()` (the single-text dimension probe) always succeeds — * a failing probe aborts the run before any chunk is written, which is a * different, already-safe path. */ class DroppingProvider implements EmbeddingProvider { readonly kind = "openai" as const; readonly modelId = "stub-embed"; readonly dim = 4; batchTextsSeen = 0; constructor(private readonly failBatchTextsFrom = Number.POSITIVE_INFINITY) {} getModelId(): string { return this.modelId; } getDimensions(): number | undefined { return this.dim; } async healthcheck(): Promise { return { ok: true, model: this.modelId, dimensions: this.dim }; } async embed(text: string): Promise { return { embedding: this.fakeEmbed(text), model: this.modelId }; } async embedBatch(texts: string[]): Promise<(ProviderEmbedding | null)[]> { return texts.map((t) => { const index = this.batchTextsSeen++; if (index >= this.failBatchTextsFrom) return null; return { embedding: this.fakeEmbed(t), model: this.modelId }; }); } async dispose(): Promise {} private fakeEmbed(text: string): number[] { return Array.from({ length: this.dim }, (_, i) => (text.length + i) * 0.01); } } let workDir: string; let store: Store; /** Long enough to chunk into several pieces (~900 tokens/chunk, ~3 chars/token). */ function multiChunkBody(marker: string): string { const paragraph = `${marker} paragraph about indexing, retrieval and embeddings. `.repeat(40); return Array.from({ length: 8 }, (_, i) => `## Section ${i}\n\n${paragraph}`).join("\n\n"); } function vectorCount(hash: string): number { return (store.db.prepare(`SELECT COUNT(*) as n FROM content_vectors WHERE hash = ?`).get(hash) as { n: number }).n; } function recordedChunkCount(hash: string): number | undefined { const row = store.db.prepare(`SELECT chunks FROM document_chunk_counts WHERE hash = ?`).get(hash) as { chunks: number } | undefined; return row?.chunks; } beforeEach(() => { workDir = mkdtempSync(join(tmpdir(), "qmd-partial-chunks-test-")); process.env.INDEX_PATH = join(workDir, "index.sqlite"); store = createStore(process.env.INDEX_PATH); const now = "2026-08-15T00:00:00Z"; store.db .prepare(`INSERT INTO content (hash, doc, created_at) VALUES (?, ?, ?)`) .run("hashPartial", multiChunkBody("alpha"), now); store.db .prepare(`INSERT INTO documents (hash, collection, path, title, created_at, modified_at, active) VALUES (?, ?, ?, ?, ?, ?, ?)`) .run("hashPartial", "alpha", "a.md", "a.md", now, now, 1); }); afterEach(() => { try { store.close(); } catch { /* ignore */ } delete process.env.INDEX_PATH; rmSync(workDir, { recursive: true, force: true }); }); describe("partially embedded documents are re-selected (i-xeekgx6h)", () => { test("chunk 0 embedded + later chunks dropped => document is STILL pending", async () => { // Upstream dies after the first chunk of the document. const dropping = new DroppingProvider(1); await generateEmbeddings(store, { embedProvider: dropping }); const expected = recordedChunkCount("hashPartial"); expect(expected).toBeDefined(); expect(expected!).toBeGreaterThan(1); // the body really did chunk into several pieces expect(vectorCount("hashPartial")).toBe(1); // ... and only chunk 0 survived // THE assertion. Before the fix this was 0: chunk 0 exists, so the document // read as complete and its missing chunks were unreachable forever. expect(getHashesNeedingEmbedding(store.db)).toBe(1); }); test("a later pass with a healthy upstream fills the missing chunks", async () => { await generateEmbeddings(store, { embedProvider: new DroppingProvider(1) }); const expected = recordedChunkCount("hashPartial")!; expect(vectorCount("hashPartial")).toBe(1); const healed = await generateEmbeddings(store, { embedProvider: new DroppingProvider() }); expect(healed.docsProcessed).toBe(1); expect(vectorCount("hashPartial")).toBe(expected); expect(getHashesNeedingEmbedding(store.db)).toBe(0); }); test("CONTROL: a fully embedded document is NOT re-selected", async () => { const first = await generateEmbeddings(store, { embedProvider: new DroppingProvider() }); expect(first.docsProcessed).toBe(1); expect(getHashesNeedingEmbedding(store.db)).toBe(0); // The whole point of the control: without it, "re-select partial documents" // could be satisfied by re-selecting EVERYTHING, i.e. re-embedding the corpus // on every cron tick. const second = await generateEmbeddings(store, { embedProvider: new DroppingProvider() }); expect(second.docsProcessed).toBe(0); expect(second.chunksEmbedded).toBe(0); }); test("CONTROL: a legacy document with no recorded chunk count is left alone", async () => { await generateEmbeddings(store, { embedProvider: new DroppingProvider() }); // Simulate a document embedded before document_chunk_counts existed, whose // later chunks are missing: no expectation is recorded, so it cannot be // detected as partial. It is deliberately NOT re-selected — a whole-corpus // re-embed on upgrade would cost more than the holes it heals; `--force` // (or the next re-chunk) is the repair path. store.db.prepare(`DELETE FROM document_chunk_counts WHERE hash = ?`).run("hashPartial"); store.db.prepare(`DELETE FROM content_vectors WHERE hash = ? AND seq > 0`).run("hashPartial"); expect(vectorCount("hashPartial")).toBe(1); expect(recordedChunkCount("hashPartial")).toBeUndefined(); expect(getHashesNeedingEmbedding(store.db)).toBe(0); }); test("the collection-filtered count agrees with the unfiltered one", async () => { // The two queries drifting apart is what produced this bug: they each // carried their own copy of the pending predicate. await generateEmbeddings(store, { embedProvider: new DroppingProvider(1) }); expect(getHashesNeedingEmbedding(store.db)).toBe(1); expect(getHashesNeedingEmbedding(store.db, "alpha")).toBe(1); expect(getHashesNeedingEmbedding(store.db, "nonexistent")).toBe(0); }); });