From 0f90e7d9ad2e6332b3715503ea342cd0680d88eb Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 10 Apr 2025 17:40:58 -0700 Subject: [PATCH] [Bug Fix MSFT SSO] Use correct field for user email when using MSFT SSO (#9886) * fix openid_from_response * test_microsoft_sso_handler_openid_from_response_user_principal_name * test upsert_sso_user --- litellm/proxy/management_endpoints/ui_sso.py | 26 +----- .../proxy/management_endpoints/test_ui_sso.py | 80 +++++++++++++++++++ 2 files changed, 83 insertions(+), 23 deletions(-) diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 1e10aebedb..0365336e73 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -468,9 +468,6 @@ async def auth_callback(request: Request): # noqa: PLR0915 result=result, user_info=user_info, user_email=user_email, - user_id_models=user_id_models, - max_internal_user_budget=max_internal_user_budget, - internal_user_budget_duration=internal_user_budget_duration, user_defined_values=user_defined_values, prisma_client=prisma_client, ) @@ -831,37 +828,20 @@ class SSOAuthenticationHandler: result: Optional[Union[CustomOpenID, OpenID, dict]], user_info: Optional[Union[NewUserResponse, LiteLLM_UserTable]], user_email: Optional[str], - user_id_models: List[str], - max_internal_user_budget: Optional[float], - internal_user_budget_duration: Optional[str], user_defined_values: Optional[SSOUserDefinedValues], prisma_client: PrismaClient, ): """ Connects the SSO Users to the User Table in LiteLLM DB - - If user on LiteLLM DB, update the user_id with the SSO user_id + - If user on LiteLLM DB, update the user_email with the SSO user_email - 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 - user_defined_values = SSOUserDefinedValues( - models=getattr(user_info, "models", user_id_models), - user_id=user_info.user_id or "", - user_email=getattr(user_info, "user_email", user_email), - user_role=getattr(user_info, "user_role", None), - max_budget=getattr( - user_info, "max_budget", max_internal_user_budget - ), - budget_duration=getattr( - user_info, "budget_duration", internal_user_budget_duration - ), - ) - - # update id await prisma_client.db.litellm_usertable.update_many( - where={"user_email": user_email}, data={"user_id": user_id} # type: ignore + where={"user_id": user_id}, data={"user_email": user_email} ) else: verbose_proxy_logger.info( @@ -1045,7 +1025,7 @@ class MicrosoftSSOHandler: response = response or {} verbose_proxy_logger.debug(f"Microsoft SSO Callback Response: {response}") openid_response = CustomOpenID( - email=response.get("mail"), + email=response.get("userPrincipalName") or response.get("mail"), display_name=response.get("displayName"), provider="microsoft", id=response.get("id"), diff --git a/tests/litellm/proxy/management_endpoints/test_ui_sso.py b/tests/litellm/proxy/management_endpoints/test_ui_sso.py index ff9700393f..09e337bf84 100644 --- a/tests/litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/litellm/proxy/management_endpoints/test_ui_sso.py @@ -21,6 +21,7 @@ from litellm.proxy.management_endpoints.types import CustomOpenID from litellm.proxy.management_endpoints.ui_sso import ( GoogleSSOHandler, MicrosoftSSOHandler, + SSOAuthenticationHandler, ) from litellm.types.proxy.management_endpoints.ui_sso import ( MicrosoftGraphAPIUserGroupDirectoryObject, @@ -29,6 +30,37 @@ from litellm.types.proxy.management_endpoints.ui_sso import ( ) +def test_microsoft_sso_handler_openid_from_response_user_principal_name(): + # Arrange + # Create a mock response similar to what Microsoft SSO would return + mock_response = { + "userPrincipalName": "test@example.com", + "displayName": "Test User", + "id": "user123", + "givenName": "Test", + "surname": "User", + "some_other_field": "value", + } + expected_team_ids = ["team1", "team2"] + # Act + # Call the method being tested + result = MicrosoftSSOHandler.openid_from_response( + response=mock_response, team_ids=expected_team_ids + ) + + # Assert + + # Check that the result is a CustomOpenID object with the expected values + assert isinstance(result, CustomOpenID) + assert result.email == "test@example.com" + assert result.display_name == "Test User" + assert result.provider == "microsoft" + assert result.id == "user123" + assert result.first_name == "Test" + assert result.last_name == "User" + assert result.team_ids == expected_team_ids + + def test_microsoft_sso_handler_openid_from_response(): # Arrange # Create a mock response similar to what Microsoft SSO would return @@ -386,6 +418,54 @@ def test_get_group_ids_from_graph_api_response(): @pytest.mark.asyncio +async def test_upsert_sso_user_existing_user(): + """ + If a user_id is already in the LiteLLM DB and the user signed in with SSO. Ensure that the user_id is updated with the SSO user_email + + SSO Test + """ + # Arrange + mock_prisma = MagicMock() + mock_prisma.db = MagicMock() + mock_prisma.db.litellm_usertable = MagicMock() + mock_prisma.db.litellm_usertable.update_many = AsyncMock() + + # Create a mock existing user + mock_user = MagicMock() + mock_user.user_id = "existing_user_123" + mock_user.user_email = "old_email@example.com" + + # Create mock SSO response + mock_sso_response = CustomOpenID( + email="new_email@example.com", + display_name="Test User", + provider="microsoft", + id="existing_user_123", + first_name="Test", + last_name="User", + team_ids=[], + ) + + # Create mock user defined values + mock_user_defined_values = MagicMock() + + # Act + result = await SSOAuthenticationHandler.upsert_sso_user( + result=mock_sso_response, + user_info=mock_user, + user_email="new_email@example.com", + user_defined_values=mock_user_defined_values, + prisma_client=mock_prisma, + ) + + # Assert + mock_prisma.db.litellm_usertable.update_many.assert_called_once_with( + where={"user_id": "existing_user_123"}, + data={"user_email": "new_email@example.com"}, + ) + assert result == mock_user + + async def test_default_team_params(): """ When litellm.default_team_params is set, it should be used to create a new team