akadmin/headroom
1
0
Fork 0
mirror of https://github.com/headroomlabs-ai/headroom.git synced 2026-08-27 14:17:10 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix(ccr): preserve thinking blocks in buffered stream re-synthesis (#1897)

Browse source 
## Description

Closes #1876.

When CCR forces `stream: false` upstream (the buffered path for
`headroom_retrieve`), the proxy re-synthesizes an SSE stream for the
client from the buffered JSON response via
`StreamingMixin._response_to_sse`. The reported symptom was
extended-thinking responses arriving corrupted: text blocks missing, and
duplicate empty `thinking` blocks with the same timestamp/requestId.

Tracing the two functions the issue pointed at:

- `_response_to_sse()` already handles `thinking`, `redacted_thinking`,
`citations`, and `server_tool_use` blocks explicitly (added across #1451
and #1826) — a direct thinking → text → tool_use round trip through it
reconstructs correctly, so that half of the reported pointer no longer
applies on current `main`.
- `_parse_sse_to_response()`'s `content_block_stop` handling still had
the bug: it deduped appended blocks with `target not in
response["content"]`, plain whole-dict equality. That has two failure
modes: (1) two genuinely distinct blocks that happen to accumulate
identical values (e.g. two separate empty `thinking` blocks) could
collapse into one, and (2) a redelivered `content_block` lifecycle for
the *same* index (e.g. from the proxy's own HTTP/2 stream-reset retry
path) whose accumulated content differs from the first delivery — a
truncated vs. complete `thinking` block, say — produced **two**
dict-unequal entries for one logical block, i.e. exactly the
"duplicated" symptom reported.

## 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()` now
dedupes appended content blocks by block index (falling back to object
identity for the legacy no-index path) instead of whole-dict equality.
One `content_block_stop` per index is honored; a redelivered lifecycle
for an already-appended index is dropped rather than appended as a
second entry.
- `tests/test_sse_thinking_blocks.py`: added three focused regressions —
two distinct empty `thinking` blocks at different indices both survive;
a redelivered block at the same index with *different* accumulated
content collapses to one entry (this one fails on `main` before the fix
— `assert 2 == 1`); and an end-to-end `_response_to_sse` →
`_parse_sse_to_response` round trip for a buffered CCR extended-thinking
response (`thinking` → `text` → `tool_use`) confirming all three block
types survive intact and the thinking block isn't duplicated.

Adjacent open PR #1854 touches the same files for a different symptom
(preserving `stop_details`/`refusal` shape through the legacy test-only
`StreamingCCRHandler`, which isn't wired into any real request path);
this PR doesn't overlap with that change.

## Testing

- [x] Unit tests pass (`pytest`)
- [x] Linting passes (`ruff check`)
- [ ] Type checking passes (`mypy headroom`)
- [x] New tests added for new functionality
- [ ] Manual testing performed

### Test Output

```text
$ .venv/Scripts/python.exe -m pytest tests/test_sse_thinking_blocks.py -v
10 passed in 0.28s

$ .venv/Scripts/python.exe -m pytest tests/ -k "streaming or ccr or sse" -q
737 passed, 39 skipped, 7569 deselected in 145.23s
(2 pre-existing, unrelated failures reproduce identically on unmodified main:
 a CRLF/LF checkout difference in test_owned_asset_encoding.py, and an
 order-dependent CCR-store state flake in test_proxy_ccr.py that passes
 in isolation on both main and this branch.)

$ .venv/Scripts/python.exe -m ruff check headroom/proxy/handlers/streaming.py tests/test_sse_thinking_blocks.py
All checks passed!

$ .venv/Scripts/python.exe -m ruff format --check headroom/proxy/handlers/streaming.py tests/test_sse_thinking_blocks.py
2 files already formatted
```

## Real Behavior Proof

- Environment: local worktree on current `origin/main`.
- Exact command / steps: `git stash` the `streaming.py` fix, run `pytest
tests/test_sse_thinking_blocks.py::test_redelivered_block_same_index_different_content_collapses_to_one_entry`,
then `git stash pop` and rerun.
- Observed result: on unmodified `main` the test fails — `assert 2 ==
1`, with `response["content"]` holding `[{'type': 'thinking', 'index':
0, 'thinking': 'partial'}, {'type': 'thinking', 'index': 0, 'thinking':
'full retried text'}]` — two entries for one logical block index. With
the fix, the same scenario produces exactly one entry. This is the
mechanism behind the reported "duplicate empty thinking blocks" symptom.
- Not tested: a live Claude Code session reproducing the exact reported
transcript signature end-to-end (requires the CCR/retrieval
infrastructure and an extended-thinking model live). The fix is verified
at the unit level against the two functions the issue traced the
corruption to.

## 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
- [ ] I have made corresponding changes to the documentation
- [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
- [ ] I have updated the CHANGELOG.md if applicable
This commit is contained in:
Rod Boev 2026-07-08 22:16:46 -04:00 • committed by GitHub
parent 87f6e93c14
commit ede085cc11
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 163 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

26
headroom/proxy/handlers/streaming.py
View file

@ -350,6 +350,19 @@ class StreamingMixin:
# sequentially, but the index map is the source of truth.
blocks_by_index: dict[int, dict[str, Any]] = {}
current_block: dict[str, Any] | None = None
# Track which block indices have already been appended to
# `response["content"]`. Dedup used to be `target not in
# response["content"]` — plain dict-equality. Two distinct blocks
# that happen to accumulate identical values (most commonly two
# separate empty `thinking` blocks, e.g. from a retried HTTP/2
# stream reset redelivering a truncated segment) either got
# wrongly collapsed into one, or — when their partial content
# happened to differ (same index, unequal dict) — both slipped
# through as duplicates. Indexing by `index` (falling back to
# object identity for the legacy no-index path) makes dedup exact
# regardless of what the accumulated content looks like: one
# entry per block index, first `content_block_stop` wins.
appended_block_keys: set[int] = set()
for line in sse_data.split("\n"):
if not line.startswith("data: "):
@ -456,12 +469,15 @@ class StreamingMixin:
# Anthropic API.
if target.get("type") == "thinking" and "thinking_buffer" in target:
target["thinking"] = target.pop("thinking_buffer")
# Append the block exactly once. `current_block`
# may not match the indexed target if the stream
# interleaved multiple blocks; index-keyed map is
# authoritative.
if target not in response["content"]:
# Append the block exactly once, keyed by its block
# index (or object identity when no index was ever
# assigned). `current_block` may not match the
# indexed target if the stream interleaved multiple
# blocks; index-keyed map is authoritative.
block_key = idx if idx is not None else id(target)
if block_key not in appended_block_keys:
response["content"].append(target)
appended_block_keys.add(block_key)
current_block = None
elif event_type == "message_delta":

142
tests/test_sse_thinking_blocks.py
View file

@ -265,3 +265,145 @@ def test_response_to_sse_rejects_unknown_content_block() -> None:
{"content": [{"type": "future_block", "payload": "preserve me"}]},
"anthropic",
)
# Issue #1876: CCR buffered-stream re-synthesis corrupted extended-thinking
# responses — `content_block_stop` deduped appended blocks by whole-dict
# equality (`target not in response["content"]`), so two distinct blocks
# that happened to be value-identical could collapse into one, or two
# stops for the *same* index with different accumulated content (a
# retried HTTP/2 stream reset redelivering a truncated segment) could
# both slip through as duplicates. Dedup is now keyed by block index.
def test_distinct_empty_thinking_blocks_at_different_indices_both_survive() -> None:
"""Two separate empty `thinking` blocks are two blocks, not one.
Regression guard: if dedup ever regresses to dict-equality, this
collapses to a single entry since both blocks are value-identical.
"""
parser = _Parser()
events = [
{"type": "message_start", "message": {"id": "msg_1", "model": "claude-opus-4"}},
{
"type": "content_block_start",
"index": 0,
"content_block": {"type": "thinking", "thinking": ""},
},
{"type": "content_block_stop", "index": 0},
{
"type": "content_block_start",
"index": 1,
"content_block": {"type": "thinking", "thinking": ""},
},
{"type": "content_block_stop", "index": 1},
{
"type": "content_block_start",
"index": 2,
"content_block": {"type": "text", "text": ""},
},
{
"type": "content_block_delta",
"index": 2,
"delta": {"type": "text_delta", "text": "Here is my answer."},
},
{"type": "content_block_stop", "index": 2},
]
response = parser._parse_sse_to_response(_build_sse(events), "anthropic")
assert response is not None
assert len(response["content"]) == 3
assert response["content"][0]["type"] == "thinking"
assert response["content"][0]["thinking"] == ""
assert response["content"][1]["type"] == "thinking"
assert response["content"][1]["thinking"] == ""
# The text block that followed the two empty thinking blocks must not
# be dropped — this is the "text blocks are missing entirely" half of
# the reported corruption.
assert response["content"][2]["type"] == "text"
assert response["content"][2]["text"] == "Here is my answer."
def test_redelivered_block_same_index_different_content_collapses_to_one_entry() -> None:
"""A fully redelivered content_block lifecycle (start/delta/stop) for
an index that was already appended must not produce a second entry —
even though the redelivered content differs from the first, which is
exactly the case the old whole-dict-equality dedup missed. Reproduces
an HTTP/2 stream-reset retry (`_stream_response`'s retry path)
redelivering a fresh accumulation for the same block index: with the
old `target not in response["content"]` check, the two dicts have
unequal `thinking` text, so *both* slipped through as duplicate
entries for one logical block.
"""
parser = _Parser()
events = [
{"type": "message_start", "message": {"id": "msg_1", "model": "claude-opus-4"}},
{
"type": "content_block_start",
"index": 0,
"content_block": {"type": "thinking", "thinking": ""},
},
{
"type": "content_block_delta",
"index": 0,
"delta": {"type": "thinking_delta", "thinking": "partial"},
},
{"type": "content_block_stop", "index": 0},
# Full redelivery of the same index with different accumulated
# content — must be ignored, not appended as a second block.
{
"type": "content_block_start",
"index": 0,
"content_block": {"type": "thinking", "thinking": ""},
},
{
"type": "content_block_delta",
"index": 0,
"delta": {"type": "thinking_delta", "thinking": "full retried text"},
},
{"type": "content_block_stop", "index": 0},
]
response = parser._parse_sse_to_response(_build_sse(events), "anthropic")
assert response is not None
assert len(response["content"]) == 1
assert response["content"][0]["thinking"] == "partial"
def test_buffered_ccr_extended_thinking_round_trip_preserves_all_blocks() -> None:
"""End-to-end shape for issue #1876: a buffered CCR continuation
response with thinking -> text -> tool_use must reconstruct to SSE
(the re-synthesis path `anthropic.py` uses for the client-facing
stream) with the text preserved and the thinking block emitted
exactly once, unduplicated."""
parser = _Parser()
response = {
"id": "msg_final",
"model": "claude-opus-4",
"role": "assistant",
"content": [
{
"type": "thinking",
"thinking": "Now I have the context, let me answer.",
"signature": "sig_final",
},
{"type": "text", "text": "Based on the retrieved context, here is the answer."},
{"type": "tool_use", "id": "toolu_real_1", "name": "real_tool", "input": {"y": 2}},
],
"stop_reason": "tool_use",
"usage": {"input_tokens": 8, "output_tokens": 12},
}
sse_text = b"".join(parser._response_to_sse(response, "anthropic")).decode("utf-8")
assert sse_text.count('"type": "thinking"') == 1
assert "Based on the retrieved context, here is the answer." in sse_text
round_tripped = parser._parse_sse_to_response(sse_text, "anthropic")
assert round_tripped is not None
assert len(round_tripped["content"]) == 3
assert round_tripped["content"][0]["type"] == "thinking"
assert round_tripped["content"][0]["thinking"] == "Now I have the context, let me answer."
assert round_tripped["content"][1]["type"] == "text"
assert (
round_tripped["content"][1]["text"] == "Based on the retrieved context, here is the answer."
)
assert round_tripped["content"][2]["type"] == "tool_use"
assert round_tripped["content"][2]["input"] == {"y": 2}