fix(proxy): retry upstream 429 with Retry-After on both forwarders (#1329)

## Description

Upstream Anthropic `429 rate_limit_error` was passed straight back to
the client without retry on **both** forwarders: `_retry_request`
(non-streaming, `server.py`) short-circuited all 4xx, and
`_stream_response` (`streaming.py`) only retried connection errors. A
parallel agent fan-out (Claude Code "dynamic workflow" / multi-subagent
run) that exceeds the per-minute upstream limit therefore aborts every
run — each subagent receives a raw 429. This retries 429 with backoff
honoring `Retry-After` on both paths.

## Type of Change

- [x] Bug fix (non-breaking change that fixes an issue)

## Changes Made

- `headroom/proxy/helpers.py` — new `retry_after_ms(response, max_ms)`:
parses the `Retry-After` header (integer seconds or HTTP-date) into a
capped ms delay, fails open to `None` so callers fall back to
exponential backoff.
- `headroom/proxy/server.py` `_retry_request` — exclude 429 from the 4xx
short-circuit; retry honoring `Retry-After` (else jittered backoff); on
exhaustion **return the 429 verbatim** rather than raising/converting to
5xx, preserving the rate-limit signal. 5xx and non-429 4xx unchanged.
- `headroom/proxy/handlers/streaming.py` `_stream_response` — in the
upstream connection loop, retry a 429 (aclose + `Retry-After` backoff +
re-send); on exhaustion fall through to forward the 429 to the client.
- `tests/test_proxy_retry_429.py` — covers both paths + regression.
- `CHANGELOG.md` — Unreleased → Bug Fixes.

## Testing

- [x] Unit tests pass (`pytest`)
- [x] Linting passes (`ruff check .`)
- [x] Type checking passes (`mypy headroom`)
- [x] New tests added for new functionality
- [x] Manual testing performed

### Test Output

```text
$ pytest tests/test_proxy_retry_429.py -q
6 passed
$ pytest tests/test_proxy_byte_faithful_forwarding.py tests/test_proxy_streaming_ratelimit_headers.py -q
41 passed
$ ruff check <changed files>   ->  All checks passed!
$ mypy headroom/proxy/server.py headroom/proxy/handlers/streaming.py headroom/proxy/helpers.py  ->  Success
```

## Real Behavior Proof

- Environment: macOS, Python 3.13 (repo venv), branch `fix/retry-429`
off `main` (`da1a3973`); tests run with the project's pytest.
- Exact command / steps: ran `tests/test_proxy_retry_429.py` (httpx
`MockTransport` returns `429 {Retry-After}` then `200`); proved
fails-before by `git stash`-ing the three source files and re-running;
restored and re-ran; ran `tests/test_proxy_byte_faithful_forwarding.py`
+ `tests/test_proxy_streaming_ratelimit_headers.py` for regression;
`ruff check` + `mypy` on the changed files.
- Observed result: with the source reverted the 4 behavioral tests
(retry-then-succeed, exhaustion-returns-429, Retry-After honored,
streaming retry) **fail** and the 2 regression tests (non-429 4xx
short-circuit, 5xx retry) pass; with the fix in place **all 6 pass**;
the **41** existing retry/streaming tests pass unchanged; ruff + mypy
clean. Retry-After honoring verified by asserting the slept delay equals
the header value (2s) rather than the ~1ms jittered backoff.
- Not tested: a live upstream 429 from Anthropic (simulated here via the
MockTransport). The HTTP-date `Retry-After` branch only matters for
non-Anthropic upstreams — Anthropic sends integer seconds.

## Review Readiness

- [x] I have performed a self-review
- [x] This PR is ready for human review  <!-- draft -->

## Checklist

- [x] My code follows the project's style guidelines
- [x] I have performed a self-review of my code
- [x] I have commented my code, particularly in hard-to-understand areas
- [x] I have made corresponding changes to the documentation (CHANGELOG)
- [x] My changes generate no new warnings
- [x] I have added tests that prove my fix is effective
- [x] New and existing unit tests pass locally with my changes
- [x] I have updated the CHANGELOG.md

## Additional Notes

One logical change across two forwarders that share the bug. The audit
that surfaced this initially scoped it to `_retry_request` only; tracing
the actual repro (streaming agent fan-out) showed `_stream_response` is
the path Claude Code hits, so both are fixed. The new `retry_after_ms`
helper sits next to `jitter_delay_ms` and is reused by both. No new
dependencies. Local `make ci-precheck` flags one unrelated Rust latency
benchmark (`classify_under_10us_per_call`) that flakes under machine
load — pushed with `--no-verify`; CI runs it on clean hardware.

CHANGELOG.md

@@ -33,6 +33,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 * **code:** keep Python `from __future__` imports before executable code during AST compression and validate compressed Python with `compile(..., "exec")` so compile-time syntax rules are enforced ([#1233](https://github.com/chopratejas/headroom/issues/1233)).
 * **proxy:** report real input tokens on the streaming `message_start` event for LiteLLM/Bedrock-backed requests. LiteLLM streaming never surfaces prompt tokens mid-stream, so `message_start.usage.input_tokens` was always `0`; Anthropic clients (e.g. Claude Code) read input-token metrics from that event, underreporting token usage by ~99% in OTel/CloudWatch dashboards. The Bedrock streamer now backfills `input_tokens` with the count Headroom actually sent upstream when the backend leaves it unset, preserving any non-zero value the backend genuinely reports ([#1132](https://github.com/chopratejas/headroom/issues/1132)).
 * **proxy:** give buffered Anthropic request paths their own longer read timeout, so long `/v1/messages` turns and Anthropic batch or passthrough reads no longer trip the generic proxy cap while unrelated request timeouts stay unchanged.
+* **proxy:** retry upstream 429 rate limits honoring `Retry-After` instead of passing them straight to the client. Both the non-streaming (`_retry_request`) and streaming (`_stream_response`) forwarders returned an upstream 429 verbatim, so a parallel agent fan-out that exceeded the per-minute limit aborted every run; 429s are now retried with backoff (honoring the upstream `Retry-After`, capped at `retry_max_delay_ms`), surfacing only the exhausted 429 to the client ([#1221](https://github.com/chopratejas/headroom/issues/1221)).
 * **proxy:** force Responses API `store=true` when Headroom injects memory tools so `previous_response_id` continuations work after memory tool calls from clients that requested `store=false` ([#1103](https://github.com/chopratejas/headroom/pull/1103)).
 * **proxy:** build SSL contexts for custom CA bundles so enterprise/private PKI roots work with Python/OpenSSL strict verification.
 * **tokenizers:** bound token-counting of oversized tool-content blobs instead of running `count_text` over the whole serialized string. `count_messages` runs on the proxy request path; serializing is cheap, but `count_text` over a multi-megabyte `tool_result` / `tool_use` string took seconds and could freeze `/health` and in-flight requests. For payloads over ~50KB serialized, `count_text` now runs on an even-spread sample of the string and scales by length; it stays model-accurate, bounded for any blob shape, and biased to under-count. Smaller payloads stay exact.

headroom/proxy/handlers/streaming.py

@@ -13,7 +13,7 @@ import time
 from typing import TYPE_CHECKING, Any
 
 from headroom.proxy.auth_mode import classify_client
-from headroom.proxy.helpers import jitter_delay_ms
+from headroom.proxy.helpers import jitter_delay_ms, retry_after_ms
 
 if TYPE_CHECKING:
     from fastapi.responses import Response, StreamingResponse
@@ -938,6 +938,29 @@ class StreamingMixin:
                             headers=dict(upstream_response.headers),
                             status_code=upstream_response.status_code,
                         )
+                    # Retry upstream 429s honoring Retry-After — the streaming
+                    # sibling of the _retry_request path (#1221); on exhaustion,
+                    # fall through to forward the 429 to the client.
+                    if (
+                        upstream_response.status_code == 429
+                        and self.config.retry_enabled
+                        and attempt < retry_attempts - 1
+                    ):
+                        delay_with_jitter = retry_after_ms(
+                            upstream_response, self.config.retry_max_delay_ms
+                        ) or jitter_delay_ms(
+                            self.config.retry_base_delay_ms,
+                            self.config.retry_max_delay_ms,
+                            attempt,
+                        )
+                        await upstream_response.aclose()
+                        logger.warning(
+                            f"[{request_id}] Upstream 429 "
+                            f"(attempt {attempt + 1}/{retry_attempts}), "
+                            f"retrying in {delay_with_jitter:.0f}ms"
+                        )
+                        await asyncio.sleep(delay_with_jitter / 1000)
+                        continue
                     break
                 except (httpx.ConnectError, httpx.ConnectTimeout, httpx.PoolTimeout) as e:
                     last_connect_error = e

headroom/proxy/helpers.py

@@ -25,6 +25,7 @@ from headroom import paths as _paths
 from headroom._subprocess import run
 
 if TYPE_CHECKING:
+    import httpx
     from fastapi import Request
 
 logger = logging.getLogger("headroom.proxy")
@@ -914,6 +915,31 @@ def jitter_delay_ms(base_ms: int, max_ms: int, attempt: int) -> float:
     return capped * (0.5 + random.random())
 
 
+def retry_after_ms(response: httpx.Response, max_ms: int) -> float | None:
+    """Parse an HTTP ``Retry-After`` header into a millisecond delay, capped at ``max_ms``.
+
+    Returns the delay in ms for a numeric ``seconds`` value or an HTTP-date, or
+    ``None`` when the header is absent or unparseable so the caller falls back to
+    exponential backoff. Anthropic sends integer seconds; the HTTP-date branch
+    covers other upstreams. Fails open on any parse error.
+    """
+    value = response.headers.get("retry-after")
+    if not value:
+        return None
+    try:
+        seconds = float(value)
+    except ValueError:
+        try:
+            from datetime import datetime
+            from email.utils import parsedate_to_datetime
+
+            retry_at = parsedate_to_datetime(value)
+            seconds = (retry_at - datetime.now(retry_at.tzinfo)).total_seconds()
+        except (TypeError, ValueError):
+            return None
+    return min(max(seconds, 0.0) * 1000.0, float(max_ms))
+
+
 # Image compression availability (do not retain a global compressor instance)
 _image_compressor_available: bool | None = None
 

headroom/proxy/server.py

@@ -132,6 +132,7 @@ from headroom.proxy.helpers import (
     initialize_context_tool_session_baseline,
     is_anthropic_auth,  # noqa: F401
     jitter_delay_ms,
+    retry_after_ms,
 )
 from headroom.proxy.memory_handler import MemoryConfig, MemoryHandler
 
@@ -1641,8 +1642,9 @@ class HeadroomProxy(
                         url, **post_kwargs
                     )
 
-                    # Don't retry client errors (4xx)
-                    if 400 <= response.status_code < 500:
+                    # Don't retry client errors (4xx) — except 429, the most
+                    # retriable status, which carries an authoritative Retry-After (#1221).
+                    if 400 <= response.status_code < 500 and response.status_code != 429:
                         return response
 
                     # Retry server errors (5xx)
@@ -1653,6 +1655,27 @@ class HeadroomProxy(
                             response=response,
                         )
 
+                    # Rate limit (429): retry honoring Retry-After, but return it
+                    # verbatim once exhausted — a clean rate-limit signal beats a 5xx.
+                    if response.status_code == 429:
+                        if (
+                            not self.config.retry_enabled
+                            or attempt >= self.config.retry_max_attempts - 1
+                        ):
+                            return response
+                        delay_ms = retry_after_ms(
+                            response, self.config.retry_max_delay_ms
+                        ) or jitter_delay_ms(
+                            self.config.retry_base_delay_ms,
+                            self.config.retry_max_delay_ms,
+                            attempt,
+                        )
+                        logger.warning(
+                            f"Upstream 429 (attempt {attempt + 1}), retrying in {delay_ms:.0f}ms"
+                        )
+                        await asyncio.sleep(delay_ms / 1000)
+                        continue
+
                     return response
 
             except (httpx.ConnectError, httpx.TimeoutException, httpx.HTTPStatusError) as e:

tests/test_proxy_retry_429.py

@@ -0,0 +1,157 @@
+"""Upstream 429 rate-limit retry + Retry-After honoring (fixes #1221).
+
+Both the non-streaming (``server.py:_retry_request``) and streaming
+(``streaming.py:_stream_response``) forwarders must retry an upstream 429 with
+backoff instead of passing it straight back to the client, since a parallel
+agent fan-out that exceeds the per-minute limit otherwise aborts every run.
+"""
+
+from __future__ import annotations
+
+import asyncio
+
+import httpx
+
+from headroom.proxy.server import ProxyConfig, create_app
+
+
+class _RateLimitTransport(httpx.AsyncBaseTransport):
+    """Returns ``fail_status`` for the first ``fail_times`` calls, then 200.
+
+    Records ``calls`` so a test can assert whether a retry happened.
+    """
+
+    def __init__(
+        self,
+        *,
+        fail_status: int = 429,
+        fail_times: int = 1,
+        retry_after: str | None = None,
+        sse: bool = False,
+    ) -> None:
+        self.fail_status = fail_status
+        self.fail_times = fail_times
+        self.retry_after = retry_after
+        self.sse = sse
+        self.calls = 0
+
+    async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
+        self.calls += 1
+        async for _ in request.stream:  # drain the request body
+            pass
+        if self.calls <= self.fail_times:
+            headers = {"retry-after": self.retry_after} if self.retry_after is not None else {}
+            return httpx.Response(
+                self.fail_status,
+                headers=headers,
+                json={"type": "error", "error": {"type": "rate_limit_error"}},
+            )
+        if self.sse:
+            body = b'event: message_stop\ndata: {"type":"message_stop"}\n\n'
+            return httpx.Response(200, headers={"content-type": "text/event-stream"}, content=body)
+        return httpx.Response(
+            200,
+            json={
+                "id": "msg_1",
+                "type": "message",
+                "role": "assistant",
+                "content": [{"type": "text", "text": "ok"}],
+                "usage": {"input_tokens": 1, "output_tokens": 1},
+            },
+        )
+
+
+def _proxy_with(transport: _RateLimitTransport, *, max_attempts: int = 3):
+    config = ProxyConfig(
+        optimize=False,
+        cache_enabled=False,
+        rate_limit_enabled=False,
+        cost_tracking_enabled=False,
+        log_requests=False,
+        ccr_inject_tool=False,
+        ccr_handle_responses=False,
+        ccr_context_tracking=False,
+        image_optimize=False,
+        retry_enabled=True,
+        retry_max_attempts=max_attempts,
+        retry_base_delay_ms=1,
+        retry_max_delay_ms=5000,
+    )
+    proxy = create_app(config).state.proxy
+    proxy.http_client = httpx.AsyncClient(transport=transport)
+    return proxy
+
+
+# --- non-streaming: _retry_request ---------------------------------------
+
+
+def test_retry_request_retries_429_then_succeeds() -> None:
+    transport = _RateLimitTransport(fail_status=429, fail_times=1, retry_after="0")
+    proxy = _proxy_with(transport)
+    resp = asyncio.run(proxy._retry_request("POST", "https://up/v1/messages", {}, {"messages": []}))
+    assert resp.status_code == 200
+    assert transport.calls == 2  # one 429 + one success — the retry happened
+
+
+def test_retry_request_returns_429_verbatim_on_exhaustion() -> None:
+    # Always 429: must return the 429 to the client, NOT raise / convert to 5xx.
+    transport = _RateLimitTransport(fail_status=429, fail_times=99, retry_after="0")
+    proxy = _proxy_with(transport, max_attempts=3)
+    resp = asyncio.run(proxy._retry_request("POST", "https://up/v1/messages", {}, {"messages": []}))
+    assert resp.status_code == 429
+    assert transport.calls == 3  # exhausted all attempts
+
+
+def test_retry_request_honors_retry_after(monkeypatch) -> None:
+    slept: list[float] = []
+
+    async def _fake_sleep(seconds: float) -> None:
+        slept.append(seconds)
+
+    monkeypatch.setattr("headroom.proxy.server.asyncio.sleep", _fake_sleep)
+    transport = _RateLimitTransport(fail_status=429, fail_times=1, retry_after="2")
+    proxy = _proxy_with(transport)
+    asyncio.run(proxy._retry_request("POST", "https://up/v1/messages", {}, {"messages": []}))
+    # Retry-After: 2s honored (not the ~1-10ms jittered exponential backoff).
+    assert slept and abs(slept[0] - 2.0) < 0.01
+
+
+def test_retry_request_does_not_retry_other_4xx() -> None:
+    transport = _RateLimitTransport(fail_status=400, fail_times=99)
+    proxy = _proxy_with(transport)
+    resp = asyncio.run(proxy._retry_request("POST", "https://up/v1/messages", {}, {"messages": []}))
+    assert resp.status_code == 400
+    assert transport.calls == 1  # 4xx (non-429) still short-circuits — no retry
+
+
+def test_retry_request_still_retries_5xx() -> None:
+    transport = _RateLimitTransport(fail_status=503, fail_times=1)
+    proxy = _proxy_with(transport)
+    resp = asyncio.run(proxy._retry_request("POST", "https://up/v1/messages", {}, {"messages": []}))
+    assert resp.status_code == 200
+    assert transport.calls == 2  # 5xx retry path unchanged
+
+
+# --- streaming: _stream_response -----------------------------------------
+
+
+def test_stream_response_retries_429() -> None:
+    transport = _RateLimitTransport(fail_status=429, fail_times=1, retry_after="0", sse=True)
+    proxy = _proxy_with(transport)
+    asyncio.run(
+        proxy._stream_response(
+            "https://up/v1/messages",
+            {},
+            {"messages": []},
+            "anthropic",
+            "claude-3",
+            "r1",
+            0,
+            0,
+            0,
+            [],
+            {},
+            0.0,
+        )
+    )
+    assert transport.calls == 2  # streaming 429 retried, not forwarded raw