Make JWKS client cache threadsafe

Author: dandorman

tests/test_session.py

@@ -2,6 +2,7 @@
 from unittest.mock import AsyncMock, Mock, patch
 import jwt
 from datetime import datetime, timezone
+import concurrent.futures
 
 from tests.conftest import with_jwks_mock
 from workos.session import AsyncSession, Session, _get_jwks_client, _jwks_cache
@@ -22,9 +23,9 @@
 class SessionFixtures:
     @pytest.fixture(autouse=True)
     def clear_jwks_cache(self):
-        _jwks_cache._clients.clear()
+        _jwks_cache.clear()
         yield
-        _jwks_cache._clients.clear()
+        _jwks_cache.clear()
 
     @pytest.fixture
     def session_constants(self):
@@ -520,3 +521,20 @@ def test_jwks_client_caching_different_urls(self):
         # Should be different instances
         assert client1 is not client2
         assert id(client1) != id(client2)
+    
+    def test_jwks_cache_thread_safety(self):
+        url = "https://api.workos.com/sso/jwks/thread_test"
+        clients = []
+
+        def get_client():
+            return _get_jwks_client(url)
+
+        with concurrent.futures.ThreadPoolExecutor(max_workers=10) as executor:
+            futures = [executor.submit(get_client) for _ in range(10)]
+            clients = [future.result() for future in futures]
+
+        first_client = clients[0]
+        for client in clients[1:]:
+            assert (
+                client is first_client
+            ), "All concurrent calls should return the same instance"

workos/session.py

@@ -2,6 +2,7 @@
 from typing import TYPE_CHECKING, List, Protocol
 
 import json
+import threading
 from typing import Any, Dict, Optional, Union, cast
 import jwt
 from jwt import PyJWKClient
@@ -24,11 +25,28 @@
 class _JWKSClientCache:
     def __init__(self) -> None:
         self._clients: Dict[str, PyJWKClient] = {}
+        self._lock = threading.Lock()
 
     def get_client(self, jwks_url: str) -> PyJWKClient:
-        if jwks_url not in self._clients:
-            self._clients[jwks_url] = PyJWKClient(jwks_url)
-        return self._clients[jwks_url]
+        if jwks_url in self._clients:
+            return self._clients[jwks_url]
+
+        with self._lock:
+            if jwks_url in self._clients:
+                return self._clients[jwks_url]
+
+            client = PyJWKClient(jwks_url)
+            self._clients[jwks_url] = client
+            return client
+
+    def clear(self) -> None:
+        """Intended primarily for test cleanup and manual cache invalidation.
+        
+        Warning: If called concurrently with get_client(), some newly created
+        clients might be lost due to lock acquisition ordering.
+        """
+        with self._lock:
+            self._clients.clear()
 
 
 # Module-level cache instance