qmd/test/embed-partial-chunks.test.ts

/**
 * 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);
  });
});