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/app/routes/chat.py b/app/routes/chat.py index 1779d44..9b16545 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 @@ -306,6 +307,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" @@ -420,6 +436,24 @@ async def execute_chat( resolved_model = candidates[0] body.model = candidates[0] # mutate for downstream completion call + # Budget enforcement: `budget_limit_cents` is a hard lifetime cap. The check + # runs only after the request has passed every pre-dispatch validation (model + # allowlist, provider deployability), so a request we reject before touching + # an upstream never consumes budget. The real cost is only known once the + # upstream response/stream completes, so we record it atomically in + # `_settle_budget` — the `UPDATE spent = spent + actual WHERE spent + actual + # <= cap` guard makes this safe under concurrency and never lets the counter + # exceed the cap (fail-closed, never over-recorded). + if kc.budget_limit_cents is not None: + cap = kc.budget_limit_cents * MICROCENTS_PER_CENT + if await is_exhausted(db, str(kc.key_id), cap): + raise HTTPException( + status_code=429, + detail=f"API key budget exhausted ({cap} microcents lifetime cap reached).", + ) + kc._budget_cap = cap + kc._budget_spent = await read_spent(db, str(kc.key_id)) + started_perf = time.perf_counter() completion_kwargs = body.model_dump(exclude_none=True) @@ -489,6 +523,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)) @@ -509,16 +544,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: @@ -542,6 +575,7 @@ async def _log_pre_stream_failure(status: int, err_type: str | None) -> None: ) db.add(log) try: + await _settle_budget(db, 0, commit=False) await db.commit() except Exception as commit_err: logger.warning("request_log_commit_failed", error=str(commit_err)) @@ -581,6 +615,17 @@ async def sse() -> AsyncGenerator[str, None]: status_code = 200 error_type: str | None = None log_written = False + # True only once a terminal `data: [DONE]` has been emitted, i.e. the + # response was delivered in full. While False, the stream ended early + # (client disconnect / mid-stream upstream error) and the real cost is + # unknown, so the budget claim must be kept (fail-closed) rather than + # released — otherwise a client could stream tokens then hang up before + # the usage frame to bypass the cap. + stream_completed = False + # True once any usage frame has been observed in the stream. A completed + # stream with no usage frame means cost is unknown (client suppressed it + # or the provider omitted it), so the cap must still be enforced. + usage_seen = False async def _finalize() -> None: """Write the request log row exactly once. @@ -659,27 +704,46 @@ 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. + + When the real cost is unknown — the stream ended without a + terminal [DONE], or a completed stream never delivered a usage + frame (e.g. a client forced include_usage=False or a provider + omitted usage) — charge the full remaining allowance so a client + cannot suppress the usage frame to bypass the cap. + """ + actual = row_values.get("cost_microcents") or 0 + cost_unknown = (not stream_completed) or (not usage_seen) + if cost_unknown: + 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: @@ -690,9 +754,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: @@ -770,10 +835,12 @@ async def _commit_row(*, retry: bool) -> None: agg_latency = meta.get("latency_ms", agg_latency) if "usage" in d and d["usage"]: agg_usage = d["usage"] + usage_seen = True if d.get("model"): agg_model = d["model"] yield f"data: {json.dumps(d, separators=(',', ':'))}\n\n" yield "data: [DONE]\n\n" + stream_completed = True except (asyncio.CancelledError, GeneratorExit): # Client closed the connection (Ctrl+C, tab closed, browser # navigated away, proxy timeout, ...). Two distinct signals @@ -874,6 +941,12 @@ async def _commit_row(*, retry: bool) -> None: # is legal; clients reading until [DONE] still get it after # an upstream error. yield "data: [DONE]\n\n" + # The error response was delivered in full (terminal [DONE] sent), + # so settle against the actual cost only — not the full remaining + # allowance. Without this, every mid-stream provider failure would + # charge (and exhaust) the key's entire remaining budget even + # though the delivered response cost ~0. + stream_completed = True finally: # Same shielding reason as the cancel branch: ensure the # log write actually completes before we unwind, even if @@ -909,6 +982,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, @@ -950,11 +1029,35 @@ 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 + + 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(log) + await _settle_budget(db, log.cost_microcents, 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, + ) if isinstance(response, dict) and "_orca_meta" in response: response = {k: v for k, v in response.items() if k != "_orca_meta"} diff --git a/packages/auth/spend.py b/packages/auth/spend.py new file mode 100644 index 0000000..e459308 --- /dev/null +++ b/packages/auth/spend.py @@ -0,0 +1,90 @@ +"""Per-key lifetime spend tracking that enforces ``ApiKey.budget_limit_cents``. + +The cap is a hard lifetime limit on the key's total spend, in microcents +(1 cent = 10_000 microcents; 1 USD = 1_000_000 microcents, matching chat.py's +cost math). + +Actual cost is only known after the upstream call returns, so enforcement is a +single atomic ``UPDATE`` that adds the real cost and refuses to let the counter +exceed the cap:: + + UPDATE api_keys SET spent_microcents = spent_microcents + :actual + WHERE id = :id AND spent_microcents + :actual <= :cap + +Concurrent requests for the same key each add their own cost atomically; only a +request whose *own* cost alone would breach the remaining budget matches zero +rows. In that case the counter is clamped to ``cap`` so the key is correctly +maxed out and the next request is rejected — fail-closed, never over-recorded. + +This avoids both failure modes of a pre-claim design: it never records spend +past the cap (no over-spend), and it does not reserve the whole remaining budget +up front (so a key's requests are not serialized behind a single in-flight one). + +Kept free of FastAPI imports so it stays unit-testable and reusable from +non-HTTP paths (background jobs, CLI minting tools). +""" + +from __future__ import annotations + +from sqlalchemy import select, update +from sqlalchemy.ext.asyncio import AsyncSession + +from packages.db.models.api_key import ApiKey + +MICROCENTS_PER_CENT = 10_000 + + +async def read_spent(db: AsyncSession, api_key_id: str) -> int: + """Return the key's currently-recorded lifetime spend in microcents.""" + spent = ( + await db.execute(select(ApiKey.spent_microcents).where(ApiKey.id == api_key_id)) + ).scalar_one_or_none() + return int(spent or 0) + + +async def is_exhausted(db: AsyncSession, api_key_id: str, cap_microcents: int) -> bool: + """Fast pre-check: has the key already reached its lifetime cap?""" + return (await read_spent(db, api_key_id)) >= cap_microcents + + +async def charge_budget( + db: AsyncSession, + api_key_id: str, + cap_microcents: int, + actual_microcents: int, + *, + commit: bool = True, +) -> bool: + """Atomically record ``actual_microcents`` of spend, never exceeding ``cap``. + + Returns ``True`` if the cost fit under the cap (the counter advanced by + ``actual``), or ``False`` if the request alone would have breached the cap — + in which case the counter is clamped to ``cap`` so the key is maxed out and + blocked going forward. The boundary request may already have been served + upstream; it cannot be un-spent, but we never record more than the cap and we + stop the next one. Fail-closed. + + When ``commit`` is False the UPDATEs are executed but not committed, so the + caller can commit them in the same transaction as the request-log write + (atomic log + charge — no window where the log lands but the charge is lost). + """ + actual = actual_microcents or 0 + result = await db.execute( + update(ApiKey) + .where(ApiKey.id == api_key_id, ApiKey.spent_microcents + actual <= cap_microcents) + .values(spent_microcents=ApiKey.spent_microcents + actual) + ) + if result.rowcount: + if commit: + await db.commit() + return True + # Would have exceeded the cap: clamp so the counter never overshoots and the + # key is correctly reported as exhausted thereafter. + await db.execute( + update(ApiKey) + .where(ApiKey.id == api_key_id, ApiKey.spent_microcents < cap_microcents) + .values(spent_microcents=cap_microcents) + ) + if commit: + await db.commit() + return False diff --git a/packages/db/migrate.py b/packages/db/migrate.py new file mode 100644 index 0000000..6fbb452 --- /dev/null +++ b/packages/db/migrate.py @@ -0,0 +1,54 @@ +"""Idempotent startup schema migrations for columns added after the first release. + +`Base.metadata.create_all` creates new tables but never alters existing ones, so a +deployment that already ran a release (a SQLite named volume, a fly.io/Postgres +volume) keeps an `api_keys` table without the `spent_microcents` column. After an +upgrade the ORM would then `SELECT` every mapped column and hit "no such column" +on every authenticated request — a 503 for the whole API. + +`ensure_budget_columns` is run once at boot, after `create_all`, and is safe to +call on every start: it inspects the live schema and only acts when the column is +missing. +""" + +from __future__ import annotations + +from sqlalchemy import inspect, text + + +async def ensure_budget_columns(engine) -> None: + """Add `spent_microcents` to `api_keys` if absent, seeded from request history. + + Also widens `budget_limit_cents` to BIGINT on Postgres (the microcent scale + can exceed int4). Both are no-ops on a fresh database. + """ + async with engine.begin() as conn: + cols = { + c["name"] + for c in await conn.run_sync(lambda sync: inspect(sync).get_columns("api_keys")) + } + is_postgres = engine.dialect.name == "postgresql" + + if "spent_microcents" not in cols: + await conn.execute( + text( + "ALTER TABLE api_keys ADD COLUMN spent_microcents BIGINT " + "NOT NULL DEFAULT 0" + ) + ) + # Seed lifetime spend from historical request logs so an existing key's + # cap is not silently reset to zero (which would re-grant a leaked key + # a full new budget). + await conn.execute( + text( + "UPDATE api_keys SET spent_microcents = (" + " SELECT COALESCE(SUM(cost_microcents), 0) FROM requests_log " + " WHERE requests_log.api_key_id = api_keys.id" + ") WHERE spent_microcents = 0" + ) + ) + + if is_postgres and "budget_limit_cents" in cols: + await conn.execute( + text("ALTER TABLE api_keys ALTER COLUMN budget_limit_cents TYPE BIGINT") + ) 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/integration/test_budget_enforcement.py b/tests/integration/test_budget_enforcement.py new file mode 100644 index 0000000..5fa33b4 --- /dev/null +++ b/tests/integration/test_budget_enforcement.py @@ -0,0 +1,305 @@ +"""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 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 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 diff --git a/tests/unit/test_startup_guards.py b/tests/unit/test_startup_guards.py index 736d144..0c8bd38 100644 --- a/tests/unit/test_startup_guards.py +++ b/tests/unit/test_startup_guards.py @@ -128,9 +128,10 @@ async def _boom(*_a, **_kw): async def test_guard_allows_missing_table_on_fresh_sqlite(tmp_sqlite_url): """A fresh DB pre-migration where provider_keys doesn't exist counts as zero rows for sqlite (inspected via engine, not string matching).""" + from sqlalchemy.ext.asyncio import async_sessionmaker + from packages.db.engine import build_engine from packages.db.guards import assert_credential_encryption_ready - from sqlalchemy.ext.asyncio import async_sessionmaker engine = build_engine(tmp_sqlite_url) factory = async_sessionmaker(engine, expire_on_commit=False)