diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 73970aae53b..c48e1c6cc58 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -51,7 +51,6 @@ from litellm.proxy.common_utils.admin_ui_utils import ( from litellm.proxy.common_utils.html_forms.jwt_display_template import ( jwt_display_template, ) -from litellm.proxy.common_utils.html_forms.ui_login import html_form from litellm.proxy.management_endpoints.internal_user_endpoints import new_user from litellm.proxy.management_endpoints.sso_helper_utils import ( check_is_admin_only_access, @@ -77,16 +76,20 @@ router = APIRouter() @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): # noqa: PLR0915 +async def serve_login_page( + request: Request, + source: Optional[str] = None, + key: Optional[str] = None, + error: Optional[str] = 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: + Serves a unified login page with options for both normal + username/password login and SSO. """ - from litellm.proxy.proxy_server import ( - premium_user, - user_custom_ui_sso_sign_in_handler, - ) + 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) @@ -99,6 +102,318 @@ async def google_login(request: Request, source: Optional[str] = None, key: Opti if is_disabled: return admin_ui_disabled() + ####### Check if user is a Enterprise / Premium User for SSO ####### + sso_available = False + 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 True: + sso_available = True + + ####### 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 + + # Build the unified login page HTML + error_message = "" + if error == "1": + error_message = """ +
+ ⚠️ Invalid username or password. Please try again. +
+ """ + + sso_button = "" + if sso_available: + sso_button = """ +
+

or

+ + 🔐 Login with SSO + +
+ """ + + # Get the base URL for form action using proper URL construction + form_action = get_custom_url(request_base_url=str(request.base_url), route="login") + + unified_login_html = f""" + + + + + LiteLLM Login + + + + +
+
+ +
+

Login

+

Access your LiteLLM Admin UI.

+ + {error_message} + +
+
+ + + + + + Default Credentials +
+

By default, Username is admin and Password is your set LiteLLM Proxy MASTER_KEY.

+

Need to set UI credentials or SSO? Check the documentation.

+
+ + + + + + +
+ + +
+ + + {sso_button} +
+ + + + """ + + from fastapi.responses import HTMLResponse + + return HTMLResponse(content=unified_login_html, status_code=200) + + +@router.get("/sso/login", tags=["experimental"], include_in_schema=False) +async def sso_login_redirect( + request: Request, source: Optional[str] = None, key: Optional[str] = None +): + """ + Handles SSO login redirect - this is what the "Login with SSO" button points to + """ + from litellm.proxy.proxy_server import ( + premium_user, + 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 user is a Enterprise / Premium User ####### if ( microsoft_client_id is not None @@ -113,18 +428,12 @@ async def google_login(request: Request, source: Optional[str] = None, key: Opti code=status.HTTP_403_FORBIDDEN, ) - ####### 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 - ui_username = os.getenv("UI_USERNAME") - # 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", ) - + # Store CLI key in state for OAuth flow cli_state: Optional[str] = SSOAuthenticationHandler._get_cli_state( source=source, @@ -137,11 +446,14 @@ async def google_login(request: Request, source: Optional[str] = None, key: Opti from litellm_enterprise.proxy.auth.custom_sso_handler import ( 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.") + raise ValueError( + "Enterprise features are not available. Custom UI SSO sign-in requires LiteLLM Enterprise." + ) # Check if we should use SSO handler if ( @@ -160,16 +472,9 @@ async def google_login(request: Request, source: Optional[str] = None, key: Opti generic_client_id=generic_client_id, state=cli_state, ) - elif ui_username is not None: - # No Google, Microsoft SSO - # Use UI Credentials set in .env - from fastapi.responses import HTMLResponse - - return HTMLResponse(content=html_form, status_code=200) else: - from fastapi.responses import HTMLResponse - - return HTMLResponse(content=html_form, status_code=200) + # No SSO configured, redirect back to login page + return RedirectResponse(url="/sso/key/generate", status_code=303) def generic_response_convertor( @@ -525,15 +830,16 @@ async def check_and_update_if_proxy_admin_id( async def auth_callback(request: Request, state: Optional[str] = None): # noqa: PLR0915 """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 + if state and state.startswith(f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:"): # Extract the key ID from the state key_id = state.split(":", 1)[1] verbose_proxy_logger.info(f"CLI SSO callback detected for key: {key_id}") return await cli_sso_callback(request, key=key_id) - + from litellm.proxy._types import LiteLLM_JWTAuth from litellm.proxy.auth.handle_jwt import JWTHandler from litellm.proxy.proxy_server import ( @@ -608,7 +914,7 @@ async def auth_callback(request: Request, state: Optional[str] = None): # noqa: status_code=401, detail="Result not returned by SSO provider.", ) - + return await SSOAuthenticationHandler.get_redirect_response_from_openid( result=result, request=request, @@ -618,28 +924,26 @@ async def auth_callback(request: Request, state: Optional[str] = None): # noqa: ) - - async def cli_sso_callback(request: Request, key: Optional[str] = None): """CLI SSO callback - generates the key with pre-specified ID""" verbose_proxy_logger.info(f"CLI SSO callback for key: {key}") - + from litellm.proxy.management_endpoints.key_management_endpoints import ( generate_key_helper_fn, ) from litellm.proxy.proxy_server import prisma_client - - if not key or not key.startswith('sk-'): + + if not key or not key.startswith("sk-"): raise HTTPException( status_code=400, - detail="Invalid key parameter. Must be a valid key ID starting with 'sk-'" + detail="Invalid key parameter. Must be a valid key ID starting with 'sk-'", ) - + if prisma_client is None: raise HTTPException( status_code=500, detail=CommonProxyErrors.db_not_connected_error.value ) - + # Generate a simple key for CLI usage with the pre-specified key ID try: await generate_key_helper_fn( @@ -653,63 +957,57 @@ async def cli_sso_callback(request: Request, key: Optional[str] = None): table_name="key", token=key, # Use the pre-specified key ID ) - + verbose_proxy_logger.info(f"Generated CLI key: {key}") - + # Return success page from fastapi.responses import HTMLResponse from litellm.proxy.common_utils.html_forms.cli_sso_success import ( render_cli_sso_success_page, ) - + html_content = render_cli_sso_success_page() return HTMLResponse(content=html_content, status_code=200) - + except Exception as e: verbose_proxy_logger.error(f"Error generating CLI key: {e}") - raise HTTPException( - status_code=500, - detail=f"Failed to generate key: {str(e)}" - ) + raise HTTPException(status_code=500, detail=f"Failed to generate key: {str(e)}") @router.get("/sso/cli/poll/{key_id}", tags=["experimental"], include_in_schema=False) async def cli_poll_key(key_id: str): """CLI polling endpoint - checks if key exists in DB""" from litellm.proxy.proxy_server import prisma_client - - if not key_id.startswith('sk-'): - raise HTTPException( - status_code=400, - detail="Invalid key ID format" - ) - + + if not key_id.startswith("sk-"): + raise HTTPException(status_code=400, detail="Invalid key ID format") + if prisma_client is None: raise HTTPException( status_code=500, detail=CommonProxyErrors.db_not_connected_error.value ) - + try: # Check if key exists in database from litellm.proxy.utils import hash_token + hashed_token = hash_token(key_id) - + key_obj = await prisma_client.db.litellm_verificationtoken.find_unique( where={"token": hashed_token} ) - + if key_obj: verbose_proxy_logger.info(f"CLI key found: {key_id}") return {"status": "ready", "key": key_id} else: return {"status": "pending"} - + except Exception as e: verbose_proxy_logger.error(f"Error polling for CLI key: {e}") raise HTTPException( - status_code=500, - detail=f"Error checking key status: {str(e)}" + status_code=500, detail=f"Error checking key status: {str(e)}" ) @@ -811,6 +1109,7 @@ class SSOAuthenticationHandler: """ Handler for SSO Authentication across all SSO providers """ + @staticmethod async def get_sso_login_redirect( redirect_url: str, @@ -1163,7 +1462,6 @@ class SSOAuthenticationHandler: _new_team_request.update(_default_team_params) team_request = NewTeamRequest(**_new_team_request) return team_request - @staticmethod def _get_cli_state(source: Optional[str], key: Optional[str]) -> Optional[str]: @@ -1176,13 +1474,15 @@ class SSOAuthenticationHandler: LITELLM_CLI_SESSION_TOKEN_PREFIX, LITELLM_CLI_SOURCE_IDENTIFIER, ) - return f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:{key}" if source == LITELLM_CLI_SOURCE_IDENTIFIER and key else None - - + return ( + f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:{key}" + if source == LITELLM_CLI_SOURCE_IDENTIFIER and key + else None + ) @staticmethod - async def get_redirect_response_from_openid( # noqa: PLR0915 + async def get_redirect_response_from_openid( # noqa: PLR0915 result: Union[OpenID, dict, CustomOpenID], request: Request, received_response: Optional[dict] = None, @@ -1202,14 +1502,18 @@ class SSOAuthenticationHandler: ) from litellm.proxy.utils import get_custom_url, 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") + prisma_client = get_prisma_client_or_throw( + "Prisma client is None, connect a database to your proxy" + ) # User is Authe'd in - generate key for the UI to access Proxy verbose_proxy_logger.info(f"SSO callback result: {result}") user_email: Optional[str] = getattr(result, "email", None) - user_id: Optional[str] = getattr(result, "id", None) if result is not None else None + user_id: Optional[str] = ( + getattr(result, "id", None) if result is not None else None + ) if user_email is not None and os.getenv("ALLOWED_EMAIL_DOMAINS") is not None: email_domain = user_email.split("@")[1] @@ -1394,7 +1698,8 @@ class SSOAuthenticationHandler: redirect_response = RedirectResponse(url=litellm_dashboard_ui, status_code=303) redirect_response.set_cookie(key="token", value=jwt_token) return redirect_response - + + class MicrosoftSSOHandler: """ Handles Microsoft SSO callback response and returns a CustomOpenID object @@ -1894,3 +2199,28 @@ async def debug_sso_callback(request: Request): ) return HTMLResponse(content=html_content) + + +@router.post("/sso/key/generate", tags=["experimental"], include_in_schema=False) +async def process_login(request: Request): + """ + Process username/password login from the unified login page + """ + try: + # Get form data + form_data = await request.form() + username = form_data.get("username") + password = form_data.get("password") + + if not username or not password: + return RedirectResponse(url="/sso/key/generate?error=1", status_code=303) + + # Import the actual login function from proxy_server + from litellm.proxy.proxy_server import login + + # Call the real login function that handles all the authentication properly + return await login(request) + + except Exception as e: + verbose_proxy_logger.error(f"Error processing login: {e}") + return RedirectResponse(url="/sso/key/generate?error=1", status_code=303) diff --git a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py index 73a7339afd0..c7c67cbba71 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -938,10 +938,10 @@ class TestUISSO_FunctionsExistence: from litellm.proxy.management_endpoints.ui_sso import auth_callback assert callable(auth_callback) - def test_google_login_exists(self): - """Test that google_login function exists""" - from litellm.proxy.management_endpoints.ui_sso import google_login - assert callable(google_login) + def test_sso_login_redirect_exists(self): + """Test that sso_login_redirect function exists""" + from litellm.proxy.management_endpoints.ui_sso import sso_login_redirect + assert callable(sso_login_redirect) def test_sso_authentication_handler_exists(self): """Test that SSOAuthenticationHandler class exists with new methods""" @@ -1054,7 +1054,7 @@ class TestCustomUISSO: """Test that proper error is raised when enterprise module is not available""" from unittest.mock import MagicMock, patch - from litellm.proxy.management_endpoints.ui_sso import google_login + from litellm.proxy.management_endpoints.ui_sso import sso_login_redirect # Mock request mock_request = MagicMock()