mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
fix(transforms): stop ContentRouter recompressing headroom_retrieve results (#2654)
## Description `ContentRouter` (the transform actually registered in the default/proxy compression pipeline -- see `transforms/pipeline.py`) recompresses the output of its own `headroom_retrieve` tool. That tool's entire contract is returning already-retrieved, original content verbatim; recompressing it produces a new `<<ccr:hash>>` marker the caller can never redeem -- an unresolvable retrieval loop. `SmartCrusher` already has a guard against this exact failure mode (#1077), but only on its `apply()` entry point. `ContentRouter` calls the lower-level `SmartCrusher.crush()` directly, bypassing that guard entirely, since `crush()` takes a raw content string with no tool identity at all. Closes #1077 (reopens the same failure mode ContentRouter's own call path, which #1077's original fix did not cover). ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) ## Changes Made - `transforms/content_router.py`: adds an unconditional guard to all three of the places `ContentRouter` can hand a `headroom_retrieve` result to compression: the OpenAI-shape `role:"tool"`/legacy `role:"function"` string-content loop, the Anthropic-shape `tool_result` block loop, and a third, distinct shape -- top-level `{"type": "text"}` blocks under a `role:"tool"`/`"function"` message that never go through a `tool_result` wrapper (a real, already-tested wire shape in this codebase; see `test_tool_role_text_blocks_compressed_by_default`). All three use `is_tool_excluded()` (not a bare comparison) because MCP-served tools appear here under their qualified form, e.g. `mcp__headroom__headroom_retrieve`. Legacy `role:"function"` messages carry no call id in that shape, so the tool name is read directly off the message's `name` field instead of through the id-keyed `tool_name_map`. - Hoisted the per-iteration `is_tool_excluded(..., ("headroom_retrieve",))` calls into a single precomputed `ccr_retrieve_tool_ids` set, computed once alongside the existing `excluded_tool_ids` set, rather than recomputing aliases on every message/block. - `config.py`: adds `"headroom_retrieve"` to `DEFAULT_EXCLUDE_TOOLS` and `DEFAULT_VERBATIM_EXCLUDE_TOOLS` -- this also covers a third path (cross-turn message dedup, `_cross_turn_dedup_messages`) that consults the same frozensets and has no dedicated guard of its own. Also hardens `_tool_name_aliases()` against a non-string tool name (pre-existing fragility, not introduced by this PR, but shares the same call path) by returning no aliases instead of crashing on `.lower()`. - Documentation: updated `ContentRouterConfig.exclude_tools`'s field comment (was stale -- didn't mention this override is unconditional even when a caller explicitly empties `exclude_tools`), and added a comment on `DEFAULT_VERBATIM_EXCLUDE_TOOLS` noting all three real consumers. - Kept `"headroom_retrieve"` as a literal string (matching every other entry in those frozensets) rather than importing the existing `CCR_TOOL_NAME` constant from `ccr.tool_injection` into `content_router.py` -- that module is imported eagerly by `pipeline.py` (unlike `smart_crusher.py`, which imports the same constant lazily), so pulling in `headroom.ccr` there would add a new eager-import edge to a hot module for a one-line DRY win. Happy to change this if a maintainer prefers the constant. **Known, accepted tradeoff:** `is_tool_excluded()`'s alias matching strips any `mcp__<server>__` prefix before comparing, so a third-party MCP server exposing a tool literally named `headroom_retrieve` would also match. Narrowing this to headroom's own server specifically would need a bespoke check inconsistent with how every other excluded-tool entry is matched in this codebase; given how specific the name is, the collision risk is accepted rather than special-cased. ## Testing - [x] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check .`) - [x] Type checking passes (`mypy headroom`) - [x] New tests added for new functionality - [x] Manual testing performed ### Test Output ```text $ uv run pytest tests/test_transforms/ tests/test_transforms_content_router.py -q 1 failed, 420 passed, 62 skipped in 12.50s FAILED tests/test_transforms/test_kompress_compressor.py::...test_onnx_session_options_read_thread_caps (pre-existing, unrelated to this diff -- confirmed via `git stash` that it fails identically against unmodified upstream/main; an ONNX thread-cap assertion, not a compression-routing test) $ uv run ruff check headroom/config.py headroom/transforms/content_router.py \ tests/test_transforms/test_content_router_ccr_retrieve_exemption.py \ tests/test_transforms_content_router.py tests/test_transforms/test_content_router.py All checks passed! $ uv run ruff format --check <same files> 5 files already formatted $ uv run mypy headroom/config.py headroom/transforms/content_router.py Success: no issues found in 2 source files ``` - `tests/test_transforms/test_content_router_ccr_retrieve_exemption.py`: 10 tests -- MCP-qualified name (Anthropic + OpenAI shape), bare name, unconditional-even-with- `exclude_tools=frozenset()`, negative control (normal tools still compressed, asserted via the absence of the `router:excluded:ccr_retrieve` marker), the top-level-text-block shape, legacy `role:"function"`, litellm list-form content nested in a `tool_result` block, mixed retrieve+normal blocks in one turn, and a content well below the compression floor (proving the guard is size-independent). - `tests/test_transforms/test_content_router.py`: `test_anthropic_mcp_bare_tool_alias_exclude_tools` (#1822) updated to assert the new, stronger byte-verbatim guarantee for `headroom_retrieve` specifically; `test_anthropic_mcp_bare_tool_alias_exclude_tools_generic` added to keep the original #1822 general-mechanism coverage (bare-alias matching for an arbitrary, non-exempt tool). - `tests/test_transforms_content_router.py`: updated 10 pre-existing `_process_content_blocks()` unit tests for the new `ccr_retrieve_tool_ids` parameter (all pass empty sets -- none of those tests involve `headroom_retrieve`). - Verified the local installed package copy (a separate, drifted internal version) with a standalone repro script exercising the two new shapes directly against `ContentRouter.apply()` -- both correctly report `router:excluded:ccr_retrieve`. ## Real Behavior Proof - Environment: Windows 11, Python 3.13.7, `uv sync --extra dev` on this branch. - Exact command / steps: standalone repro building an assistant `tool_use` for `mcp__headroom__headroom_retrieve` paired with a large-JSON `tool_result`, through `ContentRouter().apply()`; repeated for the top-level-text-block and legacy-`function`-role shapes. - Observed result: unpatched (Anthropic `tool_result` shape, `git stash` to `upstream/main`), the retrieve output was rewritten 3680 -> 1861 bytes (mangled into a compact tabular form); patched (this branch), it is forwarded 3680 -> 3680 bytes byte-identical, no `<<ccr:` marker present. The two additional shapes fixed in this PR's second commit -- top-level text block under `role:"tool"`, and legacy OpenAI `role:"function"` -- both report `excluded=True` (protected) against this branch, where they reported `excluded=False` (recompressed) before the second commit. - Not tested: the actual `headroom mcp serve` + `headroom wrap` proxy end-to-end over a live Anthropic API call (would need API credentials); the OpenAI-chat-completions `CompressionUnit` path (out of scope, see #1176 below); the opt-in `ToolResultInterceptorTransform` path. ## Relationship to other issues/PRs - Issue #1077 (closed) is this exact bug; PR #1323 fixed it only for `SmartCrusher.apply()`'s own call path (the "legacy" pipeline path, per `smart_crusher.py`'s own comment), not `ContentRouter`, which is what the default/proxy pipeline actually uses. - Open PR #1176 addresses an adjacent, non-overlapping gap: the `CompressionUnit`-based OpenAI chat-completions path (`router.compress()` calls in `transforms/compression_units.py`/`compression_batches.py`), which has no tool-identity context at all and needs its own capture/restore mechanism. This PR does not touch that path. - Filed #2656 as a follow-up: code review on this PR found the same bug class still reachable through `SmartCrusher.apply()`'s own bare-name guard (not alias-aware, so it misses the MCP-qualified form) and through two unguarded direct `.crush()` calls in the LangGraph and Strands integrations. Both are pre-existing, narrower/separate call paths from `ContentRouter`'s primary proxy pipeline, so tracking them separately keeps this PR reviewable as one logical change. - Also not covered by this PR (flagging rather than silently omitting): `proxy/system_compaction.py`'s `router.compress(text, context="")` call, and the opt-in `ToolResultInterceptorTransform` (`HEADROOM_INTERCEPT_ENABLED=1`) -- neither was checked for CCR-awareness. ## 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) ## Screenshots (if applicable) N/A -- this is a backend compression-routing fix with no UI surface. ## Additional Notes This PR is two commits: the first commit added the initial two-loop guard; a second commit followed after code review found the guard was incomplete for two additional wire shapes (top-level text blocks, legacy `role:"function"`) and added the missing test coverage plus a few cleanup items (deduplicated guard logic, comment accuracy, a pre-existing non-string-tool-name fragility). See `Changes Made` above for the full list. Filed #2656 for the remaining out-of-scope gaps found during that same review. --------- Co-authored-by: Michael Tarleton <mtarleton@istation.com>
This commit is contained in:
michaeltarleton
GitHub
parent
13a310a00d
commit
677e09735a
5 changed files with 519 additions and 3 deletions
|
|
@ -213,6 +213,9 @@ class AnchorConfig:
|
|||
# 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.
|
||||
# headroom_retrieve: its entire contract is returning already-retrieved, original
|
||||
# CCR content verbatim. Recompressing it writes a new <<ccr:hash>> marker the
|
||||
# agent can never redeem (#1077).
|
||||
DEFAULT_EXCLUDE_TOOLS: frozenset[str] = frozenset(
|
||||
{
|
||||
"Read",
|
||||
|
|
@ -222,6 +225,7 @@ DEFAULT_EXCLUDE_TOOLS: frozenset[str] = frozenset(
|
|||
"Edit",
|
||||
"WebSearch",
|
||||
"WebFetch",
|
||||
"headroom_retrieve",
|
||||
# Lowercase variants for case-insensitive matching
|
||||
"read",
|
||||
"glob",
|
||||
|
|
@ -235,18 +239,30 @@ DEFAULT_EXCLUDE_TOOLS: frozenset[str] = frozenset(
|
|||
|
||||
# These excluded web-tool results must remain byte-faithful. Even the
|
||||
# excluded-tool lossless fold rewrites formatted JSON.
|
||||
# Three independent consumers key off this frozenset, all in
|
||||
# transforms/content_router.py: ContentRouter's two per-block CCR-retrieve
|
||||
# guards, and _cross_turn_dedup_messages's verbatim_tool_ids -- the latter has
|
||||
# no dedicated guard of its own, so removing headroom_retrieve from here would
|
||||
# silently reopen the retrieval loop for that path with cross-turn dedup on.
|
||||
DEFAULT_VERBATIM_EXCLUDE_TOOLS: frozenset[str] = frozenset(
|
||||
{
|
||||
"WebSearch",
|
||||
"WebFetch",
|
||||
"web_search",
|
||||
"web_fetch",
|
||||
"headroom_retrieve",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _tool_name_aliases(name: str) -> tuple[str, ...]:
|
||||
"""Return equivalent spellings for tool exclusion matching."""
|
||||
if not isinstance(name, str):
|
||||
# Pre-existing fragility (not introduced here): a malformed message can
|
||||
# put a non-string value in the tool-name map (see _build_tool_name_map's
|
||||
# truthy-only `if tc_id and name:` filter). Fail safe -- no aliases means
|
||||
# is_tool_excluded() returns False -- rather than crashing the pipeline.
|
||||
return ()
|
||||
aliases = [name]
|
||||
lname = name.lower()
|
||||
|
||||
|
|
|
|||
|
|
@ -1639,7 +1639,11 @@ class ContentRouterConfig:
|
|||
compress_tagged_content: bool = False
|
||||
|
||||
# Tools to exclude from compression (output passed through unmodified)
|
||||
# Set to None to use DEFAULT_EXCLUDE_TOOLS, or provide custom set
|
||||
# Set to None to use DEFAULT_EXCLUDE_TOOLS, or provide custom set.
|
||||
# NOTE: headroom_retrieve is excluded unconditionally regardless of this
|
||||
# setting, even if this is explicitly set to an empty set -- recompressing
|
||||
# its output would write a new <<ccr:hash>> marker the agent can never
|
||||
# redeem (see the ccr_retrieve_tool_ids guards in apply()).
|
||||
exclude_tools: set[str] | None = None
|
||||
|
||||
# Excluded tools are protected only from *lossy* compression. Their output
|
||||
|
|
@ -4661,6 +4665,29 @@ class ContentRouter(Transform):
|
|||
if is_tool_excluded(name, exclude_tools)
|
||||
}
|
||||
|
||||
# CCR-retrieve tool IDs, precomputed once (mirrors excluded_tool_ids
|
||||
# above, rather than calling is_tool_excluded() per message/block). A
|
||||
# headroom_retrieve result IS already-retrieved, original CCR content --
|
||||
# recompressing it writes a new <<ccr:hash>> marker the agent can never
|
||||
# redeem (unresolvable retrieval loop). is_tool_excluded() (not a bare
|
||||
# comparison) because MCP-served tools appear here under their qualified
|
||||
# name, e.g. mcp__headroom__headroom_retrieve. Consulted unconditionally,
|
||||
# ahead of the age-decay/read-protection-window logic below: a decayed
|
||||
# CCR marker is exactly as unredeemable as a fresh one, so this must
|
||||
# never fall through to compression the way excluded_tool_ids does.
|
||||
# Known, accepted tradeoff: is_tool_excluded()'s alias matching strips
|
||||
# ANY mcp__<server>__ prefix before comparing, so a third-party server
|
||||
# exposing a tool literally named headroom_retrieve would also match
|
||||
# here. Narrowing this to headroom's own server specifically would need
|
||||
# a bespoke check inconsistent with how every other excluded-tool entry
|
||||
# is matched; given the name is this specific, the collision risk is
|
||||
# accepted rather than special-cased.
|
||||
ccr_retrieve_tool_ids = {
|
||||
tool_id
|
||||
for tool_id, name in tool_name_map.items()
|
||||
if is_tool_excluded(name, ("headroom_retrieve",))
|
||||
}
|
||||
|
||||
# Read protection (HEADROOM_PROTECT_READS=1): for bash-family agents the
|
||||
# exclude-by-tool-NAME set above never catches file reads (they are `bash`
|
||||
# tool calls whose COMMAND is a cat/sed/head/...). Mark those tool_use_ids so
|
||||
|
|
@ -4800,6 +4827,7 @@ class ContentRouter(Transform):
|
|||
# Routing reason counters for summary logging
|
||||
route_counts: dict[str, int] = {
|
||||
"excluded_tool": 0,
|
||||
"ccr_retrieve": 0,
|
||||
"user_msg": 0,
|
||||
"small": 0,
|
||||
"recent_code": 0,
|
||||
|
|
@ -4913,6 +4941,7 @@ class ContentRouter(Transform):
|
|||
context,
|
||||
transforms_applied,
|
||||
excluded_tool_ids,
|
||||
ccr_retrieve_tool_ids,
|
||||
tool_name_map=tool_name_map,
|
||||
route_counts=route_counts,
|
||||
compressed_details=compressed_details,
|
||||
|
|
@ -4935,11 +4964,29 @@ class ContentRouter(Transform):
|
|||
route_counts["non_string"] += 1
|
||||
continue
|
||||
|
||||
# A headroom_retrieve result IS already-retrieved, original CCR content --
|
||||
# recompressing it writes a new <<ccr:hash>> marker the agent can never
|
||||
# redeem (unresolvable retrieval loop). Covers role:"tool" (id-keyed via
|
||||
# tool_call_id -> ccr_retrieve_tool_ids, precomputed above) and legacy
|
||||
# role:"function" (that shape carries no call id -- the tool name is on
|
||||
# the message itself via "name", per OpenAI's pre-parallel-tool-calls API).
|
||||
tool_call_id = message.get("tool_call_id", "") if role in ("tool", "function") else ""
|
||||
if role in ("tool", "function") and (
|
||||
tool_call_id in ccr_retrieve_tool_ids
|
||||
or (
|
||||
role == "function"
|
||||
and is_tool_excluded(message.get("name", ""), ("headroom_retrieve",))
|
||||
)
|
||||
):
|
||||
result_slots[i] = message
|
||||
transforms_applied.append("router:excluded:ccr_retrieve")
|
||||
route_counts["ccr_retrieve"] += 1
|
||||
continue
|
||||
|
||||
# Skip OpenAI-style tool messages for excluded tools
|
||||
# BUT: allow compression of old excluded-tool outputs beyond the
|
||||
# adaptive protection window (age-based decay).
|
||||
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):
|
||||
|
|
@ -5729,6 +5776,7 @@ class ContentRouter(Transform):
|
|||
context: str,
|
||||
transforms_applied: list[str],
|
||||
excluded_tool_ids: set[str],
|
||||
ccr_retrieve_tool_ids: set[str],
|
||||
tool_name_map: dict[str, str] | None = None,
|
||||
route_counts: dict[str, int] | None = None,
|
||||
compressed_details: list[str] | None = None,
|
||||
|
|
@ -5768,6 +5816,9 @@ class ContentRouter(Transform):
|
|||
context: Context for compression.
|
||||
transforms_applied: List to append transform names to.
|
||||
excluded_tool_ids: Tool IDs to skip compression for.
|
||||
ccr_retrieve_tool_ids: Tool IDs whose output is a headroom_retrieve result --
|
||||
always passed through verbatim (see module-level comment at the
|
||||
precompute site for why this can never fall through to compression).
|
||||
tool_name_map: Mapping from tool_call_id to tool_name for profile lookup.
|
||||
route_counts: Optional routing reason counters to update.
|
||||
compressed_details: Optional list to append compression details to.
|
||||
|
|
@ -5860,6 +5911,16 @@ class ContentRouter(Transform):
|
|||
route_counts.setdefault("read_protected", 0)
|
||||
route_counts["read_protected"] += 1
|
||||
continue
|
||||
# Mirrors the OpenAI-shape guard above (issue #1077): a headroom_retrieve
|
||||
# result IS already-retrieved, original CCR content and must never be
|
||||
# recompressed. Precomputed set (mirrors excluded_tool_ids immediately
|
||||
# below) rather than a per-iteration is_tool_excluded() call.
|
||||
if tool_use_id in ccr_retrieve_tool_ids:
|
||||
new_blocks.append(block)
|
||||
transforms_applied.append("router:excluded:ccr_retrieve")
|
||||
if route_counts is not None:
|
||||
route_counts["ccr_retrieve"] = route_counts.get("ccr_retrieve", 0) + 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):
|
||||
|
|
@ -6026,6 +6087,24 @@ class ContentRouter(Transform):
|
|||
# skipped; assistant default-skipped, opt-in via
|
||||
# `compress_assistant_text_blocks`).
|
||||
elif block_type == "text" and not protect_text_blocks:
|
||||
# Same CCR-retrieve exemption as the tool_result branch above, for the
|
||||
# top-level-text-block wire shape: a role:"tool"/"function" harness that
|
||||
# normalizes content to a block list without a tool_result wrapper (see
|
||||
# test_tool_role_text_blocks_compressed_by_default for why this shape is
|
||||
# real). role:"tool" resolves via the message's own tool_call_id/
|
||||
# tool_use_id through ccr_retrieve_tool_ids; legacy role:"function" has no
|
||||
# call id in that shape, so the tool name is read off the message directly.
|
||||
if role in ("tool", "function"):
|
||||
_msg_tool_id = message.get("tool_call_id") or message.get("tool_use_id") or ""
|
||||
if _msg_tool_id in ccr_retrieve_tool_ids or (
|
||||
role == "function"
|
||||
and is_tool_excluded(message.get("name", ""), ("headroom_retrieve",))
|
||||
):
|
||||
new_blocks.append(block)
|
||||
transforms_applied.append("router:excluded:ccr_retrieve")
|
||||
if route_counts is not None:
|
||||
route_counts["ccr_retrieve"] = route_counts.get("ccr_retrieve", 0) + 1
|
||||
continue
|
||||
text_content = block.get("text", "")
|
||||
if isinstance(text_content, str) and (
|
||||
len(text_content) > min_chars or self._has_lossless_fold(text_content)
|
||||
|
|
|
|||
|
|
@ -893,7 +893,19 @@ class TestExcludeTools:
|
|||
assert "router:excluded:lossless_json" in result.transforms_applied
|
||||
|
||||
def test_anthropic_mcp_bare_tool_alias_exclude_tools(self, tokenizer):
|
||||
"""Bare tool exclusions match custom-agent MCP wrappers (#1822)."""
|
||||
"""Bare tool exclusions match custom-agent MCP wrappers (#1822).
|
||||
|
||||
Any MCP wrapper's bare tool name can be excluded via config — this test
|
||||
uses a fictitious "HeadroomZai" server name to prove the alias match is
|
||||
server-name-agnostic. ``headroom_retrieve`` specifically is now also an
|
||||
unconditional, config-independent exclusion (see the fix for the
|
||||
ContentRouter self-recompression bug: SmartCrusher.apply() already
|
||||
guarded #1077 on its own call path, but ContentRouter called
|
||||
SmartCrusher.crush() directly, bypassing it). That guard fires before
|
||||
the config-driven `excluded_tool_ids` check below, giving byte-identical
|
||||
passthrough rather than the lossless-JSON fold a narrower custom
|
||||
`exclude_tools` used to produce for this specific tool name.
|
||||
"""
|
||||
config = ContentRouterConfig(
|
||||
min_section_tokens=10,
|
||||
exclude_tools={"headroom_retrieve"},
|
||||
|
|
@ -926,6 +938,50 @@ class TestExcludeTools:
|
|||
|
||||
result = router.apply(messages, tokenizer)
|
||||
|
||||
tool_result_block = result.messages[1]["content"][0]
|
||||
# Byte-identical, not just JSON-semantically-equal: the unconditional
|
||||
# ccr_retrieve guard passes the original block through untouched.
|
||||
assert tool_result_block["content"] == messages[1]["content"][0]["content"]
|
||||
assert "router:excluded:ccr_retrieve" in result.transforms_applied
|
||||
|
||||
def test_anthropic_mcp_bare_tool_alias_exclude_tools_generic(self, tokenizer):
|
||||
"""General #1822 coverage: bare-name alias matching through the
|
||||
config-driven ``excluded_tool_ids``/``DEFAULT_VERBATIM_EXCLUDE_TOOLS``
|
||||
path for an arbitrary tool that is NOT ``headroom_retrieve`` (which now
|
||||
has its own unconditional guard that would otherwise mask this path —
|
||||
see ``test_anthropic_mcp_bare_tool_alias_exclude_tools`` above)."""
|
||||
config = ContentRouterConfig(
|
||||
min_section_tokens=10,
|
||||
exclude_tools={"measure"},
|
||||
)
|
||||
router = ContentRouter(config)
|
||||
|
||||
messages = [
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{
|
||||
"type": "tool_use",
|
||||
"id": "toolu_measure_1",
|
||||
"name": "mcp_build123d_measure",
|
||||
"input": {"key": "abc123"},
|
||||
}
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": "toolu_measure_1",
|
||||
"content": generate_json_data(50),
|
||||
}
|
||||
],
|
||||
},
|
||||
]
|
||||
|
||||
result = router.apply(messages, tokenizer)
|
||||
|
||||
tool_result_block = result.messages[1]["content"][0]
|
||||
assert json.loads(tool_result_block["content"]) == json.loads(
|
||||
messages[1]["content"][0]["content"]
|
||||
|
|
|
|||
|
|
@ -0,0 +1,355 @@
|
|||
"""Regression tests: ContentRouter must not re-compress headroom_retrieve results.
|
||||
|
||||
Companion to ``test_smart_crusher_ccr_retrieve_exemption.py`` (issue #1077).
|
||||
SmartCrusher.apply() already guards this exact failure mode, but ContentRouter
|
||||
— the transform actually registered in the default/proxy pipeline (see
|
||||
``transforms/pipeline.py``) — calls ``SmartCrusher.crush()`` directly, bypassing
|
||||
that guard entirely. Without this fix, a ``headroom_retrieve`` tool result sent
|
||||
back to the model on the next turn gets swept up by ContentRouter's own
|
||||
compression routing like any other large tool output, producing a fresh
|
||||
``<<ccr:hash>>`` marker the agent can never redeem (an unresolvable retrieval
|
||||
loop — the original reported bug).
|
||||
|
||||
Tests use the fully-qualified MCP tool name (``mcp__headroom__headroom_retrieve``)
|
||||
that Claude Code actually sends when connected to `headroom mcp serve` — not the
|
||||
bare name — since routing the guard through ``is_tool_excluded()`` (alias-aware)
|
||||
rather than a bare string comparison is the whole point of the fix.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
|
||||
from headroom.transforms.content_router import ContentRouter, ContentRouterConfig
|
||||
|
||||
|
||||
def _get_tokenizer():
|
||||
from headroom.providers import OpenAIProvider
|
||||
from headroom.tokenizer import Tokenizer
|
||||
|
||||
provider = OpenAIProvider()
|
||||
token_counter = provider.get_token_counter("gpt-4o")
|
||||
return Tokenizer(token_counter, "gpt-4o")
|
||||
|
||||
|
||||
def _big_json() -> str:
|
||||
"""A JSON array large enough to clear the compression threshold."""
|
||||
return json.dumps([{"id": i, "value": "x" * 20, "active": i % 2 == 0} for i in range(60)])
|
||||
|
||||
|
||||
def _anthropic_messages(tool_name: str, content: str) -> list[dict]:
|
||||
return [
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{
|
||||
"type": "tool_use",
|
||||
"id": "toolu_ccr_1",
|
||||
"name": tool_name,
|
||||
"input": {"hash": "abc123def456"},
|
||||
}
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": "toolu_ccr_1",
|
||||
"content": content,
|
||||
}
|
||||
],
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
def _openai_messages(tool_name: str, content: str) -> list[dict]:
|
||||
return [
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "call_ccr_1",
|
||||
"type": "function",
|
||||
"function": {"name": tool_name, "arguments": '{"hash":"abc123def456"}'},
|
||||
}
|
||||
],
|
||||
},
|
||||
{"role": "tool", "tool_call_id": "call_ccr_1", "content": content},
|
||||
]
|
||||
|
||||
|
||||
def _tool_role_top_level_text_messages(tool_name: str, content: str) -> list[dict]:
|
||||
"""Some harnesses normalize a tool-role message's content to a top-level
|
||||
list of {"type": "text"} blocks, without the {"type": "tool_result", ...}
|
||||
wrapper the primary Anthropic-shape tests exercise. This is a distinct
|
||||
code path in _process_content_blocks (the generic `text` block branch)."""
|
||||
return [
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{
|
||||
"type": "tool_use",
|
||||
"id": "toolu_ccr_toplevel",
|
||||
"name": tool_name,
|
||||
"input": {"hash": "abc123def456"},
|
||||
}
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "tool",
|
||||
"tool_call_id": "toolu_ccr_toplevel",
|
||||
"content": [{"type": "text", "text": content}],
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
def _legacy_function_messages(tool_name: str, content: str) -> list[dict]:
|
||||
"""Legacy (pre-parallel-tool-calls) OpenAI function-calling: the request
|
||||
uses a singular top-level `function_call` field (not `tool_calls`), and
|
||||
the response is a role:"function" message carrying the name directly --
|
||||
this shape has no call id at all, unlike role:"tool"."""
|
||||
return [
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"function_call": {"name": tool_name, "arguments": '{"hash":"abc123def456"}'},
|
||||
},
|
||||
{"role": "function", "name": tool_name, "content": content},
|
||||
]
|
||||
|
||||
|
||||
def _anthropic_list_form_messages(tool_name: str, content: str) -> list[dict]:
|
||||
"""litellm/OpenAI-compatible gateways sometimes wrap tool_result content
|
||||
as a list of {"type": "text"} blocks even in Anthropic-shaped messages
|
||||
(see content_router.py's _tr_list_form_early flatten, fix-7)."""
|
||||
return [
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{
|
||||
"type": "tool_use",
|
||||
"id": "toolu_ccr_listform",
|
||||
"name": tool_name,
|
||||
"input": {"hash": "abc123def456"},
|
||||
}
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": "toolu_ccr_listform",
|
||||
"content": [{"type": "text", "text": content}],
|
||||
}
|
||||
],
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
class TestContentRouterCcrRetrieveExemption:
|
||||
def test_anthropic_qualified_mcp_name_not_recompressed(self):
|
||||
"""The MCP-qualified name Claude Code actually sends must be exempted."""
|
||||
content = _big_json()
|
||||
router = ContentRouter(ContentRouterConfig(min_section_tokens=10))
|
||||
tokenizer = _get_tokenizer()
|
||||
|
||||
messages = _anthropic_messages("mcp__headroom__headroom_retrieve", content)
|
||||
result = router.apply(messages, tokenizer)
|
||||
|
||||
tool_result_block = result.messages[1]["content"][0]
|
||||
assert tool_result_block["content"] == content, (
|
||||
"headroom_retrieve result was recompressed via ContentRouter "
|
||||
"(unresolvable retrieval loop)"
|
||||
)
|
||||
assert "<<ccr:" not in tool_result_block["content"]
|
||||
assert "router:excluded:ccr_retrieve" in result.transforms_applied
|
||||
|
||||
def test_openai_qualified_mcp_name_not_recompressed(self):
|
||||
"""Same guarantee for the OpenAI/litellm tool-message shape."""
|
||||
content = _big_json()
|
||||
router = ContentRouter(ContentRouterConfig(min_section_tokens=10))
|
||||
tokenizer = _get_tokenizer()
|
||||
|
||||
messages = _openai_messages("mcp__headroom__headroom_retrieve", content)
|
||||
result = router.apply(messages, tokenizer)
|
||||
|
||||
tool_msg = next(m for m in result.messages if m.get("tool_call_id") == "call_ccr_1")
|
||||
assert tool_msg["content"] == content
|
||||
assert "<<ccr:" not in tool_msg["content"]
|
||||
assert "router:excluded:ccr_retrieve" in result.transforms_applied
|
||||
|
||||
def test_bare_tool_name_also_protected(self):
|
||||
"""The proxy's own internally-injected retrieval tool uses the bare
|
||||
name (not MCP-qualified) — must be protected too, matching the
|
||||
existing SmartCrusher.apply() coverage for this call shape."""
|
||||
content = _big_json()
|
||||
router = ContentRouter(ContentRouterConfig(min_section_tokens=10))
|
||||
tokenizer = _get_tokenizer()
|
||||
|
||||
messages = _anthropic_messages("headroom_retrieve", content)
|
||||
result = router.apply(messages, tokenizer)
|
||||
|
||||
tool_result_block = result.messages[1]["content"][0]
|
||||
assert tool_result_block["content"] == content
|
||||
assert "router:excluded:ccr_retrieve" in result.transforms_applied
|
||||
|
||||
def test_unconditional_even_with_empty_exclude_tools(self):
|
||||
"""The guard is config-independent: even a caller that explicitly
|
||||
empties exclude_tools (disabling every default exclusion) must not
|
||||
be able to recompress a headroom_retrieve result. This is the
|
||||
defense-in-depth half of the fix — config.py's DEFAULT_EXCLUDE_TOOLS/
|
||||
DEFAULT_VERBATIM_EXCLUDE_TOOLS additions alone would not survive this
|
||||
override, since ContentRouter replaces (not merges) exclude_tools
|
||||
when the caller sets it explicitly."""
|
||||
content = _big_json()
|
||||
router = ContentRouter(
|
||||
ContentRouterConfig(min_section_tokens=10, exclude_tools=frozenset())
|
||||
)
|
||||
tokenizer = _get_tokenizer()
|
||||
|
||||
messages = _anthropic_messages("mcp__headroom__headroom_retrieve", content)
|
||||
result = router.apply(messages, tokenizer)
|
||||
|
||||
tool_result_block = result.messages[1]["content"][0]
|
||||
assert tool_result_block["content"] == content, (
|
||||
"headroom_retrieve must stay protected even when exclude_tools is "
|
||||
"explicitly emptied — the guard must not depend on config"
|
||||
)
|
||||
assert "router:excluded:ccr_retrieve" in result.transforms_applied
|
||||
|
||||
def test_non_retrieve_tool_still_compressed(self):
|
||||
"""Exemption is narrow: an ordinary tool with the same large JSON
|
||||
content is still compressed — this isn't a blanket JSON passthrough."""
|
||||
content = _big_json()
|
||||
router = ContentRouter(ContentRouterConfig(min_section_tokens=10))
|
||||
tokenizer = _get_tokenizer()
|
||||
|
||||
messages = _anthropic_messages("Bash", content)
|
||||
result = router.apply(messages, tokenizer)
|
||||
|
||||
assert "router:excluded:ccr_retrieve" not in result.transforms_applied
|
||||
tool_result_block = result.messages[1]["content"][0]
|
||||
assert tool_result_block["content"] != content or result.tokens_after < result.tokens_before
|
||||
|
||||
def test_top_level_text_block_under_tool_role_not_recompressed(self):
|
||||
"""A harness that normalizes tool-role content to a top-level text
|
||||
block (not wrapped in tool_result) must be protected too -- this is
|
||||
a distinct code path (_process_content_blocks' generic `text` branch)
|
||||
from the tool_result branch the other Anthropic-shape tests exercise."""
|
||||
content = _big_json()
|
||||
router = ContentRouter(ContentRouterConfig(min_section_tokens=10))
|
||||
tokenizer = _get_tokenizer()
|
||||
|
||||
messages = _tool_role_top_level_text_messages("mcp__headroom__headroom_retrieve", content)
|
||||
result = router.apply(messages, tokenizer)
|
||||
|
||||
tool_msg = next(m for m in result.messages if m.get("role") == "tool")
|
||||
assert tool_msg["content"][0]["text"] == content
|
||||
assert "<<ccr:" not in tool_msg["content"][0]["text"]
|
||||
assert "router:excluded:ccr_retrieve" in result.transforms_applied
|
||||
|
||||
def test_legacy_openai_function_role_not_recompressed(self):
|
||||
"""Legacy OpenAI function-calling (role:"function", no tool_call_id --
|
||||
the tool name is carried directly on the message) must be protected
|
||||
too. This shape predates parallel tool_calls and bypasses
|
||||
_build_tool_name_map's id-keyed lookup entirely."""
|
||||
content = _big_json()
|
||||
router = ContentRouter(ContentRouterConfig(min_section_tokens=10))
|
||||
tokenizer = _get_tokenizer()
|
||||
|
||||
messages = _legacy_function_messages("headroom_retrieve", content)
|
||||
result = router.apply(messages, tokenizer)
|
||||
|
||||
fn_msg = next(m for m in result.messages if m.get("role") == "function")
|
||||
assert fn_msg["content"] == content
|
||||
assert "<<ccr:" not in fn_msg["content"]
|
||||
assert "router:excluded:ccr_retrieve" in result.transforms_applied
|
||||
|
||||
def test_litellm_list_form_tool_result_not_recompressed(self):
|
||||
"""litellm/OpenAI-style list-form content (a list of {"type":"text"}
|
||||
blocks) nested inside an Anthropic tool_result block must be
|
||||
protected too -- distinct from the plain-string tool_result content
|
||||
the other Anthropic-shape tests exercise (see fix-7 flatten)."""
|
||||
content = _big_json()
|
||||
router = ContentRouter(ContentRouterConfig(min_section_tokens=10))
|
||||
tokenizer = _get_tokenizer()
|
||||
|
||||
messages = _anthropic_list_form_messages("mcp__headroom__headroom_retrieve", content)
|
||||
result = router.apply(messages, tokenizer)
|
||||
|
||||
tool_result_block = result.messages[1]["content"][0]
|
||||
assert tool_result_block["content"] == [{"type": "text", "text": content}]
|
||||
assert "router:excluded:ccr_retrieve" in result.transforms_applied
|
||||
|
||||
def test_mixed_retrieve_and_normal_in_one_turn_only_normal_compressed(self):
|
||||
"""Per-block precision: a single turn carrying both a headroom_retrieve
|
||||
tool_result and a normal tool_result must protect only the former --
|
||||
parity with test_smart_crusher_ccr_retrieve_exemption.py's
|
||||
test_mixed_retrieve_and_normal_only_normal_compressed."""
|
||||
ccr_content = _big_json()
|
||||
normal_content = _big_json()
|
||||
router = ContentRouter(ContentRouterConfig(min_section_tokens=10))
|
||||
tokenizer = _get_tokenizer()
|
||||
|
||||
messages = [
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{
|
||||
"type": "tool_use",
|
||||
"id": "toolu_ccr_mixed",
|
||||
"name": "mcp__headroom__headroom_retrieve",
|
||||
"input": {"hash": "abc123def456"},
|
||||
},
|
||||
{
|
||||
"type": "tool_use",
|
||||
"id": "toolu_normal_mixed",
|
||||
"name": "Bash",
|
||||
"input": {"command": "ls"},
|
||||
},
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": "toolu_ccr_mixed",
|
||||
"content": ccr_content,
|
||||
},
|
||||
{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": "toolu_normal_mixed",
|
||||
"content": normal_content,
|
||||
},
|
||||
],
|
||||
},
|
||||
]
|
||||
result = router.apply(messages, tokenizer)
|
||||
|
||||
blocks = result.messages[1]["content"]
|
||||
ccr_block = next(b for b in blocks if b["tool_use_id"] == "toolu_ccr_mixed")
|
||||
normal_block = next(b for b in blocks if b["tool_use_id"] == "toolu_normal_mixed")
|
||||
|
||||
assert ccr_block["content"] == ccr_content
|
||||
assert normal_block["content"] != normal_content
|
||||
assert "router:excluded:ccr_retrieve" in result.transforms_applied
|
||||
|
||||
def test_small_headroom_retrieve_content_still_marked_excluded(self):
|
||||
"""The guard is size-independent: even content well below the
|
||||
compression floor must still get the router:excluded:ccr_retrieve
|
||||
marker, proving the exemption fires unconditionally rather than
|
||||
happening to survive because it was too small to compress anyway."""
|
||||
content = "tiny retrieved value"
|
||||
router = ContentRouter(ContentRouterConfig(min_section_tokens=10))
|
||||
tokenizer = _get_tokenizer()
|
||||
|
||||
messages = _anthropic_messages("mcp__headroom__headroom_retrieve", content)
|
||||
result = router.apply(messages, tokenizer)
|
||||
|
||||
tool_result_block = result.messages[1]["content"][0]
|
||||
assert tool_result_block["content"] == content
|
||||
assert "router:excluded:ccr_retrieve" in result.transforms_applied
|
||||
|
|
@ -932,6 +932,7 @@ def test_text_block_cache_control_protected_with_assistant_optin(
|
|||
"",
|
||||
[],
|
||||
set(),
|
||||
set(),
|
||||
route_counts=counts,
|
||||
compress_assistant_text_blocks=True,
|
||||
)
|
||||
|
|
@ -964,6 +965,7 @@ def test_tool_result_cache_control_protected(monkeypatch: pytest.MonkeyPatch) ->
|
|||
"",
|
||||
[],
|
||||
set(),
|
||||
set(),
|
||||
)
|
||||
# cache_control hard-skip applies to tool_result too
|
||||
assert result["content"][0]["content"] == long_text
|
||||
|
|
@ -981,6 +983,7 @@ def test_assistant_text_blocks_skipped_by_default(
|
|||
"",
|
||||
[],
|
||||
set(),
|
||||
set(),
|
||||
)
|
||||
# Default OFF: assistant text untouched, restoring pre-#431 cache safety
|
||||
assert result["content"][0]["text"] == long_text
|
||||
|
|
@ -998,6 +1001,7 @@ def test_assistant_text_blocks_opt_in_compresses(
|
|||
"",
|
||||
[],
|
||||
set(),
|
||||
set(),
|
||||
compress_assistant_text_blocks=True,
|
||||
)
|
||||
assert "[compressed]" in result["content"][0]["text"]
|
||||
|
|
@ -1015,6 +1019,7 @@ def test_user_text_blocks_never_compressed_even_with_assistant_optin(
|
|||
"",
|
||||
[],
|
||||
set(),
|
||||
set(),
|
||||
compress_assistant_text_blocks=True, # MUST NOT bleed into user
|
||||
)
|
||||
assert result["content"][0]["text"] == long_text
|
||||
|
|
@ -1032,6 +1037,7 @@ def test_system_text_blocks_skipped_when_skip_system_true(
|
|||
"",
|
||||
[],
|
||||
set(),
|
||||
set(),
|
||||
skip_system=True,
|
||||
compress_assistant_text_blocks=True,
|
||||
)
|
||||
|
|
@ -1050,6 +1056,7 @@ def test_tool_role_text_blocks_compressed_by_default(
|
|||
"",
|
||||
[],
|
||||
set(),
|
||||
set(),
|
||||
)
|
||||
# tool role ≈ tool output — compress freely
|
||||
assert "[compressed]" in result["content"][0]["text"]
|
||||
|
|
@ -1067,6 +1074,7 @@ def test_unknown_role_text_blocks_skipped_for_safety(
|
|||
"",
|
||||
[],
|
||||
set(),
|
||||
set(),
|
||||
compress_assistant_text_blocks=True,
|
||||
)
|
||||
# Unknown role: be safe, don't compress
|
||||
|
|
@ -1083,6 +1091,7 @@ def test_min_chars_gates_short_blocks(monkeypatch: pytest.MonkeyPatch) -> None:
|
|||
"",
|
||||
[],
|
||||
set(),
|
||||
set(),
|
||||
min_chars=500,
|
||||
)
|
||||
assert result["content"][0]["text"] == short_text
|
||||
|
|
@ -1098,6 +1107,7 @@ def test_pinning_skips_already_compressed(monkeypatch: pytest.MonkeyPatch) -> No
|
|||
"",
|
||||
[],
|
||||
set(),
|
||||
set(),
|
||||
)
|
||||
# Already-compressed marker keeps proxy idempotent across turns
|
||||
assert result["content"][0]["text"] == pinned
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue