fix(compression): stop an aborted rotation from growing the parent it could not publish

The rotation path flushes its un-persisted transcript to the parent (#47202)
and only then calls publish_compression_child. The abort handler rolls back
the in-memory transcript and keeps agent.session_id on the parent - its own
comment says "keep the parent live and discard the stale compacted snapshot" -
but the rows the flush just wrote are not part of what it discards. Every
failed rotation therefore leaves the parent transcript longer than it found
it, whatever the failure was.

That is survivable for a one-off failure and pathological for a sticky one.
A parent row carrying ended_at fails the publish on every attempt and nothing
in this path clears it, so each auto-compaction appends another copy of the
current turn to the transcript it was supposed to shrink. Worse, the growth
then satisfies conversation_compression's own len(durable_parent) >
len(messages) check, so the next attempt adopts the inflated snapshot as if it
were genuine concurrent activity and the in-memory transcript doubles too.

Check that one precondition before writing. It is a plain read of the row the
publish is about to read anyway, and it raises the publish's own message, so
split_status=aborted, failure_class=session_split_failed and the rollback path
are all unchanged; a live parent reaches the flush exactly as before.
Deliberately not extended to the compression lease, which is re-acquirable - a
transient miss there would abort a rotation that would otherwise have
committed. old_session_id moves above the flush so a failure raised from here
takes the same in-memory rollback as any other pre-publish failure.

Scope: this fixes the amplification for every abort cause. It does not fix
what marks a live session as ended in the first place (#88197 Bug 1), which
needs a maintainer decision on end-reason taxonomy and is tracked on the
issue; an affected session still aborts every attempt, it just stops making
itself larger while it does.

Refs #88197
This commit is contained in:
Jack Lau
2026-08-17 02:21:09 -05:00
committed by kshitij
parent 979ca57a50
commit cf64ca20c5
2 changed files with 136 additions and 1 deletions
+45 -1
View File
@@ -3476,6 +3476,13 @@ def compress_context(
# the current-turn user message before preflight runs, so
# messages[:idx] is exactly the persisted prefix; only the
# current turn's new messages get written.
#
# Bound to old_session_id, hoisted above the flush: the
# ``except`` handler below keys its in-memory rollback off
# this name, so anything that fails from here on rolls the
# transcript back instead of leaving the failed attempt's
# compacted snapshot in place.
old_session_id = agent.session_id
current_idx = getattr(agent, "_persist_user_message_idx", None)
persisted_history = (
messages[:current_idx]
@@ -3483,6 +3490,44 @@ def compress_context(
and 0 <= current_idx <= len(messages)
else None
)
# The #47202 flush below is a DURABLE append to the parent
# and it is NOT undone when the rotation aborts: the
# ``except`` handler restores the in-memory transcript and
# keeps agent.session_id on the parent, but the rows it just
# wrote stay. Survivable for a one-off failure; pathological
# for a STICKY one. A parent row that already carries
# ``ended_at`` fails publish_compression_child on every
# attempt and nothing in this path clears it, so each
# auto-compaction appends another copy of the current turn to
# the transcript it was supposed to shrink — the session grows
# until the provider rejects the request outright (#88197:
# 303 unique messages stored as 2,611 rows after 7 aborted
# attempts, ~1.66M tokens, HTTP 400).
#
# So check that one precondition BEFORE writing. It is a plain
# read of the row the publish is about to read anyway, it
# raises the publish's own message so the log line, telemetry
# and rollback path are all unchanged, and it cannot mask a
# real rotation — a live parent reaches the flush exactly as
# before. Deliberately NOT extended to the compression lease:
# a lease is re-acquirable, so a transient miss here would
# abort a rotation that would otherwise have committed.
_parent_row_reader = getattr(agent._session_db, "get_session", None)
_parent_already_ended = False
if callable(_parent_row_reader):
try:
_parent_row = _parent_row_reader(old_session_id) or {}
_parent_already_ended = (
_parent_row.get("ended_at") is not None
)
except Exception:
# Fail OPEN: an unreadable row must not turn a cheap
# guard into a new way to lose compression.
_parent_already_ended = False
if _parent_already_ended:
raise RuntimeError(
f"Compression parent already ended: {old_session_id}"
)
# Foreign-tail ceiling (#75316): the flush below writes OUR
# OWN input transcript to the parent — those rows are
# already represented in the compacted handoff and must
@@ -3523,7 +3568,6 @@ def compress_context(
except Exception:
_profile_for_child = None
old_title = agent._session_db.get_session_title(agent.session_id)
old_session_id = agent.session_id
new_session_id = (
f"{datetime.now().strftime('%Y%m%d_%H%M%S')}_"
f"{uuid.uuid4().hex[:6]}"
@@ -736,3 +736,94 @@ class TestArchivedParentActivityLabelsCleared:
assert not prov or prov == ActivityProvenance.UNKNOWN.value, (
f"archived parent kept terminal provenance {prov!r}"
)
class TestAbortedRotationDoesNotGrowParent:
"""#88197 — a rotation that cannot publish must not have already written.
The rotation flushes its un-persisted current-turn transcript to the parent
(#47202) and only then calls ``publish_compression_child``. The abort
handler rolls back memory but not that flush, so every failed rotation
leaves the parent transcript longer than it found it. When the failure is
STICKY -- a parent row stamped ``ended_at`` by something that ended the
process rather than the conversation, e.g. the TUI gateway's
``_shutdown_sessions`` stamping ``end_reason='tui_shutdown'`` while the
agent keeps running -- every subsequent auto-compaction repeats it, and the
session grows instead of shrinking until the provider rejects the request.
"""
@staticmethod
def _durable_len(db: SessionDB, session_id: str) -> int:
return len(db.get_messages_as_conversation(session_id))
def test_ended_parent_aborts_before_the_prepublish_flush(self, tmp_path: Path):
db = SessionDB(db_path=tmp_path / "state.db")
parent = "PARENT_ENDED_NO_GROWTH"
db.create_session(parent, source="cli")
agent = _build_agent_with_db(db, parent)
# The lie at the heart of #88197: the row says ended, the agent is live.
# ``tui_shutdown`` is not a lineage boundary -- nothing forked off this
# session -- so durable writes to it are still permitted, which is
# exactly why the flush lands and the publish still refuses.
db.end_session(parent, "tui_shutdown")
assert db.get_session(parent)["ended_at"] is not None
before = self._durable_len(db, parent)
# Three consecutive auto-compactions, as the reported incident saw.
for attempt in range(1, 4):
original = _msgs()
returned, _sp = agent._compress_context(
original, "sys", approx_tokens=120_000
)
assert self._durable_len(db, parent) == before, (
f"attempt {attempt} appended to the parent it could not "
"publish; repeated attempts grow the transcript compression "
"exists to shrink"
)
# Rotation refused: the agent stays on the parent with its
# transcript intact, same contract as any other publish failure.
assert agent.session_id == parent
assert returned is original
assert [(m["role"], m["content"]) for m in returned] == [
(m["role"], m["content"]) for m in _msgs()
]
assert db.find_live_compression_child(parent) is None
def test_live_parent_still_gets_the_prepublish_flush(self, tmp_path: Path):
"""The guard must not cost a real rotation its #47202 tail."""
db = SessionDB(db_path=tmp_path / "state.db")
parent = "PARENT_LIVE_FLUSH"
db.create_session(parent, source="cli")
agent = _build_agent_with_db(db, parent)
agent._compress_context(_msgs(), "sys", approx_tokens=120_000)
assert agent.session_id != parent # rotation happened
# The current-turn messages survive in the preserved parent transcript.
assert self._durable_len(db, parent) >= len(_msgs())
def test_unreadable_parent_row_fails_open(self, tmp_path: Path):
"""A guard that cannot read the row must not become a way to lose
compression -- an unreadable parent rotates exactly as before."""
db = SessionDB(db_path=tmp_path / "state.db")
parent = "PARENT_UNREADABLE_ROW"
db.create_session(parent, source="cli")
agent = _build_agent_with_db(db, parent)
real_get_session = db.get_session
calls = {"n": 0}
def _flaky_get_session(session_id: str):
if session_id == parent and calls["n"] == 0:
calls["n"] += 1
raise RuntimeError("simulated read failure")
return real_get_session(session_id)
with patch.object(db, "get_session", side_effect=_flaky_get_session):
agent._compress_context(_msgs(), "sys", approx_tokens=120_000)
assert calls["n"] == 1, "the pre-flush guard never read the parent row"
assert agent.session_id != parent # rotation still happened