Skip to content

Commit a502fb2

Browse files
authored
fix(oauth): take over a lapsed refresh lease, share overlapping service-account key reads, and label unsearched accounts (#9537)
* fix(search): keep a valid token when a refresh lease is orphaned, share overlapping service-account reads, and label unsearched accounts * fix(oauth): serve a stored token only when a refresh wait ran out on a held lease - A failed blocking refresh now falls back to the stored token only on a distinct lease-held outcome: this caller followed another lease holder for its whole wait, its last read still found the chain it started from, and no leader flagged it dead. A dead flag, a terminal refresh error, a vanished account or a leader exception stay failures. - The in-process refresh settles after the follower's full wait plus headroom instead of the default 30s, which fired before the follower's last read could count; a leader hung in this process still fails its callers, now at that bound. * fix(oauth): let a cancelled service-account mint stop waiting on a shared key read * docs(oauth): say why the stored-token fallback floor sits below the refresh window * fix(oauth): take over a lapsed refresh lease instead of serving the stored token - A follower whose wait ends with nothing now tries the lock once more: a leader whose process died leaves its lease to lapse, and the follower leads the refresh instead of failing. Losing that try still returns nothing. - Drops the stored-token fallback and its lease-held outcome, which served tokens below the five-minute floor a served token keeps. - The leader rereads the dead flag under the lock, so a follower taking over after a leader that failed terminally does not call the provider again. - The in-process refresh settles after the follower wait plus a full leader run and headroom.
1 parent f16b71a commit a502fb2

8 files changed

Lines changed: 636 additions & 274 deletions

File tree

‎apps/sim/lib/concurrency/leader-lock.test.ts‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -119,6 +119,32 @@ describe('withLeaderLock', () => {
119119
}
120120
})
121121

122+
describe('when the follower wait ends with nothing', () => {
123+
const waitedOut = {
124+
key: 'k',
125+
pollIntervalMs: 5,
126+
maxWaitMs: 20,
127+
onFollower: async () => null,
128+
}
129+
130+
it('leads once a leader that vanished leaves its lease to lapse', async () => {
131+
redisConfigMockFns.mockAcquireLock.mockResolvedValueOnce(false)
132+
const onLeader = vi.fn(async () => 'took-over')
133+
await expect(withLeaderLock<string>({ ...waitedOut, onLeader })).resolves.toBe('took-over')
134+
expect(onLeader).toHaveBeenCalledOnce()
135+
expect(redisConfigMockFns.mockAcquireLock).toHaveBeenCalledTimes(2)
136+
expect(redisConfigMockFns.mockReleaseLock).toHaveBeenCalledOnce()
137+
})
138+
139+
it('returns the follower result when a live holder still has the lease', async () => {
140+
redisConfigMockFns.mockAcquireLock.mockResolvedValue(false)
141+
const onLeader = vi.fn(async () => 'should-not-run')
142+
await expect(withLeaderLock<string>({ ...waitedOut, onLeader })).resolves.toBeNull()
143+
expect(onLeader).not.toHaveBeenCalled()
144+
expect(redisConfigMockFns.mockAcquireLock).toHaveBeenCalledTimes(2)
145+
})
146+
})
147+
122148
it('only one of N concurrent callers acquires the lock', async () => {
123149
// Track which calls won the lock: first one returns true, rest return false.
124150
let acquired = false

‎apps/sim/lib/concurrency/leader-lock.ts‎

Lines changed: 24 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,13 @@ export interface LeaderLockOptions<T> {
1919
onFollower: () => Promise<T | null>
2020
}
2121

22+
/**
23+
* Runs `onLeader` in the one caller that takes the lock for `key`; every other caller polls
24+
* `onFollower` until it reports the leader's result or `maxWaitMs` passes. A follower whose wait
25+
* ends with nothing tries the lock once more: a leader whose process died without releasing it
26+
* leaves the lease to lapse, and that follower then leads instead of failing. Losing that try
27+
* means a live holder has the lease, and the follower returns `null`.
28+
*/
2229
export async function withLeaderLock<T>(opts: LeaderLockOptions<T>): Promise<T | null> {
2330
const {
2431
key,
@@ -42,7 +49,7 @@ export async function withLeaderLock<T>(opts: LeaderLockOptions<T>): Promise<T |
4249
return onLeader()
4350
}
4451

45-
if (acquired) {
52+
const leadHoldingLock = async () => {
4653
try {
4754
return await onLeader()
4855
} finally {
@@ -57,6 +64,8 @@ export async function withLeaderLock<T>(opts: LeaderLockOptions<T>): Promise<T |
5764
}
5865
}
5966

67+
if (acquired) return leadHoldingLock()
68+
6069
const deadline = Date.now() + maxWaitMs
6170
while (Date.now() < deadline) {
6271
await sleep(pollIntervalMs)
@@ -68,6 +77,20 @@ export async function withLeaderLock<T>(opts: LeaderLockOptions<T>): Promise<T |
6877
const lastChance = await onFollower()
6978
if (lastChance !== null) return lastChance
7079

80+
let tookOver = false
81+
try {
82+
tookOver = await acquireLock(key, ownerToken, ttlSec)
83+
} catch (error) {
84+
logger.warn('Lock takeover failed after the follower wait', {
85+
key,
86+
error: toError(error).message,
87+
})
88+
}
89+
if (tookOver) {
90+
logger.warn('Follower took over a lease its leader left to lapse', { key, maxWaitMs })
91+
return leadHoldingLock()
92+
}
93+
7194
logger.warn('Follower timed out waiting for leader', { key, maxWaitMs })
7295
return null
7396
}
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

0 commit comments

Comments
 (0)