fix: app_roles missing from jwt payload (#16448)

* fix: jwt app_roles missing

* add test

Author: mubashir1osmani

litellm/proxy/management_endpoints/ui_sso.py

@@ -564,10 +564,12 @@ def apply_user_info_values_to_sso_user_defined_values(
     if user_info is not None and user_info.user_id is not None:
         user_defined_values["user_id"] = user_info.user_id
 
-    if user_info is None or user_info.user_role is None:
-        user_defined_values["user_role"] = LitellmUserRoles.INTERNAL_USER_VIEW_ONLY
-    else:
-        user_defined_values["user_role"] = user_info.user_role
+    # Check if user_role already exists in user_defined_values (from JWT/SSO response)
+    if user_defined_values.get("user_role") is None:
+        if user_info is None or user_info.user_role is None:
+            user_defined_values["user_role"] = LitellmUserRoles.INTERNAL_USER_VIEW_ONLY
+        else:
+            user_defined_values["user_role"] = user_info.user_role
 
     # Preserve the user's existing models from the database
     if user_info is not None and hasattr(user_info, "models") and user_info.models:
@@ -1581,7 +1583,7 @@ async def get_redirect_response_from_openid(  # noqa: PLR0915
                 user_id=user_id,
                 user_email=user_email,
                 max_budget=max_internal_user_budget,
-                user_role=None,
+                user_role=user_role,
                 budget_duration=internal_user_budget_duration,
             )
 

tests/test_litellm/proxy/management_endpoints/test_entraid_app_roles.py

@@ -1,55 +1,92 @@
-"""
-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
+from litellm.proxy.management_endpoints.types import get_litellm_user_role
+from litellm.proxy._types import LitellmUserRoles
+
+
+def test_extracts_proxy_admin_role_from_jwt():
+    """Ensure supported app roles like 'proxy_admin' are extracted from the id_token."""
+    payload = {
+        "sub": "user123",
+        "email": "[email protected]",
+        "app_roles": ["proxy_admin"],
+        "aud": "litellm-app",
+        "iss": "https://login.microsoftonline.com/tenant-id/v2.0",
+        "exp": 9999999999,
+    }
+
+    token = jwt.encode(payload, "secret", algorithm="HS256")
+    roles = MicrosoftSSOHandler.get_app_roles_from_id_token(token)
+
+    assert roles == ["proxy_admin"]
+
+
+def test_maps_internal_user_role():
+    """Ensure internal_user role is correctly mapped to LitellmUserRoles."""
+    payload = {
+        "sub": "user456",
+        "email": "[email protected]",
+        "app_roles": ["internal_user"],
+        "aud": "litellm-app",
+        "iss": "https://login.microsoftonline.com/tenant-id/v2.0",
+        "exp": 9999999999,
+    }
+
+    token = jwt.encode(payload, "secret", algorithm="HS256")
+    roles = MicrosoftSSOHandler.get_app_roles_from_id_token(token)
+
+    # Map to LitellmUserRoles
+    chosen = None
+    for r in roles:
+        mapped = get_litellm_user_role(r)
+        if mapped is not None:
+            chosen = mapped
+            break
+
+    assert chosen == LitellmUserRoles.INTERNAL_USER
+
+
+def test_maps_proxy_admin_viewer_role():
+    """Ensure proxy_admin_viewer role is correctly mapped."""
+    payload = {
+        "sub": "user789",
+        "email": "[email protected]",
+        "app_roles": ["proxy_admin_viewer"],
+        "aud": "litellm-app",
+        "iss": "https://login.microsoftonline.com/tenant-id/v2.0",
+        "exp": 9999999999,
+    }
+
+    token = jwt.encode(payload, "secret", algorithm="HS256")
+    roles = MicrosoftSSOHandler.get_app_roles_from_id_token(token)
+
+    chosen = None
+    for r in roles:
+        mapped = get_litellm_user_role(r)
+        if mapped is not None:
+            chosen = mapped
+            break
+
+    assert chosen == LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY
+
+
+def test_defaults_to_internal_user_viewer_when_no_role():
+    """Ensure default role is internal_user_viewer when no app role is present."""
+    payload = {
+        "sub": "user_no_role",
+        "email": "[email protected]",
+        "aud": "litellm-app",
+        "iss": "https://login.microsoftonline.com/tenant-id/v2.0",
+        "exp": 9999999999,
+    }
+
+    token = jwt.encode(payload, "secret", algorithm="HS256")
+    roles = MicrosoftSSOHandler.get_app_roles_from_id_token(token)
+
+    assert roles == []
 
+    # Default role would be internal_user_viewer
+    default_role = LitellmUserRoles.INTERNAL_USER_VIEW_ONLY
+    assert default_role.value == "internal_user_viewer"
 
-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": "[email protected]",
-            "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": "[email protected]",
-            "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