diff --git a/agent/relay_llm.py b/agent/relay_llm.py index 8f560c1854..58a5c1bd5e 100644 --- a/agent/relay_llm.py +++ b/agent/relay_llm.py @@ -55,6 +55,13 @@ def execute( def invoke(next_request: Any) -> Any: nonlocal callback_error + + def guarded(final: dict[str, Any]) -> Any: + # Nested relay calls inside a managed provider callback must run + # unmanaged (#77244) — see relay_runtime.managed_callback_guard. + with relay_runtime.managed_callback_guard(): + return callback(final) + try: final_request = _provider_request( request, @@ -63,7 +70,7 @@ def execute( codec_baseline_body=codec_baseline_body, metadata=metadata, ) - raw = callback_context.copy().run(callback, final_request) + raw = callback_context.copy().run(guarded, final_request) except BaseException as exc: callback_error = exc raise @@ -149,7 +156,10 @@ async def execute_async( metadata=metadata, ) async def call_provider() -> Any: - return await callback(final_request) + # Nested relay calls inside a managed provider callback must + # run unmanaged (#77244). + with relay_runtime.managed_callback_guard(): + return await callback(final_request) task = callback_context.copy().run( asyncio.create_task, @@ -396,7 +406,14 @@ class ManagedLlmStream(Iterator[Any]): def run_callback(callback: Callable[..., Any], *args: Any) -> Any: # Relay can invoke stream surfaces while another callback still # owns the captured Context. A fresh copy is safe to enter. - return callback_context.copy().run(callback, *args) + def guarded() -> Any: + # Hermes-side callbacks run while the native pipeline drives + # this stream; nested relay calls they make must bypass + # managed execution (#77244). + with relay_runtime.managed_callback_guard(): + return callback(*args) + + return callback_context.copy().run(guarded) runtime, session, parent = relay_runtime.resolve_execution_context(session_id) if ( diff --git a/agent/relay_runtime.py b/agent/relay_runtime.py index a0a7315796..095fbb8315 100644 --- a/agent/relay_runtime.py +++ b/agent/relay_runtime.py @@ -492,6 +492,34 @@ _CURRENT_TURN: contextvars.ContextVar[RelayTurnContext | None] = contextvars.Con "hermes_relay_turn", default=None ) +# Depth of managed Relay callbacks executing on the current logical call path. +# Set >0 while the native Relay pipeline is mid-dispatch of a Hermes callback +# (tool or LLM). Nested managed execution inside that window is structurally +# broken — the native pipeline binds its Futures to the outer, blocked event +# loop — so resolve_execution_context() bypasses Relay while the flag is set. +# ContextVar so the marker follows contextvars.copy_context() into the worker +# threads / per-thread loops that tools use for their internal async work. +_MANAGED_CALLBACK_DEPTH: contextvars.ContextVar[int] = contextvars.ContextVar( + "hermes_relay_managed_callback_depth", default=0 +) + + +class managed_callback_guard: + """Mark the current context as inside a managed Relay callback. + + Synchronous context manager used by the relay adapters around the + ``invoke()`` callbacks they hand to the native pipeline. Everything the + callback transitively calls (including work it forwards to worker threads + via ``contextvars.copy_context()``) sees the marker and runs unmanaged. + """ + + def __enter__(self) -> "managed_callback_guard": + self._token = _MANAGED_CALLBACK_DEPTH.set(_MANAGED_CALLBACK_DEPTH.get() + 1) + return self + + def __exit__(self, *exc_info: Any) -> None: + _MANAGED_CALLBACK_DEPTH.reset(self._token) + class RelaySessionCoordinator: """Own semantic conversation and turn lifetimes for Hermes core.""" @@ -843,6 +871,20 @@ def resolve_execution_context( session_id: str, ) -> tuple[RelayRuntime | None, RelaySession | None, Any]: """Resolve one active turn/session parent for managed Relay execution.""" + if _MANAGED_CALLBACK_DEPTH.get() > 0: + # A managed Relay callback is already executing on this logical call + # path (e.g. the native ``tools.execute`` pipeline is mid-dispatch of + # a Hermes tool). Nested managed execution here is structurally + # impossible: the native pipeline binds its Futures to the OUTER + # call's event loop, which is blocked inside the synchronous tool + # callback until the tool returns. A nested managed LLM call (the + # vision_analyze auxiliary path) therefore awaits a foreign-loop + # Future that can never complete — "attached to a different loop" + # at best, deadlock at worst, and "Event loop is closed" during + # shutdown when the orphaned Future is completed late (#77244). + # Run nested calls unmanaged; the outer tool scope still records + # the tool-level event for observability. + return None, None, None inherited_turn = current_turn() if inherited_turn is not None and ( not inherited_turn.relay_enabled or inherited_turn.closed diff --git a/agent/relay_tools.py b/agent/relay_tools.py index f2177f1a3b..91c0926b7e 100644 --- a/agent/relay_tools.py +++ b/agent/relay_tools.py @@ -36,8 +36,17 @@ def execute( def invoke(next_args: Any) -> Any: nonlocal callback_error, observed_args observed_args = next_args if isinstance(next_args, dict) else args + + def guarded(final_args: dict[str, Any]) -> Any: + # Everything the tool transitively calls (including auxiliary LLM + # calls it forwards to worker threads) must bypass managed Relay + # execution — the native pipeline's Futures bind to THIS loop, + # which is blocked until the tool returns (#77244). + with relay_runtime.managed_callback_guard(): + return callback(final_args) + try: - result = callback_context.copy().run(callback, observed_args) + result = callback_context.copy().run(guarded, observed_args) except BaseException as exc: callback_error = exc raise diff --git a/tests/agent/test_relay_nested_execution.py b/tests/agent/test_relay_nested_execution.py new file mode 100644 index 0000000000..45671922bd --- /dev/null +++ b/tests/agent/test_relay_nested_execution.py @@ -0,0 +1,233 @@ +"""Regression tests for nested managed Relay execution (#77244). + +The native Relay pipeline binds its Futures to the event loop that entered +``run_in_session_async``. While a managed tool callback is executing, that +loop is blocked until the callback returns — so any NESTED managed relay call +made from inside the callback (e.g. vision_analyze's auxiliary LLM call on a +worker-thread loop) awaits a Future that can never complete: +``RuntimeError: ... attached to a different loop``, or a deadlock, or +``Event loop is closed`` at shutdown. + +The fix: ``relay_runtime.managed_callback_guard`` marks the callback's +context (a ContextVar, so it propagates through ``contextvars.copy_context()`` +into tool worker threads); ``resolve_execution_context`` returns the +no-relay triple while the marker is set, so nested calls run unmanaged. +""" + +from __future__ import annotations + +import asyncio +import contextvars +import threading + +import pytest + +pytest.importorskip("nemo_relay") + +from agent import relay_llm, relay_runtime, relay_tools + + +@pytest.fixture() +def relay_turn(tmp_path, monkeypatch): + monkeypatch.setenv("HERMES_HOME", str(tmp_path / "profile")) + relay_runtime._reset_for_tests() + lease = relay_runtime.SESSION_COORDINATOR.acquire_conversation( + profile_key=relay_runtime.current_profile_key(), + session_id="session-1", + platform="cli", + ) + turn = relay_runtime.SESSION_COORDINATOR.begin_turn( + lease, + turn_id="turn-1", + task_id="task-1", + ) + lease.host.retain_managed_execution("test.nested_relay") + try: + yield lease.host + finally: + lease.host.release_managed_execution("test.nested_relay") + relay_runtime.SESSION_COORDINATOR.end_turn(turn, outcome="success") + relay_runtime.SESSION_COORDINATOR.release_conversation(lease) + relay_runtime._reset_for_tests() + + +def _nested_aux_llm_call_from_worker_thread() -> dict: + """Mimic vision_analyze: aux LLM call via a worker thread's own loop.""" + + async def aux_call(): + async def provider(request): + await asyncio.sleep(0) + return { + "id": "aux-1", + "object": "chat.completion", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "nested"}, + "finish_reason": "stop", + } + ], + } + + return await relay_llm.execute_current_async( + {"messages": [{"role": "user", "content": "look"}], "model": "m"}, + provider, + name="nested-prov", + model_name="m", + metadata={ + "api_mode": "chat_completions", + "api_request_id": "req-nested", + "call_role": "auxiliary:vision", + }, + ) + + holder: dict = {} + + def run() -> None: + loop = asyncio.new_event_loop() + try: + holder["result"] = loop.run_until_complete(aux_call()) + except BaseException as exc: # pragma: no cover - assertion payload + holder["error"] = exc + finally: + loop.close() + + ctx = contextvars.copy_context() + thread = threading.Thread(target=lambda: ctx.run(run)) + thread.start() + thread.join(timeout=30) + assert not thread.is_alive(), "nested aux call deadlocked (#77244)" + if "error" in holder: + raise holder["error"] + return holder["result"] + + +def test_nested_aux_llm_call_inside_managed_tool_does_not_cross_loops(relay_turn): + """The #77244 shape: managed tool -> worker-thread aux LLM call.""" + host = relay_turn + managed_llm_names: list[str] = [] + original_execute = host.relay.llm.execute + + def counting_execute(name, *args, **kwargs): + managed_llm_names.append(name) + return original_execute(name, *args, **kwargs) + + host.relay.llm.execute = counting_execute + try: + def the_tool(args): + result = _nested_aux_llm_call_from_worker_thread() + return {"analysis": result["choices"][0]["message"]["content"]} + + result, _final_args = relay_tools.execute( + "vision_analyze", + {"image_url": "/tmp/x.png"}, + the_tool, + session_id="session-1", + metadata={"api_request_id": "req-tool"}, + ) + finally: + host.relay.llm.execute = original_execute + + assert "nested" in str(result) + # The nested call must have bypassed the managed pipeline entirely. + assert "nested-prov" not in managed_llm_names + + +def test_main_turn_llm_call_stays_managed(relay_turn): + """The guard must not disable relay for top-level (non-nested) calls.""" + host = relay_turn + managed_llm_names: list[str] = [] + original_execute = host.relay.llm.execute + + def counting_execute(name, *args, **kwargs): + managed_llm_names.append(name) + return original_execute(name, *args, **kwargs) + + host.relay.llm.execute = counting_execute + try: + out = relay_llm.execute( + {"messages": [{"role": "user", "content": "hi"}], "model": "m"}, + lambda request: { + "id": "main-1", + "object": "chat.completion", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "main"}, + "finish_reason": "stop", + } + ], + }, + session_id="session-1", + name="main-prov", + model_name="m", + metadata={ + "api_mode": "chat_completions", + "api_request_id": "req-main", + "call_role": "primary", + }, + ) + finally: + host.relay.llm.execute = original_execute + + assert out is not None + assert "main-prov" in managed_llm_names + + +def test_guard_resets_after_managed_callback_returns(relay_turn): + """After the tool returns, subsequent calls are managed again.""" + host = relay_turn + managed_llm_names: list[str] = [] + original_execute = host.relay.llm.execute + + def counting_execute(name, *args, **kwargs): + managed_llm_names.append(name) + return original_execute(name, *args, **kwargs) + + host.relay.llm.execute = counting_execute + try: + relay_tools.execute( + "noop_tool", + {}, + lambda args: {"ok": True}, + session_id="session-1", + metadata={"api_request_id": "req-tool2"}, + ) + assert relay_runtime._MANAGED_CALLBACK_DEPTH.get() == 0 + relay_llm.execute( + {"messages": [{"role": "user", "content": "hi"}], "model": "m"}, + lambda request: { + "id": "after-1", + "object": "chat.completion", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "after"}, + "finish_reason": "stop", + } + ], + }, + session_id="session-1", + name="after-prov", + model_name="m", + metadata={ + "api_mode": "chat_completions", + "api_request_id": "req-after", + "call_role": "primary", + }, + ) + finally: + host.relay.llm.execute = original_execute + + assert "after-prov" in managed_llm_names + + +def test_resolve_execution_context_bypasses_inside_guard(relay_turn): + with relay_runtime.managed_callback_guard(): + runtime, session, parent = relay_runtime.resolve_execution_context( + "session-1" + ) + assert runtime is None and session is None and parent is None + # Outside the guard the context resolves normally again. + runtime, session, _parent = relay_runtime.resolve_execution_context("session-1") + assert runtime is not None and session is not None