feat(budget): enforce budget_limit_cents as a hard lifetime spend cap

app/main.py

@@ -69,6 +69,13 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
     async with engine.begin() as conn:
         await conn.run_sync(Base.metadata.create_all)
 
+    # create_all only makes missing tables, never alters existing ones. Bring
+    # existing deployments (SQLite volume, Postgres) up to date with columns added
+    # after their initial release so they don't 503 on the new ORM columns.
+    from packages.db.migrate import ensure_budget_columns
+
+    await ensure_budget_columns(engine)
+
     # Fail closed before any traffic can be served: refuse to boot when
     # provider credentials are (or would be) sealed with the publicly-known
     # dev encryption key. Runs after create_all so a fresh database's empty

app/routes/chat.py

@@ -31,6 +31,7 @@
 from app.protocols.sse import AdapterError
 from app.quality_scores import resolve_model_metrics
 from app.schemas import ChatCompletionRequest
+from packages.auth.spend import MICROCENTS_PER_CENT, charge_budget, is_exhausted, read_spent
 from packages.auth.types import KeyContext
 from packages.db.models.request_log import RequestLog
 from packages.litellm_adapter.catalog import CATALOG, CATALOG_BY_ID
@@ -306,6 +307,21 @@ async def execute_chat(
             detail=f"Model '{body.model}' is not allowed for this API key",
         )
 
+    async def _settle_budget(session, actual_microcents: int, *, commit: bool = True) -> None:
+        """Record `actual_microcents` of spend against the cap, if any.
+
+        No-op when the key has no budget cap. When `commit` is False the UPDATE is
+        executed but not committed, so the caller commits it in the same
+        transaction as the request-log write — making the row and the charge one
+        atomic unit. Idempotency across retries comes from the row's trace_id
+        (a persisted trace_id proves the charge also landed), not from a
+        process-local flag.
+        """
+        cap = getattr(kc, "_budget_cap", None)
+        if cap is None:
+            return
+        await charge_budget(session, str(kc.key_id), cap, actual_microcents, commit=commit)
+
     client = await router_cache.get_router(db)
     raw_strategy = getattr(client, "strategy", None)
     strategy = raw_strategy if isinstance(raw_strategy, str) and raw_strategy else "balanced"
@@ -420,6 +436,24 @@ async def execute_chat(
         resolved_model = candidates[0]
         body.model = candidates[0]  # mutate for downstream completion call
 
+    # Budget enforcement: `budget_limit_cents` is a hard lifetime cap. The check
+    # runs only after the request has passed every pre-dispatch validation (model
+    # allowlist, provider deployability), so a request we reject before touching
+    # an upstream never consumes budget. The real cost is only known once the
+    # upstream response/stream completes, so we record it atomically in
+    # `_settle_budget` — the `UPDATE spent = spent + actual WHERE spent + actual
+    # <= cap` guard makes this safe under concurrency and never lets the counter
+    # exceed the cap (fail-closed, never over-recorded).
+    if kc.budget_limit_cents is not None:
+        cap = kc.budget_limit_cents * MICROCENTS_PER_CENT
+        if await is_exhausted(db, str(kc.key_id), cap):
+            raise HTTPException(
+                status_code=429,
+                detail=f"API key budget exhausted ({cap} microcents lifetime cap reached).",
+            )
+        kc._budget_cap = cap
+        kc._budget_spent = await read_spent(db, str(kc.key_id))
+
     started_perf = time.perf_counter()
     completion_kwargs = body.model_dump(exclude_none=True)
 
@@ -489,6 +523,7 @@ async def execute_chat(
         log.cost_microcents = 0
         db.add(log)
         try:
+            await _settle_budget(db, 0, commit=False)
             await db.commit()
         except Exception as commit_err:
             logger.warning("request_log_commit_failed", error=str(commit_err))
@@ -509,16 +544,14 @@ async def execute_chat(
     # mid-flight cascade is impossible — we have to surface the error and let
     # the client decide what to do.
     if body.stream:
-        # Auto-inject `stream_options.include_usage=True` if the client
-        # didn't set it. Without this, OpenAI/LiteLLM streaming responses
-        # omit the `usage` field entirely — chunks have no token counts,
-        # so our log row gets input=0, output=0 and the cost calculation
-        # rounds to zero. Almost no client knows to opt-in to this flag,
-        # which would silently zero out streaming spend in the dashboard.
-        # Honor an explicit `include_usage=False` from the client if they
-        # really want to disable it (e.g. wire-format compatibility tests).
+        # Auto-inject `stream_options.include_usage=True` if the client didn't set
+        # it, so streaming responses carry token counts and we bill correctly.
+        # A budgeted key MUST receive usage so its spend is measured: a
+        # client-supplied `include_usage=False` would otherwise record zero cost
+        # and let a capped key stream for free, so force it on for any budgeted key
+        # regardless of the client's preference.
         existing_so = completion_kwargs.get("stream_options") or {}
-        if "include_usage" not in existing_so:
+        if getattr(kc, "_budget_cap", None) is not None or "include_usage" not in existing_so:
             completion_kwargs["stream_options"] = {**existing_so, "include_usage": True}
 
         async def _log_pre_stream_failure(status: int, err_type: str | None) -> None:
@@ -542,6 +575,7 @@ async def _log_pre_stream_failure(status: int, err_type: str | None) -> None:
             )
             db.add(log)
             try:
+                await _settle_budget(db, 0, commit=False)
                 await db.commit()
             except Exception as commit_err:
                 logger.warning("request_log_commit_failed", error=str(commit_err))
@@ -581,6 +615,17 @@ async def sse() -> AsyncGenerator[str, None]:
             status_code = 200
             error_type: str | None = None
             log_written = False
+            # True only once a terminal `data: [DONE]` has been emitted, i.e. the
+            # response was delivered in full. While False, the stream ended early
+            # (client disconnect / mid-stream upstream error) and the real cost is
+            # unknown, so the budget claim must be kept (fail-closed) rather than
+            # released — otherwise a client could stream tokens then hang up before
+            # the usage frame to bypass the cap.
+            stream_completed = False
+            # True once any usage frame has been observed in the stream. A completed
+            # stream with no usage frame means cost is unknown (client suppressed it
+            # or the provider omitted it), so the cap must still be enforced.
+            usage_seen = False
 
             async def _finalize() -> None:
                 """Write the request log row exactly once.
@@ -659,27 +704,46 @@ async def _already_persisted(s) -> bool:
                         select(RequestLog.id).where(RequestLog.trace_id == row_values["trace_id"])
                     )) is not None
 
+                def _settlement_amount() -> int:
+                    """Budget charge for this request, in microcents.
+
+                    When the real cost is unknown — the stream ended without a
+                    terminal [DONE], or a completed stream never delivered a usage
+                    frame (e.g. a client forced include_usage=False or a provider
+                    omitted usage) — charge the full remaining allowance so a client
+                    cannot suppress the usage frame to bypass the cap.
+                    """
+                    actual = row_values.get("cost_microcents") or 0
+                    cost_unknown = (not stream_completed) or (not usage_seen)
+                    if cost_unknown:
+                        actual = max(
+                            actual,
+                            (getattr(kc, "_budget_cap", 0) or 0)
+                            - (getattr(kc, "_budget_spent", 0) or 0),
+                        )
+                    return actual
+
                 async def _commit_row(*, retry: bool) -> None:
-                    """INSERT + COMMIT the row on a session of its own.
-
-                    Only a failing `commit()` propagates; a failure while
-                    closing the session AFTER the commit returned is
-                    swallowed — the row is already in. A retry is
-                    idempotent: it first looks the trace_id up, so a COMMIT
-                    that landed but whose ack was lost on the wire
-                    (PostgreSQL, connection dropped mid-ack) is not
-                    inserted a second time — and the shared primary key
-                    would reject a duplicate anyway.
+                    """Persist the request-log row and charge the budget in ONE commit.
+
+                    The INSERT and the budget charge share a single transaction. If it
+                    commits, both are durable; if it fails, both roll back and the
+                    retry re-runs both. Because the charge lands in the same commit as
+                    the row, a persisted trace_id proves the charge also landed — so a
+                    retry returns without re-charging. The charge is therefore applied
+                    exactly once per request: never doubled (on a commit-ack-loss
+                    retry) and never dropped.
                     """
                     log = RequestLog(**row_values)
                     if session_mod._session_factory is None:
                         # Test-only fallback (the app always installs a
                         # factory): the request-scoped session has to be
                         # rolled back before a retry can reuse it.
-                        if retry and await _already_persisted(db):
+                        if retry and (await _already_persisted(db)):
                             return
                         db.add(log)
                         try:
+                            await _settle_budget(db, _settlement_amount(), commit=False)
                             await db.commit()
                         except Exception:
                             try:
@@ -690,9 +754,10 @@ async def _commit_row(*, retry: bool) -> None:
                         return
                     s = session_mod._session_factory()
                     try:
-                        if retry and await _already_persisted(s):
+                        if retry and (await _already_persisted(s)):
                             return
                         s.add(log)
+                        await _settle_budget(s, _settlement_amount(), commit=False)
                         await s.commit()
                     finally:
                         try:
@@ -770,10 +835,12 @@ async def _commit_row(*, retry: bool) -> None:
                         agg_latency = meta.get("latency_ms", agg_latency)
                     if "usage" in d and d["usage"]:
                         agg_usage = d["usage"]
+                        usage_seen = True
                     if d.get("model"):
                         agg_model = d["model"]
                     yield f"data: {json.dumps(d, separators=(',', ':'))}\n\n"
                 yield "data: [DONE]\n\n"
+                stream_completed = True
             except (asyncio.CancelledError, GeneratorExit):
                 # Client closed the connection (Ctrl+C, tab closed, browser
                 # navigated away, proxy timeout, ...). Two distinct signals
@@ -874,6 +941,12 @@ async def _commit_row(*, retry: bool) -> None:
                 # is legal; clients reading until [DONE] still get it after
                 # an upstream error.
                 yield "data: [DONE]\n\n"
+                # The error response was delivered in full (terminal [DONE] sent),
+                # so settle against the actual cost only — not the full remaining
+                # allowance. Without this, every mid-stream provider failure would
+                # charge (and exhaust) the key's entire remaining budget even
+                # though the delivered response cost ~0.
+                stream_completed = True
             finally:
                 # Same shielding reason as the cancel branch: ensure the
                 # log write actually completes before we unwind, even if
@@ -909,6 +982,12 @@ async def _commit_row(*, retry: bool) -> None:
     response: dict = {}
     actual_resolved: str | None = None
     try:
+        # A budgeted key must receive usage so its spend is measured. Force
+        # include_usage on for budgeted keys even if the client omitted it.
+        if getattr(kc, "_budget_cap", None) is not None:
+            existing_so = completion_kwargs.get("stream_options") or {}
+            if existing_so.get("include_usage") is not True:
+                completion_kwargs["stream_options"] = {**existing_so, "include_usage": True}
         response = await client.acompletion(
             **completion_kwargs,
             fallbacks=fallbacks_arg,
@@ -950,11 +1029,35 @@ async def _commit_row(*, retry: bool) -> None:
             # _build_log_row would otherwise default to via requested_model).
             actual_resolved=actual_resolved or resolved_model,
         )
-        db.add(log)
-        try:
-            await db.commit()
-        except Exception as commit_err:
-            logger.warning("request_log_commit_failed", error=str(commit_err))
+        # Persist the log row and the budget charge atomically (same transaction),
+        # retrying transient commit failures so a budgeted key is never under-
+        # charged when the DB is stressed — mirroring the streaming path. A
+        # persisted trace_id proves both landed, so a retry skips rather than
+        # double-charging.
+        from sqlalchemy import select
+
+        max_attempts = len(_LOG_COMMIT_BACKOFF_S) + 1
+        for attempt in range(1, max_attempts + 1):
+            try:
+                if attempt > 1 and (
+                    await db.scalar(
+                        select(RequestLog.id).where(RequestLog.trace_id == log.trace_id)
+                    )
+                ) is not None:
+                    break  # already durable (log + charge committed)
+                db.add(log)
+                await _settle_budget(db, log.cost_microcents, commit=False)
+                await db.commit()
+                break
+            except Exception as commit_err:
+                try:
+                    await db.rollback()
+                except Exception:
+                    pass
+                if attempt == max_attempts:
+                    logger.warning(
+                        "request_log_commit_failed", error=str(commit_err), attempts=attempt,
+                    )
 
     if isinstance(response, dict) and "_orca_meta" in response:
         response = {k: v for k, v in response.items() if k != "_orca_meta"}

packages/auth/spend.py

@@ -0,0 +1,90 @@
+"""Per-key lifetime spend tracking that enforces ``ApiKey.budget_limit_cents``.
+
+The cap is a hard lifetime limit on the key's total spend, in microcents
+(1 cent = 10_000 microcents; 1 USD = 1_000_000 microcents, matching chat.py's
+cost math).
+
+Actual cost is only known after the upstream call returns, so enforcement is a
+single atomic ``UPDATE`` that adds the real cost and refuses to let the counter
+exceed the cap::
+
+    UPDATE api_keys SET spent_microcents = spent_microcents + :actual
+    WHERE id = :id AND spent_microcents + :actual <= :cap
+
+Concurrent requests for the same key each add their own cost atomically; only a
+request whose *own* cost alone would breach the remaining budget matches zero
+rows. In that case the counter is clamped to ``cap`` so the key is correctly
+maxed out and the next request is rejected — fail-closed, never over-recorded.
+
+This avoids both failure modes of a pre-claim design: it never records spend
+past the cap (no over-spend), and it does not reserve the whole remaining budget
+up front (so a key's requests are not serialized behind a single in-flight one).
+
+Kept free of FastAPI imports so it stays unit-testable and reusable from
+non-HTTP paths (background jobs, CLI minting tools).
+"""
+
+from __future__ import annotations
+
+from sqlalchemy import select, update
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from packages.db.models.api_key import ApiKey
+
+MICROCENTS_PER_CENT = 10_000
+
+
+async def read_spent(db: AsyncSession, api_key_id: str) -> int:
+    """Return the key's currently-recorded lifetime spend in microcents."""
+    spent = (
+        await db.execute(select(ApiKey.spent_microcents).where(ApiKey.id == api_key_id))
+    ).scalar_one_or_none()
+    return int(spent or 0)
+
+
+async def is_exhausted(db: AsyncSession, api_key_id: str, cap_microcents: int) -> bool:
+    """Fast pre-check: has the key already reached its lifetime cap?"""
+    return (await read_spent(db, api_key_id)) >= cap_microcents
+
+
+async def charge_budget(
+    db: AsyncSession,
+    api_key_id: str,
+    cap_microcents: int,
+    actual_microcents: int,
+    *,
+    commit: bool = True,
+) -> bool:
+    """Atomically record ``actual_microcents`` of spend, never exceeding ``cap``.
+
+    Returns ``True`` if the cost fit under the cap (the counter advanced by
+    ``actual``), or ``False`` if the request alone would have breached the cap —
+    in which case the counter is clamped to ``cap`` so the key is maxed out and
+    blocked going forward. The boundary request may already have been served
+    upstream; it cannot be un-spent, but we never record more than the cap and we
+    stop the next one. Fail-closed.
+
+    When ``commit`` is False the UPDATEs are executed but not committed, so the
+    caller can commit them in the same transaction as the request-log write
+    (atomic log + charge — no window where the log lands but the charge is lost).
+    """
+    actual = actual_microcents or 0
+    result = await db.execute(
+        update(ApiKey)
+        .where(ApiKey.id == api_key_id, ApiKey.spent_microcents + actual <= cap_microcents)
+        .values(spent_microcents=ApiKey.spent_microcents + actual)
+    )
+    if result.rowcount:
+        if commit:
+            await db.commit()
+        return True
+    # Would have exceeded the cap: clamp so the counter never overshoots and the
+    # key is correctly reported as exhausted thereafter.
+    await db.execute(
+        update(ApiKey)
+        .where(ApiKey.id == api_key_id, ApiKey.spent_microcents < cap_microcents)
+        .values(spent_microcents=cap_microcents)
+    )
+    if commit:
+        await db.commit()
+    return False