From 1685c2a9bc53c5a2e1089c93ab2c8a9ddd492b40 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Fri, 5 Jun 2026 10:18:22 -0700 Subject: [PATCH] feat(proxy): typed RequestAuthContext - the auth_v2 contract for downstream stages Introduces the single, type-safe object every post-auth stage reads instead of poking at request.state with untyped attribute access. auth_v2 populates it once (identity, principal, auth_method, route, end_user_id) after authn+authz and publishes it via set_auth_context; budget/limit hooks, the end-user resolver, and the telemetry span consume it through get_auth_context / try_get_auth_context. - RequestAuthContext is a frozen dataclass; attach_end_user replaces rather than mutates, so one stage can't clobber another's view - identity is annotated under TYPE_CHECKING so the contract is statically typed without coupling the decoupled core to the heavy litellm import - authenticators now advertise an AuthMethod and the chain returns a typed AuthResult, so how a request authenticated is recorded for telemetry - context is set in every branch (control, model-call, loud-open) Verified type-safe with mypy. This is the linchpin for moving budgets, end-user resolution, and telemetry out of the auth gate into their own stages. --- litellm/proxy/auth/v2/__init__.py | 18 +++- litellm/proxy/auth/v2/authenticators.py | 36 +++++++- litellm/proxy/auth/v2/context.py | 65 ++++++++++++++ litellm/proxy/auth/v2/entry.py | 41 ++++++--- .../proxy/auth/v2/test_authenticators.py | 12 +++ .../proxy/auth/v2/test_context.py | 84 +++++++++++++++++++ 6 files changed, 241 insertions(+), 15 deletions(-) create mode 100644 litellm/proxy/auth/v2/context.py create mode 100644 tests/test_litellm/proxy/auth/v2/test_context.py diff --git a/litellm/proxy/auth/v2/__init__.py b/litellm/proxy/auth/v2/__init__.py index 8f7bb6ff767..a65822b6ebe 100644 --- a/litellm/proxy/auth/v2/__init__.py +++ b/litellm/proxy/auth/v2/__init__.py @@ -1,3 +1,19 @@ +from .context import ( + AuthMethod, + RequestAuthContext, + attach_end_user, + get_auth_context, + set_auth_context, + try_get_auth_context, +) from .entry import user_api_key_auth_v2 -__all__ = ["user_api_key_auth_v2"] +__all__ = [ + "user_api_key_auth_v2", + "RequestAuthContext", + "AuthMethod", + "get_auth_context", + "try_get_auth_context", + "set_auth_context", + "attach_end_user", +] diff --git a/litellm/proxy/auth/v2/authenticators.py b/litellm/proxy/auth/v2/authenticators.py index 5dea06b0ef4..efb3a6494b9 100644 --- a/litellm/proxy/auth/v2/authenticators.py +++ b/litellm/proxy/auth/v2/authenticators.py @@ -1,11 +1,27 @@ import secrets -from typing import Any, List, Optional, Protocol, runtime_checkable +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, List, Optional, Protocol, runtime_checkable from fastapi import HTTPException, status +from .context import AuthMethod + +if TYPE_CHECKING: + from litellm.proxy._types import UserAPIKeyAuth + + +@dataclass(frozen=True) +class AuthResult: + """The output of the authenticator chain: the resolved identity and how.""" + + identity: "UserAPIKeyAuth" + method: AuthMethod + @runtime_checkable class Authenticator(Protocol): + method: AuthMethod + def can_handle(self, api_key: Optional[str]) -> bool: ... async def authenticate(self, api_key: str, ctx: "AuthContext") -> Any: ... @@ -35,6 +51,8 @@ class MasterKeyAuthenticator: downstream; a stable alias stands in for it. """ + method = AuthMethod.MASTER_KEY + def can_handle(self, api_key: Optional[str]) -> bool: from litellm.proxy.proxy_server import master_key @@ -60,6 +78,8 @@ class MasterKeyAuthenticator: class VirtualKeyAuthenticator: """Resolves a ``sk-`` virtual key to its identity via the existing key store.""" + method = AuthMethod.VIRTUAL_KEY + def can_handle(self, api_key: Optional[str]) -> bool: return isinstance(api_key, str) and api_key.startswith("sk-") @@ -113,6 +133,8 @@ def _jwt_is_configured() -> bool: class JWTAuthenticator: """Verifies a bearer JWT with authlib and maps its claims to an identity.""" + method = AuthMethod.JWT + def can_handle(self, api_key: Optional[str]) -> bool: # Only claim JWT-shaped tokens when JWT auth is actually configured, so an # unconfigured deployment falls through to a clean 401 instead of a 500. @@ -184,6 +206,8 @@ def _load_introspection_settings() -> Any: class OAuth2IntrospectionAuthenticator: """Validates an opaque bearer token via an RFC 7662 introspection endpoint.""" + method = AuthMethod.OAUTH2 + def can_handle(self, api_key: Optional[str]) -> bool: if not isinstance(api_key, str) or not api_key: return False @@ -248,11 +272,15 @@ AUTHENTICATORS: List[Authenticator] = [ ] -async def authenticate(api_key: Optional[str], ctx: AuthContext) -> Any: +async def authenticate(api_key: Optional[str], ctx: AuthContext) -> AuthResult: """Dispatch by credential shape to the first authenticator that handles it.""" for authenticator in AUTHENTICATORS: - if authenticator.can_handle(api_key): - return await authenticator.authenticate(api_key, ctx) + # The isinstance narrowing is redundant with can_handle at runtime (every + # can_handle requires a str) but makes the str guarantee explicit to the + # type checker before dispatching. + if authenticator.can_handle(api_key) and isinstance(api_key, str): + identity = await authenticator.authenticate(api_key, ctx) + return AuthResult(identity=identity, method=authenticator.method) raise HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, detail="auth_v2: no authenticator for the supplied credential", diff --git a/litellm/proxy/auth/v2/context.py b/litellm/proxy/auth/v2/context.py new file mode 100644 index 00000000000..16baa017262 --- /dev/null +++ b/litellm/proxy/auth/v2/context.py @@ -0,0 +1,65 @@ +from dataclasses import dataclass, replace +from enum import Enum +from typing import TYPE_CHECKING, Optional + +from starlette.requests import Request + +from .principal import Principal + +if TYPE_CHECKING: + from litellm.proxy._types import UserAPIKeyAuth + + +class AuthMethod(str, Enum): + """How a request authenticated. Recorded for telemetry and debugging.""" + + VIRTUAL_KEY = "virtual_key" + MASTER_KEY = "master_key" + JWT = "jwt" + OAUTH2 = "oauth2" + ANONYMOUS = "anonymous" + + +@dataclass(frozen=True) +class RequestAuthContext: + """The single, typed auth contract every downstream stage reads. + + auth_v2 populates this once, after authentication and authorization, and + stores it on ``request.state``. Budget/limit hooks, the end-user resolver, and + the telemetry span all read it instead of reaching into ``request.state`` with + untyped attribute access. Frozen so a stage can't mutate another stage's view; + :func:`attach_end_user` produces a replaced copy. + """ + + identity: "UserAPIKeyAuth" + principal: Principal + auth_method: AuthMethod + route: str + end_user_id: Optional[str] = None + + +_STATE_ATTR = "auth_v2_context" + + +def set_auth_context(request: Request, context: RequestAuthContext) -> None: + setattr(request.state, _STATE_ATTR, context) + + +def get_auth_context(request: Request) -> RequestAuthContext: + """Return the context, or raise if auth hasn't run for this request.""" + context: Optional[RequestAuthContext] = getattr(request.state, _STATE_ATTR, None) + if context is None: + raise LookupError("auth_v2 context is not set on this request") + return context + + +def try_get_auth_context(request: Request) -> Optional[RequestAuthContext]: + """Return the context, or None when auth_v2 did not run (e.g. v1 path).""" + return getattr(request.state, _STATE_ATTR, None) + + +def attach_end_user(request: Request, end_user_id: Optional[str]) -> RequestAuthContext: + """Record the resolved end-user on the context (the end-user stage's output).""" + updated = replace(get_auth_context(request), end_user_id=end_user_id) + set_auth_context(request, updated) + return updated diff --git a/litellm/proxy/auth/v2/entry.py b/litellm/proxy/auth/v2/entry.py index f6a336496f8..1a7740004a5 100644 --- a/litellm/proxy/auth/v2/entry.py +++ b/litellm/proxy/auth/v2/entry.py @@ -1,9 +1,10 @@ -from typing import Any, Optional +from typing import Any, Optional, Tuple from fastapi import HTTPException, Request, status -from .authenticators import AuthContext, authenticate +from .authenticators import AuthContext, AuthResult, authenticate from .authorizer import AuthorizationDenied, authorize +from .context import AuthMethod, RequestAuthContext, set_auth_context from .enforcer import CasbinEnforcer from .policy_store import load_policy_snapshot from .principal import Principal, build_principal @@ -36,7 +37,7 @@ async def _build_enforcer(principal: Principal, prisma_client: Any) -> CasbinEnf ) -async def _best_effort_identity(api_key: Optional[str], ctx: AuthContext) -> Any: +async def _best_effort_identity(api_key: Optional[str], ctx: AuthContext) -> AuthResult: """On loud-open routes, use the real identity if a usable key is present, otherwise fall back to an anonymous principal. Never fails the request.""" if isinstance(api_key, str) and api_key.startswith("sk-"): @@ -44,7 +45,26 @@ async def _best_effort_identity(api_key: Optional[str], ctx: AuthContext) -> Any return await authenticate(api_key, ctx) except Exception: pass - return await _anonymous_identity(api_key) + identity = await _anonymous_identity(api_key) + return AuthResult(identity=identity, method=AuthMethod.ANONYMOUS) + + +def _establish_context( + request: Request, result: AuthResult, route: str +) -> Tuple[Principal, Any]: + """Derive the principal and publish the typed auth context for downstream + stages (budget/limit hooks, end-user resolver, telemetry span).""" + principal = build_principal(result.identity) + set_auth_context( + request, + RequestAuthContext( + identity=result.identity, + principal=principal, + auth_method=result.method, + route=route, + ), + ) + return principal, result.identity async def user_api_key_auth_v2( @@ -73,9 +93,9 @@ async def user_api_key_auth_v2( rule = match_route(route, request.method) if rule is not None: # Control plane: RBAC over management resources. - identity = await authenticate(token, ctx) + result = await authenticate(token, ctx) request_data = await _read_request_body(request=request) - principal = build_principal(identity) + principal, identity = _establish_context(request, result, route) enforcer = await _build_enforcer(principal, prisma_client) try: @@ -90,13 +110,13 @@ async def user_api_key_auth_v2( # Model calls are a permission like any other: the `call` action on the # `model:` object, decided by the same role system. The legacy # key.models / access-group mechanism is intentionally not consulted. - identity = await authenticate(token, ctx) + result = await authenticate(token, ctx) request_data = await _read_request_body(request=request) + principal, identity = _establish_context(request, result, route) requested_model = ( request_data.get("model") if isinstance(request_data, dict) else None ) if requested_model: - principal = build_principal(identity) enforcer = await _build_enforcer(principal, prisma_client) if not enforcer.enforce( principal.subject, principal.domain, f"model:{requested_model}", "call" @@ -109,8 +129,9 @@ async def user_api_key_auth_v2( return identity # Loud-open: route v2 doesn't yet govern. No identity required. - identity = await _best_effort_identity(token, ctx) - authorize(build_principal(identity), route, None, _DENY_ALL, request.method) + result = await _best_effort_identity(token, ctx) + principal, identity = _establish_context(request, result, route) + authorize(principal, route, None, _DENY_ALL, request.method) identity.request_route = route return identity diff --git a/tests/test_litellm/proxy/auth/v2/test_authenticators.py b/tests/test_litellm/proxy/auth/v2/test_authenticators.py index 45401cc534d..c1b31f44981 100644 --- a/tests/test_litellm/proxy/auth/v2/test_authenticators.py +++ b/tests/test_litellm/proxy/auth/v2/test_authenticators.py @@ -1,13 +1,25 @@ import pytest from litellm.proxy.auth.v2.authenticators import ( + OAuth2IntrospectionAuthenticator, JWTAuthenticator, MasterKeyAuthenticator, VirtualKeyAuthenticator, ) +from litellm.proxy.auth.v2.context import AuthMethod JWT_SHAPED = "header.payload.signature" + +def test_each_authenticator_advertises_its_method(): + # The method is recorded on the auth context for telemetry, so the chain can + # tag how a request authenticated without re-deriving it. + assert MasterKeyAuthenticator().method is AuthMethod.MASTER_KEY + assert VirtualKeyAuthenticator().method is AuthMethod.VIRTUAL_KEY + assert JWTAuthenticator().method is AuthMethod.JWT + assert OAuth2IntrospectionAuthenticator().method is AuthMethod.OAUTH2 + + MASTER = "sk-master-secret-123" diff --git a/tests/test_litellm/proxy/auth/v2/test_context.py b/tests/test_litellm/proxy/auth/v2/test_context.py new file mode 100644 index 00000000000..03e108cd53a --- /dev/null +++ b/tests/test_litellm/proxy/auth/v2/test_context.py @@ -0,0 +1,84 @@ +import dataclasses +from types import SimpleNamespace + +import pytest + +from litellm.proxy.auth.v2.context import ( + AuthMethod, + RequestAuthContext, + attach_end_user, + get_auth_context, + set_auth_context, + try_get_auth_context, +) +from litellm.proxy.auth.v2.principal import Principal + +PRINCIPAL = Principal(subject="user:u1", domain="*", groupings=[]) + + +def _request(): + # request.state is the only surface the accessors touch; a namespace stands in. + return SimpleNamespace(state=SimpleNamespace()) + + +def _context(**overrides): + base = dict( + identity=SimpleNamespace(user_id="u1"), + principal=PRINCIPAL, + auth_method=AuthMethod.VIRTUAL_KEY, + route="/chat/completions", + ) + base.update(overrides) + return RequestAuthContext(**base) + + +def test_set_then_get_roundtrips(): + request = _request() + ctx = _context() + set_auth_context(request, ctx) + assert get_auth_context(request) is ctx + + +def test_get_without_context_raises(): + with pytest.raises(LookupError): + get_auth_context(_request()) + + +def test_try_get_without_context_returns_none(): + assert try_get_auth_context(_request()) is None + + +def test_context_is_frozen(): + ctx = _context() + with pytest.raises(dataclasses.FrozenInstanceError): + ctx.route = "/mutated" # type: ignore[misc] + + +def test_attach_end_user_replaces_without_mutating_original(): + request = _request() + original = _context() + set_auth_context(request, original) + + updated = attach_end_user(request, "customer-42") + + assert updated.end_user_id == "customer-42" + assert get_auth_context(request).end_user_id == "customer-42" + # The original frozen instance is untouched; attach produced a copy. + assert original.end_user_id is None + # Everything else carries over unchanged. + assert updated.identity is original.identity + assert updated.auth_method is AuthMethod.VIRTUAL_KEY + + +def test_attach_end_user_requires_existing_context(): + with pytest.raises(LookupError): + attach_end_user(_request(), "customer-42") + + +def test_auth_method_values_are_stable(): + # Recorded on spend logs / telemetry, so the string values are a contract. + assert AuthMethod.VIRTUAL_KEY.value == "virtual_key" + assert AuthMethod.MASTER_KEY.value == "master_key" + assert AuthMethod.JWT.value == "jwt" + assert AuthMethod.OAUTH2.value == "oauth2" + assert AuthMethod.ANONYMOUS.value == "anonymous"