diff --git a/.github/workflows/docker.yml b/.github/workflows/docker.yml index 5e5c19bdf3b..6f45592f936 100644 --- a/.github/workflows/docker.yml +++ b/.github/workflows/docker.yml @@ -88,6 +88,11 @@ jobs: # --------------------------------------------------------------------- - name: Install uv (for docker tests) uses: astral-sh/setup-uv@fac544c07dec837d0ccb6301d7b5580bf5edae39 # 8.2.0 + with: + # Pinned: unpinned setup-uv fetches a 'latest' manifest from + # raw.githubusercontent.com every job; transient fetch failures + # fail the job (2026-07-28 incident). Keep in sync with tests.yml. + version: "0.9.28" - name: Set up Python 3.11 (for docker tests) run: uv python install 3.11 diff --git a/.github/workflows/e2e-desktop.yml b/.github/workflows/e2e-desktop.yml index e9131c72522..bdaaa6d6cbb 100644 --- a/.github/workflows/e2e-desktop.yml +++ b/.github/workflows/e2e-desktop.yml @@ -52,6 +52,11 @@ jobs: - name: Install uv uses: astral-sh/setup-uv@fac544c07dec837d0ccb6301d7b5580bf5edae39 # 8.2.0 with: + # Pin the uv version: unpinned, setup-uv resolves "latest" by + # fetching a manifest from raw.githubusercontent.com on EVERY job — + # a transient fetch failure fails the whole job (2026-07-28 slice-5 + # incident). Pinned, the binary downloads directly; no manifest hop. + version: "0.9.28" enable-cache: true cache-dependency-glob: | pyproject.toml diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml index 670b6f2a44a..3ae120f71f3 100644 --- a/.github/workflows/lint.yml +++ b/.github/workflows/lint.yml @@ -40,6 +40,11 @@ jobs: - name: Install uv uses: astral-sh/setup-uv@fac544c07dec837d0ccb6301d7b5580bf5edae39 # 8.2.0 + with: + # Pinned: unpinned setup-uv fetches a 'latest' manifest from + # raw.githubusercontent.com every job; transient fetch failures + # fail the job (2026-07-28 incident). Keep in sync with tests.yml. + version: "0.9.28" - name: Install ruff + ty uses: ./.github/actions/retry @@ -129,6 +134,11 @@ jobs: - name: Install uv uses: astral-sh/setup-uv@fac544c07dec837d0ccb6301d7b5580bf5edae39 # 8.2.0 + with: + # Pinned: unpinned setup-uv fetches a 'latest' manifest from + # raw.githubusercontent.com every job; transient fetch failures + # fail the job (2026-07-28 incident). Keep in sync with tests.yml. + version: "0.9.28" - name: Install ruff uses: ./.github/actions/retry diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index cdae2e037a5..bd398512178 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -74,6 +74,11 @@ jobs: - name: Install uv uses: astral-sh/setup-uv@fac544c07dec837d0ccb6301d7b5580bf5edae39 # 8.2.0 with: + # Pin the uv version: unpinned, setup-uv resolves "latest" by + # fetching a manifest from raw.githubusercontent.com on EVERY job — + # a transient fetch failure fails the whole job (2026-07-28 slice-5 + # incident). Pinned, the binary downloads directly; no manifest hop. + version: "0.9.28" # Persist uv's download/wheel cache (~/.cache/uv) across runs. # Keyed on the dependency manifests, so the cache is reused until # pyproject.toml or uv.lock changes. `uv sync` still runs every @@ -188,6 +193,11 @@ jobs: - name: Install uv uses: astral-sh/setup-uv@fac544c07dec837d0ccb6301d7b5580bf5edae39 # 8.2.0 with: + # Pin the uv version: unpinned, setup-uv resolves "latest" by + # fetching a manifest from raw.githubusercontent.com on EVERY job — + # a transient fetch failure fails the whole job (2026-07-28 slice-5 + # incident). Pinned, the binary downloads directly; no manifest hop. + version: "0.9.28" # Persist uv's download/wheel cache (~/.cache/uv) across runs. # Keyed on the dependency manifests, so the cache is reused until # pyproject.toml or uv.lock changes. `uv sync` still runs every diff --git a/.github/workflows/uv-lockfile-check.yml b/.github/workflows/uv-lockfile-check.yml index aff4f0eb8cb..e0ba3ea7d50 100644 --- a/.github/workflows/uv-lockfile-check.yml +++ b/.github/workflows/uv-lockfile-check.yml @@ -70,6 +70,11 @@ jobs: - name: Install uv uses: astral-sh/setup-uv@fac544c07dec837d0ccb6301d7b5580bf5edae39 # 8.2.0 + with: + # Pinned: unpinned setup-uv fetches a 'latest' manifest from + # raw.githubusercontent.com every job; transient fetch failures + # fail the job (2026-07-28 incident). Keep in sync with tests.yml. + version: "0.9.28" # `uv lock --check` re-resolves the project from pyproject.toml and # compares the result to uv.lock, exiting non-zero if they disagree. diff --git a/Dockerfile b/Dockerfile index 6b76075e980..6d6e9ac2d91 100644 --- a/Dockerfile +++ b/Dockerfile @@ -200,6 +200,22 @@ RUN npm install --prefer-offline --no-audit --fetch-retries=5 && \ done && \ npm cache clean --force +# ---------- Photon iMessage sidecar deps (baked, NS-606) ---------- +# The photon plugin's Node sidecar needs its own node_modules +# (spectrum-ts). The install tree is immutable at runtime, so a lazy +# `npm ci` on first connect would hit EROFS — bake the deps here instead +# (deterministic installs, NS-559). The patch script is copied alongside +# the manifests because package.json's postinstall runs it, which also +# means the spectrum-ts patch is applied at build time. Layer-cached: +# only re-runs when the sidecar manifests/patch change. +COPY plugins/platforms/photon/sidecar/package.json \ + plugins/platforms/photon/sidecar/package-lock.json \ + plugins/platforms/photon/sidecar/patch-spectrum-mixed-attachments.mjs \ + plugins/platforms/photon/sidecar/ +RUN cd plugins/platforms/photon/sidecar && \ + npm ci --no-audit --fetch-retries=5 && \ + npm cache clean --force + # ---------- Layer-cached Python dependency install ---------- # Copy only pyproject.toml + uv.lock so the Python dep resolve + wheel # download + native-extension compile layer is cached unless those inputs diff --git a/agent/agent_init.py b/agent/agent_init.py index a99a0de8991..ea473632c6a 100644 --- a/agent/agent_init.py +++ b/agent/agent_init.py @@ -118,7 +118,7 @@ def _provider_default_routes(provider: str) -> set[str]: from hermes_cli.providers import HERMES_OVERLAYS, get_provider overlay = HERMES_OVERLAYS.get(provider) - provider_def = get_provider(provider) + provider_def = get_provider(provider, allow_network=False) for value in ( getattr(overlay, "base_url_override", ""), getattr(provider_def, "base_url", ""), @@ -887,8 +887,10 @@ def init_agent( # report cumulative micros spent. Surfaced behind HERMES_DEV_CREDITS. agent._credits_state = None agent._credits_session_start_micros = None - # Threshold-notice latch (L4): active sticky-notice keys + the warn90 crossing gate. - agent._credits_latch = {"active": set(), "seen_below_90": False, "usage_band": None} + # Threshold-notice latch (L4): active sticky-notice keys + the crossing gates. + from agent.credits_tracker import new_credits_latch + + agent._credits_latch = new_credits_latch() # OpenRouter response cache hit counter — incremented when # X-OpenRouter-Cache-Status: HIT is seen in streaming response headers. @@ -2276,7 +2278,18 @@ def init_agent( # AFTER the custom_providers branch so per-model overrides aren't lost. agent._config_context_length = _config_context_length - agent._ensure_lmstudio_runtime_loaded(_config_context_length) + _lmstudio_runtime_context_length = agent._ensure_lmstudio_runtime_loaded( + _config_context_length + ) + if agent._lmstudio_load_was_unverified(_lmstudio_runtime_context_length): + _ra().logger.warning( + "LM Studio model activation was rejected or completed without a " + "verifiable active context length; falling back to configured context" + ) + _effective_context_length = agent._effective_lmstudio_context_length( + _config_context_length, + _lmstudio_runtime_context_length, + ) @@ -2353,7 +2366,7 @@ def init_agent( agent.model, base_url=agent.base_url, api_key=getattr(agent, "api_key", ""), - config_context_length=_config_context_length, + config_context_length=_effective_context_length, provider=agent.provider, custom_providers=_custom_providers, ) @@ -2388,7 +2401,7 @@ def init_agent( quiet_mode=agent.quiet_mode, base_url=agent.base_url, api_key=getattr(agent, "api_key", ""), - config_context_length=_config_context_length, + config_context_length=_effective_context_length, provider=agent.provider, api_mode=agent.api_mode, abort_on_summary_failure=compression_abort_on_summary_failure, @@ -2417,7 +2430,13 @@ def init_agent( # Reject models whose context window is below the minimum required # for reliable tool-calling workflows (64K tokens). _ctx = getattr(agent.context_compressor, "context_length", 0) - if _ctx and _ctx < MINIMUM_CONTEXT_LENGTH: + _allow_lmstudio_explicit_below_floor = ( + str(getattr(agent, "provider", "") or "").strip().lower() == "lmstudio" + and isinstance(agent._config_context_length, int) + and not isinstance(agent._config_context_length, bool) + and agent._config_context_length > 0 + ) + if _ctx and _ctx < MINIMUM_CONTEXT_LENGTH and not _allow_lmstudio_explicit_below_floor: raise ValueError( f"Model {agent.model} has a context window of {_ctx:,} tokens, " f"which is below the minimum {MINIMUM_CONTEXT_LENGTH:,} required " diff --git a/agent/agent_runtime_helpers.py b/agent/agent_runtime_helpers.py index d6522d07f48..6192556d2cb 100644 --- a/agent/agent_runtime_helpers.py +++ b/agent/agent_runtime_helpers.py @@ -2137,6 +2137,16 @@ def switch_model(agent, new_model, new_provider, api_key='', base_url='', api_mo agent, "_credential_pool_entry_id", _MISSING ) + def _restore_snapshot() -> None: + for _name, _value in _snapshot.items(): + if _value is _MISSING: + # Attribute did not exist before the swap — don't fabricate it. + continue + try: + setattr(agent, _name, _value) + except Exception: # noqa: BLE001 + pass + try: # Clear the per-config context_length override so the new model's # actual context window is resolved via get_model_context_length() @@ -2305,16 +2315,42 @@ def switch_model(agent, new_model, new_provider, api_key='', base_url='', api_mo # caller's exception handler can surface a meaningful warning. The # exception is re-raised; cli.py / gateway/run.py / tui_gateway catch # it and print "Agent swap failed; change applied to next session". - for _name, _value in _snapshot.items(): - if _value is _MISSING: - # Attribute did not exist before the swap — don't fabricate it. - continue - try: - setattr(agent, _name, _value) - except Exception: # noqa: BLE001 - pass + _restore_snapshot() raise + # ── LM Studio: preload before probing context length ── + _sm_custom_providers = None + try: + from hermes_cli.config import ( + get_compatible_custom_providers, + get_custom_provider_context_length, + load_config, + ) + + _sm_cfg = load_config() + _sm_custom_providers = get_compatible_custom_providers(_sm_cfg) + _destination_context_intent = get_custom_provider_context_length( + model=agent.model, + base_url=agent.base_url, + custom_providers=_sm_custom_providers, + ) + except Exception: + _destination_context_intent = None + agent._config_context_length = _destination_context_intent + _runtime_context_length = agent._ensure_lmstudio_runtime_loaded( + _destination_context_intent + ) + if agent._lmstudio_load_was_unverified(_runtime_context_length): + logger.warning( + "LM Studio model activation was rejected or completed without a " + "verifiable active context length during model switch; continuing " + "with configured context" + ) + _effective_context_length = agent._effective_lmstudio_context_length( + _destination_context_intent, + _runtime_context_length, + ) + # ── Re-evaluate prompt caching ── agent._use_prompt_caching, agent._use_native_cache_layout = ( agent._anthropic_prompt_cache_policy( @@ -2325,22 +2361,15 @@ def switch_model(agent, new_model, new_provider, api_key='', base_url='', api_mo ) ) - # ── LM Studio: preload before probing context length ── - agent._ensure_lmstudio_runtime_loaded() - # ── Update context compressor ── if hasattr(agent, "context_compressor") and agent.context_compressor: from agent.model_metadata import get_model_context_length - # Re-read custom_providers from live config so per-model - # context_length overrides are honored when switching to a - # custom provider mid-session (closes #15779). - _sm_custom_providers = None - try: - from hermes_cli.config import load_config, get_compatible_custom_providers - _sm_cfg = load_config() - _sm_custom_providers = get_compatible_custom_providers(_sm_cfg) - except Exception: - _sm_custom_providers = None + if _sm_custom_providers is None: + try: + from hermes_cli.config import get_compatible_custom_providers, load_config + _sm_custom_providers = get_compatible_custom_providers(load_config()) + except Exception: + _sm_custom_providers = None # ``agent.api_key`` may be a callable (Azure Foundry Entra ID # token provider). ``get_model_context_length`` expects a # string for its live-probe paths; for Foundry the context @@ -2352,7 +2381,7 @@ def switch_model(agent, new_model, new_provider, api_key='', base_url='', api_mo base_url=agent.base_url, api_key=_ctx_api_key, provider=agent.provider, - config_context_length=getattr(agent, "_config_context_length", None), + config_context_length=_effective_context_length, custom_providers=_sm_custom_providers, ) agent.context_compressor.update_model( @@ -2475,7 +2504,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 +2684,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 +2696,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 fa5bc6b80d7..42647956042 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 @@ -534,7 +536,7 @@ def _get_aux_model_for_provider(provider_id: str) -> str: # plus providers we intentionally keep pinned here (e.g. Anthropic predates # profiles). New providers should set default_aux_model on their profile instead. _API_KEY_PROVIDER_AUX_MODELS_FALLBACK: Dict[str, str] = { - "gemini": "gemini-3-flash-preview", + "gemini": "gemini-3.6-flash", "zai": "glm-4.5-flash", "kimi-coding": "kimi-k2-turbo-preview", "stepfun": "step-3.5-flash", @@ -543,7 +545,7 @@ _API_KEY_PROVIDER_AUX_MODELS_FALLBACK: Dict[str, str] = { "anthropic": "claude-haiku-4-5-20251001", "opencode-zen": "gemini-3-flash", "opencode-go": "glm-5", - "kilocode": "google/gemini-3-flash-preview", + "kilocode": "google/gemini-3.6-flash", "ollama-cloud": "nemotron-3-nano:30b", "tencent-tokenhub": "hy3-preview", # NB: no "deepinfra" entry — its aux model lives on the ProviderProfile @@ -758,8 +760,8 @@ NOUS_EXTRA_BODY = _nous_extra_body() auxiliary_is_nous: bool = False # Default auxiliary models per provider -_OPENROUTER_MODEL = "google/gemini-3-flash-preview" -_NOUS_MODEL = "google/gemini-3-flash-preview" +_OPENROUTER_MODEL = "google/gemini-3.6-flash" +_NOUS_MODEL = "google/gemini-3.6-flash" _NOUS_DEFAULT_BASE_URL = "https://inference-api.nousresearch.com/v1" _ANTHROPIC_DEFAULT_BASE_URL = "https://api.anthropic.com" _AUTH_JSON_PATH = get_hermes_home() / "auth.json" @@ -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 5fecd0e26cf..fac42cc3cf7 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. @@ -519,10 +526,16 @@ def direct_api_call(agent, api_kwargs: dict): def _abort_active_request(reason: str) -> None: """Abort the inline request from a watchdog/interrupt thread.""" + # Abort while still holding the holder lock: the instant it is + # released, the inline finally may pop + cache the client for reuse + # and the NEXT call check it out — a late abort would then poison + # the slot and shut down an innocent in-flight request's sockets + # (same atomicity contract as _close_request_client_once in the + # interruptible variants; the abort itself never blocks). with request_client_lock: request_client = request_client_holder["client"] - if request_client is not None: - agent._abort_request_openai_client(request_client, reason=reason) + if request_client is not None: + agent._abort_request_openai_client(request_client, reason=reason) def _make_client(reason: str, kind: str = "openai"): # direct_api_call only runs for OpenAI-wire chat_completions cron @@ -535,6 +548,10 @@ def direct_api_call(agent, api_kwargs: dict): agent._active_request_abort = _abort_active_request return client + # Only a clean return may report the reuse reason (request_complete): + # after an error or interrupt the wire client is really closed so the + # retry builds a fresh pool (see _REQUEST_CLIENT_REUSE_REASONS). + succeeded = False try: response = _dispatch_nonstreaming_api_request( agent, api_kwargs, make_client=_make_client @@ -547,6 +564,7 @@ def direct_api_call(agent, api_kwargs: dict): if getattr(agent, "_interrupt_requested", False): raise InterruptedError("Agent interrupted during API call") _reset_stale_streak(agent) + succeeded = True return response finally: if getattr(agent, "_active_request_abort", None) is _abort_active_request: @@ -555,7 +573,10 @@ def direct_api_call(agent, api_kwargs: dict): request_client = request_client_holder["client"] request_client_holder["client"] = None if request_client is not None: - agent._close_request_openai_client(request_client, reason="request_complete") + agent._close_request_openai_client( + request_client, + reason="request_complete" if succeeded else "request_error_cleanup", + ) def interruptible_api_call(agent, api_kwargs: dict): @@ -633,20 +654,28 @@ def interruptible_api_call(agent, api_kwargs: dict): and owner_tid is not None and owner_tid != threading.get_ident() ) - if not stranger_thread: - # Owning thread (or no recorded owner) → pop and fully close. - request_client_holder["client"] = None - request_client_holder["owner_tid"] = None + if stranger_thread: + # Abort while still holding the holder lock: the instant it + # is released, the worker's finally may pop + cache the client + # for reuse and the NEXT call check it out — an abort landing + # after that would poison the slot and shut down an innocent + # in-flight request's sockets. The abort itself never blocks + # (socket shutdown + slot poison), so holding the lock across + # it only delays the racing pop, never the data path. + if request_client_kind.get("value", "openai") == "anthropic_messages": + agent._abort_request_anthropic_client( + request_client, reason=reason + ) + else: + agent._abort_request_openai_client(request_client, reason=reason) + return + # Owning thread (or no recorded owner) → pop and fully close. + request_client_holder["client"] = None + request_client_holder["owner_tid"] = None if request_client is None: return - kind = request_client_kind.get("value", "openai") - if kind == "anthropic_messages": - if stranger_thread: - agent._abort_request_anthropic_client(request_client, reason=reason) - else: - agent._close_request_anthropic_client(request_client, reason=reason) - elif stranger_thread: - agent._abort_request_openai_client(request_client, reason=reason) + if request_client_kind.get("value", "openai") == "anthropic_messages": + agent._close_request_anthropic_client(request_client, reason=reason) else: agent._close_request_openai_client(request_client, reason=reason) @@ -683,7 +712,15 @@ def interruptible_api_call(agent, api_kwargs: dict): return result["error"] = e finally: - _close_request_client_once("request_complete") + # Reuse reason only on a clean response; any other outcome — + # error, or the cancel-swallow return above (which leaves both + # result slots None) — really closes so the next attempt builds + # a fresh pool (see _REQUEST_CLIENT_REUSE_REASONS). + _close_request_client_once( + "request_complete" + if result["response"] is not None + else "request_error_cleanup" + ) # ── Stale-call timeout (mirrors streaming stale detector) ──────── # Non-streaming calls return nothing until the full response is @@ -810,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(): @@ -1989,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, " @@ -2174,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() @@ -2192,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." @@ -2206,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: @@ -2229,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() @@ -2237,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." @@ -2246,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 @@ -2404,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, @@ -2413,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() @@ -2465,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) @@ -2642,25 +2787,33 @@ def interruptible_streaming_api_call(agent, api_kwargs: dict, *, on_first_delta= and owner_tid is not None and owner_tid != threading.get_ident() ) - if not stranger_thread: - request_client_holder["client"] = None - request_client_holder["owner_tid"] = None + if stranger_thread: + # Abort under the holder lock — see the non-streaming variant + # for why the holder read and the abort must be atomic (a late + # abort would otherwise hit the NEXT request's checkout). + if request_client_kind.get("value", "openai") == "anthropic_messages": + agent._abort_request_anthropic_client( + request_client, reason=reason + ) + else: + agent._abort_request_openai_client(request_client, reason=reason) + return + request_client_holder["client"] = None + request_client_holder["owner_tid"] = None if request_client is None: return + # Stranger threads returned under the lock above, so only the owner + # (or an any-thread-safe stream handle) reaches the close dispatch. if request_kind == "stream": _close_request_stream_handle(request_client, reason) elif request_kind == "anthropic_messages": - if stranger_thread: - agent._abort_request_anthropic_client(request_client, reason=reason) - else: - agent._close_request_anthropic_client(request_client, reason=reason) - elif stranger_thread: - agent._abort_request_openai_client(request_client, reason=reason) + agent._close_request_anthropic_client(request_client, reason=reason) else: agent._close_request_openai_client(request_client, reason=reason) 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. @@ -2680,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: @@ -2794,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() @@ -2909,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") @@ -2944,6 +3120,26 @@ def interruptible_streaming_api_call(agent, api_kwargs: dict, *, on_first_delta= pass if agent._interrupt_requested: + # Abandoning a half-read SSE response leaves its connection + # permanently checked out of the httpx pool — and the partial + # response built below makes the worker's finally report a + # reuse-reason close, which would cache the client together + # with the leaked connection (each interrupt leaking one more + # until the pool exhausts). Close the stream here, on the + # owning thread, so the connection is released first. + try: + stream.close() + except Exception: + # Connection may still be checked out — poison the slot so + # 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). + 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): @@ -3075,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 @@ -3241,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": @@ -3315,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: @@ -3331,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 @@ -3407,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 @@ -3446,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 @@ -3695,7 +3950,15 @@ def interruptible_streaming_api_call(agent, api_kwargs: dict, *, on_first_delta= result["error"] = e return finally: - _close_request_client_once("stream_request_complete") + _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). + _close_request_client_once( + "stream_request_complete" + if result["response"] is not None + else "stream_error_cleanup" + ) # Provider-configured stale timeout takes priority over env default. _cfg_stale = get_provider_stale_timeout(agent.provider, agent.model) @@ -3757,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 ee75f4190e6..edff776536e 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 da3bc4f9569..0e0b87b2196 100644 --- a/agent/codex_runtime.py +++ b/agent/codex_runtime.py @@ -74,7 +74,10 @@ def _record_codex_app_server_usage(agent, turn) -> dict[str, Any]: try: if not agent._session_db_created: agent._ensure_db_session() - agent._session_db.update_token_counts( + # Enqueued for the SessionDB background writer — keeps the + # per-call accounting write off the turn thread (see + # conversation_loop's queue_token_counts call). + agent._session_db.queue_token_counts( agent.session_id, model=agent.model, billing_provider=agent.provider, @@ -154,7 +157,8 @@ def _record_codex_app_server_usage(agent, turn) -> dict[str, Any]: try: if not agent._session_db_created: agent._ensure_db_session() - agent._session_db.update_token_counts( + # Enqueued for the SessionDB background writer (see above). + agent._session_db.queue_token_counts( agent.session_id, input_tokens=canonical_usage.input_tokens, output_tokens=canonical_usage.output_tokens, @@ -1236,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. @@ -1260,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, @@ -1320,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( @@ -1330,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( @@ -1347,7 +1403,20 @@ def run_codex_stream(agent, api_kwargs: dict, client: Any = None, on_first_delta try: close_fn() except Exception: - pass + # A failed close can leave this response's connection + # checked out of the httpx pool while the caller's finally + # reports a reuse-reason close (e.g. interrupt_check broke + # the event loop with collected output) — caching the + # client with the leaked connection. Poison the slot so + # that close really closes the pool (owner-thread abort; + # mirrors the chat-streaming interrupt-break handling). + # ``client is None`` means the shared primary client, + # which is never reuse-cached and must not have its + # sockets force-shut here. + if client is not None: + agent._abort_request_openai_client( + active_client, reason="codex_stream_close_failed" + ) def run_codex_create_stream_fallback(agent, api_kwargs: dict, client: Any = None): diff --git a/agent/context_compressor.py b/agent/context_compressor.py index 73fa1e36f21..4837925ce62 100644 --- a/agent/context_compressor.py +++ b/agent/context_compressor.py @@ -1354,6 +1354,125 @@ class ContextCompressor(ContextEngine): previous = telemetry.get("aux_call_duration_ms") or 0 telemetry["aux_call_duration_ms"] = previous + max(0, int(duration_ms)) + def _emit_init_summary_once(self) -> None: + """Emit the informative startup line once, on first resolution. + + Deferred out of ``__init__`` (#32221): the line reports resolved token + budgets, so emitting it there would force the synchronous + ``get_model_context_length()`` probe during construction. Reads via + the properties below are safe here because + ``_resolved_context_length`` is already set. + """ + if not getattr(self, "_log_init_summary", False): + return + self._log_init_summary = False + logger.info( + "Context compressor initialized: model=%s context_length=%d " + "threshold=%d (%.0f%%) target_ratio=%.0f%% tail_budget=%d " + "provider=%s base_url=%s", + self.model, self._resolved_context_length, self.threshold_tokens, + self.threshold_percent * 100, self.summary_target_ratio * 100, + self.tail_token_budget, + self.provider or "none", self.base_url or "none", + ) + + def _resolve_context_length(self) -> int: + """Resolve and cache the model's context length on first access.""" + if self._resolved_context_length is None: + self._resolved_context_length = get_model_context_length( + self.model, + base_url=self.base_url, + api_key=self.api_key, + config_context_length=self._config_context_length, + provider=self.provider, + ) + # Small-context threshold floor: models under 512K trigger at + # >=75% so compaction doesn't fire with half the window still + # free. Raise-only; must run AFTER context_length is resolved + # and BEFORE threshold_tokens is derived (deferred here from + # __init__ along with the resolution itself, #32221). + # _base_threshold_percent already has the per-model override + # applied, so the floor stacks on top of it. + self.threshold_percent = self._effective_threshold_percent( + self._resolved_context_length, self._base_threshold_percent, + ) + self._emit_init_summary_once() + return self._resolved_context_length + + @property + def context_length(self) -> int: + return self._resolve_context_length() + + @context_length.setter + def context_length(self, value: int) -> None: + # No-op guard: repeated assignment of the SAME window (e.g. the codex + # app-server usage callback re-reports the window on every response) + # must not invalidate the derived budgets — that would wipe runtime + # corrections applied directly to threshold_tokens/tail_token_budget + # (see conversation_compression's aux-context threshold sync), which + # persisted on main's eager-init behavior. + if value == getattr(self, "_resolved_context_length", None): + return + self._resolved_context_length = value + # Re-apply the small-context floor (raise-only) for the genuinely new + # window so the invalidated budgets below recompute coherently — + # percent and tokens must derive from the same window. Skipped on + # bare test instances built via object.__new__ that never ran + # __init__ (no _base_threshold_percent). + _base = getattr(self, "_base_threshold_percent", None) + if _base is not None: + self.threshold_percent = self._effective_threshold_percent( + value, _base, + ) + self._threshold_tokens = None + self._tail_token_budget = None + self._max_summary_tokens = None + self._emit_init_summary_once() + + @property + def threshold_tokens(self) -> int: + if self._threshold_tokens is None: + # Resolve the window FIRST (may apply the small-context floor to + # threshold_percent as a side effect) so the percent read below + # is the floored value regardless of argument evaluation order. + _ctx = self.context_length + # Floor: never compress below MINIMUM_CONTEXT_LENGTH tokens even + # if the percentage would suggest a lower value (#14690 handles + # the degenerate small-window case inside the helper). + self._threshold_tokens = self._compute_threshold_tokens( + _ctx, self.threshold_percent, self.max_tokens, + ) + # Apply absolute token cap (compression.threshold_tokens) — + # takes the lower of the ratio-based threshold and the cap. + self._apply_threshold_tokens_cap() + return self._threshold_tokens + + @threshold_tokens.setter + def threshold_tokens(self, value: int) -> None: + self._threshold_tokens = value + + @property + def tail_token_budget(self) -> int: + if self._tail_token_budget is None: + self._tail_token_budget = int(self.threshold_tokens * self.summary_target_ratio) + return self._tail_token_budget + + @tail_token_budget.setter + def tail_token_budget(self, value: int) -> None: + self._tail_token_budget = value + + @property + def max_summary_tokens(self) -> int: + if self._max_summary_tokens is None: + self._max_summary_tokens = min( + int(self.context_length * 0.05), _SUMMARY_TOKENS_CEILING, + ) + return self._max_summary_tokens + + @max_summary_tokens.setter + def max_summary_tokens(self, value: int) -> None: + self._max_summary_tokens = value + def on_session_end(self, session_id: str, messages: List[Dict[str, Any]]) -> None: """Clear all per-session compaction state at a real session boundary. @@ -1988,56 +2107,30 @@ class ContextCompressor(ContextEngine): # deterministic "summary unavailable" handoff and drop the middle window. self.abort_on_summary_failure = abort_on_summary_failure - self.context_length = get_model_context_length( - model, base_url=base_url, api_key=api_key, - config_context_length=config_context_length, - provider=provider, - ) - # Small-context threshold floor: models under 512K trigger at >=75% - # so compaction doesn't fire with half the window still free (the - # incompressible floor makes 50%-triggered compaction thrash on - # 128K-262K models). Raise-only; must run AFTER context_length is - # resolved and BEFORE threshold_tokens is derived. The pre-floor - # value is kept so update_model() can re-derive for a new window - # (switching small -> large must drop back to the configured value). - # Note: _base_threshold_percent already has the per-model override - # applied, so the floor stacks on top of any model-specific threshold. + # Defer context-length resolution to first access (#32221): + # get_model_context_length() can issue a synchronous /models HTTP + # probe, which must not block AIAgent construction. The small-context + # threshold floor and the absolute threshold cap both need the + # resolved window, so they are applied on first resolution (see + # _resolve_context_length / the threshold_tokens property) instead + # of here. update_model() re-derives the floor for a new window from + # _config_threshold_percent (the raw config value snapshotted above), + # so switching small -> large correctly drops back to the configured + # value. + self._config_context_length = config_context_length self._configured_threshold_percent = self.threshold_percent - self.threshold_percent = self._effective_threshold_percent( - self.context_length, self._base_threshold_percent, - ) - threshold_percent = self.threshold_percent - # Floor: never compress below MINIMUM_CONTEXT_LENGTH tokens even if - # the percentage would suggest a lower value. This prevents premature - # compression on large-context models at 50% while keeping the % sane - # for models right at the minimum. _compute_threshold_tokens also - # guards the degenerate case where the floor would equal/exceed the - # window (small models), so auto-compression can still fire (#14690). - self.threshold_tokens = self._compute_threshold_tokens( - self.context_length, threshold_percent, self.max_tokens, - ) - # Apply absolute token cap (compression.threshold_tokens) — takes - # the lower of the ratio-based threshold and the cap. - self._apply_threshold_tokens_cap() + self._resolved_context_length: int | None = None + self._threshold_tokens: int | None = None + self._tail_token_budget: int | None = None + self._max_summary_tokens: int | None = None self.compression_count = 0 - # Derive token budgets: ratio is relative to the threshold, not total context - target_tokens = int(self.threshold_tokens * self.summary_target_ratio) - self.tail_token_budget = target_tokens - self.max_summary_tokens = min( - int(self.context_length * 0.05), _SUMMARY_TOKENS_CEILING, - ) - - if not quiet_mode: - logger.info( - "Context compressor initialized: model=%s context_length=%d " - "threshold=%d (%.0f%%) target_ratio=%.0f%% tail_budget=%d " - "provider=%s base_url=%s", - model, self.context_length, self.threshold_tokens, - threshold_percent * 100, self.summary_target_ratio * 100, - self.tail_token_budget, - provider or "none", base_url or "none", - ) + # The "initialized" log reports resolved token budgets, which would + # force the deferred get_model_context_length() probe to run inside + # __init__ and re-introduce the exact synchronous blocking this change + # removes (#32221). Emit it on first context-length resolution instead + # so construction stays non-blocking on every path (not just quiet). + self._log_init_summary = not quiet_mode self._context_probed = False # True after a step-down from context error self.last_prompt_tokens = 0 diff --git a/agent/conversation_compression.py b/agent/conversation_compression.py index 506181f2d39..8308b618da7 100644 --- a/agent/conversation_compression.py +++ b/agent/conversation_compression.py @@ -2721,16 +2721,28 @@ def try_shrink_image_parts_in_messages( media_type = "image/jpeg" return f"data:{media_type};base64,{data}" - def _write_data_url_to_source(source: dict, data_url: str) -> None: + def _write_data_url_to_source(source: dict, data_url: str) -> dict: + """Return a NEW source dict carrying the re-encoded payload. + + Copy-on-write: content parts on the per-call ``api_messages`` list may + be shared references into the persistent conversation history (the + per-message copy is shallow, and cache decoration only deep-copies the + marked messages). Mutating the existing dict would rewrite the stored + transcript with the degraded image — so the caller replaces the part, + never edits it in place. + """ header, _, data = data_url.partition(",") media_type = "image/jpeg" if header.startswith("data:"): candidate = header[len("data:"):].split(";", 1)[0].strip() if candidate.startswith("image/"): media_type = candidate - source["type"] = "base64" - source["media_type"] = media_type - source["data"] = data + return { + **source, + "type": "base64", + "media_type": media_type, + "data": data, + } for msg in api_messages: if not isinstance(msg, dict): @@ -2738,7 +2750,13 @@ def try_shrink_image_parts_in_messages( content = msg.get("content") if not isinstance(content, list): continue - for part in content: + # Copy-on-write per message: never mutate part/source dicts in place — + # they can alias the stored conversation history (see + # _write_data_url_to_source). Build a replacement content list on the + # first shrunken part and reassign msg["content"] (a top-level write on + # the per-call message copy, which never reaches history). + new_content: list | None = None + for part_idx, part in enumerate(content): if not isinstance(part, dict): continue ptype = part.get("type") @@ -2747,7 +2765,12 @@ def try_shrink_image_parts_in_messages( url = _source_to_data_url(source) resized, unshrinkable = _shrink_data_url(url or "") if resized and isinstance(source, dict): - _write_data_url_to_source(source, resized) + if new_content is None: + new_content = list(content) + new_content[part_idx] = { + **part, + "source": _write_data_url_to_source(source, resized), + } changed_count += 1 elif unshrinkable: unshrinkable_oversized += 1 @@ -2761,17 +2784,26 @@ def try_shrink_image_parts_in_messages( url = image_value.get("url", "") resized, unshrinkable = _shrink_data_url(url) if resized: - image_value["url"] = resized + if new_content is None: + new_content = list(content) + new_content[part_idx] = { + **part, + "image_url": {**image_value, "url": resized}, + } changed_count += 1 elif unshrinkable: unshrinkable_oversized += 1 elif isinstance(image_value, str): resized, unshrinkable = _shrink_data_url(image_value) if resized: - part["image_url"] = resized + if new_content is None: + new_content = list(content) + new_content[part_idx] = {**part, "image_url": resized} changed_count += 1 elif unshrinkable: unshrinkable_oversized += 1 + if new_content is not None: + msg["content"] = new_content if changed_count: logger.info( diff --git a/agent/conversation_loop.py b/agent/conversation_loop.py index 7df0e44db8e..06504616bac 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 @@ -3217,7 +3239,12 @@ def run_conversation( _cost_delta = (_cost_delta or 0.0) + float(_moa_ref_cost) except (TypeError, ValueError): # pragma: no cover pass - agent._session_db.update_token_counts( + # Enqueued, not written: the background writer + # applies the delta off the turn thread (a cold + # state.db UPDATE here stalled the tool loop for + # up to hundreds of ms per API call). Drained at + # turn finalize via _persist_session. + agent._session_db.queue_token_counts( agent.session_id, input_tokens=canonical_usage.input_tokens, output_tokens=canonical_usage.output_tokens, @@ -3283,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 @@ -5427,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, ) @@ -6760,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 929bc34d326..b47c3f274ed 100644 --- a/agent/credits_tracker.py +++ b/agent/credits_tracker.py @@ -170,6 +170,27 @@ CREDITS_USAGE_BANDS: tuple[tuple[float, str, int], ...] = ( ) CREDITS_USAGE_KEY = "credits.usage" # single key for the escalating usage notice +# Minimum subscription balance that counts as "grant not yet spent" for the +# grant_spent crossing gate (see evaluate_credits_notices). 1¢: portal-seeded +# states derive micros from float dollars and can carry sub-cent residue where +# the inference headers report exactly 0 — without this floor such a seed +# opens the gate and the first header re-creates the at-open nag. +GRANT_UNSPENT_MIN_MICROS = 10_000 + + +def new_credits_latch() -> dict: + """Fresh notice latch in the shape :func:`evaluate_credits_notices` expects. + + The policy owns this schema — every producer (agent build, lazy re-init, + tests) must build the latch through here so a new gate key lands everywhere + at once instead of drifting across hand-rolled literals.""" + return { + "active": set(), + "seen_below_90": False, + "usage_band": None, + "seen_grant_unspent": False, + } + # ── AgentNotice (out-of-band notice payload; driver-agnostic) ──────────────── @@ -250,7 +271,8 @@ def evaluate_credits_notices( ) -> tuple[list[AgentNotice], list[str]]: """Reconcile credits notices against the latch. Mutates ``latch`` IN PLACE. - latch = {"active": set[str], "seen_below_90": bool, "usage_band": Optional[int]}. + latch = {"active": set[str], "seen_below_90": bool, "usage_band": Optional[int], + "seen_grant_unspent": bool}. ``model_is_free``: True when the session's active model is a Nous free-tier model (see :func:`is_free_tier_model`). Suppresses the ``credits.depleted`` @@ -277,6 +299,18 @@ def evaluate_credits_notices( if uf is not None and uf < _lowest_band: latch["seen_below_90"] = True # gate opened: usage-band notices may now fire + # Grant-spent crossing gate: grant_spent may fire only after this session + # has OBSERVED the grant meaningfully unspent (≥1¢ left — see + # GRANT_UNSPENT_MIN_MICROS). Opening at grant-spent is a steady STATE, not + # an event — /usage carries it; only a live in-session crossing announces. + # Unlike seen_below_90, seeds must NOT prime this gate. + if ( + uf is not None + and uf < 1.0 + and state.subscription_micros >= GRANT_UNSPENT_MIN_MICROS + ): + latch["seen_grant_unspent"] = True + active = latch["active"] # ── Conditions ─────────────────────────────────────────────────────────── @@ -341,7 +375,17 @@ def evaluate_credits_notices( latch["usage_band"] = target_band # ── grant_spent ────────────────────────────────────────────────────────── - if grant_cond and "credits.grant_spent" not in active: + # The crossing gate guards only the SHOW and is CONSUMED by it — one + # announcement per crossing. A header flicker (uf → None → back to 1.0) + # clears the sticky line via grant_cond but cannot re-announce; only a + # renewal that re-opens the gate (a fresh ≥1¢ observation) arms the next + # announcement. .get(): default closed for any hand-built latch missing + # the key, so a first observation can never fire this notice. + if ( + grant_cond + and "credits.grant_spent" not in active + and latch.get("seen_grant_unspent", False) + ): to_show.append( AgentNotice( text=f"• Grant spent · ${state.purchased_usd} top-up left", @@ -352,6 +396,7 @@ def evaluate_credits_notices( ) ) active.add("credits.grant_spent") + latch["seen_grant_unspent"] = False elif "credits.grant_spent" in active and not grant_cond: to_clear.append("credits.grant_spent") active.discard("credits.grant_spent") @@ -627,7 +672,8 @@ _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( # uf == 1.0 + purchased>0 → SILENT at open (crossing-gated); + # flip healthy → grant_exhausted via the fixture-file path to see credits.grant_spent 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", @@ -741,6 +787,9 @@ def _hydrate_seed_state(agent, state) -> None: agent._credits_session_start_micros = state.remaining_micros _latch = getattr(agent, "_credits_latch", None) if isinstance(_latch, dict) and state.used_fraction is not None: + # Prime ONLY seen_below_90 (open-high band warnings are wanted at open). + # Never prime seen_grant_unspent here: a seed observing grant-spent is a + # steady state, and priming it would revive the every-session nag. _latch["seen_below_90"] = True emit = getattr(agent, "_emit_credits_notices", None) if callable(emit): diff --git a/agent/gemini_native_adapter.py b/agent/gemini_native_adapter.py index bb53e32b2ce..cde63f15fc1 100644 --- a/agent/gemini_native_adapter.py +++ b/agent/gemini_native_adapter.py @@ -73,7 +73,7 @@ def probe_gemini_tier( api_key: str, base_url: str = DEFAULT_GEMINI_BASE_URL, *, - model: str = "gemini-2.5-flash", + model: str = "gemini-3.6-flash", timeout: float = 10.0, ) -> str: """Probe a Google AI Studio API key and return its tier. @@ -154,8 +154,8 @@ def is_free_tier_quota_error(error_message: str) -> bool: _FREE_TIER_GUIDANCE = ( - "\n\nYour Google API key is on the free tier (<= 250 requests/day for " - "gemini-2.5-flash). Hermes typically makes 3-10 API calls per user turn, " + "\n\nYour Google API key is on the free tier (a few hundred requests/day " + "for Gemini Flash models). Hermes typically makes 3-10 API calls per user turn, " "so the free tier is exhausted in a handful of messages and cannot sustain " "an agent session. Enable billing on your Google Cloud project and " "regenerate the key in a billing-enabled project: " @@ -321,9 +321,13 @@ def _translate_tool_result_to_gemini( ) -> Dict[str, Any]: tool_name_by_call_id = tool_name_by_call_id or {} tool_call_id = str(message.get("tool_call_id") or "") + # A tool result can carry the unwrapped internal tool name (for example, + # an MCP tool invoked through the `tool_call` bridge). Gemini requires + # functionResponse.name to echo the matching functionCall.name, so the + # call-id mapping must take precedence over the internal result name. name = str( - message.get("name") - or tool_name_by_call_id.get(tool_call_id) + tool_name_by_call_id.get(tool_call_id) + or message.get("name") or tool_call_id or "tool" ) @@ -976,7 +980,7 @@ class GeminiNativeClient: def _create_chat_completion( self, *, - model: str = "gemini-2.5-flash", + model: str = "gemini-3.6-flash", messages: Optional[List[Dict[str, Any]]] = None, stream: bool = False, tools: Any = None, diff --git a/agent/insights.py b/agent/insights.py index 086150c279e..9d148a15446 100644 --- a/agent/insights.py +++ b/agent/insights.py @@ -113,6 +113,13 @@ class InsightsEngine: """ cutoff = time.time() - (days * 86400) + # Token/cost totals may still sit on the SessionDB's async + # accounting queue; drain so the report reflects exact counters. + # (self.db may be a raw sqlite3 connection in tests — guard.) + flush = getattr(self.db, "flush_token_counts", None) + if callable(flush): + flush() + # Gather raw data sessions = self._get_sessions(cutoff, source) tool_usage = self._get_tool_usage(cutoff, source) diff --git a/agent/lsp/eventlog.py b/agent/lsp/eventlog.py index b38627504b4..f118ccf0ace 100644 --- a/agent/lsp/eventlog.py +++ b/agent/lsp/eventlog.py @@ -40,7 +40,7 @@ from __future__ import annotations import logging import os import threading -from typing import Tuple +from typing import List, Tuple # Dedicated logger name so the documented grep recipe survives a # ``logging.getLogger(__name__)`` rename of any internal module. @@ -188,6 +188,25 @@ def log_spawn_failed(server_id: str, workspace_root: str, exc: BaseException) -> ) +def log_reaped(keys: List[Tuple[str, str]], idle_timeout: float) -> None: + """Idle clients were shut down by the reaper. INFO — one line per + sweep so users can correlate memory drops with LSP activity. + + Also clears the ``log_active`` announce cache for the reaped keys so + a later respawn re-announces at INFO instead of logging a misleading + DEBUG "reused client". + """ + with _announce_lock: + for key in keys: + _announced_active.discard(key) + summary = ", ".join(f"{sid} ({root})" for sid, root in keys) + _emit( + "reaper", + logging.INFO, + f"reaped {len(keys)} idle client(s) after {idle_timeout:.0f}s: {summary}", + ) + + def reset_announce_caches() -> None: """Test-only: clear the dedup caches. Production code never calls this.""" with _announce_lock: @@ -209,5 +228,6 @@ __all__ = [ "log_timeout", "log_server_error", "log_spawn_failed", + "log_reaped", "reset_announce_caches", ] diff --git a/agent/lsp/manager.py b/agent/lsp/manager.py index d3b4244790b..757f7ebc7b9 100644 --- a/agent/lsp/manager.py +++ b/agent/lsp/manager.py @@ -59,6 +59,7 @@ from agent.lsp.workspace import ( logger = logging.getLogger("agent.lsp.manager") DEFAULT_IDLE_TIMEOUT = 600 # seconds; servers idle for >10min get reaped +MIN_IDLE_TIMEOUT = 30 # floor for config values; must exceed any per-op wait budget class _BackgroundLoop: @@ -176,6 +177,7 @@ class LSPService: self._spawning: Dict[Tuple[str, str], asyncio.Future] = {} self._last_used: Dict[Tuple[str, str], float] = {} self._state_lock = threading.Lock() + self._idle_reaper_task: Optional[asyncio.Task] = None # Delta baseline: file path → snapshot of diagnostics taken # immediately before a write. ``get_diagnostics_sync`` filters @@ -183,6 +185,9 @@ class LSPService: # introduced by the current edit. self._delta_baseline: Dict[str, List[Dict[str, Any]]] = {} + if self._enabled and self._idle_timeout > 0: + self._loop.run(self._start_idle_reaper(), timeout=2.0) + @classmethod def create_from_config(cls) -> Optional["LSPService"]: """Build a service from ``hermes_cli.config`` settings. @@ -205,6 +210,16 @@ class LSPService: wait_mode = lsp_cfg.get("wait_mode", "document") wait_timeout = float(lsp_cfg.get("wait_timeout", DIAGNOSTICS_DOCUMENT_WAIT)) install_strategy = lsp_cfg.get("install_strategy", "auto") + try: + idle_timeout = float(lsp_cfg.get("idle_timeout", DEFAULT_IDLE_TIMEOUT)) + except (TypeError, ValueError): + idle_timeout = DEFAULT_IDLE_TIMEOUT + if 0 < idle_timeout < MIN_IDLE_TIMEOUT: + # A timeout below the per-operation wait budget could reap a + # client mid-flight; the resulting outer timeout would then + # mark the (server, workspace) pair broken for the process + # lifetime. Clamp to a safe floor (0 still disables). + idle_timeout = MIN_IDLE_TIMEOUT servers_cfg = lsp_cfg.get("servers") or {} disabled = [] binary_overrides: Dict[str, List[str]] = {} @@ -235,6 +250,7 @@ class LSPService: env_overrides=env_overrides, init_overrides=init_overrides, disabled_servers=disabled, + idle_timeout=idle_timeout, ) # ------------------------------------------------------------------ @@ -434,6 +450,7 @@ class LSPService: # ``_clients`` with a half-initialized state. with self._state_lock: client = self._clients.pop(key, None) + self._last_used.pop(key, None) if client is not None: try: # Fire-and-forget shutdown — give it a second to cleanup, @@ -470,7 +487,7 @@ class LSPService: except Exception as e: # noqa: BLE001 logger.debug("snapshot open/wait failed: %s", e) return [] - self._last_used[(client.server_id, client.workspace_root)] = time.time() + self._touch(client) if not fresh: # No fresh data for the pre-edit content — an empty baseline # is safe: worst case the delta filter removes less, never @@ -499,7 +516,7 @@ class LSPService: except Exception as e: # noqa: BLE001 logger.debug("open/wait failed for %s: %s", file_path, e) return None - self._last_used[(client.server_id, client.workspace_root)] = time.time() + self._touch(client) if not fresh: return None return list(client.diagnostics_for(file_path, fresh_only=True)) @@ -539,6 +556,7 @@ class LSPService: with self._state_lock: client = self._clients.get(key) if client is not None and client.is_running: + self._last_used[key] = time.time() eventlog.log_active(srv.server_id, per_server_root) return client spawning = self._spawning.get(key) @@ -589,7 +607,7 @@ class LSPService: return None with self._state_lock: self._clients[key] = client - self._last_used[key] = time.time() + self._last_used[key] = time.time() eventlog.log_active(srv.server_id, per_server_root) spawn_future.set_result(client) return client @@ -597,7 +615,63 @@ class LSPService: with self._state_lock: self._spawning.pop(key, None) + async def _start_idle_reaper(self) -> None: + self._idle_reaper_task = asyncio.create_task(self._idle_reaper_loop()) + + def _touch(self, client: LSPClient) -> None: + """Refresh the last-used timestamp for a client we just used. + + Guarded on membership so a reaped-mid-operation client can't + resurrect an orphan ``_last_used`` entry after the reaper popped + the key. All writers and the reaper run on the background loop + thread; the lock keeps this consistent with the reader anyway. + """ + key = (client.server_id, client.workspace_root) + with self._state_lock: + if key in self._clients: + self._last_used[key] = time.time() + + async def _idle_reaper_loop(self) -> None: + interval = min(60.0, self._idle_timeout) + while True: + await asyncio.sleep(interval) + try: + await self._reap_idle_once() + except asyncio.CancelledError: + raise + except Exception as e: # noqa: BLE001 + # A transient sweep error must not kill the reaper — + # otherwise one bad shutdown permanently re-opens the + # unbounded-accumulation leak this loop exists to fix. + logger.debug("LSP idle reaper sweep error: %s", e) + + async def _reap_idle_once(self) -> None: + cutoff = time.time() - self._idle_timeout + with self._state_lock: + idle_keys = [ + key + for key in self._clients + if self._last_used.get(key, 0) < cutoff + ] + clients = [self._clients.pop(key) for key in idle_keys] + for key in idle_keys: + self._last_used.pop(key, None) + if clients: + eventlog.log_reaped( + [(c.server_id, c.workspace_root) for c in clients], + self._idle_timeout, + ) + await asyncio.gather( + *(client.shutdown() for client in clients), + return_exceptions=True, + ) + async def _shutdown_async(self) -> None: + reaper = self._idle_reaper_task + self._idle_reaper_task = None + if reaper is not None: + reaper.cancel() + await asyncio.gather(reaper, return_exceptions=True) with self._state_lock: clients = list(self._clients.values()) self._clients.clear() diff --git a/agent/model_metadata.py b/agent/model_metadata.py index 296fe0aedca..757a306cde7 100644 --- a/agent/model_metadata.py +++ b/agent/model_metadata.py @@ -2429,21 +2429,27 @@ def get_model_context_length( if context_length is not None: return context_length if not _is_known_provider_base_url(base_url): - # 2b. Ollama native /api/show — any URL might be an Ollama server - # (local, cloud, or custom hosting). Non-Ollama servers return - # 404/405 quickly. Fall through on failure. - ctx = _query_ollama_api_show(model, base_url, api_key=api_key) - if ctx is not None: - if not _skip_persistent_context_cache(base_url, provider): - save_context_length(model, base_url, ctx) - return ctx - # 3. Try querying local server directly + # For local endpoints, run the probe that respects configured + # Modelfile context values first. _query_local_context_length + # prefers num_ctx from Modelfile, while _query_ollama_api_show + # returns the GGUF training max first which can be larger and + # would create a false-safe window for compression (#63122). + # Non-local endpoints preserve the existing GGUF-first behavior. if is_local_endpoint(base_url): local_ctx = _query_local_context_length(model, base_url, api_key=api_key) if local_ctx and local_ctx > 0: if not _skip_persistent_context_cache(base_url, provider): _maybe_cache_local_context_length(model, base_url, local_ctx) return local_ctx + # 2b. Ollama native /api/show — non-local endpoints preserve + # the existing generic /api/show GGUF-first behavior. + # Non-Ollama servers return 404/405 quickly. + ctx = _query_ollama_api_show(model, base_url, api_key=api_key) + if ctx is not None: + if not _skip_persistent_context_cache(base_url, provider): + save_context_length(model, base_url, ctx) + return ctx + # 3. Probe-down fallback after endpoint-specific detection failed logger.info( "Could not detect context length for model %r at %s — " "defaulting to %s tokens (probe-down). Set model.context_length " diff --git a/agent/models_dev.py b/agent/models_dev.py index 590f77806ab..54030c0965c 100644 --- a/agent/models_dev.py +++ b/agent/models_dev.py @@ -8,11 +8,15 @@ of 4000+ models across 109+ providers. Provides: (reasoning, tools, vision, PDF, audio), modalities, knowledge cutoff, open-weights flag, family grouping, deprecation status -Data resolution order (like TypeScript OpenCode): - 1. Bundled snapshot (ships with the package — offline-first) - 2. Disk cache (~/.hermes/models_dev_cache.json) - 3. Network fetch (https://models.dev/api.json) - 4. Background refresh every 60 minutes +Data resolution order: + 1. In-memory cache (fresh, or stale served immediately while a single + background daemon thread refreshes) + 2. Disk cache (~/.hermes/models_dev_cache.json — any age; stale data is + served rather than blocking callers on the network) + 3. Network fetch (https://models.dev/api.json) — only when no cache + exists at all; failed refreshes back off for 5 minutes process-wide +Latency-sensitive callers (gateway route-identity checks) pass +``allow_network=False`` and never touch the network. Other modules should import the dataclasses and query functions from here rather than parsing the raw JSON themselves. @@ -20,6 +24,7 @@ rather than parsing the raw JSON themselves. import json import logging +import threading import time from dataclasses import dataclass from pathlib import Path @@ -33,10 +38,15 @@ logger = logging.getLogger(__name__) MODELS_DEV_URL = "https://models.dev/api.json" _MODELS_DEV_CACHE_TTL = 3600 # 1 hour in-memory +_MODELS_DEV_RETRY_DELAY = 300 # 5 minutes after a failed refresh # In-memory cache _models_dev_cache: Dict[str, Any] = {} _models_dev_cache_time: float = 0 +_models_dev_retry_after: float = 0 +_models_dev_fetch_lock = threading.Lock() +_models_dev_refresh_lock = threading.Lock() +_models_dev_refresh_in_flight = False # --------------------------------------------------------------------------- @@ -237,27 +247,152 @@ def _save_disk_cache(data: Dict[str, Any]) -> None: logger.debug("Failed to save models.dev disk cache: %s", e) -def fetch_models_dev(force_refresh: bool = False) -> Dict[str, Any]: +def _fetch_models_dev_from_network() -> Dict[str, Any]: + """Fetch the live models.dev registry without touching local caches. + + Raises on network errors and on an empty/invalid registry payload. + """ + response = requests.get(MODELS_DEV_URL, timeout=15) + response.raise_for_status() + data = response.json() + if not isinstance(data, dict) or not data: + raise ValueError("models.dev returned an empty or invalid registry") + return data + + +def _mark_stale_cache_grace() -> None: + """Give stale cache data a short in-memory grace before retrying refresh. + + Only ever moves the timestamp forward: if a background refresh completed + between the caller's staleness check and this call, the fresh timestamp + is preserved instead of being rewound to a 5-minute grace. + """ + global _models_dev_cache_time + grace_time = time.time() - _MODELS_DEV_CACHE_TTL + _MODELS_DEV_RETRY_DELAY + if grace_time > _models_dev_cache_time: + _models_dev_cache_time = grace_time + + +def _commit_registry(data: Dict[str, Any], *, where: str) -> None: + """Persist a freshly fetched registry: disk + in-mem + clear backoff. + + Callers must hold ``_models_dev_fetch_lock`` so a failing refresh on one + path can never stomp the state a succeeding refresh on the other path + just committed (e.g. a failing background worker re-arming the backoff + immediately after a successful ``force_refresh``). + """ + global _models_dev_cache, _models_dev_cache_time, _models_dev_retry_after + _save_disk_cache(data) + _models_dev_cache = data + _models_dev_cache_time = time.time() + _models_dev_retry_after = 0 + logger.debug( + "Refreshed models.dev registry (%s): %d providers, %d total models", + where, + len(data), + sum(len(p.get("models", {})) for p in data.values() if isinstance(p, dict)), + ) + + +def _note_refresh_failure(exc: Exception, *, where: str) -> None: + """Record a failed refresh: arm the process-wide 5-minute backoff. + + Callers must hold ``_models_dev_fetch_lock`` (see ``_commit_registry``). + """ + global _models_dev_retry_after + _models_dev_retry_after = time.time() + _MODELS_DEV_RETRY_DELAY + logger.debug( + "models.dev refresh failed (%s); retry suppressed for %ds: %s", + where, + _MODELS_DEV_RETRY_DELAY, + exc, + ) + + +def _background_refresh_models_dev() -> None: + """Best-effort refresh after serving stale cache data.""" + global _models_dev_refresh_in_flight + try: + data = _fetch_models_dev_from_network() + with _models_dev_fetch_lock: + _commit_registry(data, where="background") + except Exception as e: + with _models_dev_fetch_lock: + _note_refresh_failure(e, where="background") + finally: + with _models_dev_refresh_lock: + _models_dev_refresh_in_flight = False + + +def _start_background_refresh_models_dev() -> None: + """Start one daemon refresh worker if none is already running. + + Honors the process-wide failure backoff: after a failed refresh, + no new background worker is spawned until ``_models_dev_retry_after``. + """ + global _models_dev_refresh_in_flight + if time.time() < _models_dev_retry_after: + return + with _models_dev_refresh_lock: + if _models_dev_refresh_in_flight: + return + _models_dev_refresh_in_flight = True + thread = threading.Thread( + target=_background_refresh_models_dev, + name="models-dev-refresh", + daemon=True, + ) + try: + thread.start() + except Exception as e: + # Thread/fd exhaustion: clear the flag so refresh isn't disabled + # for the rest of the process lifetime. Callers still get stale data. + with _models_dev_refresh_lock: + _models_dev_refresh_in_flight = False + logger.debug("Failed to start models.dev refresh thread: %s", e) + + +def fetch_models_dev( + force_refresh: bool = False, *, allow_network: bool = True +) -> Dict[str, Any]: """Fetch models.dev registry. Cache hierarchy: in-mem → disk → network. Returns the full registry dict keyed by provider ID, or empty dict on failure. Cache hierarchy (when ``force_refresh=False``): - 1. In-memory cache, populated and < TTL old → return immediately. - 2. **Disk cache file < TTL old by mtime → load, populate in-mem, return.** - No network call. Saves ~500 ms per cold-start agent construction; - ``models.dev`` only changes when providers add new models, so a - 1 hour staleness window is acceptable (same TTL as in-mem cache). - 3. Network fetch → on success, save to disk + in-mem and return. - 4. Network fails → fall back to ANY available disk cache (even stale) - with a short 5 min in-mem grace period before retrying network. + 1. Fresh in-memory cache → return immediately. + 2. Stale in-memory cache → return immediately and refresh in a single + background daemon thread. Callers never block on the network while + any cache exists; ``models.dev`` only changes when providers add + new models, so stale data is preferable to a foreground timeout. + 3. Disk cache file (any age) → load, populate in-mem, return + immediately. Stale disk caches trigger the same background refresh. + 4. No cache at all → singleflight foreground network fetch. On + success, save to disk + in-mem and return. + 5. Any failed refresh (foreground or background) suppresses further + automatic refreshes for 5 minutes process-wide. When ``force_refresh=True`` (used by ``hermes config refresh``, the - \"refresh model catalog\" code path), stages 1 and 2 are skipped. The - function always hits the network and only falls back to disk if the - network call fails. + \"refresh model catalog\" code path), cache fast paths and the failure + backoff are bypassed; the function hits the network and only falls back + to cached data if the call fails. When ``allow_network=False``, any + memory or disk cache is returned regardless of age and no request is + made — used by latency-sensitive paths (gateway route-identity checks) + that must never wait on the network. """ - global _models_dev_cache, _models_dev_cache_time + global _models_dev_cache, _models_dev_cache_time, _models_dev_retry_after + + if not allow_network: + if _models_dev_cache: + return _models_dev_cache + disk_data = _load_disk_cache() + if disk_data: + _models_dev_cache = disk_data + disk_age = _disk_cache_age_seconds() + _models_dev_cache_time = ( + time.time() - disk_age if disk_age is not None else 0 + ) + return _models_dev_cache # Stage 1: fresh in-memory cache wins. This is the hot path on # long-lived processes — no I/O, no system calls. @@ -268,54 +403,82 @@ def fetch_models_dev(force_refresh: bool = False) -> Dict[str, Any]: ): return _models_dev_cache - # Stage 2: fresh-by-mtime disk cache short-circuits the network call. - # Only kicks in on cold-start processes (in-mem cache is empty or - # expired) and only when the user hasn't asked for a forced refresh. - # Skipped if the disk cache file is missing, unreadable, or older - # than _MODELS_DEV_CACHE_TTL. + # Stage 2: stale in-memory cache is still better than blocking provider + # resolution on a foreground network timeout. Refresh it in the background. + if not force_refresh and _models_dev_cache: + _mark_stale_cache_grace() + _start_background_refresh_models_dev() + logger.debug( + "Using stale in-memory models.dev cache; refreshing in background" + ) + return _models_dev_cache + + # Stage 3: disk cache short-circuits the network call. + # Only kicks in on cold-start processes (in-mem cache is empty) and only + # when the user hasn't asked for a forced refresh. A stale disk cache is + # deliberately usable: provider/model resolution should not hang just + # because models.dev is unreachable. if not force_refresh: disk_age = _disk_cache_age_seconds() - if disk_age is not None and disk_age < _MODELS_DEV_CACHE_TTL: + if disk_age is not None: disk_data = _load_disk_cache() if disk_data: _models_dev_cache = disk_data - # Anchor in-mem TTL to the disk file's age so we don't - # extend an already-aging cache by another full hour. - _models_dev_cache_time = time.time() - disk_age - logger.debug( - "Loaded models.dev from fresh disk cache " - "(%d providers, age=%.0fs)", len(disk_data), disk_age, - ) + if disk_age < _MODELS_DEV_CACHE_TTL: + # Anchor in-mem TTL to the disk file's age so we don't + # extend an already-aging cache by another full hour. + _models_dev_cache_time = time.time() - disk_age + logger.debug( + "Loaded models.dev from fresh disk cache " + "(%d providers, age=%.0fs)", len(disk_data), disk_age, + ) + else: + _mark_stale_cache_grace() + _start_background_refresh_models_dev() + logger.debug( + "Using stale models.dev disk cache (age=%.0fs); " + "refreshing in background", + disk_age, + ) return _models_dev_cache - # Stage 3: network fetch. - try: - response = requests.get(MODELS_DEV_URL, timeout=15) - response.raise_for_status() - data = response.json() - if isinstance(data, dict) and data: - _models_dev_cache = data - _models_dev_cache_time = time.time() - _save_disk_cache(data) - logger.debug( - "Fetched models.dev registry: %d providers, %d total models", - len(data), - sum(len(p.get("models", {})) for p in data.values() if isinstance(p, dict)), - ) + # Failed automatic refreshes are process-wide. Avoid making every caller + # retry the same unreachable endpoint while no usable cache exists. + if not force_refresh and time.time() < _models_dev_retry_after: + return _models_dev_cache + + # Stage 4: singleflight foreground network fetch — only reached when no + # memory or disk cache exists (or on force_refresh). Recheck state after + # acquiring the lock because another caller may have refreshed or + # established backoff while we waited. + with _models_dev_fetch_lock: + now = time.time() + if not force_refresh: + if _models_dev_cache: + return _models_dev_cache + if now < _models_dev_retry_after: + return _models_dev_cache + + try: + data = _fetch_models_dev_from_network() + _commit_registry(data, where="foreground") return data - except Exception as e: - logger.debug("Failed to fetch models.dev: %s", e) + except Exception as e: + _note_refresh_failure(e, where="foreground") - # Stage 4: network failed — fall back to whatever disk cache exists, - # even if it's stale. Give it a short 5 min in-mem TTL so we retry - # the network soon instead of serving stale data for a full hour. - if not _models_dev_cache: - _models_dev_cache = _load_disk_cache() - if _models_dev_cache: - _models_dev_cache_time = time.time() - _MODELS_DEV_CACHE_TTL + 300 - logger.debug("Loaded models.dev from disk cache (%d providers)", len(_models_dev_cache)) + # Stage 5: network failed — return any stale memory/disk cache. Cache + # freshness remains expired; the retry-after timestamp controls when + # the next automatic request is allowed. + if not _models_dev_cache: + _models_dev_cache = _load_disk_cache() + _models_dev_cache_time = 0 + if _models_dev_cache: + logger.debug( + "Loaded stale models.dev disk cache (%d providers)", + len(_models_dev_cache), + ) - return _models_dev_cache + return _models_dev_cache def lookup_models_dev_context(provider: str, model: str) -> Optional[int]: @@ -671,7 +834,9 @@ def _parse_provider_info(provider_id: str, raw: Dict[str, Any]) -> ProviderInfo: # Provider-level queries # --------------------------------------------------------------------------- -def get_provider_info(provider_id: str) -> Optional[ProviderInfo]: +def get_provider_info( + provider_id: str, *, allow_network: bool = True +) -> Optional[ProviderInfo]: """Get full provider metadata from models.dev. Accepts either a Hermes provider ID (e.g. "kilocode") or a models.dev @@ -680,7 +845,14 @@ def get_provider_info(provider_id: str) -> Optional[ProviderInfo]: # Resolve Hermes ID → models.dev ID mdev_id = PROVIDER_TO_MODELS_DEV.get(provider_id, provider_id) - data = fetch_models_dev() + # NOTE: keep the zero-argument call on the default path. Dozens of test + # sites monkeypatch fetch_models_dev with zero-arg lambdas; passing the + # kwarg unconditionally would break them all (they raise TypeError). + data = ( + fetch_models_dev() + if allow_network + else fetch_models_dev(allow_network=False) + ) raw = data.get(mdev_id) if not isinstance(raw, dict): return None diff --git a/agent/prompt_caching.py b/agent/prompt_caching.py index 2e606cd377c..1a0326c79df 100644 --- a/agent/prompt_caching.py +++ b/agent/prompt_caching.py @@ -187,17 +187,18 @@ def apply_anthropic_cache_control( is retained. Returns: - Deep copy of messages with cache_control breakpoints injected. + Shallow copy of message list with selective deep copies of modified messages. """ - messages = copy.deepcopy(api_messages) - if not messages: - return messages + if not api_messages: + return api_messages + messages = list(api_messages) marker = _build_marker(cache_ttl) breakpoints_used = 0 if messages[0].get("role") == "system": + messages[0] = copy.deepcopy(messages[0]) breakpoints_used = _apply_system_cache_markers( messages[0], marker, @@ -213,6 +214,7 @@ def apply_anthropic_cache_control( and _can_carry_marker(messages[i], native_anthropic=native_anthropic) ] for idx in non_sys[-remaining:]: + messages[idx] = copy.deepcopy(messages[idx]) _apply_cache_marker(messages[idx], marker, native_anthropic=native_anthropic) return messages diff --git a/agent/reasoning_timeouts.py b/agent/reasoning_timeouts.py index 9c5fc202015..da7fcd2fcbe 100644 --- a/agent/reasoning_timeouts.py +++ b/agent/reasoning_timeouts.py @@ -146,19 +146,18 @@ _REASONING_STALE_TIMEOUT_FLOORS: tuple[tuple[str, int], ...] = ( # so we accept that community forks inheriting the same prefix are # treated as reasoning models (a reasonable default — the upstream # gateway timing is the same). -_PATTERN_CACHE: dict[str, re.Pattern[str]] = {} - - -def _get_pattern(slug: str) -> re.Pattern[str]: - compiled = _PATTERN_CACHE.get(slug) - if compiled is None: - compiled = re.compile( - r"^" - + re.escape(slug) - + r"(?:$|[\-._])" - ) - _PATTERN_CACHE[slug] = compiled - return compiled +# Pre-compile all patterns at module load time to avoid per-call regex +# compilation and thread-safety issues with the mutable _PATTERN_CACHE. +# The list is built once at import and never mutated afterwards, so it is +# safe for free-threaded Python 3.13+ without any locking. The slug is kept +# in each entry for debuggability (log/inspection), even though _match_any +# only consumes floor + pattern. +_SORTED_REASONING_FLOORS: list[tuple[str, float, re.Pattern[str]]] = [ + (slug, floor, re.compile(r"^" + re.escape(slug) + r"(?:$|[\-._])")) + for slug, floor in sorted( + _REASONING_STALE_TIMEOUT_FLOORS, key=lambda kv: -len(kv[0]) + ) +] def _match_any(model_lower: str) -> Optional[float]: @@ -169,13 +168,8 @@ def _match_any(model_lower: str) -> Optional[float]: order is irrelevant: longest slug wins (so ``o3-mini`` beats ``o3`` on a model like ``openai/o3-mini``). """ - # Sort by slug length descending so longer / more-specific slugs - # win on shared prefixes (o3-mini beats o3). - sorted_floors = sorted( - _REASONING_STALE_TIMEOUT_FLOORS, key=lambda kv: -len(kv[0]) - ) - for slug, floor in sorted_floors: - if _get_pattern(slug).search(model_lower): + for _slug, floor, pattern in _SORTED_REASONING_FLOORS: + if pattern.search(model_lower): return float(floor) return None diff --git a/agent/relay_llm.py b/agent/relay_llm.py new file mode 100644 index 00000000000..3481ca07237 --- /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 00000000000..533604791a8 --- /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 00000000000..5023df1bcf9 --- /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 d32fe99c0c5..1dfcf1ebf6e 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 94dc1660d1a..e080d6a5d96 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 4e2d318b2e2..0f636c1dd24 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 9825e4a27bd..b7982e46ab0 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/backend-probes.test.ts b/apps/desktop/electron/backend-probes.test.ts index bddc61152f8..2ff97814e61 100644 --- a/apps/desktop/electron/backend-probes.test.ts +++ b/apps/desktop/electron/backend-probes.test.ts @@ -14,7 +14,10 @@ import { test } from 'vitest' import { canImportHermesCli, + DEFAULT_PROBE_TIMEOUT_MS, hermesRuntimeImportProbe, + PROBE_TIMEOUT_MS, + resolveProbeTimeoutMs, shouldTrustHermesOverride, verifyHermesCli } from './backend-probes' @@ -100,9 +103,25 @@ test('verifyHermesCli returns true when --version exits 0', () => { }) test('verifyHermesCli swallows timeouts (does not throw)', () => { - // We can't easily provoke a real 5s hang in CI without slowing the + // We can't easily provoke a real hang in CI without slowing the // suite, but we CAN confirm that an invocation that DOES throw // (because the binary is missing) returns false rather than // propagating. Same code path the timeout case takes. assert.equal(verifyHermesCli('/definitely/not/a/real/binary/anywhere'), false) }) + +test('default probe timeout is 15s (not the old 5s death-loop value)', () => { + assert.equal(DEFAULT_PROBE_TIMEOUT_MS, 15_000) + // Module constant uses process.env at load time; with no override it + // matches the default (tests run without HERMES_PROBE_TIMEOUT_MS). + assert.equal(PROBE_TIMEOUT_MS, DEFAULT_PROBE_TIMEOUT_MS) +}) + +test('resolveProbeTimeoutMs honours HERMES_PROBE_TIMEOUT_MS', () => { + assert.equal(resolveProbeTimeoutMs({}), DEFAULT_PROBE_TIMEOUT_MS) + assert.equal(resolveProbeTimeoutMs({ HERMES_PROBE_TIMEOUT_MS: '30000' }), 30_000) + assert.equal(resolveProbeTimeoutMs({ HERMES_PROBE_TIMEOUT_MS: '0' }), DEFAULT_PROBE_TIMEOUT_MS) + assert.equal(resolveProbeTimeoutMs({ HERMES_PROBE_TIMEOUT_MS: 'nope' }), DEFAULT_PROBE_TIMEOUT_MS) + // Cap runaway values + assert.equal(resolveProbeTimeoutMs({ HERMES_PROBE_TIMEOUT_MS: '999999' }), 120_000) +}) diff --git a/apps/desktop/electron/backend-probes.ts b/apps/desktop/electron/backend-probes.ts index f196f5a7b77..7133342141e 100644 --- a/apps/desktop/electron/backend-probes.ts +++ b/apps/desktop/electron/backend-probes.ts @@ -20,8 +20,9 @@ * actually works. * * Both probes are deliberately fast and forgiving: - * - 5s timeout (a hung interpreter beats forever, but we still give - * slow disks / cold caches room to breathe) + * - default 15s timeout (5s was too short on cold Windows disks / AV; + * issue #61764 death-loop) with HERMES_PROBE_TIMEOUT_MS override + * - one automatic retry after a timeout before declaring the runtime dead * - stdio ignored (we only care about exit code; stdout/stderr are * not surfaced to the user, just to recentHermesLog for forensics * via the caller's catch block if it chooses) @@ -34,7 +35,82 @@ import { execFileSync } from 'node:child_process' -const PROBE_TIMEOUT_MS = 5000 +/** Default probe budget. 5s false-negativeed healthy Windows cold starts (#61764). */ +const DEFAULT_PROBE_TIMEOUT_MS = 15_000 + +/** + * Resolve the backend probe timeout (ms). + * Honours HERMES_PROBE_TIMEOUT_MS when it parses as a positive integer. + */ +function resolveProbeTimeoutMs(env: NodeJS.ProcessEnv = process.env): number { + const raw = env.HERMES_PROBE_TIMEOUT_MS + + if (raw == null || raw === '') { + return DEFAULT_PROBE_TIMEOUT_MS + } + + const n = Number.parseInt(String(raw), 10) + + if (!Number.isFinite(n) || n <= 0) { + return DEFAULT_PROBE_TIMEOUT_MS + } + + // Clamp absurd values (ms) so a typo can't hang startup forever. + return Math.min(n, 120_000) +} + +const PROBE_TIMEOUT_MS = resolveProbeTimeoutMs() + +function isTimeoutError(err: unknown): boolean { + if (!err || typeof err !== 'object') { + return false + } + + const e = err as { code?: string; killed?: boolean; signal?: string } + + if (e.killed === true) { + return true + } + + if (e.code === 'ETIMEDOUT') { + return true + } + + // Node marks timed-out execFileSync with SIGTERM on some platforms. + if (e.signal === 'SIGTERM') { + return true + } + + return false +} + +/** + * Run execFileSync; on timeout only, retry once before failing. + * Non-timeout failures (ENOENT, non-zero exit) fail immediately. + */ +function execProbeSync( + command: string, + args: string[], + options: { + cwd?: string + env?: NodeJS.ProcessEnv + stdio: 'ignore' + timeout: number + shell?: boolean + windowsHide?: boolean + } +): void { + try { + execFileSync(command, args, options) + } catch (err) { + if (!isTimeoutError(err)) { + throw err + } + + // One cold-cache / AV miss should not force hermes-setup --update (#61764). + execFileSync(command, args, options) + } +} /** * Return the Python snippet used to verify Hermes can import far enough to @@ -71,7 +147,7 @@ function canImportHermesCli(pythonPath: string, opts: { env?: Record com.apple.security.device.audio-input + com.apple.security.device.camera + diff --git a/apps/desktop/electron/entitlements.mac.plist b/apps/desktop/electron/entitlements.mac.plist index 53fdf0fc437..a3defc5f8cf 100644 --- a/apps/desktop/electron/entitlements.mac.plist +++ b/apps/desktop/electron/entitlements.mac.plist @@ -10,5 +10,7 @@ com.apple.security.device.audio-input + com.apple.security.device.camera + diff --git a/apps/desktop/electron/hardening.test.ts b/apps/desktop/electron/hardening.test.ts index 1a5852f720a..3acf649e6d5 100644 --- a/apps/desktop/electron/hardening.test.ts +++ b/apps/desktop/electron/hardening.test.ts @@ -7,8 +7,13 @@ import { pathToFileURL } from 'node:url' import { test } from 'vitest' import { + ATTACHMENT_UPLOAD_DEFAULT_MAX_BYTES, + clampDataUrlReadMaxMb, + DATA_URL_READ_DEFAULT_MAX_MB, + dataUrlReadMaxBytesFromMb, DEFAULT_FETCH_TIMEOUT_MS, encryptDesktopSecret, + readFileDataUrlForIpc, resolveDirectoryForIpc, resolveReadableFileForIpc, resolveRequestedPathForIpc, @@ -24,6 +29,49 @@ async function rejectsWithCode(promise, code: string) { }) } +test('clampDataUrlReadMaxMb defaults and bounds the attach size preference', () => { + assert.equal(clampDataUrlReadMaxMb(undefined), DATA_URL_READ_DEFAULT_MAX_MB) + assert.equal(clampDataUrlReadMaxMb(0), 1) + assert.equal(clampDataUrlReadMaxMb(256), 256) + assert.equal(clampDataUrlReadMaxMb(99999), 4096) + assert.equal(dataUrlReadMaxBytesFromMb(16), 16 * 1024 * 1024) +}) + +test('attachment upload cap is bounded above the preview default', () => { + assert.equal(ATTACHMENT_UPLOAD_DEFAULT_MAX_BYTES, 256 * 1024 * 1024) + assert.ok(ATTACHMENT_UPLOAD_DEFAULT_MAX_BYTES > dataUrlReadMaxBytesFromMb(DATA_URL_READ_DEFAULT_MAX_MB)) +}) + +test('attachment data URL helper reads bytes above the preview default without changing that limit', async () => { + const tempDir = fs.mkdtempSync(path.join(os.tmpdir(), 'hermes-desktop-large-attachment-')) + const source = path.join(tempDir, 'large.bin') + const previewLimit = dataUrlReadMaxBytesFromMb(DATA_URL_READ_DEFAULT_MAX_MB) + const content = Buffer.alloc(previewLimit + 1024, 0x5a) + + try { + fs.writeFileSync(source, content) + + await assert.rejects( + resolveReadableFileForIpc(source, { + maxBytes: previewLimit, + purpose: 'File preview' + }), + /file is too large/ + ) + + const dataUrl = await readFileDataUrlForIpc(source, { + maxBytes: ATTACHMENT_UPLOAD_DEFAULT_MAX_BYTES, + mimeType: 'application/octet-stream', + purpose: 'Attachment upload' + }) + + assert.match(dataUrl, /^data:application\/octet-stream;base64,/) + assert.deepEqual(Buffer.from(dataUrl.slice(dataUrl.indexOf(',') + 1), 'base64'), content) + } finally { + fs.rmSync(tempDir, { recursive: true, force: true }) + } +}) + test('resolveTimeoutMs falls back to defaults and accepts overrides', () => { assert.equal(resolveTimeoutMs(undefined), DEFAULT_FETCH_TIMEOUT_MS) assert.equal(resolveTimeoutMs(0), DEFAULT_FETCH_TIMEOUT_MS) diff --git a/apps/desktop/electron/hardening.ts b/apps/desktop/electron/hardening.ts index 2d6b5331001..e671dc259b7 100644 --- a/apps/desktop/electron/hardening.ts +++ b/apps/desktop/electron/hardening.ts @@ -4,9 +4,34 @@ import path from 'node:path' import { fileURLToPath } from 'node:url' const DEFAULT_FETCH_TIMEOUT_MS = 15_000 -const DATA_URL_READ_MAX_BYTES = 16 * 1024 * 1024 +// Default / floor / ceiling for Desktop's data-URL file load (composer attach, +// image preview, etc.). The whole file is base64-buffered in main, so this is +// a memory guard — not a model limit. Settings → Chat takes a free-form MB +// value; 16 MB ships as default. The ceiling is only a typo guard (very large +// values can OOM / crash the app). +const DATA_URL_READ_DEFAULT_MAX_MB = 16 +const DATA_URL_READ_MIN_MAX_MB = 1 +const DATA_URL_READ_MAX_MAX_MB = 4096 +// Remote file.attach sends one base64 JSON-RPC frame. Cap the dedicated attach +// reader so the payload still fits uvicorn's raised ws_max_size (384 MiB) +// after base64 + framing. Preview stays on the Settings-configurable path. +const ATTACHMENT_UPLOAD_DEFAULT_MAX_BYTES = 256 * 1024 * 1024 const TEXT_PREVIEW_SOURCE_MAX_BYTES = 64 * 1024 * 1024 +function clampDataUrlReadMaxMb(value) { + const parsed = Number(value) + + if (!Number.isFinite(parsed)) { + return DATA_URL_READ_DEFAULT_MAX_MB + } + + return Math.min(DATA_URL_READ_MAX_MAX_MB, Math.max(DATA_URL_READ_MIN_MAX_MB, Math.round(parsed))) +} + +function dataUrlReadMaxBytesFromMb(maxMb) { + return clampDataUrlReadMaxMb(maxMb) * 1024 * 1024 +} + const SAFE_ENV_SUFFIXES = new Set(['dist', 'example', 'sample', 'template']) const SENSITIVE_EXTENSIONS = new Set(['.kdbx', '.p12', '.pem', '.pfx']) @@ -303,10 +328,34 @@ async function resolveReadableFileForIpc( return { realPath, resolvedPath, stat } } +async function readFileDataUrlForIpc( + filePath, + options: { + purpose?: string + baseDir?: fs.PathOrFileDescriptor + fs?: typeof fs + blockSensitive?: boolean + maxBytes?: number + mimeType: string + } +): Promise { + const fsImpl = options.fs || fs + const { resolvedPath } = await resolveReadableFileForIpc(filePath, options) + const data = await fsImpl.promises.readFile(resolvedPath) + + return `data:${options.mimeType};base64,${data.toString('base64')}` +} + export { - DATA_URL_READ_MAX_BYTES, + ATTACHMENT_UPLOAD_DEFAULT_MAX_BYTES, + clampDataUrlReadMaxMb, + DATA_URL_READ_DEFAULT_MAX_MB, + DATA_URL_READ_MAX_MAX_MB, + DATA_URL_READ_MIN_MAX_MB, + dataUrlReadMaxBytesFromMb, DEFAULT_FETCH_TIMEOUT_MS, encryptDesktopSecret, + readFileDataUrlForIpc, rejectUnsafePathSyntax, resolveDirectoryForIpc, resolveReadableFileForIpc, diff --git a/apps/desktop/electron/main.ts b/apps/desktop/electron/main.ts index 3dc41b78498..344b149d0d1 100644 --- a/apps/desktop/electron/main.ts +++ b/apps/desktop/electron/main.ts @@ -37,7 +37,13 @@ import { dashboardFallbackArgs, sourceDeclaresServe } from './backend-command' import { createBackendConnectionState } from './backend-connection-state' import { buildDesktopBackendEnv, normalizeHermesHomeRoot } from './backend-env' import { isReauthRequiredError, waitForHermesReady } from './backend-health' -import { canImportHermesCli, shouldTrustHermesOverride, verifyHermesCli } from './backend-probes' +import { + canImportHermesCli, + execProbeSync, + PROBE_TIMEOUT_MS, + shouldTrustHermesOverride, + verifyHermesCli +} from './backend-probes' import { waitForDashboardPortAnnouncement } from './backend-ready' import { shouldLatchBackendStartFailure, shouldLatchRemoteReauthFailure } from './backend-start-failure' import { detectRemoteDisplay, isWindowsBinaryPathInWsl, isWslEnvironment } from './bootstrap-platform' @@ -115,9 +121,13 @@ import { switchBranch } from './git-worktree-ops' import { - DATA_URL_READ_MAX_BYTES, + ATTACHMENT_UPLOAD_DEFAULT_MAX_BYTES, + clampDataUrlReadMaxMb, + DATA_URL_READ_DEFAULT_MAX_MB, + dataUrlReadMaxBytesFromMb, DEFAULT_FETCH_TIMEOUT_MS, encryptDesktopSecret as encryptDesktopSecretStrict, + readFileDataUrlForIpc, resolveReadableFileForIpc, resolveRequestedPathForIpc, resolveTimeoutMs, @@ -175,6 +185,7 @@ import { } from './ssh-connection' import { nativeOverlayWidth as computeNativeOverlayWidth, macTitleBarOverlayHeight } from './titlebar-overlay-width' import { resolveBehindCount, shouldCountCommits } from './update-count' +import { waitForUpdateClearance } from './update-gate' import { readLiveUpdateMarker, writeUpdateMarker } from './update-marker' import { runRebuildWithRetry } from './update-rebuild' import { @@ -188,6 +199,7 @@ import { } from './update-relaunch' import { isOfficialSshRemote, OFFICIAL_REPO_HTTPS_URL } from './update-remote' import { spawnUpdaterProcess } from './updater-process' +import { formatBlockerMessage, formatProbeFailedMessage, scanVenvBlockers } from './venv-blocker-scan' import { fetchMarketplaceThemes, searchMarketplaceThemes } from './vscode-marketplace' import { computeWindowOptions, @@ -960,7 +972,7 @@ app.setAboutPanelOptions({ // Custom scheme for streaming local media (video/audio) into the renderer. // Reading large media through `readFileDataUrl` failed: it base64-loads the -// whole file into memory and is hard-capped at DATA_URL_READ_MAX_BYTES (16 MB), +// whole file into memory and is hard-capped (default 16 MB, Settings → Chat), // so any non-trivial video silently refused to load. Streaming via a protocol // handler removes the size cap and gives the