open-brainy/tests/unit/hnsw/update-item-atomic.test.ts

/**
 * @module tests/unit/hnsw/update-item-atomic
 * @description Guard for the atomic vector-index update: a row must NEVER be
 * absent from vector search during an update. The historical update path
 * staged a remove followed by an add as two separately-awaited transaction
 * operations — between them the row was in NEITHER index (dark to semantic
 * recall while perfectly visible to metadata reads; observed as seconds-long
 * flicker in a production deployment). The structural cure verified here:
 *
 *   1. `JsHnswVectorIndex.updateItem` — same vector (element-wise) is a pure
 *      no-op (the production flicker shape: a type-only update re-indexing an
 *      UNCHANGED vector); a changed vector swaps in place, the node never
 *      leaving the map (white-box probe at the first internal step after the
 *      synchronous swap), including when the node IS the entry point.
 *   2. `ReplaceInVectorIndexOperation` — one transaction leg that prefers the
 *      provider's in-place `updateItem`, with a remove+add-ADJACENT fallback
 *      for providers that have not shipped it; rollback restores the declared
 *      before-vector on both branches.
 *   3. The brain's update path — with the JS index carrying `updateItem`,
 *      `removeItem` is never called during `brain.update()`, for the
 *      type-only shape AND for a genuine vector change.
 */
import { describe, it, expect, vi } from 'vitest'
import { JsHnswVectorIndex } from '../../../src/hnsw/hnswIndex.js'
import { ReplaceInVectorIndexOperation } from '../../../src/transaction/operations/IndexOperations.js'
import type { VectorIndexProvider } from '../../../src/plugin.js'
import type { Vector, VectorDocument } from '../../../src/coreTypes.js'
import { euclideanDistance } from '../../../src/utils/index.js'
import { MemoryStorage } from '../../../src/storage/adapters/memoryStorage.js'
import { Brainy } from '../../../src/brainy'
import { createAddParams, createTestConfig } from '../../helpers/test-factory'

const DIM = 8

function seededRand(seed: number): () => number {
  let s = seed >>> 0
  return () => {
    s = (s + 0x6d2b79f5) | 0
    let t = Math.imul(s ^ (s >>> 15), 1 | s)
    t = (t + Math.imul(t ^ (t >>> 7), 61 | t)) ^ t
    return ((t ^ (t >>> 14)) >>> 0) / 4294967296
  }
}

/** A deterministic vector pointing in a pseudo-random direction (well-connected graph). */
function vec(idx: number): number[] {
  const rand = seededRand(idx + 1)
  return Array.from({ length: DIM }, () => rand() * 2 - 1)
}

type Noun = { id: string; vector: number[]; connections: Map<number, Set<string>>; level: number }

function nounsOf(index: JsHnswVectorIndex): Map<string, Noun> {
  return (index as unknown as { nouns: Map<string, Noun> }).nouns
}

/** Flatten a reverse index to sorted `target|level|source` triples. */
function triplesFromIncoming(inc: Map<string, Map<number, Set<string>>>): string[] {
  const out: string[] = []
  for (const [target, byLevel] of inc) {
    for (const [level, sources] of byLevel) {
      for (const source of sources) out.push(`${target}|${level}|${source}`)
    }
  }
  return out.sort()
}

/** Derive the ground-truth reverse index directly from the live forward adjacency. */
function triplesFromAdjacency(nouns: Map<string, Noun>): string[] {
  const out: string[] = []
  for (const [nodeId, node] of nouns) {
    for (const [level, targets] of node.connections) {
      for (const target of targets) out.push(`${target}|${level}|${nodeId}`)
    }
  }
  return out.sort()
}

function assertReverseIndexConsistent(index: JsHnswVectorIndex): void {
  const live = (
    index as unknown as { ensureIncoming: () => Map<string, Map<number, Set<string>>> }
  ).ensureIncoming()
  expect(triplesFromIncoming(live)).toEqual(triplesFromAdjacency(nounsOf(index)))
}

function assertNoSelfLoops(index: JsHnswVectorIndex, id: string): void {
  const node = nounsOf(index).get(id)!
  for (const [level, targets] of node.connections) {
    expect(targets.has(id), `self-loop at level ${level}`).toBe(false)
  }
}

function makeIndex(M = 16): JsHnswVectorIndex {
  return new JsHnswVectorIndex(
    { M, efConstruction: 200, efSearch: 64, ml: 16 },
    euclideanDistance,
    { useParallelization: false, storage: new MemoryStorage() }
  )
}

async function fillIndex(index: JsHnswVectorIndex, count: number): Promise<void> {
  for (let i = 0; i < count; i++) {
    await index.addItem({ id: `n-${i}`, vector: vec(i) })
  }
}

describe('JsHnswVectorIndex.updateItem — atomic in-place vector update', () => {
  it('same vector (element-wise equal, fresh array) is a pure no-op: no remove, no relink, still searchable', async () => {
    const index = makeIndex()
    await fillIndex(index, 30)

    const target = 'n-7'
    const sameVector = [...vec(7)] // fresh array, identical elements

    const before = await index.search(vec(7), 1)
    expect(before[0][0]).toBe(target)

    const removeSpy = vi.spyOn(index, 'removeItem')
    const nodeBefore = nounsOf(index).get(target)!
    const connectionsBefore = nodeBefore.connections // reference — a relink replaces it

    await index.updateItem({ id: target, vector: sameVector })

    expect(removeSpy).not.toHaveBeenCalled()
    expect(index.size()).toBe(30)
    // No relink happened: the connections map is the SAME object, untouched.
    expect(nounsOf(index).get(target)!.connections).toBe(connectionsBefore)

    const after = await index.search(vec(7), 1)
    expect(after[0][0]).toBe(target)
    expect(after[0][1]).toBeCloseTo(0, 10)

    removeSpy.mockRestore()
  })

  it('changed vector: node never leaves the map (probe fires after the synchronous swap), removeItem never called, findable by the NEW vector', async () => {
    const index = makeIndex()
    await fillIndex(index, 40)

    const target = 'n-5'
    const newVector = vec(500)

    // White-box probe: ensureIncoming is the FIRST internal step of the unlink
    // walk, i.e. the first thing updateItem does after the synchronous vector
    // swap. At that instant the node must (a) still be in the map and (b)
    // already carry the NEW vector — the visibility-atomic ordering.
    const inner = index as unknown as {
      nouns: Map<string, Noun>
      ensureIncoming: () => Map<string, Map<number, Set<string>>>
    }
    const origEnsure = inner.ensureIncoming.bind(index)
    let probed = false
    let presentDuring = false
    let swappedFirst = false
    ;(index as any).ensureIncoming = function () {
      if (!probed) {
        probed = true
        presentDuring = inner.nouns.has(target)
        swappedFirst = inner.nouns.get(target)?.vector === newVector
      }
      return origEnsure()
    }

    const removeSpy = vi.spyOn(index, 'removeItem')
    await index.updateItem({ id: target, vector: newVector })
    delete (index as any).ensureIncoming // restore the prototype method

    expect(probed).toBe(true)
    expect(presentDuring).toBe(true)
    expect(swappedFirst).toBe(true)
    expect(removeSpy).not.toHaveBeenCalled()
    expect(index.size()).toBe(40)
    expect(nounsOf(index).has(target)).toBe(true)

    // Findable by search with the NEW vector, at distance ~0.
    const got = await index.search(newVector, 1)
    expect(got[0][0]).toBe(target)
    expect(got[0][1]).toBeCloseTo(0, 10)

    // The relink left the graph bookkeeping exactly consistent.
    assertNoSelfLoops(index, target)
    assertReverseIndexConsistent(index)

    removeSpy.mockRestore()
  })

  it('keeps the node at its existing level (never releveled by an update)', async () => {
    const index = makeIndex()
    await fillIndex(index, 30)

    const target = 'n-3'
    const levelBefore = nounsOf(index).get(target)!.level

    await index.updateItem({ id: target, vector: vec(600) })

    expect(nounsOf(index).get(target)!.level).toBe(levelBefore)
    expect(index.getMaxLevel()).toBeGreaterThanOrEqual(levelBefore)
  })

  it('updating the ENTRY POINT in place keeps it valid — entry id and maxLevel unchanged, graph never stranded', async () => {
    const index = makeIndex()
    await fillIndex(index, 40)

    const entryId = index.getEntryPointId()!
    const maxLevelBefore = index.getMaxLevel()
    const newVector = vec(700)

    await index.updateItem({ id: entryId, vector: newVector })

    // Entry-point bookkeeping must not regress.
    expect(index.getEntryPointId()).toBe(entryId)
    expect(index.getMaxLevel()).toBe(maxLevelBefore)
    expect(index.size()).toBe(40)

    // The entry point itself is findable by its new vector...
    const gotEntry = await index.search(newVector, 1)
    expect(gotEntry[0][0]).toBe(entryId)

    // ...and the REST of the graph is still reachable through it (a stranded,
    // edgeless entry point would make every other node invisible).
    const otherId = [...nounsOf(index).keys()].find((id) => id !== entryId)!
    const otherIdx = Number(otherId.slice(2))
    const gotOther = await index.search(vec(otherIdx), 1)
    expect(gotOther[0][0]).toBe(otherId)

    assertNoSelfLoops(index, entryId)
    assertReverseIndexConsistent(index)
  })

  it('absent id delegates to addItem (plain insert)', async () => {
    const index = makeIndex()
    await fillIndex(index, 10)

    await index.updateItem({ id: 'fresh', vector: vec(900) })

    expect(index.size()).toBe(11)
    const got = await index.search(vec(900), 1)
    expect(got[0][0]).toBe('fresh')
  })
})

describe('ReplaceInVectorIndexOperation — one atomic transaction leg', () => {
  it('uses the provider updateItem path and rolls back to the old vector in place', async () => {
    const index = makeIndex()
    await fillIndex(index, 30)

    const target = 'n-9'
    const oldVector = vec(9)
    const newVector = vec(800)

    const removeSpy = vi.spyOn(index, 'removeItem')
    const op = new ReplaceInVectorIndexOperation(index, target, oldVector, newVector)
    expect(op.name).toBe('ReplaceInVectorIndex(hnsw-js)')

    const rollback = await op.execute()
    expect(removeSpy).not.toHaveBeenCalled()
    expect((await index.search(newVector, 1))[0][0]).toBe(target)

    await rollback()
    expect(removeSpy).not.toHaveBeenCalled()
    expect(index.size()).toBe(30)

    // Old vector restored, element-wise, and searchable again.
    const restored = nounsOf(index).get(target)!.vector
    expect(restored.length).toBe(oldVector.length)
    for (let i = 0; i < oldVector.length; i++) {
      expect(restored[i]).toBe(oldVector[i])
    }
    const back = await index.search(oldVector, 1)
    expect(back[0][0]).toBe(target)
    expect(back[0][1]).toBeCloseTo(0, 10)

    removeSpy.mockRestore()
  })

  it('falls back to remove+add ADJACENT within the single op for a provider without updateItem, and rolls back the same way', async () => {
    // A provider that has not shipped updateItem — the temporary seam: the
    // pair stays adjacent inside ONE op (no other transaction operation can
    // interleave), until the provider ships its own in-place updateItem.
    const calls: string[] = []
    const store = new Map<string, Vector>()
    const legacyProvider = {
      name: 'legacy-native',
      addItem: async (item: VectorDocument) => {
        calls.push(`add:${item.id}`)
        store.set(item.id, item.vector)
        return item.id
      },
      removeItem: async (id: string) => {
        calls.push(`remove:${id}`)
        return store.delete(id)
      },
      search: async () => [],
      size: () => store.size,
      clear: () => store.clear(),
      rebuild: async () => {},
      flush: async () => 0,
      getPersistMode: () => 'immediate' as const
    } as unknown as VectorIndexProvider

    store.set('x', [1, 0])
    const op = new ReplaceInVectorIndexOperation(legacyProvider, 'x', [1, 0], [0, 1])

    const rollback = await op.execute()
    expect(calls).toEqual(['remove:x', 'add:x'])
    expect(store.get('x')).toEqual([0, 1])

    await rollback()
    expect(calls).toEqual(['remove:x', 'add:x', 'remove:x', 'add:x'])
    expect(store.get('x')).toEqual([1, 0])
  })
})

describe('brain.update() — the update path stages ONE atomic vector-index leg', () => {
  it('a type-only update (unchanged vector — the production flicker shape) never calls removeItem on the vector index', async () => {
    const brain = new Brainy(createTestConfig())
    await brain.init()
    try {
      const id = await brain.add(
        createAddParams({ data: 'atomic flicker guard entity', type: 'thing' })
      )

      const index = (brain as unknown as { index: JsHnswVectorIndex }).index
      const removeSpy = vi.spyOn(index, 'removeItem')
      const sizeBefore = index.size()

      await brain.update({ id, type: 'document' })

      expect(removeSpy).not.toHaveBeenCalled()
      expect(index.size()).toBe(sizeBefore)

      const updated = await brain.get(id)
      expect(updated).not.toBeNull()
      expect(updated!.type).toBe('document')

      removeSpy.mockRestore()
    } finally {
      await brain.close()
    }
  })

  it('a genuine vector change on update also never calls removeItem (in-place replace)', async () => {
    const brain = new Brainy(createTestConfig())
    await brain.init()
    try {
      const id = await brain.add(
        createAddParams({ data: 'vector change stays visible', type: 'thing' })
      )
      const existing = await brain.get(id, { includeVectors: true })
      // Same dimensionality, guaranteed-different content.
      const changed = existing!.vector.map((x: number, i: number) => (i === 0 ? x + 0.25 : x))

      const index = (brain as unknown as { index: JsHnswVectorIndex }).index
      const removeSpy = vi.spyOn(index, 'removeItem')

      await brain.update({ id, vector: changed })

      expect(removeSpy).not.toHaveBeenCalled()
      expect(nounsOf(index).has(id)).toBe(true)

      removeSpy.mockRestore()
    } finally {
      await brain.close()
    }
  })
})