fix(proxy): reassemble server_tool_use.input from streamed partial_json (#2449)

## Description

Under `--target-ratio 0.4` a session died mid-run with a fatal Anthropic
400:

```
messages.13.content.0.server_tool_use.input: Input should be an object
```

Root cause is not compression of the request: the request path passes
structured blocks through byte-for-byte. It is **SSE stream
reconstruction**. When the proxy rebuilds a full Anthropic message from
the streamed response (non-stream retry, buffered, and CCR round-trip
paths), the `content_block_stop` handler parsed the accumulated
`_partial_json` into `input` only for blocks whose type was exactly
`tool_use`. A `server_tool_use` block streams its input identically via
`input_json_delta`, so its input was never reassembled: the block kept
the empty start-event `input: {}` and leaked the internal
`_partial_json` scratch key. That reconstructed block becomes assistant
history, and on the next turn the client replays it, so Anthropic
rejects `server_tool_use.input`. `--target-ratio` only makes the
buffered/reconstructed path more likely; it does not itself rewrite the
block.

Refs #2438 (Finding 2). Findings 1 (prompt-cache regression) and 3
(compression not engaging) are architectural and tracked separately.

## Type of Change

- [x] Bug fix (non-breaking change that fixes an issue)
- [ ] New feature (non-breaking change that adds functionality)
- [ ] Breaking change (fix or feature that would cause existing
functionality to change)
- [ ] Documentation update
- [ ] Performance improvement
- [ ] Code refactoring (no functional changes)

## Changes Made

- `headroom/proxy/handlers/streaming.py` (`_parse_sse_to_response`):
gate the `content_block_stop` `_partial_json` → `input` parse on the
presence of `_partial_json`, not `type == "tool_use"`, so
`server_tool_use` (and any future tool-ish block) is reassembled. Always
strip the scratch key; `input` is always a parsed object (`{}` on
malformed/empty JSON).
- `headroom/ccr/response_handler.py`
(`StreamingCCRHandler._reconstruct_anthropic_response`): same
stop-handler fix, and relax the `input_json_delta` accumulator that was
likewise gated on `type == "tool_use"` so server_tool_use partial JSON
is accumulated at all.
- Regression tests in `tests/test_sse_thinking_blocks.py` and
`tests/test_ccr_response_handler_extra.py`: a `server_tool_use` whose
input arrives via `input_json_delta` must reconstruct to the parsed
object with no `_partial_json` leak.
- Leave `CHANGELOG.md` untouched, release-please generates it.

## Testing

- [x] Unit tests pass (`python -m pytest
tests/test_sse_thinking_blocks.py
tests/test_ccr_response_handler_extra.py -q`)
- [x] Linting passes (`ruff check`, `ruff format --check` on the four
changed files)
- [x] Type checking passes (`mypy headroom/proxy/handlers/streaming.py
headroom/ccr/response_handler.py --ignore-missing-imports`)
- [x] New tests added for new functionality

### Test Output

```text
$ python -m pytest tests/test_sse_thinking_blocks.py tests/test_ccr_response_handler_extra.py -q
26 passed in 3.36s

$ ruff check <changed files>
All checks passed!
```

## Real Behavior Proof

- Environment: Windows 11, Python 3.13, local dev checkout on a branch
off upstream/main
- Exact command / steps: Fed a synthetic Anthropic SSE stream with a
`server_tool_use` block whose `input` arrives as `input_json_delta`
partial JSON through both reconstructors (`_parse_sse_to_response`,
`_reconstruct_anthropic_response`); then temporarily restored the `type
== "tool_use"` guard and re-ran.
- Observed result: With the fix, the reconstructed block has `input ==
{"query": ...}` and no `_partial_json` key. With the old guard the test
fails, `input` stays `{}` and the scratch key leaks, reproducing the
malformed block that Anthropic rejects on replay.
- Not tested: End-to-end multi-turn `--target-ratio` session against the
live Anthropic API from this environment, reproduced at the
reconstruction seam instead; the reporter observed the 400 on real
traffic.

## Review Readiness

- [x] I have performed a self-review
- [x] This PR is ready for human review

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
Parideboy 2026-07-20 20:04:16 +02:00 committed by GitHub
parent bec4cce8a9
commit 8c8fae0d0b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 102 additions and 14 deletions

View file

@ -811,9 +811,11 @@ class StreamingCCRHandler:
if dtype == "text_delta": if dtype == "text_delta":
target["text"] = target.get("text", "") + delta.get("text", "") target["text"] = target.get("text", "") + delta.get("text", "")
elif dtype == "input_json_delta": elif dtype == "input_json_delta":
if target.get("type") == "tool_use": # Accumulate for any block streaming input (tool_use AND
partial = delta.get("partial_json", "") # server_tool_use); the stop handler parses it into `input`
target["_partial_json"] = target.get("_partial_json", "") + partial # (#2438).
partial = delta.get("partial_json", "")
target["_partial_json"] = target.get("_partial_json", "") + partial
elif dtype == "thinking_delta": elif dtype == "thinking_delta":
target["thinking_buffer"] = target.get("thinking_buffer", "") + delta.get( target["thinking_buffer"] = target.get("thinking_buffer", "") + delta.get(
"thinking", "" "thinking", ""
@ -830,13 +832,17 @@ class StreamingCCRHandler:
idx = event.get("index") idx = event.get("index")
target = (blocks_by_index.get(idx) if idx is not None else None) or current_block target = (blocks_by_index.get(idx) if idx is not None else None) or current_block
if target is not None: if target is not None:
if target.get("type") == "tool_use" and "_partial_json" in target: # Parse streamed `_partial_json` into `input` for any block
partial = target.pop("_partial_json", "") # that carried input_json_delta — tool_use AND
if partial: # server_tool_use — not just tool_use. The narrow type gate
try: # left server_tool_use.input malformed and leaked the scratch
target["input"] = json.loads(partial) # key into replayed history (#2438). Always strip the key.
except json.JSONDecodeError: if "_partial_json" in target:
target["input"] = {} partial = target.pop("_partial_json")
try:
target["input"] = json.loads(partial) if partial else {}
except json.JSONDecodeError:
target["input"] = {}
if target.get("type") == "thinking" and "thinking_buffer" in target: if target.get("type") == "thinking" and "thinking_buffer" in target:
target["thinking"] = target.pop("thinking_buffer") target["thinking"] = target.pop("thinking_buffer")
if target not in response["content"]: if target not in response["content"]:

View file

@ -494,13 +494,20 @@ class StreamingMixin:
idx = data.get("index") idx = data.get("index")
target = (blocks_by_index.get(idx) if idx is not None else None) or current_block target = (blocks_by_index.get(idx) if idx is not None else None) or current_block
if target is not None: if target is not None:
# Parse accumulated JSON for tool_use blocks. # Parse accumulated JSON into `input` for any block that
if target.get("type") == "tool_use" and "_partial_json" in target: # streamed `input_json_delta` — tool_use AND server_tool_use
# (and future tool-ish blocks). Gating on the block type
# missed server_tool_use, leaving its `input` at the empty
# start-event value and leaking the `_partial_json` scratch
# key into replayed assistant history, which Anthropic then
# rejects with `server_tool_use.input: Input should be an
# object` (#2438). Always strip the scratch key.
if "_partial_json" in target:
raw = target.pop("_partial_json")
try: try:
target["input"] = json.loads(target["_partial_json"]) target["input"] = json.loads(raw) if raw else {}
except json.JSONDecodeError: except json.JSONDecodeError:
target["input"] = {} target["input"] = {}
del target["_partial_json"]
# Materialize the thinking buffer into the # Materialize the thinking buffer into the
# canonical `thinking` field expected by the # canonical `thinking` field expected by the
# Anthropic API. # Anthropic API.

View file

@ -425,3 +425,36 @@ async def test_response_to_sse_preserves_anthropic_shape() -> None:
assert parsed["content"][1]["data"] == "ENC:abc" assert parsed["content"][1]["data"] == "ENC:abc"
assert parsed["stop_reason"] == "refusal" assert parsed["stop_reason"] == "refusal"
assert parsed["stop_details"] == stop_details assert parsed["stop_details"] == stop_details
def test_reconstruct_server_tool_use_input_from_partial_json() -> None:
# StreamingCCRHandler._reconstruct_anthropic_response must parse streamed
# input_json_delta into `input` for server_tool_use, not only tool_use.
# The narrow type gate left server_tool_use.input malformed and leaked the
# `_partial_json` scratch key into replayed assistant history → Anthropic
# 400 `server_tool_use.input: Input should be an object` (#2438).
handler = StreamingCCRHandler(CCRResponseHandler(), provider="anthropic")
events = [
{"type": "message_start", "message": {"id": "msg_1", "model": "claude-opus-4"}},
{
"type": "content_block_start",
"index": 0,
"content_block": {
"type": "server_tool_use",
"id": "srvtoolu_1",
"name": "web_search",
"input": {},
},
},
{
"type": "content_block_delta",
"index": 0,
"delta": {"type": "input_json_delta", "partial_json": '{"query": "x"}'},
},
{"type": "content_block_stop", "index": 0},
]
response = handler._reconstruct_anthropic_response(events)
block = response["content"][0]
assert block["type"] == "server_tool_use"
assert block["input"] == {"query": "x"}
assert "_partial_json" not in block

View file

@ -499,3 +499,45 @@ def test_buffered_ccr_extended_thinking_round_trip_preserves_all_blocks() -> Non
) )
assert round_tripped["content"][2]["type"] == "tool_use" assert round_tripped["content"][2]["type"] == "tool_use"
assert round_tripped["content"][2]["input"] == {"y": 2} assert round_tripped["content"][2]["input"] == {"y": 2}
def test_server_tool_use_input_reassembled_from_partial_json() -> None:
# server_tool_use streams its input via input_json_delta exactly like
# tool_use: the content_block_start carries an empty input, the real args
# arrive as partial_json, and content_block_stop must parse them into an
# object. Gating the parse on type == "tool_use" left server_tool_use.input
# empty and leaked the `_partial_json` scratch key into replayed assistant
# history, which Anthropic rejects on the next turn (#2438).
parser = _Parser()
events = [
{"type": "message_start", "message": {"id": "msg_1", "model": "claude-opus-4"}},
{
"type": "content_block_start",
"index": 0,
"content_block": {
"type": "server_tool_use",
"id": "srvtoolu_1",
"name": "web_search",
"input": {},
},
},
{
"type": "content_block_delta",
"index": 0,
"delta": {"type": "input_json_delta", "partial_json": '{"query": "hea'},
},
{
"type": "content_block_delta",
"index": 0,
"delta": {"type": "input_json_delta", "partial_json": 'droom proxy"}'},
},
{"type": "content_block_stop", "index": 0},
]
sse = _build_sse(events)
response = parser._parse_sse_to_response(sse, "anthropic")
assert response is not None
block = response["content"][0]
assert block["type"] == "server_tool_use"
assert block["input"] == {"query": "headroom proxy"}
# Scratch key must never leak into a block that gets replayed as history.
assert "_partial_json" not in block