Cache JWKS clients per URL

Author: dandorman

tests/test_session.py

@@ -4,7 +4,7 @@
 from datetime import datetime, timezone
 
 from tests.conftest import with_jwks_mock
-from workos.session import AsyncSession, Session
+from workos.session import AsyncSession, Session, _get_jwks_client, _jwks_cache
 from workos.types.user_management.authentication_response import (
     RefreshTokenAuthenticationResponse,
 )
@@ -20,6 +20,12 @@
 
 
 class SessionFixtures:
+    @pytest.fixture(autouse=True)
+    def clear_jwks_cache(self):
+        _jwks_cache._clients.clear()
+        yield
+        _jwks_cache._clients.clear()
+
     @pytest.fixture
     def session_constants(self):
         # Generate RSA key pair for testing
@@ -491,3 +497,26 @@ async def test_refresh_success_with_aud_claim(
         response = await session.refresh()
 
         assert isinstance(response, RefreshWithSessionCookieSuccessResponse)
+
+
+class TestJWKSCaching:
+    def test_jwks_client_caching_same_url(self):
+        url = "https://api.workos.com/sso/jwks/test"
+
+        client1 = _get_jwks_client(url)
+        client2 = _get_jwks_client(url)
+
+        # Should be the exact same instance
+        assert client1 is client2
+        assert id(client1) == id(client2)
+
+    def test_jwks_client_caching_different_urls(self):
+        url1 = "https://api.workos.com/sso/jwks/client1"
+        url2 = "https://api.workos.com/sso/jwks/client2"
+
+        client1 = _get_jwks_client(url1)
+        client2 = _get_jwks_client(url2)
+
+        # Should be different instances
+        assert client1 is not client2
+        assert id(client1) != id(client2)

workos/session.py

@@ -21,6 +21,24 @@
     from workos.user_management import AsyncUserManagement, UserManagement
 
 
+class _JWKSClientCache:
+    def __init__(self) -> None:
+        self._clients: Dict[str, PyJWKClient] = {}
+
+    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]
+
+
+# Module-level cache instance
+_jwks_cache = _JWKSClientCache()
+
+
+def _get_jwks_client(jwks_url: str) -> PyJWKClient:
+    return _jwks_cache.get_client(jwks_url)
+
+
 class SessionModule(Protocol):
     user_management: "UserManagementModule"
     client_id: str
@@ -46,7 +64,7 @@ def __init__(
         self.session_data = session_data
         self.cookie_password = cookie_password
 
-        self.jwks = PyJWKClient(self.user_management.get_jwks_url())
+        self.jwks = _get_jwks_client(self.user_management.get_jwks_url())
 
         # Algorithms are hardcoded for security reasons. See https://pyjwt.readthedocs.io/en/stable/algorithms.html#specifying-an-algorithm
         self.jwk_algorithms = ["RS256"]
@@ -164,7 +182,7 @@ def __init__(
         self.session_data = session_data
         self.cookie_password = cookie_password
 
-        self.jwks = PyJWKClient(self.user_management.get_jwks_url())
+        self.jwks = _get_jwks_client(self.user_management.get_jwks_url())
 
         # Algorithms are hardcoded for security reasons. See https://pyjwt.readthedocs.io/en/stable/algorithms.html#specifying-an-algorithm
         self.jwk_algorithms = ["RS256"]
@@ -254,7 +272,7 @@ def __init__(
         self.session_data = session_data
         self.cookie_password = cookie_password
 
-        self.jwks = PyJWKClient(self.user_management.get_jwks_url())
+        self.jwks = _get_jwks_client(self.user_management.get_jwks_url())
 
         # Algorithms are hardcoded for security reasons. See https://pyjwt.readthedocs.io/en/stable/algorithms.html#specifying-an-algorithm
         self.jwk_algorithms = ["RS256"]