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:
teknium1
2026-09-19 01:26:11 -07:00
committed by Teknium
parent e1f46b8518
commit 1947f74cc4
3 changed files with 96 additions and 3 deletions

View File

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

View File

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

View File

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