diff --git a/app/routes/chat.py b/app/routes/chat.py index 1779d44..e1564df 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 budget_exceeded, get_lifetime_spend_microcents 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,34 @@ async def execute_chat( detail=f"Model '{body.model}' is not allowed for this API key", ) + # Budget enforcement: `budget_limit_cents` is a lifetime cap on this + # key's total spend (sum of `cost_microcents` for all non-deleted + # request-log rows, including streaming 499/503 rows that already + # incurred provider billing). Checked before any routing, resolution, + # or cache work so an exhausted key costs the operator nothing — no + # upstream attempt, no cache fill. + # + # NOTE: This is a best-effort soft limit, not a hard atomic cap. + # Spend is read before the request and the current request's cost is + # only written after the response/stream completes. N concurrent + # requests from the same budgeted key all observe the same + # pre-request total and may all pass the check, exceeding the cap by + # up to N× per-request cost in a burst. A hard cap would require a + # reservation/claim or row-level lock before dispatch; the current + # design trades strictness for simplicity and avoids holding a DB + # transaction across the upstream call. See spend.py for aggregation + # semantics. + if kc.budget_limit_cents is not None: + spend = await get_lifetime_spend_microcents(db, str(kc.key_id)) + if budget_exceeded(spend, kc.budget_limit_cents): + raise HTTPException( + status_code=429, + detail=( + "API key budget exhausted " + f"({spend} of {kc.budget_limit_cents * 10_000} microcents spent)." + ), + ) + 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" diff --git a/packages/auth/spend.py b/packages/auth/spend.py new file mode 100644 index 0000000..cb5cc07 --- /dev/null +++ b/packages/auth/spend.py @@ -0,0 +1,48 @@ +"""Per-key spend lookup used to enforce `ApiKey.budget_limit_cents`. + +Semantics: `budget_limit_cents` is a lifetime cap on the key's total +spend — sum of `cost_microcents` for all non-deleted request-log rows. +1 cent = 10,000 microcents (1 USD = 1,000,000 microcents, matching +chat.py's cost math). + +Rows are counted regardless of HTTP status because the streaming path +records billable token usage even when the final status is 499 (client +disconnect, chat.py:596) or 503 (mid-stream upstream failure, +chat.py:646); filtering on `status_code < 400` would exclude those and +make the cap bypassable by closing the stream early after reading the +usage chunk. + +Kept free of FastAPI imports so it stays unit-testable and reusable from +non-HTTP contexts (background jobs, CLI minting tools). +""" + +from __future__ import annotations + +from sqlalchemy import func, select +from sqlalchemy.ext.asyncio import AsyncSession + +from packages.db.models.request_log import RequestLog + +MICROCENTS_PER_CENT = 10_000 + + +async def get_lifetime_spend_microcents( + session: AsyncSession, api_key_id: str +) -> int: + """Sum of all non-deleted spend ever recorded for this key. + + Counts every row with `cost_microcents` regardless of `status_code` + so that streaming disconnect (499) and mid-stream upstream failure + (503) costs — which already incurred provider billing — are not + excluded from the budget. Failed requests with zero cost contribute + nothing to the sum regardless. + """ + stmt = select(func.coalesce(func.sum(RequestLog.cost_microcents), 0)).where( + RequestLog.api_key_id == api_key_id, + RequestLog.is_deleted == 0, + ) + return int((await session.execute(stmt)).scalar_one()) + + +def budget_exceeded(spend_microcents: int, budget_limit_cents: int) -> bool: + return spend_microcents >= budget_limit_cents * MICROCENTS_PER_CENT 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..cf75aed --- /dev/null +++ b/tests/integration/test_budget_enforcement.py @@ -0,0 +1,208 @@ +"""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.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, + )) + 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() diff --git a/tests/unit/test_budget_spend.py b/tests/unit/test_budget_spend.py new file mode 100644 index 0000000..2dfe25a --- /dev/null +++ b/tests/unit/test_budget_spend.py @@ -0,0 +1,106 @@ +"""Unit tests for packages.auth.spend — lifetime spend aggregation.""" + +import pytest + +from packages.auth.spend import ( + MICROCENTS_PER_CENT, + budget_exceeded, + get_lifetime_spend_microcents, +) + + +@pytest.fixture +async def seeded_log(db_session): + """Two keys with a mix of billable / streaming-failure / soft-deleted rows.""" + from packages.db.models.api_key import ApiKey + from packages.db.models.request_log import RequestLog + + k1 = ApiKey(workspace_id="default", name="a", key_hash="h-a", key_prefix="p-a") + k2 = ApiKey(workspace_id="default", name="b", key_hash="h-b", key_prefix="p-b") + db_session.add_all([k1, k2]) + await db_session.flush() + + rows = [ + RequestLog( + workspace_id="default", api_key_id=k1.id, model_requested="m", + model_resolved="m", provider="openai", input_tokens=1, output_tokens=1, + cost_microcents=1000, status_code=200, routing_strategy="balanced", latency_ms=10, trace_id="t-1", + ), + RequestLog( + workspace_id="default", api_key_id=k1.id, model_requested="m", + model_resolved="m", provider="openai", input_tokens=1, output_tokens=1, + cost_microcents=500, status_code=200, routing_strategy="balanced", latency_ms=10, trace_id="t-1", + ), + # Streaming failures ARE billable when they carry cost — the + # provider billed tokens even though the final status is 503 + # (mid-stream upstream failure) or 499 (client disconnect). + # Filtering on status_code < 400 would exclude these and make the + # budget bypassable, so they must be counted. + RequestLog( + workspace_id="default", api_key_id=k1.id, model_requested="m", + model_resolved="m", provider="openai", input_tokens=9, output_tokens=9, + cost_microcents=999_999, status_code=503, routing_strategy="balanced", latency_ms=10, trace_id="t-3", + ), + # soft-deleted rows must never count, even with non-zero cost + RequestLog( + workspace_id="default", api_key_id=k1.id, model_requested="m", + model_resolved="m", provider="openai", input_tokens=1, output_tokens=1, + cost_microcents=12345, status_code=200, routing_strategy="balanced", latency_ms=10, trace_id="t-5", + is_deleted=1, + ), + # another key's spend must not leak in + RequestLog( + workspace_id="default", api_key_id=k2.id, model_requested="m", + model_resolved="m", provider="openai", input_tokens=2, output_tokens=2, + cost_microcents=777_777, status_code=200, routing_strategy="balanced", latency_ms=10, trace_id="t-4", + ), + ] + db_session.add_all(rows) + await db_session.commit() + return k1, k2 + + +async def test_spend_sums_only_billable_rows_for_the_key(db_session, seeded_log): + k1, _k2 = seeded_log + spend = await get_lifetime_spend_microcents(db_session, k1.id) + # 1000 + 500 + 999_999 (503 failure with cost is now counted) = 1,001,499; + # soft-deleted 12345 is excluded. + assert spend == 1_001_499 + + +async def test_spend_counts_stream_disconnect_and_mid_stream_failure(db_session, seeded_log): + """Regression for P1: 499/503 streaming costs must count toward budget.""" + from packages.db.models.request_log import RequestLog + + k1, _k2 = seeded_log + # Add explicit 499 disconnect row with cost + db_session.add( + RequestLog( + workspace_id="default", api_key_id=k1.id, model_requested="m", + model_resolved="m", provider="openai", input_tokens=2, output_tokens=2, + cost_microcents=42_000, status_code=499, routing_strategy="balanced", latency_ms=10, trace_id="t-6", + ) + ) + await db_session.commit() + spend = await get_lifetime_spend_microcents(db_session, k1.id) + assert spend == 1_001_499 + 42_000 + + +async def test_empty_history_is_zero(db_session, seeded_log): + _k1, k2 = seeded_log + from packages.db.models.api_key import ApiKey + + fresh = ApiKey(workspace_id="default", name="c", key_hash="h-c", key_prefix="p-c") + db_session.add(fresh) + await db_session.commit() + assert await get_lifetime_spend_microcents(db_session, fresh.id) == 0 + + +def test_budget_exceeded_boundary(): + assert budget_exceeded(10_000 - 1, 1) is False # just under 1 cent + assert budget_exceeded(10_000, 1) is True # exactly at the cap blocks + assert budget_exceeded(0, 1) is False + + +def test_microcent_conversion_constant(): + assert MICROCENTS_PER_CENT == 10_000