From 60784fd8dd558d7347709074abe2d95d856e76e7 Mon Sep 17 00:00:00 2001 From: Vincent Koc Date: Wed, 12 Aug 2026 21:44:05 +0800 Subject: [PATCH] fix(auth): retain OAuth refresh ownership through CAS Punchcard-Session: golden-meadow-cedar-dv --- .../auth-profiles/credential-state.test.ts | 14 + src/agents/auth-profiles/credential-state.ts | 12 + .../auth-profiles/oauth-manager.test.ts | 292 ++++++++++++++---- src/agents/auth-profiles/oauth-manager.ts | 248 ++++++++------- .../auth-profiles/oauth-refresh-failure.ts | 1 + src/agents/auth-profiles/oauth-shared.test.ts | 78 +++++ src/agents/auth-profiles/oauth-shared.ts | 36 ++- src/agents/auth-profiles/oauth.ts | 5 +- ...t-helpers.formatassistanterrortext.test.ts | 5 +- .../provider-runtime-failure.test.ts | 4 +- .../provider-runtime-failure.ts | 7 +- src/agents/model-auth.profiles.test.ts | 1 - .../auth-storage.oauth-refresh.test.ts | 112 +++++++ src/agents/sessions/auth-storage.test.ts | 9 +- src/agents/sessions/auth-storage.ts | 207 ++++++------- ...rovider-usage.auth.normalizes-keys.test.ts | 1 - 16 files changed, 724 insertions(+), 308 deletions(-) create mode 100644 src/agents/sessions/auth-storage.oauth-refresh.test.ts diff --git a/src/agents/auth-profiles/credential-state.test.ts b/src/agents/auth-profiles/credential-state.test.ts index dd5dc5658a3f..1d88b1ef3571 100644 --- a/src/agents/auth-profiles/credential-state.test.ts +++ b/src/agents/auth-profiles/credential-state.test.ts @@ -7,6 +7,7 @@ import { describe, expect, it } from "vitest"; import { DEFAULT_OAUTH_REFRESH_MARGIN_MS, evaluateStoredCredentialEligibility, + hasOAuthTokenMaterialChanged, hasUsableOAuthCredential, resolveTokenExpiryState, } from "./credential-state.js"; @@ -65,6 +66,19 @@ describe("hasUsableOAuthCredential", () => { }); }); +describe("hasOAuthTokenMaterialChanged", () => { + const base = { access: "access", refresh: "refresh", expires: 100 }; + + it.each([ + ["unchanged", base, false], + ["access", { ...base, access: "next-access" }, true], + ["refresh", { ...base, refresh: "next-refresh" }, true], + ["expires", { ...base, expires: 200 }, true], + ])("classifies %s token material", (_name, current, expected) => { + expect(hasOAuthTokenMaterialChanged(base, current)).toBe(expected); + }); +}); + describe("evaluateStoredCredentialEligibility", () => { const now = 1_700_000_000_000; diff --git a/src/agents/auth-profiles/credential-state.ts b/src/agents/auth-profiles/credential-state.ts index a0afa11c054b..069885de5623 100644 --- a/src/agents/auth-profiles/credential-state.ts +++ b/src/agents/auth-profiles/credential-state.ts @@ -73,6 +73,18 @@ export function hasUsableOAuthCredential( ); } +/** Returns true when provider-issued OAuth token material differs. */ +export function hasOAuthTokenMaterialChanged( + previous: Pick, + current: Pick, +): boolean { + return ( + previous.access !== current.access || + previous.refresh !== current.refresh || + previous.expires !== current.expires + ); +} + // SecretRef and literal secret strings are both valid configured credentials; // unresolved refs are classified separately so callers can surface useful copy. function hasConfiguredSecretRef(value: unknown): boolean { diff --git a/src/agents/auth-profiles/oauth-manager.test.ts b/src/agents/auth-profiles/oauth-manager.test.ts index 50e8e85882e4..d1e9ebcde7a4 100644 --- a/src/agents/auth-profiles/oauth-manager.test.ts +++ b/src/agents/auth-profiles/oauth-manager.test.ts @@ -23,7 +23,7 @@ import { ensureAuthProfileStoreWithoutExternalProfiles, saveAuthProfileStore, } from "./store.js"; -import type { AuthProfileStore, OAuthCredential } from "./types.js"; +import type { AuthProfileStore, OAuthCredential, OAuthCredentials } from "./types.js"; function createCredential(overrides: Partial = {}): OAuthCredential { return { @@ -36,6 +36,12 @@ function createCredential(overrides: Partial = {}): OAuthCreden }; } +function prepareRefresh( + refresh: (credential: OAuthCredential, signal: AbortSignal) => Promise, +) { + return async () => refresh; +} + const tempDirs: string[] = []; async function withOAuthTempRoot( @@ -265,7 +271,7 @@ describe("createOAuthManager", () => { const buildApiKey = vi.fn(async (_provider, value: OAuthCredential) => value.access); const manager = createOAuthManager({ buildApiKey, - refreshCredential: vi.fn(async () => null), + prepareRefresh: prepareRefresh(vi.fn(async () => null)), readBootstrapCredential: () => null, isRefreshTokenReusedError: () => false, }); @@ -347,7 +353,7 @@ describe("createOAuthManager", () => { }); const manager = createOAuthManager({ buildApiKey: async (_provider, credential) => credential.access, - refreshCredential, + prepareRefresh: prepareRefresh(refreshCredential), readBootstrapCredential: () => null, isRefreshTokenReusedError: () => false, }); @@ -400,7 +406,7 @@ describe("createOAuthManager", () => { }); const manager = createOAuthManager({ buildApiKey: async (_provider, credential) => credential.access, - refreshCredential, + prepareRefresh: prepareRefresh(refreshCredential), readBootstrapCredential: () => null, isRefreshTokenReusedError: () => false, }); @@ -451,14 +457,16 @@ describe("createOAuthManager", () => { const manager = createOAuthManager({ buildApiKey: async (_provider, credential) => credential.access, - refreshCredential: vi.fn(async (credential) => { - expect(credential.refresh).toBe("external-refresh"); - return { - access: "rotated-access", - refresh: "rotated-refresh", - expires: Date.now() + 60_000, - }; - }), + prepareRefresh: prepareRefresh( + vi.fn(async (credential) => { + expect(credential.refresh).toBe("external-refresh"); + return { + access: "rotated-access", + refresh: "rotated-refresh", + expires: Date.now() + 60_000, + }; + }), + ), readBootstrapCredential: () => createCredential({ provider: "minimax-portal", @@ -509,7 +517,7 @@ describe("createOAuthManager", () => { const refreshCredential = vi.fn(async () => null); const manager = createOAuthManager({ buildApiKey: async (_provider, value) => value.access, - refreshCredential, + prepareRefresh: prepareRefresh(refreshCredential), readBootstrapCredential: () => null, isRefreshTokenReusedError: () => false, }); @@ -552,28 +560,30 @@ describe("createOAuthManager", () => { const manager = createOAuthManager({ buildApiKey: async (_provider, credential) => credential.access, - refreshCredential: vi.fn(async () => { - saveAuthProfileStore( - { - version: 1, - profiles: { - [profileId]: createCredential({ - access: "stale-race-access", - refresh: "consumed-race-refresh", - expires: Date.now() + 10 * 60_000, - accountId: "acct-123", - }), + prepareRefresh: prepareRefresh( + vi.fn(async () => { + saveAuthProfileStore( + { + version: 1, + profiles: { + [profileId]: createCredential({ + access: "stale-race-access", + refresh: "consumed-race-refresh", + expires: Date.now() + 10 * 60_000, + accountId: "acct-123", + }), + }, }, - }, - agentDir, - { filterExternalAuthProfiles: false }, - ); - return { - access: "rotated-access", - refresh: "rotated-refresh", - expires: Date.now() + 60_000, - }; - }), + agentDir, + { filterExternalAuthProfiles: false }, + ); + return { + access: "rotated-access", + refresh: "rotated-refresh", + expires: Date.now() + 60_000, + }; + }), + ), readBootstrapCredential: () => null, isRefreshTokenReusedError: () => false, }); @@ -600,6 +610,86 @@ describe("createOAuthManager", () => { }); }); + it.each([ + { + name: "matching identity", + refreshedIdentity: { accountId: "acct-123" }, + expectedAccess: "rotated-access", + }, + { + name: "omitted identity", + refreshedIdentity: {}, + expectedAccess: "rotated-access", + }, + { + name: "different identity", + refreshedIdentity: { accountId: "acct-456" }, + expectedError: "OAuth credential identity changed during refresh; sign in again", + }, + ])( + "validates $name on an ordinary refresh before persistence", + async ({ name, refreshedIdentity, expectedAccess, expectedError }) => { + await withOAuthTempRoot(`oauth-manager-refresh-identity-${name}-`, async (tempRoot) => { + const agentDir = path.join(tempRoot, "agents", "main", "agent"); + await fs.mkdir(agentDir, { recursive: true }); + const profileId = "openai:oauth"; + const expired = createCredential({ + access: "expired-access", + refresh: "expired-refresh", + expires: 1, + accountId: "acct-123", + }); + saveAuthProfileStore({ version: 1, profiles: { [profileId]: expired } }, agentDir, { + filterExternalAuthProfiles: false, + }); + const manager = createOAuthManager({ + buildApiKey: async (_provider, credential) => credential.access, + prepareRefresh: prepareRefresh(async () => ({ + access: "rotated-access", + refresh: "rotated-refresh", + expires: Date.now() + 60_000, + ...refreshedIdentity, + })), + readBootstrapCredential: () => null, + isRefreshTokenReusedError: () => false, + }); + const input = { + store: ensureAuthProfileStoreWithoutExternalProfiles(agentDir), + profileId, + credential: expired, + agentDir, + }; + + if (expectedError) { + try { + await manager.resolveOAuthAccess(input); + throw new Error("expected refresh failure"); + } catch (error) { + expect(error).toBeInstanceOf(OAuthManagerRefreshError); + expect((error as OAuthManagerRefreshError).cause).toMatchObject({ + message: expectedError, + }); + } + expect( + ensureAuthProfileStoreWithoutExternalProfiles(agentDir).profiles[profileId], + ).toEqual(expired); + return; + } + + await expect(manager.resolveOAuthAccess(input)).resolves.toMatchObject({ + apiKey: expectedAccess, + }); + expect( + ensureAuthProfileStoreWithoutExternalProfiles(agentDir).profiles[profileId], + ).toMatchObject({ + access: "rotated-access", + refresh: "rotated-refresh", + accountId: "acct-123", + }); + }); + }, + ); + it("uses a different-identity stored credential after a CAS race", async () => { await withOAuthTempRoot("oauth-manager-cas-different-identity-", async (tempRoot) => { const mainAgentDir = path.join(tempRoot, "agents", "main", "agent"); @@ -632,23 +722,25 @@ describe("createOAuthManager", () => { const manager = createOAuthManager({ buildApiKey: async (_provider, credential) => credential.access, - refreshCredential: vi.fn(async () => { - saveAuthProfileStore( - { - version: 1, - profiles: { - [profileId]: relogged, + prepareRefresh: prepareRefresh( + vi.fn(async () => { + saveAuthProfileStore( + { + version: 1, + profiles: { + [profileId]: relogged, + }, }, - }, - agentDir, - { filterExternalAuthProfiles: false }, - ); - return { - access: "rotated-access", - refresh: "rotated-refresh", - expires: Date.now() + 60_000, - }; - }), + agentDir, + { filterExternalAuthProfiles: false }, + ); + return { + access: "rotated-access", + refresh: "rotated-refresh", + expires: Date.now() + 60_000, + }; + }), + ), readBootstrapCredential: () => null, isRefreshTokenReusedError: () => false, }); @@ -697,9 +789,11 @@ describe("createOAuthManager", () => { ); const manager = createOAuthManager({ buildApiKey: async (_provider, credential) => credential.access, - refreshCredential: vi.fn(async () => { - throw new Error("refresh rejected managed profile"); - }), + prepareRefresh: prepareRefresh( + vi.fn(async () => { + throw new Error("refresh rejected managed profile"); + }), + ), readBootstrapCredential: () => null, isRefreshTokenReusedError: () => false, }); @@ -748,11 +842,13 @@ describe("createOAuthManager", () => { const manager = createOAuthManager({ buildApiKey: async (_provider, credential) => credential.access, - refreshCredential: vi.fn(async () => { - throw new Error( - "refresh rejected external-attempt-access external-attempt-refresh external-attempt-id-token", - ); - }), + prepareRefresh: prepareRefresh( + vi.fn(async () => { + throw new Error( + "refresh rejected external-attempt-access external-attempt-refresh external-attempt-id-token", + ); + }), + ), readBootstrapCredential: () => externalCredential, isRefreshTokenReusedError: () => false, }); @@ -781,4 +877,84 @@ describe("createOAuthManager", () => { } }); }); + + it("persists late refresh success and reuses it after the caller deadline", async () => { + await withOAuthAgentDirs("oauth-manager-late-success-", async ({ agentDir }) => { + const profileId = "openai:oauth"; + const credential = createCredential({ expires: 1 }); + saveAuthProfileStore({ version: 1, profiles: { [profileId]: credential } }, agentDir, { + filterExternalAuthProfiles: false, + }); + const stalled = Promise.withResolvers(); + const refreshCredential = vi.fn(async () => await stalled.promise); + const manager = createOAuthManager({ + buildApiKey: async (_provider, value) => value.access, + prepareRefresh: prepareRefresh(refreshCredential), + readBootstrapCredential: () => null, + isRefreshTokenReusedError: () => false, + refreshTimeoutMs: 10, + }); + const input = () => { + const store = ensureAuthProfileStoreWithoutExternalProfiles(agentDir); + return { + store, + profileId, + credential: store.profiles[profileId] as OAuthCredential, + agentDir, + }; + }; + + await expect(manager.resolveOAuthAccess(input())).rejects.toThrow("exceeded caller deadline"); + stalled.resolve({ + access: "late-access", + refresh: "late-refresh", + expires: Date.now() + 10 * 60_000, + }); + await vi.waitFor(() => { + expect( + ensureAuthProfileStoreWithoutExternalProfiles(agentDir).profiles[profileId], + ).toMatchObject({ access: "late-access", refresh: "late-refresh" }); + }); + await expect(manager.resolveOAuthAccess(input())).resolves.toMatchObject({ + apiKey: "late-access", + }); + expect(refreshCredential).toHaveBeenCalledOnce(); + }); + }); + + it("times out queued callers independently without invoking an expired follower", async () => { + await withOAuthAgentDirs("oauth-manager-follower-timeout-", async ({ agentDir }) => { + const profileId = "openai:oauth"; + const credential = createCredential({ expires: 1 }); + saveAuthProfileStore({ version: 1, profiles: { [profileId]: credential } }, agentDir, { + filterExternalAuthProfiles: false, + }); + const stalled = Promise.withResolvers(); + const refreshCredential = vi.fn(async () => await stalled.promise); + const prepareRefreshCall = vi.fn(async () => refreshCredential); + const manager = createOAuthManager({ + buildApiKey: async (_provider, value) => value.access, + prepareRefresh: prepareRefreshCall, + readBootstrapCredential: () => null, + isRefreshTokenReusedError: () => false, + refreshTimeoutMs: 10, + }); + const input = () => ({ + store: ensureAuthProfileStoreWithoutExternalProfiles(agentDir), + profileId, + credential, + agentDir, + }); + + const first = manager.resolveOAuthAccess(input()); + const firstAssertion = expect(first).rejects.toThrow("exceeded caller deadline"); + await vi.waitFor(() => expect(refreshCredential).toHaveBeenCalledOnce()); + const follower = manager.resolveOAuthAccess(input()); + const followerAssertion = expect(follower).rejects.toThrow("exceeded caller deadline"); + await Promise.all([firstAssertion, followerAssertion]); + stalled.reject(new Error("late provider failure")); + await vi.waitFor(() => expect(prepareRefreshCall).toHaveBeenCalledTimes(2)); + expect(refreshCredential).toHaveBeenCalledOnce(); + }); + }); }); diff --git a/src/agents/auth-profiles/oauth-manager.ts b/src/agents/auth-profiles/oauth-manager.ts index 559d7dd50ab5..914713315dc0 100644 --- a/src/agents/auth-profiles/oauth-manager.ts +++ b/src/agents/auth-profiles/oauth-manager.ts @@ -10,23 +10,27 @@ import { formatErrorMessage } from "../../infra/errors.js"; import { withFileLock } from "../../infra/file-lock.js"; import { redactSensitiveText } from "../../logging/redact.js"; import { KeyedAsyncQueue } from "../../plugin-sdk/keyed-async-queue.js"; +import { createDeferredCore } from "../../shared/deferred.js"; import { OAUTH_REFRESH_CALL_TIMEOUT_MS, OAUTH_REFRESH_LOCK_OPTIONS, authProfilesLog, } from "./constants.js"; -import { hasUsableOAuthCredential } from "./credential-state.js"; +import { hasOAuthTokenMaterialChanged, hasUsableOAuthCredential } from "./credential-state.js"; import { shouldMirrorRefreshedOAuthCredential } from "./oauth-identity.js"; -import { OAuthRefreshFailureError } from "./oauth-refresh-failure.js"; +import { + OAUTH_REFRESH_CALLER_DEADLINE_MESSAGE, + OAuthRefreshFailureError, +} from "./oauth-refresh-failure.js"; import { buildRefreshContentionError, isGlobalRefreshLockTimeoutError, } from "./oauth-refresh-lock-errors.js"; import { areOAuthCredentialsEquivalent, - hasMatchingOAuthIdentity, isSafeToAdoptBootstrapOAuthIdentity, isSafeToAdoptMainStoreOAuthIdentity, + resolveOAuthRefreshConflict, shouldBootstrapFromExternalCliCredential, shouldReplaceStoredOAuthCredential, } from "./oauth-shared.js"; @@ -46,23 +50,66 @@ type OAuthManagerAdapter = { credentials: OAuthCredential, context: { cfg?: OpenClawConfig; agentDir?: string }, ) => Promise; - refreshCredential: ( + prepareRefresh: ( credential: OAuthCredential, - context: { cfg?: OpenClawConfig; agentDir?: string }, - ) => Promise; + context: { cfg?: OpenClawConfig; agentDir?: string; signal: AbortSignal }, + ) => Promise; readBootstrapCredential: (params: { store: AuthProfileStore; profileId: string; credential: OAuthCredential; }) => OAuthCredential | null; isRefreshTokenReusedError: (error: unknown) => boolean; + refreshTimeoutMs?: number; }; +export type PreparedOAuthRefresh = ( + credential: OAuthCredential, + signal: AbortSignal, +) => Promise; + type ResolvedOAuthAccess = { apiKey: string; credential: OAuthCredential; }; +/** Bound one caller while retaining refresh ownership until the operation settles. */ +export function runRetainedOAuthRefreshOperation(params: { + timeoutMs: number; + run: (signal: AbortSignal) => Promise; +}): Promise { + const caller = createDeferredCore(); + const controller = new AbortController(); + let timedOut = false; + const timeoutHandle = setTimeout(() => { + timedOut = true; + const error = new Error(`${OAUTH_REFRESH_CALLER_DEADLINE_MESSAGE} (${params.timeoutMs}ms)`); + controller.abort(error); + caller.reject(error); + }, params.timeoutMs); + let owner: Promise; + try { + owner = params.run(controller.signal); + } catch (error: unknown) { + owner = Promise.reject(error instanceof Error ? error : new Error(String(error))); + } + void owner.then( + (result) => { + clearTimeout(timeoutHandle); + if (!timedOut) { + caller.resolve(result); + } + }, + (error: unknown) => { + clearTimeout(timeoutHandle); + if (!timedOut) { + caller.reject(error instanceof Error ? error : new Error(String(error))); + } + }, + ); + return caller.promise; +} + /** Refresh failure that preserves a redacted refreshed store and credential. */ export class OAuthManagerRefreshError extends OAuthRefreshFailureError { override readonly profileId: string; @@ -139,23 +186,12 @@ export class OAuthManagerRefreshError extends OAuthRefreshFailureError { } } -function hasOAuthCredentialChanged( - previous: Pick, - current: Pick, -): boolean { - return ( - previous.access !== current.access || - previous.refresh !== current.refresh || - previous.expires !== current.expires - ); -} - function canReuseOAuthCredentialAfterRefreshFailure(params: { forceRefresh?: boolean; attempted: Pick; candidate: OAuthCredential; }): boolean { - return !params.forceRefresh || hasOAuthCredentialChanged(params.attempted, params.candidate); + return !params.forceRefresh || hasOAuthTokenMaterialChanged(params.attempted, params.candidate); } function collectOAuthCredentialSecrets( @@ -254,7 +290,7 @@ async function loadFreshStoredOAuthCredential(params: { if ( params.requireChange && params.previous && - !hasOAuthCredentialChanged(params.previous, reloaded) + !hasOAuthTokenMaterialChanged(params.previous, reloaded) ) { return null; } @@ -362,26 +398,6 @@ export function createOAuthManager(adapter: OAuthManagerAdapter) { return `${provider}\u0000${profileId}`; } - async function withRefreshCallTimeout( - label: string, - timeoutMs: number, - fn: () => Promise, - ): Promise { - let timeoutHandle: NodeJS.Timeout | undefined; - try { - return await new Promise((resolve, reject) => { - timeoutHandle = setTimeout(() => { - reject(new Error(`OAuth refresh call "${label}" exceeded hard timeout (${timeoutMs}ms)`)); - }, timeoutMs); - fn().then(resolve, reject); - }); - } finally { - if (timeoutHandle) { - clearTimeout(timeoutHandle); - } - } - } - async function mirrorRefreshedCredentialIntoMainStore(params: { profileId: string; refreshed: OAuthCredential; @@ -427,17 +443,37 @@ export function createOAuthManager(adapter: OAuthManagerAdapter) { async function saveOAuthCredentialWithStoreLock(params: { agentDir?: string; profileId: string; - expected: OAuthCredential | OAuthCredential[]; + expected?: OAuthCredential | OAuthCredential[]; + attempted?: OAuthCredential; credential: OAuthCredential; - }): Promise { - let saved = false; + }): Promise { + const input = params.attempted ?? params.credential; + resolveOAuthRefreshConflict({ + authoritative: input, + attempted: input, + refreshed: params.credential, + }); + let selected: OAuthCredential | null = null; const result = await updateAuthProfileStoreWithLock({ agentDir: params.agentDir, updater: (store) => { const existing = store.profiles[params.profileId]; - const expectedCredentials = Array.isArray(params.expected) - ? params.expected - : [params.expected]; + if (params.attempted) { + const decision = resolveOAuthRefreshConflict({ + authoritative: existing, + attempted: params.attempted, + refreshed: params.credential, + }); + selected = decision?.credential ?? null; + if (!decision?.persist) { + return false; + } + // A refresh token may rotate before persistence. Same-identity CAS + // losers must persist the rotation or the token family is bricked. + store.profiles[params.profileId] = { ...decision.credential }; + return true; + } + const expectedCredentials = params.expected ? [params.expected].flat() : []; if ( existing?.type !== "oauth" || !expectedCredentials.some((expected) => areOAuthCredentialsEquivalent(existing, expected)) @@ -457,41 +493,11 @@ export function createOAuthManager(adapter: OAuthManagerAdapter) { return false; } store.profiles[params.profileId] = { ...params.credential }; - saved = true; + selected = params.credential; return true; }, }); - return result !== null && saved; - } - - async function resolveOAuthCredentialAfterPersistMiss(params: { - agentDir?: string; - profileId: string; - refreshed: OAuthCredential; - }): Promise { - // Single locked pass decides both outcomes so no relog can slip between a - // pre-read and the update: same identity persists the rotation, different - // identity adopts the stored (re-logged) credential for this call. - let adopted: OAuthCredential | null = null; - const result = await updateAuthProfileStoreWithLock({ - agentDir: params.agentDir, - updater: (store) => { - const existing = store.profiles[params.profileId]; - if (existing?.type !== "oauth" || existing.provider !== params.refreshed.provider) { - return false; - } - // Refresh tokens rotate server-side before persist. Same-identity CAS - // losers must win the store or the token family is bricked. - if (hasMatchingOAuthIdentity(existing, params.refreshed)) { - store.profiles[params.profileId] = { ...params.refreshed }; - adopted = params.refreshed; - return true; - } - adopted = hasUsableOAuthCredential(existing) ? existing : null; - return false; - }, - }); - return result === null ? null : adopted; + return result === null ? null : selected; } async function doRefreshOAuthTokenWithLock(params: { @@ -501,6 +507,8 @@ export function createOAuthManager(adapter: OAuthManagerAdapter) { cfg?: OpenClawConfig; forceRefresh?: boolean; attemptedCredentials?: OAuthCredential[]; + refreshCredential: PreparedOAuthRefresh; + signal: AbortSignal; }): Promise { const ownerAgentDir = resolvePersistedAuthProfileOwnerAgentDir(params); const authPath = resolveAuthProfileDatabasePath(ownerAgentDir); @@ -508,6 +516,7 @@ export function createOAuthManager(adapter: OAuthManagerAdapter) { try { return await withFileLock(globalRefreshLockPath, OAUTH_REFRESH_LOCK_OPTIONS, async () => { + params.signal.throwIfAborted(); const store = loadStoredOAuthRefreshStore(ownerAgentDir); const cred = store.profiles[params.profileId]; if (!cred || cred.type !== "oauth") { @@ -622,24 +631,15 @@ export function createOAuthManager(adapter: OAuthManagerAdapter) { if (normalizeSecretInputString(credentialToRefresh.refresh) === undefined) { return null; } - const refreshedCredentials = await withRefreshCallTimeout( - `refreshOAuthCredential(${cred.provider})`, - OAUTH_REFRESH_CALL_TIMEOUT_MS, - async () => { - params.attemptedCredentials?.push(credentialToRefresh); - const refreshed = await adapter.refreshCredential(credentialToRefresh, { - cfg: params.cfg, - agentDir: params.agentDir, - }); - return refreshed - ? ({ - ...credentialToRefresh, - ...refreshed, - type: "oauth", - } satisfies OAuthCredential) - : null; - }, - ); + params.attemptedCredentials?.push(credentialToRefresh); + const refreshed = await params.refreshCredential(credentialToRefresh, params.signal); + const refreshedCredentials = refreshed + ? ({ + ...credentialToRefresh, + ...refreshed, + type: "oauth", + } satisfies OAuthCredential) + : null; if (!refreshedCredentials) { return null; } @@ -647,30 +647,20 @@ export function createOAuthManager(adapter: OAuthManagerAdapter) { const persisted = await saveOAuthCredentialWithStoreLock({ agentDir: ownerAgentDir, profileId: params.profileId, - expected: - credentialToRefresh === cred || areOAuthCredentialsEquivalent(credentialToRefresh, cred) - ? credentialToRefresh - : [credentialToRefresh, cred], + attempted: credentialToRefresh, credential: refreshedCredentials, }); if (!persisted) { - const recovered = await resolveOAuthCredentialAfterPersistMiss({ - agentDir: ownerAgentDir, - profileId: params.profileId, - refreshed: refreshedCredentials, - }); - if (!recovered) { - throw new Error("Failed to persist refreshed OAuth credential"); - } - if (recovered !== refreshedCredentials) { - return { - apiKey: await adapter.buildApiKey(recovered.provider, recovered, { - cfg: params.cfg, - agentDir: params.agentDir, - }), - credential: recovered, - }; - } + throw new Error("Failed to persist refreshed OAuth credential"); + } + if (persisted !== refreshedCredentials) { + return { + apiKey: await adapter.buildApiKey(persisted.provider, persisted, { + cfg: params.cfg, + agentDir: params.agentDir, + }), + credential: persisted, + }; } if (ownerAgentDir) { const mainPath = resolveAuthProfileDatabasePath(undefined); @@ -702,6 +692,7 @@ export function createOAuthManager(adapter: OAuthManagerAdapter) { } async function refreshOAuthTokenWithLock(params: { + credential: OAuthCredential; profileId: string; provider: string; agentDir?: string; @@ -710,7 +701,26 @@ export function createOAuthManager(adapter: OAuthManagerAdapter) { attemptedCredentials?: OAuthCredential[]; }): Promise { const key = refreshQueueKey(params.provider, params.profileId); - return await refreshQueue.enqueue(key, () => doRefreshOAuthTokenWithLock(params)); + return runRetainedOAuthRefreshOperation({ + timeoutMs: adapter.refreshTimeoutMs ?? OAUTH_REFRESH_CALL_TIMEOUT_MS, + run: async (signal) => { + signal.throwIfAborted(); + const refreshCredential = await adapter.prepareRefresh(params.credential, { + cfg: params.cfg, + agentDir: params.agentDir, + signal, + }); + signal.throwIfAborted(); + return await refreshQueue.enqueue(key, async () => { + signal.throwIfAborted(); + return await doRefreshOAuthTokenWithLock({ + ...params, + refreshCredential, + signal, + }); + }); + }, + }); } async function resolveOAuthAccess(params: { @@ -748,6 +758,7 @@ export function createOAuthManager(adapter: OAuthManagerAdapter) { try { const refreshed = await refreshOAuthTokenWithLock({ + credential: effectiveCredential, profileId: params.profileId, provider: params.credential.provider, agentDir: params.agentDir, @@ -780,7 +791,7 @@ export function createOAuthManager(adapter: OAuthManagerAdapter) { adapter.isRefreshTokenReusedError(error) && refreshed?.type === "oauth" && refreshed.provider === params.credential.provider && - hasOAuthCredentialChanged(params.credential, refreshed) + hasOAuthTokenMaterialChanged(params.credential, refreshed) ) { const recovered = await loadFreshStoredOAuthCredential({ profileId: params.profileId, @@ -800,6 +811,7 @@ export function createOAuthManager(adapter: OAuthManagerAdapter) { } try { const retried = await refreshOAuthTokenWithLock({ + credential: effectiveCredential, profileId: params.profileId, provider: params.credential.provider, agentDir: params.agentDir, diff --git a/src/agents/auth-profiles/oauth-refresh-failure.ts b/src/agents/auth-profiles/oauth-refresh-failure.ts index b63602814a39..80c76897bdfe 100644 --- a/src/agents/auth-profiles/oauth-refresh-failure.ts +++ b/src/agents/auth-profiles/oauth-refresh-failure.ts @@ -9,6 +9,7 @@ import { formatCliCommand } from "../../cli/command-format.js"; import { formatInlineCodeSpan } from "../../shared/markdown-code.js"; import type { AuthProfileFailureReason } from "./types.js"; +export const OAUTH_REFRESH_CALLER_DEADLINE_MESSAGE = "OAuth refresh call exceeded caller deadline"; export type OAuthRefreshFailureReason = | "refresh_token_reused" | "invalid_grant" diff --git a/src/agents/auth-profiles/oauth-shared.test.ts b/src/agents/auth-profiles/oauth-shared.test.ts index 4e9a1bd5bc6f..699385288725 100644 --- a/src/agents/auth-profiles/oauth-shared.test.ts +++ b/src/agents/auth-profiles/oauth-shared.test.ts @@ -9,6 +9,7 @@ import { MAX_DATE_TIMESTAMP_MS } from "@openclaw/normalization-core/number-coerc import { describe, expect, it, vi } from "vitest"; import { overlayRuntimeExternalOAuthProfiles, + resolveOAuthRefreshConflict, shouldReplaceStoredOAuthCredential, } from "./oauth-shared.js"; import type { AuthProfileStore, OAuthCredential } from "./types.js"; @@ -182,3 +183,80 @@ describe("overlayRuntimeExternalOAuthProfiles", () => { expect(shouldReplaceStoredOAuthCredential(existing, incoming)).toBe(true); }); }); + +describe("resolveOAuthRefreshConflict", () => { + const attempted: OAuthCredential = { + type: "oauth", + provider: "openai", + access: "attempted-access", + refresh: "attempted-refresh", + expires: 1, + accountId: "acct-1", + email: "user@example.com", + }; + + it.each([ + { + name: "matching account", + refreshed: { accountId: "acct-1", email: "other@example.com" }, + }, + { + name: "matching email", + refreshed: { email: "USER@example.com" }, + }, + ])("accepts a refreshed $name identity", ({ refreshed }) => { + expect( + resolveOAuthRefreshConflict({ + authoritative: attempted, + attempted, + refreshed: { ...attempted, ...refreshed, access: "refreshed-access" }, + }), + ).toMatchObject({ credential: { access: "refreshed-access" }, persist: true }); + }); + + it("accepts identity learned while refreshing an identity-less credential", () => { + const identityLess = { ...attempted, accountId: undefined, email: undefined }; + expect( + resolveOAuthRefreshConflict({ + authoritative: identityLess, + attempted: identityLess, + refreshed: { ...identityLess, accountId: "acct-1" }, + }), + ).toMatchObject({ credential: { accountId: "acct-1" }, persist: true }); + }); + + it.each([ + { + name: "provider", + refreshed: { provider: "anthropic" }, + message: "OAuth credential identity changed during refresh; sign in again", + }, + { + name: "account", + refreshed: { accountId: "acct-2" }, + message: "OAuth credential identity changed during refresh; sign in again", + }, + { + name: "email", + attempted: { ...attempted, accountId: undefined }, + refreshed: { accountId: undefined, email: "other@example.com" }, + message: "OAuth credential identity changed during refresh; sign in again", + }, + { + name: "missing identity", + refreshed: { accountId: undefined, email: undefined }, + message: "OAuth credential identity changed during refresh; sign in again", + }, + ])( + "rejects a refreshed $name mismatch", + ({ attempted: input = attempted, refreshed, message }) => { + expect(() => + resolveOAuthRefreshConflict({ + authoritative: input, + attempted: input, + refreshed: { ...input, ...refreshed, access: "refreshed-access" }, + }), + ).toThrow(message); + }, + ); +}); diff --git a/src/agents/auth-profiles/oauth-shared.ts b/src/agents/auth-profiles/oauth-shared.ts index feef98fa8c28..9df79e26f38f 100644 --- a/src/agents/auth-profiles/oauth-shared.ts +++ b/src/agents/auth-profiles/oauth-shared.ts @@ -11,7 +11,7 @@ import { normalizeAuthEmailToken, normalizeAuthIdentityToken, } from "./oauth-identity.js"; -import type { AuthProfileStore, OAuthCredential } from "./types.js"; +import type { AuthProfileCredential, AuthProfileStore, OAuthCredential } from "./types.js"; export { normalizeAuthEmailToken, normalizeAuthIdentityToken } from "./oauth-identity.js"; @@ -91,6 +91,40 @@ export function hasMatchingOAuthIdentity( return hasOAuthIdentity(existing) && isSafeToCopyOAuthIdentity(existing, incoming); } +/** + * Resolve a refresh result against the credential authoritative at commit time. + * Same-identity rotations persist; only a usable different login supersedes them. + */ +export function resolveOAuthRefreshConflict(params: { + authoritative: AuthProfileCredential | undefined; + attempted: OAuthCredential; + refreshed: OAuthCredential; + now?: number; +}): { credential: OAuthCredential; persist: boolean } | null { + const { authoritative, attempted, refreshed } = params; + if ( + refreshed.provider !== attempted.provider || + !isSafeToCopyOAuthIdentity(attempted, refreshed) + ) { + throw new Error("OAuth credential identity changed during refresh; sign in again"); + } + if (authoritative?.type !== "oauth") { + return null; + } + if (authoritative.provider !== attempted.provider) { + return null; + } + if ( + areOAuthCredentialsEquivalent(authoritative, attempted) || + hasMatchingOAuthIdentity(authoritative, refreshed) + ) { + return { credential: refreshed, persist: true }; + } + return hasUsableOAuthCredential(authoritative, { now: params.now }) + ? { credential: authoritative, persist: false } + : null; +} + // Different adoption paths have different safety thresholds. Bootstrap can // adopt missing identities, while stored overwrite requires an identity match. type OAuthIdentitySafetyPolicy = { diff --git a/src/agents/auth-profiles/oauth.ts b/src/agents/auth-profiles/oauth.ts index 4e1a25039119..5dc33178caee 100644 --- a/src/agents/auth-profiles/oauth.ts +++ b/src/agents/auth-profiles/oauth.ts @@ -230,7 +230,10 @@ export async function refreshOAuthCredentialForRuntime(params: { const oauthManager = createOAuthManager({ buildApiKey: buildOAuthApiKey, - refreshCredential: refreshOAuthCredential, + prepareRefresh: async (_credential, context) => async (credential, signal) => { + signal.throwIfAborted(); + return await refreshOAuthCredential(credential, { cfg: context.cfg }); + }, readBootstrapCredential: ({ store, profileId, credential }) => readExternalCliBootstrapCredential({ store, diff --git a/src/agents/embedded-agent-helpers.formatassistanterrortext.test.ts b/src/agents/embedded-agent-helpers.formatassistanterrortext.test.ts index 5a2330f67e50..c70b9645b8ab 100644 --- a/src/agents/embedded-agent-helpers.formatassistanterrortext.test.ts +++ b/src/agents/embedded-agent-helpers.formatassistanterrortext.test.ts @@ -114,9 +114,8 @@ describe("formatAssistantErrorText", () => { expected: "Authentication refresh failed. Re-authenticate this provider and try again.", }, { - title: "returns a timeout-specific message for OAuth refresh hard timeouts", - errorText: - 'OAuth refresh call "refreshProviderOAuthCredentialWithPlugin(openai)" exceeded hard timeout (120000ms)', + title: "returns a timeout-specific message for OAuth refresh caller deadlines", + errorText: "OAuth refresh call exceeded caller deadline (120000ms)", expected: "Authentication refresh timed out before the provider completed. Retry in a moment; re-authenticate only if it keeps failing.", }, diff --git a/src/agents/embedded-agent-helpers/provider-runtime-failure.test.ts b/src/agents/embedded-agent-helpers/provider-runtime-failure.test.ts index fbc519819d3a..b5a2d11fc910 100644 --- a/src/agents/embedded-agent-helpers/provider-runtime-failure.test.ts +++ b/src/agents/embedded-agent-helpers/provider-runtime-failure.test.ts @@ -112,9 +112,7 @@ describe("classifyProviderRuntimeFailureKind", () => { it("classifies OAuth refresh timeouts and lock contention distinctly", () => { expect( - classifyProviderRuntimeFailureKind( - 'OAuth refresh call "refreshProviderOAuthCredentialWithPlugin(openai)" exceeded hard timeout (120000ms)', - ), + classifyProviderRuntimeFailureKind("OAuth refresh call exceeded caller deadline (120000ms)"), ).toBe("refresh_timeout"); expect( classifyProviderRuntimeFailureKind("file lock timeout for /tmp/openclaw-oauth-refresh.lock"), diff --git a/src/agents/embedded-agent-helpers/provider-runtime-failure.ts b/src/agents/embedded-agent-helpers/provider-runtime-failure.ts index 94327e996277..c25fe48c4ebc 100644 --- a/src/agents/embedded-agent-helpers/provider-runtime-failure.ts +++ b/src/agents/embedded-agent-helpers/provider-runtime-failure.ts @@ -1,6 +1,9 @@ import { normalizeLowercaseStringOrEmpty } from "@openclaw/normalization-core/string-coerce"; import { extractLeadingHttpStatus } from "../../shared/assistant-error-format.js"; -import { classifyOAuthRefreshFailure } from "../auth-profiles/oauth-refresh-failure.js"; +import { + classifyOAuthRefreshFailure, + OAUTH_REFRESH_CALLER_DEADLINE_MESSAGE, +} from "../auth-profiles/oauth-refresh-failure.js"; import { formatExecDeniedUserMessage } from "../exec-approval-result.js"; import { inferSignalStatus, @@ -151,7 +154,7 @@ function isTimeoutTransportErrorMessage(raw: string, status?: number): boolean { return false; } function isOAuthRefreshTimeoutMessage(raw: string): boolean { - return /\boauth refresh call\b.*\bexceeded hard timeout\b/i.test(raw); + return raw.toLowerCase().includes(OAUTH_REFRESH_CALLER_DEADLINE_MESSAGE.toLowerCase()); } function isOAuthRefreshContentionMessage(raw: string): boolean { return ( diff --git a/src/agents/model-auth.profiles.test.ts b/src/agents/model-auth.profiles.test.ts index 68219fef6d1a..004b43b4968d 100644 --- a/src/agents/model-auth.profiles.test.ts +++ b/src/agents/model-auth.profiles.test.ts @@ -215,7 +215,6 @@ vi.mock("../plugins/provider-runtime.js", () => ({ return undefined; }, formatProviderAuthProfileApiKeyWithPlugin: async () => undefined, - refreshProviderOAuthCredentialWithPlugin: async () => null, resolveProviderSyntheticAuthWithPlugin: (params: { provider: string; context: { providerConfig?: { api?: string; baseUrl?: string; models?: unknown[] } }; diff --git a/src/agents/sessions/auth-storage.oauth-refresh.test.ts b/src/agents/sessions/auth-storage.oauth-refresh.test.ts new file mode 100644 index 000000000000..4d1193975f96 --- /dev/null +++ b/src/agents/sessions/auth-storage.oauth-refresh.test.ts @@ -0,0 +1,112 @@ +import { describe, expect, it, vi } from "vitest"; + +vi.mock("../auth-profiles/constants.js", async () => { + const actual = await vi.importActual( + "../auth-profiles/constants.js", + ); + return { ...actual, OAUTH_REFRESH_CALL_TIMEOUT_MS: 10 }; +}); + +import type { OAuthCredentials } from "../../llm/utils/oauth/types.js"; +import { getAuthStorageOAuthProviderRegistry } from "./auth-storage-oauth-registry.js"; +import { AuthStorage } from "./auth-storage.js"; + +function createStorage( + refreshToken: () => Promise, + credential: Partial = {}, +) { + const storage = AuthStorage.inMemory({ + "test-oauth": { + type: "oauth", + access: "expired-access", + refresh: "expired-refresh", + expires: 1, + ...credential, + }, + }); + getAuthStorageOAuthProviderRegistry(storage).register({ + id: "test-oauth", + name: "Test OAuth", + async login() { + throw new Error("not used"); + }, + refreshToken, + getApiKey(credentials) { + return credentials.access; + }, + }); + return storage; +} + +describe("AuthStorage OAuth refresh ownership", () => { + it("persists late success and reuses it after the caller deadline", async () => { + const stalled = Promise.withResolvers<{ + access: string; + refresh: string; + expires: number; + }>(); + const refreshToken = vi.fn(async () => await stalled.promise); + const storage = createStorage(refreshToken); + + await expect(storage.getApiKey("test-oauth")).resolves.toBeUndefined(); + expect(storage.drainErrors()[0]?.message).toContain("exceeded caller deadline"); + stalled.resolve({ + access: "late-access", + refresh: "late-refresh", + expires: Date.now() + 10 * 60_000, + }); + await vi.waitFor(() => { + expect(storage.get("test-oauth")).toMatchObject({ + access: "late-access", + refresh: "late-refresh", + }); + }); + + await expect(storage.getApiKey("test-oauth")).resolves.toBe("late-access"); + expect(refreshToken).toHaveBeenCalledOnce(); + }); + + it("records identity mismatch without changing storage or falling back", async () => { + const refreshToken = vi.fn(async () => ({ + access: "rotated-access", + refresh: "rotated-refresh", + expires: Date.now() + 10 * 60_000, + accountId: "acct-2", + })); + const fallback = vi.fn(() => "fallback-key"); + const storage = createStorage(refreshToken, { accountId: "acct-1" }); + storage.setFallbackResolver(fallback); + + await expect(storage.getApiKey("test-oauth")).resolves.toBeUndefined(); + + expect(storage.get("test-oauth")).toMatchObject({ + access: "expired-access", + refresh: "expired-refresh", + accountId: "acct-1", + }); + expect(fallback).not.toHaveBeenCalled(); + expect(storage.drainErrors()).toEqual([ + expect.objectContaining({ + message: "OAuth credential identity changed during refresh; sign in again", + }), + ]); + }); + + it("times out a queued follower without a second provider invocation", async () => { + const stalled = Promise.withResolvers<{ + access: string; + refresh: string; + expires: number; + }>(); + const refreshToken = vi.fn(async () => await stalled.promise); + const storage = createStorage(refreshToken); + + const first = storage.getApiKey("test-oauth"); + await vi.waitFor(() => expect(refreshToken).toHaveBeenCalledOnce()); + const follower = storage.getApiKey("test-oauth"); + await expect(Promise.all([first, follower])).resolves.toEqual([undefined, undefined]); + stalled.reject(new Error("late provider failure")); + expect(storage.drainErrors()).toHaveLength(2); + expect(refreshToken).toHaveBeenCalledOnce(); + }); +}); diff --git a/src/agents/sessions/auth-storage.test.ts b/src/agents/sessions/auth-storage.test.ts index 238de9172130..f3e6cd02ecea 100644 --- a/src/agents/sessions/auth-storage.test.ts +++ b/src/agents/sessions/auth-storage.test.ts @@ -638,7 +638,7 @@ describe("SQLite auth storage", () => { it("throws without changing memory when the durable write fails", () => { const writeError = new Error("simulated durable write failure"); - let persisted = "{}"; + const persisted = "{}"; const backend: AuthStorageBackend = { withLock: (fn) => { const update = fn(persisted); @@ -647,13 +647,6 @@ describe("SQLite auth storage", () => { } return update.result; }, - withLockAsync: async (fn) => { - const update = await fn(persisted); - if (update.next !== undefined) { - persisted = update.next; - } - return update.result; - }, }; const storage = AuthStorage.fromStorage(backend); diff --git a/src/agents/sessions/auth-storage.ts b/src/agents/sessions/auth-storage.ts index b2db8de3891d..ccce9118ddd6 100644 --- a/src/agents/sessions/auth-storage.ts +++ b/src/agents/sessions/auth-storage.ts @@ -16,21 +16,27 @@ import type { OAuthLoginCallbacks, OAuthProviderId, } from "../../llm/utils/oauth/types.js"; +import { KeyedAsyncQueue } from "../../plugin-sdk/keyed-async-queue.js"; import { OAuthProviderConfiguredUnavailableError } from "../../plugins/provider-runtime.errors.js"; import type { OpenClawAgentDatabase } from "../../state/openclaw-agent-db.js"; -import { AUTH_STORE_VERSION, OAUTH_REFRESH_LOCK_OPTIONS } from "../auth-profiles/constants.js"; +import { + AUTH_STORE_VERSION, + OAUTH_REFRESH_CALL_TIMEOUT_MS, + OAUTH_REFRESH_LOCK_OPTIONS, +} from "../auth-profiles/constants.js"; import { assertAuthProfileMigrationReady, AuthProfileMigrationRequiredError, AuthProfileStoreUnreadableError, } from "../auth-profiles/legacy-source-diagnostic.js"; +import { runRetainedOAuthRefreshOperation } from "../auth-profiles/oauth-manager.js"; +import { resolveOAuthRefreshConflict } from "../auth-profiles/oauth-shared.js"; import { resolveOAuthRefreshLockPath } from "../auth-profiles/paths.js"; import { loadPersistedAuthProfileStore } from "../auth-profiles/persisted.js"; import { getRuntimeAuthProfileStoreSnapshotCore } from "../auth-profiles/runtime-snapshots.js"; import { inspectPersistedAuthProfileStateRaw, inspectPersistedAuthProfileStoreRaw, - resolveAuthProfileDatabasePath, runAuthProfileWriteTransaction, } from "../auth-profiles/sqlite.js"; import { loadPersistedAuthProfileState } from "../auth-profiles/state.js"; @@ -38,7 +44,7 @@ import { loadAuthProfileStoreForSecretsRuntime, saveAuthProfileStore, } from "../auth-profiles/store.js"; -import type { AuthProfileStore } from "../auth-profiles/types.js"; +import type { AuthProfileCredential, AuthProfileStore } from "../auth-profiles/types.js"; import { getAgentDir } from "../config.js"; import { getAuthStorageOAuthProviderRegistry, @@ -128,7 +134,6 @@ type LockResult = { export interface AuthStorageBackend { readonly migrationOwnerAgentDir?: string; withLock(fn: (current: string | undefined) => LockResult): T; - withLockAsync(fn: (current: string | undefined) => Promise>): Promise; } function projectAuthStorageData(store: AuthProfileStore | null): AuthStorageData { @@ -303,11 +308,6 @@ class SqliteAuthStorageBackend implements AuthStorageBackend { return current ? [current] : this.preparedStore ? [this.preparedStore] : []; } - private readRaw(): AuthProfileStore { - assertAuthProfileMigrationReady(this.agentDir); - return loadSqliteAuthStorageStore(this.agentDir); - } - withLock(fn: (current: string | undefined) => LockResult): T { assertAuthProfileMigrationReady(this.agentDir); const snapshots = this.resolveMaterializedRuntimeStores(); @@ -331,46 +331,6 @@ class SqliteAuthStorageBackend implements AuthStorageBackend { return result; }); } - - async withLockAsync(fn: (current: string | undefined) => Promise>): Promise { - assertAuthProfileMigrationReady(this.agentDir); - return await withFileLock( - resolveAuthProfileDatabasePath(this.agentDir), - OAUTH_REFRESH_LOCK_OPTIONS, - async () => { - const initialRaw = this.readRaw(); - const initialData = projectAuthoritativeAuthStorageData( - initialRaw, - this.resolveMaterializedRuntimeStores(), - ); - const { result, next } = await fn(JSON.stringify(initialData)); - if (next === undefined) { - return result; - } - assertAuthProfileMigrationReady(this.agentDir); - runAuthProfileWriteTransaction(this.agentDir, (database) => { - const authoritative = loadSqliteAuthStorageStore(this.agentDir, database); - if (!isDeepStrictEqual(authoritative.profiles, initialRaw.profiles)) { - throw new AuthStoragePersistenceError( - "Cannot update auth storage because its SQLite credentials changed concurrently.", - undefined, - ); - } - saveAuthProfileStore( - applyAuthStorageData(authoritative, JSON.parse(next) as AuthStorageData, initialData), - this.agentDir, - { - filterExternalAuthProfiles: false, - preserveStateProfileIds: collectStateOnlyAuthProfileIds(authoritative), - syncExternalCli: false, - }, - database, - ); - }); - return result; - }, - ); - } } /** @@ -399,10 +359,6 @@ export class FileAuthStorageBackend implements AuthStorageBackend { withLock(fn: (current: string | undefined) => LockResult): T { return this.delegate.withLock(fn); } - - async withLockAsync(fn: (current: string | undefined) => Promise>): Promise { - return await this.delegate.withLockAsync(fn); - } } export class InMemoryAuthStorageBackend implements AuthStorageBackend { @@ -415,14 +371,6 @@ export class InMemoryAuthStorageBackend implements AuthStorageBackend { } return result; } - - async withLockAsync(fn: (current: string | undefined) => Promise>): Promise { - const { result, next } = await fn(this.value); - if (next !== undefined) { - this.value = next; - } - return result; - } } /** @@ -436,6 +384,7 @@ export class AuthStorage { private errors: Error[] = []; private storage: AuthStorageBackend; private migrationOwnerAgentDir?: string; + private oauthRefreshQueue = new KeyedAsyncQueue(); private constructor(storage: AuthStorageBackend, migrationOwnerAgentDir?: string) { this.storage = storage; @@ -698,62 +647,96 @@ export class AuthStorage { private async refreshOAuthTokenWithLock( providerId: OAuthProviderId, ): Promise<{ apiKey: string; newCredentials: OAuthCredentials } | null> { - const provider = getAuthStorageOAuthProviderRegistry(this).get(providerId); - - const refresh = async () => - await this.storage.withLockAsync(async (current) => { - const currentData = this.parseStorageData(current); - this.data = currentData; - this.loadError = null; - - const cred = currentData[providerId]; - if (cred?.type !== "oauth") { - return { result: null }; - } - - if (Date.now() < cred.expires) { - if (provider) { - return { result: { apiKey: provider.getApiKey(cred), newCredentials: cred } }; + return runRetainedOAuthRefreshOperation({ + timeoutMs: OAUTH_REFRESH_CALL_TIMEOUT_MS, + run: async (signal) => { + const provider = getAuthStorageOAuthProviderRegistry(this).get(providerId); + const resolveCredential = async (credential: OAuthCredentials, forceRefresh = false) => { + signal.throwIfAborted(); + if (!provider) { + return await resolveAuthStoragePluginOAuthCredential( + providerId, + credential, + forceRefresh, + ); } - return { result: await resolveAuthStoragePluginOAuthCredential(providerId, cred, false) }; - } - - const oauthCreds: Record = {}; - for (const [key, value] of Object.entries(currentData)) { - if (value.type === "oauth") { - oauthCreds[key] = value; + if (!forceRefresh) { + return { apiKey: provider.getApiKey(credential), newCredentials: credential }; } - } - - const refreshed = provider - ? await getAuthStorageOAuthProviderRegistry(this).getApiKey(providerId, oauthCreds) - : await resolveAuthStoragePluginOAuthCredential(providerId, cred, true); - if (!refreshed) { - return { result: null }; - } - - const refreshedCredential: OAuthCredential = { - type: "oauth", - ...refreshed.newCredentials, + const refreshed = await provider.refreshToken(credential); + return { apiKey: provider.getApiKey(refreshed), newCredentials: refreshed }; }; - const merged: AuthStorageData = { - ...currentData, - [providerId]: refreshedCredential, - }; - this.data = merged; - this.loadError = null; - return { result: refreshed, next: JSON.stringify(merged, null, 2) }; - }); + signal.throwIfAborted(); + return await this.oauthRefreshQueue.enqueue(providerId, async () => { + signal.throwIfAborted(); + const refresh = async () => { + signal.throwIfAborted(); + const snapshot = this.storage.withLock((current) => { + const currentData = this.parseStorageData(current); + return { result: { currentData, credential: currentData[providerId] } }; + }); + this.data = snapshot.currentData; + this.loadError = null; + if (snapshot.credential?.type !== "oauth") { + return null; + } + const credential = snapshot.credential; + if (Date.now() < credential.expires) { + return await resolveCredential(credential); + } - const result = this.migrationOwnerAgentDir - ? await withFileLock( - resolveOAuthRefreshLockPath(providerId, `${providerId}:default`), - OAUTH_REFRESH_LOCK_OPTIONS, - refresh, - ) - : await refresh(); + const refreshed = await resolveCredential(credential, true); + if (!refreshed) { + return null; + } + const refreshedCredential = { + ...credential, + ...refreshed.newCredentials, + type: "oauth", + provider: providerId, + } satisfies Extract; + if (Date.now() >= refreshedCredential.expires) { + throw new Error("OAuth provider returned an expired credential"); + } - return result; + const persisted = this.storage.withLock((current) => { + const data = this.parseStorageData(current); + const decision = resolveOAuthRefreshConflict({ + authoritative: data[providerId] + ? ({ ...data[providerId], provider: providerId } as AuthProfileCredential) + : undefined, + attempted: { ...credential, provider: providerId }, + refreshed: refreshedCredential, + }); + if (!decision?.persist) { + return { result: { data, credential: decision?.credential ?? null } }; + } + const nextData = { ...data, [providerId]: decision.credential }; + return { + result: { data: nextData, credential: decision.credential }, + next: JSON.stringify(nextData, null, 2), + }; + }); + this.data = persisted.data; + this.loadError = null; + if (!persisted.credential) { + return null; + } + return persisted.credential === refreshedCredential + ? { apiKey: refreshed.apiKey, newCredentials: persisted.credential } + : await resolveCredential(persisted.credential); + }; + + return this.migrationOwnerAgentDir + ? await withFileLock( + resolveOAuthRefreshLockPath(providerId, `${providerId}:default`), + OAUTH_REFRESH_LOCK_OPTIONS, + refresh, + ) + : await refresh(); + }); + }, + }); } /** diff --git a/src/infra/provider-usage.auth.normalizes-keys.test.ts b/src/infra/provider-usage.auth.normalizes-keys.test.ts index df9d9907d2c8..c10abdd10a41 100644 --- a/src/infra/provider-usage.auth.normalizes-keys.test.ts +++ b/src/infra/provider-usage.auth.normalizes-keys.test.ts @@ -143,7 +143,6 @@ const providerRuntimeMocks = vi.hoisted(() => ({ prepareProviderDynamicModel: vi.fn(async () => {}), prepareProviderExtraParams: vi.fn(() => undefined), prepareProviderRuntimeAuth: vi.fn(async () => undefined), - refreshProviderOAuthCredentialWithPlugin: vi.fn(async () => undefined), resolveProviderBinaryThinking: vi.fn(() => undefined), resolveProviderCacheTtlEligibility: vi.fn(() => undefined), resolveProviderCapabilitiesWithPlugin: vi.fn(() => undefined),