fix(state): display reads take removed rows' identities, never their bodies

RT-39's display reads dropped the active clause so the dedupe could hide a
message whose newest copy was removed (rewound, or replaced by a newer
copy). That fetched the full body of every `active=0, compacted=0` row
(each in-place compaction leaves one per carried tail row) and discarded
it; `assert_resume_safe` counts none of them, so nothing bounded it.

The display set is read as before (`active OR compacted`), and the removed
rows are read separately with only what the dedupe needs: role, timestamp,
tool-call columns, uids and the stored `display_identity` (content only for
a legacy row without one). The dedupe keys a removed row by its stored
identity. `worker_history` (a worker's display read) now applies the same
removed-copy rule, so the two readers agree.

Review finding by @kshitijk4poor on #140971
Review finding by @dokterdok on #140971
Refs #134889 (RT-39)
Refs #140719

hermes_state_messages.py

@@ -1214,18 +1214,20 @@ class SessionMessagesMixin(SessionDetachedResultsMixin):
         removed: dict[Any, bool] = {}  # copy identity -> its newest copy is removed (rows come in id order)
         identities: dict[tuple[Any, ...], list[Any]] = {}
         for row in rows:
-            copies = None
+            copies, key = None, None
             if drop_removed:
-                key = self._display_dedupe_key(row)
-                copies = [key, *(("uid", uid) for uid in (row["message_uid"],
-                                                          *_uid_list(row["absorbed_message_uids"])) if uid)]
+                # A removed row is read without its body (``_removed_identity_rows``): its stored display identity.
                 gone = not row["active"] and not row["compacted"]
+                stored = row["display_identity"] if gone else None
+                key = None if stored else self._display_dedupe_key(row)
+                copies = [stored or self._display_identity(key), *(
+                    ("uid", uid) for uid in (row["message_uid"], *_uid_list(row["absorbed_message_uids"])) if uid)]
                 removed.update(dict.fromkeys(copies, gone))
                 if gone:
                     continue
             if self._is_model_only_row(row):
                 continue
-            key = copies[0] if copies else self._display_dedupe_key(row)
+            key = key or self._display_dedupe_key(row)
             cur = seen.get(key)
             if cur is None or (row["active"], row["id"]) > (cur["active"], cur["id"]):
                 seen[key] = row
@@ -1430,9 +1432,10 @@ class SessionMessagesMixin(SessionDetachedResultsMixin):
             # (display_order is per-segment and not comparable across segments, so the per-segment
             # display_order paging below does not apply across a lineage.)
             display = include_compacted and not include_inactive  # removed rows too: see drop_removed
-            rows = self._dedupe_display_generations(self._read_all(
-                f"SELECT * FROM messages WHERE session_id IN ({_placeholders(session_ids)})"
-                + ("" if display else active_clause) + " ORDER BY id ASC", session_ids), drop_removed=display)
+            rows = self._read_all(f"SELECT * FROM messages WHERE session_id IN ({_placeholders(session_ids)})"
+                                  + active_clause + " ORDER BY id ASC", session_ids)
+            rows = self._dedupe_display_generations(
+                self._with_removed_identities(rows, session_ids) if display else rows, drop_removed=display)
             rows = rows[::-1][offset:][:limit][::-1] if latest else rows[offset:][:limit]
         elif include_compacted and not include_inactive and self._ensure_display_order(session_id):
             # _read_retrying_ioerr: mode=ro pooled readers see a transient IOERR mid-checkpoint (#100871).
@@ -1531,12 +1534,26 @@ class SessionMessagesMixin(SessionDetachedResultsMixin):
     def _fetch_conversation_rows(self, session_ids: list[str], active_clause: str, *, with_session_id: bool,
                                  with_compacted: bool = False):
         """``_CONVERSATION_ROW_COLUMNS`` rows for *session_ids* ORDER BY id (timestamps are not monotonic
-        and would break tool-call adjacency). ``with_compacted``: also the ``compacted`` column."""
-        return self._read_all(
+        and would break tool-call adjacency). ``with_compacted``: the display set (``active_clause`` is
+        ignored) with the ``compacted`` column, plus the removed rows' identities for ``drop_removed``."""
+        rows = self._read_all(
             f"SELECT {'session_id, ' if with_session_id else ''}{'compacted, ' if with_compacted else ''}"
             f"{self._CONVERSATION_ROW_COLUMNS} "
             f"FROM messages WHERE session_id IN ({_placeholders(session_ids)})"
-            f"{active_clause} ORDER BY id", tuple(session_ids))
+            f"{_DISPLAY_ACTIVE_CLAUSE if with_compacted else active_clause} ORDER BY id", tuple(session_ids))
+        return self._with_removed_identities(rows, session_ids) if with_compacted else rows
+
+    def _with_removed_identities(self, rows, session_ids: list[str]) -> list[Any]:
+        """*rows* (ordered by id) merged with the identities of the lineage's removed rows (``active=0,
+        compacted=0``), which ``_dedupe_display_generations(drop_removed=True)`` needs only to hide a copy, never
+        to show: no body (a legacy row without a stored display identity keeps its content for the key)."""
+        removed = self._read_all(
+            "SELECT session_id, id, role, timestamp, tool_call_id, tool_calls, tool_name, display_kind, display_metadata, "
+            "message_uid, absorbed_message_uids, display_identity, active, compacted, "
+            "CASE WHEN display_identity IS NULL THEN content END AS content "
+            f"FROM messages WHERE session_id IN ({_placeholders(session_ids)}) AND active = 0 AND compacted = 0 "
+            "ORDER BY id", tuple(session_ids))
+        return sorted([*rows, *removed], key=lambda row: row["id"]) if removed else list(rows)
 
     def get_messages_as_conversation(self, session_id: str, include_ancestors: bool = False,
                                      include_inactive: bool = False, repair_alternation: bool = False,
@@ -1550,8 +1567,7 @@ class SessionMessagesMixin(SessionDetachedResultsMixin):
         display = include_compacted and not include_inactive  # removed rows too: see drop_removed
         rows = self._fetch_conversation_rows(
             self._resume_lineage_ids(session_id) if include_ancestors else [session_id],
-            "" if display else self._active_clause(include_inactive, include_compacted), with_session_id=False,
-            with_compacted=display)
+            self._active_clause(include_inactive, include_compacted), with_session_id=False, with_compacted=display)
         if include_compacted:
             rows = self._dedupe_display_generations(rows, drop_removed=display)
         return self._rows_to_conversation(rows, session_id=session_id, include_ancestors=include_ancestors,

hermes_state_worker_compression.py

@@ -189,10 +189,14 @@ def worker_history(db, conn, sid, payload):
     view.authorize(sid, target)
     ids = view._resume_lineage_ids(target) if payload['include_ancestors'] else [target]
     active = db._active_clause(payload['include_inactive'], payload['include_compacted'])
-    rows = conn.execute(f'SELECT {db._CONVERSATION_ROW_COLUMNS} FROM messages '
+    display = payload['include_compacted'] and not payload['include_inactive']
+    rows = conn.execute(f'SELECT {"compacted, " if display else ""}{db._CONVERSATION_ROW_COLUMNS} FROM messages '
                         f'WHERE session_id IN ({",".join("?" for _ in ids)}){active} ORDER BY id', ids).fetchall()
+    if display:  # the owner's display reads' rule: a message whose newest copy was removed stays hidden
+        from hermes_state_messages import SessionMessagesMixin
+        rows = SessionMessagesMixin._with_removed_identities(view, rows, ids)
     if payload['include_compacted']:
-        rows = db._dedupe_display_generations(rows)
+        rows = db._dedupe_display_generations(rows, drop_removed=display)
     return {'messages': db._rows_to_conversation(rows, session_id=target,
         include_ancestors=payload['include_ancestors'], repair_alternation=payload['repair_alternation'],
         include_row_ids=payload['include_row_ids'])}

tests/hermes_state/test_display_read_removed_rows.py

@@ -0,0 +1,88 @@
+"""Display reads hide a message whose newest copy was removed (rewound, or replaced by a newer copy), so
+they need the removed rows' identities, but never their bodies: ``assert_resume_safe`` does not count
+those rows, and repeated in-place compaction leaves one removed copy of the carried tail per round
+(review of #140971). Real SessionDB, real ``archive_and_compact``."""
+import sqlite3
+
+from hermes_state import SessionDB
+
+BIG = 'x' * 100_000
+
+
+def _compacted_repeatedly(db, rounds):
+    db.create_session('s', source='test')
+    db.append_messages_batch('s', [{'role': 'user', 'content': 'U1'}, {'role': 'assistant', 'content': 'A1'},
+                                   {'role': 'user', 'content': 'U2'}, {'role': 'assistant', 'content': BIG}])
+    for n in range(rounds):
+        # /compress with nothing new since: the same verbatim tail (U2 + the big reply) is carried again.
+        db.archive_and_compact('s', [{'role': 'user', 'content': f'[CONTEXT SUMMARY] round {n}'},
+                                     {'role': 'user', 'content': 'U2'}, {'role': 'assistant', 'content': BIG}],
+                               tail_count=2)
+
+
+def _big_bodies_read(db, monkeypatch, read):
+    seen = []
+    real = db._read_all
+
+    def spy(sql, params=()):
+        rows = real(sql, params)
+        seen.extend(r['content'] for r in rows
+                    if isinstance(r, sqlite3.Row) and 'content' in r.keys() and r['content'] == BIG)
+        return rows
+    monkeypatch.setattr(db, '_read_all', spy)
+    shown = read()
+    monkeypatch.undo()
+    return shown, len(seen)
+
+
+def test_display_reads_fetch_no_body_of_a_removed_copy(tmp_path, monkeypatch):
+    with SessionDB(db_path=tmp_path / 'state.db') as db:
+        _compacted_repeatedly(db, 6)
+        removed = db._read_one("SELECT COUNT(*) FROM messages WHERE content = ? AND active = 0 AND compacted = 0",
+                               (BIG,))[0]
+        assert removed >= 5  # the fixture really leaves removed copies behind
+        archived = db._read_one("SELECT COUNT(*) FROM messages WHERE content = ? AND compacted = 1", (BIG,))[0]
+        for read in (
+            lambda: db.get_messages_as_conversation('s', include_ancestors=True, include_compacted=True),
+            lambda: db.get_resume_conversations('s')[1],
+        ):
+            shown, bodies = _big_bodies_read(db, monkeypatch, read)
+            assert sum(1 for m in shown if m.get('content') == BIG) == 1
+            # Only the live copy and the compaction-archived ones (what the resume guard counts) carry a body.
+            assert bodies <= 1 + archived, (bodies, archived, removed)
+
+
+def test_a_rewound_turn_stays_hidden_without_its_body(tmp_path):
+    with SessionDB(db_path=tmp_path / 'state.db') as db:
+        db.create_session('s', source='test')
+        db.append_messages_batch('s', [{'role': 'user', 'content': 'KEEP'}, {'role': 'assistant', 'content': 'OK'},
+                                       {'role': 'user', 'content': 'GONE'}, {'role': 'assistant', 'content': BIG}])
+        db.rewind_to_message('s', next(m['id'] for m in db.get_messages('s') if m['content'] == 'GONE'))
+        shown = db.get_messages_as_conversation('s', include_ancestors=True, include_compacted=True)
+        assert [m['content'] for m in shown] == ['KEEP', 'OK']
+
+
+def test_the_worker_history_read_agrees_with_the_owner_display_read(tmp_path):
+    """``worker_history`` (a worker's display read) applies the same removed-copy rule as the owner's."""
+    from hermes_state_worker_compression import worker_history
+    with SessionDB(db_path=tmp_path / 'state.db') as db:
+        db.create_session('s', source='test')
+        db.append_messages_batch('s', [{'role': 'user', 'content': 'KEEP'}, {'role': 'assistant', 'content': 'OK'},
+                                       {'role': 'user', 'content': 'GONE'}, {'role': 'assistant', 'content': 'NO'}])
+        # Rotated: the parent keeps its copies active, the child's copies of GONE/NO are rewound.
+        assert db.try_acquire_compression_lock('s', 'fixture')
+        try:
+            db.publish_compression_child(parent_session_id='s', child_session_id='c', source='test', messages=[
+                {'role': 'user', 'content': 'SUMMARY'}, {'role': 'user', 'content': 'GONE'},
+                {'role': 'assistant', 'content': 'NO'}], compression_lock_holder='fixture')
+        finally:
+            db.release_compression_lock('s', 'fixture')
+        db.rewind_to_message('c', next(m['id'] for m in db.get_messages('c') if m['content'] == 'GONE'))
+        owner = [m['content'] for m in db.get_messages_as_conversation('c', include_ancestors=True,
+                                                                       include_compacted=True)]
+        with db._read_ctx() as conn:
+            worker = worker_history(db, conn, 'c', {
+                'target': 'c', 'include_ancestors': True, 'include_inactive': False, 'repair_alternation': False,
+                'include_row_ids': False, 'include_compacted': True})['messages']
+        assert 'GONE' not in owner
+        assert [m['content'] for m in worker] == owner