diff --git a/litellm/_auth/AGENTS.md b/litellm/_auth/AGENTS.md index dba12c1c412..b6e54cac36f 100644 --- a/litellm/_auth/AGENTS.md +++ b/litellm/_auth/AGENTS.md @@ -12,4 +12,5 @@ - 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 checks in `validate_resolved_virtual_key` and the centralized common checks still run as projection after `Allow(VirtualKey(...))`. 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 +- 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 diff --git a/litellm/_auth/decision.py b/litellm/_auth/decision.py index 648877f18ed..9cee892f79c 100644 --- a/litellm/_auth/decision.py +++ b/litellm/_auth/decision.py @@ -8,6 +8,7 @@ from typing import Protocol, TypeAlias class RequestFacts: route: str token: str | None + is_llm_api_route: bool @dataclass(frozen=True, slots=True) @@ -41,6 +42,8 @@ class CachedAdminKey: @dataclass(frozen=True, slots=True) class VirtualKey: loaded_from_db: bool + budget_checks_skipped: bool + budget_throttle_pct: float | None Grant: TypeAlias = NoMasterKey | UserAuthRoute | MasterKey | CachedAdminKey | VirtualKey @@ -65,6 +68,7 @@ class UserAuthDisabled: class ExpiredKey: expires: datetime now: datetime + key_label: str @dataclass(frozen=True, slots=True) @@ -82,8 +86,42 @@ class NotAVirtualKey: pass +@dataclass(frozen=True, slots=True) +class KeyBlocked: + pass + + +@dataclass(frozen=True, slots=True) +class OwnerDeactivated: + user_id: str + + +@dataclass(frozen=True, slots=True) +class TeamMemberBudgetExceeded: + user_id: str + team_id: str + spend: float + max_budget: float + + +@dataclass(frozen=True, slots=True) +class KeyBudgetExceeded: + spend: float + max_budget: float + + Denial: TypeAlias = ( - MissingKey | MalformedKey | UserAuthDisabled | ExpiredKey | MasterKeyOnlyRoute | NoDatabase | NotAVirtualKey + MissingKey + | MalformedKey + | UserAuthDisabled + | ExpiredKey + | MasterKeyOnlyRoute + | NoDatabase + | NotAVirtualKey + | KeyBlocked + | OwnerDeactivated + | TeamMemberBudgetExceeded + | KeyBudgetExceeded ) @@ -113,6 +151,28 @@ class NonAdminKey: PresentedKey: TypeAlias = AdminKey | NonAdminKey +@dataclass(frozen=True, slots=True) +class KeyFacts: + blocked: bool + expires: datetime | None + max_budget: float | None + team_member_spend: float | None + user_id: str | None + team_id: str | None + budget_throttle_pct: float | None + + +@dataclass(frozen=True, slots=True) +class OwnerFacts: + deactivated: bool + + +@dataclass(frozen=True, slots=True) +class TeamMemberBudget: + spend: float + max_budget: float + + @dataclass(frozen=True, slots=True) class ResolveEndUser: prefetch_key: bool @@ -137,9 +197,85 @@ class LoadKey: resume: "Callable[[], Step]" -Effect: TypeAlias = ResolveEndUser | ResolvePresentedKey | EvictPresentedKey | LoadKey +@dataclass(frozen=True, slots=True) +class PrepareKey: + loaded_from_db: bool + resume: "Callable[[], Step]" + + +@dataclass(frozen=True, slots=True) +class InspectKey: + resume: "Callable[[KeyFacts], Step]" + + +@dataclass(frozen=True, slots=True) +class EnforceModelAccess: + resume: "Callable[[], Step]" + + +@dataclass(frozen=True, slots=True) +class LoadKeyOwner: + user_id: str + resume: "Callable[[OwnerFacts | None], Step]" + + +@dataclass(frozen=True, slots=True) +class CheckBudgetExemption: + resume: "Callable[[bool], Step]" + + +@dataclass(frozen=True, slots=True) +class ReadTeamMemberBudget: + user_id: str + team_id: str + resume: "Callable[[TeamMemberBudget | None], Step]" + + +@dataclass(frozen=True, slots=True) +class AlertTeamMemberBudget: + user_id: str + team_id: str + budget: TeamMemberBudget + resume: "Callable[[], Step]" + + +@dataclass(frozen=True, slots=True) +class SendKeyBudgetAlerts: + resume: "Callable[[], Step]" + + +@dataclass(frozen=True, slots=True) +class ReadKeySpend: + max_budget: float + resume: "Callable[[float], Step]" + + +@dataclass(frozen=True, slots=True) +class AlertKeyBudget: + spend: float + resume: "Callable[[], Step]" + + +Effect: TypeAlias = ( + ResolveEndUser + | ResolvePresentedKey + | EvictPresentedKey + | LoadKey + | PrepareKey + | InspectKey + | EnforceModelAccess + | LoadKeyOwner + | CheckBudgetExemption + | ReadTeamMemberBudget + | AlertTeamMemberBudget + | SendKeyBudgetAlerts + | ReadKeySpend + | AlertKeyBudget +) Step: TypeAlias = Outcome | Effect class AuthCore(Protocol): def start(self, facts: RequestFacts, settings: Settings, now: datetime) -> Step: ... + + def start_resolved_key(self, facts: RequestFacts, settings: Settings, now: datetime) -> Step: ... diff --git a/litellm/_auth/driver.py b/litellm/_auth/driver.py index 1e8b0804f98..9fd1e7cc47d 100644 --- a/litellm/_auth/driver.py +++ b/litellm/_auth/driver.py @@ -1,15 +1,29 @@ from typing import Protocol, assert_never from litellm._auth.decision import ( + AlertKeyBudget, + AlertTeamMemberBudget, Allow, + CheckBudgetExemption, Deny, + Effect, + EnforceModelAccess, EvictPresentedKey, + InspectKey, + KeyFacts, LoadKey, + LoadKeyOwner, Outcome, + OwnerFacts, + PrepareKey, PresentedKey, + ReadKeySpend, + ReadTeamMemberBudget, ResolveEndUser, ResolvePresentedKey, + SendKeyBudgetAlerts, Step, + TeamMemberBudget, ) @@ -22,21 +36,72 @@ class EffectHandler(Protocol): async def load_key(self, hashed_token: str) -> None: ... + async def prepare_key(self, loaded_from_db: bool) -> None: ... + + async def inspect_key(self) -> KeyFacts: ... + + async def enforce_model_access(self) -> None: ... + + async def load_key_owner(self, user_id: str) -> OwnerFacts | None: ... + + async def check_budget_exemption(self) -> bool: ... + + async def read_team_member_budget(self, user_id: str, team_id: str) -> TeamMemberBudget | None: ... + + async def alert_team_member_budget(self, user_id: str, team_id: str, budget: TeamMemberBudget) -> None: ... + + async def send_key_budget_alerts(self) -> None: ... + + async def read_key_spend(self, max_budget: float) -> float: ... + + async def alert_key_budget(self, spend: float) -> None: ... + + +async def _answer(step: Effect, handler: EffectHandler) -> Step: + match step: + case ResolveEndUser(prefetch_key=prefetch_key, resume=resume): + await handler.resolve_end_user(prefetch_key) + return resume() + case ResolvePresentedKey(hashed_token=hashed_token, resume=resume): + return resume(await handler.resolve_presented_key(hashed_token)) + case EvictPresentedKey(hashed_token=hashed_token, resume=resume): + await handler.evict_presented_key(hashed_token) + return resume() + case LoadKey(hashed_token=hashed_token, resume=resume): + await handler.load_key(hashed_token) + return resume() + case PrepareKey(loaded_from_db=loaded_from_db, resume=resume): + await handler.prepare_key(loaded_from_db) + return resume() + case InspectKey(resume=resume): + return resume(await handler.inspect_key()) + case EnforceModelAccess(resume=resume): + await handler.enforce_model_access() + return resume() + case LoadKeyOwner(user_id=user_id, resume=resume): + return resume(await handler.load_key_owner(user_id)) + case CheckBudgetExemption(resume=resume): + return resume(await handler.check_budget_exemption()) + 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): + await handler.alert_team_member_budget(user_id, team_id, budget) + return resume() + case SendKeyBudgetAlerts(resume=resume): + await handler.send_key_budget_alerts() + return resume() + case ReadKeySpend(max_budget=max_budget, resume=resume): + return resume(await handler.read_key_spend(max_budget)) + case AlertKeyBudget(spend=spend, resume=resume): + await handler.alert_key_budget(spend) + return resume() + case _: + assert_never(step) + async def drive(step: Step, handler: EffectHandler) -> Outcome: match step: case Allow() | Deny(): return step - case ResolveEndUser(prefetch_key=prefetch_key, resume=resume): - await handler.resolve_end_user(prefetch_key) - return await drive(resume(), handler) - case ResolvePresentedKey(hashed_token=hashed_token, resume=resume): - return await drive(resume(await handler.resolve_presented_key(hashed_token)), handler) - case EvictPresentedKey(hashed_token=hashed_token, resume=resume): - await handler.evict_presented_key(hashed_token) - return await drive(resume(), handler) - case LoadKey(hashed_token=hashed_token, resume=resume): - await handler.load_key(hashed_token) - return await drive(resume(), handler) case _: - assert_never(step) + return await drive(await _answer(step, handler), handler) diff --git a/litellm/_auth/python_core.py b/litellm/_auth/python_core.py index bb160e214e5..20bc1d7e086 100644 --- a/litellm/_auth/python_core.py +++ b/litellm/_auth/python_core.py @@ -1,14 +1,25 @@ +import math +from dataclasses import dataclass from datetime import datetime from typing import Final from litellm._auth.decision import ( AdminKey, + AlertKeyBudget, + AlertTeamMemberBudget, Allow, CachedAdminKey, + CheckBudgetExemption, Deny, + EnforceModelAccess, EvictPresentedKey, ExpiredKey, + InspectKey, + KeyBlocked, + KeyBudgetExceeded, + KeyFacts, LoadKey, + LoadKeyOwner, MalformedKey, MasterKey, MasterKeyOnlyRoute, @@ -17,12 +28,20 @@ from litellm._auth.decision import ( NoMasterKey, NonAdminKey, NotAVirtualKey, + OwnerDeactivated, + OwnerFacts, + PrepareKey, PresentedKey, + ReadKeySpend, + ReadTeamMemberBudget, RequestFacts, ResolveEndUser, ResolvePresentedKey, + SendKeyBudgetAlerts, Settings, Step, + TeamMemberBudget, + TeamMemberBudgetExceeded, UserAuthDisabled, UserAuthRoute, VirtualKey, @@ -55,6 +74,11 @@ 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 "") + ) + def _after_presented_key( presented: PresentedKey | None, @@ -67,7 +91,8 @@ def _after_presented_key( match presented: case AdminKey(expires=expires) if expires is not None and expires < now: return EvictPresentedKey( - hashed_token=hashed_token, resume=lambda: Deny(ExpiredKey(expires=expires, now=now)) + hashed_token=hashed_token, + resume=lambda: Deny(ExpiredKey(expires=expires, now=now, key_label=token)), ) case AdminKey(): return Allow(CachedAdminKey()) @@ -78,7 +103,114 @@ def _after_presented_key( if not settings.has_db: return Deny(NoDatabase()) if presented is not None: - return Allow(VirtualKey(loaded_from_db=False)) + cached: Final = _KeyCheck(facts=facts, settings=settings, now=now, loaded_from_db=False, key_label=token) + return PrepareKey(loaded_from_db=False, resume=lambda: _check_key(cached)) if not token.startswith(_VIRTUAL_KEY_PREFIX): return Deny(NotAVirtualKey()) - return LoadKey(hashed_token=hashed_token, resume=lambda: Allow(VirtualKey(loaded_from_db=True))) + loaded: Final = _KeyCheck(facts=facts, settings=settings, now=now, loaded_from_db=True, 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)) + + +def _after_inspect(check: _KeyCheck, key: KeyFacts) -> Step: + if key.blocked: + return Deny(KeyBlocked()) + return EnforceModelAccess(resume=lambda: _check_owner(check, key)) + + +def _check_owner(check: _KeyCheck, key: KeyFacts) -> Step: + user_id: Final = key.user_id + if user_id is None: + return _check_budget_exemption(check, key) + return LoadKeyOwner(user_id=user_id, resume=lambda owner: _after_owner(check, key, user_id, owner)) + + +def _after_owner(check: _KeyCheck, key: KeyFacts, user_id: str, owner: OwnerFacts | None) -> Step: + if owner is not None and owner.deactivated: + return Deny(OwnerDeactivated(user_id=user_id)) + return _check_budget_exemption(check, key) + + +def _check_budget_exemption(check: _KeyCheck, key: KeyFacts) -> Step: + return CheckBudgetExemption(resume=lambda exempt: _check_team_member_budget(check, key, exempt)) + + +def _check_team_member_budget(check: _KeyCheck, key: KeyFacts, exempt: bool) -> Step: + user_id: Final = key.user_id + team_id: Final = key.team_id + if exempt or key.team_member_spend is None or not check.settings.has_db or user_id is None or team_id is None: + return _check_expiry(check, key, exempt) + return ReadTeamMemberBudget( + user_id=user_id, + team_id=team_id, + resume=lambda budget: _after_team_member_budget(check, key, user_id, team_id, budget), + ) + + +def _after_team_member_budget( + check: _KeyCheck, key: KeyFacts, user_id: str, team_id: str, budget: TeamMemberBudget | None +) -> Step: + if budget is None or budget.spend < budget.max_budget: + return _check_expiry(check, key, exempt=False) + return AlertTeamMemberBudget( + user_id=user_id, + team_id=team_id, + budget=budget, + resume=lambda: Deny( + TeamMemberBudgetExceeded(user_id=user_id, team_id=team_id, spend=budget.spend, max_budget=budget.max_budget) + ), + ) + + +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 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) + return ReadKeySpend( + max_budget=max_budget, + resume=lambda spend: AlertKeyBudget( + spend=spend, resume=lambda: _after_key_spend(check, key, spend, max_budget) + ), + ) + + +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) + if key.budget_throttle_pct is not None: + return _allow(check, budget_checks_skipped=False, budget_throttle_pct=key.budget_throttle_pct) + return Deny(KeyBudgetExceeded(spend=spend, max_budget=max_budget)) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index af60a315c60..74e7367365d 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -6707,6 +6707,53 @@ def _apply_budget_exceeded_throttle(valid_token: UserAPIKeyAuth) -> bool: return True +def send_token_budget_alert( + valid_token: UserAPIKeyAuth, + spend: float, + proxy_logging_obj: ProxyLogging, + user_obj: LiteLLM_UserTable | None, +) -> None: + call_info: Final = CallInfo( + token=valid_token.token, + spend=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, + organization_id=valid_token.org_id, + user_email=user_obj.user_email if user_obj is not None else None, + key_alias=valid_token.key_alias, + event_group=Litellm_EntityType.KEY, + ) + asyncio.create_task( + proxy_logging_obj.budget_alerts( + type="token_budget", + user_info=call_info, + ) + ) + + +def key_budget_exceeded_error( + valid_token: UserAPIKeyAuth, spend: float, max_budget: float +) -> litellm.BudgetExceededError: + # This message is returned to the caller, and key_name has no enforced + # shape (a direct DB write bypasses abbreviate_api_key), so echo it only + # when it still looks masked and fall back to the alias otherwise. + key_label: Final = valid_token.key_alias or "key" + key_descriptor: Final = ( + f"{key_label} ({valid_token.key_name})" + if valid_token.key_name and _MASKED_KEY_NAME_RE.fullmatch(valid_token.key_name) + else key_label + ) + return litellm.BudgetExceededError( + current_cost=spend, + max_budget=max_budget, + message=f"Budget has been exceeded! Key={key_descriptor} Current cost: {spend}, Max budget: {max_budget}", + entity_type=Litellm_EntityType.KEY.value, + entity_id=valid_token.token, + ) + + async def virtual_key_max_budget_check( valid_token: UserAPIKeyAuth, proxy_logging_obj: ProxyLogging, @@ -6721,47 +6768,13 @@ async def virtual_key_max_budget_check( if valid_token.max_budget is not None: from litellm.proxy.proxy_server import get_current_spend - fallback_spend: Final = valid_token.spend or 0.0 - counter_key: Final = f"spend:key:{valid_token.token}" - # Read spend from cross-pod counter (Redis-first) or cached object (fallback) spend: Final = await get_current_spend( - counter_key=counter_key, - fallback_spend=fallback_spend, + counter_key=f"spend:key:{valid_token.token}", + fallback_spend=valid_token.spend or 0.0, max_budget=valid_token.max_budget, ) - - #################################### - # collect information for alerting # - #################################### - - user_email = None - # Check if the token has any user id information - if user_obj is not None: - user_email = user_obj.user_email - - call_info: Final = CallInfo( - token=valid_token.token, - spend=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, - organization_id=valid_token.org_id, - user_email=user_email, - key_alias=valid_token.key_alias, - event_group=Litellm_EntityType.KEY, - ) - asyncio.create_task( - proxy_logging_obj.budget_alerts( - type="token_budget", - user_info=call_info, - ) - ) - - #################################### - # collect information for alerting # - #################################### + send_token_budget_alert(valid_token, spend, proxy_logging_obj, user_obj) # Defense-in-depth (GHSA-2rv4-xv66-fpjg): spend >= NaN is always False, # so a NaN max_budget would silently disable enforcement. Treat a @@ -6769,22 +6782,7 @@ async def virtual_key_max_budget_check( if math.isfinite(valid_token.max_budget) and spend >= valid_token.max_budget: if _apply_budget_exceeded_throttle(valid_token): return - # This message is returned to the caller, and key_name has no enforced - # shape (a direct DB write bypasses abbreviate_api_key), so echo it only - # when it still looks masked and fall back to the alias otherwise. - key_label: Final = valid_token.key_alias or "key" - key_descriptor: Final = ( - f"{key_label} ({valid_token.key_name})" - if valid_token.key_name and _MASKED_KEY_NAME_RE.fullmatch(valid_token.key_name) - else key_label - ) - raise litellm.BudgetExceededError( - current_cost=spend, - max_budget=valid_token.max_budget, - message=f"Budget has been exceeded! Key={key_descriptor} Current cost: {spend}, Max budget: {valid_token.max_budget}", - entity_type=Litellm_EntityType.KEY.value, - entity_id=valid_token.token, - ) + raise key_budget_exceeded_error(valid_token, spend, valid_token.max_budget) _virtual_key_max_budget_check: Final = virtual_key_max_budget_check diff --git a/litellm/proxy/auth/auth_denial.py b/litellm/proxy/auth/auth_denial.py index 69dece93d99..ae678bfbf8e 100644 --- a/litellm/proxy/auth/auth_denial.py +++ b/litellm/proxy/auth/auth_denial.py @@ -2,18 +2,24 @@ from typing import Final, assert_never from fastapi import HTTPException, status +import litellm from litellm._auth.decision import ( Denial, ExpiredKey, + KeyBlocked, + KeyBudgetExceeded, MalformedKey, MasterKeyOnlyRoute, MissingKey, NoDatabase, NotAVirtualKey, + OwnerDeactivated, + TeamMemberBudgetExceeded, UserAuthDisabled, ) from litellm.constants import INVALID_VIRTUAL_KEY_ERROR_MARKER, INVALID_VIRTUAL_KEY_ERROR_MESSAGE -from litellm.proxy._types import ProxyErrorTypes, ProxyException +from litellm.proxy._types import Litellm_EntityType, ProxyErrorTypes, ProxyException, UserAPIKeyAuth +from litellm.proxy.auth.auth_checks import key_budget_exceeded_error from litellm.proxy.auth.auth_utils import abbreviate_api_key from litellm.types.proxy.auth.user_api_key_auth import is_jwt @@ -37,8 +43,12 @@ def _not_a_virtual_key(token: str, jwt_auth_enabled: bool) -> HTTPException: def denial_exception( - denial: Denial, token: str | None, route: str, jwt_auth_enabled: bool -) -> ProxyException | HTTPException: + denial: Denial, + token: str | None, + route: str, + jwt_auth_enabled: bool, + key: UserAPIKeyAuth | None, +) -> ProxyException | HTTPException | litellm.BudgetExceededError: match denial: case MissingKey(): return ProxyException( @@ -59,12 +69,12 @@ def denial_exception( status_code=status.HTTP_403_FORBIDDEN, detail="'allow_user_auth' not set or set to False", ) - case ExpiredKey(expires=expires, now=now): + case ExpiredKey(expires=expires, now=now, key_label=key_label): return ProxyException( message=f"Authentication Error - Expired Key. Key Expiry time {expires} and current time {now}", type=ProxyErrorTypes.expired_key, code=status.HTTP_401_UNAUTHORIZED, - param=abbreviate_api_key(api_key=token or ""), + param=abbreviate_api_key(api_key=key_label), ) case MasterKeyOnlyRoute(): return ProxyException( @@ -82,5 +92,35 @@ def denial_exception( ) case NotAVirtualKey(): return _not_a_virtual_key(token or "", jwt_auth_enabled) + case KeyBlocked(): + return ProxyException( + message="Authentication Error, Key is blocked. Update via `/key/unblock` if you're an admin.", + type=ProxyErrorTypes.auth_error, + param=None, + code=status.HTTP_401_UNAUTHORIZED, + ) + case OwnerDeactivated(user_id=user_id): + return ProxyException( + message=( + f"Authentication Error, User={user_id} has been deactivated via SCIM. " + "Keys owned by this user cannot be used." + ), + type=ProxyErrorTypes.auth_error, + param=None, + code=status.HTTP_401_UNAUTHORIZED, + ) + case TeamMemberBudgetExceeded(user_id=user_id, team_id=team_id, spend=spend, max_budget=max_budget): + return litellm.BudgetExceededError( + current_cost=spend, + max_budget=max_budget, + message=( + f"Budget has been exceeded! TeamMember={user_id}:{team_id} " + f"Current cost: {spend}, Max budget: {max_budget}" + ), + entity_type=Litellm_EntityType.TEAM_MEMBER.value, + entity_id=f"{user_id}:{team_id}", + ) + case KeyBudgetExceeded(spend=spend, max_budget=max_budget): + return key_budget_exceeded_error(key or UserAPIKeyAuth(), spend, max_budget) 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 e61c501ed4b..92e14f9149b 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -36,12 +36,15 @@ from litellm._auth.decision import ( AuthCore, CachedAdminKey, Deny, + KeyFacts, MasterKey, NoMasterKey, NonAdminKey, + OwnerFacts, PresentedKey, RequestFacts, Settings, + TeamMemberBudget, UserAuthRoute, VirtualKey, ) @@ -114,6 +117,7 @@ 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_token_budget_alert, stamp_matched_model_access_groups, team_member_max_budget_alert_check, virtual_key_max_budget_alert_check, @@ -145,6 +149,7 @@ from litellm.proxy.auth.auth_utils import ( request_dispatched_to_provider_pass_through, route_in_additonal_public_routes, ) +from litellm.proxy.auth.budget_throttle import budget_throttle_percentage, should_throttle_budget_exceeded from litellm.proxy.auth.handle_jwt import JWTAuthManager, JWTHandler from litellm.proxy.auth.network import TrustedProxyConfig, resolve_network_context from litellm.proxy.auth.oauth2_check import Oauth2Handler @@ -2288,9 +2293,24 @@ async def _master_key_auth( return with_end_user +def _proxy_model_list() -> list[object] | None: + from litellm.proxy.proxy_server import llm_model_list + + return llm_model_list + + _MASTER_KEY_ONLY_ROUTES: Final = frozenset(LiteLLMRoutes.master_key_only_routes.value) +def _core_settings(master_key: str | None, allow_user_auth: bool, prisma_client: PrismaClient | None) -> Settings: + return Settings( + master_key=master_key, + allow_user_auth=allow_user_auth, + has_db=prisma_client is not None, + master_key_only_routes=_MASTER_KEY_ONLY_ROUTES, + ) + + def _aware_expiry(expires: datetime | str) -> datetime: expiry: Final = expires if isinstance(expires, datetime) else datetime.fromisoformat(expires) return ( @@ -2363,7 +2383,7 @@ class _BuilderEffects: self, *, request: Request, - request_data: Mapping[str, object], + request_data: dict[str, object], # mutable-ok: model access checks rewrite the request model in place token: str, preresolved_key: UserAPIKeyAuth | None, prisma_client: PrismaClient | None, @@ -2383,6 +2403,7 @@ class _BuilderEffects: self._route: Final = route self.end_user: _ResolvedEndUser = _ResolvedEndUser(end_user_id=None, end_user_object=None, budget=None) self.resolved_key: UserAPIKeyAuth | None = preresolved_key + self.owner: LiteLLM_UserTable | None = None @property def key(self) -> UserAPIKeyAuth: @@ -2463,6 +2484,156 @@ class _BuilderEffects: await self._identity_store(check_cache_only=False).resolve(hashed_token=hashed_token) ) + async def prepare_key(self, loaded_from_db: bool) -> None: + if loaded_from_db: + _reset_end_user_fields(self.key, self.end_user) + return + await _refresh_team_fields_from_cache( + 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 inspect_key(self) -> KeyFacts: + key: Final = _update_key_budget_with_temp_budget_increase(self.key) + self.resolved_key = key + return KeyFacts( + blocked=key.blocked is True, + expires=None if key.expires is None else _aware_expiry(key.expires), + max_budget=key.max_budget, + team_member_spend=key.team_member_spend, + user_id=key.user_id, + team_id=key.team_id, + budget_throttle_pct=budget_throttle_percentage() if should_throttle_budget_exceeded(key) else None, + ) + + async def enforce_model_access(self) -> None: + from litellm.proxy.proxy_server import llm_router + + await enforce_key_and_fallback_model_access( + valid_token=self.key, + request_data=self._request_data, + route=self._route, + request=self._request, + llm_model_list=_proxy_model_list(), + llm_router=llm_router, + ) + await _prefetch_referenced_auth_objects( + self.key, + end_user_id=self.end_user.end_user_id, + user_api_key_cache=self._user_api_key_cache, + prisma_client=self._prisma_client, + ) + + async def load_key_owner(self, user_id: str) -> OwnerFacts | None: + try: + with tracer.trace("litellm.proxy.auth.get_user_object"): + owner: Final = await get_user_object( + user_id=user_id, + prisma_client=self._prisma_client, + user_api_key_cache=self._user_api_key_cache, + user_id_upsert=False, + parent_otel_span=self._parent_otel_span, + proxy_logging_obj=self._proxy_logging_obj, + ) + except Exception as e: + verbose_logger.debug("Unable to get user %s from db/cache, continuing without it: %s", user_id, e) + return None + self.owner = owner + if owner is None: + return None + key: Final = self.key + # The joint verification-token view carries the key's columns only, so the user's own per-model budget + # reaches enforcement and the post-call increment through the row fetched here. + key.user_model_max_budget = owner.model_max_budget + return OwnerFacts(deactivated=isinstance(owner.metadata, dict) and owner.metadata.get("scim_active") is False) + + async def check_budget_exemption(self) -> bool: + from litellm.proxy.proxy_server import llm_router + + model: Final = _get_model_from_request_context( + request_data=self._request_data, + route=self._route, + request=self._request, + llm_router=llm_router, + team_id=self.key.team_id, + ) + return request_skips_budget_checks( + route=self._route, model=model, llm_router=llm_router, request_body=self._request_data + ) + + async def read_team_member_budget(self, user_id: str, team_id: str) -> TeamMemberBudget | None: + from litellm.proxy.proxy_server import get_current_spend + + if self._prisma_client is None: + return None + cache_key: Final = team_membership_auth_cache_key(team_id=team_id, user_id=user_id) + cached: Final = await self._user_api_key_cache.async_get_cache(key=cache_key, model_type=LiteLLM_TeamMembership) + membership: Final = cached if cached is not None else await self._load_team_membership(user_id, team_id) + if membership is None or membership.litellm_budget_table is None: + return None + max_budget: Final = membership.litellm_budget_table.effective_max_budget(now=datetime.now(timezone.utc)) + if max_budget is None or max_budget <= 0: + return None + spend: Final = await get_current_spend( + counter_key=f"spend:team_member:{user_id}:{team_id}", + fallback_spend=self.key.team_member_spend or 0.0, + max_budget=max_budget, + ) + return TeamMemberBudget(spend=spend, max_budget=max_budget) + + async def _load_team_membership(self, user_id: str, team_id: str) -> LiteLLM_TeamMembership | None: + if self._prisma_client is None: + return None + row: Final = await TeamMembershipRepository(self._prisma_client).table.find_first( + where={"user_id": user_id, "team_id": team_id}, + include={"litellm_budget_table": True}, + ) + if row is None: + return None + membership: Final = LiteLLM_TeamMembership.model_validate(row.model_dump()) + await self._user_api_key_cache.async_set_cache( + key=team_membership_auth_cache_key(team_id=team_id, user_id=user_id), + value=membership, + model_type=LiteLLM_TeamMembership, + ttl=5, + ) + return membership + + async def alert_team_member_budget(self, user_id: str, team_id: str, budget: TeamMemberBudget) -> None: + # common_checks sends this alert on requests that get past here, so only the rejected request sends it here. + team_member_max_budget_alert_check( + team_id=team_id, + team_alias=self.key.team_alias, + team_metadata=self.key.team_metadata, + organization_id=self.key.org_id, + user_id=user_id, + user_email=self.owner.user_email if self.owner is not None else None, + proxy_logging_obj=self._proxy_logging_obj, + spend=budget.spend, + max_budget=budget.max_budget, + ) + + async def send_key_budget_alerts(self) -> None: + await virtual_key_max_budget_alert_check( + valid_token=self.key, proxy_logging_obj=self._proxy_logging_obj, user_obj=self.owner + ) + + async def read_key_spend(self, max_budget: float) -> float: + from litellm.proxy.proxy_server import get_current_spend + + # Redis-first cross-pod counter, falling back to the key object's recorded spend. + return await get_current_spend( + counter_key=f"spend:key:{self.key.token}", + fallback_spend=self.key.spend or 0.0, + max_budget=max_budget, + ) + + async def alert_key_budget(self, spend: float) -> None: + send_token_budget_alert(self.key, spend, self._proxy_logging_obj, self.owner) + async def user_api_key_auth_builder( request: Request, @@ -2885,23 +3056,28 @@ async def user_api_key_auth_builder( proxy_logging_obj=proxy_logging_obj, route=route, ) - outcome: Final = await drive( - _AUTH_CORE.start( - RequestFacts(route=route, token=api_key), - Settings( - master_key=master_key, - allow_user_auth=general_settings.get("allow_user_auth", False) is True, - has_db=prisma_client is not None, - master_key_only_routes=_MASTER_KEY_ONLY_ROUTES, + try: + outcome: Final = await drive( + _AUTH_CORE.start( + RequestFacts( + route=route, token=api_key, is_llm_api_route=RouteChecks.is_llm_api_route(route=route) + ), + _core_settings( + master_key=master_key, + allow_user_auth=general_settings.get("allow_user_auth", False) is True, + prisma_client=prisma_client, + ), + datetime.now(timezone.utc), ), - datetime.now(timezone.utc), - ), - effects, - ) - valid_token = effects.resolved_key # rebind-ok: the exception handler below reports the resolved identity + effects, + ) + 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) + 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): @@ -2925,28 +3101,19 @@ async def user_api_key_auth_builder( ) case CachedAdminKey(): return _with_end_user(effects.key, effects.end_user, parent_otel_span) - case VirtualKey(loaded_from_db=loaded_from_db): - if not loaded_from_db: - await _refresh_team_fields_from_cache( - valid_token=effects.key, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - ) - return await validate_resolved_virtual_key( + 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=( - _reset_end_user_fields(effects.key, effects.end_user) if loaded_from_db else effects.key - ), - api_key=hash_token(api_key) if loaded_from_db else api_key, + valid_token=effects.key, route=route, start_time=start_time, parent_otel_span=parent_otel_span, - end_user_id=effects.end_user.end_user_id, - end_user_params=effects.end_user.params, - _end_user_object=effects.end_user.end_user_object, + end_user=effects.end_user, + user_obj=effects.owner, + skip_budget_checks=budget_checks_skipped, ) except Exception as e: return await UserAPIKeyAuthExceptionHandler.handle_authentication_error( @@ -2963,417 +3130,203 @@ async def user_api_key_auth_builder( _user_api_key_auth_builder: Final = user_api_key_auth_builder -async def validate_resolved_virtual_key( # noqa: C901 # Preserve ordering of existing shared authorization checks +async def _check_soft_and_model_budgets( request: Request, - request_data: dict[str, object], - valid_token: UserAPIKeyAuth | None, - api_key: str, + 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_id: str | None, - end_user_params: dict[str, object], - _end_user_object: LiteLLM_EndUserTable | None, -) -> UserAPIKeyAuth: - from litellm.proxy.proxy_server import ( - litellm_proxy_admin_name, - llm_model_list, - llm_router, - model_max_budget_limiter, - prisma_client, - proxy_logging_obj, - user_api_key_cache, + user_obj: LiteLLM_UserTable | None, +) -> None: + from litellm.proxy.proxy_server import llm_router, model_max_budget_limiter, prisma_client, proxy_logging_obj + + # 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, ) - if valid_token is not None: - valid_token = _update_key_budget_with_temp_budget_increase(valid_token) + # 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) - user_obj: LiteLLM_UserTable | None = None - valid_token_dict: dict = {} - if valid_token is not None: - # Got Valid Token from Cache, DB - # Run checks for - # 1. If token can call model - ## 1a. If token can call fallback models (if client-side fallbacks given) - # 2. If user_id for this token is in budget - # 3. If the user spend within their own team is within budget - # 4. If 'user' passed to /chat/completions, /embeddings endpoint is in budget - # 5. If token is expired - # 6. If token spend is under Budget for the token - # 7. If token spend per model is under budget per model - # 8. If token spend is under team budget - # 9. If team spend is under team budget - - ## base case ## key is disabled - if valid_token.blocked is True: - raise ProxyException( - message="Authentication Error, Key is blocked. Update via `/key/unblock` if you're an admin.", - type=ProxyErrorTypes.auth_error, - param=None, - code=status.HTTP_401_UNAUTHORIZED, + 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, + 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, ) - await enforce_key_and_fallback_model_access( - valid_token=valid_token, - request_data=request_data, - route=route, - request=request, - llm_model_list=llm_model_list, - llm_router=llm_router, - ) - await _prefetch_referenced_auth_objects( - valid_token, end_user_id=end_user_id, user_api_key_cache=user_api_key_cache, prisma_client=prisma_client - ) - # Check 2. If user_id for this token is in budget - done in common_checks() - if valid_token.user_id is not None: - try: - with tracer.trace("litellm.proxy.auth.get_user_object"): - user_obj = await get_user_object( - user_id=valid_token.user_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - user_id_upsert=False, - parent_otel_span=parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - ) - except Exception as e: - verbose_logger.debug( - "litellm.proxy.auth.user_api_key_auth.py::user_api_key_auth() - Unable to get user from db/cache. Setting user_obj to None. Exception received - %s", - e, - ) - user_obj = None - - if user_obj is not None: - # The joint verification-token view carries the key's columns only, so the - # user's own per-model budget reaches enforcement and the post-call - # increment through the row fetched here. - valid_token.user_model_max_budget = user_obj.model_max_budget - - if ( - user_obj is not None - and isinstance(user_obj.metadata, dict) - and user_obj.metadata.get("scim_active") is False - ): - raise ProxyException( - message=( - f"Authentication Error, User={valid_token.user_id} has been deactivated via SCIM. " - "Keys owned by this user cannot be used." - ), - type=ProxyErrorTypes.auth_error, - param=None, - code=status.HTTP_401_UNAUTHORIZED, - ) - - # Check 2a. Check if model has zero cost - if so, skip all budget checks - model = _get_model_from_request_context( + # 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, ) - skip_budget_checks: Final = request_skips_budget_checks( - route=route, model=model, llm_router=llm_router, request_body=request_data + 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( + valid_token=valid_token, + model_max_budget_limiter=model_max_budget_limiter, + models=current_models, ) - # Check 3. Check if user is in their team budget - if not skip_budget_checks and valid_token.team_member_spend is not None: - _user_id: Final = valid_token.user_id - _team_id: Final = valid_token.team_id - if prisma_client is not None and _user_id is not None and _team_id is not None: - _cache_key: Final = team_membership_auth_cache_key(team_id=_team_id, user_id=_user_id) - team_member_info = await user_api_key_cache.async_get_cache( - key=_cache_key, - model_type=LiteLLM_TeamMembership, - ) - if team_member_info is None: - # read from DB - _db_member: Final = await TeamMembershipRepository(prisma_client).table.find_first( - where={ - "user_id": _user_id, - "team_id": _team_id, - }, - include={"litellm_budget_table": True}, - ) - if _db_member is not None: - team_member_info = LiteLLM_TeamMembership.model_validate(_db_member.model_dump()) - await user_api_key_cache.async_set_cache( - key=_cache_key, - value=team_member_info, - model_type=LiteLLM_TeamMembership, - ttl=5, - ) +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, + ) - if team_member_info is not None and team_member_info.litellm_budget_table is not None: - team_member_budget: Final = team_member_info.litellm_budget_table.effective_max_budget( - now=datetime.now(timezone.utc), - ) - if team_member_budget is not None and team_member_budget > 0: - # Read from cross-pod counter (Redis-first) if available - from litellm.proxy.proxy_server import get_current_spend - - team_member_spend = valid_token.team_member_spend - if valid_token.user_id is not None and valid_token.team_id is not None: - team_member_spend = await get_current_spend( - counter_key=f"spend:team_member:{valid_token.user_id}:{valid_token.team_id}", - fallback_spend=team_member_spend, - max_budget=team_member_budget, - ) - if team_member_spend >= team_member_budget: - # common_checks sends this alert on requests that get past here, so only the - # request rejected here sends it from the builder. - team_member_max_budget_alert_check( - team_id=_team_id, - team_alias=valid_token.team_alias, - team_metadata=valid_token.team_metadata, - organization_id=valid_token.org_id, - user_id=_user_id, - user_email=user_obj.user_email if user_obj is not None else None, - proxy_logging_obj=proxy_logging_obj, - spend=team_member_spend, - max_budget=team_member_budget, - ) - _entity_id: Final = f"{valid_token.user_id}:{valid_token.team_id}" - raise litellm.BudgetExceededError( - current_cost=team_member_spend, - max_budget=team_member_budget, - message=( - f"Budget has been exceeded! TeamMember={_entity_id} " - f"Current cost: {team_member_spend}, Max budget: {team_member_budget}" - ), - entity_type=Litellm_EntityType.TEAM_MEMBER.value, - entity_id=_entity_id, - ) - - # Check 3. If token is expired - if valid_token.expires is not None: - current_time = datetime.now(timezone.utc) - if isinstance(valid_token.expires, datetime): - expiry_time = valid_token.expires - else: - expiry_time = datetime.fromisoformat(valid_token.expires) - if expiry_time.tzinfo is None or expiry_time.tzinfo.utcoffset(expiry_time) is None: - expiry_time = expiry_time.replace(tzinfo=timezone.utc) - verbose_proxy_logger.debug( - "Checking if token expired, expiry time %s and current time %s", expiry_time, current_time + 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 ) - if expiry_time < current_time: - # Token exists but is expired. - raise ProxyException( - message=f"Authentication Error - Expired Key. Key Expiry time {expiry_time} and current time {current_time}", - type=ProxyErrorTypes.expired_key, - code=status.HTTP_401_UNAUTHORIZED, - param=abbreviate_api_key(api_key=api_key), - ) - if not skip_budget_checks: - with tracer.trace("litellm.proxy.auth.budget_checks"): - # Check 4. Max Budget Alert Check (runs before budget enforcement - # so multi-threshold 100% alerts fire on the request that crosses - # max_budget, before BudgetExceededError is raised below) - await virtual_key_max_budget_alert_check( - valid_token=valid_token, - proxy_logging_obj=proxy_logging_obj, - user_obj=user_obj, - ) - - # Check 5. Token Spend is under budget - if RouteChecks.is_llm_api_route(route=route): - await virtual_key_max_budget_check( - valid_token=valid_token, - proxy_logging_obj=proxy_logging_obj, - user_obj=user_obj, - ) - - # 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, + # 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, - ) - current_models = _get_model_names_for_budget_checks(model=current_model) - - if ( - max_budget_per_model is not None - and isinstance(max_budget_per_model, dict) - and len(max_budget_per_model) > 0 - 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, - request_data=request_data, - request=request, - llm_model_list=cast( # cast-ok: router model lists are deployment sequences - Sequence[object] | None, llm_model_list - ), - llm_router=llm_router, - ) - - # 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( - valid_token=valid_token, - model_max_budget_limiter=model_max_budget_limiter, - models=current_models, - ) - - # 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 - - 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, - 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 - - 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, 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 - 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, - ) - asyncio.create_task( - proxy_logging_obj.budget_alerts( - type="proxy_budget", - user_info=call_info, - ) - ) - # Token passed all checks - if valid_token is None: - raise HTTPException(401, detail="Invalid API key") - if valid_token.token is None: - raise HTTPException(401, detail="Invalid API key, no token associated") - api_key = valid_token.token + 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 - valid_token_dict = 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_object is not None or litellm.max_end_user_budget_id is not None: - valid_token_dict.update(end_user_params) - if _end_user_object is not None: - valid_token_dict["end_user_object_permission"] = _end_user_object.object_permission - - # check if token is from litellm-ui, litellm ui makes keys to allow users to login with sso. These keys can only be used for LiteLLM UI functions - # sso/login, ui/login, /key functions and /user functions - # this will never be allowed to call /chat/completions - - if valid_token is None: - # No token was found when looking up in the DB - raise ProxyException( - message="Authentication Error, Invalid proxy server token passed", - type=ProxyErrorTypes.auth_error, - param=None, - code=status.HTTP_401_UNAUTHORIZED, + # 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, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, ) - if valid_token_dict is not None: - virtual_key_auth_obj: Final = await _return_user_api_key_auth_obj( - user_obj=user_obj, - api_key=api_key, - parent_otel_span=parent_otel_span, - valid_token_dict=valid_token_dict, - route=route, - start_time=start_time, - ) - virtual_key_auth_obj.via_virtual_key = True - return virtual_key_auth_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, + prisma_client=prisma_client, + ) + + 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, + ) + asyncio.create_task( + proxy_logging_obj.budget_alerts( + type="proxy_budget", + user_info=call_info, + ) + ) + if valid_token.token is None: + raise HTTPException(401, detail="Invalid API key, no token associated") + + 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 + + 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, + route=route, + start_time=start_time, + ) + virtual_key_auth_obj.via_virtual_key = True + return virtual_key_auth_obj async def _safe_fetch(label: str, awaitable): @@ -4832,6 +4785,61 @@ async def _run_post_custom_auth_checks( return valid_token +async def _authorize_resolved_key( + identity: UserAPIKeyAuth, + key_hash: str, + request: Request, + request_data: dict[str, object], # mutable-ok: model access checks rewrite the request model in place + route: str, +) -> UserAPIKeyAuth: + from litellm.proxy.proxy_server import ( + master_key, + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) + + start_time: Final = datetime.now(timezone.utc) + effects: Final = _BuilderEffects( + request=request, + request_data=request_data, + token=key_hash, + preresolved_key=identity, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=None, + proxy_logging_obj=proxy_logging_obj, + route=route, + ) + outcome: Final = await drive( + _AUTH_CORE.start_resolved_key( + RequestFacts(route=route, token=key_hash, is_llm_api_route=RouteChecks.is_llm_api_route(route=route)), + _core_settings(master_key=master_key, allow_user_auth=False, prisma_client=prisma_client), + start_time, + ), + 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) + + async def authorize_internal_virtual_key( key_hash: str, request: Request, request_data: dict[str, object] ) -> UserAPIKeyAuth: @@ -4845,17 +4853,8 @@ async def authorize_internal_virtual_key( ) route: Final = get_request_route(request=request) await pre_db_read_auth_checks(request_data=request_data, request=request, route=route) - auth: Final = await validate_resolved_virtual_key( - request=request, - request_data=request_data, - valid_token=identity, - api_key=key_hash, - route=route, - start_time=datetime.now(timezone.utc), - parent_otel_span=None, - end_user_id=None, - end_user_params={}, - _end_user_object=None, + auth: Final = await _authorize_resolved_key( + identity=identity, key_hash=key_hash, request=request, request_data=request_data, route=route ) auth.budget_reservation = None recovered: Final = await _authorize_authenticated_request( diff --git a/tests/unit/_auth/test_python_core.py b/tests/unit/_auth/test_python_core.py index 753e8ef19ca..d2245199c98 100644 --- a/tests/unit/_auth/test_python_core.py +++ b/tests/unit/_auth/test_python_core.py @@ -1,4 +1,5 @@ -from dataclasses import dataclass +import math +from dataclasses import dataclass, replace from datetime import datetime, timedelta, timezone from typing import Final @@ -10,6 +11,9 @@ from litellm._auth.decision import ( CachedAdminKey, Deny, ExpiredKey, + KeyBlocked, + KeyBudgetExceeded, + KeyFacts, MalformedKey, MasterKey, MasterKeyOnlyRoute, @@ -19,9 +23,14 @@ from litellm._auth.decision import ( NonAdminKey, NotAVirtualKey, Outcome, + OwnerDeactivated, + OwnerFacts, PresentedKey, RequestFacts, Settings, + Step, + TeamMemberBudget, + TeamMemberBudgetExceeded, UserAuthDisabled, UserAuthRoute, VirtualKey, @@ -38,25 +47,95 @@ SETTINGS: Final = Settings( has_db=True, master_key_only_routes=frozenset({"/global/spend/reset"}), ) +PLAIN_KEY: Final = KeyFacts( + blocked=False, + expires=None, + max_budget=None, + team_member_spend=None, + user_id=None, + team_id=None, + budget_throttle_pct=None, +) +CHECKED: Final = ( + "prepare_key", + "inspect_key", + "enforce_model_access", + "check_budget_exemption", + "send_key_budget_alerts", +) + +Call = tuple[str, object] @dataclass class RecordingHandler: presented: PresentedKey | None = None - calls: tuple[tuple[str, object], ...] = () + key: KeyFacts = PLAIN_KEY + owner: OwnerFacts | None = None + exempt: bool = False + team_member_budget: TeamMemberBudget | None = None + key_spend: float = 0.0 + calls: tuple[Call, ...] = () + + def _record(self, name: str, payload: object = None) -> None: + self.calls = (*self.calls, (name, payload)) async def resolve_end_user(self, prefetch_key: bool) -> None: - self.calls = (*self.calls, ("resolve_end_user", prefetch_key)) + self._record("resolve_end_user", prefetch_key) async def resolve_presented_key(self, hashed_token: str) -> PresentedKey | None: - self.calls = (*self.calls, ("resolve_presented_key", hashed_token)) + self._record("resolve_presented_key", hashed_token) return self.presented async def evict_presented_key(self, hashed_token: str) -> None: - self.calls = (*self.calls, ("evict_presented_key", hashed_token)) + self._record("evict_presented_key", hashed_token) async def load_key(self, hashed_token: str) -> None: - self.calls = (*self.calls, ("load_key", hashed_token)) + self._record("load_key", hashed_token) + + async def prepare_key(self, loaded_from_db: bool) -> None: + self._record("prepare_key", loaded_from_db) + + async def inspect_key(self) -> KeyFacts: + self._record("inspect_key") + return self.key + + async def enforce_model_access(self) -> None: + self._record("enforce_model_access") + + async def load_key_owner(self, user_id: str) -> OwnerFacts | None: + self._record("load_key_owner", user_id) + return self.owner + + async def check_budget_exemption(self) -> bool: + self._record("check_budget_exemption") + return self.exempt + + async def read_team_member_budget(self, user_id: str, team_id: str) -> TeamMemberBudget | None: + self._record("read_team_member_budget", (user_id, team_id)) + return self.team_member_budget + + async def alert_team_member_budget(self, user_id: str, team_id: str, budget: TeamMemberBudget) -> None: + self._record("alert_team_member_budget", budget) + + async def send_key_budget_alerts(self) -> None: + self._record("send_key_budget_alerts") + + async def read_key_spend(self, max_budget: float) -> float: + self._record("read_key_spend", max_budget) + return self.key_spend + + async def alert_key_budget(self, spend: float) -> None: + self._record("alert_key_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) + + +async def _run(step: Step, handler: RecordingHandler) -> tuple[Outcome, tuple[str, ...]]: + outcome: Final = await drive(step, handler) + return outcome, tuple(name for name, _ in handler.calls) async def _decide( @@ -64,11 +143,34 @@ async def _decide( *, route: str = "/chat/completions", settings: Settings = SETTINGS, - presented: PresentedKey | None = None, -) -> tuple[Outcome, tuple[tuple[str, object], ...]]: - handler: Final = RecordingHandler(presented=presented) - outcome: Final = await drive(PythonCore().start(RequestFacts(route=route, token=token), settings, NOW), handler) - return outcome, handler.calls + handler: RecordingHandler | None = None, +) -> tuple[Outcome, tuple[str, ...]]: + return await _run(PythonCore().start(_facts(token, route), settings, NOW), handler or RecordingHandler()) + + +async def _check( + key: KeyFacts, + *, + is_llm_api_route: bool = True, + settings: Settings = SETTINGS, + owner: OwnerFacts | None = None, + exempt: bool = False, + team_member_budget: TeamMemberBudget | None = None, + key_spend: float = 0.0, +) -> tuple[Outcome, tuple[str, ...]]: + handler: Final = RecordingHandler( + key=key, owner=owner, exempt=exempt, team_member_budget=team_member_budget, key_spend=key_spend + ) + facts: Final = _facts("hash", is_llm_api_route=is_llm_api_route) + 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 + ) + ) @pytest.mark.asyncio @@ -92,71 +194,206 @@ async def test_credential_decisions_do_no_io(route: str, token: str | None, sett @pytest.mark.asyncio @pytest.mark.parametrize("route", ["/chat/completions", "/global/spend/reset"]) async def test_master_key_resolves_end_user_without_key_prefetch(route: str): - assert await _decide(MASTER, route=route) == (Allow(MasterKey()), (("resolve_end_user", False),)) + handler: Final = RecordingHandler() + assert await _decide(MASTER, route=route, handler=handler) == (Allow(MasterKey()), ("resolve_end_user",)) + assert handler.calls == (("resolve_end_user", False),) @pytest.mark.asyncio -async def test_unknown_virtual_key_is_loaded_from_the_db_after_a_cache_miss(): - hashed: Final = hash_token("sk-x") - assert await _decide("sk-x") == ( - Allow(VirtualKey(loaded_from_db=True)), - (("resolve_end_user", True), ("resolve_presented_key", hashed), ("load_key", hashed)), - ) +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 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) @pytest.mark.asyncio -async def test_cached_virtual_key_is_not_loaded_again(): - assert await _decide("sk-x", presented=NonAdminKey()) == ( - Allow(VirtualKey(loaded_from_db=False)), - (("resolve_end_user", True), ("resolve_presented_key", hash_token("sk-x"))), - ) +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 calls == ("resolve_end_user", "resolve_presented_key", *CHECKED) + assert handler.calls[2] == ("prepare_key", False) @pytest.mark.asyncio @pytest.mark.parametrize("expires", [None, NOW + timedelta(seconds=1), NOW]) -async def test_unexpired_cached_admin_key_skips_route_and_db_checks(expires: datetime | None): +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"})) - outcome, _ = await _decide( - "sk-admin", route="/global/spend/reset", settings=no_db, presented=AdminKey(expires=expires) + outcome, calls = await _decide( + "sk-admin", + route="/global/spend/reset", + settings=no_db, + handler=RecordingHandler(presented=AdminKey(expires=expires)), ) - assert outcome == Allow(CachedAdminKey()) + assert (outcome, calls) == (Allow(CachedAdminKey()), ("resolve_end_user", "resolve_presented_key")) @pytest.mark.asyncio -async def test_expired_cached_admin_key_is_evicted_then_denied(): +async def test_expired_cached_admin_key_is_evicted_then_denied_with_the_presented_key(): expires: Final = NOW - timedelta(seconds=1) - hashed: Final = hash_token("sk-admin") - assert await _decide("sk-admin", presented=AdminKey(expires=expires)) == ( - Deny(ExpiredKey(expires=expires, now=NOW)), - (("resolve_end_user", True), ("resolve_presented_key", hashed), ("evict_presented_key", hashed)), - ) + handler: Final = RecordingHandler(presented=AdminKey(expires=expires)) + outcome, calls = await _decide("sk-admin", handler=handler) + assert outcome == Deny(ExpiredKey(expires=expires, now=NOW, key_label="sk-admin")) + assert calls == ("resolve_end_user", "resolve_presented_key", "evict_presented_key") + assert handler.calls[-1] == ("evict_presented_key", hash_token("sk-admin")) @pytest.mark.asyncio @pytest.mark.parametrize("presented", [None, NonAdminKey()]) async def test_non_admin_key_cannot_use_master_key_only_routes(presented: PresentedKey | None): - outcome, calls = await _decide("sk-x", route="/global/spend/reset", presented=presented) - assert outcome == Deny(MasterKeyOnlyRoute()) - assert "load_key" not in {name for name, _ in calls} + outcome, calls = await _decide("sk-x", route="/global/spend/reset", handler=RecordingHandler(presented=presented)) + assert (outcome, calls) == (Deny(MasterKeyOnlyRoute()), ("resolve_end_user", "resolve_presented_key")) @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()) - outcome, _ = await _decide("sk-x", settings=no_db, presented=presented) + outcome, _ = await _decide("sk-x", settings=no_db, handler=RecordingHandler(presented=presented)) assert outcome == Deny(NoDatabase()) @pytest.mark.asyncio @pytest.mark.parametrize("token", ["eyJhbGciOiJIUzI1NiJ9.e30.sig", "a" * 64, "sk"]) async def test_uncached_token_without_virtual_key_prefix_is_rejected_before_the_db(token: str): - outcome, calls = await _decide(token) - assert outcome == Deny(NotAVirtualKey()) - assert [name for name, _ in calls] == ["resolve_end_user", "resolve_presented_key"] + assert await _decide(token) == (Deny(NotAVirtualKey()), ("resolve_end_user", "resolve_presented_key")) @pytest.mark.asyncio -async def test_cached_token_without_virtual_key_prefix_is_accepted(): - outcome, _ = await _decide("ui-login-blob", presented=NonAdminKey()) - assert outcome == Allow(VirtualKey(loaded_from_db=False)) +async def test_expired_virtual_key_is_labelled_by_the_hash_it_was_loaded_with(): + expires: Final = NOW - timedelta(seconds=1) + outcome, _ = await _decide("sk-x", handler=RecordingHandler(key=replace(PLAIN_KEY, expires=expires))) + assert outcome == Deny(ExpiredKey(expires=expires, now=NOW, key_label=hash_token("sk-x"))) + + +@pytest.mark.asyncio +async def test_blocked_key_is_denied_before_model_access(): + assert await _check(replace(PLAIN_KEY, blocked=True)) == (Deny(KeyBlocked()), ("inspect_key",)) + + +@pytest.mark.asyncio +async def test_deactivated_owner_is_denied_after_model_access(): + outcome, calls = await _check(replace(PLAIN_KEY, user_id="u1"), owner=OwnerFacts(deactivated=True)) + assert (outcome, calls) == ( + Deny(OwnerDeactivated(user_id="u1")), + ("inspect_key", "enforce_model_access", "load_key_owner"), + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("owner", [None, OwnerFacts(deactivated=False)]) +async def test_active_or_missing_owner_continues(owner: OwnerFacts | None): + outcome, _ = await _check(replace(PLAIN_KEY, user_id="u1"), owner=owner) + assert outcome == _allowed() + + +@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) + assert await _check(key, exempt=True) == ( + _allowed(budget_checks_skipped=True), + ("inspect_key", "enforce_model_access", "load_key_owner", "check_budget_exemption"), + ) + expired: Final = replace(key, expires=NOW - timedelta(seconds=1)) + outcome, _ = await _check(expired, exempt=True) + assert outcome == Deny(ExpiredKey(expires=NOW - timedelta(seconds=1), now=NOW, key_label="hash")) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("budget", "denied"), + [ + (TeamMemberBudget(spend=10.0, max_budget=10.0), True), + (TeamMemberBudget(spend=10.1, max_budget=10.0), True), + (TeamMemberBudget(spend=9.9, max_budget=10.0), False), + (None, False), + ], +) +async def test_team_member_budget_alerts_then_denies_at_the_limit(budget: TeamMemberBudget | None, denied: bool): + key: Final = replace(PLAIN_KEY, user_id="u1", team_id="t1", team_member_spend=0.0) + outcome, calls = await _check(key, team_member_budget=budget) + if denied: + assert budget is not None + assert outcome == Deny( + TeamMemberBudgetExceeded(user_id="u1", team_id="t1", spend=budget.spend, max_budget=budget.max_budget) + ) + assert calls[-2:] == ("read_team_member_budget", "alert_team_member_budget") + else: + assert outcome == _allowed() + assert "alert_team_member_budget" not in calls + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "key", + [ + replace(PLAIN_KEY, user_id="u1", team_id="t1", team_member_spend=None), + replace(PLAIN_KEY, user_id=None, team_id="t1", team_member_spend=0.0), + replace(PLAIN_KEY, user_id="u1", team_id=None, team_member_spend=0.0), + ], +) +async def test_team_member_budget_is_read_only_for_a_member_with_tracked_spend(key: KeyFacts): + _, calls = await _check(key, team_member_budget=TeamMemberBudget(spend=1.0, max_budget=0.5)) + assert "read_team_member_budget" not in calls + + +@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()) + _, 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 + + +@pytest.mark.asyncio +async def test_team_member_budget_is_checked_before_expiry(): + key: Final = replace( + PLAIN_KEY, user_id="u1", team_id="t1", team_member_spend=0.0, expires=NOW - timedelta(seconds=1) + ) + outcome, _ = await _check(key, team_member_budget=TeamMemberBudget(spend=1.0, max_budget=0.5)) + assert outcome == Deny(TeamMemberBudgetExceeded(user_id="u1", team_id="t1", spend=1.0, max_budget=0.5)) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("spend", "max_budget", "throttle", "expected"), + [ + (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))), + ( + 1e9, + math.inf, + None, + Allow(VirtualKey(loaded_from_db=True, budget_checks_skipped=False, budget_throttle_pct=None)), + ), + ( + 1e9, + math.nan, + None, + Allow(VirtualKey(loaded_from_db=True, budget_checks_skipped=False, budget_throttle_pct=None)), + ), + ], +) +async def test_key_budget_denies_or_throttles_at_the_limit( + spend: float, max_budget: float, throttle: float | None, expected: Outcome +): + key: Final = replace(PLAIN_KEY, max_budget=max_budget, budget_throttle_pct=throttle) + 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) + + +@pytest.mark.asyncio +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" 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 cea01a7cf69..d7e58974b04 100644 --- a/tests/unit/proxy/auth/test_user_api_key_auth.py +++ b/tests/unit/proxy/auth/test_user_api_key_auth.py @@ -1600,7 +1600,9 @@ 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.validate_resolved_virtual_key)) + + textwrap.dedent(inspect.getsource(auth_module._finish_virtual_key_auth)) + + "\n" + + textwrap.dedent(inspect.getsource(auth_module._BuilderEffects)) ) def calls_before_each_return(node): @@ -1645,7 +1647,9 @@ 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.validate_resolved_virtual_key)) + + textwrap.dedent(inspect.getsource(auth_module._finish_virtual_key_auth)) + + "\n" + + textwrap.dedent(inspect.getsource(auth_module._BuilderEffects)) ) assignments = [ @@ -1663,7 +1667,7 @@ def test_every_jwt_branch_carries_the_user_model_budget(): assert "auto_registered" in targets, ( f"the auto-registered JWT token must carry the user's model budget; only these are populated: {sorted(targets)}" ) - assert "valid_token" in targets, "the virtual-key path must carry the user's model budget" + assert "key" in targets, "the virtual-key path must carry the user's model budget" @pytest.mark.asyncio @@ -1782,7 +1786,9 @@ def test_zero_cost_models_skip_the_user_budget_check_on_every_path(): tree = ast.parse( textwrap.dedent(inspect.getsource(auth_module._user_api_key_auth_builder)) + "\n" - + textwrap.dedent(inspect.getsource(auth_module.validate_resolved_virtual_key)) + + textwrap.dedent(inspect.getsource(auth_module._finish_virtual_key_auth)) + + "\n" + + textwrap.dedent(inspect.getsource(auth_module._BuilderEffects)) ) def guarded_by_skip(node: ast.AST, target: ast.AST) -> bool: @@ -1803,7 +1809,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 == "_check_user_model_budget" + 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"} ] assert len(calls) == 2, f"expected the JWT and virtual-key call sites, found {len(calls)}" @@ -1927,7 +1935,9 @@ 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.validate_resolved_virtual_key)) + + textwrap.dedent(inspect.getsource(auth_module._finish_virtual_key_auth)) + + "\n" + + textwrap.dedent(inspect.getsource(auth_module._BuilderEffects)) ) # Half one: the shared block copies the user row's budget onto the token.