open-brainy/tests/unit/hnsw/hnsw-watermark.test.ts

/**
 * @module tests/unit/hnsw/hnsw-watermark
 * @description Watermark-stamp pins for the JS HNSW vector projection.
 *
 * THE LAW under test: the persisted HNSW artifact (per-node records + the
 * entryPoint/maxLevel system record) carries a stamp asserting "this state
 * reflects every committed generation ≤ W and nothing above W" — written
 * AFTER every byte it certifies is durable — and rebuild() computes the
 * three-way verdict: stamped==committed → 'adopt' · stamped<committed →
 * 'catchup' (gap reported) · stamped>committed OR unstamped → 'rescan',
 * LOUDLY. Vector-bearing stamps carry the model identity this module can
 * honestly assert: dimensions only (no embedding-model id is reachable from
 * the index module).
 *
 * The verdict is COMPUTED AND EXPOSED only — no rebuild trigger changed.
 */
import { describe, it, expect, vi, afterEach } from 'vitest'
import { v4 as uuidv4 } from 'uuid'
import { JsHnswVectorIndex, HNSW_INDEX_STAMP_KEY } from '../../../src/hnsw/hnswIndex.js'
import { euclideanDistance } from '../../../src/utils/index.js'
import { MemoryStorage } from '../../../src/storage/adapters/memoryStorage.js'
import { prodLog } from '../../../src/utils/logger.js'

const DIM = 8

function randomVector(dim: number): number[] {
  return Array.from({ length: dim }, () => Math.random() * 2 - 1)
}

async function makeStorage(committed: number | null): Promise<MemoryStorage> {
  const storage = new MemoryStorage()
  await storage.init()
  if (committed !== null) {
    vi.spyOn(storage, 'committedGeneration').mockReturnValue(committed)
  }
  return storage
}

function setCommitted(storage: MemoryStorage, committed: number): void {
  vi.spyOn(storage, 'committedGeneration').mockReturnValue(committed)
}

function makeIndex(storage: MemoryStorage): JsHnswVectorIndex {
  return new JsHnswVectorIndex(
    { M: 4, efConstruction: 50, efSearch: 20 },
    euclideanDistance,
    { useParallelization: false, storage, persistMode: 'deferred' }
  )
}

/** Session 1: insert nodes, optionally stamp, flush — the durable artifact. */
async function writeArtifact(storage: MemoryStorage, stamp: number | null): Promise<void> {
  const index = makeIndex(storage)
  for (let i = 0; i < 3; i++) {
    await index.addItem({ id: uuidv4(), vector: randomVector(DIM) })
  }
  if (stamp !== null) index.stampWatermark(stamp)
  await index.flush()
}

/** Session 2: reopen on the same storage via the load path (rebuild). */
async function reopen(storage: MemoryStorage): Promise<JsHnswVectorIndex> {
  const index = makeIndex(storage)
  await index.rebuild()
  return index
}

afterEach(() => {
  vi.restoreAllMocks()
})

describe('JS HNSW index — watermark stamp + three-way load verdict', () => {
  it("save-with-stamp then reopen at the same committed generation → 'adopt'", async () => {
    const storage = await makeStorage(5)
    await writeArtifact(storage, 5)

    const index = await reopen(storage)
    expect(index.watermarkVerdict()).toBe('adopt')
    expect(index.watermark()).toBe(5)
    expect(index.watermarkGap()).toBeNull()
  })

  it("stamp BEHIND the committed generation → 'catchup' with the exact gap reported", async () => {
    const storage = await makeStorage(5)
    await writeArtifact(storage, 5)

    setCommitted(storage, 9)

    const index = await reopen(storage)
    expect(index.watermarkVerdict()).toBe('catchup')
    expect(index.watermark()).toBe(5)
    expect(index.watermarkGap()).toEqual({ from: 5, to: 9 })
  })

  it("stamp ABOVE the committed generation → 'rescan', said out loud", async () => {
    const storage = await makeStorage(9)
    await writeArtifact(storage, 9)

    setCommitted(storage, 4)

    const warnSpy = vi.spyOn(prodLog, 'warn')
    const index = await reopen(storage)
    expect(index.watermarkVerdict()).toBe('rescan')
    expect(index.watermarkGap()).toBeNull()
    const said = warnSpy.mock.calls.map(c => String(c[0])).join('\n')
    expect(said).toContain('RESCAN')
    expect(said).toContain('ABOVE')
  })

  it("legacy unstamped artifact on a stamped store → 'rescan', LOUD — never a silent adopt", async () => {
    const storage = await makeStorage(3)
    await writeArtifact(storage, null) // pre-stamp index: data flushed, no stamp

    expect(await storage.getMetadata(HNSW_INDEX_STAMP_KEY)).toBeNull()

    const warnSpy = vi.spyOn(prodLog, 'warn')
    const index = await reopen(storage)
    expect(index.watermarkVerdict()).toBe('rescan')
    expect(index.watermark()).toBeNull()
    const said = warnSpy.mock.calls.map(c => String(c[0])).join('\n')
    expect(said).toContain('RESCAN')
    expect(said).toContain('unstamped')
  })

  it("a store with no committed-generation capability keeps pre-stamp behavior → 'adopt'", async () => {
    const storage = await makeStorage(null)
    await writeArtifact(storage, null)

    const index = await reopen(storage)
    expect(index.watermarkVerdict()).toBe('adopt')
    expect(index.watermark()).toBeNull()
  })

  it('STAMP-AFTER-DATA: the stamp lands after every node record and the system record', async () => {
    const storage = await makeStorage(2)
    const index = makeIndex(storage)
    for (let i = 0; i < 3; i++) {
      await index.addItem({ id: uuidv4(), vector: randomVector(DIM) })
    }

    // One shared op log across all three write surfaces pins global order.
    const ops: string[] = []
    const origNode = storage.saveVectorIndexData.bind(storage)
    vi.spyOn(storage, 'saveVectorIndexData').mockImplementation(async (id, data) => {
      ops.push(`node:${id}`)
      return origNode(id, data)
    })
    const origSystem = storage.saveHNSWSystem.bind(storage)
    vi.spyOn(storage, 'saveHNSWSystem').mockImplementation(async data => {
      ops.push('system')
      return origSystem(data)
    })
    const origMeta = storage.saveMetadata.bind(storage)
    vi.spyOn(storage, 'saveMetadata').mockImplementation(async (id, metadata) => {
      ops.push(`meta:${id}`)
      return origMeta(id, metadata)
    })

    index.stampWatermark(2)
    await index.flush()

    const stampAt = ops.indexOf(`meta:${HNSW_INDEX_STAMP_KEY}`)
    expect(stampAt, 'stamp record was written').toBeGreaterThanOrEqual(0)
    expect(stampAt, 'stamp is the FINAL write of the flush').toBe(ops.length - 1)
    expect(ops.filter(o => o.startsWith('node:')).length).toBeGreaterThan(0)
    expect(ops.indexOf('system')).toBeLessThan(stampAt)
  })

  it('the stamp record carries {watermark, formatVersion, stampedAt} + modelIdentity (dims only)', async () => {
    const storage = await makeStorage(7)
    await writeArtifact(storage, 7)

    const record = (await storage.getMetadata(HNSW_INDEX_STAMP_KEY)) as {
      watermark: number
      formatVersion: number
      stampedAt: number
      modelIdentity: { embedModelId?: string; dimensions: number | null }
    }
    expect(record.watermark).toBe(7)
    expect(record.formatVersion).toBe(1)
    expect(typeof record.stampedAt).toBe('number')
    // The JS index never sees the embedder — dimensions are the only vector-
    // space identity it can honestly assert.
    expect(record.modelIdentity).toEqual({ dimensions: DIM })
  })

  it('a pending stamp still lands when nothing is dirty (already-durable bytes, stamp-after-data trivially holds)', async () => {
    const storage = await makeStorage(4)
    const index = makeIndex(storage)
    await index.addItem({ id: uuidv4(), vector: randomVector(DIM) })
    await index.flush() // data durable, no stamp yet

    index.stampWatermark(4)
    await index.flush() // nothing dirty — the stamp must still be written

    const record = (await storage.getMetadata(HNSW_INDEX_STAMP_KEY)) as { watermark: number }
    expect(record?.watermark).toBe(4)
    expect(index.watermark()).toBe(4)
  })
})