diff --git a/gateway/session_context.py b/gateway/session_context.py index cdd1a8bfafe..2214c2a5688 100644 --- a/gateway/session_context.py +++ b/gateway/session_context.py @@ -79,6 +79,13 @@ _SESSION_USER_ID: ContextVar = ContextVar("HERMES_SESSION_USER_ID", default=_UNS _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 @@ _VAR_MAP = { "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, ): diff --git a/tests/test_tui_gateway_server.py b/tests/test_tui_gateway_server.py index 50a9fa8d5ce..cc6ff4f5b72 100644 --- a/tests/test_tui_gateway_server.py +++ b/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): diff --git a/tests/tools/test_async_delegation.py b/tests/tools/test_async_delegation.py index 0cbd9313cfb..9a473394b51 100644 --- a/tests/tools/test_async_delegation.py +++ b/tests/tools/test_async_delegation.py @@ -289,6 +289,68 @@ def test_delegate_task_background_routes_async_and_does_not_block(monkeypatch): 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 diff --git a/tools/async_delegation.py b/tools/async_delegation.py index f28156e2f57..710ab5c1d3b 100644 --- a/tools/async_delegation.py +++ b/tools/async_delegation.py @@ -130,6 +130,7 @@ def dispatch_async_delegation( model: Optional[str], session_key: str, runner: Callable[[], Dict[str, Any]], + origin_ui_session_id: str = "", interrupt_fn: Optional[Callable[[], None]] = None, max_async_children: int = _DEFAULT_MAX_ASYNC_CHILDREN, ) -> Dict[str, Any]: @@ -172,6 +173,7 @@ def dispatch_async_delegation( "role": role, "model": model, "session_key": session_key, + "origin_ui_session_id": origin_ui_session_id, "status": "running", "dispatched_at": dispatched_at, "completed_at": None, @@ -282,6 +284,7 @@ def _push_completion_event( # session_key routes the completion back to the originating gateway # session; empty string => CLI (single-session) path. "session_key": record.get("session_key", ""), + "origin_ui_session_id": record.get("origin_ui_session_id", ""), "goal": record.get("goal", ""), "context": record.get("context"), "toolsets": record.get("toolsets"), @@ -317,6 +320,7 @@ def dispatch_async_delegation_batch( model: Optional[str], session_key: str, runner: Callable[[], Dict[str, Any]], + origin_ui_session_id: str = "", interrupt_fn: Optional[Callable[[], None]] = None, max_async_children: int = _DEFAULT_MAX_ASYNC_CHILDREN, ) -> Dict[str, Any]: @@ -356,6 +360,7 @@ def dispatch_async_delegation_batch( "role": role, "model": model, "session_key": session_key, + "origin_ui_session_id": origin_ui_session_id, "status": "running", "dispatched_at": dispatched_at, "completed_at": None, @@ -453,6 +458,7 @@ def _finalize_batch( "type": "async_delegation", "delegation_id": delegation_id, "session_key": event_record.get("session_key", ""), + "origin_ui_session_id": event_record.get("origin_ui_session_id", ""), "goal": event_record.get("goal", ""), "goals": event_record.get("goals"), "context": event_record.get("context"), diff --git a/tools/delegate_tool.py b/tools/delegate_tool.py index b989b7983b1..4179926091a 100644 --- a/tools/delegate_tool.py +++ b/tools/delegate_tool.py @@ -2796,6 +2796,25 @@ def delegate_task( return json.dumps(_sync_result, ensure_ascii=False) _session_key = get_current_session_key(default="") + _origin_ui_session_id = "" + try: + from gateway.session_context import get_session_env + + _source = get_session_env("HERMES_SESSION_SOURCE", "") + _origin_ui_session_id = get_session_env("HERMES_UI_SESSION_ID", "") + # In desktop/TUI, the routable session key is the durable + # AIAgent.session_id. Context compression can rotate that id during + # the same turn before the TUI-side session dict is re-anchored; + # if we capture the stale approval/session context key here, the + # async completion becomes an orphan and any desktop poller may + # consume it. Gateway chats are different: their session_key is the + # platform conversation key (agent:main:...), so keep it there. + if _source == "tui": + _agent_session_id = str(getattr(parent_agent, "session_id", "") or "") + if _agent_session_id: + _session_key = _agent_session_id + except Exception: + _origin_ui_session_id = "" _child_agents = [c for (_, _, c) in children] # Detach every child from the parent's interrupt-propagation list — the @@ -2836,6 +2855,7 @@ def delegate_task( role=top_role, model=creds["model"], session_key=_session_key, + origin_ui_session_id=_origin_ui_session_id, runner=_batch_runner, interrupt_fn=_batch_interrupt, max_async_children=_get_max_async_children(), diff --git a/tui_gateway/server.py b/tui_gateway/server.py index 8092147bc79..29c7dd42d20 100644 --- a/tui_gateway/server.py +++ b/tui_gateway/server.py @@ -1946,7 +1946,12 @@ def _cwd_for_session_key(session_key: str) -> str: return "" -def _set_session_context(session_key: str, cwd: str | None = None) -> list: +def _set_session_context( + session_key: str, + cwd: str | None = None, + *, + ui_session_id: str = "", +) -> list: try: from gateway.session_context import set_session_vars @@ -1961,7 +1966,12 @@ def _set_session_context(session_key: str, cwd: str | None = None) -> list: if sess.get("session_key") == session_key: source = _session_source(sess) break - return set_session_vars(session_key=session_key, source=source, cwd=resolved) + return set_session_vars( + session_key=session_key, + source=source, + cwd=resolved, + ui_session_id=ui_session_id, + ) except Exception: return [] @@ -8451,22 +8461,76 @@ def _(rid, params: dict) -> dict: return _ok(rid, {"status": "streaming"}) -def _notification_event_belongs_elsewhere(session: dict, evt: dict) -> bool: +def _notification_event_belongs_elsewhere(sid: str, session: dict, evt: dict) -> bool: """True if ``evt`` is owned by a *different* live session. - Background-process events carry the ``session_key`` of the session that - started the process. Since all desktop sessions share one process-wide - completion queue, each poller must skip events it doesn't own so a - background job's completion surfaces in the session that launched it — not - whichever poller happened to dequeue first. Orphaned events (owner gone) - and global/system events (empty ``session_key``) return False so the - current poller still handles them rather than losing them. + Background completions carry the ``session_key`` of the session that started + the work. Async delegation completions from the desktop also carry + ``origin_ui_session_id``: the live TUI tab/window that commissioned them. + Since all desktop sessions share one process-wide completion queue, each + poller must skip events it doesn't own so a detached result surfaces in the + launching session, not whichever poller happened to dequeue first. """ + evt_ui_sid = str(evt.get("origin_ui_session_id") or "") + if evt_ui_sid: + if evt_ui_sid == str(sid or "") and not session.get("_finalized"): + return False + try: + with _sessions_lock: + owner_live = evt_ui_sid in _sessions and not _sessions[evt_ui_sid].get("_finalized") + except Exception: + owner_live = False + if owner_live: + return True + # If the exact UI tab is gone, fall through to durable session_key + # routing. That avoids wrong-session delivery while still allowing a + # resumed continuation with the same durable key/lineage to claim it. + evt_key = str(evt.get("session_key") or "") if not evt_key: return False - if evt_key == str(session.get("session_key") or ""): + + current_keys = { + str(session.get("session_key") or ""), + _session_lookup_key(session, fallback=sid), + } + + # Compression can rotate AIAgent.session_id while the detached child is + # still running. Resolve the event's original key to its continuation tip so + # an event captured before or after compression still maps to the same live + # desktop session instead of becoming an orphan that any poller may consume. + resolved_key = evt_key + try: + db = _get_db() + if db is not None: + resolved_key = db.resolve_resume_session_id(evt_key) or evt_key + except Exception: + resolved_key = evt_key + + # If the key has a live continuation, prefer that continuation over the + # compressed parent. Otherwise a stale parent tab could consume the event + # before the real current conversation sees it. + if resolved_key != evt_key: + if resolved_key in current_keys: + return False + try: + with _sessions_lock: + continuation_live = any( + not s.get("_finalized") + and ( + str(s.get("session_key") or "") == resolved_key + or _session_lookup_key(s, fallback="") == resolved_key + ) + for s in _sessions.values() + ) + except Exception: + continuation_live = False + if continuation_live: + return True + + if evt_key in current_keys: return False + try: with _sessions_lock: snapshot = list(_sessions.values()) @@ -8476,7 +8540,12 @@ def _notification_event_belongs_elsewhere(session: dict, evt: dict) -> bool: return False return any( - s is not session and str(s.get("session_key") or "") == evt_key + s is not session + and not s.get("_finalized") + and ( + str(s.get("session_key") or "") in {evt_key, resolved_key} + or _session_lookup_key(s, fallback="") in {evt_key, resolved_key} + ) for s in snapshot ) @@ -8546,7 +8615,7 @@ def _notification_poller_loop( # process started in session A would surface its completion in whichever # session's poller happened to wake first (Ben's "reported in a # different session" bug). Leave foreign events for their owner. - if _notification_event_belongs_elsewhere(session, evt): + if _notification_event_belongs_elsewhere(sid, session, evt): process_registry.completion_queue.put(evt) time.sleep(0.1) continue @@ -8604,7 +8673,7 @@ def _notification_poller_loop( evt = process_registry.completion_queue.get_nowait() except Exception: break - if _notification_event_belongs_elsewhere(session, evt): + if _notification_event_belongs_elsewhere(sid, session, evt): deferred.append(evt) continue _evt_sid = evt.get("session_id", "") @@ -8729,7 +8798,10 @@ def _run_prompt_submit(rid, sid: str, session: dict, text: Any) -> None: ) approval_token = set_current_session_key(session["session_key"]) - session_tokens = _set_session_context(session["session_key"]) + session_tokens = _set_session_context( + session["session_key"], + ui_session_id=sid, + ) _profile_home_str = session.get("profile_home") if _profile_home_str: home_token = set_hermes_home_override(_profile_home_str)