diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 14f273f3269..3031508adeb 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -116,6 +116,14 @@ async def serve_login_page( missing_env_vars = show_missing_vars_in_env() if missing_env_vars is not None: return missing_env_vars + ######################################################### + # Construct Redirect URL + base_url_to_redirect_to: Optional[str] = None + base_url_to_redirect_to = os.getenv("PROXY_BASE_URL", "") + server_root_path = os.getenv("SERVER_ROOT_PATH", "") + if server_root_path != "": + base_url_to_redirect_to += server_root_path + ######################################################### # Build the unified login page HTML error_message = "" @@ -137,7 +145,7 @@ async def serve_login_page( sso_button = "" if sso_available: - sso_button = """ + sso_button = f"""
or

- 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 c7c67cbba71..245f350be1b 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -1245,3 +1245,47 @@ class TestCustomUISSO: # Verify the result is the redirect response assert result == mock_redirect_response assert result.status_code == 303 + + +@pytest.mark.asyncio +async def test_serve_login_page_server_root_path(): + """ + Test that serve_login_page includes SERVER_ROOT_PATH in the SSO login URL + when SERVER_ROOT_PATH is set. + """ + # Arrange + mock_request = MagicMock(spec=Request) + captured_html = "" + + # Mock environment variables + env_vars = { + "PROXY_BASE_URL": "https://example.com", + "SERVER_ROOT_PATH": "/api/v1", + "GOOGLE_CLIENT_ID": "mock_google_client_id", # Enable SSO + "DATABASE_URL": "mock_db_url", # Satisfy show_missing_vars_in_env + "LITELLM_MASTER_KEY": "mock_master_key", # Satisfy show_missing_vars_in_env + } + + # Patch HTMLResponse to capture the content + def mock_html_response(content, status_code=200): + nonlocal captured_html + captured_html = content + return MagicMock() + + with patch.dict(os.environ, env_vars): + with patch("litellm.proxy.proxy_server.premium_user", True): + with patch("litellm.proxy.proxy_server.prisma_client", MagicMock()): + with patch("litellm.proxy.proxy_server.master_key", "mock_master_key"): + with patch("fastapi.responses.HTMLResponse", side_effect=mock_html_response): + # Import the function to test + from litellm.proxy.management_endpoints.ui_sso import ( + serve_login_page, + ) + + # Act + result = await serve_login_page(request=mock_request) + + # Assert + assert result is not None + expected_url = "https://example.com/api/v1/sso/login" + assert expected_url in captured_html, f"Expected URL '{expected_url}' not found in HTML content"