diff --git a/mempalace/backends/chroma.py b/mempalace/backends/chroma.py index 426f21cb5e..780490e73b 100644 --- a/mempalace/backends/chroma.py +++ b/mempalace/backends/chroma.py @@ -1383,12 +1383,17 @@ def _client(self, palace_path: str): ) if cached is None or inode_changed or mtime_changed or mtime_appeared: - # An inode swap means we are reopening a different physical DB - # (post-restore, fresh palace at the same path, etc.); drop the - # per-process gate so the quarantine pre-checks run again - # against the new disk state instead of trusting cached "we - # already cleaned this path" credit from the prior inode. - if inode_changed: + # Drop the per-process quarantine gate so the HNSW pre-checks + # run again against the new disk state. An inode swap means a + # different physical DB (post-restore, fresh palace at the same + # path); an mtime/appearance change means an external in-place + # write (closet_llm, mine, compress) that may have drifted the + # HNSW index while this process was running. + if ( + inode_changed + or mtime_changed + or (mtime_appeared and palace_path in self._freshness) + ): ChromaBackend._quarantined_paths.discard(palace_path) ChromaBackend._prepare_palace_for_open(palace_path) cached = chromadb.PersistentClient(path=palace_path) @@ -1405,17 +1410,12 @@ def _client(self, palace_path: str): # Per-process record of palaces that have already had the cold-start # quarantine invoked at least once. The proactive HNSW checks are a - # *cold-start* protection — they catch segments that arrive stale relative + # *cold-start* protection -- they catch segments that arrive stale relative # to ``chroma.sqlite3`` or invalid on disk (e.g. cross-machine replication, - # partial restore, crashed-mid-write). Once a long-running process has - # opened the palace cleanly, re-firing the stale check on every reconnect - # is a *runtime thrash*: the daemon's own writes bump sqlite mtime but HNSW - # flushes batch on chromadb's internal cadence, so the mtime gap naturally - # exceeds the threshold under steady write load even though nothing is - # corrupt. - # Real runtime drift is still handled — palace-daemon's ``_auto_repair`` - # calls :func:`quarantine_stale_hnsw` directly on observed HNSW errors, - # which bypasses this gate. + # partial restore, crashed-mid-write). The gate is cleared whenever the + # palace changes on disk (inode swap, mtime bump, or file appearance), so + # external writes that drift HNSW segments are caught on the next open + # without requiring a full process restart. # # Thread-safety: this set is mutated without a lock. Two concurrent # ``make_client()`` calls for the same palace can both pass the @@ -1442,12 +1442,12 @@ def _prepare_palace_for_open(palace_path: str) -> None: ``index_metadata.pickle`` that fails to load, so chromadb opens against an empty index instead of crashing on the unloadable pickle (#1266 / PR #1285). - 4. ``quarantine_stale_hnsw`` — also gated by :attr:`_quarantined_paths` - so it fires once per palace per process. This is the SIGSEGV - prevention path for stale HNSW segments (see #1121, #1132, #1263); - wiring it through this helper means CLI mining, search, repair, - and status all benefit, not just the legacy ``make_client`` - callers. + 4. ``quarantine_stale_hnsw`` -- gated by :attr:`_quarantined_paths` + so it fires once per palace until the gate is re-armed by a + disk change. This is the SIGSEGV prevention path for stale + HNSW segments (see #1121, #1132, #1263); wiring it through + this helper means CLI mining, search, repair, and status all + benefit, not just the legacy ``make_client`` callers. Idempotent: safe to call from any code path that is about to open or re-open a palace. The ``_quarantined_paths`` gate prevents thrash on @@ -1468,9 +1468,8 @@ def make_client(palace_path: str): own client cache. New code should obtain a collection through :meth:`get_collection` which manages caching internally. - Quarantines HNSW segments **once per palace per process**. See - :attr:`_quarantined_paths` for the rationale (cold-start protection - vs. runtime thrash on steady-write daemons). + Quarantines HNSW segments on first open and after any detected + disk change. See :attr:`_quarantined_paths` for the gate logic. """ ChromaBackend._prepare_palace_for_open(palace_path) return chromadb.PersistentClient(path=palace_path) diff --git a/mempalace/mcp_server.py b/mempalace/mcp_server.py index 33dafd2196..e0599305c8 100644 --- a/mempalace/mcp_server.py +++ b/mempalace/mcp_server.py @@ -475,11 +475,13 @@ def _get_client(): mtime_changed = current_mtime != 0.0 and abs(current_mtime - _palace_db_mtime) > 0.01 if _client_cache is None or inode_changed or mtime_changed: - # Run the HNSW capacity probe BEFORE chromadb opens the segment — + # Run the HNSW capacity probe BEFORE chromadb opens the segment -- # if the index is severely undersized, segment load can segfault # the whole MCP server (#1222). The probe is pure sqlite + - # metadata-pickle read; never touches the HNSW binary files. + # metadata read; never touches the HNSW binary files. _refresh_vector_disabled_flag() + if inode_changed or mtime_changed: + ChromaBackend._quarantined_paths.discard(_config.palace_path) _client_cache = ChromaBackend.make_client(_config.palace_path) _collection_cache = None _metadata_cache = None @@ -497,9 +499,9 @@ def _get_collection(create=False): cached client/collection went stale — typically after the chromadb rust bindings invalidated a handle following an out-of-band write — leaving the LLM with no diagnostic and no recovery path. The retry - forces ``_get_client()`` to rebuild from scratch (which re-runs - ``quarantine_stale_hnsw`` per #1322), so the second attempt heals the - common stale-handle / stale-HNSW case automatically. + forces ``_get_client()`` to rebuild the chromadb client from + scratch, so the second attempt heals the common stale-handle case + automatically. """ global _client_cache, _collection_cache, _metadata_cache, _metadata_cache_time for attempt in range(2): @@ -570,9 +572,9 @@ def _get_collection(create=False): ) if attempt == 0: # Reset all caches so the next attempt forces _get_client() - # to rebuild the chromadb client from scratch — that path - # re-runs quarantine_stale_hnsw (#1322) and reopens the - # collection cleanly, healing the common stale-handle case. + # to rebuild the chromadb client from scratch, reopening + # the collection cleanly and healing the common + # stale-handle case. _client_cache = None _collection_cache = None _metadata_cache = None @@ -1916,6 +1918,7 @@ def tool_reconnect(): _collection_cache = None _palace_db_inode = 0 _palace_db_mtime = 0.0 + ChromaBackend._quarantined_paths.discard(_config.palace_path) # Force probe re-run on next _get_client by clearing the flag now; # _refresh_vector_disabled_flag will re-set it if the divergence # still applies after the reconnect. diff --git a/tests/test_backends.py b/tests/test_backends.py index 95ad69f1ff..0125b75ca1 100644 --- a/tests/test_backends.py +++ b/tests/test_backends.py @@ -1127,6 +1127,42 @@ def _spy(path, stale_seconds=300.0): ) +def test_client_rearms_quarantine_on_mtime_change(tmp_path, monkeypatch): + """When the DB file's mtime changes between ``_client()`` calls (external + in-place write), the quarantine gate re-arms so HNSW checks run again. + + Before #1573, the gate was only cleared on *inode* change (full palace + replacement); mtime-only changes left the gate armed, so long-running + processes were blind to external drift.""" + palace_path = str(tmp_path / "palace") + os.makedirs(palace_path, exist_ok=True) + db_file = Path(palace_path) / "chroma.sqlite3" + db_file.write_text("") + + monkeypatch.setattr(ChromaBackend, "_quarantined_paths", set()) + + calls: list[str] = [] + + def _spy(path, stale_seconds=300.0): + calls.append(path) + return [] + + monkeypatch.setattr("mempalace.backends.chroma.quarantine_stale_hnsw", _spy) + + backend = ChromaBackend() + try: + backend._client(palace_path) + assert len(calls) == 1, "quarantine should fire on first open" + + _, cached_mtime = backend._freshness[palace_path] + os.utime(str(db_file), (cached_mtime + 1.0, cached_mtime + 1.0)) + + backend._client(palace_path) + assert len(calls) == 2, "quarantine should re-fire after mtime change (gate re-armed)" + finally: + backend.close() + + # ── _pin_hnsw_threads (per-process retrofit, separate from this PR's gate) ── @@ -1426,7 +1462,9 @@ class DummyClient: ] -def test_chroma_backend_stale_quarantine_is_cold_start_only_on_refresh(tmp_path, monkeypatch): +def test_chroma_backend_quarantine_rearms_on_mtime_refresh(tmp_path, monkeypatch): + """When the DB mtime changes between ``_client()`` calls, the quarantine + gate re-arms and the HNSW safety checks run again (#1573).""" palace = tmp_path / "palace" palace.mkdir() (palace / "chroma.sqlite3").write_text("") @@ -1470,6 +1508,8 @@ class DummyClient: ("stale", str(palace)), ("collection_type", str(palace)), ("blob", str(palace)), + ("invalid", str(palace)), + ("stale", str(palace)), ] diff --git a/tests/test_mcp_server.py b/tests/test_mcp_server.py index 1870de55a3..1563d40a7f 100644 --- a/tests/test_mcp_server.py +++ b/tests/test_mcp_server.py @@ -2718,6 +2718,57 @@ def close(self): assert mcp_server._kg_by_path == {} + def test_tool_reconnect_rearms_quarantine_gate(self, monkeypatch): + """``tool_reconnect`` must clear the per-process quarantine gate so + HNSW safety checks re-run on the next open (#1573).""" + from mempalace import mcp_server + from mempalace.backends.chroma import ChromaBackend + + palace_path = "/test/palace/quarantine_rearm" + gate = {palace_path} + monkeypatch.setattr(ChromaBackend, "_quarantined_paths", gate) + monkeypatch.setattr(mcp_server, "_config", type("C", (), {"palace_path": palace_path})()) + monkeypatch.setattr(mcp_server, "_get_collection", lambda: None) + + mcp_server.tool_reconnect() + + assert palace_path not in gate, ( + "tool_reconnect should clear quarantine gate for the palace path" + ) + + def test_get_client_rearms_quarantine_on_reconnect(self, monkeypatch, config, palace_path, kg): + """``_get_client`` must clear the quarantine gate before calling + ``make_client`` so HNSW safety checks re-run on reconnect (#1573).""" + _patch_mcp_server(monkeypatch, config, kg) + from mempalace import mcp_server + from mempalace.backends.chroma import ChromaBackend + + _client, _col = _get_collection(palace_path, create=True) + del _client + + mcp_server._get_collection() + + assert config.palace_path in ChromaBackend._quarantined_paths + + old_mtime = mcp_server._palace_db_mtime + monkeypatch.setattr(mcp_server, "_palace_db_mtime", old_mtime - 10.0) + + quarantine_calls: list[str] = [] + original_prepare = ChromaBackend._prepare_palace_for_open + + @staticmethod + def spy_prepare(path): + quarantine_calls.append(path) + original_prepare(path) + + monkeypatch.setattr(ChromaBackend, "_prepare_palace_for_open", spy_prepare) + + mcp_server._get_client() + + assert len(quarantine_calls) == 1, ( + "_get_client should call _prepare_palace_for_open on reconnect" + ) + def test_call_kg_retries_after_concurrent_close(self, monkeypatch): """A KG closed mid-handler must trigger a one-shot retry with a fresh instance — not surface a -32000 to the MCP client."""