diff --git a/agent/relay_llm.py b/agent/relay_llm.py index 6f6a3358f4..df52bc269e 100644 --- a/agent/relay_llm.py +++ b/agent/relay_llm.py @@ -56,7 +56,9 @@ class _ManagedAttempt: cls, session_id: str | None, 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.""" + """Return the managed attempt for ``session_id`` (None: the inherited turn's), or None to run unmanaged.""" + if session_id is None: + session_id = _current_session_id() if not session_id: return None runtime, session, parent = relay_runtime.resolve_execution_context(session_id) @@ -68,12 +70,9 @@ class _ManagedAttempt: 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.runtime, self.session, self.request, self.metadata = runtime, session, request, metadata 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( @@ -109,28 +108,28 @@ class _ManagedAttempt: self.raw_response["json"] = _jsonable(raw) return self.raw_response["json"] - def invoke(self, callback: Callable[..., Any], next_request: Any) -> Any: - """Provider callback handed to Relay: run ``callback`` on Relay's (possibly rewritten) request.""" + @contextlib.contextmanager + def _recording_errors(self) -> Iterator[None]: try: - raw = self.run_callback(callback, self.provider_request(next_request)) + yield except BaseException as exc: self.raw_response["error"] = exc raise + + def invoke(self, callback: Callable[..., Any], next_request: Any) -> Any: + """Provider callback handed to Relay: run ``callback`` on Relay's (possibly rewritten) request.""" + with self._recording_errors(): + raw = self.run_callback(callback, self.provider_request(next_request)) return self._record(raw) async def invoke_async(self, callback: Callable[..., Any], next_request: Any) -> Any: - try: + async def call_provider() -> Any: + with relay_runtime.managed_callback_guard(): # nested relay calls run unmanaged + return await callback(final_request) + + with self._recording_errors(): final_request = self.provider_request(next_request) - - async def call_provider() -> Any: - # Nested relay calls run unmanaged — see relay_runtime.managed_callback_guard. - with relay_runtime.managed_callback_guard(): - return await callback(final_request) - raw = await self.context.copy().run(asyncio.create_task, call_provider()) - except BaseException as exc: - self.raw_response["error"] = exc - raise return self._record(raw) def run_managed(self, relay_call: Callable[..., Any], *callbacks: Any) -> Any: @@ -177,8 +176,6 @@ def execute( ) -> Any: """Run one non-streaming physical provider attempt through Relay. ``session_id`` defaults to the inherited Hermes turn's session (unmanaged when there is none).""" - if session_id is None: - session_id = _current_session_id() attempt = _ManagedAttempt.resolve(session_id, request, metadata, name=name, model_name=model_name) if attempt is None: return callback(request) @@ -196,8 +193,6 @@ async def execute_async( session_id: str | None = None, metadata: dict[str, Any] | None = None, defer_logical_completion: bool = False, ) -> Any: """Async ``execute``.""" - if session_id is None: - session_id = _current_session_id() attempt = _ManagedAttempt.resolve(session_id, request, metadata, name=name, model_name=model_name) if attempt is None: return await callback(request) @@ -214,11 +209,9 @@ execute_current_async = execute_async def _has_running_event_loop() -> bool: - try: - asyncio.get_running_loop() - except RuntimeError: - return False - return True + with contextlib.suppress(RuntimeError): + return asyncio.get_running_loop() is not None + return False def stream_current( @@ -227,15 +220,13 @@ def stream_current( defer_logical_completion: bool = False, completed_response_predicate: Callable[[Any], bool] | None = None, ) -> Any: """Run a provider stream under the inherited Hermes turn when present. - With ``completed_response_predicate`` set, a factory that ignores ``stream=True`` and - returns a complete response is unwrapped and returned directly (pre-Relay behavior) - instead of staying trapped as ``final_response``. Detecting that primes the lazy - pipeline: a genuine first chunk is buffered, but provider latency and pre-first-yield + With ``completed_response_predicate`` set, a factory that ignores ``stream=True`` and returns a + complete response is unwrapped and returned directly (pre-Relay behavior). Detecting that primes + the lazy pipeline: a genuine first chunk is buffered, but provider latency and pre-first-yield errors may surface before this returns.""" session_id = _current_session_id() - # On the Relay session's loop (inside a managed callback) a nested ManagedLlmStream would - # be iterated synchronously on that loop, which asyncio forbids; the outer managed stream - # already tracks this attempt. + # Inside a managed callback (on the Relay session's loop) a nested ManagedLlmStream would be + # iterated synchronously on that loop, which asyncio forbids; the outer stream tracks this attempt. if session_id is None or _has_running_event_loop(): return stream_factory(request) managed = stream( @@ -244,12 +235,9 @@ 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 pull; prime once so a - # completed response surfaces (a real first chunk is buffered). + # Relay may defer the provider callback until the first pull; prime once (a real first chunk is buffered). managed._prime_completed_response() - if managed.final_response is not None: - return managed.final_response - return managed + return managed.final_response if managed.final_response is not None else managed def _aclose_on_loop(loop: asyncio.AbstractEventLoop, stream: Any) -> None: @@ -258,7 +246,7 @@ def _aclose_on_loop(loop: asyncio.AbstractEventLoop, stream: Any) -> None: if not callable(close): return - async def close_stream() -> None: + async def close_stream() -> None: # create the coroutine on ``loop``, not the caller's thread await close() loop.run_until_complete(close_stream()) @@ -287,17 +275,12 @@ class ManagedLlmStream(Iterator[Any]): self._defer_logical_completion = defer_logical_completion # 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._on_chunk = on_chunk - self._chunk_adapter = chunk_adapter or _namespace - self._accept_chunk = accept_chunk + self._logical_model_name, self._logical_provider_name = (model_name, name) if auxiliary else (None, None) + self._on_chunk, self._chunk_adapter, self._accept_chunk = on_chunk, chunk_adapter or _namespace, accept_chunk + self._stream_factory, self._on_stream_created, self._finalizer = stream_factory, on_stream_created, finalizer + self._completed_response_predicate = completed_response_predicate self._raw_chunks: list[tuple[Any, Any]] = [] self._prefetched_chunks: list[Any] = [] - self._stream_factory = stream_factory - self._on_stream_created = on_stream_created - self._completed_response_predicate = completed_response_predicate - self._finalizer = finalizer attempt = _ManagedAttempt.resolve(session_id, request, metadata, name=name, model_name=model_name) if attempt is None: self._start_unmanaged(request) @@ -355,8 +338,8 @@ class ManagedLlmStream(Iterator[Any]): raise def _relay_finalizer(self, attempt: _ManagedAttempt) -> Any: - # Relay may call this while unwinding a provider-stream failure; keep the - # original error instead of a secondary "missing terminal response". + # Relay may call this while unwinding a provider-stream failure; keep the original + # error instead of a secondary "missing terminal response". if self._callback_error is not None: return None try: @@ -410,7 +393,7 @@ class ManagedLlmStream(Iterator[Any]): def _recoverable_relay_failure(self, exc: BaseException) -> bool: """Relay post-processing failed after the provider already succeeded.""" - recoverable = (isinstance(exc, Exception) and self._provider_completed and self._callback_error is None) + recoverable = isinstance(exc, Exception) and self._provider_completed and self._callback_error is None if recoverable: logger.warning( "NeMo Relay stream post-processing failed after provider success; preserving the provider result", @@ -435,12 +418,8 @@ class ManagedLlmStream(Iterator[Any]): if self._prefetched_chunks: return self._prefetched_chunks.pop() if self._loop is None: - try: - chunk = next(self._stream) - except StopIteration: - self._close(logical_outcome="cancelled") - raise - if self._accept_chunk is not None and not self._accept_chunk(chunk): + chunk = next(self._stream, self) # self: exhausted sentinel + if chunk is self or (self._accept_chunk is not None and not self._accept_chunk(chunk)): self._close(logical_outcome="cancelled") raise StopIteration return chunk @@ -510,10 +489,9 @@ class ManagedLlmStream(Iterator[Any]): self._raw_stream_resource = None for resource in resources.values(): close = getattr(resource, "close", None) - if not callable(close): - continue try: - close() + if callable(close): + close() except Exception as exc: self._keep_first_close_error(exc) logger.debug("Provider stream cleanup failed", exc_info=True) @@ -524,18 +502,17 @@ class ManagedLlmStream(Iterator[Any]): self._closed = True self._prefetched_chunks.clear() try: - loop = self._loop - self._loop = None + loop, self._loop = self._loop, None if loop is None: self._close_provider_resources() - self._finish_logical(logical_outcome) - return - try: - _aclose_on_loop(loop, self._stream) - except Exception as exc: - self._keep_first_close_error(exc) + else: + try: + _aclose_on_loop(loop, self._stream) + except Exception as exc: + self._keep_first_close_error(exc) self._finish_logical(logical_outcome) - loop.close() + if loop is not None: + loop.close() finally: self._release_runtime_lease() @@ -612,10 +589,9 @@ class AnthropicStreamAccumulator: for block in blocks: partial = block.pop("_partial_json", None) if partial is not None: - try: - block["input"] = json.loads(partial) - except (TypeError, ValueError): - block["input"] = partial + with contextlib.suppress(TypeError, ValueError): + partial = json.loads(partial) + block["input"] = partial return {**self._message, "content": blocks} def response(self, base: Any = None) -> Any: @@ -724,9 +700,8 @@ def _provider_request( 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) - # 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. + # Typed codecs may not represent provider-specific fields: overlay only values that changed + # from the codec-facing baseline so unrelated intercepts cannot delete/normalize unknown arguments. for key in baseline.keys() | intercepted.keys(): if key not in intercepted: final.pop(key, None) @@ -805,11 +780,11 @@ def _codec_round_trip_request_body( try: encoded = codec.encode(codec.decode(relay_request), relay_request) content = getattr(encoded, "content", encoded) - if isinstance(content, dict): - return _provider_request_body(content, metadata) except Exception: logger.warning("NeMo Relay request codec baseline failed; ignoring request rewrites", exc_info=True) return None + if isinstance(content, dict): + return _provider_request_body(content, metadata) logger.warning("NeMo Relay request codec returned an unsupported baseline; ignoring request rewrites") return None @@ -826,10 +801,7 @@ def _provider_request_body(content: dict[str, Any], metadata: dict[str, Any] | N def _codec(relay: Any, metadata: dict[str, Any] | None) -> Any: protocol = _RELAY_PROTOCOL_BY_API_MODE.get(_api_mode(metadata)) - codecs = getattr(relay, "codecs", None) - if protocol is None or codecs is None: - return None - codec = getattr(codecs, protocol[1], None) + codec = getattr(getattr(relay, "codecs", None), protocol[1], None) if protocol is not None else None return codec() if callable(codec) else None diff --git a/agent/relay_runtime.py b/agent/relay_runtime.py index 2f28d8c21b..6b709b376e 100644 --- a/agent/relay_runtime.py +++ b/agent/relay_runtime.py @@ -6,6 +6,7 @@ import atexit import asyncio import contextlib import contextvars +import functools import importlib import inspect import logging @@ -15,7 +16,7 @@ import tomllib import uuid from concurrent.futures import TimeoutError as FuturesTimeoutError from dataclasses import dataclass, field -from enum import Enum, auto +from enum import Enum from pathlib import Path from typing import Any, Callable @@ -37,8 +38,34 @@ _PROFILE_KEY_CACHE: dict[str, str] = {} # lost span, never a blocked agent. _SCOPE_OP_TIMEOUT = 10.0 -_SCOPE_OP_EXECUTOR: Any = None -_SCOPE_OP_EXECUTOR_LOCK = threading.Lock() + + +class _Lazy: + """Double-checked, thread-safe once-only factory; ``reset()`` re-arms it (tests).""" + + def __init__(self, factory: Callable[[], Any]) -> None: + self._factory, self._value, self._lock = factory, None, threading.Lock() + + def get(self) -> Any: + if self._value is None: + with self._lock: + if self._value is None: + self._value = self._factory() + return self._value + + def reset(self) -> None: + self._value = None + + +def _new_scope_op_executor() -> Any: + # Daemon workers: a wedged call abandoned at timeout cannot block interpreter exit; + # ``Future.result(timeout=...)`` still bounds callers when every worker is wedged. + from tools.daemon_pool import DaemonThreadPoolExecutor + return DaemonThreadPoolExecutor(max_workers=8, thread_name_prefix="relay-scope-op") + + +_SCOPE_OP_EXECUTOR = _Lazy(_new_scope_op_executor) +_scope_op_executor = _SCOPE_OP_EXECUTOR.get def runtime_metadata(runtime_id: str, **extra: Any) -> dict[str, Any]: @@ -46,25 +73,11 @@ def runtime_metadata(runtime_id: str, **extra: Any) -> dict[str, Any]: return {RUNTIME_SCHEMA_KEY: RUNTIME_SCHEMA_VERSION, RUNTIME_INSTANCE_KEY: runtime_id, **extra} -def _scope_op_executor(): - """Shared daemon executor for bounded native scope ops. - Daemon workers so a wedged call abandoned at timeout cannot block interpreter exit; - ``Future.result(timeout=...)`` still bounds callers when every worker is wedged.""" - global _SCOPE_OP_EXECUTOR - if _SCOPE_OP_EXECUTOR is None: - with _SCOPE_OP_EXECUTOR_LOCK: - if _SCOPE_OP_EXECUTOR is None: - from tools.daemon_pool import DaemonThreadPoolExecutor - _SCOPE_OP_EXECUTOR = DaemonThreadPoolExecutor(max_workers=8, thread_name_prefix="relay-scope-op") - return _SCOPE_OP_EXECUTOR - - 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. - With ``timeout`` a still-running worker is abandoned with ``TimeoutError`` (daemon: - cannot block interpreter exit).""" + With ``timeout`` a still-running worker is abandoned with ``TimeoutError`` (daemon: cannot block exit).""" outcome: dict[str, Any] = {} def _target() -> None: @@ -87,8 +100,7 @@ def pop_relay_scope(relay: Any, handle: Any, *, output: Any = None, metadata: An """Pop a Relay scope, forwarding only the kwargs the live binding accepts. ``scope.pop`` gained ``metadata`` in nemo-relay 0.4+; older wheels raise TypeError.""" pop = relay.scope.pop - candidates = (("output", output), ("metadata", metadata), ("timestamp", timestamp)) - kwargs = {key: value for key, value in candidates if value is not None} + kwargs = {k: v for k, v in (("output", output), ("metadata", metadata), ("timestamp", timestamp)) if v is not None} try: params = inspect.signature(pop).parameters except (TypeError, ValueError): @@ -100,8 +112,7 @@ def pop_relay_scope(relay: Any, handle: Any, *, output: Any = None, metadata: An 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. + # Prefer scope.get_handle(): get_scope_stack() may return a native ScopeStack that scope.pop rejects. get_handle = getattr(getattr(relay, "scope", None), "get_handle", None) if callable(get_handle): with contextlib.suppress(Exception): @@ -119,14 +130,8 @@ def _same_handle(a: Any, b: Any) -> bool: return a is b or a == b or (a_uuid is not None and a_uuid == getattr(b, "uuid", None)) -class _RelayPluginConfigurationState(Enum): - """Process-wide result shared by every currently hosted profile.""" - - UNINITIALIZED = auto() - DISABLED = auto() - ACTIVE = auto() - FOREIGN = auto() - FAILED = auto() +# Process-wide plugin-configuration result shared by every currently hosted profile. +_RelayPluginConfigurationState = Enum("_RelayPluginConfigurationState", "UNINITIALIZED DISABLED ACTIVE FOREIGN FAILED") class _RelayPluginConfigurationLoadError(RuntimeError): @@ -143,23 +148,16 @@ class RelaySession: closing: bool = False handle: Any = None context: contextvars.Context | None = None - # Session-span segmentation: rotation closes the current session scope and pushes - # segment N+1 at a turn boundary (the only LIFO-safe point). + # Segmentation: rotation closes the session scope and pushes segment N+1 at a turn + # boundary (the only LIFO-safe point). Flags are set by compaction, consumed at that 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 - # Rotating compaction landed while a turn was live here; closing now would pop the - # session scope under the live turn, so end_turn consumes this instead. - close_pending: bool = False - - -# gateway.telemetry.session_segments, cached at first read. Both defaults OFF => -# rotation never fires and the scope lifecycle is unchanged. -_SEGMENTS_CONFIG: dict[str, Any] | None = None -_SEGMENTS_CONFIG_LOCK = threading.Lock() + rotate_pending: bool = False # consumed at next begin_turn + close_pending: bool = False # rotating compaction hit a live turn; end_turn consumes it def _load_segments_config() -> dict[str, Any]: + """gateway.telemetry.session_segments; both defaults OFF => rotation never fires.""" segments: dict[str, Any] = {} with contextlib.suppress(Exception): # config absence must not crash from gateway.run import _load_gateway_config # late import @@ -172,27 +170,16 @@ def _load_segments_config() -> dict[str, Any]: return {"on_compaction": bool(segments.get("on_compaction", False)), "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: - _SEGMENTS_CONFIG = _load_segments_config() - return _SEGMENTS_CONFIG - - -def _reset_segments_config_for_tests() -> None: - global _SEGMENTS_CONFIG - _SEGMENTS_CONFIG = None +_SEGMENTS_CONFIG = _Lazy(_load_segments_config) # cached at first read +_segments_config = _SEGMENTS_CONFIG.get +_reset_segments_config_for_tests = _SEGMENTS_CONFIG.reset class RelayOperationLease: """Keep process-wide Relay plugins alive across a deferred operation.""" def __init__(self, runtime: "RelayRuntime") -> None: - self._lock = threading.Lock() - self._runtime: RelayRuntime | None = runtime + self._lock, self._runtime = threading.Lock(), runtime def run_in_session(self, session: RelaySession, callback: Callable[..., Any], *args: Any, **kwargs: Any) -> Any: """Run cleanup while this lease still owns the runtime lifetime.""" @@ -254,16 +241,14 @@ class _ProcessRelayPluginConfiguration: existing_report = relay.plugin.report() except Exception: logger.warning( - "Hermes could not determine whether a process-global Relay " - "plugin configuration is already active; refusing to replace it", - exc_info=True, + "Hermes could not determine whether a process-global Relay plugin configuration is already " + "active; refusing to replace it", exc_info=True, ) return _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" + "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 _RelayPluginConfigurationState.FOREIGN return None @@ -291,10 +276,9 @@ class _ProcessRelayPluginConfiguration: def release(self, owner: Any) -> None: """Release one host and clear Relay after the final host exits.""" with self._lock: - if id(owner) not in self._owners: - return - self._owners.remove(id(owner)) - self.retry_pending_cleanup() + if id(owner) in self._owners: + self._owners.remove(id(owner)) + self.retry_pending_cleanup() def reset_for_tests(self) -> None: """Clear process-global state left by directly constructed test hosts.""" @@ -312,21 +296,23 @@ class _ProcessRelayPluginConfiguration: relay, activation = self._relay, self._activation if relay is None: return True - try: - _resolve_plugin_awaitable(relay.subscribers.flush_async()) - except Exception: - logger.warning("Hermes Relay plugin subscriber flush failed", exc_info=True) - return False - try: + + def close_configuration() -> Any: if activation is None: - _resolve_plugin_awaitable(relay.plugin.clear_async()) - elif callable(close := getattr(activation, "close", None)): - _resolve_plugin_awaitable(close()) - else: - raise RuntimeError("NeMo Relay dynamic plugin activation has no close method") - except Exception: - logger.warning("Hermes Relay plugin configuration cleanup failed", exc_info=True) - return False + return _resolve_plugin_awaitable(relay.plugin.clear_async()) + if callable(close := getattr(activation, "close", None)): + return _resolve_plugin_awaitable(close()) + raise RuntimeError("NeMo Relay dynamic plugin activation has no close method") + + for what, step in ( + ("subscriber flush", lambda: _resolve_plugin_awaitable(relay.subscribers.flush_async())), + ("configuration cleanup", close_configuration), + ): + try: + step() + except Exception: + logger.warning("Hermes Relay plugin %s failed", what, exc_info=True) + return False self._relay = self._activation = None return True @@ -342,17 +328,15 @@ class RelayRuntime: self.relay = relay or _load_nemo_relay() self.profile_key = profile_key or current_profile_key() self.runtime_id = uuid.uuid4().hex - self._sessions_lock = threading.RLock() + self._sessions_lock, self._execution_consumers_lock = threading.RLock(), threading.RLock() self._sessions: dict[str, RelaySession] = {} self._subagent_parents: dict[str, str] = {} self._subagent_parent_handles: dict[str, Any] = {} + self._execution_consumers: set[str] = set() self._closing = self._shutdown_started = False - self._shutdown_complete = threading.Event() - self._operations_idle = threading.Event() + self._shutdown_complete, self._operations_idle = threading.Event(), threading.Event() self._operations_idle.set() 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) # Cleared (with the atexit hook) by the first successful _finish_shutdown. self._plugin_configuration_registered = True @@ -385,9 +369,8 @@ class RelayRuntime: exit_fallback: bool = False, **push_kwargs: Any, ) -> None: """Push a fresh SESSION_SCOPE for ``session`` (bounded by ``_SCOPE_OP_TIMEOUT``); record handle + context. - Subagents parent under their spawning turn/session handle; ``resolve_parent`` creates the parent - session when its handle is unknown. ``exit_fallback``: at interpreter shutdown the executor refuses - new futures; push synchronously instead (no agent turn waits at exit).""" + Subagents parent under their spawning turn/session handle (``resolve_parent`` creates the parent when + unknown). ``exit_fallback``: at interpreter shutdown the executor refuses futures; push synchronously.""" parent_handle = None if session.parent_session_id: with self._sessions_lock: @@ -421,28 +404,22 @@ 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) - self._sessions[session_id] = session + session = self._sessions[session_id] = RelaySession( + session_id=session_id, parent_session_id=self._subagent_parents.get(session_id, "") + ) with session.lock: if session.closing: return None if session.handle is None: - try: - 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 + self._open_session_scope( + session, {**(metadata or {}), **runtime_metadata(self.runtime_id)}, + resolve_parent=True, data=data, exit_fallback=True, + ) 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: the stack is LIFO and rotating under a live child - would close a parent out of order. Bookkeeping advances even when a native call - fails so a degraded rotation cannot retry on every turn.""" + """Close the current session scope and open the next segment (turn boundary ONLY: LIFO). + Bookkeeping advances even when a native call fails so a degraded rotation cannot retry every turn.""" with session.lock: if session.closing or session.handle is None: return @@ -482,12 +459,8 @@ class RelayRuntime: return None parent = self.ensure_session({"session_id": parent_session_id}) parent_handle = None if parent is None else parent.handle - turn = active_turn(parent_session_id) - # active_turn() already proved liveness and, for a RelayRuntime host, an open session. - if ( - turn is not None and turn.handle is not None and turn.lease.host is self - and turn.lease.session.session_id == parent_session_id - ): + turn = active_turn(parent_session_id) # already proved liveness and (RelayRuntime host) an open session + if turn is not None and turn.handle is not None and turn.lease.host is self: parent_handle = turn.handle with self._sessions_lock: if self._closing: @@ -524,14 +497,13 @@ class RelayRuntime: return None if session.closing else session def _session_context(self, session: RelaySession, *, allow_closing: bool) -> contextvars.Context: - """Copy the current context and overlay the session's saved Relay vars.""" + """Copy the current context and overlay the session's saved Relay vars (a copy: re-entrant from callbacks).""" 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 lets a helper inside a Relay callback re-enter the session's Context. context = contextvars.copy_context() for variable, value in relay_context.items(): context.run(variable.set, value) @@ -542,9 +514,8 @@ class RelayRuntime: allow_closing: bool = False, timeout: float | None = None, **kwargs: Any, ) -> Any: """Run a Relay operation against a session's isolated scope stack. - ``timeout`` bounds the native call on the daemon executor (``TimeoutError`` on - breach); ``None`` runs synchronously. Lifecycle ops gating turn/session completion - pass ``_SCOPE_OP_TIMEOUT``: a wedged pipeline must cost one span, never the agent.""" + ``timeout`` bounds the native call on the daemon executor (``TimeoutError`` on breach); ``None`` + runs synchronously. Lifecycle ops pass ``_SCOPE_OP_TIMEOUT``: a wedged pipeline costs one span.""" with self._operation(): return self._run_in_session_untracked( session, callback, *args, allow_closing=allow_closing, timeout=timeout, **kwargs @@ -567,8 +538,8 @@ class RelayRuntime: 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 — still bounded so a wedged call cannot block exit. + # Interpreter shutdown: the executor refuses futures but the atexit close path must + # still flush — still bounded so a wedged call cannot block exit. return _run_on_daemon_thread( lambda: context.run(invoke), name="relay-scope-op-exit", timeout=timeout, timeout_message=f"{exceeded} during interpreter shutdown; abandoning the native " @@ -628,9 +599,8 @@ class RelayRuntime: def apply_tool_request_intercepts(self, *, session_id: str, tool_name: str, args: dict[str, Any]) -> dict[str, Any]: """Apply Relay request rewriting before Hermes authorizes a tool call.""" request_intercepts = getattr(getattr(self.relay, "tools", None), "request_intercepts", None) - if not self.managed_execution_enabled() or not callable(request_intercepts): - return args - session = self.ensure_session({"session_id": session_id}) + managed = self.managed_execution_enabled() and callable(request_intercepts) + session = self.ensure_session({"session_id": session_id}) if managed else None if session is None: return args result = self.run_in_session(session, request_intercepts, tool_name, args) @@ -639,9 +609,8 @@ class RelayRuntime: 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.""" + """Pop ``handle``; if that fails, drain orphans above it and retry once; return the retry's error. + Must run inside ONE ``run_in_session`` callback so ContextVar stack views stay consistent.""" with contextlib.suppress(Exception): pop_relay_scope(self.relay, handle, output=output, metadata=metadata) return None @@ -672,10 +641,9 @@ class RelayRuntime: self, session: RelaySession, handle: Any, *, output: dict[str, Any] | None = None, allow_closing: bool = False, failure_label: str = "scope close failed", drain_limit: int = 32, operation_already_held: bool = False, ) -> 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. Drain+close is bounded so a wedged pipeline - never blocks turn/session completion. Returns a failure string or None.""" + """Pop ``handle``, draining orphaned children in the same session context; failure string or None. + Relay scopes are strict LIFO; empty-stream retries + interrupt can abandon a physical LLM scope + above TURN/SESSION. Drain+close is bounded so a wedged pipeline never blocks completion.""" if handle is None: return None run_in_session = (self._run_in_session_untracked if operation_already_held else self.run_in_session) @@ -690,15 +658,9 @@ class RelayRuntime: 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.""" - try: - self._begin_operation() - except RuntimeError: - return - try: + """Close one session scope and remove it from the core registry (no-op once shutting down).""" + with contextlib.suppress(RuntimeError), self._operation(): # _close_session itself never raises self._close_session(event) - finally: - self._end_operation() def _close_session(self, event: dict[str, Any]) -> None: """Close one session already admitted by the host lifecycle gate.""" @@ -717,8 +679,8 @@ class RelayRuntime: session, session.handle, output={}, allow_closing=True, failure_label="session scope close failed", operation_already_held=True, ) - # No subscriber flush here: it is process-wide, may wait on other sessions' - # publications and can deadlock an asyncio loop; final plugin teardown flushes once. + # No subscriber flush here: process-wide, may wait on other sessions and deadlock an asyncio + # loop; final plugin teardown flushes once. with self._sessions_lock: if self._sessions.get(session_id) is session: del self._sessions[session_id] @@ -814,8 +776,7 @@ class RelayHostRegistry: def shutdown_all(self) -> None: with self._lock: - hosts = list(self._hosts.values()) - self._hosts.clear() + hosts, self._hosts = list(self._hosts.values()), {} for host in hosts: host.shutdown() @@ -837,8 +798,7 @@ class ConversationLease: def live_runtime(self) -> RelayRuntime | None: """Return the real Relay host when this lease owns an open session.""" - host = self.host - return host if isinstance(host, RelayRuntime) and self.session is not None else None + return self.host if isinstance(self.host, RelayRuntime) and self.session is not None else None @dataclass @@ -862,28 +822,24 @@ _CURRENT_TURN: contextvars.ContextVar[RelayTurnContext | None] = contextvars.Con "hermes_relay_turn", default=None ) -# Depth of managed Relay callbacks on the current call path (>0 while the native pipeline -# is mid-dispatch of a Hermes tool/LLM callback). Nested managed execution there is -# structurally broken: the pipeline binds its Futures to the OUTER call's event loop, which -# is blocked inside the synchronous callback (wrong loop / deadlock / "Event loop is -# closed"), so resolve_execution_context() bypasses Relay while set. A ContextVar so the -# marker follows copy_context() into worker threads / per-thread loops. +# >0 while the native pipeline is mid-dispatch of a Hermes tool/LLM callback. Nested managed +# execution there is structurally broken (the pipeline binds its Futures to the OUTER call's +# loop, blocked inside the synchronous callback), so resolve_execution_context() bypasses Relay. +# A ContextVar so the marker follows copy_context() into worker threads / per-thread loops. _MANAGED_CALLBACK_DEPTH: contextvars.ContextVar[int] = contextvars.ContextVar( "hermes_relay_managed_callback_depth", default=0 ) -class managed_callback_guard: - """Mark the current context as inside a managed Relay callback. - Wrap the ``invoke()`` callbacks handed to the native pipeline; everything they - transitively call (incl. work forwarded via copy_context()) runs unmanaged.""" - - def __enter__(self) -> "managed_callback_guard": - self._token = _MANAGED_CALLBACK_DEPTH.set(_MANAGED_CALLBACK_DEPTH.get() + 1) - return self - - def __exit__(self, *exc_info: Any) -> None: - _MANAGED_CALLBACK_DEPTH.reset(self._token) +@contextlib.contextmanager +def managed_callback_guard(): + """Mark the current context as inside a managed Relay callback: everything the wrapped ``invoke()`` + transitively calls (incl. work forwarded via copy_context()) runs unmanaged.""" + token = _MANAGED_CALLBACK_DEPTH.set(_MANAGED_CALLBACK_DEPTH.get() + 1) + try: + yield + finally: + _MANAGED_CALLBACK_DEPTH.reset(token) def _warn_on_error(what: str, callback: Callable[..., Any], *args: Any, **kwargs: Any) -> Any: @@ -895,6 +851,16 @@ def _warn_on_error(what: str, callback: Callable[..., Any], *args: Any, **kwargs return None +def _fail_open(what: str) -> Callable[[Callable[..., Any]], Callable[..., Any]]: + """Decorator form of ``_warn_on_error`` for telemetry hooks that must never block the caller.""" + def wrap(fn: Callable[..., Any]) -> Callable[..., Any]: + @functools.wraps(fn) + def guarded(*args: Any, **kwargs: Any) -> Any: + return _warn_on_error(what, fn, *args, **kwargs) + return guarded + return wrap + + def _flag_open_session(session: RelaySession, flag: str) -> None: """Set a pending-rotation/close flag unless the session is already closing.""" with session.lock: @@ -929,9 +895,7 @@ class RelaySessionCoordinator: def acquire_conversation( self, *, profile_key: str, session_id: str, platform: str, parent_session_id: str = "", model: str = "", ) -> ConversationLease: - host = self.registry.for_profile(profile_key) or NoopRelayRuntime( - profile_key, "Relay host creation was disabled" - ) + host = self.registry.for_profile(profile_key) or NoopRelayRuntime(profile_key, "Relay host creation was disabled") session = None if isinstance(host, RelayRuntime): context = { @@ -949,9 +913,8 @@ class RelaySessionCoordinator: session_id, parent_session_id = context["session_id"], context["parent_session_id"] metadata = {"hermes.execution_surface": context["platform"] or "unknown"} if parent_session_id and parent_session_id != session_id: - return host.register_subagent( - {"parent_session_id": parent_session_id, "child_session_id": session_id}, metadata=metadata, - ) + event = {"parent_session_id": parent_session_id, "child_session_id": session_id} + return host.register_subagent(event, metadata=metadata) return host.ensure_session({"session_id": session_id}, metadata=metadata) def begin_turn(self, lease: ConversationLease, *, turn_id: str, task_id: str) -> RelayTurnContext: @@ -961,8 +924,7 @@ class RelaySessionCoordinator: key = (lease.profile_key, lease.session_id) with self._active_turns_lock: if self._active_turns.get(key): - # One physical scope stack per session; concurrent turns would create - # sibling scopes whose completion order is not LIFO. + # One physical scope stack per session; concurrent turns' sibling scopes would not close LIFO. turn.relay_enabled = False logger.warning( "Skipping Relay instrumentation for concurrent Hermes turn %s in session %s", @@ -973,8 +935,7 @@ class RelaySessionCoordinator: turn._active_registered = True host = lease.live_runtime() if turn.relay_enabled else None if host is not None: - # Segment rotation happens HERE — the only point with no live turn scope on - # the stack, so the session scope can close/reopen without breaking LIFO. + # Rotation happens HERE: no live turn scope on the stack, so the session scope can close/reopen LIFO. _warn_on_error("segment rotation", self._maybe_rotate_segment, host, lease.session) turn.handle = _warn_on_error( "turn initialization", host.run_in_session, lease.session, host.relay.scope.push, @@ -1011,8 +972,8 @@ class RelaySessionCoordinator: with contextlib.suppress(Exception), lease.session.lock: # accounting never blocks lease.session.segment_turns += 1 # max_turns rotation trigger try: - # Delegated agents own one turn: close their conversation while the - # active-turn guard is held so a parent timeout fallback cannot race it. + # Delegated agents own one turn: close their conversation while the active-turn + # guard is held so a parent timeout fallback cannot race it. if lease.parent_session_id and isinstance(lease.host, RelayRuntime): _warn_on_error( "child conversation finalization", lease.host.unregister_subagent, @@ -1032,15 +993,10 @@ class RelaySessionCoordinator: 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`` when the old session had a live - turn (closing then breaks LIFO). The last live turn consumes it here after its own - scope popped and it left the active-turn table.""" - # Telemetry must never block end_turn. - _warn_on_error("deferred session close", self._consume_deferred_close_unguarded, lease) - - def _consume_deferred_close_unguarded(self, lease: ConversationLease) -> None: + @_fail_open("deferred session close") + def _consume_deferred_close(self, lease: ConversationLease) -> None: + """Close a session whose rotating-compaction close was deferred (``close_pending``). + The last live turn consumes it after its own scope popped and it left the active-turn table.""" host = lease.live_runtime() if host is None: return @@ -1049,19 +1005,12 @@ class RelaySessionCoordinator: if pending and not self.has_active_turn(profile_key=lease.profile_key, session_id=lease.session_id): host.close_session({"session_id": lease.session_id}) + @_fail_open("compaction notification") def notify_session_compacted(self, *, profile_key: str, session_id: str, old_session_id: str = "") -> None: - """React to a completed compaction, per compaction mode. - In-place (``old_session_id`` empty/equal): flag rotation for the next turn boundary - — never rotate immediately, a turn may be live and rotating under it breaks LIFO. - Rotating (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.""" - # Telemetry must never block compaction. - _warn_on_error( - "compaction notification", self._notify_session_compacted_unguarded, profile_key, session_id, old_session_id - ) - - def _notify_session_compacted_unguarded(self, profile_key: str, session_id: str, old_session_id: str) -> None: + """React to a completed compaction, per compaction mode; unknown sessions / disabled config are no-ops. + In-place (``old_session_id`` empty/equal): flag rotation for the next turn boundary — never rotate + immediately, a live turn under it breaks LIFO. Rotating (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.""" if not _segments_config()["on_compaction"]: return host = self.registry.for_profile(profile_key) @@ -1121,8 +1070,7 @@ class RelaySessionCoordinator: logical_calls.pop() continue with turn.logical_llm_lock: - # Stack-owned scopes: if the newest handle cannot close even after orphan - # drain, older ones cannot either — retain the unclosed prefix. + # Stack-owned: if the newest handle cannot close even after drain, older ones cannot either. for pending_request_id, pending_handle in logical_calls: turn.logical_llm_calls.setdefault(pending_request_id, pending_handle) logger.warning("Hermes Relay logical LLM finalization failed: %s", failure) @@ -1183,11 +1131,9 @@ def active_turn(session_id: str | None = None) -> RelayTurnContext | None: def resolve_execution_context(session_id: str) -> tuple[RelayRuntime | None, RelaySession | None, Any]: """Resolve one active turn/session parent for managed Relay execution.""" - if _MANAGED_CALLBACK_DEPTH.get() > 0: - # Nested managed execution is impossible (see _MANAGED_CALLBACK_DEPTH); the - # outer scope still records the tool-level event. - return None, None, None - if not relay_instrumentation_enabled(): + # Nested managed execution is impossible (see _MANAGED_CALLBACK_DEPTH); the outer scope + # still records the tool-level event. + if _MANAGED_CALLBACK_DEPTH.get() > 0 or not relay_instrumentation_enabled(): return None, None, None turn = active_turn(session_id) host = turn.lease.live_runtime() if turn is not None else None @@ -1205,9 +1151,7 @@ def resolve_execution_context(session_id: str) -> tuple[RelayRuntime | None, Rel def apply_tool_request_intercepts(*, 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: - return args - runtime = get_runtime(create=False) + runtime = get_runtime(create=False) if session_id else None if runtime is None: return args return runtime.apply_tool_request_intercepts(session_id=session_id, tool_name=tool_name, args=args) @@ -1219,12 +1163,9 @@ def _is_relay_wrapped_callback_error(relay_error: BaseException, callback_error: return True if not isinstance(relay_error, RuntimeError): return False - callback_type = callback_error.__class__ - type_names = { - callback_type.__name__, callback_type.__qualname__, f"{callback_type.__module__}.{callback_type.__qualname__}", - } - message = str(relay_error) - return any(message.startswith(f"internal error: {type_name}: {callback_error}") for type_name in type_names) + kind = type(callback_error) + type_names = {kind.__name__, kind.__qualname__, f"{kind.__module__}.{kind.__qualname__}"} + return any(str(relay_error).startswith(f"internal error: {name}: {callback_error}") for name in type_names) def get_runtime(*, create: bool = True, profile_key: str | None = None) -> RelayRuntime | None: @@ -1238,9 +1179,7 @@ def current_profile_key() -> str: home = get_hermes_home().expanduser() if not home.is_absolute(): return str(home.resolve()) - raw = str(home) - cached = _PROFILE_KEY_CACHE.get(raw) - return cached if cached is not None else _PROFILE_KEY_CACHE.setdefault(raw, str(home.resolve())) + return _PROFILE_KEY_CACHE.get(str(home)) or _PROFILE_KEY_CACHE.setdefault(str(home), str(home.resolve())) def _load_nemo_relay() -> Any: @@ -1254,11 +1193,9 @@ def _configured_plugin_inputs(relay: Any) -> tuple[dict[str, Any], list[Any]] | if not configured: if legacy_vars := configured_legacy_relay_env_vars(os.environ): logger.warning( - "Legacy NeMo Relay exporter variables are set but no %s was " - "provided. %s no longer activate Relay exporters; migrate the " - "exporter configuration to a Relay plugins.toml file.", - RELAY_PLUGINS_CONFIG_ENV, - ", ".join(legacy_vars), + "Legacy NeMo Relay exporter variables are set but no %s was provided. %s no longer activate " + "Relay exporters; migrate the exporter configuration to a Relay plugins.toml file.", + RELAY_PLUGINS_CONFIG_ENV, ", ".join(legacy_vars), ) return None config_path = Path(configured).expanduser() @@ -1266,17 +1203,12 @@ def _configured_plugin_inputs(relay: Any) -> tuple[dict[str, Any], list[Any]] | with config_path.open("rb") as config_file: config = tomllib.load(config_file) if "dynamic_plugins" in config: - raise ValueError( - "Hermes [[dynamic_plugins]] records are unsupported; use Relay [[plugins.dynamic]] records" - ) - dynamic_plugins: list[Any] = [] - if "plugins" in config: - dynamic_plugins = relay.plugin.load_dynamic_plugin_activation_specs(config_path) + raise ValueError("Hermes [[dynamic_plugins]] records are unsupported; use Relay [[plugins.dynamic]] records") + dynamic_plugins = relay.plugin.load_dynamic_plugin_activation_specs(config_path) if "plugins" in config else [] return {k: v for k, v in config.items() if k != "plugins"}, dynamic_plugins except Exception as exc: raise _RelayPluginConfigurationLoadError( - "Hermes Relay plugin configuration could not be loaded from " - f"{config_path}; continuing without Relay plugins" + f"Hermes Relay plugin configuration could not be loaded from {config_path}; continuing without Relay plugins" ) from exc @@ -1284,11 +1216,10 @@ def _resolve_plugin_awaitable(value: Any) -> Any: """Resolve Relay's async plugin API from synchronous host construction.""" if not inspect.isawaitable(value): return value - try: + with contextlib.suppress(RuntimeError): asyncio.get_running_loop() - except RuntimeError: - return asyncio.run(value) - return _run_on_daemon_thread(lambda: asyncio.run(value), name="hermes-nemo-relay-plugin-lifecycle") + return _run_on_daemon_thread(lambda: asyncio.run(value), name="hermes-nemo-relay-plugin-lifecycle") + return asyncio.run(value) def _session_id(event: dict[str, Any]) -> str: diff --git a/agent/relay_tools.py b/agent/relay_tools.py index 42f4e45bc7..54ccb74bd4 100644 --- a/agent/relay_tools.py +++ b/agent/relay_tools.py @@ -41,8 +41,7 @@ def execute( except BaseException as exc: callback_error = exc raise - raw_result["value"] = result - raw_result["json"] = _jsonable(result) + raw_result.update(value=result, json=_jsonable(result)) return raw_result["json"] try: diff --git a/agent/transports/anthropic.py b/agent/transports/anthropic.py index b00469d3af..c80ad07099 100644 --- a/agent/transports/anthropic.py +++ b/agent/transports/anthropic.py @@ -82,26 +82,18 @@ class AnthropicTransport(ProviderTransport): elif block.type in _THINKING_TYPES: if block.type == "thinking": reasoning_parts.append(block.thinking) - # Sanitized block preferred; raw only if sanitize dropped it. - if isinstance(clean_block, dict): - reasoning_details.append(clean_block) - elif isinstance(block_dict, dict): - reasoning_details.append(block_dict) + detail = clean_block if clean_block is not None else block_dict # raw only if sanitize dropped it + if isinstance(detail, dict): + reasoning_details.append(detail) elif block.type == "tool_use": name = block.name if strip_tool_prefix and name.startswith(_MCP_PREFIX): name = _unprefix_oauth_tool_name(name) tool_calls.append(ToolCall(id=block.id, name=name, arguments=json.dumps(block.input))) - provider_data = {} - if reasoning_details: - provider_data["reasoning_details"] = reasoning_details + provider_data = {"reasoning_details": reasoning_details} if reasoning_details else {} # Ordered channel only for the shape the parallel lists reconstruct wrongly. - kinds = {b.get("type") for b in ordered_blocks if isinstance(b, dict)} - signed = any( - b.get("type") in _THINKING_TYPES and (b.get("signature") or b.get("data")) - for b in ordered_blocks if isinstance(b, dict) - ) - if signed and "tool_use" in kinds: + signed = any(b.get("type") in _THINKING_TYPES and (b.get("signature") or b.get("data")) for b in ordered_blocks) + if signed and any(b.get("type") == "tool_use" for b in ordered_blocks): provider_data["anthropic_content_blocks"] = ordered_blocks return NormalizedResponse( content="\n".join(text_parts) if text_parts else None, tool_calls=tool_calls or None, @@ -114,9 +106,9 @@ class AnthropicTransport(ProviderTransport): """Structural check; empty content is legitimate for ``end_turn``/``refusal`` (retrying either would loop forever).""" content_blocks = getattr(response, "content", None) - if not isinstance(content_blocks, list): - return False - return bool(content_blocks) or getattr(response, "stop_reason", None) in {"end_turn", "refusal"} + return isinstance(content_blocks, list) and ( + bool(content_blocks) or getattr(response, "stop_reason", None) in {"end_turn", "refusal"} + ) def extract_cache_stats(self, response: Any) -> Optional[Dict[str, int]]: """Anthropic cache_read / cache_creation token counts."""