Skip to content

Commit edd0621

Browse files
committed
fix(search): keep a valid token when a refresh lease is orphaned, share overlapping service-account reads, and label unsearched accounts
1 parent 04b66d4 commit edd0621

6 files changed

Lines changed: 285 additions & 36 deletions

File tree

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,50 @@
1+
import { flushMacrotask } from '@sim/testing/helpers/async'
2+
import { createDeferred } from '@sim/testing/helpers/deferred'
3+
import { describe, expect, it } from 'vitest'
4+
import { createConcurrencyLimit } from '@/lib/core/utils/concurrency'
5+
6+
describe('createConcurrencyLimit', () => {
7+
it('runs at most the limit at once and starts waiting tasks in call order', async () => {
8+
const limit = createConcurrencyLimit(2)
9+
const gates = Array.from({ length: 5 }, () => createDeferred<void>())
10+
const started: number[] = []
11+
let active = 0
12+
let peak = 0
13+
const runs = gates.map((gate, index) =>
14+
limit(async () => {
15+
started.push(index)
16+
peak = Math.max(peak, ++active)
17+
await gate.promise
18+
active--
19+
return index
20+
})
21+
)
22+
await flushMacrotask()
23+
expect(started).toEqual([0, 1])
24+
for (const index of [1, 0, 4, 3, 2]) {
25+
gates[index].resolve()
26+
await flushMacrotask()
27+
}
28+
await expect(Promise.all(runs)).resolves.toEqual([0, 1, 2, 3, 4])
29+
expect(started).toEqual([0, 1, 2, 3, 4])
30+
expect(peak).toBe(2)
31+
})
32+
33+
it('releases the slot of a task that throws synchronously or rejects', async () => {
34+
const limit = createConcurrencyLimit(1)
35+
const error = new Error('task failed')
36+
const throwsSynchronously = (): Promise<never> => {
37+
throw error
38+
}
39+
const thrown = limit(throwsSynchronously)
40+
const rejected = limit(async () => Promise.reject(error))
41+
let next: string | undefined
42+
void limit(async () => 'next').then((value) => {
43+
next = value
44+
})
45+
await expect(thrown).rejects.toBe(error)
46+
await expect(rejected).rejects.toBe(error)
47+
await flushMacrotask()
48+
expect(next).toBe('next')
49+
})
50+
})

‎apps/sim/lib/knowledge/search/diagnostics.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -148,6 +148,8 @@ const FAILURE_REASONS = [
148148
'mcp_payload_too_large',
149149
'mcp_payload_unreadable',
150150
'mcp_account_required',
151+
/** The account never got to search: no slot freed with enough of the search deadline left. */
152+
'not_started',
151153
] as const
152154
export type SearchFailureReason = (typeof FAILURE_REASONS)[number]
153155

‎apps/sim/lib/oauth/credential-service.test.ts‎

Lines changed: 117 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ import {
99
} from '@sim/testing'
1010
import { encryptionMock, encryptionMockFns } from '@sim/testing/mocks/encryption.mock'
1111
import { getMockLogger } from '@sim/testing/mocks/logger.mock'
12+
import { redisConfigMockFns } from '@sim/testing/mocks/redis-config.mock'
1213
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
1314

1415
const hoisted = vi.hoisted(() => ({
@@ -71,6 +72,7 @@ vi.mock('@/lib/oauth/refresh-error-codes', () => ({
7172
isTerminalRefreshError: vi.fn(() => false),
7273
}))
7374

75+
import { CredentialRevokedError } from '@/lib/oauth/credential-revoked'
7476
import {
7577
getCredentialTerminalRefreshError,
7678
getOAuthToken,
@@ -83,7 +85,10 @@ import {
8385
import { isInstagramProvider, shouldProactivelyRefreshInstagramToken } from '@/lib/oauth/instagram'
8486
import { isMicrosoftProvider } from '@/lib/oauth/microsoft'
8587
import { getOAuthRefreshCoordinationIdentity } from '@/lib/oauth/refresh-coordination'
86-
import { isTerminalRefreshError } from '@/lib/oauth/refresh-error-codes'
88+
import {
89+
isCredentialRevocationError,
90+
isTerminalRefreshError,
91+
} from '@/lib/oauth/refresh-error-codes'
8792
import { fanOutSlackTokenChain } from '@/lib/oauth/slack'
8893
import { markCredentialDead } from '@/lib/oauth/terminal-errors'
8994
import { GOOGLE_SERVICE_ACCOUNT_PROVIDER_ID } from '@/lib/oauth/types'
@@ -434,6 +439,7 @@ describe('OAuth access-token refresh headroom', () => {
434439
afterEach(() => {
435440
vi.useRealTimers()
436441
vi.mocked(isTerminalRefreshError).mockReturnValue(false)
442+
vi.mocked(isCredentialRevocationError).mockReturnValue(false)
437443
vi.mocked(isInstagramProvider).mockReturnValue(false)
438444
vi.mocked(shouldProactivelyRefreshInstagramToken).mockReturnValue(false)
439445
vi.mocked(isMicrosoftProvider).mockReturnValue(false)
@@ -783,6 +789,98 @@ describe('OAuth access-token refresh headroom', () => {
783789
await expect(concurrent).resolves.toEqual({ accessToken: 'refreshed-access-token' })
784790
expect(mocks.refreshOAuthToken).toHaveBeenCalledOnce()
785791
})
792+
describe('when a refresh lease is held by a refresh no process is running', () => {
793+
beforeEach(async () => {
794+
const { withLeaderLock } = await vi.importActual<
795+
typeof import('@/lib/concurrency/leader-lock')
796+
>('@/lib/concurrency/leader-lock')
797+
mocks.withLeaderLock.mockImplementation(withLeaderLock)
798+
redisConfigMockFns.mockAcquireLock.mockResolvedValueOnce(false)
799+
})
800+
801+
it.each([
802+
{
803+
label: 'serves a Google token still valid after the wait',
804+
row: { providerId: 'google-drive', remainingMs: 100_000 },
805+
token: 'original-access-token',
806+
},
807+
{
808+
label: 'fails a Google token with under 30 seconds left after the wait',
809+
row: { providerId: 'google-drive', remainingMs: 50_000 },
810+
token: null,
811+
},
812+
{
813+
label: 'fails a provider whose refresh can invalidate the stored chain',
814+
row: { providerId: 'confluence', remainingMs: 100_000 },
815+
token: null,
816+
},
817+
])('$label', async ({ row, token }) => {
818+
queueCredentialAccount({ ...createOAuthAccount(row.remainingMs), providerId: row.providerId })
819+
const pending = resolveCredentialTokenBundle(RAW_CREDENTIAL_ID, RAW_USER_ID, 'test')
820+
await vi.advanceTimersByTimeAsync(31_000)
821+
await expect(pending).resolves.toEqual(token === null ? null : { accessToken: token })
822+
expect(mocks.refreshOAuthToken).not.toHaveBeenCalled()
823+
})
824+
})
825+
826+
it('fails a Google chain the provider revokes even while its access token is valid', async () => {
827+
vi.mocked(isTerminalRefreshError).mockReturnValue(true)
828+
vi.mocked(isCredentialRevocationError).mockReturnValue(true)
829+
mocks.refreshOAuthToken.mockResolvedValue({ ok: false, errorCode: 'invalid_grant' })
830+
const row = createOAuthAccount(100_000)
831+
queueTableRows(account, [row])
832+
queueTableRows(account, [row])
833+
await expect(refreshTokenIfNeeded('test', row, RAW_ACCOUNT_ID)).rejects.toBeInstanceOf(
834+
CredentialRevokedError
835+
)
836+
expect(dbChainMockFns.set).toHaveBeenCalledWith(
837+
expect.objectContaining({ refreshRevokedCode: 'invalid_grant' })
838+
)
839+
})
840+
841+
it('answers a Microsoft caller from its valid token while refresh-token aging refreshes', async () => {
842+
vi.mocked(isMicrosoftProvider).mockReturnValue(true)
843+
const provider = createDeferred<unknown>()
844+
mocks.refreshOAuthToken.mockReturnValue(provider.promise)
845+
const row = {
846+
...createOAuthAccount(3_600_000),
847+
providerId: 'microsoft-teams',
848+
refreshTokenExpiresAt: new Date(now.getTime() + 24 * 60 * 60_000),
849+
}
850+
await expect(refreshTokenIfNeeded('test', row, RAW_ACCOUNT_ID)).resolves.toEqual({
851+
accessToken: 'original-access-token',
852+
refreshed: false,
853+
})
854+
await vi.waitFor(() => expect(mocks.refreshOAuthToken).toHaveBeenCalledOnce())
855+
expect(dbChainMockFns.set).not.toHaveBeenCalled()
856+
provider.resolve({
857+
ok: true,
858+
accessToken: 'refreshed-access-token',
859+
refreshToken: 'rotated-refresh-token',
860+
expiresIn: 3600,
861+
})
862+
await vi.waitFor(() =>
863+
expect(dbChainMockFns.set).toHaveBeenCalledWith(
864+
expect.objectContaining({ accessToken: 'refreshed-access-token' })
865+
)
866+
)
867+
})
868+
869+
it('spreads refresh-ahead start across accounts expiring at the same moment', async () => {
870+
const row = createOAuthAccount(270_000)
871+
/** Chosen because their hashes fall on opposite sides of the 30s jitter cut, so a hash-rule change fails loudly. */
872+
await expect(refreshTokenIfNeeded('test', row, 'account-2')).resolves.toEqual({
873+
accessToken: 'original-access-token',
874+
refreshed: false,
875+
})
876+
await vi.waitFor(() => expect(mocks.refreshOAuthToken).toHaveBeenCalledOnce())
877+
await expect(refreshTokenIfNeeded('test', row, 'account-c')).resolves.toEqual({
878+
accessToken: 'original-access-token',
879+
refreshed: false,
880+
})
881+
await flushMicrotasks(20)
882+
expect(mocks.refreshOAuthToken).toHaveBeenCalledOnce()
883+
})
786884
})
787885

788886
describe('Google service-account token minting', () => {
@@ -797,7 +895,11 @@ describe('Google service-account token minting', () => {
797895
const driveScope = 'https://www.googleapis.com/auth/drive.readonly'
798896
const now = new Date('2026-09-09T18:00:00.000Z')
799897

800-
beforeEach(() => {
898+
beforeEach(async () => {
899+
const { coalesceLocally } = await vi.importActual<
900+
typeof import('@/lib/concurrency/singleflight')
901+
>('@/lib/concurrency/singleflight')
902+
mocks.coalesceLocally.mockImplementation(coalesceLocally)
801903
resetDbChainMock()
802904
vi.useFakeTimers()
803905
vi.setSystemTime(now)
@@ -896,6 +998,19 @@ describe('Google service-account token minting', () => {
896998
expect(fetchMock).toHaveBeenCalledTimes(1)
897999
})
8981000

1001+
it('shares one key read across concurrent mints', async () => {
1002+
queueTableRows(credential, [row])
1003+
await expect(
1004+
Promise.all(
1005+
Array.from({ length: 6 }, (_, index) =>
1006+
getServiceAccountToken('credential-1', [driveScope], `member-${index}@example.com`)
1007+
)
1008+
)
1009+
).resolves.toEqual(Array(6).fill('google-access-token'))
1010+
expect(dbChainMockFns.select).toHaveBeenCalledTimes(1)
1011+
expect(fetchMock).toHaveBeenCalledTimes(6)
1012+
})
1013+
8991014
it('preserves cancellation through the service-account resolver registry', async () => {
9001015
const controller = new AbortController()
9011016
const reason = new DOMException('Caller cancelled', 'AbortError')

‎apps/sim/lib/oauth/credential-service.ts‎

Lines changed: 60 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -225,6 +225,38 @@ const SA_DIAGNOSTIC_ERROR_CODES = new Set([
225225
'unauthorized_client',
226226
])
227227

228+
/**
229+
* The decrypted key of a live Google service-account credential. Reads running at once (one per
230+
* delegated subject) share one query and decrypt; a later read queries again, so it sees a
231+
* revocation recorded since.
232+
*/
233+
function readGoogleServiceAccountKey(credentialId: string): Promise<string> {
234+
return coalesceLocally(`sa-key:${credentialId}`, async () => {
235+
const [credentialRow] = await db
236+
.select({
237+
type: credential.type,
238+
providerId: credential.providerId,
239+
revokedAt: credential.revokedAt,
240+
encryptedServiceAccountKey: credential.encryptedServiceAccountKey,
241+
})
242+
.from(credential)
243+
.where(eq(credential.id, credentialId))
244+
.limit(1)
245+
246+
if (
247+
credentialRow?.type !== 'service_account' ||
248+
credentialRow.providerId !== GOOGLE_SERVICE_ACCOUNT_PROVIDER_ID ||
249+
credentialRow.revokedAt !== null ||
250+
!credentialRow.encryptedServiceAccountKey
251+
) {
252+
throw new Error('Google service account credential is unavailable')
253+
}
254+
255+
const { decrypted } = await decryptSecret(credentialRow.encryptedServiceAccountKey)
256+
return decrypted
257+
})
258+
}
259+
228260
/**
229261
* Generates a short-lived access token for a Google service account credential
230262
* using the two-legged OAuth JWT flow (RFC 7523).
@@ -241,27 +273,7 @@ export async function getServiceAccountToken(
241273
options?: CredentialTokenResolutionOptions
242274
): Promise<string> {
243275
options?.signal?.throwIfAborted()
244-
const [credentialRow] = await db
245-
.select({
246-
type: credential.type,
247-
providerId: credential.providerId,
248-
revokedAt: credential.revokedAt,
249-
encryptedServiceAccountKey: credential.encryptedServiceAccountKey,
250-
})
251-
.from(credential)
252-
.where(eq(credential.id, credentialId))
253-
.limit(1)
254-
255-
if (
256-
credentialRow?.type !== 'service_account' ||
257-
credentialRow.providerId !== GOOGLE_SERVICE_ACCOUNT_PROVIDER_ID ||
258-
credentialRow.revokedAt !== null ||
259-
!credentialRow.encryptedServiceAccountKey
260-
) {
261-
throw new Error('Google service account credential is unavailable')
262-
}
263-
264-
const { decrypted } = await decryptSecret(credentialRow.encryptedServiceAccountKey)
276+
const decrypted = await readGoogleServiceAccountKey(credentialId)
265277
const keyData = JSON.parse(decrypted) as {
266278
client_email: string
267279
private_key: string
@@ -1003,6 +1015,10 @@ function isOAuthAccessTokenExpiring(
10031015
return expiresAt != null && expiresAt.getTime() <= now.getTime() + refreshWindowMs
10041016
}
10051017

1018+
function remainingValidityMs(expiresAt: Date, now = new Date()): number {
1019+
return expiresAt.getTime() - now.getTime()
1020+
}
1021+
10061022
/**
10071023
* Where a provider can refresh ahead of expiry, a caller is served the stored token with this much
10081024
* validity left or less while a background refresh replaces it. It starts above
@@ -1016,12 +1032,16 @@ const OAUTH_REFRESH_AHEAD_WINDOW_MS = 10 * 60 * 1000
10161032
* the two windows so every account still refreshes ahead before callers have to wait.
10171033
*/
10181034
const OAUTH_REFRESH_AHEAD_JITTER_MS = 60 * 1000
1035+
/**
1036+
* The validity a stored token must still have, after a failed blocking refresh, to be served
1037+
* instead of failing the request.
1038+
*/
1039+
const OAUTH_STORED_TOKEN_FALLBACK_MIN_MS = 30 * 1000
10191040

10201041
/**
1021-
* Google and Microsoft keep a refreshed-away access token valid until its own expiry and keep the
1022-
* old refresh token usable (Google never rotates it; Microsoft rotates without revoking the old
1023-
* one), so a caller can use the old token while the refresh runs and an abandoned refresh loses
1024-
* nothing.
1042+
* A Google refresh does not rotate the refresh token, and Microsoft does not revoke the old one
1043+
* when issuing a new one (learn.microsoft.com, "Refresh tokens in the Microsoft identity
1044+
* platform"); serving the stored access token meanwhile relies on it staying valid until expiry.
10251045
*/
10261046
function refreshesAheadOfExpiry(providerId: string): boolean {
10271047
return (
@@ -1077,7 +1097,7 @@ function planAccessTokenRefresh(
10771097
if (
10781098
refreshesAhead &&
10791099
accessTokenExpiresAt &&
1080-
accessTokenExpiresAt.getTime() - now.getTime() <=
1100+
remainingValidityMs(accessTokenExpiresAt, now) <=
10811101
OAUTH_REFRESH_AHEAD_WINDOW_MS - refreshAheadJitterMs(accountId)
10821102
) {
10831103
return { kind: 'refresh-ahead', refreshToken }
@@ -1616,8 +1636,20 @@ async function resolveStoredAccessToken(
16161636
if (fresh?.kind === 'token') {
16171637
return { kind: 'token', accessToken: fresh.accessToken, refreshed: true }
16181638
}
1619-
if (plan.kind === 'refresh-or-reuse' && stored.accessToken) {
1620-
logger.info(`${prefix}Refresh unavailable; reusing still-valid access token`)
1639+
/**
1640+
* A refresh ahead abandoned with its process (Trigger.dev runs have no keep-alive) holds its
1641+
* lease until the TTL, so a follower can wait it out with nothing; a chain that survives
1642+
* refreshes still serves its stored token if enough validity is left after that wait.
1643+
*/
1644+
const orphaned =
1645+
fresh === null &&
1646+
refreshesAheadOfExpiry(refresh.providerId) &&
1647+
stored.accessTokenExpiresAt != null &&
1648+
remainingValidityMs(stored.accessTokenExpiresAt) > OAUTH_STORED_TOKEN_FALLBACK_MIN_MS
1649+
if (stored.accessToken && (plan.kind === 'refresh-or-reuse' || orphaned)) {
1650+
logger.info(
1651+
`${prefix}${orphaned ? 'Refresh produced no token' : 'Refresh unavailable'}; reusing still-valid access token`
1652+
)
16211653
return { kind: 'token', accessToken: stored.accessToken, refreshed: false }
16221654
}
16231655
return { kind: 'failed', outcome: fresh }

0 commit comments

Comments
 (0)