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,22 @@ 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:
+        """Atomically record `actual_microcents` of spend against the cap.
+
+        Idempotent in-process: only the first successful charge sets the flag.
+        `session` is the DB session to run the charge on (the request session, or
+        the dedicated log session for the streaming path). When `commit` is False
+        the UPDATE is executed but not committed, so the caller can commit it in
+        the same transaction as the request-log write — closing the fail-open
+        window where the log landed but the charge was lost.
+        """
+        cap = getattr(kc, "_budget_cap", None)
+        if cap is None or getattr(kc, "_budget_settled", False):
+            return
+        await charge_budget(session, str(kc.key_id), cap, actual_microcents, commit=commit)
+        kc._budget_settled = True
+
     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 +437,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 +524,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))
@@ -542,6 +578,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 +618,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.
@@ -659,28 +703,50 @@ 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.
+
+                    On a stream that ended without a terminal [DONE] the real cost
+                    is unknown, so charge the full remaining allowance (fail-closed)
+                    instead of releasing the budget and letting a client bypass the
+                    cap by hanging up before the usage frame. For a completed stream
+                    the actual recorded cost is charged.
+                    """
+                    actual = row_values.get("cost_microcents") or 0
+                    if not stream_completed:
+                        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, then charge the budget.
+
+                    The log row is committed first so it is durable even if the
+                    subsequent budget charge hits a transient error. If the charge
+                    fails, the retry path (reached with the log already persisted)
+                    re-attempts ONLY the charge — it never re-inserts the row and
+                    never skips a still-outstanding charge, so spend is never
+                    silently dropped (fail-open). `_settle_budget` sets its settled
+                    flag only after a successful charge, so a failed charge stays
+                    retryable.
                     """
                     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)):
+                            if getattr(kc, "_budget_settled", False):
+                                return
+                            await _settle_budget(db, _settlement_amount())
                             return
                         db.add(log)
                         try:
                             await db.commit()
+                            await _settle_budget(db, _settlement_amount())
                         except Exception:
                             try:
                                 await db.rollback()
@@ -690,10 +756,15 @@ 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)):
+                            if getattr(kc, "_budget_settled", False):
+                                return
+                            # Log already durable; re-run only the charge.
+                            await _settle_budget(s, _settlement_amount())
                             return
                         s.add(log)
                         await s.commit()
+                        await _settle_budget(s, _settlement_amount())
                     finally:
                         try:
                             await s.close()
@@ -774,6 +845,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
@@ -874,6 +946,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
@@ -952,6 +1030,7 @@ async def _commit_row(*, retry: bool) -> None:
         )
         db.add(log)
         try:
+            await _settle_budget(db, log.cost_microcents, commit=False)
             await db.commit()
         except Exception as commit_err:
             logger.warning("request_log_commit_failed", error=str(commit_err))

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

packages/db/migrate.py

@@ -0,0 +1,54 @@
+"""Idempotent startup schema migrations for columns added after the first release.
+
+`Base.metadata.create_all` creates new tables but never alters existing ones, so a
+deployment that already ran a release (a SQLite named volume, a fly.io/Postgres
+volume) keeps an `api_keys` table without the `spent_microcents` column. After an
+upgrade the ORM would then `SELECT` every mapped column and hit "no such column"
+on every authenticated request — a 503 for the whole API.
+
+`ensure_budget_columns` is run once at boot, after `create_all`, and is safe to
+call on every start: it inspects the live schema and only acts when the column is
+missing.
+"""
+
+from __future__ import annotations
+
+from sqlalchemy import inspect, text
+
+
+async def ensure_budget_columns(engine) -> None:
+    """Add `spent_microcents` to `api_keys` if absent, seeded from request history.
+
+    Also widens `budget_limit_cents` to BIGINT on Postgres (the microcent scale
+    can exceed int4). Both are no-ops on a fresh database.
+    """
+    async with engine.begin() as conn:
+        cols = {
+            c["name"]
+            for c in await conn.run_sync(lambda sync: inspect(sync).get_columns("api_keys"))
+        }
+        is_postgres = engine.dialect.name == "postgresql"
+
+        if "spent_microcents" not in cols:
+            await conn.execute(
+                text(
+                    "ALTER TABLE api_keys ADD COLUMN spent_microcents BIGINT "
+                    "NOT NULL DEFAULT 0"
+                )
+            )
+            # Seed lifetime spend from historical request logs so an existing key's
+            # cap is not silently reset to zero (which would re-grant a leaked key
+            # a full new budget).
+            await conn.execute(
+                text(
+                    "UPDATE api_keys SET spent_microcents = ("
+                    "  SELECT COALESCE(SUM(cost_microcents), 0) FROM requests_log "
+                    "  WHERE requests_log.api_key_id = api_keys.id"
+                    ") WHERE spent_microcents = 0"
+                )
+            )
+
+        if is_postgres and "budget_limit_cents" in cols:
+            await conn.execute(
+                text("ALTER TABLE api_keys ALTER COLUMN budget_limit_cents TYPE BIGINT")
+            )

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)