From 7836aea2be8577b0f593c9cfcff1da5f5e0ee3c8 Mon Sep 17 00:00:00 2001 From: Rocker Zhang Date: Thu, 9 Jul 2026 23:07:26 +0100 Subject: [PATCH] fix(proxy): preserve upstream 5xx status on retry exhaustion (#1570) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## What When the upstream returns a retryable 5xx (529 Overloaded, 503), `_retry_request` retried up to the cap and then raised, which the caller collapsed into a generic 502. That hides the retryable signal: clients see a 502 and give up instead of applying their own overload backoff. On exhaustion, return the last upstream response (preserving its status and body) when one is available. Connection and timeout errors still raise — only an `HTTPStatusError` carrying a real upstream response is surfaced. ## Why this scope `_retry_request` is provider-agnostic, so this fix applies uniformly to all providers (no per-handler change needed). It is purely a returned-status correctness fix and does not touch request accounting — a separate change handles counting an exhausted 5xx as a failed request across all provider handlers. ## Verification `tests/test_retry_preserve_upstream_status.py`: 529/503 status+body preservation, 4xx no-retry, connect-error still raises, success passthrough. Against unpatched main the 503-preservation test fails (collapses to 502); with the fix all pass. Addresses #1568. --- headroom/proxy/server.py | 10 ++ tests/test_retry_preserve_upstream_status.py | 113 +++++++++++++++++++ 2 files changed, 123 insertions(+) create mode 100644 tests/test_retry_preserve_upstream_status.py diff --git a/headroom/proxy/server.py b/headroom/proxy/server.py index ca7127f69..fe01960e4 100644 --- a/headroom/proxy/server.py +++ b/headroom/proxy/server.py @@ -1885,6 +1885,16 @@ class HeadroomProxy( last_error = e if not self.config.retry_enabled or attempt >= self.config.retry_max_attempts - 1: + # On exhaustion, preserve the upstream 5xx status (e.g. 503 + # Service Unavailable, 500, 502, 504) so the client can apply + # its own retry/backoff. Collapsing every exhausted 5xx into a + # generic 502 hides the retryable signal and makes clients give + # up. The 429/529 overload statuses are already returned + # verbatim by the RETRYABLE_OVERLOAD_STATUSES branch above and + # never reach here. ConnectError/TimeoutException carry no + # response, so those still raise. + if isinstance(e, httpx.HTTPStatusError) and e.response is not None: + return e.response raise # Exponential backoff with jitter diff --git a/tests/test_retry_preserve_upstream_status.py b/tests/test_retry_preserve_upstream_status.py new file mode 100644 index 000000000..4afc89bc2 --- /dev/null +++ b/tests/test_retry_preserve_upstream_status.py @@ -0,0 +1,113 @@ +"""Retry-exhaustion should preserve the upstream 5xx status (500/502/503/504) +instead of collapsing every exhausted 5xx into a generic 502. + +When upstream keeps returning a server error, the proxy used to retry, then +mask the final failure as a 502, which hides the real status from the client. +It should now surface the real 5xx so the client can apply its own backoff. + +The 429/529 overload statuses take a separate path (RETRYABLE_OVERLOAD_STATUSES, +returned verbatim by #1495's branch) and never reach the exhaustion handler +exercised here, so these tests use non-overload 5xx codes to hit it. +""" + +import asyncio +import types + +import httpx +import pytest + +from headroom.proxy.server import HeadroomProxy + + +def _make_proxy(http_client, *, retry_max_attempts=3, retry_enabled=True): + proxy = HeadroomProxy.__new__(HeadroomProxy) + proxy.http_client = http_client + proxy.config = types.SimpleNamespace( + retry_enabled=retry_enabled, + retry_max_attempts=retry_max_attempts, + retry_base_delay_ms=1, + retry_max_delay_ms=2, + ) + return proxy + + +class _FakeClient: + def __init__(self, *, response=None, exc=None): + self._response = response + self._exc = exc + self.calls = 0 + + async def post(self, url, content=None, headers=None): + self.calls += 1 + if self._exc is not None: + raise self._exc + return self._response + + +def _resp(status, body=b'{"type":"error"}'): + req = httpx.Request("POST", "https://api.anthropic.com/v1/messages") + return httpx.Response(status_code=status, request=req, content=body) + + +def _call(proxy): + return asyncio.run( + proxy._retry_request( + "POST", + "https://api.anthropic.com/v1/messages", + {}, + {}, + stream=False, + ) + ) + + +def test_exhausted_500_preserves_status(): + client = _FakeClient(response=_resp(500)) + out = _call(_make_proxy(client, retry_max_attempts=3)) + assert out.status_code == 500 # not collapsed to 502 + assert client.calls == 3 # retried up to the cap + + +def test_exhausted_503_preserves_status(): + client = _FakeClient(response=_resp(503)) + out = _call(_make_proxy(client, retry_max_attempts=2)) + assert out.status_code == 503 + assert client.calls == 2 + + +def test_exhausted_5xx_preserves_body(): + client = _FakeClient( + response=_resp(503, body=b'{"type":"error","error":{"type":"overloaded_error"}}') + ) + out = _call(_make_proxy(client, retry_max_attempts=2)) + assert out.status_code == 503 + assert out.json()["error"]["type"] == "overloaded_error" # body survives, not just status + + +def test_retry_disabled_returns_5xx_on_first_attempt(): + # Use a non-overload 5xx (503) so this exercises the HTTPStatusError + # exhaustion branch rather than the 429/529 overload fast path. + client = _FakeClient(response=_resp(503)) + out = _call(_make_proxy(client, retry_enabled=False)) + assert out.status_code == 503 + assert client.calls == 1 # no retry, but status still preserved instead of raised + + +def test_4xx_returned_without_retry(): + client = _FakeClient(response=_resp(400)) + out = _call(_make_proxy(client)) + assert out.status_code == 400 + assert client.calls == 1 # client errors are not retried + + +def test_connect_error_still_raises(): + client = _FakeClient(exc=httpx.ConnectError("boom")) + with pytest.raises(httpx.ConnectError): + _call(_make_proxy(client)) + + +def test_success_passes_through(): + client = _FakeClient(response=_resp(200)) + out = _call(_make_proxy(client)) + assert out.status_code == 200 + assert client.calls == 1