fix: refuse overflowing server retry delays

Author: markstuart-oai

src/openai/_base_client.py

@@ -759,24 +759,24 @@ def _parse_retry_after_header(self, response_headers: Optional[httpx2.Headers] =
         if response_headers is None:
             return None
 
-        # First, try the non-standard `retry-after-ms` header for milliseconds,
-        # which is more precise than integer-seconds `retry-after`
-        try:
-            retry_ms_header = response_headers.get("retry-after-ms", None)
-            return float(retry_ms_header) / 1000
-        except (TypeError, ValueError):
-            pass
-
-        # Next, try parsing `retry-after` header as seconds (allowing nonstandard floats).
-        retry_header = response_headers.get("retry-after")
-        try:
-            # note: the spec indicates that this should only ever be an integer
-            # but if someone sends a float there's no reason for us to not respect it
-            return float(retry_header)
-        except (TypeError, ValueError):
-            pass
+        # Prefer milliseconds, then seconds (allowing nonstandard floats).
+        for header, divisor in (("retry-after-ms", 1000), ("retry-after", 1)):
+            value = response_headers.get(header)
+            if value is None:
+                continue
+            try:
+                delay = float(value)
+            except ValueError:
+                continue
+            if delay == math.inf and value.strip().lower() not in ("inf", "+inf", "infinity", "+infinity"):
+                # Numeric overflow is an excessive server delay, not a malformed
+                # infinity literal. Keep a finite sentinel so retry eligibility
+                # refuses it instead of falling back to a shorter wait.
+                return MAX_RETRY_AFTER_DELAY + 1
+            return delay / divisor
 
         # Last, try parsing `retry-after` as a date.
+        retry_header = response_headers.get("retry-after")
         try:
             retry_date_tuple = email.utils.parsedate_tz(retry_header)
             if retry_date_tuple is None:

tests/test_retry_after_overflow.py

@@ -0,0 +1,65 @@
+from __future__ import annotations
+
+from unittest import mock
+
+import httpx2
+import pytest
+
+from openai import OpenAI, AsyncOpenAI, APIStatusError
+
+
+@pytest.mark.parametrize("async_mode", [False, True])
+@pytest.mark.parametrize(
+    "status,headers,delay",
+    [
+        (429, {"retry-after": "1e999"}, None),
+        (503, {"retry-after": "9" * 400}, None),
+        (429, {"retry-after-ms": "1e999"}, None),
+        (503, {"retry-after-ms": "9" * 400}, None),
+        (401, {"retry-after": "1e999", "x-should-retry": "true"}, None),
+        (429, {"retry-after": "inf"}, 0.5),
+        (503, {"retry-after": " +Infinity "}, 0.5),
+        (429, {"retry-after": "NaN"}, 0.5),
+        (429, {"retry-after": "-1e999"}, 0.5),
+        (429, {"retry-after-ms": "inf", "retry-after": "90"}, 0.5),
+        (429, {"retry-after": "1e2"}, 100.0),
+        (503, {"retry-after-ms": "1e5"}, 100.0),
+    ],
+)
+async def test_retry_after_numeric_overflow(
+    async_mode: bool, status: int, headers: dict[str, str], delay: float | None
+) -> None:
+    attempts = 0
+    body = {"message": "Synthetic retry error", "type": "synthetic_error", "code": "synthetic_code"}
+
+    def handle(request: httpx2.Request) -> httpx2.Response:
+        nonlocal attempts
+        attempts += 1
+        assert request.url.path == "/models/test"
+        return httpx2.Response(status, headers={**headers, "x-request-id": "synthetic-id"}, json={"error": body})
+
+    with mock.patch("time.sleep") as sync_sleep, mock.patch("anyio.sleep") as async_sleep:
+        with mock.patch("openai._base_client.random", return_value=0), pytest.raises(APIStatusError) as exc:
+            if async_mode:
+                async with AsyncOpenAI(
+                    api_key="synthetic-key",
+                    base_url="https://retry.test",
+                    max_retries=1,
+                    http_client=httpx2.AsyncClient(transport=httpx2.MockTransport(handle), trust_env=False),
+                ) as async_client:
+                    await async_client.models.retrieve("test")
+            else:
+                with OpenAI(
+                    api_key="synthetic-key",
+                    base_url="https://retry.test",
+                    max_retries=1,
+                    http_client=httpx2.Client(transport=httpx2.MockTransport(handle), trust_env=False),
+                ) as client:
+                    client.models.retrieve("test")
+
+    assert attempts == (1 if delay is None else 2)
+    assert (async_sleep if async_mode else sync_sleep).call_args_list == ([] if delay is None else [mock.call(delay)])
+    assert exc.value.status_code == status
+    assert exc.value.body == body
+    assert exc.value.request_id == "synthetic-id"
+    assert all(exc.value.response.headers[key] == value for key, value in headers.items())

tests/test_x509_workload_identity_hardening.py

@@ -322,7 +322,15 @@ def test_x509_rejects_data_residency_without_a_confirmed_mtls_endpoint(
         client.with_options(data_residency="ae")
 
 
-@pytest.mark.parametrize("headers", [{"x-should-retry": "false"}, {"retry-after-ms": "120001"}])
+@pytest.mark.parametrize(
+    "headers",
+    [
+        {"x-should-retry": "false"},
+        {"retry-after-ms": "120001"},
+        {"retry-after": "1e999"},
+        {"retry-after-ms": "9" * 400},
+    ],
+)
 def test_sync_x509_token_exchange_honors_server_retry_refusals(headers: dict[str, str]) -> None:
     requests: list[httpx2.Request] = []
 
@@ -339,7 +347,15 @@ def handler(request: httpx2.Request) -> httpx2.Response:
     assert len(requests) == 1
 
 
-@pytest.mark.parametrize("headers", [{"x-should-retry": "false"}, {"retry-after-ms": "120001"}])
+@pytest.mark.parametrize(
+    "headers",
+    [
+        {"x-should-retry": "false"},
+        {"retry-after-ms": "120001"},
+        {"retry-after": "1e999"},
+        {"retry-after-ms": "9" * 400},
+    ],
+)
 async def test_async_x509_token_exchange_honors_server_retry_refusals(headers: dict[str, str]) -> None:
     requests: list[httpx2.Request] = []