feat(ui_sso.py): support mapping app roles from azure entra id to litellm user roles

Closes LIT-1228

Author: krrishdholakia

litellm/proxy/management_endpoints/types.py

@@ -4,10 +4,48 @@
 Might include fastapi/proxy requirements.txt related imports
 """
 
-from typing import List
+from typing import List, Optional, cast
 
 from fastapi_sso.sso.base import OpenID
 
+from litellm.proxy._types import LitellmUserRoles
+
+
+def is_valid_litellm_user_role(role_str: str) -> bool:
+    """
+    Check if a string is a valid LitellmUserRoles enum value (case-insensitive).
+
+    Args:
+        role_str: String to validate (e.g., "proxy_admin", "PROXY_ADMIN", "internal_user")
+
+    Returns:
+        True if the string matches a valid LitellmUserRoles value, False otherwise
+    """
+    try:
+        # Use _value2member_map_ for O(1) lookup, case-insensitive
+        return role_str.lower() in LitellmUserRoles._value2member_map_
+    except Exception:
+        return False
+
+
+def get_litellm_user_role(role_str: str) -> Optional[LitellmUserRoles]:
+    """
+    Convert a string to a LitellmUserRoles enum if valid (case-insensitive).
+
+    Args:
+        role_str: String to convert (e.g., "proxy_admin", "PROXY_ADMIN", "internal_user")
+
+    Returns:
+        LitellmUserRoles enum if valid, None otherwise
+    """
+    try:
+        # Use _value2member_map_ for O(1) lookup, case-insensitive
+        result = LitellmUserRoles._value2member_map_.get(role_str.lower())
+        return cast(Optional[LitellmUserRoles], result)
+    except Exception:
+        return None
+
 
 class CustomOpenID(OpenID):
     team_ids: List[str]
+    user_role: Optional[LitellmUserRoles] = None

litellm/proxy/management_endpoints/ui_sso.py

@@ -58,7 +58,7 @@
     has_admin_ui_access,
 )
 from litellm.proxy.management_endpoints.team_endpoints import new_team, team_member_add
-from litellm.proxy.management_endpoints.types import CustomOpenID
+from litellm.proxy.management_endpoints.types import CustomOpenID, get_litellm_user_role
 from litellm.proxy.utils import (
     PrismaClient,
     ProxyLogging,
@@ -277,6 +277,7 @@ def generic_response_convertor(
         last_name=response.get(generic_user_last_name_attribute_name),
         provider=response.get(generic_provider_attribute_name),
         team_ids=all_teams,
+        user_role=None,
     )
 
 
@@ -1145,7 +1146,7 @@ def get_redirect_url_for_sso(
     ) -> str:
         """
         Get the redirect URL for SSO
-        
+
         Note: existing_key is not added to the URL to avoid changing the callback URL.
         It should be passed via the state parameter instead.
         """
@@ -1348,7 +1349,7 @@ def _get_cli_state(
         Checks the request 'source' if a cli state token was passed in
 
         This is used to authenticate through the CLI login flow.
-        
+
         The state parameter format is: {PREFIX}:{key}:{existing_key}
         - If existing_key is provided, it's included in the state
         - The state parameter is used to pass data through the OAuth flow without changing the callback URL
@@ -1673,22 +1674,49 @@ async def get_microsoft_callback_response(
             access_token=microsoft_sso.access_token
         )
 
+        # Extract app roles from the id_token JWT
+        app_roles = MicrosoftSSOHandler.get_app_roles_from_id_token(
+            id_token=microsoft_sso.id_token
+        )
+        verbose_proxy_logger.debug(f"Extracted app roles from id_token: {app_roles}")
+
+        # Combine groups and app roles
+        user_role: Optional[LitellmUserRoles] = None
+        if app_roles:
+            # Check if any app role is a valid LitellmUserRoles
+            for role_str in app_roles:
+                role = get_litellm_user_role(role_str)
+                if role is not None:
+                    user_role = role
+                    verbose_proxy_logger.debug(
+                        f"Found valid LitellmUserRoles '{role.value}' in app_roles"
+                    )
+                    break
+
+        verbose_proxy_logger.debug(
+            f"Combined team_ids (groups + app roles): {user_team_ids}"
+        )
+
         # 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["app_roles"] = app_roles
             return original_msft_result or {}
 
         result = MicrosoftSSOHandler.openid_from_response(
             response=original_msft_result,
             team_ids=user_team_ids,
+            user_role=user_role,
         )
         return result
 
     @staticmethod
     def openid_from_response(
-        response: Optional[dict], team_ids: List[str]
+        response: Optional[dict],
+        team_ids: List[str],
+        user_role: Optional[LitellmUserRoles],
     ) -> CustomOpenID:
         response = response or {}
         verbose_proxy_logger.debug(f"Microsoft SSO Callback Response: {response}")
@@ -1700,10 +1728,54 @@ def openid_from_response(
             first_name=response.get("givenName"),
             last_name=response.get("surname"),
             team_ids=team_ids,
+            user_role=user_role,
         )
         verbose_proxy_logger.debug(f"Microsoft SSO OpenID Response: {openid_response}")
         return openid_response
 
+    @staticmethod
+    def get_app_roles_from_id_token(id_token: Optional[str]) -> List[str]:
+        """
+        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.
+
+        Args:
+            id_token (Optional[str]): The JWT id_token from Microsoft SSO
+
+        Returns:
+            List[str]: List of app role names assigned to the user
+        """
+        if not id_token:
+            verbose_proxy_logger.debug("No id_token provided for app role extraction")
+            return []
+
+        try:
+            import jwt
+
+            # Decode the JWT without signature verification
+            # (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", [])
+
+            if roles and isinstance(roles, list):
+                verbose_proxy_logger.debug(
+                    f"Found {len(roles)} app role(s) in id_token: {roles}"
+                )
+                return roles
+            else:
+                verbose_proxy_logger.debug(
+                    "No app roles found in id_token or roles claim is not a list"
+                )
+                return []
+
+        except Exception as e:
+            verbose_proxy_logger.error(f"Error extracting app roles from id_token: {e}")
+            return []
+
     @staticmethod
     async def get_user_groups_from_graph_api(
         access_token: Optional[str] = None,