From 2d8f5111af202896d9148fde4a96dfb564e409c2 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Tue, 16 Dec 2025 13:14:39 -0800 Subject: [PATCH] add /sso/readiness route --- litellm/proxy/management_endpoints/ui_sso.py | 85 +++++++ .../proxy/management_endpoints/test_ui_sso.py | 233 ++++++++++++++++++ 2 files changed, 318 insertions(+) diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 123f58ab89..d4dfd86744 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -1126,6 +1126,91 @@ async def get_ui_settings(request: Request): } +@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 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 500fc67de8..a08fc2cba6 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -3043,3 +3043,236 @@ class TestAddMissingTeamMember: assert set(added_teams) == set( expected_teams_added ), f"Expected teams {expected_teams_added}, but got {added_teams}" + + +class TestSSOReadinessEndpoint: + """Test the /sso/readiness endpoint""" + + @pytest.mark.asyncio + async def test_sso_readiness_no_sso_configured(self): + """Test that readiness returns healthy when no SSO is configured""" + from fastapi.testclient import TestClient + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + from litellm.proxy.proxy_server import app + + mock_user_auth = UserAPIKeyAuth( + user_id="test-user-123", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + app.dependency_overrides[user_api_key_auth] = lambda: mock_user_auth + + try: + client = TestClient(app) + + with patch.dict(os.environ, {}, clear=True): + response = client.get("/sso/readiness") + + assert response.status_code == 200 + data = response.json() + assert data["status"] == "healthy" + assert data["sso_configured"] is False + assert data["message"] == "No SSO provider configured" + finally: + app.dependency_overrides.clear() + + @pytest.mark.asyncio + async def test_sso_readiness_google_fully_configured(self): + """Test that readiness returns healthy when Google SSO is fully configured""" + from fastapi.testclient import TestClient + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + from litellm.proxy.proxy_server import app + + mock_user_auth = UserAPIKeyAuth( + user_id="test-user-123", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + app.dependency_overrides[user_api_key_auth] = lambda: mock_user_auth + + try: + client = TestClient(app) + + with patch.dict( + os.environ, + { + "GOOGLE_CLIENT_ID": "test-google-client-id", + "GOOGLE_CLIENT_SECRET": "test-google-secret", + }, + clear=True, + ): + response = client.get("/sso/readiness") + + assert response.status_code == 200 + data = response.json() + assert data["status"] == "healthy" + assert data["sso_configured"] is True + assert data["provider"] == "google" + assert "Google SSO is properly configured" in data["message"] + finally: + app.dependency_overrides.clear() + + @pytest.mark.asyncio + async def test_sso_readiness_google_missing_secret(self): + """Test that readiness returns unhealthy when Google SSO is missing GOOGLE_CLIENT_SECRET""" + from fastapi.testclient import TestClient + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + from litellm.proxy.proxy_server import app + + mock_user_auth = UserAPIKeyAuth( + user_id="test-user-123", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + app.dependency_overrides[user_api_key_auth] = lambda: mock_user_auth + + try: + client = TestClient(app) + + with patch.dict( + os.environ, + {"GOOGLE_CLIENT_ID": "test-google-client-id"}, + clear=True, + ): + response = client.get("/sso/readiness") + + assert response.status_code == 503 + data = response.json()["detail"] + assert data["status"] == "unhealthy" + assert data["sso_configured"] is True + assert data["provider"] == "google" + assert "GOOGLE_CLIENT_SECRET" in data["missing_environment_variables"] + assert "Google SSO is configured but missing required environment variables" in data["message"] + finally: + app.dependency_overrides.clear() + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "env_vars,expected_status,expected_provider,expected_missing_vars", + [ + ( + { + "MICROSOFT_CLIENT_ID": "test-microsoft-client-id", + "MICROSOFT_CLIENT_SECRET": "test-microsoft-secret", + "MICROSOFT_TENANT": "test-tenant", + }, + 200, + "microsoft", + [], + ), + ( + {"MICROSOFT_CLIENT_ID": "test-microsoft-client-id"}, + 503, + "microsoft", + ["MICROSOFT_CLIENT_SECRET", "MICROSOFT_TENANT"], + ), + ], + ) + async def test_sso_readiness_microsoft_configurations( + self, env_vars, expected_status, expected_provider, expected_missing_vars + ): + """Test Microsoft SSO readiness with both fully configured and missing variables""" + from fastapi.testclient import TestClient + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + from litellm.proxy.proxy_server import app + + mock_user_auth = UserAPIKeyAuth( + user_id="test-user-123", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + app.dependency_overrides[user_api_key_auth] = lambda: mock_user_auth + + try: + client = TestClient(app) + + with patch.dict(os.environ, env_vars, clear=True): + response = client.get("/sso/readiness") + + assert response.status_code == expected_status + + if expected_status == 200: + data = response.json() + assert data["sso_configured"] is True + assert data["provider"] == expected_provider + assert data["status"] == "healthy" + assert "Microsoft SSO is properly configured" in data["message"] + else: + data = response.json()["detail"] + assert data["sso_configured"] is True + assert data["provider"] == expected_provider + assert data["status"] == "unhealthy" + assert set(data["missing_environment_variables"]) == set( + expected_missing_vars + ) + finally: + app.dependency_overrides.clear() + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "env_vars,expected_status,expected_provider,expected_missing_vars", + [ + ( + { + "GENERIC_CLIENT_ID": "test-generic-client-id", + "GENERIC_CLIENT_SECRET": "test-generic-secret", + "GENERIC_AUTHORIZATION_ENDPOINT": "https://auth.example.com/authorize", + "GENERIC_TOKEN_ENDPOINT": "https://auth.example.com/token", + "GENERIC_USERINFO_ENDPOINT": "https://auth.example.com/userinfo", + }, + 200, + "generic", + [], + ), + ( + {"GENERIC_CLIENT_ID": "test-generic-client-id"}, + 503, + "generic", + [ + "GENERIC_CLIENT_SECRET", + "GENERIC_AUTHORIZATION_ENDPOINT", + "GENERIC_TOKEN_ENDPOINT", + "GENERIC_USERINFO_ENDPOINT", + ], + ), + ], + ) + async def test_sso_readiness_generic_configurations( + self, env_vars, expected_status, expected_provider, expected_missing_vars + ): + """Test Generic SSO readiness with both fully configured and missing variables""" + from fastapi.testclient import TestClient + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + from litellm.proxy.proxy_server import app + + mock_user_auth = UserAPIKeyAuth( + user_id="test-user-123", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + app.dependency_overrides[user_api_key_auth] = lambda: mock_user_auth + + try: + client = TestClient(app) + + with patch.dict(os.environ, env_vars, clear=True): + response = client.get("/sso/readiness") + + assert response.status_code == expected_status + + if expected_status == 200: + data = response.json() + assert data["sso_configured"] is True + assert data["provider"] == expected_provider + assert data["status"] == "healthy" + assert "Generic SSO is properly configured" in data["message"] + else: + data = response.json()["detail"] + assert data["sso_configured"] is True + assert data["provider"] == expected_provider + assert data["status"] == "unhealthy" + assert set(data["missing_environment_variables"]) == set( + expected_missing_vars + ) + finally: + app.dependency_overrides.clear()