diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 278f3bdaafd..ecad221aad9 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -1375,50 +1375,22 @@ async def insert_sso_user( if user_defined_values is None: raise ValueError("user_defined_values is None") - # Check if role_mappings is configured in SSO settings - role_mappings_configured = False - 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" - ) - - # Get SSO config from dedicated table - sso_db_record = await prisma_client.db.litellm_ssoconfig.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") - role_mappings_configured = role_mappings_data is not None - generic_user_role_mappings = os.getenv("GENERIC_USER_ROLE_MAPPINGS", None) - if generic_user_role_mappings is not None: - role_mappings_configured = True - - except Exception as e: - # If we can't check role_mappings, continue with existing logic - verbose_proxy_logger.debug( - f"Could not check role_mappings configuration: {e}. Using default behavior." - ) - # Apply default_internal_user_params if litellm.default_internal_user_params: - # If role_mappings is configured and user_role is already set from SSO, preserve it - if ( - role_mappings_configured - and user_defined_values.get("user_role") is not None - ): + # 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 = user_defined_values.get("user_role") + 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}' (role_mappings configured)" + f"Preserved SSO-extracted role '{preserved_role}'" ) else: - # Default behavior: update all values including role + # 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 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 16f80826798..3dd2782c58e 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -1532,10 +1532,16 @@ class TestSSOHandlerIntegration: SSOAuthenticationHandler.should_use_sso_handler(None, None, None) is False ) + @patch.dict(os.environ, {}, clear=False) def test_get_redirect_url_for_sso(self): """Test the redirect URL generation for SSO""" from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler + # Remove env vars that override request base_url so the test is + # isolated from local settings. + os.environ.pop("PROXY_BASE_URL", None) + os.environ.pop("SERVER_ROOT_PATH", None) + # Mock request object mock_request = MagicMock() mock_request.base_url = "https://test.litellm.ai/" @@ -3400,19 +3406,6 @@ async def test_role_mappings_override_default_internal_user_params(): "models": [], } - # Mock Prisma client with SSO config that has role_mappings configured - mock_prisma = MagicMock() - mock_sso_config = MagicMock() - mock_sso_config.sso_settings = { - "role_mappings": { - "Admin": "proxy_admin", - "User": "internal_user", - } - } - mock_prisma.db.litellm_ssoconfig.find_unique = AsyncMock( - return_value=mock_sso_config - ) - # Mock new_user function mock_new_user_response = NewUserResponse( user_id="test-user-123", @@ -3421,9 +3414,6 @@ async def test_role_mappings_override_default_internal_user_params(): ) with patch( - "litellm.proxy.utils.get_prisma_client_or_throw", - return_value=mock_prisma, - ), patch( "litellm.proxy.management_endpoints.ui_sso.new_user", return_value=mock_new_user_response, ) as mock_new_user: @@ -3450,11 +3440,6 @@ async def test_role_mappings_override_default_internal_user_params(): assert ( new_user_request.budget_duration == "30d" ), "budget_duration from default_internal_user_params should be applied" - - # Note: models are applied via _update_internal_new_user_params inside new_user, - # not in insert_sso_user, so we verify user_defined_values was updated correctly - # by checking that the function completed successfully and other defaults were applied - # The models will be applied when new_user processes the request finally: # Restore original default_internal_user_params @@ -3465,6 +3450,91 @@ async def test_role_mappings_override_default_internal_user_params(): delattr(litellm, "default_internal_user_params") +@pytest.mark.asyncio +async def test_sso_role_preserved_without_role_mappings(): + """ + Test that SSO-extracted role is preserved even when role_mappings is NOT configured. + + This covers the case where the role comes from Microsoft app_roles or + GENERIC_USER_ROLE_ATTRIBUTE (not from LiteLLM's role_mappings feature). + Previously, the role was only preserved when role_mappings was configured, + causing admin users to be downgraded to internal_user. + """ + from litellm.proxy._types import NewUserResponse, SSOUserDefinedValues + from litellm.proxy.management_endpoints.ui_sso import insert_sso_user + + original_default_params = getattr(litellm, "default_internal_user_params", None) + + try: + # Set default_internal_user_params (as most deployments do) + litellm.default_internal_user_params = { + "user_role": "internal_user", + "max_budget": 50, + } + + # Mock SSO result from Microsoft with app_roles-derived admin role + mock_result_openid = CustomOpenID( + id="msft-user-456", + email="admin@company.com", + display_name="Admin User", + provider="microsoft", + team_ids=["group-1"], + user_role=None, # role is in user_defined_values, not on the OpenID result + ) + + # User defined values with role from Microsoft app_roles (NOT role_mappings) + user_defined_values: SSOUserDefinedValues = { + "user_id": "msft-user-456", + "user_email": "admin@company.com", + "user_role": "proxy_admin", # Role from Microsoft app_roles + "max_budget": None, + "budget_duration": None, + "models": [], + } + + mock_new_user_response = NewUserResponse( + user_id="msft-user-456", + key="sk-xxxxx", + teams=None, + ) + + # No role_mappings configured anywhere - the role came from app_roles + with patch( + "litellm.proxy.management_endpoints.ui_sso.new_user", + return_value=mock_new_user_response, + ) as mock_new_user: + _ = await insert_sso_user( + result_openid=mock_result_openid, + user_defined_values=user_defined_values, + ) + + mock_new_user.assert_called_once() + call_args = mock_new_user.call_args + new_user_request = call_args.kwargs["data"] + + # SSO role should be preserved even without role_mappings configured + assert ( + new_user_request.user_role == "proxy_admin" + ), "SSO role from app_roles should not be overwritten by default_internal_user_params" + + # Other defaults should still apply + assert ( + new_user_request.max_budget == 50 + ), "max_budget from default_internal_user_params should be applied" + + # Verify user_defined_values was also updated (it's mutated in-place) + assert ( + user_defined_values["user_role"] == "proxy_admin" + ), "user_defined_values should retain the SSO role after insert_sso_user" + + finally: + if original_default_params is not None: + litellm.default_internal_user_params = original_default_params + else: + if hasattr(litellm, "default_internal_user_params"): + delattr(litellm, "default_internal_user_params") + + class TestSSOReadinessEndpoint: """Test the /sso/readiness endpoint"""