diff --git a/headroom/integrations/langchain/langgraph.py b/headroom/integrations/langchain/langgraph.py index 6fcdeff62..2e6734795 100644 --- a/headroom/integrations/langchain/langgraph.py +++ b/headroom/integrations/langchain/langgraph.py @@ -47,6 +47,8 @@ except ImportError: BaseMessage = object # type: ignore[misc,assignment] ToolMessage = object # type: ignore[misc,assignment] +from headroom.ccr.tool_injection import CCR_TOOL_NAME +from headroom.config import is_tool_excluded from headroom.transforms.smart_crusher import SmartCrusher, SmartCrusherConfig logger = logging.getLogger(__name__) @@ -156,6 +158,7 @@ def _get_crusher(min_tokens: int) -> SmartCrusher: def _should_skip( content: str, config: CompressToolMessagesConfig, + tool_name: str | None = None, ) -> str | None: """Check if a ToolMessage should skip compression. @@ -164,6 +167,9 @@ def _should_skip( if not content: return "empty_content" + if tool_name and is_tool_excluded(tool_name, (CCR_TOOL_NAME,)): + return "tool_excluded" + tokens = _estimate_tokens(content) if tokens < config.min_tokens_to_compress: return f"below_threshold:{tokens}<{config.min_tokens_to_compress}" @@ -176,6 +182,18 @@ def _should_skip( return None +def _tool_names_by_id(messages: list[BaseMessage]) -> dict[str, str]: # type: ignore[type-arg] + """Index tool-call names so results without a copied name remain classifiable.""" + names: dict[str, str] = {} + for message in messages: + for tool_call in getattr(message, "tool_calls", ()) or (): + tool_call_id = tool_call.get("id") + tool_name = tool_call.get("name") + if tool_call_id and tool_name: + names[tool_call_id] = tool_name + return names + + def compress_tool_messages( messages: list[BaseMessage], # type: ignore[type-arg] *, @@ -223,6 +241,7 @@ def compress_tool_messages( crusher = _get_crusher(config.min_tokens_to_compress) result_messages: list[BaseMessage] = [] metrics: list[ToolMessageCompressionMetrics] = [] + tool_names = _tool_names_by_id(messages) for msg in messages: if not isinstance(msg, ToolMessage): @@ -233,7 +252,8 @@ def compress_tool_messages( request_id = str(uuid4()) # Check if we should skip - skip_reason = _should_skip(content, config) + tool_name = getattr(msg, "name", None) or tool_names.get(getattr(msg, "tool_call_id", "")) + skip_reason = _should_skip(content, config, tool_name) if skip_reason: result_messages.append(msg) tokens = _estimate_tokens(content) diff --git a/headroom/integrations/strands/hooks.py b/headroom/integrations/strands/hooks.py index 29dc32a03..275b6f64d 100644 --- a/headroom/integrations/strands/hooks.py +++ b/headroom/integrations/strands/hooks.py @@ -48,6 +48,8 @@ except ImportError: ToolResult = dict # type: ignore[misc,assignment] from headroom import HeadroomConfig +from headroom.ccr.tool_injection import CCR_TOOL_NAME +from headroom.config import is_tool_excluded from headroom.transforms.smart_crusher import SmartCrusher, SmartCrusherConfig logger = logging.getLogger(__name__) @@ -314,7 +316,7 @@ class HeadroomHookProvider(HookProvider): # type: ignore[misc] tool_use_id = event.tool_use.get("toolUseId", "unknown") # Check if compression should be skipped - skip_reason = self._should_skip_compression(result) + skip_reason = self._should_skip_compression(result, tool_name) if skip_reason: self._record_metrics( request_id=request_id, @@ -421,11 +423,15 @@ class HeadroomHookProvider(HookProvider): # type: ignore[misc] tokens_before, ) - def _should_skip_compression(self, result: ToolResult) -> str | None: + def _should_skip_compression( + self, result: ToolResult, tool_name: str | None = None + ) -> str | None: """Check if compression should be skipped for this result. Args: result: The tool result to check. + tool_name: The name of the tool that produced the result, used to + honor the shared tool-exclusion set. Returns: Skip reason string if should skip, None if should compress. @@ -434,6 +440,9 @@ class HeadroomHookProvider(HookProvider): # type: ignore[misc] if not self.compress_tool_outputs: return "compression_disabled" + if tool_name and is_tool_excluded(tool_name, (CCR_TOOL_NAME,)): + return "tool_excluded" + # Skip error results if preserve_errors is True if self.preserve_errors and result.get("status") == "error": return "error_result_preserved" diff --git a/headroom/transforms/smart_crusher.py b/headroom/transforms/smart_crusher.py index 6e4ae6d1c..4dde88736 100644 --- a/headroom/transforms/smart_crusher.py +++ b/headroom/transforms/smart_crusher.py @@ -52,7 +52,7 @@ from dataclasses import dataclass from typing import Any from ..ccr.tool_injection import CCR_TOOL_NAME -from ..config import CCRConfig, TransformResult +from ..config import CCRConfig, TransformResult, is_tool_excluded from ..tokenizer import Tokenizer from ..utils import compute_short_hash, create_tool_digest_marker, deep_copy_messages from .base import Transform @@ -1280,7 +1280,10 @@ class SmartCrusher(Transform): # unresolvable retrieval loop. # ponytail: ceiling is tool_call_id lookup; if the id is missing we # compress (conservative: unknown tool names don't get a free pass). - if tool_names_by_id.get(msg.get("tool_call_id") or "") == CCR_TOOL_NAME: + if is_tool_excluded( + tool_names_by_id.get(msg.get("tool_call_id") or "") or "", + (CCR_TOOL_NAME,), + ): continue content = msg.get("content", "") if isinstance(content, str): @@ -1310,7 +1313,10 @@ class SmartCrusher(Transform): # would produce a new <> marker the agent cannot # redeem (infinite retrieval loop). # ponytail: ceiling is tool_use_id lookup; unknown ids pass through. - if tool_names_by_id.get(block.get("tool_use_id") or "") == CCR_TOOL_NAME: + if is_tool_excluded( + tool_names_by_id.get(block.get("tool_use_id") or "") or "", + (CCR_TOOL_NAME,), + ): continue tool_content = block.get("content", "") if not isinstance(tool_content, str): diff --git a/tests/integrations/test_langgraph.py b/tests/integrations/test_langgraph.py new file mode 100644 index 000000000..6f648a247 --- /dev/null +++ b/tests/integrations/test_langgraph.py @@ -0,0 +1,90 @@ +"""Regression tests for qualified CCR retrieval tool names in LangGraph.""" + +from __future__ import annotations + +import json + +import pytest + +pytest.importorskip("headroom._core") + +try: + from langchain_core.messages import AIMessage, ToolMessage +except ImportError: + pytest.skip("LangChain not installed", allow_module_level=True) + +from headroom.integrations.langchain.langgraph import compress_tool_messages + + +def _large_output() -> str: + return json.dumps([{"id": i, "name": f"item_{i}", "value": "x" * 30} for i in range(200)]) + + +def _messages(tool_name: str) -> list: + return [ + AIMessage(content="", tool_calls=[{"id": "call_1", "name": tool_name, "args": {}}]), + ToolMessage(content=_large_output(), tool_call_id="call_1"), + ] + + +@pytest.mark.parametrize( + "tool_name", + ["mcp__Headroom__headroom_retrieve", "mcp_Headroom_headroom_retrieve"], +) +def test_qualified_ccr_retrieval_message_is_preserved(tool_name: str) -> None: + messages = _messages(tool_name) + original = messages[1].content + + result = compress_tool_messages(messages) + + assert result.messages[1].content == original + assert result.metrics[0].skip_reason == "tool_excluded" + + +def test_incomplete_tool_calls_do_not_hide_later_qualified_name() -> None: + messages = [ + AIMessage( + content="", + tool_calls=[ + {"id": None, "name": "incomplete", "args": {}}, + {"id": "ignored", "name": "", "args": {}}, + { + "id": "call_1", + "name": "mcp__Headroom__headroom_retrieve", + "args": {}, + }, + ], + ), + ToolMessage(content=_large_output(), tool_call_id="call_1"), + ] + original = messages[1].content + + result = compress_tool_messages(messages) + + assert result.messages[1].content == original + assert result.metrics[0].skip_reason == "tool_excluded" + + +def test_near_match_ccr_tool_name_is_not_excluded() -> None: + messages = _messages("mcp__Headroom__headroom_retrieve_extra") + original = messages[1].content + + result = compress_tool_messages(messages) + + assert result.metrics[0].skip_reason != "tool_excluded" + assert result.messages[1].content != original + + +@pytest.mark.parametrize( + "tool_name", + ["mcp__Headroom__headroom_retrieve", "mcp_Headroom_headroom_retrieve"], +) +def test_qualified_name_on_the_tool_message_is_enough(tool_name: str) -> None: + """`ToolNode` populates `ToolMessage.name`, so the id index is only a fallback.""" + messages = [ToolMessage(content=_large_output(), tool_call_id="call_1", name=tool_name)] + original = messages[0].content + + result = compress_tool_messages(messages) + + assert result.messages[0].content == original + assert result.metrics[0].skip_reason == "tool_excluded" diff --git a/tests/integrations/test_strands/test_ccr_exclusion.py b/tests/integrations/test_strands/test_ccr_exclusion.py new file mode 100644 index 000000000..56c1d05d7 --- /dev/null +++ b/tests/integrations/test_strands/test_ccr_exclusion.py @@ -0,0 +1,91 @@ +"""Regression tests for qualified CCR names in Strands compression.""" + +from __future__ import annotations + +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest + +pytest.importorskip("headroom._core") + +from headroom.integrations.strands import hooks + + +@pytest.fixture +def hook(monkeypatch: pytest.MonkeyPatch) -> hooks.HeadroomHookProvider: + monkeypatch.setattr(hooks, "STRANDS_AVAILABLE", True) + provider = hooks.HeadroomHookProvider(min_tokens_to_compress=0) + provider._crusher = Mock() + provider._crusher.crush.return_value = SimpleNamespace( + compressed="compressed", was_modified=True + ) + return provider + + +class _Registry: + """Stand-in for Strands' HookRegistry that dispatches like the real one.""" + + def __init__(self) -> None: + self._callbacks: dict[object, list] = {} + + def add_callback(self, event_type: object, callback) -> None: # noqa: ANN001 + self._callbacks.setdefault(event_type, []).append(callback) + + def dispatch(self, event_type: object, event: object) -> None: + for callback in self._callbacks[event_type]: + callback(event) + + +def _event(tool_name: str, content: str) -> SimpleNamespace: + return SimpleNamespace( + tool_use={"name": tool_name, "toolUseId": "tool_1"}, + result={"content": [{"text": content}]}, + ) + + +def test_qualified_ccr_result_is_preserved_through_the_registered_hook( + hook: hooks.HeadroomHookProvider, +) -> None: + """Drive the seam Strands actually drives: register_hooks, then dispatch.""" + registry = _Registry() + hook.register_hooks(registry) + + content = "x" * 400 + event = _event("mcp__headroom__headroom_retrieve", content) + registry.dispatch(hooks.AfterToolCallEvent, event) + + assert event.result["content"][0]["text"] == content + assert hook.metrics_history[0].skip_reason == "tool_excluded" + hook._crusher.crush.assert_not_called() + + +@pytest.mark.parametrize( + "tool_name", + [ + "mcp__headroom__headroom_retrieve", + "mcp_headroom_headroom_retrieve", + "headroom_retrieve", + ], +) +def test_qualified_ccr_tool_result_is_preserved( + hook: hooks.HeadroomHookProvider, tool_name: str +) -> None: + content = "x" * 400 + event = _event(tool_name, content) + + hook._compress_tool_result(event) + + assert event.result["content"][0]["text"] == content + assert hook.metrics_history[0].skip_reason == "tool_excluded" + hook._crusher.crush.assert_not_called() + + +def test_near_match_tool_name_still_compresses(hook: hooks.HeadroomHookProvider) -> None: + event = _event("mcp__headroom__headroom_retrieve_extra", "x" * 400) + + hook._compress_tool_result(event) + + assert event.result["content"][0]["text"] == "compressed" + assert hook.metrics_history[0].was_compressed is True + hook._crusher.crush.assert_called_once() diff --git a/tests/integrations/test_strands/test_hooks_unit.py b/tests/integrations/test_strands/test_hooks_unit.py index b059b93c5..8a18a1c90 100644 --- a/tests/integrations/test_strands/test_hooks_unit.py +++ b/tests/integrations/test_strands/test_hooks_unit.py @@ -254,6 +254,30 @@ class TestShouldSkipCompression: skip_reason = hook._should_skip_compression(result) assert skip_reason is None + @pytest.mark.parametrize( + "tool_name", + ["mcp__Headroom__headroom_retrieve", "mcp_Headroom_headroom_retrieve"], + ) + def test_skip_qualified_ccr_retrieval_results(self, tool_name): + """Qualified CCR retrieval names use the shared exclusion authority.""" + from headroom.integrations.strands import HeadroomHookProvider + + hook = HeadroomHookProvider() + result = {"content": [{"text": "retrieved content"}]} + + assert hook._should_skip_compression(result, tool_name) == "tool_excluded" + + def test_near_match_ccr_tool_name_is_not_excluded(self): + """A similar qualified name remains eligible for compression.""" + from headroom.integrations.strands import HeadroomHookProvider + + hook = HeadroomHookProvider() + result = {"content": [{"text": "retrieved content"}]} + + assert ( + hook._should_skip_compression(result, "mcp__Headroom__headroom_retrieve_extra") is None + ) + class TestCompressToolResult: """Tests for _compress_tool_result hook handler.""" @@ -323,6 +347,24 @@ class TestCompressToolResult: assert metrics.was_compressed is False assert metrics.skip_reason == "compression_disabled" + def test_qualified_ccr_retrieval_result_is_preserved(self): + """The production hook skips qualified CCR retrieval results.""" + from headroom.integrations.strands import HeadroomHookProvider + + hook = HeadroomHookProvider(min_tokens_to_compress=1) + original = json.dumps([{"id": i, "data": "x" * 50} for i in range(50)]) + mock_event = MagicMock() + mock_event.tool_use = { + "name": "mcp__Headroom__headroom_retrieve", + "toolUseId": "tool-ccr", + } + mock_event.result = {"content": [{"text": original}]} + + hook._compress_tool_result(mock_event) + + assert mock_event.result["content"][0]["text"] == original + assert hook.metrics_history[-1].skip_reason == "tool_excluded" + class TestMetricsTracking: """Tests for metrics tracking and aggregation.""" diff --git a/tests/test_smart_crusher.py b/tests/test_smart_crusher.py new file mode 100644 index 000000000..0ba6d13cb --- /dev/null +++ b/tests/test_smart_crusher.py @@ -0,0 +1,141 @@ +"""Regression tests for qualified CCR retrieval tool names.""" + +from __future__ import annotations + +import json + +import pytest + +from headroom import OpenAIProvider, Tokenizer +from headroom.ccr.tool_injection import CCR_TOOL_NAME +from headroom.config import SmartCrusherConfig + +try: + from headroom._core import SmartCrusher as _RustSmartCrusher # noqa: F401 +except ImportError: + pytest.skip("headroom._core not built", allow_module_level=True) + +from headroom.transforms.smart_crusher import SmartCrusher + + +def _big_content() -> str: + return json.dumps([{"id": i, "value": "x" * 20} for i in range(60)]) + + +def _apply_for_tool(tool_name: str): + messages = [ + { + "role": "assistant", + "tool_calls": [ + { + "id": "call_1", + "function": {"name": tool_name, "arguments": "{}"}, + } + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": _big_content()}, + ] + tokenizer = Tokenizer(OpenAIProvider().get_token_counter("gpt-4o"), "gpt-4o") + result = SmartCrusher(config=SmartCrusherConfig(min_tokens_to_crush=0)).apply( + messages, tokenizer + ) + return messages[1]["content"], result + + +@pytest.mark.parametrize( + "tool_name", + ["mcp__Headroom__headroom_retrieve", "mcp_Headroom_headroom_retrieve"], +) +def test_qualified_ccr_retrieval_result_is_preserved(tool_name: str) -> None: + original, result = _apply_for_tool(tool_name) + + assert result.messages[1]["content"] == original + assert not any("smart_crush" in transform for transform in result.transforms_applied) + + +def test_near_match_ccr_tool_name_still_compresses() -> None: + original, result = _apply_for_tool("mcp__Headroom__headroom_retrieve_extra") + + assert result.messages[1]["content"] != original or result.tokens_after < result.tokens_before + + +def test_bare_ccr_tool_name_remains_preserved() -> None: + original, result = _apply_for_tool(CCR_TOOL_NAME) + + assert result.messages[1]["content"] == original + + +def _apply_anthropic_for_tool(tool_name: str): + """Anthropic block shape: tool_use in the assistant turn, tool_result in the user turn.""" + messages = [ + { + "role": "assistant", + "content": [ + {"type": "tool_use", "id": "tu_1", "name": tool_name, "input": {}}, + ], + }, + { + "role": "user", + "content": [ + {"type": "tool_result", "tool_use_id": "tu_1", "content": _big_content()}, + ], + }, + ] + tokenizer = Tokenizer(OpenAIProvider().get_token_counter("gpt-4o"), "gpt-4o") + result = SmartCrusher(config=SmartCrusherConfig(min_tokens_to_crush=0)).apply( + messages, tokenizer + ) + return messages[1]["content"][0]["content"], result + + +@pytest.mark.parametrize( + "tool_name", + [ + "mcp__Headroom__headroom_retrieve", + "mcp_Headroom_headroom_retrieve", + CCR_TOOL_NAME, + ], +) +def test_qualified_ccr_tool_result_block_is_preserved(tool_name: str) -> None: + original, result = _apply_anthropic_for_tool(tool_name) + + assert result.messages[1]["content"][0]["content"] == original + assert not any("smart" in transform for transform in result.transforms_applied) + + +@pytest.mark.parametrize( + "tool_name", + ["mcp__Headroom__headroom_retrieve", "mcp_Headroom_headroom_retrieve", CCR_TOOL_NAME], +) +def test_mcp_compressor_preserves_qualified_ccr_output(tool_name: str) -> None: + """`HeadroomMCPCompressor.compress` is the production entry point issue #2656 names. + + It drives `SmartCrusher.apply` with a `role=tool` message, so the guard has to + hold through that wrapper and not only on a directly built message list. + """ + from headroom.integrations.mcp.server import HeadroomMCPCompressor + + content = json.dumps({"results": [{"id": i, "value": "x" * 40} for i in range(80)]}) + result = HeadroomMCPCompressor().compress(content, tool_name=tool_name) + + assert result.compressed_content == content + + +def test_mcp_compressor_still_compresses_a_near_match_name() -> None: + from headroom.integrations.mcp.server import HeadroomMCPCompressor + + content = json.dumps({"results": [{"id": i, "value": "x" * 40} for i in range(80)]}) + result = HeadroomMCPCompressor().compress( + content, tool_name="mcp__Headroom__headroom_retrieve_extra" + ) + + assert result.compressed_content != content + + +def test_near_match_ccr_tool_result_block_still_compresses() -> None: + original, result = _apply_anthropic_for_tool("mcp__Headroom__headroom_retrieve_extra") + + assert ( + result.messages[1]["content"][0]["content"] != original + or result.tokens_after < result.tokens_before + )