fix(delegation): route async results to their origin session across compression (#55578)

gateway/session_context.py

@@ -79,6 +79,13 @@ def session_context_engaged() -> bool:
 _SESSION_USER_NAME: ContextVar = ContextVar("HERMES_SESSION_USER_NAME", default=_UNSET)
 _SESSION_KEY: ContextVar = ContextVar("HERMES_SESSION_KEY", default=_UNSET)
 _SESSION_ID: ContextVar = ContextVar("HERMES_SESSION_ID", default=_UNSET)
+# In-process UI session/window id for multi-session desktop/TUI hosts. This is
+# intentionally separate from HERMES_SESSION_ID: the latter is the durable
+# conversation/session-db id, while the UI id is the live frontend tab/window
+# that commissioned a detached completion. Background completions use it as a
+# precise return address so a stale/rotated durable session key cannot be
+# consumed by whichever desktop poller wakes first.
+_SESSION_UI_SESSION_ID: ContextVar = ContextVar("HERMES_UI_SESSION_ID", default=_UNSET)
 # ID of the message that triggered the current turn. Used as a reply anchor
 # so background-process notifications stay inside the originating Telegram
 # private-chat topic (those lanes route only with thread id + reply anchor).
@@ -123,6 +130,7 @@ def session_context_engaged() -> bool:
     "HERMES_SESSION_USER_NAME": _SESSION_USER_NAME,
     "HERMES_SESSION_KEY": _SESSION_KEY,
     "HERMES_SESSION_ID": _SESSION_ID,
+    "HERMES_UI_SESSION_ID": _SESSION_UI_SESSION_ID,
     "HERMES_SESSION_MESSAGE_ID": _SESSION_MESSAGE_ID,
     "HERMES_SESSION_PROFILE": _SESSION_PROFILE,
     "HERMES_CRON_AUTO_DELIVER_PLATFORM": _CRON_AUTO_DELIVER_PLATFORM,
@@ -160,6 +168,7 @@ def set_session_vars(
     profile: str = "",
     cwd: str = "",
     async_delivery: bool = True,
+    ui_session_id: str = "",
 ) -> list:
     """Set all session context variables and return reset tokens.
 
@@ -191,6 +200,7 @@ def set_session_vars(
         _SESSION_USER_NAME.set(user_name),
         _SESSION_KEY.set(session_key),
         _SESSION_ID.set(session_id),
+        _SESSION_UI_SESSION_ID.set(ui_session_id),
         _SESSION_MESSAGE_ID.set(message_id),
         _SESSION_PROFILE.set(profile),
         _SESSION_ASYNC_DELIVERY.set(bool(async_delivery)),
@@ -225,6 +235,7 @@ def clear_session_vars(tokens: list) -> None:
         _SESSION_USER_NAME,
         _SESSION_KEY,
         _SESSION_ID,
+        _SESSION_UI_SESSION_ID,
         _SESSION_MESSAGE_ID,
         _SESSION_PROFILE,
     ):

tests/test_tui_gateway_server.py

@@ -2112,14 +2112,85 @@ def test_notification_event_routing_by_session_key(monkeypatch):
     monkeypatch.setattr(server, "_sessions", {"a": mine, "b": other})
 
     # My own event → handle it.
-    assert server._notification_event_belongs_elsewhere(mine, {"session_key": "mine"}) is False
+    assert server._notification_event_belongs_elsewhere("a", mine, {"session_key": "mine"}) is False
     # Global/system event with no owner → handle it.
-    assert server._notification_event_belongs_elsewhere(mine, {"session_key": ""}) is False
-    assert server._notification_event_belongs_elsewhere(mine, {}) is False
+    assert server._notification_event_belongs_elsewhere("a", mine, {"session_key": ""}) is False
+    assert server._notification_event_belongs_elsewhere("a", mine, {}) is False
     # Owned by another *live* session → defer to that session's poller.
-    assert server._notification_event_belongs_elsewhere(mine, {"session_key": "other"}) is True
+    assert server._notification_event_belongs_elsewhere("a", mine, {"session_key": "other"}) is True
     # Owner is gone (not in _sessions) → handle as fallback so it isn't lost.
-    assert server._notification_event_belongs_elsewhere(mine, {"session_key": "ghost"}) is False
+    assert server._notification_event_belongs_elsewhere("a", mine, {"session_key": "ghost"}) is False
+
+
+def test_async_delegation_event_prefers_origin_ui_session(monkeypatch):
+    """Detached subagent completions return to the commissioning TUI tab.
+
+    Regression: when the durable session key was stale/orphaned, whichever
+    desktop poller woke first could consume the async result and inject it into
+    an unrelated session.
+    """
+    mine = _session(session_key="current-key")
+    other = _session(session_key="unrelated-key")
+    monkeypatch.setattr(server, "_sessions", {"origin-sid": mine, "other-sid": other})
+    monkeypatch.setattr(server, "_get_db", lambda: None)
+    evt = {
+        "type": "async_delegation",
+        "session_key": "stale-or-rotated-key",
+        "origin_ui_session_id": "origin-sid",
+    }
+
+    assert server._notification_event_belongs_elsewhere("other-sid", other, evt) is True
+    assert server._notification_event_belongs_elsewhere("origin-sid", mine, evt) is False
+
+
+def test_notification_event_follows_compression_continuation(monkeypatch):
+    """Events keyed to a compressed parent route to the live continuation."""
+    old_parent = _session(session_key="old-parent")
+    live_tip = _session(session_key="new-tip")
+    monkeypatch.setattr(server, "_sessions", {"old-sid": old_parent, "tip-sid": live_tip})
+
+    class _DB:
+        def resolve_resume_session_id(self, session_id):
+            return "new-tip" if session_id == "old-parent" else session_id
+
+    monkeypatch.setattr(server, "_get_db", lambda: _DB())
+    evt = {"type": "async_delegation", "session_key": "old-parent"}
+
+    assert server._notification_event_belongs_elsewhere("old-sid", old_parent, evt) is True
+    assert server._notification_event_belongs_elsewhere("tip-sid", live_tip, evt) is False
+    # A third session must leave it alone for the continuation's poller.
+    third = _session(session_key="third")
+    monkeypatch.setattr(
+        server,
+        "_sessions",
+        {"old-sid": old_parent, "tip-sid": live_tip, "third-sid": third},
+    )
+    assert server._notification_event_belongs_elsewhere("third-sid", third, evt) is True
+
+
+def test_finalized_origin_ui_session_falls_back_to_live_continuation(monkeypatch):
+    """A closed origin tab must not steal its resumed continuation's result."""
+    finalized_origin = _session(session_key="old-parent", _finalized=True)
+    live_tip = _session(session_key="new-tip")
+    monkeypatch.setattr(
+        server,
+        "_sessions",
+        {"origin-sid": finalized_origin, "tip-sid": live_tip},
+    )
+
+    class _DB:
+        def resolve_resume_session_id(self, session_id):
+            return "new-tip" if session_id == "old-parent" else session_id
+
+    monkeypatch.setattr(server, "_get_db", lambda: _DB())
+    evt = {
+        "type": "async_delegation",
+        "session_key": "old-parent",
+        "origin_ui_session_id": "origin-sid",
+    }
+
+    assert server._notification_event_belongs_elsewhere("origin-sid", finalized_origin, evt) is True
+    assert server._notification_event_belongs_elsewhere("tip-sid", live_tip, evt) is False
 
 
 def test_prompt_submit_rejects_negative_truncate_ordinal(monkeypatch):

tests/tools/test_async_delegation.py

@@ -289,6 +289,68 @@ def slow_child(task_index, goal, child=None, parent_agent=None, **kw):
     assert "the real task" in text
 
 
+def test_delegate_task_background_uses_live_tui_agent_session_id(monkeypatch):
+    """TUI async delegation must route to the live/compressed agent id.
+
+    Regression: delegate_task captured the stale approval/session context key
+    after compression rotated parent_agent.session_id. The resulting completion
+    was orphaned and could be consumed by an unrelated desktop session poller.
+    """
+    import json
+    from unittest.mock import MagicMock
+    import tools.delegate_tool as dt
+    from gateway.session_context import clear_session_vars, set_session_vars
+    from tools.approval import reset_current_session_key, set_current_session_key
+
+    parent = MagicMock()
+    parent._delegate_depth = 0
+    parent.session_id = "post-compress-tip"
+    parent._interrupt_requested = False
+    parent._active_children = []
+    parent._active_children_lock = None
+    fake_child = MagicMock()
+    fake_child._delegate_role = "leaf"
+
+    creds = {
+        "model": "m", "provider": None, "base_url": None, "api_key": None,
+        "api_mode": None, "command": None, "args": None,
+    }
+    monkeypatch.setattr(dt, "_build_child_agent", lambda **kw: fake_child)
+    monkeypatch.setattr(dt, "_resolve_delegation_credentials", lambda *a, **k: creds)
+    monkeypatch.setattr(
+        dt,
+        "_run_single_child",
+        lambda *a, **k: {
+            "task_index": 0,
+            "status": "completed",
+            "summary": "done",
+            "api_calls": 1,
+            "duration_seconds": 0.1,
+            "model": "m",
+            "exit_reason": "completed",
+        },
+    )
+
+    approval_token = set_current_session_key("pre-compress-parent")
+    session_tokens = set_session_vars(
+        source="tui",
+        session_key="pre-compress-parent",
+        ui_session_id="origin-tab",
+    )
+    try:
+        out = dt.delegate_task(goal="bg task", background=True, parent_agent=parent)
+        assert json.loads(out)["status"] == "dispatched"
+        evt = _drain_one()
+    finally:
+        reset_current_session_key(approval_token)
+        clear_session_vars(session_tokens)
+
+    assert evt is not None
+    assert evt["type"] == "async_delegation"
+    assert evt["session_key"] == "post-compress-tip"
+    assert evt["origin_ui_session_id"] == "origin-tab"
+
+
 def test_delegate_task_background_batch_runs_as_one_unit(monkeypatch):
     """A multi-item batch with background=True dispatches the WHOLE fan-out as
     ONE background unit (one handle, one async slot). The children run in

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"