diff --git a/mempalace/mcp_server.py b/mempalace/mcp_server.py index 1f3cef2..99ae542 100644 --- a/mempalace/mcp_server.py +++ b/mempalace/mcp_server.py @@ -9,7 +9,7 @@ Tools (read): mempalace_list_wings — all wings with drawer counts mempalace_list_rooms — rooms within a wing mempalace_get_taxonomy — full wing → room → count tree - mempalace_search — semantic search, optional wing/room filter + mempalace_search — semantic search, optional wing/room/source_file filter mempalace_check_duplicate — check if content already exists before filing Tools (write): @@ -1152,6 +1152,41 @@ def _sanitize_optional_name(value: str = None, field_name: str = "name") -> str: return sanitize_name(value, field_name) +# Bounds the whole stored source_file string (often an absolute path), so it is +# Linux PATH_MAX rather than the 128-char wing/room NAME limit. +_MAX_SOURCE_FILE_LENGTH = 4096 + + +def _sanitize_optional_source_file(value: str = None) -> str: + """Validate an optional source_file search filter (#1815). + + Unlike wing/room, a source_file is a path: ``/``, ``\\`` and ``.`` are + legal, so it is NOT run through ``sanitize_name`` (which rejects path + characters as traversal attempts). The value is matched verbatim as a + ChromaDB metadata-equality / parameterized-SQL value — never used as a + filesystem path — so there is no traversal risk to guard against. A null + byte or a pathological length can still upset the backend (chromadb + add/upsert chokes on null bytes / lone surrogates, #1235), so guard those + for parity with ``sanitize_name``. Blank / whitespace-only is "no filter". + """ + if value is None: + return None + if not isinstance(value, str): + raise ValueError("source_file must be a string") + value = value.strip() + if not value: + return None + if "\x00" in value: + raise ValueError("source_file contains null bytes") + if value != strip_lone_surrogates(value): + raise ValueError("source_file contains invalid surrogate characters") + if len(value) > _MAX_SOURCE_FILE_LENGTH: + raise ValueError( + f"source_file exceeds maximum length of {_MAX_SOURCE_FILE_LENGTH} characters" + ) + return value + + # ==================== READ TOOLS ==================== @@ -1590,6 +1625,7 @@ def tool_search( limit: int = 5, wing: str = None, room: str = None, + source_file: str = None, max_distance: float = 1.5, min_similarity: float = None, context: str = None, @@ -1598,6 +1634,7 @@ def tool_search( try: wing = _sanitize_optional_name(wing, "wing") room = _sanitize_optional_name(room, "room") + source_file = _sanitize_optional_source_file(source_file) except ValueError as e: return {"error": str(e)} # Backwards compat: accept old name @@ -1616,6 +1653,7 @@ def tool_search( palace_path=_config.palace_path, wing=wing, room=room, + source_file=source_file, n_results=limit, max_distance=dist, vector_disabled=_vector_disabled, @@ -1633,6 +1671,7 @@ def tool_search( palace_path=_config.palace_path, wing=wing, room=room, + source_file=source_file, n_results=limit, max_distance=dist, vector_disabled=_vector_disabled, @@ -3483,6 +3522,15 @@ TOOLS = { }, "wing": {"type": "string", "description": "Filter by wing (optional)"}, "room": {"type": "string", "description": "Filter by room (optional)"}, + "source_file": { + "type": "string", + "description": ( + "Filter to one exact source_file (optional). Matches the full " + "stored path exactly (leading/trailing whitespace trimmed); no " + "glob or basename matching. Pass the value from a result's " + "'source_path' field; the displayed 'source_file' is only a basename." + ), + }, "max_distance": { "type": "number", "description": "Max cosine distance threshold (0=identical, 2=opposite). Results further than this are dropped. Lower = stricter. Default 1.5. Set to 0 to disable.", diff --git a/mempalace/searcher.py b/mempalace/searcher.py index 239367b..eac488f 100644 --- a/mempalace/searcher.py +++ b/mempalace/searcher.py @@ -226,15 +226,24 @@ def _hybrid_rank( return results -def build_where_filter(wing: str = None, room: str = None) -> dict: - """Build ChromaDB where filter for wing/room filtering.""" - if wing and room: - return {"$and": [{"wing": wing}, {"room": room}]} - elif wing: - return {"wing": wing} - elif room: - return {"room": room} - return {} +def build_where_filter(wing: str = None, room: str = None, source_file: str = None) -> dict: + """Build a ChromaDB where filter from optional wing/room/source_file. + + ChromaDB needs a ``$and`` only when ≥2 clauses are present; a single + clause is returned bare and zero clauses yield an empty filter (#1815). + """ + clauses = [] + if wing: + clauses.append({"wing": wing}) + if room: + clauses.append({"room": room}) + if source_file: + clauses.append({"source_file": source_file}) + if not clauses: + return {} + if len(clauses) == 1: + return clauses[0] + return {"$and": clauses} def _extract_drawer_ids_from_closet(closet_doc: str) -> list: @@ -476,6 +485,7 @@ def _bm25_only_via_sqlite( palace_path: str, wing: str = None, room: str = None, + source_file: str = None, n_results: int = 5, max_candidates: int = 500, _include_internal: bool = False, @@ -511,7 +521,7 @@ def _bm25_only_via_sqlite( def _metadata_filter_sql(row_id_expr: str) -> tuple[str, list[str]]: clauses = [] params = [] - for key, value in (("wing", wing), ("room", room)): + for key, value in (("wing", wing), ("room", room), ("source_file", source_file)): if not value: continue clauses.append( @@ -624,7 +634,7 @@ def _bm25_only_via_sqlite( if not candidate_ids: return { "query": query, - "filters": {"wing": wing, "room": room}, + "filters": {"wing": wing, "room": room, "source_file": source_file}, "total_before_filter": 0, "results": [], "fallback": "bm25_only_via_sqlite", @@ -660,6 +670,8 @@ def _bm25_only_via_sqlite( continue if room and meta.get("room") != room: continue + if source_file and meta.get("source_file") != source_file: + continue full_source = meta.get("source_file", "") or "" candidates.append( { @@ -667,6 +679,7 @@ def _bm25_only_via_sqlite( "wing": meta.get("wing", "unknown"), "room": meta.get("room", "unknown"), "source_file": Path(full_source).name if full_source else "?", + "source_path": full_source, "created_at": meta.get("filed_at", "unknown"), # No vector distance available in BM25-only mode. "similarity": None, @@ -702,7 +715,7 @@ def _bm25_only_via_sqlite( return { "query": query, - "filters": {"wing": wing, "room": room}, + "filters": {"wing": wing, "room": room, "source_file": source_file}, "total_before_filter": len(candidates), "results": hits, "fallback": "bm25_only_via_sqlite", @@ -718,6 +731,7 @@ def _merge_bm25_union_candidates( room: str, n_results: int, max_distance: float = 0.0, + source_file: str = None, ) -> None: """Append top-K backend lexical candidates into ``hits`` in place. @@ -744,7 +758,7 @@ def _merge_bm25_union_candidates( if max_distance > 0.0: return - where = build_where_filter(wing, room) + where = build_where_filter(wing, room, source_file) try: lexical = drawers_col.lexical_search( query=query, @@ -767,6 +781,7 @@ def _merge_bm25_union_candidates( "wing": meta.get("wing", "unknown"), "room": meta.get("room", "unknown"), "source_file": Path(full_source).name if full_source else "?", + "source_path": full_source, "created_at": meta.get("filed_at", "unknown"), "similarity": None, "distance": None, @@ -832,6 +847,7 @@ def _apply_candidate_strategy( room: str, n_results: int, max_distance: float = 0.0, + source_file: str = None, ) -> None: """Dispatch to the registered merger for ``strategy``. @@ -840,7 +856,16 @@ def _apply_candidate_strategy( """ merger = _CANDIDATE_MERGERS[strategy] if merger is not None: - merger(hits, drawers_col, query, wing, room, n_results, max_distance=max_distance) + merger( + hits, + drawers_col, + query, + wing, + room, + n_results, + max_distance=max_distance, + source_file=source_file, + ) def _finalize_candidate_hits( @@ -853,6 +878,7 @@ def _finalize_candidate_hits( room: str, n_results: int, max_distance: float, + source_file: str = None, ) -> tuple: try: _apply_candidate_strategy( @@ -864,6 +890,7 @@ def _finalize_candidate_hits( room, n_results, max_distance=max_distance, + source_file=source_file, ) except UnsupportedCapabilityError: return [], { @@ -905,6 +932,7 @@ def _vector_disabled_search( room: str, n_results: int, collection_name: str, + source_file: str = None, ) -> dict: try: backend_name = resolve_backend_name(palace_path) @@ -924,6 +952,7 @@ def _vector_disabled_search( palace_path, wing=wing, room=room, + source_file=source_file, n_results=n_results, collection_name=collection_name, ) @@ -957,7 +986,9 @@ def _open_search_collection(palace_path: str, collection_name: str): } -def _query_drawers_with_filter_fallback(drawers_col, dkwargs, query, n_results, wing, room): +def _query_drawers_with_filter_fallback( + drawers_col, dkwargs, query, n_results, wing, room, source_file=None +): """Run the filtered drawer query, falling back to an unfiltered query plus a Python-side post-filter when ChromaDB raises on the filtered query. @@ -965,7 +996,7 @@ def _query_drawers_with_filter_fallback(drawers_col, dkwargs, query, n_results, "Error finding id" even when unfiltered search works fine — it happens when drawers are ingested via two different paths (e.g. bulk import vs MCP tool calls), leaving the vector index inconsistent with the metadata store. We - retry unfiltered (over-fetching) and re-apply the wing/room filter in Python. + retry unfiltered (over-fetching) and re-apply the wing/room/source_file filter in Python. See #1245 / #1035. """ where = dkwargs.get("where") @@ -994,6 +1025,8 @@ def _query_drawers_with_filter_fallback(drawers_col, dkwargs, query, n_results, continue if room and meta.get("room") != room: continue + if source_file and meta.get("source_file") != source_file: + continue fdocs.append(doc) fmetas.append(meta) fdists.append(dist) @@ -1005,6 +1038,7 @@ def search_memories( palace_path: str, wing: str = None, room: str = None, + source_file: str = None, n_results: int = 5, max_distance: float = 0.0, vector_disabled: bool = False, @@ -1020,6 +1054,8 @@ def search_memories( palace_path: Path to the ChromaDB palace directory. wing: Optional wing filter. room: Optional room filter. + source_file: Optional exact source_file filter. Matches the full + stored source_file value verbatim (#1815). n_results: Max results to return. max_distance: Max cosine distance threshold. The palace collection uses cosine distance (hnsw:space=cosine) — 0 = identical, 2 = opposite. @@ -1059,6 +1095,7 @@ def search_memories( room=room, n_results=n_results, collection_name=collection_name, + source_file=source_file, ) drawers_col, open_error = _open_search_collection(palace_path, collection_name) @@ -1066,7 +1103,7 @@ def search_memories( return open_error metric = _metric_for_collection(drawers_col) - where = build_where_filter(wing, room) + where = build_where_filter(wing, room, source_file) # Hybrid retrieval: always query drawers directly (the floor), then use # closet hits to boost rankings. Closets are a ranking SIGNAL, never a @@ -1084,7 +1121,7 @@ def search_memories( if where: dkwargs["where"] = where drawer_results = _query_drawers_with_filter_fallback( - drawers_col, dkwargs, query, n_results, wing, room + drawers_col, dkwargs, query, n_results, wing, room, source_file ) except Exception as e: return {"error": f"Search error: {e}"} @@ -1156,7 +1193,10 @@ def search_memories( "text": doc, "wing": meta.get("wing", "unknown"), "room": meta.get("room", "unknown"), + # source_file is the basename (display); source_path is the full + # stored value, the round-trippable key for the source_file filter. "source_file": Path(source).name if source else "?", + "source_path": source, "created_at": meta.get("filed_at", "unknown"), "similarity": round(_distance_to_similarity(effective_dist, metric), 3), "distance": round(dist, 4), @@ -1254,13 +1294,14 @@ def search_memories( room=room, n_results=n_results, max_distance=max_distance, + source_file=source_file, ) if strategy_error: return strategy_error return { "query": query, - "filters": {"wing": wing, "room": room}, + "filters": {"wing": wing, "room": room, "source_file": source_file}, "total_before_filter": len(_first_or_empty(drawer_results, "documents")), "results": hits, } diff --git a/tests/test_hybrid_candidate_union.py b/tests/test_hybrid_candidate_union.py index 0771001..7c1a129 100644 --- a/tests/test_hybrid_candidate_union.py +++ b/tests/test_hybrid_candidate_union.py @@ -213,6 +213,26 @@ class TestCandidateUnion: f"(basename collision would drop one); got sources={sources}" ) + def test_union_respects_source_file_filter(self, tmp_path): + """Union pulls BM25 candidates from sqlite FTS5 directly; the + source_file filter must constrain that pool too, not just the vector + path — otherwise union silently re-injects other sources (#1815).""" + palace = str(tmp_path / "palace") + _seed_drawers(palace) + result = search_memories( + _NARRATIVE_QUERY, + palace, + n_results=5, + candidate_strategy="union", + source_file="ticket_D2.md", + ) + sources = {h["source_file"] for h in result["results"]} + assert sources <= {"ticket_D2.md"}, ( + f"union must honor source_file on the BM25 pool; got {sources}" + ) + # The BM25-strong brand-voice doc must NOT leak past the filter. + assert "brand_voice_D4.md" not in sources + class TestHybridRankTolerantOfMissingDistance: """``_hybrid_rank`` accepts ``distance=None`` — required for BM25-only diff --git a/tests/test_hybrid_search.py b/tests/test_hybrid_search.py index a2672de..35aa579 100644 --- a/tests/test_hybrid_search.py +++ b/tests/test_hybrid_search.py @@ -133,3 +133,43 @@ class TestClosetMetadata: assert h["matched_via"] == "drawer" assert "closet_preview" not in h assert h["closet_boost"] == 0.0 + + +# ── source_file filter scopes both drawer and closet queries (#1815) ────── + + +class TestSourceFileFilter: + def test_source_file_filter_excludes_other_sources(self, tmp_path): + palace = str(tmp_path / "palace") + _seed_drawers(palace) + result = search_memories( + "Kafka consumer rebalance timeout", + palace, + n_results=5, + source_file="fixture_D4.md", + ) + ids = [h["source_file"] for h in result["results"]] + assert ids, "the matching source_file drawer should be returned" + assert set(ids) == {"fixture_D4.md"} + + def test_source_file_filter_overrides_closet_boost_for_other_source(self, tmp_path): + # A strong closet pointing at D1 must NOT leak D1 in when the search + # is scoped to a different source_file — the where clause is applied + # to the closet query too, not just the drawer query. + palace = str(tmp_path / "palace") + _seed_drawers(palace) + _seed_strong_closet_for( + palace, + drawer_id="D1", + source_file="fixture_D1.md", + topics=["Kafka queue tuning", "consumer rebalance config"], + ) + result = search_memories( + "Kafka consumer rebalance", + palace, + n_results=5, + source_file="fixture_D4.md", + ) + ids = [h["source_file"] for h in result["results"]] + assert "fixture_D1.md" not in ids + assert set(ids) <= {"fixture_D4.md"} diff --git a/tests/test_mcp_server.py b/tests/test_mcp_server.py index 7e745a4..f232c88 100644 --- a/tests/test_mcp_server.py +++ b/tests/test_mcp_server.py @@ -1105,6 +1105,91 @@ class TestSearchTool: result = tool_search(query="database", room="backend") assert all(r["room"] == "backend" for r in result["results"]) + def test_search_with_source_file_filter( + self, monkeypatch, config, palace_path, seeded_collection, kg + ): + _patch_mcp_server(monkeypatch, config, kg) + from mempalace.mcp_server import tool_search + + result = tool_search(query="authentication module", source_file="auth.py") + assert result["results"] + assert all(r["source_file"] == "auth.py" for r in result["results"]) + assert result["filters"]["source_file"] == "auth.py" + + def test_search_source_file_allows_path_separators( + self, monkeypatch, config, palace_path, seeded_collection, kg + ): + # Unlike wing/room, a source_file is a path — '/' must NOT be rejected + # as a path-traversal attempt the way sanitize_name() would. + _patch_mcp_server(monkeypatch, config, kg) + from mempalace.mcp_server import tool_search + + result = tool_search(query="authentication", source_file="/abs/path/to/auth.py") + assert "error" not in result + + def test_search_blank_source_file_ignored( + self, monkeypatch, config, palace_path, seeded_collection, kg + ): + _patch_mcp_server(monkeypatch, config, kg) + from mempalace.mcp_server import tool_search + + result = tool_search(query="JWT authentication", source_file=" ") + assert "results" in result + assert result["filters"]["source_file"] is None + + def test_search_rejects_null_byte_source_file( + self, monkeypatch, config, palace_path, seeded_collection, kg + ): + # A null byte in a metadata where-value can crash chromadb add/upsert + # (#1235 lineage); reject it cleanly the way sanitize_name does. + _patch_mcp_server(monkeypatch, config, kg) + from mempalace.mcp_server import tool_search + + result = tool_search(query="JWT", source_file="bad\x00null") + assert "error" in result + + def test_search_rejects_overlong_source_file( + self, monkeypatch, config, palace_path, seeded_collection, kg + ): + _patch_mcp_server(monkeypatch, config, kg) + from mempalace.mcp_server import tool_search + + result = tool_search(query="JWT", source_file="x" * 5000) + assert "error" in result + + def test_search_rejects_non_string_source_file( + self, monkeypatch, config, palace_path, seeded_collection, kg + ): + # A non-string source_file (e.g. a JSON number, which the schema's + # string type does not coerce) must yield a clean validation error, + # not an unhandled AttributeError from .strip(). + _patch_mcp_server(monkeypatch, config, kg) + from mempalace.mcp_server import tool_search + + result = tool_search(query="JWT", source_file=42) + assert "error" in result + + def test_search_rejects_lone_surrogate_source_file( + self, monkeypatch, config, palace_path, seeded_collection, kg + ): + # A lone UTF-16 surrogate can crash chromadb (#1235); reject it for + # parity with sanitize_name rather than letting it reach the backend. + _patch_mcp_server(monkeypatch, config, kg) + from mempalace.mcp_server import tool_search + + result = tool_search(query="JWT", source_file="bad\udc80surrogate") + assert "error" in result + + def test_search_accepts_source_file_at_length_boundary( + self, monkeypatch, config, palace_path, seeded_collection, kg + ): + # Exactly _MAX_SOURCE_FILE_LENGTH is allowed (the cap is a strict '>'). + _patch_mcp_server(monkeypatch, config, kg) + from mempalace.mcp_server import _MAX_SOURCE_FILE_LENGTH, tool_search + + result = tool_search(query="JWT", source_file="x" * _MAX_SOURCE_FILE_LENGTH) + assert "error" not in result + def test_search_min_similarity_backwards_compat( self, monkeypatch, config, palace_path, seeded_collection, kg ): diff --git a/tests/test_searcher.py b/tests/test_searcher.py index 236d5b9..c99ae65 100644 --- a/tests/test_searcher.py +++ b/tests/test_searcher.py @@ -9,7 +9,49 @@ from unittest.mock import MagicMock, patch import pytest -from mempalace.searcher import SearchError, search, search_memories +from mempalace.searcher import SearchError, build_where_filter, search, search_memories + + +# ── build_where_filter (unit) ────────────────────────────────────────── + + +class TestBuildWhereFilter: + """build_where_filter composes a ChromaDB where clause from optional + wing / room / source_file constraints (#1815). ChromaDB needs a ``$and`` + only when ≥2 clauses are present; a single clause is returned bare and + zero clauses yield an empty filter.""" + + def test_no_filters_returns_empty(self): + assert build_where_filter() == {} + + def test_wing_only(self): + assert build_where_filter(wing="backend") == {"wing": "backend"} + + def test_room_only(self): + assert build_where_filter(room="auth") == {"room": "auth"} + + def test_wing_and_room(self): + assert build_where_filter(wing="backend", room="auth") == { + "$and": [{"wing": "backend"}, {"room": "auth"}] + } + + def test_source_file_only(self): + assert build_where_filter(source_file="auth.py") == {"source_file": "auth.py"} + + def test_wing_and_source_file(self): + assert build_where_filter(wing="backend", source_file="auth.py") == { + "$and": [{"wing": "backend"}, {"source_file": "auth.py"}] + } + + def test_room_and_source_file(self): + assert build_where_filter(room="auth", source_file="auth.py") == { + "$and": [{"room": "auth"}, {"source_file": "auth.py"}] + } + + def test_wing_room_and_source_file(self): + assert build_where_filter(wing="backend", room="auth", source_file="auth.py") == { + "$and": [{"wing": "backend"}, {"room": "auth"}, {"source_file": "auth.py"}] + } # ── search_memories (API) ────────────────────────────────────────────── @@ -34,6 +76,68 @@ class TestSearchMemories: result = search_memories("code", palace_path, wing="project", room="frontend") assert all(r["wing"] == "project" and r["room"] == "frontend" for r in result["results"]) + def test_source_file_filter(self, palace_path, seeded_collection): + result = search_memories("authentication module", palace_path, source_file="auth.py") + assert result["results"], "exact source_file match should return its drawer" + assert all(r["source_file"] == "auth.py" for r in result["results"]) + + def test_source_file_with_wing_filter(self, palace_path, seeded_collection): + result = search_memories("database", palace_path, wing="project", source_file="db.py") + assert result["results"] + assert all( + r["source_file"] == "db.py" and r["wing"] == "project" for r in result["results"] + ) + + def test_nonmatching_source_file_returns_empty_not_error(self, palace_path, seeded_collection): + result = search_memories("authentication", palace_path, source_file="nope.md") + assert "error" not in result + assert result["results"] == [] + + def test_filters_envelope_includes_source_file(self, palace_path, seeded_collection): + result = search_memories("authentication", palace_path, source_file="auth.py") + assert result["filters"]["source_file"] == "auth.py" + + def test_result_exposes_full_source_path(self, palace_path, seeded_collection): + # The displayed source_file is a basename; source_path carries the full + # stored value so a caller can round-trip it back into a source_file filter. + result = search_memories("authentication module", palace_path) + hit = result["results"][0] + assert hit["source_file"] == "auth.py" + assert hit["source_path"] == "auth.py" + + def test_source_file_filter_matches_full_path_not_basename(self, palace_path): + from mempalace.palace import get_collection + + col = get_collection(palace_path, create=True) + col.upsert( + ids=["fp1"], + documents=["The deploy script restarts the gunicorn workers nightly."], + metadatas=[{"wing": "ops", "room": "deploy", "source_file": "/srv/app/deploy.sh"}], + ) + # The full stored path matches and round-trips via source_path. + hit = search_memories( + "deploy gunicorn workers", palace_path, source_file="/srv/app/deploy.sh" + ) + assert [h["source_path"] for h in hit["results"]] == ["/srv/app/deploy.sh"] + assert [h["source_file"] for h in hit["results"]] == ["deploy.sh"] + # The basename does NOT match — exact full-path semantics only (issue v1). + miss = search_memories("deploy gunicorn workers", palace_path, source_file="deploy.sh") + assert miss["results"] == [] + + def test_source_file_filter_honored_in_bm25_fallback(self, palace_path, seeded_collection): + # vector_disabled routes through _bm25_only_via_sqlite (#1222); the + # source_file filter must hold there too, not silently no-op. + result = search_memories( + "authentication module", + palace_path, + source_file="auth.py", + vector_disabled=True, + collection_name="mempalace_drawers", + ) + assert "error" not in result + assert result["results"], "BM25 fallback should still find the auth drawer" + assert all(r["source_file"] == "auth.py" for r in result["results"]) + def test_n_results_limit(self, palace_path, seeded_collection): result = search_memories("code", palace_path, n_results=2) assert len(result["results"]) <= 2