Treat chunked Chroma rows as one memory

mempalace_dashboard/server.py

@@ -794,6 +794,65 @@ def _row_dicts(cursor: sqlite3.Cursor) -> list[dict]:
     return [dict(row) for row in cursor.fetchall()]
 
 
+_CHUNKED_DRAWER_ID_RE = re.compile(r"^(?P<parent>.+)_chunk_(?P<ordinal>\d{6})$")
+
+
+def logical_drawer_id(drawer_id: str) -> tuple[str, int | None]:
+    """Return the public drawer id and chunk ordinal for a Chroma row.
+
+    MemPalace stores oversized drawers as ``<drawer_id>_chunk_000000``
+    records.  They are one drawer from the user's perspective: exposing the
+    physical IDs makes a delete or notification dismissal affect only one
+    fragment.  IDs that do not use the exact six-digit suffix remain normal
+    one-row drawers.
+    """
+    match = _CHUNKED_DRAWER_ID_RE.match(drawer_id)
+    if not match:
+        return drawer_id, None
+    return match.group("parent"), int(match.group("ordinal"))
+
+
+def _like_prefix(value: str) -> str:
+    """Escape a SQLite LIKE prefix while retaining its final wildcard."""
+    return value.replace("\\", "\\\\").replace("%", "\\%").replace("_", "\\_") + "%"
+
+
+def _group_physical_drawers(by_id: dict[int, dict], *, light: bool) -> list[dict]:
+    """Collapse Chroma chunk rows into their parent drawer representation."""
+    groups: dict[str, list[tuple[int, int, dict]]] = {}
+    for position, item in enumerate(by_id.values()):
+        parent_id, ordinal = logical_drawer_id(item["drawer_id"])
+        # Normal rows retain their existing ordering. Chunk ordering comes
+        # from the explicit ordinal rather than Chroma's internal numeric id.
+        sort_key = ordinal if ordinal is not None else -1
+        groups.setdefault(parent_id, []).append((position, sort_key, item))
+
+    drawers: list[dict] = []
+    for parent_id, members in groups.items():
+        members.sort(key=lambda member: (member[1], member[0]))
+        first = members[0][2]
+        if len(members) == 1 and members[0][1] == -1:
+            drawers.append(first)
+            continue
+
+        # Metadata is duplicated on every physical Chroma chunk. Take it
+        # from the first logical chunk, while concatenating chunk documents
+        # exactly in their authored ordinal order.
+        item = dict(first)
+        item["drawer_id"] = parent_id
+        item["content"] = "".join(member[2].get("content", "") for member in members)
+        if light:
+            item["truncated"] = len(item["content"]) > PALACE_PREVIEW_CHARS
+            item["content"] = item["content"][:PALACE_PREVIEW_CHARS]
+        else:
+            item["etag"] = content_etag(item["content"])
+        # Light-mode titles must be derived from the same preview body the
+        # client receives; full-mode titles keep the complete body behaviour.
+        item["title"] = extract_title(item["content"], parent_id)
+        drawers.append(item)
+    return drawers
+
+
 def count_drawers() -> int:
     """Cheap COUNT(*) of drawers — used to decide light vs full payload
     without materializing every body. Returns 0 if the DB is missing or
@@ -838,8 +897,34 @@ def read_drawers(light: bool = False, ids: list[int] | None = None) -> list[dict
     id_filter = ""
     params: tuple = ()
     if ids is not None:
-        id_filter = f" and e.id in ({','.join('?' for _ in ids)})"
-        params = tuple(ids)
+        # A full-text hit on one chunk must materialize the whole logical
+        # drawer, otherwise lazy-open/search would show only its matching
+        # fragment. The numeric id branch preserves normal-row behaviour.
+        parent_ids: set[str] = set()
+        try:
+            lookup = sqlite3.connect(PALACE_DB)
+            try:
+                placeholders = ",".join("?" for _ in ids)
+                selected = lookup.execute(
+                    f"select embedding_id from embeddings where id in ({placeholders})",
+                    tuple(ids),
+                ).fetchall()
+            finally:
+                lookup.close()
+            parent_ids = {
+                parent for (embedding_id,) in selected
+                for parent, ordinal in [logical_drawer_id(str(embedding_id))]
+                if ordinal is not None
+            }
+        except sqlite3.OperationalError:
+            return []
+        clauses = [f"e.id in ({','.join('?' for _ in ids)})"]
+        query_params: list[object] = list(ids)
+        if parent_ids:
+            clauses.extend("e.embedding_id like ? escape '\\'" for _ in parent_ids)
+            query_params.extend(_like_prefix(f"{parent}_chunk_") for parent in parent_ids)
+        id_filter = " and (" + " or ".join(clauses) + ")"
+        params = tuple(query_params)
     try:
         con = sqlite3.connect(PALACE_DB)
         con.row_factory = sqlite3.Row
@@ -907,7 +992,7 @@ def read_drawers(light: bool = False, ids: list[int] | None = None) -> list[dict
             item.setdefault("truncated", item.get("truncated", False))
         else:
             item["etag"] = content_etag(item.get("content", ""))
-    drawers = list(by_id.values())
+    drawers = _group_physical_drawers(by_id, light=light)
     enrich_drawers_with_updated_at(drawers)
     return drawers
 
@@ -929,39 +1014,38 @@ def read_single_drawer(drawer_id: str) -> dict | None:
                 from embeddings e
                 join embedding_metadata em on em.id = e.id
                 where e.embedding_id = ?
+                   or e.embedding_id like ? escape '\\'
                 """,
-                (drawer_id,),
+                (drawer_id, _like_prefix(f"{drawer_id}_chunk_")),
             )
         )
         con.close()
     except sqlite3.OperationalError:
         return None
+    rows = [
+        row for row in rows
+        if logical_drawer_id(str(row["embedding_id"]))[0] == drawer_id
+    ]
     if not rows:
         return None
-    item = {
-        "id": rows[0]["id"],
-        "drawer_id": drawer_id,
-        "wing": "unknown",
-        "room": "unknown",
-        "title": "Untitled",
-        "content": "",
-        "source_file": "",
-        "filed_at": "",
-        "added_by": "",
-        "truncated": False,
-        "metadata": {},
-    }
+    by_id: dict[int, dict] = {}
     for row in rows:
+        item = by_id.setdefault(row["id"], {
+            "id": row["id"], "drawer_id": row["embedding_id"],
+            "wing": "unknown", "room": "unknown", "title": "Untitled",
+            "content": "", "source_file": "", "filed_at": "",
+            "added_by": "", "truncated": False, "metadata": {},
+        })
         key = row["key"]
         value = row["value"]
         if key == "chroma:document":
             item["content"] = value or ""
-            item["title"] = extract_title(item["content"], item["drawer_id"])
         elif key in item:
             item[key] = value or ""
         else:
             item["metadata"][key] = value
-    item["etag"] = content_etag(item.get("content", ""))
+    item = _group_physical_drawers(by_id, light=False)[0]
+    item["truncated"] = False
     enrich_drawers_with_updated_at([item])
     return item
 

tests/test_chunk_identity.py

@@ -0,0 +1,113 @@
+from __future__ import annotations
+
+import sqlite3
+import tempfile
+import unittest
+from contextlib import ExitStack
+from pathlib import Path
+from unittest import mock
+
+from mempalace_dashboard import server
+
+
+class ChunkIdentityTests(unittest.TestCase):
+    """Keep Chroma's physical chunk representation out of dashboard state."""
+
+    def setUp(self) -> None:
+        self._tmp = tempfile.TemporaryDirectory()
+        self.addCleanup(self._tmp.cleanup)
+        self._patches = ExitStack()
+        self.addCleanup(self._patches.close)
+        self.db = Path(self._tmp.name) / "chroma.sqlite3"
+        self._patches.enter_context(mock.patch.object(server, "PALACE_DB", self.db))
+        self._patches.enter_context(mock.patch.object(server, "PALACE_PREVIEW_CHARS", 8))
+        self._patches.enter_context(mock.patch.object(server, "enrich_drawers_with_updated_at"))
+        self._create_db()
+
+    def _create_db(self) -> None:
+        con = sqlite3.connect(self.db)
+        con.executescript("""
+            create table embeddings (id integer primary key, embedding_id text not null);
+            create table embedding_metadata (
+                id integer not null,
+                key text not null,
+                string_value text,
+                int_value integer,
+                float_value real,
+                bool_value integer
+            );
+        """)
+        # Deliberately store chunk 1 before chunk 0: dashboard order must use
+        # the chunk suffix, not Chroma's internal row id.
+        con.executemany("insert into embeddings values (?, ?)", [
+            (10, "drawer_story_chunk_000001"),
+            (11, "drawer_story_chunk_000000"),
+            (12, "drawer_normal"),
+        ])
+        metadata = [
+            (10, "chroma:document", "world", None, None, None),
+            (10, "wing", "archive", None, None, None),
+            (10, "room", "stories", None, None, None),
+            (11, "chroma:document", "# Hello\n", None, None, None),
+            (11, "wing", "archive", None, None, None),
+            (11, "room", "stories", None, None, None),
+            (11, "source_file", "import.md", None, None, None),
+            (12, "chroma:document", "# Normal\n\nbody", None, None, None),
+            (12, "wing", "notes", None, None, None),
+            (12, "room", "general", None, None, None),
+        ]
+        con.executemany("insert into embedding_metadata values (?, ?, ?, ?, ?, ?)", metadata)
+        con.commit()
+        con.close()
+
+    def test_list_groups_chunks_as_one_logical_drawer(self) -> None:
+        drawers = server.read_drawers()
+
+        self.assertEqual([drawer["drawer_id"] for drawer in drawers], ["drawer_story", "drawer_normal"])
+        story = drawers[0]
+        self.assertEqual(story["content"], "# Hello\nworld")
+        self.assertEqual(story["title"], "Hello")
+        self.assertEqual(story["wing"], "archive")
+        self.assertEqual(story["source_file"], "import.md")
+        self.assertEqual(story["etag"], server.content_etag("# Hello\nworld"))
+
+    def test_lazy_read_search_and_delete_scope_keep_parent_identity(self) -> None:
+        # A hit on the second physical row must return both chunks as the
+        # logical parent, otherwise opening the search result loses content.
+        matched = server.read_drawers(light=True, ids=[10])
+        self.assertEqual(len(matched), 1)
+        self.assertEqual(matched[0]["drawer_id"], "drawer_story")
+        self.assertEqual(matched[0]["content"], "# Hello\n")
+        self.assertTrue(matched[0]["truncated"])
+
+        drawer = server.read_single_drawer("drawer_story")
+        self.assertIsNotNone(drawer)
+        self.assertEqual(drawer["drawer_id"], "drawer_story")
+        self.assertEqual(drawer["content"], "# Hello\nworld")
+        self.assertFalse(drawer["truncated"])
+
+        label, targets = server.drawers_for_delete({"scope": "drawer", "drawer_id": "drawer_story"})
+        self.assertEqual(label, "memory")
+        self.assertEqual([target["drawer_id"] for target in targets], ["drawer_story"])
+
+        with (
+            mock.patch.object(server, "mempalace_update_drawer", return_value={"success": True}) as update,
+            mock.patch.object(server, "_mark_drawer_self_seen"),
+        ):
+            server.update_memory({
+                "drawer_id": "drawer_story",
+                "content": "# Hello\nupdated world",
+                "etag": drawer["etag"],
+            })
+        update.assert_called_once_with("drawer_story", "# Hello\nupdated world", None, None)
+
+        with (
+            mock.patch.object(server, "mempalace_delete_drawer", return_value={"success": True}) as delete,
+            mock.patch.object(server, "log_version"),
+        ):
+            server.delete_memories({"scope": "drawer", "drawer_id": "drawer_story", "confirm": "DELETE"})
+        delete.assert_called_once_with("drawer_story")
+
+
+if __name__ == "__main__":
+    unittest.main()