From 1c6683e8e0bd27aa704ab264c656acafcac4f2e2 Mon Sep 17 00:00:00 2001 From: Xipong <217837358+Xipong@users.noreply.github.com> Date: Wed, 9 Sep 2026 02:34:53 +0300 Subject: [PATCH] fix: index compacted display identity writes --- hermes_state_common.py | 61 ++++++- hermes_state_messages.py | 76 +++++---- .../test_get_messages_include_compacted.py | 153 ++++++++++++++++++ 3 files changed, 252 insertions(+), 38 deletions(-) diff --git a/hermes_state_common.py b/hermes_state_common.py index 935640dad4..dcb2f77d49 100644 --- a/hermes_state_common.py +++ b/hermes_state_common.py @@ -364,6 +364,7 @@ CREATE TABLE IF NOT EXISTS messages ( api_content TEXT, display_kind TEXT, display_metadata TEXT, + display_identity BLOB, display_order INTEGER ); @@ -518,18 +519,70 @@ CREATE INDEX IF NOT EXISTS idx_messages_display_page ON messages(session_id, display_order, active DESC, id DESC) WHERE active = 1 OR compacted = 1; CREATE INDEX IF NOT EXISTS idx_messages_display_backfill - ON messages(session_id) WHERE display_order IS NULL AND (active = 1 OR compacted = 1); + ON messages(session_id) WHERE (display_order IS NULL OR display_identity IS NULL) + AND (active = 1 OR compacted = 1); +CREATE INDEX IF NOT EXISTS idx_messages_display_identity + ON messages(session_id, display_identity, display_order) + WHERE display_identity IS NOT NULL AND (active = 1 OR compacted = 1); +DROP TRIGGER IF EXISTS messages_display_order_insert; CREATE TRIGGER IF NOT EXISTS messages_display_order_insert AFTER INSERT ON messages WHEN new.display_order IS NULL BEGIN UPDATE messages SET display_order = COALESCE(( - SELECT MIN(COALESCE(display_order, id)) FROM messages + SELECT display_order FROM messages WHERE session_id = new.session_id AND id <> new.id AND (active = 1 OR compacted = 1) - AND role IS new.role AND content IS new.content AND timestamp IS new.timestamp - AND tool_call_id IS new.tool_call_id AND tool_calls IS new.tool_calls AND tool_name IS new.tool_name + AND display_identity = new.display_identity AND display_order IS NOT NULL + ORDER BY display_order LIMIT 1 ), new.id) WHERE id = new.id; END; +DROP TRIGGER IF EXISTS messages_display_visibility_update; +CREATE TRIGGER IF NOT EXISTS messages_display_visibility_update +AFTER UPDATE OF active, compacted ON messages +WHEN (new.active = 1 OR new.compacted = 1) <> (old.active = 1 OR old.compacted = 1) +BEGIN + UPDATE messages SET display_order = MIN(new.id, COALESCE(( + SELECT display_order FROM messages + WHERE session_id = new.session_id AND id <> new.id + AND (active = 1 OR compacted = 1) + AND display_identity = new.display_identity AND display_order IS NOT NULL + ORDER BY display_order LIMIT 1 + ), new.id)) WHERE id = new.id + AND (new.active = 1 OR new.compacted = 1); + UPDATE messages SET display_order = (SELECT display_order FROM messages WHERE id = new.id) + WHERE session_id = new.session_id AND id <> new.id AND (active = 1 OR compacted = 1) + AND display_identity = new.display_identity + AND (new.active = 1 OR new.compacted = 1); + UPDATE messages SET display_order = ( + SELECT MIN(peer.id) FROM messages AS peer + WHERE peer.session_id = old.session_id AND (peer.active = 1 OR peer.compacted = 1) + AND peer.display_identity = old.display_identity + ) WHERE session_id = old.session_id AND (active = 1 OR compacted = 1) + AND display_identity = old.display_identity + AND NOT (new.active = 1 OR new.compacted = 1); +END; +DROP TRIGGER IF EXISTS messages_display_identity_update; +CREATE TRIGGER IF NOT EXISTS messages_display_identity_update +AFTER UPDATE OF role, content, timestamp, tool_call_id, tool_calls, tool_name, + display_kind, display_metadata ON messages +BEGIN + UPDATE messages SET display_identity = NULL, display_order = NULL + WHERE id = new.id OR ( + session_id = old.session_id AND display_identity = old.display_identity + AND (active = 1 OR compacted = 1) + ); +END; +DROP TRIGGER IF EXISTS messages_display_identity_delete; +CREATE TRIGGER IF NOT EXISTS messages_display_identity_delete +AFTER DELETE ON messages WHEN old.active = 1 OR old.compacted = 1 +BEGIN + UPDATE messages SET display_order = ( + SELECT MIN(peer.id) FROM messages AS peer + WHERE peer.session_id = old.session_id AND (peer.active = 1 OR peer.compacted = 1) + AND peer.display_identity = old.display_identity + ) WHERE session_id = old.session_id AND (active = 1 OR compacted = 1) + AND display_identity = old.display_identity; +END; CREATE INDEX IF NOT EXISTS idx_messages_active_null ON messages(active) WHERE active IS NULL; CREATE INDEX IF NOT EXISTS idx_sessions_session_key diff --git a/hermes_state_messages.py b/hermes_state_messages.py index d3767d8ca5..dd1c84f9ed 100644 --- a/hermes_state_messages.py +++ b/hermes_state_messages.py @@ -3,6 +3,7 @@ replayed-user dedupe. Mixin bound via the MRO, built on SessionDB's _read_ctx/_e from __future__ import annotations +import hashlib import json import logging import time @@ -21,8 +22,9 @@ logger = logging.getLogger("hermes_state") # caplog tests pin the origin module _INSERT_MESSAGE_SQL = """INSERT INTO messages (session_id, role, content, tool_call_id, tool_calls, tool_name, effect_disposition, timestamp, token_count, finish_reason, reasoning, reasoning_content, reasoning_details, codex_reasoning_items, - codex_message_items, platform_message_id, observed, _compressed_summary, active, api_content, display_kind, display_metadata) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""" + codex_message_items, platform_message_id, observed, _compressed_summary, active, api_content, display_kind, + display_metadata, display_identity) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""" _BUMP_GENERATION_SQL = """ INSERT INTO conversation_generations (source, session_key, generation) VALUES (?, ?, 1) @@ -56,7 +58,11 @@ def _coerce_timestamp(value: Any, default: float) -> float: if value is None: return default try: - return float(value.timestamp()) if hasattr(value, "timestamp") else float(value) + result = float(value.timestamp()) if hasattr(value, "timestamp") else float(value) + # SQLite reads both signed zero spellings back as 0.0. Hash the value + # in that stored form so -0.0 and 0.0 retain the historical SQL/Python + # identity equality. + return 0.0 if result == 0.0 else result except (TypeError, ValueError): logger.debug("Ignoring invalid explicit message timestamp: %r", value) return default @@ -238,8 +244,18 @@ class SessionMessagesMixin: ``message_id`` (yuanbao's message-dict convention).""" _str_or_none = lambda v: _scrub_surrogates(v) if isinstance(v, str) else None # noqa: E731 _reasoning = lambda key: msg.get(key) if keep_reasoning else None # noqa: E731 - return (session_id, role, self._encode_content(msg.get("content")), msg.get("tool_call_id"), - json.dumps(tool_calls) if tool_calls else None, _scrub_surrogates(msg.get("tool_name")), + encoded_content = self._encode_content(msg.get("content")) + encoded_tool_calls = json.dumps(tool_calls) if tool_calls else None + encoded_tool_name = _scrub_surrogates(msg.get("tool_name")) + display_metadata = self._encode_display_metadata(msg.get("display_metadata")) + identity_row = { + "role": role, "content": encoded_content, "timestamp": message_timestamp, + "tool_call_id": msg.get("tool_call_id"), "tool_calls": encoded_tool_calls, + "tool_name": encoded_tool_name, "display_kind": msg.get("display_kind"), + "display_metadata": display_metadata, + } + return (session_id, role, encoded_content, msg.get("tool_call_id"), + encoded_tool_calls, encoded_tool_name, msg.get("effect_disposition"), message_timestamp, msg.get("token_count"), msg.get("finish_reason"), _scrub_surrogates(_reasoning("reasoning")), _scrub_surrogates(_reasoning("reasoning_content")), *(self._reasoning_json_text(_reasoning(k)) @@ -247,24 +263,7 @@ class SessionMessagesMixin: msg.get("platform_message_id") or msg.get("message_id"), 1 if msg.get("observed") else 0, 1 if msg.get("_compressed_summary") else 0, 1, _str_or_none(msg.get("api_content")), _str_or_none(msg.get("display_kind")), - self._encode_display_metadata(msg.get("display_metadata"))) - - def _inherit_composite_display_order(self, conn, row_id: int, session_id: str, msg: Dict[str, Any], - message_timestamp: float, tool_calls: Any) -> None: - handoff, live_view = split_user_originated_turn(msg) - if handoff is None or live_view is None: - return - encoded_content = self._encode_content(live_view.get("content")) - encoded_tool_calls = json.dumps(tool_calls) if tool_calls else None - row = conn.execute( - "SELECT MIN(COALESCE(display_order, id)) FROM messages " - "WHERE session_id = ? AND id <> ? AND (active = 1 OR compacted = 1) " - "AND role = 'user' AND content IS ? AND timestamp IS ? AND tool_call_id IS ? " - "AND tool_calls IS ? AND tool_name IS ?", - (session_id, row_id, encoded_content, message_timestamp, msg.get("tool_call_id"), - encoded_tool_calls, msg.get("tool_name"))).fetchone() - if row is not None and row[0] is not None: - conn.execute("UPDATE messages SET display_order = ? WHERE id = ?", (row[0], row_id)) + display_metadata, self._display_identity(self._display_dedupe_key(identity_row))) @staticmethod def _bump_session_counters(conn, session_id: str, inserted: int, tool_calls: int, *, unit: bool) -> None: @@ -302,9 +301,6 @@ class SessionMessagesMixin: self._check_transcript_write_guards(conn, session_id, compression_lock_holder, turn_lease_holder=turn_lease_holder, turn_lease_ttl_seconds=turn_lease_ttl_seconds) msg_id = conn.execute(_INSERT_MESSAGE_SQL, params).lastrowid - if msg_id is not None: - self._inherit_composite_display_order( - conn, int(msg_id), session_id, msg, message_timestamp, tool_calls) self._bump_session_counters(conn, session_id, 1, _tool_calls_count(tool_calls), unit=True) return msg_id # THE critical write (failure aborts the turn): long patience so a sibling legitimately @@ -452,8 +448,6 @@ class SessionMessagesMixin: session_id, role, msg, tool_calls, message_timestamp, keep_reasoning=role == "assistant")) if cur.lastrowid is not None: msg["_row_id"] = cur.lastrowid - self._inherit_composite_display_order( - conn, int(cur.lastrowid), session_id, msg, message_timestamp, tool_calls) inserted += 1 tool_calls_total += _tool_calls_count(tool_calls) now_ts = max(now_ts, message_timestamp) + 1e-6 @@ -517,10 +511,13 @@ class SessionMessagesMixin: return [int(r["id"]) for r in rows], sum(_tool_calls_len(r["tool_calls"]) for r in rows) def _clone_message_rows(self, conn, tail_ids: List[int], *, session_id: Optional[str] = None) -> None: - """Pure-SQL clone of *tail_ids* as fresh live rows (new id, active=1, compacted=0, all else byte-exact; - FTS triggers index the clones), into *session_id* when given.""" + """Pure-SQL clone of *tail_ids* as fresh live rows (new id/display order, active=1, compacted=0; + message payload columns stay byte-exact and FTS triggers index the clones), into *session_id* when given.""" retarget = session_id is not None - skip = ("id", "active", "compacted") + (("session_id",) if retarget else ()) + # A clone is a newly positioned display generation. Copy its indexed + # identity, but let the insert trigger assign order from rows that are + # still display-visible (the source may just have become rewind-only). + skip = ("id", "active", "compacted", "display_order") + (("session_id",) if retarget else ()) col_list = ", ".join(c for c in self._message_column_names(conn) if c not in skip) conn.execute( f"INSERT INTO messages ({col_list}, {'session_id, ' if retarget else ''}active, compacted) " @@ -659,6 +656,11 @@ class SessionMessagesMixin: return (row["role"], dedupe_content, row["timestamp"], row["tool_call_id"], row["tool_calls"], row["tool_name"]) + @staticmethod + def _display_identity(key: Tuple[Any, ...]) -> bytes: + """Fixed-width durable identity for indexed display-generation lookup.""" + return hashlib.sha256(repr(key).encode("utf-8", "surrogatepass")).digest() + def _dedupe_display_generations(self, rows): """Collapse compaction generations so each logical message appears once (the protected tail is copied into each generation: same role/content/timestamp, different ``active``/id); prefer the live row, then @@ -677,9 +679,13 @@ class SessionMessagesMixin: def _ensure_display_order(self, session_id: str) -> bool: """Backfill one legacy session once, preserving the pre-index display identity exactly.""" + with self._read_ctx() as conn: + columns = set(self._message_column_names(conn)) + if not {"display_order", "display_identity"} <= columns: + return False missing_sql = ( "SELECT 1 FROM messages WHERE session_id = ? AND (active = 1 OR compacted = 1) " - "AND display_order IS NULL LIMIT 1") + "AND (display_order IS NULL OR display_identity IS NULL) LIMIT 1") if self._read_one(missing_sql, (session_id,)) is None: return True if getattr(self, "read_only", False): @@ -699,8 +705,8 @@ class SessionMessagesMixin: first_id[key] = min(first_id.get(key, row["id"]), row["id"]) keyed_rows.append((row["id"], key)) conn.executemany( - "UPDATE messages SET display_order = ? WHERE id = ?", - [(first_id[key], row_id) for row_id, key in keyed_rows]) + "UPDATE messages SET display_order = ?, display_identity = ? WHERE id = ?", + [(first_id[key], self._display_identity(key), row_id) for row_id, key in keyed_rows]) return True return bool(self._execute_write(_do)) @@ -709,6 +715,8 @@ class SessionMessagesMixin: """``dict(row)`` with content/tool_calls/display_metadata decoded; *summary_flag* keeps ``_compressed_summary`` only as ``True``.""" msg = dict(row) + msg.pop("display_identity", None) + msg.pop("display_order", None) if summary_flag and msg.pop("_compressed_summary", 0): msg["_compressed_summary"] = True msg["content"] = self._decode_content(msg["content"]) diff --git a/tests/hermes_state/test_get_messages_include_compacted.py b/tests/hermes_state/test_get_messages_include_compacted.py index 5e41f4a7b6..59609398c0 100644 --- a/tests/hermes_state/test_get_messages_include_compacted.py +++ b/tests/hermes_state/test_get_messages_include_compacted.py @@ -12,6 +12,9 @@ soft-deleted Undo/Rewind rows (``active=0, compacted=0``) — that remains the job of ``include_inactive`` (audit / debug reads). """ +import json +import sqlite3 + import pytest from agent.context_compressor import ( @@ -179,6 +182,61 @@ class TestDisplayDedupe: large_steps = progress_steps("large") assert large_steps < small_steps * 3 + def test_append_display_identity_work_is_bounded_by_an_index(self, db): + def seed(sid, count): + db.create_session(sid, source="desktop") + db._execute_write(lambda conn: conn.executemany( + "INSERT INTO messages (session_id, role, content, timestamp, display_order) " + "VALUES (?, 'assistant', ?, ?, ?)", + [(sid, f"row-{index}", 100_000.0, index + 1) for index in range(count)], + )) + + seed("write-small", 2_000) + seed("write-large", 20_000) + db._wal_active = False + + def append_steps(sid): + callbacks = 0 + + def progress(): + nonlocal callbacks + callbacks += 1 + return 0 + + db._conn.set_progress_handler(progress, 10) + try: + db.append_message(sid, role="user", content="fresh", timestamp=100_000.0) + finally: + db._conn.set_progress_handler(None, 0) + return callbacks + + small_steps = append_steps("write-small") + large_steps = append_steps("write-large") + assert large_steps < small_steps * 3 + + def test_display_identity_normalizes_sql_values_and_stays_internal(self, db, tmp_path): + sid = "stored-values" + db.create_session(sid, source="desktop") + db.append_message(sid, role="assistant", content="same", timestamp=-0.0) + newest_id = db.append_message(sid, role="assistant", content="same", timestamp=0.0) + + messages = db.get_messages(sid, include_compacted=True) + assert [message["id"] for message in messages] == [newest_id] + assert "display_identity" not in messages[0] + assert "display_order" not in messages[0] + json.dumps(messages) + + path = tmp_path / "current-read-only.db" + current = SessionDB(path) + current.create_session("current", source="desktop") + current.append_message("current", role="user", content="serializable") + current.close() + reader = SessionDB(path, read_only=True) + try: + json.dumps(reader.get_messages("current", include_compacted=True)) + finally: + reader.close() + def test_composite_handoff_keeps_live_turn_identity_and_first_position(self, db): sid = "composite" db.create_session(sid, source="desktop") @@ -206,6 +264,101 @@ class TestDisplayDedupe: assert messages[0]["content"] == carrier assert original_id not in [message["id"] for message in messages] + @pytest.mark.parametrize("carrier_first", [False, True]) + def test_composite_identity_is_symmetric_and_tracks_identity_updates(self, db, carrier_first): + sid = f"composite-{'carrier' if carrier_first else 'raw'}-first" + db.create_session(sid, source="desktop") + carrier = ( + f"{_MERGED_PRIOR_CONTEXT_HEADER}\n" + "live ask\n\n" + f"{_MERGED_SUMMARY_DELIMITER}\n\n" + f"{SUMMARY_PREFIX}\n\n" + f"{HISTORICAL_TASK_HEADING}\nold work\n\n" + f"{_SUMMARY_END_MARKER}" + ) + contents = [carrier, "live ask"] if carrier_first else ["live ask", carrier] + first_id = db.append_message(sid, role="user", content=contents[0], timestamp=100.0) + middle_id = db.append_message(sid, role="assistant", content="middle", timestamp=200.0) + db._execute_write(lambda conn: conn.execute( + "UPDATE messages SET active = 0, compacted = 0 WHERE id = ?", (first_id,))) + second_id = db.append_message(sid, role="user", content=contents[1], timestamp=100.0) + assert [message["id"] for message in db.get_messages( + sid, include_compacted=True)] == [middle_id, second_id] + + # Making the first generation display-visible later must not split the + # identity that was assigned while it was hidden. + db._execute_write(lambda conn: conn.execute( + "UPDATE messages SET compacted = 1 WHERE id = ?", (first_id,))) + assert [message["id"] for message in db.get_messages( + sid, include_compacted=True)] == [second_id, middle_id] + + # Hiding the earliest generation again must advance the group's order + # to its first still-visible row. + db._execute_write(lambda conn: conn.execute( + "UPDATE messages SET compacted = 0 WHERE id = ?", (first_id,))) + assert [message["id"] for message in db.get_messages( + sid, include_compacted=True)] == [middle_id, second_id] + + # Identity-field updates invalidate the durable grouping; the next + # display read rebuilds it from the canonical key before paging. + db._execute_write(lambda conn: conn.execute( + "UPDATE messages SET role = 'assistant', content = 'middle', timestamp = 200.0 " + "WHERE id = ?", (second_id,))) + assert [message["id"] for message in db.get_messages( + sid, include_compacted=True)] == [second_id] + db._execute_write(lambda conn: conn.execute( + "DELETE FROM messages WHERE id = ?", (second_id,))) + assert [message["id"] for message in db.get_messages( + sid, include_compacted=True)] == [middle_id] + + def test_legacy_store_is_readable_then_lazily_migrated(self, tmp_path): + path = tmp_path / "legacy.db" + writer = SessionDB(path) + writer.create_session("legacy", source="desktop") + writer.append_message("legacy", role="assistant", content="same", timestamp=1.0) + writer.append_message("legacy", role="assistant", content="same", timestamp=1.0) + new_store_missing = writer._read_one( + "SELECT COUNT(*) FROM messages " + "WHERE display_order IS NULL OR display_identity IS NULL") + assert new_store_missing is not None and new_store_missing[0] == 0 + writer.close() + + conn = sqlite3.connect(path) + conn.execute("DROP TRIGGER IF EXISTS messages_display_order_insert") + conn.execute("DROP TRIGGER IF EXISTS messages_display_visibility_update") + conn.execute("DROP TRIGGER IF EXISTS messages_display_identity_update") + conn.execute("DROP TRIGGER IF EXISTS messages_display_identity_delete") + conn.execute("DROP INDEX IF EXISTS idx_messages_display_page") + conn.execute("DROP INDEX IF EXISTS idx_messages_display_backfill") + conn.execute("DROP INDEX IF EXISTS idx_messages_display_identity") + columns = {row[1] for row in conn.execute("PRAGMA table_info(messages)")} + for column in ("display_order", "display_identity"): + if column in columns: + conn.execute(f"ALTER TABLE messages DROP COLUMN {column}") + conn.commit() + conn.close() + + reader = SessionDB(path, read_only=True) + try: + legacy_messages = reader.get_messages("legacy", include_compacted=True) + assert len(legacy_messages) == 1 + json.dumps(legacy_messages) + assert "display_order" not in { + row[1] for row in reader._conn.execute("PRAGMA table_info(messages)")} + finally: + reader.close() + + migrated = SessionDB(path) + try: + assert len(migrated.get_messages("legacy", include_compacted=True)) == 1 + columns = {row[1] for row in migrated._conn.execute("PRAGMA table_info(messages)")} + assert {"display_order", "display_identity"} <= columns + assert migrated._conn.execute( + "SELECT COUNT(*) FROM messages " + "WHERE display_order IS NULL OR display_identity IS NULL").fetchone()[0] == 0 + finally: + migrated.close() + def test_copied_protected_tail_is_surfaced_once(self, db): """A message copied across compaction epochs appears exactly once.""" sid = "s1"