diff --git a/mempalace/searcher.py b/mempalace/searcher.py index ca0ba46ad1..c1a69cd52f 100644 --- a/mempalace/searcher.py +++ b/mempalace/searcher.py @@ -202,6 +202,31 @@ def _extract_drawer_ids_from_closet(closet_doc: str) -> list: return list(seen.keys()) +def _scoped_source_filter(source_file: str, parent_drawer_id=None) -> dict: + """Build a Chroma ``where`` clause that scopes a query to ``source_file``, + additionally constrained by ``parent_drawer_id`` when one is supplied. + + Two unrelated oversized ``tool_add_drawer`` writes (chunked path from + #1539) can pass the same ``source_file`` (e.g. two pastes tagged + ``"chat.log"``); each call stores its own ``parent_drawer_id`` group + of chunks but the bare ``source_file`` filter pulls chunks from both + groups as if they were siblings (#1580). When the matched chunk + carries a ``parent_drawer_id`` the filter narrows to that logical + group. Otherwise (pre-#1539 drawers, single-chunk writes, and + ``diary_ingest`` drawers grouped by real file path) the original + file-global shape is preserved. Mirrors the conditional-``$and`` + precedent in ``build_where_filter``. + """ + if parent_drawer_id: + return { + "$and": [ + {"source_file": source_file}, + {"parent_drawer_id": parent_drawer_id}, + ] + } + return {"source_file": source_file} + + def _expand_with_neighbors(drawers_col, matched_doc: str, matched_meta: dict, radius: int = 1): """Expand a matched drawer with its ±radius sibling chunks in the same source file. @@ -225,15 +250,20 @@ def _expand_with_neighbors(drawers_col, matched_doc: str, matched_meta: dict, ra if not src or not isinstance(chunk_idx, int): return {"text": matched_doc, "drawer_index": chunk_idx, "total_drawers": None} + # Narrow by ``parent_drawer_id`` when present so chunks from unrelated + # logical drawers sharing ``source_file`` do not stitch (#1580). See + # ``_scoped_source_filter`` for the contract. + parent_id = matched_meta.get("parent_drawer_id") target_indexes = [chunk_idx + offset for offset in range(-radius, radius + 1)] + neighbor_clauses = [ + {"source_file": src}, + {"chunk_index": {"$in": target_indexes}}, + ] + if parent_id: + neighbor_clauses.append({"parent_drawer_id": parent_id}) try: neighbors = drawers_col.get( - where={ - "$and": [ - {"source_file": src}, - {"chunk_index": {"$in": target_indexes}}, - ] - }, + where={"$and": neighbor_clauses}, include=["documents", "metadatas"], ) except Exception: @@ -251,10 +281,16 @@ def _expand_with_neighbors(drawers_col, matched_doc: str, matched_meta: dict, ra else: combined_text = "\n\n".join(doc for _, doc in indexed_docs) - # Cheap total_drawers lookup: metadata-only scan of the source file. + # Cheap total_drawers lookup. When ``parent_drawer_id`` is present the + # count is scoped to that group so the returned number matches the + # text the caller gets back. Without a parent id, the legacy + # file-global count is preserved. total_drawers = None try: - all_meta = drawers_col.get(where={"source_file": src}, include=["metadatas"]) + all_meta = drawers_col.get( + where=_scoped_source_filter(src, parent_id), + include=["metadatas"], + ) total_drawers = len(all_meta.ids) if all_meta.ids else None except Exception: logger.debug("total_drawers lookup failed for %s", src, exc_info=True) @@ -791,6 +827,7 @@ def _finalize_candidate_hits( h.pop("_sort_key", None) h.pop("_source_file_full", None) h.pop("_chunk_index", None) + h.pop("_parent_drawer_id", None) return hits, None @@ -1082,6 +1119,7 @@ def search_memories( "_sort_key": effective_dist, "_source_file_full": source, "_chunk_index": meta.get("chunk_index"), + "_parent_drawer_id": meta.get("parent_drawer_id"), } if closet_preview: entry["closet_preview"] = closet_preview @@ -1102,9 +1140,11 @@ def search_memories( full_source = h.get("_source_file_full") or "" if not full_source: continue + # Narrow by ``parent_drawer_id`` when present so unrelated + # chunked drawers sharing ``source_file`` do not stitch (#1580). try: source_drawers = drawers_col.get( - where={"source_file": full_source}, + where=_scoped_source_filter(full_source, h.get("_parent_drawer_id")), include=["documents", "metadatas"], ) except Exception: diff --git a/tests/test_closets.py b/tests/test_closets.py index e57bf34b33..7ba19f4417 100644 --- a/tests/test_closets.py +++ b/tests/test_closets.py @@ -1518,3 +1518,304 @@ def test_hybrid_search_enrichment_populates_drawer_index_and_total(self, palace_ # Enriched text must include the grep-best chunk plus one neighbor # on each side (chunk boundary may clip). assert "chunk_" in top["text"] + + def test_expand_isolates_chunks_by_parent_drawer_id_when_source_file_shared(self, palace_path): + """Regression for #1580. After #1539 the chunked ``tool_add_drawer`` + path stores per-chunk drawers tagged with a ``parent_drawer_id`` + linking them to the logical group. If two unrelated logical + drawers happen to share the same ``source_file`` (e.g. two pastes + labelled ``source_file="chat.log"``), filtering only by + ``source_file + chunk_index`` pulls chunks from both groups as if + they were sequential neighbors, corrupting the enriched text. + Scoping by ``parent_drawer_id`` when present keeps each logical + group isolated. (``tool_diary_write`` chunks tag a different key + (``parent_entry_id``) and are written without ``source_file``, so + they never reach this enrichment path.) + """ + col = get_collection(palace_path) + source = "shared.log" + # Group A: 2 chunks under parent_drawer_id="drawer_A". + col.upsert( + ids=["drawer_A_chunk_000000", "drawer_A_chunk_000001"], + documents=["alpha-A-chunk-0 content", "alpha-A-chunk-1 content"], + metadatas=[ + { + "wing": "w", + "room": "r", + "source_file": source, + "chunk_index": 0, + "parent_drawer_id": "drawer_A", + "filed_at": "2026-04-13T00:00:00", + }, + { + "wing": "w", + "room": "r", + "source_file": source, + "chunk_index": 1, + "parent_drawer_id": "drawer_A", + "filed_at": "2026-04-13T00:00:00", + }, + ], + ) + # Group B: 2 chunks under the SAME source_file but a different + # parent_drawer_id. Chunk indices intentionally collide with A. + col.upsert( + ids=["drawer_B_chunk_000000", "drawer_B_chunk_000001"], + documents=["bravo-B-chunk-0 content", "bravo-B-chunk-1 content"], + metadatas=[ + { + "wing": "w", + "room": "r", + "source_file": source, + "chunk_index": 0, + "parent_drawer_id": "drawer_B", + "filed_at": "2026-04-13T00:00:00", + }, + { + "wing": "w", + "room": "r", + "source_file": source, + "chunk_index": 1, + "parent_drawer_id": "drawer_B", + "filed_at": "2026-04-13T00:00:00", + }, + ], + ) + + matched_doc = "alpha-A-chunk-0 content" + matched_meta = { + "source_file": source, + "chunk_index": 0, + "parent_drawer_id": "drawer_A", + } + out = _expand_with_neighbors(col, matched_doc, matched_meta, radius=1) + text = out["text"] + # Group A's chunks are returned in chunk_index order. + assert "alpha-A-chunk-0" in text + assert "alpha-A-chunk-1" in text + # No leakage of group B's chunks through the shared source_file key. + assert "bravo-B-chunk-0" not in text + assert "bravo-B-chunk-1" not in text + # total_drawers is scoped to the parent group so the caller sees a + # count consistent with the text returned (2 chunks in group A), + # not 4 (every row sharing the source_file key). + assert out["total_drawers"] == 2 + assert out["drawer_index"] == 0 + + def test_expand_backwards_compat_no_parent_drawer_id_returns_all_source_neighbors( + self, palace_path + ): + """Drawers without a ``parent_drawer_id`` (single-chunk writes, + legacy palaces, ``diary_ingest`` chunks grouped by real file path) + must take the 2-clause fallback (``source_file + chunk_index``) + unchanged, so neighbor expansion still works file-globally for + those callers. + """ + col, _ = self._seed_source_file(palace_path, "/proj/legacy.md", n_chunks=5) + matched_meta = {"source_file": "/proj/legacy.md", "chunk_index": 2} + out = _expand_with_neighbors( + col, "chunk_2 content about topic alpha", matched_meta, radius=1 + ) + # Same expectations as test_expand_returns_matched_plus_neighbors: + # no parent_drawer_id anywhere, so behavior is unchanged. + assert out["total_drawers"] == 5 + assert out["drawer_index"] == 2 + text = out["text"] + assert "chunk_1" in text + assert "chunk_2" in text + assert "chunk_3" in text + + def test_hybrid_search_enrichment_isolates_chunks_across_drawers_sharing_source_file( + self, palace_path + ): + """End-to-end for #1580. Two oversized add_drawer-shape groups + share a ``source_file``, a closet boosts that source, and the + ranked hit lands on group A. The enrichment step in + ``search_memories`` must return only group A's text, not a mix + of A and B chunks stitched as if they were sequential context. + """ + col = get_collection(palace_path) + source = "/proj/shared_log.md" + # Group A: 2 chunks under parent_drawer_id "drawer_proj_log_aaa". + col.upsert( + ids=[ + "drawer_proj_log_aaa_chunk_000000", + "drawer_proj_log_aaa_chunk_000001", + ], + documents=[ + "alpha JWT authentication flow", + "alpha continues the auth narrative", + ], + metadatas=[ + { + "wing": "proj", + "room": "log", + "source_file": source, + "chunk_index": 0, + "parent_drawer_id": "drawer_proj_log_aaa", + "filed_at": "2026-04-13T00:00:00", + }, + { + "wing": "proj", + "room": "log", + "source_file": source, + "chunk_index": 1, + "parent_drawer_id": "drawer_proj_log_aaa", + "filed_at": "2026-04-13T00:00:00", + }, + ], + ) + # Group B: 2 chunks under the SAME source_file but a different + # parent_drawer_id, with content unrelated to the JWT query. + col.upsert( + ids=[ + "drawer_proj_log_bbb_chunk_000000", + "drawer_proj_log_bbb_chunk_000001", + ], + documents=[ + "bravo unrelated topic about database migrations", + "bravo continues with PostgreSQL specifics", + ], + metadatas=[ + { + "wing": "proj", + "room": "log", + "source_file": source, + "chunk_index": 0, + "parent_drawer_id": "drawer_proj_log_bbb", + "filed_at": "2026-04-13T00:00:00", + }, + { + "wing": "proj", + "room": "log", + "source_file": source, + "chunk_index": 1, + "parent_drawer_id": "drawer_proj_log_bbb", + "filed_at": "2026-04-13T00:00:00", + }, + ], + ) + # Closet pointing at group A's first chunk for this source. + closets = get_closets_collection(palace_path) + closets.upsert( + ids=["closet_proj_log_aaa_01"], + documents=["JWT auth|;|→drawer_proj_log_aaa_chunk_000000"], + metadatas=[{"wing": "proj", "room": "log", "source_file": source}], + ) + + result = search_memories("JWT authentication", palace_path) + assert result["results"] + boosted = [h for h in result["results"] if h["matched_via"] == "drawer+closet"] + assert boosted, "hybrid search should mark the closet-agreeing source" + top = boosted[0] + text = top["text"] + # Group A's content is present. + assert "alpha" in text + # Group B's content must not leak in through the shared source_file + # key. The enrichment loop fetches sibling chunks for the matched + # source, and prior to #1580 that fetch ignored parent_drawer_id. + assert "bravo" not in text, ( + "neighbor enrichment leaked group B's chunks through the shared " + "source_file key (see #1580)" + ) + # total_drawers on the enriched hit is scoped to the matched + # parent group (2 chunks in group A), not the full source_file + # row count (4 across both groups). Pins the scoping contract on + # the live enrichment path, not just the helper. + assert top["total_drawers"] == 2 + # Internal scoring-loop keys must be scrubbed before results are + # returned to MCP callers. ``_parent_drawer_id`` is added during + # the #1580 fix and popped in the final cleanup loop alongside + # the existing internal keys. + for h in result["results"]: + assert "_parent_drawer_id" not in h + assert "_source_file_full" not in h + assert "_chunk_index" not in h + assert "_sort_key" not in h + + def test_expand_isolates_asymmetric_groups_under_shared_source_file(self, palace_path): + """Asymmetric coverage: group A has 1 chunk, group B has 3 chunks + under the shared ``source_file``. Catches a regression where + ``total_drawers`` accidentally drifts back to the unscoped + file-global count (4) when one group dominates the row mix. + """ + col = get_collection(palace_path) + source = "asym.log" + col.upsert( + ids=["drawer_solo_chunk_000000"], + documents=["solo-A-chunk-0 content"], + metadatas=[ + { + "wing": "w", + "room": "r", + "source_file": source, + "chunk_index": 0, + "parent_drawer_id": "drawer_solo", + "filed_at": "2026-04-13T00:00:00", + } + ], + ) + col.upsert( + ids=[ + "drawer_trio_chunk_000000", + "drawer_trio_chunk_000001", + "drawer_trio_chunk_000002", + ], + documents=[ + "trio-B-chunk-0 content", + "trio-B-chunk-1 content", + "trio-B-chunk-2 content", + ], + metadatas=[ + { + "wing": "w", + "room": "r", + "source_file": source, + "chunk_index": i, + "parent_drawer_id": "drawer_trio", + "filed_at": "2026-04-13T00:00:00", + } + for i in range(3) + ], + ) + + out = _expand_with_neighbors( + col, + "solo-A-chunk-0 content", + { + "source_file": source, + "chunk_index": 0, + "parent_drawer_id": "drawer_solo", + }, + radius=1, + ) + # Singleton group A: text is the matched chunk, total_drawers == 1. + assert "solo-A-chunk-0" in out["text"] + assert "trio-B-chunk" not in out["text"] + assert out["total_drawers"] == 1 + assert out["drawer_index"] == 0 + + def test_expand_empty_string_parent_drawer_id_treated_as_absent(self, palace_path): + """Contract pin: an empty-string ``parent_drawer_id`` value + degrades to the 2-clause file-global filter (matches the + ``if not src`` empty-string handling for ``source_file`` at + ``searcher.py:239``). Writers in the codebase never emit an + empty parent id, but pinning the contract guards against a + future migration that does and avoids a silent narrow-then- + miss surprise. + """ + col, _ = self._seed_source_file(palace_path, "/proj/empty_parent.md", n_chunks=3) + matched_meta = { + "source_file": "/proj/empty_parent.md", + "chunk_index": 1, + "parent_drawer_id": "", + } + out = _expand_with_neighbors( + col, "chunk_1 content about topic alpha", matched_meta, radius=1 + ) + # Empty parent_drawer_id is treated as absent; full file-global + # neighborhood is returned. Mirrors backwards-compat behavior. + assert out["total_drawers"] == 3 + assert "chunk_0" in out["text"] + assert "chunk_1" in out["text"] + assert "chunk_2" in out["text"]