mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
fix(proxy): retry upstream 529 overloaded like 429 on both forwarders (#1495)
## Description Upstream **HTTP 529** (`overloaded_error`) is not retried consistently, so it leaks to clients even though the sibling 429 path was fixed in #1221. - **Streaming forwarder** (`_stream_response`) special-cased only `status_code == 429`. A `529` falls through to `break` and is forwarded to the client with **zero retries** — interactive (streaming) Claude Code sessions see "Overloaded" immediately on a transient Anthropic overload. - **Non-streaming forwarder** (`_retry_request`) retried `529` only via the generic `>= 500` path: it **ignores `Retry-After`** and **raises** an `HTTPStatusError` on exhaustion instead of returning the clean `529` verbatim (inconsistent with how 429 is handled right above it). `529` is documented by Anthropic as the transient "overloaded" status — semantically identical to 429 for retry purposes ("try again shortly"). This PR routes both through one shared, `Retry-After`-honoring branch. Related: #1221 (added the 429 retry this extends). ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) - [ ] New feature (non-breaking change that adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to change) - [ ] Documentation update - [ ] Performance improvement - [ ] Code refactoring (no functional changes) ## Changes Made - Add `RETRYABLE_OVERLOAD_STATUSES = frozenset({429, 529})` to `proxy/helpers.py` as the single source of truth shared by both forwarders. - `streaming.py`: retry when `status_code in RETRYABLE_OVERLOAD_STATUSES` (was `== 429`); log line now interpolates the actual status. - `server.py` `_retry_request`: handle `429`/`529` in one `Retry-After`-honoring branch that returns the status verbatim once `retry_max_attempts` is exhausted (529 no longer goes through the 5xx raise path). Other 4xx/5xx behavior is unchanged. - No new dependencies; no config/API surface changes. Retry volume stays bounded by the existing `retry_max_attempts` / `retry_*_delay_ms` config. ## 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 Reproduced the CI `lint` + `commitlint` jobs exactly (pinned `ruff==0.15.17`, `mypy==1.20.2`, `@commitlint/config-conventional`), plus the affected proxy test subset: ```text # New tests in tests/test_proxy_retry_429.py — 3 of 4 fail on main, all pass here # BEFORE (source reverted, new tests kept): FAILED ::test_retry_request_returns_529_verbatim_on_exhaustion - httpx.HTTPStatusError: Server error: 529 (raised, not returned verbatim) FAILED ::test_retry_request_honors_retry_after_on_529 - slept ~0.001s (jitter), ignored Retry-After: 2 FAILED ::test_stream_response_retries_529 - assert 1 == 2 (streaming 529 forwarded raw, no retry) 3 failed, 7 passed # AFTER (this branch): 10 passed in 2.53s # Adjacent proxy suites (regression check) — retry + streaming resilience + ratelimit headers + handler helpers + request logger: 79 passed in 6.61s $ ruff check . -> All checks passed! $ ruff format --check . -> 1005 files already formatted $ mypy headroom --ignore-missing-imports Success: no issues found in 400 source files $ commitlint --from <base> --to HEAD ✔ found 0 problems, 0 warnings ``` ## Real Behavior Proof - Environment: Linux, Python 3.14.0; the proxy running **from this branch** (`headroom proxy --mode token --backend anthropic --no-optimize ...`) in front of a fake Anthropic upstream that returns a real HTTP 529 (`{"error":{"type":"overloaded_error"}}`, `Retry-After: 0`) on request #1 then a 200 SSE stream on request #2. Real proxy process over real sockets (a synthetic upstream is used because real Anthropic 529s cannot be induced on demand). - Exact command / steps: started the fake upstream on `:9911` and the branch proxy on `:9912` with `--anthropic-api-url http://127.0.0.1:9911`, then sent a streaming request: `curl -sN -X POST http://127.0.0.1:9912/v1/messages -H 'x-api-key: …' -H 'anthropic-version: 2023-06-01' -H 'content-type: application/json' -d '{"model":"claude-3-5-sonnet-20241022","max_tokens":16,"stream":true,"messages":[{"role":"user","content":"hi"}]}'` (full scripts in the code block below). - Observed result: the client received `HTTP/1.1 200 OK` and the complete SSE stream (`message_start … "hello" … message_stop`), and the fake upstream logged **two** calls — `call #1` returned 529, `call #2` returned 200 — i.e. the proxy transparently retried the 529 and the overload never reached the client. On `main` the streaming path forwards the 529 on call #1 with no retry, exactly what `test_stream_response_retries_529` pins at `calls == 1`. - Not tested: a real (non-synthetic) Anthropic 529 (cannot induce on demand); the full sharded `pytest tests scripts/tests` job (needs CI model/torch infra) — ran the proxy suite subset above instead; the Rust jobs and non-Anthropic backends (unchanged by this PR). ```bash # fake_upstream.py: 529 (Retry-After: 0) on call #1, then 200 SSE; logs each call python fake_upstream.py & # :9911 headroom proxy --host 127.0.0.1 --port 9912 \ --anthropic-api-url http://127.0.0.1:9911 \ --mode token --backend anthropic \ --no-optimize --no-cache --no-rate-limit & # :9912 (this branch) curl -sN -D - -X POST http://127.0.0.1:9912/v1/messages \ -H 'x-api-key: sk-ant-test' -H 'anthropic-version: 2023-06-01' \ -H 'content-type: application/json' \ -d '{"model":"claude-3-5-sonnet-20241022","max_tokens":16,"stream":true, "messages":[{"role":"user","content":"hi"}]}' # -> HTTP/1.1 200 OK + full SSE; upstream log: "call #1" (529) then "call #2" (200) ``` ## Review Readiness - [x] I have performed a self-review - [x] This PR is ready for human review ## 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 - [ ] I have made corresponding changes to the documentation (N/A — no doc/config surface change) - [x] My changes generate no new warnings - [x] I have added tests that prove my fix is effective or that my feature works - [x] New and existing unit tests pass locally with my changes - [x] I have updated the CHANGELOG.md if applicable ## Additional Notes - Replicated the CI `lint` job exactly (fresh venv, pinned `ruff==0.15.17` + `mypy==1.20.2`, `ruff check .` / `ruff format --check .` / `mypy headroom --ignore-missing-imports`) and `commitlint` (`@commitlint/config-conventional`) — all clean. The full `test` shards (model/torch) and Rust jobs were not run locally (no GPU/model cache / Rust toolchain in this environment); they are unaffected by this Python-only change. - `CHANGELOG.md`'s `## Unreleased` section currently contains unresolved merge-conflict markers on `main` (`<<<<<<< … >>>>>>>`) unrelated to this PR; I added my entry to the clean `### Bug Fixes` list above that region without touching the conflicts.
This commit is contained in:
Rick van Hattem
GitHub
parent
17c7347402
commit
547b15dab2
5 changed files with 104 additions and 23 deletions
|
|
@ -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).
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue