nlm-memory 0.4.0

package/tests/integration/cite-memo.test.ts

@@ -0,0 +1,87 @@
+import { mkdtempSync, readdirSync, rmSync, writeFileSync } from "node:fs";
+import { tmpdir } from "node:os";
+import { join } from "node:path";
+import { afterEach, beforeEach, describe, expect, it } from "vitest";
+import {
+  clearCited,
+  loadCited,
+  recordCited,
+} from "../../src/core/hook/cite-memo.js";
+
+describe("cite-memo", () => {
+  let tmp: string;
+
+  beforeEach(() => {
+    tmp = mkdtempSync(join(tmpdir(), "nlm-cite-memo-"));
+    process.env["NLM_HOOK_STATE_DIR"] = tmp;
+  });
+
+  afterEach(() => {
+    delete process.env["NLM_HOOK_STATE_DIR"];
+    rmSync(tmp, { recursive: true, force: true });
+  });
+
+  it("loadCited returns empty set when no memo file exists", () => {
+    expect(loadCited("conv-x").size).toBe(0);
+  });
+
+  it("recordCited persists ids; loadCited returns them on next call", () => {
+    recordCited("conv-x", ["cc_a", "cc_b"]);
+    expect(loadCited("conv-x")).toEqual(new Set(["cc_a", "cc_b"]));
+  });
+
+  it("recordCited unions across calls (does not overwrite)", () => {
+    recordCited("conv-x", ["cc_a"]);
+    recordCited("conv-x", ["cc_b", "cc_a"]);
+    expect(loadCited("conv-x")).toEqual(new Set(["cc_a", "cc_b"]));
+  });
+
+  it("recordCited with empty list is a no-op (no file created)", () => {
+    recordCited("conv-empty", []);
+    expect(readdirSync(tmp).filter((f) => f.startsWith("conv-empty"))).toEqual([]);
+  });
+
+  it("clearCited removes the file and returns true; second call returns false", () => {
+    recordCited("conv-x", ["cc_a"]);
+    expect(clearCited("conv-x")).toBe(true);
+    expect(clearCited("conv-x")).toBe(false);
+    expect(loadCited("conv-x").size).toBe(0);
+  });
+
+  it("uses .cited.json filename suffix — parallel to surfaced memo's .json", () => {
+    recordCited("conv-x", ["cc_a"]);
+    const files = readdirSync(tmp);
+    expect(files).toContain("conv-x.cited.json");
+  });
+
+  it("treats corrupt JSON as empty without throwing", () => {
+    writeFileSync(join(tmp, "conv-bad.cited.json"), "not json", "utf8");
+    expect(loadCited("conv-bad").size).toBe(0);
+  });
+
+  it("treats non-array JSON as empty without throwing", () => {
+    writeFileSync(
+      join(tmp, "conv-obj.cited.json"),
+      JSON.stringify({ cc_a: 1 }),
+      "utf8",
+    );
+    expect(loadCited("conv-obj").size).toBe(0);
+  });
+
+  it("filters out non-string entries from the persisted array", () => {
+    writeFileSync(
+      join(tmp, "conv-mixed.cited.json"),
+      JSON.stringify(["cc_a", 42, null, "cc_b"]),
+      "utf8",
+    );
+    expect(loadCited("conv-mixed")).toEqual(new Set(["cc_a", "cc_b"]));
+  });
+
+  it("sanitizes unsafe conversation IDs so the path stays inside the state dir", () => {
+    recordCited("../escape/attempt", ["cc_a"]);
+    const files = readdirSync(tmp);
+    // No file at ../escape/attempt should exist; conversion replaces unsafe chars.
+    expect(files.some((f) => f.endsWith(".cited.json"))).toBe(true);
+    expect(files).not.toContain("..");
+  });
+});

package/tests/integration/db-restore.test.ts

@@ -0,0 +1,153 @@
+/**
+ * Backup + restore integration. Real SQLite stores, VACUUM INTO snapshot,
+ * candidate validation, staged restore, and boot-time promotion.
+ */
+
+import { existsSync, mkdtempSync, readdirSync, rmSync, statSync, writeFileSync } from "node:fs";
+import { tmpdir } from "node:os";
+import { join, resolve } from "node:path";
+import { afterEach, beforeEach, describe, expect, it } from "vitest";
+import type Database from "better-sqlite3";
+import { SqliteSessionStore } from "../../src/core/storage/sqlite-session-store.js";
+import {
+  PENDING_SUFFIX,
+  applyPendingRestore,
+  snapshotScratchPath,
+  stageRestore,
+  validateRestoreCandidate,
+  vacuumSnapshot,
+} from "../../src/core/storage/db-restore.js";
+
+const MIGRATIONS_DIR = resolve(__dirname, "../../migrations");
+
+function insertMarkerSource(db: Database.Database, name: string, runtimeLabel: string): void {
+  db.prepare(
+    "INSERT INTO sources (kind, name, path_or_url, runtime_label, parse_config, enabled) " +
+    "VALUES ('webhook', ?, NULL, ?, '{}', 1)",
+  ).run(name, runtimeLabel);
+}
+
+describe("db-restore", () => {
+  let tmp: string;
+  let dbPath: string;
+
+  beforeEach(() => {
+    tmp = mkdtempSync(join(tmpdir(), "nlm-restore-"));
+    dbPath = join(tmp, "canonical.sqlite");
+  });
+
+  afterEach(() => {
+    rmSync(tmp, { recursive: true, force: true });
+  });
+
+  function freshStore(path: string): SqliteSessionStore {
+    return new SqliteSessionStore({ dbPath: path, migrationsDir: MIGRATIONS_DIR });
+  }
+
+  it("vacuumSnapshot writes a valid standalone copy", () => {
+    const store = freshStore(dbPath);
+    const snap = snapshotScratchPath(dbPath);
+    const bytes = vacuumSnapshot(store.rawDb(), snap);
+    store.close();
+
+    expect(bytes).toBeGreaterThan(0);
+    expect(statSync(snap).size).toBe(bytes);
+    expect(validateRestoreCandidate(snap).ok).toBe(true);
+    rmSync(snap, { force: true });
+  });
+
+  it("validateRestoreCandidate rejects a non-SQLite file", () => {
+    const junk = join(tmp, "junk.sqlite");
+    writeFileSync(junk, "this is not a database");
+    const result = validateRestoreCandidate(junk);
+    expect(result.ok).toBe(false);
+    expect(result.error).toBeTruthy();
+  });
+
+  it("validateRestoreCandidate rejects a SQLite file lacking nlm tables", () => {
+    const bare = join(tmp, "bare.sqlite");
+    const store = freshStore(bare);
+    store.rawDb().prepare("DROP TABLE sessions").run();
+    store.close();
+    const result = validateRestoreCandidate(bare);
+    expect(result.ok).toBe(false);
+    expect(result.error).toMatch(/sessions/);
+  });
+
+  it("validateRestoreCandidate reports session count and schema version", () => {
+    const store = freshStore(dbPath);
+    const snap = snapshotScratchPath(dbPath);
+    vacuumSnapshot(store.rawDb(), snap);
+    store.close();
+
+    const result = validateRestoreCandidate(snap);
+    expect(result.ok).toBe(true);
+    expect(result.sessions).toBe(0);
+    expect(result.schemaVersion).toBeGreaterThanOrEqual(0);
+    rmSync(snap, { force: true });
+  });
+
+  it("stageRestore parks a valid candidate at the pending path", () => {
+    const store = freshStore(dbPath);
+    const snap = snapshotScratchPath(dbPath);
+    vacuumSnapshot(store.rawDb(), snap);
+    store.close();
+
+    const result = stageRestore(dbPath, snap);
+    expect(result.ok).toBe(true);
+    expect(existsSync(dbPath + PENDING_SUFFIX)).toBe(true);
+    expect(existsSync(snap)).toBe(false); // candidate was renamed, not copied
+  });
+
+  it("stageRestore removes the candidate when validation fails", () => {
+    const junk = join(tmp, "junk.sqlite");
+    writeFileSync(junk, "not a database");
+    const result = stageRestore(dbPath, junk);
+    expect(result.ok).toBe(false);
+    expect(existsSync(junk)).toBe(false);
+    expect(existsSync(dbPath + PENDING_SUFFIX)).toBe(false);
+  });
+
+  it("applyPendingRestore is a no-op when nothing is staged", () => {
+    freshStore(dbPath).close();
+    const result = applyPendingRestore(dbPath);
+    expect(result.applied).toBe(false);
+  });
+
+  it("applyPendingRestore promotes the staged DB and archives the current one", () => {
+    // Current DB: one source seeded so we can tell the two stores apart.
+    const current = freshStore(dbPath);
+    insertMarkerSource(current.rawDb(), "marker-current", "current");
+    current.close();
+
+    // Staged DB: built elsewhere, carries a different marker source.
+    const stagedSrc = join(tmp, "staged-src.sqlite");
+    const staged = freshStore(stagedSrc);
+    insertMarkerSource(staged.rawDb(), "marker-staged", "staged");
+    const snap = snapshotScratchPath(dbPath);
+    vacuumSnapshot(staged.rawDb(), snap);
+    staged.close();
+    expect(stageRestore(dbPath, snap).ok).toBe(true);
+
+    const result = applyPendingRestore(dbPath);
+    expect(result.applied).toBe(true);
+    expect(result.archivedTo).toBeTruthy();
+    expect(existsSync(result.archivedTo!)).toBe(true);
+    expect(existsSync(dbPath + PENDING_SUFFIX)).toBe(false);
+
+    // The promoted DB is the staged one.
+    const reopened = freshStore(dbPath);
+    const names = reopened
+      .rawDb()
+      .prepare<[], { name: string }>("SELECT name FROM sources")
+      .all()
+      .map((r) => r.name);
+    reopened.close();
+    expect(names).toContain("marker-staged");
+    expect(names).not.toContain("marker-current");
+
+    // Exactly one pre-restore archive was created.
+    const archives = readdirSync(tmp).filter((f) => f.includes(".pre-restore-"));
+    expect(archives.length).toBe(1);
+  });
+});

package/tests/integration/embed-backfill.test.ts

@@ -0,0 +1,176 @@
+/**
+ * Integration tests for embed-backfill + embed-normalize against a real
+ * SQLite + sqlite-vec store. No network: a deterministic fake LLMClient
+ * stands in for Ollama.
+ */
+
+import { mkdtempSync, rmSync, existsSync } from "node:fs";
+import { tmpdir } from "node:os";
+import { join, resolve } from "node:path";
+import { afterEach, beforeEach, describe, expect, it } from "vitest";
+import Database from "better-sqlite3";
+import * as sqliteVec from "sqlite-vec";
+import { SqliteSessionStore } from "../../src/core/storage/sqlite-session-store.js";
+import { reembedCorpus } from "../../src/core/embedding/embed-backfill.js";
+import { normalizeEmbeddings } from "../../src/core/embedding/embed-normalize.js";
+import type { EmbedResult, LLMClient } from "../../src/ports/llm-client.js";
+import type { Session } from "../../src/shared/types.js";
+import { makeSession } from "../fixtures/sessions.js";
+
+const MIGRATIONS_DIR = resolve(__dirname, "../../migrations");
+
+function unitWithLeading(value: number): Float32Array {
+  const v = new Float32Array(768);
+  v[0] = value;
+  let sum = 0;
+  for (let i = 0; i < v.length; i++) sum += (v[i] ?? 0) ** 2;
+  const n = Math.sqrt(sum) || 1;
+  const out = new Float32Array(768);
+  for (let i = 0; i < v.length; i++) out[i] = (v[i] ?? 0) / n;
+  return out;
+}
+
+class DeterministicEmbedder implements LLMClient {
+  calls = 0;
+  async embed(): Promise<EmbedResult> {
+    this.calls += 1;
+    // Stable, distinct, unit-length vectors per call
+    return { vector: unitWithLeading(this.calls), model: "fake" };
+  }
+  async classify(): Promise<never> {
+    throw new Error("not used");
+  }
+}
+
+const seed: ReadonlyArray<Session> = [
+  makeSession({ id: "s_a", label: "Hono setup", body: "wired Hono routes" }),
+  makeSession({ id: "s_b", label: "pgvector plan", body: "drafted pgvector swap" }),
+  makeSession({ id: "s_c", label: "tx tax county", body: "ingested county directory" }),
+];
+
+describe("reembedCorpus", () => {
+  let tmp: string;
+  let dbPath: string;
+  let statePath: string;
+
+  beforeEach(() => {
+    tmp = mkdtempSync(join(tmpdir(), "nlm-emb-"));
+    dbPath = join(tmp, "canonical.sqlite");
+    statePath = join(tmp, "state.json");
+    const store = new SqliteSessionStore({ dbPath, migrationsDir: MIGRATIONS_DIR });
+    for (const s of seed) {
+      store.insertSessionForTest(s);
+      // seed each with a non-normalized vector so backfill has something to replace
+      store.insertEmbeddingForTest(s.id, new Float32Array(768).fill(0.5));
+    }
+    store.close();
+  });
+
+  afterEach(() => rmSync(tmp, { recursive: true, force: true }));
+
+  it("replaces every embedding and writes a state file", async () => {
+    const embedder = new DeterministicEmbedder();
+    const report = await reembedCorpus({ dbPath, embedder, statePath });
+    expect(report.dbMissing).toBe(false);
+    expect(report.total).toBe(3);
+    expect(report.succeeded).toBe(3);
+    expect(report.failed).toBe(0);
+    expect(report.skippedAlreadyDone).toBe(0);
+    expect(embedder.calls).toBe(3);
+    expect(existsSync(statePath)).toBe(true);
+  });
+
+  it("is resumable — second run skips ids already in state", async () => {
+    const embedder1 = new DeterministicEmbedder();
+    await reembedCorpus({ dbPath, embedder: embedder1, statePath });
+    const embedder2 = new DeterministicEmbedder();
+    const report = await reembedCorpus({ dbPath, embedder: embedder2, statePath });
+    expect(report.skippedAlreadyDone).toBe(3);
+    expect(report.succeeded).toBe(0);
+    expect(embedder2.calls).toBe(0);
+  });
+
+  it("respects --limit", async () => {
+    const embedder = new DeterministicEmbedder();
+    const report = await reembedCorpus({ dbPath, embedder, statePath, limit: 2 });
+    expect(report.total).toBe(2);
+    expect(report.succeeded).toBe(2);
+  });
+});
+
+describe("normalizeEmbeddings", () => {
+  let tmp: string;
+  let dbPath: string;
+
+  beforeEach(() => {
+    tmp = mkdtempSync(join(tmpdir(), "nlm-norm-"));
+    dbPath = join(tmp, "canonical.sqlite");
+    const store = new SqliteSessionStore({ dbPath, migrationsDir: MIGRATIONS_DIR });
+    store.insertSessionForTest(makeSession({ id: "raw" }));
+    store.insertSessionForTest(makeSession({ id: "already" }));
+    store.insertSessionForTest(makeSession({ id: "zero" }));
+    store.close();
+    // embed-normalize operates on the legacy session_embeddings table that
+    // migration 003 still creates (left in place for rollback safety after
+    // the chunk + max-pool migration). Seed it directly via raw SQL — the
+    // session store's helpers now target session_embedding_chunks.
+    const db = new Database(dbPath);
+    sqliteVec.load(db);
+    const ins = db.prepare(
+      "INSERT INTO session_embeddings (session_id, embedding) VALUES (?, ?)",
+    );
+    const toBlob = (v: Float32Array): Buffer =>
+      Buffer.from(v.buffer, v.byteOffset, v.byteLength);
+    // raw: non-unit (||v|| = sqrt(768 * 0.25) ≈ 13.86)
+    ins.run("raw", toBlob(new Float32Array(768).fill(0.5)));
+    // already: unit (one component at 1.0)
+    const unit = new Float32Array(768);
+    unit[0] = 1;
+    ins.run("already", toBlob(unit));
+    // zero: zero vector
+    ins.run("zero", toBlob(new Float32Array(768)));
+    db.close();
+  });
+
+  afterEach(() => rmSync(tmp, { recursive: true, force: true }));
+
+  it("rewrites only the non-unit non-zero row", () => {
+    const report = normalizeEmbeddings({ dbPath });
+    expect(report.total).toBe(3);
+    expect(report.rewritten).toBe(1);
+    expect(report.alreadyNormalized).toBe(1);
+    expect(report.zeroVector).toBe(1);
+  });
+
+  it("dry-run reports the same counts without writing", () => {
+    const beforeDb = new Database(dbPath);
+    sqliteVec.load(beforeDb);
+    const beforeBlob = beforeDb
+      .prepare<[string], { embedding: Buffer }>(
+        "SELECT embedding FROM session_embeddings WHERE session_id = ?",
+      )
+      .get("raw")!.embedding;
+    beforeDb.close();
+
+    const report = normalizeEmbeddings({ dbPath, dryRun: true });
+    expect(report.rewritten).toBe(1);
+    expect(report.dryRun).toBe(true);
+
+    const afterDb = new Database(dbPath);
+    sqliteVec.load(afterDb);
+    const afterBlob = afterDb
+      .prepare<[string], { embedding: Buffer }>(
+        "SELECT embedding FROM session_embeddings WHERE session_id = ?",
+      )
+      .get("raw")!.embedding;
+    afterDb.close();
+    expect(afterBlob.equals(beforeBlob)).toBe(true);
+  });
+
+  it("is idempotent — second run rewrites nothing", () => {
+    normalizeEmbeddings({ dbPath });
+    const report = normalizeEmbeddings({ dbPath });
+    expect(report.rewritten).toBe(0);
+    expect(report.alreadyNormalized).toBe(2); // raw is now unit too
+  });
+});