fix(index-persistence): shard persisted search snapshots

[Rokurolize]

Note: This PR description was drafted with AI assistance. I personally experienced the underlying issue. If any AI-slop-like shortcomings remain, they are due to my prompting, and I welcome suggestions on how to communicate my requirements more clearly to Codex GPT-5.5 xhigh.

Fixes #762.

This PR does not close #309 because that issue is a broader, already-closed SQLite/FTS5 architecture discussion. It implements the narrower chunked-persistence stopgap discussed there for the existing IndexPersistence path.

It also does not supersede #300. #300 is an open sqlite-vec vector-backend direction, opt-in through VECTOR_BACKEND=sqlite-vec; as checked immediately before filing on 2026-06-01 at 22:56 UTC, it is still unmerged and reports mergeable_state=dirty. If #300 lands, it may reduce vector pressure for that opt-in backend, but it does not make BM25 snapshot persistence bounded and does not remove the need for safe legacy/default snapshot publication semantics in IndexPersistence.

Summary

Large BM25/vector search snapshots are currently saved as monolithic strings in the mem:index:bm25 state scope. On large corpora, that turns rebuild persistence into one very large state value for BM25 and one for vectors. The evidence is intentionally separated: upstream source shows both snapshots are monolithic; #309 contains a broad large-corpus persistence-timeout report after #204 made timeout failures nonfatal, plus a concrete vector report where the serialized vectors value reached about 27 MB and timed out. This PR saves each snapshot as bounded shards and publishes a compact manifest only after all shards for that snapshot are written.

What was happening

The current upstream save path is:

• SearchIndex.serialize() returns one JSON string containing all BM25 entries, inverted terms, document term counts, and total document length.
• VectorIndex.serialize() returns one JSON string containing all vector entries and base64-encoded embedding data.
• IndexPersistence.save() writes those two strings with kv.set(KV.bm25Index, "data", ...) and kv.set(KV.bm25Index, "vectors", ...).
• KV.bm25Index is the mem:index:bm25 state scope.

Current upstream source at fd9e3bd: src/state/index-persistence.ts lines 31-40, src/state/search-index.ts lines 194-210, src/state/vector-index.ts lines 124-136, and src/state/schema.ts lines 3-13.

As the corpus grows, those two strings can become large enough that persistence fails even though the in-memory rebuild completed. Since save() catches and logs persistence failures, the process keeps running, but the rebuilt snapshot is not durably published. A later start can therefore load no rebuilt snapshot or the last older snapshot that was saved successfully.

That is the narrower problem this PR addresses. It is related to #204 / #209, but those changes only prevented state::set timeout rejections from crashing the process; they did not reduce persisted value size or make search snapshot publication robust for larger serialized indexes.

The strongest detailed public size evidence is vector-specific. BM25 is included because current source gives it the same single-value persistence boundary, its JSON grows with indexed entries and postings, and the guarded-KV regression test proves the storage-granularity failure for both generated snapshots. If maintainers prefer a narrower first step, the same approach can be split into vector-memory-backend sharding first, with BM25 snapshot sharding as a follow-up.

Fix

• Split BM25 and vector snapshots into bounded string shards.
• Write shards under generation-specific shard scopes before changing the active manifest.
• Publish the new snapshot with a compact manifest that records exact shard scopes and character counts.
• Load legacy monolithic snapshots when no manifest exists.
• Load any valid manifest that records exact shard scopes, even if optional metadata is absent.
• Delete partial new shards if a save fails before the manifest is published.
• Leave new shards in place if the manifest write reports failure after it was attempted, because state persistence may have committed the manifest before returning an error.
• Delete legacy keys and old shards only after the new manifest is published; cleanup failure is nonfatal.
• Return null and log a warning if a published shard is missing or has a length mismatch.

BM25 and vector snapshots are still saved independently, matching the previous save order. If vector saving fails after BM25 is published, the new BM25 snapshot remains active and the previous vector snapshot remains active.

Reproduction and tests

The regression test avoids one oversized state::set string payload for persisted indexes wraps the StateKV used by IndexPersistence so any string value over a small limit is rejected. It then saves generated BM25/vector indexes whose serialized forms are larger than that limit.

That test reproduces the storage-granularity bug on current upstream behavior because save() tries to write one oversized data string and one oversized vectors string. With this PR, the same logical indexes are written as bounded shards and load back successfully.

Additional coverage:

• saves BM25 index shards outside the BM25 metadata scope
• saves vector index shards outside the BM25 scope
• loads legacy monolithic BM25 and vector snapshots
• loads sharded manifests that omit optional generation metadata
• keeps the previous generation when a shard write fails before manifest commit
• keeps the previous generation when manifest set rejects before commit
• keeps a generation loadable when manifest set commits before rejecting
• deletes a shard that committed before set rejected
• loads the new generation even when old generation cleanup fails
• keeps the previous vector generation when vector save fails after BM25 publish
• fails closed when a manifest shard is missing
• fails closed when a manifest shard length mismatches

Verification:

• npm test -- test/index-persistence.test.ts
• npm run build
• HOME=/tmp/agentmemory-pr-verify-home-20260601 npm test on origin/main plus this two-file patch

Review notes

Please review the manifest publication order and failure handling carefully. The intended invariant is that failed writes do not corrupt or delete the last successfully published snapshot for that index, while invalid published manifests fail closed instead of silently loading partial or stale data.

Summary by CodeRabbit

• Improvements

  • Index persistence now uses sharded storage with per-index manifests for more reliable, resumable saves; safer commit/rollback, validation, and legacy snapshot support prevent partial or corrupted restores. Cleared vector snapshots persist as empty entries to avoid unintended reloads. Audit entries include a new operation type for index persistence.

• Tests

  • Expanded coverage for sharded persistence, manifest/rollback failure modes, load validation, vector behavior, and save/stop scheduling.

[vercel]

@Rokurolize is attempting to deploy a commit to the rohitg00's projects Team on Vercel.

A member of the Team first needs to authorize it.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: 7ff20788-802d-422a-a466-0b006fd82cc2

📥 Commits

Reviewing files that changed from the base of the PR and between 0914ad6 and b7d6b68.

📒 Files selected for processing (1)

• src/state/index-persistence.ts

🚧 Files skipped from review as they are similar to previous changes (1)

• src/state/index-persistence.ts

---

📝 Walkthrough

Walkthrough

IndexPersistence now shards BM25 and vector snapshots into deterministic per-generation chunks, publishes a manifest of shard descriptors, validates and reconstructs shards on load, deletes legacy monolithic payloads after commit, and includes extensive tests for normal and failure scenarios.

Changes

Sharded Index Persistence

Layer / File(s)	Summary
Schema, helpers, and constructor src/state/index-persistence.ts	Adds shard/manifest types, sharding configuration and helpers, generation creation, validation helpers, and extends IndexPersistence constructor to accept IndexPersistenceOptions.
Sharded save implementation src/state/index-persistence.ts	save() now chunks serialized BM25 and vector payloads, writes deterministic shard chunks, audits writes, publishes a manifest with generation and shard descriptors, reconciles manifest-publish failures, rolls back uncommitted shards when needed, deletes legacy monolithic keys, and cleans up obsolete shards from the previous manifest.
Sharded load and validation src/state/index-persistence.ts	load() prefers manifest-based reconstruction: reads and validates manifest, fetches all referenced shard chunks, verifies per-shard and total lengths, concatenates shards into the full serialized payload, and falls back to legacy monolithic reads only when the manifest value is null.
Types update src/types.ts	Adds "index_persist" to AuditEntry["operation"].
Test constants, types, and helpers test/index-persistence.test.ts	Defines BM25/vector scope and manifest key constants, TestIndexShardManifest type, typed MockKV alias, and builders/helpers (makeBm25, makeVector, getBm25Manifest) for tests.
Core sharding behavior tests test/index-persistence.test.ts	Tests verify shards are written outside BM25 metadata scope with manifest-driven discovery, legacy monolithic snapshots still load, manifests may omit optional generation metadata, and vector shards behave correctly (including empty snapshots for cleared vectors).
Robustness and edge-case tests test/index-persistence.test.ts	Guards against oversized kv.set string payloads during persistence and ensures fractional shardChars sizing does not produce zero-width shards.
Failure-mode and regression tests test/index-persistence.test.ts	Extensive tests for shard-write failures, manifest commit/rejection ordering, rollback deletion of orphaned shards, preserving new generation on old-generation cleanup errors, cross-index failure interactions, and fail-closed behavior for missing or length-mismatched shards.
Debounce / stop test updates test/index-persistence.test.ts	Debounce and stop() timer tests updated to assert against the BM25 manifest key and scope constants.

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~75 minutes

Poem

I'm a rabbit splitting giant strings,
I hop through shards and write small things.
I map the pieces, shout the plan,
old keys drift while new ones stand.
Hop—safe save—index dreams take wing! 🐰

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title 'fix(index-persistence): shard persisted search snapshots' clearly and concisely summarizes the main change—converting monolithic index persistence to sharded storage.
Linked Issues check	✅ Passed	The PR implementation comprehensively addresses all coding-level objectives from #762: sharded persistence with bounded chunks, manifest-based publication, legacy monolithic support, failure handling (partial shard deletion, manifest commit verification), and fail-closed load semantics.
Out of Scope Changes check	✅ Passed	All changes are in-scope: core persistence logic refactoring in index-persistence.ts, comprehensive test coverage in index-persistence.test.ts, and a minimal type addition in types.ts for audit logging are all directly required by #762.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Actionable comments posted: 4

🧹 Nitpick comments (2)

src/state/index-persistence.ts (2)

243-265: ⚡ Quick win

Load shard chunks in parallel.

Each kv.get here is independent, but the serial await loop makes restore latency and timeout exposure scale with shard count. A Promise.all over the validated shard list preserves ordering while following the repo's KV guidance. As per coding guidelines: src/**/*.ts: Use parallel operations where possible (Promise.all for independent kv writes/reads).

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/state/index-persistence.ts` around lines 243 - 265, The loop over
manifest.shards does serial awaits which slows restore; replace it by mapping
shards to promises via this.kv.get and use Promise.all to fetch all chunks in
parallel (preserving shard order), then iterate the results to perform the same
type and length validations, logging via logger.warn and returning null if any
shard is missing or length-mismatched; ensure you still populate chunks (in the
original order) and update chars only after successful validation; reference
manifest.shards, this.kv.get, chunks, chars, label, and logger.warn when making
the change.

---

40-41: ⚡ Quick win

Use the repo ID helper for shard generations.

generation is a unique ID, so building it with randomUUID() bypasses the project-standard helper and its shared format/behavior. Please switch this to generateId(). As per coding guidelines: src/**/*.ts: Use fingerprintId() for content-addressable dedup and generateId() for unique IDs.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/state/index-persistence.ts` around lines 40 - 41, Replace the ad-hoc
randomUUID() usage in createIndexGeneration with the project helper
generateId(): update the function to combine Date.now().toString(36) with
generateId() (dropping the .replace(/-/g,"") step), and add the necessary import
for generateId from the helper module while removing reliance on randomUUID so
the generation follows the repo-standard format; keep the existing timestamp
prefix `${Date.now().toString(36)}-` and append generateId().

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@src/state/index-persistence.ts`:
- Around line 73-75: The current branch only persists when this.vector.size > 0,
so if the in-memory index shrinks to 0 we must clear the previously persisted
index to avoid resurrecting deleted vectors in loadVectorData(); update the
logic around saveVectorIndex (and any related persistence API) so that when
this.vector exists but this.vector.size === 0 you explicitly remove or overwrite
the persisted index (e.g., call a delete/clear persistence method or call
saveVectorIndex with an empty/cleared manifest) and when size > 0 continue
saving this.vector.serialize(); reference saveVectorIndex, this.vector,
this.vector.serialize(), and loadVectorData when implementing the cleanup.
- Around line 145-192: The saveShardedIndex path writes shards, publishes the
manifest, deletes the legacy manifest, and cleans up old shards without calling
recordAudit(); add recordAudit() calls around each state-changing operation so
the audit trail captures: (1) each successful shard write inside the loop where
kv.set(scope, INDEX_SHARD_KEY, chunk) is called (include
shard.scope/shard.key/chunk length), (2) the manifest publication where
kv.set<IndexShardManifest>(KV.bm25Index, manifestKey, ...) is invoked (include
generation, manifestKey, total chars), (3) the legacy deletion
kv.delete(KV.bm25Index, legacyKey) (include legacyKey and result), and (4) each
deleteShards(...) call used for rollback and for cleaning previous.shards
(include the shard identifiers); use the existing recordAudit(...) helper in
this module/class so all these operations are recorded with the same metadata
format as other state changes.
- Around line 233-246: Validate each shard descriptor in the manifest before
calling kv.get: inside the loop that iterates manifest.shards (the code that
builds chunks and accumulates chars), check that shard.scope and shard.key are
non-empty strings and shard.chars is a finite non-negative integer (or >0 if
expected). If any shard fails validation, log a warning with the label and
return null (fail closed) instead of calling this.kv.get(shard.scope,
shard.key); otherwise proceed to await this.kv.get(...). This ensures corrupted
manifest entries never trigger KV reads.
- Around line 215-224: The manifest read currently swallows all errors via
catch(() => null) when calling this.kv.get for KV.bm25Index with manifestKey,
which can mask transient/backend failures and wrongly fall back to legacyKey;
update the logic to call this.kv.get(manifestKey) and only treat a confirmed
"not found" sentinel/error from the StateKV implementation as a miss (then call
loadManifestData(manifest, label) or continue to legacy), but for any other
error log the error (include context: KV.bm25Index, manifestKey) and return null
instead of falling back; inspect the StateKV.get contract to use its specific
sentinel type or error class/shape rather than a blanket catch, and apply the
same stricter handling for the legacyKey path.

---

Nitpick comments:
In `@src/state/index-persistence.ts`:
- Around line 243-265: The loop over manifest.shards does serial awaits which
slows restore; replace it by mapping shards to promises via this.kv.get and use
Promise.all to fetch all chunks in parallel (preserving shard order), then
iterate the results to perform the same type and length validations, logging via
logger.warn and returning null if any shard is missing or length-mismatched;
ensure you still populate chunks (in the original order) and update chars only
after successful validation; reference manifest.shards, this.kv.get, chunks,
chars, label, and logger.warn when making the change.
- Around line 40-41: Replace the ad-hoc randomUUID() usage in
createIndexGeneration with the project helper generateId(): update the function
to combine Date.now().toString(36) with generateId() (dropping the
.replace(/-/g,"") step), and add the necessary import for generateId from the
helper module while removing reliance on randomUUID so the generation follows
the repo-standard format; keep the existing timestamp prefix
`${Date.now().toString(36)}-` and append generateId().

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: ae7699af-dc04-4f80-838e-baf2444eb0d1

📥 Commits

Reviewing files that changed from the base of the PR and between fd9e3bd and 33a23d3.

📒 Files selected for processing (2)

• src/state/index-persistence.ts
• test/index-persistence.test.ts

[coderabbitai]

⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Clear persisted vectors when the in-memory index becomes empty.

If a non-empty vector snapshot was saved earlier and later this.vector.size drops to 0, this branch skips any write or cleanup. After a restart, loadVectorData() will still load the previous manifest, so deleted vectors reappear.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/state/index-persistence.ts` around lines 73 - 75, The current branch only
persists when this.vector.size > 0, so if the in-memory index shrinks to 0 we
must clear the previously persisted index to avoid resurrecting deleted vectors
in loadVectorData(); update the logic around saveVectorIndex (and any related
persistence API) so that when this.vector exists but this.vector.size === 0 you
explicitly remove or overwrite the persisted index (e.g., call a delete/clear
persistence method or call saveVectorIndex with an empty/cleared manifest) and
when size > 0 continue saving this.vector.serialize(); reference
saveVectorIndex, this.vector, this.vector.serialize(), and loadVectorData when
implementing the cleanup.

[coderabbitai]

🧹 Nitpick comments (1)

src/state/index-persistence.ts (1)

183-206: 💤 Low value

Optional: Consider parallelizing shard writes with tracked rollback.

The sequential shard writes are correct and simplify rollback (failed shard means previous shards in shards array need deletion). However, for large indexes with many shards, parallel writes could improve performance.

Alternative approach using Promise.allSettled:

1. Write all shards in parallel
2. Collect fulfilled/rejected results
3. On any rejection, delete only fulfilled shards

This is optional since sequential writes are functionally correct and the current approach aligns with the rollback semantics. As per coding guidelines, src/**/*.ts: Use parallel operations where possible (Promise.all for independent kv writes/reads).

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/state/index-persistence.ts` around lines 183 - 206, The current loop
writes shards sequentially which is safe but slow; change to parallel writes
using Promise.allSettled: build an array of shard descriptors (using
scopePrefix, generation, INDEX_SHARD_KEY, chunkChars and the computed scope and
chunk from serialized.slice) and an array of promises that call
this.kv.set(scope, INDEX_SHARD_KEY, chunk) followed by
this.auditIndexPersistence("shard_write", [statePath(scope, INDEX_SHARD_KEY)], {
scope, key: INDEX_SHARD_KEY, manifestKey, generation, chars: chunk.length });
run Promise.allSettled on those promises, collect the indices/scopes of
fulfilled tasks and on any rejected result call this.deleteShards(filledShards,
"shard_write_rollback") and rethrow the first error; ensure shards array
reflects only successfully written shards so deleteShards targets the right
scopes.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@src/state/index-persistence.ts`:
- Around line 183-206: The current loop writes shards sequentially which is safe
but slow; change to parallel writes using Promise.allSettled: build an array of
shard descriptors (using scopePrefix, generation, INDEX_SHARD_KEY, chunkChars
and the computed scope and chunk from serialized.slice) and an array of promises
that call this.kv.set(scope, INDEX_SHARD_KEY, chunk) followed by
this.auditIndexPersistence("shard_write", [statePath(scope, INDEX_SHARD_KEY)], {
scope, key: INDEX_SHARD_KEY, manifestKey, generation, chars: chunk.length });
run Promise.allSettled on those promises, collect the indices/scopes of
fulfilled tasks and on any rejected result call this.deleteShards(filledShards,
"shard_write_rollback") and rethrow the first error; ensure shards array
reflects only successfully written shards so deleteShards targets the right
scopes.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: 98a7cf1d-8925-407a-983c-33194fad3e78

📥 Commits

Reviewing files that changed from the base of the PR and between 80088ef and 0914ad6.

📒 Files selected for processing (3)

• src/state/index-persistence.ts
• src/types.ts
• test/index-persistence.test.ts

🚧 Files skipped from review as they are similar to previous changes (1)

• test/index-persistence.test.ts

[rohitg00]

Merged. Thanks @Rokurolize — sharded persistence + manifest commit/rollback semantics are exactly what #762 needed, and the regression coverage (oversized-payload guard, manifest-commits-after-error, length-mismatch fail-closed) is the highest-quality test addition the persistence path has seen.