refactor(auth): resolve virtual keys through AuthCore effects

The Python core now decides the whole virtual key resolution path. It asks
the proxy to resolve the end user, the presented key (preresolved, cached or
UI login), evict an expired cached admin key and load an uncached key from
the DB, each as a typed effect with its own continuation. It decides the
admin key expiry, master-key-only routes, the no-DB case and the sk- prefix
check itself, so the builder no longer has a legacy fallback.

The budget and team checks in validate_resolved_virtual_key still run after
Allow(VirtualKey) and move into the core in later passes

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Yujong Lee 2026-11-06 07:58:45 +09:00
parent 4b5b3abfce
commit 8683d4a77e
7 changed files with 647 additions and 290 deletions

View file

@ -8,8 +8,8 @@
## The swap point
- `decision.py` is the interface. `AuthCore.decide(RequestFacts, Settings) -> Outcome` 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 and settings always give the same outcome
- `Defer` means the core does not handle this request yet, and the legacy builder in `litellm/proxy/auth/user_api_key_auth.py` takes it from there. Move one decision at a time out of the builder into the core, and delete the builder code it replaces in the same change
- Projection stays in the proxy. `Allow` grants become `UserAPIKeyAuth` in the builder, and `litellm/proxy/auth/auth_denial.py` maps every `DenyReason` to the existing error contract
- When a decision needs a lookup, add it as an effect the core yields and the proxy answers, so both cores share the same IO code
- `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 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

View file

@ -1,5 +1,6 @@
from collections.abc import Callable
from dataclasses import dataclass
from enum import Enum
from datetime import datetime
from typing import Protocol, TypeAlias
@ -13,6 +14,8 @@ class RequestFacts:
class Settings:
master_key: str | None
allow_user_auth: bool
has_db: bool
master_key_only_routes: frozenset[str]
@dataclass(frozen=True, slots=True)
@ -30,13 +33,58 @@ class MasterKey:
pass
Grant: TypeAlias = NoMasterKey | UserAuthRoute | MasterKey
@dataclass(frozen=True, slots=True)
class CachedAdminKey:
pass
class DenyReason(Enum):
MISSING_KEY = "missing_key"
MALFORMED_KEY = "malformed_key"
USER_AUTH_DISABLED = "user_auth_disabled"
@dataclass(frozen=True, slots=True)
class VirtualKey:
loaded_from_db: bool
Grant: TypeAlias = NoMasterKey | UserAuthRoute | MasterKey | CachedAdminKey | VirtualKey
@dataclass(frozen=True, slots=True)
class MissingKey:
pass
@dataclass(frozen=True, slots=True)
class MalformedKey:
pass
@dataclass(frozen=True, slots=True)
class UserAuthDisabled:
pass
@dataclass(frozen=True, slots=True)
class ExpiredKey:
expires: datetime
now: datetime
@dataclass(frozen=True, slots=True)
class MasterKeyOnlyRoute:
pass
@dataclass(frozen=True, slots=True)
class NoDatabase:
pass
@dataclass(frozen=True, slots=True)
class NotAVirtualKey:
pass
Denial: TypeAlias = (
MissingKey | MalformedKey | UserAuthDisabled | ExpiredKey | MasterKeyOnlyRoute | NoDatabase | NotAVirtualKey
)
@dataclass(frozen=True, slots=True)
@ -46,16 +94,52 @@ class Allow:
@dataclass(frozen=True, slots=True)
class Deny:
reason: DenyReason
denial: Denial
Outcome: TypeAlias = Allow | Deny
@dataclass(frozen=True, slots=True)
class Defer:
class AdminKey:
expires: datetime | None
@dataclass(frozen=True, slots=True)
class NonAdminKey:
pass
Outcome: TypeAlias = Allow | Deny | Defer
PresentedKey: TypeAlias = AdminKey | NonAdminKey
@dataclass(frozen=True, slots=True)
class ResolveEndUser:
prefetch_key: bool
resume: "Callable[[], Step]"
@dataclass(frozen=True, slots=True)
class ResolvePresentedKey:
hashed_token: str
resume: "Callable[[PresentedKey | None], Step]"
@dataclass(frozen=True, slots=True)
class EvictPresentedKey:
hashed_token: str
resume: "Callable[[], Step]"
@dataclass(frozen=True, slots=True)
class LoadKey:
hashed_token: str
resume: "Callable[[], Step]"
Effect: TypeAlias = ResolveEndUser | ResolvePresentedKey | EvictPresentedKey | LoadKey
Step: TypeAlias = Outcome | Effect
class AuthCore(Protocol):
def decide(self, facts: RequestFacts, settings: Settings) -> Outcome: ...
def start(self, facts: RequestFacts, settings: Settings, now: datetime) -> Step: ...

42
litellm/_auth/driver.py Normal file
View file

@ -0,0 +1,42 @@
from typing import Protocol, assert_never
from litellm._auth.decision import (
Allow,
Deny,
EvictPresentedKey,
LoadKey,
Outcome,
PresentedKey,
ResolveEndUser,
ResolvePresentedKey,
Step,
)
class EffectHandler(Protocol):
async def resolve_end_user(self, prefetch_key: bool) -> None: ...
async def resolve_presented_key(self, hashed_token: str) -> PresentedKey | None: ...
async def evict_presented_key(self, hashed_token: str) -> None: ...
async def load_key(self, hashed_token: str) -> None: ...
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)

View file

@ -1,32 +1,84 @@
from datetime import datetime
from typing import Final
from litellm._auth.decision import (
AdminKey,
Allow,
Defer,
CachedAdminKey,
Deny,
DenyReason,
EvictPresentedKey,
ExpiredKey,
LoadKey,
MalformedKey,
MasterKey,
MasterKeyOnlyRoute,
MissingKey,
NoDatabase,
NoMasterKey,
Outcome,
NonAdminKey,
NotAVirtualKey,
PresentedKey,
RequestFacts,
ResolveEndUser,
ResolvePresentedKey,
Settings,
Step,
UserAuthDisabled,
UserAuthRoute,
VirtualKey,
)
from litellm._auth.key import is_master_key
from litellm._auth.key import hash_token, is_master_key
_USER_AUTH_ROUTE: Final = "/user/auth"
_VIRTUAL_KEY_PREFIX: Final = "sk-"
class PythonCore:
def decide(self, facts: RequestFacts, settings: Settings) -> Outcome:
def start(self, facts: RequestFacts, settings: Settings, now: datetime) -> Step:
token: Final = facts.token
if settings.master_key is None:
return Allow(NoMasterKey(token=facts.token))
if facts.token is None:
return Deny(DenyReason.MISSING_KEY)
if facts.token == "":
return Deny(DenyReason.MALFORMED_KEY)
return Allow(NoMasterKey(token=token))
if token is None:
return Deny(MissingKey())
if token == "":
return Deny(MalformedKey())
if facts.route == _USER_AUTH_ROUTE:
return Allow(UserAuthRoute()) if settings.allow_user_auth else Deny(DenyReason.USER_AUTH_DISABLED)
if is_master_key(facts.token, settings.master_key):
return Allow(MasterKey())
return Defer()
return Allow(UserAuthRoute()) if settings.allow_user_auth else Deny(UserAuthDisabled())
if is_master_key(token, settings.master_key):
return ResolveEndUser(prefetch_key=False, resume=lambda: Allow(MasterKey()))
hashed_token: Final = hash_token(token)
return ResolveEndUser(
prefetch_key=True,
resume=lambda: ResolvePresentedKey(
hashed_token=hashed_token,
resume=lambda presented: _after_presented_key(presented, token, hashed_token, facts, settings, now),
),
)
def _after_presented_key(
presented: PresentedKey | None,
token: str,
hashed_token: str,
facts: RequestFacts,
settings: Settings,
now: datetime,
) -> Step:
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))
)
case AdminKey():
return Allow(CachedAdminKey())
case NonAdminKey() | None:
pass
if facts.route in settings.master_key_only_routes:
return Deny(MasterKeyOnlyRoute())
if not settings.has_db:
return Deny(NoDatabase())
if presented is not None:
return Allow(VirtualKey(loaded_from_db=False))
if not token.startswith(_VIRTUAL_KEY_PREFIX):
return Deny(NotAVirtualKey())
return LoadKey(hashed_token=hashed_token, resume=lambda: Allow(VirtualKey(loaded_from_db=True)))

View file

@ -1,31 +1,86 @@
from typing import assert_never
from typing import Final, assert_never
from fastapi import HTTPException, status
from litellm._auth.decision import DenyReason
from litellm._auth.decision import (
Denial,
ExpiredKey,
MalformedKey,
MasterKeyOnlyRoute,
MissingKey,
NoDatabase,
NotAVirtualKey,
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.auth.auth_utils import abbreviate_api_key
from litellm.types.proxy.auth.user_api_key_auth import is_jwt
_JWT_AUTH_DISABLED_HINT: Final = (
" This key has the structure of a JWT, but JWT auth is not enabled on this proxy, so it was treated as a"
" virtual key. Set `enable_jwt_auth: true` under `general_settings` in your proxy config to authenticate"
" with JWTs."
)
def denial_exception(reason: DenyReason) -> ProxyException | HTTPException:
match reason:
case DenyReason.MISSING_KEY:
def _not_a_virtual_key(token: str, jwt_auth_enabled: bool) -> HTTPException:
masked: Final = f"{token[:4]}****{token[-4:]}" if len(token) > 8 else "****"
hint: Final = _JWT_AUTH_DISABLED_HINT if not jwt_auth_enabled and is_jwt(token=token) else ""
error: Final = HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail=f"{INVALID_VIRTUAL_KEY_ERROR_MESSAGE}. Received={masked}, expected to start with 'sk-'.{hint}",
)
# Log routing classifies this 401 by where it was raised, never by its message text.
setattr(error, INVALID_VIRTUAL_KEY_ERROR_MARKER, True)
return error
def denial_exception(
denial: Denial, token: str | None, route: str, jwt_auth_enabled: bool
) -> ProxyException | HTTPException:
match denial:
case MissingKey():
return ProxyException(
message="Authentication Error, No api key passed in.",
type=ProxyErrorTypes.auth_error,
param=None,
code=status.HTTP_401_UNAUTHORIZED,
)
case DenyReason.MALFORMED_KEY:
case MalformedKey():
return ProxyException(
message="Authentication Error, Malformed API Key passed in. Ensure Key has `Bearer ` prefix.",
type=ProxyErrorTypes.auth_error,
param=None,
code=status.HTTP_401_UNAUTHORIZED,
)
case DenyReason.USER_AUTH_DISABLED:
case UserAuthDisabled():
return HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
detail="'allow_user_auth' not set or set to False",
)
case ExpiredKey(expires=expires, now=now):
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 ""),
)
case MasterKeyOnlyRoute():
return ProxyException(
message=f"Authentication Error, Tried to access route={route}, which is only for MASTER KEY",
type=ProxyErrorTypes.auth_error,
param=None,
code=status.HTTP_401_UNAUTHORIZED,
)
case NoDatabase():
return ProxyException(
message="No connected db.",
type=ProxyErrorTypes.no_db_connection,
code=400,
param=None,
)
case NotAVirtualKey():
return _not_a_virtual_key(token or "", jwt_auth_enabled)
case _:
assert_never(reason)
assert_never(denial)

View file

@ -31,16 +31,21 @@ from litellm._auth.credential import (
select_websocket_api_key,
)
from litellm._auth.decision import (
AdminKey,
Allow,
AuthCore,
Defer,
CachedAdminKey,
Deny,
MasterKey,
NoMasterKey,
NonAdminKey,
PresentedKey,
RequestFacts,
Settings,
UserAuthRoute,
VirtualKey,
)
from litellm._auth.driver import drive
from litellm._auth.python_core import PythonCore
from litellm._internal_context import service_target
from litellm._logging import verbose_logger, verbose_proxy_logger
@ -49,8 +54,6 @@ from litellm.caching.redis_cache import RedisCache
from litellm.constants import (
CLIENT_REQUESTED_MODEL_SCOPE_KEY,
GLOBAL_PROXY_SPEND_CACHE_KEY,
INVALID_VIRTUAL_KEY_ERROR_MARKER,
INVALID_VIRTUAL_KEY_ERROR_MESSAGE,
LITELLM_PROXY_BUDGET_NAME,
LITELLM_PROXY_MASTER_KEY_ALIAS,
MODEL_GROUP_ALIAS_RESOLVED_SCOPE_KEY,
@ -1555,12 +1558,6 @@ async def check_api_key_for_custom_headers_or_pass_through_endpoints(
# the lookup and return None (caller proceeds to auth_builder).
_JWT_PROXY_ADMIN_SENTINEL: Final = "__JWT_PROXY_ADMIN__"
_JWT_AUTH_DISABLED_HINT = (
" This key has the structure of a JWT, but JWT auth is not enabled on this proxy, so it was treated as a"
" virtual key. Set `enable_jwt_auth: true` under `general_settings` in your proxy config to authenticate"
" with JWTs."
)
class _JWTKeyMappingFailure(Enum):
NO_ELIGIBLE_KEY = "JWT Key Mapping: No eligible existing virtual key for this identity. Access denied."
@ -2196,7 +2193,7 @@ class _ResolvedEndUser(NamedTuple):
async def _resolve_end_user(
request: Request,
request_data: dict[str, object],
request_data: Mapping[str, object],
token_end_user_id: str | None,
api_key: str,
key_is_resolved: bool,
@ -2291,6 +2288,182 @@ async def _master_key_auth(
return with_end_user
_MASTER_KEY_ONLY_ROUTES: Final = frozenset(LiteLLMRoutes.master_key_only_routes.value)
def _aware_expiry(expires: datetime | str) -> datetime:
expiry: Final = expires if isinstance(expires, datetime) else datetime.fromisoformat(expires)
return (
expiry if expiry.tzinfo is not None and expiry.utcoffset() is not None else expiry.replace(tzinfo=timezone.utc)
)
def _presented_key(key: UserAPIKeyAuth) -> PresentedKey:
if key.user_role != LitellmUserRoles.PROXY_ADMIN:
return NonAdminKey()
return AdminKey(expires=None if key.expires is None else _aware_expiry(key.expires))
def _with_end_user(key: UserAPIKeyAuth, end_user: _ResolvedEndUser, parent_otel_span: Span | None) -> UserAPIKeyAuth:
with_end_user: Final = update_valid_token_with_end_user_params(valid_token=key, end_user_params=end_user.params)
with_end_user.parent_otel_span = parent_otel_span
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 _reset_end_user_fields(key: UserAPIKeyAuth, end_user: _ResolvedEndUser) -> UserAPIKeyAuth:
key.end_user_id = end_user.end_user_id or None
key.end_user_tpm_limit = end_user.budget.tpm_limit if end_user.budget is not None else None
key.end_user_rpm_limit = end_user.budget.rpm_limit if end_user.budget is not None else None
key.end_user_tpd_limit = end_user.budget.tpd_limit if end_user.budget is not None else None
key.allowed_model_region = (
end_user.end_user_object.allowed_model_region if end_user.end_user_object is not None else None
)
return key
async def _refresh_team_fields_from_cache(
valid_token: UserAPIKeyAuth,
prisma_client: PrismaClient | None,
user_api_key_cache: UserApiKeyCache,
parent_otel_span: Span | None,
proxy_logging_obj: ProxyLogging,
) -> None:
if valid_token.team_id is None or valid_token.team_id == UI_TEAM_ID:
return
## UPDATE TEAM VALUES BASED ON CACHED TEAM OBJECT - allows `/team/update` values to work for cached token
try:
team_obj: Final[LiteLLM_TeamTableCachedObj] = 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,
check_cache_only=True,
)
if (
team_obj.last_refreshed_at is not None
and valid_token.last_refreshed_at is not None
and team_obj.last_refreshed_at > valid_token.last_refreshed_at
):
team_obj_dict: Final[Mapping[str, object]] = MappingProxyType(team_obj.__dict__)
for k, v in team_obj_dict.items():
field_name = f"team_{k}"
if field_name in valid_token.__fields__:
setattr(valid_token, field_name, v)
except Exception as e:
verbose_logger.debug(e) # moving from .warning to .debug as it spams logs when team missing from cache.
class _BuilderEffects:
def __init__(
self,
*,
request: Request,
request_data: Mapping[str, object],
token: str,
preresolved_key: UserAPIKeyAuth | None,
prisma_client: PrismaClient | None,
user_api_key_cache: UserApiKeyCache,
parent_otel_span: Span | None,
proxy_logging_obj: ProxyLogging,
route: str,
) -> None:
self._request: Final = request
self._request_data: Final = request_data
self._token: Final = token
self._preresolved_key: Final = preresolved_key
self._prisma_client: Final = prisma_client
self._user_api_key_cache: Final = user_api_key_cache
self._parent_otel_span: Final = parent_otel_span
self._proxy_logging_obj: Final = proxy_logging_obj
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
@property
def key(self) -> UserAPIKeyAuth:
if self.resolved_key is None:
raise RuntimeError("the auth core granted a key it never resolved")
return self.resolved_key
def _identity_store(self, *, check_cache_only: bool) -> IdentityStore:
return IdentityStore(
self._prisma_client,
self._user_api_key_cache,
parent_otel_span=self._parent_otel_span,
proxy_logging_obj=self._proxy_logging_obj,
check_cache_only=check_cache_only,
)
async def resolve_end_user(self, prefetch_key: bool) -> None:
self.end_user = await _resolve_end_user(
request=self._request,
request_data=self._request_data,
token_end_user_id=self._preresolved_key.end_user_id if self._preresolved_key is not None else None,
api_key=self._token,
key_is_resolved=not prefetch_key or self._preresolved_key is not None,
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,
route=self._route,
)
async def _cached_key(self, hashed_token: str) -> UserAPIKeyAuth | None:
try:
with tracer.trace("litellm.proxy.auth.get_key_object_check_cache"):
cached: Final = IdentityStore.key_from_principal(
await self._identity_store(check_cache_only=True).resolve(hashed_token=hashed_token)
)
except Exception:
verbose_logger.debug("api key not found in cache.")
return None
# Key-cache entries are written only after the proxy validated a virtual key or the master key, but
# via_virtual_key is exclude=True so serialization drops it. Restore it at this trusted boundary.
cached.via_virtual_key = True
return cached
def _ui_login_key(self) -> UserAPIKeyAuth | None:
# Decryption fails closed for anything that is not a genuine UI login blob.
if self._token.startswith("sk-") or get_secret_bool("EXPERIMENTAL_UI_LOGIN") is False:
return None
return ExperimentalUIJWTToken.get_key_object_from_ui_hash_key(self._token)
async def resolve_presented_key(self, hashed_token: str) -> PresentedKey | None:
found: Final = self._preresolved_key or await self._cached_key(hashed_token) or self._ui_login_key()
if found is None:
return None
self.resolved_key = (
await _refresh_session_token_grants(
valid_token=found,
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,
)
if found.is_session_token and self._prisma_client is not None
else found
)
return _presented_key(self.resolved_key)
async def evict_presented_key(self, hashed_token: str) -> None:
await delete_cache_key_object(
hashed_token=hashed_token,
user_api_key_cache=self._user_api_key_cache,
proxy_logging_obj=self._proxy_logging_obj,
)
async def load_key(self, hashed_token: str) -> None:
with tracer.trace("litellm.proxy.auth.get_key_object_from_db"):
self.resolved_key = IdentityStore.key_from_principal(
await self._identity_store(check_cache_only=False).resolve(hashed_token=hashed_token)
)
async def user_api_key_auth_builder(
request: Request,
api_key: str,
@ -2701,10 +2874,34 @@ async def user_api_key_auth_builder(
status_code=500,
detail={f"Master key must be a valid string. Current type={type(master_key)}"},
)
match _AUTH_CORE.decide(
RequestFacts(route=route, token=api_key),
Settings(master_key=master_key, allow_user_auth=general_settings.get("allow_user_auth", False) is True),
):
effects: Final = _BuilderEffects(
request=request,
request_data=request_data,
token=api_key,
preresolved_key=valid_token,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=parent_otel_span,
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,
),
datetime.now(timezone.utc),
),
effects,
)
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)
case Allow(grant=grant):
match grant:
case NoMasterKey(token=token):
@ -2718,18 +2915,7 @@ async def user_api_key_auth_builder(
case MasterKey():
return await _master_key_auth(
token=api_key,
end_user=await _resolve_end_user(
request=request,
request_data=request_data,
token_end_user_id=valid_token.end_user_id if valid_token is not None else None,
api_key=api_key,
key_is_resolved=True,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
route=route,
),
end_user=effects.end_user,
admin_user_id=litellm_proxy_admin_name,
parent_otel_span=parent_otel_span,
route=route,
@ -2737,212 +2923,31 @@ async def user_api_key_auth_builder(
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
case Deny(reason=reason):
raise denial_exception(reason)
case Defer():
pass
end_user: Final = await _resolve_end_user(
request=request,
request_data=request_data,
token_end_user_id=valid_token.end_user_id if valid_token is not None else None,
api_key=api_key,
key_is_resolved=valid_token is not None,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
route=route,
)
if valid_token is None:
## Check CACHE
try:
with tracer.trace("litellm.proxy.auth.get_key_object_check_cache"):
valid_token = IdentityStore.key_from_principal(
await IdentityStore(
prisma_client,
user_api_key_cache,
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(
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,
route=route,
start_time=start_time,
parent_otel_span=parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
check_cache_only=True,
).resolve(hashed_token=hash_token(api_key))
)
# Key-cache entries are written only after the proxy validated a
# virtual key or the master key, but via_virtual_key is exclude=True
# so serialization drops it; restore it at this trusted boundary.
# The UI-login JWT fallback below constructs its token from a
# decrypted blob, not this cache, and stays unmarked.
if isinstance(valid_token, UserAPIKeyAuth):
valid_token.via_virtual_key = True
except Exception:
verbose_logger.debug("api key not found in cache.")
valid_token = None
## Check UI/CLI Hash Key
# Attempt decryption for non-sk- tokens unless the operator has
# explicitly set EXPERIMENTAL_UI_LOGIN=false to disable it.
# Unset (None) keeps the new default of always attempting decryption;
# decryption fails closed for anything that is not a genuine blob.
if (
valid_token is None
and not api_key.startswith("sk-")
and get_secret_bool("EXPERIMENTAL_UI_LOGIN") is not False
):
valid_token = ExperimentalUIJWTToken.get_key_object_from_ui_hash_key(api_key)
if valid_token is not None and valid_token.is_session_token and prisma_client is not None:
valid_token = await _refresh_session_token_grants( # rebind-ok: later checks read this name
valid_token=valid_token,
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 (
valid_token is not None
and isinstance(valid_token, UserAPIKeyAuth)
and valid_token.user_role == LitellmUserRoles.PROXY_ADMIN
):
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)
if expiry_time < current_time:
await delete_cache_key_object(
hashed_token=hash_token(api_key),
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
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),
)
valid_token = update_valid_token_with_end_user_params(
valid_token=valid_token, end_user_params=end_user.params
)
valid_token.parent_otel_span = parent_otel_span
if end_user.end_user_object is not None:
valid_token.end_user_object_permission = end_user.end_user_object.object_permission
return valid_token
if (
valid_token is not None
and isinstance(valid_token, UserAPIKeyAuth)
and valid_token.team_id is not None
and valid_token.team_id != UI_TEAM_ID
):
## UPDATE TEAM VALUES BASED ON CACHED TEAM OBJECT - allows `/team/update` values to work for cached token
try:
team_obj: Final[LiteLLM_TeamTableCachedObj] = 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,
check_cache_only=True,
)
if (
team_obj.last_refreshed_at is not None
and valid_token.last_refreshed_at is not None
and team_obj.last_refreshed_at > valid_token.last_refreshed_at
):
team_obj_dict: Final[Mapping[str, object]] = MappingProxyType(team_obj.__dict__)
for k, v in team_obj_dict.items():
field_name = f"team_{k}"
if field_name in valid_token.__fields__:
setattr(valid_token, field_name, v)
except Exception as e:
verbose_logger.debug(e) # moving from .warning to .debug as it spams logs when team missing from cache.
## IF it's not a master key
## Route should not be in master_key_only_routes
if route in LiteLLMRoutes.master_key_only_routes.value:
raise ProxyException(
message=f"Authentication Error, Tried to access route={route}, which is only for MASTER KEY",
type=ProxyErrorTypes.auth_error,
param=None,
code=status.HTTP_401_UNAUTHORIZED,
)
## Check DB
if (
prisma_client is None
): # if both master key + user key submitted, and user key != master key, and no db connected, raise an error
raise ProxyException(
message="No connected db.",
type=ProxyErrorTypes.no_db_connection,
code=400,
param=None,
)
if valid_token is None:
if isinstance(api_key, str): # if generated token, make sure it starts with sk-.
_masked_key: Final = f"{api_key[:4]}****{api_key[-4:]}" if len(api_key) > 8 else "****"
if not api_key.startswith("sk-"):
_hint = _JWT_AUTH_DISABLED_HINT if not enable_jwt_auth and JWTHandler.is_jwt(token=api_key) else ""
_malformed_key_error = HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail=(
f"{INVALID_VIRTUAL_KEY_ERROR_MESSAGE}. Received={_masked_key}, "
f"expected to start with 'sk-'.{_hint}"
),
) # prevent token hashes from being used
# Stamp provenance here so log routing classifies this 401 by
# where it was raised, never by its message text.
setattr(_malformed_key_error, INVALID_VIRTUAL_KEY_ERROR_MARKER, True)
raise _malformed_key_error
else:
verbose_logger.warning(
"litellm.proxy.proxy_server.user_api_key_auth(): Warning - Key is not a string. Got type={}".format(
type(api_key) if api_key is not None else "None"
)
)
if api_key.startswith("sk-"):
api_key = hash_token(token=api_key)
with tracer.trace("litellm.proxy.auth.get_key_object_from_db"):
valid_token = IdentityStore.key_from_principal(
await IdentityStore(
prisma_client,
user_api_key_cache,
parent_otel_span=parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
).resolve(hashed_token=api_key)
)
# update end-user params on valid token
# These can change per request - it's important to update them here
valid_token.end_user_id = end_user.end_user_id or None
valid_token.end_user_tpm_limit = end_user.budget.tpm_limit if end_user.budget is not None else None
valid_token.end_user_rpm_limit = end_user.budget.rpm_limit if end_user.budget is not None else None
valid_token.end_user_tpd_limit = end_user.budget.tpd_limit if end_user.budget is not None else None
valid_token.allowed_model_region = (
end_user.end_user_object.allowed_model_region if end_user.end_user_object is not None else None
)
return await validate_resolved_virtual_key(
request=request,
request_data=request_data,
valid_token=valid_token,
api_key=api_key,
route=route,
start_time=start_time,
parent_otel_span=parent_otel_span,
end_user_id=end_user.end_user_id,
end_user_params=end_user.params,
_end_user_object=end_user.end_user_object,
)
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,
)
except Exception as e:
return await UserAPIKeyAuthExceptionHandler.handle_authentication_error(
e=e,

View file

@ -1,43 +1,162 @@
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from typing import Final
import pytest
from litellm._auth.decision import (
AdminKey,
Allow,
Defer,
CachedAdminKey,
Deny,
DenyReason,
ExpiredKey,
MalformedKey,
MasterKey,
MasterKeyOnlyRoute,
MissingKey,
NoDatabase,
NoMasterKey,
NonAdminKey,
NotAVirtualKey,
Outcome,
PresentedKey,
RequestFacts,
Settings,
UserAuthDisabled,
UserAuthRoute,
VirtualKey,
)
from litellm._auth.driver import drive
from litellm._auth.key import hash_token
from litellm._auth.python_core import PythonCore
MASTER: str = "sk-master"
MASTER: Final = "sk-master"
NOW: Final = datetime(2026, 11, 6, 12, tzinfo=timezone.utc)
SETTINGS: Final = Settings(
master_key=MASTER,
allow_user_auth=False,
has_db=True,
master_key_only_routes=frozenset({"/global/spend/reset"}),
)
@dataclass
class RecordingHandler:
presented: PresentedKey | None = None
calls: tuple[tuple[str, object], ...] = ()
async def resolve_end_user(self, prefetch_key: bool) -> None:
self.calls = (*self.calls, ("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))
return self.presented
async def evict_presented_key(self, hashed_token: str) -> None:
self.calls = (*self.calls, ("evict_presented_key", hashed_token))
async def load_key(self, hashed_token: str) -> None:
self.calls = (*self.calls, ("load_key", hashed_token))
async def _decide(
token: str | None,
*,
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
@pytest.mark.asyncio
@pytest.mark.parametrize(
("route", "token", "settings", "expected"),
[
("/chat/completions", None, Settings(master_key=None, allow_user_auth=False), Allow(NoMasterKey(token=None))),
(
"/chat/completions",
"sk-x",
Settings(master_key=None, allow_user_auth=False),
Allow(NoMasterKey(token="sk-x")),
),
("/chat/completions", None, Settings(master_key=MASTER, allow_user_auth=False), Deny(DenyReason.MISSING_KEY)),
("/chat/completions", "", Settings(master_key=MASTER, allow_user_auth=False), Deny(DenyReason.MALFORMED_KEY)),
("/user/auth", "", Settings(master_key=MASTER, allow_user_auth=True), Deny(DenyReason.MALFORMED_KEY)),
("/user/auth", "sk-x", Settings(master_key=MASTER, allow_user_auth=True), Allow(UserAuthRoute())),
("/user/auth", "sk-x", Settings(master_key=MASTER, allow_user_auth=False), Deny(DenyReason.USER_AUTH_DISABLED)),
("/user/auth", MASTER, Settings(master_key=MASTER, allow_user_auth=False), Deny(DenyReason.USER_AUTH_DISABLED)),
("/chat/completions", MASTER, Settings(master_key=MASTER, allow_user_auth=False), Allow(MasterKey())),
("/key/generate", MASTER, Settings(master_key=MASTER, allow_user_auth=False), Allow(MasterKey())),
("/chat/completions", "sk-master2", Settings(master_key=MASTER, allow_user_auth=False), Defer()),
("/chat/completions", "sk-x", Settings(master_key=MASTER, allow_user_auth=False), Defer()),
("/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, 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", "sk-x", SETTINGS, Deny(UserAuthDisabled())),
("/user/auth", MASTER, SETTINGS, Deny(UserAuthDisabled())),
],
)
def test_decide(route: str, token: str | None, settings: Settings, expected: Outcome):
assert PythonCore().decide(RequestFacts(route=route, token=token), settings) == expected
async def test_credential_decisions_do_no_io(route: str, token: str | None, settings: Settings, expected: Outcome):
assert await _decide(token, route=route, settings=settings) == (expected, ())
@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),))
@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)),
)
@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"))),
)
@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):
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)
)
assert outcome == Allow(CachedAdminKey())
@pytest.mark.asyncio
async def test_expired_cached_admin_key_is_evicted_then_denied():
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)),
)
@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}
@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)
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"]
@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))