diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index d85ad173434..23fe7af9994 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -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", diff --git a/litellm/proxy/management_endpoints/management_v1/__init__.py b/litellm/proxy/management_endpoints/management_v1/__init__.py index 257de66130b..a06c6b2591c 100644 --- a/litellm/proxy/management_endpoints/management_v1/__init__.py +++ b/litellm/proxy/management_endpoints/management_v1/__init__.py @@ -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"] diff --git a/litellm/proxy/management_endpoints/management_v1/budgets.py b/litellm/proxy/management_endpoints/management_v1/budgets.py new file mode 100644 index 00000000000..bc1521caf0b --- /dev/null +++ b/litellm/proxy/management_endpoints/management_v1/budgets.py @@ -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.", + ) + ) diff --git a/litellm/proxy/management_endpoints/management_v1/list_framework.py b/litellm/proxy/management_endpoints/management_v1/list_framework.py index 3e4b9131d1e..8f25c45016b 100644 --- a/litellm/proxy/management_endpoints/management_v1/list_framework.py +++ b/litellm/proxy/management_endpoints/management_v1/list_framework.py @@ -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, ...]]: diff --git a/tests/e2e/coverage_registry/mgmt.yaml b/tests/e2e/coverage_registry/mgmt.yaml index 68c5ef6b31d..2a0fc5c9f29 100644 --- a/tests/e2e/coverage_registry/mgmt.yaml +++ b/tests/e2e/coverage_registry/mgmt.yaml @@ -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)"} diff --git a/tests/e2e/management/test_budget_customer_user_org_e2e.py b/tests/e2e/management/test_budget_customer_user_org_e2e.py index 54cc18b228b..12372bb7cc1 100644 --- a/tests/e2e/management/test_budget_customer_user_org_e2e.py +++ b/tests/e2e/management/test_budget_customer_user_org_e2e.py @@ -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 ---------- diff --git a/tests/test_litellm/proxy/auth/test_admin_viewer_handler_access.py b/tests/test_litellm/proxy/auth/test_admin_viewer_handler_access.py index b0d2595e48c..36f656a7adc 100644 --- a/tests/test_litellm/proxy/auth/test_admin_viewer_handler_access.py +++ b/tests/test_litellm/proxy/auth/test_admin_viewer_handler_access.py @@ -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"}) diff --git a/tests/test_litellm/proxy/auth/test_route_checks.py b/tests/test_litellm/proxy/auth/test_route_checks.py index a6d4dc63697..06764139eda 100644 --- a/tests/test_litellm/proxy/auth/test_route_checks.py +++ b/tests/test_litellm/proxy/auth/test_route_checks.py @@ -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", diff --git a/tests/test_litellm/proxy/management_endpoints/management_v1/test_budgets.py b/tests/test_litellm/proxy/management_endpoints/management_v1/test_budgets.py new file mode 100644 index 00000000000..40473f1a25a --- /dev/null +++ b/tests/test_litellm/proxy/management_endpoints/management_v1/test_budgets.py @@ -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 + ) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index e7c2b3a5a54..9a301f1c474 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -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: {