perf(index): store vectors in buckets and rebuild BM25 from content

.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

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.

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.
 
 <!-- AUTOGEN:env START - generated by scripts/skills/generate.ts, do not edit by hand -->
-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`

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;
-}

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<void>;
-} | null = null;
+  status?: () => IndexPersistenceStatus;
+};
+
+let indexPersistence: IndexPersistenceHook | null = null;
 
-export function setIndexPersistence(
-  p: { scheduleSave: () => void; save: () => Promise<void> } | 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<number> {
   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<void> {
+  return new Promise((resolve) => setTimeout(resolve, ms))
+}
+
+async function listSessionsWithRetry(kv: StateKV): Promise<Session[] | null> {
+  let lastErr: unknown
+  for (let attempt = 0; attempt < SESSIONS_LIST_RETRY_ATTEMPTS; attempt++) {
+    try {
+      return await kv.list<Session>(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<Memory>(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<CompressedObservation>(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<number> {
+  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', {