fix(persistence): keep live multimodal content and hash stored rows
The row-addressed repair stamped the decoded durable row on every resolved message, so after our own rewrite the sync copied the lossy durable projection (image parts -> "text\n[screenshot]") back onto the live dict: multimodal user/tool messages lost their images and the prompt-cache prefix changed. Adopt the DB row only when another writer won (digest mismatch) or on the legacy assistant path, as BASE did. The insert-time digest hashed Python bind values, but SQLite affinity rewrites them on storage (int message_id -> TEXT, float token_count -> INTEGER), so live and DB digests never matched and in-place edits were silently dropped. Hash the stored rows instead, only on the append_messages_batch flush path that reads the digest (one SELECT per batch), incrementally (type tag + length prefix) instead of via JSON. Also skip the no-op UPDATE, fix the _write_columns comment/spacing and drop the duplicate top-level Optional import (F811).
This commit is contained in:
kshitijk4poor
@@ -5,9 +5,8 @@ clone lookup) and sync markers after commit."""
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
import sqlite3
|
||||
from typing import Any, Callable, Dict, List, Mapping, Optional
|
||||
from typing import Any, Callable, Dict, List, Mapping
|
||||
|
||||
from agent.context_compressor import _DB_PERSISTED_MARKER
|
||||
from hermes_state_common import _placeholders
|
||||
@@ -15,8 +14,11 @@ from hermes_state_common import _placeholders
|
||||
|
||||
_DB_ROW_SNAPSHOT = "_db_row_snapshot"
|
||||
_CANONICAL_ROW = "_canonical_row"
|
||||
|
||||
|
||||
def _write_columns() -> tuple:
|
||||
# Late import: hermes_state_messages imports this module lazily inside its methods.
|
||||
# Function-local import: hermes_state_messages imports this module lazily, so a top-level import
|
||||
# here would create an import cycle for callers that load agent.transcript_repair first.
|
||||
from hermes_state_messages import _MESSAGE_WRITE_COLUMNS
|
||||
|
||||
return _MESSAGE_WRITE_COLUMNS
|
||||
@@ -29,24 +31,43 @@ _REPAIR_COLUMNS = tuple(c for c in _write_columns() if c not in _NON_PAYLOAD_COL
|
||||
_SYNC_FIELDS = ("role", "message_id") + _REPAIR_COLUMNS
|
||||
|
||||
|
||||
def _canonical_value(value: Any) -> Any:
|
||||
return value.hex() if isinstance(value, (bytes, bytearray, memoryview)) else value
|
||||
|
||||
|
||||
def transcript_row_snapshot(row: Mapping[str, Any]) -> Optional[str]:
|
||||
def transcript_row_snapshot(row: Mapping[str, Any]) -> str | None:
|
||||
"""Fixed-size digest of the durable repair columns (the CAS version), or ``None`` for a partial SELECT.
|
||||
|
||||
A digest rather than a row copy: the value rides on live message dicts, so a full copy would double
|
||||
transcript memory and anything that prices dict bytes.
|
||||
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
|
||||
would never match the stored version. A digest rather than a row copy: the value rides on live message
|
||||
dicts, so a full copy would double transcript memory and anything that prices dict bytes.
|
||||
"""
|
||||
keys = set(row.keys())
|
||||
if not set(_REPAIR_COLUMNS) <= keys:
|
||||
return None
|
||||
canonical = json.dumps(
|
||||
[_canonical_value(row[column]) for column in _REPAIR_COLUMNS],
|
||||
ensure_ascii=True, separators=(",", ":"), default=str,
|
||||
)
|
||||
return hashlib.blake2b(canonical.encode("utf-8"), digest_size=16).hexdigest()
|
||||
digest = hashlib.blake2b(digest_size=16)
|
||||
for column in _REPAIR_COLUMNS:
|
||||
value = row[column]
|
||||
if value is None:
|
||||
digest.update(b"N")
|
||||
continue
|
||||
if isinstance(value, (bytes, bytearray, memoryview)):
|
||||
tag, data = b"B", bytes(value)
|
||||
elif isinstance(value, str):
|
||||
tag, data = b"S", value.encode("utf-8", "surrogatepass")
|
||||
else:
|
||||
tag, data = (b"F" if isinstance(value, float) else b"I"), repr(value).encode("ascii")
|
||||
digest.update(tag + len(data).to_bytes(8, "big") + data)
|
||||
return digest.hexdigest()
|
||||
|
||||
|
||||
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)."""
|
||||
by_id = {msg["_row_id"]: msg for msg in messages if isinstance(msg.get("_row_id"), int)}
|
||||
ids = list(by_id)
|
||||
for start in range(0, len(ids), 500):
|
||||
chunk = ids[start:start + 500]
|
||||
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)
|
||||
|
||||
|
||||
def is_content_blank(content: Any) -> bool:
|
||||
@@ -91,9 +112,16 @@ def resolve_and_repair_transcript_batch(
|
||||
serialized = serialize_message_fn(msg, float(target_row["timestamp"]))
|
||||
expected = msg.get(_DB_ROW_SNAPSHOT)
|
||||
has_snapshot = isinstance(expected, str)
|
||||
# Adopt the durable row onto the live dict only when another writer's version wins (digest mismatch)
|
||||
# or on the legacy assistant path. After our own rewrite the live dict is the source of truth: the DB
|
||||
# holds its lossy durable projection (multimodal parts -> text), which must never flow back.
|
||||
adopt = role == "assistant" and not has_snapshot
|
||||
if has_snapshot and transcript_row_snapshot(target_row) == expected:
|
||||
# The caller holds BEGIN IMMEDIATE, so the row cannot change between this compare and the UPDATE.
|
||||
_rewrite_row(conn, session_id, target_row, serialized)
|
||||
if any(target_row[column] != serialized[column] for column in _REPAIR_COLUMNS):
|
||||
_rewrite_row(conn, session_id, target_row, serialized)
|
||||
elif has_snapshot:
|
||||
adopt = True
|
||||
if not has_snapshot and role == "assistant" and is_content_blank(decode_content_fn(target_row["content"])):
|
||||
# Blank assistant rows are the pre-existing interrupted-stream repair path. Keep its narrow
|
||||
# content-only CAS for live dicts that predate durable row snapshots.
|
||||
@@ -107,7 +135,10 @@ def resolve_and_repair_transcript_batch(
|
||||
).fetchone()
|
||||
msg["timestamp"] = final_row["timestamp"]
|
||||
msg[_DB_ROW_SNAPSHOT] = transcript_row_snapshot(final_row)
|
||||
msg[_CANONICAL_ROW] = decode_row_fn(final_row)
|
||||
if adopt:
|
||||
msg[_CANONICAL_ROW] = decode_row_fn(final_row)
|
||||
else:
|
||||
msg.pop(_CANONICAL_ROW, None)
|
||||
return inserted_rows
|
||||
|
||||
|
||||
|
||||
@@ -432,7 +432,7 @@ class SessionMessagesMixin:
|
||||
def _do(conn):
|
||||
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
|
||||
from agent.transcript_repair import resolve_and_repair_transcript_batch, stamp_inserted_row_snapshots
|
||||
inserted_rows = resolve_and_repair_transcript_batch(
|
||||
conn,
|
||||
session_id,
|
||||
@@ -445,6 +445,8 @@ 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)
|
||||
@@ -655,8 +657,6 @@ class SessionMessagesMixin:
|
||||
Never touches sessions.* counters (callers reconcile differently); reasoning kept for assistant rows.
|
||||
A caller that re-archives some rows afterwards passes ``prune_checkpoints=False`` and prunes once they
|
||||
are archived again (:meth:`_prune_shadowed_checkpoints`)."""
|
||||
from agent.transcript_repair import _DB_ROW_SNAPSHOT, transcript_row_snapshot
|
||||
|
||||
now_ts = time.time()
|
||||
inserted = tool_calls_total = 0
|
||||
for msg in messages:
|
||||
@@ -672,7 +672,6 @@ class SessionMessagesMixin:
|
||||
msg["timestamp"] = message_timestamp
|
||||
if cur.lastrowid is not None:
|
||||
msg["_row_id"] = cur.lastrowid
|
||||
msg[_DB_ROW_SNAPSHOT] = transcript_row_snapshot(serialized)
|
||||
inserted += 1
|
||||
tool_calls_total += _tool_calls_count(tool_calls)
|
||||
now_ts = max(now_ts, message_timestamp) + 1e-6
|
||||
|
||||
@@ -762,7 +762,10 @@ def test_flush_sanitized_active_user_and_tool_rows_do_not_append_duplicates(tmp_
|
||||
session_id = "sess-sanitized-active-rows"
|
||||
db = _attach_real_session_db(agent, tmp_path / "state.db", session_id)
|
||||
messages = [
|
||||
{"role": "user", "content": "hi \ud800 there " + "x" * 4000},
|
||||
{"role": "user", "content": [
|
||||
{"type": "text", "text": "hi \ud800 there " + "x" * 4000},
|
||||
{"type": "image_url", "image_url": {"url": "data:image/png;base64,AAAA"}},
|
||||
]},
|
||||
{"role": "assistant", "content": "ok"},
|
||||
{"role": "tool", "tool_call_id": "c1", "name": "terminal", "content": "r \ud800"},
|
||||
]
|
||||
@@ -787,6 +790,8 @@ 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")
|
||||
# Our own rewrite never copies the lossy durable projection back: the live image part survives.
|
||||
assert messages[0]["content"][1]["type"] == "image_url"
|
||||
assert rows[2]["content"] == "winner"
|
||||
assert messages[2]["content"] == "winner"
|
||||
|
||||
|
||||
Reference in New Issue
Block a user