From 42f0dd91a2254eac6a59e923380d3c91b7d283c7 Mon Sep 17 00:00:00 2001 From: WallaceNodded <268681894+WallaceNodded@users.noreply.github.com> Date: Wed, 22 Jul 2026 16:26:14 +0800 Subject: [PATCH 1/2] fix(agent): recover primary before fast transport fallback --- agent/conversation_loop.py | 31 ++-- ...fast_transport_recovery_before_fallback.py | 161 ++++++++++++++++++ 2 files changed, 180 insertions(+), 12 deletions(-) create mode 100644 tests/run_agent/test_fast_transport_recovery_before_fallback.py diff --git a/agent/conversation_loop.py b/agent/conversation_loop.py index 517c9d86ed7e8..f0add54575272 100644 --- a/agent/conversation_loop.py +++ b/agent/conversation_loop.py @@ -3358,21 +3358,28 @@ def _perform_api_call(next_api_kwargs): # Fall through to normal error handling if compression # is exhausted or didn't help. - # Eager fallback for rate-limit errors (429 or quota exhaustion) - # and transport errors (connection failure / timeout / provider - # overloaded). Rate limits and billing: switch immediately — - # the primary provider won't recover within the retry window. - # Transport errors: allow 1 retry first (transient hiccups - # recover), then fall back if the provider is truly unreachable. + # Eager fallback for rate-limit errors (429 or quota exhaustion), + # provider overload, and stale-stream failures. Ordinary fast + # connection errors deliberately complete the normal retry cycle + # so the primary-client recovery path below can rebuild the + # transport and grant one fresh retry cycle before fallback. + # Stale streams remain eager after one retry: each attempt can + # consume the full multi-minute stale timeout (#22277). is_rate_limited = classified.reason in { FailoverReason.rate_limit, FailoverReason.billing, FailoverReason.upstream_rate_limit, } - _is_transport_failure = classified.reason in { - FailoverReason.timeout, - FailoverReason.overloaded, - } + _is_overload_failure = ( + classified.reason == FailoverReason.overloaded + ) + _is_stale_timeout = ( + classified.reason == FailoverReason.timeout + and getattr(agent, "_consecutive_stale_streams", 0) > 0 + ) + _is_eager_transport_failure = ( + _is_overload_failure or _is_stale_timeout + ) # Z.AI Coding Plan GLM-5.2 overload 429s classify as # `overloaded` (to spare the credential pool), but `overloaded` # is excluded from `is_rate_limited` — the gate for the adaptive @@ -3387,7 +3394,7 @@ def _perform_api_call(next_api_kwargs): max_retries = max(max_retries, zai_coding_overload_retry_ceiling()) _should_fallback = ( is_rate_limited - or (_is_transport_failure and retry_count >= 2) + or (_is_eager_transport_failure and retry_count >= 2) ) if _should_fallback and agent._fallback_index < len(agent._fallback_chain): # Don't eagerly fallback if credential pool rotation may @@ -3418,7 +3425,7 @@ def _perform_api_call(next_api_kwargs): agent._buffer_status( "⚠️ Billing or credits exhausted — switching to fallback provider..." ) - elif _is_transport_failure: + elif _is_eager_transport_failure: agent._buffer_status( "⚠️ Provider unreachable — switching to fallback provider..." ) diff --git a/tests/run_agent/test_fast_transport_recovery_before_fallback.py b/tests/run_agent/test_fast_transport_recovery_before_fallback.py new file mode 100644 index 0000000000000..3abcaa7ada55b --- /dev/null +++ b/tests/run_agent/test_fast_transport_recovery_before_fallback.py @@ -0,0 +1,161 @@ +"""Regression coverage for #69186. + +Fast SDK/socket connection errors should reach the existing primary-client +recovery cycle before fallback. Stale-stream failures remain eligible for +bounded eager fallback so #22277 does not regress. +""" + +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +from run_agent import AIAgent + + +class APIConnectionError(Exception): + pass + + +def _tool_defs(): + return [ + { + "type": "function", + "function": { + "name": "web_search", + "description": "search", + "parameters": {"type": "object", "properties": {}}, + }, + } + ] + + +def _response(content: str): + message = SimpleNamespace(content=content, tool_calls=None) + choice = SimpleNamespace(message=message, finish_reason="stop") + return SimpleNamespace(choices=[choice], model="test-model", usage=None) + + +def _agent(): + fallback_chain = [ + { + "provider": "deepseek", + "model": "deepseek-chat", + "base_url": "https://api.deepseek.com", + } + ] + with ( + patch("run_agent.get_tool_definitions", return_value=_tool_defs()), + patch("run_agent.check_toolset_requirements", return_value={}), + patch("run_agent.OpenAI", return_value=MagicMock()), + ): + agent = AIAgent( + api_key="primary-key-abcdef12", + base_url="https://primary.example.com/v1", + provider="custom", + model="primary-model", + quiet_mode=True, + skip_context_files=True, + skip_memory=True, + fallback_model=fallback_chain, + ) + agent.client = MagicMock() + agent._api_max_retries = 3 + return agent + + +def _fallback_client(): + client = MagicMock() + client.api_key = "fallback-key-abcdef12" + client.base_url = "https://api.deepseek.com" + client._custom_headers = None + client.default_headers = None + return client + + +def _run(agent, api_call, *, fallback_client=None): + patches = [ + patch.object(agent, "_interruptible_api_call", side_effect=api_call), + patch.object(agent, "_persist_session"), + patch.object(agent, "_save_trajectory"), + patch.object(agent, "_cleanup_task_resources"), + patch("run_agent.OpenAI", return_value=MagicMock()), + patch("agent.agent_runtime_helpers.time.sleep"), + patch("agent.conversation_loop.time.sleep"), + patch("agent.model_metadata.get_model_context_length", return_value=200000), + ] + if fallback_client is not None: + patches.extend([ + patch( + "agent.auxiliary_client.resolve_provider_client", + return_value=(fallback_client, "deepseek-chat"), + ), + patch( + "hermes_cli.model_normalize.normalize_model_for_provider", + side_effect=lambda model, provider: model, + ), + ]) + + entered = [] + try: + for ctx in patches: + entered.append(ctx) + ctx.__enter__() + return agent.run_conversation("hello") + finally: + for ctx in reversed(entered): + ctx.__exit__(None, None, None) + + +def test_fast_connection_errors_recover_primary_before_fallback(): + """Three fast failures exhaust retries, then a rebuilt primary succeeds.""" + agent = _agent() + calls = [] + + def api_call(_kwargs): + calls.append((agent.provider, agent.model)) + if len(calls) <= 3: + raise APIConnectionError("Connection error.") + return _response("Recovered on rebuilt primary") + + with ( + patch.object( + agent, + "_try_recover_primary_transport", + wraps=agent._try_recover_primary_transport, + ) as recover, + patch.object( + agent, + "_try_activate_fallback", + wraps=agent._try_activate_fallback, + ) as fallback, + ): + result = _run(agent, api_call, fallback_client=_fallback_client()) + + assert result["completed"] is True + assert result["final_response"] == "Recovered on rebuilt primary" + assert calls == [("custom", "primary-model")] * 4 + recover.assert_called_once() + fallback.assert_not_called() + + +def test_stale_stream_timeout_keeps_bounded_eager_fallback(): + """A stale-detector-derived timeout still falls back after one retry.""" + agent = _agent() + agent._consecutive_stale_streams = 1 + calls = [] + + def api_call(_kwargs): + calls.append((agent.provider, agent.model)) + if agent.provider == "custom": + raise APIConnectionError("Connection closed after stale-stream kill.") + return _response("Recovered via fallback") + + result = _run(agent, api_call, fallback_client=_fallback_client()) + + assert result["completed"] is True + assert result["final_response"] == "Recovered via fallback" + assert calls == [ + ("custom", "primary-model"), + ("custom", "primary-model"), + ("deepseek", "deepseek-chat"), + ] + assert agent._fallback_activated is True From 331c6c9a01b82ca561ee29e37c2bee77b0c23060 Mon Sep 17 00:00:00 2001 From: WallaceNodded <268681894+WallaceNodded@users.noreply.github.com> Date: Wed, 29 Jul 2026 19:35:37 +0800 Subject: [PATCH 2/2] fix(agent): preserve stale provenance for Codex watchdogs --- agent/chat_completion_helpers.py | 8 ++++++++ tests/agent/test_codex_ttfb_watchdog.py | 2 ++ 2 files changed, 10 insertions(+) diff --git a/agent/chat_completion_helpers.py b/agent/chat_completion_helpers.py index b2e5c8653a483..6f09d0f489507 100644 --- a/agent/chat_completion_helpers.py +++ b/agent/chat_completion_helpers.py @@ -839,6 +839,11 @@ def _call(): agent._touch_activity( f"codex stream killed after {int(_elapsed)}s with no first byte" ) + # This is a stale-derived timeout just like the generic stale-call + # kill below. Record its provenance before returning to the retry + # loop so bounded eager fallback and the stale give-up breaker + # continue to apply to Codex TTFB stalls. + _bump_stale_streak(agent) # Wait briefly for the worker to notice the closed connection. t.join(timeout=2.0) if result["error"] is None and result["response"] is None: @@ -885,6 +890,9 @@ def _call(): agent._touch_activity( f"codex stream killed after {int(_event_stale_elapsed)}s with no SSE events" ) + # Preserve stale-timeout provenance for the retry/fallback gate + # and the cross-turn stale-stream circuit breaker. + _bump_stale_streak(agent) t.join(timeout=2.0) if result["error"] is None and result["response"] is None: result["error"] = TimeoutError( diff --git a/tests/agent/test_codex_ttfb_watchdog.py b/tests/agent/test_codex_ttfb_watchdog.py index d685f4ba3bacd..8fcbe354dba38 100644 --- a/tests/agent/test_codex_ttfb_watchdog.py +++ b/tests/agent/test_codex_ttfb_watchdog.py @@ -96,6 +96,7 @@ def fake_hang(api_kwargs, client=None, on_first_delta=None): elapsed = time.time() - t0 assert "TTFB" in str(excinfo.value) assert "codex_ttfb_kill" in closes + assert agent._consecutive_stale_streams == 1 # ~1s cutoff + 2s join grace; must be far under the 60s stale timeout. assert elapsed < 15, f"TTFB watchdog took {elapsed:.1f}s" finally: @@ -312,6 +313,7 @@ def fake_stream(api_kwargs, client=None, on_first_delta=None): assert "after first byte" in str(excinfo.value) assert "codex_stream_idle_kill" in closes assert "codex_ttfb_kill" not in closes + assert agent._consecutive_stale_streams == 1 finally: stop["flag"] = True