From 3f2e54da66b8f5da19d1dd7a8f1fa67c342133f5 Mon Sep 17 00:00:00 2001 From: Konstantin Khlopkov <47825603+kokhlo@users.noreply.github.com> Date: Sun, 30 Aug 2026 12:43:26 +0300 Subject: [PATCH 1/3] fix(agent): default relay completions to the progress-hook wrapper --- agent/auxiliary_client.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/agent/auxiliary_client.py b/agent/auxiliary_client.py index 9f77e631a1d73..fd14e8ef5377a 100644 --- a/agent/auxiliary_client.py +++ b/agent/auxiliary_client.py @@ -2443,7 +2443,12 @@ def _relay_sync_completion( from agent.auxiliary_wire import prepare_chat_messages kwargs = prepare_chat_messages(client, kwargs) - callback = create or (lambda request: client.chat.completions.create(**request)) + # The progress hook is installed per task, so every attempt for that + # task must go through _create_with_progress — including auxiliary + # retries and provider fallbacks. Defaulting at this seam covers every + # call site that predates the hook or forgets to thread create= through + # (18 of 24 today); call sites that pass their own create= keep it. + callback = create or (lambda request: _create_with_progress(client, request)) route = _relay_auxiliary_metadata(provider=provider, api_mode=api_mode) # Isolate only the provider callback so the owning thread can unwind its lease/DB # transaction on hard cancel without touching the shared client. From 65828fbddca7645b7b3f03bd31cb51dbeb1e3e08 Mon Sep 17 00:00:00 2001 From: Teknium <127238744+teknium1@users.noreply.github.com> Date: Wed, 9 Sep 2026 13:34:23 -0700 Subject: [PATCH 2/3] fix(auxiliary): async progress-hook wrapper + route the async primary through it (#98466) Complete the seam default from the salvaged commit: _relay_async_completion needs an async twin of _create_with_progress, and the async primary attempt previously streamed only for stream-only providers, so a hooked async compression call ticked the watchdog zero times. Adds one invariant test: with a hook installed, BOTH relay defaults stream and tick per chunk; without a hook they are byte-identical plain creates. Red on origin/main. --- agent/auxiliary_client.py | 50 +++++++++++---- tests/agent/test_aux_relay_progress_seam.py | 64 +++++++++++++++++++ .../test_auxiliary_explicit_cancellation.py | 3 +- 3 files changed, 103 insertions(+), 14 deletions(-) create mode 100644 tests/agent/test_aux_relay_progress_seam.py diff --git a/agent/auxiliary_client.py b/agent/auxiliary_client.py index fd14e8ef5377a..047e0d6b5896a 100644 --- a/agent/auxiliary_client.py +++ b/agent/auxiliary_client.py @@ -2443,11 +2443,8 @@ def _relay_sync_completion( from agent.auxiliary_wire import prepare_chat_messages kwargs = prepare_chat_messages(client, kwargs) - # The progress hook is installed per task, so every attempt for that - # task must go through _create_with_progress — including auxiliary - # retries and provider fallbacks. Defaulting at this seam covers every - # call site that predates the hook or forgets to thread create= through - # (18 of 24 today); call sites that pass their own create= keep it. + # The progress hook is installed per TASK, so every attempt (retries, recovery rungs, fallbacks) + # must stream through _create_with_progress or the compression watchdog sees silence (#98466). callback = create or (lambda request: _create_with_progress(client, request)) route = _relay_auxiliary_metadata(provider=provider, api_mode=api_mode) # Isolate only the provider callback so the owning thread can unwind its lease/DB @@ -2470,7 +2467,8 @@ async def _relay_async_completion( from agent.auxiliary_wire import prepare_chat_messages kwargs = prepare_chat_messages(client, kwargs) - callback = create or (lambda request: client.chat.completions.create(**request)) + # Async twin of the seam default above (#98466). + callback = create or (lambda request: _acreate_with_progress(client, request)) route = _relay_auxiliary_metadata(provider=provider, api_mode=api_mode) if route is None: return await callback(kwargs) @@ -6496,6 +6494,37 @@ async def _acreate_with_stream(client: Any, kwargs: Dict[str, Any], task: Option return await _aggregate_chat_stream_async(chunks, model=model, total_ceiling=total_ceiling) +def _async_client_streams_internally(client: Any) -> bool: + """Async twin of :func:`_client_streams_internally` (the async adapters are separate classes).""" + return isinstance(client, (AsyncCodexAuxiliaryClient, AsyncAnthropicAuxiliaryClient, AsyncBedrockAuxiliaryClient)) + + +async def _acreate_with_progress( + client: Any, kwargs: Dict[str, Any], task: Optional[str] = None, *, force_stream: bool = False +) -> Any: + """Async :func:`_create_with_progress`: stream + re-aggregate (ticking the hook per substantive + chunk) when a progress hook is active or the provider is stream-only; plain create otherwise.""" + _notify_aux_dispatch() + _notify_aux_progress() + if (not _aux_progress_active() and not force_stream) or _async_client_streams_internally(client): + response = await client.chat.completions.create(**kwargs) + if not _async_client_streams_internally(client): + _notify_aux_provider_response() + return response + try: + return await _acreate_with_stream(client, kwargs, task) + except Exception as exc: + if (force_stream or _is_transient_transport_error(exc) or _is_auth_error(exc) + or _is_payment_error(exc) or _is_rate_limit_error(exc)): + raise + logger.debug("Auxiliary %s: streamed async request failed (%s); retrying non-streaming", + task or "call", exc) + _notify_aux_dispatch() + response = await client.chat.completions.create(**kwargs) + _notify_aux_provider_response() + return response + + # Shared request head + recovery ladder for call_llm / async_call_llm: the entry points differ # only in how a request is awaited, so route resolution and the ordered recovery ladder are # written once. The ladder is a generator yielding ``_LadderStep`` requests and receiving the @@ -7342,15 +7371,10 @@ async def _async_call_llm_impl( try: # Retry ONCE on the same provider for a transient blip before fallback (see call_llm()). # (PR #16587) - _force_stream_async = ( - _provider_requires_stream(request_provider, req.base_info or req.resolved_base_url) - and not isinstance(client, ( - AsyncCodexAuxiliaryClient, AsyncAnthropicAuxiliaryClient, AsyncBedrockAuxiliaryClient))) + _force_stream_async = _provider_requires_stream(request_provider, req.base_info or req.resolved_base_url) async def _acreate(_kwargs: Dict[str, Any]) -> Any: - if _force_stream_async: - return await _acreate_with_stream(client, _kwargs, task) - return await client.chat.completions.create(**_kwargs) + return await _acreate_with_progress(client, _kwargs, task, force_stream=_force_stream_async) async def _primary(**validate_kw: Any) -> Any: return _validate_llm_response( diff --git a/tests/agent/test_aux_relay_progress_seam.py b/tests/agent/test_aux_relay_progress_seam.py new file mode 100644 index 0000000000000..9e14151600d4d --- /dev/null +++ b/tests/agent/test_aux_relay_progress_seam.py @@ -0,0 +1,64 @@ +"""Every relay attempt (retries, recovery rungs, fallbacks) streams through the progress hook (#98466). + +Before the seam default, only the primary attempt passed ``create=``; the 18 unwrapped relay sites went +out non-streaming and ticked the compression watchdog zero times, so a healthy still-generating summary +was killed at the idle deadline ("timed out after 120.0s with no output from the summary model"). +""" +import asyncio +from types import SimpleNamespace + +from agent import auxiliary_client as aux + + +def _chunk(text): + return SimpleNamespace(id="r1", model="m", usage=None, choices=[SimpleNamespace( + finish_reason=None, delta=SimpleNamespace(content=text, reasoning=None, reasoning_content=None, + reasoning_details=None, tool_calls=None))]) + + +class _SyncClient: + def __init__(self): + self.wire = [] + self.chat = SimpleNamespace(completions=SimpleNamespace(create=self._create)) + self.base_url = "https://example.test/v1" + + def _create(self, **kwargs): + self.wire.append(kwargs.get("stream")) + if kwargs.get("stream"): + return iter([_chunk("hello"), _chunk(" world")]) + return SimpleNamespace(choices=[SimpleNamespace(message=SimpleNamespace(content="plain"))]) + + +class _AsyncClient(_SyncClient): + async def _create(self, **kwargs): + self.wire.append(kwargs.get("stream")) + if kwargs.get("stream"): + async def agen(): + yield _chunk("hello") + yield _chunk(" world") + return agen() + return SimpleNamespace(choices=[SimpleNamespace(message=SimpleNamespace(content="plain"))]) + + +def test_relay_default_callback_streams_and_ticks_hook_sync_and_async(): + ticks = [] + with aux.aux_progress_hook(lambda: ticks.append(1)): + sync_client = _SyncClient() + resp = aux._relay_sync_completion(sync_client, {"model": "m", "messages": []}) + assert sync_client.wire == [True] + assert resp.choices[0].message.content == "hello world" + sync_ticks = len(ticks) + assert sync_ticks >= 2 # one per substantive chunk, on top of the dispatch tick + + async_client = _AsyncClient() + resp = asyncio.run(aux._relay_async_completion(async_client, {"model": "m", "messages": []})) + assert async_client.wire == [True] + assert resp.choices[0].message.content == "hello world" + assert len(ticks) - sync_ticks >= 2 + + +def test_relay_default_callback_is_plain_create_without_hook(): + client = _SyncClient() + resp = aux._relay_sync_completion(client, {"model": "m", "messages": []}) + assert client.wire == [None] + assert resp.choices[0].message.content == "plain" diff --git a/tests/agent/test_auxiliary_explicit_cancellation.py b/tests/agent/test_auxiliary_explicit_cancellation.py index 8330ab5806136..831250e853dec 100644 --- a/tests/agent/test_auxiliary_explicit_cancellation.py +++ b/tests/agent/test_auxiliary_explicit_cancellation.py @@ -549,7 +549,8 @@ def _create(**_kwargs: Any) -> Any: assert observed["protected"] is True assert observed["thread"] != caller - assert progress == ["tick"] + # The hook must reach the isolated worker thread; the seam wrapper adds its own dispatch ticks. + assert "tick" in progress def test_isolated_provider_worker_inherits_caller_contextvars() -> None: From 86c15dde708be00fd207f4775f2775d2bcefcdf5 Mon Sep 17 00:00:00 2001 From: Teknium <127238744+teknium1@users.noreply.github.com> Date: Wed, 9 Sep 2026 13:34:23 -0700 Subject: [PATCH 3/3] fix(auxiliary): compression sticks to the session's OpenAI endpoint; foreign-host rejections don't kill the key `auxiliary..provider: openai` was rewritten to custom + api.openai.com/v1 unconditionally, and the api-key discovery rung used the registry default endpoint, so proxy/gateway users (OPENAI_BASE_URL or providers.openai) had compression hop to the public endpoint with a proxy-issued key -> 401 -> the key/pool was quarantined and the session sat over the compression threshold. - _expand_direct_api_alias: a providers.openai entry keeps its name (named-custom branch applies its base_url/key); otherwise OPENAI_BASE_URL wins over the public default. - _resolve_api_key_provider: when the bound session runtime is that provider, use its endpoint + key. - _recoverable_pool_provider: a rejection at a host other than the session's configured endpoint for the same provider is an endpoint mismatch, not a dead key -> no rotation / unhealthy mark on the pool. --- agent/auxiliary_client.py | 44 ++++++++++++++++--- .../test_aux_session_endpoint_affinity.py | 26 +++++++++++ 2 files changed, 64 insertions(+), 6 deletions(-) create mode 100644 tests/agent/test_aux_session_endpoint_affinity.py diff --git a/agent/auxiliary_client.py b/agent/auxiliary_client.py index 047e0d6b5896a..422803425b25d 100644 --- a/agent/auxiliary_client.py +++ b/agent/auxiliary_client.py @@ -2041,6 +2041,14 @@ def _resolve_api_key_provider() -> Tuple[Optional[OpenAI], Optional[str]]: continue raw_base_url = str(creds.get("base_url", "")).strip().rstrip("/") or pconfig.inference_base_url via = "" + # The session's own endpoint wins for its provider: the key was issued for that gateway, and + # sending it to the registry default 401s, then quarantines the provider the main model is on. + runtime = _normalize_main_runtime(None) + if runtime.get("provider") == provider_id and runtime.get("base_url"): + raw_base_url = runtime["base_url"].rstrip("/") + if isinstance(runtime.get("api_key"), str) and runtime["api_key"]: + api_key = runtime["api_key"] + via = " (session endpoint)" model = _get_aux_model_for_provider(provider_id) or None if model is None: continue # skip provider if we don't know a valid aux model @@ -3284,24 +3292,39 @@ def _provider_for_host(base_url: str, table: Tuple[Tuple[str, str], ...]) -> Opt def _recoverable_pool_provider( resolved_provider: str, client: Any, main_runtime: Optional[Dict[str, Any]] = None ) -> Optional[str]: - """Infer which provider pool can recover the current auxiliary client.""" + """Infer which provider pool can recover the current auxiliary client. + None when the client targets a different host than the session's configured endpoint for that + provider: a rejection there says nothing about the key, so rotating/quarantining it would kill a + working credential (Miho report — proxy users).""" normalized = _normalize_aux_provider(resolved_provider) + base = str(getattr(client, "base_url", "") or "") + runtime = _normalize_main_runtime(main_runtime) + rt_base = str(runtime.get("base_url") or "") + if (base and rt_base and normalized == runtime.get("provider") + and not base_url_host_matches(base, base_url_hostname(rt_base))): + logger.info("Auxiliary: %s rejected at %s, but the session's %s endpoint is %s — " + "endpoint mismatch, not a dead key; skipping credential rotation", + normalized, base_url_hostname(base), normalized, base_url_hostname(rt_base)) + return None if normalized not in {"", "auto", "custom"}: return normalized - base = str(getattr(client, "base_url", "") or "") known = _provider_for_host(base, _POOL_PROVIDER_BY_HOST) if known is not None: return known # Providers outside the table (e.g. opencode-go): match base URL against registered # api_key providers so pool rotation works for them too. if main_runtime: - rt_provider = _normalize_main_runtime(main_runtime).get("provider", "") + runtime = _normalize_main_runtime(main_runtime) + rt_provider = runtime.get("provider", "") if rt_provider and rt_provider not in {"", "auto", "custom"}: with contextlib.suppress(Exception): from hermes_cli.auth import PROVIDER_REGISTRY pconfig = PROVIDER_REGISTRY.get(rt_provider) if pconfig and getattr(pconfig, "auth_type", None) == "api_key": - rt_base = str(getattr(pconfig, "inference_base_url", "") or "").rstrip("/") + # The pool's key was issued for the endpoint the main runtime actually uses; a + # rejection at any other host (registry default vs configured proxy) says nothing + # about that key, so it must not be marked exhausted. + rt_base = str(runtime.get("base_url") or getattr(pconfig, "inference_base_url", "") or "").rstrip("/") if rt_base and base_url_host_matches(base, base_url_hostname(rt_base)): return rt_provider return None @@ -5476,13 +5499,22 @@ def _unwrap_moa_provider(prov: str, mdl: Optional[str]) -> Tuple[str, Optional[s def _expand_direct_api_alias(prov: Optional[str], existing_base: Optional[str]) -> Tuple[Optional[str], Optional[str]]: - """``provider: openai`` → custom + api.openai.com/v1; a user base_url is kept but the provider still becomes custom.""" + """``provider: openai`` → custom + the user's OpenAI endpoint, api.openai.com/v1 only as the last resort. + + A ``providers.openai`` entry keeps the provider name so the named-custom branch applies its base_url and + key; otherwise ``OPENAI_BASE_URL`` (a proxy/gateway the OPENAI_API_KEY was issued for) wins over the + public endpoint — sending the proxy key to api.openai.com 401s and then quarantines a valid key. + """ if not prov: return prov, existing_base target_base = _AUX_DIRECT_API_BASE_URLS.get(prov.strip().lower()) if target_base is None: return prov, existing_base - return "custom", existing_base or target_base + with contextlib.suppress(Exception): + from hermes_cli.runtime_provider import _get_named_custom_provider + if _get_named_custom_provider(prov) is not None: + return prov, existing_base + return "custom", existing_base or os.getenv("OPENAI_BASE_URL", "").strip().rstrip("/") or target_base def _preserve_provider_with_base_url(prov: Optional[str]) -> bool: diff --git a/tests/agent/test_aux_session_endpoint_affinity.py b/tests/agent/test_aux_session_endpoint_affinity.py new file mode 100644 index 0000000000000..92dbeaf66339d --- /dev/null +++ b/tests/agent/test_aux_session_endpoint_affinity.py @@ -0,0 +1,26 @@ +"""Auxiliary routing sticks to the session's configured OpenAI endpoint; a rejection elsewhere is not a dead key. + +Proxy users (`OPENAI_BASE_URL` / `providers.openai` pointing at a corporate gateway) saw compression +hop to api.openai.com, 401 with the proxy-issued key, and then have that key quarantined. +""" +from types import SimpleNamespace + +from agent import auxiliary_client as aux + + +def test_openai_alias_prefers_configured_endpoint_over_public_default(monkeypatch): + monkeypatch.setenv("OPENAI_BASE_URL", "https://llm-proxy.corp.example/v1") + provider, base = aux._expand_direct_api_alias("openai", None) + assert provider == "custom" + assert base == "https://llm-proxy.corp.example/v1" + monkeypatch.delenv("OPENAI_BASE_URL") + assert aux._expand_direct_api_alias("openai", None) == ("custom", "https://api.openai.com/v1") + + +def test_rejection_at_foreign_host_does_not_name_the_session_pool(): + runtime = {"provider": "openai-api", "model": "gpt-5.4", + "base_url": "https://llm-proxy.corp.example/v1", "api_key": "sk-proxy"} + foreign = SimpleNamespace(base_url="https://api.openai.com/v1/", api_key="sk-proxy") + same = SimpleNamespace(base_url="https://llm-proxy.corp.example/v1/", api_key="sk-proxy") + assert aux._recoverable_pool_provider("openai-api", foreign, main_runtime=runtime) is None + assert aux._recoverable_pool_provider("openai-api", same, main_runtime=runtime) == "openai-api"