fix(state): preserve summarized tool rows in micro-compaction

Micro-compaction carries a non-contiguous prefix and suffix around its summary marker. Using tail_count=len(result)-1 incorrectly marked summarized assistant/tool rows as rewind-only active=0, compacted=0.

Pass the exact unchanged carried messages instead and resolve their durable originals transactionally by row id or unique identity+timestamp, preserving summarized rows as compacted history.

Fixes #118481
This commit is contained in:
Brian Fernstrom
2026-09-22 13:19:52 -07:00
committed by Teknium
parent c07501ec41
commit 91df54184d
4 changed files with 180 additions and 8 deletions

View File

@@ -352,8 +352,17 @@ class MicroCompactionMixin:
if not session_db or not session_id:
return
try:
# Every row except the marker is a carried-forward original: archive rewind-style.
session_db.archive_and_compact(session_id, compacted_messages, tail_count=max(0, len(compacted_messages) - 1))
# Micro-compaction is prefix + marker + suffix, not a contiguous tail. Identify the exact
# byte-identical originals by their persistence marker; the state transaction resolves each
# one by row id or durable identity+timestamp. In-place mutations deliberately pop
# _DB_PERSISTED_MARKER, and the fresh summary marker never has one. A positional tail_count
# can otherwise classify the summarized assistant/tool rows as rewind-only (#118481).
carried_messages = [
message for message in compacted_messages
if isinstance(message, dict) and message.get(_cc()._DB_PERSISTED_MARKER)
]
session_db.archive_and_compact(
session_id, compacted_messages, carried_messages=carried_messages)
# Shared post-commit stamp site with batch commit and proactive prune.
# See #98450.
_cc().stamp_db_persisted_markers(compacted_messages)

View File

@@ -662,18 +662,72 @@ class SessionMessagesMixin:
f"WHERE id IN ({_placeholders(tail_ids)}) ORDER BY id",
[session_id, *tail_ids] if retarget else tail_ids)
def _resolve_carried_row_ids(
self, conn, session_id: str, carried_messages: List[Dict[str, Any]],
) -> List[int]:
"""Resolve byte-identical carried-forward live dicts to their ACTIVE durable originals.
_row_id is authoritative when the message carries it and the stored identity still matches.
Resume surfaces that intentionally omit row ids fall back to a UNIQUE
(role/content/tool identity, timestamp) match. Ambiguous or timestamp-less fallbacks are left
as compacted history rather than risking a false rewind classification.
"""
if not carried_messages:
return []
rows = conn.execute(
"SELECT id, role, content, tool_call_id, tool_calls, timestamp FROM messages "
"WHERE session_id = ? AND active = 1 ORDER BY id",
(session_id,),
).fetchall()
by_id: Dict[int, Tuple[Any, ...]] = {}
by_key: Dict[Tuple[Any, ...], List[int]] = {}
for row in rows:
row_id = int(row["id"])
identity = self._row_identity(
row["role"], self._decode_content(row["content"]), row["tool_call_id"],
_parse_tool_calls(row["tool_calls"]))
by_id[row_id] = identity
timestamp = coerce_epoch(row["timestamp"], field="message timestamp")
if timestamp is not None:
by_key.setdefault((*identity, timestamp), []).append(row_id)
resolved: List[int] = []
for message in carried_messages:
if not isinstance(message, dict):
continue
role = message.get("role", "unknown")
identity = self._row_identity(
role, message.get("content"), message.get("tool_call_id"),
_parse_tool_calls(message.get("tool_calls")))
row_id = message.get("_row_id")
if (isinstance(row_id, int) and not isinstance(row_id, bool)
and row_id > 0 and by_id.get(row_id) == identity):
resolved.append(row_id)
continue
timestamp = coerce_epoch(message.get("timestamp"), field="message timestamp")
if timestamp is None:
continue
matches = by_key.get((*identity, timestamp), [])
if len(matches) == 1:
resolved.append(matches[0])
return list(dict.fromkeys(resolved))
def archive_and_compact(self, session_id: str, compacted_messages: List[Dict[str, Any]],
model_config_patch: Optional[Dict[str, Any]] = None, watermark: Optional[int] = None,
lock_holder: Optional[str] = None, tail_count: int = 0) -> int:
lock_holder: Optional[str] = None, tail_count: int = 0,
carried_messages: Optional[List[Dict[str, Any]]] = None) -> int:
"""Non-destructive in-place compaction under ONE session id: soft-archive the active rows (``active=0,
compacted=1``: summarized away, still searchable) and insert *compacted_messages* as fresh active
rows, atomically; returns the new ACTIVE count (= ``message_count``). *watermark* (compression
START): rows ``id > watermark`` arrived during the slow summary and are re-sequenced after the
compacted set by a pure-SQL clone (fresh ids); ``None`` archives everything. *lock_holder*: verified
in-txn so a reclaimed lease fails instead of clobbering the winner. *tail_count*: the LAST N compacted
rows are the verbatim carried tail; their originals and the clones' originals are superseded
duplicates and get rewind flags (``active=0, compacted=0``) so search doesn't return each carried
message once per compaction. ``model_config_patch`` merges in the same txn (``None`` removes a key).
rows are the verbatim carried tail; *carried_messages* names exact durable originals carried forward
verbatim when they are not a contiguous suffix (micro-compaction's prefix + marker + suffix shape).
They are resolved inside this transaction by row id when present, else by unique durable identity +
timestamp. Those originals and the clones' originals are superseded duplicates and get rewind flags
(``active=0, compacted=0``) so search doesn't return each carried message once per compaction.
``model_config_patch`` merges in the same txn (``None`` removes a key).
Concurrent-append safety (#75316): when *watermark* is provided (the value of
:meth:`get_active_message_watermark` captured at compression START), rows that arrived during the
@@ -699,14 +753,16 @@ class SessionMessagesMixin:
(session_id, int(watermark)))
# Rewind targets sit AT/BELOW the watermark (all the compressor saw); unbounded, a
# concurrent append would steal a LIMIT slot.
rewind_ids: list[int] = []
rewind_ids: list[int] = self._resolve_carried_row_ids(
conn, session_id, carried_messages or [])
if tail_count > 0:
bound = watermark is not None
rewind_ids = [int(row["id"]) for row in conn.execute(
rewind_ids += [int(row["id"]) for row in conn.execute(
f"SELECT id FROM messages WHERE session_id = ? AND active = 1{' AND id <= ?' if bound else ''} "
"ORDER BY id DESC LIMIT ?",
(session_id, *((int(watermark),) if bound else ()), int(tail_count))).fetchall()]
rewind_ids += tail_ids
rewind_ids = list(dict.fromkeys(rewind_ids))
if rewind_ids:
placeholders = _placeholders(rewind_ids)
conn.execute("UPDATE messages SET active = 0, compacted = 0 "

View File

@@ -733,6 +733,48 @@ class TestMicroCompaction:
"micro-compaction gate must check agent._persist_disabled"
)
def test_db_sync_passes_exact_carried_messages(self):
"""Micro-compaction carries a prefix and suffix around its summary marker.
The persistence layer must receive exact durable ids, not len(result)-1:
a tail count reaches backward across the removed assistant/tool exchange and
turns summarized rows into rewind-only active=0, compacted=0 debris
(#118481).
"""
from agent.context_compressor import _DB_PERSISTED_MARKER
cc = _compressor()
captured = {}
class _DB:
def archive_and_compact(self, session_id, messages, **kwargs):
captured["session_id"] = session_id
captured["messages"] = messages
captured["kwargs"] = kwargs
return len(messages)
cc._session_db = _DB()
cc._session_id = "sess"
compacted = [
{"role": "user", "content": "prefix", "_row_id": 11, _DB_PERSISTED_MARKER: True},
{
"role": "assistant",
"content": "summary",
COMPRESSED_SUMMARY_METADATA_KEY: True,
},
{"role": "user", "content": "suffix", "_row_id": 15, _DB_PERSISTED_MARKER: True},
# Content was rewritten in-place: the mutation contract deliberately
# popped _DB_PERSISTED_MARKER, so its old row is NOT byte-identical.
{"role": "assistant", "content": "rewritten", "_row_id": 16},
]
cc._sync_micro_compact_to_db(compacted)
assert captured["session_id"] == "sess"
carried = captured["kwargs"].get("carried_messages")
assert [message.get("_row_id") for message in carried] == [11, 15]
assert "tail_count" not in captured["kwargs"]
def test_splice_preserves_db_persisted_stamps(self):
"""Surviving messages keep their _db_persisted stamps through a splice.

View File

@@ -138,6 +138,71 @@ class TestTailCountArchivesAsRewindSemantics:
f"{[h.get('snippet') for h in _recall(db, 'turn 1')]}"
)
def test_explicit_carried_messages_leave_summarized_tool_results_discoverable(self, db: SessionDB) -> None:
"""A non-contiguous carried set must not classify a summarized tool block as rewind-only.
Micro-compaction rewrites a prefix plus assistant/tool exchange plus suffix into
a prefix plus summary marker plus suffix. The carried rows therefore live on BOTH
sides of the removed exchange; using a positional tail count hid the removed
tool result as active=0, compacted=0 (#118481).
"""
db.append_message("sess1", role="user", content="question 0")
db.append_message(
"sess1",
role="assistant",
content="",
tool_calls=[{
"id": "call-1",
"type": "function",
"function": {"name": "read_file", "arguments": "{}"},
}],
)
db.append_message(
"sess1",
role="tool",
content="tool result that was summarized",
tool_call_id="call-1",
tool_name="read_file",
)
db.append_message("sess1", role="user", content="question 1")
db.append_message("sess1", role="assistant", content="answer 1")
carried = [
message for message in db.get_messages_as_conversation("sess1")
if message.get("content") in {"question 0", "question 1", "answer 1"}
]
assert carried and all("_row_id" not in message for message in carried), (
"this regression must exercise the row-id-less resume fallback"
)
db.archive_and_compact(
"sess1",
[
{"role": "user", "content": "question 0"},
{"role": "assistant", "content": "[CONTEXT COMPACTION] summarized tool exchange"},
{"role": "user", "content": "question 1"},
{"role": "assistant", "content": "answer 1"},
],
carried_messages=carried,
)
rows = _rows(db)
summarized_tool = [r for r in rows if r["content"] == "tool result that was summarized"]
assert len(summarized_tool) == 1
assert summarized_tool[0]["active"] == 0
assert summarized_tool[0]["compacted"] == 1, (
"summarized tool results must stay discoverable as compacted history"
)
carried_original = [
r for r in rows
if r["content"] == "question 1" and r["active"] == 0
]
assert len(carried_original) == 1
assert carried_original[0]["compacted"] == 0, (
"only exact carried-forward originals should be hidden as superseded duplicates"
)
def test_default_zero_keeps_archive_everything(self, db: SessionDB) -> None:
"""Without tail_count the historical behavior is untouched."""
_seed(db)