From f337631f434b666f4ee390e3fac08b2609f50184 Mon Sep 17 00:00:00 2001 From: kshitijk4poor <82637225+kshitijk4poor@users.noreply.github.com> Date: Sun, 27 Sep 2026 16:00:33 +0530 Subject: [PATCH] fix(persistence): restore caller row state when any transcript insert rolls back _insert_message_rows stamps _row_id and the stored-row digest onto the caller's dicts inside the write transaction. Only append_messages_batch restored that state on rollback. archive_and_compact, replace_messages and the rotation handoff left the rolled-back id + digest on the dicts; SQLite reuses the id, so a later flush found a digest mismatch on the foreign row, adopted it and silently dropped the user's message. Move the capture/restore into _execute_transcript_write, used by every caller that inserts caller-owned dicts: each attempt starts from the caller's state and a final failure restores it before re-raising. (Rewind replacement and import insert dicts built inside the txn.) Also: bind _message_row_params directly on insert instead of the serialized-dict round-trip, import the public DB_ROW_SNAPSHOT / CANONICAL_ROW names, set adopt=False once, and reuse target_row instead of re-SELECTing when nothing was written. --- agent/session_persistence.py | 7 +++-- agent/transcript_repair.py | 32 +++++++++++----------- hermes_state_compression.py | 2 +- hermes_state_messages.py | 51 +++++++++++++++++++++++------------- 4 files changed, 53 insertions(+), 39 deletions(-) diff --git a/agent/session_persistence.py b/agent/session_persistence.py index 03d7f2d692..4025256c01 100644 --- a/agent/session_persistence.py +++ b/agent/session_persistence.py @@ -23,8 +23,7 @@ from agent.memory_manager import sanitize_context from agent.tool_dispatch_helpers import _is_multimodal_tool_result, _multimodal_text_summary from agent.trajectory import save_trajectory as _save_trajectory_to_file -from agent.message_metadata import REPAIR_BOOKKEEPING_FIELDS -from agent.message_metadata import DB_ROW_SNAPSHOT as _DB_ROW_SNAPSHOT +from agent.message_metadata import DB_ROW_SNAPSHOT, REPAIR_BOOKKEEPING_FIELDS from agent.transcript_repair import sync_flushed_message_markers @@ -227,8 +226,8 @@ def _db_flush_row(agent, msg: Dict, is_current_turn_user: bool) -> Dict[str, Any } if isinstance(msg.get("_row_id"), int): row["_row_id"] = msg["_row_id"] - if isinstance(msg.get(_DB_ROW_SNAPSHOT), str): - row[_DB_ROW_SNAPSHOT] = msg[_DB_ROW_SNAPSHOT] + if isinstance(msg.get(DB_ROW_SNAPSHOT), str): + row[DB_ROW_SNAPSHOT] = msg[DB_ROW_SNAPSHOT] return row diff --git a/agent/transcript_repair.py b/agent/transcript_repair.py index bedc2e8ed5..c1c9e92f5c 100644 --- a/agent/transcript_repair.py +++ b/agent/transcript_repair.py @@ -9,8 +9,7 @@ import sqlite3 from typing import Any, Callable, Dict, List, Mapping from agent.context_compressor import _DB_PERSISTED_MARKER -from agent.message_metadata import CANONICAL_ROW as _CANONICAL_ROW -from agent.message_metadata import DB_ROW_SNAPSHOT as _DB_ROW_SNAPSHOT +from agent.message_metadata import CANONICAL_ROW, DB_ROW_SNAPSHOT from hermes_state_common import _id_chunks, _placeholders from hermes_state_messages import _MESSAGE_WRITE_COLUMNS @@ -62,13 +61,14 @@ def transcript_row_snapshot(row: Mapping[str, Any]) -> str: def stamp_inserted_row_snapshots(conn: sqlite3.Connection, session_id: str, messages: List[Dict[str, Any]]) -> None: - """Stamp the stored-row digest on freshly inserted live dicts (flush path only; one SELECT per batch).""" + """Stamp the stored-row digest on freshly inserted dicts (every ``_insert_message_rows`` caller; one SELECT + per batch).""" by_id = {msg["_row_id"]: msg for msg in messages if isinstance(msg.get("_row_id"), int)} for chunk in _id_chunks(by_id): for row in conn.execute( f"SELECT * FROM messages WHERE session_id = ? AND id IN ({_placeholders(chunk)})", (session_id, *chunk) ).fetchall(): - by_id[int(row["id"])][_DB_ROW_SNAPSHOT] = transcript_row_snapshot(row) + by_id[int(row["id"])][DB_ROW_SNAPSHOT] = transcript_row_snapshot(row) def is_content_blank(content: Any) -> bool: @@ -110,8 +110,9 @@ def resolve_and_repair_transcript_batch( target_id = int(target_row["id"]) msg["_row_id"] = target_id - expected = msg.get(_DB_ROW_SNAPSHOT) + expected = msg.get(DB_ROW_SNAPSHOT) canonical = None + adopt = wrote = False if isinstance(expected, str): # The digest covers only the columns we own, so it answers "is the row still what we last # committed?". Match: the live dict is the source of truth (the DB holds its lossy durable @@ -124,6 +125,7 @@ def resolve_and_repair_transcript_batch( serialized = serialize_message_fn(msg, float(target_row["timestamp"])) if any(target_row[column] != serialized[column] for column in _OWNED_COLUMNS): _rewrite_row(conn, session_id, target_row, serialized) + wrote = True missing = {c: target_row[c] for c in _LIVE_MISSING_METADATA if msg.get(c) is None} if any(value is not None for value in missing.values()): decoded = decode_row_fn(target_row) @@ -132,31 +134,29 @@ def resolve_and_repair_transcript_batch( elif role == "assistant" and is_content_blank(decode_content_fn(target_row["content"])): # Legacy dict (no digest) over a blank assistant row: the interrupted-stream repair. Fill the row # from live content with a content-only CAS and never adopt the blank row onto the live dict. - adopt = False - conn.execute( + wrote = conn.execute( "UPDATE messages SET content = ? WHERE id = ? AND session_id = ? AND content IS ?", (encode_content_fn(msg.get("content")), target_id, session_id, target_row["content"]), - ) + ).rowcount > 0 else: # Legacy dict (no digest: a resumed or cloned dict) over a non-blank assistant row: another writer # already filled it. Adopt its content only, never the whole row: the live tool_calls / # reasoning* / codex_* fields may be sanitizer-fixed while the durable JSON still holds the raw # escaped surrogate, and live-only fields must survive. - adopt = False if role == "assistant": canonical = {"content": decode_content_fn(target_row["content"]), _CONTENT_ONLY: True} final_row = conn.execute( "SELECT * FROM messages WHERE id = ? AND session_id = ?", (target_id, session_id) - ).fetchone() + ).fetchone() if wrote else target_row msg["timestamp"] = final_row["timestamp"] - msg[_DB_ROW_SNAPSHOT] = transcript_row_snapshot(final_row) + msg[DB_ROW_SNAPSHOT] = transcript_row_snapshot(final_row) if adopt: canonical = decode_row_fn(final_row) if canonical: - msg[_CANONICAL_ROW] = canonical + msg[CANONICAL_ROW] = canonical else: - msg.pop(_CANONICAL_ROW, None) + msg.pop(CANONICAL_ROW, None) return inserted_rows @@ -273,9 +273,9 @@ def sync_flushed_message_markers(batch_msgs: List[Dict[str, Any]], batch_rows: L written["_row_id"] = row["_row_id"] if isinstance(row.get("timestamp"), (int, float)): written["timestamp"] = row["timestamp"] - if isinstance(row.get(_DB_ROW_SNAPSHOT), str): - written[_DB_ROW_SNAPSHOT] = row[_DB_ROW_SNAPSHOT] - canonical = row.get(_CANONICAL_ROW) + if isinstance(row.get(DB_ROW_SNAPSHOT), str): + written[DB_ROW_SNAPSHOT] = row[DB_ROW_SNAPSHOT] + canonical = row.get(CANONICAL_ROW) if isinstance(canonical, dict) and canonical.get(_METADATA_ONLY): # Our own row: only hand over presentation metadata the live dict lacks, never payload. for key in _LIVE_MISSING_METADATA: diff --git a/hermes_state_compression.py b/hermes_state_compression.py index 1228515c88..f0a82bc601 100644 --- a/hermes_state_compression.py +++ b/hermes_state_compression.py @@ -308,7 +308,7 @@ class SessionCompressionMixin: "WHERE id = ? AND ended_at IS NULL", (time.time(), parent_session_id)) if updated.rowcount != 1: raise RuntimeError(f"Compression parent changed during publication: {parent_session_id}") - self._execute_write(_do) + self._execute_transcript_write(_do, messages) def _write_sql_logged(self, op: str, session_id: str, sql: str, params) -> None: """``_write_sql`` that logs (never raises) on ``sqlite3.Error``.""" diff --git a/hermes_state_messages.py b/hermes_state_messages.py index f035b142b4..ab342aa366 100644 --- a/hermes_state_messages.py +++ b/hermes_state_messages.py @@ -430,20 +430,7 @@ class SessionMessagesMixin: compression_lock_holder=compression_lock_holder, turn_lease_holder=turn_lease_holder, turn_lease_ttl_seconds=turn_lease_ttl_seconds) for start in range(0, len(messages), chunk_rows)) - # _execute_write re-runs _do after a rollback: every attempt must start from the caller's state, or a - # rolled-back attempt's stamped _row_id could resolve to a row another writer took meanwhile. - _absent = object() - pre_state = [{k: m.get(k, _absent) for k in ("_row_id", DB_ROW_SNAPSHOT, "timestamp")} - for m in messages] - def _do(conn): - for msg, state in zip(messages, pre_state): - msg.pop(CANONICAL_ROW, None) - for key, value in state.items(): - if value is _absent: - msg.pop(key, None) - else: - msg[key] = value 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) from agent.transcript_repair import resolve_and_repair_transcript_batch @@ -461,7 +448,35 @@ class SessionMessagesMixin: inserted, tool_calls_total = self._insert_message_rows(conn, session_id, inserted_rows) self._bump_session_counters(conn, session_id, inserted, tool_calls_total, unit=False) return inserted - return self._execute_write(_do, patience_s=self._TRANSCRIPT_WRITE_PATIENCE_S) + return self._execute_transcript_write(_do, messages, patience_s=self._TRANSCRIPT_WRITE_PATIENCE_S) + + _ROW_STATE_KEYS = ("_row_id", DB_ROW_SNAPSHOT, "timestamp") + + def _execute_transcript_write(self, fn, messages: List[Dict[str, Any]], **kwargs): + """``_execute_write(fn)`` for callbacks that stamp row state onto the caller's *messages* (every + :meth:`_insert_message_rows` caller). Each attempt, and a final failure, restores the caller's + ``_row_id`` / digest / timestamp: a rolled-back insert's id is reused by SQLite, so a stale stamp + would make a later flush adopt another writer's row and drop this message.""" + _absent = object() + pre_state = [tuple(m.get(k, _absent) for k in self._ROW_STATE_KEYS) for m in messages] + + def _restore() -> None: + for msg, state in zip(messages, pre_state): + msg.pop(CANONICAL_ROW, None) + for key, value in zip(self._ROW_STATE_KEYS, state): + if value is _absent: + msg.pop(key, None) + else: + msg[key] = value + + def _attempt(conn): + _restore() + return fn(conn) + try: + return self._execute_write(_attempt, **kwargs) + except BaseException: + _restore() + raise def set_latest_matching_message_display_kind(self, session_id: str, *, role: str, content: str, display_kind: str, @@ -675,8 +690,8 @@ class SessionMessagesMixin: role = msg.get("role", "unknown") tool_calls = _parse_tool_calls(msg.get("tool_calls")) message_timestamp = _coerce_timestamp(msg.get("timestamp"), now_ts) - serialized = self._serialized_message_row(session_id, msg, message_timestamp) - cur = conn.execute(_INSERT_MESSAGE_SQL, tuple(serialized[column] for column in _MESSAGE_WRITE_COLUMNS)) + cur = conn.execute(_INSERT_MESSAGE_SQL, self._message_row_params( + session_id, role, msg, tool_calls, message_timestamp, keep_reasoning=role == "assistant")) # Keep the caller's live row aligned with the durable identity. Rows created without an explicit # timestamp (notably mid-turn steers) may be carried through several compaction generations; if # the generated timestamp exists only in SQLite, every copy receives a new identity and renders @@ -764,7 +779,7 @@ class SessionMessagesMixin: inserted, inserted_tool_calls = self._insert_message_rows(conn, session_id, messages[kept:]) conn.execute(f"{_SET_COUNTERS_SQL} WHERE id = ?", (kept + inserted, kept_tool_calls + inserted_tool_calls, session_id)) - self._execute_write(_do) + self._execute_transcript_write(_do, messages) @classmethod def _row_identity(cls, role: str, content: Any, tool_call_id: Any, tool_calls: Any) -> tuple: @@ -1062,7 +1077,7 @@ class SessionMessagesMixin: conn.execute(f"{_SET_COUNTERS_SQL}{', model_config = ?' if patch else ''} WHERE id = ?", (inserted, tool_calls_total, *((patched_model_config,) if patch else ()), session_id)) return inserted - return self._execute_write(_do) + return self._execute_transcript_write(_do, compacted_messages) def _message_column_names(self, conn) -> List[str]: """Column names of the messages table, cached per-connection era."""