diff --git a/agent/relay_llm.py b/agent/relay_llm.py index 2d8c37a18f..b49d527ee5 100644 --- a/agent/relay_llm.py +++ b/agent/relay_llm.py @@ -17,12 +17,11 @@ from agent import relay_runtime logger = logging.getLogger(__name__) -_PROVIDER_MESSAGE_EXTENSION_KEYS = frozenset( - {"reasoning_content", "reasoning_details"} -) -_RELAY_INTERNAL_PROVIDER_HEADERS = frozenset( - {"x-dynamo-parent-session-id", "x-dynamo-session-id"} -) +_PROVIDER_MESSAGE_EXTENSION_KEYS = frozenset({"reasoning_content", "reasoning_details"}) +_RELAY_INTERNAL_PROVIDER_HEADERS = frozenset({"x-dynamo-parent-session-id", "x-dynamo-session-id"}) +_LogicalCall = tuple[relay_runtime.RelayTurnContext, Any, str] + + @dataclass(frozen=True, slots=True) class _RelayProtocol: operation: str @@ -30,27 +29,20 @@ class _RelayProtocol: _RELAY_PROTOCOL_BY_API_MODE = { - "chat_completions": _RelayProtocol( - operation="openai.chat_completions", - codec_class="OpenAIChatCodec", - ), - "codex_responses": _RelayProtocol( - operation="openai.responses", - codec_class="OpenAIResponsesCodec", - ), - "anthropic_messages": _RelayProtocol( - operation="anthropic.messages", - codec_class="AnthropicMessagesCodec", - ), + "chat_completions": _RelayProtocol("openai.chat_completions", "OpenAIChatCodec"), + "codex_responses": _RelayProtocol("openai.responses", "OpenAIResponsesCodec"), + "anthropic_messages": _RelayProtocol("anthropic.messages", "AnthropicMessagesCodec"), } +def _api_mode(metadata: dict[str, Any] | None) -> str: + return str((metadata or {}).get("api_mode") or "") + + def _relay_protocol(metadata: dict[str, Any] | None) -> _RelayProtocol | None: """Return Relay's operation and codec descriptor for an API mode.""" api_mode = (metadata or {}).get("api_mode") - if not isinstance(api_mode, str): - return None - return _RELAY_PROTOCOL_BY_API_MODE.get(api_mode) + return _RELAY_PROTOCOL_BY_API_MODE.get(api_mode) if isinstance(api_mode, str) else None def _relay_operation_name(provider_name: str, metadata: dict[str, Any] | None) -> str: @@ -59,17 +51,141 @@ def _relay_operation_name(provider_name: str, metadata: dict[str, Any] | None) - return protocol.operation if protocol is not None else provider_name -def _relay_metadata( - provider_name: str, metadata: dict[str, Any] | None -) -> dict[str, Any]: +def _relay_metadata(provider_name: str, metadata: dict[str, Any] | None) -> dict[str, Any]: """Preserve the physical provider when the operation name is canonicalized.""" - relay_metadata = _jsonable(metadata or {}) - if not isinstance(relay_metadata, dict): - relay_metadata = {} + relay_metadata = _jsonable_dict(metadata or {}) relay_metadata.setdefault("hermes.provider", provider_name) return relay_metadata +class _ManagedAttempt: + """Relay request state shared by the sync, async, and streaming adapters.""" + + @classmethod + def resolve( + cls, + session_id: str, + request: dict[str, Any], + metadata: dict[str, Any] | None, + *, + name: str, + model_name: str, + ) -> "_ManagedAttempt | None": + """Return the managed attempt for ``session_id``, or None to run unmanaged.""" + runtime, session, parent = relay_runtime.resolve_execution_context(session_id) + if runtime is None or session is None or not runtime.managed_execution_enabled(): + return None + return cls(runtime, session, parent, request, metadata, name=name, model_name=model_name) + + def __init__( + self, + runtime: relay_runtime.RelayRuntime, + session: Any, + parent: Any, + request: dict[str, Any], + metadata: dict[str, Any] | None, + *, + name: str, + model_name: str, + ) -> None: + self.runtime = runtime + self.session = session + self.logical = _logical_parent(runtime, session, parent, metadata) + self.parent = self.logical[1] if self.logical is not None else parent + self.request = request + self.metadata = metadata + self.body = _relay_request_body(request, metadata) + self.relay_request = runtime.relay.LLMRequest({}, self.body) + self.codec_baseline = _codec_round_trip_request_body( + runtime.relay, self.relay_request, relay_request_body=self.body, metadata=metadata + ) + self.operation = _relay_operation_name(name, metadata) + self.relay_kwargs = { + "handle": self.parent, + "metadata": _relay_metadata(name, metadata), + "model_name": model_name, + "codec": _codec(runtime.relay, metadata), + "response_codec": _codec(runtime.relay, metadata), + } + # Provider callback bookkeeping: "value"/"json" once it returned, "error" if it raised. + self.raw_response: dict[str, Any] = {} + self.context = contextvars.copy_context() + + def provider_request(self, next_request: Any) -> dict[str, Any]: + return _provider_request( + self.request, + next_request, + relay_request_body=self.body, + codec_baseline_body=self.codec_baseline, + metadata=self.metadata, + ) + + def run_callback(self, callback: Callable[..., Any], *args: Any) -> Any: + """Run a Hermes callback in a fresh copy of the captured context. + + Relay can invoke callbacks while another one still owns the captured + Context, hence the copy. Nested relay calls inside a managed provider + callback must run unmanaged — see relay_runtime.managed_callback_guard. + """ + def guarded() -> Any: + with relay_runtime.managed_callback_guard(): + return callback(*args) + + return self.context.copy().run(guarded) + + def record(self, raw: Any) -> Any: + self.raw_response["value"] = raw + self.raw_response["json"] = _jsonable(raw) + return self.raw_response["json"] + + def fail(self, exc: BaseException) -> None: + self.raw_response["error"] = exc + + def run_managed(self, relay_call: Callable[..., Any], *callbacks: Any) -> Any: + """Return the awaitable running ``relay_call`` inside the session context.""" + return self.runtime.run_in_session_async( + self.session, + relay_call, + self.operation, + self.relay_request, + *callbacks, + **self.relay_kwargs, + ) + + def resolve_failure(self, exc: BaseException, defer_logical_completion: bool) -> Any: + """Re-raise the provider's own error, or recover a completed provider result. + + Must be called from the ``except`` handling ``exc`` (bare ``raise``). + """ + callback_error = self.raw_response.get("error") + if ( + callback_error is not None + and relay_runtime._is_relay_wrapped_callback_error(exc, callback_error) + ): + raise callback_error + if ( + not isinstance(exc, Exception) + or callback_error is not None + or "value" not in self.raw_response + ): + raise + logger.warning( + "NeMo Relay LLM post-processing failed after provider success; " + "returning the provider response", + exc_info=True, + ) + if not defer_logical_completion: + _complete_logical(self.logical, outcome="success") + return self.raw_response["value"] + + def result(self, managed: Any, defer_logical_completion: bool) -> Any: + if not defer_logical_completion: + _complete_logical(self.logical, outcome="success") + if "value" in self.raw_response and _json_equal(managed, self.raw_response["json"]): + return self.raw_response["value"] + return _namespace(managed) + + def execute( request: dict[str, Any], callback: Callable[[dict[str, Any]], Any], @@ -81,85 +197,25 @@ def execute( defer_logical_completion: bool = False, ) -> Any: """Run one non-streaming physical provider attempt through Relay.""" - runtime, session, parent = relay_runtime.resolve_execution_context(session_id) - if runtime is None or session is None or not runtime.managed_execution_enabled(): - return callback(request) - logical = _logical_parent(runtime, session, parent, metadata) - parent = logical[1] if logical is not None else parent - - relay_request_body = _relay_request_body(request, metadata) - relay_request = runtime.relay.LLMRequest({}, relay_request_body) - codec_baseline_body = _codec_round_trip_request_body( - runtime.relay, - relay_request, - relay_request_body=relay_request_body, - metadata=metadata, + attempt = _ManagedAttempt.resolve( + session_id, request, metadata, name=name, model_name=model_name ) - raw_response: dict[str, Any] = {} - callback_error: BaseException | None = None - callback_context = contextvars.copy_context() + if attempt is None: + return callback(request) 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, - next_request, - relay_request_body=relay_request_body, - codec_baseline_body=codec_baseline_body, - metadata=metadata, - ) - raw = callback_context.copy().run(guarded, final_request) + raw = attempt.run_callback(callback, attempt.provider_request(next_request)) except BaseException as exc: - callback_error = exc + attempt.fail(exc) raise - raw_response["value"] = raw - raw_response["json"] = _jsonable(raw) - return raw_response["json"] + return attempt.record(raw) try: - managed = _run_awaitable( - runtime.run_in_session_async( - session, - runtime.relay.llm.execute, - _relay_operation_name(name, metadata), - relay_request, - invoke, - handle=parent, - metadata=_relay_metadata(name, metadata), - model_name=model_name, - codec=_codec(runtime.relay, metadata), - response_codec=_codec(runtime.relay, metadata), - ) - ) + managed = _run_awaitable(attempt.run_managed(attempt.runtime.relay.llm.execute, invoke)) except BaseException as exc: - if ( - callback_error is not None - and relay_runtime._is_relay_wrapped_callback_error(exc, callback_error) - ): - raise callback_error - if _recover_successful_callback( - raw_response, - relay_error=exc, - callback_error=callback_error, - logical=logical, - defer_logical_completion=defer_logical_completion, - ): - return raw_response["value"] - raise - - if not defer_logical_completion: - _complete_logical(logical, outcome="success") - if "value" in raw_response and _json_equal(managed, raw_response["json"]): - return raw_response["value"] - return _namespace(managed) + return attempt.resolve_failure(exc, defer_logical_completion) + return attempt.result(managed, defer_logical_completion) async def execute_async( @@ -173,86 +229,39 @@ async def execute_async( defer_logical_completion: bool = False, ) -> Any: """Run one asynchronous physical provider attempt through Relay.""" - runtime, session, parent = relay_runtime.resolve_execution_context(session_id) - if runtime is None or session is None or not runtime.managed_execution_enabled(): - return await callback(request) - logical = _logical_parent(runtime, session, parent, metadata) - parent = logical[1] if logical is not None else parent - - relay_request_body = _relay_request_body(request, metadata) - relay_request = runtime.relay.LLMRequest({}, relay_request_body) - codec_baseline_body = _codec_round_trip_request_body( - runtime.relay, - relay_request, - relay_request_body=relay_request_body, - metadata=metadata, + attempt = _ManagedAttempt.resolve( + session_id, request, metadata, name=name, model_name=model_name ) - raw_response: dict[str, Any] = {} - callback_error: BaseException | None = None - callback_context = contextvars.copy_context() + if attempt is None: + return await callback(request) async def invoke(next_request: Any) -> Any: - nonlocal callback_error try: - final_request = _provider_request( - request, - next_request, - relay_request_body=relay_request_body, - codec_baseline_body=codec_baseline_body, - metadata=metadata, - ) + final_request = attempt.provider_request(next_request) + async def call_provider() -> Any: # Nested relay calls inside a managed provider callback must - # run unmanaged (#77244). + # run unmanaged — see relay_runtime.managed_callback_guard. with relay_runtime.managed_callback_guard(): return await callback(final_request) - task = callback_context.copy().run( - asyncio.create_task, - call_provider(), - ) - raw = await task + raw = await attempt.context.copy().run(asyncio.create_task, call_provider()) except BaseException as exc: - callback_error = exc + attempt.fail(exc) raise - raw_response["value"] = raw - raw_response["json"] = _jsonable(raw) - return raw_response["json"] + return attempt.record(raw) try: - managed = await runtime.run_in_session_async( - session, - runtime.relay.llm.execute, - _relay_operation_name(name, metadata), - relay_request, - invoke, - handle=parent, - metadata=_relay_metadata(name, metadata), - model_name=model_name, - codec=_codec(runtime.relay, metadata), - response_codec=_codec(runtime.relay, metadata), - ) + managed = await attempt.run_managed(attempt.runtime.relay.llm.execute, invoke) except BaseException as exc: - if ( - callback_error is not None - and relay_runtime._is_relay_wrapped_callback_error(exc, callback_error) - ): - raise callback_error - if _recover_successful_callback( - raw_response, - relay_error=exc, - callback_error=callback_error, - logical=logical, - defer_logical_completion=defer_logical_completion, - ): - return raw_response["value"] - raise + return attempt.resolve_failure(exc, defer_logical_completion) + return attempt.result(managed, defer_logical_completion) - if not defer_logical_completion: - _complete_logical(logical, outcome="success") - if "value" in raw_response and _json_equal(managed, raw_response["json"]): - return raw_response["value"] - return _namespace(managed) + +def _current_session_id() -> str | None: + """Return the inherited Hermes turn's session id, or None outside a live turn.""" + turn = relay_runtime.active_turn() + return None if turn is None else turn.lease.session_id def execute_current( @@ -265,13 +274,13 @@ def execute_current( defer_logical_completion: bool = False, ) -> Any: """Run a provider attempt under the inherited Hermes turn when present.""" - turn = relay_runtime.active_turn() - if turn is None: + session_id = _current_session_id() + if session_id is None: return callback(request) return execute( request, callback, - session_id=turn.lease.session_id, + session_id=session_id, name=name, model_name=model_name, metadata=metadata, @@ -289,13 +298,13 @@ async def execute_current_async( defer_logical_completion: bool = False, ) -> Any: """Run an async provider attempt under the inherited turn when present.""" - turn = relay_runtime.active_turn() - if turn is None: + session_id = _current_session_id() + if session_id is None: return await callback(request) return await execute_async( request, callback, - session_id=turn.lease.session_id, + session_id=session_id, name=name, model_name=model_name, metadata=metadata, @@ -324,39 +333,27 @@ def stream_current( ) -> Any: """Run a provider stream under the inherited Hermes turn when present. - When ``completed_response_predicate`` is set and the stream_factory returns - a complete response instead of an iterator (e.g. AnthropicAuxiliaryClient - and other shims that ignore ``stream=True``), unwrap and return the - completed response directly. This mirrors the pre-Relay behavior where - ``call_llm(stream=True)`` returned the raw response and the consumer's - own ``hasattr(stream, "choices")`` check handled it (#11732, #55933) — - without the unwrap the response stays trapped as ``final_response`` on the - inner ManagedLlmStream and the outer consumer sees an empty stream. - - Determining that return shape requires starting the lazy managed pipeline, - and Relay may read ahead internally while satisfying that first pull. A - genuine first returned chunk remains buffered, while provider work, - latency, and pre-first-yield errors may surface before this function - returns. + With ``completed_response_predicate`` set, a factory that ignores + ``stream=True`` and returns a complete response is unwrapped and returned + directly (pre-Relay ``call_llm(stream=True)`` behavior); otherwise it + would stay trapped as ``final_response`` on the inner ManagedLlmStream. + Detecting that shape starts the lazy managed pipeline: a genuine first + chunk is buffered, but provider latency and pre-first-yield errors may + surface before this function returns. """ - turn = relay_runtime.active_turn() - if turn is None: + session_id = _current_session_id() + if session_id is None: return stream_factory(request) if _has_running_event_loop(): - # Managed provider callbacks execute on the Relay session's event - # loop. A nested ManagedLlmStream built here would be synchronously - # iterated on that same loop thread, which asyncio forbids - # ("Cannot run the event loop while another loop is running"). - # Return the raw factory result instead: the outer managed stream - # already provides Relay tracking for the enclosing attempt, and its - # own completed_response_predicate traps a completed response (e.g. - # the MoA facade's auxiliary ``call_llm(stream=True)`` returning a - # full response when an adapter ignores ``stream=True``). + # Managed provider callbacks run on the Relay session's event loop; a + # nested ManagedLlmStream would be iterated synchronously on that same + # loop thread, which asyncio forbids. The outer managed stream already + # tracks the enclosing attempt and traps a completed response itself. return stream_factory(request) managed = stream( request, stream_factory, - session_id=turn.lease.session_id, + session_id=session_id, name=name, model_name=model_name, finalizer=finalizer, @@ -365,9 +362,8 @@ def stream_current( completed_response_predicate=completed_response_predicate, ) if completed_response_predicate is not None: - # Relay may defer the provider callback until the first stream pull. - # Prime once so adapters that ignore stream=True can still return their - # completed response directly. A real first chunk is buffered. + # Relay may defer the provider callback until the first pull; prime + # once so a completed response surfaces. A real first chunk is buffered. managed._prime_completed_response() completed = getattr(managed, "final_response", None) if completed is not None: @@ -409,6 +405,18 @@ def stream( ) +def _aclose_on_loop(loop: asyncio.AbstractEventLoop, stream: Any) -> None: + """Await ``stream.aclose()`` on ``loop`` when the stream exposes one.""" + close = getattr(stream, "aclose", None) + if not callable(close): + return + + async def close_stream() -> None: + await close() + + loop.run_until_complete(close_stream()) + + class ManagedLlmStream(Iterator[Any]): """Drive Relay's async stream from Hermes's provider worker thread.""" @@ -437,16 +445,13 @@ class ManagedLlmStream(Iterator[Any]): self._runtime_lease: relay_runtime.RelayOperationLease | None = None self._close_error: BaseException | None = None self._callback_error: BaseException | None = None - self._logical: tuple[relay_runtime.RelayTurnContext, Any, str] | None = None + self._logical: _LogicalCall | None = None self._defer_logical_completion = defer_logical_completion - if str((metadata or {}).get("call_role") or "").startswith("auxiliary:"): - self._logical_model_name: str | None = model_name - self._logical_provider_name: str | None = name - self._logical_response_model_name: str | None = None - else: - self._logical_model_name = None - self._logical_provider_name = None - self._logical_response_model_name = None + # Only auxiliary calls report model/provider on their logical scope. + auxiliary = str((metadata or {}).get("call_role") or "").startswith("auxiliary:") + self._logical_model_name: str | None = model_name if auxiliary else None + self._logical_provider_name: str | None = name if auxiliary else None + self._logical_response_model_name: str | None = None self._on_chunk = on_chunk self._chunk_adapter = chunk_adapter or _namespace self._accept_chunk = accept_chunk @@ -455,70 +460,53 @@ class ManagedLlmStream(Iterator[Any]): self._raw_chunks: list[tuple[Any, Any]] = [] self._prefetched_chunks: list[Any] = [] self.output_modified = False - callback_context = contextvars.copy_context() - - 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. - 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 ( - runtime is None - or session is None - or not runtime.managed_execution_enabled() - ): - raw_stream = stream_factory(request) - if completed_response_predicate is not None and completed_response_predicate( - raw_stream - ): - self.final_response = raw_stream - self._stream = iter(()) - else: - self._raw_stream_resource = raw_stream - if on_stream_created is not None: - on_stream_created(raw_stream) - self._stream = iter(raw_stream) - return - - self._logical = _logical_parent(runtime, session, parent, metadata) - if self._logical is not None: - parent = self._logical[1] - relay_request_body = _relay_request_body(request, metadata) - relay_request = runtime.relay.LLMRequest({}, relay_request_body) - codec_baseline_body = _codec_round_trip_request_body( - runtime.relay, - relay_request, - relay_request_body=relay_request_body, - metadata=metadata, + attempt = _ManagedAttempt.resolve( + session_id, request, metadata, name=name, model_name=model_name ) + if attempt is None: + self._start_unmanaged( + request, stream_factory, on_stream_created, completed_response_predicate + ) + return + self._logical = attempt.logical + self._start_managed( + attempt, stream_factory, on_stream_created, completed_response_predicate, finalizer + ) + + def _start_unmanaged( + self, + request: dict[str, Any], + stream_factory: Callable[[dict[str, Any]], Any], + on_stream_created: Callable[[Any], None] | None, + completed_response_predicate: Callable[[Any], bool] | None, + ) -> None: + raw_stream = stream_factory(request) + if completed_response_predicate is not None and completed_response_predicate(raw_stream): + self.final_response = raw_stream + self._stream = iter(()) + return + self._raw_stream_resource = raw_stream + if on_stream_created is not None: + on_stream_created(raw_stream) + self._stream = iter(raw_stream) + + def _start_managed( + self, + attempt: _ManagedAttempt, + stream_factory: Callable[[dict[str, Any]], Any], + on_stream_created: Callable[[Any], None] | None, + completed_response_predicate: Callable[[Any], bool] | None, + finalizer: Callable[[], Any], + ) -> None: + """Open Relay's stream on a private event loop owned by this iterator.""" + run_callback = attempt.run_callback async def provider_stream(next_request: Any): raw_stream = None try: - raw_stream = run_callback( - stream_factory, - _provider_request( - request, - next_request, - relay_request_body=relay_request_body, - codec_baseline_body=codec_baseline_body, - metadata=metadata, - ) - ) - if ( - completed_response_predicate is not None - and run_callback( - completed_response_predicate, - raw_stream, - ) + raw_stream = run_callback(stream_factory, attempt.provider_request(next_request)) + if completed_response_predicate is not None and run_callback( + completed_response_predicate, raw_stream ): self.final_response = raw_stream self._provider_completed = True @@ -532,8 +520,7 @@ class ManagedLlmStream(Iterator[Any]): except StopIteration: break if self._accept_chunk is not None and not run_callback( - self._accept_chunk, - chunk, + self._accept_chunk, chunk ): break encoded_chunk = _jsonable(chunk) @@ -558,14 +545,13 @@ class ManagedLlmStream(Iterator[Any]): def relay_finalizer() -> Any: # Relay can invoke the finalizer while unwinding a provider-stream - # failure. Preserve that original callback error instead of - # replacing it with a secondary "missing terminal response" error. + # failure; keep that original error instead of a secondary + # "missing terminal response" error. if self._callback_error is not None: return None try: - if self.final_response is not None: - response = self.final_response - else: + response = self.final_response + if response is None: response = run_callback(finalizer) if self._logical_model_name is not None: self._logical_response_model_name = _response_model_name(response) @@ -574,7 +560,7 @@ class ManagedLlmStream(Iterator[Any]): self._callback_error = exc raise - self._runtime_lease = runtime.acquire_operation_lease() + self._runtime_lease = attempt.runtime.acquire_operation_lease() try: loop = asyncio.new_event_loop() except BaseException: @@ -584,44 +570,18 @@ class ManagedLlmStream(Iterator[Any]): self._relay_observes_chunks = True try: self._stream = loop.run_until_complete( - runtime.run_in_session_async( - session, - runtime.relay.llm.stream_execute, - _relay_operation_name(name, metadata), - relay_request, + attempt.run_managed( + attempt.runtime.relay.llm.stream_execute, provider_stream, observe_chunk, relay_finalizer, - handle=parent, - metadata=_relay_metadata(name, metadata), - model_name=model_name, - codec=_codec(runtime.relay, metadata), - response_codec=_codec(runtime.relay, metadata), ) ) except BaseException as exc: - if ( - isinstance(exc, Exception) - and self._provider_completed - and self._callback_error is None - ): - logger.warning( - "NeMo Relay stream post-processing failed after provider success; " - "preserving the provider result", - exc_info=True, - ) + if self._recoverable_relay_failure(exc): self._preserve_pending_provider_chunks() return - if not self._defer_logical_completion: - _complete_logical( - self._logical, - outcome="cancelled" if _is_cancellation(exc) else "failed", - model_name=self._logical_model_name, - provider_name=self._logical_provider_name, - response_model_name=self._logical_response_model_name, - operation_lease=self._runtime_lease, - ) - self._logical = None + self._finish_logical("cancelled" if _is_cancellation(exc) else "failed") try: loop.close() finally: @@ -641,6 +601,33 @@ class ManagedLlmStream(Iterator[Any]): except StopIteration: pass + def _recoverable_relay_failure(self, exc: BaseException) -> bool: + """Relay post-processing failed after the provider already succeeded.""" + if ( + isinstance(exc, Exception) and self._provider_completed and self._callback_error is None + ): + logger.warning( + "NeMo Relay stream post-processing failed after provider success; " + "preserving the provider result", + exc_info=True, + ) + return True + return False + + def _finish_logical(self, outcome: str) -> None: + """Complete the logical LLM scope unless the caller deferred it.""" + if self._defer_logical_completion: + return + _complete_logical( + self._logical, + outcome=outcome, + model_name=self._logical_model_name, + provider_name=self._logical_provider_name, + response_model_name=self._logical_response_model_name, + operation_lease=self._runtime_lease, + ) + self._logical = None + def __next__(self) -> Any: if self._closed: raise StopIteration @@ -665,16 +652,7 @@ class ManagedLlmStream(Iterator[Any]): except StopAsyncIteration: if self._raw_chunks: self.output_modified = True - if not self._defer_logical_completion: - _complete_logical( - self._logical, - outcome="success", - model_name=self._logical_model_name, - provider_name=self._logical_provider_name, - response_model_name=self._logical_response_model_name, - operation_lease=self._runtime_lease, - ) - self._logical = None + self._finish_logical("success") self._close(logical_outcome="cancelled") raise StopIteration from None except BaseException as exc: @@ -685,21 +663,10 @@ class ManagedLlmStream(Iterator[Any]): ): self._close(logical_outcome="failed") raise callback_error - if ( - isinstance(exc, Exception) - and self._provider_completed - and callback_error is None - ): - logger.warning( - "NeMo Relay stream post-processing failed after provider success; " - "preserving the provider result", - exc_info=True, - ) + if self._recoverable_relay_failure(exc): self._preserve_pending_provider_chunks() return next(self) - self._close( - logical_outcome="cancelled" if _is_cancellation(exc) else "failed" - ) + self._close(logical_outcome="cancelled" if _is_cancellation(exc) else "failed") raise if not self._relay_observes_chunks and self._on_chunk is not None: self._on_chunk(chunk) @@ -732,33 +699,37 @@ class ManagedLlmStream(Iterator[Any]): self._accept_chunk = None try: if loop is not None: - close = getattr(relay_stream, "aclose", None) - if callable(close): - - async def close_stream() -> None: - await close() - - try: - loop.run_until_complete(close_stream()) - except Exception: - logger.debug( - "Relay stream cleanup failed during provider fallback", - exc_info=True, - ) + try: + _aclose_on_loop(loop, relay_stream) + except Exception: + logger.debug( + "Relay stream cleanup failed during provider fallback", exc_info=True + ) loop.close() - if not self._defer_logical_completion: - _complete_logical( - self._logical, - outcome="success", - model_name=self._logical_model_name, - provider_name=self._logical_provider_name, - response_model_name=self._logical_response_model_name, - operation_lease=self._runtime_lease, - ) - self._logical = None + self._finish_logical("success") finally: self._release_runtime_lease() + def _close_provider_resources(self) -> None: + """Close the unmanaged provider stream/resource once each (they may be the same object).""" + resources = (self._stream, self._raw_stream_resource) + self._stream = None + self._raw_stream_resource = None + closed_ids: set[int] = set() + for resource in resources: + if resource is None or id(resource) in closed_ids: + continue + closed_ids.add(id(resource)) + close = getattr(resource, "close", None) + if not callable(close): + continue + try: + close() + except Exception as exc: + if self._close_error is None: + self._close_error = exc + logger.debug("Provider stream cleanup failed", exc_info=True) + def _close(self, *, logical_outcome: str) -> None: if self._closed: return @@ -768,57 +739,15 @@ class ManagedLlmStream(Iterator[Any]): loop = self._loop self._loop = None if loop is None: - resources = (self._stream, self._raw_stream_resource) - self._stream = None - self._raw_stream_resource = None - closed_ids: set[int] = set() - for resource in resources: - if resource is None or id(resource) in closed_ids: - continue - closed_ids.add(id(resource)) - close = getattr(resource, "close", None) - if callable(close): - try: - close() - except Exception as exc: - if self._close_error is None: - self._close_error = exc - logger.debug( - "Provider stream cleanup failed", - exc_info=True, - ) - if not self._defer_logical_completion: - _complete_logical( - self._logical, - outcome=logical_outcome, - model_name=self._logical_model_name, - provider_name=self._logical_provider_name, - response_model_name=self._logical_response_model_name, - operation_lease=self._runtime_lease, - ) - self._logical = None + self._close_provider_resources() + self._finish_logical(logical_outcome) return - close = getattr(self._stream, "aclose", None) - if callable(close): - - async def close_stream() -> None: - await close() - - try: - loop.run_until_complete(close_stream()) - except Exception as exc: - if self._close_error is None: - self._close_error = exc - if not self._defer_logical_completion: - _complete_logical( - self._logical, - outcome=logical_outcome, - model_name=self._logical_model_name, - provider_name=self._logical_provider_name, - response_model_name=self._logical_response_model_name, - operation_lease=self._runtime_lease, - ) - self._logical = None + try: + _aclose_on_loop(loop, self._stream) + except Exception as exc: + if self._close_error is None: + self._close_error = exc + self._finish_logical(logical_outcome) loop.close() finally: self._release_runtime_lease() @@ -833,6 +762,11 @@ class ManagedLlmStream(Iterator[Any]): self._close(logical_outcome="cancelled") +_ANTHROPIC_APPEND_DELTAS = { + "text_delta": "text", "thinking_delta": "thinking", "signature_delta": "signature" +} + + class AnthropicStreamAccumulator: """Rebuild an Anthropic Message from post-intercept SSE events.""" @@ -844,60 +778,59 @@ class AnthropicStreamAccumulator: payload = _jsonable(event) if not isinstance(payload, dict): return - event_type = payload.get("type") - if event_type == "message_start": - message = payload.get("message") - if isinstance(message, dict): - for key in ("id", "type", "role", "model", "usage"): - if key in message: - self._message[key] = message[key] + handler = self._EVENT_HANDLERS.get(payload.get("type")) + if handler is not None: + handler(self, payload) + + def _on_message_start(self, payload: dict[str, Any]) -> None: + message = payload.get("message") + if isinstance(message, dict): + for key in ("id", "type", "role", "model", "usage"): + if key in message: + self._message[key] = message[key] + + def _on_content_block_start(self, payload: dict[str, Any]) -> None: + index = payload.get("index") + block = payload.get("content_block") + if isinstance(index, int) and isinstance(block, dict): + self._blocks[index] = dict(block) + + def _on_content_block_delta(self, payload: dict[str, Any]) -> None: + index = payload.get("index") + delta = payload.get("delta") + if not isinstance(index, int) or not isinstance(delta, dict): return - if event_type == "content_block_start": - index = payload.get("index") - block = payload.get("content_block") - if isinstance(index, int) and isinstance(block, dict): - self._blocks[index] = dict(block) - return - if event_type == "content_block_delta": - index = payload.get("index") - delta = payload.get("delta") - if not isinstance(index, int) or not isinstance(delta, dict): - return - block = self._blocks.setdefault(index, {}) - delta_type = delta.get("type") - if delta_type == "text_delta": - block["text"] = str(block.get("text") or "") + str( - delta.get("text") or "" - ) - elif delta_type == "thinking_delta": - block["thinking"] = str(block.get("thinking") or "") + str( - delta.get("thinking") or "" - ) - elif delta_type == "signature_delta": - block["signature"] = str(block.get("signature") or "") + str( - delta.get("signature") or "" - ) - elif delta_type == "input_json_delta": - partial = str(block.pop("_partial_json", "")) + str( - delta.get("partial_json") or "" - ) - block["_partial_json"] = partial - elif delta_type == "citations_delta" and "citation" in delta: - block.setdefault("citations", []).append(delta["citation"]) - return - if event_type == "message_delta": - delta = payload.get("delta") - if isinstance(delta, dict): - for key in ("stop_reason", "stop_sequence"): - if key in delta: - self._message[key] = delta[key] - if "usage" in payload: - usage = payload["usage"] - current_usage = self._message.get("usage") - if isinstance(current_usage, dict) and isinstance(usage, dict): - self._message["usage"] = {**current_usage, **usage} - else: - self._message["usage"] = usage + block = self._blocks.setdefault(index, {}) + delta_type = delta.get("type") + field = _ANTHROPIC_APPEND_DELTAS.get(delta_type) + if field is not None: + block[field] = str(block.get(field) or "") + str(delta.get(field) or "") + elif delta_type == "input_json_delta": + block["_partial_json"] = str(block.pop("_partial_json", "")) + str( + delta.get("partial_json") or "" + ) + elif delta_type == "citations_delta" and "citation" in delta: + block.setdefault("citations", []).append(delta["citation"]) + + def _on_message_delta(self, payload: dict[str, Any]) -> None: + delta = payload.get("delta") + if isinstance(delta, dict): + for key in ("stop_reason", "stop_sequence"): + if key in delta: + self._message[key] = delta[key] + if "usage" in payload: + usage = payload["usage"] + current_usage = self._message.get("usage") + if isinstance(current_usage, dict) and isinstance(usage, dict): + usage = {**current_usage, **usage} + self._message["usage"] = usage + + _EVENT_HANDLERS = { + "message_start": _on_message_start, + "content_block_start": _on_content_block_start, + "content_block_delta": _on_content_block_delta, + "message_delta": _on_message_delta, + } def finalize(self) -> dict[str, Any]: blocks = [] @@ -915,9 +848,7 @@ class AnthropicStreamAccumulator: def response(self, base: Any = None) -> Any: """Return the attribute-shaped response consumed by Hermes.""" assembled = self.finalize() - base_payload = _jsonable(base) - if not isinstance(base_payload, dict): - base_payload = {} + base_payload = _jsonable_dict(base) content = assembled.pop("content", []) merged = {**base_payload, **assembled} if content or "content" not in merged: @@ -926,11 +857,9 @@ class AnthropicStreamAccumulator: def _logical_parent( - runtime: relay_runtime.RelayRuntime, - session: Any, - parent: Any, - metadata: dict[str, Any] | None, -) -> tuple[relay_runtime.RelayTurnContext, Any, str] | None: + runtime: relay_runtime.RelayRuntime, session: Any, parent: Any, metadata: dict[str, Any] | None +) -> _LogicalCall | None: + """Return (turn, handle, request_id) for the turn's logical LLM scope, pushing it once.""" turn = relay_runtime.active_turn(session.session_id) request_id = str((metadata or {}).get("api_request_id") or "") if turn is None or not request_id or turn.lease.host is not runtime: @@ -948,20 +877,17 @@ def _logical_parent( runtime.relay.ScopeType.Function, handle=parent, input={}, - metadata={ - relay_runtime.RUNTIME_SCHEMA_KEY: relay_runtime.RUNTIME_SCHEMA_VERSION, - relay_runtime.RUNTIME_INSTANCE_KEY: runtime.runtime_id, - "hermes.call_role": str( - (metadata or {}).get("call_role") or "primary" - ), - }, + metadata=relay_runtime.runtime_metadata( + runtime.runtime_id, + **{"hermes.call_role": str((metadata or {}).get("call_role") or "primary")}, + ), ) turn.logical_llm_calls[request_id] = handle return turn, handle, request_id def _complete_logical( - logical: tuple[relay_runtime.RelayTurnContext, Any, str] | None, + logical: _LogicalCall | None, *, outcome: str, model_name: str | None = None, @@ -996,53 +922,20 @@ def _complete_logical( lease.host.relay, handle, output=output, - metadata={ - relay_runtime.RUNTIME_SCHEMA_KEY: relay_runtime.RUNTIME_SCHEMA_VERSION, - relay_runtime.RUNTIME_INSTANCE_KEY: lease.host.runtime_id, - }, + metadata=relay_runtime.runtime_metadata(lease.host.runtime_id), ) except Exception: # The provider result is authoritative. Retain the handle so turn # finalization can retry cleanup without changing that result. - logger.warning( - "Hermes Relay logical LLM finalization failed", - exc_info=True, - ) + logger.warning("Hermes Relay logical LLM finalization failed", exc_info=True) return with turn.logical_llm_lock: if turn.logical_llm_calls.get(request_id) is handle: turn.logical_llm_calls.pop(request_id, None) -def _recover_successful_callback( - raw_response: dict[str, Any], - *, - relay_error: BaseException, - callback_error: BaseException | None, - logical: tuple[relay_runtime.RelayTurnContext, Any, str] | None, - defer_logical_completion: bool, -) -> bool: - if ( - not isinstance(relay_error, Exception) - or callback_error is not None - or "value" not in raw_response - ): - return False - logger.warning( - "NeMo Relay LLM post-processing failed after provider success; " - "returning the provider response", - exc_info=True, - ) - if not defer_logical_completion: - _complete_logical(logical, outcome="success") - return True - - def _is_cancellation(error: BaseException) -> bool: - return isinstance( - error, - (asyncio.CancelledError, InterruptedError, KeyboardInterrupt), - ) + return isinstance(error, (asyncio.CancelledError, InterruptedError, KeyboardInterrupt)) def complete_logical_call( @@ -1089,28 +982,20 @@ def _provider_request( content = getattr(request, "content", request) if not isinstance(content, dict): content = relay_request_body - if codec_baseline_body is None or _json_equal(content, relay_request_body): - final = dict(original) - else: + final = dict(original) + if codec_baseline_body is not None and not _json_equal(content, relay_request_body): baseline = codec_baseline_body intercepted = _provider_request_body(content, metadata) - final = dict(original) # Typed codecs may not represent provider-specific fields. Overlay only # values that changed from the codec-facing baseline so unrelated # intercepts cannot delete or normalize unknown provider arguments. for key in baseline.keys() | intercepted.keys(): if key not in intercepted: final.pop(key, None) - elif key not in baseline or not _json_equal( - intercepted[key], - baseline[key], - ): + elif key not in baseline or not _json_equal(intercepted[key], baseline[key]): final[key] = intercepted[key] _restore_provider_message_extensions( - original, - final, - baseline=baseline, - intercepted=intercepted, + original, final, baseline=baseline, intercepted=intercepted ) headers = getattr(request, "headers", None) if isinstance(headers, dict): @@ -1120,40 +1005,27 @@ def _provider_request( if str(key).lower() not in _RELAY_INTERNAL_PROVIDER_HEADERS } if headers: - final["extra_headers"] = { - **dict(final.get("extra_headers") or {}), - **headers, - } + final["extra_headers"] = {**dict(final.get("extra_headers") or {}), **headers} return final -def _relay_request_body( - request: dict[str, Any], metadata: dict[str, Any] | None -) -> dict[str, Any]: - body = _jsonable(request) - if not isinstance(body, dict): - return {} - # ``timeout`` configures the provider SDK client, not a provider wire - # protocol. Preserve it for the original callback request, but never pass - # it to Relay intercepts or routed transports. +def _relay_request_body(request: dict[str, Any], metadata: dict[str, Any] | None) -> dict[str, Any]: + body = _jsonable_dict(request) + # ``timeout`` configures the provider SDK client, not a wire protocol: + # keep it on the original callback request, never on Relay intercepts. body.pop("timeout", None) - # The Responses SDK accepts ``tools=None`` as "no tools", while Relay's - # typed Responses codec correctly expects either an array or an absent - # field. Normalize only the codec-facing copy; the original provider - # request is restored when no interceptor changes it. - if str((metadata or {}).get("api_mode") or "") == "codex_responses": - body = dict(body) + api_mode = _api_mode(metadata) + if api_mode == "codex_responses": + # The Responses SDK accepts ``tools=None`` as "no tools" while Relay's + # typed codec expects an array or an absent field; normalize only the + # codec-facing copy (the original request is restored when unchanged). if body.get("tools") is None: body.pop("tools", None) elif isinstance(body.get("tools"), list): body["tools"] = [ { "type": "function", - "function": { - key: value - for key, value in tool.items() - if key != "type" - }, + "function": {key: value for key, value in tool.items() if key != "type"}, } if isinstance(tool, dict) and tool.get("type") == "function" @@ -1161,15 +1033,12 @@ def _relay_request_body( else tool for tool in body["tools"] ] - elif str((metadata or {}).get("api_mode") or "") == "chat_completions": + elif api_mode == "chat_completions": tools = body.get("tools") if isinstance(tools, list): - body = dict(body) body["tools"] = [ {"type": "function", **tool} - if isinstance(tool, dict) - and "function" in tool - and "type" not in tool + if isinstance(tool, dict) and "function" in tool and "type" not in tool else tool for tool in tools ] @@ -1184,44 +1053,15 @@ def _restore_provider_message_extensions( intercepted: dict[str, Any], ) -> None: """Restore provider wire fields that Relay's typed codec cannot represent.""" - original_messages = original.get("messages") - final_messages = final.get("messages") - baseline_messages = baseline.get("messages") - intercepted_messages = intercepted.get("messages") - if not all( - isinstance(messages, list) - for messages in ( - original_messages, - final_messages, - baseline_messages, - intercepted_messages, - ) - ): + message_lists = tuple(body.get("messages") for body in (original, final, baseline, intercepted)) + if not all(isinstance(messages, list) for messages in message_lists): return - if not ( - len(original_messages) - == len(final_messages) - == len(baseline_messages) - == len(intercepted_messages) - ): + if len({len(messages) for messages in message_lists}) != 1: return - for original_message, final_message, baseline_message, intercepted_message in zip( - original_messages, - final_messages, - baseline_messages, - intercepted_messages, - strict=True, - ): - if not all( - isinstance(message, dict) - for message in ( - original_message, - final_message, - baseline_message, - intercepted_message, - ) - ): + for messages in zip(*message_lists, strict=True): + if not all(isinstance(message, dict) for message in messages): continue + original_message, final_message, baseline_message, intercepted_message = messages for key in _PROVIDER_MESSAGE_EXTENSION_KEYS: if ( key in original_message @@ -1251,13 +1091,11 @@ def _codec_round_trip_request_body( return _provider_request_body(content, metadata) except Exception: logger.warning( - "NeMo Relay request codec baseline failed; ignoring request rewrites", - exc_info=True, + "NeMo Relay request codec baseline failed; ignoring request rewrites", exc_info=True ) return None logger.warning( - "NeMo Relay request codec returned an unsupported baseline; " - "ignoring request rewrites" + "NeMo Relay request codec returned an unsupported baseline; ignoring request rewrites" ) return None @@ -1266,16 +1104,13 @@ def _provider_request_body( content: dict[str, Any], metadata: dict[str, Any] | None ) -> dict[str, Any]: body = dict(content) - if str((metadata or {}).get("api_mode") or "") != "codex_responses": + if _api_mode(metadata) != "codex_responses": return body tools = body.get("tools") if not isinstance(tools, list): return body body["tools"] = [ - { - "type": "function", - **dict(tool["function"]), - } + {"type": "function", **dict(tool["function"])} if isinstance(tool, dict) and tool.get("type") == "function" and isinstance(tool.get("function"), dict) @@ -1304,10 +1139,8 @@ def _jsonable(value: Any) -> Any: model_dump = getattr(type(value), "model_dump", None) if callable(model_dump): try: - # warnings=False: SDK stream events (e.g. the Anthropic - # ParsedMessage inside message_stop) carry generic-union content - # blocks that pydantic serializes fine but warns about — and the - # warning leaks to the user's terminal mid-response (#82xxx). + # warnings=False: pydantic warns on generic-union SDK stream events + # and that warning would leak to the user's terminal mid-response. try: return _jsonable(value.model_dump(mode="json", warnings=False)) except TypeError: @@ -1317,20 +1150,22 @@ def _jsonable(value: Any) -> Any: pass try: attributes = { - str(key): item - for key, item in vars(value).items() - if not str(key).startswith("_") + str(key): item for key, item in vars(value).items() if not str(key).startswith("_") } except (TypeError, AttributeError): return str(value) return _jsonable(attributes) if attributes else str(value) +def _jsonable_dict(value: Any) -> dict[str, Any]: + """``_jsonable`` for values that must be a JSON object; anything else becomes ``{}``.""" + payload = _jsonable(value) + return payload if isinstance(payload, dict) else {} + + def _namespace(value: Any) -> Any: if isinstance(value, dict): - return SimpleNamespace(**{ - str(key): _namespace(item) for key, item in value.items() - }) + return SimpleNamespace(**{str(key): _namespace(item) for key, item in value.items()}) if isinstance(value, list): return [_namespace(item) for item in value] return value @@ -1352,6 +1187,4 @@ def _run_awaitable(value: Any) -> Any: asyncio.get_running_loop() except RuntimeError: return asyncio.run(value) - raise RuntimeError( - "Synchronous Relay LLM execution cannot run on an event-loop thread" - ) + raise RuntimeError("Synchronous Relay LLM execution cannot run on an event-loop thread") diff --git a/agent/relay_runtime.py b/agent/relay_runtime.py index 3dbbd397a0..06db7a287f 100644 --- a/agent/relay_runtime.py +++ b/agent/relay_runtime.py @@ -20,8 +20,7 @@ from typing import Any, Callable from hermes_constants import get_hermes_home from hermes_cli.relay_plugin_cutover import ( - RELAY_PLUGINS_CONFIG_ENV, - configured_legacy_relay_env_vars, + RELAY_PLUGINS_CONFIG_ENV, configured_legacy_relay_env_vars ) logger = logging.getLogger(__name__) @@ -35,26 +34,25 @@ RUNTIME_INSTANCE_KEY = "hermes.relay.runtime_instance" RELAY_PLUGINS_EXECUTION_CONSUMER = "hermes.nemo_relay.plugins" _PROFILE_KEY_CACHE: dict[str, str] = {} -# Bound for native scope lifecycle operations (push/pop/flush) that gate -# turn/session completion. Healthy operations complete in microseconds; -# only a wedged native pipeline breaches this, and the correct trade there -# is one lost span, never a blocked agent (2026-08-10 delegation stall). +# Bound for native scope lifecycle ops (push/pop/flush) gating turn/session completion. +# Healthy ops take microseconds; a wedged pipeline costs one lost span, never a blocked agent. _SCOPE_OP_TIMEOUT = 10.0 _SCOPE_OP_EXECUTOR: Any = None _SCOPE_OP_EXECUTOR_LOCK = threading.Lock() +def runtime_metadata(runtime_id: str, **extra: Any) -> dict[str, Any]: + """Return the scope metadata that stamps every Hermes-owned Relay scope.""" + return {RUNTIME_SCHEMA_KEY: RUNTIME_SCHEMA_VERSION, RUNTIME_INSTANCE_KEY: runtime_id, **extra} + + def _scope_op_executor(): """Shared daemon executor for bounded native scope operations. - Daemon workers (tools.daemon_pool) so a wedged native call abandoned at - timeout cannot block interpreter exit. Sized generously: workers are - only consumed for the duration of healthy (microsecond) operations plus - any wedged calls, and ``Future.result(timeout=...)`` bounds callers even - when every worker is consumed by wedged calls — an unstarted future - still honors the result timeout, so exhaustion degrades to fast - timeouts, never a new hang. + Daemon workers (tools.daemon_pool) so a wedged call abandoned at timeout + cannot block interpreter exit. ``Future.result(timeout=...)`` bounds callers + even when every worker is wedged, so exhaustion degrades to fast timeouts. """ global _SCOPE_OP_EXECUTOR if _SCOPE_OP_EXECUTOR is None: @@ -63,21 +61,18 @@ def _scope_op_executor(): from tools.daemon_pool import DaemonThreadPoolExecutor _SCOPE_OP_EXECUTOR = DaemonThreadPoolExecutor( - max_workers=8, - thread_name_prefix="relay-scope-op", + max_workers=8, thread_name_prefix="relay-scope-op" ) return _SCOPE_OP_EXECUTOR -def _run_bounded_on_exit_thread(fn: Callable[[], Any], timeout: float) -> Any: - """Bounded fallback lane for interpreter shutdown. +def _run_on_daemon_thread( + fn: Callable[[], Any], *, name: str, timeout: float | None = None, timeout_message: str = "" +) -> Any: + """Run ``fn`` on a fresh daemon thread; re-raise its error or return its result. - When the shared executor refuses new futures (interpreter shutdown), - the operation still must not run unbounded on the calling thread: a - wedged native call would block process exit forever — the same defect - class this module exists to prevent, on the exit lane. Run it on a - fresh daemon thread with a bounded join; on breach the daemon worker - is abandoned exactly like the executor lane abandons its worker. + With ``timeout`` the join is bounded and a still-running worker is abandoned + with ``TimeoutError`` — a daemon thread cannot block interpreter exit. """ result: list[Any] = [] error: list[BaseException] = [] @@ -88,55 +83,68 @@ def _run_bounded_on_exit_thread(fn: Callable[[], Any], timeout: float) -> Any: except BaseException as exc: # noqa: BLE001 - propagated below error.append(exc) - worker = threading.Thread( - target=_target, daemon=True, name="relay-scope-op-exit" - ) + worker = threading.Thread(target=_target, daemon=True, name=name) worker.start() worker.join(timeout) if worker.is_alive(): - raise TimeoutError( - f"Relay scope operation exceeded {timeout}s during interpreter " - "shutdown; abandoning the native call so process exit can proceed" - ) + raise TimeoutError(timeout_message) if error: raise error[0] return result[0] if result else None def pop_relay_scope( - relay: Any, - handle: Any, - *, - output: Any = None, - metadata: Any = None, - timestamp: Any = None, + relay: Any, handle: Any, *, output: Any = None, metadata: Any = None, timestamp: Any = None ) -> Any: - """Pop a Relay scope without passing kwargs the binding rejects. + """Pop a Relay scope, forwarding only the kwargs the live binding accepts. - NeMo Relay ``scope.pop`` gained ``metadata`` in 0.4+. Older wheels (e.g. - 0.3.x) raise ``TypeError: pop() got an unexpected keyword argument - 'metadata'`` when Hermes finalization forwards runtime metadata. Filter to - parameters the live binding accepts so turn/session close can complete. + ``scope.pop`` gained ``metadata`` in nemo-relay 0.4+; older wheels raise + TypeError on it, which would wedge turn/session close. """ pop = relay.scope.pop - kwargs: dict[str, Any] = {} - if output is not None: - kwargs["output"] = output - if metadata is not None: - kwargs["metadata"] = metadata - if timestamp is not None: - kwargs["timestamp"] = timestamp + kwargs = { + key: value + for key, value in (("output", output), ("metadata", metadata), ("timestamp", timestamp)) + if value is not None + } try: params = inspect.signature(pop).parameters except (TypeError, ValueError): params = {} - if params and not any( - param.kind == inspect.Parameter.VAR_KEYWORD for param in params.values() - ): + if params and not any(param.kind == inspect.Parameter.VAR_KEYWORD for param in params.values()): kwargs = {key: value for key, value in kwargs.items() if key in params} return pop(handle, **kwargs) +def _current_top(relay: Any) -> Any: + """Return the current top-of-stack scope handle, or None.""" + # Prefer ``scope.get_handle()``: ``get_scope_stack()`` may return a native + # ScopeStack object that ``scope.pop`` rejects, so never treat it as a handle. + get_handle = getattr(getattr(relay, "scope", None), "get_handle", None) + if callable(get_handle): + try: + return get_handle() + except Exception: + pass + top = relay.get_scope_stack() + # Some builds return the live stack (list), others the top handle directly + # (including tuple handles from test fakes): only unwrap real lists. + if isinstance(top, list): + return top[-1] if top else None + return top + + +def _same_handle(a: Any, b: Any) -> bool: + # Native ScopeHandle has no value __eq__; compare by uuid when both expose one. + if a is None or b is None: + return a is b + if a is b or a == b: + return True + a_uuid = getattr(a, "uuid", None) + b_uuid = getattr(b, "uuid", None) + return a_uuid is not None and a_uuid == b_uuid + + class _RelayPluginConfigurationState(Enum): """Process-wide result shared by every currently hosted profile.""" @@ -161,59 +169,49 @@ class RelaySession: closing: bool = False handle: Any = None context: contextvars.Context | None = None - # --- session-span segmentation (continuous sessions) --- - # Segment index of the CURRENT session scope (0 = first). Rotation - # closes the current scope and pushes segment N+1 at a turn boundary. - segment: int = 0 - # Turns completed within the current segment (max_turns accounting). - segment_turns: int = 0 - # Set by compaction notification; consumed at the next begin_turn. - rotate_pending: bool = False - # Rotating compaction landed while a turn was live on THIS session: - # closing now would pop the session scope under a live turn scope - # (LIFO violation). end_turn consumes this and closes the session. + # Session-span segmentation (continuous sessions): rotation closes the + # current session scope and pushes segment N+1 at a turn boundary. + segment: int = 0 # index of the CURRENT session scope (0 = first) + segment_turns: int = 0 # turns completed within the current segment + rotate_pending: bool = False # set by compaction; consumed at next begin_turn + # A rotating compaction landed while a turn was live on THIS session; closing + # now would pop the session scope under the live turn (LIFO violation), so + # end_turn consumes this and closes the session. close_pending: bool = False -# --------------------------------------------------------------------------- -# Session-span segmentation config (gateway.telemetry.session_segments). -# Cached at first read; both defaults OFF => rotation never fires and the -# scope lifecycle is byte-identical to the pre-segmentation behavior. -# --------------------------------------------------------------------------- - +# Segmentation config (gateway.telemetry.session_segments), cached at first read. +# Both defaults OFF => rotation never fires and the scope lifecycle is identical +# to the pre-segmentation behavior. _SEGMENTS_CONFIG: dict[str, Any] | None = None _SEGMENTS_CONFIG_LOCK = threading.Lock() +def _load_segments_config() -> dict[str, Any]: + on_compaction = False + max_turns = 0 + try: + from gateway.run import _load_gateway_config # late import + + telemetry = (_load_gateway_config().get("gateway") or {}).get("telemetry") or {} + segments = telemetry.get("session_segments") or {} + on_compaction = bool(segments.get("on_compaction", False)) + try: + max_turns = max(0, int(segments.get("max_turns", 0) or 0)) + except (TypeError, ValueError): + max_turns = 0 + except Exception: # noqa: BLE001 - config absence must not crash + pass + return {"on_compaction": on_compaction, "max_turns": max_turns} + + def _segments_config() -> dict[str, Any]: """Resolve session-segmentation settings; inert defaults when unset.""" global _SEGMENTS_CONFIG if _SEGMENTS_CONFIG is None: with _SEGMENTS_CONFIG_LOCK: if _SEGMENTS_CONFIG is None: - on_compaction = False - max_turns = 0 - try: - from gateway.run import _load_gateway_config # late import - - telemetry = ( - (_load_gateway_config().get("gateway") or {}).get( - "telemetry" - ) - or {} - ) - segments = telemetry.get("session_segments") or {} - on_compaction = bool(segments.get("on_compaction", False)) - try: - max_turns = max(0, int(segments.get("max_turns", 0) or 0)) - except (TypeError, ValueError): - max_turns = 0 - except Exception: # noqa: BLE001 - config absence must not crash - pass - _SEGMENTS_CONFIG = { - "on_compaction": on_compaction, - "max_turns": max_turns, - } + _SEGMENTS_CONFIG = _load_segments_config() return _SEGMENTS_CONFIG @@ -230,23 +228,14 @@ class RelayOperationLease: self._runtime: RelayRuntime | None = runtime def run_in_session( - self, - session: RelaySession, - callback: Callable[..., Any], - *args: Any, - **kwargs: Any, + self, session: RelaySession, callback: Callable[..., Any], *args: Any, **kwargs: Any ) -> Any: """Run cleanup while this lease still owns the runtime lifetime.""" with self._lock: runtime = self._runtime if runtime is None: raise RuntimeError("Hermes Relay operation lease is released") - return runtime._run_in_session_untracked( - session, - callback, - *args, - **kwargs, - ) + return runtime._run_in_session_untracked(session, callback, *args, **kwargs) def release(self) -> None: """Release this lease exactly once.""" @@ -268,11 +257,7 @@ class _ProcessRelayPluginConfiguration: self._relay: Any = None self._activation: Any = None - def acquire( - self, - owner: Any, - relay: Any, - ) -> _RelayPluginConfigurationState: + def acquire(self, owner: Any, relay: Any) -> _RelayPluginConfigurationState: """Join the process configuration, initializing it for the first host.""" owner_id = id(owner) with self._lock: @@ -286,10 +271,7 @@ class _ProcessRelayPluginConfiguration: "Hermes Relay plugin cleanup is still pending; refusing to " "replace the process-global configuration" ) - return self._remember( - owner_id, - _RelayPluginConfigurationState.FAILED, - ) + return self._remember(owner_id, _RelayPluginConfigurationState.FAILED) try: existing_report = relay.plugin.report() @@ -299,87 +281,68 @@ class _ProcessRelayPluginConfiguration: "plugin configuration is already active; refusing to replace it", exc_info=True, ) - return self._remember( - owner_id, - _RelayPluginConfigurationState.FAILED, - ) + return self._remember(owner_id, _RelayPluginConfigurationState.FAILED) if existing_report is not None: logger.warning( "A process-global Relay plugin configuration is already active " "outside Hermes native ownership; leaving it unchanged and " "disabling Hermes-managed Relay middleware for this process" ) - return self._remember( - owner_id, - _RelayPluginConfigurationState.FOREIGN, - ) + return self._remember(owner_id, _RelayPluginConfigurationState.FOREIGN) try: - configured_inputs = _configured_plugin_inputs(relay) - if configured_inputs is None: - return self._remember( - owner_id, - _RelayPluginConfigurationState.DISABLED, - ) - plugin_config, dynamic_plugins = configured_inputs - if dynamic_plugins: - try: - activation = _resolve_plugin_awaitable( - relay.plugin.initialize_with_dynamic_plugins( - plugin_config, - dynamic_plugins, - ) - ) - if activation is None: - raise RuntimeError( - "NeMo Relay dynamic plugin initialization " - "returned no activation handle" - ) - self._activation = activation - except Exception as exc: - raise RuntimeError( - "Hermes Relay dynamic plugin activation failed" - ) from exc - - if self._activation is None: - # Hermes only enters Relay's initialization path after an - # explicit opt-in. Relay currently owns any subsequent ambient - # layering; a future discovery=False API can make this exact. - _resolve_plugin_awaitable(relay.plugin.initialize(plugin_config)) + if not self._initialize(relay): + return self._remember(owner_id, _RelayPluginConfigurationState.DISABLED) except Exception as exc: self._activation = None - logger.warning( - "Hermes Relay plugin initialization failed: %s", - exc, - exc_info=True, - ) - return self._remember( - owner_id, - _RelayPluginConfigurationState.FAILED, - ) + logger.warning("Hermes Relay plugin initialization failed: %s", exc, exc_info=True) + return self._remember(owner_id, _RelayPluginConfigurationState.FAILED) self._active = True self._relay = relay - state = self._remember( - owner_id, - _RelayPluginConfigurationState.ACTIVE, - ) + state = self._remember(owner_id, _RelayPluginConfigurationState.ACTIVE) logger.info( "Relay plugins are active process-wide and apply to all profiles " "hosted by this Hermes process." ) return state + def _initialize(self, relay: Any) -> bool: + """Initialize Relay from the selected plugins.toml; False when none is selected.""" + configured_inputs = _configured_plugin_inputs(relay) + if configured_inputs is None: + return False + plugin_config, dynamic_plugins = configured_inputs + if dynamic_plugins: + try: + activation = _resolve_plugin_awaitable( + relay.plugin.initialize_with_dynamic_plugins(plugin_config, dynamic_plugins) + ) + if activation is None: + raise RuntimeError( + "NeMo Relay dynamic plugin initialization returned no activation handle" + ) + self._activation = activation + except Exception as exc: + raise RuntimeError("Hermes Relay dynamic plugin activation failed") from exc + if self._activation is None: + # Hermes only enters Relay's initialization path after an + # explicit opt-in; Relay owns any subsequent ambient layering. + _resolve_plugin_awaitable(relay.plugin.initialize(plugin_config)) + return True + def _remember( - self, - owner_id: int, - state: _RelayPluginConfigurationState, + self, owner_id: int, state: _RelayPluginConfigurationState ) -> _RelayPluginConfigurationState: """Retain one process decision for all concurrently hosted profiles.""" self._owners.add(owner_id) self._state = state return state + def _reset_if_cleared(self) -> None: + if self._clear_active(): + self._state = _RelayPluginConfigurationState.UNINITIALIZED + def release(self, owner: Any) -> None: """Release one host and clear Relay after the final host exits.""" owner_id = id(owner) @@ -387,54 +350,41 @@ class _ProcessRelayPluginConfiguration: if owner_id not in self._owners: return self._owners.remove(owner_id) - if self._owners: - return - if self._clear_active(): - self._state = _RelayPluginConfigurationState.UNINITIALIZED + if not self._owners: + self._reset_if_cleared() def reset_for_tests(self) -> None: """Clear process-global state left by directly constructed test hosts.""" with self._lock: self._owners.clear() - if self._clear_active(): - self._state = _RelayPluginConfigurationState.UNINITIALIZED + self._reset_if_cleared() def retry_pending_cleanup(self) -> None: """Retry a failed final cleanup without disrupting live owners.""" with self._lock: if not self._owners: - if self._clear_active(): - self._state = _RelayPluginConfigurationState.UNINITIALIZED + self._reset_if_cleared() def _clear_active(self) -> bool: relay = self._relay activation = self._activation - active = self._active - if not active or relay is None: + if not self._active or relay is None: return True try: - _flush_relay_subscribers(relay) + _resolve_plugin_awaitable(relay.subscribers.flush_async()) except Exception: - logger.warning( - "Hermes Relay plugin subscriber flush failed", - exc_info=True, - ) + logger.warning("Hermes Relay plugin subscriber flush failed", exc_info=True) return False try: if activation is not None: close = getattr(activation, "close", None) if not callable(close): - raise RuntimeError( - "NeMo Relay dynamic plugin activation has no close method" - ) + raise RuntimeError("NeMo Relay dynamic plugin activation has no close method") _resolve_plugin_awaitable(close()) else: - _clear_relay_plugins(relay) + _resolve_plugin_awaitable(relay.plugin.clear_async()) except Exception: - logger.warning( - "Hermes Relay plugin configuration cleanup failed", - exc_info=True, - ) + logger.warning("Hermes Relay plugin configuration cleanup failed", exc_info=True) return False self._active = False self._relay = None @@ -465,19 +415,16 @@ class RelayRuntime: self._active_operations = 0 self._execution_consumers_lock = threading.RLock() self._execution_consumers: set[str] = set() - self._plugin_configuration_state = _PLUGIN_CONFIGURATION.acquire( - self, - self.relay, - ) + self._plugin_configuration_state = _PLUGIN_CONFIGURATION.acquire(self, self.relay) self._plugin_configuration_registered = True - if ( - self._plugin_configuration_state - is _RelayPluginConfigurationState.ACTIVE - ): + if self._plugins_active(): self.retain_managed_execution(RELAY_PLUGINS_EXECUTION_CONSUMER) self._shutdown_registered = True atexit.register(self.shutdown) + def _plugins_active(self) -> bool: + return self._plugin_configuration_state is _RelayPluginConfigurationState.ACTIVE + def retain_managed_execution(self, consumer: str) -> None: """Keep managed LLM and tool execution active for one consumer.""" if not consumer: @@ -495,12 +442,59 @@ class RelayRuntime: with self._execution_consumers_lock: return bool(self._execution_consumers) - def ensure_session( + def _subagent_parent_handle(self, session: RelaySession) -> Any: + with self._sessions_lock: + return self._subagent_parent_handles.get(session.session_id) + + def _push_session_scope( + self, context: contextvars.Context, *, exit_fallback: bool = False, **push_kwargs: Any + ) -> Any: + """Push a SESSION_SCOPE Agent scope inside ``context``, bounded by ``_SCOPE_OP_TIMEOUT``. + + ``exit_fallback``: at interpreter shutdown the executor refuses new futures + (RuntimeError); push synchronously instead, since no agent turn waits at exit. + """ + args = (self.relay.scope.push, SESSION_SCOPE, self.relay.ScopeType.Agent) + try: + return ( + _scope_op_executor() + .submit(context.run, *args, input={}, **push_kwargs) + .result(timeout=_SCOPE_OP_TIMEOUT) + ) + except RuntimeError: + if not exit_fallback: + raise + return context.run(*args, input={}, **push_kwargs) + + def _open_session_scope( self, - event: dict[str, Any], + session: RelaySession, + scope_metadata: dict[str, Any], *, - data: Any = None, - metadata: dict[str, Any] | None = None, + resolve_parent: bool, + **push_kwargs: Any, + ) -> None: + """Push a fresh session scope for ``session`` and record its handle + context. + + Subagent sessions parent under their spawning turn/session handle; + ``resolve_parent`` creates the parent session when its handle is unknown. + """ + parent_handle = None + if session.parent_session_id: + parent_handle = self._subagent_parent_handle(session) + if parent_handle is None and resolve_parent: + parent = self.ensure_session({"session_id": session.parent_session_id}) + if parent is not None: + parent_handle = parent.handle + scope_metadata["nemo_relay_scope_role"] = "subagent" + context = contextvars.Context() + session.handle = self._push_session_scope( + context, handle=parent_handle, metadata=scope_metadata, **push_kwargs + ) + session.context = context + + def ensure_session( + self, event: dict[str, Any], *, data: Any = None, metadata: dict[str, Any] | None = None ) -> RelaySession | None: """Return the existing session scope or create it once.""" session_id = _session_id(event) @@ -511,72 +505,36 @@ class RelayRuntime: return None session = self._sessions.get(session_id) if session is None: - parent_session_id = self._subagent_parents.get(session_id, "") session = RelaySession( session_id=session_id, - parent_session_id=parent_session_id, + parent_session_id=self._subagent_parents.get(session_id, ""), ) self._sessions[session_id] = session with session.lock: if session.closing: return None if session.handle is None: - parent_handle = None - scope_metadata = { - **(metadata or {}), - RUNTIME_SCHEMA_KEY: RUNTIME_SCHEMA_VERSION, - RUNTIME_INSTANCE_KEY: self.runtime_id, - } - if session.parent_session_id: - with self._sessions_lock: - parent_handle = self._subagent_parent_handles.get(session_id) - if parent_handle is None: - parent = self.ensure_session({ - "session_id": session.parent_session_id - }) - if parent is not None: - parent_handle = parent.handle - scope_metadata["nemo_relay_scope_role"] = "subagent" - context = contextvars.Context() try: - try: - session.handle = _scope_op_executor().submit( - context.run, - self.relay.scope.push, - SESSION_SCOPE, - self.relay.ScopeType.Agent, - handle=parent_handle, - data=data, - input={}, - metadata=scope_metadata, - ).result(timeout=_SCOPE_OP_TIMEOUT) - except RuntimeError: - # Interpreter shutdown: executor refuses new futures; - # push synchronously (no agent turn waits at exit). - session.handle = context.run( - self.relay.scope.push, - SESSION_SCOPE, - self.relay.ScopeType.Agent, - handle=parent_handle, - data=data, - input={}, - metadata=scope_metadata, - ) + self._open_session_scope( + session, + {**(metadata or {}), **runtime_metadata(self.runtime_id)}, + resolve_parent=True, + data=data, + exit_fallback=True, + ) except Exception: session.context = None raise - session.context = context return session def rotate_session_scope(self, session: RelaySession, *, reason: str) -> None: """Close the current session scope and open the next segment. - Called ONLY at a turn boundary (before the turn scope pushes), never - mid-turn: the scope stack is LIFO and rotating under a live child - would close a parent out of order. Both native calls are bounded by - ``_SCOPE_OP_TIMEOUT`` — a wedged rotation costs one segment span, - never the agent. Segment bookkeeping advances even when a native - call fails, so a degraded rotation cannot retry on every turn. + Called ONLY at a turn boundary (before the turn scope pushes): the scope + stack is LIFO and rotating under a live child would close a parent out + of order. Both native calls are bounded by ``_SCOPE_OP_TIMEOUT``, and + segment bookkeeping advances even when a native call fails so a degraded + rotation cannot retry on every turn. """ with session.lock: if session.closing or session.handle is None: @@ -593,10 +551,7 @@ class RelayRuntime: self.relay.scope.pop, old_handle, output={"hermes.session.segment_reason": reason}, - metadata={ - RUNTIME_SCHEMA_KEY: RUNTIME_SCHEMA_VERSION, - RUNTIME_INSTANCE_KEY: self.runtime_id, - }, + metadata=runtime_metadata(self.runtime_id), timeout=_SCOPE_OP_TIMEOUT, ) except Exception: @@ -607,31 +562,15 @@ class RelayRuntime: session.segment - 1, exc_info=True, ) - scope_metadata = { - RUNTIME_SCHEMA_KEY: RUNTIME_SCHEMA_VERSION, - RUNTIME_INSTANCE_KEY: self.runtime_id, - "hermes.session.segment": session.segment, - "hermes.session.segment_reason": reason, - } - parent_handle = None - if session.parent_session_id: - with self._sessions_lock: - parent_handle = self._subagent_parent_handles.get( - session.session_id - ) - scope_metadata["nemo_relay_scope_role"] = "subagent" - context = contextvars.Context() + scope_metadata = runtime_metadata( + self.runtime_id, + **{ + "hermes.session.segment": session.segment, + "hermes.session.segment_reason": reason, + }, + ) try: - session.handle = _scope_op_executor().submit( - context.run, - self.relay.scope.push, - SESSION_SCOPE, - self.relay.ScopeType.Agent, - handle=parent_handle, - input={}, - metadata=scope_metadata, - ).result(timeout=_SCOPE_OP_TIMEOUT) - session.context = context + self._open_session_scope(session, scope_metadata, resolve_parent=False) except Exception: logger.warning( "Hermes Relay segment open failed (session=%s segment=%d); " @@ -642,19 +581,12 @@ class RelayRuntime: ) def register_subagent( - self, - event: dict[str, Any], - *, - metadata: dict[str, Any] | None = None, + self, event: dict[str, Any], *, metadata: dict[str, Any] | None = None ) -> RelaySession | None: """Open a child Agent scope under its spawning turn when available.""" parent_session_id = str(event.get("parent_session_id") or "") child_session_id = str(event.get("child_session_id") or "") - if ( - not parent_session_id - or not child_session_id - or parent_session_id == child_session_id - ): + if not parent_session_id or not child_session_id or parent_session_id == child_session_id: return None parent = self.ensure_session({"session_id": parent_session_id}) parent_handle = None if parent is None else parent.handle @@ -674,10 +606,7 @@ class RelayRuntime: self._subagent_parents[child_session_id] = parent_session_id if parent_handle is not None: self._subagent_parent_handles[child_session_id] = parent_handle - return self.ensure_session( - {"session_id": child_session_id}, - metadata=metadata, - ) + return self.ensure_session({"session_id": child_session_id}, metadata=metadata) def unregister_subagent(self, event: dict[str, Any]) -> None: """Close a delegated session and forget its parent relationship.""" @@ -685,25 +614,38 @@ class RelayRuntime: if not child_session_id: return self.close_session({"session_id": child_session_id}) + self._forget_subagent(child_session_id) + + def _forget_subagent(self, session_id: str) -> None: with self._sessions_lock: - self._subagent_parents.pop(child_session_id, None) - self._subagent_parent_handles.pop(child_session_id, None) + self._subagent_parents.pop(session_id, None) + self._subagent_parent_handles.pop(session_id, None) def get_session(self, session_id: str) -> RelaySession | None: """Return an active Hermes Relay session without creating one.""" with self._sessions_lock: - if self._closing: - return None - session = self._sessions.get(str(session_id or "")) + session = None if self._closing else self._sessions.get(str(session_id or "")) if session is None: return None with session.lock: return None if session.closing else session - def get_session_handle(self, session_id: str) -> Any: - """Return the Relay parent handle for a Hermes session, if active.""" - session = self.get_session(session_id) - return None if session is None else session.handle + def _session_context( + self, session: RelaySession, *, allow_closing: bool + ) -> contextvars.Context: + """Copy the current context and overlay the session's saved Relay vars.""" + with session.lock: + if session.closing and not allow_closing: + raise RuntimeError("Hermes Relay session is closing") + if session.context is None or session.handle is None: + raise RuntimeError("Hermes Relay session context is unavailable") + relay_context = session.context.copy() + # A copy permits a helper called by an existing Relay callback to + # re-enter the same logical session without re-entering Context. + context = contextvars.copy_context() + for variable, value in relay_context.items(): + context.run(variable.set, value) + return context def run_in_session( self, @@ -716,26 +658,17 @@ class RelayRuntime: ) -> Any: """Run a Relay operation against a session's isolated scope stack. - ``timeout`` (seconds) bounds the native call by running it on a - shared daemon executor; ``TimeoutError`` propagates to the caller's - existing exception handling on breach. ``None`` (default) preserves - the historical synchronous behavior. Scope lifecycle operations - that gate turn/session completion pass ``_SCOPE_OP_TIMEOUT``: the - native binding's ``scope.pop`` "returns after the scope is closed - successfully" — unbounded — and a wedged native pipeline (proven - live 2026-08-10 in the delegation topology) must cost at most one - span, never the agent. The abandoned daemon worker cannot block - process exit (tools.daemon_pool contract). + ``timeout`` (seconds) bounds the native call on a shared daemon + executor; ``TimeoutError`` propagates on breach. ``None`` keeps the + synchronous behavior. Scope lifecycle ops that gate turn/session + completion pass ``_SCOPE_OP_TIMEOUT``: the native ``scope.pop`` is + unbounded, and a wedged pipeline must cost at most one span, never the + agent (the abandoned daemon worker cannot block process exit). """ self._begin_operation() try: return self._run_in_session_untracked( - session, - callback, - *args, - allow_closing=allow_closing, - timeout=timeout, - **kwargs, + session, callback, *args, allow_closing=allow_closing, timeout=timeout, **kwargs ) finally: self._end_operation() @@ -750,35 +683,28 @@ class RelayRuntime: **kwargs: Any, ) -> Any: """Run inside a session whose host-level lifetime is already held.""" - with session.lock: - if session.closing and not allow_closing: - raise RuntimeError("Hermes Relay session is closing") - if session.context is None or session.handle is None: - raise RuntimeError("Hermes Relay session context is unavailable") - relay_context = session.context.copy() - - context = contextvars.copy_context() - for variable, value in relay_context.items(): - context.run(variable.set, value) + context = self._session_context(session, allow_closing=allow_closing) def invoke() -> Any: self.relay.get_scope_stack() return callback(*args, **kwargs) - # A copy permits a helper called by an existing Relay callback to - # re-enter the same logical session without re-entering Context. if timeout is None: return context.run(invoke) try: future = _scope_op_executor().submit(context.run, invoke) except RuntimeError: - # Interpreter shutdown: the executor refuses new futures, but - # the atexit close path must still flush cleanly. Still - # bounded — a wedged native call must not block process exit - # (the CI runner hang, 2026-08-12: 6 tests passed in 4s, then - # the file-timeout SIGKILL'd a process stuck in this lane). - return _run_bounded_on_exit_thread( - lambda: context.run(invoke), timeout + # Interpreter shutdown: the executor refuses new futures, but the + # atexit close path must still flush — and still bounded, since a + # wedged native call must not block process exit. + return _run_on_daemon_thread( + lambda: context.run(invoke), + name="relay-scope-op-exit", + timeout=timeout, + timeout_message=( + f"Relay scope operation exceeded {timeout}s during interpreter " + "shutdown; abandoning the native call so process exit can proceed" + ), ) try: return future.result(timeout=timeout) @@ -800,16 +726,7 @@ class RelayRuntime: """Create and await an operation inside the session's saved context.""" self._begin_operation() try: - with session.lock: - if session.closing and not allow_closing: - raise RuntimeError("Hermes Relay session is closing") - if session.context is None or session.handle is None: - raise RuntimeError("Hermes Relay session context is unavailable") - relay_context = session.context.copy() - - context = contextvars.copy_context() - for variable, value in relay_context.items(): - context.run(variable.set, value) + context = self._session_context(session, allow_closing=allow_closing) async def invoke() -> Any: self.relay.get_scope_stack() @@ -843,12 +760,7 @@ class RelayRuntime: return RelayOperationLease(self) def emit_mark( - self, - name: str, - event: dict[str, Any], - *, - data: Any = None, - metadata: Any = None, + self, name: str, event: dict[str, Any], *, data: Any = None, metadata: Any = None ) -> bool: """Emit a mark parented to the Hermes session identified by ``event``.""" session = self.ensure_session(event) @@ -865,33 +777,72 @@ class RelayRuntime: return True def apply_tool_request_intercepts( - self, - *, - session_id: str, - tool_name: str, - args: dict[str, Any], + self, *, session_id: str, tool_name: str, args: dict[str, Any] ) -> dict[str, Any]: """Apply Relay request rewriting before Hermes authorizes a tool call.""" if not self.managed_execution_enabled(): return args - request_intercepts = getattr( - getattr(self.relay, "tools", None), - "request_intercepts", - None, - ) + request_intercepts = getattr(getattr(self.relay, "tools", None), "request_intercepts", None) if not callable(request_intercepts): return args session = self.ensure_session({"session_id": session_id}) if session is None: return args - result = self.run_in_session( - session, - request_intercepts, - tool_name, - args, - ) + result = self.run_in_session(session, request_intercepts, tool_name, args) return result if isinstance(result, dict) else args + def _pop_with_drain( + self, + handle: Any, + *, + output: dict[str, Any], + metadata: dict[str, Any], + session_root: Any, + drain_limit: int, + ) -> BaseException | None: + """Pop ``handle``; if that fails, drain orphans above it and retry once. + + Returns the retry's error (None on success). Must run inside ONE + ``run_in_session`` callback so ContextVar stack views stay consistent. + """ + try: + pop_relay_scope(self.relay, handle, output=output, metadata=metadata) + return None + except Exception: + pass + drained = 0 + for _ in range(drain_limit): + top = _current_top(self.relay) + if top is None or _same_handle(top, handle): + break + # Never pop the session root while draining for a nested handle. + if ( + session_root is not None + and _same_handle(top, session_root) + and handle is not session_root + ): + break + try: + pop_relay_scope( + self.relay, + top, + output={"outcome": "cancelled", "hermes.orphan_drain": True}, + metadata=metadata, + ) + drained += 1 + except Exception: + logger.warning("Hermes Relay orphaned scope drain failed", exc_info=True) + break + if drained: + logger.warning( + "Hermes Relay drained %d orphaned scope(s) before closing %s", drained, handle + ) + try: + pop_relay_scope(self.relay, handle, output=output, metadata=metadata) + return None + except Exception as retry_exc: + return retry_exc + def _close_scope_handle( self, session: RelaySession, @@ -905,140 +856,31 @@ class RelayRuntime: ) -> str | None: """Pop ``handle``, draining orphaned children in the same session context. - Relay scopes are strict LIFO. Empty-stream retries + interrupt can - abandon a physical LLM scope above TURN/SESSION (#81521). Drain and - close must run inside one ``run_in_session`` callback so ContextVar - stack views stay consistent across pops. + Relay scopes are strict LIFO; empty-stream retries + interrupt can + abandon a physical LLM scope above TURN/SESSION. The whole drain+close + is bounded like the direct pops it replaced: a wedged native pipeline + must never block turn/session completion. Returns a failure string. """ if handle is None: return None - metadata = { - RUNTIME_SCHEMA_KEY: RUNTIME_SCHEMA_VERSION, - RUNTIME_INSTANCE_KEY: self.runtime_id, - } - close_output = output or {} - session_root = session.handle - drained_holder = {"count": 0} - error_holder: dict[str, BaseException] = {} - - def close_with_drain() -> None: - def current_top() -> Any: - # Version-correct accessor first: the pinned nemo-relay - # binding exposes ``scope.get_handle()`` returning the - # current top-of-stack ScopeHandle. Its - # ``get_scope_stack()`` returns a native ScopeStack object - # that ``scope.pop`` rejects with TypeError, so it must - # never be treated as a handle (#81601 review). - get_handle = getattr( - getattr(self.relay, "scope", None), "get_handle", None - ) - if callable(get_handle): - try: - return get_handle() - except Exception: - pass - top = self.relay.get_scope_stack() - # Some Relay builds return the live stack (list). Others - # return the top handle directly — including tuple handles - # like ("scope", name, serial) from the test fake. Only - # unwrap real list stacks; never index a handle tuple. - if isinstance(top, list): - return top[-1] if top else None - return top - - def same_handle(a: Any, b: Any) -> bool: - # Native ScopeHandle instances do not implement __eq__ by - # value — two handles for the same scope compare unequal — - # so compare by uuid when both sides expose one. - if a is None or b is None: - return a is b - if a is b or a == b: - return True - a_uuid = getattr(a, "uuid", None) - b_uuid = getattr(b, "uuid", None) - return a_uuid is not None and a_uuid == b_uuid - - try: - pop_relay_scope( - self.relay, - handle, - output=close_output, - metadata=metadata, - ) - return - except Exception as first_exc: - error_holder["first"] = first_exc - - for _ in range(drain_limit): - top = current_top() - if top is None or same_handle(top, handle): - break - # Never pop the session root while draining for a nested handle. - if ( - session_root is not None - and same_handle(top, session_root) - and handle is not session_root - ): - break - try: - pop_relay_scope( - self.relay, - top, - output={ - "outcome": "cancelled", - "hermes.orphan_drain": True, - }, - metadata=metadata, - ) - drained_holder["count"] += 1 - except Exception as drain_exc: - error_holder["drain"] = drain_exc - logger.warning( - "Hermes Relay orphaned scope drain failed", - exc_info=True, - ) - break - - if drained_holder["count"]: - logger.warning( - "Hermes Relay drained %d orphaned scope(s) before closing %s", - drained_holder["count"], - handle, - ) - try: - pop_relay_scope( - self.relay, - handle, - output=close_output, - metadata=metadata, - ) - error_holder.pop("first", None) - error_holder.pop("drain", None) - except Exception as retry_exc: - error_holder["retry"] = retry_exc - + run_in_session = ( + self._run_in_session_untracked if operation_already_held else self.run_in_session + ) try: - run_in_session = ( - self._run_in_session_untracked - if operation_already_held - else self.run_in_session - ) - run_in_session( + failure = run_in_session( session, - close_with_drain, + self._pop_with_drain, + handle, + output=output or {}, + metadata=runtime_metadata(self.runtime_id), + session_root=session.handle, + drain_limit=drain_limit, allow_closing=allow_closing, - # Bound the whole drain+close like the direct pops it - # replaced: a wedged native pipeline must cost at most one - # span, never block turn/session completion (see - # tests/agent/test_relay_runtime_bounded_scope_ops.py). timeout=_SCOPE_OP_TIMEOUT, ) except Exception as exc: return f"{failure_label}: {exc}" - retry_exc = error_holder.get("retry") or error_holder.get("first") - if retry_exc is not None: - return f"{failure_label}: {retry_exc}" - return None + return None if failure is None else f"{failure_label}: {failure}" def close_session(self, event: dict[str, Any]) -> None: """Close one session scope and remove it from the core registry.""" @@ -1057,11 +899,9 @@ class RelayRuntime: with self._sessions_lock: session = self._sessions.get(session_id) if session is None: - with self._sessions_lock: - self._subagent_parents.pop(session_id, None) - self._subagent_parent_handles.pop(session_id, None) + self._forget_subagent(session_id) return - failures: list[str] = [] + failure = None with session.lock: if session.closing: return @@ -1075,22 +915,15 @@ class RelayRuntime: failure_label="session scope close failed", operation_already_held=True, ) - if failure: - failures.append(failure) # Subscriber flushing is process-wide and may wait for publications - # owned by other sessions. Final plugin teardown flushes once after all - # tracked operations drain; doing it here can deadlock an asyncio loop. + # owned by other sessions; final plugin teardown flushes once after all + # tracked operations drain. Flushing here can deadlock an asyncio loop. with self._sessions_lock: if self._sessions.get(session_id) is session: self._sessions.pop(session_id, None) - self._subagent_parents.pop(session_id, None) - self._subagent_parent_handles.pop(session_id, None) - if failures: - logger.warning( - "Hermes Relay session %s closed with errors: %s", - session_id, - "; ".join(failures), - ) + self._forget_subagent(session_id) + if failure: + logger.warning("Hermes Relay session %s closed with errors: %s", session_id, failure) def shutdown(self) -> None: """Close core scopes and release process plugin configuration.""" @@ -1111,10 +944,7 @@ class RelayRuntime: except Exception: with self._sessions_lock: self._shutdown_started = False - logger.warning( - "Hermes Relay deferred shutdown could not start", - exc_info=True, - ) + logger.warning("Hermes Relay deferred shutdown could not start", exc_info=True) return self._finish_shutdown() @@ -1129,20 +959,12 @@ class RelayRuntime: for session_id in session_ids: self._safe(self._close_session, {"session_id": session_id}) if self._plugin_configuration_registered: - if ( - self._plugin_configuration_state - is _RelayPluginConfigurationState.ACTIVE - ): - self.release_managed_execution( - RELAY_PLUGINS_EXECUTION_CONSUMER - ) + if self._plugins_active(): + self.release_managed_execution(RELAY_PLUGINS_EXECUTION_CONSUMER) _PLUGIN_CONFIGURATION.release(self) self._plugin_configuration_registered = False if self._shutdown_registered: - try: - atexit.unregister(self.shutdown) - except Exception: - pass + self._safe(atexit.unregister, self.shutdown, quiet=True) self._shutdown_registered = False except Exception: with self._sessions_lock: @@ -1153,11 +975,12 @@ class RelayRuntime: self._shutdown_complete.set() @staticmethod - def _safe(callback: Callable[..., Any], *args: Any, **kwargs: Any) -> Any: + def _safe(callback: Callable[..., Any], *args: Any, quiet: bool = False, **kwargs: Any) -> Any: try: return callback(*args, **kwargs) except Exception: - logger.warning("Hermes Relay runtime operation failed", exc_info=True) + if not quiet: + logger.warning("Hermes Relay runtime operation failed", exc_info=True) return None @@ -1168,16 +991,8 @@ class NoopRelayRuntime: profile_key: str reason: str - @property - def available(self) -> bool: - return False - def apply_tool_request_intercepts( - self, - *, - session_id: str, - tool_name: str, - args: dict[str, Any], + self, *, session_id: str, tool_name: str, args: dict[str, Any] ) -> dict[str, Any]: del session_id, tool_name return args @@ -1186,9 +1001,7 @@ class NoopRelayRuntime: def retain_managed_execution(consumer: str) -> None: del consumer - @staticmethod - def release_managed_execution(consumer: str) -> None: - del consumer + release_managed_execution = retain_managed_execution @staticmethod def managed_execution_enabled() -> bool: @@ -1209,10 +1022,7 @@ class RelayHostRegistry: self._hosts: dict[str, RelayHost] = {} def for_profile( - self, - profile_key: str | None = None, - *, - create: bool = True, + self, profile_key: str | None = None, *, create: bool = True ) -> RelayHost | None: key = profile_key or current_profile_key() host = self._hosts.get(key) @@ -1225,19 +1035,11 @@ class RelayHostRegistry: try: host = RelayRuntime(profile_key=key) except Exception as exc: - logger.warning( - "Hermes Relay runtime initialization failed", exc_info=True - ) + logger.warning("Hermes Relay runtime initialization failed", exc_info=True) host = NoopRelayRuntime(profile_key=key, reason=str(exc)) self._hosts[key] = host return host - def shutdown_profile(self, profile_key: str) -> None: - with self._lock: - host = self._hosts.pop(profile_key, None) - if host is not None: - host.shutdown() - def shutdown_all(self) -> None: with self._lock: hosts = list(self._hosts.values()) @@ -1261,6 +1063,12 @@ class ConversationLease: parent_session_id: str = "" released: bool = False + def live_runtime(self) -> RelayRuntime | None: + """Return the real Relay host when this lease owns an open session.""" + if isinstance(self.host, RelayRuntime) and self.session is not None: + return self.host + return None + @dataclass class RelayTurnContext: @@ -1271,14 +1079,8 @@ class RelayTurnContext: task_id: str handle: Any = None logical_llm_calls: dict[str, Any] = field(default_factory=dict, repr=False) - logical_llm_lock: threading.RLock = field( - default_factory=threading.RLock, - repr=False, - ) - finalize_lock: threading.RLock = field( - default_factory=threading.RLock, - repr=False, - ) + logical_llm_lock: threading.RLock = field(default_factory=threading.RLock, repr=False) + finalize_lock: threading.RLock = field(default_factory=threading.RLock, repr=False) _previous_turn: RelayTurnContext | None = field(default=None, repr=False) _active_registered: bool = field(default=False, repr=False) relay_enabled: bool = True @@ -1289,13 +1091,13 @@ _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. +# Depth of managed Relay callbacks on the current logical call path (>0 while the +# native pipeline is mid-dispatch of a Hermes tool/LLM callback). Nested managed +# execution there is structurally broken: the native pipeline binds its Futures to +# the OUTER call's event loop, which is blocked inside the synchronous callback +# ("attached to a different loop" at best, deadlock or "Event loop is closed" at +# worst), so resolve_execution_context() bypasses Relay while set. A ContextVar so +# the marker follows contextvars.copy_context() into worker threads / per-thread loops. _MANAGED_CALLBACK_DEPTH: contextvars.ContextVar[int] = contextvars.ContextVar( "hermes_relay_managed_callback_depth", default=0 ) @@ -1304,10 +1106,9 @@ _MANAGED_CALLBACK_DEPTH: contextvars.ContextVar[int] = contextvars.ContextVar( 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. + Wrap the ``invoke()`` callbacks handed to the native pipeline; everything + they transitively call (including work forwarded via copy_context()) sees + the marker and runs unmanaged. """ def __enter__(self) -> "managed_callback_guard": @@ -1324,43 +1125,25 @@ class RelaySessionCoordinator: def __init__(self, registry: RelayHostRegistry = HOST_REGISTRY) -> None: self.registry = registry self._initializer_lock = threading.RLock() - self._session_initializers: dict[ - str, - Callable[[RelayRuntime, dict[str, Any]], None], - ] = {} + self._session_initializers: dict[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, - name: str, - callback: Callable[[RelayRuntime, dict[str, Any]], None], + self, name: str, callback: Callable[[RelayRuntime, dict[str, Any]], None] ) -> None: """Register idempotent profile/session preparation before scope creation.""" with self._initializer_lock: self._session_initializers[name] = callback - def unregister_session_initializer(self, name: str) -> None: - """Remove a previously registered session initializer.""" - with self._initializer_lock: - self._session_initializers.pop(name, None) - - def _prepare_session( - self, - host: RelayRuntime, - context: dict[str, Any], - ) -> None: + def _prepare_session(self, host: RelayRuntime, context: dict[str, Any]) -> None: with self._initializer_lock: initializers = list(self._session_initializers.items()) for name, callback in initializers: try: callback(host, context) except Exception: - logger.warning( - "Hermes Relay session initializer failed: %s", - name, - exc_info=True, - ) + logger.warning("Hermes Relay session initializer failed: %s", name, exc_info=True) def acquire_conversation( self, @@ -1377,33 +1160,23 @@ class RelaySessionCoordinator: session = None if isinstance(host, RelayRuntime): try: - session_context = { + self._prepare_session(host, { "profile_key": profile_key, "session_id": session_id, "platform": platform, "parent_session_id": parent_session_id, "model": model, - } - self._prepare_session(host, session_context) + }) metadata = {"hermes.execution_surface": platform or "unknown"} if parent_session_id and parent_session_id != session_id: session = host.register_subagent( - { - "parent_session_id": parent_session_id, - "child_session_id": session_id, - }, + {"parent_session_id": parent_session_id, "child_session_id": session_id}, metadata=metadata, ) else: - session = host.ensure_session( - {"session_id": session_id}, - metadata=metadata, - ) + session = host.ensure_session({"session_id": session_id}, metadata=metadata) except Exception: - logger.warning( - "Hermes Relay conversation initialization failed", - exc_info=True, - ) + logger.warning("Hermes Relay conversation initialization failed", exc_info=True) return ConversationLease( profile_key=profile_key, session_id=session_id, @@ -1414,22 +1187,16 @@ class RelaySessionCoordinator: ) def begin_turn( - self, - lease: ConversationLease, - *, - turn_id: str, - task_id: str, + self, lease: ConversationLease, *, turn_id: str, task_id: str ) -> RelayTurnContext: if lease.released: raise RuntimeError("Hermes Relay conversation lease is released") turn = RelayTurnContext(lease=lease, turn_id=turn_id, task_id=task_id) key = (lease.profile_key, lease.session_id) with self._active_turns_lock: - active = self._active_turns.get(key) - if active: - # A Relay session owns one physical scope stack. Concurrent - # Hermes turns would create sibling scopes on that stack, but - # their completion order is not guaranteed to be LIFO. + if self._active_turns.get(key): + # A Relay session owns one physical scope stack; concurrent turns + # would create sibling scopes whose completion order is not LIFO. turn.relay_enabled = False logger.warning( "Skipping Relay instrumentation for concurrent Hermes turn " @@ -1440,45 +1207,26 @@ class RelaySessionCoordinator: else: self._active_turns[key] = {id(turn)} turn._active_registered = True - if ( - turn.relay_enabled - and isinstance(lease.host, RelayRuntime) - and lease.session is not None - ): - # Session-span segmentation: consume a pending rotation (set by - # compaction) or the max_turns cap HERE — the only point where - # no turn scope is live on this session's stack, so the session - # scope can close/reopen without violating LIFO order. + host = lease.live_runtime() if turn.relay_enabled else None + if host is not None: + # Segment rotation (pending compaction flag or max_turns cap) happens + # HERE — the only point with no live turn scope on the session's + # stack, so the session scope can close/reopen without breaking LIFO. try: - config = _segments_config() - session = lease.session - cap = config["max_turns"] - if (config["on_compaction"] and session.rotate_pending) or ( - cap > 0 and session.segment_turns >= cap - ): - reason = ( - "compaction" - if config["on_compaction"] and session.rotate_pending - else "max_turns" - ) - lease.host.rotate_session_scope(session, reason=reason) + self._maybe_rotate_segment(host, lease.session) except Exception: - logger.warning( - "Hermes Relay segment rotation failed", exc_info=True - ) + logger.warning("Hermes Relay segment rotation failed", exc_info=True) try: - turn.handle = lease.host.run_in_session( + turn.handle = host.run_in_session( lease.session, - lease.host.relay.scope.push, + host.relay.scope.push, TURN_SCOPE, - lease.host.relay.ScopeType.Function, + host.relay.ScopeType.Function, handle=lease.session.handle, input={}, - metadata={ - RUNTIME_SCHEMA_KEY: RUNTIME_SCHEMA_VERSION, - RUNTIME_INSTANCE_KEY: lease.host.runtime_id, - "hermes.execution_surface": lease.platform or "unknown", - }, + metadata=runtime_metadata( + host.runtime_id, **{"hermes.execution_surface": lease.platform or "unknown"} + ), timeout=_SCOPE_OP_TIMEOUT, ) except Exception: @@ -1487,122 +1235,101 @@ class RelaySessionCoordinator: _CURRENT_TURN.set(turn) return turn - def end_turn( - self, - turn: RelayTurnContext, - *, - outcome: str, - ) -> None: + @staticmethod + def _maybe_rotate_segment(host: RelayRuntime, session: RelaySession) -> None: + """Rotate the session scope when compaction flagged it or the turn cap is hit.""" + config = _segments_config() + cap = config["max_turns"] + compaction = config["on_compaction"] and session.rotate_pending + if compaction or (cap > 0 and session.segment_turns >= cap): + host.rotate_session_scope(session, reason="compaction" if compaction else "max_turns") + + def end_turn(self, turn: RelayTurnContext, *, outcome: str) -> None: with turn.finalize_lock: if turn.closed: self._reset_turn_context(turn) return turn.closed = True lease = turn.lease + host = lease.live_runtime() try: - if isinstance(lease.host, RelayRuntime) and lease.session is not None: - self._finish_logical_calls(turn, outcome=outcome) - if turn.handle is not None: - failure = lease.host._close_scope_handle( - lease.session, - turn.handle, - output={"outcome": outcome}, - failure_label="turn scope close failed", - ) - if failure: - logger.warning( - "Hermes Relay turn finalization failed: %s", - failure, - ) + if host is not None: + self._close_turn_scope(host, turn, outcome=outcome) finally: try: # Segment turn accounting (max_turns rotation trigger). - if ( - turn._active_registered - and isinstance(lease.host, RelayRuntime) - and lease.session is not None - ): + if turn._active_registered and host is not None: with lease.session.lock: lease.session.segment_turns += 1 except Exception: # noqa: BLE001 - accounting must never block pass try: - # Delegated agents own one turn. Close their conversation + # Delegated agents own one turn: close their conversation # while the active-turn guard is still held so a parent # timeout fallback cannot race this terminal boundary. - if ( - lease.parent_session_id - and isinstance(lease.host, RelayRuntime) - ): - lease.host.unregister_subagent({ - "child_session_id": lease.session_id - }) + if lease.parent_session_id and isinstance(lease.host, RelayRuntime): + lease.host.unregister_subagent({"child_session_id": lease.session_id}) except Exception: logger.warning( - "Hermes Relay child conversation finalization failed", - exc_info=True, + "Hermes Relay child conversation finalization failed", exc_info=True ) finally: self._unregister_active_turn(turn) self._reset_turn_context(turn) self._consume_deferred_close(lease) + def _close_turn_scope( + self, host: RelayRuntime, turn: RelayTurnContext, *, outcome: str + ) -> None: + """Pop the turn's logical LLM children, then the turn scope itself (LIFO).""" + self._finish_logical_calls(turn, outcome=outcome) + if turn.handle is None: + return + failure = host._close_scope_handle( + turn.lease.session, + turn.handle, + output={"outcome": outcome}, + failure_label="turn scope close failed", + ) + if failure: + logger.warning("Hermes Relay turn finalization failed: %s", failure) + def _consume_deferred_close(self, lease: Any) -> None: """Close a session whose rotating-compaction close was deferred. - ``notify_session_compacted`` sets ``close_pending`` instead of - closing when the old session still has a live turn (closing then - would pop the session scope under the live turn scope — LIFO - violation). The turn that was live consumes the flag here, after - its own turn scope popped and it unregistered from the - active-turn table. Skips when another turn is still live on the - same session; that turn's end_turn will consume it instead. + ``notify_session_compacted`` sets ``close_pending`` when the old session + still had a live turn (closing then would break LIFO). The turn that was + live consumes it here, after its own scope popped and it left the + active-turn table; if another turn is still live, that turn's end_turn + consumes it instead. """ try: - if not ( - isinstance(lease.host, RelayRuntime) and lease.session is not None - ): + host = lease.live_runtime() + if host is None: return session = lease.session with session.lock: pending = session.close_pending and not session.closing if not pending: return - if self.has_active_turn( - profile_key=lease.profile_key, session_id=lease.session_id - ): + if self.has_active_turn(profile_key=lease.profile_key, session_id=lease.session_id): return - lease.host.close_session({"session_id": lease.session_id}) + host.close_session({"session_id": lease.session_id}) except Exception: # noqa: BLE001 - telemetry must never block end_turn - logger.warning( - "Hermes Relay deferred session close failed", exc_info=True - ) + logger.warning("Hermes Relay deferred session close failed", exc_info=True) def notify_session_compacted( - self, - *, - profile_key: str, - session_id: str, - old_session_id: str = "", + self, *, profile_key: str, session_id: str, old_session_id: str = "" ) -> None: """React to a completed compaction, per compaction mode. - In-place compaction (``old_session_id`` empty or equal to - ``session_id``): flag the session for segment rotation at its next - turn boundary. Never rotates immediately — a compaction can - complete while a turn is live, and rotation under a live turn - scope would violate the scope stack's LIFO order; ``begin_turn`` - consumes the flag. - - Legacy rotating compaction (``old_session_id`` differs): the next - turn acquires a fresh Relay session under the new id on its own, - but the OLD session's scope would stay open forever — an - unexported orphan. Close it now so the pre-compaction segment - exports. - - Unknown sessions and disabled config are silent no-ops; this - method must never add work or failure modes to the compaction - critical path. + In-place compaction (``old_session_id`` empty or equal): flag the session + for rotation at its next turn boundary — never rotate immediately, since + a compaction can finish while a turn is live and rotating under it would + break LIFO; ``begin_turn`` consumes the flag. Rotating compaction (ids + differ): the next turn gets a fresh session under the new id, so close + the OLD session now or its scope stays an unexported orphan. Unknown + sessions and disabled config are silent no-ops. """ try: if not _segments_config()["on_compaction"]: @@ -1611,11 +1338,8 @@ class RelaySessionCoordinator: if not isinstance(host, RelayRuntime): return if old_session_id and old_session_id != session_id: - # Rotating compaction: export the orphaned pre-compaction - # session scope (close_session is already bounded). If a - # turn is still LIVE on the old session, closing now would - # pop the session scope under the live turn scope (LIFO - # violation) — defer to that turn's end_turn instead. + # If a turn is still LIVE on the old session, closing now would + # pop the session scope under it (LIFO) — defer to its end_turn. with host._sessions_lock: old_session = host._sessions.get(old_session_id) if old_session is not None and self.has_active_turn( @@ -1635,15 +1359,12 @@ class RelaySessionCoordinator: if not session.closing: session.rotate_pending = True except Exception: # noqa: BLE001 - telemetry must never block compaction - logger.warning( - "Hermes Relay compaction notification failed", exc_info=True - ) + logger.warning("Hermes Relay compaction notification failed", exc_info=True) 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)) + return bool(self._active_turns.get((profile_key, session_id))) def _unregister_active_turn(self, turn: RelayTurnContext) -> None: if not turn._active_registered: @@ -1657,16 +1378,7 @@ class RelaySessionCoordinator: 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, - *, - outcome: str, - ) -> None: + def finish_logical_calls(self, turn: RelayTurnContext, *, outcome: str) -> None: """Close logical LLM children before sibling task aggregation scopes.""" with turn.finalize_lock: if turn.closed: @@ -1674,20 +1386,17 @@ class RelaySessionCoordinator: self._finish_logical_calls(turn, outcome=outcome) @staticmethod - def _finish_logical_calls( - turn: RelayTurnContext, - *, - outcome: str, - ) -> None: + def _finish_logical_calls(turn: RelayTurnContext, *, outcome: str) -> None: lease = turn.lease - if not isinstance(lease.host, RelayRuntime) or lease.session is None: + host = lease.live_runtime() + if host is None: return with turn.logical_llm_lock: logical_calls = list(turn.logical_llm_calls.items()) turn.logical_llm_calls.clear() for index in range(len(logical_calls) - 1, -1, -1): request_id, logical_handle = logical_calls[index] - failure = lease.host._close_scope_handle( + failure = host._close_scope_handle( lease.session, logical_handle, output={"outcome": outcome}, @@ -1696,17 +1405,11 @@ class RelaySessionCoordinator: if failure is None: continue with turn.logical_llm_lock: - # Relay scopes are stack-owned. If the newest remaining - # handle cannot close even after orphan drain, older - # handles cannot close safely either — retain the - # unclosed prefix for diagnostics (#81521). - for pending_request_id, pending_handle in logical_calls[ - : index + 1 - ]: - turn.logical_llm_calls.setdefault( - pending_request_id, - pending_handle, - ) + # Relay scopes are stack-owned: if the newest remaining handle + # cannot close even after orphan drain, older ones cannot close + # safely either — retain the unclosed prefix for diagnostics. + for pending_request_id, pending_handle in logical_calls[: index + 1]: + turn.logical_llm_calls.setdefault(pending_request_id, pending_handle) logger.warning("Hermes Relay logical LLM finalization failed: %s", failure) break @@ -1730,19 +1433,11 @@ class RelaySessionCoordinator: """Release a caller lease without closing a resumable conversation.""" lease.released = True - def finalize_conversation( - self, - *, - profile_key: str, - session_id: str, - ) -> None: + def finalize_conversation(self, *, profile_key: str, session_id: str) -> None: host = self.registry.for_profile(profile_key, create=False) if isinstance(host, RelayRuntime): host.close_session({"session_id": session_id}) - def shutdown_profile(self, profile_key: str) -> None: - self.registry.shutdown_profile(profile_key) - SESSION_COORDINATOR = RelaySessionCoordinator() @@ -1761,12 +1456,7 @@ def relay_instrumentation_enabled() -> bool: def active_turn(session_id: str | None = None) -> RelayTurnContext | None: """Return a live turn only when it belongs to the active profile/session.""" turn = current_turn() - if ( - turn is None - or not turn.relay_enabled - or turn.closed - or turn.lease.released - ): + if turn is None or not turn.relay_enabled or turn.closed or turn.lease.released: return None if turn.lease.profile_key != current_profile_key(): return None @@ -1785,39 +1475,23 @@ def resolve_execution_context( ) -> 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. + # Inside a managed Relay callback: nested managed execution is impossible + # (see _MANAGED_CALLBACK_DEPTH). Run unmanaged; the outer 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 - ): + if inherited_turn is not None and (not inherited_turn.relay_enabled or inherited_turn.closed): return None, None, None turn = active_turn(session_id) - if ( - turn is not None - and isinstance(turn.lease.host, RelayRuntime) - and turn.lease.session is not None - ): + host = turn.lease.live_runtime() if turn is not None else None + if host is not None: session = turn.lease.session - return turn.lease.host, session, turn.handle or session.handle + return host, session, turn.handle or session.handle # Managed-execution consumers create and retain the profile host before - # reaching an out-of-turn adapter. Do not initialize Relay for the default + # reaching an out-of-turn adapter; never initialize Relay for the default # no-consumer path. runtime = get_runtime(create=False) - if runtime is None: - return None, None, None - if not runtime.managed_execution_enabled(): + if runtime is None or not runtime.managed_execution_enabled(): return None, None, None session = runtime.get_session(session_id) if session is None: @@ -1825,34 +1499,20 @@ def resolve_execution_context( return runtime, session, None if session is None else session.handle -def emit_mark( - name: str, - *, - session_id: str, - data: Any = None, - metadata: Any = None, -) -> bool: +def emit_mark(name: str, *, session_id: str, data: Any = None, metadata: Any = None) -> bool: """Emit a fail-open Relay mark under a Hermes session.""" runtime = get_runtime(create=False) if runtime is None: return False try: - return runtime.emit_mark( - name, - {"session_id": session_id}, - data=data, - metadata=metadata, - ) + return runtime.emit_mark(name, {"session_id": session_id}, data=data, metadata=metadata) except Exception: logger.warning("Hermes Relay mark failed: %s", name, exc_info=True) return False def apply_tool_request_intercepts( - *, - session_id: str, - tool_name: str, - args: dict[str, Any], + *, session_id: str, tool_name: str, args: dict[str, Any] ) -> dict[str, Any]: """Return Relay-rewritten arguments at Hermes's authorization boundary.""" if not session_id: @@ -1861,69 +1521,12 @@ def apply_tool_request_intercepts( if runtime is None: return args return runtime.apply_tool_request_intercepts( - session_id=session_id, - tool_name=tool_name, - args=args, + session_id=session_id, tool_name=tool_name, args=args ) -def ensure_session(*, session_id: str, **context: Any) -> RelaySession | None: - """Create or return the shared Relay session used by Hermes core.""" - runtime = get_runtime() - if runtime is None: - return None - try: - return runtime.ensure_session({"session_id": session_id, **context}) - except Exception: - logger.warning("Hermes Relay session initialization failed", exc_info=True) - return None - - -def run_in_session( - session_id: str, - callback: Callable[..., Any], - *args: Any, - **kwargs: Any, -) -> Any: - """Run a scope, LLM, or tool API against a shared Hermes session.""" - runtime = get_runtime() - if runtime is None: - raise RuntimeError("Hermes Relay runtime is unavailable") - session = runtime.get_session(session_id) - if session is None: - session = runtime.ensure_session({"session_id": session_id}) - if session is None: - raise RuntimeError("Hermes Relay session is unavailable") - return runtime.run_in_session(session, callback, *args, **kwargs) - - -async def run_in_session_async( - session_id: str, - callback: Callable[..., Any], - *args: Any, - **kwargs: Any, -) -> Any: - """Await a Relay operation inside a shared Hermes session context.""" - runtime = get_runtime() - if runtime is None: - raise RuntimeError("Hermes Relay runtime is unavailable") - session = runtime.get_session(session_id) - if session is None: - session = runtime.ensure_session({"session_id": session_id}) - if session is None: - raise RuntimeError("Hermes Relay session is unavailable") - return await runtime.run_in_session_async(session, callback, *args, **kwargs) - - -def get_session_handle(session_id: str) -> Any: - """Return the shared Relay handle for direct core instrumentation.""" - runtime = get_runtime(create=False) - return None if runtime is None else runtime.get_session_handle(session_id) - - def _is_relay_wrapped_callback_error( - relay_error: BaseException, - callback_error: BaseException, + relay_error: BaseException, callback_error: BaseException ) -> bool: """Match Relay's native callback wrapper without masking policy errors.""" if relay_error is callback_error: @@ -1943,25 +1546,12 @@ def _is_relay_wrapped_callback_error( ) -def get_runtime( - *, - create: bool = True, - profile_key: str | None = None, -) -> RelayRuntime | None: +def get_runtime(*, create: bool = True, profile_key: str | None = None) -> RelayRuntime | None: """Return the Relay host for the active Hermes profile.""" host = HOST_REGISTRY.for_profile(profile_key, create=create) return host if isinstance(host, RelayRuntime) else None -def get_host( - *, - create: bool = True, - profile_key: str | None = None, -) -> RelayHost | None: - """Return the explicit real or reduced-capability host for a profile.""" - return HOST_REGISTRY.for_profile(profile_key, create=create) - - def current_profile_key() -> str: """Return the canonical profile identity used for runtime isolation.""" home = get_hermes_home().expanduser() @@ -1980,9 +1570,7 @@ def _load_nemo_relay() -> Any: return importlib.import_module("nemo_relay") -def _configured_plugin_inputs( - relay: Any, -) -> tuple[dict[str, Any], list[Any]] | None: +def _configured_plugin_inputs(relay: Any) -> tuple[dict[str, Any], list[Any]] | None: """Load selected plugin inputs, or return ``None`` when none were selected.""" configured = os.environ.get(RELAY_PLUGINS_CONFIG_ENV, "").strip() if not configured: @@ -2008,9 +1596,7 @@ def _configured_plugin_inputs( ) dynamic_plugins: list[Any] = [] if "plugins" in config: - dynamic_plugins = relay.plugin.load_dynamic_plugin_activation_specs( - config_path - ) + dynamic_plugins = relay.plugin.load_dynamic_plugin_activation_specs(config_path) plugin_config = dict(config) plugin_config.pop("plugins", None) return plugin_config, dynamic_plugins @@ -2021,16 +1607,6 @@ def _configured_plugin_inputs( ) from exc -def _flush_relay_subscribers(relay: Any) -> None: - """Flush Relay without blocking an asyncio event-loop thread.""" - _resolve_plugin_awaitable(relay.subscribers.flush_async()) - - -def _clear_relay_plugins(relay: Any) -> None: - """Clear Relay plugins without blocking an asyncio event-loop thread.""" - _resolve_plugin_awaitable(relay.plugin.clear_async()) - - def _resolve_plugin_awaitable(value: Any) -> Any: """Resolve Relay's async plugin API from synchronous host construction.""" if not inspect.isawaitable(value): @@ -2039,26 +1615,9 @@ def _resolve_plugin_awaitable(value: Any) -> Any: asyncio.get_running_loop() except RuntimeError: return asyncio.run(value) - - result: dict[str, Any] = {} - error: dict[str, BaseException] = {} - - def _runner() -> None: - try: - result["value"] = asyncio.run(value) - except BaseException as exc: # pragma: no cover - re-raised below - error["exc"] = exc - - thread = threading.Thread( - target=_runner, - name="hermes-nemo-relay-plugin-lifecycle", - daemon=True, + return _run_on_daemon_thread( + lambda: asyncio.run(value), name="hermes-nemo-relay-plugin-lifecycle" ) - thread.start() - thread.join() - if "exc" in error: - raise error["exc"] - return result.get("value") def _session_id(event: dict[str, Any]) -> str: @@ -2067,7 +1626,8 @@ 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() + with SESSION_COORDINATOR._active_turns_lock: + SESSION_COORDINATOR._active_turns.clear() HOST_REGISTRY.shutdown_all() _PLUGIN_CONFIGURATION.reset_for_tests() _PROFILE_KEY_CACHE.clear() diff --git a/agent/relay_tools.py b/agent/relay_tools.py index 91c0926b7e..c38d44068e 100644 --- a/agent/relay_tools.py +++ b/agent/relay_tools.py @@ -72,11 +72,7 @@ def execute( and relay_runtime._is_relay_wrapped_callback_error(exc, callback_error) ): raise callback_error - if ( - isinstance(exc, Exception) - and callback_error is None - and "value" in raw_result - ): + if isinstance(exc, Exception) and callback_error is None and "value" in raw_result: logger.warning( "NeMo Relay tool post-processing failed after dispatch success; " "returning the Hermes tool result", diff --git a/tests/hermes_cli/test_relay_shared_metrics_runtime.py b/tests/hermes_cli/test_relay_shared_metrics_runtime.py index 998fae07b2..c800862ba0 100644 --- a/tests/hermes_cli/test_relay_shared_metrics_runtime.py +++ b/tests/hermes_cli/test_relay_shared_metrics_runtime.py @@ -939,7 +939,7 @@ def test_execution_adapters_do_not_create_relay_host_without_a_consumer( assert result is tool_result assert observed_args is tool_args - assert relay_runtime.get_host(create=False) is None + assert relay_runtime.HOST_REGISTRY.for_profile(create=False) is None assert imports == [] @@ -958,7 +958,7 @@ def test_core_runtime_is_fail_open_without_a_published_binding(monkeypatch, capl monkeypatch.setattr(relay_runtime.importlib, "import_module", missing_relay) assert relay_runtime.get_runtime() is None - host = relay_runtime.get_host() + host = relay_runtime.HOST_REGISTRY.for_profile() assert isinstance(host, relay_runtime.NoopRelayRuntime) assert host.profile_key == relay_runtime.current_profile_key() assert "nemo_relay" in host.reason