fix(_memory): make interrupted reindex resumable, verified by model identity - #1792
Conversation
…dentity Memory.initialize() used to embed every document in one pass when a reindex was needed (embedding model changed) and only persist the result at the very end. Interrupting that pass -- a crash, a restart, a slow or shared embeddings backend timing out mid-batch -- threw away every already-embedded document; the next attempt started over from zero. On a large memory store against a slow backend this can cost hours of redone work every time the process is interrupted. Reindexing now embeds in batches (Memory._REBUILD_BATCH_SIZE) and checkpoints to disk after each one (index.rebuilding.faiss/.pkl), so an interrupted rebuild resumes from where it left off instead of restarting. The checkpoint is not trusted blindly: it's tagged with the embedding model it was built under (index.rebuilding.json), and a resume is only used if that matches the model we're about to (re)index with. If the target model changed again while a rebuild was interrupted, the stale checkpoint is discarded and a fresh rebuild starts -- otherwise resuming would add new-model vectors into an old-model FAISS index, which fails a dimension assertion on the very next batch. Adds tests/test_memory_rebuild_resume.py covering the checkpoint save/load/clear helpers directly and Memory.initialize() end-to-end for both the resume-under-same-model and discard-under-changed-model cases.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: c4dca4d4ed
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| if ( | ||
| meta.get("model_provider") != model_config.provider | ||
| or meta.get("model_name") != model_config.name | ||
| ): |
There was a problem hiding this comment.
Include embedding kwargs in checkpoint identity
When a rebuild is resumed after the user changes an embedding option while keeping the same provider/name (for example OpenAI dimensions, or api_base pointing to a different compatible backend), this check still accepts the old checkpoint even though initialize() constructs the embedder from model_config.build_kwargs(). The next batch can then append vectors from a different dimension/space to the checkpoint, causing a FAISS dimension assertion or a mixed index instead of the fresh rebuild intended for the new configuration; include the effective non-secret embedding signature (or at least the vector dimension) in the checkpoint validation.
AGENTS.md reference: plugins/_memory/AGENTS.md:L17-L19
Useful? React with 👍 / 👎.
| PrintStyle.standard(f"Indexed {done}/{total} memories...") | ||
| Memory._write_rebuild_checkpoint(db, db_dir, model_config) | ||
|
|
||
| Memory._clear_rebuild_checkpoint(db_dir) |
There was a problem hiding this comment.
Clear the rebuild checkpoint only after the final save
Because the checkpoint is removed before the rebuilt index.* files and the new embedding.json are durably written, a kill/restart during _save_db_file() or in the gap just before it leaves no index.rebuilding.* files to resume from while the persisted index is still stale or only partially written. For large stores this reintroduces the full re-embed-on-retry failure mode that the checkpoint is meant to avoid; move the cleanup after the final index and metadata writes succeed.
Useful? React with 👍 / 👎.
…e ordering Two gaps flagged by review on the resumable-reindex fix: 1. The rebuild checkpoint's model-identity check only compared provider/name. A kwarg that changes the embedder's actual output (e.g. an OpenAI-style `dimensions` override, or an api_base pointed at a different backend) could change while provider/name stayed the same, letting a checkpoint built under the old kwargs get resumed into under the new ones. The signature now includes build_kwargs() (minus api_key, which is deliberately excluded since it's persisted to disk in the memory dir and a credential rotation alone isn't "a different model"). 2. The rebuild checkpoint was cleared right after the batch loop, before the final index/embedding.json were durably written. A crash in that gap left neither a valid final index nor a resumable checkpoint, reintroducing the full re-embed-from-scratch failure mode the checkpoint exists to prevent. Clearing now happens strictly after the final save succeeds. Adds 3 tests: api_key exclusion, kwarg-change rejection, and checkpoint-survives-a-crash-during-final-save.
Summary
When the configured embedding model no longer matches the one a memory store was built with,
Memory.initialize()triggers a full reindex — but it embedded every document in a single pass and only persisted the result at the very end. Interrupting that pass (a crash, a restart, a slow/shared embeddings backend timing out mid-batch) threw away every already-embedded document, so the next attempt started over from zero. Against a slow or shared embeddings backend this can cost hours of redone work on every interruption — hit this in production on a small local GPU box.Memory._REBUILD_BATCH_SIZE, 20) and checkpoints to disk after each batch (index.rebuilding.faiss/.pkl), so an interrupted rebuild resumes where it left off.index.rebuilding.json) and is only trusted if that matches the model we're about to (re)index with. If the target model changed again while a rebuild was interrupted, the stale checkpoint is discarded and a fresh rebuild starts — otherwise resuming would add new-model vectors into an old-model FAISS index, failing a dimension assertion on the next batch.No change to the existing reindex-trigger logic (comparing
embedding.json's recorded provider/name against the current config) — this only changes how the reindex, once triggered, is executed.Test plan
tests/test_memory_rebuild_resume.py— 9 tests covering the checkpoint save/load/clear helpers directly, plusMemory.initialize()end-to-end for both the resume-under-same-model and discard-under-changed-model cases. All pass.mainuntouched: a stale import inemail_parser_test.py, andrate_limiter_test.pyrequiring live OpenRouter credentials).🤖 Generated with Claude Code