test: add e2e tests for spend, budgets and llms #30790
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,93 @@ | ||
| # Budget Code Matrix | ||
|
Contributor
There was a problem hiding this comment.
Four Rule Used: Prevent documentation from being added - needs to ... (source)
Contributor
There was a problem hiding this comment. Temp until v1 to document the gaps |
||
|
|
||
| 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 | | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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/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. |
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Medium: Unscoped spend tag aggregation
view_spend_tagsonly requiresuser_api_key_authand does not pass the caller into this query, so an authenticated non-admin key can call/spend/tagsand get tag names, counts, and spend totals for every row inLiteLLM_SpendLogs. With the normalization in this change, this includes the normal double-encoded tag rows; either restrict this endpoint to admin/view-only admin roles or add the same user/team ownership predicates used by/spend/logsbefore aggregating.