diff --git a/litellm/_auth/AGENTS.md b/litellm/_auth/AGENTS.md index b6e54cac36f..04ffcebf3b6 100644 --- a/litellm/_auth/AGENTS.md +++ b/litellm/_auth/AGENTS.md @@ -11,6 +11,6 @@ - `decision.py` is the interface. `AuthCore.start(RequestFacts, Settings, now) -> Step` is the only thing a core implements, and `python_core.PythonCore` is the Python one. A Rust core will implement the same protocol through pyo3, and the builder picks one without changing anything else - A core does no IO and reads no globals. The same facts, settings, time and effect answers always give the same outcome - A `Step` is either an `Outcome` or an effect. Each effect carries its typed `resume` continuation, and `driver.drive` answers effects through an `EffectHandler` until an outcome comes back. The proxy's handler (`_BuilderEffects` in `litellm/proxy/auth/user_api_key_auth.py`) owns all cache and DB IO, so both cores share it -- Projection stays in the proxy. `Allow` grants become `UserAPIKeyAuth` in the builder, and `litellm/proxy/auth/auth_denial.py` maps every `Denial` to the existing error contract -- The key checks run in the core, in the order the proxy has always applied them: blocked key, model access, SCIM-deactivated owner, budget exemption, team member budget, expiry, budget alerts, then the key's own max budget (or its throttle). `start_resolved_key` runs the same checks for keys the proxy resolved itself, such as internal jobs -- The soft budget alert, per-model budgets, the user's model budget and the team, project and global spend lookups still run as projection in `_finish_virtual_key_auth`. Move them into the core one at a time, as effects plus pure decisions, and delete the builder code each one replaces in the same change +- Projection stays in the proxy. `_authenticated` in the builder module turns an `Allow` grant into `UserAPIKeyAuth`, and `litellm/proxy/auth/auth_denial.py` maps every `Denial` to the existing error contract +- The key checks run in the core, in the order the proxy has always applied them: blocked key, model access, SCIM-deactivated owner, budget exemption, team member budget, expiry, budget alerts, the key's own max budget (or its throttle), the soft budget alert, per-model budgets, the team and project context, the proxy budget alert, then the key's token. A budget-exempt request skips everything from the team member budget to per-model budgets except expiry. `start_resolved_key` runs the same checks for keys the proxy resolved itself, such as internal jobs +- Per-model budgets are still one opaque effect (`EnforceModelBudgets`), because the key, team member and custom-auth paths share their fallback logic. Move that logic into the core for every path at once rather than forking it per path diff --git a/litellm/_auth/decision.py b/litellm/_auth/decision.py index 9cee892f79c..9e697de7fad 100644 --- a/litellm/_auth/decision.py +++ b/litellm/_auth/decision.py @@ -17,6 +17,7 @@ class Settings: allow_user_auth: bool has_db: bool master_key_only_routes: frozenset[str] + proxy_max_budget: float @dataclass(frozen=True, slots=True) @@ -41,8 +42,6 @@ class CachedAdminKey: @dataclass(frozen=True, slots=True) class VirtualKey: - loaded_from_db: bool - budget_checks_skipped: bool budget_throttle_pct: float | None @@ -110,6 +109,11 @@ class KeyBudgetExceeded: max_budget: float +@dataclass(frozen=True, slots=True) +class KeyWithoutToken: + pass + + Denial: TypeAlias = ( MissingKey | MalformedKey @@ -122,6 +126,7 @@ Denial: TypeAlias = ( | OwnerDeactivated | TeamMemberBudgetExceeded | KeyBudgetExceeded + | KeyWithoutToken ) @@ -153,8 +158,11 @@ PresentedKey: TypeAlias = AdminKey | NonAdminKey @dataclass(frozen=True, slots=True) class KeyFacts: + has_token: bool blocked: bool expires: datetime | None + spend: float + soft_budget: float | None max_budget: float | None team_member_spend: float | None user_id: str | None @@ -256,7 +264,33 @@ class AlertKeyBudget: resume: "Callable[[], Step]" -Effect: TypeAlias = ( +@dataclass(frozen=True, slots=True) +class AlertKeySoftBudget: + resume: "Callable[[], Step]" + + +@dataclass(frozen=True, slots=True) +class EnforceModelBudgets: + resume: "Callable[[], Step]" + + +@dataclass(frozen=True, slots=True) +class LoadKeyContext: + resume: "Callable[[], Step]" + + +@dataclass(frozen=True, slots=True) +class ReadProxySpend: + resume: "Callable[[float | None], Step]" + + +@dataclass(frozen=True, slots=True) +class AlertProxyBudget: + spend: float + resume: "Callable[[], Step]" + + +IdentityEffect: TypeAlias = ( ResolveEndUser | ResolvePresentedKey | EvictPresentedKey @@ -266,12 +300,20 @@ Effect: TypeAlias = ( | EnforceModelAccess | LoadKeyOwner | CheckBudgetExemption - | ReadTeamMemberBudget + | LoadKeyContext +) +BudgetEffect: TypeAlias = ( + ReadTeamMemberBudget | AlertTeamMemberBudget | SendKeyBudgetAlerts | ReadKeySpend | AlertKeyBudget + | AlertKeySoftBudget + | EnforceModelBudgets + | ReadProxySpend + | AlertProxyBudget ) +Effect: TypeAlias = IdentityEffect | BudgetEffect Step: TypeAlias = Outcome | Effect diff --git a/litellm/_auth/driver.py b/litellm/_auth/driver.py index 9fd1e7cc47d..e709e095b2a 100644 --- a/litellm/_auth/driver.py +++ b/litellm/_auth/driver.py @@ -2,22 +2,29 @@ from typing import Protocol, assert_never from litellm._auth.decision import ( AlertKeyBudget, + AlertKeySoftBudget, + AlertProxyBudget, AlertTeamMemberBudget, Allow, + BudgetEffect, CheckBudgetExemption, Deny, Effect, EnforceModelAccess, + EnforceModelBudgets, EvictPresentedKey, + IdentityEffect, InspectKey, KeyFacts, LoadKey, + LoadKeyContext, LoadKeyOwner, Outcome, OwnerFacts, PrepareKey, PresentedKey, ReadKeySpend, + ReadProxySpend, ReadTeamMemberBudget, ResolveEndUser, ResolvePresentedKey, @@ -56,8 +63,37 @@ class EffectHandler(Protocol): async def alert_key_budget(self, spend: float) -> None: ... + async def alert_key_soft_budget(self) -> None: ... + + async def enforce_model_budgets(self) -> None: ... + + async def load_key_context(self) -> None: ... + + async def read_proxy_spend(self) -> float | None: ... + + async def alert_proxy_budget(self, spend: float) -> None: ... + async def _answer(step: Effect, handler: EffectHandler) -> Step: + match step: + case ( + ResolveEndUser() + | ResolvePresentedKey() + | EvictPresentedKey() + | LoadKey() + | PrepareKey() + | InspectKey() + | EnforceModelAccess() + | LoadKeyOwner() + | CheckBudgetExemption() + | LoadKeyContext() + ): + return await _answer_identity(step, handler) + case _: + return await _answer_budget(step, handler) + + +async def _answer_identity(step: IdentityEffect, handler: EffectHandler) -> Step: match step: case ResolveEndUser(prefetch_key=prefetch_key, resume=resume): await handler.resolve_end_user(prefetch_key) @@ -82,6 +118,15 @@ async def _answer(step: Effect, handler: EffectHandler) -> Step: return resume(await handler.load_key_owner(user_id)) case CheckBudgetExemption(resume=resume): return resume(await handler.check_budget_exemption()) + case LoadKeyContext(resume=resume): + await handler.load_key_context() + return resume() + case _: + assert_never(step) + + +async def _answer_budget(step: BudgetEffect, handler: EffectHandler) -> Step: + match step: case ReadTeamMemberBudget(user_id=user_id, team_id=team_id, resume=resume): return resume(await handler.read_team_member_budget(user_id, team_id)) case AlertTeamMemberBudget(user_id=user_id, team_id=team_id, budget=budget, resume=resume): @@ -95,6 +140,17 @@ async def _answer(step: Effect, handler: EffectHandler) -> Step: case AlertKeyBudget(spend=spend, resume=resume): await handler.alert_key_budget(spend) return resume() + case AlertKeySoftBudget(resume=resume): + await handler.alert_key_soft_budget() + return resume() + case EnforceModelBudgets(resume=resume): + await handler.enforce_model_budgets() + return resume() + case ReadProxySpend(resume=resume): + return resume(await handler.read_proxy_spend()) + case AlertProxyBudget(spend=spend, resume=resume): + await handler.alert_proxy_budget(spend) + return resume() case _: assert_never(step) diff --git a/litellm/_auth/python_core.py b/litellm/_auth/python_core.py index 20bc1d7e086..c57577c8227 100644 --- a/litellm/_auth/python_core.py +++ b/litellm/_auth/python_core.py @@ -6,19 +6,24 @@ from typing import Final from litellm._auth.decision import ( AdminKey, AlertKeyBudget, + AlertKeySoftBudget, + AlertProxyBudget, AlertTeamMemberBudget, Allow, CachedAdminKey, CheckBudgetExemption, Deny, EnforceModelAccess, + EnforceModelBudgets, EvictPresentedKey, ExpiredKey, InspectKey, KeyBlocked, KeyBudgetExceeded, KeyFacts, + KeyWithoutToken, LoadKey, + LoadKeyContext, LoadKeyOwner, MalformedKey, MasterKey, @@ -33,6 +38,7 @@ from litellm._auth.decision import ( PrepareKey, PresentedKey, ReadKeySpend, + ReadProxySpend, ReadTeamMemberBudget, RequestFacts, ResolveEndUser, @@ -52,6 +58,14 @@ _USER_AUTH_ROUTE: Final = "/user/auth" _VIRTUAL_KEY_PREFIX: Final = "sk-" +@dataclass(frozen=True, slots=True) +class _KeyCheck: + settings: Settings + now: datetime + is_llm_api_route: bool + key_label: str + + class PythonCore: def start(self, facts: RequestFacts, settings: Settings, now: datetime) -> Step: token: Final = facts.token @@ -75,9 +89,7 @@ class PythonCore: ) def start_resolved_key(self, facts: RequestFacts, settings: Settings, now: datetime) -> Step: - return _check_key( - _KeyCheck(facts=facts, settings=settings, now=now, loaded_from_db=True, key_label=facts.token or "") - ) + return _check_key(_KeyCheck(settings, now, facts.is_llm_api_route, key_label=facts.token or "")) def _after_presented_key( @@ -103,36 +115,17 @@ def _after_presented_key( if not settings.has_db: return Deny(NoDatabase()) if presented is not None: - cached: Final = _KeyCheck(facts=facts, settings=settings, now=now, loaded_from_db=False, key_label=token) + cached: Final = _KeyCheck(settings, now, facts.is_llm_api_route, key_label=token) return PrepareKey(loaded_from_db=False, resume=lambda: _check_key(cached)) if not token.startswith(_VIRTUAL_KEY_PREFIX): return Deny(NotAVirtualKey()) - loaded: Final = _KeyCheck(facts=facts, settings=settings, now=now, loaded_from_db=True, key_label=hashed_token) + loaded: Final = _KeyCheck(settings, now, facts.is_llm_api_route, key_label=hashed_token) return LoadKey( hashed_token=hashed_token, resume=lambda: PrepareKey(loaded_from_db=True, resume=lambda: _check_key(loaded)), ) -@dataclass(frozen=True, slots=True) -class _KeyCheck: - facts: RequestFacts - settings: Settings - now: datetime - loaded_from_db: bool - key_label: str - - -def _allow(check: _KeyCheck, budget_checks_skipped: bool, budget_throttle_pct: float | None) -> Step: - return Allow( - VirtualKey( - loaded_from_db=check.loaded_from_db, - budget_checks_skipped=budget_checks_skipped, - budget_throttle_pct=budget_throttle_pct, - ) - ) - - def _check_key(check: _KeyCheck) -> Step: return InspectKey(resume=lambda key: _after_inspect(check, key)) @@ -191,14 +184,14 @@ def _check_expiry(check: _KeyCheck, key: KeyFacts, exempt: bool) -> Step: if key.expires is not None and key.expires < check.now: return Deny(ExpiredKey(expires=key.expires, now=check.now, key_label=check.key_label)) if exempt: - return _allow(check, budget_checks_skipped=True, budget_throttle_pct=None) + return _load_context(check, key, budget_throttle_pct=None) return SendKeyBudgetAlerts(resume=lambda: _check_key_budget(check, key)) def _check_key_budget(check: _KeyCheck, key: KeyFacts) -> Step: max_budget: Final = key.max_budget - if not check.facts.is_llm_api_route or max_budget is None: - return _allow(check, budget_checks_skipped=False, budget_throttle_pct=None) + if not check.is_llm_api_route or max_budget is None: + return _check_soft_budget(check, key, budget_throttle_pct=None) return ReadKeySpend( max_budget=max_budget, resume=lambda spend: AlertKeyBudget( @@ -210,7 +203,38 @@ def _check_key_budget(check: _KeyCheck, key: KeyFacts) -> Step: def _after_key_spend(check: _KeyCheck, key: KeyFacts, spend: float, max_budget: float) -> Step: # A NaN max_budget compares false against any spend, so a non-finite budget means no limit, never a bypass. if not math.isfinite(max_budget) or spend < max_budget: - return _allow(check, budget_checks_skipped=False, budget_throttle_pct=None) + return _check_soft_budget(check, key, budget_throttle_pct=None) if key.budget_throttle_pct is not None: - return _allow(check, budget_checks_skipped=False, budget_throttle_pct=key.budget_throttle_pct) + return _check_soft_budget(check, key, budget_throttle_pct=key.budget_throttle_pct) return Deny(KeyBudgetExceeded(spend=spend, max_budget=max_budget)) + + +def _check_soft_budget(check: _KeyCheck, key: KeyFacts, budget_throttle_pct: float | None) -> Step: + def model_budgets() -> Step: + return EnforceModelBudgets(resume=lambda: _load_context(check, key, budget_throttle_pct)) + + if key.soft_budget and key.spend >= key.soft_budget: + return AlertKeySoftBudget(resume=model_budgets) + return model_budgets() + + +def _load_context(check: _KeyCheck, key: KeyFacts, budget_throttle_pct: float | None) -> Step: + return LoadKeyContext(resume=lambda: _check_proxy_budget(check, key, budget_throttle_pct)) + + +def _check_proxy_budget(check: _KeyCheck, key: KeyFacts, budget_throttle_pct: float | None) -> Step: + if not (check.settings.proxy_max_budget > 0 and check.settings.has_db): + return _grant(key, budget_throttle_pct) + return ReadProxySpend( + resume=lambda spend: ( + _grant(key, budget_throttle_pct) + if spend is None + else AlertProxyBudget(spend=spend, resume=lambda: _grant(key, budget_throttle_pct)) + ) + ) + + +def _grant(key: KeyFacts, budget_throttle_pct: float | None) -> Step: + if not key.has_token: + return Deny(KeyWithoutToken()) + return Allow(VirtualKey(budget_throttle_pct=budget_throttle_pct)) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 74e7367365d..7ab5fb63dde 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -6843,32 +6843,39 @@ async def virtual_key_soft_budget_check( """ if valid_token.soft_budget and valid_token.spend >= valid_token.soft_budget: - verbose_proxy_logger.debug( - "Crossed Soft Budget for token %s, spend %s, soft_budget %s", - valid_token.token, - valid_token.spend, - valid_token.soft_budget, - ) - call_info: Final = CallInfo( - token=valid_token.token, - spend=valid_token.spend, - max_budget=valid_token.max_budget, - soft_budget=valid_token.soft_budget, - user_id=valid_token.user_id, - team_id=valid_token.team_id, - team_alias=valid_token.team_alias, - organization_id=valid_token.org_id, - user_email=user_obj.user_email if user_obj else None, - key_alias=valid_token.key_alias, - event_group=Litellm_EntityType.KEY, - ) + send_soft_budget_alert(valid_token, proxy_logging_obj, user_obj) - asyncio.create_task( - proxy_logging_obj.budget_alerts( - type="soft_budget", - user_info=call_info, - ) + +def send_soft_budget_alert( + valid_token: UserAPIKeyAuth, + proxy_logging_obj: ProxyLogging, + user_obj: LiteLLM_UserTable | None, +) -> None: + verbose_proxy_logger.debug( + "Crossed Soft Budget for token %s, spend %s, soft_budget %s", + valid_token.token, + valid_token.spend, + valid_token.soft_budget, + ) + call_info: Final = CallInfo( + token=valid_token.token, + spend=valid_token.spend, + max_budget=valid_token.max_budget, + soft_budget=valid_token.soft_budget, + user_id=valid_token.user_id, + team_id=valid_token.team_id, + team_alias=valid_token.team_alias, + organization_id=valid_token.org_id, + user_email=user_obj.user_email if user_obj else None, + key_alias=valid_token.key_alias, + event_group=Litellm_EntityType.KEY, + ) + asyncio.create_task( + proxy_logging_obj.budget_alerts( + type="soft_budget", + user_info=call_info, ) + ) _virtual_key_soft_budget_check: Final = virtual_key_soft_budget_check diff --git a/litellm/proxy/auth/auth_denial.py b/litellm/proxy/auth/auth_denial.py index ae678bfbf8e..acbff4996a8 100644 --- a/litellm/proxy/auth/auth_denial.py +++ b/litellm/proxy/auth/auth_denial.py @@ -8,6 +8,7 @@ from litellm._auth.decision import ( ExpiredKey, KeyBlocked, KeyBudgetExceeded, + KeyWithoutToken, MalformedKey, MasterKeyOnlyRoute, MissingKey, @@ -122,5 +123,9 @@ def denial_exception( ) case KeyBudgetExceeded(spend=spend, max_budget=max_budget): return key_budget_exceeded_error(key or UserAPIKeyAuth(), spend, max_budget) + case KeyWithoutToken(): + return HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid API key, no token associated" + ) case _: assert_never(denial) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 92e14f9149b..f58757b5d7e 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -13,7 +13,7 @@ from collections.abc import Mapping, Sequence from datetime import datetime, timezone from enum import Enum from types import MappingProxyType -from typing import Final, NamedTuple, Protocol, Union, cast +from typing import Final, NamedTuple, Protocol, Union, assert_never, cast import fastapi from fastapi import HTTPException, Request, WebSocket, status @@ -36,10 +36,12 @@ from litellm._auth.decision import ( AuthCore, CachedAdminKey, Deny, + Grant, KeyFacts, MasterKey, NoMasterKey, NonAdminKey, + Outcome, OwnerFacts, PresentedKey, RequestFacts, @@ -117,12 +119,13 @@ from litellm.proxy.auth.auth_checks import ( # noqa: F401 # legacy module expo resolve_and_validate_end_user_id, resolve_default_end_user_budget, route_skips_budget_checks, + send_soft_budget_alert, send_token_budget_alert, stamp_matched_model_access_groups, team_member_max_budget_alert_check, virtual_key_max_budget_alert_check, virtual_key_max_budget_check, - virtual_key_soft_budget_check, + virtual_key_soft_budget_check, # pyright: ignore[reportUnusedImport] # backwards-compatible package export ) from litellm.proxy.auth.auth_denial import denial_exception from litellm.proxy.auth.auth_exception_handler import UserAPIKeyAuthExceptionHandler @@ -2255,44 +2258,6 @@ async def _resolve_end_user( return _ResolvedEndUser(end_user_id=end_user_id, end_user_object=None, budget=None) -async def _master_key_auth( - token: str, - end_user: _ResolvedEndUser, - admin_user_id: str, - parent_otel_span: Span | None, - route: str, - start_time: datetime, - user_api_key_cache: UserApiKeyCache, - proxy_logging_obj: ProxyLogging, -) -> UserAPIKeyAuth: - # A stable alias replaces the raw master key so neither the key nor its hash reaches spend logs, - # metrics labels, audit trails or rate-limit buckets. - master_key_auth: Final = await _return_user_api_key_auth_obj( - user_obj=None, - user_role=LitellmUserRoles.PROXY_ADMIN, - api_key=LITELLM_PROXY_MASTER_KEY_ALIAS, - parent_otel_span=parent_otel_span, - valid_token_dict={**end_user.params, "user_id": admin_user_id}, - route=route, - start_time=start_time, - ) - asyncio.create_task( - cache_key_object( - hashed_token=hash_token(token), - user_api_key_obj=master_key_auth, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - ) - with_end_user: Final = update_valid_token_with_end_user_params( - valid_token=master_key_auth, end_user_params=end_user.params - ) - with_end_user.via_virtual_key = True - if end_user.end_user_object is not None: - with_end_user.end_user_object_permission = end_user.end_user_object.object_permission - return with_end_user - - def _proxy_model_list() -> list[object] | None: from litellm.proxy.proxy_server import llm_model_list @@ -2308,6 +2273,7 @@ def _core_settings(master_key: str | None, allow_user_auth: bool, prisma_client: allow_user_auth=allow_user_auth, has_db=prisma_client is not None, master_key_only_routes=_MASTER_KEY_ONLY_ROUTES, + proxy_max_budget=litellm.max_budget, ) @@ -2405,6 +2371,22 @@ class _BuilderEffects: self.resolved_key: UserAPIKeyAuth | None = preresolved_key self.owner: LiteLLM_UserTable | None = None + @property + def token(self) -> str: + return self._token + + @property + def parent_otel_span(self) -> Span | None: + return self._parent_otel_span + + @property + def user_api_key_cache(self) -> UserApiKeyCache: + return self._user_api_key_cache + + @property + def proxy_logging_obj(self) -> ProxyLogging: + return self._proxy_logging_obj + @property def key(self) -> UserAPIKeyAuth: if self.resolved_key is None: @@ -2500,8 +2482,11 @@ class _BuilderEffects: key: Final = _update_key_budget_with_temp_budget_increase(self.key) self.resolved_key = key return KeyFacts( + has_token=key.token is not None, blocked=key.blocked is True, expires=None if key.expires is None else _aware_expiry(key.expires), + spend=key.spend, + soft_budget=key.soft_budget, max_budget=key.max_budget, team_member_spend=key.team_member_spend, user_id=key.user_id, @@ -2634,6 +2619,47 @@ class _BuilderEffects: async def alert_key_budget(self, spend: float) -> None: send_token_budget_alert(self.key, spend, self._proxy_logging_obj, self.owner) + async def alert_key_soft_budget(self) -> None: + send_soft_budget_alert(self.key, self._proxy_logging_obj, self.owner) + + async def enforce_model_budgets(self) -> None: + with tracer.trace("litellm.proxy.auth.budget_checks"): + await _check_model_budgets( + request=self._request, request_data=self._request_data, valid_token=self.key, route=self._route + ) + + async def load_key_context(self) -> None: + await _load_team_and_project( + valid_token=self.key, + prisma_client=self._prisma_client, + user_api_key_cache=self._user_api_key_cache, + parent_otel_span=self._parent_otel_span, + proxy_logging_obj=self._proxy_logging_obj, + ) + + async def read_proxy_spend(self) -> float | None: + if self._prisma_client is None: + return None + with tracer.trace("litellm.proxy.auth.get_global_proxy_spend"): + return await fetch_global_spend_with_event_coordination( + cache_key=GLOBAL_PROXY_SPEND_CACHE_KEY, + user_api_key_cache=self._user_api_key_cache, + prisma_client=self._prisma_client, + ) + + async def alert_proxy_budget(self, spend: float) -> None: + from litellm.proxy.proxy_server import litellm_proxy_admin_name + + call_info: Final = CallInfo( + token=self.key.token, + spend=spend, + max_budget=litellm.max_budget, + user_id=litellm_proxy_admin_name, + team_id=self.key.team_id, + event_group=Litellm_EntityType.PROXY, + ) + asyncio.create_task(self._proxy_logging_obj.budget_alerts(type="proxy_budget", user_info=call_info)) + async def user_api_key_auth_builder( request: Request, @@ -3073,48 +3099,14 @@ async def user_api_key_auth_builder( ) finally: valid_token = effects.resolved_key # rebind-ok: the exception handler below reports the resolved identity - match outcome: - case Deny(denial=denial): - raise denial_exception( - denial, token=api_key, route=route, jwt_auth_enabled=enable_jwt_auth, key=effects.resolved_key - ) - case Allow(grant=grant): - match grant: - case NoMasterKey(token=token): - return UserAPIKeyAuth( - api_key=token, - user_role=LitellmUserRoles.INTERNAL_USER, - parent_otel_span=parent_otel_span, - ) - case UserAuthRoute(): - return UserAPIKeyAuth() - case MasterKey(): - return await _master_key_auth( - token=api_key, - end_user=effects.end_user, - admin_user_id=litellm_proxy_admin_name, - parent_otel_span=parent_otel_span, - route=route, - start_time=start_time, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - case CachedAdminKey(): - return _with_end_user(effects.key, effects.end_user, parent_otel_span) - case VirtualKey(budget_checks_skipped=budget_checks_skipped, budget_throttle_pct=throttle_pct): - if throttle_pct is not None: - effects.key.budget_throttle_pct = throttle_pct - return await _finish_virtual_key_auth( - request=request, - request_data=request_data, - valid_token=effects.key, - route=route, - start_time=start_time, - parent_otel_span=parent_otel_span, - end_user=effects.end_user, - user_obj=effects.owner, - skip_budget_checks=budget_checks_skipped, - ) + return await _authenticated( + outcome, + effects, + admin_user_id=litellm_proxy_admin_name, + route=route, + start_time=start_time, + jwt_auth_enabled=enable_jwt_auth, + ) except Exception as e: return await UserAPIKeyAuthExceptionHandler.handle_authentication_error( e=e, @@ -3130,198 +3122,233 @@ async def user_api_key_auth_builder( _user_api_key_auth_builder: Final = user_api_key_auth_builder -async def _check_soft_and_model_budgets( +async def _check_model_budgets( request: Request, request_data: dict[str, object], # mutable-ok: per-model budget fallbacks rewrite the request model in place valid_token: UserAPIKeyAuth, route: str, - user_obj: LiteLLM_UserTable | None, ) -> None: - from litellm.proxy.proxy_server import llm_router, model_max_budget_limiter, prisma_client, proxy_logging_obj + from litellm.proxy.proxy_server import llm_router, model_max_budget_limiter, prisma_client - # Check 6. Soft Budget Check - await virtual_key_soft_budget_check( - valid_token=valid_token, - proxy_logging_obj=proxy_logging_obj, - user_obj=user_obj, - ) - - # Check 5. Token Model Spend is under Model budget - max_budget_per_model: Final = valid_token.model_max_budget - current_model = _get_model_from_request_context( - request_data=request_data, - route=route, - request=request, - llm_router=llm_router, - team_id=valid_token.team_id, - ) - current_models = _get_model_names_for_budget_checks(model=current_model) - - if max_budget_per_model and prisma_client is not None and current_models and valid_token.token is not None: - ## GET THE SPEND FOR THIS MODEL - for model_name in current_models: - await _check_key_model_budget_with_fallback( - valid_token=valid_token, - model_max_budget_limiter=model_max_budget_limiter, - model_name=model_name, + def requested_models() -> list[str]: + return _get_model_names_for_budget_checks( + model=_get_model_from_request_context( request_data=request_data, + route=route, request=request, - llm_model_list=cast( # cast-ok: router model lists are deployment sequences - Sequence[object] | None, _proxy_model_list() - ), llm_router=llm_router, + team_id=valid_token.team_id, ) - - # Recompute after a potential budget-fallback rewrite so - # the user check below validates the final model - current_model = _get_model_from_request_context( - request_data=request_data, - route=route, - request=request, - llm_router=llm_router, - team_id=valid_token.team_id, ) - current_models = _get_model_names_for_budget_checks(model=current_model) - # Check 5a. Internal user model_max_budget - if current_models: - await _check_user_model_budget( + models: Final = requested_models() + checks_key_budgets: Final = ( + bool(valid_token.model_max_budget) + and prisma_client is not None + and bool(models) + and valid_token.token is not None + ) + for model_name in models if checks_key_budgets else (): + await _check_key_model_budget_with_fallback( valid_token=valid_token, model_max_budget_limiter=model_max_budget_limiter, - models=current_models, + model_name=model_name, + request_data=request_data, + request=request, + llm_model_list=cast( # cast-ok: router model lists are deployment sequences + Sequence[object] | None, _proxy_model_list() + ), + llm_router=llm_router, ) - - -async def _finish_virtual_key_auth( - request: Request, - request_data: dict[str, object], # mutable-ok: per-model budget fallbacks rewrite the request model in place - valid_token: UserAPIKeyAuth, - route: str, - start_time: datetime, - parent_otel_span: Span | None, - end_user: _ResolvedEndUser, - user_obj: LiteLLM_UserTable | None, - skip_budget_checks: bool, -) -> UserAPIKeyAuth: - from litellm.proxy.proxy_server import ( - litellm_proxy_admin_name, - prisma_client, - proxy_logging_obj, - user_api_key_cache, + # A budget fallback may have rewritten the request model, so the user check validates the final one. + await _check_user_model_budget( + valid_token=valid_token, + model_max_budget_limiter=model_max_budget_limiter, + models=requested_models() if checks_key_budgets else models, ) - if not skip_budget_checks: - with tracer.trace("litellm.proxy.auth.budget_checks"): - await _check_soft_and_model_budgets( - request=request, request_data=request_data, valid_token=valid_token, route=route, user_obj=user_obj - ) - # Check 6: Additional Common Checks across jwt + key auth - if valid_token.team_id is not None: - try: - if valid_token.team_id == UI_TEAM_ID: - raise TeamNotFoundError(team_id=UI_TEAM_ID) - with tracer.trace("litellm.proxy.auth.get_team_object"): - _team_obj = await get_team_object( - team_id=valid_token.team_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - ) - except HTTPException: - token_team_models: Final = _token_team_models(valid_token) - _team_obj = LiteLLM_TeamTableCachedObj( - team_id=valid_token.team_id, - max_budget=valid_token.team_max_budget, - soft_budget=valid_token.team_soft_budget, - model_max_budget=valid_token.team_model_max_budget, - spend=valid_token.team_spend, - tpm_limit=valid_token.team_tpm_limit, - rpm_limit=valid_token.team_rpm_limit, - tpd_limit=valid_token.team_tpd_limit, - blocked=valid_token.team_blocked, - models=token_team_models, - metadata=valid_token.team_metadata, - object_permission_id=valid_token.team_object_permission_id, - object_permission=await _resolve_object_permission_for_unresolvable_team( - object_permission_id=valid_token.team_object_permission_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - ), - ) - else: - _team_obj = None +async def _load_team_and_project( + valid_token: UserAPIKeyAuth, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Span | None, + proxy_logging_obj: ProxyLogging, +) -> None: + team_obj: Final = ( + None + if valid_token.team_id is None + else await _key_team_object( + valid_token, valid_token.team_id, prisma_client, user_api_key_cache, parent_otel_span, proxy_logging_obj + ) + ) + valid_token.team_object_permission = team_obj.object_permission if team_obj is not None else None + if team_obj is not None: + # Guardrails or other metadata added after the key was cached apply without a cache eviction. + valid_token.team_metadata = team_obj.metadata + if valid_token.project_id is None: + return + project_obj: Final = await get_project_object( + project_id=valid_token.project_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + if project_obj is not None: + valid_token.project_metadata = project_obj.metadata + valid_token.project_alias = project_obj.project_alias - if _team_obj is not None: - valid_token.team_object_permission = _team_obj.object_permission - # Keep team_metadata in sync with the freshly fetched team so that - # guardrails (or any other metadata) added after the key was cached - # are picked up on subsequent requests without a cache eviction. - valid_token.team_metadata = _team_obj.metadata - else: - valid_token.team_object_permission = None - # Fetch project object if key belongs to a project - _project_obj = None - if valid_token.project_id is not None: - _project_obj = await get_project_object( - project_id=valid_token.project_id, +async def _key_team_object( + valid_token: UserAPIKeyAuth, + team_id: str, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Span | None, + proxy_logging_obj: ProxyLogging, +) -> LiteLLM_TeamTableCachedObj: + fetched: Final = ( + None + if team_id == UI_TEAM_ID + else await _team_object_or_none(team_id, prisma_client, user_api_key_cache, parent_otel_span, proxy_logging_obj) + ) + if fetched is not None: + return fetched + return LiteLLM_TeamTableCachedObj( + team_id=team_id, + max_budget=valid_token.team_max_budget, + soft_budget=valid_token.team_soft_budget, + model_max_budget=valid_token.team_model_max_budget, + spend=valid_token.team_spend, + tpm_limit=valid_token.team_tpm_limit, + rpm_limit=valid_token.team_rpm_limit, + tpd_limit=valid_token.team_tpd_limit, + blocked=valid_token.team_blocked, + models=_token_team_models(valid_token), + metadata=valid_token.team_metadata, + object_permission_id=valid_token.team_object_permission_id, + object_permission=await _resolve_object_permission_for_unresolvable_team( + object_permission_id=valid_token.team_object_permission_id, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, proxy_logging_obj=proxy_logging_obj, - ) - if _project_obj is not None: - valid_token.project_metadata = _project_obj.metadata - valid_token.project_alias = _project_obj.project_alias + ), + ) - global_proxy_spend = None - if litellm.max_budget > 0 and prisma_client is not None: # user set proxy max budget - cache_key: Final = GLOBAL_PROXY_SPEND_CACHE_KEY - with tracer.trace("litellm.proxy.auth.get_global_proxy_spend"): - global_proxy_spend = await fetch_global_spend_with_event_coordination( # rebind-ok: pre-existing rebinding on a rename-only line - cache_key=cache_key, - user_api_key_cache=user_api_key_cache, + +async def _team_object_or_none( + team_id: str, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Span | None, + proxy_logging_obj: ProxyLogging, +) -> LiteLLM_TeamTableCachedObj | None: + try: + with tracer.trace("litellm.proxy.auth.get_team_object"): + return await get_team_object( + team_id=team_id, prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, ) + except HTTPException: + return None - if global_proxy_spend is not None: - call_info: Final = CallInfo( - token=valid_token.token, - spend=global_proxy_spend, - max_budget=litellm.max_budget, - user_id=litellm_proxy_admin_name, - team_id=valid_token.team_id, - event_group=Litellm_EntityType.PROXY, + +async def _master_key_auth( + effects: _BuilderEffects, admin_user_id: str, route: str, start_time: datetime +) -> UserAPIKeyAuth: + # A stable alias replaces the raw master key so neither the key nor its hash reaches spend logs, + # metrics labels, audit trails or rate-limit buckets. + master_key_auth: Final = await _return_user_api_key_auth_obj( + user_obj=None, + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key=LITELLM_PROXY_MASTER_KEY_ALIAS, + parent_otel_span=effects.parent_otel_span, + valid_token_dict={**effects.end_user.params, "user_id": admin_user_id}, + route=route, + start_time=start_time, + ) + asyncio.create_task( + cache_key_object( + hashed_token=hash_token(effects.token), + user_api_key_obj=master_key_auth, + user_api_key_cache=effects.user_api_key_cache, + proxy_logging_obj=effects.proxy_logging_obj, + ) + ) + with_end_user: Final = _with_end_user(master_key_auth, effects.end_user, effects.parent_otel_span) + with_end_user.via_virtual_key = True + return with_end_user + + +async def _authenticated( + outcome: Outcome, + effects: _BuilderEffects, + admin_user_id: str, + route: str, + start_time: datetime, + jwt_auth_enabled: bool, +) -> UserAPIKeyAuth: + match outcome: + case Deny(denial=denial): + raise denial_exception( + denial, token=effects.token, route=route, jwt_auth_enabled=jwt_auth_enabled, key=effects.resolved_key ) - asyncio.create_task( - proxy_logging_obj.budget_alerts( - type="proxy_budget", - user_info=call_info, - ) + case Allow(grant=grant): + return await _granted(grant, effects, admin_user_id, route, start_time) + + +async def _granted( + grant: Grant, effects: _BuilderEffects, admin_user_id: str, route: str, start_time: datetime +) -> UserAPIKeyAuth: + match grant: + case NoMasterKey(token=token): + return UserAPIKeyAuth( + api_key=token, user_role=LitellmUserRoles.INTERNAL_USER, parent_otel_span=effects.parent_otel_span ) - if valid_token.token is None: - raise HTTPException(401, detail="Invalid API key, no token associated") + case UserAuthRoute(): + return UserAPIKeyAuth() + case MasterKey(): + return await _master_key_auth(effects, admin_user_id, route, start_time) + case CachedAdminKey(): + return _with_end_user(effects.key, effects.end_user, effects.parent_otel_span) + case VirtualKey(budget_throttle_pct=budget_throttle_pct): + return await _virtual_key_auth(effects, budget_throttle_pct, route, start_time) + case _: + assert_never(grant) - valid_token_dict: Final = valid_token.model_dump(exclude_none=True) - valid_token_dict.pop("token", None) - # budget_throttle_pct is excluded from model_dump (it must not leak - # into serialized responses), so carry the request-scoped decision - # forward by hand to the auth object the rate limiter receives. - if valid_token.budget_throttle_pct is not None: - valid_token_dict["budget_throttle_pct"] = valid_token.budget_throttle_pct - if end_user.end_user_object is not None or litellm.max_end_user_budget_id is not None: - valid_token_dict.update(end_user.params) - if end_user.end_user_object is not None: - valid_token_dict["end_user_object_permission"] = end_user.end_user_object.object_permission +async def _virtual_key_auth( + effects: _BuilderEffects, budget_throttle_pct: float | None, route: str, start_time: datetime +) -> UserAPIKeyAuth: + key: Final = effects.key + token: Final = key.token + if token is None: + raise RuntimeError("the auth core granted a virtual key without a token") + end_user: Final = effects.end_user + # budget_throttle_pct is excluded from model_dump so it never leaks into serialized responses. + throttle: Final = {} if budget_throttle_pct is None else {"budget_throttle_pct": budget_throttle_pct} + end_user_params: Final = ( + end_user.params if end_user.end_user_object is not None or litellm.max_end_user_budget_id is not None else {} + ) + end_user_permission: Final = ( + {} + if end_user.end_user_object is None + else {"end_user_object_permission": end_user.end_user_object.object_permission} + ) virtual_key_auth_obj: Final = await _return_user_api_key_auth_obj( - user_obj=user_obj, - api_key=valid_token.token, - parent_otel_span=parent_otel_span, - valid_token_dict=valid_token_dict, + user_obj=effects.owner, + api_key=token, + parent_otel_span=effects.parent_otel_span, + valid_token_dict={ + **{k: v for k, v in key.model_dump(exclude_none=True).items() if k != "token"}, + **throttle, + **end_user_params, + **end_user_permission, + }, route=route, start_time=start_time, ) @@ -4793,6 +4820,7 @@ async def _authorize_resolved_key( route: str, ) -> UserAPIKeyAuth: from litellm.proxy.proxy_server import ( + litellm_proxy_admin_name, master_key, prisma_client, proxy_logging_obj, @@ -4819,25 +4847,14 @@ async def _authorize_resolved_key( ), effects, ) - match outcome: - case Allow(grant=VirtualKey(budget_checks_skipped=budget_checks_skipped, budget_throttle_pct=throttle_pct)): - if throttle_pct is not None: - effects.key.budget_throttle_pct = throttle_pct - return await _finish_virtual_key_auth( - request=request, - request_data=request_data, - valid_token=effects.key, - route=route, - start_time=start_time, - parent_otel_span=None, - end_user=effects.end_user, - user_obj=effects.owner, - skip_budget_checks=budget_checks_skipped, - ) - case Allow(grant=grant): - raise RuntimeError(f"a resolved key check granted {grant}") - case Deny(denial=denial): - raise denial_exception(denial, token=key_hash, route=route, jwt_auth_enabled=False, key=effects.key) + return await _authenticated( + outcome, + effects, + admin_user_id=litellm_proxy_admin_name, + route=route, + start_time=start_time, + jwt_auth_enabled=False, + ) async def authorize_internal_virtual_key( diff --git a/tests/unit/_auth/test_python_core.py b/tests/unit/_auth/test_python_core.py index d2245199c98..d48d0545d0c 100644 --- a/tests/unit/_auth/test_python_core.py +++ b/tests/unit/_auth/test_python_core.py @@ -14,6 +14,7 @@ from litellm._auth.decision import ( KeyBlocked, KeyBudgetExceeded, KeyFacts, + KeyWithoutToken, MalformedKey, MasterKey, MasterKeyOnlyRoute, @@ -46,10 +47,14 @@ SETTINGS: Final = Settings( allow_user_auth=False, has_db=True, master_key_only_routes=frozenset({"/global/spend/reset"}), + proxy_max_budget=0.0, ) PLAIN_KEY: Final = KeyFacts( + has_token=True, blocked=False, expires=None, + spend=0.0, + soft_budget=None, max_budget=None, team_member_spend=None, user_id=None, @@ -62,6 +67,8 @@ CHECKED: Final = ( "enforce_model_access", "check_budget_exemption", "send_key_budget_alerts", + "enforce_model_budgets", + "load_key_context", ) Call = tuple[str, object] @@ -75,6 +82,7 @@ class RecordingHandler: exempt: bool = False team_member_budget: TeamMemberBudget | None = None key_spend: float = 0.0 + proxy_spend: float | None = None calls: tuple[Call, ...] = () def _record(self, name: str, payload: object = None) -> None: @@ -128,6 +136,22 @@ class RecordingHandler: async def alert_key_budget(self, spend: float) -> None: self._record("alert_key_budget", spend) + async def alert_key_soft_budget(self) -> None: + self._record("alert_key_soft_budget") + + async def enforce_model_budgets(self) -> None: + self._record("enforce_model_budgets") + + async def load_key_context(self) -> None: + self._record("load_key_context") + + async def read_proxy_spend(self) -> float | None: + self._record("read_proxy_spend") + return self.proxy_spend + + async def alert_proxy_budget(self, spend: float) -> None: + self._record("alert_proxy_budget", spend) + def _facts(token: str | None, route: str = "/chat/completions", is_llm_api_route: bool = True) -> RequestFacts: return RequestFacts(route=route, token=token, is_llm_api_route=is_llm_api_route) @@ -165,24 +189,20 @@ async def _check( return await _run(PythonCore().start_resolved_key(facts, settings, NOW), handler) -def _allowed(budget_checks_skipped: bool = False, budget_throttle_pct: float | None = None) -> Outcome: - return Allow( - VirtualKey( - loaded_from_db=True, budget_checks_skipped=budget_checks_skipped, budget_throttle_pct=budget_throttle_pct - ) - ) +def _allowed(budget_throttle_pct: float | None = None) -> Outcome: + return Allow(VirtualKey(budget_throttle_pct=budget_throttle_pct)) @pytest.mark.asyncio @pytest.mark.parametrize( ("route", "token", "settings", "expected"), [ - ("/chat/completions", None, Settings(None, False, True, frozenset()), Allow(NoMasterKey(token=None))), - ("/chat/completions", "sk-x", Settings(None, False, False, frozenset()), Allow(NoMasterKey(token="sk-x"))), + ("/chat/completions", None, Settings(None, False, True, frozenset(), 0.0), Allow(NoMasterKey(token=None))), + ("/chat/completions", "sk-x", Settings(None, False, False, frozenset(), 0.0), Allow(NoMasterKey(token="sk-x"))), ("/chat/completions", None, SETTINGS, Deny(MissingKey())), ("/chat/completions", "", SETTINGS, Deny(MalformedKey())), - ("/user/auth", "", Settings(MASTER, True, True, frozenset()), Deny(MalformedKey())), - ("/user/auth", "sk-x", Settings(MASTER, True, True, frozenset()), Allow(UserAuthRoute())), + ("/user/auth", "", Settings(MASTER, True, True, frozenset(), 0.0), Deny(MalformedKey())), + ("/user/auth", "sk-x", Settings(MASTER, True, True, frozenset(), 0.0), Allow(UserAuthRoute())), ("/user/auth", "sk-x", SETTINGS, Deny(UserAuthDisabled())), ("/user/auth", MASTER, SETTINGS, Deny(UserAuthDisabled())), ], @@ -203,7 +223,7 @@ async def test_master_key_resolves_end_user_without_key_prefetch(route: str): async def test_unknown_virtual_key_is_loaded_from_the_db_then_checked(): handler: Final = RecordingHandler() outcome, calls = await _decide("sk-x", handler=handler) - assert outcome == Allow(VirtualKey(loaded_from_db=True, budget_checks_skipped=False, budget_throttle_pct=None)) + assert outcome == _allowed() assert calls == ("resolve_end_user", "resolve_presented_key", "load_key", *CHECKED) assert handler.calls[2] == ("load_key", hash_token("sk-x")) assert handler.calls[3] == ("prepare_key", True) @@ -213,7 +233,7 @@ async def test_unknown_virtual_key_is_loaded_from_the_db_then_checked(): async def test_cached_virtual_key_is_checked_without_a_db_load(): handler: Final = RecordingHandler(presented=NonAdminKey()) outcome, calls = await _decide("sk-x", handler=handler) - assert outcome == Allow(VirtualKey(loaded_from_db=False, budget_checks_skipped=False, budget_throttle_pct=None)) + assert outcome == _allowed() assert calls == ("resolve_end_user", "resolve_presented_key", *CHECKED) assert handler.calls[2] == ("prepare_key", False) @@ -221,7 +241,7 @@ async def test_cached_virtual_key_is_checked_without_a_db_load(): @pytest.mark.asyncio @pytest.mark.parametrize("expires", [None, NOW + timedelta(seconds=1), NOW]) async def test_unexpired_cached_admin_key_skips_route_db_and_key_checks(expires: datetime | None): - no_db: Final = Settings(MASTER, False, False, frozenset({"/global/spend/reset"})) + no_db: Final = Settings(MASTER, False, False, frozenset({"/global/spend/reset"}), 0.0) outcome, calls = await _decide( "sk-admin", route="/global/spend/reset", @@ -251,7 +271,7 @@ async def test_non_admin_key_cannot_use_master_key_only_routes(presented: Presen @pytest.mark.asyncio @pytest.mark.parametrize("presented", [None, NonAdminKey()]) async def test_virtual_key_needs_a_database(presented: PresentedKey | None): - no_db: Final = Settings(MASTER, False, False, frozenset()) + no_db: Final = Settings(MASTER, False, False, frozenset(), 0.0) outcome, _ = await _decide("sk-x", settings=no_db, handler=RecordingHandler(presented=presented)) assert outcome == Deny(NoDatabase()) @@ -292,10 +312,12 @@ async def test_active_or_missing_owner_continues(owner: OwnerFacts | None): @pytest.mark.asyncio async def test_exempt_request_skips_every_budget_check_but_not_expiry(): - key: Final = replace(PLAIN_KEY, user_id="u1", team_id="t1", team_member_spend=99.0, max_budget=1.0) + key: Final = replace( + PLAIN_KEY, user_id="u1", team_id="t1", team_member_spend=99.0, max_budget=1.0, spend=5.0, soft_budget=1.0 + ) assert await _check(key, exempt=True) == ( - _allowed(budget_checks_skipped=True), - ("inspect_key", "enforce_model_access", "load_key_owner", "check_budget_exemption"), + _allowed(), + ("inspect_key", "enforce_model_access", "load_key_owner", "check_budget_exemption", "load_key_context"), ) expired: Final = replace(key, expires=NOW - timedelta(seconds=1)) outcome, _ = await _check(expired, exempt=True) @@ -343,7 +365,7 @@ async def test_team_member_budget_is_read_only_for_a_member_with_tracked_spend(k @pytest.mark.asyncio async def test_team_member_budget_needs_a_database(): key: Final = replace(PLAIN_KEY, user_id="u1", team_id="t1", team_member_spend=0.0) - no_db: Final = Settings(MASTER, False, False, frozenset()) + no_db: Final = Settings(MASTER, False, False, frozenset(), 0.0) _, calls = await _check(key, settings=no_db, team_member_budget=TeamMemberBudget(spend=1.0, max_budget=0.5)) assert "read_team_member_budget" not in calls @@ -363,20 +385,20 @@ async def test_team_member_budget_is_checked_before_expiry(): [ (5.0, 5.0, None, Deny(KeyBudgetExceeded(spend=5.0, max_budget=5.0))), (6.0, 5.0, None, Deny(KeyBudgetExceeded(spend=6.0, max_budget=5.0))), - (4.9, 5.0, None, Allow(VirtualKey(loaded_from_db=True, budget_checks_skipped=False, budget_throttle_pct=None))), - (6.0, 5.0, 0.5, Allow(VirtualKey(loaded_from_db=True, budget_checks_skipped=False, budget_throttle_pct=0.5))), - (4.9, 5.0, 0.5, Allow(VirtualKey(loaded_from_db=True, budget_checks_skipped=False, budget_throttle_pct=None))), + (4.9, 5.0, None, _allowed()), + (6.0, 5.0, 0.5, _allowed(budget_throttle_pct=0.5)), + (4.9, 5.0, 0.5, _allowed()), ( 1e9, math.inf, None, - Allow(VirtualKey(loaded_from_db=True, budget_checks_skipped=False, budget_throttle_pct=None)), + _allowed(), ), ( 1e9, math.nan, None, - Allow(VirtualKey(loaded_from_db=True, budget_checks_skipped=False, budget_throttle_pct=None)), + _allowed(), ), ], ) @@ -387,8 +409,9 @@ async def test_key_budget_denies_or_throttles_at_the_limit( handler: Final = RecordingHandler(key=key, key_spend=spend) outcome, calls = await _run(PythonCore().start_resolved_key(_facts("hash"), SETTINGS, NOW), handler) assert outcome == expected - assert calls[-3:] == ("send_key_budget_alerts", "read_key_spend", "alert_key_budget") - assert handler.calls[-1] == ("alert_key_budget", spend) + assert calls[3:6] == ("send_key_budget_alerts", "read_key_spend", "alert_key_budget") + assert handler.calls[5] == ("alert_key_budget", spend) + assert ("enforce_model_budgets" in calls) == isinstance(expected, Allow) @pytest.mark.asyncio @@ -396,4 +419,72 @@ async def test_key_budget_is_enforced_only_on_llm_routes(): key: Final = replace(PLAIN_KEY, max_budget=1.0) outcome, calls = await _check(key, is_llm_api_route=False, key_spend=2.0) assert outcome == _allowed() - assert calls[-1] == "send_key_budget_alerts" + assert "read_key_spend" not in calls + assert "enforce_model_budgets" in calls + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("spend", "soft_budget", "alerted"), + [ + (5.0, 5.0, True), + (6.0, 5.0, True), + (4.9, 5.0, False), + (10.0, None, False), + (10.0, 0.0, False), + (10.0, math.nan, False), + ], +) +async def test_soft_budget_alerts_at_the_threshold_before_model_budgets( + spend: float, soft_budget: float | None, alerted: bool +): + outcome, calls = await _check(replace(PLAIN_KEY, spend=spend, soft_budget=soft_budget)) + assert outcome == _allowed() + soft_alert: Final = ("alert_key_soft_budget",) if alerted else () + assert calls[3:] == ("send_key_budget_alerts", *soft_alert, "enforce_model_budgets", "load_key_context") + + +@pytest.mark.asyncio +async def test_throttled_key_still_gets_soft_and_model_budget_checks(): + key: Final = replace(PLAIN_KEY, max_budget=5.0, budget_throttle_pct=0.5, spend=6.0, soft_budget=1.0) + outcome, calls = await _check(key, key_spend=6.0) + assert outcome == _allowed(budget_throttle_pct=0.5) + assert calls[-4:] == ("alert_key_budget", "alert_key_soft_budget", "enforce_model_budgets", "load_key_context") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("exempt", [False, True]) +@pytest.mark.parametrize( + ("proxy_max_budget", "has_db", "proxy_spend", "tail"), + [ + (0.0, True, 50.0, (("load_key_context", None),)), + (-1.0, True, 50.0, (("load_key_context", None),)), + (math.nan, True, 50.0, (("load_key_context", None),)), + (10.0, False, 50.0, (("load_key_context", None),)), + (10.0, True, None, (("load_key_context", None), ("read_proxy_spend", None))), + ( + 10.0, + True, + 50.0, + (("load_key_context", None), ("read_proxy_spend", None), ("alert_proxy_budget", 50.0)), + ), + ], +) +async def test_proxy_budget_only_alerts_and_runs_for_exempt_requests_too( + exempt: bool, proxy_max_budget: float, has_db: bool, proxy_spend: float | None, tail: tuple[Call, ...] +): + settings: Final = replace(SETTINGS, proxy_max_budget=proxy_max_budget, has_db=has_db) + handler: Final = RecordingHandler(exempt=exempt, proxy_spend=proxy_spend) + outcome, _ = await _run(PythonCore().start_resolved_key(_facts("hash"), settings, NOW), handler) + assert outcome == _allowed() + assert handler.calls[-len(tail) :] == tail + + +@pytest.mark.asyncio +@pytest.mark.parametrize("exempt", [False, True]) +async def test_key_without_token_is_denied_after_its_context_loads(exempt: bool): + settings: Final = replace(SETTINGS, proxy_max_budget=10.0) + handler: Final = RecordingHandler(key=replace(PLAIN_KEY, has_token=False), exempt=exempt, proxy_spend=1.0) + outcome, calls = await _run(PythonCore().start_resolved_key(_facts("hash"), settings, NOW), handler) + assert outcome == Deny(KeyWithoutToken()) + assert calls[-3:] == ("load_key_context", "read_proxy_spend", "alert_proxy_budget") diff --git a/tests/unit/proxy/auth/test_user_api_key_auth.py b/tests/unit/proxy/auth/test_user_api_key_auth.py index d7e58974b04..efa0ee912d8 100644 --- a/tests/unit/proxy/auth/test_user_api_key_auth.py +++ b/tests/unit/proxy/auth/test_user_api_key_auth.py @@ -1600,8 +1600,6 @@ def test_jwt_path_enforces_the_user_model_budget_before_returning(): tree = ast.parse( textwrap.dedent(inspect.getsource(auth_module._user_api_key_auth_builder)) + "\n" - + textwrap.dedent(inspect.getsource(auth_module._finish_virtual_key_auth)) - + "\n" + textwrap.dedent(inspect.getsource(auth_module._BuilderEffects)) ) @@ -1647,8 +1645,6 @@ def test_every_jwt_branch_carries_the_user_model_budget(): tree = ast.parse( textwrap.dedent(inspect.getsource(auth_module._user_api_key_auth_builder)) + "\n" - + textwrap.dedent(inspect.getsource(auth_module._finish_virtual_key_auth)) - + "\n" + textwrap.dedent(inspect.getsource(auth_module._BuilderEffects)) ) @@ -1774,8 +1770,10 @@ def test_zero_cost_models_skip_the_user_budget_check_on_every_path(): whether the caller used a JWT or a virtual key, and makes that log a lie. Structural rather than behavioural on purpose: the defect is a call site - sitting outside a guard, and driving both auth paths to a zero-cost model - would prove it for the two requests exercised rather than for every site. + sitting outside a guard, and driving the JWT path to a zero-cost model + would prove it for the request exercised rather than for every site. The + virtual-key path's exemption is decided by the auth core and covered in + tests/unit/_auth/test_python_core.py. """ import ast import inspect @@ -1783,13 +1781,7 @@ def test_zero_cost_models_skip_the_user_budget_check_on_every_path(): from litellm.proxy.auth import user_api_key_auth as auth_module - tree = ast.parse( - textwrap.dedent(inspect.getsource(auth_module._user_api_key_auth_builder)) - + "\n" - + textwrap.dedent(inspect.getsource(auth_module._finish_virtual_key_auth)) - + "\n" - + textwrap.dedent(inspect.getsource(auth_module._BuilderEffects)) - ) + tree = ast.parse(textwrap.dedent(inspect.getsource(auth_module._user_api_key_auth_builder))) def guarded_by_skip(node: ast.AST, target: ast.AST) -> bool: for parent in ast.walk(node): @@ -1809,11 +1801,9 @@ def test_zero_cost_models_skip_the_user_budget_check_on_every_path(): calls = [ node for node in ast.walk(tree) - if isinstance(node, ast.Call) - and isinstance(node.func, ast.Name) - and node.func.id in {"_check_user_model_budget", "_check_soft_and_model_budgets"} + if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == "_check_user_model_budget" ] - assert len(calls) == 2, f"expected the JWT and virtual-key call sites, found {len(calls)}" + assert len(calls) == 1, f"expected the JWT call site, found {len(calls)}" unguarded = [c for c in calls if not guarded_by_skip(tree, c)] assert not unguarded, ( @@ -1935,8 +1925,6 @@ def test_mapped_key_jwt_falls_through_to_the_shared_user_budget_attach(): tree = ast.parse( textwrap.dedent(inspect.getsource(auth_module._user_api_key_auth_builder)) + "\n" - + textwrap.dedent(inspect.getsource(auth_module._finish_virtual_key_auth)) - + "\n" + textwrap.dedent(inspect.getsource(auth_module._BuilderEffects)) )