From 87e9975721f7c85f55a165a577e1fe48f19c5bae Mon Sep 17 00:00:00 2001 From: codeasier Date: Mon, 28 Sep 2026 20:20:57 +0800 Subject: [PATCH] perf(index): store vectors in buckets and rebuild BM25 from content Port of upstream #1415 + #1416 (cfc1934 + bcf4f0d, combined end state; bcf4f0d's parent IS cfc1934) to this line, replacing the ledger/generation persistence design (this line's #1115, submitted upstream as #1256, not merged; #1416 credits its analysis) with upstream's bucket design: - saves run at most once per AGENTMEMORY_INDEX_SAVE_INTERVAL_MS (default 10 min); explicit saves still immediate; single-flight with one queued follow-up - vectors live in fixed bucket scopes, one key per vector; a save writes only the changed keys (AGENTMEMORY_VECTOR_BUCKET_SIZE=500, fill-and-roll) - BM25 is no longer persisted: rebuilt at boot from stored content; the boot log reports docs + duration; the sharded snapshot and the ledger/generation reclaim machinery are removed with the AGENTMEMORY_INDEX_RECLAIM_BOOT_MAX_SHARDS knob - whole-store vector backfill requires AGENTMEMORY_VECTOR_BACKFILL=all, capped per boot by AGENTMEMORY_VECTOR_BACKFILL_MAX=500, progress saved - /agentmemory/status reports per-leg save state, buckets, pending backfill; new problem codes bm25-rebuild-incomplete, index-save-failing, index-save-stale, index-vector-backfill-pending, index-vector-count-shortfall - first boot migrates existing stores: old sharded/monolithic snapshot read once, written into buckets, then old vector and BM25 shards and manifests removed; unreadable snapshots left untouched Line-required adaptations (engine 0.11.2): - ISdk instead of IIIClient (search.ts import + registerSearchFunction) - search-index.ts gains the upstream dedup-on-add line from #1396 (d547883) as a prerequisite - index.ts ctor new IndexPersistence(kv, vectorIndex); fork's reclaimOnLoad/bootReclaimMaxDeletes options removed along with the knob and its evict.test.ts coverage - api.ts keeps this line's consolidation-status import beside the expanded search.js import Verification: npm test 1790 passed / 1 skipped (baseline 1804; net -14 = ledger tests removed with the machinery, upstream rewrite plus new throttle/keyword-rebuild tests in); npm run build; npm run skills:gen + skills:check (17 skills, REFERENCE.md 43 -> 47 vars); npx tsc --noEmit one error FEWER than baseline (upstream fixes this line's TS18046 in isValidShardDescriptor). --- .env.example | 5 +- CHANGELOG.md | 4 + plugin/skills/agentmemory-config/REFERENCE.md | 7 +- src/config.ts | 35 +- src/functions/search.ts | 184 +- src/functions/status.ts | 79 + src/index.ts | 118 +- src/state/index-persistence.ts | 1017 +++++------ src/state/search-index.ts | 1 + src/state/vector-index.ts | 55 +- src/triggers/api.ts | 12 +- test/evict.test.ts | 18 - test/index-keyword-rebuild.test.ts | 240 +++ test/index-persistence-throttle.test.ts | 248 +++ test/index-persistence.test.ts | 1547 ++++------------- test/search.test.ts | 17 + test/status.test.ts | 24 + 17 files changed, 1651 insertions(+), 1960 deletions(-) create mode 100644 test/index-keyword-rebuild.test.ts create mode 100644 test/index-persistence-throttle.test.ts diff --git a/.env.example b/.env.example index c3e3cdfe5a..ce8d203fb7 100644 --- a/.env.example +++ b/.env.example @@ -132,9 +132,12 @@ # AGENTMEMORY_OBSERVE_PROMPT_CHARS=8000 # Server-side prompt bound; 0 disables # AGENTMEMORY_OBSERVE_RAW_CHARS=16000 # Server-side raw envelope bound; 0 disables # AGENTMEMORY_GRAPH_INCREMENTAL_EXTRACT=true # Extract only the observation delta per turn; set false for whole-session LLM entity merging -# AGENTMEMORY_INDEX_RECLAIM_BOOT_MAX_SHARDS=200 # Maximum orphan index shards deleted per background boot pass # AGENTMEMORY_DROP_STALE_INDEX=true # Drop on-disk BM25 / vector index on startup if dim guard fires (#248). Recovery toggle for stuck-state debugging. # AGENTMEMORY_IMAGE_EMBEDDINGS=true # Enable image embeddings when an image provider is present (experimental). +# AGENTMEMORY_INDEX_SAVE_INTERVAL_MS=600000 # Minimum time between search index saves; shutdown and deletes still save immediately. +# AGENTMEMORY_VECTOR_BUCKET_SIZE=500 # Vectors held per storage bucket before a new bucket opens; a save only rewrites buckets it changed. +# AGENTMEMORY_VECTOR_BACKFILL_MAX=500 # Cap on vectors embedded per boot to backfill ones missing since the last save. +# AGENTMEMORY_VECTOR_BACKFILL=all # Required to backfill an entire store's vectors at once (no prior vector save, or a stale legacy manifest); unset means that case waits for an explicit opt-in. # ----------------------------------------------------------------------------- # 6. CLI / runtime knobs diff --git a/CHANGELOG.md b/CHANGELOG.md index 0a55a5b548..39c98db4ee 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,10 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), ## [Unreleased] +### Changed + +- **Vector indexes are stored in buckets, BM25 is rebuilt from content at boot, and saves are throttled (aligned with upstream #1415, #1416).** Every observation, compression and memory save scheduled a full index save 5 s later, and each save re-serialized the whole BM25 and vector index and rewrote every shard of a new generation — on a busy store a full index rewrite almost continuously. Saves now run at most once per `AGENTMEMORY_INDEX_SAVE_INTERVAL_MS` (default 10 minutes, measured from the start of the previous save); explicit saves (deletes, shutdown) still run immediately, only one save runs at a time, and saves requested while one is running collapse into a single follow-up. Vectors live in fixed bucket scopes with one key per vector (`AGENTMEMORY_VECTOR_BUCKET_SIZE`, default 500, fill-and-roll), so a save writes or deletes only the changed keys and nothing builds a string of the whole index. BM25 is no longer persisted: it is rebuilt at boot from stored observations and latest memories — the boot log reports the document count and duration, and on the pinned 0.11.2 engine the rebuild adds one state list per session — which also drops entries for content that no longer exists. This removes this line's ledger/generation reclaim machinery (`AGENTMEMORY_INDEX_RECLAIM_BOOT_MAX_SHARDS` is gone): its threat model is met by construction, since unchanged buckets are never rewritten and BM25 has no shards to strand. Existing stores migrate on first boot: the old sharded or monolithic snapshot is read once, written into buckets, and only then are the old vector and BM25 shards and manifests removed; if the old snapshot cannot be read it is left untouched and the next save retries. This line's two legacy `:gc` ledger keys are not part of that cleanup and can be deleted by hand once migration completes. Documents added after the last vector save that have no vector are embedded in the background; a whole-store embed never starts silently — it requires `AGENTMEMORY_VECTOR_BACKFILL=all` and is capped per boot by `AGENTMEMORY_VECTOR_BACKFILL_MAX` (default 500), with progress saved every few batches. `/agentmemory/status` now shows each leg's save state, the bucket count and unsaved vector changes, reports an incomplete boot rebuild as an error (`bm25-rebuild-incomplete`), a failing vector save as an error, a pending backfill as info, and a vector count shortfall against the last save's record as a warning with a bounded backfill. + ### Fixed - **Observation capture no longer grows the persisted live-stream store.** Raw and compressed observation events use `stream::send` instead of `stream::set`; the viewer still receives live updates and REST remains the history source, but `data/stream_store` no longer accumulates a write-only second and third copy of every observation. diff --git a/plugin/skills/agentmemory-config/REFERENCE.md b/plugin/skills/agentmemory-config/REFERENCE.md index 491108a5dc..dee70cfd6f 100644 --- a/plugin/skills/agentmemory-config/REFERENCE.md +++ b/plugin/skills/agentmemory-config/REFERENCE.md @@ -3,7 +3,7 @@ Generated by scanning `src/` for `AGENTMEMORY_*` usage. Do not edit the block below by hand; run `npm run skills:gen` after adding or removing a variable. Internal markers ending in two underscores are excluded. -Configuration is read from the environment and from `~/.agentmemory/.env` (no `export` prefix). 52 recognized variables: +Configuration is read from the environment and from `~/.agentmemory/.env` (no `export` prefix). 55 recognized variables: - `AGENTMEMORY_AGENT_SCOPE` - `AGENTMEMORY_ALLOW_AGENT_SDK` @@ -33,7 +33,7 @@ Configuration is read from the environment and from `~/.agentmemory/.env` (no `e - `AGENTMEMORY_III_VERSION` - `AGENTMEMORY_IMAGE_EMBEDDINGS` - `AGENTMEMORY_IMAGE_STORE_MAX_BYTES` -- `AGENTMEMORY_INDEX_RECLAIM_BOOT_MAX_SHARDS` +- `AGENTMEMORY_INDEX_SAVE_INTERVAL_MS` - `AGENTMEMORY_INJECT_CONTEXT` - `AGENTMEMORY_LLM_NOTHINK` - `AGENTMEMORY_LLM_TIMEOUT_MS` @@ -54,6 +54,9 @@ Configuration is read from the environment and from `~/.agentmemory/.env` (no `e - `AGENTMEMORY_TOOLS` - `AGENTMEMORY_URL` - `AGENTMEMORY_USE_DOCKER` +- `AGENTMEMORY_VECTOR_BACKFILL` +- `AGENTMEMORY_VECTOR_BACKFILL_MAX` +- `AGENTMEMORY_VECTOR_BUCKET_SIZE` - `AGENTMEMORY_VERBOSE` - `AGENTMEMORY_VIEWER_HOST` - `AGENTMEMORY_VIEWER_URL` diff --git a/src/config.ts b/src/config.ts index b47948eb74..a8e0538008 100644 --- a/src/config.ts +++ b/src/config.ts @@ -474,6 +474,34 @@ export function getConsolidationCooldownMs(): number { return raw >= 0 ? raw : CONSOLIDATION_COOLDOWN_DEFAULT_MS; } +export const INDEX_SAVE_INTERVAL_DEFAULT_MS = 600_000; + +export function getIndexSaveIntervalMs(): number { + const raw = safeParseInt( + getMergedEnv()["AGENTMEMORY_INDEX_SAVE_INTERVAL_MS"], + INDEX_SAVE_INTERVAL_DEFAULT_MS, + ); + return raw > 0 ? raw : INDEX_SAVE_INTERVAL_DEFAULT_MS; +} + +export const VECTOR_BUCKET_SIZE_DEFAULT = 500; + +export function getVectorBucketSize(): number { + const raw = safeParseInt(getMergedEnv()["AGENTMEMORY_VECTOR_BUCKET_SIZE"], VECTOR_BUCKET_SIZE_DEFAULT); + return raw > 0 ? raw : VECTOR_BUCKET_SIZE_DEFAULT; +} + +export const VECTOR_BACKFILL_MAX_DEFAULT = 500; + +export function getVectorBackfillMax(): number { + const raw = safeParseInt(getMergedEnv()["AGENTMEMORY_VECTOR_BACKFILL_MAX"], VECTOR_BACKFILL_MAX_DEFAULT); + return raw > 0 ? raw : VECTOR_BACKFILL_MAX_DEFAULT; +} + +export function isVectorBackfillAllEnabled(): boolean { + return getMergedEnv()["AGENTMEMORY_VECTOR_BACKFILL"] === "all"; +} + export function isStandaloneMcp(): boolean { return getMergedEnv()["STANDALONE_MCP"] === "true"; } @@ -548,10 +576,3 @@ export function getEvictionIntervalMs(): number { 86400000, ); } - -export function getIndexReclaimBootMaxDeletes(): number { - const raw = getEnvVar("AGENTMEMORY_INDEX_RECLAIM_BOOT_MAX_SHARDS"); - if (raw === undefined || !/^\d+$/.test(raw.trim())) return 200; - const value = Number(raw.trim()); - return Number.isSafeInteger(value) && value >= 0 ? value : 200; -} diff --git a/src/functions/search.ts b/src/functions/search.ts index 9b0ae70f2f..312a185629 100644 --- a/src/functions/search.ts +++ b/src/functions/search.ts @@ -1,4 +1,5 @@ import type { ISdk } from 'iii-sdk' +import type { IndexPersistenceStatus } from "../state/index-persistence.js"; import type { CompactSearchResult, CompressedObservation, Memory, SearchResult, Session } from '../types.js' import { KV } from '../state/schema.js' import { StateKV } from '../state/kv.js' @@ -8,7 +9,12 @@ import type { EmbeddingProvider } from '../types.js' import { memoryToObservation } from '../state/memory-utils.js' import { recordAccessBatch } from './access-tracker.js' import { logger } from "../logger.js"; -import { getAgentId, isAgentScopeIsolated } from "../config.js"; +import { + getAgentId, + getVectorBackfillMax, + isAgentScopeIsolated, + isVectorBackfillAllEnabled, +} from "../config.js"; let index: SearchIndex | null = null let vectorIndex: VectorIndex | null = null @@ -44,6 +50,19 @@ export function isMemoryIndexReady(): boolean { return memoryIndexReady } +let bm25RebuildIncomplete = false +export function isBm25RebuildIncomplete(): boolean { + return bm25RebuildIncomplete +} + +let pendingVectorBackfill = 0 +export function getPendingVectorBackfillCount(): number { + return pendingVectorBackfill +} +export function setPendingVectorBackfillCount(n: number): void { + pendingVectorBackfill = Math.max(0, n) +} + export function getSearchIndex(): SearchIndex { if (!index) index = new SearchIndex() return index @@ -75,17 +94,22 @@ export function vectorIndexRemove(id: string): void { // Wired by src/index.ts after IndexPersistence is constructed; no-op // until then so unit tests that exercise the delete paths in // isolation don't need to wire persistence. -let indexPersistence: { +type IndexPersistenceHook = { scheduleSave: () => void; save: () => Promise; -} | null = null; + status?: () => IndexPersistenceStatus; +}; + +let indexPersistence: IndexPersistenceHook | null = null; -export function setIndexPersistence( - p: { scheduleSave: () => void; save: () => Promise } | null, -): void { +export function setIndexPersistence(p: IndexPersistenceHook | null): void { indexPersistence = p; } +export function getIndexPersistenceStatus(): IndexPersistenceStatus | null { + return indexPersistence?.status?.() ?? null; +} + export function scheduleIndexSave(): void { indexPersistence?.scheduleSave(); } @@ -385,6 +409,146 @@ export async function rebuildIndex(kv: StateKV): Promise { return indexed } +export type VectorBackfillJob = { + id: string + sessionId: string + text: string + context: { kind: "memory" | "observation" | "synthetic"; logId: string } +} + +const SESSIONS_LIST_RETRY_ATTEMPTS = 3 +const SESSIONS_LIST_RETRY_BASE_MS = 250 + +function delay(ms: number): Promise { + return new Promise((resolve) => setTimeout(resolve, ms)) +} + +async function listSessionsWithRetry(kv: StateKV): Promise { + let lastErr: unknown + for (let attempt = 0; attempt < SESSIONS_LIST_RETRY_ATTEMPTS; attempt++) { + try { + return await kv.list(KV.sessions) + } catch (err) { + lastErr = err + if (attempt < SESSIONS_LIST_RETRY_ATTEMPTS - 1) { + await delay(SESSIONS_LIST_RETRY_BASE_MS * (attempt + 1)) + } + } + } + logger.warn('rebuildKeywordIndex: failed to load sessions after retries', { + attempts: SESSIONS_LIST_RETRY_ATTEMPTS, + error: lastErr instanceof Error ? lastErr.message : String(lastErr), + }) + return null +} + +export async function rebuildKeywordIndex( + kv: StateKV, + vectorBackfillSince?: string | null, +): Promise<{ documents: number; vectorJobs: VectorBackfillJob[]; fullBackfillPending: number }> { + const idx = getSearchIndex() + idx.clear() + memoryIndexReady = false + bm25RebuildIncomplete = false + const vi = vectorIndex + const backfillEligible = Boolean(vi && currentEmbeddingProvider) && vectorBackfillSince !== undefined + const wholeStoreBackfill = vectorBackfillSince === null + const backfillGated = backfillEligible && wholeStoreBackfill && !isVectorBackfillAllEnabled() + const backfillActive = backfillEligible && !backfillGated + const cutoff = typeof vectorBackfillSince === 'string' ? Date.parse(vectorBackfillSince) : Number.NaN + const backfillCap = backfillActive ? getVectorBackfillMax() : 0 + const vectorJobs: VectorBackfillJob[] = [] + let fullBackfillPending = 0 + const consider = ( + id: string, + sessionId: string, + text: string, + timestamp: string | undefined, + kind: "memory" | "observation", + ): void => { + if (!backfillEligible || vi?.has(id)) return + if (!Number.isNaN(cutoff) && !(Date.parse(timestamp ?? '') > cutoff)) return + if (backfillGated) { + fullBackfillPending++ + return + } + if (!backfillActive || vectorJobs.length >= backfillCap) return + vectorJobs.push({ id, sessionId, text, context: { kind, logId: id } }) + } + + let documents = 0 + let memoriesLoaded = false + try { + const memories = await kv.list(KV.memories) + memoriesLoaded = true + for (const memory of memories) { + if (memory.isLatest === false) continue + if (!memory.title || !memory.content) continue + idx.add(memoryToObservation(memory)) + consider(memory.id, memory.sessionIds?.[0] ?? 'memory', memory.title + ' ' + memory.content, memory.createdAt, "memory") + documents++ + } + } catch (err) { + logger.warn('rebuildKeywordIndex: failed to load memories', { + error: err instanceof Error ? err.message : String(err), + }) + } + + const sessions = await listSessionsWithRetry(kv) + if (sessions === null) { + bm25RebuildIncomplete = true + if (memoriesLoaded) memoryIndexReady = true + return { documents, vectorJobs, fullBackfillPending } + } + const failedSessions: string[] = [] + for (let batch = 0; batch < sessions.length; batch += 10) { + const chunk = sessions.slice(batch, batch + 10) + const results = await Promise.all( + chunk.map(async (s) => { + try { + return await kv.list(KV.observations(s.id)) + } catch { + failedSessions.push(s.id) + return [] as CompressedObservation[] + } + }) + ) + for (const obs of results.flat()) { + if (!obs.title || !obs.narrative) continue + idx.add(obs) + consider(obs.id, obs.sessionId, obs.title + ' ' + obs.narrative, obs.timestamp, "observation") + documents++ + } + } + if (failedSessions.length > 0) { + bm25RebuildIncomplete = true + logger.warn('rebuildKeywordIndex: failed to load observations for sessions', { failedSessions }) + } + if (memoriesLoaded) memoryIndexReady = true + return { documents, vectorJobs, fullBackfillPending } +} + +const BACKFILL_SAVE_EVERY_BATCHES = 10 + +export async function backfillVectors(jobs: VectorBackfillJob[]): Promise { + const batchSize = getRebuildEmbedBatchSize() + let added = 0 + let batchesSinceSave = 0 + setPendingVectorBackfillCount(jobs.length) + for (let offset = 0; offset < jobs.length; offset += batchSize) { + const { ok } = await vectorIndexAddBatchGuarded(jobs.slice(offset, offset + batchSize)) + added += ok + setPendingVectorBackfillCount(jobs.length - Math.min(jobs.length, offset + batchSize)) + batchesSinceSave++ + if (batchesSinceSave >= BACKFILL_SAVE_EVERY_BATCHES) { + await flushIndexSave() + batchesSinceSave = 0 + } + } + if (added > 0) await flushIndexSave() + return added +} + export function registerSearchFunction(sdk: ISdk, kv: StateKV): void { sdk.registerFunction( 'mem::search', @@ -466,10 +630,10 @@ export function registerSearchFunction(sdk: ISdk, kv: StateKV): void { // Share one rebuild across concurrent cold-start queries so they // don't each walk the whole corpus and saturate the pool. if (!rebuildPromise) { - rebuildPromise = rebuildIndex(kv) - .then((count) => { - logger.info('Search index rebuilt', { entries: count }) - return count + rebuildPromise = rebuildKeywordIndex(kv) + .then((result) => { + logger.info('Search index rebuilt', { entries: result.documents }) + return result.documents }) .catch((err) => { logger.warn('Index rebuild failed', { diff --git a/src/functions/status.ts b/src/functions/status.ts index 552fe6c07a..c6577c44e6 100644 --- a/src/functions/status.ts +++ b/src/functions/status.ts @@ -1,3 +1,5 @@ +import type { IndexLegStatus, IndexPersistenceStatus } from "../state/index-persistence.js"; + export type StatusLevel = "ok" | "info" | "warn" | "error"; export interface StatusProblem { @@ -55,9 +57,12 @@ export interface StatusInputs { observationsIndexed: number; missingObservations: number | null; sessions: number | null; + bm25Incomplete: boolean; + pendingVectorBackfill: number; }; graph: GraphStatsInput | null; graphExtractionEnabled: boolean; + indexPersistence?: IndexPersistenceStatus | null; } export interface StatusReport { @@ -72,6 +77,7 @@ export interface StatusReport { health: StatusInputs["health"]; provider: { llm: string; embeddings: string; circuitBreaker: StatusInputs["circuitBreaker"] }; index: StatusInputs["index"]; + indexPersistence: IndexPersistenceStatus | null; graph: (GraphStatsInput & { ageSeconds: number | null; extractionEnabled: boolean }) | null; functions: Array; flags: StatusFlag[]; @@ -168,6 +174,55 @@ export function evaluateStatus(input: StatusInputs): StatusReport { }); } + if (input.index.bm25Incomplete) { + problems.push({ + level: "error", + code: "bm25-rebuild-incomplete", + message: "The keyword index could not load every session at boot, so some observations are missing from search until the next successful rebuild.", + fix: "Check the server log for the session listing failure. Restarting or the next cold-start search retries automatically.", + }); + } + + if (input.index.pendingVectorBackfill > 0) { + problems.push({ + level: "info", + code: "index-vector-backfill-pending", + message: `${input.index.pendingVectorBackfill} documents are waiting for a vector embedding.`, + fix: "This runs in the background and is capped per boot by AGENTMEMORY_VECTOR_BACKFILL_MAX. If it is stuck at a nonzero count with no embedding provider errors, set AGENTMEMORY_VECTOR_BACKFILL=all to run a full backfill.", + }); + } + + const persistence = input.indexPersistence ?? null; + const vectorLeg = persistence?.vector ?? null; + if (persistence && vectorLeg) { + if (vectorLeg.lastError) { + problems.push({ + level: "error", + code: "index-save-failing", + message: `The vector index could not be saved: ${vectorLeg.lastError}. Search keeps working from memory, but vectors added since the last save are lost on restart.`, + fix: "Check the server log for the failing state write. The next save retries automatically.", + }); + } + const dirtyAge = secondsBetween(input.now, vectorLeg.dirtySince ?? undefined); + if (!vectorLeg.lastError && dirtyAge !== null && dirtyAge * 1000 > 2 * persistence.saveIntervalMs) { + problems.push({ + level: "warn", + code: "index-save-stale", + message: `The vector index has unsaved changes from ${formatDuration(dirtyAge)} ago.`, + fix: "Saves run at most once per AGENTMEMORY_INDEX_SAVE_INTERVAL_MS. Check the server log for save errors or a save that never finishes.", + }); + } + } + if (persistence?.vectorCountShortfall) { + const { expected, loaded } = persistence.vectorCountShortfall; + problems.push({ + level: "warn", + code: "index-vector-count-shortfall", + message: `Only ${loaded} of ${expected} vectors loaded from the last save; a bounded backfill re-embeds the rest.`, + fix: "No action needed — missing vectors are queued for backfill up to the per-boot cap.", + }); + } + let graph: StatusReport["graph"] = null; if (input.graph) { const ageSeconds = secondsBetween(input.now, input.graph.updatedAt); @@ -218,6 +273,7 @@ export function evaluateStatus(input: StatusInputs): StatusReport { circuitBreaker: input.circuitBreaker, }, index: input.index, + indexPersistence: persistence, graph, functions, flags: input.flags, @@ -248,6 +304,26 @@ function row(label: string, value: string): string { return `${escapeHtml(label)}${value}`; } +function legSummary(report: StatusReport, leg: IndexLegStatus): string { + if (leg.lastError) return `failing: ${leg.lastError}`; + const saved = leg.lastSavedAt ? `saved ${formatDuration(secondsBetween(new Date(report.checkedAt), leg.lastSavedAt))} ago` : "not saved since start"; + return leg.dirtySince ? `${saved}, unsaved changes pending` : saved; +} + +function indexPersistenceRows(report: StatusReport): string { + const persistence = report.indexPersistence; + if (!persistence) return ""; + let rows = row("BM25 index", "rebuilt from stored content at boot"); + if (persistence.vector) { + rows += row("Vector save", escapeHtml(legSummary(report, persistence.vector))); + rows += row( + "Vector storage", + escapeHtml(`${persistence.buckets} buckets, ${persistence.pendingChanges} unsaved changes`), + ); + } + return rows; +} + export function renderStatusHtml(report: StatusReport, styleNonce: string): string { const problems = report.problems.length ? report.problems @@ -330,6 +406,9 @@ ${row("Vector documents", escapeHtml(idx.vectorDocuments ?? "vector search off") ${row("Observations indexed", escapeHtml(idx.observationsIndexed))} ${row("Missing from index", escapeHtml(idx.missingObservations ?? "not checked"))} ${row("Sessions", escapeHtml(idx.sessions ?? "unknown"))} +${row("BM25 rebuild", idx.bm25Incomplete ? 'incomplete' : "complete")} +${row("Pending vector backfill", escapeHtml(idx.pendingVectorBackfill))} +${indexPersistenceRows(report)}

Knowledge graph

${graph diff --git a/src/index.ts b/src/index.ts index 5d0a43c4fe..9e49359a38 100644 --- a/src/index.ts +++ b/src/index.ts @@ -15,7 +15,6 @@ import { isDropStaleIndexEnabled, isEvictionEnabled, getEvictionIntervalMs, - getIndexReclaimBootMaxDeletes, } from "./config.js"; import { createProvider, @@ -24,7 +23,6 @@ import { createImageEmbeddingProvider, } from "./providers/index.js"; import { StateKV } from "./state/kv.js"; -import { KV } from "./state/schema.js"; import { VectorIndex } from "./state/vector-index.js"; import { HybridSearch } from "./state/hybrid-search.js"; import { IndexPersistence } from "./state/index-persistence.js"; @@ -37,12 +35,14 @@ import { registerDiskSizeManager } from "./functions/disk-size-manager.js"; import { registerCompressFunction } from "./functions/compress.js"; import { registerSearchFunction, - rebuildIndex, + backfillVectors, + rebuildKeywordIndex, getSearchIndex, setVectorIndex, setEmbeddingProvider, setIndexPersistence, setHybridRanker, + setPendingVectorBackfillCount, } from "./functions/search.js"; import { registerContextFunction } from "./functions/context.js"; import { registerSummarizeFunction } from "./functions/summarize.js"; @@ -405,26 +405,13 @@ async function main() { const healthMonitor = registerHealthMonitor(sdk, kv); - const indexPersistence = new IndexPersistence(kv, bm25Index, vectorIndex, { - reclaimOnLoad: "background", - bootReclaimMaxDeletes: getIndexReclaimBootMaxDeletes(), - }); - // Wire the persistence hook so delete paths can flush BM25/vector - // index mutations to disk. Without this, an in-memory remove can be - // lost across a hard process exit and the persisted snapshot - // restores the deleted entry at next boot. + const indexPersistence = new IndexPersistence(kv, vectorIndex); setIndexPersistence(indexPersistence); const loaded = await indexPersistence.load().catch((err) => { - console.warn(`[agentmemory] Failed to load persisted index:`, err); + console.warn(`[agentmemory] Failed to load persisted vector index:`, err); return null; }); - if (loaded?.bm25 && loaded.bm25.size > 0) { - bm25Index.restoreFrom(loaded.bm25); - bootLog( - `Loaded persisted BM25 index (${bm25Index.size} docs)`, - ); - } if (loaded?.vector && vectorIndex && loaded.vector.size > 0) { // Persisted vectors carry whatever dimension the provider had when // they were written. If the active provider declares a different @@ -449,6 +436,7 @@ async function main() { const distinct = Array.from(seenDimensions).sort((a, b) => a - b).join(", "); const dropStale = isDropStaleIndexEnabled(); if (dropStale) { + for (const [obsId] of loaded.vector.entries()) vectorIndex.markRemoved(obsId); console.warn( `[agentmemory] Persisted vector index has ${mismatches.length} of ` + `${loaded.vector.size} vectors with the wrong dimension. Active ` + @@ -479,68 +467,42 @@ async function main() { } } - const needsRebuild = bm25Index.size === 0; - - if (needsRebuild) { - // Fire-and-forget. rebuildIndex iterates every observation across - // every session and AWAITS an embedding-provider call per record. - // On a large corpus + rate-limited embedding endpoint that can - // take HOURS; awaiting it here blocks every subsequent boot step - // (including startViewerServer below, leaving the viewer port - // unbound for the duration). The index lazily fills in over time - // and search degrades gracefully — partial coverage > no viewer - // for hours. Errors still surface via the inner .catch. - void rebuildIndex(kv) - .then((indexCount) => { - if (indexCount > 0) { - bootLog(`Search index rebuilt: ${indexCount} entries`); - indexPersistence.scheduleSave(); - } - }) - .catch((err) => { - console.warn(`[agentmemory] Failed to rebuild search index:`, err); - }); - } else { - // Backfill memories into BM25 for users upgrading from <0.9.5: prior - // versions of mem::remember never indexed memories, so the persisted - // BM25 covers observations only and `memory_smart_search` returns - // empty for everything saved via memory_save (#257). Walk KV.memories - // and add the ones missing from the restored index. Idempotent on - // re-runs because SearchIndex.has() short-circuits already-indexed - // ids. - try { - const memories = await kv.list(KV.memories); - let backfilled = 0; - for (const memory of memories) { - if (memory.isLatest === false) continue; - if (!memory.title || !memory.content) continue; - if (bm25Index.has(memory.id)) continue; - bm25Index.add({ - id: memory.id, - sessionId: memory.sessionIds?.[0] ?? "memory", - timestamp: memory.createdAt, - type: "decision", - title: memory.title, - facts: [memory.content], - narrative: memory.content, - concepts: memory.concepts, - files: memory.files, - importance: memory.strength, - }); - backfilled++; - } - if (backfilled > 0) { - bootLog( - `Backfilled ${backfilled} memories into BM25 (legacy index gap)`, - ); - indexPersistence.scheduleSave(); - } - } catch (err) { - console.warn( - `[agentmemory] Failed to backfill memories into BM25:`, - err, + const vectorCountShortfall = + Boolean(loaded?.vector) && + loaded?.expectedCount !== undefined && + loaded.vector!.size < loaded.expectedCount; + const vectorBackfillSince = + !loaded || loaded.state === "unavailable" + ? undefined + : loaded.state === "none" || vectorCountShortfall + ? null + : loaded.savedAt; + const keywordStart = Date.now(); + try { + const keyword = await rebuildKeywordIndex(kv, vectorBackfillSince); + bootLog( + `Rebuilt BM25 index from stored content (${keyword.documents} docs in ${Date.now() - keywordStart} ms)`, + ); + setPendingVectorBackfillCount(keyword.vectorJobs.length + keyword.fullBackfillPending); + if (keyword.fullBackfillPending > 0) { + bootLog( + `Vector backfill needs ${keyword.fullBackfillPending} embeddings but a full backfill was not started ` + + `(set AGENTMEMORY_VECTOR_BACKFILL=all to opt in). See /agentmemory/status.`, ); } + if (keyword.vectorJobs.length > 0) { + bootLog(`Backfilling ${keyword.vectorJobs.length} missing vectors in the background`); + void backfillVectors(keyword.vectorJobs) + .then((count) => { + setPendingVectorBackfillCount(keyword.fullBackfillPending); + if (count > 0) bootLog(`Vector index backfilled: ${count} entries`); + }) + .catch((err) => { + console.warn(`[agentmemory] Failed to backfill vectors:`, err); + }); + } + } catch (err) { + console.warn(`[agentmemory] Failed to rebuild the BM25 index:`, err); } // Ready / Endpoints lines are emitted via `bootLog` so they're diff --git a/src/state/index-persistence.ts b/src/state/index-persistence.ts index 0970137d5a..d046d0dd0f 100644 --- a/src/state/index-persistence.ts +++ b/src/state/index-persistence.ts @@ -1,30 +1,21 @@ -import { SearchIndex } from "./search-index.js"; -import { VectorIndex } from "./vector-index.js"; +import { VectorIndex, base64ToFloat32, float32ToBase64 } from "./vector-index.js"; import type { StateKV } from "./kv.js"; -import { KV, generateId } from "./schema.js"; +import { KV } from "./schema.js"; import { logger } from "../logger.js"; import { safeAudit } from "../functions/audit.js"; +import { getIndexSaveIntervalMs, getVectorBucketSize } from "../config.js"; -const DEBOUNCE_MS = 5000; const FAILURE_LOG_THROTTLE_MS = 60_000; const INDEX_PERSISTENCE_FUNCTION_ID = "mem::index-persistence"; -const BM25_KEY = "data"; -const BM25_MANIFEST_KEY = "data:manifest"; -const BM25_SHARD_SCOPE_PREFIX = `${KV.bm25Index}:bm25:`; -const VECTOR_KEY = "vectors"; -const VECTOR_MANIFEST_KEY = "vectors:manifest"; -const VECTOR_SHARD_SCOPE_PREFIX = `${KV.bm25Index}:vectors:`; -const INDEX_SHARD_KEY = "data"; -const DEFAULT_INDEX_SHARD_CHARS = 2_000_000; -const DEFAULT_BOOT_RECLAIM_MAX_DELETES = 200; - -// mem:audit exists to record structural deletions of user data — that is -// the policy stated at the top of src/functions/audit.ts. Index shard -// writes and manifest publishes remove no user rows, so they fall outside -// it, yet a single save() emits three of them: on a real store they -// reached 59876 of 84028 entries (71%), which is what makes the audit log -// slow to query and bloats startup. Off by default; set -// AGENTMEMORY_AUDIT_INDEX_PERSIST=1 when debugging index persistence. +const LEGACY_BM25_KEY = "data"; +const LEGACY_BM25_MANIFEST_KEY = "data:manifest"; +const LEGACY_VECTOR_KEY = "vectors"; +const LEGACY_VECTOR_MANIFEST_KEY = "vectors:manifest"; +const VECTOR_META_KEY = "vectors:meta"; +const VECTOR_BUCKET_SCOPE_PREFIX = `${KV.bm25Index}:vec:`; +const WRITE_CONCURRENCY = 32; +const LOAD_CONCURRENCY = 8; + function auditIndexPersistEnabled(): boolean { const raw = process.env.AGENTMEMORY_AUDIT_INDEX_PERSIST; if (!raw) return false; @@ -39,45 +30,57 @@ type IndexShardManifest = { chars: number; }; -// Suffix for the reclaim ledger that sits beside each manifest. -const GC_LEDGER_SUFFIX = ":gc"; - -// Stands in for a manifest written before generations were recorded, whose -// `generation` field is absent. Cannot collide with createIndexGeneration(). -const PRE_LEDGER_GENERATION = "pre-ledger"; +type VectorMeta = { + v: 3; + bucketCount: number; + savedAt: string; + count: number; +}; -// Every generation whose shards may still be on disk, live one included. The -// manifest alone cannot answer that: it names only the generation that is -// current, so a generation stranded by a failed read, a throw after commit, or -// a kill mid-cleanup becomes unreachable the moment the next manifest replaces -// it. Nothing else enumerates shard scopes — StateKV lists keys within a scope, -// not scopes by prefix — so what is not recorded here can never be found again. -type IndexGcLedger = { - v: 1; - generations: Array<{ - generation: string; - shards: Array<{ scope: string; key: string }>; - }>; +type PersistedVector = { + id: string; + s: string; + e: string; }; type IndexPersistenceOptions = { - shardChars?: number; - createGeneration?: () => string; - reclaimOnLoad?: "blocking" | "background"; - bootReclaimMaxDeletes?: number; + saveIntervalMs?: number; + bucketSize?: number; + now?: () => number; }; -function shardChars(options: IndexPersistenceOptions): number { - const configured = options.shardChars; - if (typeof configured !== "number" || !Number.isFinite(configured)) { - return DEFAULT_INDEX_SHARD_CHARS; - } - const wholeChars = Math.floor(configured); - return wholeChars >= 1 ? wholeChars : DEFAULT_INDEX_SHARD_CHARS; +export interface IndexLegStatus { + lastSavedAt: string | null; + lastError: string | null; + lastErrorAt: string | null; + dirtySince: string | null; +} + +export interface VectorCountShortfall { + expected: number; + loaded: number; +} + +export interface IndexPersistenceStatus { + saveIntervalMs: number; + saving: boolean; + buckets: number; + pendingChanges: number; + vector: IndexLegStatus | null; + vectorCountShortfall: VectorCountShortfall | null; +} + +export type VectorLoadState = "buckets" | "migrated" | "none" | "unavailable"; + +export interface VectorLoadResult { + vector: VectorIndex | null; + state: VectorLoadState; + savedAt: string | null; + expectedCount?: number; } -function createIndexGeneration(): string { - return generateId("idx"); +export function vectorBucketScope(bucket: number): string { + return `${VECTOR_BUCKET_SCOPE_PREFIX}${String(bucket).padStart(4, "0")}`; } function statePath(scope: string, key: string): string { @@ -88,6 +91,13 @@ function errorMessage(err: unknown): string { return err instanceof Error ? err.message : String(err); } +function generationTime(generation: string | undefined): string | null { + const stamp = generation?.split("_")[1]; + if (!stamp) return null; + const ms = parseInt(stamp, 36); + return Number.isFinite(ms) && ms > 0 ? new Date(ms).toISOString() : null; +} + function isValidShardDescriptor( shard: unknown, ): shard is IndexShardManifest["shards"][number] { @@ -99,522 +109,386 @@ function isValidShardDescriptor( typeof candidate.key === "string" && candidate.key.length > 0 && Number.isInteger(candidate.chars) && - candidate.chars >= 0 + (candidate.chars as number) >= 0 ); } +function isPersistedVector(row: unknown): row is PersistedVector { + if (!row || typeof row !== "object") return false; + const candidate = row as Partial; + return typeof candidate.id === "string" && typeof candidate.s === "string" && typeof candidate.e === "string"; +} + +async function inBatches(items: T[], size: number, run: (item: T) => Promise): Promise { + const failures: unknown[] = []; + for (let offset = 0; offset < items.length; offset += size) { + const results = await Promise.allSettled(items.slice(offset, offset + size).map(run)); + for (const result of results) { + if (result.status === "rejected") failures.push(result.reason); + } + } + return failures; +} + +function emptyLegStatus(): IndexLegStatus { + return { lastSavedAt: null, lastError: null, lastErrorAt: null, dirtySince: null }; +} + export class IndexPersistence { private timer: ReturnType | null = null; - private lastFailureLogAt = new Map(); - private queue: Promise = Promise.resolve(); - /** - * A save that is queued but has not started running yet. - * - * It will serialise the index as it stands at the moment it runs, so it - * already covers every caller that arrives before then. Handing them this - * promise instead of queueing another identical save is what keeps a delete - * from waiting out an unbounded line of saves ahead of it. - */ - private pendingSave: Promise | null = null; + private lastFailureLogAt = 0; + private running: Promise | null = null; + private queued: Promise | null = null; + private stopped = false; + private lastSaveAt: number; + private dirtyEpoch = 0; + private markedDuringRunAt: number | null = null; + private metaWritten = false; + private leg: IndexLegStatus = emptyLegStatus(); + private readonly saveIntervalMs: number; + private readonly bucketSize: number; + private readonly now: () => number; + private bucketOfId: Map = new Map(); + private bucketCounts: Map = new Map(); + private highestBucket = 0; + private hasOpenBucket = false; + private vectorCountShortfall: VectorCountShortfall | null = null; constructor( private kv: StateKV, - private bm25: SearchIndex, private vector: VectorIndex | null, - private options: IndexPersistenceOptions = {}, - ) {} + options: IndexPersistenceOptions = {}, + ) { + this.now = options.now ?? Date.now; + const interval = options.saveIntervalMs; + this.saveIntervalMs = + typeof interval === "number" && Number.isFinite(interval) && interval > 0 + ? interval + : getIndexSaveIntervalMs(); + const bucketSize = options.bucketSize; + this.bucketSize = + typeof bucketSize === "number" && Number.isInteger(bucketSize) && bucketSize > 0 + ? bucketSize + : getVectorBucketSize(); + this.lastSaveAt = this.now(); + } scheduleSave(): void { - if (this.timer) clearTimeout(this.timer); - // setTimeout discards the returned promise, so any rejection inside - // save() would surface as unhandledRejection and crash the process - // under sustained iii-engine write timeouts (issue #204). Funnel - // rejections through logFailure() instead. + if (this.stopped || !this.vector) return; + const now = this.now(); + this.dirtyEpoch++; + if (this.running && this.markedDuringRunAt === null) this.markedDuringRunAt = now; + if (this.leg.dirtySince === null) this.leg.dirtySince = new Date(now).toISOString(); + if (this.timer) return; + const delay = Math.max(0, this.lastSaveAt + this.saveIntervalMs - now); this.timer = setTimeout(() => { - this.save().catch((err) => this.logFailure("index", err)); - }, DEBOUNCE_MS); + this.timer = null; + this.save().catch((err) => this.logFailure(err)); + }, delay); } - async save(): Promise { - if (this.timer) { - clearTimeout(this.timer); - this.timer = null; + save(): Promise { + this.clearTimer(); + if (this.queued) return this.queued; + if (this.running) { + const queued = this.running.then(() => { + this.queued = null; + return this.startRun(); + }); + this.queued = queued; + return queued; } - // Coalesce onto a save that is queued but has not started. It has not taken - // its snapshot yet, so it will include whatever this caller just changed. - // Only a save that is ALREADY RUNNING must not be handed back: its snapshot - // predates this caller's mutation, and reporting success for it would tell - // a delete it was durable when it was not. - if (this.pendingSave) return this.pendingSave; - - const pending = this.enqueue(() => { - // Running now. Its snapshot is fixed from here, so the next caller needs - // a save of its own. - if (this.pendingSave === pending) this.pendingSave = null; - return this.runSave(); - }); - this.pendingSave = pending; - return pending; + return this.startRun(); } - /** - * Serialise everything that read-modify-writes the gc ledger. - * - * Two saves genuinely overlap here: scheduleSave fires save() unawaited from - * a timer, flushIndexSave awaits save() on every delete path - * (src/functions/search.ts), and stop() clears the timer without awaiting a - * save already running. Overlapping saves drop each other's ledger entries, - * and worse, a save that publishes second reclaims the shards the first is - * still writing — leaving the first to publish a manifest naming data that - * is already gone, which fails the next load closed. - * - * Queue, never coalesce. flushIndexSave awaits this to make a delete - * durable, so handing back an in-flight promise that started before the - * delete would report success for a snapshot that does not contain it. The - * cost is real: a delete-path flush waits out the save ahead of it. - */ - private enqueue(work: () => Promise): Promise { - // runSave never rejects and the tail below always fulfils, so the queue - // cannot be left rejected and needs no rejection handler here. - const run = this.queue.then(work); - this.queue = run.then( - () => undefined, - () => undefined, - ); - return run; + status(): IndexPersistenceStatus { + return { + saveIntervalMs: this.saveIntervalMs, + saving: this.running !== null, + buckets: this.bucketCountInUse(), + pendingChanges: this.vector?.pendingChanges ?? 0, + vector: this.vector ? { ...this.leg } : null, + vectorCountShortfall: this.vectorCountShortfall, + }; } - private async runSave(): Promise { - // Each index fails on its own. One try around both would let a BM25 - // failure stop the vector index persisting at all, and a lost vector index - // is never rebuilt: both rebuild triggers key on the BM25 size - // (src/index.ts, src/functions/search.ts). + stop(): void { + this.stopped = true; + this.clearTimer(); + } + + async load(): Promise { + if (!this.vector) { + await this.removeLegacyBm25Snapshot(); + return { vector: null, state: "none", savedAt: null }; + } + + let meta: VectorMeta | null; try { - await this.saveBm25Index(this.bm25.serialize()); + meta = await this.kv.get(KV.bm25Index, VECTOR_META_KEY); } catch (err) { - this.logFailure("BM25", err); + logger.warn("index persistence: vector metadata read failed", { message: errorMessage(err) }); + return { vector: null, state: "unavailable", savedAt: null }; } - if (this.vector) { - try { - await this.saveVectorIndex(this.vector.serialize()); - } catch (err) { - this.logFailure("vector", err); + + if (meta && meta.v === 3 && Number.isInteger(meta.bucketCount) && meta.bucketCount >= 0) { + const loaded = await this.loadBuckets(meta.bucketCount); + if (!loaded) return { vector: null, state: "unavailable", savedAt: null }; + this.metaWritten = true; + await this.removeLegacyVectorSnapshotIfPresent(); + await this.removeLegacyBm25Snapshot(); + const expectedCount = Number.isInteger(meta.count) ? meta.count : undefined; + this.vectorCountShortfall = + expectedCount !== undefined && loaded.size < expectedCount + ? { expected: expectedCount, loaded: loaded.size } + : null; + if (this.vectorCountShortfall) { + logger.warn("index persistence: loaded fewer vectors than the last save recorded", { + ...this.vectorCountShortfall, + }); } + return { vector: loaded, state: "buckets", savedAt: meta.savedAt ?? null, expectedCount }; } - } - async load(): Promise<{ - bm25: SearchIndex | null; - vector: VectorIndex | null; - }> { - let bm25: SearchIndex | null = null; - let vector: VectorIndex | null = null; + return this.migrateLegacy(); + } - const bm25Data = await this.loadBm25Data(); - if (bm25Data && typeof bm25Data === "string") { - bm25 = SearchIndex.deserialize(bm25Data); - } + private bucketCountInUse(): number { + return this.hasOpenBucket ? this.highestBucket + 1 : 0; + } - const vecData = await this.loadVectorData(); - if (vecData && typeof vecData === "string") { - vector = VectorIndex.deserialize(vecData); + private assignBucket(id: string): number { + const existing = this.bucketOfId.get(id); + if (existing !== undefined) return existing; + if (!this.hasOpenBucket) { + this.hasOpenBucket = true; + this.highestBucket = 0; + } else if ((this.bucketCounts.get(this.highestBucket) ?? 0) >= this.bucketSize) { + this.highestBucket++; } + this.bucketOfId.set(id, this.highestBucket); + this.bucketCounts.set(this.highestBucket, (this.bucketCounts.get(this.highestBucket) ?? 0) + 1); + return this.highestBucket; + } - return { bm25, vector }; + private releaseBucket(id: string): void { + const bucket = this.bucketOfId.get(id); + if (bucket === undefined) return; + this.bucketOfId.delete(id); + const remaining = (this.bucketCounts.get(bucket) ?? 1) - 1; + if (remaining > 0) this.bucketCounts.set(bucket, remaining); + else this.bucketCounts.delete(bucket); } - stop(): void { + private clearTimer(): void { if (this.timer) { clearTimeout(this.timer); this.timer = null; } } - private logFailure(index: string, err: unknown): void { - const now = Date.now(); - // Throttle: persistence failures under load arrive in bursts - // (iii-engine queue pressure). Logging every debounce flush adds - // noise without information. Throttled PER INDEX, so a vector failure - // right after a BM25 one is not swallowed — the two fail independently - // now, and at 3am you need to know which one stopped persisting. - const lastAt = this.lastFailureLogAt.get(index) ?? 0; - if (now - lastAt < FAILURE_LOG_THROTTLE_MS) return; - this.lastFailureLogAt.set(index, now); - const code = (err as { code?: string })?.code; - const message = err instanceof Error ? err.message : String(err); - logger.warn(`index persistence: failed to save ${index} index`, { - code, - message, - hint: - code === "TIMEOUT" - ? "iii-engine state::set timed out; recent index updates remain in memory and will retry on the next debounce flush" - : undefined, + private startRun(): Promise { + const run = this.runSave().finally(() => { + if (this.running === run) this.running = null; }); + this.running = run; + return run; } - private async saveBm25Index(serialized: string): Promise { - await this.saveShardedIndex( - serialized, - BM25_MANIFEST_KEY, - BM25_KEY, - BM25_SHARD_SCOPE_PREFIX, - ); - } - - private async saveVectorIndex(serialized: string): Promise { - await this.saveShardedIndex( - serialized, - VECTOR_MANIFEST_KEY, - VECTOR_KEY, - VECTOR_SHARD_SCOPE_PREFIX, - ); - } - - private async saveShardedIndex( - serialized: string, - manifestKey: string, - legacyKey: string, - scopePrefix: string, - ): Promise { - const previous = await this.kv - .get(KV.bm25Index, manifestKey) - .catch(() => null); - const generation = - this.options.createGeneration?.() ?? createIndexGeneration(); - const chunkChars = shardChars(this.options); - const shards: IndexShardManifest["shards"] = []; - const chunks: string[] = []; - - for (let offset = 0; offset < serialized.length; offset += chunkChars) { - const shardIndex = shards.length; - const scope = `${scopePrefix}${generation}:${String(shardIndex).padStart( - 5, - "0", - )}`; - const chunk = serialized.slice(offset, offset + chunkChars); - shards.push({ scope, key: INDEX_SHARD_KEY, chars: chunk.length }); - chunks.push(chunk); - } - - // Record the generation BEFORE the first shard write. A kill anywhere from - // here to the manifest publish would otherwise leave shards on disk that - // nothing references and nothing can enumerate. - const tracked = await this.trackGeneration( - manifestKey, - generation, - shards, - previous, - ); - - const writeResults = await Promise.allSettled( - shards.map(async (shard, index) => { - const chunk = chunks[index] ?? ""; - await this.kv.set(shard.scope, shard.key, chunk); - await this.auditIndexPersistence("shard_write", [ - statePath(shard.scope, shard.key), - ], { - scope: shard.scope, - key: shard.key, - manifestKey, - generation, - chars: chunk.length, - }); - }), - ); - const failedWrite = writeResults.find( - (result): result is PromiseRejectedResult => result.status === "rejected", - ); - if (failedWrite) { - const allGone = await this.deleteShards(shards, "shard_write_rollback"); - // Drop the entry only if every shard actually went. A shard that survived - // its delete is unnameable once its entry is gone, which is the contract - // this file states for reclaim: a delete failure costs a retry, never a - // stranded generation. - if (tracked && allGone) { - await this.untrackGeneration(manifestKey, generation); + private async runSave(): Promise { + const vector = this.vector; + if (!vector) return; + const epoch = this.dirtyEpoch; + this.markedDuringRunAt = null; + this.lastSaveAt = this.now(); + try { + await this.writeChanges(vector); + this.leg.lastSavedAt = new Date(this.now()).toISOString(); + this.leg.lastError = null; + this.leg.lastErrorAt = null; + if (this.dirtyEpoch === epoch || vector.pendingChanges === 0) { + this.leg.dirtySince = null; + } else if (this.markedDuringRunAt !== null) { + this.leg.dirtySince = new Date(this.markedDuringRunAt).toISOString(); } - throw failedWrite.reason; + } catch (err) { + this.leg.lastError = errorMessage(err); + this.leg.lastErrorAt = new Date(this.now()).toISOString(); + if (this.leg.dirtySince === null) this.leg.dirtySince = this.leg.lastErrorAt; + this.logFailure(err); } + } - const nextManifest: IndexShardManifest = { - v: 1, - generation, - shards, - chars: serialized.length, - }; - try { - await this.kv.set( - KV.bm25Index, - manifestKey, - nextManifest, - ); - await this.auditIndexPersistence("manifest_publish", [ - statePath(KV.bm25Index, manifestKey), - ], { - manifestKey, - generation, - chars: serialized.length, - shards: shards.length, - result: "committed", - }); - } catch (err) { - if (await this.isManifestPublished(manifestKey, nextManifest)) { - await this.auditIndexPersistence("manifest_publish", [ - statePath(KV.bm25Index, manifestKey), - ], { - manifestKey, - generation, - chars: serialized.length, - shards: shards.length, - result: "committed_after_error", - error: errorMessage(err), - }); - } else { - const allGone = await this.deleteShards( - shards, - "manifest_publish_rollback", - ); - if (tracked && allGone) { - await this.untrackGeneration(manifestKey, generation); + private async writeChanges(vector: VectorIndex): Promise { + const changes = vector.takeChanges(); + if (changes.size === 0 && this.metaWritten) return; + const failed = new Map(); + const failures = await inBatches([...changes], WRITE_CONCURRENCY, async ([id, present]) => { + const entry = present ? vector.get(id) : undefined; + try { + if (entry) { + const bucket = this.assignBucket(id); + await this.kv.set(vectorBucketScope(bucket), id, { + id, + s: entry.sessionId, + e: float32ToBase64(entry.embedding), + }); + } else { + const bucket = this.bucketOfId.get(id); + if (bucket !== undefined) { + await this.kv.delete(vectorBucketScope(bucket), id); + this.releaseBucket(id); + } } + } catch (err) { + failed.set(id, present); + throw err; } - throw err; - } - - await this.deleteKey(KV.bm25Index, legacyKey, "legacy_cleanup"); - // Only reclaim when this generation is actually in the ledger. Reclaiming - // against a ledger that does not list us would treat live shards as dead. - if (tracked) { - await this.reclaimGenerations(manifestKey, generation); - } else if (previous?.v === 1 && Array.isArray(previous.shards)) { - // The ledger was unusable this cycle, so nothing above will ever revisit - // `previous`. Fall back to the pre-ledger cleanup: the manifest just - // published supersedes it, and saves are serialised, so its shards are - // dead. Without this, a run of unusable cycles orphans one generation - // each — strictly worse than the code this replaced, which always had - // this path. - const liveIds = new Set( - shards.map((shard) => `${shard.scope}\u0000${shard.key}`), - ); - await this.deleteShards( - previous.shards.filter( - (shard) => - isValidShardDescriptor(shard) && - !liveIds.has(`${shard.scope}\u0000${shard.key}`), - ), - "previous_generation_cleanup", + }); + if (failures.length > 0) { + vector.returnChanges(failed); + throw new Error( + `${failures.length} of ${changes.size} vector writes failed: ${errorMessage(failures[0])}`, ); } + await this.kv.set(KV.bm25Index, VECTOR_META_KEY, { + v: 3, + bucketCount: this.bucketCountInUse(), + savedAt: new Date(this.now()).toISOString(), + count: vector.size, + }); + this.metaWritten = true; } - /** Drop a generation's entry after its shards have been rolled back. */ - private async untrackGeneration( - manifestKey: string, - generation: string, - ): Promise { - try { - const ledger = await this.readLedger(manifestKey); - const generations = ledger.generations.filter( - (entry) => entry?.generation !== generation, - ); - if (generations.length === ledger.generations.length) return; - await this.kv.set(KV.bm25Index, this.gcKey(manifestKey), { - v: 1, - generations, + private async loadBuckets(bucketCount: number): Promise { + const loaded = new VectorIndex(); + const bucketOfId = new Map(); + const bucketCounts = new Map(); + const bucketIds = Array.from({ length: bucketCount }, (_, bucket) => bucket); + const failures = await inBatches(bucketIds, LOAD_CONCURRENCY, async (bucket) => { + const rows = await this.kv.list(vectorBucketScope(bucket)); + let count = 0; + for (const row of rows) { + if (!isPersistedVector(row)) continue; + try { + loaded.loadPersisted(row.id, row.s, base64ToFloat32(row.e)); + } catch { + continue; + } + bucketOfId.set(row.id, bucket); + count++; + } + if (count > 0) bucketCounts.set(bucket, count); + }); + if (failures.length > 0) { + logger.warn("index persistence: vector bucket read failed", { + failed: failures.length, + message: errorMessage(failures[0]), }); - } catch { - // Best effort. A surviving entry costs a retry on the next reclaim, and - // this runs while a save is already failing — never make that worse. + return null; } + this.bucketOfId = bucketOfId; + this.bucketCounts = bucketCounts; + this.highestBucket = bucketCount > 0 ? bucketCount - 1 : 0; + this.hasOpenBucket = bucketCount > 0; + return loaded; } - private gcKey(manifestKey: string): string { - return `${manifestKey}${GC_LEDGER_SUFFIX}`; - } - - private async readLedger(manifestKey: string): Promise { - // Throws rather than returning a blank ledger, because a blank one would - // be written straight back over whatever is stored. Callers catch it and - // skip tracking for that cycle; they must never treat it as "no ledger". - const stored = await this.kv.get( + private async migrateLegacy(): Promise { + const manifest = await this.readIndexValue( KV.bm25Index, - this.gcKey(manifestKey), + LEGACY_VECTOR_MANIFEST_KEY, + "vector", + "manifest", ); - if (stored == null) return { v: 1, generations: [] }; - // Present but unrecognised: a newer version, or a rollback to this build - // after one that wrote a different shape. Overwriting drops every - // generation it tracked, so leave it exactly where it is. - if (stored.v !== 1 || !Array.isArray(stored.generations)) { - throw new Error( - `index gc ledger ${this.gcKey(manifestKey)} has an unrecognised shape ` + - `(v=${String((stored as { v?: unknown }).v)}); refusing to overwrite it`, - ); - } - return stored; - } - - private async trackGeneration( - manifestKey: string, - generation: string, - shards: IndexShardManifest["shards"], - previous: IndexShardManifest | null, - ): Promise { - try { - // The WHOLE body runs under this guard, not just the ledger read. A - // malformed ledger entry or manifest shard would otherwise throw a - // TypeError out of saveShardedIndex before any shard write, leaving one - // throttled log line per 60s as the only trace while BM25 stopped - // persisting for good. isValidShardDescriptor exists in this file - // because a per-shard-malformed manifest is already its threat model. - return await this.recordGeneration( - manifestKey, - generation, - shards, - previous, - ); - } catch (err) { - // A state::get brownout is the exact condition this bug appears under, - // so this path is not rare. Aborting the save here would stop persisting - // the index at all, which is worse than the leak being fixed. Writing a - // fresh ledger would drop every generation the stored one lists. Do - // neither: let the shards and manifest land, skip this cycle's tracking - // and its reclaim, and leak at most one generation instead of one per - // brownout. - // - // That leak can OUTLIVE the unreadable window. This generation publishes - // untracked; the next save re-seeds it from `previous`, but that read is - // itself `.catch(() => null)`, so if it also fails the generation is in - // no ledger and reclaim only ever looks at what precedes the live entry. - // With no scope enumeration in StateKV, nothing can find it again. - // Throttled for the same reason every other failure log here is: an - // unusable ledger stays unusable, so this fires on every debounce. - const throttleKey = `gc:${manifestKey}`; - const now = Date.now(); - if (now - (this.lastFailureLogAt.get(throttleKey) ?? 0) >= FAILURE_LOG_THROTTLE_MS) { - this.lastFailureLogAt.set(throttleKey, now); - logger.warn( - "index persistence: gc ledger unavailable, skipping reclaim", - { manifestKey, message: errorMessage(err) }, - ); - } - return false; + const legacyKey = await this.readIndexValue(KV.bm25Index, LEGACY_VECTOR_KEY, "vector", "legacy"); + if (!manifest.ok || !legacyKey.ok) return { vector: null, state: "unavailable", savedAt: null }; + const legacyPresent = manifest.value != null || (typeof legacyKey.value === "string" && legacyKey.value.length > 0); + if (!legacyPresent) { + await this.removeLegacyBm25Snapshot(); + return { vector: null, state: "none", savedAt: null }; } - } - private async recordGeneration( - manifestKey: string, - generation: string, - shards: IndexShardManifest["shards"], - previous: IndexShardManifest | null, - ): Promise { - const ledger = await this.readLedger(manifestKey); - // Optional-chained so one malformed entry costs that entry, not the whole - // cycle's tracking. reclaimGenerations tolerates them the same way. - const known = new Set(ledger.generations.map((entry) => entry?.generation)); - - // Seed the generation the pre-ledger code left live, so upgrading does not - // strand it. Only reachable while `previous` is readable, which is exactly - // the case the old cleanup already handled. - // - // `generation` is optional on the manifest and older stores really do omit - // it, so fall back to a sentinel rather than skipping the seed. The - // sentinel cannot collide with createIndexGeneration()'s `idx_` ids, and - // since every manifest written from here on carries a generation, the - // seeded entry always ends up preceding a live one and gets reclaimed. - if (previous?.v === 1 && Array.isArray(previous.shards)) { - const previousGeneration = previous.generation ?? PRE_LEDGER_GENERATION; - if (!known.has(previousGeneration)) { - ledger.generations.push({ - generation: previousGeneration, - shards: previous.shards - .filter(isValidShardDescriptor) - .map(({ scope, key }) => ({ scope, key })), - }); - known.add(previousGeneration); - } + const data = await this.loadShardedData(LEGACY_VECTOR_KEY, LEGACY_VECTOR_MANIFEST_KEY, "vector"); + if (typeof data !== "string") { + logger.warn("index persistence: legacy vector snapshot is unreadable; keeping it and skipping migration"); + return { vector: null, state: "unavailable", savedAt: null }; } + const legacy = VectorIndex.deserialize(data); + const savedAt = + (manifest.value && typeof manifest.value === "object" ? generationTime(manifest.value.generation) : null) ?? + new Date(this.now()).toISOString(); - if (!known.has(generation)) { - ledger.generations.push({ - generation, - shards: shards.map(({ scope, key }) => ({ scope, key })), + const staging = new VectorIndex(); + for (const [id, entry] of legacy.entries()) staging.add(id, entry.sessionId, entry.embedding); + try { + await this.writeChanges(staging); + } catch (err) { + logger.warn("index persistence: migrating the legacy vector snapshot failed; it stays in place and the next save retries", { + message: errorMessage(err), }); + legacy.markAllChanged(); + return { vector: legacy, state: "unavailable", savedAt }; } - await this.kv.set( - KV.bm25Index, - this.gcKey(manifestKey), - ledger, - ); - return true; + await this.removeLegacyVectorSnapshot(manifest.value); + await this.removeLegacyBm25Snapshot(); + await this.auditIndexPersistence("migrate", [statePath(KV.bm25Index, VECTOR_META_KEY)], { + vectors: legacy.size, + buckets: this.bucketCountInUse(), + }); + logger.info("index persistence: migrated the vector index to bucketed storage", { + vectors: legacy.size, + buckets: this.bucketCountInUse(), + }); + const migrated = new VectorIndex(); + for (const [id, entry] of legacy.entries()) migrated.loadPersisted(id, entry.sessionId, entry.embedding); + return { vector: migrated, state: "migrated", savedAt }; } - private async reclaimGenerations( - manifestKey: string, - liveGeneration: string | undefined, - maxDeletes = Number.POSITIVE_INFINITY, - ): Promise { - if (!liveGeneration) return; - const ledger = await this.readLedger(manifestKey).catch(() => null); - if (!ledger) return; - - const liveIndex = ledger.generations.findIndex( - (entry) => entry?.generation === liveGeneration, - ); - if (liveIndex < 0) return; - - // Entries at or after the live generation may belong to an in-flight save; only earlier entries are superseded. - const reclaimedPaths: string[] = []; - let failed = 0; - let remainingBudget = Math.max(0, Math.floor(maxDeletes)); - const survivors: IndexGcLedger["generations"] = []; - for (const [index, entry] of ledger.generations.entries()) { - if (index >= liveIndex || !entry || !Array.isArray(entry.shards)) { - survivors.push(entry); - continue; - } - const stranded: IndexGcLedger["generations"][number]["shards"] = []; - for (const shard of entry.shards) { - if (remainingBudget <= 0) { - stranded.push(shard); - continue; - } - remainingBudget -= 1; - try { - await this.kv.delete(shard.scope, shard.key); - reclaimedPaths.push(statePath(shard.scope, shard.key)); - } catch { - failed += 1; - stranded.push(shard); - } - } - if (stranded.length > 0) { - survivors.push({ generation: entry.generation, shards: stranded }); - } + private async removeLegacyVectorSnapshot(manifest: IndexShardManifest | null): Promise { + if (manifest && Array.isArray(manifest.shards)) { + await this.deleteShards(manifest.shards.filter(isValidShardDescriptor), "legacy_vector_cleanup"); } + await this.deleteKey(KV.bm25Index, LEGACY_VECTOR_MANIFEST_KEY, "legacy_vector_cleanup"); + await this.deleteKey(KV.bm25Index, LEGACY_VECTOR_KEY, "legacy_vector_cleanup"); + } - if (reclaimedPaths.length > 0) { - await this.auditIndexPersistence("delete", reclaimedPaths, { - manifestKey, - reason: "generation_reclaim", - liveGeneration, - evicted: reclaimedPaths.length, - failed, - deferred: survivors.reduce( - (count, entry) => - count + (Array.isArray(entry?.shards) ? entry.shards.length : 0), - 0, - ), - }); + private async removeLegacyVectorSnapshotIfPresent(): Promise { + const manifest = await this.readIndexValue(KV.bm25Index, LEGACY_VECTOR_MANIFEST_KEY, "vector", "manifest"); + const legacy = await this.readIndexValue(KV.bm25Index, LEGACY_VECTOR_KEY, "vector", "legacy"); + if (!manifest.ok || !legacy.ok) return; + if (manifest.value == null && legacy.value == null) return; + await this.removeLegacyVectorSnapshot(manifest.value ?? null); + } + + private async removeLegacyBm25Snapshot(): Promise { + const manifestRead = await this.readIndexValue(KV.bm25Index, LEGACY_BM25_MANIFEST_KEY, "bm25", "manifest"); + const legacyRead = await this.readIndexValue(KV.bm25Index, LEGACY_BM25_KEY, "bm25", "legacy"); + if (!manifestRead.ok || !legacyRead.ok) return; + const manifest = manifestRead.value ?? null; + const legacy = legacyRead.value ?? null; + if (manifest == null && legacy == null) return; + if (manifest && Array.isArray(manifest.shards)) { + await this.deleteShards(manifest.shards.filter(isValidShardDescriptor), "legacy_bm25_cleanup"); } + if (manifest != null) await this.deleteKey(KV.bm25Index, LEGACY_BM25_MANIFEST_KEY, "legacy_bm25_cleanup"); + if (legacy != null) await this.deleteKey(KV.bm25Index, LEGACY_BM25_KEY, "legacy_bm25_cleanup"); + } - if (reclaimedPaths.length === 0) return; - await this.kv - .set(KV.bm25Index, this.gcKey(manifestKey), { - v: 1, - generations: survivors, - }) - .catch(() => undefined); + private logFailure(err: unknown): void { + const now = this.now(); + if (now - this.lastFailureLogAt < FAILURE_LOG_THROTTLE_MS) return; + this.lastFailureLogAt = now; + const code = (err as { code?: string })?.code; + logger.warn("index persistence: failed to save the vector index", { + code, + message: errorMessage(err), + hint: + code === "TIMEOUT" + ? "iii-engine state::set timed out; unsaved vectors stay in memory and retry on the next save" + : undefined, + }); } private async auditIndexPersistence( @@ -623,139 +497,41 @@ export class IndexPersistence { details: Record, ): Promise { if (!auditIndexPersistEnabled()) return; - await safeAudit( - this.kv, - "index_persist", - INDEX_PERSISTENCE_FUNCTION_ID, - targetIds, - { action, ...details }, - ); + await safeAudit(this.kv, "index_persist", INDEX_PERSISTENCE_FUNCTION_ID, targetIds, { + action, + ...details, + }); } - /** Reports whether the delete landed, so the reclaim path can retry the rest. */ - private async deleteKey( - scope: string, - key: string, - reason: string, - ): Promise { - let ok = true; + private async deleteKey(scope: string, key: string, reason: string): Promise { + let result = "deleted"; let error: string | undefined; try { await this.kv.delete(scope, key); } catch (err) { - ok = false; + result = "failed"; error = errorMessage(err); } await this.auditIndexPersistence("delete", [statePath(scope, key)], { scope, key, reason, - result: ok ? "deleted" : "failed", + result, error, }); - return ok; - } - - private async deleteShards( - shards: IndexShardManifest["shards"], - reason: string, - ): Promise { - let allGone = true; - for (const shard of shards) { - if (!(await this.deleteKey(shard.scope, shard.key, reason))) { - allGone = false; - } - } - return allGone; } - private async isManifestPublished( - manifestKey: string, - expected: IndexShardManifest, - ): Promise { - const published = await this.kv - .get(KV.bm25Index, manifestKey) - .catch(() => null); - if ( - published?.v !== 1 || - published.generation !== expected.generation || - published.chars !== expected.chars || - !Array.isArray(published.shards) || - published.shards.length !== expected.shards.length - ) { - return false; - } - return published.shards.every((shard, index) => { - const expectedShard = expected.shards[index]; - if (!expectedShard) return false; - return ( - shard.scope === expectedShard.scope && - shard.key === expectedShard.key && - shard.chars === expectedShard.chars - ); - }); + private async deleteShards(shards: IndexShardManifest["shards"], reason: string): Promise { + await inBatches(shards, WRITE_CONCURRENCY, (shard) => this.deleteKey(shard.scope, shard.key, reason)); } - private async loadBm25Data(): Promise { - return this.loadShardedData(BM25_KEY, BM25_MANIFEST_KEY, "BM25"); - } - - private async loadVectorData(): Promise { - return this.loadShardedData(VECTOR_KEY, VECTOR_MANIFEST_KEY, "vector"); - } - - private async loadShardedData( - legacyKey: string, - manifestKey: string, - label: string, - ): Promise { - const manifest = await this.readIndexValue( - KV.bm25Index, - manifestKey, - label, - "manifest", - ); + private async loadShardedData(legacyKey: string, manifestKey: string, label: string): Promise { + const manifest = await this.readIndexValue(KV.bm25Index, manifestKey, label, "manifest"); if (!manifest.ok) return null; - // #797: some iii-state adapters return `undefined` (not `null`) for - // a missing key. The previous `value !== null` check passed - // undefined through to loadManifestData, which then crashed on - // `manifest.v` with TypeError. Treat both null and undefined as - // "no manifest" and fall through to the legacy path. The shape - // check stays so a malformed-but-present row still fails closed. - if ( - manifest.value != null && - typeof manifest.value === "object" - ) { - const data = await this.loadManifestData(manifest.value, label); - // Boot is the only point that sees a generation stranded by a kill: the - // save path only ever inspects its own predecessor. Reclaim once the live - // generation has actually loaded — a failed load must not authorise - // deleting anything. - if (data !== null) { - const live = manifest.value.generation; - const reclaim = this.enqueue(() => - this.reclaimGenerations( - manifestKey, - live, - this.options.bootReclaimMaxDeletes ?? - DEFAULT_BOOT_RECLAIM_MAX_DELETES, - ), - ).catch(() => undefined); - if (this.options.reclaimOnLoad === "background") { - void reclaim; - } else { - await reclaim; - } - } - return data; + if (manifest.value != null && typeof manifest.value === "object") { + return this.loadManifestData(manifest.value, label); } - - const legacy = await this.readIndexValue( - KV.bm25Index, - legacyKey, - label, - "legacy", - ); + const legacy = await this.readIndexValue(KV.bm25Index, legacyKey, label, "legacy"); if (!legacy.ok) return null; if (legacy.value && typeof legacy.value === "string") return legacy.value; return null; @@ -779,10 +555,7 @@ export class IndexPersistence { } } - private async loadManifestData( - manifest: IndexShardManifest, - label: string, - ): Promise { + private async loadManifestData(manifest: IndexShardManifest, label: string): Promise { if ( manifest.v !== 1 || !Array.isArray(manifest.shards) || @@ -809,10 +582,7 @@ export class IndexPersistence { let chars = 0; for (const { shard, chunk } of loadedShards) { if (typeof chunk !== "string") { - logger.warn(`index persistence: ${label} shard missing`, { - scope: shard.scope, - key: shard.key, - }); + logger.warn(`index persistence: ${label} shard missing`, { scope: shard.scope, key: shard.key }); return null; } if (chunk.length !== shard.chars) { @@ -828,10 +598,7 @@ export class IndexPersistence { chars += chunk.length; } if (chars !== manifest.chars) { - logger.warn(`index persistence: ${label} total length mismatch`, { - expected: manifest.chars, - actual: chars, - }); + logger.warn(`index persistence: ${label} total length mismatch`, { expected: manifest.chars, actual: chars }); return null; } return chunks.join(""); diff --git a/src/state/search-index.ts b/src/state/search-index.ts index eae4b55b36..a1ae6cf278 100644 --- a/src/state/search-index.ts +++ b/src/state/search-index.ts @@ -20,6 +20,7 @@ export class SearchIndex { private readonly b = 0.75; add(obs: CompressedObservation): void { + if (this.entries.has(obs.id)) this.remove(obs.id); const terms = this.extractTerms(obs); const termFreq = new Map(); let termCount = 0; diff --git a/src/state/vector-index.ts b/src/state/vector-index.ts index d4b8bda760..2425bef8e1 100644 --- a/src/state/vector-index.ts +++ b/src/state/vector-index.ts @@ -5,13 +5,13 @@ // pool. Same risk on the encode side if the input Float32Array is itself // a sliced view. Reported as a phantom "2048 dimensions on disk" crash // in #455 / #469 / #584 / #587. -function float32ToBase64(arr: Float32Array): string { +export function float32ToBase64(arr: Float32Array): string { return Buffer.from(arr.buffer, arr.byteOffset, arr.byteLength).toString( "base64", ); } -function base64ToFloat32(b64: string): Float32Array { +export function base64ToFloat32(b64: string): Float32Array { const buf = Buffer.from(b64, "base64"); return new Float32Array( buf.buffer, @@ -34,16 +34,59 @@ function cosineSimilarity(a: Float32Array, b: Float32Array): number { return denom === 0 ? 0 : dot / denom; } +export type VectorEntry = { embedding: Float32Array; sessionId: string }; + export class VectorIndex { - private vectors: Map = - new Map(); + private vectors: Map = new Map(); + private changes: Map = new Map(); add(obsId: string, sessionId: string, embedding: Float32Array): void { this.vectors.set(obsId, { embedding, sessionId }); + this.changes.set(obsId, true); } remove(obsId: string): void { - this.vectors.delete(obsId); + if (this.vectors.delete(obsId)) this.changes.set(obsId, false); + } + + has(obsId: string): boolean { + return this.vectors.has(obsId); + } + + get(obsId: string): VectorEntry | undefined { + return this.vectors.get(obsId); + } + + entries(): IterableIterator<[string, VectorEntry]> { + return this.vectors.entries(); + } + + loadPersisted(obsId: string, sessionId: string, embedding: Float32Array): void { + this.vectors.set(obsId, { embedding, sessionId }); + } + + get pendingChanges(): number { + return this.changes.size; + } + + takeChanges(): Map { + const taken = this.changes; + this.changes = new Map(); + return taken; + } + + returnChanges(changes: Map): void { + for (const [obsId, present] of changes) { + if (!this.changes.has(obsId)) this.changes.set(obsId, present); + } + } + + markRemoved(obsId: string): void { + if (!this.vectors.has(obsId)) this.changes.set(obsId, false); + } + + markAllChanged(): void { + for (const obsId of this.vectors.keys()) this.changes.set(obsId, true); } search( @@ -104,6 +147,7 @@ export class VectorIndex { } clear(): void { + for (const obsId of this.vectors.keys()) this.changes.set(obsId, false); this.vectors.clear(); } @@ -119,6 +163,7 @@ export class VectorIndex { sessionId: entry.sessionId, }); } + this.changes = new Map(other.changes); } serialize(): string { diff --git a/src/triggers/api.ts b/src/triggers/api.ts index c18d99e8d8..d2b587aafe 100644 --- a/src/triggers/api.ts +++ b/src/triggers/api.ts @@ -11,7 +11,14 @@ import type { ResilientProvider } from "../providers/resilient.js"; import { III_PINNED_VERSION, VERSION } from "../version.js"; import { UNINDEXED_SCAN_REUSE_MS, evaluateStatus, prefersHtml, renderStatusHtml, singleFlight, type GraphStatsInput } from "../functions/status.js"; import { CONSOLIDATION_COUNTS_REUSE_MS, CONSOLIDATION_LAST_RUN_KEY, PROCEDURAL_MIN_SESSIONS_PER_PATTERN, describeConsolidation, type ConsolidationRunRecord } from "../functions/consolidation-status.js"; -import { findUnindexedObservations, getSearchIndex, getVectorIndex } from "../functions/search.js"; +import { + findUnindexedObservations, + getIndexPersistenceStatus, + getPendingVectorBackfillCount, + getSearchIndex, + getVectorIndex, + isBm25RebuildIncomplete, +} from "../functions/search.js"; import { timingSafeCompare } from "../auth.js"; import { isSlotsEnabled, isReflectEnabled } from "../functions/slots.js"; import { renderViewerDocument } from "../viewer/document.js"; @@ -336,9 +343,12 @@ export function registerApiTriggers( observationsIndexed, missingObservations: unindexed ? unindexed.missing.length : null, sessions: unindexed ? unindexed.sessions : null, + bm25Incomplete: isBm25RebuildIncomplete(), + pendingVectorBackfill: getPendingVectorBackfillCount(), }, graph, graphExtractionEnabled: isGraphExtractionEnabled(), + indexPersistence: getIndexPersistenceStatus(), }); const accept = req.headers?.["accept"] ?? req.headers?.["Accept"]; const format = req.query_params?.["format"]; diff --git a/test/evict.test.ts b/test/evict.test.ts index 8654ead610..37422c737c 100644 --- a/test/evict.test.ts +++ b/test/evict.test.ts @@ -2,7 +2,6 @@ import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { readFileSync } from "node:fs"; import { getEvictionIntervalMs, - getIndexReclaimBootMaxDeletes, isEvictionEnabled, parsePositiveIntervalMs, TIMER_MAX_INTERVAL_MS, @@ -408,23 +407,6 @@ describe("eviction scheduling", () => { expect(getEvictionIntervalMs()).toBe(3600000); }); - it("bounds boot index-reclaim deletes to a safe non-negative integer", () => { - delete process.env.AGENTMEMORY_INDEX_RECLAIM_BOOT_MAX_SHARDS; - expect(getIndexReclaimBootMaxDeletes()).toBe(200); - - process.env.AGENTMEMORY_INDEX_RECLAIM_BOOT_MAX_SHARDS = "50"; - expect(getIndexReclaimBootMaxDeletes()).toBe(50); - - process.env.AGENTMEMORY_INDEX_RECLAIM_BOOT_MAX_SHARDS = "0"; - expect(getIndexReclaimBootMaxDeletes()).toBe(0); - - for (const bogus of ["abc", "-5", "1.5", "9999999999999999999999"]) { - process.env.AGENTMEMORY_INDEX_RECLAIM_BOOT_MAX_SHARDS = bogus; - expect(getIndexReclaimBootMaxDeletes()).toBe(200); - } - delete process.env.AGENTMEMORY_INDEX_RECLAIM_BOOT_MAX_SHARDS; - }); - it("registers an unref'd interval with completion logging and an overlap guard", () => { expect(src).toMatch(/if\s*\(\s*isEvictionEnabled\(\)\s*\)/); expect(src).toMatch(/const\s+evictionTimer\s*=\s*setInterval/); diff --git a/test/index-keyword-rebuild.test.ts b/test/index-keyword-rebuild.test.ts new file mode 100644 index 0000000000..f0080675ba --- /dev/null +++ b/test/index-keyword-rebuild.test.ts @@ -0,0 +1,240 @@ +import { describe, it, expect, afterEach } from "vitest"; +import { + backfillVectors, + getPendingVectorBackfillCount, + getSearchIndex, + isBm25RebuildIncomplete, + isMemoryIndexReady, + rebuildKeywordIndex, + setEmbeddingProvider, + setIndexPersistence, + setVectorIndex, +} from "../src/functions/search.js"; +import { VectorIndex } from "../src/state/vector-index.js"; +import type { CompressedObservation, EmbeddingProvider, Memory, Session } from "../src/types.js"; + +function mockKV() { + const store = new Map>(); + return { + store, + list: async (scope: string): Promise => Array.from(store.get(scope)?.values() ?? []) as T[], + get: async (scope: string, key: string): Promise => (store.get(scope)?.get(key) as T) ?? null, + set: async (scope: string, key: string, data: T): Promise => { + if (!store.has(scope)) store.set(scope, new Map()); + store.get(scope)!.set(key, data); + return data; + }, + delete: async (scope: string, key: string): Promise => { + store.get(scope)?.delete(key); + }, + }; +} + +const WORDS = ["auth", "cache", "index", "vector", "session", "memory", "graph", "retry", "queue", "schema", "deploy", "worker"]; + +function makeObs(id: string, sessionId: string, i: number, timestamp: string): CompressedObservation { + const a = WORDS[i % WORDS.length]; + const b = WORDS[(i * 7) % WORDS.length]; + return { + id, + sessionId, + timestamp, + type: "file_edit", + title: `${a} change ${i}`, + facts: [`touched ${b} path`], + narrative: `Updated the ${a} module so ${b} handling works for case ${i}`, + concepts: [a, b], + files: [`src/${a}/${b}.ts`], + importance: 5, + }; +} + +function seed(kv: ReturnType, sessions: number, perSession: number, timestamp: (i: number) => string) { + let n = 0; + for (let s = 0; s < sessions; s++) { + const sessionId = `ses_${s}`; + kv.store.set("mem:sessions", kv.store.get("mem:sessions") ?? new Map()); + kv.store.get("mem:sessions")!.set(sessionId, { + id: sessionId, + project: "p", + cwd: "/tmp", + startedAt: timestamp(0), + status: "completed", + observationCount: perSession, + } satisfies Session); + const scope = `mem:obs:${sessionId}`; + kv.store.set(scope, new Map()); + for (let o = 0; o < perSession; o++) { + const id = `obs_${n}`; + kv.store.get(scope)!.set(id, makeObs(id, sessionId, n, timestamp(n))); + n++; + } + } + const memories = new Map(); + memories.set("mem_1", { + id: "mem_1", + createdAt: timestamp(0), + updatedAt: timestamp(0), + type: "fact", + title: "Deploy uses blue green", + content: "Deploys switch traffic after health checks pass", + concepts: ["deploy"], + files: [], + sessionIds: ["ses_0"], + strength: 7, + version: 1, + isLatest: true, + } as Memory); + memories.set("mem_old", { + id: "mem_old", + createdAt: timestamp(0), + updatedAt: timestamp(0), + type: "fact", + title: "Superseded", + content: "Old version", + concepts: [], + files: [], + sessionIds: [], + strength: 1, + version: 1, + isLatest: false, + } as Memory); + kv.store.set("mem:memories", memories); + kv.store.get("mem:obs:ses_0")!.set("raw_1", { id: "raw_1", sessionId: "ses_0", timestamp: timestamp(0) }); + return n; +} + +function stubProvider(): EmbeddingProvider { + return { + name: "stub", + dimensions: 3, + embed: async () => new Float32Array([0.1, 0.2, 0.3]), + embedBatch: async (texts: string[]) => texts.map(() => new Float32Array([0.1, 0.2, 0.3])), + } as EmbeddingProvider; +} + +describe("rebuildKeywordIndex", () => { + afterEach(() => { + getSearchIndex().clear(); + setVectorIndex(null); + setEmbeddingProvider(null); + setIndexPersistence(null); + }); + + it("rebuilds BM25 from stored observations and latest memories, quickly", async () => { + const kv = mockKV(); + const total = seed(kv, 100, 30, () => "2026-09-01T00:00:00.000Z"); + getSearchIndex().add(makeObs("ghost", "ses_gone", 0, "2026-09-01T00:00:00.000Z")); + + const started = performance.now(); + const result = await rebuildKeywordIndex(kv as never); + const elapsed = performance.now() - started; + + expect(total).toBe(3000); + expect(result.documents).toBe(3001); + expect(result.vectorJobs).toEqual([]); + const idx = getSearchIndex(); + expect(idx.size).toBe(3001); + expect(idx.has("ghost")).toBe(false); + expect(idx.has("mem_1")).toBe(true); + expect(idx.has("mem_old")).toBe(false); + expect(idx.has("raw_1")).toBe(false); + expect(idx.search("blue green").map((r) => r.obsId)).toContain("mem_1"); + expect(isMemoryIndexReady()).toBe(true); + expect(elapsed).toBeLessThan(10_000); + }); + + it("keeps the vector index untouched and reports an incomplete rebuild when the sessions listing fails", async () => { + const kv = mockKV(); + seed(kv, 1, 3, () => "2026-09-01T00:00:00.000Z"); + const vector = new VectorIndex(); + vector.add("obs_0", "ses_0", new Float32Array([0.1, 0.2, 0.3])); + setVectorIndex(vector); + setEmbeddingProvider(stubProvider()); + const failingKv = { + ...kv, + list: async (scope: string): Promise => { + if (scope === "mem:sessions") throw new Error("engine unreachable"); + return kv.list(scope); + }, + }; + + const result = await rebuildKeywordIndex(failingKv as never, null); + + expect(result.vectorJobs).toEqual([]); + expect(vector.size).toBe(1); + expect(vector.has("obs_0")).toBe(true); + expect(isBm25RebuildIncomplete()).toBe(true); + }, 10_000); + + it("queues embeddings only for documents newer than the last vector save", async () => { + const kv = mockKV(); + seed(kv, 2, 5, (i) => (i < 5 ? "2026-09-01T00:00:00.000Z" : "2026-09-20T00:00:00.000Z")); + const vector = new VectorIndex(); + vector.add("obs_6", "ses_1", new Float32Array([0.1, 0.2, 0.3])); + setVectorIndex(vector); + setEmbeddingProvider(stubProvider()); + + const result = await rebuildKeywordIndex(kv as never, "2026-09-10T00:00:00.000Z"); + + expect(result.vectorJobs.map((j) => j.id).sort()).toEqual(["obs_5", "obs_7", "obs_8", "obs_9"]); + }); + + it("withholds a whole-store backfill by default and reports it as pending, and queues none when storage was unavailable", async () => { + const kv = mockKV(); + seed(kv, 1, 3, () => "2026-09-01T00:00:00.000Z"); + setVectorIndex(new VectorIndex()); + setEmbeddingProvider(stubProvider()); + + const fresh = await rebuildKeywordIndex(kv as never, null); + expect(fresh.vectorJobs).toEqual([]); + expect(fresh.fullBackfillPending).toBe(4); + + const unavailable = await rebuildKeywordIndex(kv as never, undefined); + expect(unavailable.vectorJobs).toEqual([]); + expect(unavailable.fullBackfillPending).toBe(0); + }); + + it("runs the whole-store backfill, capped, once AGENTMEMORY_VECTOR_BACKFILL=all is set", async () => { + const kv = mockKV(); + seed(kv, 1, 3, () => "2026-09-01T00:00:00.000Z"); + setVectorIndex(new VectorIndex()); + setEmbeddingProvider(stubProvider()); + process.env.AGENTMEMORY_VECTOR_BACKFILL = "all"; + try { + const opted = await rebuildKeywordIndex(kv as never, null); + expect(opted.vectorJobs.map((j) => j.id).sort()).toEqual(["mem_1", "obs_0", "obs_1", "obs_2"]); + expect(opted.fullBackfillPending).toBe(0); + + process.env.AGENTMEMORY_VECTOR_BACKFILL_MAX = "2"; + const capped = await rebuildKeywordIndex(kv as never, null); + expect(capped.vectorJobs.length).toBe(2); + } finally { + delete process.env.AGENTMEMORY_VECTOR_BACKFILL; + delete process.env.AGENTMEMORY_VECTOR_BACKFILL_MAX; + } + }); + + it("backfills queued vectors in batches and saves progress", async () => { + const kv = mockKV(); + seed(kv, 1, 3, () => "2026-09-01T00:00:00.000Z"); + const vector = new VectorIndex(); + setVectorIndex(vector); + setEmbeddingProvider(stubProvider()); + let saved = 0; + setIndexPersistence({ scheduleSave: () => {}, save: async () => void saved++ }); + process.env.AGENTMEMORY_VECTOR_BACKFILL = "all"; + + try { + const { vectorJobs } = await rebuildKeywordIndex(kv as never, null); + const added = await backfillVectors(vectorJobs); + + expect(added).toBe(4); + expect(vector.size).toBe(4); + expect(saved).toBe(1); + expect(getPendingVectorBackfillCount()).toBe(0); + } finally { + delete process.env.AGENTMEMORY_VECTOR_BACKFILL; + } + }); +}); diff --git a/test/index-persistence-throttle.test.ts b/test/index-persistence-throttle.test.ts new file mode 100644 index 0000000000..bc03834d41 --- /dev/null +++ b/test/index-persistence-throttle.test.ts @@ -0,0 +1,248 @@ +import { describe, it, expect, beforeEach, afterEach, vi } from "vitest"; +import { IndexPersistence } from "../src/state/index-persistence.js"; +import { VectorIndex } from "../src/state/vector-index.js"; +import { evaluateStatus, renderStatusHtml, type StatusInputs } from "../src/functions/status.js"; + +const META_KEY = "vectors:meta"; + +function mockKV() { + const store = new Map>(); + const sets: Array<{ scope: string; key: string }> = []; + return { + sets, + get: async (scope: string, key: string): Promise => (store.get(scope)?.get(key) as T) ?? null, + set: async (scope: string, key: string, data: T): Promise => { + sets.push({ scope, key }); + if (!store.has(scope)) store.set(scope, new Map()); + store.get(scope)!.set(key, data); + return data; + }, + delete: async (scope: string, key: string): Promise => { + store.get(scope)?.delete(key); + }, + list: async (scope: string): Promise => Array.from(store.get(scope)?.values() ?? []) as T[], + }; +} + +function metaSaves(kv: ReturnType): number { + return kv.sets.filter((s) => s.key === META_KEY).length; +} + +function touch(vector: VectorIndex, id: string): void { + vector.add(id, "ses_1", new Float32Array([Math.random(), 0.2, 0.3])); +} + +describe("IndexPersistence save throttling", () => { + let kv: ReturnType; + + beforeEach(() => { + vi.useFakeTimers(); + kv = mockKV(); + }); + + afterEach(() => { + vi.useRealTimers(); + }); + + it("saves at most once per interval however often changes are scheduled", async () => { + const vector = new VectorIndex(); + const persistence = new IndexPersistence(kv as never, vector, { saveIntervalMs: 60_000, bucketSize: 16 }); + + for (let i = 0; i < 50; i++) { + touch(vector, `obs_${i}`); + persistence.scheduleSave(); + } + await vi.advanceTimersByTimeAsync(59_000); + expect(metaSaves(kv)).toBe(0); + + await vi.advanceTimersByTimeAsync(1_000); + expect(metaSaves(kv)).toBe(1); + + for (let i = 0; i < 50; i++) { + touch(vector, `obs_more_${i}`); + persistence.scheduleSave(); + } + await vi.advanceTimersByTimeAsync(30_000); + expect(metaSaves(kv)).toBe(1); + await vi.advanceTimersByTimeAsync(30_000); + expect(metaSaves(kv)).toBe(2); + }); + + it("an explicit save runs immediately and cancels the pending timer", async () => { + const vector = new VectorIndex(); + const persistence = new IndexPersistence(kv as never, vector, { saveIntervalMs: 60_000, bucketSize: 16 }); + + touch(vector, "obs_1"); + persistence.scheduleSave(); + await persistence.save(); + expect(metaSaves(kv)).toBe(1); + + await vi.advanceTimersByTimeAsync(120_000); + expect(metaSaves(kv)).toBe(1); + }); + + it("never runs two saves at once and coalesces requests made during a save into one", async () => { + let release: () => void = () => undefined; + let inFlight = 0; + let maxInFlight = 0; + let gate: Promise | null = new Promise((resolve) => { + release = resolve; + }); + const slowKv = { + ...kv, + set: async (scope: string, key: string, data: T): Promise => { + inFlight++; + maxInFlight = Math.max(maxInFlight, inFlight); + if (gate) await gate; + const result = await kv.set(scope, key, data); + inFlight--; + return result; + }, + }; + const vector = new VectorIndex(); + touch(vector, "obs_1"); + const persistence = new IndexPersistence(slowKv as never, vector, { saveIntervalMs: 60_000, bucketSize: 16 }); + + const first = persistence.save(); + touch(vector, "obs_2"); + const second = persistence.save(); + const third = persistence.save(); + expect(second).toBe(third); + expect(persistence.status().saving).toBe(true); + + gate = null; + release(); + await Promise.all([first, second, third]); + + expect(metaSaves(kv)).toBe(2); + expect(maxInFlight).toBe(1); + expect(persistence.status().saving).toBe(false); + }); + + it("keeps the index dirty when a change arrives while it is being saved", async () => { + let release: () => void = () => undefined; + const gate = new Promise((resolve) => { + release = resolve; + }); + let held = false; + const slowKv = { + ...kv, + set: async (scope: string, key: string, data: T): Promise => { + if (!held) { + held = true; + await gate; + } + return kv.set(scope, key, data); + }, + }; + const vector = new VectorIndex(); + touch(vector, "obs_1"); + const persistence = new IndexPersistence(slowKv as never, vector, { saveIntervalMs: 60_000, bucketSize: 16 }); + + const saving = persistence.save(); + touch(vector, "obs_2"); + persistence.scheduleSave(); + release(); + await saving; + + expect(persistence.status().vector?.lastSavedAt).not.toBeNull(); + expect(persistence.status().vector?.dirtySince).not.toBeNull(); + + await persistence.save(); + expect(persistence.status().vector?.dirtySince).toBeNull(); + }); + + it("stop prevents later scheduled saves", async () => { + const vector = new VectorIndex(); + touch(vector, "obs_1"); + const persistence = new IndexPersistence(kv as never, vector, { saveIntervalMs: 1_000, bucketSize: 16 }); + persistence.stop(); + persistence.scheduleSave(); + await vi.advanceTimersByTimeAsync(10_000); + expect(metaSaves(kv)).toBe(0); + }); +}); + +function statusInputs(overrides: Partial = {}): StatusInputs { + return { + now: new Date("2026-09-25T12:00:00.000Z"), + version: "0.9.29", + engineVersion: "0.22.1", + uptimeSeconds: 100, + ports: { rest: 3111, streams: 3112, viewer: 3113 }, + health: { status: "healthy", alerts: [], notes: [], connectionState: "connected" }, + circuitBreaker: null, + functionMetrics: [], + provider: "llm", + embeddingProvider: "embeddings", + flags: [], + index: { + bm25Documents: 1, + vectorDocuments: 1, + observationsIndexed: 1, + missingObservations: 0, + sessions: 1, + bm25Incomplete: false, + pendingVectorBackfill: 0, + }, + graph: null, + graphExtractionEnabled: false, + ...overrides, + }; +} + +const cleanLeg = { + lastSavedAt: "2026-09-25T11:59:00.000Z", + lastError: null, + lastErrorAt: null, + dirtySince: null, +}; + +function persistenceStatus(vector: typeof cleanLeg | null, pendingChanges = 0) { + return { saveIntervalMs: 600_000, saving: false, buckets: 256, pendingChanges, vector, vectorCountShortfall: null }; +} + +describe("status reports index persistence", () => { + it("is quiet when the vector index saved cleanly", () => { + const report = evaluateStatus(statusInputs({ indexPersistence: persistenceStatus(cleanLeg) })); + expect(report.problems).toEqual([]); + expect(report.indexPersistence?.buckets).toBe(256); + }); + + it("reports a failing vector save as an error", () => { + const report = evaluateStatus( + statusInputs({ + indexPersistence: persistenceStatus({ + ...cleanLeg, + lastError: "3 of 10 vector writes failed: timed out", + lastErrorAt: "2026-09-25T11:59:30.000Z", + } as never), + }), + ); + expect(report.status).toBe("error"); + expect(report.problems.map((p) => p.code)).toEqual(["index-save-failing"]); + expect(report.problems[0].message).toContain("3 of 10 vector writes failed"); + }); + + it("warns when unsaved changes are older than twice the save interval", () => { + const report = evaluateStatus( + statusInputs({ indexPersistence: persistenceStatus({ ...cleanLeg, dirtySince: "2026-09-25T11:30:00.000Z" } as never, 4) }), + ); + expect(report.problems.map((p) => p.code)).toEqual(["index-save-stale"]); + }); + + it("does not report vector persistence when vector search is off", () => { + const report = evaluateStatus(statusInputs({ indexPersistence: persistenceStatus(null) })); + expect(report.problems).toEqual([]); + }); + + it("shows save state and bucket storage on the status page", () => { + const report = evaluateStatus( + statusInputs({ indexPersistence: persistenceStatus({ ...cleanLeg, dirtySince: "2026-09-25T11:59:50.000Z" } as never, 3) }), + ); + const html = renderStatusHtml(report, "n"); + expect(html).toContain(""); + expect(html).toContain(""); + expect(html).toContain(""); + }); +}); diff --git a/test/index-persistence.test.ts b/test/index-persistence.test.ts index 34f72bebe3..568d180006 100644 --- a/test/index-persistence.test.ts +++ b/test/index-persistence.test.ts @@ -1,1172 +1,475 @@ import { describe, it, expect, beforeEach, afterEach, vi } from "vitest"; -import { IndexPersistence } from "../src/state/index-persistence.js"; +import { IndexPersistence, vectorBucketScope } from "../src/state/index-persistence.js"; import { SearchIndex } from "../src/state/search-index.js"; import { VectorIndex } from "../src/state/vector-index.js"; import type { CompressedObservation } from "../src/types.js"; -const BM25_SCOPE = "mem:index:bm25"; -const BM25_LEGACY_KEY = "data"; -const BM25_MANIFEST_KEY = "data:manifest"; -const VECTOR_LEGACY_KEY = "vectors"; -const VECTOR_MANIFEST_KEY = "vectors:manifest"; - -type TestIndexShardManifest = { - v: 1; - generation?: string; - shards: Array<{ scope: string; key: string; chars: number }>; - chars: number; -}; +const INDEX_SCOPE = "mem:index:bm25"; +const META_KEY = "vectors:meta"; function mockKV() { const store = new Map>(); + const ops: Array<{ op: "set" | "delete" | "list"; scope: string; key?: string }> = []; return { - get: async (scope: string, key: string): Promise => { - return (store.get(scope)?.get(key) as T) ?? null; - }, + store, + ops, + get: async (scope: string, key: string): Promise => (store.get(scope)?.get(key) as T) ?? null, set: async (scope: string, key: string, data: T): Promise => { + ops.push({ op: "set", scope, key }); if (!store.has(scope)) store.set(scope, new Map()); store.get(scope)!.set(key, data); return data; }, delete: async (scope: string, key: string): Promise => { + ops.push({ op: "delete", scope, key }); store.get(scope)?.delete(key); }, list: async (scope: string): Promise => { - const entries = store.get(scope); - return entries ? (Array.from(entries.values()) as T[]) : []; + ops.push({ op: "list", scope }); + return Array.from(store.get(scope)?.values() ?? []) as T[]; }, }; } type MockKV = ReturnType; -function makeObs( - overrides: Partial = {}, -): CompressedObservation { +function vec(values: number[]): Float32Array { + return new Float32Array(values); +} + +function vectorWith(entries: Array<[string, number[]]>): VectorIndex { + const index = new VectorIndex(); + for (const [id, values] of entries) index.add(id, `ses_${id}`, vec(values)); + return index; +} + +function bucketSets(kv: MockKV): Array<{ scope: string; key?: string }> { + return kv.ops.filter((o) => o.op === "set" && o.scope.startsWith(`${INDEX_SCOPE}:vec:`)); +} + +function obs(id: string): CompressedObservation { return { - id: "obs_1", + id, sessionId: "ses_1", timestamp: new Date().toISOString(), type: "file_edit", - title: "Edit auth middleware", - subtitle: "JWT validation", - facts: ["Added token check"], - narrative: "Modified the auth middleware to validate JWT tokens", - concepts: ["authentication", "jwt"], - files: ["src/middleware/auth.ts"], - importance: 7, - ...overrides, + title: `title ${id}`, + facts: [], + narrative: `narrative ${id}`, + concepts: [], + files: [], + importance: 5, }; } -function makeBm25(id: string, title: string): SearchIndex { - const bm25 = new SearchIndex(); - bm25.add(makeObs({ id, title, narrative: `${title} narrative` })); - return bm25; +async function writeLegacyVectorSnapshot( + kv: MockKV, + vector: VectorIndex, + options: { generation?: string; monolithic?: boolean; shardChars?: number } = {}, +): Promise { + const serialized = vector.serialize(); + if (options.monolithic) { + await kv.set(INDEX_SCOPE, "vectors", serialized); + return; + } + const generation = options.generation ?? "idx_mfabcd12_aaaaaaaaaaaa"; + const size = options.shardChars ?? 40; + const shards: Array<{ scope: string; key: string; chars: number }> = []; + for (let offset = 0, i = 0; offset < serialized.length; offset += size, i++) { + const scope = `${INDEX_SCOPE}:vectors:${generation}:${String(i).padStart(5, "0")}`; + const chunk = serialized.slice(offset, offset + size); + await kv.set(scope, "data", chunk); + shards.push({ scope, key: "data", chars: chunk.length }); + } + await kv.set(INDEX_SCOPE, "vectors:manifest", { v: 1, generation, shards, chars: serialized.length }); } -function makeVector(id = "obs_1"): VectorIndex { - const vector = new VectorIndex(); - vector.add(id, "ses_1", new Float32Array([0.1, 0.2, 0.3])); - return vector; +async function writeLegacyBm25Snapshot(kv: MockKV): Promise { + const bm25 = new SearchIndex(); + bm25.add(obs("obs_legacy")); + const serialized = bm25.serialize(); + const scope = `${INDEX_SCOPE}:bm25:idx_mfabcd12_bbbbbbbbbbbb:00000`; + await kv.set(scope, "data", serialized); + await kv.set(INDEX_SCOPE, "data:manifest", { + v: 1, + generation: "idx_mfabcd12_bbbbbbbbbbbb", + shards: [{ scope, key: "data", chars: serialized.length }], + chars: serialized.length, + }); } -async function getBm25Manifest(kv: MockKV): Promise { - const manifest = await kv.get( - BM25_SCOPE, - BM25_MANIFEST_KEY, - ); - expect(manifest).not.toBeNull(); - return manifest!; +function expectSameVectors(actual: VectorIndex | null, expected: VectorIndex): void { + expect(actual).not.toBeNull(); + expect(actual!.size).toBe(expected.size); + for (const [id, entry] of expected.entries()) { + const loaded = actual!.get(id); + expect(loaded?.sessionId).toBe(entry.sessionId); + expect(Array.from(loaded!.embedding)).toEqual(Array.from(entry.embedding)); + } } -describe("IndexPersistence", () => { - let kv: ReturnType; +describe("IndexPersistence bucketed vector storage", () => { + let kv: MockKV; beforeEach(() => { - vi.useFakeTimers(); kv = mockKV(); }); - afterEach(() => { - vi.useRealTimers(); - }); - - it("saves and loads BM25 index round-trip", async () => { - const bm25 = new SearchIndex(); - bm25.add(makeObs({ id: "obs_1", title: "auth handler" })); - - const persistence = new IndexPersistence(kv as never, bm25, null); + it("fills buckets in insertion order and rolls to a new one at the cap", async () => { + const vector = vectorWith([ + ["obs_a", [0.1, 0.2, 0.3]], + ["obs_b", [0.4, 0.5, 0.6]], + ["obs_c", [0.7, 0.8, 0.9]], + ]); + await new IndexPersistence(kv as never, vector, { bucketSize: 2 }).save(); + + expect(kv.store.get(vectorBucketScope(0))?.has("obs_a")).toBe(true); + expect(kv.store.get(vectorBucketScope(0))?.has("obs_b")).toBe(true); + expect(kv.store.get(vectorBucketScope(1))?.has("obs_c")).toBe(true); + const meta = await kv.get<{ v: number; bucketCount: number; count: number }>(INDEX_SCOPE, META_KEY); + expect(meta).toMatchObject({ v: 3, bucketCount: 2, count: 3 }); + + const loaded = await new IndexPersistence(kv as never, new VectorIndex(), { bucketSize: 2 }).load(); + expect(loaded.state).toBe("buckets"); + expectSameVectors(loaded.vector, vector); + expect(loaded.vector!.pendingChanges).toBe(0); + }); + + it("writes only the bucket entry of a single added vector", async () => { + const vector = vectorWith([ + ["obs_a", [0.1, 0.2, 0.3]], + ["obs_b", [0.4, 0.5, 0.6]], + ]); + const persistence = new IndexPersistence(kv as never, vector, { bucketSize: 16 }); await persistence.save(); + kv.ops.length = 0; - const loaded = await persistence.load(); - expect(loaded.bm25).not.toBeNull(); - expect(loaded.bm25!.size).toBe(1); - const results = loaded.bm25!.search("auth"); - expect(results.length).toBe(1); - }); - - it("saves BM25 index shards outside the BM25 metadata scope", async () => { - const bm25 = new SearchIndex(); - bm25.add( - makeObs({ - id: "obs_1", - title: "auth handler ".repeat(40), - narrative: "JWT middleware validation ".repeat(40), - }), - ); - - const persistence = new IndexPersistence(kv as never, bm25, null, { - shardChars: 80, - createGeneration: () => "gen_bm25", - }); + vector.add("obs_new", "ses_new", vec([1, 0, 0])); await persistence.save(); - const manifest = await getBm25Manifest(kv); - expect(manifest.generation).toBe("gen_bm25"); - expect(manifest.shards.length).toBeGreaterThan(1); - expect(manifest.shards[0].scope).toContain(":gen_bm25:"); - await expect(kv.get(BM25_SCOPE, BM25_LEGACY_KEY)).resolves.toBeNull(); - await expect( - kv.get(manifest.shards[0].scope, manifest.shards[0].key), - ).resolves.toEqual(expect.any(String)); - - const loaded = await persistence.load(); - expect(loaded.bm25).not.toBeNull(); - expect(loaded.bm25!.search("auth").length).toBe(1); - }); - - it("loads legacy monolithic BM25 and vector snapshots", async () => { - const bm25 = makeBm25("obs_1", "legacy auth handler"); - const vector = makeVector("obs_1"); - await kv.set(BM25_SCOPE, BM25_LEGACY_KEY, bm25.serialize()); - await kv.set(BM25_SCOPE, VECTOR_LEGACY_KEY, vector.serialize()); - - const loaded = await new IndexPersistence( - kv as never, - new SearchIndex(), - null, - ).load(); - - expect(loaded.bm25).not.toBeNull(); - expect(loaded.bm25!.search("legacy").length).toBe(1); - expect(loaded.vector).not.toBeNull(); - expect(loaded.vector!.size).toBe(1); + const writes = bucketSets(kv); + expect(writes).toEqual([{ op: "set", scope: vectorBucketScope(0), key: "obs_new" }]); }); - it("fails closed instead of falling back when manifest reads fail", async () => { - const legacy = makeBm25("obs_legacy", "legacy stale snapshot"); - await kv.set(BM25_SCOPE, BM25_LEGACY_KEY, legacy.serialize()); - const failingKv = { - ...kv, - get: vi.fn(async (scope: string, key: string): Promise => { - if (scope === BM25_SCOPE && key === BM25_MANIFEST_KEY) { - throw new Error("manifest backend unavailable"); - } - return kv.get(scope, key); - }), - }; - - const loaded = await new IndexPersistence( - failingKv as never, - new SearchIndex(), - null, - ).load(); - - expect(loaded.bm25).toBeNull(); - }); - - it("fails closed when legacy snapshot reads fail", async () => { - const failingKv = { - ...kv, - get: vi.fn(async (scope: string, key: string): Promise => { - if (scope === BM25_SCOPE && key === BM25_LEGACY_KEY) { - throw new Error("legacy backend unavailable"); - } - return kv.get(scope, key); - }), - }; - - const loaded = await new IndexPersistence( - failingKv as never, - new SearchIndex(), - null, - ).load(); + it("does not write when nothing changed", async () => { + const vector = vectorWith([["obs_a", [0.1, 0.2, 0.3]]]); + const persistence = new IndexPersistence(kv as never, vector, { bucketSize: 16 }); + await persistence.save(); + kv.ops.length = 0; - expect(loaded.bm25).toBeNull(); - }); + await persistence.save(); - it("loads sharded manifests that omit optional generation metadata", async () => { - const bm25 = makeBm25("obs_1", "deterministic shard auth"); - const serialized = bm25.serialize(); - const chunks = [serialized.slice(0, 50), serialized.slice(50)]; - await kv.set("mem:index:bm25:bm25:00000", "data", chunks[0]); - await kv.set("mem:index:bm25:bm25:00001", "data", chunks[1]); - await kv.set(BM25_SCOPE, BM25_MANIFEST_KEY, { - v: 1, - chars: serialized.length, - shards: [ - { - scope: "mem:index:bm25:bm25:00000", - key: "data", - chars: chunks[0].length, - }, - { - scope: "mem:index:bm25:bm25:00001", - key: "data", - chars: chunks[1].length, - }, - ], - }); - - const loaded = await new IndexPersistence( - kv as never, - new SearchIndex(), - null, - ).load(); - - expect(loaded.bm25).not.toBeNull(); - expect(loaded.bm25!.search("deterministic").length).toBe(1); + expect(kv.ops.filter((o) => o.op !== "list")).toEqual([]); }); - it("saves and loads vector index round-trip", async () => { - const bm25 = new SearchIndex(); - const vector = makeVector(); + it("removes deleted vectors from their bucket", async () => { + const vector = vectorWith([ + ["obs_a", [0.1, 0.2, 0.3]], + ["obs_b", [0.4, 0.5, 0.6]], + ]); + const persistence = new IndexPersistence(kv as never, vector, { bucketSize: 16 }); + await persistence.save(); - const persistence = new IndexPersistence(kv as never, bm25, vector); + vector.remove("obs_a"); await persistence.save(); - const loaded = await persistence.load(); - expect(loaded.vector).not.toBeNull(); - expect(loaded.vector!.size).toBe(1); + expect(kv.ops).toContainEqual({ op: "delete", scope: vectorBucketScope(0), key: "obs_a" }); + const loaded = await new IndexPersistence(kv as never, new VectorIndex(), { bucketSize: 16 }).load(); + expect(loaded.vector!.has("obs_a")).toBe(false); + expect(loaded.vector!.has("obs_b")).toBe(true); }); - it("saves vector index shards outside the BM25 scope", async () => { - const bm25 = new SearchIndex(); - const vector = new VectorIndex(); - vector.add( - "obs_1", - "ses_1", - new Float32Array(Array.from({ length: 32 }, (_, i) => i)), + it("a delete touches only the bucket that owns the id, in a multi-bucket store", async () => { + const vector = vectorWith( + Array.from({ length: 5 }, (_, i) => [`obs_${i}`, [i, i + 1, i + 2]] as [string, number[]]), ); - - const persistence = new IndexPersistence(kv as never, bm25, vector, { - shardChars: 40, - createGeneration: () => "gen_vector", - }); + const persistence = new IndexPersistence(kv as never, vector, { bucketSize: 2 }); await persistence.save(); + kv.ops.length = 0; - const manifest = await kv.get( - BM25_SCOPE, - VECTOR_MANIFEST_KEY, - ); - expect(manifest).not.toBeNull(); - expect(manifest!.generation).toBe("gen_vector"); - expect(manifest!.shards.length).toBeGreaterThan(1); - expect(manifest!.shards[0].scope).toContain(":gen_vector:"); - await expect(kv.get(BM25_SCOPE, VECTOR_LEGACY_KEY)).resolves.toBeNull(); - await expect( - kv.get(manifest!.shards[0].scope, manifest!.shards[0].key), - ).resolves.toEqual(expect.any(String)); + vector.remove("obs_2"); + await persistence.save(); - const loaded = await persistence.load(); - expect(loaded.vector).not.toBeNull(); - expect(loaded.vector!.size).toBe(1); + const touched = new Set(kv.ops.filter((o) => o.scope.startsWith(`${INDEX_SCOPE}:vec:`)).map((o) => o.scope)); + expect(touched).toEqual(new Set([vectorBucketScope(1)])); + expect(kv.ops).toContainEqual({ op: "delete", scope: vectorBucketScope(1), key: "obs_2" }); }); - it("persists empty vector snapshots so cleared vectors do not reload", async () => { - const previousBm25 = makeBm25("obs_old", "alpha previous snapshot"); - const previousVector = makeVector("obs_old"); - await new IndexPersistence(kv as never, previousBm25, previousVector, { - shardChars: 80, - createGeneration: () => "gen_old", - }).save(); - - const nextBm25 = makeBm25("obs_new", "bravo new snapshot"); - const emptyVector = new VectorIndex(); - await new IndexPersistence(kv as never, nextBm25, emptyVector, { - shardChars: 80, - createGeneration: () => "gen_empty", - }).save(); - - const vectorManifest = await kv.get( - BM25_SCOPE, - VECTOR_MANIFEST_KEY, - ); - expect(vectorManifest).not.toBeNull(); - expect(vectorManifest!.generation).toBe("gen_empty"); - const loaded = await new IndexPersistence( - kv as never, - new SearchIndex(), - null, - ).load(); - expect(loaded.bm25!.search("bravo").length).toBe(1); - expect(loaded.vector).not.toBeNull(); - expect(loaded.vector!.size).toBe(0); - }); + it("clearing the index deletes every persisted vector", async () => { + const vector = vectorWith([ + ["obs_a", [0.1, 0.2, 0.3]], + ["obs_b", [0.4, 0.5, 0.6]], + ]); + const persistence = new IndexPersistence(kv as never, vector, { bucketSize: 16 }); + await persistence.save(); - it("avoids one oversized state::set string payload for persisted indexes", async () => { - const maxStringPayloadChars = 80; - const bm25 = new SearchIndex(); - bm25.add( - makeObs({ - id: "obs_1", - title: "large persisted snapshot ".repeat(40), - narrative: "oversized state set reproduction ".repeat(40), - }), - ); - const vector = new VectorIndex(); - vector.add( - "obs_1", - "ses_1", - new Float32Array(Array.from({ length: 64 }, (_, i) => i / 10)), - ); - const guardedKv = { - ...kv, - set: vi.fn(async (scope: string, key: string, data: T): Promise => { - if ( - typeof data === "string" && - data.length > maxStringPayloadChars - ) { - throw new Error(`oversized state::set payload: ${scope}/${key}`); - } - return kv.set(scope, key, data); - }), - }; + vector.clear(); + vector.add("obs_c", "ses_c", vec([0, 1, 0])); + await persistence.save(); - await new IndexPersistence(guardedKv as never, bm25, vector, { - shardChars: maxStringPayloadChars, - createGeneration: () => "gen_payload_limit", - }).save(); - - const loaded = await new IndexPersistence( - kv as never, - new SearchIndex(), - null, - ).load(); - expect(loaded.bm25!.search("oversized").length).toBe(1); - expect(loaded.vector!.size).toBe(1); + const loaded = await new IndexPersistence(kv as never, new VectorIndex(), { bucketSize: 16 }).load(); + expect([...loaded.vector!.entries()].map(([id]) => id)).toEqual(["obs_c"]); }); - it("falls back to the default shard size for fractional values below one", async () => { - const bm25 = makeBm25("obs_fraction", "fractional shard config"); - let newShardWrites = 0; - const guardedKv = { + it("keeps failed writes pending and retries them on the next save", async () => { + const vector = vectorWith([["obs_a", [0.1, 0.2, 0.3]]]); + let fail = true; + const flakyKv = { ...kv, - set: vi.fn(async (scope: string, key: string, data: T): Promise => { - if (scope.includes(":gen_fraction:")) { - newShardWrites += 1; - if (newShardWrites > 3) { - throw new Error("fractional shard size caused zero-width shards"); - } - } + set: async (scope: string, key: string, data: T): Promise => { + if (fail && scope.startsWith(`${INDEX_SCOPE}:vec:`)) throw new Error("state::set timed out"); return kv.set(scope, key, data); - }), + }, }; + const persistence = new IndexPersistence(flakyKv as never, vector, { bucketSize: 16 }); - await new IndexPersistence(guardedKv as never, bm25, null, { - shardChars: 0.5, - createGeneration: () => "gen_fraction", - }).save(); + await expect(persistence.save()).resolves.toBeUndefined(); + const failed = persistence.status(); + expect(failed.vector?.lastError).toBe("1 of 1 vector writes failed: state::set timed out"); + expect(failed.pendingChanges).toBe(1); + expect(await kv.get(INDEX_SCOPE, META_KEY)).toBeNull(); - const manifest = await getBm25Manifest(kv); - expect(manifest.generation).toBe("gen_fraction"); - expect(manifest.shards.length).toBe(1); - expect(newShardWrites).toBe(1); + fail = false; + await persistence.save(); + expect(persistence.status().vector?.lastError).toBeNull(); + expect(persistence.status().pendingChanges).toBe(0); + const loaded = await new IndexPersistence(kv as never, new VectorIndex(), { bucketSize: 16 }).load(); + expectSameVectors(loaded.vector, vector); }); - it("keeps the previous generation when a shard write fails before manifest commit", async () => { - const previous = makeBm25("obs_old", "alpha previous snapshot"); - await new IndexPersistence(kv as never, previous, null, { - shardChars: 80, - createGeneration: () => "gen_old", - }).save(); - const previousManifest = await getBm25Manifest(kv); + it("a save after many vectors added over several waves touches only the currently open bucket", async () => { + const vector = new VectorIndex(); + const persistence = new IndexPersistence(kv as never, vector, { bucketSize: 100 }); - let newShardWrites = 0; - const failingKv = { - ...kv, - set: vi.fn(async (scope: string, key: string, data: T): Promise => { - if (scope.includes(":gen_new:")) { - newShardWrites += 1; - if (newShardWrites === 2) throw new Error("shard write failed"); - } - return kv.set(scope, key, data); - }), - }; + for (let i = 0; i < 950; i++) vector.add(`obs_${i}`, "ses_1", vec([i, 0, 0])); + await persistence.save(); + expect(kv.store.get(vectorBucketScope(9))?.size).toBe(50); + for (let bucket = 0; bucket <= 8; bucket++) expect(kv.store.get(vectorBucketScope(bucket))?.size).toBe(100); + kv.ops.length = 0; - const next = makeBm25("obs_new", "bravo new snapshot"); - await new IndexPersistence(failingKv as never, next, null, { - shardChars: 80, - createGeneration: () => "gen_new", - }).save(); + for (let i = 950; i < 1000; i++) vector.add(`obs_${i}`, "ses_1", vec([i, 0, 0])); + await persistence.save(); + let touched = new Set(kv.ops.filter((o) => o.scope.startsWith(`${INDEX_SCOPE}:vec:`)).map((o) => o.scope)); + expect(touched).toEqual(new Set([vectorBucketScope(9)])); + expect(kv.store.get(vectorBucketScope(9))?.size).toBe(100); + kv.ops.length = 0; - await expect(kv.get(BM25_SCOPE, BM25_MANIFEST_KEY)).resolves.toEqual( - previousManifest, - ); - await expect( - kv.get("mem:index:bm25:bm25:gen_new:00000", "data"), - ).resolves.toBeNull(); - const loaded = await new IndexPersistence( - kv as never, - new SearchIndex(), - null, - ).load(); - expect(loaded.bm25!.search("alpha").length).toBe(1); - expect(loaded.bm25!.search("bravo").length).toBe(0); + vector.add("obs_1000", "ses_1", vec([1000, 0, 0])); + await persistence.save(); + touched = new Set(kv.ops.filter((o) => o.scope.startsWith(`${INDEX_SCOPE}:vec:`)).map((o) => o.scope)); + expect(touched).toEqual(new Set([vectorBucketScope(10)])); }); - it("keeps the previous generation when manifest set rejects before commit", async () => { - const previous = makeBm25("obs_old", "alpha previous snapshot"); - await new IndexPersistence(kv as never, previous, null, { - shardChars: 80, - createGeneration: () => "gen_old", - }).save(); - const previousManifest = await getBm25Manifest(kv); - - const failingKv = { - ...kv, - set: vi.fn(async (scope: string, key: string, data: T): Promise => { - if (scope === BM25_SCOPE && key === BM25_MANIFEST_KEY) { - throw new Error("manifest write failed"); - } - return kv.set(scope, key, data); - }), - }; - - const next = makeBm25("obs_new", "bravo new snapshot"); - await new IndexPersistence(failingKv as never, next, null, { - shardChars: 80, - createGeneration: () => "gen_new", - }).save(); - - await expect(kv.get(BM25_SCOPE, BM25_MANIFEST_KEY)).resolves.toEqual( - previousManifest, + it("a smaller bucket size configured later opens a new bucket instead of moving existing data", async () => { + const vector = vectorWith( + Array.from({ length: 4 }, (_, i) => [`obs_${i}`, [i, i + 1, i + 2]] as [string, number[]]), ); - await expect( - kv.get("mem:index:bm25:bm25:gen_new:00000", "data"), - ).resolves.toBeNull(); - const loaded = await new IndexPersistence( - kv as never, - new SearchIndex(), - null, - ).load(); - expect(loaded.bm25!.search("alpha").length).toBe(1); - expect(loaded.bm25!.search("bravo").length).toBe(0); - }); + await new IndexPersistence(kv as never, vector, { bucketSize: 4 }).save(); + for (let i = 0; i < 4; i++) expect(kv.store.get(vectorBucketScope(0))?.has(`obs_${i}`)).toBe(true); - it("keeps a generation loadable when manifest set commits before rejecting", async () => { - const previous = makeBm25("obs_old", "alpha previous snapshot"); - await new IndexPersistence(kv as never, previous, null, { - shardChars: 80, - createGeneration: () => "gen_old", - }).save(); + const reloaded = new VectorIndex(); + const persistence = new IndexPersistence(kv as never, reloaded, { bucketSize: 2 }); + const loaded = await persistence.load(); + reloaded.restoreFrom(loaded.vector!); + kv.ops.length = 0; - const failingKv = { - ...kv, - set: vi.fn(async (scope: string, key: string, data: T): Promise => { - if (scope === BM25_SCOPE && key === BM25_MANIFEST_KEY) { - await kv.set(scope, key, data); - throw new Error("manifest write timed out after commit"); - } - return kv.set(scope, key, data); - }), - }; + reloaded.add("obs_new", "ses_new", vec([9, 9, 9])); + await persistence.save(); - const next = makeBm25("obs_new", "bravo new snapshot"); - await new IndexPersistence(failingKv as never, next, null, { - shardChars: 80, - createGeneration: () => "gen_new", - }).save(); - - const manifest = await getBm25Manifest(kv); - expect(manifest.generation).toBe("gen_new"); - await expect( - kv.get("mem:index:bm25:bm25:gen_new:00000", "data"), - ).resolves.toEqual(expect.any(String)); - const loaded = await new IndexPersistence( - kv as never, - new SearchIndex(), - null, - ).load(); - expect(loaded.bm25!.search("bravo").length).toBe(1); + const touched = new Set(kv.ops.filter((o) => o.scope.startsWith(`${INDEX_SCOPE}:vec:`)).map((o) => o.scope)); + expect(touched).toEqual(new Set([vectorBucketScope(1)])); + for (let i = 0; i < 4; i++) expect(kv.store.get(vectorBucketScope(0))?.has(`obs_${i}`)).toBe(true); }); - it("deletes a shard that committed before set rejected", async () => { - const previous = makeBm25("obs_old", "alpha previous snapshot"); - await new IndexPersistence(kv as never, previous, null, { - shardChars: 80, - createGeneration: () => "gen_old", - }).save(); - const previousManifest = await getBm25Manifest(kv); - - const failingKv = { - ...kv, - set: vi.fn(async (scope: string, key: string, data: T): Promise => { - if (scope === "mem:index:bm25:bm25:gen_new:00000") { - await kv.set(scope, key, data); - throw new Error("state::set timed out after commit"); - } - return kv.set(scope, key, data); - }), - }; - - const next = makeBm25("obs_new", "bravo new snapshot"); - await new IndexPersistence(failingKv as never, next, null, { - shardChars: 80, - createGeneration: () => "gen_new", - }).save(); - - await expect(kv.get(BM25_SCOPE, BM25_MANIFEST_KEY)).resolves.toEqual( - previousManifest, + it("detects a partial bucket set against the saved count and reports it in status", async () => { + const vector = vectorWith( + Array.from({ length: 4 }, (_, i) => [`obs_${i}`, [i, i + 1, i + 2]] as [string, number[]]), ); - await expect( - kv.get("mem:index:bm25:bm25:gen_new:00000", "data"), - ).resolves.toBeNull(); - const loaded = await new IndexPersistence( - kv as never, - new SearchIndex(), - null, - ).load(); - expect(loaded.bm25!.search("alpha").length).toBe(1); - expect(loaded.bm25!.search("bravo").length).toBe(0); - }); + await new IndexPersistence(kv as never, vector, { bucketSize: 2 }).save(); + kv.store.get(vectorBucketScope(1))?.clear(); - it("loads the new generation even when old generation cleanup fails", async () => { - const previous = makeBm25("obs_old", "alpha previous snapshot"); - await new IndexPersistence(kv as never, previous, null, { - shardChars: 80, - createGeneration: () => "gen_old", - }).save(); - - const cleanupKv = { - ...kv, - delete: vi.fn(async () => { - throw new Error("cleanup failed"); - }), - }; - const next = makeBm25("obs_new", "bravo new snapshot"); - await new IndexPersistence(cleanupKv as never, next, null, { - shardChars: 80, - createGeneration: () => "gen_new", - }).save(); - - const manifest = await getBm25Manifest(kv); - expect(manifest.generation).toBe("gen_new"); - expect(cleanupKv.delete).toHaveBeenCalled(); - const loaded = await new IndexPersistence( - kv as never, - new SearchIndex(), - null, - ).load(); - expect(loaded.bm25!.search("bravo").length).toBe(1); - expect(loaded.bm25!.search("alpha").length).toBe(0); - }); - - it("reclaims the previous generation when the previous manifest read fails (#1115)", async () => { - const previous = makeBm25("obs_old", "alpha previous snapshot"); - await new IndexPersistence(kv as never, previous, null, { - shardChars: 80, - createGeneration: () => "gen_old", - }).save(); - const oldShardScope = "mem:index:bm25:bm25:gen_old:00000"; - await expect(kv.get(oldShardScope, "data")).resolves.not.toBeNull(); - - // The manifest read that opens saveShardedIndex times out. It used to be - // swallowed into `previous = null`, which skipped the cleanup guard and - // stranded gen_old's shards with nothing left to ever revisit them. - const readFailsKv = { - ...kv, - get: vi.fn(async (scope: string, key: string): Promise => { - if (scope === BM25_SCOPE && key === BM25_MANIFEST_KEY) { - throw new Error("Invocation timeout after 180000ms: state::get"); - } - return kv.get(scope, key); - }), - }; - - const next = makeBm25("obs_new", "bravo new snapshot"); - await new IndexPersistence(readFailsKv as never, next, null, { - shardChars: 80, - createGeneration: () => "gen_new", - }).save(); + const persistence = new IndexPersistence(kv as never, new VectorIndex(), { bucketSize: 2 }); + const loaded = await persistence.load(); - const manifest = await getBm25Manifest(kv); - expect(manifest.generation).toBe("gen_new"); - await expect(kv.get(oldShardScope, "data")).resolves.toBeNull(); + expect(loaded.state).toBe("buckets"); + expect(loaded.vector!.size).toBe(2); + expect(loaded.expectedCount).toBe(4); + expect(persistence.status().vectorCountShortfall).toEqual({ expected: 4, loaded: 2 }); }); - it("reclaims a generation stranded by a failed cleanup on the next load (#1115)", async () => { - const previous = makeBm25("obs_old", "alpha previous snapshot"); - await new IndexPersistence(kv as never, previous, null, { - shardChars: 80, - createGeneration: () => "gen_old", - }).save(); - const oldShardScope = "mem:index:bm25:bm25:gen_old:00000"; - - const cleanupKv = { + it("reports storage as unavailable when the metadata read fails", async () => { + const failingKv = { ...kv, - delete: vi.fn(async () => { - throw new Error("cleanup failed"); - }), + get: async () => { + throw new Error("engine down"); + }, }; - const next = makeBm25("obs_new", "bravo new snapshot"); - await new IndexPersistence(cleanupKv as never, next, null, { - shardChars: 80, - createGeneration: () => "gen_new", - }).save(); - - // Cleanup failed, so gen_old is still on disk. Nothing in the save path - // will revisit it — a later save only ever inspects its own predecessor. - await expect(kv.get(oldShardScope, "data")).resolves.not.toBeNull(); - - const loaded = await new IndexPersistence( - kv as never, - new SearchIndex(), - null, - ).load(); - expect(loaded.bm25!.search("bravo").length).toBe(1); - await expect(kv.get(oldShardScope, "data")).resolves.toBeNull(); + const loaded = await new IndexPersistence(failingKv as never, new VectorIndex(), { bucketSize: 16 }).load(); + expect(loaded).toEqual({ vector: null, state: "unavailable", savedAt: null }); }); - it("reclaims the previous vector generation when the vector manifest read fails (#1115)", async () => { - await new IndexPersistence( - kv as never, - makeBm25("obs_old", "alpha previous snapshot"), - makeVector("obs_old"), - { shardChars: 80, createGeneration: () => "gen_old" }, - ).save(); - const oldVectorScope = "mem:index:bm25:vectors:gen_old:00000"; - await expect(kv.get(oldVectorScope, "data")).resolves.not.toBeNull(); - - const readFailsKv = { - ...kv, - get: vi.fn(async (scope: string, key: string): Promise => { - if (scope === BM25_SCOPE && key === VECTOR_MANIFEST_KEY) { - throw new Error("Invocation timeout after 180000ms: state::get"); - } - return kv.get(scope, key); - }), - }; - - await new IndexPersistence( - readFailsKv as never, - makeBm25("obs_new", "bravo new snapshot"), - makeVector("obs_new"), - { shardChars: 80, createGeneration: () => "gen_new" }, - ).save(); - - await expect(kv.get(oldVectorScope, "data")).resolves.toBeNull(); + it("treats a missing store as empty, including adapters that return undefined", async () => { + const undefinedKv = { ...kv, get: async () => undefined }; + const loaded = await new IndexPersistence(undefinedKv as never, new VectorIndex(), { bucketSize: 16 }).load(); + expect(loaded).toEqual({ vector: null, state: "none", savedAt: null }); }); +}); - it("leaves a generation recorded after the live one untouched on load (#1115)", async () => { - await new IndexPersistence( - kv as never, - makeBm25("obs_old", "alpha previous snapshot"), - null, - { shardChars: 80, createGeneration: () => "gen_live" }, - ).save(); - - // A concurrent save has recorded its generation and is mid-write, but has - // not published its manifest yet. setIndexPersistence runs before load() - // in src/index.ts, so a request arriving during boot produces exactly this. - const inflightScope = "mem:index:bm25:bm25:gen_inflight:00000"; - await kv.set(inflightScope, "data", "partial shard"); - const gcKey = `${BM25_MANIFEST_KEY}:gc`; - const ledger = await kv.get<{ - v: 1; - generations: Array<{ - generation: string; - shards: Array<{ scope: string; key: string }>; - }>; - }>(BM25_SCOPE, gcKey); - ledger!.generations.push({ - generation: "gen_inflight", - shards: [{ scope: inflightScope, key: "data" }], - }); - await kv.set(BM25_SCOPE, gcKey, ledger); - - await new IndexPersistence(kv as never, new SearchIndex(), null).load(); - - // Deleting it would leave the in-flight save publishing a manifest whose - // shards are already gone, which fails closed on the next load. - await expect(kv.get(inflightScope, "data")).resolves.toBe("partial shard"); - }); +describe("IndexPersistence migration from the single-string format", () => { + let kv: MockKV; - it("reclaims a pre-ledger manifest that carries no generation (#1115)", async () => { - const legacyScope = "mem:index:bm25:bm25:gen_legacy:00000"; - await kv.set(legacyScope, "data", "legacy shard"); - await kv.set(BM25_SCOPE, BM25_MANIFEST_KEY, { - v: 1, - shards: [{ scope: legacyScope, key: "data", chars: 12 }], - chars: 12, - }); - - await new IndexPersistence( - kv as never, - makeBm25("obs_new", "bravo new snapshot"), - null, - { shardChars: 80, createGeneration: () => "gen_new" }, - ).save(); - - await expect(kv.get(legacyScope, "data")).resolves.toBeNull(); + beforeEach(() => { + kv = mockKV(); }); - it("keeps a published manifest whole when two saves overlap (#1115)", async () => { - // scheduleSave fires save() unawaited from a timer while flushIndexSave - // awaits save() on every delete path, so two saves on ONE instance is the - // normal shape, not a contrivance. Without the queue, whichever published - // second reclaimed the other's shards and left it naming data already gone. - vi.useRealTimers(); - const store = new Map>(); - const slowKv = { - get: async (scope: string, key: string): Promise => - (store.get(scope)?.get(key) as T) ?? null, - set: async (scope: string, key: string, data: T): Promise => { - // Stall one of the first generation's shard writes so the second save - // overtakes it. - if (scope.includes(":gen_0:") && scope.endsWith("00005")) { - await new Promise((resolve) => setTimeout(resolve, 200)); - } - if (!store.has(scope)) store.set(scope, new Map()); - store.get(scope)!.set(key, data); - return data; - }, - delete: async (scope: string, key: string): Promise => { - store.get(scope)?.delete(key); - }, - list: async (scope: string): Promise => - Array.from((store.get(scope)?.values() ?? []) as Iterable), - }; - - const bm25 = new SearchIndex(); - for (let i = 0; i < 30; i++) { - bm25.add( - makeObs({ id: `obs_${i}`, title: `lorem ipsum dolor sit amet ${i}` }), - ); + it("migrates a sharded vector snapshot into buckets and removes the old shards", async () => { + const legacy = vectorWith([ + ["obs_a", [0.1, 0.2, 0.3]], + ["obs_b", [0.4, 0.5, 0.6]], + ["obs_c", [0.7, 0.8, 0.9]], + ]); + await writeLegacyVectorSnapshot(kv, legacy, { generation: "idx_mfabcd12_aaaaaaaaaaaa" }); + await writeLegacyBm25Snapshot(kv); + + const loaded = await new IndexPersistence(kv as never, new VectorIndex(), { bucketSize: 16 }).load(); + + expect(loaded.state).toBe("migrated"); + expect(loaded.savedAt).toBe(new Date(parseInt("mfabcd12", 36)).toISOString()); + expectSameVectors(loaded.vector, legacy); + expect(loaded.vector!.pendingChanges).toBe(0); + expect(await kv.get(INDEX_SCOPE, "vectors:manifest")).toBeNull(); + expect(await kv.get(INDEX_SCOPE, "data:manifest")).toBeNull(); + for (const [scope, entries] of kv.store) { + if (scope.includes(":vectors:") || scope.includes(":bm25:idx_")) expect(entries.size).toBe(0); } - let generation = 0; - const persistence = new IndexPersistence(slowKv as never, bm25, null, { - shardChars: 400, - createGeneration: () => `gen_${generation++}`, - }); - - // Let the first save get into its shard writes before the second starts, - // so the second is the one that publishes. - const first = persistence.save(); - await new Promise((resolve) => setTimeout(resolve, 20)); - const second = persistence.save(); - await Promise.all([first, second]); - - const manifest = await slowKv.get( - BM25_SCOPE, - BM25_MANIFEST_KEY, - ); - const missing: string[] = []; - for (const shard of manifest!.shards) { - if ((await slowKv.get(shard.scope, shard.key)) === null) { - missing.push(shard.scope); - } - } - expect(missing).toEqual([]); - - const loaded = await new IndexPersistence( - slowKv as never, - new SearchIndex(), - null, - ).load(); - expect(loaded.bm25).not.toBeNull(); - }); - it("still persists the index when the gc ledger is unreadable (#1115)", async () => { - // state::get timing out is the condition this bug appears under. Aborting - // the save there would stop persisting entirely, which is worse than the - // leak: a lost BM25 index costs a full-corpus rebuild. - const failingKv = { - ...kv, - get: vi.fn(async (scope: string, key: string): Promise => { - if (scope === BM25_SCOPE && key === `${BM25_MANIFEST_KEY}:gc`) { - throw new Error("Invocation timeout after 180000ms: state::get"); - } - return kv.get(scope, key); - }), - }; - - await new IndexPersistence( - failingKv as never, - makeBm25("obs_new", "bravo new snapshot"), - null, - { shardChars: 80, createGeneration: () => "gen_new" }, - ).save(); - - const manifest = await getBm25Manifest(kv); - expect(manifest.generation).toBe("gen_new"); - await expect( - kv.get(manifest.shards[0].scope, manifest.shards[0].key), - ).resolves.not.toBeNull(); - }); - - it("still persists when the gc ledger holds a malformed entry (#1115)", async () => { - // A null entry used to throw a TypeError out of saveShardedIndex before any - // shard write, leaving one throttled log line per 60s while BM25 silently - // stopped persisting for good. - await kv.set(BM25_SCOPE, `${BM25_MANIFEST_KEY}:gc`, { - v: 1, - generations: [null], - }); - - await new IndexPersistence( - kv as never, - makeBm25("obs_new", "bravo new snapshot"), - null, - { shardChars: 80, createGeneration: () => "gen_new" }, - ).save(); - - const manifest = await getBm25Manifest(kv); - expect(manifest.generation).toBe("gen_new"); - await expect( - kv.get(manifest.shards[0].scope, manifest.shards[0].key), - ).resolves.not.toBeNull(); + const reloaded = await new IndexPersistence(kv as never, new VectorIndex(), { bucketSize: 16 }).load(); + expect(reloaded.state).toBe("buckets"); + expectSameVectors(reloaded.vector, legacy); }); - it("still reclaims when the gc ledger holds a malformed entry (#1115)", async () => { - // Publishing is not enough. A malformed entry that silently disabled - // reclaim would leak a generation every cycle while this test stayed green. - await new IndexPersistence(kv as never, makeBm25("obs_old", "alpha"), null, { - shardChars: 80, - createGeneration: () => "gen_old", - }).save(); - const oldShardScope = "mem:index:bm25:bm25:gen_old:00000"; - - const gcKey = `${BM25_MANIFEST_KEY}:gc`; - const ledger = await kv.get<{ v: 1; generations: unknown[] }>( - BM25_SCOPE, - gcKey, - ); - ledger!.generations.unshift(null); - await kv.set(BM25_SCOPE, gcKey, ledger); + it("removes the old BM25 shards when missing keys read back as undefined, as the engine SDK returns them", async () => { + const baseGet = kv.get; + kv.get = (async (scope: string, key: string) => { + const value = await baseGet(scope, key); + return value === null ? undefined : value; + }) as typeof kv.get; + const legacy = vectorWith([["obs_a", [0.1, 0.2, 0.3]]]); + await writeLegacyVectorSnapshot(kv, legacy, { generation: "idx_mfabcd12_bbbbbbbbbbbb" }); + await writeLegacyBm25Snapshot(kv); - await new IndexPersistence(kv as never, makeBm25("obs_new", "bravo"), null, { - shardChars: 80, - createGeneration: () => "gen_new", - }).save(); + const loaded = await new IndexPersistence(kv as never, new VectorIndex(), { bucketSize: 16 }).load(); - await expect(kv.get(oldShardScope, "data")).resolves.toBeNull(); + expect(loaded.state).toBe("migrated"); + expect(await kv.get(INDEX_SCOPE, "data:manifest")).toBeUndefined(); + expect(await kv.get(INDEX_SCOPE, "vectors:manifest")).toBeUndefined(); + for (const [scope, entries] of kv.store) { + if (scope.includes(":vectors:") || scope.includes(":bm25:idx_")) expect(entries.size).toBe(0); + } }); - it("falls back to previous-generation cleanup when the ledger is unusable (#1115)", async () => { - // Measured regression guard: with no fallback, an unusable ledger leaked - // one generation per cycle — strictly worse than the code this replaced, - // which always had this path. Without this test the `else if` can be - // deleted with the whole suite green. - await new IndexPersistence(kv as never, makeBm25("obs_old", "alpha"), null, { - shardChars: 80, - createGeneration: () => "gen_old", - }).save(); - const oldShardScope = "mem:index:bm25:bm25:gen_old:00000"; - await expect(kv.get(oldShardScope, "data")).resolves.not.toBeNull(); - - // Ledger read fails, manifest read succeeds. Both are state::get with - // independent timeouts, so this split is ordinary, not contrived. - const ledgerFailsKv = { - ...kv, - get: vi.fn(async (scope: string, key: string): Promise => { - if (scope === BM25_SCOPE && key === `${BM25_MANIFEST_KEY}:gc`) { - throw new Error("Invocation timeout after 180000ms: state::get"); - } - return kv.get(scope, key); - }), - }; - - await new IndexPersistence( - ledgerFailsKv as never, - makeBm25("obs_new", "bravo"), - null, - { shardChars: 80, createGeneration: () => "gen_new" }, - ).save(); - - const manifest = await getBm25Manifest(kv); - expect(manifest.generation).toBe("gen_new"); - await expect(kv.get(oldShardScope, "data")).resolves.toBeNull(); - }); + it("migrates a monolithic vector snapshot", async () => { + const legacy = vectorWith([["obs_a", [0.1, 0.2, 0.3]]]); + await writeLegacyVectorSnapshot(kv, legacy, { monolithic: true }); - it("keeps the ledger entry when a rollback delete fails (#1115)", async () => { - // Shards that survived their rollback delete must keep the entry that - // names them, or nothing can ever reclaim them. - let shardWrites = 0; - const failingKv = { - ...kv, - set: vi.fn(async (scope: string, key: string, data: T): Promise => { - if (scope.includes(":gen_new:")) { - shardWrites += 1; - if (shardWrites === 2) throw new Error("shard write failed"); - } - return kv.set(scope, key, data); - }), - delete: vi.fn(async () => { - throw new Error("rollback delete failed"); - }), - }; + const loaded = await new IndexPersistence(kv as never, new VectorIndex(), { bucketSize: 16 }).load(); - await new IndexPersistence( - failingKv as never, - makeBm25("obs_new", "bravo new snapshot"), - null, - { shardChars: 80, createGeneration: () => "gen_new" }, - ).save(); - - const ledger = await kv.get<{ - v: 1; - generations: Array<{ generation: string }>; - }>(BM25_SCOPE, `${BM25_MANIFEST_KEY}:gc`); - expect(ledger!.generations.map((entry) => entry.generation)).toContain( - "gen_new", - ); + expect(loaded.state).toBe("migrated"); + expectSameVectors(loaded.vector, legacy); + expect(await kv.get(INDEX_SCOPE, "vectors")).toBeNull(); }); - it("refuses to overwrite a gc ledger it does not recognise (#1115)", async () => { - const gcKey = `${BM25_MANIFEST_KEY}:gc`; - const future = { v: 2, generations: [], writtenByANewerBuild: true }; - await kv.set(BM25_SCOPE, gcKey, future); - - await new IndexPersistence( - kv as never, - makeBm25("obs_new", "bravo new snapshot"), - null, - { shardChars: 80, createGeneration: () => "gen_new" }, - ).save(); + it("leaves an unreadable legacy snapshot in place and writes nothing", async () => { + const legacy = vectorWith([["obs_a", [0.1, 0.2, 0.3]]]); + await writeLegacyVectorSnapshot(kv, legacy); + const manifest = await kv.get<{ shards: Array<{ scope: string; key: string }> }>(INDEX_SCOPE, "vectors:manifest"); + await kv.delete(manifest!.shards[0].scope, manifest!.shards[0].key); + kv.ops.length = 0; - // Rewriting it as a v1 ledger would drop every generation it tracked. - await expect(kv.get(BM25_SCOPE, gcKey)).resolves.toEqual(future); - }); + const loaded = await new IndexPersistence(kv as never, new VectorIndex(), { bucketSize: 16 }).load(); - it("bounds background boot reclaim and resumes on a later pass", async () => { - const liveScope = "mem:index:bm25:bm25:gen_live:00000"; - const orphanScopes = [ - "mem:index:bm25:bm25:gen_old_a:00000", - "mem:index:bm25:bm25:gen_old_a:00001", - "mem:index:bm25:bm25:gen_old_b:00000", - "mem:index:bm25:bm25:gen_old_b:00001", - ]; - await kv.set(liveScope, "data", "live shard"); - for (const scope of orphanScopes) await kv.set(scope, "data", "orphan shard"); - await kv.set(BM25_SCOPE, BM25_MANIFEST_KEY, { - v: 1, - generation: "gen_live", - shards: [{ scope: liveScope, key: "data", chars: 10 }], - chars: 10, - }); - await kv.set(BM25_SCOPE, `${BM25_MANIFEST_KEY}:gc`, { - v: 1, - generations: [ - { - generation: "gen_old_a", - shards: orphanScopes.slice(0, 2).map((scope) => ({ scope, key: "data" })), - }, - { - generation: "gen_old_b", - shards: orphanScopes.slice(2).map((scope) => ({ scope, key: "data" })), - }, - { generation: "gen_live", shards: [{ scope: liveScope, key: "data" }] }, - ], - }); - - await new IndexPersistence(kv as never, new SearchIndex(), null, { - reclaimOnLoad: "background", - bootReclaimMaxDeletes: 2, - }).load(); - - await vi.waitFor(async () => { - const ledger = await kv.get<{ - generations: Array<{ generation: string }>; - }>(BM25_SCOPE, `${BM25_MANIFEST_KEY}:gc`); - expect(ledger?.generations.map((entry) => entry.generation)).not.toContain( - "gen_old_a", - ); - }); - await expect(kv.get(orphanScopes[2], "data")).resolves.toBe("orphan shard"); - await expect(kv.get(orphanScopes[3], "data")).resolves.toBe("orphan shard"); - - await new IndexPersistence(kv as never, new SearchIndex(), null, { - bootReclaimMaxDeletes: 10, - }).load(); - await expect(kv.get(orphanScopes[2], "data")).resolves.toBeNull(); - await expect(kv.get(orphanScopes[3], "data")).resolves.toBeNull(); - await expect(kv.get(liveScope, "data")).resolves.toBe("live shard"); + expect(loaded.state).toBe("unavailable"); + expect(await kv.get(INDEX_SCOPE, "vectors:manifest")).not.toBeNull(); + expect(kv.ops.filter((o) => o.op === "set" || o.op === "delete")).toEqual([]); }); - it("keeps the previous vector generation when vector save fails after BM25 publish", async () => { - const previousBm25 = makeBm25("obs_old", "alpha previous snapshot"); - const previousVector = makeVector("obs_old"); - await new IndexPersistence(kv as never, previousBm25, previousVector, { - shardChars: 80, - createGeneration: () => "gen_old", - }).save(); - - const failingKv = { + it("keeps the legacy snapshot when migration writes fail and finishes on a later save", async () => { + const legacy = vectorWith([ + ["obs_a", [0.1, 0.2, 0.3]], + ["obs_b", [0.4, 0.5, 0.6]], + ]); + await writeLegacyVectorSnapshot(kv, legacy); + let fail = true; + const flakyKv = { ...kv, - set: vi.fn(async (scope: string, key: string, data: T): Promise => { - if (scope === BM25_SCOPE && key === VECTOR_MANIFEST_KEY) { - throw new Error("vector manifest write failed"); - } + set: async (scope: string, key: string, data: T): Promise => { + if (fail && scope.startsWith(`${INDEX_SCOPE}:vec:`)) throw new Error("write failed"); return kv.set(scope, key, data); - }), + }, }; - const nextBm25 = makeBm25("obs_new", "bravo new snapshot"); - const nextVector = new VectorIndex(); - nextVector.add("obs_new", "ses_1", new Float32Array([0.4, 0.5, 0.6])); - - await new IndexPersistence(failingKv as never, nextBm25, nextVector, { - shardChars: 80, - createGeneration: () => "gen_new", - }).save(); - - await expect( - kv.get("mem:index:bm25:vectors:gen_new:00000", "data"), - ).resolves.toBeNull(); - const loaded = await new IndexPersistence( - kv as never, - new SearchIndex(), - null, - ).load(); - expect(loaded.bm25!.search("bravo").length).toBe(1); - expect(loaded.vector!.size).toBe(1); - expect( - loaded.vector!.search(new Float32Array([0.1, 0.2, 0.3]))[0]?.obsId, - ).toBe("obs_old"); - }); - - it("fails closed when a manifest shard is missing", async () => { - const bm25 = makeBm25("obs_1", "alpha sharded snapshot"); - await new IndexPersistence(kv as never, bm25, null, { - shardChars: 80, - createGeneration: () => "gen_missing", - }).save(); - const manifest = await getBm25Manifest(kv); - await kv.delete(manifest.shards[0].scope, manifest.shards[0].key); - - const loaded = await new IndexPersistence( - kv as never, - new SearchIndex(), - null, - ).load(); - - expect(loaded.bm25).toBeNull(); - }); - - it("fails closed when a manifest shard length mismatches", async () => { - const bm25 = makeBm25("obs_1", "alpha sharded snapshot"); - await new IndexPersistence(kv as never, bm25, null, { - shardChars: 80, - createGeneration: () => "gen_mismatch", - }).save(); - const manifest = await getBm25Manifest(kv); - const firstShard = manifest.shards[0]; - const chunk = await kv.get(firstShard.scope, firstShard.key); - await kv.set(firstShard.scope, firstShard.key, `${chunk}x`); - - const loaded = await new IndexPersistence( - kv as never, - new SearchIndex(), - null, - ).load(); - - expect(loaded.bm25).toBeNull(); - }); - it("fails closed before reading invalid shard descriptors", async () => { - await kv.set(BM25_SCOPE, BM25_MANIFEST_KEY, { - v: 1, - chars: 10, - shards: [{ scope: "", key: "data", chars: 10 }], - }); - const guardedKv = { - ...kv, - get: vi.fn(async (scope: string, key: string): Promise => { - if (scope === "") { - throw new Error("invalid shard descriptor was read"); - } - return kv.get(scope, key); - }), - }; + const live = new VectorIndex(); + const persistence = new IndexPersistence(flakyKv as never, live, { bucketSize: 16 }); + const loaded = await persistence.load(); + expect(loaded.state).toBe("unavailable"); + expectSameVectors(loaded.vector, legacy); + expect(await kv.get(INDEX_SCOPE, "vectors:manifest")).not.toBeNull(); + expect(await kv.get(INDEX_SCOPE, META_KEY)).toBeNull(); - const loaded = await new IndexPersistence( - guardedKv as never, - new SearchIndex(), - null, - ).load(); + live.restoreFrom(loaded.vector!); + fail = false; + await persistence.save(); - expect(loaded.bm25).toBeNull(); - expect(guardedKv.get).not.toHaveBeenCalledWith("", "data"); + const next = await new IndexPersistence(kv as never, new VectorIndex(), { bucketSize: 16 }).load(); + expect(next.state).toBe("buckets"); + expectSameVectors(next.vector, legacy); + expect(await kv.get(INDEX_SCOPE, "vectors:manifest")).toBeNull(); }); - it("scheduleSave debounces multiple calls", async () => { - const bm25 = new SearchIndex(); - const persistence = new IndexPersistence(kv as never, bm25, null); - - persistence.scheduleSave(); - persistence.scheduleSave(); - persistence.scheduleSave(); - - await expect(kv.get(BM25_SCOPE, BM25_MANIFEST_KEY)).resolves.toBeNull(); + it("without vector search, removes only the old BM25 snapshot", async () => { + const legacy = vectorWith([["obs_a", [0.1, 0.2, 0.3]]]); + await writeLegacyVectorSnapshot(kv, legacy); + await writeLegacyBm25Snapshot(kv); - vi.advanceTimersByTime(5000); - await vi.runAllTimersAsync(); + const loaded = await new IndexPersistence(kv as never, null, { bucketSize: 16 }).load(); - const saved = await kv.get(BM25_SCOPE, BM25_MANIFEST_KEY); - expect(saved).not.toBeNull(); + expect(loaded).toEqual({ vector: null, state: "none", savedAt: null }); + expect(await kv.get(INDEX_SCOPE, "data:manifest")).toBeNull(); + expect(await kv.get(INDEX_SCOPE, "vectors:manifest")).not.toBeNull(); }); +}); - it("stop clears the pending timer", async () => { - const bm25 = new SearchIndex(); - bm25.add(makeObs({ id: "obs_1", title: "auth handler" })); - const persistence = new IndexPersistence(kv as never, bm25, null); - - persistence.scheduleSave(); - persistence.stop(); +describe("IndexPersistence scheduling", () => { + let kv: MockKV; - vi.advanceTimersByTime(10000); - const saved = await kv.get(BM25_SCOPE, BM25_MANIFEST_KEY); - expect(saved).toBeNull(); + beforeEach(() => { + vi.useFakeTimers(); + kv = mockKV(); }); - it("returns null indexes when nothing has been saved", async () => { - const bm25 = new SearchIndex(); - const persistence = new IndexPersistence(kv as never, bm25, null); - - const loaded = await persistence.load(); - expect(loaded.bm25).toBeNull(); - expect(loaded.vector).toBeNull(); + afterEach(() => { + vi.useRealTimers(); }); - it("scheduled save swallows kv.set rejection without unhandledRejection (#204)", async () => { + it("scheduled saves never raise unhandled rejections", async () => { const failingKv = { - ...mockKV(), + ...kv, set: vi.fn(async () => { - const err = new Error( - "TIMEOUT: invocation timed out after 30000ms", - ) as Error & { code?: string; function_id?: string }; + const err = new Error("TIMEOUT: invocation timed out after 30000ms") as Error & { code?: string }; err.code = "TIMEOUT"; - err.function_id = "state::set"; throw err; }), }; - const bm25 = new SearchIndex(); - bm25.add(makeObs({ id: "obs_1", title: "auth handler" })); - const persistence = new IndexPersistence(failingKv as never, bm25, null); - + const vector = vectorWith([["obs_a", [0.1, 0.2, 0.3]]]); + const persistence = new IndexPersistence(failingKv as never, vector, { saveIntervalMs: 5000, bucketSize: 16 }); let unhandled = false; const onUnhandled = () => { unhandled = true; }; process.on("unhandledRejection", onUnhandled); - try { persistence.scheduleSave(); - vi.advanceTimersByTime(5000); - await vi.runAllTimersAsync(); - // give microtasks a chance to flush - await Promise.resolve(); + await vi.advanceTimersByTimeAsync(5000); expect(failingKv.set).toHaveBeenCalled(); expect(unhandled).toBe(false); } finally { @@ -1174,194 +477,12 @@ describe("IndexPersistence", () => { } }); - it("save() does not throw when kv.set rejects (#204)", async () => { - const failingKv = { - ...mockKV(), - set: vi.fn(async () => { - throw new Error("TIMEOUT"); - }), - }; - const bm25 = new SearchIndex(); - bm25.add(makeObs({ id: "obs_1", title: "auth handler" })); - const persistence = new IndexPersistence(failingKv as never, bm25, null); - - await expect(persistence.save()).resolves.toBeUndefined(); - }); - - // #797: first run after upgrading to 0.9.25 crashed with - // 'TypeError: Cannot read properties of undefined (reading "v")' - // because some iii-state adapters return `undefined` (not `null`) - // for a missing key. The load path's `value !== null` check passed - // undefined to loadManifestData, which then read `undefined.v`. - it("load() returns null instead of crashing when kv.get returns undefined for the manifest (#797)", async () => { - const undefinedKv = { - ...mockKV(), - get: vi.fn(async () => undefined), - }; - const persistence = new IndexPersistence( - undefinedKv as never, - new SearchIndex(), - null, - ); - - const loaded = await persistence.load(); - expect(loaded.bm25).toBeNull(); - expect(loaded.vector).toBeNull(); - }); - - it("load() does not crash when a manifest row value is the wrong shape (#797)", async () => { - const wrongShapeKv = { - ...mockKV(), - get: vi.fn(async () => "not-a-manifest"), - }; - const persistence = new IndexPersistence( - wrongShapeKv as never, - new SearchIndex(), - null, - ); - - await expect(persistence.load()).resolves.toBeDefined(); - }); -}); - -describe("IndexPersistence save coalescing", () => { - let kv: ReturnType; - - beforeEach(() => { - vi.useFakeTimers(); - kv = mockKV(); - }); - afterEach(() => vi.useRealTimers()); - - it("does not queue a second identical save behind one that has not started", async () => { - // flushIndexSave awaits save() on every delete path. Serialising every one - // of them means a delete waits out every save ahead of it, which is fine - // when a save takes a second and is an outage when saves are timing out at - // 180s. A save that is queued but not yet running will serialise the index - // as it stands when it runs, so it already covers these callers. - let manifestWrites = 0; - let releaseFirst!: () => void; - const firstStarted = new Promise((r) => { - releaseFirst = r; - }); - let started = 0; - const slowKv = { - ...kv, - set: async (scope: string, key: string, data: T): Promise => { - if (scope === BM25_SCOPE && key === BM25_MANIFEST_KEY) { - manifestWrites++; - if (++started === 1) await firstStarted; - } - return kv.set(scope, key, data); - }, - }; - - const persistence = new IndexPersistence( - slowKv as never, - makeBm25("obs_1", "alpha"), - null, - { shardChars: 400 }, - ); - - const first = persistence.save(); - // Five more delete-path flushes arrive while the first is still running. - const rest = Array.from({ length: 5 }, () => persistence.save()); - releaseFirst(); - await Promise.all([first, ...rest]); - - // One running save plus at most one queued behind it covers all six - // callers. Six serialised saves would be the regression. - expect(manifestWrites).toBeLessThanOrEqual(2); - }); -}); - -describe("index_persist audit gating", () => { - let kv: ReturnType; - let previousFlag: string | undefined; - - beforeEach(() => { - // AGENTMEMORY_* variables are documented as living in - // ~/.agentmemory/.env, so a developer running the suite on a - // configured machine can inherit this one. Clear it going in and put - // whatever was there back on the way out. - previousFlag = process.env.AGENTMEMORY_AUDIT_INDEX_PERSIST; - delete process.env.AGENTMEMORY_AUDIT_INDEX_PERSIST; - vi.useFakeTimers(); - kv = mockKV(); - }); - - afterEach(() => { - vi.useRealTimers(); - if (previousFlag === undefined) { - delete process.env.AGENTMEMORY_AUDIT_INDEX_PERSIST; - } else { - process.env.AGENTMEMORY_AUDIT_INDEX_PERSIST = previousFlag; - } - }); - - async function indexPersistEntries(): Promise> { - const entries = await kv.list<{ operation: string }>("mem:audit"); - return entries.filter((entry) => entry.operation === "index_persist"); - } - - it("writes no index_persist audit entries by default", async () => { - const persistence = new IndexPersistence( - kv as never, - makeBm25("obs_1", "auth handler"), - null, - ); - - await persistence.save(); - - expect(await indexPersistEntries()).toEqual([]); - }); - - it.each(["1", " 1 ", "true", "TRUE", " true "])( - "writes index_persist audit entries when set to %j", - async (value) => { - process.env.AGENTMEMORY_AUDIT_INDEX_PERSIST = value; - const persistence = new IndexPersistence( - kv as never, - makeBm25("obs_1", "auth handler"), - null, - ); - - await persistence.save(); - - expect((await indexPersistEntries()).length).toBeGreaterThan(0); - }, - ); - - // Anything that is not an affirmative stays off. "0" and "false" are the - // ones an operator is likely to reach for to disable it, and they must - // not read as "present, therefore enabled". - it.each(["0", "false", "yes", "", " "])( - "keeps auditing off when set to %j", - async (value) => { - process.env.AGENTMEMORY_AUDIT_INDEX_PERSIST = value; - const persistence = new IndexPersistence( - kv as never, - makeBm25("obs_1", "auth handler"), - null, - ); - - await persistence.save(); - - expect(await indexPersistEntries()).toEqual([]); - }, - ); - - it("still persists the index when auditing is off", async () => { - const persistence = new IndexPersistence( - kv as never, - makeBm25("obs_1", "auth handler"), - null, - ); - - await persistence.save(); - - const loaded = await persistence.load(); - expect(loaded.bm25).not.toBeNull(); - expect(loaded.bm25!.size).toBe(1); + it("stop clears the pending timer", async () => { + const vector = vectorWith([["obs_a", [0.1, 0.2, 0.3]]]); + const persistence = new IndexPersistence(kv as never, vector, { saveIntervalMs: 5000, bucketSize: 16 }); + persistence.scheduleSave(); + persistence.stop(); + await vi.advanceTimersByTimeAsync(10_000); + expect(await kv.get(INDEX_SCOPE, META_KEY)).toBeNull(); }); }); diff --git a/test/search.test.ts b/test/search.test.ts index 84b8134e3a..99b7f4d6d3 100644 --- a/test/search.test.ts +++ b/test/search.test.ts @@ -209,4 +209,21 @@ describe("mem::search", () => { setVectorIndex(null); setEmbeddingProvider(null); }); + + it("a cold-start search rebuilds BM25 without touching the persisted vector index", async () => { + const vector = new VectorIndex(); + vector.add("obs_persisted", "ses_1", new Float32Array([0.1, 0.2, 0.3])); + setVectorIndex(vector); + getSearchIndex().clear(); + + const result = (await sdk.trigger("mem::search", { + query: "auth middleware", + })) as { results: Array<{ observation: CompressedObservation }> }; + + expect(result.results[0]?.observation.id).toBe("obs_a"); + expect(getVectorIndex()?.has("obs_persisted")).toBe(true); + expect(getVectorIndex()?.size).toBe(1); + + setVectorIndex(null); + }); }); diff --git a/test/status.test.ts b/test/status.test.ts index afed9076d5..9efee6d91a 100644 --- a/test/status.test.ts +++ b/test/status.test.ts @@ -27,6 +27,8 @@ function inputs(overrides: Partial = {}): StatusInputs { observationsIndexed: 110, missingObservations: 0, sessions: 95, + bm25Incomplete: false, + pendingVectorBackfill: 0, }, graph: { totalNodes: 55, @@ -93,6 +95,28 @@ describe("evaluateStatus", () => { expect(report.status).toBe("info"); }); + it("reports an incomplete BM25 rebuild as an error", () => { + const report = evaluateStatus(inputs({ index: { ...inputs().index, bm25Incomplete: true } })); + expect(report.status).toBe("error"); + expect(codes(report)).toContain("bm25-rebuild-incomplete"); + }); + + it("notes a pending vector backfill without raising the overall status", () => { + const report = evaluateStatus(inputs({ index: { ...inputs().index, pendingVectorBackfill: 42 } })); + expect(report.status).toBe("info"); + const problem = report.problems.find((p) => p.code === "index-vector-backfill-pending"); + expect(problem?.message).toBe("42 documents are waiting for a vector embedding."); + }); + + it("reports a vector count shortfall from the last save as a warning", () => { + const report = evaluateStatus( + inputs({ indexPersistence: { saveIntervalMs: 600_000, saving: false, buckets: 3, pendingChanges: 0, vector: null, vectorCountShortfall: { expected: 100, loaded: 40 } } }), + ); + expect(report.status).toBe("warn"); + const problem = report.problems.find((p) => p.code === "index-vector-count-shortfall"); + expect(problem?.message).toContain("40 of 100 vectors"); + }); + it("checks snapshot presence when extraction is on, and snapshot age whenever a snapshot exists", () => { const noSnapshot = { totalNodes: 0, totalEdges: 0, fromSnapshot: false }; expect(codes(evaluateStatus(inputs({ graph: noSnapshot })))).toEqual(["graph-no-snapshot"]);
BM25 indexrebuilt from stored content at bootVector savesaved 1m ago, unsaved changes pendingVector storage256 buckets, 3 unsaved changes