fix(cache): extract tool_result content from list-of-blocks format (#2092)

## Description

Closes #2053

Modern Claude Code sends `tool_result` content as a list of typed blocks
(`[{"type": "text", "text": "..."}]`) instead of a plain string.
`_extract_tool_result_content` and `_swap_tool_result_content` in
`compression_cache.py` only handled the plain string case, so every
tool_result was skipped before compression — zero savings for all
`headroom wrap claude` users.

Fix: extract text from list-of-blocks content in
`_extract_tool_result_content`, and collapse the list to a single text
block in `_swap_tool_result_content` when replacing with compressed
content.

## Type of Change

- [x] Bug fix (non-breaking)
- [ ] New feature (non-breaking)
- [ ] Breaking change
- [ ] Documentation update

## Changes Made

- `headroom/cache/compression_cache.py`:
- Added `_extract_text_from_blocks()` helper to extract joined text from
Anthropic list-of-blocks format
- Updated `_extract_tool_result_content()` to handle list content in
both Anthropic tool_result blocks and OpenAI role=tool messages
- Updated `_swap_tool_result_content()` to collapse list-of-blocks to a
single text block when replacing content (preserves structure, prevents
multi-block join mismatch)
- `tests/test_token_headroom_mode.py`: Added 12 tests covering Anthropic
list-of-blocks, OpenAI list, mixed blocks, empty list, non-mutation,
missing-text-block fallback, and non-tool messages

## Testing

- [x] Existing tests pass
- [x] New tests cover the fix
- [x] PBT round-trip property verified (6 properties × 850+ random
examples)
- [x] Adversarial edge case tests pass (50 cases across 3 functions)

```
tests/test_token_headroom_mode.py ....................... 32 passed (0.35s)
tests/test_transforms_content_router.py ................. 37 passed (1.08s)
tests/test_backend_bugs.py ............................... 37 passed (4.86s)
```

## Real Behavior Proof

- Environment: headroom main (upstream/main at time of PR), uv-managed
Python 3.12
- Exact command / steps:
1. `uv run python -m pytest tests/test_token_headroom_mode.py -x -q
--no-header`
2. `uv run python /tmp/pbt_tool_result_content.py` (PBT round-trip, 200
examples each property)
3. `uv run python /tmp/adversarial_lens_security.py` (50 edge case
checks)
4. Design scan: grep'd codebase for `isinstance(content, str)` near
tool_result — 3 sibling functions (`_to_text`, `_block_text`) already
handle list-of-blocks
- Observed result: All 106 tests pass. PBT confirmed extract+swap
round-trip invariant holds for 850+ random inputs. 50 adversarial edge
cases (null/None/nested/unicode/100K chars/1000 blocks) produce no
crashes or wrong output.
- Not tested: live proxy with real Claude Code traffic (the
ContentRouter already handles list content correctly since v0.32.0,
confirmed by code audit)

## Review Readiness

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

---------

Co-authored-by: lennney <lennney@users.noreply.github.com>
Co-authored-by: JD Davis <mxjerrett@gmail.com>
This commit is contained in:
GUOHAO LIU 2026-07-14 11:51:01 +08:00 committed by GitHub
parent e9000863fc
commit 3bcef2be37
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 189 additions and 2 deletions

View file

@ -40,12 +40,29 @@ def _is_tool_result_message(msg: dict) -> bool:
return False
def _extract_text_from_blocks(blocks: list) -> str | None:
"""Extract joined text from a list-of-blocks content (e.g. Anthropic list-of-text-blocks).
Modern Claude Code sends ``tool_result`` content as a list of typed
blocks (``[{"type": "text", "text": "..."}]``) instead of a plain
string. This helper extracts text from ``type == "text"`` blocks and
joins them.
"""
texts = [b.get("text", "") for b in blocks if isinstance(b, dict) and b.get("type") == "text"]
return "\n".join(t for t in texts if t != "") or None
def _extract_tool_result_content(msg: dict) -> str | None:
"""Extract text content from a tool result message (both formats)."""
# OpenAI format
if msg.get("role") == "tool":
content = msg.get("content")
return content if isinstance(content, str) else None
if isinstance(content, str):
return content
# OpenAI content can also be a list of content parts
if isinstance(content, list):
return _extract_text_from_blocks(content)
return None
# Anthropic format
content = msg.get("content")
if isinstance(content, list):
@ -54,6 +71,8 @@ def _extract_tool_result_content(msg: dict) -> str | None:
inner = block.get("content")
if isinstance(inner, str):
return inner
if isinstance(inner, list):
return _extract_text_from_blocks(inner)
return None
@ -69,6 +88,15 @@ def _swap_tool_result_content(msg: dict, new_content: str) -> dict:
if isinstance(content, list):
for block in content:
if isinstance(block, dict) and block.get("type") == "tool_result":
inner = block.get("content")
if isinstance(inner, list):
# Collapse list-of-blocks to a single text block.
# The compressed content is a single string; preserving
# multiple text blocks would produce a different joined
# output on re-extraction (first text block replaced,
# remaining text blocks still joined).
block["content"] = [{"type": "text", "text": new_content}]
else:
block["content"] = new_content
break
return new_msg

View file

@ -400,3 +400,162 @@ class TestProseFormatLiveZoneInvariant:
]
# Walk: user (stable, 1), tool (cached, 2). Cap → 1.
assert cache.compute_frozen_count(messages) == 1
# ── List-of-blocks tool_result content (Claude Code modern format) ──────────
def _make_tool_result_list_content_msg(tool_id: str, texts: list[str]) -> dict:
"""Anthropic-format tool result with list-of-blocks content.
Modern Claude Code sends ``tool_result`` content as a list of typed
blocks instead of a plain string.
"""
return {
"role": "user",
"content": [
{
"type": "tool_result",
"tool_use_id": tool_id,
"content": [{"type": "text", "text": t} for t in texts],
}
],
}
def _make_openai_tool_list_content_msg(tool_call_id: str, texts: list[str]) -> dict:
"""OpenAI-format tool message with list-of-blocks content."""
return {
"role": "tool",
"tool_call_id": tool_call_id,
"content": [{"type": "text", "text": t} for t in texts],
}
class TestExtractToolResultListContent:
"""_extract_tool_result_content handles list-of-blocks content."""
def test_anthropic_string_content_preserved(self):
"""Plain string content in Anthropic format still works."""
from headroom.cache.compression_cache import _extract_tool_result_content as f
msg = _make_tool_result_msg("t1", "hello world")
assert f(msg) == "hello world"
def test_anthropic_list_content_extracted(self):
"""List-of-blocks content is extracted and joined."""
from headroom.cache.compression_cache import _extract_tool_result_content as f
msg = _make_tool_result_list_content_msg("t1", ["Line 1", "Line 2"])
assert f(msg) == "Line 1\nLine 2"
def test_anthropic_mixed_blocks(self):
"""Non-text blocks (e.g. image) are skipped, only text blocks joined."""
from headroom.cache.compression_cache import _extract_tool_result_content as f
msg = {
"role": "user",
"content": [
{
"type": "tool_result",
"tool_use_id": "t1",
"content": [
{"type": "text", "text": "Hello"},
{"type": "image", "source": {"type": "base64", "data": "..."}},
{"type": "text", "text": "World"},
],
}
],
}
assert f(msg) == "Hello\nWorld"
def test_anthropic_list_empty_returns_none(self):
"""Empty text-only list returns None."""
from headroom.cache.compression_cache import _extract_tool_result_content as f
msg = {
"role": "user",
"content": [{"type": "tool_result", "tool_use_id": "t1", "content": []}],
}
assert f(msg) is None
def test_openai_list_content_extracted(self):
"""OpenAI format tool message with list content is extracted."""
from headroom.cache.compression_cache import _extract_tool_result_content as f
msg = _make_openai_tool_list_content_msg("tc1", ["Result 1", "Result 2"])
assert f(msg) == "Result 1\nResult 2"
def test_openai_string_content_still_works(self):
"""OpenAI format with plain string content is unchanged."""
from headroom.cache.compression_cache import _extract_tool_result_content as f
msg = _make_openai_tool_msg("tc1", "plain result")
assert f(msg) == "plain result"
def test_non_tool_msg_returns_none(self):
"""Regular user/assistant messages return None."""
from headroom.cache.compression_cache import _extract_tool_result_content as f
assert f(_make_user_msg("hello")) is None
assert f(_make_assistant_msg("response")) is None
class TestSwapToolResultListContent:
"""_swap_tool_result_content preserves list-of-blocks structure."""
def test_swap_anthropic_list_content_preserves_structure(self):
"""Swap on list-of-blocks content replaces text in place."""
from headroom.cache.compression_cache import _swap_tool_result_content
msg = _make_tool_result_list_content_msg("t1", ["original"])
swapped = _swap_tool_result_content(msg, "compressed")
inner = swapped["content"][0]["content"]
assert isinstance(inner, list)
assert inner[0]["type"] == "text"
assert inner[0]["text"] == "compressed"
def test_swap_anthropic_string_content_preserved(self):
"""Swap on plain-string content still works."""
from headroom.cache.compression_cache import _swap_tool_result_content
msg = _make_tool_result_msg("t1", "original")
swapped = _swap_tool_result_content(msg, "compressed")
assert swapped["content"][0]["content"] == "compressed"
def test_swap_openai_list_content(self):
"""Swap on OpenAI list-content replaces text."""
from headroom.cache.compression_cache import _swap_tool_result_content
msg = _make_openai_tool_list_content_msg("tc1", ["original"])
swapped = _swap_tool_result_content(msg, "compressed")
assert swapped["content"] == "compressed"
def test_swap_does_not_mutate_original(self):
"""_swap_tool_result_content performs a deep copy."""
from headroom.cache.compression_cache import _swap_tool_result_content
msg = _make_tool_result_list_content_msg("t1", ["original"])
_swap_tool_result_content(msg, "compressed")
assert msg["content"][0]["content"][0]["text"] == "original"
def test_swap_list_content_adds_text_block_when_missing(self):
"""When list has no text block, collapses to a single text block."""
from headroom.cache.compression_cache import _swap_tool_result_content
msg = {
"role": "user",
"content": [
{
"type": "tool_result",
"tool_use_id": "t1",
"content": [{"type": "image", "source": {"type": "base64", "data": "..."}}],
}
],
}
swapped = _swap_tool_result_content(msg, "compressed")
inner = swapped["content"][0]["content"]
assert isinstance(inner, list)
assert len(inner) == 1
assert inner[0]["type"] == "text"
assert inner[0]["text"] == "compressed"