Skip to content

fix(knowledge): make KB configuration updates atomic and race-safe#620

Open
huytg2610 wants to merge 5 commits into
HKUDS:devfrom
huytg2610:fix/kb-config-atomic-store
Open

fix(knowledge): make KB configuration updates atomic and race-safe#620
huytg2610 wants to merge 5 commits into
HKUDS:devfrom
huytg2610:fix/kb-config-atomic-store

Conversation

@huytg2610

@huytg2610 huytg2610 commented Jul 11, 2026

Copy link
Copy Markdown

Description

kb_config.json is mutated by two independent writers — KnowledgeBaseManager and KnowledgeBaseConfigService — and neither held a lock across its full read-modify-write cycle. Concretely:

  • Lost updates. Each writer loaded the file, mutated its in-memory snapshot, and dumped the whole snapshot back. If writer B committed while writer A was between its read and its write, A's commit erased B's change entirely. This affects status updates, registrations, provider/default changes, orphan pruning, and deletions.
  • Torn reads / truncate-before-lock. _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.
  • Silent corruption recovery. Any parse error was swallowed into a mutable default payload, so a damaged-but-recoverable file could be overwritten by the next write.

This PR routes every mutation of kb_config.json through a single new KBConfigStore (deeptutor/knowledge/config_store.py, stdlib-only):

  • One exclusive lock on a stable sidecar (.kb_config.json.lock) held across read → mutate → write. A sidecar is required because os.replace swaps 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 (TimeoutError on expiry), which also maps cleanly onto Windows' non-blocking msvcrt.locking(LK_NBLCK) semantics.
  • Commits write a unique same-directory temp file + 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.
  • A file that exists but does not parse raises a dedicated KBConfigCorruptionError and 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 KnowledgeBaseManager mutators (status updates, all five registration paths, orphan pruning + auto-discovery in list_knowledge_bases, deletion, load-time migration) and all KnowledgeBaseConfigService mutators (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 in remove_asset (which shares its kb_config.json with the partner-side manager) — were moved onto transactions as well; no direct writer to any kb_config.json remains. 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.rmtree in 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.json writers (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 canonical defaults.default_kb (the legacy top-level default migration 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 slow rmtree cannot 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), and resolve_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 through KBConfigStore.read(): a missing file falls back exactly as before; a file that exists but does not parse raises KBConfigCorruptionError. 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

  • agents
  • api
  • config
  • core
  • knowledge
  • logging
  • services
  • tools
  • utils
  • web (Frontend)
  • docs (Documentation)
  • scripts
  • tests
  • Other: ...

Checklist

  • I have read and followed the contribution guidelines.
  • My code follows the project's coding standards.
  • I have run pre-commit run --all-files and fixed any issues.
  • I have added relevant tests for my changes.
  • I have updated the documentation (if necessary).
  • My changes do not introduce any new security vulnerabilities.

Additional Notes

tests/knowledge/test_kb_config_store.py pins 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 / failed os.replace leaving 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.py covers 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.py and test_linked_folder_sync.py add 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.py gains two regressions for the remove_asset prune. 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 on dev and stem from missing optional deps, not this change.

Notes for reviewers: _deleted_knowledge_bases entries 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 transient os.replace failures on Windows (small follow-up candidate).

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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant