fix(state): prune shadowed compaction checkpoints when a newer carrier row is persisted
`Fixes #102374` was only true for the compacted transcript: the prune ran inside local `compress()` / `salvage_grown_transcript`, and under native compaction the server compacts LOCAL_TRIGGER_SAFETY_MARGIN below the local trigger, so local compaction rarely fires. Every assistant response still persisted its ~120 KB checkpoint to `messages.codex_reasoning_items` and the older rows' shadowed copies were never rewritten — the field path in the issue (1,967 checkpoint rows, 12.9 GiB in one lineage). `SessionDB._insert_message_rows` (the one INSERT path: append batch, replace, compact, import) now rewrites the older ACTIVE assistant rows of the same session to drop their `type: "compaction"` items once a newer carrier row lands, keeping every non-checkpoint item; wire-neutral because `prune_pre_checkpoint_items` never replays a shadowed checkpoint. The flush mirrors it on the live transcript (`drop_shadowed_checkpoints`) so forks and compaction built from memory carry one checkpoint too, and the marker contract holds: the dicts read exactly as their rows.
This commit is contained in:
@@ -13,6 +13,8 @@ from agent.context_compressor import (
|
||||
COMPRESSED_SUMMARY_METADATA_KEY,
|
||||
_DB_PERSISTED_MARKER,
|
||||
ContextCompressor,
|
||||
_newest_checkpoint_carrier,
|
||||
drop_shadowed_checkpoints,
|
||||
user_originated_turn_view,
|
||||
)
|
||||
from agent.lazy_forward import forward as _forward, forward_static as _forward_static
|
||||
@@ -217,7 +219,7 @@ def _db_flush_collect(agent, messages: List[Dict], conversation_history: Optiona
|
||||
return batch_rows, batch_msgs
|
||||
|
||||
|
||||
def _db_flush_write(agent, batch_rows: List[Dict[str, Any]], batch_msgs: List[Dict]) -> None:
|
||||
def _db_flush_write(agent, batch_rows: List[Dict[str, Any]], batch_msgs: List[Dict], messages: List[Dict]) -> None:
|
||||
"""One transaction for the turn's new rows: on failure nothing lands and no markers are stamped."""
|
||||
if not batch_rows:
|
||||
return
|
||||
@@ -228,6 +230,11 @@ def _db_flush_write(agent, batch_rows: List[Dict[str, Any]], batch_msgs: List[Di
|
||||
turn_lease_ttl_seconds=getattr(agent, "_active_session_turn_lease_ttl_seconds", 300.0) or 300.0,
|
||||
)
|
||||
sync_flushed_message_markers(batch_msgs, batch_rows)
|
||||
if _newest_checkpoint_carrier(batch_msgs, "codex_reasoning_items") >= 0:
|
||||
# The insert already rewrote the older rows (SessionDB._drop_shadowed_checkpoint_rows); mirror it on
|
||||
# the live transcript so forks/compaction built from memory carry one checkpoint too. Markers stay:
|
||||
# the rows are durable exactly as the dicts now read.
|
||||
drop_shadowed_checkpoints(messages)
|
||||
|
||||
|
||||
def _db_flush_adopt_compression_tip(agent) -> bool:
|
||||
@@ -371,7 +378,7 @@ class SessionPersistenceMixin:
|
||||
if not self._session_db_created: # retry row creation if the earlier attempt failed transiently
|
||||
self._ensure_db_session()
|
||||
batch_rows, batch_msgs = _db_flush_collect(self, messages, conversation_history)
|
||||
_db_flush_write(self, batch_rows, batch_msgs)
|
||||
_db_flush_write(self, batch_rows, batch_msgs, messages)
|
||||
# Markers are now the sole truth; reset the one-shot seed so no id() outlives this flush.
|
||||
self._flushed_db_message_ids = set()
|
||||
self._last_flushed_db_idx = len(messages)
|
||||
|
||||
@@ -9,7 +9,9 @@ import logging
|
||||
import time
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
|
||||
from agent.context_compressor import _DB_PERSISTED_MARKER as _DB_PERSISTED_MARKER_KEY, split_user_originated_turn
|
||||
from agent.context_compressor import (
|
||||
_DB_PERSISTED_MARKER as _DB_PERSISTED_MARKER_KEY, _is_checkpoint_item, _newest_checkpoint_carrier,
|
||||
split_user_originated_turn)
|
||||
from agent.memory_manager import sanitize_context
|
||||
from agent.message_sanitization import _sanitize_surrogates
|
||||
from hermes_cli.timefmt import coerce_epoch
|
||||
@@ -42,6 +44,9 @@ _SET_COUNTERS_SQL = "UPDATE sessions SET message_count = ?, tool_call_count = ?"
|
||||
_RESET_COUNTERS_SQL = "UPDATE sessions SET message_count = 0, tool_call_count = 0 WHERE id = ?"
|
||||
_SET_DISPLAY_META_SQL = "UPDATE messages SET display_metadata = ? WHERE id = ?"
|
||||
_ARCHIVE_ACTIVE_SQL = "UPDATE messages SET active = 0, compacted = 1 WHERE session_id = ? AND active = 1"
|
||||
_SHADOWED_CHECKPOINT_ROWS_SQL = ("SELECT id, codex_reasoning_items FROM messages WHERE session_id = ? AND active = 1 "
|
||||
"AND role = 'assistant' AND id < ? AND codex_reasoning_items LIKE '%\"compaction\"%'")
|
||||
_SET_CODEX_REASONING_SQL = "UPDATE messages SET codex_reasoning_items = ? WHERE id = ?"
|
||||
_INVALID = object() # _json_or sentinel where the fallback must be distinguishable from JSON null
|
||||
|
||||
|
||||
@@ -506,8 +511,27 @@ class SessionMessagesMixin:
|
||||
inserted += 1
|
||||
tool_calls_total += _tool_calls_count(tool_calls)
|
||||
now_ts = max(now_ts, message_timestamp) + 1e-6
|
||||
carrier = _newest_checkpoint_carrier(messages, "codex_reasoning_items")
|
||||
if carrier >= 0 and isinstance(messages[carrier].get("_row_id"), int):
|
||||
self._drop_shadowed_checkpoint_rows(conn, session_id, messages[carrier]["_row_id"])
|
||||
return inserted, tool_calls_total
|
||||
|
||||
def _drop_shadowed_checkpoint_rows(self, conn, session_id: str, carrier_row_id: int) -> int:
|
||||
"""Rewrite older active assistant rows so only the row *carrier_row_id* keeps a ``type: "compaction"``
|
||||
checkpoint (durable twin of ``context_compressor.drop_shadowed_checkpoints``). Under native compaction
|
||||
every assistant response re-persists a ~120 KB checkpoint the wire builder will never replay once a
|
||||
newer one lands, and local compaction — the only other prune site — rarely fires (#102374).
|
||||
Non-checkpoint items stay; returns rows rewritten."""
|
||||
rewritten = 0
|
||||
for row_id, raw in conn.execute(_SHADOWED_CHECKPOINT_ROWS_SQL, (session_id, carrier_row_id)).fetchall():
|
||||
items = _json_or(raw, None, "Ignoring malformed codex_reasoning_items on message row")
|
||||
if not isinstance(items, list) or not any(_is_checkpoint_item(item) for item in items):
|
||||
continue
|
||||
kept = [item for item in items if not _is_checkpoint_item(item)]
|
||||
conn.execute(_SET_CODEX_REASONING_SQL, (self._reasoning_json_text(kept), row_id))
|
||||
rewritten += 1
|
||||
return rewritten
|
||||
|
||||
def replace_messages(self, session_id: str, messages: List[Dict[str, Any]], active_only: bool = False,
|
||||
archive_dropped: bool = False, reject_active_turn_lease: bool = False) -> None:
|
||||
"""Atomically replace a session's messages (/retry, /undo, /compress). DESTRUCTIVE by default (rows
|
||||
|
||||
@@ -190,3 +190,65 @@ class TestAppendMessagesBatch:
|
||||
db.append_messages_batch("sess-batch", msgs)
|
||||
raw = db._conn.execute("SELECT tool_calls FROM messages").fetchone()[0]
|
||||
assert json.loads(raw) == [{"name": "t", "arguments": "{}"}]
|
||||
|
||||
|
||||
class TestShadowedCheckpointRowsArePruned:
|
||||
"""Under native compaction every assistant response persists a fresh ``type: "compaction"`` checkpoint
|
||||
and local compaction (the only other prune site) rarely fires, so older rows kept ~120 KB of ciphertext
|
||||
the wire builder never replays (#102374). Landing a newer carrier row rewrites the older active rows."""
|
||||
|
||||
@staticmethod
|
||||
def _checkpoint(tag):
|
||||
return {"type": "compaction", "encrypted_content": f"ckpt-{tag}"}
|
||||
|
||||
@staticmethod
|
||||
def _reasoning(tag):
|
||||
return {"type": "reasoning", "encrypted_content": f"rs-{tag}", "id": f"rs_{tag}"}
|
||||
|
||||
@staticmethod
|
||||
def _agent(db):
|
||||
from agent.session_persistence import SessionPersistenceMixin
|
||||
|
||||
class _Agent(SessionPersistenceMixin):
|
||||
pass
|
||||
|
||||
agent = _Agent()
|
||||
agent._session_db, agent._session_db_created, agent.session_id = db, True, "sess-batch"
|
||||
agent._last_flushed_db_idx, agent._flushed_db_message_ids = 0, set()
|
||||
agent._flushed_db_message_session_id, agent._persist_disabled = None, False
|
||||
return agent
|
||||
|
||||
def _durable_items(self, db):
|
||||
return [
|
||||
(row["id"], json.loads(row["codex_reasoning_items"]) if row["codex_reasoning_items"] else None)
|
||||
for row in db._conn.execute(
|
||||
"SELECT id, codex_reasoning_items FROM messages WHERE session_id = ? AND role = 'assistant' "
|
||||
"AND active = 1 ORDER BY id", ("sess-batch",)).fetchall()
|
||||
]
|
||||
|
||||
def test_newer_carrier_row_prunes_the_older_rows_checkpoints(self, db):
|
||||
"""The production flush path: turn 1 lands a carrier, turn 2 lands a newer one -> only the newest
|
||||
row still holds a checkpoint, durably and in the live transcript; reasoning items are untouched."""
|
||||
agent = self._agent(db)
|
||||
messages = [
|
||||
{"role": "user", "content": "u0"},
|
||||
{"role": "assistant", "content": "a0", "codex_reasoning_items": [self._reasoning(0), self._checkpoint(0)]},
|
||||
]
|
||||
assert agent._flush_messages_to_session_db(messages) is True
|
||||
assert self._durable_items(db) == [(2, [self._reasoning(0), self._checkpoint(0)])]
|
||||
|
||||
messages += [
|
||||
{"role": "user", "content": "u1"},
|
||||
{"role": "assistant", "content": "a1", "codex_reasoning_items": [self._reasoning(1), self._checkpoint(1)]},
|
||||
]
|
||||
assert agent._flush_messages_to_session_db(messages) is True
|
||||
|
||||
assert self._durable_items(db) == [
|
||||
(2, [self._reasoning(0)]),
|
||||
(4, [self._reasoning(1), self._checkpoint(1)]),
|
||||
]
|
||||
# The live dicts match the rows they were persisted as (the marker contract), so no re-write is queued.
|
||||
assert messages[1]["codex_reasoning_items"] == [self._reasoning(0)]
|
||||
assert messages[3]["codex_reasoning_items"] == [self._reasoning(1), self._checkpoint(1)]
|
||||
assert agent._flush_messages_to_session_db(messages) is True
|
||||
assert db.message_count("sess-batch") == 4
|
||||
|
||||
Reference in New Issue
Block a user