diff --git a/CHANGELOG.md b/CHANGELOG.md index 704272785..b0e87771a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -42,6 +42,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 * **proxy:** stop discarding a finished compression on very large requests. After the transform pipeline completed, a telemetry-only waste-signal re-parse of the *original* messages ran on the critical path; on huge Claude Code transcripts (~400k tokens) that parse could exceed the Anthropic compression timeout, so the proxy failed open and forwarded the uncompressed request despite "Pipeline complete" logging real savings (`tokens_saved: 0`, `transforms_applied: []`, ~31s latency). Waste-signal detection is now skipped above `MAX_WASTE_SIGNAL_DETECTION_TOKENS` (100k) so the compression result stays on the critical path ([#296](https://github.com/chopratejas/headroom/issues/296)). * **codex:** retag existing Codex threads when `headroom init` injects the `headroom` provider, so Codex Desktop history stays visible. Codex filters its sidebar/search by the active `model_provider`; the init path set `model_provider = "headroom"` without retagging, so existing native `openai` threads disappeared from the menu (data was never deleted, only hidden). `_ensure_codex_provider` now reconciles thread tags openai→headroom, matching what the install and `wrap` paths already do; `headroom unwrap codex` handles the revert direction ([#961](https://github.com/chopratejas/headroom/issues/961)). * **install:** stop duplicating the container ENTRYPOINT in the `persistent-docker` runtime command. The published image already runs `headroom proxy` as its ENTRYPOINT, but `build_runtime_command` re-added `headroom proxy` after the image name, so the container ran `headroom proxy headroom proxy --host 0.0.0.0 …` and Click aborted with "Got unexpected extra arguments (headroom proxy)" — the deployment never became ready and rollback left nothing running. The runtime command now appends only the proxy flags ([#833](https://github.com/chopratejas/headroom/issues/833)). +* **proxy:** retry upstream `529 overloaded_error` like a 429 on both the streaming and non-streaming forwarders, honoring `Retry-After`. The streaming path previously surfaced a 529 straight to the client with no retry (interactive sessions saw "Overloaded" immediately), and `_retry_request` retried it only via the generic 5xx path — raising on exhaustion instead of returning the 529 verbatim, and ignoring `Retry-After`. A shared `RETRYABLE_OVERLOAD_STATUSES = {429, 529}` keeps the two forwarders in agreement (extends [#1221](https://github.com/headroomlabs-ai/headroom/issues/1221)). * **gemini:** run compression off the asyncio event loop. The Gemini handlers (`generateContent`, Cloud Code stream, `countTokens`) ran the CPU-bound compression pipeline (Magika detection plus ML compression) synchronously on the loop, stalling every concurrent request for the duration of each Gemini request's compression. They now offload it via the shared compression executor, matching the existing OpenAI and Anthropic paths. * **proxy:** queue mid-turn user messages on non-Bedrock streaming path instead of silently dropping them — closes [#902](https://github.com/headroomlabs-ai/headroom/issues/902). * **proxy:** add `--protect-tool-results` / `HEADROOM_PROTECT_TOOL_RESULTS` to prevent lossy compression of exact-output tool results (e.g. `Bash cat`/`grep` results) — closes [#1307](https://github.com/headroomlabs-ai/headroom/issues/1307). diff --git a/headroom/proxy/handlers/streaming.py b/headroom/proxy/handlers/streaming.py index b301d00d6..7292ea41e 100644 --- a/headroom/proxy/handlers/streaming.py +++ b/headroom/proxy/handlers/streaming.py @@ -13,7 +13,11 @@ import time from typing import TYPE_CHECKING, Any from headroom.proxy.auth_mode import classify_client -from headroom.proxy.helpers import jitter_delay_ms, retry_after_ms +from headroom.proxy.helpers import ( + RETRYABLE_OVERLOAD_STATUSES, + jitter_delay_ms, + retry_after_ms, +) if TYPE_CHECKING: from fastapi.responses import Response, StreamingResponse @@ -997,11 +1001,12 @@ 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. + # Retry transient overloads (429 rate-limit, 529 overloaded) + # honoring Retry-After — the streaming sibling of the + # _retry_request path (#1221); on exhaustion, fall through to + # forward the status to the client. if ( - upstream_response.status_code == 429 + upstream_response.status_code in RETRYABLE_OVERLOAD_STATUSES and self.config.retry_enabled and attempt < retry_attempts - 1 ): @@ -1014,7 +1019,7 @@ class StreamingMixin: ) await upstream_response.aclose() logger.warning( - f"[{request_id}] Upstream 429 " + f"[{request_id}] Upstream {upstream_response.status_code} " f"(attempt {attempt + 1}/{retry_attempts}), " f"retrying in {delay_with_jitter:.0f}ms" ) diff --git a/headroom/proxy/helpers.py b/headroom/proxy/helpers.py index e2914c289..38be067bd 100644 --- a/headroom/proxy/helpers.py +++ b/headroom/proxy/helpers.py @@ -952,6 +952,14 @@ def retry_after_ms(response: httpx.Response, max_ms: int) -> float | None: return min(max(seconds, 0.0) * 1000.0, float(max_ms)) +# Transient upstream statuses worth retrying with backoff: 429 (rate limit) and +# 529 (Anthropic ``overloaded_error``). Both mean "the server is temporarily +# limiting/overloaded — try again shortly", unlike other 4xx which signal a +# problem with the request itself. Single source of truth so the streaming and +# non-streaming forwarders agree on what is retriable. +RETRYABLE_OVERLOAD_STATUSES: frozenset[int] = frozenset({429, 529}) + + # Image compression availability (do not retain a global compressor instance) _image_compressor_available: bool | None = None diff --git a/headroom/proxy/server.py b/headroom/proxy/server.py index 34f6fa126..03749011b 100644 --- a/headroom/proxy/server.py +++ b/headroom/proxy/server.py @@ -128,6 +128,7 @@ from headroom.proxy.helpers import ( MAX_MESSAGE_ARRAY_LENGTH, # noqa: F401 MAX_REQUEST_BODY_SIZE, # noqa: F401 MAX_SSE_BUFFER_SIZE, # noqa: F401 + RETRYABLE_OVERLOAD_STATUSES, _get_context_tool_stats, _get_image_compressor, # noqa: F401 _get_rtk_stats, # noqa: F401 @@ -1759,22 +1760,11 @@ class HeadroomProxy( url, **post_kwargs ) - # 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) - if response.status_code >= 500: - raise httpx.HTTPStatusError( - f"Server error: {response.status_code}", - request=response.request, - 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: + # Transient overloads (429 rate-limit, 529 overloaded): + # retry honoring Retry-After, but return verbatim once + # exhausted — a clean overload signal beats a synthesized 5xx + # (extends #1221 to 529, Anthropic's overloaded_error). + if response.status_code in RETRYABLE_OVERLOAD_STATUSES: if ( not self.config.retry_enabled or attempt >= self.config.retry_max_attempts - 1 @@ -1788,11 +1778,24 @@ class HeadroomProxy( attempt, ) logger.warning( - f"Upstream 429 (attempt {attempt + 1}), retrying in {delay_ms:.0f}ms" + f"Upstream {response.status_code} (attempt {attempt + 1}), " + f"retrying in {delay_ms:.0f}ms" ) await asyncio.sleep(delay_ms / 1000) continue + # Don't retry other client errors (4xx) + if 400 <= response.status_code < 500: + return response + + # Retry other server errors (5xx) + if response.status_code >= 500: + raise httpx.HTTPStatusError( + f"Server error: {response.status_code}", + request=response.request, + response=response, + ) + return response except (httpx.ConnectError, httpx.TimeoutException, httpx.HTTPStatusError) as e: diff --git a/tests/test_proxy_retry_429.py b/tests/test_proxy_retry_429.py index cf4733d2f..67f145c6e 100644 --- a/tests/test_proxy_retry_429.py +++ b/tests/test_proxy_retry_429.py @@ -155,3 +155,67 @@ def test_stream_response_retries_429() -> None: ) ) assert transport.calls == 2 # streaming 429 retried, not forwarded raw + + +# --- 529 overloaded: same transient-retry path as 429 -------------------- +# +# 529 is Anthropic's ``overloaded_error``. Like 429 it means "try again +# shortly", so both forwarders must retry it honoring Retry-After. Before this +# fix the streaming path forwarded a 529 to the client raw (zero retries), and +# _retry_request retried it only via the generic 5xx path — raising on +# exhaustion instead of returning the 529 verbatim, and ignoring Retry-After. + + +def test_retry_request_retries_529_then_succeeds() -> None: + transport = _RateLimitTransport(fail_status=529, 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 529 + one success — the retry happened + + +def test_retry_request_returns_529_verbatim_on_exhaustion() -> None: + # Always 529: must return the 529 to the client, NOT raise / convert to 5xx. + transport = _RateLimitTransport(fail_status=529, 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 == 529 + assert transport.calls == 3 # exhausted all attempts, returned verbatim + + +def test_retry_request_honors_retry_after_on_529(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=529, 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 for 529 just like 429. + assert slept and abs(slept[0] - 2.0) < 0.01 + + +def test_stream_response_retries_529() -> None: + # The gap this PR closes: an interactive (streaming) session hitting a 529 + # used to get "Overloaded" surfaced immediately, with no retry. + transport = _RateLimitTransport(fail_status=529, 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 529 retried, not forwarded raw