perf(session): cache JWKS per client with 5-minute TTL

Previously every authenticate() call issued a live HTTP GET to the JWKS
endpoint, making each session check dependent on an external round-trip
and inflating latency. Add an in-memory cache on SessionManager keyed by
client ID with a 300-second TTL, plus a force-refresh path on `kid` miss
so newly-rotated signing keys are still discovered without waiting for
TTL expiry.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: gjtorikian

lib/SessionManager.php

@@ -13,6 +13,23 @@
 
 class SessionManager
 {
+    /**
+     * In-memory JWKS cache, keyed by client ID. Values are
+     * `['keys' => array, 'fetched_at' => int]`. Cache lives for the
+     * lifetime of the SessionManager instance and is bypassed when a
+     * token's `kid` isn't found, so key rotation still resolves quickly.
+     *
+     * @var array<string, array{keys: array, fetched_at: int}>
+     */
+    private array $jwksCache = [];
+
+    /**
+     * JWKS cache TTL in seconds. WorkOS rotates signing keys on the order
+     * of weeks, so a few minutes is plenty to absorb traffic spikes
+     * without making session checks dependent on a live JWKS round-trip.
+     */
+    private const JWKS_CACHE_TTL_SECONDS = 300;
+
     public function __construct(
         private readonly HttpClient $client,
     ) {
@@ -302,6 +319,33 @@ public function fetchJwks(string $clientId): array
         );
     }
 
+    /**
+     * Return the JWKS for `$clientId`, served from an in-memory cache
+     * with a {@see JWKS_CACHE_TTL_SECONDS}-second TTL. Set
+     * `$forceRefresh` to bypass the cache after a `kid` miss, which
+     * lets newly-rotated keys be discovered without waiting for TTL
+     * expiry.
+     *
+     * @return array<string, mixed>
+     */
+    private function getCachedJwks(string $clientId, bool $forceRefresh = false): array
+    {
+        $now = time();
+        $entry = $this->jwksCache[$clientId] ?? null;
+        if (
+            !$forceRefresh
+            && $entry !== null
+            && ($now - $entry['fetched_at']) < self::JWKS_CACHE_TTL_SECONDS
+        ) {
+            return $entry['keys'];
+        }
+
+        $keys = $this->fetchJwks($clientId);
+        $this->jwksCache[$clientId] = ['keys' => $keys, 'fetched_at' => $now];
+
+        return $keys;
+    }
+
     /**
      * Algorithms permitted on the JWS header. WorkOS access tokens are signed
      * with RS256; no other algorithm is accepted, in particular `none` is
@@ -370,8 +414,14 @@ private function decodeAccessToken(
             throw new \InvalidArgumentException('JWT header missing kid');
         }
 
-        $jwks = $this->fetchJwks($clientId);
+        // Try the cached JWKS first; if the `kid` isn't present, force a
+        // refresh once to handle key rotation, then fail if still unknown.
+        $jwks = $this->getCachedJwks($clientId);
         $jwk = self::findJwkByKid($jwks, $kid);
+        if ($jwk === null) {
+            $jwks = $this->getCachedJwks($clientId, forceRefresh: true);
+            $jwk = self::findJwkByKid($jwks, $kid);
+        }
         if ($jwk === null) {
             throw new \InvalidArgumentException('No JWKS key matches JWT kid');
         }

tests/SessionManagerTest.php

@@ -289,4 +289,80 @@ public function testAuthenticateRejectsUnknownKid(): void
         $this->assertFalse($result['authenticated']);
         $this->assertSame('invalid_jwt', $result['reason']);
     }
+
+    public function testAuthenticateCachesJwksAcrossCalls(): void
+    {
+        [$jwks, $jwt] = $this->buildSignedJwt([
+            'sid' => 'session_test',
+            'exp' => time() + 3600,
+        ]);
+
+        $sealed = SessionManager::sealSessionFromAuthResponse(
+            accessToken: $jwt,
+            refreshToken: 'ref_test',
+            cookiePassword: $this->cookiePassword,
+        );
+
+        // Only ONE JWKS response is queued; a second fetch would make
+        // the MockHandler throw, so successful back-to-back authenticate
+        // calls prove the cache served the second one.
+        $client = $this->createMockClient([['status' => 200, 'body' => $jwks]]);
+        $sessionManager = $client->sessionManager();
+
+        $first = $sessionManager->authenticate(
+            sessionData: $sealed,
+            cookiePassword: $this->cookiePassword,
+            clientId: 'client_123',
+        );
+        $second = $sessionManager->authenticate(
+            sessionData: $sealed,
+            cookiePassword: $this->cookiePassword,
+            clientId: 'client_123',
+        );
+
+        $this->assertTrue($first['authenticated']);
+        $this->assertTrue($second['authenticated']);
+        $this->assertCount(1, $this->requestHistory);
+    }
+
+    public function testAuthenticateRefreshesJwksOnUnknownKid(): void
+    {
+        // Seed the cache with a JWKS that only knows `kid_old`, then
+        // present a token signed with `kid_new`. The kid-miss path
+        // should force a refresh and find `kid_new` in the second
+        // JWKS response.
+        [$oldJwks] = $this->buildSignedJwt(
+            ['sid' => 'session_old', 'exp' => time() + 3600],
+            'RS256',
+            'kid_old',
+        );
+
+        [$newJwks, $newJwt] = $this->buildSignedJwt(
+            ['sid' => 'session_new', 'exp' => time() + 3600],
+            'RS256',
+            'kid_new',
+        );
+
+        $sealedNew = SessionManager::sealSessionFromAuthResponse(
+            accessToken: $newJwt,
+            refreshToken: 'ref_test',
+            cookiePassword: $this->cookiePassword,
+        );
+
+        $client = $this->createMockClient([
+            ['status' => 200, 'body' => $oldJwks],
+            ['status' => 200, 'body' => $newJwks],
+        ]);
+        $sessionManager = $client->sessionManager();
+
+        $result = $sessionManager->authenticate(
+            sessionData: $sealedNew,
+            cookiePassword: $this->cookiePassword,
+            clientId: 'client_123',
+        );
+
+        $this->assertTrue($result['authenticated']);
+        $this->assertSame('session_new', $result['session_id']);
+        $this->assertCount(2, $this->requestHistory);
+    }
 }