Skip to content

Commit a1985a7

Browse files
committed
fix(desktop): renew an import's lease only after its claim, and give a capped call up without reading its lease
The first heartbeat now comes one beat in, after the desktop's bounded claim, so every renewal follows the claim and a refusal always means the call was stopped, settled, or lapsed. A call at its ceiling is given up before its lease is read, and the force-fail log names whether the budget, the cap, a lapsed lease, or failed lease lookups ended the wait.
1 parent 0b20373 commit a1985a7

4 files changed

Lines changed: 130 additions & 113 deletions

File tree

‎apps/sim/lib/mothership/request/lifecycle/run.test.ts‎

Lines changed: 28 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import { resetEnvFlagsMock, resetEnvironmentUtilsMock, setEnvFlags } from '@sim/testing'
2+
import { getMockLogger } from '@sim/testing/mocks/logger.mock'
23
import {
34
mothershipAgentUrlMock,
45
mothershipAgentUrlMockFns,
@@ -3607,7 +3608,12 @@ describe('runCopilotLifecycle', () => {
36073608
if (!streamContext) throw new Error('The turn did not start its stream')
36083609
return streamContext
36093610
}
3610-
return { bodies, lifecycle, resumedWithLostImport, context }
3611+
/** The reasons the turn logged for force-failing its calls. */
3612+
const forceFailures = () =>
3613+
getMockLogger('CopilotLifecycle')
3614+
.error.mock.calls.map(([message]) => String(message))
3615+
.filter((message) => message.endsWith('force-failing'))
3616+
return { bodies, lifecycle, resumedWithLostImport, context, forceFailures }
36113617
}
36123618

36133619
it('waits on a chat-view import while its lease is renewed, and fails it once the lease lapses', async () => {
@@ -3625,6 +3631,9 @@ describe('runCopilotLifecycle', () => {
36253631
await vi.advanceTimersByTimeAsync(52_000)
36263632
expect((await turn.lifecycle).success).toBe(true)
36273633
expect(turn.resumedWithLostImport()).toBe(true)
3634+
expect(turn.forceFailures()).toEqual([
3635+
'Pending tool execution has no live lease past its budget; force-failing',
3636+
])
36283637
} finally {
36293638
mothershipAsyncRunsMockFns.mockGetChatViewDesktopLeaseRemainingMs.mockReset()
36303639
vi.useRealTimers()
@@ -3646,6 +3655,9 @@ describe('runCopilotLifecycle', () => {
36463655
await vi.advanceTimersByTimeAsync(32_000)
36473656
expect((await turn.lifecycle).success).toBe(true)
36483657
expect(turn.resumedWithLostImport()).toBe(true)
3658+
expect(turn.forceFailures()).toEqual([
3659+
'Pending tool execution has no live lease past its budget; force-failing',
3660+
])
36493661
} finally {
36503662
mothershipAsyncRunsMockFns.mockGetChatViewDesktopLeaseRemainingMs.mockReset()
36513663
vi.useRealTimers()
@@ -3660,9 +3672,18 @@ describe('runCopilotLifecycle', () => {
36603672
const turn = runImportTurn()
36613673
await vi.advanceTimersByTimeAsync(CLIENT_TOOL_RESULT_TIMEOUT_MS - 1_000)
36623674
expect(turn.bodies).toHaveLength(1)
3675+
const leaseReads =
3676+
mothershipAsyncRunsMockFns.mockGetChatViewDesktopLeaseRemainingMs.mock.calls.length
36633677
await vi.advanceTimersByTimeAsync(2_000)
3678+
// At the cap the call is given up without reading its lease.
3679+
expect(
3680+
mothershipAsyncRunsMockFns.mockGetChatViewDesktopLeaseRemainingMs.mock.calls.length
3681+
).toBe(leaseReads)
36643682
expect((await turn.lifecycle).success).toBe(true)
36653683
expect(turn.resumedWithLostImport()).toBe(true)
3684+
expect(turn.forceFailures()).toEqual([
3685+
'Pending tool execution reached the client tool result cap; force-failing',
3686+
])
36663687
} finally {
36673688
mothershipAsyncRunsMockFns.mockGetChatViewDesktopLeaseRemainingMs.mockReset()
36683689
vi.useRealTimers()
@@ -3682,6 +3703,9 @@ describe('runCopilotLifecycle', () => {
36823703
await vi.advanceTimersByTimeAsync(10_000)
36833704
expect((await turn.lifecycle).success).toBe(true)
36843705
expect(turn.resumedWithLostImport()).toBe(true)
3706+
expect(turn.forceFailures()).toEqual([
3707+
'Pending tool execution lease could not be read for a whole lease; force-failing',
3708+
])
36853709
} finally {
36863710
mothershipAsyncRunsMockFns.mockGetChatViewDesktopLeaseRemainingMs.mockReset()
36873711
vi.useRealTimers()
@@ -3702,6 +3726,9 @@ describe('runCopilotLifecycle', () => {
37023726
await vi.advanceTimersByTimeAsync(20_000)
37033727
expect((await turn.lifecycle).success).toBe(true)
37043728
expect(turn.resumedWithLostImport()).toBe(true)
3729+
expect(turn.forceFailures()).toEqual([
3730+
'Pending tool execution lease could not be read for a whole lease; force-failing',
3731+
])
37053732
} finally {
37063733
mothershipAsyncRunsMockFns.mockGetChatViewDesktopLeaseRemainingMs.mockReset()
37073734
vi.useRealTimers()

‎apps/sim/lib/mothership/request/lifecycle/run.ts‎

Lines changed: 73 additions & 62 deletions
Original file line numberDiff line numberDiff line change
@@ -103,6 +103,33 @@ const LEASE_LOOKUP_RETRY_MS = 5_000
103103
/** How long one lease lookup may take before it counts as failed. */
104104
const LEASE_LOOKUP_TIMEOUT_MS = 5_000
105105

106+
/** The resume wait's watch over one pending tool call. */
107+
interface PendingToolWatchdog {
108+
promise: Promise<AsyncCompletionSignal>
109+
settlement: Promise<{ toolCallId: string; promise: Promise<AsyncCompletionSignal> }>
110+
deadlineAt: number
111+
waitBudgetMs: number
112+
/** When this wait began. */
113+
startedAt: number
114+
/** Past this, the call is given up however its lease stands: any client tool's cap. */
115+
ceilingAt: number
116+
/** Why the deadline was moved past the plain budget, if it was. */
117+
extendedBy?: 'lease' | 'lease_lookup'
118+
/** Until when a failing lease lookup is retried before the call is given up. */
119+
leaseLookupRetryUntil?: number
120+
}
121+
122+
/** Why a pending call's wait ran out, in the words its force-fail is logged with. */
123+
function pendingToolExpiryMessage(watchdog: PendingToolWatchdog, now: number): string {
124+
if (!watchdog.extendedBy)
125+
return 'Pending tool execution exceeded its resume wait budget; force-failing'
126+
if (now >= watchdog.ceilingAt)
127+
return 'Pending tool execution reached the client tool result cap; force-failing'
128+
if (watchdog.leaseLookupRetryUntil !== undefined)
129+
return 'Pending tool execution lease could not be read for a whole lease; force-failing'
130+
return 'Pending tool execution has no live lease past its budget; force-failing'
131+
}
132+
106133
/** A pending call's chat-view lease, or why it could not be read within `timeoutMs`. */
107134
async function readLeaseWithin(
108135
toolCallId: string,
@@ -1390,26 +1417,7 @@ async function runCheckpointLoop(
13901417
})
13911418
let maximumWaitBudgetMs = 0
13921419
let timedOutCount = 0
1393-
const pendingWatchdogs = new Map<
1394-
string,
1395-
{
1396-
promise: Promise<AsyncCompletionSignal>
1397-
settlement: Promise<{
1398-
toolCallId: string
1399-
promise: Promise<AsyncCompletionSignal>
1400-
}>
1401-
deadlineAt: number
1402-
waitBudgetMs: number
1403-
/** When this wait began. */
1404-
startedAt: number
1405-
/** Past this, the call is given up however its lease stands: any client tool's cap. */
1406-
ceilingAt: number
1407-
/** Why the deadline was moved past the plain budget, if it was. */
1408-
extendedBy?: 'lease' | 'lease_lookup'
1409-
/** Until when a failing lease lookup is retried before the call is given up. */
1410-
leaseLookupRetryUntil?: number
1411-
}
1412-
>()
1420+
const pendingWatchdogs = new Map<string, PendingToolWatchdog>()
14131421

14141422
/**
14151423
* A long-running approval must not lend its deadline to an unrelated
@@ -1450,59 +1458,62 @@ async function runCheckpointLoop(
14501458
context.pendingToolPromises.get(toolCallId) === watchdog.promise
14511459
)
14521460
// A desktop import the chat view is running renews its lease while it works: its budget
1453-
// runs to the end of that lease, and only a lapsed lease fails it.
1461+
// runs to the end of that lease, and only a lapsed lease fails it. At its ceiling a call is
1462+
// given up however its lease stands, so its lease is not read.
14541463
// A lookup that fails says nothing about the lease, so it is retried, for at most one lease.
14551464
// Each lookup is bounded, so a stalled read cannot hold the wait past its deadlines or Stop.
1465+
const leaseChecked = overdueTools.filter(([, watchdog]) => now < watchdog.ceilingAt)
14561466
const leases = await Promise.all(
1457-
overdueTools.map(([toolCallId]) =>
1467+
leaseChecked.map(([toolCallId]) =>
14581468
readLeaseWithin(toolCallId, LEASE_LOOKUP_TIMEOUT_MS, options.abortSignal)
14591469
)
14601470
)
14611471
if (isAborted(options, context)) break
1462-
const expiredTools = overdueTools.filter(([toolCallId, watchdog], index) => {
1463-
// A call replaced while its lease was read belongs to its new watchdog.
1464-
if (context.pendingToolPromises.get(toolCallId) !== watchdog.promise) return false
1465-
const lease = leases[index]
1466-
const checkedAt = Date.now()
1467-
if (checkedAt >= watchdog.ceilingAt) return true
1468-
const extendTo = (deadlineAt: number, reason: 'lease' | 'lease_lookup') => {
1469-
watchdog.deadlineAt = Math.min(deadlineAt, watchdog.ceilingAt)
1470-
watchdog.extendedBy = reason
1471-
maximumWaitBudgetMs = Math.max(
1472-
maximumWaitBudgetMs,
1473-
watchdog.deadlineAt - watchdog.startedAt
1474-
)
1475-
return false
1476-
}
1477-
if ('error' in lease) {
1478-
watchdog.leaseLookupRetryUntil ??= checkedAt + SIM_TOOL_EXECUTION_LEASE_SECONDS * 1000
1479-
if (checkedAt >= watchdog.leaseLookupRetryUntil) return true
1480-
logger.warn('Could not read a pending tool call lease; checking again', {
1481-
toolCallId,
1482-
error: getErrorMessage(lease.error),
1483-
})
1484-
return extendTo(checkedAt + LEASE_LOOKUP_RETRY_MS, 'lease_lookup')
1485-
}
1486-
watchdog.leaseLookupRetryUntil = undefined
1487-
if (typeof lease.remainingMs !== 'number' || lease.remainingMs <= 0) return true
1488-
return extendTo(checkedAt + lease.remainingMs + LEASE_RECHECK_SLACK_MS, 'lease')
1489-
})
1472+
const expiredTools = [
1473+
...overdueTools.filter(
1474+
([toolCallId, watchdog]) =>
1475+
now >= watchdog.ceilingAt &&
1476+
context.pendingToolPromises.get(toolCallId) === watchdog.promise
1477+
),
1478+
...leaseChecked.filter(([toolCallId, watchdog], index) => {
1479+
// A call replaced while its lease was read belongs to its new watchdog.
1480+
if (context.pendingToolPromises.get(toolCallId) !== watchdog.promise) return false
1481+
const lease = leases[index]
1482+
const checkedAt = Date.now()
1483+
const extendTo = (deadlineAt: number, reason: 'lease' | 'lease_lookup') => {
1484+
watchdog.deadlineAt = Math.min(deadlineAt, watchdog.ceilingAt)
1485+
watchdog.extendedBy = reason
1486+
maximumWaitBudgetMs = Math.max(
1487+
maximumWaitBudgetMs,
1488+
watchdog.deadlineAt - watchdog.startedAt
1489+
)
1490+
return false
1491+
}
1492+
if ('error' in lease) {
1493+
watchdog.leaseLookupRetryUntil ??= checkedAt + SIM_TOOL_EXECUTION_LEASE_SECONDS * 1000
1494+
if (checkedAt >= watchdog.leaseLookupRetryUntil) return true
1495+
logger.warn('Could not read a pending tool call lease; checking again', {
1496+
toolCallId,
1497+
error: getErrorMessage(lease.error),
1498+
})
1499+
return extendTo(checkedAt + LEASE_LOOKUP_RETRY_MS, 'lease_lookup')
1500+
}
1501+
watchdog.leaseLookupRetryUntil = undefined
1502+
if (typeof lease.remainingMs !== 'number' || lease.remainingMs <= 0) return true
1503+
return extendTo(checkedAt + lease.remainingMs + LEASE_RECHECK_SLACK_MS, 'lease')
1504+
}),
1505+
]
14901506
if (expiredTools.length > 0) {
14911507
await Promise.all(
14921508
expiredTools.map(async ([toolCallId, watchdog]) => {
14931509
const waitedMs = Date.now() - watchdog.startedAt
1494-
logger.error(
1495-
watchdog.extendedBy
1496-
? 'Pending tool execution outlived its renewed lease or its cap; force-failing'
1497-
: 'Pending tool execution exceeded its resume wait budget; force-failing',
1498-
{
1499-
checkpointId: continuation.checkpointId,
1500-
toolCallId,
1501-
waitBudgetMs: watchdog.waitBudgetMs,
1502-
waitedMs,
1503-
...(watchdog.extendedBy ? { extendedBy: watchdog.extendedBy } : {}),
1504-
}
1505-
)
1510+
logger.error(pendingToolExpiryMessage(watchdog, now), {
1511+
checkpointId: continuation.checkpointId,
1512+
toolCallId,
1513+
waitBudgetMs: watchdog.waitBudgetMs,
1514+
waitedMs,
1515+
...(watchdog.extendedBy ? { extendedBy: watchdog.extendedBy } : {}),
1516+
})
15061517
await failPendingToolCall(toolCallId, context, execContext)
15071518
if (context.pendingToolPromises.get(toolCallId) === watchdog.promise) {
15081519
context.pendingToolPromises.delete(toolCallId)

‎apps/sim/lib/mothership/tools/client/native-files.test.ts‎

Lines changed: 19 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -158,18 +158,18 @@ describe('an import keeps its lease while it runs', () => {
158158
const LEASE_MS = 60_000
159159
/**
160160
* The server's side of the lease, by the rules the lease route applies: the claim (which the
161-
* desktop makes as it starts building the manifest) takes a lease, and a renewal extends it only
162-
* while it is live; a refused renewal answers 410.
161+
* desktop makes before it scans, within its 8 s authorization timeout) takes a lease, and a
162+
* renewal extends it only while it is live; a refused renewal answers 410.
163163
*/
164164
let server: {
165165
leaseUntil: number
166166
claimed: boolean
167167
stopped: boolean
168168
renewalsReceived: number
169-
transient: number[]
170-
/** How long the server takes to answer the next renewal it receives. */
171-
nextAnswerDelayMs: number
169+
/** How the server answers the renewals it receives next: a failure that may pass, or as usual. */
170+
answers: (number | 'usual')[]
172171
}
172+
let claimMs: number
173173
let scanMs: number
174174
let finishUpload: () => void
175175
const leaseLive = () => Date.now() < server.leaseUntil
@@ -183,13 +183,14 @@ describe('an import keeps its lease while it runs', () => {
183183
claimed: false,
184184
stopped: false,
185185
renewalsReceived: 0,
186-
transient: [],
187-
nextAnswerDelayMs: 0,
186+
answers: [],
188187
}
188+
claimMs = 0
189189
scanMs = 0
190190
mocks.invoke.mockImplementation(async (request: { operation: string }) => {
191191
if (request.operation !== 'manifest')
192192
return { ok: true, data: { kind: 'chunk', bytes: new Uint8Array([65, 66, 67]), eof: true } }
193+
await sleep(claimMs)
193194
server.claimed = true
194195
server.leaseUntil = Date.now() + LEASE_MS
195196
await sleep(scanMs)
@@ -204,14 +205,10 @@ describe('an import keeps its lease while it runs', () => {
204205
mocks.json.mockImplementation(async (contract: unknown) => {
205206
if (contract !== renewDesktopToolLeaseContract) return { folder: { id: 'created-folder' } }
206207
server.renewalsReceived += 1
207-
const delayMs = server.nextAnswerDelayMs
208-
server.nextAnswerDelayMs = 0
209-
const transient = server.transient.shift()
210-
const refused = server.stopped || !server.claimed || !leaseLive()
211-
if (!transient && !refused) server.leaseUntil = Date.now() + LEASE_MS
212-
await sleep(delayMs)
213-
if (transient) return answer(transient)
214-
if (refused) return answer(410)
208+
const transient = server.answers.shift()
209+
if (typeof transient === 'number') return answer(transient)
210+
if (server.stopped || !server.claimed || !leaseLive()) return answer(410)
211+
server.leaseUntil = Date.now() + LEASE_MS
215212
return { renewed: true }
216213
})
217214
})
@@ -233,11 +230,8 @@ describe('an import keeps its lease while it runs', () => {
233230
expect(leaseLive()).toBe(false)
234231
})
235232

236-
it('a refusal of a renewal sent before the claim does not stop it, however late it arrives', async () => {
237-
// The first renewal goes out before the desktop claims the call; its 410 arrives only after
238-
// the manifest confirmed the claim.
239-
server.nextAnswerDelayMs = 5_000
240-
scanMs = 1_000
233+
it('keeps renewing when the desktop takes most of its authorization timeout to claim', async () => {
234+
claimMs = 7_000
241235
const run = executeNativeFileTool('tool', 'import_local_files')
242236
await vi.advanceTimersByTimeAsync(90_000)
243237
expect(leaseLive()).toBe(true)
@@ -259,9 +253,12 @@ describe('an import keeps its lease while it runs', () => {
259253
})
260254

261255
it('keeps renewing through failures that may pass', async () => {
262-
server.transient = [401, 429, 503]
256+
// Beats at 20 s, 60 s and 100 s fail; the beats between them renew. Stopping at any of the
257+
// failures would let the lease lapse by 140 s.
258+
server.answers = [401, 'usual', 429, 'usual', 503]
263259
const run = executeNativeFileTool('tool', 'import_local_files')
264-
await vi.advanceTimersByTimeAsync(90_000)
260+
await vi.advanceTimersByTimeAsync(150_000)
261+
expect(server.renewalsReceived).toBe(7)
265262
expect(leaseLive()).toBe(true)
266263
finishUpload()
267264
await run

‎apps/sim/lib/mothership/tools/client/native-files.ts‎

Lines changed: 10 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -112,24 +112,19 @@ export async function importNativeFiles(
112112
}
113113

114114
/**
115-
* Renews an import's lease from the moment its manifest is requested, every heartbeat, the way a
115+
* Renews an import's lease every heartbeat from the moment its manifest is requested, the way a
116116
* desktop renews a bound call, so a slow directory scan cannot outlast the lease the claim took.
117-
* The claim lands while the manifest is being built, so a refusal (410) counts only once the
118-
* manifest confirmed the claim: then the call was stopped, settled, or its lease lapsed, and
117+
* The desktop claims the call before it scans, within its authorization timeout, which is well
118+
* inside one heartbeat: every renewal follows the claim, and the claim's own lease covers the
119+
* first beat. A refusal (410) means the call was stopped, settled, or its lease lapsed, and
119120
* renewing stops. Any other failure may pass, and the next beat tries again.
120121
*/
121-
function keepImportLeased(toolCallId: string): { claimed(): void; stop(): void } {
122-
let claimed = false
123-
let stopped = false
124-
const renew = () => {
125-
if (stopped) return
126-
// Whether the claim was confirmed when this renewal was sent: a refusal of one sent before it
127-
// may arrive after, and says nothing about the running import.
128-
const sentAfterClaim = claimed
122+
function keepImportLeased(toolCallId: string): { stop(): void } {
123+
const timer = setInterval(() => {
129124
requestJson(renewDesktopToolLeaseContract, { body: { toolCallId, chatView: true } }).catch(
130125
(error) => {
131-
if (sentAfterClaim && error instanceof ApiClientError && error.status === 410) {
132-
stop()
126+
if (error instanceof ApiClientError && error.status === 410) {
127+
clearInterval(timer)
133128
return
134129
}
135130
logger.warn('Could not renew the import lease; trying again next beat', {
@@ -138,20 +133,8 @@ function keepImportLeased(toolCallId: string): { claimed(): void; stop(): void }
138133
})
139134
}
140135
)
141-
}
142-
const timer = setInterval(renew, SIM_TOOL_EXECUTION_HEARTBEAT_MS)
143-
const stop = () => {
144-
stopped = true
145-
clearInterval(timer)
146-
}
147-
renew()
148-
return {
149-
claimed() {
150-
claimed = true
151-
renew()
152-
},
153-
stop,
154-
}
136+
}, SIM_TOOL_EXECUTION_HEARTBEAT_MS)
137+
return { stop: () => clearInterval(timer) }
155138
}
156139

157140
/** The server claims imports before reading their manifest, preventing replayed uploads. */
@@ -192,7 +175,6 @@ export async function executeNativeFileTool(
192175
if (response.data.kind === 'chunk') throw new Error('Unexpected chunk outside an import.')
193176
let completion
194177
if (response.data.kind === 'manifest') {
195-
lease?.claimed()
196178
completion = localFileImportCompletion(
197179
await importNativeFiles(toolCallId, response.data, signal)
198180
)

0 commit comments

Comments
 (0)