From 83ef72a76a84328d2c81b2bd2bbe9f16142f5bb9 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Wed, 19 Aug 2026 18:11:51 -0700 Subject: [PATCH] fix(key budgets): stop the report inventing denials and leaking keys through the URL Round-2 review fixes on the budgets endpoint. Making the route an info route also enrolled it in proxy-only error logging, which hands the request path to every failure callback. That gate exists to stop management endpoints leaking temporary keys, and /key/info is immune only because its key is a query param. The path form now rejects a plaintext key and asks for the hash, since a key in a URL also reaches access logs and span names regardless of this gate. The reservation layer builds no counter for a non-positive cap, so reporting its tightened operator there turned a team, tag or end user sitting at max_budget 0.0 into "exceeded" when both layers admit the request. The tightening now needs a positive cap. Per-model budgets are enforced per request model, not per cap, so reporting the highest counter under a cap claimed a denial that no request would hit. Each request model that maps onto a cap gets its own row, identified by the model whose counter it reports. Deployment names join the candidates, because routing straight at a deployment keys the counter on that name rather than on a model group. The reservation note claimed more than the layer delivers: a request it cannot price up front is gated by the read-time check alone. Reworded rather than re-implemented, since introspection cannot know the request. Severity is now the fallback for a code a client has not been taught yet, not a ranking: warning means the numbers may be incomplete or misread, info means they are accurate. Also dropped the token end-user cap plumbing, which could only ever attribute the caller's request-scoped cap to the inspected key, and memoised is_info_route, which the pattern change had put 25 uncached regex builds per request behind. --- litellm/proxy/auth/route_checks.py | 6 + .../proxy/hooks/model_max_budget_limiter.py | 3 +- .../key_budget_resolver.py | 134 ++++++--------- .../key_management_endpoints.py | 26 ++- .../key_management_endpoints.py | 9 +- .../proxy/auth/test_info_routes.py | 14 ++ .../test_key_management_endpoints.py | 154 ++++++++++++++++-- 7 files changed, 233 insertions(+), 113 deletions(-) diff --git a/litellm/proxy/auth/route_checks.py b/litellm/proxy/auth/route_checks.py index 9758162e753..be697ffb792 100644 --- a/litellm/proxy/auth/route_checks.py +++ b/litellm/proxy/auth/route_checks.py @@ -1,4 +1,5 @@ import re +from functools import lru_cache from typing import Final from fastapi import HTTPException, Request, status @@ -437,12 +438,17 @@ class RouteChecks: return RouteChecks.check_route_access(route=route, allowed_routes=LiteLLMRoutes.management_routes.value) @staticmethod + @lru_cache(maxsize=2048) def is_info_route(route: str) -> bool: """ Check if route is an info route Pattern-aware, like ``is_management_route``, so an info route carrying a path parameter is as reachable as one without: the incoming route holds a resolved id, never the ``{...}`` template. + + Cached because this runs per request off ``_check_end_user_budget`` and the allowlist is a + module constant, so an uncached call rebuilds and matches 25 regexes to answer the same + question. ``normalize_request_route`` in ``auth_utils`` is bounded the same way. """ return RouteChecks.check_route_access(route=route, allowed_routes=LiteLLMRoutes.info_routes.value) diff --git a/litellm/proxy/hooks/model_max_budget_limiter.py b/litellm/proxy/hooks/model_max_budget_limiter.py index 85f96480016..9fdd3140742 100644 --- a/litellm/proxy/hooks/model_max_budget_limiter.py +++ b/litellm/proxy/hooks/model_max_budget_limiter.py @@ -1,4 +1,5 @@ import json +from collections.abc import Mapping from typing import Final import litellm @@ -202,7 +203,7 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): return _current_spend def get_request_model_budget_key( - self, model: str, internal_model_max_budget: GenericBudgetConfigType + self, model: str, internal_model_max_budget: Mapping[str, BudgetConfig] ) -> str | None: """ Which `model_max_budget` key a request model is charged against. diff --git a/litellm/proxy/management_endpoints/key_budget_resolver.py b/litellm/proxy/management_endpoints/key_budget_resolver.py index f2297e784d9..5a594b550a3 100644 --- a/litellm/proxy/management_endpoints/key_budget_resolver.py +++ b/litellm/proxy/management_endpoints/key_budget_resolver.py @@ -103,15 +103,15 @@ _ALERT_ONLY_NOTE: Final = KeyBudgetNote( ) _ROLLING_WINDOW_NOTE: Final = KeyBudgetNote( code="rolling_window", - severity="warning", + severity="info", text="rolling window; the start moves with reset_at so consecutive windows can overlap", ) _MODEL_BUDGET_NOTE: Final = KeyBudgetNote( code="per_model_counters", severity="warning", text=( - "per-model spend is counted separately for every request model that maps onto this cap, " - "and each counter is compared against it on its own, so the highest is reported" + "one row per request model that maps onto this cap, because each model's counter is compared " + "against it alone; a model this proxy cannot enumerate has a counter that is not reported here" ), ) _MODEL_BUDGET_COLD_NOTE: Final = KeyBudgetNote( @@ -121,15 +121,15 @@ _MODEL_BUDGET_COLD_NOTE: Final = KeyBudgetNote( ) _RESERVATION_NOTE: Final = KeyBudgetNote( code="reservation_blocks_at_limit", - severity="warning", + severity="info", text=( - "the reservation layer blocks this scope as soon as spend reaches the limit, " - "before the read-time check would trip" + "the reservation layer usually blocks this scope as soon as spend reaches the limit, ahead of the " + "read-time check; requests it cannot price up front are still gated by the read-time check alone" ), ) _PROJECT_SPEND_NOTE: Final = KeyBudgetNote( code="project_spend_not_tracked", - severity="info", + severity="warning", text="project spend is never incremented today, so this budget cannot trip", ) _TAG_NOTE: Final = KeyBudgetNote( @@ -177,7 +177,7 @@ class SpendReader(Protocol): class ModelSpendReader(Protocol): - async def __call__(self, *, entity_id: str, models: Sequence[str], budget_config: BudgetConfig) -> float | None: ... + async def __call__(self, *, entity_id: str, model: str, budget_config: BudgetConfig) -> float | None: ... class ModelBudgetKeyMatcher(Protocol): @@ -207,62 +207,41 @@ async def _read_counter_spend( ) -def _highest(spends: Sequence[float | None]) -> float | None: - """Each request model has its own counter compared against the same cap, so the highest is the one that blocks.""" - found: Final = tuple(spend for spend in spends if spend is not None) - return max(found) if found else None - - -async def _read_key_model_spend(*, entity_id: str, models: Sequence[str], budget_config: BudgetConfig) -> float | None: +async def _read_key_model_spend(*, entity_id: str, model: str, budget_config: BudgetConfig) -> float | None: from litellm.proxy.proxy_server import model_max_budget_limiter - return _highest( - await asyncio.gather( - *( - model_max_budget_limiter.get_virtual_key_spend_for_model( - user_api_key_hash=entity_id, - model=model, - key_budget_config=budget_config, - ) - for model in models - ) - ) + return await model_max_budget_limiter.get_virtual_key_spend_for_model( + user_api_key_hash=entity_id, model=model, key_budget_config=budget_config ) -async def _read_end_user_model_spend( - *, entity_id: str, models: Sequence[str], budget_config: BudgetConfig -) -> float | None: +async def _read_end_user_model_spend(*, entity_id: str, model: str, budget_config: BudgetConfig) -> float | None: from litellm.proxy.proxy_server import model_max_budget_limiter - return _highest( - await asyncio.gather( - *( - model_max_budget_limiter.get_end_user_spend_for_model( - end_user_id=entity_id, - model=model, - key_budget_config=budget_config, - ) - for model in models - ) - ) + return await model_max_budget_limiter.get_end_user_spend_for_model( + end_user_id=entity_id, model=model, key_budget_config=budget_config ) def _match_model_budget_key(*, model: str, configured: Mapping[str, BudgetConfig]) -> str | None: from litellm.proxy.proxy_server import model_max_budget_limiter - return model_max_budget_limiter.get_request_model_budget_key( - model=model, internal_model_max_budget=dict(configured) - ) + return model_max_budget_limiter.get_request_model_budget_key(model=model, internal_model_max_budget=configured) def _request_models(allowed_models: tuple[str, ...]) -> tuple[str, ...]: - """Model names a request could carry, since per-model counters are keyed by the request model.""" + """ + Model names a request could carry, since per-model counters are keyed by the request model. + + Deployment names are in here because routing straight at one bypasses the model group, and a + wildcard or non-router deployment can still produce a counter no enumeration can predict. + """ from litellm.proxy.proxy_server import llm_router - router_models: Final = tuple(llm_router.get_model_names()) if llm_router is not None else () - return tuple(dict.fromkeys((*allowed_models, *router_models))) + if llm_router is None: + return allowed_models + routable: Final = (*llm_router.get_model_names(), *llm_router.deployment_names) + return tuple(dict.fromkeys((*allowed_models, *routable))) @dataclass(frozen=True, slots=True) @@ -296,14 +275,14 @@ class _RecordedSpend: @dataclass(frozen=True, slots=True) class _KeyModelSpend: key_hash: str - models: tuple[str, ...] + model: str budget_config: BudgetConfig @dataclass(frozen=True, slots=True) class _EndUserModelSpend: end_user_id: str - models: tuple[str, ...] + model: str budget_config: BudgetConfig @@ -457,7 +436,6 @@ class _KeyBudgetContext: valid_token: UserAPIKeyAuth token_inputs: _TokenBudgetInputs end_user_id: str | None - token_end_user_max_budget: float | None custom_auth_enabled: bool request_models: tuple[str, ...] match_model_budget_key: ModelBudgetKeyMatcher @@ -478,15 +456,9 @@ async def resolve_key_budgets( valid_token: UserAPIKeyAuth, end_user_id: str | None, deps: KeyBudgetResolverDeps, - token_end_user_max_budget: float | None = None, ) -> tuple[KeyBudgetEntry, ...]: """Every budget that applies to ``valid_token``, configured or not, with live spend.""" - context: Final = await _load_context( - valid_token=valid_token, - end_user_id=end_user_id, - deps=deps, - token_end_user_max_budget=token_end_user_max_budget, - ) + context: Final = await _load_context(valid_token=valid_token, end_user_id=end_user_id, deps=deps) plans: Final = _plan_budgets(context) readings: Final = await asyncio.gather(*(_read_spend(plan=plan, deps=deps) for plan in plans)) reservation_enabled: Final = deps.general_settings.get("disable_budget_reservation") is not True @@ -523,13 +495,13 @@ async def _read_spend(plan: _PlannedBudget, deps: KeyBudgetResolverDeps) -> _Spe case _KeyModelSpend(): return _model_reading( await deps.read_key_model_spend( - entity_id=source.key_hash, models=source.models, budget_config=source.budget_config + entity_id=source.key_hash, model=source.model, budget_config=source.budget_config ) ) case _EndUserModelSpend(): return _model_reading( await deps.read_end_user_model_spend( - entity_id=source.end_user_id, models=source.models, budget_config=source.budget_config + entity_id=source.end_user_id, model=source.model, budget_config=source.budget_config ) ) case _: @@ -543,9 +515,14 @@ def _effective_comparison(plan: _PlannedBudget, reservation_enabled: bool) -> Bu """ Reservation runs ahead of the read-time check and blocks once spend reaches the cap, so for every scope it covers it, not the read-time comparison, decides when a request stops going through. + + It builds no counter at all for a non-positive cap, though, which leaves the read-time operator + in charge there and is the one case where reporting the tightened one would invent a denial. """ if plan.enforcement == "soft" or not reservation_enabled: return plan.comparison + if plan.max_budget is None or plan.max_budget <= 0: + return plan.comparison if plan.scope not in _RESERVATION_COVERED_SCOPES: return plan.comparison return ">=" @@ -592,7 +569,6 @@ async def _load_context( valid_token: UserAPIKeyAuth, end_user_id: str | None, deps: KeyBudgetResolverDeps, - token_end_user_max_budget: float | None, ) -> _KeyBudgetContext: token_inputs: Final = _token_budget_inputs(valid_token) proxy, team, user, project, tags, end_user = await asyncio.gather( @@ -620,7 +596,6 @@ async def _load_context( valid_token=valid_token, token_inputs=token_inputs, end_user_id=end_user_id, - token_end_user_max_budget=token_end_user_max_budget, custom_auth_enabled=deps.custom_auth_enabled, request_models=_request_models(token_inputs.models), match_model_budget_key=deps.match_model_budget_key, @@ -954,21 +929,18 @@ def _plan_key_models(context: _KeyBudgetContext) -> tuple[_PlannedBudget, ...]: return tuple( _PlannedBudget( scope="key_model", - entity_id=model, - entity_label=token.key_alias, + entity_id=request_model, + entity_label=cap_key, enforcement="hard", max_budget=_positive_or_none(config.max_budget), comparison=">", - source=f"key.model_max_budget[{model}]", - spend_source=_KeyModelSpend( - key_hash=token.token or "", - models=_model_counter_candidates(config_key=model, configured=configured, context=context), - budget_config=config, - ), + source=f"key.model_max_budget[{cap_key}]", + spend_source=_KeyModelSpend(key_hash=token.token or "", model=request_model, budget_config=config), budget_duration=config.budget_duration, notes=(_MODEL_BUDGET_NOTE,), ) - for model, config in configured.items() + for cap_key, config in configured.items() + for request_model in _model_counter_candidates(config_key=cap_key, configured=configured, context=context) ) @@ -1190,17 +1162,14 @@ def _plan_end_user(context: _KeyBudgetContext) -> tuple[_PlannedBudget, ...]: if end_user is not None and end_user.budget_id is not None else "litellm_settings.max_end_user_budget_id" ) - # Reservation reads the request token's cap first and only falls back to the row, so mirror that order. - token_max_budget: Final = context.token_end_user_max_budget - source: Final = "token.end_user_max_budget" if token_max_budget is not None else row_source primary: Final = _PlannedBudget( scope="end_user", entity_id=end_user_id, entity_label=end_user.alias if end_user is not None else None, enforcement="hard", - max_budget=token_max_budget if token_max_budget is not None else (budget.max_budget if budget else None), + max_budget=budget.max_budget if budget is not None else None, comparison=">", - source=source, + source=row_source, spend_source=_CounterSpend( counter_key=end_user_spend_counter(end_user_id), fallback_spend=(end_user.spend or 0.0) if end_user is not None else 0.0, @@ -1210,7 +1179,7 @@ def _plan_end_user(context: _KeyBudgetContext) -> tuple[_PlannedBudget, ...]: budget_reset_at=meta.budget_reset_at, notes=( (_END_USER_ROUTE_NOTE, _CUSTOM_AUTH_END_USER_NOTE) - if context.custom_auth_enabled and token_max_budget is None + if context.custom_auth_enabled else (_END_USER_ROUTE_NOTE,) ), ) @@ -1218,20 +1187,17 @@ def _plan_end_user(context: _KeyBudgetContext) -> tuple[_PlannedBudget, ...]: per_model: Final = tuple( _PlannedBudget( scope="end_user_model", - entity_id=model, - entity_label=end_user_id, + entity_id=request_model, + entity_label=cap_key, enforcement="hard", max_budget=_positive_or_none(config.max_budget), comparison=">", - source=f"{row_source}.model_max_budget[{model}]", - spend_source=_EndUserModelSpend( - end_user_id=end_user_id, - models=_model_counter_candidates(config_key=model, configured=configured, context=context), - budget_config=config, - ), + source=f"{row_source}.model_max_budget[{cap_key}]", + spend_source=_EndUserModelSpend(end_user_id=end_user_id, model=request_model, budget_config=config), budget_duration=config.budget_duration, notes=(_MODEL_BUDGET_NOTE,), ) - for model, config in configured.items() + for cap_key, config in configured.items() + for request_model in _model_counter_candidates(config_key=cap_key, configured=configured, context=context) ) return (primary, *per_model) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 0e7b1103a5f..aadffb62320 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -3802,8 +3802,9 @@ async def key_budgets_fn( can be ruled out without opening every object. Parameters: - - key_id: str | None (path parameter) - The key to inspect. Accepts the plaintext key or its - hash. Defaults to the key in the Authorization header when omitted (`GET /key/budgets`). + - key_id: str | None (path parameter) - The hash of the key to inspect. The key itself is + rejected here, because a URL path reaches access logs, tracing spans and error-logging + callbacks. Defaults to the key in the Authorization header when omitted (`GET /key/budgets`). - end_user_id: str | None (query parameter) - Also report the budgets that would apply to this end user. Omitted end users produce no `end_user` rows, because nothing binds an end user to a key outside a request. Proxy admins only, since end users are a proxy-global namespace with @@ -3830,11 +3831,13 @@ async def key_budgets_fn( - source: str - Where the limit is configured, e.g. `key.max_budget`, `budget_table:` - status: str - `unlimited`, `ok` or `exceeded` - notes: list - Caveats worth knowing before trusting the row, each with a stable `code` - to branch on, a `severity` of `info` or `warning`, and human-facing `text` + to branch on, a `severity` of `info` or `warning` for codes a client does not know yet, + and human-facing `text` that is free to be reworded. Ordered most to least specific to + this row's numbers, and empty rather than null when there is nothing to say Example Curl: ``` - curl -X GET "http://0.0.0.0:4000/key/sk-test-example-key-123/budgets" \ + curl -X GET "http://0.0.0.0:4000/key/a1b2c3d4e5f6a7b8c9d0e1f2a3b4c5d6e7f8a9b0c1d2e3f4a5b6c7d8e9f0a1b2/budgets" \ -H "Authorization: Bearer sk-1234" ``` @@ -3867,6 +3870,16 @@ async def key_budgets_fn( ), ) + if key_id is not None and key_id.startswith("sk-"): + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail=( + "Pass the key's hash in the path, not the key itself. A URL path is recorded by access logs, " + "tracing spans and error-logging callbacks, so a key placed there does not stay secret. " + "Call GET /key/budgets with the key in the Authorization header to inspect your own key." + ), + ) + key: Final = key_id or user_api_key_dict.api_key if key is None: raise _key_not_found_error() @@ -3908,11 +3921,6 @@ async def key_budgets_fn( general_settings=general_settings, custom_auth_enabled=user_custom_auth is not None, ), - token_end_user_max_budget=( - user_api_key_dict.end_user_max_budget - if end_user_id is not None and end_user_id == user_api_key_dict.end_user_id - else None - ), ) return KeyBudgetsResponse(key=key, budgets=budgets) except Exception as e: # noqa: BLE001 # every management handler maps unexpected failures onto the proxy error contract diff --git a/litellm/types/proxy/management_endpoints/key_management_endpoints.py b/litellm/types/proxy/management_endpoints/key_management_endpoints.py index cc81742cbde..4e96787524a 100644 --- a/litellm/types/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/types/proxy/management_endpoints/key_management_endpoints.py @@ -152,7 +152,14 @@ BudgetSpendState = Literal["live", "no_counter", "unavailable"] class KeyBudgetNote(BaseModel): - """One caveat about a budget row. Branch on ``code``; ``text`` is free to be reworded.""" + """ + One caveat about a budget row. + + ``code`` is the contract: map it to whatever treatment the caveat deserves. ``text`` is free to be + reworded and must not be matched on. ``severity`` exists for the code a client has not been taught + yet, since this union grows: ``warning`` means the row's numbers may be incomplete or read as + something they are not, and ``info`` means they are accurate and the note is only context. + """ model_config = ConfigDict(frozen=True) diff --git a/tests/test_litellm/proxy/auth/test_info_routes.py b/tests/test_litellm/proxy/auth/test_info_routes.py index 79a73f663e5..b2512607c9f 100644 --- a/tests/test_litellm/proxy/auth/test_info_routes.py +++ b/tests/test_litellm/proxy/auth/test_info_routes.py @@ -194,3 +194,17 @@ def test_templated_route_matching_does_not_widen_unrelated_key_routes(): assert RouteChecks.is_info_route("/key/some-hash/regenerate") is False assert RouteChecks.is_info_route("/key/some-hash/delete") is False assert RouteChecks.is_info_route("/key/some-hash/budgets") is True + + +def test_info_route_matching_is_cached_without_answering_for_the_wrong_route(): + """Pattern matching runs per request off the end-user budget check, so it is memoised; the cache must not blur routes.""" + RouteChecks.is_info_route.cache_clear() + + assert RouteChecks.is_info_route("/key/hash-a/budgets") is True + assert RouteChecks.is_info_route("/key/hash-a/regenerate") is False + assert RouteChecks.is_info_route("/chat/completions") is False + assert RouteChecks.is_info_route("/key/hash-b/budgets") is True + assert RouteChecks.is_info_route.cache_info().hits == 0, "four distinct routes cannot share an answer" + + assert RouteChecks.is_info_route("/key/hash-a/regenerate") is False + assert RouteChecks.is_info_route.cache_info().hits == 1 diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 77b047e06ca..db9b2426298 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -16462,7 +16462,11 @@ from litellm.models.team import LiteLLM_TeamTable # noqa: E402 from litellm.proxy._types import LiteLLM_ProjectTableCachedObj # noqa: E402 from litellm.proxy.auth.auth_checks import TeamMemberBudget # noqa: E402 from litellm.proxy.management_endpoints.key_budget_resolver import ( # noqa: E402 + _match_model_budget_key, _MODEL_BUDGET_COLD_NOTE, + _read_end_user_model_spend, + _read_key_model_spend, + _request_models, _RESERVATION_NOTE, KeyBudgetResolverDeps, resolve_key_budgets, @@ -16470,6 +16474,10 @@ from litellm.proxy.management_endpoints.key_budget_resolver import ( # noqa: E4 from litellm.types.proxy.management_endpoints.key_management_endpoints import ( # noqa: E402 KeyBudgetEntry, ) +from litellm.types.utils import BudgetConfig as _BudgetsBudgetConfig # noqa: E402 +from types import MappingProxyType # noqa: E402 + +BudgetConfig = _BudgetsBudgetConfig _BUDGETS_RESOLVER = "litellm.proxy.management_endpoints.key_budget_resolver" _BUDGETS_KEY_HASH = "hash-of-the-budgets-key" @@ -16516,10 +16524,9 @@ class _RecordingModelSpendReader: self.spend_by_model = default if spend_by_model is None else spend_by_model self.probed = [] - async def __call__(self, *, entity_id, models, budget_config): - self.probed.append(tuple(models)) - found = [self.spend_by_model[model] for model in models if model in self.spend_by_model] - return max(found) if found else None + async def __call__(self, *, entity_id, model, budget_config): + self.probed.append(model) + return self.spend_by_model.get(model) def _budgets_deps(read_spend=None, model_spend=None, match_model_budget_key=None, general_settings=None): @@ -17076,10 +17083,14 @@ async def test_key_budgets_probe_every_request_model_that_maps_onto_a_per_model_ ), ) - entry = next(e for e in budgets if e.scope == "key_model") - assert "openai/gpt-5" in model_reader.probed[0], "the request model's counter was never probed" - assert entry.spend == 6.0 - assert entry.status == "exceeded" + rows = {e.entity_id: e for e in budgets if e.scope == "key_model"} + assert "openai/gpt-5" in model_reader.probed, "the request model's counter was never probed" + assert set(rows) == {"gpt-5", "openai/gpt-5"}, "each counter under the cap needs its own row" + assert all(row.entity_label == "gpt-5" and row.source == "key.model_max_budget[gpt-5]" for row in rows.values()) + assert rows["openai/gpt-5"].spend == 6.0 + assert rows["openai/gpt-5"].status == "exceeded" + assert rows["gpt-5"].spend is None + assert rows["gpt-5"].status == "ok", "a model with no counter of its own is not over the cap" @pytest.mark.asyncio @@ -17141,19 +17152,18 @@ async def test_key_budgets_keep_the_valid_model_caps_when_one_model_entry_is_mal @pytest.mark.asyncio -async def test_key_budgets_prefer_the_request_token_end_user_cap_over_the_row(): - """Reservation reads the token's cap first, so reporting the row's number would understate the block.""" +async def test_key_budgets_report_the_end_user_cap_from_the_budget_row_not_the_calling_request(): + """A request-scoped cap belongs to whoever is calling, so attributing it to the inspected key is a lie.""" with _budgets_world(**_fully_populated_world()): budgets = await resolve_key_budgets( - valid_token=_budgets_token(), + valid_token=_budgets_token(end_user_max_budget=25.0), end_user_id="end-user-budgets", deps=_budgets_deps(), - token_end_user_max_budget=25.0, ) entry = next(e for e in budgets if e.scope == "end_user") - assert entry.max_budget == 25.0 - assert entry.source == "token.end_user_max_budget" + assert entry.max_budget == 800.0 + assert entry.source == "budget_table:budget-end-user" @pytest.mark.asyncio @@ -17230,10 +17240,11 @@ async def test_key_budgets_separate_caveats_that_rule_a_row_out_from_caveats_tha budgets = await _budgets_at_limit() severity_by_code = {note.code: note.severity for entry in budgets for note in entry.notes} - assert severity_by_code["project_spend_not_tracked"] == "info" + assert severity_by_code["project_spend_not_tracked"] == "warning", "a budget that cannot trip is not a live row" + assert severity_by_code["request_tags_add_budgets"] == "warning", "the list of tag budgets is incomplete" + assert severity_by_code["reservation_blocks_at_limit"] == "info", "`comparison` already carries this" + assert severity_by_code["rolling_window"] == "info", "the numbers are right, the window just moves" assert severity_by_code["alert_only"] == "info" - assert severity_by_code["reservation_blocks_at_limit"] == "warning" - assert severity_by_code["rolling_window"] == "warning" @pytest.mark.asyncio @@ -17259,7 +17270,9 @@ async def test_key_budgets_emit_each_caveat_as_its_own_note_instead_of_one_joine "end_user_route_only", "custom_auth_may_override_end_user_cap", ] - assert all("; " not in note.text for note in end_user.notes) + assert all( + other.text not in note.text for note in end_user.notes for other in end_user.notes if other is not note + ), "one note swallowing another is the joined string coming back" @pytest.mark.asyncio @@ -17280,6 +17293,111 @@ async def test_key_budgets_always_send_notes_as_a_list_so_a_client_never_branche ) +@pytest.mark.parametrize("scope", ["team", "tag", "end_user"]) +@pytest.mark.asyncio +async def test_key_budgets_do_not_invent_a_denial_for_a_zero_cap(scope): + """The reservation builds no counter for a non-positive cap, so the read-time `>` is still the gate.""" + world = _fully_populated_world() + world["team"] = LiteLLM_TeamTable(team_id="team-budgets", team_alias="Reporting Team", spend=0.0, max_budget=0.0) + world["tags"] = { + "prod": LiteLLM_TagTable( + tag_name="prod", + spend=0.0, + budget_id="budget-tag", + litellm_budget_table=LiteLLM_BudgetTable(budget_id="budget-tag", max_budget=0.0), + ) + } + world["end_user"] = LiteLLM_EndUserTable( + user_id="end-user-budgets", + blocked=False, + alias="End User", + spend=0.0, + budget_id="budget-end-user", + litellm_budget_table=LiteLLM_BudgetTable(budget_id="budget-end-user", max_budget=0.0), + ) + with _budgets_world(**world): + budgets = await resolve_key_budgets( + valid_token=_budgets_token(), end_user_id="end-user-budgets", deps=_budgets_deps() + ) + + entry = next(e for e in budgets if e.scope == scope and e.enforcement == "hard") + assert entry.max_budget == 0.0, scope + assert entry.comparison == ">", scope + assert entry.status == "ok", scope + assert _RESERVATION_NOTE not in entry.notes, scope + + +@pytest.mark.asyncio +async def test_key_budgets_still_tighten_a_positive_cap_in_the_same_scope(): + """The zero-cap carve-out must not disarm the operator fix for every cap that does reserve.""" + budgets = await _budgets_at_limit() + + team = next(e for e in budgets if e.scope == "team" and e.enforcement == "hard") + assert team.comparison == ">=" + assert _RESERVATION_NOTE in team.notes + + +@pytest.mark.asyncio +async def test_key_budgets_probe_a_deployment_routed_at_directly(): + """Routing straight at a deployment keys the counter on its name, which is not a model group.""" + router = MagicMock() + router.get_model_names.return_value = ["gpt-5"] + router.deployment_names = ["azure/gpt-5-prod"] + with patch("litellm.proxy.proxy_server.llm_router", router): + assert _request_models(()) == ("gpt-5", "azure/gpt-5-prod") + + with patch("litellm.proxy.proxy_server.llm_router", None): + assert _request_models(("only-the-key-models",)) == ("only-the-key-models",) + + +@pytest.mark.asyncio +async def test_key_budgets_read_each_per_model_counter_from_the_enforcing_limiter(): + """These readers exist to reuse enforcement's cache lookup; a fake in every other test hides a wrong call.""" + limiter = MagicMock() + limiter.get_virtual_key_spend_for_model = AsyncMock(return_value=4.0) + limiter.get_end_user_spend_for_model = AsyncMock(return_value=9.0) + config = BudgetConfig(max_budget=5.0, budget_duration="1d") + with patch("litellm.proxy.proxy_server.model_max_budget_limiter", limiter): + key_spend = await _read_key_model_spend(entity_id="hash-1", model="openai/gpt-5", budget_config=config) + end_user_spend = await _read_end_user_model_spend(entity_id="cust-1", model="gpt-5", budget_config=config) + + assert key_spend == 4.0 + assert limiter.get_virtual_key_spend_for_model.await_args.kwargs == { + "user_api_key_hash": "hash-1", + "model": "openai/gpt-5", + "key_budget_config": config, + } + assert end_user_spend == 9.0 + assert limiter.get_end_user_spend_for_model.await_args.kwargs == { + "end_user_id": "cust-1", + "model": "gpt-5", + "key_budget_config": config, + } + + +@pytest.mark.parametrize( + ("model", "expected"), + [("gpt-5", "gpt-5"), ("openai/gpt-5", "gpt-5"), ("claude-sonnet-4-5", None)], +) +def test_key_budgets_match_a_request_model_to_a_cap_the_way_enforcement_does(model, expected): + """Introspection and enforcement must agree on which cap a request model is charged against.""" + configured = MappingProxyType({"gpt-5": BudgetConfig(max_budget=5.0, budget_duration="1d")}) + + assert _match_model_budget_key(model=model, configured=configured) == expected + + +@pytest.mark.asyncio +async def test_key_budgets_route_refuses_a_plaintext_key_in_the_url_path(): + """An info route's path reaches error-logging callbacks and access logs, so a key must not sit in it.""" + caller = UserAPIKeyAuth(api_key="sk-admin", user_role=LitellmUserRoles.PROXY_ADMIN.value) + with _budgets_route_world(key_row=_budgets_key_row(), caller=caller) as resolver: + response = client.get("/key/sk-a-real-looking-secret/budgets") + + assert response.status_code == 400 + assert "hash" in json.dumps(response.json()).lower() + resolver.assert_not_awaited() + + @pytest.mark.parametrize( "role", [