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.
This commit is contained in:
Ryan Crabbe 2026-03-26 13:06:01 -07:00 • committed by Ryan Crabbe
parent 8f425ec3ff
commit 2ece79930b
No known key found for this signature in database

View file

@ -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