fix(desktop): bound lease lookups, never fail a replaced call, and renew from the start of an import

- Each lease lookup gets 5 s (and Stop) before it counts as failed, so a stalled read cannot hold
  the wait past its deadlines
- A call replaced while its lease was read is left to its new watchdog
- The page renews from the moment it asks for the manifest; a refusal counts only once the claim
  is confirmed, and renewing stops on every exit
- Heartbeat tests check the lease a fake server holds, not request counts

Author: waleedlatif1

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

@@ -3549,6 +3549,7 @@ describe('runCopilotLifecycle', () => {
   /** Runs a turn whose import never settles, recording each leg's request body. */
   function runImportTurn() {
     const bodies: Record<string, unknown>[] = []
+    let streamContext: StreamingContext | undefined
     mockForceFailHungToolCall.mockImplementation(
       async (toolCallId: string, context: StreamingContext) => {
         const tool = context.toolCalls.get(toolCallId)
@@ -3562,6 +3563,7 @@ describe('runCopilotLifecycle', () => {
     mockRunStreamLoop.mockImplementationOnce(
       async (_url: string, fetchOptions: RequestInit, context: StreamingContext) => {
         bodies.push(JSON.parse(String(fetchOptions.body)))
+        streamContext = context
         context.toolCalls.set('tool-import', {
           id: 'tool-import',
           name: 'import_local_files',
@@ -3601,7 +3603,11 @@ describe('runCopilotLifecycle', () => {
       bodies.length === 2 &&
       JSON.stringify(bodies[1].results).includes('"callId":"tool-import"') &&
       JSON.stringify(bodies[1].results).includes('"success":false')
-    return { bodies, lifecycle, resumedWithLostImport }
+    const context = () => {
+      if (!streamContext) throw new Error('The turn did not start its stream')
+      return streamContext
+    }
+    return { bodies, lifecycle, resumedWithLostImport, context }
   }
 
   it('waits on a chat-view import while its lease is renewed, and fails it once the lease lapses', async () => {
@@ -3682,6 +3688,67 @@ describe('runCopilotLifecycle', () => {
     }
   })
 
+  it('a lease lookup that stalls counts as failed, and gives the import up after one lease', async () => {
+    vi.useFakeTimers()
+    try {
+      mothershipAsyncRunsMockFns.mockGetChatViewDesktopLeaseRemainingMs.mockImplementation(
+        () => new Promise<number | null>(() => {})
+      )
+      const turn = runImportTurn()
+      // The lookup at the default budget (90 s) stalls; each attempt is cut off after 5 s and
+      // tried again, for one lease.
+      await vi.advanceTimersByTimeAsync(150_000)
+      expect(turn.bodies).toHaveLength(1)
+      await vi.advanceTimersByTimeAsync(20_000)
+      expect((await turn.lifecycle).success).toBe(true)
+      expect(turn.resumedWithLostImport()).toBe(true)
+    } finally {
+      mothershipAsyncRunsMockFns.mockGetChatViewDesktopLeaseRemainingMs.mockReset()
+      vi.useRealTimers()
+    }
+  })
+
+  it('does not fail a call replaced while its lease was being read', async () => {
+    vi.useFakeTimers()
+    try {
+      let finishReplacement = () => {}
+      const turn = runImportTurn()
+      mothershipAsyncRunsMockFns.mockGetChatViewDesktopLeaseRemainingMs.mockImplementationOnce(
+        async () => {
+          const context = turn.context()
+          context.pendingToolPromises.set(
+            'tool-import',
+            new Promise<{ status: 'success' }>((resolve) => {
+              finishReplacement = () => {
+                const tool = context.toolCalls.get('tool-import')
+                if (tool) {
+                  tool.status = MothershipStreamV1ToolOutcome.success
+                  tool.endTime = Date.now()
+                  tool.result = { success: true, output: { imported: true } }
+                }
+                context.pendingToolPromises.delete('tool-import')
+                resolve({ status: 'success' })
+              }
+            })
+          )
+          // The old promise's lease has lapsed, but the call now belongs to the replacement.
+          return null
+        }
+      )
+      await vi.advanceTimersByTimeAsync(91_000)
+      expect(turn.bodies).toHaveLength(1)
+      // The replacement is still running: nothing has settled the call as failed.
+      expect(turn.context().toolCalls.get('tool-import')?.status).toBe('executing')
+      finishReplacement()
+      await vi.advanceTimersByTimeAsync(0)
+      expect((await turn.lifecycle).success).toBe(true)
+      expect(JSON.stringify(turn.bodies[1].results)).toContain('"success":true')
+    } finally {
+      mothershipAsyncRunsMockFns.mockGetChatViewDesktopLeaseRemainingMs.mockReset()
+      vi.useRealTimers()
+    }
+  })
+
   it('force-fails each hung tool on its own budget while awaiting a long approval', async () => {
     vi.useFakeTimers()
     try {

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

@@ -100,6 +100,32 @@ const logger = createLogger('CopilotLifecycle')
 const LEASE_RECHECK_SLACK_MS = 1_000
 /** How soon a failed lease lookup is tried again. */
 const LEASE_LOOKUP_RETRY_MS = 5_000
+/** How long one lease lookup may take before it counts as failed. */
+const LEASE_LOOKUP_TIMEOUT_MS = 5_000
+
+/** A pending call's chat-view lease, or why it could not be read within `timeoutMs`. */
+async function readLeaseWithin(
+  toolCallId: string,
+  timeoutMs: number,
+  abortSignal: AbortSignal | undefined
+): Promise<{ remainingMs: number | null } | { error: unknown }> {
+  const lookup = getChatViewDesktopLeaseRemainingMs(toolCallId).then(
+    (remainingMs) => ({ remainingMs }),
+    (error: unknown) => ({ error })
+  )
+  const giveUp = new AbortController()
+  const signal = abortSignal ? AbortSignal.any([abortSignal, giveUp.signal]) : giveUp.signal
+  try {
+    return await Promise.race([
+      lookup,
+      interruptibleSleep(timeoutMs, signal).then(() => ({
+        error: new Error(`The lease lookup took longer than ${timeoutMs} ms`),
+      })),
+    ])
+  } finally {
+    giveUp.abort()
+  }
+}
 
 const COPILOT_MODEL_CONTENT_PROJECTION_ERROR = 'Copilot model input could not be safely projected'
 
@@ -1426,15 +1452,16 @@ 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.
+        // Each lookup is bounded, so a stalled read cannot hold the wait past its deadlines or Stop.
         const leases = await Promise.all(
           overdueTools.map(([toolCallId]) =>
-            getChatViewDesktopLeaseRemainingMs(toolCallId).then(
-              (remainingMs) => ({ remainingMs }),
-              (error: unknown) => ({ error })
-            )
+            readLeaseWithin(toolCallId, LEASE_LOOKUP_TIMEOUT_MS, options.abortSignal)
           )
         )
+        if (isAborted(options, context)) break
         const expiredTools = overdueTools.filter(([toolCallId, watchdog], index) => {
+          // A call replaced while its lease was read belongs to its new watchdog.
+          if (context.pendingToolPromises.get(toolCallId) !== watchdog.promise) return false
           const lease = leases[index]
           const checkedAt = Date.now()
           if (checkedAt >= watchdog.ceilingAt) return true

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

@@ -4,6 +4,7 @@ import {
   apiClientRequestMockFns,
 } from '@sim/testing/mocks/api-client-request.mock'
 import { libDesktopMock, libDesktopMockFns } from '@sim/testing/mocks/lib-desktop.mock'
+import { sleep } from '@sim/utils/helpers'
 import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
 
 const hoisted = vi.hoisted(() => ({
@@ -154,21 +155,37 @@ it('a read the user stopped while the desktop was reading reports nothing', asyn
 })
 
 describe('an import keeps its lease while it runs', () => {
-  /** The import's upload, which the test finishes when it chooses. */
+  const LEASE_MS = 60_000
+  /**
+   * The server's side of the lease, by the rules the lease route applies: the claim (which the
+   * desktop makes as it starts building the manifest) takes a lease, and a renewal extends it only
+   * while it is live; a refused renewal answers 410.
+   */
+  let server: {
+    leaseUntil: number
+    claimed: boolean
+    stopped: boolean
+    renewalsReceived: number
+    transient: number[]
+  }
+  let scanMs: number
   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
+  const leaseLive = () => Date.now() < server.leaseUntil
+  const answer = (status: number) =>
+    Promise.reject(new ApiClientError({ message: `HTTP ${status}`, status, body: {} }))
 
   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 } }
-    )
+    server = { leaseUntil: 0, claimed: false, stopped: false, renewalsReceived: 0, transient: [] }
+    scanMs = 0
+    mocks.invoke.mockImplementation(async (request: { operation: string }) => {
+      if (request.operation !== 'manifest')
+        return { ok: true, data: { kind: 'chunk', bytes: new Uint8Array([65, 66, 67]), eof: true } }
+      server.claimed = true
+      server.leaseUntil = Date.now() + LEASE_MS
+      await sleep(scanMs)
+      return { ok: true, data: manifest }
+    })
     mocks.upload.mockImplementation(
       () =>
         new Promise((resolve) => {
@@ -177,44 +194,50 @@ describe('an import keeps its lease while it runs', () => {
     )
     mocks.json.mockImplementation(async (contract: unknown) => {
       if (contract !== renewDesktopToolLeaseContract) return { folder: { id: 'created-folder' } }
-      const answer = renewals.shift()
-      return answer ? answer() : { renewed: true }
+      server.renewalsReceived += 1
+      const transient = server.transient.shift()
+      if (transient) return answer(transient)
+      if (server.stopped || !server.claimed || !leaseLive()) return answer(410)
+      server.leaseUntil = Date.now() + LEASE_MS
+      return { 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 () => {
+  it('keeps the lease live through a slow scan and a long upload, and lets it lapse after', async () => {
+    scanMs = 70_000
     const run = executeNativeFileTool('tool', 'import_local_files')
-    await vi.advanceTimersByTimeAsync(0)
-    expect(renewalsSent()).toBe(1)
-    await vi.advanceTimersByTimeAsync(40_000)
-    expect(renewalsSent()).toBe(3)
+    await vi.advanceTimersByTimeAsync(70_000)
+    expect(leaseLive()).toBe(true)
+    await vi.advanceTimersByTimeAsync(100_000)
+    expect(leaseLive()).toBe(true)
     finishUpload()
     await run
-    await vi.advanceTimersByTimeAsync(60_000)
-    expect(renewalsSent()).toBe(3)
+    await vi.advanceTimersByTimeAsync(LEASE_MS + 1_000)
+    expect(leaseLive()).toBe(false)
   })
 
-  it('stops renewing once the server refuses the call', async () => {
-    renewals.push(refused(410))
+  it('stops renewing once the server refuses the claimed call', async () => {
     const run = executeNativeFileTool('tool', 'import_local_files')
+    await vi.advanceTimersByTimeAsync(1_000)
+    expect(leaseLive()).toBe(true)
+    server.stopped = true
+    await vi.advanceTimersByTimeAsync(20_000)
+    const received = server.renewalsReceived
     await vi.advanceTimersByTimeAsync(60_000)
-    expect(renewalsSent()).toBe(1)
+    expect(server.renewalsReceived).toBe(received)
     finishUpload()
     await run
   })
 
   it('keeps renewing through failures that may pass', async () => {
-    renewals.push(refused(401), refused(429), refused(503))
+    server.transient = [401, 429, 503]
     const run = executeNativeFileTool('tool', 'import_local_files')
-    await vi.advanceTimersByTimeAsync(60_000)
-    expect(renewalsSent()).toBe(4)
+    await vi.advanceTimersByTimeAsync(90_000)
+    expect(leaseLive()).toBe(true)
     finishUpload()
     await run
   })

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

@@ -112,19 +112,20 @@ export async function importNativeFiles(
 }
 
 /**
- * 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.
+ * Renews an import's lease from the moment its manifest is requested, every heartbeat, the way a
+ * desktop renews a bound call, so a slow directory scan cannot outlast the lease the claim took.
+ * The claim lands while the manifest is being built, so a refusal (410) counts only once the
+ * manifest confirmed the claim: then the call was stopped, settled, or its lease lapsed, and
+ * renewing stops. Any other failure may pass, and the next beat tries again.
  */
-function keepImportLeased(toolCallId: string): { stop(): void } {
+function keepImportLeased(toolCallId: string): { claimed(): void; stop(): void } {
+  let claimed = false
   let stopped = false
   const renew = () => {
     if (stopped) return
     requestJson(renewDesktopToolLeaseContract, { body: { toolCallId, chatView: true } }).catch(
       (error) => {
-        // Any other failure may pass: keep renewing, as the lease outlasts a couple of missed beats.
-        if (error instanceof ApiClientError && error.status === 410) {
+        if (claimed && error instanceof ApiClientError && error.status === 410) {
           stop()
           return
         }
@@ -141,7 +142,13 @@ function keepImportLeased(toolCallId: string): { stop(): void } {
     clearInterval(timer)
   }
   renew()
-  return { stop }
+  return {
+    claimed() {
+      claimed = true
+      renew()
+    },
+    stop,
+  }
 }
 
 /** The server claims imports before reading their manifest, preventing replayed uploads. */
@@ -166,6 +173,10 @@ export async function executeNativeFileTool(
       )
   }
   window.addEventListener('pagehide', onPageHide)
+  // An import's claim takes a lease under this session: keep it renewed while the import runs, so
+  // the turn waits for it however long it takes, and no longer than a lease once this page stops
+  // renewing (closed, crashed, or signed out).
+  const lease = toolName === 'import_local_files' ? keepImportLeased(toolCallId) : null
   try {
     const response = await invoke(
       { operation: toolName === 'read_local_file' ? 'read' : 'manifest', toolCallId },
@@ -178,17 +189,10 @@ export async function executeNativeFileTool(
     if (response.data.kind === 'chunk') throw new Error('Unexpected chunk outside an import.')
     let completion
     if (response.data.kind === 'manifest') {
-      // The import's claim took a lease under this session: keep it renewed while files transfer,
-      // so the turn waits for the import however long it takes, and no longer than a lease once
-      // this page stops renewing (closed, crashed, or signed out).
-      const lease = keepImportLeased(toolCallId)
-      try {
-        completion = localFileImportCompletion(
-          await importNativeFiles(toolCallId, response.data, signal)
-        )
-      } finally {
-        lease.stop()
-      }
+      lease?.claimed()
+      completion = localFileImportCompletion(
+        await importNativeFiles(toolCallId, response.data, signal)
+      )
     } else completion = localFileReadCompletion(response)
     // Cancelled by the user's Stop: Stop settles the call. Cancelled by signing out: nobody reports
     // it, and the server's resume watchdog settles it once its budget or lease runs out. Either way
@@ -220,6 +224,7 @@ export async function executeNativeFileTool(
     )
     settled = true
   } finally {
+    lease?.stop()
     window.removeEventListener('pagehide', onPageHide)
   }
 }