0ef5fcb1c5
Security / Dependency audit (pip-audit) (push) Has been cancelled
Security / CodeQL (javascript-typescript) (push) Has been cancelled
Security / CodeQL (python) (push) Has been cancelled
Security / Secret scan (gitleaks) (push) Has been cancelled
rust / test (ubuntu) (push) Has been cancelled
rust / simulator e2e (macos-latest) (push) Has been cancelled
rust / simulator e2e (ubuntu-latest) (push) Has been cancelled
rust / simulator e2e (windows-latest) (push) Has been cancelled
rust / wheels (aarch64-apple-darwin) (push) Has been cancelled
rust / wheels (x86_64-unknown-linux-gnu) (push) Has been cancelled
rust / wheels (x86_64-apple-darwin) (push) Has been cancelled
rust / audit (push) Has been cancelled
rust / parity (nightly, allowed to fail during Phase 0) (push) Has been cancelled
CI / commitlint (push) Has been skipped
Dev Containers / validate (.devcontainer/devcontainer.json, default) (push) Failing after 0s
Dev Containers / validate (.devcontainer/memory-stack/devcontainer.json, memory-stack) (push) Failing after 0s
Dev Containers / validate-worktree (push) Failing after 0s
CI / changes (push) Failing after 4s
Deploy Documentation / validate (push) Has been skipped
Deploy Documentation / deploy (push) Failing after 1s
Init Native E2E / init-native (ubuntu-latest, claude) (push) Failing after 1s
Init Native E2E / init-native (ubuntu-latest, codex) (push) Failing after 1s
Install Native E2E / install-native (ubuntu-latest) (push) Failing after 1s
OpenCode Plugin / typecheck + build + test (push) Failing after 1s
Init Native E2E / init-native (ubuntu-latest, copilot) (push) Failing after 1s
Release Please / release-please (push) Failing after 1s
Wrap E2E / docker-wrap-e2e (push) Failing after 1s
Wrap Native E2E / wrap-native (ubuntu-latest) (push) Failing after 1s
Init E2E / docker-init-e2e (push) Failing after 4s
Merge Conflicts / merge-conflicts (push) Failing after 4s
CI / lint (push) Has been cancelled
CI / build-wheel (push) Has been cancelled
CI / build-wheel-windows (push) Has been cancelled
CI / prefetch-model (push) Has been cancelled
CI / test-dashboard-ui (push) Has been cancelled
CI / test (1) (push) Has been cancelled
CI / test (2) (push) Has been cancelled
CI / test (3) (push) Has been cancelled
CI / test (4) (push) Has been cancelled
CI / test-extras (push) Has been cancelled
CI / test-agno (push) Has been cancelled
CI / build (push) Has been cancelled
CI / workflow-validation (push) Has been cancelled
CI / docker-native-e2e (push) Has been cancelled
CI / windows-native-wrapper (push) Has been cancelled
CI / macos-native-wrapper (push) Has been cancelled
Docker / docker-build (map[name:arm64 platform:linux/arm64 runs_on:ubuntu-24.04-arm], map[bake_target:runtime-code-nonroot name:code-nonroot]) (push) Has been cancelled
Docker / docker-build (map[name:arm64 platform:linux/arm64 runs_on:ubuntu-24.04-arm], map[bake_target:runtime-code-slim name:code-slim]) (push) Has been cancelled
Docker / docker-build (map[name:arm64 platform:linux/arm64 runs_on:ubuntu-24.04-arm], map[bake_target:runtime-code-slim-nonroot name:code-slim-nonroot]) (push) Has been cancelled
Docker / docker-build (map[name:arm64 platform:linux/arm64 runs_on:ubuntu-24.04-arm], map[bake_target:runtime-nonroot name:nonroot]) (push) Has been cancelled
Docker / docker-build (map[name:arm64 platform:linux/arm64 runs_on:ubuntu-24.04-arm], map[bake_target:runtime-slim name:slim]) (push) Has been cancelled
Docker / docker-build (map[name:arm64 platform:linux/arm64 runs_on:ubuntu-24.04-arm], map[bake_target:runtime-slim-nonroot name:slim-nonroot]) (push) Has been cancelled
Docker / docker-manifest (map[bake_target:runtime name:]) (push) Has been cancelled
Docker / docker-manifest (map[bake_target:runtime-code name:code]) (push) Has been cancelled
Docker / docker-manifest (map[bake_target:runtime-code-nonroot name:code-nonroot]) (push) Has been cancelled
Docker / docker-manifest (map[bake_target:runtime-code-slim name:code-slim]) (push) Has been cancelled
Docker / docker-manifest (map[bake_target:runtime-code-slim-nonroot name:code-slim-nonroot]) (push) Has been cancelled
Docker / docker-manifest (map[bake_target:runtime-nonroot name:nonroot]) (push) Has been cancelled
Docker / docker-manifest (map[bake_target:runtime-slim name:slim]) (push) Has been cancelled
Docker / docker-manifest (map[bake_target:runtime-slim-nonroot name:slim-nonroot]) (push) Has been cancelled
Docker / docker-build (map[name:amd64 platform:linux/amd64 runs_on:ubuntu-24.04], map[bake_target:runtime name:]) (push) Has been cancelled
Docker / docker-build (map[name:amd64 platform:linux/amd64 runs_on:ubuntu-24.04], map[bake_target:runtime-code name:code]) (push) Has been cancelled
Docker / docker-build (map[name:amd64 platform:linux/amd64 runs_on:ubuntu-24.04], map[bake_target:runtime-code-nonroot name:code-nonroot]) (push) Has been cancelled
Docker / docker-build (map[name:amd64 platform:linux/amd64 runs_on:ubuntu-24.04], map[bake_target:runtime-code-slim name:code-slim]) (push) Has been cancelled
Docker / docker-build (map[name:amd64 platform:linux/amd64 runs_on:ubuntu-24.04], map[bake_target:runtime-code-slim-nonroot name:code-slim-nonroot]) (push) Has been cancelled
Docker / docker-build (map[name:amd64 platform:linux/amd64 runs_on:ubuntu-24.04], map[bake_target:runtime-nonroot name:nonroot]) (push) Has been cancelled
Docker / docker-build (map[name:amd64 platform:linux/amd64 runs_on:ubuntu-24.04], map[bake_target:runtime-slim name:slim]) (push) Has been cancelled
Docker / docker-build (map[name:amd64 platform:linux/amd64 runs_on:ubuntu-24.04], map[bake_target:runtime-slim-nonroot name:slim-nonroot]) (push) Has been cancelled
Docker / docker-build (map[name:arm64 platform:linux/arm64 runs_on:ubuntu-24.04-arm], map[bake_target:runtime name:]) (push) Has been cancelled
Docker / docker-build (map[name:arm64 platform:linux/arm64 runs_on:ubuntu-24.04-arm], map[bake_target:runtime-code name:code]) (push) Has been cancelled
Docker / promote-latest (push) Has been cancelled
Init Native E2E / init-native (macos-latest, claude) (push) Has been cancelled
Init Native E2E / init-native (macos-latest, codex) (push) Has been cancelled
Init Native E2E / init-native (macos-latest, copilot) (push) Has been cancelled
Install Native E2E / install-native (macos-latest) (push) Has been cancelled
Wrap Native E2E / wrap-native (macos-latest) (push) Has been cancelled
1345 lines
50 KiB
Python
1345 lines
50 KiB
Python
"""Tests for CLI proxy env variable handling and backend validation.
|
|
|
|
Verifies that:
|
|
1. Provider target URL env vars are read by `headroom proxy`
|
|
2. litellm-* backends are accepted by both CLI and argparse paths
|
|
3. HEADROOM_WRAP_PROXY_TIMEOUT controls `headroom wrap` proxy readiness waits
|
|
"""
|
|
|
|
import os
|
|
from unittest.mock import patch
|
|
|
|
import pytest
|
|
|
|
click = pytest.importorskip("click")
|
|
pytest.importorskip("fastapi")
|
|
|
|
from click.testing import CliRunner # noqa: E402
|
|
|
|
from headroom.cli import wrap as wrap_mod # noqa: E402
|
|
from headroom.cli.main import main # noqa: E402
|
|
|
|
|
|
@pytest.fixture
|
|
def runner():
|
|
return CliRunner()
|
|
|
|
|
|
class _FakeProxyProcess:
|
|
returncode = None
|
|
|
|
def __init__(self):
|
|
self.killed = False
|
|
|
|
def poll(self):
|
|
return None
|
|
|
|
def kill(self):
|
|
self.killed = True
|
|
|
|
|
|
class TestCLIWrapProxyTimeout:
|
|
"""Test wrap proxy readiness timeout configuration."""
|
|
|
|
def test_default_timeout_stays_current_without_ml_extras(self, monkeypatch):
|
|
monkeypatch.delenv(wrap_mod._WRAP_PROXY_TIMEOUT_ENV, raising=False)
|
|
monkeypatch.setattr(wrap_mod, "_ml_wrap_extras_detected", lambda: False)
|
|
|
|
assert (
|
|
wrap_mod._resolve_wrap_proxy_timeout_seconds()
|
|
== wrap_mod._WRAP_PROXY_TIMEOUT_DEFAULT_SECONDS
|
|
)
|
|
|
|
def test_default_timeout_is_longer_when_ml_extras_detected(self, monkeypatch):
|
|
monkeypatch.delenv(wrap_mod._WRAP_PROXY_TIMEOUT_ENV, raising=False)
|
|
monkeypatch.setattr(wrap_mod, "_ml_wrap_extras_detected", lambda: True)
|
|
|
|
assert (
|
|
wrap_mod._resolve_wrap_proxy_timeout_seconds()
|
|
== wrap_mod._WRAP_PROXY_TIMEOUT_ML_DEFAULT_SECONDS
|
|
)
|
|
|
|
def test_start_proxy_succeeds_when_ready_within_default_timeout(self, monkeypatch, tmp_path):
|
|
fake_proc = _FakeProxyProcess()
|
|
sleeps = []
|
|
|
|
monkeypatch.delenv(wrap_mod._WRAP_PROXY_TIMEOUT_ENV, raising=False)
|
|
monkeypatch.setattr(wrap_mod, "_ml_wrap_extras_detected", lambda: False)
|
|
monkeypatch.setattr(wrap_mod, "_get_log_path", lambda: tmp_path / "proxy.log")
|
|
monkeypatch.setattr(wrap_mod, "_check_proxy", lambda _port: True)
|
|
monkeypatch.setattr(wrap_mod.time, "sleep", lambda seconds: sleeps.append(seconds))
|
|
monkeypatch.setattr(wrap_mod.subprocess, "Popen", lambda *args, **kwargs: fake_proc)
|
|
|
|
proc = wrap_mod._start_proxy(8787, agent_type="codex")
|
|
|
|
assert proc is fake_proc
|
|
assert sleeps == [1]
|
|
assert fake_proc.killed is False
|
|
|
|
def test_start_proxy_passes_resolved_copilot_api_url_to_proxy(self, monkeypatch, tmp_path):
|
|
fake_proc = _FakeProxyProcess()
|
|
captured: dict[str, object] = {}
|
|
|
|
monkeypatch.delenv(wrap_mod._WRAP_PROXY_TIMEOUT_ENV, raising=False)
|
|
monkeypatch.setattr(wrap_mod, "_ml_wrap_extras_detected", lambda: False)
|
|
monkeypatch.setattr(wrap_mod, "_get_log_path", lambda: tmp_path / "proxy.log")
|
|
monkeypatch.setattr(wrap_mod, "_check_proxy", lambda _port: True)
|
|
monkeypatch.setattr(wrap_mod.time, "sleep", lambda _seconds: None)
|
|
|
|
def fake_popen(*args, **kwargs): # noqa: ANN002, ANN003
|
|
captured["args"] = args
|
|
captured["kwargs"] = kwargs
|
|
return fake_proc
|
|
|
|
monkeypatch.setattr(wrap_mod.subprocess, "Popen", fake_popen)
|
|
|
|
proc = wrap_mod._start_proxy(
|
|
8787,
|
|
agent_type="copilot",
|
|
openai_api_url="https://copilot-api.acme.ghe.com",
|
|
copilot_api_token="copilot-api-token",
|
|
)
|
|
|
|
assert proc is fake_proc
|
|
env = captured["kwargs"]["env"]
|
|
assert env["OPENAI_TARGET_API_URL"] == "https://copilot-api.acme.ghe.com"
|
|
assert env["GITHUB_COPILOT_API_URL"] == "https://copilot-api.acme.ghe.com"
|
|
assert env["GITHUB_COPILOT_API_TOKEN"] == "copilot-api-token"
|
|
|
|
def test_start_proxy_redirects_subprocess_stdio_to_standalone_log(self, monkeypatch, tmp_path):
|
|
fake_proc = _FakeProxyProcess()
|
|
captured: dict[str, object] = {}
|
|
logs: list[str] = []
|
|
|
|
monkeypatch.delenv(wrap_mod._WRAP_PROXY_TIMEOUT_ENV, raising=False)
|
|
monkeypatch.setattr(wrap_mod, "_get_log_path", lambda: tmp_path / "proxy.log")
|
|
monkeypatch.setattr(wrap_mod, "_check_proxy", lambda _port: True)
|
|
monkeypatch.setattr(wrap_mod.time, "sleep", lambda _seconds: None)
|
|
monkeypatch.setattr(wrap_mod, "_ml_wrap_extras_detected", lambda: False)
|
|
monkeypatch.setattr(wrap_mod.click, "echo", lambda message: logs.append(str(message)))
|
|
|
|
def fake_popen(*args, **kwargs): # noqa: ANN002, ANN003
|
|
captured["kwargs"] = kwargs
|
|
return fake_proc
|
|
|
|
monkeypatch.setattr(wrap_mod.subprocess, "Popen", fake_popen)
|
|
|
|
proc = wrap_mod._start_proxy(8787, agent_type="codex")
|
|
|
|
assert proc is fake_proc
|
|
assert captured["kwargs"]["stdout"] is captured["kwargs"]["stderr"]
|
|
assert captured["kwargs"]["stdout"].name == str(tmp_path / "proxy-stdio.log")
|
|
assert captured["kwargs"]["stdout"].name != str(tmp_path / "proxy.log")
|
|
assert f" Logs: {tmp_path / 'proxy.log'}" in logs
|
|
|
|
def test_env_timeout_allows_slow_start_proxy_to_succeed(self, monkeypatch, tmp_path):
|
|
fake_proc = _FakeProxyProcess()
|
|
sleeps = []
|
|
checks = []
|
|
|
|
monkeypatch.setenv(wrap_mod._WRAP_PROXY_TIMEOUT_ENV, "4")
|
|
monkeypatch.setattr(wrap_mod, "_get_log_path", lambda: tmp_path / "proxy.log")
|
|
monkeypatch.setattr(wrap_mod.time, "sleep", lambda seconds: sleeps.append(seconds))
|
|
monkeypatch.setattr(wrap_mod.subprocess, "Popen", lambda *args, **kwargs: fake_proc)
|
|
|
|
def ready_on_fourth_check(port):
|
|
checks.append(port)
|
|
return len(checks) == 4
|
|
|
|
monkeypatch.setattr(wrap_mod, "_check_proxy", ready_on_fourth_check)
|
|
|
|
proc = wrap_mod._start_proxy(8787, agent_type="codex")
|
|
|
|
assert proc is fake_proc
|
|
assert checks == [8787, 8787, 8787, 8787]
|
|
assert sleeps == [1, 1, 1, 1]
|
|
assert fake_proc.killed is False
|
|
|
|
def test_start_proxy_tail_reads_standalone_stdio_log_on_process_exit(
|
|
self, monkeypatch, tmp_path
|
|
):
|
|
fake_proc = _FakeProxyProcess()
|
|
fake_proc.returncode = 1
|
|
fake_proc.poll = lambda: fake_proc.returncode
|
|
|
|
monkeypatch.setenv(wrap_mod._WRAP_PROXY_TIMEOUT_ENV, "2")
|
|
monkeypatch.setattr(wrap_mod, "_get_log_path", lambda: tmp_path / "proxy.log")
|
|
monkeypatch.setattr(wrap_mod, "_check_proxy", lambda _port: False)
|
|
monkeypatch.setattr(wrap_mod.time, "sleep", lambda _seconds: None)
|
|
monkeypatch.setattr(wrap_mod.subprocess, "Popen", lambda *args, **kwargs: fake_proc)
|
|
|
|
(tmp_path / "proxy.log").write_text("canonical runtime log output")
|
|
(tmp_path / "proxy-stdio.log").write_text("proxy stdio startup output")
|
|
|
|
with pytest.raises(RuntimeError) as excinfo:
|
|
wrap_mod._start_proxy(8787, agent_type="codex")
|
|
|
|
message = str(excinfo.value)
|
|
assert "Proxy exited with code 1" in message
|
|
assert "proxy stdio startup output" in message
|
|
assert "canonical runtime log output" not in message
|
|
|
|
def test_timeout_error_names_configured_timeout_and_env_var(self, monkeypatch, tmp_path):
|
|
fake_proc = _FakeProxyProcess()
|
|
|
|
monkeypatch.setenv(wrap_mod._WRAP_PROXY_TIMEOUT_ENV, "2")
|
|
monkeypatch.setattr(wrap_mod, "_get_log_path", lambda: tmp_path / "proxy.log")
|
|
monkeypatch.setattr(wrap_mod, "_check_proxy", lambda _port: False)
|
|
monkeypatch.setattr(wrap_mod.time, "sleep", lambda _seconds: None)
|
|
monkeypatch.setattr(wrap_mod.subprocess, "Popen", lambda *args, **kwargs: fake_proc)
|
|
|
|
with pytest.raises(RuntimeError) as excinfo:
|
|
wrap_mod._start_proxy(8787, agent_type="codex")
|
|
|
|
message = str(excinfo.value)
|
|
assert "within 2 seconds" in message
|
|
assert wrap_mod._WRAP_PROXY_TIMEOUT_ENV in message
|
|
assert fake_proc.killed is True
|
|
|
|
|
|
class TestCLIProxyEnvVars:
|
|
"""Test that the CLI proxy command reads API URL env vars."""
|
|
|
|
def test_headroom_host_from_env(self, runner):
|
|
"""HEADROOM_HOST env var should be passed to ProxyConfig."""
|
|
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_HOST": "0.0.0.0"},
|
|
catch_exceptions=False,
|
|
)
|
|
|
|
assert result.exit_code == 0, result.output
|
|
assert captured_config["config"].host == "0.0.0.0"
|
|
|
|
def test_headroom_port_from_env(self, runner):
|
|
"""HEADROOM_PORT env var should be passed to ProxyConfig."""
|
|
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_PORT": "9797"},
|
|
catch_exceptions=False,
|
|
)
|
|
|
|
assert result.exit_code == 0, result.output
|
|
assert captured_config["config"].port == 9797
|
|
|
|
def test_headroom_min_tokens_from_env(self, runner):
|
|
"""HEADROOM_MIN_TOKENS env var should be passed to ProxyConfig."""
|
|
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_MIN_TOKENS": "120"},
|
|
catch_exceptions=False,
|
|
)
|
|
|
|
assert result.exit_code == 0, result.output
|
|
assert captured_config["config"].min_tokens_to_crush == 120
|
|
|
|
def test_headroom_min_tokens_zero_is_preserved(self, runner):
|
|
"""HEADROOM_MIN_TOKENS=0 is a legitimate value ("crush everything") and
|
|
must not be discarded by an `or 500` fallback (regression)."""
|
|
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_MIN_TOKENS": "0", "HEADROOM_MAX_ITEMS": "0"},
|
|
catch_exceptions=False,
|
|
)
|
|
|
|
assert result.exit_code == 0, result.output
|
|
assert captured_config["config"].min_tokens_to_crush == 0
|
|
assert captured_config["config"].max_items_after_crush == 0
|
|
|
|
def test_headroom_budget_from_env(self, runner):
|
|
"""HEADROOM_BUDGET env var should be passed to ProxyConfig."""
|
|
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_BUDGET": "100.5"},
|
|
catch_exceptions=False,
|
|
)
|
|
|
|
assert result.exit_code == 0, result.output
|
|
assert captured_config["config"].budget_limit_usd == 100.5
|
|
|
|
def test_budget_period_flag_and_env(self, runner):
|
|
"""--budget-period and HEADROOM_BUDGET_PERIOD should reach ProxyConfig."""
|
|
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", "--budget", "50", "--budget-period", "monthly"],
|
|
catch_exceptions=False,
|
|
)
|
|
|
|
assert result.exit_code == 0, result.output
|
|
assert captured_config["config"].budget_period == "monthly"
|
|
|
|
with patch("headroom.proxy.server.run_server", mock_run_server):
|
|
result = runner.invoke(
|
|
main,
|
|
["proxy"],
|
|
env={"HEADROOM_BUDGET_PERIOD": "hourly"},
|
|
catch_exceptions=False,
|
|
)
|
|
|
|
assert result.exit_code == 0, result.output
|
|
assert captured_config["config"].budget_period == "hourly"
|
|
|
|
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, **kwargs):
|
|
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_true(self, runner):
|
|
"""Without HEADROOM_CODE_AWARE_ENABLED, code-aware defaults ON (coding
|
|
posture; consistent with the argparse server path). It degrades to a no-op
|
|
when tree-sitter isn't installed, so defaulting it on is safe."""
|
|
captured_config = {}
|
|
|
|
def mock_run_server(config, **kwargs):
|
|
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 True
|
|
|
|
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, **kwargs):
|
|
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_disable_kompress_from_env(self, runner):
|
|
"""HEADROOM_DISABLE_KOMPRESS should be passed to ProxyConfig."""
|
|
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_DISABLE_KOMPRESS": "1"},
|
|
catch_exceptions=False,
|
|
)
|
|
|
|
assert result.exit_code == 0, result.output
|
|
assert captured_config["config"].disable_kompress is True
|
|
|
|
def test_disable_kompress_from_cli_flag(self, runner):
|
|
"""--disable-kompress should disable Kompress ML compression."""
|
|
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", "--disable-kompress"],
|
|
catch_exceptions=False,
|
|
)
|
|
|
|
assert result.exit_code == 0, result.output
|
|
assert captured_config["config"].disable_kompress 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, **kwargs):
|
|
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 = {}
|
|
|
|
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={"OPENAI_TARGET_API_URL": "http://my-vllm:4000"},
|
|
catch_exceptions=False,
|
|
)
|
|
|
|
assert result.exit_code == 0, result.output
|
|
assert captured_config["config"].openai_api_url == "http://my-vllm:4000"
|
|
|
|
def test_gemini_target_api_url_from_env(self, runner):
|
|
"""GEMINI_TARGET_API_URL env var should be passed to ProxyConfig."""
|
|
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={"GEMINI_TARGET_API_URL": "http://my-gemini:5000"},
|
|
catch_exceptions=False,
|
|
)
|
|
|
|
assert result.exit_code == 0, result.output
|
|
assert captured_config["config"].gemini_api_url == "http://my-gemini:5000"
|
|
|
|
def test_vertex_target_api_url_from_env(self, runner):
|
|
"""VERTEX_TARGET_API_URL env var should be passed to ProxyConfig."""
|
|
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={"VERTEX_TARGET_API_URL": "https://europe-west4-aiplatform.googleapis.com"},
|
|
catch_exceptions=False,
|
|
)
|
|
|
|
assert result.exit_code == 0, result.output
|
|
assert (
|
|
captured_config["config"].vertex_api_url
|
|
== "https://europe-west4-aiplatform.googleapis.com"
|
|
)
|
|
|
|
def test_openai_api_url_cli_flag(self, runner):
|
|
"""--openai-api-url CLI flag should take precedence."""
|
|
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", "--openai-api-url", "http://from-cli:4000"],
|
|
catch_exceptions=False,
|
|
)
|
|
|
|
assert result.exit_code == 0, result.output
|
|
assert captured_config["config"].openai_api_url == "http://from-cli:4000"
|
|
|
|
def test_vertex_api_url_cli_flag(self, runner):
|
|
"""--vertex-api-url CLI flag should take precedence."""
|
|
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", "--vertex-api-url", "https://us-east5-aiplatform.googleapis.com"],
|
|
catch_exceptions=False,
|
|
)
|
|
|
|
assert result.exit_code == 0, result.output
|
|
assert (
|
|
captured_config["config"].vertex_api_url == "https://us-east5-aiplatform.googleapis.com"
|
|
)
|
|
|
|
def test_cli_flag_overrides_env_var(self, runner):
|
|
"""CLI flag should take precedence over env var."""
|
|
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", "--openai-api-url", "http://from-cli:4000"],
|
|
env={"OPENAI_TARGET_API_URL": "http://from-env:4000"},
|
|
catch_exceptions=False,
|
|
)
|
|
|
|
assert result.exit_code == 0, result.output
|
|
assert captured_config["config"].openai_api_url == "http://from-cli:4000"
|
|
|
|
def test_no_env_var_defaults_to_none(self, runner):
|
|
"""Without env var or flag, openai_api_url should be None."""
|
|
captured_config = {}
|
|
|
|
def mock_run_server(config, **kwargs):
|
|
captured_config["config"] = config
|
|
|
|
# Ensure the env var is not set
|
|
env = {k: v for k, v in os.environ.items() if k != "OPENAI_TARGET_API_URL"}
|
|
|
|
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"].openai_api_url is None
|
|
|
|
def test_both_api_urls_from_env(self, runner):
|
|
"""Both OPENAI and GEMINI target URLs can be set via env."""
|
|
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={
|
|
"OPENAI_TARGET_API_URL": "http://my-vllm:4000",
|
|
"GEMINI_TARGET_API_URL": "http://my-gemini:5000",
|
|
},
|
|
catch_exceptions=False,
|
|
)
|
|
|
|
assert result.exit_code == 0, result.output
|
|
assert captured_config["config"].openai_api_url == "http://my-vllm:4000"
|
|
assert captured_config["config"].gemini_api_url == "http://my-gemini:5000"
|
|
|
|
@pytest.mark.parametrize("timeout", [-1, 0, 1, 10000])
|
|
def test_request_timeout_cli_flags(self, runner, timeout):
|
|
"""Fast-fail CLI flags should map into ProxyConfig."""
|
|
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", "--request-timeout-seconds", f"{timeout}"],
|
|
catch_exceptions=False,
|
|
)
|
|
|
|
assert result.exit_code == 0, result.output
|
|
assert (
|
|
captured_config["config"].request_timeout_seconds == timeout
|
|
if timeout and timeout > 0
|
|
else 300
|
|
)
|
|
|
|
@pytest.mark.parametrize("timeout", [-1, 0, 1, 10000])
|
|
def test_request_timeout_from_env(self, runner, timeout):
|
|
"""HEADROOM_REQUEST_TIMEOUT env var should be passed to ProxyConfig."""
|
|
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_REQUEST_TIMEOUT": f"{timeout}"},
|
|
catch_exceptions=False,
|
|
)
|
|
|
|
assert result.exit_code == 0, result.output
|
|
assert (
|
|
captured_config["config"].request_timeout_seconds == timeout
|
|
if timeout and timeout > 0
|
|
else 300
|
|
)
|
|
|
|
def test_retry_and_connect_timeout_cli_flags(self, runner):
|
|
"""Fast-fail CLI flags should map into ProxyConfig."""
|
|
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",
|
|
"--retry-max-attempts",
|
|
"1",
|
|
"--connect-timeout-seconds",
|
|
"3",
|
|
],
|
|
catch_exceptions=False,
|
|
)
|
|
|
|
assert result.exit_code == 0, result.output
|
|
assert captured_config["config"].retry_max_attempts == 1
|
|
assert captured_config["config"].connect_timeout_seconds == 3
|
|
|
|
def test_production_scaling_env_vars(self, runner):
|
|
captured = {}
|
|
|
|
def mock_run_server(config, **kwargs):
|
|
captured["config"] = config
|
|
captured["kwargs"] = kwargs
|
|
|
|
with patch("headroom.proxy.server.run_server", mock_run_server):
|
|
result = runner.invoke(
|
|
main,
|
|
["proxy"],
|
|
env={
|
|
"HEADROOM_WORKERS": "4",
|
|
"HEADROOM_LIMIT_CONCURRENCY": "250",
|
|
"HEADROOM_MAX_CONNECTIONS": "200",
|
|
"HEADROOM_MAX_KEEPALIVE": "50",
|
|
},
|
|
catch_exceptions=False,
|
|
)
|
|
|
|
assert result.exit_code == 0, result.output
|
|
assert captured["config"].max_connections == 200
|
|
assert captured["config"].max_keepalive_connections == 50
|
|
# Click CLI also passes `print_banner=False` to suppress the legacy
|
|
# run_server banner (cli/proxy.py prints its own). Assert the
|
|
# production-scaling keys we care about, not the full kwargs dict.
|
|
assert captured["kwargs"]["workers"] == 4
|
|
assert captured["kwargs"]["limit_concurrency"] == 250
|
|
assert captured["kwargs"].get("print_banner") is False
|
|
|
|
def test_keepalive_expiry_env_var(self, runner):
|
|
captured = {}
|
|
|
|
def mock_run_server(config, **kwargs):
|
|
captured["config"] = config
|
|
captured["kwargs"] = kwargs
|
|
|
|
with patch("headroom.proxy.server.run_server", mock_run_server):
|
|
result = runner.invoke(
|
|
main,
|
|
["proxy"],
|
|
env={"HEADROOM_KEEPALIVE_EXPIRY": "45"},
|
|
catch_exceptions=False,
|
|
)
|
|
|
|
assert result.exit_code == 0, result.output
|
|
assert captured["config"].keepalive_expiry == 45.0
|
|
|
|
def test_production_scaling_cli_flags_override_env_vars(self, runner):
|
|
captured = {}
|
|
|
|
def mock_run_server(config, **kwargs):
|
|
captured["config"] = config
|
|
captured["kwargs"] = kwargs
|
|
|
|
with patch("headroom.proxy.server.run_server", mock_run_server):
|
|
result = runner.invoke(
|
|
main,
|
|
[
|
|
"proxy",
|
|
"--workers",
|
|
"3",
|
|
"--limit-concurrency",
|
|
"125",
|
|
"--max-connections",
|
|
"150",
|
|
"--max-keepalive",
|
|
"25",
|
|
],
|
|
env={
|
|
"HEADROOM_WORKERS": "4",
|
|
"HEADROOM_LIMIT_CONCURRENCY": "250",
|
|
"HEADROOM_MAX_CONNECTIONS": "200",
|
|
"HEADROOM_MAX_KEEPALIVE": "50",
|
|
},
|
|
catch_exceptions=False,
|
|
)
|
|
|
|
assert result.exit_code == 0, result.output
|
|
assert captured["config"].max_connections == 150
|
|
assert captured["config"].max_keepalive_connections == 25
|
|
# Click CLI also passes `print_banner=False`. Assert production
|
|
# scaling keys explicitly rather than the full kwargs dict.
|
|
assert captured["kwargs"]["workers"] == 3
|
|
assert captured["kwargs"]["limit_concurrency"] == 125
|
|
assert captured["kwargs"].get("print_banner") is False
|
|
|
|
|
|
class TestCLIProxyBackend:
|
|
"""Test that litellm-* backends are accepted by the CLI."""
|
|
|
|
def test_litellm_hosted_vllm_backend_accepted(self, runner):
|
|
"""--backend litellm-hosted_vllm should be accepted (not rejected)."""
|
|
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", "--backend", "litellm-hosted_vllm"],
|
|
catch_exceptions=False,
|
|
)
|
|
|
|
assert result.exit_code == 0, result.output
|
|
assert captured_config["config"].backend == "litellm-hosted_vllm"
|
|
|
|
def test_litellm_vertex_backend_accepted(self, runner):
|
|
"""--backend litellm-vertex should be accepted."""
|
|
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", "--backend", "litellm-vertex"],
|
|
catch_exceptions=False,
|
|
)
|
|
|
|
assert result.exit_code == 0, result.output
|
|
assert captured_config["config"].backend == "litellm-vertex"
|
|
|
|
def test_litellm_backend_with_openai_url(self, runner):
|
|
"""Full vLLM setup: litellm backend + OPENAI_TARGET_API_URL."""
|
|
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",
|
|
"--backend",
|
|
"litellm-hosted_vllm",
|
|
"--openai-api-url",
|
|
"http://my-vllm:4000",
|
|
],
|
|
catch_exceptions=False,
|
|
)
|
|
|
|
assert result.exit_code == 0, result.output
|
|
assert captured_config["config"].backend == "litellm-hosted_vllm"
|
|
assert captured_config["config"].openai_api_url == "http://my-vllm:4000"
|
|
|
|
|
|
class TestCLIAnyllmProviderEnv:
|
|
"""Test that HEADROOM_ANYLLM_PROVIDER env var is read by the CLI."""
|
|
|
|
def test_anyllm_provider_from_env(self, runner):
|
|
"""HEADROOM_ANYLLM_PROVIDER env var should override the default."""
|
|
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", "--backend", "anyllm"],
|
|
env={"HEADROOM_ANYLLM_PROVIDER": "llamacpp"},
|
|
catch_exceptions=False,
|
|
)
|
|
|
|
assert result.exit_code == 0, result.output
|
|
assert captured_config["config"].anyllm_provider == "llamacpp"
|
|
|
|
def test_anyllm_provider_cli_flag_works(self, runner):
|
|
"""--anyllm-provider flag should still work."""
|
|
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", "--backend", "anyllm", "--anyllm-provider", "groq"],
|
|
catch_exceptions=False,
|
|
)
|
|
|
|
assert result.exit_code == 0, result.output
|
|
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_flag(self, runner):
|
|
"""--no-ccr disables BOTH the retrieve-tool injection and the markers."""
|
|
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"], 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
|
|
# Unrelated CCR knob stays on.
|
|
assert cfg.ccr_proactive_expansion is True
|
|
|
|
def test_compression_only_all_flags(self, runner):
|
|
"""--no-ccr + --no-ccr-proactive-expansion yields 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",
|
|
"--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_from_env(self, runner):
|
|
"""HEADROOM_NO_CCR env var disables both markers and tool 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": "1"},
|
|
catch_exceptions=False,
|
|
)
|
|
|
|
assert result.exit_code == 0, result.output
|
|
cfg = captured_config["config"]
|
|
assert cfg.ccr_inject_marker is False
|
|
assert cfg.ccr_inject_tool is False
|
|
|
|
|
|
class TestNoCcrMarkerCompressors:
|
|
"""Verify --no-ccr actually suppresses <<ccr:...>> markers
|
|
from every compressor, not just SmartCrusher (#1022)."""
|
|
|
|
def test_content_router_propagates_ccr_inject_marker_false_to_compressors(self):
|
|
"""#1022: ContentRouter must pass enable_ccr=False to compressors
|
|
when ccr_inject_marker=False. Before the fix, only SmartCrusher
|
|
received the flag — Search/Log/Diff compressors always got
|
|
enable_ccr=True (the default)."""
|
|
from headroom.transforms.content_router import (
|
|
ContentRouter,
|
|
ContentRouterConfig,
|
|
)
|
|
|
|
router = ContentRouter(ContentRouterConfig(ccr_inject_marker=False, ccr_enabled=True))
|
|
|
|
# search compressor
|
|
sc = router._get_search_compressor()
|
|
assert sc is not None
|
|
assert sc.config.enable_ccr is False, (
|
|
f"SearchCompressor enable_ccr={sc.config.enable_ccr}, expected False"
|
|
)
|
|
|
|
# log compressor
|
|
lc = router._get_log_compressor()
|
|
assert lc is not None
|
|
assert lc.config.enable_ccr is False, (
|
|
f"LogCompressor enable_ccr={lc.config.enable_ccr}, expected False"
|
|
)
|
|
|
|
# diff compressor
|
|
dc = router._get_diff_compressor()
|
|
assert dc is not None
|
|
assert dc.config.enable_ccr is False, (
|
|
f"DiffCompressor enable_ccr={dc.config.enable_ccr}, expected False"
|
|
)
|
|
|
|
# SmartCrusher already works (regression guard)
|
|
sc2 = router._get_smart_crusher()
|
|
assert sc2 is not None
|
|
# SmartCrusher uses inject_retrieval_marker, not enable_ccr
|
|
|
|
def test_content_router_default_ccr_inject_marker_true(self):
|
|
"""Default config (ccr_inject_marker=True) should give enable_ccr=True."""
|
|
from headroom.transforms.content_router import (
|
|
ContentRouter,
|
|
ContentRouterConfig,
|
|
)
|
|
|
|
router = ContentRouter(ContentRouterConfig())
|
|
sc = router._get_search_compressor()
|
|
assert sc.config.enable_ccr is True
|
|
|
|
lc = router._get_log_compressor()
|
|
assert lc.config.enable_ccr is True
|
|
|
|
dc = router._get_diff_compressor()
|
|
assert dc.config.enable_ccr is True
|
|
|
|
def test_search_compressor_suppresses_markers_with_enable_ccr_false(self):
|
|
"""SearchCompressor with enable_ccr=False must not emit <<ccr: markers."""
|
|
from headroom.transforms.search_compressor import (
|
|
SearchCompressor,
|
|
SearchCompressorConfig,
|
|
)
|
|
|
|
compressor = SearchCompressor(
|
|
SearchCompressorConfig(
|
|
enable_ccr=False,
|
|
min_matches_for_ccr=1,
|
|
context_keywords=["error"],
|
|
)
|
|
)
|
|
content = "\n".join(
|
|
f"src/file{i}.py:{line}: error: something went wrong here"
|
|
for i in range(20)
|
|
for line in range(1, 11)
|
|
)
|
|
result = compressor.compress(content)
|
|
assert "<<ccr:" not in result.compressed, (
|
|
f"SearchCompressor emitted marker when enable_ccr=False: {result.compressed[:300]!r}"
|
|
)
|
|
|
|
def test_log_compressor_suppresses_markers_with_enable_ccr_false(self):
|
|
"""LogCompressor with enable_ccr=False must not emit <<ccr: markers."""
|
|
from headroom.transforms.log_compressor import (
|
|
LogCompressor,
|
|
LogCompressorConfig,
|
|
)
|
|
|
|
npm_lines = ["npm WARN deprecated x"] * 30 + ["npm ERR! something broke"] * 5
|
|
content = "\n".join(npm_lines)
|
|
compressor = LogCompressor(LogCompressorConfig(enable_ccr=False, min_lines_for_ccr=3))
|
|
result = compressor.compress(content)
|
|
assert "<<ccr:" not in result.compressed, (
|
|
f"LogCompressor emitted marker when enable_ccr=False: {result.compressed[:300]!r}"
|
|
)
|
|
|
|
def test_diff_compressor_suppresses_markers_with_enable_ccr_false(self):
|
|
"""DiffCompressor with enable_ccr=False must not emit <<ccr: markers."""
|
|
from headroom.transforms.diff_compressor import (
|
|
DiffCompressor,
|
|
DiffCompressorConfig,
|
|
)
|
|
|
|
compressor = DiffCompressor(DiffCompressorConfig(enable_ccr=False, min_lines_for_ccr=10))
|
|
diff_lines = []
|
|
for i in range(30):
|
|
diff_lines.append(f"diff --git a/src/file{i}.py b/src/file{i}.py")
|
|
diff_lines.append(f"--- a/src/file{i}.py")
|
|
diff_lines.append(f"+++ b/src/file{i}.py")
|
|
for line in range(1, 6):
|
|
diff_lines.append(f"+added line {line} in file {i}")
|
|
diff_lines.append(f"-removed line {line} in file {i}")
|
|
content = "\n".join(diff_lines)
|
|
result = compressor.compress(content)
|
|
assert "<<ccr:" not in result.compressed, (
|
|
f"DiffCompressor emitted marker when enable_ccr=False: {result.compressed[:300]!r}"
|
|
)
|
|
|
|
def test_code_compressor_suppresses_markers_with_enable_ccr_false(self):
|
|
"""CodeAwareCompressor with enable_ccr=False must not emit <<ccr:
|
|
markers when tree-sitter is available (#1022 coverage gap)."""
|
|
from headroom.transforms.code_compressor import (
|
|
CodeAwareCompressor,
|
|
CodeCompressorConfig,
|
|
_check_tree_sitter_available,
|
|
)
|
|
|
|
if not _check_tree_sitter_available():
|
|
pytest.skip("tree-sitter not available in this environment")
|
|
|
|
# Code that would compress with tree-sitter (enough to trigger CCR)
|
|
func_template = (
|
|
"def func_{i}(x: int) -> int:\n"
|
|
' """Docstring for func_{i}."""\n'
|
|
" # Line {j}\n"
|
|
" result = x + {j}\n"
|
|
" result *= 2\n"
|
|
" return result\n"
|
|
)
|
|
content = "\n".join(func_template.format(i=i, j=j) for i in range(30) for j in range(1, 6))
|
|
|
|
compressor = CodeAwareCompressor(
|
|
CodeCompressorConfig(enable_ccr=False, min_tokens_for_compression=1)
|
|
)
|
|
result = compressor.compress(content)
|
|
assert "<<ccr:" not in result.compressed, (
|
|
f"CodeAwareCompressor emitted marker when enable_ccr=False: {result.compressed[:300]!r}"
|
|
)
|
|
|
|
|
|
class TestArgparseBackendValidation:
|
|
"""Test that the argparse path (python -m headroom.proxy.server) accepts litellm-* backends."""
|
|
|
|
def test_argparse_accepts_litellm_backend(self):
|
|
"""The argparse --backend should accept litellm-hosted_vllm (no choices restriction)."""
|
|
import argparse
|
|
|
|
# Recreate the parser matching server.py's main() argparse setup
|
|
# We just need to verify argparse doesn't reject litellm-* values
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--backend", default="anthropic")
|
|
args = parser.parse_args(["--backend", "litellm-hosted_vllm"])
|
|
assert args.backend == "litellm-hosted_vllm"
|
|
|
|
def test_proxy_config_from_env_reads_disable_kompress(self):
|
|
"""The direct server env path should honor HEADROOM_DISABLE_KOMPRESS."""
|
|
from headroom.proxy.server import _proxy_config_from_env
|
|
|
|
with patch.dict(os.environ, {"HEADROOM_DISABLE_KOMPRESS": "1"}):
|
|
config = _proxy_config_from_env()
|
|
|
|
assert config.disable_kompress is True
|
|
|
|
def test_proxy_config_from_env_reads_disable_kompress_fallback(self):
|
|
"""The direct server env path should honor HEADROOM_DISABLE_KOMPRESS_FALLBACK."""
|
|
from headroom.proxy.server import _proxy_config_from_env
|
|
|
|
with patch.dict(os.environ, {"HEADROOM_DISABLE_KOMPRESS_FALLBACK": "1"}):
|
|
config = _proxy_config_from_env()
|
|
|
|
assert config.disable_kompress_fallback is True
|
|
|
|
def test_argparse_registers_keepalive_expiry_flag(self):
|
|
"""The argparse path (python -m headroom.proxy.server) must register
|
|
--keepalive-expiry as a float flag, so it can override the
|
|
HEADROOM_KEEPALIVE_EXPIRY fallback. A bad value makes argparse exit
|
|
before the server boots, which both proves the flag exists and keeps
|
|
the test fast.
|
|
"""
|
|
import subprocess
|
|
import sys
|
|
|
|
result = subprocess.run(
|
|
[sys.executable, "-m", "headroom.proxy.server", "--keepalive-expiry", "notafloat"],
|
|
capture_output=True,
|
|
text=True,
|
|
creationflags=getattr(subprocess, "CREATE_NO_WINDOW", 0),
|
|
)
|
|
|
|
assert result.returncode == 2, result.stderr
|
|
# "invalid float value" only appears if --keepalive-expiry is a registered
|
|
# float arg; a missing flag would instead say "unrecognized arguments".
|
|
assert "--keepalive-expiry" in result.stderr
|
|
assert "invalid float value" in result.stderr
|
|
|
|
|
|
class TestCLIProxyExcludeToolsEnvVar:
|
|
"""HEADROOM_EXCLUDE_TOOLS and HEADROOM_TOOL_PROFILES must reach ProxyConfig via the Click path.
|
|
|
|
Regression coverage for issue #825: the Click entrypoint (headroom/cli/proxy.py)
|
|
previously built ProxyConfig without calling _parse_exclude_tools or
|
|
_parse_tool_profiles, so those env vars were silently ignored for all
|
|
shared/deployed services that launch via `headroom proxy`.
|
|
"""
|
|
|
|
def test_exclude_tools_single_name_from_env(self, runner):
|
|
"""HEADROOM_EXCLUDE_TOOLS=WebSearch propagates to ProxyConfig.exclude_tools."""
|
|
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_EXCLUDE_TOOLS": "WebSearch"},
|
|
catch_exceptions=False,
|
|
)
|
|
|
|
assert result.exit_code == 0, result.output
|
|
cfg = captured_config["config"]
|
|
assert cfg.exclude_tools is not None
|
|
assert "WebSearch" in cfg.exclude_tools
|
|
|
|
def test_exclude_tools_multi_name_from_env(self, runner):
|
|
"""HEADROOM_EXCLUDE_TOOLS=WebSearch,WebFetch yields both names (and lowercased) in result."""
|
|
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_EXCLUDE_TOOLS": "WebSearch,WebFetch"},
|
|
catch_exceptions=False,
|
|
)
|
|
|
|
assert result.exit_code == 0, result.output
|
|
cfg = captured_config["config"]
|
|
assert cfg.exclude_tools is not None
|
|
assert "WebSearch" in cfg.exclude_tools
|
|
assert "WebFetch" in cfg.exclude_tools
|
|
assert "websearch" in cfg.exclude_tools
|
|
assert "webfetch" in cfg.exclude_tools
|
|
|
|
def test_exclude_tools_unset_leaves_none(self, runner):
|
|
"""Without HEADROOM_EXCLUDE_TOOLS, exclude_tools stays None (DEFAULT_EXCLUDE_TOOLS used)."""
|
|
captured_config = {}
|
|
|
|
def mock_run_server(config, **kwargs):
|
|
captured_config["config"] = config
|
|
|
|
env = {k: v for k, v in os.environ.items() if k != "HEADROOM_EXCLUDE_TOOLS"}
|
|
|
|
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"].exclude_tools is None
|
|
|
|
def test_tool_profiles_from_env(self, runner):
|
|
"""HEADROOM_TOOL_PROFILES=Grep:conservative propagates to ProxyConfig.tool_profiles."""
|
|
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_TOOL_PROFILES": "Grep:conservative"},
|
|
catch_exceptions=False,
|
|
)
|
|
|
|
assert result.exit_code == 0, result.output
|
|
cfg = captured_config["config"]
|
|
assert cfg.tool_profiles is not None
|
|
assert "Grep" in cfg.tool_profiles
|
|
|
|
def test_tool_profiles_unset_leaves_none(self, runner):
|
|
"""Without HEADROOM_TOOL_PROFILES, tool_profiles stays None."""
|
|
captured_config = {}
|
|
|
|
def mock_run_server(config, **kwargs):
|
|
captured_config["config"] = config
|
|
|
|
env = {k: v for k, v in os.environ.items() if k != "HEADROOM_TOOL_PROFILES"}
|
|
|
|
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"].tool_profiles is None
|
|
|
|
|
|
class TestCLIProxyRpmTpm:
|
|
"""--rpm/--tpm flags and HEADROOM_RPM/HEADROOM_TPM env vars must reach ProxyConfig."""
|
|
|
|
def test_rpm_default(self, runner):
|
|
"""Without --rpm, rate_limit_requests_per_minute defaults to 60."""
|
|
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
|
|
assert captured_config["config"].rate_limit_requests_per_minute == 60
|
|
|
|
def test_rpm_flag(self, runner):
|
|
"""--rpm 30 should set rate_limit_requests_per_minute to 30."""
|
|
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", "--rpm", "30"], catch_exceptions=False)
|
|
|
|
assert result.exit_code == 0, result.output
|
|
assert captured_config["config"].rate_limit_requests_per_minute == 30
|
|
|
|
def test_rpm_env_var(self, runner):
|
|
"""HEADROOM_RPM=20 should set rate_limit_requests_per_minute to 20."""
|
|
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_RPM": "20"},
|
|
catch_exceptions=False,
|
|
)
|
|
|
|
assert result.exit_code == 0, result.output
|
|
assert captured_config["config"].rate_limit_requests_per_minute == 20
|
|
|
|
def test_tpm_default(self, runner):
|
|
"""Without --tpm, rate_limit_tokens_per_minute defaults to 100000."""
|
|
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
|
|
assert captured_config["config"].rate_limit_tokens_per_minute == 100000
|
|
|
|
def test_tpm_flag(self, runner):
|
|
"""--tpm 50000 should set rate_limit_tokens_per_minute to 50000."""
|
|
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", "--tpm", "50000"], catch_exceptions=False)
|
|
|
|
assert result.exit_code == 0, result.output
|
|
assert captured_config["config"].rate_limit_tokens_per_minute == 50000
|
|
|
|
def test_tpm_env_var(self, runner):
|
|
"""HEADROOM_TPM=80000 should set rate_limit_tokens_per_minute to 80000."""
|
|
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_TPM": "80000"},
|
|
catch_exceptions=False,
|
|
)
|
|
|
|
assert result.exit_code == 0, result.output
|
|
assert captured_config["config"].rate_limit_tokens_per_minute == 80000
|