mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
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:
parent
677e09735a
commit
dcb674b5e4
7 changed files with 405 additions and 6 deletions
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
90
tests/integrations/test_langgraph.py
Normal file
90
tests/integrations/test_langgraph.py
Normal 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"
|
||||
91
tests/integrations/test_strands/test_ccr_exclusion.py
Normal file
91
tests/integrations/test_strands/test_ccr_exclusion.py
Normal 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()
|
||||
|
|
@ -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
141
tests/test_smart_crusher.py
Normal 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
|
||||
)
|
||||
Loading…
Add table
Add a link
Reference in a new issue