fix(compression): honor qualified CCR names across integrations (#2698)

## Description

Three compression consumers compare tool names against the bare literal
`headroom_retrieve`, so the qualified forms MCP clients actually send
(`mcp__Headroom__headroom_retrieve`, `mcp_Headroom_headroom_retrieve`)
slip past the guard and get recompressed. `SmartCrusher.apply` has the
bare comparison at both its OpenAI `role=tool` site and its Anthropic
`tool_result` block site; the LangGraph compressor and the Strands hook
have no tool-name check at all. Recompressing already-retrieved CCR
content mints a new `<<ccr:hash>>` marker the agent cannot redeem.

`headroom.config.is_tool_excluded` already owns alias resolution,
including the MCP wrapper forms. This routes all three consumers through
it instead of adding a second name matcher. Closes #2656.

## 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

- `SmartCrusher.apply` routes both its `role=tool` and its Anthropic
`tool_result` guards through `is_tool_excluded`
- `_should_skip` in the LangGraph compressor takes the tool name and
skips excluded tools; tool-call names are indexed by id so a
`ToolMessage` without a copied `name` is still classifiable
- `_should_skip_compression` in the Strands hook takes the tool name and
skips excluded tools, recording `tool_excluded`
- regressions for the qualified and bare names across all three
consumers, the Anthropic block shape, the MCP wrapper entry point, and a
near-match name that must still compress
- a LangGraph regression for incomplete tool-call metadata that
continues to a later qualified call

## Testing

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

### Test Output

`pytest tests/test_smart_crusher.py tests/integrations/test_langgraph.py
tests/integrations/test_strands
tests/test_transforms/test_smart_crusher_ccr_retrieve_exemption.py -q`

```text
tests\test_smart_crusher.py ............                                     [ 10%]
tests\integrations\test_langgraph.py .....                                   [ 15%]
tests\integrations\test_strands\test_ccr_exclusion.py .....                  [ 19%]
tests\integrations\test_strands\test_hooks.py sssssssss                      [ 27%]
tests\integrations\test_strands\test_hooks_unit.py ssssssssssssssssssssssssssssssssss [ 57%]
tests\integrations\test_strands\test_model.py ssssssssssssssss               [ 71%]
tests\integrations\test_strands\test_model_unit.py sssssssssssssssssssssssssss [ 95%]
tests\test_transforms\test_smart_crusher_ccr_retrieve_exemption.py .....      [100%]

28 passed, 86 skipped in the focused invariant suite
```

## Real Behavior Proof

- Environment: Windows 11, Python 3.12.13, `headroom._core` built
- Exact command / steps: `uv run pytest tests/test_smart_crusher.py
tests/integrations/test_langgraph.py tests/integrations/test_strands
-q`, and the same suite run against the pre-change implementation with
the new tests in place
- Observed result: before the change, five regressions fail.
`SmartCrusher` returns non-byte-identical content for a
`mcp__Headroom__headroom_retrieve` result, the LangGraph compressor
replaces the message content, and the Strands hook returns
`"compressed"` in place of the tool output. After the change all three
preserve the content byte-for-byte, incomplete LangGraph tool-call
metadata is ignored while the later qualified call remains indexed, the
Strands hook records `tool_excluded` and never calls the crusher, and
`HeadroomMCPCompressor.compress` returns the payload unchanged.
`mcp__Headroom__headroom_retrieve_extra` still compresses in all three,
and the Kompress and ContentRouter suites are unchanged.
- Not tested: the optional Strands package, so the additions to
`tests/integrations/test_strands/test_hooks_unit.py` skip locally

## 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
- [x] I did **not** edit `CHANGELOG.md` — it is generated by
release-please from my Conventional Commit PR title (a CI guard enforces
this)

## Additional Notes

One deliberate divergence from the issue: the suggested snippet passes
`DEFAULT_VERBATIM_EXCLUDE_TOOLS` to `is_tool_excluded`, but that
constant holds only `WebSearch`, `WebFetch`, `web_search`, `web_fetch`.
Applied literally it would drop `headroom_retrieve` from the comparison
entirely and delete the #1077 guard these two SmartCrusher sites exist
to enforce. This passes `(CCR_TOOL_NAME,)` so each guard keeps doing the
one thing it documents. If you'd rather these paths also honor the
verbatim-exclude set, the tuple can become `(CCR_TOOL_NAME,
*DEFAULT_VERBATIM_EXCLUDE_TOOLS)` — the CCR name has to stay in it
either way.

Adjacent work: PR #2654 covers `ContentRouter` only.
This commit is contained in:
Rod Boev 2026-08-03 23:17:39 -04:00 committed by GitHub
parent 677e09735a
commit dcb674b5e4
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 405 additions and 6 deletions

View file

@ -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)

View file

@ -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"

View file

@ -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 <<ccr:hash>> 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):

View file

@ -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"

View file

@ -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()

View file

@ -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."""

141
tests/test_smart_crusher.py Normal file
View file

@ -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
)