mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
fix(proxy): add --protect-tool-results to prevent lossy compression of exact-output Bash results (#1374)
## 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 <mxjerrett@gmail.com>
This commit is contained in:
parent
e06b61671f
commit
51d4bcfc11
6 changed files with 213 additions and 1 deletions
|
|
@ -8,7 +8,6 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
|
||||||
|
|
||||||
## Unreleased
|
## Unreleased
|
||||||
|
|
||||||
|
|
||||||
### Changed
|
### 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.
|
* **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)).
|
* **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)).
|
* **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)).
|
* **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).
|
* **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).
|
* **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)).
|
* **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)).
|
||||||
|
|
|
||||||
|
|
@ -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-optimize", is_flag=True, help="Disable optimization (passthrough mode)")
|
||||||
@click.option("--no-cache", is_flag=True, help="Disable semantic caching")
|
@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("--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(
|
@click.option(
|
||||||
"--rpm",
|
"--rpm",
|
||||||
default=None,
|
default=None,
|
||||||
|
|
@ -807,6 +817,7 @@ def proxy(
|
||||||
no_optimize: bool,
|
no_optimize: bool,
|
||||||
no_cache: bool,
|
no_cache: bool,
|
||||||
no_rate_limit: bool,
|
no_rate_limit: bool,
|
||||||
|
protect_tool_results: str | None,
|
||||||
rpm: int | None,
|
rpm: int | None,
|
||||||
tpm: int | None,
|
tpm: int | None,
|
||||||
no_ccr_inject_tool: bool,
|
no_ccr_inject_tool: bool,
|
||||||
|
|
@ -890,6 +901,7 @@ def proxy(
|
||||||
try:
|
try:
|
||||||
from headroom.proxy.server import (
|
from headroom.proxy.server import (
|
||||||
ProxyConfig,
|
ProxyConfig,
|
||||||
|
_parse_csv_tools,
|
||||||
_parse_exclude_tools,
|
_parse_exclude_tools,
|
||||||
_parse_tool_profiles,
|
_parse_tool_profiles,
|
||||||
run_server,
|
run_server,
|
||||||
|
|
@ -1013,6 +1025,9 @@ def proxy(
|
||||||
min_tokens_to_crush=_get_env_int_optional("HEADROOM_MIN_TOKENS") or 500,
|
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,
|
max_items_after_crush=_get_env_int_optional("HEADROOM_MAX_ITEMS") or 50,
|
||||||
exclude_tools=_parse_exclude_tools(None) or None,
|
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,
|
tool_profiles=_parse_tool_profiles([]) or None,
|
||||||
smart_crusher_with_compaction=_get_env_bool_optional("HEADROOM_SMART_CRUSHER_COMPACTION"),
|
smart_crusher_with_compaction=_get_env_bool_optional("HEADROOM_SMART_CRUSHER_COMPACTION"),
|
||||||
savings_profile=os.environ.get("HEADROOM_SAVINGS_PROFILE") or None,
|
savings_profile=os.environ.get("HEADROOM_SAVINGS_PROFILE") or None,
|
||||||
|
|
|
||||||
|
|
@ -210,6 +210,8 @@ class AnchorConfig:
|
||||||
# Read/Glob/Grep contain exact file contents/search results the agent needs for edits.
|
# 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.
|
# 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.
|
# 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(
|
DEFAULT_EXCLUDE_TOOLS: frozenset[str] = frozenset(
|
||||||
{
|
{
|
||||||
"Read",
|
"Read",
|
||||||
|
|
|
||||||
|
|
@ -198,6 +198,11 @@ class ProxyConfig:
|
||||||
# CLI: --exclude-tools <name1,name2>; env: HEADROOM_EXCLUDE_TOOLS=<name1,name2>
|
# CLI: --exclude-tools <name1,name2>; env: HEADROOM_EXCLUDE_TOOLS=<name1,name2>
|
||||||
exclude_tools: set[str] | None = None
|
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 <name1,name2>; env: HEADROOM_PROTECT_TOOL_RESULTS=<name1,name2>
|
||||||
|
protect_tool_results: frozenset[str] = field(default_factory=frozenset)
|
||||||
|
|
||||||
# Read lifecycle management
|
# Read lifecycle management
|
||||||
read_lifecycle: bool = True
|
read_lifecycle: bool = True
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -647,12 +647,23 @@ class HeadroomProxy(
|
||||||
# ContentRouter, so merge rather than assign.
|
# ContentRouter, so merge rather than assign.
|
||||||
if config.exclude_tools:
|
if config.exclude_tools:
|
||||||
router_config.exclude_tools = set(DEFAULT_EXCLUDE_TOOLS) | 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,
|
# Token mode: allow compression of older excluded-tool results,
|
||||||
# and emit search results grouped by file (path once per file
|
# and emit search results grouped by file (path once per file
|
||||||
# instead of repeated on every match line).
|
# instead of repeated on every match line).
|
||||||
if is_token_mode(config.mode):
|
if is_token_mode(config.mode):
|
||||||
router_config.protect_recent_reads_fraction = 0.3
|
router_config.protect_recent_reads_fraction = 0.3
|
||||||
router_config.search_group_by_file = True
|
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.
|
# `--compress-user-messages` flips the router's default skip rule.
|
||||||
# Off by default for prefix-cache safety; enabled for workloads where
|
# Off by default for prefix-cache safety; enabled for workloads where
|
||||||
# user-message content dominates input (OpenAI/Azure chat with pasted
|
# 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
|
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]:
|
def _parse_tool_profiles(cli_profiles: list[str]) -> dict[str, Any]:
|
||||||
"""Parse tool profiles from CLI args and HEADROOM_TOOL_PROFILES env var.
|
"""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. "
|
"Entries may use glob patterns, e.g. 'mcp__*' to exclude every MCP tool. "
|
||||||
"Also settable via HEADROOM_EXCLUDE_TOOLS env var.",
|
"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
|
# Caching
|
||||||
parser.add_argument("--no-cache", action="store_true", help="Disable 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)
|
tool_profiles = _parse_tool_profiles(args.tool_profile)
|
||||||
# Parse extra never-compress tools from CLI and env var
|
# Parse extra never-compress tools from CLI and env var
|
||||||
exclude_tools = _parse_exclude_tools(args.exclude_tools)
|
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(
|
config = ProxyConfig(
|
||||||
host=_get_env_str("HEADROOM_HOST", args.host),
|
host=_get_env_str("HEADROOM_HOST", args.host),
|
||||||
|
|
@ -4507,6 +4541,9 @@ if __name__ == "__main__":
|
||||||
),
|
),
|
||||||
tool_profiles=tool_profiles if tool_profiles else None,
|
tool_profiles=tool_profiles if tool_profiles else None,
|
||||||
exclude_tools=exclude_tools if exclude_tools 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)),
|
mode=normalize_proxy_mode(_get_env_str("HEADROOM_MODE", PROXY_MODE_TOKEN)),
|
||||||
compress_user_messages=args.compress_user_messages
|
compress_user_messages=args.compress_user_messages
|
||||||
or _get_env_bool("HEADROOM_COMPRESS_USER_MESSAGES", False),
|
or _get_env_bool("HEADROOM_COMPRESS_USER_MESSAGES", False),
|
||||||
|
|
|
||||||
153
tests/test_content_router_exclude_tools.py
Normal file
153
tests/test_content_router_exclude_tools.py
Normal file
|
|
@ -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
|
||||||
Loading…
Add table
Add a link
Reference in a new issue