diff --git a/tests/tui_gateway/test_delegation_session_lifecycle.py b/tests/tui_gateway/test_delegation_session_lifecycle.py new file mode 100644 index 00000000000..57a96469e23 --- /dev/null +++ b/tests/tui_gateway/test_delegation_session_lifecycle.py @@ -0,0 +1,158 @@ +"""Fail-closed ownership + session-scoped delegation lifecycle (#55578). + +Covers the two hardening rules layered on top of the origin-routing salvage: + +1. ``_session_owns_notification_event`` — positive-proof ownership. An + async-delegation completion may only be injected into a session that + PROVABLY commissioned it (origin UI id, or session-key/lineage match). + Orphans are never adopted by a foreign chat. + +2. ``interrupt_for_session`` — a session's in-flight async delegations end + with the session. ``_finalize_session`` interrupts delegations owned by + the closing session (by origin UI id always; by durable key only when the + TUI owns the lifecycle). +""" + +import threading +from unittest.mock import MagicMock, patch + +import pytest + +import tools.async_delegation as ad +from tui_gateway.server import ( + _finalize_session, + _session_owns_notification_event, +) + + +@pytest.fixture(autouse=True) +def _reset_async_delegation(): + ad._reset_for_tests() + yield + ad._reset_for_tests() + + +class TestSessionOwnsNotificationEvent: + def _session(self, key="sess_key_1"): + return {"session_key": key, "_finalized": False} + + def test_origin_ui_match_owns(self): + evt = {"type": "async_delegation", "origin_ui_session_id": "tab1", "session_key": "other"} + assert _session_owns_notification_event("tab1", self._session(), evt) is True + + def test_session_key_match_owns(self): + evt = {"type": "async_delegation", "origin_ui_session_id": "", "session_key": "sess_key_1"} + assert _session_owns_notification_event("tabX", self._session("sess_key_1"), evt) is True + + def test_orphan_is_not_owned(self): + """No origin match, no key match, owner gone → NOT ours (fail closed).""" + evt = {"type": "async_delegation", "origin_ui_session_id": "dead_tab", "session_key": "gone_key"} + assert _session_owns_notification_event("tab1", self._session(), evt) is False + + def test_empty_key_and_origin_not_owned(self): + """A delegation event with no return address at all is never adopted.""" + evt = {"type": "async_delegation", "origin_ui_session_id": "", "session_key": ""} + assert _session_owns_notification_event("tab1", self._session(), evt) is False + + def test_finalized_session_owns_nothing(self): + evt = {"type": "async_delegation", "origin_ui_session_id": "tab1", "session_key": "sess_key_1"} + sess = self._session() + sess["_finalized"] = True + assert _session_owns_notification_event("tab1", sess, evt) is False + + def test_compression_chain_resolution_owns(self): + evt = {"type": "async_delegation", "origin_ui_session_id": "", "session_key": "parent_key"} + db = MagicMock() + db.resolve_resume_session_id.return_value = "child_key" + with patch("tui_gateway.server._get_db", return_value=db): + assert _session_owns_notification_event("tabX", self._session("child_key"), evt) is True + + +class TestInterruptForSession: + def _seed_record(self, delegation_id, session_key="", origin_ui_session_id="", status="running"): + fn = MagicMock() + with ad._records_lock: + ad._records[delegation_id] = { + "delegation_id": delegation_id, + "status": status, + "session_key": session_key, + "origin_ui_session_id": origin_ui_session_id, + "interrupt_fn": fn, + } + return fn + + def test_interrupts_only_matching_session(self): + mine = self._seed_record("d1", session_key="sess_A") + other = self._seed_record("d2", session_key="sess_B") + n = ad.interrupt_for_session(session_key="sess_A") + assert n == 1 + mine.assert_called_once() + other.assert_not_called() + + def test_matches_by_origin_ui_session_id(self): + mine = self._seed_record("d1", origin_ui_session_id="tab1") + other = self._seed_record("d2", origin_ui_session_id="tab2") + n = ad.interrupt_for_session(origin_ui_session_id="tab1") + assert n == 1 + mine.assert_called_once() + other.assert_not_called() + + def test_no_selector_is_noop(self): + fn = self._seed_record("d1", session_key="sess_A") + assert ad.interrupt_for_session() == 0 + fn.assert_not_called() + + def test_completed_records_untouched(self): + fn = self._seed_record("d1", session_key="sess_A", status="completed") + assert ad.interrupt_for_session(session_key="sess_A") == 0 + fn.assert_not_called() + + +class TestFinalizeInterruptsOwnDelegations: + def _make_session(self, session_key="sess_A", sid="tab1"): + agent = MagicMock() + agent.session_id = session_key + agent._session_messages = None + agent.model = "m" + agent.platform = "tui" + return { + "agent": agent, + "history": [{"role": "user", "content": "x"}], + "history_lock": threading.Lock(), + "session_key": session_key, + "_finalized": False, + "_sid": sid, + } + + @patch("tui_gateway.server._get_db") + def test_finalize_interrupts_sessions_delegations(self, mock_get_db): + mock_db = MagicMock() + mock_db.get_session.return_value = {"source": "tui"} + mock_get_db.return_value = mock_db + + with patch("tools.async_delegation.interrupt_for_session") as mock_int: + _finalize_session(self._make_session(), end_reason="tui_close") + + mock_int.assert_called_once() + kwargs = mock_int.call_args.kwargs + assert kwargs["session_key"] == "sess_A" + assert kwargs["origin_ui_session_id"] == "tab1" + + @patch("tui_gateway.server._get_db") + def test_viewer_of_gateway_session_only_interrupts_by_origin(self, mock_get_db): + """Closing a TUI viewer tab on a live gateway session must not kill + the gateway's own background work — key-based interrupt is skipped, + origin-id interrupt (this tab's own dispatches) still applies.""" + mock_db = MagicMock() + mock_db.get_session.return_value = {"source": "telegram"} + mock_get_db.return_value = mock_db + + with patch("tools.async_delegation.interrupt_for_session") as mock_int: + _finalize_session( + self._make_session(session_key="agent:main:telegram:dm:123", sid="tab9"), + end_reason="ws_orphan_reap", + ) + + kwargs = mock_int.call_args.kwargs + assert kwargs["session_key"] == "" + assert kwargs["origin_ui_session_id"] == "tab9" diff --git a/tools/async_delegation.py b/tools/async_delegation.py index 710ab5c1d3b..cc22a7f6e9a 100644 --- a/tools/async_delegation.py +++ b/tools/async_delegation.py @@ -525,6 +525,54 @@ def interrupt_all(reason: str = "shutdown") -> int: return count +def interrupt_for_session( + session_key: str = "", + origin_ui_session_id: str = "", + reason: str = "session_end", +) -> int: + """Signal running async delegations owned by ONE session to stop. + + A delegation's lifecycle is bound to the session that spawned it: when + that session ends, its in-flight background subagents must end with it — + a completed orphan would otherwise sit on the shared completion queue + with no live owner, either leaking into another chat or burning tokens + with no one listening (#55578). + + Matches on ``origin_ui_session_id`` (the live UI session that + commissioned the work) and/or the durable ``session_key``; either + matching field claims the record. Returns how many were interrupted. + """ + if not session_key and not origin_ui_session_id: + return 0 + count = 0 + with _records_lock: + targets = [ + r for r in _records.values() + if r.get("status") == "running" + and ( + (origin_ui_session_id and str(r.get("origin_ui_session_id") or "") == origin_ui_session_id) + or (session_key and str(r.get("session_key") or "") == session_key) + ) + ] + for r in targets: + fn = r.get("interrupt_fn") + if callable(fn): + try: + fn() + count += 1 + except Exception as exc: + logger.debug( + "interrupt_for_session: %s interrupt failed: %s", + r.get("delegation_id"), exc, + ) + if count: + logger.info( + "Interrupted %d async delegation(s) for ending session (%s)", + count, reason, + ) + return count + + def _reset_for_tests() -> None: """Test-only: clear all state and tear down the executor.""" global _executor, _executor_max_workers diff --git a/tui_gateway/server.py b/tui_gateway/server.py index 29c7dd42d20..80fa9cea95e 100644 --- a/tui_gateway/server.py +++ b/tui_gateway/server.py @@ -625,6 +625,7 @@ def _finalize_session(session: dict | None, end_reason: str = "tui_close") -> No # Use session_id (from agent.session_id) not session_key — after compression, # session_key may be stale (the ended parent) while session_id is the live # continuation. Fix for #20001. + _tui_owns_lifecycle = True if session_id: try: db = _get_db() @@ -638,11 +639,40 @@ def _finalize_session(session: dict | None, end_reason: str = "tui_close") -> No # repeats on every inbound message. (#60609) row = db.get_session(session_id) source = (row or {}).get("source", "") - if not _is_gateway_owned_source(source): + _tui_owns_lifecycle = not _is_gateway_owned_source(source) + if _tui_owns_lifecycle: db.end_session(session_id, end_reason) except Exception: pass + # A session's in-flight async delegations end WITH the session (#55578): + # once nobody owns the return address, a still-running background subagent + # can only burn tokens and park an orphaned completion on the shared + # queue. Always interrupt delegations commissioned by THIS live UI session + # (its sid); additionally interrupt by durable session_key, but only when + # the TUI owns the lifecycle — closing a viewer tab on a live gateway + # session must not kill the gateway's own background work. + try: + from tools.async_delegation import interrupt_for_session + + _own_sid = str(session.get("_sid") or "") + if not _own_sid: + try: + with _sessions_lock: + for _cand_sid, _cand in _sessions.items(): + if _cand is session: + _own_sid = _cand_sid + break + except Exception: + _own_sid = "" + interrupt_for_session( + session_key=str(session_key or "") if _tui_owns_lifecycle else "", + origin_ui_session_id=_own_sid, + reason=end_reason, + ) + except Exception: + pass + # Close the slash-worker subprocess as part of finalize itself, not just # in the callers. Defense-in-depth: every session-end path goes through # _finalize_session (it's the single ``_finalized``-guarded chokepoint), so @@ -712,6 +742,10 @@ def _close_session_by_id(sid: str, *, end_reason: str = "tui_close") -> bool: session = _sessions.pop(sid, None) if session is None: return False + # The session is already out of _sessions here, so downstream teardown + # (e.g. _finalize_session's per-session async-delegation interrupt) can't + # recover its live id by scanning the dict — stamp it on the record. + session["_sid"] = sid _teardown_session(session, end_reason=end_reason) return True @@ -8550,6 +8584,40 @@ def _notification_event_belongs_elsewhere(sid: str, session: dict, evt: dict) -> ) +def _session_owns_notification_event(sid: str, session: dict, evt: dict) -> bool: + """True iff *this* session PROVABLY owns ``evt``. + + Positive ownership — the mirror of ``_notification_event_belongs_elsewhere`` + minus its orphan-adoption fallback. An event owns-matches when its + ``origin_ui_session_id`` is this live session, or its ``session_key`` + (raw or resolved through the compression chain) matches this session's + key/lineage. Used as a fail-closed gate for async-delegation payloads: + "not provably elsewhere" is NOT good enough to inject a conversation + payload into this chat (#55578). + """ + if session.get("_finalized"): + return False + if str(evt.get("origin_ui_session_id") or "") == str(sid or ""): + return True + evt_key = str(evt.get("session_key") or "") + if not evt_key: + return False + current_keys = { + str(session.get("session_key") or ""), + _session_lookup_key(session, fallback=sid), + } + if evt_key in current_keys: + return True + try: + db = _get_db() + resolved_key = ( + db.resolve_resume_session_id(evt_key) if db is not None else evt_key + ) or evt_key + except Exception: + resolved_key = evt_key + return resolved_key in current_keys + + def _notification_event_dedup_key(evt: dict) -> tuple: """Return the UI-emission identity for a process notification event. @@ -8620,6 +8688,32 @@ def _notification_poller_loop( time.sleep(0.1) continue + # Fail closed for async-delegation results (#55578): these carry a + # conversation payload, and injecting one into any chat other than the + # one that commissioned it is a hard cross-session leak. The + # belongs-elsewhere check above already re-queued events owned by + # another LIVE session; what reaches here is either ours or an + # orphan whose owner is gone. Orphaned delegation payloads are + # DROPPED, not adopted — the subagent's summary is already persisted + # in the delegation records/output store, so nothing is lost, whereas + # a wrong-chat injection is unrecoverable. Non-delegation events + # (background process completions etc.) keep the historical + # adopt-orphans behavior. + if evt.get("type") == "async_delegation" and not _session_owns_notification_event( + sid, session, evt + ): + logger.warning( + "async-delegation completion %s has no live owner " + "(origin=%r key=%r); dropping from injection instead of " + "delivering to session %s (#55578 fail-closed; result " + "remains in the delegation records)", + evt.get("delegation_id", "?"), + str(evt.get("origin_ui_session_id") or ""), + str(evt.get("session_key") or ""), + sid, + ) + continue + _evt_sid = evt.get("session_id", "") if evt.get("type") == "completion" and process_registry.is_completion_consumed(_evt_sid): continue @@ -8676,6 +8770,15 @@ def _notification_poller_loop( if _notification_event_belongs_elsewhere(sid, session, evt): deferred.append(evt) continue + # Same fail-closed rule as the live loop: an orphaned async-delegation + # payload is never adopted by a foreign session — defer it (a later + # resume of the owner's lineage can still claim it) rather than + # injecting another chat's conversation here (#55578). + if evt.get("type") == "async_delegation" and not _session_owns_notification_event( + sid, session, evt + ): + deferred.append(evt) + continue _evt_sid = evt.get("session_id", "") if evt.get("type") == "completion" and process_registry.is_completion_consumed(_evt_sid): continue