Cache JWKS across isolates and serve stale keys on refresh failure (#1665)

* Cache JWKS across isolates and serve stale keys on refresh failure

The module-scope JWKS cache only helps while an isolate stays warm. When
workos.session.local_verify started missing ~92% of the time, every miss
paid a fresh upstream fetch (p50 3s) and the tail crossed the 5s timeout,
so a slow key server surfaced as AbortError -> 500 on authenticated API
routes and 503 on /mcp.

Add a cross-isolate store (Workers Cache API by default, injectable, null
to disable) so a cold isolate reads keys colo-locally instead of going
upstream, and make the resolver stale-while-revalidate: past ttlMs a
usable key set is served immediately and refreshed in the background, and
a failed refresh keeps serving the last good keys until staleMaxMs rather
than failing the verify. Only a fully cold path blocks on the network.

Bounded by staleMaxMs (24h) and unchanged forced-refresh on unknown kid,
so a retired key cannot be honoured indefinitely.

* Keep JWKS latency attribution honest under stale-while-revalidate

Background revalidation moves fetchCount without the caller waiting on it,
so jwks.fetched_during_verify would report true for verifies that paid
nothing. Track blockingFetchCount (fetches a verify awaited) and
storeHitCount (cold-isolate reads served by the cross-isolate store), and
attribute the span off the blocking counter.

Author: RhysSullivan

apps/cloud/src/auth/jwks-cache.node.test.ts

@@ -9,7 +9,7 @@ import {
   type KeyLike,
 } from "jose";
 
-import { createCachedRemoteJWKSet } from "./jwks-cache";
+import { createCachedRemoteJWKSet, type JwksStore, type StoredJwks } from "./jwks-cache";
 
 const issuer = "https://test-authkit.example.com";
 const audience = "client_test_fixture";
@@ -65,6 +65,40 @@ const makeFetchHarness = (initialKeys: ReadonlyArray<JWK>): FetchHarness => {
   };
 };
 
+interface StoreHarness extends JwksStore {
+  readonly seed: (stored: StoredJwks) => void;
+  readonly reads: () => number;
+  readonly writes: () => number;
+}
+
+/** Stands in for the Workers Cache API: shared across "isolates", in memory. */
+const makeStoreHarness = (): StoreHarness => {
+  const entries = new Map<string, StoredJwks>();
+  let reads = 0;
+  let writes = 0;
+  return {
+    get: async (url) => {
+      reads++;
+      return entries.get(url.toString()) ?? null;
+    },
+    put: async (url, stored) => {
+      writes++;
+      entries.set(url.toString(), stored);
+    },
+    seed: (stored) => {
+      entries.set(jwksUrl.toString(), stored);
+    },
+    reads: () => reads,
+    writes: () => writes,
+  };
+};
+
+/** A fetch that always fails, standing in for a slow/down key server. */
+const failingFetch: typeof globalThis.fetch = async () => {
+  // oxlint-disable-next-line executor/no-try-catch-or-throw, executor/no-error-constructor -- test double: `fetch` signals an unreachable key server by rejecting
+  throw new Error("JWKS endpoint unreachable");
+};
+
 describe("createCachedRemoteJWKSet", () => {
   it("FAILING-WITHOUT-CACHE: N verifications hit JWKS endpoint only once within TTL", async () => {
     const kp = await generateRotatableKeypair("k1");
@@ -168,4 +202,151 @@ describe("createCachedRemoteJWKSet", () => {
     await jwtVerify(t1, jwks, { issuer, audience });
     expect(harness.callCount()).toBe(2);
   });
+
+  // -------------------------------------------------------------------------
+  // Cross-isolate store — the cold-isolate path that caused the 2026-08-18
+  // 500s: no module-scope entry, so every verify paid an upstream fetch.
+  // -------------------------------------------------------------------------
+
+  it("a cold isolate serves from the store instead of going upstream", async () => {
+    const kp = await generateRotatableKeypair("k1");
+    const store = makeStoreHarness();
+
+    // First isolate: cold everywhere, so it fetches and populates the store.
+    const warm = makeFetchHarness([kp.publicJwk]);
+    const first = createCachedRemoteJWKSet(jwksUrl, { fetch: warm.fetch, store });
+    const t1 = await sign(kp);
+    await jwtVerify(t1, first, { issuer, audience });
+    expect(warm.callCount()).toBe(1);
+    expect(store.writes()).toBe(1);
+
+    // A brand-new isolate (fresh module scope). If it reaches upstream at all
+    // the fetch fails, so verifying proves the store answered.
+    const second = createCachedRemoteJWKSet(jwksUrl, { fetch: failingFetch, store });
+    const { payload } = await jwtVerify(t1, second, { issuer, audience });
+    expect(payload.sub).toBe("user_test");
+    expect(store.reads()).toBeGreaterThan(0);
+  });
+
+  it("keeps serving the last good keys when the key server is down", async () => {
+    const kp = await generateRotatableKeypair("k1");
+    const harness = makeFetchHarness([kp.publicJwk]);
+    let upstreamUp = true;
+    const flaky: typeof globalThis.fetch = (...args) =>
+      upstreamUp ? harness.fetch(...args) : failingFetch(...args);
+
+    // ttl short enough that the next call is past it, stale window generous.
+    const jwks = createCachedRemoteJWKSet(jwksUrl, {
+      fetch: flaky,
+      store: null,
+      ttlMs: 10,
+      staleMaxMs: 60_000,
+    });
+
+    const token = await sign(kp);
+    await jwtVerify(token, jwks, { issuer, audience });
+
+    upstreamUp = false;
+    await new Promise((r) => setTimeout(r, 20));
+
+    // Past the TTL with a dead upstream: the cached keys still verify.
+    const { payload } = await jwtVerify(token, jwks, { issuer, audience });
+    expect(payload.sub).toBe("user_test");
+  });
+
+  it("stops serving stale keys once the stale window closes", async () => {
+    const kp = await generateRotatableKeypair("k1");
+    const harness = makeFetchHarness([kp.publicJwk]);
+    let upstreamUp = true;
+    const flaky: typeof globalThis.fetch = (...args) =>
+      upstreamUp ? harness.fetch(...args) : failingFetch(...args);
+
+    const jwks = createCachedRemoteJWKSet(jwksUrl, {
+      fetch: flaky,
+      store: null,
+      ttlMs: 5,
+      staleMaxMs: 10,
+    });
+
+    const token = await sign(kp);
+    await jwtVerify(token, jwks, { issuer, audience });
+
+    upstreamUp = false;
+    await new Promise((r) => setTimeout(r, 30));
+
+    // Beyond staleMaxMs the keys are no longer trustworthy — fail, don't
+    // silently honour a key set we can no longer confirm.
+    await expect(jwtVerify(token, jwks, { issuer, audience })).rejects.toThrow();
+  });
+
+  it("attributes background revalidation to fetchCount but not blockingFetchCount", async () => {
+    const kp = await generateRotatableKeypair("k1");
+    const harness = makeFetchHarness([kp.publicJwk]);
+    const jwks = createCachedRemoteJWKSet(jwksUrl, {
+      fetch: harness.fetch,
+      store: null,
+      ttlMs: 10,
+      staleMaxMs: 60_000,
+    });
+
+    const token = await sign(kp);
+    await jwtVerify(token, jwks, { issuer, audience });
+    expect(jwks.inspect().blockingFetchCount).toBe(1);
+
+    await new Promise((r) => setTimeout(r, 20));
+
+    // Past the TTL: served from the stale entry, revalidated behind it. The
+    // caller waited on nothing, so latency must not be attributed to it.
+    const before = jwks.inspect();
+    await jwtVerify(token, jwks, { issuer, audience });
+    const after = jwks.inspect();
+
+    expect(after.blockingFetchCount).toBe(before.blockingFetchCount);
+    expect(after.fetchCount).toBeGreaterThan(before.fetchCount);
+  });
+
+  it("records a cold-isolate store read as a store hit, not a fetch", async () => {
+    const kp = await generateRotatableKeypair("k1");
+    const store = makeStoreHarness();
+    const warm = makeFetchHarness([kp.publicJwk]);
+
+    const first = createCachedRemoteJWKSet(jwksUrl, { fetch: warm.fetch, store });
+    const token = await sign(kp);
+    await jwtVerify(token, first, { issuer, audience });
+
+    const second = createCachedRemoteJWKSet(jwksUrl, { fetch: failingFetch, store });
+    await jwtVerify(token, second, { issuer, audience });
+
+    const stats = second.inspect();
+    expect(stats.storeHitCount).toBe(1);
+    expect(stats.blockingFetchCount).toBe(0);
+  });
+
+  it("does not block a request on revalidation once keys are cached", async () => {
+    const kp = await generateRotatableKeypair("k1");
+    const harness = makeFetchHarness([kp.publicJwk]);
+    let hang = false;
+    const slowAfterFirst: typeof globalThis.fetch = async (...args) => {
+      if (hang) await new Promise((r) => setTimeout(r, 5_000));
+      return harness.fetch(...args);
+    };
+
+    const jwks = createCachedRemoteJWKSet(jwksUrl, {
+      fetch: slowAfterFirst,
+      store: null,
+      ttlMs: 10,
+      staleMaxMs: 60_000,
+    });
+
+    const token = await sign(kp);
+    await jwtVerify(token, jwks, { issuer, audience });
+
+    hang = true;
+    await new Promise((r) => setTimeout(r, 20));
+
+    // The revalidation behind this call hangs for 5s; the verify must not.
+    const startedAt = Date.now();
+    await jwtVerify(token, jwks, { issuer, audience });
+    expect(Date.now() - startedAt).toBeLessThan(1_000);
+  });
 });