diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 59a93f3c486..c8f58713cda 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -562,7 +562,7 @@ 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. " @@ -572,6 +572,41 @@ def _should_use_role_from_sso_response(sso_role: Optional[str]) -> bool: 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": 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 + def apply_user_info_values_to_sso_user_defined_values( user_info: Optional[Union[LiteLLM_UserTable, NewUserResponse]], @@ -1341,14 +1376,20 @@ class SSOAuthenticationHandler: """ Connects the SSO Users to the User Table in LiteLLM DB - - If user on LiteLLM DB, update the user_email with the SSO user_email + - 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 prisma_client.db.litellm_usertable.update_many( - where={"user_id": user_id}, data={"user_email": user_email} + where={"user_id": user_id}, data=update_data ) else: verbose_proxy_logger.info( 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 8d7aa51fa0f..710d27e072c 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -572,6 +572,205 @@ def test_apply_user_info_values_sso_role_takes_precedence(): assert sso_user_defined_values["models"] == ["model-1"] +def test_build_sso_user_update_data_with_valid_role(): + """ + Test that _build_sso_user_update_data includes role when SSO provides a valid role. + """ + from litellm.proxy._types import LitellmUserRoles + from litellm.proxy.management_endpoints.types import CustomOpenID + from litellm.proxy.management_endpoints.ui_sso import _build_sso_user_update_data + + sso_result = CustomOpenID( + id="test-user-123", + email="test@example.com", + display_name="Test User", + provider="microsoft", + team_ids=[], + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + + update_data = _build_sso_user_update_data( + result=sso_result, + user_email="test@example.com", + user_id="test-user-123", + ) + + assert update_data["user_email"] == "test@example.com" + assert update_data["user_role"] == "proxy_admin" + + +def test_build_sso_user_update_data_without_role(): + """ + Test that _build_sso_user_update_data only includes email when SSO has no role. + """ + from litellm.proxy.management_endpoints.types import CustomOpenID + from litellm.proxy.management_endpoints.ui_sso import _build_sso_user_update_data + + sso_result = CustomOpenID( + id="test-user-456", + email="test@example.com", + display_name="Test User", + provider="microsoft", + team_ids=[], + user_role=None, + ) + + update_data = _build_sso_user_update_data( + result=sso_result, + user_email="test@example.com", + user_id="test-user-456", + ) + + assert update_data["user_email"] == "test@example.com" + assert "user_role" not in update_data + + +@pytest.mark.asyncio +async def test_upsert_sso_user_updates_role_for_existing_user(): + """ + Test that upsert_sso_user updates the user role in database when SSO provides a valid role. + + When a user's role is updated in the SSO provider (e.g., Azure), the role should be + updated in the LiteLLM database on subsequent logins, not just at initial user creation. + """ + from litellm.proxy._types import LiteLLM_UserTable, LitellmUserRoles + from litellm.proxy.management_endpoints.types import CustomOpenID + from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler + + # Mock prisma client + mock_prisma = MagicMock() + mock_prisma.db.litellm_usertable.update_many = AsyncMock() + + # Existing user in DB with old role + existing_user = LiteLLM_UserTable( + user_id="test-user-123", + user_email="test@example.com", + user_role="internal_user", + models=["model-1"], + ) + + # SSO result with new role (e.g., user was promoted to admin in Azure) + sso_result = CustomOpenID( + id="test-user-123", + email="test@example.com", + display_name="Test User", + provider="microsoft", + team_ids=["team-1"], + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + + # Act + await SSOAuthenticationHandler.upsert_sso_user( + result=sso_result, + user_info=existing_user, + user_email="test@example.com", + user_defined_values=None, + prisma_client=mock_prisma, + ) + + # Assert - verify database was updated with both email and role + mock_prisma.db.litellm_usertable.update_many.assert_called_once() + call_args = mock_prisma.db.litellm_usertable.update_many.call_args + assert call_args.kwargs["where"] == {"user_id": "test-user-123"} + assert call_args.kwargs["data"]["user_email"] == "test@example.com" + assert call_args.kwargs["data"]["user_role"] == "proxy_admin" + + +@pytest.mark.asyncio +async def test_upsert_sso_user_does_not_update_invalid_role(): + """ + Test that upsert_sso_user does not update the role if SSO provides an invalid role. + + If the SSO returns a role that is not a valid LiteLLM role, it should be ignored + and only the email should be updated. + """ + from litellm.proxy._types import LiteLLM_UserTable + from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler + + # Mock prisma client + mock_prisma = MagicMock() + mock_prisma.db.litellm_usertable.update_many = AsyncMock() + + # Existing user in DB + existing_user = LiteLLM_UserTable( + user_id="test-user-456", + user_email="test@example.com", + user_role="internal_user", + models=[], + ) + + # SSO result with invalid role - use MagicMock to bypass validation + # This simulates a raw SSO response that has an invalid role string + sso_result = MagicMock() + sso_result.user_role = "invalid_role_not_in_enum" + + # Act + await SSOAuthenticationHandler.upsert_sso_user( + result=sso_result, + user_info=existing_user, + user_email="test@example.com", + user_defined_values=None, + prisma_client=mock_prisma, + ) + + # Assert - verify only email was updated, not role + mock_prisma.db.litellm_usertable.update_many.assert_called_once() + call_args = mock_prisma.db.litellm_usertable.update_many.call_args + assert call_args.kwargs["where"] == {"user_id": "test-user-456"} + assert call_args.kwargs["data"]["user_email"] == "test@example.com" + assert "user_role" not in call_args.kwargs["data"] + + +@pytest.mark.asyncio +async def test_upsert_sso_user_no_role_in_sso_response(): + """ + Test that upsert_sso_user only updates email when SSO response has no role. + + When the SSO provider does not return a role, only the email should be updated. + """ + from litellm.proxy._types import LiteLLM_UserTable + from litellm.proxy.management_endpoints.types import CustomOpenID + from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler + + # Mock prisma client + mock_prisma = MagicMock() + mock_prisma.db.litellm_usertable.update_many = AsyncMock() + + # Existing user in DB + existing_user = LiteLLM_UserTable( + user_id="test-user-789", + user_email="old@example.com", + user_role="internal_user", + models=[], + ) + + # SSO result without role + sso_result = CustomOpenID( + id="test-user-789", + email="new@example.com", + display_name="Test User", + provider="microsoft", + team_ids=[], + user_role=None, + ) + + # Act + await SSOAuthenticationHandler.upsert_sso_user( + result=sso_result, + user_info=existing_user, + user_email="new@example.com", + user_defined_values=None, + prisma_client=mock_prisma, + ) + + # Assert - verify only email was updated + mock_prisma.db.litellm_usertable.update_many.assert_called_once() + call_args = mock_prisma.db.litellm_usertable.update_many.call_args + assert call_args.kwargs["where"] == {"user_id": "test-user-789"} + assert call_args.kwargs["data"]["user_email"] == "new@example.com" + assert "user_role" not in call_args.kwargs["data"] + + def test_get_user_email_and_id_extracts_microsoft_role(): """ Test that _get_user_email_and_id_from_result extracts user_role from Microsoft SSO.