headroom/tests/test_runtime_env.py
Tejas Chopra 6904d47a01
feat(proxy): hot-reload live env knobs so a reused proxy picks them up without a restart (#1090)
## Description

A small class of env vars is read by the proxy **live, per request** —
the output-shaper family (`HEADROOM_OUTPUT_SHAPER`,
`HEADROOM_VERBOSITY_LEVEL`, `HEADROOM_EFFORT_ROUTER`,
`HEADROOM_MECHANICAL_EFFORT`, `HEADROOM_VERBOSITY_AUTOTUNE`,
`HEADROOM_OUTPUT_HOLDOUT`), or captured at import
(`HEADROOM_INTERCEPT_READ_MIN_CHARS`). The proxy reads them from its own
process environment, fixed at launch. But `headroom wrap` reuses an
already-running proxy (it restarts only on startup-config drift), so a
value exported *after* the proxy started silently no-op'd — e.g. `export
HEADROOM_OUTPUT_SHAPER=1` had zero effect on a reused proxy on `:8787`.

This PR makes those live knobs **hot-reloadable**: `headroom wrap`
pushes them to the running proxy, which applies them in memory — no
restart (a restart would cold-start the ML stack, drop in-flight
requests, and lose CCR/router caches).

_No linked issue._

## Type of Change

- [x] 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)
- [x] Documentation update
- [ ] Performance improvement
- [ ] Code refactoring (no functional changes)

## Changes Made

- `headroom/proxy/runtime_env.py` (new): single source of truth
registering the live knobs + a thread-safe process-global override
store. `getenv()` (override-then-env) is a drop-in for `os.environ.get`;
behaviour is byte-identical when no override is set.
- Readers rerouted through `runtime_env.getenv`: `output_shaper.py`, the
anthropic holdout read, and the ast-grep threshold (now a live read, not
an import-time constant).
- Proxy: loopback-only `POST /admin/runtime-env` applies overrides in
memory; `/health` → `config.runtime_env` surfaces the live values so
reuse is observable.
- `wrap`: after attaching to a proxy (all call sites), best-effort push
of the session's **explicitly-set** knobs. No-ops if nothing is set,
`--no-proxy`, the proxy is unreachable, or it predates the endpoint
(404). Only explicitly-set knobs are pushed, so a session never clobbers
another with a default it never asked for.
- Docs: README + output-token-reduction guide document the
global-override caveat.

## Testing

- [x] Unit tests pass (`pytest`)
- [x] Linting passes (`ruff check .`)
- [x] Type checking passes (`mypy headroom`)
- [x] New tests added for new functionality
- [ ] Manual testing performed

### Test Output

```text
$ python -m pytest tests/test_runtime_env.py -q
16 passed

$ python -m pytest tests/test_runtime_env.py tests/test_output_shaper.py -q
50 passed

$ ruff check headroom/proxy/runtime_env.py headroom/proxy/output_shaper.py headroom/proxy/handlers/anthropic.py headroom/proxy/interceptors/astgrep.py headroom/proxy/server.py headroom/cli/wrap.py
All checks passed!

$ mypy headroom/proxy/runtime_env.py
Success: no issues found in 1 source file
```

## Real Behavior Proof

- Environment: local macOS, Python 3.12 `.venv`, branch
`fix/runtime-env-hot-reload` at the PR head.
- Exact command / steps: ran the test suites above. The 16 new
`test_runtime_env` tests exercise the registry/store, overrides reaching
the shaper + the ast-grep threshold, the `POST /admin/runtime-env` apply
+ `/health` reflect + loopback-only 404 + 400-on-non-object, and the
wrap push payload / no-op / error-swallow paths.
- Observed result: 50 passed; ruff + mypy clean on the changed modules;
an override set via the endpoint is read by `getenv()` at the shaper and
surfaced in `/health` config.
- Not tested: a literal two-terminal manual session (start a proxy,
`headroom wrap` a second session, `export HEADROOM_OUTPUT_SHAPER=1`,
confirm the reused proxy picks it up). The behaviour is covered by the
endpoint + wrap-push integration tests, but was not exercised by hand
here.

## 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
- [ ] I have updated the CHANGELOG.md if applicable

## Additional Notes

- **Inherent caveat (documented):** overrides are global to the proxy —
one process serves every attached wrapper, so the last explicit setting
wins. No mechanism (restart or hot-reload) can give two sessions on one
shared proxy different output-shaper settings.
- **Scope:** startup-captured settings (`HEADROOM_TARGET_RATIO` etc.)
are intentionally out of scope — a fresh proxy already gets them and
they ride the existing `/health` config channel.
- **Merge blocker:** this branch is currently **CONFLICTING with
`main`** and needs a rebase/merge before it can land.
- CHANGELOG.md left unchanged — releases are managed by release-please
from conventional commits.
2026-06-18 09:50:50 -07:00

255 lines
8.6 KiB
Python

"""Tests for the live runtime-env registry, override store, hot-reload endpoint,
and the wrap-side push that keeps a reused proxy in sync without a restart.
"""
from __future__ import annotations
import pytest
from headroom.proxy import runtime_env as rt
pytest.importorskip("fastapi")
pytest.importorskip("httpx")
from fastapi.testclient import TestClient # noqa: E402
from headroom.proxy.server import ProxyConfig, create_app # noqa: E402
@pytest.fixture(autouse=True)
def _clean_runtime_env(monkeypatch):
"""Each test starts with no overrides and no knob env vars set."""
for knob in rt.RUNTIME_ENV_KNOBS:
monkeypatch.delenv(knob.env, raising=False)
rt.clear_overrides()
yield
rt.clear_overrides()
# ---------------------------------------------------------------------------
# Registry + override store
# ---------------------------------------------------------------------------
def test_getenv_falls_back_to_environment(monkeypatch):
monkeypatch.setenv("HEADROOM_OUTPUT_SHAPER", "1")
assert rt.getenv("HEADROOM_OUTPUT_SHAPER") == "1"
assert rt.getenv("HEADROOM_VERBOSITY_LEVEL", "2") == "2" # unset -> default
assert rt.getenv("HEADROOM_VERBOSITY_LEVEL") is None
def test_getenv_override_wins_over_environment(monkeypatch):
monkeypatch.setenv("HEADROOM_OUTPUT_SHAPER", "0")
rt.set_overrides({"HEADROOM_OUTPUT_SHAPER": "1"})
assert rt.getenv("HEADROOM_OUTPUT_SHAPER") == "1"
def test_set_overrides_ignores_unknown_keys_and_non_strings():
applied = rt.set_overrides(
{
"HEADROOM_OUTPUT_SHAPER": "1",
"NOT_A_KNOB": "x",
"HEADROOM_VERBOSITY_LEVEL": 3, # non-string ignored
}
)
assert applied == {"HEADROOM_OUTPUT_SHAPER": "1"}
assert rt.getenv("NOT_A_KNOB") is None
# The rejected non-string did not become an override.
assert rt.getenv("HEADROOM_VERBOSITY_LEVEL") is None
def test_explicit_env_returns_only_explicitly_set_knobs():
environ = {
"HEADROOM_OUTPUT_SHAPER": "1",
"HEADROOM_MECHANICAL_EFFORT": "low",
"HEADROOM_VERBOSITY_LEVEL": " ", # blank -> not "explicitly set"
"PATH": "/usr/bin", # not a knob
}
assert rt.explicit_env(environ) == {
"HEADROOM_OUTPUT_SHAPER": "1",
"HEADROOM_MECHANICAL_EFFORT": "low",
}
def test_effective_runtime_env_reports_override_or_none(monkeypatch):
monkeypatch.setenv("HEADROOM_EFFORT_ROUTER", "0")
rt.set_overrides({"HEADROOM_OUTPUT_SHAPER": "1"})
eff = rt.effective_runtime_env()
assert eff["HEADROOM_OUTPUT_SHAPER"] == "1" # from override
assert eff["HEADROOM_EFFORT_ROUTER"] == "0" # from env
assert eff["HEADROOM_VERBOSITY_LEVEL"] is None # unset
# Every registered knob is reported.
assert set(eff) == {knob.env for knob in rt.RUNTIME_ENV_KNOBS}
def test_clear_overrides_resets(monkeypatch):
rt.set_overrides({"HEADROOM_OUTPUT_SHAPER": "1"})
rt.clear_overrides()
assert rt.getenv("HEADROOM_OUTPUT_SHAPER") is None
# ---------------------------------------------------------------------------
# Overrides reach the live readers (the whole point)
# ---------------------------------------------------------------------------
def test_override_enables_output_shaper_without_env():
from headroom.proxy.output_shaper import OutputShaperSettings
assert OutputShaperSettings.from_env().enabled is False
rt.set_overrides({"HEADROOM_OUTPUT_SHAPER": "1", "HEADROOM_VERBOSITY_LEVEL": "3"})
settings = OutputShaperSettings.from_env()
assert settings.enabled is True
assert settings.verbosity_level == 3
def test_override_changes_astgrep_threshold_without_env():
from headroom.proxy.interceptors import astgrep
assert astgrep._min_chars_to_rewrite() == 500
rt.set_overrides({"HEADROOM_INTERCEPT_READ_MIN_CHARS": "999"})
assert astgrep._min_chars_to_rewrite() == 999
# Bad value falls back to the documented default rather than raising.
rt.set_overrides({"HEADROOM_INTERCEPT_READ_MIN_CHARS": "not-an-int"})
assert astgrep._min_chars_to_rewrite() == 500
# ---------------------------------------------------------------------------
# /health surface + /admin/runtime-env hot-reload endpoint
# ---------------------------------------------------------------------------
@pytest.fixture
def loopback_client(monkeypatch):
monkeypatch.setenv("HEADROOM_SKIP_UPSTREAM_CHECK", "1")
config = ProxyConfig(
optimize=False,
cache_enabled=False,
rate_limit_enabled=False,
cost_tracking_enabled=False,
)
app = create_app(config)
with TestClient(app, base_url="http://127.0.0.1", client=("127.0.0.1", 12345)) as c:
yield c
def test_health_exposes_runtime_env(loopback_client):
config = loopback_client.get("/health").json()["config"]
assert "runtime_env" in config
assert set(config["runtime_env"]) == {knob.env for knob in rt.RUNTIME_ENV_KNOBS}
assert config["runtime_env"]["HEADROOM_OUTPUT_SHAPER"] is None
def test_admin_runtime_env_applies_and_reflects_in_health(loopback_client):
resp = loopback_client.post(
"/admin/runtime-env",
json={"HEADROOM_OUTPUT_SHAPER": "1", "HEADROOM_VERBOSITY_LEVEL": "3", "BOGUS": "x"},
)
assert resp.status_code == 200
body = resp.json()
assert body["applied"] == {"HEADROOM_OUTPUT_SHAPER": "1", "HEADROOM_VERBOSITY_LEVEL": "3"}
assert body["runtime_env"]["HEADROOM_OUTPUT_SHAPER"] == "1"
# And it is observable on the live /health surface.
health = loopback_client.get("/health").json()["config"]["runtime_env"]
assert health["HEADROOM_OUTPUT_SHAPER"] == "1"
assert health["HEADROOM_VERBOSITY_LEVEL"] == "3"
def test_admin_runtime_env_rejects_non_object(loopback_client):
resp = loopback_client.post("/admin/runtime-env", json=["not", "a", "dict"])
assert resp.status_code == 400
def test_admin_runtime_env_is_loopback_only():
config = ProxyConfig(
optimize=False,
cache_enabled=False,
rate_limit_enabled=False,
cost_tracking_enabled=False,
)
app = create_app(config)
with TestClient(app, base_url="http://127.0.0.1", client=("10.0.0.1", 54321)) as external:
resp = external.post("/admin/runtime-env", json={"HEADROOM_OUTPUT_SHAPER": "1"})
assert resp.status_code == 404 # invisible to non-loopback callers
assert rt.getenv("HEADROOM_OUTPUT_SHAPER") is None # nothing applied
# ---------------------------------------------------------------------------
# wrap-side push
# ---------------------------------------------------------------------------
def test_push_runtime_env_posts_explicit_env(monkeypatch):
import urllib.request
from headroom.cli import wrap
monkeypatch.setenv("HEADROOM_OUTPUT_SHAPER", "1")
monkeypatch.setenv("HEADROOM_VERBOSITY_LEVEL", "3")
captured = {}
class _Resp:
def __enter__(self):
return self
def __exit__(self, *a):
return False
def read(self):
return b"{}"
def fake_urlopen(request, timeout=None):
captured["url"] = request.full_url
captured["body"] = request.data
return _Resp()
monkeypatch.setattr(urllib.request, "urlopen", fake_urlopen)
wrap._push_runtime_env(8787, no_proxy=False)
assert captured["url"] == "http://127.0.0.1:8787/admin/runtime-env"
import json
assert json.loads(captured["body"]) == {
"HEADROOM_OUTPUT_SHAPER": "1",
"HEADROOM_VERBOSITY_LEVEL": "3",
}
def test_push_runtime_env_noop_when_nothing_set(monkeypatch):
import urllib.request
from headroom.cli import wrap
def boom(*a, **k): # must never be called
raise AssertionError("should not POST when nothing is explicitly set")
monkeypatch.setattr(urllib.request, "urlopen", boom)
wrap._push_runtime_env(8787, no_proxy=False) # no env set -> no-op
def test_push_runtime_env_noop_when_no_proxy(monkeypatch):
import urllib.request
from headroom.cli import wrap
monkeypatch.setenv("HEADROOM_OUTPUT_SHAPER", "1")
monkeypatch.setattr(
urllib.request, "urlopen", lambda *a, **k: (_ for _ in ()).throw(AssertionError("no POST"))
)
wrap._push_runtime_env(8787, no_proxy=True) # --no-proxy -> no-op
def test_push_runtime_env_swallows_unreachable_proxy(monkeypatch):
import urllib.request
from headroom.cli import wrap
monkeypatch.setenv("HEADROOM_OUTPUT_SHAPER", "1")
def refused(*a, **k):
raise OSError("connection refused")
monkeypatch.setattr(urllib.request, "urlopen", refused)
# Best-effort: an unreachable / old proxy must not raise.
wrap._push_runtime_env(8787, no_proxy=False)