fix(proxy): restore the buffered-CCR heartbeat behind a grace window (#3091)
## Description Closes #3079 #2997 removed the buffered-CCR keepalive preamble from both provider handlers. That preamble was added by #2479 to close #2465, so `main` is back to the condition #2465 described while #2465 stays closed. Confirmed against the tags: ``` v0.35.0 (released): keepalive_deadline = loop.time() + 1.0 + b'event: ping...' main (-> 0.36.0): neither ``` Since #2997 is queued in #3067, 0.36.0 would ship this. The justification left in the code does not hold. It reads *"clients budget minutes for a turn (Claude Code sends `x-stainless-timeout: 600`), so waiting is free"* — but `x-stainless-timeout` is the **total request** budget, and #2465 was about the **stream idle** watchdog, a separate timer. Total-budget headroom says nothing about idle-budget headroom, and not every client sends 600. The reporter's buffered turns routinely run 15-25s, all of it silent. The status-ordering half of #2997 is correct and is kept. What was wrong was treating the two properties as a trade. ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) ## Changes Made - New `headroom/proxy/buffered_ccr_response.py` holding one implementation of the buffered-CCR ASGI wrapper. Both handlers previously carried ~100 duplicated lines each, which is how the OpenAI twin drifted from the Anthropic one; now only the error wire format differs. - A buffered turn holds out for `buffered_ccr_grace_seconds` before committing. Inside the window nothing is sent and the result is relayed untouched, so a fast 4xx — or a 429/529 that resolves once `_retry_request` has honored `Retry-After` — keeps its real status and headers. That is #2997's property. - Past the window the response is committed as SSE and a heartbeat starts, so a first byte always precedes the client's idle watchdog. That is #2479's property. - A failure landing after the commit can no longer carry an HTTP status, so it is translated into the provider's own **typed** SSE error (`rate_limit_error`, `overloaded_error`, ...) carrying the upstream's own message where there is one, rather than a generic `api_error`. **This is the part worth reviewing.** #2997 was right that early commits broke client backoff — but that was a consequence of degrading every post-commit failure to a bare `api_error`, not of committing itself. - New `buffered_ccr_grace_seconds` on `ProxyConfig`, default 5s, env `HEADROOM_BUFFERED_CCR_GRACE_SECONDS`. Setting it to `0` restores current `main` behavior exactly. ## Testing - [x] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check .`) - [x] New tests added for new functionality - [x] Manual testing performed **#2997's own tests pass unchanged.** `test_buffered_ccr_preserves_late_failure_status_and_headers` and `test_buffered_ccr_withholds_output_until_delayed_upstream_resolves` resolve at 1.1s, comfortably inside the 5s window, so everything #2997 bought for the cases it tested is intact. New `tests/test_buffered_ccr_grace_window.py` pins both constraints @JerrettDavis asked for on #2959 — a late failure keeping its real status and headers, and a slow success producing a first byte before the ceiling — plus the typed-error mapping, the zero-grace escape hatch, and the OpenAI wire format. ### Test Output ```text $ pytest tests/test_buffered_ccr_grace_window.py -q 9 passed in 2.01s $ pytest tests/test_proxy/test_anthropic_streaming_ccr_retrieve.py tests/test_ccr_buffered_stream_signed_thinking.py -q 16 passed, 1 warning in 8.19s $ pytest tests/ -q 3 failed, 11157 passed, 581 skipped in 334.97s (0:05:34) Same 3 failures as a clean-main baseline run on this machine: tests/test_graceful_shutdown.py::test_run_server_installs_cancelled_error_filter tests/test_learn/test_integration.py::TestCodexIntegration::test_full_pipeline tests/test_release_workflows.py::test_no_native_tls_in_wheel_build_tree (missing local `cargo` toolchain; a stale tool-name fixture; a logging test that loses to global-state pollution in a full run — all present on main.) $ ruff check . && ruff format --check . All checks passed! ``` ## Real Behavior Proof - Environment: this branch, the wrapper driven directly over ASGI with a stubbed buffered operation standing in for upstream; macOS arm64, Python 3.12. - Exact command / steps: drove three scenarios — a 429 resolving at 0.05s with a 5s window; a success released only after the window with a 0.1s window; a 429 resolving at 0.3s with a 0.05s window — recording every ASGI message sent. - Observed result: (1) client receives **HTTP 429** with `retry-after: 30` and zero bytes beforehand; (2) first byte (`200 text/event-stream` + ping) arrives **before** the upstream resolves, body follows intact; (3) committed 200, then `event: error` typed `rate_limit_error` carrying the upstream's own message rather than a generic one. - Not tested: a live client idle-timeout against a real 15-25s turn. #3079 notes this depends on whether the client's idle timer starts at request send or at first response byte — the grace window makes Headroom correct under either reading, but confirming the original symptom is gone needs the reporter's fleet. ## Runtime Rollout Safety - Rollout-managed feature(s): none — this is a fix to an existing code path, not behind a rollout channel. - Minimum rollout channel: n/a (ships to stable with the fix). - Stable/default behavior changed: yes. A buffered-CCR turn slower than 5s now emits SSE headers plus keepalives instead of staying silent. Turns resolving under 5s are byte-identical to current `main`. - Kill switch / disable path: `HEADROOM_BUFFERED_CCR_GRACE_SECONDS=0` restores current `main` behavior exactly (never commit early, no heartbeat). Covered by `test_a_zero_grace_window_never_commits_early`. - Unsafe override required: no. - Qualification impact: none — no qualification-gated surface is touched. - Rollback path: revert this commit, or set the env var to `0` without a redeploy of code. ## 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 - [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 did **not** edit `CHANGELOG.md` ## Additional Notes Documentation update is marked N/A: the new env var is documented in the module docstring and the `ProxyConfig` field comment, matching how `HEADROOM_ANTHROPIC_BUFFERED_REQUEST_TIMEOUT_SECONDS` is handled. Type checking (`mypy headroom`) was not run separately; `ruff` is the gate this repo's CI enforces. Related: #2465, #2479, #2959, #2968, #2997, #3067. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-authored-by: Tejas Chopra <tejas@Tejass-MacBook-Pro.local> Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
Tejas Chopra
@@ -0,0 +1,312 @@
|
||||
"""The ASGI wrapper for a buffered-CCR turn, shared by both provider handlers.
|
||||
|
||||
Server-side CCR retrieval needs the whole upstream reply in hand before it can
|
||||
answer, so a ``stream: true`` turn is flipped to ``stream: false`` upstream and
|
||||
resynthesized as SSE on the way out. That leaves a window — the entire
|
||||
generation — where the proxy is holding a request open with nothing to say yet,
|
||||
and two constraints pull in opposite directions across it:
|
||||
|
||||
* **Status fidelity.** Committing ``200 text/event-stream`` before the outcome
|
||||
is known destroys it. Any reply that then fails to become SSE reaches the
|
||||
client as a 200 with no ``message_start`` ("API returned an empty or
|
||||
malformed response (HTTP 200)"), and the real status goes with it, so
|
||||
client-side 429/5xx backoff never fires. This is what #2997 fixed by never
|
||||
committing early.
|
||||
|
||||
* **Liveness.** Sending nothing at all for the whole wait trips the client's
|
||||
*stream-idle* watchdog. That timer is separate from the total-request budget
|
||||
(``x-stainless-timeout``), and it is the one #2465 was about. #2479 fixed it
|
||||
with a keepalive preamble, which #2997 removed as collateral — putting #2465's
|
||||
condition back (#3079).
|
||||
|
||||
Neither property is worth trading for the other, so this keeps both:
|
||||
|
||||
1. For ``grace_seconds`` nothing is sent, and anything resolving inside that
|
||||
window is handed to the client untouched — real status, real headers. Fast
|
||||
failures (a 4xx, or a 429/529 that resolves once ``_retry_request`` has
|
||||
honored ``Retry-After``) land here.
|
||||
2. Past the window the response is committed as SSE and a heartbeat starts, so
|
||||
a first byte always precedes any client idle watchdog.
|
||||
3. A failure arriving *after* the commit can no longer carry an HTTP status, so
|
||||
it is translated into the provider's own typed SSE error instead of a generic
|
||||
one. A rate limit still reads as a rate limit, and client backoff still
|
||||
fires — the property that made early commits harmful in the first place.
|
||||
|
||||
Set ``grace_seconds`` to 0 or less to disable the heartbeat entirely and always
|
||||
wait for full fidelity.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
from collections.abc import Awaitable, Callable
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
logger = logging.getLogger("headroom.proxy")
|
||||
|
||||
DEFAULT_BUFFERED_CCR_GRACE_SECONDS = 5.0
|
||||
"""Seconds to hold out for full status fidelity before committing to SSE.
|
||||
|
||||
Comfortably under any client stream-idle watchdog, while still covering the
|
||||
fast failures whose status matters most.
|
||||
"""
|
||||
|
||||
_HEARTBEAT_INTERVAL_SECONDS = 0.25
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class BufferedCCRErrorFormat:
|
||||
"""The provider-specific error shapes this wrapper has to emit.
|
||||
|
||||
Anthropic and OpenAI disagree on both the JSON envelope and the SSE framing,
|
||||
and the difference is pure formatting — the decision logic above is shared.
|
||||
"""
|
||||
|
||||
provider: str
|
||||
#: Build the pre-commit JSON body for a 502.
|
||||
json_body: Callable[[str], bytes]
|
||||
#: Build a post-commit SSE error event from an error type and message.
|
||||
sse_event: Callable[[str, str], bytes]
|
||||
#: Map an upstream HTTP status onto this provider's error type string.
|
||||
error_type_for_status: Callable[[int], str]
|
||||
#: The keepalive frame to emit while waiting, once committed.
|
||||
heartbeat: bytes
|
||||
|
||||
|
||||
def _anthropic_error_type(status: int) -> str:
|
||||
# The wire types Anthropic documents; clients switch their retry behaviour
|
||||
# on these, so a 429 must not arrive labelled `api_error`.
|
||||
return {
|
||||
400: "invalid_request_error",
|
||||
401: "authentication_error",
|
||||
403: "permission_error",
|
||||
404: "not_found_error",
|
||||
413: "request_too_large",
|
||||
429: "rate_limit_error",
|
||||
529: "overloaded_error",
|
||||
}.get(status, "api_error")
|
||||
|
||||
|
||||
def _openai_error_type(status: int) -> str:
|
||||
return {
|
||||
400: "invalid_request_error",
|
||||
401: "authentication_error",
|
||||
403: "permission_error",
|
||||
404: "not_found_error",
|
||||
413: "invalid_request_error",
|
||||
429: "rate_limit_error",
|
||||
}.get(status, "server_error")
|
||||
|
||||
|
||||
ANTHROPIC_ERROR_FORMAT = BufferedCCRErrorFormat(
|
||||
provider="anthropic",
|
||||
json_body=lambda message: json.dumps(
|
||||
{"type": "error", "error": {"type": "api_error", "message": message}}
|
||||
).encode(),
|
||||
sse_event=lambda error_type, message: (
|
||||
"event: error\ndata: "
|
||||
+ json.dumps({"type": "error", "error": {"type": error_type, "message": message}})
|
||||
+ "\n\n"
|
||||
).encode(),
|
||||
error_type_for_status=_anthropic_error_type,
|
||||
# Anthropic's stream carries a real `ping` event type.
|
||||
heartbeat=b'event: ping\ndata: {"type":"ping"}\n\n',
|
||||
)
|
||||
|
||||
OPENAI_ERROR_FORMAT = BufferedCCRErrorFormat(
|
||||
provider="openai",
|
||||
json_body=lambda message: json.dumps(
|
||||
{"error": {"message": message, "type": "server_error", "code": "proxy_error"}}
|
||||
).encode(),
|
||||
sse_event=lambda error_type, message: (
|
||||
"data: " + json.dumps({"error": {"message": message, "type": error_type}}) + "\n\n"
|
||||
).encode(),
|
||||
error_type_for_status=_openai_error_type,
|
||||
# OpenAI has no ping event; an SSE comment keeps the socket warm without
|
||||
# putting a frame the client would try to parse on the wire.
|
||||
heartbeat=b": ping\n\n",
|
||||
)
|
||||
|
||||
_GENERIC_FAILURE_MESSAGE = "An error occurred while processing your request. Please try again."
|
||||
|
||||
|
||||
def _upstream_message(body: bytes | None, fallback: str) -> str:
|
||||
"""Prefer the upstream's own error text over a synthesized one.
|
||||
|
||||
A committed stream cannot carry the upstream status, so the message is the
|
||||
only place its detail survives. Falls back whenever the body is not a
|
||||
recognizable error envelope.
|
||||
"""
|
||||
|
||||
if not body:
|
||||
return fallback
|
||||
try:
|
||||
parsed = json.loads(body)
|
||||
except (ValueError, TypeError):
|
||||
return fallback
|
||||
if not isinstance(parsed, dict):
|
||||
return fallback
|
||||
error = parsed.get("error")
|
||||
if isinstance(error, dict):
|
||||
message = error.get("message")
|
||||
if isinstance(message, str) and message.strip():
|
||||
return message.strip()
|
||||
message = parsed.get("message")
|
||||
if isinstance(message, str) and message.strip():
|
||||
return message.strip()
|
||||
return fallback
|
||||
|
||||
|
||||
async def _send_committed_failure(
|
||||
send: Callable[[dict[str, Any]], Awaitable[None]],
|
||||
fmt: BufferedCCRErrorFormat,
|
||||
*,
|
||||
status: int,
|
||||
body: bytes | None,
|
||||
) -> None:
|
||||
"""Emit a typed SSE error for a failure that arrived after the commit."""
|
||||
|
||||
error_type = fmt.error_type_for_status(status)
|
||||
message = _upstream_message(body, _GENERIC_FAILURE_MESSAGE)
|
||||
await send(
|
||||
{
|
||||
"type": "http.response.body",
|
||||
"body": fmt.sse_event(error_type, message),
|
||||
"more_body": False,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
async def _forward_after_commit(
|
||||
result: Any,
|
||||
send: Callable[[dict[str, Any]], Awaitable[None]],
|
||||
fmt: BufferedCCRErrorFormat,
|
||||
*,
|
||||
request_id: str,
|
||||
) -> None:
|
||||
"""Relay a resolved result once SSE headers are already on the wire."""
|
||||
|
||||
status = int(getattr(result, "status_code", 200) or 200)
|
||||
body_iterator = getattr(result, "body_iterator", None)
|
||||
|
||||
if body_iterator is not None:
|
||||
# Streaming results are already SSE — both the success path and the
|
||||
# 502 CCR-failure paths, whose bodies are error events. Forwarding the
|
||||
# bytes keeps whatever detail they carry.
|
||||
async for chunk in body_iterator:
|
||||
await send({"type": "http.response.body", "body": chunk, "more_body": True})
|
||||
await send({"type": "http.response.body", "body": b"", "more_body": False})
|
||||
return
|
||||
|
||||
# A non-streaming result here is an upstream reply that never became SSE.
|
||||
# Its status cannot reach the client anymore, so preserve its meaning in
|
||||
# the error type instead of degrading to a generic failure (#3079).
|
||||
body = getattr(result, "body", None)
|
||||
if status != 200:
|
||||
logger.warning(
|
||||
f"[{request_id}] CCR: buffered upstream returned {status} after the response "
|
||||
"was committed as SSE; relaying it as a typed stream error"
|
||||
)
|
||||
await _send_committed_failure(send, fmt, status=status, body=body)
|
||||
|
||||
|
||||
def buffered_ccr_asgi_call(
|
||||
*,
|
||||
operation: asyncio.Task,
|
||||
fmt: BufferedCCRErrorFormat,
|
||||
grace_seconds: float,
|
||||
record_failed: Callable[..., Awaitable[None]],
|
||||
request_id: str,
|
||||
) -> Callable[[Any, Any, Any], Awaitable[None]]:
|
||||
"""Build the ``__call__`` body for a buffered-CCR ASGI response.
|
||||
|
||||
Returned rather than subclassed so both handlers can keep their existing
|
||||
``Response`` shells and differ only in the error format they pass in.
|
||||
"""
|
||||
|
||||
async def __call__(scope: Any, receive: Any, send: Any) -> None: # noqa: ANN401
|
||||
loop = asyncio.get_running_loop()
|
||||
committed = False
|
||||
deadline = loop.time() + grace_seconds
|
||||
heartbeat_enabled = grace_seconds > 0
|
||||
try:
|
||||
while True:
|
||||
if heartbeat_enabled:
|
||||
timeout = (
|
||||
_HEARTBEAT_INTERVAL_SECONDS
|
||||
if committed
|
||||
else max(0.0, deadline - loop.time())
|
||||
)
|
||||
else:
|
||||
timeout = None
|
||||
done, _pending = await asyncio.wait({operation}, timeout=timeout)
|
||||
|
||||
if done:
|
||||
try:
|
||||
result = operation.result()
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except Exception as exc:
|
||||
await record_failed(provider=fmt.provider)
|
||||
logger.error(f"[{request_id}] Request failed: {type(exc).__name__}: {exc}")
|
||||
if committed:
|
||||
await _send_committed_failure(send, fmt, status=502, body=None)
|
||||
return
|
||||
await send(
|
||||
{
|
||||
"type": "http.response.start",
|
||||
"status": 502,
|
||||
"headers": [(b"content-type", b"application/json")],
|
||||
}
|
||||
)
|
||||
await send(
|
||||
{
|
||||
"type": "http.response.body",
|
||||
"body": fmt.json_body(_GENERIC_FAILURE_MESSAGE),
|
||||
"more_body": False,
|
||||
}
|
||||
)
|
||||
return
|
||||
|
||||
if not committed:
|
||||
# Nothing is on the wire yet, so the result keeps its
|
||||
# own status and headers. This is the fidelity #2997
|
||||
# was protecting.
|
||||
await result(scope, receive, send)
|
||||
return
|
||||
|
||||
await _forward_after_commit(result, send, fmt, request_id=request_id)
|
||||
return
|
||||
|
||||
if not committed:
|
||||
# The grace window expired without an outcome. Commit now so
|
||||
# a first byte beats the client's stream-idle watchdog.
|
||||
await send(
|
||||
{
|
||||
"type": "http.response.start",
|
||||
"status": 200,
|
||||
"headers": [(b"content-type", b"text/event-stream")],
|
||||
}
|
||||
)
|
||||
committed = True
|
||||
logger.debug(
|
||||
f"[{request_id}] CCR: buffered turn exceeded the {grace_seconds}s "
|
||||
"grace window; committing SSE and starting the heartbeat"
|
||||
)
|
||||
await send({"type": "http.response.body", "body": fmt.heartbeat, "more_body": True})
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
finally:
|
||||
if not operation.done():
|
||||
operation.cancel()
|
||||
try:
|
||||
await operation
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
return __call__
|
||||
@@ -34,6 +34,11 @@ from headroom.proxy.auth_mode import (
|
||||
classify_client,
|
||||
supports_mid_turn_coalescing,
|
||||
)
|
||||
from headroom.proxy.buffered_ccr_response import (
|
||||
ANTHROPIC_ERROR_FORMAT,
|
||||
DEFAULT_BUFFERED_CCR_GRACE_SECONDS,
|
||||
buffered_ccr_asgi_call,
|
||||
)
|
||||
from headroom.proxy.compression_decision import CompressionDecision
|
||||
from headroom.proxy.forwarded_headers import resolve_client_ip
|
||||
from headroom.proxy.handlers._debug_dump import _debug_dump_mode, _redact_debug_value
|
||||
@@ -4415,54 +4420,26 @@ class AnthropicHandlerMixin:
|
||||
|
||||
if buffered_stream_ccr:
|
||||
operation = asyncio.create_task(_buffered_ccr_operation())
|
||||
record_failed = self.metrics.record_failed
|
||||
|
||||
# Holds out for the real status, then keeps the stream alive
|
||||
# once waiting silently would risk the client's idle
|
||||
# watchdog. Both halves live in one shared place so the
|
||||
# OpenAI twin cannot drift from it (#3079).
|
||||
_buffered_call = buffered_ccr_asgi_call(
|
||||
operation=operation,
|
||||
fmt=ANTHROPIC_ERROR_FORMAT,
|
||||
grace_seconds=getattr(
|
||||
self.config,
|
||||
"buffered_ccr_grace_seconds",
|
||||
DEFAULT_BUFFERED_CCR_GRACE_SECONDS,
|
||||
),
|
||||
record_failed=self.metrics.record_failed,
|
||||
request_id=request_id,
|
||||
)
|
||||
|
||||
class _BufferedCCRResponse(Response):
|
||||
async def __call__(self, scope, receive, send): # noqa: ANN001
|
||||
# Send nothing until the buffered operation resolves.
|
||||
# The previous keepalive preamble committed
|
||||
# `200 text/event-stream` after 1s, i.e. before the
|
||||
# outcome was known: any upstream reply that then
|
||||
# failed to become SSE (non-200, unparseable body)
|
||||
# reached the client as a 200 whose body carried no
|
||||
# `message_start`, which Claude Code reports as "API
|
||||
# returned an empty or malformed response (HTTP 200) —
|
||||
# check for a proxy or gateway intercepting the
|
||||
# request". The real status was lost with it, so
|
||||
# client-side 429/5xx backoff never fired. Clients
|
||||
# budget minutes for a turn (Claude Code sends
|
||||
# `x-stainless-timeout: 600`), so waiting is free.
|
||||
try:
|
||||
result = await operation
|
||||
except Exception as e:
|
||||
await record_failed(provider=provider_name)
|
||||
logger.error(
|
||||
f"[{request_id}] Request failed: {type(e).__name__}: {e}"
|
||||
)
|
||||
await send(
|
||||
{
|
||||
"type": "http.response.start",
|
||||
"status": 502,
|
||||
"headers": [(b"content-type", b"application/json")],
|
||||
}
|
||||
)
|
||||
await send(
|
||||
{
|
||||
"type": "http.response.body",
|
||||
"body": json.dumps(
|
||||
{
|
||||
"type": "error",
|
||||
"error": {
|
||||
"type": "api_error",
|
||||
"message": "An error occurred while processing your request. Please try again.",
|
||||
},
|
||||
}
|
||||
).encode(),
|
||||
"more_body": False,
|
||||
}
|
||||
)
|
||||
return
|
||||
await result(scope, receive, send)
|
||||
await _buffered_call(scope, receive, send)
|
||||
|
||||
return _BufferedCCRResponse(media_type="text/event-stream")
|
||||
return await _buffered_ccr_operation()
|
||||
|
||||
@@ -71,6 +71,11 @@ from headroom.proxy.auth_mode import (
|
||||
classify_client,
|
||||
should_stamp_codex_client,
|
||||
)
|
||||
from headroom.proxy.buffered_ccr_response import (
|
||||
DEFAULT_BUFFERED_CCR_GRACE_SECONDS,
|
||||
OPENAI_ERROR_FORMAT,
|
||||
buffered_ccr_asgi_call,
|
||||
)
|
||||
from headroom.proxy.compression_decision import CompressionDecision
|
||||
from headroom.proxy.cost import header_safe_transforms
|
||||
from headroom.proxy.handlers._debug_dump import _debug_dump_mode, _redact_debug_value
|
||||
@@ -6207,48 +6212,24 @@ class OpenAIHandlerMixin:
|
||||
|
||||
if buffered_stream_ccr:
|
||||
operation = asyncio.create_task(_buffered_ccr_operation())
|
||||
record_failed = self.metrics.record_failed
|
||||
|
||||
# Same wrapper as the Anthropic twin; only the error wire format
|
||||
# differs. See headroom/proxy/buffered_ccr_response.py (#3079).
|
||||
_buffered_call = buffered_ccr_asgi_call(
|
||||
operation=operation,
|
||||
fmt=OPENAI_ERROR_FORMAT,
|
||||
grace_seconds=getattr(
|
||||
self.config,
|
||||
"buffered_ccr_grace_seconds",
|
||||
DEFAULT_BUFFERED_CCR_GRACE_SECONDS,
|
||||
),
|
||||
record_failed=self.metrics.record_failed,
|
||||
request_id=request_id,
|
||||
)
|
||||
|
||||
class _BufferedCCRResponse(Response):
|
||||
async def __call__(self, scope, receive, send): # noqa: ANN001
|
||||
# Send nothing until the buffered operation resolves —
|
||||
# see the AnthropicHandler twin for the full rationale.
|
||||
# Committing `200 text/event-stream` on a keepalive timer,
|
||||
# before the outcome is known, turns every non-200 or
|
||||
# unparseable upstream reply into a 200 with no usable
|
||||
# body and discards the status the client needs to back
|
||||
# off on.
|
||||
try:
|
||||
result = await operation
|
||||
except Exception as e:
|
||||
await record_failed(provider="openai")
|
||||
logger.error(
|
||||
f"[{request_id}] OpenAI responses request failed: {type(e).__name__}: {e}"
|
||||
)
|
||||
await send(
|
||||
{
|
||||
"type": "http.response.start",
|
||||
"status": 502,
|
||||
"headers": [(b"content-type", b"application/json")],
|
||||
}
|
||||
)
|
||||
await send(
|
||||
{
|
||||
"type": "http.response.body",
|
||||
"body": json.dumps(
|
||||
{
|
||||
"error": {
|
||||
"message": "An error occurred while processing your request. Please try again.",
|
||||
"type": "server_error",
|
||||
"code": "proxy_error",
|
||||
}
|
||||
}
|
||||
).encode(),
|
||||
"more_body": False,
|
||||
}
|
||||
)
|
||||
return
|
||||
await result(scope, receive, send)
|
||||
await _buffered_call(scope, receive, send)
|
||||
|
||||
return _BufferedCCRResponse(media_type="text/event-stream")
|
||||
return await _buffered_ccr_operation()
|
||||
|
||||
@@ -14,6 +14,7 @@ from typing import Any, Literal
|
||||
|
||||
from headroom.memory import qdrant_env
|
||||
from headroom.providers.registry import ProviderApiOverrides
|
||||
from headroom.proxy.buffered_ccr_response import DEFAULT_BUFFERED_CCR_GRACE_SECONDS
|
||||
from headroom.proxy.model_router import ModelRouterConfig
|
||||
from headroom.rollout import RolloutSnapshot, resolve_rollout
|
||||
|
||||
@@ -370,6 +371,12 @@ class ProxyConfig:
|
||||
# Anthropic buffered reads can legitimately run longer than the generic
|
||||
# proxy request cap. Keep the generic timeout unchanged elsewhere.
|
||||
anthropic_buffered_request_timeout_seconds: int = 600
|
||||
# How long a buffered-CCR turn holds out for full status fidelity before it
|
||||
# commits to SSE and starts a keepalive. Under the window, failures keep
|
||||
# their real HTTP status; past it, the client gets a first byte before its
|
||||
# stream-idle watchdog fires. 0 or less disables the keepalive entirely.
|
||||
# See headroom/proxy/buffered_ccr_response.py (#3079).
|
||||
buffered_ccr_grace_seconds: float = DEFAULT_BUFFERED_CCR_GRACE_SECONDS
|
||||
|
||||
# Connection pool
|
||||
max_connections: int = 500
|
||||
|
||||
@@ -114,6 +114,7 @@ from headroom.proxy.audit import is_auditable_path, record_admin_action
|
||||
from headroom.proxy.auth_mode import should_stamp_codex_client
|
||||
from headroom.proxy.background_compression import BackgroundCompressor
|
||||
from headroom.proxy.budget_basis_policy import resolve_estimated_basis_policy
|
||||
from headroom.proxy.buffered_ccr_response import DEFAULT_BUFFERED_CCR_GRACE_SECONDS
|
||||
|
||||
# =============================================================================
|
||||
# Extracted modules (re-exported for backward compatibility)
|
||||
@@ -5208,6 +5209,10 @@ def _proxy_config_from_env() -> ProxyConfig:
|
||||
600,
|
||||
min_value=1,
|
||||
),
|
||||
buffered_ccr_grace_seconds=_get_env_float(
|
||||
"HEADROOM_BUFFERED_CCR_GRACE_SECONDS",
|
||||
DEFAULT_BUFFERED_CCR_GRACE_SECONDS,
|
||||
),
|
||||
vertex_api_url=os.environ.get("VERTEX_TARGET_API_URL"),
|
||||
backend=_get_env_str("HEADROOM_BACKEND", "anthropic"),
|
||||
bedrock_region=_get_env_str("HEADROOM_BEDROCK_REGION", "us-west-2"),
|
||||
|
||||
@@ -0,0 +1,314 @@
|
||||
"""The buffered-CCR grace window: keep status fidelity *and* liveness (#3079).
|
||||
|
||||
A buffered CCR turn holds a request open for the whole generation. Two things
|
||||
have to be true across that window, and the history here is of each being fixed
|
||||
by breaking the other:
|
||||
|
||||
* #2465 → #2479 added a keepalive so the client's *stream-idle* watchdog saw a
|
||||
first byte.
|
||||
* #2997 removed that keepalive, because committing ``200 text/event-stream``
|
||||
before the outcome was known destroyed the real status: a later 429 reached
|
||||
the client as a 200 with no ``message_start`` and no ``retry-after``, so
|
||||
client backoff never fired.
|
||||
* Removing it put #2465's condition back — zero bytes for 15-25s turns.
|
||||
|
||||
The grace window keeps both: full fidelity while there is still time to send a
|
||||
real status, a heartbeat once there is not, and a *typed* stream error for a
|
||||
failure that lands after the commit so backoff still works.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
|
||||
import pytest
|
||||
|
||||
from headroom.proxy.buffered_ccr_response import (
|
||||
ANTHROPIC_ERROR_FORMAT,
|
||||
OPENAI_ERROR_FORMAT,
|
||||
buffered_ccr_asgi_call,
|
||||
)
|
||||
|
||||
|
||||
class _Recorder:
|
||||
"""Collects ASGI messages, exposing when the response was committed."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.messages: list[dict] = []
|
||||
self.committed = asyncio.Event()
|
||||
|
||||
async def send(self, message: dict) -> None:
|
||||
self.messages.append(message)
|
||||
if message["type"] == "http.response.start":
|
||||
self.committed.set()
|
||||
|
||||
@property
|
||||
def status(self) -> int | None:
|
||||
for m in self.messages:
|
||||
if m["type"] == "http.response.start":
|
||||
return m["status"]
|
||||
return None
|
||||
|
||||
@property
|
||||
def headers(self) -> dict[str, str]:
|
||||
for m in self.messages:
|
||||
if m["type"] == "http.response.start":
|
||||
return {k.decode(): v.decode() for k, v in m.get("headers", [])}
|
||||
return {}
|
||||
|
||||
@property
|
||||
def body(self) -> bytes:
|
||||
return b"".join(
|
||||
m.get("body", b"") for m in self.messages if m["type"] == "http.response.body"
|
||||
)
|
||||
|
||||
@property
|
||||
def pings(self) -> int:
|
||||
return self.body.count(b"ping")
|
||||
|
||||
|
||||
class _FakeResult:
|
||||
"""Stands in for a resolved buffered result (a Starlette response)."""
|
||||
|
||||
def __init__(self, status: int, body: bytes, headers: dict[str, str] | None = None) -> None:
|
||||
self.status_code = status
|
||||
self.body = body
|
||||
self._headers = headers or {}
|
||||
|
||||
async def __call__(self, scope, receive, send) -> None: # noqa: ANN001
|
||||
await send(
|
||||
{
|
||||
"type": "http.response.start",
|
||||
"status": self.status_code,
|
||||
"headers": [(k.encode(), v.encode()) for k, v in self._headers.items()],
|
||||
}
|
||||
)
|
||||
await send({"type": "http.response.body", "body": self.body, "more_body": False})
|
||||
|
||||
|
||||
class _FakeStreamingResult:
|
||||
"""A resolved result that is already SSE, like the success path."""
|
||||
|
||||
def __init__(self, chunks: list[bytes], status: int = 200) -> None:
|
||||
self.status_code = status
|
||||
self._chunks = chunks
|
||||
|
||||
@property
|
||||
def body_iterator(self): # noqa: ANN201
|
||||
async def _gen():
|
||||
for chunk in self._chunks:
|
||||
yield chunk
|
||||
|
||||
return _gen()
|
||||
|
||||
async def __call__(self, scope, receive, send) -> None: # noqa: ANN001
|
||||
await send(
|
||||
{
|
||||
"type": "http.response.start",
|
||||
"status": self.status_code,
|
||||
"headers": [(b"content-type", b"text/event-stream")],
|
||||
}
|
||||
)
|
||||
for chunk in self._chunks:
|
||||
await send({"type": "http.response.body", "body": chunk, "more_body": True})
|
||||
await send({"type": "http.response.body", "body": b"", "more_body": False})
|
||||
|
||||
|
||||
async def _drive(*, produce, grace: float, fmt=ANTHROPIC_ERROR_FORMAT) -> _Recorder:
|
||||
"""Run the wrapper against a coroutine standing in for the buffered work."""
|
||||
recorder = _Recorder()
|
||||
failures: list[str] = []
|
||||
|
||||
async def record_failed(provider: str) -> None:
|
||||
failures.append(provider)
|
||||
|
||||
operation = asyncio.create_task(produce())
|
||||
call = buffered_ccr_asgi_call(
|
||||
operation=operation,
|
||||
fmt=fmt,
|
||||
grace_seconds=grace,
|
||||
record_failed=record_failed,
|
||||
request_id="req-test",
|
||||
)
|
||||
await call({"type": "http"}, None, recorder.send)
|
||||
recorder.failures = failures # type: ignore[attr-defined]
|
||||
return recorder
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Property 1 — fidelity: anything inside the window keeps its real status
|
||||
# --------------------------------------------------------------------------- #
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_failure_inside_the_window_keeps_its_status_and_headers() -> None:
|
||||
"""What #2997 bought. A 429 must arrive as a 429, with retry-after."""
|
||||
|
||||
async def produce():
|
||||
await asyncio.sleep(0.05)
|
||||
return _FakeResult(
|
||||
429,
|
||||
b'{"type":"error","error":{"type":"rate_limit_error","message":"slow down"}}',
|
||||
{"content-type": "application/json", "retry-after": "30"},
|
||||
)
|
||||
|
||||
rec = await _drive(produce=produce, grace=5.0)
|
||||
|
||||
assert rec.status == 429
|
||||
assert rec.headers.get("retry-after") == "30"
|
||||
assert rec.pings == 0, "nothing should have been sent before the outcome was known"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_success_inside_the_window_is_relayed_untouched() -> None:
|
||||
async def produce():
|
||||
await asyncio.sleep(0.05)
|
||||
return _FakeStreamingResult([b'event: message_start\ndata: {"x":1}\n\n'])
|
||||
|
||||
rec = await _drive(produce=produce, grace=5.0)
|
||||
|
||||
assert rec.status == 200
|
||||
assert b"message_start" in rec.body
|
||||
assert rec.pings == 0
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Property 2 — liveness: a slow turn gets a first byte before the ceiling
|
||||
# --------------------------------------------------------------------------- #
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_slow_success_produces_a_first_byte_before_the_ceiling() -> None:
|
||||
"""What #2479 bought, and what #2997 removed (#3079).
|
||||
|
||||
The client's stream-idle watchdog cannot distinguish "still generating"
|
||||
from "dead socket", so silence for the whole generation is what tripped it.
|
||||
"""
|
||||
release = asyncio.Event()
|
||||
|
||||
async def produce():
|
||||
await release.wait()
|
||||
return _FakeStreamingResult([b'event: message_start\ndata: {"x":1}\n\n'])
|
||||
|
||||
recorder = _Recorder()
|
||||
|
||||
async def record_failed(provider: str) -> None: # pragma: no cover - not hit
|
||||
raise AssertionError("a success must not be recorded as a failure")
|
||||
|
||||
operation = asyncio.create_task(produce())
|
||||
call = buffered_ccr_asgi_call(
|
||||
operation=operation,
|
||||
fmt=ANTHROPIC_ERROR_FORMAT,
|
||||
grace_seconds=0.1,
|
||||
record_failed=record_failed,
|
||||
request_id="req-slow",
|
||||
)
|
||||
driver = asyncio.create_task(call({"type": "http"}, None, recorder.send))
|
||||
|
||||
# The first byte must arrive from the heartbeat alone, well before the
|
||||
# upstream resolves.
|
||||
await asyncio.wait_for(recorder.committed.wait(), timeout=2.0)
|
||||
assert recorder.status == 200
|
||||
assert recorder.headers["content-type"] == "text/event-stream"
|
||||
assert recorder.pings >= 1
|
||||
|
||||
release.set()
|
||||
await asyncio.wait_for(driver, timeout=2.0)
|
||||
assert b"message_start" in recorder.body
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Property 3 — a failure past the commit stays actionable
|
||||
# --------------------------------------------------------------------------- #
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_late_rate_limit_survives_as_a_typed_stream_error() -> None:
|
||||
"""The status is gone once committed, so the *meaning* must survive.
|
||||
|
||||
Degrading every post-commit failure to a generic ``api_error`` is what made
|
||||
early commits harmful: the client cannot tell a rate limit from a bug, so
|
||||
it does not back off. A typed error keeps that behaviour reachable.
|
||||
"""
|
||||
|
||||
async def produce():
|
||||
await asyncio.sleep(0.3)
|
||||
return _FakeResult(
|
||||
429,
|
||||
b'{"type":"error","error":{"type":"rate_limit_error","message":"slow down"}}',
|
||||
{"content-type": "application/json", "retry-after": "30"},
|
||||
)
|
||||
|
||||
rec = await _drive(produce=produce, grace=0.05)
|
||||
|
||||
assert rec.status == 200, "already committed; the status can no longer change"
|
||||
assert rec.pings >= 1
|
||||
payload = json.loads(rec.body.split(b"event: error\ndata: ")[-1].strip())
|
||||
assert payload["error"]["type"] == "rate_limit_error"
|
||||
assert payload["error"]["message"] == "slow down"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_late_overload_maps_to_the_overloaded_type() -> None:
|
||||
async def produce():
|
||||
await asyncio.sleep(0.3)
|
||||
return _FakeResult(529, b"", {})
|
||||
|
||||
rec = await _drive(produce=produce, grace=0.05)
|
||||
|
||||
payload = json.loads(rec.body.split(b"event: error\ndata: ")[-1].strip())
|
||||
assert payload["error"]["type"] == "overloaded_error"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_late_exception_is_reported_on_the_committed_stream() -> None:
|
||||
async def produce():
|
||||
await asyncio.sleep(0.3)
|
||||
raise RuntimeError("upstream exploded")
|
||||
|
||||
rec = await _drive(produce=produce, grace=0.05)
|
||||
|
||||
assert rec.status == 200
|
||||
assert b"event: error" in rec.body
|
||||
assert rec.failures == ["anthropic"] # type: ignore[attr-defined]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_an_early_exception_still_gets_a_real_502() -> None:
|
||||
async def produce():
|
||||
raise RuntimeError("upstream exploded")
|
||||
|
||||
rec = await _drive(produce=produce, grace=5.0)
|
||||
|
||||
assert rec.status == 502
|
||||
assert json.loads(rec.body)["error"]["type"] == "api_error"
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# The escape hatch
|
||||
# --------------------------------------------------------------------------- #
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_zero_grace_window_never_commits_early() -> None:
|
||||
"""Operators who want #2997's behaviour verbatim can still have it."""
|
||||
|
||||
async def produce():
|
||||
await asyncio.sleep(0.3)
|
||||
return _FakeResult(429, b"{}", {"retry-after": "30"})
|
||||
|
||||
rec = await _drive(produce=produce, grace=0)
|
||||
|
||||
assert rec.status == 429
|
||||
assert rec.headers.get("retry-after") == "30"
|
||||
assert rec.pings == 0
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# OpenAI wire format
|
||||
# --------------------------------------------------------------------------- #
|
||||
@pytest.mark.asyncio
|
||||
async def test_openai_late_failures_use_the_openai_error_shape() -> None:
|
||||
async def produce():
|
||||
await asyncio.sleep(0.3)
|
||||
return _FakeResult(429, b'{"error":{"message":"too many"}}', {})
|
||||
|
||||
rec = await _drive(produce=produce, grace=0.05, fmt=OPENAI_ERROR_FORMAT)
|
||||
|
||||
assert b": ping" in rec.body, "OpenAI keepalives are SSE comments, not ping events"
|
||||
payload = json.loads(rec.body.split(b"data: ")[-1].strip())
|
||||
assert payload["error"]["type"] == "rate_limit_error"
|
||||
assert payload["error"]["message"] == "too many"
|
||||
Reference in New Issue
Block a user