diff --git a/headroom/cli/proxy.py b/headroom/cli/proxy.py index a45f1bf33..8fe639f02 100644 --- a/headroom/cli/proxy.py +++ b/headroom/cli/proxy.py @@ -159,6 +159,31 @@ def _selected_context_tool() -> str: @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( + "--no-ccr-inject-tool", + is_flag=True, + envvar="HEADROOM_NO_CCR_INJECT_TOOL", + help=( + "Don't inject the CCR headroom_retrieve tool. Run compression-only — " + "for streaming / non-MCP clients that can't resolve the retrieve tool " + "and would otherwise error on it. Env: HEADROOM_NO_CCR_INJECT_TOOL." + ), +) +@click.option( + "--no-ccr-marker", + is_flag=True, + envvar="HEADROOM_NO_CCR_MARKER", + help=("Don't add CCR retrieval markers to compressed content. Env: HEADROOM_NO_CCR_MARKER."), +) +@click.option( + "--no-ccr-proactive-expansion", + is_flag=True, + envvar="HEADROOM_NO_CCR_PROACTIVE_EXPANSION", + help=( + "Disable proactive expansion of previously compressed content. " + "Env: HEADROOM_NO_CCR_PROACTIVE_EXPANSION." + ), +) @click.option( "--proxy-extension", "proxy_extension", @@ -555,6 +580,9 @@ def proxy( no_optimize: bool, no_cache: bool, no_rate_limit: bool, + no_ccr_inject_tool: bool, + no_ccr_marker: bool, + no_ccr_proactive_expansion: bool, proxy_extension: tuple[str, ...], no_subscription_tracking: bool, subscription_poll_interval: int | None, @@ -729,6 +757,12 @@ def proxy( optimize=not no_optimize, cache_enabled=not no_cache, rate_limit_enabled=not no_rate_limit, + # CCR opt-outs for compression-only deployments (streaming / non-MCP + # clients that can't resolve the injected retrieve tool). Defaults keep + # CCR fully on; each flag flips one dataclass default to False. + ccr_inject_tool=not no_ccr_inject_tool, + ccr_inject_marker=not no_ccr_marker, + ccr_proactive_expansion=not no_ccr_proactive_expansion, # Flatten repeat-flag tuple AND any comma-separated values inside it. # `--proxy-extension a,b --proxy-extension c` and `HEADROOM_PROXY_EXTENSIONS=a,b,c` # both yield ["a", "b", "c"]. None when nothing was supplied. diff --git a/headroom/proxy/models.py b/headroom/proxy/models.py index 246df1f25..125da55a0 100644 --- a/headroom/proxy/models.py +++ b/headroom/proxy/models.py @@ -118,6 +118,10 @@ class ProxyConfig: # CCR Tool Injection ccr_inject_tool: bool = True ccr_inject_system_instructions: bool = False + # Proxy-level mirror of ContentRouterConfig.ccr_inject_marker, so retrieval + # markers can be toggled from the CLI (--no-ccr-marker). Threaded into the + # router in server.py; default preserves current behavior. + ccr_inject_marker: bool = True # CCR Response Handling ccr_handle_responses: bool = True diff --git a/headroom/proxy/server.py b/headroom/proxy/server.py index afac7fa9d..eda36ff6f 100644 --- a/headroom/proxy/server.py +++ b/headroom/proxy/server.py @@ -364,6 +364,7 @@ class HeadroomProxy( enable_code_aware=config.code_aware_enabled, tool_profiles=config.tool_profiles, read_lifecycle=ReadLifecycleConfig(enabled=config.read_lifecycle), + ccr_inject_marker=config.ccr_inject_marker, ) # A non-None exclude_tools replaces DEFAULT_EXCLUDE_TOOLS in # ContentRouter, so merge rather than assign. diff --git a/tests/test_cli_proxy_env.py b/tests/test_cli_proxy_env.py index 289f1700c..78ff22482 100644 --- a/tests/test_cli_proxy_env.py +++ b/tests/test_cli_proxy_env.py @@ -596,6 +596,90 @@ class TestCLIAnyllmProviderEnv: assert captured_config["config"].anyllm_provider == "groq" +class TestCLICompressionOnlyFlags: + """The CCR opt-out flags must flip the corresponding ProxyConfig fields. + + These enable a compression-only deployment for streaming / non-MCP clients + that can't resolve the injected headroom_retrieve tool (issue #645). + """ + + def test_ccr_defaults_on(self, runner): + """Without flags, all three CCR toggles stay enabled (no behavior change).""" + captured_config = {} + + def mock_run_server(config, **kwargs): + captured_config["config"] = config + + with patch("headroom.proxy.server.run_server", mock_run_server): + result = runner.invoke(main, ["proxy"], catch_exceptions=False) + + assert result.exit_code == 0, result.output + cfg = captured_config["config"] + assert cfg.ccr_inject_tool is True + assert cfg.ccr_inject_marker is True + assert cfg.ccr_proactive_expansion is True + + def test_no_ccr_inject_tool_flag(self, runner): + """--no-ccr-inject-tool disables retrieve-tool injection only.""" + captured_config = {} + + def mock_run_server(config, **kwargs): + captured_config["config"] = config + + with patch("headroom.proxy.server.run_server", mock_run_server): + result = runner.invoke(main, ["proxy", "--no-ccr-inject-tool"], catch_exceptions=False) + + assert result.exit_code == 0, result.output + cfg = captured_config["config"] + assert cfg.ccr_inject_tool is False + # Untouched flags remain on. + assert cfg.ccr_inject_marker is True + assert cfg.ccr_proactive_expansion is True + + def test_compression_only_all_flags(self, runner): + """All three flags together yield a compression-only config.""" + captured_config = {} + + def mock_run_server(config, **kwargs): + captured_config["config"] = config + + with patch("headroom.proxy.server.run_server", mock_run_server): + result = runner.invoke( + main, + [ + "proxy", + "--no-ccr-inject-tool", + "--no-ccr-marker", + "--no-ccr-proactive-expansion", + ], + catch_exceptions=False, + ) + + assert result.exit_code == 0, result.output + cfg = captured_config["config"] + assert cfg.ccr_inject_tool is False + assert cfg.ccr_inject_marker is False + assert cfg.ccr_proactive_expansion is False + + def test_no_ccr_marker_from_env(self, runner): + """HEADROOM_NO_CCR_MARKER env var disables marker injection.""" + captured_config = {} + + def mock_run_server(config, **kwargs): + captured_config["config"] = config + + with patch("headroom.proxy.server.run_server", mock_run_server): + result = runner.invoke( + main, + ["proxy"], + env={"HEADROOM_NO_CCR_MARKER": "1"}, + catch_exceptions=False, + ) + + assert result.exit_code == 0, result.output + assert captured_config["config"].ccr_inject_marker is False + + class TestArgparseBackendValidation: """Test that the argparse path (python -m headroom.proxy.server) accepts litellm-* backends."""