"""Byte-faithful Python forwarder tests for PR-A3 (P0-2 fix). The Python forwarder layer (server.py:_retry_request, streaming.py, openai.py:_ws_http_fallback, batch.py) historically re-serialized every request body via httpx's default JSON encoder, drifting separators (``, `` vs ``,``) and ASCII-escaping non-ASCII text. Every such request collapsed Anthropic prompt-cache hit-rate. PR-A3 makes every forwarder byte-faithful: * unmutated body β†’ forward original ``await request.body()`` verbatim; * mutated body β†’ re-serialize once via ``serialize_body_canonical`` (compact separators, ``ensure_ascii=False``). The legacy behavior is still reachable via ``HEADROOM_PROXY_PYTHON_FORWARDER_MODE=legacy_json_kwarg`` for emergency rollback (operator opt-in, not a fallback). """ from __future__ import annotations import gzip import hashlib import json import logging from types import SimpleNamespace from unittest.mock import AsyncMock import httpx import pytest from fastapi.testclient import TestClient from headroom.pipeline import PipelineStage from headroom.proxy.body_forwarding import ( BodyMutationTracker, OutboundBody, get_python_forwarder_mode, outbound_body_is_client_bytes, prepare_outbound_body_bytes, select_outbound_body, serialize_body_canonical, ) from headroom.proxy.helpers import ( _reset_session_beta_tracker_for_test, append_text_to_latest_user_chat_message, get_session_beta_tracker, log_outbound_request, ) from headroom.proxy.server import ProxyConfig, create_app pytest.importorskip("fastapi") @pytest.fixture(autouse=True) def _disable_output_shaper(monkeypatch: pytest.MonkeyPatch) -> None: # Isolate this suite from the opt-in HEADROOM_OUTPUT_SHAPER a developer shell # may export, which otherwise perturbs the byte-faithful assertions. monkeypatch.delenv("HEADROOM_OUTPUT_SHAPER", raising=False) # --------------------------------------------------------------------------- # Unit tests for serializer + tracker # --------------------------------------------------------------------------- def test_serialize_canonical_compact_separators() -> None: """``serialize_body_canonical`` must use compact ``,``/``:`` (no spaces).""" body = {"a": 1, "b": 2} out = serialize_body_canonical(body) assert out == b'{"a":1,"b":2}', repr(out) def test_serialize_canonical_unicode_passthrough() -> None: """UTF-8 must survive β€” no ``\\uXXXX`` ASCII escaping.""" body = {"emoji": "πŸ”₯", "cjk": "ζ—₯本θͺž", "mixed": "hello β†’ δΈ–η•Œ"} out = serialize_body_canonical(body) # Each non-ASCII char appears as raw UTF-8 bytes, never as a \uXXXX literal. assert b"\\u" not in out, repr(out) parsed = json.loads(out.decode("utf-8")) assert parsed == body def test_serialize_canonical_preserves_dict_insertion_order() -> None: """Dict insertion order is preserved (Python 3.7+ guarantee).""" body = {"z": 1, "a": 2, "m": 3} out = serialize_body_canonical(body) assert out.startswith(b'{"z":1,"a":2,"m":3'), repr(out) def test_mutation_tracker_records_reason_memory_injection() -> None: tracker = BodyMutationTracker() assert tracker.mutated is False assert tracker.reasons == [] tracker.mark_mutated("memory_injection") assert tracker.mutated is True assert tracker.reasons == ["memory_injection"] def test_mutation_tracker_records_reason_compression() -> None: tracker = BodyMutationTracker() tracker.mark_mutated("compression_smart_crusher") assert tracker.mutated is True assert tracker.reasons == ["compression_smart_crusher"] def test_mutation_tracker_dedupes_reasons() -> None: tracker = BodyMutationTracker() tracker.mark_mutated("memory_injection") tracker.mark_mutated("memory_injection") tracker.mark_mutated("compression") assert tracker.reasons == ["memory_injection", "compression"] def test_mutation_tracker_rejects_empty_reason() -> None: tracker = BodyMutationTracker() with pytest.raises(ValueError): tracker.mark_mutated("") def test_mutation_tracker_reasons_is_a_copy() -> None: """Caller-mutating the returned list must not affect the tracker.""" tracker = BodyMutationTracker() tracker.mark_mutated("a") out = tracker.reasons out.append("b") assert tracker.reasons == ["a"] # --------------------------------------------------------------------------- # prepare_outbound_body_bytes mode selection # --------------------------------------------------------------------------- def test_prepare_outbound_unmutated_returns_passthrough_bytes() -> None: original = b'{"a":1,"b":"\xf0\x9f\x94\xa5"}' out, source = prepare_outbound_body_bytes( body={"a": 1, "b": "πŸ”₯"}, original_body_bytes=original, body_mutated=False, forwarder_mode="byte_faithful", ) assert out == original assert source == "passthrough" def test_select_outbound_body_returns_value_object() -> None: original = b'{"a":1}' outbound = select_outbound_body( body={"a": 1}, original_body_bytes=original, body_mutated=False, forwarder_mode="byte_faithful", ) assert outbound == OutboundBody(content=original, source="passthrough") def test_helpers_preserve_body_forwarding_compatibility_exports() -> None: from headroom.proxy import helpers assert helpers.BodyMutationTracker is BodyMutationTracker assert helpers.get_python_forwarder_mode is get_python_forwarder_mode assert helpers.prepare_outbound_body_bytes is prepare_outbound_body_bytes assert helpers.serialize_body_canonical is serialize_body_canonical def test_prepare_outbound_mutated_uses_canonical() -> None: out, source = prepare_outbound_body_bytes( body={"a": 1, "b": "πŸ”₯"}, original_body_bytes=b'{"a": 1, "b": "\xf0\x9f\x94\xa5"}', # spaces in original body_mutated=True, forwarder_mode="byte_faithful", ) assert out == b'{"a":1,"b":"\xf0\x9f\x94\xa5"}' assert source == "canonical" @pytest.mark.parametrize("block_type", ["thinking", "redacted_thinking"]) def test_signed_thinking_history_with_original_bytes_uses_passthrough( block_type: str, ) -> None: body = { "model": "claude-sonnet-4-5", "messages": [ {"role": "user", "content": "Solve this"}, { "role": "assistant", "content": [ { "type": block_type, "thinking": "private reasoning", "signature": "sig123", }, {"type": "text", "text": "The answer is 42."}, ], }, {"role": "user", "content": "Continue"}, ], } original = json.dumps(body, indent=2).encode("utf-8") outbound = select_outbound_body( body=body, original_body_bytes=original, body_mutated=True, forwarder_mode="byte_faithful", ) assert outbound.source == "passthrough" assert outbound.content == original def test_signed_thinking_history_without_original_bytes_uses_canonical() -> None: body = { "messages": [ { "role": "assistant", "content": [ { "type": "thinking", "thinking": "private reasoning", "signature": "sig123", } ], } ] } outbound = select_outbound_body( body=body, original_body_bytes=None, body_mutated=True, forwarder_mode="byte_faithful", ) assert outbound.source == "canonical" assert outbound.content == serialize_body_canonical(body) def test_signed_thinking_history_overrides_legacy_encoder() -> None: body = { "messages": [ { "role": "assistant", "content": [{"type": "thinking", "signature": "sig123"}], } ] } original = json.dumps(body, indent=2).encode("utf-8") outbound = select_outbound_body( body=body, original_body_bytes=original, body_mutated=True, forwarder_mode="legacy_json_kwarg", ) assert outbound == OutboundBody(content=original, source="passthrough", dropped_mutations=True) def test_signed_thinking_passthrough_reports_the_mutations_it_discarded() -> None: """Passthrough silently winning over a mutated body is what hid #2952.""" body = { "stream": False, "messages": [ { "role": "assistant", "content": [{"type": "thinking", "signature": "sig123"}], } ], } original = json.dumps({**body, "stream": True}).encode("utf-8") outbound = select_outbound_body( body=body, original_body_bytes=original, body_mutated=True, forwarder_mode="byte_faithful", mutation_reasons=["ccr_streaming_retrieve_buffered_non_stream"], ) assert outbound.source == "passthrough" assert outbound.dropped_mutations is True assert outbound.dropped_mutation_reasons == ("ccr_streaming_retrieve_buffered_non_stream",) def test_original_signed_thinking_still_locks_when_mutation_removed_the_block() -> None: original_body = { "messages": [ { "role": "assistant", "content": [{"type": "thinking", "signature": "sig123"}], } ] } mutated_body = {"messages": [{"role": "assistant", "content": "rewritten"}]} original = json.dumps(original_body, indent=2).encode() outbound = select_outbound_body( body=mutated_body, original_body_bytes=original, body_mutated=True, forwarder_mode="byte_faithful", mutation_reasons=["compression"], ) assert outbound.content == original assert outbound.source == "passthrough" assert outbound.dropped_mutation_reasons == ("compression",) def test_signed_thinking_passthrough_reports_nothing_when_body_unmutated() -> None: body = { "messages": [ { "role": "assistant", "content": [{"type": "thinking", "signature": "sig123"}], } ] } original = json.dumps(body).encode("utf-8") outbound = select_outbound_body( body=body, original_body_bytes=original, body_mutated=False, forwarder_mode="byte_faithful", mutation_reasons=["irrelevant"], ) assert outbound.source == "passthrough" assert outbound.dropped_mutations is False assert outbound.dropped_mutation_reasons == () def test_canonical_path_reports_no_dropped_mutations() -> None: body = {"messages": [{"role": "user", "content": "hi"}]} outbound = select_outbound_body( body=body, original_body_bytes=b'{"messages": []}', body_mutated=True, forwarder_mode="byte_faithful", mutation_reasons=["compression"], ) assert outbound.source == "canonical" assert outbound.dropped_mutations is False assert outbound.dropped_mutation_reasons == () @pytest.mark.parametrize( ("original_body_bytes", "expected"), [(b'{"messages": []}', True), (None, False)], ) def test_outbound_body_is_client_bytes_matches_selection( original_body_bytes: bytes | None, expected: bool ) -> None: """Handlers gate on this before mutating a body for their own upstream call.""" body = { "messages": [ { "role": "assistant", "content": [{"type": "thinking", "signature": "sig123"}], } ] } assert ( outbound_body_is_client_bytes(body=body, original_body_bytes=original_body_bytes) is expected ) outbound = select_outbound_body( body=body, original_body_bytes=original_body_bytes, body_mutated=True, forwarder_mode="byte_faithful", ) assert (outbound.source == "passthrough") is expected def test_outbound_body_is_client_bytes_false_without_thinking_blocks() -> None: assert ( outbound_body_is_client_bytes( body={"messages": [{"role": "user", "content": "hi"}]}, original_body_bytes=b'{"messages": []}', ) is False ) def test_prepare_outbound_no_original_bytes_uses_canonical() -> None: out, source = prepare_outbound_body_bytes( body={"a": 1}, original_body_bytes=None, body_mutated=False, forwarder_mode="byte_faithful", ) assert out == b'{"a":1}' assert source == "canonical" def test_legacy_json_kwarg_mode_falls_back() -> None: """legacy_json_kwarg is an explicit operator opt-in β€” produces the historical bytes. This is NOT a silent fallback (build constraint #4). It is reachable only via env var and exists for emergency rollback validation. """ out, source = prepare_outbound_body_bytes( body={"a": 1, "b": "πŸ”₯"}, original_body_bytes=b'{"a":1}', body_mutated=False, forwarder_mode="legacy_json_kwarg", ) # Old httpx default: spaces after `,` and `:`, ascii escaping. assert out == b'{"a": 1, "b": "\\ud83d\\udd25"}', repr(out) assert source == "legacy" def test_python_forwarder_mode_default_is_byte_faithful( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.delenv("HEADROOM_PROXY_PYTHON_FORWARDER_MODE", raising=False) assert get_python_forwarder_mode() == "byte_faithful" def test_python_forwarder_mode_invalid_raises( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setenv("HEADROOM_PROXY_PYTHON_FORWARDER_MODE", "garbage") with pytest.raises(ValueError, match="HEADROOM_PROXY_PYTHON_FORWARDER_MODE"): get_python_forwarder_mode() def test_python_forwarder_mode_legacy_value_accepted( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setenv("HEADROOM_PROXY_PYTHON_FORWARDER_MODE", "legacy_json_kwarg") assert get_python_forwarder_mode() == "legacy_json_kwarg" # --------------------------------------------------------------------------- # log_outbound_request structured log content # --------------------------------------------------------------------------- def test_log_outbound_request_emits_structured_fields() -> None: """Capture the structured log line via a temporary handler. We attach a memory handler directly to the proxy logger so the test is independent of whether ``_setup_file_logging`` has set ``propagate=False`` (which it does in the live proxy). """ import logging proxy_logger = logging.getLogger("headroom.proxy") records: list[logging.LogRecord] = [] class _ListHandler(logging.Handler): def emit(self, record: logging.LogRecord) -> None: records.append(record) handler = _ListHandler(level=logging.INFO) prev_level = proxy_logger.level proxy_logger.addHandler(handler) proxy_logger.setLevel(logging.INFO) try: log_outbound_request( forwarder="server", method="POST", path="/v1/messages", body_bytes_count=42, body_mutated=False, mutation_reasons=[], request_id="hr_test_1", source="passthrough", ) finally: proxy_logger.removeHandler(handler) proxy_logger.setLevel(prev_level) matching = [r for r in records if "outbound_request" in r.getMessage()] assert matching, f"no outbound_request log emitted; records={records!r}" msg = matching[-1].getMessage() assert "event=outbound_request" in msg assert "forwarder=server" in msg assert "path=/v1/messages" in msg assert "body_bytes=42" in msg assert "body_mutated=false" in msg assert "source=passthrough" in msg assert "request_id=hr_test_1" in msg # Never log auth / body content. assert "Authorization" not in msg assert "x-api-key" not in msg.lower() # --------------------------------------------------------------------------- # httpx-mock end-to-end byte-faithful checks # --------------------------------------------------------------------------- class _CapturingTransport(httpx.AsyncBaseTransport): """An httpx transport that records the exact bytes received.""" def __init__(self) -> None: self.captured_body: bytes | None = None self.captured_headers: dict[str, str] | None = None self.captured_url: str | None = None async def handle_async_request(self, request: httpx.Request) -> httpx.Response: # Eagerly read the request body so streaming bodies are captured too. body = b"" async for chunk in request.stream: body += chunk self.captured_body = body self.captured_headers = dict(request.headers.items()) self.captured_url = str(request.url) return httpx.Response( 200, json={ "id": "msg_1", "type": "message", "role": "assistant", "content": [{"type": "text", "text": "ok"}], "usage": { "input_tokens": 10, "output_tokens": 3, "cache_read_input_tokens": 0, "cache_creation_input_tokens": 0, }, }, ) class _FakePrefixTracker: def __init__(self, frozen_count: int = 0): self._frozen_count = frozen_count self._cached_token_count = 0 self._last_original_messages: list = [] self._last_forwarded_messages: list = [] def get_frozen_message_count(self) -> int: return self._frozen_count def get_last_original_messages(self): # noqa: ANN201 return list(self._last_original_messages) def get_last_forwarded_messages(self): # noqa: ANN201 return list(self._last_forwarded_messages) def update_from_response(self, **kwargs): # noqa: ANN003 self._last_original_messages = kwargs.get("original_messages", kwargs.get("messages", [])) self._last_forwarded_messages = kwargs.get("messages", []) return None class _SortedEmptyToolsPreSendExtension: def on_pipeline_event(self, event): # noqa: ANN001 if event.stage is PipelineStage.PRE_SEND: event.tools = [] return None def _make_anthropic_app(*, optimize: bool) -> tuple[TestClient, _CapturingTransport]: """Boot an Anthropic proxy with a capturing transport.""" config = ProxyConfig( optimize=optimize, 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) transport = _CapturingTransport() proxy = app.state.proxy proxy.http_client = httpx.AsyncClient(transport=transport) # Pin a stable session tracker so the prefix walker doesn't re-read # turn 0 on every run. fake_tracker = _FakePrefixTracker(frozen_count=0) proxy.session_tracker_store.compute_session_id = lambda request, model, messages: "s1" proxy.session_tracker_store.get_or_create = lambda session_id, provider: fake_tracker return TestClient(app), transport def _make_no_optimize_app() -> tuple[TestClient, _CapturingTransport]: """Boot a proxy with all transforms disabled and a capturing transport.""" return _make_anthropic_app(optimize=False) def test_signed_thinking_discarded_mutation_uses_wire_truth_for_all_accounting() -> None: 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 transport = _CapturingTransport() proxy.http_client = httpx.AsyncClient(transport=transport) proxy._record_request_outcome = AsyncMock(wraps=proxy._record_request_outcome) tracker = _FakePrefixTracker(frozen_count=0) proxy.session_tracker_store.compute_session_id = lambda request, model, messages: "signed" proxy.session_tracker_store.get_or_create = lambda session_id, provider: tracker inbound = { "model": "claude-opus-5", "max_tokens": 64, "messages": [ {"role": "user", "content": "Solve this."}, { "role": "assistant", "content": [ { "type": "thinking", "thinking": "private", "signature": "sig123", }, {"type": "text", "text": "Working."}, ], }, {"role": "user", "content": "Continue."}, ], "tools": [ { "name": "lookup", "description": " Look up a value. ", "input_schema": { "$schema": "https://json-schema.org/draft/2020-12/schema", "type": "object", "properties": {"key": {"type": "string"}}, }, } ], } inbound_bytes = json.dumps(inbound, indent=2).encode() response = TestClient(app).post( "/v1/messages", headers={ "x-api-key": "test-key", "anthropic-version": "2023-06-01", "content-type": "application/json", }, content=inbound_bytes, ) assert response.status_code == 200 assert transport.captured_body == inbound_bytes assert response.headers["x-headroom-tokens-saved"] == "0" assert "x-headroom-transforms" not in response.headers outcome = proxy._record_request_outcome.await_args.args[0] assert outcome.tokens_saved == 0 assert outcome.optimized_tokens == outcome.original_tokens assert outcome.transforms_applied == () assert outcome.tags["wire_mutations_discarded"] > 0 assert "anthropic:tool_schema_compaction" not in outcome.transforms_applied assert "tool_search_deferred_tokens" not in outcome.tags assert outcome.tags.get("_headroom_savings_attribution") == [] assert proxy.metrics.tokens_saved_total == 0 assert proxy.metrics.tool_search_saved_total == 0 assert tracker._last_forwarded_messages[: len(inbound["messages"])] == inbound["messages"] def _openai_responses_body_bytes(*, stream: bool) -> bytes: payload = { "model": "gpt-5.5", "input": [ { "type": "message", "role": "user", "content": [ { "type": "input_text", "text": "hello πŸ”₯ with spaces preserved", } ], } ], "stream": stream, } return json.dumps(payload, ensure_ascii=False, indent=2).encode("utf-8") def _openai_responses_codex_headers(content_encoding: str) -> dict[str, str]: return { "authorization": "Bearer test-token", "chatgpt-account-id": "acct_test", "originator": "Codex Desktop", "content-type": "application/json", "content-encoding": content_encoding, "accept": "text/event-stream", } def _start_proxy_log_capture() -> tuple[ logging.Logger, logging.Handler, int, list[logging.LogRecord], ]: proxy_logger = logging.getLogger("headroom.proxy") records: list[logging.LogRecord] = [] class _ListHandler(logging.Handler): def emit(self, record: logging.LogRecord) -> None: records.append(record) handler = _ListHandler(level=logging.INFO) prev_level = proxy_logger.level proxy_logger.addHandler(handler) proxy_logger.setLevel(logging.INFO) return proxy_logger, handler, prev_level, records def _stop_proxy_log_capture( proxy_logger: logging.Logger, handler: logging.Handler, prev_level: int, ) -> None: proxy_logger.removeHandler(handler) proxy_logger.setLevel(prev_level) def _assert_openai_responses_encoded_passthrough( transport: _CapturingTransport, decoded_body: bytes, ) -> None: assert transport.captured_body == decoded_body assert transport.captured_headers is not None captured_headers = {key.lower(): value for key, value in transport.captured_headers.items()} assert "content-encoding" not in captured_headers assert captured_headers.get("content-length") == str(len(decoded_body)) def _assert_outbound_passthrough_log( records: list[logging.LogRecord], *, forwarder: str, ) -> None: messages = [record.getMessage() for record in records] assert any( "event=outbound_request" in message and f"forwarder={forwarder}" in message and "body_mutated=false" in message and "source=passthrough" in message for message in messages ), messages def test_passthrough_no_mutation_byte_equal_sha256() -> None: """No transform β†’ upstream SHA-256 equals client-sent SHA-256.""" client, transport = _make_no_optimize_app() # Compact JSON, simulating Claude Code / Codex CLI byte format. inbound_dict = { "model": "claude-sonnet-4-6", "max_tokens": 64, "messages": [{"role": "user", "content": "hello"}], } inbound_bytes = serialize_body_canonical(inbound_dict) response = client.post( "/v1/messages", headers={ "x-api-key": "test-key", "anthropic-version": "2023-06-01", "content-type": "application/json", }, content=inbound_bytes, ) assert response.status_code == 200, response.text assert transport.captured_body is not None inbound_sha = hashlib.sha256(inbound_bytes).hexdigest() upstream_sha = hashlib.sha256(transport.captured_body).hexdigest() assert inbound_sha == upstream_sha, ( f"Byte-faithful invariant broken: inbound {inbound_sha} vs upstream " f"{upstream_sha}; upstream body={transport.captured_body!r}" ) def test_compression_off_unicode_preserved() -> None: """Emoji + CJK content survives forwarding without ``\\uXXXX`` escaping.""" client, transport = _make_no_optimize_app() inbound_dict = { "model": "claude-sonnet-4-6", "max_tokens": 64, "messages": [ {"role": "user", "content": "Hello πŸ”₯ β€” δΈ–η•Œ β€” emoji is πŸš€"}, ], } inbound_bytes = serialize_body_canonical(inbound_dict) response = client.post( "/v1/messages", headers={ "x-api-key": "test-key", "anthropic-version": "2023-06-01", "content-type": "application/json", }, content=inbound_bytes, ) assert response.status_code == 200 upstream = transport.captured_body or b"" assert upstream == inbound_bytes assert b"\\u" not in upstream, repr(upstream) assert "πŸ”₯".encode() in upstream assert "δΈ–η•Œ".encode() in upstream def test_compression_off_numeric_precision_preserved() -> None: """Floats with trailing zero stay floats; large ints preserve precision.""" client, transport = _make_no_optimize_app() inbound_bytes = b'{"model":"claude-sonnet-4-6","max_tokens":64,"temperature":1.0,"seed":12345678901234567,"messages":[{"role":"user","content":"hi"}]}' response = client.post( "/v1/messages", headers={ "x-api-key": "test-key", "anthropic-version": "2023-06-01", "content-type": "application/json", }, content=inbound_bytes, ) assert response.status_code == 200 upstream = transport.captured_body or b"" # Unmutated β†’ byte-faithful: exact bytes preserved. assert upstream == inbound_bytes # Forward coverage only; the PRE_SEND case below is the base-fails proof for this fix. def test_anthropic_tools_canonical_order_preserves_byte_faithful_request() -> None: client, transport = _make_no_optimize_app() inbound_dict = { "model": "claude-sonnet-4-6", "max_tokens": 64, "messages": [{"role": "user", "content": "plan test"}], "tools": [ {"name": "alpha"}, {"name": "zeta", "description": "later"}, ], } inbound_bytes = serialize_body_canonical(inbound_dict) response = client.post( "/v1/messages", headers={ "x-api-key": "test-key", "anthropic-version": "2023-06-01", "content-type": "application/json", }, content=inbound_bytes, ) assert response.status_code == 200, response.text upstream = transport.captured_body or b"" assert upstream == inbound_bytes, ( f"Expected byte-faithful passthrough for canonical tools; upstream={upstream!r}" ) def test_anthropic_tools_unsorted_order_preserves_byte_faithful_request() -> None: client, transport = _make_no_optimize_app() inbound_dict = { "model": "claude-sonnet-4-6", "max_tokens": 64, "messages": [{"role": "user", "content": "plan test"}], "tools": [ {"name": "zeta", "description": "later"}, {"name": "alpha"}, ], } inbound_bytes = serialize_body_canonical(inbound_dict) response = client.post( "/v1/messages", headers={ "x-api-key": "test-key", "anthropic-version": "2023-06-01", "content-type": "application/json", }, content=inbound_bytes, ) assert response.status_code == 200, response.text upstream = transport.captured_body or b"" assert upstream == inbound_bytes forwarded = json.loads(upstream.decode("utf-8")) assert [tool["name"] for tool in forwarded["tools"]] == ["zeta", "alpha"] def test_anthropic_tools_unsorted_reordered_and_canonicalized_when_optimized() -> None: client, transport = _make_anthropic_app(optimize=True) proxy = client.app.state.proxy proxy.config.mode = "token" def _fake_apply(**kwargs): return SimpleNamespace( messages=kwargs["messages"], transforms_applied=[], timing={}, tokens_before=100, tokens_after=100, waste_signals=None, ) proxy.anthropic_pipeline.apply = _fake_apply inbound_dict = { "model": "claude-sonnet-4-6", "max_tokens": 64, "messages": [{"role": "user", "content": "plan test"}], "tools": [ {"name": "zeta", "description": "later"}, {"name": "alpha"}, ], } expected_dict = { **inbound_dict, "tools": [ inbound_dict["tools"][1], inbound_dict["tools"][0], ], } inbound_bytes = serialize_body_canonical(inbound_dict) expected_bytes = serialize_body_canonical(expected_dict) response = client.post( "/v1/messages", headers={ "x-api-key": "test-key", "anthropic-version": "2023-06-01", "content-type": "application/json", }, content=inbound_bytes, ) assert response.status_code == 200, response.text upstream = transport.captured_body or b"" assert upstream == expected_bytes assert upstream != inbound_bytes def test_anthropic_presend_sorted_empty_tools_keeps_body_unmutated() -> None: inbound_dict = { "model": "claude-sonnet-4-6", "max_tokens": 64, "messages": [{"role": "user", "content": "plan test"}], } inbound_bytes = serialize_body_canonical(inbound_dict) 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, pipeline_extensions=[_SortedEmptyToolsPreSendExtension()], discover_pipeline_extensions=False, ) app = create_app(config) client = TestClient(app) captured: dict[str, object] = {} async def _fake_retry( method: str, # noqa: ARG001 url: str, # noqa: ARG001 headers: dict[str, str], # noqa: ARG001 body: dict[str, object], # noqa: ARG001 body_mutated: bool, mutation_reasons: list[str], **kwargs: object, # noqa: ANN003 ) -> httpx.Response: # noqa: ANN201 captured["body_mutated"] = body_mutated captured["mutation_reasons"] = mutation_reasons captured["body"] = body return httpx.Response( 200, json={ "id": "msg_1", "type": "message", "role": "assistant", "content": [{"type": "text", "text": "ok"}], "usage": { "input_tokens": 10, "output_tokens": 3, "cache_read_input_tokens": 0, "cache_creation_input_tokens": 0, }, }, ) app.state.proxy._retry_request = _fake_retry # type: ignore[assignment] response = client.post( "/v1/messages", headers={ "x-api-key": "test-key", "anthropic-version": "2023-06-01", "content-type": "application/json", }, content=inbound_bytes, ) assert response.status_code == 200, response.text assert captured["body_mutated"] is False assert captured["mutation_reasons"] == [] forwarded = captured["body"] assert isinstance(forwarded, dict) assert "tools" not in forwarded def test_legacy_json_kwarg_mode_yields_drifted_bytes( monkeypatch: pytest.MonkeyPatch, ) -> None: """Operator opt-in produces the OLD drifted bytes (rollback validation).""" monkeypatch.setenv("HEADROOM_PROXY_PYTHON_FORWARDER_MODE", "legacy_json_kwarg") client, transport = _make_no_optimize_app() inbound_dict = { "model": "claude-sonnet-4-6", "max_tokens": 64, "messages": [{"role": "user", "content": "πŸ”₯ hi"}], } inbound_bytes = serialize_body_canonical(inbound_dict) response = client.post( "/v1/messages", headers={ "x-api-key": "test-key", "anthropic-version": "2023-06-01", "content-type": "application/json", }, content=inbound_bytes, ) assert response.status_code == 200 upstream = transport.captured_body or b"" # Legacy mode: spaces after separators + ASCII escaping β†’ bytes drift. assert upstream != inbound_bytes assert b", " in upstream or b": " in upstream assert b"\\u" in upstream # ASCII escaping confirms legacy path. # --------------------------------------------------------------------------- # A2 follow-up: OpenAI Chat Completions memory routes to user-tail # --------------------------------------------------------------------------- def test_append_text_to_latest_user_chat_message_string_content() -> None: msgs = [ {"role": "system", "content": "you are helpful"}, {"role": "user", "content": "previous"}, {"role": "assistant", "content": "ack"}, {"role": "user", "content": "latest"}, ] new_msgs, appended = append_text_to_latest_user_chat_message(msgs, "MEMCTX") assert appended > 0 assert new_msgs[0] == msgs[0] assert new_msgs[1] == msgs[1] assert new_msgs[2] == msgs[2] assert new_msgs[3]["content"] == "latest\n\nMEMCTX" # Original list untouched. assert msgs[3]["content"] == "latest" def test_append_text_to_latest_user_chat_message_list_content() -> None: msgs = [ { "role": "user", "content": [ {"type": "text", "text": "first text"}, {"type": "image_url", "image_url": {"url": "..."}}, ], } ] new_msgs, appended = append_text_to_latest_user_chat_message(msgs, "MEM") assert appended > 0 parts = new_msgs[0]["content"] assert parts[0]["text"] == "first text\n\nMEM" assert parts[1] == msgs[0]["content"][1] def test_append_text_to_latest_user_chat_message_no_user_returns_zero() -> None: msgs = [{"role": "system", "content": "sys"}] new_msgs, appended = append_text_to_latest_user_chat_message(msgs, "MEM") assert appended == 0 assert new_msgs == msgs def test_openai_chat_memory_routes_to_user_tail_not_system() -> None: """A2 follow-up: Chat Completions memory injection lives in user tail.""" 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 proxy.memory_handler = SimpleNamespace( config=SimpleNamespace(inject_context=True, inject_tools=False), search_and_format_context=AsyncMock(return_value="MEMCTX_OAI"), has_memory_tool_calls=lambda resp, provider: False, ) captured: dict[str, object] = {} async def _fake_retry(method, url, headers, body, stream=False, **kwargs): # noqa: ANN001 captured["body"] = body return httpx.Response( 200, json={ "id": "chatcmpl_1", "object": "chat.completion", "model": "gpt-4o", "choices": [ { "index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop", } ], "usage": {"prompt_tokens": 10, "completion_tokens": 1, "total_tokens": 11}, }, ) proxy._retry_request = _fake_retry # type: ignore[attr-defined] client = TestClient(app) resp = client.post( "/v1/chat/completions", headers={ "authorization": "Bearer sk-test", "x-headroom-user-id": "u1", }, json={ "model": "gpt-4o", "messages": [ {"role": "system", "content": "you are helpful"}, {"role": "user", "content": "what is up"}, ], }, ) assert resp.status_code == 200, resp.text sent = captured.get("body") assert isinstance(sent, dict), captured sent_msgs = sent["messages"] # System message must NOT be mutated. assert sent_msgs[0]["role"] == "system" assert sent_msgs[0]["content"] == "you are helpful", "system message must remain byte-equal" # No injected system message at the start (legacy prepend retired). # The ONLY new content is in the latest user message tail. user_msgs = [m for m in sent_msgs if m.get("role") == "user"] assert len(user_msgs) == 1 assert user_msgs[-1]["content"].endswith("MEMCTX_OAI") # No additional system messages either (memory must not prepend). system_msgs = [m for m in sent_msgs if m.get("role") == "system"] assert len(system_msgs) == 1 def test_openai_chat_memory_disabled_mode_no_op( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setenv("HEADROOM_MEMORY_INJECTION_MODE", "disabled") 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 proxy.memory_handler = SimpleNamespace( config=SimpleNamespace(inject_context=True, inject_tools=False), search_and_format_context=AsyncMock(return_value="WOULD_NOT_INJECT"), has_memory_tool_calls=lambda resp, provider: False, ) captured: dict[str, object] = {} async def _fake_retry(method, url, headers, body, stream=False, **kwargs): # noqa: ANN001 captured["body"] = body return httpx.Response( 200, json={ "id": "c1", "object": "chat.completion", "model": "gpt-4o", "choices": [ { "index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop", } ], "usage": {"prompt_tokens": 5, "completion_tokens": 1, "total_tokens": 6}, }, ) proxy._retry_request = _fake_retry client = TestClient(app) resp = client.post( "/v1/chat/completions", headers={ "authorization": "Bearer sk-test", "x-headroom-user-id": "u1", }, json={ "model": "gpt-4o", "messages": [ {"role": "system", "content": "sys"}, {"role": "user", "content": "hi"}, ], }, ) assert resp.status_code == 200 sent = captured["body"] assert isinstance(sent, dict) assert sent["messages"][1]["content"] == "hi" # --------------------------------------------------------------------------- # Streaming forwarder byte-faithfulness # --------------------------------------------------------------------------- class _StreamingCapturingTransport(httpx.AsyncBaseTransport): def __init__(self) -> None: self.captured_body: bytes | None = None self.captured_headers: dict[str, str] | None = None async def handle_async_request(self, request: httpx.Request) -> httpx.Response: body = b"" async for chunk in request.stream: body += chunk self.captured_body = body self.captured_headers = dict(request.headers.items()) return httpx.Response( 200, headers={"content-type": "text/event-stream"}, 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( 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 # Pin session tracker so the cache-stable delta path is a no-op. fake_tracker = _FakePrefixTracker(frozen_count=0) proxy.session_tracker_store.compute_session_id = lambda request, model, messages: "s_stream" 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) inbound_bytes = ( '{"model":"claude-sonnet-4-6","max_tokens":16,"stream":true,' '"messages":[{"role":"user","content":"hi πŸ”₯"}]}' ).encode() with client.stream( "POST", "/v1/messages", headers={ "x-api-key": "test-key", "anthropic-version": "2023-06-01", "content-type": "application/json", }, content=inbound_bytes, ) as resp: # Drain the response. for _ in resp.iter_bytes(): pass upstream = transport.captured_body or b"" inbound_sha = hashlib.sha256(inbound_bytes).hexdigest() upstream_sha = hashlib.sha256(upstream).hexdigest() assert upstream_sha == inbound_sha, ( f"Streaming byte-faithfulness broken: inbound {inbound_sha} vs " f"upstream {upstream_sha}; upstream={upstream!r}" ) 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) encoded_body = gzip.compress(decoded_body) proxy_logger, handler, prev_level, records = _start_proxy_log_capture() try: response = client.post( "/v1/responses", headers=_openai_responses_codex_headers("gzip"), content=encoded_body, ) finally: _stop_proxy_log_capture(proxy_logger, handler, prev_level) assert response.status_code == 200, response.text _assert_openai_responses_encoded_passthrough(transport, decoded_body) _assert_outbound_passthrough_log(records, forwarder="openai_responses") def test_openai_responses_gzip_stream_passthrough_strips_content_encoding() -> None: client, transport = _make_no_optimize_app() decoded_body = _openai_responses_body_bytes(stream=True) encoded_body = gzip.compress(decoded_body) proxy_logger, handler, prev_level, records = _start_proxy_log_capture() try: with client.stream( "POST", "/v1/responses", headers=_openai_responses_codex_headers("gzip"), content=encoded_body, ) as response: assert response.status_code == 200 for _ in response.iter_bytes(): pass finally: _stop_proxy_log_capture(proxy_logger, handler, prev_level) _assert_openai_responses_encoded_passthrough(transport, decoded_body) _assert_outbound_passthrough_log(records, forwarder="streaming") def test_openai_responses_codex_desktop_zstd_stream_passthrough_strips_content_encoding() -> None: zstandard = pytest.importorskip("zstandard") client, transport = _make_no_optimize_app() decoded_body = _openai_responses_body_bytes(stream=True) encoded_body = zstandard.ZstdCompressor().compress(decoded_body) proxy_logger, handler, prev_level, records = _start_proxy_log_capture() try: with client.stream( "POST", "/v1/responses", headers=_openai_responses_codex_headers("zstd"), content=encoded_body, ) as response: assert response.status_code == 200 for _ in response.iter_bytes(): pass finally: _stop_proxy_log_capture(proxy_logger, handler, prev_level) _assert_openai_responses_encoded_passthrough(transport, decoded_body) _assert_outbound_passthrough_log(records, forwarder="streaming") # --------------------------------------------------------------------------- # Batch forwarder byte-faithfulness (passthrough variant) # --------------------------------------------------------------------------- def test_batch_passthrough_byte_faithful() -> None: """OpenAI batch passthrough forwards original bytes verbatim.""" 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 transport = _CapturingTransport() proxy.http_client = httpx.AsyncClient(transport=transport) client = TestClient(app) # Use a non-chat-completions endpoint so the handler routes directly to # _batch_passthrough (bypassing the file-compression flow). inbound = { "input_file_id": "file-abc", "endpoint": "/v1/embeddings", "completion_window": "24h", } inbound_bytes = serialize_body_canonical(inbound) resp = client.post( "/v1/batches", headers={ "authorization": "Bearer sk-test", "content-type": "application/json", }, content=inbound_bytes, ) # The capturing transport returns a Message JSON, which is fine; the # status code may vary depending on routing. What matters is that when # the upstream did receive bytes, they are byte-equal to what the # client sent (passthrough case, no body mutation). assert resp.status_code in (200, 400, 401, 404, 422, 500), resp.text if transport.captured_body is not None: assert transport.captured_body == inbound_bytes, ( f"Batch passthrough bytes drifted: " f"sent={inbound_bytes!r} upstream={transport.captured_body!r}" ) # --------------------------------------------------------------------------- # WSβ†’HTTP fallback: just exercises the helper resolution # --------------------------------------------------------------------------- def test_ws_http_fallback_uses_canonical_serializer() -> None: """WSβ†’HTTP fallback resynthesizes the body, so canonical bytes apply. We can't easily exercise the full WS path in a TestClient without a Codex client; instead we assert the helper choice yields the expected bytes when a tracker reports mutation. """ body = {"model": "gpt-5", "input": [{"role": "user", "content": "hi πŸš€"}]} out, source = prepare_outbound_body_bytes( body=body, original_body_bytes=None, body_mutated=True, ) assert source == "canonical" assert b"\\u" not in out # Round-trip equality via JSON parse. assert json.loads(out.decode("utf-8")) == body