|
|
@@ -0,0 +1,176 @@
|
|
|
+/**
|
|
|
+ * 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<ProviderHealth> {
|
|
|
+ return { ok: true, model: this.modelId, dimensions: this.dim };
|
|
|
+ }
|
|
|
+ async embed(text: string): Promise<ProviderEmbedding | null> {
|
|
|
+ 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<void> {}
|
|
|
+ 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);
|
|
|
+ });
|
|
|
+});
|
Claude