diff --git a/litellm/proxy/auth/impersonation.py b/litellm/proxy/auth/impersonation.py new file mode 100644 index 00000000000..1ebef1387ff --- /dev/null +++ b/litellm/proxy/auth/impersonation.py @@ -0,0 +1,120 @@ +"""Kubernetes-style impersonation for service keys: a key with the `impersonate` permission acts as a named user.""" + +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Final, Protocol, TypeAlias + +from pydantic import BaseModel, ConfigDict + +from litellm.proxy._types import LiteLLM_UserTable, LitellmUserRoles +from litellm.proxy.auth.auth_checks import get_user_role +from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth + +IMPERSONATE_USER_HEADER: Final = "X-LiteLLM-Impersonate-User" +IMPERSONATE_EMAIL_HEADER: Final = "X-LiteLLM-Impersonate-Email" +IMPERSONATE_PERMISSION: Final = "impersonate" + + +@dataclass(frozen=True, slots=True) +class ImpersonationRequest: + user_id: str + email: str | None + + +@dataclass(frozen=True, slots=True) +class ImpersonationForbidden: + detail: str + + +@dataclass(frozen=True, slots=True) +class ImpersonatedUserNotFound: + user_id: str + + +ImpersonationError: TypeAlias = ImpersonationForbidden | ImpersonatedUserNotFound + + +class _KeyPermissions(BaseModel): + """Typed view of the key's `permissions` JSON column, read off the auth object by attribute.""" + + model_config = ConfigDict(frozen=True, from_attributes=True, extra="ignore") + + permissions: Mapping[str, object] = {} + + +class _ActorRoutes(BaseModel): + """Typed view of the service key's route allowlist, kept as a ceiling on the impersonated request.""" + + model_config = ConfigDict(frozen=True, from_attributes=True, extra="ignore") + + allowed_routes: tuple[str, ...] | None = None + + +class _SubjectGrants(BaseModel): + """Typed view of the untyped list/dict columns on the user row.""" + + model_config = ConfigDict(frozen=True, from_attributes=True, extra="ignore") + + models: tuple[str, ...] = () + model_max_budget: Mapping[str, object] | None = None + + +class UserLookup(Protocol): + async def __call__(self, user_id: str) -> LiteLLM_UserTable | None: ... + + +def read_impersonation_request(headers: Mapping[str, str]) -> ImpersonationRequest | None: + user_id: Final = headers.get(IMPERSONATE_USER_HEADER, "").strip() + email: Final = headers.get(IMPERSONATE_EMAIL_HEADER, "").strip() + if not user_id: + return None + return ImpersonationRequest(user_id=user_id, email=email or None) + + +def key_may_impersonate(permissions: Mapping[str, object]) -> bool: + return permissions.get(IMPERSONATE_PERMISSION) is True + + +def actor_label(actor: UserAPIKeyAuth) -> str: + return actor.key_alias or actor.token or "unknown-key" + + +async def impersonate( + actor: UserAPIKeyAuth, + request: ImpersonationRequest, + lookup: UserLookup, +) -> UserAPIKeyAuth | ImpersonationError: + if not key_may_impersonate(_KeyPermissions.model_validate(actor).permissions): + return ImpersonationForbidden(f"key {actor_label(actor)} lacks the {IMPERSONATE_PERMISSION!r} permission") + subject: Final = await lookup(request.user_id) + if subject is None or not _email_matches(request.email, subject.user_email): + return ImpersonatedUserNotFound(request.user_id) + return _as_subject(actor, subject) + + +def _email_matches(claimed: str | None, actual: str | None) -> bool: + return claimed is None or actual is None or claimed.casefold() == actual.casefold() + + +def _as_subject(actor: UserAPIKeyAuth, subject: LiteLLM_UserTable) -> UserAPIKeyAuth: + """Only the key's identity and its route allowlist survive, the allowlist stays a ceiling on what the control + plane may do here. Team, models and permissions come from the subject, so a user never inherits what the + service key may do. Budgets stay off the key fields, the user budget check reads the + subject's own row by ``user_id``, a shared service key must not pool every subject's spend.""" + grants: Final = _SubjectGrants.model_validate(subject) + routes: Final = _ActorRoutes.model_validate(actor).allowed_routes + return UserAPIKeyAuth( + api_key=actor.api_key, + token=actor.token, + key_alias=actor.key_alias, + allowed_routes=None if routes is None else list(routes), + user_id=subject.user_id, + user_email=subject.user_email, + user_role=get_user_role(subject) or LitellmUserRoles.INTERNAL_USER, + models=list(grants.models), + user_tpm_limit=subject.tpm_limit, + user_rpm_limit=subject.rpm_limit, + user_model_max_budget=grants.model_max_budget, + parent_otel_span=actor.parent_otel_span, + impersonated_by=actor_label(actor), + ) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index cdb3e685c20..75f6d0908d0 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -18,6 +18,7 @@ from typing import Final, NamedTuple, Protocol, Union, cast import fastapi from fastapi import HTTPException, Request, WebSocket, status from fastapi.security.api_key import APIKeyHeader +from starlette.datastructures import Headers from starlette.exceptions import WebSocketException import litellm @@ -117,6 +118,15 @@ from litellm.proxy.auth.auth_utils import ( route_in_additonal_public_routes, ) from litellm.proxy.auth.handle_jwt import JWTAuthManager, JWTHandler +from litellm.proxy.auth.impersonation import ( + IMPERSONATE_USER_HEADER, + ImpersonatedUserNotFound, + ImpersonationForbidden, + ImpersonationRequest, + actor_label, + impersonate, + read_impersonation_request, +) from litellm.proxy.auth.network import TrustedProxyConfig, resolve_network_context from litellm.proxy.auth.oauth2_check import Oauth2Handler from litellm.proxy.auth.oauth2_proxy_hook import handle_oauth2_proxy_request @@ -185,6 +195,7 @@ from litellm.router_utils.fallback_event_handlers import ( get_fallback_model_group_for_lookup_groups, ) from litellm.secret_managers.main import get_secret_bool +from litellm.types.proxy.auth.auth_checks import UserNotFoundError from litellm.types.services import ServiceTypes try: @@ -3210,7 +3221,63 @@ async def validate_resolved_virtual_key( # noqa: C901 # Preserve ordering of e start_time=start_time, ) virtual_key_auth_obj.via_virtual_key = True - return virtual_key_auth_obj + impersonation: Final = read_impersonation_request( + Headers(scope=request.scope) if "headers" in request.scope else Headers() + ) + if impersonation is None: + return virtual_key_auth_obj + return await _impersonate_or_raise( + actor=virtual_key_auth_obj, + impersonation=impersonation, + 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, + ) + + +async def _impersonate_or_raise( + actor: UserAPIKeyAuth, + impersonation: ImpersonationRequest, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Span | None, + proxy_logging_obj: ProxyLogging, + route: str, +) -> UserAPIKeyAuth: + async def lookup(user_id: str) -> LiteLLM_UserTable | None: + try: + return await get_user_object( + user_id=user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + user_id_upsert=False, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + except UserNotFoundError: + return None + + result: Final = await impersonate(actor, impersonation, lookup) + match result: + case UserAPIKeyAuth(): + verbose_proxy_logger.info( + "impersonation: actor=%s subject=%s route=%s", actor_label(actor), result.user_id, route + ) + result.via_virtual_key = actor.via_virtual_key + return result + case ImpersonationForbidden(detail=detail): + raise ProxyException( + message=detail, type=ProxyErrorTypes.auth_error, param=IMPERSONATE_USER_HEADER, code=403 + ) + case ImpersonatedUserNotFound(user_id=user_id): + raise ProxyException( + message=f"impersonated user {user_id!r} does not exist on this proxy", + type=ProxyErrorTypes.auth_error, + param=IMPERSONATE_USER_HEADER, + code=401, + ) async def _safe_fetch(label: str, awaitable): diff --git a/litellm/proxy/control_plane/app.py b/litellm/proxy/control_plane/app.py index 7c450ed6fe7..2fab0c49557 100644 --- a/litellm/proxy/control_plane/app.py +++ b/litellm/proxy/control_plane/app.py @@ -9,6 +9,8 @@ from litellm.proxy.control_plane.auth.router import get_login_service from litellm.proxy.control_plane.auth.router import router as auth_router from litellm.proxy.control_plane.common.exceptions import install_exception_handlers from litellm.proxy.control_plane.container import Container +from litellm.proxy.control_plane.dataplanes.dependencies import get_dataplane_client, get_registry +from litellm.proxy.control_plane.dataplanes.router import router as dataplanes_router from litellm.proxy.control_plane.health.router import router as health_router TITLE: Final = "LiteLLM Control Plane" @@ -36,7 +38,10 @@ def create_app(container: Container, license_check: License) -> FastAPI: ) app.dependency_overrides[get_token_verifier] = lambda: container.token_verifier app.dependency_overrides[get_login_service] = lambda: container.login_service + app.dependency_overrides[get_registry] = lambda: container.registry + app.dependency_overrides[get_dataplane_client] = lambda: container.dataplane_client install_exception_handlers(app) app.include_router(health_router) app.include_router(auth_router, prefix=container.settings.api_prefix) + app.include_router(dataplanes_router, prefix=container.settings.api_prefix) return app diff --git a/litellm/proxy/control_plane/common/exceptions.py b/litellm/proxy/control_plane/common/exceptions.py index 7e4d62ce43c..0c6dd7486a1 100644 --- a/litellm/proxy/control_plane/common/exceptions.py +++ b/litellm/proxy/control_plane/common/exceptions.py @@ -39,6 +39,10 @@ def too_many_requests(detail: str) -> ProblemDetail: return problem("too-many-requests", "Too Many Requests", 429, detail) +def bad_gateway(detail: str) -> ProblemDetail: + return problem("bad-gateway", "Bad Gateway", 502, detail) + + class ControlPlaneException(ManagementProblem): """A problem raised from a dependency or handler, rendered by the app-level handler below.""" diff --git a/litellm/proxy/control_plane/container.py b/litellm/proxy/control_plane/container.py index e1caa26f485..c62836704e3 100644 --- a/litellm/proxy/control_plane/container.py +++ b/litellm/proxy/control_plane/container.py @@ -7,6 +7,12 @@ from datetime import timedelta from types import MappingProxyType from typing import Final +import httpx + +from litellm.llms.custom_httpx.http_handler import ( + AsyncHTTPHandler, + get_async_httpx_client, # pyright: ignore[reportUnknownVariableType] # params is an untyped dict upstream +) from litellm.proxy.control_plane.auth.jwt import JwksKeyLookup, SessionTokenCodec, TokenVerifier, build_token_verifier from litellm.proxy.control_plane.auth.login import LoginService, LoginThrottle from litellm.proxy.control_plane.auth.rbac import ( @@ -18,7 +24,11 @@ from litellm.proxy.control_plane.auth.rbac import ( build_policy, parse_role_definition, ) +from litellm.proxy.control_plane.dataplanes.client import DataplaneClient, HttpTransport, RetryPolicy +from litellm.proxy.control_plane.dataplanes.credentials import TokenExchanger, http_token_exchanger +from litellm.proxy.control_plane.dataplanes.registry import DataplaneRegistry, build_registry from litellm.proxy.control_plane.settings import ControlPlaneSettings +from litellm.types.llms.custom_http import httpxSpecialProvider @dataclass(frozen=True, slots=True) @@ -27,6 +37,8 @@ class Container: policy: RolePolicy token_verifier: TokenVerifier login_service: LoginService + registry: DataplaneRegistry + dataplane_client: DataplaneClient class InvalidRolePolicy(ValueError): @@ -57,9 +69,52 @@ def build_role_policy(settings: ControlPlaneSettings) -> RolePolicy: raise InvalidRolePolicy(policy) -def build_container(settings: ControlPlaneSettings, jwks: JwksKeyLookup | None = None) -> Container: +def _shared_http_handler(settings: ControlPlaneSettings) -> AsyncHTTPHandler: + return get_async_httpx_client( + llm_provider=httpxSpecialProvider.ControlPlane, params={"timeout": settings.dataplane_timeout_seconds} + ) + + +@dataclass(frozen=True, slots=True) +class HandlerTransport: + """Plain httpx requests over the shared handler's pool, re-read per call because the handler cache closes and + recreates the pool on eviction. No handler retries, only httpx errors.""" + + handler: AsyncHTTPHandler + + async def request( + self, + method: str, + url: str, + *, + params: dict[str, str] | None = None, # mutable-ok: httpx takes dict + json: dict[str, object] | None = None, # mutable-ok: httpx takes dict + headers: dict[str, str] | None = None, # mutable-ok: httpx takes dict + timeout: float, + ) -> httpx.Response: + return await self.handler.client.request( + method, url, params=params, json=json, headers=headers, timeout=timeout, follow_redirects=False + ) + + +def _shared_http_client(settings: ControlPlaneSettings) -> HttpTransport: + return HandlerTransport(_shared_http_handler(settings)) + + +def build_container( + settings: ControlPlaneSettings, + jwks: JwksKeyLookup | None = None, + transport: HttpTransport | None = None, + exchanger: TokenExchanger | None = None, +) -> Container: master_key: Final = settings.master_key.get_secret_value() policy: Final = build_role_policy(settings) + http: Final = transport if transport is not None else _shared_http_client(settings) + exchange: Final = ( + exchanger + if exchanger is not None + else http_token_exchanger(_shared_http_handler(settings), settings.dataplane_timeout_seconds) + ) sessions: Final = SessionTokenCodec( master_key=master_key, issuer=settings.session_issuer, @@ -80,4 +135,10 @@ def build_container(settings: ControlPlaneSettings, jwks: JwksKeyLookup | None = admin_password=None if password is None else password.get_secret_value(), throttle=LoginThrottle(time.monotonic), ), + registry=build_registry(settings.dataplanes, exchange), + dataplane_client=DataplaneClient( + transport=http, + timeout_seconds=settings.dataplane_timeout_seconds, + retry=RetryPolicy(attempts=settings.dataplane_retry_attempts), + ), ) diff --git a/litellm/proxy/control_plane/dataplanes/__init__.py b/litellm/proxy/control_plane/dataplanes/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/proxy/control_plane/dataplanes/client.py b/litellm/proxy/control_plane/dataplanes/client.py new file mode 100644 index 00000000000..b677d449a1f --- /dev/null +++ b/litellm/proxy/control_plane/dataplanes/client.py @@ -0,0 +1,170 @@ +"""One HTTP call to one dataplane: credential injection, timeout, bounded retry with jitter, failure as a value. + +The transport is a plain ``httpx.AsyncClient`` (the pool behind ``get_async_httpx_client``), not ``AsyncHTTPHandler``: +the handler retries writes on connection errors, replays PATCH as POST and raises ``litellm.Timeout``, none of which +a management passthrough can tolerate. Retrying is this module's job and only ever happens for GET.""" + +import asyncio +import random +from collections.abc import Awaitable, Callable, Mapping +from dataclasses import dataclass +from typing import Final, Literal, Protocol, TypeAlias + +import httpx +from pydantic import TypeAdapter + +from litellm.proxy.control_plane.dataplanes.credentials import Caller, CredentialKind, CredentialUnavailable +from litellm.proxy.control_plane.dataplanes.registry import Dataplane + +Method: TypeAlias = Literal["GET", "POST", "PUT", "PATCH", "DELETE"] +Sleeper: TypeAlias = Callable[[float], Awaitable[None]] +RETRYABLE_STATUSES: Final = frozenset({502, 503, 504}) +_JSON: Final = TypeAdapter(object) + + +class HttpTransport(Protocol): + """The slice of ``httpx.AsyncClient.request`` the client needs, so tests inject a fake. Raises only httpx errors + and never retries on its own.""" + + async def request( + self, + method: str, + url: str, + *, + params: dict[str, str] | None = None, # mutable-ok: httpx takes dict + json: dict[str, object] | None = None, # mutable-ok: httpx takes dict + headers: dict[str, str] | None = None, # mutable-ok: httpx takes dict + timeout: float, + ) -> httpx.Response: ... + + +@dataclass(frozen=True, slots=True) +class DataplaneReply: + status: int + body: object + credential: CredentialKind + + +@dataclass(frozen=True, slots=True) +class DataplaneUnreachable: + detail: str + attempts: int + + +CallResult: TypeAlias = DataplaneReply | DataplaneUnreachable | CredentialUnavailable + + +@dataclass(frozen=True, slots=True) +class _RawReply: + status: int + body: object + + +@dataclass(frozen=True, slots=True) +class RetryPolicy: + attempts: int = 3 + base_delay_seconds: float = 0.2 + max_delay_seconds: float = 2.0 + + def delay(self, attempt: int, jitter: float) -> float: + capped: Final = min(self.max_delay_seconds, self.base_delay_seconds * (2.0**attempt)) + return capped * (0.5 + jitter / 2) + + +def _parse_body(response: httpx.Response) -> object: + try: + return _JSON.validate_python(response.json()) + except ValueError: + return response.text + + +@dataclass(frozen=True, slots=True) +class DataplaneClient: + transport: HttpTransport + timeout_seconds: float + retry: RetryPolicy = RetryPolicy() + sleep: Sleeper = asyncio.sleep + jitter: Callable[[], float] = random.random + + async def get( + self, + dataplane: Dataplane, + path: str, + caller: Caller, + params: Mapping[str, str] | None = None, + ) -> CallResult: + return await self._call("GET", dataplane, path, caller, params=params, json=None) + + async def post(self, dataplane: Dataplane, path: str, caller: Caller, json: Mapping[str, object]) -> CallResult: + return await self._call("POST", dataplane, path, caller, params=None, json=json) + + async def _call( + self, + method: Method, + dataplane: Dataplane, + path: str, + caller: Caller, + params: Mapping[str, str] | None, + json: Mapping[str, object] | None, + ) -> CallResult: + credential: Final = await dataplane.credential.headers(caller) + if isinstance(credential, CredentialUnavailable): + return credential + url: Final = f"{dataplane.url}/{path.lstrip('/')}" + attempts: Final = self.retry.attempts if method == "GET" else 1 + outcome: Final = await self._attempt( + method, url, credential.headers, params, json, attempt=0, attempts=attempts + ) + return ( + outcome + if isinstance(outcome, DataplaneUnreachable) + else DataplaneReply(outcome.status, outcome.body, credential.kind) + ) + + async def _attempt( + self, + method: Method, + url: str, + headers: Mapping[str, str], + params: Mapping[str, str] | None, + json: Mapping[str, object] | None, + attempt: int, + attempts: int, + ) -> _RawReply | DataplaneUnreachable: + outcome: Final = await self._send(method, url, headers, params, json) + last: Final = attempt + 1 >= attempts + if last or not _retryable(outcome): + return outcome if isinstance(outcome, _RawReply) else DataplaneUnreachable(outcome.detail, attempt + 1) + await self.sleep(self.retry.delay(attempt, self.jitter())) + return await self._attempt(method, url, headers, params, json, attempt + 1, attempts) + + async def _send( + self, + method: Method, + url: str, + headers: Mapping[str, str], + params: Mapping[str, str] | None, + json: Mapping[str, object] | None, + ) -> _RawReply | DataplaneUnreachable: + try: + response: Final = await self.transport.request( + method, + url, + params=None if params is None else dict(params), + json=None if json is None else dict(json), + headers=dict(headers), + timeout=self.timeout_seconds, + ) + except httpx.HTTPError as error: + return DataplaneUnreachable(_describe(error), attempts=1) + return _RawReply(response.status_code, _parse_body(response)) + + +def _describe(error: httpx.HTTPError) -> str: + """Never echo the request URL, its query string may carry a secret the audit log must not keep.""" + message: Final = str(error) + return type(error).__name__ if not message or "://" in message else f"{type(error).__name__}: {message}" + + +def _retryable(outcome: _RawReply | DataplaneUnreachable) -> bool: + return isinstance(outcome, DataplaneUnreachable) or outcome.status in RETRYABLE_STATUSES diff --git a/litellm/proxy/control_plane/dataplanes/credentials.py b/litellm/proxy/control_plane/dataplanes/credentials.py new file mode 100644 index 00000000000..dc5bb75f624 --- /dev/null +++ b/litellm/proxy/control_plane/dataplanes/credentials.py @@ -0,0 +1,136 @@ +"""How the control plane authenticates to a dataplane without losing who the caller is.""" + +from collections.abc import Awaitable, Callable, Mapping +from dataclasses import dataclass +from types import MappingProxyType +from typing import Final, Literal, Protocol, TypeAlias + +from litellm.exceptions import APIConnectionError, Timeout +from litellm.litellm_core_utils.oauth_token_exchange import ( + OAuthTokenExchangeConfig, + OAuthTokenExchangeError, + aexchange_token, +) +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler +from litellm.proxy.auth.impersonation import IMPERSONATE_EMAIL_HEADER, IMPERSONATE_USER_HEADER +from litellm.proxy.control_plane.auth.rbac import Principal + +CredentialKind: TypeAlias = Literal["oidc_exchange", "service_token"] +TokenExchanger: TypeAlias = Callable[[OAuthTokenExchangeConfig, str], Awaitable[str]] + + +@dataclass(frozen=True, slots=True) +class Caller: + """The authenticated control-plane principal plus the raw bearer it presented, needed for token exchange.""" + + principal: Principal + bearer: str + + +@dataclass(frozen=True, slots=True) +class CredentialUnavailable: + kind: CredentialKind + reason: str + + +@dataclass(frozen=True, slots=True) +class IssuedCredential: + kind: CredentialKind + headers: Mapping[str, str] + + +CredentialResult: TypeAlias = IssuedCredential | CredentialUnavailable + + +class DataplaneCredential(Protocol): + async def headers(self, caller: Caller) -> CredentialResult: ... + + +def _bearer(token: str) -> Mapping[str, str]: + return MappingProxyType({"Authorization": f"Bearer {token}"}) + + +def _issued(kind: CredentialKind, headers: Mapping[str, str]) -> IssuedCredential: + return IssuedCredential(kind, MappingProxyType(dict(headers))) + + +@dataclass(frozen=True, slots=True) +class OidcTokenExchange: + """RFC 8693 (or jwt-bearer on-behalf-of) exchange of the caller's IdP token for a dataplane-audience token.""" + + config: OAuthTokenExchangeConfig + exchange: TokenExchanger + + async def headers(self, caller: Caller) -> CredentialResult: + kind: Final = caller.principal.token_kind + if kind != "oidc": + return CredentialUnavailable( + "oidc_exchange", f"token exchange needs an IdP token, caller presented a {kind} token" + ) + try: + return _issued("oidc_exchange", _bearer(await self.exchange(self.config, caller.bearer))) + except OAuthTokenExchangeError as error: + return CredentialUnavailable("oidc_exchange", f"token exchange failed with status {error.status_code}") + + +@dataclass(frozen=True, slots=True) +class ServiceTokenImpersonation: + """A dataplane virtual key with the `impersonate` permission, asserting the caller via X-LiteLLM-Impersonate-*.""" + + token: str + + async def headers(self, caller: Caller) -> CredentialResult: + principal: Final = caller.principal + if principal.token_kind == "master_key": + return _issued("service_token", _bearer(self.token)) + email: Final = {} if principal.email is None else {IMPERSONATE_EMAIL_HEADER: principal.email} + return _issued("service_token", {**_bearer(self.token), IMPERSONATE_USER_HEADER: principal.subject, **email}) + + +@dataclass(frozen=True, slots=True) +class CredentialChain: + """Tries each strategy in order and returns the first one that can authenticate this caller.""" + + strategies: tuple[DataplaneCredential, ...] + + async def headers(self, caller: Caller) -> CredentialResult: + return await _first_available(self.strategies, caller, ()) + + @property + def kinds(self) -> tuple[CredentialKind, ...]: + return tuple(kind_of(strategy) for strategy in self.strategies) + + +async def _first_available( + strategies: tuple[DataplaneCredential, ...], + caller: Caller, + failures: tuple[CredentialUnavailable, ...], +) -> CredentialResult: + if not strategies: + reasons: Final = "; ".join(f"{failure.kind}: {failure.reason}" for failure in failures) + return CredentialUnavailable( + failures[-1].kind if failures else "service_token", reasons or "no credential configured" + ) + result: Final = await strategies[0].headers(caller) + if not isinstance(result, CredentialUnavailable): + return result + return await _first_available(strategies[1:], caller, (*failures, result)) + + +def kind_of(strategy: DataplaneCredential) -> CredentialKind: + return "oidc_exchange" if isinstance(strategy, OidcTokenExchange) else "service_token" + + +def http_token_exchanger(client: AsyncHTTPHandler, timeout: float) -> TokenExchanger: + """``AsyncHTTPHandler.post`` turns transport failures into litellm exceptions, map them back to the exchange + error so the credential chain can fall through instead of surfacing a 500.""" + + async def exchange(config: OAuthTokenExchangeConfig, subject_token: str) -> str: + try: + return await aexchange_token(client, config, subject_token, timeout=timeout) + except (Timeout, APIConnectionError) as error: + raise OAuthTokenExchangeError( + status_code=502, message=f"token exchange transport failed: {type(error).__name__}" + ) from None + + return exchange diff --git a/litellm/proxy/control_plane/dataplanes/dependencies.py b/litellm/proxy/control_plane/dataplanes/dependencies.py new file mode 100644 index 00000000000..c9ea7441c37 --- /dev/null +++ b/litellm/proxy/control_plane/dataplanes/dependencies.py @@ -0,0 +1,43 @@ +"""FastAPI glue for the dataplanes domain: registry and client injection, per-dataplane access checks.""" + +from typing import Annotated, Final + +from fastapi import Depends +from fastapi.security import HTTPAuthorizationCredentials + +from litellm.proxy.control_plane.auth.dependencies import bearer_scheme, get_principal +from litellm.proxy.control_plane.auth.rbac import Principal +from litellm.proxy.control_plane.common.exceptions import ControlPlaneException, forbidden, not_found, unauthorized +from litellm.proxy.control_plane.dataplanes.client import DataplaneClient +from litellm.proxy.control_plane.dataplanes.credentials import Caller +from litellm.proxy.control_plane.dataplanes.registry import Dataplane, DataplaneRegistry + + +def get_registry() -> DataplaneRegistry: + raise RuntimeError("get_registry is not bound, build the app with create_app") + + +def get_dataplane_client() -> DataplaneClient: + raise RuntimeError("get_dataplane_client is not bound, build the app with create_app") + + +async def get_caller( + credentials: Annotated[HTTPAuthorizationCredentials | None, Depends(bearer_scheme)], + principal: Annotated[Principal, Depends(get_principal)], +) -> Caller: + if credentials is None: + raise ControlPlaneException(unauthorized("Missing bearer credentials.")) + return Caller(principal=principal, bearer=credentials.credentials) + + +async def dataplane_access( + dataplane_id: str, + principal: Annotated[Principal, Depends(get_principal)], + registry: Annotated[DataplaneRegistry, Depends(get_registry)], +) -> Dataplane: + dataplane: Final = registry.get(dataplane_id) + if dataplane is None: + raise ControlPlaneException(not_found(f"Unknown dataplane {dataplane_id!r}.")) + if not principal.dataplanes.allows(dataplane_id): + raise ControlPlaneException(forbidden(f"Not bound to dataplane {dataplane_id!r}.")) + return dataplane diff --git a/litellm/proxy/control_plane/dataplanes/registry.py b/litellm/proxy/control_plane/dataplanes/registry.py new file mode 100644 index 00000000000..c7672367c82 --- /dev/null +++ b/litellm/proxy/control_plane/dataplanes/registry.py @@ -0,0 +1,101 @@ +"""Which dataplanes exist and how to reach each one. Phase 1 is config-backed and read-only.""" + +from collections.abc import Mapping +from dataclasses import dataclass +from types import MappingProxyType +from typing import Final, Protocol + +from litellm.litellm_core_utils.oauth_token_exchange import OAuthTokenExchangeConfig, OAuthTokenExchangeError +from litellm.proxy.control_plane.auth.rbac import Principal +from litellm.proxy.control_plane.dataplanes.credentials import ( + CredentialChain, + CredentialKind, + DataplaneCredential, + OidcTokenExchange, + ServiceTokenImpersonation, + TokenExchanger, +) +from litellm.proxy.control_plane.settings import DataplaneSpec + + +@dataclass(frozen=True, slots=True) +class Dataplane: + id: str + name: str + url: str + credential: DataplaneCredential + credential_kinds: tuple[CredentialKind, ...] + + +class DataplaneRegistry(Protocol): + def get(self, dataplane_id: str) -> Dataplane | None: ... + + def all(self) -> tuple[Dataplane, ...]: ... + + +@dataclass(frozen=True, slots=True) +class StaticRegistry: + entries: Mapping[str, Dataplane] + + def get(self, dataplane_id: str) -> Dataplane | None: + return self.entries.get(dataplane_id) + + def all(self) -> tuple[Dataplane, ...]: + return tuple(self.entries.values()) + + +class InvalidDataplaneConfig(ValueError): + def __init__(self, dataplane_id: str, reason: str) -> None: + super().__init__(f"dataplane {dataplane_id!r}: {reason}") + + +def visible_to(registry: DataplaneRegistry, principal: Principal) -> tuple[Dataplane, ...]: + return tuple(dataplane for dataplane in registry.all() if principal.dataplanes.allows(dataplane.id)) + + +def _strategies(dataplane_id: str, spec: DataplaneSpec, exchanger: TokenExchanger) -> tuple[DataplaneCredential, ...]: + oidc: Final = spec.oidc_exchange + service: Final = spec.service_token + if oidc is None and service is None: + raise InvalidDataplaneConfig(dataplane_id, "set oidc_exchange, service_token or both") + try: + exchange: Final = ( + () + if oidc is None + else ( + OidcTokenExchange( + OAuthTokenExchangeConfig( + token_endpoint=oidc.token_endpoint, + client_id=oidc.client_id, + client_secret=oidc.client_secret.get_secret_value(), + profile=oidc.profile, + scope=oidc.scope, + audience=oidc.audience, + ), + exchanger, + ), + ) + ) + except OAuthTokenExchangeError as error: + raise InvalidDataplaneConfig(dataplane_id, error.message) from error + impersonation: Final = () if service is None else (ServiceTokenImpersonation(service.token.get_secret_value()),) + return (*exchange, *impersonation) + + +def build_dataplane(dataplane_id: str, spec: DataplaneSpec, exchanger: TokenExchanger) -> Dataplane: + chain: Final = CredentialChain(_strategies(dataplane_id, spec, exchanger)) + return Dataplane( + id=dataplane_id, + name=spec.name or dataplane_id, + url=spec.url.rstrip("/"), + credential=chain, + credential_kinds=chain.kinds, + ) + + +def build_registry(specs: Mapping[str, DataplaneSpec], exchanger: TokenExchanger) -> StaticRegistry: + return StaticRegistry( + MappingProxyType( + {dataplane_id: build_dataplane(dataplane_id, spec, exchanger) for dataplane_id, spec in specs.items()} + ) + ) diff --git a/litellm/proxy/control_plane/dataplanes/router.py b/litellm/proxy/control_plane/dataplanes/router.py new file mode 100644 index 00000000000..9522ae056ba --- /dev/null +++ b/litellm/proxy/control_plane/dataplanes/router.py @@ -0,0 +1,130 @@ +from http import HTTPStatus +from typing import Annotated, Final + +from fastapi import APIRouter, Depends, Security +from pydantic import BaseModel, ConfigDict, ValidationError + +from litellm.proxy.control_plane.auth.dependencies import require_scopes +from litellm.proxy.control_plane.auth.rbac import Principal, Scope +from litellm.proxy.control_plane.common.exceptions import ControlPlaneException, bad_gateway +from litellm.proxy.control_plane.dataplanes.client import DataplaneClient, DataplaneReply, DataplaneUnreachable +from litellm.proxy.control_plane.dataplanes.credentials import Caller, CredentialUnavailable +from litellm.proxy.control_plane.dataplanes.dependencies import ( + dataplane_access, + get_caller, + get_dataplane_client, + get_registry, +) +from litellm.proxy.control_plane.dataplanes.registry import Dataplane, DataplaneRegistry, visible_to +from litellm.types.proxy.control_plane.dataplanes import ( + DataplaneHealthResponse, + DataplaneIdentityResponse, + DataplaneListResponse, + DataplaneResponse, +) + +router: Final = APIRouter(prefix="/dataplanes", tags=["dataplanes"]) + +LIVENESS_PATH: Final = "/health/liveliness" +USER_INFO_PATH: Final = "/user/info" + + +class _UserInfo(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + user_id: str | None = None + user_email: str | None = None + user_role: str | None = None + + +class _UserInfoBody(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + user_id: str | None = None + user_info: _UserInfo | None = None + + +def _to_response(dataplane: Dataplane) -> DataplaneResponse: + return DataplaneResponse( + id=dataplane.id, name=dataplane.name, url=dataplane.url, credentials=dataplane.credential_kinds + ) + + +@router.get("", response_model=DataplaneListResponse) +async def list_dataplanes( + principal: Annotated[Principal, Security(require_scopes, scopes=[Scope.DATAPLANES_READ.value])], + registry: Annotated[DataplaneRegistry, Depends(get_registry)], +) -> DataplaneListResponse: + return DataplaneListResponse(items=tuple(_to_response(dataplane) for dataplane in visible_to(registry, principal))) + + +@router.get("/{dataplane_id}", response_model=DataplaneResponse) +async def get_dataplane( + _principal: Annotated[Principal, Security(require_scopes, scopes=[Scope.DATAPLANES_READ.value])], + dataplane: Annotated[Dataplane, Depends(dataplane_access)], +) -> DataplaneResponse: + return _to_response(dataplane) + + +@router.get("/{dataplane_id}/health", response_model=DataplaneHealthResponse) +async def dataplane_health( + _principal: Annotated[Principal, Security(require_scopes, scopes=[Scope.DATAPLANES_READ.value])], + dataplane: Annotated[Dataplane, Depends(dataplane_access)], + caller: Annotated[Caller, Depends(get_caller)], + client: Annotated[DataplaneClient, Depends(get_dataplane_client)], +) -> DataplaneHealthResponse: + result: Final = await client.get(dataplane, LIVENESS_PATH, caller) + match result: + case DataplaneReply(status=status) if HTTPStatus.OK <= status < HTTPStatus.MULTIPLE_CHOICES: + return DataplaneHealthResponse(id=dataplane.id, status="ok", http_status=status, detail=None) + case DataplaneReply(status=status, body=body): + return DataplaneHealthResponse(id=dataplane.id, status="degraded", http_status=status, detail=str(body)) + case DataplaneUnreachable(detail=detail): + return DataplaneHealthResponse(id=dataplane.id, status="unreachable", http_status=None, detail=detail) + case CredentialUnavailable(reason=reason): + return DataplaneHealthResponse(id=dataplane.id, status="degraded", http_status=None, detail=reason) + + +@router.get("/{dataplane_id}/whoami", response_model=DataplaneIdentityResponse) +async def dataplane_whoami( + _principal: Annotated[Principal, Security(require_scopes, scopes=[Scope.DATAPLANES_READ.value])], + dataplane: Annotated[Dataplane, Depends(dataplane_access)], + caller: Annotated[Caller, Depends(get_caller)], + client: Annotated[DataplaneClient, Depends(get_dataplane_client)], +) -> DataplaneIdentityResponse: + """Ask the dataplane who it authenticated, proving the caller's identity survived the hop.""" + result: Final = await client.get(dataplane, USER_INFO_PATH, caller) + match result: + case DataplaneReply(status=status, body=body, credential=credential) if status < 400: + info: Final = _identity_body(body) + if info is None: + raise ControlPlaneException(bad_gateway("Dataplane returned a malformed /user/info body.")) + user: Final = info.user_info or _UserInfo() + return DataplaneIdentityResponse( + id=dataplane.id, + credential=credential, + user_id=info.user_id or user.user_id, + user_email=user.user_email, + user_role=user.user_role, + ) + case DataplaneReply(status=HTTPStatus.NOT_FOUND, credential=credential): + return DataplaneIdentityResponse( + id=dataplane.id, credential=credential, user_id=None, user_email=None, user_role=None + ) + case DataplaneReply(status=status, credential=credential): + raise ControlPlaneException( + bad_gateway(f"Dataplane rejected the {credential} credential with HTTP {status}.") + ) + case DataplaneUnreachable(detail=detail, attempts=attempts): + raise ControlPlaneException(bad_gateway(f"Dataplane unreachable after {attempts} attempt(s): {detail}")) + case CredentialUnavailable(reason=reason): + raise ControlPlaneException(bad_gateway(f"No usable credential for this dataplane: {reason}")) + + +def _identity_body(body: object) -> _UserInfoBody | None: + if not isinstance(body, dict): + return None + try: + return _UserInfoBody.model_validate(body) + except ValidationError: + return None diff --git a/litellm/proxy/control_plane/settings.py b/litellm/proxy/control_plane/settings.py index 0e7f9812644..660c8c5862d 100644 --- a/litellm/proxy/control_plane/settings.py +++ b/litellm/proxy/control_plane/settings.py @@ -2,11 +2,13 @@ from collections.abc import Mapping from dataclasses import dataclass -from typing import Final +from typing import Final, TypeAlias -from pydantic import BaseModel, ConfigDict, Field, SecretStr +import yaml +from pydantic import BaseModel, ConfigDict, Field, SecretStr, TypeAdapter, model_validator from pydantic_settings import BaseSettings, SettingsConfigDict +from litellm.litellm_core_utils.oauth_token_exchange import TokenExchangeProfile from litellm.types.proxy.control_plane.auth import PASSWORD_MAX_LENGTH, USERNAME_MAX_LENGTH DEFAULT_SESSION_TTL_SECONDS: Final = 12 * 60 * 60 @@ -36,6 +38,51 @@ class RoleSpec(BaseModel): inherits: tuple[str, ...] = () +class OidcExchangeSpec(BaseModel): + """IdP token-exchange client the control plane uses to mint a token for this dataplane's audience.""" + + model_config = ConfigDict(frozen=True, extra="forbid") + + token_endpoint: str + client_id: str + client_secret: SecretStr + profile: TokenExchangeProfile = "rfc8693" + audience: str | None = None + scope: str | None = None + + +class ServiceTokenSpec(BaseModel): + """A dataplane virtual key holding the `impersonate` permission, the fallback when no IdP exchange is possible.""" + + model_config = ConfigDict(frozen=True, extra="forbid") + + token: SecretStr + + +class DataplaneSpec(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + url: str = Field(pattern=r"^https?://") + name: str | None = None + oidc_exchange: OidcExchangeSpec | None = None + service_token: ServiceTokenSpec | None = None + allow_insecure_http: bool = False + + @model_validator(mode="after") + def _https_unless_opted_out(self) -> "DataplaneSpec": + if self.url.startswith("http://") and not self.allow_insecure_http: + raise ValueError( + f"dataplane url {self.url!r} is plain http, the bearer credential would travel unencrypted. " + "Use https or set allow_insecure_http: true for a local or in-cluster dataplane." + ) + return self + + +DataplaneSpecs: TypeAlias = Mapping[str, DataplaneSpec] +_DATAPLANE_SPECS: Final = TypeAdapter(DataplaneSpecs) +_YAML_DOCUMENT: Final = TypeAdapter(dict[str, object]) + + class ControlPlaneSettings(BaseSettings): model_config = SettingsConfigDict(frozen=True, env_prefix="CONTROL_PLANE_", extra="ignore") @@ -55,7 +102,10 @@ class ControlPlaneSettings(BaseSettings): oidc_dataplanes_claim: str = "litellm_dataplanes" oidc_email_claim: str = "email" oidc_leeway_seconds: int = 30 + dataplanes: DataplaneSpecs = Field(default_factory=dict) + dataplanes_file: str | None = None dataplane_timeout_seconds: float = DEFAULT_DATAPLANE_TIMEOUT_SECONDS + dataplane_retry_attempts: int = 3 fanout_concurrency: int = 16 api_prefix: str = "/api/v1" @@ -75,5 +125,15 @@ class ControlPlaneSettings(BaseSettings): ) +def read_dataplanes_file(path: str) -> DataplaneSpecs: + """A yaml file with a top-level `dataplanes:` map keyed by dataplane id, the same shape as the env JSON.""" + with open(path, encoding="utf-8") as handle: + document: Final = _YAML_DOCUMENT.validate_python(yaml.safe_load(handle)) + return _DATAPLANE_SPECS.validate_python(document.get("dataplanes", {})) + + def load_settings() -> ControlPlaneSettings: - return ControlPlaneSettings() # pyright: ignore[reportCallIssue] # pydantic-settings fills fields from env + from_env: Final = ControlPlaneSettings() # pyright: ignore[reportCallIssue] # pydantic-settings fills fields from env + if from_env.dataplanes_file is None: + return from_env + return from_env.model_copy(update={"dataplanes": read_dataplanes_file(from_env.dataplanes_file)}) diff --git a/litellm/types/llms/custom_http.py b/litellm/types/llms/custom_http.py index 47f80c52845..d1c1ec1c3f9 100644 --- a/litellm/types/llms/custom_http.py +++ b/litellm/types/llms/custom_http.py @@ -16,6 +16,7 @@ class httpxSpecialProvider(str, Enum): GuardrailCallback = "guardrail_callback" Caching = "caching" Oauth2Check = "oauth2_check" + ControlPlane = "control_plane" Oauth2Register = "oauth2_register" SecretManager = "secret_manager" PassThroughEndpoint = "pass_through_endpoint" diff --git a/litellm/types/proxy/auth/user_api_key_auth.py b/litellm/types/proxy/auth/user_api_key_auth.py index 18dd4e53701..d8a2fe79a69 100644 --- a/litellm/types/proxy/auth/user_api_key_auth.py +++ b/litellm/types/proxy/auth/user_api_key_auth.py @@ -171,6 +171,7 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response ob mcp_session_resource_server_id: str | None = Field(default=None, exclude=True) mcp_toolset_id: str | None = Field(default=None, exclude=True) authenticated_by_custom_auth: bool = Field(default=False, exclude=True) + impersonated_by: str | None = None via_virtual_key: bool = Field( default=False, exclude=True, diff --git a/litellm/types/proxy/control_plane/dataplanes.py b/litellm/types/proxy/control_plane/dataplanes.py new file mode 100644 index 00000000000..a62d103878d --- /dev/null +++ b/litellm/types/proxy/control_plane/dataplanes.py @@ -0,0 +1,44 @@ +from typing import Literal, TypeAlias + +from pydantic import ConfigDict + +from litellm.types.utils import LiteLLMBaseModel + +CredentialKindName: TypeAlias = Literal["oidc_exchange", "service_token"] +DataplaneHealthStatus: TypeAlias = Literal["ok", "unreachable", "degraded"] + + +class DataplaneResponse(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + id: str + name: str + url: str + credentials: tuple[CredentialKindName, ...] + + +class DataplaneListResponse(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + items: tuple[DataplaneResponse, ...] + + +class DataplaneHealthResponse(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + id: str + status: DataplaneHealthStatus + http_status: int | None + detail: str | None + + +class DataplaneIdentityResponse(LiteLLMBaseModel): + """Who the dataplane thinks the caller is, the end-to-end proof that identity survived the hop.""" + + model_config = ConfigDict(frozen=True) + + id: str + credential: CredentialKindName | None + user_id: str | None + user_email: str | None + user_role: str | None diff --git a/tests/unit/proxy/auth/test_impersonation.py b/tests/unit/proxy/auth/test_impersonation.py new file mode 100644 index 00000000000..b78f4d9acd8 --- /dev/null +++ b/tests/unit/proxy/auth/test_impersonation.py @@ -0,0 +1,117 @@ +import pytest + +from litellm.proxy._types import LiteLLM_UserTable, LitellmUserRoles +from litellm.proxy.auth.impersonation import ( + IMPERSONATE_EMAIL_HEADER, + IMPERSONATE_USER_HEADER, + ImpersonatedUserNotFound, + ImpersonationForbidden, + ImpersonationRequest, + impersonate, + read_impersonation_request, +) +from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth + +ALICE = LiteLLM_UserTable(user_id="alice", user_email="alice@x.test", user_role=LitellmUserRoles.PROXY_ADMIN.value) + + +def _service_key(permissions: dict[str, object]) -> UserAPIKeyAuth: + return UserAPIKeyAuth( + api_key="hashed", + token="hashed", + key_alias="control-plane", + user_id="svc", + team_id="platform", + permissions=permissions, + ) + + +async def _directory(user_id: str) -> LiteLLM_UserTable | None: + return ALICE if user_id == "alice" else None + + +def test_headers_without_a_user_are_not_an_impersonation_request() -> None: + assert read_impersonation_request({}) is None + assert read_impersonation_request({IMPERSONATE_USER_HEADER: " ", IMPERSONATE_EMAIL_HEADER: "a@x"}) is None + + +def test_headers_parse_user_and_optional_email() -> None: + assert read_impersonation_request({IMPERSONATE_USER_HEADER: "alice"}) == ImpersonationRequest("alice", None) + assert read_impersonation_request({IMPERSONATE_USER_HEADER: "alice", IMPERSONATE_EMAIL_HEADER: "a@x"}) == ( + ImpersonationRequest("alice", "a@x") + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("permissions", [{}, {"impersonate": False}, {"impersonate": "true"}]) +async def test_key_without_the_permission_is_forbidden(permissions: dict[str, object]) -> None: + result = await impersonate(_service_key(permissions), ImpersonationRequest("alice", None), _directory) + assert isinstance(result, ImpersonationForbidden) + assert "control-plane" in result.detail + + +@pytest.mark.asyncio +async def test_unknown_subject_is_not_found_not_created() -> None: + result = await impersonate(_service_key({"impersonate": True}), ImpersonationRequest("ghost", None), _directory) + assert result == ImpersonatedUserNotFound("ghost") + + +@pytest.mark.asyncio +async def test_permitted_key_becomes_the_subject_and_records_the_actor() -> None: + result = await impersonate(_service_key({"impersonate": True}), ImpersonationRequest("alice", None), _directory) + assert isinstance(result, UserAPIKeyAuth) + assert (result.user_id, result.user_email, result.user_role) == ( + "alice", + "alice@x.test", + LitellmUserRoles.PROXY_ADMIN, + ) + assert result.impersonated_by == "control-plane" + assert result.api_key == "hashed" + + +@pytest.mark.asyncio +async def test_subject_does_not_inherit_the_service_keys_team_models_or_permissions() -> None: + key = _service_key({"impersonate": True}).model_copy( + update={"models": ["gpt-4o"], "max_budget": 1000.0, "metadata": {"tier": "platform"}} + ) + result = await impersonate(key, ImpersonationRequest("alice", None), _directory) + assert isinstance(result, UserAPIKeyAuth) + assert result.model_dump(include={"team_id", "models", "max_budget", "permissions", "metadata"}) == { + "team_id": None, + "models": [], + "max_budget": None, + "permissions": {}, + "metadata": {}, + } + + +@pytest.mark.asyncio +async def test_email_header_naming_another_user_is_not_found() -> None: + mismatch = await impersonate( + _service_key({"impersonate": True}), ImpersonationRequest("alice", "bob@x.test"), _directory + ) + assert mismatch == ImpersonatedUserNotFound("alice") + same = await impersonate( + _service_key({"impersonate": True}), ImpersonationRequest("alice", "ALICE@x.test"), _directory + ) + assert isinstance(same, UserAPIKeyAuth) + + +@pytest.mark.asyncio +async def test_subject_budget_never_lands_on_the_shared_key_fields() -> None: + budgeted = ALICE.model_copy(update={"max_budget": 5.0, "spend": 4.0}) + + async def directory(user_id: str) -> LiteLLM_UserTable | None: + return budgeted + + result = await impersonate(_service_key({"impersonate": True}), ImpersonationRequest("alice", None), directory) + assert isinstance(result, UserAPIKeyAuth) + assert (result.max_budget, result.spend) == (None, 0.0) + + +@pytest.mark.asyncio +async def test_the_service_keys_route_allowlist_stays_a_ceiling_on_the_subject() -> None: + key = _service_key({"impersonate": True}).model_copy(update={"allowed_routes": ["/user/info", "/key/list"]}) + result = await impersonate(key, ImpersonationRequest("alice", None), _directory) + assert isinstance(result, UserAPIKeyAuth) + assert result.allowed_routes == ["/user/info", "/key/list"] diff --git a/tests/unit/proxy/auth/test_user_api_key_auth.py b/tests/unit/proxy/auth/test_user_api_key_auth.py index da8b25c2cee..e8057724a16 100644 --- a/tests/unit/proxy/auth/test_user_api_key_auth.py +++ b/tests/unit/proxy/auth/test_user_api_key_auth.py @@ -1283,6 +1283,7 @@ async def test_user_api_key_from_query_param(): "type": "http", "path": "/v1beta/models/gemini:streamGenerateContent", "query_string": f"alt=sse&key={user_key}".encode(), + "headers": [], } ) request._url = URL(url=f"/v1beta/models/gemini:streamGenerateContent?alt=sse&key={user_key}") diff --git a/tests/unit/proxy/control_plane/dataplanes/__init__.py b/tests/unit/proxy/control_plane/dataplanes/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/proxy/control_plane/dataplanes/conftest.py b/tests/unit/proxy/control_plane/dataplanes/conftest.py new file mode 100644 index 00000000000..76e2789c31b --- /dev/null +++ b/tests/unit/proxy/control_plane/dataplanes/conftest.py @@ -0,0 +1,39 @@ +from typing import Final + +import pytest + +from litellm.litellm_core_utils.oauth_token_exchange import OAuthTokenExchangeConfig +from litellm.proxy.control_plane.auth.rbac import ( + ADMIN_ROLE, + Principal, + RolePolicy, + TokenKind, + build_policy, + build_principal, +) +from litellm.proxy.control_plane.dataplanes.credentials import Caller + +POLICY: Final = build_policy({}, include_defaults=True) +assert isinstance(POLICY, RolePolicy) + + +def principal( + kind: TokenKind, subject: str = "alice", email: str | None = "alice@x.test", *dataplanes: str +) -> Principal: + return build_principal(POLICY, subject, email, frozenset({ADMIN_ROLE}), None, frozenset(dataplanes), kind) + + +def caller(kind: TokenKind, bearer: str = "cp-bearer", *dataplanes: str) -> Caller: + return Caller(principal=principal(kind, "alice", "alice@x.test", *dataplanes), bearer=bearer) + + +@pytest.fixture +def exchange_config() -> OAuthTokenExchangeConfig: + return OAuthTokenExchangeConfig( + token_endpoint="https://idp.test/token", + client_id="cp", + client_secret="s", + profile="rfc8693", + scope=None, + audience="dp", + ) diff --git a/tests/unit/proxy/control_plane/dataplanes/test_client.py b/tests/unit/proxy/control_plane/dataplanes/test_client.py new file mode 100644 index 00000000000..a2bb909a98f --- /dev/null +++ b/tests/unit/proxy/control_plane/dataplanes/test_client.py @@ -0,0 +1,133 @@ +from collections.abc import Iterator +from dataclasses import dataclass, field + +import httpx +import pytest + +from litellm.proxy.control_plane.dataplanes.client import ( + DataplaneClient, + DataplaneReply, + DataplaneUnreachable, + RetryPolicy, +) +from litellm.proxy.control_plane.dataplanes.credentials import ( + CredentialChain, + CredentialUnavailable, + ServiceTokenImpersonation, +) +from litellm.proxy.control_plane.dataplanes.registry import Dataplane +from tests.unit.proxy.control_plane.dataplanes.conftest import caller + + +@dataclass +class FakeTransport: # mutable-ok: records calls for assertions + outcomes: Iterator[httpx.Response | Exception] + calls: list[tuple[str, str, dict[str, str] | None]] = field(default_factory=list) + + async def request( + self, + method: str, + url: str, + *, + params: dict[str, str] | None = None, + json: dict[str, object] | None = None, + headers: dict[str, str] | None = None, + timeout: float, + ) -> httpx.Response: + self.calls.append((method, url, headers)) + outcome = next(self.outcomes) + if isinstance(outcome, Exception): + raise outcome + return outcome + + +@dataclass +class Clock: # mutable-ok: records sleeps for assertions + slept: list[float] = field(default_factory=list) + + async def sleep(self, seconds: float) -> None: + self.slept.append(seconds) + + +def _response(status: int, body: object = None) -> httpx.Response: + return httpx.Response(status, json=body, request=httpx.Request("GET", "https://eu.test/x")) + + +def _dataplane(token: str | None = "sk-svc") -> Dataplane: + chain = CredentialChain(() if token is None else (ServiceTokenImpersonation(token),)) + return Dataplane(id="eu", name="eu", url="https://eu.test", credential=chain, credential_kinds=chain.kinds) + + +def _client(transport: FakeTransport, clock: Clock) -> DataplaneClient: + return DataplaneClient( + transport, + timeout_seconds=1.0, + retry=RetryPolicy(attempts=3, base_delay_seconds=0.1), + sleep=clock.sleep, + jitter=lambda: 1.0, + ) + + +@pytest.mark.asyncio +async def test_get_injects_the_credential_and_parses_json() -> None: + transport, clock = FakeTransport(iter([_response(200, {"ok": True})])), Clock() + reply = await _client(transport, clock).get(_dataplane(), "/user/info", caller("session")) + assert reply == DataplaneReply(200, {"ok": True}, "service_token") + method, url, headers = transport.calls[0] + assert (method, url) == ("GET", "https://eu.test/user/info") + assert ( + headers is not None + and headers["Authorization"] == "Bearer sk-svc" + and headers["X-LiteLLM-Impersonate-User"] == "alice" + ) + assert clock.slept == [] + + +@pytest.mark.asyncio +async def test_get_retries_transient_failures_with_capped_exponential_backoff() -> None: + transport, clock = FakeTransport(iter([httpx.ConnectError("down"), _response(503), _response(200, "up")])), Clock() + reply = await _client(transport, clock).get(_dataplane(), "/health/liveliness", caller("session")) + assert reply == DataplaneReply(200, "up", "service_token") + assert clock.slept == [pytest.approx(0.1), pytest.approx(0.2)] + + +@pytest.mark.asyncio +async def test_get_gives_up_after_the_configured_attempts() -> None: + transport, clock = FakeTransport(iter([httpx.ConnectTimeout("t")] * 3)), Clock() + result = await _client(transport, clock).get(_dataplane(), "/x", caller("session")) + assert result == DataplaneUnreachable("ConnectTimeout: t", attempts=3) + assert len(transport.calls) == 3 + + +@pytest.mark.asyncio +async def test_post_is_never_retried_and_client_errors_are_replies_not_failures() -> None: + transport, clock = FakeTransport(iter([_response(502), _response(200)])), Clock() + result = await _client(transport, clock).post(_dataplane(), "/key/generate", caller("session"), json={"a": 1}) + assert result == DataplaneReply(502, "", "service_token") + assert len(transport.calls) == 1 and clock.slept == [] + forbidden = FakeTransport(iter([_response(403, {"error": "no"})])) + denied = await _client(forbidden, clock).get(_dataplane(), "/x", caller("session")) + assert denied == DataplaneReply(403, {"error": "no"}, "service_token") + assert len(forbidden.calls) == 1, "a 4xx is the dataplane's answer, not a reason to retry" + + +@pytest.mark.asyncio +async def test_missing_credential_short_circuits_before_any_request() -> None: + transport, clock = FakeTransport(iter([])), Clock() + result = await _client(transport, clock).get(_dataplane(token=None), "/x", caller("session")) + assert isinstance(result, CredentialUnavailable) and transport.calls == [] + + +def test_retry_delay_is_jittered_within_half_to_full_of_the_capped_backoff() -> None: + policy = RetryPolicy(attempts=5, base_delay_seconds=0.5, max_delay_seconds=2.0) + assert policy.delay(0, jitter=0.0) == pytest.approx(0.25) and policy.delay(0, jitter=1.0) == pytest.approx(0.5) + assert policy.delay(10, jitter=1.0) == pytest.approx(2.0) + + +@pytest.mark.asyncio +async def test_unreachable_detail_never_echoes_the_request_url() -> None: + leaky = httpx.ConnectError("failed for https://eu.test/key/list?token=sk-secret") + transport, clock = FakeTransport(iter([leaky])), Clock() + result = await _client(transport, clock).post(_dataplane(), "/key/list", caller("session"), json={}) + assert isinstance(result, DataplaneUnreachable) + assert result.detail == "ConnectError" and "sk-secret" not in result.detail diff --git a/tests/unit/proxy/control_plane/dataplanes/test_credentials.py b/tests/unit/proxy/control_plane/dataplanes/test_credentials.py new file mode 100644 index 00000000000..ff29c262d27 --- /dev/null +++ b/tests/unit/proxy/control_plane/dataplanes/test_credentials.py @@ -0,0 +1,104 @@ +from typing import cast + +import httpx +import pytest + +from litellm.exceptions import Timeout +from litellm.litellm_core_utils.oauth_token_exchange import OAuthTokenExchangeConfig, OAuthTokenExchangeError +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler +from litellm.proxy.auth.impersonation import IMPERSONATE_EMAIL_HEADER, IMPERSONATE_USER_HEADER +from litellm.proxy.control_plane.dataplanes.credentials import ( + CredentialChain, + CredentialUnavailable, + IssuedCredential, + OidcTokenExchange, + ServiceTokenImpersonation, + http_token_exchanger, +) +from tests.unit.proxy.control_plane.dataplanes.conftest import caller + + +async def _exchange(config: OAuthTokenExchangeConfig, subject_token: str) -> str: + return f"exchanged({subject_token})->{config.audience}" + + +async def _failing_exchange(config: OAuthTokenExchangeConfig, subject_token: str) -> str: + raise OAuthTokenExchangeError(status_code=401, message="nope") + + +@pytest.mark.asyncio +async def test_service_token_asserts_the_caller_in_impersonation_headers() -> None: + issued = await ServiceTokenImpersonation("sk-svc").headers(caller("session")) + assert issued == IssuedCredential( + "service_token", + {"Authorization": "Bearer sk-svc", IMPERSONATE_USER_HEADER: "alice", IMPERSONATE_EMAIL_HEADER: "alice@x.test"}, + ) + + +@pytest.mark.asyncio +async def test_master_key_caller_uses_the_service_token_as_itself() -> None: + issued = await ServiceTokenImpersonation("sk-svc").headers(caller("master_key")) + assert isinstance(issued, IssuedCredential) + assert dict(issued.headers) == {"Authorization": "Bearer sk-svc"} + + +@pytest.mark.asyncio +async def test_oidc_exchange_needs_an_idp_token(exchange_config: OAuthTokenExchangeConfig) -> None: + result = await OidcTokenExchange(exchange_config, _exchange).headers(caller("session")) + assert isinstance(result, CredentialUnavailable) and result.kind == "oidc_exchange" + assert "session" in result.reason + + +@pytest.mark.asyncio +async def test_oidc_exchange_trades_the_callers_bearer_for_a_dataplane_token( + exchange_config: OAuthTokenExchangeConfig, +) -> None: + issued = await OidcTokenExchange(exchange_config, _exchange).headers(caller("oidc", "idp-token")) + assert issued == IssuedCredential("oidc_exchange", {"Authorization": "Bearer exchanged(idp-token)->dp"}) + + +@pytest.mark.asyncio +async def test_failed_exchange_is_a_value_with_the_idp_status(exchange_config: OAuthTokenExchangeConfig) -> None: + result = await OidcTokenExchange(exchange_config, _failing_exchange).headers(caller("oidc")) + assert result == CredentialUnavailable("oidc_exchange", "token exchange failed with status 401") + + +@pytest.mark.asyncio +async def test_chain_falls_back_to_the_service_token_when_exchange_cannot_serve_the_caller( + exchange_config: OAuthTokenExchangeConfig, +) -> None: + chain = CredentialChain((OidcTokenExchange(exchange_config, _exchange), ServiceTokenImpersonation("sk-svc"))) + assert chain.kinds == ("oidc_exchange", "service_token") + via_session = await chain.headers(caller("session")) + via_oidc = await chain.headers(caller("oidc")) + assert isinstance(via_session, IssuedCredential) and via_session.kind == "service_token" + assert isinstance(via_oidc, IssuedCredential) and via_oidc.kind == "oidc_exchange" + + +@pytest.mark.asyncio +async def test_chain_reports_every_failure_when_nothing_can_authenticate( + exchange_config: OAuthTokenExchangeConfig, +) -> None: + result = await CredentialChain((OidcTokenExchange(exchange_config, _failing_exchange),)).headers(caller("oidc")) + assert isinstance(result, CredentialUnavailable) + assert result.reason == "oidc_exchange: token exchange failed with status 401" + assert isinstance(await CredentialChain(()).headers(caller("oidc")), CredentialUnavailable) + + +class _TimingOutHandler: + """Stands in for AsyncHTTPHandler, whose post raises litellm.Timeout instead of an httpx error.""" + + async def post(self, url: str, *, data: dict[str, str], timeout: float) -> httpx.Response: + raise Timeout("timed out", model=None, llm_provider="oidc") + + +@pytest.mark.asyncio +async def test_exchange_transport_timeout_is_a_credential_failure_not_a_crash( + exchange_config: OAuthTokenExchangeConfig, +) -> None: + exchanger = http_token_exchanger( + cast(AsyncHTTPHandler, _TimingOutHandler()), # cast-ok: the handler protocol slice the exchange uses + timeout=1.0, + ) + result = await OidcTokenExchange(exchange_config, exchanger).headers(caller("oidc")) + assert result == CredentialUnavailable("oidc_exchange", "token exchange failed with status 502") diff --git a/tests/unit/proxy/control_plane/dataplanes/test_registry.py b/tests/unit/proxy/control_plane/dataplanes/test_registry.py new file mode 100644 index 00000000000..13a2654f47d --- /dev/null +++ b/tests/unit/proxy/control_plane/dataplanes/test_registry.py @@ -0,0 +1,57 @@ +import pytest +from pydantic import ValidationError + +from litellm.litellm_core_utils.oauth_token_exchange import OAuthTokenExchangeConfig +from litellm.proxy.control_plane.dataplanes.registry import InvalidDataplaneConfig, build_registry, visible_to +from litellm.proxy.control_plane.settings import DataplaneSpec +from tests.unit.proxy.control_plane.dataplanes.conftest import principal + + +async def _exchange(config: OAuthTokenExchangeConfig, subject_token: str) -> str: + return subject_token + + +def _spec(**fields: object) -> DataplaneSpec: + return DataplaneSpec.model_validate({"url": "https://eu.litellm.test/", **fields}) + + +OIDC = {"token_endpoint": "https://idp.test/token", "client_id": "cp", "client_secret": "s", "audience": "dp"} + + +def test_registry_exposes_credential_order_and_normalizes_urls() -> None: + registry = build_registry( + {"eu": _spec(oidc_exchange=OIDC, service_token={"token": "sk"}), "us": _spec(service_token={"token": "sk"})}, + _exchange, + ) + eu = registry.get("eu") + assert eu is not None + assert (eu.name, eu.url, eu.credential_kinds) == ( + "eu", + "https://eu.litellm.test", + ("oidc_exchange", "service_token"), + ) + assert tuple(dataplane.id for dataplane in registry.all()) == ("eu", "us") + assert registry.get("apac") is None + + +def test_dataplane_without_any_credential_fails_at_startup() -> None: + with pytest.raises(InvalidDataplaneConfig, match="eu"): + build_registry({"eu": _spec()}, _exchange) + + +def test_plain_http_token_endpoint_fails_at_startup() -> None: + with pytest.raises(InvalidDataplaneConfig, match="HTTPS"): + build_registry({"eu": _spec(oidc_exchange={**OIDC, "token_endpoint": "http://idp.test/token"})}, _exchange) + + +def test_visible_dataplanes_follow_the_principals_binding() -> None: + registry = build_registry({name: _spec(service_token={"token": "sk"}) for name in ("eu", "us", "apac")}, _exchange) + assert tuple(d.id for d in visible_to(registry, principal("session"))) == ("eu", "us", "apac") + assert tuple(d.id for d in visible_to(registry, principal("session", "a", None, "us", "apac"))) == ("us", "apac") + + +def test_plain_http_dataplane_needs_an_explicit_opt_in() -> None: + with pytest.raises(ValidationError, match="unencrypted"): + DataplaneSpec.model_validate({"url": "http://dp.internal:4000"}) + assert DataplaneSpec.model_validate({"url": "http://dp.internal:4000", "allow_insecure_http": True}).url + assert DataplaneSpec.model_validate({"url": "https://dp.example.test"}).allow_insecure_http is False diff --git a/tests/unit/proxy/control_plane/dataplanes/test_router.py b/tests/unit/proxy/control_plane/dataplanes/test_router.py new file mode 100644 index 00000000000..7fbb68d4e9e --- /dev/null +++ b/tests/unit/proxy/control_plane/dataplanes/test_router.py @@ -0,0 +1,180 @@ +from collections.abc import Iterator +from dataclasses import dataclass, field +from datetime import timedelta +from pathlib import Path +from typing import Final + +import httpx +from fastapi.testclient import TestClient + +from litellm.litellm_core_utils.oauth_token_exchange import OAuthTokenExchangeConfig +from litellm.proxy.control_plane.app import create_app +from litellm.proxy.control_plane.auth.jwt import SessionTokenCodec +from litellm.proxy.control_plane.auth.rbac import USAGE_VIEWER_ROLE, VIEWER_ROLE, RolePolicy, build_principal +from litellm.proxy.control_plane.container import build_container +from litellm.proxy.control_plane.settings import ControlPlaneSettings, read_dataplanes_file + +MASTER_KEY: Final = "sk-cp-master" +DATAPLANES: Final = { + "eu": {"url": "https://eu.test", "service_token": {"token": "sk-eu"}}, + "us": {"url": "https://us.test", "name": "US East", "service_token": {"token": "sk-us"}}, +} + + +class Licensed: + def is_premium(self) -> bool: + return True + + +@dataclass +class FakeTransport: # mutable-ok: records calls for assertions + outcomes: Iterator[httpx.Response | Exception] + calls: list[tuple[str, dict[str, str] | None]] = field(default_factory=list) + + async def request( + self, + method: str, + url: str, + *, + params: dict[str, str] | None = None, + json: dict[str, object] | None = None, + headers: dict[str, str] | None = None, + timeout: float, + ) -> httpx.Response: + assert method == "GET", "router never writes" + self.calls.append((url, headers)) + outcome = next(self.outcomes) + if isinstance(outcome, Exception): + raise outcome + return outcome + + +async def _no_exchange(config: OAuthTokenExchangeConfig, subject_token: str) -> str: + raise AssertionError("no oidc dataplane configured") + + +def _response(status: int, body: object = None) -> httpx.Response: + return httpx.Response(status, json=body, request=httpx.Request("GET", "https://x")) + + +@dataclass(frozen=True) +class _Minter: + codec: SessionTokenCodec + policy: RolePolicy + + +def _rig(*outcomes: httpx.Response | Exception) -> tuple[TestClient, FakeTransport, _Minter]: + settings = ControlPlaneSettings.model_validate({"LITELLM_MASTER_KEY": MASTER_KEY, "dataplanes": DATAPLANES}) + transport = FakeTransport(iter(outcomes)) + container = build_container(settings, transport=transport, exchanger=_no_exchange) + codec = SessionTokenCodec(MASTER_KEY, settings.session_issuer, timedelta(hours=1), container.policy) + return TestClient(create_app(container, Licensed())), transport, _Minter(codec, container.policy) + + +def _viewer(minter: _Minter, *dataplanes: str, role: str = VIEWER_ROLE) -> dict[str, str]: + principal = build_principal( + minter.policy, "alice", "alice@x.test", frozenset({role}), None, frozenset(dataplanes), "session" + ) + return {"Authorization": f"Bearer {minter.codec.mint(principal)}"} + + +def test_list_shows_only_the_dataplanes_the_caller_is_bound_to() -> None: + client, _, codec = _rig() + fleet = client.get("/api/v1/dataplanes", headers={"Authorization": f"Bearer {MASTER_KEY}"}).json() + assert [item["id"] for item in fleet["items"]] == ["eu", "us"] + assert fleet["items"][1] == { + "id": "us", + "name": "US East", + "url": "https://us.test", + "credentials": ["service_token"], + } + bound = client.get("/api/v1/dataplanes", headers=_viewer(codec, "us")).json() + assert [item["id"] for item in bound["items"]] == ["us"] + + +def test_unbound_dataplane_is_403_and_unknown_is_404() -> None: + client, _, codec = _rig() + assert client.get("/api/v1/dataplanes/eu", headers=_viewer(codec, "us")).status_code == 403 + missing = client.get("/api/v1/dataplanes/apac", headers=_viewer(codec)) + assert missing.status_code == 404 and missing.json()["type"] == "urn:litellm:error:not-found" + + +def test_usage_viewer_lacks_the_dataplanes_read_scope() -> None: + client, _, codec = _rig() + denied = client.get("/api/v1/dataplanes", headers=_viewer(codec, role=USAGE_VIEWER_ROLE)) + assert denied.status_code == 403 and "dataplanes:read" in denied.json()["detail"] + + +def test_whoami_forwards_the_callers_identity_and_reports_what_the_dataplane_saw() -> None: + client, transport, codec = _rig( + _response(200, {"user_id": "alice", "user_info": {"user_email": "alice@x.test", "user_role": "proxy_admin"}}) + ) + body = client.get("/api/v1/dataplanes/us/whoami", headers=_viewer(codec, "us")).json() + assert body == { + "id": "us", + "credential": "service_token", + "user_id": "alice", + "user_email": "alice@x.test", + "user_role": "proxy_admin", + } + url, headers = transport.calls[0] + assert url == "https://us.test/user/info" + assert headers is not None + assert ( + headers["Authorization"], + headers["X-LiteLLM-Impersonate-User"], + headers["X-LiteLLM-Impersonate-Email"], + ) == ("Bearer sk-us", "alice", "alice@x.test") + + +def test_whoami_with_a_subjectless_caller_reports_the_credential_and_no_user() -> None: + client, transport, _ = _rig(_response(404, {"error": {"message": "User None not found"}})) + body = client.get("/api/v1/dataplanes/eu/whoami", headers={"Authorization": f"Bearer {MASTER_KEY}"}).json() + assert body == {"id": "eu", "credential": "service_token", "user_id": None, "user_email": None, "user_role": None} + _, headers = transport.calls[0] + assert headers is not None and "X-LiteLLM-Impersonate-User" not in headers + + +def test_dataplane_rejection_and_outage_surface_as_502_problems() -> None: + client, _, codec = _rig(_response(403, {"error": "lacks impersonate"}), *([httpx.ConnectError("down")] * 3)) + rejected = client.get("/api/v1/dataplanes/eu/whoami", headers=_viewer(codec)) + assert rejected.status_code == 502 and "HTTP 403" in rejected.json()["detail"] + down = client.get("/api/v1/dataplanes/eu/whoami", headers=_viewer(codec)) + assert down.status_code == 502 and "3 attempt" in down.json()["detail"] + + +def test_health_reports_ok_degraded_and_unreachable_without_failing_the_request() -> None: + client, _, codec = _rig( + _response(200, {"status": "healthy"}), _response(500, "warming"), *([httpx.ReadTimeout("slow")] * 3) + ) + headers = _viewer(codec) + assert client.get("/api/v1/dataplanes/eu/health", headers=headers).json()["status"] == "ok" + degraded = client.get("/api/v1/dataplanes/eu/health", headers=headers).json() + assert (degraded["status"], degraded["http_status"]) == ("degraded", 500) + unreachable = client.get("/api/v1/dataplanes/eu/health", headers=headers).json() + assert (unreachable["status"], unreachable["http_status"]) == ( + "unreachable", + None, + ) and "ReadTimeout" in unreachable["detail"] + + +def test_dataplanes_file_has_the_same_shape_as_the_env_json(tmp_path: Path) -> None: + path = tmp_path / "dataplanes.yaml" + path.write_text("dataplanes:\n eu:\n url: https://eu.test\n service_token:\n token: sk-eu\n") + specs = read_dataplanes_file(str(path)) + assert list(specs) == ["eu"] and specs["eu"].service_token is not None + assert specs["eu"].service_token.token.get_secret_value() == "sk-eu" + + +def test_a_redirect_or_not_modified_liveness_reply_is_degraded_not_ok() -> None: + client, _, codec = _rig(_response(304), _response(204)) + not_modified = client.get("/api/v1/dataplanes/eu/health", headers=_viewer(codec)).json() + assert (not_modified["status"], not_modified["http_status"]) == ("degraded", 304) + assert client.get("/api/v1/dataplanes/eu/health", headers=_viewer(codec)).json()["status"] == "ok" + + +def test_malformed_user_info_body_is_a_502_problem_not_a_crash() -> None: + client, _, codec = _rig(_response(200, {"user_info": []}), _response(200, ["alice"])) + for _ in range(2): + malformed = client.get("/api/v1/dataplanes/eu/whoami", headers=_viewer(codec)) + assert malformed.status_code == 502 and "malformed" in malformed.json()["detail"] diff --git a/tests/unit/proxy/control_plane/test_container.py b/tests/unit/proxy/control_plane/test_container.py new file mode 100644 index 00000000000..485e509ef4e --- /dev/null +++ b/tests/unit/proxy/control_plane/test_container.py @@ -0,0 +1,55 @@ +from dataclasses import dataclass, field +from typing import cast + +import httpx +import pytest + +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler +from litellm.proxy.control_plane.container import HandlerTransport + + +@dataclass +class _RecyclingHandler: # mutable-ok: counts how often the pool is re-read + """Stands in for AsyncHTTPHandler, whose client property hands out a fresh pool after cache eviction.""" + + pools: list[httpx.AsyncClient] = field(default_factory=list) + + @property + def client(self) -> httpx.AsyncClient: + pool = httpx.AsyncClient( + transport=httpx.MockTransport(lambda request: httpx.Response(200, json={"n": len(self.pools)})) + ) + self.pools.append(pool) + return pool + + +@pytest.mark.asyncio +async def test_transport_rereads_the_handlers_pool_on_every_request() -> None: + handler = _RecyclingHandler() + transport = HandlerTransport(cast(AsyncHTTPHandler, handler)) # cast-ok: only the client property is used + first = await transport.request("GET", "https://dp.test/health", timeout=1.0) + second = await transport.request("GET", "https://dp.test/health", timeout=1.0) + assert (first.json(), second.json()) == ({"n": 1}, {"n": 2}), "a pool closed by eviction must not be reused" + + +@dataclass +class _RedirectingHandler: # mutable-ok: counts the hops a following pool would take + hops: list[str] = field(default_factory=list) + + @property + def client(self) -> httpx.AsyncClient: + def route(request: httpx.Request) -> httpx.Response: + self.hops.append(str(request.url)) + if request.url.host == "dp.test": + return httpx.Response(302, headers={"Location": "http://127.0.0.1:9/secret"}) + return httpx.Response(200, json={"leaked": True}) + + return httpx.AsyncClient(transport=httpx.MockTransport(route), follow_redirects=True) + + +@pytest.mark.asyncio +async def test_transport_never_follows_a_dataplane_redirect_even_on_a_following_pool() -> None: + handler = _RedirectingHandler() + transport = HandlerTransport(cast(AsyncHTTPHandler, handler)) # cast-ok: only the client property is used + reply = await transport.request("GET", "https://dp.test/v1/mcp/server/oauth/x/authorize", timeout=1.0) + assert reply.status_code == 302 and handler.hops == ["https://dp.test/v1/mcp/server/oauth/x/authorize"] diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index adf020ca96a..fc5cf48a4b0 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -48446,6 +48446,8 @@ export interface components { end_user_tpm_limit?: number | null; /** Expires */ expires?: string | null; + /** Impersonated By */ + impersonated_by?: string | null; /** * Is Session Token * @default false