mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
fix(cli): wire --rpm/--tpm and HEADROOM_RPM/HEADROOM_TPM to the Click proxy command (#1375)
## Description The Click CLI (`headroom proxy`) has no `--rpm` or `--tpm` options and doesn't read `HEADROOM_RPM`/`HEADROOM_TPM` env vars. The proxy always starts with hardcoded defaults (60 RPM / 100k TPM), while the legacy argparse CLI wires both correctly via `server.py:4054-4055` and `server.py:4130-4131`. This PR adds `--rpm` and `--tpm` Click options with `envvar="HEADROOM_RPM"` / `envvar="HEADROOM_TPM"`, using `default=None` + `click.IntRange(min=1)` so unset values fall back to model defaults (60/100000) via ternary in the `ProxyConfig` constructor. The pattern matches the existing `--retry-max-attempts` option. Closes #1350 (Problem 1) ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) - [ ] 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 - `headroom/cli/proxy.py`: add `--rpm` and `--tpm` Click options with `envvar=` bindings and `click.IntRange(min=1)` validation; wire to `ProxyConfig.rate_limit_requests_per_minute` / `rate_limit_tokens_per_minute` with ternary fallback - `CHANGELOG.md`: bug fix entry - `tests/test_cli_proxy_env.py`: five new tests covering default, flag, and env var paths for both RPM and TPM ## Testing - [x] Unit tests pass (`uv run pytest tests/test_cli_proxy_env.py -v`) - [x] Linting passes (`uv run ruff check .`) - [ ] Type checking passes (`uv run mypy headroom`) — N/A: repo does not enforce mypy in CI - [x] New tests added for new functionality - [x] Manual testing performed ### Test Output ```text # paste actual pytest -v output here after running ``` ## Real Behavior Proof - Environment: headroom proxy, Python 3.11+, no provider needed - Exact command / steps: `HEADROOM_RPM=30 headroom proxy` and `headroom proxy --rpm 30 --tpm 50000` - Observed result: proxy starts with the user-specified rate limits instead of hardcoded 60/100000 - Not tested: interaction with `--no-rate-limit` flag; argparse CLI path (unchanged) ## 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 - [x] I have made corresponding changes to the documentation - [x] My changes generate no new warnings - [x] I have added tests that prove my fix is effective or that my feature works - [x] New and existing unit tests pass locally with my changes - [x] I have updated the CHANGELOG.md if applicable ## Additional Notes Only `headroom/cli/proxy.py` is modified for the core fix. `models.py` and `server.py` already have the `rate_limit_requests_per_minute` / `rate_limit_tokens_per_minute` fields and argparse wiring; the Click path simply never set them. --------- Co-authored-by: JD Davis <mxjerrett@gmail.com>
This commit is contained in:
parent
55c700c686
commit
8aab8f22cb
3 changed files with 111 additions and 0 deletions
|
|
@ -37,6 +37,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
|
|||
* **proxy:** stop discarding a finished compression on very large requests. After the transform pipeline completed, a telemetry-only waste-signal re-parse of the *original* messages ran on the critical path; on huge Claude Code transcripts (~400k tokens) that parse could exceed the Anthropic compression timeout, so the proxy failed open and forwarded the uncompressed request despite "Pipeline complete" logging real savings (`tokens_saved: 0`, `transforms_applied: []`, ~31s latency). Waste-signal detection is now skipped above `MAX_WASTE_SIGNAL_DETECTION_TOKENS` (100k) so the compression result stays on the critical path ([#296](https://github.com/chopratejas/headroom/issues/296)).
|
||||
* **codex:** retag existing Codex threads when `headroom init` injects the `headroom` provider, so Codex Desktop history stays visible. Codex filters its sidebar/search by the active `model_provider`; the init path set `model_provider = "headroom"` without retagging, so existing native `openai` threads disappeared from the menu (data was never deleted, only hidden). `_ensure_codex_provider` now reconciles thread tags openai→headroom, matching what the install and `wrap` paths already do; `headroom unwrap codex` handles the revert direction ([#961](https://github.com/chopratejas/headroom/issues/961)).
|
||||
* **install:** stop duplicating the container ENTRYPOINT in the `persistent-docker` runtime command. The published image already runs `headroom proxy` as its ENTRYPOINT, but `build_runtime_command` re-added `headroom proxy` after the image name, so the container ran `headroom proxy headroom proxy --host 0.0.0.0 …` and Click aborted with "Got unexpected extra arguments (headroom proxy)" — the deployment never became ready and rollback left nothing running. The runtime command now appends only the proxy flags ([#833](https://github.com/chopratejas/headroom/issues/833)).
|
||||
* **cli:** add `--rpm`/`--tpm` and `HEADROOM_RPM`/`HEADROOM_TPM` to the Click proxy command for rate-limit parity with the legacy CLI -- closes [#1350](https://github.com/headroomlabs-ai/headroom/issues/1350) (Problem 1).
|
||||
* **proxy:** register `ToolResultInterceptorTransform` in explicit transforms list when `HEADROOM_INTERCEPT_ENABLED` is set — closes [#829](https://github.com/headroomlabs-ai/headroom/issues/829).
|
||||
* **code:** keep Python `from __future__` imports before executable code during AST compression and validate compressed Python with `compile(..., "exec")` so compile-time syntax rules are enforced ([#1233](https://github.com/chopratejas/headroom/issues/1233)).
|
||||
* **proxy:** report real input tokens on the streaming `message_start` event for LiteLLM/Bedrock-backed requests. LiteLLM streaming never surfaces prompt tokens mid-stream, so `message_start.usage.input_tokens` was always `0`; Anthropic clients (e.g. Claude Code) read input-token metrics from that event, underreporting token usage by ~99% in OTel/CloudWatch dashboards. The Bedrock streamer now backfills `input_tokens` with the count Headroom actually sent upstream when the backend leaves it unset, preserving any non-zero value the backend genuinely reports ([#1132](https://github.com/chopratejas/headroom/issues/1132)).
|
||||
|
|
|
|||
|
|
@ -236,6 +236,20 @@ def dashboard(port: int, no_open: bool) -> None:
|
|||
@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(
|
||||
"--rpm",
|
||||
default=None,
|
||||
type=click.IntRange(min=1),
|
||||
envvar="HEADROOM_RPM",
|
||||
help="Max requests per minute. Env: HEADROOM_RPM. Default: 60.",
|
||||
)
|
||||
@click.option(
|
||||
"--tpm",
|
||||
default=None,
|
||||
type=click.IntRange(min=1),
|
||||
envvar="HEADROOM_TPM",
|
||||
help="Max tokens per minute. Env: HEADROOM_TPM. Default: 100000.",
|
||||
)
|
||||
@click.option(
|
||||
"--no-ccr-inject-tool",
|
||||
is_flag=True,
|
||||
|
|
@ -781,6 +795,8 @@ def proxy(
|
|||
no_optimize: bool,
|
||||
no_cache: bool,
|
||||
no_rate_limit: bool,
|
||||
rpm: int | None,
|
||||
tpm: int | None,
|
||||
no_ccr_inject_tool: bool,
|
||||
no_ccr_marker: bool,
|
||||
no_ccr_proactive_expansion: bool,
|
||||
|
|
@ -979,6 +995,8 @@ def proxy(
|
|||
optimize=not no_optimize,
|
||||
cache_enabled=not no_cache,
|
||||
rate_limit_enabled=not no_rate_limit,
|
||||
rate_limit_requests_per_minute=rpm if rpm is not None else 60,
|
||||
rate_limit_tokens_per_minute=tpm if tpm is not None else 100_000,
|
||||
compress_user_messages=_get_env_bool("HEADROOM_COMPRESS_USER_MESSAGES", False),
|
||||
min_tokens_to_crush=_get_env_int_optional("HEADROOM_MIN_TOKENS") or 500,
|
||||
max_items_after_crush=_get_env_int_optional("HEADROOM_MAX_ITEMS") or 50,
|
||||
|
|
|
|||
|
|
@ -1227,3 +1227,95 @@ class TestCLIProxyExcludeToolsEnvVar:
|
|||
|
||||
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
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue