fix(copilot): interrupt pending tool waits on stop

Author: waleedlatif1

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

@@ -2063,6 +2063,73 @@ describe('runCopilotLifecycle', () => {
     }
   })
 
+  it('cancels promptly while a sequential tool promise remains unsettled', async () => {
+    vi.useFakeTimers()
+    try {
+      const controller = new AbortController()
+      const fetchUrls: string[] = []
+      let capturedContext: StreamingContext | null = null
+      const executionContext: ExecutionContext = {
+        userId: 'user-1',
+        workflowId: '',
+        workspaceId: 'ws-1',
+        chatId: 'chat-1',
+      }
+
+      mockPendingToolWaitBudgetMs.mockReturnValue(3_600_000)
+      mockRunStreamLoop.mockImplementationOnce(
+        async (fetchUrl: string, _fetchOptions: RequestInit, context: StreamingContext) => {
+          fetchUrls.push(fetchUrl)
+          capturedContext = context
+          context.toolCalls.set('tool-hung', {
+            id: 'tool-hung',
+            name: 'terminal',
+            status: 'awaiting_approval',
+          })
+          context.pendingToolPromises.set('tool-hung', new Promise(() => {}))
+          context.awaitingAsyncContinuation = {
+            checkpointId: 'ckpt-1',
+            pendingToolCallIds: ['tool-hung'],
+          }
+        }
+      )
+
+      const lifecycle = runCopilotLifecycle(
+        { message: 'hello', messageId: 'stream-aborted-tool-wait' },
+        {
+          userId: 'user-1',
+          workspaceId: 'ws-1',
+          chatId: 'chat-1',
+          executionId: 'exec-1',
+          runId: 'run-1',
+          executionContext,
+          abortSignal: controller.signal,
+        }
+      )
+
+      await vi.advanceTimersByTimeAsync(0)
+      expect(mockPendingToolWaitBudgetMs).toHaveBeenCalled()
+      controller.abort('user_stop')
+      await vi.advanceTimersByTimeAsync(0)
+      const result = await lifecycle
+
+      expect(result.success).toBe(false)
+      expect(result.cancelled).toBe(true)
+      expect(fetchUrls).toEqual(['http://mothership.test/api/copilot'])
+      expect(mockForceFailHungToolCall).not.toHaveBeenCalled()
+      expect(capturedContext?.toolCalls.get('tool-hung')).toMatchObject({
+        status: MothershipStreamV1ToolOutcome.cancelled,
+        error: 'Stopped by user',
+      })
+
+      await vi.advanceTimersByTimeAsync(3_700_000)
+      expect(mockForceFailHungToolCall).not.toHaveBeenCalled()
+      expect(fetchUrls).toEqual(['http://mothership.test/api/copilot'])
+    } finally {
+      vi.useRealTimers()
+    }
+  })
+
   it('force-fails each hung tool on its own budget while awaiting a long approval', async () => {
     vi.useFakeTimers()
     try {
@@ -2420,6 +2487,57 @@ describe('runCopilotLifecycle', () => {
     expect(result.success).toBe(true)
   })
 
+  it('cancels promptly while a per-subagent tool promise remains unsettled', async () => {
+    const controller = new AbortController()
+    const addAbortListener = vi.spyOn(controller.signal, 'addEventListener')
+    const fetchUrls: string[] = []
+    let capturedContext: StreamingContext | null = null
+    mockRunStreamLoop.mockImplementationOnce(
+      async (fetchUrl: string, _fetchOptions: RequestInit, context: StreamingContext) => {
+        fetchUrls.push(fetchUrl)
+        capturedContext = context
+        context.toolCalls.set('tool-hung', {
+          id: 'tool-hung',
+          name: 'read',
+          status: 'executing',
+        })
+        context.pendingToolPromises.set('tool-hung', new Promise(() => {}))
+        context.awaitingAsyncContinuation = {
+          checkpointId: 'cp-root',
+          pendingToolCallIds: ['tool-hung'],
+          frames: [
+            {
+              parentToolCallId: 'subagent-file',
+              parentToolName: 'file',
+              pendingToolIds: ['tool-hung'],
+              checkpointId: 'cp-file',
+            },
+          ],
+        }
+      }
+    )
+
+    const lifecycle = runCopilotLifecycle(
+      { message: 'hello', messageId: 'stream-aborted-subagent-wait' },
+      { userId: 'user-1', workspaceId: 'ws-1', abortSignal: controller.signal }
+    )
+
+    await vi.waitFor(() => {
+      expect(addAbortListener).toHaveBeenCalledWith('abort', expect.any(Function), { once: true })
+    })
+    controller.abort('user_stop')
+    const result = await lifecycle
+
+    expect(result.success).toBe(false)
+    expect(result.cancelled).toBe(true)
+    expect(fetchUrls).toEqual(['http://mothership.test/api/copilot'])
+    expect(mockForceFailHungToolCall).not.toHaveBeenCalled()
+    expect(capturedContext?.toolCalls.get('tool-hung')).toMatchObject({
+      status: MothershipStreamV1ToolOutcome.cancelled,
+      error: 'Stopped by user',
+    })
+  })
+
   it('classifies a Stop landing during a subagent fanout as cancelled', async () => {
     // Guards the trap in the fanout fix: `wasAborted` is now isolated per leg, so
     // a user Stop must still reach the turn — via the abort signal or the folded

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

@@ -538,13 +538,34 @@ export function mergeResumeLegOutputs(
   if (leg.completionStatus) context.completionStatus = leg.completionStatus
 }
 
-async function waitForToolIds(context: StreamingContext, toolIds: string[]): Promise<void> {
+async function waitForToolIds(
+  context: StreamingContext,
+  toolIds: string[],
+  abortSignal?: AbortSignal
+): Promise<boolean> {
   const promises: Promise<unknown>[] = []
   for (const id of toolIds) {
     const p = context.pendingToolPromises.get(id)
     if (p) promises.push(p)
   }
-  if (promises.length > 0) await Promise.allSettled(promises)
+  if (promises.length === 0) return true
+  if (!abortSignal) {
+    await Promise.allSettled(promises)
+    return true
+  }
+  if (abortSignal.aborted) return false
+
+  let onAbort = () => {}
+  const aborted = new Promise<false>((resolve) => {
+    onAbort = () => resolve(false)
+    abortSignal.addEventListener('abort', onAbort, { once: true })
+    if (abortSignal.aborted) onAbort()
+  })
+  try {
+    return await Promise.race([Promise.allSettled(promises).then(() => true as const), aborted])
+  } finally {
+    abortSignal.removeEventListener('abort', onAbort)
+  }
 }
 
 interface ResumeToolResult {
@@ -689,7 +710,8 @@ async function driveOneChildChain(
   for (;;) {
     if (isAborted(options, context)) return null
 
-    await waitForToolIds(context, toolIds)
+    const toolsSettled = await waitForToolIds(context, toolIds, options.abortSignal)
+    if (!toolsSettled || isAborted(options, context)) return null
     const results = collectResultsForToolIds(context, toolIds, checkpointId)
 
     const leg = makeResumeLegContext(context)
@@ -1007,7 +1029,15 @@ async function runCheckpointLoop(
           next = null
           break
         }
-        await waitForToolIds(context, next.pendingToolCallIds)
+        const toolsSettled = await waitForToolIds(
+          context,
+          next.pendingToolCallIds,
+          options.abortSignal
+        )
+        if (!toolsSettled || isAborted(options, context)) {
+          next = null
+          break
+        }
         next = await driveSubagentChains(
           next,
           context,
@@ -1018,7 +1048,10 @@ async function runCheckpointLoop(
           hostedBillingRequest
         )
       }
-      if (!next) break
+      if (!next) {
+        if (isAborted(options, context)) cancelPendingTools(context)
+        break
+      }
       continuation = next
     }
 
@@ -1049,6 +1082,7 @@ async function runCheckpointLoop(
        * sibling. Unchanged promises retain their absolute deadlines.
        */
       while (context.pendingToolPromises.size > 0) {
+        if (isAborted(options, context)) break
         const now = Date.now()
         for (const [toolCallId, watchdog] of pendingWatchdogs) {
           if (context.pendingToolPromises.get(toolCallId) !== watchdog.promise) {
@@ -1116,28 +1150,36 @@ async function runCheckpointLoop(
         })
 
         const watchdogController = new AbortController()
+        const waitSignal = options.abortSignal
+          ? AbortSignal.any([watchdogController.signal, options.abortSignal])
+          : watchdogController.signal
         try {
           const wake = await Promise.race([
             ...activeWatchdogs.map(([, watchdog]) => watchdog.settlement),
-            interruptibleSleep(
-              Math.max(0, nextDeadlineAt - Date.now()),
-              watchdogController.signal
-            ).then(() => null),
+            interruptibleSleep(Math.max(0, nextDeadlineAt - Date.now()), waitSignal).then(
+              () => null
+            ),
           ])
+          if (isAborted(options, context)) break
           if (wake && context.pendingToolPromises.get(wake.toolCallId) === wake.promise) {
             context.pendingToolPromises.delete(wake.toolCallId)
           }
         } finally {
           watchdogController.abort()
         }
       }
+      const waitWasAborted = isAborted(options, context)
       waitSpan.attributes = {
         ...waitSpan.attributes,
         waitBudgetMs: maximumWaitBudgetMs,
         timedOutCount,
-        settledInTime: timedOutCount === 0,
+        aborted: waitWasAborted,
+        settledInTime: timedOutCount === 0 && !waitWasAborted,
       }
-      context.trace.endSpan(waitSpan)
+      context.trace.endSpan(
+        waitSpan,
+        waitWasAborted ? RequestTraceV1SpanStatus.cancelled : RequestTraceV1SpanStatus.ok
+      )
     }
 
     if (isAborted(options, context)) {