diff --git a/auth_proxy.py b/auth_proxy.py index 0af14ab..40f4330 100644 --- a/auth_proxy.py +++ b/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:` / `k:`) 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")) diff --git a/control_plane_client.py b/control_plane_client.py index 4d018cc..136509a 100644 --- a/control_plane_client.py +++ b/control_plane_client.py @@ -12,6 +12,10 @@ "rate_per_min": int|null, "burst": int|null} GET {CONTROL_PLANE_URL}/internal/tenants//provider-env -> {"env": {ENV_NAME: secret, ...}} + GET {CONTROL_PLANE_URL}/internal/tenants/[/projects/

/environments/]/routes/ + -> the published route, optionally with the key's limits: + "key": {"id", "subject": "g:"|"k:", "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 diff --git a/docs/saas-project-scopes.md b/docs/saas-project-scopes.md index 6483ab9..98ec284 100644 --- a/docs/saas-project-scopes.md +++ b/docs/saas-project-scopes.md @@ -84,3 +84,62 @@ Cloud's production settings also reject plaintext router/ingress URLs. Existing operator-only deployments are unaffected because their control-plane integration is disabled. Existing HTTP control-plane deployments must configure HTTPS or explicitly opt into their existing insecure transport before adopting this release. + +## Key limits, budgets and reconciliation + +The route response may carry the calling key's identity and limits: + +```json +"key": {"id": 42, "subject": "g:7", "labels": {"team": "search"}}, +"limits": {"monthly_budget_usd": "40.00", "rate_per_min": 120, "burst": 40} +``` + +`subject` (`g:` or `k:`) is the unit that limits apply to; any +limit may be `null`. Both objects absent means no key limits (older control +planes). A present but malformed object makes the route unavailable (`503 +route_unavailable`): a limit is never silently dropped. + +After the plan rate check the ingress enforces, in order: + +1. the subject token bucket `"{caller}|{subject}"` (`429 key_rate_limit` with + `Retry-After`); +2. an atomic PostgreSQL reservation of `CLOUD_BUDGET_RESERVATION_USD` (default + 0.05, capped at the budget itself) against `(tenant, subject, UTC month)`, + admitted only if `spent + reserved + reservation <= budget` + (`402 key_budget_exhausted` with `budget_usd` and `spent_usd`); +3. the per-pod capacity permit. + +An unavailable store fails closed (`503 key_rate_limit_unavailable` / +`key_budget_unavailable`). The reservation is settled synchronously when the +call finishes on every path (success, upstream error, exception, capacity +rejection, stream end or client disconnect): it is released and the actual cost +(unknown cost counts as 0) is added. A reservation whose holder died expires +after 15 minutes. Overshoot is bounded by concurrent in-flight calls times +(actual cost - reservation). The recorded `routing_summary` includes `key_id` +and `subject`. + +Because the scoped route lookup is uncached and carries the key digest, a key +revoked in the control plane is refused on its very next request even while the +key resolution itself is still cached. + +Reconciliation endpoints (same shared secret; `caller` and the complete +project/environment pair behave as for `/internal/usage`, filtering before +aggregation or paging): + +* `GET /internal/usage?...&until_ts=&group_by=key|model|route|day|hour` adds a + `groups` list (`key` groups by the full key digest). Windows are half-open + (`since_ts <= ts < until_ts`). Every response carries `watermark_ts` + (`now - LEDGER_WATERMARK_LAG_S`, capped by the oldest call still queued in + this replica's ledger writer); a window with `until_ts <= watermark_ts` is + closed. +* `GET /internal/usage/export?caller=&after_id=&limit=<=5000` returns ledger + rows in id order with `next_after_id` and `watermark_ts`. Rows are delivered + once they were inserted at least `LEDGER_WATERMARK_LAG_S` (default 120) ago, so + the id cursor cannot pass a lower id still committing on another replica. + `watermark_ts` is `now - LEDGER_WATERMARK_LAG_S`, capped by the oldest call + still queued in this replica's ledger writer and by the oldest matching row + not yet delivered: every call with `ts < watermark_ts` is in a page at or + before `next_after_id`. Other replicas' writer backlogs are covered only by + the lag. +* `GET /internal/budgets?tenant_id=&period=YYYY-MM&subject=...` (repeatable, at + most 500) returns `spent_usd` and live `reserved_usd` per subject. diff --git a/host_store.py b/host_store.py index 6f87547..a62b78e 100644 --- a/host_store.py +++ b/host_store.py @@ -310,6 +310,42 @@ def _retention_days() -> int: updated_at BIGINT, PRIMARY KEY (provider_id, model_family) )""", + # Insert time of a ledger row (DB clock, epoch seconds). `ts` is the CALL + # time stamped by the ingress, which can precede the insert by the writer + # queue's backlog; the usage export holds back freshly inserted rows so its + # id cursor never passes a lower id that is still committing on another pod. + # Nullable + a separate SET DEFAULT: a volatile default on ADD COLUMN would + # rewrite the whole table; pre-existing rows stay NULL (= long settled). + "ALTER TABLE calls ADD COLUMN IF NOT EXISTS recorded_at DOUBLE PRECISION", + "ALTER TABLE calls ALTER COLUMN recorded_at SET DEFAULT extract(epoch FROM clock_timestamp())", + # Monthly budget per control-plane key subject (`g:` / `k:`), + # enforced by atomic reservation BEFORE forwarding. `spent_usd` accumulates + # settled actual cost; `reserved_usd` mirrors the sum of live reservations + # (admission recomputes it from the rows, so it cannot drift). Money state: + # never pruned by the retention sweep. + """CREATE TABLE IF NOT EXISTS subject_budget_usage ( + tenant_id BIGINT NOT NULL, + subject TEXT NOT NULL, + period TEXT NOT NULL, + spent_usd DOUBLE PRECISION NOT NULL DEFAULT 0, + reserved_usd DOUBLE PRECISION NOT NULL DEFAULT 0, + updated_at BIGINT NOT NULL, + PRIMARY KEY (tenant_id, subject, period) + )""", + # One row per in-flight call. `released` = swept after expiry (the holder + # died): it no longer counts as reserved, but the row survives so a late + # settle still books the actual cost exactly once. + """CREATE TABLE IF NOT EXISTS subject_budget_reservations ( + id TEXT PRIMARY KEY, + tenant_id BIGINT NOT NULL, + subject TEXT NOT NULL, + period TEXT NOT NULL, + amount_usd DOUBLE PRECISION NOT NULL, + expires_at DOUBLE PRECISION NOT NULL, + released BOOLEAN NOT NULL DEFAULT FALSE + )""", + "CREATE INDEX IF NOT EXISTS idx_subject_budget_reservations_subject" + " ON subject_budget_reservations(tenant_id, subject, period, expires_at)", ] _pool_lock = threading.Lock() @@ -389,7 +425,7 @@ def routing_summary(trace) -> dict | None: or the full candidate catalog in the tenant activity ledger.""" if not isinstance(trace, dict): return None - summary = {key: str(trace[key])[:100] for key in ('route', 'route_revision', 'policy_id', 'routing_preference', 'project_id', 'environment_id') + summary = {key: str(trace[key])[:100] for key in ('route', 'route_revision', 'policy_id', 'routing_preference', 'project_id', 'environment_id', 'key_id', 'subject') if trace.get(key) is not None} summary['attempts'] = [ {key: str(step[key])[:160] for key in ('provider_id', 'model_family', 'error_kind') @@ -489,6 +525,9 @@ def _prune() -> None: (int(now * 1000) - _RETENTION_DAYS * 86400 * 1000,)) conn.execute("DELETE FROM login_history WHERE ts < %s", (int(now) - _RETENTION_DAYS * 86400,)) + # Swept reservations only wait for a late settle; a day is plenty. + conn.execute("DELETE FROM subject_budget_reservations" + " WHERE released AND expires_at < %s", (now - 86400,)) except Exception as exc: # noqa: BLE001 _log.warning("host_store prune failed: %s", exc) @@ -499,16 +538,38 @@ def _prune() -> None: # backlog — best-effort telemetry, so a dropped row is acceptable. The queue holds # thunks so it serves both the call ledger and the per-attempt route observations. _write_q: "queue.Queue" = queue.Queue(maxsize=_WRITE_QUEUE_MAX) +# Call timestamps of ledger rows queued OR being written by this process's +# writer (ts -> count). The ledger is incomplete at the oldest of them, which +# bounds ledger_watermark(); a job leaves only after its write returned. +_pending_ts: dict[int, int] = {} +_pending_lock = threading.Lock() + + +def _pending_add(ts: "int | None") -> None: + if ts is not None: + with _pending_lock: + _pending_ts[ts] = _pending_ts.get(ts, 0) + 1 + + +def _pending_done(ts: "int | None") -> None: + if ts is not None: + with _pending_lock: + left = _pending_ts.get(ts, 0) - 1 + if left > 0: + _pending_ts[ts] = left + else: + _pending_ts.pop(ts, None) def _writer_loop() -> None: while True: - job = _write_q.get() + job, ts = _write_q.get() try: job() except Exception as exc: # noqa: BLE001 — a bad row must not kill the writer _log.warning("host_store background write failed: %s", exc) finally: + _pending_done(ts) _write_q.task_done() @@ -523,18 +584,21 @@ def _writer_loop() -> None: _SYNC_WRITES = bool(os.getenv("HOST_STORE_SYNC_WRITES")) -def _enqueue(job) -> None: +def _enqueue(job, ts: "int | None" = None) -> None: if _SYNC_WRITES: try: job() except Exception as exc: # noqa: BLE001 — mirror the writer loop's tolerance _log.warning("host_store sync write failed: %s", exc) return + _pending_add(ts) try: - _write_q.put_nowait(job) + _write_q.put_nowait((job, ts)) except queue.Full: + _pending_done(ts) _log.warning("host_store: write queue full (%d); dropping row", _WRITE_QUEUE_MAX) except Exception as exc: # noqa: BLE001 — never break a request + _pending_done(ts) _log.warning("host_store enqueue failed: %s", exc) @@ -542,7 +606,31 @@ def insert_call_async(row: dict[str, Any]) -> None: """Record a call WITHOUT blocking the caller — enqueue a SNAPSHOT for the background writer. Drops the row (best-effort) if the queue is full.""" snap = dict(row) - _enqueue(lambda: insert_call(snap)) + # Stamp ts NOW (insert_call would default it at write time) so the pending + # watermark bound and the stored row agree on the call time. + snap["ts"] = int(snap.get("ts") or time.time()) + _enqueue(lambda: insert_call(snap), snap["ts"]) + + +def _watermark_lag_s() -> float: + raw = os.getenv("LEDGER_WATERMARK_LAG_S", "120") + try: + return max(0.0, float(raw)) + except (TypeError, ValueError): + return 120.0 + + +def ledger_watermark(now: "float | None" = None) -> int: + """Epoch second up to which (exclusive) this ledger is complete. + + `now - LEDGER_WATERMARK_LAG_S` covers other replicas' writer queues and + in-flight commits, which this process cannot see; it is further capped by + the oldest call still queued or being written by THIS process's writer + (a backlog longer than the lag would otherwise be settled as complete).""" + mark = int(float(time.time() if now is None else now) - _watermark_lag_s()) + with _pending_lock: + oldest = min(_pending_ts) if _pending_ts else None + return mark if oldest is None else min(mark, oldest) def observe_route_call_async(row: dict[str, Any]) -> None: @@ -1297,7 +1385,8 @@ def cost_by_route(window_s: int = 86_400) -> list[dict[str, Any]]: def _usage_where(since_ts: "int | None" = None, caller: "str | None" = None, caller_is_null: bool = False, consumer_sha: "str | None" = None, provider: "str | None" = None, - model_family: "str | None" = None, project_id=None, environment_id=None) -> "tuple[str, list[Any]]": + model_family: "str | None" = None, project_id=None, environment_id=None, + until_ts: "int | None" = None) -> "tuple[str, list[Any]]": """The shared WHERE for every calls-derived usage read. "all" (since_ts=None) is still bounded to the retention horizon so the read is ALWAYS time-bounded (never a bare table scan): rows older than retention are pruned anyway, so the @@ -1310,6 +1399,8 @@ def _usage_where(since_ts: "int | None" = None, caller: "str | None" = None, if since_ts is not None: floor = max(int(since_ts), floor) clauses, params = ["ts >= %s"], [floor] + if until_ts is not None: # half-open window [since_ts, until_ts) + clauses.append("ts < %s"); params.append(int(until_ts)) if caller is not None: clauses.append("caller = %s"); params.append(caller) elif caller_is_null: @@ -1460,7 +1551,8 @@ def _empty_usage_aggregate() -> dict[str, Any]: def usage_aggregate(since_ts: "int | None" = None, caller: "str | None" = None, caller_is_null: bool = False, consumer_sha: "str | None" = None, provider: "str | None" = None, - model_family: "str | None" = None, *, project_id=None, environment_id=None, strict=False) -> dict[str, Any]: + model_family: "str | None" = None, *, project_id=None, environment_id=None, + until_ts: "int | None" = None, strict=False) -> dict[str, Any]: """Every dashboard stats aggregate in ONE window scan: overall totals plus the by_caller / by_provider / by_model_family / by_route / by_served_model / by_status / by_day breakdowns, as raw counters (see _agg_counter). Fail-soft @@ -1469,7 +1561,8 @@ def usage_aggregate(since_ts: "int | None" = None, caller: "str | None" = None, where, params = _usage_where(since_ts=since_ts, caller=caller, caller_is_null=caller_is_null, consumer_sha=consumer_sha, provider=provider, - model_family=model_family, project_id=project_id, environment_id=environment_id) + model_family=model_family, project_id=project_id, environment_id=environment_id, + until_ts=until_ts) sql = ("SELECT grouping(caller_k, provider_k, family_k, route_k," " served_k, status_k, day_k) AS gset," " caller_k, provider_k, family_k, route_k, served_k, status_k," @@ -1504,14 +1597,16 @@ def usage_aggregate(since_ts: "int | None" = None, caller: "str | None" = None, def usage_totals(since_ts: "int | None" = None, - caller: "str | None" = None, *, project_id=None, environment_id=None, strict=False) -> dict[str, Any]: + caller: "str | None" = None, *, project_id=None, environment_id=None, + until_ts: "int | None" = None, strict=False) -> 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, project_id=project_id, environment_id=environment_id) + where, params = _usage_where(since_ts=since_ts, caller=caller, project_id=project_id, + environment_id=environment_id, until_ts=until_ts) sql = ( "SELECT count(*)," " count(*) FILTER (WHERE COALESCE(status,0) >= 400)," @@ -1538,6 +1633,191 @@ def usage_totals(since_ts: "int | None" = None, return zeros +# The metering group keys. `key` is the FULL key digest (the control plane maps +# it to its key/labels); `route` prefers the published route recorded in the +# routing summary over the raw requested model. +_USAGE_GROUPS = { + "key": "COALESCE(NULLIF(consumer_sha,''),'unknown')", + "model": "COALESCE(NULLIF(model_family,''),NULLIF(requested_model,''),'unknown')", + "route": "COALESCE(NULLIF(routing_summary->>'route',''),NULLIF(requested_model,''),'unknown')", + "day": "to_char(to_timestamp(ts) AT TIME ZONE 'UTC','YYYY-MM-DD')", + "hour": "to_char(to_timestamp(ts) AT TIME ZONE 'UTC','YYYY-MM-DD\"T\"HH24:00:00\"Z\"')", +} +_TOKENS_TOTAL_SQL = ("CASE WHEN COALESCE(tokens_total,0) <> 0 THEN tokens_total" + " ELSE COALESCE(tokens_in,0)+COALESCE(tokens_out,0) END") + + +def usage_groups(group_by: str, since_ts: "int | None" = None, caller: "str | None" = None, *, + until_ts: "int | None" = None, project_id=None, environment_id=None) -> list[dict[str, Any]]: + """Window usage per `group_by` (key|model|route|day|hour), ordered by key. + The scope filter is part of the WHERE, so it applies before aggregation. + Raises on a store error (metering reads must not answer a false zero).""" + expr = _USAGE_GROUPS[group_by] + where, params = _usage_where(since_ts=since_ts, caller=caller, project_id=project_id, + environment_id=environment_id, until_ts=until_ts) + sql = (f"SELECT {expr} AS k, count(*)," + " count(*) FILTER (WHERE COALESCE(status,0) >= 400)," + " COALESCE(sum(COALESCE(tokens_in,0)),0), COALESCE(sum(COALESCE(tokens_out,0)),0)," + f" COALESCE(sum(COALESCE(tokens_cached,0)),0), COALESCE(sum({_TOKENS_TOTAL_SQL}),0)," + " round(COALESCE(sum(GREATEST(cost_usd,0)),0)::numeric,6)::float8" + f" FROM calls{where} GROUP BY 1 ORDER BY 1") + with _get_pool().connection() as conn: + _set_dashboard_statement_timeout(conn) + rows = conn.execute(sql, params).fetchall() + return [{"key": k, "runs": int(n), "errors": int(e), "tokens_in": int(ti), + "tokens_out": int(to), "tokens_cached": int(tc), "tokens_total": int(tt), + "cost_usd": float(c)} for k, n, e, ti, to, tc, tt, c in rows] + + +_EXPORT_COLS = ("id", "ts", "key_sha256", "subject", "route", "model_family", "provider", + "status", "tokens_in", "tokens_out", "tokens_cached", "tokens_total", + "cost_usd", "cost_basis") + + +def usage_export(caller: str, after_id: int = 0, limit: int = 1000, *, project_id=None, + environment_id=None, now: "float | None" = None) -> dict[str, Any]: + """One page of raw ledger rows in id order, for reconciliation. + + Only rows inserted at least LEDGER_WATERMARK_LAG_S ago are returned: ids are + allocated at INSERT but become visible at COMMIT, so on a multi-writer + deployment a lower id can appear after a higher one; holding back fresh rows + keeps the cursor from skipping it. `watermark_ts` is additionally capped by + the oldest call NOT yet delivered (held back or beyond this page), so + every call with ts < watermark_ts is in a page at or before next_after_id. + Raises on a store error.""" + stamp = float(time.time() if now is None else now) + where, params = _usage_where(caller=caller, project_id=project_id, + environment_id=environment_id) + after_id, limit = max(0, int(after_id)), max(1, int(limit)) + sql = ("SELECT id, ts, consumer_sha, routing_summary->>'subject'," + " COALESCE(NULLIF(routing_summary->>'route',''),requested_model)," + " model_family, provider_id, status, tokens_in, tokens_out, tokens_cached," + f" {_TOKENS_TOTAL_SQL}, cost_usd, cost_basis FROM calls{where}" + " AND id > %s AND COALESCE(recorded_at,0) <= %s ORDER BY id LIMIT %s") + with _get_pool().connection() as conn: + _set_dashboard_statement_timeout(conn) + rows = conn.execute(sql, params + [after_id, stamp - _watermark_lag_s(), limit]).fetchall() + next_after_id = int(rows[-1][0]) if rows else after_id + oldest = conn.execute(f"SELECT min(ts) FROM calls{where} AND id > %s", + params + [next_after_id]).fetchone()[0] + watermark = ledger_watermark(stamp) + if oldest is not None: + watermark = min(watermark, int(oldest)) + return {"rows": [dict(zip(_EXPORT_COLS, row)) for row in rows], + "next_after_id": next_after_id, "watermark_ts": watermark} + + +# ---- per-subject monthly budgets (control-plane key limits) --------------------- + +def reserve_subject_budget(tenant_id: int, subject: str, period: str, limit_usd: float, + amount_usd: float, reservation_id: str, ttl_s: float = 900.0, *, + now: "float | None" = None) -> "tuple[bool, bool, float, float]": + """Atomically admit one in-flight call against a subject's monthly budget. + + Returns ``(allowed, store_ok, spent_usd, reserved_usd)`` (reserved as seen + after this call). The subject's usage row is locked for one short + transaction, so concurrent replicas serialize: expired reservations are + swept first (their holder died), then the call is admitted only if + spent + live reservations + amount <= limit.""" + try: + stamp = float(time.time() if now is None else now) + key = (int(tenant_id), str(subject), str(period)) + with _get_pool().connection() as conn: + conn.execute( + "INSERT INTO subject_budget_usage(tenant_id,subject,period,updated_at)" + " VALUES (%s,%s,%s,%s) ON CONFLICT DO NOTHING", key + (int(stamp),)) + spent = float(conn.execute( + "SELECT spent_usd FROM subject_budget_usage" + " WHERE tenant_id=%s AND subject=%s AND period=%s FOR UPDATE", key).fetchone()[0]) + conn.execute( + "UPDATE subject_budget_reservations SET released=TRUE" + " WHERE tenant_id=%s AND subject=%s AND period=%s AND NOT released" + " AND expires_at <= %s", key + (stamp,)) + # Recomputed from the live rows (not the counter) so float drift or a + # lost update can never under-count what is in flight. + reserved = float(conn.execute( + "SELECT COALESCE(sum(amount_usd),0) FROM subject_budget_reservations" + " WHERE tenant_id=%s AND subject=%s AND period=%s AND NOT released", + key).fetchone()[0]) + amount = max(0.0, float(amount_usd)) + allowed = spent + reserved + amount <= float(limit_usd) + 1e-9 + if allowed: + conn.execute( + "INSERT INTO subject_budget_reservations" + "(id,tenant_id,subject,period,amount_usd,expires_at) VALUES (%s,%s,%s,%s,%s,%s)", + (str(reservation_id),) + key + (amount, stamp + float(ttl_s))) + reserved += amount + conn.execute( + "UPDATE subject_budget_usage SET reserved_usd=%s, updated_at=%s" + " WHERE tenant_id=%s AND subject=%s AND period=%s", + (reserved, int(stamp)) + key) + return allowed, True, spent, reserved + except Exception as exc: # noqa: BLE001 + _log.warning("host_store reserve_subject_budget failed: %s", exc) + return False, False, 0.0, 0.0 + + +def settle_subject_budget(reservation_id: str, cost_usd: "float | None") -> "tuple[bool, bool]": + """Release a reservation and book the call's actual cost (unknown or + negative -> 0). Idempotent: the reservation row is consumed, so a second + settle finds nothing and books nothing. A reservation already swept after + expiry still books its cost once. Returns ``(settled, store_ok)``. + + Called synchronously on the request's exit path, so the connection wait + and statement are bounded: on a store outage the reservation simply + expires instead of stalling the event loop.""" + try: + cost = float(cost_usd) if cost_usd is not None else 0.0 + if not math.isfinite(cost) or cost < 0: + cost = 0.0 + with _get_pool().connection(timeout=2.0) as conn: + conn.execute("SELECT set_config('statement_timeout', '2000', true)") + found = conn.execute( + "SELECT tenant_id,subject,period FROM subject_budget_reservations WHERE id=%s", + (str(reservation_id),)).fetchone() + if found is None: + return False, True + key = tuple(found) + # Lock order matches reserve (usage row, then reservation). + conn.execute("SELECT 1 FROM subject_budget_usage WHERE tenant_id=%s AND subject=%s" + " AND period=%s FOR UPDATE", key) + if conn.execute("DELETE FROM subject_budget_reservations WHERE id=%s RETURNING id", + (str(reservation_id),)).fetchone() is None: + return False, True # a concurrent settle won + conn.execute( + "UPDATE subject_budget_usage SET spent_usd=spent_usd+%s," + " reserved_usd=(SELECT COALESCE(sum(amount_usd),0) FROM subject_budget_reservations r" + " WHERE r.tenant_id=%s AND r.subject=%s AND r.period=%s AND NOT r.released)," + " updated_at=%s WHERE tenant_id=%s AND subject=%s AND period=%s", + (cost,) + key + (int(time.time()),) + key) + return True, True + except Exception as exc: # noqa: BLE001 + _log.warning("host_store settle_subject_budget failed: %s", exc) + return False, False + + +def subject_budgets(tenant_id: int, period: str, subjects: list[str], *, + now: "float | None" = None) -> list[dict[str, Any]]: + """Spent and live-reserved USD per requested subject for one tenant/period + (zeros for a subject with no activity). Reserved excludes expired holders + even before a sweep. Raises on a store error.""" + stamp = float(time.time() if now is None else now) + subjects = list(dict.fromkeys(str(s) for s in subjects)) + with _get_pool().connection() as conn: + _set_dashboard_statement_timeout(conn) + spent = dict(conn.execute( + "SELECT subject, spent_usd FROM subject_budget_usage" + " WHERE tenant_id=%s AND period=%s AND subject = ANY(%s)", + (int(tenant_id), str(period), subjects)).fetchall()) + reserved = dict(conn.execute( + "SELECT subject, sum(amount_usd) FROM subject_budget_reservations" + " WHERE tenant_id=%s AND period=%s AND subject = ANY(%s)" + " AND NOT released AND expires_at > %s GROUP BY subject", + (int(tenant_id), str(period), subjects, stamp)).fetchall()) + return [{"subject": s, "spent_usd": round(float(spent.get(s) or 0.0), 6), + "reserved_usd": round(float(reserved.get(s) or 0.0), 6)} for s in subjects] + + def policy_backtest_groups(since_ts: "int | None" = None, caller: "str | None" = None, limit: int = 50) -> dict[str, Any]: @@ -2352,4 +2632,5 @@ def truncate_all_for_tests() -> None: " consumer_keys, consumer_key_digests, consumer_rate_buckets," " peer_concurrency_leases," " consumer_budget_usage, analytics_hourly, analytics_rollup_state, peer_offers, buyer_status, route_observations," - " login_history, provider_prices, wallet_ops") + " login_history, provider_prices, wallet_ops," + " subject_budget_usage, subject_budget_reservations") diff --git a/internal_api.py b/internal_api.py index 939c3d0..8ef4808 100644 --- a/internal_api.py +++ b/internal_api.py @@ -5,14 +5,23 @@ 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?caller=&since_ts=[&until_ts=] + [&bucket=day][&group_by=key|model|route|day|hour] GET /internal/usage/recent?caller=&limit= + GET /internal/usage/export?caller=&after_id=&limit= + GET /internal/budgets?tenant_id=&period=YYYY-MM&subject=... + +The usage endpoints accept project_id+environment_id (both, positive) and then +filter by the recorded routing scope BEFORE aggregating or paging. Windows are +half-open: since_ts <= ts < until_ts. Usage and export responses carry +`watermark_ts` (host_store.ledger_watermark): the ledger is complete below it. Rows come from the `calls` ledger, which only records LLM calls that reached -the router — ingress-level rejects (401/429) are not counted here. +the router — ingress-level rejects (401/402/429) are not counted here. """ from __future__ import annotations +import re import time import asyncio from typing import Any @@ -26,6 +35,10 @@ router = APIRouter() _RECENT_LIMIT_MAX = 500 +_EXPORT_LIMIT_MAX = 5000 +_BUDGET_SUBJECTS_MAX = 500 +_SUBJECT_RE = re.compile(r"[gk]:[1-9][0-9]{0,17}") +_PERIOD_RE = re.compile(r"[0-9]{4}-(0[1-9]|1[0-2])") def _gate(request: Request) -> JSONResponse | None: @@ -36,31 +49,46 @@ def _gate(request: Request) -> JSONResponse | None: return None +def _scope(project_id: int | None, environment_id: int | None) -> dict | None: + """{} unscoped, the scope when both are positive, None when invalid.""" + if project_id is None and environment_id is None: + return {} + if project_id is None or environment_id is None or min(project_id, environment_id) <= 0: + return None + return {"project_id": project_id, "environment_id": environment_id} + + @router.get("/internal/usage") async def internal_usage(request: Request, caller: str = "", - since_ts: int | None = None, - bucket: str | None = None, + since_ts: int | None = None, until_ts: int | None = None, + bucket: str | None = None, group_by: str | None = None, project_id: int | None = None, environment_id: int | None = None) -> JSONResponse: denied = _gate(request) if denied is not None: return denied - scope = {} - if project_id is not None or environment_id is not None: - if project_id is None or environment_id is None or min(project_id, environment_id) <= 0: - return JSONResponse({"error": "invalid_scope"}, status_code=400) - scope = {"project_id": project_id, "environment_id": environment_id} + scope = _scope(project_id, environment_id) + if scope is None: + return JSONResponse({"error": "invalid_scope"}, status_code=400) caller = caller.strip() if not caller: return JSONResponse({"error": "caller_required"}, status_code=400) + if group_by is not None and group_by not in host_store._USAGE_GROUPS: + return JSONResponse({"error": "invalid_group_by"}, status_code=400) + window = {"since_ts": since_ts, "until_ts": until_ts} try: - totals = await asyncio.to_thread(host_store.usage_totals, since_ts=since_ts, caller=caller, strict=True, **scope) - by_day = (await asyncio.to_thread(host_store.usage_aggregate, since_ts=since_ts, caller=caller, strict=True, **scope))["by_day"] if bucket == "day" else {} + totals = await asyncio.to_thread(host_store.usage_totals, caller=caller, strict=True, **window, **scope) + by_day = (await asyncio.to_thread(host_store.usage_aggregate, caller=caller, strict=True, **window, **scope))["by_day"] if bucket == "day" else {} + groups = (await asyncio.to_thread(host_store.usage_groups, group_by, caller=caller, **window, **scope) + if group_by else None) except Exception: return JSONResponse({"error": "ledger_unavailable"}, status_code=503) out: dict[str, Any] = { "caller": caller, **({"scope_version": 2, **scope} if scope else {}), - "window": {"since_ts": since_ts, "until_ts": int(time.time())}, + "window": {"since_ts": since_ts, "until_ts": until_ts if until_ts is not None else int(time.time())}, + # Calls with ts < watermark_ts are all in the ledger: a window whose + # until_ts <= watermark_ts is closed and will not change. + "watermark_ts": host_store.ledger_watermark(), "runs": totals["requests"], "errors": totals["errors"], "tokens_in": totals["tokens_in"], @@ -74,6 +102,9 @@ async def internal_usage(request: Request, caller: str = "", {"date": day, "runs": counter["requests"], "cost_usd": counter["cost_usd"]} for day, counter in sorted(by_day.items()) ] + if groups is not None: + out["group_by"] = group_by + out["groups"] = groups return JSONResponse(out) @@ -84,11 +115,9 @@ async def internal_usage_recent(request: Request, caller: str = "", denied = _gate(request) if denied is not None: return denied - scope = {} - if project_id is not None or environment_id is not None: - if project_id is None or environment_id is None or min(project_id, environment_id) <= 0: - return JSONResponse({"error": "invalid_scope"}, status_code=400) - scope = {"project_id": project_id, "environment_id": environment_id} + scope = _scope(project_id, environment_id) + if scope is None: + return JSONResponse({"error": "invalid_scope"}, status_code=400) caller = caller.strip() if not caller: return JSONResponse({"error": "caller_required"}, status_code=400) @@ -117,3 +146,52 @@ async def internal_usage_recent(request: Request, caller: str = "", "key_sha256_prefix": (row.get("consumer_sha") or "")[:12] or None, }) return JSONResponse({"caller": caller, "calls": calls, **({"scope_version": 2, **scope} if scope else {})}) + + +@router.get("/internal/usage/export") +async def internal_usage_export(request: Request, caller: str = "", after_id: int = 0, + limit: int = 1000, project_id: int | None = None, + environment_id: int | None = None) -> JSONResponse: + """Raw ledger rows by id cursor — the reconciliation source of truth. Page + until `rows` is empty; periods ending before `watermark_ts` are complete.""" + denied = _gate(request) + if denied is not None: + return denied + scope = _scope(project_id, environment_id) + if scope is None: + return JSONResponse({"error": "invalid_scope"}, status_code=400) + caller = caller.strip() + if not caller: + return JSONResponse({"error": "caller_required"}, status_code=400) + if after_id < 0: + return JSONResponse({"error": "invalid_cursor"}, status_code=400) + limit = max(1, min(int(limit), _EXPORT_LIMIT_MAX)) + try: + page = await asyncio.to_thread(host_store.usage_export, caller, after_id, limit, **scope) + except Exception: + return JSONResponse({"error": "ledger_unavailable"}, status_code=503) + return JSONResponse({"caller": caller, **({"scope_version": 2, **scope} if scope else {}), **page}) + + +@router.get("/internal/budgets") +async def internal_budgets(request: Request, tenant_id: int | None = None, + period: str = "") -> JSONResponse: + """Spent/reserved per key subject for one tenant and UTC month. Subjects + are namespaced by tenant, so another tenant's identical subject never + leaks into this answer.""" + denied = _gate(request) + if denied is not None: + return denied + subjects = request.query_params.getlist("subject") + if tenant_id is None or tenant_id <= 0: + return JSONResponse({"error": "tenant_required"}, status_code=400) + if not _PERIOD_RE.fullmatch(period): + return JSONResponse({"error": "invalid_period"}, status_code=400) + if (not subjects or len(subjects) > _BUDGET_SUBJECTS_MAX + or any(not _SUBJECT_RE.fullmatch(s) for s in subjects)): + return JSONResponse({"error": "invalid_subject"}, status_code=400) + try: + rows = await asyncio.to_thread(host_store.subject_budgets, tenant_id, period, subjects) + except Exception: + return JSONResponse({"error": "budget_store_unavailable"}, status_code=503) + return JSONResponse({"tenant_id": tenant_id, "period": period, "subjects": rows}) diff --git a/tests/test_key_subject_limits.py b/tests/test_key_subject_limits.py new file mode 100644 index 0000000..e83b724 --- /dev/null +++ b/tests/test_key_subject_limits.py @@ -0,0 +1,698 @@ +"""Control-plane key limits: subject rate/budget enforcement at the ingress, +the reservation ledger, and the reconciliation surface (/internal/usage +group_by + until_ts, /internal/usage/export, /internal/budgets). + +The control plane is faked at its HTTP client (or at resolve_route), the +router at auth_proxy._client; budget/ledger state is the shared test Postgres. +""" +from __future__ import annotations + +import asyncio +import hashlib +import json +import os +import sys +import threading +import time +from pathlib import Path + +import httpx +import pytest +from fastapi.testclient import TestClient +from starlette.requests import Request + +ROOT = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(ROOT)) + +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 + +SECRET = {"x-internal-secret": "s3cret"} +ROUTE = {"route": "route:production", "revision": 2, "policy_id": "a" * 64, + "policy_ir": ["policy"], "execution": {}} +KEY = {"id": 42, "subject": "g:7", "labels": {"customer": "c-1"}} + + +@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() + + +# ---- fakes ------------------------------------------------------------------- + +class _JsonResp: + def __init__(self, cost=0.02, status_code=200): + self.status_code = status_code + self.headers = {"content-type": "application/json"} + self._body = json.dumps({"ok": True, "x_router": {"cost_usd": cost, "provider": "p"}}).encode() + + async def aread(self): + return self._body + + async def aclose(self): + pass + + +class _StreamResp: + status_code = 200 + headers = {"content-type": "text/event-stream"} + + def __init__(self, cost=0.03, fail_mid_stream=False): + self.cost, self.fail = cost, fail_mid_stream + self.closed = False + + async def aiter_raw(self): + yield b'data: {"choices":[]}\n\n' + if self.fail: + raise httpx.ReadError("upstream reset") + yield ('data: ' + json.dumps({"usage": {"prompt_tokens": 1, "completion_tokens": 1}, + "x_router": {"cost_usd": self.cost}}) + "\n\n").encode() + yield b"data: [DONE]\n\n" + + async def aclose(self): + self.closed = True + + +class _Upstream: + def __init__(self, make=lambda: _JsonResp()): + self.make, self.requests = make, [] + + def build_request(self, method, url, content=None, headers=None): + self.requests.append({"headers": headers or {}, "body": content}) + return object() + + async def send(self, req, stream=True): + resp = self.make() + if isinstance(resp, Exception): + raise resp + return resp + + +def _saas(monkeypatch, *, limits, key=KEY, upstream=None, tenant=5, caller="app-a"): + """A resolved SaaS key whose published route carries `key`/`limits`.""" + async def auth(token): + return {"ok": True, "caller": caller, "tenant_id": tenant, + "digest": hashlib.sha256(token.encode()).hexdigest(), "meta": {}} + + async def route(*args, **kwargs): + return {**ROUTE, "key": key, "limits": limits} + monkeypatch.setattr(auth_proxy, "_caller_auth_async", auth) + monkeypatch.setattr(cpc, "resolve_route", route) + up = upstream or _Upstream() + monkeypatch.setattr(auth_proxy, "_client", up) + return up + + +def _post(client, token="tok-a"): + return client.post("/v1/chat/completions", headers={"Authorization": f"Bearer {token}"}, + json={"model": "route:production", "messages": []}) + + +def _period(): + return time.strftime("%Y-%m", time.gmtime()) + + +def _budget(subject="g:7", tenant=5): + auth_proxy.drain_settles() + return host_store.subject_budgets(tenant, _period(), [subject])[0] + + +def _live_reservations(): + auth_proxy.drain_settles() + with host_store._get_pool().connection() as conn: + return conn.execute("SELECT count(*) FROM subject_budget_reservations" + " WHERE NOT released").fetchone()[0] + + +# ---- resolve_route: key/limits parsing ----------------------------------------- + +def _resolve_with(monkeypatch, extra): + def handler(request): + return httpx.Response(200, json={**ROUTE, **extra}) + monkeypatch.setattr(cpc, "_client", httpx.AsyncClient(transport=httpx.MockTransport(handler))) + return asyncio.run(cpc.resolve_route(3, "production")) + + +def test_resolve_route_parses_key_limits(monkeypatch): + data = _resolve_with(monkeypatch, {"key": KEY, "limits": { + "monthly_budget_usd": "40.00", "rate_per_min": 120, "burst": 40}}) + assert data["key"] == KEY + assert data["limits"] == {"monthly_budget_usd": 40.0, "rate_per_min": 120, "burst": 40} + data = _resolve_with(monkeypatch, {"key": {"id": 9, "subject": "k:9"}, "limits": { + "monthly_budget_usd": 2.5, "rate_per_min": None, "burst": None}}) + assert data["key"] == {"id": 9, "subject": "k:9", "labels": {}} + assert data["limits"] == {"monthly_budget_usd": 2.5, "rate_per_min": None, "burst": None} + + +def test_resolve_route_without_key_limits_is_unlimited(monkeypatch): + data = _resolve_with(monkeypatch, {}) + assert data["key"] is None and data["limits"] is None + + +@pytest.mark.parametrize("extra", [ + {"limits": {"rate_per_min": 5}}, # limits, no key + {"key": {**KEY, "subject": "x:7"}, "limits": {}}, + {"key": {**KEY, "subject": "g:07"}, "limits": {}}, + {"key": {**KEY, "subject": "g:0"}, "limits": {}}, + {"key": {**KEY, "subject": "g:1234567890123456789"}, "limits": {}}, # 19 digits + {"key": {**KEY, "id": "42"}, "limits": {}}, + {"key": {**KEY, "id": True}, "limits": {}}, + {"key": {**KEY, "labels": {"a": 1}}, "limits": {}}, + {"key": {**KEY, "labels": {f"k{i}": "v" for i in range(33)}}, "limits": {}}, + {"key": {**KEY, "labels": {"k": "v" * 257}}, "limits": {}}, + {"key": KEY, "limits": {"monthly_budget_usd": "0"}}, + {"key": KEY, "limits": {"monthly_budget_usd": "-1"}}, + {"key": KEY, "limits": {"monthly_budget_usd": "1e3"}}, + {"key": KEY, "limits": {"monthly_budget_usd": True}}, + {"key": KEY, "limits": {"monthly_budget_usd": 0}}, + {"key": KEY, "limits": {"rate_per_min": 0}}, + {"key": KEY, "limits": {"rate_per_min": 1.5}}, + {"key": KEY, "limits": {"burst": "10"}}, + {"key": KEY, "limits": ["bad"]}, + {"key": "g:7"}, +]) +def test_resolve_route_invalid_key_limits_fail_closed(monkeypatch, extra): + with pytest.raises(cpc.RouteUnavailable): + _resolve_with(monkeypatch, extra) + + +def test_invalid_limits_reject_request_before_router(monkeypatch): + require_host_store() + up = _Upstream() + monkeypatch.setattr(auth_proxy, "_client", up) + + async def auth(token): + return {"ok": True, "caller": "app-a", "tenant_id": 5, "digest": "a" * 64, "meta": {}} + monkeypatch.setattr(auth_proxy, "_caller_auth_async", auth) + + def handler(request): + return httpx.Response(200, json={**ROUTE, "key": {**KEY, "subject": "bogus"}}) + monkeypatch.setattr(cpc, "_client", httpx.AsyncClient(transport=httpx.MockTransport(handler))) + r = _post(TestClient(auth_proxy.app)) + assert r.status_code == 503 and r.json()["error"]["code"] == "route_unavailable" + assert up.requests == [] + + +# ---- ingress enforcement ------------------------------------------------------- + +def test_subject_rate_limit_429_after_plan_rate(monkeypatch): + require_host_store() + _saas(monkeypatch, limits={"monthly_budget_usd": None, "rate_per_min": 1, "burst": None}) + client = TestClient(auth_proxy.app) + assert _post(client).status_code == 200 + r = _post(client) + assert r.status_code == 429 + assert r.json()["error"]["code"] == "key_rate_limit" + assert int(r.headers["retry-after"]) >= 1 + # The bucket is per caller+subject: another subject is unaffected. + _saas(monkeypatch, limits={"monthly_budget_usd": None, "rate_per_min": 1, "burst": None}, + key={**KEY, "subject": "k:43"}) + assert _post(client).status_code == 200 + + +def test_subject_rate_store_unavailable_is_503(monkeypatch): + require_host_store() + up = _Upstream() + _saas(monkeypatch, limits={"monthly_budget_usd": None, "rate_per_min": 5, "burst": 5}, upstream=up) + real = host_store.consume_rate_token + monkeypatch.setattr(host_store, "consume_rate_token", + lambda consumer, *a, **k: (False, False, 0.0) if "|" in consumer else real(consumer, *a, **k)) + r = _post(TestClient(auth_proxy.app)) + assert r.status_code == 503 and r.json()["error"]["code"] == "key_rate_limit_unavailable" + assert up.requests == [] + + +def test_budget_settles_actual_cost_and_records_subject(monkeypatch): + require_host_store() + _saas(monkeypatch, limits={"monthly_budget_usd": "1.00", "rate_per_min": None, "burst": None}) + assert _post(TestClient(auth_proxy.app)).status_code == 200 + assert _budget() == {"subject": "g:7", "spent_usd": 0.02, "reserved_usd": 0.0} + assert _live_reservations() == 0 + host_store._write_q.join() + summary = host_store.recent_calls(caller="app-a")[0]["routing_summary"] + assert summary["key_id"] == "42" and summary["subject"] == "g:7" + + +def test_budget_exhausted_is_402_without_forwarding(monkeypatch): + require_host_store() + up = _Upstream(lambda: _JsonResp(cost=0.04)) + _saas(monkeypatch, limits={"monthly_budget_usd": "0.10", "rate_per_min": None, "burst": None}, + upstream=up) + client = TestClient(auth_proxy.app) + # 0.04 + 0.04 spent; a third call needs 0.08 + 0.05 > 0.10. + assert _post(client).status_code == 200 + assert _post(client).status_code == 200 + r = _post(client) + assert r.status_code == 402 + err = r.json()["error"] + assert err["code"] == "key_budget_exhausted" and err["type"] == "budget_error" + assert err["budget_usd"] == 0.1 and err["spent_usd"] == pytest.approx(0.08) + assert len(up.requests) == 2 + # Budgets are per tenant: the same subject in another tenant is fresh. + _saas(monkeypatch, limits={"monthly_budget_usd": "0.10", "rate_per_min": None, "burst": None}, + upstream=up, tenant=6) + assert _post(client).status_code == 200 + + +def test_budget_store_unavailable_is_503(monkeypatch): + require_host_store() + up = _Upstream() + _saas(monkeypatch, limits={"monthly_budget_usd": "1", "rate_per_min": None, "burst": None}, upstream=up) + monkeypatch.setattr(host_store, "reserve_subject_budget", lambda *a, **k: (False, False, 0.0, 0.0)) + r = _post(TestClient(auth_proxy.app)) + assert r.status_code == 503 and r.json()["error"]["code"] == "key_budget_unavailable" + assert up.requests == [] + + +def test_reservation_released_on_upstream_error_status(monkeypatch): + require_host_store() + _saas(monkeypatch, limits={"monthly_budget_usd": "1", "rate_per_min": None, "burst": None}, + upstream=_Upstream(lambda: _JsonResp(cost=None, status_code=500))) + assert _post(TestClient(auth_proxy.app)).status_code == 500 + assert _live_reservations() == 0 + assert _budget()["spent_usd"] == 0.0 + + +def test_reservation_released_on_upstream_exception(monkeypatch): + require_host_store() + _saas(monkeypatch, limits={"monthly_budget_usd": "1", "rate_per_min": None, "burst": None}, + upstream=_Upstream(lambda: httpx.ConnectError("down"))) + r = _post(TestClient(auth_proxy.app)) + assert r.status_code == 502 + assert _live_reservations() == 0 and _budget()["reserved_usd"] == 0.0 + + +def test_reservation_released_on_capacity_failure(monkeypatch): + require_host_store() + up = _Upstream() + _saas(monkeypatch, limits={"monthly_budget_usd": "1", "rate_per_min": None, "burst": None}, upstream=up) + + async def full(): + return False + monkeypatch.setattr(auth_proxy, "_capacity_acquire", full) + r = _post(TestClient(auth_proxy.app)) + assert r.status_code == 503 and r.json()["error"]["code"] == "router_overloaded" + assert _live_reservations() == 0 and up.requests == [] + + +def test_reservation_released_when_request_building_raises(monkeypatch): + require_host_store() + up = _Upstream() + _saas(monkeypatch, limits={"monthly_budget_usd": "1", "rate_per_min": None, "burst": None}, upstream=up) + + def boom(*a, **k): + raise RuntimeError("bad request build") + monkeypatch.setattr(up, "build_request", boom) + assert _post(TestClient(auth_proxy.app)).status_code == 502 + assert _live_reservations() == 0 + + +def test_streaming_settles_after_stream_with_streamed_cost(monkeypatch): + require_host_store() + _saas(monkeypatch, limits={"monthly_budget_usd": "1", "rate_per_min": None, "burst": None}, + upstream=_Upstream(lambda: _StreamResp(cost=0.03))) + r = _post(TestClient(auth_proxy.app)) + assert r.status_code == 200 and b"[DONE]" in r.content + assert _live_reservations() == 0 + assert _budget()["spent_usd"] == pytest.approx(0.03) + + +def test_streaming_upstream_failure_mid_stream_releases(monkeypatch): + require_host_store() + _saas(monkeypatch, limits={"monthly_budget_usd": "1", "rate_per_min": None, "burst": None}, + upstream=_Upstream(lambda: _StreamResp(fail_mid_stream=True))) + with pytest.raises(httpx.ReadError): + _post(TestClient(auth_proxy.app)) + assert _live_reservations() == 0 + + +def _request(token="tok-a"): + body = json.dumps({"model": "route:production", "messages": []}).encode() + + async def receive(): + return {"type": "http.request", "body": body, "more_body": False} + return Request({"type": "http", "method": "POST", "path": "/v1/chat/completions", + "headers": [(b"authorization", f"Bearer {token}".encode()), + (b"content-type", b"application/json")], + "query_string": b"", "client": ("test", 1), "server": ("test", 80), + "scheme": "http", "root_path": "", "http_version": "1.1"}, receive) + + +@pytest.mark.parametrize("before_first_chunk", [False, True]) +def test_client_disconnect_during_stream_releases(monkeypatch, before_first_chunk): + """Starlette closes the body iterator on disconnect (or only runs the + background task if the client left before the first chunk).""" + require_host_store() + stream = _StreamResp(cost=0.03) + _saas(monkeypatch, limits={"monthly_budget_usd": "1", "rate_per_min": None, "burst": None}, + upstream=_Upstream(lambda: stream)) + + async def run(): + resp = await auth_proxy.proxy("v1/chat/completions", _request()) + assert _live_reservations() == 1 # held while the stream is open + if not before_first_chunk: + await resp.body_iterator.__anext__() + await resp.body_iterator.aclose() + await resp.background() + asyncio.run(run()) + assert stream.closed and _live_reservations() == 0 + assert _budget()["reserved_usd"] == 0.0 + + +def test_routes_without_limits_take_no_reservation(monkeypatch): + require_host_store() + _saas(monkeypatch, limits=None, key=None) + assert _post(TestClient(auth_proxy.app)).status_code == 200 + with host_store._get_pool().connection() as conn: + assert conn.execute("SELECT count(*) FROM subject_budget_usage").fetchone()[0] == 0 + assert conn.execute("SELECT count(*) FROM consumer_rate_buckets" + " WHERE consumer LIKE '%|%'").fetchone()[0] == 0 + + +# ---- revocation ---------------------------------------------------------------- + +def test_revoked_scoped_key_rejected_on_next_request_despite_cached_resolve(monkeypatch): + """/internal/keys/resolve is cached positive for CP_RESOLVE_TTL_S, but the + scoped route lookup carries the key digest and is never cached: once the + control plane answers 403 for the (revoked) key, the next call is refused.""" + require_host_store() + state = {"revoked": False, "resolves": 0, "routes": 0} + + def handler(request): + if request.url.path == "/internal/keys/resolve": + state["resolves"] += 1 + return httpx.Response(200, json={"active": True, "consumer": "app-r", "tenant_id": 5, + "scope_version": 2, "project_id": 8, "environment_id": 9}) + state["routes"] += 1 + assert request.url.params["key_sha256"] == hashlib.sha256(b"tok-rev").hexdigest() + if state["revoked"]: + return httpx.Response(403, json={"error": "key_revoked"}) + return httpx.Response(200, json={**ROUTE, "scope_version": 2, "tenant_id": 5, + "project_id": 8, "environment_id": 9}) + monkeypatch.setattr(cpc, "_client", httpx.AsyncClient(transport=httpx.MockTransport(handler))) + up = _Upstream() + monkeypatch.setattr(auth_proxy, "_client", up) + client = TestClient(auth_proxy.app) + assert _post(client, "tok-rev").status_code == 200 + state["revoked"] = True + r = _post(client, "tok-rev") + assert r.status_code == 503 and r.json()["error"]["code"] == "route_unavailable" + assert state == {"revoked": True, "resolves": 1, "routes": 2} # resolve still cached + assert len(up.requests) == 1 + + +# ---- reservation ledger -------------------------------------------------------- + +def test_reservation_admission_and_settle_idempotent(host_store_clean): + hs = host_store_clean + assert hs.reserve_subject_budget(1, "k:1", "2030-01", 0.10, 0.05, "r1") == (True, True, 0.0, 0.05) + assert hs.reserve_subject_budget(1, "k:1", "2030-01", 0.10, 0.05, "r2")[:2] == (True, True) + allowed, ok, spent, reserved = hs.reserve_subject_budget(1, "k:1", "2030-01", 0.10, 0.05, "r3") + assert (allowed, ok, spent, reserved) == (False, True, 0.0, pytest.approx(0.10)) + assert hs.settle_subject_budget("r1", 0.07) == (True, True) + assert hs.settle_subject_budget("r1", 0.07) == (False, True) # no double count + assert hs.settle_subject_budget("r2", None) == (True, True) # unknown cost -> 0 + assert hs.settle_subject_budget("missing", 1.0) == (False, True) + assert hs.subject_budgets(1, "2030-01", ["k:1", "k:2"]) == [ + {"subject": "k:1", "spent_usd": 0.07, "reserved_usd": 0.0}, + {"subject": "k:2", "spent_usd": 0.0, "reserved_usd": 0.0}] + # spent 0.07 + 0.05 > 0.10 + assert hs.reserve_subject_budget(1, "k:1", "2030-01", 0.10, 0.05, "r4")[0] is False + assert hs.settle_subject_budget("neg", -5) == (False, True) + + +def test_expired_reservations_are_swept_and_still_settle_once(host_store_clean): + hs = host_store_clean + t = 1_900_000_000.0 + assert hs.reserve_subject_budget(2, "g:3", "2030-02", 0.05, 0.05, "old", ttl_s=10, now=t)[0] + assert hs.reserve_subject_budget(2, "g:3", "2030-02", 0.05, 0.05, "blocked", ttl_s=10, now=t + 5)[0] is False + # The reads ignore an expired holder even before a sweep. + assert hs.subject_budgets(2, "2030-02", ["g:3"], now=t + 11)[0]["reserved_usd"] == 0.0 + allowed, ok, spent, reserved = hs.reserve_subject_budget(2, "g:3", "2030-02", 0.05, 0.05, "new", + ttl_s=10, now=t + 11) + assert (allowed, reserved) == (True, pytest.approx(0.05)) + # The dead holder's late settle books its cost once without un-reserving "new". + assert hs.settle_subject_budget("old", 0.01) == (True, True) + assert hs.settle_subject_budget("old", 0.01) == (False, True) + with hs._get_pool().connection() as conn: + spent, reserved = conn.execute("SELECT spent_usd, reserved_usd FROM subject_budget_usage" + " WHERE tenant_id=2").fetchone() + assert spent == pytest.approx(0.01) and reserved == pytest.approx(0.05) + + +def test_concurrent_reservations_never_exceed_limit(host_store_clean): + hs = host_store_clean + results, barrier = [], threading.Barrier(40) + + def worker(i): + barrier.wait() + results.append(hs.reserve_subject_budget(3, "g:9", "2030-03", 1.0, 0.05, f"c{i}")) + threads = [threading.Thread(target=worker, args=(i,)) for i in range(40)] + for t in threads: + t.start() + for t in threads: + t.join() + assert all(ok for _a, ok, _s, _r in results) + assert sum(1 for allowed, *_ in results if allowed) == 20 + assert hs.subject_budgets(3, "2030-03", ["g:9"])[0]["reserved_usd"] == pytest.approx(1.0) + + +def test_routing_summary_keeps_key_identity(): + summary = host_store.routing_summary({"route": "route:x", "key_id": 42, "subject": "g:7", + "prompt": "secret", "labels": {"a": "b"}}) + assert summary["key_id"] == "42" and summary["subject"] == "g:7" + assert "prompt" not in summary and "labels" not in summary + + +# ---- /internal/usage group_by + until_ts --------------------------------------- + +def _seed_usage(now): + rows = [ # (caller, key, project, env, ts, tokens_in, cost, status) + ("acme", "a" * 64, 1, 10, now - 7200, 10, 0.01, 200), + ("acme", "a" * 64, 1, 10, now - 60, 20, 0.02, 200), + ("acme", "b" * 64, 1, 10, now - 60, 40, 0.04, 500), + ("acme", "b" * 64, 1, 11, now - 60, 800, 0.8, 200), # other environment + ("other", "a" * 64, 1, 10, now - 60, 900, 0.9, 200), # other caller + ] + for caller, key, project, env, ts, tin, cost, status in rows: + host_store.insert_call({"ts": ts, "caller": caller, "key_sha256": key, "status": status, + "tokens_in": tin, "tokens_out": 1, "cost_usd": cost, + "model_family": "fam", "requested_model": "route:production", + "decision_trace": {"route": "route:production", "project_id": project, + "environment_id": env, "subject": "g:7", + "key_id": 42}}) + + +def test_usage_group_by_key_is_scoped_before_aggregation(): + require_host_store() + now = int(time.time()) + _seed_usage(now) + client = TestClient(auth_proxy.app) + params = {"caller": "acme", "project_id": 1, "environment_id": 10, "group_by": "key"} + data = client.get("/internal/usage", params=params, headers=SECRET).json() + assert data["runs"] == 3 and data["group_by"] == "key" + groups = {g["key"]: g for g in data["groups"]} + assert set(groups) == {"a" * 64, "b" * 64} # full digests + assert groups["a" * 64]["runs"] == 2 and groups["a" * 64]["cost_usd"] == pytest.approx(0.03) + assert groups["b" * 64] == {"key": "b" * 64, "runs": 1, "errors": 1, "tokens_in": 40, + "tokens_out": 1, "tokens_cached": 0, "tokens_total": 41, + "cost_usd": 0.04} + # until_ts is exclusive and applies to totals and groups alike. + data = client.get("/internal/usage", params={**params, "until_ts": now - 60}, headers=SECRET).json() + assert data["runs"] == 1 and [g["runs"] for g in data["groups"]] == [1] + assert data["window"]["until_ts"] == now - 60 + # since_ts is inclusive: the boundary row at exactly since_ts counts. + data = client.get("/internal/usage", params={**params, "since_ts": now - 60}, headers=SECRET).json() + assert data["runs"] == 2 + + +def test_usage_reports_ledger_watermark(monkeypatch): + require_host_store() + monkeypatch.setenv("LEDGER_WATERMARK_LAG_S", "120") + client = TestClient(auth_proxy.app) + for extra in ({}, {"group_by": "key"}, {"bucket": "day"}): + before = int(time.time()) - 120 + data = client.get("/internal/usage", params={"caller": "acme", **extra}, headers=SECRET).json() + assert before <= data["watermark_ts"] <= int(time.time()) - 120 + host_store._pending_add(1_000) # a queued write holds the watermark back + try: + data = client.get("/internal/usage", params={"caller": "acme"}, headers=SECRET).json() + assert data["watermark_ts"] == 1_000 + finally: + host_store._pending_done(1_000) + + +@pytest.mark.parametrize("group_by,expected", [ + ("model", {"fam"}), ("route", {"route:production"})]) +def test_usage_group_by_model_and_route(group_by, expected): + require_host_store() + _seed_usage(int(time.time())) + data = TestClient(auth_proxy.app).get("/internal/usage", params={ + "caller": "acme", "group_by": group_by}, headers=SECRET).json() + assert {g["key"] for g in data["groups"]} == expected + assert sum(g["runs"] for g in data["groups"]) == 4 + + +def test_usage_group_by_day_and_hour_are_utc_buckets(): + require_host_store() + now = int(time.time()) + _seed_usage(now) + client = TestClient(auth_proxy.app) + hours = client.get("/internal/usage", params={"caller": "acme", "group_by": "hour"}, + headers=SECRET).json()["groups"] + expect = {time.strftime("%Y-%m-%dT%H:00:00Z", time.gmtime(ts)) for ts in (now - 7200, now - 60)} + assert {g["key"] for g in hours} == expect + days = client.get("/internal/usage", params={"caller": "acme", "group_by": "day"}, + headers=SECRET).json()["groups"] + assert sum(g["runs"] for g in days) == 4 + assert all(len(g["key"]) == 10 for g in days) + assert client.get("/internal/usage", params={"caller": "acme", "group_by": "caller"}, + headers=SECRET).status_code == 400 + + +# ---- /internal/usage/export ---------------------------------------------------- + +def test_export_cursor_pages_in_id_order_with_watermark(monkeypatch): + require_host_store() + monkeypatch.setenv("LEDGER_WATERMARK_LAG_S", "0") + now = int(time.time()) + _seed_usage(now) + client = TestClient(auth_proxy.app) + scope = {"caller": "acme", "project_id": 1, "environment_id": 10} + page = client.get("/internal/usage/export", params={**scope, "limit": 2}, headers=SECRET).json() + assert [r["tokens_in"] for r in page["rows"]] == [10, 20] + row = page["rows"][0] + assert set(row) == {"id", "ts", "key_sha256", "subject", "route", "model_family", "provider", + "status", "tokens_in", "tokens_out", "tokens_cached", "tokens_total", + "cost_usd", "cost_basis"} + assert row["key_sha256"] == "a" * 64 and row["subject"] == "g:7" and row["route"] == "route:production" + # The third scoped row (ts now-60) is not delivered yet, so the watermark + # cannot pass it even though the lag is zero. + assert page["watermark_ts"] <= now - 60 + page2 = client.get("/internal/usage/export", params={ + **scope, "after_id": page["next_after_id"], "limit": 2}, headers=SECRET).json() + assert [r["tokens_in"] for r in page2["rows"]] == [40] + assert page2["rows"][0]["id"] > page["next_after_id"] + assert page2["watermark_ts"] >= now - 1 + page3 = client.get("/internal/usage/export", params={ + **scope, "after_id": page2["next_after_id"]}, headers=SECRET).json() + assert page3["rows"] == [] and page3["next_after_id"] == page2["next_after_id"] + assert page3["scope_version"] == 2 + + +def test_export_holds_back_fresh_rows_and_caps_watermark(monkeypatch): + require_host_store() + monkeypatch.setenv("LEDGER_WATERMARK_LAG_S", "120") + now = int(time.time()) + _seed_usage(now) + page = TestClient(auth_proxy.app).get("/internal/usage/export", params={"caller": "acme"}, + headers=SECRET).json() + assert page["rows"] == [] and page["next_after_id"] == 0 + assert page["watermark_ts"] <= now - 7200 # the oldest undelivered call + + +def test_ledger_watermark_waits_for_queued_writes(monkeypatch): + monkeypatch.setenv("LEDGER_WATERMARK_LAG_S", "10") + now = 2_000_000_000 + assert host_store.ledger_watermark(now) == now - 10 + host_store._pending_add(now - 500) + try: + assert host_store.ledger_watermark(now) == now - 500 + finally: + host_store._pending_done(now - 500) + assert host_store.ledger_watermark(now) == now - 10 + + +def test_export_validation_and_auth(monkeypatch): + client = TestClient(auth_proxy.app) + assert client.get("/internal/usage/export", params={"caller": "a"}, + headers={"x-internal-secret": "no"}).status_code == 403 + assert client.get("/internal/usage/export", params={"caller": "a", "after_id": -1}, + headers=SECRET).status_code == 400 + assert client.get("/internal/usage/export", headers=SECRET).status_code == 400 + for invalid in ({"project_id": 1}, {"project_id": 1, "environment_id": 0}): + assert client.get("/internal/usage/export", params={"caller": "a", **invalid}, + headers=SECRET).status_code == 400 + monkeypatch.setattr(cpc, "CONTROL_PLANE_INTERNAL_SECRET", "") + assert client.get("/internal/usage/export", params={"caller": "a"}).status_code == 404 + + +def test_export_limit_is_capped(monkeypatch): + seen = {} + + def fake(caller, after_id, limit, **scope): + seen.update(limit=limit) + return {"rows": [], "next_after_id": after_id, "watermark_ts": 0} + monkeypatch.setattr(host_store, "usage_export", fake) + TestClient(auth_proxy.app).get("/internal/usage/export", params={"caller": "a", "limit": 99999}, + headers=SECRET) + assert seen["limit"] == 5000 + + +# ---- /internal/budgets --------------------------------------------------------- + +def test_budgets_endpoint_auth_validation_and_tenant_scoping(monkeypatch): + require_host_store() + period = "2030-04" + host_store.reserve_subject_budget(5, "g:7", period, 10, 0.05, "b1") + host_store.settle_subject_budget("b1", 1.25) + host_store.reserve_subject_budget(5, "g:7", period, 10, 0.05, "b2") # in flight + host_store.reserve_subject_budget(6, "g:7", period, 10, 0.05, "other") # other tenant + host_store.settle_subject_budget("other", 9.0) + client = TestClient(auth_proxy.app) + r = client.get("/internal/budgets", params=[("tenant_id", 5), ("period", period), + ("subject", "g:7"), ("subject", "k:1")], headers=SECRET) + assert r.status_code == 200 + assert r.json() == {"tenant_id": 5, "period": period, "subjects": [ + {"subject": "g:7", "spent_usd": 1.25, "reserved_usd": 0.05}, + {"subject": "k:1", "spent_usd": 0.0, "reserved_usd": 0.0}]} + ok = {"tenant_id": 5, "period": period, "subject": "g:7"} + assert client.get("/internal/budgets", params=ok, headers={"x-internal-secret": "no"}).status_code == 403 + for bad in ({"tenant_id": 0}, {"tenant_id": None}, {"period": "2030-13"}, {"period": "2030-4"}, + {"subject": "x:1"}, {"subject": None}): + params = {k: v for k, v in {**ok, **bad}.items() if v is not None} + assert client.get("/internal/budgets", params=params, headers=SECRET).status_code == 400, bad + too_many = [("tenant_id", 5), ("period", period)] + [("subject", f"k:{i + 1}") for i in range(501)] + assert client.get("/internal/budgets", params=too_many, headers=SECRET).status_code == 400 + + def down(*a, **k): + raise RuntimeError("db down") + monkeypatch.setattr(host_store, "subject_budgets", down) + assert client.get("/internal/budgets", params=ok, headers=SECRET).status_code == 503 + monkeypatch.setattr(cpc, "CONTROL_PLANE_INTERNAL_SECRET", "") + assert client.get("/internal/budgets", params=ok).status_code == 404 + + +def test_successful_call_with_unknown_cost_books_the_reservation(monkeypatch): + """An unpriced model, or a stream closed before its usage chunk, must not + make a successful call free against the budget.""" + require_host_store() + _saas(monkeypatch, limits={"monthly_budget_usd": "1", "rate_per_min": None, "burst": None}, + upstream=_Upstream(lambda: _JsonResp(cost=None))) + assert _post(TestClient(auth_proxy.app)).status_code == 200 + assert _live_reservations() == 0 + assert _budget()["spent_usd"] == pytest.approx(auth_proxy.CLOUD_BUDGET_RESERVATION_USD) + + +def test_failed_settle_is_retried(monkeypatch): + require_host_store() + monkeypatch.setattr(auth_proxy, "SETTLE_RETRY_DELAYS_S", (0.0, 0.0, 0.0)) + real, calls = host_store.settle_subject_budget, [] + + def flaky(reservation_id, cost): + calls.append(reservation_id) + return (False, False) if len(calls) < 3 else real(reservation_id, cost) + monkeypatch.setattr(host_store, "settle_subject_budget", flaky) + _saas(monkeypatch, limits={"monthly_budget_usd": "1.00", "rate_per_min": None, "burst": None}) + assert _post(TestClient(auth_proxy.app)).status_code == 200 + assert _budget()["spent_usd"] == 0.02 and len(calls) == 3