From cd81242ee6b97c6ee425d3e019d9185eac63480d Mon Sep 17 00:00:00 2001 From: Dhruv Date: Thu, 16 Jul 2026 23:04:35 +0530 Subject: [PATCH] Async delegation completion is dropped after context compression and marked delivered --- gateway/run.py | 90 ++++++++---- .../test_async_delegation_session_binding.py | 132 ++++++++++++------ 2 files changed, 157 insertions(+), 65 deletions(-) diff --git a/gateway/run.py b/gateway/run.py index e5ca409706ad..536f58fc437c 100644 --- a/gateway/run.py +++ b/gateway/run.py @@ -11114,6 +11114,65 @@ def _get_cached_session_source(self, session_key: str): pass return source + async def _resolve_pinned_delegation_session(self, pinned_session_id: str) -> Optional[str]: + """Resolve an async-delegation completion's pinned session id (#55578/#65779). + + Returns the session id the completion should be injected into, or + ``None`` to fail closed (drop the injection; the subagent's output + remains in the delegation records): + + - pinned row live -> pin to it as-is; + - pinned row ended for ``compression`` -> that is a continuation + boundary, not a user-ended conversation: follow + ``get_compression_tip()`` and accept only a distinct, live tip; + - anything else (unknown row, /new reset, orphan reap, ended tip) + -> fail closed, as switch_session() re-opens ended sessions and + pinning blindly would RESURRECT a conversation the user explicitly + ended — the same illicit-revival class as the ws_orphan_reap loop + (#60609). + """ + pinned_row = None + try: + if self._session_db is not None: + # AsyncSessionDB already offloads to a thread. + pinned_row = await self._session_db.get_session(pinned_session_id) + except Exception: + pinned_row = None + if pinned_row is not None and not pinned_row.get("ended_at"): + return pinned_session_id + if pinned_row is not None and pinned_row.get("end_reason") == "compression": + try: + tip_id = await self._session_db.get_compression_tip(pinned_session_id) + except Exception: + logger.debug( + "compression-tip lookup failed for pinned session %s", + pinned_session_id, exc_info=True, + ) + tip_id = None + if tip_id and tip_id != pinned_session_id: + try: + tip_row = await self._session_db.get_session(tip_id) + except Exception: + tip_row = None + if tip_row is not None and not tip_row.get("ended_at"): + logger.info( + "Async-delegation completion pinned to compression-ended " + "session %s — rerouting to live continuation tip %s " + "(#65779)", + pinned_session_id, + tip_id, + ) + return tip_id + logger.warning( + "Async-delegation completion pinned to session %s, which is " + "%s — dropping injection instead of resurrecting it " + "(#55578 fail-closed; result remains in the delegation " + "records).", + pinned_session_id, + "unknown" if pinned_row is None else "ended", + ) + return None + async def _handle_message_with_agent(self, event, source, _quick_key: str, run_generation: int): """Inner handler that runs under the _running_agents sentinel guard.""" _msg_start_time = time.time() @@ -11149,31 +11208,14 @@ async def _handle_message_with_agent(self, event, source, _quick_key: str, run_g (getattr(event, "metadata", None) or {}).get("gateway_session_id") or "" ).strip() if pinned_session_id and pinned_session_id != session_entry.session_id: - # Fail closed (#55578): the spawning session may have ENDED since - # dispatch (user /new-reset, compression rotation whose parent was - # closed). switch_session() re-opens ended sessions, so pinning - # blindly would RESURRECT a conversation the user explicitly - # ended and inject into it — the same illicit-revival class as - # the ws_orphan_reap loop (#60609). A completion whose spawning - # session is dead is dropped from injection; the subagent's - # output remains in the delegation records. - pinned_row = None - try: - if self._session_db is not None: - # AsyncSessionDB already offloads to a thread. - pinned_row = await self._session_db.get_session(pinned_session_id) - except Exception: - pinned_row = None - if pinned_row is None or pinned_row.get("ended_at"): - logger.warning( - "Async-delegation completion pinned to session %s, which is " - "%s — dropping injection instead of resurrecting it " - "(#55578 fail-closed; result remains in the delegation " - "records).", - pinned_session_id, - "unknown" if pinned_row is None else "ended", - ) + # Fail closed (#55578) with one carve-out: a parent ended for + # compression is a continuation boundary, so the completion is + # rerouted to the live compression tip (#65779). All other ended/ + # unknown pinned sessions drop the injection. + resolved_session_id = await self._resolve_pinned_delegation_session(pinned_session_id) + if resolved_session_id is None: return + pinned_session_id = resolved_session_id prior_session_id = session_entry.session_id switched = await self.async_session_store.switch_session(session_key, pinned_session_id) if switched is not None: diff --git a/tests/gateway/test_async_delegation_session_binding.py b/tests/gateway/test_async_delegation_session_binding.py index c48b37e47dfa..7683acecffc2 100644 --- a/tests/gateway/test_async_delegation_session_binding.py +++ b/tests/gateway/test_async_delegation_session_binding.py @@ -64,59 +64,109 @@ def test_reset_interrupts_by_key_and_parent(self): class TestGatewayPinningFailsClosed: """The gateway injection path must never resurrect an ended session.""" - def _make_runner(self, pinned_row): + def _make_runner(self, rows, tip=None): from gateway.run import GatewayRunner runner = object.__new__(GatewayRunner) db = MagicMock() - db.get_session = AsyncMock(return_value=pinned_row) + db.get_session = AsyncMock(side_effect=lambda sid: rows.get(sid)) + db.get_compression_tip = AsyncMock(return_value=tip) runner._session_db = db + return runner - entry = MagicMock() - entry.session_key = "agent:main:telegram:dm:1" - entry.session_id = "sess_current" - runner.session_store = MagicMock() - runner.session_store.get_or_create_session.return_value = entry - runner.session_store.switch_session.return_value = entry - return runner, entry - - def _run_pinning_prefix(self, runner, pinned_session_id): - """Execute the pinning guard logic exactly as _handle_message does.""" - - async def _go(): - event = MagicMock() - event.metadata = {"gateway_session_id": pinned_session_id} - session_entry = runner.session_store.get_or_create_session(MagicMock()) - pinned = str((getattr(event, "metadata", None) or {}).get("gateway_session_id") or "").strip() - if pinned and pinned != session_entry.session_id: - pinned_row = None - try: - if runner._session_db is not None: - pinned_row = await runner._session_db.get_session(pinned) - except Exception: - pinned_row = None - if pinned_row is None or pinned_row.get("ended_at"): - return "dropped" - switched = runner.session_store.switch_session(session_entry.session_key, pinned) - if switched is not None: - return "pinned" - return "default" - - return asyncio.run(_go()) + def _resolve(self, runner, pinned_session_id): + return asyncio.run(runner._resolve_pinned_delegation_session(pinned_session_id)) def test_live_spawning_session_pins(self): - runner, _ = self._make_runner({"id": "sess_old", "ended_at": None}) - assert self._run_pinning_prefix(runner, "sess_old") == "pinned" + runner = self._make_runner({"sess_old": {"id": "sess_old", "ended_at": None}}) + assert self._resolve(runner, "sess_old") == "sess_old" def test_ended_spawning_session_drops(self): - runner, _ = self._make_runner({"id": "sess_old", "ended_at": "2026-07-08T00:00:00"}) - assert self._run_pinning_prefix(runner, "sess_old") == "dropped" - runner.session_store.switch_session.assert_not_called() + runner = self._make_runner( + {"sess_old": {"id": "sess_old", "ended_at": "2026-07-08T00:00:00"}} + ) + assert self._resolve(runner, "sess_old") is None def test_unknown_spawning_session_drops(self): - runner, _ = self._make_runner(None) - assert self._run_pinning_prefix(runner, "sess_gone") == "dropped" - runner.session_store.switch_session.assert_not_called() + runner = self._make_runner({}) + assert self._resolve(runner, "sess_gone") is None + + def test_handler_routes_through_resolver(self): + """The inner handler must gate pinning on the resolver, not raw ended_at.""" + import inspect + from gateway.run import GatewayRunner + + src = inspect.getsource(GatewayRunner._handle_message_with_agent) + assert "_resolve_pinned_delegation_session" in src + + +class TestCompressionContinuationRerouting: + """#65779: a compression-ended parent is a continuation boundary, not a + user-ended conversation — reroute the completion to the live tip.""" + + _make_runner = TestGatewayPinningFailsClosed._make_runner + _resolve = TestGatewayPinningFailsClosed._resolve + + def test_compression_parent_reroutes_to_live_tip(self): + runner = self._make_runner( + { + "sess_parent": { + "id": "sess_parent", + "ended_at": "2026-07-16T11:09:33", + "end_reason": "compression", + }, + "sess_tip": {"id": "sess_tip", "ended_at": None}, + }, + tip="sess_tip", + ) + assert self._resolve(runner, "sess_parent") == "sess_tip" + + def test_compression_parent_with_ended_tip_drops(self): + runner = self._make_runner( + { + "sess_parent": { + "id": "sess_parent", + "ended_at": "2026-07-16T11:09:33", + "end_reason": "compression", + }, + "sess_tip": { + "id": "sess_tip", + "ended_at": "2026-07-16T12:00:00", + "end_reason": "session_reset", + }, + }, + tip="sess_tip", + ) + assert self._resolve(runner, "sess_parent") is None + + def test_compression_parent_without_continuation_drops(self): + # get_compression_tip returns the input id when no continuation exists. + runner = self._make_runner( + { + "sess_parent": { + "id": "sess_parent", + "ended_at": "2026-07-16T11:09:33", + "end_reason": "compression", + }, + }, + tip="sess_parent", + ) + assert self._resolve(runner, "sess_parent") is None + + def test_explicit_reset_still_drops(self): + """Security assertion: /new stays fail-closed and never walks the chain.""" + runner = self._make_runner( + { + "sess_old": { + "id": "sess_old", + "ended_at": "2026-07-16T11:09:33", + "end_reason": "session_reset", + }, + }, + tip="sess_child", + ) + assert self._resolve(runner, "sess_old") is None + runner._session_db.get_compression_tip.assert_not_called() class TestResetHandlerInterruptsDelegations: