perf(docs): persist the doc-search BM25 index across sessions - #55
perf(docs): persist the doc-search BM25 index across sessions#55pradeepmouli wants to merge 4 commits into
Conversation
… docs_embeddings.bin
murari316
left a comment
There was a problem hiding this comment.
Automated review (Fable, adversarially verified). Recommendation: request-changes.
The PR persists the doc-search BM25 index to .infigraph/docs_bm25_cache.bin with a versioned, bounds-checked binary format, mtime freshness anchored on docs_embeddings.bin, and graceful fallback-to-rebuild on any load failure. The core design is solid and well-tested (roundtrip, corrupt-cache, stale-cache tests), but it misses cleanup integration (DocIndex::clean leaves the full-text cache on disk) and has a query-time/reindex race that can pin a stale cache as fresh until the next reindex.
[MAJOR] data-retention — crates/infigraph-docs/src/lib.rs:127-129
DocIndex::clean() removes docs.kuzu, docs_embeddings.bin, and the HNSW files but the PR does not add removal of docs_bm25_cache.bin, and the cache file contains the verbatim full text of every indexed chunk. clean_docs is supposed to wipe doc-index artifacts; the new artifact survives it.
Failure scenario: User indexes internal Confluence pages, runs one doc search (cache written), then runs clean_docs to purge the doc index. docs.kuzu and embeddings are deleted, but .infigraph/docs_bm25_cache.bin still holds the complete text of every chunk on disk indefinitely. It is never served (anchor missing -> not fresh), so this is orphaned sensitive data plus disk bloat, not wrong results. One-line fix: remove_file the cache alongside lib.rs:127-129.
[MAJOR] concurrency — crates/infigraph-docs/src/search.rs:233-249 (load_or_build_doc_index, PR hunk)
TOCTOU between reading chunks and saving the cache: if a doc reindex rewrites docs_embeddings.bin while a search is between store.get_all_chunks() and idx.save(), the cache built from pre-reindex chunks lands with an mtime newer than the new anchor, so doc_bm25_cache_fresh (c >= a) treats the stale cache as fresh from then on.
Failure scenario: Doc watcher fires a reindex (update_doc_embeddings rewrites the anchor) while a search_docs call is mid-flight rebuilding the cache. The stale cache is saved after the anchor and wins the mtime comparison. Every subsequent search silently drops chunks deleted in the reindex (hydration lookup returns None -> filter_map discards) and never BM25-matches newly added docs — persisting until the NEXT reindex bumps the anchor, which on a quiet repo can be days. Fix: capture the anchor mtime before get_all_chunks and skip the save (or set the cache mtime to the captured anchor mtime) if the anchor changed.
[MINOR] concurrency — crates/infigraph-docs/src/search.rs:128-131 (save, PR hunk)
The temp-file name is per-process (path.with_extension("{pid}.tmp")), but the MCP server handles concurrent tool calls on multiple threads in ONE process — two simultaneous cache misses share the same tmp path, so the code comment 'concurrent writers can't interleave' is wrong for the in-process case.
Failure scenario: Two search_docs requests hit the MCP server concurrently after a reindex; both miss the cache and call save(). Thread A's fs::write truncates/overwrites the tmp file while thread B is mid-write, and the rename publishes a torn file. Self-healing because load() validates and falls back to rebuild, so impact is wasted rebuilds — but the comment claims a guarantee the code doesn't provide. Use tempfile::NamedTempFile or add a per-call unique suffix.
[MINOR] error-handling — crates/infigraph-docs/src/search.rs:129-131 (save, PR hunk)
save() leaks the .tmp file when the write succeeds but the rename fails — the error is returned without cleaning up, and per-PID names mean leaked files accumulate rather than being overwritten.
Failure scenario: On Windows (CI runs there per repo invariants), std::fs::rename over a destination another process currently has open can fail with permission denied. Each failed save from a fresh process leaves a distinct docs_bm25_cache..tmp in .infigraph, accumulating over time. Add a remove_file(&tmp) on the rename error path (and arguably on the write error path).
[MINOR] correctness — crates/infigraph-docs/src/search.rs:read_str closure in load (PR hunk)
read_str uses String::from_utf8_lossy, silently accepting invalid UTF-8, which contradicts the function's own doc ('Any structural problem is an Err ... callers treat a bad cache as a miss and rebuild').
Failure scenario: Bit-rot flips bytes inside a chunk-id string without changing lengths: all bounds checks pass, load() returns Ok, but the mangled id (with U+FFFD) fails the get_chunk_details hydration lookup, so matching results are silently dropped from search output instead of the corrupt cache triggering a rebuild. Use String::from_utf8(...).map_err(...) so corruption becomes an Err.
[MINOR] test-coverage — crates/infigraph-docs/tests/doc_bm25_cache.rs:125-155 (stale_cache_is_rebuilt_when_embeddings_are_newer)
No test verifies end-to-end invalidation semantics: the stale-cache test only asserts the cache file's mtime advanced, not that search results actually reflect content indexed after the cache was created.
Failure scenario: A future refactor could rewrite the cache file (mtime advances, test passes) while still serving the old in-memory/cached doc set — e.g., saving the loaded stale index back out. Extend the test to index_doc a new file with a distinctive term after the first search, then assert a search for that term returns the new chunk.
[NIT] correctness — crates/infigraph-docs/src/search.rs:load (PR hunk)
load() never checks that pos == data.len() after parsing, so files with trailing garbage (e.g., the torn-write outcome from the tmp-path collision where a shorter payload was renamed over a longer partial one) load as valid.
Failure scenario: A structurally-complete prefix followed by junk bytes parses successfully and the junk is ignored; combined with the lossy-UTF-8 issue this weakens the 'validated load' safety net. Add anyhow::ensure!(pos == data.len()).
[NIT] correctness — crates/infigraph-docs/src/search.rs:doc_bm25_cache_fresh (PR hunk)
Freshness uses c >= a on mtimes; on filesystems with coarse mtime granularity a reindex that rewrites docs_embeddings.bin within the same timestamp tick as a cache save leaves the stale cache 'fresh'. Consistent with the existing code-search bm25_cache pattern, so noting rather than blocking; the relatedly mtime-based test's 50ms sleep could be flaky on coarse-granularity filesystems.
Generated with Claude Code (Fable). Findings verified by an adversarial refute pass; false positives removed.
Summary
Doc search (
hybrid_doc_search_in_dir) currently rebuilds its entireDocBM25Indexin memory on every query — loading all chunks from the store and re-tokenizing the whole corpus — while code search already persists its BM25 index to.infigraph/bm25_cache.bin. This PR gives doc search the same treatment:DocBM25Indexgainssave/load/docs(). The on-disk format mirrors code search'sBM25Index::saveidiom (leading version byte + length-prefixed LE binary).hybrid_doc_search_in_dirloadsdocs_bm25_cache.binwhen it is at least as new asdocs_embeddings.bin— the same mtime freshness anchor code search uses, so a doc reindex invalidates the cache automatically. On a cache hit thestore.get_all_chunks()DB read is skipped too (the cache carries the chunk texts).Two deliberate hardenings over the idiom being mirrored:
loadis fully bounds-checked and returnsErr(never panics) on truncated/corrupt input, since a bad cache must degrade to a rebuild.Group/combined doc stores get their own cache automatically (keyed off the store's artifact dir; combined stores publish per-generation directories, so no cross-generation staleness is possible).
Testing
DocStore: first search writes the cache and repeat results are identical; corrupt cache falls back silently and is rewritten; stale cache (newer embeddings) is rebuilt.infigraph-docssuite: 93/93 passing on this branch (rebased on currentmain);cargo fmt/clippy -D warnings/cargo check --workspaceclean.🤖 Generated with Claude Code
https://claude.ai/code/session_01SjmvwHuwV5r7ZeZpJLp5oR