From 4f560bccc7c4c96cddb7a45a47492d092cb78b4d Mon Sep 17 00:00:00 2001 From: Tejas Chopra Date: Tue, 30 Jun 2026 15:30:22 -0700 Subject: [PATCH] feat(proxy): add --force-kompress-all to route all content through kompress-v2-base (#1613) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Description Adds an opt-in flag that routes **all** compressible content through Kompress (`kompress-v2-base`), bypassing per-type compressor selection (SmartCrusher / CodeAware / log / diff / html / tabular / search). For deployments that prefer a single uniform compressor over the per-type set, at a deliberate cost of per-type structural fidelity. The mechanism already existed: `ContentRouter` reads a `force_kompress` runtime kwarg but nothing turned it on. This PR wires it to user-facing config (CLI + env), defaulting off. Closes # N/A ## Type of Change - [ ] Bug fix (non-breaking change that fixes an issue) - [x] New feature (non-breaking change that adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to change) - [ ] Documentation update - [ ] Performance improvement - [ ] Code refactoring (no functional changes) ## Changes Made - Add `force_kompress_all` to `ProxyConfig` (`headroom/proxy/models.py`) and `ContentRouterConfig` (`headroom/transforms/content_router.py`). - Default the existing `force_kompress` runtime path from config: `kwargs.get("force_kompress", self.config.force_kompress_all)` — a per-request kwarg still overrides. - Expose `--force-kompress-all` CLI flag and `HEADROOM_FORCE_KOMPRESS_ALL=1` env, mirroring the existing `--disable-kompress` pattern (both the env factory and the `__main__` CLI path). - Add `tests/test_force_kompress_all.py`. **Safety preserved:** the flag changes *strategy selection only*. The Read/Glob/Grep exclusion (`excluded_tool_ids`) runs *before* any compressor, and the tool-output reversibility gate (`#1307`/`#1479`) runs *after* — neither is reachable from the strategy choice. So tool ground truth stays verbatim. ## Testing - [x] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check`) - [x] Type checking passes (`mypy`) - [x] New tests added for new functionality - [ ] Manual testing performed (see Real Behavior Proof → Not tested) ### Test Output ```text $ ruff check headroom/proxy/models.py headroom/transforms/content_router.py headroom/proxy/server.py tests/test_force_kompress_all.py All checks passed! $ ruff format --check 4 files already formatted $ mypy headroom/proxy/models.py headroom/transforms/content_router.py headroom/proxy/server.py tests/test_force_kompress_all.py Success: no issues found in 4 source files $ pytest tests/test_force_kompress_all.py tests/test_content_router_exclude_tools.py -q tests/test_force_kompress_all.py .... [ 44%] tests/test_content_router_exclude_tools.py ..... [100%] ============================== 9 passed in 1.31s =============================== ``` ## Real Behavior Proof - **Environment:** macOS (Darwin 25.4.0), Python 3.12.6, project `.venv`. - **Exact command / steps:** Constructed `ContentRouter(ContentRouterConfig(force_kompress_all=True))` and drove the real `apply()` entry point (see `tests/test_force_kompress_all.py`) to verify: (1) the config resolves the runtime flag on; (2) an explicit `force_kompress=False` kwarg overrides it; (3) a `Read` tool_result is passed through **verbatim** with the flag on (`router:excluded:tool` marker present). Plus the full ruff/mypy/pytest suite above. - **Observed result:** 9 tests pass. Read tool output is unchanged (byte-for-byte) under `force_kompress_all=True`; the per-request kwarg override works; the existing exclude-tools suite still passes through `HeadroomProxy` (which now builds `ContentRouterConfig(force_kompress_all=...)`). - **Not tested:** Live proxy end-to-end against a real upstream with the `kompress-v2-base` ONNX model compressing real traffic; aggregate savings/accuracy deltas on a real workload. The unit tests assert the **routing decision and the Read/Glob carve-out**, not model output quality or ratio. ## 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 — N/A (documented inline via config docstring + `--help`; see Additional Notes) - [x] My changes generate no new warnings - [x] I have added tests that prove my feature works - [x] New and existing unit tests pass locally with my changes - [ ] I have updated the CHANGELOG.md if applicable — N/A (Release Please generates it from the `feat(proxy):` commit) ## Additional Notes - **Accuracy tradeoff (intentional):** forcing Kompress on all types trades per-type structural fidelity (and possibly compression ratio, since SmartCrusher/CodeAware can beat a general model on their native type) for a single uniform compressor. Off by default; opt-in per deployment. Correctness is *not* affected — excluded tools and reversibility-gated tool ground truth are never touched. - **Docs:** behavior is documented inline (CLI `--help` text + `ProxyConfig` docstring). Happy to add a README/wiki note if maintainers want one. --- headroom/proxy/models.py | 9 +++ headroom/proxy/server.py | 16 +++++ headroom/transforms/content_router.py | 7 +- tests/test_force_kompress_all.py | 92 +++++++++++++++++++++++++++ 4 files changed, 123 insertions(+), 1 deletion(-) create mode 100644 tests/test_force_kompress_all.py 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