From ff8b9239dfb299c7f28d926ec39c8beb0c587fd0 Mon Sep 17 00:00:00 2001 From: yassin Date: Tue, 13 Oct 2026 20:43:35 +0000 Subject: [PATCH] feat(control_plane): dataplane registry, credential strategies and client Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/agent_endpoints/kill_switch.py | 33 +- litellm/proxy/auth/auth_checks.py | 2 +- litellm/proxy/auth/auth_utils.py | 11 +- litellm/proxy/auth/impersonation.py | 169 ++++++++ litellm/proxy/auth/user_api_key_auth.py | 194 ++++++++- litellm/proxy/control_plane/app.py | 5 + .../proxy/control_plane/common/exceptions.py | 4 + litellm/proxy/control_plane/container.py | 67 ++- .../control_plane/dataplanes/__init__.py | 0 .../proxy/control_plane/dataplanes/client.py | 170 ++++++++ .../control_plane/dataplanes/credentials.py | 133 ++++++ .../control_plane/dataplanes/dependencies.py | 41 ++ .../control_plane/dataplanes/registry.py | 99 +++++ .../proxy/control_plane/dataplanes/router.py | 129 ++++++ litellm/proxy/control_plane/settings.py | 73 +++- .../proxy/hooks/key_management_event_hooks.py | 111 ++--- .../hooks/user_management_event_hooks.py | 30 +- .../cache_settings_endpoints.py | 24 +- .../config_override_endpoints.py | 24 +- .../coordination_redis_endpoints.py | 24 +- .../internal_user_endpoints.py | 30 +- .../key_management_endpoints.py | 64 +-- .../password_endpoints.py | 2 +- .../tag_management_endpoints.py | 2 +- .../team_callback_endpoints.py | 24 +- .../management_endpoints/team_endpoints.py | 88 ++-- .../workflow_management_endpoints.py | 14 +- .../proxy/management_helpers/audit_logs.py | 51 ++- litellm/types/llms/custom_http.py | 1 + litellm/types/proxy/auth/user_api_key_auth.py | 2 + .../types/proxy/control_plane/dataplanes.py | 42 ++ tests/unit/proxy/auth/test_impersonation.py | 384 ++++++++++++++++++ .../unit/proxy/auth/test_user_api_key_auth.py | 173 ++++++-- .../control_plane/dataplanes/__init__.py | 0 .../control_plane/dataplanes/conftest.py | 39 ++ .../control_plane/dataplanes/test_client.py | 136 +++++++ .../dataplanes/test_credentials.py | 157 +++++++ .../control_plane/dataplanes/test_registry.py | 66 +++ .../control_plane/dataplanes/test_router.py | 199 +++++++++ .../proxy/control_plane/test_container.py | 63 +++ .../unit/proxy/control_plane/test_settings.py | 11 +- .../hooks/test_key_management_event_hooks.py | 91 ++++- .../test_key_management_endpoints.py | 29 ++ .../test_tag_management_endpoints.py | 69 ++-- .../test_workflow_management_endpoints.py | 152 +++---- .../test_audit_logs_proxy.py | 43 ++ 46 files changed, 2883 insertions(+), 392 deletions(-) create mode 100644 litellm/proxy/auth/impersonation.py create mode 100644 litellm/proxy/control_plane/dataplanes/__init__.py create mode 100644 litellm/proxy/control_plane/dataplanes/client.py create mode 100644 litellm/proxy/control_plane/dataplanes/credentials.py create mode 100644 litellm/proxy/control_plane/dataplanes/dependencies.py create mode 100644 litellm/proxy/control_plane/dataplanes/registry.py create mode 100644 litellm/proxy/control_plane/dataplanes/router.py create mode 100644 litellm/types/proxy/control_plane/dataplanes.py create mode 100644 tests/unit/proxy/auth/test_impersonation.py create mode 100644 tests/unit/proxy/control_plane/dataplanes/__init__.py create mode 100644 tests/unit/proxy/control_plane/dataplanes/conftest.py create mode 100644 tests/unit/proxy/control_plane/dataplanes/test_client.py create mode 100644 tests/unit/proxy/control_plane/dataplanes/test_credentials.py create mode 100644 tests/unit/proxy/control_plane/dataplanes/test_registry.py create mode 100644 tests/unit/proxy/control_plane/dataplanes/test_router.py create mode 100644 tests/unit/proxy/control_plane/test_container.py diff --git a/litellm/proxy/agent_endpoints/kill_switch.py b/litellm/proxy/agent_endpoints/kill_switch.py index 120c1bf9259..8bb8d16e155 100644 --- a/litellm/proxy/agent_endpoints/kill_switch.py +++ b/litellm/proxy/agent_endpoints/kill_switch.py @@ -19,7 +19,11 @@ from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, # pyright: ignore[reportUnknownVariableType] # its params arg is a bare dict in http_handler ) from litellm.proxy._types import LiteLLM_AuditLogs, LitellmTableNames, UserAPIKeyAuth -from litellm.proxy.management_helpers.audit_logs import create_audit_log_for_update, get_audit_log_changed_by +from litellm.proxy.management_helpers.audit_logs import ( + create_audit_log_for_update, + get_audit_log_changed_by, + with_audit_actor, +) from litellm.types.agents import ( AgentKillSwitchApiKeyAuth, AgentKillSwitchAuth, @@ -167,19 +171,22 @@ def build_kill_switch_audit_log( user_api_key_dict: UserAPIKeyAuth, litellm_proxy_admin_name: str | None, ) -> LiteLLM_AuditLogs: - return LiteLLM_AuditLogs( - id=str(uuid.uuid4()), - updated_at=datetime.now(timezone.utc), - changed_by=get_audit_log_changed_by( - litellm_changed_by=None, - user_api_key_dict=user_api_key_dict, - litellm_proxy_admin_name=litellm_proxy_admin_name, + return with_audit_actor( + LiteLLM_AuditLogs( + id=str(uuid.uuid4()), + updated_at=datetime.now(timezone.utc), + changed_by=get_audit_log_changed_by( + litellm_changed_by=None, + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name=litellm_proxy_admin_name, + ), + changed_by_api_key=user_api_key_dict.api_key, + table_name=LitellmTableNames.AGENT_TABLE_NAME, + object_id=result.agent_id, + action="kill_switch_fired", + updated_values=result.model_dump_json(exclude_none=True), ), - changed_by_api_key=user_api_key_dict.api_key, - table_name=LitellmTableNames.AGENT_TABLE_NAME, - object_id=result.agent_id, - action="kill_switch_fired", - updated_values=result.model_dump_json(exclude_none=True), + user_api_key_dict, ) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index c7b698b7f18..7b6093a07c4 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -1149,7 +1149,7 @@ def request_skips_budget_checks( async def common_checks( - request_body: dict, + request_body: dict[str, object], team_object: LiteLLM_TeamTable | None, user_object: LiteLLM_UserTable | None, end_user_object: LiteLLM_EndUserTable | None, diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 7d0ca935cde..c4c8716d2a5 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -2639,7 +2639,7 @@ def _cache_prediction_model_candidates( def _extract_model_candidates_from_request( - request_data: dict, + request_data: Mapping[str, object], route: str, request_headers: Mapping[str, object] | None = None, request_query_params: Mapping[str, object] | None = None, @@ -2687,10 +2687,11 @@ def _extract_model_candidates_from_request( _append_model_candidates(candidates, request_data.get("target_model_names")) if _route_matches_any_marker(route=route, markers=_MODEL_ROUTING_SESSION_MODEL_ROUTE_MARKERS): session: Final = request_data.get("session") - if isinstance(session, dict): + if isinstance(session, Mapping): _append_model_candidates(candidates, session.get("model")) - if uses_completion_model_sources and isinstance(request_data.get("completion"), dict): - _append_model_candidates(candidates, request_data["completion"].get("model")) + completion: Final = request_data.get("completion") + if uses_completion_model_sources and isinstance(completion, Mapping): + _append_model_candidates(candidates, completion.get("model")) if uses_model_routing_sources: if uses_header_or_query_model_sources: @@ -2760,7 +2761,7 @@ def request_dispatched_to_provider_pass_through(request: Request) -> bool: def get_model_from_request( - request_data: dict, + request_data: Mapping[str, object], route: str, request_headers: Mapping[str, object] | None = None, request_query_params: Mapping[str, object] | None = None, diff --git a/litellm/proxy/auth/impersonation.py b/litellm/proxy/auth/impersonation.py new file mode 100644 index 00000000000..132a181c083 --- /dev/null +++ b/litellm/proxy/auth/impersonation.py @@ -0,0 +1,169 @@ +"""Kubernetes-style impersonation for service keys: a key with the `impersonate` permission acts as a named user.""" + +from base64 import b64decode, b64encode +from collections.abc import Mapping +from dataclasses import dataclass +from types import MappingProxyType +from typing import Final, Protocol, TypeAlias + +from fastapi import HTTPException +from pydantic import BaseModel, ConfigDict, field_validator + +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_ENCODING_HEADER: Final = "X-LiteLLM-Impersonate-Encoding" +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] = {} + + @field_validator("permissions", mode="before") + @classmethod + def _default_permissions(cls, value: object) -> object: + return {} if value is None else value + + +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 + metadata: Mapping[str, object] | None = None + + @field_validator("models", mode="before") + @classmethod + def _default_models(cls, value: object) -> object: + return () if value is None else value + + @field_validator("model_max_budget", mode="before") + @classmethod + def _default_model_max_budget(cls, value: object) -> object: + return None if value is None else value + + @field_validator("metadata", mode="before") + @classmethod + def _default_metadata(cls, value: object) -> object: + return None if value is None else value + + +class UserLookup(Protocol): + async def __call__(self, user_id: str) -> LiteLLM_UserTable | None: ... + + +def impersonation_headers(user_id: str, email: str | None) -> Mapping[str, str]: + encoded: Final = any( + bool(value) and (not value.isascii() or not value.isprintable()) for value in (user_id, email or "") + ) + return MappingProxyType( + { + IMPERSONATE_USER_HEADER: b64encode(user_id.encode("utf-8")).decode("ascii") if encoded else user_id, + **( + {IMPERSONATE_EMAIL_HEADER: b64encode(email.encode("utf-8")).decode("ascii") if encoded else email} + if email is not None + else {} + ), + **({IMPERSONATE_ENCODING_HEADER: "base64-utf8"} if encoded else {}), + } + ) + + +def _read_identity(value: str, encoding: str | None) -> str: + if encoding is None: + return value + if encoding == "base64-utf8": + try: + return b64decode(value, validate=True).decode("utf-8") + except (ValueError, UnicodeDecodeError): + raise HTTPException(status_code=400, detail="Invalid impersonation identity encoding") from None + raise HTTPException(status_code=400, detail="Unsupported impersonation identity encoding") + + +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 + encoding: Final = headers.get(IMPERSONATE_ENCODING_HEADER) + return ImpersonationRequest(_read_identity(user_id, encoding), _read_identity(email, encoding) 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, + *, + actor_is_master_key: bool, +) -> UserAPIKeyAuth | ImpersonationError: + if not actor_is_master_key and actor.team_blocked: + return ImpersonationForbidden(f"team {actor.team_id!r} is blocked") + if not (actor_is_master_key or 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) + grants: Final = _SubjectGrants.model_validate(subject) + if grants.metadata is not None and grants.metadata.get("scim_active") is False: + return ImpersonationForbidden(f"user {request.user_id!r} is deactivated via SCIM") + return _as_subject(actor, subject, grants) + + +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, grants: _SubjectGrants) -> UserAPIKeyAuth: + """Only the key identity survives, while subject grants and budgets come from the user row.""" + return UserAPIKeyAuth( + api_key=None, + token=actor.token, + key_alias=actor.key_alias, + allowed_routes=None, + 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_max_budget=subject.max_budget, + user_spend=subject.spend, + user_model_max_budget=grants.model_max_budget, + parent_otel_span=actor.parent_otel_span, + ).model_copy(update={"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 02d0e8736e3..bec734ff2cb 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -22,6 +22,7 @@ import fastapi from fastapi import HTTPException, Request, WebSocket, status from fastapi.security.api_key import APIKeyHeader from pydantic import TypeAdapter +from starlette.datastructures import Headers from starlette.exceptions import WebSocketException import litellm @@ -128,6 +129,18 @@ 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.impersonation import ( + UserLookup as ImpersonationUserLookup, +) 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 @@ -214,6 +227,7 @@ from litellm.router_api import ( request_fallback_chain, ) 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: @@ -276,17 +290,19 @@ azure_apim_header: Final = APIKeyHeader( def _get_model_from_request_context( - request_data: dict, + request_data: Mapping[str, object], route: str, request: Request | None, llm_router: litellm.Router | None = None, team_id: str | None = None, ) -> str | list[str] | None: + request_headers: Final[Mapping[str, object]] = safe_get_request_headers(request=request) + request_query_params: Final[Mapping[str, object]] = safe_get_request_query_params(request=request) return get_model_from_request( request_data=request_data, route=route, - request_headers=safe_get_request_headers(request=request), - request_query_params=safe_get_request_query_params(request=request), + request_headers=request_headers, + request_query_params=request_query_params, llm_router=llm_router, request=request, team_id=team_id, @@ -2921,7 +2937,22 @@ async def user_api_key_auth_builder( if _end_user_object is not None: valid_token.end_user_object_permission = _end_user_object.object_permission - return valid_token + return await _apply_impersonation_if_requested( + actor=valid_token, + request=request, + request_data=request_data, + route=route, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + llm_router=llm_router, + model_max_budget_limiter=model_max_budget_limiter, + actor_is_master_key=( + secrets.compare_digest(valid_token.api_key or "", hash_token(master_key)) + or secrets.compare_digest(api_key.encode("utf-8"), master_key.encode("utf-8")) + ), + ) if ( valid_token is not None @@ -2997,7 +3028,19 @@ async def user_api_key_auth_builder( ) _user_api_key_obj.via_virtual_key = True - return _user_api_key_obj + return await _apply_impersonation_if_requested( + actor=_user_api_key_obj, + request=request, + request_data=request_data, + route=route, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + llm_router=llm_router, + model_max_budget_limiter=model_max_budget_limiter, + actor_is_master_key=True, + ) ## IF it's not a master key ## Route should not be in master_key_only_routes @@ -3504,7 +3547,146 @@ async def validate_resolved_virtual_key( # noqa: C901 # Preserve ordering of e ) virtual_key_auth_obj.via_virtual_key = True request.state.requires_key_model_check = True # rebind-ok: defer until central grants - return virtual_key_auth_obj + return await _apply_impersonation_if_requested( + actor=virtual_key_auth_obj, + request=request, + request_data=request_data, + route=route, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + llm_router=llm_router, + model_max_budget_limiter=model_max_budget_limiter, + actor_is_master_key=False, + ) + + +async def _apply_impersonation_if_requested( + actor: UserAPIKeyAuth, + request: Request, + request_data: Mapping[str, object], + route: str, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Span | None, + proxy_logging_obj: ProxyLogging, + llm_router: litellm.Router | None, + model_max_budget_limiter: _UserModelBudgetLimiter, + actor_is_master_key: bool, +) -> UserAPIKeyAuth: + impersonation: Final = read_impersonation_request( + Headers(scope=request.scope) if "headers" in request.scope else Headers() + ) + if impersonation is None: + return actor + + team: Final = ( + await get_team_object( + team_id=actor.team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + if actor.team_id is not None and actor.team_id != UI_TEAM_ID and not actor_is_master_key + else None + ) + + 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 + + return await _impersonate_or_raise( + actor=actor if team is None else actor.model_copy(update={"team_blocked": team.blocked}), + impersonation=impersonation, + lookup=lookup, + request=request, + request_data=request_data, + llm_router=llm_router, + model_max_budget_limiter=model_max_budget_limiter, + route=route, + actor_is_master_key=actor_is_master_key, + ) + + +async def _impersonate_or_raise( + actor: UserAPIKeyAuth, + impersonation: ImpersonationRequest, + lookup: ImpersonationUserLookup, + request: Request, + request_data: Mapping[str, object], + llm_router: litellm.Router | None, + model_max_budget_limiter: _UserModelBudgetLimiter, + route: str, + actor_is_master_key: bool = False, +) -> UserAPIKeyAuth: + if RouteChecks.is_llm_api_route(route=route): + raise ProxyException( + message=f"Impersonation is not allowed on LLM API route {route}", + type=ProxyErrorTypes.auth_error, + param=IMPERSONATE_USER_HEADER, + code=403, + ) + if not RouteChecks.is_virtual_key_allowed_to_call_route(route=route, valid_token=actor, request=request): + raise ProxyException( + message=f"key {actor_label(actor)} is not allowed to impersonate on route {route}", + type=ProxyErrorTypes.auth_error, + param=IMPERSONATE_USER_HEADER, + code=403, + ) + result: Final = await impersonate(actor, impersonation, lookup, actor_is_master_key=actor_is_master_key) + match result: + case UserAPIKeyAuth(): + if result.user_role != LitellmUserRoles.PROXY_ADMIN: + RouteChecks.non_proxy_admin_allowed_routes_check( + user_obj=await lookup(impersonation.user_id), + _user_role=result.user_role, + route=route, + request=request, + valid_token=result, + request_data=dict(request_data), + ) + model: Final = _get_model_from_request_context( + request_data=request_data, + route=route, + request=request, + llm_router=llm_router, + team_id=result.team_id, + ) + models: Final = _get_model_names_for_budget_checks(model=model) + if not request_skips_budget_checks( + route=route, model=model, llm_router=llm_router, request_body=request_data + ): + await _check_user_model_budget( + valid_token=result, + model_max_budget_limiter=model_max_budget_limiter, + models=models, + ) + verbose_proxy_logger.info( + "impersonation: actor=%s subject=%s route=%s", actor_label(actor), result.user_id, route + ) + return result.model_copy(update={"via_virtual_key": actor.via_virtual_key}) + 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 6aa52e2618f..d8e5fc093b3 100644 --- a/litellm/proxy/control_plane/app.py +++ b/litellm/proxy/control_plane/app.py @@ -15,6 +15,8 @@ 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 from litellm.proxy.middleware.request_size_limit_middleware import RequestSizeLimitMiddleware @@ -79,7 +81,10 @@ def create_app(container: Container, license_check: License) -> FastAPI: app.add_middleware(ControlPlaneRequestSizeLimitMiddleware) 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 d069f087537..097a3f74bd2 100644 --- a/litellm/proxy/control_plane/common/exceptions.py +++ b/litellm/proxy/control_plane/common/exceptions.py @@ -50,6 +50,10 @@ class ControlPlaneException(Exception): super().__init__(body.detail) +def bad_gateway(detail: str) -> ErrorBody: + return error_body("bad-gateway", "Bad Gateway", 502, detail) + + def _render(body: ErrorBody) -> JSONResponse: response: Final = error_response(body) if body.status == 401: diff --git a/litellm/proxy/control_plane/container.py b/litellm/proxy/control_plane/container.py index e1caa26f485..a45268e9d9c 100644 --- a/litellm/proxy/control_plane/container.py +++ b/litellm/proxy/control_plane/container.py @@ -1,12 +1,18 @@ """Long-lived objects built once per process and handed to the app. Tests build their own.""" import time -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from dataclasses import dataclass 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,54 @@ 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, "follow_redirects": False}, + ) + + +@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: Sequence[tuple[str, str]] | None = None, + 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: + httpx_params: Final = tuple(params) if params is not None else None + return await self.handler.client.request( + method, url, params=httpx_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 +137,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..c662fadffbf --- /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, Sequence +from dataclasses import dataclass +from typing import Final, Literal, Protocol, TypeAlias + +import httpx +from pydantic import TypeAdapter + +from litellm._logging import verbose_proxy_logger +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: Sequence[tuple[str, str]] | None = None, + 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: Sequence[tuple[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: Sequence[tuple[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: Sequence[tuple[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: Sequence[tuple[str, str]] | None, + json: Mapping[str, object] | None, + ) -> _RawReply | DataplaneUnreachable: + try: + response: Final = await self.transport.request( + method, + url, + params=params, + json=None if json is None else dict(json), + headers=dict(headers), + timeout=self.timeout_seconds, + ) + except httpx.HTTPError as error: + verbose_proxy_logger.debug("Control-plane dataplane transport failed: %s", error) + return DataplaneUnreachable(_describe(error), attempts=1) + return _RawReply(response.status_code, _parse_body(response)) + + +def _describe(error: httpx.HTTPError) -> str: + return type(error).__name__ + + +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..d4e251b0993 --- /dev/null +++ b/litellm/proxy/control_plane/dataplanes/credentials.py @@ -0,0 +1,133 @@ +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 impersonation_headers +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)) + return _issued( + "service_token", {**_bearer(self.token), **impersonation_headers(principal.subject, principal.email)} + ) + + +@dataclass(frozen=True, slots=True) +class CredentialChain: + 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..bf081c8c719 --- /dev/null +++ b/litellm/proxy/control_plane/dataplanes/dependencies.py @@ -0,0 +1,41 @@ +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..b5a78881600 --- /dev/null +++ b/litellm/proxy/control_plane/dataplanes/registry.py @@ -0,0 +1,99 @@ +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..92328058f2f --- /dev/null +++ b/litellm/proxy/control_plane/dataplanes/router.py @@ -0,0 +1,129 @@ +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: + result: Final = await client.get(dataplane, USER_INFO_PATH, caller) + match result: + case DataplaneReply(status=status, body=body, credential=credential) if 200 <= status < 300: + 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 6ff7929f0b6..90490960a68 100644 --- a/litellm/proxy/control_plane/settings.py +++ b/litellm/proxy/control_plane/settings.py @@ -1,12 +1,15 @@ """Process-level configuration for the control plane, read once at startup.""" +import os 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, field_validator, model_validator +import yaml +from pydantic import BaseModel, ConfigDict, Field, SecretStr, TypeAdapter, field_validator, model_validator from pydantic_settings import BaseSettings, SettingsConfigDict +from litellm.litellm_core_utils.oauth_token_exchange import TokenExchangeProfile from litellm.proxy.control_plane.auth.rbac import ADMIN_ROLE from litellm.types.proxy.control_plane.auth import PASSWORD_MAX_LENGTH, USERNAME_MAX_LENGTH @@ -37,6 +40,53 @@ 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 "@" in self.url.partition("://")[2].partition("/")[0]: + raise ValueError("dataplane url must not contain userinfo; credentials belong in a credential strategy.") + 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") @@ -56,7 +106,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" @@ -94,5 +147,19 @@ 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 + if os.environ.get("CONTROL_PLANE_DATAPLANES") is not None: + raise ValueError("CONTROL_PLANE_DATAPLANES and CONTROL_PLANE_DATAPLANES_FILE cannot both be set.") + return ControlPlaneSettings.model_validate( + {**from_env.model_dump(by_alias=True), "dataplanes": read_dataplanes_file(from_env.dataplanes_file)} + ) diff --git a/litellm/proxy/hooks/key_management_event_hooks.py b/litellm/proxy/hooks/key_management_event_hooks.py index 8784c02ac41..297786d9ea7 100644 --- a/litellm/proxy/hooks/key_management_event_hooks.py +++ b/litellm/proxy/hooks/key_management_event_hooks.py @@ -22,6 +22,7 @@ from litellm.proxy._types import ( UpdateKeyRequest, UserAPIKeyAuth, ) +from litellm.proxy.management_helpers.audit_logs import with_audit_actor from litellm.proxy.utils import ( # noqa: F401 # legacy module exports _hash_token_if_needed, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export hash_token_if_needed, @@ -66,20 +67,23 @@ class KeyManagementEventHooks: _updated_values: Final = response.model_dump_json(exclude_none=True) asyncio.create_task( create_audit_log_for_update( - request_data=LiteLLM_AuditLogs( - id=str(uuid.uuid4()), - updated_at=datetime.now(timezone.utc), - changed_by=get_audit_log_changed_by( - litellm_changed_by=litellm_changed_by, - user_api_key_dict=user_api_key_dict, - litellm_proxy_admin_name=litellm_proxy_admin_name, + request_data=with_audit_actor( + LiteLLM_AuditLogs( + id=str(uuid.uuid4()), + updated_at=datetime.now(timezone.utc), + changed_by=get_audit_log_changed_by( + litellm_changed_by=litellm_changed_by, + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name=litellm_proxy_admin_name, + ), + changed_by_api_key=user_api_key_dict.api_key, + table_name=LitellmTableNames.KEY_TABLE_NAME, + object_id=response.token_id or "", + action="created", + updated_values=_updated_values, + before_value=None, ), - changed_by_api_key=user_api_key_dict.api_key, - table_name=LitellmTableNames.KEY_TABLE_NAME, - object_id=response.token_id or "", - action="created", - updated_values=_updated_values, - before_value=None, + user_api_key_dict, ) ) ) @@ -128,20 +132,23 @@ class KeyManagementEventHooks: **data.model_dump(exclude_none=True), **({"project_id": data.project_id} if "project_id" in data.model_fields_set else {}), } - audit_log: Final = LiteLLM_AuditLogs( - id=str(uuid.uuid4()), - updated_at=datetime.now(timezone.utc), - changed_by=get_audit_log_changed_by( - litellm_changed_by=litellm_changed_by, - user_api_key_dict=user_api_key_dict, - litellm_proxy_admin_name=litellm_proxy_admin_name, + audit_log: Final = with_audit_actor( + LiteLLM_AuditLogs( + id=str(uuid.uuid4()), + updated_at=datetime.now(timezone.utc), + changed_by=get_audit_log_changed_by( + litellm_changed_by=litellm_changed_by, + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name=litellm_proxy_admin_name, + ), + changed_by_api_key=user_api_key_dict.api_key, + table_name=LitellmTableNames.KEY_TABLE_NAME, + object_id=hash_token_if_needed(data.key), + action="updated", + updated_values=json.dumps(updated_fields, default=str), + before_value=json.dumps(existing_key_row.json(exclude_none=True), default=str), ), - changed_by_api_key=user_api_key_dict.api_key, - table_name=LitellmTableNames.KEY_TABLE_NAME, - object_id=hash_token_if_needed(data.key), - action="updated", - updated_values=json.dumps(updated_fields, default=str), - before_value=json.dumps(existing_key_row.json(exclude_none=True), default=str), + user_api_key_dict, ) masked_values: Final = TypeAdapter(dict[str, object]).validate_json(str(audit_log.updated_values)) request_data: Final = ( @@ -180,20 +187,23 @@ class KeyManagementEventHooks: if is_audit_logging_enabled() and existing_key_row.token is not None: asyncio.create_task( create_audit_log_for_update( - request_data=LiteLLM_AuditLogs( - id=str(uuid.uuid4()), - updated_at=datetime.now(timezone.utc), - changed_by=get_audit_log_changed_by( - litellm_changed_by=litellm_changed_by, - user_api_key_dict=user_api_key_dict, - litellm_proxy_admin_name=litellm_proxy_admin_name, + request_data=with_audit_actor( + LiteLLM_AuditLogs( + id=str(uuid.uuid4()), + updated_at=datetime.now(timezone.utc), + changed_by=get_audit_log_changed_by( + litellm_changed_by=litellm_changed_by, + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name=litellm_proxy_admin_name, + ), + changed_by_api_key=user_api_key_dict.token, + table_name=LitellmTableNames.KEY_TABLE_NAME, + object_id=existing_key_row.token, + action="rotated", + updated_values=response.model_dump_json(exclude_none=True), + before_value=existing_key_row.model_dump_json(exclude_none=True), ), - changed_by_api_key=user_api_key_dict.token, - table_name=LitellmTableNames.KEY_TABLE_NAME, - object_id=existing_key_row.token, - action="rotated", - updated_values=response.model_dump_json(exclude_none=True), - before_value=existing_key_row.model_dump_json(exclude_none=True), + user_api_key_dict, ) ) ) @@ -277,16 +287,19 @@ class KeyManagementEventHooks: LiteLLM_VerificationToken.model_validate(key, from_attributes=True) for key in keys_being_deleted ) audit_rows: Final = tuple( - LiteLLM_AuditLogs( - id=str(uuid.uuid4()), - updated_at=datetime.now(timezone.utc), - changed_by=changed_by, - changed_by_api_key=user_api_key_dict.token, - table_name=LitellmTableNames.KEY_TABLE_NAME, - object_id=key_row.token, - action="deleted", - updated_values="{}", - before_value=key_row.model_dump_json(exclude_none=True), + with_audit_actor( + LiteLLM_AuditLogs( + id=str(uuid.uuid4()), + updated_at=datetime.now(timezone.utc), + changed_by=changed_by, + changed_by_api_key=user_api_key_dict.token, + table_name=LitellmTableNames.KEY_TABLE_NAME, + object_id=key_row.token, + action="deleted", + updated_values="{}", + before_value=key_row.model_dump_json(exclude_none=True), + ), + user_api_key_dict, ) for key_row in key_rows if key_row.token is not None diff --git a/litellm/proxy/hooks/user_management_event_hooks.py b/litellm/proxy/hooks/user_management_event_hooks.py index 6618af70c0a..b0b9643f932 100644 --- a/litellm/proxy/hooks/user_management_event_hooks.py +++ b/litellm/proxy/hooks/user_management_event_hooks.py @@ -24,6 +24,7 @@ from litellm.proxy.management_helpers.audit_logs import ( create_audit_log_for_update, is_audit_logging_enabled, track_audit_task, + with_audit_actor, ) from litellm.repositories.user_repository import UserRepository @@ -229,19 +230,22 @@ class UserManagementEventHooks: ) await create_audit_log_for_update( - request_data=LiteLLM_AuditLogs( - id=str(uuid.uuid4()), - updated_at=datetime.now(timezone.utc), - changed_by=get_audit_log_changed_by( - litellm_changed_by=litellm_changed_by, - user_api_key_dict=user_api_key_dict, - litellm_proxy_admin_name=litellm_proxy_admin_name, + request_data=with_audit_actor( + LiteLLM_AuditLogs( + id=str(uuid.uuid4()), + updated_at=datetime.now(timezone.utc), + changed_by=get_audit_log_changed_by( + litellm_changed_by=litellm_changed_by, + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name=litellm_proxy_admin_name, + ), + changed_by_api_key=user_api_key_dict.api_key, + table_name=LitellmTableNames.USER_TABLE_NAME, + object_id=user_id, + action=action, + updated_values=after_value, + before_value=before_value, ), - changed_by_api_key=user_api_key_dict.api_key, - table_name=LitellmTableNames.USER_TABLE_NAME, - object_id=user_id, - action=action, - updated_values=after_value, - before_value=before_value, + user_api_key_dict, ) ) diff --git a/litellm/proxy/management_endpoints/cache_settings_endpoints.py b/litellm/proxy/management_endpoints/cache_settings_endpoints.py index a09045829b8..f0c96c241ca 100644 --- a/litellm/proxy/management_endpoints/cache_settings_endpoints.py +++ b/litellm/proxy/management_endpoints/cache_settings_endpoints.py @@ -30,6 +30,7 @@ from litellm.proxy._types import ( ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.db.exception_handler import call_with_db_reconnect_retry +from litellm.proxy.management_helpers.audit_logs import with_audit_actor from litellm.repositories.table_repositories import CacheConfigRepository from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.management_endpoints import ( @@ -318,16 +319,19 @@ async def _emit_cache_settings_audit_log( task: Final = asyncio.create_task( create_audit_log_for_update( - request_data=LiteLLM_AuditLogs( - id=str(uuid.uuid4()), - updated_at=datetime.now(timezone.utc), - changed_by=litellm_changed_by or user_api_key_dict.user_id or litellm_proxy_admin_name, - changed_by_api_key=user_api_key_dict.api_key, - table_name=LitellmTableNames.CACHE_CONFIG_TABLE_NAME, - object_id="cache_config", - action=action, - updated_values=json.dumps({"settings": _redact_settings(after_settings)}, default=str), - before_value=json.dumps({"settings": _redact_settings(before_settings)}, default=str), + request_data=with_audit_actor( + LiteLLM_AuditLogs( + id=str(uuid.uuid4()), + updated_at=datetime.now(timezone.utc), + changed_by=litellm_changed_by or user_api_key_dict.user_id or litellm_proxy_admin_name, + changed_by_api_key=user_api_key_dict.api_key, + table_name=LitellmTableNames.CACHE_CONFIG_TABLE_NAME, + object_id="cache_config", + action=action, + updated_values=json.dumps({"settings": _redact_settings(after_settings)}, default=str), + before_value=json.dumps({"settings": _redact_settings(before_settings)}, default=str), + ), + user_api_key_dict, ) ) ) diff --git a/litellm/proxy/management_endpoints/config_override_endpoints.py b/litellm/proxy/management_endpoints/config_override_endpoints.py index 7bc6b63ccc2..5e699481ce9 100644 --- a/litellm/proxy/management_endpoints/config_override_endpoints.py +++ b/litellm/proxy/management_endpoints/config_override_endpoints.py @@ -13,6 +13,7 @@ from typing_extensions import ReadOnly, TypedDict from litellm._uuid import uuid from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.safe_json_loads import safe_json_loads +from litellm.proxy.management_helpers.audit_logs import with_audit_actor try: from prisma.errors import RecordNotFoundError @@ -116,16 +117,19 @@ async def _emit_config_override_audit_log( task: Final = asyncio.create_task( create_audit_log_for_update( - request_data=LiteLLM_AuditLogs( - id=str(uuid.uuid4()), - updated_at=datetime.now(timezone.utc), - changed_by=litellm_changed_by or user_api_key_dict.user_id or litellm_proxy_admin_name, - changed_by_api_key=user_api_key_dict.api_key, - table_name=LitellmTableNames.CONFIG_OVERRIDES_TABLE_NAME, - object_id=object_id, - action=action, - updated_values=json.dumps({"config": _redact_config(after_config)}, default=str), - before_value=json.dumps({"config": _redact_config(before_config)}, default=str), + request_data=with_audit_actor( + LiteLLM_AuditLogs( + id=str(uuid.uuid4()), + updated_at=datetime.now(timezone.utc), + changed_by=litellm_changed_by or user_api_key_dict.user_id or litellm_proxy_admin_name, + changed_by_api_key=user_api_key_dict.api_key, + table_name=LitellmTableNames.CONFIG_OVERRIDES_TABLE_NAME, + object_id=object_id, + action=action, + updated_values=json.dumps({"config": _redact_config(after_config)}, default=str), + before_value=json.dumps({"config": _redact_config(before_config)}, default=str), + ), + user_api_key_dict, ) ) ) diff --git a/litellm/proxy/management_endpoints/coordination_redis_endpoints.py b/litellm/proxy/management_endpoints/coordination_redis_endpoints.py index 663e22c4312..8a844bdba29 100644 --- a/litellm/proxy/management_endpoints/coordination_redis_endpoints.py +++ b/litellm/proxy/management_endpoints/coordination_redis_endpoints.py @@ -36,6 +36,7 @@ from litellm.proxy._types import ( user_api_key_has_admin_view, ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.management_helpers.audit_logs import with_audit_actor from litellm.proxy.utils import invalidate_config_param from litellm.repositories.config_repository import ConfigRepository from litellm.secret_managers.main import get_secret_str @@ -255,16 +256,19 @@ async def _emit_coordination_redis_audit_log( task: Final = asyncio.create_task( create_audit_log_for_update( - request_data=LiteLLM_AuditLogs( - id=str(uuid.uuid4()), - updated_at=datetime.now(timezone.utc), - changed_by=litellm_changed_by or user_api_key_dict.user_id or litellm_proxy_admin_name, - changed_by_api_key=user_api_key_dict.api_key, - table_name=LitellmTableNames.CONFIG_TABLE_NAME, - object_id=_COORDINATION_REDIS_KEY, - action=action, - updated_values=json.dumps({"settings": _redact_all_values(after_settings)}, default=str), - before_value=json.dumps({"settings": _redact_all_values(before_settings)}, default=str), + request_data=with_audit_actor( + LiteLLM_AuditLogs( + id=str(uuid.uuid4()), + updated_at=datetime.now(timezone.utc), + changed_by=litellm_changed_by or user_api_key_dict.user_id or litellm_proxy_admin_name, + changed_by_api_key=user_api_key_dict.api_key, + table_name=LitellmTableNames.CONFIG_TABLE_NAME, + object_id=_COORDINATION_REDIS_KEY, + action=action, + updated_values=json.dumps({"settings": _redact_all_values(after_settings)}, default=str), + before_value=json.dumps({"settings": _redact_all_values(before_settings)}, default=str), + ), + user_api_key_dict, ) ) ) diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index f973fadb53c..f1c01587425 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -83,6 +83,7 @@ from litellm.proxy.management_endpoints.key_management_endpoints import ( # noq prepare_metadata_fields, raise_on_logging_credential_change, ) +from litellm.proxy.management_helpers.audit_logs import with_audit_actor from litellm.proxy.management_helpers.bulk_user_deletion import delete_users_from_current_teams from litellm.proxy.management_helpers.object_permission_utils import ( # noqa: F401 # legacy module exports _set_object_permission, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export @@ -2626,20 +2627,23 @@ async def delete_user( asyncio.create_task( create_audit_log_for_update( - request_data=LiteLLM_AuditLogs( - id=str(uuid.uuid4()), - updated_at=datetime.now(timezone.utc), - changed_by=get_audit_log_changed_by( - litellm_changed_by=litellm_changed_by, - user_api_key_dict=user_api_key_dict, - litellm_proxy_admin_name=litellm_proxy_admin_name, + request_data=with_audit_actor( + LiteLLM_AuditLogs( + id=str(uuid.uuid4()), + updated_at=datetime.now(timezone.utc), + changed_by=get_audit_log_changed_by( + litellm_changed_by=litellm_changed_by, + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name=litellm_proxy_admin_name, + ), + changed_by_api_key=user_api_key_dict.api_key, + table_name=LitellmTableNames.USER_TABLE_NAME, + object_id=user_id, + action="deleted", + updated_values="{}", + before_value=_user_row, ), - changed_by_api_key=user_api_key_dict.api_key, - table_name=LitellmTableNames.USER_TABLE_NAME, - object_id=user_id, - action="deleted", - updated_values="{}", - before_value=_user_row, + user_api_key_dict, ) ) ) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 3bd17c4b5d8..6b683d87526 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -137,6 +137,7 @@ from litellm.proxy.management_helpers.access_group_key_sync import ( sync_key_regeneration_access_group_membership, sync_key_update_access_group_membership, ) +from litellm.proxy.management_helpers.audit_logs import with_audit_actor from litellm.proxy.management_helpers.key_settings_audit import with_settings_updated_at from litellm.proxy.management_helpers.managed_resource_owner_sync import repoint_key_owned_managed_resources from litellm.proxy.management_helpers.object_permission_utils import ( # noqa: F401 # legacy module exports @@ -5175,12 +5176,13 @@ async def info_key_fn( ) try: + if user_api_key_dict.impersonated_by is not None and key is None: + raise HTTPException(status_code=400, detail="A key must be specified when impersonating a user.") if prisma_client is None: raise Exception( "Database not connected. Connect a database to your proxy - https://docs.litellm.ai/docs/simple_proxy#managing-auth---virtual-keys" ) - # default to using Auth token if no key is passed in key = key or user_api_key_dict.api_key if not key: raise ProxyException( @@ -8207,20 +8209,23 @@ async def block_key( if is_audit_logging_enabled(): asyncio.create_task( create_audit_log_for_update( - request_data=LiteLLM_AuditLogs( - id=str(uuid.uuid4()), - updated_at=datetime.now(timezone.utc), - changed_by=get_audit_log_changed_by( - litellm_changed_by=litellm_changed_by, - user_api_key_dict=user_api_key_dict, - litellm_proxy_admin_name=litellm_proxy_admin_name, + request_data=with_audit_actor( + LiteLLM_AuditLogs( + id=str(uuid.uuid4()), + updated_at=datetime.now(timezone.utc), + changed_by=get_audit_log_changed_by( + litellm_changed_by=litellm_changed_by, + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name=litellm_proxy_admin_name, + ), + changed_by_api_key=user_api_key_dict.api_key, + table_name=LitellmTableNames.KEY_TABLE_NAME, + object_id=hashed_token, + action="blocked", + updated_values="{}", + before_value=existing_record.model_dump_json(), ), - changed_by_api_key=user_api_key_dict.api_key, - table_name=LitellmTableNames.KEY_TABLE_NAME, - object_id=hashed_token, - action="blocked", - updated_values="{}", - before_value=existing_record.model_dump_json(), + user_api_key_dict, ) ) ) @@ -8321,20 +8326,23 @@ async def unblock_key( if is_audit_logging_enabled(): asyncio.create_task( create_audit_log_for_update( - request_data=LiteLLM_AuditLogs( - id=str(uuid.uuid4()), - updated_at=datetime.now(timezone.utc), - changed_by=get_audit_log_changed_by( - litellm_changed_by=litellm_changed_by, - user_api_key_dict=user_api_key_dict, - litellm_proxy_admin_name=litellm_proxy_admin_name, + request_data=with_audit_actor( + LiteLLM_AuditLogs( + id=str(uuid.uuid4()), + updated_at=datetime.now(timezone.utc), + changed_by=get_audit_log_changed_by( + litellm_changed_by=litellm_changed_by, + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name=litellm_proxy_admin_name, + ), + changed_by_api_key=user_api_key_dict.api_key, + table_name=LitellmTableNames.KEY_TABLE_NAME, + object_id=hashed_token, + action="unblocked", + updated_values="{}", + before_value=existing_record.model_dump_json(), ), - changed_by_api_key=user_api_key_dict.api_key, - table_name=LitellmTableNames.KEY_TABLE_NAME, - object_id=hashed_token, - action="unblocked", - updated_values="{}", - before_value=existing_record.model_dump_json(), + user_api_key_dict, ) ) ) @@ -8452,7 +8460,7 @@ async def _can_user_query_key_info( user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value or user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value ) - or user_api_key_dict.api_key == key + or (user_api_key_dict.impersonated_by is None and user_api_key_dict.api_key == key) or key_info.user_id == user_api_key_dict.user_id or await TeamMemberPermissionChecks.user_belongs_to_keys_team( user_api_key_dict=user_api_key_dict, diff --git a/litellm/proxy/management_endpoints/password_endpoints.py b/litellm/proxy/management_endpoints/password_endpoints.py index 21172e30639..64dc249bb2a 100644 --- a/litellm/proxy/management_endpoints/password_endpoints.py +++ b/litellm/proxy/management_endpoints/password_endpoints.py @@ -154,7 +154,7 @@ async def change_password( await revoke_ui_session_keys( user_id=user_id, user_api_key_dict=user_api_key_dict, - keep_hashed_token=user_api_key_dict.token, + keep_hashed_token=None if user_api_key_dict.impersonated_by is not None else user_api_key_dict.token, ) verbose_proxy_logger.info("Password changed via /user/password/change for user_id=%s", user_id) diff --git a/litellm/proxy/management_endpoints/tag_management_endpoints.py b/litellm/proxy/management_endpoints/tag_management_endpoints.py index 0b7a7ca43cd..f04750d6871 100644 --- a/litellm/proxy/management_endpoints/tag_management_endpoints.py +++ b/litellm/proxy/management_endpoints/tag_management_endpoints.py @@ -168,7 +168,7 @@ async def _get_internal_user_api_keys( return [] user_api_keys: Final = set() - if user_api_key_dict.api_key: + if user_api_key_dict.api_key and user_api_key_dict.impersonated_by is None: user_api_keys.add(user_api_key_dict.api_key) user_id: Final = user_api_key_dict.user_id diff --git a/litellm/proxy/management_endpoints/team_callback_endpoints.py b/litellm/proxy/management_endpoints/team_callback_endpoints.py index 24675406567..206396b8416 100644 --- a/litellm/proxy/management_endpoints/team_callback_endpoints.py +++ b/litellm/proxy/management_endpoints/team_callback_endpoints.py @@ -58,6 +58,7 @@ from litellm.proxy.management_endpoints.team_endpoints import ( # noqa: F401 # _refresh_cached_team, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export refresh_cached_team, ) +from litellm.proxy.management_helpers.audit_logs import with_audit_actor from litellm.proxy.management_helpers.utils import management_endpoint_wrapper from litellm.repositories.team_repository import TeamRepository from litellm.types.utils import OTEL_DESTINATION_VARS @@ -228,16 +229,19 @@ async def _emit_team_callback_audit_log( task: Final = asyncio.create_task( create_audit_log_for_update( - request_data=LiteLLM_AuditLogs( - id=str(uuid.uuid4()), - updated_at=datetime.now(timezone.utc), - changed_by=litellm_changed_by or user_api_key_dict.user_id or litellm_proxy_admin_name, - changed_by_api_key=user_api_key_dict.api_key, - table_name=LitellmTableNames.TEAM_TABLE_NAME, - object_id=team_id, - action="updated", - updated_values=json.dumps({"metadata": redacted_after}, default=str), - before_value=json.dumps({"metadata": redacted_before}, default=str), + request_data=with_audit_actor( + LiteLLM_AuditLogs( + id=str(uuid.uuid4()), + updated_at=datetime.now(timezone.utc), + changed_by=litellm_changed_by or user_api_key_dict.user_id or litellm_proxy_admin_name, + changed_by_api_key=user_api_key_dict.api_key, + table_name=LitellmTableNames.TEAM_TABLE_NAME, + object_id=team_id, + action="updated", + updated_values=json.dumps({"metadata": redacted_after}, default=str), + before_value=json.dumps({"metadata": redacted_before}, default=str), + ), + user_api_key_dict, ) ) ) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index f3d3c01b3a4..fbae739bb73 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -208,6 +208,7 @@ from litellm.proxy.management_helpers.access_group_team_sync import ( reconcile_team_access_group_membership, sync_team_access_group_membership, ) +from litellm.proxy.management_helpers.audit_logs import with_audit_actor from litellm.proxy.management_helpers.object_permission_utils import ( # noqa: F401 # legacy module exports _set_object_permission, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export enforce_all_proxy_mcp_servers_grant_is_admin_only, @@ -2087,20 +2088,23 @@ async def new_team( asyncio.create_task( create_audit_log_for_update( - request_data=LiteLLM_AuditLogs( - id=str(uuid.uuid4()), - updated_at=datetime.now(timezone.utc), - changed_by=get_audit_log_changed_by( - litellm_changed_by=litellm_changed_by, - user_api_key_dict=user_api_key_dict, - litellm_proxy_admin_name=litellm_proxy_admin_name, + request_data=with_audit_actor( + LiteLLM_AuditLogs( + id=str(uuid.uuid4()), + updated_at=datetime.now(timezone.utc), + changed_by=get_audit_log_changed_by( + litellm_changed_by=litellm_changed_by, + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name=litellm_proxy_admin_name, + ), + changed_by_api_key=user_api_key_dict.api_key, + table_name=LitellmTableNames.TEAM_TABLE_NAME, + object_id=data.team_id, + action="created", + updated_values=_updated_values, + before_value=None, ), - changed_by_api_key=user_api_key_dict.api_key, - table_name=LitellmTableNames.TEAM_TABLE_NAME, - object_id=data.team_id, - action="created", - updated_values=_updated_values, - before_value=None, + user_api_key_dict, ) ) ) @@ -2143,20 +2147,23 @@ async def _create_team_update_audit_log( asyncio.create_task( create_audit_log_for_update( - request_data=LiteLLM_AuditLogs( - id=str(uuid.uuid4()), - updated_at=datetime.now(timezone.utc), - changed_by=get_audit_log_changed_by( - litellm_changed_by=litellm_changed_by, - user_api_key_dict=user_api_key_dict, - litellm_proxy_admin_name=litellm_proxy_admin_name, + request_data=with_audit_actor( + LiteLLM_AuditLogs( + id=str(uuid.uuid4()), + updated_at=datetime.now(timezone.utc), + changed_by=get_audit_log_changed_by( + litellm_changed_by=litellm_changed_by, + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name=litellm_proxy_admin_name, + ), + changed_by_api_key=user_api_key_dict.api_key, + table_name=LitellmTableNames.TEAM_TABLE_NAME, + object_id=team_id, + action="updated", + updated_values=_after_value, + before_value=_before_value, ), - changed_by_api_key=user_api_key_dict.api_key, - table_name=LitellmTableNames.TEAM_TABLE_NAME, - object_id=team_id, - action="updated", - updated_values=_after_value, - before_value=_before_value, + user_api_key_dict, ) ) ) @@ -5056,20 +5063,23 @@ def _schedule_deleted_team_audit_rows( for team in teams: asyncio.create_task( create_audit_log_for_update( - request_data=LiteLLM_AuditLogs( - id=str(uuid.uuid4()), - updated_at=datetime.now(timezone.utc), - changed_by=get_audit_log_changed_by( - litellm_changed_by=litellm_changed_by, - user_api_key_dict=user_api_key_dict, - litellm_proxy_admin_name=litellm_proxy_admin_name, + request_data=with_audit_actor( + LiteLLM_AuditLogs( + id=str(uuid.uuid4()), + updated_at=datetime.now(timezone.utc), + changed_by=get_audit_log_changed_by( + litellm_changed_by=litellm_changed_by, + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name=litellm_proxy_admin_name, + ), + changed_by_api_key=user_api_key_dict.api_key, + table_name=LitellmTableNames.TEAM_TABLE_NAME, + object_id=team.team_id, + action="deleted", + updated_values="{}", + before_value=team.model_dump_json(exclude_none=True), ), - changed_by_api_key=user_api_key_dict.api_key, - table_name=LitellmTableNames.TEAM_TABLE_NAME, - object_id=team.team_id, - action="deleted", - updated_values="{}", - before_value=team.model_dump_json(exclude_none=True), + user_api_key_dict, ) ) ) diff --git a/litellm/proxy/management_endpoints/workflow_management_endpoints.py b/litellm/proxy/management_endpoints/workflow_management_endpoints.py index 649941bc233..431cb4f4859 100644 --- a/litellm/proxy/management_endpoints/workflow_management_endpoints.py +++ b/litellm/proxy/management_endpoints/workflow_management_endpoints.py @@ -57,8 +57,8 @@ def _read_scope_caller(user_api_key_dict: UserAPIKeyAuth) -> UserAPIKeyAuth | No def _caller_key(user_api_key_dict: UserAPIKeyAuth) -> str | None: - """Return the hashed key token that identifies this caller, or None for master key.""" - return user_api_key_dict.token + """Return the hashed key token when it identifies the caller's own scope.""" + return None if user_api_key_dict.impersonated_by is not None else user_api_key_dict.token # Status transitions driven by event_type @@ -237,10 +237,13 @@ async def create_workflow_run( if prisma_client is None: raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value) + caller: Final = _caller_key(user_api_key_dict) + if caller is None: + raise HTTPException(status_code=403, detail="Impersonated callers cannot create workflow runs.") try: create_data: Final[_RunCreateData] = { "workflow_type": data.workflow_type, - "created_by": _caller_key(user_api_key_dict), + "created_by": caller, } if data.input is not None: create_data["input"] = _json(data.input) @@ -283,8 +286,9 @@ async def list_workflow_runs( # Non-admin callers are scoped to their own key. if not user_api_key_has_admin_view(user_api_key_dict): caller: Final = _caller_key(user_api_key_dict) - if caller: - where["created_by"] = caller + if caller is None: + return {"runs": (), "count": 0} + where["created_by"] = caller try: runs: Final[Sequence[object]] = await WorkflowRunRepository(prisma_client).table.find_many( diff --git a/litellm/proxy/management_helpers/audit_logs.py b/litellm/proxy/management_helpers/audit_logs.py index aa31f44dd13..674b178a2d7 100644 --- a/litellm/proxy/management_helpers/audit_logs.py +++ b/litellm/proxy/management_helpers/audit_logs.py @@ -8,6 +8,8 @@ from collections.abc import Awaitable, Sequence from datetime import datetime, timezone from typing import Final, TypeVar +from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter + import litellm from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid @@ -23,6 +25,32 @@ from litellm.types.utils import StandardAuditLogPayload _audit_log_callback_cache: Final[dict[str, CustomLogger]] = {} ALLOW_LITELLM_CHANGED_BY_HEADER_METADATA_KEY: Final = "allow_litellm_changed_by_header" +_AUDIT_JSON: Final = TypeAdapter(JsonValue) + + +class _AuditValues(BaseModel): + model_config = ConfigDict(frozen=True, from_attributes=True) + + updated_values: str | None = None + + +def with_audit_actor(record: LiteLLM_AuditLogs, auth: UserAPIKeyAuth) -> LiteLLM_AuditLogs: + if auth.impersonated_by is None: + return record + updated_values: Final = _AuditValues.model_validate(record).updated_values + values: Final = _AUDIT_JSON.validate_json(updated_values) if updated_values is not None else {} + return record.model_copy( + update={ + "changed_by": auth.user_id, + "changed_by_api_key": auth.token, + "updated_values": json.dumps( + { + **(values if isinstance(values, dict) else {"value": values}), + "_litellm_impersonation": {"impersonated_by": auth.impersonated_by, "subject": auth.user_id}, + } + ), + } + ) def is_audit_logging_enabled(store_audit_logs: bool | None = None) -> bool: @@ -245,16 +273,19 @@ async def create_object_audit_log( ) await create_audit_log_for_update( - request_data=LiteLLM_AuditLogs( - id=str(uuid.uuid4()), - updated_at=datetime.now(timezone.utc), - changed_by=_changed_by, - changed_by_api_key=user_api_key_dict.api_key, - table_name=table_name, - object_id=object_id, - action=action, - updated_values=after_value, - before_value=before_value, + request_data=with_audit_actor( + LiteLLM_AuditLogs( + id=str(uuid.uuid4()), + updated_at=datetime.now(timezone.utc), + changed_by=_changed_by, + changed_by_api_key=user_api_key_dict.api_key, + table_name=table_name, + object_id=object_id, + action=action, + updated_values=after_value, + before_value=before_value, + ), + user_api_key_dict, ) ) diff --git a/litellm/types/llms/custom_http.py b/litellm/types/llms/custom_http.py index 4220a67bddb..d200d4f0f72 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 a0dd405ac18..0cd44d0a9c0 100644 --- a/litellm/types/proxy/auth/user_api_key_auth.py +++ b/litellm/types/proxy/auth/user_api_key_auth.py @@ -189,6 +189,7 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response ob mcp_forwards_caller_authorization: bool = Field(default=False, 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 = Field(default=None, exclude=True) via_virtual_key: bool = Field( default=False, exclude=True, @@ -253,6 +254,7 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response ob values.pop("mcp_toolset_id", None) values.pop("via_virtual_key", None) values.pop("authenticated_by_custom_auth", None) + values.pop("impersonated_by", None) values.pop("agent_caller", None) values.pop("managed_agent_context", None) values.pop("managed_agent_policy", None) diff --git a/litellm/types/proxy/control_plane/dataplanes.py b/litellm/types/proxy/control_plane/dataplanes.py new file mode 100644 index 00000000000..012a7bf0748 --- /dev/null +++ b/litellm/types/proxy/control_plane/dataplanes.py @@ -0,0 +1,42 @@ +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): + 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..9088b71cee7 --- /dev/null +++ b/tests/unit/proxy/auth/test_impersonation.py @@ -0,0 +1,384 @@ +from collections.abc import Mapping +from typing import Final + +import pytest +from fastapi import HTTPException, Request + +from litellm.proxy._types import LiteLLM_UserTable, LitellmUserRoles, ProxyException +from litellm.proxy.auth.impersonation import ( + IMPERSONATE_EMAIL_HEADER, + IMPERSONATE_ENCODING_HEADER, + IMPERSONATE_USER_HEADER, + ImpersonatedUserNotFound, + ImpersonationForbidden, + ImpersonationRequest, + UserLookup, + read_impersonation_request, +) +from litellm.proxy.auth.impersonation import ( + impersonate as _impersonate, +) +from litellm.proxy.auth.resolvers.models import PrincipalType +from litellm.proxy.auth.user_api_key_auth import _impersonate_or_raise, _resolve_request_principal +from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth + + +async def impersonate( + actor: UserAPIKeyAuth, + request: ImpersonationRequest, + lookup: UserLookup, +) -> UserAPIKeyAuth | ImpersonationForbidden | ImpersonatedUserNotFound: + return await _impersonate(actor, request, lookup, actor_is_master_key=False) + + +ALICE: Final = LiteLLM_UserTable( + user_id="alice", user_email="alice@x.test", user_role=LitellmUserRoles.PROXY_ADMIN.value +) + + +class _BudgetLimiter: + async def is_user_within_model_budget( + self, + user_id: str, + user_model_max_budget: Mapping[str, object], + model: str, + team_id: str | None = None, + ) -> bool: + raise RuntimeError("subject model budget exceeded") + + +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.parametrize("encoding,value", [("unknown", "alice"), ("base64-utf8", "!!!!"), ("base64-utf8", "/w==")]) +def test_invalid_identity_encoding_is_a_bad_request(encoding: str, value: str) -> None: + with pytest.raises(HTTPException) as denied: + read_impersonation_request({IMPERSONATE_USER_HEADER: value, IMPERSONATE_ENCODING_HEADER: encoding}) + assert denied.value.status_code == 400 + + +@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: Final = 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_key_without_permission_is_rejected_with_http_403() -> None: + with pytest.raises(ProxyException) as exc: + await _impersonate_or_raise( + actor=_service_key({}), + impersonation=ImpersonationRequest("alice", None), + lookup=_directory, + request=Request(scope={"type": "http", "method": "GET", "path": "/key/info", "headers": []}), + request_data={}, + llm_router=None, + model_max_budget_limiter=_BudgetLimiter(), + route="/key/info", + ) + assert exc.value.code == "403" + + +@pytest.mark.asyncio +async def test_blocked_team_key_is_rejected_with_http_403() -> None: + actor: Final = _service_key({"impersonate": True}).model_copy(update={"team_blocked": True}) + with pytest.raises(ProxyException) as exc: + await _impersonate_or_raise( + actor=actor, + impersonation=ImpersonationRequest("alice", None), + lookup=_directory, + request=Request(scope={"type": "http", "method": "GET", "path": "/key/info", "headers": []}), + request_data={}, + llm_router=None, + model_max_budget_limiter=_BudgetLimiter(), + route="/key/info", + ) + assert exc.value.code == "403" + + +@pytest.mark.asyncio +async def test_unknown_subject_is_not_found_not_created() -> None: + result: Final = 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: Final = 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 is None + assert result.token == "hashed" + + +@pytest.mark.asyncio +async def test_null_permissions_and_subject_grants_use_safe_defaults() -> None: + actor: Final = _service_key({"impersonate": True}).model_copy(update={"permissions": None}) + denied: Final = await impersonate(actor, ImpersonationRequest("alice", None), _directory) + assert isinstance(denied, ImpersonationForbidden) + + subject: Final = ALICE.model_copy(update={"models": None, "model_max_budget": None, "metadata": None}) + result: Final = await _impersonate( + _service_key({"impersonate": True}), + ImpersonationRequest("alice", None), + lambda user_id: _return_subject(subject, user_id), + actor_is_master_key=False, + ) + assert isinstance(result, UserAPIKeyAuth) + assert result.models == [] + + +async def _return_subject(subject: LiteLLM_UserTable, user_id: str) -> LiteLLM_UserTable | None: + return subject if user_id == subject.user_id else None + + +@pytest.mark.asyncio +async def test_impersonated_subject_principal_keeps_the_presented_service_token() -> None: + result: Final = await impersonate( + _service_key({"impersonate": True}), ImpersonationRequest("alice", None), _directory + ) + assert isinstance(result, UserAPIKeyAuth) + request: Final = Request( + { + "type": "http", + "http_version": "1.1", + "method": "GET", + "path": "/user/info", + "raw_path": b"/user/info", + "query_string": b"", + "headers": [], + "client": ("127.0.0.1", 1), + "server": ("testserver", 80), + "scheme": "http", + } + ) + principal: Final = _resolve_request_principal(request, result) + assert principal.principal_type == PrincipalType.HUMAN + assert principal.credential_ref.token_id == "hashed" + + +@pytest.mark.asyncio +async def test_subject_does_not_inherit_the_service_keys_team_models_or_permissions() -> None: + key: Final = _service_key({"impersonate": True}).model_copy( + update={"models": ["gpt-4o"], "max_budget": 1000.0, "metadata": {"tier": "platform"}} + ) + result: Final = 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: Final = await impersonate( + _service_key({"impersonate": True}), ImpersonationRequest("alice", "bob@x.test"), _directory + ) + assert mismatch == ImpersonatedUserNotFound("alice") + same: Final = await impersonate( + _service_key({"impersonate": True}), ImpersonationRequest("alice", "ALICE@x.test"), _directory + ) + assert isinstance(same, UserAPIKeyAuth) + + +@pytest.mark.asyncio +async def test_subject_budget_is_copied_to_user_budget_fields_only() -> None: + budgeted: Final = ALICE.model_copy(update={"max_budget": 5.0, "spend": 4.0}) + + async def directory(user_id: str) -> LiteLLM_UserTable | None: + return budgeted + + result: Final = await impersonate( + _service_key({"impersonate": True}), ImpersonationRequest("alice", None), directory + ) + assert isinstance(result, UserAPIKeyAuth) + assert (result.max_budget, result.spend) == (None, 0.0) + assert (result.user_max_budget, result.user_spend) == (5.0, 4.0) + assert result.impersonated_by == "control-plane" + + +def test_impersonated_by_is_not_set_from_auth_input() -> None: + auth: Final = UserAPIKeyAuth(impersonated_by="forged") + assert auth.impersonated_by is None + + +@pytest.mark.asyncio +async def test_subject_does_not_inherit_the_service_keys_route_allowlist() -> None: + key: Final = _service_key({"impersonate": True}).model_copy(update={"allowed_routes": ["/user/info", "/key/list"]}) + result: Final = await impersonate(key, ImpersonationRequest("alice", None), _directory) + assert isinstance(result, UserAPIKeyAuth) + assert result.allowed_routes is None + + +@pytest.mark.asyncio +async def test_scim_deactivated_subject_is_forbidden() -> None: + inactive: Final = ALICE.model_copy(update={"metadata": {"scim_active": False}}) + + async def directory(user_id: str) -> LiteLLM_UserTable | None: + return inactive if user_id == "alice" else None + + result: Final = await impersonate( + _service_key({"impersonate": True}), ImpersonationRequest("alice", None), directory + ) + assert isinstance(result, ImpersonationForbidden) + assert "deactivated via SCIM" in result.detail + + +@pytest.mark.asyncio +@pytest.mark.parametrize("master", [False, True]) +async def test_view_only_subject_cannot_use_a_service_keys_write_allowlist(master: bool) -> None: + subject: Final = LiteLLM_UserTable( + user_id="alice", + user_email="alice@x.test", + user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value, + ) + actor: Final = _service_key({"impersonate": True}).model_copy( + update={"allowed_routes": ["/user/new"], "user_role": LitellmUserRoles.PROXY_ADMIN} + ) + + async def directory(user_id: str) -> LiteLLM_UserTable | None: + return subject if user_id == "alice" else None + + with pytest.raises(HTTPException) as denied: + await _impersonate_or_raise( + actor=actor, + impersonation=ImpersonationRequest("alice", None), + lookup=directory, + route="/user/new", + request=Request(scope={"type": "http", "method": "POST", "path": "/user/new", "headers": []}), + request_data={}, + llm_router=None, + model_max_budget_limiter=_BudgetLimiter(), + actor_is_master_key=master, + ) + assert denied.value.status_code == 403 + + +@pytest.mark.asyncio +async def test_admin_subject_stays_within_the_service_keys_route_ceiling() -> None: + actor: Final = _service_key({"impersonate": True}).model_copy(update={"allowed_routes": ["/user/info"]}) + + with pytest.raises(HTTPException) as exc: + await _impersonate_or_raise( + actor=actor, + impersonation=ImpersonationRequest("alice", None), + lookup=_directory, + request=Request(scope={"type": "http", "method": "POST", "path": "/user/new", "headers": []}), + request_data={}, + llm_router=None, + model_max_budget_limiter=_BudgetLimiter(), + route="/user/new", + ) + assert exc.value.status_code == 403 + + +@pytest.mark.asyncio +async def test_impersonation_is_forbidden_on_chat_completions() -> None: + with pytest.raises(ProxyException) as exc: + await _impersonate_or_raise( + actor=_service_key({"impersonate": True}), + impersonation=ImpersonationRequest("alice", None), + lookup=_directory, + request=Request(scope={"type": "http", "method": "POST", "path": "/chat/completions", "headers": []}), + request_data={}, + llm_router=None, + model_max_budget_limiter=_BudgetLimiter(), + route="/chat/completions", + ) + assert exc.value.code == "403" + + +@pytest.mark.asyncio +async def test_impersonation_remains_available_on_key_info() -> None: + result: Final = await _impersonate_or_raise( + actor=_service_key({"impersonate": True}), + impersonation=ImpersonationRequest("alice", None), + lookup=_directory, + request=Request(scope={"type": "http", "method": "GET", "path": "/key/info", "headers": []}), + request_data={}, + llm_router=None, + model_max_budget_limiter=_BudgetLimiter(), + route="/key/info", + ) + assert result.user_id == "alice" + + +@pytest.mark.asyncio +async def test_llm_route_group_does_not_allow_impersonation() -> None: + actor: Final = _service_key({"impersonate": True}).model_copy(update={"allowed_routes": ["llm_api_routes"]}) + request: Final = Request(scope={"type": "http", "method": "POST", "path": "/v1/chat/completions", "headers": []}) + + with pytest.raises(ProxyException) as exc: + await _impersonate_or_raise( + actor=actor, + impersonation=ImpersonationRequest("alice", None), + lookup=_directory, + request=request, + request_data={"model": "gpt-4o"}, + llm_router=None, + model_max_budget_limiter=_BudgetLimiter(), + route="/v1/chat/completions", + ) + assert exc.value.code == "403" + + +@pytest.mark.asyncio +async def test_impersonation_is_denied_before_subject_budget_check() -> None: + subject: Final = ALICE.model_copy( + update={"model_max_budget": {"gpt-4o": {"budget_limit": 1.0, "time_period": "1mo"}}} + ) + + async def directory(user_id: str) -> LiteLLM_UserTable | None: + return subject if user_id == "alice" else None + + with pytest.raises(ProxyException) as exc: + await _impersonate_or_raise( + actor=_service_key({"impersonate": True}), + impersonation=ImpersonationRequest("alice", None), + lookup=directory, + request=Request(scope={"type": "http", "method": "POST", "path": "/chat/completions", "headers": []}), + request_data={"model": "gpt-4o"}, + llm_router=None, + model_max_budget_limiter=_BudgetLimiter(), + route="/chat/completions", + ) + assert exc.value.code == "403" 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 d39c8cb8c85..472409cefa6 100644 --- a/tests/unit/proxy/auth/test_user_api_key_auth.py +++ b/tests/unit/proxy/auth/test_user_api_key_auth.py @@ -9,6 +9,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi import HTTPException, WebSocket, status +from fastapi import Request as FastAPIRequest from starlette.datastructures import URL from starlette.requests import Request as StarletteRequest from starlette.types import Message @@ -17,7 +18,13 @@ import litellm import litellm.proxy import litellm.proxy.proxy_server from litellm._logging import verbose_proxy_logger -from litellm.proxy._types import LiteLLM_AccessGroupTable, LiteLLM_UserTable, LitellmUserRoles, ProxyException +from litellm.proxy._types import ( + LiteLLM_AccessGroupTable, + LiteLLM_TeamTableCachedObj, + LiteLLM_UserTable, + LitellmUserRoles, + ProxyException, +) from litellm.proxy.auth.user_api_key_auth import ( UserAPIKeyAuth, _authorize_authenticated_request, @@ -25,6 +32,9 @@ from litellm.proxy.auth.user_api_key_auth import ( user_api_key_auth, ) from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.auth.impersonation import IMPERSONATE_USER_HEADER +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth_builder +from litellm.proxy.proxy_server import hash_token from tests._master_key import MASTER_KEY @@ -286,8 +296,8 @@ async def test_team_object_has_object_permission_id(): from fastapi import Request from starlette.datastructures import URL - from litellm.proxy.proxy_server import hash_token, user_api_key_cache from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + from litellm.proxy.proxy_server import hash_token, user_api_key_cache team_id = "team-vector" permission_id = "perm-vector-123" @@ -330,9 +340,10 @@ async def test_team_object_has_object_permission_id(): ) @pytest.mark.asyncio async def test_returned_user_api_key_auth(user_role, expected_role): + from datetime import datetime + from litellm.proxy._types import LiteLLM_UserTable, LitellmUserRoles from litellm.proxy.auth.user_api_key_auth import _return_user_api_key_auth_obj - from datetime import datetime new_obj = await _return_user_api_key_auth_obj( user_obj=LiteLLM_UserTable(user_role=user_role, user_id="", max_budget=None, user_email=""), @@ -368,8 +379,8 @@ async def test_aaauser_personal_budgets(key_ownership): from fastapi import Request from starlette.datastructures import URL - import litellm + import litellm from litellm.proxy._types import ( LiteLLM_UserTable, ProxyErrorTypes, @@ -530,8 +541,8 @@ def test_ui_token_route_access(route, user_role, should_be_allowed): RBAC checks as API tokens. Non-admin dashboard users must not be able to access admin-only routes like /config/update. """ - from litellm.proxy.auth.auth_checks import _is_api_route_allowed from litellm.proxy._types import LiteLLM_UserTable, UserAPIKeyAuth + from litellm.proxy.auth.auth_checks import _is_api_route_allowed user_obj = LiteLLM_UserTable( user_id="3b803c0e-666e-4e99-bd5c-6e534c07e297", @@ -553,8 +564,8 @@ def test_ui_token_route_access(route, user_role, should_be_allowed): user_role=user_role, ) - from starlette.datastructures import URL from fastapi import Request + from starlette.datastructures import URL request = Request(scope={"type": "http", "method": "POST", "path": route, "headers": []}) request._url = URL(url=route) @@ -586,8 +597,8 @@ def test_ui_token_route_access(route, user_role, should_be_allowed): ], ) def test_is_api_route_allowed(route, user_role, expected_result): - from litellm.proxy.auth.auth_checks import _is_api_route_allowed from litellm.proxy._types import LiteLLM_UserTable + from litellm.proxy.auth.auth_checks import _is_api_route_allowed user_obj = LiteLLM_UserTable( user_id="3b803c0e-666e-4e99-bd5c-6e534c07e297", @@ -690,8 +701,6 @@ def test_get_api_key_from_custom_header_different_casing(): ) - - @pytest.mark.parametrize( "user_role, auth_user_id, requested_user_id, expected_result", [ @@ -703,8 +712,8 @@ def test_get_api_key_from_custom_header_different_casing(): ], ) def test_allowed_route_inside_route(user_role, auth_user_id, requested_user_id, expected_result): + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.auth_checks import allowed_route_check_inside_route - from litellm.proxy._types import UserAPIKeyAuth, LitellmUserRoles assert ( allowed_route_check_inside_route( @@ -716,9 +725,10 @@ def test_allowed_route_inside_route(user_role, auth_user_id, requested_user_id, def test_read_request_body(): - from litellm.proxy.common_utils.http_parsing_utils import read_request_body from fastapi import Request + from litellm.proxy.common_utils.http_parsing_utils import read_request_body + payload = "()" * 1000000 request = Request(scope={"type": "http"}) @@ -739,10 +749,11 @@ async def test_auth_with_form_data_and_model(): """ from fastapi import Request from starlette.datastructures import URL, FormData + from litellm.proxy.proxy_server import ( hash_token, - user_api_key_cache, user_api_key_auth, + user_api_key_cache, ) # Setup @@ -861,10 +872,11 @@ async def test_soft_budget_alert(): def test_is_allowed_route(): - from litellm.proxy.auth.auth_checks import _is_api_route_allowed - from litellm.proxy._types import UserAPIKeyAuth import datetime + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.auth.auth_checks import _is_api_route_allowed + request = MagicMock() args = { @@ -1143,10 +1155,11 @@ async def test_accepted_websocket_reauthorizes_the_frame_model_without_closing() @pytest.mark.parametrize("enforce_rbac", [True, False]) @pytest.mark.asyncio async def test_jwt_user_api_key_auth_builder_enforce_rbac(enforce_rbac, monkeypatch): - from litellm.proxy.auth.handle_jwt import JWTHandler, JWTAuthManager - from unittest.mock import patch, Mock - from litellm.proxy._types import LiteLLM_JWTAuth + from unittest.mock import Mock, patch + from litellm.caching import DualCache + from litellm.proxy._types import LiteLLM_JWTAuth + from litellm.proxy.auth.handle_jwt import JWTAuthManager, JWTHandler monkeypatch.setenv("JWT_PUBLIC_KEY_URL", "my-fake-url") monkeypatch.setenv("JWT_AUDIENCE", "api://LiteLLM_Proxy-dev") @@ -1225,8 +1238,8 @@ def test_user_api_key_auth_end_user_str(): def test_can_rbac_role_call_model(): - from litellm.proxy.auth.handle_jwt import JWTAuthManager from litellm.proxy._types import RoleBasedPermissions + from litellm.proxy.auth.handle_jwt import JWTAuthManager roles_based_permissions = [ RoleBasedPermissions( @@ -1346,12 +1359,14 @@ async def test_jwt_non_admin_team_route_access(monkeypatch): """ Test that a non-admin JWT user cannot access team management routes """ - from fastapi import Request, HTTPException - from starlette.datastructures import URL - from unittest.mock import patch - from litellm.proxy.auth.user_api_key_auth import user_api_key_auth import json + from unittest.mock import patch + + from fastapi import HTTPException, Request + from starlette.datastructures import URL + from litellm.proxy._types import ProxyException + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth mock_jwt_response = { "is_proxy_admin": False, @@ -1382,8 +1397,8 @@ async def test_jwt_non_admin_team_route_access(monkeypatch): # Initialize jwt_handler with a default LiteLLM_JWTAuth so that the # virtual_key_claim_field check in user_api_key_auth doesn't fail with # "JWTHandler has no attribute 'litellm_jwtauth'" - from litellm.proxy._types import LiteLLM_JWTAuth from litellm.caching.dual_cache import DualCache + from litellm.proxy._types import LiteLLM_JWTAuth litellm.proxy.proxy_server.jwt_handler.update_environment( prisma_client=None, @@ -1586,6 +1601,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}") @@ -1693,7 +1709,7 @@ async def test_user_model_budget_is_enforced_through_user_api_key_auth(over_budg from fastapi import Request from starlette.datastructures import URL - from litellm.proxy._types import LiteLLM_UserTable, Litellm_EntityType, UserAPIKeyAuth + from litellm.proxy._types import Litellm_EntityType, LiteLLM_UserTable, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.hooks.model_max_budget_limiter import model_budget_spend_cache_key from litellm.proxy.proxy_server import ( @@ -2647,3 +2663,112 @@ async def test_non_object_metadata_is_served_like_unset_mappings_whatever_the_ba ) assert mapped == unset assert mapped[1] == "app-deployment" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("blocked", [False, True]) +@pytest.mark.parametrize("cached_team", [False, True]) +@pytest.mark.parametrize("last_refreshed_at", [None, 1]) +async def test_cached_admin_service_key_applies_impersonation_header( + monkeypatch: pytest.MonkeyPatch, blocked: bool, cached_team: bool, last_refreshed_at: float | None +) -> None: + key: Final = "sk-cached-service-impersonation" + subject_id: Final = "cached-impersonated-user" + token: Final = hash_token(key) + actor: Final = UserAPIKeyAuth( + token=token, + user_id="service-key-owner", + user_role=LitellmUserRoles.PROXY_ADMIN, + key_alias="cached-service", + permissions={"impersonate": True}, + team_id=f"impersonation-admin-team-{blocked}-{cached_team}-{last_refreshed_at}", + team_blocked=not blocked, + last_refreshed_at=last_refreshed_at, + ) + subject: Final = LiteLLM_UserTable( + user_id=subject_id, + user_email="subject@example.com", + user_role=LitellmUserRoles.PROXY_ADMIN.value, + ) + cache: Final = UserApiKeyCache() + team: Final = LiteLLM_TeamTableCachedObj(team_id=actor.team_id, blocked=blocked, last_refreshed_at=2) + database: Final = MagicMock() + database.db.litellm_teamtable.find_unique = AsyncMock(return_value=team) + monkeypatch.setattr(litellm.proxy.proxy_server, "prisma_client", database) + monkeypatch.setattr(litellm.proxy.proxy_server, "user_api_key_cache", cache) + monkeypatch.setattr(litellm.proxy.proxy_server, "master_key", "sk-cached-service-master") + await cache.async_set_cache(key=token, value=actor, model_type=UserAPIKeyAuth) + await cache.async_set_cache(key=subject_id, value=subject, model_type=LiteLLM_UserTable) + if cached_team: + await cache.async_set_cache(key=f"team_id:{actor.team_id}", value=team, model_type=LiteLLM_TeamTableCachedObj) + + request: Final = FastAPIRequest( + scope={ + "type": "http", + "method": "GET", + "path": "/user/info", + "headers": [(IMPERSONATE_USER_HEADER.lower().encode(), subject_id.encode())], + } + ) + request._url = URL(url="/user/info") + + async def authenticate() -> UserAPIKeyAuth: + return await user_api_key_auth_builder( + request=request, + api_key=f"Bearer {key}", + azure_api_key_header="", + anthropic_api_key_header=None, + google_ai_studio_api_key_header=None, + azure_apim_header=None, + request_data={}, + ) + + if blocked: + with pytest.raises(ProxyException) as denied: + await authenticate() + assert denied.value.code == "403" + else: + result: Final = await authenticate() + assert (result.user_id, result.impersonated_by) == (subject_id, "cached-service") + + +@pytest.mark.asyncio +async def test_master_key_impersonates_on_uncached_and_cached_requests(monkeypatch: pytest.MonkeyPatch) -> None: + key: Final = MASTER_KEY + subject_id: Final = "master-key-impersonated-user" + token: Final = hash_token(key) + subject: Final = LiteLLM_UserTable( + user_id=subject_id, + user_email="subject@example.com", + user_role=LitellmUserRoles.PROXY_ADMIN.value, + ) + cache: Final = litellm.proxy.proxy_server.user_api_key_cache + monkeypatch.setattr(litellm.proxy.proxy_server, "prisma_client", "present") + monkeypatch.setattr(litellm.proxy.proxy_server, "master_key", key) + await cache.async_delete_cache(key=token) + await cache.async_set_cache(key=subject_id, value=subject, model_type=LiteLLM_UserTable) + + async def authenticate() -> UserAPIKeyAuth: + request: Final = FastAPIRequest( + scope={ + "type": "http", + "method": "GET", + "path": "/user/info", + "headers": [(IMPERSONATE_USER_HEADER.lower().encode(), subject_id.encode())], + } + ) + request._url = URL(url="/user/info") + return await user_api_key_auth_builder( + request=request, + api_key=f"Bearer {key}", + azure_api_key_header="", + anthropic_api_key_header=None, + google_ai_studio_api_key_header=None, + azure_apim_header=None, + request_data={}, + ) + + first: Final = await authenticate() + second: Final = await authenticate() + assert (first.user_id, first.impersonated_by) == (subject_id, "litellm_proxy_master_key") + assert (second.user_id, second.impersonated_by) == (subject_id, "litellm_proxy_master_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..04117798460 --- /dev/null +++ b/tests/unit/proxy/control_plane/dataplanes/test_client.py @@ -0,0 +1,136 @@ +from collections.abc import Iterator, Sequence +from dataclasses import dataclass, field +from typing import Final + +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: Sequence[tuple[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: Final = 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: Final = 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: Final = 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: Final = 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: Final = await _client(transport, clock).get(_dataplane(), "/x", caller("session")) + assert result == DataplaneUnreachable("ConnectTimeout", 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: Final = 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: Final = FakeTransport(iter([_response(403, {"error": "no"})])) + denied: Final = 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: Final = 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: Final = 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_only_contains_the_transport_exception_type() -> None: + leaky: Final = httpx.ConnectError("failed with bearer secret-token") + transport, clock = FakeTransport(iter([leaky])), Clock() + result: Final = await _client(transport, clock).post(_dataplane(), "/key/list", caller("session"), json={}) + assert isinstance(result, DataplaneUnreachable) + assert result.detail == "ConnectError" and "secret-token" 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..3b34cbedde0 --- /dev/null +++ b/tests/unit/proxy/control_plane/dataplanes/test_credentials.py @@ -0,0 +1,157 @@ +from typing import Final, cast + +import httpx +import pytest +from starlette.datastructures import Headers + +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, + read_impersonation_request, +) +from litellm.proxy.control_plane.dataplanes.credentials import ( + Caller, + CredentialChain, + CredentialUnavailable, + IssuedCredential, + OidcTokenExchange, + ServiceTokenImpersonation, + http_token_exchanger, +) +from tests.unit.proxy.control_plane.dataplanes.conftest import caller, principal + + +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: Final = 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 +@pytest.mark.parametrize( + "subject,email", + [ + ("用户", "user@example.test"), + ("alice", "élise@example.test"), + ("用户", None), + ("alice%20smith", "alice@example.test"), + ], +) +async def test_impersonation_identities_survive_http_headers(subject: str, email: str | None) -> None: + issued: Final = await ServiceTokenImpersonation("sk-svc").headers( + Caller(principal("session", subject, email), "cp") + ) + assert isinstance(issued, IssuedCredential) + + def respond(request: httpx.Request) -> httpx.Response: + identity: Final = read_impersonation_request(Headers(headers=request.headers)) + assert identity is not None + return httpx.Response(200, json={"subject": identity.user_id, "email": identity.email}) + + async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as client: + response: Final = await client.get("https://dataplane.test/user/info", headers=issued.headers) + assert response.json() == {"subject": subject, "email": email} + + +@pytest.mark.asyncio +async def test_master_key_caller_uses_the_service_token_as_itself() -> None: + issued: Final = 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: Final = 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: Final = 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_oidc_exchange_does_not_replay_a_secret_to_a_redirect_target( + exchange_config: OAuthTokenExchangeConfig, +) -> None: + def respond(request: httpx.Request) -> httpx.Response: + assert str(request.url) == exchange_config.token_endpoint + return httpx.Response(307, headers={"Location": "https://redirect.test/token"}) + + handler: Final = AsyncHTTPHandler(transport=httpx.MockTransport(respond), follow_redirects=False) + try: + result: Final = await OidcTokenExchange(exchange_config, http_token_exchanger(handler, timeout=1.0)).headers( + caller("oidc", "idp-secret") + ) + finally: + await handler.client.aclose() + + assert result == CredentialUnavailable("oidc_exchange", "token exchange failed with status 307") + + +@pytest.mark.asyncio +async def test_failed_exchange_is_a_value_with_the_idp_status(exchange_config: OAuthTokenExchangeConfig) -> None: + result: Final = 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: Final = CredentialChain((OidcTokenExchange(exchange_config, _exchange), ServiceTokenImpersonation("sk-svc"))) + assert chain.kinds == ("oidc_exchange", "service_token") + via_session: Final = await chain.headers(caller("session")) + via_oidc: Final = 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: Final = 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: Final = http_token_exchanger( + cast(AsyncHTTPHandler, _TimingOutHandler()), # cast-ok: the handler protocol slice the exchange uses + timeout=1.0, + ) + result: Final = 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..c26d7583abe --- /dev/null +++ b/tests/unit/proxy/control_plane/dataplanes/test_registry.py @@ -0,0 +1,66 @@ +from typing import Final + +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: Final = build_registry( + {"eu": _spec(oidc_exchange=OIDC, service_token={"token": "sk"}), "us": _spec(service_token={"token": "sk"})}, + _exchange, + ) + eu: Final = 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: Final = 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 + + +def test_dataplane_url_rejects_userinfo() -> None: + with pytest.raises(ValidationError, match="must not contain userinfo"): + DataplaneSpec.model_validate({"url": "https://operator:secret@dp.example.test"}) 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..889ea660cae --- /dev/null +++ b/tests/unit/proxy/control_plane/dataplanes/test_router.py @@ -0,0 +1,199 @@ +from collections.abc import Iterator +from dataclasses import dataclass, field +from datetime import timedelta +from pathlib import Path +from typing import Final + +import httpx +import pytest +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 = "01234567890123456789012345678901" +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: Final = 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: Final = ControlPlaneSettings.model_validate({"LITELLM_MASTER_KEY": MASTER_KEY, "dataplanes": DATAPLANES}) + transport: Final = FakeTransport(iter(outcomes)) + container: Final = build_container(settings, transport=transport, exchanger=_no_exchange) + codec: Final = 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: Final = 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: Final = 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: Final = 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: Final = 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: Final = 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: Final = 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: Final = 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 + + +@pytest.mark.parametrize("status", [300, 301, 302, 303, 304, 307, 308]) +def test_whoami_does_not_accept_redirects_as_identity_proof(status: int) -> None: + client, _, codec = _rig(*(_response(status, {"user_id": "alice"}) for _ in range(3))) + reply: Final = client.get("/api/v1/dataplanes/eu/whoami", headers=_viewer(codec)) + assert reply.status_code == 502 + assert str(status) in reply.json()["detail"] + + +def test_dataplane_rejection_and_outage_surface_as_502_errors() -> None: + client, _, codec = _rig(_response(403, {"error": "lacks impersonate"}), *([httpx.ConnectError("down")] * 3)) + rejected: Final = client.get("/api/v1/dataplanes/eu/whoami", headers=_viewer(codec)) + assert rejected.status_code == 502 and "HTTP 403" in rejected.json()["detail"] + down: Final = client.get("/api/v1/dataplanes/eu/whoami", headers=_viewer(codec)) + assert down.status_code == 502 and "3 attempt" in down.json()["detail"] + + +def test_transport_errors_do_not_expose_bearer_details_from_health_or_whoami() -> None: + bearer: Final = "Bearer secret-service-token" + error: Final = f"connection failed with {bearer}" + client, _, codec = _rig(*(httpx.ConnectError(error) for _ in range(6))) + health: Final = client.get("/api/v1/dataplanes/eu/health", headers=_viewer(codec)) + whoami: Final = client.get("/api/v1/dataplanes/eu/whoami", headers=_viewer(codec)) + assert health.status_code == 200 and bearer not in health.text + assert whoami.status_code == 502 and bearer not in whoami.text + + +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: Final = _viewer(codec) + assert client.get("/api/v1/dataplanes/eu/health", headers=headers).json()["status"] == "ok" + degraded: Final = client.get("/api/v1/dataplanes/eu/health", headers=headers).json() + assert (degraded["status"], degraded["http_status"]) == ("degraded", 500) + unreachable: Final = 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: Final = tmp_path / "dataplanes.yaml" + path.write_text("dataplanes:\n eu:\n url: https://eu.test\n service_token:\n token: sk-eu\n") + specs: Final = 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: Final = 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_error_not_a_crash() -> None: + client, _, codec = _rig(_response(200, {"user_info": []}), _response(200, ["alice"])) + for _ in range(2): + malformed: Final = 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..1688df093ee --- /dev/null +++ b/tests/unit/proxy/control_plane/test_container.py @@ -0,0 +1,63 @@ +from dataclasses import dataclass, field +from typing import Final, cast + +import httpx +import pytest + +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler +from litellm.proxy.control_plane.container import HandlerTransport, _shared_http_handler +from litellm.proxy.control_plane.settings import ControlPlaneSettings + + +def test_shared_http_handler_does_not_follow_redirects() -> None: + handler: Final = _shared_http_handler( + ControlPlaneSettings.model_validate({"LITELLM_MASTER_KEY": "01234567890123456789012345678901"}) + ) + assert handler.follow_redirects is False + + +@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: Final = 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: Final = _RecyclingHandler() + transport: Final = HandlerTransport(cast(AsyncHTTPHandler, handler)) # cast-ok: only the client property is used + first: Final = await transport.request("GET", "https://dp.test/health", timeout=1.0) + second: Final = 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: Final = _RedirectingHandler() + transport: Final = HandlerTransport(cast(AsyncHTTPHandler, handler)) # cast-ok: only the client property is used + reply: Final = 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/tests/unit/proxy/control_plane/test_settings.py b/tests/unit/proxy/control_plane/test_settings.py index 71ece6c8e56..3f441f605bb 100644 --- a/tests/unit/proxy/control_plane/test_settings.py +++ b/tests/unit/proxy/control_plane/test_settings.py @@ -3,7 +3,7 @@ from typing import Final import pytest from pydantic import ValidationError -from litellm.proxy.control_plane.settings import ControlPlaneSettings +from litellm.proxy.control_plane.settings import ControlPlaneSettings, load_settings from litellm.types.proxy.control_plane.auth import PASSWORD_MAX_LENGTH, USERNAME_MAX_LENGTH MASTER_KEY: Final = "01234567890123456789012345678901" @@ -51,3 +51,12 @@ def test_credentials_at_the_bound_are_accepted() -> None: } ) assert settings.ui_username is not None and len(settings.ui_username) == USERNAME_MAX_LENGTH + + +def test_dataplanes_environment_and_file_cannot_both_be_configured(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("LITELLM_MASTER_KEY", MASTER_KEY) + monkeypatch.setenv("CONTROL_PLANE_DATAPLANES", "{}") + monkeypatch.setenv("CONTROL_PLANE_DATAPLANES_FILE", "dataplanes.yaml") + + with pytest.raises(ValueError, match="CONTROL_PLANE_DATAPLANES and CONTROL_PLANE_DATAPLANES_FILE"): + load_settings() diff --git a/tests/unit/proxy/hooks/test_key_management_event_hooks.py b/tests/unit/proxy/hooks/test_key_management_event_hooks.py index 4ba6fb95a9d..047dd9acb87 100644 --- a/tests/unit/proxy/hooks/test_key_management_event_hooks.py +++ b/tests/unit/proxy/hooks/test_key_management_event_hooks.py @@ -12,8 +12,61 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest - +from litellm.proxy._types import ( + GenerateKeyRequest, + GenerateKeyResponse, + LiteLLM_VerificationToken, + UpdateKeyRequest, + UserAPIKeyAuth, +) from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks +from litellm.proxy.management_helpers.audit_logs import drain_audit_tasks +from litellm.proxy.utils import hash_token + + +@pytest.mark.asyncio +@pytest.mark.parametrize("action", ["created", "updated", "rotated", "deleted"]) +async def test_impersonated_key_writes_preserve_actor_attribution(action: str) -> None: + actor_hash: Final = hash_token("sk-original-audit-actor") + auth: Final = UserAPIKeyAuth(api_key=None, token=actor_hash, user_id="subject").model_copy( + update={"impersonated_by": "service-key"} + ) + database: Final = MagicMock() + database.db.litellm_auditlog.create = AsyncMock() + database.db.litellm_auditlog.create_many = AsyncMock() + existing: Final = LiteLLM_VerificationToken(token="target-hash") + response: Final = GenerateKeyResponse(key="sk-new", token_id="new-hash") + with ( + patch("litellm.store_audit_logs", True), + patch("litellm.proxy.proxy_server.premium_user", True), + patch("litellm.proxy.proxy_server.prisma_client", database), + patch("litellm.audit_log_callbacks", []), + ): + if action == "created": + await KeyManagementEventHooks.async_key_generated_hook(GenerateKeyRequest(), response, auth) + elif action == "updated": + await KeyManagementEventHooks.async_key_updated_hook( + UpdateKeyRequest(key="target-hash", max_budget=42), existing, None, auth + ) + elif action == "rotated": + await KeyManagementEventHooks.async_key_rotated_hook(None, existing, response, auth) + else: + KeyManagementEventHooks.create_key_deleted_audit_logs((existing,), auth) + await asyncio.sleep(0) + await drain_audit_tasks() + row: Final = ( + database.db.litellm_auditlog.create_many.call_args.kwargs["data"][0] + if action == "deleted" + else database.db.litellm_auditlog.create.call_args.kwargs["data"] + ) + assert row["action"] == action + assert row["changed_by_api_key"] == actor_hash + assert row["changed_by"] == "subject" + assert json.loads(row["updated_values"])["_litellm_impersonation"] == { + "impersonated_by": "service-key", + "subject": "subject", + } + assert auth.api_key is None class TestKeyManagementEventHooksIndependentOperations: @@ -200,12 +253,12 @@ class TestRotateVirtualKeyInSecretManager: @pytest.mark.asyncio async def test_rotate_virtual_key_with_team_id(self): """Test that team_id is passed to async_rotate_secret.""" - from litellm.types.secret_managers.main import ( - KeyManagementSystem, - KeyManagementSettings, - ) - from litellm.secret_managers.base_secret_manager import BaseSecretManager import litellm + from litellm.secret_managers.base_secret_manager import BaseSecretManager + from litellm.types.secret_managers.main import ( + KeyManagementSettings, + KeyManagementSystem, + ) # Setup - Create a mock that inherits from BaseSecretManager mock_secret_manager = MagicMock(spec=BaseSecretManager) @@ -276,12 +329,12 @@ class TestRotateVirtualKeyInSecretManager: @pytest.mark.asyncio async def test_rotate_virtual_key_without_team_id(self): """Test that None team_id is handled correctly.""" - from litellm.types.secret_managers.main import ( - KeyManagementSystem, - KeyManagementSettings, - ) - from litellm.secret_managers.base_secret_manager import BaseSecretManager import litellm + from litellm.secret_managers.base_secret_manager import BaseSecretManager + from litellm.types.secret_managers.main import ( + KeyManagementSettings, + KeyManagementSystem, + ) # Setup - Create a mock that inherits from BaseSecretManager mock_secret_manager = MagicMock(spec=BaseSecretManager) @@ -340,16 +393,16 @@ class TestRotateVirtualKeyInSecretManager: @pytest.mark.asyncio async def test_rotate_virtual_key_in_key_rotated_hook(self): """Test that async_key_rotated_hook passes team_id to _rotate_virtual_key_in_secret_manager.""" + import litellm from litellm.proxy._types import ( - LiteLLM_VerificationToken, GenerateKeyResponse, + LiteLLM_VerificationToken, RegenerateKeyRequest, ) from litellm.types.secret_managers.main import ( - KeyManagementSystem, KeyManagementSettings, + KeyManagementSystem, ) - import litellm # Setup mock_secret_manager = MagicMock() @@ -459,11 +512,11 @@ class TestRotateVirtualKeyInSecretManager: @pytest.mark.asyncio async def test_rotate_virtual_key_when_store_virtual_keys_disabled(self): """Test that rotation is skipped when store_virtual_keys is False.""" - from litellm.types.secret_managers.main import ( - KeyManagementSystem, - KeyManagementSettings, - ) import litellm + from litellm.types.secret_managers.main import ( + KeyManagementSettings, + KeyManagementSystem, + ) # Setup mock_secret_manager = MagicMock() @@ -489,8 +542,8 @@ class TestRotateVirtualKeyInSecretManager: @pytest.mark.asyncio async def test_rotate_virtual_key_when_secret_manager_not_set(self): """Test that rotation is skipped when secret_manager_client is None.""" - from litellm.types.secret_managers.main import KeyManagementSettings import litellm + from litellm.types.secret_managers.main import KeyManagementSettings # Setup litellm.secret_manager_client = None diff --git a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py index c67cb974583..527b10c5c28 100644 --- a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py @@ -60,6 +60,7 @@ from litellm.proxy.common_utils.user_api_key_cache import ( from litellm.proxy.db.db_lookup_gate import db_lookup_stall_tracker from litellm.proxy.management_endpoints.key_management_endpoints import ( KeyProjectBindingError, + _can_user_query_key_info, _check_key_project_team, _check_key_project_team_on_mutation, _check_org_key_limits, @@ -95,6 +96,7 @@ from litellm.proxy.management_endpoints.key_management_endpoints import ( generate_key_fn, generate_key_helper_fn, generate_service_account_key_fn, + info_key_fn, key_aliases, key_generation_check, list_key_helper, @@ -120,6 +122,33 @@ from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters client = TestClient(app) +@pytest.mark.asyncio +async def test_impersonated_caller_cannot_use_the_service_key_for_key_info() -> None: + auth: Final = UserAPIKeyAuth(api_key="service-key", user_id="subject").model_copy( + update={"impersonated_by": "service-user"} + ) + key_info: Final = LiteLLM_VerificationToken(token="service-key", user_id="service-user") + + assert ( + await _can_user_query_key_info( + user_api_key_dict=auth, + key=auth.api_key, + key_info=key_info, + ) + is False + ) + + +@pytest.mark.asyncio +async def test_impersonated_key_info_requires_an_explicit_key() -> None: + auth: Final = UserAPIKeyAuth(user_id="subject").model_copy(update={"impersonated_by": "service-user"}) + + with pytest.raises(ProxyException) as error: + await info_key_fn(key=None, user_api_key_dict=auth) + + assert error.value.code == "400" + + @pytest.mark.asyncio async def test_list_keys(): mock_prisma_client = AsyncMock() diff --git a/tests/unit/proxy/management_endpoints/test_tag_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_tag_management_endpoints.py index 813875bce7c..844b4f41d7e 100644 --- a/tests/unit/proxy/management_endpoints/test_tag_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_tag_management_endpoints.py @@ -50,9 +50,7 @@ class FakeVerificationTokenTable: self.calls: list[dict[str, object]] = [] async def find_many(self, **kwargs: object) -> tuple[Mock, ...]: - inspect.signature(LiteLLM_VerificationTokenActions.find_many).bind( - self, **kwargs - ) + inspect.signature(LiteLLM_VerificationTokenActions.find_many).bind(self, **kwargs) self.calls.append(kwargs) return self._records @@ -77,9 +75,7 @@ async def test_create_and_get_tag(): with ( patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch("litellm.proxy.proxy_server.llm_router"), - patch( - "litellm.proxy.proxy_server.litellm_proxy_admin_name", "default_user_id" - ), + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "default_user_id"), patch( "litellm.proxy.management_endpoints.tag_management_endpoints.get_deployments_by_model" ) as mock_get_deployments, @@ -172,9 +168,7 @@ async def test_update_tag(): try: with ( patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, - patch( - "litellm.proxy.proxy_server.litellm_proxy_admin_name", "default_user_id" - ), + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "default_user_id"), ): # Setup prisma mocks mock_db = Mock() @@ -709,9 +703,7 @@ async def test_new_tag_invalidates_tag_and_registry_caches(): _tag_cache_doubles() as (recording_cache, mock_publish), patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch("litellm.proxy.proxy_server.llm_router"), - patch( - "litellm.proxy.proxy_server.litellm_proxy_admin_name", "default_user_id" - ), + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "default_user_id"), patch( "litellm.proxy.management_endpoints.tag_management_endpoints.get_deployments_by_model" ) as mock_get_deployments, @@ -764,9 +756,7 @@ async def test_update_tag_invalidates_only_the_tag_cache(): with ( _tag_cache_doubles() as (recording_cache, mock_publish), patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, - patch( - "litellm.proxy.proxy_server.litellm_proxy_admin_name", "default_user_id" - ), + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "default_user_id"), ): mock_db = Mock() mock_prisma.db = mock_db @@ -912,9 +902,7 @@ async def test_list_tags_with_dynamic_tags(): assert "dynamic-tag-2" in tag_names # Verify dynamic tags include created_at/updated_at - dynamic_tags = { - t["name"]: t for t in result if t["name"].startswith("dynamic-") - } + dynamic_tags = {t["name"]: t for t in result if t["name"].startswith("dynamic-")} assert dynamic_tags["dynamic-tag-1"]["created_at"] is not None assert dynamic_tags["dynamic-tag-1"]["updated_at"] is not None @@ -1036,9 +1024,7 @@ async def test_internal_user_list_tags_only_returns_tags_used_by_their_keys(): "stored-owned-tag", "dynamic-owned-tag", ] - assert fake_token_table.calls == [ - {"where": {"user_id": "internal-user-123"}} - ] + assert fake_token_table.calls == [{"where": {"user_id": "internal-user-123"}}] mock_db.litellm_dailytagspend.group_by.assert_awaited_once_with( by=["tag"], where={ @@ -1092,15 +1078,11 @@ async def test_internal_user_list_tags_does_not_500_on_unsupported_prisma_kwarg( mock_db.litellm_dailytagspend.group_by = AsyncMock(return_value=[]) mock_db.litellm_tagtable.find_many = AsyncMock(return_value=[]) - response = client.get( - "/tag/list", headers={"Authorization": "Bearer new-user-key"} - ) + response = client.get("/tag/list", headers={"Authorization": "Bearer new-user-key"}) assert response.status_code == 200, response.text assert response.json() == [] - assert fake_token_table.calls == [ - {"where": {"user_id": "brand-new-internal-user"}} - ] + assert fake_token_table.calls == [{"where": {"user_id": "brand-new-internal-user"}}] finally: app.dependency_overrides.clear() @@ -1277,9 +1259,7 @@ async def test_internal_user_tag_daily_activity_scopes_to_current_key_without_us assert result == "daily-activity-response" assert fake_token_table.calls == [] mock_get_daily_activity.assert_awaited_once() - assert mock_get_daily_activity.await_args.kwargs["api_key"] == [ - "current-owned-key" - ] + assert mock_get_daily_activity.await_args.kwargs["api_key"] == ["current-owned-key"] @pytest.mark.asyncio @@ -1556,9 +1536,7 @@ async def test_add_tag_to_deployment_preserves_encrypted_fields(): await _add_tag_to_deployment(deployment, "test-tag") # Verify find_unique was called - mock_db.litellm_proxymodeltable.find_unique.assert_called_once_with( - where={"model_id": "model-123"} - ) + mock_db.litellm_proxymodeltable.find_unique.assert_called_once_with(where={"model_id": "model-123"}) # Verify update was called with preserved encrypted fields update_call = mock_db.litellm_proxymodeltable.update.call_args @@ -1780,3 +1758,28 @@ def test_tag_info_and_tag_list_return_the_stored_model_info_decoded( assert info_response.json()["routed-tag"]["model_info"] == returned_model_info assert list_response.status_code == 200 assert [tag["model_info"] for tag in list_response.json()] == [returned_model_info] + + +@pytest.mark.asyncio +async def test_impersonated_internal_user_tag_activity_filter_excludes_actor_key() -> None: + from litellm.proxy.management_endpoints.tag_management_endpoints import get_tag_daily_activity_api_key_filter + + actor_token = "service-key-hash" + subject_token = "subject-key-hash" + prisma = cast( + PrismaClient, + SimpleNamespace( + db=SimpleNamespace(litellm_verificationtoken=FakeVerificationTokenTable((Mock(token=subject_token),))) + ), + ) + auth = UserAPIKeyAuth( + api_key=actor_token, + token=actor_token, + user_id="internal-user", + user_role=LitellmUserRoles.INTERNAL_USER, + ).model_copy(update={"impersonated_by": "service-key"}) + + api_keys = await get_tag_daily_activity_api_key_filter(prisma, auth, None) + + assert api_keys == [subject_token] + assert actor_token not in api_keys diff --git a/tests/unit/proxy/management_endpoints/test_workflow_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_workflow_management_endpoints.py index 0c3d5107cb6..c6fa24c6d64 100644 --- a/tests/unit/proxy/management_endpoints/test_workflow_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_workflow_management_endpoints.py @@ -13,14 +13,12 @@ from fastapi import FastAPI, HTTPException from fastapi.testclient import TestClient from prisma.errors import UniqueViolationError - from litellm.proxy.management_endpoints.workflow_management_endpoints import ( _read_scope_caller, _require_run, router, ) - # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- @@ -87,15 +85,11 @@ def _make_tx(event_return=None, run_return=None, msg_return=None) -> MagicMock: """Build an async context-manager mock for prisma_client.db.tx().""" tx = MagicMock() tx.litellm_workflowevent = MagicMock() - tx.litellm_workflowevent.create = AsyncMock( - return_value=event_return or _make_event() - ) + tx.litellm_workflowevent.create = AsyncMock(return_value=event_return or _make_event()) tx.litellm_workflowrun = MagicMock() tx.litellm_workflowrun.update = AsyncMock(return_value=run_return or _make_run()) tx.litellm_workflowmessage = MagicMock() - tx.litellm_workflowmessage.create = AsyncMock( - return_value=msg_return or _make_message() - ) + tx.litellm_workflowmessage.create = AsyncMock(return_value=msg_return or _make_message()) tx.__aenter__ = AsyncMock(return_value=tx) tx.__aexit__ = AsyncMock(return_value=False) return tx @@ -215,9 +209,7 @@ class TestListWorkflowRuns: @patch("litellm.proxy.proxy_server.prisma_client") def test_list_returns_runs(self, mock_pc): mock_pc.db = self._prisma.db - self._prisma.db.litellm_workflowrun.find_many = AsyncMock( - return_value=[_make_run()] - ) + self._prisma.db.litellm_workflowrun.find_many = AsyncMock(return_value=[_make_run()]) resp = self.client.get("/v1/workflows/runs") assert resp.status_code == 200 @@ -257,9 +249,7 @@ class TestGetWorkflowRun: @patch("litellm.proxy.proxy_server.prisma_client") def test_get_existing_run(self, mock_pc): mock_pc.db = self._prisma.db - self._prisma.db.litellm_workflowrun.find_unique = AsyncMock( - return_value=_make_run() - ) + self._prisma.db.litellm_workflowrun.find_unique = AsyncMock(return_value=_make_run()) resp = self.client.get("/v1/workflows/runs/run-1") assert resp.status_code == 200 @@ -285,15 +275,11 @@ class TestUpdateWorkflowRun: @patch("litellm.proxy.proxy_server.prisma_client") def test_update_status(self, mock_pc): mock_pc.db = self._prisma.db - self._prisma.db.litellm_workflowrun.find_unique = AsyncMock( - return_value=_make_run() - ) + self._prisma.db.litellm_workflowrun.find_unique = AsyncMock(return_value=_make_run()) updated = _make_run(status="completed") self._prisma.db.litellm_workflowrun.update = AsyncMock(return_value=updated) - resp = self.client.patch( - "/v1/workflows/runs/run-1", json={"status": "completed"} - ) + resp = self.client.patch("/v1/workflows/runs/run-1", json={"status": "completed"}) assert resp.status_code == 200 self._prisma.db.litellm_workflowrun.update.assert_awaited_once() @@ -317,13 +303,9 @@ class TestAppendWorkflowEvent: def test_append_event_updates_run_status(self, mock_pc): mock_pc.db = self._prisma.db # _require_run check - self._prisma.db.litellm_workflowrun.find_unique = AsyncMock( - return_value=_make_run() - ) + self._prisma.db.litellm_workflowrun.find_unique = AsyncMock(return_value=_make_run()) self._prisma.db.litellm_workflowevent.find_many = AsyncMock(return_value=[]) - tx = _make_tx( - event_return=_make_event(), run_return=_make_run(status="running") - ) + tx = _make_tx(event_return=_make_event(), run_return=_make_run(status="running")) self._prisma.db.tx = MagicMock(return_value=tx) resp = self.client.post( @@ -339,9 +321,7 @@ class TestAppendWorkflowEvent: @patch("litellm.proxy.proxy_server.prisma_client") def test_append_event_no_status_update_for_unknown_type(self, mock_pc): mock_pc.db = self._prisma.db - self._prisma.db.litellm_workflowrun.find_unique = AsyncMock( - return_value=_make_run() - ) + self._prisma.db.litellm_workflowrun.find_unique = AsyncMock(return_value=_make_run()) self._prisma.db.litellm_workflowevent.find_many = AsyncMock(return_value=[]) tx = _make_tx(event_return=_make_event(event_type="custom.event")) self._prisma.db.tx = MagicMock(return_value=tx) @@ -357,13 +337,9 @@ class TestAppendWorkflowEvent: @patch("litellm.proxy.proxy_server.prisma_client") def test_sequence_number_increments(self, mock_pc): mock_pc.db = self._prisma.db - self._prisma.db.litellm_workflowrun.find_unique = AsyncMock( - return_value=_make_run() - ) + self._prisma.db.litellm_workflowrun.find_unique = AsyncMock(return_value=_make_run()) existing = _make_event(sequence_number=4) - self._prisma.db.litellm_workflowevent.find_many = AsyncMock( - return_value=[existing] - ) + self._prisma.db.litellm_workflowevent.find_many = AsyncMock(return_value=[existing]) tx = _make_tx(event_return=_make_event(sequence_number=5)) self._prisma.db.tx = MagicMock(return_value=tx) @@ -389,18 +365,14 @@ class TestAppendWorkflowEvent: def test_sequence_collision_retries_and_succeeds(self, mock_pc): """UniqueViolationError on first attempt triggers retry; second attempt succeeds.""" mock_pc.db = self._prisma.db - self._prisma.db.litellm_workflowrun.find_unique = AsyncMock( - return_value=_make_run() - ) + self._prisma.db.litellm_workflowrun.find_unique = AsyncMock(return_value=_make_run()) self._prisma.db.litellm_workflowevent.find_many = AsyncMock(return_value=[]) # First tx raises UniqueViolationError; second succeeds. tx_fail = _make_tx() tx_fail.__aenter__ = AsyncMock(return_value=tx_fail) tx_fail.litellm_workflowevent.create = AsyncMock( - side_effect=UniqueViolationError( - {"user_facing_error": {"message": "unique"}} - ) + side_effect=UniqueViolationError({"user_facing_error": {"message": "unique"}}) ) tx_fail.__aexit__ = AsyncMock(return_value=False) @@ -427,13 +399,9 @@ class TestWorkflowMessages: @patch("litellm.proxy.proxy_server.prisma_client") def test_append_message(self, mock_pc): mock_pc.db = self._prisma.db - self._prisma.db.litellm_workflowrun.find_unique = AsyncMock( - return_value=_make_run() - ) + self._prisma.db.litellm_workflowrun.find_unique = AsyncMock(return_value=_make_run()) self._prisma.db.litellm_workflowmessage.find_many = AsyncMock(return_value=[]) - self._prisma.db.litellm_workflowmessage.create = AsyncMock( - return_value=_make_message() - ) + self._prisma.db.litellm_workflowmessage.create = AsyncMock(return_value=_make_message()) resp = self.client.post( "/v1/workflows/runs/run-1/messages", @@ -455,9 +423,7 @@ class TestWorkflowMessages: @patch("litellm.proxy.proxy_server.prisma_client") def test_list_messages_ordered(self, mock_pc): mock_pc.db = self._prisma.db - self._prisma.db.litellm_workflowrun.find_unique = AsyncMock( - return_value=_make_run() - ) + self._prisma.db.litellm_workflowrun.find_unique = AsyncMock(return_value=_make_run()) self._prisma.db.litellm_workflowmessage.find_many = AsyncMock( return_value=[ _make_message(sequence_number=0), @@ -475,9 +441,7 @@ class TestWorkflowMessages: @patch("litellm.proxy.proxy_server.prisma_client") def test_list_messages_respects_limit(self, mock_pc): mock_pc.db = self._prisma.db - self._prisma.db.litellm_workflowrun.find_unique = AsyncMock( - return_value=_make_run() - ) + self._prisma.db.litellm_workflowrun.find_unique = AsyncMock(return_value=_make_run()) self._prisma.db.litellm_workflowmessage.find_many = AsyncMock(return_value=[]) resp = self.client.get("/v1/workflows/runs/run-1/messages?limit=25") @@ -498,9 +462,7 @@ class TestListWorkflowEvents: @patch("litellm.proxy.proxy_server.prisma_client") def test_list_events_ordered(self, mock_pc): mock_pc.db = self._prisma.db - self._prisma.db.litellm_workflowrun.find_unique = AsyncMock( - return_value=_make_run() - ) + self._prisma.db.litellm_workflowrun.find_unique = AsyncMock(return_value=_make_run()) self._prisma.db.litellm_workflowevent.find_many = AsyncMock( return_value=[ _make_event(sequence_number=0), @@ -518,9 +480,7 @@ class TestListWorkflowEvents: @patch("litellm.proxy.proxy_server.prisma_client") def test_list_events_respects_limit(self, mock_pc): mock_pc.db = self._prisma.db - self._prisma.db.litellm_workflowrun.find_unique = AsyncMock( - return_value=_make_run() - ) + self._prisma.db.litellm_workflowrun.find_unique = AsyncMock(return_value=_make_run()) self._prisma.db.litellm_workflowevent.find_many = AsyncMock(return_value=[]) resp = self.client.get("/v1/workflows/runs/run-1/events?limit=10") @@ -553,9 +513,7 @@ class TestTenantIsolation: token = "tok-owner" client = self._make_app_with_auth(lambda: _override_auth_user_with_token(token)) mock_pc.db = self._prisma.db - self._prisma.db.litellm_workflowrun.create = AsyncMock( - return_value=_make_run(created_by=token) - ) + self._prisma.db.litellm_workflowrun.create = AsyncMock(return_value=_make_run(created_by=token)) resp = client.post("/v1/workflows/runs", json={"workflow_type": "test"}) assert resp.status_code == 200 @@ -574,6 +532,44 @@ class TestTenantIsolation: call_kwargs = self._prisma.db.litellm_workflowrun.find_many.call_args[1] assert call_kwargs["where"].get("created_by") == token + @patch("litellm.proxy.proxy_server.prisma_client") + def test_impersonated_list_does_not_use_service_key_as_owner(self, mock_pc): + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + + actor_token = "tok-service-actor" + auth = UserAPIKeyAuth( + api_key=None, + token=actor_token, + user_id="subject", + user_role=LitellmUserRoles.INTERNAL_USER, + ).model_copy(update={"impersonated_by": "service-key"}) + client = self._make_app_with_auth(lambda: auth) + mock_pc.db = self._prisma.db + self._prisma.db.litellm_workflowrun.find_many = AsyncMock(return_value=[]) + + resp = client.get("/v1/workflows/runs") + assert resp.status_code == 200 + assert resp.json() == {"runs": [], "count": 0} + self._prisma.db.litellm_workflowrun.find_many.assert_not_awaited() + + @patch("litellm.proxy.proxy_server.prisma_client") + def test_impersonated_create_is_forbidden(self, mock_pc): + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + + auth = UserAPIKeyAuth( + api_key=None, + token="tok-service-actor", + user_id="subject", + user_role=LitellmUserRoles.INTERNAL_USER, + ).model_copy(update={"impersonated_by": "service-key"}) + client = self._make_app_with_auth(lambda: auth) + mock_pc.db = self._prisma.db + self._prisma.db.litellm_workflowrun.create = AsyncMock(return_value=_make_run()) + + resp = client.post("/v1/workflows/runs", json={"workflow_type": "test"}) + assert resp.status_code == 403 + self._prisma.db.litellm_workflowrun.create.assert_not_awaited() + @patch("litellm.proxy.proxy_server.prisma_client") def test_admin_list_not_scoped(self, mock_pc): client = self._make_app_with_auth(_override_auth_admin) @@ -603,9 +599,7 @@ class TestTenantIsolation: token = "tok-caller" client = self._make_app_with_auth(lambda: _override_auth_user_with_token(token)) mock_pc.db = self._prisma.db - self._prisma.db.litellm_workflowrun.find_unique = AsyncMock( - return_value=_make_run(created_by=None) - ) + self._prisma.db.litellm_workflowrun.find_unique = AsyncMock(return_value=_make_run(created_by=None)) resp = client.get("/v1/workflows/runs/run-1") assert resp.status_code == 404 @@ -615,12 +609,8 @@ class TestTenantIsolation: token = "tok-caller" client = self._make_app_with_auth(lambda: _override_auth_user_with_token(token)) mock_pc.db = self._prisma.db - self._prisma.db.litellm_workflowrun.find_unique = AsyncMock( - return_value=_make_run(created_by=None) - ) - self._prisma.db.litellm_workflowrun.update = AsyncMock( - return_value=_make_run(status="completed") - ) + self._prisma.db.litellm_workflowrun.find_unique = AsyncMock(return_value=_make_run(created_by=None)) + self._prisma.db.litellm_workflowrun.update = AsyncMock(return_value=_make_run(status="completed")) resp = client.patch("/v1/workflows/runs/run-1", json={"status": "completed"}) assert resp.status_code == 404 @@ -631,9 +621,7 @@ class TestTenantIsolation: token = "tok-caller" client = self._make_app_with_auth(lambda: _override_auth_user_with_token(token)) mock_pc.db = self._prisma.db - self._prisma.db.litellm_workflowrun.find_unique = AsyncMock( - return_value=_make_run(created_by=token) - ) + self._prisma.db.litellm_workflowrun.find_unique = AsyncMock(return_value=_make_run(created_by=token)) resp = client.get("/v1/workflows/runs/run-1") assert resp.status_code == 200 @@ -685,9 +673,7 @@ class TestAdminViewerReadParity: self._prisma.db.litellm_workflowrun.find_unique = AsyncMock( return_value=_make_run(created_by="tok-other-owner") ) - self._prisma.db.litellm_workflowevent.find_many = AsyncMock( - return_value=[_make_event(sequence_number=0)] - ) + self._prisma.db.litellm_workflowevent.find_many = AsyncMock(return_value=[_make_event(sequence_number=0)]) resp = client.get("/v1/workflows/runs/run-1/events") assert resp.status_code == 200 @@ -700,9 +686,7 @@ class TestAdminViewerReadParity: self._prisma.db.litellm_workflowrun.find_unique = AsyncMock( return_value=_make_run(created_by="tok-other-owner") ) - self._prisma.db.litellm_workflowmessage.find_many = AsyncMock( - return_value=[_make_message(sequence_number=0)] - ) + self._prisma.db.litellm_workflowmessage.find_many = AsyncMock(return_value=[_make_message(sequence_number=0)]) resp = client.get("/v1/workflows/runs/run-1/messages") assert resp.status_code == 200 @@ -716,9 +700,7 @@ class TestAdminViewerReadParity: self._prisma.db.litellm_workflowrun.find_unique = AsyncMock( return_value=_make_run(created_by="tok-other-owner") ) - self._prisma.db.litellm_workflowrun.update = AsyncMock( - return_value=_make_run(status="completed") - ) + self._prisma.db.litellm_workflowrun.update = AsyncMock(return_value=_make_run(status="completed")) resp = client.patch("/v1/workflows/runs/run-1", json={"status": "completed"}) assert resp.status_code == 404 @@ -727,9 +709,7 @@ class TestAdminViewerReadParity: def test_require_run_still_scopes_when_handed_a_viewer(self): """Only read callers pass None; the helper itself never loosened.""" prisma = _make_prisma_client() - prisma.db.litellm_workflowrun.find_unique = AsyncMock( - return_value=_make_run(created_by="tok-other-owner") - ) + prisma.db.litellm_workflowrun.find_unique = AsyncMock(return_value=_make_run(created_by="tok-other-owner")) with pytest.raises(HTTPException) as exc_info: asyncio.run(_require_run(prisma, "run-1", _override_auth_admin_viewer())) diff --git a/tests/unit/proxy/management_helpers/test_audit_logs_proxy.py b/tests/unit/proxy/management_helpers/test_audit_logs_proxy.py index 325275ab391..a423b9bec3d 100644 --- a/tests/unit/proxy/management_helpers/test_audit_logs_proxy.py +++ b/tests/unit/proxy/management_helpers/test_audit_logs_proxy.py @@ -42,6 +42,7 @@ from starlette.datastructures import URL from litellm.proxy.management_helpers.audit_logs import ( create_audit_log_for_update, + create_object_audit_log, drain_audit_tasks, get_audit_log_changed_by, track_audit_task, @@ -58,6 +59,48 @@ import json from tests._master_key import MASTER_KEY +@pytest.mark.asyncio +@pytest.mark.parametrize("impersonated", [False, True]) +@pytest.mark.parametrize("after_value", [None, '{"max_budget":42}', '"deleted"']) +async def test_object_audit_preserves_the_actor_and_subject_in_storage_and_exports( + impersonated: bool, after_value: str | None +) -> None: + actor_hash: Final = hash_token("sk-original-audit-actor") + auth: Final = UserAPIKeyAuth( + api_key=None if impersonated else actor_hash, + token=actor_hash, + user_id="subject", + ).model_copy(update={"impersonated_by": "service-key" if impersonated else None}) + database: Final = MagicMock() + database.db.litellm_auditlog.create = AsyncMock() + callback: Final = CustomLogger() + export: Final = AsyncMock() + callback.async_log_audit_log_event = export + with ( + patch("litellm.store_audit_logs", True), + patch("litellm.proxy.proxy_server.premium_user", True), + patch("litellm.proxy.proxy_server.prisma_client", database), + patch("litellm.audit_log_callbacks", [callback]), + ): + await create_object_audit_log( + "object", "updated", None, auth, "admin", LitellmTableNames.KEY_TABLE_NAME, after_value=after_value + ) + await drain_audit_tasks() + stored: Final = database.db.litellm_auditlog.create.call_args.kwargs["data"] + exported: Final = export.call_args.args[0] + assert stored["changed_by"] == exported["changed_by"] == "subject" + assert stored["changed_by_api_key"] == exported["changed_by_api_key"] == actor_hash + assert stored.get("updated_values") == exported["updated_values"] + if impersonated: + assert json.loads(stored["updated_values"])["_litellm_impersonation"] == { + "impersonated_by": "service-key", + "subject": "subject", + } + assert auth.api_key is None + else: + assert "_litellm_impersonation" not in (stored.get("updated_values") or "") + + def test_get_audit_log_changed_by_prefers_authenticated_user(): user_api_key_dict = UserAPIKeyAuth( api_key="test-key",