fix(transport): Retry-After honored for 503/504/408 + jitter applied to Retry-After + success-body size cap

kalshi/_base_client.py

@@ -46,6 +46,15 @@
 MAX_ERROR_BODY_BYTES = 16 * 1024
 MAX_ERROR_MESSAGE_CHARS = 1024
 
+# #323: parallel cap on the success path. ``response.json()`` (and any
+# caller-supplied ``rest_json_loads``) materialises the full body into a
+# Python object graph that typically inflates 5-10x over wire size. A
+# compromised reverse proxy, a hijacked DNS entry, or a backend regression
+# streaming an oversize list-orders payload could OOM a serverless function
+# before the parse completes. 64 MiB is comfortably above the largest
+# legitimate Kalshi list endpoint response while still bounding memory.
+MAX_RESPONSE_BODY_BYTES = 64 * 1024 * 1024
+
 
 def _map_error(response: httpx.Response) -> KalshiError:
     """Map an HTTP error response to the appropriate SDK exception."""
@@ -90,59 +99,62 @@ def _map_error(response: httpx.Response) -> KalshiError:
         return KalshiNotFoundError(message=message, status_code=status)
     if status == 409:
         return KalshiConflictError(message=message, status_code=status)
-    if status == 429:
-        retry_after = response.headers.get("Retry-After")
-        retry_after_val: float | None = None
-        if retry_after:
-            try:
-                retry_after_val = float(retry_after)
-                # Reject non-finite (NaN/inf) which would survive the
-                # transport caller's min() cap. Negative deltas — a server
-                # signalling "retry now" or a clock-skewed proxy emitting a
-                # past timestamp — clamp to 0.0 so both Retry-After forms
-                # behave the same (#267): a negative delta-seconds value and
-                # a past HTTP-date below now both produce ``retry_after=0.0``
-                # rather than diverging into None vs 0.0.
-                if not math.isfinite(retry_after_val):
-                    retry_after_val = None
-                elif retry_after_val < 0:
-                    retry_after_val = 0.0
-            except ValueError:
-                # RFC 7231 §7.1.3: Retry-After may also be an HTTP-date.
-                try:
-                    dt = email.utils.parsedate_to_datetime(retry_after)
-                except (TypeError, ValueError):
-                    dt = None
-                if dt is None:
-                    logger.debug(
-                        "Retry-After %r is neither delta-seconds nor HTTP-date; "
-                        "falling back to computed backoff",
-                        retry_after,
-                    )
-                else:
-                    # RFC 5322 dates without a tz are interpreted as UTC by
-                    # parsedate_to_datetime; ensure aware before subtracting.
-                    if dt.tzinfo is None:
-                        dt = dt.replace(tzinfo=UTC)
-                    delta = (dt - datetime.now(tz=UTC)).total_seconds()
-                    # Clamp negatives (date in the past) to 0 — retry immediately.
-                    # The transport caller caps at config.retry_max_delay.
-                    retry_after_val = max(0.0, delta)
-        return KalshiRateLimitError(
+    if status in (408, 429, 504) or status >= 500:
+        retry_after_val = _parse_retry_after(response)
+        if status == 429:
+            return KalshiRateLimitError(
+                message=message, status_code=status, retry_after=retry_after_val
+            )
+        # #251: 408 Request Timeout and 504 Gateway Timeout carry the same
+        # "may have committed" semantic as a transport-level timeout. Route
+        # them to KalshiTimeoutError so callers can branch on it (e.g.,
+        # reconcile via client_order_id before retrying an order create).
+        if status in (408, 504):
+            return KalshiTimeoutError(
+                message=message, status_code=status, retry_after=retry_after_val
+            )
+        return KalshiServerError(
             message=message, status_code=status, retry_after=retry_after_val
         )
-    # #251: 408 Request Timeout and 504 Gateway Timeout carry the same
-    # "may have committed" semantic as a transport-level timeout. Route them
-    # to KalshiTimeoutError so callers can branch on it (e.g., reconcile via
-    # client_order_id before retrying an order create).
-    if status in (408, 504):
-        return KalshiTimeoutError(message=message, status_code=status)
-    if status >= 500:
-        return KalshiServerError(message=message, status_code=status)
 
     return KalshiError(message=message, status_code=status)
 
 
+def _parse_retry_after(response: httpx.Response) -> float | None:
+    """Parse RFC 7231 §7.1.3 Retry-After to seconds; ``None`` if absent/unparseable (#267, #322)."""
+    raw = response.headers.get("Retry-After")
+    if not raw:
+        return None
+    try:
+        val = float(raw)
+        if not math.isfinite(val):
+            return None
+        if val < 0:
+            return 0.0
+        return val
+    except ValueError:
+        pass
+    # RFC 7231 §7.1.3: Retry-After may also be an HTTP-date.
+    try:
+        dt = email.utils.parsedate_to_datetime(raw)
+    except (TypeError, ValueError):
+        dt = None
+    if dt is None:
+        logger.debug(
+            "Retry-After %r is neither delta-seconds nor HTTP-date; "
+            "falling back to computed backoff",
+            raw,
+        )
+        return None
+    # RFC 5322 dates without a tz are interpreted as UTC by
+    # parsedate_to_datetime; ensure aware before subtracting.
+    if dt.tzinfo is None:
+        dt = dt.replace(tzinfo=UTC)
+    delta = (dt - datetime.now(tz=UTC)).total_seconds()
+    # Clamp negatives (date in the past) to 0 — retry immediately.
+    return max(0.0, delta)
+
+
 def _compute_backoff(attempt: int, config: KalshiConfig) -> float:
     """Exponential backoff with AWS "Full Jitter".
 
@@ -387,10 +399,21 @@ def request(
             if not should_retry:
                 raise error
 
-            # Use Retry-After header if available for 429.
-            # `is not None` — not truthy — so Retry-After: 0 ("retry immediately") is honored.
-            if isinstance(error, KalshiRateLimitError) and error.retry_after is not None:
-                delay = min(error.retry_after, self._config.retry_max_delay)
+            # #322: Retry-After is honored for any retryable status that
+            # advertises it (429 + 408/503/504/other 5xx). #321: apply Full
+            # Jitter on top of the server's hint — a synchronized client
+            # fleet hitting the same Retry-After would otherwise re-stampede
+            # the rate limit on the same tick. Floor is the server hint
+            # (clamped at retry_max_delay); ceiling is retry_max_delay.
+            # ``is not None`` — not truthy — so Retry-After: 0 ("retry
+            # immediately") still adds jitter on top instead of falling
+            # through to plain backoff.
+            if error.retry_after is not None:
+                floor = min(error.retry_after, self._config.retry_max_delay)
+                delay = min(
+                    floor + _compute_backoff(attempt, self._config),
+                    self._config.retry_max_delay,
+                )
             else:
                 delay = _compute_backoff(attempt, self._config)
 
@@ -598,9 +621,14 @@ async def request(
             if not should_retry:
                 raise error
 
-            # `is not None` so Retry-After: 0 ("retry immediately") is honored.
-            if isinstance(error, KalshiRateLimitError) and error.retry_after is not None:
-                delay = min(error.retry_after, self._config.retry_max_delay)
+            # #321/#322: jittered Retry-After on any retryable status that
+            # carries the header. See SyncTransport for the rationale.
+            if error.retry_after is not None:
+                floor = min(error.retry_after, self._config.retry_max_delay)
+                delay = min(
+                    floor + _compute_backoff(attempt, self._config),
+                    self._config.retry_max_delay,
+                )
             else:
                 delay = _compute_backoff(attempt, self._config)
 

kalshi/errors.py

@@ -6,10 +6,23 @@
 
 
 class KalshiError(Exception):
-    """Base exception for all Kalshi SDK errors."""
+    """Base exception for all Kalshi SDK errors.
 
-    def __init__(self, message: str, status_code: int | None = None) -> None:
+    ``retry_after`` carries the server's RFC 7231 §7.1.3 Retry-After hint
+    (parsed to seconds) when present. Populated for 408/429/503/504 and
+    other retryable 5xx responses; ``None`` otherwise. The transport
+    retry loop reads it via plain attribute access regardless of error
+    subclass (#322).
+    """
+
+    def __init__(
+        self,
+        message: str,
+        status_code: int | None = None,
+        retry_after: float | None = None,
+    ) -> None:
         self.status_code = status_code
+        self.retry_after = retry_after
         super().__init__(message)
 
 
@@ -52,15 +65,6 @@ def __init__(
 class KalshiRateLimitError(KalshiError):
     """Rate limit exceeded (429). Check retry_after for backoff hint."""
 
-    def __init__(
-        self,
-        message: str,
-        status_code: int | None = None,
-        retry_after: float | None = None,
-    ) -> None:
-        self.retry_after = retry_after
-        super().__init__(message, status_code)
-
 
 class KalshiConflictError(KalshiError):
     """409 Conflict (e.g., duplicate ``client_order_id``)."""

kalshi/resources/_base.py

@@ -9,10 +9,32 @@
 import httpx
 from pydantic import BaseModel
 
-from kalshi._base_client import AsyncTransport, SyncTransport
+from kalshi._base_client import MAX_RESPONSE_BODY_BYTES, AsyncTransport, SyncTransport
 from kalshi.errors import AuthRequiredError, KalshiError
 from kalshi.models.common import Page
 
+
+def _enforce_response_body_cap(response: httpx.Response) -> None:
+    """Raise ``KalshiError`` when a 2xx body exceeds ``MAX_RESPONSE_BODY_BYTES`` (#323)."""
+    content_length = response.headers.get("content-length")
+    if content_length:
+        try:
+            if int(content_length) > MAX_RESPONSE_BODY_BYTES:
+                raise KalshiError(
+                    f"Response body advertises {content_length} bytes, exceeds "
+                    f"max_response_bytes={MAX_RESPONSE_BODY_BYTES}",
+                )
+        except ValueError:
+            # Malformed Content-Length — fall through to the post-buffer check.
+            pass
+    body_len = len(response.content)
+    if body_len > MAX_RESPONSE_BODY_BYTES:
+        raise KalshiError(
+            f"Response body is {body_len} bytes, exceeds "
+            f"max_response_bytes={MAX_RESPONSE_BODY_BYTES}",
+        )
+
+
 T = TypeVar("T", bound=BaseModel)
 
 
@@ -145,6 +167,7 @@ def _load_json(self, response: httpx.Response) -> Any:
         Falls back to ``response.json()`` (stdlib) when no custom loader is
         configured. Custom loaders receive ``response.content`` (bytes).
         """
+        _enforce_response_body_cap(response)
         loader = self._transport._config.rest_json_loads
         if loader is None:
             return response.json()
@@ -391,6 +414,7 @@ def _load_json(self, response: httpx.Response) -> Any:
         Falls back to ``response.json()`` (stdlib) when no custom loader is
         configured. Custom loaders receive ``response.content`` (bytes).
         """
+        _enforce_response_body_cap(response)
         loader = self._transport._config.rest_json_loads
         if loader is None:
             return response.json()

tests/test_async_client.py

@@ -12,7 +12,7 @@
 import pytest
 import respx
 
-from kalshi._base_client import AsyncTransport
+from kalshi._base_client import AsyncTransport, _map_error
 from kalshi.async_client import AsyncKalshiClient
 from kalshi.auth import KalshiAuth
 from kalshi.config import DEMO_BASE_URL, PRODUCTION_BASE_URL, KalshiConfig
@@ -93,8 +93,10 @@ async def fake_sleep(d: float) -> None:
         )
         resp = await transport.request("GET", "/markets")
         assert resp.status_code == 200
-        assert sleeps == [0.0], (
-            f"Expected one sleep of 0.0s honoring Retry-After: 0, got {sleeps!r}"
+        # Floor=0 (Retry-After: 0), jitter in [0, retry_base_delay=0.01) (#321).
+        assert len(sleeps) == 1, f"Expected one sleep, got {sleeps!r}"
+        assert 0.0 <= sleeps[0] <= 0.01, (
+            f"Expected sleep in [0.0, 0.01] (Retry-After: 0 + Full Jitter), got {sleeps!r}"
         )
 
     @respx.mock
@@ -277,6 +279,90 @@ async def test_post_not_retried_on_timeout(self, transport: AsyncTransport) -> N
             await transport.request("POST", "/portfolio/orders", json={"ticker": "T"})
         assert route.call_count == 1
 
+    @respx.mock
+    @pytest.mark.asyncio
+    async def test_issue_321_retry_after_with_jitter_async(
+        self, transport: AsyncTransport, monkeypatch: pytest.MonkeyPatch
+    ) -> None:
+        """#321 async: Retry-After is floor; Full Jitter on top, capped at retry_max_delay."""
+        monkeypatch.setattr(
+            "kalshi._base_client.random.uniform",
+            lambda lo, hi: lo + (hi - lo) / 2,
+        )
+        sleeps: list[float] = []
+
+        async def fake_sleep(d: float) -> None:
+            sleeps.append(d)
+
+        monkeypatch.setattr("asyncio.sleep", fake_sleep)
+        respx.get("https://test.kalshi.com/trade-api/v2/markets").mock(
+            side_effect=[
+                httpx.Response(429, headers={"Retry-After": "0.03"}, json={"message": "rl"}),
+                httpx.Response(200, json={"markets": []}),
+            ]
+        )
+        resp = await transport.request("GET", "/markets")
+        assert resp.status_code == 200
+        # floor = min(0.03, 0.1) = 0.03; backoff cap at attempt=0 = min(0.01, 0.1) = 0.01;
+        # midpoint jitter = 0.005; delay = min(0.03 + 0.005, 0.1) = 0.035.
+        assert sleeps == pytest.approx([0.035]), (
+            f"Expected one sleep of 0.035s (floor 0.03 + 0.005 jitter), got {sleeps!r}"
+        )
+
+    @respx.mock
+    @pytest.mark.asyncio
+    async def test_issue_322_retry_after_honored_on_5xx_and_408_async(
+        self, transport: AsyncTransport, monkeypatch: pytest.MonkeyPatch
+    ) -> None:
+        """#322 async: Retry-After parsed for 408/503/504, not only 429."""
+        sleeps: list[float] = []
+
+        async def fake_sleep(d: float) -> None:
+            sleeps.append(d)
+
+        monkeypatch.setattr("asyncio.sleep", fake_sleep)
+        monkeypatch.setattr(
+            "kalshi._base_client.random.uniform",
+            lambda lo, hi: 0.0,  # zero jitter: assert the floor is honored
+        )
+
+        for status, exc_type in (
+            (408, KalshiTimeoutError),
+            (503, KalshiServerError),
+            (504, KalshiTimeoutError),
+        ):
+            sleeps.clear()
+            mapped = _map_error(
+                httpx.Response(
+                    status,
+                    headers={"Retry-After": "0.04"},
+                    json={"message": "slow"},
+                )
+            )
+            assert isinstance(mapped, exc_type), (
+                f"status {status} mapped to {type(mapped).__name__}, expected {exc_type.__name__}"
+            )
+            assert mapped.retry_after == 0.04, (
+                f"status {status}: retry_after={mapped.retry_after!r}, expected 0.04"
+            )
+
+            respx.get("https://test.kalshi.com/trade-api/v2/markets").mock(
+                side_effect=[
+                    httpx.Response(
+                        status,
+                        headers={"Retry-After": "0.04"},
+                        json={"message": "slow"},
+                    ),
+                    httpx.Response(200, json={"markets": []}),
+                ]
+            )
+            resp = await transport.request("GET", "/markets")
+            assert resp.status_code == 200
+            assert sleeps == pytest.approx([0.04]), (
+                f"status {status}: expected floor=0.04 with zero jitter, got {sleeps!r}"
+            )
+            respx.reset()
+
 
 class TestAsyncTransportContextManager:
     @pytest.mark.asyncio