diff --git a/src/workos/_base_client.py b/src/workos/_base_client.py index 7c6f7e64..15c5f4c1 100644 --- a/src/workos/_base_client.py +++ b/src/workos/_base_client.py @@ -3,6 +3,7 @@ import asyncio import json +import math import os import platform import time @@ -49,6 +50,10 @@ INITIAL_RETRY_DELAY = 1 MAX_RETRY_DELAY = 30 RETRY_MULTIPLIER = 2 +# Upper bound on a server-provided Retry-After sleep, in seconds. Matches the +# Node SDK so an aggressive proxy or a far-future HTTP-date cannot stall the +# caller. The raw header value is still reported on RateLimitExceededError. +MAX_RETRY_AFTER = 60 # Segments that URL normalization treats structurally; never valid identifiers. @@ -129,16 +134,24 @@ def build_url( @staticmethod def _parse_retry_after(retry_after: Optional[str]) -> Optional[float]: - """Parse Retry-After as seconds or an HTTP-date.""" + """Parse Retry-After as RFC 9110 delay-seconds or an HTTP-date. + + Returns ``None`` for absent or unparseable values so callers fall back + to exponential backoff. Delay-seconds must be a plain non-negative + integer, which rejects forms ``float()`` would otherwise accept such as + ``NaN``, ``Infinity`` or ``1e309``. The result is not capped here so + :class:`RateLimitExceededError` reports what the server actually sent; + :meth:`_calculate_retry_delay` bounds the sleep. + """ if not retry_after: return None value = retry_after.strip() if not value: return None - try: - return max(float(value), 0.0) - except ValueError: - pass + if value.isdecimal(): + seconds = float(value) + # A digit string long enough to overflow float parses as inf. + return seconds if math.isfinite(seconds) else None try: retry_at = parsedate_to_datetime(value) except (TypeError, ValueError, IndexError, OverflowError): @@ -151,10 +164,15 @@ def _parse_retry_after(retry_after: Optional[str]) -> Optional[float]: def _calculate_retry_delay( attempt: int, retry_after: Optional[str] = None ) -> float: - """Calculate retry delay with exponential backoff and jitter.""" + """Calculate retry delay with exponential backoff and jitter. + + A valid Retry-After header wins over backoff but is capped at + ``MAX_RETRY_AFTER`` so a huge value cannot overflow ``time.sleep`` or + leave an async retry waiting indefinitely. + """ parsed_retry_after = _BaseWorkOSClient._parse_retry_after(retry_after) if parsed_retry_after is not None: - return parsed_retry_after + return min(parsed_retry_after, float(MAX_RETRY_AFTER)) delay = min(INITIAL_RETRY_DELAY * (RETRY_MULTIPLIER**attempt), MAX_RETRY_DELAY) return delay * (0.5 + random.random()) diff --git a/tests/test_generated_client.py b/tests/test_generated_client.py index 3dcdd66f..09cd20d6 100644 --- a/tests/test_generated_client.py +++ b/tests/test_generated_client.py @@ -2,6 +2,9 @@ """Client tests: retries, errors, context manager, idempotency.""" +from datetime import datetime, timedelta, timezone +from email.utils import format_datetime + import httpx2 as httpx import pytest @@ -30,6 +33,27 @@ ServerError, ) +# Retry-After values that must never reach time.sleep/asyncio.sleep unbounded, +# paired with the sleep range the client must use instead: non-finite or +# non-integer forms fall back to the first backoff step (INITIAL_RETRY_DELAY +# with +/-50% jitter), huge finite values and far-future dates are capped at +# MAX_RETRY_AFTER. +_BACKOFF_FIRST_STEP = ( + generated_client_module.INITIAL_RETRY_DELAY * 0.5, + generated_client_module.INITIAL_RETRY_DELAY * 1.5, +) +_CAPPED = ( + float(generated_client_module.MAX_RETRY_AFTER), + float(generated_client_module.MAX_RETRY_AFTER), +) +UNSAFE_RETRY_AFTER_CASES = [ + ("NaN", _BACKOFF_FIRST_STEP), + ("Infinity", _BACKOFF_FIRST_STEP), + ("1e309", _BACKOFF_FIRST_STEP), + ("99999999999", _CAPPED), + ("Fri, 31 Dec 9999 23:59:59 GMT", _CAPPED), +] + class TestWorkOSClient: def test_missing_credentials_raise(self): @@ -381,6 +405,54 @@ def test_empty_body_sends_json(self, httpx_mock): def test_calculate_retry_delay_uses_retry_after_seconds(self): assert WorkOSClient._calculate_retry_delay(1, "30") == 30.0 + @pytest.mark.parametrize( + "retry_after", ["3600", "99999999999", "Fri, 31 Dec 9999 23:59:59 GMT"] + ) + def test_calculate_retry_delay_caps_retry_after(self, retry_after): + assert WorkOSClient._calculate_retry_delay(0, retry_after) == float( + generated_client_module.MAX_RETRY_AFTER + ) + + @pytest.mark.parametrize( + "retry_after", + ["NaN", "Infinity", "1e309", "-1", "1.5", "0x10", "1" * 400, "not a date"], + ) + def test_parse_retry_after_rejects_non_delay_seconds(self, retry_after): + assert WorkOSClient._parse_retry_after(retry_after) is None + + def test_parse_retry_after_keeps_raw_server_value(self): + assert WorkOSClient._parse_retry_after("3600") == 3600.0 + + def test_parse_retry_after_http_date(self): + retry_at = datetime.now(timezone.utc) + timedelta(seconds=30) + parsed = WorkOSClient._parse_retry_after(format_datetime(retry_at, usegmt=True)) + assert parsed is not None + assert 28.0 <= parsed <= 30.0 + + def test_parse_retry_after_past_http_date_is_zero(self): + assert WorkOSClient._parse_retry_after("Thu, 01 Jan 1970 00:00:00 GMT") == 0.0 + + @pytest.mark.parametrize( + ("retry_after", "expected_range"), UNSAFE_RETRY_AFTER_CASES + ) + def test_retry_sleep_is_finite_and_bounded( + self, httpx_mock, monkeypatch, retry_after, expected_range + ): + delays: list[float] = [] + monkeypatch.setattr(generated_client_module.time, "sleep", delays.append) + httpx_mock.add_response( + status_code=503, + headers={"Retry-After": retry_after}, + json={"message": "down"}, + ) + httpx_mock.add_response(json={"ok": True}) + client = WorkOSClient(api_key="sk_test_123", client_id="client_test") + assert client.request("GET", ("test",)) == {"ok": True} + client.close() + assert len(delays) == 1 + low, high = expected_range + assert low <= delays[0] <= high + def test_retry_exhaustion_raises_rate_limit(self, httpx_mock, monkeypatch): monkeypatch.setattr(generated_client_module.time, "sleep", lambda _: None) for _ in range(4): @@ -410,6 +482,37 @@ def test_rate_limit_retry_after_is_parsed(self, httpx_mock): assert exc_info.value.retry_after == 30.0 client.close() + def test_rate_limit_retry_after_is_not_capped_on_error(self, httpx_mock): + httpx_mock.add_response( + status_code=429, + headers={"Retry-After": "3600"}, + json={"message": "Slow down"}, + ) + client = WorkOSClient( + api_key="sk_test_123", client_id="client_test", max_retries=0 + ) + with pytest.raises(RateLimitExceededError) as exc_info: + client.request("GET", ("test",)) + assert exc_info.value.retry_after == 3600.0 + client.close() + + @pytest.mark.parametrize("retry_after", ["NaN", "Infinity", "1e309"]) + def test_rate_limit_retry_after_is_none_for_invalid_header( + self, httpx_mock, retry_after + ): + httpx_mock.add_response( + status_code=429, + headers={"Retry-After": retry_after}, + json={"message": "Slow down"}, + ) + client = WorkOSClient( + api_key="sk_test_123", client_id="client_test", max_retries=0 + ) + with pytest.raises(RateLimitExceededError) as exc_info: + client.request("GET", ("test",)) + assert exc_info.value.retry_after is None + client.close() + def test_timeout_error_is_wrapped(self, httpx_mock, monkeypatch): monkeypatch.setattr(generated_client_module.time, "sleep", lambda _: None) httpx_mock.add_exception(httpx.TimeoutException("timed out")) @@ -705,6 +808,31 @@ async def test_raises_429(self, httpx_mock): await client.request("GET", ("test",)) await client.close() + @pytest.mark.parametrize( + ("retry_after", "expected_range"), UNSAFE_RETRY_AFTER_CASES + ) + async def test_retry_sleep_is_finite_and_bounded( + self, httpx_mock, monkeypatch, retry_after, expected_range + ): + delays: list[float] = [] + + async def _sleep(delay: float) -> None: + delays.append(delay) + + monkeypatch.setattr(generated_client_module.asyncio, "sleep", _sleep) + httpx_mock.add_response( + status_code=503, + headers={"Retry-After": retry_after}, + json={"message": "down"}, + ) + httpx_mock.add_response(json={"ok": True}) + client = AsyncWorkOSClient(api_key="sk_test_123", client_id="client_test") + assert await client.request("GET", ("test",)) == {"ok": True} + await client.close() + assert len(delays) == 1 + low, high = expected_range + assert low <= delays[0] <= high + async def test_raises_500(self, httpx_mock): httpx_mock.add_response( status_code=500,