Repository navigation
fix(index-persistence): shard persisted search snapshots - #763
Rokurolize wants to merge 1 commit into
Conversation
|
@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. |
📝 WalkthroughWalkthroughIndexPersistence now persists BM25 and vector indexes as bounded shards with manifest metadata instead of single monolithic values. Save writes per-generation shard chunks, publishes a manifest, and cleans up legacy/stale data. Load reads manifests with integrity validation and falls back to legacy single-key data when manifests are missing or invalid. The constructor accepts configurable shard chunk size and generation factory options. ChangesIndex Sharding Implementation
🎯 3 (Moderate) | ⏱️ ~25 minutes
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
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. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (2)
src/state/index-persistence.ts (2)
1-1: ⚡ Quick winUse the repo-standard ID/timestamp helpers for generation strings.
This new path rolls its own generation format with
randomUUID()andDate.now(). Please switch it to the shared helpers so persistence generations follow the same ID and timestamp conventions as the rest ofsrc/**/*.ts.As per coding guidelines,
src/**/*.ts: UsegenerateId()for unique IDs, and capture timestamps once withnew Date().toISOString()and reuse.Also applies to: 40-42
🤖 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` at line 1, The code currently imports and uses randomUUID() and Date.now() to build persistence generation strings; replace randomUUID() with the repo-standard generateId() helper and capture the timestamp once using const ts = new Date().toISOString() and reuse ts wherever Date.now() was used so IDs and timestamps follow project conventions; update the import to remove randomUUID and import generateId from the shared helper module, and change any occurrences that build generation strings to use generateId() and ts (e.g., wherever generation or generationId is composed).
243-265: ⚡ Quick winFetch manifest shards in parallel.
These
kv.get()calls are independent, but the current loop serializes them one-by-one. On larger manifests that directly inflates load latency even thoughPromise.all()can preserve shard order for the final join.Suggested refactor
- const chunks: string[] = []; - let chars = 0; - for (const shard of manifest.shards) { - const chunk = await this.kv - .get<string>(shard.scope, shard.key) - .catch(() => null); + const loadedShards = await Promise.all( + manifest.shards.map(async (shard) => ({ + shard, + chunk: await this.kv.get<string>(shard.scope, shard.key).catch(() => null), + })), + ); + const chunks: string[] = []; + 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, });As per coding guidelines,
src/**/*.ts: Use parallel operations where possible (Promise.allfor 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 kv reads (this.kv.get) causing unnecessary latency; change to fetch shards in parallel using Promise.all over manifest.shards.map(...) to call this.kv.get for each shard, preserve the original shard order when awaiting results, then validate each returned chunk (type string and length === shard.chars) and log the same warnings via logger.warn for missing/mismatched shards, pushing valid chunks into chunks and summing chars as before; reference manifest.shards, this.kv.get, chunks, chars, and logger.warn when locating and updating the code.
🤖 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 31-37: The shardChars function currently accepts positive
fractions (e.g. 0.5) which get floored to 0 and causes save() to never advance
offset; update the validation in shardChars to reject any configured value less
than 1 before flooring (e.g. require Number.isFinite(configured) && configured
>= 1) and then return Math.floor(configured), otherwise fall back to
DEFAULT_INDEX_SHARD_CHARS; reference the shardChars function and the
save()/offset behavior so reviewers can verify the fix.
- Around line 175-180: Wrap the manifest write
(this.kv.set<IndexShardManifest>(KV.bm25Index, manifestKey, {...})) in a
try/catch; on catch, re-read the manifest via this.kv.get(KV.bm25Index,
manifestKey) and check whether the new generation was actually published
(compare returned manifest.generation or v/generation to the local generation).
If the manifest does not reflect the new generation, delete the freshly written
shards (iterate the shards array and remove each shard key via this.kv.delete or
the existing shard-delete helper) to roll back, then rethrow the original error;
if the manifest does reflect the new generation, leave shards intact and rethrow
or return as appropriate. Ensure you reference generation, shards, manifestKey
and IndexShardManifest in the fix.
---
Nitpick comments:
In `@src/state/index-persistence.ts`:
- Line 1: The code currently imports and uses randomUUID() and Date.now() to
build persistence generation strings; replace randomUUID() with the
repo-standard generateId() helper and capture the timestamp once using const ts
= new Date().toISOString() and reuse ts wherever Date.now() was used so IDs and
timestamps follow project conventions; update the import to remove randomUUID
and import generateId from the shared helper module, and change any occurrences
that build generation strings to use generateId() and ts (e.g., wherever
generation or generationId is composed).
- Around line 243-265: The loop over manifest.shards does serial kv reads
(this.kv.get) causing unnecessary latency; change to fetch shards in parallel
using Promise.all over manifest.shards.map(...) to call this.kv.get for each
shard, preserve the original shard order when awaiting results, then validate
each returned chunk (type string and length === shard.chars) and log the same
warnings via logger.warn for missing/mismatched shards, pushing valid chunks
into chunks and summing chars as before; reference manifest.shards, this.kv.get,
chunks, chars, and logger.warn when locating and updating the code.
🪄 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: 4bd46c34-2d1a-47ab-ba6d-26f1d1ab11b1
📒 Files selected for processing (2)
src/state/index-persistence.tstest/index-persistence.test.ts
| function shardChars(options: IndexPersistenceOptions): number { | ||
| const configured = options.shardChars; | ||
| return typeof configured === "number" && | ||
| Number.isFinite(configured) && | ||
| configured > 0 | ||
| ? Math.floor(configured) | ||
| : DEFAULT_INDEX_SHARD_CHARS; |
There was a problem hiding this comment.
Reject fractional shard sizes below 1 before flooring.
A value like shardChars: 0.5 passes the current guard, gets floored to 0, and then Line 159 never advances offset. That turns save() into an infinite loop on a reachable misconfiguration path.
Suggested fix
function shardChars(options: IndexPersistenceOptions): number {
const configured = options.shardChars;
- return typeof configured === "number" &&
- Number.isFinite(configured) &&
- configured > 0
- ? Math.floor(configured)
- : DEFAULT_INDEX_SHARD_CHARS;
+ 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;
}📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| function shardChars(options: IndexPersistenceOptions): number { | |
| const configured = options.shardChars; | |
| return typeof configured === "number" && | |
| Number.isFinite(configured) && | |
| configured > 0 | |
| ? Math.floor(configured) | |
| : DEFAULT_INDEX_SHARD_CHARS; | |
| 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; | |
| } |
🤖 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 31 - 37, The shardChars function
currently accepts positive fractions (e.g. 0.5) which get floored to 0 and
causes save() to never advance offset; update the validation in shardChars to
reject any configured value less than 1 before flooring (e.g. require
Number.isFinite(configured) && configured >= 1) and then return
Math.floor(configured), otherwise fall back to DEFAULT_INDEX_SHARD_CHARS;
reference the shardChars function and the save()/offset behavior so reviewers
can verify the fix.
| await this.kv.set<IndexShardManifest>(KV.bm25Index, manifestKey, { | ||
| v: 1, | ||
| generation, | ||
| shards, | ||
| chars: serialized.length, | ||
| }); |
There was a problem hiding this comment.
Roll back newly written shards when manifest publication fails before commit.
This path currently leaves the freshly written generation in KV whenever the manifest set() rejects before publishing. That breaks the PR's rollback invariant: only shard-write failures are cleaned up, but manifest-write failures leak unreachable shards.
A small try/catch here can preserve the intended behavior: after a manifest write error, re-read manifestKey; if the new generation is not published, delete shards and rethrow, otherwise leave them in place because the manifest may already have committed.
Suggested fix
- await this.kv.set<IndexShardManifest>(KV.bm25Index, manifestKey, {
+ const nextManifest: IndexShardManifest = {
v: 1,
generation,
shards,
chars: serialized.length,
- });
+ };
+ try {
+ await this.kv.set<IndexShardManifest>(
+ KV.bm25Index,
+ manifestKey,
+ nextManifest,
+ );
+ } catch (err) {
+ const published = await this.kv
+ .get<IndexShardManifest>(KV.bm25Index, manifestKey)
+ .catch(() => null);
+ const committed =
+ published?.v === 1 &&
+ published.generation === generation &&
+ published.chars === nextManifest.chars &&
+ published.shards.length === nextManifest.shards.length &&
+ published.shards.every((shard, index) => {
+ const expected = nextManifest.shards[index];
+ return (
+ shard.scope === expected?.scope &&
+ shard.key === expected.key &&
+ shard.chars === expected.chars
+ );
+ });
+ if (!committed) {
+ await this.deleteShards(shards);
+ }
+ throw err;
+ }📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| await this.kv.set<IndexShardManifest>(KV.bm25Index, manifestKey, { | |
| v: 1, | |
| generation, | |
| shards, | |
| chars: serialized.length, | |
| }); | |
| const nextManifest: IndexShardManifest = { | |
| v: 1, | |
| generation, | |
| shards, | |
| chars: serialized.length, | |
| }; | |
| try { | |
| await this.kv.set<IndexShardManifest>( | |
| KV.bm25Index, | |
| manifestKey, | |
| nextManifest, | |
| ); | |
| } catch (err) { | |
| const published = await this.kv | |
| .get<IndexShardManifest>(KV.bm25Index, manifestKey) | |
| .catch(() => null); | |
| const committed = | |
| published?.v === 1 && | |
| published.generation === generation && | |
| published.chars === nextManifest.chars && | |
| published.shards.length === nextManifest.shards.length && | |
| published.shards.every((shard, index) => { | |
| const expected = nextManifest.shards[index]; | |
| return ( | |
| shard.scope === expected?.scope && | |
| shard.key === expected.key && | |
| shard.chars === expected.chars | |
| ); | |
| }); | |
| if (!committed) { | |
| await this.deleteShards(shards); | |
| } | |
| throw err; | |
| } |
🤖 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 175 - 180, Wrap the manifest
write (this.kv.set<IndexShardManifest>(KV.bm25Index, manifestKey, {...})) in a
try/catch; on catch, re-read the manifest via this.kv.get(KV.bm25Index,
manifestKey) and check whether the new generation was actually published
(compare returned manifest.generation or v/generation to the local generation).
If the manifest does not reflect the new generation, delete the freshly written
shards (iterate the shards array and remove each shard key via this.kv.delete or
the existing shard-delete helper) to roll back, then rethrow the original error;
if the manifest does reflect the new generation, leave shards intact and rethrow
or return as appropriate. Ensure you reference generation, shards, manifestKey
and IndexShardManifest in the fix.
|
Note: This comment 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. Closing this PR in favor of #764. After opening this PR, I re-read
#764 replaces this PR with the same focused two-file change from a branch named |
1 participant
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
IndexPersistencepath.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 reportsmergeable_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 inIndexPersistence.Summary
Large BM25/vector search snapshots are currently saved as monolithic strings in the
mem:index:bm25state 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 serializedvectorsvalue 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 withkv.set(KV.bm25Index, "data", ...)andkv.set(KV.bm25Index, "vectors", ...).KV.bm25Indexis themem:index:bm25state scope.Current upstream source at
fd9e3bd:src/state/index-persistence.tslines 31-40,src/state/search-index.tslines 194-210,src/state/vector-index.tslines 124-136, andsrc/state/schema.tslines 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::settimeout 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
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 indexeswraps theStateKVused byIndexPersistenceso 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 oversizeddatastring and one oversizedvectorsstring. 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 scopesaves vector index shards outside the BM25 scopeloads legacy monolithic BM25 and vector snapshotsloads sharded manifests that omit optional generation metadatakeeps the previous generation when a shard write fails before manifest commitkeeps the previous generation when manifest set rejects before commitkeeps a generation loadable when manifest set commits before rejectingdeletes a shard that committed before set rejectedloads the new generation even when old generation cleanup failskeeps the previous vector generation when vector save fails after BM25 publishfails closed when a manifest shard is missingfails closed when a manifest shard length mismatchesVerification:
npm test -- test/index-persistence.test.tsnpm run buildHOME=/tmp/agentmemory-pr-verify-home-20260601 npm testonorigin/mainplus this two-file patchReview 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
Release Notes
New Features
Bug Fixes