feat(budget): enforce ApiKey.budget_limit_cents on the chat path

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"

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

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)

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()