From 2934804f80ffae8d2421771b8fa129a83a3bd3af Mon Sep 17 00:00:00 2001 From: Hasit Bhatt Date: Thu, 24 Sep 2026 20:41:25 -0700 Subject: [PATCH 01/23] feat(budget): add spent accounting, startup migration, and atomic charge --- app/main.py | 7 ++ packages/auth/spend.py | 91 +++++++++++++++++ packages/db/migrate.py | 98 ++++++++++++++++++ packages/db/models/api_key.py | 13 ++- packages/db/models/request_log.py | 1 + tests/unit/test_budget_migration.py | 151 ++++++++++++++++++++++++++++ tests/unit/test_budget_spend.py | 90 +++++++++++++++++ 7 files changed, 449 insertions(+), 2 deletions(-) create mode 100644 packages/auth/spend.py create mode 100644 packages/db/migrate.py create mode 100644 tests/unit/test_budget_migration.py create mode 100644 tests/unit/test_budget_spend.py diff --git a/app/main.py b/app/main.py index ff0d3d3..d2d2960 100644 --- a/app/main.py +++ b/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 diff --git a/packages/auth/spend.py b/packages/auth/spend.py new file mode 100644 index 0000000..44a799f --- /dev/null +++ b/packages/auth/spend.py @@ -0,0 +1,91 @@ +"""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?""" + 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). + """ + 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 diff --git a/packages/db/migrate.py b/packages/db/migrate.py new file mode 100644 index 0000000..c5dddce --- /dev/null +++ b/packages/db/migrate.py @@ -0,0 +1,98 @@ +"""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, only acts when the change is +missing, and tolerates another process racing it to the same change. +""" + +from __future__ import annotations + +from sqlalchemy import inspect, text +from sqlalchemy.exc import DBAPIError + + +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 + + +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: + """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) and creates the `ix_requests_log_api_key_spend` index that + create_all only builds on fresh databases. All steps 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 _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). + 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 _apply_ddl( + conn, "ALTER TABLE api_keys ALTER COLUMN budget_limit_cents TYPE BIGINT" + ) + + # 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. 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)", + ) diff --git a/packages/db/models/api_key.py b/packages/db/models/api_key.py index a96d99a..0dfd8b7 100644 --- a/packages/db/models/api_key.py +++ b/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) diff --git a/packages/db/models/request_log.py b/packages/db/models/request_log.py index 9871610..84c7c06 100644 --- a/packages/db/models/request_log.py +++ b/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) diff --git a/tests/unit/test_budget_migration.py b/tests/unit/test_budget_migration.py new file mode 100644 index 0000000..2200267 --- /dev/null +++ b/tests/unit/test_budget_migration.py @@ -0,0 +1,151 @@ +"""Upgrade-path coverage for ensure_budget_columns (packages/db/migrate.py). + +create_all never alters existing tables, so an upgraded deployment starts from +a legacy schema: api_keys without spent_microcents and requests_log without the +spend index. These tests build that legacy state by creating the real schema, +seeding rows, then dropping exactly what the pre-budget release lacked — and +pin that the startup migration restores it: (a) the column seeded from +historical request-log spend, (b) the composite index, (c) idempotency across +repeated boots, (d) the ORM (and thus auth) working again. +""" + +from __future__ import annotations + +from sqlalchemy import inspect as sa_inspect +from sqlalchemy import select, text +from sqlalchemy.ext.asyncio import async_sessionmaker + +from packages.db.engine import build_engine +from packages.db.migrate import ensure_budget_columns +from packages.db.models.api_key import ApiKey +from packages.db.models.base import Base +from packages.db.models.request_log import RequestLog + + +async def _legacy_deploy_engine(tmp_sqlite_url): + """Engine over a DB shaped like the last released schema, with one + budgeted key that already burned 2500 microcents of history.""" + engine = build_engine(tmp_sqlite_url) + async with engine.begin() as conn: + await conn.run_sync(Base.metadata.create_all) + factory = async_sessionmaker(engine, expire_on_commit=False) + async with factory() as s: + row = ApiKey( + workspace_id="w1", name="leaked-then-capped", key_hash="h", + key_prefix="p", budget_limit_cents=100, spent_microcents=2500, + ) + s.add(row) + await s.commit() + await s.refresh(row) + s.add(RequestLog( + workspace_id="w1", api_key_id=row.id, + trace_id="t1", model_requested="gpt-4o-mini", + model_resolved="gpt-4o-mini", provider="openai", + routing_strategy="balanced", input_tokens=5, output_tokens=2, + cost_microcents=2500, latency_ms=10, status_code=200, + )) + await s.commit() + # Downgrade to the pre-budget schema: drop the column and the index that + # only the new release's metadata declares. + async with engine.begin() as conn: + await conn.execute(text("DROP INDEX IF EXISTS ix_requests_log_api_key_spend")) + await conn.execute(text("ALTER TABLE api_keys DROP COLUMN spent_microcents")) + await engine.dispose() + return build_engine(tmp_sqlite_url) + + +async def test_ensure_budget_columns_upgrades_legacy_schema(tmp_sqlite_url): + engine = await _legacy_deploy_engine(tmp_sqlite_url) + try: + await ensure_budget_columns(engine) + + async with engine.connect() as conn: + # Seeded from history: a key that already burned 2500 microcents + # must not get a fresh full budget on upgrade. + spent = await conn.scalar( + text("SELECT spent_microcents FROM api_keys WHERE workspace_id = 'w1'") + ) + assert spent == 2500 + + idx_names = { + i["name"] + for i in await conn.run_sync( + lambda sync: sa_inspect(sync).get_indexes("requests_log") + ) + } + assert "ix_requests_log_api_key_spend" in idx_names + + # Idempotent across restarts: a second boot changes nothing. + await ensure_budget_columns(engine) + async with engine.connect() as conn: + assert await conn.scalar( + text("SELECT spent_microcents FROM api_keys WHERE workspace_id = 'w1'") + ) == 2500 + finally: + await engine.dispose() + + +async def test_orm_reads_work_after_upgrade(tmp_sqlite_url): + # The original 503 failure mode: the ORM SELECTs every mapped column, so + # without the migration every authenticated request broke on the upgraded + # database. After ensure_budget_columns the ApiKey select must succeed. + engine = await _legacy_deploy_engine(tmp_sqlite_url) + try: + await ensure_budget_columns(engine) + factory = async_sessionmaker(engine, expire_on_commit=False) + async with factory() as s: + row = (await s.execute(select(ApiKey))).scalar_one() + assert row.spent_microcents == 2500 + assert row.budget_limit_cents == 100 + finally: + await engine.dispose() + + +async def test_ensure_budget_columns_survives_a_racing_boot(tmp_sqlite_url, monkeypatch): + """The loser of a concurrent-boot ALTER still boots, and still seeds. + + Every worker runs this at startup, and the first boot after an upgrade + starts several of them at once against one database. Both inspect the + schema before either alters it, so the loser's ALTER meets a column that + appeared in between and the driver rejects it with "duplicate column + name" — which used to escape into the lifespan and keep that worker down. + """ + import packages.db.migrate as migrate + + engine = await _legacy_deploy_engine(tmp_sqlite_url) + try: + # The other boot gets there first: the column is committed, so this + # process's inspection is now stale relative to the schema. + async with engine.begin() as conn: + await conn.execute( + text( + "ALTER TABLE api_keys ADD COLUMN spent_microcents BIGINT " + "NOT NULL DEFAULT 0" + ) + ) + inspector = sa_inspect(engine.sync_engine) + + class _StaleSchema: + def __init__(self, inner): + self._inner = inner + + def get_columns(self, table_name): + return [ + c for c in self._inner.get_columns(table_name) + if c["name"] != "spent_microcents" + ] + + def __getattr__(self, name): + return getattr(self._inner, name) + + monkeypatch.setattr(migrate, "inspect", lambda _sync: _StaleSchema(inspector)) + await ensure_budget_columns(engine) + + async with engine.connect() as conn: + # It went on to seed the column the winner added — before the fix + # the boot died on the ALTER and never reached this. + assert await conn.scalar( + text("SELECT spent_microcents FROM api_keys WHERE workspace_id = 'w1'") + ) == 2500 + finally: + await engine.dispose() diff --git a/tests/unit/test_budget_spend.py b/tests/unit/test_budget_spend.py new file mode 100644 index 0000000..19d6cee --- /dev/null +++ b/tests/unit/test_budget_spend.py @@ -0,0 +1,90 @@ +"""Unit tests for packages.auth.spend — atomic budget charge under a hard cap.""" + +import asyncio + +import pytest + +from packages.auth.spend import ( + MICROCENTS_PER_CENT, + charge_budget, + is_exhausted, + read_spent, +) + + +@pytest.fixture +async def key(db_session): + from packages.db.models.api_key import ApiKey + + k = ApiKey(workspace_id="default", name="a", key_hash="h-a", key_prefix="p-a") + db_session.add(k) + await db_session.flush() + return k + + +async def test_charge_within_cap_advances_counter(db_session, key): + cap = 10_000 + assert await charge_budget(db_session, key.id, cap, 300) is True + assert await read_spent(db_session, key.id) == 300 + + +async def test_charge_past_cap_clamps_and_reports_false(db_session, key): + cap = 10_000 + # A single request whose cost exceeds the remaining budget must not push the + # counter past the cap; it is clamped and reported as over-budget. + assert await charge_budget(db_session, key.id, cap, 50_000) is False + assert await read_spent(db_session, key.id) == cap + assert await is_exhausted(db_session, key.id, cap) is True + + +async def test_is_exhausted_false_below_cap(db_session, key): + cap = 10_000 + await charge_budget(db_session, key.id, cap, 9_000) + assert await is_exhausted(db_session, key.id, cap) is False + await charge_budget(db_session, key.id, cap, 2_000) # clamps at 10_000 + assert await is_exhausted(db_session, key.id, cap) is True + + +async def test_concurrent_charges_never_exceed_cap(db_session, key): + """Two simultaneous charges that together would exceed the cap are bounded. + + Build two independent sessions against the same engine so the atomic + `UPDATE ... WHERE spent + actual <= cap` guard is exercised for real. + Exactly one fits; the other is clamped. The counter ends at `cap`, never + above it. + """ + from sqlalchemy.ext.asyncio import async_sessionmaker + + from packages.db.engine import build_engine + from packages.db.models.base import Base + + engine = build_engine("sqlite+aiosqlite:///:memory:") + async with engine.begin() as conn: + await conn.run_sync(Base.metadata.create_all) + factory = async_sessionmaker(engine, expire_on_commit=False) + async with factory() as s: + from packages.db.models.api_key import ApiKey + + k = ApiKey(workspace_id="default", name="race", key_hash="h-race", key_prefix="p-race") + s.add(k) + await s.commit() + await s.refresh(k) + + cap = 10_000 + # Each request costs 6_000; both cannot fit under a 10_000 cap. Use two + # independent sessions so the atomic `UPDATE ... WHERE spent + actual <= cap` + # guard is exercised for real. + async with factory() as s1, factory() as s2: + r1, r2 = await asyncio.gather( + charge_budget(s1, k.id, cap, 6_000), + charge_budget(s2, k.id, cap, 6_000), + ) + final = (await read_spent(s1, k.id)) or (await read_spent(s2, k.id)) + await engine.dispose() + # One succeeds, the other is clamped — but the counter never exceeds cap. + assert (r1 is True) ^ (r2 is True) or (r1 is False and r2 is False) + assert final <= cap + + +def test_microcent_conversion_constant(): + assert MICROCENTS_PER_CENT == 10_000 From 228396e0a01fced212283bb706e16332328952bf Mon Sep 17 00:00:00 2001 From: Hasit Bhatt Date: Thu, 24 Sep 2026 22:16:57 -0700 Subject: [PATCH 02/23] fix(budget): let the boot seed use the index it aggregates through MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The upgrade seed is a correlated SUM over requests_log, and the index that would serve it was created further down the same function — so the one boot that runs the seed was also the one that could not use the index, and later boots do neither. Build the index first and pin the order with a test. Also states the two contracts the schema leaves implicit: the seed counts soft-deleted request rows on purpose, because restoring an accrued lifetime total can only ever tighten a cap, and cap_microcents is budget_limit_cents scaled to microcents rather than the column itself. --- packages/auth/spend.py | 10 +++++-- packages/db/migrate.py | 44 +++++++++++++++++------------ packages/db/models/api_key.py | 10 ++++--- tests/unit/test_budget_migration.py | 31 ++++++++++++++++++++ 4 files changed, 71 insertions(+), 24 deletions(-) diff --git a/packages/auth/spend.py b/packages/auth/spend.py index 44a799f..730b255 100644 --- a/packages/auth/spend.py +++ b/packages/auth/spend.py @@ -2,7 +2,8 @@ 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). +cost math). `ApiKey.budget_limit_cents` is stored in cents, so every +`cap_microcents` argument below is that column scaled by MICROCENTS_PER_CENT. 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 @@ -43,7 +44,12 @@ async def read_spent(db: AsyncSession, api_key_id: str) -> int: 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?""" + """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 diff --git a/packages/db/migrate.py b/packages/db/migrate.py index c5dddce..26bffd5 100644 --- a/packages/db/migrate.py +++ b/packages/db/migrate.py @@ -57,6 +57,27 @@ async def ensure_budget_columns(engine) -> None: } 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, @@ -65,7 +86,11 @@ async def ensure_budget_columns(engine) -> None: ) # 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). + # 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. await conn.execute( text( "UPDATE api_keys SET spent_microcents = (" @@ -79,20 +104,3 @@ async def ensure_budget_columns(engine) -> None: await _apply_ddl( conn, "ALTER TABLE api_keys ALTER COLUMN budget_limit_cents TYPE BIGINT" ) - - # 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. 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)", - ) diff --git a/packages/db/models/api_key.py b/packages/db/models/api_key.py index 0dfd8b7..e34eb14 100644 --- a/packages/db/models/api_key.py +++ b/packages/db/models/api_key.py @@ -21,10 +21,12 @@ class ApiKey(Base, UUIDMixin, TimestampMixin, SoftDeleteMixin): # 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. + # Running lifetime spend in microcents. `spend.charge_budget` is the only + # writer: 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. This column is the state that protocol needs — + # a caller enforces by checking `is_exhausted` before dispatch and charging + # after, so the cap is only as live as the paths that route through it. spent_microcents: Mapped[int] = mapped_column( BigInteger, nullable=False, server_default="0", default=0 ) diff --git a/tests/unit/test_budget_migration.py b/tests/unit/test_budget_migration.py index 2200267..8ca92ea 100644 --- a/tests/unit/test_budget_migration.py +++ b/tests/unit/test_budget_migration.py @@ -101,6 +101,37 @@ async def test_orm_reads_work_after_upgrade(tmp_sqlite_url): await engine.dispose() +async def test_seed_runs_after_the_index_it_aggregates_through(tmp_sqlite_url): + """The first post-upgrade boot must not full-scan requests_log per key. + + The seed is a correlated SUM over requests_log, and the only thing that + makes it cheap is ix_requests_log_api_key_spend. Building that index after + the seed means the one boot that runs the seed — this function's whole + reason to exist — is also the one that cannot use the index, and every + later boot skips both. + """ + from sqlalchemy import event + + engine = await _legacy_deploy_engine(tmp_sqlite_url) + order: list[str] = [] + try: + def _record(conn, cursor, statement, parameters, context, executemany): + if "ix_requests_log_api_key_spend" in statement: + order.append("index") + elif "UPDATE api_keys SET spent_microcents" in statement: + order.append("seed") + + event.listen(engine.sync_engine, "before_cursor_execute", _record) + try: + await ensure_budget_columns(engine) + finally: + event.remove(engine.sync_engine, "before_cursor_execute", _record) + + assert order == ["index", "seed"] + finally: + await engine.dispose() + + async def test_ensure_budget_columns_survives_a_racing_boot(tmp_sqlite_url, monkeypatch): """The loser of a concurrent-boot ALTER still boots, and still seeds. From 033c637ba926af00d284b2261418cb1bd58bf0df Mon Sep 17 00:00:00 2001 From: Hasit Bhatt Date: Fri, 25 Sep 2026 02:10:11 -0700 Subject: [PATCH 03/23] fix(budget): make the lifetime-spend seed a repair, not a one-shot MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The seed ran only inside the branch that added the column, and the two statements do not vouch for each other: on SQLite the ALTER is durable the instant it executes while the seed 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. Gating on the column's absence made that half-applied boot the only one that could ever have seeded, so every key predating the release kept a full fresh allowance forever, silently, which is exactly the outcome the seed exists to prevent. It runs on every boot now, restricted to keys that hold a cap, and the statement is idempotent because a log row and its charge are one commit — a key already holding spend has nothing to restore. Three more from the same review: - Gate the Postgres BIGINT widen on the reflected type rather than the column's name, which was present forever and so took ACCESS EXCLUSIVE on api_keys at every start. - Correct the model comment justifying that widen with a client-supplied budget no route accepts, in the wrong unit. - Make the upgrade tests able to fail: the legacy fixture had one key and one log row, so a seed that dropped its correlation predicate and stamped every key with the table total passed. It now has three keys with three histories, pins the half-applied boot above, and the concurrency test runs over a file instead of `:memory:`'s StaticPool, where two "independent" sessions shared one connection and the atomic guard never met a concurrent writer. --- packages/db/migrate.py | 68 +++++++++------ packages/db/models/api_key.py | 5 +- tests/unit/test_budget_migration.py | 124 ++++++++++++++++++++-------- tests/unit/test_budget_spend.py | 31 +++---- 4 files changed, 151 insertions(+), 77 deletions(-) diff --git a/packages/db/migrate.py b/packages/db/migrate.py index 26bffd5..d934163 100644 --- a/packages/db/migrate.py +++ b/packages/db/migrate.py @@ -6,14 +6,15 @@ 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, only acts when the change is -missing, and tolerates another process racing it to the same change. +`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 inspect, text +from sqlalchemy import BigInteger, inspect, text from sqlalchemy.exc import DBAPIError @@ -43,16 +44,20 @@ async def _apply_ddl(conn, statement: str) -> None: async def ensure_budget_columns(engine) -> None: - """Add `spent_microcents` to `api_keys` if absent, seeded from request history. + """Make `api_keys.spent_microcents` correct, whatever schema it started from. - Also widens `budget_limit_cents` to BIGINT on Postgres (the microcent scale - can exceed int4) and creates the `ix_requests_log_api_key_spend` index that - create_all only builds on fresh databases. All steps are no-ops on a fresh - database. + 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) 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 seed is a single indexed UPDATE that matches no row. """ async with engine.begin() as conn: cols = { - c["name"] + 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" @@ -84,23 +89,36 @@ async def ensure_budget_columns(engine) -> None: "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. - 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" - ) + + # 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 COALESCE(SUM(cost_microcents), 0) FROM requests_log " + " WHERE requests_log.api_key_id = api_keys.id" + ") WHERE spent_microcents = 0 AND budget_limit_cents IS NOT NULL" ) + ) - if is_postgres and "budget_limit_cents" in cols: + 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" ) diff --git a/packages/db/models/api_key.py b/packages/db/models/api_key.py index e34eb14..b93668f 100644 --- a/packages/db/models/api_key.py +++ b/packages/db/models/api_key.py @@ -18,8 +18,9 @@ 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) - # 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. + # 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. `spend.charge_budget` is the only # writer: a single atomic UPDATE adds the actual cost and refuses to let the diff --git a/tests/unit/test_budget_migration.py b/tests/unit/test_budget_migration.py index 8ca92ea..dc6c19a 100644 --- a/tests/unit/test_budget_migration.py +++ b/tests/unit/test_budget_migration.py @@ -4,9 +4,10 @@ a legacy schema: api_keys without spent_microcents and requests_log without the spend index. These tests build that legacy state by creating the real schema, seeding rows, then dropping exactly what the pre-budget release lacked — and -pin that the startup migration restores it: (a) the column seeded from -historical request-log spend, (b) the composite index, (c) idempotency across -repeated boots, (d) the ORM (and thus auth) working again. +pin that the startup migration restores it: (a) each key's counter seeded from +its own request-log history, (b) the composite index, (c) idempotency across +repeated boots, (d) the ORM (and thus auth) working again, and (e) a seed that +survives the boot that was supposed to run it dying halfway. """ from __future__ import annotations @@ -23,28 +24,37 @@ async def _legacy_deploy_engine(tmp_sqlite_url): - """Engine over a DB shaped like the last released schema, with one - budgeted key that already burned 2500 microcents of history.""" + """Engine over a DB shaped like the last released schema. + + Three keys, each with its own request history, so the seed's own predicates + are the only thing that can make the numbers come out right: `w1` is capped + and burned 2500 microcents, `w2` is capped and burned 700 (drop the + `api_key_id` correlation and both read the table total), and `w3` is uncapped + with 400 of history that must stay uncounted — nothing ever charges an + uncapped key, so its counter means nothing until it gets a cap. + """ engine = build_engine(tmp_sqlite_url) async with engine.begin() as conn: await conn.run_sync(Base.metadata.create_all) factory = async_sessionmaker(engine, expire_on_commit=False) async with factory() as s: - row = ApiKey( - workspace_id="w1", name="leaked-then-capped", key_hash="h", - key_prefix="p", budget_limit_cents=100, spent_microcents=2500, - ) - s.add(row) - await s.commit() - await s.refresh(row) - s.add(RequestLog( - workspace_id="w1", api_key_id=row.id, - trace_id="t1", model_requested="gpt-4o-mini", - model_resolved="gpt-4o-mini", provider="openai", - routing_strategy="balanced", input_tokens=5, output_tokens=2, - cost_microcents=2500, latency_ms=10, status_code=200, - )) - await s.commit() + for ws, cents, burned in (("w1", 100, 2500), ("w2", 50, 700), ("w3", None, 400)): + row = ApiKey( + workspace_id=ws, name=f"{ws}-key", key_hash=f"h-{ws}", + key_prefix=f"p-{ws}", budget_limit_cents=cents, + spent_microcents=burned, + ) + s.add(row) + await s.commit() + await s.refresh(row) + s.add(RequestLog( + workspace_id=ws, api_key_id=row.id, + trace_id=f"t-{ws}", model_requested="gpt-4o-mini", + model_resolved="gpt-4o-mini", provider="openai", + routing_strategy="balanced", input_tokens=5, output_tokens=2, + cost_microcents=burned, latency_ms=10, status_code=200, + )) + await s.commit() # Downgrade to the pre-budget schema: drop the column and the index that # only the new release's metadata declares. async with engine.begin() as conn: @@ -54,19 +64,29 @@ async def _legacy_deploy_engine(tmp_sqlite_url): return build_engine(tmp_sqlite_url) +async def _spent(engine, workspace_id: str) -> int: + async with engine.connect() as conn: + return await conn.scalar( + text("SELECT spent_microcents FROM api_keys WHERE workspace_id = :ws"), + {"ws": workspace_id}, + ) + + async def test_ensure_budget_columns_upgrades_legacy_schema(tmp_sqlite_url): engine = await _legacy_deploy_engine(tmp_sqlite_url) try: await ensure_budget_columns(engine) - async with engine.connect() as conn: - # Seeded from history: a key that already burned 2500 microcents - # must not get a fresh full budget on upgrade. - spent = await conn.scalar( - text("SELECT spent_microcents FROM api_keys WHERE workspace_id = 'w1'") - ) - assert spent == 2500 + # Seeded from each key's own history: a key that already burned 2500 + # microcents must not get a fresh full budget on upgrade, and the key + # next to it must not be billed for its neighbour's traffic. + assert await _spent(engine, "w1") == 2500 + assert await _spent(engine, "w2") == 700 + # The uncapped key is never charged, so restoring a total for it would + # aggregate the one unbounded table to compute a number nothing reads. + assert await _spent(engine, "w3") == 0 + async with engine.connect() as conn: idx_names = { i["name"] for i in await conn.run_sync( @@ -77,10 +97,38 @@ async def test_ensure_budget_columns_upgrades_legacy_schema(tmp_sqlite_url): # Idempotent across restarts: a second boot changes nothing. await ensure_budget_columns(engine) - async with engine.connect() as conn: - assert await conn.scalar( - text("SELECT spent_microcents FROM api_keys WHERE workspace_id = 'w1'") - ) == 2500 + assert await _spent(engine, "w1") == 2500 + assert await _spent(engine, "w2") == 700 + assert await _spent(engine, "w3") == 0 + finally: + await engine.dispose() + + +async def test_seed_repairs_a_boot_that_died_before_it(tmp_sqlite_url): + """A half-applied upgrade reseeds, on the next boot, by itself. + + The column and the seed are two statements and one is not proof of the + other: on SQLite the `ALTER` is durable the instant it executes while the + seed is DML in the transaction a kill, a dropped connection, or the 5s + `busy_timeout` this aggregate can hit rolls back. Gating the seed on the + column's absence — which is what the first version did — makes that boot the + only one that ever could have seeded, so every key that predates the release + keeps a full fresh allowance forever, silently, with no repair path. + """ + engine = await _legacy_deploy_engine(tmp_sqlite_url) + try: + # The state that boot left behind: column present, every counter zero, + # history intact. + async with engine.begin() as conn: + await conn.execute( + text("ALTER TABLE api_keys ADD COLUMN spent_microcents BIGINT " + "NOT NULL DEFAULT 0") + ) + assert await _spent(engine, "w1") == 0 + + await ensure_budget_columns(engine) + assert await _spent(engine, "w1") == 2500 + assert await _spent(engine, "w2") == 700 finally: await engine.dispose() @@ -94,7 +142,9 @@ async def test_orm_reads_work_after_upgrade(tmp_sqlite_url): await ensure_budget_columns(engine) factory = async_sessionmaker(engine, expire_on_commit=False) async with factory() as s: - row = (await s.execute(select(ApiKey))).scalar_one() + row = ( + await s.execute(select(ApiKey).where(ApiKey.workspace_id == "w1")) + ).scalar_one() assert row.spent_microcents == 2500 assert row.budget_limit_cents == 100 finally: @@ -106,9 +156,13 @@ async def test_seed_runs_after_the_index_it_aggregates_through(tmp_sqlite_url): The seed is a correlated SUM over requests_log, and the only thing that makes it cheap is ix_requests_log_api_key_spend. Building that index after - the seed means the one boot that runs the seed — this function's whole - reason to exist — is also the one that cannot use the index, and every - later boot skips both. + the seed means the boot that restores every pre-release counter — and every + later boot that finds a half-applied upgrade to repair — does it with a full + scan of the one table that grows without bound. + + The exact-list assertion also pins that the seed is emitted once per boot: + it runs on every start now, so a second copy is a second aggregate over the + same table for no reason. """ from sqlalchemy import event diff --git a/tests/unit/test_budget_spend.py b/tests/unit/test_budget_spend.py index 19d6cee..fe62e65 100644 --- a/tests/unit/test_budget_spend.py +++ b/tests/unit/test_budget_spend.py @@ -45,45 +45,46 @@ async def test_is_exhausted_false_below_cap(db_session, key): assert await is_exhausted(db_session, key.id, cap) is True -async def test_concurrent_charges_never_exceed_cap(db_session, key): +async def test_concurrent_charges_never_exceed_cap(tmp_sqlite_url): """Two simultaneous charges that together would exceed the cap are bounded. - Build two independent sessions against the same engine so the atomic - `UPDATE ... WHERE spent + actual <= cap` guard is exercised for real. - Exactly one fits; the other is clamped. The counter ends at `cap`, never - above it. + A file-backed URL, because `:memory:` hands back a `StaticPool`: both + sessions would then share one DBAPI connection, the statements would + serialise inside it, and the atomic `UPDATE ... WHERE spent + actual <= cap` + guard would never meet a concurrent writer. Over a file each session gets + its own connection, and SQLite's single writer plus the pool's busy timeout + still makes the outcome deterministic — one charge fits, the other's guard + matches no row and its clamp fills the counter to exactly `cap`. """ from sqlalchemy.ext.asyncio import async_sessionmaker from packages.db.engine import build_engine + from packages.db.models.api_key import ApiKey from packages.db.models.base import Base - engine = build_engine("sqlite+aiosqlite:///:memory:") + engine = build_engine(tmp_sqlite_url) async with engine.begin() as conn: await conn.run_sync(Base.metadata.create_all) factory = async_sessionmaker(engine, expire_on_commit=False) async with factory() as s: - from packages.db.models.api_key import ApiKey - k = ApiKey(workspace_id="default", name="race", key_hash="h-race", key_prefix="p-race") s.add(k) await s.commit() await s.refresh(k) cap = 10_000 - # Each request costs 6_000; both cannot fit under a 10_000 cap. Use two - # independent sessions so the atomic `UPDATE ... WHERE spent + actual <= cap` - # guard is exercised for real. async with factory() as s1, factory() as s2: r1, r2 = await asyncio.gather( charge_budget(s1, k.id, cap, 6_000), charge_budget(s2, k.id, cap, 6_000), ) - final = (await read_spent(s1, k.id)) or (await read_spent(s2, k.id)) + # Read the winner's outcome from a session that took part in neither charge. + async with factory() as reader: + final = await read_spent(reader, k.id) await engine.dispose() - # One succeeds, the other is clamped — but the counter never exceeds cap. - assert (r1 is True) ^ (r2 is True) or (r1 is False and r2 is False) - assert final <= cap + + assert (r1 is True) + (r2 is True) == 1 + assert final == cap def test_microcent_conversion_constant(): From d36b960829c28c362859ed623dac12bc6e2da940 Mon Sep 17 00:00:00 2001 From: Hasit Bhatt Date: Fri, 25 Sep 2026 03:48:58 -0700 Subject: [PATCH 04/23] fix(budget): cast migration seed to BIGINT and clarify schema library scope --- packages/auth/spend.py | 10 +++++++--- packages/db/migrate.py | 2 +- packages/db/models/api_key.py | 10 ++++------ 3 files changed, 12 insertions(+), 10 deletions(-) diff --git a/packages/auth/spend.py b/packages/auth/spend.py index 730b255..4d6e0a0 100644 --- a/packages/auth/spend.py +++ b/packages/auth/spend.py @@ -1,11 +1,15 @@ -"""Per-key lifetime spend tracking that enforces ``ApiKey.budget_limit_cents``. +"""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 is that column scaled by MICROCENTS_PER_CENT. +`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 enforcement is a +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:: diff --git a/packages/db/migrate.py b/packages/db/migrate.py index d934163..3751c6b 100644 --- a/packages/db/migrate.py +++ b/packages/db/migrate.py @@ -111,7 +111,7 @@ async def ensure_budget_columns(engine) -> None: await conn.execute( text( "UPDATE api_keys SET spent_microcents = (" - " SELECT COALESCE(SUM(cost_microcents), 0) FROM requests_log " + " 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" ) diff --git a/packages/db/models/api_key.py b/packages/db/models/api_key.py index b93668f..19b33a9 100644 --- a/packages/db/models/api_key.py +++ b/packages/db/models/api_key.py @@ -22,12 +22,10 @@ class ApiKey(Base, UUIDMixin, TimestampMixin, SoftDeleteMixin): # 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. `spend.charge_budget` is the only - # writer: 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. This column is the state that protocol needs — - # a caller enforces by checking `is_exhausted` before dispatch and charging - # after, so the cap is only as live as the paths that route through it. + # 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 ) From f530d9b4f65365fc7bdcb25fcf67ef65aa090478 Mon Sep 17 00:00:00 2001 From: Hasit Bhatt Date: Thu, 24 Sep 2026 20:48:54 -0700 Subject: [PATCH 05/23] feat(budget): enforce hard cap with atomic settlement and fair stream pricing --- app/routes/chat.py | 277 +++++- tests/integration/test_budget_enforcement.py | 799 ++++++++++++++++++ .../unit/test_unmeasured_stream_settlement.py | 45 + 3 files changed, 1095 insertions(+), 26 deletions(-) create mode 100644 tests/integration/test_budget_enforcement.py create mode 100644 tests/unit/test_unmeasured_stream_settlement.py diff --git a/app/routes/chat.py b/app/routes/chat.py index 5bb34c1..fa9884a 100644 --- a/app/routes/chat.py +++ b/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 @@ -58,6 +59,54 @@ # retries hold the (already [DONE]) stream open. Tests shrink this. _LOG_COMMIT_BACKOFF_S: tuple[float, ...] = (0.1, 0.4) +# Crude character→token divisor used only to price a delivery the provider +# never measured (see `_settle_unmeasured_stream`). +_CHARS_PER_TOKEN = 4 + + +def _text_chars(content) -> int: + """Character count of a message's text, across str and content-part lists.""" + if isinstance(content, str): + return len(content) + if isinstance(content, list): + return sum( + len(part["text"]) + for part in content + if isinstance(part, dict) and isinstance(part.get("text"), str) + ) + return 0 + + +def _settle_unmeasured_stream( + agg_usage: dict, agg_output_chars: int, body, *, caller_bailed: bool = False +) -> dict: + """Token counts for a stream that delivered content but reported no usage. + + Reached when the stream ends without a usage frame after content had + already been forwarded — a provider failing mid-generation or a client + hanging up. The prompt was billed upstream and the delivered text is real, + so settling such a stream at zero would let a flaky provider be streamed + for free against a capped key. Character counts divided by 4 under-count + code and CJK on purpose — an estimate must not over-bill for a failure the + caller cannot steer. + + `caller_bailed` also prices a delivery that carried nothing. An empty + delivery from a provider failure is evidence of an empty cost, so there the + two are the same thing; a client that hung up *caused* the empty delivery, + after the prompt had already gone upstream. Without the flag, disconnecting + at the first byte would settle every request at zero and the cap would stop + moving for exactly the client choosing not to wait. + """ + if agg_usage: + return agg_usage + if not agg_output_chars and not caller_bailed: + return agg_usage + prompt_chars = sum(_text_chars(m.content) for m in body.messages) + return { + "prompt_tokens": max(1, prompt_chars // _CHARS_PER_TOKEN), + "completion_tokens": max(1, agg_output_chars // _CHARS_PER_TOKEN), + } + def _chunk_to_dict(chunk) -> dict: """Normalize a litellm chunk (Pydantic model or dict) into a plain dict. @@ -368,6 +417,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, *, commit: bool = True) -> None: + """Record `actual_microcents` of spend against the cap, if any. + + No-op when the key has no budget cap. When `commit` is False the UPDATE is + executed but not committed, so the caller commits it in the same + transaction as the request-log write — making the row and the charge one + atomic unit. Idempotency across retries comes from the row's trace_id + (a persisted trace_id proves the charge also landed), not from a + process-local flag. + """ + cap = getattr(kc, "_budget_cap", None) + if cap is None: + return + await charge_budget(session, str(kc.key_id), cap, actual_microcents, commit=commit) + 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" @@ -482,6 +546,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) @@ -552,6 +634,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)) @@ -572,16 +655,14 @@ async def execute_chat( # mid-flight cascade is impossible — we have to surface the error and let # the client decide what to do. if body.stream: - # Auto-inject `stream_options.include_usage=True` if the client - # didn't set it. Without this, OpenAI/LiteLLM streaming responses - # omit the `usage` field entirely — chunks have no token counts, - # so our log row gets input=0, output=0 and the cost calculation - # rounds to zero. Almost no client knows to opt-in to this flag, - # which would silently zero out streaming spend in the dashboard. - # Honor an explicit `include_usage=False` from the client if they - # really want to disable it (e.g. wire-format compatibility tests). + # Auto-inject `stream_options.include_usage=True` if the client didn't set + # it, so streaming responses carry token counts and we bill correctly. + # A budgeted key MUST receive usage so its spend is measured: a + # client-supplied `include_usage=False` would otherwise record zero cost + # and let a capped key stream for free, so force it on for any budgeted key + # regardless of the client's preference. existing_so = completion_kwargs.get("stream_options") or {} - if "include_usage" not in existing_so: + if getattr(kc, "_budget_cap", None) is not None or "include_usage" not in existing_so: completion_kwargs["stream_options"] = {**existing_so, "include_usage": True} async def _log_pre_stream_failure(status: int, err_type: str | None) -> None: @@ -605,6 +686,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: # Roll back so the request-scoped session is not left in a @@ -648,12 +730,26 @@ async def sse() -> AsyncGenerator[str, None]: agg_provider = "unknown" agg_fallback = False agg_latency = 0 + # Characters of assistant text handed to the client — the only + # measure of what a stream delivered when the provider never + # reported usage (see `_settle_unmeasured_stream`). + agg_output_chars = 0 # The first chunk's `model` field tells us what LiteLLM actually # served (could be a cascaded fallback, not the resolved primary). agg_model: str | None = None status_code = 200 error_type: str | None = None log_written = False + # The usage frame is the billing signal: True once one has been + # observed. A stream that ends without it — client hung up before + # the usage frame, suppressed it, or the provider omitted it — has + # an unknown cost and is settled fail-closed against the key's + # remaining allowance so the cap cannot be bypassed. With a usage + # frame delivered the cost is known even if the client then + # disconnects, and charging more would over-bill a quantity the + # row already accounts for (and break charged == row.cost, the + # invariant the trace-id idempotence relies on). + usage_seen = False async def _finalize() -> None: """Write the request log row exactly once. @@ -751,27 +847,47 @@ 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. + + The usage frame is the billing signal. When no usage frame + was ever observed — the stream ended early, the client + suppressed the frame, or the provider omitted it — the real + cost is unknown and the full remaining allowance is charged + (fail-closed) so no client-side choice can bypass the cap. + Once a usage frame was delivered the cost is known — even + if the stream then died — and the recorded cost is charged. + """ + actual = row_values.get("cost_microcents") or 0 + if not usage_seen: + 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 and charge the budget in ONE commit. + + The INSERT and the budget charge share a single transaction. If it + commits, both are durable; if it fails, both roll back and the + retry re-runs both. Because the charge lands in the same commit as + the row, a persisted trace_id proves the charge also landed — so a + retry returns without re-charging. The charge is therefore applied + exactly once per request: never doubled (on a commit-ack-loss + retry) and never dropped. """ 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)): return db.add(log) try: + await _settle_budget(db, _settlement_amount(), commit=False) await db.commit() except Exception: try: @@ -782,9 +898,10 @@ 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)): return s.add(log) + await _settle_budget(s, _settlement_amount(), commit=False) await s.commit() finally: try: @@ -864,8 +981,13 @@ async def _commit_row(*, retry: bool) -> None: agg_fallback = True if "usage" in d and d["usage"]: agg_usage = d["usage"] + usage_seen = True if d.get("model"): agg_model = d["model"] + for choice in d.get("choices") or []: + if isinstance(choice, dict): + delta = choice.get("delta") or {} + agg_output_chars += _text_chars(delta.get("content")) yield f"data: {json.dumps(d, separators=(',', ':'))}\n\n" yield "data: [DONE]\n\n" except (asyncio.CancelledError, GeneratorExit): @@ -899,6 +1021,20 @@ async def _commit_row(*, retry: bool) -> None: # awaits inside the shielded scope ignore outer cancellation # and run to completion. The scope exits normally and we # re-raise the original CancelledError below. + # + # Settle the delivery before unwinding. The usage frame is the + # last chunk, so a hangup almost always means it never arrived + # and the cost is unknown — but unknown is not licence to bill + # the whole remaining allowance for a few sentences the user + # chose to stop reading. Price what reached the client, the way + # the provider-error branch does. A bail is the one unmeasured + # ending where even an empty delivery costs the prompt: the + # client chose the emptiness, and `acompletion` had already + # sent the prompt upstream. + agg_usage = _settle_unmeasured_stream( + agg_usage, agg_output_chars, body, caller_bailed=True, + ) + usage_seen = True aclose = getattr(stream_obj, "aclose", None) with anyio.CancelScope(shield=True): if aclose is not None: @@ -968,6 +1104,20 @@ 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" + # Mark the settlement known: the error response was delivered + # in full (terminal [DONE] sent), so charge the recorded cost + # rather than the full remaining allowance. Otherwise every + # transient mid-stream provider failure (rate limit, 5xx, + # network drop) would charge (and exhaust) the key's entire + # remaining budget. What the provider never measured is priced + # from what actually reached the client, so an unmeasured + # partial stream still costs something proportional to the + # delivery instead of nothing — a capped key cannot stream for + # free behind a flaky provider. Client disconnects never reach + # this branch (GeneratorExit is not an Exception); the cancel + # branch prices them from the same delivery estimate. + agg_usage = _settle_unmeasured_stream(agg_usage, agg_output_chars, body) + usage_seen = True finally: # Same shielding reason as the cancel branch: ensure the # log write actually completes before we unwind, even if @@ -1008,6 +1158,12 @@ async def _commit_row(*, retry: bool) -> None: response: dict = {} actual_resolved: str | None = None try: + # A budgeted key must receive usage so its spend is measured. Force + # include_usage on for budgeted keys even if the client omitted it. + if getattr(kc, "_budget_cap", None) is not None: + existing_so = completion_kwargs.get("stream_options") or {} + if existing_so.get("include_usage") is not True: + completion_kwargs["stream_options"] = {**existing_so, "include_usage": True} response = await client.acompletion( **completion_kwargs, fallbacks=fallbacks_arg, @@ -1058,11 +1214,80 @@ async def _commit_row(*, retry: bool) -> None: # _build_log_row would otherwise default to via requested_model). actual_resolved=actual_resolved or resolved_model, ) - db.add(log) - try: - await db.commit() - except Exception as commit_err: - logger.warning("request_log_commit_failed", error=str(commit_err)) + # Persist the log row and the budget charge atomically (same transaction), + # retrying transient commit failures so a budgeted key is never under- + # charged when the DB is stressed — mirroring the streaming path. A + # persisted trace_id proves both landed, so a retry skips rather than + # double-charging. + from sqlalchemy import select + + settle_amount = log.cost_microcents + # Fail-closed mirror of the streaming path's cost-unknown rule: a + # budgeted key whose successful response carries no usage (provider + # ignored the forced include_usage) has an unknown cost — charge the + # full remaining allowance so a delivered completion can never cost + # nothing. Gated on having actually received a completion dict: a + # request that failed before the upstream answered (response == {}, + # e.g. the re-raised HTTPException above, whose status_code never + # left 200) charges its recorded ~0 cost instead — mirroring the + # cache-hit and pre-stream-failure paths. + if ( + getattr(kc, "_budget_cap", None) is not None + and status_code < 400 + and isinstance(response, dict) + and response + and not response.get("usage") + ): + settle_amount = max( + log.cost_microcents or 0, + kc._budget_cap - (getattr(kc, "_budget_spent", 0) or 0), + ) + + # Values are snapshotted once (latency is measured in _build_log_row, + # before any commit attempt, so retry backoff never inflates it) and + # each attempt inserts a fresh ORM object carrying the same id/trace_id + # — mirroring the streaming path, so a retry works regardless of what + # the rollback left the old object as. + log_values = { + c.key: getattr(log, c.key) + for c in RequestLog.__table__.columns + if getattr(log, c.key) is not None + } + max_attempts = len(_LOG_COMMIT_BACKOFF_S) + 1 + for attempt in range(1, max_attempts + 1): + try: + if attempt > 1 and ( + await db.scalar( + select(RequestLog.id).where(RequestLog.trace_id == log.trace_id) + ) + ) is not None: + break # already durable (log + charge committed) + db.add(RequestLog(**log_values)) + await _settle_budget(db, settle_amount, commit=False) + await db.commit() + break + except Exception as commit_err: + try: + await db.rollback() + except Exception: + pass + if attempt == max_attempts: + logger.warning( + "request_log_commit_failed", error=str(commit_err), attempts=attempt, + ) + break + logger.info( + "request_log_commit_retry", error=str(commit_err), attempt=attempt, + ) + try: + await asyncio.sleep(_LOG_COMMIT_BACKOFF_S[attempt - 1]) + except BaseException: + # Cancelled during the backoff: nothing is in flight and the + # row is given up on — say so, then propagate like the arm above. + logger.warning( + "request_log_commit_failed", error=str(commit_err), attempts=attempt, + ) + raise hosted_fallback = _meta_hosted_fallback(response) if isinstance(response, dict) and "_orca_meta" in response: diff --git a/tests/integration/test_budget_enforcement.py b/tests/integration/test_budget_enforcement.py new file mode 100644 index 0000000..6f65877 --- /dev/null +++ b/tests/integration/test_budget_enforcement.py @@ -0,0 +1,799 @@ +"""Budget enforcement on /v1/chat/completions. + +`budget_limit_cents` was loaded into KeyContext but never enforced anywhere — +a leaked key meant unbounded spend. These tests pin the new behavior: an +exhausted key gets 429 before any routing / cache / upstream work and +unbudgeted keys are unaffected. Provisioning of budgeted/allowlisted keys +is covered in the keys-authz PR. +""" + +from __future__ import annotations + +import asyncio +import time +from unittest.mock import AsyncMock + +import pytest + + +@pytest.fixture +async def budget_env(tmp_sqlite_url, monkeypatch): + """Full app + seeded root key, with the router client mocked out. + + Yields (make_client, fake_client, session_factory, root_key). + """ + monkeypatch.setenv("DATABASE_URL", tmp_sqlite_url) + monkeypatch.setenv("OPENAI_API_KEY", "sk-test-openai") + + from app import config as cfg + cfg.get_settings.cache_clear() + + from packages.db.engine import build_engine + from packages.db.models.base import Base + + engine = build_engine(tmp_sqlite_url) + async with engine.begin() as conn: + await conn.run_sync(Base.metadata.create_all) + + from sqlalchemy.ext.asyncio import async_sessionmaker + + from packages.db import session as session_mod + factory = async_sessionmaker(engine, expire_on_commit=False) + session_mod._session_factory = factory + + from app.seed import seed_initial_state + async with factory() as s: + seed = await seed_initial_state(s) + + fake_client = AsyncMock() + fake_client.acompletion = AsyncMock( + return_value={ + "id": "chatcmpl-budget-test", + "model": "gpt-4o-mini", + "object": "chat.completion", + "created": int(time.time()), + "choices": [{ + "index": 0, + "message": {"role": "assistant", "content": "Hello!"}, + "finish_reason": "stop", + }], + "usage": {"prompt_tokens": 5, "completion_tokens": 2, "total_tokens": 7}, + "_orca_meta": { + "provider": "openai", + "litellm_model": "openai/gpt-4o-mini", + "latency_ms": 42, + }, + } + ) + + from app import router_cache + router_cache.invalidate_router() + + async def _fake_get_router(_session): + return fake_client + + monkeypatch.setattr(router_cache, "get_router", _fake_get_router) + + from httpx import ASGITransport, AsyncClient + + from app.main import create_app + app = create_app() + + async def make_client(api_key: str): + return AsyncClient( + transport=ASGITransport(app=app), + base_url="http://t", + headers={"Authorization": f"Bearer {api_key}"}, + ) + + yield make_client, fake_client, factory, seed.api_key + + await engine.dispose() + session_mod._session_factory = None + + +async def _make_budgeted_key( + factory, *, budget_limit_cents: int | None +) -> tuple[str, str]: + """Insert a budgeted child key; return (plaintext_key, key_id).""" + from packages.auth.hashing import generate_api_key + from packages.db.models.api_key import ApiKey + + full_key, key_hash, key_prefix = generate_api_key() + async with factory() as s: + row = ApiKey( + workspace_id="default", + name="budgeted", + key_hash=key_hash, + key_prefix=key_prefix, + budget_limit_cents=budget_limit_cents, + ) + s.add(row) + await s.commit() + await s.refresh(row) + return full_key, row.id + + +async def _add_billable_spend(factory, key_id: str, microcents: int) -> None: + from packages.db.models.api_key import ApiKey + from packages.db.models.request_log import RequestLog + + async with factory() as s: + s.add(RequestLog( + workspace_id="default", + api_key_id=key_id, + trace_id="budget-test-trace", + model_requested="gpt-4o-mini", + model_resolved="gpt-4o-mini", + provider="openai", + routing_strategy="balanced", + input_tokens=5, + output_tokens=2, + cost_microcents=microcents, + latency_ms=10, + status_code=200, + )) + # The budget counter lives on the key, not the request-log rows, so + # pre-load it directly to simulate prior spend. + await s.execute( + ApiKey.__table__.update() + .where(ApiKey.id == key_id) + .values(spent_microcents=ApiKey.spent_microcents + microcents) + ) + await s.commit() + + +async def test_exhausted_budget_returns_429_without_upstream_call(budget_env): + make_client, fake, factory, _root = budget_env + key, key_id = await _make_budgeted_key(factory, budget_limit_cents=1) + # Pre-load spend past the 1-cent cap (10_000 microcents). + await _add_billable_spend(factory, key_id, microcents=20_000) + + async with await make_client(key) as c: + r = await c.post( + "/v1/chat/completions", + json={"model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hi"}]}, + ) + + assert r.status_code == 429, r.text + assert r.json()["error"]["type"] == "rate_limit_error" + fake.acompletion.assert_not_awaited() + + +async def test_blocked_request_writes_no_log_row(budget_env): + make_client, _fake, factory, _root = budget_env + key, key_id = await _make_budgeted_key(factory, budget_limit_cents=1) + await _add_billable_spend(factory, key_id, microcents=99_999) + + async with await make_client(key) as c: + await c.post( + "/v1/chat/completions", + json={"model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hi"}]}, + ) + + from sqlalchemy import func, select + + from packages.db.models.request_log import RequestLog + + async with factory() as s: + count = ( + await s.execute( + select(func.count()).select_from(RequestLog).where( + RequestLog.api_key_id == key_id + ) + ) + ).scalar_one() + assert count == 1 # only the pre-loaded history row + + +async def test_under_budget_key_serves_normally(budget_env): + make_client, fake, factory, _root = budget_env + key, _key_id = await _make_budgeted_key(factory, budget_limit_cents=100) + + async with await make_client(key) as c: + r = await c.post( + "/v1/chat/completions", + json={"model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hi"}]}, + ) + + assert r.status_code == 200, r.text + fake.acompletion.assert_awaited_once() + + +async def test_unbudgeted_root_key_unaffected(budget_env): + make_client, fake, _factory, root = budget_env + + async with await make_client(root) as c: + r = await c.post( + "/v1/chat/completions", + json={"model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hi"}]}, + ) + + assert r.status_code == 200, r.text + fake.acompletion.assert_awaited_once() + + +async def _budgeted_stream(budget_env, *, chunks, budget_limit_cents=10): + """Drive a streaming request for a budgeted key and return its final spend.""" + make_client, fake, factory, _root = budget_env + key, key_id = await _make_budgeted_key(factory, budget_limit_cents=budget_limit_cents) + + async def _stream(): + for ch in chunks: + yield ch + + fake.acompletion = AsyncMock(return_value=_stream()) + + async with await make_client(key) as c: + async with c.stream( + "POST", + "/v1/chat/completions", + json={ + "model": "gpt-4o-mini", + "stream": True, + "messages": [{"role": "user", "content": "hi"}], + "stream_options": {"include_usage": False}, + }, + ) as r: + async for _ in r.aiter_lines(): + pass + + from sqlalchemy import select + + from packages.db.models.api_key import ApiKey + + async with factory() as s: + return ( + await s.execute(select(ApiKey.spent_microcents).where(ApiKey.id == key_id)) + ).scalar_one(), fake.acompletion.call_args + + +async def test_budgeted_stream_without_usage_charges_remaining(budget_env): + # A completed stream that never delivers a usage frame (client forced + # include_usage=False, provider ignored it) must NOT bill zero — that would + # let a capped key stream for free. Fail-closed: charge the full remaining cap. + spent, call_args = await _budgeted_stream( + budget_env, + chunks=[ + {"choices": [{"delta": {"content": "hi"}, "finish_reason": None}]}, + {"choices": [{"delta": {}, "finish_reason": "stop"}]}, + ], + ) + # Even though the client demanded include_usage=False, the budgeted key forces it. + assert call_args.kwargs["stream_options"]["include_usage"] is True + # No usage frame observed -> full cap charged. + assert spent == 100_000 + + +async def test_budgeted_stream_with_usage_frame_charges_actual(budget_env): + # A usage frame was observed, so only the real (tiny) cost is charged, not the + # full remaining allowance. + spent, _call_args = await _budgeted_stream( + budget_env, + budget_limit_cents=100, + chunks=[ + {"choices": [{"delta": {"content": "hi"}, "finish_reason": None}]}, + { + "usage": {"prompt_tokens": 5000, "completion_tokens": 2000, "total_tokens": 7000}, + "choices": [{"delta": {}, "finish_reason": "stop"}], + }, + ], + ) + assert 0 <= spent < 100_000 + + +async def test_budgeted_blocking_forces_include_usage(budget_env): + # Non-streaming budgeted request also forces include_usage on, even when the + # client omits it. + make_client, fake, factory, _root = budget_env + key, _key_id = await _make_budgeted_key(factory, budget_limit_cents=100) + + async with await make_client(key) as c: + r = await c.post( + "/v1/chat/completions", + json={ + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hi"}], + "stream_options": {"include_usage": False}, + }, + ) + + assert r.status_code == 200, r.text + assert fake.acompletion.call_args.kwargs["stream_options"]["include_usage"] is True + + +async def _get_spent(factory, key_id: str) -> int: + from sqlalchemy import select + + from packages.db.models.api_key import ApiKey + + async with factory() as s: + return ( + await s.execute(select(ApiKey.spent_microcents).where(ApiKey.id == key_id)) + ).scalar_one() + + +async def test_budgeted_stream_midstream_error_charges_actual_only(budget_env): + # A mid-stream provider error is delivered as a complete error response + # (SSE error frame + terminal [DONE]); the log row records its ~0 cost, so + # settlement is KNOWN and must charge the actual cost only. Before the fix, + # usage_seen stayed False in that branch and every transient provider + # failure permanently exhausted the key (charged cap - spent). + make_client, fake, factory, _root = budget_env + key, key_id = await _make_budgeted_key(factory, budget_limit_cents=10) + + def _failing_stream(): + async def _gen(): + yield {"choices": [{"delta": {"content": "partial"}, "finish_reason": None}]} + raise RuntimeError("upstream exploded") + return _gen() + + fake.acompletion = AsyncMock(return_value=_failing_stream()) + + async with await make_client(key) as c: + async with c.stream( + "POST", + "/v1/chat/completions", + json={ + "model": "gpt-4o-mini", + "stream": True, + "messages": [{"role": "user", "content": "hi"}], + }, + ) as r: + text = "\n".join([line async for line in r.aiter_lines()]) + + # The error response was delivered in full. + assert "Upstream provider error" in text + assert "[DONE]" in text + # Only the recorded (~0) cost is charged — not the 100_000-microcent cap. + assert await _get_spent(factory, key_id) == 0 + + # The key is NOT exhausted: a follow-up streaming request is still served. + fake.acompletion = AsyncMock(return_value=_ok_stream()) + async with await make_client(key) as c: + async with c.stream( + "POST", + "/v1/chat/completions", + json={ + "model": "gpt-4o-mini", + "stream": True, + "messages": [{"role": "user", "content": "hi again"}], + }, + ) as r2: + assert r2.status_code == 200 + async for _ in r2.aiter_lines(): + pass + + +def _ok_stream(): + async def _gen(): + yield {"choices": [{"delta": {"content": "hi"}, "finish_reason": None}]} + yield { + "usage": {"prompt_tokens": 5, "completion_tokens": 2, "total_tokens": 7}, + "choices": [{"delta": {}, "finish_reason": "stop"}], + } + return _gen() + + +async def test_budgeted_stream_error_after_unmeasured_content_charges_estimate(budget_env): + """Partial content the provider never measured must still cost something. + + The upstream dies mid-generation after a long delivery and no usage frame + ever arrives, so nothing measures it. Charging zero — what the recorded cost + says — would let a capped key stream unbounded tokens free of charge behind + a flaky provider; charging the whole remaining allowance would exhaust the + key for a failure it cannot steer. Settlement is therefore priced from the + delivered characters, and the same number lands on the row and on the key. + """ + make_client, fake, factory, _root = budget_env + key, key_id = await _make_budgeted_key(factory, budget_limit_cents=100) + + delivered = "the quick brown fox " * 2_000 # ~40k chars ≈ 10k tokens + + def _failing_stream(): + async def _gen(): + yield {"choices": [{"delta": {"content": delivered}, "finish_reason": None}]} + raise RuntimeError("upstream exploded mid-generation") + return _gen() + + fake.acompletion = AsyncMock(return_value=_failing_stream()) + + async with await make_client(key) as c: + async with c.stream( + "POST", + "/v1/chat/completions", + json={ + "model": "gpt-4o-mini", + "stream": True, + "messages": [{"role": "user", "content": "say it again " * 400}], + }, + ) as r: + text = "\n".join([line async for line in r.aiter_lines()]) + + assert "[DONE]" in text + + from sqlalchemy import select + + from packages.db.models.request_log import RequestLog + + async with factory() as s: + row = ( + await s.execute( + select(RequestLog).where(RequestLog.api_key_id == key_id) + ) + ).scalars().one() + assert row.output_tokens > 0 # the delivery is recorded, not erased + spent = await _get_spent(factory, key_id) + assert spent == row.cost_microcents # charged == accounted + assert 0 < spent < 1_000_000 # not free, and not the 100-cent cap + + +async def test_budgeted_blocking_without_usage_charges_remaining(budget_env): + # A budgeted key whose provider ignores the forced include_usage and returns + # a usage-less completion has an unknown cost. Mirroring the streaming rule, + # the blocking path must fail closed and charge the full remaining allowance + # — otherwise the delivered completion costs nothing and the cap is bypassed. + make_client, fake, factory, _root = budget_env + key, key_id = await _make_budgeted_key(factory, budget_limit_cents=10) + + fake.acompletion = AsyncMock(return_value={ + "id": "chatcmpl-no-usage", + "model": "gpt-4o-mini", + "object": "chat.completion", + "created": int(time.time()), + "choices": [{ + "index": 0, + "message": {"role": "assistant", "content": "Hello!"}, + "finish_reason": "stop", + }], + }) + + async with await make_client(key) as c: + r = await c.post( + "/v1/chat/completions", + json={"model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hi"}]}, + ) + + assert r.status_code == 200, r.text + assert await _get_spent(factory, key_id) == 100_000 # 10 cents, fail-closed + + +async def test_budgeted_blocking_httpexception_charges_recorded_cost(budget_env): + # A budgeted blocking request whose upstream call raised HTTPException never + # received a completion (response == {}, status_code never left 200). The + # fail-closed remaining-charge rule applies only to *delivered* usage-less + # completions — charging the cap here would repeat the mid-stream-error bug + # class on the blocking path. The key must be charged its recorded ~0 cost. + from fastapi import HTTPException + + make_client, fake, factory, _root = budget_env + key, key_id = await _make_budgeted_key(factory, budget_limit_cents=10) + + fake.acompletion = AsyncMock( + side_effect=HTTPException(status_code=429, detail="upstream rate limit") + ) + + async with await make_client(key) as c: + r = await c.post( + "/v1/chat/completions", + json={"model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hi"}]}, + ) + + assert r.status_code == 429, r.text + assert await _get_spent(factory, key_id) == 0 + + +async def test_budgeted_blocking_with_usage_charges_actual(budget_env): + # Control for the test above: a blocking response WITH usage must charge only + # the recorded cost (never the remaining allowance) — no over-charging. + make_client, fake, factory, _root = budget_env + key, key_id = await _make_budgeted_key(factory, budget_limit_cents=100) + + async with await make_client(key) as c: + r = await c.post( + "/v1/chat/completions", + json={"model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hi"}]}, + ) + + assert r.status_code == 200, r.text # fixture response carries usage + + from sqlalchemy import select + + from packages.db.models.request_log import RequestLog + + async with factory() as s: + row_cost = ( + await s.execute( + select(RequestLog.cost_microcents).where( + RequestLog.api_key_id == key_id + ) + ) + ).scalar_one() + assert await _get_spent(factory, key_id) == row_cost + + +async def test_budgeted_stream_disconnect_after_usage_charges_actual(budget_env): + """Measured spend must not be re-opened by a later hang-up. + + The usage frame arrives, then the client disconnects. Cost is therefore + KNOWN (the row records it), so settlement charges that cost. Keying the + fail-closed rule on stream completion instead charged the whole remaining + allowance for a request whose tokens were already accounted for. + """ + make_client, fake, factory, _root = budget_env + key, key_id = await _make_budgeted_key(factory, budget_limit_cents=100) + + class _CancelAfterUsage: + def __init__(self): + self._n = 0 + + def __aiter__(self): + return self + + async def __anext__(self): + self._n += 1 + if self._n == 1: + return {"choices": [{"delta": {"content": "hi"}, + "finish_reason": None}]} + if self._n == 2: + return { + "usage": { + "prompt_tokens": 100_000, + "completion_tokens": 50_000, + "total_tokens": 150_000, + }, + "choices": [{"delta": {}, "finish_reason": "stop"}], + } + # Mirrors Starlette cancelling the response task on http.disconnect. + raise asyncio.CancelledError() + + async def aclose(self): + pass + + fake.acompletion = AsyncMock(return_value=_CancelAfterUsage()) + + async with await make_client(key) as c: + try: + async with c.stream( + "POST", + "/v1/chat/completions", + json={ + "model": "gpt-4o-mini", + "stream": True, + "messages": [{"role": "user", "content": "hi"}], + }, + ) as r: + async for _ in r.aiter_lines(): + pass + except Exception: + pass # the injected cancel may surface to the test transport + + from sqlalchemy import select + + from packages.db.models.request_log import RequestLog + + async with factory() as s: + row = ( + await s.execute( + select(RequestLog).where(RequestLog.api_key_id == key_id) + ) + ).scalars().one() + assert row.status_code == 499 + assert row.error_type == "client_disconnect" + assert row.cost_microcents > 0 + spent = await _get_spent(factory, key_id) + assert spent == row.cost_microcents + # The disconnect is not a cost-unknown bail: it must not exhaust the key. + assert spent < 1_000_000 # cap is 100 cents = 1_000_000 microcents + + +async def test_budgeted_stream_disconnect_before_usage_charges_delivery(budget_env): + """A hangup before the usage frame must cost the delivery, not the cap. + + The usage frame is the last chunk, so a user pressing stop mid-answer means + it never arrives and nothing measured the cost. Treating that as + cost-unknown-and-therefore-max charged the entire remaining allowance for + reading a few sentences, which is a normal action and bricks the key + permanently. The delivery is priced from the characters that reached the + client, exactly as the provider-error branch does it. + """ + make_client, fake, factory, _root = budget_env + key, key_id = await _make_budgeted_key(factory, budget_limit_cents=100) + + delivered = "the quick brown fox " * 200 # 4000 chars, no usage frame ever + + class _CancelBeforeUsage: + def __init__(self): + self._n = 0 + + def __aiter__(self): + return self + + async def __anext__(self): + self._n += 1 + if self._n == 1: + return {"choices": [{"delta": {"content": delivered}, + "finish_reason": None}]} + # The client hangs up mid-answer: no usage frame was ever produced. + raise asyncio.CancelledError() + + async def aclose(self): + pass + + fake.acompletion = AsyncMock(return_value=_CancelBeforeUsage()) + + async with await make_client(key) as c: + try: + async with c.stream( + "POST", + "/v1/chat/completions", + json={ + "model": "gpt-4o-mini", + "stream": True, + "messages": [{"role": "user", "content": "hi"}], + }, + ) as r: + async for _ in r.aiter_lines(): + pass + except Exception: + pass # the injected cancel may surface to the test transport + + from sqlalchemy import select + + from packages.db.models.request_log import RequestLog + + async with factory() as s: + row = ( + await s.execute( + select(RequestLog).where(RequestLog.api_key_id == key_id) + ) + ).scalars().one() + assert row.status_code == 499 + assert row.error_type == "client_disconnect" + + spent = await _get_spent(factory, key_id) + # Something real was delivered, so it is not free... + assert spent > 0 + # ...and it is the delivery, not the 1_000_000-microcent remainder. + assert spent == row.cost_microcents + assert spent < 1_000_000 + + # The key still works: a follow-up streaming request is served. + fake.acompletion = AsyncMock(return_value=_ok_stream()) + async with await make_client(key) as c: + async with c.stream( + "POST", + "/v1/chat/completions", + json={ + "model": "gpt-4o-mini", + "stream": True, + "messages": [{"role": "user", "content": "hi again"}], + }, + ) as r2: + assert r2.status_code == 200 + async for _ in r2.aiter_lines(): + pass + + +async def test_budgeted_stream_hangup_before_first_chunk_still_costs_the_prompt(budget_env): + """Bailing at the first byte must not be a way to read a capped key for free. + + Pricing an unmeasured stream from what was delivered is right for a failure + the caller cannot steer, but a disconnect is the caller's own choice, and it + happens after `acompletion` has already sent the prompt upstream. Settling an + empty delivery at zero made "hang up immediately" the cheapest request of + all: the counter never moved, so the lifetime cap stopped applying entirely + and every retry was dispatched upstream and billed there. + """ + make_client, fake, factory, _root = budget_env + key, key_id = await _make_budgeted_key(factory, budget_limit_cents=100) + + class _CancelBeforeAnyChunk: + def __aiter__(self): + return self + + async def __anext__(self): + # The client is gone before the first delta is forwarded. + raise asyncio.CancelledError() + + async def aclose(self): + pass + + fake.acompletion = AsyncMock(return_value=_CancelBeforeAnyChunk()) + prompt = "a very long prompt " * 400 # 7600 chars, priced upstream + + async with await make_client(key) as c: + try: + async with c.stream( + "POST", + "/v1/chat/completions", + json={ + "model": "gpt-4o-mini", + "stream": True, + "messages": [{"role": "user", "content": prompt}], + }, + ) as r: + async for _ in r.aiter_lines(): + pass + except Exception: + pass + + spent = await _get_spent(factory, key_id) + # Nothing was delivered, yet the request was not free. + assert spent > 0 + # And it is the prompt, not the whole remaining allowance. + assert spent < 1_000_000 + + # Repeating the trick keeps charging, so the cap still closes. + for _ in range(3): + fake.acompletion = AsyncMock(return_value=_CancelBeforeAnyChunk()) + async with await make_client(key) as c: + try: + async with c.stream( + "POST", + "/v1/chat/completions", + json={ + "model": "gpt-4o-mini", + "stream": True, + "messages": [{"role": "user", "content": prompt}], + }, + ) as r: + async for _ in r.aiter_lines(): + pass + except Exception: + pass + assert await _get_spent(factory, key_id) > spent + spent = await _get_spent(factory, key_id) + + +async def test_budgeted_blocking_commit_failure_persists_row_and_charge(budget_env): + """A transient write failure must drop neither the row nor the charge. + + The blocking path retries with a FRESH ORM object (the failed attempt's + INSERT was rolled back) and skips the retry when the trace_id is already + durable, so the atomic row+charge unit lands exactly once. + """ + from sqlalchemy import event, select + + from packages.db.models.request_log import RequestLog + + make_client, fake, factory, _root = budget_env + key, key_id = await _make_budgeted_key(factory, budget_limit_cents=100) + + # Fail the log INSERT once, at the cursor: by the time commit runs, the + # row is already flushed (the budget UPDATE autoflushes it), so this is the + # only seam that reproduces a real "database is locked" mid-write. + sync_engine = factory.kw["bind"].sync_engine + failures = {"n": 0} + + def _fail_first_log_insert(conn, cursor, statement, parameters, context, executemany): + if "INSERT INTO requests_log" in statement and failures["n"] == 0: + failures["n"] += 1 + raise RuntimeError("database is locked") + + event.listen(sync_engine, "before_cursor_execute", _fail_first_log_insert) + try: + async with await make_client(key) as c: + r = await c.post( + "/v1/chat/completions", + json={"model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hi"}]}, + ) + finally: + event.remove(sync_engine, "before_cursor_execute", _fail_first_log_insert) + + assert r.status_code == 200, r.text + assert failures["n"] == 1 # the retry is what saved the write + async with factory() as s: + rows = ( + await s.execute(select(RequestLog).where(RequestLog.api_key_id == key_id)) + ).scalars().all() + assert len(rows) == 1 # never doubled + assert await _get_spent(factory, key_id) == rows[0].cost_microcents diff --git a/tests/unit/test_unmeasured_stream_settlement.py b/tests/unit/test_unmeasured_stream_settlement.py new file mode 100644 index 0000000..e295bc3 --- /dev/null +++ b/tests/unit/test_unmeasured_stream_settlement.py @@ -0,0 +1,45 @@ +"""`_settle_unmeasured_stream` only prices what nothing else measured.""" + +from __future__ import annotations + +from types import SimpleNamespace + +from app.routes.chat import _settle_unmeasured_stream + + +def _body(prompt: str): + return SimpleNamespace(messages=[SimpleNamespace(content=prompt)]) + + +def test_measured_usage_is_never_replaced_by_an_estimate(): + usage = {"prompt_tokens": 11, "completion_tokens": 22} + got = _settle_unmeasured_stream(usage, 90_000, _body("x" * 400)) + assert got == usage + + +def test_nothing_delivered_stays_unbilled(): + # A failure before the first content chunk delivered no tokens; inventing a + # prompt charge for it would repeat the over-charge this estimate replaces. + assert _settle_unmeasured_stream({}, 0, _body("x" * 400)) == {} + + +def test_empty_client_bail_still_costs_the_prompt(): + # A disconnect is the caller's choice after the prompt went upstream, so an + # empty delivery is priced from the prompt rather than settled at zero. + assert _settle_unmeasured_stream({}, 0, _body("x" * 400), caller_bailed=True) == { + "prompt_tokens": 100, + "completion_tokens": 1, + } + + +def test_estimate_prices_prompt_and_delivery_at_char_quarter(): + got = _settle_unmeasured_stream({}, 4_000, _body("y" * 400)) + assert got == {"prompt_tokens": 100, "completion_tokens": 1000} + + +def test_content_part_lists_count_their_text(): + body = SimpleNamespace(messages=[SimpleNamespace( + content=[{"type": "text", "text": "z" * 800}, {"type": "image_url"}] + )]) + got = _settle_unmeasured_stream({}, 400, body) + assert got == {"prompt_tokens": 200, "completion_tokens": 100} From a3a34ce8c37177ee0b397289463ed0f4ea4032c0 Mon Sep 17 00:00:00 2001 From: Hasit Bhatt Date: Thu, 24 Sep 2026 22:28:53 -0700 Subject: [PATCH 06/23] fix(budget): stop sending stream_options on blocking requests, and bill what the log says MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A capped key forced `stream_options.include_usage` onto the blocking request too. That parameter only decides whether the last frame of a stream reports usage — a non-streaming completion always carries it — and LiteLLM forwards it without looking at `stream`, so OpenAI rejected every budgeted blocking request outright. The cap made the endpoint unusable rather than enforced. The fail-closed settlement also moved the counter without moving the row it was charging for, so a key could be exhausted by an amount no query over its request history reproduced. Both paths now record the charged amount on the row. --- app/routes/chat.py | 39 ++++++++++-------- tests/integration/test_budget_enforcement.py | 42 +++++++++++++++++--- 2 files changed, 60 insertions(+), 21 deletions(-) diff --git a/app/routes/chat.py b/app/routes/chat.py index fa9884a..464ca3c 100644 --- a/app/routes/chat.py +++ b/app/routes/chat.py @@ -857,6 +857,11 @@ def _settlement_amount() -> int: (fail-closed) so no client-side choice can bypass the cap. Once a usage frame was delivered the cost is known — even if the stream then died — and the recorded cost is charged. + + The raised amount is written back into the row, not just + into the counter: a charge only the counter saw would leave + a key exhausted by an amount nothing in its own request + history accounts for. """ actual = row_values.get("cost_microcents") or 0 if not usage_seen: @@ -865,6 +870,7 @@ def _settlement_amount() -> int: (getattr(kc, "_budget_cap", 0) or 0) - (getattr(kc, "_budget_spent", 0) or 0), ) + row_values["cost_microcents"] = actual return actual async def _commit_row(*, retry: bool) -> None: @@ -878,6 +884,10 @@ async def _commit_row(*, retry: bool) -> None: exactly once per request: never doubled (on a commit-ack-loss retry) and never dropped. """ + # Resolved before the row is built: a fail-closed charge + # raises `row_values["cost_microcents"]`, and the object + # inserted must carry the amount the counter will move by. + settlement = _settlement_amount() log = RequestLog(**row_values) if session_mod._session_factory is None: # Test-only fallback (the app always installs a @@ -887,7 +897,7 @@ async def _commit_row(*, retry: bool) -> None: return db.add(log) try: - await _settle_budget(db, _settlement_amount(), commit=False) + await _settle_budget(db, settlement, commit=False) await db.commit() except Exception: try: @@ -901,7 +911,7 @@ async def _commit_row(*, retry: bool) -> None: if retry and (await _already_persisted(s)): return s.add(log) - await _settle_budget(s, _settlement_amount(), commit=False) + await _settle_budget(s, settlement, commit=False) await s.commit() finally: try: @@ -1158,12 +1168,6 @@ async def _commit_row(*, retry: bool) -> None: response: dict = {} actual_resolved: str | None = None try: - # A budgeted key must receive usage so its spend is measured. Force - # include_usage on for budgeted keys even if the client omitted it. - if getattr(kc, "_budget_cap", None) is not None: - existing_so = completion_kwargs.get("stream_options") or {} - if existing_so.get("include_usage") is not True: - completion_kwargs["stream_options"] = {**existing_so, "include_usage": True} response = await client.acompletion( **completion_kwargs, fallbacks=fallbacks_arg, @@ -1223,14 +1227,16 @@ async def _commit_row(*, retry: bool) -> None: settle_amount = log.cost_microcents # Fail-closed mirror of the streaming path's cost-unknown rule: a - # budgeted key whose successful response carries no usage (provider - # ignored the forced include_usage) has an unknown cost — charge the - # full remaining allowance so a delivered completion can never cost - # nothing. Gated on having actually received a completion dict: a - # request that failed before the upstream answered (response == {}, - # e.g. the re-raised HTTPException above, whose status_code never - # left 200) charges its recorded ~0 cost instead — mirroring the - # cache-hit and pre-stream-failure paths. + # budgeted key whose successful response carries no usage (a provider + # that answered without the field at all) has an unknown cost — charge + # the full remaining allowance so a delivered completion can never cost + # nothing, and record that amount on the row: a charge only the counter + # saw would leave a key exhausted by an amount nothing in its own + # request history accounts for. Gated on having actually received a + # completion dict: a request that failed before the upstream answered + # (response == {}, e.g. the re-raised HTTPException above, whose + # status_code never left 200) charges its recorded ~0 cost instead — + # mirroring the cache-hit and pre-stream-failure paths. if ( getattr(kc, "_budget_cap", None) is not None and status_code < 400 @@ -1242,6 +1248,7 @@ async def _commit_row(*, retry: bool) -> None: log.cost_microcents or 0, kc._budget_cap - (getattr(kc, "_budget_spent", 0) or 0), ) + log.cost_microcents = settle_amount # Values are snapshotted once (latency is measured in _build_log_row, # before any commit attempt, so retry backoff never inflates it) and diff --git a/tests/integration/test_budget_enforcement.py b/tests/integration/test_budget_enforcement.py index 6f65877..ebd8487 100644 --- a/tests/integration/test_budget_enforcement.py +++ b/tests/integration/test_budget_enforcement.py @@ -286,9 +286,16 @@ async def test_budgeted_stream_with_usage_frame_charges_actual(budget_env): assert 0 <= spent < 100_000 -async def test_budgeted_blocking_forces_include_usage(budget_env): - # Non-streaming budgeted request also forces include_usage on, even when the - # client omits it. +async def test_budgeted_blocking_request_sends_no_stream_options(budget_env): + """A cap must not put a streaming-only parameter on a blocking request. + + `include_usage` only decides whether the last frame of a *stream* reports + usage — a non-streaming completion always carries it. LiteLLM forwards the + parameter without looking at `stream`, and OpenAI rejects it on a request + where stream is false, so forcing it here turned every request for a + budgeted key into an upstream 400: the cap made the endpoint unusable + instead of enforced. + """ make_client, fake, factory, _root = budget_env key, _key_id = await _make_budgeted_key(factory, budget_limit_cents=100) @@ -298,12 +305,37 @@ async def test_budgeted_blocking_forces_include_usage(budget_env): json={ "model": "gpt-4o-mini", "messages": [{"role": "user", "content": "hi"}], - "stream_options": {"include_usage": False}, }, ) assert r.status_code == 200, r.text - assert fake.acompletion.call_args.kwargs["stream_options"]["include_usage"] is True + assert "stream_options" not in fake.acompletion.call_args.kwargs + + +async def test_fail_closed_charge_is_recorded_on_the_row_it_charges(budget_env): + """The counter and the request history are one quantity and may not diverge. + + `spent_microcents` is seeded from, and reconciled against, the sum of + `cost_microcents`, so a fail-closed charge that only the counter saw leaves + a key exhausted by an amount no query over its requests can reproduce. + """ + spent, _call_args = await _budgeted_stream( + budget_env, + chunks=[ + {"choices": [{"delta": {"content": "hi"}, "finish_reason": None}]}, + {"choices": [{"delta": {}, "finish_reason": "stop"}]}, + ], + ) + assert spent == 100_000 # the full remaining allowance + + _make_client, _fake, factory, _root = budget_env + from sqlalchemy import select + + from packages.db.models.request_log import RequestLog + + async with factory() as s: + rows = (await s.execute(select(RequestLog.cost_microcents))).scalars().all() + assert rows == [spent] async def _get_spent(factory, key_id: str) -> int: From 994ac27f8f3327cde933cde927bed2ed163ffdbc Mon Sep 17 00:00:00 2001 From: Hasit Bhatt Date: Fri, 25 Sep 2026 02:12:30 -0700 Subject: [PATCH 07/23] fix(budget): read the spend counter once per pre-dispatch check Every request for a budgeted key issued two identical `SELECT spent_microcents WHERE id = ?` round trips: `is_exhausted` loaded the counter to decide the 429 and the snapshot of the remaining allowance loaded it again straight after, on the same session with nothing written in between. The fix is a seam rather than an inlined comparison because the number and the boolean cannot both come from one read otherwise, and the route's single call is where a pre-check that has more to do than read a column will hang it. --- app/routes/chat.py | 7 ++++--- packages/auth/spend.py | 17 +++++++++++++++-- 2 files changed, 19 insertions(+), 5 deletions(-) diff --git a/app/routes/chat.py b/app/routes/chat.py index 464ca3c..e0732d2 100644 --- a/app/routes/chat.py +++ b/app/routes/chat.py @@ -31,7 +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.spend import MICROCENTS_PER_CENT, budget_precheck, charge_budget from packages.auth.types import KeyContext from packages.db.models.request_log import RequestLog from packages.litellm_adapter.catalog import CATALOG, CATALOG_BY_ID @@ -556,13 +556,14 @@ async def _settle_budget(session, actual_microcents: int, *, commit: bool = True # 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): + spent = await budget_precheck(db, str(kc.key_id), cap) + if spent >= 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)) + kc._budget_spent = spent started_perf = time.perf_counter() completion_kwargs = body.model_dump(exclude_none=True) diff --git a/packages/auth/spend.py b/packages/auth/spend.py index 4d6e0a0..1179a0e 100644 --- a/packages/auth/spend.py +++ b/packages/auth/spend.py @@ -47,6 +47,20 @@ async def read_spent(db: AsyncSession, api_key_id: str) -> int: return int(spent or 0) +async def budget_precheck(db: AsyncSession, api_key_id: str, cap_microcents: int) -> int: + """The key's spend as of this pre-check, in microcents. + + One read behind both halves of the caller's decision: whether to reject the + request, and what allowance is left for a cost that is not known yet. + ``cap_microcents`` is ``ApiKey.budget_limit_cents`` scaled by + ``MICROCENTS_PER_CENT``, not the column itself. It is unused here and + becomes load-bearing the moment this function has a parked obligation to + fold before it answers; keeping it in the signature is what lets the caller + hold on to a single pre-check call instead of reading the counter twice. + """ + return await read_spent(db, api_key_id) + + 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? @@ -54,8 +68,7 @@ async def is_exhausted(db: AsyncSession, api_key_id: str, cap_microcents: int) - ``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 + return await budget_precheck(db, api_key_id, cap_microcents) >= cap_microcents async def charge_budget( From f6059fb212cc6416b6ae4b9157ddc4c465826fae Mon Sep 17 00:00:00 2001 From: Hasit Bhatt Date: Fri, 25 Sep 2026 03:54:35 -0700 Subject: [PATCH 08/23] fix(budget): settle unmeasured stream before yielding error frame --- app/routes/chat.py | 22 ++++++++-------------- 1 file changed, 8 insertions(+), 14 deletions(-) diff --git a/app/routes/chat.py b/app/routes/chat.py index e0732d2..6c1acb1 100644 --- a/app/routes/chat.py +++ b/app/routes/chat.py @@ -1104,6 +1104,14 @@ async def _commit_row(*, retry: bool) -> None: "chat_completion_stream_error", error=str(exc), error_type=error_type, ) + # Mark the settlement known: compute the delivery estimate + # and mark usage_seen BEFORE yielding the error frame and sentinel. + # If the client disconnects during yield, GeneratorExit unwinds + # directly through finally without executing lines below the yield; + # settling first ensures _finalize never charges the key's full + # remaining allowance for a client disconnect during error delivery. + agg_usage = _settle_unmeasured_stream(agg_usage, agg_output_chars, body) + usage_seen = True err_body = { "error": { "message": f"Upstream provider error: {exc}", @@ -1115,20 +1123,6 @@ 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" - # Mark the settlement known: the error response was delivered - # in full (terminal [DONE] sent), so charge the recorded cost - # rather than the full remaining allowance. Otherwise every - # transient mid-stream provider failure (rate limit, 5xx, - # network drop) would charge (and exhaust) the key's entire - # remaining budget. What the provider never measured is priced - # from what actually reached the client, so an unmeasured - # partial stream still costs something proportional to the - # delivery instead of nothing — a capped key cannot stream for - # free behind a flaky provider. Client disconnects never reach - # this branch (GeneratorExit is not an Exception); the cancel - # branch prices them from the same delivery estimate. - agg_usage = _settle_unmeasured_stream(agg_usage, agg_output_chars, body) - usage_seen = True finally: # Same shielding reason as the cancel branch: ensure the # log write actually completes before we unwind, even if From 9f9d06df2e1304e4ca2d836ebe580ffdc15b1f94 Mon Sep 17 00:00:00 2001 From: Hasit Bhatt Date: Thu, 24 Sep 2026 21:01:03 -0700 Subject: [PATCH 09/23] feat(budget): park lost settlements durably and fold them at pre-check --- app/routes/chat.py | 140 ++++++- packages/auth/spend.py | 223 ++++++++++- packages/db/migrate.py | 20 +- packages/db/models/__init__.py | 2 + packages/db/models/budget_park.py | 28 ++ tests/integration/test_budget_enforcement.py | 399 +++++++++++++++++++ tests/unit/test_budget_migration.py | 39 ++ tests/unit/test_budget_spend.py | 147 +++++++ tests/unit/test_give_up_settlement.py | 74 ++++ 9 files changed, 1041 insertions(+), 31 deletions(-) create mode 100644 packages/db/models/budget_park.py create mode 100644 tests/unit/test_give_up_settlement.py diff --git a/app/routes/chat.py b/app/routes/chat.py index 6c1acb1..e398d4e 100644 --- a/app/routes/chat.py +++ b/app/routes/chat.py @@ -11,7 +11,7 @@ import json import time import uuid -from collections.abc import AsyncGenerator, AsyncIterable, Callable +from collections.abc import AsyncGenerator, AsyncIterable, Awaitable, Callable import anyio import structlog @@ -31,7 +31,12 @@ 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, budget_precheck, charge_budget +from packages.auth.spend import ( + MICROCENTS_PER_CENT, + budget_precheck, + charge_budget, + record_unsettled_spend, +) from packages.auth.types import KeyContext from packages.db.models.request_log import RequestLog from packages.litellm_adapter.catalog import CATALOG, CATALOG_BY_ID @@ -108,6 +113,82 @@ def _settle_unmeasured_stream( } +async def _give_up_settlement( + kc: KeyContext, + trace_id: str, + amount: int, + attempts: int, + error: BaseException, + persisted: Callable[[], Awaitable[bool]], +) -> None: + """Last resort for a settlement that is not durable and will not be retried. + + The log row dies with the charge (one transaction), so nothing anywhere + remembers this cost. Park it — one idempotent row keyed by this + settlement's `trace_id` — against the key's cap, rather than leaving the + cap open for whoever reads the warning (see `packages.auth.spend`). + + `persisted()` is asked first, because the last attempt has no retry left to + run the trace_id check: a commit that applied but whose ack was lost (or a + cancellation that landed after it) leaves the charge already in + `spent_microcents`, and parking on top of it would bill one delivery twice. + """ + logger.warning("request_log_commit_failed", error=str(error), attempts=attempts) + if getattr(kc, "_budget_cap", None) is None: + return + try: + durable = await persisted() + except asyncio.CancelledError: + # Torn down mid-probe, so the outcome is still unknown: park before + # propagating rather than letting this cost vanish with the coroutine. + await record_unsettled_spend( + trace_id=trace_id, api_key_id=str(kc.key_id), microcents=amount + ) + raise + if not durable: + await record_unsettled_spend( + trace_id=trace_id, api_key_id=str(kc.key_id), microcents=amount + ) + + +async def _trace_is_persisted(session: AsyncSession, trace_id: str) -> bool: + from sqlalchemy import select + + return ( + await session.scalar( + select(RequestLog.id).where(RequestLog.trace_id == trace_id) + ) + ) is not None + + +async def _settlement_is_durable(db: AsyncSession, trace_id: str) -> bool: + """Whether the request-log row for this trace_id is committed. + + The row and the budget charge are one transaction, so the row is proof the + spend is already counted. It reads on its own session because the failed + attempt's session is closed or rolled back by the time the give-up runs, and + answers False when it cannot read at all: during a real outage parking is the + only thing keeping the cap honest, so an unreadable DB must look + not-durable rather than durable. + """ + from packages.db import session as session_mod + + try: + if session_mod._session_factory is None: + # Test-only fallback (the app always installs a factory). + return await _trace_is_persisted(db, trace_id) + s = session_mod._session_factory() + try: + return await _trace_is_persisted(s, trace_id) + finally: + try: + await s.close() + except Exception as close_err: + logger.debug("request_log_session_close_failed", error=str(close_err)) + except Exception: + return False + + def _chunk_to_dict(chunk) -> dict: """Normalize a litellm chunk (Pydantic model or dict) into a plain dict. @@ -848,6 +929,9 @@ async def _already_persisted(s) -> bool: select(RequestLog.id).where(RequestLog.trace_id == row_values["trace_id"]) )) is not None + async def _durable() -> bool: + return await _settlement_is_durable(db, row_values["trace_id"]) + def _settlement_amount() -> int: """Budget charge for this request, in microcents. @@ -952,15 +1036,15 @@ async def _commit_row(*, retry: bool) -> None: # Cancelled during the backoff: nothing is # in flight, the row is given up on — say # so, then propagate like the arm below. - logger.warning( - "request_log_commit_failed", - error=str(commit_err), attempts=attempt, + await _give_up_settlement( + kc, row_values["trace_id"], _settlement_amount(), + attempt, commit_err, _durable, ) raise continue - logger.warning( - "request_log_commit_failed", - error=str(commit_err), attempts=attempt, + await _give_up_settlement( + kc, row_values["trace_id"], _settlement_amount(), + attempt, commit_err, _durable, ) except BaseException: # CancelledError aimed at us, not at the commit — @@ -971,9 +1055,9 @@ async def _commit_row(*, retry: bool) -> None: try: await commit_task except Exception as commit_err: - logger.warning( - "request_log_commit_failed", - error=str(commit_err), attempts=attempt, + await _give_up_settlement( + kc, row_values["trace_id"], _settlement_amount(), + attempt, commit_err, _durable, ) except BaseException: pass @@ -1256,6 +1340,12 @@ async def _commit_row(*, retry: bool) -> None: if getattr(log, c.key) is not None } max_attempts = len(_LOG_COMMIT_BACKOFF_S) + 1 + + async def _durable() -> bool: + # The snapshot, not `log`: the give-up can run after a rollback has + # expired the session's state. + return await _settlement_is_durable(db, log_values["trace_id"]) + for attempt in range(1, max_attempts + 1): try: if attempt > 1 and ( @@ -1274,8 +1364,9 @@ async def _commit_row(*, retry: bool) -> None: except Exception: pass if attempt == max_attempts: - logger.warning( - "request_log_commit_failed", error=str(commit_err), attempts=attempt, + await _give_up_settlement( + kc, log_values["trace_id"], settle_amount, + attempt, commit_err, _durable, ) break logger.info( @@ -1286,10 +1377,29 @@ async def _commit_row(*, retry: bool) -> None: except BaseException: # Cancelled during the backoff: nothing is in flight and the # row is given up on — say so, then propagate like the arm above. - logger.warning( - "request_log_commit_failed", error=str(commit_err), attempts=attempt, + await _give_up_settlement( + kc, log_values["trace_id"], settle_amount, + attempt, commit_err, _durable, ) raise + except BaseException as cancel_err: + # Cancelled while the write was in flight. Unlike the streaming + # path there is no detached task left to land it: the transaction + # dies with this coroutine. Release it first — it holds the + # pending write's locks, which the durability read below needs + # past — then park unless it did commit, or this request's + # delivered cost would vanish with the coroutine. + try: + await db.rollback() + except BaseException: + # A second cancellation here must not skip the park: the + # give-up below propagates the original either way. + pass + await _give_up_settlement( + kc, log_values["trace_id"], settle_amount, + attempt, cancel_err, _durable, + ) + raise hosted_fallback = _meta_hosted_fallback(response) if isinstance(response, dict) and "_orca_meta" in response: diff --git a/packages/auth/spend.py b/packages/auth/spend.py index 1179a0e..da9c782 100644 --- a/packages/auth/spend.py +++ b/packages/auth/spend.py @@ -31,13 +31,199 @@ from __future__ import annotations -from sqlalchemy import select, update +import asyncio + +from sqlalchemy import delete, func, select, update +from sqlalchemy.exc import IntegrityError from sqlalchemy.ext.asyncio import AsyncSession from packages.db.models.api_key import ApiKey +from packages.db.models.budget_park import BudgetPark MICROCENTS_PER_CENT = 10_000 +# 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 the counter never moves. The obligation is parked here — one row per +# settlement, keyed by its `trace_id` — and keeps counting against the cap until +# a budget pre-check folds it into `spent_microcents`. +# +# The park is a database table, not process memory, because the deployment +# stops its machine whenever it goes idle: an in-memory obligation is lost on +# the next cold start, which reopens the cap for exactly the key the failure +# was about to protect. `_unsettled` below only holds an amount while the +# database itself is unreachable — the same outage that caused the park — and a +# later pre-check re-files it once a write goes through again. +_unsettled: dict[tuple[str, str], int] = {} + + +class _FoldConflict(Exception): + """A concurrent worker folded the same park first; the loser retries later.""" + + +async def _insert_park(*, trace_id: str, api_key_id: str, microcents: int) -> bool: + """Persist one parked obligation. Returns True when it is durable. + + Idempotent on `trace_id`: a commit that applied but whose ack was lost + retries into the same primary key instead of recording the obligation a + second time. Returns False when the database is unavailable (or this is a + unit test with no session factory), leaving the caller to hold the amount + in memory. A cancellation propagates with the memory copy still held. + """ + from packages.db import session as session_mod + + factory = session_mod._session_factory + if factory is None: + return False + try: + async with factory() as s: + s.add( + BudgetPark( + trace_id=trace_id, api_key_id=api_key_id, microcents=microcents, + ) + ) + await s.commit() + return True + except IntegrityError: + # The obligation is already parked — either by our own retried write + # after an ack loss, or by a concurrent give-up for the same trace. + # Either way there is exactly one durable copy, so report success. + return True + except asyncio.CancelledError: + raise + except Exception: + return False + + +async def record_unsettled_spend( + *, trace_id: str, api_key_id: str, microcents: int +) -> None: + """Keep a settlement that gave up counting against the key's cap. + + The caller is unwinding a failure, so this never loses the amount: a park + that cannot be written durably is held in memory under its `trace_id` for + the next pre-check to re-file, and a cancellation holds it before + propagating rather than taking the obligation with it. + """ + if microcents <= 0 or not trace_id or not api_key_id: + return + key = (str(api_key_id), str(trace_id)) + try: + if await _insert_park( + trace_id=key[1], api_key_id=key[0], microcents=microcents + ): + return + except asyncio.CancelledError: + # `_insert_park` only raises the cancellation unwinding the caller — + # and the amount still has to be held before it propagates. + _unsettled[key] = microcents + raise + _unsettled[key] = microcents + + +async def pending_parked_spend(api_key_id: str) -> int: + """The outstanding park for a key: durable rows plus whatever is memory-only.""" + from packages.db import session as session_mod + + key = str(api_key_id) + total = sum( + amount for (held_key, _trace), amount in _unsettled.items() if held_key == key + ) + factory = session_mod._session_factory + if factory is None: + return total + try: + async with factory() as s: + stored = ( + await s.execute( + select(func.sum(BudgetPark.microcents)).where( + BudgetPark.api_key_id == key + ) + ) + ).scalar() + total += int(stored or 0) + except Exception: + pass + return total + + +async def settle_parked_spend(api_key_id: str, cap_microcents: int) -> int: + """Fold a key's parked obligations into its recorded spend. Returns what moved. + + The park exists because a charge could not be recorded; leaving it parked + forever would mean a key at its cap is rejected by an amount that never + settles and never clears, so every pre-dispatch check tries to move it. + Only whole obligations that fit under the remaining allowance move, so the + counter still cannot overshoot the cap; a park larger than the remainder + stays parked in full — that over-claim is the fail-closed policy, and it + stays visible as a row rather than being written off. + + The charge and the row deletions share one transaction with + compare-and-swap guards: two workers folding the same park cannot + double-bill it, because the loser's UPDATE or DELETE matches nothing and + its next request folds what the winner left. + """ + from packages.db import session as session_mod + + key = str(api_key_id) + factory = session_mod._session_factory + if factory is None: + return 0 + for (held_key, trace_id), amount in list(_unsettled.items()): + if held_key != key: + continue + # A cancellation here propagates with the entry still held; a later + # pre-check re-files it, and the `trace_id` key keeps the retry from + # duplicating it. + if await _insert_park( + trace_id=trace_id, api_key_id=key, microcents=amount + ): + _unsettled.pop((held_key, trace_id), None) + try: + async with factory() as s: + async with s.begin(): + spent = ( + await s.execute( + select(ApiKey.spent_microcents).where(ApiKey.id == key) + ) + ).scalar_one_or_none() + if spent is None: + return 0 + spent = int(spent) + rows = ( + await s.execute( + select(BudgetPark.trace_id, BudgetPark.microcents) + .where(BudgetPark.api_key_id == key) + .order_by(BudgetPark.created_at) + ) + ).all() + move = 0 + settling: list[str] = [] + for trace_id, microcents in rows: + microcents = int(microcents) + if spent + move + microcents <= cap_microcents: + move += microcents + settling.append(trace_id) + else: + break + if move <= 0: + return 0 + charged = await s.execute( + update(ApiKey) + .where(ApiKey.id == key, ApiKey.spent_microcents == spent) + .values(spent_microcents=spent + move) + ) + if charged.rowcount != 1: + raise _FoldConflict + cleared = await s.execute( + delete(BudgetPark).where(BudgetPark.trace_id.in_(settling)) + ) + if cleared.rowcount != len(settling): + raise _FoldConflict + return move + except _FoldConflict: + return 0 + async def read_spent(db: AsyncSession, api_key_id: str) -> int: """Return the key's currently-recorded lifetime spend in microcents.""" @@ -51,23 +237,34 @@ async def budget_precheck(db: AsyncSession, api_key_id: str, cap_microcents: int """The key's spend as of this pre-check, in microcents. One read behind both halves of the caller's decision: whether to reject the - request, and what allowance is left for a cost that is not known yet. - ``cap_microcents`` is ``ApiKey.budget_limit_cents`` scaled by - ``MICROCENTS_PER_CENT``, not the column itself. It is unused here and - becomes load-bearing the moment this function has a parked obligation to - fold before it answers; keeping it in the signature is what lets the caller - hold on to a single pre-check call instead of reading the counter twice. - """ - return await read_spent(db, api_key_id) - - -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? + request, and what allowance is left for a cost that is not known yet. It + folds a parked obligation into the counter before answering, so a write + outage is neither a window of free requests nor a park that can never clear, + and the result adds whatever is still pending rather than reading only the + counter: when the fold could not commit, the obligation still has to block + dispatch. ``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. """ + key = str(api_key_id) + spent = await read_spent(db, key) + pending = await pending_parked_spend(key) + if pending: + try: + await settle_parked_spend(key, cap_microcents) + except Exception: + # The fold runs on sessions of its own, so there is nothing to + # roll back here — and the re-read below still counts the park. + pass + spent = await read_spent(db, key) + pending = await pending_parked_spend(key) + return spent + pending + + +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 budget_precheck(db, api_key_id, cap_microcents) >= cap_microcents diff --git a/packages/db/migrate.py b/packages/db/migrate.py index 3751c6b..50f290c 100644 --- a/packages/db/migrate.py +++ b/packages/db/migrate.py @@ -17,6 +17,8 @@ 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.""" @@ -50,18 +52,30 @@ async def ensure_budget_columns(engine) -> None: 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) 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 seed 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 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` keeps a racing boot from + # failing when the winner creates it first. + await conn.run_sync(BudgetPark.__table__.create, 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 diff --git a/packages/db/models/__init__.py b/packages/db/models/__init__.py index ca728ff..f8fc00a 100644 --- a/packages/db/models/__init__.py +++ b/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", diff --git a/packages/db/models/budget_park.py b/packages/db/models/budget_park.py new file mode 100644 index 0000000..a2943f1 --- /dev/null +++ b/packages/db/models/budget_park.py @@ -0,0 +1,28 @@ +"""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. A row is deleted in the +same transaction that folds it into `spent_microcents`, so a crash between the +two is impossible — the obligation is either still parked or already billed, +never both and never neither. +""" + +from sqlalchemy import BigInteger, 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) diff --git a/tests/integration/test_budget_enforcement.py b/tests/integration/test_budget_enforcement.py index ebd8487..ba5fc5e 100644 --- a/tests/integration/test_budget_enforcement.py +++ b/tests/integration/test_budget_enforcement.py @@ -14,6 +14,7 @@ from unittest.mock import AsyncMock import pytest +from sqlalchemy.ext.asyncio import AsyncSession @pytest.fixture @@ -829,3 +830,401 @@ def _fail_first_log_insert(conn, cursor, statement, parameters, context, execute ).scalars().all() assert len(rows) == 1 # never doubled assert await _get_spent(factory, key_id) == rows[0].cost_microcents + + +# ── Durable recovery: the park outlives the process that lost it ────── + +class _AckLossSession(AsyncSession): + """Commit for real, then report failure as if the ack never came back. + + Armed for settlement commits only — a commit carrying a `RequestLog` row. + The row is durable while its caller still sees an exception: the case the + retry loops' trace_id check exists for. + """ + + drop_ack = False + drops = 0 + + async def commit(self): + settles = self.info.pop("settles_request_log", False) + await super().commit() + if settles and _AckLossSession.drop_ack: + _AckLossSession.drop_ack = False + _AckLossSession.drops += 1 + raise ConnectionError("connection dropped mid-ack") + + async def rollback(self): + self.info.pop("settles_request_log", None) + await super().rollback() + + +def _note_settlement_flush(session, flush_context, instances): + # `commit()` runs after autoflush has already emptied `session.new`, so + # the settlement commit is recognised here, while the row is still new. + from packages.db.models.request_log import RequestLog + + if any(isinstance(o, RequestLog) for o in session.new): + session.info["settles_request_log"] = True + + +from sqlalchemy import event as _sa_event +from sqlalchemy.orm import Session as _SyncSession + +_sa_event.listen(_SyncSession, "before_flush", _note_settlement_flush) + + +class _WriteBlackout: + """Fail every settlement write at the cursor — a sustained write outage. + + Reads still work, which is what makes this the dangerous shape: the key + keeps being served on its pre-check while nothing it spends can be + recorded — not the charge, and not even the park. + """ + + _MATCHES = ( + "INSERT INTO requests_log", + "UPDATE api_keys SET spent_microcents", + "INSERT INTO budget_parks", + ) + + def __init__(self, factory): + self.active = False + self._engine = factory.kw["bind"].sync_engine + from sqlalchemy import event + + event.listen(self._engine, "before_cursor_execute", self._handle) + + def _handle(self, conn, cursor, statement, parameters, context, executemany): + if self.active and any(m in statement for m in self._MATCHES): + raise RuntimeError("database is locked") + + def close(self): + from sqlalchemy import event + + event.remove(self._engine, "before_cursor_execute", self._handle) + + +def _completion(text: str, *, usage: dict | None = None) -> dict: + response = { + "id": "chatcmpl-blackout", + "model": "gpt-4o-mini", + "object": "chat.completion", + "created": int(time.time()), + "choices": [{ + "index": 0, + "message": {"role": "assistant", "content": text}, + "finish_reason": "stop", + }], + "_orca_meta": {"provider": "openai", "litellm_model": "openai/gpt-4o-mini", "latency_ms": 42}, + } + if usage: + response["usage"] = usage + return response + + +async def _lossy_factory(factory): + """A session factory on the same engine whose settlement commits drop acks.""" + from sqlalchemy.ext.asyncio import async_sessionmaker + + return async_sessionmaker( + factory.kw["bind"], expire_on_commit=False, class_=_AckLossSession, + ) + + +async def test_parked_spend_survives_a_real_restart(tmp_sqlite_url): + """A lost settlement must outlive the process that lost it. + + A new engine on the same file, with no process memory carried over, still + folds the park exactly once: the obligation lives in the database, not in + the worker that recorded it. + """ + from sqlalchemy.ext.asyncio import async_sessionmaker + + from packages.auth import spend as spend_mod + from packages.auth.hashing import generate_api_key + from packages.auth.spend import ( + charge_budget, + is_exhausted, + pending_parked_spend, + read_spent, + record_unsettled_spend, + ) + from packages.db import session as session_mod + from packages.db.engine import build_engine + from packages.db.models.api_key import ApiKey + from packages.db.models.base import Base + + cap = 10_000 + engine1 = build_engine(tmp_sqlite_url) + try: + async with engine1.begin() as conn: + await conn.run_sync(Base.metadata.create_all) + factory1 = async_sessionmaker(engine1, expire_on_commit=False) + session_mod._session_factory = factory1 + + full_key, key_hash, key_prefix = generate_api_key() + async with factory1() as s: + row = ApiKey( + workspace_id="default", name="restart", + key_hash=key_hash, key_prefix=key_prefix, + budget_limit_cents=1, + ) + s.add(row) + await s.commit() + await s.refresh(row) + key_id = row.id + await charge_budget(s, key_id, cap, 4_000) + + await record_unsettled_spend( + trace_id="t-restart", api_key_id=key_id, microcents=3_000 + ) + assert await pending_parked_spend(key_id) == 3_000 + finally: + await engine1.dispose() + + spend_mod._unsettled.clear() # the machine stopped: no memory survives + session_mod._session_factory = None + + engine2 = build_engine(tmp_sqlite_url) + try: + factory2 = async_sessionmaker(engine2, expire_on_commit=False) + session_mod._session_factory = factory2 + async with factory2() as s: + assert await is_exhausted(s, key_id, cap) is False # 7_000 of 10_000 + async with factory2() as s: + assert await read_spent(s, key_id) == 7_000 + assert await pending_parked_spend(key_id) == 0 + # A second pre-check must not move the same microcents a second time. + async with factory2() as s: + assert await is_exhausted(s, key_id, cap) is False + async with factory2() as s: + assert await read_spent(s, key_id) == 7_000 + finally: + session_mod._session_factory = None + await engine2.dispose() + + +async def test_park_is_visible_to_another_worker(budget_env, monkeypatch): + """A park recorded on one worker blocks and folds on another.""" + from sqlalchemy.ext.asyncio import async_sessionmaker + + from packages.auth.spend import ( + charge_budget, + is_exhausted, + pending_parked_spend, + read_spent, + record_unsettled_spend, + ) + from packages.db import session as session_mod + + _make_client, _fake, factory, _root = budget_env + _key, key_id = await _make_budgeted_key(factory, budget_limit_cents=1) + cap = 10_000 + async with factory() as s: + await charge_budget(s, key_id, cap, 4_000) + await record_unsettled_spend( + trace_id="t-worker", api_key_id=key_id, microcents=3_000 + ) + + worker_b = async_sessionmaker( + factory.kw["bind"], expire_on_commit=False, + ) + monkeypatch.setattr(session_mod, "_session_factory", worker_b) + async with worker_b() as s: + assert await is_exhausted(s, key_id, cap) is False + assert await pending_parked_spend(key_id) == 0 + async with worker_b() as s: + assert await read_spent(s, key_id) == 7_000 + + +async def test_budgeted_blocking_write_outage_still_bills_the_delivery(budget_env): + """Spend a write outage could not record must not simply disappear. + + Three requests settle while every write — charges and park inserts alike — + fails, so no row and no durable park survives anywhere; the obligations sit + in memory. When the DB recovers, the next settlement re-files and pays for + what was already delivered as well. + """ + from sqlalchemy import select + + from packages.auth.spend import pending_parked_spend + from packages.db.models.request_log import RequestLog + + make_client, fake, factory, _root = budget_env + key, key_id = await _make_budgeted_key(factory, budget_limit_cents=100) + usage = {"prompt_tokens": 10_000, "completion_tokens": 5_000, "total_tokens": 15_000} + fake.acompletion = AsyncMock(side_effect=lambda **kw: _completion("hello", usage=usage)) + + async def _ask(i: int): + async with await make_client(key) as c: + r = await c.post( + "/v1/chat/completions", + json={"model": "gpt-4o-mini", + "messages": [{"role": "user", "content": f"hi {i}"}]}, + ) + assert r.status_code == 200, r.text + + blackout = _WriteBlackout(factory) + blackout.active = True + try: + for i in range(3): + await _ask(i) + assert await _get_spent(factory, key_id) == 0 # nothing was recordable + parked = await pending_parked_spend(key_id) + assert parked > 0 # held in memory: even the park writes failed + blackout.active = False + await _ask(3) + finally: + blackout.close() + + async with factory() as s: + rows = ( + await s.execute(select(RequestLog).where(RequestLog.api_key_id == key_id)) + ).scalars().all() + assert len(rows) == 1 # the three lost settlements left no rows, no doubles + cost = rows[0].cost_microcents + assert cost > 0 + assert await pending_parked_spend(key_id) == 0 + assert await _get_spent(factory, key_id) == 4 * cost + + +async def test_budgeted_stream_write_outage_still_blocks_the_next_request(budget_env): + """The streaming loop's give-up must clamp the next request too. + + A budgeted stream with no usage frame settles fail-closed at the whole + remaining allowance; if that commit is impossible, the amount has to keep + counting, or the outage leaves the key uncapped and the very next request + is served for free. + """ + from packages.auth.spend import pending_parked_spend + + make_client, fake, factory, _root = budget_env + key, key_id = await _make_budgeted_key(factory, budget_limit_cents=10) + + async def _no_usage(): + yield {"choices": [{"delta": {"content": "hi"}, "finish_reason": None}]} + yield {"choices": [{"delta": {}, "finish_reason": "stop"}]} + + fake.acompletion = AsyncMock(return_value=_no_usage()) + + payload = { + "model": "gpt-4o-mini", + "stream": True, + "messages": [{"role": "user", "content": "hi"}], + } + blackout = _WriteBlackout(factory) + blackout.active = True + try: + async with await make_client(key) as c: + async with c.stream("POST", "/v1/chat/completions", json=payload) as r: + assert r.status_code == 200 + async for _ in r.aiter_lines(): + pass + await asyncio.sleep(1.0) # the bounded retries run out after the response + finally: + blackout.close() + + assert await _get_spent(factory, key_id) == 0 + assert await pending_parked_spend(key_id) == 100_000 # the full remainder, held + fake.acompletion = AsyncMock(return_value=_completion( + "hello", usage={"prompt_tokens": 10_000, "completion_tokens": 5_000, "total_tokens": 15_000}, + )) + async with await make_client(key) as c: + r = await c.post( + "/v1/chat/completions", + json={"model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hi again"}]}, + ) + assert r.status_code == 429, r.text + assert r.json()["error"]["type"] == "rate_limit_error" + + +async def test_budgeted_blocking_commit_ack_loss_bills_the_delivery_once( + budget_env, monkeypatch, +): + """A commit that lands but loses its ack must bill once and park nothing.""" + from sqlalchemy import select + + from packages.auth.spend import pending_parked_spend + from packages.db import session as session_mod + from packages.db.models.request_log import RequestLog + + make_client, fake, factory, _root = budget_env + key, key_id = await _make_budgeted_key(factory, budget_limit_cents=100) + monkeypatch.setattr(session_mod, "_session_factory", await _lossy_factory(factory)) + + _AckLossSession.drop_ack = True + _AckLossSession.drops = 0 + async with await make_client(key) as c: + r = await c.post( + "/v1/chat/completions", + json={"model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hi"}]}, + ) + assert r.status_code == 200, r.text + assert _AckLossSession.drops == 1 # the scenario actually dropped the ack + + async with factory() as s: + rows = ( + await s.execute(select(RequestLog).where(RequestLog.api_key_id == key_id)) + ).scalars().all() + cost = rows[0].cost_microcents + assert cost > 0 + assert await _get_spent(factory, key_id) == cost + assert await pending_parked_spend(key_id) == 0 # durable, so nothing to park + + await _ask_blocking_once(make_client, key) + assert await _get_spent(factory, key_id) == 2 * cost # not three + + +async def _ask_blocking_once(make_client, key: str) -> None: + async with await make_client(key) as c: + r = await c.post( + "/v1/chat/completions", + json={"model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hi again"}]}, + ) + assert r.status_code == 200, r.text + + +async def test_budgeted_stream_commit_ack_loss_bills_the_delivery_once( + budget_env, monkeypatch, +): + """A streaming commit that lands but loses its ack bills once, parks nothing.""" + from sqlalchemy import select + + from packages.auth.spend import pending_parked_spend + from packages.db import session as session_mod + from packages.db.models.request_log import RequestLog + + make_client, fake, factory, _root = budget_env + key, key_id = await _make_budgeted_key(factory, budget_limit_cents=100) + fake.acompletion = AsyncMock(return_value=_ok_stream()) + monkeypatch.setattr(session_mod, "_session_factory", await _lossy_factory(factory)) + + _AckLossSession.drop_ack = True + _AckLossSession.drops = 0 + async with await make_client(key) as c: + async with c.stream( + "POST", + "/v1/chat/completions", + json={ + "model": "gpt-4o-mini", + "stream": True, + "messages": [{"role": "user", "content": "hi"}], + }, + ) as r: + assert r.status_code == 200 + async for _ in r.aiter_lines(): + pass + await asyncio.sleep(0.5) # the shielded retry runs out after the response + assert _AckLossSession.drops == 1 + + async with factory() as s: + rows = ( + await s.execute(select(RequestLog).where(RequestLog.api_key_id == key_id)) + ).scalars().all() + cost = rows[0].cost_microcents + assert cost > 0 + assert await _get_spent(factory, key_id) == cost + assert await pending_parked_spend(key_id) == 0 diff --git a/tests/unit/test_budget_migration.py b/tests/unit/test_budget_migration.py index dc6c19a..35a8d8a 100644 --- a/tests/unit/test_budget_migration.py +++ b/tests/unit/test_budget_migration.py @@ -20,6 +20,7 @@ from packages.db.migrate import ensure_budget_columns from packages.db.models.api_key import ApiKey from packages.db.models.base import Base +from packages.db.models.budget_park import BudgetPark from packages.db.models.request_log import RequestLog @@ -234,3 +235,41 @@ def __getattr__(self, name): ) == 2500 finally: await engine.dispose() + + +async def test_ensure_budget_columns_creates_budget_parks_table(tmp_sqlite_url): + """A deployment that predates durable recovery gets the park table. + + `create_all` covers fresh databases, but an upgraded SQLite volume keeps + its old schema — without this step the first give-up would have nowhere + durable to park, and the cap would silently reopen after every restart. + """ + engine = await _legacy_deploy_engine(tmp_sqlite_url) + try: + async with engine.begin() as conn: + await conn.execute(text("DROP TABLE IF EXISTS budget_parks")) + await ensure_budget_columns(engine) + + async with engine.connect() as conn: + tables = await conn.run_sync( + lambda sync: sa_inspect(sync).get_table_names() + ) + assert "budget_parks" in tables + + # The table the migration created actually holds a parked obligation. + factory = async_sessionmaker(engine, expire_on_commit=False) + async with factory() as s: + s.add(BudgetPark(trace_id="park-t1", api_key_id="k1", microcents=900)) + await s.commit() + async with factory() as s: + row = ( + await s.execute( + select(BudgetPark).where(BudgetPark.trace_id == "park-t1") + ) + ).scalar_one() + assert row.microcents == 900 + + # Idempotent across restarts: a second boot changes nothing. + await ensure_budget_columns(engine) + finally: + await engine.dispose() diff --git a/tests/unit/test_budget_spend.py b/tests/unit/test_budget_spend.py index fe62e65..0616b52 100644 --- a/tests/unit/test_budget_spend.py +++ b/tests/unit/test_budget_spend.py @@ -8,7 +8,10 @@ MICROCENTS_PER_CENT, charge_budget, is_exhausted, + pending_parked_spend, read_spent, + record_unsettled_spend, + settle_parked_spend, ) @@ -89,3 +92,147 @@ async def test_concurrent_charges_never_exceed_cap(tmp_sqlite_url): def test_microcent_conversion_constant(): assert MICROCENTS_PER_CENT == 10_000 + + +@pytest.fixture +async def parked_env(tmp_sqlite_url): + """Engine + global session factory, so the park ledger is durable here.""" + from sqlalchemy.ext.asyncio import async_sessionmaker + + from packages.db import session as session_mod + from packages.db.engine import build_engine + from packages.db.models.base import Base + + engine = build_engine(tmp_sqlite_url) + async with engine.begin() as conn: + await conn.run_sync(Base.metadata.create_all) + factory = async_sessionmaker(engine, expire_on_commit=False) + old, session_mod._session_factory = session_mod._session_factory, factory + try: + yield factory + finally: + session_mod._session_factory = old + await engine.dispose() + + +async def _parked_key(factory, *, spent: int = 0): + from packages.db.models.api_key import ApiKey + + async with factory() as s: + k = ApiKey( + workspace_id="default", name="p", key_hash="h-p", key_prefix="p-p", + ) + s.add(k) + await s.commit() + await s.refresh(k) + if spent: + await charge_budget(s, k.id, 10_000_000, spent) + return k.id + + +async def test_recorded_park_folds_exactly_once(parked_env): + """A lost settlement is billed once, by the next pre-check — never twice.""" + from packages.auth import spend as spend_mod + + key_id = await _parked_key(parked_env, spent=4_000) + cap = 10_000 + await record_unsettled_spend(trace_id="t-fold", api_key_id=key_id, microcents=3_000) + assert await pending_parked_spend(key_id) == 3_000 # durable, not just memory + + spend_mod._unsettled.clear() # the machine stopped and cold-started + + async with parked_env() as s: + assert await is_exhausted(s, key_id, cap) is False # 7_000 of 10_000 + async with parked_env() as s: + assert await read_spent(s, key_id) == 7_000 + assert await pending_parked_spend(key_id) == 0 + + # A second pre-check must not move the same microcents a second time. + async with parked_env() as s: + assert await is_exhausted(s, key_id, cap) is False + async with parked_env() as s: + assert await read_spent(s, key_id) == 7_000 + + +async def test_duplicate_park_record_is_idempotent(parked_env): + """Retrying a park write after an ack loss must not record it twice.""" + key_id = await _parked_key(parked_env) + await record_unsettled_spend(trace_id="t-dup", api_key_id=key_id, microcents=900) + # The ack never came back, so the caller retries the identical record. + await record_unsettled_spend(trace_id="t-dup", api_key_id=key_id, microcents=900) + assert await pending_parked_spend(key_id) == 900 + + +async def test_park_beyond_the_remainder_stays_parked_and_exhausted(parked_env): + """A park larger than the remainder cannot move without over-recording. + + Moving the whole 2_000 row would push the counter past the cap, and + clamping the counter without consuming the row would record spend that was + never billed. So the row stays parked — visible, not written off — and the + key is exhausted by the obligation it still owes. + """ + key_id = await _parked_key(parked_env, spent=9_000) + cap = 10_000 + await record_unsettled_spend(trace_id="t-big", api_key_id=key_id, microcents=2_000) + + async with parked_env() as s: + assert await is_exhausted(s, key_id, cap) is True + async with parked_env() as s: + assert await read_spent(s, key_id) == 9_000 # untouched, never overshot + assert await pending_parked_spend(key_id) == 2_000 # the debt stays visible + + +async def test_concurrent_folds_bill_the_park_once(parked_env): + """Two workers folding the same park move it exactly once.""" + key_id = await _parked_key(parked_env) + await record_unsettled_spend(trace_id="t-race", api_key_id=key_id, microcents=3_000) + + moved = await asyncio.gather( + settle_parked_spend(key_id, 10_000), settle_parked_spend(key_id, 10_000), + ) + assert sorted(moved) == [0, 3_000] + async with parked_env() as s: + assert await read_spent(s, key_id) == 3_000 + assert await pending_parked_spend(key_id) == 0 + + +async def test_memory_fallback_refiles_once_the_database_recovers(parked_env): + """An amount held in memory during an outage becomes exactly one park.""" + from packages.db import session as session_mod + + key_id = await _parked_key(parked_env) + old = session_mod._session_factory + session_mod._session_factory = None + try: + # No session factory: the database might as well be down. + await record_unsettled_spend(trace_id="t-mem", api_key_id=key_id, microcents=700) + assert await pending_parked_spend(key_id) == 700 + finally: + session_mod._session_factory = old + + async with parked_env() as s: + assert await is_exhausted(s, key_id, 10_000) is False + assert await pending_parked_spend(key_id) == 0 + async with parked_env() as s: + assert await read_spent(s, key_id) == 700 + + +async def test_cancelled_refile_keeps_the_memory_copy(parked_env, monkeypatch): + """Cancelling a memory-to-database refile must not drop the obligation.""" + from packages.auth import spend as spend_mod + + key_id = await _parked_key(parked_env) + spend_mod._unsettled[(key_id, "t-cancel")] = 500 + + async def _cancelled(**kwargs): + raise asyncio.CancelledError + + monkeypatch.setattr(spend_mod, "_insert_park", _cancelled) + with pytest.raises(asyncio.CancelledError): + await settle_parked_spend(key_id, 10_000) + assert spend_mod._unsettled[(key_id, "t-cancel")] == 500 + + monkeypatch.undo() + assert await settle_parked_spend(key_id, 10_000) == 500 + async with parked_env() as s: + assert await read_spent(s, key_id) == 500 diff --git a/tests/unit/test_give_up_settlement.py b/tests/unit/test_give_up_settlement.py new file mode 100644 index 0000000..cd484d4 --- /dev/null +++ b/tests/unit/test_give_up_settlement.py @@ -0,0 +1,74 @@ +"""The give-up's durability gate (app/routes/chat.py::_give_up_settlement). + +Parking is the last resort for a settlement the database never accepted: it +keeps a delivered response counting against the key's cap until a budget +pre-check can fold it back into the counter. But the final attempt has no retry +left to run the trace_id check, so it gives up on the exception alone — and an +exception can follow a commit that applied (ack lost, or a cancellation that +landed after it). Parking that cost again would bill one delivery twice. + +These run without a session factory, so the park lands in the process-memory +overflow `record_unsettled_spend` falls back to; the durable ledger itself is +covered in tests/unit/test_budget_spend.py. +""" + +from __future__ import annotations + +import asyncio + +import pytest + +from app.routes.chat import _give_up_settlement +from packages.auth.spend import pending_parked_spend + + +class _Kc: + """The two attributes the give-up reads off a KeyContext.""" + + def __init__(self, key_id: str, *, cap: int = 10_000): + self.key_id = key_id + self._budget_cap = cap + + +def _failed(error: str = "connection dropped mid-ack") -> RuntimeError: + return RuntimeError(error) + + +async def test_durable_settlement_is_not_parked_again(): + """The row is committed, so its charge already counts: park nothing.""" + kc = _Kc("durable-key") + + async def persisted() -> bool: + return True + + await _give_up_settlement(kc, "trace-durable", 900, 3, _failed(), persisted) + assert await pending_parked_spend(kc.key_id) == 0 + + +async def test_lost_settlement_parks_its_own_cost(): + """Nothing is durable: the delivery stays on the cap until it is folded.""" + kc = _Kc("lost-key") + + async def persisted() -> bool: + return False + + await _give_up_settlement( + kc, "trace-lost", 900, 3, _failed("database is locked"), persisted + ) + assert await pending_parked_spend(kc.key_id) == 900 + + +async def test_probe_cancelled_parks_before_propagating(): + """Torn down mid-read: the outcome is unknown, so park and re-raise. + + Dropping it here would be fail-open — the cancellation would take this + request's cost out with the coroutine. + """ + kc = _Kc("cancelled-probe-key") + + async def persisted() -> bool: + raise asyncio.CancelledError + + with pytest.raises(asyncio.CancelledError): + await _give_up_settlement(kc, "trace-cancelled", 900, 1, _failed(), persisted) + assert await pending_parked_spend(kc.key_id) == 900 From c4ac608272a181a26d27f2c1178e907f7b04e13d Mon Sep 17 00:00:00 2001 From: Hasit Bhatt Date: Thu, 24 Sep 2026 23:03:22 -0700 Subject: [PATCH 10/23] fix(budget): fold the park queue down to the cap instead of freezing on it A park larger than the remaining allowance moved nowhere and stopped the scan, so the counter could sit below the cap on a key refused by a row nothing would ever shrink. Apply the allowance oldest-debt-first and rewrite the oversized row to its remainder, which keeps the invariant the fold exists to hold: either the queue is empty or the counter is exactly on the cap. The over-claim stays a row, so it still blocks and folds for free if the cap is raised. created_at alone is not a total order (second resolution on SQLite, ties on Postgres), and two workers computing the same fold have to agree on which row is the partial one, so the scan breaks ties on trace_id. --- packages/auth/spend.py | 67 ++++++++++++++++++++++--------- packages/db/models/budget_park.py | 8 ++-- tests/unit/test_budget_spend.py | 50 ++++++++++++++++++----- 3 files changed, 93 insertions(+), 32 deletions(-) diff --git a/packages/auth/spend.py b/packages/auth/spend.py index da9c782..4dfae65 100644 --- a/packages/auth/spend.py +++ b/packages/auth/spend.py @@ -152,16 +152,21 @@ async def settle_parked_spend(api_key_id: str, cap_microcents: int) -> int: The park exists because a charge could not be recorded; leaving it parked forever would mean a key at its cap is rejected by an amount that never - settles and never clears, so every pre-dispatch check tries to move it. - Only whole obligations that fit under the remaining allowance move, so the - counter still cannot overshoot the cap; a park larger than the remainder - stays parked in full — that over-claim is the fail-closed policy, and it - stays visible as a row rather than being written off. - - The charge and the row deletions share one transaction with - compare-and-swap guards: two workers folding the same park cannot - double-bill it, because the loser's UPDATE or DELETE matches nothing and - its next request folds what the winner left. + settles and never clears, so every pre-check tries to move it. The + remaining allowance is applied oldest-obligation-first and a park larger + than it bills what fits and is rewritten to its remainder, rather than + staying parked whole. That keeps the invariant the fold exists to hold: + either the queue is empty, or the counter sits exactly on the cap. Without + it a key can be refused at a lifetime spend below its limit with a row that + nothing will ever shrink, which is the state this function is supposed to + drain. The remainder is still a real debt — the over-claim is the + fail-closed policy — so it stays visible and keeps `is_exhausted` blocking; + it is never written off, and it folds for free the moment the cap is raised. + + The charge and the row writes share one transaction with compare-and-swap + guards on each: two workers folding the same park cannot double-bill it, + because the loser's UPDATE or DELETE matches nothing and its next request + folds what the winner left. """ from packages.db import session as session_mod @@ -194,18 +199,29 @@ async def settle_parked_spend(api_key_id: str, cap_microcents: int) -> int: await s.execute( select(BudgetPark.trace_id, BudgetPark.microcents) .where(BudgetPark.api_key_id == key) - .order_by(BudgetPark.created_at) + # Oldest debt first. `created_at` alone is not a total + # order — it is second-resolution on SQLite and ties on + # Postgres — and two workers computing the same fold + # have to agree on which row is the partial one, so + # `trace_id` breaks the tie. + .order_by(BudgetPark.created_at, BudgetPark.trace_id) ) ).all() move = 0 + room = cap_microcents - spent settling: list[str] = [] + trim: tuple[str, int, int] | None = None for trace_id, microcents in rows: microcents = int(microcents) - if spent + move + microcents <= cap_microcents: + if microcents <= room: + room -= microcents move += microcents settling.append(trace_id) - else: - break + continue + if room > 0: + trim = (trace_id, microcents, microcents - room) + move += room + break if move <= 0: return 0 charged = await s.execute( @@ -215,11 +231,24 @@ async def settle_parked_spend(api_key_id: str, cap_microcents: int) -> int: ) if charged.rowcount != 1: raise _FoldConflict - cleared = await s.execute( - delete(BudgetPark).where(BudgetPark.trace_id.in_(settling)) - ) - if cleared.rowcount != len(settling): - raise _FoldConflict + if settling: + cleared = await s.execute( + delete(BudgetPark).where(BudgetPark.trace_id.in_(settling)) + ) + if cleared.rowcount != len(settling): + raise _FoldConflict + if trim is not None: + trace_id, whole, remainder = trim + trimmed = await s.execute( + update(BudgetPark) + .where( + BudgetPark.trace_id == trace_id, + BudgetPark.microcents == whole, + ) + .values(microcents=remainder) + ) + if trimmed.rowcount != 1: + raise _FoldConflict return move except _FoldConflict: return 0 diff --git a/packages/db/models/budget_park.py b/packages/db/models/budget_park.py index a2943f1..a233157 100644 --- a/packages/db/models/budget_park.py +++ b/packages/db/models/budget_park.py @@ -8,10 +8,10 @@ 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. A row is deleted in the -same transaction that folds it into `spent_microcents`, so a crash between the -two is impossible — the obligation is either still parked or already billed, -never both and never neither. +primary key instead of recording the obligation twice. 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`, so the obligation is never both +parked and billed, and never neither. """ from sqlalchemy import BigInteger, String diff --git a/tests/unit/test_budget_spend.py b/tests/unit/test_budget_spend.py index 0616b52..d5a23a4 100644 --- a/tests/unit/test_budget_spend.py +++ b/tests/unit/test_budget_spend.py @@ -163,13 +163,14 @@ async def test_duplicate_park_record_is_idempotent(parked_env): assert await pending_parked_spend(key_id) == 900 -async def test_park_beyond_the_remainder_stays_parked_and_exhausted(parked_env): - """A park larger than the remainder cannot move without over-recording. - - Moving the whole 2_000 row would push the counter past the cap, and - clamping the counter without consuming the row would record spend that was - never billed. So the row stays parked — visible, not written off — and the - key is exhausted by the obligation it still owes. +async def test_park_beyond_the_remainder_bills_what_fits(parked_env): + """An oversized park converges the counter on the cap instead of freezing. + + Moving the whole 2_000 row would push the counter past the cap, and leaving + it parked whole would refuse the key at a lifetime spend below its limit on + the strength of a row nothing ever shrinks. So the 1_000 the cap can absorb + bills and the row keeps the rest: the over-claim stays visible and still + blocks, and with a park outstanding the counter is now exactly on the cap. """ key_id = await _parked_key(parked_env, spent=9_000) cap = 10_000 @@ -178,8 +179,39 @@ async def test_park_beyond_the_remainder_stays_parked_and_exhausted(parked_env): async with parked_env() as s: assert await is_exhausted(s, key_id, cap) is True async with parked_env() as s: - assert await read_spent(s, key_id) == 9_000 # untouched, never overshot - assert await pending_parked_spend(key_id) == 2_000 # the debt stays visible + assert await read_spent(s, key_id) == 10_000 # the cap, never past it + assert await pending_parked_spend(key_id) == 1_000 # the debt stays visible + + # A second pre-check must not bill the microcents the first one moved, and + # the remainder must not clear on its own. + async with parked_env() as s: + assert await is_exhausted(s, key_id, cap) is True + async with parked_env() as s: + assert await read_spent(s, key_id) == 10_000 + assert await pending_parked_spend(key_id) == 1_000 + + +async def test_parked_queue_drains_oldest_first_as_the_cap_opens(parked_env): + """What the cap could not absorb waits, and folds the moment it can. + + The oversized head takes the whole allowance, so the younger park behind it + waits — not lost, just queued. Raising the cap reopens room, and the fold + keeps its promise that the counter reaches `min(cap, spent + debt)`. + """ + key_id = await _parked_key(parked_env, spent=9_000) + await record_unsettled_spend(trace_id="t-old", api_key_id=key_id, microcents=2_000) + await record_unsettled_spend(trace_id="t-new", api_key_id=key_id, microcents=500) + + assert await settle_parked_spend(key_id, 10_000) == 1_000 + assert await settle_parked_spend(key_id, 10_000) == 0 # no room left + assert await pending_parked_spend(key_id) == 1_500 + + assert await settle_parked_spend(key_id, 11_000) == 1_000 + assert await pending_parked_spend(key_id) == 500 + assert await settle_parked_spend(key_id, 12_000) == 500 + assert await pending_parked_spend(key_id) == 0 + async with parked_env() as s: + assert await read_spent(s, key_id) == 11_500 async def test_concurrent_folds_bill_the_park_once(parked_env): From b3f6ad2706279d42cd55c1d2dc9abf428f628cfe Mon Sep 17 00:00:00 2001 From: Hasit Bhatt Date: Thu, 24 Sep 2026 23:03:44 -0700 Subject: [PATCH 11/23] fix(migrations): let a boot that lost the park-table race go on booting MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `checkfirst` asks the catalog and then creates, and the ask cannot see another worker's uncommitted CREATE, so two boots racing an upgrade both issue it and one is rejected. That was the only startup DDL outside `_apply_ddl`, and on Postgres the error aborts the transaction the rest of startup runs in: the worker never comes up, and never reaches the fold this table exists to feed. Generalise `_apply_ddl` to a callable so dialect-generated DDL goes through the same savepoint, and teach `_already_applied` the shape this collision actually takes in Postgres — a unique violation on the catalog, not an "already exists". --- packages/db/migrate.py | 34 ++++++++++++++--- tests/unit/test_budget_migration.py | 58 +++++++++++++++++++++++++++++ 2 files changed, 86 insertions(+), 6 deletions(-) diff --git a/packages/db/migrate.py b/packages/db/migrate.py index 50f290c..2ab4c41 100644 --- a/packages/db/migrate.py +++ b/packages/db/migrate.py @@ -14,6 +14,9 @@ from __future__ import annotations +from collections.abc import Callable +from typing import Any + from sqlalchemy import BigInteger, inspect, text from sqlalchemy.exc import DBAPIError @@ -23,10 +26,16 @@ 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 + 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 @@ -36,10 +45,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 @@ -72,9 +87,16 @@ async def ensure_budget_columns(engine) -> None: if BudgetPark.__tablename__ not in tables: # `create_all` covers fresh databases; this covers upgrades whose - # schema predates the table. `checkfirst` keeps a racing boot from - # failing when the winner creates it first. - await conn.run_sync(BudgetPark.__table__.create, checkfirst=True) + # 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 diff --git a/tests/unit/test_budget_migration.py b/tests/unit/test_budget_migration.py index 35a8d8a..ea77005 100644 --- a/tests/unit/test_budget_migration.py +++ b/tests/unit/test_budget_migration.py @@ -14,6 +14,7 @@ from sqlalchemy import inspect as sa_inspect from sqlalchemy import select, text +from sqlalchemy.exc import OperationalError from sqlalchemy.ext.asyncio import async_sessionmaker from packages.db.engine import build_engine @@ -273,3 +274,60 @@ async def test_ensure_budget_columns_creates_budget_parks_table(tmp_sqlite_url): await ensure_budget_columns(engine) finally: await engine.dispose() + + +async def test_losing_the_park_table_race_still_boots(tmp_sqlite_url, monkeypatch): + """The one startup DDL a racing boot used to die on. + + `checkfirst` asks the catalog and then creates, so two workers inspecting an + un-migrated schema both hear "no" and one is rejected anyway — and Postgres + reports that collision as a catalog unique violation rather than an + "already exists" message. The create has to land in the same + tolerate-and-continue path as the ALTER, because a worker that dies here + never gets far enough to fold the obligation this table holds. + """ + engine = await _legacy_deploy_engine(tmp_sqlite_url) + async with engine.begin() as conn: + await conn.execute(text("DROP TABLE IF EXISTS budget_parks")) + + attempts: list[str] = [] + + def _lose_the_race(*args, **kwargs): + attempts.append("create") + raise OperationalError( + "CREATE TABLE budget_parks (...)", + {}, + Exception( + 'duplicate key value violates unique constraint ' + '"pg_class_relname_nsp_index"' + ), + ) + + monkeypatch.setattr(BudgetPark.__table__, "create", _lose_the_race) + try: + await ensure_budget_columns(engine) + + assert attempts == ["create"] # it really did take the race + async with engine.connect() as conn: + # and it went on to do the rest of startup. + assert await conn.scalar( + text("SELECT spent_microcents FROM api_keys WHERE workspace_id = 'w1'") + ) == 2500 + finally: + await engine.dispose() + + +def test_catalog_collision_reads_as_already_applied(): + """Only a collision on the object's own name counts as someone winning.""" + from packages.db.migrate import _already_applied + + def _err(msg: str) -> OperationalError: + return OperationalError("CREATE TABLE budget_parks (...)", {}, Exception(msg)) + + assert _already_applied( + _err('duplicate key value violates unique constraint "pg_class_relname_nsp_index"') + ) + assert _already_applied(_err("table budget_parks already exists")) + assert _already_applied(_err("duplicate column name: spent_microcents")) + assert not _already_applied(_err("permission denied to create relation")) + assert not _already_applied(_err('near "TABL": syntax error')) From afa9e4108556bd2ced4d6bc4d18416344c27b850 Mon Sep 17 00:00:00 2001 From: Hasit Bhatt Date: Thu, 24 Sep 2026 23:03:52 -0700 Subject: [PATCH 12/23] fix(budget): keep a cancelled rollback from skipping the give-up MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The retry handler rolls the failed attempt back before deciding, and that is an await — so it is where a cancellation aimed at the request can land. An exception raised inside a handler is not caught by the same try's other arms, so the cancellation escaped past every give-up below it: the row was already lost, the park never happened, and the key's cap reopened for the cost it had just been served. Treat it like the other arm's rollback failure and carry on to the give-up. The next attempt then still finds the session poisoned, so what lands is one park; what matters is that the settlement is accounted for rather than dropped. --- app/routes/chat.py | 11 +- tests/integration/test_budget_enforcement.py | 165 +++++++++++++++++++ 2 files changed, 174 insertions(+), 2 deletions(-) diff --git a/app/routes/chat.py b/app/routes/chat.py index e398d4e..b99bc6c 100644 --- a/app/routes/chat.py +++ b/app/routes/chat.py @@ -1361,7 +1361,11 @@ async def _durable() -> bool: except Exception as commit_err: try: await db.rollback() - except Exception: + except BaseException: + # A cancellation here must not escape the handler: the + # give-ups below are the only thing that keeps this cost + # counting, and an exception raised from inside a handler is + # not caught by this `try`'s other arms. pass if attempt == max_attempts: await _give_up_settlement( @@ -1376,7 +1380,10 @@ async def _durable() -> bool: await asyncio.sleep(_LOG_COMMIT_BACKOFF_S[attempt - 1]) except BaseException: # Cancelled during the backoff: nothing is in flight and the - # row is given up on — say so, then propagate like the arm above. + # row is given up on — say so, then propagate. An exception + # raised from a handler is not caught by this `try`'s other + # arms, so the give-up below does not run a second time for + # this settlement. await _give_up_settlement( kc, log_values["trace_id"], settle_amount, attempt, commit_err, _durable, diff --git a/tests/integration/test_budget_enforcement.py b/tests/integration/test_budget_enforcement.py index ba5fc5e..8e5140c 100644 --- a/tests/integration/test_budget_enforcement.py +++ b/tests/integration/test_budget_enforcement.py @@ -832,6 +832,171 @@ def _fail_first_log_insert(conn, cursor, statement, parameters, context, execute assert await _get_spent(factory, key_id) == rows[0].cost_microcents +async def test_cancelled_during_backoff_gives_up_exactly_once( + budget_env, monkeypatch +): + """A request torn down between retries parks its cost once. + + The write failed, the backoff is cancelled, and the transaction is abandoned + with the completion already delivered: the obligation has to be parked, and + parked exactly once. The `raise` out of the give-up is what keeps the + write-in-flight arm from running a second give-up for the same settlement — + two warnings for one abandoned row, and a second durability probe while the + process is going away. + + The row is a usage-less completion for a 10-cent key, so the obligation is + the fail-closed 100_000 microcents. + """ + from sqlalchemy import event, select + + import app.routes.chat as chat + from packages.auth.spend import pending_parked_spend + from packages.db.models.request_log import RequestLog + + make_client, fake, factory, _root = budget_env + key, key_id = await _make_budgeted_key(factory, budget_limit_cents=10) + fake.acompletion = AsyncMock(return_value=_completion("Hello!")) + + give_ups: list[int] = [] + real_give_up = chat._give_up_settlement + + async def _count_give_ups(*args, **kwargs): + give_ups.append(1) + return await real_give_up(*args, **kwargs) + + monkeypatch.setattr(chat, "_give_up_settlement", _count_give_ups) + + # Aim the cancellation at the backoff and nowhere earlier: a task is + # cancelled at its next suspension, and the handler's rollback is one, so + # make that call a coroutine that never yields. The sleep is then the only + # place the cancellation can land. + from sqlalchemy.ext.asyncio import AsyncSession + + async def _suspendless_rollback(self, *args, **kwargs): + return None + + monkeypatch.setattr(AsyncSession, "rollback", _suspendless_rollback) + + in_retry = asyncio.Event() + failures = {"n": 0} + + def _fail_the_write(conn, cursor, statement, parameters, context, executemany): + if "INSERT INTO requests_log" in statement and failures["n"] == 0: + failures["n"] += 1 + in_retry.set() + raise RuntimeError("database is locked") + + sync_engine = factory.kw["bind"].sync_engine + event.listen(sync_engine, "before_cursor_execute", _fail_the_write) + + async def _request(): + async with await make_client(key) as c: + return await c.post( + "/v1/chat/completions", + json={"model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hi"}]}, + ) + + try: + task = asyncio.ensure_future(_request()) + await in_retry.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + finally: + event.remove(sync_engine, "before_cursor_execute", _fail_the_write) + + assert failures["n"] == 1 # the write really did fail and start backing off + assert give_ups == [1] # abandoned once, not twice and not never + async with factory() as s: + assert not ( + await s.execute(select(RequestLog.id).where(RequestLog.api_key_id == key_id)) + ).all() + assert await pending_parked_spend(key_id) == 100_000 + + +async def test_cancelled_rollback_does_not_skip_the_give_up( + budget_env, monkeypatch +): + """A cancellation inside the retry handler still accounts for the cost. + + Same teardown, other delivery point: the rollback that opens the handler is + itself an await, so it can be where the cancellation lands. An exception + raised from inside a handler is not caught by this `try`'s other arms, so it + used to escape straight past the give-up below — the write lost, the + obligation unparked, and the cap reopened for exactly the key whose write + had just failed. + """ + from sqlalchemy.ext.asyncio import AsyncSession + + import app.routes.chat as chat + from packages.auth.spend import pending_parked_spend + + make_client, fake, factory, _root = budget_env + key, key_id = await _make_budgeted_key(factory, budget_limit_cents=10) + fake.acompletion = AsyncMock(return_value=_completion("Hello!")) + + give_ups: list[int] = [] + real_give_up = chat._give_up_settlement + + async def _count_give_ups(*args, **kwargs): + give_ups.append(1) + return await real_give_up(*args, **kwargs) + + rollbacks = {"n": 0} + real_rollback = AsyncSession.rollback + + async def _cancel_the_first_rollback(self, *args, **kwargs): + rollbacks["n"] += 1 + if rollbacks["n"] == 1: + raise asyncio.CancelledError + return await real_rollback(self, *args, **kwargs) + + failures = {"n": 0} + + def _fail_the_write(conn, cursor, statement, parameters, context, executemany): + if "INSERT INTO requests_log" in statement and failures["n"] == 0: + failures["n"] += 1 + raise RuntimeError("database is locked") + + monkeypatch.setattr(chat, "_give_up_settlement", _count_give_ups) + monkeypatch.setattr(AsyncSession, "rollback", _cancel_the_first_rollback) + sync_engine = factory.kw["bind"].sync_engine + from sqlalchemy import event + + event.listen(sync_engine, "before_cursor_execute", _fail_the_write) + try: + async with await make_client(key) as c: + await c.post( + "/v1/chat/completions", + json={"model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hi"}]}, + ) + finally: + event.remove(sync_engine, "before_cursor_execute", _fail_the_write) + + assert rollbacks["n"] >= 1 # the cancellation really did land there + assert failures["n"] == 1 + # The cancellation is swallowed rather than escaping, so the loop keeps + # going: the next attempt still finds the session poisoned (the failed flush + # was never rolled back), and the real rollback at the head of that handler + # clears it, so a later attempt lands the row and its charge. Pinned: the + # cost is accounted for once, and no give-up was needed to do it — before + # this the request died here with nothing billed and nothing parked. + assert give_ups == [] + assert await _get_spent(factory, key_id) == 100_000 + assert await pending_parked_spend(key_id) == 0 + from sqlalchemy import select + + from packages.db.models.request_log import RequestLog + + async with factory() as s: + rows = ( + await s.execute(select(RequestLog.id).where(RequestLog.api_key_id == key_id)) + ).all() + assert len(rows) == 1 + + # ── Durable recovery: the park outlives the process that lost it ────── class _AckLossSession(AsyncSession): From 93bb71f8ec12196b2079895a730e4b3bf965e057 Mon Sep 17 00:00:00 2001 From: Hasit Bhatt Date: Fri, 25 Sep 2026 02:38:28 -0700 Subject: [PATCH 13/23] fix(budget): stop a park from being held durably and in memory at once MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `_insert_park` answered "not durable" to every failure except a unique violation, including a COMMIT that applied and lost its ack. The caller then stored the amount in `_unsettled` beside the row it had just written, another worker folded and deleted the row, and the stale memory copy re-filed as a new park under a `trace_id` that no longer collided — billing one delivery twice and leaving the key pinned on its cap with no debt left to fold. Ask the database what landed instead of guessing, and drop the memory hold for any row a fold actually billed. An unreadable park ledger folded into the total as zero, which dispatched a key whose cap was held shut only by parks this worker could not see. The two reads are not the same connection either: the counter rides the request's, the ledger opens a new one, so a checkout timeout hid every park while the request worked fine. `pending_parked_spend` now answers `None` for unknown and the pre-check maps unknown onto the cap. Oldest-debt-first was decided by `trace_id`, because `created_at` came from `CURRENT_TIMESTAMP` — one second wide on SQLite, and a recovered outage re-files a whole batch inside a single pre-check. The row stamps itself Python-side at sub-second resolution so the tiebreak is the rare path it is documented as. --- packages/auth/spend.py | 98 +++++++++++---- packages/db/models/budget_park.py | 33 ++++- tests/unit/test_budget_spend.py | 196 ++++++++++++++++++++++++++++-- 3 files changed, 289 insertions(+), 38 deletions(-) diff --git a/packages/auth/spend.py b/packages/auth/spend.py index 4dfae65..e7112ec 100644 --- a/packages/auth/spend.py +++ b/packages/auth/spend.py @@ -44,16 +44,16 @@ # 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 the counter never moves. The obligation is parked here — one row per +# back and the counter never moves. The obligation is parked — one row per # settlement, keyed by its `trace_id` — and keeps counting against the cap until # a budget pre-check folds it into `spent_microcents`. # # The park is a database table, not process memory, because the deployment # stops its machine whenever it goes idle: an in-memory obligation is lost on # the next cold start, which reopens the cap for exactly the key the failure -# was about to protect. `_unsettled` below only holds an amount while the -# database itself is unreachable — the same outage that caused the park — and a -# later pre-check re-files it once a write goes through again. +# was about to protect. `_unsettled` below is the hold for when even that write +# cannot be made — it maps `(api_key_id, trace_id)` to the amount, and a later +# pre-check re-files it once a write goes through again. _unsettled: dict[tuple[str, str], int] = {} @@ -61,14 +61,42 @@ class _FoldConflict(Exception): """A concurrent worker folded the same park first; the loser retries later.""" +async def _park_is_durable(trace_id: str) -> bool: + """Whether a park row for this `trace_id` is committed. + + A commit that applied but whose ack never came back raises exactly like a + failure, and holding a memory copy beside the durable row puts one + obligation in both ledgers — the double-bill the `trace_id` key exists to + absorb. It reads on a fresh session because the failed one is closed by the + time this runs, and answers False when it cannot read at all: with the + database truly unreachable the memory hold is all that keeps the cap honest. + """ + from packages.db import session as session_mod + + factory = session_mod._session_factory + if factory is None: + return False + try: + async with factory() as s: + return ( + await s.scalar( + select(BudgetPark.trace_id).where(BudgetPark.trace_id == trace_id) + ) + ) is not None + except Exception: + return False + + async def _insert_park(*, trace_id: str, api_key_id: str, microcents: int) -> bool: """Persist one parked obligation. Returns True when it is durable. Idempotent on `trace_id`: a commit that applied but whose ack was lost retries into the same primary key instead of recording the obligation a - second time. Returns False when the database is unavailable (or this is a - unit test with no session factory), leaving the caller to hold the amount - in memory. A cancellation propagates with the memory copy still held. + second time — and when the retry does not reach the server to be told that, + the probe below asks the database directly rather than reporting a write that + landed as one that did not. Returns False only when the obligation is + genuinely not durable, leaving the caller to hold the amount in memory. A + cancellation propagates with the memory copy still held. """ from packages.db import session as session_mod @@ -92,7 +120,7 @@ async def _insert_park(*, trace_id: str, api_key_id: str, microcents: int) -> bo except asyncio.CancelledError: raise except Exception: - return False + return await _park_is_durable(trace_id) async def record_unsettled_spend( @@ -121,8 +149,16 @@ async def record_unsettled_spend( _unsettled[key] = microcents -async def pending_parked_spend(api_key_id: str) -> int: - """The outstanding park for a key: durable rows plus whatever is memory-only.""" +async def pending_parked_spend(api_key_id: str) -> int | None: + """The outstanding park for a key: durable rows plus whatever is memory-only. + + ``None`` means the durable ledger could not be read, which is a different + answer from ``0``: the table is the only record of a park written by another + worker, or by this one before it stopped, so folding a failed read into the + total reopens the cap for exactly the key the park exists to hold shut. The + memory half still counts on the way to a real durable read failing, because + that half is what an outage is expected to lose and what recovers after it. + """ from packages.db import session as session_mod key = str(api_key_id) @@ -141,10 +177,9 @@ async def pending_parked_spend(api_key_id: str) -> int: ) ) ).scalar() - total += int(stored or 0) except Exception: - pass - return total + return None + return total + int(stored or 0) async def settle_parked_spend(api_key_id: str, cap_microcents: int) -> int: @@ -184,6 +219,9 @@ async def settle_parked_spend(api_key_id: str, cap_microcents: int) -> int: trace_id=trace_id, api_key_id=key, microcents=amount ): _unsettled.pop((held_key, trace_id), None) + move = 0 + settling: list[str] = [] + trim: tuple[str, int, int] | None = None try: async with factory() as s: async with s.begin(): @@ -199,18 +237,15 @@ async def settle_parked_spend(api_key_id: str, cap_microcents: int) -> int: await s.execute( select(BudgetPark.trace_id, BudgetPark.microcents) .where(BudgetPark.api_key_id == key) - # Oldest debt first. `created_at` alone is not a total - # order — it is second-resolution on SQLite and ties on - # Postgres — and two workers computing the same fold - # have to agree on which row is the partial one, so - # `trace_id` breaks the tie. + # Oldest debt first. `created_at` is stamped + # Python-side with sub-second resolution, but two + # workers computing the same fold still have to agree on + # which row is the partial one, so `trace_id` breaks any + # tie rather than letting the order depend on a race. .order_by(BudgetPark.created_at, BudgetPark.trace_id) ) ).all() - move = 0 room = cap_microcents - spent - settling: list[str] = [] - trim: tuple[str, int, int] | None = None for trace_id, microcents in rows: microcents = int(microcents) if microcents <= room: @@ -249,9 +284,17 @@ async def settle_parked_spend(api_key_id: str, cap_microcents: int) -> int: ) if trimmed.rowcount != 1: raise _FoldConflict - return move except _FoldConflict: return 0 + # Past the commit, so everything in `settling` is billed and gone from the + # table. A memory hold for one of those rows is now worse than stale: the + # next pre-check re-files it as a brand-new park — the `trace_id` no longer + # collides, this transaction deleted the row — and the same delivery is + # charged twice against a cap that has no idea it moved. Only `settling`: + # a trimmed row is still owed, and the hold on it has to stay. + for _settled in settling: + _unsettled.pop((key, _settled), None) + return move async def read_spent(db: AsyncSession, api_key_id: str) -> int: @@ -271,7 +314,8 @@ async def budget_precheck(db: AsyncSession, api_key_id: str, cap_microcents: int outage is neither a window of free requests nor a park that can never clear, and the result adds whatever is still pending rather than reading only the counter: when the fold could not commit, the obligation still has to block - dispatch. + dispatch. A park ledger that cannot be read at all is answered as the cap + rather than as no debt. ``cap_microcents`` is ``ApiKey.budget_limit_cents`` scaled by ``MICROCENTS_PER_CENT``, not the column itself — passing the raw cents value @@ -280,6 +324,12 @@ async def budget_precheck(db: AsyncSession, api_key_id: str, cap_microcents: int key = str(api_key_id) spent = await read_spent(db, key) pending = await pending_parked_spend(key) + if pending is None: + # Unknown debt is answered as full debt, the way the durability probe + # that wrote the park reads an unreadable database as "not settled". + # The alternative is a key whose cap is held shut only by parked rows + # dispatching freely while the ledger is down. + return cap_microcents if pending: try: await settle_parked_spend(key, cap_microcents) @@ -289,6 +339,8 @@ async def budget_precheck(db: AsyncSession, api_key_id: str, cap_microcents: int pass spent = await read_spent(db, key) pending = await pending_parked_spend(key) + if pending is None: + return cap_microcents return spent + pending diff --git a/packages/db/models/budget_park.py b/packages/db/models/budget_park.py index a233157..4a55f26 100644 --- a/packages/db/models/budget_park.py +++ b/packages/db/models/budget_park.py @@ -8,13 +8,24 @@ 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. 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`, so the obligation is never both -parked and billed, and never neither. +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 sqlalchemy import BigInteger, String +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 @@ -26,3 +37,15 @@ class BudgetPark(Base, UUIDMixin, TimestampMixin): 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, + ) diff --git a/tests/unit/test_budget_spend.py b/tests/unit/test_budget_spend.py index d5a23a4..32f8351 100644 --- a/tests/unit/test_budget_spend.py +++ b/tests/unit/test_budget_spend.py @@ -1,11 +1,16 @@ """Unit tests for packages.auth.spend — atomic budget charge under a hard cap.""" import asyncio +import contextlib import pytest +from sqlalchemy import select +from sqlalchemy.exc import OperationalError +from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker from packages.auth.spend import ( MICROCENTS_PER_CENT, + budget_precheck, charge_budget, is_exhausted, pending_parked_spend, @@ -13,6 +18,17 @@ record_unsettled_spend, settle_parked_spend, ) +from packages.db.models.budget_park import BudgetPark + + +@pytest.fixture(autouse=True) +def _isolated_memory_holds(): + """`_unsettled` is process state, so one test's leftover hold is another's bug.""" + from packages.auth import spend as spend_mod + + spend_mod._unsettled.clear() + yield + spend_mod._unsettled.clear() @pytest.fixture @@ -59,8 +75,6 @@ async def test_concurrent_charges_never_exceed_cap(tmp_sqlite_url): still makes the outcome deterministic — one charge fits, the other's guard matches no row and its clamp fills the counter to exactly `cap`. """ - from sqlalchemy.ext.asyncio import async_sessionmaker - from packages.db.engine import build_engine from packages.db.models.api_key import ApiKey from packages.db.models.base import Base @@ -97,8 +111,34 @@ def test_microcent_conversion_constant(): @pytest.fixture async def parked_env(tmp_sqlite_url): """Engine + global session factory, so the park ledger is durable here.""" - from sqlalchemy.ext.asyncio import async_sessionmaker + async with _park_ledger(tmp_sqlite_url) as factory: + yield factory + + +class _LostAckParkSession(AsyncSession): + """``AsyncSession`` whose park COMMIT applies and then reports failure. + Losing the ack is the state that matters: a commit that landed and a commit + that failed look identical to the caller, and treating the first as the + second is what puts one obligation in the table and in memory at once. Only + a session inserting a park is rigged, so the fixture's own bookkeeping + commits run untouched. + """ + + lost_acks_left = 0 + + async def commit(self): + if type(self).lost_acks_left > 0 and any( + isinstance(o, BudgetPark) for o in self.sync_session.new + ): + type(self).lost_acks_left -= 1 + await super().commit() + raise OperationalError("COMMIT", {}, Exception("connection reset before ack")) + return await super().commit() + + +@contextlib.asynccontextmanager +async def _park_ledger(tmp_sqlite_url, session_class=AsyncSession): from packages.db import session as session_mod from packages.db.engine import build_engine from packages.db.models.base import Base @@ -106,7 +146,7 @@ async def parked_env(tmp_sqlite_url): engine = build_engine(tmp_sqlite_url) async with engine.begin() as conn: await conn.run_sync(Base.metadata.create_all) - factory = async_sessionmaker(engine, expire_on_commit=False) + factory = async_sessionmaker(engine, expire_on_commit=False, class_=session_class) old, session_mod._session_factory = session_mod._session_factory, factory try: yield factory @@ -115,6 +155,19 @@ async def parked_env(tmp_sqlite_url): await engine.dispose() +async def _parks(factory, api_key_id: str) -> dict[str, int]: + """What is still parked for a key, by `trace_id`.""" + async with factory() as s: + rows = ( + await s.execute( + select(BudgetPark.trace_id, BudgetPark.microcents).where( + BudgetPark.api_key_id == api_key_id + ) + ) + ).all() + return {trace_id: int(microcents) for trace_id, microcents in rows} + + async def _parked_key(factory, *, spent: int = 0): from packages.db.models.api_key import ApiKey @@ -154,13 +207,92 @@ async def test_recorded_park_folds_exactly_once(parked_env): assert await read_spent(s, key_id) == 7_000 -async def test_duplicate_park_record_is_idempotent(parked_env): - """Retrying a park write after an ack loss must not record it twice.""" +async def test_a_park_that_lost_its_ack_is_not_also_held_in_memory( + tmp_sqlite_url, monkeypatch +): + """A commit that applied must not report itself as one that did not. + + The old test of this name called `record_unsettled_spend` twice against a + working database, which only exercises the retried insert — the durability + probe never ran. Here the COMMIT lands and the ack does not, so `_insert_park` + has to ask the table instead of trusting its own failure. Reporting it as + not durable is what left the obligation parked *and* held in memory: another + worker folds and deletes the row, this process then re-files its stale copy + under a `trace_id` that no longer collides, and the same delivery is billed + twice against a key that is now pinned on its cap with no debt left to fold. + """ + from packages.auth import spend as spend_mod + + async with _park_ledger(tmp_sqlite_url, _LostAckParkSession) as factory: + key_id = await _parked_key(factory) + monkeypatch.setattr(_LostAckParkSession, "lost_acks_left", 1) + await record_unsettled_spend( + trace_id="t-ack", api_key_id=key_id, microcents=900 + ) + + assert spend_mod._unsettled == {} # the probe found the durable row + assert await _parks(factory, key_id) == {"t-ack": 900} + assert await pending_parked_spend(key_id) == 900 + + assert await settle_parked_spend(key_id, 10_000) == 900 + assert await pending_parked_spend(key_id) == 0 + async with factory() as s: + assert await read_spent(s, key_id) == 900 + + +async def test_folding_a_row_this_process_also_holds_clears_the_hold( + parked_env, monkeypatch +): + """The fold that bills a durable row drops the memory hold for it too. + + A path the durability probe does not close: the probe reads on a session of + its own and can fail while the row is real, and this same call then picks + that row up and bills it. Leaving the hold behind means the next pre-check + re-files it as a new park — the row it mirrored is gone, so nothing collides. + Only fully-billed rows are dropped; a trimmed one is still owed. + """ + from packages.auth import spend as spend_mod + key_id = await _parked_key(parked_env) - await record_unsettled_spend(trace_id="t-dup", api_key_id=key_id, microcents=900) - # The ack never came back, so the caller retries the identical record. - await record_unsettled_spend(trace_id="t-dup", api_key_id=key_id, microcents=900) - assert await pending_parked_spend(key_id) == 900 + await record_unsettled_spend(trace_id="t-both", api_key_id=key_id, microcents=1_200) + spend_mod._unsettled[(key_id, "t-both")] = 1_200 + + async def _unreachable(**kwargs): + return False + + monkeypatch.setattr(spend_mod, "_insert_park", _unreachable) + assert await settle_parked_spend(key_id, 10_000) == 1_200 + assert (key_id, "t-both") not in spend_mod._unsettled + assert await pending_parked_spend(key_id) == 0 + + monkeypatch.undo() + assert await settle_parked_spend(key_id, 10_000) == 0 # nothing re-files + async with parked_env() as s: + assert await read_spent(s, key_id) == 1_200 + + +async def test_a_fold_that_overshoots_the_cap_keeps_its_memory_hold( + parked_env, monkeypatch +): + """A trimmed row is still owed, so the hold beside it must stay. + + The reconcile in the test above is deliberately narrow: running it over the + trimmed row as well would write off the part of a delivery the cap could not + absorb. + """ + from packages.auth import spend as spend_mod + + key_id = await _parked_key(parked_env, spent=9_000) + await record_unsettled_spend(trace_id="t-trim", api_key_id=key_id, microcents=2_000) + spend_mod._unsettled[(key_id, "t-trim")] = 2_000 + + async def _unreachable(**kwargs): + return False + + monkeypatch.setattr(spend_mod, "_insert_park", _unreachable) + assert await settle_parked_spend(key_id, 10_000) == 1_000 + assert spend_mod._unsettled[(key_id, "t-trim")] == 2_000 + assert await _parks(parked_env, key_id) == {"t-trim": 1_000} async def test_park_beyond_the_remainder_bills_what_fits(parked_env): @@ -197,16 +329,24 @@ async def test_parked_queue_drains_oldest_first_as_the_cap_opens(parked_env): The oversized head takes the whole allowance, so the younger park behind it waits — not lost, just queued. Raising the cap reopens room, and the fold keeps its promise that the counter reaches `min(cap, spent + debt)`. + + Which row shrinks is the assertion that matters: the totals come out the + same either way, so only the ledger shows whether the head of the queue was + the older obligation or whichever `trace_id` sorts first. """ key_id = await _parked_key(parked_env, spent=9_000) await record_unsettled_spend(trace_id="t-old", api_key_id=key_id, microcents=2_000) await record_unsettled_spend(trace_id="t-new", api_key_id=key_id, microcents=500) assert await settle_parked_spend(key_id, 10_000) == 1_000 + # `t-old` absorbed the whole allowance and kept its remainder; the younger, + # smaller park behind it is untouched. + assert await _parks(parked_env, key_id) == {"t-old": 1_000, "t-new": 500} assert await settle_parked_spend(key_id, 10_000) == 0 # no room left assert await pending_parked_spend(key_id) == 1_500 assert await settle_parked_spend(key_id, 11_000) == 1_000 + assert await _parks(parked_env, key_id) == {"t-new": 500} assert await pending_parked_spend(key_id) == 500 assert await settle_parked_spend(key_id, 12_000) == 500 assert await pending_parked_spend(key_id) == 0 @@ -214,6 +354,42 @@ async def test_parked_queue_drains_oldest_first_as_the_cap_opens(parked_env): assert await read_spent(s, key_id) == 11_500 +async def test_an_unreadable_park_ledger_does_not_read_as_no_debt(parked_env): + """A failed durable SUM has to block dispatch, not clear the key's cap. + + The durable table is the only record of what another worker, or this process + before it stopped, still owes — and it is a *different* read from the rest of + the pre-check: the counter comes in on the request's own connection, while + the ledger opens a fresh one from the factory. So a pool checkout that times + out can hide every park while the request itself would otherwise work, and + folding that failure into the total as zero dispatches the exact key the park + exists to hold shut. + + The factory is swapped by hand rather than with `monkeypatch`: that teardown + runs after the fixture's and would restore the rigged one, leaving every + later test in the session reading a disposed engine. + """ + from packages.db import session as session_mod + + key_id = await _parked_key(parked_env, spent=1_000) + cap = 10_000 + async with parked_env() as s: + assert await is_exhausted(s, key_id, cap) is False + + def _blinded(): + raise TimeoutError("database connection checkout timed out") + + real = session_mod._session_factory + try: + session_mod._session_factory = _blinded + assert await pending_parked_spend(key_id) is None + async with parked_env() as s: + assert await budget_precheck(s, key_id, cap) == cap + assert await is_exhausted(s, key_id, cap) is True + finally: + session_mod._session_factory = real + + async def test_concurrent_folds_bill_the_park_once(parked_env): """Two workers folding the same park move it exactly once.""" key_id = await _parked_key(parked_env) From eac6f2ac93c7ddd4e0ce4f3e85892b8259960ea2 Mon Sep 17 00:00:00 2001 From: Hasit Bhatt Date: Fri, 25 Sep 2026 09:52:37 -0700 Subject: [PATCH 14/23] fix(budget): refresh session snapshot after folding parked spend in precheck --- packages/auth/spend.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/packages/auth/spend.py b/packages/auth/spend.py index e7112ec..247d8f0 100644 --- a/packages/auth/spend.py +++ b/packages/auth/spend.py @@ -337,6 +337,12 @@ async def budget_precheck(db: AsyncSession, api_key_id: str, cap_microcents: int # The fold runs on sessions of its own, so there is nothing to # roll back here — and the re-read below still counts the park. pass + # End db's read transaction so its SQLite/Postgres snapshot doesn't stay + # pinned to the pre-fold spend counter, then re-read on a fresh snapshot. + try: + await db.rollback() + except Exception: + pass spent = await read_spent(db, key) pending = await pending_parked_spend(key) if pending is None: From 56606e1982786d67f201e7274aea560267d59a6a Mon Sep 17 00:00:00 2001 From: Hasit Bhatt Date: Fri, 25 Sep 2026 12:38:02 -0700 Subject: [PATCH 15/23] fix(budget): drop a park whose charge already landed before folding A commit whose acknowledgement was lost, combined with a durability probe that also failed, parks an obligation whose charge is already in spent_microcents. Folding it re-charges the same delivery, and because the fold deletes the park row there is nothing left to correct it. The request-log row and the charge share one transaction, so the log row proves the charge landed: clear those parks without billing them, and drop the matching in-memory hold so a later pre-check cannot re-file them. --- packages/auth/spend.py | 28 ++++++++++++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/packages/auth/spend.py b/packages/auth/spend.py index 247d8f0..8ebe87b 100644 --- a/packages/auth/spend.py +++ b/packages/auth/spend.py @@ -39,6 +39,7 @@ from packages.db.models.api_key import ApiKey from packages.db.models.budget_park import BudgetPark +from packages.db.models.request_log import RequestLog MICROCENTS_PER_CENT = 10_000 @@ -245,6 +246,31 @@ async def settle_parked_spend(api_key_id: str, cap_microcents: int) -> int: .order_by(BudgetPark.created_at, BudgetPark.trace_id) ) ).all() + # A log row and its budget charge are committed atomically. + # If a commit acknowledgement was lost and the durability + # probe also failed, the matching park is only a fallback + # record; the request-log row proves the charge already landed. + logged = set( + ( + await s.scalars( + select(RequestLog.trace_id).where( + RequestLog.trace_id.in_( + [trace_id for trace_id, _amount in rows] + ) + ) + ) + ).all() + ) if rows else set() + already_charged = [trace_id for trace_id, _amount in rows if trace_id in logged] + if already_charged: + cleared = await s.execute( + delete(BudgetPark).where( + BudgetPark.trace_id.in_(already_charged) + ) + ) + if cleared.rowcount != len(already_charged): + raise _FoldConflict + rows = [row for row in rows if row[0] not in logged] room = cap_microcents - spent for trace_id, microcents in rows: microcents = int(microcents) @@ -294,6 +320,8 @@ async def settle_parked_spend(api_key_id: str, cap_microcents: int) -> int: # a trimmed row is still owed, and the hold on it has to stay. for _settled in settling: _unsettled.pop((key, _settled), None) + for _already_charged in already_charged: + _unsettled.pop((key, _already_charged), None) return move From 9b82eb2e4efd7be5b43d2194062916f9be4c931f Mon Sep 17 00:00:00 2001 From: Hasit Bhatt Date: Fri, 25 Sep 2026 19:52:21 -0700 Subject: [PATCH 16/23] fix(budget): price adapter-fault delivery instead of charging remaining budget --- app/routes/chat.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/app/routes/chat.py b/app/routes/chat.py index b99bc6c..4a7a8b0 100644 --- a/app/routes/chat.py +++ b/app/routes/chat.py @@ -1154,6 +1154,13 @@ async def _commit_row(*, retry: bool) -> None: logger.warning( "chat_completion_stream_adapter_error", served_model=agg_model, ) + # An adapter fault is our bug, not a choice the caller made: + # price what reached the client instead of leaving the + # settlement unknown, which would charge a budgeted key its + # entire remaining budget. Nothing delivered settles at the + # 0 the row already records. + agg_usage = _settle_unmeasured_stream(agg_usage, agg_output_chars, body) + usage_seen = True aclose = getattr(stream_obj, "aclose", None) with anyio.CancelScope(shield=True): if aclose is not None: From d26d14f698e95becd29159fd6e41d013f55e4591 Mon Sep 17 00:00:00 2001 From: Hasit Bhatt Date: Fri, 25 Sep 2026 21:00:41 -0700 Subject: [PATCH 17/23] fix(budget): fail closed when a delivered completion has tokens but no price --- app/routes/chat.py | 64 ++++++++++++++++++++++++++++++++++++++++++++-- 1 file changed, 62 insertions(+), 2 deletions(-) diff --git a/app/routes/chat.py b/app/routes/chat.py index 4a7a8b0..606e14e 100644 --- a/app/routes/chat.py +++ b/app/routes/chat.py @@ -445,6 +445,28 @@ def _lookup_priced_model(model_id: str | None): return None +def _has_known_price( + *, + litellm_cost_usd: float | None, + model_id: str | None, + fallback_model: str | None = None, +) -> bool: + """True when a delivered completion's cost is measurable, even if it is 0. + + Mirrors `_compute_cost_microcents`' two tiers without computing: an + authoritative LiteLLM cost, or any catalog entry. A 0.0/0.0 entry is a + known-free model, not an unknown cost. Tokens with neither are + unpriceable — a custom upstream LiteLLM can't cost, or a model absent + from our catalog — and must not settle at 0 for a budgeted key, or the + lifetime cap would stand still while the upstream still bills us. + """ + if litellm_cost_usd is not None and litellm_cost_usd > 0: + return True + return ( + _lookup_priced_model(model_id) or _lookup_priced_model(fallback_model) + ) is not None + + @router.post("/chat/completions") async def chat_completions( body: ChatCompletionRequest, @@ -942,6 +964,12 @@ def _settlement_amount() -> int: (fail-closed) so no client-side choice can bypass the cap. Once a usage frame was delivered the cost is known — even if the stream then died — and the recorded cost is charged. + A frame that carries tokens but no price (a custom upstream + LiteLLM can't cost, or a model absent from our catalog) is + unknown too: charging the recorded 0 would let the cap + stand still while the upstream still bills us. A + catalog-listed free model, or an empty delivery, is + known-zero and still settles at the 0 the row records. The raised amount is written back into the row, not just into the counter: a charge only the counter saw would leave @@ -956,6 +984,24 @@ def _settlement_amount() -> int: - (getattr(kc, "_budget_spent", 0) or 0), ) row_values["cost_microcents"] = actual + elif ( + not actual + and ( + row_values.get("input_tokens") + or row_values.get("output_tokens") + ) + and not _has_known_price( + litellm_cost_usd=(agg_usage or {}).get("cost_usd"), + model_id=row_values.get("model_resolved"), + fallback_model=row_values.get("model_requested"), + ) + ): + actual = max( + actual, + (getattr(kc, "_budget_cap", 0) or 0) + - (getattr(kc, "_budget_spent", 0) or 0), + ) + row_values["cost_microcents"] = actual return actual async def _commit_row(*, retry: bool) -> None: @@ -1318,7 +1364,10 @@ async def _commit_row(*, retry: bool) -> None: # the full remaining allowance so a delivered completion can never cost # nothing, and record that amount on the row: a charge only the counter # saw would leave a key exhausted by an amount nothing in its own - # request history accounts for. Gated on having actually received a + # request history accounts for. Usage with tokens but no price is the + # same unknown (a custom upstream LiteLLM can't cost, or a model + # absent from our catalog); a catalog-listed free model is known-zero + # and keeps its 0. Gated on having actually received a # completion dict: a request that failed before the upstream answered # (response == {}, e.g. the re-raised HTTPException above, whose # status_code never left 200) charges its recorded ~0 cost instead — @@ -1328,7 +1377,18 @@ async def _commit_row(*, retry: bool) -> None: and status_code < 400 and isinstance(response, dict) and response - and not response.get("usage") + and ( + not response.get("usage") + or ( + not (log.cost_microcents or 0) + and not _has_known_price( + litellm_cost_usd=(response.get("usage") or {}).get("cost_usd") + or (response.get("_orca_meta") or {}).get("cost_usd"), + model_id=log.model_resolved, + fallback_model=log.model_requested, + ) + ) + ) ): settle_amount = max( log.cost_microcents or 0, From f0eb0ce60f747d857f05091164e5df6429308247 Mon Sep 17 00:00:00 2001 From: Hasit Bhatt Date: Fri, 25 Sep 2026 22:36:15 -0700 Subject: [PATCH 18/23] fix(budget): drop comments in adapter-error settlement branch --- app/routes/chat.py | 10 ---------- 1 file changed, 10 deletions(-) diff --git a/app/routes/chat.py b/app/routes/chat.py index 606e14e..f03469a 100644 --- a/app/routes/chat.py +++ b/app/routes/chat.py @@ -1190,21 +1190,11 @@ async def _commit_row(*, retry: bool) -> None: # Re-raise so asyncio/Starlette see proper cancel propagation. raise except AdapterError: - # The protocol adapter downstream of us failed and threw - # this in rather than closing us: our own fault, not the - # caller's. Without this branch the close would be - # indistinguishable from a disconnect and every adapter bug - # would be filed as 499/client_disconnect. error_type = "adapter_error" status_code = 500 logger.warning( "chat_completion_stream_adapter_error", served_model=agg_model, ) - # An adapter fault is our bug, not a choice the caller made: - # price what reached the client instead of leaving the - # settlement unknown, which would charge a budgeted key its - # entire remaining budget. Nothing delivered settles at the - # 0 the row already records. agg_usage = _settle_unmeasured_stream(agg_usage, agg_output_chars, body) usage_seen = True aclose = getattr(stream_obj, "aclose", None) From 8fa485065f0ad17cfadf902a3463bfdb406950d4 Mon Sep 17 00:00:00 2001 From: Hasit Bhatt Date: Fri, 25 Sep 2026 23:00:25 -0700 Subject: [PATCH 19/23] fix(budget): gate the unpriced-stream arm on status and read parks before spend --- app/routes/chat.py | 15 ++++++++++----- packages/auth/spend.py | 13 ++++++++++++- 2 files changed, 22 insertions(+), 6 deletions(-) diff --git a/app/routes/chat.py b/app/routes/chat.py index f03469a..3473a39 100644 --- a/app/routes/chat.py +++ b/app/routes/chat.py @@ -966,10 +966,14 @@ def _settlement_amount() -> int: if the stream then died — and the recorded cost is charged. A frame that carries tokens but no price (a custom upstream LiteLLM can't cost, or a model absent from our catalog) is - unknown too: charging the recorded 0 would let the cap - stand still while the upstream still bills us. A - catalog-listed free model, or an empty delivery, is - known-zero and still settles at the 0 the row records. + unknown too — but only on a stream that completed + normally: charging the recorded 0 would let the cap stand + still while the upstream still bills us. An error ending + (disconnect, upstream or adapter fault) is priced by its + delivery estimate and keeps that estimate even at 0, so + this arm is gated on status_code < 400. A catalog-listed + free model, or an empty delivery, is known-zero and still + settles at the 0 the row records. The raised amount is written back into the row, not just into the counter: a charge only the counter saw would leave @@ -985,7 +989,8 @@ def _settlement_amount() -> int: ) row_values["cost_microcents"] = actual elif ( - not actual + status_code < 400 + and not actual and ( row_values.get("input_tokens") or row_values.get("output_tokens") diff --git a/packages/auth/spend.py b/packages/auth/spend.py index 8ebe87b..3c5d9c8 100644 --- a/packages/auth/spend.py +++ b/packages/auth/spend.py @@ -350,7 +350,6 @@ async def budget_precheck(db: AsyncSession, api_key_id: str, cap_microcents: int asks whether the key has spent a ten-thousandth of its budget. """ key = str(api_key_id) - spent = await read_spent(db, key) pending = await pending_parked_spend(key) if pending is None: # Unknown debt is answered as full debt, the way the durability probe @@ -358,6 +357,18 @@ async def budget_precheck(db: AsyncSession, api_key_id: str, cap_microcents: int # The alternative is a key whose cap is held shut only by parked rows # dispatching freely while the ledger is down. return cap_microcents + # The park ledger is read before the counter, and the counter on a fresh + # snapshot: a fold both moves `spent_microcents` and empties the park + # queue, so a counter read first can pair a pre-fold spend with a zero + # pending — understating the spend by exactly the folded amount with no + # evidence left to trigger the re-read. Parks-first keeps the evidence: + # a fold racing these reads leaves `pending > 0`, which takes the + # re-read below. + try: + await db.rollback() + except Exception: + pass + spent = await read_spent(db, key) if pending: try: await settle_parked_spend(key, cap_microcents) From 12cc7fc6144fed784501ee5eabe41e6d9fafd972 Mon Sep 17 00:00:00 2001 From: Hasit Bhatt Date: Sat, 26 Sep 2026 00:02:04 -0700 Subject: [PATCH 20/23] fix(budget): swallow cancellation on pre-check rollbacks so the request survives --- packages/auth/spend.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/packages/auth/spend.py b/packages/auth/spend.py index 3c5d9c8..7c21efa 100644 --- a/packages/auth/spend.py +++ b/packages/auth/spend.py @@ -363,10 +363,13 @@ async def budget_precheck(db: AsyncSession, api_key_id: str, cap_microcents: int # pending — understating the spend by exactly the folded amount with no # evidence left to trigger the re-read. Parks-first keeps the evidence: # a fold racing these reads leaves `pending > 0`, which takes the - # re-read below. + # re-read below. The rollback ends whatever pending work the request + # session carries; a cancellation landing on it must not escape — the + # pre-check has no retry of its own and the caller's loop owns the + # session, so swallow it and keep reading on the (stale) snapshot. try: await db.rollback() - except Exception: + except BaseException: pass spent = await read_spent(db, key) if pending: @@ -380,7 +383,7 @@ async def budget_precheck(db: AsyncSession, api_key_id: str, cap_microcents: int # pinned to the pre-fold spend counter, then re-read on a fresh snapshot. try: await db.rollback() - except Exception: + except BaseException: pass spent = await read_spent(db, key) pending = await pending_parked_spend(key) From 99b4326e6456500a7ac7016e5fa742e9cc2f22ae Mon Sep 17 00:00:00 2001 From: Hasit Bhatt Date: Sat, 26 Sep 2026 01:45:53 -0700 Subject: [PATCH 21/23] fix(migrations): make lifetime-spend seed monotonic so racing charges cannot drop history --- packages/db/migrate.py | 21 ++++++++---- tests/unit/test_budget_migration.py | 52 +++++++++++++++++++++++++++++ 2 files changed, 67 insertions(+), 6 deletions(-) diff --git a/packages/db/migrate.py b/packages/db/migrate.py index 2ab4c41..2e4d0c0 100644 --- a/packages/db/migrate.py +++ b/packages/db/migrate.py @@ -139,17 +139,26 @@ async def ensure_budget_columns(engine) -> None: # 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. + # would then never retry. + # + # The repair is monotonic: it only updates a key whose recorded counter + # lags behind the sum of its request logs (`spent_microcents < computed SUM`). + # A live charge racing the boot moves `spent_microcents` and inserts its log + # row in one transaction; gating on `spent_microcents < SUM` rather than + # `spent_microcents = 0` guarantees that a key which took traffic before + # this seed ran is still brought up to its true historical total instead of + # permanently dropping pre-upgrade spend. Steady-state boots match zero + # rows. 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" + ") WHERE budget_limit_cents IS NOT NULL AND spent_microcents < (" + " SELECT CAST(COALESCE(SUM(cost_microcents), 0) AS BIGINT) FROM requests_log " + " WHERE requests_log.api_key_id = api_keys.id" + ")" ) ) diff --git a/tests/unit/test_budget_migration.py b/tests/unit/test_budget_migration.py index ea77005..4f2c694 100644 --- a/tests/unit/test_budget_migration.py +++ b/tests/unit/test_budget_migration.py @@ -135,6 +135,58 @@ async def test_seed_repairs_a_boot_that_died_before_it(tmp_sqlite_url): await engine.dispose() +async def test_seed_repairs_a_key_racing_live_charges(tmp_sqlite_url): + """A racing request that moved spent_microcents before the seed ran is repaired. + + If live traffic on another worker commits a charge before the seed statement + evaluates, `spent_microcents` leaves 0. Gating on `spent_microcents = 0` would + skip the key and permanently drop its pre-upgrade history. The monotonic + `spent_microcents < computed SUM` gate detects that the counter lags behind + the key's request-log total and restores the true sum. + """ + engine = await _legacy_deploy_engine(tmp_sqlite_url) + try: + # Simulate: column added, and worker A immediately committed a 300 charge + # with its request log row for w1 (which had 2500 historical spend). + # Counter is at 300, but request log total is 2500 + 300 = 2800. + async with engine.begin() as conn: + await conn.execute( + text( + "ALTER TABLE api_keys ADD COLUMN spent_microcents BIGINT " + "NOT NULL DEFAULT 0" + ) + ) + await conn.execute( + text("UPDATE api_keys SET spent_microcents = 300 WHERE workspace_id = 'w1'") + ) + # Add the 300 request log row that worker A committed with the charge + w1_id = await conn.scalar( + text("SELECT id FROM api_keys WHERE workspace_id = 'w1'") + ) + await conn.execute( + text( + "INSERT INTO requests_log (id, workspace_id, api_key_id, trace_id, " + "model_requested, model_resolved, provider, routing_strategy, " + "input_tokens, output_tokens, cost_microcents, latency_ms, status_code) " + "VALUES ('r-race-1', 'w1', :k, 't-race-1', 'gpt-4o-mini', 'gpt-4o-mini', " + "'openai', 'balanced', 1, 1, 300, 10, 200)" + ), + {"k": w1_id}, + ) + + assert await _spent(engine, "w1") == 300 + + # Boot runs ensure_budget_columns + await ensure_budget_columns(engine) + + # w1 counter is restored to the full cumulative total (2500 + 300 = 2800) + assert await _spent(engine, "w1") == 2800 + assert await _spent(engine, "w2") == 700 + assert await _spent(engine, "w3") == 0 + finally: + await engine.dispose() + + async def test_orm_reads_work_after_upgrade(tmp_sqlite_url): # The original 503 failure mode: the ORM SELECTs every mapped column, so # without the migration every authenticated request broke on the upgraded From c387d54654d625f6a5b8d86095be2ed2b9aa9851 Mon Sep 17 00:00:00 2001 From: Hasit Bhatt Date: Sat, 26 Sep 2026 18:43:06 -0700 Subject: [PATCH 22/23] fix(budget): gate the blocking cost-unknown arm on delivered tokens --- app/routes/chat.py | 13 ++- tests/integration/test_budget_enforcement.py | 96 ++++++++++++++++++++ 2 files changed, 105 insertions(+), 4 deletions(-) diff --git a/app/routes/chat.py b/app/routes/chat.py index 3473a39..72cad80 100644 --- a/app/routes/chat.py +++ b/app/routes/chat.py @@ -1362,9 +1362,13 @@ async def _commit_row(*, retry: bool) -> None: # request history accounts for. Usage with tokens but no price is the # same unknown (a custom upstream LiteLLM can't cost, or a model # absent from our catalog); a catalog-listed free model is known-zero - # and keeps its 0. Gated on having actually received a - # completion dict: a request that failed before the upstream answered - # (response == {}, e.g. the re-raised HTTPException above, whose + # and keeps its 0. As in the streaming arm, the unpriceable branch is + # gated on the usage having carried tokens: a usage of {0, 0} is an + # empty delivery, whose cost is known to be zero, and charging the whole + # remaining allowance for it would let an upstream that reports no tokens + # on a 200 permanently exhaust a capped key. Gated on having actually + # received a completion dict: a request that failed before the upstream + # answered (response == {}, e.g. the re-raised HTTPException above, whose # status_code never left 200) charges its recorded ~0 cost instead — # mirroring the cache-hit and pre-stream-failure paths. if ( @@ -1375,7 +1379,8 @@ async def _commit_row(*, retry: bool) -> None: and ( not response.get("usage") or ( - not (log.cost_microcents or 0) + (log.input_tokens or log.output_tokens) + and not (log.cost_microcents or 0) and not _has_known_price( litellm_cost_usd=(response.get("usage") or {}).get("cost_usd") or (response.get("_orca_meta") or {}).get("cost_usd"), diff --git a/tests/integration/test_budget_enforcement.py b/tests/integration/test_budget_enforcement.py index 8e5140c..09e19d2 100644 --- a/tests/integration/test_budget_enforcement.py +++ b/tests/integration/test_budget_enforcement.py @@ -1393,3 +1393,99 @@ async def test_budgeted_stream_commit_ack_loss_bills_the_delivery_once( assert cost > 0 assert await _get_spent(factory, key_id) == cost assert await pending_parked_spend(key_id) == 0 + + +async def _blocking_with_usage(budget_env, monkeypatch, *, usage: dict) -> tuple[int, int]: + """Drive a blocking request whose model has no price; return (key spend, row cost). + + `_lookup_priced_model` is the single oracle both the cost tier and + `_has_known_price` consult, so forcing it to miss makes the response + genuinely "usage with no price" — the shape a custom upstream produces — + without depending on which models the catalog happens to list today. + """ + from app.routes import chat + from sqlalchemy import select + + from packages.db.models.api_key import ApiKey + from packages.db.models.request_log import RequestLog + + monkeypatch.setattr(chat, "_lookup_priced_model", lambda model_id: None) + + make_client, fake, factory, _root = budget_env + key, key_id = await _make_budgeted_key(factory, budget_limit_cents=10) + + fake.acompletion = AsyncMock(return_value={ + "id": "chatcmpl-usage-shape", + "model": "gpt-4o-mini", + "object": "chat.completion", + "created": int(time.time()), + "choices": [{ + "index": 0, + "message": {"role": "assistant", "content": "ok"}, + "finish_reason": "stop", + }], + "usage": usage, + "_orca_meta": { + "provider": "openai", + "litellm_model": "openai/gpt-4o-mini", + "latency_ms": 10, + }, + }) + + async with await make_client(key) as c: + r = await c.post( + "/v1/chat/completions", + json={"model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hi"}]}, + ) + assert r.status_code == 200, r.text + + async with factory() as s: + spent = ( + await s.execute(select(ApiKey.spent_microcents).where(ApiKey.id == key_id)) + ).scalar_one() + cost = ( + await s.execute( + select(RequestLog.cost_microcents).where(RequestLog.api_key_id == key_id) + ) + ).scalar_one() + return spent, cost + + +async def test_budgeted_blocking_empty_usage_is_not_billed_the_cap(budget_env, monkeypatch): + """An empty delivery costs nothing, blocking or streaming. + + The blocking fail-closed arm is the streaming path's mirror but lost its token + guard, so any 200 with no price charged the key's ENTIRE remaining allowance + regardless of usage — a provider that reports zero tokens on an empty prompt, + or a custom upstream that leaves the field unset, exhausted a capped key with + one request and 429-blocked it for good. A usage frame that carried no tokens + is a measured zero; only tokens without a price are an unknown cost. + """ + spent, cost = await _blocking_with_usage( + budget_env, + monkeypatch, + usage={"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}, + ) + assert spent == 0 + assert cost == 0 + + +async def test_budgeted_blocking_unpriced_usage_with_tokens_bills_the_cap( + budget_env, monkeypatch, +): + """The token guard must not disarm the fail-closed arm standing beside it. + + Same unpriceable model, but the usage carried tokens: the cost is unknown and + a budgeted key must not get a delivered completion for free. This is the case + the arm exists for, so it is pinned next to the empty-delivery case — together + they fix the arm's meaning to exactly the streaming path's. + """ + spent, cost = await _blocking_with_usage( + budget_env, + monkeypatch, + usage={"prompt_tokens": 5, "completion_tokens": 2, "total_tokens": 7}, + ) + assert spent == 100_000 # the full remaining allowance (10 cents) + assert cost == 100_000 # recorded on the row it charges, not just the counter + From 6f8a2a47ebd03a4c5d2faac6e5edd9ae2bb5a930 Mon Sep 17 00:00:00 2001 From: Hasit Bhatt Date: Sat, 26 Sep 2026 19:01:45 -0700 Subject: [PATCH 23/23] style(test): order imports in the blocking settlement helper so ruff I001 passes --- tests/integration/test_budget_enforcement.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/integration/test_budget_enforcement.py b/tests/integration/test_budget_enforcement.py index 09e19d2..5eca319 100644 --- a/tests/integration/test_budget_enforcement.py +++ b/tests/integration/test_budget_enforcement.py @@ -1403,9 +1403,9 @@ async def _blocking_with_usage(budget_env, monkeypatch, *, usage: dict) -> tuple genuinely "usage with no price" — the shape a custom upstream produces — without depending on which models the catalog happens to list today. """ - from app.routes import chat from sqlalchemy import select + from app.routes import chat from packages.db.models.api_key import ApiKey from packages.db.models.request_log import RequestLog