Merge pull request #60034 from kshitijk4poor/salvage/59523-zai-overload-backoff

fix(agent): run Z.AI Coding overload adaptive backoff on the overloaded path

agent/conversation_loop.py

@@ -58,7 +58,12 @@ from agent.model_metadata import (
 )
 from agent.process_bootstrap import _install_safe_stdio
 from agent.prompt_caching import apply_anthropic_cache_control
-from agent.retry_utils import adaptive_rate_limit_backoff, jittered_backoff
+from agent.retry_utils import (
+    adaptive_rate_limit_backoff,
+    is_zai_coding_overload_error,
+    jittered_backoff,
+    zai_coding_overload_retry_ceiling,
+)
 from agent.trajectory import has_incomplete_scratchpad
 from agent.usage_pricing import estimate_usage_cost, normalize_usage
 from hermes_constants import PARTIAL_STREAM_STUB_ID
@@ -3142,6 +3147,18 @@ def run_conversation(
                     FailoverReason.timeout,
                     FailoverReason.overloaded,
                 }
+                # 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
+                # Z.AI backoff below. Detect the overload directly so its
+                # long-backoff schedule runs, and raise the retry ceiling so the
+                # long tier (30/60/90/120s) is reachable. See
+                # zai_coding_overload_retry_ceiling() for the ceiling rationale.
+                _is_zai_coding_overload = is_zai_coding_overload_error(
+                    base_url=str(_base), model=_model, error=api_error
+                )
+                if _is_zai_coding_overload:
+                    max_retries = max(max_retries, zai_coding_overload_retry_ceiling())
                 _should_fallback = (
                     is_rate_limited
                     or (_is_transport_failure and retry_count >= 2)
@@ -4092,7 +4109,7 @@ def run_conversation(
                                 pass
                 wait_time = _retry_after if _retry_after else jittered_backoff(retry_count, base_delay=2.0, max_delay=60.0)
                 _backoff_policy = None
-                if is_rate_limited and not _retry_after:
+                if (is_rate_limited or _is_zai_coding_overload) and not _retry_after:
                     wait_time, _backoff_policy = adaptive_rate_limit_backoff(
                         retry_count,
                         base_url=str(_base),
@@ -4100,13 +4117,14 @@ def run_conversation(
                         error=api_error,
                         default_wait=wait_time,
                     )
-                if is_rate_limited:
+                if is_rate_limited or _is_zai_coding_overload:
                     _policy_note = ""
                     if _backoff_policy == "zai_coding_overload_long":
                         _policy_note = " (Z.AI Coding overload adaptive long backoff)"
                     elif _backoff_policy == "zai_coding_overload_short":
                         _policy_note = " (Z.AI Coding overload short retry)"
-                    _rate_limit_status = f"⏱️ Rate limited. Waiting {wait_time:.1f}s (attempt {retry_count + 1}/{max_retries}){_policy_note}..."
+                    _wait_reason = "Provider overloaded" if _is_zai_coding_overload and not is_rate_limited else "Rate limited"
+                    _rate_limit_status = f"⏱️ {_wait_reason}. Waiting {wait_time:.1f}s (attempt {retry_count + 1}/{max_retries}){_policy_note}..."
                     # Normal retries are buffered to avoid noisy transient chatter. Long
                     # Z.AI Coding waits are different: they can last minutes, so surface
                     # progress immediately instead of making the TUI look frozen.

agent/retry_utils.py

@@ -24,6 +24,14 @@ _jitter_lock = threading.Lock()
 # not sit silent for 20+ minutes.
 _ZAI_CODING_OVERLOAD_LONG_BACKOFF = (30.0, 60.0, 90.0, 120.0)
 
+# Number of initial short retries before the adaptive long-backoff tier kicks
+# in. Shared by ``adaptive_rate_limit_backoff`` (which walks the long table
+# starting at attempt ``short_attempts + 1``) and
+# ``zai_coding_overload_retry_ceiling`` (which sizes the retry loop so every
+# long-tier entry is reachable). Keeping it a single module constant prevents
+# the two from silently desyncing if the short-retry count is ever tuned.
+_ZAI_CODING_OVERLOAD_SHORT_ATTEMPTS = 3
+
 
 def jittered_backoff(
     attempt: int,
@@ -104,7 +112,7 @@ def adaptive_rate_limit_backoff(
     model: str | None,
     error: Any,
     default_wait: float,
-    short_attempts: int = 3,
+    short_attempts: int = _ZAI_CODING_OVERLOAD_SHORT_ATTEMPTS,
 ) -> tuple[float, str | None]:
     """Provider-aware rate-limit backoff.
 
@@ -127,3 +135,20 @@ def adaptive_rate_limit_backoff(
     # A smaller jitter ratio keeps long waits readable while still avoiding
     # synchronized retry storms across concurrent Hermes sessions.
     return jittered_backoff(1, base_delay=base_delay, max_delay=base_delay, jitter_ratio=0.2), "zai_coding_overload_long"
+
+
+def zai_coding_overload_retry_ceiling(short_attempts: int = _ZAI_CODING_OVERLOAD_SHORT_ATTEMPTS) -> int:
+    """Retry-loop ceiling needed for the full Z.AI overload backoff schedule.
+
+    The adaptive policy runs ``short_attempts`` short retries, then walks the
+    long-backoff table one entry per subsequent attempt. The retry loop gives
+    up as soon as ``retry_count >= ceiling`` — and that check runs *before* the
+    attempt's backoff is computed — so the ceiling must sit one past the final
+    long-backoff entry for every long tier to actually execute.
+
+    With the default ``api_max_retries`` (3) equal to ``short_attempts`` (3),
+    the loop always gave up before reaching the long tier, leaving the whole
+    long-backoff schedule as dead code. Callers extend the ceiling to this
+    value for Z.AI Coding overload 429s so the 30/60/90/120s waits run.
+    """
+    return short_attempts + len(_ZAI_CODING_OVERLOAD_LONG_BACKOFF) + 1

tests/test_retry_utils.py

@@ -210,3 +210,51 @@ def test_non_zai_backoff_returns_default_wait():
     )
     assert wait == 12.0
     assert policy is None
+
+
+def test_zai_overload_retry_ceiling_exceeds_short_attempts():
+    """Invariant: the ceiling must sit above the short-retry threshold, or the
+    long-backoff tier is unreachable and the whole schedule is dead code
+    (the original bug: default api_max_retries == short_attempts == 3)."""
+    from agent.retry_utils import (
+        zai_coding_overload_retry_ceiling,
+        _ZAI_CODING_OVERLOAD_LONG_BACKOFF,
+    )
+
+    short_attempts = 3
+    ceiling = zai_coding_overload_retry_ceiling(short_attempts)
+    assert ceiling > short_attempts
+    # Invariant (not a formula mirror): the loop's give-up check
+    # (retry_count >= ceiling) runs *before* the attempt's backoff, so the
+    # ceiling must leave headroom for every long-backoff entry to execute —
+    # i.e. the largest attempt the loop still computes backoff for
+    # (ceiling - 1) must reach the final long-tier index.
+    last_attempt_with_backoff = ceiling - 1
+    assert last_attempt_with_backoff - short_attempts >= len(_ZAI_CODING_OVERLOAD_LONG_BACKOFF)
+
+
+def test_zai_overload_ceiling_makes_long_tier_reachable(monkeypatch):
+    """End-to-end over the attempt range the retry loop actually walks: with the
+    extended ceiling, at least one attempt reaches the long-backoff tier and the
+    full 30/60/90/120s schedule is exercised."""
+    monkeypatch.setattr(retry_utils, "jittered_backoff", lambda *a, **kw: kw["base_delay"])
+    from agent.retry_utils import zai_coding_overload_retry_ceiling
+
+    err = _zai_overload_error()
+    ceiling = zai_coding_overload_retry_ceiling()
+
+    long_waits = []
+    # The loop computes backoff for attempts 1..ceiling-1 (it gives up at ceiling).
+    for attempt in range(1, ceiling):
+        _wait, policy = adaptive_rate_limit_backoff(
+            attempt,
+            base_url="https://api.z.ai/api/coding/paas/v4",
+            model="glm-5.2",
+            error=err,
+            default_wait=1.0,
+        )
+        if policy == "zai_coding_overload_long":
+            long_waits.append(_wait)
+
+    assert long_waits, "long-backoff tier never reached within the retry ceiling"
+    assert long_waits == [30.0, 60.0, 90.0, 120.0]