mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
fix(proxy): make third-party extensions opt-in
Previously, any package registered under the headroom.proxy_extension
entry-point group auto-loaded at proxy startup. A user pip-installing a
plugin (or pulling one in transitively) would get its middleware running
in front of all their LLM traffic with zero opt-in or visibility — the
same mechanism that masked the Shield Enterprise streaming bug.
Change: install_all() now takes an explicit enabled set (or reads
HEADROOM_PROXY_EXTENSIONS). Discovery still runs to enumerate what's
available, but only names the operator opted into actually install.
The literal '*' is a wildcard for trusted environments.
CLI: headroom proxy --proxy-extension shield_enterprise
headroom proxy --proxy-extension shield_enterprise,mypkg
headroom proxy --proxy-extension '*'
Env: HEADROOM_PROXY_EXTENSIONS=shield_enterprise
The startup banner now shows discovered + enabled extensions:
Extensions: discovered=shield_enterprise (opt-in: --proxy-extension ...)
Extensions: ENABLED shield_enterprise (available: shield_enterprise)
Extensions: ENABLED (wildcard) shield_enterprise
Names that were requested but not found are logged as warnings.
Adds proxy_extensions: list[str] | None to ProxyConfig. Plumbs it
through CLI -> ProxyConfig -> install_all(enabled=...).
This is a behavior change for users who relied on auto-loading.
Existing Shield/extension users must add --proxy-extension or set
HEADROOM_PROXY_EXTENSIONS to keep their middleware running.
This commit is contained in:
chopratejas
parent
15877fb63f
commit
f28f697310
4 changed files with 135 additions and 9 deletions
|
|
@ -57,6 +57,17 @@ from .main import main
|
|||
@click.option("--no-optimize", is_flag=True, help="Disable optimization (passthrough mode)")
|
||||
@click.option("--no-cache", is_flag=True, help="Disable semantic caching")
|
||||
@click.option("--no-rate-limit", is_flag=True, help="Disable rate limiting")
|
||||
@click.option(
|
||||
"--proxy-extension",
|
||||
"proxy_extension",
|
||||
multiple=True,
|
||||
envvar="HEADROOM_PROXY_EXTENSIONS",
|
||||
help=(
|
||||
"Enable a registered proxy extension by entry-point name (opt-in). "
|
||||
"Repeat the flag or pass a comma-separated list. Use '*' to enable "
|
||||
"every discovered extension. Env: HEADROOM_PROXY_EXTENSIONS."
|
||||
),
|
||||
)
|
||||
@click.option(
|
||||
"--no-subscription-tracking",
|
||||
is_flag=True,
|
||||
|
|
@ -274,6 +285,7 @@ def proxy(
|
|||
no_optimize: bool,
|
||||
no_cache: bool,
|
||||
no_rate_limit: bool,
|
||||
proxy_extension: tuple[str, ...],
|
||||
no_subscription_tracking: bool,
|
||||
subscription_poll_interval: int | None,
|
||||
retry_max_attempts: int | None,
|
||||
|
|
@ -403,6 +415,13 @@ def proxy(
|
|||
optimize=not no_optimize,
|
||||
cache_enabled=not no_cache,
|
||||
rate_limit_enabled=not no_rate_limit,
|
||||
# Flatten repeat-flag tuple AND any comma-separated values inside it.
|
||||
# `--proxy-extension a,b --proxy-extension c` and `HEADROOM_PROXY_EXTENSIONS=a,b,c`
|
||||
# both yield ["a", "b", "c"]. None when nothing was supplied.
|
||||
proxy_extensions=(
|
||||
[part.strip() for chunk in proxy_extension for part in chunk.split(",") if part.strip()]
|
||||
or None
|
||||
),
|
||||
subscription_tracking_enabled=not no_subscription_tracking,
|
||||
subscription_poll_interval_s=(
|
||||
subscription_poll_interval if subscription_poll_interval is not None else 300
|
||||
|
|
@ -536,6 +555,33 @@ Memory (Multi-Provider):
|
|||
else:
|
||||
telemetry_line = " Telemetry: DISABLED"
|
||||
|
||||
# Discover proxy extensions (third-party packages registered via the
|
||||
# `headroom.proxy_extension` entry-point group). Surfaced in the banner
|
||||
# so operators can see what's available + what's currently opted-in.
|
||||
# Discovery does NOT run extension code; only the explicitly-enabled
|
||||
# set in config.proxy_extensions actually installs.
|
||||
try:
|
||||
from headroom.proxy.extensions import discover as _discover_extensions
|
||||
|
||||
_ext_available = sorted(name for name, _ in _discover_extensions())
|
||||
except Exception: # noqa: BLE001 — banner must never crash startup
|
||||
_ext_available = []
|
||||
_ext_enabled = config.proxy_extensions or []
|
||||
if not _ext_available:
|
||||
extensions_line = " Extensions: (none discovered)"
|
||||
elif not _ext_enabled:
|
||||
extensions_line = (
|
||||
f" Extensions: discovered={','.join(_ext_available)} "
|
||||
f"(opt-in: --proxy-extension <name> or HEADROOM_PROXY_EXTENSIONS=<n>)"
|
||||
)
|
||||
elif "*" in _ext_enabled:
|
||||
extensions_line = f" Extensions: ENABLED (wildcard) {','.join(_ext_available)}"
|
||||
else:
|
||||
extensions_line = (
|
||||
f" Extensions: ENABLED {','.join(sorted(_ext_enabled))} "
|
||||
f"(available: {','.join(_ext_available)})"
|
||||
)
|
||||
|
||||
click.echo(f"""
|
||||
╔═══════════════════════════════════════════════════════════════════════╗
|
||||
║ HEADROOM PROXY ║
|
||||
|
|
@ -551,6 +597,7 @@ Starting proxy server...
|
|||
Rate Limit: {"ENABLED" if config.rate_limit_enabled else "DISABLED"}
|
||||
Memory: {memory_status}
|
||||
License: {license_status}
|
||||
{extensions_line}
|
||||
{stateless_line}{telemetry_line}
|
||||
{backend_section}
|
||||
Routing:
|
||||
|
|
|
|||
|
|
@ -17,6 +17,18 @@ Each ``install`` callable is invoked with the FastAPI ``app`` and the
|
|||
OSS makes no assumptions about what extensions do. The interface is
|
||||
deliberately minimal; extensions own the complexity behind it.
|
||||
|
||||
**Extensions are opt-in.** Discovery enumerates every registered extension,
|
||||
but ``install_all`` only invokes those explicitly enabled by the operator.
|
||||
This protects users from silent behavior changes when a package they didn't
|
||||
audit gets installed in the same environment (e.g., as a transitive dep).
|
||||
|
||||
Enabling extensions:
|
||||
|
||||
* CLI: ``headroom proxy --proxy-extension shield_enterprise,mypkg``
|
||||
* Env: ``HEADROOM_PROXY_EXTENSIONS=shield_enterprise,mypkg``
|
||||
* Wildcard: ``--proxy-extension '*'`` enables every discovered extension
|
||||
(use only when you trust everything in your environment).
|
||||
|
||||
Stability contract: this module is load-bearing for the Enterprise build and
|
||||
any third-party extensions. Changes to the signature of ``install(app, config)``
|
||||
or the entry-point group name require a deprecation cycle.
|
||||
|
|
@ -26,12 +38,14 @@ from __future__ import annotations
|
|||
|
||||
import importlib.metadata
|
||||
import logging
|
||||
from collections.abc import Callable, Iterator
|
||||
import os
|
||||
from collections.abc import Callable, Iterable, Iterator
|
||||
from typing import Any
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
ENTRY_POINT_GROUP = "headroom.proxy_extension"
|
||||
ENV_VAR = "HEADROOM_PROXY_EXTENSIONS"
|
||||
|
||||
ProxyExtension = Callable[[Any, Any], None]
|
||||
"""Signature: ``install(app: FastAPI, config: ProxyConfig) -> None``."""
|
||||
|
|
@ -58,17 +72,73 @@ def discover() -> Iterator[tuple[str, ProxyExtension]]:
|
|||
yield entry.name, install
|
||||
|
||||
|
||||
def install_all(app: Any, config: Any) -> list[str]:
|
||||
"""Run every discovered extension's ``install(app, config)``.
|
||||
def _resolve_enabled(enabled: Iterable[str] | None) -> set[str]:
|
||||
"""Resolve the set of enabled extension names.
|
||||
|
||||
Precedence: explicit ``enabled`` argument > ``HEADROOM_PROXY_EXTENSIONS``
|
||||
env var > empty (no extensions). Empty strings and whitespace are
|
||||
stripped. The literal ``*`` enables all discovered extensions.
|
||||
"""
|
||||
raw: Iterable[str]
|
||||
if enabled is not None:
|
||||
raw = enabled
|
||||
else:
|
||||
raw = (os.environ.get(ENV_VAR) or "").split(",")
|
||||
out: set[str] = set()
|
||||
for n in raw:
|
||||
n = n.strip()
|
||||
if n:
|
||||
out.add(n)
|
||||
return out
|
||||
|
||||
|
||||
def install_all(
|
||||
app: Any,
|
||||
config: Any,
|
||||
enabled: Iterable[str] | None = None,
|
||||
) -> list[str]:
|
||||
"""Run only the explicitly-enabled extensions' ``install(app, config)``.
|
||||
|
||||
Discovery still runs so we can log the universe of available extensions,
|
||||
but only those whose entry-point ``name`` is in ``enabled`` are invoked.
|
||||
The literal ``"*"`` in ``enabled`` is a wildcard that enables every
|
||||
discovered extension.
|
||||
|
||||
Returns the names of successfully installed extensions. If an extension
|
||||
raises inside ``install()``, the exception propagates — this is the
|
||||
documented fail-closed signal (e.g., a Shield Enterprise license check
|
||||
failing should abort startup rather than silently run without protection).
|
||||
documented fail-closed signal (e.g., a license check failing should
|
||||
abort startup rather than silently run without protection).
|
||||
"""
|
||||
enabled_set = _resolve_enabled(enabled)
|
||||
discovered = list(discover())
|
||||
discovered_names = [n for n, _ in discovered]
|
||||
|
||||
if not enabled_set:
|
||||
if discovered_names:
|
||||
log.info(
|
||||
"proxy extensions discovered but disabled (opt-in): %s. "
|
||||
"Enable with --proxy-extension <name> or %s=<name1,name2>.",
|
||||
",".join(discovered_names),
|
||||
ENV_VAR,
|
||||
)
|
||||
return []
|
||||
|
||||
wildcard = "*" in enabled_set
|
||||
installed: list[str] = []
|
||||
for name, install in discover():
|
||||
for name, install in discovered:
|
||||
if not wildcard and name not in enabled_set:
|
||||
continue
|
||||
install(app, config)
|
||||
installed.append(name)
|
||||
log.info("proxy extension installed: %s", name)
|
||||
|
||||
# Warn about names the user asked for that weren't found.
|
||||
if not wildcard:
|
||||
missing = enabled_set - set(discovered_names)
|
||||
if missing:
|
||||
log.warning(
|
||||
"proxy extensions requested but not found: %s (available: %s)",
|
||||
",".join(sorted(missing)),
|
||||
",".join(discovered_names) or "<none>",
|
||||
)
|
||||
return installed
|
||||
|
|
|
|||
|
|
@ -171,6 +171,12 @@ class ProxyConfig:
|
|||
log_file: str | None = None
|
||||
log_full_messages: bool = False
|
||||
|
||||
# Third-party proxy extensions (opt-in only). List of entry-point names
|
||||
# to enable from the `headroom.proxy_extension` group, or `["*"]` for
|
||||
# wildcard. Empty/None means no extensions run, even if installed.
|
||||
# CLI: --proxy-extension <name1,name2>; env: HEADROOM_PROXY_EXTENSIONS.
|
||||
proxy_extensions: list[str] | None = None
|
||||
|
||||
# Fallback
|
||||
fallback_enabled: bool = False
|
||||
fallback_provider: str | None = None
|
||||
|
|
|
|||
|
|
@ -1332,11 +1332,14 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
|
|||
return await call_next(request)
|
||||
|
||||
# Third-party proxy extensions (Enterprise, custom plugins). Discovered via
|
||||
# the `headroom.proxy_extension` entry-point group. An extension that raises
|
||||
# from its install() is a deliberate fail-closed signal and aborts startup.
|
||||
# the `headroom.proxy_extension` entry-point group, but **opt-in only**:
|
||||
# only names listed in config.proxy_extensions (CLI: --proxy-extension,
|
||||
# env: HEADROOM_PROXY_EXTENSIONS) actually get installed. Discovery alone
|
||||
# never runs third-party code. An extension that raises from its install()
|
||||
# is a deliberate fail-closed signal and aborts startup.
|
||||
from headroom.proxy.extensions import install_all as _install_extensions
|
||||
|
||||
_install_extensions(app, config)
|
||||
_install_extensions(app, config, enabled=getattr(config, "proxy_extensions", None))
|
||||
|
||||
# Health & Metrics
|
||||
@app.get("/livez")
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue