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.
2026-05-25 20:11:31 +02:00
parent 4cd86383c0
commit 37fe5b93b8
2 changed files with 106 additions and 0 deletions
@@ -642,6 +642,7 @@ export class AuthStorage {
 	#generation = 1;
 	#generationListeners: Set<(generation: number) => void> = new Set();
 	#oauthRefreshInFlight: Map<number, Promise<AuthCredentialSnapshotEntry>> = new Map();
+	#oauthCredentialRefreshInFlight: Map<number, Promise<OAuthCredentials>> = new Map();
 	#closed = false;

 	constructor(store: AuthCredentialStore, options: AuthStorageOptions = {}) {
@@ -2611,7 +2612,27 @@ export class AuthStorage {
 		credentialId: number | undefined,
 		signal?: AbortSignal,
 	): Promise<OAuthCredentials> {
+		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<OAuthCredentials> {
 		let refreshPromise: Promise<OAuthCredentials>;
 		// Caller override > store-level hook > local per-provider refresh.
 		// `RemoteAuthCredentialStore` exposes the hook so a broker-backed gateway
@@ -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<void>();
+		const allowRefresh = Promise.withResolvers<void>();
+		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");