fix(sandbox): keep a cancellation during the timed-out output listing a cancellation

Author: waleedlatif1

apps/sim/lib/execution/remote-sandbox/index.ts

@@ -279,6 +279,7 @@ async function throwIfRunTimedOut(
     lease: SandboxLease
     outputSandboxDir: string | undefined
     timeoutMs: number
+    signal: AbortSignal
   }
 ): Promise<void> {
   if (!result.timedOut) return
@@ -287,9 +288,12 @@ async function throwIfRunTimedOut(
   const listed = outputSandboxDir
     ? await withinDeadline(
         () => sandbox.listFiles(outputSandboxDir, { depth: SANDBOX_OUTPUT_DIR_MAX_DEPTH }),
-        Date.now() + TIMED_OUT_OUTPUT_LISTING_MS
+        Date.now() + TIMED_OUT_OUTPUT_LISTING_MS,
+        run.signal
       ).catch(() => [])
     : []
+  // A cancellation that arrives while listing stays a cancellation.
+  throwIfAborted(run.signal)
   throw new SandboxExecutionTimeoutError(sandboxTimeoutProgress(streams, run.timeoutMs, listed))
 }
 
@@ -1025,6 +1029,7 @@ async function executeInSandboxWithinBudget(
       lease,
       outputSandboxDir: req.outputSandboxDir,
       timeoutMs: runTimeoutMs,
+      signal,
     })
 
     if (execution.error) {
@@ -1236,6 +1241,7 @@ async function executeShellInSandboxWithinBudget(
       lease,
       outputSandboxDir: req.outputSandboxDir,
       timeoutMs: runTimeoutMs,
+      signal,
     })
 
     const stdout = [result.stdout, result.stderr].filter(Boolean).join('\n')

apps/sim/lib/execution/remote-sandbox/session-sandbox.test.ts

@@ -1038,6 +1038,34 @@ describe('session sandbox lease', () => {
       }
     })
 
+    it('keeps a cancellation that arrives while listing the outputs a cancellation', async () => {
+      const { handle } = fakeSandbox('timeout-cancelled')
+      mockFindSessionSandbox.mockResolvedValue(handle)
+      const controller = new AbortController()
+      handle.runCommand = async () => ({
+        stdout: 'step 1',
+        stderr: '',
+        exitCode: 124,
+        timedOut: true,
+      })
+      handle.listFiles = async () => {
+        controller.abort(new Error('User Stop'))
+        return new Promise<never>(() => {})
+      }
+      const error = await executeShellInSandbox({
+        ...CODE_REQUEST,
+        envs: {},
+        session: { key: 'timeout-cancelled' },
+        outputSandboxDir: OUTPUT_DIR,
+        signal: controller.signal,
+      }).then(
+        () => undefined,
+        (rejection: unknown) => rejection
+      )
+      expect(error).not.toBeInstanceOf(SandboxExecutionTimeoutError)
+      expect(isTimeoutAbortReason(error)).toBe(false)
+    })
+
     it('withholds the output when the workbench secret provenance is unavailable', async () => {
       const error = await timedOutRun(
         'shell',