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

@@ -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",
+]

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",

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

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:<id>` 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
 

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"
 
 

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"