diff --git a/mempalace/searcher.py b/mempalace/searcher.py index 417ec7c95a..a38cad119a 100644 --- a/mempalace/searcher.py +++ b/mempalace/searcher.py @@ -407,11 +407,109 @@ def _warn_if_legacy_metric(col) -> None: ) +def _hnsw_capacity_diverged(palace_path: str) -> bool: + """Return True if HNSW divergence is severe enough to crash ChromaDB. + + Thin, exception-safe wrapper around + :func:`mempalace.backends.chroma.hnsw_capacity_status`. Used by the + CLI search path to short-circuit to the BM25-only fallback before + opening a Chroma client. Client construction and collection identity + checks can themselves touch the damaged index, so guarding only + ``col.query()`` is too late (#1222 covers the MCP path via the module-level + ``_vector_disabled`` flag; this covers the CLI path). + + A probe that raises falls through to ``False`` so the caller proceeds + to the normal vector path — the underlying query then either succeeds + (probe was a false negative) or raises its own diagnostic error. The + probe itself must never be the thing that crashes search. + """ + try: + from .backends.chroma import hnsw_capacity_status + from .config import get_configured_collection_name + + info = hnsw_capacity_status(palace_path, get_configured_collection_name()) + return bool(info.get("diverged")) + except Exception: + logger.debug("HNSW capacity probe raised; proceeding to vector path", exc_info=True) + return False + + +def _print_search_results_bm25_only( + query: str, palace_path: str, wing: str, room: str, n_results: int +) -> None: + """CLI fallback printer for when HNSW divergence fences off vector search. + + Mirrors the vector-path output shape so users get lexical matches in + the format they expect, plus a clear notice pointing at + ``mempalace repair``. Replaces the silent SIGBUS users otherwise hit + when the CLI called ``col.query()`` against a diverged segment. + """ + result = _bm25_only_via_sqlite( + query=query, + palace_path=palace_path, + wing=wing, + room=room, + n_results=n_results, + ) + hits = result.get("results", []) + + print( + "\n NOTICE: vector search disabled — HNSW index has diverged from SQLite.\n" + " Showing BM25-only results. Run `mempalace repair` to restore " + "vector search.\n" + ) + print(f"{'=' * 60}") + print(f' Results for: "{query}"') + if wing: + print(f" Wing: {wing}") + if room: + print(f" Room: {room}") + print(f"{'=' * 60}\n") + + if not hits: + print(f' No results found for: "{query}"') + return + + for i, hit in enumerate(hits, 1): + bm25 = hit.get("bm25_score", 0.0) + wing_name = hit.get("wing", "?") + room_name = hit.get("room", "?") + source = Path(hit.get("source_file", "?")).name + + print(f" [{i}] {wing_name} / {room_name}") + print(f" Source: {source}") + print(f" Match: bm25={bm25} (vector disabled)") + print() + for line in (hit.get("text", "") or "").strip().split("\n"): + print(f" {line}") + print() + print(f" {'─' * 56}") + + print() + + def search(query: str, palace_path: str, wing: str = None, room: str = None, n_results: int = 5): """ Search the palace. Returns verbatim drawer content. Optionally filter by wing (project) or room (aspect). """ + # Probe a Chroma palace before get_collection(). Opening the client can + # load native index state, and embedder-identity enforcement may call + # collection.count(); both happen before the old query-only guard and can + # hit the same native crash. Non-Chroma backends never use Chroma's HNSW + # files or sqlite-specific fallback and proceed normally. + try: + backend_name = resolve_backend_name(palace_path) + except (BackendMismatchError, KeyError): + # Preserve _open_collection_or_explain's state-specific diagnostics + # for mixed artifacts and unknown backend selections. This probe is + # only an early Chroma safety fence; it must not become a second, + # less-helpful backend validation path. + backend_name = None + + if backend_name == "chroma" and _hnsw_capacity_diverged(palace_path): + return _print_search_results_bm25_only(query, palace_path, wing, room, n_results) + col = _open_collection_or_explain(palace_path, opener=get_collection) if col is None: if not os.path.isdir(palace_path): diff --git a/tests/test_searcher.py b/tests/test_searcher.py index bd24638cfa..c2fe5b9683 100644 --- a/tests/test_searcher.py +++ b/tests/test_searcher.py @@ -11,7 +11,14 @@ from _chroma_palace_helper import make_minimal_chroma_sqlite -from mempalace.searcher import SearchError, build_where_filter, search, search_memories +from mempalace.backends import BackendMismatchError +from mempalace.searcher import ( + SearchError, + build_where_filter, + get_collection, + search, + search_memories, +) # ── build_where_filter (unit) ────────────────────────────────────────── @@ -519,3 +526,119 @@ def test_search_handles_none_document_without_crash(self, fake_palace_path, caps captured = capsys.readouterr() assert "[1]" in captured.out assert "[2]" in captured.out + + def test_search_routes_to_bm25_when_hnsw_diverged(self, fake_palace_path, capsys): + """Regression: `mempalace search` on a diverged HNSW segment must not + segfault ChromaDB's Rust bindings. + + The MCP path gates this via ``_vector_disabled`` (#1222); the CLI + path was missing the gate, so any query into a diverged palace + exited 139 (SIGBUS) at ``chromadb/api/rust.py:_query`` with zero + diagnostic output. This test verifies the CLI now probes + ``hnsw_capacity_status`` and routes to the BM25-only sqlite + fallback instead of calling ``col.query()`` against a segment + that would crash it. + """ + bm25_result = { + "query": "anything", + "filters": {}, + "total_before_filter": 1, + "results": [ + { + "text": "diary entry that matches the query", + "wing": "wing_test", + "room": "diary", + "source_file": "test.jsonl", + "bm25_score": 1.5, + "distance": None, + } + ], + "fallback": "bm25_only_via_sqlite", + "fallback_reason": "vector_search_disabled", + } + with ( + patch("mempalace.searcher.resolve_backend_name", return_value="chroma"), + patch( + "mempalace.backends.chroma.hnsw_capacity_status", + return_value={"diverged": True, "message": "test divergence"}, + ), + patch("mempalace.searcher._bm25_only_via_sqlite", return_value=bm25_result), + patch("mempalace.searcher.get_collection") as mock_get_collection, + ): + search("anything", fake_palace_path) + captured = capsys.readouterr() + # Routed to BM25 before opening Chroma at all. Client construction and + # identity enforcement can touch the same damaged native index, so a + # query-only guard is insufficient. + mock_get_collection.assert_not_called() + # User got actionable output, not a silent crash. + assert "mempalace repair" in captured.out + assert "diary entry that matches" in captured.out + + def test_search_proceeds_to_vector_when_hnsw_healthy(self, fake_palace_path, capsys): + """Paired guard: when HNSW is healthy, the divergence probe must NOT + short-circuit to BM25 — vector search proceeds normally. + + Prevents a regression where the gate accidentally always fires. + """ + mock_col = MagicMock() + mock_col.metadata = {"hnsw:space": "cosine"} + mock_col.query.return_value = { + "documents": [["a matching doc"]], + "metadatas": [[{"source_file": "a.md", "wing": "w", "room": "r"}]], + "distances": [[0.1]], + } + with ( + patch("mempalace.searcher.resolve_backend_name", return_value="chroma"), + patch( + "mempalace.backends.chroma.hnsw_capacity_status", + return_value={"diverged": False, "status": "ok"}, + ), + patch("mempalace.searcher._bm25_only_via_sqlite") as mock_bm25, + patch("mempalace.searcher.get_collection", return_value=mock_col), + ): + search("anything", fake_palace_path) + captured = capsys.readouterr() + # Vector path ran. + mock_col.query.assert_called_once() + # BM25 fallback was NOT invoked. + mock_bm25.assert_not_called() + assert "a matching doc" in captured.out + + def test_search_does_not_run_chroma_probe_for_other_backends(self, fake_palace_path, capsys): + """The HNSW guard is Chroma-specific and must not fence other backends.""" + mock_col = MagicMock() + mock_col.query.return_value = { + "documents": [["backend-native result"]], + "metadatas": [[{"source_file": "native.md", "wing": "w", "room": "r"}]], + "distances": [[0.1]], + } + with ( + patch("mempalace.searcher.resolve_backend_name", return_value="sqlite_exact"), + patch("mempalace.backends.chroma.hnsw_capacity_status") as mock_probe, + patch("mempalace.searcher.get_collection", return_value=mock_col), + ): + search("anything", fake_palace_path) + + mock_probe.assert_not_called() + mock_col.query.assert_called_once() + assert "backend-native result" in capsys.readouterr().out + + @pytest.mark.parametrize( + "resolution_error", + [BackendMismatchError("mixed backend artifacts"), KeyError("unknown_backend")], + ) + def test_search_delegates_backend_resolution_errors_to_open_diagnostic( + self, fake_palace_path, resolution_error + ): + """The early HNSW fence must not replace normal CLI diagnostics.""" + with ( + patch("mempalace.searcher.resolve_backend_name", side_effect=resolution_error), + patch("mempalace.searcher._hnsw_capacity_diverged") as mock_probe, + patch("mempalace.searcher._open_collection_or_explain", return_value=None) as mock_open, + ): + with pytest.raises(SearchError): + search("anything", fake_palace_path) + + mock_probe.assert_not_called() + mock_open.assert_called_once_with(fake_palace_path, opener=get_collection)