Import AITURK IDE 1.0.0-beta.1 from Hermes 63279301; preserve MIT license

tests/test_retry_utils.py

@@ -0,0 +1,210 @@
+"""Tests for agent.retry_utils jittered backoff."""
+
+import threading
+
+import agent.retry_utils as retry_utils
+from types import SimpleNamespace
+
+from agent.retry_utils import adaptive_rate_limit_backoff, is_zai_coding_overload_error, jittered_backoff
+
+
+def test_backoff_is_exponential():
+    """Base delay should double each attempt (before jitter)."""
+    for attempt in (1, 2, 3, 4):
+        delays = [jittered_backoff(attempt, base_delay=5.0, max_delay=120.0, jitter_ratio=0.0) for _ in range(100)]
+        expected = min(5.0 * (2 ** (attempt - 1)), 120.0)
+        mean = sum(delays) / len(delays)
+        assert abs(mean - expected) < 0.01, f"attempt {attempt}: expected {expected}, got {mean}"
+
+
+def test_backoff_respects_max_delay():
+    """Even with high attempt numbers, delay should not exceed max_delay."""
+    for attempt in (10, 20, 100):
+        delay = jittered_backoff(attempt, base_delay=5.0, max_delay=60.0, jitter_ratio=0.0)
+        assert delay <= 60.0, f"attempt {attempt}: delay {delay} exceeds max 60s"
+
+
+
+
+def test_backoff_attempt_1_is_base():
+    """First attempt delay should equal base_delay (with no jitter)."""
+    delay = jittered_backoff(1, base_delay=3.0, max_delay=120.0, jitter_ratio=0.0)
+    assert delay == 3.0
+
+
+
+
+
+
+
+
+def test_backoff_thread_safety():
+    """Concurrent calls should generally produce different delays."""
+    results = []
+    barrier = threading.Barrier(8)
+
+    def _call_backoff():
+        barrier.wait()
+        results.append(jittered_backoff(1, base_delay=10.0, max_delay=120.0, jitter_ratio=0.5))
+
+    threads = [threading.Thread(target=_call_backoff) for _ in range(8)]
+    for t in threads:
+        t.start()
+    for t in threads:
+        t.join(timeout=5)
+
+    assert len(results) == 8
+    unique = len(set(results))
+    assert unique >= 6, f"Expected mostly unique delays, got {unique}/8 unique"
+
+
+def test_backoff_uses_locked_tick_for_seed(monkeypatch):
+    """Seed derivation should use per-call tick captured under lock."""
+    import time
+
+    monkeypatch.setattr(retry_utils, "_jitter_counter", 0)
+
+    recorded_seeds = []
+
+    class _RecordingRandom:
+        def __init__(self, seed):
+            recorded_seeds.append(seed)
+
+        def uniform(self, a, b):
+            return 0.0
+
+    monkeypatch.setattr(retry_utils.random, "Random", _RecordingRandom)
+
+    fixed_time_ns = 123456789
+
+    def _time_ns_wait_for_two_ticks():
+        deadline = time.time() + 2.0
+        while retry_utils._jitter_counter < 2 and time.time() < deadline:
+            time.sleep(0.001)
+        return fixed_time_ns
+
+    monkeypatch.setattr(retry_utils.time, "time_ns", _time_ns_wait_for_two_ticks)
+
+    barrier = threading.Barrier(2)
+
+    def _call():
+        barrier.wait()
+        jittered_backoff(1, base_delay=10.0, max_delay=120.0, jitter_ratio=0.5)
+
+    threads = [threading.Thread(target=_call) for _ in range(2)]
+    for t in threads:
+        t.start()
+    for t in threads:
+        t.join(timeout=5)
+
+    assert len(recorded_seeds) == 2
+    assert len(set(recorded_seeds)) == 2, f"Expected unique seeds, got {recorded_seeds}"
+
+
+def _zai_overload_error():
+    return SimpleNamespace(
+        status_code=429,
+        body={
+            "error": {
+                "code": "1305",
+                "message": "The service may be temporarily overloaded, please try again later",
+            }
+        },
+    )
+
+
+
+
+
+
+
+
+
+
+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]
+
+
+# ---------------------------------------------------------------------------
+# parse_retry_after_seconds — shared Retry-After parser
+# ---------------------------------------------------------------------------
+
+
+class TestParseRetryAfterSeconds:
+    def test_numeric_string(self):
+        from agent.retry_utils import parse_retry_after_seconds
+        assert parse_retry_after_seconds("120") == 120.0
+        assert parse_retry_after_seconds(" 4.5 ") == 4.5
+
+    def test_numeric_value(self):
+        from agent.retry_utils import parse_retry_after_seconds
+        assert parse_retry_after_seconds(45) == 45.0
+        assert parse_retry_after_seconds(3.25) == 3.25
+
+
+    def test_http_date(self):
+        from datetime import datetime, timedelta, timezone
+        from email.utils import format_datetime
+        from agent.retry_utils import parse_retry_after_seconds
+
+        future = datetime.now(timezone.utc) + timedelta(seconds=90)
+        seconds = parse_retry_after_seconds(format_datetime(future, usegmt=True))
+        assert seconds is not None and 80 <= seconds <= 91
+
+        past = datetime.now(timezone.utc) - timedelta(seconds=90)
+        assert parse_retry_after_seconds(format_datetime(past, usegmt=True)) == 0.0
+
+
+
+    def test_headers_get_raises(self):
+        from agent.retry_utils import parse_retry_after_seconds
+
+        class Explosive:
+            def get(self, _key):
+                raise RuntimeError("boom")
+
+        assert parse_retry_after_seconds(Explosive()) is None