Switch to lru_cache

Author: dandorman

tests/test_session.py

@@ -5,7 +5,7 @@
 import concurrent.futures
 
 from tests.conftest import with_jwks_mock
-from workos.session import AsyncSession, Session, _get_jwks_client, _jwks_cache
+from workos.session import AsyncSession, Session, _get_jwks_client
 from workos.types.user_management.authentication_response import (
     RefreshTokenAuthenticationResponse,
 )
@@ -23,9 +23,9 @@
 class SessionFixtures:
     @pytest.fixture(autouse=True)
     def clear_jwks_cache(self):
-        _jwks_cache.clear()
+        _get_jwks_client.cache_clear()
         yield
-        _jwks_cache.clear()
+        _get_jwks_client.cache_clear()
 
     @pytest.fixture
     def session_constants(self):
@@ -521,7 +521,7 @@ 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 = []

workos/session.py

@@ -1,8 +1,8 @@
 from __future__ import annotations
 from typing import TYPE_CHECKING, List, Protocol
 
+from functools import lru_cache
 import json
-import threading
 from typing import Any, Dict, Optional, Union, cast
 import jwt
 from jwt import PyJWKClient
@@ -22,39 +22,9 @@
     from workos.user_management import AsyncUserManagement, UserManagement
 
 
-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 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
-_jwks_cache = _JWKSClientCache()
-
-
+@lru_cache(maxsize=None)
 def _get_jwks_client(jwks_url: str) -> PyJWKClient:
-    return _jwks_cache.get_client(jwks_url)
+    return PyJWKClient(jwks_url)
 
 
 class SessionModule(Protocol):