diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 76f00846c21..4f2e970e97a 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -1,4315 +1 @@ -""" -Has all /sso/* routes - -/sso/key/generate - handles user signing in with SSO and redirects to /sso/callback -/sso/callback - returns JWT Redirect Response that redirects to LiteLLM UI - -/sso/debug/login - handles user signing in with SSO and redirects to /sso/debug/callback -/sso/debug/callback - returns the OpenID object returned by the SSO provider -""" - -import asyncio -import base64 -import hashlib -import inspect -import os -import re -import secrets -from copy import deepcopy -from html import escape -from typing import ( - TYPE_CHECKING, - Any, - Dict, - List, - Literal, - NoReturn, - Optional, - Tuple, - Union, - cast, -) -from urllib.parse import parse_qs, urlencode, urlparse - -if TYPE_CHECKING: - import httpx - -import jwt -from fastapi import APIRouter, Depends, Header, HTTPException, Request, status -from fastapi.responses import RedirectResponse - -import litellm -from litellm._logging import verbose_proxy_logger -from litellm._uuid import uuid -from litellm.caching.dual_cache import DualCache -from litellm.constants import ( - CLI_SSO_CLAIM_MAP, - CLI_SSO_CLAIM_MAX_SCALAR_LENGTH, - CLI_SSO_SESSION_CACHE_KEY_PREFIX, - CLI_SSO_SESSION_TTL_SECONDS, - LITELLM_CLI_SOURCE_IDENTIFIER, - LITELLM_UI_SESSION_DURATION, - MAX_SPENDLOG_ROWS_TO_QUERY, - MICROSOFT_USER_DISPLAY_NAME_ATTRIBUTE, - MICROSOFT_USER_EMAIL_ATTRIBUTE, - MICROSOFT_USER_FIRST_NAME_ATTRIBUTE, - MICROSOFT_USER_ID_ATTRIBUTE, - MICROSOFT_USER_LAST_NAME_ATTRIBUTE, -) -from litellm.litellm_core_utils.dot_notation_indexing import get_nested_value -from litellm.llms.custom_httpx.http_handler import ( - AsyncHTTPHandler, - get_async_httpx_client, - httpxSpecialProvider, -) -from litellm.proxy._types import ( - CommonProxyErrors, - LiteLLM_UserTable, - LitellmUserRoles, - Member, - NewTeamRequest, - NewUserRequest, - NewUserResponse, - ProxyErrorTypes, - ProxyException, - SSOUserDefinedValues, - TeamMemberAddRequest, - UserAPIKeyAuth, -) -from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken, get_user_object -from litellm.proxy.auth.auth_utils import ( - _get_request_ip_address, - _has_user_setup_sso, -) -from litellm.proxy.auth.handle_jwt import JWTHandler -from litellm.proxy.auth.user_api_key_auth import user_api_key_auth -from litellm.proxy.common_utils.admin_ui_utils import ( - admin_ui_disabled, - show_missing_vars_in_env, -) -from litellm.proxy.common_utils.html_forms.jwt_display_template import ( - jwt_display_template, -) -from litellm.proxy.common_utils.html_forms.ui_login import build_ui_login_form -from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache -from litellm.proxy.management_endpoints.internal_user_endpoints import new_user -from litellm.proxy.management_endpoints.sso import CustomMicrosoftSSO -from litellm.proxy.management_endpoints.sso_helper_utils import ( - check_is_admin_only_access, - has_admin_ui_access, -) -from litellm.proxy.management_endpoints.team_endpoints import new_team, team_member_add -from litellm.proxy.management_endpoints.types import ( - CustomOpenID, - get_litellm_user_role, - is_valid_litellm_user_role, -) -from litellm.proxy.utils import ( - PrismaClient, - ProxyLogging, - get_custom_url, - get_server_root_path, -) -from litellm.repositories.table_repositories import SSOConfigRepository -from litellm.repositories.team_repository import TeamRepository -from litellm.repositories.user_repository import UserRepository -from litellm.secret_managers.main import get_secret_bool, str_to_bool -from litellm.types.proxy.management_endpoints.ui_sso import * # noqa: F403 -from litellm.types.proxy.management_endpoints.ui_sso import ( - DefaultTeamSSOParams, - MicrosoftGraphAPIUserGroupDirectoryObject, - MicrosoftGraphAPIUserGroupResponse, - MicrosoftServicePrincipalTeam, - RoleMappings, - TeamMappings, -) -from litellm.types.proxy.ui_sso import ParsedOpenIDResult - -if TYPE_CHECKING: - from fastapi_sso.sso.base import OpenID -else: - from typing import Any as OpenID - -router = APIRouter() - -# OAuth bearer credential fields that must not appear in SSO debug responses -# (received_response is included in restricted-group error messages). -# Metadata fields (token_type, expires_in, scope) are intentionally kept so -# response convertors see the same fields in the PKCE path as in the non-PKCE path. -_OAUTH_TOKEN_FIELDS = frozenset({"access_token", "id_token", "refresh_token"}) -_CLI_SSO_FLOW_CACHE_KEY_PREFIX = f"{CLI_SSO_SESSION_CACHE_KEY_PREFIX}:flow" -_CLI_SSO_START_RATE_LIMIT_CACHE_KEY_PREFIX = f"{_CLI_SSO_FLOW_CACHE_KEY_PREFIX}:start_rate_limit" -_CLI_SSO_START_RATE_LIMIT_WINDOW_SECONDS = 60 -_CLI_SSO_START_RATE_LIMIT_MAX_ATTEMPTS = 30 -_CLI_SSO_USER_CODE_ALPHABET = "ABCDEFGHJKLMNPQRSTUVWXYZ23456789" -_CLI_SSO_LOGIN_ID_RE = re.compile(r"^cli-[A-Za-z0-9_-]{12,124}$") -_CLI_SSO_USER_CODE_RE = re.compile(rf"^[{_CLI_SSO_USER_CODE_ALPHABET}]{{4}}-[{_CLI_SSO_USER_CODE_ALPHABET}]{{4}}$") -_CLI_SSO_SCALAR_TYPES = (str, int, float, bool) -_CLI_SSO_DEST_KEY_RE = re.compile(r"^[A-Za-z0-9_.-]+$") -_CLI_SSO_SECRET_KEY_FRAGMENTS = frozenset( - { - "access_token", - "api_key", - "client_secret", - "id_token", - "password", - "private_key", - "refresh_token", - "secret", - } -) - - -def _hash_cli_sso_secret(secret: str) -> str: - return hashlib.sha256(secret.encode("utf-8")).hexdigest() - - -def _normalize_cli_sso_user_code(user_code: str) -> str: - return "".join(ch for ch in user_code.upper() if ch.isalnum()) - - -def _generate_cli_sso_user_code() -> str: - user_code = "".join(secrets.choice(_CLI_SSO_USER_CODE_ALPHABET) for _ in range(8)) - return f"{user_code[:4]}-{user_code[4:]}" - - -def _get_cli_sso_flow_cache_key(login_id: str) -> str: - return f"{_CLI_SSO_FLOW_CACHE_KEY_PREFIX}:{login_id}" - - -def _is_valid_cli_sso_login_id(login_id: Optional[str]) -> bool: - return isinstance(login_id, str) and bool(_CLI_SSO_LOGIN_ID_RE.fullmatch(login_id)) - - -def _is_valid_cli_sso_user_code(user_code: str | None) -> bool: - return isinstance(user_code, str) and bool(_CLI_SSO_USER_CODE_RE.fullmatch(user_code)) - - -def _cli_sso_verification_uri_complete_enabled() -> bool: - from litellm.proxy.proxy_server import general_settings - - return bool(general_settings.get("allow_cli_sso_verification_uri_complete", False)) - - -def _cli_sso_start_response_body( - *, - login_id: str, - poll_secret: str, - user_code: str, - verification_uri_complete: str | None, -) -> dict[str, str | int]: - if verification_uri_complete is None: - return { - "login_id": login_id, - "poll_secret": poll_secret, - "user_code": user_code, - "expires_in": CLI_SSO_SESSION_TTL_SECONDS, - } - return { - "login_id": login_id, - "poll_secret": poll_secret, - "user_code": user_code, - "verification_uri_complete": verification_uri_complete, - "expires_in": CLI_SSO_SESSION_TTL_SECONDS, - } - - -def _get_cli_sso_start_rate_limit_cache_key(request: Request, use_x_forwarded_for: Optional[bool] = False) -> str: - client_ip = _get_request_ip_address(request=request, use_x_forwarded_for=use_x_forwarded_for) or "unknown" - client_ip_hash = _hash_cli_sso_secret(client_ip) - return f"{_CLI_SSO_START_RATE_LIMIT_CACHE_KEY_PREFIX}:{client_ip_hash}" - - -def _check_cli_sso_start_rate_limit( - request: Request, - cache: DualCache, - use_x_forwarded_for: Optional[bool] = False, -) -> None: - rate_limit_cache_key = _get_cli_sso_start_rate_limit_cache_key( - request=request, use_x_forwarded_for=use_x_forwarded_for - ) - current_attempts = cache.increment_cache( - key=rate_limit_cache_key, - value=1, - ttl=_CLI_SSO_START_RATE_LIMIT_WINDOW_SECONDS, - ) - if current_attempts > _CLI_SSO_START_RATE_LIMIT_MAX_ATTEMPTS: - raise HTTPException( - status_code=429, - detail="Too many CLI login attempts. Try again later.", - ) - - -def _get_cli_sso_flow_or_raise(login_id: Optional[str], cache: DualCache) -> dict: - if not _is_valid_cli_sso_login_id(login_id): - raise HTTPException(status_code=400, detail="Invalid CLI login session") - - cache_key = _get_cli_sso_flow_cache_key(cast(str, login_id)) - flow = cache.get_cache(key=cache_key) - if not isinstance(flow, dict) or "poll_secret_hash" not in flow: - raise HTTPException(status_code=400, detail="Invalid CLI login session") - return flow - - -def _set_cli_sso_flow(login_id: str, cache: DualCache, flow: dict) -> None: - cache.set_cache( - key=_get_cli_sso_flow_cache_key(login_id), - value=flow, - ttl=CLI_SSO_SESSION_TTL_SECONDS, - ) - - -def _verify_cli_sso_poll_secret(flow: dict, poll_secret: Optional[str]) -> bool: - expected_poll_secret_hash = flow.get("poll_secret_hash") - if not isinstance(expected_poll_secret_hash, str) or not isinstance(poll_secret, str): - return False - supplied_poll_secret_hash = _hash_cli_sso_secret(poll_secret) - return secrets.compare_digest(supplied_poll_secret_hash, expected_poll_secret_hash) - - -def _parse_cli_sso_claim_map() -> List[Tuple[str, str]]: - """ - Parse CLI_SSO_CLAIM_MAP / LITELLM_CLI_SSO_CLAIM_MAP. - - Format: comma-separated ``source_claim->metadata_key`` pairs, e.g. - ``employment_type->acme_employment_type,org_info.department->department``. - Destination keys may use an optional ``metadata.`` prefix; values are stored - on the LiteLLM user's ``metadata`` JSON column. - """ - claim_map_raw = CLI_SSO_CLAIM_MAP.strip() - if not claim_map_raw: - return [] - - parsed: List[Tuple[str, str]] = [] - for entry in claim_map_raw.split(","): - entry = entry.strip() - if not entry or "->" not in entry: - continue - source_claim, dest_key = entry.split("->", 1) - source_claim = source_claim.strip() - dest_key = dest_key.strip() - if dest_key.startswith("metadata."): - dest_key = dest_key[len("metadata.") :] - if source_claim and dest_key: - parsed.append((source_claim, dest_key)) - return parsed - - -def _is_safe_cli_sso_metadata_dest_key(dest_key: str) -> bool: - if not dest_key or not _CLI_SSO_DEST_KEY_RE.fullmatch(dest_key): - return False - lowered = dest_key.lower() - return not any(fragment in lowered for fragment in _CLI_SSO_SECRET_KEY_FRAGMENTS) - - -def _is_safe_cli_sso_scalar_claim_value(value: Any) -> bool: - if not isinstance(value, _CLI_SSO_SCALAR_TYPES): - return False - if isinstance(value, str): - if len(value) > CLI_SSO_CLAIM_MAX_SCALAR_LENGTH: - return False - if value.startswith("eyJ") and value.count(".") >= 2: - return False - return True - - -def _sso_result_to_dict(result: Union[CustomOpenID, OpenID, dict]) -> Dict[str, Any]: - if isinstance(result, dict): - return result - if hasattr(result, "model_dump"): - dumped = result.model_dump() - if isinstance(dumped, dict): - return cast(Dict[str, Any], dumped) - return {} - - -def _get_nested_claim_value(data: Dict[str, Any], claim_path: str) -> Any: - """Resolve a dot-notation claim path against an SSO result dict. - - Unlike ``get_nested_value``, this does not strip a leading ``metadata.`` - prefix, since OIDC claims may legitimately use ``metadata`` as a top-level - key. - """ - if not claim_path: - return None - if claim_path in data: - return data[claim_path] - placeholder = "\x00" - parts = claim_path.replace("\\.", placeholder).split(".") - parts = [p.replace(placeholder, ".") for p in parts] - current: Any = data - for part in parts: - if isinstance(current, dict) and part in current: - current = current[part] - else: - return None - return current - - -def _extract_sso_claim_value(result: Union[CustomOpenID, OpenID, dict], claim_path: str) -> Any: - extra_fields = getattr(result, "extra_fields", None) - if isinstance(extra_fields, dict): - if claim_path in extra_fields: - return extra_fields[claim_path] - nested = _get_nested_claim_value(extra_fields, claim_path) - if nested is not None: - return nested - - if isinstance(result, dict): - return _get_nested_claim_value(result, claim_path) - - result_dict = _sso_result_to_dict(result) - return _get_nested_claim_value(result_dict, claim_path) - - -def _set_nested_metadata_value(metadata: Dict[str, Any], key_path: str, value: Any) -> None: - placeholder = "\x00" - parts = key_path.replace("\\.", placeholder).split(".") - parts = [p.replace(placeholder, ".") for p in parts] - current: Any = metadata - for part in parts[:-1]: - existing = current.get(part) - if not isinstance(existing, dict): - existing = {} - current[part] = existing - current = existing - current[parts[-1]] = value - - -def _flatten_cli_sso_metadata_for_poll( - metadata: Dict[str, Any], -) -> Dict[str, Union[str, int, float, bool]]: - """Expose scalar attribution metadata as a flat dict for CLI poll responses.""" - flattened: Dict[str, Union[str, int, float, bool]] = {} - stack: List[Tuple[str, Any]] = [("", metadata)] - while stack: - prefix, value = stack.pop() - if isinstance(value, dict): - for key, nested in value.items(): - nested_prefix = f"{prefix}.{key}" if prefix else key - stack.append((nested_prefix, nested)) - elif _is_safe_cli_sso_scalar_claim_value(value): - flattened[prefix] = value - return flattened - - -def build_cli_sso_attribution_metadata( - result: Union[CustomOpenID, OpenID, dict], -) -> Dict[str, Any]: - """ - Build allowlisted, non-secret scalar attribution metadata from an SSO result. - - Sources are configured via CLI_SSO_CLAIM_MAP / LITELLM_CLI_SSO_CLAIM_MAP and - may include claims captured by GENERIC_USER_EXTRA_ATTRIBUTES on CustomOpenID. - """ - claim_map = _parse_cli_sso_claim_map() - if not claim_map: - return {} - - metadata: Dict[str, Any] = {} - for source_claim, dest_key in claim_map: - if not _is_safe_cli_sso_metadata_dest_key(dest_key): - verbose_proxy_logger.debug(f"Skipping unsafe CLI SSO metadata destination key: {dest_key}") - continue - - raw_value = _extract_sso_claim_value(result=result, claim_path=source_claim) - if not _is_safe_cli_sso_scalar_claim_value(raw_value): - continue - - _set_nested_metadata_value(metadata=metadata, key_path=dest_key, value=raw_value) - - return metadata - - -def _merge_cli_sso_attribution_metadata( - existing_metadata: Dict[str, Any], attribution_metadata: Dict[str, Any] -) -> Dict[str, Any]: - """Merge attribution metadata into existing user metadata in-place. - - Preserves original value types (in particular, string claim values that - happen to look numeric are NOT coerced to ``int``/``float``). Nested dicts - are merged iteratively so attribution claims do not clobber unrelated keys - under the same parent. - """ - pending: List[Tuple[Dict[str, Any], Dict[str, Any]]] = [(existing_metadata, attribution_metadata)] - while pending: - target, source = pending.pop() - for key, value in source.items(): - if value is None: - continue - existing_value = target.get(key) - if isinstance(value, dict) and isinstance(existing_value, dict): - pending.append((existing_value, value)) - else: - target[key] = value - return existing_metadata - - -async def _persist_cli_sso_user_metadata( - prisma_client: PrismaClient, - user_id: str, - attribution_metadata: Dict[str, Any], -) -> None: - if not attribution_metadata: - return - - try: - user_row = await UserRepository(prisma_client).table.find_unique(where={"user_id": user_id}) - existing_metadata: Dict[str, Any] = {} - if user_row is not None: - row_metadata = user_row.metadata - if isinstance(row_metadata, dict): - existing_metadata = deepcopy(row_metadata) - - merged_metadata = _merge_cli_sso_attribution_metadata( - existing_metadata=existing_metadata, - attribution_metadata=attribution_metadata, - ) - await UserRepository(prisma_client).table.update_many( - where={"user_id": user_id}, - data={"metadata": merged_metadata}, - ) - verbose_proxy_logger.info( - f"Persisted CLI SSO attribution metadata for user {user_id}: " - f"{list(_flatten_cli_sso_metadata_for_poll(attribution_metadata).keys())}" - ) - except Exception as e: - verbose_proxy_logger.error(f"Failed to persist CLI SSO attribution metadata for user {user_id}: {e}") - - -def _cli_poll_attribution_metadata_from_session( - session_data: Dict[str, Any], -) -> Dict[str, Union[str, int, float, bool]]: - stored = session_data.get("attribution_metadata") - if isinstance(stored, dict): - return _flatten_cli_sso_metadata_for_poll(stored) - return {} - - -def _render_cli_sso_verification_page( - verify_url: str, - browser_complete_token: str, - prefill_user_code: str | None = None, -) -> str: - escaped_verify_url = escape(verify_url, quote=True) - escaped_browser_complete_token = escape(browser_complete_token, quote=True) - user_code_value_attr = f' value="{escape(prefill_user_code, quote=True)}"' if prefill_user_code else "" - instructions = ( - "Confirm the verification code below to finish this login." - if prefill_user_code - else "Enter the verification code shown in your terminal to finish this login." - ) - return f""" - - - - LiteLLM CLI Login - - - -
-

Complete CLI Login

-

{instructions}

-
- - - - -
-
- - - """ - - -@router.post("/sso/cli/start", tags=["experimental"], include_in_schema=False) -async def cli_sso_start(request: Request): - from litellm.proxy.proxy_server import general_settings, user_api_key_cache - - _check_cli_sso_start_rate_limit( - request=request, - cache=user_api_key_cache, - use_x_forwarded_for=bool((general_settings or {}).get("use_x_forwarded_for", False)), - ) - - login_id = f"cli-{secrets.token_urlsafe(24)}" - poll_secret = secrets.token_urlsafe(32) - user_code = _generate_cli_sso_user_code() - - flow = { - "poll_secret_hash": _hash_cli_sso_secret(poll_secret), - "user_code_hash": _hash_cli_sso_secret(_normalize_cli_sso_user_code(user_code)), - "sso_complete": False, - "user_code_verified": False, - "session_data": None, - } - _set_cli_sso_flow(login_id=login_id, cache=user_api_key_cache, flow=flow) - - verification_uri_complete: str | None = ( - ( - get_custom_url(request_base_url=str(request.base_url), route="sso/key/generate") - + "?" - + urlencode( - { - "source": LITELLM_CLI_SOURCE_IDENTIFIER, - "key": login_id, - "user_code": user_code, - } - ) - ) - if _cli_sso_verification_uri_complete_enabled() - else None - ) - return _cli_sso_start_response_body( - login_id=login_id, - poll_secret=poll_secret, - user_code=user_code, - verification_uri_complete=verification_uri_complete, - ) - - -@router.post("/sso/cli/complete/{login_id}", tags=["experimental"], include_in_schema=False) -async def cli_sso_complete(request: Request, login_id: str): - from fastapi.responses import HTMLResponse - - from litellm.proxy.common_utils.html_forms.cli_sso_success import ( - render_cli_sso_success_page, - ) - from litellm.proxy.proxy_server import user_api_key_cache - - flow = _get_cli_sso_flow_or_raise(login_id=login_id, cache=user_api_key_cache) - if not flow.get("sso_complete") or not flow.get("session_data"): - raise HTTPException(status_code=400, detail="CLI login is not ready") - - body = (await request.body()).decode("utf-8") - form_values = parse_qs(body) - supplied_user_code = (form_values.get("user_code") or [""])[0] - supplied_browser_complete_token = (form_values.get("browser_complete_token") or [""])[0] - supplied_user_code_hash = _hash_cli_sso_secret(_normalize_cli_sso_user_code(supplied_user_code)) - supplied_browser_complete_token_hash = _hash_cli_sso_secret(supplied_browser_complete_token) - - expected_user_code_hash = flow.get("user_code_hash") - if not isinstance(expected_user_code_hash, str) or not secrets.compare_digest( - supplied_user_code_hash, expected_user_code_hash - ): - raise HTTPException(status_code=400, detail="Invalid verification code") - - expected_browser_complete_token_hash = flow.get("browser_complete_token_hash") - if not isinstance(expected_browser_complete_token_hash, str) or not secrets.compare_digest( - supplied_browser_complete_token_hash, expected_browser_complete_token_hash - ): - raise HTTPException(status_code=400, detail="Invalid verification code") - - flow["user_code_verified"] = True - _set_cli_sso_flow(login_id=login_id, cache=user_api_key_cache, flow=flow) - - html_content = render_cli_sso_success_page() - return HTMLResponse(content=html_content, status_code=200) - - -def normalize_email(email: Optional[str]) -> Optional[str]: - """ - Normalize email address to lowercase for consistent storage and comparison. - - Email addresses should be treated as case-insensitive for SSO purposes, - even though RFC 5321 technically allows case-sensitive local parts. - This prevents issues where SSO providers return emails with different casing - than what's stored in the database. - - Args: - email: Email address to normalize, can be None - - Returns: - Lowercased email address, or None if input is None - """ - if email is None: - return None - return email.lower() if isinstance(email, str) else email - - -def determine_role_from_groups( - user_groups: List[str], - role_mappings: "RoleMappings", -) -> Optional[LitellmUserRoles]: - """ - Determine the highest privilege role for a user based on their groups. - - Role hierarchy (highest to lowest): - - proxy_admin - - proxy_admin_viewer - - internal_user - - internal_user_viewer - - Args: - user_groups: List of group names from the SSO token - role_mappings: RoleMappings configuration object - - Returns: - The highest privilege role found, or default_role if no matches, or None - """ - if not role_mappings.roles: - # No role mappings configured, return default_role - return role_mappings.default_role - - # Role hierarchy (highest to lowest) - role_hierarchy = [ - LitellmUserRoles.PROXY_ADMIN, - LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, - LitellmUserRoles.INTERNAL_USER, - LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, - ] - - # Convert user_groups to a set for efficient lookup - user_groups_set = set(user_groups) if isinstance(user_groups, list) else set() - - # Find the highest privilege role the user belongs to - for role in role_hierarchy: - if role in role_mappings.roles: - role_groups = role_mappings.roles[role] - if isinstance(role_groups, list) and user_groups_set.intersection(set(role_groups)): - verbose_proxy_logger.debug( - f"User groups {user_groups} matched role '{role.value}' via groups: {role_groups}" - ) - return role - - # No matching groups found, return default_role - verbose_proxy_logger.debug( - f"User groups {user_groups} did not match any role mappings, using default_role: {role_mappings.default_role}" - ) - return role_mappings.default_role - - -def process_sso_jwt_access_token( - access_token_str: Optional[str], - sso_jwt_handler: Optional[JWTHandler], - result: Union[OpenID, dict, None], - role_mappings: Optional["RoleMappings"] = None, -) -> Optional[dict]: - """ - Process SSO JWT access token and extract team IDs and user role if available. - - This function decodes the JWT access token and extracts team IDs and user - role, then sets them on the result object. Role extraction from the access - token is needed because some SSO providers (e.g., Keycloak) do not include - role claims in the UserInfo endpoint response. - - Args: - access_token_str: The JWT access token string - sso_jwt_handler: SSO-specific JWT handler for team ID extraction - result: The SSO result object to update with team IDs and role - role_mappings: Optional role mappings configuration for group-based role determination - - Returns: - The decoded access token payload dict, or None if decoding failed or - inputs were missing. Callers can pass this to _sync_user_role_from_jwt_role_map - so it has access to custom role claims (e.g. custom_roles) that are - encoded inside the JWT but stripped from received_response. - """ - if access_token_str and result: - import jwt - - try: - access_token_payload = jwt.decode(access_token_str, options={"verify_signature": False}) - except jwt.exceptions.DecodeError: - verbose_proxy_logger.debug( - "Access token is not a valid JWT (possibly an opaque token), skipping JWT-based extraction" - ) - return None - - # Extract team IDs from access token if sso_jwt_handler is available - if sso_jwt_handler: - if isinstance(result, dict): - result_team_ids: Optional[List[str]] = result.get("team_ids", []) - if not result_team_ids: - team_ids = sso_jwt_handler.get_team_ids_from_jwt(access_token_payload) - result["team_ids"] = team_ids - else: - result_team_ids = getattr(result, "team_ids", []) if result else [] - if not result_team_ids: - team_ids = sso_jwt_handler.get_team_ids_from_jwt(access_token_payload) - setattr(result, "team_ids", team_ids) - - # Extract user role from access token if not already set from UserInfo - existing_role = result.get("user_role") if isinstance(result, dict) else getattr(result, "user_role", None) - if existing_role is None: - user_role: Optional[LitellmUserRoles] = None - - # Try role_mappings first (group-based role determination) - if role_mappings is not None and role_mappings.roles: - group_claim = role_mappings.group_claim - user_groups_raw: Any = get_nested_value(access_token_payload, group_claim) - - user_groups: List[str] = [] - if isinstance(user_groups_raw, list): - user_groups = [str(g) for g in user_groups_raw] - elif isinstance(user_groups_raw, str): - user_groups = [g.strip() for g in user_groups_raw.split(",") if g.strip()] - elif user_groups_raw is not None: - user_groups = [str(user_groups_raw)] - - if user_groups: - user_role = determine_role_from_groups(user_groups, role_mappings) - verbose_proxy_logger.debug( - f"Determined role '{user_role}' from access token groups '{user_groups}' using role_mappings" - ) - elif role_mappings.default_role: - user_role = role_mappings.default_role - - # Fallback: try GENERIC_USER_ROLE_ATTRIBUTE on the access token payload - if user_role is None: - generic_user_role_attribute_name = os.getenv("GENERIC_USER_ROLE_ATTRIBUTE", "role") - user_role_from_token = get_nested_value(access_token_payload, generic_user_role_attribute_name) - if user_role_from_token is not None: - user_role = get_litellm_user_role(user_role_from_token) - verbose_proxy_logger.debug( - f"Extracted role '{user_role}' from access token field '{generic_user_role_attribute_name}'" - ) - - if user_role is not None: - if isinstance(result, dict): - result["user_role"] = user_role - else: - setattr(result, "user_role", user_role) - verbose_proxy_logger.debug(f"Set user_role='{user_role}' from JWT access token") - - return access_token_payload - - return None - - -@router.get("/sso/key/generate", tags=["experimental"], include_in_schema=False) -async def google_login( - request: Request, - source: Optional[str] = None, - key: Optional[str] = None, - existing_key: Optional[str] = None, - return_to: Optional[str] = None, - user_code: str | None = None, -): - """ - Create Proxy API Keys using Google Workspace SSO. Requires setting PROXY_BASE_URL in .env - PROXY_BASE_URL should be the your deployed proxy endpoint, e.g. PROXY_BASE_URL="https://litellm-production-7002.up.railway.app/" - Example: - """ - from litellm.proxy.proxy_server import ( - general_settings, - premium_user, - prisma_client, - user_api_key_cache, - user_custom_ui_sso_sign_in_handler, - ) - - microsoft_client_id = os.getenv("MICROSOFT_CLIENT_ID", None) - google_client_id = os.getenv("GOOGLE_CLIENT_ID", None) - generic_client_id = os.getenv("GENERIC_CLIENT_ID", None) - - ####### Check if UI is disabled ####### - _disable_ui_flag = os.getenv("DISABLE_ADMIN_UI") - if _disable_ui_flag is not None: - is_disabled = str_to_bool(value=_disable_ui_flag) - if is_disabled: - return admin_ui_disabled() - - ####### Check if user is a Enterprise / Premium User ####### - if microsoft_client_id is not None or google_client_id is not None or generic_client_id is not None: - if premium_user is not True: - if prisma_client is None: - raise ProxyException( - message=CommonProxyErrors.db_not_connected_error.value, - type=ProxyErrorTypes.auth_error, - param="premium_user", - code=status.HTTP_403_FORBIDDEN, - ) - await _enforce_free_sso_user_limit( - prisma_client=prisma_client, - premium_user=premium_user, - block_at_limit=False, - ) - - ####### Detect DB + MASTER KEY in .env ####### - missing_env_vars = show_missing_vars_in_env() - if missing_env_vars is not None: - return missing_env_vars - - # get url from request - always use regular callback, but set state for CLI - redirect_url = SSOAuthenticationHandler.get_redirect_url_for_sso( - request=request, - sso_callback_route="sso/callback", - ) - - if source == LITELLM_CLI_SOURCE_IDENTIFIER: - _get_cli_sso_flow_or_raise(login_id=key, cache=user_api_key_cache) - - # Store CLI login handle in state for OAuth flow - cli_state: Optional[str] = SSOAuthenticationHandler._get_cli_state( - source=source, - key=key, - user_code=(user_code if _cli_sso_verification_uri_complete_enabled() else None), - ) - - # check if user defined a custom auth sso sign in handler, if yes, use it - if user_custom_ui_sso_sign_in_handler is not None: - try: - from litellm_enterprise.proxy.auth.custom_sso_handler import ( # type: ignore[import-untyped] - EnterpriseCustomSSOHandler, - ) - - return await EnterpriseCustomSSOHandler.handle_custom_ui_sso_sign_in( - request=request, - ) - except ImportError: - raise ValueError( - "Enterprise features are not available. Custom UI SSO sign-in requires LiteLLM Enterprise." - ) - - # Check if we should use SSO handler - if ( - SSOAuthenticationHandler.should_use_sso_handler( - microsoft_client_id=microsoft_client_id, - google_client_id=google_client_id, - generic_client_id=generic_client_id, - ) - is True - ): - verbose_proxy_logger.info(f"Redirecting to SSO login for {redirect_url}") - sso_redirect = await SSOAuthenticationHandler.get_sso_login_redirect( - redirect_url=redirect_url, - microsoft_client_id=microsoft_client_id, - google_client_id=google_client_id, - generic_client_id=generic_client_id, - state=cli_state, - request=request, - ) - if return_to is not None and sso_redirect is not None: - if SSOAuthenticationHandler._validate_return_to(return_to): - sso_redirect.set_cookie( - key="litellm_cp_return_to", - value=return_to, - max_age=600, - httponly=True, - samesite="lax", - ) - return sso_redirect - - from fastapi.responses import HTMLResponse - - hide_default_credentials_hint = ( - os.getenv("LITELLM_HIDE_DEFAULT_CREDENTIALS_HINT", "false").lower() == "true" - or general_settings.get("hide_default_credentials_hint", False) is True - ) - return HTMLResponse( - content=build_ui_login_form( - show_deprecation_banner=True, - hide_default_credentials_hint=hide_default_credentials_hint, - ), - status_code=200, - ) - - -def generic_response_convertor( - response, - jwt_handler: JWTHandler, - sso_jwt_handler: Optional[JWTHandler] = None, - role_mappings: Optional["RoleMappings"] = None, - team_mappings: Optional["TeamMappings"] = None, -) -> CustomOpenID: - generic_user_id_attribute_name = os.getenv("GENERIC_USER_ID_ATTRIBUTE", "preferred_username") - generic_user_display_name_attribute_name = os.getenv("GENERIC_USER_DISPLAY_NAME_ATTRIBUTE", "sub") - generic_user_email_attribute_name = os.getenv("GENERIC_USER_EMAIL_ATTRIBUTE", "email") - - generic_user_first_name_attribute_name = os.getenv("GENERIC_USER_FIRST_NAME_ATTRIBUTE", "first_name") - generic_user_last_name_attribute_name = os.getenv("GENERIC_USER_LAST_NAME_ATTRIBUTE", "last_name") - - generic_provider_attribute_name = os.getenv("GENERIC_USER_PROVIDER_ATTRIBUTE", "provider") - - generic_user_role_attribute_name = os.getenv("GENERIC_USER_ROLE_ATTRIBUTE", "role") - - generic_user_extra_attributes = os.getenv("GENERIC_USER_EXTRA_ATTRIBUTES", None) - - verbose_proxy_logger.debug( - f" generic_user_id_attribute_name: {generic_user_id_attribute_name}\n generic_user_email_attribute_name: {generic_user_email_attribute_name}" - ) - - all_teams = [] - if sso_jwt_handler is not None: - team_ids = sso_jwt_handler.get_all_jwt_team_ids(cast(dict, response)) - all_teams.extend(team_ids) - - if team_mappings is not None and team_mappings.team_ids_jwt_field is not None: - team_ids_from_db_mapping: Optional[List[str]] = get_nested_value( - data=cast(dict, response), - key_path=team_mappings.team_ids_jwt_field, - default=[], - ) - if team_ids_from_db_mapping: - all_teams.extend(team_ids_from_db_mapping) - verbose_proxy_logger.debug( - f"Loaded team_ids from DB team_mappings.team_ids_jwt_field='{team_mappings.team_ids_jwt_field}': {team_ids_from_db_mapping}" - ) - else: - team_ids = jwt_handler.get_all_jwt_team_ids(cast(dict, response)) - all_teams.extend(team_ids) - - # Determine user role based on role_mappings if available - # Only apply role_mappings for GENERIC SSO provider - user_role: Optional[LitellmUserRoles] = None - - if role_mappings is not None and role_mappings.provider.lower() in [ - "generic", - "okta", - ]: - # Use role_mappings to determine role from groups - group_claim = role_mappings.group_claim - user_groups_raw: Any = get_nested_value(response, group_claim) - - # Handle different formats: could be a list, string (comma-separated), or single value - user_groups: List[str] = [] - if isinstance(user_groups_raw, list): - user_groups = [str(g) for g in user_groups_raw] - elif isinstance(user_groups_raw, str): - # Handle comma-separated string - user_groups = [g.strip() for g in user_groups_raw.split(",") if g.strip()] - elif user_groups_raw is not None: - # Single value - user_groups = [str(user_groups_raw)] - - if user_groups: - user_role = determine_role_from_groups(user_groups, role_mappings) - verbose_proxy_logger.debug( - f"Determined role '{user_role.value if user_role else None}' from groups '{user_groups}' using role_mappings" - ) - else: - # No groups found, use default_role - user_role = role_mappings.default_role - verbose_proxy_logger.debug( - f"No groups found in '{group_claim}', using default_role: {role_mappings.default_role}" - ) - - # Fallback to existing logic if role_mappings not used - if user_role is None: - user_role_from_sso = get_nested_value(response, generic_user_role_attribute_name) - if user_role_from_sso is not None: - role = get_litellm_user_role(user_role_from_sso) - if role is not None: - user_role = role - verbose_proxy_logger.debug( - f"Found valid LitellmUserRoles '{role.value}' from SSO attribute '{generic_user_role_attribute_name}'" - ) - - # Build extra_fields dict from GENERIC_USER_EXTRA_ATTRIBUTES if specified - extra_fields: Optional[Dict[str, Any]] = None - if generic_user_extra_attributes: - extra_fields = {} - for attr_name in generic_user_extra_attributes.split(","): - attr_name = attr_name.strip() - extra_fields[attr_name] = get_nested_value(response, attr_name) - - return CustomOpenID( - id=get_nested_value(response, generic_user_id_attribute_name), - display_name=get_nested_value(response, generic_user_display_name_attribute_name), - email=normalize_email(get_nested_value(response, generic_user_email_attribute_name)), - first_name=get_nested_value(response, generic_user_first_name_attribute_name), - last_name=get_nested_value(response, generic_user_last_name_attribute_name), - provider=get_nested_value(response, generic_provider_attribute_name), - team_ids=all_teams, - user_role=user_role, - extra_fields=extra_fields, - ) - - -def _setup_generic_sso_env_vars( - generic_client_id: str, redirect_url: str -) -> Tuple[str, List[str], str, str, str, bool]: - """Setup and validate Generic SSO environment variables.""" - generic_client_secret = os.getenv("GENERIC_CLIENT_SECRET", None) - generic_scope = os.getenv("GENERIC_SCOPE", "openid email profile").split(" ") - generic_authorization_endpoint = os.getenv("GENERIC_AUTHORIZATION_ENDPOINT", None) - generic_token_endpoint = os.getenv("GENERIC_TOKEN_ENDPOINT", None) - generic_userinfo_endpoint = os.getenv("GENERIC_USERINFO_ENDPOINT", None) - generic_include_client_id = os.getenv("GENERIC_INCLUDE_CLIENT_ID", "false").lower() == "true" - - # Validate required environment variables - if generic_client_secret is None: - raise ProxyException( - message="GENERIC_CLIENT_SECRET not set. Set it in .env file", - type=ProxyErrorTypes.auth_error, - param="GENERIC_CLIENT_SECRET", - code=status.HTTP_500_INTERNAL_SERVER_ERROR, - ) - if generic_authorization_endpoint is None: - raise ProxyException( - message="GENERIC_AUTHORIZATION_ENDPOINT not set. Set it in .env file", - type=ProxyErrorTypes.auth_error, - param="GENERIC_AUTHORIZATION_ENDPOINT", - code=status.HTTP_500_INTERNAL_SERVER_ERROR, - ) - if generic_token_endpoint is None: - raise ProxyException( - message="GENERIC_TOKEN_ENDPOINT not set. Set it in .env file", - type=ProxyErrorTypes.auth_error, - param="GENERIC_TOKEN_ENDPOINT", - code=status.HTTP_500_INTERNAL_SERVER_ERROR, - ) - if generic_userinfo_endpoint is None: - raise ProxyException( - message="GENERIC_USERINFO_ENDPOINT not set. Set it in .env file", - type=ProxyErrorTypes.auth_error, - param="GENERIC_USERINFO_ENDPOINT", - code=status.HTTP_500_INTERNAL_SERVER_ERROR, - ) - - verbose_proxy_logger.debug( - f"authorization_endpoint: {generic_authorization_endpoint}\ntoken_endpoint: {generic_token_endpoint}\nuserinfo_endpoint: {generic_userinfo_endpoint}" - ) - verbose_proxy_logger.debug(f"GENERIC_REDIRECT_URI: {redirect_url}\nGENERIC_CLIENT_ID: {generic_client_id}\n") - - return ( - generic_client_secret, - generic_scope, - generic_authorization_endpoint, - generic_token_endpoint, - generic_userinfo_endpoint, - generic_include_client_id, - ) - - -async def _setup_team_mappings() -> Optional["TeamMappings"]: - """Setup team mappings from SSO database settings.""" - team_mappings: Optional["TeamMappings"] = None - try: - from litellm.proxy.utils import get_prisma_client_or_throw - - prisma_client = get_prisma_client_or_throw("Prisma client is None, connect a database to your proxy") - - sso_db_record = await SSOConfigRepository(prisma_client).table.find_unique(where={"id": "sso_config"}) - - if sso_db_record and sso_db_record.sso_settings: - sso_settings_dict = dict(sso_db_record.sso_settings) - team_mappings_data = sso_settings_dict.get("team_mappings") - - if team_mappings_data: - from litellm.types.proxy.management_endpoints.ui_sso import TeamMappings - - if isinstance(team_mappings_data, dict): - team_mappings = TeamMappings(**team_mappings_data) - elif isinstance(team_mappings_data, TeamMappings): - team_mappings = team_mappings_data - - if team_mappings and team_mappings.team_ids_jwt_field: - verbose_proxy_logger.debug( - f"Loaded team_mappings with team_ids_jwt_field: '{team_mappings.team_ids_jwt_field}'" - ) - except Exception as e: - verbose_proxy_logger.debug( - f"Could not load team_mappings from database: {e}. Continuing with config-based team mapping." - ) - - return team_mappings - - -async def _setup_role_mappings() -> Optional["RoleMappings"]: - """Setup role mappings from SSO database settings.""" - role_mappings: Optional["RoleMappings"] = None - try: - from litellm.proxy.utils import get_prisma_client_or_throw - - prisma_client = get_prisma_client_or_throw("Prisma client is None, connect a database to your proxy") - - sso_db_record = await SSOConfigRepository(prisma_client).table.find_unique(where={"id": "sso_config"}) - - if sso_db_record and sso_db_record.sso_settings: - sso_settings_dict = dict(sso_db_record.sso_settings) - role_mappings_data = sso_settings_dict.get("role_mappings") - - if role_mappings_data: - from litellm.types.proxy.management_endpoints.ui_sso import RoleMappings - - if isinstance(role_mappings_data, dict): - role_mappings = RoleMappings(**role_mappings_data) - elif isinstance(role_mappings_data, RoleMappings): - role_mappings = role_mappings_data - - if role_mappings: - verbose_proxy_logger.debug(f"Loaded role_mappings for provider '{role_mappings.provider}'") - except Exception as e: - verbose_proxy_logger.debug( - f"Could not load role_mappings from database: {e}. Continuing with existing role logic." - ) - - generic_role_mappings = os.getenv("GENERIC_ROLE_MAPPINGS_ROLES", None) - generic_role_mappings_group_claim = os.getenv("GENERIC_ROLE_MAPPINGS_GROUP_CLAIM", None) - generic_role_mappings_default_role = os.getenv("GENERIC_ROLE_MAPPINGS_DEFAULT_ROLE", None) - if generic_role_mappings is not None: - verbose_proxy_logger.debug("Found role_mappings for generic provider in environment variables") - import ast - - try: - generic_user_role_mappings_data: Dict[LitellmUserRoles, List[str]] = ast.literal_eval(generic_role_mappings) - if isinstance(generic_user_role_mappings_data, dict): - from litellm.types.proxy.management_endpoints.ui_sso import RoleMappings - - role_mappings_data = { - "provider": "generic", - "group_claim": generic_role_mappings_group_claim, - "default_role": generic_role_mappings_default_role, - "roles": generic_user_role_mappings_data, - } - - role_mappings = RoleMappings(**role_mappings_data) - verbose_proxy_logger.debug( - f"Loaded role_mappings from environments for provider '{role_mappings.provider}'." - ) - return role_mappings - except TypeError as e: - verbose_proxy_logger.warning( - f"Error decoding role mappings from environment variables: {e}. Continuing with existing role logic." - ) - return role_mappings - - -def _parse_generic_sso_headers() -> dict: - """Parse comma-separated GENERIC_SSO_HEADERS env var into a dict.""" - raw = os.getenv("GENERIC_SSO_HEADERS", None) - if raw is None: - return {} - result: Dict[str, str] = {} - for header in raw.split(","): - header = header.strip() - if header: - key, value = header.split("=") - result[key] = value - return result - - -def _handle_generic_sso_error( - e: Exception, - generic_authorization_endpoint: Optional[str], - generic_token_endpoint: Optional[str], - additional_headers: dict, -) -> NoReturn: - """Handle errors from generic SSO verify_and_process. Always re-raises.""" - error_message = str(e) - - # Surface a helpful PKCE misconfiguration hint only when: - # 1. The error mentions PKCE/code verifier, AND - # 2. PKCE is not currently configured (GENERIC_CLIENT_USE_PKCE != true) - pkce_configured = os.getenv("GENERIC_CLIENT_USE_PKCE", "false").lower() == "true" - if not pkce_configured and ("PKCE" in error_message or "code verifier" in error_message.lower()): - is_okta = (generic_authorization_endpoint and "okta" in generic_authorization_endpoint.lower()) or ( - generic_token_endpoint and "okta" in generic_token_endpoint.lower() - ) - provider_name = "Okta" if is_okta else "Your OAuth provider" - - detailed_message = ( - f"SSO authentication failed: {provider_name} requires PKCE (Proof Key for Code Exchange) " - f"but it's not enabled in your LiteLLM configuration.\n\n" - f"SOLUTION: Add this environment variable and restart your proxy:\n" - f" GENERIC_CLIENT_USE_PKCE=true\n\n" - ) - if is_okta: - detailed_message += ( - "For AWS ECS: Add the environment variable to your task definition.\n" - "For Docker: Add -e GENERIC_CLIENT_USE_PKCE=true to your docker run command.\n" - "For .env file: Add GENERIC_CLIENT_USE_PKCE=true to your .env file.\n\n" - ) - detailed_message += f"Original error: {error_message}" - - raise ProxyException( - message=detailed_message, - type=ProxyErrorTypes.auth_error, - param="GENERIC_CLIENT_USE_PKCE", - code=status.HTTP_401_UNAUTHORIZED, - ) - - if isinstance(e, ProxyException): - verbose_proxy_logger.error( - "SSO authentication failed: %s. Passed in headers: %s", - e, - additional_headers, - ) - else: - verbose_proxy_logger.exception( - "Error verifying and processing generic SSO: %s. Passed in headers: %s", - e, - additional_headers, - ) - raise e - - -async def get_generic_sso_response( - request: Request, - jwt_handler: JWTHandler, - sso_jwt_handler: Optional[JWTHandler], # sso specific jwt handler - used for restricted sso group access control - generic_client_id: str, - redirect_url: str, -) -> Tuple[Union[OpenID, dict], Optional[dict], Optional[dict]]: # (result, received_response, access_token_payload) - # make generic sso provider - from fastapi_sso.sso.base import DiscoveryDocument - from fastapi_sso.sso.generic import create_provider - - received_response: Optional[dict] = None - - # Setup environment variables - ( - generic_client_secret, - generic_scope, - generic_authorization_endpoint, - generic_token_endpoint, - generic_userinfo_endpoint, - generic_include_client_id, - ) = _setup_generic_sso_env_vars(generic_client_id, redirect_url) - - discovery = DiscoveryDocument( - authorization_endpoint=generic_authorization_endpoint, - token_endpoint=generic_token_endpoint, - userinfo_endpoint=generic_userinfo_endpoint, - ) - - role_mappings = await _setup_role_mappings() - team_mappings = await _setup_team_mappings() - - def response_convertor(response, client): - nonlocal received_response # return for user debugging - received_response = response - return generic_response_convertor( - response=response, - jwt_handler=jwt_handler, - sso_jwt_handler=sso_jwt_handler, - role_mappings=role_mappings, - team_mappings=team_mappings, - ) - - SSOProvider = create_provider( - name="oidc", - discovery_document=discovery, - response_convertor=response_convertor, - ) - generic_sso = SSOProvider( - client_id=generic_client_id, - client_secret=generic_client_secret, - redirect_uri=redirect_url, - allow_insecure_http=True, - scope=generic_scope, - ) - verbose_proxy_logger.debug("calling generic_sso.verify_and_process") - additional_generic_sso_headers_dict = _parse_generic_sso_headers() - - code_verifier: Optional[str] = None # assigned inside try; initialized for type tracking - access_token_payload: Optional[dict] = None # decoded JWT access token claims - - try: - token_exchange_params = await SSOAuthenticationHandler.prepare_token_exchange_parameters( - request=request, - generic_include_client_id=generic_include_client_id, - ) - - # Extract code_verifier (and the cache key for deferred deletion) before calling fastapi-sso - code_verifier = token_exchange_params.pop("code_verifier", None) - pkce_cache_key = token_exchange_params.pop("_pkce_cache_key", None) - - # Get authorization code from query params (only used in the PKCE path below; - # the non-PKCE path delegates to verify_and_process which handles OAuth error - # callbacks — user-denied, CSRF mismatch — internally). - authorization_code = request.query_params.get("code") - - if code_verifier: - # State-to-session-cookie binding. The non-PKCE branch below - # delegates to fastapi-sso's ``verify_and_process``, which - # performs its own session-cookie check. The PKCE branch - # bypasses that helper, so we validate the URL ``state`` - # against the ``litellm_oauth_state`` cookie set on the - # redirect response — without this an attacker can pre-mint - # a state + cached PKCE verifier and hijack a victim's auth - # code (Login-CSRF / token theft). - url_state = request.query_params.get("state") - cookie_state = request.cookies.get("litellm_oauth_state") - if not url_state or not cookie_state or not secrets.compare_digest(url_state, cookie_state): - raise ProxyException( - message=("Invalid OAuth state parameter — does not match the browser-bound state cookie."), - type=ProxyErrorTypes.auth_error, - param="state", - code=status.HTTP_400_BAD_REQUEST, - ) - if not authorization_code: - raise ProxyException( - message="Missing authorization code in callback", - type=ProxyErrorTypes.auth_error, - param="code", - code=status.HTTP_400_BAD_REQUEST, - ) - if not generic_client_id: - raise ProxyException( - message="GENERIC_CLIENT_ID must be set when PKCE is enabled", - type=ProxyErrorTypes.auth_error, - param="GENERIC_CLIENT_ID", - code=status.HTTP_401_UNAUTHORIZED, - ) - if not generic_token_endpoint: - raise ProxyException( - message="GENERIC_TOKEN_ENDPOINT must be set when PKCE is enabled", - type=ProxyErrorTypes.auth_error, - param="GENERIC_TOKEN_ENDPOINT", - code=status.HTTP_401_UNAUTHORIZED, - ) - # All guards above raise, so authorization_code is a non-empty str here. - # Use an explicit type guard rather than assert (assert is a no-op with -O). - if not isinstance(authorization_code, str): - raise ProxyException( - message="Missing authorization code in callback", - type=ProxyErrorTypes.auth_error, - param="code", - code=status.HTTP_400_BAD_REQUEST, - ) - combined_response = await SSOAuthenticationHandler._pkce_token_exchange( - authorization_code=authorization_code, - code_verifier=code_verifier, - client_id=generic_client_id, - client_secret=generic_client_secret, - token_endpoint=generic_token_endpoint, - userinfo_endpoint=generic_userinfo_endpoint, - include_client_id=generic_include_client_id, - redirect_url=redirect_url, - additional_headers=additional_generic_sso_headers_dict, - ) - # Pass the full response so custom response_convertor implementations - # can access all fields (including id_token for claim extraction). - result = response_convertor(combined_response, generic_sso) - # Strip bearer credentials from combined_response before storing in - # received_response. received_response may appear in restricted-group - # error messages — bearer tokens (access_token, id_token, refresh_token) - # must not be exposed to callers. - # Assign directly rather than relying on nonlocal mutation so that Pyright - # can track that received_response is non-None from this point on. - received_response = {k: v for k, v in combined_response.items() if k not in _OAUTH_TOKEN_FIELDS} - # In the PKCE path verify_and_process is skipped, so generic_sso.access_token - # is never set. Read the token directly from the exchange response instead so - # process_sso_jwt_access_token can extract JWT-embedded roles/teams. - access_token_str: Optional[str] = combined_response.get("access_token") - else: - result = await generic_sso.verify_and_process( - request, - params=token_exchange_params, - headers=additional_generic_sso_headers_dict, - ) - access_token_str = generic_sso.access_token - - access_token_payload = process_sso_jwt_access_token( - access_token_str, sso_jwt_handler, result, role_mappings=role_mappings - ) - # Delete the single-use PKCE verifier only after all downstream processing - # (response_convertor and process_sso_jwt_access_token) has completed - # successfully. Deleting earlier would consume the verifier on a transient - # failure, forcing the user to restart the entire OAuth flow from scratch. - if pkce_cache_key: - await SSOAuthenticationHandler._delete_pkce_verifier(pkce_cache_key) - - except Exception as e: - _handle_generic_sso_error( - e, - generic_authorization_endpoint, - generic_token_endpoint, - additional_generic_sso_headers_dict, - ) - verbose_proxy_logger.debug("generic result: %s", result) - return result or {}, received_response, access_token_payload - - -async def create_team_member_add_task(team_id, user_info): - """Create a task for adding a member to a team.""" - try: - member = Member(user_id=user_info.user_id, role="user") - team_member_add_request = TeamMemberAddRequest( - member=member, - team_id=team_id, - ) - return await team_member_add( - data=team_member_add_request, - user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), - ) - except Exception as e: - verbose_proxy_logger.debug(f"[Non-Blocking] Error trying to add sso user to db: {e}") - - -async def add_missing_team_member(user_info: Union[NewUserResponse, LiteLLM_UserTable], sso_teams: List[str]): - """ - - Get missing teams (diff b/w user_info.team_ids and sso_teams) - - Add missing user to missing teams - """ - # Handle None as empty list for new users - user_teams = user_info.teams if user_info.teams is not None else [] - missing_teams = set(sso_teams) - set(user_teams) - missing_teams_list = list(missing_teams) - tasks = [] - tasks = [create_team_member_add_task(team_id, user_info) for team_id in missing_teams_list] - - try: - await asyncio.gather(*tasks) - except Exception as e: - verbose_proxy_logger.debug(f"[Non-Blocking] Error trying to add sso user to db: {e}") - - -def get_disabled_non_admin_personal_key_creation(): - key_generation_settings = litellm.key_generation_settings - if key_generation_settings is None: - return False - personal_key_generation = key_generation_settings.get("personal_key_generation") or {} - allowed_user_roles = personal_key_generation.get("allowed_user_roles") or [] - return bool("proxy_admin" in allowed_user_roles) - - -def _free_tier_sso_user_limit_error() -> ProxyException: - return ProxyException( - message="You must be a LiteLLM Enterprise user to use SSO for more than 5 users. If you have a license please set `LITELLM_LICENSE` in your env. If you want to obtain a license meet with us here: https://enterprise.litellm.ai/demo You are seeing this error message because You set one of `MICROSOFT_CLIENT_ID`, `GOOGLE_CLIENT_ID`, or `GENERIC_CLIENT_ID` in your env. Please unset this", - type=ProxyErrorTypes.auth_error, - param="premium_user", - code=status.HTTP_403_FORBIDDEN, - ) - - -async def _enforce_free_sso_user_limit( - prisma_client: PrismaClient | None, - premium_user: bool, - block_at_limit: bool, -) -> None: - if premium_user or prisma_client is None: - return - total_users = await prisma_client.db.litellm_usertable.count() - if total_users is None: - return - if total_users > 5 or (block_at_limit and total_users >= 5): - raise _free_tier_sso_user_limit_error() - - -async def get_existing_user_info_from_db( - user_id: Optional[str], - user_email: Optional[str], - prisma_client: PrismaClient, - user_api_key_cache: UserApiKeyCache, - proxy_logging_obj: ProxyLogging, -) -> Optional[LiteLLM_UserTable]: - try: - user_info = await get_user_object( - user_id=user_id, - user_email=user_email, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - user_id_upsert=False, - parent_otel_span=None, - proxy_logging_obj=proxy_logging_obj, - sso_user_id=user_id, - ) - except Exception as e: - verbose_proxy_logger.debug(f"Error getting user object: {e}") - user_info = None - - return user_info - - -async def get_user_info_from_db( - result: Union[CustomOpenID, OpenID, dict], - prisma_client: PrismaClient, - user_api_key_cache: UserApiKeyCache, - proxy_logging_obj: ProxyLogging, - user_email: Optional[str], - user_defined_values: Optional[SSOUserDefinedValues], - alternate_user_id: Optional[str] = None, -) -> Optional[Union[LiteLLM_UserTable, NewUserResponse]]: - try: - potential_user_ids = [] - if alternate_user_id is not None: - potential_user_ids.append(alternate_user_id) - if not isinstance(result, dict): - _id = getattr(result, "id", None) - if _id is not None and isinstance(_id, str): - potential_user_ids.append(_id) - else: - _id = result.get("id", None) - if _id is not None and isinstance(_id, str): - potential_user_ids.append(_id) - - user_email = normalize_email( - getattr(result, "email", None) if not isinstance(result, dict) else result.get("email", None) - ) - - user_info: Optional[Union[LiteLLM_UserTable, NewUserResponse]] = None - - for user_id in potential_user_ids: - user_info = await get_existing_user_info_from_db( - user_id=user_id, - user_email=user_email, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - if user_info is not None: - break - - verbose_proxy_logger.debug( - f"user_info: {user_info}; litellm.default_internal_user_params: {litellm.default_internal_user_params}" - ) - - # Upsert SSO User to LiteLLM DB - user_info = await SSOAuthenticationHandler.upsert_sso_user( - result=result, - user_info=user_info, - user_email=user_email, - user_defined_values=user_defined_values, - prisma_client=prisma_client, - ) - - await SSOAuthenticationHandler.add_user_to_teams_from_sso_response( - result=result, - user_info=user_info, - ) - - return user_info - except ProxyException: - raise - except Exception as e: - verbose_proxy_logger.exception(f"[Non-Blocking] Error trying to add sso user to db: {e}") - - return None - - -def _should_use_role_from_sso_response(sso_role: Optional[str]) -> bool: - """returns true if SSO upsert should use the 'role' defined on the SSO response""" - if sso_role is None: - return False - - if not is_valid_litellm_user_role(sso_role): - verbose_proxy_logger.debug( - f"SSO role '{sso_role}' is not a valid LiteLLM user role. " - "Ignoring role from SSO response. See LitellmUserRoles enum for valid roles." - ) - return False - return True - - -def _build_sso_user_update_data( - result: Optional[Union["CustomOpenID", OpenID, dict]], - user_email: Optional[str], - user_id: Optional[str], -) -> dict: - """ - Build the update data dictionary for SSO user upsert. - - Args: - result: The SSO response containing user information - user_email: The user's email from SSO - user_id: The user's ID for logging purposes - - Returns: - dict: Update data containing user_email and optionally user_role if valid - """ - update_data: dict = {"user_email": normalize_email(user_email)} - - # Get SSO role from result and include if valid - sso_role = getattr(result, "user_role", None) - if sso_role is not None: - # Convert enum to string if needed - sso_role_str = sso_role.value if isinstance(sso_role, LitellmUserRoles) else sso_role - - # Only include if it's a valid LiteLLM role - if _should_use_role_from_sso_response(sso_role_str): - update_data["user_role"] = sso_role_str - verbose_proxy_logger.info(f"Updating user {user_id} role from SSO: {sso_role_str}") - - return update_data - - -async def _sync_user_role_from_jwt_role_map( - jwt_handler: Optional[JWTHandler], - received_response: Optional[dict], - user_info: Optional[Union[LiteLLM_UserTable, NewUserResponse]], - prisma_client: PrismaClient, - user_api_key_cache: UserApiKeyCache, - user_defined_values: Optional[SSOUserDefinedValues], -) -> None: - """ - Apply jwt_litellm_role_map during SSO login. - - When jwt_litellm_role_map is configured with sync_user_role_and_teams=True, - this ensures SSO users get the same role mapping as API/JWT users. Without - this, the SSO path falls back to INTERNAL_USER_VIEW_ONLY for roles that - don't directly match LitellmUserRoles enum values. - """ - if jwt_handler is None or received_response is None: - return - if not jwt_handler.litellm_jwtauth.sync_user_role_and_teams: - return - if not jwt_handler.litellm_jwtauth.jwt_litellm_role_map: - return - - mapped_role = jwt_handler.map_jwt_role_to_litellm_role(received_response) - if mapped_role is None: - return - - verbose_proxy_logger.info(f"SSO jwt_litellm_role_map matched role: {mapped_role.value}") - - # Update user_defined_values so downstream code uses the mapped role - if user_defined_values is not None: - user_defined_values["user_role"] = mapped_role.value - - # Update existing DB record if role differs - if user_info is not None and user_info.user_role != mapped_role.value: - await UserRepository(prisma_client).table.update( - where={"user_id": user_info.user_id}, - data={"user_role": mapped_role.value}, - ) - user_info.user_role = mapped_role.value - await user_api_key_cache.async_set_cache( - key=user_info.user_id, - value=user_info, - model_type=LiteLLM_UserTable, - ) - - -def apply_user_info_values_to_sso_user_defined_values( - user_info: Optional[Union[LiteLLM_UserTable, NewUserResponse]], - user_defined_values: Optional[SSOUserDefinedValues], -) -> Optional[SSOUserDefinedValues]: - if user_defined_values is None: - return None - if user_info is not None and user_info.user_id is not None: - user_defined_values["user_id"] = user_info.user_id - - # SSO role takes precedence - only use DB role if SSO didn't provide one - # This ensures SSO is the authoritative source for user roles - sso_role = user_defined_values.get("user_role") - db_role = user_info.user_role if user_info else None - - if _should_use_role_from_sso_response(sso_role): - # SSO provided a valid role, keep it and log that we're using it - verbose_proxy_logger.info(f"Using SSO role: {sso_role} (DB role was: {db_role})") - else: - # SSO didn't provide a valid role, fall back to DB role or default - if user_info is None or user_info.user_role is None: - user_defined_values["user_role"] = LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value - verbose_proxy_logger.debug("No SSO or DB role found, using default: INTERNAL_USER_VIEW_ONLY") - else: - user_defined_values["user_role"] = user_info.user_role - verbose_proxy_logger.debug(f"Using DB role: {user_info.user_role}") - - # Preserve the user's existing models from the database - if user_info is not None and hasattr(user_info, "models") and user_info.models: - user_defined_values["models"] = user_info.models - - return user_defined_values - - -async def check_and_update_if_proxy_admin_id(user_role: str, user_id: str, prisma_client: Optional[PrismaClient]): - """ - - Check if user role in DB is admin - - If not, update user role in DB to admin role - """ - proxy_admin_id = os.getenv("PROXY_ADMIN_ID") - if proxy_admin_id is not None and proxy_admin_id == user_id: - if user_role and user_role == LitellmUserRoles.PROXY_ADMIN.value: - return user_role - - if prisma_client: - await UserRepository(prisma_client).table.update( - where={"user_id": user_id}, - data={"user_role": LitellmUserRoles.PROXY_ADMIN.value}, - ) - - user_role = LitellmUserRoles.PROXY_ADMIN.value - - return user_role - - -@router.get("/sso/callback", tags=["experimental"], include_in_schema=False) -async def auth_callback(request: Request, state: Optional[str] = None): - """Verify login""" - verbose_proxy_logger.info(f"Starting SSO callback with state: {state}") - - # Check if this is a CLI login (state starts with our CLI prefix) - from litellm.constants import LITELLM_CLI_SESSION_TOKEN_PREFIX - from litellm.proxy._types import LiteLLM_JWTAuth - from litellm.proxy.auth.handle_jwt import JWTHandler - from litellm.proxy.proxy_server import ( - general_settings, - jwt_handler, - master_key, - prisma_client, - user_api_key_cache, - ) - - if prisma_client is None: - raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value) - - sso_jwt_handler: Optional[JWTHandler] = None - ui_access_mode = general_settings.get("ui_access_mode", None) - if ui_access_mode is not None and isinstance(ui_access_mode, dict): - sso_jwt_handler = JWTHandler() - sso_jwt_handler.update_environment( - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - litellm_jwtauth=LiteLLM_JWTAuth( - team_ids_jwt_field=general_settings.get("ui_access_mode", {}).get("sso_group_jwt_field", None), - ), - leeway=0, - ) - - microsoft_client_id = os.getenv("MICROSOFT_CLIENT_ID", None) - google_client_id = os.getenv("GOOGLE_CLIENT_ID", None) - generic_client_id = os.getenv("GENERIC_CLIENT_ID", None) - received_response: Optional[dict] = None - access_token_payload: Optional[dict] = None - # get url from request - if master_key is None: - raise ProxyException( - message="Master Key not set for Proxy. Please set Master Key to use Admin UI. Set `LITELLM_MASTER_KEY` in .env or set general_settings:master_key in config.yaml. https://docs.litellm.ai/docs/proxy/virtual_keys. If set, use `--detailed_debug` to debug issue.", - type=ProxyErrorTypes.auth_error, - param="master_key", - code=status.HTTP_500_INTERNAL_SERVER_ERROR, - ) - redirect_url = SSOAuthenticationHandler.get_redirect_url_for_sso(request=request, sso_callback_route="sso/callback") - - verbose_proxy_logger.info(f"Redirecting to {redirect_url}") - result = None - if google_client_id is not None: - result = await GoogleSSOHandler.get_google_callback_response( - request=request, - google_client_id=google_client_id, - redirect_url=redirect_url, - ) - elif microsoft_client_id is not None: - result = await MicrosoftSSOHandler.get_microsoft_callback_response( - request=request, - microsoft_client_id=microsoft_client_id, - redirect_url=redirect_url, - ) - - elif generic_client_id is not None: - ( - result, - received_response, - access_token_payload, - ) = await get_generic_sso_response( - request=request, - jwt_handler=jwt_handler, - generic_client_id=generic_client_id, - redirect_url=redirect_url, - sso_jwt_handler=sso_jwt_handler, - ) - - if result is None: - raise HTTPException( - status_code=401, - detail="Result not returned by SSO provider.", - ) - - if state and state.startswith(f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:"): - # State format: {PREFIX}:{login_id}[:{user_code}] - state_parts = state.split(":", 2) - key_id = state_parts[1] if len(state_parts) > 1 else None - prefill_user_code = state_parts[2] if len(state_parts) > 2 else None - - verbose_proxy_logger.info("CLI SSO callback detected") - return await cli_sso_callback( - request=request, - key=key_id, - prefill_user_code=prefill_user_code, - result=result, - received_response=received_response, - ) - - # Control-plane cross-origin: read return_to from cookie. - # Starlette's cookie_parser already handles RFC 2109 unquoting. - cp_return_to: Optional[str] = request.cookies.get("litellm_cp_return_to") - - return await SSOAuthenticationHandler.get_redirect_response_from_openid( - result=result, - request=request, - received_response=received_response, - generic_client_id=generic_client_id, - ui_access_mode=ui_access_mode, - access_token_payload=access_token_payload, - jwt_handler=jwt_handler, - return_to=cp_return_to, - ) - - -async def _build_cli_sso_user_defined_values( - result: Union[OpenID, dict], - parsed_openid_result: ParsedOpenIDResult, -) -> Optional[SSOUserDefinedValues]: - from litellm.proxy.proxy_server import user_custom_sso - - user_id = parsed_openid_result.get("user_id") - if user_custom_sso is not None: - if inspect.iscoroutinefunction(user_custom_sso): - return await user_custom_sso(result) # type: ignore - raise ValueError("user_custom_sso must be a coroutine function") - if user_id is None: - return None - return SSOUserDefinedValues( - models=[], - user_id=user_id, - user_email=parsed_openid_result.get("user_email"), - max_budget=litellm.max_internal_user_budget, - user_role=parsed_openid_result.get("user_role"), - budget_duration=litellm.internal_user_budget_duration, - ) - - -async def _fetch_cli_sso_team_details( - prisma_client: PrismaClient, - teams: List[str], -) -> List[Dict[str, Any]]: - team_details: List[Dict[str, Any]] = [] - try: - if teams: - prisma_teams = await TeamRepository(prisma_client).table.find_many(where={"team_id": {"in": teams}}) - for team_row in prisma_teams: - team_dict = team_row.model_dump() - team_details.append( - { - "team_id": team_dict.get("team_id"), - "team_alias": team_dict.get("team_alias"), - } - ) - except Exception as e: - verbose_proxy_logger.error(f"Error fetching team details for CLI SSO session: {e}") - return team_details - - -async def _complete_cli_sso_callback_session( - *, - request: Request, - key: str, - flow: dict, - result: Union[OpenID, dict], - parsed_openid_result: ParsedOpenIDResult, - user_defined_values: Optional[SSOUserDefinedValues], - prisma_client: PrismaClient, - user_api_key_cache: UserApiKeyCache, - proxy_logging_obj: ProxyLogging, - prefill_user_code: str | None = None, -): - from fastapi.responses import HTMLResponse - - user_id = parsed_openid_result.get("user_id") - user_email = parsed_openid_result.get("user_email") - user_info = await get_user_info_from_db( - result=result, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - user_email=user_email, - user_defined_values=user_defined_values, - alternate_user_id=user_id, - ) - if user_info is None: - raise HTTPException(status_code=500, detail="Failed to retrieve user information from SSO") - if not user_info.user_id: - raise HTTPException(status_code=500, detail="Failed to retrieve user information from SSO") - - teams: List[str] = [] - if hasattr(user_info, "teams") and user_info.teams: - teams = user_info.teams if isinstance(user_info.teams, list) else [] - - team_details = await _fetch_cli_sso_team_details(prisma_client=prisma_client, teams=teams) - attribution_metadata = build_cli_sso_attribution_metadata(result=result) - if attribution_metadata: - await _persist_cli_sso_user_metadata( - prisma_client=prisma_client, - user_id=cast(str, user_info.user_id), - attribution_metadata=attribution_metadata, - ) - - flow["session_data"] = { - "user_id": cast(str, user_info.user_id), - "user_role": user_info.user_role, - "models": user_info.models if hasattr(user_info, "models") else [], - "user_email": user_email, - "teams": teams, - "team_details": team_details, - "attribution_metadata": attribution_metadata, - } - flow["sso_complete"] = True - browser_complete_token = secrets.token_urlsafe(32) - flow["browser_complete_token_hash"] = _hash_cli_sso_secret(browser_complete_token) - _set_cli_sso_flow(login_id=key, cache=user_api_key_cache, flow=flow) - - verbose_proxy_logger.info( - f"Stored CLI SSO session for user: {user_info.user_id}, teams: {teams}, num_teams: {len(teams)}" - ) - verify_url = get_custom_url( - request_base_url=str(request.base_url), - route=f"sso/cli/complete/{key}", - ) - return HTMLResponse( - content=_render_cli_sso_verification_page( - verify_url=verify_url, - browser_complete_token=browser_complete_token, - prefill_user_code=prefill_user_code, - ), - status_code=200, - ) - - -async def cli_sso_callback( - request: Request, - key: Optional[str] = None, - result: Optional[Union[OpenID, dict]] = None, - received_response: Optional[dict] = None, - prefill_user_code: str | None = None, -): - """CLI SSO callback - stores session info for JWT generation on polling""" - verbose_proxy_logger.info("CLI SSO callback") - - from litellm.proxy.proxy_server import ( - general_settings, - prisma_client, - proxy_logging_obj, - user_api_key_cache, - ) - - flow = _get_cli_sso_flow_or_raise(login_id=key, cache=user_api_key_cache) - - if prisma_client is None: - raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value) - - if result is None: - raise HTTPException( - status_code=500, - detail="SSO authentication failed - no result returned from provider", - ) - - # After None check, cast to non-None type for type checker - result_non_none: Union[OpenID, dict] = cast(Union[OpenID, dict], result) - - try: - parsed_openid_result = SSOAuthenticationHandler._get_user_email_and_id_from_result( - result=result_non_none, - generic_client_id=os.getenv("GENERIC_CLIENT_ID", None), - ) - verbose_proxy_logger.debug(f"parsed_openid_result: {parsed_openid_result}") - user_defined_values = await _build_cli_sso_user_defined_values( - result=result_non_none, - parsed_openid_result=parsed_openid_result, - ) - - SSOAuthenticationHandler.verify_user_in_restricted_sso_group( - general_settings=general_settings, - result=result_non_none, - received_response=received_response, - ) - - return await _complete_cli_sso_callback_session( - request=request, - key=cast(str, key), - flow=flow, - result=result_non_none, - parsed_openid_result=parsed_openid_result, - user_defined_values=user_defined_values, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - prefill_user_code=prefill_user_code, - ) - except ProxyException: - raise - except HTTPException: - raise - except Exception as e: - verbose_proxy_logger.error(f"Error with CLI SSO callback: {e}") - raise HTTPException(status_code=500, detail=f"Failed to process CLI SSO: {str(e)}") - - -@router.get("/sso/cli/poll/{key_id}", tags=["experimental"], include_in_schema=False) -async def cli_poll_key( - key_id: str, - team_id: Optional[str] = None, - x_litellm_cli_poll_secret: Optional[str] = Header(default=None), -): - """ - CLI polling endpoint - retrieves session from cache and generates JWT. - - Flow: - 1. First poll (no team_id): Returns teams list without generating JWT - 2. Second poll (with team_id): Generates JWT with selected team and deletes session - - Args: - key_id: The CLI login session ID - team_id: Optional team ID to assign to the JWT. If provided, must be one of user's teams. - """ - from litellm.proxy.auth.auth_checks import ( - ExperimentalUIJWTToken, - get_team_object, - get_user_object, - ) - from litellm.proxy.proxy_server import prisma_client, user_api_key_cache - - try: - flow = _get_cli_sso_flow_or_raise(login_id=key_id, cache=user_api_key_cache) - if not _verify_cli_sso_poll_secret(flow=flow, poll_secret=x_litellm_cli_poll_secret): - raise HTTPException(status_code=403, detail="Invalid CLI polling secret") - - if not flow.get("sso_complete") or not flow.get("user_code_verified"): - return {"status": "pending"} - - session_data = flow.get("session_data") - - if isinstance(session_data, dict): - user_teams = session_data.get("teams", []) - user_team_details = session_data.get("team_details") - user_id = session_data["user_id"] - - verbose_proxy_logger.info( - f"CLI poll: user={user_id}, team_id={team_id}, user_teams={user_teams}, num_teams={len(user_teams)}" - ) - - # If no team_id provided and user has teams, return teams list for selection - # Don't generate JWT yet - let CLI select a team first. For newer - # clients we return rich team details (id + alias); older clients - # can continue to rely on the simple "teams" list. - if team_id is None and len(user_teams) > 1: - verbose_proxy_logger.info(f"Returning teams list for user {user_id} to select from: {user_teams}") - # Best-effort construction of team_details if it wasn't - # already cached for some reason. - team_details_response: Optional[List[Dict[str, Any]]] = None - if isinstance(user_team_details, list) and user_team_details: - team_details_response = user_team_details - elif user_teams: - team_details_response = [{"team_id": t, "team_alias": None} for t in user_teams] - poll_response: Dict[str, Any] = { - "status": "ready", - "user_id": user_id, - "teams": user_teams, - "team_details": team_details_response, - "requires_team_selection": True, - } - attribution_metadata = _cli_poll_attribution_metadata_from_session(session_data) - if attribution_metadata: - poll_response["attribution_metadata"] = attribution_metadata - return poll_response - - # Validate team_id if provided - if team_id is not None: - if team_id not in user_teams: - raise HTTPException( - status_code=403, - detail=f"User does not belong to team: {team_id}. Available teams: {user_teams}", - ) - else: - # If no team_id provided and user has 0 or 1 team, use first team (or None) - team_id = user_teams[0] if len(user_teams) > 0 else None - - team_alias = None - if team_id and isinstance(user_team_details, list): - team_alias = next( - (team.get("team_alias") for team in user_team_details if team.get("team_id") == team_id), - None, - ) - - user_info = LiteLLM_UserTable( - user_id=user_id, - user_role=session_data["user_role"], - models=session_data.get("models", []), - ) - - user_db_obj = await get_user_object( - user_id=user_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - user_id_upsert=False, - ) - user_budget = user_db_obj.max_budget if user_db_obj is not None else None - - team_budget: Optional[float] = None - team_budget_resolved = False - if team_id is not None: - try: - team_obj = await get_team_object( - team_id=team_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - ) - team_budget = team_obj.max_budget - team_budget_resolved = True - except Exception: - pass - - session_max_budget = ( - litellm.max_ui_session_budget - if user_budget is None and (team_id is None or (team_budget_resolved and team_budget is None)) - else None - ) - - jwt_token = ExperimentalUIJWTToken.get_cli_jwt_auth_token( - user_info=user_info, - team_id=team_id, - team_alias=team_alias, - max_budget=session_max_budget, - ) - - # Delete cache entry (single-use) - user_api_key_cache.delete_cache(key=_get_cli_sso_flow_cache_key(key_id)) - - verbose_proxy_logger.info(f"CLI JWT generated for user: {user_id}, team: {team_id}") - poll_response = { - "status": "ready", - "key": jwt_token, - "user_id": user_id, - "team_id": team_id, - "teams": user_teams, - # Echo back any team details we have so clients can - # present nicer information if needed. - "team_details": user_team_details, - } - attribution_metadata = _cli_poll_attribution_metadata_from_session(session_data) - if attribution_metadata: - poll_response["attribution_metadata"] = attribution_metadata - return poll_response - else: - return {"status": "pending"} - - except HTTPException: - raise - except Exception as e: - verbose_proxy_logger.error(f"Error polling for CLI JWT: {e}") - raise HTTPException(status_code=500, detail=f"Error checking session status: {str(e)}") - - -async def insert_sso_user( - result_openid: Optional[Union[OpenID, dict]], - user_defined_values: Optional[SSOUserDefinedValues] = None, - prisma_client: PrismaClient | None = None, -) -> NewUserResponse: - """ - Helper function to create a New User in LiteLLM DB after a successful SSO login - - Args: - result_openid (OpenID): User information in OpenID format if the login was successful. - user_defined_values (Optional[SSOUserDefinedValues], optional): LiteLLM SSOValues / fields that were read - prisma_client (Optional[PrismaClient], optional): Prisma client instance. When provided, - uses this instance directly instead of re-importing from proxy_server. - - Returns: - Tuple[str, str]: User ID and User Role - """ - verbose_proxy_logger.debug( - f"Inserting SSO user into DB. User values: {user_defined_values}" - ) - from litellm.proxy.proxy_server import premium_user - - - if result_openid is None: - raise ValueError("result_openid is None") - if isinstance(result_openid, dict): - result_openid = OpenID(**result_openid) - - if user_defined_values is None: - raise ValueError("user_defined_values is None") - - await _enforce_free_sso_user_limit( - prisma_client=prisma_client, - premium_user=premium_user, - block_at_limit=True, - ) - # Apply default_internal_user_params - if litellm.default_internal_user_params: - # Preserve the SSO-extracted role if it's a valid LiteLLM role, - # regardless of how it was determined (role_mappings, Microsoft app_roles, - # GENERIC_USER_ROLE_ATTRIBUTE, custom SSO handler, etc.) - sso_role = user_defined_values.get("user_role") - if _should_use_role_from_sso_response(sso_role): - # Preserve the SSO-extracted role, but apply other defaults - preserved_role = sso_role - user_defined_values.update(litellm.default_internal_user_params) # type: ignore - user_defined_values["user_role"] = preserved_role # Restore preserved role - verbose_proxy_logger.debug(f"Preserved SSO-extracted role '{preserved_role}'") - else: - # SSO didn't provide a valid role, apply all defaults including role - user_defined_values.update(litellm.default_internal_user_params) # type: ignore - - # Set budget for internal users - if user_defined_values.get("user_role") == LitellmUserRoles.INTERNAL_USER.value: - if user_defined_values.get("max_budget") is None: - user_defined_values["max_budget"] = litellm.max_internal_user_budget - if user_defined_values.get("budget_duration") is None: - user_defined_values["budget_duration"] = litellm.internal_user_budget_duration - - if user_defined_values["user_role"] is None: - user_defined_values["user_role"] = LitellmUserRoles.INTERNAL_USER_VIEW_ONLY - - new_user_request = NewUserRequest( - user_id=user_defined_values["user_id"], - user_email=normalize_email(user_defined_values["user_email"]), - user_role=user_defined_values["user_role"], # type: ignore - max_budget=user_defined_values["max_budget"], - budget_duration=user_defined_values["budget_duration"], - sso_user_id=user_defined_values["user_id"], - auto_create_key=False, - ) - - if result_openid and hasattr(result_openid, "provider"): - new_user_request.metadata = {"auth_provider": getattr(result_openid, "provider")} - - response = await new_user( - data=new_user_request, - user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), - ) - - return response - - -@router.get( - "/sso/get/ui_settings", - tags=["experimental"], - include_in_schema=False, - dependencies=[Depends(user_api_key_auth)], -) -async def get_ui_settings(request: Request): - from litellm.proxy.proxy_server import general_settings, proxy_state - - _proxy_base_url = os.getenv("PROXY_BASE_URL", None) - _logout_url = os.getenv("PROXY_LOGOUT_URL", None) - _api_doc_base_url = os.getenv("LITELLM_UI_API_DOC_BASE_URL", None) - _is_sso_enabled = _has_user_setup_sso() - disable_expensive_db_queries = ( - proxy_state.get_proxy_state_variable("spend_logs_row_count") > MAX_SPENDLOG_ROWS_TO_QUERY - ) - default_team_disabled = general_settings.get("default_team_disabled", False) - if "PROXY_DEFAULT_TEAM_DISABLED" in os.environ: - if os.environ["PROXY_DEFAULT_TEAM_DISABLED"].lower() == "true": - default_team_disabled = True - - return { - "PROXY_BASE_URL": _proxy_base_url, - "PROXY_LOGOUT_URL": _logout_url, - "LITELLM_UI_API_DOC_BASE_URL": _api_doc_base_url, - "DEFAULT_TEAM_DISABLED": default_team_disabled, - "SSO_ENABLED": _is_sso_enabled, - "NUM_SPEND_LOGS_ROWS": proxy_state.get_proxy_state_variable("spend_logs_row_count"), - "DISABLE_EXPENSIVE_DB_QUERIES": disable_expensive_db_queries, - } - - -@router.get( - "/sso/readiness", - tags=["experimental"], - dependencies=[Depends(user_api_key_auth)], -) -async def sso_readiness(): - """ - Health endpoint for checking SSO readiness. - Checks if the configured SSO provider has all required environment variables set in memory. - """ - microsoft_client_id = os.getenv("MICROSOFT_CLIENT_ID", None) - google_client_id = os.getenv("GOOGLE_CLIENT_ID", None) - generic_client_id = os.getenv("GENERIC_CLIENT_ID", None) - - # Determine which SSO provider is configured - configured_provider = None - if google_client_id is not None: - configured_provider = "google" - elif microsoft_client_id is not None: - configured_provider = "microsoft" - elif generic_client_id is not None: - configured_provider = "generic" - - # If no SSO is configured, return healthy (SSO is optional) - if configured_provider is None: - return { - "status": "healthy", - "sso_configured": False, - "message": "No SSO provider configured", - } - - # Check required environment variables for the configured provider - missing_vars = [] - - if configured_provider == "google": - google_client_secret = os.getenv("GOOGLE_CLIENT_SECRET", None) - if google_client_secret is None: - missing_vars.append("GOOGLE_CLIENT_SECRET") - - elif configured_provider == "microsoft": - microsoft_client_secret = os.getenv("MICROSOFT_CLIENT_SECRET", None) - microsoft_tenant = os.getenv("MICROSOFT_TENANT", None) - if microsoft_client_secret is None: - missing_vars.append("MICROSOFT_CLIENT_SECRET") - if microsoft_tenant is None: - missing_vars.append("MICROSOFT_TENANT") - - elif configured_provider == "generic": - generic_client_secret = os.getenv("GENERIC_CLIENT_SECRET", None) - generic_authorization_endpoint = os.getenv("GENERIC_AUTHORIZATION_ENDPOINT", None) - generic_token_endpoint = os.getenv("GENERIC_TOKEN_ENDPOINT", None) - generic_userinfo_endpoint = os.getenv("GENERIC_USERINFO_ENDPOINT", None) - if generic_client_secret is None: - missing_vars.append("GENERIC_CLIENT_SECRET") - if generic_authorization_endpoint is None: - missing_vars.append("GENERIC_AUTHORIZATION_ENDPOINT") - if generic_token_endpoint is None: - missing_vars.append("GENERIC_TOKEN_ENDPOINT") - if generic_userinfo_endpoint is None: - missing_vars.append("GENERIC_USERINFO_ENDPOINT") - - # If all required variables are present, return healthy - if len(missing_vars) == 0: - return { - "status": "healthy", - "sso_configured": True, - "provider": configured_provider, - "message": f"{configured_provider.capitalize()} SSO is properly configured", - } - - # If some variables are missing, return unhealthy - raise HTTPException( - status_code=503, - detail={ - "status": "unhealthy", - "sso_configured": True, - "provider": configured_provider, - "missing_environment_variables": missing_vars, - "message": f"{configured_provider.capitalize()} SSO is configured but missing required environment variables: {', '.join(missing_vars)}", - }, - ) - - -class SSOAuthenticationHandler: - """ - Handler for SSO Authentication across all SSO providers - """ - - @staticmethod - def _validate_return_to(return_to: str) -> bool: - """ - Validate that return_to matches the configured control_plane_url origin. - - Returns True if return_to is valid and should be used. - Returns False if control_plane_url is not configured (return_to is ignored). - Raises HTTPException(400) if return_to origin does not match control_plane_url origin. - """ - from litellm.proxy.proxy_server import general_settings - - control_plane_url = general_settings.get("control_plane_url") - if control_plane_url is None: - return False - - def _origin(url: str) -> tuple: - parsed = urlparse(url) - scheme = (parsed.scheme or "").lower() - hostname = (parsed.hostname or "").lower() - default_port = 443 if scheme == "https" else 80 - port = parsed.port if parsed.port is not None else default_port - return (scheme, hostname, port) - - if _origin(return_to) != _origin(control_plane_url): - raise HTTPException( - status_code=400, - detail="return_to does not match the configured control_plane_url", - ) - - return True - - @staticmethod - async def get_sso_login_redirect( - redirect_url: str, - google_client_id: Optional[str] = None, - microsoft_client_id: Optional[str] = None, - generic_client_id: Optional[str] = None, - state: Optional[str] = None, - request: Optional[Request] = None, - ) -> Optional[RedirectResponse]: - """ - Step 1. Call Get Login Redirect for the SSO provider. Send the redirect response to `redirect_url` - - Args: - redirect_url (str): The URL to redirect the user to after login - google_client_id (Optional[str], optional): The Google Client ID. Defaults to None. - microsoft_client_id (Optional[str], optional): The Microsoft Client ID. Defaults to None. - generic_client_id (Optional[str], optional): The Generic Client ID. Defaults to None. - request: Optional FastAPI request, used to drive the ``Secure`` - attribute on the ``litellm_oauth_state`` CSRF cookie. - - Returns: - RedirectResponse: The redirect response from the SSO provider. - """ - # Google SSO Auth - if google_client_id is not None: - from fastapi_sso.sso.google import GoogleSSO - - google_client_secret = os.getenv("GOOGLE_CLIENT_SECRET", None) - if google_client_secret is None: - raise ProxyException( - message="GOOGLE_CLIENT_SECRET not set. Set it in .env file", - type=ProxyErrorTypes.auth_error, - param="GOOGLE_CLIENT_SECRET", - code=status.HTTP_500_INTERNAL_SERVER_ERROR, - ) - google_sso = GoogleSSO( - client_id=google_client_id, - client_secret=google_client_secret, - redirect_uri=redirect_url, - ) - verbose_proxy_logger.info( - f"In /google-login/key/generate, \nGOOGLE_REDIRECT_URI: {redirect_url}\nGOOGLE_CLIENT_ID: {google_client_id}" - ) - with google_sso: - return await google_sso.get_login_redirect(state=state) - # Microsoft SSO Auth - elif microsoft_client_id is not None: - microsoft_client_secret = os.getenv("MICROSOFT_CLIENT_SECRET", None) - microsoft_tenant = os.getenv("MICROSOFT_TENANT", None) - if microsoft_client_secret is None: - raise ProxyException( - message="MICROSOFT_CLIENT_SECRET not set. Set it in .env file", - type=ProxyErrorTypes.auth_error, - param="MICROSOFT_CLIENT_SECRET", - code=status.HTTP_500_INTERNAL_SERVER_ERROR, - ) - microsoft_sso = CustomMicrosoftSSO( - client_id=microsoft_client_id, - client_secret=microsoft_client_secret, - tenant=microsoft_tenant, - redirect_uri=redirect_url, - allow_insecure_http=True, - ) - with microsoft_sso: - return await microsoft_sso.get_login_redirect(state=state) - elif generic_client_id is not None: - from fastapi_sso.sso.base import DiscoveryDocument - from fastapi_sso.sso.generic import create_provider - - generic_client_secret = os.getenv("GENERIC_CLIENT_SECRET", None) - generic_scope = os.getenv("GENERIC_SCOPE", "openid email profile").split(" ") - generic_authorization_endpoint = os.getenv("GENERIC_AUTHORIZATION_ENDPOINT", None) - generic_token_endpoint = os.getenv("GENERIC_TOKEN_ENDPOINT", None) - generic_userinfo_endpoint = os.getenv("GENERIC_USERINFO_ENDPOINT", None) - if generic_client_secret is None: - raise ProxyException( - message="GENERIC_CLIENT_SECRET not set. Set it in .env file", - type=ProxyErrorTypes.auth_error, - param="GENERIC_CLIENT_SECRET", - code=status.HTTP_500_INTERNAL_SERVER_ERROR, - ) - if generic_authorization_endpoint is None: - raise ProxyException( - message="GENERIC_AUTHORIZATION_ENDPOINT not set. Set it in .env file", - type=ProxyErrorTypes.auth_error, - param="GENERIC_AUTHORIZATION_ENDPOINT", - code=status.HTTP_500_INTERNAL_SERVER_ERROR, - ) - if generic_token_endpoint is None: - raise ProxyException( - message="GENERIC_TOKEN_ENDPOINT not set. Set it in .env file", - type=ProxyErrorTypes.auth_error, - param="GENERIC_TOKEN_ENDPOINT", - code=status.HTTP_500_INTERNAL_SERVER_ERROR, - ) - if generic_userinfo_endpoint is None: - raise ProxyException( - message="GENERIC_USERINFO_ENDPOINT not set. Set it in .env file", - type=ProxyErrorTypes.auth_error, - param="GENERIC_USERINFO_ENDPOINT", - code=status.HTTP_500_INTERNAL_SERVER_ERROR, - ) - verbose_proxy_logger.debug( - f"authorization_endpoint: {generic_authorization_endpoint}\ntoken_endpoint: {generic_token_endpoint}\nuserinfo_endpoint: {generic_userinfo_endpoint}" - ) - verbose_proxy_logger.debug( - f"GENERIC_REDIRECT_URI: {redirect_url}\nGENERIC_CLIENT_ID: {generic_client_id}\n" - ) - discovery = DiscoveryDocument( - authorization_endpoint=generic_authorization_endpoint, - token_endpoint=generic_token_endpoint, - userinfo_endpoint=generic_userinfo_endpoint, - ) - SSOProvider = create_provider(name="oidc", discovery_document=discovery) - generic_sso = SSOProvider( - client_id=generic_client_id, - client_secret=generic_client_secret, - redirect_uri=redirect_url, - allow_insecure_http=True, - scope=generic_scope, - ) - return await SSOAuthenticationHandler.get_generic_sso_redirect_response( - generic_sso=generic_sso, - state=state, - generic_authorization_endpoint=generic_authorization_endpoint, - request=request, - ) - raise ValueError( - "Unknown SSO provider. Please setup SSO with client IDs https://docs.litellm.ai/docs/proxy/admin_ui_sso" - ) - - @staticmethod - async def get_generic_sso_redirect_response( - generic_sso: Any, - state: Optional[str] = None, - generic_authorization_endpoint: Optional[str] = None, - request: Optional[Request] = None, - ) -> Optional[RedirectResponse]: - """ - Get the redirect response for Generic SSO - """ - from urllib.parse import parse_qs, urlencode, urlparse, urlunparse - - from litellm.proxy.proxy_server import redis_usage_cache, user_api_key_cache - - with generic_sso: - # State is bound to the caller's browser via a ``litellm_oauth_state`` - # HttpOnly cookie set on the redirect response below; the SSO - # callback validates the URL ``state`` against that cookie before - # completing the PKCE token exchange. Without this binding, an - # attacker who pre-mints a state + a cached PKCE verifier can hand - # the link to a victim and capture the resulting access token - # (Login CSRF / token theft). - ( - redirect_params, - code_verifier, - ) = SSOAuthenticationHandler._get_generic_sso_redirect_params( - state=state, - generic_authorization_endpoint=generic_authorization_endpoint, - ) - - # Separate PKCE params from state params (fastapi-sso doesn't accept code_challenge) - pkce_params = {} - state_only_params = {} - for key, value in redirect_params.items(): - if key in ("code_challenge", "code_challenge_method"): - pkce_params[key] = value - else: - state_only_params[key] = value - - # Get the redirect response from fastapi-sso with only state param - redirect_response = await generic_sso.get_login_redirect(**state_only_params) # type: ignore - - # If PKCE is enabled, add PKCE parameters to the redirect URL - if code_verifier and "state" in redirect_params: - # Store code_verifier in cache (10 min TTL). Wrap in dict for proper - # JSON serialization in Redis. Use Redis when available so callbacks - # landing on another pod can retrieve it (multi-pod SSO). - cache_key = f"pkce_verifier:{redirect_params['state']}" - if redis_usage_cache is not None: - await redis_usage_cache.async_set_cache( - key=cache_key, - value={"code_verifier": code_verifier}, - ttl=600, - ) - else: - await user_api_key_cache.async_set_cache( - key=cache_key, - value={"code_verifier": code_verifier}, - ttl=600, - ) - verbose_proxy_logger.debug("PKCE code_verifier stored in cache (TTL: 600s)") - - # Add PKCE parameters to the authorization URL - if pkce_params: - parsed_url = urlparse(str(redirect_response.headers["location"])) - query_params = parse_qs(parsed_url.query) - - # Add PKCE parameters - for key, value in pkce_params.items(): - query_params[key] = [value] - - # Reconstruct the URL with PKCE parameters - new_query = urlencode(query_params, doseq=True) - new_url = urlunparse( - ( - parsed_url.scheme, - parsed_url.netloc, - parsed_url.path, - parsed_url.params, - new_query, - parsed_url.fragment, - ) - ) - - # Update the redirect response - redirect_response.headers["location"] = new_url - - # Bind state to the user's browser session. The /callback - # handler validates the URL ``state`` against this cookie via - # ``secrets.compare_digest`` before exchanging the PKCE - # code_verifier. Only set the cookie when PKCE is in use - # (i.e. inside this ``code_verifier`` branch) so two - # concurrent SSO sessions — one PKCE, one plain — cannot - # overwrite each other's state cookie. - state_value = redirect_params.get("state") - if state_value and redirect_response is not None: - # Production-safe default: require HTTPS for the - # CSRF-protection cookie unless we can prove the - # incoming request is HTTP (local dev). Without - # ``Secure`` the cookie is sent over plain HTTP, - # letting a network observer read and replay the - # state value and bypass this protection. - secure_flag = request is None or request.url.scheme == "https" - redirect_response.set_cookie( - key="litellm_oauth_state", - value=state_value, - max_age=600, - httponly=True, - samesite="lax", - secure=secure_flag, - ) - return redirect_response - - @staticmethod - def _get_generic_sso_redirect_params( - state: Optional[str] = None, - generic_authorization_endpoint: Optional[str] = None, - ) -> Tuple[dict, Optional[str]]: - """ - Get redirect parameters for Generic SSO with proper state priority handling. - Optionally generates PKCE parameters if GENERIC_CLIENT_USE_PKCE is enabled. - - Priority order: - 1. CLI state (if provided) - 2. GENERIC_CLIENT_STATE environment variable - 3. Generated UUID (required by Okta and most OAuth providers) - - - Args: - state: Optional state parameter (e.g., CLI state) - generic_authorization_endpoint: Authorization endpoint URL - - Returns: - Tuple[dict, Optional[str]]: - - Redirect parameters for SSO login (may include PKCE params) - - code_verifier (if PKCE is enabled, None otherwise) - """ - redirect_params = {} - code_verifier: Optional[str] = None - - if state: - # CLI state takes priority - # the litellm proxy cli sends the "state" parameter to the proxy server for auth. We should maintain the state parameter for the cli if it is provided - redirect_params["state"] = state - else: - generic_client_state = os.getenv("GENERIC_CLIENT_STATE", None) - if generic_client_state: - redirect_params["state"] = generic_client_state - else: - redirect_params["state"] = uuid.uuid4().hex - - # Handle PKCE (Proof Key for Code Exchange) if enabled - # Set GENERIC_CLIENT_USE_PKCE=true to enable PKCE for enhanced OAuth security - use_pkce = os.getenv("GENERIC_CLIENT_USE_PKCE", "false").lower() == "true" - - if use_pkce: - ( - code_verifier, - code_challenge, - ) = SSOAuthenticationHandler.generate_pkce_params() - redirect_params["code_challenge"] = code_challenge - redirect_params["code_challenge_method"] = "S256" - verbose_proxy_logger.debug("PKCE enabled for authorization request") - - return redirect_params, code_verifier - - @staticmethod - def should_use_sso_handler( - google_client_id: Optional[str] = None, - microsoft_client_id: Optional[str] = None, - generic_client_id: Optional[str] = None, - ) -> bool: - if google_client_id is not None or microsoft_client_id is not None or generic_client_id is not None: - return True - return False - - @staticmethod - def get_redirect_url_for_sso( - request: Request, - sso_callback_route: str, - existing_key: Optional[str] = None, - ) -> str: - """ - Get the redirect URL for SSO - - Note: existing_key is not added to the URL to avoid changing the callback URL. - It should be passed via the state parameter instead. - """ - from litellm.proxy.utils import get_custom_url - - redirect_url = get_custom_url(request_base_url=str(request.base_url)) - if redirect_url.endswith("/"): - redirect_url += sso_callback_route - else: - redirect_url += "/" + sso_callback_route - - return redirect_url - - @staticmethod - async def upsert_sso_user( - result: Optional[Union[CustomOpenID, OpenID, dict]], - user_info: Optional[Union[NewUserResponse, LiteLLM_UserTable]], - user_email: Optional[str], - user_defined_values: Optional[SSOUserDefinedValues], - prisma_client: PrismaClient, - ): - """ - Connects the SSO Users to the User Table in LiteLLM DB - - - If user on LiteLLM DB, update the user_email and user_role (if SSO provides valid role) with the SSO values - - If user not on LiteLLM DB, insert the user into LiteLLM DB - """ - try: - if user_info is not None: - user_id = user_info.user_id - update_data = _build_sso_user_update_data( - result=result, - user_email=user_email, - user_id=user_id, - ) - - await UserRepository(prisma_client).table.update_many(where={"user_id": user_id}, data=update_data) - else: - verbose_proxy_logger.info("user not in DB, inserting user into LiteLLM DB") - # user not in DB, insert User into LiteLLM DB - user_info = await insert_sso_user( - result_openid=result, - user_defined_values=user_defined_values, - prisma_client=prisma_client, - ) - return user_info - except ProxyException: - raise - except Exception as e: - verbose_proxy_logger.exception(f"Error upserting SSO user into LiteLLM DB: {e}") - return user_info - - @staticmethod - async def add_user_to_teams_from_sso_response( - result: Optional[Union[CustomOpenID, OpenID, dict]], - user_info: Optional[Union[NewUserResponse, LiteLLM_UserTable]], - ): - """ - Adds the user as a team member to the teams specified in the SSO responses `team_ids` field - - - The `team_ids` field is populated by litellm after processing the SSO response - """ - if user_info is None: - verbose_proxy_logger.debug("User not found in LiteLLM DB, skipping team member addition") - return - sso_teams = getattr(result, "team_ids", []) - await add_missing_team_member(user_info=user_info, sso_teams=sso_teams) - - @staticmethod - def verify_user_in_restricted_sso_group( - general_settings: Dict, - result: Optional[Union[CustomOpenID, OpenID, dict]], - received_response: Optional[dict], - ) -> Literal[True]: - """ - when ui_access_mode.type == "restricted_sso_group": - - - result.team_ids should contain the restricted_sso_group - - if not, raise a ProxyException - - if so, return True - - if result.team_ids is None, return False - - if result.team_ids is an empty list, return False - - if result.team_ids is a list, return True if the restricted_sso_group is in the list, otherwise return False - """ - - ui_access_mode = cast(Optional[Union[Dict, str]], general_settings.get("ui_access_mode")) - - if ui_access_mode is None: - return True - if isinstance(ui_access_mode, str): - return True - team_ids = getattr(result, "team_ids", []) - - if ui_access_mode.get("type") == "restricted_sso_group": - restricted_sso_group = ui_access_mode.get("restricted_sso_group") - if restricted_sso_group not in team_ids: - raise ProxyException( - message=f"User is not in the restricted SSO group: {restricted_sso_group}. User groups: {team_ids}. Received SSO response: {received_response}", - type=ProxyErrorTypes.auth_error, - param="restricted_sso_group", - code=status.HTTP_403_FORBIDDEN, - ) - return True - - @staticmethod - async def create_litellm_team_from_sso_group( - litellm_team_id: str, - litellm_team_name: Optional[str] = None, - ): - """ - Creates a Litellm Team from a SSO Group ID - - Your SSO provider might have groups that should be created on LiteLLM - - Use this helper to create a Litellm Team from a SSO Group ID - - Args: - litellm_team_id (str): The ID of the Litellm Team - litellm_team_name (Optional[str]): The name of the Litellm Team - """ - from litellm.proxy.proxy_server import prisma_client - - if prisma_client is None: - raise ProxyException( - message="Prisma client not found. Set it in the proxy_server.py file", - type=ProxyErrorTypes.auth_error, - param="prisma_client", - code=status.HTTP_500_INTERNAL_SERVER_ERROR, - ) - try: - team_obj = await TeamRepository(prisma_client).table.find_first(where={"team_id": litellm_team_id}) - verbose_proxy_logger.debug(f"Team object: {team_obj}") - - # only create a new team if it doesn't exist - if team_obj: - verbose_proxy_logger.debug(f"Team already exists: {litellm_team_id} - {litellm_team_name}") - return - - team_request: NewTeamRequest = NewTeamRequest( - team_id=litellm_team_id, - team_alias=litellm_team_name, - ) - if litellm.default_team_params: - team_request = SSOAuthenticationHandler._cast_and_deepcopy_litellm_default_team_params( - default_team_params=litellm.default_team_params, - litellm_team_id=litellm_team_id, - litellm_team_name=litellm_team_name, - team_request=team_request, - ) - - await new_team( - data=team_request, - # params used for Audit Logging - http_request=Request(scope={"type": "http", "method": "POST"}), - user_api_key_dict=UserAPIKeyAuth( - token="", - key_alias=f"litellm.{MicrosoftSSOHandler.__name__}", - ), - ) - except Exception as e: - verbose_proxy_logger.exception(f"Error creating Litellm Team: {e}") - - @staticmethod - def _cast_and_deepcopy_litellm_default_team_params( - default_team_params: Union[DefaultTeamSSOParams, Dict], - team_request: NewTeamRequest, - litellm_team_id: str, - litellm_team_name: Optional[str] = None, - ) -> NewTeamRequest: - """ - Casts and deepcopies the litellm.default_team_params to a NewTeamRequest object - - - Ensures we create a new DefaultTeamSSOParams object - - Handle the case where litellm.default_team_params is a dict or a DefaultTeamSSOParams object - - Adds the litellm_team_id and litellm_team_name to the DefaultTeamSSOParams object - """ - if isinstance(default_team_params, dict): - _team_request = deepcopy(default_team_params) - _team_request["team_id"] = litellm_team_id - _team_request["team_alias"] = litellm_team_name - team_request = NewTeamRequest(**_team_request) - elif isinstance(litellm.default_team_params, DefaultTeamSSOParams): - _default_team_params = deepcopy(litellm.default_team_params) - _new_team_request = team_request.model_dump() - _new_team_request.update(_default_team_params) - team_request = NewTeamRequest(**_new_team_request) - return team_request - - @staticmethod - def _get_cli_state( - source: str | None, - key: str | None, - existing_key: str | None = None, - user_code: str | None = None, - ) -> Optional[str]: - """ - Checks the request 'source' if a cli state token was passed in - - This is used to authenticate through the CLI login flow. - - The state parameter format is: {PREFIX}:{login_id}[:{user_code}] - - The state parameter is used to pass data through the OAuth flow without changing the callback URL - - user_code is appended only for the opt-in verification_uri_complete flow so the verify page can pre-fill it - """ - from litellm.constants import ( - LITELLM_CLI_SESSION_TOKEN_PREFIX, - ) - - if source == LITELLM_CLI_SOURCE_IDENTIFIER and key: - if _is_valid_cli_sso_user_code(user_code): - return f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:{key}:{user_code}" - return f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:{key}" - else: - return None - - @staticmethod - def _get_user_email_and_id_from_result( - result: Optional[Union[OpenID, dict]], - generic_client_id: Optional[str] = None, - ) -> ParsedOpenIDResult: - """ - Gets the user email and id from the OpenID result after validating the email domain - """ - user_email: Optional[str] = normalize_email(getattr(result, "email", None)) - user_id: Optional[str] = getattr(result, "id", None) if result is not None else None - user_role: Optional[str] = None - - if user_email is not None and os.getenv("ALLOWED_EMAIL_DOMAINS") is not None: - email_domain = user_email.split("@")[1] - allowed_domains = os.getenv("ALLOWED_EMAIL_DOMAINS").split(",") # type: ignore - if email_domain not in allowed_domains: - raise HTTPException( - status_code=401, - detail={ - "message": "The email domain={}, is not an allowed email domain={}. Contact your admin to change this.".format( - email_domain, allowed_domains - ) - }, - ) - - # Extract user_role from result (works for all SSO providers) - if result is not None: - _user_role = getattr(result, "user_role", None) - if _user_role is not None: - # Convert enum to string if needed - user_role = _user_role.value if isinstance(_user_role, LitellmUserRoles) else _user_role - verbose_proxy_logger.debug(f"Extracted user_role from SSO result: {user_role}") - - # generic client id - override with custom attribute name if specified - if generic_client_id is not None and result is not None: - generic_user_role_attribute_name = os.getenv("GENERIC_USER_ROLE_ATTRIBUTE", "role") - user_id = getattr(result, "id", None) - user_email = normalize_email(getattr(result, "email", None)) - if user_role is None: - _role_from_attr = getattr(result, generic_user_role_attribute_name, None) # type: ignore - if _role_from_attr is not None: - # Convert enum to string if needed - user_role = ( - _role_from_attr.value if isinstance(_role_from_attr, LitellmUserRoles) else _role_from_attr - ) - - if user_id is None and result is not None: - _first_name = getattr(result, "first_name", "") or "" - _last_name = getattr(result, "last_name", "") or "" - user_id = _first_name + _last_name - - if user_email is not None and (user_id is None or len(user_id) == 0): - user_id = user_email - - return ParsedOpenIDResult( - user_email=user_email, - user_id=user_id, - user_role=user_role, - ) - - @staticmethod - async def get_redirect_response_from_openid( - result: Union[OpenID, dict, CustomOpenID], - request: Request, - received_response: Optional[dict] = None, - generic_client_id: Optional[str] = None, - ui_access_mode: Optional[Dict] = None, - access_token_payload: Optional[dict] = None, - jwt_handler: Optional[JWTHandler] = None, - return_to: Optional[str] = None, - ) -> RedirectResponse: - import jwt - - from litellm.proxy.proxy_server import ( - general_settings, - generate_key_helper_fn, - master_key, - premium_user, - proxy_logging_obj, - redis_usage_cache, - user_api_key_cache, - user_custom_sso, - ) - from litellm.proxy.utils import get_prisma_client_or_throw - from litellm.types.proxy.ui_sso import ReturnedUITokenObject - - prisma_client = get_prisma_client_or_throw( - "Prisma client is None, connect a database to your proxy" - ) - await _enforce_free_sso_user_limit( - prisma_client=prisma_client, - premium_user=premium_user, - block_at_limit=False, - ) - - # User is Authe'd in - generate key for the UI to access Proxy - parsed_openid_result = SSOAuthenticationHandler._get_user_email_and_id_from_result( - result=result, generic_client_id=generic_client_id - ) - user_email = parsed_openid_result.get("user_email") - user_id = parsed_openid_result.get("user_id") - user_role = parsed_openid_result.get("user_role") - verbose_proxy_logger.info(f"SSO callback result: {result}") - - user_info = None - user_id_models: List = [] - max_internal_user_budget = litellm.max_internal_user_budget - internal_user_budget_duration = litellm.internal_user_budget_duration - - # User might not be already created on first generation of key - # But if it is, we want their models preferences - default_ui_key_values: Dict[str, Any] = { - "duration": LITELLM_UI_SESSION_DURATION, - "key_max_budget": litellm.max_ui_session_budget, - "aliases": {}, - "config": {}, - "spend": 0, - "team_id": "litellm-dashboard", - } - user_defined_values: Optional[SSOUserDefinedValues] = None - - if user_custom_sso is not None: - if inspect.iscoroutinefunction(user_custom_sso): - user_defined_values = await user_custom_sso(result) # type: ignore - else: - raise ValueError("user_custom_sso must be a coroutine function") - elif user_id is not None: - user_defined_values = SSOUserDefinedValues( - models=user_id_models, - user_id=user_id, - user_email=user_email, - max_budget=max_internal_user_budget, - user_role=user_role, - budget_duration=internal_user_budget_duration, - ) - - # (IF SET) Verify user is in restricted SSO group - SSOAuthenticationHandler.verify_user_in_restricted_sso_group( - general_settings=general_settings, - result=result, - received_response=received_response, - ) - - user_info = await get_user_info_from_db( - result=result, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - user_email=user_email, - user_defined_values=user_defined_values, - alternate_user_id=user_id, - ) - - # Sync user role from JWT claims via jwt_litellm_role_map (if configured). - # This ensures SSO users get the same role mapping as API/JWT users. - # Use the decoded access_token_payload (not received_response) because - # custom role claims (e.g. custom_roles) are encoded inside the JWT - # access token, which is stripped from received_response. - await _sync_user_role_from_jwt_role_map( - jwt_handler=jwt_handler, - received_response=access_token_payload or received_response, - user_info=user_info, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - user_defined_values=user_defined_values, - ) - - user_defined_values = apply_user_info_values_to_sso_user_defined_values( - user_info=user_info, user_defined_values=user_defined_values - ) - - if user_defined_values is None: - raise Exception( - "Unable to map user identity to known values. 'user_defined_values' is None. File an issue - https://github.com/BerriAI/litellm/issues" - ) - - verbose_proxy_logger.info(f"user_defined_values for creating ui key: {user_defined_values}") - - default_ui_key_values.update(user_defined_values) - default_ui_key_values["request_type"] = "key" - response = await generate_key_helper_fn( - **default_ui_key_values, # type: ignore - table_name="key", - ) - - key = response["token"] # type: ignore - user_id = response["user_id"] # type: ignore - - user_role = user_defined_values["user_role"] or LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value - if user_id and isinstance(user_id, str): - user_role = await check_and_update_if_proxy_admin_id( - user_role=user_role, user_id=user_id, prisma_client=prisma_client - ) - - verbose_proxy_logger.debug(f"user_role: {user_role}; ui_access_mode: {ui_access_mode}") - ## CHECK IF ROLE ALLOWED TO USE PROXY ## - is_admin_only_access = check_is_admin_only_access(ui_access_mode or {}) - if is_admin_only_access: - has_access = has_admin_ui_access(user_role or "") - if not has_access: - raise HTTPException( - status_code=401, - detail={ - "error": f"User not allowed to access proxy. User role={user_role}, proxy mode={ui_access_mode}" - }, - ) - - disabled_non_admin_personal_key_creation = get_disabled_non_admin_personal_key_creation() - litellm_dashboard_ui = get_custom_url(request_base_url=str(request.base_url), route="ui/") - - if get_secret_bool("EXPERIMENTAL_UI_LOGIN"): - _user_info: Optional[LiteLLM_UserTable] = None - if user_defined_values is not None and user_defined_values["user_id"] is not None: - _user_info = LiteLLM_UserTable( - user_id=user_defined_values["user_id"], - user_role=user_defined_values["user_role"] or user_role, - models=[], - max_budget=litellm.max_ui_session_budget, - ) - if _user_info is None: - raise HTTPException( - status_code=401, - detail={"error": "User Information is required for experimental UI login"}, - ) - - key = ExperimentalUIJWTToken.get_experimental_ui_login_jwt_auth_token(_user_info) - - returned_ui_token_object = ReturnedUITokenObject( - user_id=cast(str, user_id), - key=key, - user_email=user_email, - user_role=user_role or LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value, - login_method="sso", - premium_user=premium_user, - auth_header_name=general_settings.get("litellm_key_header_name", "Authorization"), - disabled_non_admin_personal_key_creation=disabled_non_admin_personal_key_creation, - server_root_path=get_server_root_path(), - ) - - jwt_token = jwt.encode( - cast(dict, returned_ui_token_object), - master_key or "", - algorithm="HS256", - ) - - # Control-plane cross-origin: store JWT behind a single-use opaque - # code (60s TTL) so the token never appears in browser history / logs. - # The control plane redeems it via POST /v3/login/exchange. - if return_to is not None and SSOAuthenticationHandler._validate_return_to(return_to): - code = secrets.token_urlsafe(32) - cache_key = f"login_code:{code}" - cache_value = {"token": jwt_token, "redirect_url": return_to} - if redis_usage_cache is not None: - await redis_usage_cache.async_set_cache(key=cache_key, value=cache_value, ttl=60) - else: - await user_api_key_cache.async_set_cache(key=cache_key, value=cache_value, ttl=60) - - separator = "&" if "?" in return_to else "?" - redirect_url = return_to + separator + urlencode({"login": "success", "code": code}) - verbose_proxy_logger.info("Cross-origin SSO: redirecting to control plane with login code") - redirect_response = RedirectResponse(url=redirect_url, status_code=303) - redirect_response.delete_cookie("litellm_cp_return_to") - return redirect_response - - if user_id is not None and isinstance(user_id, str): - litellm_dashboard_ui += "?login=success" - verbose_proxy_logger.info(f"Redirecting to {litellm_dashboard_ui}") - redirect_response = RedirectResponse(url=litellm_dashboard_ui, status_code=303) - redirect_response.set_cookie(key="token", value=jwt_token) - return redirect_response - - @staticmethod - async def prepare_token_exchange_parameters( - request: Request, - generic_include_client_id: bool, - ) -> dict: - """ - Prepare token exchange parameters for Generic SSO. - - Args: - request: Request object - generic_include_client_id: Generic OAuth Client ID - - Returns: - dict: Token exchange parameters - """ - # Prepare token exchange parameters (may add code_verifier: str later) - token_params: Dict[str, Any] = {"include_client_id": generic_include_client_id} - - # Retrieve PKCE code_verifier if PKCE was used in authorization. - # Gate on GENERIC_CLIENT_USE_PKCE to avoid an unnecessary Redis round-trip - # on every non-PKCE SSO callback. - query_params = dict(request.query_params) - state = query_params.get("state") - - use_pkce = os.getenv("GENERIC_CLIENT_USE_PKCE", "false").lower() == "true" - - if use_pkce and not state: - verbose_proxy_logger.warning( - "PKCE is enabled (GENERIC_CLIENT_USE_PKCE=true) but no 'state' parameter " - "was found in the callback. The PKCE verifier cannot be retrieved without " - "a state value — the token exchange will proceed without code_verifier, " - "which the provider may reject. Ensure your OAuth provider returns 'state' " - "in the callback redirect." - ) - - if state and use_pkce: - from litellm.proxy.proxy_server import redis_usage_cache, user_api_key_cache - - cache_key = f"pkce_verifier:{state}" - if redis_usage_cache is not None: - cached_data = await redis_usage_cache.async_get_cache(key=cache_key) - else: - cached_data = await user_api_key_cache.async_get_cache(key=cache_key) - - code_verifier = None - # Track why code_verifier is absent for accurate strict-mode diagnostics. - _empty_value_in_dict = False # dict format correct but value is empty/null - - if cached_data: - # Extract code_verifier from dict (stored as dict for JSON serialization) - if isinstance(cached_data, dict) and "code_verifier" in cached_data: - code_verifier = cached_data["code_verifier"] - if not code_verifier: - # Dict format is correct but value is empty or null. This is - # a distinct case from an unrecognized format — the entry exists - # but was stored with an empty/null verifier (data integrity issue). - _empty_value_in_dict = True - verbose_proxy_logger.warning( - "PKCE verifier dict for state '%s' has an empty/null code_verifier " - "value — may indicate a storage bug. Treating as a cache miss.", - state, - ) - else: - verbose_proxy_logger.debug("PKCE code_verifier retrieved from cache") - elif isinstance(cached_data, str): - # Handle legacy format (plain string) for backward compatibility - code_verifier = cached_data - verbose_proxy_logger.warning( - "Retrieved code_verifier in legacy plain-string format. Future storage will use dict format." - ) - else: - # Defer the detailed ERROR log to the strict-mode branch below - # (which includes state and a diagnostic message). Log at DEBUG - # here to avoid duplicate ERROR entries in the same request. - verbose_proxy_logger.debug( - "Unexpected PKCE verifier cache format (type=%s); skipping.", - type(cached_data).__name__, - ) - - if code_verifier: - # Add code_verifier to token exchange parameters. - token_params["code_verifier"] = code_verifier - # Return the cache key so the caller can delete it *after* a - # successful token exchange (avoids losing the verifier on retry - # if the exchange fails partway through). - token_params["_pkce_cache_key"] = cache_key - else: - await SSOAuthenticationHandler._handle_missing_pkce_verifier( - state=state, - cache_key=cache_key, - cached_data=cached_data, - empty_value_in_dict=_empty_value_in_dict, - redis_usage_cache=redis_usage_cache, - user_api_key_cache=user_api_key_cache, - ) - return token_params - - @staticmethod - async def _handle_missing_pkce_verifier( - state: Optional[str], - cache_key: str, - cached_data: object, - empty_value_in_dict: bool, - redis_usage_cache: object, - user_api_key_cache: object, - ) -> None: - """Handle the case where PKCE verifier could not be extracted from cache. - - In strict mode (PKCE_STRICT_CACHE_MISS=true) raises ProxyException. - Otherwise logs a warning and returns (token exchange proceeds without verifier). - """ - active_cache = redis_usage_cache if redis_usage_cache is not None else user_api_key_cache - strict_cache_miss = os.getenv("PKCE_STRICT_CACHE_MISS", "false").lower() == "true" - if strict_cache_miss: - if empty_value_in_dict: - await SSOAuthenticationHandler._delete_pkce_verifier(cache_key) - raise ProxyException( - message=( - f"PKCE verifier for state '{state}' was found in cache but " - f"has an empty or null code_verifier value — possible storage bug." - ), - type=ProxyErrorTypes.auth_error, - param="PKCE_CACHE_MISS", - code=status.HTTP_401_UNAUTHORIZED, - ) - elif cached_data is not None: - await SSOAuthenticationHandler._delete_pkce_verifier(cache_key) - verbose_proxy_logger.error( - "PKCE verifier for state '%s' has an unrecognized format (type=%s); " - "treating as a cache miss. Investigate the cached value — it may be " - "a corrupt or stale entry.", - state, - type(cached_data).__name__, - ) - raise ProxyException( - message=( - f"PKCE verifier for state '{state}' has an unrecognized format " - f"(type={type(cached_data).__name__}). The cached entry may be corrupt." - ), - type=ProxyErrorTypes.auth_error, - param="PKCE_CACHE_MISS", - code=status.HTTP_401_UNAUTHORIZED, - ) - else: - if redis_usage_cache is not None: - cause = ( - "The authorization and callback were likely handled by different " - "instances — the verifier was stored on one pod but not found on another." - ) - else: - cause = ( - "The verifier may have expired (TTL), been lost on a pod restart, " - "or the PKCE authorization step was never completed. " - "Configure Redis so all proxy instances share the PKCE verifier." - ) - verbose_proxy_logger.error( - "PKCE is enabled but no verifier found in cache for state '%s'. %s Cache type: %s.", - state, - cause, - type(active_cache).__name__, - ) - raise ProxyException( - message=f"PKCE verifier not found in cache for state '{state}'. {cause}", - type=ProxyErrorTypes.auth_error, - param="PKCE_CACHE_MISS", - code=status.HTTP_401_UNAUTHORIZED, - ) - else: - if cached_data is not None: - await SSOAuthenticationHandler._delete_pkce_verifier(cache_key) - verbose_proxy_logger.warning( - "PKCE is enabled but verifier not found in cache for state '%s' " - "(cache type: %s, raw data present: %s). " - "Continuing without code_verifier — set PKCE_STRICT_CACHE_MISS=true to fail fast instead.", - state, - type(active_cache).__name__, - cached_data is not None, - ) - - @staticmethod - async def _delete_pkce_verifier(cache_key: str) -> None: - """Delete a single-use PKCE verifier from cache after a successful exchange. - - Failure is non-fatal: a leftover verifier is a minor security concern - (unused key in cache) but not worth aborting an otherwise-successful login. - """ - from litellm.proxy.proxy_server import redis_usage_cache, user_api_key_cache - - try: - if redis_usage_cache is not None: - await redis_usage_cache.async_delete_cache(key=cache_key) - else: - await user_api_key_cache.async_delete_cache(key=cache_key) - except Exception as exc: - verbose_proxy_logger.warning( - "PKCE: failed to delete verifier cache key '%s' (best-effort cleanup): %s", - cache_key, - exc, - ) - - @staticmethod - def generate_pkce_params() -> Tuple[str, str]: - """ - Generate PKCE (Proof Key for Code Exchange) parameters for OAuth 2.0. - - Returns: - Tuple[str, str]: (code_verifier, code_challenge) - - code_verifier: Random 43-128 character string (we use 43 for efficiency) - - code_challenge: Base64-URL-encoded SHA256 hash of the code_verifier - - Reference: https://datatracker.ietf.org/doc/html/rfc7636 - """ - # Generate a cryptographically random code_verifier (43 characters) - # Using 32 random bytes which becomes 43 characters when base64-url-encoded - code_verifier = base64.urlsafe_b64encode(secrets.token_bytes(32)).decode("utf-8").rstrip("=") - - # Generate code_challenge using S256 method (SHA256) - code_challenge_bytes = hashlib.sha256(code_verifier.encode("utf-8")).digest() - code_challenge = base64.urlsafe_b64encode(code_challenge_bytes).decode("utf-8").rstrip("=") - - return code_verifier, code_challenge - - @staticmethod - def _validate_token_response(response: "httpx.Response") -> dict: - """ - Parse and validate the token endpoint response. - - Ensures the response is valid JSON, a dict, and contains a non-null - access_token string. Raises ProxyException on any validation failure. - """ - try: - token_response_raw = response.json() - except Exception as json_err: - verbose_proxy_logger.error( - "Failed to parse token response as JSON: %s. Body: %s", - json_err, - response.text[:500], - ) - raise ProxyException( - message=f"Token endpoint returned invalid JSON: {json_err}", - type=ProxyErrorTypes.auth_error, - param="token_exchange", - code=status.HTTP_401_UNAUTHORIZED, - ) - - if not isinstance(token_response_raw, dict): - verbose_proxy_logger.error( - "Token endpoint returned non-dict JSON (type=%s). Body: %s", - type(token_response_raw).__name__, - response.text[:500], - ) - raise ProxyException( - message=( - f"Token endpoint returned unexpected response format " - f"(expected JSON object, got {type(token_response_raw).__name__})" - ), - type=ProxyErrorTypes.auth_error, - param="token_exchange", - code=status.HTTP_401_UNAUTHORIZED, - ) - token_response: dict = token_response_raw - - access_token_val = token_response.get("access_token") - if not isinstance(access_token_val, str) or not access_token_val: - error = token_response.get("error") - error_desc = token_response.get("error_description", "") - if error: - detail = f"{error} - {error_desc}" if error_desc else error - else: - detail = ( - "token endpoint returned HTTP 200 but no access_token " - f"(response keys: {sorted(token_response.keys())})" - ) - verbose_proxy_logger.error("Token response missing or null access_token. detail=%s", detail) - raise ProxyException( - message=f"Token exchange failed: {detail}", - type=ProxyErrorTypes.auth_error, - param="token_exchange", - code=status.HTTP_401_UNAUTHORIZED, - ) - - return token_response - - @staticmethod - async def _pkce_token_exchange( - authorization_code: str, - code_verifier: str, - client_id: str, - client_secret: Optional[str], - token_endpoint: str, - userinfo_endpoint: Optional[str], - include_client_id: bool, - redirect_url: Optional[str], - additional_headers: Dict[str, str], - ) -> dict: - """ - Performs a direct OAuth token exchange including the PKCE code_verifier. - - fastapi-sso does not forward code_verifier, so when PKCE is enabled we - bypass it and call the token endpoint ourselves, then fetch user info. - - Returns a combined dict of the token response and user info, suitable - for passing to a response_convertor. - """ - verbose_proxy_logger.debug( - "PKCE: performing direct token exchange (code_verifier length=%d)", - len(code_verifier), - ) - - token_data: Dict[str, str] = { - "grant_type": "authorization_code", - "code": authorization_code, - "code_verifier": code_verifier, - } - # Only include redirect_uri when set — omitting it avoids sending the - # literal string "None" to the provider if the env var is missing. - if redirect_url: - token_data["redirect_uri"] = redirect_url - - request_headers = { - **additional_headers, - "Content-Type": "application/x-www-form-urlencoded", # must not be overridden - "Accept": "application/json", - } - - if not include_client_id: - # Use Basic Auth only when a secret is available; public PKCE clients omit it. - if client_secret: - credentials = base64.b64encode(f"{client_id}:{client_secret}".encode()).decode() - request_headers["Authorization"] = f"Basic {credentials}" - else: - token_data["client_id"] = client_id - else: - token_data["client_id"] = client_id - if client_secret: - token_data["client_secret"] = client_secret - - http_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.SSO_HANDLER) - try: - response = await http_client.post( - url=token_endpoint, - data=token_data, - headers=request_headers, - timeout=30.0, - ) - except Exception as exc: - # Catch network-level errors (SSL, DNS, TCP, timeout, etc.) and - # wrap them as a clean ProxyException rather than leaking raw - # httpx or OS exceptions to callers. - verbose_proxy_logger.error("PKCE token endpoint unreachable: %s", exc) - raise ProxyException( - message=f"Token endpoint request failed: {exc}", - type=ProxyErrorTypes.auth_error, - param="token_exchange", - code=status.HTTP_401_UNAUTHORIZED, - ) from exc - if response.status_code != 200: - verbose_proxy_logger.error( - "PKCE token exchange failed. status=%s body=%s", - response.status_code, - response.text[:500], - ) - raise ProxyException( - message=f"Token exchange failed: {response.status_code} - {response.text[:500]}", - type=ProxyErrorTypes.auth_error, - param="token_exchange", - code=status.HTTP_401_UNAUTHORIZED, - ) - - token_response = SSOAuthenticationHandler._validate_token_response(response) - - verbose_proxy_logger.debug( - "PKCE token exchange successful. id_token_present=%s", - bool(token_response.get("id_token")), - ) - # Bearer credentials (access_token, id_token, refresh_token) are always sourced - # from token_response — not from userinfo — in the merge step below. - userinfo = await SSOAuthenticationHandler._get_pkce_userinfo( - access_token=token_response["access_token"], - id_token=token_response.get("id_token"), - userinfo_endpoint=userinfo_endpoint, - additional_headers=additional_headers, - ) - - # Merge: userinfo takes precedence for identity claims (sub, email, name, …) per - # the OpenID Connect spec (userinfo is the authoritative source for identity). - # Bearer credentials (access_token, id_token, refresh_token) from the token endpoint - # take precedence over same-named fields in userinfo — non-standard providers sometimes - # include token fields in userinfo, which must not shadow the real bearer token. - # If a bearer field is absent from the token response, any userinfo-provided value - # is preserved as a fallback (useful for non-standard providers that omit id_token - # from the token response but include it in userinfo). - # - # Three-way merge semantics for each bearer-credential field: - # 1. token_response has a non-null value → use it (token endpoint is authoritative) - # 2. token_response explicitly sent null → remove the key so callers get a clean - # absence signal; the null from the token endpoint overrides userinfo too - # 3. field absent from token_response → leave whatever userinfo provided as-is - # (e.g. userinfo-provided id_token from a non-standard provider) - merged = {**token_response, **userinfo} - for field in _OAUTH_TOKEN_FIELDS: - if token_response.get(field) is not None: - # Case 1: non-null in token_response — restore authoritative value. - merged[field] = token_response[field] - elif field in token_response: - # Case 2: key exists but value is explicitly null — remove from merged. - merged.pop(field, None) - # Case 3: field absent from token_response — leave userinfo value as-is. - return merged - - @staticmethod - async def _get_pkce_userinfo( - access_token: str, - id_token: Optional[str], - userinfo_endpoint: Optional[str], - additional_headers: Dict[str, str], - ) -> dict: - """ - Fetches user info from the userinfo endpoint. - Falls back to decoding the id_token if the endpoint is unavailable. - """ - # None = request not yet attempted, failed, or returned empty/null (treated as failure - # so the id_token fallback can be attempted instead of returning a session with no claims). - userinfo: Optional[dict] = None - - if userinfo_endpoint: - try: - client = get_async_httpx_client(llm_provider=httpxSpecialProvider.SSO_HANDLER) - resp = await client.get( - url=userinfo_endpoint, - headers={ - **additional_headers, - "Authorization": f"Bearer {access_token}", # must not be overridden - }, - ) - if resp.status_code == 200: - try: - userinfo_raw = resp.json() - if not userinfo_raw: - # JSON null (None) or empty dict ({}) — no identity claims. - # Treat as failure so id_token fallback can be attempted. - verbose_proxy_logger.warning( - "Userinfo endpoint returned an empty or null response " - "(type=%s); treating as failure and attempting id_token fallback. " - "Check your provider's userinfo endpoint configuration.", - type(userinfo_raw).__name__, - ) - userinfo = None - else: - userinfo = userinfo_raw - except Exception as json_err: - verbose_proxy_logger.warning( - "Userinfo endpoint returned non-JSON response (status 200): %s", - json_err, - ) - else: - verbose_proxy_logger.warning( - "Userinfo endpoint returned %s (body: %s), falling back to id_token", - resp.status_code, - resp.text[:500], - ) - except Exception as e: - verbose_proxy_logger.warning("Userinfo endpoint error: %s, falling back to id_token", e) - - # Only fall back to id_token when the userinfo request failed (None). - # Empty dict ({}) and JSON null are both treated as failure (set to None above) since - # they contain no identity claims — id_token fallback is attempted in that case too. - # Explicitly check for a non-empty string to avoid attempting JWT decode on - # a blank or non-string id_token field from a misbehaving provider. - if userinfo is None and isinstance(id_token, str) and id_token: - try: - userinfo = jwt.decode(id_token, options={"verify_signature": False}) - if not userinfo: - # jwt.decode returned an empty dict (payload-free JWT or provider bug). - # Treat this the same as a missing userinfo — the session would have no - # identity claims, which is equivalent to a broken session. - verbose_proxy_logger.warning("id_token decoded to an empty payload — treating as failure.") - userinfo = None - except Exception as decode_err: - verbose_proxy_logger.error("Failed to decode id_token: %s", decode_err) - raise ProxyException( - message=f"Failed to decode id_token JWT: {decode_err}", - type=ProxyErrorTypes.auth_error, - param="userinfo", - code=status.HTTP_401_UNAUTHORIZED, - ) - - if userinfo is None: - id_token_attempted = isinstance(id_token, str) and bool(id_token) - if userinfo_endpoint: - if id_token_attempted: - detail = ( - "userinfo endpoint failed and id_token was present but " - "decoded to an empty payload — no identity claims available" - ) - else: - detail = "userinfo endpoint failed and no id_token was present in the token response" - else: - if id_token_attempted: - detail = ( - "no userinfo endpoint is configured (GENERIC_USERINFO_ENDPOINT) " - "and id_token decoded to an empty payload — no identity claims available" - ) - else: - detail = ( - "no userinfo endpoint is configured (GENERIC_USERINFO_ENDPOINT) and no id_token was present" - ) - raise ProxyException( - message=f"SSO user info unavailable: {detail}.", - type=ProxyErrorTypes.auth_error, - param="userinfo", - code=status.HTTP_401_UNAUTHORIZED, - ) - - return userinfo - - -class MicrosoftSSOHandler: - """ - Handles Microsoft SSO callback response and returns a CustomOpenID object - """ - - graph_api_base_url = "https://graph.microsoft.com/v1.0" - graph_api_user_groups_endpoint = f"{graph_api_base_url}/me/memberOf" - - """ - Constants - """ - MAX_GRAPH_API_PAGES = 200 - - # used for debugging to show the user groups litellm found from Graph API - GRAPH_API_RESPONSE_KEY = "graph_api_user_groups" - - @staticmethod - async def get_microsoft_callback_response( - request: Request, - microsoft_client_id: str, - redirect_url: str, - return_raw_sso_response: bool = False, - ) -> Union[CustomOpenID, OpenID, dict]: - """ - Get the Microsoft SSO callback response - - Args: - return_raw_sso_response: If True, return the raw SSO response - """ - microsoft_client_secret = os.getenv("MICROSOFT_CLIENT_SECRET", None) - microsoft_tenant = os.getenv("MICROSOFT_TENANT", None) - if microsoft_client_secret is None: - raise ProxyException( - message="MICROSOFT_CLIENT_SECRET not set. Set it in .env file", - type=ProxyErrorTypes.auth_error, - param="MICROSOFT_CLIENT_SECRET", - code=status.HTTP_500_INTERNAL_SERVER_ERROR, - ) - if microsoft_tenant is None: - raise ProxyException( - message="MICROSOFT_TENANT not set. Set it in .env file", - type=ProxyErrorTypes.auth_error, - param="MICROSOFT_TENANT", - code=status.HTTP_500_INTERNAL_SERVER_ERROR, - ) - microsoft_sso = CustomMicrosoftSSO( - client_id=microsoft_client_id, - client_secret=microsoft_client_secret, - tenant=microsoft_tenant, - redirect_uri=redirect_url, - allow_insecure_http=True, - ) - original_msft_result = ( - await microsoft_sso.verify_and_process( - request=request, - convert_response=False, # type: ignore - ) - or {} - ) - - user_team_ids = await MicrosoftSSOHandler.get_user_groups_from_graph_api( - access_token=microsoft_sso.access_token - ) - - # Extract app roles from the id_token JWT - app_roles = MicrosoftSSOHandler.get_app_roles_from_id_token(id_token=microsoft_sso.id_token) - verbose_proxy_logger.debug(f"Extracted app roles from id_token: {app_roles}") - - # Combine groups and app roles - user_role: Optional[LitellmUserRoles] = None - if app_roles: - # Check if any app role is a valid LitellmUserRoles - for role_str in app_roles: - role = get_litellm_user_role(role_str) - if role is not None: - user_role = role - verbose_proxy_logger.debug(f"Found valid LitellmUserRoles '{role.value}' in app_roles") - break - - verbose_proxy_logger.debug(f"Combined team_ids (groups + app roles): {user_team_ids}") - - # if user is trying to get the raw sso response for debugging, return the raw sso response - if return_raw_sso_response: - original_msft_result[MicrosoftSSOHandler.GRAPH_API_RESPONSE_KEY] = user_team_ids - original_msft_result["app_roles"] = app_roles - return original_msft_result or {} - - result = MicrosoftSSOHandler.openid_from_response( - response=original_msft_result, - team_ids=user_team_ids, - user_role=user_role, - ) - return result - - @staticmethod - def openid_from_response( - response: Optional[dict], - team_ids: List[str], - user_role: Optional[LitellmUserRoles], - ) -> CustomOpenID: - response = response or {} - verbose_proxy_logger.debug(f"Microsoft SSO Callback Response: {response}") - openid_response = CustomOpenID( - email=normalize_email(response.get(MICROSOFT_USER_EMAIL_ATTRIBUTE) or response.get("mail")), - display_name=response.get(MICROSOFT_USER_DISPLAY_NAME_ATTRIBUTE), - provider="microsoft", - id=response.get(MICROSOFT_USER_ID_ATTRIBUTE), - first_name=response.get(MICROSOFT_USER_FIRST_NAME_ATTRIBUTE), - last_name=response.get(MICROSOFT_USER_LAST_NAME_ATTRIBUTE), - team_ids=team_ids, - user_role=user_role, - ) - verbose_proxy_logger.debug(f"Microsoft SSO OpenID Response: {openid_response}") - return openid_response - - @staticmethod - def get_app_roles_from_id_token(id_token: Optional[str]) -> List[str]: - """ - Extract app roles from the Microsoft Entra ID (Azure AD) id_token JWT. - - App roles are assigned in the Azure AD Enterprise Application and appear - in the 'app_roles' claim of the id_token. - - Args: - id_token (Optional[str]): The JWT id_token from Microsoft SSO - - Returns: - List[str]: List of app role names assigned to the user - """ - if not id_token: - verbose_proxy_logger.debug("No id_token provided for app role extraction") - return [] - - try: - import jwt - - # Decode the JWT without signature verification - # (signature is already verified by fastapi_sso) - decoded_token = jwt.decode(id_token, options={"verify_signature": False}) - - # Extract app_roles claim from the token - ## check for both 'roles' and 'app_roles' claims - roles = decoded_token.get("app_roles", []) or decoded_token.get("roles", []) - - if roles and isinstance(roles, list): - verbose_proxy_logger.debug(f"Found {len(roles)} app role(s) in id_token: {roles}") - return roles - else: - verbose_proxy_logger.debug("No app roles found in id_token or roles claim is not a list") - return [] - - except Exception as e: - verbose_proxy_logger.error(f"Error extracting app roles from id_token: {e}") - return [] - - @staticmethod - async def get_user_groups_from_graph_api( - access_token: Optional[str] = None, - ) -> List[str]: - """ - Returns a list of `team_ids` the user belongs to from the Microsoft Graph API - - Args: - access_token (Optional[str]): Microsoft Graph API access token - - Returns: - List[str]: List of group IDs the user belongs to - """ - try: - async_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.SSO_HANDLER) - - # Handle MSFT Enterprise Application Groups - service_principal_id = os.getenv("MICROSOFT_SERVICE_PRINCIPAL_ID", None) - service_principal_group_ids: Optional[List[str]] = [] - service_principal_teams: Optional[List[MicrosoftServicePrincipalTeam]] = [] - if service_principal_id: - ( - service_principal_group_ids, - service_principal_teams, - ) = await MicrosoftSSOHandler.get_group_ids_from_service_principal( - service_principal_id=service_principal_id, - async_client=async_client, - access_token=access_token, - ) - verbose_proxy_logger.debug(f"Service principal group IDs: {service_principal_group_ids}") - if len(service_principal_group_ids) > 0: - await MicrosoftSSOHandler.create_litellm_teams_from_service_principal_team_ids( - service_principal_teams=service_principal_teams, - ) - - # Fetch user membership from Microsoft Graph API - all_group_ids = [] - next_link: Optional[str] = MicrosoftSSOHandler.graph_api_user_groups_endpoint - auth_headers = {"Authorization": f"Bearer {access_token}"} - page_count = 0 - - while next_link is not None and page_count < MicrosoftSSOHandler.MAX_GRAPH_API_PAGES: - group_ids, next_link = await MicrosoftSSOHandler.fetch_and_parse_groups( - url=next_link, headers=auth_headers, async_client=async_client - ) - all_group_ids.extend(group_ids) - page_count += 1 - - if next_link is not None and page_count >= MicrosoftSSOHandler.MAX_GRAPH_API_PAGES: - verbose_proxy_logger.warning( - f"Reached maximum page limit of {MicrosoftSSOHandler.MAX_GRAPH_API_PAGES}. Some groups may not be included." - ) - - # If service_principal_group_ids is not empty, only return group_ids that are in both all_group_ids and service_principal_group_ids - if service_principal_group_ids and len(service_principal_group_ids) > 0: - all_group_ids = [group_id for group_id in all_group_ids if group_id in service_principal_group_ids] - - return all_group_ids - - except Exception as e: - verbose_proxy_logger.error(f"Error getting user groups from Microsoft Graph API: {e}") - return [] - - @staticmethod - async def fetch_and_parse_groups( - url: str, headers: dict, async_client: AsyncHTTPHandler - ) -> Tuple[List[str], Optional[str]]: - """Helper function to fetch and parse group data from a URL""" - response = await async_client.get(url, headers=headers) - response_json = response.json() - response_typed = await MicrosoftSSOHandler._cast_graph_api_response_dict(response=response_json) - group_ids = MicrosoftSSOHandler._get_group_ids_from_graph_api_response(response=response_typed) - return group_ids, response_typed.get("odata_nextLink") - - @staticmethod - def _get_group_ids_from_graph_api_response( - response: MicrosoftGraphAPIUserGroupResponse, - ) -> List[str]: - group_ids = [] - for _object in response.get("value", []) or []: - _group_id = _object.get("id") - if _group_id is not None: - group_ids.append(_group_id) - return group_ids - - @staticmethod - async def _cast_graph_api_response_dict( - response: dict, - ) -> MicrosoftGraphAPIUserGroupResponse: - directory_objects: List[MicrosoftGraphAPIUserGroupDirectoryObject] = [] - for _object in response.get("value", []): - directory_objects.append( - MicrosoftGraphAPIUserGroupDirectoryObject( - odata_type=_object.get("@odata.type"), - id=_object.get("id"), - deletedDateTime=_object.get("deletedDateTime"), - description=_object.get("description"), - displayName=_object.get("displayName"), - roleTemplateId=_object.get("roleTemplateId"), - ) - ) - return MicrosoftGraphAPIUserGroupResponse( - odata_context=response.get("@odata.context"), - odata_nextLink=response.get("@odata.nextLink"), - value=directory_objects, - ) - - @staticmethod - async def get_group_ids_from_service_principal( - service_principal_id: str, - async_client: AsyncHTTPHandler, - access_token: Optional[str] = None, - ) -> Tuple[List[str], List[MicrosoftServicePrincipalTeam]]: - """ - Gets the groups belonging to the Service Principal Application - - Service Principal Id is an `Enterprise Application` in Azure AD - - Users use Enterprise Applications to manage Groups and Users on Microsoft Entra ID - """ - base_url = "https://graph.microsoft.com/v1.0" - # Endpoint to get app role assignments for the given service principal - endpoint = f"/servicePrincipals/{service_principal_id}/appRoleAssignedTo" - url = base_url + endpoint - - headers = { - "Authorization": f"Bearer {access_token}", - "Content-Type": "application/json", - } - - response = await async_client.get(url, headers=headers) - response_json = response.json() - verbose_proxy_logger.debug(f"Response from service principal app role assigned to: {response_json}") - group_ids: List[str] = [] - service_principal_teams: List[MicrosoftServicePrincipalTeam] = [] - - for _object in response_json.get("value", []): - if _object.get("principalType") == "Group": - # Append the group ID to the list - group_ids.append(_object.get("principalId")) - # Append the service principal team to the list - service_principal_teams.append( - MicrosoftServicePrincipalTeam( - principalDisplayName=_object.get("principalDisplayName"), - principalId=_object.get("principalId"), - ) - ) - - return group_ids, service_principal_teams - - @staticmethod - async def create_litellm_teams_from_service_principal_team_ids( - service_principal_teams: List[MicrosoftServicePrincipalTeam], - ): - """ - Creates Litellm Teams from the Service Principal Group IDs - - When a user sets a `SERVICE_PRINCIPAL_ID` in the env, litellm will fetch groups under that service principal and create Litellm Teams from them - """ - verbose_proxy_logger.debug(f"Creating Litellm Teams from Service Principal Teams: {service_principal_teams}") - for service_principal_team in service_principal_teams: - litellm_team_id: Optional[str] = service_principal_team.get("principalId") - litellm_team_name: Optional[str] = service_principal_team.get("principalDisplayName") - if not litellm_team_id: - verbose_proxy_logger.debug( - f"Skipping team creation for {litellm_team_name} because it has no principalId" - ) - continue - - await SSOAuthenticationHandler.create_litellm_team_from_sso_group( - litellm_team_id=litellm_team_id, - litellm_team_name=litellm_team_name, - ) - - -class GoogleSSOHandler: - """ - Handles Google SSO callback response and returns a CustomOpenID object - """ - - @staticmethod - async def get_google_callback_response( - request: Request, - google_client_id: str, - redirect_url: str, - return_raw_sso_response: bool = False, - ) -> Union[OpenID, dict]: - """ - Get the Google SSO callback response - - Args: - return_raw_sso_response: If True, return the raw SSO response - """ - from fastapi_sso.sso.google import GoogleSSO - - google_client_secret = os.getenv("GOOGLE_CLIENT_SECRET", None) - if google_client_secret is None: - raise ProxyException( - message="GOOGLE_CLIENT_SECRET not set. Set it in .env file", - type=ProxyErrorTypes.auth_error, - param="GOOGLE_CLIENT_SECRET", - code=status.HTTP_500_INTERNAL_SERVER_ERROR, - ) - google_sso = GoogleSSO( - client_id=google_client_id, - redirect_uri=redirect_url, - client_secret=google_client_secret, - ) - - # if user is trying to get the raw sso response for debugging, return the raw sso response - if return_raw_sso_response: - return ( - await google_sso.verify_and_process( - request=request, - convert_response=False, # type: ignore - ) - or {} - ) - - result = await google_sso.verify_and_process(request) - return result or {} - - -@router.get("/sso/debug/login", tags=["experimental"], include_in_schema=False) -async def debug_sso_login(request: Request): - """ - Create Proxy API Keys using Google Workspace SSO. Requires setting PROXY_BASE_URL in .env - PROXY_BASE_URL should be the your deployed proxy endpoint, e.g. PROXY_BASE_URL="https://litellm-production-7002.up.railway.app/" - Example: - """ - from litellm.proxy.proxy_server import premium_user - - microsoft_client_id = os.getenv("MICROSOFT_CLIENT_ID", None) - google_client_id = os.getenv("GOOGLE_CLIENT_ID", None) - generic_client_id = os.getenv("GENERIC_CLIENT_ID", None) - - ####### Check if user is a Enterprise / Premium User ####### - if microsoft_client_id is not None or google_client_id is not None or generic_client_id is not None: - if premium_user is not True: - raise ProxyException( - message="You must be a LiteLLM Enterprise user to use SSO. If you have a license please set `LITELLM_LICENSE` in your env. If you want to obtain a license meet with us here: https://enterprise.litellm.ai/demo You are seeing this error message because You set one of `MICROSOFT_CLIENT_ID`, `GOOGLE_CLIENT_ID`, or `GENERIC_CLIENT_ID` in your env. Please unset this", - type=ProxyErrorTypes.auth_error, - param="premium_user", - code=status.HTTP_403_FORBIDDEN, - ) - - # get url from request - redirect_url = SSOAuthenticationHandler.get_redirect_url_for_sso( - request=request, - sso_callback_route="sso/debug/callback", - ) - - # Check if we should use SSO handler - if ( - SSOAuthenticationHandler.should_use_sso_handler( - microsoft_client_id=microsoft_client_id, - google_client_id=google_client_id, - generic_client_id=generic_client_id, - ) - is True - ): - return await SSOAuthenticationHandler.get_sso_login_redirect( - redirect_url=redirect_url, - microsoft_client_id=microsoft_client_id, - google_client_id=google_client_id, - generic_client_id=generic_client_id, - request=request, - ) - - -@router.get("/sso/debug/callback", tags=["experimental"], include_in_schema=False) -async def debug_sso_callback(request: Request): - """ - Returns the OpenID object returned by the SSO provider - """ - import json - - from fastapi.responses import HTMLResponse - - from litellm.proxy._types import LiteLLM_JWTAuth - from litellm.proxy.auth.handle_jwt import JWTHandler - from litellm.proxy.proxy_server import ( - general_settings, - jwt_handler, - prisma_client, - user_api_key_cache, - ) - - sso_jwt_handler: Optional[JWTHandler] = None - ui_access_mode = general_settings.get("ui_access_mode", None) - if ui_access_mode is not None and isinstance(ui_access_mode, dict): - sso_jwt_handler = JWTHandler() - sso_jwt_handler.update_environment( - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - litellm_jwtauth=LiteLLM_JWTAuth( - team_ids_jwt_field=general_settings.get("ui_access_mode", {}).get("sso_group_jwt_field", None), - ), - leeway=0, - ) - - microsoft_client_id = os.getenv("MICROSOFT_CLIENT_ID", None) - google_client_id = os.getenv("GOOGLE_CLIENT_ID", None) - generic_client_id = os.getenv("GENERIC_CLIENT_ID", None) - - redirect_url = os.getenv("PROXY_BASE_URL", str(request.base_url)) - if redirect_url.endswith("/"): - redirect_url += "sso/debug/callback" - else: - redirect_url += "/sso/debug/callback" - - result = None - received_response: Optional[dict] = None - access_token_payload: Optional[dict] = None - if google_client_id is not None: - result = await GoogleSSOHandler.get_google_callback_response( - request=request, - google_client_id=google_client_id, - redirect_url=redirect_url, - return_raw_sso_response=True, - ) - elif microsoft_client_id is not None: - result = await MicrosoftSSOHandler.get_microsoft_callback_response( - request=request, - microsoft_client_id=microsoft_client_id, - redirect_url=redirect_url, - return_raw_sso_response=True, - ) - - elif generic_client_id is not None: - ( - result, - received_response, - access_token_payload, - ) = await get_generic_sso_response( - request=request, - jwt_handler=jwt_handler, - generic_client_id=generic_client_id, - redirect_url=redirect_url, - sso_jwt_handler=sso_jwt_handler, - ) - - # If result is None, return a basic error message - if result is None: - return HTMLResponse( - content="

SSO Authentication Failed

No data was returned from the SSO provider.

", - status_code=400, - ) - - # Convert the OpenID object to a dictionary - if hasattr(result, "__dict__"): - result_dict = result.__dict__ - else: - result_dict = dict(result) - - # Filter out any None values and convert to JSON serializable format - filtered_result = {} - for key, value in result_dict.items(): - if value is not None and not key.startswith("_"): - if isinstance(value, (str, int, float, bool)) or value is None: - filtered_result[key] = value - else: - try: - # Try to convert to string or another JSON serializable format - filtered_result[key] = str(value) - except Exception as e: - filtered_result[key] = f"Complex value (not displayable): {str(e)}" - - # Defense-in-depth: ensure no bearer tokens leak into the rendered HTML even if - # a non-conforming IdP places them in its userinfo response. - safe_raw_claims = {k: v for k, v in (received_response or {}).items() if k not in _OAUTH_TOKEN_FIELDS} - safe_access_token_claims = {k: v for k, v in (access_token_payload or {}).items() if k not in _OAUTH_TOKEN_FIELDS} - - sso_payload = { - "parsed_by_proxy": filtered_result, - "raw_claims": safe_raw_claims, - "access_token_claims": safe_access_token_claims, - } - - # Replace the placeholder in the template with the actual data - sso_payload_json = json.dumps(sso_payload, indent=2, default=str).replace("