scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-13 08:52:02 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 174

refactor(auth): finish the virtual-key decision chain in AuthCore

Browse source 
The core now decides the soft budget alert, when per-model budgets run, the
team and project context load, the proxy budget alert and the missing-token
denial, so the builder only projects the grant. One projection serves both
the request builder and internal jobs.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Yujong Lee 2026-11-06 09:52:33 +09:00
parent 08fbc4ac4c
commit 87a1064e90
9 changed files with 596 additions and 366 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
litellm/_auth/AGENTS.md
View file

@ -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

50
litellm/_auth/decision.py
View file

@ -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

56
litellm/_auth/driver.py
View file

@ -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)

82
litellm/_auth/python_core.py
View file

@ -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))

55
litellm/proxy/auth/auth_checks.py
View file

@ -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

5
litellm/proxy/auth/auth_denial.py
View file

@ -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)

539
litellm/proxy/auth/user_api_key_auth.py
View file

@ -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(

143
tests/unit/_auth/test_python_core.py
View file

@ -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")

26
tests/unit/proxy/auth/test_user_api_key_auth.py
View file

@ -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))
)