Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 18 additions & 2 deletions headroom/proxy/body_forwarding.py
Original file line number Diff line number Diff line change
Expand Up @@ -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."""

Expand Down Expand Up @@ -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",
Expand Down Expand Up @@ -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)
)
70 changes: 69 additions & 1 deletion headroom/proxy/handlers/anthropic.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@
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

Expand Down Expand Up @@ -46,6 +47,23 @@
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.

Expand Down Expand Up @@ -2990,7 +3008,19 @@ def _count_tool_tokens(value: object) -> int:
# 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
Expand Down Expand Up @@ -3371,6 +3401,44 @@ def _count_tool_tokens(value: object) -> int:
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,
Expand Down
7 changes: 5 additions & 2 deletions headroom/proxy/handlers/streaming.py
Original file line number Diff line number Diff line change
Expand Up @@ -1121,18 +1121,20 @@ async def _stream_response_inner(
# 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",
Expand All @@ -1143,6 +1145,7 @@ async def _stream_response_inner(
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
Expand Down
70 changes: 59 additions & 11 deletions headroom/proxy/helpers.py
Original file line number Diff line number Diff line change
Expand Up @@ -871,16 +871,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
Expand All @@ -890,9 +889,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

Expand Down
107 changes: 107 additions & 0 deletions tests/test_proxy_byte_faithful_forwarding.py
Original file line number Diff line number Diff line change
Expand Up @@ -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": [
Expand Down Expand Up @@ -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",
Expand Down
Loading
Loading