diff --git a/headroom/cache/compression_cache.py b/headroom/cache/compression_cache.py index a5ec4bf53..71beb4091 100644 --- a/headroom/cache/compression_cache.py +++ b/headroom/cache/compression_cache.py @@ -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,7 +88,16 @@ 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": - block["content"] = new_content + 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 diff --git a/tests/test_token_headroom_mode.py b/tests/test_token_headroom_mode.py index 300d8a48a..2198aba39 100644 --- a/tests/test_token_headroom_mode.py +++ b/tests/test_token_headroom_mode.py @@ -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"