From b91b5bd967e59a3076f89ece16e68dd01281ea00 Mon Sep 17 00:00:00 2001 From: Jack Lau <72348727+jackulau@users.noreply.github.com> Date: Mon, 17 Aug 2026 02:21:09 -0500 Subject: [PATCH] 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 --- agent/conversation_compression.py | 46 +++++++++- .../agent/test_compression_rotation_state.py | 91 +++++++++++++++++++ 2 files changed, 136 insertions(+), 1 deletion(-) diff --git a/agent/conversation_compression.py b/agent/conversation_compression.py index a910a05c5f45c..97a9e3dbf2d3f 100644 --- a/agent/conversation_compression.py +++ b/agent/conversation_compression.py @@ -3476,6 +3476,13 @@ def _release_lock() -> None: # 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 _release_lock() -> None: 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 _release_lock() -> None: 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]}" diff --git a/tests/agent/test_compression_rotation_state.py b/tests/agent/test_compression_rotation_state.py index f87baa78acf01..490eee248309a 100644 --- a/tests/agent/test_compression_rotation_state.py +++ b/tests/agent/test_compression_rotation_state.py @@ -736,3 +736,94 @@ def test_parent_labels_cleared_after_rotation_child_lineage_intact( 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