citation-evidence/src/engine/persistence.test.ts

import { beforeEach, describe, expect, it, vi } from "vitest";
import type { Document, DocumentRepresentation } from "@shared/document";
import type { DocumentId, RepresentationId } from "@shared/ids";
import {
  attachPersister,
  captureSnapshot,
  createEngine,
  restoreFromStorage,
  restoreSnapshot,
  sanitizeDocumentForPersistence,
  type Engine,
  type EngineEvent,
  type EngineSnapshot,
} from "./index";

function fakeDocAndRep(suffix: string): {
  document: Document;
  representation: DocumentRepresentation;
} {
  const docId = `doc_${suffix}` as DocumentId;
  const repId = `rep_${suffix}` as RepresentationId;
  return {
    document: {
      id: docId,
      mediaType: "application/pdf",
      title: `Doc ${suffix}`,
      createdAt: "2026-05-25T00:00:00.000Z",
      updatedAt: "2026-05-25T00:00:00.000Z",
    },
    representation: {
      id: repId,
      documentId: docId,
      representationType: "pdf-text",
      contentHash: `hash-${suffix}`,
      canonicalText: "The quick brown fox.",
      pageMap: [{ page: 1, width: 100, height: 100 }],
      offsetMap: [{ page: 1, globalStart: 0, globalEnd: 20, pageLength: 20 }],
      generatedAt: "2026-05-25T00:00:00.000Z",
    },
  };
}

function memoryStorage(): Pick<Storage, "getItem" | "setItem" | "removeItem"> {
  const map = new Map<string, string>();
  return {
    getItem: (k) => map.get(k) ?? null,
    setItem: (k, v) => void map.set(k, v),
    removeItem: (k) => void map.delete(k),
  };
}

function seed(engine: Engine, suffix: string) {
  const { document, representation } = fakeDocAndRep(suffix);
  engine.documents.register({ document, representation });
  const ann = engine.annotations.create({
    documentId: document.id,
    representationId: representation.id,
    selectors: [{ type: "TextQuoteSelector", exact: "brown fox" }],
    quote: "brown fox",
  });
  const item = engine.evidence.create({
    annotationIds: [ann.id],
    commentary: `commentary-${suffix}`,
  });
  return { document, representation, ann, item };
}

describe("captureSnapshot + restoreSnapshot", () => {
  it("round-trips documents, representations, annotations and evidence items", () => {
    const src = createEngine();
    const a = seed(src, "a");
    const b = seed(src, "b");
    const snap = captureSnapshot(src);
    expect(snap.documents).toHaveLength(2);
    expect(snap.representations).toHaveLength(2);
    expect(snap.annotations).toHaveLength(2);
    expect(snap.evidenceItems).toHaveLength(2);

    const dst = createEngine();
    restoreSnapshot(dst, snap);
    expect(dst.documents.get(a.document.id)?.title).toBe("Doc a");
    expect(dst.documents.get(b.document.id)?.title).toBe("Doc b");
    expect(dst.annotations.get(a.ann.id)?.quote).toBe("brown fox");
    expect(dst.evidence.get(a.item.id)?.commentary).toBe("commentary-a");
  });

  it("restoreSnapshot does NOT emit *Created events (events would loop the persister)", () => {
    const src = createEngine();
    seed(src, "x");
    const snap = captureSnapshot(src);

    const dst = createEngine();
    const seen: EngineEvent["type"][] = [];
    dst.bus.onAny((e) => seen.push(e.type));
    restoreSnapshot(dst, snap);
    expect(seen).toEqual([]);
  });

  it("rejects a snapshot with a mismatching version", () => {
    const dst = createEngine();
    expect(() =>
      restoreSnapshot(dst, {
        version: 999,
        documents: [],
        representations: [],
        annotations: [],
        evidenceItems: [],
      } as EngineSnapshot),
    ).toThrow(/version/);
  });
});

describe("attachPersister", () => {
  let storage: ReturnType<typeof memoryStorage>;
  let engine: Engine;
  const KEY = "ce-test-snap";

  beforeEach(() => {
    storage = memoryStorage();
    engine = createEngine();
  });

  it("writes a snapshot to storage on every mutating event", () => {
    const off = attachPersister(engine, { key: KEY, storage });
    expect(storage.getItem(KEY)).toBeNull();
    seed(engine, "z");
    const raw = storage.getItem(KEY);
    expect(raw).not.toBeNull();
    const snap = JSON.parse(raw!) as EngineSnapshot;
    expect(snap.documents).toHaveLength(1);
    expect(snap.evidenceItems).toHaveLength(1);
    off();
  });

  it("stops writing after the unsubscribe is called", () => {
    const off = attachPersister(engine, { key: KEY, storage });
    seed(engine, "q");
    const after = storage.getItem(KEY);
    off();
    seed(engine, "r");
    expect(storage.getItem(KEY)).toBe(after);
  });

  it("survives a JSON.stringify failure without throwing into the caller", () => {
    const warn = vi.spyOn(console, "warn").mockImplementation(() => {});
    const failing = { ...memoryStorage(), setItem: () => { throw new Error("boom"); } };
    attachPersister(engine, { key: KEY, storage: failing });
    expect(() => seed(engine, "k")).not.toThrow();
    expect(warn).toHaveBeenCalled();
    warn.mockRestore();
  });
});

describe("restoreFromStorage", () => {
  it("returns {restored: false} when the key is empty", () => {
    const storage = memoryStorage();
    const engine = createEngine();
    const result = restoreFromStorage(engine, { key: "missing", storage });
    expect(result.restored).toBe(false);
  });

  it("hydrates the engine when storage holds a valid snapshot", () => {
    const src = createEngine();
    seed(src, "rs");
    const storage = memoryStorage();
    storage.setItem("snap", JSON.stringify(captureSnapshot(src)));

    const dst = createEngine();
    const result = restoreFromStorage(dst, { key: "snap", storage });
    expect(result.restored).toBe(true);
    expect(dst.documents.list()).toHaveLength(1);
  });

  it("strips blob: URIs from persisted documents", () => {
    const engine = createEngine();
    const docId = "doc_blob" as DocumentId;
    engine.documents.register({
      document: {
        id: docId,
        mediaType: "application/pdf",
        title: "upload.pdf",
        uri: "blob:http://localhost/dead",
        createdAt: "2026-06-07T00:00:00.000Z",
        updatedAt: "2026-06-07T00:00:00.000Z",
      },
      representation: fakeDocAndRep("blob").representation,
    });
    const snap = captureSnapshot(engine);
    expect(snap.documents[0]?.uri).toBeUndefined();
    expect(sanitizeDocumentForPersistence({
      id: docId,
      mediaType: "application/pdf",
      uri: "blob:x",
      createdAt: "x",
      updatedAt: "x",
    }).uri).toBeUndefined();
  });

  it("ignores malformed JSON without throwing", () => {
    const warn = vi.spyOn(console, "warn").mockImplementation(() => {});
    const storage = memoryStorage();
    storage.setItem("snap", "not-json");
    const engine = createEngine();
    const result = restoreFromStorage(engine, { key: "snap", storage });
    expect(result.restored).toBe(false);
    expect(warn).toHaveBeenCalled();
    warn.mockRestore();
  });
});