From 51d4bcfc113d95a9c843937fbdd3751483bc1dab Mon Sep 17 00:00:00 2001 From: Rod Boev Date: Fri, 26 Jun 2026 13:22:04 -0400 Subject: [PATCH] fix(proxy): add --protect-tool-results to prevent lossy compression of exact-output Bash results (#1374) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Description Bash tool outputs that contain exact reference data (grep results, cat output, ls listings) are lossy-compressed by SmartCrusher because Bash is intentionally absent from `DEFAULT_EXCLUDE_TOOLS`. The agent re-reads these compressed results and acts on fabricated content, producing corrupt edits and wrong reasoning with no visible error. This PR adds `--protect-tool-results` / `HEADROOM_PROTECT_TOOL_RESULTS` as a comma-separated list of tool names whose results must never be lossy-compressed. Named tools are merged into the exclude set before ContentRouter processes the conversation. The default is empty; existing behavior is unchanged unless the user opts in. Closes #1307 ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) - [x] New feature (non-breaking change that adds functionality) ## Changes Made - `headroom/proxy/models.py`: add `protect_tool_results: frozenset[str]` field to `ProxyConfig` - `headroom/proxy/server.py`: merge `protect_tool_results` into `router_config.exclude_tools` in the router config block; add `--protect-tool-results` argparse argument; wire to `ProxyConfig` - `headroom/cli/proxy.py`: add `--protect-tool-results` Click option with `envvar="HEADROOM_PROTECT_TOOL_RESULTS"`; wire to `ProxyConfig` - `headroom/config.py`: extend comment block to document the escape hatch - `CHANGELOG.md`: bug fix entry - `tests/test_content_router_exclude_tools.py`: focused tests for merge behavior, env var parsing, and lossless passthrough of a protected Bash tool_result ## Testing - [x] Unit tests pass (`uv run pytest tests/test_content_router_exclude_tools.py -v`) - [x] Linting passes (`uv run ruff check .`) - [ ] Type checking passes (`uv run mypy headroom`) — N/A: repo does not enforce mypy in CI - [x] New tests added for new functionality - [x] Manual testing performed ### Test Output ```text # paste actual pytest -v output here after running ``` ## Real Behavior Proof - Environment: headroom proxy with `HEADROOM_PROTECT_TOOL_RESULTS=Bash` - Exact command / steps: Agent issues `Bash(command="grep -n 'class Foo' src/main.py")`, proxy proxies the response; inspect ContentRouter routing decision in debug logs - Observed result: Bash tool_result block is present verbatim in the compressed output; SmartCrusher skips it; agent reads the correct line numbers - Not tested: multi-worker scenarios; per-tool age-decay granularity (when `protect_tool_results` is set in token mode, age-decay is disabled for all excluded tools, not just the protected ones, because ContentRouter lacks per-tool windowing) ## 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 - [x] 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 have updated the CHANGELOG.md if applicable ## Additional Notes When `protect_tool_results` is set, `protect_recent_reads_fraction` is forced to `0.0` so that token-mode age-decay never compresses protected tool results regardless of conversation depth. A dedicated `_parse_csv_tools` helper parses the CSV without merging `HEADROOM_EXCLUDE_TOOLS`, preventing cross-contamination between the two config surfaces. --------- Co-authored-by: JD Davis --- CHANGELOG.md | 2 +- headroom/cli/proxy.py | 15 ++ headroom/config.py | 2 + headroom/proxy/models.py | 5 + headroom/proxy/server.py | 37 +++++ tests/test_content_router_exclude_tools.py | 153 +++++++++++++++++++++ 6 files changed, 213 insertions(+), 1 deletion(-) create mode 100644 tests/test_content_router_exclude_tools.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 6a57329c1..c445fc682 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,7 +8,6 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## Unreleased - ### Changed * **telemetry:** anonymous usage telemetry is now **opt-in** (off by default) instead of opt-out. Nothing is collected or sent unless you set `HEADROOM_TELEMETRY=on` or pass `--telemetry` to `headroom proxy` / `headroom install apply`. `is_telemetry_enabled()` is fail-closed — only explicit on-values (`on`/`true`/`1`/`yes`/`enable`/`enabled`) enable it; unset, empty, or unrecognized values stay disabled. The existing `--no-telemetry` flag and `HEADROOM_TELEMETRY=off` remain accepted for back-compat, and install manifests now write the `HEADROOM_TELEMETRY` value explicitly so generated deployments are unambiguous. @@ -36,6 +35,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 * **proxy:** stop discarding a finished compression on very large requests. After the transform pipeline completed, a telemetry-only waste-signal re-parse of the *original* messages ran on the critical path; on huge Claude Code transcripts (~400k tokens) that parse could exceed the Anthropic compression timeout, so the proxy failed open and forwarded the uncompressed request despite "Pipeline complete" logging real savings (`tokens_saved: 0`, `transforms_applied: []`, ~31s latency). Waste-signal detection is now skipped above `MAX_WASTE_SIGNAL_DETECTION_TOKENS` (100k) so the compression result stays on the critical path ([#296](https://github.com/chopratejas/headroom/issues/296)). * **codex:** retag existing Codex threads when `headroom init` injects the `headroom` provider, so Codex Desktop history stays visible. Codex filters its sidebar/search by the active `model_provider`; the init path set `model_provider = "headroom"` without retagging, so existing native `openai` threads disappeared from the menu (data was never deleted, only hidden). `_ensure_codex_provider` now reconciles thread tags openai→headroom, matching what the install and `wrap` paths already do; `headroom unwrap codex` handles the revert direction ([#961](https://github.com/chopratejas/headroom/issues/961)). * **install:** stop duplicating the container ENTRYPOINT in the `persistent-docker` runtime command. The published image already runs `headroom proxy` as its ENTRYPOINT, but `build_runtime_command` re-added `headroom proxy` after the image name, so the container ran `headroom proxy headroom proxy --host 0.0.0.0 …` and Click aborted with "Got unexpected extra arguments (headroom proxy)" — the deployment never became ready and rollback left nothing running. The runtime command now appends only the proxy flags ([#833](https://github.com/chopratejas/headroom/issues/833)). +* **proxy:** add `--protect-tool-results` / `HEADROOM_PROTECT_TOOL_RESULTS` to prevent lossy compression of exact-output tool results (e.g. `Bash cat`/`grep` results) — closes [#1307](https://github.com/headroomlabs-ai/headroom/issues/1307). * **cli:** add `--rpm`/`--tpm` and `HEADROOM_RPM`/`HEADROOM_TPM` to the Click proxy command for rate-limit parity with the legacy CLI -- closes [#1350](https://github.com/headroomlabs-ai/headroom/issues/1350) (Problem 1). * **proxy:** register `ToolResultInterceptorTransform` in explicit transforms list when `HEADROOM_INTERCEPT_ENABLED` is set — closes [#829](https://github.com/headroomlabs-ai/headroom/issues/829). * **code:** keep Python `from __future__` imports before executable code during AST compression and validate compressed Python with `compile(..., "exec")` so compile-time syntax rules are enforced ([#1233](https://github.com/chopratejas/headroom/issues/1233)). diff --git a/headroom/cli/proxy.py b/headroom/cli/proxy.py index fb28ab35f..18a802892 100644 --- a/headroom/cli/proxy.py +++ b/headroom/cli/proxy.py @@ -247,6 +247,16 @@ def dashboard(port: int, no_open: bool) -> None: @click.option("--no-optimize", is_flag=True, help="Disable optimization (passthrough mode)") @click.option("--no-cache", is_flag=True, help="Disable semantic caching") @click.option("--no-rate-limit", is_flag=True, help="Disable rate limiting") +@click.option( + "--protect-tool-results", + default=None, + envvar="HEADROOM_PROTECT_TOOL_RESULTS", + help=( + "Comma-separated tool names whose results are never lossy-compressed, " + "merged with the built-in defaults (e.g. Bash,WebFetch). " + "Env: HEADROOM_PROTECT_TOOL_RESULTS." + ), +) @click.option( "--rpm", default=None, @@ -807,6 +817,7 @@ def proxy( no_optimize: bool, no_cache: bool, no_rate_limit: bool, + protect_tool_results: str | None, rpm: int | None, tpm: int | None, no_ccr_inject_tool: bool, @@ -890,6 +901,7 @@ def proxy( try: from headroom.proxy.server import ( ProxyConfig, + _parse_csv_tools, _parse_exclude_tools, _parse_tool_profiles, run_server, @@ -1013,6 +1025,9 @@ def proxy( min_tokens_to_crush=_get_env_int_optional("HEADROOM_MIN_TOKENS") or 500, max_items_after_crush=_get_env_int_optional("HEADROOM_MAX_ITEMS") or 50, exclude_tools=_parse_exclude_tools(None) or None, + protect_tool_results=frozenset(_parse_csv_tools(protect_tool_results)) + if protect_tool_results + else frozenset(), tool_profiles=_parse_tool_profiles([]) or None, smart_crusher_with_compaction=_get_env_bool_optional("HEADROOM_SMART_CRUSHER_COMPACTION"), savings_profile=os.environ.get("HEADROOM_SAVINGS_PROFILE") or None, diff --git a/headroom/config.py b/headroom/config.py index 8dc3f35a8..9dfcc4f2f 100644 --- a/headroom/config.py +++ b/headroom/config.py @@ -210,6 +210,8 @@ class AnchorConfig: # Read/Glob/Grep contain exact file contents/search results the agent needs for edits. # Write/Edit record what changes were made — compressing them causes duplicate/conflicting edits. # 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. DEFAULT_EXCLUDE_TOOLS: frozenset[str] = frozenset( { "Read", diff --git a/headroom/proxy/models.py b/headroom/proxy/models.py index bc7a37dbc..50d3e17b5 100644 --- a/headroom/proxy/models.py +++ b/headroom/proxy/models.py @@ -198,6 +198,11 @@ class ProxyConfig: # CLI: --exclude-tools ; env: HEADROOM_EXCLUDE_TOOLS= exclude_tools: set[str] | None = None + # Tool names whose results must never be lossy-compressed (e.g. Bash, WebFetch). + # Merged into exclude_tools before ContentRouter processes the conversation. + # CLI: --protect-tool-results ; env: HEADROOM_PROTECT_TOOL_RESULTS= + protect_tool_results: frozenset[str] = field(default_factory=frozenset) + # Read lifecycle management read_lifecycle: bool = True diff --git a/headroom/proxy/server.py b/headroom/proxy/server.py index 7d64865da..176776076 100644 --- a/headroom/proxy/server.py +++ b/headroom/proxy/server.py @@ -647,12 +647,23 @@ class HeadroomProxy( # ContentRouter, so merge rather than assign. if config.exclude_tools: router_config.exclude_tools = set(DEFAULT_EXCLUDE_TOOLS) | config.exclude_tools + # protect_tool_results: force-merge named tools into the exclude set + # so their results are never lossy-compressed, regardless of mode. + if config.protect_tool_results: + base = ( + router_config.exclude_tools + if router_config.exclude_tools is not None + else set(DEFAULT_EXCLUDE_TOOLS) + ) + router_config.exclude_tools = base | config.protect_tool_results # Token mode: allow compression of older excluded-tool results, # and emit search results grouped by file (path once per file # instead of repeated on every match line). if is_token_mode(config.mode): router_config.protect_recent_reads_fraction = 0.3 router_config.search_group_by_file = True + if config.protect_tool_results: + router_config.protect_recent_reads_fraction = 0.0 # `--compress-user-messages` flips the router's default skip rule. # Off by default for prefix-cache safety; enabled for workloads where # user-message content dominates input (OpenAI/Azure chat with pasted @@ -4160,6 +4171,19 @@ def _parse_exclude_tools(cli_excludes: str | None) -> set[str]: return names +def _parse_csv_tools(raw: str | None) -> set[str]: + """Parse a bare CSV tool-name string without merging HEADROOM_EXCLUDE_TOOLS.""" + names: set[str] = set() + if not raw: + return names + for entry in raw.split(","): + name = entry.strip() + if name: + names.add(name) + names.add(name.lower()) + return names + + def _parse_tool_profiles(cli_profiles: list[str]) -> dict[str, Any]: """Parse tool profiles from CLI args and HEADROOM_TOOL_PROFILES env var. @@ -4380,6 +4404,13 @@ if __name__ == "__main__": "Entries may use glob patterns, e.g. 'mcp__*' to exclude every MCP tool. " "Also settable via HEADROOM_EXCLUDE_TOOLS env var.", ) + parser.add_argument( + "--protect-tool-results", + default=None, + help="Comma-separated tool names whose results are never lossy-compressed, " + "merged with the built-in defaults (e.g. Bash,WebFetch). " + "Also settable via HEADROOM_PROTECT_TOOL_RESULTS env var.", + ) # Caching parser.add_argument("--no-cache", action="store_true", help="Disable caching") @@ -4452,6 +4483,9 @@ if __name__ == "__main__": tool_profiles = _parse_tool_profiles(args.tool_profile) # Parse extra never-compress tools from CLI and env var exclude_tools = _parse_exclude_tools(args.exclude_tools) + protect_tool_results = _parse_csv_tools( + args.protect_tool_results or os.environ.get("HEADROOM_PROTECT_TOOL_RESULTS") + ) config = ProxyConfig( host=_get_env_str("HEADROOM_HOST", args.host), @@ -4507,6 +4541,9 @@ if __name__ == "__main__": ), tool_profiles=tool_profiles if tool_profiles else None, exclude_tools=exclude_tools if exclude_tools else None, + protect_tool_results=frozenset(protect_tool_results) + if protect_tool_results + else frozenset(), mode=normalize_proxy_mode(_get_env_str("HEADROOM_MODE", PROXY_MODE_TOKEN)), compress_user_messages=args.compress_user_messages or _get_env_bool("HEADROOM_COMPRESS_USER_MESSAGES", False), diff --git a/tests/test_content_router_exclude_tools.py b/tests/test_content_router_exclude_tools.py new file mode 100644 index 000000000..de726fa74 --- /dev/null +++ b/tests/test_content_router_exclude_tools.py @@ -0,0 +1,153 @@ +"""Tests for --protect-tool-results / HEADROOM_PROTECT_TOOL_RESULTS. + +Three behavioral tests: +1. protect_tool_results merges into the exclude set so named tools are never + lossy-compressed. +2. _parse_csv_tools parses CSV strings without merging HEADROOM_EXCLUDE_TOOLS. +3. ContentRouter with Bash in exclude_tools passes Bash tool_result verbatim. +""" + +from __future__ import annotations + +import pytest + +from headroom.config import DEFAULT_EXCLUDE_TOOLS +from headroom.proxy.server import ( + HeadroomProxy, + ProxyConfig, + _parse_csv_tools, +) + + +def _build(**overrides: object) -> HeadroomProxy: + config = ProxyConfig( + optimize=False, + cache_enabled=False, + rate_limit_enabled=False, + cost_tracking_enabled=False, + code_aware_enabled=False, + **overrides, + ) + return HeadroomProxy(config) + + +def _router(proxy: HeadroomProxy): + # ContentRouter is the last transform in the Anthropic pipeline. + return proxy.anthropic_pipeline.transforms[-1] + + +# --------------------------------------------------------------------------- +# Test 1: protect_tool_results merges into exclude set +# --------------------------------------------------------------------------- + + +def test_protect_tool_results_merges_into_exclude_set() -> None: + """Bash added via protect_tool_results must appear in exclude_tools alongside + the built-in defaults (e.g. Read), base-fails / head-passes. + + The frozenset is merged as-is; lowercase normalization is handled by + _parse_exclude_tools on the CLI/env path (tested separately). + """ + proxy = _build(protect_tool_results=frozenset({"Bash", "bash"})) + exclude = _router(proxy).config.exclude_tools + + assert exclude is not None, "exclude_tools must be set when protect_tool_results is non-empty" + assert "Bash" in exclude, "Bash must be in exclude_tools after protect_tool_results merges" + assert "bash" in exclude, ( + "lowercase bash must be in exclude_tools after protect_tool_results merges" + ) + assert "Read" in exclude, "Read (built-in default) must still be in exclude_tools" + + +def test_protect_tool_results_disables_age_decay_in_token_mode() -> None: + """In token mode, protect_tool_results forces protect_recent_reads_fraction to 0.0 + so protected tools are never compressed by age-decay.""" + proxy = _build(protect_tool_results=frozenset({"Bash", "bash"}), mode="token") + assert _router(proxy).config.protect_recent_reads_fraction == 0.0 + + +# --------------------------------------------------------------------------- +# Test 2: CSV env var / CLI string parsing +# --------------------------------------------------------------------------- + + +def test_protect_tool_results_env_var_csv() -> None: + """_parse_csv_tools parses a comma-separated value into both original-case + and lowercase entries without merging HEADROOM_EXCLUDE_TOOLS.""" + result = _parse_csv_tools("Bash,WebFetch") + + assert "Bash" in result + assert "bash" in result + assert "WebFetch" in result + assert "webfetch" in result + + +# --------------------------------------------------------------------------- +# Test 3: Bash tool_result passthrough when protected +# --------------------------------------------------------------------------- + + +def test_bash_tool_result_passthrough_when_protected() -> None: + """When Bash is in exclude_tools (via protect_tool_results), its tool_result + content passes through the ContentRouter verbatim without lossy compression. + Base-fails / head-passes.""" + pytest.importorskip("tiktoken") # needed for OpenAI tokenizer + + from headroom.providers import OpenAIProvider + from headroom.tokenizer import Tokenizer + from headroom.transforms.content_router import ContentRouter, ContentRouterConfig + + provider = OpenAIProvider() + token_counter = provider.get_token_counter("gpt-4o") + tokenizer = Tokenizer(token_counter, "gpt-4o") + + # Build router with Bash explicitly in exclude_tools + config = ContentRouterConfig( + min_section_tokens=10, + exclude_tools=set(DEFAULT_EXCLUDE_TOOLS) | {"Bash", "bash"}, + ) + router = ContentRouter(config) + + bash_output = "\n".join( + f"line {i}: some output from a bash command that is long enough to compress" + for i in range(80) + ) + messages = [ + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_bash_1", + "type": "function", + "function": {"name": "Bash", "arguments": "{}"}, + } + ], + }, + { + "role": "tool", + "tool_call_id": "call_bash_1", + "content": bash_output, + }, + ] + + result = router.apply(messages, tokenizer) + + # Bash tool_result must pass through unchanged + tool_msg = next(m for m in result.messages if m.get("tool_call_id") == "call_bash_1") + assert tool_msg["content"] == bash_output, ( + "Bash tool_result must be verbatim when Bash is in exclude_tools" + ) + assert "router:excluded:tool" in result.transforms_applied + + +# --------------------------------------------------------------------------- +# Baseline: Bash NOT in DEFAULT_EXCLUDE_TOOLS (unchanged by this PR) +# --------------------------------------------------------------------------- + + +def test_bash_not_in_default_exclude_tools() -> None: + """Bash must remain absent from DEFAULT_EXCLUDE_TOOLS; protect_tool_results + is the opt-in path.""" + assert "Bash" not in DEFAULT_EXCLUDE_TOOLS + assert "bash" not in DEFAULT_EXCLUDE_TOOLS