diff --git a/litellm/proxy/auth/v2/__init__.py b/litellm/proxy/auth/v2/__init__.py index a65822b6ebe..8a65af74f3e 100644 --- a/litellm/proxy/auth/v2/__init__.py +++ b/litellm/proxy/auth/v2/__init__.py @@ -6,7 +6,9 @@ from .context import ( set_auth_context, try_get_auth_context, ) +from .end_user import resolve_end_user from .entry import user_api_key_auth_v2 +from .telemetry import identity_span_attributes __all__ = [ "user_api_key_auth_v2", @@ -16,4 +18,6 @@ __all__ = [ "try_get_auth_context", "set_auth_context", "attach_end_user", + "resolve_end_user", + "identity_span_attributes", ] diff --git a/litellm/proxy/auth/v2/end_user.py b/litellm/proxy/auth/v2/end_user.py new file mode 100644 index 00000000000..d2eb74940e1 --- /dev/null +++ b/litellm/proxy/auth/v2/end_user.py @@ -0,0 +1,42 @@ +from typing import Awaitable, Callable, Optional + +from starlette.requests import Request + +from .context import attach_end_user + +# Extraction reuses the existing request-body/header logic; validation (the +# customer-table lookup) is injected so this stage is testable without a DB and +# so callers choose whether to validate. +Extractor = Callable[[dict, Optional[dict]], Optional[str]] +Validator = Callable[[str], Awaitable[Optional[str]]] + + +def _default_extractor( + request_data: dict, request_headers: Optional[dict] +) -> Optional[str]: + from litellm.proxy.auth.auth_utils import get_end_user_id_from_request_body + + return get_end_user_id_from_request_body(request_data, request_headers) + + +async def resolve_end_user( + request: Request, + request_data: Optional[dict], + request_headers: Optional[dict] = None, + *, + extractor: Optional[Extractor] = None, + validator: Optional[Validator] = None, +) -> Optional[str]: + """Resolve the request's end-user (customer) and record it on the auth context. + + A dedicated stage rather than auth-gate work: it reads the already-published + context via :func:`attach_end_user`. Returns the resolved id, or None when the + request carries no end-user. + """ + extract = extractor or _default_extractor + raw = extract(request_data or {}, request_headers) + if raw is None: + return None + resolved = await validator(raw) if validator is not None else raw + attach_end_user(request, resolved) + return resolved diff --git a/litellm/proxy/auth/v2/entry.py b/litellm/proxy/auth/v2/entry.py index 1a7740004a5..96e813c8759 100644 --- a/litellm/proxy/auth/v2/entry.py +++ b/litellm/proxy/auth/v2/entry.py @@ -5,6 +5,7 @@ from fastapi import HTTPException, Request, status from .authenticators import AuthContext, AuthResult, authenticate from .authorizer import AuthorizationDenied, authorize from .context import AuthMethod, RequestAuthContext, set_auth_context +from .end_user import resolve_end_user from .enforcer import CasbinEnforcer from .policy_store import load_policy_snapshot from .principal import Principal, build_principal @@ -125,6 +126,7 @@ async def user_api_key_auth_v2( status_code=status.HTTP_403_FORBIDDEN, detail=f"auth_v2: not permitted to call model '{requested_model}'", ) + await resolve_end_user(request, request_data, dict(request.headers)) identity.request_route = route return identity diff --git a/litellm/proxy/auth/v2/telemetry.py b/litellm/proxy/auth/v2/telemetry.py new file mode 100644 index 00000000000..18e8d1a8bf8 --- /dev/null +++ b/litellm/proxy/auth/v2/telemetry.py @@ -0,0 +1,29 @@ +from typing import Any, Dict + +from .context import RequestAuthContext + + +def identity_span_attributes(context: RequestAuthContext) -> Dict[str, Any]: + """OTel attributes describing the request's identity, for the route span to set. + + Telemetry lives on the route (OTel v2) and reads the auth context here instead + of auth seeding the span itself. Only present fields are emitted so spans stay + free of empty attributes. + """ + identity = context.identity + attributes: Dict[str, Any] = { + "litellm.auth.method": context.auth_method.value, + "litellm.auth.subject": context.principal.subject, + } + for attr_name, span_key in ( + ("user_id", "litellm.user_id"), + ("team_id", "litellm.team_id"), + ("key_alias", "litellm.key_alias"), + ("org_id", "litellm.org_id"), + ): + value = getattr(identity, attr_name, None) + if value: + attributes[span_key] = value + if context.end_user_id: + attributes["litellm.end_user_id"] = context.end_user_id + return attributes diff --git a/tests/test_litellm/proxy/auth/v2/test_end_user.py b/tests/test_litellm/proxy/auth/v2/test_end_user.py new file mode 100644 index 00000000000..e1211c396d9 --- /dev/null +++ b/tests/test_litellm/proxy/auth/v2/test_end_user.py @@ -0,0 +1,61 @@ +from types import SimpleNamespace + +import pytest + +from litellm.proxy.auth.v2.context import ( + AuthMethod, + RequestAuthContext, + get_auth_context, + set_auth_context, +) +from litellm.proxy.auth.v2.end_user import resolve_end_user +from litellm.proxy.auth.v2.principal import Principal + + +def _request_with_context(): + request = SimpleNamespace(state=SimpleNamespace()) + set_auth_context( + request, + RequestAuthContext( + identity=SimpleNamespace(user_id="u1"), + principal=Principal(subject="user:u1", domain="*", groupings=[]), + auth_method=AuthMethod.VIRTUAL_KEY, + route="/chat/completions", + ), + ) + return request + + +@pytest.mark.asyncio +async def test_resolved_end_user_is_attached_to_context(): + request = _request_with_context() + result = await resolve_end_user( + request, {"user": "cust-1"}, extractor=lambda data, headers: data.get("user") + ) + assert result == "cust-1" + assert get_auth_context(request).end_user_id == "cust-1" + + +@pytest.mark.asyncio +async def test_no_end_user_returns_none_and_leaves_context_unset(): + request = _request_with_context() + result = await resolve_end_user(request, {}, extractor=lambda data, headers: None) + assert result is None + assert get_auth_context(request).end_user_id is None + + +@pytest.mark.asyncio +async def test_validator_can_transform_or_reject(): + request = _request_with_context() + + async def validator(raw): + return f"validated:{raw}" + + result = await resolve_end_user( + request, + {"user": "cust-9"}, + extractor=lambda data, headers: data.get("user"), + validator=validator, + ) + assert result == "validated:cust-9" + assert get_auth_context(request).end_user_id == "validated:cust-9" diff --git a/tests/test_litellm/proxy/auth/v2/test_telemetry.py b/tests/test_litellm/proxy/auth/v2/test_telemetry.py new file mode 100644 index 00000000000..466bdf8eb4d --- /dev/null +++ b/tests/test_litellm/proxy/auth/v2/test_telemetry.py @@ -0,0 +1,44 @@ +from types import SimpleNamespace + +from litellm.proxy.auth.v2.context import AuthMethod, RequestAuthContext +from litellm.proxy.auth.v2.principal import Principal +from litellm.proxy.auth.v2.telemetry import identity_span_attributes + +PRINCIPAL = Principal(subject="user:u1", domain="team:eng", groupings=[]) + + +def _context(identity, end_user_id=None): + return RequestAuthContext( + identity=identity, + principal=PRINCIPAL, + auth_method=AuthMethod.JWT, + route="/chat/completions", + end_user_id=end_user_id, + ) + + +def test_present_identity_fields_become_span_attributes(): + identity = SimpleNamespace( + user_id="u1", team_id="t1", key_alias="prod-key", org_id="o1" + ) + attrs = identity_span_attributes(_context(identity, end_user_id="cust-1")) + assert attrs["litellm.auth.method"] == "jwt" + assert attrs["litellm.auth.subject"] == "user:u1" + assert attrs["litellm.user_id"] == "u1" + assert attrs["litellm.team_id"] == "t1" + assert attrs["litellm.key_alias"] == "prod-key" + assert attrs["litellm.org_id"] == "o1" + assert attrs["litellm.end_user_id"] == "cust-1" + + +def test_absent_fields_are_omitted_not_emitted_empty(): + identity = SimpleNamespace(user_id="u1", team_id=None, key_alias=None, org_id=None) + attrs = identity_span_attributes(_context(identity)) + # Always-present anchors. + assert attrs["litellm.auth.method"] == "jwt" + assert attrs["litellm.user_id"] == "u1" + # Empty/missing fields must not appear at all. + assert "litellm.team_id" not in attrs + assert "litellm.key_alias" not in attrs + assert "litellm.org_id" not in attrs + assert "litellm.end_user_id" not in attrs