mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
feat: add HEADROOM_KEEPALIVE_EXPIRY to keep upstream connections warm (#1124)
## Description The Python proxy's `httpx.AsyncClient` (in `server.py`) sets `max_connections` and `max_keepalive_connections` but never `keepalive_expiry`, so httpx's default of **5 seconds** applies. Idle upstream connections are dropped after 5s, and any request after a >5s gap pays a fresh TCP + TLS handshake — costly on high-RTT upstream paths. The **Rust** `crates/headroom-proxy` reqwest client already hardcodes `pool_idle_timeout(Duration::from_secs(90))`; the Python path silently differs at 5s. This PR closes that gap. Closes # ## Type of Change - [ ] Bug fix (non-breaking change that fixes an issue) - [x] New feature (non-breaking change that adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to change) - [ ] Documentation update - [x] Performance improvement - [ ] Code refactoring (no functional changes) ## Changes Made - `ProxyConfig.keepalive_expiry: float = 90.0` (`headroom/proxy/models.py`) - Wired into `httpx.Limits(keepalive_expiry=...)` (`headroom/proxy/server.py`) - `HEADROOM_KEEPALIVE_EXPIRY` env in both env-based config builders (`headroom/proxy/server.py`) - CLI `--keepalive-expiry` (env `HEADROOM_KEEPALIVE_EXPIRY`) following the existing `--max-keepalive` option pattern (`headroom/cli/proxy.py`) - Docs row in `configuration.mdx` + a CLI env test in `tests/test_cli_proxy_env.py` - Default of 90s matches the Rust path; operators can override (e.g. back to `5`). ## Testing - [ ] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check .`) - [ ] Type checking passes (`mypy headroom`) - [x] New tests added for new functionality - [x] Manual testing performed ### Test Output ```text $ ruff check headroom/proxy/models.py headroom/proxy/server.py headroom/cli/proxy.py tests/test_cli_proxy_env.py All checks passed! $ ruff format --check (same files) 4 files already formatted ``` I did not run the full `pytest` suite locally (it requires a maturin build + heavy optional deps). The added test mirrors the existing `test_cli_proxy_env.py` patterns and the CLI option follows the adjacent `--max-keepalive` exactly. ## Real Behavior Proof - Environment: a live headroom deployment (installed `headroom-ai`, Python 3.11) reaching an upstream over a high-RTT tunnel. - Exact command / steps: applied the same field change, restarted the proxy, then inspected the live config. - Observed result: `ProxyConfig.keepalive_expiry == 90.0` at runtime; proxy serves normally; sparse upstream requests no longer re-handshake within the 90s window (the ~300ms cold-handshake penalty that previously recurred after the 5s default expiry is gone). - Not tested: full `pytest`/`mypy` suite locally (maturin build). ## Review Readiness - [x] I have performed a self-review - [x] This PR is ready for human review ## Additional Notes Default changes from httpx's implicit 5s to 90s to reach parity with the Rust `pool_idle_timeout(90s)`; this is the intended behavior alignment rather than a silent regression. CHANGELOG not touched (no entry pattern for proxy knobs observed); happy to add one if preferred. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
parent
487aa71a3c
commit
85786b33a3
5 changed files with 62 additions and 0 deletions
|
|
@ -223,6 +223,7 @@ headroom proxy --learn --min-evidence 3
|
|||
| `HEADROOM_LIMIT_CONCURRENCY` | Maximum concurrent connections before 503 | `1000` |
|
||||
| `HEADROOM_MAX_CONNECTIONS` | Maximum upstream HTTP connections | `500` |
|
||||
| `HEADROOM_MAX_KEEPALIVE` | Maximum upstream keep-alive connections | `100` |
|
||||
| `HEADROOM_KEEPALIVE_EXPIRY` | Seconds an idle upstream keep-alive connection is kept open | `90` |
|
||||
| `HEADROOM_BUDGET` | Daily budget limit in USD | -- |
|
||||
| `HEADROOM_TELEMETRY` | Set to `on` to opt in to anonymous telemetry | `off` (opt-in) |
|
||||
| `HEADROOM_STATELESS` | Set to `true` to disable filesystem writes | `false` |
|
||||
|
|
|
|||
|
|
@ -137,6 +137,14 @@ def _selected_context_tool() -> str:
|
|||
envvar="HEADROOM_MAX_KEEPALIVE",
|
||||
help="Maximum upstream keep-alive connections (default: 100, env: HEADROOM_MAX_KEEPALIVE)",
|
||||
)
|
||||
@click.option(
|
||||
"--keepalive-expiry",
|
||||
"keepalive_expiry",
|
||||
default=90.0,
|
||||
type=click.FloatRange(min=0),
|
||||
envvar="HEADROOM_KEEPALIVE_EXPIRY",
|
||||
help="Seconds an idle upstream keep-alive connection is kept open (default: 90, env: HEADROOM_KEEPALIVE_EXPIRY)",
|
||||
)
|
||||
@click.option(
|
||||
"--mode",
|
||||
default=None,
|
||||
|
|
@ -672,6 +680,7 @@ def proxy(
|
|||
limit_concurrency: int,
|
||||
max_connections: int,
|
||||
max_keepalive_connections: int,
|
||||
keepalive_expiry: float,
|
||||
intercept_tool_results: bool,
|
||||
no_optimize: bool,
|
||||
no_cache: bool,
|
||||
|
|
@ -906,6 +915,7 @@ def proxy(
|
|||
else 10,
|
||||
max_connections=max_connections,
|
||||
max_keepalive_connections=max_keepalive_connections,
|
||||
keepalive_expiry=keepalive_expiry,
|
||||
log_file=None if is_stateless else log_file,
|
||||
log_full_messages=log_messages
|
||||
or os.environ.get("HEADROOM_LOG_MESSAGES", "").lower() in ("true", "1", "yes", "on"),
|
||||
|
|
|
|||
|
|
@ -246,6 +246,7 @@ class ProxyConfig:
|
|||
# Connection pool
|
||||
max_connections: int = 500
|
||||
max_keepalive_connections: int = 100
|
||||
keepalive_expiry: float = 90.0
|
||||
http2: bool = True
|
||||
|
||||
# Memory System
|
||||
|
|
|
|||
|
|
@ -1188,6 +1188,7 @@ class HeadroomProxy(
|
|||
"limits": httpx.Limits(
|
||||
max_connections=self.config.max_connections,
|
||||
max_keepalive_connections=self.config.max_keepalive_connections,
|
||||
keepalive_expiry=self.config.keepalive_expiry,
|
||||
),
|
||||
"verify": _ca_bundle if _ca_bundle is not None else True,
|
||||
}
|
||||
|
|
@ -3773,6 +3774,7 @@ def _proxy_config_from_env() -> ProxyConfig:
|
|||
disable_kompress_openai=_get_env_optional_bool("HEADROOM_DISABLE_KOMPRESS_OPENAI"),
|
||||
max_connections=_get_env_int("HEADROOM_MAX_CONNECTIONS", 500),
|
||||
max_keepalive_connections=_get_env_int("HEADROOM_MAX_KEEPALIVE", 100),
|
||||
keepalive_expiry=_get_env_float("HEADROOM_KEEPALIVE_EXPIRY", 90.0),
|
||||
http2=_get_env_bool("HEADROOM_HTTP2", True),
|
||||
periodic_toin_stats_enabled=_get_env_bool("HEADROOM_PERIODIC_TOIN_STATS", True),
|
||||
mode=normalize_proxy_mode(_get_env_str("HEADROOM_MODE", PROXY_MODE_TOKEN)),
|
||||
|
|
@ -4107,6 +4109,12 @@ if __name__ == "__main__":
|
|||
parser.add_argument(
|
||||
"--max-keepalive", type=int, default=100, help="Max keepalive connections (default: 100)"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--keepalive-expiry",
|
||||
type=float,
|
||||
default=90.0,
|
||||
help="Seconds an idle upstream keep-alive connection is kept open (default: 90)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--no-http2",
|
||||
action="store_true",
|
||||
|
|
@ -4307,6 +4315,7 @@ if __name__ == "__main__":
|
|||
# Connection pool settings
|
||||
max_connections=_get_env_int("HEADROOM_MAX_CONNECTIONS", args.max_connections),
|
||||
max_keepalive_connections=_get_env_int("HEADROOM_MAX_KEEPALIVE", args.max_keepalive),
|
||||
keepalive_expiry=_get_env_float("HEADROOM_KEEPALIVE_EXPIRY", args.keepalive_expiry),
|
||||
http2=not args.no_http2 and _get_env_bool("HEADROOM_HTTP2", True),
|
||||
tool_profiles=tool_profiles if tool_profiles else None,
|
||||
exclude_tools=exclude_tools if exclude_tools else None,
|
||||
|
|
|
|||
|
|
@ -611,6 +611,24 @@ class TestCLIProxyEnvVars:
|
|||
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 = {}
|
||||
|
||||
|
|
@ -859,6 +877,29 @@ class TestArgparseBackendValidation:
|
|||
|
||||
assert config.disable_kompress 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.
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue