mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
fix(proxy): preserve streaming passthrough beta headers (#1783)
## Description Anthropic-compatible custom upstreams can reject streaming passthrough requests when Headroom expands the client's `anthropic-beta` header with sticky session tokens. The request body is still forwarded byte-faithfully, but the header no longer matches the direct request that succeeds against the same upstream. This change keeps sticky beta learning intact while preserving the direct client beta header for the custom-upstream streaming passthrough path that owns the 503. Closes #1724 ## 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 - Preserves client `anthropic-beta` headers on Vertex `:streamRawPredict` and custom Anthropic API URL streaming passthrough requests. - Keeps sticky beta tracking and adjacent sticky-header behavior for non-hazard paths. - Adds focused regression coverage that captures outgoing streaming headers and preserves existing byte-faithful body checks. ## Testing - [x] Unit tests pass (`uv run pytest tests/test_proxy_byte_faithful_forwarding.py tests/test_anthropic_beta_session_sticky.py -q`) - [x] Linting passes (`uvx ruff==0.15.17 check .` and `uvx ruff==0.15.17 format --check .`) - [ ] Type checking passes (`uv run mypy headroom`) - [x] New tests added for new functionality - [ ] Manual testing performed ### Test Output ```text pytest base: exit_code=1, stdout excerpt: AssertionError: assert 'sticky-beta-2024-01-01,claude-code-20250219' == 'claude-code-20250219' pytest head: exit_code=0, stdout excerpt: 64 passed, 1 warning in 3.57s ruff: exit_code=0, stdout excerpt: All checks passed! / 1044 files already formatted ``` ## Real Behavior Proof - Environment: Windows, focused local proxy tests through the headless runner. - Exact command / steps: Pre-seed sticky beta state, send streaming Vertex `:streamRawPredict` and `/v1/messages` requests through custom upstream routing with `anthropic-beta: claude-code-20250219`, and capture the outgoing request headers. Run the same focused pytest command on the base checkout, then on the fixed checkout. Run pinned Ruff 0.15.17 check and format validation against the final branch. - Observed result: The base checkout expands the streaming custom-upstream beta header, and the fixed checkout preserves the direct client beta header for both streaming routes while adjacent non-streaming custom-upstream requests still carry the sticky union. - Not tested: The reporter's live MaaS upstream and the full test suite. ## 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 - [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 - [ ] I have updated the CHANGELOG.md if applicable ## Additional Notes Documentation is left unchecked because the fix preserves the existing passthrough contract rather than adding a new user-facing option. The changelog box is left unchecked because Headroom generates changelog entries from conventional commits.
This commit is contained in:
Rod Boev
GitHub
parent
be51008c70
commit
0f553a8ebb
2 changed files with 163 additions and 5 deletions
|
|
@ -1055,6 +1055,7 @@ class AnthropicHandlerMixin:
|
|||
headroom_added=[],
|
||||
request_id=request_id,
|
||||
)
|
||||
_headroom_beta_added = False
|
||||
|
||||
# In cache mode, avoid rewriting any message body bytes. The latest user
|
||||
# turn becomes historical on the next request, so even "latest turn only"
|
||||
|
|
@ -1955,6 +1956,7 @@ class AnthropicHandlerMixin:
|
|||
continue
|
||||
existing_value = headers.get(key, "")
|
||||
required_tokens = [t.strip() for t in value.split(",") if t.strip()]
|
||||
_headroom_beta_added = True
|
||||
merged = merge_anthropic_beta(existing_value, required_tokens)
|
||||
_existing_count = (
|
||||
len([t for t in existing_value.split(",") if t.strip()])
|
||||
|
|
@ -2118,6 +2120,17 @@ class AnthropicHandlerMixin:
|
|||
except (json.JSONDecodeError, ValueError):
|
||||
body_mutation_tracker.mark_mutated("original_unparseable")
|
||||
|
||||
if (
|
||||
(upstream_base_url or self.ANTHROPIC_API_URL != "https://api.anthropic.com")
|
||||
and stream
|
||||
and _client_beta_value
|
||||
and _sticky_beta_value
|
||||
and _sticky_beta_value != _client_beta_value
|
||||
and not body_mutation_tracker.mutated
|
||||
and not _headroom_beta_added
|
||||
):
|
||||
headers["anthropic-beta"] = _client_beta_value
|
||||
|
||||
# Forward request - use Bedrock backend if configured, otherwise direct API
|
||||
if self.anthropic_backend is not None:
|
||||
# Route through Bedrock backend
|
||||
|
|
|
|||
|
|
@ -32,8 +32,10 @@ from fastapi.testclient import TestClient
|
|||
from headroom.pipeline import PipelineStage
|
||||
from headroom.proxy.helpers import (
|
||||
BodyMutationTracker,
|
||||
_reset_session_beta_tracker_for_test,
|
||||
append_text_to_latest_user_chat_message,
|
||||
get_python_forwarder_mode,
|
||||
get_session_beta_tracker,
|
||||
log_outbound_request,
|
||||
prepare_outbound_body_bytes,
|
||||
serialize_body_canonical,
|
||||
|
|
@ -928,17 +930,19 @@ class _StreamingCapturingTransport(httpx.AsyncBaseTransport):
|
|||
self.captured_body = body
|
||||
self.captured_headers = dict(request.headers.items())
|
||||
|
||||
async def _empty_sse(): # pragma: no cover - generator
|
||||
yield b'event: message_start\ndata: {"type":"message_start","message":{"id":"msg_s","type":"message","role":"assistant","model":"claude","usage":{"input_tokens":1,"output_tokens":0,"cache_read_input_tokens":0,"cache_creation_input_tokens":0}}}\n\n'
|
||||
yield b'event: message_stop\ndata: {"type":"message_stop"}\n\n'
|
||||
|
||||
return httpx.Response(
|
||||
200,
|
||||
headers={"content-type": "text/event-stream"},
|
||||
stream=httpx.AsyncByteStream(_empty_sse()), # type: ignore[arg-type]
|
||||
stream=_SSEByteStream(),
|
||||
)
|
||||
|
||||
|
||||
class _SSEByteStream(httpx.AsyncByteStream):
|
||||
async def __aiter__(self):
|
||||
yield b'event: message_start\ndata: {"type":"message_start","message":{"id":"msg_s","type":"message","role":"assistant","model":"claude","usage":{"input_tokens":1,"output_tokens":0,"cache_read_input_tokens":0,"cache_creation_input_tokens":0}}}\n\n'
|
||||
yield b'event: message_stop\ndata: {"type":"message_stop"}\n\n'
|
||||
|
||||
|
||||
def test_streaming_forwarder_byte_faithful() -> None:
|
||||
"""Streaming forwarder uses the same byte-faithful path as non-streaming."""
|
||||
config = ProxyConfig(
|
||||
|
|
@ -992,6 +996,147 @@ def test_streaming_forwarder_byte_faithful() -> None:
|
|||
)
|
||||
|
||||
|
||||
def test_vertex_stream_rawpredict_preserves_client_beta_header_on_passthrough() -> None:
|
||||
_reset_session_beta_tracker_for_test()
|
||||
try:
|
||||
client, transport = _make_anthropic_app(optimize=False)
|
||||
get_session_beta_tracker().record_and_get_sticky_betas(
|
||||
provider="anthropic",
|
||||
session_id="s1",
|
||||
client_value="sticky-beta-2024-01-01",
|
||||
)
|
||||
|
||||
inbound_bytes = (
|
||||
b'{"model":"claude-sonnet-4-6","stream":true,'
|
||||
b'"messages":[{"role":"user","content":"hi"}]}'
|
||||
)
|
||||
client_beta = "claude-code-20250219"
|
||||
|
||||
with client.stream(
|
||||
"POST",
|
||||
"/projects/p/locations/us-central1/publishers/anthropic/models/"
|
||||
"claude-sonnet-4-6:streamRawPredict",
|
||||
headers={
|
||||
"x-api-key": "test-key",
|
||||
"x-headroom-session-id": "vertex-stream-beta-1",
|
||||
"anthropic-version": "2023-06-01",
|
||||
"anthropic-beta": client_beta,
|
||||
"content-type": "application/json",
|
||||
},
|
||||
content=inbound_bytes,
|
||||
) as resp:
|
||||
response_body = b"".join(resp.iter_bytes())
|
||||
assert resp.status_code == 200, response_body
|
||||
|
||||
assert transport.captured_body == inbound_bytes
|
||||
assert transport.captured_headers is not None
|
||||
captured_headers = {key.lower(): value for key, value in transport.captured_headers.items()}
|
||||
assert captured_headers["anthropic-beta"] == client_beta
|
||||
finally:
|
||||
_reset_session_beta_tracker_for_test()
|
||||
|
||||
|
||||
def test_messages_custom_upstream_stream_preserves_client_beta_header() -> None:
|
||||
_reset_session_beta_tracker_for_test()
|
||||
old_anthropic_url = None
|
||||
try:
|
||||
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,
|
||||
)
|
||||
app = create_app(config)
|
||||
proxy = app.state.proxy
|
||||
old_anthropic_url = type(proxy).ANTHROPIC_API_URL
|
||||
type(proxy).ANTHROPIC_API_URL = "https://custom.example"
|
||||
fake_tracker = _FakePrefixTracker(frozen_count=0)
|
||||
proxy.session_tracker_store.compute_session_id = lambda request, model, messages: (
|
||||
"custom-stream-beta-1"
|
||||
)
|
||||
proxy.session_tracker_store.get_or_create = lambda session_id, provider: fake_tracker
|
||||
|
||||
transport = _StreamingCapturingTransport()
|
||||
proxy.http_client = httpx.AsyncClient(transport=transport)
|
||||
client = TestClient(app)
|
||||
|
||||
get_session_beta_tracker().record_and_get_sticky_betas(
|
||||
provider="anthropic",
|
||||
session_id="custom-stream-beta-1",
|
||||
client_value="sticky-beta-2024-01-01",
|
||||
)
|
||||
|
||||
inbound_bytes = (
|
||||
b'{"model":"claude-sonnet-4-6","max_tokens":16,"stream":true,'
|
||||
b'"messages":[{"role":"user","content":"hi"}]}'
|
||||
)
|
||||
client_beta = "claude-code-20250219"
|
||||
|
||||
with client.stream(
|
||||
"POST",
|
||||
"/v1/messages",
|
||||
headers={
|
||||
"x-api-key": "test-key",
|
||||
"x-headroom-session-id": "custom-stream-beta-1",
|
||||
"anthropic-version": "2023-06-01",
|
||||
"anthropic-beta": client_beta,
|
||||
"content-type": "application/json",
|
||||
},
|
||||
content=inbound_bytes,
|
||||
) as resp:
|
||||
response_body = b"".join(resp.iter_bytes())
|
||||
assert resp.status_code == 200, response_body
|
||||
|
||||
assert transport.captured_body == inbound_bytes
|
||||
assert transport.captured_headers is not None
|
||||
captured_headers = {key.lower(): value for key, value in transport.captured_headers.items()}
|
||||
assert captured_headers["anthropic-beta"] == client_beta
|
||||
finally:
|
||||
if old_anthropic_url is not None:
|
||||
type(proxy).ANTHROPIC_API_URL = old_anthropic_url
|
||||
_reset_session_beta_tracker_for_test()
|
||||
|
||||
|
||||
def test_vertex_rawpredict_keeps_sticky_beta_union_on_non_stream_passthrough() -> None:
|
||||
_reset_session_beta_tracker_for_test()
|
||||
try:
|
||||
client, transport = _make_anthropic_app(optimize=False)
|
||||
get_session_beta_tracker().record_and_get_sticky_betas(
|
||||
provider="anthropic",
|
||||
session_id="s1",
|
||||
client_value="sticky-beta-2024-01-01",
|
||||
)
|
||||
|
||||
inbound_bytes = b'{"model":"claude-sonnet-4-6","messages":[{"role":"user","content":"hi"}]}'
|
||||
client_beta = "claude-code-20250219"
|
||||
|
||||
response = client.post(
|
||||
"/projects/p/locations/us-central1/publishers/anthropic/models/"
|
||||
"claude-sonnet-4-6:rawPredict",
|
||||
headers={
|
||||
"x-api-key": "test-key",
|
||||
"x-headroom-session-id": "vertex-raw-beta-1",
|
||||
"anthropic-version": "2023-06-01",
|
||||
"anthropic-beta": client_beta,
|
||||
"content-type": "application/json",
|
||||
},
|
||||
content=inbound_bytes,
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert transport.captured_body == inbound_bytes
|
||||
assert transport.captured_headers is not None
|
||||
captured_headers = {key.lower(): value for key, value in transport.captured_headers.items()}
|
||||
assert captured_headers["anthropic-beta"] == "sticky-beta-2024-01-01,claude-code-20250219"
|
||||
finally:
|
||||
_reset_session_beta_tracker_for_test()
|
||||
|
||||
|
||||
def test_openai_responses_gzip_nonstream_passthrough_strips_content_encoding() -> None:
|
||||
client, transport = _make_no_optimize_app()
|
||||
decoded_body = _openai_responses_body_bytes(stream=False)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue