diff --git a/headroom/proxy/body_forwarding.py b/headroom/proxy/body_forwarding.py index 4648fce6c..51a9011b2 100644 --- a/headroom/proxy/body_forwarding.py +++ b/headroom/proxy/body_forwarding.py @@ -81,6 +81,16 @@ def has_signed_thinking_blocks(body: dict[str, Any]) -> bool: return False +def _original_body_has_signed_thinking_blocks(original_body_bytes: bytes | None) -> bool: + if original_body_bytes is None: + return False + try: + original = json.loads(original_body_bytes) + except (json.JSONDecodeError, UnicodeDecodeError, ValueError, MemoryError, RecursionError): + return False + return isinstance(original, dict) and has_signed_thinking_blocks(original) + + class BodyMutationTracker: """Records whether a request body was mutated and why.""" @@ -123,7 +133,10 @@ def select_outbound_body( upstream instead of silently claiming the edit landed. """ mode = forwarder_mode if forwarder_mode is not None else get_python_forwarder_mode() - if original_body_bytes is not None and has_signed_thinking_blocks(body): + if original_body_bytes is not None and ( + has_signed_thinking_blocks(body) + or _original_body_has_signed_thinking_blocks(original_body_bytes) + ): return OutboundBody( content=original_body_bytes, source="passthrough", @@ -180,4 +193,7 @@ def outbound_body_is_client_bytes( Mirrors the first branch of :func:`select_outbound_body`; the forwarder mode is deliberately not consulted because that branch overrides it too. """ - return original_body_bytes is not None and has_signed_thinking_blocks(body) + return original_body_bytes is not None and ( + has_signed_thinking_blocks(body) + or _original_body_has_signed_thinking_blocks(original_body_bytes) + ) diff --git a/headroom/proxy/handlers/anthropic.py b/headroom/proxy/handlers/anthropic.py index c41f861fd..09ef90d70 100644 --- a/headroom/proxy/handlers/anthropic.py +++ b/headroom/proxy/handlers/anthropic.py @@ -14,6 +14,7 @@ import time import uuid from datetime import datetime from typing import TYPE_CHECKING, Any +from urllib.parse import urlsplit from headroom.proxy.stage_timer import StageTimer, emit_stage_timings_log @@ -50,6 +51,23 @@ from headroom.proxy.outcome import RequestOutcome logger = logging.getLogger("headroom.proxy") +def _is_googleapis_endpoint(value: object) -> bool: + """Return whether *value* targets Google APIs by parsed hostname. + + A substring check would also trust attacker-controlled hosts such as + ``googleapis.com.example.test``. URL parsing plus a label-boundary suffix + check accepts Google API subdomains without widening the route gate. + """ + raw = str(value).strip() + if not raw: + return False + try: + hostname = (urlsplit(raw).hostname or "").rstrip(".").lower() + except ValueError: + return False + return hostname == "googleapis.com" or hostname.endswith(".googleapis.com") + + class _AnthropicTurnHookUsage: """Usage from hook-triggered Anthropic calls the main response omits. @@ -3012,7 +3030,19 @@ class AnthropicHandlerMixin: # the top-level 'system' parameter ..."), so relocate it back to the # top-level ``system`` parameter as the last step before forwarding. relocated_messages, relocated_system, system_relocated = ( - relocate_system_messages_to_top_level(body["messages"], body.get("system")) + relocate_system_messages_to_top_level( + body["messages"], + body.get("system"), + ( + str(model) + if ( + not upstream_base_url + or getattr(self, "anthropic_backend", None) is not None + or _is_googleapis_endpoint(upstream_base_url) + ) + else None + ), + ) ) if system_relocated: body["messages"] = relocated_messages @@ -3393,6 +3423,44 @@ class AnthropicHandlerMixin: tools = _ttl_tools body_mutation_tracker.mark_mutated("cache_control_ttl_order") + # Signed thinking locks the request to the client's original + # bytes. Once all mutation sites have run, make every downstream + # observer use that same wire body and neutralize savings from + # edits that will not be sent (#2990). This covers PERF, /stats, + # durable savings, response headers, pipeline events, and the + # prefix tracker rather than fixing only one reporting surface. + if outbound_locked_to_client_bytes and body_mutation_tracker.mutated: + discarded_reasons = body_mutation_tracker.reasons + try: + wire_body = json.loads(original_body_bytes or b"") + except (json.JSONDecodeError, UnicodeDecodeError, ValueError, RecursionError): + wire_body = None + if not isinstance(wire_body, dict): + raise ValueError( + "signed-thinking passthrough could not reconstruct its client wire body" + ) + + from headroom.proxy.savings_attribution import SAVINGS_ATTRIBUTION_TAG + from headroom.proxy.tool_schema_savings_policy import ( + TOOL_SCHEMA_SAVINGS_TAGS, + ) + + attribution = tags.get(SAVINGS_ATTRIBUTION_TAG) + if isinstance(attribution, list): + attribution.clear() + for savings_tag in TOOL_SCHEMA_SAVINGS_TAGS: + tags.pop(savings_tag, None) + tags.pop("tool_search_deferred_tools", None) + tags["wire_mutations_discarded"] = len(discarded_reasons) + tags["wire_mutation_reasons"] = ",".join(discarded_reasons) + + body = wire_body + optimized_messages = body.get("messages", []) + tools = body.get("tools") + optimized_tokens = original_tokens + tokens_saved = 0 + transforms_applied = [] + log_cache_breakpoints( request_id=request_id, inbound=inbound_breakpoints, diff --git a/headroom/proxy/handlers/streaming.py b/headroom/proxy/handlers/streaming.py index 07d4f70be..805b8e623 100644 --- a/headroom/proxy/handlers/streaming.py +++ b/headroom/proxy/handlers/streaming.py @@ -1121,18 +1121,20 @@ class StreamingMixin: # bytes once before entering the connection-retry loop. When a # transform mutated the body we re-serialize canonically; otherwise # we forward the original client bytes verbatim. - from headroom.proxy.body_forwarding import prepare_outbound_body_bytes + from headroom.proxy.body_forwarding import select_outbound_body from headroom.proxy.helpers import ( capture_codex_wire_debug, codex_wire_debug_enabled, log_outbound_request, ) - outbound_bytes, outbound_source = prepare_outbound_body_bytes( + outbound = select_outbound_body( body=body, original_body_bytes=original_body_bytes, body_mutated=body_mutated, + mutation_reasons=list(mutation_reasons or []), ) + outbound_bytes, outbound_source = outbound.content, outbound.source outbound_headers = {**headers, "content-type": "application/json"} log_outbound_request( forwarder="streaming", @@ -1143,6 +1145,7 @@ class StreamingMixin: mutation_reasons=list(mutation_reasons or []), request_id=request_id, source=outbound_source, + dropped_mutation_reasons=outbound.dropped_mutation_reasons, ) _codex_wire_debug = ( codex_wire_debug_enabled() and provider == "openai" and "/responses" in url diff --git a/headroom/proxy/helpers.py b/headroom/proxy/helpers.py index a20801c2e..c5349a878 100644 --- a/headroom/proxy/helpers.py +++ b/headroom/proxy/helpers.py @@ -905,16 +905,15 @@ def _system_message_to_blocks(message: dict[str, Any]) -> list[Any]: def relocate_system_messages_to_top_level( messages: list[dict[str, Any]], system: Any, + model: str | None = None, ) -> tuple[list[dict[str, Any]], Any, bool]: - """Move any ``role="system"`` entries out of ``messages`` into ``system``. + """Relocate only system messages invalid for the selected Anthropic model. - Anthropic's Messages API rejects a ``system`` role inside ``messages`` with - HTTP 400 ("messages.0: use the top-level 'system' parameter for the initial - system prompt"). Internal transforms / pipeline extensions can leave a stray - system message in the list (e.g. a relocated harness system block during - compression). This is the Anthropic forwarder's last line of defense: it - guarantees the forwarded body never violates the wire contract, regardless - of which transform introduced the entry. + Supported models accept mid-conversation system sections after a user turn + (or an assistant server-tool result) when followed by an assistant turn or + placed at the end. Hoisting those changes semantics and invalidates the + cached prefix. The initial/invalid forms are still moved to the top-level + field as the issue-765 last-line wire-contract guard. The relocated content is appended after any existing top-level ``system`` so wire order (system prompt, then conversation) is preserved and no content @@ -924,9 +923,58 @@ def relocate_system_messages_to_top_level( message is present the inputs pass through unchanged (``changed=False``) so the common path is untouched. """ - system_indices = { - i for i, m in enumerate(messages) if isinstance(m, dict) and m.get("role") == _ROLE_SYSTEM - } + model_id = str(model or "").lower() + supports_mid_conversation = any( + family in model_id + for family in ( + "claude-fable-5", + "claude-mythos-5", + "claude-opus-4-8", + "claude-opus-5", + "claude-sonnet-5", + ) + ) + + def _assistant_ends_in_server_tool_result(message: object) -> bool: + if not isinstance(message, dict) or message.get("role") != "assistant": + return False + content = message.get("content") + if not isinstance(content, list) or not content: + return False + final = content[-1] + if not isinstance(final, dict): + return False + block_type = str(final.get("type") or "") + return block_type == "server_tool_use" or block_type.endswith("_tool_result") + + system_indices: set[int] = set() + index = 0 + while index < len(messages): + message = messages[index] + if not isinstance(message, dict) or message.get("role") != _ROLE_SYSTEM: + index += 1 + continue + + section_start = index + while ( + index + 1 < len(messages) + and isinstance(messages[index + 1], dict) + and messages[index + 1].get("role") == _ROLE_SYSTEM + ): + index += 1 + section_end = index + + previous = messages[section_start - 1] if section_start > 0 else None + following = messages[section_end + 1] if section_end + 1 < len(messages) else None + valid_previous = ( + isinstance(previous, dict) and previous.get("role") == "user" + ) or _assistant_ends_in_server_tool_result(previous) + valid_following = following is None or ( + isinstance(following, dict) and following.get("role") == "assistant" + ) + if not (supports_mid_conversation and valid_previous and valid_following): + system_indices.update(range(section_start, section_end + 1)) + index += 1 if not system_indices: return messages, system, False diff --git a/tests/test_proxy_byte_faithful_forwarding.py b/tests/test_proxy_byte_faithful_forwarding.py index eacfdc0bd..72fd6056e 100644 --- a/tests/test_proxy_byte_faithful_forwarding.py +++ b/tests/test_proxy_byte_faithful_forwarding.py @@ -282,6 +282,31 @@ def test_signed_thinking_passthrough_reports_the_mutations_it_discarded() -> Non 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": [ @@ -566,6 +591,88 @@ def _make_no_optimize_app() -> tuple[TestClient, _CapturingTransport]: 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", diff --git a/tests/test_proxy_handler_helpers.py b/tests/test_proxy_handler_helpers.py index e0e2a90ea..15b2cd869 100644 --- a/tests/test_proxy_handler_helpers.py +++ b/tests/test_proxy_handler_helpers.py @@ -8,9 +8,10 @@ from types import SimpleNamespace from unittest.mock import patch import httpx +import pytest from fastapi.responses import StreamingResponse -from headroom.proxy.handlers.anthropic import AnthropicHandlerMixin +from headroom.proxy.handlers.anthropic import AnthropicHandlerMixin, _is_googleapis_endpoint from headroom.proxy.handlers.openai import ( OpenAIHandlerMixin, _decode_openai_bearer_payload, @@ -34,6 +35,23 @@ def _jwt(payload: object) -> str: return f"{encode(header)}.{encode(payload)}." +@pytest.mark.parametrize( + ("url", "expected"), + [ + ("https://us-central1-aiplatform.googleapis.com/v1", True), + ("https://googleapis.com/v1", True), + ("https://AIPLATFORM.GOOGLEAPIS.COM./v1", True), + ("https://googleapis.com.example.test/v1", False), + ("https://notgoogleapis.com/v1", False), + ("https://googleapis.com@attacker.test/v1", False), + ("not a url", False), + ("", False), + ], +) +def test_googleapis_endpoint_gate_uses_hostname_boundary(url: str, expected: bool) -> None: + assert _is_googleapis_endpoint(url) is expected + + class _ImageCompressor: def __init__(self, compressed_message): self._compressed_message = compressed_message @@ -312,6 +330,80 @@ def test_relocate_system_messages_noop_without_system_entry() -> None: assert system == "A" +def test_relocate_system_messages_preserves_valid_mid_conversation_section() -> None: + messages = [ + {"role": "user", "content": "Run the tests."}, + { + "role": "system", + "content": "The user added: update the changelog too.", + }, + {"role": "assistant", "content": "I will do both."}, + ] + + clean, system, changed = relocate_system_messages_to_top_level( + messages, "base", "claude-opus-5" + ) + + assert changed is False + assert clean is messages + assert system == "base" + + +def test_relocate_system_messages_preserves_consecutive_valid_section_at_end() -> None: + messages = [ + {"role": "user", "content": [{"type": "tool_result", "content": "ok"}]}, + {"role": "system", "content": "First update."}, + {"role": "system", "content": "Second update."}, + ] + + clean, system, changed = relocate_system_messages_to_top_level( + messages, None, "global.anthropic.claude-sonnet-5-v1:0" + ) + + assert changed is False + assert clean is messages + assert system is None + + +@pytest.mark.parametrize( + "messages", + [ + [ + {"role": "assistant", "content": "answer"}, + {"role": "system", "content": "bad predecessor"}, + ], + [ + {"role": "user", "content": "question"}, + {"role": "system", "content": "bad successor"}, + {"role": "user", "content": "another question"}, + ], + ], +) +def test_relocate_system_messages_still_moves_invalid_mid_conversation_placement( + messages: list[dict], +) -> None: + clean, system, changed = relocate_system_messages_to_top_level(messages, None, "claude-fable-5") + + assert changed is True + assert all(message.get("role") != "system" for message in clean) + assert system == [{"type": "text", "text": messages[1]["content"]}] + + +def test_relocate_system_messages_moves_valid_shape_for_unsupported_model() -> None: + messages = [ + {"role": "user", "content": "question"}, + {"role": "system", "content": "mid-turn instruction"}, + ] + + clean, system, changed = relocate_system_messages_to_top_level( + messages, None, "claude-sonnet-4-6" + ) + + assert changed is True + assert clean == [{"role": "user", "content": "question"}] + assert system == [{"type": "text", "text": "mid-turn instruction"}] + + def test_headroom_bypass_helper_is_transport_neutral() -> None: assert _headroom_bypass_enabled({"x-headroom-bypass": "true"}) is True assert _headroom_bypass_enabled({"x-headroom-bypass": " TRUE "}) is True