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

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) -> None:
+        """Atomically record `actual_microcents` of spend against the cap.
+
+        Idempotent: only the first call for a request does work. `session` is
+        the DB session to run the charge on (the request session, or the
+        dedicated log session for the streaming path). `charge_budget` makes the
+        UPDATE refuse to let the counter exceed the cap, so a request can never
+        over-spend even under concurrency.
+        """
+        cap = getattr(kc, "_budget_cap", None)
+        if cap is None or getattr(kc, "_budget_settled", False):
+            return
+        kc._budget_settled = True
+        await charge_budget(session, str(kc.key_id), cap, actual_microcents)
+
     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)
 
@@ -492,6 +526,7 @@ async def execute_chat(
             await db.commit()
         except Exception as commit_err:
             logger.warning("request_log_commit_failed", error=str(commit_err))
+        await _settle_budget(db, 0)
         return JSONResponse(
             content=cache_hit_response,
             headers={
@@ -545,6 +580,7 @@ async def _log_pre_stream_failure(status: int, err_type: str | None) -> None:
                 await db.commit()
             except Exception as commit_err:
                 logger.warning("request_log_commit_failed", error=str(commit_err))
+            await _settle_budget(db, 0)
 
         try:
             stream_obj = await client.acompletion(
@@ -581,6 +617,13 @@ 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
 
             async def _finalize() -> None:
                 """Write the request log row exactly once.
@@ -681,6 +724,18 @@ async def _commit_row(*, retry: bool) -> None:
                         db.add(log)
                         try:
                             await db.commit()
+                            actual_cost = row_values.get("cost_microcents") or 0
+                            if not stream_completed:
+                                # Stream ended without [DONE]: real cost unknown,
+                                # so charge the full remaining allowance to keep
+                                # the budget consumed (fail-closed) rather than
+                                # releasing it and letting a client bypass the cap
+                                # by hanging up before the usage frame.
+                                actual_cost = max(
+                                    actual_cost,
+                                    (kc._budget_cap or 0) - (kc._budget_spent or 0),
+                                )
+                            await _settle_budget(db, actual_cost)
                         except Exception:
                             try:
                                 await db.rollback()
@@ -694,6 +749,16 @@ async def _commit_row(*, retry: bool) -> None:
                             return
                         s.add(log)
                         await s.commit()
+                        actual_cost = row_values.get("cost_microcents") or 0
+                        if not stream_completed:
+                            # Stream ended without [DONE]: real cost unknown,
+                            # so charge the full remaining allowance to keep the
+                            # budget consumed (fail-closed).
+                            actual_cost = max(
+                                actual_cost,
+                                (kc._budget_cap or 0) - (kc._budget_spent or 0),
+                            )
+                        await _settle_budget(s, actual_cost)
                     finally:
                         try:
                             await s.close()
@@ -774,6 +839,7 @@ async def _commit_row(*, retry: bool) -> None:
                         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
@@ -955,6 +1021,7 @@ async def _commit_row(*, retry: bool) -> None:
             await db.commit()
         except Exception as commit_err:
             logger.warning("request_log_commit_failed", error=str(commit_err))
+        await _settle_budget(db, log.cost_microcents)
 
     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,79 @@
+"""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
+) -> 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.
+    """
+    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:
+        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)
+    )
+    await db.commit()
+    return False

packages/db/models/api_key.py

@@ -2,7 +2,7 @@
 
 from datetime import datetime
 
-from sqlalchemy import JSON, Boolean, DateTime, ForeignKey, Integer, String
+from sqlalchemy import JSON, BigInteger, Boolean, DateTime, ForeignKey, String
 from sqlalchemy.orm import Mapped, mapped_column
 
 from packages.db.models.base import Base, SoftDeleteMixin, TimestampMixin, UUIDMixin
@@ -18,7 +18,16 @@ class ApiKey(Base, UUIDMixin, TimestampMixin, SoftDeleteMixin):
     key_hash: Mapped[str] = mapped_column(String(64), unique=True, nullable=False)
     key_prefix: Mapped[str] = mapped_column(String(20), nullable=False)
     model_allowlist: Mapped[list[str] | None] = mapped_column(JSON, nullable=True)
-    budget_limit_cents: Mapped[int | None] = mapped_column(Integer, nullable=True)
+    # BIGINT (not Integer): a client-supplied value up to the microcent scale
+    # can exceed a 32-bit int4 on Postgres, which would otherwise 500 on insert.
+    budget_limit_cents: Mapped[int | None] = mapped_column(BigInteger, nullable=True)
+    # Running lifetime spend in microcents. Maintained transactionally by
+    # spend.charge_budget: a single atomic UPDATE adds the actual cost and
+    # refuses to let the counter exceed budget_limit_cents, so the cap holds
+    # even under concurrent requests for the same key.
+    spent_microcents: Mapped[int] = mapped_column(
+        BigInteger, nullable=False, server_default="0", default=0
+    )
     is_active: Mapped[bool] = mapped_column(Boolean, server_default="true")
     last_used_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
     revoked_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)

packages/db/models/request_log.py

@@ -12,6 +12,7 @@ class RequestLog(Base, UUIDMixin, SoftDeleteMixin):
     __tablename__ = "requests_log"
     __table_args__ = (
         Index("ix_requests_log_ws_created", "workspace_id", "created_at"),
+        Index("ix_requests_log_api_key_spend", "api_key_id", "is_deleted"),
     )
 
     workspace_id: Mapped[str] = mapped_column(String(36), nullable=False, index=True)