From a6fbff1d7ecee0e9f8dfeb12f6ef9ae0ef9f469c Mon Sep 17 00:00:00 2001 From: Alex Fournier Date: Thu, 23 Jul 2026 13:22:50 -0700 Subject: [PATCH] fix(relay): retain active child sessions Signed-off-by: Alex Fournier --- agent/relay_runtime.py | 31 ++++++ .../test_relay_shared_metrics_runtime.py | 48 ++++++++++ tests/tools/test_zombie_process_cleanup.py | 94 ++++++++++++++++++- tools/delegate_tool.py | 6 +- 4 files changed, 175 insertions(+), 4 deletions(-) diff --git a/agent/relay_runtime.py b/agent/relay_runtime.py index 0b1c341d34c9..2b49725cbcfc 100644 --- a/agent/relay_runtime.py +++ b/agent/relay_runtime.py @@ -473,6 +473,7 @@ class RelayTurnContext: default=None, repr=False, ) + _active_registered: bool = field(default=False, repr=False) closed: bool = False @@ -491,6 +492,8 @@ class RelaySessionCoordinator: str, Callable[[RelayRuntime, dict[str, Any]], None], ] = {} + self._active_turns_lock = threading.RLock() + self._active_turns: dict[tuple[str, str], set[int]] = {} def register_session_initializer( self, @@ -602,6 +605,10 @@ class RelaySessionCoordinator: except Exception: logger.warning("Hermes Relay turn initialization failed", exc_info=True) turn._token = _CURRENT_TURN.set(turn) + key = (lease.profile_key, lease.session_id) + with self._active_turns_lock: + self._active_turns.setdefault(key, set()).add(id(turn)) + turn._active_registered = True return turn def end_turn( @@ -636,8 +643,31 @@ class RelaySessionCoordinator: "Hermes Relay turn finalization failed", exc_info=True ) finally: + self._unregister_active_turn(turn) self._reset_turn_context(turn) + def has_active_turn(self, *, profile_key: str, session_id: str) -> bool: + """Return whether a turn is still running for one profile/session.""" + key = (profile_key, session_id) + with self._active_turns_lock: + return bool(self._active_turns.get(key)) + + def _unregister_active_turn(self, turn: RelayTurnContext) -> None: + if not turn._active_registered: + return + key = (turn.lease.profile_key, turn.lease.session_id) + with self._active_turns_lock: + active = self._active_turns.get(key) + if active is not None: + active.discard(id(turn)) + if not active: + self._active_turns.pop(key, None) + turn._active_registered = False + + def _reset_active_turns_for_tests(self) -> None: + with self._active_turns_lock: + self._active_turns.clear() + def finish_logical_calls( self, turn: RelayTurnContext, @@ -913,4 +943,5 @@ def _session_id(event: dict[str, Any]) -> str: def _reset_for_tests() -> None: """Reset all profile-scoped Relay hosts for isolated tests.""" + SESSION_COORDINATOR._reset_active_turns_for_tests() HOST_REGISTRY.shutdown_all() diff --git a/tests/hermes_cli/test_relay_shared_metrics_runtime.py b/tests/hermes_cli/test_relay_shared_metrics_runtime.py index 36f89bccc44f..7b39d0973c69 100644 --- a/tests/hermes_cli/test_relay_shared_metrics_runtime.py +++ b/tests/hermes_cli/test_relay_shared_metrics_runtime.py @@ -1535,6 +1535,54 @@ def test_concurrent_subagents_inherit_parent_turn_and_close_independently( ) +def test_coordinator_tracks_active_turns_across_threads(direct_runtime): + del direct_runtime + coordinator = relay_runtime.SESSION_COORDINATOR + profile_key = relay_runtime.current_profile_key() + lease = coordinator.acquire_conversation( + profile_key=profile_key, + session_id="cross-thread-child", + platform="subagent", + ) + + assert not coordinator.has_active_turn( + profile_key=profile_key, + session_id="cross-thread-child", + ) + + turn = coordinator.begin_turn( + lease, + turn_id="child-turn", + task_id="child-task", + ) + + observed = [] + thread = threading.Thread( + target=lambda: observed.append( + coordinator.has_active_turn( + profile_key=profile_key, + session_id="cross-thread-child", + ) + ) + ) + thread.start() + thread.join(timeout=5) + + assert observed == [True] + + coordinator.end_turn(turn, outcome="success") + + assert not coordinator.has_active_turn( + profile_key=profile_key, + session_id="cross-thread-child", + ) + coordinator.release_conversation(lease) + coordinator.finalize_conversation( + profile_key=profile_key, + session_id="cross-thread-child", + ) + + def test_core_runtime_ignores_self_parenting_subagent_event(direct_runtime): runtime = relay_runtime.get_runtime() assert runtime is not None diff --git a/tests/tools/test_zombie_process_cleanup.py b/tests/tools/test_zombie_process_cleanup.py index 6207c159da68..b4679ffbe3ac 100644 --- a/tests/tools/test_zombie_process_cleanup.py +++ b/tests/tools/test_zombie_process_cleanup.py @@ -342,7 +342,6 @@ class TestDelegationCleanup: raise RuntimeError("test abort") child.run_conversation.side_effect = run_conversation - child._relay_pending_turn_id = None relay_host = MagicMock() monkeypatch.setattr(relay_runtime, "get_runtime", lambda **kwargs: relay_host) @@ -379,12 +378,16 @@ class TestDelegationCleanup: parent._active_children_lock = threading.Lock() child = MagicMock() child.session_id = "active-child-session" - child._relay_pending_turn_id = "active-child-turn" child._delegate_saved_tool_names = ["tool1"] child.run_conversation.side_effect = RuntimeError("test abort") parent._active_children.append(child) relay_host = MagicMock() monkeypatch.setattr(relay_runtime, "get_runtime", lambda **kwargs: relay_host) + monkeypatch.setattr( + relay_runtime.SESSION_COORDINATOR, + "has_active_turn", + lambda **_kwargs: True, + ) result = _run_single_child( task_index=0, @@ -395,3 +398,90 @@ class TestDelegationCleanup: assert result["status"] == "error" relay_host.unregister_subagent.assert_not_called() + + def test_timed_out_child_keeps_relay_session_until_its_turn_exits( + self, monkeypatch, tmp_path + ): + from unittest.mock import MagicMock + + from agent import relay_runtime + from hermes_constants import ( + reset_hermes_home_override, + set_hermes_home_override, + ) + from tools.delegate_tool import _run_single_child + + relay_runtime._reset_for_tests() + profile_home = tmp_path / "profile-timeout" + profile_token = set_hermes_home_override(profile_home) + child_started = threading.Event() + release_child = threading.Event() + child_finished = threading.Event() + parent = MagicMock() + parent._active_children = [] + parent._active_children_lock = threading.Lock() + child = MagicMock() + child.session_id = "timed-out-child" + child._delegate_saved_tool_names = ["tool1"] + child.get_activity_summary.return_value = {"api_call_count": 1} + parent._active_children.append(child) + relay_host = MagicMock() + monkeypatch.setattr(relay_runtime, "get_runtime", lambda **_kwargs: relay_host) + monkeypatch.setattr("tools.delegate_tool._get_child_timeout", lambda: 0.1) + + def run_conversation(**kwargs): + lease = relay_runtime.SESSION_COORDINATOR.acquire_conversation( + profile_key=relay_runtime.current_profile_key(), + session_id=child.session_id, + platform="subagent", + ) + turn = relay_runtime.SESSION_COORDINATOR.begin_turn( + lease, + turn_id="timed-out-child-turn", + task_id=kwargs["task_id"], + ) + child_started.set() + try: + release_child.wait(timeout=5) + return { + "final_response": "late result", + "completed": True, + "interrupted": False, + "api_calls": 1, + "messages": [], + } + finally: + relay_runtime.SESSION_COORDINATOR.end_turn( + turn, + outcome="cancelled", + ) + relay_runtime.SESSION_COORDINATOR.release_conversation(lease) + child_finished.set() + + child.run_conversation.side_effect = run_conversation + try: + result = _run_single_child( + task_index=0, + goal="test timed-out turn cleanup", + child=child, + parent_agent=parent, + ) + + assert child_started.is_set() + assert result["status"] == "timeout" + assert relay_runtime.SESSION_COORDINATOR.has_active_turn( + profile_key=str(profile_home), + session_id=child.session_id, + ) + relay_host.unregister_subagent.assert_not_called() + + release_child.set() + assert child_finished.wait(timeout=5) + assert not relay_runtime.SESSION_COORDINATOR.has_active_turn( + profile_key=str(profile_home), + session_id=child.session_id, + ) + finally: + release_child.set() + reset_hermes_home_override(profile_token) + relay_runtime._reset_for_tests() diff --git a/tools/delegate_tool.py b/tools/delegate_tool.py index 5326a06f81a1..dbe69aed433a 100644 --- a/tools/delegate_tool.py +++ b/tools/delegate_tool.py @@ -2411,8 +2411,10 @@ def _run_single_child( runtime = relay_runtime.get_runtime(create=False) child_session_id = str(getattr(child, "session_id", "") or "") - pending_turn = getattr(child, "_relay_pending_turn_id", None) - child_turn_is_active = isinstance(pending_turn, str) and bool(pending_turn) + child_turn_is_active = relay_runtime.SESSION_COORDINATOR.has_active_turn( + profile_key=relay_runtime.current_profile_key(), + session_id=child_session_id, + ) if runtime is not None and child_session_id and not child_turn_is_active: runtime.unregister_subagent({"child_session_id": child_session_id}) except Exception: