diff --git a/litellm/proxy/auth/auth_exception_handler.py b/litellm/proxy/auth/auth_exception_handler.py index 603e72463bc..ff9bb2c9de0 100644 --- a/litellm/proxy/auth/auth_exception_handler.py +++ b/litellm/proxy/auth/auth_exception_handler.py @@ -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 diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index a105bf19458..e2a756d840a 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -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 = [ diff --git a/litellm/proxy/hooks/model_max_budget_limiter.py b/litellm/proxy/hooks/model_max_budget_limiter.py index 9fdd3140742..186851e9b90 100644 --- a/litellm/proxy/hooks/model_max_budget_limiter.py +++ b/litellm/proxy/hooks/model_max_budget_limiter.py @@ -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] diff --git a/litellm/proxy/management_endpoints/key_budget_resolver.py b/litellm/proxy/management_endpoints/key_budget_resolver.py index 601c17f4ba5..d94742c2dc1 100644 --- a/litellm/proxy/management_endpoints/key_budget_resolver.py +++ b/litellm/proxy/management_endpoints/key_budget_resolver.py @@ -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, diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 0c4fc6658c3..f95e35a3df2 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -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 diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index a743526e975..ca4c662d7e2 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -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, diff --git a/litellm/types/proxy/management_endpoints/key_management_endpoints.py b/litellm/types/proxy/management_endpoints/key_management_endpoints.py index 574ac81290c..2a2b1bc5182 100644 --- a/litellm/types/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/types/proxy/management_endpoints/key_management_endpoints.py @@ -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", diff --git a/tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py b/tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py index 6642cc1a2c3..4344a3ce122 100644 --- a/tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py +++ b/tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py @@ -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 + ) diff --git a/tests/test_litellm/proxy/auth/test_auth_exception_handler.py b/tests/test_litellm/proxy/auth/test_auth_exception_handler.py index 27798ec0bff..2c5ac683cfa 100644 --- a/tests/test_litellm/proxy/auth/test_auth_exception_handler.py +++ b/tests/test_litellm/proxy/auth/test_auth_exception_handler.py @@ -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" diff --git a/tests/test_litellm/proxy/auth/test_auth_utils.py b/tests/test_litellm/proxy/auth/test_auth_utils.py index 5becd05b8e8..13e9cef951d 100644 --- a/tests/test_litellm/proxy/auth/test_auth_utils.py +++ b/tests/test_litellm/proxy/auth/test_auth_utils.py @@ -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}" 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 717c0150b77..a843121f01c 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 @@ -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 `/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, } diff --git a/tests/test_litellm/proxy/test_proxy_utils.py b/tests/test_litellm/proxy/test_proxy_utils.py index d6cf0e30139..f23b15660b1 100644 --- a/tests/test_litellm/proxy/test_proxy_utils.py +++ b/tests/test_litellm/proxy/test_proxy_utils.py @@ -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"