diff --git a/headroom/config.py b/headroom/config.py index 2a61e0abe..e673823e8 100644 --- a/headroom/config.py +++ b/headroom/config.py @@ -209,6 +209,7 @@ class AnchorConfig: # Tool outputs that are reference data and must NOT be compressed. # Read/Glob/Grep contain exact file contents/search results the agent needs for edits. # Write/Edit record what changes were made — compressing them causes duplicate/conflicting edits. +# WebSearch/WebFetch results are large reference payloads that must remain verbatim. # Bash is NOT excluded — its outputs (build logs, test output) are ideal compression targets. # To protect Bash or other non-excluded tools from lossy compression, use # HEADROOM_PROTECT_TOOL_RESULTS=Bash or --protect-tool-results Bash. @@ -219,12 +220,27 @@ DEFAULT_EXCLUDE_TOOLS: frozenset[str] = frozenset( "Grep", "Write", "Edit", + "WebSearch", + "WebFetch", # Lowercase variants for case-insensitive matching "read", "glob", "grep", "write", "edit", + "web_search", + "web_fetch", + } +) + +# These excluded web-tool results must remain byte-faithful. Even the +# excluded-tool lossless fold rewrites formatted JSON. +DEFAULT_VERBATIM_EXCLUDE_TOOLS: frozenset[str] = frozenset( + { + "WebSearch", + "WebFetch", + "web_search", + "web_fetch", } ) diff --git a/headroom/proxy/handlers/openai.py b/headroom/proxy/handlers/openai.py index 1e6fa37bc..88de6c1d9 100644 --- a/headroom/proxy/handlers/openai.py +++ b/headroom/proxy/handlers/openai.py @@ -805,6 +805,7 @@ def _dedup_responses_output_items( items: list[dict[str, Any]], output_types: frozenset[str], count_tokens: Any = None, + protected_call_ids: set[str] | None = None, ) -> tuple[int, int]: """Cross-turn verbatim de-dup over Responses tool-output items (mutates in place). @@ -834,7 +835,17 @@ def _dedup_responses_output_items( out = item.get("output") if isinstance(out, str) and out: locs.append(i) - blocks.append(DedupBlock(text=out, turn=i, protected=False)) + blocks.append( + DedupBlock( + text=out, + turn=i, + protected=bool( + isinstance(item.get("call_id"), str) + and protected_call_ids + and item.get("call_id") in protected_call_ids + ), + ) + ) if len(blocks) < 2: return 0, 0 @@ -1372,7 +1383,11 @@ class OpenAIHandlerMixin: # mirroring ContentRouter's policy. exclude_tools already contains both # original and lowercased name variants (see _parse_exclude_tools), but # we also test the lowercased name defensively for case-insensitivity. - from headroom.config import DEFAULT_EXCLUDE_TOOLS, is_tool_excluded + from headroom.config import ( + DEFAULT_EXCLUDE_TOOLS, + DEFAULT_VERBATIM_EXCLUDE_TOOLS, + is_tool_excluded, + ) router_exclude_tools = getattr(router.config, "exclude_tools", None) effective_exclude_tools = ( @@ -1383,6 +1398,11 @@ class OpenAIHandlerMixin: for call_id, fn_name in function_name_by_call_id.items() if is_tool_excluded(fn_name, effective_exclude_tools) } + verbatim_excluded_call_ids: set[str] = { + call_id + for call_id, fn_name in function_name_by_call_id.items() + if is_tool_excluded(fn_name, DEFAULT_VERBATIM_EXCLUDE_TOOLS) + } timing_sink: dict[str, float] = timing if timing is not None else {} @@ -1428,6 +1448,20 @@ class OpenAIHandlerMixin: ) continue if isinstance(call_id, str) and call_id in excluded_call_ids: + if call_id in verbatim_excluded_call_ids: + if debug_enabled: + extraction_debug.append( + { + "index": idx, + "eligible": False, + "reason": "exclude_tools_verbatim", + "item_type": item_type, + "call_id": call_id, + "tool_name": function_name_by_call_id.get(call_id), + "item": item, + } + ) + continue # Protected from lossy compression — but grep/log/json output # can still be losslessly compacted. Reuse the router helper # so the Responses path matches the chat/Anthropic behavior. @@ -1824,7 +1858,10 @@ class OpenAIHandlerMixin: # chat path (ContentRouter._cross_turn_dedup_messages runs last there too). if getattr(router, "_cross_turn_dedup_enabled", False): dd_folded, dd_saved = _dedup_responses_output_items( - updated_items, self.OPENAI_RESPONSES_OUTPUT_TYPES, tokenizer.count_text + updated_items, + self.OPENAI_RESPONSES_OUTPUT_TYPES, + tokenizer.count_text, + protected_call_ids=verbatim_excluded_call_ids, ) if dd_folded: modified = True diff --git a/headroom/transforms/content_router.py b/headroom/transforms/content_router.py index 3ca85b0c0..5a697446e 100644 --- a/headroom/transforms/content_router.py +++ b/headroom/transforms/content_router.py @@ -51,6 +51,7 @@ from typing import Any from ..config import ( DEFAULT_EXCLUDE_TOOLS, + DEFAULT_VERBATIM_EXCLUDE_TOOLS, ReadLifecycleConfig, RelevanceScorerConfig, TransformResult, @@ -3567,6 +3568,12 @@ class ContentRouter(Transform): if role == "tool": tool_call_id = message.get("tool_call_id", "") if tool_call_id in excluded_tool_ids: + tool_name = tool_name_map.get(tool_call_id, "") + if tool_name and is_tool_excluded(tool_name, DEFAULT_VERBATIM_EXCLUDE_TOOLS): + result_slots[i] = message + transforms_applied.append("router:excluded:tool") + route_counts["excluded_tool"] += 1 + continue if messages_from_end <= read_protection_window: # Protected from lossy compression — but grep/log/json # output can still be losslessly compacted. @@ -4165,6 +4172,12 @@ class ContentRouter(Transform): locs: list[tuple[int, int | None, int | None]] = [] dblocks: list[DedupBlock] = [] + tool_name_map = self._build_tool_name_map(messages) + verbatim_tool_ids = { + tool_id + for tool_id, name in tool_name_map.items() + if is_tool_excluded(name, DEFAULT_VERBATIM_EXCLUDE_TOOLS) + } def _is_user_read_observation(idx: int) -> bool: # A file read can land in a plain ``role:user`` STRING (text @@ -4191,7 +4204,11 @@ class ContentRouter(Transform): if not isinstance(block, dict) or block.get("type") != "tool_result": continue tc = block.get("content") - protected = frozen or ("cache_control" in block) + protected = ( + frozen + or ("cache_control" in block) + or block.get("tool_use_id") in verbatim_tool_ids + ) if isinstance(tc, str) and tc: locs.append((i, bidx, None)) dblocks.append(DedupBlock(text=tc, turn=i, protected=protected)) @@ -4225,7 +4242,11 @@ class ContentRouter(Transform): if role in ("tool", "function") or ( role == "user" and _is_user_read_observation(i) ): - protected = frozen or ("cache_control" in msg) + protected = ( + frozen + or ("cache_control" in msg) + or msg.get("tool_call_id") in verbatim_tool_ids + ) locs.append((i, None, None)) dblocks.append(DedupBlock(text=content, turn=i, protected=protected)) @@ -4414,6 +4435,13 @@ class ContentRouter(Transform): route_counts["read_protected"] += 1 continue if tool_use_id in excluded_tool_ids: + tool_name = tool_name_map.get(tool_use_id, "") if tool_name_map else "" + if tool_name and is_tool_excluded(tool_name, DEFAULT_VERBATIM_EXCLUDE_TOOLS): + new_blocks.append(block) + transforms_applied.append("router:excluded:tool") + if route_counts is not None: + route_counts["excluded_tool"] += 1 + continue if messages_from_end <= read_protection_window: # Protected from lossy compression — but grep/log/json # output can still be losslessly compacted. diff --git a/tests/test_openai_responses_compression_units.py b/tests/test_openai_responses_compression_units.py index 9ac2bc162..c2ae3f294 100644 --- a/tests/test_openai_responses_compression_units.py +++ b/tests/test_openai_responses_compression_units.py @@ -620,6 +620,47 @@ def test_openai_responses_adapter_excludes_tool_case_insensitively_with_debug(mo assert new_payload == payload +def test_openai_responses_adapter_keeps_websearch_output_verbatim(): + """Default-excluded web tools must bypass both lossy and lossless rewrites.""" + router = ContentRouter() + + def compress(self, content: str, **_kwargs): + return RouterCompressionResult( + compressed="should not be used", + original=content, + strategy_used=CompressionStrategy.KOMPRESS, + ) + + router.compress = MethodType(compress, router) + handler = _handler_with_router(router) + output = ( + "{\n" + ' "results": [\n' + ' {"title": "Headroom", "snippet": "structured web payload with spacing that must remain verbatim"}\n' + " ]\n" + "}" + ) + payload = { + "model": "gpt-5", + "input": [ + {"type": "function_call", "call_id": "call_1", "name": "WebSearch", "arguments": "{}"}, + {"type": "function_call_output", "call_id": "call_1", "output": output}, + ], + } + + new_payload, modified, saved, *_ = ( + handler._compress_openai_responses_live_text_units_with_router( + payload, + model="gpt-5", + request_id="req_test", + ) + ) + + assert modified is False + assert saved == 0 + assert new_payload == payload + + def test_openai_responses_adapter_compresses_non_excluded_tool_outputs(): """Only excluded tools are protected; other tool outputs still compress.""" router = ContentRouter() diff --git a/tests/test_responses_cross_turn_dedup.py b/tests/test_responses_cross_turn_dedup.py index 276531f61..490a93a39 100644 --- a/tests/test_responses_cross_turn_dedup.py +++ b/tests/test_responses_cross_turn_dedup.py @@ -97,6 +97,32 @@ def test_single_read_does_not_fold(): assert items[1]["output"] == _wrap("492f0f", "0.0000") + BODY +def test_protected_websearch_outputs_do_not_fold(): + items = [ + { + "type": "function_call_output", + "call_id": "c1", + "output": '{\n "results": [\n {"title": "Headroom"}\n ]\n}', + }, + { + "type": "function_call_output", + "call_id": "c2", + "output": '{\n "results": [\n {"title": "Headroom"}\n ]\n}', + }, + ] + folded, saved = _dedup_responses_output_items( + items, + _RESPONSES_OUTPUT_ITEM_TYPES, + count_tokens=len, + protected_call_ids={"c1", "c2"}, + ) + + assert folded == 0 + assert saved == 0 + assert items[0]["output"].endswith('{"title": "Headroom"}\n ]\n}') + assert items[1]["output"].endswith('{"title": "Headroom"}\n ]\n}') + + def test_non_output_items_untouched(): # A duplicated MESSAGE (not a tool output) must never fold — only output # items are eligible. diff --git a/tests/test_websearch_tool_result_protection.py b/tests/test_websearch_tool_result_protection.py new file mode 100644 index 000000000..a270b9a24 --- /dev/null +++ b/tests/test_websearch_tool_result_protection.py @@ -0,0 +1,184 @@ +"""Regression tests for web-tool result passthrough.""" + +from __future__ import annotations + +from headroom.config import DEFAULT_EXCLUDE_TOOLS +from headroom.proxy.server import HeadroomProxy, ProxyConfig +from headroom.transforms.content_detector import ContentType +from headroom.transforms.content_router import ( + CompressionStrategy, + ContentRouter, + RouterCompressionResult, + RoutingDecision, +) + + +class _Tokenizer: + def count_text(self, text: str) -> int: + return max(1, len(text) // 4) + + +def _messages(tool_name: str, payload: str) -> list[dict[str, object]]: + return [ + { + "role": "assistant", + "content": [ + { + "type": "tool_use", + "id": "tool-1", + "name": tool_name, + "input": {}, + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "tool-1", + "content": payload, + } + ], + }, + ] + + +def _router() -> ContentRouter: + proxy = HeadroomProxy( + ProxyConfig( + optimize=False, + cache_enabled=False, + rate_limit_enabled=False, + cost_tracking_enabled=False, + code_aware_enabled=False, + mode="token", + ) + ) + router = proxy.anthropic_pipeline.transforms[-1] + assert isinstance(router, ContentRouter) + router.config.min_section_tokens = 1 + router.config.min_chars_for_block_compression = 1 + return router + + +def test_web_tools_are_default_exclusions() -> None: + assert {"WebSearch", "WebFetch", "web_search", "web_fetch"} <= DEFAULT_EXCLUDE_TOOLS + + +def test_web_tool_results_bypass_compressor() -> None: + router = _router() + calls = 0 + + def fake_compress(*args: object, **kwargs: object) -> RouterCompressionResult: + nonlocal calls + calls += 1 + content = str(args[0]) + return RouterCompressionResult( + compressed="mutated", + original=content, + strategy_used=CompressionStrategy.TEXT, + routing_log=[ + RoutingDecision(ContentType.PLAIN_TEXT, CompressionStrategy.TEXT, 100, 10) + ], + ) + + router.compress = fake_compress # type: ignore[method-assign] + payload = ( + "{\n" + ' "results": [\n' + ' {"title": "Headroom", "snippet": "reference payload reference payload reference payload"},\n' + ' {"title": "Docs", "snippet": "structured web payload with spacing that must remain verbatim"}\n' + " ],\n" + ' "source": "web"\n' + "}" + ) + + for tool_name in ("WebSearch", "WebFetch", "web_search", "web_fetch"): + messages = _messages(tool_name, payload) + result = router.apply(messages, _Tokenizer()) + + tool_result = result.messages[1]["content"][0] # type: ignore[index] + assert tool_result["content"] == payload # type: ignore[index] + assert "router:excluded:tool" in result.transforms_applied + + assert calls == 0 + + +def test_web_tool_results_stay_verbatim_outside_token_age_window() -> None: + router = _router() + calls = 0 + + def fake_compress(*args: object, **kwargs: object) -> RouterCompressionResult: + nonlocal calls + calls += 1 + content = str(args[0]) + return RouterCompressionResult( + compressed="mutated", + original=content, + strategy_used=CompressionStrategy.TEXT, + routing_log=[ + RoutingDecision(ContentType.PLAIN_TEXT, CompressionStrategy.TEXT, 100, 10) + ], + ) + + router.compress = fake_compress # type: ignore[method-assign] + payload = ( + "{\n" + ' "results": [\n' + ' {"title": "Headroom", "snippet": "reference payload reference payload reference payload"}\n' + " ]\n" + "}" + ) + messages = _messages("WebSearch", payload) + messages.extend({"role": "user", "content": f"later turn {i}"} for i in range(18)) + + result = router.apply(messages, _Tokenizer()) + tool_result = result.messages[1]["content"][0] # type: ignore[index] + + assert tool_result["content"] == payload # type: ignore[index] + assert calls == 0 + + +def test_web_tool_results_skip_cross_turn_dedup() -> None: + router = _router() + payload = ( + "{\n" + ' "results": [\n' + ' {"title": "Headroom", "snippet": "structured web payload with spacing that must remain verbatim"}\n' + " ]\n" + "}" + ) + messages = _messages("WebSearch", payload) + _messages("WebSearch", payload) + + result = router.apply(messages, _Tokenizer()) + + first = result.messages[1]["content"][0] # type: ignore[index] + second = result.messages[3]["content"][0] # type: ignore[index] + assert first["content"] == payload # type: ignore[index] + assert second["content"] == payload # type: ignore[index] + + +def test_bash_remains_compressible() -> None: + router = _router() + calls = 0 + + def fake_compress(*args: object, **kwargs: object) -> RouterCompressionResult: + nonlocal calls + calls += 1 + content = str(args[0]) + return RouterCompressionResult( + compressed="compressed bash output", + original=content, + strategy_used=CompressionStrategy.TEXT, + routing_log=[ + RoutingDecision(ContentType.PLAIN_TEXT, CompressionStrategy.TEXT, 100, 10) + ], + ) + + router.compress = fake_compress # type: ignore[method-assign] + payload = "bash output " * 100 + result = router.apply(_messages("Bash", payload), _Tokenizer()) + + assert calls == 1 + assert "router:excluded:tool" not in result.transforms_applied