From 011e63e5de3396ba86d18fd767f3eed45001993f Mon Sep 17 00:00:00 2001 From: Igor Lins e Silva <4753812+igorls@users.noreply.github.com> Date: Sat, 8 Aug 2026 09:06:53 -0300 Subject: [PATCH] fix(mcp): resolve chunked diary entries by their entry_id (#2185) `mempalace_diary_write` returns an `entry_id` for every diary entry, but for entries large enough to be chunked that id was unusable: get_drawer, update_drawer and delete_drawer all answered "Drawer not found", and list_drawers showed the entry as N unrelated chunk rows. Two metadata conventions never met. The diary chunking path stamped `parent_entry_id` on each chunk, while the logical-id read paths added in #1782 query only `parent_drawer_id`. Both keys mean the same thing -- "physical chunk of this logical drawer" -- so chunk groups written by diary_write were invisible to logical-id resolution. Same bug class as #1763, which #1782 fixed for `add_drawer` drawers only. Read paths now resolve either key via `_PARENT_ID_KEYS`: - `_logical_chunk_group()` matches both with an `$or` (fixes get / update / delete). All four backends support `$or`. - `_collapse_drawer_rows()` groups on either (fixes list_drawers, which the `$or` alone does not cover). - `searcher._result_drawer_id()` resolves either, so a hit on a chunked diary entry reports the id that fetches the whole entry rather than the single chunk that matched. New diary writes also stamp `parent_drawer_id` alongside `parent_entry_id` so the two conventions converge going forward. Because the read paths still accept the `parent_entry_id`-only shape, palaces written before this fix are repaired with no data migration. Diary chunks are written without `source_file`, so neighbor expansion (#1580) returns early on them and is unaffected by the added key. Also drops the comment telling callers to iterate `chunk_ids` (it documented the bug as intended behavior) and a stale claim that search rejoins chunks via `parent_entry_id` -- no search code read that key. --- mempalace/mcp_server.py | 58 ++++++++++++++---- mempalace/searcher.py | 12 +++- tests/test_closets.py | 7 ++- tests/test_mcp_server.py | 126 +++++++++++++++++++++++++++++++++++++++ tests/test_searcher.py | 30 ++++++++++ 5 files changed, 218 insertions(+), 15 deletions(-) diff --git a/mempalace/mcp_server.py b/mempalace/mcp_server.py index bb19a0fae..10f9b6a2b 100644 --- a/mempalace/mcp_server.py +++ b/mempalace/mcp_server.py @@ -2630,10 +2630,44 @@ def _single_drawer_record(col, drawer_id: str): } +# Two write paths stamp the logical-group id under different keys: +# ``tool_add_drawer`` chunks carry ``parent_drawer_id`` (#1539, resolved as a +# logical drawer by #1782) while ``tool_diary_write`` chunks carry +# ``parent_entry_id`` (#1539). Both mean the same thing -- "physical chunk of +# this logical drawer" -- so every read path must resolve either one, or the +# id a write path hands back is unusable with get/update/delete (#2185). +# New diary writes stamp both keys; the read paths below still accept the +# ``parent_entry_id``-only shape so palaces written before this fix keep +# working with no data migration. +_PARENT_ID_KEYS = ("parent_drawer_id", "parent_entry_id") + + +def _logical_parent_id(meta): + """Return the logical-group id from chunk metadata, whichever key holds it. + + Returns ``None`` for rows that are not chunks of a larger drawer. + """ + for key in _PARENT_ID_KEYS: + value = (meta or {}).get(key) + if value: + return value + return None + + +def _logical_parent_where(drawer_id: str) -> dict: + """Chroma ``where`` matching every chunk of ``drawer_id`` under either key. + + A chunk carrying both keys (diary writes after #2185) matches both + branches of the ``$or`` but is still returned once -- Chroma dedupes by + physical id -- so the joined content never repeats a chunk. + """ + return {"$or": [{key: drawer_id} for key in _PARENT_ID_KEYS]} + + def _logical_chunk_group(col, drawer_id: str): try: result = col.get( - where={"parent_drawer_id": drawer_id}, + where=_logical_parent_where(drawer_id), include=["documents", "metadatas"], ) except Exception: @@ -2741,7 +2775,7 @@ def _collapse_drawer_rows(ids, documents, metadatas): for idx, drawer_id in enumerate(ids): doc = documents[idx] if idx < len(documents) else "" meta = _safe_meta(metadatas[idx] if idx < len(metadatas) else {}) - parent_id = meta.get("parent_drawer_id") + parent_id = _logical_parent_id(meta) if parent_id: groups.setdefault(parent_id, []).append( @@ -3861,14 +3895,17 @@ def tool_diary_write(agent_name: str, entry: str, topic: str = "general", wing: # Oversized entry: split into bounded per-chunk drawers so the # embedding model never sees a document above ``chunk_size``. - # Every chunk carries ``parent_entry_id`` so search can rejoin - # them and ``chunk_index`` for ordered reconstruction (#1539). - # Note on ``entry_id`` in the return value: for the chunked - # path the returned ``entry_id`` is the LOGICAL group handle - # (no drawer is stored under that exact id). The physical - # drawer ids are in ``chunk_ids``. Callers wanting to fetch - # by id should iterate ``chunk_ids``; callers wanting to - # query by metadata can filter on ``parent_entry_id``. + # Every chunk carries ``chunk_index`` for ordered reconstruction + # and the group id under BOTH ``parent_entry_id`` (the original + # #1539 key, kept so existing readers and palaces are unaffected) + # and ``parent_drawer_id`` (the key the logical-id read paths were + # built around in #1782) -- see ``_PARENT_ID_KEYS`` (#2185). + # Note on ``entry_id`` in the return value: for the chunked path + # the returned ``entry_id`` is the LOGICAL group handle -- no + # drawer is stored under that exact id, but it resolves through + # ``mempalace_get_drawer`` / ``update_drawer`` / ``delete_drawer`` + # to the whole entry, exactly as an oversized ``add_drawer`` id + # does. The physical drawer ids remain available in ``chunk_ids``. # Use a single batched ``add`` so the embedding pass either # commits all chunks or none — avoids a half-written palace # if the embedding model fails mid-loop. ``col.add`` (not @@ -3890,6 +3927,7 @@ def tool_diary_write(agent_name: str, entry: str, topic: str = "general", wing: **base_metadata, "chunk_index": chunk_idx, "parent_entry_id": entry_id, + "parent_drawer_id": entry_id, } ) col.add(ids=chunk_ids, documents=chunk_docs, metadatas=chunk_metas) diff --git a/mempalace/searcher.py b/mempalace/searcher.py index 44fc4dc9e..71f139fc4 100644 --- a/mempalace/searcher.py +++ b/mempalace/searcher.py @@ -77,8 +77,16 @@ def _aligned_query_ids(results, document_count: int) -> list: def _result_drawer_id(meta, stored_drawer_id): - """Return the ID that round-trips through ``mempalace_get_drawer``.""" - return (meta or {}).get("parent_drawer_id") or stored_drawer_id + """Return the ID that round-trips through ``mempalace_get_drawer``. + + Chunk metadata carries the logical-group id under ``parent_drawer_id`` + (``tool_add_drawer``) or ``parent_entry_id`` (``tool_diary_write``); + resolving both means a hit on a chunked diary entry reports the id that + fetches the WHOLE entry rather than the one chunk that matched (#2185). + Kept in sync with ``mcp_server._PARENT_ID_KEYS``. + """ + meta = meta or {} + return meta.get("parent_drawer_id") or meta.get("parent_entry_id") or stored_drawer_id def _tokenize(text: str, stop_words: frozenset = frozenset()) -> list: diff --git a/tests/test_closets.py b/tests/test_closets.py index ba017b823..6f894629c 100644 --- a/tests/test_closets.py +++ b/tests/test_closets.py @@ -1535,9 +1535,10 @@ def test_expand_isolates_chunks_by_parent_drawer_id_when_source_file_shared(self ``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.) + group isolated. (``tool_diary_write`` chunks are written without + ``source_file``, so they never reach this enrichment path at all -- + it returns early on the missing key -- regardless of which + parent-id key they carry.) """ col = get_collection(palace_path) source = "shared.log" diff --git a/tests/test_mcp_server.py b/tests/test_mcp_server.py index 131854e0b..9b5693b2d 100644 --- a/tests/test_mcp_server.py +++ b/tests/test_mcp_server.py @@ -3076,6 +3076,132 @@ def test_update_drawer_chunked_logical_id_rewrites_group(monkeypatch, config, pa assert listed["drawers"][0]["drawer_id"] == logical_id +def test_diary_write_chunked_logical_id_fetches_deletes_and_lists_as_one( + monkeypatch, config, palace_path, kg +): + """Regression for #2185: the ``entry_id`` returned by a chunked + ``tool_diary_write`` must behave like any other logical drawer id. + + Before the fix the diary chunking path stamped only ``parent_entry_id`` + while logical-id resolution queried only ``parent_drawer_id``, so + get/update/delete answered "Drawer not found" for the one id the diary + tools ever hand to MCP clients, and ``list_drawers`` showed the entry as + N unrelated chunk rows. Mirrors the ``tool_add_drawer`` contract locked + in by #1782. + """ + _patch_mcp_server(monkeypatch, config, kg) + _client, _col = _get_collection(palace_path, create=True) + del _client + + from mempalace.mcp_server import ( + tool_delete_drawer, + tool_diary_write, + tool_get_drawer, + tool_list_drawers, + ) + + oversized = "Z" * 5000 + written = tool_diary_write(agent_name="TestAgent", entry=oversized, topic="general") + assert written["success"] is True + assert written["chunks"] > 1 + + entry_id = written["entry_id"] + + fetched = tool_get_drawer(entry_id) + assert "error" not in fetched + assert fetched["drawer_id"] == entry_id + assert fetched["content"] == oversized, "must return the entry verbatim, not one chunk" + assert fetched["chunks"] == written["chunks"] + assert fetched["chunk_ids"] == written["chunk_ids"] + + listed = tool_list_drawers(wing="wing_testagent", room="diary") + assert listed["total"] == 1, "a chunked entry is ONE logical drawer, not N chunk rows" + assert listed["drawers"][0]["drawer_id"] == entry_id + assert listed["drawers"][0]["chunks"] == written["chunks"] + + deleted = tool_delete_drawer(entry_id) + assert deleted["success"] is True + assert deleted["chunks_deleted"] == written["chunks"] + + missing = tool_get_drawer(entry_id) + assert "error" in missing + + +def test_diary_write_chunked_logical_id_updates_group(monkeypatch, config, palace_path, kg): + """Regression for #2185: updating a chunked diary entry by its + ``entry_id`` must rewrite the whole underlying chunk group.""" + _patch_mcp_server(monkeypatch, config, kg) + _client, _col = _get_collection(palace_path, create=True) + del _client + + from mempalace.mcp_server import ( + tool_diary_write, + tool_get_drawer, + tool_update_drawer, + ) + + written = tool_diary_write(agent_name="TestAgent", entry="A" * 4000, topic="general") + assert written["chunks"] > 1 + entry_id = written["entry_id"] + + updated = tool_update_drawer(entry_id, content="B" * 2600) + assert updated["success"] is True + assert updated["drawer_id"] == entry_id + + fetched = tool_get_drawer(entry_id) + assert fetched["content"] == "B" * 2600 + _client2, col = _get_collection(palace_path) + del _client2 + assert "".join(col.get()["documents"]) == "B" * 2600, "stale chunks must not survive" + + +def test_legacy_diary_chunks_resolve_without_parent_drawer_id(monkeypatch, config, palace_path, kg): + """Regression for #2185: palaces written BEFORE this fix carry diary + chunks tagged only with ``parent_entry_id``. The read paths must resolve + that shape too, so existing palaces are repaired with no data migration. + """ + _patch_mcp_server(monkeypatch, config, kg) + _client, col = _get_collection(palace_path, create=True) + del _client + + from mempalace.mcp_server import ( + tool_delete_drawer, + tool_get_drawer, + tool_list_drawers, + ) + + entry_id = "diary_wing_lily_20260808_142113121027_3e4c74763d73" + # Exactly what mempalace 3.6.0 wrote: parent_entry_id only. + col.upsert( + ids=[f"{entry_id}_chunk_{i:06d}" for i in range(3)], + documents=["legacy-0 ", "legacy-1 ", "legacy-2"], + metadatas=[ + { + "wing": "wing_lily", + "room": "diary", + "type": "diary_entry", + "chunk_index": i, + "parent_entry_id": entry_id, + "filed_at": "2026-08-08T14:21:13", + } + for i in range(3) + ], + ) + + fetched = tool_get_drawer(entry_id) + assert "error" not in fetched, f"legacy diary chunks must resolve; got {fetched}" + assert fetched["content"] == "legacy-0 legacy-1 legacy-2" + assert fetched["chunks"] == 3 + + listed = tool_list_drawers(wing="wing_lily", room="diary") + assert listed["total"] == 1 + assert listed["drawers"][0]["drawer_id"] == entry_id + + deleted = tool_delete_drawer(entry_id) + assert deleted["success"] is True + assert deleted["chunks_deleted"] == 3 + + # ── Delete by source (#1722) ──────────────────────────────────────────── diff --git a/tests/test_searcher.py b/tests/test_searcher.py index 48654f05d..4ab8b525a 100644 --- a/tests/test_searcher.py +++ b/tests/test_searcher.py @@ -15,6 +15,7 @@ from mempalace.backends import BackendMismatchError from mempalace.searcher import ( SearchError, + _result_drawer_id, build_where_filter, get_collection, search, @@ -1017,3 +1018,32 @@ def _hybrid_spy(results, query, **kwargs): searcher.search(query="cat", palace_path=str(tmp_path)) assert captured["stop_words"] == frozenset({"the"}) + + +# ── _result_drawer_id logical-id resolution (#2185) ──────────────────── + + +class TestResultDrawerId: + """The drawer_id on a search hit must round-trip through + ``mempalace_get_drawer``. Chunk rows carry their logical-group id under + ``parent_drawer_id`` (``tool_add_drawer``) or ``parent_entry_id`` + (``tool_diary_write``); both must resolve to the logical id (#2185). + """ + + def test_plain_drawer_returns_stored_id(self): + assert _result_drawer_id({"wing": "w"}, "drawer_abc") == "drawer_abc" + + def test_parent_drawer_id_wins_over_stored_chunk_id(self): + meta = {"parent_drawer_id": "drawer_abc", "chunk_index": 2} + assert _result_drawer_id(meta, "drawer_abc_chunk_000002") == "drawer_abc" + + def test_parent_entry_id_resolves_for_diary_chunks(self): + """Regression for #2185: before the fix a chunked diary hit reported + the physical chunk id, so fetching it returned one chunk of the entry + instead of the whole entry.""" + meta = {"parent_entry_id": "diary_wing_lily_20260808_1", "chunk_index": 3} + stored = "diary_wing_lily_20260808_1_chunk_000003" + assert _result_drawer_id(meta, stored) == "diary_wing_lily_20260808_1" + + def test_missing_metadata_falls_back_to_stored_id(self): + assert _result_drawer_id(None, "drawer_abc") == "drawer_abc"