"""Direct NeMo Relay integration for Hermes shared client metrics.""" from __future__ import annotations import atexit import contextvars import logging import threading from dataclasses import dataclass, field from time import monotonic_ns from typing import Any, Callable from agent import relay_runtime from hermes_cli import __version__ from .shared_metrics import SharedMetricsStore from .shared_metrics_contract import ( MODEL_CALL_SCOPE, SCHEMA_KEY, SCHEMA_VERSION, SUBSCRIBER_NAME, TASK_SCOPE, model_call_fields, model_call_outcome, task_start_fields, task_terminal_fields, ) from .shared_metrics_subscriber import SharedMetricsSubscriber logger = logging.getLogger(__name__) HANDLED_HOOKS = frozenset({ "on_session_start", "on_session_end", "on_session_finalize", "on_session_reset", "pre_llm_call", "pre_api_request", "post_tool_call", "post_api_request", "api_request_error", "subagent_stop", }) _RUNTIME_FAILED = object() _RUNTIMES: dict[str, _Runtime | object] = {} _RUNTIME_LOCK = threading.RLock() def _retry_ordinal(event: dict[str, Any]) -> int | None: value = event.get("retry_count") if isinstance(value, int) and not isinstance(value, bool) and value >= 0: return value return None @dataclass class _ModelCall: handle: Any task_id: str fields: dict[str, str] retry_ordinal: int | None = None @dataclass class _TaskRun: handle: Any context: contextvars.Context started_ns: int start_fields: dict[str, str] model_call_ids: set[str] = field(default_factory=set) tool_call_ids: set[str] = field(default_factory=set) turn_ids: set[str] = field(default_factory=set) unidentified_tool_calls: int = 0 retry_count: int = 0 @dataclass class _MetricsSession: session_id: str relay_session: relay_runtime.RelaySession lock: threading.RLock = field(default_factory=threading.RLock, repr=False) closing: bool = False model_calls: dict[str, _ModelCall] = field(default_factory=dict) tasks: dict[str, _TaskRun] = field(default_factory=dict) class _Runtime: """Own shared-metrics state layered on the Hermes core Relay host.""" def __init__(self, host: relay_runtime.RelayRuntime | None = None) -> None: resolved_host = host or relay_runtime.get_runtime() if resolved_host is None: raise RuntimeError("Hermes core Relay runtime is unavailable") self.host: relay_runtime.RelayRuntime = resolved_host self.relay = self.host.relay self._sessions_lock = threading.RLock() self._active = True self._sessions: dict[str, _MetricsSession] = {} self._task_creation_lock = threading.RLock() self._task_sessions_lock = threading.RLock() self._task_sessions: dict[tuple[str, str], _MetricsSession] = {} self._turn_sessions: dict[tuple[str, str], _MetricsSession] = {} self._subscriber_name = f"{SUBSCRIBER_NAME}.{self.host.runtime_id}" self.subscriber = SharedMetricsSubscriber( SharedMetricsStore(), __version__, runtime_id=self.host.runtime_id, ) self.relay.subscribers.register(self._subscriber_name, self.subscriber) self.host.retain_managed_execution(self._subscriber_name) self._registered = True atexit.register(self.shutdown) def ensure_session(self, event: dict[str, Any]) -> _MetricsSession | None: session_id = str(event.get("session_id") or "") if not session_id: return None with self._sessions_lock: if not self._active: return None relay_session = self.host.ensure_session(event) if relay_session is None: return None session = self._sessions.get(session_id) if session is None: session = _MetricsSession( session_id=session_id, relay_session=relay_session, ) self._sessions[session_id] = session with session.lock: if session.closing: return None return session def _run_in_session( self, session: _MetricsSession, callback: Callable[..., Any], *args: Any, **kwargs: Any, ) -> Any: return self.host.run_in_session( session.relay_session, callback, *args, **kwargs, ) def start_task(self, event: dict[str, Any]) -> _TaskRun | None: """Open one Relay function scope for a Hermes task run.""" task_key = self._task_key(event) if task_key is None: return None _, task_id = task_key with self._task_creation_lock: owner = self._task_session(event) if owner is not None: with owner.lock: if owner.closing: return None task = owner.tasks.get(task_id) if task is not None: self._remember_turn(owner, task, event) return task session = self.ensure_session(event) if session is None: return None with session.lock: if session.closing or session.relay_session.context is None: return None task_context = session.relay_session.context.copy() start_fields = task_start_fields(event) active_turn = relay_runtime.active_turn(session.session_id) parent_handle = session.relay_session.handle if ( active_turn is not None and active_turn.lease.session_id == session.session_id and active_turn.task_id == task_id and active_turn.handle is not None ): parent_handle = active_turn.handle def push_task() -> Any: self.relay.get_scope_stack() return self.relay.scope.push( TASK_SCOPE, self.relay.ScopeType.Function, handle=parent_handle, input=start_fields, metadata=self._event_metadata(), ) handle = task_context.run(push_task) task = _TaskRun( handle=handle, context=task_context, started_ns=monotonic_ns(), start_fields=start_fields, ) session.tasks[task_id] = task with self._task_sessions_lock: self._task_sessions[task_key] = session self._remember_turn(session, task, event) return task def _run_in_task( self, task: _TaskRun, callback: Callable[..., Any], *args: Any, **kwargs: Any, ) -> Any: def invoke() -> Any: self.relay.get_scope_stack() return callback(*args, **kwargs) return task.context.copy().run(invoke) def start_model_call(self, event: dict[str, Any]) -> None: task_id = str(event.get("task_id") or "") session = self._task_session(event, allow_task_id_fallback=True) task = session.tasks.get(task_id) if session is not None else None if task is None: task = self.start_task(event) session = self._task_session(event) if task is not None else None if session is None: session = self.ensure_session(event) if session is None: return request_id = str(event.get("api_request_id") or "") if not request_id: return fields = model_call_fields(event) retry_ordinal = _retry_ordinal(event) model_family = fields["model_family"] with session.lock: if session.closing: return if task is not None: self._remember_turn(session, task, event) existing = session.model_calls.get(request_id) if existing is not None: existing.fields = fields if task is not None: if retry_ordinal is None or existing.retry_ordinal is None: task.retry_count += 1 elif retry_ordinal > existing.retry_ordinal: task.retry_count += retry_ordinal - existing.retry_ordinal if retry_ordinal is not None: existing.retry_ordinal = max( existing.retry_ordinal or 0, retry_ordinal, ) return if task is not None: task.model_call_ids.add(request_id) if retry_ordinal is not None and retry_ordinal > 0: # A real Hermes retry can advance api_request_id while # carrying the retry ordinal. Count that physical attempt. task.retry_count += 1 handle = self._run_in_task( task, self.relay.llm.call, MODEL_CALL_SCOPE, self.relay.LLMRequest({}, {}), handle=task.handle, metadata=self._event_metadata(), model_name=model_family, ) else: handle = self._run_in_session( session, self.relay.llm.call, MODEL_CALL_SCOPE, self.relay.LLMRequest({}, {}), handle=session.relay_session.handle, metadata=self._event_metadata(), model_name=model_family, ) session.model_calls[request_id] = _ModelCall( handle=handle, task_id=str(event.get("task_id") or ""), fields=fields, retry_ordinal=retry_ordinal, ) def record_tool_call(self, event: dict[str, Any]) -> None: """Count one unique tool invocation under its owning task.""" task_id = str(event.get("task_id") or "") session = self._task_session(event, allow_task_id_fallback=True) task = session.tasks.get(task_id) if session is not None else None if task is None: task = self.start_task(event) session = self._task_session(event) if task is not None else None if session is None or task is None: return tool_call_id = str(event.get("tool_call_id") or "") with session.lock: if session.closing: return self._remember_turn(session, task, event) if tool_call_id: task.tool_call_ids.add(tool_call_id) else: task.unidentified_tool_calls += 1 def end_model_call(self, event: dict[str, Any], outcome: str | None = None) -> None: session = self._task_session(event, allow_task_id_fallback=True) if session is None: session = self._session(event) if session is None: return request_id = str(event.get("api_request_id") or "") with session.lock: if session.closing: return model_call = session.model_calls.get(request_id) if model_call is None: return fields = model_call_fields(event) model_call.fields = fields self._finish_model_call( session, request_id, outcome or model_call_outcome(event), ) def end_pending_model_calls(self, event: dict[str, Any]) -> None: session = self._task_session(event, allow_task_id_fallback=True) if session is None: session = self._session(event) if session is None: return with session.lock: if session.closing: return self._end_pending_model_calls(session, event) def finish_task(self, event: dict[str, Any]) -> None: """Close one task scope exactly once with bounded terminal fields.""" task_id = str(event.get("task_id") or "") session = self._task_session( event, allow_task_id_fallback=True, ) or self._session(event) if session is None: return with session.lock: if session.closing: return finished = self._finish_task(session, task_id, event) if finished: try: self.relay.subscribers.flush() except Exception: logger.warning( "Hermes shared-metrics task flush failed", exc_info=True, ) else: self._export() def close_session(self, event: dict[str, Any]) -> None: session = self._session(event) if session is None: return failures: list[str] = [] with session.lock: if session.closing: return session.closing = True for task_id in list(session.tasks): self._finish_task( session, task_id, { **event, "task_id": task_id, "completed": False, "failed": True, "interrupted": False, "turn_exit_reason": "system_aborted", }, ) self._end_pending_model_calls(session, event) try: self.relay.subscribers.flush() except Exception as exc: failures.append(f"subscriber flush failed: {exc}") else: self._export() with self._sessions_lock: if self._sessions.get(session.session_id) is session: self._sessions.pop(session.session_id, None) if failures: logger.warning( "Hermes shared-metrics session %s closed with errors: %s", session.session_id, "; ".join(failures), ) def shutdown(self) -> None: with self._sessions_lock: self._active = False session_ids = list(self._sessions) for session_id in session_ids: self._safe(self.close_session, {"session_id": session_id}) if not self._registered: return try: self.relay.subscribers.flush() except Exception: logger.warning( "Hermes shared-metrics shutdown flush failed", exc_info=True, ) else: self._export() self._safe(self.relay.subscribers.deregister, self._subscriber_name) self.host.release_managed_execution(self._subscriber_name) self._registered = False try: atexit.unregister(self.shutdown) except Exception: pass def deactivate(self) -> None: """Stop collection without exporting locally aggregated metrics.""" with self._sessions_lock: self._active = False self.subscriber.deactivate() if self._registered: self._safe(self.relay.subscribers.deregister, self._subscriber_name) self.host.release_managed_execution(self._subscriber_name) self._registered = False with self._sessions_lock: sessions = list(self._sessions.values()) for session in sessions: with session.lock: if session.closing: continue session.closing = True for task_id in list(session.tasks): self._finish_task( session, task_id, { "session_id": session.session_id, "task_id": task_id, "failed": True, "turn_exit_reason": "system_aborted", }, ) self._end_pending_model_calls(session, {}) with self._sessions_lock: self._sessions.clear() with self._task_sessions_lock: self._task_sessions.clear() self._turn_sessions.clear() try: atexit.unregister(self.shutdown) except Exception: pass def _session(self, event: dict[str, Any]) -> _MetricsSession | None: session_id = str(event.get("session_id") or "") with self._sessions_lock: return self._sessions.get(session_id) @staticmethod def _task_key(event: dict[str, Any]) -> tuple[str, str] | None: session_id = str(event.get("session_id") or "") task_id = str(event.get("task_id") or "") if not session_id or not task_id: return None return session_id, task_id def _task_session( self, event: dict[str, Any], *, allow_task_id_fallback: bool = False, ) -> _MetricsSession | None: task_key = self._task_key(event) if task_key is None: return None turn_key = self._turn_key(event) with self._task_sessions_lock: if turn_key is not None: owner = self._turn_sessions.get(turn_key) if owner is not None: return owner owner = self._task_sessions.get(task_key) if owner is not None or not allow_task_id_fallback: return owner task_id = task_key[1] candidates: list[_MetricsSession] = [] for (_, candidate_task_id), session in self._task_sessions.items(): if candidate_task_id != task_id: continue if not any(candidate is session for candidate in candidates): candidates.append(session) return candidates[0] if len(candidates) == 1 else None @staticmethod def _turn_key(event: dict[str, Any]) -> tuple[str, str] | None: session_id = str(event.get("session_id") or "") turn_id = str(event.get("turn_id") or "") if not session_id or not turn_id: return None return session_id, turn_id def _remember_turn( self, session: _MetricsSession, task: _TaskRun, event: dict[str, Any], ) -> None: turn_id = str(event.get("turn_id") or "") if not turn_id: return task.turn_ids.add(turn_id) with self._task_sessions_lock: self._turn_sessions[(session.session_id, turn_id)] = session def _finish_model_call( self, session: _MetricsSession, request_id: str, outcome: str, ) -> None: model_call = session.model_calls.pop(request_id, None) if model_call is None: return try: task = session.tasks.get(model_call.task_id) if task is not None: self._run_in_task( task, self.relay.llm.call_end, model_call.handle, {**model_call.fields, "outcome": outcome}, metadata=self._event_metadata(), ) else: self._run_in_session( session, self.relay.llm.call_end, model_call.handle, {**model_call.fields, "outcome": outcome}, metadata=self._event_metadata(), ) except Exception: logger.warning( "Hermes shared-metrics model call close failed", exc_info=True ) def _end_pending_model_calls( self, session: _MetricsSession, event: dict[str, Any], ) -> None: task_id = str(event.get("task_id") or "") request_ids = [ request_id for request_id, model_call in session.model_calls.items() if not task_id or model_call.task_id == task_id ] outcome = "cancelled" if event.get("interrupted") else "failed" for request_id in request_ids: self._finish_model_call(session, request_id, outcome) def _finish_task( self, session: _MetricsSession, task_id: str, event: dict[str, Any], ) -> bool: task = session.tasks.get(task_id) if task is None: return False self._end_pending_model_calls(session, {**event, "task_id": task_id}) fields = task_terminal_fields( {**task.start_fields, **event}, duration_ms=max(0, (monotonic_ns() - task.started_ns) // 1_000_000), model_call_count=len(task.model_call_ids), tool_call_count=len(task.tool_call_ids) + task.unidentified_tool_calls, retry_count=task.retry_count, ) try: self._run_in_task( task, self.relay.scope.pop, task.handle, output=fields, metadata=self._event_metadata(), ) except Exception: logger.warning("Hermes shared-metrics task close failed", exc_info=True) finally: session.tasks.pop(task_id, None) with self._task_sessions_lock: task_key = (session.session_id, task_id) if self._task_sessions.get(task_key) is session: self._task_sessions.pop(task_key, None) for turn_id in task.turn_ids: turn_key = (session.session_id, turn_id) if self._turn_sessions.get(turn_key) is session: self._turn_sessions.pop(turn_key, None) return True def _export(self) -> None: self._safe(self.subscriber.store.create_and_export_package_if_due) def _event_metadata(self) -> dict[str, str]: return { SCHEMA_KEY: SCHEMA_VERSION, relay_runtime.RUNTIME_INSTANCE_KEY: self.host.runtime_id, } @staticmethod def _safe(callback: Callable[..., Any], *args: Any, **kwargs: Any) -> Any: try: return callback(*args, **kwargs) except Exception: logger.warning("Hermes shared metrics operation failed", exc_info=True) return None def enabled() -> bool: """Return the shared-metrics policy for the active Hermes profile.""" profile_key = relay_runtime.current_profile_key() try: from hermes_cli.config import read_raw_config_readonly # Collection consent is profile-owned. Managed config overlays may # control runtime policy, but cannot opt a profile into or out of # shared metrics. Read-only fast path: this gate runs 2-3x per agent # turn, and the mutable read_raw_config() paid a full config deepcopy # on every call. config = read_raw_config_readonly() or {} except Exception: logger.debug("Unable to read Hermes shared-metrics policy", exc_info=True) value = False else: telemetry = config.get("telemetry") if isinstance(config, dict) else None shared_metrics = ( telemetry.get("shared_metrics") if isinstance(telemetry, dict) else None ) value = ( isinstance(shared_metrics, dict) and shared_metrics.get("enabled") is True ) if value: return True with _RUNTIME_LOCK: runtime = _RUNTIMES.pop(profile_key, None) if isinstance(runtime, _Runtime): runtime.deactivate() return False def handles_hook(hook_name: str) -> bool: return hook_name in HANDLED_HOOKS and enabled() def observe_lifecycle(hook_name: str, **kwargs: Any) -> None: """Project one Hermes lifecycle event into the core Relay integration.""" if not handles_hook(hook_name): return runtime = _get_runtime() if runtime is None: return try: if hook_name == "on_session_start": runtime.ensure_session(kwargs) elif hook_name == "pre_llm_call": runtime.start_task(kwargs) elif hook_name == "pre_api_request": runtime.start_model_call(kwargs) elif hook_name == "post_tool_call": runtime.record_tool_call(kwargs) elif hook_name == "post_api_request": runtime.end_model_call(kwargs, "success") elif hook_name == "api_request_error": if kwargs.get("retryable") is False: runtime.end_model_call(kwargs, "failed") elif hook_name == "on_session_end": runtime.finish_task(kwargs) elif hook_name == "subagent_stop": child_session_id = str(kwargs.get("child_session_id") or "") if child_session_id: runtime.close_session({"session_id": child_session_id}) elif hook_name in {"on_session_finalize", "on_session_reset"}: runtime.close_session(kwargs) except Exception: logger.warning( "Hermes shared metrics hook failed: %s", hook_name, exc_info=True ) def prepare_session_start() -> None: """Register the subscriber before any producer opens the session scope.""" if enabled(): _get_runtime(retry_failed=True) def _prepare_core_session( host: relay_runtime.RelayRuntime, context: dict[str, Any], ) -> None: """Prepare the profile subscriber before the coordinator opens a scope.""" del context if host.profile_key == relay_runtime.current_profile_key(): if enabled(): _get_runtime(retry_failed=True, host=host) def start_task_run( *, session_id: str, task_id: str, platform: str, parent_session_id: str = "", ) -> None: """Start task metrics at the outer Hermes execution boundary.""" if not enabled(): return runtime = _get_runtime(retry_failed=True) if runtime is None: return runtime._safe( runtime.start_task, { "session_id": session_id, "task_id": task_id, "platform": platform, "parent_session_id": parent_session_id, }, ) def finish_task_run( *, session_id: str, task_id: str, platform: str, result: dict[str, Any] | None = None, error: BaseException | None = None, ) -> None: """Finish task metrics for every return or exception path.""" if not enabled(): return runtime = _get_runtime() if runtime is None: return terminal = result if isinstance(result, dict) else {} interrupted = terminal.get("interrupted") is True completed = terminal.get("completed") is True failed = terminal.get("failed") is True reason = str( terminal.get("turn_exit_reason") or terminal.get("failure_reason") or "" ) if error is not None: interrupted = isinstance(error, (KeyboardInterrupt, InterruptedError)) or ( type(error).__name__ == "CancelledError" ) timed_out = isinstance(error, TimeoutError) completed = False failed = not interrupted if interrupted: reason = "interrupted_by_user" elif timed_out: reason = "timed_out" else: reason = "system_aborted" elif not reason: reason = "failed" if failed else "unknown" runtime._safe( runtime.finish_task, { "session_id": session_id, "task_id": task_id, "platform": platform, "completed": completed, "failed": failed, "interrupted": interrupted, "turn_exit_reason": reason, }, ) def _get_runtime( *, retry_failed: bool = False, host: relay_runtime.RelayRuntime | None = None, ) -> _Runtime | None: profile_key = relay_runtime.current_profile_key() with _RUNTIME_LOCK: runtime = _RUNTIMES.get(profile_key) if isinstance(runtime, _Runtime): if host is None or runtime.host is host: return runtime runtime.deactivate() _RUNTIMES.pop(profile_key, None) if runtime is _RUNTIME_FAILED and not retry_failed: return None if runtime is _RUNTIME_FAILED: _RUNTIMES.pop(profile_key, None) try: runtime = _Runtime(host=host) except Exception: logger.warning("Hermes shared metrics initialization failed", exc_info=True) _RUNTIMES[profile_key] = _RUNTIME_FAILED return None _RUNTIMES[profile_key] = runtime return runtime relay_runtime.SESSION_COORDINATOR.register_session_initializer( SUBSCRIBER_NAME, _prepare_core_session, ) def _reset_for_tests() -> None: """Reset all profile-scoped shared-metrics state for isolated tests.""" with _RUNTIME_LOCK: runtimes = list(_RUNTIMES.values()) _RUNTIMES.clear() for runtime in runtimes: if isinstance(runtime, _Runtime): runtime.shutdown()