diff --git a/agent/agent_runtime_helpers.py b/agent/agent_runtime_helpers.py index d6522d07f488..092a87da95ef 100644 --- a/agent/agent_runtime_helpers.py +++ b/agent/agent_runtime_helpers.py @@ -2475,7 +2475,8 @@ def invoke_tool(agent, function_name: str, function_args: dict, effective_task_i tool_call_id: Optional[str] = None, messages: list = None, pre_tool_block_checked: bool = False, skip_tool_request_middleware: bool = False, - tool_request_middleware_trace: Optional[List[Dict[str, Any]]] = None) -> str: + tool_request_middleware_trace: Optional[List[Dict[str, Any]]] = None, + skip_tool_execution_middleware: bool = False) -> str: """Invoke a single tool and return the result string. No display logic. Handles both agent-level tools (todo, memory, etc.) and registry-dispatched @@ -2654,8 +2655,7 @@ def invoke_tool(agent, function_name: str, function_args: dict, effective_task_i return _finish_agent_tool(agent._dispatch_delegate_task(next_args), next_args) else: def _execute(next_args: dict) -> Any: - return _ra().handle_function_call( - function_name, next_args, effective_task_id, + dispatch_kwargs = dict( tool_call_id=tool_call_id, session_id=agent.session_id or "", turn_id=getattr(agent, "_current_turn_id", "") or "", @@ -2667,6 +2667,17 @@ def invoke_tool(agent, function_name: str, function_args: dict, effective_task_i disabled_toolsets=getattr(agent, "disabled_toolsets", None), tool_request_middleware_trace=list(_tool_middleware_trace), ) + if skip_tool_execution_middleware: + dispatch_kwargs["skip_tool_execution_middleware"] = True + return _ra().handle_function_call( + function_name, + next_args, + effective_task_id, + **dispatch_kwargs, + ) + + if skip_tool_execution_middleware: + return _execute(function_args) from hermes_cli.middleware import run_tool_execution_middleware diff --git a/agent/auxiliary_client.py b/agent/auxiliary_client.py index fa5bc6b80d7e..da85f03659e8 100644 --- a/agent/auxiliary_client.py +++ b/agent/auxiliary_client.py @@ -42,6 +42,7 @@ Payment / credit exhaustion fallback: import contextlib import contextvars +import functools import hashlib import inspect import json @@ -50,9 +51,10 @@ import os import re import threading import time +import uuid from pathlib import Path # noqa: F401 — used by test mocks from types import SimpleNamespace -from typing import Any, Dict, List, Optional, Tuple, TYPE_CHECKING +from typing import Any, Callable, Dict, List, Optional, Tuple, TYPE_CHECKING from urllib.parse import urlparse, parse_qs, urlunparse # NOTE: `from openai import OpenAI` is deliberately NOT at module top — the @@ -2492,6 +2494,167 @@ _RUNTIME_MAIN_AUTH_MODE: str = "" _RUNTIME_MAIN_CONTEXT: contextvars.ContextVar[Optional[Dict[str, Any]]] = ( contextvars.ContextVar("auxiliary_runtime_main", default=None) ) + +_RELAY_AUX_CALL_CONTEXT: contextvars.ContextVar[Optional[Dict[str, Any]]] = ( + contextvars.ContextVar("auxiliary_relay_call", default=None) +) + + +def _relay_auxiliary_call(callback): + """Give every physical retry in one auxiliary call a shared Relay identity.""" + + @functools.wraps(callback) + def wrapped(*args, **kwargs): + task = args[0] if args else kwargs.get("task") + token = _RELAY_AUX_CALL_CONTEXT.set({ + "task": str(task or "unknown"), + "request_id": f"aux-{uuid.uuid4().hex}", + "attempt_count": 0, + "provider": "", + "model": "", + "api_mode": "chat_completions", + }) + try: + return callback(*args, **kwargs) + except BaseException: + _fail_relay_auxiliary_call() + raise + finally: + _RELAY_AUX_CALL_CONTEXT.reset(token) + + return wrapped + + +def _relay_auxiliary_call_async(callback): + """Async counterpart to :func:`_relay_auxiliary_call`.""" + + @functools.wraps(callback) + async def wrapped(*args, **kwargs): + task = args[0] if args else kwargs.get("task") + token = _RELAY_AUX_CALL_CONTEXT.set({ + "task": str(task or "unknown"), + "request_id": f"aux-{uuid.uuid4().hex}", + "attempt_count": 0, + "provider": "", + "model": "", + "api_mode": "chat_completions", + }) + try: + return await callback(*args, **kwargs) + except BaseException: + _fail_relay_auxiliary_call() + raise + finally: + _RELAY_AUX_CALL_CONTEXT.reset(token) + + return wrapped + + +def _set_relay_auxiliary_route( + provider: str | None, + model: str | None, + api_mode: str | None, +) -> None: + context = _RELAY_AUX_CALL_CONTEXT.get() + if context is None: + return + context["provider"] = str(provider or "auxiliary") + context["model"] = str(model or "unknown") + context["api_mode"] = str(api_mode or "chat_completions") + + +def _relay_auxiliary_metadata( + *, + provider: str | None = None, + api_mode: str | None = None, +) -> tuple[str, str, dict[str, Any]] | None: + context = _RELAY_AUX_CALL_CONTEXT.get() + if context is None: + return None + attempt_count = int(context.get("attempt_count") or 0) + context["attempt_count"] = attempt_count + 1 + provider_name = str(provider or context.get("provider") or "auxiliary") + model_name = str(context.get("model") or "unknown") + return provider_name, model_name, { + "api_mode": str(api_mode or context.get("api_mode") or "chat_completions"), + "api_request_id": str(context["request_id"]), + "call_role": f"auxiliary:{context['task']}", + "retry_count": attempt_count, + "auxiliary_task": str(context["task"]), + } + + +def _relay_sync_completion( + client: Any, + kwargs: dict[str, Any], + *, + provider: str | None = None, + api_mode: str | None = None, + create: Callable[[dict[str, Any]], Any] | None = None, +) -> Any: + callback = create or (lambda request: client.chat.completions.create(**request)) + route = _relay_auxiliary_metadata(provider=provider, api_mode=api_mode) + if route is None: + return callback(kwargs) + provider_name, fallback_model, metadata = route + from agent import relay_llm + + return relay_llm.execute_current( + kwargs, + callback, + name=provider_name, + model_name=str(kwargs.get("model") or fallback_model), + metadata=metadata, + defer_logical_completion=True, + ) + + +async def _relay_async_completion( + client: Any, + kwargs: dict[str, Any], + *, + provider: str | None = None, + api_mode: str | None = None, + create: Callable[[dict[str, Any]], Any] | None = None, +) -> Any: + callback = create or (lambda request: client.chat.completions.create(**request)) + route = _relay_auxiliary_metadata(provider=provider, api_mode=api_mode) + if route is None: + return await callback(kwargs) + provider_name, fallback_model, metadata = route + from agent import relay_llm + + return await relay_llm.execute_current_async( + kwargs, + callback, + name=provider_name, + model_name=str(kwargs.get("model") or fallback_model), + metadata=metadata, + defer_logical_completion=True, + ) + + +def _relay_sync_stream( + client: Any, + kwargs: dict[str, Any], + *, + provider: str | None = None, + api_mode: str | None = None, +) -> Any: + route = _relay_auxiliary_metadata(provider=provider, api_mode=api_mode) + if route is None: + return client.chat.completions.create(**kwargs) + provider_name, fallback_model, metadata = route + from agent import relay_llm + + return relay_llm.stream_current( + kwargs, + lambda request: client.chat.completions.create(**request), + name=provider_name, + model_name=str(kwargs.get("model") or fallback_model), + finalizer=dict, + metadata=metadata, + ) _RUNTIME_MAIN_COMPAT_SNAPSHOT: Tuple[Any, ...] = ("", "", "", "", "", "") _RUNTIME_MAIN_COMPAT_LOCK = threading.Lock() @@ -3801,7 +3964,13 @@ def _retry_same_provider_sync( if _is_anthropic_compat_endpoint(resolved_provider, retry_base): retry_kwargs["messages"] = _convert_openai_images_to_anthropic(retry_kwargs["messages"]) return _validate_llm_response( - retry_client.chat.completions.create(**retry_kwargs), task, + _relay_sync_completion( + retry_client, + retry_kwargs, + provider=resolved_provider, + api_mode=resolved_api_mode, + ), + task, ) @@ -3866,7 +4035,13 @@ async def _retry_same_provider_async( if _is_anthropic_compat_endpoint(resolved_provider, retry_base): retry_kwargs["messages"] = _convert_openai_images_to_anthropic(retry_kwargs["messages"]) return _validate_llm_response( - await retry_client.chat.completions.create(**retry_kwargs), task, + await _relay_async_completion( + retry_client, + retry_kwargs, + provider=resolved_provider, + api_mode=resolved_api_mode, + ), + task, ) @@ -4071,7 +4246,7 @@ def _call_fallback_candidate_sync( base_url=fb_base, task=task) try: return _validate_llm_response( - fb_client.chat.completions.create(**fb_kwargs), task) + _relay_sync_completion(fb_client, fb_kwargs, provider=fb_label), task) except Exception as fb_err: if not _is_auth_error(fb_err): raise @@ -4088,7 +4263,13 @@ def _call_fallback_candidate_sync( base_url=str(getattr(retry_client, "base_url", "") or fb_base), task=task) try: return _validate_llm_response( - retry_client.chat.completions.create(**retry_kwargs), task) + _relay_sync_completion( + retry_client, + retry_kwargs, + provider=fb_provider, + ), + task, + ) except Exception as retry_err: if not _is_auth_error(retry_err): raise @@ -4137,7 +4318,13 @@ async def _call_fallback_candidate_async( base_url=fb_base, task=task) try: return _validate_llm_response( - await fb_client.chat.completions.create(**fb_kwargs), task) + await _relay_async_completion( + fb_client, + fb_kwargs, + provider=fb_label, + ), + task, + ) except Exception as fb_err: if not _is_auth_error(fb_err): raise @@ -4155,7 +4342,13 @@ async def _call_fallback_candidate_async( base_url=str(getattr(retry_client, "base_url", "") or fb_base), task=task) try: return _validate_llm_response( - await retry_client.chat.completions.create(**retry_kwargs), task) + await _relay_async_completion( + retry_client, + retry_kwargs, + provider=fb_provider, + ), + task, + ) except Exception as retry_err: if not _is_auth_error(retry_err): raise @@ -7270,6 +7463,7 @@ def _validate_llm_response( except (AttributeError, TypeError, IndexError) as exc: recovered = _recover_aux_response_message(response) if recovered is not None: + _complete_relay_auxiliary_call() return recovered response_type = type(response).__name__ response_preview = str(response)[:120] @@ -7279,9 +7473,34 @@ def _validate_llm_response( f"Expected object with .choices[0].message — check provider " f"adapter or custom endpoint compatibility." ) from exc + _complete_relay_auxiliary_call() return response +def _complete_relay_auxiliary_call(*, outcome: str = "success") -> None: + """Close one auxiliary logical call after acceptance or terminal failure.""" + context = _RELAY_AUX_CALL_CONTEXT.get() + if context is None: + return + from agent import relay_llm + + relay_llm.complete_logical_call( + str(context.get("request_id") or ""), + outcome=outcome, + ) + + +def _fail_relay_auxiliary_call() -> None: + """Close a terminally failed call without replacing its original error.""" + try: + _complete_relay_auxiliary_call(outcome="failed") + except Exception: + logger.warning( + "Relay auxiliary failure finalization failed", + exc_info=True, + ) + + def _recover_aux_response_message(response: Any) -> Optional[Any]: """Synthesize chat-completions shape from Responses-style text fields. @@ -7680,6 +7899,7 @@ async def _acreate_with_stream( ) +@_relay_auxiliary_call def call_llm( task: str = None, *, @@ -7820,6 +8040,11 @@ def call_llm( f"Run: hermes setup") effective_timeout = _effective_aux_timeout(task, timeout) + _set_relay_auxiliary_route( + resolved_provider, + final_model, + resolved_api_mode, + ) # Log what we're about to do — makes auxiliary operations visible _base_info = str(getattr(client, "base_url", resolved_base_url) or "") @@ -7857,7 +8082,12 @@ def call_llm( kwargs["stream"] = True if stream_options: kwargs["stream_options"] = stream_options - return client.chat.completions.create(**kwargs) + return _relay_sync_stream( + client, + kwargs, + provider=resolved_provider, + api_mode=resolved_api_mode, + ) # Handle unsupported temperature, max_tokens vs max_completion_tokens retry, # then payment fallback. @@ -7880,10 +8110,18 @@ def call_llm( # for the transient retry every auxiliary task shares. (PR #16587) try: return _validate_llm_response( - _create_with_progress( - client, kwargs, task, - force_stream=_provider_requires_stream( - resolved_provider, _base_info or resolved_base_url, + _relay_sync_completion( + client, + kwargs, + provider=resolved_provider, + api_mode=resolved_api_mode, + create=lambda request: _create_with_progress( + client, + request, + task, + force_stream=_provider_requires_stream( + resolved_provider, _base_info or resolved_base_url, + ), ), ), task, @@ -7919,10 +8157,19 @@ def call_llm( time.sleep(_backoff) try: return _validate_llm_response( - _create_with_progress( - client, kwargs, task, - force_stream=_provider_requires_stream( - resolved_provider, _base_info or resolved_base_url, + _relay_sync_completion( + client, + kwargs, + provider=resolved_provider, + api_mode=resolved_api_mode, + create=lambda request: _create_with_progress( + client, + request, + task, + force_stream=_provider_requires_stream( + resolved_provider, + _base_info or resolved_base_url, + ), ), ), task) @@ -7942,7 +8189,12 @@ def call_llm( ) try: return _validate_llm_response( - client.chat.completions.create(**retry_kwargs), task) + _relay_sync_completion( + client, + retry_kwargs, + provider=resolved_provider, + api_mode=resolved_api_mode, + ), task) except Exception as retry_err: retry_err_str = str(retry_err) # If retry still fails, fall through to the max_tokens / @@ -7980,7 +8232,12 @@ def call_llm( kwargs.pop("max_completion_tokens", None) try: return _validate_llm_response( - client.chat.completions.create(**kwargs), task) + _relay_sync_completion( + client, + kwargs, + provider=resolved_provider, + api_mode=resolved_api_mode, + ), task) except Exception as retry_err: # If the max_tokens retry also hits a payment or connection # error, fall through to the fallback chain below. @@ -8010,7 +8267,12 @@ def call_llm( kwargs["model"] = healed_model try: return _validate_llm_response( - client.chat.completions.create(**kwargs), task) + _relay_sync_completion( + client, + kwargs, + provider=resolved_provider, + api_mode=resolved_api_mode, + ), task) except Exception as retry_err: first_err = retry_err @@ -8043,7 +8305,12 @@ def call_llm( kwargs["model"] = refreshed_model try: return _validate_llm_response( - refreshed_client.chat.completions.create(**kwargs), task) + _relay_sync_completion( + refreshed_client, + kwargs, + provider=resolved_provider, + api_mode=resolved_api_mode, + ), task) except Exception as retry_err: if not ( _is_auth_error(retry_err) @@ -8071,7 +8338,12 @@ def call_llm( if refreshed_model and refreshed_model != kwargs.get("model"): kwargs["model"] = refreshed_model return _validate_llm_response( - refreshed_client.chat.completions.create(**kwargs), task) + _relay_sync_completion( + refreshed_client, + kwargs, + provider=resolved_provider, + api_mode=resolved_api_mode, + ), task) # ── Auth refresh retry ─────────────────────────────────────── auth_refresh_provider = _auth_refresh_provider_for_route( @@ -8121,7 +8393,12 @@ def call_llm( if _is_rate_limit_error(first_err) and not _is_payment_error(first_err): try: return _validate_llm_response( - client.chat.completions.create(**kwargs), task) + _relay_sync_completion( + client, + kwargs, + provider=resolved_provider, + api_mode=resolved_api_mode, + ), task) except Exception as retry_err: if not (_is_auth_error(retry_err) or _is_payment_error(retry_err) or _is_rate_limit_error(retry_err)): raise @@ -8379,6 +8656,7 @@ def extract_content_or_reasoning(response) -> str: return "" +@_relay_auxiliary_call_async async def async_call_llm( task: str = None, *, @@ -8470,6 +8748,11 @@ async def async_call_llm( f"Run: hermes setup") effective_timeout = _effective_aux_timeout(task, timeout) + _set_relay_auxiliary_route( + resolved_provider, + final_model, + resolved_api_mode, + ) # Pass the client's actual base_url (not just resolved_base_url) so # endpoint-specific temperature overrides can distinguish @@ -8508,7 +8791,14 @@ async def async_call_llm( try: return _validate_llm_response( - await _acreate(kwargs), task, + await _relay_async_completion( + client, + kwargs, + provider=resolved_provider, + api_mode=resolved_api_mode, + create=_acreate, + ), + task, provider=resolved_provider, base_url=_client_base) except Exception as transient_err: if not _is_transient_transport_error(transient_err): @@ -8529,7 +8819,14 @@ async def async_call_llm( task or "call", transient_err, ) return _validate_llm_response( - await _acreate(kwargs), task) + await _relay_async_completion( + client, + kwargs, + provider=resolved_provider, + api_mode=resolved_api_mode, + create=_acreate, + ), + task) except Exception as first_err: if "temperature" in kwargs and _is_unsupported_temperature_error(first_err): retry_kwargs = dict(kwargs) @@ -8540,7 +8837,12 @@ async def async_call_llm( ) try: return _validate_llm_response( - await client.chat.completions.create(**retry_kwargs), task) + await _relay_async_completion( + client, + retry_kwargs, + provider=resolved_provider, + api_mode=resolved_api_mode, + ), task) except Exception as retry_err: retry_err_str = str(retry_err) if not ( @@ -8574,7 +8876,12 @@ async def async_call_llm( kwargs.pop("max_completion_tokens", None) try: return _validate_llm_response( - await client.chat.completions.create(**kwargs), task) + await _relay_async_completion( + client, + kwargs, + provider=resolved_provider, + api_mode=resolved_api_mode, + ), task) except Exception as retry_err: # If the max_tokens retry also hits a payment or connection # error, fall through to the fallback chain below. @@ -8603,7 +8910,12 @@ async def async_call_llm( kwargs["model"] = healed_model try: return _validate_llm_response( - await client.chat.completions.create(**kwargs), task) + await _relay_async_completion( + client, + kwargs, + provider=resolved_provider, + api_mode=resolved_api_mode, + ), task) except Exception as retry_err: first_err = retry_err @@ -8635,7 +8947,12 @@ async def async_call_llm( kwargs["model"] = refreshed_model try: return _validate_llm_response( - await refreshed_client.chat.completions.create(**kwargs), task) + await _relay_async_completion( + refreshed_client, + kwargs, + provider=resolved_provider, + api_mode=resolved_api_mode, + ), task) except Exception as retry_err: if not ( _is_auth_error(retry_err) @@ -8662,7 +8979,12 @@ async def async_call_llm( if refreshed_model and refreshed_model != kwargs.get("model"): kwargs["model"] = refreshed_model return _validate_llm_response( - await refreshed_client.chat.completions.create(**kwargs), task) + await _relay_async_completion( + refreshed_client, + kwargs, + provider=resolved_provider, + api_mode=resolved_api_mode, + ), task) # ── Auth refresh retry (mirrors sync call_llm) ─────────────── auth_refresh_provider = _auth_refresh_provider_for_route( @@ -8706,7 +9028,12 @@ async def async_call_llm( if _is_rate_limit_error(first_err) and not _is_payment_error(first_err): try: return _validate_llm_response( - await client.chat.completions.create(**kwargs), task) + await _relay_async_completion( + client, + kwargs, + provider=resolved_provider, + api_mode=resolved_api_mode, + ), task) except Exception as retry_err: if not (_is_auth_error(retry_err) or _is_payment_error(retry_err) or _is_rate_limit_error(retry_err)): raise diff --git a/agent/chat_completion_helpers.py b/agent/chat_completion_helpers.py index f4bca4fa1755..fac42cc3cf78 100644 --- a/agent/chat_completion_helpers.py +++ b/agent/chat_completion_helpers.py @@ -15,6 +15,7 @@ sites unchanged. Symbols that tests patch on ``run_agent`` (e.g. from __future__ import annotations +import contextvars import json import logging import math @@ -57,6 +58,12 @@ _OPENROUTER_PROVIDER_SORT_VALUES = {"throughput", "latency", "price"} _FALLBACK_EXHAUSTED_COOLDOWN_S = 5.0 +def _context_thread_target(callback): + """Bind a no-argument thread target to the caller's ContextVars.""" + context = contextvars.copy_context() + return lambda: context.run(callback) + + def _ra(): """Lazy ``run_agent`` reference. @@ -840,7 +847,7 @@ def interruptible_api_call(agent, api_kwargs: dict): _call_start = time.time() agent._touch_activity("waiting for non-streaming API response") - t = threading.Thread(target=_call, daemon=True) + t = threading.Thread(target=_context_thread_target(_call), daemon=True) t.start() _poll_count = 0 while t.is_alive(): @@ -2019,6 +2026,28 @@ def handle_max_iterations(agent, messages: list, api_call_count: int) -> str: """Request a summary when max iterations are reached. Returns the final response text.""" print(f"⚠️ Reached maximum iterations ({agent.max_iterations}). Requesting summary...") + summary_api_request_id = f"iteration-summary:{uuid.uuid4()}" + summary_call_outcome = "failed" + + def _managed_summary_call(request, callback, *, retry_count: int): + from agent import relay_llm + + return relay_llm.execute_current( + request, + callback, + name=str(getattr(agent, "provider", "") or "provider"), + model_name=str(getattr(agent, "model", "") or ""), + metadata={ + "api_mode": str( + getattr(agent, "api_mode", "") or "chat_completions" + ), + "api_request_id": summary_api_request_id, + "call_role": "iteration_summary", + "retry_count": retry_count, + }, + defer_logical_completion=True, + ) + summary_request = ( "You've reached the maximum number of tool-calling iterations allowed. " "Please provide a final response summarizing what you've found and accomplished so far, " @@ -2204,17 +2233,33 @@ def handle_max_iterations(agent, messages: list, api_call_count: int) -> str: if agent.api_mode == "anthropic_messages": _tsum = agent._get_transport() - _ant_kw = _tsum.build_kwargs(model=agent.model, messages=api_messages, tools=None, - max_tokens=agent.max_tokens, reasoning_config=agent.reasoning_config, - is_oauth=agent._is_anthropic_oauth, - preserve_dots=agent._anthropic_preserve_dots(), - base_url=getattr(agent, "_anthropic_base_url", None)) + _ant_kw = _tsum.build_kwargs( + model=agent.model, + messages=api_messages, + tools=None, + max_tokens=agent.max_tokens, + reasoning_config=agent.reasoning_config, + is_oauth=agent._is_anthropic_oauth, + preserve_dots=agent._anthropic_preserve_dots(), + base_url=getattr(agent, "_anthropic_base_url", None), + ) _ant_kw = _merge_nous_portal_messages_extra_body(agent, _ant_kw) - summary_response = agent._anthropic_messages_create(_ant_kw) + summary_response = _managed_summary_call( + _ant_kw, + agent._anthropic_messages_create, + retry_count=0, + ) _summary_result = _tsum.normalize_response(summary_response, strip_tool_prefix=agent._is_anthropic_oauth) final_response = (_summary_result.content or "").strip() else: - summary_response = agent._ensure_primary_openai_client(reason="iteration_limit_summary").chat.completions.create(**summary_kwargs) + summary_client = agent._ensure_primary_openai_client( + reason="iteration_limit_summary" + ) + summary_response = _managed_summary_call( + summary_kwargs, + lambda request: summary_client.chat.completions.create(**request), + retry_count=0, + ) _summary_result = agent._get_transport().normalize_response(summary_response) final_response = (_summary_result.content or "").strip() @@ -2222,6 +2267,7 @@ def handle_max_iterations(agent, messages: list, api_call_count: int) -> str: if "" in final_response: final_response = re.sub(r'.*?\s*', '', final_response, flags=re.DOTALL).strip() if final_response: + summary_call_outcome = "success" messages.append({"role": "assistant", "content": final_response}) else: final_response = "I reached the iteration limit and couldn't generate a summary." @@ -2236,13 +2282,22 @@ def handle_max_iterations(agent, messages: list, api_call_count: int) -> str: final_response = (_cnr_retry.content or "").strip() elif agent.api_mode == "anthropic_messages": _tretry = agent._get_transport() - _ant_kw2 = _tretry.build_kwargs(model=agent.model, messages=api_messages, tools=None, - is_oauth=agent._is_anthropic_oauth, - max_tokens=agent.max_tokens, reasoning_config=agent.reasoning_config, - preserve_dots=agent._anthropic_preserve_dots(), - base_url=getattr(agent, "_anthropic_base_url", None)) + _ant_kw2 = _tretry.build_kwargs( + model=agent.model, + messages=api_messages, + tools=None, + is_oauth=agent._is_anthropic_oauth, + max_tokens=agent.max_tokens, + reasoning_config=agent.reasoning_config, + preserve_dots=agent._anthropic_preserve_dots(), + base_url=getattr(agent, "_anthropic_base_url", None), + ) _ant_kw2 = _merge_nous_portal_messages_extra_body(agent, _ant_kw2) - retry_response = agent._anthropic_messages_create(_ant_kw2) + retry_response = _managed_summary_call( + _ant_kw2, + agent._anthropic_messages_create, + retry_count=1, + ) _retry_result = _tretry.normalize_response(retry_response, strip_tool_prefix=agent._is_anthropic_oauth) final_response = (_retry_result.content or "").strip() else: @@ -2259,7 +2314,14 @@ def handle_max_iterations(agent, messages: list, api_call_count: int) -> str: if summary_extra_body: summary_kwargs["extra_body"] = summary_extra_body - summary_response = agent._ensure_primary_openai_client(reason="iteration_limit_summary_retry").chat.completions.create(**summary_kwargs) + summary_client = agent._ensure_primary_openai_client( + reason="iteration_limit_summary_retry" + ) + summary_response = _managed_summary_call( + summary_kwargs, + lambda request: summary_client.chat.completions.create(**request), + retry_count=1, + ) _retry_result = agent._get_transport().normalize_response(summary_response) final_response = (_retry_result.content or "").strip() @@ -2267,6 +2329,7 @@ def handle_max_iterations(agent, messages: list, api_call_count: int) -> str: if "" in final_response: final_response = re.sub(r'.*?\s*', '', final_response, flags=re.DOTALL).strip() if final_response: + summary_call_outcome = "success" messages.append({"role": "assistant", "content": final_response}) else: final_response = "I reached the iteration limit and couldn't generate a summary." @@ -2276,6 +2339,13 @@ def handle_max_iterations(agent, messages: list, api_call_count: int) -> str: except Exception as e: logger.warning(f"Failed to get summary response: {e}") final_response = f"I reached the maximum iterations ({agent.max_iterations}) but couldn't summarize. Error: {str(e)}" + finally: + from agent import relay_llm + + relay_llm.complete_logical_call( + summary_api_request_id, + outcome=summary_call_outcome, + ) return final_response @@ -2434,7 +2504,9 @@ def interruptible_streaming_api_call(agent, api_kwargs: dict, *, on_first_delta= pass def _bedrock_call(): + stream = None try: + from agent import relay_llm from agent.bedrock_adapter import ( _get_bedrock_runtime_client, invalidate_runtime_client, @@ -2443,44 +2515,40 @@ def interruptible_streaming_api_call(agent, api_kwargs: dict, *, on_first_delta= normalize_converse_response, stream_converse_with_callbacks, ) - region = api_kwargs.pop("__bedrock_region__", "us-east-1") - api_kwargs.pop("__bedrock_converse__", None) - client = _get_bedrock_runtime_client(region) - try: - raw_response = client.converse_stream(**api_kwargs) - except Exception as _bedrock_exc: - # IAM policies scoped to bedrock:InvokeModel only (no - # InvokeModelWithResponseStream) reject converse_stream() - # with AccessDeniedException. That denial is permanent for - # the session — fall back to the non-streaming converse() - # inline (it maps to bedrock:InvokeModel) and disable - # streaming for subsequent calls so we don't re-fail every - # turn. - if is_streaming_access_denied_error(_bedrock_exc): - agent._disable_streaming = True - agent._safe_print( - "\n⚠ AWS IAM denied bedrock:InvokeModelWithResponseStream — " - "falling back to non-streaming InvokeModel.\n" - " Grant that action to restore streaming output.\n" - ) - logger.info( - "bedrock: converse_stream denied by IAM (%s) — " - "using non-streaming converse() for this session.", - type(_bedrock_exc).__name__, - ) - result["response"] = normalize_converse_response( - client.converse(**api_kwargs) - ) - return - # Evict the cached client on stale-connection failures - # so the outer retry loop builds a fresh client/pool. - if is_stale_connection_error(_bedrock_exc): - invalidate_runtime_client(region) - raise + intercepted_events = [] + writer_token = {"value": None} - # Claim the delta sink for this bedrock stream (#65991) so a - # superseded attempt's callbacks are fenced by the sink guard. - claim_stream_writer(agent) + def _open_bedrock_stream(next_api_kwargs: dict[str, Any]): + final_kwargs = dict(next_api_kwargs) + region = final_kwargs.pop("__bedrock_region__", "us-east-1") + final_kwargs.pop("__bedrock_converse__", None) + client = _get_bedrock_runtime_client(region) + try: + raw_response = client.converse_stream(**final_kwargs) + except Exception as _bedrock_exc: + # InvokeModel-only policies cannot open a stream. Keep + # the fallback inside the same managed Relay attempt so + # the real provider request and terminal response still + # share one lifecycle boundary. + if is_streaming_access_denied_error(_bedrock_exc): + agent._disable_streaming = True + agent._safe_print( + "\n⚠ AWS IAM denied bedrock:InvokeModelWithResponseStream — " + "falling back to non-streaming InvokeModel.\n" + " Grant that action to restore streaming output.\n" + ) + logger.info( + "bedrock: converse_stream denied by IAM (%s) — " + "using non-streaming converse() for this session.", + type(_bedrock_exc).__name__, + ) + return normalize_converse_response( + client.converse(**final_kwargs) + ) + if is_stale_connection_error(_bedrock_exc): + invalidate_runtime_client(region) + raise + return raw_response.get("stream", []) def _on_text(text): _fire_first() @@ -2495,18 +2563,65 @@ def interruptible_streaming_api_call(agent, api_kwargs: dict, *, on_first_delta= _fire_first() agent._fire_reasoning_delta(text) - result["response"] = stream_converse_with_callbacks( - raw_response, + def _finalize_bedrock_stream(): + return stream_converse_with_callbacks( + {"stream": list(intercepted_events)} + ) + + def _bedrock_stream_created(_stream: Any) -> None: + writer_token["value"] = claim_stream_writer(agent) + + def _accept_bedrock_event(_event: Any) -> bool: + token = writer_token["value"] + return token is None or stream_writer_is_current(agent, token) + + stream = relay_llm.stream( + dict(api_kwargs), + _open_bedrock_stream, + session_id=str(getattr(agent, "session_id", "") or ""), + name=str(getattr(agent, "provider", "") or "bedrock"), + model_name=str(getattr(agent, "model", "") or ""), + finalizer=_finalize_bedrock_stream, + on_stream_created=_bedrock_stream_created, + on_chunk=intercepted_events.append, + chunk_adapter=lambda chunk: chunk, + accept_chunk=_accept_bedrock_event, + completed_response_predicate=lambda response: bool( + getattr(response, "choices", None) + ), + metadata={ + "api_mode": "custom", + "api_request_id": getattr( + agent, "_current_api_request_id", None + ), + "call_role": ( + "delegated" + if getattr(agent, "is_subagent", False) + else "fallback" + if int(getattr(agent, "_fallback_index", 0) or 0) > 0 + else "primary" + ), + }, + defer_logical_completion=True, + ) + streamed_response = stream_converse_with_callbacks( + {"stream": stream}, on_text_delta=_on_text if agent._has_stream_consumers() else None, on_tool_start=_on_tool, on_reasoning_delta=_on_reasoning if agent.reasoning_callback or agent.stream_delta_callback else None, on_interrupt_check=lambda: agent._interrupt_requested, on_event=lambda: _bedrock_last_event.__setitem__("t", time.time()), ) + result["response"] = stream.final_response or streamed_response except Exception as e: result["error"] = e + finally: + if stream is not None: + stream.close() - t = threading.Thread(target=_bedrock_call, daemon=True) + t = threading.Thread( + target=_context_thread_target(_bedrock_call), daemon=True + ) t.start() while t.is_alive(): t.join(timeout=0.3) @@ -2698,6 +2813,7 @@ def interruptible_streaming_api_call(agent, api_kwargs: dict, *, on_first_delta= first_delta_fired = {"done": False} deltas_were_sent = {"yes": False} # Track if any deltas were fired (for fallback) + provider_tool_in_flight = {"yes": False} # Wall-clock timestamp of the last real streaming chunk. The outer # poll loop uses this to detect stale connections that keep receiving # SSE keep-alive pings but no actual data. @@ -2717,11 +2833,29 @@ def interruptible_streaming_api_call(agent, api_kwargs: dict, *, on_first_delta= "discarded_chunks": 0, "discarded_bytes": 0, } + managed_stream_holder = {"stream": None} + + def _set_managed_stream(stream: Any) -> Any: + managed_stream_holder["stream"] = stream + return stream + + def _close_managed_stream() -> None: + stream = managed_stream_holder.pop("stream", None) + if stream is None: + return + close = getattr(stream, "close", None) + if callable(close): + try: + close() + except Exception: + logger.debug("Managed provider stream cleanup failed", exc_info=True) def _start_stream_attempt() -> int: with stream_attempt_lock: stream_attempt_state["current"] += 1 - return int(stream_attempt_state["current"]) + attempt_id = int(stream_attempt_state["current"]) + provider_tool_in_flight["yes"] = False + return attempt_id def _cancel_current_stream_attempt(reason: str) -> None: with stream_attempt_lock: @@ -2831,107 +2965,6 @@ def interruptible_streaming_api_call(agent, api_kwargs: dict, *, on_first_delta= # Cap connect/pool at 60s even when provider timeout is higher. # connect/pool cover TCP handshake, not model inference. _conn_cap = min(_base_timeout, 60.0) if _provider_timeout_cfg is not None else 30.0 - stream_kwargs = { - **api_kwargs, - "stream": True, - "timeout": _httpx.Timeout( - connect=_conn_cap, - read=_stream_read_timeout, - write=_base_timeout, - pool=_conn_cap, - ), - } - # OpenAI's `stream_options={"include_usage": True}` drives usage - # accounting on OpenAI-compatible endpoints (incl. the Gemini OpenAI - # compat shim and aggregators like OpenRouter). Google's *native* - # Gemini REST endpoint rejects the keyword outright - # (`Completions.create() got an unexpected keyword argument - # 'stream_options'`), so omit it only for that endpoint. - if not is_native_gemini_base_url(agent.base_url): - stream_kwargs["stream_options"] = {"include_usage": True} - request_client = _set_request_client( - agent._create_request_openai_client( - reason="chat_completion_stream_request", - api_kwargs=stream_kwargs, - ) - ) - # Reset stale-stream timer so the detector measures from this - # attempt's start, not a previous attempt's last chunk. - last_chunk_time["t"] = time.time() - agent._touch_activity("waiting for provider response (streaming)") - # Initialize per-attempt stream diagnostics so the retry block can - # reach for them after the stream dies. Lives on - # ``request_client_holder["diag"]`` for closure access. - _diag = agent._stream_diag_init() - request_client_holder["diag"] = _diag - stream = request_client.chat.completions.create(**stream_kwargs) - if agent.provider == "moa": - # The MoA facade is a shared singleton — abort/close of the - # registered client is a no-op, so register the stream handle - # itself for interrupt teardown (#57354). - stream = _set_request_stream_handle(stream) - # Claim the delta sink for THIS attempt (#65991). If a prior attempt's - # stream is somehow still alive (a stale-stream reconnect whose socket - # abort raced), this claim supersedes it so its late chunks are fenced - # out of the turn instead of interleaving with ours. - _writer_token = claim_stream_writer(agent) - - # Some OpenAI-compatible adapters (for example copilot-acp, and the MoA - # openai-codex aggregator) accept stream=True but still return a - # completed response object rather than an iterator of chunks. Treat - # that as "streaming unsupported" for the rest of this session instead - # of crashing on ``for chunk in stream`` with ``'types.SimpleNamespace' - # object is not iterable`` (#11732, #55933). - # - # Discriminate on the mere PRESENCE of a ``choices`` attribute, not on - # it being a non-empty list: an adapter may hand back a completed - # response whose ``choices`` is ``None`` or empty (an error / - # content-filter / terminal frame), and every such shape is still a - # whole response — not a token stream — that would crash iteration just - # the same. A genuine provider stream (SDK ``Stream`` object, - # generator) exposes no ``choices`` attribute, so it is left untouched. - if hasattr(stream, "choices"): - logger.info( - "Streaming request returned a final response object instead of " - "an iterator; switching %s/%s to non-streaming for this session.", - agent.provider or "unknown", - agent.model or "unknown", - ) - agent._disable_streaming = True - # An empty/None ``choices`` carries no message to surface; return the - # completed object as-is so the outer loop's normal invalid-response - # validation (conversation_loop.py) handles it via the retry path, - # never ``for chunk in stream``. - choices = stream.choices - first_choice = choices[0] if isinstance(choices, (list, tuple)) and choices else None - message = getattr(first_choice, "message", None) - if message is not None: - reasoning_text = ( - getattr(message, "reasoning_content", None) - or getattr(message, "reasoning", None) - ) - if isinstance(reasoning_text, str) and reasoning_text: - _fire_first_delta() - agent._fire_reasoning_delta(reasoning_text) - content = getattr(message, "content", None) - if isinstance(content, str) and content: - _fire_first_delta() - agent._fire_stream_delta(content) - return stream - - # Capture rate limit headers from the initial HTTP response. - # The OpenAI SDK Stream object exposes the underlying httpx - # response via .response before any chunks are consumed. - agent._capture_rate_limits(getattr(stream, "response", None)) - agent._capture_credits(getattr(stream, "response", None)) - # Snapshot diagnostic headers (cf-ray, x-openrouter-provider, etc.) - # so they survive even when the stream dies before any chunk - # arrives. Best-effort; never raises. - agent._stream_diag_capture_response(_diag, getattr(stream, "response", None)) - - # Log OpenRouter response cache status when present. - agent._check_openrouter_cache_status(getattr(stream, "response", None)) - content_parts: list = [] tool_calls_acc: dict = {} tool_gen_notified: set = set() @@ -2946,19 +2979,125 @@ def interruptible_streaming_api_call(agent, api_kwargs: dict, *, on_first_delta= role = "assistant" reasoning_parts: list = [] usage_obj = None - for chunk in stream: - # Stop the moment a newer attempt has claimed the delta sink - # (#65991): this attempt has been superseded, so it must neither - # fire deltas (incl. the tool-suppressed raw-callback path below) - # nor keep consuming a stream that would interleave into the turn. - if not stream_writer_is_current(agent, _writer_token): + _diag = agent._stream_diag_init() + request_client_holder["diag"] = _diag + _writer_token = {"value": None} + attempt_request_client = {"value": None} + + def _open_stream(next_api_kwargs: dict[str, Any]): + stream_kwargs = { + **next_api_kwargs, + "stream": True, + "timeout": _httpx.Timeout( + connect=_conn_cap, + read=_stream_read_timeout, + write=_base_timeout, + pool=_conn_cap, + ), + } + # Native Gemini rejects OpenAI's usage-streaming extension. + if not is_native_gemini_base_url(agent.base_url): + stream_kwargs["stream_options"] = {"include_usage": True} + request_client = _set_request_client( + agent._create_request_openai_client( + reason="chat_completion_stream_request", + api_kwargs=stream_kwargs, + ) + ) + attempt_request_client["value"] = request_client + last_chunk_time["t"] = time.time() + agent._touch_activity("waiting for provider response (streaming)") + return request_client.chat.completions.create(**stream_kwargs) + + def _stream_created(raw_stream: Any) -> None: + response = getattr(raw_stream, "response", None) + agent._capture_rate_limits(response) + agent._capture_credits(response) + agent._stream_diag_capture_response(_diag, response) + agent._check_openrouter_cache_status(response) + _writer_token["value"] = claim_stream_writer(agent) + + def _accept_stream_chunk(_chunk: Any) -> bool: + # A stale-attempt fence can win while Relay is handing an + # already-received tool-call chunk back to Hermes. Preserve only + # the fact that a tool call was in flight so retry policy does not + # misclassify the attempt as a partial text response. The chunk + # itself is still rejected below and never reaches callbacks. + try: + choices = getattr(_chunk, "choices", None) + delta = getattr(choices[0], "delta", None) if choices else None + if getattr(delta, "tool_calls", None): + provider_tool_in_flight["yes"] = True + except Exception: + pass + if not _stream_attempt_is_active(stream_attempt_id): + return False + token = _writer_token["value"] + if token is not None and not stream_writer_is_current(agent, token): logger.warning( "Streaming attempt superseded by a newer stream; stopping " "consumption to preserve the single-writer invariant " "(model=%s).", api_kwargs.get("model", "unknown"), ) - break + return False + # Record provider activity before Relay processes the chunk. This + # prevents the stale watchdog from cancelling a live stream while + # an interceptor or codec is still handling an already-received + # event. + last_chunk_time["t"] = time.time() + return True + + def _relay_final_response() -> dict[str, Any]: + tool_calls = [tool_calls_acc[index] for index in sorted(tool_calls_acc)] + return { + "model": model_name, + "choices": [ + { + "message": { + "role": role, + "content": "".join(content_parts) or None, + "reasoning_content": "".join(reasoning_parts) or None, + "tool_calls": tool_calls or None, + }, + "finish_reason": finish_reason or "stop", + } + ], + "usage": usage_obj, + } + + from agent import relay_llm + + stream = _set_managed_stream( + relay_llm.stream( + api_kwargs, + _open_stream, + session_id=str(getattr(agent, "session_id", "") or ""), + name=str(getattr(agent, "provider", "") or "provider"), + model_name=str(getattr(agent, "model", "") or ""), + finalizer=_relay_final_response, + on_stream_created=_stream_created, + accept_chunk=_accept_stream_chunk, + completed_response_predicate=lambda value: hasattr(value, "choices"), + metadata={ + "api_mode": "chat_completions", + "api_request_id": getattr(agent, "_current_api_request_id", None), + "call_role": ( + "delegated" + if getattr(agent, "is_subagent", False) + else "fallback" + if int(getattr(agent, "_fallback_index", 0) or 0) > 0 + else "primary" + ), + }, + defer_logical_completion=True, + ) + ) + if agent.provider == "moa": + # Hermes interrupts the managed stream; Relay retains sole + # ownership of closing the underlying provider stream. + _set_request_stream_handle(stream) + for chunk in stream: last_chunk_time["t"] = time.time() agent._touch_activity("receiving stream response") @@ -2995,9 +3134,12 @@ def interruptible_streaming_api_call(agent, api_kwargs: dict, *, on_first_delta= # the finally's close really closes the pool instead of # caching it (owner-thread abort: shutdown is safe, and the # FD release still happens in the finally below). - agent._abort_request_openai_client( - request_client, reason="interrupt_stream_close_failed" - ) + request_client = attempt_request_client["value"] + if request_client is not None: + agent._abort_request_openai_client( + request_client, + reason="interrupt_stream_close_failed", + ) break if not _stream_attempt_is_active(stream_attempt_id): @@ -3129,11 +3271,46 @@ def interruptible_streaming_api_call(agent, api_kwargs: dict, *, on_first_delta= if hasattr(chunk, "usage") and chunk.usage: usage_obj = chunk.usage + _close_managed_stream() + if _stream_attempt_was_cancelled(stream_attempt_id): raise _httpx.RemoteProtocolError( f"stream attempt {stream_attempt_id} was superseded" ) + # Some OpenAI-compatible adapters accept ``stream=True`` but return a + # completed response. Relay records that attempt while Hermes preserves + # its existing switch-to-non-streaming behavior for later calls. + if stream.final_response is not None: + final_response = stream.final_response + logger.info( + "Streaming request returned a final response object instead of " + "an iterator; switching %s/%s to non-streaming for this session.", + agent.provider or "unknown", + agent.model or "unknown", + ) + agent._disable_streaming = True + choices = final_response.choices + first_choice = ( + choices[0] + if isinstance(choices, (list, tuple)) and choices + else None + ) + message = getattr(first_choice, "message", None) + if message is not None: + reasoning_text = ( + getattr(message, "reasoning_content", None) + or getattr(message, "reasoning", None) + ) + if isinstance(reasoning_text, str) and reasoning_text: + _fire_first_delta() + agent._fire_reasoning_delta(reasoning_text) + content = getattr(message, "content", None) + if isinstance(content, str) and content: + _fire_first_delta() + agent._fire_stream_delta(content) + return final_response + # Build mock response matching non-streaming shape full_content = "".join(content_parts) or None mock_tool_calls = None @@ -3295,72 +3472,95 @@ def interruptible_streaming_api_call(agent, api_kwargs: dict, *, on_first_delta= # fabricated "successful" empty turn. saw_stream_event = False - # Reset stale-stream timer for this attempt last_chunk_time["t"] = time.time() - # Per-attempt diagnostic dict for the retry block to consume. _diag = agent._stream_diag_init() request_client_holder["diag"] = _diag - # Defensive: strip Responses-only kwargs (instructions, input, ...) - # that can leak in under an api_mode-flip race. The Anthropic SDK - # raises a non-retryable TypeError on them, killing the turn. See - # #31673 / sanitize_anthropic_kwargs(). + _writer_token = {"value": None} + _stream_context = {"manager": None, "stream": None} + base_final_message = None + + from agent import relay_llm from agent.anthropic_adapter import sanitize_anthropic_kwargs - sanitize_anthropic_kwargs( - api_kwargs, log_prefix=getattr(agent, "log_prefix", "") - ) - # Use the Anthropic SDK's streaming context manager - with request_client.messages.stream(**api_kwargs) as stream: + + accumulator = relay_llm.AnthropicStreamAccumulator() + + def _open_anthropic_stream(next_api_kwargs: dict[str, Any]): + final_kwargs = dict(next_api_kwargs) + sanitize_anthropic_kwargs( + final_kwargs, + log_prefix=getattr(agent, "log_prefix", ""), + ) + manager = request_client.messages.stream(**final_kwargs) + _stream_context["manager"] = manager + return manager.__enter__() + + def _anthropic_stream_created(raw_stream: Any) -> None: + _stream_context["stream"] = raw_stream # The Anthropic SDK exposes the raw httpx response on - # ``stream.response``. Snapshot diagnostic headers - # immediately so they survive a stream that dies before the - # first event. + # ``stream.response``. Snapshot diagnostics immediately so they + # survive a stream that dies before the first event. try: agent._stream_diag_capture_response( - _diag, getattr(stream, "response", None) + _diag, + getattr(raw_stream, "response", None), ) except Exception: pass - # Claim the delta sink for THIS attempt (#65991) — parity with the - # chat_completions path so a superseded anthropic stream is fenced. - _writer_token = claim_stream_writer(agent) + _writer_token["value"] = claim_stream_writer(agent) + + def _accept_anthropic_event(_event: Any) -> bool: + token = _writer_token["value"] + if token is None or stream_writer_is_current(agent, token): + return True + logger.warning( + "Anthropic streaming attempt superseded by a newer stream; " + "stopping consumption to preserve the single-writer " + "invariant (model=%s).", + api_kwargs.get("model", "unknown"), + ) + return False + + stream = _set_managed_stream( + relay_llm.stream( + api_kwargs, + _open_anthropic_stream, + session_id=str(getattr(agent, "session_id", "") or ""), + name=str(getattr(agent, "provider", "") or "anthropic"), + model_name=str(getattr(agent, "model", "") or ""), + finalizer=accumulator.finalize, + on_stream_created=_anthropic_stream_created, + on_chunk=accumulator.observe, + accept_chunk=_accept_anthropic_event, + metadata={ + "api_mode": "anthropic_messages", + "api_request_id": getattr(agent, "_current_api_request_id", None), + "call_role": ( + "delegated" + if getattr(agent, "is_subagent", False) + else "fallback" + if int(getattr(agent, "_fallback_index", 0) or 0) > 0 + else "primary" + ), + }, + defer_logical_completion=True, + ) + ) + try: for event in stream: - # Bail the instant a newer attempt supersedes this one so a - # stale stream can't interleave tokens into the turn. - if not stream_writer_is_current(agent, _writer_token): - logger.warning( - "Anthropic streaming attempt superseded by a newer " - "stream; stopping consumption to preserve the " - "single-writer invariant (model=%s).", - api_kwargs.get("model", "unknown"), - ) - break saw_stream_event = True - # Update stale-stream timer on every event so the - # outer poll loop knows data is flowing. Without - # this, the detector kills healthy long-running - # Opus streams after 180 s even when events are - # actively arriving (the chat_completions path - # already does this at the top of its chunk loop). last_chunk_time["t"] = time.time() agent._touch_activity("receiving stream response") - - # Update per-attempt diagnostic counters (best-effort). try: _diag["chunks"] = int(_diag.get("chunks", 0)) + 1 if _diag.get("first_chunk_at") is None: _diag["first_chunk_at"] = last_chunk_time["t"] - try: - _diag["bytes"] = int(_diag.get("bytes", 0)) + len(repr(event)) - except Exception: - pass + _diag["bytes"] = int(_diag.get("bytes", 0)) + len(repr(event)) except Exception: pass - if agent._interrupt_requested: break event_type = getattr(event, "type", None) - if event_type == "content_block_start": block = getattr(event, "content_block", None) if block and getattr(block, "type", None) == "tool_use": @@ -3369,7 +3569,6 @@ def interruptible_streaming_api_call(agent, api_kwargs: dict, *, on_first_delta= if tool_name: _fire_first_delta() agent._fire_tool_gen_started(tool_name) - elif event_type == "content_block_delta": delta = getattr(event, "delta", None) if delta: @@ -3385,48 +3584,49 @@ def interruptible_streaming_api_call(agent, api_kwargs: dict, *, on_first_delta= if thinking_text: _fire_first_delta() agent._fire_reasoning_delta(thinking_text) - - # Return the native Anthropic Message for downstream processing. - # If the stream was interrupted (the event loop broke out above on - # agent._interrupt_requested), do NOT call get_final_message() — on - # a partially-consumed stream the SDK may hang draining remaining - # events or return a Message with incomplete tool_use blocks (partial - # JSON in `input`). The outer poll loop raises InterruptedError, so - # this return value is discarded anyway. - if agent._interrupt_requested: - return None - # Zero-event guard (parity with the chat_completions zero-chunk - # guard above). Real SDK: an eventless stream has no - # message_start, so get_final_message() raises AssertionError - # (final-message snapshot is None) — normalize that to - # EmptyStreamError so it gets the transient retry budget - # instead of surfacing raw. + if not agent._interrupt_requested: + raw_stream = _stream_context["stream"] + if raw_stream is not None: + try: + base_final_message = raw_stream.get_final_message() + except AssertionError: + if not saw_stream_event: + raise EmptyStreamError( + "Provider returned an empty stream with no events " + "(possible upstream error or malformed event stream)." + ) from None + raise + finally: try: - _final_message = stream.get_final_message() - except AssertionError: - if not saw_stream_event: - raise EmptyStreamError( - "Provider returned an empty stream with no events " - "(possible upstream error or malformed event stream)." - ) from None - raise - # Shim variants of the same failure: an OpenAI-compat adapter - # may fabricate a contentless Message with no stop_reason, or - # return None where the SDK assert would have fired (e.g. - # ``python -O``). A real completed response always carries a - # stop_reason, so this cannot fire on legitimate turns. - if not saw_stream_event and ( - _final_message is None - or ( - not getattr(_final_message, "content", None) - and getattr(_final_message, "stop_reason", None) is None - ) - ): - raise EmptyStreamError( - "Provider returned an empty stream with no stop_reason " - "(possible upstream error or malformed event stream)." - ) - return _final_message + _close_managed_stream() + finally: + manager = _stream_context["manager"] + if manager is not None: + manager.__exit__(None, None, None) + + if agent._interrupt_requested: + return None + if ( + base_final_message is not None + and not getattr(base_final_message, "content", None) + and getattr(base_final_message, "stop_reason", None) is None + ): + raise EmptyStreamError( + "Provider returned an empty stream with no stop_reason " + "(possible upstream error or malformed event stream)." + ) + if base_final_message is not None and not stream.output_modified: + return base_final_message + final_message = accumulator.response(base_final_message) + if ( + not getattr(final_message, "content", None) + and getattr(final_message, "stop_reason", None) is None + ): + raise EmptyStreamError( + "Provider returned an empty stream with no stop_reason " + "(possible upstream error or malformed event stream)." + ) + return final_message def _call(): import httpx as _httpx @@ -3461,6 +3661,7 @@ def interruptible_streaming_api_call(agent, api_kwargs: dict, *, on_first_delta= result["response"] = _call_chat_completions(stream_attempt_id) return # success except Exception as e: + _close_managed_stream() # If the main poll loop force-closed this request because # of an interrupt, the resulting transport error is the # expected consequence of our own close — NOT a transient @@ -3500,7 +3701,7 @@ def interruptible_streaming_api_call(agent, api_kwargs: dict, *, on_first_delta= if deltas_were_sent["yes"]: _partial_tool_in_flight = bool( result.get("partial_tool_names") - ) + ) or provider_tool_in_flight["yes"] _is_sse_conn_err_preview = False if not _is_timeout and not _is_conn_err: from openai import APIError as _APIError @@ -3749,6 +3950,7 @@ def interruptible_streaming_api_call(agent, api_kwargs: dict, *, on_first_delta= result["error"] = e return finally: + _close_managed_stream() # Reuse reason only on a clean stream; any other outcome (error, # cancel-swallow) really closes so the next attempt builds a # fresh pool (see _REQUEST_CLIENT_REUSE_REASONS). @@ -3818,7 +4020,7 @@ def interruptible_streaming_api_call(agent, api_kwargs: dict, *, on_first_delta= if _reasoning_floor is not None: _stream_stale_timeout = max(_stream_stale_timeout, _reasoning_floor) - t = threading.Thread(target=_call, daemon=True) + t = threading.Thread(target=_context_thread_target(_call), daemon=True) t.start() _last_heartbeat = time.time() _HEARTBEAT_INTERVAL = 30.0 # seconds between gateway activity touches diff --git a/agent/codex_responses_adapter.py b/agent/codex_responses_adapter.py index ee75f4190e6d..edff776536e7 100644 --- a/agent/codex_responses_adapter.py +++ b/agent/codex_responses_adapter.py @@ -191,6 +191,28 @@ def _deterministic_call_id(fn_name: str, arguments: str, index: int = 0) -> str: return f"call_{digest}" +def _clamp_responses_call_id(call_id: str) -> str: + """Keep a ``call_id`` within the Responses API's 64-char limit (#73492). + + The codex app-server namespaces MCP tool call ids as + ``codex_mcp_____``; with an ``exec-`` + component the built-in ``hermes-tools`` server already overflows 64 chars, + and the Responses API rejects the whole payload with a non-retryable HTTP + 400 that then replays every turn — permanently bricking the session. + + Sibling defect to #10788 (which clamped ``input[*].id``), applied here to + ``call_id``. The surrogate is a pure, deterministic function of the + original, so the ``function_call`` and its matching ``function_call_output`` + — which carry the same original id — map to the same surrogate and stay + paired without correlating the two items. Short ids pass through unchanged, + preserving prompt-cache prefixes. + """ + if len(call_id) <= _MAX_RESPONSES_ITEM_ID_LENGTH: + return call_id + digest = hashlib.sha256(call_id.encode("utf-8", errors="replace")).hexdigest()[:32] + return f"call_{digest}" + + def _split_responses_tool_id(raw_id: Any) -> tuple[Optional[str], Optional[str]]: """Split a stored tool id into (call_id, response_item_id).""" if not isinstance(raw_id, str): @@ -546,7 +568,7 @@ def _chat_messages_to_responses_input( items.append({ "type": "function_call", - "call_id": call_id, + "call_id": _clamp_responses_call_id(call_id), "name": fn_name, "arguments": arguments, }) @@ -589,7 +611,7 @@ def _chat_messages_to_responses_input( items.append({ "type": "function_call_output", - "call_id": call_id, + "call_id": _clamp_responses_call_id(call_id), "output": output_value, }) diff --git a/agent/codex_runtime.py b/agent/codex_runtime.py index f7d3d3c9f181..0e0b87b21960 100644 --- a/agent/codex_runtime.py +++ b/agent/codex_runtime.py @@ -1240,6 +1240,8 @@ def run_codex_stream(agent, api_kwargs: dict, client: Any = None, on_first_delta """ import httpx as _httpx + from agent import relay_llm + active_client = client or agent._ensure_primary_openai_client(reason="codex_stream_direct") max_stream_retries = 1 # Accumulate streamed text so callers / compat shims can read it. @@ -1264,48 +1266,88 @@ def run_codex_stream(agent, api_kwargs: dict, client: Any = None, on_first_delta if agent._interrupt_requested: raise InterruptedError("Agent interrupted before Codex stream retry") - stream_kwargs = dict(api_kwargs) - stream_kwargs["stream"] = True + intercepted_events = [] + writer_token = {"value": None} + + def _open_codex_stream(next_api_kwargs: dict[str, Any]): + stream_kwargs = dict(next_api_kwargs) + stream_kwargs["stream"] = True + return active_client.responses.create(**stream_kwargs) + + def _codex_stream_created(_raw_stream: Any) -> None: + # Claim the delta sink for THIS physical attempt. A newer attempt + # supersedes this token and fences late deltas out of the turn. + writer_token["value"] = claim_stream_writer(agent) + + def _accept_codex_chunk(_chunk: Any) -> bool: + token = writer_token["value"] + if token is None or stream_writer_is_current(agent, token): + return True + logger.warning( + "Codex streaming attempt superseded by a newer stream; " + "stopping consumption to preserve the single-writer " + "invariant (model=%s).", + api_kwargs.get("model", "unknown"), + ) + return False + + def _finalize_codex_stream() -> Any: + return _consume_codex_event_stream( + list(intercepted_events), + model=api_kwargs.get("model"), + ) try: - event_stream = active_client.responses.create(**stream_kwargs) - except (_httpx.RemoteProtocolError, _httpx.ReadTimeout, _httpx.ConnectError, ConnectionError) as exc: + event_stream = relay_llm.stream( + dict(api_kwargs), + _open_codex_stream, + session_id=str(getattr(agent, "session_id", "") or ""), + name=str(getattr(agent, "provider", "") or "codex"), + model_name=str(api_kwargs.get("model") or ""), + finalizer=_finalize_codex_stream, + on_stream_created=_codex_stream_created, + on_chunk=intercepted_events.append, + chunk_adapter=lambda chunk: chunk, + accept_chunk=_accept_codex_chunk, + completed_response_predicate=lambda response: bool( + hasattr(response, "output") and not hasattr(response, "__iter__") + ), + metadata={ + "api_mode": "codex_responses", + "api_request_id": getattr(agent, "_current_api_request_id", None), + "call_role": ( + "delegated" + if getattr(agent, "is_subagent", False) + else "fallback" + if int(getattr(agent, "_fallback_index", 0) or 0) > 0 + else "primary" + ), + "retry_count": attempt, + }, + defer_logical_completion=True, + ) + except ( + _httpx.RemoteProtocolError, + _httpx.ReadTimeout, + _httpx.ConnectError, + ConnectionError, + ) as exc: if attempt < max_stream_retries: logger.debug( - "Codex Responses stream connect failed (attempt %s/%s); retrying. %s error=%s", - attempt + 1, max_stream_retries + 1, - agent._client_log_context(), exc, + "Codex Responses stream connect failed (attempt %s/%s); " + "retrying. %s error=%s", + attempt + 1, + max_stream_retries + 1, + agent._client_log_context(), + exc, ) continue raise - # Claim the delta sink for THIS attempt (#65991) — parity with the - # chat_completions/anthropic/bedrock paths. If a prior attempt's - # stream is somehow still alive, this claim supersedes it so its - # late deltas are fenced out of the turn; conversely, a newer - # attempt supersedes us and the interrupt_check below stops our - # consumption immediately. - _writer_token = claim_stream_writer(agent) - - def _interrupt_or_superseded(_tok=_writer_token) -> bool: - if agent._interrupt_requested: - return True - if not stream_writer_is_current(agent, _tok): - logger.warning( - "Codex streaming attempt superseded by a newer stream; " - "stopping consumption to preserve the single-writer " - "invariant (model=%s).", - api_kwargs.get("model", "unknown"), - ) - return True - return False + def _interrupt_or_superseded() -> bool: + return bool(agent._interrupt_requested) try: - # Compatibility: some mocks/providers return a concrete response - # instead of an iterable. Pass it straight through. - if hasattr(event_stream, "output") and not hasattr(event_stream, "__iter__"): - return event_stream - try: final = _consume_codex_event_stream( event_stream, @@ -1324,6 +1366,12 @@ def run_codex_stream(agent, api_kwargs: dict, client: Any = None, on_first_delta on_event=_on_event, interrupt_check=_interrupt_or_superseded, ) + # The terminal SSE frame is contractually last. Request the + # end-of-stream marker so Relay can run its response finalizer + # and close the physical attempt scope before Hermes returns. + if not agent._interrupt_requested: + for _ignored in event_stream: + pass except (_httpx.RemoteProtocolError, _httpx.ReadTimeout, _httpx.ConnectError, ConnectionError) as exc: if attempt < max_stream_retries: logger.debug( @@ -1334,6 +1382,10 @@ def run_codex_stream(agent, api_kwargs: dict, client: Any = None, on_first_delta ) continue raise + except RuntimeError: + if event_stream.final_response is not None: + return event_stream.final_response + raise if final.status in {"incomplete", "failed"}: logger.warning( diff --git a/agent/conversation_loop.py b/agent/conversation_loop.py index 38b6b1b31a67..06504616bac0 100644 --- a/agent/conversation_loop.py +++ b/agent/conversation_loop.py @@ -520,7 +520,7 @@ def _restore_or_build_system_prompt(agent, system_message, conversation_history) # session is created (not on continuation). Plugins can use this # to initialise session-scoped state (e.g. warm a memory cache). try: - from hermes_cli.plugins import invoke_hook as _invoke_hook + from hermes_cli.lifecycle import invoke_hook as _invoke_hook _invoke_hook( "on_session_start", session_id=agent.session_id, @@ -2100,7 +2100,7 @@ def run_conversation( _llm_middleware_trace = [] try: - from hermes_cli.plugins import ( + from hermes_cli.lifecycle import ( has_hook, invoke_hook as _invoke_hook, ) @@ -2141,6 +2141,7 @@ def run_conversation( base_url=agent.base_url, api_mode=agent.api_mode, api_call_count=api_call_count, + retry_count=retry_count, request_messages=list(request_messages) if isinstance(request_messages, list) else [], @@ -2230,7 +2231,28 @@ def run_conversation( return agent._interruptible_streaming_api_call( next_api_kwargs, on_first_delta=_stop_spinner ) - return agent._interruptible_api_call(next_api_kwargs) + from agent import relay_llm + + return relay_llm.execute( + next_api_kwargs, + agent._interruptible_api_call, + session_id=str(agent.session_id or ""), + name=str(agent.provider or "provider"), + model_name=str(agent.model or ""), + metadata={ + "api_mode": agent.api_mode, + "api_request_id": api_request_id, + "call_role": ( + "delegated" + if getattr(agent, "is_subagent", False) + else "fallback" + if int(getattr(agent, "_fallback_index", 0) or 0) > 0 + else "primary" + ), + "retry_count": retry_count, + }, + defer_logical_completion=True, + ) from hermes_cli.middleware import run_llm_execution_middleware @@ -3288,6 +3310,12 @@ def run_conversation( clear_nous_rate_limit() except Exception: pass + from agent import relay_llm + + relay_llm.complete_logical_call( + api_request_id, + outcome="success", + ) agent._touch_activity(f"API call #{api_call_count} completed") break # Success, exit retry loop @@ -5432,7 +5460,7 @@ def run_conversation( assistant_message.content = str(raw) try: - from hermes_cli.plugins import ( + from hermes_cli.lifecycle import ( has_hook, invoke_hook as _invoke_hook, ) @@ -6765,7 +6793,8 @@ def run_conversation( _attempt = getattr(agent, "_pre_verify_nudges", 0) try: from agent.verify_hooks import max_verify_nudges - from hermes_cli.plugins import get_pre_verify_continue_message, has_hook + from hermes_cli.lifecycle import has_hook + from hermes_cli.plugins import get_pre_verify_continue_message if _edited and has_hook("pre_verify") and _attempt < max_verify_nudges(): # Posture is fixed for the session — resolve once + cache. diff --git a/agent/credits_tracker.py b/agent/credits_tracker.py index 929bc34d3262..dfd39e4290fb 100644 --- a/agent/credits_tracker.py +++ b/agent/credits_tracker.py @@ -288,20 +288,12 @@ def evaluate_credits_notices( current_band = band # Top-up suppression: when the account holds purchased (top-up) credits, # the subscription-cap gauge is the wrong denominator — warning "90% used" - # at a user sitting on $50 of top-up is noise (and it previously stuck - # PERMANENTLY alongside grant_spent at >=100%). Suppress the usage band - # entirely; the cap-reached case is covered by the grant_spent info notice - # below, which already names the remaining top-up balance. A top-up landing - # mid-session flips current_band → None and the clear path below removes - # any showing band line. + # at a user sitting on $50 of top-up is noise. Suppress the usage band + # entirely; /usage still reports the full balance breakdown. A top-up + # landing mid-session flips current_band → None and the clear path below + # removes any showing band line. if state.purchased_micros > 0: current_band = None - grant_cond = ( - state.denominator_kind == "subscription_cap" - and uf is not None - and uf >= 1.0 - and state.purchased_micros > 0 - ) depleted_cond = not state.paid_access # ── usage gauge (escalating single notice: 50 → 75 → 90) ────────────────── @@ -340,22 +332,6 @@ def evaluate_credits_notices( active.add(CREDITS_USAGE_KEY) latch["usage_band"] = target_band - # ── grant_spent ────────────────────────────────────────────────────────── - if grant_cond and "credits.grant_spent" not in active: - to_show.append( - AgentNotice( - text=f"• Grant spent · ${state.purchased_usd} top-up left", - level="info", - kind=CREDITS_NOTICE_KIND, - key="credits.grant_spent", - id="credits.grant_spent", - ) - ) - active.add("credits.grant_spent") - elif "credits.grant_spent" in active and not grant_cond: - to_clear.append("credits.grant_spent") - active.discard("credits.grant_spent") - # ── depleted ───────────────────────────────────────────────────────────── # Suppressed while the active model is free: inference still works there, # so the error banner would just alarm users (free-tier users especially, @@ -627,7 +603,7 @@ _DEV_FIXTURES: dict[str, dict] = { subscription_limit_micros=20_000_000, subscription_limit_usd="20.00", denominator_kind="subscription_cap", paid_access=True, ), - "grant_exhausted": dict( # used_fraction == 1.0 + purchased>0 → credits.grant_spent + "grant_exhausted": dict( # cap reached + purchased>0 → no notice (gauge suppressed by top-up) remaining_micros=12_340_000, remaining_usd="12.34", subscription_micros=0, subscription_usd="0.00", subscription_limit_micros=20_000_000, subscription_limit_usd="20.00", diff --git a/agent/relay_llm.py b/agent/relay_llm.py new file mode 100644 index 000000000000..3481ca072376 --- /dev/null +++ b/agent/relay_llm.py @@ -0,0 +1,1130 @@ +"""Core NeMo Relay adapters for physical Hermes provider attempts.""" + +from __future__ import annotations + +import asyncio +import contextvars +import inspect +import json +import logging +from collections.abc import Callable, Iterator +from types import SimpleNamespace +from typing import Any + +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"} +) + + +def execute( + request: dict[str, Any], + callback: Callable[[dict[str, Any]], Any], + *, + session_id: str, + name: str, + model_name: str, + metadata: dict[str, Any] | None = None, + 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, + ) + raw_response: dict[str, Any] = {} + callback_error: BaseException | None = None + callback_context = contextvars.copy_context() + + 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, + ) + raw = callback_context.copy().run(callback, final_request) + except BaseException as exc: + callback_error = exc + raise + raw_response["value"] = raw + raw_response["json"] = _jsonable(raw) + return raw_response["json"] + + try: + managed = _run_awaitable( + runtime.run_in_session_async( + session, + runtime.relay.llm.execute, + name, + relay_request, + invoke, + handle=parent, + metadata=_jsonable(metadata or {}), + model_name=model_name, + codec=_codec(runtime.relay, metadata), + response_codec=_codec(runtime.relay, metadata), + ) + ) + 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) + + +async def execute_async( + request: dict[str, Any], + callback: Callable[[dict[str, Any]], Any], + *, + session_id: str, + name: str, + model_name: str, + metadata: dict[str, Any] | None = None, + 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, + ) + raw_response: dict[str, Any] = {} + callback_error: BaseException | None = None + callback_context = contextvars.copy_context() + + 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, + ) + async def call_provider() -> Any: + return await callback(final_request) + + task = callback_context.copy().run( + asyncio.create_task, + call_provider(), + ) + raw = await task + except BaseException as exc: + callback_error = exc + raise + raw_response["value"] = raw + raw_response["json"] = _jsonable(raw) + return raw_response["json"] + + try: + managed = await runtime.run_in_session_async( + session, + runtime.relay.llm.execute, + name, + relay_request, + invoke, + handle=parent, + metadata=_jsonable(metadata or {}), + model_name=model_name, + codec=_codec(runtime.relay, metadata), + response_codec=_codec(runtime.relay, metadata), + ) + 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) + + +def execute_current( + request: dict[str, Any], + callback: Callable[[dict[str, Any]], Any], + *, + name: str, + model_name: str, + metadata: dict[str, Any] | None = None, + 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: + return callback(request) + return execute( + request, + callback, + session_id=turn.lease.session_id, + name=name, + model_name=model_name, + metadata=metadata, + defer_logical_completion=defer_logical_completion, + ) + + +async def execute_current_async( + request: dict[str, Any], + callback: Callable[[dict[str, Any]], Any], + *, + name: str, + model_name: str, + metadata: dict[str, Any] | None = None, + 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: + return await callback(request) + return await execute_async( + request, + callback, + session_id=turn.lease.session_id, + name=name, + model_name=model_name, + metadata=metadata, + defer_logical_completion=defer_logical_completion, + ) + + +def stream_current( + request: dict[str, Any], + stream_factory: Callable[[dict[str, Any]], Any], + *, + name: str, + model_name: str, + finalizer: Callable[[], Any], + metadata: dict[str, Any] | None = None, + defer_logical_completion: bool = False, +) -> Any: + """Run a provider stream under the inherited Hermes turn when present.""" + turn = relay_runtime.active_turn() + if turn is None: + return stream_factory(request) + return stream( + request, + stream_factory, + session_id=turn.lease.session_id, + name=name, + model_name=model_name, + finalizer=finalizer, + metadata=metadata, + defer_logical_completion=defer_logical_completion, + ) + + +def stream( + request: dict[str, Any], + stream_factory: Callable[[dict[str, Any]], Any], + *, + session_id: str, + name: str, + model_name: str, + finalizer: Callable[[], Any], + on_stream_created: Callable[[Any], None] | None = None, + on_chunk: Callable[[Any], None] | None = None, + chunk_adapter: Callable[[Any], Any] | None = None, + accept_chunk: Callable[[Any], bool] | None = None, + completed_response_predicate: Callable[[Any], bool] | None = None, + metadata: dict[str, Any] | None = None, + defer_logical_completion: bool = False, +) -> "ManagedLlmStream": + """Return a synchronous view of one Relay-managed provider stream.""" + return ManagedLlmStream( + request, + stream_factory, + session_id=session_id, + name=name, + model_name=model_name, + finalizer=finalizer, + on_stream_created=on_stream_created, + on_chunk=on_chunk, + chunk_adapter=chunk_adapter, + accept_chunk=accept_chunk, + completed_response_predicate=completed_response_predicate, + metadata=metadata, + defer_logical_completion=defer_logical_completion, + ) + + +class ManagedLlmStream(Iterator[Any]): + """Drive Relay's async stream from Hermes's provider worker thread.""" + + def __init__( + self, + request: dict[str, Any], + stream_factory: Callable[[dict[str, Any]], Any], + *, + session_id: str, + name: str, + model_name: str, + finalizer: Callable[[], Any], + on_stream_created: Callable[[Any], None] | None, + on_chunk: Callable[[Any], None] | None, + chunk_adapter: Callable[[Any], Any] | None, + accept_chunk: Callable[[Any], bool] | None, + completed_response_predicate: Callable[[Any], bool] | None, + metadata: dict[str, Any] | None, + defer_logical_completion: bool, + ) -> None: + self.final_response: Any = None + self._loop: asyncio.AbstractEventLoop | None = None + self._stream: Any = None + self._raw_stream_resource: Any = None + self._closed = False + self._close_error: BaseException | None = None + self._callback_error: BaseException | None = None + self._logical: tuple[relay_runtime.RelayTurnContext, Any, str] | None = None + self._defer_logical_completion = defer_logical_completion + self._on_chunk = on_chunk + self._chunk_adapter = chunk_adapter or _namespace + self._accept_chunk = accept_chunk + self._relay_observes_chunks = False + self._provider_completed = False + self._raw_chunks: list[tuple[Any, 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. + return callback_context.copy().run(callback, *args) + + 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, + ) + + 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, + ) + ): + self.final_response = raw_stream + self._provider_completed = True + return + if on_stream_created is not None: + run_callback(on_stream_created, raw_stream) + raw_iterator = run_callback(iter, raw_stream) + while True: + try: + chunk = run_callback(next, raw_iterator) + except StopIteration: + break + if self._accept_chunk is not None and not run_callback( + self._accept_chunk, + chunk, + ): + break + encoded_chunk = _jsonable(chunk) + self._raw_chunks.append((encoded_chunk, chunk)) + yield encoded_chunk + self._provider_completed = True + except BaseException as exc: + self._callback_error = exc + raise + finally: + close = getattr(raw_stream, "close", None) + if callable(close): + try: + run_callback(close) + except BaseException as exc: + self._close_error = exc + raise + + def observe_chunk(chunk: Any) -> None: + if self._on_chunk is not None: + run_callback(self._on_chunk, _jsonable(chunk)) + + 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. + if self._callback_error is not None: + return None + try: + if self.final_response is not None: + return _jsonable(self.final_response) + return _jsonable(run_callback(finalizer)) + except BaseException as exc: + self._callback_error = exc + raise + + loop = asyncio.new_event_loop() + self._loop = loop + self._relay_observes_chunks = True + try: + self._stream = loop.run_until_complete( + runtime.run_in_session_async( + session, + runtime.relay.llm.stream_execute, + name, + relay_request, + provider_stream, + observe_chunk, + relay_finalizer, + handle=parent, + metadata=_jsonable(metadata or {}), + 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, + ) + self._preserve_pending_provider_chunks() + return + if not self._defer_logical_completion: + _complete_logical( + self._logical, + outcome="cancelled" if _is_cancellation(exc) else "failed", + ) + self._logical = None + loop.close() + self._loop = None + raise + + def __iter__(self) -> "ManagedLlmStream": + return self + + def __next__(self) -> Any: + if self._closed: + raise StopIteration + 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): + self._close(logical_outcome="cancelled") + raise StopIteration + return chunk + + async def next_chunk() -> Any: + return await anext(self._stream) + + try: + chunk = self._loop.run_until_complete(next_chunk()) + except StopAsyncIteration: + if self._raw_chunks: + self.output_modified = True + if not self._defer_logical_completion: + _complete_logical(self._logical, outcome="success") + self._logical = None + self._close(logical_outcome="cancelled") + raise StopIteration from None + except BaseException as exc: + callback_error = self._callback_error + if ( + callback_error is not None + and relay_runtime._is_relay_wrapped_callback_error(exc, callback_error) + ): + 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, + ) + self._preserve_pending_provider_chunks() + return next(self) + 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) + for index, (encoded, raw) in enumerate(self._raw_chunks): + if _json_equal(chunk, encoded): + if index > 0: + self.output_modified = True + del self._raw_chunks[: index + 1] + return raw + self.output_modified = True + return self._chunk_adapter(chunk) + + def close(self) -> None: + """Close an explicitly abandoned stream and cancel its logical call.""" + self._close(logical_outcome="cancelled") + close_error = self._close_error + self._close_error = None + if close_error is not None: + raise close_error + + def _preserve_pending_provider_chunks(self) -> None: + """Switch a failed Relay stream to its undelivered provider chunks.""" + pending = [raw for _encoded, raw in self._raw_chunks] + self._raw_chunks.clear() + loop = self._loop + relay_stream = self._stream + self._loop = None + self._stream = iter(pending) + self._raw_stream_resource = None + self._accept_chunk = None + 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, + ) + loop.close() + if not self._defer_logical_completion: + _complete_logical(self._logical, outcome="success") + self._logical = None + + def _close(self, *, logical_outcome: str) -> None: + if self._closed: + return + self._closed = True + 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) + self._logical = None + 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) + self._logical = None + loop.close() + + def __del__(self) -> None: + self._close(logical_outcome="cancelled") + + +class AnthropicStreamAccumulator: + """Rebuild an Anthropic Message from post-intercept SSE events.""" + + def __init__(self) -> None: + self._message: dict[str, Any] = {} + self._blocks: dict[int, dict[str, Any]] = {} + + def observe(self, event: Any) -> None: + 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] + 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 + + def finalize(self) -> dict[str, Any]: + blocks = [] + for index in sorted(self._blocks): + block = dict(self._blocks[index]) + partial = block.pop("_partial_json", None) + if partial is not None: + try: + block["input"] = json.loads(partial) + except (TypeError, ValueError): + block["input"] = partial + blocks.append(block) + return {**self._message, "content": blocks} + + 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 = {} + content = assembled.pop("content", []) + merged = {**base_payload, **assembled} + if content or "content" not in merged: + merged["content"] = content + return _namespace(merged) + + +def _logical_parent( + runtime: relay_runtime.RelayRuntime, + session: Any, + parent: Any, + metadata: dict[str, Any] | None, +) -> tuple[relay_runtime.RelayTurnContext, Any, str] | None: + 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: + return None + with turn.finalize_lock: + if turn.closed: + return None + with turn.logical_llm_lock: + handle = turn.logical_llm_calls.get(request_id) + if handle is None: + handle = runtime.run_in_session( + session, + runtime.relay.scope.push, + relay_runtime.LOGICAL_LLM_SCOPE, + 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" + ), + }, + ) + turn.logical_llm_calls[request_id] = handle + return turn, handle, request_id + + +def _complete_logical( + logical: tuple[relay_runtime.RelayTurnContext, Any, str] | None, + *, + outcome: str, +) -> None: + if logical is None: + return + turn, handle, request_id = logical + lease = turn.lease + if not isinstance(lease.host, relay_runtime.RelayRuntime): + return + with turn.finalize_lock: + with turn.logical_llm_lock: + if turn.logical_llm_calls.get(request_id) is not handle: + return + if lease.session is None: + return + try: + lease.host.run_in_session( + lease.session, + lease.host.relay.scope.pop, + handle, + output={"outcome": outcome}, + metadata={ + relay_runtime.RUNTIME_SCHEMA_KEY: relay_runtime.RUNTIME_SCHEMA_VERSION, + relay_runtime.RUNTIME_INSTANCE_KEY: 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, + ) + 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), + ) + + +def complete_logical_call(api_request_id: str, *, outcome: str) -> None: + """Complete the active turn's logical LLM call after caller validation.""" + turn = relay_runtime.active_turn() + if turn is None or not api_request_id: + return + with turn.logical_llm_lock: + handle = turn.logical_llm_calls.get(api_request_id) + if handle is not None: + _complete_logical((turn, handle, api_request_id), outcome=outcome) + + +def _provider_request( + original: dict[str, Any], + request: Any, + *, + relay_request_body: dict[str, Any], + codec_baseline_body: dict[str, Any] | None, + metadata: dict[str, Any] | None, +) -> dict[str, Any]: + 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: + 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], + ): + final[key] = intercepted[key] + _restore_provider_message_extensions( + original, + final, + baseline=baseline, + intercepted=intercepted, + ) + headers = getattr(request, "headers", None) + if isinstance(headers, dict): + headers = { + key: value + for key, value in headers.items() + if str(key).lower() not in _RELAY_INTERNAL_PROVIDER_HEADERS + } + if 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 {} + # 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) + 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" + }, + } + if isinstance(tool, dict) + and tool.get("type") == "function" + and "function" not in tool + else tool + for tool in body["tools"] + ] + elif str((metadata or {}).get("api_mode") or "") == "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 + else tool + for tool in tools + ] + return body + + +def _restore_provider_message_extensions( + original: dict[str, Any], + final: dict[str, Any], + *, + baseline: dict[str, Any], + 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, + ) + ): + return + if not ( + len(original_messages) + == len(final_messages) + == len(baseline_messages) + == len(intercepted_messages) + ): + 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, + ) + ): + continue + for key in _PROVIDER_MESSAGE_EXTENSION_KEYS: + if ( + key in original_message + and key not in baseline_message + and key not in intercepted_message + and key not in final_message + ): + final_message[key] = original_message[key] + + +def _codec_round_trip_request_body( + relay: Any, + relay_request: Any, + *, + relay_request_body: dict[str, Any], + metadata: dict[str, Any] | None, +) -> dict[str, Any] | None: + """Return the codec-only request shape used to identify real rewrites.""" + codec = _codec(relay, metadata) + if codec is None: + return _provider_request_body(relay_request_body, metadata) + try: + annotated = codec.decode(relay_request) + encoded = codec.encode(annotated, 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 + logger.warning( + "NeMo Relay request codec returned an unsupported baseline; " + "ignoring request rewrites" + ) + return None + + +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": + return body + tools = body.get("tools") + if not isinstance(tools, list): + return body + body["tools"] = [ + { + "type": "function", + **dict(tool["function"]), + } + if isinstance(tool, dict) + and tool.get("type") == "function" + and isinstance(tool.get("function"), dict) + else tool + for tool in tools + ] + return body + + +def _codec(relay: Any, metadata: dict[str, Any] | None) -> Any: + api_mode = str((metadata or {}).get("api_mode") or "") + codecs = getattr(relay, "codecs", None) + if codecs is None: + return None + if api_mode == "chat_completions": + codec = getattr(codecs, "OpenAIChatCodec", None) + elif api_mode == "anthropic_messages": + codec = getattr(codecs, "AnthropicMessagesCodec", None) + elif api_mode == "codex_responses": + codec = getattr(codecs, "OpenAIResponsesCodec", None) + else: + codec = None + return codec() if callable(codec) else None + + +def _jsonable(value: Any) -> Any: + if value is None or isinstance(value, (str, int, float, bool)): + return value + if isinstance(value, dict): + return {str(key): _jsonable(item) for key, item in value.items()} + if isinstance(value, (list, tuple, set)): + return [_jsonable(item) for item in value] + model_dump = getattr(type(value), "model_dump", None) + if callable(model_dump): + try: + return _jsonable(value.model_dump(mode="json")) + except Exception: + pass + try: + attributes = { + 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 _namespace(value: Any) -> Any: + if isinstance(value, dict): + 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 + + +def _json_equal(left: Any, right: Any) -> bool: + try: + return json.dumps( + _jsonable(left), sort_keys=True, separators=(",", ":") + ) == json.dumps(_jsonable(right), sort_keys=True, separators=(",", ":")) + except (TypeError, ValueError): + return False + + +def _run_awaitable(value: Any) -> Any: + if not inspect.isawaitable(value): + return value + try: + asyncio.get_running_loop() + except RuntimeError: + return asyncio.run(value) + 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 new file mode 100644 index 000000000000..533604791a86 --- /dev/null +++ b/agent/relay_runtime.py @@ -0,0 +1,1002 @@ +"""Profile-scoped NeMo Relay runtimes owned by the Hermes agent core.""" + +from __future__ import annotations + +import atexit +import asyncio +import contextvars +import importlib +import inspect +import logging +import threading +import uuid +from dataclasses import dataclass, field +from typing import Any, Callable + +from hermes_constants import get_hermes_home + +logger = logging.getLogger(__name__) + +SESSION_SCOPE = "hermes.session" +TURN_SCOPE = "hermes.turn" +LOGICAL_LLM_SCOPE = "hermes.logical_llm_call" +RUNTIME_SCHEMA_KEY = "hermes.relay.schema_version" +RUNTIME_SCHEMA_VERSION = "hermes.relay.runtime.v1" +RUNTIME_INSTANCE_KEY = "hermes.relay.runtime_instance" +_PROFILE_KEY_CACHE: dict[str, str] = {} + + +@dataclass +class RelaySession: + """One isolated Relay scope stack owned by a Hermes session.""" + + session_id: str + parent_session_id: str = "" + lock: threading.RLock = field(default_factory=threading.RLock, repr=False) + closing: bool = False + handle: Any = None + context: contextvars.Context | None = None + + +class RelayRuntime: + """Own Relay session scopes independently of any exporter or plugin.""" + + def __init__(self, relay: Any = None, *, profile_key: str | None = None) -> None: + 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: dict[str, RelaySession] = {} + self._subagent_parents: dict[str, str] = {} + self._subagent_parent_handles: dict[str, Any] = {} + self._execution_consumers_lock = threading.RLock() + self._execution_consumers: set[str] = set() + self._shutdown_registered = True + atexit.register(self.shutdown) + + def retain_managed_execution(self, consumer: str) -> None: + """Keep managed LLM and tool execution active for one consumer.""" + if not consumer: + raise ValueError("Relay managed-execution consumer must not be empty") + with self._execution_consumers_lock: + self._execution_consumers.add(consumer) + + def release_managed_execution(self, consumer: str) -> None: + """Release a consumer's managed-execution requirement.""" + with self._execution_consumers_lock: + self._execution_consumers.discard(consumer) + + def managed_execution_enabled(self) -> bool: + """Return whether a Hermes-managed consumer needs the Relay pipeline.""" + with self._execution_consumers_lock: + return bool(self._execution_consumers) + + 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) + if not session_id: + return None + with self._sessions_lock: + 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 + 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: + session.handle = context.run( + self.relay.scope.push, + SESSION_SCOPE, + self.relay.ScopeType.Agent, + handle=parent_handle, + data=data, + input={}, + metadata=scope_metadata, + ) + except Exception: + session.context = None + raise + session.context = context + return session + + def register_subagent( + 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 + ): + 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) + if ( + turn is not None + and not turn.closed + and turn.handle is not None + and turn.lease.host is self + and turn.lease.session is not None + and turn.lease.session.session_id == parent_session_id + ): + parent_handle = turn.handle + with self._sessions_lock: + 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, + ) + + def unregister_subagent(self, event: dict[str, Any]) -> None: + """Close a delegated session and forget its parent relationship.""" + child_session_id = str(event.get("child_session_id") or "") + if not child_session_id: + return + self.close_session({"session_id": child_session_id}) + with self._sessions_lock: + self._subagent_parents.pop(child_session_id, None) + self._subagent_parent_handles.pop(child_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: + session = 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 run_in_session( + self, + session: RelaySession, + callback: Callable[..., Any], + *args: Any, + allow_closing: bool = False, + **kwargs: Any, + ) -> Any: + """Run a Relay operation against a session's isolated scope stack.""" + 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) + + 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. + return context.run(invoke) + + async def run_in_session_async( + self, + session: RelaySession, + callback: Callable[..., Any], + *args: Any, + allow_closing: bool = False, + **kwargs: Any, + ) -> Any: + """Create and await an operation inside the session's saved context.""" + 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) + + async def invoke() -> Any: + self.relay.get_scope_stack() + result = callback(*args, **kwargs) + if inspect.isawaitable(result): + return await result + return result + + task = context.run(asyncio.create_task, invoke()) + return await task + + def emit_mark( + 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) + if session is None: + return False + self.run_in_session( + session, + self.relay.scope.event, + name, + handle=session.handle, + data=data, + metadata=metadata, + ) + return True + + 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.""" + if not self.managed_execution_enabled(): + return args + 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, + ) + return result if isinstance(result, dict) else args + + def close_session(self, event: dict[str, Any]) -> None: + """Close one session scope and remove it from the core registry.""" + session_id = _session_id(event) + 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) + return + failures: list[str] = [] + with session.lock: + if session.closing: + return + session.closing = True + if session.handle is not None: + try: + self.run_in_session( + session, + self.relay.scope.pop, + session.handle, + output={}, + metadata={ + RUNTIME_SCHEMA_KEY: RUNTIME_SCHEMA_VERSION, + RUNTIME_INSTANCE_KEY: self.runtime_id, + }, + allow_closing=True, + ) + except Exception as exc: + failures.append(f"session scope close failed: {exc}") + try: + self.relay.subscribers.flush() + except Exception as exc: + failures.append(f"subscriber flush failed: {exc}") + 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), + ) + + def shutdown(self) -> None: + """Close all core-owned Relay session scopes.""" + with self._sessions_lock: + session_ids = list(self._sessions) + for session_id in session_ids: + self._safe(self.close_session, {"session_id": session_id}) + if self._shutdown_registered: + try: + atexit.unregister(self.shutdown) + except Exception: + pass + self._shutdown_registered = False + + @staticmethod + def _safe(callback: Callable[..., Any], *args: Any, **kwargs: Any) -> Any: + try: + return callback(*args, **kwargs) + except Exception: + logger.warning("Hermes Relay runtime operation failed", exc_info=True) + return None + + +@dataclass(frozen=True) +class NoopRelayRuntime: + """Explicit reduced-capability host for platforms without Relay wheels.""" + + 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], + ) -> dict[str, Any]: + del session_id, tool_name + return args + + @staticmethod + def retain_managed_execution(consumer: str) -> None: + del consumer + + @staticmethod + def release_managed_execution(consumer: str) -> None: + del consumer + + @staticmethod + def managed_execution_enabled() -> bool: + return False + + def shutdown(self) -> None: + """No resources are allocated on unsupported platforms.""" + + +RelayHost = RelayRuntime | NoopRelayRuntime + + +class RelayHostRegistry: + """Own exactly one Relay host for each canonical Hermes profile.""" + + def __init__(self) -> None: + self._lock = threading.RLock() + self._hosts: dict[str, RelayHost] = {} + + def for_profile( + self, + profile_key: str | None = None, + *, + create: bool = True, + ) -> RelayHost | None: + key = profile_key or current_profile_key() + host = self._hosts.get(key) + if host is not None or not create: + return host + with self._lock: + host = self._hosts.get(key) + if host is not None or not create: + return host + try: + host = RelayRuntime(profile_key=key) + except Exception as exc: + 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()) + self._hosts.clear() + for host in hosts: + host.shutdown() + + +HOST_REGISTRY = RelayHostRegistry() + + +@dataclass +class ConversationLease: + """A resumable reference to one profile-scoped conversation scope.""" + + profile_key: str + session_id: str + platform: str + host: RelayHost + session: RelaySession | None + parent_session_id: str = "" + released: bool = False + + +@dataclass +class RelayTurnContext: + """Runtime-only context for one Hermes turn or top-level task.""" + + lease: ConversationLease + turn_id: str + 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, + ) + _token: contextvars.Token[RelayTurnContext | None] | None = field( + default=None, + repr=False, + ) + _active_registered: bool = field(default=False, repr=False) + closed: bool = False + + +_CURRENT_TURN: contextvars.ContextVar[RelayTurnContext | None] = contextvars.ContextVar( + "hermes_relay_turn", default=None +) + + +class RelaySessionCoordinator: + """Own semantic conversation and turn lifetimes for Hermes core.""" + + 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._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], + ) -> 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: + 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, + ) + + 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) + if host is None: + host = NoopRelayRuntime(profile_key, "Relay host creation was disabled") + session = None + if isinstance(host, RelayRuntime): + try: + session_context = { + "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, + }, + metadata=metadata, + ) + else: + session = host.ensure_session( + {"session_id": session_id}, + metadata=metadata, + ) + except Exception: + logger.warning( + "Hermes Relay conversation initialization failed", + exc_info=True, + ) + return ConversationLease( + profile_key=profile_key, + session_id=session_id, + platform=platform, + host=host, + session=session, + parent_session_id=parent_session_id, + ) + + def begin_turn( + 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) + if isinstance(lease.host, RelayRuntime) and lease.session is not None: + try: + turn.handle = lease.host.run_in_session( + lease.session, + lease.host.relay.scope.push, + TURN_SCOPE, + lease.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", + }, + ) + except Exception: + logger.warning("Hermes Relay turn initialization failed", exc_info=True) + turn._token = _CURRENT_TURN.set(turn) + key = (lease.profile_key, lease.session_id) + with self._active_turns_lock: + self._active_turns.setdefault(key, set()).add(id(turn)) + turn._active_registered = True + return turn + + def end_turn( + 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 + 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: + try: + lease.host.run_in_session( + lease.session, + lease.host.relay.scope.pop, + turn.handle, + output={"outcome": outcome}, + metadata={ + RUNTIME_SCHEMA_KEY: RUNTIME_SCHEMA_VERSION, + RUNTIME_INSTANCE_KEY: lease.host.runtime_id, + }, + ) + except Exception: + logger.warning( + "Hermes Relay turn finalization failed", exc_info=True + ) + finally: + try: + # 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 + }) + except Exception: + logger.warning( + "Hermes Relay child conversation finalization failed", + exc_info=True, + ) + finally: + self._unregister_active_turn(turn) + self._reset_turn_context(turn) + + def has_active_turn(self, *, profile_key: str, session_id: str) -> bool: + """Return whether a turn is still running for one profile/session.""" + key = (profile_key, session_id) + with self._active_turns_lock: + return bool(self._active_turns.get(key)) + + def _unregister_active_turn(self, turn: RelayTurnContext) -> None: + if not turn._active_registered: + return + key = (turn.lease.profile_key, turn.lease.session_id) + with self._active_turns_lock: + active = self._active_turns.get(key) + if active is not None: + active.discard(id(turn)) + if not active: + self._active_turns.pop(key, None) + turn._active_registered = False + + def _reset_active_turns_for_tests(self) -> None: + with self._active_turns_lock: + self._active_turns.clear() + + def finish_logical_calls( + self, + turn: RelayTurnContext, + *, + outcome: str, + ) -> None: + """Close logical LLM children before sibling task aggregation scopes.""" + with turn.finalize_lock: + if turn.closed: + return + self._finish_logical_calls(turn, outcome=outcome) + + @staticmethod + def _finish_logical_calls( + turn: RelayTurnContext, + *, + outcome: str, + ) -> None: + lease = turn.lease + if not isinstance(lease.host, RelayRuntime) or lease.session 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] + try: + lease.host.run_in_session( + lease.session, + lease.host.relay.scope.pop, + logical_handle, + output={"outcome": outcome}, + metadata={ + RUNTIME_SCHEMA_KEY: RUNTIME_SCHEMA_VERSION, + RUNTIME_INSTANCE_KEY: lease.host.runtime_id, + }, + ) + except Exception: + with turn.logical_llm_lock: + # Relay scopes are stack-owned. If the newest remaining + # handle cannot close, older handles cannot close safely + # either, so 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", + exc_info=True, + ) + break + + @staticmethod + def _reset_turn_context(turn: RelayTurnContext) -> None: + """Reset the originating ContextVar token when called in that context.""" + if turn._token is None: + return + try: + _CURRENT_TURN.reset(turn._token) + except ValueError: + # A copied async/thread context may own terminal cleanup. Keep the + # token so the originating context can clear its stale reference. + return + turn._token = None + + @staticmethod + def release_conversation(lease: ConversationLease) -> None: + """Release a caller lease without closing a resumable conversation.""" + lease.released = True + + 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() + + +def current_turn() -> RelayTurnContext | None: + """Return the turn context inherited by current async and thread work.""" + return _CURRENT_TURN.get() + + +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 turn.closed or turn.lease.released: + return None + if turn.lease.profile_key != current_profile_key(): + return None + if session_id is not None and turn.lease.session_id != session_id: + return None + if isinstance(turn.lease.host, RelayRuntime): + if turn.lease.session is None: + return None + if turn.lease.host.get_session(turn.lease.session_id) is not turn.lease.session: + return None + return turn + + +def resolve_execution_context( + session_id: str, +) -> tuple[RelayRuntime | None, RelaySession | None, Any]: + """Resolve one active turn/session parent for managed Relay execution.""" + turn = active_turn(session_id) + if ( + turn is not None + and isinstance(turn.lease.host, RelayRuntime) + and turn.lease.session is not None + ): + session = turn.lease.session + return turn.lease.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 + # no-consumer path. + runtime = get_runtime(create=False) + if runtime is None: + return None, None, None + if not runtime.managed_execution_enabled(): + return None, None, None + session = runtime.get_session(session_id) + if session is None: + session = runtime.ensure_session({"session_id": session_id}) + 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: + """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, + ) + 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], +) -> dict[str, Any]: + """Return Relay-rewritten arguments at Hermes's authorization boundary.""" + if not session_id: + return args + runtime = get_runtime(create=False) + if runtime is None: + return args + return runtime.apply_tool_request_intercepts( + 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, +) -> bool: + """Match Relay's native callback wrapper without masking policy errors.""" + if relay_error is 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 + ) + + +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() + if not home.is_absolute(): + return str(home.resolve()) + raw = str(home) + cached = _PROFILE_KEY_CACHE.get(raw) + if cached is not None: + return cached + resolved = str(home.resolve()) + return _PROFILE_KEY_CACHE.setdefault(raw, resolved) + + +def _load_nemo_relay() -> Any: + """Load the binding only when a producer or consumer needs Relay.""" + return importlib.import_module("nemo_relay") + + +def _session_id(event: dict[str, Any]) -> str: + return str(event.get("session_id") or "") + + +def _reset_for_tests() -> None: + """Reset all profile-scoped Relay hosts for isolated tests.""" + SESSION_COORDINATOR._reset_active_turns_for_tests() + HOST_REGISTRY.shutdown_all() + _PROFILE_KEY_CACHE.clear() diff --git a/agent/relay_tools.py b/agent/relay_tools.py new file mode 100644 index 000000000000..5023df1bcf92 --- /dev/null +++ b/agent/relay_tools.py @@ -0,0 +1,123 @@ +"""Core NeMo Relay adapter for Hermes tool execution.""" + +from __future__ import annotations + +import asyncio +import contextvars +import inspect +import json +import logging +from collections.abc import Callable +from typing import Any + +from agent import relay_runtime + +logger = logging.getLogger(__name__) + + +def execute( + tool_name: str, + args: dict[str, Any], + callback: Callable[[dict[str, Any]], Any], + *, + session_id: str, + metadata: dict[str, Any] | None = None, +) -> tuple[Any, dict[str, Any]]: + """Run one tool call through Relay and return its final arguments.""" + 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(args), args + + observed_args = args + raw_result: dict[str, Any] = {} + callback_error: BaseException | None = None + callback_context = contextvars.copy_context() + + def invoke(next_args: Any) -> Any: + nonlocal callback_error, observed_args + observed_args = next_args if isinstance(next_args, dict) else args + try: + result = callback_context.copy().run(callback, observed_args) + except BaseException as exc: + callback_error = exc + raise + raw_result["value"] = result + raw_result["json"] = _jsonable(result) + return raw_result["json"] + + try: + managed = _run_awaitable( + runtime.run_in_session_async( + session, + runtime.relay.tools.execute, + tool_name, + _jsonable(args), + invoke, + handle=parent, + metadata=_jsonable(metadata or {}), + ) + ) + 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 ( + 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", + exc_info=True, + ) + return raw_result["value"], observed_args + raise + + if "value" in raw_result and _json_equal(managed, raw_result["json"]): + return raw_result["value"], observed_args + if isinstance(managed, str): + return managed, observed_args + return json.dumps(_jsonable(managed), ensure_ascii=False), observed_args + + +def _jsonable(value: Any) -> Any: + if value is None or isinstance(value, (str, int, float, bool)): + return value + if isinstance(value, dict): + return {str(key): _jsonable(item) for key, item in value.items()} + if isinstance(value, (list, tuple, set)): + return [_jsonable(item) for item in value] + model_dump = getattr(value, "model_dump", None) + if callable(model_dump): + try: + return _jsonable(model_dump(mode="json")) + except Exception: + pass + try: + return _jsonable(vars(value)) + except (TypeError, AttributeError): + return str(value) + + +def _json_equal(left: Any, right: Any) -> bool: + try: + return json.dumps( + _jsonable(left), sort_keys=True, separators=(",", ":") + ) == json.dumps(_jsonable(right), sort_keys=True, separators=(",", ":")) + except (TypeError, ValueError): + return left == right + + +def _run_awaitable(value: Any) -> Any: + if not inspect.isawaitable(value): + return value + try: + asyncio.get_running_loop() + except RuntimeError: + return asyncio.run(value) + raise RuntimeError( + "Synchronous Hermes Relay tool execution cannot run on an active event-loop thread" + ) diff --git a/agent/tool_executor.py b/agent/tool_executor.py index d32fe99c0c55..1dfcf1ebf6ef 100644 --- a/agent/tool_executor.py +++ b/agent/tool_executor.py @@ -20,6 +20,7 @@ import os import random import threading import time +from dataclasses import dataclass from typing import Any, Optional from agent.display import ( @@ -292,31 +293,63 @@ def _tool_search_scoped_names(agent) -> frozenset: return names -def _apply_tool_request_middleware_for_agent( - agent, - *, - function_name: str, - function_args: dict, - effective_task_id: str, - tool_call_id: str, -) -> tuple[dict, list[dict[str, Any]]]: - try: - from hermes_cli.middleware import apply_tool_request_middleware +@dataclass +class _ManagedToolResult: + result: Any + args: dict[str, Any] + middleware_trace: list[dict[str, Any]] + blocked: bool - result = apply_tool_request_middleware( - function_name, - function_args, - task_id=effective_task_id or "", - session_id=getattr(agent, "session_id", "") or "", - tool_call_id=tool_call_id or "", - turn_id=getattr(agent, "_current_turn_id", "") or "", - api_request_id=getattr(agent, "_current_api_request_id", "") or "", - ) - payload = result.payload if isinstance(result.payload, dict) else function_args - return payload, list(result.trace) - except Exception as exc: - logger.debug("tool_request middleware error: %s", exc) - return function_args, [] + +class _ConcurrentToolAuthorizationGate: + """Serialize policy prompts and exclude their queue from batch deadlines.""" + + def __init__(self) -> None: + self._serialization_lock = threading.Lock() + self._state_lock = threading.Lock() + self._pending = 0 + self._window_started: float | None = None + self._excluded_seconds = 0.0 + + def run(self, callback): + now = time.monotonic() + with self._state_lock: + if self._pending == 0: + self._window_started = now + self._pending += 1 + try: + with self._serialization_lock: + return callback() + finally: + now = time.monotonic() + with self._state_lock: + self._pending -= 1 + if self._pending == 0: + if self._window_started is not None: + self._excluded_seconds += max( + 0.0, now - self._window_started + ) + self._window_started = None + + def excluded_seconds(self) -> float: + """Return completed plus currently active authorization wait time.""" + now = time.monotonic() + with self._state_lock: + excluded = self._excluded_seconds + if self._window_started is not None: + excluded += max(0.0, now - self._window_started) + return excluded + + +def _managed_values( + outcome: _ManagedToolResult, +) -> tuple[Any, dict[str, Any], list[dict[str, Any]], bool]: + return ( + outcome.result, + outcome.args, + outcome.middleware_trace, + outcome.blocked, + ) def _run_agent_tool_execution_middleware( @@ -327,28 +360,271 @@ def _run_agent_tool_execution_middleware( effective_task_id: str, tool_call_id: str, execute, -) -> tuple[Any, dict]: - observed_args = function_args + scope_block: str | None = None, + display_index: int | None = None, + middleware_trace: list[dict[str, Any]] | None = None, + begin_execution=None, + authorization_gate: _ConcurrentToolAuthorizationGate | None = None, +) -> _ManagedToolResult: + """Run Relay rewrites before Hermes policy and dispatch exactly once.""" + from agent import relay_tools + from hermes_cli.middleware import ( + apply_tool_request_middleware, + run_tool_execution_middleware, + ) - def _execute(next_args: dict) -> Any: - nonlocal observed_args - observed_args = next_args if isinstance(next_args, dict) else function_args - return execute(observed_args) + trace = middleware_trace if middleware_trace is not None else [] + state = { + "args": function_args, + "middleware_trace": trace, + "blocked": False, + "dispatched": False, + } + dispatch_lock = threading.Lock() - from hermes_cli.middleware import run_tool_execution_middleware + def _authorized_dispatch(final_args: dict[str, Any]) -> Any: + with dispatch_lock: + if state["dispatched"]: + raise RuntimeError( + "Hermes tool execution callback invoked more than once" + ) + state["dispatched"] = True + state["blocked"] = False + state["args"] = final_args - result = run_tool_execution_middleware( + def _begin() -> None: + _begin_tool_execution( + agent, + function_name=function_name, + function_args=final_args, + effective_task_id=effective_task_id, + tool_call_id=tool_call_id, + display_index=display_index, + ) + + def _advance_start_order(callback=None) -> None: + if begin_execution is None: + if callback is not None: + callback() + return + begin_execution(callback) + + block_message = scope_block + block_error_type = "tool_scope_block" + if block_message is None: + block_error_type = "plugin_block" + + def _resolve_pre_tool_block(): + try: + from hermes_cli.plugins import resolve_pre_tool_block + + return resolve_pre_tool_block( + function_name, + final_args, + task_id=effective_task_id or "", + session_id=getattr(agent, "session_id", "") or "", + tool_call_id=tool_call_id or "", + turn_id=getattr(agent, "_current_turn_id", "") or "", + api_request_id=getattr(agent, "_current_api_request_id", "") + or "", + middleware_trace=list(state["middleware_trace"]), + ) + except Exception: + return None + + block_message = ( + _resolve_pre_tool_block() + if authorization_gate is None + else authorization_gate.run(_resolve_pre_tool_block) + ) + + guardrail_decision = None + if block_message is None: + guardrail_decision = agent._tool_guardrails.before_call( + function_name, final_args + ) + if guardrail_decision.allows_execution: + guardrail_decision = None + + if block_message is not None or guardrail_decision is not None: + _advance_start_order() + state["blocked"] = True + if block_message is not None: + result = json.dumps({"error": block_message}, ensure_ascii=False) + error_type = block_error_type + error_message = block_message + else: + result = agent._guardrail_block_result(guardrail_decision) + error_type = "guardrail_block" + error_message = ( + getattr(guardrail_decision, "message", None) + or "Tool blocked by guardrail policy" + ) + _emit_terminal_post_tool_call( + agent, + function_name=function_name, + function_args=final_args, + result=result, + effective_task_id=effective_task_id, + tool_call_id=tool_call_id, + status="blocked", + error_type=error_type, + error_message=error_message, + middleware_trace=list(state["middleware_trace"]), + ) + return result + + if function_name == "memory": + agent._turns_since_memory = 0 + elif function_name == "skill_manage": + agent._iters_since_skill = 0 + + _advance_start_order(_begin) + return execute(final_args) + + def _hermes_pipeline(relay_args: dict[str, Any]) -> Any: + request_result = apply_tool_request_middleware( + function_name, + relay_args, + skip_relay=True, + task_id=effective_task_id or "", + session_id=getattr(agent, "session_id", "") or "", + tool_call_id=tool_call_id or "", + turn_id=getattr(agent, "_current_turn_id", "") or "", + api_request_id=getattr(agent, "_current_api_request_id", "") or "", + ) + request_args = ( + request_result.payload + if isinstance(request_result.payload, dict) + else relay_args + ) + trace.clear() + trace.extend(request_result.trace) + return run_tool_execution_middleware( + function_name, + request_args, + lambda next_args: _authorized_dispatch( + next_args if isinstance(next_args, dict) else request_args + ), + original_args=function_args, + task_id=effective_task_id or "", + session_id=getattr(agent, "session_id", "") or "", + tool_call_id=tool_call_id or "", + turn_id=getattr(agent, "_current_turn_id", "") or "", + api_request_id=getattr(agent, "_current_api_request_id", "") or "", + ) + + result, _relay_args = relay_tools.execute( function_name, function_args, - _execute, - original_args=function_args, - task_id=effective_task_id or "", - session_id=getattr(agent, "session_id", "") or "", - tool_call_id=tool_call_id or "", - turn_id=getattr(agent, "_current_turn_id", "") or "", - api_request_id=getattr(agent, "_current_api_request_id", "") or "", + _hermes_pipeline, + session_id=str(getattr(agent, "session_id", "") or ""), + metadata={ + "task_id": effective_task_id or "", + "turn_id": getattr(agent, "_current_turn_id", "") or "", + "api_request_id": getattr(agent, "_current_api_request_id", "") or "", + "tool_call_id": tool_call_id or "", + }, ) - return result, observed_args + return _ManagedToolResult( + result=result, + args=state["args"], + middleware_trace=state["middleware_trace"], + blocked=bool(state["blocked"]), + ) + + +def _begin_tool_execution( + agent, + *, + function_name: str, + function_args: dict[str, Any], + effective_task_id: str, + tool_call_id: str, + display_index: int | None, +) -> None: + """Run user-visible and checkpoint preflight on final tool arguments.""" + if not agent.quiet_mode and getattr(agent, "tool_progress_mode", "all") != "off": + display_args = ( + _redact_tool_args_for_display(function_name, function_args) or function_args + ) + args_str = json.dumps(display_args, ensure_ascii=False) + prefix = f"Tool {display_index}" if display_index is not None else "Tool" + if agent.verbose_logging: + print(f" 📞 {prefix}: {function_name}({list(display_args.keys())})") + print( + agent._wrap_verbose( + "Args: ", json.dumps(display_args, indent=2, ensure_ascii=False) + ) + ) + else: + args_preview = ( + args_str[: agent.log_prefix_chars] + "..." + if len(args_str) > agent.log_prefix_chars + else args_str + ) + print( + f" 📞 {prefix}: {function_name}({list(function_args.keys())}) - " + f"{args_preview}" + ) + + agent._current_tool = function_name + agent._touch_activity(f"executing tool: {function_name}") + try: + from tools.environments.base import set_activity_callback + + set_activity_callback(agent._touch_activity) + except Exception: + pass + + if agent.tool_progress_callback: + try: + display_args = ( + _redact_tool_args_for_display(function_name, function_args) + or function_args + ) + preview = _build_tool_preview(function_name, display_args) + agent.tool_progress_callback( + "tool.started", function_name, preview, display_args + ) + except Exception as callback_error: + logging.debug("Tool progress callback error: %s", callback_error) + + if agent.tool_start_callback: + try: + display_args = ( + _redact_tool_args_for_display(function_name, function_args) + or function_args + ) + agent.tool_start_callback( + tool_call_id, function_name, display_args + ) + except Exception as callback_error: + logging.debug("Tool start callback error: %s", callback_error) + + if function_name in {"write_file", "patch"} and agent._checkpoint_mgr.enabled: + try: + _ensure_file_checkpoint( + agent, + function_name, + function_args, + effective_task_id, + ) + except Exception: + pass + + if function_name == "terminal" and agent._checkpoint_mgr.enabled: + try: + command = function_args.get("command", "") + if _is_destructive_command(command): + cwd = function_args.get("workdir") or os.getenv( + "TERMINAL_CWD", os.getcwd() + ) + agent._checkpoint_mgr.ensure_checkpoint( + cwd, f"before terminal: {command[:60]}" + ) + except Exception: + pass def execute_tool_calls_concurrent(agent, assistant_message, messages: list, effective_task_id: str, api_call_count: int = 0, *, finalize: bool = True) -> None: @@ -386,7 +662,9 @@ def execute_tool_calls_concurrent(agent, assistant_message, messages: list, effe return # ── Parse args + pre-execution bookkeeping ─────────────────────── - parsed_calls = [] # list of (tool_call, function_name, function_args, middleware_trace, block_result, blocked_by_guardrail) + # (tool call, resolved name, parsed args, middleware trace, parse error, + # tool-search scope block) + parsed_calls = [] for tool_call in tool_calls: function_name = tool_call.function.name @@ -402,17 +680,11 @@ def execute_tool_calls_concurrent(agent, assistant_message, messages: list, effe function_args, [], malformed_args_result, - False, + None, ) ) continue - # Reset nudge counters only for a structurally valid invocation. - if function_name == "memory": - agent._turns_since_memory = 0 - elif function_name == "skill_manage": - agent._iters_since_skill = 0 - # ── Tool Search unwrap ──────────────────────────────────────── # When the model invokes the tool_call bridge, peel it open so # every downstream check (checkpointing, guardrails, plugin @@ -446,167 +718,58 @@ def execute_tool_calls_concurrent(agent, assistant_message, messages: list, effe function_name = _underlying function_args = _underlying_args else: - _ts_scope_block = json.dumps({ - "error": ( - f"'{_underlying}' is not available in this session. " - "Use tool_search to find tools you can call." - ), - }, ensure_ascii=False) + _ts_scope_block = ( + f"'{_underlying}' is not available in this session. " + "Use tool_search to find tools you can call." + ) except Exception: pass - function_args, middleware_trace = _apply_tool_request_middleware_for_agent( - agent, - function_name=function_name, - function_args=function_args, - effective_task_id=effective_task_id, - tool_call_id=getattr(tool_call, "id", "") or "", + parsed_calls.append( + (tool_call, function_name, function_args, [], None, _ts_scope_block) ) - # ── Block evaluation (BEFORE checkpoint preflight) ─────────── - # We must know whether the tool will execute before touching - # checkpoint state (dedup slot, real snapshots). - block_result = None - blocked_by_guardrail = False - if _ts_scope_block is not None: - # Out-of-scope tool_call: reject before hooks/guardrails/dispatch. - block_result = _ts_scope_block - _emit_terminal_post_tool_call( - agent, - function_name=function_name, - function_args=function_args, - result=block_result, - effective_task_id=effective_task_id, - tool_call_id=getattr(tool_call, "id", "") or "", - status="blocked", - error_type="tool_scope_block", - error_message=_ts_scope_block, - middleware_trace=list(middleware_trace), - ) - else: - try: - from hermes_cli.plugins import resolve_pre_tool_block - block_message = resolve_pre_tool_block( - function_name, - function_args, - task_id=effective_task_id or "", - session_id=getattr(agent, "session_id", "") or "", - tool_call_id=getattr(tool_call, "id", "") or "", - turn_id=getattr(agent, "_current_turn_id", "") or "", - api_request_id=getattr(agent, "_current_api_request_id", "") or "", - middleware_trace=list(middleware_trace), - ) - except Exception: - block_message = None - - if block_message is not None: - block_result = json.dumps({"error": block_message}, ensure_ascii=False) - _emit_terminal_post_tool_call( - agent, - function_name=function_name, - function_args=function_args, - result=block_result, - effective_task_id=effective_task_id, - tool_call_id=getattr(tool_call, "id", "") or "", - status="blocked", - error_type="plugin_block", - error_message=block_message, - middleware_trace=list(middleware_trace), - ) - else: - guardrail_decision = agent._tool_guardrails.before_call(function_name, function_args) - if not guardrail_decision.allows_execution: - block_result = agent._guardrail_block_result(guardrail_decision) - blocked_by_guardrail = True - _emit_terminal_post_tool_call( - agent, - function_name=function_name, - function_args=function_args, - result=block_result, - effective_task_id=effective_task_id, - tool_call_id=getattr(tool_call, "id", "") or "", - status="blocked", - error_type="guardrail_block", - error_message=getattr(guardrail_decision, "message", None) or "Tool blocked by guardrail policy", - middleware_trace=list(middleware_trace), - ) - - # ── Checkpoint preflight (only for tools that will execute) ── - if block_result is None: - # Checkpoint for file-mutating tools - if function_name in {"write_file", "patch"} and agent._checkpoint_mgr.enabled: - try: - _ensure_file_checkpoint( - agent, - function_name, - function_args, - effective_task_id, - ) - except Exception: - pass - - # Checkpoint before destructive terminal commands - if function_name == "terminal" and agent._checkpoint_mgr.enabled: - try: - cmd = function_args.get("command", "") - if _is_destructive_command(cmd): - cwd = function_args.get("workdir") or os.getenv("TERMINAL_CWD", os.getcwd()) - agent._checkpoint_mgr.ensure_checkpoint( - cwd, f"before terminal: {cmd[:60]}" - ) - except Exception: - pass - - parsed_calls.append((tool_call, function_name, function_args, middleware_trace, block_result, blocked_by_guardrail)) - # ── Logging / callbacks ────────────────────────────────────────── tool_names_str = ", ".join(name for _, name, _, _, _, _ in parsed_calls) if not agent.quiet_mode and getattr(agent, "tool_progress_mode", "all") != "off": print(f" ⚡ Concurrent: {num_tools} tool calls — {tool_names_str}") - for i, (tc, name, args, middleware_trace, block_result, blocked_by_guardrail) in enumerate(parsed_calls, 1): - display_args = _redact_tool_args_for_display(name, args) or args - args_str = json.dumps(display_args, ensure_ascii=False) - if agent.verbose_logging: - print(f" 📞 Tool {i}: {name}({list(display_args.keys())})") - print(agent._wrap_verbose("Args: ", json.dumps(display_args, indent=2, ensure_ascii=False))) - else: - args_preview = args_str[:agent.log_prefix_chars] + "..." if len(args_str) > agent.log_prefix_chars else args_str - print(f" 📞 Tool {i}: {name}({list(args.keys())}) - {args_preview}") - - for tc, name, args, middleware_trace, block_result, blocked_by_guardrail in parsed_calls: - if block_result is not None: - continue - if agent.tool_progress_callback: - try: - display_args = _redact_tool_args_for_display(name, args) or args - preview = _build_tool_preview(name, display_args) - agent.tool_progress_callback("tool.started", name, preview, display_args) - except Exception as cb_err: - logging.debug(f"Tool progress callback error: {cb_err}") - - for tc, name, args, middleware_trace, block_result, blocked_by_guardrail in parsed_calls: - if block_result is not None: - continue - if agent.tool_start_callback: - try: - display_args = _redact_tool_args_for_display(name, args) or args - agent.tool_start_callback(tc.id, name, display_args) - except Exception as cb_err: - logging.debug(f"Tool start callback error: {cb_err}") # ── Concurrent execution ───────────────────────────────────────── # Each slot holds (function_name, function_args, function_result, duration, error_flag, blocked_flag, middleware_trace) results = [None] * num_tools - for i, (tc, name, args, middleware_trace, block_result, blocked_by_guardrail) in enumerate(parsed_calls): + for i, (tc, name, args, middleware_trace, block_result, _scope_block) in enumerate(parsed_calls): if block_result is not None: results[i] = (name, args, block_result, 0.0, True, True, middleware_trace) + start_condition = threading.Condition() + next_start_order = 0 + authorization_gate = _ConcurrentToolAuthorizationGate() + + def _begin_in_order(order: int, callback=None) -> None: + nonlocal next_start_order + with start_condition: + start_condition.wait_for(lambda: order == next_start_order) + try: + if callback is not None: + callback() + finally: + next_start_order += 1 + start_condition.notify_all() + # Touch activity before launching workers so the gateway knows # we're executing tools (not stuck). agent._current_tool = tool_names_str agent._touch_activity(f"executing {num_tools} tools concurrently: {tool_names_str}") - def _run_tool(index, tool_call, function_name, function_args, middleware_trace): + def _run_tool( + index, + tool_call, + function_name, + function_args, + middleware_trace, + scope_block, + start_order, + ): """Worker function executed in a thread.""" # Register this worker tid so the agent can fan out an interrupt # to it — see AIAgent.interrupt(). Must happen first thing, and @@ -636,18 +799,50 @@ def execute_tool_calls_concurrent(agent, assistant_message, messages: list, effe # ContextVars are propagated by propagate_context_to_thread() at the # submit site below (GHSA-qg5c-hvr5-hjgr, #13617). start = time.time() + blocked = False + start_advanced = False + + def _advance_start(callback=None) -> None: + nonlocal start_advanced + if start_advanced: + return + try: + _begin_in_order(start_order, callback) + finally: + start_advanced = True + try: try: - result = agent._invoke_tool( - function_name, - function_args, - effective_task_id, - tool_call.id, - messages=messages, - pre_tool_block_checked=True, - skip_tool_request_middleware=True, - tool_request_middleware_trace=list(middleware_trace), + def _execute(next_args: dict[str, Any]) -> Any: + return agent._invoke_tool( + function_name, + next_args, + effective_task_id, + tool_call.id, + messages=messages, + pre_tool_block_checked=True, + skip_tool_request_middleware=True, + skip_tool_execution_middleware=True, + tool_request_middleware_trace=list(middleware_trace), + ) + + managed = _run_agent_tool_execution_middleware( + agent, + function_name=function_name, + function_args=function_args, + effective_task_id=effective_task_id, + tool_call_id=getattr(tool_call, "id", "") or "", + execute=_execute, + scope_block=scope_block, + display_index=index + 1, + middleware_trace=middleware_trace, + begin_execution=_advance_start, + authorization_gate=authorization_gate, ) + result = managed.result + function_args = managed.args + middleware_trace = managed.middleware_trace + blocked = managed.blocked except KeyboardInterrupt: try: agent.interrupt("keyboard interrupt") @@ -664,7 +859,15 @@ def execute_tool_calls_concurrent(agent, assistant_message, messages: list, effe ) duration = time.time() - start logger.info("tool %s cancelled (%.2fs)", function_name, duration) - results[index] = (function_name, function_args, result, duration, True, False, middleware_trace) + results[index] = ( + function_name, + function_args, + result, + duration, + True, + False, + middleware_trace, + ) return except Exception as tool_error: result = f"Error executing tool '{function_name}': {tool_error}" @@ -675,8 +878,17 @@ def execute_tool_calls_concurrent(agent, assistant_message, messages: list, effe logger.info("tool %s failed (%.2fs): %s", function_name, duration, result[:200]) else: logger.info("tool %s completed (%.2fs, %d chars)", function_name, duration, len(result)) - results[index] = (function_name, function_args, result, duration, is_error, False, middleware_trace) + results[index] = ( + function_name, + function_args, + result, + duration, + is_error, + blocked, + middleware_trace, + ) finally: + _advance_start() # Tear down worker-tid tracking. Clear any interrupt bit we may # have set so the next task scheduled onto this recycled tid # starts with a clean slate. This MUST be in a finally block @@ -699,9 +911,11 @@ def execute_tool_calls_concurrent(agent, assistant_message, messages: list, effe try: runnable_calls = [ - (i, tc, name, args) - for i, (tc, name, args, middleware_trace, block_result, blocked_by_guardrail) in enumerate(parsed_calls) - if block_result is None + (i, tc, name, args, scope_block) + for i, (tc, name, args, _trace, parse_error, scope_block) in enumerate( + parsed_calls + ) + if parse_error is None ] futures = [] future_to_index = {} @@ -719,13 +933,22 @@ def execute_tool_calls_concurrent(agent, assistant_message, messages: list, effe executor = DaemonThreadPoolExecutor(max_workers=max_workers) abandon_executor = False try: - for submit_index, (i, tc, name, args) in enumerate(runnable_calls): + for submit_index, (i, tc, name, args, scope_block) in enumerate( + runnable_calls + ): # Propagate the agent turn's ContextVars (e.g. # _approval_session_key) AND thread-local approval/sudo # callbacks into the worker thread; clears callbacks on exit. try: f = executor.submit( - propagate_context_to_thread(_run_tool), i, tc, name, args, parsed_calls[i][3] + propagate_context_to_thread(_run_tool), + i, + tc, + name, + args, + parsed_calls[i][3], + scope_block, + submit_index, ) except RuntimeError as submit_error: if not _is_interpreter_shutdown_submit_error(submit_error): @@ -736,7 +959,13 @@ def execute_tool_calls_concurrent(agent, assistant_message, messages: list, effe "skipping %d unsubmitted tool(s)", len(skipped_calls), ) - for skipped_i, _tc, skipped_name, skipped_args in skipped_calls: + for ( + skipped_i, + _tc, + skipped_name, + skipped_args, + _scope_block, + ) in skipped_calls: if results[skipped_i] is None: middleware_trace = parsed_calls[skipped_i][3] result = ( @@ -766,7 +995,10 @@ def execute_tool_calls_concurrent(agent, assistant_message, messages: list, effe while True: wait_timeout = 5.0 if deadline is not None: - remaining = deadline - time.monotonic() + effective_deadline = ( + deadline + authorization_gate.excluded_seconds() + ) + remaining = effective_deadline - time.monotonic() if remaining <= 0: done, not_done = set(), { f for f in futures if not f.done() @@ -783,7 +1015,11 @@ def execute_tool_calls_concurrent(agent, assistant_message, messages: list, effe if not not_done: break - if deadline is not None and time.monotonic() >= deadline: + if ( + deadline is not None + and time.monotonic() + >= deadline + authorization_gate.excluded_seconds() + ): abandon_executor = True timed_out_indices = { future_to_index[f] @@ -863,7 +1099,9 @@ def execute_tool_calls_concurrent(agent, assistant_message, messages: list, effe spinner.stop(f"⚡ {completed}/{num_tools} tools completed in {total_dur:.1f}s total") # ── Post-execution: display per-tool results ───────────────────── - for i, (tc, name, args, middleware_trace, block_result, blocked_by_guardrail) in enumerate(parsed_calls): + for i, (tc, name, args, middleware_trace, _parse_error, _scope_block) in enumerate( + parsed_calls + ): r = results[i] blocked = False is_error = True @@ -923,6 +1161,8 @@ def execute_tool_calls_concurrent(agent, assistant_message, messages: list, effe tool_duration = 0.0 else: function_name, function_args, function_result, tool_duration, is_error, blocked, middleware_trace = r + name = function_name + args = function_args progress_function_name = function_name if blocked: effect_disposition = "none" @@ -1169,153 +1409,12 @@ def execute_tool_calls_sequential(agent, assistant_message, messages: list, effe except Exception: pass - function_args, middleware_trace = _apply_tool_request_middleware_for_agent( - agent, - function_name=function_name, - function_args=function_args, - effective_task_id=effective_task_id, - tool_call_id=getattr(tool_call, "id", "") or "", - ) - - # Check plugin hooks for a block directive before executing. - _block_msg: Optional[str] = None - _block_error_type = "plugin_block" - if _ts_scope_block is not None: - _block_msg = _ts_scope_block - _block_error_type = "tool_scope_block" - else: - try: - from hermes_cli.plugins import resolve_pre_tool_block - _block_msg = resolve_pre_tool_block( - function_name, - function_args, - task_id=effective_task_id or "", - session_id=getattr(agent, "session_id", "") or "", - tool_call_id=getattr(tool_call, "id", "") or "", - turn_id=getattr(agent, "_current_turn_id", "") or "", - api_request_id=getattr(agent, "_current_api_request_id", "") or "", - middleware_trace=list(middleware_trace), - ) - except Exception: - pass - - _guardrail_block_decision: ToolGuardrailDecision | None = None - if _block_msg is None: - guardrail_decision = agent._tool_guardrails.before_call(function_name, function_args) - if not guardrail_decision.allows_execution: - _guardrail_block_decision = guardrail_decision - - _execution_blocked = _block_msg is not None or _guardrail_block_decision is not None - - if _execution_blocked: - # Tool blocked by plugin or guardrail policy — skip counters, - # callbacks, checkpointing, activity mutation, and real execution. - pass - # Reset nudge counters when the relevant tool is actually used - elif function_name == "memory": - agent._turns_since_memory = 0 - elif function_name == "skill_manage": - agent._iters_since_skill = 0 - - if not agent.quiet_mode and getattr(agent, "tool_progress_mode", "all") != "off": - display_args = _redact_tool_args_for_display(function_name, function_args) or function_args - args_str = json.dumps(display_args, ensure_ascii=False) - if agent.verbose_logging: - print(f" 📞 Tool {i}: {function_name}({list(display_args.keys())})") - print(agent._wrap_verbose("Args: ", json.dumps(display_args, indent=2, ensure_ascii=False))) - else: - args_preview = args_str[:agent.log_prefix_chars] + "..." if len(args_str) > agent.log_prefix_chars else args_str - print(f" 📞 Tool {i}: {function_name}({list(function_args.keys())}) - {args_preview}") - - if not _execution_blocked: - agent._current_tool = function_name - agent._touch_activity(f"executing tool: {function_name}") - - # Set activity callback for long-running tool execution (terminal - # commands, etc.) so the gateway's inactivity monitor doesn't kill - # the agent while a command is running. - if not _execution_blocked: - try: - from tools.environments.base import set_activity_callback - set_activity_callback(agent._touch_activity) - except Exception: - pass - - if not _execution_blocked and agent.tool_progress_callback: - try: - display_args = _redact_tool_args_for_display(function_name, function_args) or function_args - preview = _build_tool_preview(function_name, display_args) - agent.tool_progress_callback("tool.started", function_name, preview, display_args) - except Exception as cb_err: - logging.debug(f"Tool progress callback error: {cb_err}") - - if not _execution_blocked and agent.tool_start_callback: - try: - display_args = _redact_tool_args_for_display(function_name, function_args) or function_args - agent.tool_start_callback(tool_call.id, function_name, display_args) - except Exception as cb_err: - logging.debug(f"Tool start callback error: {cb_err}") - - # Checkpoint: snapshot working dir before file-mutating tools - if not _execution_blocked and function_name in {"write_file", "patch"} and agent._checkpoint_mgr.enabled: - try: - _ensure_file_checkpoint( - agent, - function_name, - function_args, - effective_task_id, - ) - except Exception: - pass # never block tool execution - - # Checkpoint before destructive terminal commands - if not _execution_blocked and function_name == "terminal" and agent._checkpoint_mgr.enabled: - try: - cmd = function_args.get("command", "") - if _is_destructive_command(cmd): - cwd = function_args.get("workdir") or os.getenv("TERMINAL_CWD", os.getcwd()) - agent._checkpoint_mgr.ensure_checkpoint( - cwd, f"before terminal: {cmd[:60]}" - ) - except Exception: - pass # never block tool execution + middleware_trace: list[dict[str, Any]] = [] + _execution_blocked = False tool_start_time = time.time() - if _block_msg is not None: - # Tool blocked by plugin policy — return error without executing. - function_result = json.dumps({"error": _block_msg}, ensure_ascii=False) - tool_duration = 0.0 - _emit_terminal_post_tool_call( - agent, - function_name=function_name, - function_args=function_args, - result=function_result, - effective_task_id=effective_task_id, - tool_call_id=getattr(tool_call, "id", "") or "", - status="blocked", - error_type=_block_error_type, - error_message=_block_msg, - middleware_trace=list(middleware_trace), - ) - elif _guardrail_block_decision is not None: - # Tool blocked by tool-loop guardrail — synthesize exactly one - # tool result for the original tool_call_id without executing. - function_result = agent._guardrail_block_result(_guardrail_block_decision) - tool_duration = 0.0 - _emit_terminal_post_tool_call( - agent, - function_name=function_name, - function_args=function_args, - result=function_result, - effective_task_id=effective_task_id, - tool_call_id=getattr(tool_call, "id", "") or "", - status="blocked", - error_type="guardrail_block", - error_message=getattr(_guardrail_block_decision, "message", None) or "Tool blocked by guardrail policy", - middleware_trace=list(middleware_trace), - ) - elif function_name == "todo": + if function_name == "todo": def _execute(next_args: dict) -> Any: from tools.todo_tool import todo_tool as _todo_tool return _todo_tool( @@ -1323,14 +1422,16 @@ def execute_tool_calls_sequential(agent, assistant_message, messages: list, effe merge=next_args.get("merge", False), store=agent._todo_store, ) - function_result, function_args = _run_agent_tool_execution_middleware( + function_result, function_args, middleware_trace, _execution_blocked = _managed_values(_run_agent_tool_execution_middleware( agent, function_name=function_name, function_args=function_args, effective_task_id=effective_task_id, tool_call_id=getattr(tool_call, "id", "") or "", execute=_execute, - ) + scope_block=_ts_scope_block, + display_index=i, + )) tool_duration = time.time() - tool_start_time if agent._should_emit_quiet_tool_messages(): agent._vprint(f" {_get_cute_tool_message_impl('todo', function_args, tool_duration, result=function_result)}") @@ -1352,14 +1453,16 @@ def execute_tool_calls_sequential(agent, assistant_message, messages: list, effe db=session_db, current_session_id=agent.session_id, ) - function_result, function_args = _run_agent_tool_execution_middleware( + function_result, function_args, middleware_trace, _execution_blocked = _managed_values(_run_agent_tool_execution_middleware( agent, function_name=function_name, function_args=function_args, effective_task_id=effective_task_id, tool_call_id=getattr(tool_call, "id", "") or "", execute=_execute, - ) + scope_block=_ts_scope_block, + display_index=i, + )) tool_duration = time.time() - tool_start_time if agent._should_emit_quiet_tool_messages(): agent._vprint(f" {_get_cute_tool_message_impl('session_search', function_args, tool_duration, result=function_result)}") @@ -1389,14 +1492,16 @@ def execute_tool_calls_sequential(agent, assistant_message, messages: list, effe ), ) return result - function_result, function_args = _run_agent_tool_execution_middleware( + function_result, function_args, middleware_trace, _execution_blocked = _managed_values(_run_agent_tool_execution_middleware( agent, function_name=function_name, function_args=function_args, effective_task_id=effective_task_id, tool_call_id=getattr(tool_call, "id", "") or "", execute=_execute, - ) + scope_block=_ts_scope_block, + display_index=i, + )) tool_duration = time.time() - tool_start_time if agent._should_emit_quiet_tool_messages(): agent._vprint(f" {_get_cute_tool_message_impl('memory', function_args, tool_duration, result=function_result)}") @@ -1409,14 +1514,16 @@ def execute_tool_calls_sequential(agent, assistant_message, messages: list, effe multi_select=next_args.get("multi_select", False), callback=agent.clarify_callback, ) - function_result, function_args = _run_agent_tool_execution_middleware( + function_result, function_args, middleware_trace, _execution_blocked = _managed_values(_run_agent_tool_execution_middleware( agent, function_name=function_name, function_args=function_args, effective_task_id=effective_task_id, tool_call_id=getattr(tool_call, "id", "") or "", execute=_execute, - ) + scope_block=_ts_scope_block, + display_index=i, + )) tool_duration = time.time() - tool_start_time if agent._should_emit_quiet_tool_messages(): agent._vprint(f" {_get_cute_tool_message_impl('clarify', function_args, tool_duration, result=function_result)}") @@ -1428,14 +1535,16 @@ def execute_tool_calls_sequential(agent, assistant_message, messages: list, effe count=next_args.get("count"), callback=getattr(agent, "read_terminal_callback", None), ) - function_result, function_args = _run_agent_tool_execution_middleware( + function_result, function_args, middleware_trace, _execution_blocked = _managed_values(_run_agent_tool_execution_middleware( agent, function_name=function_name, function_args=function_args, effective_task_id=effective_task_id, tool_call_id=getattr(tool_call, "id", "") or "", execute=_execute, - ) + scope_block=_ts_scope_block, + display_index=i, + )) tool_duration = time.time() - tool_start_time if agent._should_emit_quiet_tool_messages(): agent._vprint(f" {_get_cute_tool_message_impl('read_terminal', function_args, tool_duration, result=function_result)}") @@ -1460,14 +1569,16 @@ def execute_tool_calls_sequential(agent, assistant_message, messages: list, effe try: def _execute(next_args: dict) -> Any: return agent._dispatch_delegate_task(next_args) - function_result, function_args = _run_agent_tool_execution_middleware( + function_result, function_args, middleware_trace, _execution_blocked = _managed_values(_run_agent_tool_execution_middleware( agent, function_name=function_name, function_args=function_args, effective_task_id=effective_task_id, tool_call_id=getattr(tool_call, "id", "") or "", execute=_execute, - ) + scope_block=_ts_scope_block, + display_index=i, + )) _delegate_result = function_result finally: agent._delegate_spinner = None @@ -1491,14 +1602,16 @@ def execute_tool_calls_sequential(agent, assistant_message, messages: list, effe try: def _execute(next_args: dict) -> Any: return agent.context_compressor.handle_tool_call(function_name, next_args, messages=messages) - function_result, function_args = _run_agent_tool_execution_middleware( + function_result, function_args, middleware_trace, _execution_blocked = _managed_values(_run_agent_tool_execution_middleware( agent, function_name=function_name, function_args=function_args, effective_task_id=effective_task_id, tool_call_id=getattr(tool_call, "id", "") or "", execute=_execute, - ) + scope_block=_ts_scope_block, + display_index=i, + )) _ce_result = function_result except Exception as tool_error: function_result = json.dumps({"error": f"Context engine tool '{function_name}' failed: {tool_error}"}) @@ -1525,14 +1638,16 @@ def execute_tool_calls_sequential(agent, assistant_message, messages: list, effe try: def _execute(next_args: dict) -> Any: return agent._memory_manager.handle_tool_call(function_name, next_args) - function_result, function_args = _run_agent_tool_execution_middleware( + function_result, function_args, middleware_trace, _execution_blocked = _managed_values(_run_agent_tool_execution_middleware( agent, function_name=function_name, function_args=function_args, effective_task_id=effective_task_id, tool_call_id=getattr(tool_call, "id", "") or "", execute=_execute, - ) + scope_block=_ts_scope_block, + display_index=i, + )) _mem_result = function_result except Exception as tool_error: function_result = json.dumps({"error": f"Memory tool '{function_name}' failed: {tool_error}"}) @@ -1555,18 +1670,46 @@ def execute_tool_calls_sequential(agent, assistant_message, messages: list, effe spinner.start() _spinner_result = None try: - function_result = _ra().handle_function_call( - function_name, function_args, effective_task_id, - tool_call_id=tool_call.id, - session_id=agent.session_id or "", - turn_id=getattr(agent, "_current_turn_id", "") or "", - api_request_id=getattr(agent, "_current_api_request_id", "") or "", - enabled_tools=list(agent.valid_tool_names) if agent.valid_tool_names else None, - skip_pre_tool_call_hook=True, - skip_tool_request_middleware=True, - enabled_toolsets=getattr(agent, "enabled_toolsets", None), - disabled_toolsets=getattr(agent, "disabled_toolsets", None), - tool_request_middleware_trace=list(middleware_trace), + def _execute(next_args: dict) -> Any: + return _ra().handle_function_call( + function_name, + next_args, + effective_task_id, + tool_call_id=tool_call.id, + session_id=agent.session_id or "", + turn_id=getattr(agent, "_current_turn_id", "") or "", + api_request_id=getattr(agent, "_current_api_request_id", "") + or "", + enabled_tools=( + list(agent.valid_tool_names) + if agent.valid_tool_names + else None + ), + skip_pre_tool_call_hook=True, + skip_tool_request_middleware=True, + skip_tool_execution_middleware=True, + tool_request_middleware_trace=list(middleware_trace), + enabled_toolsets=getattr(agent, "enabled_toolsets", None), + disabled_toolsets=getattr(agent, "disabled_toolsets", None), + ) + + ( + function_result, + function_args, + middleware_trace, + _execution_blocked, + ) = _managed_values( + _run_agent_tool_execution_middleware( + agent, + function_name=function_name, + function_args=function_args, + effective_task_id=effective_task_id, + tool_call_id=getattr(tool_call, "id", "") or "", + execute=_execute, + scope_block=_ts_scope_block, + display_index=i, + middleware_trace=middleware_trace, + ) ) _spinner_result = function_result except KeyboardInterrupt: @@ -1597,18 +1740,46 @@ def execute_tool_calls_sequential(agent, assistant_message, messages: list, effe agent._vprint(f" {cute_msg}") else: try: - function_result = _ra().handle_function_call( - function_name, function_args, effective_task_id, - tool_call_id=tool_call.id, - session_id=agent.session_id or "", - turn_id=getattr(agent, "_current_turn_id", "") or "", - api_request_id=getattr(agent, "_current_api_request_id", "") or "", - enabled_tools=list(agent.valid_tool_names) if agent.valid_tool_names else None, - skip_pre_tool_call_hook=True, - skip_tool_request_middleware=True, - enabled_toolsets=getattr(agent, "enabled_toolsets", None), - disabled_toolsets=getattr(agent, "disabled_toolsets", None), - tool_request_middleware_trace=list(middleware_trace), + def _execute(next_args: dict) -> Any: + return _ra().handle_function_call( + function_name, + next_args, + effective_task_id, + tool_call_id=tool_call.id, + session_id=agent.session_id or "", + turn_id=getattr(agent, "_current_turn_id", "") or "", + api_request_id=getattr(agent, "_current_api_request_id", "") + or "", + enabled_tools=( + list(agent.valid_tool_names) + if agent.valid_tool_names + else None + ), + skip_pre_tool_call_hook=True, + skip_tool_request_middleware=True, + skip_tool_execution_middleware=True, + tool_request_middleware_trace=list(middleware_trace), + enabled_toolsets=getattr(agent, "enabled_toolsets", None), + disabled_toolsets=getattr(agent, "disabled_toolsets", None), + ) + + ( + function_result, + function_args, + middleware_trace, + _execution_blocked, + ) = _managed_values( + _run_agent_tool_execution_middleware( + agent, + function_name=function_name, + function_args=function_args, + effective_task_id=effective_task_id, + tool_call_id=getattr(tool_call, "id", "") or "", + execute=_execute, + scope_block=_ts_scope_block, + display_index=i, + middleware_trace=middleware_trace, + ) ) except KeyboardInterrupt: _emit_cancelled_terminal_post_tool_call( diff --git a/agent/turn_context.py b/agent/turn_context.py index 94dc1660d1a7..e080d6a5d969 100644 --- a/agent/turn_context.py +++ b/agent/turn_context.py @@ -438,7 +438,12 @@ def build_turn_context( # Generate unique task_id if not provided to isolate VMs between tasks. effective_task_id = task_id or str(uuid.uuid4()) agent._current_task_id = effective_task_id - turn_id = f"{agent.session_id or 'session'}:{effective_task_id}:{uuid.uuid4().hex[:8]}" + turn_id = str(getattr(agent, "_relay_pending_turn_id", "") or "") + if not turn_id: + turn_id = ( + f"{agent.session_id or 'session'}:{effective_task_id}:{uuid.uuid4().hex[:8]}" + ) + agent._relay_pending_turn_id = None agent._current_turn_id = turn_id agent._current_api_request_id = "" # Tripwire: warn (with both turn ids) when this turn starts before the @@ -1045,7 +1050,7 @@ def build_turn_context( # Plugin hook: pre_llm_call (context injected into user message, not system prompt). plugin_user_context = "" try: - from hermes_cli.plugins import invoke_hook as _invoke_hook + from hermes_cli.lifecycle import invoke_hook as _invoke_hook _pre_results = _invoke_hook( "pre_llm_call", session_id=agent.session_id, @@ -1056,6 +1061,7 @@ def build_turn_context( is_first_turn=(not bool(conversation_history)), model=agent.model, platform=getattr(agent, "platform", None) or "", + parent_session_id=getattr(agent, "_parent_session_id", None) or "", sender_id=getattr(agent, "_user_id", None) or "", ) _ctx_parts: list[str] = [] diff --git a/agent/turn_finalizer.py b/agent/turn_finalizer.py index 4e2d318b2e2c..0f636c1dd241 100644 --- a/agent/turn_finalizer.py +++ b/agent/turn_finalizer.py @@ -488,7 +488,7 @@ def finalize_turn( # First hook to return a string wins; None/empty return leaves text unchanged. if final_response and not interrupted: try: - from hermes_cli.plugins import invoke_hook as _invoke_hook + from hermes_cli.lifecycle import invoke_hook as _invoke_hook _transform_results = _invoke_hook( "transform_llm_output", response_text=final_response, @@ -510,7 +510,7 @@ def finalize_turn( # to an external memory system). if final_response and not interrupted: try: - from hermes_cli.plugins import invoke_hook as _invoke_hook + from hermes_cli.lifecycle import invoke_hook as _invoke_hook _invoke_hook( "post_llm_call", session_id=agent.session_id, @@ -669,14 +669,16 @@ def finalize_turn( # Fired at the very end of every run_conversation call. # Plugins can use this for cleanup, flushing buffers, etc. try: - from hermes_cli.plugins import invoke_hook as _invoke_hook + from hermes_cli.lifecycle import invoke_hook as _invoke_hook _invoke_hook( "on_session_end", session_id=agent.session_id, task_id=effective_task_id, turn_id=turn_id, completed=completed, + failed=failed, interrupted=interrupted, + turn_exit_reason=_turn_exit_reason, model=agent.model, platform=getattr(agent, "platform", None) or "", ) diff --git a/agent/usage_pricing.py b/agent/usage_pricing.py index 9825e4a27bd5..b7982e46ab05 100644 --- a/agent/usage_pricing.py +++ b/agent/usage_pricing.py @@ -511,15 +511,93 @@ _OFFICIAL_DOCS_PRICING: Dict[tuple[str, str], PricingEntry] = { pricing_version="deepseek-pricing-2026-07", ), # Google Gemini + ( + "google", + "gemini-3.6-flash", + ): PricingEntry( + input_cost_per_million=Decimal("1.50"), + output_cost_per_million=Decimal("7.50"), + cache_read_cost_per_million=Decimal("0.15"), + source="official_docs_snapshot", + source_url="https://ai.google.dev/gemini-api/docs/pricing", + pricing_version="google-pricing-2026-07-28", + ), + ( + "google", + "gemini-3.5-flash", + ): PricingEntry( + input_cost_per_million=Decimal("1.50"), + output_cost_per_million=Decimal("9.00"), + cache_read_cost_per_million=Decimal("0.15"), + source="official_docs_snapshot", + source_url="https://ai.google.dev/pricing", + pricing_version="google-pricing-2026-07-07", + ), + ( + "google", + "gemini-3.5-flash-lite", + ): PricingEntry( + input_cost_per_million=Decimal("0.30"), + output_cost_per_million=Decimal("2.50"), + cache_read_cost_per_million=Decimal("0.03"), + source="official_docs_snapshot", + source_url="https://ai.google.dev/gemini-api/docs/pricing", + pricing_version="google-pricing-2026-07-28", + ), + ( + "google", + "gemini-3.1-pro", + ): PricingEntry( + input_cost_per_million=Decimal("2.00"), + output_cost_per_million=Decimal("12.00"), + cache_read_cost_per_million=Decimal("0.20"), + source="official_docs_snapshot", + source_url="https://ai.google.dev/pricing", + pricing_version="google-pricing-2026-07-07", + ), + ( + "google", + "gemini-3.1-flash-lite", + ): PricingEntry( + input_cost_per_million=Decimal("0.25"), + output_cost_per_million=Decimal("1.50"), + cache_read_cost_per_million=Decimal("0.025"), + source="official_docs_snapshot", + source_url="https://ai.google.dev/pricing", + pricing_version="google-pricing-2026-07-07", + ), + ( + "google", + "gemini-3-pro-preview", + ): PricingEntry( + input_cost_per_million=Decimal("2.00"), + output_cost_per_million=Decimal("12.00"), + cache_read_cost_per_million=Decimal("0.20"), + source="official_docs_snapshot", + source_url="https://ai.google.dev/pricing", + pricing_version="google-pricing-2026-07-07", + ), + ( + "google", + "gemini-3-flash-preview", + ): PricingEntry( + input_cost_per_million=Decimal("0.50"), + output_cost_per_million=Decimal("3.00"), + cache_read_cost_per_million=Decimal("0.05"), + source="official_docs_snapshot", + source_url="https://ai.google.dev/pricing", + pricing_version="google-pricing-2026-07-07", + ), ( "google", "gemini-2.5-pro", ): PricingEntry( input_cost_per_million=Decimal("1.25"), output_cost_per_million=Decimal("10.00"), + cache_read_cost_per_million=Decimal("0.125"), source="official_docs_snapshot", source_url="https://ai.google.dev/pricing", - pricing_version="google-pricing-2026-03-16", + pricing_version="google-pricing-2026-07-07", ), ( "google", @@ -527,9 +605,10 @@ _OFFICIAL_DOCS_PRICING: Dict[tuple[str, str], PricingEntry] = { ): PricingEntry( input_cost_per_million=Decimal("0.15"), output_cost_per_million=Decimal("0.60"), + cache_read_cost_per_million=Decimal("0.015"), source="official_docs_snapshot", source_url="https://ai.google.dev/pricing", - pricing_version="google-pricing-2026-03-16", + pricing_version="google-pricing-2026-07-07", ), ( "google", @@ -537,9 +616,10 @@ _OFFICIAL_DOCS_PRICING: Dict[tuple[str, str], PricingEntry] = { ): PricingEntry( input_cost_per_million=Decimal("0.10"), output_cost_per_million=Decimal("0.40"), + cache_read_cost_per_million=Decimal("0.01"), source="official_docs_snapshot", source_url="https://ai.google.dev/pricing", - pricing_version="google-pricing-2026-03-16", + pricing_version="google-pricing-2026-07-07", ), # AWS Bedrock — pricing per the Bedrock pricing page. # Bedrock charges the same per-token rates as the model provider but @@ -878,6 +958,18 @@ for _base_56 in ("gpt-5.6-sol", "gpt-5.6-terra", "gpt-5.6-luna"): ] del _base_56 +# The direct Gemini provider currently exposes preview IDs for these two +# models. Keep the official snapshot keyed by both their documented stable +# names and the provider's emitted IDs so a catalog selection is billable. +for _alias, _canonical in { + "gemini-3.1-pro-preview": "gemini-3.1-pro", + "gemini-3.1-flash-lite-preview": "gemini-3.1-flash-lite", +}.items(): + _OFFICIAL_DOCS_PRICING[("google", _alias)] = _OFFICIAL_DOCS_PRICING[ + ("google", _canonical) + ] +del _alias, _canonical + def _to_decimal(value: Any) -> Optional[Decimal]: if value is None: @@ -925,11 +1017,17 @@ def resolve_billing_route( return BillingRoute(provider="openai", model=model.split("/")[-1], base_url=base_url or "", billing_mode="official_docs_snapshot") if provider_name in {"minimax", "minimax-cn"}: return BillingRoute(provider=provider_name, model=model.split("/")[-1], base_url=base_url or "", billing_mode="official_docs_snapshot") - # Vertex AI hosts the same Gemini models as Google AI Studio; price them - # off the gemini official-docs snapshot. Strip the "google/" vendor prefix - # the OpenAI-compat endpoint requires so the pricing key matches. - if provider_name == "vertex" or base_url_host_matches(base_url or "", "aiplatform.googleapis.com"): - return BillingRoute(provider="gemini", model=model.split("/")[-1], base_url=base_url or "", billing_mode="official_docs_snapshot") + # Google AI Studio (Gemini) and Vertex AI host the same Gemini models. + # Price them off the official docs snapshot — the pricing keys are + # keyed on provider='google', so normalize every Google-flavored + # provider name/host onto it. Strip the "google/" vendor prefix the + # Vertex OpenAI-compat endpoint requires so the pricing key matches. + if ( + provider_name in {"google", "gemini", "vertex", "google-gemini", "google-ai-studio", "google-vertex", "vertex-ai"} + or base_url_host_matches(base_url or "", "aiplatform.googleapis.com") + or base_url_host_matches(base_url or "", "generativelanguage.googleapis.com") + ): + return BillingRoute(provider="google", model=model.split("/")[-1], base_url=base_url or "", billing_mode="official_docs_snapshot") if provider_name == "fireworks" or base_url_host_matches(base_url or "", "api.fireworks.ai"): # Fireworks model ids look like accounts/fireworks/models/; # rsplit("/", 1)[-1] yields just which is what the dict keys on. diff --git a/apps/desktop/electron/main.ts b/apps/desktop/electron/main.ts index dfb4747f0392..b9e966067a38 100644 --- a/apps/desktop/electron/main.ts +++ b/apps/desktop/electron/main.ts @@ -9272,8 +9272,8 @@ ipcMain.handle('hermes:window:openInstance', async () => { // shortcuts and the View menu. Reads and writes target the asking window. ipcMain.handle('hermes:zoom:get', event => { const window = BrowserWindow.fromWebContents(event.sender) - const level = - window && !window.isDestroyed() ? window.webContents.getZoomLevel() : DEFAULT_ZOOM_LEVEL + + const level = window && !window.isDestroyed() ? window.webContents.getZoomLevel() : DEFAULT_ZOOM_LEVEL return { level, percent: zoomLevelToPercent(level) } }) diff --git a/apps/desktop/src/app/chat/composer/hooks/use-voice-conversation.ts b/apps/desktop/src/app/chat/composer/hooks/use-voice-conversation.ts index 59ba4ee3c8ef..674b19bbb739 100644 --- a/apps/desktop/src/app/chat/composer/hooks/use-voice-conversation.ts +++ b/apps/desktop/src/app/chat/composer/hooks/use-voice-conversation.ts @@ -11,6 +11,7 @@ import { } from '@/lib/voice-playback' import { isVoiceStopCommand } from '@/lib/voice-stop-word' import { notify, notifyError } from '@/store/notifications' +import { $voicePlayback } from '@/store/voice-playback' import { useMicRecorder } from './use-mic-recorder' @@ -61,6 +62,7 @@ export function useVoiceConversation({ const speechSessionRef = useRef(null) const stopBargeMonitorRef = useRef<(() => void) | null>(null) const bargeCapturePendingRef = useRef(false) + const speechStartSequenceRef = useRef(0) const enabledRef = useRef(enabled) const mutedRef = useRef(muted) const busyRef = useRef(busy) @@ -257,7 +259,16 @@ export function useVoiceConversation({ dropSpeechSession() - if (enabledRef.current) { + // If stopVoicePlayback() was called externally (Stop button, end), the + // voice-playback sequence has advanced past what we captured at speech + // start — don't auto-start the next sentence, the user chose to stop. + const stoppedByUser = + speechStartSequenceRef.current > 0 && + $voicePlayback.get().sequence > speechStartSequenceRef.current + + speechStartSequenceRef.current = 0 + + if (enabledRef.current && !stoppedByUser) { pendingStartRef.current = true } @@ -387,6 +398,8 @@ export function useVoiceConversation({ barged = true }) + speechStartSequenceRef.current = $voicePlayback.get().sequence + void playSpeechText(response.text, { source: 'voice-conversation' }) .catch(error => notifyError(error, voiceCopy.playbackFailed)) .finally(() => { @@ -411,6 +424,7 @@ export function useVoiceConversation({ (responseId: string) => { responseIdRef.current = responseId spokenSourceLengthRef.current = 0 + speechStartSequenceRef.current = $voicePlayback.get().sequence setStatus('speaking') let barged = false diff --git a/apps/desktop/src/app/contrib/dev/credits-notice-demo.ts b/apps/desktop/src/app/contrib/dev/credits-notice-demo.ts index 1766489818b8..a2ba8080da1e 100644 --- a/apps/desktop/src/app/contrib/dev/credits-notice-demo.ts +++ b/apps/desktop/src/app/contrib/dev/credits-notice-demo.ts @@ -25,7 +25,7 @@ interface NoticeStep { } // Walks the same lifecycle the Nous credits tracker drives: usage escalates in -// place (50→75→90, one key), then grant-spent, then the depleted/restored pair. +// place (50→75→90, one key), then the depleted/restored pair. // Wraps around. These are all separate SHOW steps; the stepper auto-clears the // previous notice when the key changes, so the demo shows one toast at a time // (real usage CAN stack these, but that's noise when you're eyeballing a single @@ -34,7 +34,6 @@ const STEPS: readonly NoticeStep[] = [ { key: 'credits.usage', kind: 'sticky', level: 'info', text: "• You've used $110.00 of your $220.00 cap" }, { key: 'credits.usage', kind: 'sticky', level: 'warn', text: "⚠ You've used $165.00 of your $220.00 cap" }, { key: 'credits.usage', kind: 'sticky', level: 'warn', text: "⚠ You've used $198.00 of your $220.00 cap" }, - { key: 'credits.grant_spent', kind: 'sticky', level: 'info', text: '• Grant spent · $12.00 top-up left' }, { key: 'credits.depleted', kind: 'sticky', level: 'error', text: '✕ Credit access paused · run /topup to top up' }, { key: 'credits.restored', kind: 'ttl', level: 'success', text: '✓ Credit access restored', ttl_ms: 8000 } ] diff --git a/apps/desktop/src/app/session/hooks/use-session-actions/index.ts b/apps/desktop/src/app/session/hooks/use-session-actions/index.ts index 9ec65d78d028..266a33b65fd2 100644 --- a/apps/desktop/src/app/session/hooks/use-session-actions/index.ts +++ b/apps/desktop/src/app/session/hooks/use-session-actions/index.ts @@ -558,6 +558,7 @@ export function useSessionActions({ resetViewSync() setSelectedStoredSessionId(storedSessionId) selectedStoredSessionIdRef.current = storedSessionId + // A session is EITHER the main thread OR a tile — never both. openSessionTile // enforces this from the tile side (it refuses to tile the selected session); // this enforces it from the main side. Loading an existing session into main @@ -570,6 +571,7 @@ export function useSessionActions({ if ($sessionTiles.get().some(t => t.storedSessionId === storedSessionId)) { closeSessionTile(storedSessionId) } + // Optimistically clear any prior resume-failure latch for this session: // we're attempting a fresh resume, so the self-heal in use-route-resume // must not keep treating it as stranded. It's re-armed below only if THIS diff --git a/apps/desktop/src/app/settings/config-field.tsx b/apps/desktop/src/app/settings/config-field.tsx index aac75a2ee971..7461b4b0c08c 100644 --- a/apps/desktop/src/app/settings/config-field.tsx +++ b/apps/desktop/src/app/settings/config-field.tsx @@ -13,6 +13,7 @@ import { CONTROL_TEXT, EMPTY_SELECT_VALUE, FIELD_DESCRIPTIONS, FIELD_LABELS, FRE import { FallbackModelsField } from './fallback-models-field' import { fieldCopyForSchemaKey } from './field-copy' import { ListRow } from './primitives' +import { SearchableSelect } from './searchable-select' /** * One generic config row: label + description resolved from the i18n field @@ -93,6 +94,23 @@ export function ConfigField({ const selectOptions = enumOptions ?? (schema.type === 'select' ? (schema.options ?? []).map(String) : undefined) + // Large closed-world lists (e.g. ~590 IANA timezones) get a searchable + // Popover + cmdk combobox instead of a closed Select dropdown. The schema + // opt-in via `searchable: true` keeps this deterministic — no field + // accidentally triggers based on dynamic option count. + if (selectOptions && schema.searchable) { + return row( + onChange(next)} + options={selectOptions.filter(o => o !== '')} + placeholder={c.searchPlaceholder} + value={String(value ?? '')} + /> + ) + } + // Voice/model name fields are open-world (custom voice IDs, cloned voices, // brand-new model names) — render a free-input combobox where the known // options are datalist suggestions instead of a closed Select gate. diff --git a/apps/desktop/src/app/settings/constants.ts b/apps/desktop/src/app/settings/constants.ts index f041b2992470..3e12e08e70d7 100644 --- a/apps/desktop/src/app/settings/constants.ts +++ b/apps/desktop/src/app/settings/constants.ts @@ -564,7 +564,7 @@ export const FIELD_DESCRIPTIONS: Record = defineFieldCopy({ repoScanRoots: 'Folders to scan. Leave empty to scan your home directory.', repoScanExcludePaths: 'Folders and their descendants to skip during repository discovery.' }, - timezone: 'Used when Hermes needs local time context. Blank uses the system timezone.', + timezone: 'IANA timezone identifier. Blank uses the system timezone.', agent: { imageInputMode: 'Controls how image attachments are sent to the model.', maxTurns: 'Upper bound for tool-calling turns before Hermes stops a run.' diff --git a/apps/desktop/src/app/settings/searchable-select.test.tsx b/apps/desktop/src/app/settings/searchable-select.test.tsx new file mode 100644 index 000000000000..2e87c15c92f0 --- /dev/null +++ b/apps/desktop/src/app/settings/searchable-select.test.tsx @@ -0,0 +1,137 @@ +import { cleanup, fireEvent, render, screen } from '@testing-library/react' +import { afterEach, beforeAll, describe, expect, it, vi } from 'vitest' + +import type { ConfigFieldSchema } from '@/types/hermes' + +import { ConfigField } from './config-field' +import { rankSearchOption, SearchableSelect } from './searchable-select' + +// Radix Popover + cmdk call scrollIntoView / pointer-capture / ResizeObserver +// APIs jsdom lacks. +class TestResizeObserver { + disconnect() {} + observe() {} + unobserve() {} +} + +beforeAll(() => { + vi.stubGlobal('ResizeObserver', TestResizeObserver) + Element.prototype.scrollIntoView = vi.fn() + Element.prototype.hasPointerCapture = vi.fn(() => false) + Element.prototype.releasePointerCapture = vi.fn() +}) + +afterEach(() => { + cleanup() + vi.clearAllMocks() +}) + +describe('rankSearchOption', () => { + it('ranks a final-segment match above a mid-path match', () => { + // "york" hits the city segment of America/New_York (score 2) but only a + // mid-path segment of America/New_York/Special (score 1). + expect(rankSearchOption('America/New_York', 'york')).toBe(2) + expect(rankSearchOption('America/New_York/Special', 'york')).toBe(1) + expect(rankSearchOption('America/New_York', 'york')).toBeGreaterThan( + rankSearchOption('America/New_York/Special', 'york') + ) + }) + + it('is case-insensitive', () => { + expect(rankSearchOption('Asia/Kolkata', 'KOLKATA')).toBe(2) + expect(rankSearchOption('ASIA/KOLKATA', 'kolkata')).toBe(2) + }) + + it('scores a substring match anywhere as 1', () => { + expect(rankSearchOption('America/New_York', 'amer')).toBe(1) + }) + + it('scores a slashless option by plain substring', () => { + expect(rankSearchOption('UTC', 'ut')).toBe(1) + expect(rankSearchOption('UTC', 'xyz')).toBe(0) + }) + + it('scores a non-match as 0', () => { + expect(rankSearchOption('Europe/Berlin', 'tokyo')).toBe(0) + }) +}) + +describe('SearchableSelect', () => { + const options = ['America/New_York', 'Asia/Kolkata', 'Europe/Berlin', 'UTC'] + + it('opens, filters, and selects an option', () => { + const onChange = vi.fn() + + render() + + fireEvent.click(screen.getByRole('combobox')) + fireEvent.change(screen.getByPlaceholderText('Search…'), { target: { value: 'kolkata' } }) + fireEvent.click(screen.getByText('Asia/Kolkata')) + + expect(onChange).toHaveBeenCalledWith('Asia/Kolkata') + }) + + it('renders the clear item when clearLabel is set and selecting it resets to blank', () => { + const onChange = vi.fn() + + render() + + fireEvent.click(screen.getByRole('combobox')) + fireEvent.click(screen.getByText('System default')) + + expect(onChange).toHaveBeenCalledWith('') + }) + + it('omits the clear item without clearLabel', () => { + render() + + fireEvent.click(screen.getByRole('combobox')) + + expect(screen.queryByText('System default')).toBeNull() + }) + + it('shows the placeholder when the value is blank', () => { + render() + + expect(screen.getByRole('combobox').textContent).toContain('Search…') + }) +}) + +describe('ConfigField searchable routing', () => { + const searchableSchema: ConfigFieldSchema = { + type: 'select', + searchable: true, + clearable: true, + options: ['America/New_York', 'UTC'] + } + + it('routes searchable select schemas to SearchableSelect, not a free-text input', () => { + const { container } = render( + + ) + + // The searchable trigger renders; the generic free-text does not. + expect(container.querySelector('[data-slot="searchable-select-trigger"]')).not.toBeNull() + expect(container.querySelector('input[type="text"]')).toBeNull() + }) + + it('keeps plain string schemas on the free-text input', () => { + const { container } = render( + + ) + + expect(container.querySelector('[data-slot="searchable-select-trigger"]')).toBeNull() + expect(screen.getByDisplayValue('hello')).not.toBeNull() + }) + + it('surfaces the clear item via schema.clearable and resets to blank', () => { + const onChange = vi.fn() + + render() + + fireEvent.click(screen.getByRole('combobox')) + fireEvent.click(screen.getByText('System default')) + + expect(onChange).toHaveBeenCalledWith('') + }) +}) diff --git a/apps/desktop/src/app/settings/searchable-select.tsx b/apps/desktop/src/app/settings/searchable-select.tsx new file mode 100644 index 000000000000..8f679b48ba26 --- /dev/null +++ b/apps/desktop/src/app/settings/searchable-select.tsx @@ -0,0 +1,115 @@ +import { useCallback, useRef, useState } from 'react' + +import { Codicon } from '@/components/ui/codicon' +import { Command, CommandEmpty, CommandGroup, CommandInput, CommandItem, CommandList } from '@/components/ui/command' +import { controlVariants } from '@/components/ui/control' +import { Popover, PopoverContent, PopoverTrigger } from '@/components/ui/popover' +import { cn } from '@/lib/utils' + +/** + * cmdk filter score for one option. Case-insensitive substring match, with + * the final path segment (after the last "/") ranked above matches anywhere + * else so "york" ranks "America/New_York" over "America/New_York/Special". + * Exported for tests. + */ +export function rankSearchOption(option: string, search: string): number { + const lower = search.toLowerCase() + const itemLower = option.toLowerCase() + const slash = itemLower.lastIndexOf('/') + + if (slash !== -1 && itemLower.slice(slash + 1).includes(lower)) { + return 2 + } + + if (itemLower.includes(lower)) { + return 1 + } + + return 0 +} + +/** + * Searchable select for large option lists (e.g. ~590 IANA timezones). + * Built on Popover + cmdk Command — the same stack as Shadcn's Combobox. + * + * The trigger renders like the existing closed ` pattern of EMPTY_SELECT_VALUE + "(none)". */ + clearLabel?: string +}) { + const [open, setOpen] = useState(false) + const triggerRef = useRef(null) + + const handleSelect = useCallback( + (selected: string) => { + onChange(selected) + setOpen(false) + }, + [onChange] + ) + + const displayValue = value !== '' && value !== undefined ? value : placeholder + + return ( + + + + + + + + + {emptyMessage} + + {clearLabel && ( + handleSelect('')} value={clearLabel}> + + {clearLabel} + + )} + {options.map(option => ( + handleSelect(option)} value={option}> + + {option} + + ))} + + + + + + ) +} diff --git a/apps/desktop/src/app/shell/model-menu-panel.tsx b/apps/desktop/src/app/shell/model-menu-panel.tsx index 5fd4c9691c8a..8ca63f5afd6d 100644 --- a/apps/desktop/src/app/shell/model-menu-panel.tsx +++ b/apps/desktop/src/app/shell/model-menu-panel.tsx @@ -259,7 +259,11 @@ export function ModelMenuPanel({ gateway, onSelectModel, profile = 'default', re textValue="" > {group.provider.name} - + {!collapsed && group.families.map(family => { diff --git a/apps/desktop/src/components/assistant-ui/markdown-text.media.test.tsx b/apps/desktop/src/components/assistant-ui/markdown-text.media.test.tsx index e534da37eb93..339bfb18fdf7 100644 --- a/apps/desktop/src/components/assistant-ui/markdown-text.media.test.tsx +++ b/apps/desktop/src/components/assistant-ui/markdown-text.media.test.tsx @@ -1,9 +1,9 @@ -import { cleanup, render, screen } from '@testing-library/react' +import { cleanup, render, screen, waitFor } from '@testing-library/react' import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' import { $connection } from '@/store/session' -import { MarkdownTextContent } from './markdown-text' +import { MarkdownImage, MarkdownTextContent } from './markdown-text' const REMOTE_IMAGE_PATH = '/home/user/project/images/remote-preview.png' const REMOTE_IMAGE_DATA_URL = 'data:image/png;base64,cmVtb3RlLWltYWdl' @@ -50,3 +50,32 @@ describe('MarkdownTextContent remote images', () => { }) }) }) + +// Regression for #40896: generated media often arrives as image markdown +// (`![clip](clip.mp4)`). A raw with a video/audio source paints a +// broken-image icon even though the file is valid, so MarkdownImage must route +// video/audio sources to the proper