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