scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-07 08:26:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

fix(auth): keep a key in the path out of failure logging, and report one budget row per counter

Browse source 
The 400 rejecting a plaintext key in GET /key/{key_id}/budgets sits in the handler, which runs
after Depends(user_api_key_auth), so it never fires on the path that leaks: a wrong or expired
Authorization header rejects the request first, and the failure handler stamps the raw route onto
the server span and hands it to post_call_failure_hook, where it becomes call_type on the payload
every failure callback receives plus the OTEL llm.request.type and gen_ai.operation.name
attributes. Normalize the route instead, which is where the success path already collapsed it, and
add the three key management routes that take a key in the path to the placeholder list. Redaction
runs ahead of the memoized pass so no cache retains a live credential as a cache key. The 400
stays. Both existing routes, /key/{key_id}/regenerate and /key/{key_id}/reset_spend, had the same
hole on the span and are covered by the same change.

Per-model rows were one row per request model routed to a cap, but both readers fall back to the
provider-stripped model, so a single counter under gpt-5 was reported once each for openai/gpt-5,
azure/gpt-5 and bedrock/gpt-5, all reading the same balance and all claiming spend_state "live". A
$12 counter under a $12 cap rendered as four rows totalling $48, three of them asserting a counter
that does not exist. The lookup now reports which counter it landed on, and rows collapse on that,
so one counter is one row named after the counter rather than after whichever request model probed
it first. Two models with counters of their own still get a row each.

model_budget_fails_open always travelled with spend_state "no_counter", which made a cap whose
cache is merely cold indistinguishable from project_spend_not_tracked, a budget that can never
trip. spend_state already carries "cold", so the code is gone and the fails-open fact moves onto
per_model_counters, which rides every per-model row.

custom_auth_skips_read_time_checks is new: a custom auth callable returns its own token before any
of these checks run, and the wrapper skips common_checks too unless custom_auth_run_common_checks
is set, so on such a proxy none of these budgets are enforced at request time. Saying that only on
the end-user row understated it by a dozen scopes.

Router.deployment_names is appended to and never pruned, so a deleted deployment kept producing a
row; read the same names off the live model_list instead.
This commit is contained in:
ryan-crabbe-berri 2026-08-19 19:06:20 -07:00
parent ae60ef5a07
commit 36904f55ad
12 changed files with 713 additions and 95 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
litellm/proxy/auth/auth_exception_handler.py
View file

@ -15,7 +15,7 @@ from litellm.proxy._types import (
ProxyException,
UserAPIKeyAuth,
)
from litellm.proxy.auth.auth_utils import _get_request_ip_address
from litellm.proxy.auth.auth_utils import _get_request_ip_address, normalize_request_route
from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
from litellm.types.services import ServiceTypes
@ -86,7 +86,7 @@ class UserAPIKeyAuthExceptionHandler:
token="failed-to-connect-to-db",
user_id=DB_UNAVAILABLE_FALLBACK_USER_ID,
user_role=LitellmUserRoles.INTERNAL_USER,
request_route=route,
request_route=normalize_request_route(route),
)
else:
# raise the exception to the caller
@ -108,7 +108,7 @@ class UserAPIKeyAuthExceptionHandler:
# so the handler is side-effect-free for the caller's identity object.
user_api_key_dict = resolved_identity.model_copy() if resolved_identity is not None else UserAPIKeyAuth()
user_api_key_dict.parent_otel_span = parent_otel_span
user_api_key_dict.request_route = route
user_api_key_dict.request_route = normalize_request_route(route)
user_api_key_dict.api_key = user_api_key_dict.api_key or UserAPIKeyAuth(api_key=api_key).api_key
# Stamp identity onto the request's server span now, before the request

25
litellm/proxy/auth/auth_utils.py
View file

@ -661,7 +661,18 @@ def get_request_route_template(request: Request) -> str | None:
return None
@lru_cache(maxsize=256)
_KEY_PATH_PARAM_ROUTE: Final = re.compile(r"^/key/[^/]+/(budgets|regenerate|reset_spend)$")
def _redact_key_path_param(route: str) -> str:
"""
Replace the key in a ``/key/{key_id}/...`` path, which is a key hash or a plaintext key.
Kept out of the memoized pass below so that no cache ever retains a live credential as a cache key.
"""
return _KEY_PATH_PARAM_ROUTE.sub(r"/key/{key_id}/\1", route)
def normalize_request_route(route: str) -> str:
"""
Normalize request routes by replacing dynamic path parameters with placeholders.
@ -670,6 +681,10 @@ def normalize_request_route(route: str) -> str:
- /v1/responses/1234567890 -> /v1/responses/{response_id}
- /v1/threads/thread_123 -> /v1/threads/{thread_id}
It is also what keeps a secret out of any consumer of a request route: the key management routes
that take a key in the path are collapsed here, so a route recorded on a span, a metric label or a
logging callback carries the placeholder rather than the credential.
Args:
route: The request route path
@ -681,9 +696,17 @@ def normalize_request_route(route: str) -> str:
'/v1/responses/{response_id}'
>>> normalize_request_route("/v1/responses/abc123/cancel")
'/v1/responses/{response_id}/cancel'
>>> normalize_request_route("/key/sk-1234/budgets")
'/key/{key_id}/budgets'
>>> normalize_request_route("/chat/completions")
'/chat/completions'
"""
return _normalize_known_id_routes(_redact_key_path_param(route))
@lru_cache(maxsize=256)
def _normalize_known_id_routes(route: str) -> str:
"""Collapse the resource-id path params of the OpenAI-shaped routes."""
# Define patterns for routes with dynamic IDs
# Format: (regex_pattern, replacement_template)
patterns: Final = [

90
litellm/proxy/hooks/model_max_budget_limiter.py
View file

@ -1,6 +1,7 @@
import json
from collections.abc import Mapping
from typing import Final
from dataclasses import dataclass
from typing import Final, Literal
import litellm
from litellm._logging import verbose_proxy_logger
@ -18,6 +19,26 @@ from litellm.types.utils import (
VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX: Final = "virtual_key_spend"
END_USER_SPEND_CACHE_KEY_PREFIX: Final = "end_user_model_spend"
ModelSpendScope = Literal["virtual_key", "end_user"]
_SPEND_CACHE_KEY_PREFIX: Final[Mapping[ModelSpendScope, str]] = {
"virtual_key": VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX,
"end_user": END_USER_SPEND_CACHE_KEY_PREFIX,
}
@dataclass(frozen=True, slots=True)
class ModelSpendHit:
"""
Which model's counter a lookup landed on, and what it held.
A request model falls back to its provider-stripped form, so several request models can share one
counter. Callers that report spend rather than enforce it need to know which, or they double-count.
"""
model: str
spend: float
class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting):
"""
@ -149,27 +170,40 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting):
return True
async def resolve_model_spend(
self,
*,
scope: ModelSpendScope,
entity_id: str | None,
model: str,
budget_config: BudgetConfig,
) -> ModelSpendHit | None:
"""
The counter holding this model's spend, looked up the way enforcement looks it up.
Order is the request model first, then the same name with its ``{custom_llm_provider}/`` prefix
removed, so a counter written under ``gpt-4o`` still gates a request for ``openai/gpt-4o``.
"""
prefix: Final = _SPEND_CACHE_KEY_PREFIX[scope]
candidates: Final = dict.fromkeys((model, self._get_model_without_custom_llm_provider(model)))
for candidate in candidates:
spend = await self.dual_cache.async_get_cache(
key=f"{prefix}:{entity_id}:{candidate}:{budget_config.budget_duration}",
)
if spend is not None:
return ModelSpendHit(model=candidate, spend=spend)
return None
async def get_end_user_spend_for_model(
self,
end_user_id: str,
model: str,
key_budget_config: BudgetConfig,
) -> float | None:
# 1. model: directly look up `model`
end_user_model_spend_cache_key = (
f"{END_USER_SPEND_CACHE_KEY_PREFIX}:{end_user_id}:{model}:{key_budget_config.budget_duration}"
hit: Final = await self.resolve_model_spend(
scope="end_user", entity_id=end_user_id, model=model, budget_config=key_budget_config
)
_current_spend = await self.dual_cache.async_get_cache(
key=end_user_model_spend_cache_key,
)
if _current_spend is None:
# 2. If 1, does not exist, check if passed as {custom_llm_provider}/model
end_user_model_spend_cache_key = f"{END_USER_SPEND_CACHE_KEY_PREFIX}:{end_user_id}:{self._get_model_without_custom_llm_provider(model)}:{key_budget_config.budget_duration}"
_current_spend = await self.dual_cache.async_get_cache(
key=end_user_model_spend_cache_key,
)
return _current_spend
return hit.spend if hit is not None else None
async def get_virtual_key_spend_for_model(
self,
@ -177,30 +211,10 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting):
model: str,
key_budget_config: BudgetConfig,
) -> float | None:
"""
Get the current spend for a virtual key for a model
Lookup model in this order:
1. model: directly look up `model`
2. If 1, does not exist, check if passed as {custom_llm_provider}/model
"""
# 1. model: directly look up `model`
virtual_key_model_spend_cache_key = (
f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:{user_api_key_hash}:{model}:{key_budget_config.budget_duration}"
hit: Final = await self.resolve_model_spend(
scope="virtual_key", entity_id=user_api_key_hash, model=model, budget_config=key_budget_config
)
_current_spend = await self.dual_cache.async_get_cache(
key=virtual_key_model_spend_cache_key,
)
if _current_spend is None:
# 2. If 1, does not exist, check if passed as {custom_llm_provider}/model
# if "/" in model, remove first part before "/" - eg. openai/o1-preview -> o1-preview
virtual_key_model_spend_cache_key = f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:{user_api_key_hash}:{self._get_model_without_custom_llm_provider(model)}:{key_budget_config.budget_duration}"
_current_spend = await self.dual_cache.async_get_cache(
key=virtual_key_model_spend_cache_key,
)
return _current_spend
return hit.spend if hit is not None else None
def get_request_model_budget_key(
self, model: str, internal_model_max_budget: Mapping[str, BudgetConfig]

124
litellm/proxy/management_endpoints/key_budget_resolver.py
View file

@ -48,6 +48,7 @@ from litellm.proxy.auth.auth_checks import (
user_budget_applies_to_key,
)
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.hooks.model_max_budget_limiter import ModelSpendHit
from litellm.proxy.spend_tracking.budget_reservation import get_budget_window_start
from litellm.proxy.spend_tracking.spend_counter_keys import (
end_user_spend_counter,
@ -92,6 +93,8 @@ _ENTITY_TYPE_BY_SCOPE: Final[Mapping[BudgetScope, Litellm_EntityType]] = Mapping
}
)
_MODEL_SCOPES: Final[frozenset[BudgetScope]] = frozenset({"key_model", "end_user_model"})
_RESERVATION_COVERED_SCOPES: Final[frozenset[BudgetScope]] = frozenset(
{"key", "key_window", "team", "team_window", "team_member", "user", "organization", "tag", "end_user"}
)
@ -110,15 +113,11 @@ _MODEL_BUDGET_NOTE: Final = KeyBudgetNote(
code="per_model_counters",
severity="warning",
text=(
"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"
"one row per counter behind this cap, because each counter is compared against it alone; these "
"counters are cache-only and fail open while one is missing, and a model this proxy cannot "
"enumerate has a counter that is not reported here"
),
)
_MODEL_BUDGET_COLD_NOTE: Final = KeyBudgetNote(
code="model_budget_fails_open",
severity="info",
text="no per-model counter exists yet; these budgets are cache-only and fail open until one does",
)
_RESERVATION_NOTE: Final = KeyBudgetNote(
code="reservation_blocks_at_limit",
severity="info",
@ -147,6 +146,15 @@ _CUSTOM_AUTH_END_USER_NOTE: Final = KeyBudgetNote(
severity="warning",
text="a custom auth callable can set a request-scoped end user cap that overrides this one and is not visible here",
)
_CUSTOM_AUTH_SKIPS_CHECKS_NOTE: Final = KeyBudgetNote(
code="custom_auth_skips_read_time_checks",
severity="warning",
text=(
"a custom auth callable authenticates requests on this proxy and "
"general_settings.custom_auth_run_common_checks is not set, so none of these budgets are checked "
"at request time for the requests it authenticates; only the reservation layer still applies"
),
)
_USER_ON_TEAM_KEY_NOTE: Final = KeyBudgetNote(
code="user_budget_not_applied_to_team_key",
severity="info",
@ -180,7 +188,7 @@ class SpendReader(Protocol):
class ModelSpendReader(Protocol):
async def __call__(self, *, entity_id: str, model: str, budget_config: BudgetConfig) -> float | None: ...
async def __call__(self, *, entity_id: str, model: str, budget_config: BudgetConfig) -> ModelSpendHit | None: ...
class ModelBudgetKeyMatcher(Protocol):
@ -210,19 +218,21 @@ async def _read_counter_spend(
)
async def _read_key_model_spend(*, entity_id: str, model: str, budget_config: BudgetConfig) -> float | None:
async def _read_key_model_spend(*, entity_id: str, model: str, budget_config: BudgetConfig) -> ModelSpendHit | None:
from litellm.proxy.proxy_server import model_max_budget_limiter
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
return await model_max_budget_limiter.resolve_model_spend(
scope="virtual_key", entity_id=entity_id, model=model, budget_config=budget_config
)
async def _read_end_user_model_spend(*, entity_id: str, model: str, budget_config: BudgetConfig) -> float | None:
async def _read_end_user_model_spend(
*, entity_id: str, model: str, budget_config: BudgetConfig
) -> ModelSpendHit | None:
from litellm.proxy.proxy_server import model_max_budget_limiter
return await model_max_budget_limiter.get_end_user_spend_for_model(
end_user_id=entity_id, model=model, key_budget_config=budget_config
return await model_max_budget_limiter.resolve_model_spend(
scope="end_user", entity_id=entity_id, model=model, budget_config=budget_config
)
@ -232,6 +242,15 @@ def _match_model_budget_key(*, model: str, configured: Mapping[str, BudgetConfig
return model_max_budget_limiter.get_request_model_budget_key(model=model, internal_model_max_budget=configured)
def _deployment_models(model_list: Sequence[object]) -> tuple[str, ...]:
"""
``Router.deployment_names`` is appended to and never pruned, so it names deployments that were
deleted. ``model_list`` is the live registry, so read the same names back out of it instead.
"""
parsed: Final = (_validated(_DEPLOYMENT_ENTRY, entry, "model_list entry") for entry in model_list)
return tuple(entry.litellm_params.model for entry in parsed if entry is not None)
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.
@ -243,7 +262,7 @@ def _request_models(allowed_models: tuple[str, ...]) -> tuple[str, ...]:
if llm_router is None:
return allowed_models
routable: Final = (*llm_router.get_model_names(), *llm_router.deployment_names)
routable: Final = (*llm_router.get_model_names(), *_deployment_models(llm_router.model_list or ()))
return tuple(dict.fromkeys((*allowed_models, *routable)))
@ -298,6 +317,7 @@ class _SpendReading:
value: float | None
state: BudgetSpendState
counter_model: str | None = None
@dataclass(frozen=True, slots=True)
@ -316,6 +336,9 @@ class _PlannedBudget:
notes: tuple[KeyBudgetNote, ...] = ()
_ReadPlan = tuple[_PlannedBudget, _SpendReading]
class _MetadataTags(BaseModel):
model_config = ConfigDict(extra="ignore")
@ -348,11 +371,24 @@ class _ModelBudgetFields(BaseModel):
model_max_budget: Mapping[str, object] = MappingProxyType({})
class _DeploymentParamsFields(BaseModel):
model_config = ConfigDict(extra="ignore", protected_namespaces=())
model: str
class _DeploymentEntryFields(BaseModel):
model_config = ConfigDict(extra="ignore")
litellm_params: _DeploymentParamsFields
_T = TypeVar("_T")
_METADATA_FIELDS: Final = TypeAdapter(_MetadataFields)
_WINDOW_FIELDS: Final = TypeAdapter(_WindowFields)
_MODEL_BUDGET_FIELDS: Final = TypeAdapter(_ModelBudgetFields)
_DEPLOYMENT_ENTRY: Final = TypeAdapter(_DeploymentEntryFields)
_BUDGET_LIMIT_ENTRY: Final = TypeAdapter(BudgetLimitEntry)
_BUDGET_CONFIG: Final = TypeAdapter(BudgetConfig)
_KEY_MODELS: Final = TypeAdapter(_KeyModelsFields)
@ -440,6 +476,7 @@ class _KeyBudgetContext:
token_inputs: _TokenBudgetInputs
end_user_id: str | None
custom_auth_enabled: bool
custom_auth_skips_checks: bool
request_models: tuple[str, ...]
match_model_budget_key: ModelBudgetKeyMatcher
general_settings: Mapping[str, object]
@ -465,15 +502,45 @@ async def resolve_key_budgets(
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
proxy_notes: Final = (_CUSTOM_AUTH_SKIPS_CHECKS_NOTE,) if context.custom_auth_skips_checks else ()
return tuple(
_to_entry(plan=plan, reading=reading, reservation_enabled=reservation_enabled)
for plan, reading in zip(plans, readings, strict=True)
_to_entry(
plan=plan,
reading=reading,
reservation_enabled=reservation_enabled,
proxy_notes=proxy_notes,
)
for plan, reading in _collapse_shared_counters(tuple(zip(plans, readings, strict=True)))
)
def _model_reading(spend: float | None) -> _SpendReading:
def _counter_identity(index: int, plan: _PlannedBudget, reading: _SpendReading) -> object:
"""Two rows share an identity only when enforcement compares both against the same counter."""
if plan.scope not in _MODEL_SCOPES:
return index
return (plan.source, reading.counter_model)
def _collapse_shared_counters(pairs: tuple[_ReadPlan, ...]) -> tuple[_ReadPlan, ...]:
"""
One counter, one row.
A per-model cap is planned once per request model routed to it, but the reader falls back to the
provider-stripped name, so several of those models read one counter. Reporting each of them would
show a single balance several times over as though it were several balances.
"""
identities: Final = tuple(
_counter_identity(index=index, plan=plan, reading=reading) for index, (plan, reading) in enumerate(pairs)
)
first_index: Final = {identity: index for index, identity in reversed(tuple(enumerate(identities)))}
return tuple(pair for index, pair in enumerate(pairs) if first_index[identities[index]] == index)
def _model_reading(hit: ModelSpendHit | None) -> _SpendReading:
"""Per-model budgets are cache-only, so a missing counter means untouched rather than unreadable."""
return _SpendReading(value=spend, state="live" if spend is not None else "no_counter")
if hit is None:
return _SpendReading(value=None, state="no_counter")
return _SpendReading(value=hit.spend, state="live", counter_model=hit.model)
async def _read_spend(plan: _PlannedBudget, deps: KeyBudgetResolverDeps) -> _SpendReading:
@ -532,14 +599,18 @@ def _effective_comparison(plan: _PlannedBudget, reservation_enabled: bool) -> Bu
def _entry_notes(
plan: _PlannedBudget, reading: _SpendReading, comparison: BudgetComparison
plan: _PlannedBudget, comparison: BudgetComparison, proxy_notes: tuple[KeyBudgetNote, ...]
) -> tuple[KeyBudgetNote, ...]:
tightened: Final = () if comparison == plan.comparison else (_RESERVATION_NOTE,)
cold: Final = (_MODEL_BUDGET_COLD_NOTE,) if reading.state == "no_counter" else ()
return (*tightened, *plan.notes, *cold)
return (*tightened, *plan.notes, *proxy_notes)
def _to_entry(plan: _PlannedBudget, reading: _SpendReading, reservation_enabled: bool) -> KeyBudgetEntry:
def _to_entry(
plan: _PlannedBudget,
reading: _SpendReading,
reservation_enabled: bool,
proxy_notes: tuple[KeyBudgetNote, ...],
) -> KeyBudgetEntry:
comparison: Final = _effective_comparison(plan=plan, reservation_enabled=reservation_enabled)
spend: Final = reading.value
exceeded: Final = (
@ -551,7 +622,7 @@ def _to_entry(plan: _PlannedBudget, reading: _SpendReading, reservation_enabled:
return KeyBudgetEntry(
scope=plan.scope,
entity_type=_ENTITY_TYPE_BY_SCOPE[plan.scope],
entity_id=plan.entity_id,
entity_id=reading.counter_model or plan.entity_id,
entity_label=plan.entity_label,
enforcement=plan.enforcement,
max_budget=plan.max_budget,
@ -564,7 +635,7 @@ def _to_entry(plan: _PlannedBudget, reading: _SpendReading, reservation_enabled:
window_start=plan.window_start,
source=plan.source,
status=status,
notes=_entry_notes(plan=plan, reading=reading, comparison=comparison),
notes=_entry_notes(plan=plan, comparison=comparison, proxy_notes=proxy_notes),
)
@ -600,6 +671,9 @@ async def _load_context(
token_inputs=token_inputs,
end_user_id=end_user_id,
custom_auth_enabled=deps.custom_auth_enabled,
custom_auth_skips_checks=(
deps.custom_auth_enabled and deps.general_settings.get("custom_auth_run_common_checks") is not True
),
request_models=_request_models(token_inputs.models),
match_model_budget_key=deps.match_model_budget_key,
general_settings=deps.general_settings,

6
litellm/proxy/management_endpoints/key_management_endpoints.py
View file

@ -3817,7 +3817,11 @@ async def key_budgets_fn(
`team_member`, `user`, `organization`, `project`, `tag`, `end_user` or `end_user_model`
- entity_type: Litellm_EntityType - The entity a `BudgetExceededError` from this scope
names, so a denial message maps back to a row here
- entity_id / entity_label: str | None - Which entity is limited, and its human-facing alias
- entity_id / entity_label: str | None - Which entity is limited, and its human-facing alias.
On the per-model scopes this is one row per counter rather than per request model, so
`entity_id` is the model whose counter was read and `entity_label` is the configured cap
it is compared against; several request models can share one counter, and they are not
listed separately because their spend is not separate
- enforcement: str - `hard` blocks the request, `soft` only raises an alert, `throttled`
scales the key's rate limits down instead of denying anything. Only the key's own
`max_budget` can be `throttled`; every other scope on the same key still blocks

6
litellm/proxy/utils.py
View file

@ -104,6 +104,7 @@ from litellm.proxy._types import (
Member,
UserAPIKeyAuth,
)
from litellm.proxy.auth.auth_utils import normalize_request_route
from litellm.proxy.auth.route_checks import RouteChecks
from litellm.proxy.common_utils.config_sync_pubsub import publish_config_param_change
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
@ -2413,7 +2414,10 @@ class ProxyLogging:
)
litellm_logging_obj, data = litellm.utils.function_setup(
original_function=route or "IGNORE_THIS",
# The route becomes `call_type` on the payload every failure callback receives, and the
# OTEL span attributes derived from it, so it must not be the raw path: the key
# management routes carry a key in the path and this hook runs on the 401 they fail with.
original_function=normalize_request_route(route) if route else "IGNORE_THIS",
rules_obj=litellm.utils.Rules(),
start_time=datetime.now(),
**request_data,

2
litellm/types/proxy/management_endpoints/key_management_endpoints.py
View file

@ -135,8 +135,8 @@ BudgetStatus = Literal["unlimited", "ok", "exceeded"]
BudgetNoteCode = Literal[
"alert_only",
"custom_auth_may_override_end_user_cap",
"custom_auth_skips_read_time_checks",
"end_user_route_only",
"model_budget_fails_open",
"per_model_counters",
"project_spend_not_tracked",
"request_tags_add_budgets",

105
tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py
View file

@ -558,3 +558,108 @@ async def test_async_log_success_event_skips_redis_push_without_redis(budget_lim
kwargs, response_obj=None, start_time=None, end_time=None
)
mock_push.assert_not_awaited()
@pytest.mark.asyncio
@pytest.mark.parametrize(
("stored_under", "requested", "expected_model"),
[
("gpt-4", "gpt-4", "gpt-4"),
("gpt-4", "openai/gpt-4", "gpt-4"),
("openai/gpt-4", "openai/gpt-4", "openai/gpt-4"),
],
)
async def test_resolve_model_spend_reports_which_counter_it_landed_on(
budget_limiter, stored_under, requested, expected_model
):
"""
The lookup falls back to the provider-stripped name, so several request models share one counter.
Anything reporting spend rather than enforcing it has to know which, or it shows one balance once
per request model and appears to multiply the spend.
"""
from litellm.proxy.hooks.model_max_budget_limiter import (
ModelSpendHit,
VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX,
)
budget_config = GenericBudgetInfo(max_budget=100.0, budget_duration="1d")
await budget_limiter.dual_cache.async_set_cache(
key=f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:test-key:{stored_under}:1d", value=12.0
)
hit = await budget_limiter.resolve_model_spend(
scope="virtual_key", entity_id="test-key", model=requested, budget_config=budget_config
)
assert hit == ModelSpendHit(model=expected_model, spend=12.0)
@pytest.mark.asyncio
async def test_resolve_model_spend_returns_nothing_when_neither_candidate_has_a_counter(budget_limiter):
"""A missing counter is not a zero balance: these budgets fail open until something writes one."""
budget_config = GenericBudgetInfo(max_budget=100.0, budget_duration="1d")
assert (
await budget_limiter.resolve_model_spend(
scope="virtual_key", entity_id="test-key", model="openai/gpt-4", budget_config=budget_config
)
is None
)
@pytest.mark.asyncio
async def test_resolve_model_spend_reads_the_end_user_counters_from_their_own_prefix(budget_limiter):
"""The two scopes must not read each other's counters, which sharing one lookup makes easy to do."""
from litellm.proxy.hooks.model_max_budget_limiter import (
END_USER_SPEND_CACHE_KEY_PREFIX,
ModelSpendHit,
VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX,
)
budget_config = GenericBudgetInfo(max_budget=100.0, budget_duration="1d")
await budget_limiter.dual_cache.async_set_cache(
key=f"{END_USER_SPEND_CACHE_KEY_PREFIX}:cust-1:gpt-4:1d", value=3.0
)
await budget_limiter.dual_cache.async_set_cache(
key=f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:cust-1:gpt-4:1d", value=99.0
)
hit = await budget_limiter.resolve_model_spend(
scope="end_user", entity_id="cust-1", model="gpt-4", budget_config=budget_config
)
assert hit == ModelSpendHit(model="gpt-4", spend=3.0)
@pytest.mark.asyncio
async def test_the_public_spend_readers_still_return_a_bare_float(budget_limiter):
"""Both wrappers are called from enforcement, so widening the shared lookup must not widen them."""
from litellm.proxy.hooks.model_max_budget_limiter import (
END_USER_SPEND_CACHE_KEY_PREFIX,
VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX,
)
budget_config = GenericBudgetInfo(max_budget=100.0, budget_duration="1d")
await budget_limiter.dual_cache.async_set_cache(
key=f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:test-key:gpt-4:1d", value=7.0
)
await budget_limiter.dual_cache.async_set_cache(
key=f"{END_USER_SPEND_CACHE_KEY_PREFIX}:cust-1:gpt-4:1d", value=8.0
)
assert (
await budget_limiter.get_virtual_key_spend_for_model(
user_api_key_hash="test-key", model="openai/gpt-4", key_budget_config=budget_config
)
== 7.0
)
assert (
await budget_limiter.get_end_user_spend_for_model(
end_user_id="cust-1", model="openai/gpt-4", key_budget_config=budget_config
)
== 8.0
)
assert (
await budget_limiter.get_virtual_key_spend_for_model(
user_api_key_hash="test-key", model="claude-sonnet-4-5", key_budget_config=budget_config
)
is None
)

63
tests/test_litellm/proxy/auth/test_auth_exception_handler.py
View file

@ -511,3 +511,66 @@ async def test_auth_failure_without_resolved_identity_still_logs():
assert logged.api_key != "sk-unknown"
assert logged.api_key == UserAPIKeyAuth(api_key="sk-unknown").api_key
assert logged.request_route == "/v1/chat/completions"
@pytest.mark.asyncio
@pytest.mark.parametrize(
"route",
[
"/key/sk-victim-1234567890/budgets",
"/key/" + "a" * 64 + "/budgets",
"/key/sk-victim-1234567890/regenerate",
"/key/sk-victim-1234567890/reset_spend",
],
)
async def test_auth_failure_on_a_key_route_does_not_hand_the_key_to_the_failure_callbacks(route):
"""
Auth runs as a dependency, so a wrong Authorization header rejects the request before the handler
that guards these routes ever runs. The route it fails on is stamped onto the server span and
handed to every failure callback, and on these routes the path param is a key. Normalizing it here
is what keeps that credential out of traces and third-party logging sinks.
"""
handler = UserAPIKeyAuthExceptionHandler()
with patch(
"litellm.proxy.proxy_server.proxy_logging_obj.post_call_failure_hook",
new_callable=AsyncMock,
) as mock_post_call_failure_hook, patch(
"litellm.proxy.proxy_server.general_settings",
{"allow_requests_on_db_unavailable": False},
):
with pytest.raises(Exception):
await handler._handle_authentication_error(
HTTPException(status_code=401, detail="Invalid proxy server token passed"),
MagicMock(),
{},
route,
None,
"sk-wrong-key",
)
recorded = mock_post_call_failure_hook.call_args[1]["user_api_key_dict"].request_route
assert "sk-victim" not in recorded, f"the caller's key reached the failure callbacks as {recorded!r}"
assert "a" * 64 not in recorded, f"the key hash reached the failure callbacks as {recorded!r}"
assert recorded.startswith("/key/{key_id}/"), recorded
@pytest.mark.asyncio
async def test_auth_failure_on_a_db_outage_fallback_also_redacts_the_key_route():
"""The DB-unavailable branch builds its own token and stamps the same route onto it."""
handler = UserAPIKeyAuthExceptionHandler()
with patch(
"litellm.proxy.proxy_server.general_settings",
{"allow_requests_on_db_unavailable": True},
):
result = await handler._handle_authentication_error(
EngineConnectionError(),
MagicMock(),
{},
"/key/sk-victim-1234567890/budgets",
None,
"sk-wrong-key",
)
assert result.request_route == "/key/{key_id}/budgets"

94
tests/test_litellm/proxy/auth/test_auth_utils.py
View file

@ -3200,3 +3200,97 @@ class TestIsRequestBodySafeBlocksAwsIdentitySelectors:
)
is True
)
class TestNormalizeRequestRouteRedactsKeys:
"""
The key management routes that take a path param take a key hash, or a plaintext key when a caller
passes one. Every consumer of a normalized route records it somewhere durable, so the placeholder
has to be substituted here rather than at each of those consumers.
"""
@pytest.mark.parametrize("suffix", ["budgets", "regenerate", "reset_spend"])
@pytest.mark.parametrize(
"key",
["sk-1234567890abcdef", "a" * 64, "sk-Ab_9-xY", "not-really-a-key"],
)
def test_key_path_param_is_replaced_whatever_the_key_looks_like(self, suffix, key):
from litellm.proxy.auth.auth_utils import normalize_request_route
assert normalize_request_route(f"/key/{key}/{suffix}") == f"/key/{{key_id}}/{suffix}"
@pytest.mark.parametrize(
"route",
[
"/key/info",
"/key/budgets",
"/key/list",
"/key/sk-123/budgets/extra",
"/keyring/sk-123/budgets",
"/chat/completions",
],
)
def test_routes_without_a_key_path_param_are_left_alone(self, route):
from litellm.proxy.auth.auth_utils import normalize_request_route
assert normalize_request_route(route) == route
def test_the_memoized_pass_never_receives_a_live_key(self):
"""
An lru_cache holds its arguments for the life of the process, so a raw key reaching it would sit
in memory long after the request. Redaction runs first, which also means many distinct keys
collapse onto one cache entry instead of evicting everything else.
"""
from litellm.proxy.auth.auth_utils import _normalize_known_id_routes, normalize_request_route
_normalize_known_id_routes.cache_clear()
before = _normalize_known_id_routes.cache_info().currsize
for index in range(50):
normalize_request_route(f"/key/sk-distinct-key-{index}/budgets")
assert _normalize_known_id_routes.cache_info().currsize == before + 1
def test_normalizing_does_not_change_how_a_route_is_classified(self):
"""
Failure logging decides whether to run by matching the route, and it now matches the normalized
form. A pattern that stopped matching once its ids were replaced would silently switch off error
logging for that endpoint, which is a far quieter regression than the leak this fixes.
"""
from litellm.proxy.auth.auth_utils import normalize_request_route
from litellm.proxy.auth.route_checks import RouteChecks
routes = [
"/chat/completions",
"/v1/responses/resp_abc",
"/v1/responses/resp_abc/cancel",
"/v1/responses/resp_abc/input_items",
"/v1/threads/th_1",
"/v1/threads/th_1/messages",
"/v1/threads/th_1/runs/run_1/steps/step_1",
"/v1/vector_stores/vs_1",
"/v1/vector_stores/vs_1/files/file_1",
"/v1/vector_stores/vs_1/file_batches/batch_1",
"/v1/assistants/asst_1",
"/v1/files/file_1",
"/v1/files/file_1/content",
"/v1/batches/batch_1",
"/v1/batches/batch_1/cancel",
"/v1/fine_tuning/jobs/job_1",
"/v1/fine_tuning/jobs/job_1/events",
"/v1/models/gpt-5",
"/key/sk-1234/budgets",
"/key/sk-1234/regenerate",
"/key/sk-1234/reset_spend",
"/key/info",
"/team/info",
]
changed = [
route
for route in routes
if (RouteChecks.is_llm_api_route(route), RouteChecks.is_info_route(route))
!= (
RouteChecks.is_llm_api_route(normalize_request_route(route)),
RouteChecks.is_info_route(normalize_request_route(route)),
)
]
assert changed == [], f"normalization changed the classification of {changed}"

221
tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py
View file

@ -16461,10 +16461,12 @@ from litellm.models.tag import LiteLLM_TagTable # noqa: E402
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.hooks.model_max_budget_limiter import ModelSpendHit # noqa: E402
from litellm.proxy.management_endpoints.key_budget_resolver import ( # noqa: E402
_match_model_budget_key,
_RecordedSpend as _BudgetsRecordedSpend,
_MODEL_BUDGET_COLD_NOTE,
_CUSTOM_AUTH_SKIPS_CHECKS_NOTE,
_MODEL_BUDGET_NOTE,
_THROTTLE_NOTE,
_read_end_user_model_spend,
_read_key_model_spend,
@ -16478,12 +16480,10 @@ from litellm.types.proxy.management_endpoints.key_management_endpoints import (
KeyBudgetEntry,
KeyBudgetNote,
)
from litellm.types.utils import BudgetConfig as _BudgetsBudgetConfig # noqa: E402
from litellm.types.utils import BudgetConfig # noqa: E402
from types import MappingProxyType # noqa: E402
from typing import get_args as _budgets_get_args # noqa: E402
BudgetConfig = _BudgetsBudgetConfig
_BUDGETS_RESOLVER = "litellm.proxy.management_endpoints.key_budget_resolver"
_BUDGETS_KEY_HASH = "hash-of-the-budgets-key"
_BUDGETS_RESET_AT = _budgets_datetime(2026, 9, 1, tzinfo=_budgets_timezone.utc)
@ -16522,7 +16522,12 @@ class _RecordingSpendReader:
class _RecordingModelSpendReader:
"""Stands in for the per-model cache so a test can prove which request models were probed."""
"""
Stands in for the per-model cache, falling back to the provider-stripped name the way it does.
A fake that only did `.get(model)` would report one counter under as many request models as route
to it, each claiming the full balance, which is exactly the bug this fallback causes.
"""
def __init__(self, spend_by_model=None):
default = {"gpt-5": 6.0, "claude-sonnet-4-5": 2.0}
@ -16531,16 +16536,27 @@ class _RecordingModelSpendReader:
async def __call__(self, *, entity_id, model, budget_config):
self.probed.append(model)
return self.spend_by_model.get(model)
for candidate in dict.fromkeys((model, model.split("/")[-1])):
spend = self.spend_by_model.get(candidate)
if spend is not None:
return ModelSpendHit(model=candidate, spend=spend)
return None
def _budgets_deps(read_spend=None, model_spend=None, match_model_budget_key=None, general_settings=None):
def _budgets_deps(
read_spend=None,
model_spend=None,
match_model_budget_key=None,
general_settings=None,
custom_auth_enabled=False,
):
model_reader = model_spend or _RecordingModelSpendReader()
return KeyBudgetResolverDeps(
prisma_client=MagicMock(),
user_api_key_cache=MagicMock(),
proxy_logging_obj=MagicMock(),
general_settings=general_settings if general_settings is not None else {},
custom_auth_enabled=custom_auth_enabled,
read_spend=read_spend or _RecordingSpendReader({}),
read_key_model_spend=model_reader,
read_end_user_model_spend=model_reader,
@ -17099,8 +17115,140 @@ async def test_key_budgets_probe_every_request_model_that_maps_onto_a_per_model_
@pytest.mark.asyncio
async def test_key_budgets_say_a_per_model_counter_is_missing_rather_than_claiming_a_cold_cache():
"""The stale note claimed the cache was cold even when it was warm under a different model key."""
async def test_key_budgets_report_one_row_per_counter_when_several_request_models_share_one():
"""
The reader retries without the provider prefix, so every `<provider>/gpt-5` lands on `gpt-5`'s
counter. A row each turned one $12 balance into four rows reading $12, a table that appears to
show $48 against a $12 cap while three of the four rows assert a counter that does not exist.
"""
reader = _RecordingModelSpendReader({"gpt-5": 12.0})
token = _budgets_token(
models=["openai/gpt-5", "azure/gpt-5", "bedrock/gpt-5"],
model_max_budget={"gpt-5": {"max_budget": 12.0, "budget_duration": "1d"}},
)
with _budgets_world(**_fully_populated_world()):
budgets = await resolve_key_budgets(
valid_token=token,
end_user_id=None,
deps=_budgets_deps(
model_spend=reader,
match_model_budget_key=lambda *, model, configured: "gpt-5" if model.endswith("gpt-5") else None,
),
)
rows = [e for e in budgets if e.scope == "key_model"]
assert {"openai/gpt-5", "azure/gpt-5", "bedrock/gpt-5"} <= set(reader.probed), "every request model is probed"
assert [r.entity_id for r in rows] == ["gpt-5"], "one counter is one row, and the row names the counter"
assert rows[0].entity_label == "gpt-5", "the cap it is compared against stays on the row"
assert rows[0].spend == 12.0
assert sum(r.spend for r in rows) == 12.0, "a shared balance must not be added up once per request model"
@pytest.mark.asyncio
async def test_key_budgets_name_the_counter_a_row_was_read_from_not_the_model_that_found_it():
"""
A provider-qualified cap falls back onto the bare counter, so the row that survives the collapse
is labelled by whichever request model probed first. Reporting that name would point a reader at a
counter that does not exist while showing a balance that came from a different one.
"""
reader = _RecordingModelSpendReader({"gpt-5": 7.0})
token = _budgets_token(
models=["openai/gpt-5"],
model_max_budget={"openai/gpt-5": {"max_budget": 12.0, "budget_duration": "1d"}},
)
with _budgets_world(**_fully_populated_world()):
budgets = await resolve_key_budgets(
valid_token=token,
end_user_id=None,
deps=_budgets_deps(
model_spend=reader,
match_model_budget_key=lambda *, model, configured: "openai/gpt-5" if "gpt-5" in model else None,
),
)
rows = [e for e in budgets if e.scope == "key_model"]
assert [(r.entity_id, r.entity_label) for r in rows] == [("gpt-5", "openai/gpt-5")]
assert rows[0].spend == 7.0
@pytest.mark.asyncio
async def test_key_budgets_keep_separate_rows_for_request_models_with_counters_of_their_own():
"""Collapsing on the counter must not collapse two counters that really are separate balances."""
reader = _RecordingModelSpendReader({"gpt-5": 4.0, "azure/gpt-5": 9.0})
token = _budgets_token(
models=["azure/gpt-5", "openai/gpt-5"],
model_max_budget={"gpt-5": {"max_budget": 12.0, "budget_duration": "1d"}},
)
with _budgets_world(**_fully_populated_world()):
budgets = await resolve_key_budgets(
valid_token=token,
end_user_id=None,
deps=_budgets_deps(
model_spend=reader,
match_model_budget_key=lambda *, model, configured: "gpt-5" if model.endswith("gpt-5") else None,
),
)
spend_by_model = {e.entity_id: e.spend for e in budgets if e.scope == "key_model"}
assert spend_by_model == {"gpt-5": 4.0, "azure/gpt-5": 9.0}, "openai/gpt-5 shares gpt-5's counter, azure does not"
@pytest.mark.asyncio
async def test_key_budgets_warn_on_every_row_when_custom_auth_skips_the_read_time_checks():
"""
A custom auth callable that returns its own token returns before any of these checks run, and the
wrapper skips `common_checks` too unless `custom_auth_run_common_checks` is set. Saying that on the
end-user row alone understates it by a dozen scopes, and the proxy already warns about it at boot.
"""
with _budgets_world(**_fully_populated_world()):
skipped = await resolve_key_budgets(
valid_token=_budgets_token(), end_user_id=None, deps=_budgets_deps(custom_auth_enabled=True)
)
opted_in = await resolve_key_budgets(
valid_token=_budgets_token(),
end_user_id=None,
deps=_budgets_deps(custom_auth_enabled=True, general_settings={"custom_auth_run_common_checks": True}),
)
no_custom_auth = await resolve_key_budgets(
valid_token=_budgets_token(), end_user_id=None, deps=_budgets_deps()
)
assert skipped, "the fixture must produce rows for this to assert anything"
assert all(_CUSTOM_AUTH_SKIPS_CHECKS_NOTE in e.notes for e in skipped), "every scope is unenforced, not one"
assert not any(_CUSTOM_AUTH_SKIPS_CHECKS_NOTE in e.notes for e in opted_in)
assert not any(_CUSTOM_AUTH_SKIPS_CHECKS_NOTE in e.notes for e in no_custom_auth)
@pytest.mark.asyncio
async def test_key_budgets_never_pair_a_missing_spend_state_with_a_number():
"""
`spend_state` is the only thing that stops a blank cell rendering as `$0.00` at 0% of the meter,
so a row that says the number is missing must not also carry one, or anything derived from one.
"""
async def _explode(**kwargs):
raise RuntimeError("redis is down")
with _budgets_world(**_fully_populated_world()):
budgets = await resolve_key_budgets(
valid_token=_budgets_token(),
end_user_id="end-user-budgets",
deps=_budgets_deps(read_spend=_explode, model_spend=_RecordingModelSpendReader({})),
)
states = {e.spend_state for e in budgets}
assert states == {"live", "no_counter", "unavailable"}, f"all three states must occur here, saw {states}"
for entry in budgets:
assert (entry.spend is None) == (entry.spend_state != "live"), entry
assert entry.spend is not None or entry.remaining is None, entry
@pytest.mark.asyncio
async def test_key_budgets_tell_a_cold_per_model_counter_apart_from_a_budget_that_cannot_trip():
"""
A counter nothing has written yet is a budget that will start working, so only `spend_state` may
say so. Saying it in a note instead put a live cap next to `project_spend_not_tracked`, which is
inert forever, under codes a client had no way to tell apart.
"""
warm = _RecordingModelSpendReader({"gpt-5": 6.0})
with _budgets_world(**_fully_populated_world()):
budgets = await resolve_key_budgets(
@ -17110,7 +17258,6 @@ async def test_key_budgets_say_a_per_model_counter_is_missing_rather_than_claimi
warm_entry = next(e for e in budgets if e.scope == "key_model" and e.entity_id == "gpt-5")
assert warm_entry.spend == 6.0
assert warm_entry.spend_state == "live"
assert _MODEL_BUDGET_COLD_NOTE not in warm_entry.notes
cold = _RecordingModelSpendReader({})
with _budgets_world(**_fully_populated_world()):
@ -17121,7 +17268,11 @@ async def test_key_budgets_say_a_per_model_counter_is_missing_rather_than_claimi
cold_entry = next(e for e in cold_budgets if e.scope == "key_model" and e.entity_id == "gpt-5")
assert cold_entry.spend is None
assert cold_entry.spend_state == "no_counter"
assert _MODEL_BUDGET_COLD_NOTE in cold_entry.notes
assert cold_entry.notes == warm_entry.notes, "a cold counter is the same budget, so it earns no extra caveat"
assert _MODEL_BUDGET_NOTE in cold_entry.notes
assert {note.code for note in cold_entry.notes}.isdisjoint(
{"project_spend_not_tracked"}
), "a transient cold counter must not read like a budget that can never trip"
@pytest.mark.asyncio
@ -17355,10 +17506,10 @@ def test_key_budgets_classify_every_note_code_and_leave_none_to_a_default():
"reservation_blocks_at_limit": "info",
"rolling_window": "info",
"user_budget_not_applied_to_team_key": "info",
"model_budget_fails_open": "info",
"throttled_instead_of_blocked": "info",
# only the note states it
"custom_auth_may_override_end_user_cap": "warning",
"custom_auth_skips_read_time_checks": "warning",
"end_user_route_only": "warning",
"per_model_counters": "warning",
"project_spend_not_tracked": "warning",
@ -17402,6 +17553,7 @@ async def test_key_budgets_emit_each_caveat_as_its_own_note_instead_of_one_joine
"reservation_blocks_at_limit",
"end_user_route_only",
"custom_auth_may_override_end_user_cap",
"custom_auth_skips_read_time_checks",
]
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
@ -17475,7 +17627,8 @@ 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"]
router.model_list = [{"model_name": "gpt-5", "litellm_params": {"model": "azure/gpt-5-prod"}}]
router.deployment_names = ["azure/gpt-5-prod", "azure/deleted-last-year"]
with patch("litellm.proxy.proxy_server.llm_router", router):
assert _request_models(()) == ("gpt-5", "azure/gpt-5-prod")
@ -17483,28 +17636,46 @@ async def test_key_budgets_probe_a_deployment_routed_at_directly():
assert _request_models(("only-the-key-models",)) == ("only-the-key-models",)
@pytest.mark.asyncio
async def test_key_budgets_stop_naming_a_deployment_after_it_is_deleted():
"""`Router.deployment_names` is appended to and never pruned, so it still names deleted deployments."""
router = MagicMock()
router.get_model_names.return_value = ()
router.model_list = [{"model_name": "gpt-5", "litellm_params": {"model": "azure/still-here"}}]
router.deployment_names = ["azure/still-here", "azure/deleted-last-year"]
with patch("litellm.proxy.proxy_server.llm_router", router):
assert _request_models(()) == ("azure/still-here",)
router.model_list = [{"not": "a deployment"}, {"model_name": "gpt-5", "litellm_params": {"model": "azure/ok"}}]
with patch("litellm.proxy.proxy_server.llm_router", router):
assert _request_models(()) == ("azure/ok",), "one malformed entry must not drop the rest"
@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)
limiter.resolve_model_spend = AsyncMock(return_value=ModelSpendHit(model="gpt-5", spend=4.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)
key_hit = await _read_key_model_spend(entity_id="hash-1", model="openai/gpt-5", budget_config=config)
key_kwargs = limiter.resolve_model_spend.await_args.kwargs
end_user_hit = 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",
assert key_hit == ModelSpendHit(model="gpt-5", spend=4.0)
assert key_kwargs == {
"scope": "virtual_key",
"entity_id": "hash-1",
"model": "openai/gpt-5",
"key_budget_config": config,
"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",
assert end_user_hit == ModelSpendHit(model="gpt-5", spend=4.0)
assert limiter.resolve_model_spend.await_args.kwargs == {
"scope": "end_user",
"entity_id": "cust-1",
"model": "gpt-5",
"key_budget_config": config,
"budget_config": config,
}

66
tests/test_litellm/proxy/test_proxy_utils.py
View file

@ -1522,3 +1522,69 @@ async def test_post_mcp_call_hook_skips_opted_out_guardrail(restore_callbacks):
assert guardrail.call_count == 0
assert [item.text for item in returned.content] == ["jane@example.com"]
@pytest.mark.asyncio
@pytest.mark.parametrize(
"route",
["/key/sk-victim-1234567890/budgets", "/key/" + "b" * 64 + "/budgets"],
)
async def test_proxy_only_error_log_does_not_put_a_key_path_param_in_call_type(route):
"""
The route becomes ``call_type`` on the payload every failure callback receives, and the OTEL
``llm.request.type`` and ``gen_ai.operation.name`` attributes derived from it. ``/key/{key_id}/budgets``
is an info route, so an auth failure on it reaches this hook with a key sitting in the path.
"""
from unittest.mock import patch as _patch
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.proxy._types import UserAPIKeyAuth
proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache())
captured = {}
async def _capture_async_failure(self, *args, **kwargs):
captured["call_type"] = self.call_type
return None
with _patch.object(Logging, "pre_call", lambda self, *a, **k: None), _patch.object(
Logging, "async_failure_handler", _capture_async_failure
):
await proxy_logging_obj._handle_logging_proxy_only_error(
request_data={},
user_api_key_dict=UserAPIKeyAuth(api_key="sk-bad", request_route=route),
route=route,
original_exception=Exception("Invalid proxy server token passed"),
)
assert captured["call_type"] == "/key/{key_id}/budgets", captured
assert "sk-victim" not in captured["call_type"]
assert "b" * 64 not in captured["call_type"]
@pytest.mark.asyncio
async def test_proxy_only_error_log_still_reports_the_route_for_ordinary_llm_failures():
"""Redacting the key routes must not blank out the call type every other failure is grouped by."""
from unittest.mock import patch as _patch
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.proxy._types import UserAPIKeyAuth
proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache())
captured = {}
async def _capture_async_failure(self, *args, **kwargs):
captured["call_type"] = self.call_type
return None
with _patch.object(Logging, "pre_call", lambda self, *a, **k: None), _patch.object(
Logging, "async_failure_handler", _capture_async_failure
):
await proxy_logging_obj._handle_logging_proxy_only_error(
request_data={"model": "gpt-4o", "messages": [{"role": "user", "content": "hi"}]},
user_api_key_dict=UserAPIKeyAuth(api_key="sk-bad", request_route="/v1/chat/completions"),
route="/v1/chat/completions",
original_exception=Exception("bad key"),
)
assert captured["call_type"] == "acompletion"