From 7d99a71285b3b3ec5c475c0676aa1a4a0887fbe7 Mon Sep 17 00:00:00 2001
From: Manorit Chawdhry <m-chawdhry@ti.com>
Date: Mon, 4 May 2026 22:24:11 +0530
Subject: [PATCH] fix: expose code-aware flag

headroom proxy now accepts --code-aware so callers do not need to drop to the lower-level server entrypoint.

Keep the existing env fallback so HEADROOM_CODE_AWARE_ENABLED still works when the flag is omitted.

Assisted-by: Sisyphus gpt-5.4-mini

Signed-off-by: Manorit Chawdhry <m-chawdhry@ti.com>
---
 headroom/cli/proxy.py       | 17 +++++++++
 tests/test_cli_proxy_env.py | 71 +++++++++++++++++++++++++++++++++++++
 2 files changed, 88 insertions(+)

diff --git a/headroom/cli/proxy.py b/headroom/cli/proxy.py
index 47333ec35..c38a0762f 100644
--- a/headroom/cli/proxy.py
+++ b/headroom/cli/proxy.py
@@ -12,6 +12,13 @@ from headroom.proxy.modes import PROXY_MODE_TOKEN, normalize_proxy_mode
 from .main import main
 
 
+def _get_env_bool(name: str, default: bool) -> bool:
+    val = os.environ.get(name)
+    if val is None:
+        return default
+    return val.lower() in ("true", "1", "yes", "on")
+
+
 @main.command()
 @click.option(
     "--host",
@@ -90,6 +97,14 @@ from .main import main
 @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(
+    "--code-aware",
+    is_flag=True,
+    help=(
+        "Enable code-aware compression in the proxy (env: HEADROOM_CODE_AWARE_ENABLED). "
+        "Code-aware remains off by default."
+    ),
+)
 @click.option(
     "--proxy-extension",
     "proxy_extension",
@@ -348,6 +363,7 @@ def proxy(
     no_optimize: bool,
     no_cache: bool,
     no_rate_limit: bool,
+    code_aware: bool,
     proxy_extension: tuple[str, ...],
     no_subscription_tracking: bool,
     subscription_poll_interval: int | None,
@@ -516,6 +532,7 @@ def proxy(
         budget_limit_usd=budget,
         # Code graph: live file watcher for incremental reindexing
         code_graph_watcher=code_graph,
+        code_aware_enabled=code_aware or _get_env_bool("HEADROOM_CODE_AWARE_ENABLED", False),
         # Read lifecycle: ON by default (use --no-read-lifecycle to disable)
         read_lifecycle=not no_read_lifecycle,
         # Memory System (Multi-Provider with auto-detection)
diff --git a/tests/test_cli_proxy_env.py b/tests/test_cli_proxy_env.py
index c8c910493..c5993ca02 100644
--- a/tests/test_cli_proxy_env.py
+++ b/tests/test_cli_proxy_env.py
@@ -80,6 +80,77 @@ class TestCLIProxyEnvVars:
         assert result.exit_code == 0, result.output
         assert captured_config["config"].budget_limit_usd == 100.5
 
+    def test_code_aware_enabled_from_env(self, runner):
+        """HEADROOM_CODE_AWARE_ENABLED env var should be passed to ProxyConfig."""
+        captured_config = {}
+
+        def mock_run_server(config):
+            captured_config["config"] = config
+
+        with patch("headroom.proxy.server.run_server", mock_run_server):
+            result = runner.invoke(
+                main,
+                ["proxy"],
+                env={"HEADROOM_CODE_AWARE_ENABLED": "true"},
+                catch_exceptions=False,
+            )
+
+        assert result.exit_code == 0, result.output
+        assert captured_config["config"].code_aware_enabled is True
+
+    def test_code_aware_enabled_defaults_false(self, runner):
+        """Without HEADROOM_CODE_AWARE_ENABLED, code-aware stays disabled in the wrapper."""
+        captured_config = {}
+
+        def mock_run_server(config):
+            captured_config["config"] = config
+
+        env = {k: v for k, v in os.environ.items() if k != "HEADROOM_CODE_AWARE_ENABLED"}
+
+        with (
+            patch("headroom.proxy.server.run_server", mock_run_server),
+            patch.dict(os.environ, env, clear=True),
+        ):
+            result = runner.invoke(
+                main,
+                ["proxy"],
+                catch_exceptions=False,
+            )
+
+        assert result.exit_code == 0, result.output
+        assert captured_config["config"].code_aware_enabled is False
+
+    def test_code_aware_enabled_from_cli_flag(self, runner):
+        """--code-aware should enable code-aware compression in the wrapper."""
+        captured_config = {}
+
+        def mock_run_server(config):
+            captured_config["config"] = config
+
+        with patch("headroom.proxy.server.run_server", mock_run_server):
+            result = runner.invoke(main, ["proxy", "--code-aware"], catch_exceptions=False)
+
+        assert result.exit_code == 0, result.output
+        assert captured_config["config"].code_aware_enabled is True
+
+    def test_code_aware_flag_overrides_env_var(self, runner):
+        """--code-aware should win over HEADROOM_CODE_AWARE_ENABLED=false."""
+        captured_config = {}
+
+        def mock_run_server(config):
+            captured_config["config"] = config
+
+        with patch("headroom.proxy.server.run_server", mock_run_server):
+            result = runner.invoke(
+                main,
+                ["proxy", "--code-aware"],
+                env={"HEADROOM_CODE_AWARE_ENABLED": "false"},
+                catch_exceptions=False,
+            )
+
+        assert result.exit_code == 0, result.output
+        assert captured_config["config"].code_aware_enabled is True
+
     def test_openai_target_api_url_from_env(self, runner):
         """OPENAI_TARGET_API_URL env var should be passed to ProxyConfig."""
         captured_config = {}