quak/test/library/mldata.test.ts

/**
 * Tests for the ML-data cache and its derived CLIP index (issue #49).
 *
 * Two layers are exercised:
 *
 *   1. `MLDataStore` on its own: storing one payload file per fileID (present
 *      means complete), building a `clip.f32` + `clip.json` index that reloads
 *      in a single read, rebuilding that index from the payloads when it is
 *      missing or disagrees with the files present, appending as new payloads
 *      arrive, overwriting a refetched file in place, and deciding what to
 *      (re)fetch as `updationTime` advances.
 *
 *   2. `Library` wiring: after each refresh the library fetches ML data through
 *      the metadata pool for every known file not yet cached, is incremental on
 *      later refreshes, and refetches a file whose `updationTime` advanced.
 *
 * Embedding values are chosen to be exactly representable as float32 so the
 * round-trip through `clip.f32` compares equal.
 */

import { describe, it, expect, beforeEach, afterEach, vi } from "vitest";
import { existsSync, mkdtempSync, rmSync, writeFileSync } from "node:fs";
import { tmpdir } from "node:os";
import { join } from "node:path";

import { MLDataStore } from "../../src/library/mldata.js";
import { Library } from "../../src/library/index.js";
import type { CollectionsPage, FilesPage } from "../../src/client.js";
import type { MLData } from "../../src/mldata-fetch.js";
import type { Collection, EnteFile } from "../../src/model/types.js";

// A payload shaped like Ente's: a CLIP embedding plus face data that only the
// on-disk payload carries (never the RAM index).
const payload = (embedding: number[]): MLData => ({
    face: {
        faces: [{ faceID: "f", detection: { box: { x: 0.5 } } }],
    },
    clip: { embedding },
});

describe("MLDataStore", () => {
    let dir: string;

    beforeEach(() => {
        dir = mkdtempSync(join(tmpdir(), "quak-mldata-"));
    });

    afterEach(() => {
        rmSync(dir, { recursive: true, force: true });
    });

    it("stores one payload file per fileID and builds a one-read index", async () => {
        const store = await MLDataStore.open(dir);
        const res = await store.storeFetched(
            new Map([
                [100, payload([0.5, 0.25, 0.75])],
                [200, payload([1, -2, 0.5])],
            ]),
            new Map([
                [100, 10],
                [200, 20],
            ]),
        );
        expect(res).toEqual({ stored: 2, indexed: 2 });

        // One payload file per fileID, and the derived index files.
        expect(existsSync(join(dir, "100.json"))).toBe(true);
        expect(existsSync(join(dir, "200.json"))).toBe(true);
        expect(existsSync(join(dir, "clip.f32"))).toBe(true);
        expect(existsSync(join(dir, "clip.json"))).toBe(true);

        // Reopening loads the index from disk in one read.
        const reopened = await MLDataStore.open(dir);
        const index = reopened.getIndex();
        expect(index.fileIDs).toEqual([100, 200]);
        expect(index.embeddingLength).toBe(3);
        expect([...index.embeddings]).toEqual([0.5, 0.25, 0.75, 1, -2, 0.5]);

        // The full payload (face boxes) is read back from disk on demand.
        const full = await reopened.readPayload(100);
        expect(full?.face).toBeDefined();
        expect(await reopened.readPayload(999)).toBeUndefined();
    });

    it("rebuilds the index from payloads when it is missing", async () => {
        const store = await MLDataStore.open(dir);
        await store.storeFetched(
            new Map([[100, payload([0.5, 0.25, 0.75])]]),
            new Map([[100, 10]]),
        );

        // The derived index is lost but the payloads survive.
        rmSync(join(dir, "clip.f32"));
        rmSync(join(dir, "clip.json"));

        const reopened = await MLDataStore.open(dir);
        const index = reopened.getIndex();
        expect(index.fileIDs).toEqual([100]);
        expect([...index.embeddings]).toEqual([0.5, 0.25, 0.75]);
        expect(existsSync(join(dir, "clip.f32"))).toBe(true);
    });

    it("rebuilds the index when it disagrees with the files present", async () => {
        const store = await MLDataStore.open(dir);
        await store.storeFetched(
            new Map([
                [100, payload([0.5, 0.25, 0.75])],
                [200, payload([1, -2, 0.5])],
            ]),
            new Map([
                [100, 10],
                [200, 20],
            ]),
        );

        // A payload disappears out from under the index (leaving it referencing
        // a file no longer present); the index must be rebuilt from what is
        // actually on disk.
        rmSync(join(dir, "200.json"));

        const reopened = await MLDataStore.open(dir);
        expect(reopened.getIndex().fileIDs).toEqual([100]);
    });

    it("rebuilds the index when a payload on disk is missing from it", async () => {
        const store = await MLDataStore.open(dir);
        await store.storeFetched(
            new Map([[100, payload([0.5, 0.25, 0.75])]]),
            new Map([[100, 10]]),
        );

        // A crash between storeFetched renaming a payload into place and
        // rewriting the index leaves the payload complete on disk but absent
        // from clip.json. Write a second payload directly to reproduce that
        // torn state without touching the index.
        writeFileSync(
            join(dir, "200.json"),
            JSON.stringify(payload([1, -2, 0.5])),
        );

        // Reopening self-heals with no manual delete: the index is rebuilt from
        // the payloads to include the orphaned embedding.
        const reopened = await MLDataStore.open(dir);
        const index = reopened.getIndex();
        expect(index.fileIDs).toEqual([100, 200]);
        expect([...index.embeddings]).toEqual([0.5, 0.25, 0.75, 1, -2, 0.5]);
    });

    it("appends new payloads and overwrites a refetched file in place", async () => {
        const store = await MLDataStore.open(dir);
        await store.storeFetched(
            new Map([[100, payload([0.5, 0.25, 0.75])]]),
            new Map([[100, 10]]),
        );
        // A later batch adds a new file: appended after the first.
        await store.storeFetched(
            new Map([[200, payload([1, -2, 0.5])]]),
            new Map([[200, 20]]),
        );
        // Refetching 100 (its embedding changed) updates it in place, not a
        // duplicate row.
        await store.storeFetched(
            new Map([[100, payload([9, 9, 9])]]),
            new Map([[100, 30]]),
        );

        const index = store.getIndex();
        expect(index.fileIDs).toEqual([100, 200]);
        expect([...index.embeddings]).toEqual([9, 9, 9, 1, -2, 0.5]);
    });

    it("keeps a payload without a CLIP embedding out of the index", async () => {
        const store = await MLDataStore.open(dir);
        const res = await store.storeFetched(
            new Map<number, MLData>([[100, { face: { faces: [] } }]]),
            new Map([[100, 10]]),
        );
        expect(res.stored).toBe(1);
        expect(res.indexed).toBe(0);
        // The payload is still cached (present means complete).
        expect(existsSync(join(dir, "100.json"))).toBe(true);
        expect(store.getIndex().fileIDs).toEqual([]);
    });

    it("fetches only what is missing or has a newer updationTime", async () => {
        const store = await MLDataStore.open(dir);
        await store.storeFetched(
            new Map([[100, payload([0.5, 0.25, 0.75])]]),
            new Map([[100, 10]]),
        );

        // 100 is cached and current; 200 has never been fetched.
        expect(
            store.neededFor([
                { id: 100, updationTime: 10 },
                { id: 200, updationTime: 5 },
            ]),
        ).toEqual([200]);

        // 100's updationTime advanced past what it was fetched at: refetch.
        expect(store.neededFor([{ id: 100, updationTime: 15 }])).toEqual([100]);

        // Nothing advanced: nothing to fetch.
        expect(store.neededFor([{ id: 100, updationTime: 10 }])).toEqual([]);
    });

    it("survives a corrupt index without losing the payloads", async () => {
        const store = await MLDataStore.open(dir);
        await store.storeFetched(
            new Map([[100, payload([0.5, 0.25, 0.75])]]),
            new Map([[100, 10]]),
        );
        writeFileSync(join(dir, "clip.json"), "not json");

        const reopened = await MLDataStore.open(dir);
        expect(reopened.getIndex().fileIDs).toEqual([100]);
    });
});

// --- Library wiring ---------------------------------------------------------

const USER_ID = 42;
const FAST_INTERVAL = 0.02;

const collection = (id: number, updationTime: number): Collection => ({
    id,
    ownerID: USER_ID,
    key: new Uint8Array([id & 0xff]),
    name: `album-${id}`,
    type: "album",
    updationTime,
    isShared: false,
});

const file = (
    id: number,
    collectionID: number,
    updationTime: number,
): EnteFile => ({
    id,
    collectionID,
    ownerID: USER_ID,
    key: new Uint8Array([id & 0xff]),
    metadata: {
        title: `file-${id}.jpg`,
        fileType: "image",
        creationTime: updationTime,
        modificationTime: updationTime,
    },
    file: { decryptionHeader: "aGVhZGVy" },
    thumbnail: { decryptionHeader: "dGh1bWI=" },
    updationTime,
});

// A mock client that serves scripted collection/file pages and per-file ML
// payloads, recording every ML fetch request so incremental behaviour is
// provable.
class MLMockClient {
    userID = USER_ID;
    collectionsQueue: CollectionsPage[] = [];
    filesByCollection = new Map<number, FilesPage[]>();
    mlByFile = new Map<number, MLData>();
    mlFetchCalls: number[][] = [];

    whoami(): { email: string; userID: number } {
        return { email: "user@example.com", userID: this.userID };
    }

    async collectionsSince(args: {
        sinceTime: number;
    }): Promise<CollectionsPage> {
        return (
            this.collectionsQueue.shift() ?? {
                collections: [],
                deleted: [],
                cursor: args.sinceTime,
            }
        );
    }

    async filesSince(args: {
        collectionID: number;
        collectionKey: Uint8Array;
        sinceTime: number;
    }): Promise<FilesPage> {
        const queue = this.filesByCollection.get(args.collectionID);
        return (
            queue?.shift() ?? {
                files: [],
                deleted: [],
                cursor: args.sinceTime,
            }
        );
    }

    async fetchMLData(args: {
        fileIDs: number[];
        fileKeys: Map<number, Uint8Array>;
    }): Promise<Map<number, MLData>> {
        this.mlFetchCalls.push([...args.fileIDs]);
        const result = new Map<number, MLData>();
        for (const id of args.fileIDs) {
            const p = this.mlByFile.get(id);
            if (p) result.set(id, p);
        }
        return result;
    }

    filesFor(collectionID: number, ...pages: FilesPage[]): void {
        this.filesByCollection.set(collectionID, pages);
    }
}

describe("Library ML-data fetch on refresh", () => {
    let dir: string;
    let cacheDirectory: string;

    beforeEach(() => {
        dir = mkdtempSync(join(tmpdir(), "quak-lib-mldata-"));
        cacheDirectory = join(dir, "cache");
    });

    afterEach(() => {
        rmSync(dir, { recursive: true, force: true });
    });

    it("fetches, stores and indexes ML data for known files, then is incremental", async () => {
        const client = new MLMockClient();
        client.collectionsQueue.push({
            collections: [collection(1, 100)],
            deleted: [],
            cursor: 100,
        });
        client.filesFor(1, {
            files: [file(1001, 1, 90), file(1002, 1, 95)],
            deleted: [],
            cursor: 95,
        });
        client.mlByFile.set(1001, payload([0.5, 0.25, 0.75]));
        client.mlByFile.set(1002, payload([1, -2, 0.5]));

        const lib = await Library.open({
            client,
            cacheDirectory,
            refreshIntervalSeconds: FAST_INTERVAL,
        });
        try {
            // Wait on `lastMLFetchAt`, set only once the pass has persisted the
            // index and payloads — not on the in-RAM counts, which advance
            // before `storeFetched` writes to disk, so the reopen below reads
            // the committed index rather than racing the write.
            await vi.waitFor(
                () => {
                    expect(lib.status().lastMLFetchAt).toBeGreaterThan(0);
                    expect(lib.status().mlIndexed).toBe(2);
                    expect(lib.status().mlStored).toBe(2);
                },
                { timeout: 2000, interval: 5 },
            );

            // Both files were fetched, in one batch.
            expect(client.mlFetchCalls.flat().sort((a, b) => a - b)).toEqual([
                1001, 1002,
            ]);
            const callsAfterFirst = client.mlFetchCalls.length;

            // The index is on disk and reloads to the same shape.
            const reopened = await MLDataStore.open(
                join(cacheDirectory, "mldata"),
            );
            expect(reopened.getIndex().fileIDs).toEqual([1001, 1002]);

            // Later refreshes with nothing new must not refetch.
            await new Promise((r) => setTimeout(r, FAST_INTERVAL * 1000 * 5));
            expect(client.mlFetchCalls.length).toBe(callsAfterFirst);
        } finally {
            lib.close();
        }
    });

    it("refetches a file whose updationTime advanced", async () => {
        const client = new MLMockClient();
        client.collectionsQueue.push({
            collections: [collection(1, 100)],
            deleted: [],
            cursor: 100,
        });
        client.filesFor(1, {
            files: [file(1001, 1, 90)],
            deleted: [],
            cursor: 90,
        });
        client.mlByFile.set(1001, payload([0.5, 0.25, 0.75]));

        const lib = await Library.open({
            client,
            cacheDirectory,
            refreshIntervalSeconds: FAST_INTERVAL,
        });
        try {
            // Wait on `lastMLFetchAt`, set only after the first pass has
            // persisted, not on `mlIndexed`, which is bumped in RAM before the
            // write lands.
            await vi.waitFor(
                () => expect(lib.status().lastMLFetchAt).toBeGreaterThan(0),
                { timeout: 2000, interval: 5 },
            );
            const callsBefore = client.mlFetchCalls.length;

            // The file changes on the server (updationTime advances) with a new
            // embedding; the next refresh must refetch it.
            client.mlByFile.set(1001, payload([9, 9, 9]));
            client.collectionsQueue.push({
                collections: [collection(1, 200)],
                deleted: [],
                cursor: 200,
            });
            client.filesFor(1, {
                files: [file(1001, 1, 190)],
                deleted: [],
                cursor: 190,
            });

            // Poll the persisted index itself, not the fetch-call log: a call
            // is recorded the instant the mock is entered, but `storeFetched`
            // rewrites `clip.f32` only after it resolves, so an earlier reopen
            // would read the pre-refetch vector. Reopening reads only committed
            // (atomically renamed) files, so this sees the new embedding once —
            // and only once — the store has written it.
            await vi.waitFor(
                async () => {
                    const reopened = await MLDataStore.open(
                        join(cacheDirectory, "mldata"),
                    );
                    expect([...reopened.getIndex().embeddings]).toEqual([
                        9, 9, 9,
                    ]);
                },
                { timeout: 2000, interval: 20 },
            );

            // The refetch really went back to the server for 1001.
            expect(client.mlFetchCalls.length).toBeGreaterThan(callsBefore);
            expect(client.mlFetchCalls.flat()).toContain(1001);
        } finally {
            lib.close();
        }
    });
});