Skip to content

Commit 198a8b7

Browse files
committed
fix(desktop): keep a chat-view import alive while it works, by the lease its session renews
An import the chat view runs was failed as lost once it ran past the default tool budget (60 s plus the 30 s resume grace), though it was still uploading. Its claim now takes the execution lease under the claiming session, the chat view renews it through the existing lease route while the import runs, and the turn's wait budget runs to the end of that lease. Without renewals the lease lapses with the default budget, so a closed or crashed window still settles within about a lease.
1 parent af5501b commit 198a8b7

11 files changed

Lines changed: 662 additions & 22 deletions

File tree

‎apps/desktop/e2e/desktop-tools-live-sim.spec.ts‎

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ import {
1111
test,
1212
} from '@playwright/test'
1313
import type { SimDesktopApi } from '@sim/desktop-bridge'
14+
import { sleep } from '@sim/utils/helpers'
1415
import { generateId } from '@sim/utils/id'
1516
import { toRecord } from '@sim/utils/object'
1617
import {
@@ -34,6 +35,10 @@ import {
3435
const DESKTOP_DIR = fileURLToPath(new URL('..', import.meta.url))
3536
const config = liveSimConfig()
3637
const PICKUP_GRACE_MS = 15_000
38+
/** Longer than the default tool budget (60 s) plus the resume grace (30 s). */
39+
const LONG_IMPORT_MS = 120_000
40+
/** The execution lease a running import holds and renews (`SIM_TOOL_EXECUTION_LEASE_SECONDS`). */
41+
const LEASE_MS = 60_000
3742
/** How long a held request may take to arrive once the step that sends it ran. */
3843
const ARRIVAL_MS = 60_000
3944
/** First requests to a route compile it, which takes minutes on a cold dev app. */
@@ -472,6 +477,69 @@ test.describe('desktop tools against a live Sim', () => {
472477
expect(await db.workspaceFileNames(user.workspaceId)).toEqual(['a.txt'])
473478
})
474479

480+
/** An import whose first upload the proxy holds until released, as a large file's would take. */
481+
async function slowImport(user: SeededUser, title: string, marker: string) {
482+
const source = importSource()
483+
let callId = ''
484+
agent.script(marker, (turn) => {
485+
callId = turn.toolCall({
486+
toolName: 'import_local_files',
487+
args: { path: source, targetWorkspaceId: user.workspaceId },
488+
})
489+
turn.pause()
490+
})
491+
const page = await openApp(user, title)
492+
const firstUpload = proxy.hold(isUploadStart)
493+
await send(page, `${marker} import my reports`)
494+
await firstUpload.arrival(ARRIVAL_MS, 'The first upload')
495+
return { page, firstUpload, callId: () => callId }
496+
}
497+
498+
test('an import that runs longer than the default tool budget still completes', async () => {
499+
test.setTimeout(420_000)
500+
const user = await db.seedUser(['Long import chat', 'Other chat'])
501+
const chatId = user.chats['Long import chat']
502+
const { page, firstUpload, callId } = await slowImport(
503+
user,
504+
'Long import chat',
505+
'[long-import]'
506+
)
507+
// The import is alive and working past the 60 s default budget and its 30 s grace, while the
508+
// user is in another chat: its lease is renewed by the import, not by the chat view.
509+
await openChat(page, user, 'Other chat')
510+
await page.waitForTimeout(LONG_IMPORT_MS)
511+
expect(agent.resultFor(callId())).toBeUndefined()
512+
firstUpload.release()
513+
514+
await agent.waitForResume(() => Boolean(agent.resultFor(callId())), 120_000)
515+
const result = agent.resultFor(callId())
516+
expect(JSON.stringify(result?.data)).not.toContain('outcomeUnknown')
517+
expect(result?.success).toBe(true)
518+
await expect
519+
.poll(() => db.workspaceFileNames(user.workspaceId), { timeout: 30_000 })
520+
.toEqual(['a.txt', 'b.txt'])
521+
expect(await callState(chatId)).toMatch(/^completed/)
522+
})
523+
524+
test('a long import whose window crashed settles as outcome unknown about one lease later', async () => {
525+
test.setTimeout(420_000)
526+
const user = await db.seedUser(['Crashing import chat'])
527+
const { callId } = await slowImport(user, 'Crashing import chat', '[crash-import]')
528+
await sleep(LONG_IMPORT_MS)
529+
expect(agent.resultFor(callId())).toBeUndefined()
530+
// A crash reports nothing on its way out, unlike a closed window: only the lapse of the
531+
// lease the page was renewing tells Sim the import is gone.
532+
const crashedAt = Date.now()
533+
await app?.evaluate(({ webContents }) => {
534+
for (const contents of webContents.getAllWebContents())
535+
if (contents.getURL().includes('/workspace/')) contents.forcefullyCrashRenderer()
536+
})
537+
await agent.waitForResume(() => Boolean(agent.resultFor(callId())), 150_000)
538+
const result = agent.resultFor(callId())
539+
expect(result?.data).toMatchObject({ outcomeUnknown: true })
540+
expect((result?.at ?? 0) - crashedAt).toBeLessThan(LEASE_MS + 20_000)
541+
})
542+
475543
test('signing out ends a desktop tool still running', async () => {
476544
const user = await db.seedUser(['Import chat'])
477545
const chatId = user.chats['Import chat']

‎apps/sim/app/api/desktop/tool/authorize/route.ts‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,8 +18,10 @@ import {
1818
createUnauthorizedResponse,
1919
} from '@/lib/mothership/request/http'
2020
import {
21+
chatViewDesktopLeaseOwnerToken,
2122
getDesktopToolClaimOwner,
2223
isDesktopToolCall,
24+
isLeasedChatViewDesktopTool,
2325
isLocalReadToolCall,
2426
} from '@/lib/mothership/tools/desktop-tools'
2527

@@ -52,7 +54,7 @@ function refusedClaimResponse(
5254
* that they have not allowed.
5355
*/
5456
export const POST = withRouteHandler(async (request: NextRequest) => {
55-
const { userId, isAuthenticated } = await authenticateCopilotRequestSessionOnly()
57+
const { userId, isAuthenticated, principal } = await authenticateCopilotRequestSessionOnly()
5658
if (!isAuthenticated || !userId) {
5759
return createUnauthorizedResponse()
5860
}
@@ -114,11 +116,15 @@ export const POST = withRouteHandler(async (request: NextRequest) => {
114116
{ status: 409 }
115117
)
116118
if (toolCall.status !== 'pending') return alreadyStarted()
119+
// The import's lease is held by this session; the chat view renews it while the import runs.
117120
const { outcome } = await claimDesktopToolCall({
118121
toolCallId: toolCall.toolCallId,
119122
runId: toolCall.runId,
120123
userId,
121124
claimedBy: DESKTOP_TOOL_CLAIM_OWNER.files,
125+
...(principal && isLeasedChatViewDesktopTool(toolCall.toolName)
126+
? { chatView: { ownerToken: chatViewDesktopLeaseOwnerToken(principal.sessionId) } }
127+
: {}),
122128
})
123129
if (outcome !== 'claimed') return refusedClaimResponse(outcome, alreadyStarted)
124130
} else if (

‎apps/sim/lib/api/contracts/desktop-executor.ts‎

Lines changed: 15 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -136,11 +136,21 @@ export const claimDesktopToolContract = defineRouteContract({
136136
error: z.object({ error: z.string() }),
137137
})
138138

139-
const renewDesktopToolLeaseBodySchema = z.object({
140-
deviceId: desktopDeviceIdSchema,
141-
toolCallId: desktopToolCallIdSchema,
142-
executionToken: z.string().min(1).max(128),
143-
})
139+
/**
140+
* A device renews a call of a run bound to it under its execution token; the chat view renews an
141+
* import it claimed (`chatView`), as the session that claimed it.
142+
*/
143+
const renewDesktopToolLeaseBodySchema = z.union([
144+
z.object({
145+
deviceId: desktopDeviceIdSchema,
146+
toolCallId: desktopToolCallIdSchema,
147+
executionToken: z.string().min(1).max(128),
148+
}),
149+
z.object({
150+
toolCallId: desktopToolCallIdSchema,
151+
chatView: z.literal(true),
152+
}),
153+
])
144154
export type RenewDesktopToolLeaseBody = z.input<typeof renewDesktopToolLeaseBodySchema>
145155

146156
export const renewDesktopToolLeaseResponseSchema = z.object({ renewed: z.literal(true) })

‎apps/sim/lib/desktop/application/executor.ts‎

Lines changed: 31 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@ import {
4141
import {
4242
claimDesktopToolCall,
4343
type DesktopToolCallClaim,
44+
getAsyncToolCall,
4445
renewSimToolExecutionLease,
4546
} from '@/lib/mothership/async-runs/repository'
4647
import { sealClientToolSettlement } from '@/lib/mothership/request/tools/client-completion-seal.server'
@@ -50,6 +51,7 @@ import {
5051
settleClientToolCall,
5152
} from '@/lib/mothership/request/tools/client-settlement.server'
5253
import {
54+
chatViewDesktopLeaseOwnerToken,
5355
getDesktopExecutorClaimOwner,
5456
isDesktopToolCall,
5557
} from '@/lib/mothership/tools/desktop-tools'
@@ -329,9 +331,19 @@ interface DesktopCallTokenInput extends DeviceInput {
329331
executionToken: string
330332
}
331333

332-
/** Keeps a running call owned; failing here always means the device must stop the action. */
334+
/** A desktop call the chat view is running, renewed by the session that claimed it. */
335+
interface ChatViewCallInput {
336+
toolCallId: string
337+
chatView: true
338+
}
339+
340+
/**
341+
* Keeps a running call owned; failing here always means the device must stop the action. A
342+
* device renews a call of a run bound to it under its execution token. The chat view renews an
343+
* import it claimed on an unbound run: only the session that claimed it, while it runs.
344+
*/
333345
export const renewDesktopToolLease = defineAuthorizedCredentialUserUseCase({
334-
// permission-group-exempt: extends only a lease this device's token already holds.
346+
// permission-group-exempt: extends only a lease this device's token, or this session, already holds.
335347
operation: defineOperation({
336348
id: 'desktop.executor.calls.renew',
337349
principalKinds: ['session'],
@@ -342,8 +354,24 @@ export const renewDesktopToolLease = defineAuthorizedCredentialUserUseCase({
342354
input,
343355
}: {
344356
principal: SessionPrincipal
345-
input: DesktopCallTokenInput
357+
input: DesktopCallTokenInput | ChatViewCallInput
346358
}) {
359+
if ('chatView' in input) {
360+
const call = await getAsyncToolCall(input.toolCallId)
361+
const renewed =
362+
call !== null &&
363+
(await renewSimToolExecutionLease(
364+
{
365+
toolCallId: call.toolCallId,
366+
runId: call.runId,
367+
userId: principal.userId,
368+
ownerToken: chatViewDesktopLeaseOwnerToken(principal.sessionId),
369+
},
370+
{ chatView: true }
371+
))
372+
if (!renewed) throw new DesktopCallRevokedError()
373+
return { renewed: true as const }
374+
}
347375
await requireBoundDevice(principal, input.deviceId)
348376
await markDesktopPresent(input.deviceId)
349377
const call = await getBoundDesktopCall(

‎apps/sim/lib/mothership/async-runs/repository.ts‎

Lines changed: 46 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -779,6 +779,12 @@ export interface DesktopToolCallClaimant {
779779
* settlement or the next turn's workbench, whatever the device does after.
780780
*/
781781
executor?: { deviceId: string; ownerToken: string }
782+
/**
783+
* Set when the chat view claims a call that runs long enough to need a lease (an import): the
784+
* claim takes the execution lease under the chat view's session token, and the chat view renews
785+
* it while the import runs. Without renewals the lease lapses as the default tool budget would.
786+
*/
787+
chatView?: { ownerToken: string }
782788
}
783789

784790
export type DesktopToolCallClaim =
@@ -795,7 +801,8 @@ export type DesktopToolCallClaim =
795801
export async function claimDesktopToolCall(
796802
claimant: DesktopToolCallClaimant
797803
): Promise<DesktopToolCallClaim> {
798-
const { executor } = claimant
804+
const { executor, chatView } = claimant
805+
const leaseOwnerToken = executor?.ownerToken ?? chatView?.ownerToken
799806
return await claimUnderRunAdmission(
800807
{ ...claimant, desktopDeviceId: executor?.deviceId },
801808
claimant.claimedBy,
@@ -810,9 +817,9 @@ export async function claimDesktopToolCall(
810817
claimedBy: claimant.claimedBy,
811818
claimedAt,
812819
updatedAt: claimedAt,
813-
...(executor
820+
...(leaseOwnerToken
814821
? {
815-
executionOwnerToken: executor.ownerToken,
822+
executionOwnerToken: leaseOwnerToken,
816823
executionLeaseExpiresAt: sql`clock_timestamp() + ${SIM_TOOL_EXECUTION_LEASE_SECONDS} * interval '1 second'`,
817824
}
818825
: {}),
@@ -868,7 +875,7 @@ export async function claimDesktopToolCall(
868875
*/
869876
export async function renewSimToolExecutionLease(
870877
owner: SimToolExecutionOwner,
871-
desktop?: { deviceId: string }
878+
desktop?: { deviceId: string } | { chatView: true }
872879
): Promise<boolean> {
873880
const [renewed] = await db
874881
.update(copilotAsyncToolCalls)
@@ -887,7 +894,12 @@ export async function renewSimToolExecutionLease(
887894
desktop
888895
? and(
889896
eq(copilotAsyncToolCalls.status, ASYNC_TOOL_STATUS.running),
890-
sql`EXISTS (SELECT 1 FROM ${copilotRuns} r WHERE r.id = ${copilotAsyncToolCalls.runId} AND r.desktop_device_id = ${desktop.deviceId})`
897+
'deviceId' in desktop
898+
? sql`EXISTS (SELECT 1 FROM ${copilotRuns} r WHERE r.id = ${copilotAsyncToolCalls.runId} AND r.desktop_device_id = ${desktop.deviceId})`
899+
: and(
900+
eq(copilotAsyncToolCalls.claimedBy, DESKTOP_TOOL_CLAIM_OWNER.files),
901+
sql`EXISTS (SELECT 1 FROM ${copilotRuns} r WHERE r.id = ${copilotAsyncToolCalls.runId} AND r.desktop_device_id IS NULL)`
902+
)
891903
)
892904
: undefined
893905
)
@@ -896,6 +908,35 @@ export async function renewSimToolExecutionLease(
896908
return !!renewed
897909
}
898910

911+
/**
912+
* How much longer the chat view's renewed lease keeps a desktop call it is running alive, in ms,
913+
* on the database clock; null once the lease lapsed or the call is not one the chat view holds.
914+
*/
915+
export async function getChatViewDesktopLeaseRemainingMs(
916+
toolCallId: string
917+
): Promise<number | null> {
918+
const [row] = await db
919+
.select({
920+
remainingMs: sql<number>`extract(epoch from (${copilotAsyncToolCalls.executionLeaseExpiresAt} - clock_timestamp())) * 1000`,
921+
})
922+
.from(copilotAsyncToolCalls)
923+
.innerJoin(copilotRuns, eq(copilotRuns.id, copilotAsyncToolCalls.runId))
924+
.where(
925+
and(
926+
eq(copilotAsyncToolCalls.toolCallId, toolCallId),
927+
eq(copilotAsyncToolCalls.status, ASYNC_TOOL_STATUS.running),
928+
eq(copilotAsyncToolCalls.claimedBy, DESKTOP_TOOL_CLAIM_OWNER.files),
929+
isNotNull(copilotAsyncToolCalls.executionOwnerToken),
930+
isNull(copilotAsyncToolCalls.executionSettledAt),
931+
isNull(copilotAsyncToolCalls.executionRevokedAt),
932+
isNull(copilotRuns.desktopDeviceId),
933+
sql`${copilotAsyncToolCalls.executionLeaseExpiresAt} > clock_timestamp()`
934+
)
935+
)
936+
.limit(1)
937+
return row ? Number(row.remainingMs) : null
938+
}
939+
899940
/** Revocation ends local execution authority; recorded remote commands remain independently unsettled. */
900941
export async function revokeExpiredSimToolExecutions(input: { runId: string; userId: string }) {
901942
return db.transaction((tx) =>

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

Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3545,6 +3545,83 @@ describe('runCopilotLifecycle', () => {
35453545
}
35463546
})
35473547

3548+
it('waits on a chat-view import while its lease is renewed, and fails it once the lease lapses', async () => {
3549+
vi.useFakeTimers()
3550+
try {
3551+
const bodies: Record<string, unknown>[] = []
3552+
mockForceFailHungToolCall.mockImplementation(
3553+
async (toolCallId: string, context: StreamingContext) => {
3554+
const tool = context.toolCalls.get(toolCallId)
3555+
if (!tool) return
3556+
tool.status = MothershipStreamV1ToolOutcome.error
3557+
tool.endTime = Date.now()
3558+
tool.result = { success: false }
3559+
tool.error = 'Tool execution hung'
3560+
}
3561+
)
3562+
// Renewed once past the default budget, then the renewals stop and the lease lapses.
3563+
mothershipAsyncRunsMockFns.mockGetChatViewDesktopLeaseRemainingMs
3564+
.mockResolvedValueOnce(50_000)
3565+
.mockResolvedValue(null)
3566+
mockRunStreamLoop.mockImplementationOnce(
3567+
async (_url: string, fetchOptions: RequestInit, context: StreamingContext) => {
3568+
bodies.push(JSON.parse(String(fetchOptions.body)))
3569+
context.toolCalls.set('tool-import', {
3570+
id: 'tool-import',
3571+
name: 'import_local_files',
3572+
status: 'executing',
3573+
})
3574+
context.pendingToolPromises.set('tool-import', new Promise(() => {}))
3575+
context.awaitingAsyncContinuation = {
3576+
checkpointId: 'ckpt-1',
3577+
pendingToolCallIds: ['tool-import'],
3578+
}
3579+
}
3580+
)
3581+
mockRunStreamLoop.mockImplementationOnce(
3582+
async (_url: string, fetchOptions: RequestInit, context: StreamingContext) => {
3583+
bodies.push(JSON.parse(String(fetchOptions.body)))
3584+
context.accumulatedContent = 'Done.'
3585+
}
3586+
)
3587+
3588+
const lifecycle = runCopilotLifecycle(
3589+
{ message: 'import', messageId: 'stream-1' },
3590+
{
3591+
userId: 'user-1',
3592+
workspaceId: 'ws-1',
3593+
chatId: 'chat-1',
3594+
executionId: 'exec-1',
3595+
runId: 'run-1',
3596+
executionContext: {
3597+
userId: 'user-1',
3598+
workflowId: '',
3599+
workspaceId: 'ws-1',
3600+
chatId: 'chat-1',
3601+
},
3602+
}
3603+
)
3604+
3605+
// Past the default budget (60 s + 30 s grace) the lease is still live: no force-fail.
3606+
await vi.advanceTimersByTimeAsync(91_000)
3607+
expect(mockForceFailHungToolCall).not.toHaveBeenCalled()
3608+
// Once the lease the renewals kept alive runs out, the import is failed as lost.
3609+
await vi.advanceTimersByTimeAsync(52_000)
3610+
const result = await lifecycle
3611+
expect(mockForceFailHungToolCall).toHaveBeenCalledWith(
3612+
'tool-import',
3613+
expect.anything(),
3614+
expect.objectContaining({ userId: 'user-1' })
3615+
)
3616+
expect(bodies[1].results).toEqual([
3617+
expect.objectContaining({ callId: 'tool-import', success: false }),
3618+
])
3619+
expect(result.success).toBe(true)
3620+
} finally {
3621+
vi.useRealTimers()
3622+
}
3623+
})
3624+
35483625
it('force-fails each hung tool on its own budget while awaiting a long approval', async () => {
35493626
vi.useFakeTimers()
35503627
try {

0 commit comments

Comments
 (0)