[Feat] JWT - Sync user roles and team memberships when JWT Auth is used (#11994)

* add JWTLiteLLMRoleMap

* test_sync_user_role_and_teams

* add sync_user_role_and_teams

* test_sync_user_role_and_teams

* fix types

* Sync User Roles and Teams with IDP

* Add test for JWT role mapping to LiteLLM roles

docs/my-website/docs/proxy/token_auth.md

@@ -501,6 +501,145 @@ curl -L -X POST 'http://0.0.0.0:4000/v1/chat/completions' \
 }'
 ```
 
+## [BETA] Sync User Roles and Teams with IDP
+
+Automatically sync user roles and team memberships from your Identity Provider (IDP) to LiteLLM's database. This ensures that user permissions and team memberships in LiteLLM stay in sync with your IDP.
+
+**Note:** This is in beta and might change unexpectedly.
+
+### Use Cases
+
+- **Role Synchronization**: Automatically update user roles in LiteLLM when they change in your IDP
+- **Team Membership Sync**: Keep team memberships in sync between your IDP and LiteLLM
+- **Centralized Access Management**: Manage all user permissions through your IDP while maintaining LiteLLM functionality
+
+### Setup
+
+#### 1. Configure JWT Role Mapping
+
+Map roles from your JWT token to LiteLLM user roles:
+
+```yaml
+general_settings:
+  enable_jwt_auth: True
+  litellm_jwtauth:
+    user_id_jwt_field: "sub"
+    team_ids_jwt_field: "groups"
+    roles_jwt_field: "roles"
+    user_id_upsert: true
+    sync_user_role_and_teams: true # 👈 Enable sync functionality
+    jwt_litellm_role_map: # 👈 Map JWT roles to LiteLLM roles
+      - jwt_role: "ADMIN"
+        litellm_role: "proxy_admin"
+      - jwt_role: "USER"
+        litellm_role: "internal_user"
+      - jwt_role: "VIEWER"
+        litellm_role: "internal_user"
+```
+
+#### 2. JWT Role Mapping Spec
+
+- `jwt_role`: The role name as it appears in your JWT token. Supports wildcard patterns using `fnmatch` (e.g., `"ADMIN_*"` matches `"ADMIN_READ"`, `"ADMIN_WRITE"`, etc.)
+- `litellm_role`: The corresponding LiteLLM user role
+
+**Supported LiteLLM Roles:**
+- `proxy_admin`: Full administrative access
+- `internal_user`: Standard user access
+- `internal_user_view_only`: Read-only access
+
+#### 3. Example JWT Token
+
+```json
+{
+  "sub": "user-123",
+  "roles": ["ADMIN"],
+  "groups": ["team-alpha", "team-beta"],
+  "iat": 1234567890,
+  "exp": 1234567890
+}
+```
+
+### How It Works
+
+When a user makes a request with a JWT token:
+
+1. **Role Sync**: 
+   - LiteLLM checks if the user's role in the JWT matches their role in the database
+   - If different, the user's role is updated in LiteLLM's database
+   - Uses the `jwt_litellm_role_map` to convert JWT roles to LiteLLM roles
+
+2. **Team Membership Sync**:
+   - Compares team memberships from the JWT token with the user's current teams in LiteLLM
+   - Adds the user to new teams found in the JWT
+   - Removes the user from teams not present in the JWT
+
+3. **Database Updates**:
+   - Updates happen automatically during the authentication process
+   - No manual intervention required
+
+### Configuration Options
+
+```yaml
+general_settings:
+  enable_jwt_auth: True
+  litellm_jwtauth:
+    # Required fields
+    user_id_jwt_field: "sub"
+    team_ids_jwt_field: "groups"
+    roles_jwt_field: "roles"
+    
+    # Sync configuration
+    sync_user_role_and_teams: true
+    user_id_upsert: true
+    
+    # Role mapping
+    jwt_litellm_role_map:
+      - jwt_role: "AI_ADMIN_*"  # Wildcard pattern
+        litellm_role: "proxy_admin"
+      - jwt_role: "AI_USER"
+        litellm_role: "internal_user"
+```
+
+### Important Notes
+
+- **Performance**: Sync operations happen during authentication, which may add slight latency
+- **Database Access**: Requires database access for user and team updates
+- **Team Creation**: Teams mentioned in JWT tokens must exist in LiteLLM before sync can assign users to them
+- **Wildcard Support**: JWT role patterns support wildcard matching using `fnmatch`
+
+### Testing the Sync Feature
+
+1. **Create a test user with initial role**:
+
+```bash
+curl -X POST 'http://0.0.0.0:4000/user/new' \
+-H 'Authorization: Bearer <PROXY_MASTER_KEY>' \
+-H 'Content-Type: application/json' \
+-d '{
+    "user_id": "user-123",
+    "user_role": "internal_user"
+}'
+```
+
+2. **Make a request with JWT containing different role**:
+
+```bash
+curl -X POST 'http://0.0.0.0:4000/v1/chat/completions' \
+-H 'Content-Type: application/json' \
+-H 'Authorization: Bearer <JWT_WITH_ADMIN_ROLE>' \
+-d '{
+  "model": "claude-sonnet-4-20250514",
+  "messages": [{"role": "user", "content": "Hello"}]
+}'
+```
+
+3. **Verify the role was updated**:
+
+```bash
+curl -X GET 'http://0.0.0.0:4000/user/info?user_id=user-123' \
+-H 'Authorization: Bearer <PROXY_MASTER_KEY>'
+```
+
 ## All JWT Params
 
 [**See Code**](https://github.com/BerriAI/litellm/blob/b204f0c01c703317d812a1553363ab0cb989d5b6/litellm/proxy/_types.py#L95)

litellm/proxy/_types.py

@@ -2872,6 +2872,11 @@ class RoleMapping(BaseModel):
     internal_role: RBAC_ROLES
 
 
+class JWTLiteLLMRoleMap(BaseModel):
+    jwt_role: str
+    litellm_role: LitellmUserRoles
+
+
 class ScopeMapping(OIDCPermissions):
     scope: str
 
@@ -2943,6 +2948,11 @@ class LiteLLM_JWTAuth(LiteLLMPydanticObjectBase):
     enforce_scope_based_access: bool = False
     enforce_team_based_model_access: bool = False
     custom_validate: Optional[Callable[..., Literal[True]]] = None
+    #########################################################
+    # Fields for syncing user team membership and roles with IDP provider
+    jwt_litellm_role_map: Optional[List[JWTLiteLLMRoleMap]] = None
+    sync_user_role_and_teams: bool = False
+    #########################################################
 
     def __init__(self, **kwargs: Any) -> None:
         # get the attribute names for this Pydantic model

litellm/proxy/auth/handle_jwt.py

@@ -6,6 +6,7 @@ Currently only supports admin.
 JWT token must have 'litellm_proxy_admin' in scope.
 """
 
+import fnmatch
 import json
 import os
 from typing import Any, List, Literal, Optional, Set, Tuple, cast
@@ -258,6 +259,21 @@ class JWTHandler:
             user_roles = default_value
         return user_roles
 
+    def map_jwt_role_to_litellm_role(self, token: dict) -> Optional[LitellmUserRoles]:
+        """Map roles from JWT to LiteLLM user roles"""
+        if not self.litellm_jwtauth.jwt_litellm_role_map:
+            return None
+
+        jwt_roles = self.get_jwt_role(token=token, default_value=[])
+        if not jwt_roles:
+            return None
+
+        for mapping in self.litellm_jwtauth.jwt_litellm_role_map:
+            for role in jwt_roles:
+                if fnmatch.fnmatch(role, mapping.jwt_role):
+                    return mapping.litellm_role
+        return None
+
     def get_jwt_role(
         self, token: dict, default_value: Optional[List[str]]
     ) -> Optional[List[str]]:
@@ -913,6 +929,55 @@ class JWTAuthManager:
                 raise e
         return None
 
+    @staticmethod
+    async def sync_user_role_and_teams(
+        jwt_handler: JWTHandler,
+        jwt_valid_token: dict,
+        user_object: Optional[LiteLLM_UserTable],
+        prisma_client: Optional[PrismaClient],
+    ) -> None:
+        """
+        Sync user role and team memberships with JWT claims
+
+        The goal of this method is to ensure:
+        1. The user role on LiteLLM DB is in sync with the IDP provider role
+        2. The user is a member of the teams specified in the JWT token
+
+        This method is only called if sync_user_role_and_teams is set to True in the JWT config.
+        """
+        if not jwt_handler.litellm_jwtauth.sync_user_role_and_teams:
+            return None
+
+        if user_object is None or prisma_client is None:
+            return None
+
+        # Update user role
+        new_role = jwt_handler.map_jwt_role_to_litellm_role(jwt_valid_token)
+        if new_role and user_object.user_role != new_role.value:
+            await prisma_client.db.litellm_usertable.update(
+                where={"user_id": user_object.user_id},
+                data={"user_role": new_role.value},
+            )
+            user_object.user_role = new_role.value
+
+        # Sync team memberships
+        jwt_team_ids = set(jwt_handler.get_team_ids_from_jwt(jwt_valid_token))
+        existing_teams = set(user_object.teams or [])
+        teams_to_add = jwt_team_ids - existing_teams
+        teams_to_remove = existing_teams - jwt_team_ids
+        if teams_to_add or teams_to_remove:
+            from litellm.proxy.management_endpoints.scim.scim_v2 import (
+                patch_team_membership,
+            )
+
+            await patch_team_membership(
+                user_id=user_object.user_id,
+                teams_ids_to_add_user_to=list(teams_to_add),
+                teams_ids_to_remove_user_from=list(teams_to_remove),
+            )
+            user_object.teams = list(jwt_team_ids)
+        return None
+
     @staticmethod
     async def auth_builder(
         api_key: str,
@@ -1034,6 +1099,13 @@ class JWTAuthManager:
             proxy_logging_obj=proxy_logging_obj,
         )
 
+        await JWTAuthManager.sync_user_role_and_teams(
+            jwt_handler=jwt_handler,
+            jwt_valid_token=jwt_valid_token,
+            user_object=user_object,
+            prisma_client=prisma_client,
+        )
+
         ## MAP USER TO TEAMS
         await JWTAuthManager.map_user_to_teams(
             user_object=user_object,

tests/test_litellm/proxy/auth/test_handle_jwt.py

@@ -3,7 +3,7 @@ from unittest.mock import AsyncMock, patch
 import pytest
 
 from litellm.proxy._types import (
-    JWTAuthBuilderResult,
+    JWTLiteLLMRoleMap,
     LiteLLM_JWTAuth,
     LiteLLM_TeamTable,
     LiteLLM_UserTable,
@@ -296,3 +296,154 @@ async def test_auth_builder_non_proxy_admin_user_role():
         assert result["is_proxy_admin"] is False
         assert result["user_object"] == user_object
         assert result["user_id"] == "test_user_1"
+
+
+@pytest.mark.asyncio
+async def test_sync_user_role_and_teams():
+    from unittest.mock import MagicMock
+
+    # Create mock objects for required types
+    mock_user_api_key_cache = MagicMock()
+    mock_proxy_logging_obj = MagicMock()
+    
+    jwt_handler = JWTHandler()
+    jwt_handler.update_environment(
+        prisma_client=None,
+        user_api_key_cache=mock_user_api_key_cache,
+        litellm_jwtauth=LiteLLM_JWTAuth(
+            jwt_litellm_role_map=[
+                JWTLiteLLMRoleMap(jwt_role="ADMIN", litellm_role=LitellmUserRoles.PROXY_ADMIN)
+            ],
+            roles_jwt_field="roles",
+            team_ids_jwt_field="my_id_teams",
+            sync_user_role_and_teams=True
+        ),
+    )
+
+    token = {"roles": ["ADMIN"], "my_id_teams": ["team1", "team2"]}
+
+    user = LiteLLM_UserTable(user_id="u1", user_role=LitellmUserRoles.INTERNAL_USER.value, teams=["team2"])
+
+    prisma = AsyncMock()
+    prisma.db.litellm_usertable.update = AsyncMock()
+
+    with patch(
+        "litellm.proxy.management_endpoints.scim.scim_v2.patch_team_membership",
+        new_callable=AsyncMock,
+    ) as mock_patch:
+        await JWTAuthManager.sync_user_role_and_teams(jwt_handler, token, user, prisma)
+
+    prisma.db.litellm_usertable.update.assert_called_once()
+    mock_patch.assert_called_once()
+    assert user.user_role == LitellmUserRoles.PROXY_ADMIN.value
+    assert set(user.teams) == {"team1", "team2"}
+
+
+@pytest.mark.asyncio
+async def test_map_jwt_role_to_litellm_role():
+    """Test JWT role mapping to LiteLLM roles with various patterns"""
+    from unittest.mock import MagicMock
+
+    # Create mock objects for required types
+    mock_user_api_key_cache = MagicMock()
+    
+    jwt_handler = JWTHandler()
+    jwt_handler.update_environment(
+        prisma_client=None,
+        user_api_key_cache=mock_user_api_key_cache,
+        litellm_jwtauth=LiteLLM_JWTAuth(
+            jwt_litellm_role_map=[
+                # Exact match
+                JWTLiteLLMRoleMap(jwt_role="ADMIN", litellm_role=LitellmUserRoles.PROXY_ADMIN),
+                # Wildcard patterns
+                JWTLiteLLMRoleMap(jwt_role="user_*", litellm_role=LitellmUserRoles.INTERNAL_USER),
+                JWTLiteLLMRoleMap(jwt_role="team_?", litellm_role=LitellmUserRoles.TEAM),
+                JWTLiteLLMRoleMap(jwt_role="dev_[123]", litellm_role=LitellmUserRoles.INTERNAL_USER),
+            ],
+            roles_jwt_field="roles"
+        ),
+    )
+
+    # Test exact match
+    token = {"roles": ["ADMIN"]}
+    result = jwt_handler.map_jwt_role_to_litellm_role(token)
+    assert result == LitellmUserRoles.PROXY_ADMIN
+
+    # Test wildcard match with *
+    token = {"roles": ["user_manager"]}
+    result = jwt_handler.map_jwt_role_to_litellm_role(token)
+    assert result == LitellmUserRoles.INTERNAL_USER
+
+    token = {"roles": ["user_"]}
+    result = jwt_handler.map_jwt_role_to_litellm_role(token)
+    assert result == LitellmUserRoles.INTERNAL_USER
+
+    # Test wildcard match with ?
+    token = {"roles": ["team_1"]}
+    result = jwt_handler.map_jwt_role_to_litellm_role(token)
+    assert result == LitellmUserRoles.TEAM
+
+    token = {"roles": ["team_a"]}
+    result = jwt_handler.map_jwt_role_to_litellm_role(token)
+    assert result == LitellmUserRoles.TEAM
+
+    # Test character class match
+    token = {"roles": ["dev_1"]}
+    result = jwt_handler.map_jwt_role_to_litellm_role(token)
+    assert result == LitellmUserRoles.INTERNAL_USER
+
+    token = {"roles": ["dev_2"]}
+    result = jwt_handler.map_jwt_role_to_litellm_role(token)
+    assert result == LitellmUserRoles.INTERNAL_USER
+
+    # Test no match
+    token = {"roles": ["unknown_role"]}
+    result = jwt_handler.map_jwt_role_to_litellm_role(token)
+    assert result is None
+
+    # Test multiple roles - should return first mapping match
+    token = {"roles": ["user_test", "ADMIN"]}
+    result = jwt_handler.map_jwt_role_to_litellm_role(token)
+    assert result == LitellmUserRoles.PROXY_ADMIN  # ADMIN matches first mapping
+
+    # Test empty roles
+    token = {"roles": []}
+    result = jwt_handler.map_jwt_role_to_litellm_role(token)
+    assert result is None
+
+    # Test no roles field
+    token = {}
+    result = jwt_handler.map_jwt_role_to_litellm_role(token)
+    assert result is None
+
+    # Test no role mappings configured
+    jwt_handler.litellm_jwtauth.jwt_litellm_role_map = None
+    token = {"roles": ["ADMIN"]}
+    result = jwt_handler.map_jwt_role_to_litellm_role(token)
+    assert result is None
+
+    # Test empty role mappings
+    jwt_handler.litellm_jwtauth.jwt_litellm_role_map = []
+    token = {"roles": ["ADMIN"]}
+    result = jwt_handler.map_jwt_role_to_litellm_role(token)
+    assert result is None
+
+    # Test patterns that don't match character classes
+    jwt_handler.litellm_jwtauth.jwt_litellm_role_map = [
+        JWTLiteLLMRoleMap(jwt_role="dev_[123]", litellm_role=LitellmUserRoles.INTERNAL_USER),
+    ]
+    token = {"roles": ["dev_4"]}  # 4 is not in [123]
+    result = jwt_handler.map_jwt_role_to_litellm_role(token)
+    assert result is None
+
+    # Test ? pattern that requires exactly one character
+    jwt_handler.litellm_jwtauth.jwt_litellm_role_map = [
+        JWTLiteLLMRoleMap(jwt_role="team_?", litellm_role=LitellmUserRoles.TEAM),
+    ]
+    token = {"roles": ["team_12"]}  # More than one character after underscore
+    result = jwt_handler.map_jwt_role_to_litellm_role(token)
+    assert result is None
+
+    token = {"roles": ["team_"]}  # No character after underscore
+    result = jwt_handler.map_jwt_role_to_litellm_role(token)
+    assert result is None