Enforce Cloud key budgets and rate limits; add usage export and budget reads

auth_proxy.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import asyncio
+import concurrent.futures
 import base64
 import contextlib
 import hashlib
@@ -53,6 +54,22 @@
 MAX_INFLIGHT_REQUESTS = max(0, int(os.getenv("MAX_INFLIGHT_REQUESTS", "0")))
 MAX_PENDING_REQUESTS = max(0, int(os.getenv("MAX_PENDING_REQUESTS", "0")))
 CAPACITY_QUEUE_TIMEOUT_S = max(0.0, float(os.getenv("CAPACITY_QUEUE_TIMEOUT_S", "2")))
+
+
+def _reservation_usd() -> float:
+    raw = os.getenv("CLOUD_BUDGET_RESERVATION_USD", "0.05")
+    try:
+        value = float(raw)
+    except (TypeError, ValueError):
+        return 0.05
+    return value if math.isfinite(value) and value > 0 else 0.05
+
+
+# Held against a control-plane key subject's monthly budget per in-flight call
+# and replaced by the actual cost when the call is recorded; overshoot is
+# bounded by concurrent calls x (actual cost - reservation).
+CLOUD_BUDGET_RESERVATION_USD = _reservation_usd()
+KEY_BUDGET_RESERVATION_TTL_S = 900.0
 UPSTREAM_MAX_CONNECTIONS = max(1, int(os.getenv("UPSTREAM_MAX_CONNECTIONS", "100")))
 UPSTREAM_MAX_KEEPALIVE_CONNECTIONS = max(
     1, int(os.getenv("UPSTREAM_MAX_KEEPALIVE_CONNECTIONS", "40")))
@@ -791,6 +808,35 @@ def _rate_ok(caller: str, meta: dict[str, Any] | None = None) \
     return host_store.consume_rate_token(caller, rate_per_min, burst)
 
 
+# Settles run off the event loop: a slow store must not stall every request on
+# this worker. Until a settle lands its reservation still counts, which only
+# errs toward rejecting; if it fails, the reservation expires (TTL).
+_SETTLE_POOL = concurrent.futures.ThreadPoolExecutor(max_workers=4, thread_name_prefix="budget-settle")
+_pending_settles: set[concurrent.futures.Future] = set()
+SETTLE_RETRY_DELAYS_S = (0.5, 2.0, 5.0, 15.0, 60.0)
+
+
+def _submit_settle(reservation_id: str, cost: float | None, caller: str) -> None:
+    def settle() -> None:
+        # Retried so a brief store outage does not lose the call's cost; a
+        # reservation swept after its TTL still books its cost once.
+        for delay in SETTLE_RETRY_DELAYS_S:
+            _settled, store_ok = host_store.settle_subject_budget(reservation_id, cost)
+            if store_ok:
+                return
+            _metric_store_error("key_budget")
+            time.sleep(delay)
+        _log({"event": "key_budget_settle_failed", "caller": caller})
+    future = _SETTLE_POOL.submit(settle)
+    _pending_settles.add(future)
+    future.add_done_callback(_pending_settles.discard)
+
+
+def drain_settles(timeout: float = 10.0) -> None:
+    """Wait for in-flight settles (tests and graceful shutdown)."""
+    concurrent.futures.wait(list(_pending_settles), timeout=timeout)
+
+
 def _requested_route_from(path: str, body: bytes | None) -> str | None:
     route = None
     if body:
@@ -1296,6 +1342,7 @@ async def startup() -> None:
 @app.on_event("shutdown")
 async def shutdown() -> None:
     global _probe_task
+    await asyncio.to_thread(drain_settles)
     if _probe_task:
         _probe_task.cancel()
         with contextlib.suppress(asyncio.CancelledError):
@@ -4026,6 +4073,31 @@ async def proxy(path: str, request: Request) -> Response:
             headers={"Retry-After": str(max(1, math.ceil(retry_after_s)))},
             content={"error": {"message": "caller rate limit exceeded", "type": "rate_limit_error", "code": "caller_rate_limit"}})
 
+    # Control-plane key limits ride the (uncached) route resolution. A subject
+    # (`g:<group>` / `k:<key>`) may span keys, so its bucket and budget are
+    # shared by every key in it, inside this caller's namespace.
+    key_identity = (published_route or {}).get("key")
+    key_limits = ((published_route or {}).get("limits") or {}) if key_identity else {}
+    if key_limits.get("rate_per_min"):
+        subject_rate = int(key_limits["rate_per_min"])
+        rate_allowed, rate_store_ok, retry_after_s = await asyncio.to_thread(
+            host_store.consume_rate_token, f"{caller}|{key_identity['subject']}",
+            subject_rate, int(key_limits.get("burst") or subject_rate))
+        if not rate_store_ok:
+            _metric_store_error("rate_limit")
+            _record_reject(reason="key_rate_limit_unavailable", path="/" + path,
+                           caller=caller, status=503)
+            return JSONResponse(status_code=503, content={"error": {
+                "message": "key rate-limit state is temporarily unavailable",
+                "type": "server_error", "code": "key_rate_limit_unavailable"}})
+        if not rate_allowed:
+            _record_reject(reason="key_rate_limit", path="/" + path, caller=caller, status=429)
+            _log({"event": "reject", "reason": "key_rate_limit", "caller": caller, "path": "/" + path})
+            return JSONResponse(
+                status_code=429,
+                headers={"Retry-After": str(max(1, math.ceil(retry_after_s)))},
+                content={"error": {"message": "key rate limit exceeded", "type": "rate_limit_error", "code": "key_rate_limit"}})
+
     if path == "api/show" and request.method.upper() == "POST":
         requested_model = None
         if body:
@@ -4060,7 +4132,50 @@ async def proxy(path: str, request: Request) -> Response:
             "message": "Trusted router bridge requires HTTPS", "type": "server_error",
             "code": "bridge_transport_unavailable"}})
     assert _client is not None
-    if not await _capacity_acquire():
+
+    # Budget reservation is the LAST admission step before the capacity permit;
+    # from here on every exit path must settle it (see _settle_reservation).
+    reservation: dict[str, Any] | None = None
+    budget_limit = key_limits.get("monthly_budget_usd")
+    if budget_limit is not None:
+        reservation = {"id": secrets.token_hex(16), "done": False,
+                       # A budget below one reservation still admits a first call.
+                       "amount": min(CLOUD_BUDGET_RESERVATION_USD, float(budget_limit))}
+        allowed, store_ok, spent_usd, _reserved = await asyncio.to_thread(
+            host_store.reserve_subject_budget, int(auth["tenant_id"]), key_identity["subject"],
+            time.strftime("%Y-%m", time.gmtime()), float(budget_limit),
+            reservation["amount"], reservation["id"],
+            KEY_BUDGET_RESERVATION_TTL_S)
+        if not store_ok:
+            _metric_store_error("key_budget")
+            _record_reject(reason="key_budget_unavailable", path="/" + path,
+                           caller=caller, status=503, route=requested_route)
+            return JSONResponse(status_code=503, content={"error": {
+                "message": "key budget state is temporarily unavailable",
+                "type": "server_error", "code": "key_budget_unavailable"}})
+        if not allowed:
+            _record_reject(reason="key_budget_exhausted", path="/" + path,
+                           caller=caller, status=402, route=requested_route)
+            return JSONResponse(status_code=402, content={"error": {
+                "message": "key budget exhausted", "type": "budget_error",
+                "code": "key_budget_exhausted",
+                "budget_usd": float(budget_limit), "spent_usd": round(spent_usd, 6)}})
+
+    def _settle_reservation(cost: float | None) -> None:
+        # Idempotent per request; the store is idempotent per reservation too.
+        # A failed settle leaves the reservation to expire (TTL), never to leak.
+        if reservation is None or reservation["done"]:
+            return
+        reservation["done"] = True
+        _submit_settle(reservation["id"], cost, caller)
+
+    try:
+        acquired = await _capacity_acquire()
+    except BaseException:
+        _settle_reservation(None)
+        raise
+    if not acquired:
+        _settle_reservation(None)
         _record_reject(reason="router_overloaded", path="/" + path,
                        caller=caller, status=503, route=requested_route)
         return JSONResponse(
@@ -4070,33 +4185,6 @@ async def proxy(path: str, request: Request) -> Response:
                 "message": "router capacity is temporarily exhausted; retry shortly",
                 "type": "server_error", "code": "router_overloaded"}})
 
-    upstream_url = f"{UPSTREAM}/{path}"
-    if request.url.query:
-        upstream_url += f"?{request.url.query}"
-
-    headers = {
-        k: v for k, v in request.headers.items()
-        if k.lower() not in {"authorization", "host", "connection", "content-length",
-                             "x-llm-router-tenant", "x-internal-secret",
-                             "x-unhardcoded-scope-version", "x-unhardcoded-project", "x-unhardcoded-environment",
-                             "x-unhardcoded-route", "x-unhardcoded-revision", "x-unhardcoded-policy-id", "x-unhardcoded-preference"}
-    }
-    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"])
-        headers["x-internal-secret"] = control_plane_client.CONTROL_PLANE_INTERNAL_SECRET
-        if auth.get("scope_version") == 2:
-            headers["x-unhardcoded-scope-version"] = "2"
-            headers["x-unhardcoded-project"] = str(auth["project_id"])
-            headers["x-unhardcoded-environment"] = str(auth["environment_id"])
-        headers["x-unhardcoded-route"] = requested_route
-        headers["x-unhardcoded-revision"] = str(published_route["revision"])
-        headers["x-unhardcoded-policy-id"] = published_route["policy_id"]
-        headers['x-unhardcoded-preference'] = published_route.get('routing_preference', 'default')
-
     status = 502
     provider = None
     model_family = None
@@ -4124,21 +4212,59 @@ def _finish():
             return
         capacity_released = True
         latency_ms = round((time.perf_counter() - started) * 1000, 1)
-        if published_route:
-            decision_trace = {**(decision_trace or {}), 'route': requested_route,
-                              'route_revision': published_route['revision'],
-                              'policy_id': (decision_trace or {}).get('policy_id') or published_route['policy_id'],
-                              'routing_preference': published_route.get('routing_preference', 'default')}
-            if auth.get('scope_version') == 2:
-                decision_trace.update(project_id=auth['project_id'], environment_id=auth['environment_id'])
         try:
+            if published_route:
+                decision_trace = {**(decision_trace or {}), 'route': requested_route,
+                                  'route_revision': published_route['revision'],
+                                  'policy_id': (decision_trace or {}).get('policy_id') or published_route['policy_id'],
+                                  'routing_preference': published_route.get('routing_preference', 'default')}
+                if auth.get('scope_version') == 2:
+                    decision_trace.update(project_id=auth['project_id'], environment_id=auth['environment_id'])
+                if key_identity:
+                    decision_trace.update(key_id=key_identity['id'], subject=key_identity['subject'])
             _record_request(caller=caller, method=request.method, path="/" + path, status=status, latency_ms=latency_ms, provider=provider, model_family=model_family, served_model_id=served_model_id, served_by=served_by, requested_model=requested_model, session=session_id, tokens_in=tokens_in, tokens_out=tokens_out, tokens_total=tokens_total, tokens_cached=tokens_cached, cost_usd=cost_usd, cost_basis=cost_basis, decision_trace=decision_trace, error_type=error_type, error_code=error_code, error_message=error_message, key_sha256=auth.get("digest"))
             _metric_request(status, latency_ms)
             _log({"event": "request", "caller": caller, "method": request.method, "path": "/" + path, "status": status, "latency_ms": latency_ms, "provider": provider, "model_family": model_family})
         finally:
-            _capacity_release()
+            try:
+                # Not the droppable ledger queue: a dedicated settle, retried.
+                # A successful call whose cost is unknown (a stream the client
+                # closed before its usage chunk, or an unpriced model) books the
+                # reservation, never $0; failed calls book only what they report.
+                _settle_reservation(cost_usd if cost_usd is not None or status >= 400
+                                    else reservation["amount"] if reservation else None)
+            finally:
+                _capacity_release()
 
     try:
+        # Built inside the try so a failure here still reaches _finish, which
+        # releases the capacity permit and the budget reservation.
+        upstream_url = f"{UPSTREAM}/{path}"
+        if request.url.query:
+            upstream_url += f"?{request.url.query}"
+
+        headers = {
+            k: v for k, v in request.headers.items()
+            if k.lower() not in {"authorization", "host", "connection", "content-length",
+                                 "x-llm-router-tenant", "x-internal-secret",
+                                 "x-unhardcoded-scope-version", "x-unhardcoded-project", "x-unhardcoded-environment",
+                                 "x-unhardcoded-route", "x-unhardcoded-revision", "x-unhardcoded-policy-id", "x-unhardcoded-preference"}
+        }
+        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"])
+            headers["x-internal-secret"] = control_plane_client.CONTROL_PLANE_INTERNAL_SECRET
+            if auth.get("scope_version") == 2:
+                headers["x-unhardcoded-scope-version"] = "2"
+                headers["x-unhardcoded-project"] = str(auth["project_id"])
+                headers["x-unhardcoded-environment"] = str(auth["environment_id"])
+            headers["x-unhardcoded-route"] = requested_route
+            headers["x-unhardcoded-revision"] = str(published_route["revision"])
+            headers["x-unhardcoded-policy-id"] = published_route["policy_id"]
+            headers['x-unhardcoded-preference'] = published_route.get('routing_preference', 'default')
         if body:
             try:
                 parsed_body = json.loads(body.decode("utf-8"))

control_plane_client.py

@@ -12,6 +12,10 @@
           "rate_per_min": int|null, "burst": int|null}
   GET {CONTROL_PLANE_URL}/internal/tenants/<id>/provider-env
       -> {"env": {ENV_NAME: secret, ...}}
+  GET {CONTROL_PLANE_URL}/internal/tenants/<t>[/projects/<p>/environments/<e>]/routes/<slug>
+      -> the published route, optionally with the key's limits:
+         "key": {"id", "subject": "g:<n>"|"k:<n>", "labels"},
+         "limits": {"monthly_budget_usd", "rate_per_min", "burst"}
 
 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
@@ -25,7 +29,9 @@
 import hmac
 import json
 import logging
+import math
 import os
+import re
 import time
 from dataclasses import dataclass
 from typing import Any, Mapping
@@ -333,8 +339,68 @@ def _validate_scope(data, tenant_id, project_id, environment_id):
         raise ValueError("Mismatched project/environment scope")
 
 
+_SUBJECT_RE = re.compile(r"[gk]:[1-9][0-9]{0,17}")
+_BUDGET_RE = re.compile(r"[0-9]{1,12}(\.[0-9]{1,12})?")
+_LABELS_MAX, _LABEL_KEY_MAX, _LABEL_VALUE_MAX = 32, 64, 256
+
+
+def _positive_int_or_none(value: Any) -> int | None:
+    if value is None:
+        return None
+    if type(value) is not int or value <= 0:
+        raise ValueError("invalid limit")
+    return value
+
+
+def _parse_key_limits(data: dict) -> tuple[dict | None, dict | None]:
+    """Strict parse of the route's optional `key` identity and `limits`.
+
+    Both absent (older control plane) -> (None, None): no key limits. Anything
+    present but malformed raises ValueError so the route is UNAVAILABLE — a
+    limit silently dropped would let a capped key spend without a cap."""
+    key, limits = data.get("key"), data.get("limits")
+    if key is None and limits is None:
+        return None, None
+    if not isinstance(key, dict):
+        raise ValueError("limits without a key identity")
+    key_id, subject, labels = key.get("id"), key.get("subject"), key.get("labels")
+    if type(key_id) is not int or key_id <= 0:
+        raise ValueError("invalid key id")
+    if not isinstance(subject, str) or not _SUBJECT_RE.fullmatch(subject):
+        raise ValueError("invalid key subject")
+    labels = {} if labels is None else labels
+    if (not isinstance(labels, dict) or len(labels) > _LABELS_MAX
+            or any(not isinstance(k, str) or not isinstance(v, str) or not k
+                   or len(k) > _LABEL_KEY_MAX or len(v) > _LABEL_VALUE_MAX
+                   for k, v in labels.items())):
+        raise ValueError("invalid key labels")
+    limits = {} if limits is None else limits
+    if not isinstance(limits, dict):
+        raise ValueError("invalid key limits")
+    budget = limits.get("monthly_budget_usd")
+    if budget is not None:
+        # A decimal string ("40.00") or a JSON number; bool is an int subclass.
+        if isinstance(budget, str) and _BUDGET_RE.fullmatch(budget):
+            budget = float(budget)
+        elif type(budget) in (int, float) and math.isfinite(budget):
+            budget = float(budget)
+        else:
+            raise ValueError("invalid monthly budget")
+        if not budget > 0:
+            raise ValueError("invalid monthly budget")
+    rate, burst = _positive_int_or_none(limits.get("rate_per_min")), _positive_int_or_none(limits.get("burst"))
+    return ({"id": key_id, "subject": subject, "labels": dict(labels)},
+            {"monthly_budget_usd": budget, "rate_per_min": rate, "burst": burst})
+
+
 async def resolve_route(tenant_id: int, name: str, *, project_id=None, environment_id=None, key_digest=None) -> dict:
-    """Resolve the published revision on every call, so publish/pause is immediate."""
+    """Resolve the published revision on every call, so publish/pause is immediate.
+
+    The response may carry the calling key's identity and limits (`key`:
+    {id, subject, labels}; `limits`: {monthly_budget_usd, rate_per_min, burst});
+    they are returned normalized (budget as float, absent -> None) and enforced
+    by the ingress. Because this lookup is uncached and keyed by the digest, a
+    revoked scoped key is refused on its very next request."""
     try:
         scoped = project_id is not None or environment_id is not None
         if scoped:
@@ -358,6 +424,7 @@ async def resolve_route(tenant_id: int, name: str, *, project_id=None, environme
                 or not isinstance(data.get("revision"), int) or data["revision"] <= 0
                 or not isinstance(data.get("execution", {}), dict)):
             raise ValueError("invalid route contract")
+        data["key"], data["limits"] = _parse_key_limits(data)
         return data
     except (httpx.HTTPError, ValueError, AttributeError) as exc:
         raise RouteUnavailable("This route is unavailable or has not been published.") from exc