fix(relay): retain active child sessions

Signed-off-by: Alex Fournier <afournier@nvidia.com>
This commit is contained in:
Alex Fournier 2026-07-23 13:22:50 -07:00
parent 398a8ca580
commit a6fbff1d7e
4 changed files with 175 additions and 4 deletions

View file

@ -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()

View file

@ -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

View file

@ -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()

View file

@ -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: