diff --git a/apps/sim/lib/concurrency/leader-lock.test.ts b/apps/sim/lib/concurrency/leader-lock.test.ts index 222e523b7b1..ce83c6e39bc 100644 --- a/apps/sim/lib/concurrency/leader-lock.test.ts +++ b/apps/sim/lib/concurrency/leader-lock.test.ts @@ -119,6 +119,32 @@ describe('withLeaderLock', () => { } }) + describe('when the follower wait ends with nothing', () => { + const waitedOut = { + key: 'k', + pollIntervalMs: 5, + maxWaitMs: 20, + onFollower: async () => null, + } + + it('leads once a leader that vanished leaves its lease to lapse', async () => { + redisConfigMockFns.mockAcquireLock.mockResolvedValueOnce(false) + const onLeader = vi.fn(async () => 'took-over') + await expect(withLeaderLock({ ...waitedOut, onLeader })).resolves.toBe('took-over') + expect(onLeader).toHaveBeenCalledOnce() + expect(redisConfigMockFns.mockAcquireLock).toHaveBeenCalledTimes(2) + expect(redisConfigMockFns.mockReleaseLock).toHaveBeenCalledOnce() + }) + + it('returns the follower result when a live holder still has the lease', async () => { + redisConfigMockFns.mockAcquireLock.mockResolvedValue(false) + const onLeader = vi.fn(async () => 'should-not-run') + await expect(withLeaderLock({ ...waitedOut, onLeader })).resolves.toBeNull() + expect(onLeader).not.toHaveBeenCalled() + expect(redisConfigMockFns.mockAcquireLock).toHaveBeenCalledTimes(2) + }) + }) + it('only one of N concurrent callers acquires the lock', async () => { // Track which calls won the lock: first one returns true, rest return false. let acquired = false diff --git a/apps/sim/lib/concurrency/leader-lock.ts b/apps/sim/lib/concurrency/leader-lock.ts index 567abd1c8a1..69949ad98fa 100644 --- a/apps/sim/lib/concurrency/leader-lock.ts +++ b/apps/sim/lib/concurrency/leader-lock.ts @@ -19,6 +19,13 @@ export interface LeaderLockOptions { onFollower: () => Promise } +/** + * Runs `onLeader` in the one caller that takes the lock for `key`; every other caller polls + * `onFollower` until it reports the leader's result or `maxWaitMs` passes. A follower whose wait + * ends with nothing tries the lock once more: a leader whose process died without releasing it + * leaves the lease to lapse, and that follower then leads instead of failing. Losing that try + * means a live holder has the lease, and the follower returns `null`. + */ export async function withLeaderLock(opts: LeaderLockOptions): Promise { const { key, @@ -42,7 +49,7 @@ export async function withLeaderLock(opts: LeaderLockOptions): Promise { try { return await onLeader() } finally { @@ -57,6 +64,8 @@ export async function withLeaderLock(opts: LeaderLockOptions): Promise(opts: LeaderLockOptions): Promise { + it('runs at most the limit at once and starts waiting tasks in call order', async () => { + const limit = createConcurrencyLimit(2) + const gates = Array.from({ length: 5 }, () => createDeferred()) + const started: number[] = [] + let active = 0 + let peak = 0 + const runs = gates.map((gate, index) => + limit(async () => { + started.push(index) + peak = Math.max(peak, ++active) + await gate.promise + active-- + return index + }) + ) + await flushMacrotask() + expect(started).toEqual([0, 1]) + for (const index of [1, 0, 4, 3, 2]) { + gates[index].resolve() + await flushMacrotask() + } + await expect(Promise.all(runs)).resolves.toEqual([0, 1, 2, 3, 4]) + expect(started).toEqual([0, 1, 2, 3, 4]) + expect(peak).toBe(2) + }) + + it('releases the slot of a task that throws synchronously or rejects', async () => { + const limit = createConcurrencyLimit(1) + const error = new Error('task failed') + const throwsSynchronously = (): Promise => { + throw error + } + const thrown = limit(throwsSynchronously) + const rejected = limit(async () => Promise.reject(error)) + let next: string | undefined + void limit(async () => 'next').then((value) => { + next = value + }) + await expect(thrown).rejects.toBe(error) + await expect(rejected).rejects.toBe(error) + await flushMacrotask() + expect(next).toBe('next') + }) +}) diff --git a/apps/sim/lib/knowledge/search/diagnostics.ts b/apps/sim/lib/knowledge/search/diagnostics.ts index 6b603bd3abe..64ca7f4808a 100644 --- a/apps/sim/lib/knowledge/search/diagnostics.ts +++ b/apps/sim/lib/knowledge/search/diagnostics.ts @@ -148,6 +148,8 @@ const FAILURE_REASONS = [ 'mcp_payload_too_large', 'mcp_payload_unreadable', 'mcp_account_required', + /** The account never got to search: no slot freed with enough of the search deadline left. */ + 'not_started', ] as const export type SearchFailureReason = (typeof FAILURE_REASONS)[number] diff --git a/apps/sim/lib/oauth/credential-service.test.ts b/apps/sim/lib/oauth/credential-service.test.ts index 4e947758e1e..7043f128245 100644 --- a/apps/sim/lib/oauth/credential-service.test.ts +++ b/apps/sim/lib/oauth/credential-service.test.ts @@ -9,6 +9,7 @@ import { } from '@sim/testing' import { encryptionMock, encryptionMockFns } from '@sim/testing/mocks/encryption.mock' import { getMockLogger } from '@sim/testing/mocks/logger.mock' +import { redisConfigMockFns } from '@sim/testing/mocks/redis-config.mock' import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' const hoisted = vi.hoisted(() => ({ @@ -71,6 +72,7 @@ vi.mock('@/lib/oauth/refresh-error-codes', () => ({ isTerminalRefreshError: vi.fn(() => false), })) +import { CredentialRevokedError } from '@/lib/oauth/credential-revoked' import { getCredentialTerminalRefreshError, getOAuthToken, @@ -83,7 +85,10 @@ import { import { isInstagramProvider, shouldProactivelyRefreshInstagramToken } from '@/lib/oauth/instagram' import { isMicrosoftProvider } from '@/lib/oauth/microsoft' import { getOAuthRefreshCoordinationIdentity } from '@/lib/oauth/refresh-coordination' -import { isTerminalRefreshError } from '@/lib/oauth/refresh-error-codes' +import { + isCredentialRevocationError, + isTerminalRefreshError, +} from '@/lib/oauth/refresh-error-codes' import { fanOutSlackTokenChain } from '@/lib/oauth/slack' import { markCredentialDead } from '@/lib/oauth/terminal-errors' import { GOOGLE_SERVICE_ACCOUNT_PROVIDER_ID } from '@/lib/oauth/types' @@ -434,6 +439,7 @@ describe('OAuth access-token refresh headroom', () => { afterEach(() => { vi.useRealTimers() vi.mocked(isTerminalRefreshError).mockReturnValue(false) + vi.mocked(isCredentialRevocationError).mockReturnValue(false) vi.mocked(isInstagramProvider).mockReturnValue(false) vi.mocked(shouldProactivelyRefreshInstagramToken).mockReturnValue(false) vi.mocked(isMicrosoftProvider).mockReturnValue(false) @@ -595,8 +601,11 @@ describe('OAuth access-token refresh headroom', () => { expect(guard).toContain('original-refresh-token') }) + /** A Google token with a minute left, so its caller waits for the refresh instead of using it. */ + const WAITED_GOOGLE_REMAINING_MS = 60_000 + it('returns no token when the rotation write finds the account gone', async () => { - queueCredentialAccount(createOAuthAccount()) + queueCredentialAccount(createOAuthAccount(WAITED_GOOGLE_REMAINING_MS)) dbChainMockFns.returning.mockResolvedValueOnce([]) queueTableRows(account, [createOAuthAccount()]) queueTableRows(account, []) @@ -623,8 +632,17 @@ describe('OAuth access-token refresh headroom', () => { expect(markCredentialDead).not.toHaveBeenCalled() }) + it('returns no token for a credential flagged dead, however long its stored token lasts', async () => { + mocks.getRecentTerminalError.mockResolvedValue('invalid_client') + queueCredentialAccount(createOAuthAccount(WAITED_GOOGLE_REMAINING_MS)) + await expect( + resolveCredentialTokenBundle(RAW_CREDENTIAL_ID, RAW_USER_ID, 'test') + ).resolves.toBeNull() + expect(mocks.withLeaderLock).not.toHaveBeenCalled() + }) + it('flags a credential dead on a terminal failure when its chain did not move', async () => { - queueCredentialAccount(createOAuthAccount()) + queueCredentialAccount(createOAuthAccount(WAITED_GOOGLE_REMAINING_MS)) vi.mocked(isTerminalRefreshError).mockReturnValue(true) mocks.refreshOAuthToken.mockResolvedValue({ ok: false, errorCode: 'invalid_client' }) queueTableRows(account, [createOAuthAccount()]) @@ -783,6 +801,135 @@ describe('OAuth access-token refresh headroom', () => { await expect(concurrent).resolves.toEqual({ accessToken: 'refreshed-access-token' }) expect(mocks.refreshOAuthToken).toHaveBeenCalledOnce() }) + describe('when a refresh lease is held by a refresh no process is running', () => { + beforeEach(async () => { + const { withLeaderLock } = await vi.importActual< + typeof import('@/lib/concurrency/leader-lock') + >('@/lib/concurrency/leader-lock') + mocks.withLeaderLock.mockImplementation(withLeaderLock) + redisConfigMockFns.mockAcquireLock.mockResolvedValueOnce(false) + }) + + it.each([ + { label: 'a Google token under the refresh floor', providerId: 'google-drive' }, + { label: 'a provider that cannot refresh ahead', providerId: 'confluence' }, + ])('refreshes $label once the follower takes over the lapsed lease', async ({ providerId }) => { + queueCredentialAccount({ ...createOAuthAccount(240_000), providerId }) + const pending = resolveCredentialTokenBundle(RAW_CREDENTIAL_ID, RAW_USER_ID, 'test') + await vi.advanceTimersByTimeAsync(31_000) + await expect(pending).resolves.toEqual({ accessToken: 'refreshed-access-token' }) + expect(mocks.refreshOAuthToken).toHaveBeenCalledOnce() + }) + + it('never serves the stored token under the refresh floor when a live holder keeps the lease', async () => { + redisConfigMockFns.mockAcquireLock.mockResolvedValueOnce(false) + queueCredentialAccount(createOAuthAccount(240_000)) + const pending = resolveCredentialTokenBundle(RAW_CREDENTIAL_ID, RAW_USER_ID, 'test') + await vi.advanceTimersByTimeAsync(31_000) + await expect(pending).resolves.toBeNull() + expect(mocks.refreshOAuthToken).not.toHaveBeenCalled() + }) + + it('does not refresh again after taking over from a leader that flagged the credential dead', async () => { + mocks.getRecentTerminalError + .mockResolvedValueOnce(null) + .mockResolvedValueOnce('invalid_client') + queueCredentialAccount(createOAuthAccount(240_000)) + const pending = resolveCredentialTokenBundle(RAW_CREDENTIAL_ID, RAW_USER_ID, 'test') + await vi.advanceTimersByTimeAsync(31_000) + await expect(pending).resolves.toBeNull() + expect(mocks.refreshOAuthToken).not.toHaveBeenCalled() + }) + + it('lets the follower take over before the in-process refresh times out', async () => { + const { coalesceLocally } = await vi.importActual< + typeof import('@/lib/concurrency/singleflight') + >('@/lib/concurrency/singleflight') + mocks.coalesceLocally.mockImplementation(coalesceLocally) + queueCredentialAccount(createOAuthAccount(240_000)) + const pending = resolveCredentialTokenBundle(RAW_CREDENTIAL_ID, RAW_USER_ID, 'test') + await vi.advanceTimersByTimeAsync(31_000) + await expect(pending).resolves.toEqual({ accessToken: 'refreshed-access-token' }) + }) + }) + + it('fails callers of a refresh hung in this process once it times out', async () => { + const { coalesceLocally } = await vi.importActual< + typeof import('@/lib/concurrency/singleflight') + >('@/lib/concurrency/singleflight') + mocks.coalesceLocally.mockImplementation(coalesceLocally) + mocks.refreshOAuthToken.mockReturnValue(new Promise(() => {})) + queueCredentialAccount(createOAuthAccount(100_000)) + let settled = false + const pending = resolveCredentialTokenBundle(RAW_CREDENTIAL_ID, RAW_USER_ID, 'test').finally( + () => { + settled = true + } + ) + await vi.advanceTimersByTimeAsync(74_000) + expect(settled).toBe(false) + await vi.advanceTimersByTimeAsync(2_000) + await expect(pending).resolves.toBeNull() + }) + + it('fails a Google chain the provider revokes even while its access token is valid', async () => { + vi.mocked(isTerminalRefreshError).mockReturnValue(true) + vi.mocked(isCredentialRevocationError).mockReturnValue(true) + mocks.refreshOAuthToken.mockResolvedValue({ ok: false, errorCode: 'invalid_grant' }) + const row = createOAuthAccount(100_000) + queueTableRows(account, [row]) + queueTableRows(account, [row]) + await expect(refreshTokenIfNeeded('test', row, RAW_ACCOUNT_ID)).rejects.toBeInstanceOf( + CredentialRevokedError + ) + expect(dbChainMockFns.set).toHaveBeenCalledWith( + expect.objectContaining({ refreshRevokedCode: 'invalid_grant' }) + ) + }) + + it('answers a Microsoft caller from its valid token while refresh-token aging refreshes', async () => { + vi.mocked(isMicrosoftProvider).mockReturnValue(true) + const provider = createDeferred() + mocks.refreshOAuthToken.mockReturnValue(provider.promise) + const row = { + ...createOAuthAccount(3_600_000), + providerId: 'microsoft-teams', + refreshTokenExpiresAt: new Date(now.getTime() + 24 * 60 * 60_000), + } + await expect(refreshTokenIfNeeded('test', row, RAW_ACCOUNT_ID)).resolves.toEqual({ + accessToken: 'original-access-token', + refreshed: false, + }) + 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' }) + ) + ) + }) + + it('spreads refresh-ahead start across accounts expiring at the same moment', async () => { + const row = createOAuthAccount(570_000) + /** Chosen because their hashes fall on opposite sides of the 30s jitter cut, so a hash-rule change fails loudly. */ + await expect(refreshTokenIfNeeded('test', row, 'account-2')).resolves.toEqual({ + accessToken: 'original-access-token', + refreshed: false, + }) + await vi.waitFor(() => expect(mocks.refreshOAuthToken).toHaveBeenCalledOnce()) + await expect(refreshTokenIfNeeded('test', row, 'account-c')).resolves.toEqual({ + accessToken: 'original-access-token', + refreshed: false, + }) + await flushMicrotasks(20) + expect(mocks.refreshOAuthToken).toHaveBeenCalledOnce() + }) }) describe('Google service-account token minting', () => { @@ -796,20 +943,25 @@ describe('Google service-account token minting', () => { } const driveScope = 'https://www.googleapis.com/auth/drive.readonly' const now = new Date('2026-09-09T18:00:00.000Z') + const decryptedKey = { + decrypted: JSON.stringify({ + client_email: 'crawler@qa-project.iam.gserviceaccount.com', + private_key: privateKey.export({ type: 'pkcs8', format: 'pem' }).toString(), + token_uri: 'https://oauth2.googleapis.com/token', + }), + } - beforeEach(() => { + beforeEach(async () => { + const { coalesceLocally } = await vi.importActual< + typeof import('@/lib/concurrency/singleflight') + >('@/lib/concurrency/singleflight') + mocks.coalesceLocally.mockImplementation(coalesceLocally) resetDbChainMock() vi.useFakeTimers() vi.setSystemTime(now) vi.stubGlobal('fetch', fetchMock) fetchMock.mockImplementation(async () => Response.json({ access_token: 'google-access-token' })) - mocks.decryptSecret.mockResolvedValue({ - decrypted: JSON.stringify({ - client_email: 'crawler@qa-project.iam.gserviceaccount.com', - private_key: privateKey.export({ type: 'pkcs8', format: 'pem' }).toString(), - token_uri: 'https://oauth2.googleapis.com/token', - }), - }) + mocks.decryptSecret.mockResolvedValue(decryptedKey) }) afterEach(() => { @@ -896,6 +1048,37 @@ describe('Google service-account token minting', () => { expect(fetchMock).toHaveBeenCalledTimes(1) }) + it('shares one key read across concurrent mints', async () => { + queueTableRows(credential, [row]) + await expect( + Promise.all( + Array.from({ length: 6 }, (_, index) => + getServiceAccountToken('credential-1', [driveScope], `member-${index}@example.com`) + ) + ) + ).resolves.toEqual(Array(6).fill('google-access-token')) + expect(dbChainMockFns.select).toHaveBeenCalledTimes(1) + expect(fetchMock).toHaveBeenCalledTimes(6) + }) + + it('stops a cancelled mint waiting on a shared key read while the others finish', async () => { + queueTableRows(credential, [row]) + const decrypted = createDeferred<{ decrypted: string }>() + mocks.decryptSecret.mockReturnValue(decrypted.promise) + const controller = new AbortController() + const reason = new DOMException('Caller cancelled', 'AbortError') + const concurrent = getServiceAccountToken('credential-1', [driveScope], 'a@example.com') + const cancelled = getServiceAccountToken('credential-1', [driveScope], 'b@example.com', { + signal: controller.signal, + }) + await vi.waitFor(() => expect(mocks.decryptSecret).toHaveBeenCalledOnce()) + controller.abort(reason) + await expect(cancelled).rejects.toBe(reason) + decrypted.resolve(decryptedKey) + await expect(concurrent).resolves.toBe('google-access-token') + expect(fetchMock).toHaveBeenCalledOnce() + }) + it('preserves cancellation through the service-account resolver registry', async () => { const controller = new AbortController() const reason = new DOMException('Caller cancelled', 'AbortError') diff --git a/apps/sim/lib/oauth/credential-service.ts b/apps/sim/lib/oauth/credential-service.ts index 2086185b60e..734c61efe60 100644 --- a/apps/sim/lib/oauth/credential-service.ts +++ b/apps/sim/lib/oauth/credential-service.ts @@ -225,6 +225,38 @@ const SA_DIAGNOSTIC_ERROR_CODES = new Set([ 'unauthorized_client', ]) +/** + * The decrypted key of a live Google service-account credential. Reads running at once (one per + * delegated subject) share one query and decrypt; a later read queries again, so it sees a + * revocation recorded since. Sharing also bounds a hung read by the coalescing settle timeout. + */ +function readGoogleServiceAccountKey(credentialId: string): Promise { + return coalesceLocally(`sa-key:${credentialId}`, async () => { + const [credentialRow] = await db + .select({ + type: credential.type, + providerId: credential.providerId, + revokedAt: credential.revokedAt, + encryptedServiceAccountKey: credential.encryptedServiceAccountKey, + }) + .from(credential) + .where(eq(credential.id, credentialId)) + .limit(1) + + if ( + credentialRow?.type !== 'service_account' || + credentialRow.providerId !== GOOGLE_SERVICE_ACCOUNT_PROVIDER_ID || + credentialRow.revokedAt !== null || + !credentialRow.encryptedServiceAccountKey + ) { + throw new Error('Google service account credential is unavailable') + } + + const { decrypted } = await decryptSecret(credentialRow.encryptedServiceAccountKey) + return decrypted + }) +} + /** * Generates a short-lived access token for a Google service account credential * using the two-legged OAuth JWT flow (RFC 7523). @@ -241,27 +273,7 @@ export async function getServiceAccountToken( options?: CredentialTokenResolutionOptions ): Promise { options?.signal?.throwIfAborted() - const [credentialRow] = await db - .select({ - type: credential.type, - providerId: credential.providerId, - revokedAt: credential.revokedAt, - encryptedServiceAccountKey: credential.encryptedServiceAccountKey, - }) - .from(credential) - .where(eq(credential.id, credentialId)) - .limit(1) - - if ( - credentialRow?.type !== 'service_account' || - credentialRow.providerId !== GOOGLE_SERVICE_ACCOUNT_PROVIDER_ID || - credentialRow.revokedAt !== null || - !credentialRow.encryptedServiceAccountKey - ) { - throw new Error('Google service account credential is unavailable') - } - - const { decrypted } = await decryptSecret(credentialRow.encryptedServiceAccountKey) + const decrypted = await waitWithAbort(readGoogleServiceAccountKey(credentialId), options?.signal) const keyData = JSON.parse(decrypted) as { client_email: string private_key: string @@ -1018,10 +1030,9 @@ const OAUTH_REFRESH_AHEAD_WINDOW_MS = 10 * 60 * 1000 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. + * A Google refresh does not rotate the refresh token, and Microsoft does not revoke the old one + * when issuing a new one (learn.microsoft.com, "Refresh tokens in the Microsoft identity + * platform"); serving the stored access token meanwhile relies on it staying valid until expiry. */ function refreshesAheadOfExpiry(providerId: string): boolean { return ( @@ -1049,7 +1060,8 @@ interface StoredOAuthTokens extends RevocableChain { * - `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). + * stored token if it fails (Instagram long-lived renewal; Microsoft refresh-token aging + * refreshes ahead instead). * - `must-refresh`: wait for the refresh; the stored token cannot serve the request. */ type AccessTokenRefreshPlan = @@ -1116,6 +1128,14 @@ function planAccessTokenRefresh( const REFRESH_LOCK_HEADROOM_MS = 15_000 const REFRESH_LOCK_TTL_SEC = Math.ceil((TOKEN_REFRESH_TIMEOUT_MS + REFRESH_LOCK_HEADROOM_MS) / 1000) const REFRESH_FOLLOWER_MAX_WAIT_MS = REFRESH_LOCK_TTL_SEC * 1000 +/** + * How long callers sharing one in-process refresh wait for it: a follower's full wait, then the + * leader run it may take over when that wait ends with a lapsed lease (budgeted like any lease), + * plus headroom for the lock round trips around them; so a follower concludes on its own, and a + * leader hung in this process still fails its callers instead of wedging them. + */ +const REFRESH_SETTLE_TIMEOUT_MS = + REFRESH_FOLLOWER_MAX_WAIT_MS + REFRESH_LOCK_TTL_SEC * 1000 + REFRESH_LOCK_HEADROOM_MS /** * The raw scope one refresh coordinates on: the account row, or the installation for a Slack @@ -1301,260 +1321,268 @@ async function performCoalescedRefresh({ const lockKey = `oauth:refresh:${scopeKey}` - const refreshPromise = coalesceLocally(lockKey, () => - withLeaderLock({ - key: lockKey, - ttlSec: REFRESH_LOCK_TTL_SEC, - maxWaitMs: REFRESH_FOLLOWER_MAX_WAIT_MS, - onLeader: async () => { - try { - /** - * Reread under the lock: another process may have recorded a revocation after the - * caller read the row, and the version guards the record against a reconnect that - * lands while the provider call is in flight. - */ - const current = await readStoredChain(accountId) - const recordedMeanwhile = recordedRevocationCode(current) - 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 - if (slackTeamId) { - const freshest = await getFreshestSlackChain(slackTeamId) - if (!freshest) { - throw new Error( - `No refresh-capable account row found for Slack installation ${slackTeamId}` - ) + const refreshPromise = coalesceLocally( + lockKey, + () => + withLeaderLock({ + key: lockKey, + ttlSec: REFRESH_LOCK_TTL_SEC, + maxWaitMs: REFRESH_FOLLOWER_MAX_WAIT_MS, + onLeader: async () => { + try { + /** + * Reread under the lock: another process may have recorded a revocation or flagged the + * credential dead after the caller read the row (a follower taking over the lock after + * a failed leader released it sees only these), and the version guards the record + * against a reconnect that lands while the provider call is in flight. + */ + const current = await readStoredChain(accountId) + const recordedMeanwhile = recordedRevocationCode(current) + if (recordedMeanwhile && !isRevocationReprobeDue(current)) { + return revokedOutcome(recordedMeanwhile) } - slackChainVersion = freshest.chainVersion - chainVersion = freshest.chainVersion - if ( - freshest.accessToken && - freshest.accessTokenExpiresAt && - !isOAuthAccessTokenExpiring(freshest.accessTokenExpiresAt, providerId) - ) { - await fanOutSlackTokenChain( - slackTeamId, - { - accessToken: freshest.accessToken, - refreshToken: freshest.refreshToken, - accessTokenExpiresAt: freshest.accessTokenExpiresAt, - }, - { ifChainUnchangedSince: freshest.chainVersion } + if (await getRecentTerminalError(scopeKey)) return null + /** + * 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 ) - logger.info('Reused freshest Slack installation token', logContext) - return tokenOutcome(freshest.accessToken) - } - refreshTokenToUse = freshest.refreshToken - } - - let quickBooksClientConfig - if (providerId === 'quickbooks') { - if (!oauthConfig) { - throw new Error('QuickBooks OAuth client configuration is missing') - } - quickBooksClientConfig = await decryptQuickBooksOAuthClientConfig(oauthConfig) - } - const result = quickBooksClientConfig - ? await refreshOAuthToken(providerId, refreshTokenToUse, quickBooksClientConfig) - : await refreshOAuthToken(providerId, refreshTokenToUse) - - if (!result.ok) { - const { errorCode } = result - const revoked = isCredentialRevocationError(errorCode, providerId) - if (!revoked) { - logger.error('Failed to refresh token', { - ...logContext, - errorCode, - message: result.message, - }) + ? usableStoredToken(current, providerId) + : null + if (storedSinceRead) { + logger.info('Using the token stored since this caller read the account', logContext) + return tokenOutcome(storedSinceRead) } - const chainScope = { accountId, slackTeamId } - if (errorCode && isTerminalRefreshError(errorCode, providerId)) { - // A refresh that lost a race with a concurrent connect or a newer - // rotation fails with a revoked/rotated-out token even though the - // account just got a live chain — recording it then would take - // down a healthy credential. + let chainVersion = current?.updatedAt + let refreshTokenToUse = refreshToken + let slackChainVersion: Date | null = null + if (slackTeamId) { + const freshest = await getFreshestSlackChain(slackTeamId) + if (!freshest) { + throw new Error( + `No refresh-capable account row found for Slack installation ${slackTeamId}` + ) + } + slackChainVersion = freshest.chainVersion + chainVersion = freshest.chainVersion if ( - slackChainVersion && - (await hasSlackChainMoved(slackTeamId!, slackChainVersion)) + freshest.accessToken && + freshest.accessTokenExpiresAt && + !isOAuthAccessTokenExpiring(freshest.accessTokenExpiresAt, providerId) ) { - logger.info('Skipping dead flag: Slack chain moved during refresh', logContext) - return null + await fanOutSlackTokenChain( + slackTeamId, + { + accessToken: freshest.accessToken, + refreshToken: freshest.refreshToken, + accessTokenExpiresAt: freshest.accessTokenExpiresAt, + }, + { ifChainUnchangedSince: freshest.chainVersion } + ) + logger.info('Reused freshest Slack installation token', logContext) + return tokenOutcome(freshest.accessToken) } - if (!slackTeamId) { - const stored = await readStoredChain(accountId) - if (stored && stored.refreshToken !== refreshToken) { - logger.info('Skipping dead flag: chain moved during refresh', logContext) - return tokenOutcome(usableStoredToken(stored, providerId)) - } + refreshTokenToUse = freshest.refreshToken + } + + let quickBooksClientConfig + if (providerId === 'quickbooks') { + if (!oauthConfig) { + throw new Error('QuickBooks OAuth client configuration is missing') } + quickBooksClientConfig = await decryptQuickBooksOAuthClientConfig(oauthConfig) + } + const result = quickBooksClientConfig + ? await refreshOAuthToken(providerId, refreshTokenToUse, quickBooksClientConfig) + : await refreshOAuthToken(providerId, refreshTokenToUse) + + if (!result.ok) { + const { errorCode } = result + const revoked = isCredentialRevocationError(errorCode, providerId) if (!revoked) { - await markCredentialDead(scopeKey, errorCode) - return null + logger.error('Failed to refresh token', { + ...logContext, + errorCode, + message: result.message, + }) } - if ( - !(await recordRefreshRevocation( + const chainScope = { accountId, slackTeamId } + if (errorCode && isTerminalRefreshError(errorCode, providerId)) { + // A refresh that lost a race with a concurrent connect or a newer + // rotation fails with a revoked/rotated-out token even though the + // account just got a live chain — recording it then would take + // down a healthy credential. + if ( + slackChainVersion && + (await hasSlackChainMoved(slackTeamId!, slackChainVersion)) + ) { + logger.info('Skipping dead flag: Slack chain moved during refresh', logContext) + return null + } + if (!slackTeamId) { + const stored = await readStoredChain(accountId) + if (stored && stored.refreshToken !== refreshToken) { + logger.info('Skipping dead flag: chain moved during refresh', logContext) + return tokenOutcome(usableStoredToken(stored, providerId)) + } + } + if (!revoked) { + await markCredentialDead(scopeKey, errorCode) + return null + } + if ( + !(await recordRefreshRevocation( + chainScope, + refreshTokenToUse, + errorCode, + chainVersion + )) + ) { + logger.info('Skipping revocation record: chain moved during refresh', logContext) + return null + } + // Logged once per rejection; later runs short-circuited by the record stay silent. + logger.warn('Provider revoked the refresh token; recorded until reconnect', { + ...logContext, + errorCode, + }) + return revokedOutcome(errorCode) + } + // A re-probe that fails for a passing reason keeps the revocation, unless the chain moved. + const stillRevoked = + recordedCode !== null && + (await recordRefreshRevocation( chainScope, refreshTokenToUse, - errorCode, + recordedCode, chainVersion )) - ) { - logger.info('Skipping revocation record: chain moved during refresh', logContext) - return null - } - // Logged once per rejection; later runs short-circuited by the record stay silent. - logger.warn('Provider revoked the refresh token; recorded until reconnect', { - ...logContext, - errorCode, - }) - return revokedOutcome(errorCode) + return stillRevoked ? revokedOutcome(recordedCode) : null } - // A re-probe that fails for a passing reason keeps the revocation, unless the chain moved. - const stillRevoked = - recordedCode !== null && - (await recordRefreshRevocation( - chainScope, - refreshTokenToUse, - recordedCode, - chainVersion - )) - return stillRevoked ? revokedOutcome(recordedCode) : null - } - const accessTokenExpiresAt = new Date(Date.now() + result.expiresIn * 1000) + const accessTokenExpiresAt = new Date(Date.now() + result.expiresIn * 1000) - if (slackTeamId) { - await fanOutSlackTokenChain( - slackTeamId, - { + if (slackTeamId) { + await fanOutSlackTokenChain( + slackTeamId, + { + accessToken: result.accessToken, + refreshToken: result.refreshToken || refreshTokenToUse, + accessTokenExpiresAt, + }, + { ifChainUnchangedSince: slackChainVersion ?? undefined, freshlyIssued: true } + ) + } else { + const updateData: Record = { accessToken: result.accessToken, - refreshToken: result.refreshToken || refreshTokenToUse, accessTokenExpiresAt, - }, - { ifChainUnchangedSince: slackChainVersion ?? undefined, freshlyIssued: true } - ) - } else { - const updateData: Record = { - accessToken: result.accessToken, - accessTokenExpiresAt, - ...CLEARED_REFRESH_REVOCATION, - updatedAt: new Date(), - } - if (result.refreshToken && result.refreshToken !== refreshToken) { - updateData.refreshToken = result.refreshToken - } - if (isMicrosoftProvider(providerId)) { - updateData.refreshTokenExpiresAt = getMicrosoftRefreshTokenExpiry() - } - if (result.refreshTokenExpiresIn) { - updateData.refreshTokenExpiresAt = new Date( - Date.now() + result.refreshTokenExpiresIn * 1000 - ) - } + ...CLEARED_REFRESH_REVOCATION, + updatedAt: new Date(), + } + if (result.refreshToken && result.refreshToken !== refreshToken) { + updateData.refreshToken = result.refreshToken + } + if (isMicrosoftProvider(providerId)) { + updateData.refreshTokenExpiresAt = getMicrosoftRefreshTokenExpiry() + } + if (result.refreshTokenExpiresIn) { + updateData.refreshTokenExpiresAt = new Date( + Date.now() + result.refreshTokenExpiresIn * 1000 + ) + } - /** - * The chain is rotated only from the refresh token this refresh started from. - * A lease is not mutual exclusion: it can expire under a slow provider or a - * paused process while the leader is still running, and an unconditional - * write would then let this refresh overwrite a newer rotation with a chain - * the provider has already retired, which the next refresh pays for as - * `invalid_grant` and, under reuse detection, as a revoked grant. When no row - * matches, another writer rotated first: its chain is the live one, so this - * caller uses what is stored and never retries the provider. - */ - const rotated = await db - .update(account) - .set(updateData) - .where( - and(eq(account.id, accountId), sealedTextEquals(account.refreshToken, refreshToken)) - ) - .returning({ id: account.id }) - if (rotated.length === 0) { - const stored = await readStoredChain(accountId) - if (!stored) { - logger.warn('Rotation write found no account; the credential is gone', logContext) - return null + /** + * The chain is rotated only from the refresh token this refresh started from. + * A lease is not mutual exclusion: it can expire under a slow provider or a + * paused process while the leader is still running, and an unconditional + * write would then let this refresh overwrite a newer rotation with a chain + * the provider has already retired, which the next refresh pays for as + * `invalid_grant` and, under reuse detection, as a revoked grant. When no row + * matches, another writer rotated first: its chain is the live one, so this + * caller uses what is stored and never retries the provider. + */ + const rotated = await db + .update(account) + .set(updateData) + .where( + and( + eq(account.id, accountId), + sealedTextEquals(account.refreshToken, refreshToken) + ) + ) + .returning({ id: account.id }) + if (rotated.length === 0) { + const stored = await readStoredChain(accountId) + if (!stored) { + logger.warn('Rotation write found no account; the credential is gone', logContext) + return null + } + logger.warn( + 'Rotation write lost to a newer chain; using the stored token', + logContext + ) + return tokenOutcome(usableStoredToken(stored, providerId)) } - logger.warn( - 'Rotation write lost to a newer chain; using the stored token', - logContext - ) - return tokenOutcome(usableStoredToken(stored, providerId)) } - } - logger.info('Successfully refreshed access token', logContext) - return { kind: 'token', accessToken: result.accessToken } - } catch (error) { - logger.error('Refresh failed inside leader path', { - ...logContext, - ...(privacyMode === 'selector' ? {} : { error: toError(error).message }), - }) - return revokedOutcome(recordedCode) - } - }, - onFollower: async () => { - try { - const [row] = await db - .select({ - accessToken: account.accessToken, - accessTokenExpiresAt: account.accessTokenExpiresAt, - refreshToken: account.refreshToken, - ...refreshRevocationColumns, + logger.info('Successfully refreshed access token', logContext) + return { kind: 'token', accessToken: result.accessToken } + } catch (error) { + logger.error('Refresh failed inside leader path', { + ...logContext, + ...(privacyMode === 'selector' ? {} : { error: toError(error).message }), }) - .from(account) - .where(eq(account.id, accountId)) - .limit(1) - if ( - row?.accessToken && - row.accessTokenExpiresAt && - !isOAuthAccessTokenExpiring(row.accessTokenExpiresAt, providerId) - ) { - logger.info('Got fresh access token from coalesced refresh', logContext) - return { kind: 'token', accessToken: row.accessToken } + return revokedOutcome(recordedCode) } - const revokedByLeader = recordedRevocationCode(row) - return revokedByLeader && !isRevocationReprobeDue(row) - ? revokedOutcome(revokedByLeader) - : null - } catch (error) { - logger.warn('Follower DB read failed during refresh poll', { - ...logContext, - ...(privacyMode === 'selector' ? {} : { error: toError(error).message }), - }) - return null - } - }, - }) + }, + onFollower: async () => { + try { + const [row] = await db + .select({ + accessToken: account.accessToken, + accessTokenExpiresAt: account.accessTokenExpiresAt, + refreshToken: account.refreshToken, + ...refreshRevocationColumns, + }) + .from(account) + .where(eq(account.id, accountId)) + .limit(1) + if ( + row?.accessToken && + row.accessTokenExpiresAt && + !isOAuthAccessTokenExpiring(row.accessTokenExpiresAt, providerId) + ) { + logger.info('Got fresh access token from coalesced refresh', logContext) + return { kind: 'token', accessToken: row.accessToken } + } + const revokedByLeader = recordedRevocationCode(row) + return revokedByLeader && !isRevocationReprobeDue(row) + ? revokedOutcome(revokedByLeader) + : null + } catch (error) { + logger.warn('Follower DB read failed during refresh poll', { + ...logContext, + ...(privacyMode === 'selector' ? {} : { error: toError(error).message }), + }) + return null + } + }, + }), + REFRESH_SETTLE_TIMEOUT_MS ) try { diff --git a/apps/sim/lib/sim-search/live/application.test.ts b/apps/sim/lib/sim-search/live/application.test.ts index 466ccd94d4c..b69d53ff555 100644 --- a/apps/sim/lib/sim-search/live/application.test.ts +++ b/apps/sim/lib/sim-search/live/application.test.ts @@ -10,6 +10,7 @@ import { knowledgeContextsMock, knowledgeContextsMockFns, } from '@sim/testing/mocks/knowledge-contexts.mock' +import { getMockLogger } from '@sim/testing/mocks/logger.mock' import { workspaceAuthzMock, workspaceAuthzMockFns } from '@sim/testing/mocks/workspace-authz.mock' import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' @@ -18,6 +19,7 @@ const mocks = vi.hoisted(() => ({ resolveAccount: vi.fn(), search: vi.fn(), mcpCall: vi.fn(), + mcpClose: vi.fn(async () => {}), read: vi.fn(), admin: vi.fn(), adminSearch: vi.fn(), @@ -47,7 +49,7 @@ vi.mock('@/lib/sim-search/live/policy-store', () => ({ livePolicyFor: vi.fn(() => defaultLiveSearchPolicy()), })) vi.mock('@/lib/sim-search/live/managed-mcp', () => ({ - createManagedSearchMcpClient: async () => ({ call: mocks.mcpCall, close: async () => {} }), + createManagedSearchMcpClient: async () => ({ call: mocks.mcpCall, close: mocks.mcpClose }), })) vi.mock('@/lib/sim-search/live/coda-mcp', () => ({ searchCodaMcp: vi.fn(), @@ -74,6 +76,7 @@ vi.mock('@/executor/utils/resolved-secret-content-projection', () => ({ projectResolvedSecretModelContent: (value: string) => ({ safe: true, value }), })) +import { withSearchDiagnostics } from '@/lib/knowledge/search/diagnostics' import { decodeLiveReference, encodeLiveReference, @@ -96,6 +99,7 @@ inputValidationMockFns.mockCreatePinnedConnectionPool.mockImplementation(() => ( destroy: mocks.destroyPool, })) +const diagnosticsLogger = getMockLogger('KnowledgeSearchDiagnostics') const principal = createSessionPrincipal({ userId: 'reader', sessionId: 's' }) const input = { workspaceId: 'workspace', query: 'launch', topK: 20 } const account = { @@ -1309,7 +1313,9 @@ describe('authorized live retrieval', () => { search.signal?.throwIfAborted() return { documents: [document] } }) - const pending = searchLiveKnowledge.execute({ principal, input }) + const pending = withSearchDiagnostics({ surface: 'copilot' }, () => + searchLiveKnowledge.execute({ principal, input }) + ) await vi.advanceTimersByTimeAsync(20_000) const statuses = (await pending).live?.accounts ?? [] /** The second wave waits 9.5s for a slot, then still has its full account deadline. */ @@ -1322,6 +1328,45 @@ describe('authorized live retrieval', () => { message: expect.stringContaining('Not searched before the search deadline'), }) expect(mocks.resolveAccount).toHaveBeenCalledTimes(12) + const skipped = diagnosticsLogger.info.mock.calls.find( + ([message, fields]) => + message === 'Live search provider completed' && fields?.accountIndex === 12 + ) + expect(skipped?.[1]?.results).toEqual([ + expect.objectContaining({ status: 'timeout', failure: { reason: 'not_started' } }), + ]) + }) + + it('frees the slot of a session that never opens, so an account waiting on it still searches', async () => { + mocks.accounts.mockResolvedValue(accounts(7)) + for (let index = 0; index < 6; index++) + mocks.service.mockReturnValueOnce(new Promise(() => {})) + let result: Awaited> | undefined + void searchLiveKnowledge.execute({ principal, input }).then((value) => { + result = value + }) + await vi.advanceTimersByTimeAsync(20_000) + expect(result?.live?.accounts.map(({ status }) => status)).toEqual([ + ...Array.from({ length: 6 }, () => 'timeout'), + 'ok', + ]) + }) + + it('closes a session that opens after its account stopped waiting for it', async () => { + mocks.accounts.mockResolvedValue([ + { ...account, provider: 'coda', providerId: 'mcp:coda', type: 'managed_mcp' }, + ]) + const opening = createDeferred() + mocks.service.mockReturnValueOnce(opening.promise) + let result: Awaited> | undefined + void searchLiveKnowledge.execute({ principal, input }).then((value) => { + result = value + }) + await vi.advanceTimersByTimeAsync(12_000) + expect(result?.live?.accounts[0]?.status).toBe('timeout') + expect(mocks.mcpClose).not.toHaveBeenCalled() + opening.resolve(undefined) + await vi.waitFor(() => expect(mocks.mcpClose).toHaveBeenCalledOnce()) }) }) }) diff --git a/apps/sim/lib/sim-search/live/application.ts b/apps/sim/lib/sim-search/live/application.ts index 07595b70402..78e9817b53f 100644 --- a/apps/sim/lib/sim-search/live/application.ts +++ b/apps/sim/lib/sim-search/live/application.ts @@ -285,9 +285,11 @@ function resultFor( /** * Accounts resolving their credential and opening their session at once. That phase reads the * database (credentials, source settings, provider rollout) along with the network calls those - * reads lead to, such as a token refresh or a Google Directory lookup and delegated token mint, - * so it is bounded well under the web database pool. Provider searches are not bounded: a search - * waits for every account, and a second wave would add a full provider round trip. + * reads lead to, such as a token refresh or a Google Directory lookup and delegated token mints, + * so it is bounded well under the web database pool. The bound is per account, not per + * connection: Google delegation resolves its subjects concurrently, and concurrent subjects share + * one service-account key read. Provider searches are not bounded: a search waits + * for every account, and a second wave would add a full provider round trip. */ const ACCOUNT_OPEN_CONCURRENCY = 6 /** Milliseconds a whole search may take. */ @@ -809,7 +811,10 @@ const liveKnowledgeSearch = defineAuthorizedKnowledgeUseCase({ failed( new NativeSearchError( 'timeout', - 'Not searched before the search deadline. Search this account alone or narrow the query.' + 'Not searched before the search deadline. Search this account alone or narrow the query.', + undefined, + undefined, + 'not_started' ), target )