feat(budget): park lost settlements durably and fold them at pre-check

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/db/migrate.py

@@ -0,0 +1,160 @@
+"""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 collections.abc import Callable
+from typing import Any
+
+from sqlalchemy import BigInteger, inspect, text
+from sqlalchemy.exc import DBAPIError
+
+from packages.db.models.budget_park import BudgetPark
+
+
+def _already_applied(err: DBAPIError) -> bool:
+    """Whether a DDL failure means someone else applied the change first."""
+    msg = str(err).lower()
+    if "already exists" in msg or "duplicate column" in msg:
+        return True
+    # Two concurrent CREATE TABLE are serialised by the catalog rather than by
+    # the wording of a complaint, so the loser gets a unique violation on
+    # pg_class/pg_type (`*_relname_nsp_index`, `*_typname_nsp_index`) instead of
+    # "already exists". Same meaning: the object is there now.
+    return "duplicate key value" in msg and "nsp_index" in msg
+
+
+async def _apply_ddl(conn, statement: str | Callable[..., Any]) -> 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.
+
+    `statement` is raw SQL, or a callable handed to `run_sync` for DDL that only
+    the dialect's own generator can emit (a `Table.create`).
+    """
+    try:
+        async with conn.begin_nested():
+            if callable(statement):
+                await conn.run_sync(statement)
+            else:
+                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 re-seeds lifetime spend
+    from 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),
+    creates the `ix_requests_log_api_key_spend` index that create_all only builds
+    on fresh databases, and creates `budget_parks` for deployments that predate
+    the durable-recovery release — a lost settlement needs somewhere every
+    worker, and every reboot, can see it. Each step costs nothing on a database
+    that needs none of it; there the seed is a single indexed UPDATE that matches
+    no row.
+    """
+    async with engine.begin() as conn:
+        tables = set(
+            await conn.run_sync(lambda sync: inspect(sync).get_table_names())
+        )
+        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"
+
+        if BudgetPark.__tablename__ not in tables:
+            # `create_all` covers fresh databases; this covers upgrades whose
+            # schema predates the table. `checkfirst` re-reads the catalog, and
+            # that read cannot see another boot's uncommitted CREATE — so two
+            # workers both get here and one loses anyway. It goes through
+            # `_apply_ddl` for the same reason the ALTER does: losing that race
+            # has to count as having won, and the savepoint is what keeps the
+            # error from aborting the transaction the rest of startup runs in.
+            await _apply_ddl(
+                conn,
+                lambda sync: BudgetPark.__table__.create(sync, checkfirst=True),
+            )
+
+        # 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.
+        #
+        # Every boot, not only beside the ALTER, 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. It is idempotent because a log row and its
+        # charge are one commit, so this SUM *is* the lifetime counter and a key
+        # already holding spend has nothing to restore. Capped keys only: an
+        # uncapped key is never charged, so its counter stays at zero by design
+        # and nothing reads it.
+        await conn.execute(
+            text(
+                "UPDATE api_keys SET spent_microcents = ("
+                "  SELECT CAST(COALESCE(SUM(cost_microcents), 0) AS BIGINT) FROM requests_log "
+                "  WHERE requests_log.api_key_id = api_keys.id"
+                ") WHERE spent_microcents = 0 AND budget_limit_cents IS NOT NULL"
+            )
+        )
+
+        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/__init__.py

@@ -2,6 +2,7 @@
 
 from packages.db.models.api_key import ApiKey
 from packages.db.models.base import Base, SoftDeleteMixin, TimestampMixin, UUIDMixin
+from packages.db.models.budget_park import BudgetPark
 from packages.db.models.provider_key import ProviderKey
 from packages.db.models.quality_score_override import QualityScoreOverride
 from packages.db.models.quality_score_snapshot import QualityScoreSnapshot
@@ -15,6 +16,7 @@
     "TimestampMixin",
     "UUIDMixin",
     "ApiKey",
+    "BudgetPark",
     "ProviderKey",
     "QualityScoreOverride",
     "QualityScoreSnapshot",

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/budget_park.py

@@ -0,0 +1,51 @@
+"""Unsettled budget obligations — delivered spend the ledger never recorded.
+
+A settlement that gives up after every retry leaves a delivered cost with no
+record anywhere: the log row and the charge are one transaction, so both roll
+back and `spent_microcents` never moves. The obligation is parked here, one row
+per settlement, so it outlives the process that lost it and is visible to every
+worker behind the same database.
+
+Rows are keyed by the settlement's `trace_id`, which makes every park write
+idempotent: a commit that applied but whose ack was lost retries into the same
+primary key instead of recording the obligation twice, and a write that fails
+outright is checked for having landed anyway before the process falls back to
+holding it in memory. A fold that bills a row either deletes it or shrinks it to
+what the cap could not absorb, in the same transaction that moves
+`spent_microcents`, and it drops the writer's memory hold for any row it fully
+billed.
+
+What that leaves is a double charge needing three faults at once — the ack lost,
+the durability probe failing alongside it, and another worker folding the row
+before this one retries — at which point the extra charge lands on a key that had
+already breached its cap. Closing that last window means a fold leaving a tombstone
+behind instead of deleting, so a re-file always collides with something; that is a
+second state the queue has to drain, and it is not worth its weight here.
+"""
+
+from datetime import datetime, timezone
+
+from sqlalchemy import BigInteger, DateTime, String
+from sqlalchemy.orm import Mapped, mapped_column
+
+from packages.db.models.base import Base, TimestampMixin, UUIDMixin
+
+
+class BudgetPark(Base, UUIDMixin, TimestampMixin):
+    __tablename__ = "budget_parks"
+
+    trace_id: Mapped[str] = mapped_column(String(36), nullable=False, unique=True)
+    api_key_id: Mapped[str] = mapped_column(String(36), nullable=False, index=True)
+    microcents: Mapped[int] = mapped_column(BigInteger, nullable=False)
+    # Overrides TimestampMixin's column, whose `server_default=func.now()` is
+    # CURRENT_TIMESTAMP: one second wide on SQLite, where a recovered outage
+    # re-files a whole batch of obligations in a single pre-check and every row
+    # in it ties. The fold bills oldest debt first, and two workers have to
+    # agree on which row the partial one is, so the stamp needs enough
+    # resolution to settle that on its own instead of falling through to the
+    # `trace_id` tiebreak — which is a uuid4, and so picks at random.
+    created_at: Mapped[datetime] = mapped_column(
+        DateTime(timezone=True),
+        default=lambda: datetime.now(timezone.utc),
+        nullable=False,
+    )

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)