From 2ece79930b1b2884b76a33582c91615f2c70a544 Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Thu, 26 Mar 2026 13:06:01 -0700 Subject: [PATCH 1/2] fix(jwt): invalidate user cache after role/team sync updates sync_user_role_and_teams updates the DB when a user's JWT role changes, but the in-memory cache retained the stale role until TTL expiry. This caused subsequent requests to see the old role for up to 60 seconds. Fix: accept user_api_key_cache parameter and re-cache the updated user object after both role and team membership DB writes. --- litellm/proxy/auth/handle_jwt.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index bfad9f0c3c..d24e710271 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -1324,6 +1324,7 @@ class JWTAuthManager: jwt_valid_token: dict, user_object: Optional[LiteLLM_UserTable], prisma_client: Optional[PrismaClient], + user_api_key_cache: Optional[DualCache] = None, ) -> None: """ Sync user role and team memberships with JWT claims @@ -1348,6 +1349,11 @@ class JWTAuthManager: data={"user_role": new_role.value}, ) user_object.user_role = new_role.value + if user_api_key_cache is not None: + await user_api_key_cache.async_set_cache( + key=user_object.user_id, + value=user_object.model_dump(), + ) # Sync team memberships jwt_team_ids = set(jwt_handler.get_team_ids_from_jwt(jwt_valid_token)) @@ -1365,6 +1371,11 @@ class JWTAuthManager: teams_ids_to_remove_user_from=list(teams_to_remove), ) user_object.teams = list(jwt_team_ids) + if user_api_key_cache is not None: + await user_api_key_cache.async_set_cache( + key=user_object.user_id, + value=user_object.model_dump(), + ) return None @staticmethod @@ -1536,6 +1547,7 @@ class JWTAuthManager: jwt_valid_token=jwt_valid_token, user_object=user_object, prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, ) ## MAP USER TO TEAMS From dd11e778529cc3582d037b7edd0d49703ae028f6 Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Fri, 27 Mar 2026 19:45:13 -0700 Subject: [PATCH 2/2] fix: add explicit TTL to cache writes and test coverage for user cache invalidation Add DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL to both async_set_cache calls in sync_user_role_and_teams for consistency with all other user cache writes. Add 3 tests covering cache invalidation on role change, team change, and no-op when nothing changes. --- litellm/proxy/auth/handle_jwt.py | 3 + .../proxy/auth/test_handle_jwt.py | 117 ++++++++++++++++++ 2 files changed, 120 insertions(+) diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index d24e710271..202f51e0cb 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -18,6 +18,7 @@ from fastapi import HTTPException from litellm._logging import verbose_proxy_logger from litellm.caching.caching import DualCache +from litellm.constants import DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL from litellm.litellm_core_utils.dot_notation_indexing import get_nested_value from litellm.llms.custom_httpx.httpx_handler import HTTPHandler from litellm.proxy._types import ( @@ -1353,6 +1354,7 @@ class JWTAuthManager: await user_api_key_cache.async_set_cache( key=user_object.user_id, value=user_object.model_dump(), + ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL, ) # Sync team memberships @@ -1375,6 +1377,7 @@ class JWTAuthManager: await user_api_key_cache.async_set_cache( key=user_object.user_id, value=user_object.model_dump(), + ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL, ) return None diff --git a/tests/test_litellm/proxy/auth/test_handle_jwt.py b/tests/test_litellm/proxy/auth/test_handle_jwt.py index 11939f0fdd..ada67fbba8 100644 --- a/tests/test_litellm/proxy/auth/test_handle_jwt.py +++ b/tests/test_litellm/proxy/auth/test_handle_jwt.py @@ -339,6 +339,123 @@ async def test_sync_user_role_and_teams(): assert set(user.teams) == {"team1", "team2"} +@pytest.mark.asyncio +async def test_sync_user_role_and_teams_cache_invalidation_on_role_change(): + """Test that user cache is updated when role changes.""" + mock_cache = AsyncMock() + + jwt_handler = JWTHandler() + jwt_handler.update_environment( + prisma_client=None, + user_api_key_cache=AsyncMock(), + 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"]} + user = LiteLLM_UserTable( + user_id="u1", + user_role=LitellmUserRoles.INTERNAL_USER.value, + teams=["team1"], # teams already match — only role differs + ) + + prisma = AsyncMock() + prisma.db.litellm_usertable.update = AsyncMock() + + await JWTAuthManager.sync_user_role_and_teams( + jwt_handler, token, user, prisma, user_api_key_cache=mock_cache + ) + + mock_cache.async_set_cache.assert_called_once() + call_kwargs = mock_cache.async_set_cache.call_args + assert call_kwargs.kwargs["key"] == "u1" + assert call_kwargs.kwargs["value"]["user_role"] == LitellmUserRoles.PROXY_ADMIN.value + + +@pytest.mark.asyncio +async def test_sync_user_role_and_teams_cache_invalidation_on_team_change(): + """Test that user cache is updated when team memberships change.""" + mock_cache = AsyncMock() + + jwt_handler = JWTHandler() + jwt_handler.update_environment( + prisma_client=None, + user_api_key_cache=AsyncMock(), + 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.PROXY_ADMIN.value, # role already matches + teams=["team2"], # teams differ + ) + + prisma = AsyncMock() + prisma.db.litellm_usertable.update = AsyncMock() + + with patch( + "litellm.proxy.management_endpoints.scim.scim_v2.patch_team_membership", + new_callable=AsyncMock, + ): + await JWTAuthManager.sync_user_role_and_teams( + jwt_handler, token, user, prisma, user_api_key_cache=mock_cache + ) + + mock_cache.async_set_cache.assert_called_once() + call_kwargs = mock_cache.async_set_cache.call_args + assert call_kwargs.kwargs["key"] == "u1" + assert set(call_kwargs.kwargs["value"]["teams"]) == {"team1", "team2"} + + +@pytest.mark.asyncio +async def test_sync_user_role_and_teams_no_cache_write_when_nothing_changes(): + """Test that cache is NOT written when role and teams already match.""" + mock_cache = AsyncMock() + + jwt_handler = JWTHandler() + jwt_handler.update_environment( + prisma_client=None, + user_api_key_cache=AsyncMock(), + 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"]} + user = LiteLLM_UserTable( + user_id="u1", + user_role=LitellmUserRoles.PROXY_ADMIN.value, + teams=["team1"], + ) + + prisma = AsyncMock() + + await JWTAuthManager.sync_user_role_and_teams( + jwt_handler, token, user, prisma, user_api_key_cache=mock_cache + ) + + mock_cache.async_set_cache.assert_not_called() + + @pytest.mark.asyncio async def test_map_jwt_role_to_litellm_role(): """Test JWT role mapping to LiteLLM roles with various patterns"""