[Fix] UI SSO - allow up-serting user role when SSO provider role changes (#17754)

* use _build_sso_user_update_data

* test_build_sso_user_update_data_with_valid_role

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(

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.