mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
## Description Signed-thinking histories force byte-faithful passthrough because re-serializing signed Anthropic blocks can invalidate their signatures. Headroom correctly forwarded the original client bytes, but continued reporting mutations, transforms, savings, response headers, and prefix state from a different body that never reached the provider. Separately, the final Anthropic guard hoisted every `role: system` message into the top-level prompt, including valid mid-conversation system sections, changing their semantics and destroying the cached prefix if that mutation ever shipped. This coupled fix makes downstream accounting use the actual wire body whenever the signed-thinking lock discards edits, and narrows system relocation to the current Anthropic model and placement contract. Closes #2990 Closes #2991 ## 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 - [x] Performance improvement - [x] Code refactoring (no functional changes) ## Changes Made - Detects signed thinking in the original request as well as the mutated body, so a transform cannot remove the block and accidentally bypass the byte lock. - Keeps the original-body signature probe best-effort under malformed, recursive, and `MemoryError` conditions. - Carries discarded mutation reasons through the streaming forwarder and emits the existing structured warning on HTTP streaming paths too. - When signed passthrough wins, resets message savings, tool-schema savings, attribution ledgers, transform labels, response headers, and prefix tracking to the original client wire body. - Adds bounded public diagnostic tags naming/counting discarded mutation reasons without exposing body content. - Preserves valid mid-conversation system sections on currently supported Claude models and official Anthropic, Bedrock, and parsed `*.googleapis.com` routes; hostname-boundary validation rejects lookalike and userinfo URLs. - Preserves consecutive system sections and enforces documented predecessor/successor placement rules. - Continues relocating initial, invalidly placed, unsupported-model, and conservative third-party-gateway system messages to avoid upstream 400s. - Includes current `main`, including #2996, #2997, #2971, #3009, #3012, and the MCP dependency cap. ## Testing - [x] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check .`) - [x] Type checking passes (`mypy headroom`) - [x] New tests added for new functionality - [x] Manual testing performed ### Test Output ```text uv run pytest -q <wire/cache/savings/system focused suite> 379 passed uv run pytest -q tests/test_proxy/test_anthropic_recount_and_reparse_safety.py tests/test_proxy_byte_faithful_forwarding.py tests/test_proxy_handler_helpers.py 99 passed pytest tests scripts/tests --splits 4 --group N --tb=short -q All four CI-shaped fresh-process groups passed locally after correcting the MemoryError regression; each completed in roughly 75-83 seconds. Post-CodeQL correction: 91 focused tests passed; all four exact-head CI-shaped shards passed in roughly 82-95 seconds. uv run ruff format --check . 1411 files already formatted uv run ruff check . All checks passed uv run mypy headroom Success: no issues found in 520 source files ``` ## Real Behavior Proof - Environment: macOS arm64, Python 3.13, branch rebased onto current `main`. - Exact command / steps: sent a signed-thinking request whose tool schema is measurably compacted inside the handler, captured the exact upstream bytes, wrapped the real outcome funnel, and inspected response headers, aggregate metrics, attribution tags, transforms, and prefix-tracker state. Exercised valid, consecutive, invalid, initial, supported-model, and unsupported-model system placements. - Observed result: upstream bytes remain byte-identical to the client; discarded edits contribute zero tokens, zero tool savings, no transform header, and no attribution while the prefix tracker stores the actual wire messages. Valid mid-conversation system sections remain in place; only out-of-contract sections relocate. - Not tested: live paid Anthropic traffic with production credentials. The placement/model contract was verified against the current official documentation and wire behavior is covered with a byte-capturing transport. ## Runtime Rollout Safety - Rollout-managed feature(s): signed-thinking wire-truth accounting and Anthropic mid-conversation system preservation. - Minimum rollout channel: normal patch release after exact-head CI is entirely green. - Stable/default behavior changed: discarded mutations no longer inflate savings; supported valid system sections are no longer hoisted into the top-level prompt. - Kill switch / disable path: no unsafe runtime override; human revert restores the previous conservative relocation/accounting behavior. - Unsafe override required: none. - Qualification impact: all Python shards, byte-forwarding, cache-prefix, outcome/savings, signed-thinking, Anthropic handler, static, Docker, and security checks must remain green. - Rollback path: fix forward through a human-reviewed corrective PR; no persisted data or configuration migration is involved. ## 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] I have made corresponding changes to the documentation — inline wire-contract documentation; no separate guide is required - [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` — it is generated by release-please from my Conventional Commit PR title (a CI guard enforces this) ## Screenshots (if applicable) Not applicable; proxy wire behavior and accounting only. ## Additional Notes Human review only. No merge or auto-merge is configured. Current provider contract reference: https://platform.claude.com/docs/en/build-with-claude/mid-conversation-system-messages
1610 lines
55 KiB
Python
1610 lines
55 KiB
Python
"""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
|