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:
@@ -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)
|
||||
|
||||
@@ -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]
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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"
|
||||
|
||||
Reference in New Issue
Block a user