Cache JWKS clients per URL

[dandorman]

Description

Previously, any time a session was initiated, a separate JWKS client was created for each instance. Even though each instance individually cached the results of fetching the JWKS, these weren't getting meaningfully re-used.

This PR creates a cache that stores JWKS clients per JWKS URL. Rather than instantiating their own instances of the JWKS client, the session instances will now fetch them from the cache.

Documentation

Does this require changes to the WorkOS Docs? E.g. the API Reference or code snippets need updates.

[ ] Yes

If yes, link a related docs PR and add a docs maintainer as a reviewer. Their approval is required.

[greptile-apps]

Greptile Summary

This PR introduces a performance optimization for JWKS (JSON Web Key Set) client management in the WorkOS Python SDK. Previously, each Session, AsyncSession, and SessionModule instance would create its own PyJWKClient when initialized, even when multiple sessions used the same JWKS URL. While these individual clients cached JWKS responses internally, this caching wasn't being shared across session instances, leading to redundant client creation and suboptimal resource usage.

The changes implement a module-level cache (_JWKSClientCache) that stores PyJWKClient instances keyed by their JWKS URL. The cache uses a simple dictionary-based approach with lazy initialization - when a JWKS client is requested for a URL, the cache either returns an existing client or creates a new one if none exists for that URL. All three session classes now use the new _get_jwks_client() function instead of directly instantiating PyJWKClient objects.

This optimization is particularly beneficial for applications that create multiple session instances, as they can now share the same underlying JWKS client and benefit from its internal caching mechanisms. The change maintains the existing API while improving memory usage and performance through better resource sharing.

Comprehensive test coverage has been added to verify the caching behavior, including a fixture that ensures proper cache cleanup between tests to maintain test isolation.

Confidence score: 4/5

• This PR is very safe to merge with minimal risk of causing production issues
• The implementation is straightforward, well-tested, and maintains backward compatibility while providing clear performance benefits
• workos/session.py needs careful review to ensure the cache implementation is thread-safe in concurrent environments

_{2 files reviewed, 1 comment}

_{Edit Code Review Bot Settings | Greptile}