From 37fe5b93b896ed2dbbfa86bb72a9eed79a116ff8 Mon Sep 17 00:00:00 2001 From: can1357 Date: Mon, 25 May 2026 20:11:31 +0200 Subject: [PATCH] fix(ai): fixed OAuth refresh race by coalescing concurrent credential refreshes - Added a per-credential in-flight refresh map in AuthStorage and returned existing promises for concurrent OAuth refreshes. - Updated OAuth API-key retrieval to share refreshes and persist refreshed candidate credentials. - Added tests covering preflight persistence across multiple credentials and ensuring concurrent refreshes coalesced to one call. --- packages/ai/src/auth-storage.ts | 21 +++++ .../auth-storage-oauth-refresh-race.test.ts | 85 +++++++++++++++++++ 2 files changed, 106 insertions(+) diff --git a/packages/ai/src/auth-storage.ts b/packages/ai/src/auth-storage.ts index 91af21e0a..58a477887 100644 --- a/packages/ai/src/auth-storage.ts +++ b/packages/ai/src/auth-storage.ts @@ -642,6 +642,7 @@ export class AuthStorage { #generation = 1; #generationListeners: Set<(generation: number) => void> = new Set(); #oauthRefreshInFlight: Map> = new Map(); + #oauthCredentialRefreshInFlight: Map> = new Map(); #closed = false; constructor(store: AuthCredentialStore, options: AuthStorageOptions = {}) { @@ -2611,7 +2612,27 @@ export class AuthStorage { credentialId: number | undefined, signal?: AbortSignal, ): Promise { + if (credentialId !== undefined) { + const existing = this.#oauthCredentialRefreshInFlight.get(credentialId); + if (existing) return raceCredentialRefreshWithSignal(existing, signal); + } if (Date.now() < credential.expires) return credential; + if (credentialId === undefined) { + return this.#refreshOAuthCredentialUnshared(provider, credential, undefined, signal); + } + const promise = this.#refreshOAuthCredentialUnshared(provider, credential, credentialId).finally(() => { + this.#oauthCredentialRefreshInFlight.delete(credentialId); + }); + this.#oauthCredentialRefreshInFlight.set(credentialId, promise); + return raceCredentialRefreshWithSignal(promise, signal); + } + + async #refreshOAuthCredentialUnshared( + provider: Provider, + credential: OAuthCredential, + credentialId: number | undefined, + signal?: AbortSignal, + ): Promise { let refreshPromise: Promise; // Caller override > store-level hook > local per-provider refresh. // `RemoteAuthCredentialStore` exposes the hook so a broker-backed gateway diff --git a/packages/ai/test/auth-storage-oauth-refresh-race.test.ts b/packages/ai/test/auth-storage-oauth-refresh-race.test.ts index 78b8134bc..25797974f 100644 --- a/packages/ai/test/auth-storage-oauth-refresh-race.test.ts +++ b/packages/ai/test/auth-storage-oauth-refresh-race.test.ts @@ -197,6 +197,91 @@ describe("AuthStorage OAuth refresh race", () => { expect(events[0]?.disabledCause).toContain("invalid_grant"); }); }); + test("persists every credential refreshed during candidate preflight", async () => { + if (!authStorage || !store) throw new Error("test setup failed"); + + const expires = Date.now() - 60_000; + const refreshedExpires = Date.now() + 60 * 60_000; + oauthUtils.registerOAuthProvider({ + id: "unit-oauth-preflight", + name: "Unit OAuth Preflight", + sourceId: "auth-storage-oauth-refresh-race-test", + async login() { + return { access: "unused", refresh: "unused", expires: refreshedExpires }; + }, + async refreshToken(credentials) { + return { + ...credentials, + access: `${credentials.access}-rotated`, + refresh: `${credentials.refresh}-rotated`, + expires: refreshedExpires, + }; + }, + getApiKey(credentials) { + return credentials.access; + }, + }); + + await authStorage.set("unit-oauth-preflight", [ + { type: "oauth", access: "access-a", refresh: "refresh-a", expires }, + { type: "oauth", access: "access-b", refresh: "refresh-b", expires }, + ]); + + const apiKey = await authStorage.getApiKey("unit-oauth-preflight"); + expect(apiKey).toBe("access-a-rotated"); + + const stored = store.listAuthCredentials("unit-oauth-preflight"); + expect(stored).toHaveLength(2); + const oauth = stored.map(entry => entry.credential).filter(credential => credential.type === "oauth"); + expect(oauth.map(credential => credential.refresh).sort()).toEqual(["refresh-a-rotated", "refresh-b-rotated"]); + }); + + test("coalesces concurrent refreshes for the same credential", async () => { + if (!authStorage) throw new Error("test setup failed"); + + const expires = Date.now() - 60_000; + const refreshedExpires = Date.now() + 60 * 60_000; + const refreshStarted = Promise.withResolvers(); + const allowRefresh = Promise.withResolvers(); + let refreshCalls = 0; + + oauthUtils.registerOAuthProvider({ + id: "unit-oauth-mutex", + name: "Unit OAuth Mutex", + sourceId: "auth-storage-oauth-refresh-race-test", + async login() { + return { access: "unused", refresh: "unused", expires: refreshedExpires }; + }, + async refreshToken(credentials) { + refreshCalls += 1; + refreshStarted.resolve(); + await allowRefresh.promise; + return { + ...credentials, + access: "access-rotated", + refresh: "refresh-rotated", + expires: refreshedExpires, + }; + }, + getApiKey(credentials) { + return credentials.access; + }, + }); + + await authStorage.set("unit-oauth-mutex", [ + { type: "oauth", access: "access-old", refresh: "refresh-old", expires }, + ]); + + const first = authStorage.getApiKey("unit-oauth-mutex", "same-session"); + const second = authStorage.getApiKey("unit-oauth-mutex", "same-session"); + + await refreshStarted.promise; + allowRefresh.resolve(); + + await expect(first).resolves.toBe("access-rotated"); + await expect(second).resolves.toBe("access-rotated"); + expect(refreshCalls).toBe(1); + }); test("invalidating a session-sticky OAuth credential rotates the retry to another active credential", async () => { if (!authStorage) throw new Error("test setup failed");