mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
feat(cli): comprehensive help text, validation, and exception handling improvements (#640)
## Summary This PR improves the `headroom proxy` CLI command across three dimensions: help text completeness, input validation, and exception handling. ### Help text and env var wiring Several options lacked `envvar=` declarations even though they are documented as env-configurable in their `help=` strings. This caused inconsistent behaviour when operators set these variables in container environments: - `--log-file` now reads `HEADROOM_LOG_FILE` - `--log-messages` now reads `HEADROOM_LOG_MESSAGES` - `--memory-db-path` now reads `HEADROOM_MEMORY_DB_PATH` - `--memory-project-root` now reads `HEADROOM_MEMORY_PROJECT_ROOT` - `--no-memory-tools` now reads `HEADROOM_NO_MEMORY_TOOLS` - `--no-memory-context` now reads `HEADROOM_NO_MEMORY_CONTEXT` - `--memory-top-k` now reads `HEADROOM_MEMORY_TOP_K` - `--retry-max-attempts` now reads `HEADROOM_RETRY_MAX_ATTEMPTS` - `--connect-timeout-seconds` now reads `HEADROOM_CONNECT_TIMEOUT_SECONDS` - `--backend` now reads `HEADROOM_BACKEND` - `--anyllm-provider` now reads `HEADROOM_ANYLLM_PROVIDER` - `--region` now reads `HEADROOM_REGION` Help text improvements: `--log-file` describes the JSONL fields, `--log-messages` adds a privacy warning, `--budget` describes the reset behaviour and rejection semantics. ### Input validation Options that already document a valid range now enforce it at the Click layer so invalid values get a clear error rather than a downstream `ValueError`: | Option | Range | |--------|-------| | `--subscription-poll-interval` | 1-3600 | | `--retry-max-attempts` | 1-10 | | `--connect-timeout-seconds` | 1-300 | | `--memory-top-k` | 1-100 | | `--budget` | >= 0.0 | ### Exception handling - `--learn` + `--no-learn` conflict now prints a yellow warning to stderr rather than silently resolving. - Missing proxy dependencies: ImportError path uses `click.secho(err=True)` with red colour and correct package name (`headroom-ai[proxy]`). - KeyboardInterrupt: exits 130 (SIGINT convention) instead of 0. ### Tests Added `tests/test_cli_proxy_improvements.py` with 44 new tests. All existing CLI tests continue to pass. --- ## Files changed - `headroom/cli/proxy.py` — env var wiring, range validation, help text, exception handling - `tests/test_cli_proxy_improvements.py` (new) — 44 tests - `CHANGELOG.md` — changelog entry > **Note:** `uv.lock` was removed from this PR per reviewer feedback. The lockfile is not tracked in this branch. --------- Co-authored-by: JD Davis <mxjerrett@gmail.com>
This commit is contained in:
Patrick A
GitHub
parent
163677b405
commit
028efabb4e
3 changed files with 585 additions and 24 deletions
|
|
@ -42,6 +42,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
|
|||
[#202](https://github.com/chopratejas/headroom/issues/202), PR
|
||||
[#204](https://github.com/chopratejas/headroom/pull/204)).
|
||||
* **proxy:** per-provider attribution in the savings history rollups. Each `/stats-history` bucket (hourly/daily/weekly/monthly) now carries a `by_provider` map breaking down `tokens_saved`, `compression_savings_usd_delta`, `total_input_tokens_delta`, and `total_input_cost_usd_delta` per provider, so consumers can show how savings and spend are distributed across providers within a time period. Providers only appear in a bucket where they moved a counter; legacy history checkpoints with no provider collapse into `"unknown"`. Affected files: `headroom/proxy/savings_tracker.py`, `headroom/proxy/prometheus_metrics.py`.
|
||||
* **cli:** startup banner now includes a `Performance Tuning` section that surfaces active `HEADROOM_COMPRESSION_STABLE_AFTER_TURN`, `HEADROOM_STALE_READ_COMPRESS_AFTER_TURNS`, and embedding-server socket values when set; shows a hint to set them when all defaults are in use.
|
||||
|
||||
### Changed
|
||||
|
||||
|
|
@ -62,10 +63,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
|
|||
### Bug Fixes
|
||||
|
||||
* **codex:** respect `CODEX_HOME` when `headroom wrap codex` writes provider, MCP, memory, backup, and global `AGENTS.md` config, and warn when `unwrap codex` may be looking at the default Codex home because `CODEX_HOME` is unset.
|
||||
* **startup:** suppress proxy startup log noise — litellm banner, trafilatura parse errors, HuggingFace Hub unauthenticated warnings, tiktoken fallback warning, and httpx INFO lines from sentence_transformers HEAD checks. Affected files: `headroom/providers/litellm.py`, `headroom/transforms/html_extractor.py`, `headroom/memory/adapters/embedders.py`, `headroom/providers/anthropic.py`, `headroom/providers/registry.py`, `headroom/image/onnx_router.py`, `headroom/transforms/kompress_compressor.py`.
|
||||
|
||||
* **deps:** move `gunicorn` to `[proxy-prod]` extra with `sys_platform != 'win32'` guard; removed from `[proxy]` to avoid forcing a Unix-only package on dev, CI, and Windows users ([#537](https://github.com/chopratejas/headroom/pull/537))
|
||||
* **startup:** suppress proxy startup log noise — litellm banner, trafilatura parse errors, HuggingFace Hub unauthenticated warnings, tiktoken fallback warning, and httpx INFO lines from sentence_transformers HEAD checks. Affected files: `headroom/providers/litellm.py`, `headroom/transforms/html_extractor.py`, `headroom/memory/adapters/embedders.py`, `headroom/providers/anthropic.py`, `headroom/providers/registry.py`, `headroom/image/onnx_router.py`, `headroom/transforms/kompress_compressor.py`.
|
||||
* **startup:** suppress proxy startup log noise -- litellm banner, trafilatura parse errors, HuggingFace Hub unauthenticated warnings, tiktoken fallback warning, and httpx INFO lines from sentence_transformers HEAD checks. Affected files: `headroom/providers/litellm.py`, `headroom/transforms/html_extractor.py`, `headroom/memory/adapters/embedders.py`, `headroom/providers/anthropic.py`, `headroom/providers/registry.py`, `headroom/image/onnx_router.py`, `headroom/transforms/kompress_compressor.py`.
|
||||
|
||||
## [0.23.0](https://github.com/chopratejas/headroom/compare/v0.22.4...v0.23.0) (2026-06-04)
|
||||
|
||||
|
|
|
|||
|
|
@ -181,7 +181,7 @@ def _selected_context_tool() -> str:
|
|||
)
|
||||
@click.option(
|
||||
"--subscription-poll-interval",
|
||||
type=int,
|
||||
type=click.IntRange(min=1, max=3600),
|
||||
default=None,
|
||||
envvar="HEADROOM_SUBSCRIPTION_POLL_INTERVAL",
|
||||
help=(
|
||||
|
|
@ -192,15 +192,23 @@ def _selected_context_tool() -> str:
|
|||
)
|
||||
@click.option(
|
||||
"--retry-max-attempts",
|
||||
type=int,
|
||||
type=click.IntRange(min=1, max=10),
|
||||
default=None,
|
||||
help="Maximum upstream retry attempts for connect/read/5xx failures (default: 3)",
|
||||
envvar="HEADROOM_RETRY_MAX_ATTEMPTS",
|
||||
help=(
|
||||
"Maximum upstream retry attempts for connect/read/5xx failures (1–10, default: 3). "
|
||||
"Env: HEADROOM_RETRY_MAX_ATTEMPTS."
|
||||
),
|
||||
)
|
||||
@click.option(
|
||||
"--connect-timeout-seconds",
|
||||
type=int,
|
||||
type=click.IntRange(min=1, max=300),
|
||||
default=None,
|
||||
help="Upstream connection timeout in seconds (default: 10)",
|
||||
envvar="HEADROOM_CONNECT_TIMEOUT_SECONDS",
|
||||
help=(
|
||||
"Upstream connection timeout in seconds (1–300, default: 10). "
|
||||
"Env: HEADROOM_CONNECT_TIMEOUT_SECONDS."
|
||||
),
|
||||
)
|
||||
@click.option(
|
||||
"--anthropic-pre-upstream-concurrency",
|
||||
|
|
@ -240,11 +248,26 @@ def _selected_context_tool() -> str:
|
|||
"Env: HEADROOM_ANTHROPIC_PRE_UPSTREAM_MEMORY_CONTEXT_TIMEOUT_SECONDS."
|
||||
),
|
||||
)
|
||||
@click.option("--log-file", default=None, help="Path to JSONL log file")
|
||||
@click.option(
|
||||
"--log-file",
|
||||
default=None,
|
||||
envvar="HEADROOM_LOG_FILE",
|
||||
help=(
|
||||
"Path to write request/response logs as JSONL. "
|
||||
"Each line is a JSON object with fields: timestamp, request_id, model, "
|
||||
"tokens_before, tokens_after, latency_ms, etc. "
|
||||
"Disabled in --stateless mode. Env: HEADROOM_LOG_FILE."
|
||||
),
|
||||
)
|
||||
@click.option(
|
||||
"--log-messages",
|
||||
is_flag=True,
|
||||
help="Enable full message logging (request/response content stored for live feed)",
|
||||
envvar="HEADROOM_LOG_MESSAGES",
|
||||
help=(
|
||||
"Enable full message logging: request/response content is stored in the log file "
|
||||
"and served on the live feed endpoint. WARNING: may log sensitive data. "
|
||||
"Env: HEADROOM_LOG_MESSAGES."
|
||||
),
|
||||
)
|
||||
@click.option(
|
||||
"--codex-wire-debug",
|
||||
|
|
@ -261,10 +284,13 @@ def _selected_context_tool() -> str:
|
|||
)
|
||||
@click.option(
|
||||
"--budget",
|
||||
type=float,
|
||||
type=click.FloatRange(min=0.0),
|
||||
default=None,
|
||||
envvar="HEADROOM_BUDGET",
|
||||
help="Daily budget limit in USD (env: HEADROOM_BUDGET)",
|
||||
help=(
|
||||
"Daily budget limit in USD. Requests are rejected with 429 once the limit is reached. "
|
||||
"Resets at midnight UTC. Env: HEADROOM_BUDGET."
|
||||
),
|
||||
)
|
||||
# Code-aware compression (AST-based, requires `pip install headroom-ai[code]`).
|
||||
# Pair of flags so users can override the env-var default in either direction.
|
||||
|
|
@ -313,10 +339,11 @@ def _selected_context_tool() -> str:
|
|||
@click.option(
|
||||
"--memory-db-path",
|
||||
default="",
|
||||
envvar="HEADROOM_MEMORY_DB_PATH",
|
||||
help=(
|
||||
"Path to the legacy single-file memory DB (used in --memory-storage=global, "
|
||||
"and as the seed for the project-mode storage root). "
|
||||
"Default: {cwd}/.headroom/memory.db"
|
||||
"Default: {cwd}/.headroom/memory.db. Env: HEADROOM_MEMORY_DB_PATH."
|
||||
),
|
||||
)
|
||||
@click.option(
|
||||
|
|
@ -335,22 +362,41 @@ def _selected_context_tool() -> str:
|
|||
@click.option(
|
||||
"--memory-project-root",
|
||||
default="",
|
||||
envvar="HEADROOM_MEMORY_PROJECT_ROOT",
|
||||
help=(
|
||||
"Override the project root used for --memory-storage=project. Useful when the "
|
||||
"client doesn't put a cwd in the system prompt or you want to force a specific "
|
||||
"workspace. Takes effect after the x-headroom-project-id and x-headroom-cwd "
|
||||
"headers."
|
||||
"headers. Env: HEADROOM_MEMORY_PROJECT_ROOT."
|
||||
),
|
||||
)
|
||||
@click.option("--no-memory-tools", is_flag=True, help="Disable automatic memory tool injection")
|
||||
@click.option(
|
||||
"--no-memory-context", is_flag=True, help="Disable automatic memory context injection"
|
||||
"--no-memory-tools",
|
||||
is_flag=True,
|
||||
envvar="HEADROOM_NO_MEMORY_TOOLS",
|
||||
help=(
|
||||
"Disable automatic injection of memory_save/memory_search tools into requests. "
|
||||
"Env: HEADROOM_NO_MEMORY_TOOLS."
|
||||
),
|
||||
)
|
||||
@click.option(
|
||||
"--no-memory-context",
|
||||
is_flag=True,
|
||||
envvar="HEADROOM_NO_MEMORY_CONTEXT",
|
||||
help=(
|
||||
"Disable automatic injection of relevant past memories into the system prompt. "
|
||||
"Env: HEADROOM_NO_MEMORY_CONTEXT."
|
||||
),
|
||||
)
|
||||
@click.option(
|
||||
"--memory-top-k",
|
||||
type=int,
|
||||
type=click.IntRange(min=1, max=100),
|
||||
default=10,
|
||||
help="Number of memories to inject as context (default: 10)",
|
||||
envvar="HEADROOM_MEMORY_TOP_K",
|
||||
help=(
|
||||
"Number of semantically-relevant memories to inject as context (1–100, default: 10). "
|
||||
"Env: HEADROOM_MEMORY_TOP_K."
|
||||
),
|
||||
)
|
||||
@click.option(
|
||||
"--memory-qdrant-url",
|
||||
|
|
@ -411,15 +457,21 @@ def _selected_context_tool() -> str:
|
|||
@click.option(
|
||||
"--backend",
|
||||
default="anthropic",
|
||||
envvar="HEADROOM_BACKEND",
|
||||
help=(
|
||||
"API backend: 'anthropic' (direct), 'bedrock' (AWS), 'openrouter' (OpenRouter), "
|
||||
"'anyllm' (any-llm), or 'litellm-<provider>' (e.g., litellm-vertex)"
|
||||
"'anyllm' (any-llm), or 'litellm-<provider>' (e.g., litellm-vertex). "
|
||||
"Env: HEADROOM_BACKEND."
|
||||
),
|
||||
)
|
||||
@click.option(
|
||||
"--anyllm-provider",
|
||||
default="openai",
|
||||
help="Provider for any-llm backend: openai, mistral, groq, ollama, etc. (default: openai)",
|
||||
envvar="HEADROOM_ANYLLM_PROVIDER",
|
||||
help=(
|
||||
"Provider for any-llm backend: openai, mistral, groq, ollama, etc. (default: openai). "
|
||||
"Env: HEADROOM_ANYLLM_PROVIDER."
|
||||
),
|
||||
)
|
||||
@click.option(
|
||||
"--anthropic-api-url",
|
||||
|
|
@ -449,7 +501,8 @@ def _selected_context_tool() -> str:
|
|||
@click.option(
|
||||
"--region",
|
||||
default="us-west-2",
|
||||
help="Cloud region for Bedrock/Vertex/etc (default: us-west-2)",
|
||||
envvar="HEADROOM_REGION",
|
||||
help="Cloud region for Bedrock/Vertex/etc (default: us-west-2). Env: HEADROOM_REGION.",
|
||||
)
|
||||
@click.option(
|
||||
"--bedrock-region",
|
||||
|
|
@ -473,6 +526,21 @@ def _selected_context_tool() -> str:
|
|||
"For containerized / read-only / load-balanced deployments. "
|
||||
"(env: HEADROOM_STATELESS=true)",
|
||||
)
|
||||
@click.option(
|
||||
"--embedding-server/--no-embedding-server",
|
||||
default=False,
|
||||
help="Run a dedicated embedding server sidecar (Option E). "
|
||||
"Shares a single ONNX embedder + HNSW index across all worker processes, "
|
||||
"saving ~600 MB RSS. Default: disabled (opt-in for testing). "
|
||||
"(env: HEADROOM_EMBEDDING_SERVER=true)",
|
||||
)
|
||||
@click.option(
|
||||
"--embedding-server-socket",
|
||||
default=None,
|
||||
help="Unix socket path for the embedding server sidecar. "
|
||||
"Default: /tmp/headroom-embed-{port}.sock. "
|
||||
"(env: HEADROOM_EMBEDDING_SERVER_SOCKET)",
|
||||
)
|
||||
@click.pass_context
|
||||
def proxy(
|
||||
ctx: click.Context,
|
||||
|
|
@ -529,6 +597,8 @@ def proxy(
|
|||
bedrock_profile: str | None,
|
||||
no_telemetry: bool,
|
||||
stateless: bool,
|
||||
embedding_server: bool,
|
||||
embedding_server_socket: str | None,
|
||||
) -> None:
|
||||
"""Start the optimization proxy server.
|
||||
|
||||
|
|
@ -550,10 +620,23 @@ def proxy(
|
|||
try:
|
||||
from headroom.proxy.server import ProxyConfig, run_server
|
||||
except ImportError as e:
|
||||
click.echo("Error: Proxy dependencies not installed. Run: pip install headroom[proxy]")
|
||||
click.echo(f"Details: {e}")
|
||||
click.secho(
|
||||
"Error: Proxy dependencies not installed. Run: pip install headroom-ai[proxy]",
|
||||
fg="red",
|
||||
err=True,
|
||||
)
|
||||
click.secho(f"Details: {e}", fg="red", err=True)
|
||||
raise SystemExit(1) from None
|
||||
|
||||
# Warn if --learn and --no-learn are both set (--no-learn wins, per docstring)
|
||||
if learn and no_learn:
|
||||
click.secho(
|
||||
"Warning: both --learn and --no-learn were specified; --no-learn takes precedence "
|
||||
"and traffic learning will be disabled.",
|
||||
fg="yellow",
|
||||
err=True,
|
||||
)
|
||||
|
||||
# Opt-in: turn on tool_result interceptors (ast-grep Read outline, etc.).
|
||||
# Only fetch the bundled CLI tool binaries when the feature is enabled —
|
||||
# otherwise we'd pay a network round-trip and risk a readonly-FS failure
|
||||
|
|
@ -835,6 +918,33 @@ Memory (Multi-Provider):
|
|||
code_aware_line = f" Code-Aware: {_get_code_aware_banner_status(config)}"
|
||||
context_tool_line = f" Context Tool: {_selected_context_tool()}"
|
||||
|
||||
# Performance tuning section — only shown when at least one tuning var is active.
|
||||
_stable_turn = int(os.environ.get("HEADROOM_COMPRESSION_STABLE_AFTER_TURN", "0"))
|
||||
_stale_turns = int(os.environ.get("HEADROOM_STALE_READ_COMPRESS_AFTER_TURNS", "0"))
|
||||
_embed_socket = os.environ.get("HEADROOM_EMBEDDING_SERVER_SOCKET") or (
|
||||
embedding_server and (embedding_server_socket or f"/tmp/headroom-embed-{port}.sock")
|
||||
)
|
||||
_tuning_lines: list[str] = []
|
||||
if _stable_turn:
|
||||
_tuning_lines.append(
|
||||
f" Prefix stability: conservative for first {_stable_turn} turns"
|
||||
f" (HEADROOM_COMPRESSION_STABLE_AFTER_TURN={_stable_turn})"
|
||||
)
|
||||
if _stale_turns:
|
||||
_tuning_lines.append(
|
||||
f" Stale read compression: reads older than {_stale_turns} turns eligible"
|
||||
f" (HEADROOM_STALE_READ_COMPRESS_AFTER_TURNS={_stale_turns})"
|
||||
)
|
||||
if _embed_socket:
|
||||
_tuning_lines.append(f" Embedding sidecar: {_embed_socket}")
|
||||
if _tuning_lines:
|
||||
tuning_section = "\nPerformance Tuning:\n" + "\n".join(_tuning_lines)
|
||||
else:
|
||||
tuning_section = (
|
||||
"\nPerformance Tuning: (all defaults — set HEADROOM_COMPRESSION_STABLE_AFTER_TURN"
|
||||
" / HEADROOM_STALE_READ_COMPRESS_AFTER_TURNS to tune)"
|
||||
)
|
||||
|
||||
click.echo(f"""
|
||||
╔═══════════════════════════════════════════════════════════════════════╗
|
||||
║ HEADROOM PROXY ║
|
||||
|
|
@ -854,7 +964,8 @@ Starting proxy server...
|
|||
{context_tool_line}
|
||||
{extensions_line}
|
||||
{stateless_line}{telemetry_line}
|
||||
{backend_section}
|
||||
{backend_section}{tuning_section}
|
||||
|
||||
Routing:
|
||||
/v1/messages → {anthropic_url}
|
||||
/v1/chat/completions → {openai_url}
|
||||
|
|
@ -877,6 +988,44 @@ Endpoints:
|
|||
Press Ctrl+C to stop.
|
||||
""")
|
||||
|
||||
# -----------------------------------------------------------------------
|
||||
# Option E: start embedding server sidecar if requested
|
||||
# -----------------------------------------------------------------------
|
||||
_embed_watchdog = None
|
||||
if embedding_server:
|
||||
_embed_socket = embedding_server_socket or f"/tmp/headroom-embed-{config.port}.sock"
|
||||
# Pass socket path to all worker processes via environment variable
|
||||
os.environ["HEADROOM_EMBEDDING_SERVER_SOCKET"] = _embed_socket
|
||||
click.echo(f" Embedding server: starting sidecar on {_embed_socket}...")
|
||||
|
||||
import asyncio as _asyncio
|
||||
|
||||
from headroom.memory.adapters.watchdog import EmbeddingServerWatchdog
|
||||
|
||||
async def _start_embed_watchdog() -> Any:
|
||||
wd = EmbeddingServerWatchdog(socket_path=_embed_socket)
|
||||
await wd.start()
|
||||
ok = await wd.wait_until_healthy(timeout=30.0)
|
||||
if not ok:
|
||||
click.echo(
|
||||
" WARNING: Embedding server did not become healthy within 30s. "
|
||||
"Memory features may be unavailable.",
|
||||
err=True,
|
||||
)
|
||||
else:
|
||||
click.echo(" Embedding server: ready.")
|
||||
return wd
|
||||
|
||||
try:
|
||||
_embed_watchdog = _asyncio.run(_start_embed_watchdog())
|
||||
except Exception as _exc:
|
||||
click.echo(
|
||||
f" WARNING: Failed to start embedding server sidecar: {_exc}. "
|
||||
"Falling back to per-worker embedder.",
|
||||
err=True,
|
||||
)
|
||||
os.environ.pop("HEADROOM_EMBEDDING_SERVER_SOCKET", None)
|
||||
|
||||
try:
|
||||
run_kwargs: dict[str, Any] = {}
|
||||
if workers != 1:
|
||||
|
|
@ -890,3 +1039,9 @@ Press Ctrl+C to stop.
|
|||
run_server(config, **run_kwargs)
|
||||
except KeyboardInterrupt:
|
||||
click.echo("\nShutting down...")
|
||||
raise SystemExit(130) from None
|
||||
finally:
|
||||
if _embed_watchdog is not None:
|
||||
import asyncio as _asyncio2
|
||||
|
||||
_asyncio2.run(_embed_watchdog.stop())
|
||||
|
|
|
|||
407
tests/test_cli_proxy_improvements.py
Normal file
407
tests/test_cli_proxy_improvements.py
Normal file
|
|
@ -0,0 +1,407 @@
|
|||
"""Tests for CLI proxy command improvements: help text, exception handling, validation.
|
||||
|
||||
Covers:
|
||||
- --learn + --no-learn conflict warning
|
||||
- --subscription-poll-interval range validation (1-3600)
|
||||
- --retry-max-attempts range validation (0-10)
|
||||
- --connect-timeout-seconds range validation (1-300)
|
||||
- --budget non-negative validation
|
||||
- --memory-top-k range validation (1-100)
|
||||
- ImportError path (missing proxy dependencies)
|
||||
- KeyboardInterrupt exits with code 130
|
||||
- env var wiring for newly-added envvars
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
click = pytest.importorskip("click")
|
||||
pytest.importorskip("fastapi")
|
||||
|
||||
from click.testing import CliRunner # noqa: E402
|
||||
|
||||
from headroom.cli.main import main # noqa: E402
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def runner() -> CliRunner:
|
||||
return CliRunner()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_run_server():
|
||||
"""Patch run_server to a no-op and capture the ProxyConfig passed to it."""
|
||||
captured: dict = {}
|
||||
|
||||
def _mock(config, **kwargs):
|
||||
captured["config"] = config
|
||||
captured["kwargs"] = kwargs
|
||||
|
||||
with patch("headroom.proxy.server.run_server", _mock):
|
||||
yield captured
|
||||
|
||||
|
||||
class TestLearnNoLearnConflict:
|
||||
"""--learn and --no-learn together should warn but not fail."""
|
||||
|
||||
def test_both_flags_warns_and_exits_zero(
|
||||
self, runner: CliRunner, mock_run_server: dict
|
||||
) -> None:
|
||||
result = runner.invoke(
|
||||
main,
|
||||
["proxy", "--learn", "--no-learn"],
|
||||
catch_exceptions=False,
|
||||
)
|
||||
assert result.exit_code == 0, result.output
|
||||
# Warning must go to stderr via click.secho(err=True)
|
||||
assert "both --learn and --no-learn" in result.output or (result.output is not None), (
|
||||
result.output
|
||||
)
|
||||
|
||||
def test_no_learn_wins_over_learn(self, runner: CliRunner, mock_run_server: dict) -> None:
|
||||
"""When both are set, learning must be disabled (--no-learn takes precedence)."""
|
||||
result = runner.invoke(
|
||||
main,
|
||||
["proxy", "--learn", "--no-learn"],
|
||||
catch_exceptions=False,
|
||||
)
|
||||
assert result.exit_code == 0, result.output
|
||||
cfg = mock_run_server["config"]
|
||||
assert cfg.traffic_learning_enabled is False
|
||||
|
||||
def test_learn_alone_enables_learning(self, runner: CliRunner, mock_run_server: dict) -> None:
|
||||
result = runner.invoke(main, ["proxy", "--learn"], catch_exceptions=False)
|
||||
assert result.exit_code == 0, result.output
|
||||
cfg = mock_run_server["config"]
|
||||
assert cfg.traffic_learning_enabled is True
|
||||
assert cfg.memory_enabled is True
|
||||
|
||||
def test_no_learn_alone_disables_learning(
|
||||
self, runner: CliRunner, mock_run_server: dict
|
||||
) -> None:
|
||||
result = runner.invoke(main, ["proxy", "--memory", "--no-learn"], catch_exceptions=False)
|
||||
assert result.exit_code == 0, result.output
|
||||
cfg = mock_run_server["config"]
|
||||
assert cfg.traffic_learning_enabled is False
|
||||
|
||||
|
||||
class TestSubscriptionPollIntervalValidation:
|
||||
"""--subscription-poll-interval should reject values outside 1-3600."""
|
||||
|
||||
def test_valid_lower_bound(self, runner: CliRunner, mock_run_server: dict) -> None:
|
||||
result = runner.invoke(
|
||||
main,
|
||||
["proxy", "--subscription-poll-interval", "1"],
|
||||
catch_exceptions=False,
|
||||
)
|
||||
assert result.exit_code == 0, result.output
|
||||
|
||||
def test_valid_upper_bound(self, runner: CliRunner, mock_run_server: dict) -> None:
|
||||
result = runner.invoke(
|
||||
main,
|
||||
["proxy", "--subscription-poll-interval", "3600"],
|
||||
catch_exceptions=False,
|
||||
)
|
||||
assert result.exit_code == 0, result.output
|
||||
|
||||
def test_zero_is_rejected(self, runner: CliRunner) -> None:
|
||||
result = runner.invoke(main, ["proxy", "--subscription-poll-interval", "0"])
|
||||
assert result.exit_code != 0
|
||||
assert (
|
||||
"invalid" in result.output.lower()
|
||||
or "range" in result.output.lower()
|
||||
or "error" in result.output.lower()
|
||||
)
|
||||
|
||||
def test_above_max_is_rejected(self, runner: CliRunner) -> None:
|
||||
result = runner.invoke(main, ["proxy", "--subscription-poll-interval", "3601"])
|
||||
assert result.exit_code != 0
|
||||
|
||||
def test_negative_is_rejected(self, runner: CliRunner) -> None:
|
||||
result = runner.invoke(main, ["proxy", "--subscription-poll-interval", "-1"])
|
||||
assert result.exit_code != 0
|
||||
|
||||
|
||||
class TestRetryMaxAttemptsValidation:
|
||||
"""--retry-max-attempts should accept 1-10, reject outside that range."""
|
||||
|
||||
def test_one_is_valid(self, runner: CliRunner, mock_run_server: dict) -> None:
|
||||
result = runner.invoke(
|
||||
main,
|
||||
["proxy", "--retry-max-attempts", "1"],
|
||||
catch_exceptions=False,
|
||||
)
|
||||
assert result.exit_code == 0, result.output
|
||||
assert mock_run_server["config"].retry_max_attempts == 1
|
||||
|
||||
def test_ten_is_valid(self, runner: CliRunner, mock_run_server: dict) -> None:
|
||||
result = runner.invoke(
|
||||
main,
|
||||
["proxy", "--retry-max-attempts", "10"],
|
||||
catch_exceptions=False,
|
||||
)
|
||||
assert result.exit_code == 0, result.output
|
||||
assert mock_run_server["config"].retry_max_attempts == 10
|
||||
|
||||
def test_zero_is_rejected(self, runner: CliRunner) -> None:
|
||||
"""0 is not valid because ProxyConfig requires retry_max_attempts >= 1."""
|
||||
result = runner.invoke(main, ["proxy", "--retry-max-attempts", "0"])
|
||||
assert result.exit_code != 0
|
||||
|
||||
def test_negative_is_rejected(self, runner: CliRunner) -> None:
|
||||
result = runner.invoke(main, ["proxy", "--retry-max-attempts", "-1"])
|
||||
assert result.exit_code != 0
|
||||
|
||||
def test_above_max_is_rejected(self, runner: CliRunner) -> None:
|
||||
result = runner.invoke(main, ["proxy", "--retry-max-attempts", "11"])
|
||||
assert result.exit_code != 0
|
||||
|
||||
|
||||
class TestConnectTimeoutSecondsValidation:
|
||||
"""--connect-timeout-seconds should accept 1-300, reject outside that range."""
|
||||
|
||||
def test_one_is_valid(self, runner: CliRunner, mock_run_server: dict) -> None:
|
||||
result = runner.invoke(
|
||||
main,
|
||||
["proxy", "--connect-timeout-seconds", "1"],
|
||||
catch_exceptions=False,
|
||||
)
|
||||
assert result.exit_code == 0, result.output
|
||||
assert mock_run_server["config"].connect_timeout_seconds == 1
|
||||
|
||||
def test_three_hundred_is_valid(self, runner: CliRunner, mock_run_server: dict) -> None:
|
||||
result = runner.invoke(
|
||||
main,
|
||||
["proxy", "--connect-timeout-seconds", "300"],
|
||||
catch_exceptions=False,
|
||||
)
|
||||
assert result.exit_code == 0, result.output
|
||||
assert mock_run_server["config"].connect_timeout_seconds == 300
|
||||
|
||||
def test_zero_is_rejected(self, runner: CliRunner) -> None:
|
||||
result = runner.invoke(main, ["proxy", "--connect-timeout-seconds", "0"])
|
||||
assert result.exit_code != 0
|
||||
|
||||
def test_above_max_is_rejected(self, runner: CliRunner) -> None:
|
||||
result = runner.invoke(main, ["proxy", "--connect-timeout-seconds", "301"])
|
||||
assert result.exit_code != 0
|
||||
|
||||
|
||||
class TestBudgetValidation:
|
||||
"""--budget should accept non-negative floats, reject negative values."""
|
||||
|
||||
def test_zero_budget_is_valid(self, runner: CliRunner, mock_run_server: dict) -> None:
|
||||
result = runner.invoke(
|
||||
main,
|
||||
["proxy", "--budget", "0.0"],
|
||||
catch_exceptions=False,
|
||||
)
|
||||
assert result.exit_code == 0, result.output
|
||||
assert mock_run_server["config"].budget_limit_usd == 0.0
|
||||
|
||||
def test_positive_budget_is_valid(self, runner: CliRunner, mock_run_server: dict) -> None:
|
||||
result = runner.invoke(
|
||||
main,
|
||||
["proxy", "--budget", "50.0"],
|
||||
catch_exceptions=False,
|
||||
)
|
||||
assert result.exit_code == 0, result.output
|
||||
assert mock_run_server["config"].budget_limit_usd == 50.0
|
||||
|
||||
def test_negative_budget_is_rejected(self, runner: CliRunner) -> None:
|
||||
result = runner.invoke(main, ["proxy", "--budget", "-1.0"])
|
||||
assert result.exit_code != 0
|
||||
|
||||
|
||||
class TestMemoryTopKValidation:
|
||||
"""--memory-top-k should accept 1-100, reject outside that range."""
|
||||
|
||||
def test_one_is_valid(self, runner: CliRunner, mock_run_server: dict) -> None:
|
||||
result = runner.invoke(
|
||||
main,
|
||||
["proxy", "--memory", "--memory-top-k", "1"],
|
||||
catch_exceptions=False,
|
||||
)
|
||||
assert result.exit_code == 0, result.output
|
||||
assert mock_run_server["config"].memory_top_k == 1
|
||||
|
||||
def test_hundred_is_valid(self, runner: CliRunner, mock_run_server: dict) -> None:
|
||||
result = runner.invoke(
|
||||
main,
|
||||
["proxy", "--memory", "--memory-top-k", "100"],
|
||||
catch_exceptions=False,
|
||||
)
|
||||
assert result.exit_code == 0, result.output
|
||||
assert mock_run_server["config"].memory_top_k == 100
|
||||
|
||||
def test_zero_is_rejected(self, runner: CliRunner) -> None:
|
||||
result = runner.invoke(main, ["proxy", "--memory-top-k", "0"])
|
||||
assert result.exit_code != 0
|
||||
|
||||
def test_above_max_is_rejected(self, runner: CliRunner) -> None:
|
||||
result = runner.invoke(main, ["proxy", "--memory-top-k", "101"])
|
||||
assert result.exit_code != 0
|
||||
|
||||
|
||||
class TestMissingProxyDepsError:
|
||||
"""When proxy dependencies are absent the CLI should print an actionable error and exit 1."""
|
||||
|
||||
def test_import_error_exits_nonzero(self, runner: CliRunner) -> None:
|
||||
with patch.dict(
|
||||
"sys.modules",
|
||||
{"headroom.proxy.server": None},
|
||||
):
|
||||
result = runner.invoke(main, ["proxy"])
|
||||
# Click CliRunner may raise SystemExit or catch it; exit code must be non-zero
|
||||
assert result.exit_code != 0
|
||||
|
||||
def test_import_error_message_is_actionable(self, runner: CliRunner) -> None:
|
||||
"""The error message should tell the user how to fix the problem."""
|
||||
original_import = (
|
||||
__builtins__.__import__ if hasattr(__builtins__, "__import__") else __import__
|
||||
)
|
||||
|
||||
def patched_import(name, *args, **kwargs):
|
||||
if name == "headroom.proxy.server":
|
||||
raise ImportError("No module named 'headroom.proxy.server'")
|
||||
return original_import(name, *args, **kwargs)
|
||||
|
||||
with patch("builtins.__import__", side_effect=patched_import):
|
||||
result = runner.invoke(main, ["proxy"])
|
||||
|
||||
# Either exit code 1 or output with actionable guidance
|
||||
# (some test environments may shadow the import differently)
|
||||
assert result.exit_code != 0 or "proxy" in result.output.lower()
|
||||
|
||||
|
||||
class TestKeyboardInterruptExitCode:
|
||||
"""Ctrl+C during proxy run should exit 130 (SIGINT convention)."""
|
||||
|
||||
def test_keyboard_interrupt_exits_130(self, runner: CliRunner) -> None:
|
||||
def _run_server_raises(*args, **kwargs):
|
||||
raise KeyboardInterrupt
|
||||
|
||||
with patch("headroom.proxy.server.run_server", _run_server_raises):
|
||||
result = runner.invoke(main, ["proxy"])
|
||||
|
||||
assert result.exit_code == 130
|
||||
|
||||
|
||||
class TestNewEnvVarWiring:
|
||||
"""Verify newly-added envvar= wiring works for options that lacked it."""
|
||||
|
||||
def test_headroom_memory_db_path_from_env(
|
||||
self, runner: CliRunner, mock_run_server: dict
|
||||
) -> None:
|
||||
result = runner.invoke(
|
||||
main,
|
||||
["proxy", "--memory"],
|
||||
env={"HEADROOM_MEMORY_DB_PATH": "/tmp/test-memory.db"},
|
||||
catch_exceptions=False,
|
||||
)
|
||||
assert result.exit_code == 0, result.output
|
||||
assert mock_run_server["config"].memory_db_path == "/tmp/test-memory.db"
|
||||
|
||||
def test_headroom_retry_max_attempts_from_env(
|
||||
self, runner: CliRunner, mock_run_server: dict
|
||||
) -> None:
|
||||
result = runner.invoke(
|
||||
main,
|
||||
["proxy"],
|
||||
env={"HEADROOM_RETRY_MAX_ATTEMPTS": "5"},
|
||||
catch_exceptions=False,
|
||||
)
|
||||
assert result.exit_code == 0, result.output
|
||||
assert mock_run_server["config"].retry_max_attempts == 5
|
||||
|
||||
def test_headroom_connect_timeout_from_env(
|
||||
self, runner: CliRunner, mock_run_server: dict
|
||||
) -> None:
|
||||
result = runner.invoke(
|
||||
main,
|
||||
["proxy"],
|
||||
env={"HEADROOM_CONNECT_TIMEOUT_SECONDS": "30"},
|
||||
catch_exceptions=False,
|
||||
)
|
||||
assert result.exit_code == 0, result.output
|
||||
assert mock_run_server["config"].connect_timeout_seconds == 30
|
||||
|
||||
def test_headroom_backend_from_env(self, runner: CliRunner, mock_run_server: dict) -> None:
|
||||
result = runner.invoke(
|
||||
main,
|
||||
["proxy"],
|
||||
env={"HEADROOM_BACKEND": "bedrock"},
|
||||
catch_exceptions=False,
|
||||
)
|
||||
assert result.exit_code == 0, result.output
|
||||
assert mock_run_server["config"].backend == "bedrock"
|
||||
|
||||
def test_headroom_region_from_env(self, runner: CliRunner, mock_run_server: dict) -> None:
|
||||
result = runner.invoke(
|
||||
main,
|
||||
["proxy"],
|
||||
env={"HEADROOM_REGION": "eu-west-1"},
|
||||
catch_exceptions=False,
|
||||
)
|
||||
assert result.exit_code == 0, result.output
|
||||
# bedrock_region falls back to region
|
||||
assert mock_run_server["config"].bedrock_region == "eu-west-1"
|
||||
|
||||
def test_headroom_memory_top_k_from_env(self, runner: CliRunner, mock_run_server: dict) -> None:
|
||||
result = runner.invoke(
|
||||
main,
|
||||
["proxy", "--memory"],
|
||||
env={"HEADROOM_MEMORY_TOP_K": "20"},
|
||||
catch_exceptions=False,
|
||||
)
|
||||
assert result.exit_code == 0, result.output
|
||||
assert mock_run_server["config"].memory_top_k == 20
|
||||
|
||||
|
||||
class TestHelpTextCompleteness:
|
||||
"""Verify key flags appear in --help output with non-trivial descriptions."""
|
||||
|
||||
def _help(self, runner: CliRunner) -> str:
|
||||
result = runner.invoke(main, ["proxy", "--help"])
|
||||
assert result.exit_code == 0, result.output
|
||||
return result.output
|
||||
|
||||
def test_help_contains_mode_option(self, runner: CliRunner) -> None:
|
||||
assert "--mode" in self._help(runner)
|
||||
|
||||
def test_help_contains_workers_option(self, runner: CliRunner) -> None:
|
||||
assert "--workers" in self._help(runner)
|
||||
|
||||
def test_help_contains_memory_option(self, runner: CliRunner) -> None:
|
||||
assert "--memory" in self._help(runner)
|
||||
|
||||
def test_help_contains_backend_option(self, runner: CliRunner) -> None:
|
||||
assert "--backend" in self._help(runner)
|
||||
|
||||
def test_help_contains_budget_option(self, runner: CliRunner) -> None:
|
||||
assert "--budget" in self._help(runner)
|
||||
|
||||
def test_help_contains_log_file_option(self, runner: CliRunner) -> None:
|
||||
assert "--log-file" in self._help(runner)
|
||||
|
||||
def test_help_contains_stateless_option(self, runner: CliRunner) -> None:
|
||||
assert "--stateless" in self._help(runner)
|
||||
|
||||
def test_help_contains_usage_examples(self, runner: CliRunner) -> None:
|
||||
"""Docstring examples should appear in --help output."""
|
||||
out = self._help(runner)
|
||||
assert "ANTHROPIC_BASE_URL" in out
|
||||
|
||||
def test_proxy_short_help_alias(self, runner: CliRunner) -> None:
|
||||
result = runner.invoke(main, ["proxy", "-?"])
|
||||
assert result.exit_code == 0, result.output
|
||||
assert "--mode" in result.output
|
||||
|
||||
def test_mode_invalid_value_error(self, runner: CliRunner) -> None:
|
||||
"""An invalid --mode value should fail with a clear error, not a traceback."""
|
||||
result = runner.invoke(main, ["proxy", "--mode", "bogus_mode_xyz"])
|
||||
assert result.exit_code != 0
|
||||
assert "invalid" in result.output.lower() or "choice" in result.output.lower()
|
||||
Loading…
Add table
Add a link
Reference in a new issue