Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
71 changes: 68 additions & 3 deletions apps/sim/lib/oauth/credential-service.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -442,11 +442,16 @@ describe('OAuth access-token refresh headroom', () => {
describe.each(readers)('$name', ({ read, throwsOnFailure }) => {
it.each([
{ providerId: 'google-drive', remainingMs: 120_000, refresh: 'wait' },
{ providerId: 'google-drive', remainingMs: 240_000, refresh: 'ahead' },
{ providerId: 'google-drive', remainingMs: 300_001, refresh: 'none' },
{ providerId: 'google-drive', remainingMs: 240_000, refresh: 'wait' },
{ providerId: 'google-drive', remainingMs: 300_000, refresh: 'wait' },
{ providerId: 'google-drive', remainingMs: 300_001, refresh: 'ahead' },
/** Inside the refresh-ahead window for every account, whatever its jitter. */
{ providerId: 'google-drive', remainingMs: 540_000, refresh: 'ahead' },
{ providerId: 'google-drive', remainingMs: 600_001, refresh: 'none' },
{ providerId: 'google-drive', remainingMs: null, refresh: 'none' },
{ providerId: 'confluence', remainingMs: 300_000, refresh: 'wait' },
{ providerId: 'confluence', remainingMs: 300_001, refresh: 'none' },
{ providerId: 'confluence', remainingMs: 540_000, refresh: 'none' },
])(
'handles $providerId with $remainingMs milliseconds remaining',
async ({ providerId, remainingMs, refresh }) => {
Expand All @@ -463,7 +468,7 @@ describe('OAuth access-token refresh headroom', () => {
it('uses a still-valid Google token while its refresh is still with the provider', async () => {
const provider = createDeferred<unknown>()
mocks.refreshOAuthToken.mockReturnValue(provider.promise)
await expect(read(createOAuthAccount(240_000))).resolves.toBe('original-access-token')
await expect(read(createOAuthAccount(360_000))).resolves.toBe('original-access-token')
await vi.waitFor(() => expect(mocks.refreshOAuthToken).toHaveBeenCalledOnce())
expect(dbChainMockFns.set).not.toHaveBeenCalled()
provider.resolve({
Expand Down Expand Up @@ -495,6 +500,16 @@ describe('OAuth access-token refresh headroom', () => {
providerId: 'google-drive',
remainingMs: 60_000,
},
{
label: 'a Google token with less than five minutes left',
providerId: 'google-drive',
remainingMs: 240_000,
},
{
label: 'a Google token with exactly five minutes left',
providerId: 'google-drive',
remainingMs: 300_000,
},
])('waits for the provider with $label', async ({ providerId, remainingMs }) => {
const provider = createDeferred<unknown>()
mocks.refreshOAuthToken.mockReturnValue(provider.promise)
Expand Down Expand Up @@ -632,6 +647,56 @@ describe('OAuth access-token refresh headroom', () => {
expect(mocks.refreshOAuthToken).toHaveBeenCalledTimes(1)
})

it('shares one background refresh across every caller served inside the refresh-ahead window', async () => {
const { coalesceLocally } = await vi.importActual<
typeof import('@/lib/concurrency/singleflight')
>('@/lib/concurrency/singleflight')
mocks.coalesceLocally.mockImplementation(coalesceLocally)
const provider = createDeferred<unknown>()
mocks.refreshOAuthToken.mockReturnValue(provider.promise)
const row = createOAuthAccount(360_000)
await expect(refreshTokenIfNeeded('test', row, RAW_ACCOUNT_ID)).resolves.toEqual({
accessToken: 'original-access-token',
refreshed: false,
})
await expect(refreshTokenIfNeeded('test', row, RAW_ACCOUNT_ID)).resolves.toEqual({
accessToken: 'original-access-token',
refreshed: false,
})
await vi.waitFor(() => expect(mocks.refreshOAuthToken).toHaveBeenCalled())
provider.resolve({
ok: true,
accessToken: 'refreshed-access-token',
refreshToken: 'rotated-refresh-token',
expiresIn: 3600,
})
await vi.waitFor(() => expect(dbChainMockFns.set).toHaveBeenCalled())
await flushMicrotasks()
expect(mocks.refreshOAuthToken).toHaveBeenCalledOnce()
})

it('uses a token another writer stored after the caller read the account instead of refreshing again', async () => {
queueCredentialAccount(createOAuthAccount())
queueTableRows(account, [
{ ...createOAuthAccount(3_600_000), accessToken: 'stored-meanwhile-token', updatedAt: now },
])
await expect(
resolveCredentialTokenBundle(RAW_CREDENTIAL_ID, RAW_USER_ID, 'test')
).resolves.toEqual({ accessToken: 'stored-meanwhile-token' })
expect(mocks.refreshOAuthToken).not.toHaveBeenCalled()
})

it('refreshes when the token another writer stored is itself due to refresh ahead', async () => {
queueCredentialAccount(createOAuthAccount())
queueTableRows(account, [
{ ...createOAuthAccount(420_000), accessToken: 'stored-meanwhile-token', updatedAt: now },
])
await expect(
resolveCredentialTokenBundle(RAW_CREDENTIAL_ID, RAW_USER_ID, 'test')
).resolves.toEqual({ accessToken: 'refreshed-access-token' })
expect(mocks.refreshOAuthToken).toHaveBeenCalledOnce()
})

it.each([
{ remainingMs: 300_000, refresh: true },
{ remainingMs: 300_001, refresh: false },
Expand Down
82 changes: 62 additions & 20 deletions apps/sim/lib/oauth/credential-service.ts
Original file line number Diff line number Diff line change
Expand Up @@ -886,8 +886,8 @@ interface CoalescedRefreshOptions {
/** External provider account id (`account.accountId`), used to scope Slack refreshes per installation. */
providerAccountId?: string | null
oauthConfig?: string | null
/** The account row's refresh token and recorded revocation, read together. */
chain?: RevocableChain
/** The account row the caller read: its tokens, recorded revocation, and version, read together. */
chain?: StoredOAuthTokens
requestId?: string
userId?: string
privacyMode?: 'selector'
Expand Down Expand Up @@ -1003,9 +1003,18 @@ function isOAuthAccessTokenExpiring(
return expiresAt != null && expiresAt.getTime() <= now.getTime() + refreshWindowMs
}

/** Below this much remaining validity a caller waits for the refresh even where it could refresh ahead. */
const OAUTH_ACCESS_TOKEN_MUST_REFRESH_MS = 2 * 60 * 1000
/** Spreads when credentials that expire together start refreshing ahead. */
/**
* Where a provider can refresh ahead of expiry, a caller is served the stored token with this much
* validity left or less while a background refresh replaces it. It starts above
* {@link OAUTH_ACCESS_TOKEN_REFRESH_WINDOW_MS}, so a served token keeps the same headroom a caller
* that waits for its refresh gets: callers that resolve once and reuse the token through a long
* batch (the pollers) depend on it.
*/
const OAUTH_REFRESH_AHEAD_WINDOW_MS = 10 * 60 * 1000
/**
* Spreads when credentials that expire together start refreshing ahead. Kept below the gap between
* the two windows so every account still refreshes ahead before callers have to wait.
*/
const OAUTH_REFRESH_AHEAD_JITTER_MS = 60 * 1000

/**
Expand Down Expand Up @@ -1061,15 +1070,17 @@ function planAccessTokenRefresh(
: { kind: 'fresh' }
}
if (!accessToken) return { kind: 'must-refresh', refreshToken }
const refreshesAhead = refreshesAheadOfExpiry(providerId)
if (accessTokenExpiresAt && isOAuthAccessTokenExpiring(accessTokenExpiresAt, providerId, now)) {
const remainingMs = accessTokenExpiresAt.getTime() - now.getTime()
if (!refreshesAhead || remainingMs <= OAUTH_ACCESS_TOKEN_MUST_REFRESH_MS) {
return { kind: 'must-refresh', refreshToken }
}
if (remainingMs <= OAUTH_ACCESS_TOKEN_REFRESH_WINDOW_MS - refreshAheadJitterMs(accountId)) {
return { kind: 'refresh-ahead', refreshToken }
}
return { kind: 'must-refresh', refreshToken }
Comment thread
waleedlatif1 marked this conversation as resolved.
Resolved
}
const refreshesAhead = refreshesAheadOfExpiry(providerId)
if (
refreshesAhead &&
accessTokenExpiresAt &&
accessTokenExpiresAt.getTime() - now.getTime() <=
OAUTH_REFRESH_AHEAD_WINDOW_MS - refreshAheadJitterMs(accountId)
) {
return { kind: 'refresh-ahead', refreshToken }
}
/** Microsoft refresh tokens expire after 90 days without use. */
const refreshTokenAging =
Expand Down Expand Up @@ -1307,6 +1318,31 @@ async function performCoalescedRefresh({
if (recordedMeanwhile && !isRevocationReprobeDue(current)) {
return revokedOutcome(recordedMeanwhile)
}
/**
* A writer that stored a new token after the caller's read (a refresh that finished just
* before this one took the lock, or a reconnect) already answers it, so this refresh
* would only spend a provider call. The token must still clear the headroom a waiting
* caller is promised, and sit outside the refresh-ahead window so the next caller does
* not start another refresh straight away.
*/
const storedSinceRead =
!slackTeamId &&
!recordedMeanwhile &&
current &&
chain &&
current.updatedAt > chain.updatedAt &&
current.accessToken !== chain.accessToken &&
!(
refreshesAheadOfExpiry(providerId) &&
current.accessTokenExpiresAt &&
current.accessTokenExpiresAt.getTime() <= Date.now() + OAUTH_REFRESH_AHEAD_WINDOW_MS
)
? usableStoredToken(current, providerId)
: null
if (storedSinceRead) {
logger.info('Using the token stored since this caller read the account', logContext)
return tokenOutcome(storedSinceRead)
}
let chainVersion = current?.updatedAt
let refreshTokenToUse = refreshToken
let slackChainVersion: Date | null = null
Expand Down Expand Up @@ -1532,6 +1568,18 @@ async function performCoalescedRefresh({
}
}

/** Starts a refresh ahead of expiry; callers in the window share it through the refresh's own single-flight. */
function refreshInBackground(options: CoalescedRefreshOptions, logPrefix: string): void {
logger.debug(`${logPrefix}Refreshing ahead of expiry; using the still-valid access token`)
// The refresh settles and logs its own failures; this catches only a setup error before it starts.
performCoalescedRefresh(options).catch((error: unknown) => {
logger.error('Background token refresh failed to start', {
providerId: options.providerId,
...(options.privacyMode === 'selector' ? {} : { error: toError(error).message }),
})
})
}

type StoredAccessTokenResult =
| { kind: 'unrefreshable' }
| { kind: 'token'; accessToken: string | null; refreshed: boolean }
Expand Down Expand Up @@ -1561,13 +1609,7 @@ async function resolveStoredAccessToken(
}
const prefix = refresh.requestId ? `[${refresh.requestId}] ` : ''
if (plan.kind === 'refresh-ahead') {
performCoalescedRefresh(options).catch((error: unknown) => {
logger.error('Background token refresh failed', {
providerId: refresh.providerId,
...(refresh.privacyMode === 'selector' ? {} : { error: toError(error).message }),
})
})
logger.info(`${prefix}Refreshing ahead of expiry; using the still-valid access token`)
refreshInBackground(options, prefix)
return { kind: 'token', accessToken: stored.accessToken, refreshed: false }
}
const fresh = await waitWithAbort(performCoalescedRefresh(options), signal)
Expand Down