Merge pull request #15583 from TeddyAmkie/fix/entraid-app-roles-jwt-claim-clean

Fix/entraid app roles jwt claim clean

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
 

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