fix(desktop): cap a renewed import's wait at the client tool limit, renew at once, and report the extended wait

- A chat-view import's lease extends its wait only up to the cap every client tool has
  (CLIENT_TOOL_RESULT_TIMEOUT_MS), so an import that hangs with its page alive still settles
- The page renews the lease as soon as the import starts, then every heartbeat
- The force-fail log names an extended wait and how long it lasted; the wait span's budget
  includes the extension
- Tests for the cap, the bound on failed lease lookups, and the client heartbeat

Author: waleedlatif1

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

@@ -21,7 +21,6 @@ import {
   chatViewDesktopLeaseOwnerToken,
   getDesktopToolClaimOwner,
   isDesktopToolCall,
-  isLeasedChatViewDesktopTool,
   isLocalReadToolCall,
 } from '@/lib/mothership/tools/desktop-tools'
 
@@ -116,13 +115,14 @@ export const POST = withRouteHandler(async (request: NextRequest) => {
           { status: 409 }
         )
       if (toolCall.status !== 'pending') return alreadyStarted()
-      // The import's lease is held by this session; the chat view renews it while the import runs.
+      // An import runs as long as its files take, so its claim takes a lease this session holds:
+      // the chat view renews it while the import runs. Reads finish in seconds and take none.
       const { outcome } = await claimDesktopToolCall({
         toolCallId: toolCall.toolCallId,
         runId: toolCall.runId,
         userId,
         claimedBy: DESKTOP_TOOL_CLAIM_OWNER.files,
-        ...(principal && isLeasedChatViewDesktopTool(toolCall.toolName)
+        ...(principal
           ? { chatView: { ownerToken: chatViewDesktopLeaseOwnerToken(principal.sessionId) } }
           : {}),
       })

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

@@ -221,6 +221,7 @@ vi.mock('@/lib/mothership/request/enterprise-byok', () => ({
 
 import { resetUsageGateCache } from '@/lib/billing/core/usage-gate-cache'
 import { buildPersistedAssistantMessage } from '@/lib/mothership/chat/persisted-message'
+import { CLIENT_TOOL_RESULT_TIMEOUT_MS } from '@/lib/mothership/constants'
 import {
   MothershipStreamV1CompletionStatus,
   MothershipStreamV1ToolOutcome,
@@ -3619,6 +3620,7 @@ describe('runCopilotLifecycle', () => {
       expect((await turn.lifecycle).success).toBe(true)
       expect(turn.resumedWithLostImport()).toBe(true)
     } finally {
+      mothershipAsyncRunsMockFns.mockGetChatViewDesktopLeaseRemainingMs.mockReset()
       vi.useRealTimers()
     }
   })
@@ -3639,6 +3641,43 @@ describe('runCopilotLifecycle', () => {
       expect((await turn.lifecycle).success).toBe(true)
       expect(turn.resumedWithLostImport()).toBe(true)
     } finally {
+      mothershipAsyncRunsMockFns.mockGetChatViewDesktopLeaseRemainingMs.mockReset()
+      vi.useRealTimers()
+    }
+  })
+
+  it('gives up a renewed import at the client result cap, however long the renewals go on', async () => {
+    vi.useFakeTimers()
+    try {
+      // The page stays alive and renews, but the import itself never finishes.
+      mothershipAsyncRunsMockFns.mockGetChatViewDesktopLeaseRemainingMs.mockResolvedValue(50_000)
+      const turn = runImportTurn()
+      await vi.advanceTimersByTimeAsync(CLIENT_TOOL_RESULT_TIMEOUT_MS - 1_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('gives up an import whose lease cannot be read for a whole lease', async () => {
+    vi.useFakeTimers()
+    try {
+      mothershipAsyncRunsMockFns.mockGetChatViewDesktopLeaseRemainingMs.mockRejectedValue(
+        new Error('database unavailable')
+      )
+      const turn = runImportTurn()
+      // The first failed lookup comes at the default budget (90 s); retries go on for one lease.
+      await vi.advanceTimersByTimeAsync(145_000)
+      expect(turn.bodies).toHaveLength(1)
+      await vi.advanceTimersByTimeAsync(10_000)
+      expect((await turn.lifecycle).success).toBe(true)
+      expect(turn.resumedWithLostImport()).toBe(true)
+    } finally {
+      mothershipAsyncRunsMockFns.mockGetChatViewDesktopLeaseRemainingMs.mockReset()
       vi.useRealTimers()
     }
   })

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

@@ -22,7 +22,10 @@ import {
   getChatViewDesktopLeaseRemainingMs,
   updateRunStatus,
 } from '@/lib/mothership/async-runs/repository'
-import { TOOL_WATCHDOG_RESUME_GRACE_MS } from '@/lib/mothership/constants'
+import {
+  CLIENT_TOOL_RESULT_TIMEOUT_MS,
+  TOOL_WATCHDOG_RESUME_GRACE_MS,
+} from '@/lib/mothership/constants'
 import {
   type CopilotEnvironmentContext,
   prepareCopilotEnvironmentContext,
@@ -1371,6 +1374,12 @@ async function runCheckpointLoop(
           }>
           deadlineAt: number
           waitBudgetMs: number
+          /** When this wait began. */
+          startedAt: number
+          /** Past this, the call is given up however its lease stands: any client tool's cap. */
+          ceilingAt: number
+          /** Why the deadline was moved past the plain budget, if it was. */
+          extendedBy?: 'lease' | 'lease_lookup'
           /** Until when a failing lease lookup is retried before the call is given up. */
           leaseLookupRetryUntil?: number
         }
@@ -1403,6 +1412,8 @@ async function runCheckpointLoop(
             ),
             deadlineAt: now + waitBudgetMs,
             waitBudgetMs,
+            startedAt: now,
+            ceilingAt: now + Math.max(waitBudgetMs, CLIENT_TOOL_RESULT_TIMEOUT_MS),
           })
           maximumWaitBudgetMs = Math.max(maximumWaitBudgetMs, waitBudgetMs)
         }
@@ -1426,30 +1437,43 @@ async function runCheckpointLoop(
         const expiredTools = overdueTools.filter(([toolCallId, watchdog], index) => {
           const lease = leases[index]
           const checkedAt = Date.now()
+          if (checkedAt >= watchdog.ceilingAt) return true
+          const extendTo = (deadlineAt: number, reason: 'lease' | 'lease_lookup') => {
+            watchdog.deadlineAt = Math.min(deadlineAt, watchdog.ceilingAt)
+            watchdog.extendedBy = reason
+            maximumWaitBudgetMs = Math.max(
+              maximumWaitBudgetMs,
+              watchdog.deadlineAt - watchdog.startedAt
+            )
+            return false
+          }
           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
+            return extendTo(checkedAt + LEASE_LOOKUP_RETRY_MS, 'lease_lookup')
           }
           watchdog.leaseLookupRetryUntil = undefined
           if (typeof lease.remainingMs !== 'number' || lease.remainingMs <= 0) return true
-          watchdog.deadlineAt = checkedAt + lease.remainingMs + LEASE_RECHECK_SLACK_MS
-          return false
+          return extendTo(checkedAt + lease.remainingMs + LEASE_RECHECK_SLACK_MS, 'lease')
         })
         if (expiredTools.length > 0) {
           await Promise.all(
             expiredTools.map(async ([toolCallId, watchdog]) => {
+              const waitedMs = Date.now() - watchdog.startedAt
               logger.error(
-                'Pending tool execution exceeded its resume wait budget; force-failing',
+                watchdog.extendedBy
+                  ? 'Pending tool execution outlived its renewed lease or its cap; force-failing'
+                  : 'Pending tool execution exceeded its resume wait budget; force-failing',
                 {
                   checkpointId: continuation.checkpointId,
                   toolCallId,
                   waitBudgetMs: watchdog.waitBudgetMs,
+                  waitedMs,
+                  ...(watchdog.extendedBy ? { extendedBy: watchdog.extendedBy } : {}),
                 }
               )
               await failPendingToolCall(toolCallId, context, execContext)

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

@@ -4,7 +4,7 @@ import {
   apiClientRequestMockFns,
 } from '@sim/testing/mocks/api-client-request.mock'
 import { libDesktopMock, libDesktopMockFns } from '@sim/testing/mocks/lib-desktop.mock'
-import { beforeEach, expect, it, vi } from 'vitest'
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
 
 const hoisted = vi.hoisted(() => ({
   invoke: vi.fn(),
@@ -23,6 +23,8 @@ vi.mock('@/lib/mothership/tools/client/completion', () => ({
 }))
 
 import type { DesktopLocalFileManifest } from '@sim/desktop-bridge'
+import { ApiClientError } from '@/lib/api/client/errors'
+import { renewDesktopToolLeaseContract } from '@/lib/api/contracts/desktop-executor'
 import {
   executeNativeFileTool,
   importNativeFiles,
@@ -150,3 +152,70 @@ it('a read the user stopped while the desktop was reading reports nothing', asyn
   await executeNativeFileTool('tool', 'read_local_file', stop.signal)
   expect(mocks.complete).not.toHaveBeenCalled()
 })
+
+describe('an import keeps its lease while it runs', () => {
+  /** The import's upload, which the test finishes when it chooses. */
+  let finishUpload: () => void
+  /** What the server answers each lease renewal, in turn. */
+  let renewals: Array<() => Promise<unknown>>
+  const renewalsSent = () =>
+    mocks.json.mock.calls.filter(([contract]) => contract === renewDesktopToolLeaseContract).length
+
+  beforeEach(() => {
+    vi.useFakeTimers()
+    renewals = []
+    mocks.invoke.mockImplementation(async (request: { operation: string }) =>
+      request.operation === 'manifest'
+        ? { ok: true, data: manifest }
+        : { ok: true, data: { kind: 'chunk', bytes: new Uint8Array([65, 66, 67]), eof: true } }
+    )
+    mocks.upload.mockImplementation(
+      () =>
+        new Promise((resolve) => {
+          finishUpload = () => resolve({ id: 'saved-file', name: 'report.txt' })
+        })
+    )
+    mocks.json.mockImplementation(async (contract: unknown) => {
+      if (contract !== renewDesktopToolLeaseContract) return { folder: { id: 'created-folder' } }
+      const answer = renewals.shift()
+      return answer ? answer() : { renewed: true }
+    })
+  })
+
+  afterEach(() => {
+    vi.useRealTimers()
+  })
+
+  const refused = (status: number) => () =>
+    Promise.reject(new ApiClientError({ message: `HTTP ${status}`, status, body: {} }))
+
+  it('renews at once, then every heartbeat, and stops when the import ends', async () => {
+    const run = executeNativeFileTool('tool', 'import_local_files')
+    await vi.advanceTimersByTimeAsync(0)
+    expect(renewalsSent()).toBe(1)
+    await vi.advanceTimersByTimeAsync(40_000)
+    expect(renewalsSent()).toBe(3)
+    finishUpload()
+    await run
+    await vi.advanceTimersByTimeAsync(60_000)
+    expect(renewalsSent()).toBe(3)
+  })
+
+  it('stops renewing once the server refuses the call', async () => {
+    renewals.push(refused(410))
+    const run = executeNativeFileTool('tool', 'import_local_files')
+    await vi.advanceTimersByTimeAsync(60_000)
+    expect(renewalsSent()).toBe(1)
+    finishUpload()
+    await run
+  })
+
+  it('keeps renewing through failures that may pass', async () => {
+    renewals.push(refused(401), refused(429), refused(503))
+    const run = executeNativeFileTool('tool', 'import_local_files')
+    await vi.advanceTimersByTimeAsync(60_000)
+    expect(renewalsSent()).toBe(4)
+    finishUpload()
+    await run
+  })
+})

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

@@ -112,17 +112,20 @@ export async function importNativeFiles(
 }
 
 /**
- * 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.
+ * Renews an import's lease at once and then every heartbeat, until stopped or until the server
+ * refuses it (410: the call was stopped, settled, or its lease already lapsed), the way a desktop
+ * renews a bound call. The first renewal comes right after the claim, so a slow start cannot
+ * outlast the lease the claim took.
  */
 function keepImportLeased(toolCallId: string): { stop(): void } {
-  const timer = setInterval(() => {
+  let stopped = false
+  const renew = () => {
+    if (stopped) return
     requestJson(renewDesktopToolLeaseContract, { body: { toolCallId, chatView: true } }).catch(
       (error) => {
-        // 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.
+        // 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)
+          stop()
           return
         }
         logger.warn('Could not renew the import lease; trying again next beat', {
@@ -131,8 +134,14 @@ function keepImportLeased(toolCallId: string): { stop(): void } {
         })
       }
     )
-  }, SIM_TOOL_EXECUTION_HEARTBEAT_MS)
-  return { stop: () => clearInterval(timer) }
+  }
+  const timer = setInterval(renew, SIM_TOOL_EXECUTION_HEARTBEAT_MS)
+  const stop = () => {
+    stopped = true
+    clearInterval(timer)
+  }
+  renew()
+  return { stop }
 }
 
 /** The server claims imports before reading their manifest, preventing replayed uploads. */

apps/sim/lib/mothership/tools/desktop-tools.ts

@@ -95,12 +95,3 @@ export const STOPPED_WHILE_RUNNING_MESSAGE =
 export function chatViewDesktopLeaseOwnerToken(sessionId: string): string {
   return `chat-view:${sessionId}`
 }
-
-/**
- * Desktop calls the chat view keeps alive with a renewed lease while they run: imports, which
- * take as long as their files do (up to 1,000 entries of up to 64 MB each). Reads finish in
- * seconds and keep the default budget.
- */
-export function isLeasedChatViewDesktopTool(toolName: string): boolean {
-  return toolName === 'import_local_files'
-}