diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index d8f426e1ae6..54909b9712b 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -875,9 +875,9 @@ async def insert_sso_user( if user_defined_values.get("max_budget") is None: user_defined_values["max_budget"] = litellm.max_internal_user_budget if user_defined_values.get("budget_duration") is None: - user_defined_values["budget_duration"] = ( - litellm.internal_user_budget_duration - ) + user_defined_values[ + "budget_duration" + ] = litellm.internal_user_budget_duration if user_defined_values["user_role"] is None: user_defined_values["user_role"] = LitellmUserRoles.INTERNAL_USER_VIEW_ONLY @@ -1118,9 +1118,9 @@ class SSOAuthenticationHandler: generic_authorization_endpoint and "okta" in generic_authorization_endpoint ): - redirect_params["state"] = ( - uuid.uuid4().hex - ) # set state param for okta - required + redirect_params[ + "state" + ] = uuid.uuid4().hex # set state param for okta - required return redirect_params @@ -1699,9 +1699,9 @@ class MicrosoftSSOHandler: # if user is trying to get the raw sso response for debugging, return the raw sso response if return_raw_sso_response: - original_msft_result[MicrosoftSSOHandler.GRAPH_API_RESPONSE_KEY] = ( - user_team_ids - ) + original_msft_result[ + MicrosoftSSOHandler.GRAPH_API_RESPONSE_KEY + ] = user_team_ids original_msft_result["app_roles"] = app_roles return original_msft_result or {} @@ -1739,7 +1739,7 @@ class MicrosoftSSOHandler: Extract app roles from the Microsoft Entra ID (Azure AD) id_token JWT. App roles are assigned in the Azure AD Enterprise Application and appear - in the 'roles' claim of the id_token. + in the 'app_roles' claim of the id_token. Args: id_token (Optional[str]): The JWT id_token from Microsoft SSO @@ -1758,8 +1758,8 @@ class MicrosoftSSOHandler: # (signature is already verified by fastapi_sso) decoded_token = jwt.decode(id_token, options={"verify_signature": False}) - # Extract roles claim from the token - roles = decoded_token.get("roles", []) + # Extract app_roles claim from the token + roles = decoded_token.get("app_roles", []) if roles and isinstance(roles, list): verbose_proxy_logger.debug( @@ -1817,9 +1817,9 @@ class MicrosoftSSOHandler: # Fetch user membership from Microsoft Graph API all_group_ids = [] - next_link: Optional[str] = ( - MicrosoftSSOHandler.graph_api_user_groups_endpoint - ) + next_link: Optional[ + str + ] = MicrosoftSSOHandler.graph_api_user_groups_endpoint auth_headers = {"Authorization": f"Bearer {access_token}"} page_count = 0 diff --git a/tests/test_litellm/proxy/management_endpoints/test_entraid_app_roles.py b/tests/test_litellm/proxy/management_endpoints/test_entraid_app_roles.py new file mode 100644 index 00000000000..f6248d36628 --- /dev/null +++ b/tests/test_litellm/proxy/management_endpoints/test_entraid_app_roles.py @@ -0,0 +1,55 @@ +""" +Unit tests for EntraID app roles JWT claim extraction. + +This module tests the get_app_roles_from_id_token method to ensure it correctly +extracts app roles from Microsoft EntraID JWT tokens and prevents regressions. +""" + +import pytest +import jwt + +from litellm.proxy.management_endpoints.ui_sso import MicrosoftSSOHandler + + +class TestEntraIDAppRoles: + """Test EntraID app roles extraction from JWT tokens""" + + def test_get_app_roles_from_id_token_works_without_roles(self): + """Test that JWT token works fine without app_roles claim""" + # Arrange - Token without app_roles (normal user) + payload = { + "sub": "user123", + "email": "user@company.com", + "aud": "litellm-app", + "iss": "https://login.microsoftonline.com/tenant-id/v2.0", + "exp": 9999999999, + } + no_roles_token = jwt.encode(payload, "secret", algorithm="HS256") + + # Act + result = MicrosoftSSOHandler.get_app_roles_from_id_token(no_roles_token) + + # Assert - Should return empty list, not error + assert result == [] + assert len(result) == 0 + + def test_get_app_roles_from_id_token_assigns_roles_when_present(self): + """Test that valid app roles are properly assigned when present""" + # Arrange - Token with valid roles + payload = { + "sub": "user123", + "email": "admin@company.com", + "app_roles": ["proxy_admin"], + "aud": "litellm-app", + "iss": "https://login.microsoftonline.com/tenant-id/v2.0", + "exp": 9999999999, + } + valid_roles_token = jwt.encode(payload, "secret", algorithm="HS256") + + # Act + result = MicrosoftSSOHandler.get_app_roles_from_id_token(valid_roles_token) + + # Assert - Should extract the role + assert result == ["proxy_admin"] + assert len(result) == 1 + assert "proxy_admin" in result