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

packages/db/migrate.py

@@ -14,26 +14,29 @@
 
 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
 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
-    )
+    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) -> None:
+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
@@ -43,10 +46,16 @@ async def _apply_ddl(conn, statement: str) -> None:
     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():
-            await conn.execute(text(statement))
+            if callable(statement):
+                await conn.run_sync(statement)
+            else:
+                await conn.execute(text(statement))
     except DBAPIError as err:
         if not _already_applied(err):
             raise
@@ -59,18 +68,37 @@ async def ensure_budget_columns(engine) -> None:
     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
+    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. Each step costs nothing on a database that needs none of
-    it — there the repair is a single indexed UPDATE that matches no row.
+    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 repair 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

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/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,
+    )

tests/integration/test_adapter_failure_attribution.py

@@ -27,6 +27,10 @@
     _messages_payload,
     native_client,
 )
+from tests.integration.test_budget_enforcement import (  # noqa: F401
+    _get_spent,
+    _make_budgeted_key,
+)
 
 _GEMINI_PAYLOAD = {"contents": [{"role": "user", "parts": [{"text": "hi"}]}]}
 
@@ -129,3 +133,90 @@ async def _stream_router(**kwargs):
 
     assert slow.closed is True
     assert await _log_rows() == [(499, "client_disconnect", True)]
+
+
+# ── Budget settlement on an adapter fault ──────────────────────────────
+# The adapter IS the response body, so its failure unwinds through the
+# engine's SSE generator. The engine must settle that request on what it
+# delivered — leaving the settlement "unknown" would charge a budgeted key
+# its entire remaining lifetime budget for a bug of ours.
+
+_BAD_CHUNK = {
+    "id": "chatcmpl-2", "object": "chat.completion.chunk", "model": "gpt-4o-mini",
+    "created": int(time.time()), "choices": "boom",
+}
+
+
+def _content_chunk(text: str) -> dict:
+    return {
+        "id": "chatcmpl-1", "object": "chat.completion.chunk", "model": "gpt-4o-mini",
+        "created": int(time.time()),
+        "choices": [{"index": 0, "delta": {"content": text}, "finish_reason": None}],
+    }
+
+
+def _budget_stream_router(fake, chunks) -> None:
+    async def _acompletion(**kwargs):
+        assert kwargs.get("stream")
+
+        async def _gen():
+            for c in chunks:
+                yield c
+
+        return _gen()
+
+    fake.acompletion = AsyncMock(side_effect=_acompletion)
+
+
+async def _log_row_for(api_key_id: str):
+    from packages.db import session as session_mod
+    from packages.db.models.request_log import RequestLog
+
+    async with session_mod._session_factory() as s:
+        return (
+            await s.execute(select(RequestLog).where(RequestLog.api_key_id == api_key_id))
+        ).scalars().one()
+
+
+async def test_budgeted_adapter_fault_after_content_charges_delivery_not_remaining(
+    native_client,
+):
+    client, fake, _root = native_client
+    from packages.db import session as session_mod
+
+    factory = session_mod._session_factory
+    # 100 cents = 1_000_000 microcents of lifetime budget.
+    key, key_id = await _make_budgeted_key(factory, budget_limit_cents=100)
+    delivered = "the quick brown fox jumps over the lazy dog. " * 200
+    _budget_stream_router(fake, [_content_chunk(delivered), _BAD_CHUNK])
+
+    r = await client.post(
+        "/v1/messages", json=_messages_payload(stream=True), headers={"x-api-key": key},
+    )
+    assert r.status_code == 200
+    await asyncio.sleep(0.1)
+
+    assert await _log_rows() == [(500, "adapter_error", True)]
+    row = await _log_row_for(key_id)
+    spent = await _get_spent(factory, key_id)
+    assert spent == row.cost_microcents  # charged == accounted
+    assert 0 < spent < 1_000_000  # priced, not the whole remaining budget
+
+
+async def test_budgeted_adapter_fault_before_content_charges_nothing(native_client):
+    """A fault with nothing delivered cost nothing, so it must bill nothing."""
+    client, fake, _root = native_client
+    from packages.db import session as session_mod
+
+    factory = session_mod._session_factory
+    key, key_id = await _make_budgeted_key(factory, budget_limit_cents=100)
+    _budget_stream_router(fake, [_BAD_CHUNK])
+
+    r = await client.post(
+        "/v1/messages", json=_messages_payload(stream=True), headers={"x-api-key": key},
+    )
+    assert r.status_code == 200
+    await asyncio.sleep(0.1)
+
+    assert await _log_rows() == [(500, "adapter_error", True)]
+    assert await _get_spent(factory, key_id) == 0