diff --git a/agent/message_metadata.py b/agent/message_metadata.py index f7369c6b75..17ab94a240 100644 --- a/agent/message_metadata.py +++ b/agent/message_metadata.py @@ -11,9 +11,14 @@ from typing import Any, MutableMapping, Optional, TypeVar # outgoing copy and the token estimator ignores them: one set, so an estimate # never prices bytes the provider never receives (an edit's inline_diff in # display_metadata is ~9KB and would trigger premature compaction). +# Transcript-repair bookkeeping riding on batch rows / live dicts (agent/transcript_repair.py): the +# stored-row CAS digest and the durable row adopted onto the live dict. Never transcript payload. +DB_ROW_SNAPSHOT = "_db_row_snapshot" +CANONICAL_ROW = "_canonical_row" +REPAIR_BOOKKEEPING_FIELDS = frozenset({DB_ROW_SNAPSHOT, CANONICAL_ROW}) PERSISTENCE_ONLY_MESSAGE_FIELDS = frozenset( - {"timestamp", "display_kind", "display_metadata", "_row_id", "_db_row_snapshot", "_canonical_row"} -) + {"timestamp", "display_kind", "display_metadata", "_row_id"} +) | REPAIR_BOOKKEEPING_FIELDS _Message = TypeVar("_Message", bound=MutableMapping[str, Any]) diff --git a/agent/message_sanitization.py b/agent/message_sanitization.py index 79e1ab370a..1377f5de8b 100644 --- a/agent/message_sanitization.py +++ b/agent/message_sanitization.py @@ -14,6 +14,7 @@ import re from functools import partial from typing import Any, Callable +from agent.message_metadata import DB_ROW_SNAPSHOT from agent.vision_message_prep import _provider_model_key logger = logging.getLogger(__name__) @@ -24,7 +25,7 @@ _SURROGATE_RE = re.compile(r'[\ud800-\udfff]') # Keys handled explicitly by _sanitize_messages; every OTHER key is swept generically. # The durable snapshot is an immutable compare-and-swap version, not message payload. -_MESSAGE_CORE_KEYS = frozenset({"content", "name", "tool_calls", "role", "_db_row_snapshot"}) +_MESSAGE_CORE_KEYS = frozenset({"content", "name", "tool_calls", "role", DB_ROW_SNAPSHOT}) def _sanitize_surrogates(text: str) -> str: diff --git a/agent/session_persistence.py b/agent/session_persistence.py index cab6a260ff..03d7f2d692 100644 --- a/agent/session_persistence.py +++ b/agent/session_persistence.py @@ -23,7 +23,9 @@ 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, REPAIR_BOOKKEEPING_FIELDS, sync_flushed_message_markers +from agent.message_metadata import REPAIR_BOOKKEEPING_FIELDS +from agent.message_metadata import DB_ROW_SNAPSHOT as _DB_ROW_SNAPSHOT +from agent.transcript_repair import sync_flushed_message_markers logger = logging.getLogger("run_agent") # origin module's name: log records / caplog filters unchanged diff --git a/agent/transcript_repair.py b/agent/transcript_repair.py index f1fa653360..bedc2e8ed5 100644 --- a/agent/transcript_repair.py +++ b/agent/transcript_repair.py @@ -9,13 +9,12 @@ 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 hermes_state_common import _id_chunks, _placeholders from hermes_state_messages import _MESSAGE_WRITE_COLUMNS -_DB_ROW_SNAPSHOT = "_db_row_snapshot" -_CANONICAL_ROW = "_canonical_row" - # Durable payload columns a row-addressed rewrite may change: every INSERT column except row identity, # 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"}) @@ -32,9 +31,10 @@ _OWNED_COLUMNS = tuple(c for c in _REPAIR_COLUMNS if c not in _METADATA_COLUMNS) _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}) +# Canonical-row markers: a matched row hands over only missing presentation metadata; a legacy (no-digest) +# dict over a filled assistant row adopts only its content. _METADATA_ONLY = "_metadata_only" +_CONTENT_ONLY = "_content_only" def transcript_row_snapshot(row: Mapping[str, Any]) -> str: @@ -138,8 +138,13 @@ def resolve_and_repair_transcript_batch( (encode_content_fn(msg.get("content")), target_id, session_id, target_row["content"]), ) else: - # Legacy dict over a non-blank assistant row: another writer already filled it; adopt that. - adopt = role == "assistant" + # 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) @@ -276,6 +281,8 @@ def sync_flushed_message_markers(batch_msgs: List[Dict[str, Any]], batch_rows: L 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) and canonical.get(_CONTENT_ONLY): + written["content"] = canonical["content"] elif isinstance(canonical, dict): for key in _SYNC_FIELDS: if key in canonical and canonical[key] is not None: diff --git a/hermes_state_messages.py b/hermes_state_messages.py index 2510698c50..f035b142b4 100644 --- a/hermes_state_messages.py +++ b/hermes_state_messages.py @@ -14,6 +14,7 @@ from agent.context_compressor import ( _DB_PERSISTED_MARKER as _DB_PERSISTED_MARKER_KEY, MODEL_ONLY_DISPLAY_METADATA_KEY, _is_checkpoint_item, _newest_checkpoint_carrier, split_user_originated_turn) from agent.memory_manager import sanitize_context +from agent.message_metadata import CANONICAL_ROW, DB_ROW_SNAPSHOT from agent.message_sanitization import _sanitize_surrogates from hermes_cli.timefmt import coerce_epoch from hermes_state_common import ( @@ -432,12 +433,12 @@ class SessionMessagesMixin: # _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")} + 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) + msg.pop(CANONICAL_ROW, None) for key, value in state.items(): if value is _absent: msg.pop(key, None) @@ -445,7 +446,7 @@ class SessionMessagesMixin: 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 + from agent.transcript_repair import resolve_and_repair_transcript_batch inserted_rows = resolve_and_repair_transcript_batch( conn, session_id, @@ -458,8 +459,6 @@ class SessionMessagesMixin: decode_row_fn=self._decoded_repair_row, ) inserted, tool_calls_total = self._insert_message_rows(conn, session_id, inserted_rows) - # Only this flush path re-reads the digest, so only it pays for one; hash the STORED rows. - stamp_inserted_row_snapshots(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) @@ -685,14 +684,19 @@ 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) + # A new row makes any carried CAS version (a clone's parent digest) meaningless; the + # stored digest of the new row is stamped below. + 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 if prune_checkpoints: self._prune_shadowed_checkpoints(conn, session_id, messages) + # Every inserter (flush, compaction clone, replace, rotation handoff, import) gets the new rows' own + # stored-row digest, so a later re-flush of these dicts takes the versioned rewrite path instead of the + # legacy one. One batched SELECT; hash the STORED rows (column affinity rewrites bind values). + from agent.transcript_repair import stamp_inserted_row_snapshots + stamp_inserted_row_snapshots(conn, session_id, messages) return inserted, tool_calls_total def _prune_shadowed_checkpoints(self, conn, session_id: str, live_messages: List[Dict[str, Any]]) -> None: diff --git a/tests/agent/test_tool_call_incremental_persistence.py b/tests/agent/test_tool_call_incremental_persistence.py index 17dfb94218..b7b68234cf 100644 --- a/tests/agent/test_tool_call_incremental_persistence.py +++ b/tests/agent/test_tool_call_incremental_persistence.py @@ -767,7 +767,9 @@ def test_flush_sanitized_active_user_and_tool_rows_do_not_append_duplicates(tmp_ {"type": "text", "text": "hi \ud800 there " + "x" * 4000}, {"type": "image_url", "image_url": {"url": "data:image/png;base64,AAAA"}}, ]}, - {"role": "assistant", "content": "ok"}, + {"role": "assistant", "content": "ok", "tool_calls": [ + {"id": "c1", "type": "function", "function": {"name": "terminal", "arguments": '{"q": "x \ud800"}'}}, + ]}, {"role": "tool", "tool_call_id": "c1", "name": "terminal", "content": "r \ud800"}, ] estimate_before = estimate_messages_tokens_rough(messages) @@ -784,6 +786,8 @@ def test_flush_sanitized_active_user_and_tool_rows_do_not_append_duplicates(tmp_ 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" + # A resumed / cloned dict carries ``_row_id`` without a digest (legacy path) over a filled assistant row. + messages[1].pop("_db_row_snapshot") assert _sanitize_messages_surrogates(messages) is True assert not any(message.get("_db_persisted") for message in (messages[0], messages[2])) @@ -804,6 +808,10 @@ def test_flush_sanitized_active_user_and_tool_rows_do_not_append_duplicates(tmp_ assert messages[0]["display_metadata"] == rows[0]["display_metadata"] assert rows[2]["content"] == "winner" assert messages[2]["content"] == "winner" + # The legacy path adopts the durable content only: the sanitized live tool_calls never revert to the + # stored lone surrogate. + assert messages[1]["content"] == "ok" + assert messages[1]["tool_calls"][0]["function"]["arguments"] == '{"q": "x \ufffd"}' def test_flush_sanitized_archived_user_and_tool_rows_do_not_append_duplicates(tmp_path):