fix(persistence): version only owned columns so metadata writes keep row ownership

The row digest hashed every repair column, so a same-process metadata write
(reaction, display-kind stamp, api_content / codex reasoning backfill,
platform message id) made our own row look like a foreign winner. The
re-flush then adopted the stale DB row: a later live edit (the non-ASCII
strip recovery) was reverted, and unsanitized tool_calls/reasoning were
copied back onto the live dict.

The digest now covers only the owned (non-metadata) columns: it means "the
row is still what we last committed". Match -> write the live owned values
and hand over only presentation metadata the live dict lacks; mismatch ->
genuine other writer, adopt as before. The r3 "stored content equals the
durable form of live" special case is subsumed and removed.

Also:
- _insert_message_rows drops a carried digest when it assigns a new row id
  (compaction/replace/import clones carried the parent's version).
- append_messages_batch restores each message's _row_id / digest /
  timestamp and pops the adopted row at the top of every _execute_write
  attempt, so a rolled-back attempt cannot resolve to a foreign row.
- message_id is no longer synced onto live (int -> str flip, spurious
  platform_message_id).
- The JSONL divert strips both bookkeeping keys via one frozenset.
This commit is contained in:
kshitijk4poor
2026-09-27 15:03:31 +05:30
committed by kshitij
parent 43b1dd7c5f
commit 6e8aa00626
4 changed files with 66 additions and 20 deletions

View File

@@ -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)

View File

@@ -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]

View File

@@ -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

View File

@@ -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"