fix(knowledge): make KB configuration updates atomic and race-safe#620
Open
huytg2610 wants to merge 5 commits into
Open
fix(knowledge): make KB configuration updates atomic and race-safe#620huytg2610 wants to merge 5 commits into
huytg2610 wants to merge 5 commits into
Conversation
update_folder_sync_state() mutated the loaded metadata dict and returned without writing it back, so last_sync/synced_files never reached metadata.json while the API layer logged success. Future scans re-treated already-synced files as new. Write back through a same-directory temp file + os.replace so readers never observe a truncated metadata.json.
kb_config.json had two independent writers (KnowledgeBaseManager and KnowledgeBaseConfigService) doing read-modify-write with no lock spanning the cycle, so a slower writer committed its stale snapshot over anything written in between. Worse, _save_config() truncated the file via open(..., "w") before acquiring its lock: a reader entering that window saw an empty file, fell back to an empty default, and the next mutation persisted that fallback over the valid config. Invalid JSON was likewise swallowed into a mutable default payload. Route every mutation through a shared KBConfigStore transaction: exclusive lock on a stable sidecar (.kb_config.json.lock — the data file's inode changes on every commit, so it cannot carry the lock), read latest state, mutate, write a unique same-directory temp file, fsync, os.replace. Lock-free readers now see the previous or the new snapshot, never a torn one. A file that exists but does not parse raises KBConfigCorruptionError instead of becoming an empty default; a zero-byte legacy crash artifact still reads as unwritten. Slow work (directory scans, index probes, rmtree, PocketBase sync) stays outside the lock, and lock acquisition retries against a bounded deadline, matching Windows' non-blocking msvcrt semantics.
- Treat a zero-byte kb_config.json as corruption, same as invalid JSON: an empty file is the signature of the legacy truncate-then- crash write, and silently rebuilding a default over it would hide that damage. Only a missing file bootstraps the default payload. - Actually persist the legacy top-level "default" migration: the in-memory pop never marked the config as changed, so every later transaction read the raw file and wrote the stale field back. - Route the partner-workspace KB prune (remove_asset) through KBConfigStore; it was the last direct read-modify-write on a kb_config.json and could both lose a concurrent manager's update and truncate the file if interrupted mid-write.
The store and both writer classes already refuse to treat a damaged kb_config.json as an empty default, but the query-time readers still swallowed every parse error: a corrupt config silently rebound the KB to the default provider (provider_binding), dropped its configured retrieval mode (modes), and resolved a linked KB to a non-existent local folder that returns empty results (kb_paths). Read through KBConfigStore.read() instead: a missing file still falls back exactly as before, but a file that exists and does not parse now raises KBConfigCorruptionError through the resolver. The llamaindex embedding-mismatch banner keeps its never-fail-the-search contract but logs the corruption instead of hiding it.
Follow-ups from adversarial review of the atomic config store: - metadata.json writers (link/unlink folder, sync state, document hashes, initializer/reindex stamps) share the same sidecar-lock transaction and atomic-replace primitive as kb_config.json, so two metadata writers can no longer lose each other's updates or tear the file. The metadata store refuses to recreate a deleted KB directory. - The legacy top-level "default" migration preserves the value into defaults.default_kb instead of dropping it, and deleting a KB now repoints the canonical defaults.default_kb, not just the legacy field. Manager default get/set uses its own store, so a custom base_dir keeps its own default instead of the admin-scoped one. - Deletion leaves a durable tombstone (_deleted_knowledge_bases) plus per-KB creation generations: late status/progress/initializer/hash writers from a deleted or recreated generation are rejected instead of resurrecting the KB, and auto-discovery skips tombstoned orphan directories. Explicit creation (allow_recreate) clears the marker. Create/delete filesystem transitions serialize on a per-KB lifecycle lock so a slow rmtree cannot interleave with recreation, without blocking unrelated config transactions. - Store schema validation: knowledge_bases/defaults/_deleted_* must be JSON objects or readers fail loud; the memory-snapshot adapter reads through the store.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
kb_config.jsonis mutated by two independent writers —KnowledgeBaseManagerandKnowledgeBaseConfigService— and neither held a lock across its full read-modify-write cycle. Concretely:_save_config()opened the file with"w"(truncating it) before acquiring its lock, and the service's_save()had no lock at all. A reader entering that window saw an empty file,_load_config()silently fell back to{"knowledge_bases": {}}, and the next mutation persisted that empty fallback over the previously valid config. A crash between truncate and write left the same empty file behind permanently.This PR routes every mutation of
kb_config.jsonthrough a single newKBConfigStore(deeptutor/knowledge/config_store.py, stdlib-only):.kb_config.json.lock) held across read → mutate → write. A sidecar is required becauseos.replaceswaps the data file's inode on every commit, so a lock on the data file itself would not exclude writers that open the new inode. Acquisition polls a non-blocking attempt against a bounded deadline (TimeoutErroron expiry), which also maps cleanly onto Windows' non-blockingmsvcrt.locking(LK_NBLCK)semantics.fsync+os.replace. Readers — locked or not — only ever observe the previous or the new snapshot, never a truncated or partial file, and a crash mid-write cannot damage the current config.KBConfigCorruptionErrorand is left untouched. This includes a zero-byte file — the artifact the legacy truncate-crash leaves behind — since rebuilding a default over it would hide the damage. Only a missing file bootstraps the default payload.Both writer classes were refactored onto the store: all
KnowledgeBaseManagermutators (status updates, all five registration paths, orphan pruning + auto-discovery inlist_knowledge_bases, deletion, load-time migration) and allKnowledgeBaseConfigServicemutators (set_kb_config,set_provider_mode,set_global_defaults,set_default_kb,delete_kb_config) now run their read-modify-write inside one transaction, and their caches are updated from the committed snapshot rather than written from a possibly-stale one. The remaining direct writers found outside those classes —initializer.py, two spots in the knowledge router, and the partner-workspace KB prune inremove_asset(which shares itskb_config.jsonwith the partner-side manager) — were moved onto transactions as well; no direct writer to anykb_config.jsonremains. The legacy top-level"default"field migration is now actually committed, so later transactions can no longer write the stale field back.Slow work is kept outside the lock by construction: directory scans and metadata probing for auto-discovery, embedding fingerprint/index-version probes for "ready" status,
shutil.rmtreein deletion, and the best-effort PocketBase mirror all run before or after the transaction, never inside it.Adversarial review of the store surfaced a second family of lifecycle races, fixed here as well. Per-KB
metadata.jsonwriters (folder link/unlink, sync state, document hashes, initializer/reindex stamps) now share the same sidecar-lock transaction + atomic-replace primitive, so concurrent metadata writers no longer lose each other's updates or tear the file. Deleting a KB repoints the canonicaldefaults.default_kb(the legacy top-leveldefaultmigration also preserves its value now instead of dropping it), and leaves a durable tombstone plus per-KB creation generations: a late status/progress/initializer write queued before the delete can no longer resurrect the KB or write into a same-name successor, and auto-discovery skips tombstoned orphan directories. Explicit recreation clears the marker. Create/delete filesystem transitions serialize on a per-KB lifecycle lock, so a slowrmtreecannot interleave with recreation without blocking unrelated config transactions.The query-time readers were brought under the same corruption contract.
resolve_bound_provider(provider binding),resolve_kb_mode(retrieval mode), andresolve_kb_dir(linked-KB external pointer) used to swallow every parse error, so a corrupt config silently rebound the KB to the default provider, dropped its configured mode, or resolved a linked KB to a non-existent local folder with empty results. They now read throughKBConfigStore.read(): a missing file falls back exactly as before; a file that exists but does not parse raisesKBConfigCorruptionError. The llamaindex embedding-mismatch banner keeps its never-fail-the-search contract but logs the corruption instead of hiding it.Related Issues
Module(s) Affected
agentsapiconfigcoreknowledgeloggingservicestoolsutilsweb(Frontend)docs(Documentation)scriptstests...Checklist
pre-commit run --all-filesand fixed any issues.Additional Notes
tests/knowledge/test_kb_config_store.pypins the store contract: lost-update prevention across two threads and across two OS processes, mid-transaction readers seeing only complete snapshots, corruption (invalid JSON, zero-byte, and wrong-shaped sections alike) raising while preserving the damaged file byte-for-byte, the legacy"default"migration keeping its value, failed serialization / failedos.replaceleaving the original intact with no temp-file litter, bounded lock retry under contention, lock timeout, and manager+service cross-class races.tests/services/rag/test_kb_config_readers.pycovers the readers: corrupt and zero-byte configs raise through provider/mode/path resolution, a missing file still falls back, and valid configs still resolve providers, modes, and linked external paths.tests/knowledge/test_manager_delete.pyandtest_linked_folder_sync.pyadd the lifecycle regressions: delete repoints the canonical default, late status updates and stale initializers cannot resurrect or clear a tombstone, an old generation cannot write into a same-name successor, concurrent folder links keep both entries, and concurrent metadata writers preserve each other's keys.tests/services/partners/test_partner_workspace.pygains two regressions for theremove_assetprune. Existing manager/service/router tests that seeded state through the removed private_save_config()were ported to the transactional API. Full suite: 2421 passed; the 10 failures on this machine (partners/msteams/cron) reproduce identically ondevand stem from missing optional deps, not this change.Notes for reviewers:
_deleted_knowledge_basesentries persist until the same name is explicitly recreated — deliberate, since the tombstone is what keeps late background writers and orphaned-but-indexable directories from resurrecting a deleted KB; the map only accumulates names deleted and never reused. Intentionally out of scope: no retry was added around transientos.replacefailures on Windows (small follow-up candidate).