From 8aab8f22cbd11061484991262d3fee3268e95bfa Mon Sep 17 00:00:00 2001 From: Rod Boev Date: Wed, 24 Jun 2026 21:58:35 -0400 Subject: [PATCH] fix(cli): wire --rpm/--tpm and HEADROOM_RPM/HEADROOM_TPM to the Click proxy command (#1375) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## 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 --- CHANGELOG.md | 1 + headroom/cli/proxy.py | 18 ++++++++ tests/test_cli_proxy_env.py | 92 +++++++++++++++++++++++++++++++++++++ 3 files changed, 111 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 16e152153..05a74a660 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -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)). diff --git a/headroom/cli/proxy.py b/headroom/cli/proxy.py index c0d757325..336557949 100644 --- a/headroom/cli/proxy.py +++ b/headroom/cli/proxy.py @@ -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, diff --git a/tests/test_cli_proxy_env.py b/tests/test_cli_proxy_env.py index cd5f4200c..d5443f587 100644 --- a/tests/test_cli_proxy_env.py +++ b/tests/test_cli_proxy_env.py @@ -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