feat(budget): spend accounting schema, migration and library [1/4]

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

packages/auth/spend.py

@@ -0,0 +1,114 @@
+"""Per-key lifetime spend tracking schema and accounting primitives for ``ApiKey.budget_limit_cents``.
+
+This module provides the accounting schema foundation and atomic charge primitives
+(part 1 of the 4-part budget subsystem; request-path enforcement is wired in #161).
+
+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). `ApiKey.budget_limit_cents` is stored in cents, so every
+`cap_microcents` argument below must be that column scaled by MICROCENTS_PER_CENT
+(passing raw cents asks whether the key has spent a ten-thousandth of its budget).
+
+Actual cost is only known after the upstream call returns, so accounting 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
+
+# Defined in `packages.db.units` so the boot repair that clamps a counter to the
+# same cap scales by the identical number; re-exported here because this module
+# is the documented home of the budget accounting primitives.
+from packages.db.units import MICROCENTS_PER_CENT as MICROCENTS_PER_CENT
+
+
+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?
+
+    ``cap_microcents`` is ``ApiKey.budget_limit_cents`` scaled by
+    ``MICROCENTS_PER_CENT``, not the column itself — passing the raw cents value
+    asks whether the key has spent a ten-thousandth of its budget.
+    """
+    spent = await read_spent(db, api_key_id)
+    return spent >= 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).
+
+    Both UPDATEs run with ``synchronize_session=False``: the session default
+    evaluates the SET in Python against any identity-map copy of the row and
+    marks it dirty, so a caller that loaded the key before this charge (the
+    documented #161 flow: validate_api_key loaded it, same session commits the
+    charge) would flush that stale computed value over the DB's atomic result —
+    silently dropping a concurrent charge. The DB row stays the only source of
+    truth; read the counter back with ``read_spent`` if you need it.
+    """
+    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)
+        .execution_options(synchronize_session=False)
+    )
+    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)
+        .execution_options(synchronize_session=False)
+    )
+    if commit:
+        await db.commit()
+    return False

packages/db/migrate.py

@@ -0,0 +1,149 @@
+"""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` runs at boot, after `create_all`, on every start: it
+inspects the live schema, only acts where the change is missing, tolerates
+another process racing it to the same change, and re-attempts a repair that an
+earlier boot applied only halfway.
+"""
+
+from __future__ import annotations
+
+from sqlalchemy import BigInteger, inspect, text
+from sqlalchemy.exc import DBAPIError
+
+from packages.db.units import MICROCENTS_PER_CENT
+
+
+def _already_applied(err: DBAPIError) -> bool:
+    """Whether a DDL failure means someone else applied the change first."""
+    msg = str(err).lower()
+    return (
+        "already exists" in msg
+        or "duplicate column" in msg
+        # Concurrent CREATE INDEX on Postgres can lose the race at the catalog
+        # insert rather than the IF NOT EXISTS probe, surfacing as a verror on
+        # pg_class's unique index instead of the usual "already exists".
+        or "pg_class_relname_nsp_index" in msg
+    )
+
+
+async def _apply_ddl(conn, statement: str) -> None:
+    """Run one startup DDL statement, tolerating a boot that raced us to it.
+
+    Every worker runs this in its lifespan, so the first boot after an upgrade
+    has several processes inspecting a schema none of them has altered yet. Each
+    then issues the same statement and all but one fail — "column ... already
+    exists" on Postgres, "duplicate column name" on SQLite — which is success
+    from here, not a reason to keep the worker from booting. The failure is
+    caught inside a SAVEPOINT because on Postgres an error would otherwise abort
+    the whole transaction and take the rest of the startup with it.
+    """
+    try:
+        async with conn.begin_nested():
+            await conn.execute(text(statement))
+    except DBAPIError as err:
+        if not _already_applied(err):
+            raise
+
+
+async def ensure_budget_columns(engine) -> None:
+    """Make `api_keys.spent_microcents` correct, whatever schema it started from.
+
+    Adds the column when an upgraded volume lacks it, and repairs lifetime spend
+    against historical request logs on every boot so the `ALTER` is never mistaken
+    for proof that the seed ran. Also widens `budget_limit_cents` to BIGINT on
+    Postgres (the column is scaled into microcents for every comparison against
+    spend, and an int4 ceiling is about 214,748 dollars of lifetime budget) and
+    creates the `ix_requests_log_api_key_spend` index that create_all only builds
+    on fresh databases. Each step costs nothing on a database that needs none of
+    it — there the repair is a single indexed UPDATE that matches no row.
+    """
+    async with engine.begin() as conn:
+        cols = {
+            c["name"]: c["type"]
+            for c in await conn.run_sync(lambda sync: inspect(sync).get_columns("api_keys"))
+        }
+        is_postgres = engine.dialect.name == "postgresql"
+
+        # The model declares ix_requests_log_api_key_spend (api_key_id,
+        # is_deleted); create_all only builds it on fresh databases, so an
+        # upgraded deployment would drift. Built before the seed below, which is
+        # a correlated aggregate over requests_log and otherwise full-scans the
+        # one table that grows without bound here — once per key, inside the
+        # transaction that already holds the api_keys lock. is_deleted has
+        # existed since the first release (SoftDeleteMixin), so the index is
+        # always creatable.
+        idx = {
+            i["name"]
+            for i in await conn.run_sync(
+                lambda sync: inspect(sync).get_indexes("requests_log")
+            )
+        }
+        if "ix_requests_log_api_key_spend" not in idx:
+            await _apply_ddl(
+                conn,
+                "CREATE INDEX IF NOT EXISTS ix_requests_log_api_key_spend "
+                "ON requests_log (api_key_id, is_deleted)",
+            )
+
+        if "spent_microcents" not in cols:
+            await _apply_ddl(
+                conn,
+                "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). Deliberately unfiltered by is_deleted, unlike the
+        # analytics reads over the same table: this restores an accrued total,
+        # so counting a row a retention job has hidden can only ever make the
+        # cap tighter, while honouring the filter would hand a capped key
+        # back the spend it was capped for.
+        #
+        # Not a one-shot `= 0` backfill but a monotonic, cap-clamped repair run
+        # every boot: the counter rises to max(current, log total) and is never
+        # written past budget_limit_cents scaled to microcents or below its
+        # current value. Until #161 wires the per-request charge, request logs
+        # are the only recorded spend, so a counter seeded at an earlier boot
+        # goes stale as traffic flows and only re-aggregating closes the gap.
+        # Afterwards the clamp keeps this a no-op on correct state: a key whose
+        # final charge was clamped deliberately under-reports its log total, and
+        # raising it back to the SUM would undo the clamp (and overshoot the
+        # cap). The strict inequality means a steady-state boot matches no row.
+        #
+        # It runs beside every ALTER, not only with the column's own, because
+        # the ALTER is no record of the seed: on SQLite that DDL is durable the
+        # instant it executes while this is DML in the transaction a kill — or
+        # the `database is locked` this very aggregate provokes on an upgrade
+        # that overlaps the old machine's writes — rolls back, and a gate keyed
+        # on the column would then never retry. Capped keys only: an uncapped
+        # key is never charged, so its counter stays at zero by design and
+        # nothing reads it. The cap is scaled by MICROCENTS_PER_CENT, the same
+        # constant the request-path charge uses. GREATEST/LEAST are Postgres;
+        # SQLite's scalar MAX/MIN take the same arguments.
+        greatest, least = ("GREATEST", "LEAST") if is_postgres else ("MAX", "MIN")
+        repaired = (
+            f"{least}({greatest}(spent_microcents, ("
+            "SELECT CAST(COALESCE(SUM(cost_microcents), 0) AS BIGINT) FROM requests_log "
+            "WHERE requests_log.api_key_id = api_keys.id"
+            f")), CAST(budget_limit_cents AS BIGINT) * {MICROCENTS_PER_CENT})"
+        )
+        await conn.execute(
+            text(
+                f"UPDATE api_keys SET spent_microcents = {repaired} "
+                f"WHERE budget_limit_cents IS NOT NULL AND spent_microcents < {repaired}"
+            )
+        )
+
+        limit_type = cols.get("budget_limit_cents")
+        if is_postgres and limit_type is not None and not isinstance(limit_type, BigInteger):
+            await _apply_ddl(
+                conn, "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,17 @@ 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): this is the cap input, scaled by MICROCENTS_PER_CENT
+    # into microcents for every comparison against `spent_microcents` below, and
+    # a 32-bit int4 would ceiling a lifetime budget near 214,748 dollars.
+    budget_limit_cents: Mapped[int | None] = mapped_column(BigInteger, nullable=True)
+    # Running lifetime spend in microcents. Schema foundation for the budget subsystem
+    # (part 1/4; request-path enforcement wired in #161). `spend.charge_budget`
+    # records actual cost atomically and ensures the counter never exceeds
+    # the scaled budget_limit_cents.
+    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)

packages/db/units.py

@@ -0,0 +1,15 @@
+"""Shared monetary units for the budget subsystem.
+
+`ApiKey.budget_limit_cents` is stored in cents, but every comparison against
+`ApiKey.spent_microcents` and `RequestLog.cost_microcents` is done in
+microcents. Two places scale that column by the same factor: the request-path
+charge (`packages.auth.spend`) and the boot repair that clamps a counter to a
+key's cap (`packages.db.migrate`). The constant lives here, below both layers,
+so neither can drift from the other by a factor of 10_000.
+"""
+
+from __future__ import annotations
+
+# 1 cent = 10_000 microcents; 1 USD = 1_000_000 microcents, matching the cost
+# math in chat.py. Every cap_microcents value is this column times this number.
+MICROCENTS_PER_CENT = 10_000