diff --git a/agent/session_persistence.py b/agent/session_persistence.py index 738be15f1a..cab6a260ff 100644 --- a/agent/session_persistence.py +++ b/agent/session_persistence.py @@ -23,7 +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.transcript_repair import _DB_ROW_SNAPSHOT, sync_flushed_message_markers +from agent.transcript_repair import _DB_ROW_SNAPSHOT, REPAIR_BOOKKEEPING_FIELDS, sync_flushed_message_markers logger = logging.getLogger("run_agent") # origin module's name: log records / caplog filters unchanged @@ -370,9 +370,10 @@ def _db_flush_failed(agent, e: Exception, batch_rows: List[Dict[str, Any]], adop if isinstance(e, (StateDbReplacedError, StateDbCorruptError)): # A replaced/quarantined handle will not take this batch again — keep it on disk. try: - # The CAS digest is local repair bookkeeping, not transcript payload. + # The CAS digest / adopted row are local repair bookkeeping, not transcript payload. divert_session_transcript_jsonl(getattr(agent, "session_id", "") or "", - [{k: v for k, v in r.items() if k != _DB_ROW_SNAPSHOT} for r in batch_rows]) + [{k: v for k, v in r.items() if k not in REPAIR_BOOKKEEPING_FIELDS} + for r in batch_rows]) except Exception: logger.warning("JSONL divert failed after state.db %s for %s", agent._last_persistence_error_cause, getattr(agent, "session_id", None), exc_info=True) diff --git a/agent/transcript_repair.py b/agent/transcript_repair.py index f507551a9d..f1fa653360 100644 --- a/agent/transcript_repair.py +++ b/agent/transcript_repair.py @@ -20,11 +20,25 @@ _CANONICAL_ROW = "_canonical_row" # role, active flag and the ones owned by the display index / timestamp / session linkage. _NON_PAYLOAD_COLUMNS = frozenset({"session_id", "role", "timestamp", "active", "display_identity"}) _REPAIR_COLUMNS = tuple(c for c in _MESSAGE_WRITE_COLUMNS if c not in _NON_PAYLOAD_COLUMNS) -_SYNC_FIELDS = ("role", "message_id") + _REPAIR_COLUMNS +# Columns same-process writers update after our flush (reactions / display-kind stamps, api_content +# backfill, codex reasoning backfill + checkpoint pruning, platform message ids). They are not part of the +# ownership version: a metadata write must never make our own row look like a foreign winner. +_METADATA_COLUMNS = frozenset( + {"display_kind", "display_metadata", "api_content", "codex_reasoning_items", "platform_message_id"} +) +# The ownership version: what we last committed for the payload the sanitizer / live edits change. +_OWNED_COLUMNS = tuple(c for c in _REPAIR_COLUMNS if c not in _METADATA_COLUMNS) +# Presentation-only metadata a matched (ours) row may hand to a live dict that lacks it. +_LIVE_MISSING_METADATA = ("display_kind", "display_metadata") +# ``message_id`` is identity the flush derived from the live dict (int there, TEXT in SQLite): never synced. +_SYNC_FIELDS = ("role",) + _REPAIR_COLUMNS +# Local repair bookkeeping riding on batch rows / live dicts; never transcript payload. +REPAIR_BOOKKEEPING_FIELDS = frozenset({_DB_ROW_SNAPSHOT, _CANONICAL_ROW}) +_METADATA_ONLY = "_metadata_only" def transcript_row_snapshot(row: Mapping[str, Any]) -> str: - """Fixed-size digest of the durable repair columns (the CAS version) of a ``SELECT *`` messages row. + """Fixed-size digest of the owned (non-metadata) columns (the CAS version) of a ``SELECT *`` messages row. Callers pass rows read back from SQLite, never Python bind values: column affinity rewrites values on storage (int ``platform_message_id`` -> TEXT, float ``token_count`` -> INTEGER), so hashing bind values @@ -32,7 +46,7 @@ def transcript_row_snapshot(row: Mapping[str, Any]) -> str: dicts, so a full copy would double transcript memory and anything that prices dict bytes. """ digest = hashlib.blake2b(digest_size=16) - for column in _REPAIR_COLUMNS: + for column in _OWNED_COLUMNS: value = row[column] if value is None: digest.update(b"N") @@ -97,16 +111,24 @@ def resolve_and_repair_transcript_batch( target_id = int(target_row["id"]) msg["_row_id"] = target_id expected = msg.get(_DB_ROW_SNAPSHOT) + canonical = None if isinstance(expected, str): - # Digest match: the row is still the version we last committed, so our live dict is the source of - # truth (the DB holds its lossy durable projection, multimodal parts -> text, which must never flow - # back). The caller holds BEGIN IMMEDIATE, so the row cannot change between compare and UPDATE. - # Digest mismatch: someone changed the row after our flush; adopt the durable version. + # 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 + # projection, multimodal parts -> text, which must never flow back): write the live owned values + # and leave same-process metadata writes (reactions, backfills) alone. The caller holds BEGIN + # IMMEDIATE, so the row cannot change between compare and UPDATE. Mismatch: another writer changed + # the payload after our flush; adopt its durable version. adopt = transcript_row_snapshot(target_row) != expected if not adopt: serialized = serialize_message_fn(msg, float(target_row["timestamp"])) - if any(target_row[column] != serialized[column] for column in _REPAIR_COLUMNS): + if any(target_row[column] != serialized[column] for column in _OWNED_COLUMNS): _rewrite_row(conn, session_id, target_row, serialized) + 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) + canonical = {c: decoded[c] for c in missing if c in decoded} + canonical[_METADATA_ONLY] = True 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. @@ -126,11 +148,7 @@ def resolve_and_repair_transcript_batch( msg[_DB_ROW_SNAPSHOT] = transcript_row_snapshot(final_row) if adopt: canonical = decode_row_fn(final_row) - # Same-process metadata writers (reactions, display kind, api_content backfill, codex reasoning) - # change the digest without touching content. When the stored content is just the durable form - # of the live content, sync metadata only and keep the live (possibly multimodal) content. - if canonical.get("content") == decode_content_fn(encode_content_fn(msg.get("content"))): - canonical.pop("content", None) + if canonical: msg[_CANONICAL_ROW] = canonical else: msg.pop(_CANONICAL_ROW, None) @@ -154,10 +172,11 @@ def _rewrite_row( ).fetchall() ] if old_identity is not None else [] - assignments = ", ".join(f"{column} = ?" for column in _REPAIR_COLUMNS) + # Owned columns only: same-process metadata writes (a reaction, a backfill) stay as committed. + assignments = ", ".join(f"{column} = ?" for column in _OWNED_COLUMNS) conn.execute( f"UPDATE messages SET {assignments} WHERE id = ? AND session_id = ?", - [*(serialized[column] for column in _REPAIR_COLUMNS), int(target_row["id"]), session_id], + [*(serialized[column] for column in _OWNED_COLUMNS), int(target_row["id"]), session_id], ) _restore_display_index(conn, session_id, target_row, serialized, old_identity, old_peer_ids) @@ -252,7 +271,12 @@ def sync_flushed_message_markers(batch_msgs: List[Dict[str, Any]], batch_rows: L 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): + 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: + if canonical.get(key) is not None and written.get(key) is None: + written[key] = canonical[key] + elif isinstance(canonical, dict): for key in _SYNC_FIELDS: if key in canonical and canonical[key] is not None: written[key] = canonical[key] diff --git a/hermes_state_messages.py b/hermes_state_messages.py index bcff7de3a9..2510698c50 100644 --- a/hermes_state_messages.py +++ b/hermes_state_messages.py @@ -429,7 +429,20 @@ 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, stamp_inserted_row_snapshots @@ -672,6 +685,9 @@ class SessionMessagesMixin: msg["timestamp"] = message_timestamp if cur.lastrowid is not None: msg["_row_id"] = cur.lastrowid + # A new row makes any carried CAS version (a clone's parent digest) meaningless; the flush + # path restamps the stored digest right after this insert. + msg.pop("_db_row_snapshot", None) inserted += 1 tool_calls_total += _tool_calls_count(tool_calls) now_ts = max(now_ts, message_timestamp) + 1e-6 diff --git a/tests/agent/test_tool_call_incremental_persistence.py b/tests/agent/test_tool_call_incremental_persistence.py index 0e8f524409..17dfb94218 100644 --- a/tests/agent/test_tool_call_incremental_persistence.py +++ b/tests/agent/test_tool_call_incremental_persistence.py @@ -22,8 +22,8 @@ read snapshots captured at flush time, so removing any production flush call makes the corresponding assertion fail. """ -import sqlite3 import copy +import sqlite3 from types import SimpleNamespace from pathlib import Path import tempfile @@ -782,6 +782,8 @@ def test_flush_sanitized_active_user_and_tool_rows_do_not_append_duplicates(tmp_ other.close() # A same-process metadata write (reaction) changes the user row's digest but not its content. assert db.set_message_reaction(session_id, durable_ids[0], "\U0001F44D") + # ...followed by a real in-place live edit: the row is still ours, so the edit must win and be persisted. + messages[0]["content"][0]["text"] += " EDITED" assert _sanitize_messages_surrogates(messages) is True assert not any(message.get("_db_persisted") for message in (messages[0], messages[2])) @@ -793,6 +795,9 @@ def test_flush_sanitized_active_user_and_tool_rows_do_not_append_duplicates(tmp_ assert [row["id"] for row in rows] == durable_ids assert [message["_row_id"] for message in messages] == durable_ids assert rows[0]["content"].startswith("hi \ufffd there") + assert " EDITED" in rows[0]["content"] + assert messages[0]["content"][0]["text"].endswith(" EDITED") + assert messages[0]["message_id"] == 12345 and "platform_message_id" not in messages[0] # Neither our own rewrite nor a metadata-only change copies the lossy durable projection back: the live # image part survives while the reaction metadata is synced. assert messages[0]["content"][1]["type"] == "image_url"