fix(mothership): hold desktop leases, the resume wait and first claims to the call deadline

Author: waleedlatif1

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

@@ -54,8 +54,10 @@ import {
 } from '@/lib/desktop/executor/errors'
 import { isDesktopPresent } from '@/lib/desktop/executor/presence'
 import { claimDesktopToolCall, requestRunStop } from '@/lib/mothership/async-runs/repository'
+import { TOOL_WATCHDOG_RESUME_GRACE_MS } from '@/lib/mothership/constants'
 import { waitForToolConfirmation } from '@/lib/mothership/persistence/tool-confirm'
 import { unsealClientToolCompletion } from '@/lib/mothership/request/tools/client-completion-seal.server'
+import { settleAbandonedDesktopToolCalls } from '@/lib/mothership/request/tools/desktop-wait'
 
 const CAPABILITIES = { executor: 1, browser: true, terminal: true, localFiles: true }
 
@@ -556,6 +558,7 @@ describe.runIf(Boolean(redisUrl))('desktop background executor protocol', () =>
       const toolCallId = await pendingCall(run.runId)
       const { executionToken } = await claim(desktop.principal, desktop.deviceId, toolCallId)
       const input = { deviceId: desktop.deviceId, toolCallId, executionToken }
+      expect((await row(toolCallId)).executionDeadlineAt).toBeNull()
 
       await db
         .update(copilotAsyncToolCalls)
@@ -583,6 +586,41 @@ describe.runIf(Boolean(redisUrl))('desktop background executor protocol', () =>
       ).rejects.toBeInstanceOf(DesktopCallRevokedError)
     })
 
+    it('stops renewing a call once the resume gate gave up on its deadline, then settles it as outcome unknown', async () => {
+      const desktop = await signedInDesktop()
+      const run = await boundRun(desktop)
+      const toolCallId = await pendingCall(run.runId)
+      const setRow = (fields: Partial<typeof copilotAsyncToolCalls.$inferInsert>) =>
+        db
+          .update(copilotAsyncToolCalls)
+          .set(fields)
+          .where(eq(copilotAsyncToolCalls.toolCallId, toolCallId))
+      const deadlinePassedBy = (seconds: number) =>
+        setRow({ executionDeadlineAt: sql`clock_timestamp() - ${seconds} * interval '1 second'` })
+      await setRow({ executionDeadlineAt: sql`clock_timestamp() + interval '10 minutes'` })
+      const { executionToken } = await claim(desktop.principal, desktop.deviceId, toolCallId)
+      const input = { deviceId: desktop.deviceId, toolCallId, executionToken }
+
+      // Inside the gate's grace past the deadline, the device keeps the call.
+      await deadlinePassedBy(TOOL_WATCHDOG_RESUME_GRACE_MS / 1000 - 10)
+      await expect(
+        renewDesktopToolLease.execute({ principal: desktop.principal, input })
+      ).resolves.toEqual({ renewed: true })
+
+      await deadlinePassedBy(TOOL_WATCHDOG_RESUME_GRACE_MS / 1000 + 1)
+      await expect(
+        renewDesktopToolLease.execute({ principal: desktop.principal, input })
+      ).rejects.toBeInstanceOf(DesktopCallRevokedError)
+
+      // Unrenewed, the lease lapses, and the cleanup settles the call as it would any lapse.
+      await setRow({ executionLeaseExpiresAt: sql`clock_timestamp() - interval '1 second'` })
+      await settleAbandonedDesktopToolCalls(0)
+      expect(await row(toolCallId)).toMatchObject({
+        status: 'failed',
+        error: expect.stringContaining('It may already have taken effect'),
+      })
+    })
+
     it('after Stop, refuses the stopped calls, revokes the running one and lists it to cancel', async () => {
       const desktop = await signedInDesktop()
       const run = await boundRun(desktop)

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

@@ -55,6 +55,7 @@ import {
   isAwaitingToolPermission,
   SIM_TOOL_EXECUTION_VERSION,
 } from '@/lib/mothership/async-runs/lifecycle'
+import { TOOL_WATCHDOG_RESUME_GRACE_MS } from '@/lib/mothership/constants'
 import { TraceAttr } from '@/lib/mothership/generated/trace-attributes-v1'
 import { TraceSpan } from '@/lib/mothership/generated/trace-spans-v1'
 import {
@@ -897,6 +898,12 @@ const unfinishedSimExecution = allOf(
  */
 const tooLateToStart = sql`(${copilotAsyncToolCalls.executionDeadlineAt} IS NOT NULL AND ${copilotAsyncToolCalls.executionDeadlineAt} < clock_timestamp() + ${SIM_TOOL_MIN_START_BUDGET_MS} * interval '1 millisecond')`
 
+/**
+ * A call past its deadline plus the resume grace: the resume wait's ceiling for it (`run.ts`), past
+ * which nothing still waits for its result. A call without a deadline never matches.
+ */
+const pastResumeGiveUp = sql`(${copilotAsyncToolCalls.executionDeadlineAt} IS NOT NULL AND ${copilotAsyncToolCalls.executionDeadlineAt} + ${TOOL_WATCHDOG_RESUME_GRACE_MS} * interval '1 millisecond' <= clock_timestamp())`
+
 /**
  * A Sim execution whose owner died before recording a result: its lease lapsed, which a live
  * owner's heartbeat never lets happen. A takeover runs it again until it has started
@@ -1069,7 +1076,8 @@ export async function claimDesktopToolCall(
 /**
  * Expired ownership cannot be renewed, even before a follower has observed the expiry. With
  * `desktop`, the call must still be running on a run bound to that device, so a call Stop settled
- * cannot be renewed either.
+ * cannot be renewed either, nor one past its deadline plus the resume grace: its lease then lapses
+ * and the call settles as outcome unknown.
  */
 export async function renewSimToolExecutionLease(
   owner: SimToolExecutionOwner,
@@ -1092,6 +1100,7 @@ export async function renewSimToolExecutionLease(
         desktop
           ? and(
               eq(copilotAsyncToolCalls.status, ASYNC_TOOL_STATUS.running),
+              not(pastResumeGiveUp),
               'deviceId' in desktop
                 ? sql`EXISTS (SELECT 1 FROM ${copilotRuns} r WHERE r.id = ${copilotAsyncToolCalls.runId} AND r.desktop_device_id = ${desktop.deviceId})`
                 : and(

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

@@ -3672,6 +3672,41 @@ describe('runCopilotLifecycle', () => {
     }
   })
 
+  it('fails a call whose deadline already passed at that deadline plus grace, not a grace from now', async () => {
+    vi.useFakeTimers()
+    try {
+      mothershipAsyncRunsMockFns.mockGetChatViewDesktopLeaseRemainingMs.mockResolvedValue(50_000)
+      const turn = runImportTurn(Date.now() - 10_000)
+      await vi.advanceTimersByTimeAsync(19_000)
+      expect(turn.bodies).toHaveLength(1)
+      await vi.advanceTimersByTimeAsync(2_000)
+      expect((await turn.lifecycle).success).toBe(true)
+      expect(turn.resumedWithLostImport()).toBe(true)
+      expect(
+        mothershipAsyncRunsMockFns.mockGetChatViewDesktopLeaseRemainingMs
+      ).not.toHaveBeenCalled()
+    } finally {
+      mothershipAsyncRunsMockFns.mockGetChatViewDesktopLeaseRemainingMs.mockReset()
+      vi.useRealTimers()
+    }
+  })
+
+  it('a call without a deadline keeps its full budget and grace', async () => {
+    vi.useFakeTimers()
+    try {
+      mothershipAsyncRunsMockFns.mockGetChatViewDesktopLeaseRemainingMs.mockResolvedValue(null)
+      const turn = runImportTurn()
+      await vi.advanceTimersByTimeAsync(89_000)
+      expect(turn.bodies).toHaveLength(1)
+      await vi.advanceTimersByTimeAsync(2_000)
+      expect((await turn.lifecycle).success).toBe(true)
+      expect(turn.resumedWithLostImport()).toBe(true)
+    } finally {
+      mothershipAsyncRunsMockFns.mockGetChatViewDesktopLeaseRemainingMs.mockReset()
+      vi.useRealTimers()
+    }
+  })
+
   it('a lease lookup that fails is checked again instead of failing a renewed import', async () => {
     vi.useFakeTimers()
     try {

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

@@ -1450,21 +1450,22 @@ async function runCheckpointLoop(
           const waitBudgetMs =
             pendingToolWaitBudgetMs(toolCall, context.desktopDeviceId) +
             TOOL_WATCHDOG_RESUME_GRACE_MS
-          const ceilingAt = now + Math.max(waitBudgetMs, CLIENT_TOOL_RESULT_TIMEOUT_MS)
+          const clientCapAt = now + Math.max(waitBudgetMs, CLIENT_TOOL_RESULT_TIMEOUT_MS)
+          // Neither the wait nor a lease carries a call past the deadline the worker named for it.
+          const ceilingAt =
+            toolCall?.deadlineAt === undefined
+              ? clientCapAt
+              : Math.min(clientCapAt, toolCall.deadlineAt + TOOL_WATCHDOG_RESUME_GRACE_MS)
           pendingWatchdogs.set(toolCallId, {
             promise,
             settlement: promise.then(
               () => ({ toolCallId, promise }),
               () => ({ toolCallId, promise })
             ),
-            deadlineAt: now + waitBudgetMs,
+            deadlineAt: Math.min(now + waitBudgetMs, ceilingAt),
             waitBudgetMs,
             startedAt: now,
-            // A lease never carries a call past the deadline the worker named for it.
-            ceilingAt:
-              toolCall?.deadlineAt === undefined
-                ? ceilingAt
-                : Math.min(ceilingAt, toolCall.deadlineAt + TOOL_WATCHDOG_RESUME_GRACE_MS),
+            ceilingAt,
           })
           maximumWaitBudgetMs = Math.max(maximumWaitBudgetMs, waitBudgetMs)
         }

apps/sim/lib/mothership/request/tools/executor.test.ts

@@ -698,6 +698,67 @@ describe('executeToolAndReport provenance isolation', () => {
     expect(tool.error).toBe(TOOL_DEADLINE_BEFORE_START_MESSAGE)
   })
 
+  it('settles an unclaimed call past its start window as never run, without claiming it', async () => {
+    completePendingAsyncToolCall.mockImplementationOnce(async (input) => ({ ...input }))
+    const tool: ToolCallState = {
+      ...buildPendingToolCall(),
+      name: 'generate_audio',
+      deadlineAt: Date.now() + 30_000,
+    }
+    const completion = await executeToolAndReport(tool.id, buildStreamingContext(tool), {
+      userId: 'user-1',
+      workflowId: 'workflow-1',
+      resolvedSecretTraceRegistry: new ResolvedSecretTraceRegistry(),
+    })
+
+    expect(claimToolExecution).not.toHaveBeenCalled()
+    expect(executeTool).not.toHaveBeenCalled()
+    expect(completePendingAsyncToolCall).toHaveBeenCalledExactlyOnceWith(
+      expect.objectContaining({
+        toolCallId: tool.id,
+        status: 'failed',
+        error: TOOL_DEADLINE_BEFORE_START_MESSAGE,
+      })
+    )
+    expect(completion.status).toBe('error')
+    expect(tool.result?.output).toMatchObject({
+      error: TOOL_DEADLINE_BEFORE_START_MESSAGE,
+      notStarted: true,
+    })
+  })
+
+  it('leaves a call past its start window that is already claimed to its claim', async () => {
+    completePendingAsyncToolCall.mockResolvedValueOnce(null)
+    const tool: ToolCallState = {
+      ...buildPendingToolCall(),
+      name: 'generate_audio',
+      deadlineAt: Date.now() + 30_000,
+    }
+    await executeToolAndReport(tool.id, buildStreamingContext(tool), {
+      userId: 'user-1',
+      workflowId: 'workflow-1',
+      resolvedSecretTraceRegistry: new ResolvedSecretTraceRegistry(),
+    })
+
+    expect(claimToolExecution).toHaveBeenCalledOnce()
+    expect(executeTool).not.toHaveBeenCalled()
+    expect(tool.error).toBe(TOOL_DEADLINE_BEFORE_START_MESSAGE)
+  })
+
+  it('claims a call without a deadline without trying to settle it first', async () => {
+    executeTool.mockResolvedValue({ success: true, output: {} })
+    const tool: ToolCallState = { ...buildPendingToolCall(), name: 'generate_audio' }
+    await executeToolAndReport(tool.id, buildStreamingContext(tool), {
+      userId: 'user-1',
+      workflowId: 'workflow-1',
+      resolvedSecretTraceRegistry: new ResolvedSecretTraceRegistry(),
+    })
+
+    expect(completePendingAsyncToolCall).not.toHaveBeenCalled()
+    expect(claimToolExecution).toHaveBeenCalledOnce()
+    expect(executeTool).toHaveBeenCalledOnce()
+  })
+
   it('runs a call again only until the deadline recorded when it was first handed', async () => {
     vi.useFakeTimers()
     let signal: AbortSignal | undefined

apps/sim/lib/mothership/request/tools/executor.ts

@@ -754,6 +754,18 @@ async function executeToolAndReportInner(
   })
   /** A re-run after a takeover gets what remains of the call's first deadline, never a fresh one. */
   if (record?.executionDeadlineAt) toolCall.deadlineAt = record.executionDeadlineAt.getTime()
+  // Too late to start: an unclaimed call is settled as never run rather than claimed first, so a
+  // crash before its handler refuses to start cannot leave it recovered as an unknown outcome.
+  if (
+    toolStartWindowClosed(toolCall.deadlineAt, Date.now()) &&
+    (await failUnstartedToolCall(
+      toolCall.id,
+      context,
+      execContext,
+      TOOL_DEADLINE_BEFORE_START_MESSAGE
+    ))
+  )
+    return terminalCompletionFromToolCall(toolCall)
   if (context.runId) {
     const claim = await lifetime.claim(context.runId, execContext.userId)
     if (claim.outcome === 'closed') return settleCancelled('Run stopped before tool admission')