fix(mcp): resolve chunked diary entries by their entry_id (#2185) - #2186
Conversation
`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.
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 011e63e5de
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| 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]} |
There was a problem hiding this comment.
When using the Qdrant or pgvector backends, this $or makes every logical get/update/delete scan the entire collection: Qdrant's _requires_local_filter() classifies $or as local-only and _rows() consequently calls _scroll_all() without a filter, while pgvector likewise disables SQL pushdown for $or. This regresses existing chunked parent_drawer_id lookups from a server-side equality filter and can make fetching a single drawer transfer all palace rows; issue two equality-filtered queries and merge their results instead.
AGENTS.md reference: AGENTS.md:L26-L26
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Pull request overview
Fixes logical-id resolution for chunked diary entries so the entry_id returned by mempalace_diary_write can be used with get/update/delete/list (and search results round-trip through mempalace_get_drawer) by reconciling the parent_drawer_id vs parent_entry_id metadata conventions.
Changes:
- Centralizes “logical parent id” resolution in
mcp_serverand updates chunk-group lookup and list collapsing to accept either parent-id key. - Updates diary chunk writes to stamp
parent_drawer_idalongside legacyparent_entry_idfor forward consistency. - Updates search hit
drawer_idreporting to resolve via either parent-id key and adds regression/unit tests.
Reviewed changes
Copilot reviewed 5 out of 5 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
mempalace/mcp_server.py |
Adds shared parent-id-key helpers; resolves chunk groups and collapses list rows using either parent_drawer_id or parent_entry_id; stamps both keys on new diary chunks. |
mempalace/searcher.py |
Makes _result_drawer_id resolve diary chunks via parent_entry_id as well as parent_drawer_id. |
tests/test_mcp_server.py |
Adds regression tests ensuring chunked diary entry_id supports get/list/update/delete and legacy parent_entry_id-only palaces still work. |
tests/test_searcher.py |
Adds unit coverage for _result_drawer_id behavior across both parent-id keys and fallbacks. |
tests/test_closets.py |
Updates an explanatory comment about why diary chunks don’t enter the neighbor-expansion path. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| 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"], | ||
| ) |
| _client2, col = _get_collection(palace_path) | ||
| del _client2 | ||
| assert "".join(col.get()["documents"]) == "B" * 2600, "stale chunks must not survive" |
Fixes #2185.
Problem
mempalace_diary_writereturns anentry_idfor every diary entry, but for entries large enough to be chunked that id is unusable —mempalace_get_drawer,mempalace_update_drawerandmempalace_delete_drawerall answerDrawer not found, and it's the only id the diary tools ever surface to MCP clients. Single-chunk entries are unaffected (their physical id equals theentry_id).Reproduced on
develop: a 5000-char entry atchunk_size=800returns anentry_id, thenget_drawer(entry_id)→{"error": "Drawer not found: diary_wing_..."}.Root cause
Two metadata conventions that never meet. The diary chunking path stamps
parent_entry_idon each chunk, while the logical-id read paths added in #1782 query onlyparent_drawer_id. Both keys mean the same thing — "physical chunk of this logical drawer" — so chunk groups written bydiary_writeare invisible to logical-id resolution.This is the same bug class as #1763, which #1782 fixed for
add_drawerdrawers only.Fix
Read paths now resolve either key, centralized in
_PARENT_ID_KEYSso the two conventions are reconciled in one place:_logical_chunk_group()— matches both keys via$orget/update/deletebyentry_id_collapse_drawer_rows()— groups on either keylist_drawersreturned a chunked entry as N chunk rows instead of one logical drawersearcher._result_drawer_id()— resolves either keyThe last two are beyond what the issue reports, but they're the same convention split and the
$oralone does not cover them — fixing only_logical_chunk_groupwould leave list and search still fragmenting.New diary writes additionally stamp
parent_drawer_idalongsideparent_entry_id, so the conventions converge going forward and the$orbecomes pure legacy compatibility.No migration needed. The read paths still accept the
parent_entry_id-only shape, so palaces written before this fix are repaired on read. A regression test writes chunks in exactly the 3.6.0 shape to lock that in.Notes on blast radius
$oris explicitly supported by all four backends (chroma, qdrant, pgvector, milvus), so the where-clause is portable.$orbranches but is returned once — Chroma dedupes by physical id. The tests assert the joined content equals the original entry verbatim, which would fail on a duplicated chunk.source_file, so neighbor expansion (searcher: _expand_with_neighbors stitches unrelated chunks across MCP drawers sharing empty source_file #1580) returns early on them and is unaffected by the added key. Comment intest_closets.pyupdated, since it cited the parent-id key as the reason rather than the missingsource_file.mempalace_diary_readnever used the chunk/parent logic — it renders rows directly, which is why it always looked fine. Unchanged here.chunk_ids(documenting the bug as intended behavior), and one claimed search rejoins chunks viaparent_entry_id— no search code read that key.Tests
Five new regression tests, mirroring the
add_drawercontract locked in by #1782:diary_write→entry_idfetches verbatim, lists as one logical drawer, deletes every chunkdiary_write→ update byentry_idrewrites the group and leaves no stale chunksparent_entry_id-only chunks resolve for get / list / delete (no-migration proof)_result_drawer_idunit coverage for both keys, plus the plain and missing-metadata fallbacksFull suite green: 3857 passed, 31 skipped.
ruff checkandruff format --checkclean.