fix(desktop): keep renewing an import's lease through transient failures, and retry a failed lease lookup

- The chat view stops renewing only when the server refuses the call (410)
- The resume watchdog retries a failed lease lookup for up to one lease instead of treating it as a lapse
- The lifecycle tests assert what the agent is resumed with, and when

Author: waleedlatif1

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

@@ -3545,78 +3545,99 @@ describe('runCopilotLifecycle', () => {
     }
   })
 
+  /** Runs a turn whose import never settles, recording each leg's request body. */
+  function runImportTurn() {
+    const bodies: Record<string, unknown>[] = []
+    mockForceFailHungToolCall.mockImplementation(
+      async (toolCallId: string, context: StreamingContext) => {
+        const tool = context.toolCalls.get(toolCallId)
+        if (!tool) return
+        tool.status = MothershipStreamV1ToolOutcome.error
+        tool.endTime = Date.now()
+        tool.result = { success: false }
+        tool.error = 'Tool execution hung'
+      }
+    )
+    mockRunStreamLoop.mockImplementationOnce(
+      async (_url: string, fetchOptions: RequestInit, context: StreamingContext) => {
+        bodies.push(JSON.parse(String(fetchOptions.body)))
+        context.toolCalls.set('tool-import', {
+          id: 'tool-import',
+          name: 'import_local_files',
+          status: 'executing',
+        })
+        context.pendingToolPromises.set('tool-import', new Promise(() => {}))
+        context.awaitingAsyncContinuation = {
+          checkpointId: 'ckpt-1',
+          pendingToolCallIds: ['tool-import'],
+        }
+      }
+    )
+    mockRunStreamLoop.mockImplementationOnce(
+      async (_url: string, fetchOptions: RequestInit, context: StreamingContext) => {
+        bodies.push(JSON.parse(String(fetchOptions.body)))
+        context.accumulatedContent = 'Done.'
+      }
+    )
+    const lifecycle = runCopilotLifecycle(
+      { message: 'import', messageId: 'stream-1' },
+      {
+        userId: 'user-1',
+        workspaceId: 'ws-1',
+        chatId: 'chat-1',
+        executionId: 'exec-1',
+        runId: 'run-1',
+        executionContext: {
+          userId: 'user-1',
+          workflowId: '',
+          workspaceId: 'ws-1',
+          chatId: 'chat-1',
+        },
+      }
+    )
+    /** The agent was resumed with the import given up as lost. */
+    const resumedWithLostImport = () =>
+      bodies.length === 2 &&
+      JSON.stringify(bodies[1].results).includes('"callId":"tool-import"') &&
+      JSON.stringify(bodies[1].results).includes('"success":false')
+    return { bodies, lifecycle, resumedWithLostImport }
+  }
+
   it('waits on a chat-view import while its lease is renewed, and fails it once the lease lapses', async () => {
     vi.useFakeTimers()
     try {
-      const bodies: Record<string, unknown>[] = []
-      mockForceFailHungToolCall.mockImplementation(
-        async (toolCallId: string, context: StreamingContext) => {
-          const tool = context.toolCalls.get(toolCallId)
-          if (!tool) return
-          tool.status = MothershipStreamV1ToolOutcome.error
-          tool.endTime = Date.now()
-          tool.result = { success: false }
-          tool.error = 'Tool execution hung'
-        }
-      )
       // Renewed once past the default budget, then the renewals stop and the lease lapses.
       mothershipAsyncRunsMockFns.mockGetChatViewDesktopLeaseRemainingMs
         .mockResolvedValueOnce(50_000)
         .mockResolvedValue(null)
-      mockRunStreamLoop.mockImplementationOnce(
-        async (_url: string, fetchOptions: RequestInit, context: StreamingContext) => {
-          bodies.push(JSON.parse(String(fetchOptions.body)))
-          context.toolCalls.set('tool-import', {
-            id: 'tool-import',
-            name: 'import_local_files',
-            status: 'executing',
-          })
-          context.pendingToolPromises.set('tool-import', new Promise(() => {}))
-          context.awaitingAsyncContinuation = {
-            checkpointId: 'ckpt-1',
-            pendingToolCallIds: ['tool-import'],
-          }
-        }
-      )
-      mockRunStreamLoop.mockImplementationOnce(
-        async (_url: string, fetchOptions: RequestInit, context: StreamingContext) => {
-          bodies.push(JSON.parse(String(fetchOptions.body)))
-          context.accumulatedContent = 'Done.'
-        }
-      )
-
-      const lifecycle = runCopilotLifecycle(
-        { message: 'import', messageId: 'stream-1' },
-        {
-          userId: 'user-1',
-          workspaceId: 'ws-1',
-          chatId: 'chat-1',
-          executionId: 'exec-1',
-          runId: 'run-1',
-          executionContext: {
-            userId: 'user-1',
-            workflowId: '',
-            workspaceId: 'ws-1',
-            chatId: 'chat-1',
-          },
-        }
-      )
-
-      // Past the default budget (60 s + 30 s grace) the lease is still live: no force-fail.
+      const turn = runImportTurn()
+      // Past the default budget (60 s + 30 s grace) the lease is still live: the agent waits.
       await vi.advanceTimersByTimeAsync(91_000)
-      expect(mockForceFailHungToolCall).not.toHaveBeenCalled()
-      // Once the lease the renewals kept alive runs out, the import is failed as lost.
+      expect(turn.bodies).toHaveLength(1)
+      // Once that lease runs out, the agent is resumed with the import given up as lost.
       await vi.advanceTimersByTimeAsync(52_000)
-      const result = await lifecycle
-      expect(mockForceFailHungToolCall).toHaveBeenCalledWith(
-        'tool-import',
-        expect.anything(),
-        expect.objectContaining({ userId: 'user-1' })
-      )
-      expect(bodies[1].results).toEqual([
-        expect.objectContaining({ callId: 'tool-import', success: false }),
-      ])
-      expect(result.success).toBe(true)
+      expect((await turn.lifecycle).success).toBe(true)
+      expect(turn.resumedWithLostImport()).toBe(true)
+    } finally {
+      vi.useRealTimers()
+    }
+  })
+
+  it('a lease lookup that fails is checked again instead of failing a renewed import', async () => {
+    vi.useFakeTimers()
+    try {
+      mothershipAsyncRunsMockFns.mockGetChatViewDesktopLeaseRemainingMs
+        .mockRejectedValueOnce(new Error('database unavailable'))
+        .mockResolvedValueOnce(30_000)
+        .mockResolvedValue(null)
+      const turn = runImportTurn()
+      // The failed lookup at the default budget, and its retry 5 s later, leave the import running.
+      await vi.advanceTimersByTimeAsync(97_000)
+      expect(turn.bodies).toHaveLength(1)
+      // It is given up only once the lease the retry found has lapsed.
+      await vi.advanceTimersByTimeAsync(32_000)
+      expect((await turn.lifecycle).success).toBe(true)
+      expect(turn.resumedWithLostImport()).toBe(true)
     } finally {
       vi.useRealTimers()
     }

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

@@ -15,6 +15,7 @@ import {
 } from '@/lib/billing/core/billing-attribution'
 import { env } from '@/lib/core/config/env'
 import { isCopilotToolPermissionsEnabled, isHosted } from '@/lib/core/config/env-flags'
+import { SIM_TOOL_EXECUTION_LEASE_SECONDS } from '@/lib/mothership/async-runs/execution-lease'
 import type { AsyncCompletionSignal } from '@/lib/mothership/async-runs/lifecycle'
 import {
   createRunSegment,
@@ -94,6 +95,8 @@ const logger = createLogger('CopilotLifecycle')
 
 /** After a renewed lease's end, the wait looks once more before it calls the call lost. */
 const LEASE_RECHECK_SLACK_MS = 1_000
+/** How soon a failed lease lookup is tried again. */
+const LEASE_LOOKUP_RETRY_MS = 5_000
 
 const COPILOT_MODEL_CONTENT_PROJECTION_ERROR = 'Copilot model input could not be safely projected'
 
@@ -1368,6 +1371,8 @@ async function runCheckpointLoop(
           }>
           deadlineAt: number
           waitBudgetMs: number
+          /** Until when a failing lease lookup is retried before the call is given up. */
+          leaseLookupRetryUntil?: number
         }
       >()
 
@@ -1409,15 +1414,31 @@ async function runCheckpointLoop(
         )
         // A desktop import the chat view is running renews its lease while it works: its budget
         // runs to the end of that lease, and only a lapsed lease fails it.
+        // A lookup that fails says nothing about the lease, so it is retried, for at most one lease.
         const leases = await Promise.all(
           overdueTools.map(([toolCallId]) =>
-            getChatViewDesktopLeaseRemainingMs(toolCallId).catch(() => null)
+            getChatViewDesktopLeaseRemainingMs(toolCallId).then(
+              (remainingMs) => ({ remainingMs }),
+              (error: unknown) => ({ error })
+            )
           )
         )
-        const expiredTools = overdueTools.filter(([, watchdog], index) => {
-          const remainingMs = leases[index]
-          if (typeof remainingMs !== 'number' || remainingMs <= 0) return true
-          watchdog.deadlineAt = Date.now() + remainingMs + LEASE_RECHECK_SLACK_MS
+        const expiredTools = overdueTools.filter(([toolCallId, watchdog], index) => {
+          const lease = leases[index]
+          const checkedAt = Date.now()
+          if ('error' in lease) {
+            watchdog.leaseLookupRetryUntil ??= checkedAt + SIM_TOOL_EXECUTION_LEASE_SECONDS * 1000
+            if (checkedAt >= watchdog.leaseLookupRetryUntil) return true
+            logger.warn('Could not read a pending tool call lease; checking again', {
+              toolCallId,
+              error: getErrorMessage(lease.error),
+            })
+            watchdog.deadlineAt = checkedAt + LEASE_LOOKUP_RETRY_MS
+            return false
+          }
+          watchdog.leaseLookupRetryUntil = undefined
+          if (typeof lease.remainingMs !== 'number' || lease.remainingMs <= 0) return true
+          watchdog.deadlineAt = checkedAt + lease.remainingMs + LEASE_RECHECK_SLACK_MS
           return false
         })
         if (expiredTools.length > 0) {

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

@@ -112,18 +112,23 @@ export async function importNativeFiles(
 }
 
 /**
- * Renews an import's lease every heartbeat until stopped, or until the server refuses (the call
+ * Renews an import's lease every heartbeat until stopped, or until the server refuses it (the call
  * was stopped, settled, or its lease already lapsed), the way a desktop renews a bound call.
  */
 function keepImportLeased(toolCallId: string): { stop(): void } {
   const timer = setInterval(() => {
     requestJson(renewDesktopToolLeaseContract, { body: { toolCallId, chatView: true } }).catch(
       (error) => {
-        logger.warn('Could not renew the import lease; it will lapse', {
+        // 410: the server refuses the call (stopped, settled, or its lease lapsed), so stop. Any
+        // other failure may pass: keep renewing, as the lease outlasts a couple of missed beats.
+        if (error instanceof ApiClientError && error.status === 410) {
+          clearInterval(timer)
+          return
+        }
+        logger.warn('Could not renew the import lease; trying again next beat', {
           toolCallId,
           error: getErrorMessage(error),
         })
-        clearInterval(timer)
       }
     )
   }, SIM_TOOL_EXECUTION_HEARTBEAT_MS)