Merge remote-tracking branch 'origin/litellm_/management-v1-budgets' into worktree-floofy-watching-dijkstra

This commit is contained in:
Yuneng Jiang 2026-07-31 11:35:19 -07:00
commit 8e60c2e725
No known key found for this signature in database
10 changed files with 1075 additions and 10 deletions

View file

@ -841,6 +841,7 @@ class LiteLLMRoutes(enum.Enum):
"/config/list",
"/config/field/info",
"/budget/list",
"/management/v1/budgets",
"/budget/settings",
# Invitation viewing (admin viewer cannot create/delete; can read).
"/invitation/info",

View file

@ -2,11 +2,15 @@
from fastapi import APIRouter
from litellm.proxy.management_endpoints.management_v1.budgets import (
router as budgets_router,
)
from litellm.proxy.management_endpoints.management_v1.spend_logs import (
router as spend_logs_router,
)
router = APIRouter()
router.include_router(budgets_router)
router.include_router(spend_logs_router)
__all__ = ["router"]

View file

@ -0,0 +1,205 @@
"""`GET /management/v1/budgets`."""
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from datetime import datetime
from types import MappingProxyType
from typing import Annotated
from fastapi import APIRouter, Depends, Request
from pydantic import BaseModel, TypeAdapter
from litellm._logging import verbose_proxy_logger
from litellm.proxy._types import (
CommonProxyErrors,
UserAPIKeyAuth,
user_api_key_has_admin_view,
)
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.management_endpoints.management_v1.common import (
MANAGEMENT_V1_PREFIX,
PROBLEM_TYPE_BASE,
ManagementProblem,
)
from litellm.proxy.management_endpoints.management_v1.list_framework import (
FilterSpec,
ListSpec,
Predicate,
QueryPlan,
Scope,
ScopeAll,
ScopeDenied,
SortKey,
handle_list,
order_by_sql,
where_sql,
)
from litellm.proxy.utils import PrismaClient
from litellm.types.proxy.management_endpoints.management_v1 import (
ListResponse,
ProblemDetail,
)
router = APIRouter(prefix=MANAGEMENT_V1_PREFIX)
BUDGET_TABLE = '"LiteLLM_BudgetTable"'
class BudgetListItem(BaseModel):
"""One budget as the Budgets page reads it, and as it comes back off the table.
Validating the raw row through here is what makes `tpm_limit` / `rpm_limit`
numbers: they are `BigInt?` in the schema, which the query engine hands back as
decimal strings, and a quoted "60000" breaks arithmetic in the dashboard.
"""
budget_id: str
max_budget: float | None = None
soft_budget: float | None = None
tpm_limit: int | None = None
rpm_limit: int | None = None
budget_duration: str | None = None
budget_reset_at: datetime | None = None
created_at: datetime
updated_at: datetime
class _RowCount(BaseModel):
count: int
_BUDGET_ROWS = TypeAdapter(tuple[BudgetListItem, ...])
_ROW_COUNTS = TypeAdapter(tuple[_RowCount, ...])
SELECTED_COLUMNS = ", ".join(f'"{name}"' for name in BudgetListItem.model_fields)
@dataclass(frozen=True, slots=True)
class PrismaBudgetListExecutor:
"""The database half of the budgets list. Every caller-supplied value is bound to a
placeholder by `where_sql`; only the spec's own column names reach the SQL text."""
prisma_client: PrismaClient
async def count(self, where: tuple[Predicate, ...]) -> int:
clauses, params = where_sql(where)
sql = f"SELECT COUNT(*) AS count FROM {BUDGET_TABLE}" + (f" WHERE {clauses}" if clauses else "")
rows = await self.prisma_client.db.query_raw(sql, *params)
counted = _ROW_COUNTS.validate_python(rows)
return counted[0].count if counted else 0
async def find_many(self, plan: QueryPlan) -> Sequence[BudgetListItem]:
clauses, params = where_sql(plan.where)
sql = (
f"SELECT {SELECTED_COLUMNS} FROM {BUDGET_TABLE}"
+ (f" WHERE {clauses}" if clauses else "")
+ f" ORDER BY {order_by_sql(plan.order)}"
+ f" LIMIT ${len(params) + 1} OFFSET ${len(params) + 2}"
)
rows = await self.prisma_client.db.query_raw(sql, *params, plan.take, plan.skip)
return _BUDGET_ROWS.validate_python(rows)
def _serialize(row: BudgetListItem) -> BudgetListItem:
"""The row shape is the wire shape: the query selects exactly the columns served."""
return row
def _scope(caller: UserAPIKeyAuth) -> Scope:
if user_api_key_has_admin_view(caller):
return ScopeAll()
return ScopeDenied(reason="Only proxy admins can list budgets, your role={}".format(caller.user_role))
# budget_duration is deliberately absent from `sortable`: the column holds strings
# like "7d" and "30d", so a lexicographic ORDER BY puts "30d" ahead of "7d".
BUDGET_FILTERS: Mapping[str, FilterSpec] = MappingProxyType(
{ # mutable-ok: an immutable mapping has no literal form; MappingProxyType freezes this one and it never escapes
"budget_duration": FilterSpec(type=str, ops=frozenset(("in", "is_null"))),
"max_budget": FilterSpec(type=float, ops=frozenset(("gte", "lte", "is_null"))),
"created_at": FilterSpec(type=datetime, ops=frozenset(("gte", "lte"))),
}
)
BUDGETS_LIST_SPEC: ListSpec[BudgetListItem, BudgetListItem] = ListSpec(
resource="budgets",
sortable=frozenset(("budget_id", "max_budget", "tpm_limit", "rpm_limit", "created_at")),
searchable=frozenset(("budget_id",)),
filters=BUDGET_FILTERS,
default_sort=(SortKey(field="created_at", descending=True),),
default_page_size=50,
max_page_size=100,
scope=_scope,
serialize=_serialize,
tiebreaker="budget_id",
)
@router.get(
"/budgets",
tags=("budget management",),
dependencies=(Depends(user_api_key_auth),),
response_model=ListResponse[BudgetListItem],
)
async def list_budgets(
request: Request,
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
) -> ListResponse[BudgetListItem]:
"""
The budgets defined on this proxy, paged, sortable and filterable, for the
Budgets page.
Readable by a proxy admin or an admin viewer; anyone else is refused 403. The
older `/budget/list` answers with the whole table as a bare array and has no
way to page, sort or filter it.
`sort` takes a comma-separated list of `budget_id`, `max_budget`, `tpm_limit`,
`rpm_limit` or `created_at`, each optionally prefixed with `-` for descending,
and defaults to `-created_at`. `budget_id` is appended to every sort as the
tiebreaker. `q` is a case-insensitive substring match on `budget_id`.
`page_size` defaults to 50 and is capped at 100. Filters are
`filter[budget_duration][in|is_null]`, `filter[max_budget][gte|lte|is_null]`
and `filter[created_at][gte|lte]`.
Example curl:
```
curl --location --globoff 'http://0.0.0.0:4000/management/v1/budgets?sort=-max_budget&filter[budget_duration][in]=7d,30d&page_size=25' \
--header 'Authorization: Bearer sk-1234'
```
"""
try:
from litellm.proxy.proxy_server import prisma_client
if prisma_client is None:
raise ManagementProblem(
ProblemDetail(
type=f"{PROBLEM_TYPE_BASE}database-not-connected",
title="Database not connected",
status=503,
detail=CommonProxyErrors.db_not_connected_error.value,
)
)
return await handle_list(
spec=BUDGETS_LIST_SPEC,
executor=PrismaBudgetListExecutor(prisma_client=prisma_client),
request=request,
caller=user_api_key_dict,
)
except ManagementProblem:
raise
except Exception as e: # noqa: BLE001 # a driver error answers as a problem document, not the OpenAI error shape
verbose_proxy_logger.exception(
"litellm.proxy.management_endpoints.management_v1.budgets.list_budgets(): Exception occured - {}".format(
str(e)
)
)
raise ManagementProblem(
ProblemDetail(
type=f"{PROBLEM_TYPE_BASE}internal-server-error",
title="Internal server error",
status=500,
detail="Failed to list budgets.",
)
)

View file

@ -15,6 +15,7 @@ raw-SQL executor with every caller-supplied value bound to a placeholder.
from collections.abc import Callable, Mapping, Sequence
from dataclasses import dataclass
from datetime import datetime, timezone
from functools import partial, reduce
from math import ceil
from typing import Generic, Literal, Protocol, TypeVar
@ -213,12 +214,23 @@ def _sql_operator(op: ComparisonOp) -> str:
assert_never(op)
def _placeholder(index: int, value: FilterValue) -> str:
"""`$n`, cast when the bind is a datetime.
Binds cross into the query engine as JSON, so a datetime arrives as text and
Postgres refuses `timestamp >= text` outright. Prisma stores DateTime as a naive
`TIMESTAMP(3)` holding UTC, so the bind is read as an instant and then dropped to
naive UTC to match the column, the same cast `/spend/logs/ui` applies.
"""
return f"${index}::timestamptz AT TIME ZONE 'UTC'" if isinstance(value, datetime) else f"${index}"
def _render(predicate: Predicate, index: int) -> tuple[str, tuple[object, ...]]:
match predicate:
case IsNull(field=field, negated=negated):
return f'"{field}" IS {"NOT NULL" if negated else "NULL"}', ()
case Within(field=field, values=values):
placeholders = ", ".join(f"${index + offset}" for offset in range(len(values)))
placeholders = ", ".join(_placeholder(index + offset, value) for offset, value in enumerate(values))
return f'"{field}" IN ({placeholders})', values
case AnyOf(clauses=clauses):
rendered, params = _render_all(clauses, index)
@ -226,17 +238,31 @@ def _render(predicate: Predicate, index: int) -> tuple[str, tuple[object, ...]]:
case Compare(field=field, op="contains", value=value):
return f"\"{field}\" ILIKE ${index} ESCAPE '\\'", (f"%{escape_like(str(value))}%",)
case Compare(field=field, op=op, value=value):
return f'"{field}" {_sql_operator(op)} ${index}', (value,)
return f'"{field}" {_sql_operator(op)} {_placeholder(index, value)}', (value,)
case _:
assert_never(predicate)
def _render_one(
rendered: tuple[tuple[str, ...], tuple[object, ...]],
predicate: Predicate,
first_index: int,
) -> tuple[tuple[str, ...], tuple[object, ...]]:
"""Append one predicate, numbering it after the binds already consumed."""
clauses, params = rendered
clause, clause_params = _render(predicate, first_index + len(params))
return (*clauses, clause), (*params, *clause_params)
def _render_all(predicates: tuple[Predicate, ...], index: int) -> tuple[tuple[str, ...], tuple[object, ...]]:
if not predicates:
return (), ()
head, head_params = _render(predicates[0], index)
tail, tail_params = _render_all(predicates[1:], index + len(head_params))
return (head, *tail), head_params + tail_params
"""Render every predicate, numbering placeholders continuously across them.
Folded rather than self-recursive: walking a predicate list is a running index, and
recursing per predicate grew the stack with the filter count for nothing. `_render`
still re-enters here for `AnyOf`, whose clauses are plain `Compare`s from `?q=`, so
that nesting is one level deep and cannot be driven deeper by a caller.
"""
return reduce(partial(_render_one, first_index=index), predicates, ((), ()))
def where_sql(where: tuple[Predicate, ...], first_index: int = 1) -> tuple[str, tuple[object, ...]]:

View file

@ -57,6 +57,8 @@
- {id: mgmt.budget.update.persists, module: mgmt, tier: P1, surface: api, assertions: [persists], source: "budget_management_endpoints.py:155", rationale: "Limit changes apply"}
- {id: mgmt.budget.delete.persists, module: mgmt, tier: P1, surface: api, assertions: [persists], source: "budget_management_endpoints.py:280", rationale: "Clears limits"}
- {id: mgmt.budget.list.happy_path, module: mgmt, tier: P1, surface: api, assertions: [happy_path], source: "budget_management_endpoints.py:215", rationale: "Budget enumeration"}
- {id: mgmt.budget.list_v1.happy_path, module: mgmt, tier: P1, surface: api, assertions: [happy_path], source: "management_v1/budgets.py:129", rationale: "Budget enumeration the Budgets page can page, sort and filter"}
- {id: mgmt.budget.list_v1.admin_only, module: mgmt, tier: P1, surface: api, assertions: [admin_only], source: "management_v1/budgets.py:129", rationale: "A caller without admin view is refused, not served an empty page"}
- {id: mgmt.callback.list.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "callback_management_endpoints.py", rationale: "Callback config (smoke)"}
- {id: mgmt.cache_settings.update.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "cache_settings_endpoints.py", rationale: "Cache config (smoke). Deliberately uncovered: the previous test read the live settings and wrote them back, which proves nothing (identical values in, so a no-op POST still passes) while being able to break the deployment. /cache/settings persists what it receives and that row outranks YAML cache_params, re-applied on a timer, so a write that omits ssl or redis_startup_nodes turns a TLS cluster into a plaintext standalone node and every later Redis call hangs. That took out 60 of 72 tests on 2026-07-25. GET cannot round-trip it either: it resolves the stored row overlaid with REDIS_* env and never reads YAML, so on a fresh deploy it cannot see YAML ssl to echo back. A safe test needs an isolated proxy, or LIT-4816 fixed so a partial write cannot downgrade transport. Do not re-add a read-then-write-back test against a shared proxy."}
- {id: mgmt.cost_tracking.estimate.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "cost_tracking_settings.py", rationale: "Cost estimate (smoke)"}

View file

@ -19,10 +19,10 @@ import time
from collections.abc import Callable
import pytest
from pydantic import BaseModel, RootModel
from pydantic import BaseModel, Field, RootModel
from e2e_config import unique_marker
from e2e_http import NoBody, unwrap
from e2e_http import NoBody, Success, UnauthorizedError, UnknownApiError, unwrap
from lifecycle import ResourceManager
from management_client import ManagementClient
from models import KeyGenerateBody, OrgInfoParams, OrgNewBody, UserNewBody
@ -44,9 +44,11 @@ def _poll[T](client: ManagementClient, attempt: Callable[[], T | None], failure:
class BudgetNewBody(BaseModel):
max_budget: float
max_budget: float | None = None
soft_budget: float | None = None
budget_duration: str | None = None
budget_id: str | None = None
tpm_limit: int | None = None
class BudgetNewResponse(BaseModel):
@ -204,6 +206,164 @@ class TestBudgetManagement:
)
# ---------- /management/v1/budgets ----------
_BUDGETS_V1 = "/management/v1/budgets"
class BudgetPageParams(BaseModel):
"""Query for GET /management/v1/budgets. The filter fields serialize to the
bracketed keys the route reads them under, so nothing here is a raw dict."""
q: str | None = None
sort: str | None = None
page: int | None = None
page_size: int | None = None
duration_in: str | None = Field(default=None, serialization_alias="filter[budget_duration][in]")
max_budget_is_null: bool | None = Field(default=None, serialization_alias="filter[max_budget][is_null]")
not_a_parameter: str | None = Field(default=None, serialization_alias="filter[budget_id][eq]")
class BudgetPageMeta(BaseModel):
page: int
page_size: int
total_count: int
total_pages: int
class BudgetPageLinks(BaseModel):
first: str
prev: str | None = None
next: str | None = None
last: str
class BudgetPageRow(BaseModel):
budget_id: str
max_budget: float | None = None
tpm_limit: int | None = None
budget_duration: str | None = None
class BudgetPageResponse(BaseModel):
data: list[BudgetPageRow]
meta: BudgetPageMeta
links: BudgetPageLinks
def _list_budgets(client: ManagementClient, params: BudgetPageParams) -> BudgetPageResponse:
return unwrap(
client.proxy.transport.get(
_BUDGETS_V1,
headers=client.proxy.transport.master,
params=params,
response_type=BudgetPageResponse,
)
)
def _list_budget_ids(client: ManagementClient, params: BudgetPageParams) -> tuple[str, ...]:
return tuple(row.budget_id for row in _list_budgets(client, params).data)
def _list_status(client: ManagementClient, params: BudgetPageParams, key: str | None = None) -> int:
headers = client.proxy.transport.master if key is None else client.proxy.transport.bearer(key)
outcome = client.proxy.transport.get(
_BUDGETS_V1, headers=headers, params=params, response_type=BudgetPageResponse
)
match outcome:
case Success(status_code=status_code):
return status_code
case UnauthorizedError():
return 401
case UnknownApiError(status_code=status_code):
return status_code
case _:
raise AssertionError(outcome)
class TestBudgetListV1:
"""The paged, sorted, filtered budget list the Budgets page reads.
Every test tags its own budgets with a marker in the budget_id and searches on
it, so budgets left behind by other suites cannot move the assertions.
"""
@pytest.mark.covers("mgmt.budget.list_v1.happy_path")
def test_sorts_pages_and_filters_the_budgets_it_created(
self, client: ManagementClient, resources: ResourceManager
) -> None:
marker = unique_marker()
small, medium, large = (f"{marker}-small", f"{marker}-medium", f"{marker}-large")
for budget_id, max_budget, duration in (
(small, 1.0, "7d"),
(medium, 2.0, "30d"),
(large, 3.0, "30d"),
):
_create_budget(
client,
resources,
BudgetNewBody(
budget_id=budget_id, max_budget=max_budget, budget_duration=duration, tpm_limit=60000
),
)
mine = BudgetPageParams(q=marker, sort="-max_budget")
_ = _poll(
client,
lambda: mine if len(_list_budget_ids(client, mine)) == 3 else None,
f"{_BUDGETS_V1} never listed all three budgets tagged {marker}",
)
assert _list_budget_ids(client, mine) == (large, medium, small)
page_two = _list_budgets(client, BudgetPageParams(q=marker, sort="-max_budget", page=2, page_size=1))
assert [row.budget_id for row in page_two.data] == [medium]
assert page_two.meta.total_count == 3
assert page_two.meta.total_pages == 3
assert page_two.meta.page_size == 1
assert page_two.links.prev is not None and page_two.links.next is not None
assert set(_list_budget_ids(client, BudgetPageParams(q=marker, duration_in="30d"))) == {medium, large}
limits = _list_budgets(client, BudgetPageParams(q=marker, sort="budget_id")).data
assert [row.tpm_limit for row in limits] == [60000, 60000, 60000]
@pytest.mark.covers("mgmt.budget.list_v1.happy_path")
def test_is_null_finds_the_budget_left_uncapped(
self, client: ManagementClient, resources: ResourceManager
) -> None:
marker = unique_marker()
uncapped = _create_budget(client, resources, BudgetNewBody(budget_id=f"{marker}-uncapped"))
_ = _create_budget(client, resources, BudgetNewBody(budget_id=f"{marker}-capped", max_budget=4.0))
params = BudgetPageParams(q=marker, max_budget_is_null=True)
found = _poll(
client,
lambda: params if _list_budget_ids(client, params) == (uncapped,) else None,
f"{_BUDGETS_V1} never isolated the uncapped budget {uncapped}",
)
assert [row.max_budget for row in _list_budgets(client, found).data] == [None]
@pytest.mark.covers("mgmt.budget.list_v1.happy_path")
def test_refuses_a_sort_field_and_a_parameter_it_does_not_support(self, client: ManagementClient) -> None:
assert _list_status(client, BudgetPageParams(sort="budget_duration")) == 400
assert _list_status(client, BudgetPageParams(not_a_parameter="b-1")) == 400
@pytest.mark.covers("mgmt.budget.list_v1.admin_only")
def test_is_refused_for_a_non_admin_key(self, client: ManagementClient, resources: ResourceManager) -> None:
key = client.proxy.generate_key(KeyGenerateBody())
resources.defer(lambda: client.proxy.delete_key(key))
status = _list_status(client, BudgetPageParams(), key=key)
assert status in (401, 403), (
f"a non-admin key listing budgets must be refused 401/403, got {status}. Serving 200 with an "
f"empty page would read as 'this proxy has no budgets'"
)
# ---------- customer / end-user ----------

View file

@ -58,10 +58,14 @@ def admin_viewer_client(monkeypatch):
mock_config_table = MagicMock()
mock_config_table.find_first = AsyncMock(return_value=None)
# /management/v1/budgets reads through query_raw: count first, then the page.
mock_query_raw = AsyncMock(side_effect=[[{"count": 0}], []])
mock_prisma.db = types.SimpleNamespace(
litellm_budgettable=mock_budget_table,
litellm_invitationlink=mock_invitation_table,
litellm_config=mock_config_table,
query_raw=mock_query_raw,
)
monkeypatch.setattr(ps, "prisma_client", mock_prisma)
@ -106,6 +110,14 @@ def test_budget_list_allows_admin_viewer(admin_viewer_client):
assert resp.status_code == 200, resp.text
def test_management_v1_budgets_allows_admin_viewer(admin_viewer_client):
"""`/management/v1/budgets` is the paged/sortable budget list; same read tier as
`/budget/list`, and it answers 403 rather than an empty page when it refuses."""
resp = admin_viewer_client.get("/management/v1/budgets")
_assert_not_role_blocked(resp)
assert resp.status_code == 200, resp.text
def test_budget_settings_allows_admin_viewer(admin_viewer_client):
"""`/budget/settings` describes a budget's fields; read-only."""
resp = admin_viewer_client.get("/budget/settings", params={"budget_id": "b1"})

View file

@ -1960,6 +1960,7 @@ ADMIN_VIEWER_SETTINGS_ROUTES = [
"/config/field/info",
# Budgets page
"/budget/list",
"/management/v1/budgets",
"/budget/settings",
# Invitation viewing (admin viewer cannot create/delete; can read)
"/invitation/info",

View file

@ -0,0 +1,525 @@
from dataclasses import replace
from datetime import datetime, timezone
from typing import Any
from unittest.mock import AsyncMock, MagicMock
import pytest
from fastapi import FastAPI, Request
from fastapi.exceptions import RequestValidationError
from fastapi.testclient import TestClient
from litellm.proxy._types import LiteLLMRoutes, LitellmUserRoles
from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth
from litellm.proxy.management_endpoints.management_v1 import router
from litellm.proxy.management_endpoints.management_v1.budgets import (
BUDGETS_LIST_SPEC,
BudgetListItem,
)
from litellm.proxy.management_endpoints.management_v1.common import (
MANAGEMENT_V1_PREFIX,
PROBLEM_TYPE_BASE,
ManagementProblem,
problem_response,
)
from litellm.proxy.management_endpoints.management_v1.list_framework import (
Compare,
ScopeWhere,
build_query_plan,
)
from litellm.types.proxy.management_endpoints.management_v1 import ProblemDetail
app = FastAPI()
@app.exception_handler(ManagementProblem)
async def management_problem_exception_handler(request: Request, exc: ManagementProblem):
return problem_response(exc.problem)
@app.exception_handler(RequestValidationError)
async def validation_exception_handler(request: Request, exc: RequestValidationError):
return problem_response(
ProblemDetail(
type=f"{PROBLEM_TYPE_BASE}invalid-query-parameter",
title="Invalid query parameter",
status=400,
detail="The request query parameters are invalid.",
)
)
app.include_router(router)
client = TestClient(app)
BUDGETS_PATH = f"{MANAGEMENT_V1_PREFIX}/budgets"
SORTABLE = ["budget_id", "created_at", "max_budget", "rpm_limit", "tpm_limit"]
def _row(budget_id: str, **overrides: Any) -> dict[str, Any]:
return {
"budget_id": budget_id,
"max_budget": 10.0,
"soft_budget": None,
"tpm_limit": None,
"rpm_limit": None,
"budget_duration": "30d",
"budget_reset_at": None,
"created_at": "2026-07-20T12:00:00+00:00",
"updated_at": "2026-07-21T12:00:00+00:00",
**overrides,
}
@pytest.fixture
def query_raw(monkeypatch):
"""Mocks the one call the executor makes. `count` reads the first result, `find_many` the second."""
mock = AsyncMock(side_effect=[[{"count": 0}], []])
prisma_client = MagicMock()
prisma_client.db.query_raw = mock
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma_client)
return mock
def _serve(query_raw, rows: list[dict[str, Any]], total: int | None = None) -> None:
query_raw.side_effect = [[{"count": len(rows) if total is None else total}], rows]
@pytest.fixture
def as_proxy_admin():
app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(
user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN
)
yield
app.dependency_overrides.clear()
def _as_role(role: LitellmUserRoles):
original = app.dependency_overrides.copy()
app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(user_id="u", user_role=role)
return original
def _get(query: str = ""):
suffix = f"?{query}" if query else ""
return client.get(f"{BUDGETS_PATH}{suffix}", headers={"Authorization": "Bearer k"})
def _select_call(query_raw):
"""The find_many call: (sql, *params). The count call comes first."""
return query_raw.call_args_list[1].args
def test_returns_flat_rows_in_the_control_plane_envelope(query_raw, as_proxy_admin):
"""`{data, meta, links}` with flat rows; no JSON:API `{type, id, attributes}` wrapper."""
_serve(query_raw, [_row("b-1")])
response = _get()
assert response.status_code == 200
body = response.json()
assert set(body) == {"data", "meta", "links"}
assert body["data"][0]["budget_id"] == "b-1"
assert "attributes" not in body["data"][0]
def test_serves_the_columns_the_budgets_page_renders(query_raw, as_proxy_admin):
_serve(query_raw, [_row("b-1", soft_budget=5.0, budget_reset_at="2026-08-01T00:00:00+00:00")])
row = _get().json()["data"][0]
assert set(row) == {
"budget_id",
"max_budget",
"soft_budget",
"tpm_limit",
"rpm_limit",
"budget_duration",
"budget_reset_at",
"created_at",
"updated_at",
}
assert row["soft_budget"] == 5.0
assert row["budget_reset_at"].startswith("2026-08-01T00:00:00")
def test_selects_only_the_columns_it_serves(query_raw, as_proxy_admin):
"""A `SELECT *` would ship created_by/updated_by and model_max_budget to the browser."""
_serve(query_raw, [])
_get()
sql = _select_call(query_raw)[0]
assert "SELECT *" not in sql
assert '"budget_id"' in sql and '"max_budget"' in sql
assert "created_by" not in sql and "model_max_budget" not in sql
def test_defaults_to_newest_first_with_budget_id_breaking_ties(query_raw, as_proxy_admin):
"""Two budgets created in the same transaction share a created_at; without the
tiebreaker their relative order is undefined and pages can repeat or drop rows."""
_serve(query_raw, [])
_get()
assert 'ORDER BY "created_at" DESC NULLS LAST, "budget_id" ASC NULLS LAST' in _select_call(query_raw)[0]
def test_appends_the_tiebreaker_to_an_explicit_sort(query_raw, as_proxy_admin):
_serve(query_raw, [])
_get("sort=-max_budget")
assert 'ORDER BY "max_budget" DESC NULLS LAST, "budget_id" ASC NULLS LAST' in _select_call(query_raw)[0]
def test_refuses_to_sort_on_budget_duration(query_raw, as_proxy_admin):
"""The column holds "7d"/"30d", so a lexicographic ORDER BY would put "30d"
before "7d" and silently mis-order the page."""
_serve(query_raw, [])
response = _get("sort=budget_duration")
assert response.status_code == 400
assert response.headers["content-type"].startswith("application/problem+json")
body = response.json()
assert "budget_duration" in body["detail"]
assert body["allowed"] == SORTABLE
query_raw.assert_not_called()
def test_the_advertised_sort_fields_are_the_ones_that_work(query_raw, as_proxy_admin):
"""Guards the rejection above against drifting from what the spec actually accepts."""
for field in SORTABLE:
_serve(query_raw, [])
assert _get(f"sort={field}").status_code == 200, field
assert sorted(BUDGETS_LIST_SPEC.sortable) == SORTABLE
def test_rejects_an_unknown_query_parameter(query_raw, as_proxy_admin):
"""A silently ignored filter over-returns budgets, which is worse than a rejected request."""
_serve(query_raw, [])
response = _get("filtre[max_budget][gte]=5")
assert response.status_code == 400
assert response.headers["content-type"].startswith("application/problem+json")
body = response.json()
assert "filtre[max_budget][gte]" in body["detail"]
assert "filter[max_budget][gte]" in body["allowed"]
query_raw.assert_not_called()
def test_rejects_an_operator_the_filter_does_not_declare(query_raw, as_proxy_admin):
"""`max_budget` takes ranges, not `in`; accepting an undeclared operator is how a
filter starts meaning something the query planner never checked."""
for query in ("filter[max_budget][in]=5,10", "filter[created_at][is_null]=true"):
_serve(query_raw, [])
assert _get(query).status_code == 400, query
def test_omitted_page_size_serves_fifty(query_raw, as_proxy_admin):
_serve(query_raw, [])
body = _get().json()
assert body["meta"]["page_size"] == 50
assert _select_call(query_raw)[-2] == 50
def test_clamps_an_oversized_page_size_to_a_hundred(query_raw, as_proxy_admin):
"""Unclamped, one request can ask the proxy to serialize the whole budget table."""
_serve(query_raw, [])
body = _get("page_size=500").json()
assert body["meta"]["page_size"] == 100
assert _select_call(query_raw)[-2] == 100
def test_offsets_by_page(query_raw, as_proxy_admin):
_serve(query_raw, [])
_get("page=3&page_size=25")
assert _select_call(query_raw)[-2:] == (25, 50)
@pytest.mark.parametrize(
"role",
[
LitellmUserRoles.INTERNAL_USER,
LitellmUserRoles.INTERNAL_USER_VIEW_ONLY,
LitellmUserRoles.TEAM,
],
)
def test_refuses_a_caller_without_admin_view(query_raw, role):
"""Budgets are proxy-wide, so a caller who cannot read all of them must be told
so. Answering 200 with an empty list would read as "there are no budgets"."""
_serve(query_raw, [_row("b-1")])
original = _as_role(role)
try:
response = _get()
finally:
app.dependency_overrides = original
assert response.status_code == 403
assert response.headers["content-type"].startswith("application/problem+json")
assert response.json()["status"] == 403
query_raw.assert_not_called()
@pytest.mark.parametrize("role", [LitellmUserRoles.PROXY_ADMIN, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY])
def test_admins_and_admin_viewers_may_read_every_budget(query_raw, role):
_serve(query_raw, [_row("b-1")])
original = _as_role(role)
try:
response = _get()
finally:
app.dependency_overrides = original
assert response.status_code == 200
assert [row["budget_id"] for row in response.json()["data"]] == ["b-1"]
assert "WHERE" not in _select_call(query_raw)[0]
def test_a_denied_caller_stays_denied_whatever_they_filter_on(query_raw):
"""The scope decision reads the caller, never the query string."""
_serve(query_raw, [_row("b-1")])
original = _as_role(LitellmUserRoles.INTERNAL_USER)
try:
response = _get("filter[max_budget][gte]=0&q=b-")
finally:
app.dependency_overrides = original
assert response.status_code == 403
def test_a_filter_sits_behind_the_scope_predicate_instead_of_replacing_it():
"""Same spec, planned for a row-scoped caller: the scope clause has to survive, and
lead, whatever the caller filtered on. Replacing it would let a filter widen a read."""
scoped = replace(
BUDGETS_LIST_SPEC,
scope=lambda _caller: ScopeWhere(where=(Compare(field="budget_id", op="eq", value="b-1"),)),
)
plan = build_query_plan(
spec=scoped,
params={"filter[max_budget][gte]": "5"},
caller=UserAPIKeyAuth(user_id="u", user_role=LitellmUserRoles.INTERNAL_USER),
)
assert plan.where[0] == Compare(field="budget_id", op="eq", value="b-1")
assert Compare(field="max_budget", op="gte", value=5.0) in plan.where
def test_q_matches_budget_id_case_insensitively(query_raw, as_proxy_admin):
"""budget_id is the only text identity on the row; matching anything else would
return budgets whose ids do not contain what the user typed."""
_serve(query_raw, [])
_get("q=Prod")
sql, *params = _select_call(query_raw)
assert '"budget_id" ILIKE $1' in sql
assert params[0] == "%Prod%"
def test_q_escapes_like_metacharacters(query_raw, as_proxy_admin):
"""Budget ids routinely contain '_'; unescaped it is a single-character wildcard."""
_serve(query_raw, [])
_get("q=team_a%25")
assert _select_call(query_raw)[1] == r"%team\_a\%%"
def test_q_does_not_search_any_other_column(query_raw, as_proxy_admin):
_serve(query_raw, [])
_get("q=30d")
assert "budget_duration" not in _select_call(query_raw)[0].split("WHERE")[1]
assert BUDGETS_LIST_SPEC.searchable == frozenset({"budget_id"})
def test_is_null_selects_the_unlimited_budgets(query_raw, as_proxy_admin):
""""Unlimited" is max_budget IS NULL; `max_budget = 0` would be a hard zero cap."""
_serve(query_raw, [_row("b-unlimited", max_budget=None)])
body = _get("filter[max_budget][is_null]=true").json()
assert '"max_budget" IS NULL' in _select_call(query_raw)[0]
assert body["data"][0]["max_budget"] is None
def test_is_null_false_selects_the_capped_budgets(query_raw, as_proxy_admin):
_serve(query_raw, [])
_get("filter[max_budget][is_null]=false")
assert '"max_budget" IS NOT NULL' in _select_call(query_raw)[0]
def test_in_filter_binds_each_requested_duration(query_raw, as_proxy_admin):
_serve(query_raw, [])
_get("filter[budget_duration][in]=7d,30d")
sql, *params = _select_call(query_raw)
assert '"budget_duration" IN ($1, $2)' in sql
assert params[:2] == ["7d", "30d"]
def test_created_at_range_is_bound_as_a_timestamp(query_raw, as_proxy_admin):
"""The bind crosses into the query engine as JSON, so an uncast placeholder reaches
Postgres as text and `timestamp >= text` is a hard error, not a wrong answer."""
_serve(query_raw, [])
_get("filter[created_at][gte]=2026-07-01T00:00:00Z")
sql, *params = _select_call(query_raw)
assert "\"created_at\" >= $1::timestamptz AT TIME ZONE 'UTC'" in sql
assert params[0] == datetime(2026, 7, 1, tzinfo=timezone.utc)
def test_numbers_placeholders_continuously_across_predicates(query_raw, as_proxy_admin):
"""Each predicate is numbered after the binds the ones before it consumed. Restart
the count and `$1` gets read as the duration while the search string goes unbound."""
_serve(query_raw, [])
_get("filter[budget_duration][in]=7d,30d&filter[max_budget][gte]=5&q=prod")
sql, *params = _select_call(query_raw)
assert '"budget_duration" IN ($1, $2)' in sql
assert '"max_budget" >= $3' in sql
assert '"budget_id" ILIKE $4' in sql
assert params[:4] == ["7d", "30d", 5.0, "%prod%"]
def test_a_non_datetime_bind_is_not_cast(query_raw, as_proxy_admin):
"""Guards the cast above from being applied to every placeholder."""
_serve(query_raw, [])
_get("filter[max_budget][gte]=5")
sql = _select_call(query_raw)[0]
assert '"max_budget" >= $1' in sql
assert "timestamptz" not in sql
def test_an_offsetless_created_at_bound_is_read_as_utc(query_raw, as_proxy_admin):
"""The dashboard sends 'YYYY-MM-DDTHH:MM:SS' with no offset. Left naive, Postgres
would compare it in the session timezone and shift the window off the rows shown."""
_serve(query_raw, [])
_get("filter[created_at][gte]=2026-07-01T00:00:00")
bound = _select_call(query_raw)[1]
assert bound == datetime(2026, 7, 1, tzinfo=timezone.utc)
assert bound.tzinfo is not None
def test_rejects_a_filter_value_that_is_not_of_the_declared_type(query_raw, as_proxy_admin):
for query in ("filter[max_budget][gte]=lots", "filter[created_at][gte]=yesterday"):
_serve(query_raw, [])
assert _get(query).status_code == 400, query
def test_reports_the_total_and_links_every_page_on_a_middle_page(query_raw, as_proxy_admin):
"""The Budgets page renders a page count, so the total has to be the match total,
not the length of the page it just received."""
_serve(query_raw, [_row("b-3"), _row("b-4")], total=7)
body = _get("page=2&page_size=2").json()
assert body["meta"] == {"page": 2, "page_size": 2, "total_count": 7, "total_pages": 4}
links = body["links"]
assert "page=1" in links["first"] and "page_size=2" in links["first"]
assert "page=1" in links["prev"]
assert "page=3" in links["next"]
assert "page=4" in links["last"]
assert "page=2" in links["self"]
def test_the_last_page_has_no_next(query_raw, as_proxy_admin):
_serve(query_raw, [_row("b-5")], total=5)
links = _get("page=3&page_size=2").json()["links"]
assert links["next"] is None
assert "page=2" in links["prev"]
def test_the_first_page_has_no_prev(query_raw, as_proxy_admin):
_serve(query_raw, [_row("b-1")], total=5)
links = _get("page_size=2").json()["links"]
assert links["prev"] is None
assert "page=2" in links["next"]
def test_an_empty_table_still_links_a_first_and_last_page(query_raw, as_proxy_admin):
_serve(query_raw, [], total=0)
body = _get().json()
assert body["meta"]["total_count"] == 0
assert body["meta"]["total_pages"] == 0
assert "page=1" in body["links"]["first"] and "page=1" in body["links"]["last"]
def test_counts_over_the_same_predicate_it_pages(query_raw, as_proxy_admin):
"""A total counted without the caller's filter would page through rows the
filter excluded."""
_serve(query_raw, [], total=0)
_get("filter[budget_duration][in]=30d")
count_sql, *count_params = query_raw.call_args_list[0].args
assert '"budget_duration" IN ($1)' in count_sql
assert count_params == ["30d"]
assert "COUNT(*)" in count_sql
def test_bigint_limits_serialize_as_json_numbers(query_raw, as_proxy_admin):
"""tpm_limit/rpm_limit are BigInt? in Prisma; the query engine can hand them back
as decimal strings, and a quoted "60000" breaks arithmetic in the dashboard."""
_serve(query_raw, [_row("b-1", tpm_limit="60000", rpm_limit=1200)])
response = _get()
row = response.json()["data"][0]
assert row["tpm_limit"] == 60000
assert row["rpm_limit"] == 1200
assert isinstance(row["tpm_limit"], int) and not isinstance(row["tpm_limit"], bool)
assert '"tpm_limit": "60000"' not in response.text
def test_reports_a_missing_database_as_a_problem_document(monkeypatch, as_proxy_admin):
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
response = _get()
assert response.status_code == 503
assert response.headers["content-type"].startswith("application/problem+json")
def test_the_spec_serves_what_the_row_model_declares():
"""SELECTED_COLUMNS is built off the model, so a field added to one cannot go
missing from the other and produce a row the validator rejects."""
from litellm.proxy.management_endpoints.management_v1.budgets import SELECTED_COLUMNS
assert SELECTED_COLUMNS == ", ".join(f'"{name}"' for name in BudgetListItem.model_fields)
assert BUDGETS_LIST_SPEC.tiebreaker in BudgetListItem.model_fields
assert BUDGETS_LIST_SPEC.sortable <= frozenset(BudgetListItem.model_fields)
assert frozenset(BUDGETS_LIST_SPEC.filters) <= frozenset(BudgetListItem.model_fields)
def test_is_reachable_by_the_roles_that_can_open_the_budgets_page():
"""Route-level auth gate, which the dependency_overrides above bypass. The handler's
admin-view check is dead code if RouteChecks rejects the role first."""
assert BUDGETS_PATH in LiteLLMRoutes.admin_viewer_routes.value
assert ("/budget/list" in LiteLLMRoutes.admin_viewer_routes.value) == (
BUDGETS_PATH in LiteLLMRoutes.admin_viewer_routes.value
)

View file

@ -7169,6 +7169,44 @@ export interface paths {
patch?: never;
trace?: never;
};
"/management/v1/budgets": {
parameters: {
query?: never;
header?: never;
path?: never;
cookie?: never;
};
/**
* List Budgets
* @description The budgets defined on this proxy, paged, sortable and filterable, for the
* Budgets page.
*
* Readable by a proxy admin or an admin viewer; anyone else is refused 403. The
* older `/budget/list` answers with the whole table as a bare array and has no
* way to page, sort or filter it.
*
* `sort` takes a comma-separated list of `budget_id`, `max_budget`, `tpm_limit`,
* `rpm_limit` or `created_at`, each optionally prefixed with `-` for descending,
* and defaults to `-created_at`. `budget_id` is appended to every sort as the
* tiebreaker. `q` is a case-insensitive substring match on `budget_id`.
* `page_size` defaults to 50 and is capped at 100. Filters are
* `filter[budget_duration][in|is_null]`, `filter[max_budget][gte|lte|is_null]`
* and `filter[created_at][gte|lte]`.
*
* Example curl:
* ```
* curl --location --globoff 'http://0.0.0.0:4000/management/v1/budgets?sort=-max_budget&filter[budget_duration][in]=7d,30d&page_size=25' --header 'Authorization: Bearer sk-1234'
* ```
*/
get: operations["list_budgets_management_v1_budgets_get"];
put?: never;
post?: never;
delete?: never;
options?: never;
head?: never;
patch?: never;
trace?: never;
};
"/management/v1/spend_logs/end_users": {
parameters: {
query?: never;
@ -21533,6 +21571,40 @@ export interface components {
/** Reset At */
reset_at?: string | null;
};
/**
* BudgetListItem
* @description One budget as the Budgets page reads it, and as it comes back off the table.
*
* Validating the raw row through here is what makes `tpm_limit` / `rpm_limit`
* numbers: they are `BigInt?` in the schema, which the query engine hands back as
* decimal strings, and a quoted "60000" breaks arithmetic in the dashboard.
*/
BudgetListItem: {
/** Budget Duration */
budget_duration?: string | null;
/** Budget Id */
budget_id: string;
/** Budget Reset At */
budget_reset_at?: string | null;
/**
* Created At
* Format: date-time
*/
created_at: string;
/** Max Budget */
max_budget?: number | null;
/** Rpm Limit */
rpm_limit?: number | null;
/** Soft Budget */
soft_budget?: number | null;
/** Tpm Limit */
tpm_limit?: number | null;
/**
* Updated At
* Format: date-time
*/
updated_at: string;
};
/** BudgetNewRequest */
BudgetNewRequest: {
/**
@ -24928,6 +25000,36 @@ export interface components {
/** Guardrails */
guardrails: components["schemas"]["GuardrailInfoResponse"][];
};
/**
* ListLinks
* @description Page-mode counterpart to `PageLinks`. `first`/`last` are knowable here because the total count is.
*/
ListLinks: {
/** First */
first: string;
/** Last */
last: string;
/** Next */
next?: string | null;
/** Prev */
prev?: string | null;
/** Self */
self: string;
};
/**
* ListMeta
* @description Page-mode counterpart to `PageMeta`: an entity list pays for the COUNT(*) so the table can show a page count.
*/
ListMeta: {
/** Page */
page: number;
/** Page Size */
page_size: number;
/** Total Count */
total_count: number;
/** Total Pages */
total_pages: number;
};
/**
* ListPluginsResponse
* @description Response from listing plugins.
@ -24943,6 +25045,13 @@ export interface components {
/** Prompts */
prompts: components["schemas"]["PromptSpec"][];
};
/** ListResponse[BudgetListItem] */
ListResponse_BudgetListItem_: {
/** Data */
data: components["schemas"]["BudgetListItem"][];
links: components["schemas"]["ListLinks"];
meta: components["schemas"]["ListMeta"];
};
/**
* ListRunsResponse
* @description Response from listing runs
@ -43479,6 +43588,26 @@ export interface operations {
};
};
};
list_budgets_management_v1_budgets_get: {
parameters: {
query?: never;
header?: never;
path?: never;
cookie?: never;
};
requestBody?: never;
responses: {
/** @description Successful Response */
200: {
headers: {
[name: string]: unknown;
};
content: {
"application/json": components["schemas"]["ListResponse_BudgetListItem_"];
};
};
};
};
list_spend_log_end_users_management_v1_spend_logs_end_users_get: {
parameters: {
query: {