From 045644bc7cb1c123fa43d413a0c9d5ce0a0960c8 Mon Sep 17 00:00:00 2001 From: Remotework Date: Fri, 17 Apr 2026 03:09:52 -0400 Subject: [PATCH] fix: harden codex stale timeout and minimax overload retries --- agent/retry_utils.py | 38 +++++++++++++++++++++ run_agent.py | 56 ++++++++++++++++++++++++++----- tests/run_agent/test_run_agent.py | 22 ++++++++++++ tests/test_retry_utils.py | 22 +++++++++++- 4 files changed, 128 insertions(+), 10 deletions(-) diff --git a/agent/retry_utils.py b/agent/retry_utils.py index 71d6963f7b41..7774c5353789 100644 --- a/agent/retry_utils.py +++ b/agent/retry_utils.py @@ -9,6 +9,8 @@ import threading import time +from agent.error_classifier import FailoverReason + # Monotonic counter for jitter seed uniqueness within the same process. # Protected by a lock to avoid race conditions in concurrent retry paths # (e.g. multiple gateway sessions retrying simultaneously). @@ -55,3 +57,39 @@ def jittered_backoff( jitter = rng.uniform(0, jitter_ratio * delay) return delay + jitter + + +def select_retry_wait_time( + attempt: int, + *, + reason: FailoverReason | str | None = None, + retry_after: float | int | None = None, +) -> float: + """Choose a retry delay tuned to the failure class. + + Overload errors (503/529) need a noticeably longer cool-down than generic + transport failures; otherwise Hermes can burn through retries before the + upstream cluster has time to recover. Retry-After, when present, wins. + """ + + normalized_reason = getattr(reason, "value", reason) or "" + + if normalized_reason == FailoverReason.overloaded.value: + max_delay = 180.0 + if retry_after is not None: + try: + return min(float(retry_after), max_delay) + except (TypeError, ValueError): + pass + return jittered_backoff(attempt, base_delay=8.0, max_delay=max_delay) + + if normalized_reason in (FailoverReason.rate_limit.value, FailoverReason.billing.value): + max_delay = 120.0 + if retry_after is not None: + try: + return min(float(retry_after), max_delay) + except (TypeError, ValueError): + pass + return jittered_backoff(attempt, base_delay=2.0, max_delay=60.0) + + return jittered_backoff(attempt, base_delay=2.0, max_delay=60.0) diff --git a/run_agent.py b/run_agent.py index 325df9beb1d6..d0a188e61810 100644 --- a/run_agent.py +++ b/run_agent.py @@ -76,7 +76,7 @@ # Agent internals extracted to agent/ package for modularity from agent.memory_manager import build_memory_context_block, sanitize_context -from agent.retry_utils import jittered_backoff +from agent.retry_utils import jittered_backoff, select_retry_wait_time from agent.error_classifier import classify_api_error, FailoverReason from agent.prompt_builder import ( DEFAULT_AGENT_IDENTITY, PLATFORM_HINTS, @@ -5075,6 +5075,34 @@ def _anthropic_messages_create(self, api_kwargs: dict): self._try_refresh_anthropic_client_credentials() return self._anthropic_client.messages.create(**api_kwargs) + def _estimate_request_context_tokens(self, api_kwargs: dict) -> int: + """Roughly estimate request size for stale-timeout heuristics/logging. + + Chat Completions requests carry context in ``messages``. + Codex/Responses API requests carry most of it in ``input`` plus + ``instructions``. Using only ``messages`` underestimates large Codex + requests as zero, which can trigger the default 300s stale timeout even + when the request is legitimately huge. + """ + + pieces = [] + for key in ("messages", "input", "instructions", "tools"): + value = api_kwargs.get(key) + if value is None: + continue + try: + pieces.append(json.dumps(value, ensure_ascii=False, separators=(",", ":"))) + except Exception: + pieces.append(str(value)) + + if not pieces: + try: + return len(json.dumps(api_kwargs, ensure_ascii=False, separators=(",", ":"))) // 4 + except Exception: + return len(str(api_kwargs)) // 4 + + return sum(len(piece) for piece in pieces) // 4 + def _interruptible_api_call(self, api_kwargs: dict): """ Run the API call in a background thread so the main conversation loop @@ -5135,7 +5163,7 @@ def _call(): if _stale_base == 300.0 and _base_url and is_local_endpoint(_base_url): _stale_timeout = float("inf") else: - _est_tokens = sum(len(str(v)) for v in api_kwargs.get("messages", [])) // 4 + _est_tokens = self._estimate_request_context_tokens(api_kwargs) if _est_tokens > 100_000: _stale_timeout = max(_stale_base, 600.0) elif _est_tokens > 50_000: @@ -5165,7 +5193,7 @@ def _call(): # arrives within the configured timeout. _elapsed = time.time() - _call_start if _elapsed > _stale_timeout: - _est_ctx = sum(len(str(v)) for v in api_kwargs.get("messages", [])) // 4 + _est_ctx = self._estimate_request_context_tokens(api_kwargs) logger.warning( "Non-streaming API call stale for %.0fs (threshold %.0fs). " "model=%s context=~%s tokens. Killing connection.", @@ -5850,7 +5878,7 @@ def _call(): # when the context is large. Without this, the stale detector kills # healthy connections during the model's thinking phase, producing # spurious RemoteProtocolError ("peer closed connection"). - _est_tokens = sum(len(str(v)) for v in api_kwargs.get("messages", [])) // 4 + _est_tokens = self._estimate_request_context_tokens(api_kwargs) if _est_tokens > 100_000: _stream_stale_timeout = max(_stream_stale_timeout_base, 300.0) elif _est_tokens > 50_000: @@ -5886,7 +5914,7 @@ def _call(): # inner retry loop can start a fresh connection. _stale_elapsed = time.time() - last_chunk_time["t"] if _stale_elapsed > _stream_stale_timeout: - _est_ctx = sum(len(str(v)) for v in api_kwargs.get("messages", [])) // 4 + _est_ctx = self._estimate_request_context_tokens(api_kwargs) logger.warning( "Stream stale for %.0fs (threshold %.0fs) — no chunks received. " "model=%s context=~%s tokens. Killing connection.", @@ -10284,20 +10312,30 @@ def _stop_spinner(): "error": _final_summary, } - # For rate limits, respect the Retry-After header if present + # For rate limits and overloads, respect Retry-After when present. _retry_after = None - if is_rate_limited: + if classified.reason in ( + FailoverReason.rate_limit, + FailoverReason.billing, + FailoverReason.overloaded, + ): _resp_headers = getattr(getattr(api_error, "response", None), "headers", None) if _resp_headers and hasattr(_resp_headers, "get"): _ra_raw = _resp_headers.get("retry-after") or _resp_headers.get("Retry-After") if _ra_raw: try: - _retry_after = min(int(_ra_raw), 120) # Cap at 2 minutes + _retry_after = float(_ra_raw) except (TypeError, ValueError): pass - wait_time = _retry_after if _retry_after else jittered_backoff(retry_count, base_delay=2.0, max_delay=60.0) + wait_time = select_retry_wait_time( + retry_count, + reason=classified.reason, + retry_after=_retry_after, + ) if is_rate_limited: self._emit_status(f"⏱️ Rate limit reached. Waiting {wait_time}s before retry (attempt {retry_count + 1}/{max_retries})...") + elif classified.reason == FailoverReason.overloaded: + self._emit_status(f"🧯 Provider overloaded. Waiting {wait_time}s before retry (attempt {retry_count}/{max_retries})...") else: self._emit_status(f"⏳ Retrying in {wait_time}s (attempt {retry_count}/{max_retries})...") logger.warning( diff --git a/tests/run_agent/test_run_agent.py b/tests/run_agent/test_run_agent.py index 46eec2cf71df..de4f175c5ed7 100644 --- a/tests/run_agent/test_run_agent.py +++ b/tests/run_agent/test_run_agent.py @@ -3532,6 +3532,28 @@ def test_anthropic_messages_create_preflights_refresh(self): assert result is response +class TestRequestContextTokenEstimate: + def test_uses_messages_payload_for_chat_completions(self, agent): + api_kwargs = { + "messages": [ + {"role": "system", "content": "x" * 400}, + {"role": "user", "content": "y" * 400}, + ] + } + + assert agent._estimate_request_context_tokens(api_kwargs) >= 200 + + def test_counts_codex_input_and_instructions(self, agent): + api_kwargs = { + "input": [{"role": "user", "content": [{"type": "input_text", "text": "z" * 410000}]}], + "instructions": "i" * 2000, + } + + # Regression: Codex/Responses requests do not use `messages`, so the + # old estimator returned 0 and kept the stale timeout at 300s. + assert agent._estimate_request_context_tokens(api_kwargs) > 100_000 + + # =================================================================== # _streaming_api_call tests # =================================================================== diff --git a/tests/test_retry_utils.py b/tests/test_retry_utils.py index f39c3142d9fd..2a35a33a1fc6 100644 --- a/tests/test_retry_utils.py +++ b/tests/test_retry_utils.py @@ -3,7 +3,8 @@ import threading import agent.retry_utils as retry_utils -from agent.retry_utils import jittered_backoff +from agent.error_classifier import FailoverReason +from agent.retry_utils import jittered_backoff, select_retry_wait_time def test_backoff_is_exponential(): @@ -115,3 +116,22 @@ def _call(): assert len(recorded_seeds) == 2 assert len(set(recorded_seeds)) == 2, f"Expected unique seeds, got {recorded_seeds}" + + +def test_select_retry_wait_time_uses_longer_backoff_for_overloaded(monkeypatch): + monkeypatch.setattr(retry_utils, "jittered_backoff", lambda attempt, **kwargs: kwargs["base_delay"]) + + overloaded = select_retry_wait_time(1, reason=FailoverReason.overloaded) + generic = select_retry_wait_time(1, reason=FailoverReason.timeout) + + assert overloaded > generic + assert overloaded == 8.0 + assert generic == 2.0 + + +def test_select_retry_wait_time_prefers_retry_after_for_overloaded(): + assert select_retry_wait_time(2, reason=FailoverReason.overloaded, retry_after=95) == 95 + + +def test_select_retry_wait_time_caps_retry_after_for_overloaded(): + assert select_retry_wait_time(2, reason=FailoverReason.overloaded, retry_after=999) == 180.0