From eb55d0b5b2038e0e957844ae6c959999c2616488 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Sat, 31 Oct 2026 00:59:41 -0700 Subject: [PATCH] improvement(oauth): serve still-valid tokens while refreshing ahead of expiry --- apps/sim/lib/oauth/credential-service.test.ts | 120 +++++- apps/sim/lib/oauth/credential-service.ts | 342 ++++++++++-------- apps/sim/lib/sim-search/live/accounts.ts | 25 +- .../lib/sim-search/live/application.test.ts | 50 ++- apps/sim/lib/sim-search/live/application.ts | 15 +- scripts/check-explicit-any.baseline.json | 2 +- 6 files changed, 374 insertions(+), 180 deletions(-) diff --git a/apps/sim/lib/oauth/credential-service.test.ts b/apps/sim/lib/oauth/credential-service.test.ts index 0553e56cce0..45e4ccc74c2 100644 --- a/apps/sim/lib/oauth/credential-service.test.ts +++ b/apps/sim/lib/oauth/credential-service.test.ts @@ -1,6 +1,12 @@ import { generateKeyPairSync, verify } from 'node:crypto' import { account, credential } from '@sim/db/schema' -import { dbChainMockFns, queueTableRows, resetDbChainMock } from '@sim/testing' +import { + createDeferred, + dbChainMockFns, + flushMicrotasks, + queueTableRows, + resetDbChainMock, +} from '@sim/testing' import { encryptionMock, encryptionMockFns } from '@sim/testing/mocks/encryption.mock' import { getMockLogger } from '@sim/testing/mocks/logger.mock' import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' @@ -435,14 +441,77 @@ describe('OAuth access-token refresh headroom', () => { describe.each(readers)('$name', ({ read, throwsOnFailure }) => { it.each([ - { remainingMs: 300_000, refresh: true }, - { remainingMs: 300_001, refresh: false }, - { remainingMs: null, refresh: false }, - ])('handles $remainingMs milliseconds remaining', async ({ remainingMs, refresh }) => { - await expect(read(createOAuthAccount(remainingMs))).resolves.toBe( - refresh ? 'refreshed-access-token' : 'original-access-token' + { 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: null, refresh: 'none' }, + { providerId: 'confluence', remainingMs: 300_000, refresh: 'wait' }, + { providerId: 'confluence', remainingMs: 300_001, refresh: 'none' }, + ])( + 'handles $providerId with $remainingMs milliseconds remaining', + async ({ providerId, remainingMs, refresh }) => { + await expect(read({ ...createOAuthAccount(remainingMs), providerId })).resolves.toBe( + refresh === 'wait' ? 'refreshed-access-token' : 'original-access-token' + ) + if (refresh === 'ahead') { + await vi.waitFor(() => expect(dbChainMockFns.set).toHaveBeenCalled()) + } + expect(mocks.refreshOAuthToken).toHaveBeenCalledTimes(refresh === 'none' ? 0 : 1) + } + ) + + it('uses a still-valid Google token while its refresh is still with the provider', async () => { + const provider = createDeferred() + mocks.refreshOAuthToken.mockReturnValue(provider.promise) + await expect(read(createOAuthAccount(240_000))).resolves.toBe('original-access-token') + await vi.waitFor(() => expect(mocks.refreshOAuthToken).toHaveBeenCalledOnce()) + expect(dbChainMockFns.set).not.toHaveBeenCalled() + provider.resolve({ + ok: true, + accessToken: 'refreshed-access-token', + refreshToken: 'rotated-refresh-token', + expiresIn: 3600, + }) + await vi.waitFor(() => + expect(dbChainMockFns.set).toHaveBeenCalledWith( + expect.objectContaining({ + accessToken: 'refreshed-access-token', + refreshToken: 'rotated-refresh-token', + }) + ) ) - expect(mocks.refreshOAuthToken).toHaveBeenCalledTimes(refresh ? 1 : 0) + const guard = JSON.stringify(dbChainMockFns.where.mock.calls.at(-1)) + expect(guard).toContain('original-refresh-token') + }) + + it.each([ + { + label: 'a provider not verified for refresh-ahead', + providerId: 'confluence', + remainingMs: 240_000, + }, + { + label: 'a Google token under the refresh floor', + providerId: 'google-drive', + remainingMs: 60_000, + }, + ])('waits for the provider with $label', async ({ providerId, remainingMs }) => { + const provider = createDeferred() + mocks.refreshOAuthToken.mockReturnValue(provider.promise) + let settled = false + const pending = read({ ...createOAuthAccount(remainingMs), providerId }).finally(() => { + settled = true + }) + await vi.waitFor(() => expect(mocks.refreshOAuthToken).toHaveBeenCalledOnce()) + await flushMicrotasks(20) + expect(settled).toBe(false) + provider.resolve({ + ok: true, + accessToken: 'refreshed-access-token', + refreshToken: 'rotated-refresh-token', + expiresIn: 3600, + }) + await expect(pending).resolves.toBe('refreshed-access-token') }) it('still refreshes a missing access token without a known expiry', async () => { @@ -614,7 +683,40 @@ describe('OAuth access-token refresh headroom', () => { accessToken: 'original-access-token', refreshed: false, }) - expect(mocks.refreshOAuthToken).toHaveBeenCalledTimes(2) + await vi.waitFor(() => expect(mocks.refreshOAuthToken).toHaveBeenCalledTimes(2)) + }) + + it('stops a cancelled caller waiting while a concurrent caller still gets the refreshed token', async () => { + const { coalesceLocally } = await vi.importActual< + typeof import('@/lib/concurrency/singleflight') + >('@/lib/concurrency/singleflight') + mocks.coalesceLocally.mockImplementation(coalesceLocally) + const provider = createDeferred() + mocks.refreshOAuthToken.mockReturnValue(provider.promise) + const controller = new AbortController() + const reason = new DOMException('Search deadline', 'TimeoutError') + queueCredentialAccount(createOAuthAccount()) + queueCredentialAccount(createOAuthAccount()) + const cancelled = resolveCredentialTokenBundle( + RAW_CREDENTIAL_ID, + RAW_USER_ID, + 'test', + undefined, + undefined, + { signal: controller.signal } + ) + const concurrent = resolveCredentialTokenBundle(RAW_CREDENTIAL_ID, RAW_USER_ID, 'test') + await vi.waitFor(() => expect(mocks.refreshOAuthToken).toHaveBeenCalledOnce()) + controller.abort(reason) + await expect(cancelled).rejects.toBe(reason) + provider.resolve({ + ok: true, + accessToken: 'refreshed-access-token', + refreshToken: 'rotated-refresh-token', + expiresIn: 3600, + }) + await expect(concurrent).resolves.toEqual({ accessToken: 'refreshed-access-token' }) + expect(mocks.refreshOAuthToken).toHaveBeenCalledOnce() }) }) diff --git a/apps/sim/lib/oauth/credential-service.ts b/apps/sim/lib/oauth/credential-service.ts index 80ef8c8f9ad..9ffd9bdf381 100644 --- a/apps/sim/lib/oauth/credential-service.ts +++ b/apps/sim/lib/oauth/credential-service.ts @@ -1,4 +1,4 @@ -import { createHmac, createSign } from 'crypto' +import { createHash, createHmac, createSign } from 'crypto' import { db } from '@sim/db' import { account, credential } from '@sim/db/schema' import { sealedTextEquals } from '@sim/db/sealed-text' @@ -91,7 +91,10 @@ export interface CredentialTokenResolutionOptions { privacyMode?: 'selector' /** GitHub installation content tokens may only address one connector repository. */ githubRepositoryScope?: GitHubInstallationRepositoryScope - /** Cancels supported service-account token exchanges and retry waits. */ + /** + * Cancels supported service-account token exchanges and retry waits, and stops this caller + * waiting on an OAuth refresh without cancelling the refresh other callers share. + */ signal?: AbortSignal } @@ -1000,6 +1003,94 @@ 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. */ +const OAUTH_REFRESH_AHEAD_JITTER_MS = 60 * 1000 + +/** + * Google and Microsoft keep a refreshed-away access token valid until its own expiry and keep the + * old refresh token usable (Google never rotates it; Microsoft rotates without revoking the old + * one), so a caller can use the old token while the refresh runs and an abandoned refresh loses + * nothing. + */ +function refreshesAheadOfExpiry(providerId: string): boolean { + return ( + providerId === 'google' || providerId.startsWith('google-') || isMicrosoftProvider(providerId) + ) +} + +function refreshAheadJitterMs(accountId: string): number { + return ( + createHash('sha256').update(accountId).digest().readUInt32BE(0) % OAUTH_REFRESH_AHEAD_JITTER_MS + ) +} + +/** The stored token state of one OAuth account row that decides whether it refreshes. */ +interface StoredOAuthTokens extends RevocableChain { + accessToken: string | null + accessTokenExpiresAt: Date | null + refreshTokenExpiresAt?: Date | null + updatedAt: Date +} + +/** + * How a caller gets a usable access token from what an account row stores: + * - `fresh`: use the stored token (or report that none is stored). + * - `unrefreshable`: the stored token expired and nothing can refresh it. + * - `refresh-ahead`: use the still-valid stored token and refresh in the background. + * - `refresh-or-reuse`: wait for a refresh the token's own validity does not need, keeping the + * stored token if it fails (Microsoft refresh-token aging, Instagram long-lived renewal). + * - `must-refresh`: wait for the refresh; the stored token cannot serve the request. + */ +type AccessTokenRefreshPlan = + | { kind: 'fresh' } + | { kind: 'unrefreshable' } + | { kind: 'refresh-ahead' | 'refresh-or-reuse' | 'must-refresh'; refreshToken: string } + +function planAccessTokenRefresh( + accountId: string, + providerId: string, + stored: StoredOAuthTokens, + now: Date +): AccessTokenRefreshPlan { + const { accessToken, refreshToken, accessTokenExpiresAt, refreshTokenExpiresAt } = stored + if (!refreshToken) { + return accessTokenExpiresAt && accessTokenExpiresAt <= now + ? { kind: 'unrefreshable' } + : { 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 } + } + } + /** Microsoft refresh tokens expire after 90 days without use. */ + const refreshTokenAging = + isMicrosoftProvider(providerId) && + refreshTokenExpiresAt != null && + refreshTokenExpiresAt.getTime() <= + now.getTime() + PROACTIVE_REFRESH_THRESHOLD_DAYS * 24 * 60 * 60 * 1000 + /** Instagram long-lived tokens can only be refreshed while still valid. */ + const instagramRenewal = + isInstagramProvider(providerId) && + shouldProactivelyRefreshInstagramToken({ + accessTokenExpiresAt, + updatedAt: stored.updatedAt, + now, + }) + if (refreshTokenAging || instagramRenewal) { + return { kind: refreshesAhead ? 'refresh-ahead' : 'refresh-or-reuse', refreshToken } + } + return { kind: 'fresh' } +} + /** * Lock budgets sized past the provider call: the lease covers * {@link TOKEN_REFRESH_TIMEOUT_MS} plus headroom for the account read before it and @@ -1441,6 +1532,55 @@ async function performCoalescedRefresh({ } } +type StoredAccessTokenResult = + | { kind: 'unrefreshable' } + | { kind: 'token'; accessToken: string | null; refreshed: boolean } + | { kind: 'failed'; outcome: RefreshOutcome | null } + +/** + * The access token a caller can use now from what an account row stores, refreshing as + * {@link planAccessTokenRefresh} decides. A refresh ahead of expiry runs detached on stored data + * alone, never the caller's signal, so it outlives the request; if the process exits first, the + * lock lease lapses and the next caller refreshes. `signal` stops only this caller's wait on a + * refresh, which keeps running for every caller sharing it. + */ +async function resolveStoredAccessToken( + stored: StoredOAuthTokens, + refresh: Omit, + signal?: AbortSignal +): Promise { + const plan = planAccessTokenRefresh(refresh.accountId, refresh.providerId, stored, new Date()) + if (plan.kind === 'unrefreshable') return { kind: 'unrefreshable' } + if (plan.kind === 'fresh') { + return { kind: 'token', accessToken: stored.accessToken, refreshed: false } + } + const options: CoalescedRefreshOptions = { + ...refresh, + refreshToken: plan.refreshToken, + chain: stored, + } + 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`) + return { kind: 'token', accessToken: stored.accessToken, refreshed: false } + } + const fresh = await waitWithAbort(performCoalescedRefresh(options), signal) + if (fresh?.kind === 'token') { + return { kind: 'token', accessToken: fresh.accessToken, refreshed: true } + } + if (plan.kind === 'refresh-or-reuse' && stored.accessToken) { + logger.info(`${prefix}Refresh unavailable; reusing still-valid access token`) + return { kind: 'token', accessToken: stored.accessToken, refreshed: false } + } + return { kind: 'failed', outcome: fresh } +} + export async function getOAuthToken(userId: string, providerId: string): Promise { const connections = await db .select({ @@ -1449,6 +1589,7 @@ export async function getOAuthToken(userId: string, providerId: string): Promise accessToken: account.accessToken, refreshToken: account.refreshToken, accessTokenExpiresAt: account.accessTokenExpiresAt, + refreshTokenExpiresAt: account.refreshTokenExpiresAt, idToken: account.idToken, scope: account.scope, updatedAt: account.updatedAt, @@ -1466,51 +1607,28 @@ export async function getOAuthToken(userId: string, providerId: string): Promise } const credential = connections[0] - - const now = new Date() - const tokenExpiry = credential.accessTokenExpiresAt - if (!credential.refreshToken && tokenExpiry && tokenExpiry <= now) { + const result = await resolveStoredAccessToken(credential, { + accountId: credential.id, + providerId, + providerAccountId: credential.providerAccountId, + oauthConfig: credential.oauthConfig, + userId, + }) + if (result.kind === 'unrefreshable') { logger.warn('OAuth access token expired and cannot be refreshed; reconnect the account') return null } - const accessTokenNeedsRefresh = - !!credential.refreshToken && - (!credential.accessToken || isOAuthAccessTokenExpiring(tokenExpiry, providerId, now)) - const instagramNeedsProactiveRefresh = - !!credential.refreshToken && - isInstagramProvider(providerId) && - shouldProactivelyRefreshInstagramToken({ - accessTokenExpiresAt: credential.accessTokenExpiresAt, - updatedAt: credential.updatedAt, - now, - }) - - if (accessTokenNeedsRefresh || instagramNeedsProactiveRefresh) { - const fresh = await performCoalescedRefresh({ - accountId: credential.id, - providerId, - refreshToken: credential.refreshToken!, - providerAccountId: credential.providerAccountId, - oauthConfig: credential.oauthConfig, - chain: credential, - userId, - }) - if (fresh?.kind === 'token') return fresh.accessToken - if (!accessTokenNeedsRefresh && credential.accessToken) { - return credential.accessToken - } - return null - } - - if (!credential.accessToken) { + if (result.kind === 'failed') return null + if (!result.accessToken) { logger.warn( `Access token is null and no refresh attempted or available for user ${userId}, provider ${providerId}` ) return null } - - logger.info(`Found valid OAuth token for user ${userId}, provider ${providerId}`) - return credential.accessToken + if (!result.refreshed) { + logger.info(`Found valid OAuth token for user ${userId}, provider ${providerId}`) + } + return result.accessToken } /** @@ -1553,79 +1671,31 @@ export async function resolveCredentialTokenBundle( return null } - const accessTokenExpiresAt = credential.accessTokenExpiresAt - const refreshTokenExpiresAt = credential.refreshTokenExpiresAt - const now = new Date() - - if (!credential.refreshToken && accessTokenExpiresAt && accessTokenExpiresAt <= now) { - logger.warn('OAuth access token expired and cannot be refreshed; reconnect the account') - return null - } - - const accessTokenNeedsRefresh = - !!credential.refreshToken && - (!credential.accessToken || - isOAuthAccessTokenExpiring(accessTokenExpiresAt, credential.providerId, now)) - - // Check if we should proactively refresh to prevent refresh token expiry - // This applies to Microsoft providers whose refresh tokens expire after 90 days of inactivity - const proactiveRefreshThreshold = new Date( - now.getTime() + PROACTIVE_REFRESH_THRESHOLD_DAYS * 24 * 60 * 60 * 1000 - ) - const refreshTokenNeedsProactiveRefresh = - !!credential.refreshToken && - isMicrosoftProvider(credential.providerId) && - refreshTokenExpiresAt && - refreshTokenExpiresAt <= proactiveRefreshThreshold - - // Instagram long-lived tokens can only be refreshed while still valid. - const instagramNeedsProactiveRefresh = - !!credential.refreshToken && - isInstagramProvider(credential.providerId) && - shouldProactivelyRefreshInstagramToken({ - accessTokenExpiresAt, - updatedAt: credential.updatedAt, - now, - }) - - const shouldRefresh = - accessTokenNeedsRefresh || refreshTokenNeedsProactiveRefresh || instagramNeedsProactiveRefresh - - const accessToken = credential.accessToken - - if (shouldRefresh) { - const resolvedCredentialId = - (credential as { resolvedCredentialId?: string }).resolvedCredentialId ?? credentialId - - const fresh = await performCoalescedRefresh({ - accountId: resolvedCredentialId, + const result = await resolveStoredAccessToken( + credential, + { + accountId: credential.resolvedCredentialId, providerId: credential.providerId, - refreshToken: credential.refreshToken!, providerAccountId: credential.accountId, oauthConfig: credential.oauthConfig, - chain: credential, requestId, userId: credential.userId, privacyMode: options?.privacyMode, - }) - if (fresh?.kind === 'token') return { accessToken: fresh.accessToken } - - // If refresh was only triggered proactively (Microsoft refresh-token aging / - // Instagram long-lived nearing expiry), the still-valid access token is fine. - if (!accessTokenNeedsRefresh && accessToken) { - logger.info(`[${requestId}] Refresh unavailable; reusing still-valid access token`) - return { accessToken } - } + }, + options?.signal + ) + if (result.kind === 'unrefreshable') { + logger.warn('OAuth access token expired and cannot be refreshed; reconnect the account') return null } - if (!accessToken) { + if (result.kind === 'failed') return null + if (!result.accessToken) { // We have no access token and either no refresh token or not eligible to refresh logger.error(`[${requestId}] Missing access token for credential`) return null } - - logger.info(`[${requestId}] Access token is valid for credential`) - return { accessToken } + if (!result.refreshed) logger.info(`[${requestId}] Access token is valid for credential`) + return { accessToken: result.accessToken } } /** @@ -1665,71 +1735,29 @@ export async function refreshTokenIfNeeded( credential: any, credentialId: string ): Promise<{ accessToken: string; refreshed: boolean }> { - const resolvedCredentialId = credential.resolvedCredentialId ?? credentialId - - const accessTokenExpiresAt = credential.accessTokenExpiresAt - const refreshTokenExpiresAt = credential.refreshTokenExpiresAt - const now = new Date() - - if (!credential.refreshToken && accessTokenExpiresAt && accessTokenExpiresAt <= now) { - throw new Error('OAuth access token expired and cannot be refreshed; reconnect the account') - } - - const accessTokenNeedsRefresh = - !!credential.refreshToken && - (!credential.accessToken || - isOAuthAccessTokenExpiring(accessTokenExpiresAt, credential.providerId, now)) - - // Check if we should proactively refresh to prevent refresh token expiry - // This applies to Microsoft providers whose refresh tokens expire after 90 days of inactivity - const proactiveRefreshThreshold = new Date( - now.getTime() + PROACTIVE_REFRESH_THRESHOLD_DAYS * 24 * 60 * 60 * 1000 - ) - const refreshTokenNeedsProactiveRefresh = - !!credential.refreshToken && - isMicrosoftProvider(credential.providerId) && - refreshTokenExpiresAt && - refreshTokenExpiresAt <= proactiveRefreshThreshold - - // Instagram long-lived tokens can only be refreshed while still valid. - const instagramNeedsProactiveRefresh = - !!credential.refreshToken && - isInstagramProvider(credential.providerId) && - shouldProactivelyRefreshInstagramToken({ - accessTokenExpiresAt, - updatedAt: credential.updatedAt, - now, - }) - - const shouldRefresh = - accessTokenNeedsRefresh || refreshTokenNeedsProactiveRefresh || instagramNeedsProactiveRefresh - - // If token appears valid and present, return it directly - if (!shouldRefresh) { - logger.info(`[${requestId}] Access token is valid`) - return { accessToken: credential.accessToken, refreshed: false } - } - - const fresh = await performCoalescedRefresh({ - accountId: resolvedCredentialId, + const result = await resolveStoredAccessToken(credential, { + accountId: credential.resolvedCredentialId ?? credentialId, providerId: credential.providerId, - refreshToken: credential.refreshToken!, providerAccountId: credential.accountId, oauthConfig: credential.oauthConfig, - chain: credential, requestId, userId: credential.userId, }) - if (fresh?.kind === 'token') return { accessToken: fresh.accessToken, refreshed: true } - - if (!accessTokenNeedsRefresh && credential.accessToken) { - logger.info(`[${requestId}] Refresh unavailable; reusing still-valid access token`) - return { accessToken: credential.accessToken, refreshed: false } + if (result.kind === 'unrefreshable') { + throw new Error('OAuth access token expired and cannot be refreshed; reconnect the account') + } + if (result.kind === 'token') { + if (!result.refreshed) logger.info(`[${requestId}] Access token is valid`) + /** A row storing neither token returns its empty token as before; callers decide how to fail. */ + return { + accessToken: result.accessToken ?? credential.accessToken, + refreshed: result.refreshed, + } } - if (fresh?.kind === 'revoked') { + if (result.outcome?.kind === 'revoked') { throw new CredentialRevokedError( - `The ${credential.providerId} credential was revoked by the provider (${fresh.errorCode}). Reconnect it to continue.`, - { providerId: credential.providerId, errorCode: fresh.errorCode } + `The ${credential.providerId} credential was revoked by the provider (${result.outcome.errorCode}). Reconnect it to continue.`, + { providerId: credential.providerId, errorCode: result.outcome.errorCode } ) } throw new Error('Failed to refresh token') diff --git a/apps/sim/lib/sim-search/live/accounts.ts b/apps/sim/lib/sim-search/live/accounts.ts index 1764878252e..db08911c759 100644 --- a/apps/sim/lib/sim-search/live/accounts.ts +++ b/apps/sim/lib/sim-search/live/accounts.ts @@ -147,11 +147,16 @@ export async function listLiveAccounts( } /** Re-lists the member's accounts so a reference from an earlier request is checked afresh. */ -export async function resolveLiveAccount(owner: ResourceOwner, userId: string, accountId: string) { +export async function resolveLiveAccount( + owner: ResourceOwner, + userId: string, + accountId: string, + signal?: AbortSignal +) { const current = (await listLiveAccounts(owner, userId)).find((row) => row.id === accountId) if (!current) throw new NativeSearchError('reconnect', 'This search connection is no longer available.') - return resolveListedLiveAccount(owner, userId, current) + return resolveListedLiveAccount(owner, userId, current, signal) } export type ResolvedLiveAccount = Awaited> @@ -163,7 +168,8 @@ export type ResolvedLiveAccount = Awaited ({ accounts: vi.fn(), @@ -753,6 +753,54 @@ describe('authorized live retrieval', () => { for (const signal of signals) expect(signal).toBeInstanceOf(AbortSignal) }) + describe('when account token resolution never settles', () => { + beforeEach(() => { + vi.useFakeTimers() + vi.spyOn(AbortSignal, 'timeout').mockImplementation((milliseconds) => { + const controller = new AbortController() + setTimeout( + () => controller.abort(new DOMException('deadline', 'TimeoutError')), + milliseconds + ) + return controller.signal + }) + }) + afterEach(() => { + vi.mocked(AbortSignal.timeout).mockRestore() + vi.useRealTimers() + }) + + it('reports the account timed out within the search deadline', async () => { + mocks.resolveAccount.mockReturnValue(new Promise(() => {})) + const pending = searchLiveKnowledge.execute({ principal, input }) + await vi.advanceTimersByTimeAsync(20_000) + const result = await pending + expect(result.live?.accounts[0]?.status).toBe('timeout') + expect(mocks.search).not.toHaveBeenCalled() + }) + + it('fails a document read as retryable within the read deadline', async () => { + const search = await searchLiveKnowledge.execute({ principal, input }) + mocks.resolveAccount.mockReturnValue(new Promise(() => {})) + const read = readLiveDocument.execute({ + principal, + input: { + workspaceId: 'workspace', + documentId: search.results[0]?.documentId ?? '', + limit: 1, + resultSecretRegistry: new ResolvedSecretTraceRegistry([]), + }, + }) + const rejected = expect(read).rejects.toMatchObject({ + retryable: true, + message: expect.stringContaining('took too long'), + }) + await vi.advanceTimersByTimeAsync(15_000) + await rejected + expect(mocks.read).not.toHaveBeenCalled() + }) + }) + it('retries an underfilled search with any term for providers that require every term', async () => { mocks.requiresEveryTerm.mockReturnValue(true) const loose = { ...document, id: 'loose', url: 'https://docs.google.com/document/d/loose/edit' } diff --git a/apps/sim/lib/sim-search/live/application.ts b/apps/sim/lib/sim-search/live/application.ts index a068e0c9abc..ad446352c0f 100644 --- a/apps/sim/lib/sim-search/live/application.ts +++ b/apps/sim/lib/sim-search/live/application.ts @@ -27,7 +27,7 @@ import { resourceScopeKey, } from '@/lib/core/resource-scope' import { createPinnedConnectionPool } from '@/lib/core/security/input-validation.server' -import { mapWithConcurrency } from '@/lib/core/utils/concurrency' +import { mapWithConcurrency, waitWithAbort } from '@/lib/core/utils/concurrency' import { MANAGED_MCP_CONNECTORS } from '@/lib/credential-groups/managed-mcp-connectors' import { requireOrganizationSearchAvailable } from '@/lib/knowledge/access/availability' import { defineAuthorizedKnowledgeUseCase } from '@/lib/knowledge/application/authorized-knowledge-use-case' @@ -716,8 +716,9 @@ const liveKnowledgeSearch = defineAuthorizedKnowledgeUseCase({ let session: LiveAccountSession | undefined try { signal.throwIfAborted() + /** Bounds every resolution path, including those that take no signal (admin GitLab, managed OAuth). */ const resolved = await measureSearchStage('live.resolve', () => - resolveListedLiveAccount(input, userId, account) + waitWithAbort(resolveListedLiveAccount(input, userId, account, signal), signal) ) session = await measureSearchStage('live.session', () => openLiveAccountSession({ @@ -870,7 +871,7 @@ export type LiveReadInput = ResourceOwner & { signal?: AbortSignal } -/** Milliseconds one document read may spend on provider calls, after account resolution. */ +/** Milliseconds one document read may spend resolving its account and calling the provider. */ const READ_DEADLINE_MS = 15_000 function liveProviderName(provider: string): string { @@ -946,16 +947,16 @@ export const readLiveDocument = defineAuthorizedKnowledgeUseCase({ /** Cancellation by the caller passes through unclassified. */ const readFailure = (error: unknown, deadline?: AbortSignal) => input.signal?.aborted ? error : liveReadFailure(error, reference.provider, deadline) + const deadline = AbortSignal.timeout(READ_DEADLINE_MS) + const signal = input.signal ? AbortSignal.any([input.signal, deadline]) : deadline const [resolved, policies] = await Promise.all([ - resolveLiveAccount(input, userId, reference.account), + waitWithAbort(resolveLiveAccount(input, userId, reference.account, signal), signal), loadLiveSearchPolicies(input), ]).catch((error: unknown) => { - throw readFailure(error) + throw readFailure(error, deadline) }) if (resolved.account.provider !== reference.provider) throw new OrchestrationError('not_found', 'Document account changed') - const deadline = AbortSignal.timeout(READ_DEADLINE_MS) - const signal = input.signal ? AbortSignal.any([input.signal, deadline]) : deadline const pool = createPinnedConnectionPool() let document: NativeDocument let session: LiveAccountSession | undefined diff --git a/scripts/check-explicit-any.baseline.json b/scripts/check-explicit-any.baseline.json index dd2edbe420c..85770794ce5 100644 --- a/scripts/check-explicit-any.baseline.json +++ b/scripts/check-explicit-any.baseline.json @@ -1597,7 +1597,7 @@ "apps/sim/lib/mothership/tools/server/files/workspace-file.ts": 2, "apps/sim/lib/mothership/tools/server/search-sources.test.ts": 2, "apps/sim/lib/mothership/transport/control.ts": 1, - "apps/sim/lib/oauth/credential-service.ts": 4, + "apps/sim/lib/oauth/credential-service.ts": 1, "apps/sim/lib/oauth/github-installation.test.ts": 2, "apps/sim/lib/oauth/github-repositories.test.ts": 5, "apps/sim/lib/organizations/application/reads.ts": 1,