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>
This commit is contained in:
Manorit Chawdhry 2026-05-04 22:24:11 +05:30
parent 1a3098a48f
commit 7d99a71285
2 changed files with 88 additions and 0 deletions

View file

@ -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)

View file

@ -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 = {}