fix(mcp): resolve chunked diary entries by their entry_id (#2185)

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)

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:

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"

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) ────────────────────────────────────────────
 
 

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"