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