mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
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.
This commit is contained in:
Tejas Chopra
GitHub
parent
26be2c39cb
commit
6904d47a01
9 changed files with 527 additions and 15 deletions
|
|
@ -146,6 +146,15 @@ export HEADROOM_OUTPUT_SHAPER=1 # off by default
|
|||
headroom proxy --port 8787
|
||||
```
|
||||
|
||||
> **Already running a proxy?** These switches are read *live* on every request,
|
||||
> so a proxy that `headroom wrap` **reused** (rather than started) would not see
|
||||
> a value you export afterwards — its environment was snapshotted at launch.
|
||||
> `headroom wrap` now hot-syncs your current settings to the running proxy via a
|
||||
> loopback `POST /admin/runtime-env`, so they take effect immediately with **no
|
||||
> restart** (no cold start, no dropped requests, no lost caches). Set them before
|
||||
> you `wrap`. On a shared proxy these overrides are global — the last explicit
|
||||
> setting wins.
|
||||
|
||||
**Learn the right terseness for you.** People don't *say* how terse they want
|
||||
answers — they *show* it (they interrupt long replies, or move on before they
|
||||
could have read them). `headroom learn --verbosity` reads your past sessions and
|
||||
|
|
|
|||
|
|
@ -21,6 +21,15 @@ export HEADROOM_OUTPUT_SHAPER=1 # off by default
|
|||
headroom proxy --port 8787
|
||||
```
|
||||
|
||||
> **If a proxy is already running** (e.g. `headroom wrap claude` attaches to one
|
||||
> on port 8787 instead of starting a fresh one), it reads this switch from the
|
||||
> environment it was launched with — so exporting it afterwards wouldn't reach
|
||||
> it. `headroom wrap` handles this for you: it hot-syncs your current output
|
||||
> settings to the running proxy (loopback `POST /admin/runtime-env`), applied
|
||||
> immediately with no restart. Set the variables before you run `wrap`. Because
|
||||
> one proxy is shared by every session attached to it, these settings are global
|
||||
> — the most recent explicit value wins.
|
||||
|
||||
That's it. Two things now happen on every request:
|
||||
|
||||
1. **Verbosity steering** — a short "be terse, don't restate context" instruction
|
||||
|
|
|
|||
|
|
@ -1535,6 +1535,7 @@ def _run_proxy_only_watcher(
|
|||
proxy_holder[0] = _ensure_proxy(
|
||||
port, no_proxy, learn=learn, memory=memory, agent_type=agent_type
|
||||
)
|
||||
_push_runtime_env(port, no_proxy)
|
||||
click.echo()
|
||||
print_setup_lines()
|
||||
click.echo()
|
||||
|
|
@ -2175,6 +2176,44 @@ def _should_use_copilot_oauth(
|
|||
return has_oauth_auth()
|
||||
|
||||
|
||||
def _push_runtime_env(port: int, no_proxy: bool) -> None:
|
||||
"""Hot-sync this session's live env knobs to the proxy on ``port``.
|
||||
|
||||
Live knobs (the output-shaper family, the ast-grep read threshold) are read
|
||||
from the *proxy's* process environment. A proxy we reused — rather than
|
||||
started — would otherwise ignore values exported in this shell, since its
|
||||
environment was snapshotted when it first launched. Pushing them to
|
||||
``/admin/runtime-env`` applies them in memory with no disruptive restart.
|
||||
|
||||
Best-effort: a silent no-op when nothing is explicitly set, when there is no
|
||||
proxy (``--no-proxy``), when the proxy is unreachable, or when it predates
|
||||
the endpoint (older build returns 404).
|
||||
"""
|
||||
if no_proxy:
|
||||
return
|
||||
from headroom.proxy import runtime_env as _rt
|
||||
|
||||
payload = _rt.explicit_env(os.environ)
|
||||
if not payload:
|
||||
return
|
||||
|
||||
import urllib.error
|
||||
import urllib.request
|
||||
|
||||
request = urllib.request.Request(
|
||||
f"http://127.0.0.1:{port}/admin/runtime-env",
|
||||
data=json.dumps(payload).encode("utf-8"),
|
||||
method="POST",
|
||||
headers={"Content-Type": "application/json"},
|
||||
)
|
||||
try:
|
||||
with urllib.request.urlopen(request, timeout=2) as response:
|
||||
response.read()
|
||||
except (OSError, urllib.error.URLError, ValueError):
|
||||
return
|
||||
click.echo(f" Synced output settings to proxy: {', '.join(sorted(payload))}")
|
||||
|
||||
|
||||
def _ensure_proxy(
|
||||
port: int,
|
||||
no_proxy: bool,
|
||||
|
|
@ -2600,6 +2639,7 @@ def _launch_tool(
|
|||
openai_api_url=openai_api_url,
|
||||
copilot_api_token=copilot_api_token,
|
||||
)
|
||||
_push_runtime_env(port, no_proxy)
|
||||
|
||||
if code_graph:
|
||||
_setup_code_graph(verbose=False)
|
||||
|
|
@ -3047,6 +3087,7 @@ def claude(
|
|||
region=region,
|
||||
anthropic_api_url=foundry_upstream,
|
||||
)
|
||||
_push_runtime_env(port, no_proxy)
|
||||
|
||||
if not no_rtk:
|
||||
if _selected_context_tool() == _CONTEXT_TOOL_LEAN_CTX:
|
||||
|
|
|
|||
|
|
@ -1722,11 +1722,11 @@ class AnthropicHandlerMixin:
|
|||
# conversation is treatment or control. This keeps the A/B
|
||||
# comparison clean AND keeps the prefix cache stable (we
|
||||
# never flip a conversation's system-prompt tail mid-stream).
|
||||
import os as _os
|
||||
from headroom.proxy import runtime_env
|
||||
|
||||
_holdout = 0.0
|
||||
try:
|
||||
_holdout = float(_os.environ.get("HEADROOM_OUTPUT_HOLDOUT", "0") or "0")
|
||||
_holdout = float(runtime_env.getenv("HEADROOM_OUTPUT_HOLDOUT", "0") or "0")
|
||||
except ValueError:
|
||||
_holdout = 0.0
|
||||
_arm = assign_arm(conversation_key_from_body(body), _holdout)
|
||||
|
|
|
|||
|
|
@ -20,17 +20,25 @@ from pathlib import Path
|
|||
from typing import Any
|
||||
|
||||
from headroom import binaries
|
||||
from headroom.proxy import runtime_env
|
||||
|
||||
from . import base
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# Latency floor: below this size, the subprocess cost of running ast-grep
|
||||
# isn't worth the tiny win. It is NOT a semantic threshold — the framework
|
||||
# rejects any rewrite that doesn't actually shrink tokens, so we don't need
|
||||
# a "big enough to matter" check here, only a "big enough to justify the
|
||||
# fork()" check.
|
||||
MIN_CHARS_TO_REWRITE = int(os.environ.get("HEADROOM_INTERCEPT_READ_MIN_CHARS", "500"))
|
||||
# fork()" check. Read live (not as a module constant) so a hot-reload or a
|
||||
# reused proxy re-synced by ``headroom wrap`` takes effect without a restart.
|
||||
def _min_chars_to_rewrite() -> int:
|
||||
try:
|
||||
return int(runtime_env.getenv("HEADROOM_INTERCEPT_READ_MIN_CHARS", "500"))
|
||||
except (TypeError, ValueError):
|
||||
return 500
|
||||
|
||||
|
||||
# Tool_input keys that indicate the model targeted a specific line range;
|
||||
# outlining would frustrate that intent and likely cause a re-read.
|
||||
|
|
@ -93,7 +101,7 @@ class AstGrepReadOutline:
|
|||
) -> bool:
|
||||
if tool_name not in ("Read", "read_file", "view", "cat"):
|
||||
return False
|
||||
if len(tool_output) < MIN_CHARS_TO_REWRITE:
|
||||
if len(tool_output) < _min_chars_to_rewrite():
|
||||
return False
|
||||
# Respect explicit line ranges — the model wants those specific lines.
|
||||
if any(k in tool_input for k in _RANGE_KEYS):
|
||||
|
|
|
|||
|
|
@ -34,11 +34,12 @@ flags) — no content regexes or keyword patterns.
|
|||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
from dataclasses import dataclass
|
||||
from enum import Enum
|
||||
from typing import Any
|
||||
|
||||
from headroom.proxy import runtime_env
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Documented Anthropic API minimum for thinking.budget_tokens on models
|
||||
|
|
@ -106,22 +107,22 @@ class OutputShaperSettings:
|
|||
|
||||
@classmethod
|
||||
def from_env(cls) -> OutputShaperSettings:
|
||||
enabled = os.environ.get("HEADROOM_OUTPUT_SHAPER", "").lower() in (
|
||||
enabled = runtime_env.getenv("HEADROOM_OUTPUT_SHAPER", "").lower() in (
|
||||
"1",
|
||||
"true",
|
||||
"yes",
|
||||
)
|
||||
try:
|
||||
level = int(os.environ.get("HEADROOM_VERBOSITY_LEVEL", "2"))
|
||||
level = int(runtime_env.getenv("HEADROOM_VERBOSITY_LEVEL", "2"))
|
||||
except ValueError:
|
||||
level = 2
|
||||
level = max(0, min(4, level))
|
||||
router = os.environ.get("HEADROOM_EFFORT_ROUTER", "1").lower() not in (
|
||||
router = runtime_env.getenv("HEADROOM_EFFORT_ROUTER", "1").lower() not in (
|
||||
"0",
|
||||
"false",
|
||||
"no",
|
||||
)
|
||||
mech = os.environ.get("HEADROOM_MECHANICAL_EFFORT", "low")
|
||||
mech = runtime_env.getenv("HEADROOM_MECHANICAL_EFFORT", "low")
|
||||
if mech not in _EFFORT_RANK:
|
||||
mech = "low"
|
||||
return cls(
|
||||
|
|
@ -144,9 +145,7 @@ def resolve_verbosity_level(settings: OutputShaperSettings) -> tuple[int, str]:
|
|||
Returns ``(level, source)``. Kept separate from :func:`shape_request` so the
|
||||
body-mutating core stays a pure function of an explicit level.
|
||||
"""
|
||||
import os
|
||||
|
||||
if os.environ.get("HEADROOM_VERBOSITY_LEVEL"):
|
||||
if runtime_env.getenv("HEADROOM_VERBOSITY_LEVEL"):
|
||||
return settings.verbosity_level, "env"
|
||||
|
||||
try:
|
||||
|
|
@ -156,7 +155,7 @@ def resolve_verbosity_level(settings: OutputShaperSettings) -> tuple[int, str]:
|
|||
except Exception:
|
||||
return settings.verbosity_level, "default"
|
||||
|
||||
autotune = os.environ.get("HEADROOM_VERBOSITY_AUTOTUNE", "").lower() in ("1", "true", "yes")
|
||||
autotune = runtime_env.getenv("HEADROOM_VERBOSITY_AUTOTUNE", "").lower() in ("1", "true", "yes")
|
||||
if autotune:
|
||||
ctrl_path = ws / "verbosity_controller.json"
|
||||
if ctrl_path.exists():
|
||||
|
|
|
|||
151
headroom/proxy/runtime_env.py
Normal file
151
headroom/proxy/runtime_env.py
Normal file
|
|
@ -0,0 +1,151 @@
|
|||
"""Live (per-request) env knobs and a hot-reload override store.
|
||||
|
||||
Most Headroom settings are read once at proxy startup into ``Config`` and are
|
||||
visible in ``/health``. A second, smaller class of environment variables is
|
||||
read *live* — on every request (the output-shaper family) or captured at module
|
||||
import (the ast-grep read-rewrite threshold). The proxy reads these from its own
|
||||
process environment, so a *reused* proxy — one ``headroom wrap`` attaches to
|
||||
rather than starting fresh — never sees values a user exports afterwards. The
|
||||
fix without this module would be to restart the proxy, which is disruptive
|
||||
(cold-start of the ML stack, dropped in-flight requests, lost compression
|
||||
caches).
|
||||
|
||||
This module is the single source of truth for that class of knob and provides a
|
||||
process-global override store. ``headroom wrap`` pushes the values it would
|
||||
otherwise only be able to apply by restarting (``POST /admin/runtime-env``), and
|
||||
the proxy applies them in memory with no restart. Readers call :func:`getenv`
|
||||
instead of ``os.environ.get`` so an override wins over the launch-time
|
||||
environment; with no override set, behaviour is byte-for-byte identical to
|
||||
reading the environment directly.
|
||||
|
||||
Scope rule: a variable belongs here only if the proxy reads it *after* startup
|
||||
(or captures it at import) AND it is not already reflected in the ``/health``
|
||||
``config`` block that ``wrap`` compares for reuse. Startup-captured settings
|
||||
(``HEADROOM_TARGET_RATIO`` etc.) do not belong here — a fresh proxy already
|
||||
gets them and they ride the existing config channel.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import threading
|
||||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass
|
||||
from typing import overload
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Knob:
|
||||
"""One reuse-invalidating live env var.
|
||||
|
||||
``env`` is the environment variable name (also the key used in ``/health``
|
||||
and in the hot-reload payload). ``kind`` is advisory metadata for the
|
||||
``/health`` surface and validation — readers still parse the raw string
|
||||
exactly as they did when reading the environment directly, so a knob's
|
||||
parsing/clamping semantics live with its reader, not here.
|
||||
"""
|
||||
|
||||
env: str
|
||||
kind: str # "bool" | "int" | "float" | "str"
|
||||
summary: str
|
||||
|
||||
|
||||
# The registry. Adding a knob here is all it takes to make a live env var
|
||||
# hot-reloadable and visible in /health. Keep this list to genuinely live knobs
|
||||
# (see the scope rule in the module docstring).
|
||||
RUNTIME_ENV_KNOBS: tuple[Knob, ...] = (
|
||||
Knob("HEADROOM_OUTPUT_SHAPER", "bool", "Master switch for output-token shaping."),
|
||||
Knob(
|
||||
"HEADROOM_VERBOSITY_LEVEL", "int", "Verbosity steering level 0-4 (unset = learned/default)."
|
||||
),
|
||||
Knob("HEADROOM_EFFORT_ROUTER", "bool", "Lower effort on mechanical tool-result continuations."),
|
||||
Knob("HEADROOM_MECHANICAL_EFFORT", "str", "Effort value used on mechanical continuations."),
|
||||
Knob("HEADROOM_VERBOSITY_AUTOTUNE", "bool", "Use the AIMD verbosity controller state."),
|
||||
Knob(
|
||||
"HEADROOM_OUTPUT_HOLDOUT",
|
||||
"float",
|
||||
"Fraction of conversations held out for A/B measurement.",
|
||||
),
|
||||
Knob(
|
||||
"HEADROOM_INTERCEPT_READ_MIN_CHARS",
|
||||
"int",
|
||||
"Min tool-output chars before the ast-grep read rewrite.",
|
||||
),
|
||||
)
|
||||
|
||||
_KNOBS_BY_ENV: dict[str, Knob] = {k.env: k for k in RUNTIME_ENV_KNOBS}
|
||||
|
||||
# Process-global override store. Writes take the lock; reads are a plain
|
||||
# ``dict.get`` (atomic in CPython) so the per-request hot path stays lock-free.
|
||||
_lock = threading.Lock()
|
||||
_overrides: dict[str, str] = {}
|
||||
|
||||
|
||||
@overload
|
||||
def getenv(name: str, default: str) -> str: ...
|
||||
|
||||
|
||||
@overload
|
||||
def getenv(name: str, default: None = ...) -> str | None: ...
|
||||
|
||||
|
||||
def getenv(name: str, default: str | None = None) -> str | None:
|
||||
"""Return the live value for ``name``: hot-reload override, else environment.
|
||||
|
||||
Drop-in for ``os.environ.get`` at the reader site. When no override has been
|
||||
pushed, this is exactly ``os.environ.get(name, default)``. Overloaded like
|
||||
``os.environ.get`` so a string default yields ``str`` — callers can ``.lower()``
|
||||
or ``int(...)`` the result without a None-check.
|
||||
"""
|
||||
override = _overrides.get(name)
|
||||
if override is not None:
|
||||
return override
|
||||
return os.environ.get(name, default)
|
||||
|
||||
|
||||
def set_overrides(values: dict[str, object]) -> dict[str, str]:
|
||||
"""Apply hot-reload overrides for known knobs. Returns what was applied.
|
||||
|
||||
Unknown keys and non-string values are ignored (the endpoint is loopback-only
|
||||
but we still never trust the body blindly). Storing the raw string preserves
|
||||
each reader's own parsing/clamping semantics.
|
||||
"""
|
||||
applied: dict[str, str] = {}
|
||||
with _lock:
|
||||
for key, value in values.items():
|
||||
if key not in _KNOBS_BY_ENV or not isinstance(value, str):
|
||||
continue
|
||||
_overrides[key] = value
|
||||
applied[key] = value
|
||||
return applied
|
||||
|
||||
|
||||
def clear_overrides() -> None:
|
||||
"""Drop all overrides (used by tests and to reset state)."""
|
||||
with _lock:
|
||||
_overrides.clear()
|
||||
|
||||
|
||||
def explicit_env(environ: Mapping[str, str] | None = None) -> dict[str, str]:
|
||||
"""Knobs *explicitly* set (non-empty) in ``environ`` — the wrap push payload.
|
||||
|
||||
Only explicitly-set knobs are pushed so a session never clobbers another
|
||||
session's setting with a default it never asked for. To force a knob back to
|
||||
a default on a shared proxy, set it explicitly (e.g. ``HEADROOM_OUTPUT_SHAPER=0``).
|
||||
"""
|
||||
src = os.environ if environ is None else environ
|
||||
out: dict[str, str] = {}
|
||||
for knob in RUNTIME_ENV_KNOBS:
|
||||
raw = src.get(knob.env)
|
||||
if raw is not None and raw.strip() != "":
|
||||
out[knob.env] = raw
|
||||
return out
|
||||
|
||||
|
||||
def effective_runtime_env() -> dict[str, str | None]:
|
||||
"""The live value of every knob (override-or-environment) for ``/health``.
|
||||
|
||||
``None`` means the knob is unset, so the reader will fall back to its own
|
||||
default. This is what the proxy will actually use on the next request.
|
||||
"""
|
||||
return {knob.env: getenv(knob.env) for knob in RUNTIME_ENV_KNOBS}
|
||||
|
|
@ -103,6 +103,7 @@ from headroom.providers.registry import (
|
|||
format_backend_status,
|
||||
resolve_api_targets,
|
||||
)
|
||||
from headroom.proxy import runtime_env
|
||||
from headroom.proxy.auth_mode import should_stamp_codex_client
|
||||
|
||||
# =============================================================================
|
||||
|
|
@ -2057,6 +2058,10 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
|
|||
),
|
||||
"force_kompress": bool(profile_kwargs.get("force_kompress", False)),
|
||||
"accuracy_guard": config.accuracy_guard,
|
||||
# Live (per-request) env knobs the proxy reads after startup.
|
||||
# Surfaced so `headroom wrap` can see what a reused proxy is
|
||||
# actually using and hot-sync it via /admin/runtime-env.
|
||||
"runtime_env": runtime_env.effective_runtime_env(),
|
||||
"pid": os.getpid(),
|
||||
}
|
||||
return payload
|
||||
|
|
@ -2329,6 +2334,39 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
|
|||
payload["runtime"] = _runtime_payload()
|
||||
return JSONResponse(status_code=200, content=payload)
|
||||
|
||||
@app.post("/admin/runtime-env", dependencies=[Depends(_require_loopback)])
|
||||
async def admin_runtime_env(request: Request):
|
||||
"""Hot-reload live env knobs (the output-shaper family, the ast-grep
|
||||
read threshold) without restarting the proxy.
|
||||
|
||||
Live knobs are read from the proxy's *process* environment, so a proxy
|
||||
that ``headroom wrap`` reused — rather than started — never sees values
|
||||
a user exported afterwards. Instead of a disruptive restart (cold ML
|
||||
load, dropped requests, lost caches), ``wrap`` POSTs the values here and
|
||||
the proxy applies them in memory, effective on the next request.
|
||||
|
||||
Loopback-only. The body is a flat ``{ENV_NAME: "value"}`` map; unknown
|
||||
keys and non-string values are ignored. Returns what was applied plus
|
||||
the resulting live config. Last writer wins (overrides are global to the
|
||||
proxy, which is inherent — every wrapper shares one process).
|
||||
"""
|
||||
try:
|
||||
body = await request.json()
|
||||
except (ValueError, UnicodeDecodeError):
|
||||
body = None
|
||||
if not isinstance(body, dict):
|
||||
return JSONResponse(
|
||||
status_code=400,
|
||||
content={"error": "expected a JSON object of {ENV_NAME: value}"},
|
||||
)
|
||||
applied = runtime_env.set_overrides(body)
|
||||
if applied:
|
||||
logger.info("runtime-env hot-reload applied: %s", sorted(applied))
|
||||
return JSONResponse(
|
||||
status_code=200,
|
||||
content={"applied": applied, "runtime_env": runtime_env.effective_runtime_env()},
|
||||
)
|
||||
|
||||
@app.get("/dashboard", response_class=HTMLResponse)
|
||||
async def dashboard():
|
||||
"""Serve the Headroom dashboard UI."""
|
||||
|
|
@ -4109,10 +4147,12 @@ if __name__ == "__main__":
|
|||
compress_user_messages=args.compress_user_messages
|
||||
or _get_env_bool("HEADROOM_COMPRESS_USER_MESSAGES", False),
|
||||
savings_profile=os.environ.get("HEADROOM_SAVINGS_PROFILE") or None,
|
||||
# Default 0.4 keep-ratio so the Kompress text (prose/code) path compresses
|
||||
# meaningfully out of the box; HEADROOM_TARGET_RATIO overrides.
|
||||
target_ratio=(
|
||||
float(os.environ["HEADROOM_TARGET_RATIO"])
|
||||
if os.environ.get("HEADROOM_TARGET_RATIO")
|
||||
else None
|
||||
else 0.4
|
||||
),
|
||||
compress_system_messages=(
|
||||
_get_env_bool("HEADROOM_COMPRESS_SYSTEM_MESSAGES", False)
|
||||
|
|
|
|||
255
tests/test_runtime_env.py
Normal file
255
tests/test_runtime_env.py
Normal file
|
|
@ -0,0 +1,255 @@
|
|||
"""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)
|
||||
Loading…
Add table
Add a link
Reference in a new issue