test: add e2e tests for spend, budgets and llms

tests/e2e_tests/budgets/BUDGET_CODE_MATRIX.md

@@ -0,0 +1,93 @@
+# Budget Code Matrix
+
+What LiteLLM actually implements for budgets: every entity that can carry a dollar
+budget, how the limit is enforced, and where in the code it happens. This is the
+"what we support" reference; the companion `BUDGET_TEST_COVERAGE_MATRIX.md` maps
+each row to its tests and the e2e gaps.
+
+Over-budget surfaces as a `budget_exceeded` error (the live suite
+`tests/otel_tests/test_e2e_budgeting.py` asserts `type == "budget_exceeded"`,
+`code == "429"`); the underlying `BudgetExceededError` is defined in
+`litellm/exceptions.py` (`status_code=400`). Enforcement runs in `common_checks()`
+/ `auth_checks.py` at auth time, plus pre-call reservation in
+`budget_reservation.py`.
+
+Legend for "Enforced": **block** = request rejected; **filter** = router skips the
+deployment; **alert** = notify only, request proceeds.
+
+---
+
+## 1. Per-entity dollar budgets
+
+| Entity | Budget stored | Hard `max_budget` | Soft budget | Per-window | Model budget | Reset by `budget_duration` |
+|--------|---------------|-------------------|-------------|------------|--------------|----------------------------|
+| API key | `LiteLLM_VerificationToken` (direct cols + `budget_id` FK) | block (`_virtual_key_max_budget_check`) | alert (`_virtual_key_soft_budget_check`) + 80% alert | block (`_virtual_key_multi_budget_check`) | block (`model_max_budget_limiter.is_key_within_model_budget`) | keys reset job |
+| Internal user | `LiteLLM_UserTable` (direct cols) | block (`common_checks`, only when not on a team) | - | - | via `model_max_budget` json | users reset job |
+| Team | `LiteLLM_TeamTable` (direct cols) | block (`_team_max_budget_check`) | alert (`_team_soft_budget_check`) | block (`_team_multi_budget_check`) | via `model_max_budget` | teams reset job |
+| Team member | `LiteLLM_TeamMembership` -> `LiteLLM_BudgetTable` | block (`_check_team_member_budget`) | - | - | - | budget-table reset job |
+| End-user / customer | `LiteLLM_EndUserTable` -> `LiteLLM_BudgetTable` | block (`_check_end_user_budget`) | - | - | block (`is_end_user_within_model_budget`) | budget-table reset job |
+| Organization | `LiteLLM_OrganizationTable` -> `LiteLLM_BudgetTable` | block (`_organization_max_budget_check`) | - | - | via budget-table | budget-table reset job |
+| Tag | `LiteLLM_TagTable` -> `LiteLLM_BudgetTable` | block (`_tag_max_budget_check`) | - | - | via budget-table | budget-table reset job |
+| Project | `LiteLLM_ProjectTable` -> `LiteLLM_BudgetTable` | block (`_project_max_budget_check`) | alert (`_project_soft_budget_check`) | - | - | budget-table reset job |
+| Provider (router) | config `provider_budget_config` (in-memory) | filter (`router_strategy/budget_limiter`) | - | yes (time window) | - | window TTL |
+| Global proxy | `litellm.max_budget` (config) | block (`_global_proxy_budget_check`) | - | - | - | - |
+
+Notes / flags from the code:
+- **User budget only enforced off-team**: `common_checks` skips the personal-user
+  budget when the key belongs to a team (team budget governs instead).
+- **Comparison operators are inconsistent**: key/user use `>=`, team/end-user main
+  budget use `>`. Spend exactly at `max_budget` blocks a key but not a team.
+- **Provider budgets are filter-only**: an over-budget provider is removed from
+  routing; if all are over budget the router raises
+  `no_deployments_with_provider_budget_routing` (not a per-entity block).
+- **Enforcement timing differs by entity**: key / user / org / team-member / tag /
+  model enforce off real-time reservation counters (block within ~2 calls);
+  **end-user** enforcement reads `EndUserTable.spend`, which only updates on the
+  `proxy_batch_write_at` flush, so it lags by that interval (verified live).
+
+## 2. Budget mechanisms
+
+| Mechanism | What it does | Code |
+|-----------|--------------|------|
+| Pre-call reservation | Estimates max request cost, atomically reserves against redis spend counters for key/team/user/end_user/tag/team_member/org before the call; blocks if a counter would exceed | `spend_tracking/budget_reservation.py` |
+| Post-call reconciliation | Adjusts the reservation to the actual cost once known | `reconcile_budget_reservation` |
+| Read-time enforcement | Auth-time check of current spend vs `max_budget` | `auth_checks.common_checks` + per-entity `_*_max_budget_check` |
+| Soft budget / alerts | At `soft_budget` (or 80% of max) fire Slack/email alert, do not block | `_virtual_key_soft_budget_check`, `_team_soft_budget_check`, `budget_alerts` |
+| Multi-window budgets | `budget_limits` list of `{budget_duration, max_budget}`; each window enforced + reset independently | `_virtual_key_multi_budget_check`, `reset_budget_windows` |
+| Model-level budgets | `model_max_budget` dict (per model: `budget_limit` + `time_period`) on key/user/end_user | `hooks/model_max_budget_limiter.py` |
+| Reset by duration | Job zeros `spend`, recomputes `budget_reset_at = now + duration_in_seconds(budget_duration)`, invalidates redis counters | `common_utils/reset_budget_job.py`, `duration_parser.duration_in_seconds` |
+| Zero-cost bypass | Models with no configured price bypass budget reservation | `budget_reservation` zero-cost path |
+
+## 3. Budget management surface (endpoints)
+
+| Action | Endpoint | Handler |
+|--------|----------|---------|
+| Create budget | `POST /budget/new` | `new_budget` |
+| Update budget | `POST /budget/update` | `update_budget` |
+| Budget info | `POST /budget/info` (`{"budgets": [id]}`) | `info_budget` |
+| Budget settings | `GET /budget/settings` | `budget_settings` |
+| List budgets | `GET /budget/list` | `list_budget` |
+| Delete budget | `POST /budget/delete` (`{"id": id}`) | `delete_budget` |
+| Set on key | `POST /key/generate`, `/key/update` (`max_budget`, `soft_budget`, `budget_duration`, `model_max_budget`, `budget_id`) | key mgmt |
+| Set on user | `POST /user/new` (`max_budget`, `budget_duration`) | internal user |
+| Set on team | `POST /team/new` (`max_budget`, `soft_budget`, `team_member_budget`) | team |
+| Set on team member | `POST /team/member_add` (`max_budget_in_team`) | team |
+| Set on org | `POST /organization/new` (`max_budget`, `soft_budget`, `model_max_budget`) | org |
+| Set on customer | `POST /customer/new`, `/customer/update` (`max_budget`, `budget_id`) | customer |
+| Set on tag | `POST /tag/new`, `/tag/update` (`max_budget`) | tag mgmt |
+| Read budget+spend | `/key/info`, `/user/info`, `/team/info`, `/organization/info`, `/customer/info`, `/budget/info` | per-entity info |
+
+Endpoint method/shape gotchas verified live: `/organization/delete` is **DELETE**
+with `{"organization_ids": [id]}`; `/budget/info` takes `{"budgets": [id]}`;
+`model_max_budget` entries use `{"budget_limit", "time_period"}`.
+
+## 4. Config knobs
+
+| Setting | Effect |
+|---------|--------|
+| `litellm.max_budget` | proxy-wide hard cap (global proxy budget) |
+| `max_internal_user_budget` / `default_max_internal_user_budget` | default `max_budget` for internal users |
+| `internal_user_budget_duration` | default reset duration for internal users |
+| `max_end_user_budget` / `max_end_user_budget_id` | default budget for end-users |
+| `default_team_params` | default `max_budget` / `budget_duration` / limits for teams |
+| `provider_budget_config` (router) | per-provider spend caps + windows |

tests/e2e_tests/budgets/BUDGET_TEST_COVERAGE_MATRIX.md

@@ -0,0 +1,78 @@
+# Budget Test Coverage Matrix
+
+Maps every row of `BUDGET_CODE_MATRIX.md` (what LiteLLM implements) to its tests
+and level, then marks the live e2e coverage this suite adds.
+
+Levels: `unit` mocked (`AsyncMock` on `get_current_spend`/prisma); `router` live
+router with fake deployments; `live-e2e` real proxy, real key/team, real requests
+until blocked. Status: `covered` / `partial` / `gap`.
+
+Pre-existing live coverage outside this suite:
+- `tests/otel_tests/test_e2e_budgeting.py` - key + team enforcement, budget update.
+- `tests/local_testing/test_router_budget_limiter.py` - provider / tag / deployment
+  budgets at the router.
+
+This suite (`tests/e2e_tests/budgets/`) adds the missing live coverage and runs
+on the shared lifecycle (every entity it creates is deleted on teardown).
+
+---
+
+## Per-entity enforcement
+
+| Entity | Unit | Pre-existing live | This suite (live) | Status |
+|--------|------|-------------------|-------------------|--------|
+| API key | `test_budget_reservation.py`, `test_max_budget_limiter.py` | `otel_tests` | `test_budget_enforcement_e2e::test_key_budget_blocks` | **covered** |
+| Team | `test_team_budget_limits.py` | `otel_tests` | (org test builds a team) | **covered** |
+| Internal user | auth unit tests | - | `test_internal_user_budget_blocks` | **covered (new)** |
+| Team member | `test_team_member_budget.py` | - | `test_team_member_budget_blocks` | **covered (new)** |
+| End-user / customer | `test_custom_auth_end_user_budget.py` | - | `test_end_user_budget_blocks` | **covered (new)** |
+| Organization | `test_organization_budget_enforcement.py` (flagged weak) | - | `test_organization_budget_blocks` | **covered (new)** |
+| Tag (proxy-level) | - | router only | `test_tag_budget_e2e::test_tag_budget_blocks_tagged_requests` | **covered (new)** |
+| Model-level (`model_max_budget`) | `test_unit_test_max_model_budget_limiter.py` | - | `test_model_max_budget_e2e::test_model_max_budget_isolates_per_model` | **covered (new)** |
+| Provider (router) | `test_budget_limiter_hotpath.py` | `test_router_budget_limiter.py` | - | **covered** (router) |
+| Global proxy (`litellm.max_budget`) | unit | - | - | **gap** (needs a config-level cap; not key-settable) |
+
+## Budget mechanisms
+
+| Mechanism | Unit | This suite (live) | Status |
+|-----------|------|-------------------|--------|
+| Pre-call reservation | `test_budget_reservation.py` | exercised by every enforcement test | **partial** |
+| Soft budget / alerts | `SlackAlerting/test_budget_alert_types.py` | `test_soft_budget_e2e::test_soft_budget_does_not_block` | **covered (new)** (block-vs-alert; the alert side-effect itself stays unit) |
+| Budget CRUD | `test_budget_endpoints.py` | `test_budget_crud_e2e` (roundtrip + delete) | **covered (new)** |
+| Reset scheduling | `test_proxy_budget_reset.py` | `test_budget_crud_e2e::test_budget_duration_schedules_reset_on_key` | **covered (new)** (scheduling; actual zeroing is time-dependent -> unit) |
+| Multi-window budgets | `test_multi_budget_windows.py` | - | **gap** (window setup is fiddly; left to unit for now) |
+| Read budget+spend | `test_spend_management_endpoints.py` | `/key/info` asserted in CRUD + enforcement | **partial** |
+
+## Remaining gaps (intentionally not live-tested)
+
+- **Global proxy budget** (`litellm.max_budget`): set via proxy config, not a
+  per-key API, so it needs a dedicated proxy boot with that config rather than a
+  runtime-created entity. Out of scope for the per-entity suite.
+- **Multi-window budgets**: the `budget_limits` list shape and per-window reset are
+  covered by `test_multi_budget_windows.py` (unit); a live version would need to
+  wait out a short window to see the reset, which is time-dependent.
+- **Soft-budget alert delivery**: whether the Slack/email actually fires is not
+  observable from the proxy API; unit tests own that. The live test pins the
+  load-bearing behavior (soft does not block).
+- **Reset zeroing after the window elapses**: time-dependent; unit tests own the
+  reset-job logic. The live test pins that `budget_reset_at` is scheduled.
+
+## This suite's files
+
+| File | Covers |
+|------|--------|
+| `test_budget_enforcement_e2e.py` | key / internal-user / end-user / organization / team-member hard enforcement |
+| `test_model_max_budget_e2e.py` | per-model caps isolate by model |
+| `test_soft_budget_e2e.py` | soft budget alerts but does not block |
+| `test_tag_budget_e2e.py` | proxy-level tag budget blocks tagged requests, spares others |
+| `test_budget_crud_e2e.py` | `/budget/*` CRUD roundtrip + delete + `budget_reset_at` scheduling |
+
+## Pattern + timing
+
+Create the entity with a tiny `max_budget`, drive spend until a `budget_exceeded`
+block. The enforcement helper is two-phase: a fast warmup (key/user/org/member/tag/
+model block within ~2 calls off real-time counters), then a poll across the ~60s
+batch-write window (end-user enforcement reads table spend that lags). Skip on a
+non-budget error (provider down / key missing); fail if the budget is never
+enforced. Chat tests use `gpt-5.5` (the model with a working key on the reference
+proxy); swap the literal if your proxy differs.

tests/e2e_tests/budgets/budget_client.py

@@ -0,0 +1,185 @@
+"""Client for budget e2e tests: the shared ProxyClient plus budget-bearing entity
+management (user / team / team-member / org / customer / tag / budget-table) and
+info reads.
+
+Over-budget surfaces as a ``budget_exceeded`` error; ``is_budget_block`` detects it
+on a CallResult. Create methods return the new id and raise on failure; tests
+register the matching delete with ``resources.defer(...)`` for cleanup.
+"""
+
+from typing import Dict, List, Optional
+
+import requests
+
+from proxy_client import CallResult, ProxyClient, _auth, proxy_client_kwargs
+
+
+def is_budget_block(result: CallResult) -> bool:
+    """True if the call was rejected for being over budget (vs a provider error)."""
+    return not result.ok and "budget_exceeded" in result.body
+
+
+def model_budget(model: str, limit: float, period: str = "30d") -> dict:
+    """A model_max_budget dict entry: per-model cap with a reset window."""
+    return {model: {"budget_limit": limit, "time_period": period}}
+
+
+class BudgetClient(ProxyClient):
+    def _post(self, path: str, body: Dict[str, object]) -> Dict[str, object]:
+        resp = requests.post(
+            f"{self._base_url}{path}",
+            headers=_auth(self._master_key),
+            json=body,
+            timeout=self._request_timeout,
+        )
+        resp.raise_for_status()
+        return dict(resp.json()) if resp.text else {}
+
+    def _get(self, path: str, params: Dict[str, str]) -> Dict[str, object]:
+        resp = requests.get(
+            f"{self._base_url}{path}",
+            headers=_auth(self._master_key),
+            params=params,
+            timeout=self._request_timeout,
+        )
+        resp.raise_for_status()
+        return dict(resp.json())
+
+    def _delete(self, path: str, body: Dict[str, object]) -> None:
+        try:
+            requests.delete(
+                f"{self._base_url}{path}",
+                headers=_auth(self._master_key),
+                json=body,
+                timeout=self._request_timeout,
+            )
+        except requests.RequestException:
+            pass
+
+    def _post_quiet(self, path: str, body: Dict[str, object]) -> None:
+        try:
+            requests.post(
+                f"{self._base_url}{path}",
+                headers=_auth(self._master_key),
+                json=body,
+                timeout=self._request_timeout,
+            )
+        except requests.RequestException:
+            pass
+
+    # ---- internal user --------------------------------------------------
+
+    def create_user(
+        self, *, max_budget: float, budget_duration: Optional[str] = None
+    ) -> str:
+        body: Dict[str, object] = {"max_budget": max_budget}
+        if budget_duration is not None:
+            body["budget_duration"] = budget_duration
+        return str(self._post("/user/new", body)["user_id"])
+
+    def delete_user(self, user_id: str) -> None:
+        self._post_quiet("/user/delete", {"user_ids": [user_id]})
+
+    def user_info(self, user_id: str) -> Dict[str, object]:
+        return self._get("/user/info", {"user_id": user_id})
+
+    # ---- customer / end-user -------------------------------------------
+
+    def create_customer(self, customer_id: str, *, max_budget: float) -> str:
+        self._post("/customer/new", {"user_id": customer_id, "max_budget": max_budget})
+        return customer_id
+
+    def customer_info(self, customer_id: str) -> Dict[str, object]:
+        return self._get("/customer/info", {"end_user_id": customer_id})
+
+    # ---- organization ---------------------------------------------------
+
+    def create_org(self, *, max_budget: float, alias: str) -> str:
+        return str(
+            self._post(
+                "/organization/new",
+                {"organization_alias": alias, "max_budget": max_budget},
+            )["organization_id"]
+        )
+
+    def delete_org(self, org_id: str) -> None:
+        self._delete("/organization/delete", {"organization_ids": [org_id]})
+
+    # ---- team -----------------------------------------------------------
+
+    def create_team(
+        self,
+        *,
+        alias: str,
+        max_budget: Optional[float] = None,
+        organization_id: Optional[str] = None,
+        extra: Optional[Dict[str, object]] = None,
+    ) -> str:
+        body: Dict[str, object] = {"team_alias": alias}
+        if max_budget is not None:
+            body["max_budget"] = max_budget
+        if organization_id is not None:
+            body["organization_id"] = organization_id
+        if extra:
+            body.update(extra)
+        return str(self._post("/team/new", body)["team_id"])
+
+    def delete_team(self, team_id: str) -> None:
+        self._post_quiet("/team/delete", {"team_ids": [team_id]})
+
+    def team_info(self, team_id: str) -> Dict[str, object]:
+        return self._get("/team/info", {"team_id": team_id})
+
+    def add_team_member(
+        self, team_id: str, user_id: str, *, max_budget_in_team: Optional[float] = None
+    ) -> None:
+        body: Dict[str, object] = {
+            "team_id": team_id,
+            "member": {"role": "user", "user_id": user_id},
+        }
+        if max_budget_in_team is not None:
+            body["max_budget_in_team"] = max_budget_in_team
+        self._post("/team/member_add", body)
+
+    # ---- tag ------------------------------------------------------------
+
+    def create_tag(self, name: str, *, max_budget: float) -> str:
+        self._post("/tag/new", {"name": name, "max_budget": max_budget})
+        return name
+
+    def delete_tag(self, name: str) -> None:
+        self._post_quiet("/tag/delete", {"name": name})
+
+    # ---- budget table ---------------------------------------------------
+
+    def create_budget(
+        self,
+        *,
+        max_budget: float,
+        soft_budget: Optional[float] = None,
+        budget_duration: Optional[str] = None,
+    ) -> str:
+        body: Dict[str, object] = {"max_budget": max_budget}
+        if soft_budget is not None:
+            body["soft_budget"] = soft_budget
+        if budget_duration is not None:
+            body["budget_duration"] = budget_duration
+        return str(self._post("/budget/new", body)["budget_id"])
+
+    def delete_budget(self, budget_id: str) -> None:
+        self._post_quiet("/budget/delete", {"id": budget_id})
+
+    def budget_info(self, budget_id: str) -> List[Dict[str, object]]:
+        resp = requests.post(
+            f"{self._base_url}/budget/info",
+            headers=_auth(self._master_key),
+            json={"budgets": [budget_id]},
+            timeout=self._request_timeout,
+        )
+        resp.raise_for_status()
+        data = resp.json()
+        return [dict(row) for row in data] if isinstance(data, list) else []
+
+
+def build_client() -> BudgetClient:
+    return BudgetClient(**proxy_client_kwargs())

tests/e2e_tests/budgets/conftest.py

@@ -0,0 +1,16 @@
+"""Budgets suite's `client` fixture.
+
+The shared lifecycle (resources/scoped_key), proxy liveness skip, and e2e marker
+live in the parent tests/e2e_tests/conftest.py. BudgetClient subclasses
+ProxyClient, so it satisfies lifecycle.ResourceClient and the shared `resources`
+fixture cleans up keys; tests register entity deletes via `resources.defer(...)`.
+"""
+
+import pytest
+
+from budget_client import BudgetClient, build_client
+
+
+@pytest.fixture(scope="session")
+def client() -> BudgetClient:
+    return build_client()