f337631f43
_insert_message_rows stamps _row_id and the stored-row digest onto the caller's dicts inside the write transaction. Only append_messages_batch restored that state on rollback. archive_and_compact, replace_messages and the rotation handoff left the rolled-back id + digest on the dicts; SQLite reuses the id, so a later flush found a digest mismatch on the foreign row, adopted it and silently dropped the user's message. Move the capture/restore into _execute_transcript_write, used by every caller that inserts caller-owned dicts: each attempt starts from the caller's state and a final failure restores it before re-raising. (Rewind replacement and import insert dicts built inside the txn.) Also: bind _message_row_params directly on insert instead of the serialized-dict round-trip, import the public DB_ROW_SNAPSHOT / CANONICAL_ROW names, set adopt=False once, and reuse target_row instead of re-SELECTing when nothing was written.
300 lines
15 KiB
Python
300 lines
15 KiB
Python
"""Transcript repair for SessionDB batch appends: reconcile in-memory rows with committed SQLite rows
|
|
(in-place sanitizer rewrites, assistant blank-row repair, concurrent-winner adoption, watermark-compaction
|
|
clone lookup) and sync markers after commit."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import hashlib
|
|
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, DB_ROW_SNAPSHOT
|
|
from hermes_state_common import _id_chunks, _placeholders
|
|
from hermes_state_messages import _MESSAGE_WRITE_COLUMNS
|
|
|
|
|
|
# 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"})
|
|
_REPAIR_COLUMNS = tuple(c for c in _MESSAGE_WRITE_COLUMNS if c not in _NON_PAYLOAD_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
|
|
# 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:
|
|
"""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
|
|
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.
|
|
"""
|
|
digest = hashlib.blake2b(digest_size=16)
|
|
for column in _OWNED_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 dicts (every ``_insert_message_rows`` caller; one SELECT
|
|
per batch)."""
|
|
by_id = {msg["_row_id"]: msg for msg in messages if isinstance(msg.get("_row_id"), int)}
|
|
for chunk in _id_chunks(by_id):
|
|
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:
|
|
"""True when decoded message content is None, whitespace-only, or has no visible text parts."""
|
|
if content is None:
|
|
return True
|
|
if isinstance(content, str):
|
|
return not content.strip()
|
|
if isinstance(content, list):
|
|
return not "".join(p.get("text", "") for p in content if isinstance(p, dict) and p.get("type") == "text").strip()
|
|
return False
|
|
|
|
|
|
def resolve_and_repair_transcript_batch(
|
|
conn: sqlite3.Connection,
|
|
session_id: str,
|
|
messages: List[Dict[str, Any]],
|
|
encode_content_fn: Callable[[Any], Any],
|
|
decode_content_fn: Callable[[Any], Any],
|
|
serialize_message_fn: Callable[[Dict[str, Any], float], Mapping[str, Any]],
|
|
decode_row_fn: Callable[[Mapping[str, Any]], Dict[str, Any]],
|
|
) -> List[Dict[str, Any]]:
|
|
"""Resolve row-addressed rewrites without appending duplicates or replacing concurrent winners.
|
|
|
|
A durable row snapshot is a compare-and-swap version for sanitizer rewrites. Watermark-compaction clones
|
|
are matched by their copied payload identity, not timestamp alone. Legacy blank assistant rows retain the
|
|
narrow interrupted-stream content repair. Returns only rows that need fresh inserts.
|
|
"""
|
|
inserted_rows: List[Dict[str, Any]] = []
|
|
for msg in messages:
|
|
existing_row_id = msg.get("_row_id") if isinstance(msg, dict) else None
|
|
role = msg.get("role", "unknown") if isinstance(msg, dict) else "unknown"
|
|
target_row = None
|
|
if isinstance(existing_row_id, int):
|
|
target_row = _active_message_row(conn, session_id, existing_row_id, role)
|
|
if target_row is None:
|
|
inserted_rows.append(msg)
|
|
continue
|
|
|
|
target_id = int(target_row["id"])
|
|
msg["_row_id"] = target_id
|
|
expected = msg.get(DB_ROW_SNAPSHOT)
|
|
canonical = None
|
|
adopt = wrote = False
|
|
if isinstance(expected, str):
|
|
# 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 _OWNED_COLUMNS):
|
|
_rewrite_row(conn, session_id, target_row, serialized)
|
|
wrote = True
|
|
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.
|
|
wrote = conn.execute(
|
|
"UPDATE messages SET content = ? WHERE id = ? AND session_id = ? AND content IS ?",
|
|
(encode_content_fn(msg.get("content")), target_id, session_id, target_row["content"]),
|
|
).rowcount > 0
|
|
else:
|
|
# 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.
|
|
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)
|
|
).fetchone() if wrote else target_row
|
|
msg["timestamp"] = final_row["timestamp"]
|
|
msg[DB_ROW_SNAPSHOT] = transcript_row_snapshot(final_row)
|
|
if adopt:
|
|
canonical = decode_row_fn(final_row)
|
|
if canonical:
|
|
msg[CANONICAL_ROW] = canonical
|
|
else:
|
|
msg.pop(CANONICAL_ROW, None)
|
|
return inserted_rows
|
|
|
|
|
|
def _rewrite_row(
|
|
conn: sqlite3.Connection,
|
|
session_id: str,
|
|
target_row: Mapping[str, Any],
|
|
serialized: Mapping[str, Any],
|
|
) -> None:
|
|
"""Rewrite one durable payload whose digest matched the live dict's last committed version."""
|
|
old_identity = target_row["display_identity"]
|
|
old_peer_ids = [
|
|
int(row["id"])
|
|
for row in conn.execute(
|
|
"SELECT id FROM messages WHERE session_id = ? AND id != ? "
|
|
"AND (active = 1 OR compacted = 1) AND display_identity IS ?",
|
|
(session_id, target_row["id"], old_identity),
|
|
).fetchall()
|
|
] if old_identity is not None else []
|
|
|
|
# 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 _OWNED_COLUMNS), int(target_row["id"]), session_id],
|
|
)
|
|
_restore_display_index(conn, session_id, target_row, serialized, old_identity, old_peer_ids)
|
|
|
|
|
|
def _restore_display_index(
|
|
conn: sqlite3.Connection,
|
|
session_id: str,
|
|
target_row: Mapping[str, Any],
|
|
serialized: Mapping[str, Any],
|
|
old_identity: Any,
|
|
old_peer_ids: List[int],
|
|
) -> None:
|
|
"""Restore display identities/orders invalidated by the payload-update trigger."""
|
|
if old_peer_ids:
|
|
placeholders = _placeholders(old_peer_ids)
|
|
old_order = min(old_peer_ids)
|
|
conn.execute(
|
|
f"UPDATE messages SET display_identity = ?, display_order = ? "
|
|
f"WHERE session_id = ? AND id IN ({placeholders})",
|
|
(old_identity, old_order, session_id, *old_peer_ids),
|
|
)
|
|
|
|
target_id = int(target_row["id"])
|
|
new_identity = serialized["display_identity"]
|
|
visible = bool(target_row["active"] or target_row["compacted"])
|
|
if not visible:
|
|
conn.execute(
|
|
"UPDATE messages SET display_identity = ?, display_order = ? WHERE id = ? AND session_id = ?",
|
|
(new_identity, target_id, target_id, session_id),
|
|
)
|
|
return
|
|
|
|
peers = conn.execute(
|
|
"SELECT id, display_order FROM messages WHERE session_id = ? AND id != ? "
|
|
"AND (active = 1 OR compacted = 1) AND display_identity IS ?",
|
|
(session_id, target_id, new_identity),
|
|
).fetchall()
|
|
new_order = min(
|
|
[target_id]
|
|
+ [int(peer["display_order"] if peer["display_order"] is not None else peer["id"]) for peer in peers]
|
|
)
|
|
conn.execute(
|
|
"UPDATE messages SET display_identity = ?, display_order = ? WHERE id = ? AND session_id = ?",
|
|
(new_identity, new_order, target_id, session_id),
|
|
)
|
|
if peers:
|
|
conn.executemany(
|
|
"UPDATE messages SET display_order = ? WHERE id = ? AND session_id = ?",
|
|
[(new_order, int(peer["id"]), session_id) for peer in peers],
|
|
)
|
|
|
|
|
|
def _active_message_row(conn: sqlite3.Connection, session_id: str, row_id: int, role: str):
|
|
"""The same-role active clone for ``row_id``, or the addressed inactive row when no clone exists."""
|
|
row = conn.execute(
|
|
"SELECT * FROM messages WHERE id = ? AND session_id = ?", (row_id, session_id)
|
|
).fetchone()
|
|
if row is None or row["role"] != role:
|
|
return None
|
|
if int(row["active"] or 0) == 1:
|
|
return row
|
|
# Watermark compaction copies the complete durable row payload and display identity byte-for-byte.
|
|
# Timestamp alone is not an identity: externally supplied event timestamps may collide across unrelated
|
|
# messages. Legacy rows without an indexed identity cannot be resolved safely, so retain the addressed row.
|
|
if row["display_identity"] is None:
|
|
return row
|
|
payload_predicates = " AND ".join(f"{column} IS ?" for column in _REPAIR_COLUMNS)
|
|
clones = conn.execute(
|
|
"SELECT * FROM messages WHERE session_id = ? AND active = 1 AND role = ? "
|
|
"AND display_identity IS ? AND timestamp IS ? AND id != ? AND "
|
|
f"{payload_predicates} ORDER BY id DESC LIMIT 2",
|
|
(
|
|
session_id,
|
|
role,
|
|
row["display_identity"],
|
|
row["timestamp"],
|
|
row["id"],
|
|
*(row[column] for column in _REPAIR_COLUMNS),
|
|
),
|
|
).fetchall()
|
|
return clones[0] if len(clones) == 1 else row
|
|
|
|
|
|
def sync_flushed_message_markers(batch_msgs: List[Dict[str, Any]], batch_rows: List[Dict[str, Any]]) -> None:
|
|
"""Stamp persistence markers and sync canonical durable fields onto live dicts after commit."""
|
|
for written, row in zip(batch_msgs, batch_rows):
|
|
written[_DB_PERSISTED_MARKER] = True
|
|
if isinstance(row.get("_row_id"), int):
|
|
written["_row_id"] = row["_row_id"]
|
|
if isinstance(row.get("timestamp"), (int, float)):
|
|
written["timestamp"] = row["timestamp"]
|
|
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) 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) 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:
|
|
written[key] = canonical[key]
|
|
elif key not in ("role", "content"):
|
|
written.pop(key, None)
|
|
|
|
|
|
# ---- BEGIN PLUGIN-COMPAT (revert-scheduled; see COMPAT_MANIFEST.md) ----
|
|
# Names external plugins imported from this module before the Sep 2026 decomposition.
|
|
# Internal code MUST NOT use these (scripts/check_compat_pointers.py fails CI if it does).
|
|
# The whole block is removed by reverting the commit that added it.
|
|
from typing import Optional # noqa: F401,E402
|
|
# ---- END PLUGIN-COMPAT ----
|