mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-07 02:59:05 +00:00
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.
This commit is contained in:
ryan-crabbe-berri
parent
af0c016b4f
commit
83ef72a76a
7 changed files with 233 additions and 113 deletions
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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:<id>`
|
||||
- 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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
[
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue