diff --git a/headroom/proxy/models.py b/headroom/proxy/models.py index d9b3302c4..34b6e8e25 100644 --- a/headroom/proxy/models.py +++ b/headroom/proxy/models.py @@ -171,6 +171,15 @@ class ProxyConfig: disable_kompress_anthropic: bool | None = None disable_kompress_openai: bool | None = None + # Force ALL compressible content through Kompress (kompress-v2-base), + # bypassing per-type compressor selection (SmartCrusher/CodeAware/log/ + # diff/html/tabular/search). Tool ground truth stays protected: excluded + # tools (Read/Glob/Grep/...) and reversibility-gated tool output are never + # touched. Off by default; opt-in for systems that want one uniform + # compressor at the cost of per-type structural fidelity. + # CLI: --force-kompress-all; env: HEADROOM_FORCE_KOMPRESS_ALL=1. + force_kompress_all: bool = False + # Code graph live watcher (triggers incremental reindex on file changes) code_graph_watcher: bool = False diff --git a/headroom/proxy/server.py b/headroom/proxy/server.py index 5b35afd57..5d861e6d3 100644 --- a/headroom/proxy/server.py +++ b/headroom/proxy/server.py @@ -673,6 +673,7 @@ class HeadroomProxy( profile_kwargs.get("smart_crusher_with_compaction", True), ), ccr_inject_marker=config.ccr_inject_marker, + force_kompress_all=config.force_kompress_all, ) if config.disable_kompress: router_config.enable_kompress = False @@ -4078,6 +4079,7 @@ def _proxy_config_from_env() -> ProxyConfig: disable_kompress_fallback=_get_env_bool("HEADROOM_DISABLE_KOMPRESS_FALLBACK", False), disable_kompress_anthropic=_get_env_optional_bool("HEADROOM_DISABLE_KOMPRESS_ANTHROPIC"), disable_kompress_openai=_get_env_optional_bool("HEADROOM_DISABLE_KOMPRESS_OPENAI"), + force_kompress_all=_get_env_bool("HEADROOM_FORCE_KOMPRESS_ALL", False), max_connections=_get_env_int("HEADROOM_MAX_CONNECTIONS", 500), max_keepalive_connections=_get_env_int("HEADROOM_MAX_KEEPALIVE", 100), keepalive_expiry=_get_env_float("HEADROOM_KEEPALIVE_EXPIRY", 90.0), @@ -4562,6 +4564,16 @@ if __name__ == "__main__": const=False, help="Force-enable Kompress for the OpenAI/Codex pipeline, overriding --disable-kompress.", ) + parser.add_argument( + "--force-kompress-all", + action="store_true", + help=( + "Route ALL compressible content through Kompress (kompress-v2-base), " + "bypassing per-type compressor selection. Tool ground truth " + "(Read/Glob/... and reversibility-gated output) is still never touched. " + "Also settable via HEADROOM_FORCE_KOMPRESS_ALL=1." + ), + ) parser.add_argument( "--exclude-tools", default=None, @@ -4640,6 +4652,9 @@ if __name__ == "__main__": if args.disable_kompress_openai is not None else _get_env_optional_bool("HEADROOM_DISABLE_KOMPRESS_OPENAI") ) + force_kompress_all = args.force_kompress_all or _get_env_bool( + "HEADROOM_FORCE_KOMPRESS_ALL", False + ) # Set OpenRouter API key from CLI if provided if hasattr(args, "openrouter_api_key") and args.openrouter_api_key: @@ -4694,6 +4709,7 @@ if __name__ == "__main__": disable_kompress_fallback=disable_kompress_fallback, disable_kompress_anthropic=disable_kompress_anthropic, disable_kompress_openai=disable_kompress_openai, + force_kompress_all=force_kompress_all, # Connection pool settings max_connections=_get_env_int("HEADROOM_MAX_CONNECTIONS", args.max_connections), max_keepalive_connections=_get_env_int("HEADROOM_MAX_KEEPALIVE", args.max_keepalive), diff --git a/headroom/transforms/content_router.py b/headroom/transforms/content_router.py index a0e45ecee..d9ce37096 100644 --- a/headroom/transforms/content_router.py +++ b/headroom/transforms/content_router.py @@ -631,6 +631,9 @@ class ContentRouterConfig: # Routing preferences prefer_code_aware_for_code: bool = False # Disabled: let code pass through unmangled + # Route ALL compressible content to Kompress, skipping per-type selection. + # Tool exclusion (Read/Glob/...) and reversibility gates still apply. + force_kompress_all: bool = False mixed_content_threshold: int = 2 # Min types to consider mixed min_section_tokens: int = 20 # Min tokens to compress a section @@ -2468,7 +2471,9 @@ class ContentRouter(Transform): ) # Store runtime options on self for access by _route_and_compress_block self._runtime_target_ratio: float | None = kwargs.get("target_ratio") - self._runtime_force_kompress: bool = bool(kwargs.get("force_kompress", False)) + self._runtime_force_kompress: bool = bool( + kwargs.get("force_kompress", self.config.force_kompress_all) + ) self._runtime_kompress_model: str | None = kwargs.get("kompress_model") # F2.2: capture the per-request CompressionPolicy so # ``_record_to_toin`` can gate TOIN writes on diff --git a/tests/test_force_kompress_all.py b/tests/test_force_kompress_all.py new file mode 100644 index 000000000..19476e3c5 --- /dev/null +++ b/tests/test_force_kompress_all.py @@ -0,0 +1,92 @@ +"""Tests for the --force-kompress-all / HEADROOM_FORCE_KOMPRESS_ALL flag. + +force_kompress_all routes ALL compressible content through Kompress, bypassing +per-type compressor selection. Critically it must NOT change protection: excluded +tools (Read/Glob/...) stay verbatim. These tests verify the config -> runtime +wiring and that the Read/Glob carve-out still holds with the flag on. They use an +excluded tool's output so no Kompress model load is required. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import pytest + +from headroom.config import DEFAULT_EXCLUDE_TOOLS +from headroom.transforms.content_router import ContentRouter, ContentRouterConfig + +if TYPE_CHECKING: + from headroom.tokenizer import Tokenizer + + +def _tokenizer() -> Tokenizer: + from headroom.providers import OpenAIProvider + from headroom.tokenizer import Tokenizer + + provider = OpenAIProvider() + return Tokenizer(provider.get_token_counter("gpt-4o"), "gpt-4o") + + +def _read_messages() -> list[dict]: + """A Read tool_result. Read is in DEFAULT_EXCLUDE_TOOLS, so it is never compressed.""" + file_dump = "\n".join(f"line {i}: contents of a file that Read returned" for i in range(80)) + return [ + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_read_1", + "type": "function", + "function": {"name": "Read", "arguments": "{}"}, + } + ], + }, + {"role": "tool", "tool_call_id": "call_read_1", "content": file_dump}, + ] + + +def test_read_is_a_default_excluded_tool() -> None: + """Guards the carve-out's premise: Read ships in DEFAULT_EXCLUDE_TOOLS.""" + assert "Read" in DEFAULT_EXCLUDE_TOOLS + + +def test_config_sets_runtime_force_kompress() -> None: + """force_kompress_all=True in config resolves the runtime flag on; default off.""" + pytest.importorskip("tiktoken") + tokenizer = _tokenizer() + + on = ContentRouter(ContentRouterConfig(force_kompress_all=True)) + on.apply(_read_messages(), tokenizer) + assert on._runtime_force_kompress is True + + off = ContentRouter(ContentRouterConfig()) + off.apply(_read_messages(), tokenizer) + assert off._runtime_force_kompress is False + + +def test_per_request_kwarg_overrides_config() -> None: + """An explicit force_kompress kwarg still wins over the config default.""" + pytest.importorskip("tiktoken") + tokenizer = _tokenizer() + + router = ContentRouter(ContentRouterConfig(force_kompress_all=True)) + router.apply(_read_messages(), tokenizer, force_kompress=False) + assert router._runtime_force_kompress is False + + +def test_read_output_verbatim_under_force_kompress_all() -> None: + """The carve-out: with force_kompress_all on, Read output (an excluded tool) + is passed through verbatim — never routed to Kompress.""" + pytest.importorskip("tiktoken") + tokenizer = _tokenizer() + + messages = _read_messages() + original = messages[1]["content"] + router = ContentRouter(ContentRouterConfig(force_kompress_all=True, min_section_tokens=10)) + result = router.apply(messages, tokenizer) + + tool_msg = next(m for m in result.messages if m.get("tool_call_id") == "call_read_1") + assert tool_msg["content"] == original, "Read tool_result must stay verbatim" + assert "router:excluded:tool" in result.transforms_applied