From 5bbc4193886127b1077b910615563cf9861d0ccd Mon Sep 17 00:00:00 2001 From: jmlago Date: Sat, 4 Jul 2026 01:10:07 +0100 Subject: [PATCH] feat: external control-plane protocol (key resolve, usage API, tenant BYO keys) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds an optional "bring your own control plane" integration, entirely off unless CONTROL_PLANE_URL + CONTROL_PLANE_INTERNAL_SECRET are set: - control_plane_client.py (new leaf module): resolve_key with positive/ negative caching, single-flight, and stale-grace served only while the control plane is unreachable; tenant_env (allowlisted, cached, fail-soft to platform keys); env_get — a request-scoped credential chain (ContextVar tenant map, then process env). - Ingress: _caller_auth_async falls through to the control plane only on a pure local miss (local stores stay authoritative; explicit local status=inactive is an operator kill-switch for a tenant slug). Plan rate limits ride the resolve response into the existing _rate_ok meta. x-llm-router-tenant / x-internal-secret are stripped from client input; the tenant header is stamped only from the authenticated resolve. - internal_api.py (new): GET /internal/usage[?bucket=day] and /internal/usage/recent, gated by x-internal-secret (hidden 404 while unconfigured), backed by new host_store.usage_totals (includes cached tokens) and recent_calls(caller=...). - Router: per-tenant BYO provider credentials — the shim activates the tenant env from the trusted header; serve.py threads control_plane_client.env_get into the openai-compatible + anthropic + google adapters (bedrock/codex/antseed stay platform-only; background source pollers never see tenant keys). _rate_ok/_route_allowed now accept the already-resolved meta, dropping a duplicate consumer_keys read per proxied request. --- .env.example | 7 + auth_proxy.py | 78 ++++++- compose.yml | 10 + control_plane_client.py | 281 ++++++++++++++++++++++++ host_store.py | 42 +++- internal_api.py | 97 +++++++++ providers.py | 44 ++-- serve.py | 17 +- shim.py | 26 +++ tests/test_auth_proxy_control_plane.py | 287 +++++++++++++++++++++++++ tests/test_control_plane_client.py | 249 +++++++++++++++++++++ tests/test_host_store.py | 41 ++++ tests/test_tenant_env_shim.py | 223 +++++++++++++++++++ 13 files changed, 1368 insertions(+), 34 deletions(-) create mode 100644 control_plane_client.py create mode 100644 internal_api.py create mode 100644 tests/test_auth_proxy_control_plane.py create mode 100644 tests/test_control_plane_client.py create mode 100644 tests/test_tenant_env_shim.py diff --git a/.env.example b/.env.example index 9cc5682..844b129 100644 --- a/.env.example +++ b/.env.example @@ -68,3 +68,10 @@ ANTSEED_MAX_OUTPUT=1000 # to deposit/pay gas with). Defaults to a public Base endpoint; set ""/off to # disable the on-chain read (dashboard then shows escrow only). ANTSEED_WALLET_RPC_URL= + +# Optional external control plane ("bring your own control plane"): a separate +# service owning consumer keys / tenants (e.g. unhardcoded-cloud). Both must be +# set to enable; otherwise the router runs operator-only. The secret gates the +# ingress /internal/usage surface AND authenticates outbound resolve calls. +CONTROL_PLANE_URL= +CONTROL_PLANE_INTERNAL_SECRET= diff --git a/auth_proxy.py b/auth_proxy.py index dd84906..8022429 100644 --- a/auth_proxy.py +++ b/auth_proxy.py @@ -22,7 +22,9 @@ from fastapi import FastAPI, Request from fastapi.responses import HTMLResponse, JSONResponse, Response, StreamingResponse +import control_plane_client import host_store +import internal_api from env_secrets import load_env_secrets from shim import _CACHE_READ_FACTOR # the billing cache-read discount — one source @@ -63,6 +65,10 @@ SYNTHETIC_PROBE_INITIAL_DELAY_S = float(os.getenv("DASHBOARD_SYNTHETIC_PROBE_INITIAL_DELAY_S", "45")) SYNTHETIC_PROBE_TIMEOUT_S = float(os.getenv("DASHBOARD_SYNTHETIC_PROBE_TIMEOUT_S", "45")) SYNTHETIC_PROBE_CALLER = os.getenv("DASHBOARD_SYNTHETIC_PROBE_CALLER", "dashboard-probe") +# Optional namespace for control-plane-resolved callers (e.g. "cp:") so tenant +# slugs can't merge attribution with operator-minted consumer names. Default off: +# calls.caller == tenant slug, which is what the control plane's usage reads key on. +CONTROL_PLANE_CALLER_PREFIX = os.getenv("CONTROL_PLANE_CALLER_PREFIX", "") @@ -84,6 +90,10 @@ def _load_caller_map(raw: str, name: str) -> Dict[str, str]: logging.basicConfig(level=os.getenv("LOG_LEVEL", "INFO"), format="%(message)s") log = logging.getLogger("llm-router-auth-proxy") app = FastAPI(title="llm-router auth proxy", docs_url=None, redoc_url=None) +# Control-plane metering surface (/internal/usage[/recent]). Registered before +# any route definition below so Starlette matches it ahead of the catch-all +# proxy; hidden (404) unless CONTROL_PLANE_INTERNAL_SECRET is configured. +app.include_router(internal_api.router) _client: httpx.AsyncClient | None = None _probe_task: asyncio.Task[None] | None = None _windows: dict[str, deque[float]] = defaultdict(deque) @@ -477,8 +487,44 @@ def _caller_auth(token: str | None) -> dict[str, Any]: return {"ok": True, "caller": caller, "digest": digest, "storage": storage, "meta": meta} -def _rate_ok(caller: str) -> bool: +async def _caller_auth_async(token: str | None) -> dict[str, Any]: + """_caller_auth plus the external-control-plane fallback. + + Local key stores stay authoritative: a locally known key — including a + revoked/inactive one — never falls through to the control plane; only a + pure local miss consults it (and only when the feature is configured, see + control_plane_client.enabled()). On a resolve hit the caller becomes the + control-plane consumer name (tenant slug) and the plan's rate limits ride + the returned meta; an explicit local record for that same name keeps + precedence (operator kill-switch: mark the slug inactive locally).""" + auth = _caller_auth(token) + if auth.get("ok") or auth.get("caller") or not token or not control_plane_client.enabled(): + return auth + digest = hashlib.sha256(token.encode()).hexdigest() + resolved = await control_plane_client.resolve_key(digest) + if resolved is None or not resolved.active or not resolved.consumer: + return auth + caller = f"{CONTROL_PLANE_CALLER_PREFIX}{resolved.consumer}" meta = _consumer_meta(caller) + if meta.get("keys"): + # A local consumer already answers to this name — attribution merges. + control_plane_client.log_collision_once(caller) + if meta.get("status") != "active": + return {"ok": False, "caller": caller, "digest": digest, + "storage": "control_plane", "error_code": "caller_inactive"} + if meta.get("rate_per_min") is None: + meta["rate_per_min"] = resolved.rate_per_min + if meta.get("burst") is None: + meta["burst"] = resolved.burst + out = {"ok": True, "caller": caller, "digest": digest, + "storage": "control_plane", "meta": meta} + if resolved.tenant_id is not None: + out["tenant_id"] = resolved.tenant_id + return out + + +def _rate_ok(caller: str, meta: dict[str, Any] | None = None) -> bool: + meta = meta if meta is not None else _consumer_meta(caller) rate_per_min = int(meta.get("rate_per_min") or RATE_PER_MIN) burst = int(meta.get("burst") or BURST) now = time.monotonic() @@ -517,8 +563,10 @@ def _route_matches(pattern: str, route: str) -> bool: return False -def _route_allowed(caller: str, route: str | None) -> bool: - allowed_routes = _consumer_meta(caller).get("allowed_routes") or [] +def _route_allowed(caller: str, route: str | None, + meta: dict[str, Any] | None = None) -> bool: + meta = meta if meta is not None else _consumer_meta(caller) + allowed_routes = meta.get("allowed_routes") or [] if not allowed_routes: return True if not route: @@ -973,6 +1021,7 @@ async def shutdown() -> None: _probe_task = None if _client: await _client.aclose() + await control_plane_client.close() @app.get("/healthz") @@ -1730,7 +1779,7 @@ async def consumer_skill(request: Request) -> Response: and the same key auth (revoked/expired/inactive keys are rejected).""" auth = request.headers.get("authorization") or "" token = auth[7:].strip() if auth.lower().startswith("bearer ") else auth.strip() - if not _caller_auth(token).get("ok"): + if not (await _caller_auth_async(token)).get("ok"): return JSONResponse(status_code=401, content={"error": { "message": "invalid or inactive API key", "type": "auth_error", "code": "consumer_auth"}}) @@ -1749,7 +1798,7 @@ async def consumer_skill(request: Request) -> Response: # or mutate state: they admit/normalize a term, preview the ranking, or list the # field vocabulary — all derived from the live catalog the skill already shows. async def _consumer_x_proxy(request: Request, method: str, path: str) -> Response: - if not _caller_auth(_extract_token(request)).get("ok"): + if not (await _caller_auth_async(_extract_token(request))).get("ok"): return JSONResponse(status_code=401, content={"error": { "message": "invalid or inactive API key", "type": "auth_error", "code": "consumer_auth"}}) @@ -2376,7 +2425,7 @@ async def dashboard_key_usage(request: Request) -> Response: @app.get("/api/usage") async def key_usage(request: Request) -> Response: token = _extract_token(request) - auth = _caller_auth(token) + auth = await _caller_auth_async(token) if not auth.get("ok"): code = auth.get("error_code") or "caller_auth" status_code = 403 if code in {"caller_inactive", "caller_key_revoked", "caller_key_expired"} else 401 @@ -2404,7 +2453,7 @@ async def session_view(sid: str, request: Request) -> Response: itself sends as X-Unhardcoded-Session. Lets a harness (e.g. the opencode plugin) show live router economics without operator access to /x/*.""" token = _extract_token(request) - auth = _caller_auth(token) + auth = await _caller_auth_async(token) if not auth.get("ok"): code = auth.get("error_code") or "caller_auth" status_code = 403 if code in {"caller_inactive", "caller_key_revoked", "caller_key_expired"} else 401 @@ -3207,7 +3256,7 @@ async def proxy(path: str, request: Request) -> Response: return JSONResponse(status_code=404, content={"error": {"message": "not found", "type": "invalid_request_error", "code": None}}) started = time.perf_counter() token = _extract_token(request) - auth = _caller_auth(token) + auth = await _caller_auth_async(token) caller = auth.get("caller") if not auth.get("ok"): code = auth.get("error_code") or "caller_auth" @@ -3223,13 +3272,14 @@ async def proxy(path: str, request: Request) -> Response: return JSONResponse(status_code=status_code, content={"error": {"message": messages.get(code, "caller not authorized"), "type": "auth_error", "code": code}}) caller = str(caller) + auth_meta = auth.get("meta") body = await request.body() requested_route = _requested_route_from(path, body) - if not _route_allowed(caller, requested_route): + if not _route_allowed(caller, requested_route, meta=auth_meta): _record_reject(reason="route_not_allowed", path="/" + path, caller=caller, status=403, route=requested_route) _log({"event": "reject", "reason": "route_not_allowed", "caller": caller, "path": "/" + path, "route": requested_route}) return JSONResponse(status_code=403, content={"error": {"message": "caller is not allowed to use this route", "type": "auth_error", "code": "caller_route_not_allowed"}}) - if not _rate_ok(caller): + if not _rate_ok(caller, meta=auth_meta): _record_reject(reason="rate_limit", path="/" + path, caller=caller, status=429) _log({"event": "reject", "reason": "rate_limit", "caller": caller, "path": "/" + path}) return JSONResponse(status_code=429, content={"error": {"message": "caller rate limit exceeded", "type": "rate_limit_error", "code": "caller_rate_limit"}}) @@ -3255,9 +3305,15 @@ async def proxy(path: str, request: Request) -> Response: headers = { k: v for k, v in request.headers.items() - if k.lower() not in {"authorization", "host", "connection", "content-length"} + if k.lower() not in {"authorization", "host", "connection", "content-length", + "x-llm-router-tenant", "x-internal-secret"} } headers["x-llm-router-caller"] = caller + # Tenant identity for per-tenant provider credentials (BYO keys): set ONLY + # from the authenticated resolve — the client-sent header is stripped above, + # so it can never be smuggled past auth. + if auth.get("tenant_id") is not None: + headers["x-llm-router-tenant"] = str(auth["tenant_id"]) status = 502 provider = None diff --git a/compose.yml b/compose.yml index 6df7e4e..38bd6d3 100644 --- a/compose.yml +++ b/compose.yml @@ -27,6 +27,11 @@ services: ANTSEED_CONTROL_TOKEN: ${ANTSEED_CONTROL_TOKEN:-} # Operational store (router reads operator config from it). Prod = RDS. DATABASE_URL: ${DATABASE_URL:-postgresql://hoststore:hoststore@postgres:5432/hoststore} + # Optional external control plane: per-tenant BYO provider credentials + # (fetched by tenant id from the trusted x-llm-router-tenant header the + # ingress stamps). Both unset -> feature off, platform keys only. + CONTROL_PLANE_URL: ${CONTROL_PLANE_URL:-} + CONTROL_PLANE_INTERNAL_SECRET: ${CONTROL_PLANE_INTERNAL_SECRET:-} depends_on: postgres: condition: service_healthy @@ -84,6 +89,11 @@ services: DATABASE_URL: ${DATABASE_URL:-postgresql://hoststore:hoststore@postgres:5432/hoststore} RATE_PER_MIN: ${RATE_PER_MIN:-600} BURST: ${BURST:-200} + # Optional external control plane: consumer-key resolve fallback plus the + # /internal/usage metering surface. Both unset -> feature off (/internal/* + # answers 404, unknown keys just 401). + CONTROL_PLANE_URL: ${CONTROL_PLANE_URL:-} + CONTROL_PLANE_INTERNAL_SECRET: ${CONTROL_PLANE_INTERNAL_SECRET:-} LOG_LEVEL: ${LOG_LEVEL:-INFO} DASHBOARD_KEY_ENV_PATH: /run/llm-router/.env.secrets # env_file carries the HOST path; inside the container the file is the diff --git a/control_plane_client.py b/control_plane_client.py new file mode 100644 index 0000000..3b085a7 --- /dev/null +++ b/control_plane_client.py @@ -0,0 +1,281 @@ +"""Client for an external control plane ("bring your own control plane"). + +An operator can front this router with a separate control plane that owns +consumer keys and per-tenant provider credentials (any service speaking the +small HTTP contract below). The feature is OFF unless both CONTROL_PLANE_URL +and CONTROL_PLANE_INTERNAL_SECRET are set; nothing in this module assumes any +particular control-plane implementation. + +Contract (all requests carry the shared secret in `x-internal-secret`): + GET {CONTROL_PLANE_URL}/internal/keys/resolve?sha256=<64hex> + -> {"active": bool, "consumer": str, "tenant_id": int, + "rate_per_min": int|null, "burst": int|null} + GET {CONTROL_PLANE_URL}/internal/tenants//provider-env + -> {"env": {ENV_NAME: secret, ...}} + +The module is a leaf (no imports from auth_proxy/shim) shared by the ingress +(key resolution) and the router (per-tenant provider env). Secrets are never +logged — events carry env NAMES and tenant ids only. +""" +from __future__ import annotations + +import asyncio +import contextvars +import hashlib +import hmac +import json +import logging +import os +import time +from dataclasses import dataclass +from typing import Any, Mapping + +import httpx + +log = logging.getLogger("llm-router-control-plane") + +CONTROL_PLANE_URL = os.getenv("CONTROL_PLANE_URL", "").rstrip("/") +CONTROL_PLANE_INTERNAL_SECRET = os.getenv("CONTROL_PLANE_INTERNAL_SECRET", "") +RESOLVE_TTL_S = float(os.getenv("CP_RESOLVE_TTL_S", "60")) +NEGATIVE_TTL_S = float(os.getenv("CP_NEGATIVE_TTL_S", "15")) +RESOLVE_STALE_GRACE_S = float(os.getenv("CP_RESOLVE_STALE_GRACE_S", "300")) +TENANT_ENV_TTL_S = float(os.getenv("CP_TENANT_ENV_TTL_S", "120")) +TENANT_ENV_STALE_GRACE_S = float(os.getenv("CP_TENANT_ENV_STALE_GRACE_S", "600")) +ENV_ALLOWLIST = { + name.strip() + for name in os.getenv( + "CP_ENV_ALLOWLIST", + "OPENAI_API_KEY,OPENROUTER_API_KEY,ANTHROPIC_API_KEY,GEMINI_API_KEY", + ).split(",") + if name.strip() +} + +_RESOLVE_CACHE_MAX = 4096 +_TIMEOUT = httpx.Timeout(3.0, connect=1.5) + + +def enabled() -> bool: + return bool(CONTROL_PLANE_URL and CONTROL_PLANE_INTERNAL_SECRET) + + +def internal_secret_ok(headers: Mapping[str, str]) -> bool: + """Validate an inbound `x-internal-secret` header. False when the shared + secret is unconfigured — callers must treat that as 'feature hidden'.""" + if not CONTROL_PLANE_INTERNAL_SECRET: + return False + presented = headers.get("x-internal-secret") or "" + return hmac.compare_digest(presented, CONTROL_PLANE_INTERNAL_SECRET) + + +@dataclass +class ResolvedKey: + active: bool + consumer: str | None + tenant_id: int | None + rate_per_min: int | None + burst: int | None + fetched_at: float # time.monotonic() + + +_resolve_cache: dict[str, ResolvedKey] = {} # sha256 hex -> entry (positive AND negative) +_resolve_inflight: dict[str, asyncio.Future] = {} +_tenant_env_cache: dict[int, tuple[dict[str, str], float]] = {} +_TENANT_ENV: contextvars.ContextVar[dict[str, str] | None] = contextvars.ContextVar( + "cp_tenant_env", default=None +) +_client: httpx.AsyncClient | None = None +_collision_logged: set[str] = set() + + +def _get_client() -> httpx.AsyncClient: + global _client + if _client is None: + _client = httpx.AsyncClient(timeout=_TIMEOUT) + return _client + + +def sha256_hex(token: str) -> str: + return hashlib.sha256(token.encode()).hexdigest() + + +def _entry_ttl(entry: ResolvedKey) -> float: + return RESOLVE_TTL_S if entry.active else NEGATIVE_TTL_S + + +def _evict_if_full() -> None: + if len(_resolve_cache) < _RESOLVE_CACHE_MAX: + return + now = time.monotonic() + expired = [k for k, e in _resolve_cache.items() if now - e.fetched_at > _entry_ttl(e)] + for k in expired: + _resolve_cache.pop(k, None) + if len(_resolve_cache) >= _RESOLVE_CACHE_MAX: + _resolve_cache.clear() + + +def _parse_resolved(data: Any) -> ResolvedKey: + if not isinstance(data, dict): + data = {} + + def _opt_int(value: Any) -> int | None: + try: + out = int(value) + except (TypeError, ValueError): + return None + return out if out > 0 else None + + consumer = str(data.get("consumer") or "").strip() or None + active = bool(data.get("active")) and consumer is not None + return ResolvedKey( + active=active, + consumer=consumer if active else None, + tenant_id=_opt_int(data.get("tenant_id")) if active else None, + rate_per_min=_opt_int(data.get("rate_per_min")) if active else None, + burst=_opt_int(data.get("burst")) if active else None, + fetched_at=time.monotonic(), + ) + + +async def _fetch_resolve(digest: str) -> ResolvedKey | None: + """One HTTP resolve. Returns the parsed entry (positive or negative) on a + definitive control-plane answer, None on transport error / 5xx.""" + try: + resp = await _get_client().get( + f"{CONTROL_PLANE_URL}/internal/keys/resolve", + params={"sha256": digest}, + headers={"x-internal-secret": CONTROL_PLANE_INTERNAL_SECRET}, + ) + except httpx.HTTPError as exc: + log.warning(json.dumps({"event": "cp_resolve_error", "error": type(exc).__name__})) + return None + if resp.status_code >= 500: + log.warning(json.dumps({"event": "cp_resolve_error", "status": resp.status_code})) + return None + if resp.status_code != 200: + # 403 (secret mismatch) etc. — a definitive "no": cache as negative so a + # misconfigured secret can't turn into a per-request CP hammer. + log.warning(json.dumps({"event": "cp_resolve_rejected", "status": resp.status_code})) + return _parse_resolved({}) + try: + return _parse_resolved(resp.json()) + except ValueError: + return _parse_resolved({}) + + +async def resolve_key(digest: str) -> ResolvedKey | None: + """Resolve a key digest against the control plane, with caching. + + Returns None when the feature is off or the CP is unreachable with no + usable cache (the caller should 401). A stale positive entry is served for + up to RESOLVE_STALE_GRACE_S past its TTL, but ONLY when the CP is + unreachable — a definitive answer always replaces the cache. Negative + entries never get grace. + """ + if not enabled(): + return None + now = time.monotonic() + cached = _resolve_cache.get(digest) + if cached is not None and now - cached.fetched_at <= _entry_ttl(cached): + return cached + + pending = _resolve_inflight.get(digest) + if pending is not None: + return await asyncio.shield(pending) + + future: asyncio.Future = asyncio.get_running_loop().create_future() + _resolve_inflight[digest] = future + try: + fresh = await _fetch_resolve(digest) + if fresh is not None: + _evict_if_full() + _resolve_cache[digest] = fresh + result: ResolvedKey | None = fresh + elif ( + cached is not None + and cached.active + and now - cached.fetched_at <= RESOLVE_TTL_S + RESOLVE_STALE_GRACE_S + ): + log.warning(json.dumps({"event": "cp_resolve_stale_grace", "consumer": cached.consumer})) + result = cached + else: + _resolve_cache.pop(digest, None) + result = None + future.set_result(result) + return result + except BaseException as exc: + future.set_exception(exc) + raise + finally: + _resolve_inflight.pop(digest, None) + + +async def tenant_env(tenant_id: int) -> dict[str, str]: + """Cached BYO provider env for a tenant, filtered through ENV_ALLOWLIST. + Fail-soft: refetch error -> stale within grace -> {} (platform keys).""" + if not enabled(): + return {} + now = time.monotonic() + cached = _tenant_env_cache.get(tenant_id) + if cached is not None and now - cached[1] <= TENANT_ENV_TTL_S: + return cached[0] + try: + resp = await _get_client().get( + f"{CONTROL_PLANE_URL}/internal/tenants/{int(tenant_id)}/provider-env", + headers={"x-internal-secret": CONTROL_PLANE_INTERNAL_SECRET}, + ) + resp.raise_for_status() + raw = resp.json().get("env") + except (httpx.HTTPError, ValueError) as exc: + if cached is not None and now - cached[1] <= TENANT_ENV_TTL_S + TENANT_ENV_STALE_GRACE_S: + log.warning(json.dumps({"event": "tenant_env_stale_grace", "tenant_id": tenant_id})) + return cached[0] + log.warning(json.dumps({ + "event": "tenant_env_fallback", "tenant_id": tenant_id, "error": type(exc).__name__, + })) + return {} + env = { + str(k): str(v) + for k, v in (raw or {}).items() + if str(k) in ENV_ALLOWLIST and isinstance(v, str) and v + } + _tenant_env_cache[tenant_id] = (env, now) + return env + + +def activate_tenant_env(env: dict[str, str] | None) -> contextvars.Token: + return _TENANT_ENV.set(env or None) + + +def reset_tenant_env(token: contextvars.Token) -> None: + _TENANT_ENV.reset(token) + + +def env_get(name: str) -> str | None: + """Adapter credential lookup: per-request tenant map first, then process env.""" + override = _TENANT_ENV.get() + if override and name in override: + return override[name] + return os.environ.get(name) + + +def log_collision_once(consumer: str) -> None: + if consumer in _collision_logged: + return + _collision_logged.add(consumer) + log.warning(json.dumps({"event": "cp_caller_collision", "caller": consumer})) + + +async def close() -> None: + global _client + if _client is not None: + await _client.aclose() + _client = None + + +def reset_for_tests() -> None: + global _client + _resolve_cache.clear() + _resolve_inflight.clear() + _tenant_env_cache.clear() + _collision_logged.clear() + _client = None diff --git a/host_store.py b/host_store.py index 50be7ce..e0b574e 100644 --- a/host_store.py +++ b/host_store.py @@ -378,13 +378,15 @@ def observe_route_call_async(row: dict[str, Any]) -> None: _enqueue(lambda: _insert_route_observation(snap)) -def recent_calls(limit: int = 100) -> list[dict[str, Any]]: - """The most recent calls, newest first (operator view / verification).""" +def recent_calls(limit: int = 100, caller: "str | None" = None) -> list[dict[str, Any]]: + """The most recent calls, newest first (operator view / verification). + Optionally scoped to one caller (control-plane activity feed).""" try: + where, params = ("", []) if caller is None else (" WHERE caller = %s", [caller]) with _get_pool().connection() as conn: with conn.cursor(row_factory=dict_row) as cur: - cur.execute("SELECT * FROM calls ORDER BY id DESC LIMIT %s", - (int(limit),)) + cur.execute(f"SELECT * FROM calls{where} ORDER BY id DESC LIMIT %s", + params + [int(limit)]) return cur.fetchall() except Exception as exc: # noqa: BLE001 _log.warning("host_store recent_calls failed: %s", exc) @@ -912,6 +914,38 @@ def usage_aggregate(since_ts: "int | None" = None, caller: "str | None" = None, return _empty_usage_aggregate() +def usage_totals(since_ts: "int | None" = None, + caller: "str | None" = None) -> dict[str, Any]: + """One-row window totals over `calls`, including cached tokens (which the + dashboard aggregate doesn't sum) — the control-plane metering read. + Fail-soft -> zeros (same keys).""" + zeros = {"requests": 0, "errors": 0, "tokens_in": 0, "tokens_out": 0, + "tokens_cached": 0, "tokens_total": 0, "cost_usd": 0.0, "priced": 0} + try: + where, params = _usage_where(since_ts=since_ts, caller=caller) + sql = ( + "SELECT count(*)," + " count(*) FILTER (WHERE COALESCE(status,0) >= 400)," + " COALESCE(sum(COALESCE(tokens_in,0)),0)," + " COALESCE(sum(COALESCE(tokens_out,0)),0)," + " COALESCE(sum(COALESCE(tokens_cached,0)),0)," + " COALESCE(sum(CASE WHEN COALESCE(tokens_total,0) <> 0 THEN tokens_total" + " ELSE COALESCE(tokens_in,0)+COALESCE(tokens_out,0) END),0)," + " round(COALESCE(sum(GREATEST(cost_usd,0)),0)::numeric,6)::float8," + " count(cost_usd)" + f" FROM calls{where}") + with _get_pool().connection() as conn: + row = conn.execute(sql, params).fetchone() + requests, errors, tin, tout, tcached, ttotal, cost, priced = row + return {"requests": int(requests), "errors": int(errors), + "tokens_in": int(tin), "tokens_out": int(tout), + "tokens_cached": int(tcached), "tokens_total": int(ttotal), + "cost_usd": float(cost), "priced": int(priced)} + except Exception as exc: # noqa: BLE001 + _log.warning("host_store usage_totals failed: %s", exc) + return zeros + + def policy_backtest_groups(since_ts: "int | None" = None, caller: "str | None" = None, limit: int = 50) -> dict[str, Any]: diff --git a/internal_api.py b/internal_api.py new file mode 100644 index 0000000..cfa8dd4 --- /dev/null +++ b/internal_api.py @@ -0,0 +1,97 @@ +"""Internal metering API for an external control plane. + +Served by the ingress, gated by the same shared secret the control-plane +client sends outbound (`x-internal-secret`); see control_plane_client for the +overall contract. Hidden entirely (404) while the secret is unconfigured, so +the surface does not exist on operator-only deployments. + + GET /internal/usage?caller=&since_ts=[&bucket=day] + GET /internal/usage/recent?caller=&limit= + +Rows come from the `calls` ledger, which only records LLM calls that reached +the router — ingress-level rejects (401/429) are not counted here. +""" +from __future__ import annotations + +import time +from typing import Any + +from fastapi import APIRouter, Request +from fastapi.responses import JSONResponse + +import control_plane_client +import host_store + +router = APIRouter() + +_RECENT_LIMIT_MAX = 500 + + +def _gate(request: Request) -> JSONResponse | None: + if not control_plane_client.CONTROL_PLANE_INTERNAL_SECRET: + return JSONResponse({"error": "not_found"}, status_code=404) + if not control_plane_client.internal_secret_ok(request.headers): + return JSONResponse({"error": "forbidden"}, status_code=403) + return None + + +@router.get("/internal/usage") +async def internal_usage(request: Request, caller: str = "", + since_ts: int | None = None, + bucket: str | None = None) -> JSONResponse: + denied = _gate(request) + if denied is not None: + return denied + caller = caller.strip() + if not caller: + return JSONResponse({"error": "caller_required"}, status_code=400) + totals = host_store.usage_totals(since_ts=since_ts, caller=caller) + out: dict[str, Any] = { + "caller": caller, + "window": {"since_ts": since_ts, "until_ts": int(time.time())}, + "runs": totals["requests"], + "errors": totals["errors"], + "tokens_in": totals["tokens_in"], + "tokens_out": totals["tokens_out"], + "tokens_cached": totals["tokens_cached"], + "tokens_total": totals["tokens_total"], + "cost_usd": totals["cost_usd"], + } + if bucket == "day": + by_day = host_store.usage_aggregate(since_ts=since_ts, caller=caller)["by_day"] + out["buckets"] = [ + {"date": day, "runs": counter["requests"], "cost_usd": counter["cost_usd"]} + for day, counter in sorted(by_day.items()) + ] + return JSONResponse(out) + + +@router.get("/internal/usage/recent") +async def internal_usage_recent(request: Request, caller: str = "", + limit: int = 50) -> JSONResponse: + denied = _gate(request) + if denied is not None: + return denied + caller = caller.strip() + if not caller: + return JSONResponse({"error": "caller_required"}, status_code=400) + limit = max(1, min(int(limit), _RECENT_LIMIT_MAX)) + calls = [] + for row in host_store.recent_calls(limit=limit, caller=caller): + calls.append({ + "ts": row.get("ts"), + "status": row.get("status"), + "requested_model": row.get("requested_model"), + "model_family": row.get("model_family"), + "provider": row.get("provider_id"), + "served_model_id": row.get("served_model_id"), + "latency_ms": row.get("latency_ms"), + "tokens_in": row.get("tokens_in"), + "tokens_out": row.get("tokens_out"), + "tokens_total": row.get("tokens_total"), + "tokens_cached": row.get("tokens_cached"), + "cost_usd": row.get("cost_usd"), + "error_type": row.get("error_type"), + "key_sha256_prefix": (row.get("consumer_sha") or "")[:12] or None, + }) + return JSONResponse({"caller": caller, "calls": calls}) diff --git a/providers.py b/providers.py index c300d4a..2ed9412 100644 --- a/providers.py +++ b/providers.py @@ -80,34 +80,45 @@ def _bedrock_source(catalog, env_get): return BedrockSource(catalog, env_get=env_get) -def _anthropic_adapter(timeout_s): +def _anthropic_adapter(timeout_s, env_get=None): from provider_adapters.anthropic import make_anthropic_async_call_provider - return make_anthropic_async_call_provider(timeout_s=timeout_s) + if env_get is None: + return make_anthropic_async_call_provider(timeout_s=timeout_s) + return make_anthropic_async_call_provider(timeout_s=timeout_s, env_get=env_get) -def _anthropic_stream_adapter(timeout_s): +def _anthropic_stream_adapter(timeout_s, env_get=None): from provider_adapters.anthropic import stream_anthropic - return functools.partial(stream_anthropic, timeout_s=timeout_s) + if env_get is None: + return functools.partial(stream_anthropic, timeout_s=timeout_s) + return functools.partial(stream_anthropic, timeout_s=timeout_s, env_get=env_get) -def _bedrock_adapter(timeout_s): +def _bedrock_adapter(timeout_s, env_get=None): + # env_get intentionally ignored: bedrock auths via the AWS chain, and tenant + # BYO credentials never apply to it (platform-only provider). from provider_adapters.bedrock import make_bedrock_async_call_provider return make_bedrock_async_call_provider(timeout_s=timeout_s) -def _bedrock_stream_adapter(timeout_s): +def _bedrock_stream_adapter(timeout_s, env_get=None): + # env_get intentionally ignored (see _bedrock_adapter). from provider_adapters.bedrock import stream_bedrock return functools.partial(stream_bedrock, timeout_s=timeout_s) -def _google_adapter(timeout_s): +def _google_adapter(timeout_s, env_get=None): from provider_adapters.google import make_google_async_call_provider - return make_google_async_call_provider(timeout_s=timeout_s) + if env_get is None: + return make_google_async_call_provider(timeout_s=timeout_s) + return make_google_async_call_provider(timeout_s=timeout_s, env_get=env_get) -def _google_stream_adapter(timeout_s): +def _google_stream_adapter(timeout_s, env_get=None): from provider_adapters.google import stream_google - return functools.partial(stream_google, timeout_s=timeout_s) + if env_get is None: + return functools.partial(stream_google, timeout_s=timeout_s) + return functools.partial(stream_google, timeout_s=timeout_s, env_get=env_get) def _codex_source(catalog, env_get): @@ -265,18 +276,21 @@ def build_source_registry(catalog: dict, env_get=os.environ.get) -> list: return out -def native_adapter_handlers(timeout_s: float) -> "dict[str, Any]": +def native_adapter_handlers(timeout_s: float, env_get=None) -> "dict[str, Any]": """api_kind -> wire backend for the providers with a dedicated adapter - (codex is wired separately in serve.py because it needs `observe`).""" - return {p.api_kind: p.adapter(timeout_s) + (codex is wired separately in serve.py because it needs `observe`). + `env_get` threads a request-scoped credential lookup (per-tenant BYO keys) + into the adapters that resolve auth from env; platform-only adapters + (bedrock) ignore it.""" + return {p.api_kind: p.adapter(timeout_s, env_get) for p in PROVIDERS if not p.special and p.api_kind and p.adapter} -def native_streaming_adapter_handlers(timeout_s: float) -> "dict[str, Any]": +def native_streaming_adapter_handlers(timeout_s: float, env_get=None) -> "dict[str, Any]": """api_kind -> true streaming backend for native providers that support it (codex remains wired separately in serve.py because it needs `observe`).""" - return {p.api_kind: p.stream_adapter(timeout_s) + return {p.api_kind: p.stream_adapter(timeout_s, env_get) for p in PROVIDERS if not p.special and p.api_kind and p.stream_adapter} diff --git a/serve.py b/serve.py index 839019d..2d4843c 100644 --- a/serve.py +++ b/serve.py @@ -135,10 +135,17 @@ def main() -> None: # The native api_kind adapters (anthropic/bedrock/google) come from the # modular provider registry; codex is wired here because its backend takes # the `observe` hook that feeds its scarcity-price source. - _native = providers.native_adapter_handlers(args.timeout_s) + # Credential lookup for env-auth adapters: control_plane_client.env_get + # consults the request-scoped tenant BYO map first (set by the shim from the + # trusted x-llm-router-tenant header), then the process env — identical to + # os.environ.get while the control-plane feature is off. + import control_plane_client + _native = providers.native_adapter_handlers(args.timeout_s, + env_get=control_plane_client.env_get) call_async = make_api_kind_dispatcher( default=make_async_call_provider(timeout_s=args.timeout_s, - provider_rules=provider_rules), + provider_rules=provider_rules, + env_get=control_plane_client.env_get), handlers={**_native, "openai_codex": make_codex_async_call_provider(codex_auth, observe=observe)}, ) @@ -152,11 +159,13 @@ def main() -> None: stream_codex, stream_openai_compatible, ) - _native_streaming = providers.native_streaming_adapter_handlers(args.timeout_s) + _native_streaming = providers.native_streaming_adapter_handlers( + args.timeout_s, env_get=control_plane_client.env_get) streaming_call = make_streaming_dispatcher( default=functools.partial(stream_openai_compatible, timeout_s=args.timeout_s, - provider_rules=provider_rules), + provider_rules=provider_rules, + env_get=control_plane_client.env_get), # Native providers and Codex have real streaming twins; Codex also feeds # `observe` for quota/scarcity pricing. handlers={**_native_streaming, diff --git a/shim.py b/shim.py index 5509937..7778322 100644 --- a/shim.py +++ b/shim.py @@ -37,6 +37,7 @@ from fastapi.responses import JSONResponse, StreamingResponse from pydantic import BaseModel, ConfigDict +import control_plane_client import host_store _log = logging.getLogger("unhardcoded.shim") @@ -839,9 +840,31 @@ async def compact(req: CompactRequest): sealed = {"role": "system", "content": _SEAL_PREFIX + summary} return {"messages": frozen + [sealed] + recent, "compacted": True} + async def _activate_tenant(request: Request) -> None: + """Load the caller tenant's BYO provider env into the request context. + + The tenant id rides x-llm-router-tenant, stamped by the ingress ONLY + after key auth (client copies are stripped there). The fetched map is + activated on a ContextVar in THIS request's task context: asyncio tasks + created later in the request (streaming, flow nodes) copy it, and the + context dies with the request task, so no reset is needed and requests + can't see each other's credentials. Fail-soft: no/invalid header or a + fetch failure leaves the platform env in force.""" + raw = request.headers.get("x-llm-router-tenant") + if not raw or not control_plane_client.enabled(): + return + try: + tenant_id = int(raw) + except ValueError: + return + env = await control_plane_client.tenant_env(tenant_id) + if env: + control_plane_client.activate_tenant_env(env) + @app.post("/v1/chat/completions") async def chat_completions(req: ChatRequest, request: Request): _session_from_header(req, request) + await _activate_tenant(request) return await _handle_chat(req) @app.post("/{profile_name}/v1/chat/completions") @@ -854,6 +877,7 @@ async def chat_completions_profiled(profile_name: str, req: ChatRequest, request instead of a `profile:` model prefix. """ _session_from_header(req, request) + await _activate_tenant(request) return await _handle_chat(req, profile_name=profile_name) def _session_from_header(req: ChatRequest, request: Request) -> None: @@ -882,12 +906,14 @@ def _session_from_header(req: ChatRequest, request: Request) -> None: @app.post("/v1/responses") async def responses(req: ResponsesRequest, request: Request): _session_from_header(req, request) + await _activate_tenant(request) return await _handle_responses(req) @app.post("/{profile_name}/v1/responses") async def responses_profiled(profile_name: str, req: ResponsesRequest, request: Request): _session_from_header(req, request) + await _activate_tenant(request) return await _handle_responses(req, profile_name=profile_name) def _responses_object_with_router(result: dict, req: ResponsesRequest, diff --git a/tests/test_auth_proxy_control_plane.py b/tests/test_auth_proxy_control_plane.py new file mode 100644 index 0000000..70ebe49 --- /dev/null +++ b/tests/test_auth_proxy_control_plane.py @@ -0,0 +1,287 @@ +"""Ingress <-> external control plane integration (auth fallback + /internal/*). + +The control plane's HTTP side is a fake client on control_plane_client._client; +the upstream router is a fake on auth_proxy._client. Store-backed tests use the +shared Postgres fixture (host_store_clean). +""" +from __future__ import annotations + +import hashlib +import os +import sys +import time +from pathlib import Path + +import httpx +import pytest +from fastapi.testclient import TestClient + +ROOT = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(ROOT)) + +# auth_proxy reads the caller-key env at import; this module can be the first +# in the session to import it, so mirror the suite-wide fixture env here or the +# later-collected dashboard tests would see an empty CALLER_KEYS map. +os.environ.setdefault("CALLER_KEYS_JSON", '{"internal":"default"}') +os.environ.setdefault("CALLER_KEYS_SHA256_JSON", "{}") + +import auth_proxy # noqa: E402 +import control_plane_client as cpc # noqa: E402 +import host_store # noqa: E402 + +from conftest import require_host_store # noqa: E402 + + +class _FakeCPResp: + def __init__(self, status_code: int, payload): + self.status_code = status_code + self._payload = payload + + def json(self): + return self._payload + + +class _FakeCPClient: + def __init__(self, payloads): + self.payloads = list(payloads) + self.calls = 0 + + async def get(self, url, params=None, headers=None): + self.calls += 1 + item = self.payloads.pop(0) + if isinstance(item, Exception): + raise item + return _FakeCPResp(200, item) + + +class _FakeUpstreamResp: + status_code = 200 + headers = {"content-type": "application/json"} + + async def aread(self): + return b'{"ok": true}' + + async def aclose(self): + pass + + +class _FakeUpstream: + """Captures the headers the proxy forwards to the router.""" + + def __init__(self): + self.requests: list[dict] = [] + + def build_request(self, method, url, content=None, headers=None): + self.requests.append({"method": method, "url": url, "headers": headers or {}}) + return object() + + async def send(self, req, stream=True): + return _FakeUpstreamResp() + + +@pytest.fixture(autouse=True) +def _clean(monkeypatch): + cpc.reset_for_tests() + auth_proxy._windows.clear() + monkeypatch.setattr(cpc, "CONTROL_PLANE_URL", "http://cp.test") + monkeypatch.setattr(cpc, "CONTROL_PLANE_INTERNAL_SECRET", "s3cret") + yield + cpc.reset_for_tests() + auth_proxy._windows.clear() + + +def _cp(monkeypatch, payloads) -> _FakeCPClient: + fake = _FakeCPClient(payloads) + monkeypatch.setattr(cpc, "_client", fake) + return fake + + +def _upstream(monkeypatch) -> _FakeUpstream: + fake = _FakeUpstream() + monkeypatch.setattr(auth_proxy, "_client", fake) + return fake + + +def _post_chat(client, token: str, extra_headers: dict | None = None): + headers = {"Authorization": f"Bearer {token}"} + headers.update(extra_headers or {}) + return client.post("/v1/chat/completions", headers=headers, + json={"model": "profile:default", "messages": []}) + + +# ---- key resolution ---------------------------------------------------------- + +def test_feature_off_unknown_key_401_without_cp_call(monkeypatch): + require_host_store() + monkeypatch.setattr(cpc, "CONTROL_PLANE_URL", "") + fake_cp = _cp(monkeypatch, []) + _upstream(monkeypatch) + r = _post_chat(TestClient(auth_proxy.app), "tok-unknown") + assert r.status_code == 401 + assert fake_cp.calls == 0 + + +def test_cp_resolved_key_proxies_with_caller_and_tenant_headers(monkeypatch): + require_host_store() + _cp(monkeypatch, [{"active": True, "consumer": "acme", "tenant_id": 7, + "rate_per_min": 600, "burst": 200}]) + upstream = _upstream(monkeypatch) + r = _post_chat(TestClient(auth_proxy.app), "tok-tenant", + extra_headers={"x-llm-router-tenant": "999", # smuggle attempt + "x-internal-secret": "leak"}) + assert r.status_code == 200 + fwd = upstream.requests[0]["headers"] + assert fwd["x-llm-router-caller"] == "acme" + assert fwd["x-llm-router-tenant"] == "7" # authed value, not the smuggled 999 + assert "x-internal-secret" not in fwd + + +def test_second_request_served_from_resolve_cache(monkeypatch): + require_host_store() + fake_cp = _cp(monkeypatch, [{"active": True, "consumer": "acme", "tenant_id": 7}]) + _upstream(monkeypatch) + client = TestClient(auth_proxy.app) + assert _post_chat(client, "tok-cache").status_code == 200 + assert _post_chat(client, "tok-cache").status_code == 200 + assert fake_cp.calls == 1 + + +def test_cp_plan_rate_limits_enforced(monkeypatch): + require_host_store() + _cp(monkeypatch, [{"active": True, "consumer": "tiny-plan", "tenant_id": 3, + "rate_per_min": 1, "burst": 1}]) + _upstream(monkeypatch) + client = TestClient(auth_proxy.app) + assert _post_chat(client, "tok-limited").status_code == 200 + r = _post_chat(client, "tok-limited") + assert r.status_code == 429 + assert r.json()["error"]["code"] == "caller_rate_limit" + + +def test_inactive_resolve_is_401(monkeypatch): + require_host_store() + _cp(monkeypatch, [{"active": False}]) + _upstream(monkeypatch) + assert _post_chat(TestClient(auth_proxy.app), "tok-revoked").status_code == 401 + + +def test_local_plaintext_key_never_consults_cp(monkeypatch): + require_host_store() + fake_cp = _cp(monkeypatch, []) + _upstream(monkeypatch) + monkeypatch.setattr(auth_proxy, "CALLER_KEYS", {"tok-local": "operator-app"}) + r = _post_chat(TestClient(auth_proxy.app), "tok-local") + assert r.status_code == 200 + assert fake_cp.calls == 0 + + +def test_locally_revoked_hash_key_never_falls_through_to_cp(monkeypatch): + require_host_store() + fake_cp = _cp(monkeypatch, []) + _upstream(monkeypatch) + digest = hashlib.sha256(b"tok-revoked-local").hexdigest() + monkeypatch.setattr(auth_proxy, "CALLER_KEY_HASHES", {digest: "operator-app"}) + host_store.set_consumer_keys({"operator-app": { + "status": "active", "keys": [{"sha256_prefix": digest[:12], "status": "revoked"}]}}) + r = _post_chat(TestClient(auth_proxy.app), "tok-revoked-local") + assert r.status_code == 403 + assert r.json()["error"]["code"] == "caller_key_revoked" + assert fake_cp.calls == 0 + + +def test_operator_kill_switch_blocks_cp_slug(monkeypatch): + require_host_store() + _cp(monkeypatch, [{"active": True, "consumer": "banned-tenant", "tenant_id": 4}]) + _upstream(monkeypatch) + host_store.set_consumer_keys({"banned-tenant": {"status": "inactive"}}) + r = _post_chat(TestClient(auth_proxy.app), "tok-banned") + assert r.status_code == 403 + assert r.json()["error"]["code"] == "caller_inactive" + + +def test_cp_caller_lands_in_the_ledger_under_the_tenant_slug(monkeypatch): + require_host_store() + _cp(monkeypatch, [{"active": True, "consumer": "acme", "tenant_id": 7}]) + _upstream(monkeypatch) + assert _post_chat(TestClient(auth_proxy.app), "tok-ledger").status_code == 200 + host_store._write_q.join() + rows = host_store.recent_calls(caller="acme") + assert len(rows) == 1 + assert rows[0]["caller"] == "acme" + + +# ---- /internal/usage surface --------------------------------------------------- + +def _seed_calls(): + now = int(time.time()) + base = {"session": "s", "key_sha256": "c" * 64, "provider": "openrouter", + "model_family": "fam", "served_model_id": "m", "requested_model": "profile:default", + "latency_ms": 100.0} + host_store.insert_call({**base, "ts": now - 10, "caller": "acme", "status": 200, + "tokens_in": 70, "tokens_out": 30, "tokens_total": 100, + "tokens_cached": 25, "cost_usd": 0.01}) + host_store.insert_call({**base, "ts": now - 5, "caller": "acme", "status": 500, + "tokens_in": 10, "tokens_out": 0, "tokens_total": 10, + "cost_usd": 0.0}) + host_store.insert_call({**base, "ts": now - 5, "caller": "other", "status": 200, + "tokens_in": 1000, "tokens_out": 1000, "tokens_total": 2000, + "cost_usd": 9.99}) + return now + + +def test_internal_usage_hidden_without_secret(monkeypatch): + monkeypatch.setattr(cpc, "CONTROL_PLANE_INTERNAL_SECRET", "") + r = TestClient(auth_proxy.app).get("/internal/usage", params={"caller": "acme"}) + assert r.status_code == 404 + + +def test_internal_usage_wrong_secret_403(): + r = TestClient(auth_proxy.app).get("/internal/usage", params={"caller": "acme"}, + headers={"x-internal-secret": "wrong"}) + assert r.status_code == 403 + + +def test_internal_usage_totals_and_daily_buckets(): + require_host_store() + now = _seed_calls() + client = TestClient(auth_proxy.app) + r = client.get("/internal/usage", + params={"caller": "acme", "since_ts": now - 3600, "bucket": "day"}, + headers={"x-internal-secret": "s3cret"}) + assert r.status_code == 200 + data = r.json() + assert data["caller"] == "acme" + assert data["runs"] == 2 and data["errors"] == 1 + assert data["tokens_in"] == 80 and data["tokens_out"] == 30 + assert data["tokens_cached"] == 25 and data["tokens_total"] == 110 + assert data["cost_usd"] == pytest.approx(0.01) + assert data["window"]["since_ts"] == now - 3600 + assert sum(b["runs"] for b in data["buckets"]) == 2 + + # missing caller -> 400, and no cross-tenant bleed + assert client.get("/internal/usage", + headers={"x-internal-secret": "s3cret"}).status_code == 400 + + +def test_internal_usage_recent_scopes_to_caller(): + require_host_store() + _seed_calls() + r = TestClient(auth_proxy.app).get("/internal/usage/recent", + params={"caller": "acme", "limit": 10}, + headers={"x-internal-secret": "s3cret"}) + assert r.status_code == 200 + calls = r.json()["calls"] + assert len(calls) == 2 + assert {c["status"] for c in calls} == {200, 500} + assert all(c["key_sha256_prefix"] == "c" * 12 for c in calls) + assert all("consumer_sha" not in c for c in calls) # only the prefix leaves + assert calls[0]["latency_ms"] == 100.0 + + +def test_internal_usage_is_not_proxied_upstream(monkeypatch): + """Regression: the /internal router must match BEFORE the catch-all proxy.""" + upstream = _upstream(monkeypatch) + r = TestClient(auth_proxy.app).get("/internal/usage", + headers={"x-internal-secret": "s3cret"}) + assert r.status_code == 400 # caller_required, answered locally + assert upstream.requests == [] diff --git a/tests/test_control_plane_client.py b/tests/test_control_plane_client.py new file mode 100644 index 0000000..a3b0116 --- /dev/null +++ b/tests/test_control_plane_client.py @@ -0,0 +1,249 @@ +"""Unit tests for control_plane_client (no Postgres, no network). + +The HTTP boundary is a fake async client monkeypatched onto the module-level +`_client` (suite convention — no respx). Feature flags are monkeypatched module +attributes; every test starts from reset_for_tests(). +""" +from __future__ import annotations + +import asyncio +import os +import sys +from pathlib import Path + +import httpx +import pytest + +ROOT = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(ROOT)) + +import control_plane_client as cpc # noqa: E402 + + +class _FakeResp: + def __init__(self, status_code: int, payload): + self.status_code = status_code + self._payload = payload + + def json(self): + if isinstance(self._payload, Exception): + raise self._payload + return self._payload + + def raise_for_status(self): + if self.status_code >= 400: + raise httpx.HTTPStatusError("boom", request=None, response=None) + + +class _FakeClient: + """Scripted responses per URL prefix; records every call.""" + + def __init__(self, responses): + # responses: list of _FakeResp | Exception, consumed in order + self.responses = list(responses) + self.calls: list[dict] = [] + + async def get(self, url, params=None, headers=None): + self.calls.append({"url": url, "params": params, "headers": headers}) + item = self.responses.pop(0) + if isinstance(item, Exception): + raise item + return item + + +@pytest.fixture(autouse=True) +def _clean(monkeypatch): + cpc.reset_for_tests() + monkeypatch.setattr(cpc, "CONTROL_PLANE_URL", "http://cp.test") + monkeypatch.setattr(cpc, "CONTROL_PLANE_INTERNAL_SECRET", "s3cret") + yield + cpc.reset_for_tests() + + +def _install(monkeypatch, responses) -> _FakeClient: + fake = _FakeClient(responses) + monkeypatch.setattr(cpc, "_client", fake) + return fake + + +def _active(consumer="acme", tenant_id=7, rate=None, burst=None): + return {"active": True, "consumer": consumer, "tenant_id": tenant_id, + "rate_per_min": rate, "burst": burst} + + +DIGEST = "ab" * 32 + + +# ---- resolve_key ------------------------------------------------------------- + +def test_feature_off_returns_none_without_http(monkeypatch): + monkeypatch.setattr(cpc, "CONTROL_PLANE_URL", "") + fake = _install(monkeypatch, []) + out = asyncio.run(cpc.resolve_key(DIGEST)) + assert out is None + assert fake.calls == [] + + +def test_resolve_fetches_then_serves_from_cache(monkeypatch): + fake = _install(monkeypatch, [_FakeResp(200, _active(rate=60, burst=20))]) + first = asyncio.run(cpc.resolve_key(DIGEST)) + assert first.active and first.consumer == "acme" and first.tenant_id == 7 + assert first.rate_per_min == 60 and first.burst == 20 + second = asyncio.run(cpc.resolve_key(DIGEST)) + assert second is first + assert len(fake.calls) == 1 + call = fake.calls[0] + assert call["url"].endswith("/internal/keys/resolve") + assert call["params"] == {"sha256": DIGEST} + assert call["headers"] == {"x-internal-secret": "s3cret"} + + +def test_negative_answer_is_cached(monkeypatch): + fake = _install(monkeypatch, [_FakeResp(200, {"active": False})]) + first = asyncio.run(cpc.resolve_key(DIGEST)) + assert first is not None and first.active is False + second = asyncio.run(cpc.resolve_key(DIGEST)) + assert second.active is False + assert len(fake.calls) == 1 + + +def test_ttl_expiry_revalidates_and_definitive_no_replaces_positive(monkeypatch): + fake = _install(monkeypatch, [_FakeResp(200, _active()), + _FakeResp(200, {"active": False})]) + assert asyncio.run(cpc.resolve_key(DIGEST)).active is True + monkeypatch.setattr(cpc, "RESOLVE_TTL_S", 0.0) # everything positive expired + out = asyncio.run(cpc.resolve_key(DIGEST)) + assert out.active is False # revoked upstream wins + assert len(fake.calls) == 2 + + +def test_stale_grace_serves_positive_on_transport_error_only(monkeypatch): + fake = _install(monkeypatch, [_FakeResp(200, _active()), + httpx.ConnectError("down"), + httpx.ConnectError("down")]) + assert asyncio.run(cpc.resolve_key(DIGEST)).active is True + monkeypatch.setattr(cpc, "RESOLVE_TTL_S", 0.0) + # CP unreachable inside the grace window -> stale positive keeps working + assert asyncio.run(cpc.resolve_key(DIGEST)).consumer == "acme" + # grace exhausted -> None (caller 401s) + monkeypatch.setattr(cpc, "RESOLVE_STALE_GRACE_S", 0.0) + assert asyncio.run(cpc.resolve_key(DIGEST)) is None + assert len(fake.calls) == 3 + + +def test_negative_entries_get_no_grace(monkeypatch): + _install(monkeypatch, [_FakeResp(200, {"active": False}), + httpx.ConnectError("down")]) + assert asyncio.run(cpc.resolve_key(DIGEST)).active is False + monkeypatch.setattr(cpc, "NEGATIVE_TTL_S", 0.0) + assert asyncio.run(cpc.resolve_key(DIGEST)) is None + + +def test_5xx_counts_as_unreachable(monkeypatch): + _install(monkeypatch, [_FakeResp(200, _active()), _FakeResp(500, {})]) + assert asyncio.run(cpc.resolve_key(DIGEST)).active is True + monkeypatch.setattr(cpc, "RESOLVE_TTL_S", 0.0) + assert asyncio.run(cpc.resolve_key(DIGEST)).consumer == "acme" # stale grace + + +def test_single_flight_coalesces_concurrent_resolves(monkeypatch): + fake = _FakeClient([]) + started = asyncio.Event() + + async def slow_get(url, params=None, headers=None): + fake.calls.append({"url": url}) + started.set() + await asyncio.sleep(0.02) + return _FakeResp(200, _active()) + + fake.get = slow_get + monkeypatch.setattr(cpc, "_client", fake) + + async def run(): + return await asyncio.gather(cpc.resolve_key(DIGEST), cpc.resolve_key(DIGEST)) + + a, b = asyncio.run(run()) + assert a.consumer == b.consumer == "acme" + assert len(fake.calls) == 1 + + +def test_malformed_body_is_a_definitive_negative(monkeypatch): + _install(monkeypatch, [_FakeResp(200, ValueError("not json"))]) + out = asyncio.run(cpc.resolve_key(DIGEST)) + assert out is not None and out.active is False + + +# ---- tenant_env -------------------------------------------------------------- + +def test_tenant_env_filters_through_allowlist_and_caches(monkeypatch): + fake = _install(monkeypatch, [_FakeResp(200, {"env": { + "OPENAI_API_KEY": "sk-tenant", "EVIL_PATH_OVERRIDE": "x", + "ANTHROPIC_API_KEY": ""}})]) + env = asyncio.run(cpc.tenant_env(7)) + assert env == {"OPENAI_API_KEY": "sk-tenant"} # allowlist + empty dropped + assert asyncio.run(cpc.tenant_env(7)) == env + assert len(fake.calls) == 1 + assert fake.calls[0]["url"].endswith("/internal/tenants/7/provider-env") + + +def test_tenant_env_fail_soft_to_platform_keys(monkeypatch): + _install(monkeypatch, [httpx.ConnectError("down")]) + assert asyncio.run(cpc.tenant_env(9)) == {} + + +def test_tenant_env_stale_grace_then_empty(monkeypatch): + _install(monkeypatch, [_FakeResp(200, {"env": {"OPENAI_API_KEY": "sk-t"}}), + httpx.ConnectError("down"), + httpx.ConnectError("down")]) + assert asyncio.run(cpc.tenant_env(7)) == {"OPENAI_API_KEY": "sk-t"} + monkeypatch.setattr(cpc, "TENANT_ENV_TTL_S", 0.0) + assert asyncio.run(cpc.tenant_env(7)) == {"OPENAI_API_KEY": "sk-t"} # grace + monkeypatch.setattr(cpc, "TENANT_ENV_STALE_GRACE_S", 0.0) + assert asyncio.run(cpc.tenant_env(7)) == {} + + +# ---- env_get / context isolation --------------------------------------------- + +def test_env_get_prefers_active_tenant_map_then_process_env(monkeypatch): + monkeypatch.setenv("OPENAI_API_KEY", "sk-platform") + assert cpc.env_get("OPENAI_API_KEY") == "sk-platform" + token = cpc.activate_tenant_env({"OPENAI_API_KEY": "sk-tenant"}) + try: + assert cpc.env_get("OPENAI_API_KEY") == "sk-tenant" + assert cpc.env_get("OPENROUTER_API_KEY") == os.environ.get("OPENROUTER_API_KEY") + finally: + cpc.reset_tenant_env(token) + assert cpc.env_get("OPENAI_API_KEY") == "sk-platform" + + +def test_env_get_isolation_across_tasks(monkeypatch): + monkeypatch.setenv("OPENAI_API_KEY", "sk-platform") + seen: dict[str, str | None] = {} + + async def tenant_task(name: str, key: str | None): + if key is not None: + cpc.activate_tenant_env({"OPENAI_API_KEY": key}) + await asyncio.sleep(0.01) + + async def child(): + seen[name] = cpc.env_get("OPENAI_API_KEY") + + # tasks created AFTER activation copy the context (streaming/flow nodes) + await asyncio.create_task(child()) + + async def run(): + await asyncio.gather(tenant_task("a", "sk-a"), tenant_task("b", "sk-b"), + tenant_task("none", None)) + + asyncio.run(run()) + assert seen == {"a": "sk-a", "b": "sk-b", "none": "sk-platform"} + + +# ---- internal_secret_ok -------------------------------------------------------- + +def test_internal_secret_ok(monkeypatch): + assert cpc.internal_secret_ok({"x-internal-secret": "s3cret"}) is True + assert cpc.internal_secret_ok({"x-internal-secret": "wrong"}) is False + assert cpc.internal_secret_ok({}) is False + monkeypatch.setattr(cpc, "CONTROL_PLANE_INTERNAL_SECRET", "") + assert cpc.internal_secret_ok({"x-internal-secret": ""}) is False diff --git a/tests/test_host_store.py b/tests/test_host_store.py index f220c17..3cc09a3 100644 --- a/tests/test_host_store.py +++ b/tests/test_host_store.py @@ -250,3 +250,44 @@ def test_route_stats_window_excludes_old_observations(store): seed_route_obs("p", "m", "stale", ok=True, ts=now - 20 * 60 * 1000) # 20 min ago assert set(store.route_stats(window_ms=15 * 60 * 1000)) == {"p|m|fresh"} assert set(store.route_stats(window_ms=30 * 60 * 1000)) == {"p|m|fresh", "p|m|stale"} + + +# ---- control-plane metering helpers --------------------------------------- + +def test_usage_totals_window_and_caller(store): + store.insert_call(_row(usage_event_id="a", caller="acme", ts=2_000_000, + tokens_in=70, tokens_out=30, tokens_total=100, + cost_usd=0.01)) + store.insert_call(_row(usage_event_id="b", caller="acme", ts=2_000_100, + status=500, tokens_in=10, tokens_out=0, + tokens_total=0, cost_usd=None)) + store.insert_call(_row(usage_event_id="c", caller="other", ts=2_000_100, + tokens_in=999, tokens_out=999, tokens_total=1998, + cost_usd=5.0)) + store.insert_call(_row(usage_event_id="old", caller="acme", ts=1_000)) + t = store.usage_totals(since_ts=1_500_000, caller="acme") + assert t["requests"] == 2 and t["errors"] == 1 + assert t["tokens_in"] == 80 and t["tokens_out"] == 30 + # tokens_total falls back to in+out when the stamped total is 0 + assert t["tokens_total"] == 110 + assert t["cost_usd"] == pytest.approx(0.01) + assert t["priced"] == 1 # the NULL-cost row is not counted as priced + + +def test_usage_totals_includes_cached_tokens_and_fails_soft(store, monkeypatch): + store.insert_call(_row(tokens_cached=40, caller="acme")) + assert store.usage_totals(caller="acme")["tokens_cached"] == 40 + monkeypatch.setattr(hs, "_get_pool", lambda: (_ for _ in ()).throw(RuntimeError("down"))) + zeros = store.usage_totals(caller="acme") + assert zeros["requests"] == 0 and zeros["cost_usd"] == 0.0 + assert set(zeros) == {"requests", "errors", "tokens_in", "tokens_out", + "tokens_cached", "tokens_total", "cost_usd", "priced"} + + +def test_recent_calls_caller_filter(store): + store.insert_call(_row(usage_event_id="a", caller="acme", ts=1)) + store.insert_call(_row(usage_event_id="b", caller="other", ts=2)) + store.insert_call(_row(usage_event_id="c", caller="acme", ts=3)) + rows = store.recent_calls(caller="acme") + assert [r["usage_event_id"] for r in rows] == ["c", "a"] # newest first + assert [r["usage_event_id"] for r in store.recent_calls()][0] == "c" diff --git a/tests/test_tenant_env_shim.py b/tests/test_tenant_env_shim.py new file mode 100644 index 0000000..a6cb87e --- /dev/null +++ b/tests/test_tenant_env_shim.py @@ -0,0 +1,223 @@ +"""Per-tenant BYO provider credentials through the router (shim + adapters). + +Covers the full chain: the trusted x-llm-router-tenant header -> _activate_tenant +-> ContextVar -> control_plane_client.env_get inside the provider call, including +task-context propagation (streaming/timeout paths create tasks) and isolation +between concurrent tenants. The control plane is faked at the HTTP boundary. +""" +from __future__ import annotations + +import asyncio +import sys +from pathlib import Path + +import pytest +from fastapi.testclient import TestClient + +ROOT = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(ROOT)) + +import control_plane_client as cpc # noqa: E402 +from llm_router_host import LLMRouterHost # noqa: E402 +from shim import create_app # noqa: E402 + + +class _FakeCPResp: + status_code = 200 + + def __init__(self, payload): + self._payload = payload + + def json(self): + return self._payload + + def raise_for_status(self): + pass + + +class _FakeCPClient: + def __init__(self, env_by_tenant): + self.env_by_tenant = env_by_tenant + self.calls: list[str] = [] + + async def get(self, url, params=None, headers=None): + self.calls.append(url) + tenant_id = int(url.rstrip("/").split("/")[-2]) + return _FakeCPResp({"env": self.env_by_tenant.get(tenant_id, {})}) + + +@pytest.fixture(autouse=True) +def _clean(monkeypatch): + cpc.reset_for_tests() + monkeypatch.setattr(cpc, "CONTROL_PLANE_URL", "http://cp.test") + monkeypatch.setattr(cpc, "CONTROL_PLANE_INTERNAL_SECRET", "s3cret") + yield + cpc.reset_for_tests() + + +@pytest.fixture +def host(): + h = LLMRouterHost( + router_path=ROOT / "core" / "router.lua", + config_path=ROOT / "core" / "config.example.lua", + metrics_path=ROOT / "core" / "metrics.example.lua", + now_ms=lambda: 1_000_000, + # The custom call hook below would otherwise turn on auth enforcement + # and pre-disable every example provider (their env keys are unset here). + enforce_provider_auth=False, + ) + h.init() + return h + + +def _ok_result(): + return {"ok": True, "latency_ms": 10, + "response": {"text": "hi", "tool_calls": None, "finish_reason": "stop", + "tokens_in": 7, "tokens_out": 3, "tokens_total": 10, + "raw_model": "mock-model-id"}} + + +def _capture_hook(seen: list): + """An async provider hook that records what env_get resolves AT CALL TIME — + i.e. inside host.execute_async, past any create_task boundaries.""" + + async def hook(request: dict) -> dict: + await asyncio.sleep(0.005) # let concurrent requests interleave + seen.append(cpc.env_get("OPENAI_API_KEY")) + return _ok_result() + + return hook + + +def _chat(client, tenant: int | None): + headers = {"x-llm-router-tenant": str(tenant)} if tenant is not None else {} + return client.post("/v1/chat/completions", headers=headers, + json={"model": "profile:default", + "messages": [{"role": "user", "content": "hi"}]}) + + +def test_tenant_header_activates_byo_key(host, monkeypatch): + monkeypatch.setenv("OPENAI_API_KEY", "sk-platform") + monkeypatch.setattr(cpc, "_client", _FakeCPClient({7: {"OPENAI_API_KEY": "sk-tenant"}})) + seen: list = [] + host.set_async_call_hook(_capture_hook(seen)) + client = TestClient(create_app(host, default_profile="default")) + assert _chat(client, tenant=7).status_code == 200 + assert seen == ["sk-tenant"] + + +def test_no_header_uses_platform_key(host, monkeypatch): + monkeypatch.setenv("OPENAI_API_KEY", "sk-platform") + fake = _FakeCPClient({}) + monkeypatch.setattr(cpc, "_client", fake) + seen: list = [] + host.set_async_call_hook(_capture_hook(seen)) + client = TestClient(create_app(host, default_profile="default")) + assert _chat(client, tenant=None).status_code == 200 + assert seen == ["sk-platform"] + assert fake.calls == [] + + +def test_feature_off_ignores_header(host, monkeypatch): + monkeypatch.setenv("OPENAI_API_KEY", "sk-platform") + monkeypatch.setattr(cpc, "CONTROL_PLANE_URL", "") + fake = _FakeCPClient({7: {"OPENAI_API_KEY": "sk-tenant"}}) + monkeypatch.setattr(cpc, "_client", fake) + seen: list = [] + host.set_async_call_hook(_capture_hook(seen)) + client = TestClient(create_app(host, default_profile="default")) + assert _chat(client, tenant=7).status_code == 200 + assert seen == ["sk-platform"] + assert fake.calls == [] + + +def test_concurrent_tenants_are_isolated(host, monkeypatch): + monkeypatch.setenv("OPENAI_API_KEY", "sk-platform") + monkeypatch.setattr(cpc, "_client", _FakeCPClient({ + 1: {"OPENAI_API_KEY": "sk-one"}, 2: {"OPENAI_API_KEY": "sk-two"}})) + seen: list = [] + host.set_async_call_hook(_capture_hook(seen)) + app = create_app(host, default_profile="default") + + # Drive the ASGI app directly so both requests share one event loop and + # genuinely interleave inside the provider hook. + import httpx + + async def run(): + transport = httpx.ASGITransport(app=app) + async with httpx.AsyncClient(transport=transport, base_url="http://t") as c: + body = {"model": "profile:default", + "messages": [{"role": "user", "content": "hi"}]} + r1, r2, r3 = await asyncio.gather( + c.post("/v1/chat/completions", json=body, + headers={"x-llm-router-tenant": "1"}), + c.post("/v1/chat/completions", json=body, + headers={"x-llm-router-tenant": "2"}), + c.post("/v1/chat/completions", json=body), + ) + assert r1.status_code == r2.status_code == r3.status_code == 200 + + asyncio.run(run()) + assert sorted(seen, key=str) == ["sk-one", "sk-platform", "sk-two"] + + +def test_streaming_request_carries_tenant_env(host, monkeypatch): + """stream:true goes through asyncio.create_task in the shim — the context + must propagate into the task.""" + monkeypatch.setenv("OPENAI_API_KEY", "sk-platform") + monkeypatch.setattr(cpc, "_client", _FakeCPClient({7: {"OPENAI_API_KEY": "sk-tenant"}})) + seen: list = [] + + async def streaming_call(request, emit): + seen.append(cpc.env_get("OPENAI_API_KEY")) + emit({"delta": "hi"}) + return _ok_result() + + host.set_async_call_hook(_capture_hook(seen)) + client = TestClient(create_app(host, default_profile="default", + streaming_call=streaming_call)) + r = client.post("/v1/chat/completions", + headers={"x-llm-router-tenant": "7"}, + json={"model": "profile:default", "stream": True, + "messages": [{"role": "user", "content": "hi"}]}) + assert r.status_code == 200 + assert "sk-tenant" in seen + + +def test_adapter_authorization_header_uses_tenant_key(monkeypatch): + """End of the chain: the OpenAI-compatible adapter builds Authorization from + control_plane_client.env_get, so an active tenant map changes the wire key.""" + from provider_adapters.openai_compatible import make_async_call_provider + + captured: dict = {} + + class _FakeHTTPResp: + status_code = 200 + + def json(self): + return {"choices": [{"message": {"content": "ok"}, + "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, + "total_tokens": 2}} + + class _FakeHTTPClient: + async def post(self, url, json=None, headers=None, timeout=None): + captured["headers"] = headers + return _FakeHTTPResp() + + call = make_async_call_provider(env_get=cpc.env_get, client=_FakeHTTPClient()) + request = {"provider_id": "openai", "served_model_id": "gpt-x", + "base_url": "https://api.test/v1", + "auth": {"kind": "bearer", "env": "OPENAI_API_KEY"}, + "messages": [{"role": "user", "content": "hi"}]} + monkeypatch.setenv("OPENAI_API_KEY", "sk-platform") + token = cpc.activate_tenant_env({"OPENAI_API_KEY": "sk-tenant"}) + try: + result = asyncio.run(call(dict(request))) + finally: + cpc.reset_tenant_env(token) + assert result["ok"] is True + assert captured["headers"]["Authorization"] == "Bearer sk-tenant" + + result = asyncio.run(call(dict(request))) + assert captured["headers"]["Authorization"] == "Bearer sk-platform"