fix(sandbox): drop in-flight fetches once the pool gives up or the time budget passes

Author: waleedlatif1

apps/sim/lib/execution/isolated-vm.ts

@@ -190,8 +190,14 @@ interface PendingExecution {
   timeout: ReturnType<typeof setTimeout>
   ownerKey: string
   brokers?: Record<string, IsolatedVMBrokerHandler>
-  /** Set when the caller aborts. Broker dispatches and the final result stop resolving the promise. */
-  cancelled: boolean
+  /**
+   * Aborted once the pool gives up on the execution (caller abort, timeout, worker loss). Broker
+   * dispatches and the final result stop resolving the promise, and in-flight `fetch` calls are
+   * dropped, since their responses can no longer reach the isolate.
+   */
+  abort: AbortController
+  /** How long the pool waits for a result before giving up; also the longest any one `fetch` may run. */
+  hostTimeoutMs: number
   /** Number of broker calls made so far for this execution. */
   brokerCallCount: number
 }
@@ -387,7 +393,8 @@ function normalizeFetchOptions(options?: IsolatedFetchOptions): SecureFetchOptio
 async function secureFetch(
   requestId: string,
   url: string,
-  options?: IsolatedFetchOptions
+  options: IsolatedFetchOptions | undefined,
+  signal: AbortSignal
 ): Promise<string> {
   if (url.length > MAX_FETCH_URL_LENGTH) {
     return JSON.stringify({
@@ -398,7 +405,7 @@ async function secureFetch(
   try {
     const response = await secureFetchWithValidation(
       url,
-      normalizeFetchOptions(options),
+      { ...normalizeFetchOptions(options), signal },
       'fetchUrl'
     )
     const bodyResult = truncateString(await response.text(), MAX_FETCH_RESPONSE_CHARS)
@@ -758,7 +765,7 @@ function handleBrokerMessage(
     return
   }
 
-  if (pending.cancelled) {
+  if (pending.abort.signal.aborted) {
     sendResponse({ error: 'Execution cancelled' })
     return
   }
@@ -801,7 +808,7 @@ function handleBrokerMessage(
   Promise.resolve()
     .then(() => pending.runInOutboundScope(() => handler(args, brokerId)))
     .then((resultValue) => {
-      if (pending.cancelled) {
+      if (pending.abort.signal.aborted) {
         sendResponse({ error: 'Execution cancelled' })
         return
       }
@@ -864,7 +871,7 @@ function handleWorkerMessage(workerId: number, message: unknown) {
       // If the caller aborted, the outer Promise is already resolved with
       // AbortError. Still run all the bookkeeping above so pool counters stay
       // accurate; just skip re-resolving.
-      if (!pending.cancelled) {
+      if (!pending.abort.signal.aborted) {
         pending.resolve(msg.result as IsolatedVMExecutionResult)
       }
       drainQueue()
@@ -895,7 +902,7 @@ function handleWorkerMessage(workerId: number, message: unknown) {
       typeof msg.executionId === 'number'
         ? workerInfo?.pendingExecutions.get(msg.executionId)
         : undefined
-    if (!pending || pending.cancelled) {
+    if (!pending || pending.abort.signal.aborted) {
       sendFetchResponse(JSON.stringify({ error: 'Execution no longer active' }))
       return
     }
@@ -926,8 +933,14 @@ function handleWorkerMessage(workerId: number, message: unknown) {
         return
       }
     }
+    // A fetch the script left running after returning is cut off once the execution's own time
+    // budget has passed: nothing can read its response.
+    const signal = AbortSignal.any([
+      pending.abort.signal,
+      AbortSignal.timeout(pending.hostTimeoutMs),
+    ])
     pending
-      .runInOutboundScope(() => secureFetch(requestId, url, options))
+      .runInOutboundScope(() => secureFetch(requestId, url, options, signal))
       .then(sendFetchResponse)
       .catch((err) => {
         sendFetchResponse(JSON.stringify({ error: getErrorMessage(err, 'Fetch failed') }))
@@ -971,6 +984,7 @@ function cleanupWorker(workerId: number, interruptedMessage?: string) {
 
   for (const [id, pending] of workerInfo.pendingExecutions) {
     clearTimeout(pending.timeout)
+    pending.abort.abort()
     totalActiveExecutions--
     const owner = ownerStates.get(pending.ownerKey)
     if (owner) {
@@ -1316,8 +1330,11 @@ function dispatchToWorker(
     workerInfo.idleTimeout = null
   }
 
+  const abort = new AbortController()
+  const hostTimeoutMs = req.timeoutMs + 1000
   const timeout = setTimeout(() => {
     workerInfo.pendingExecutions.delete(execId)
+    abort.abort()
     // A never-settling promise escapes the isolate's own script timeout, so the worker
     // would otherwise keep the isolate and its timers alive until the worker retires.
     sendCancel(workerInfo, execId)
@@ -1345,15 +1362,16 @@ function dispatchToWorker(
       resetWorkerIdleTimeout(workerInfo.id)
     }
     drainQueue()
-  }, req.timeoutMs + 1000)
+  }, hostTimeoutMs)
 
   workerInfo.pendingExecutions.set(execId, {
     runInOutboundScope: state.runInOutboundScope,
     resolve,
     timeout,
     ownerKey: ownerState.ownerKey,
     brokers,
-    cancelled: false,
+    abort,
+    hostTimeoutMs,
     brokerCallCount: 0,
   })
   workerInfo.activeExecutions++
@@ -1687,13 +1705,13 @@ export async function executeInIsolatedVM(
           if (removed) clearTimeout(removed.queueTimeout)
           state.queueId = undefined
         }
-        // If dispatched, mark the pending entry cancelled and ask the worker to
+        // If dispatched, abort the pending entry and ask the worker to
         // dispose its isolate so the pool slot can be released. The worker will
         // emit a `result` shortly after, which runs the normal counter cleanup.
         if (state.workerId !== undefined && state.execId !== undefined) {
           const wi = workers.get(state.workerId)
           const pending = wi?.pendingExecutions.get(state.execId)
-          if (pending) pending.cancelled = true
+          pending?.abort.abort()
           if (wi) sendCancel(wi, state.execId)
         }
         resolveWithRelease({

apps/sim/scripts/test-isolated-vm-remote-e2e.ts

@@ -508,6 +508,88 @@ async function run() {
     assert.equal(next.result, 'still alive', next.error?.message)
   })
 
+  /** A server that never answers, so a fetch to it stays in flight until the app drops it. */
+  const stalledServer = async () => {
+    const reached = Promise.withResolvers<void>()
+    const dropped = Promise.withResolvers<void>()
+    const server = createHttpServer((req) => {
+      req.socket.once('close', () => dropped.resolve())
+      reached.resolve()
+    })
+    await new Promise<void>((resolve) => server.listen(0, '127.0.0.1', resolve))
+    const address = server.address()
+    assert(address && typeof address === 'object')
+    return {
+      url: `http://127.0.0.1:${address.port}/`,
+      reached: reached.promise,
+      /** Whether the app closed the request within `withinMs`. */
+      droppedSoon: (withinMs = 2000) =>
+        Promise.race([dropped.promise.then(() => true), sleep(withinMs).then(() => false)]),
+      close: () => {
+        server.closeAllConnections()
+        server.close()
+      },
+    }
+  }
+  type StalledServer = Awaited<ReturnType<typeof stalledServer>>
+
+  const checkFetchDropped = (
+    name: string,
+    endExecution: (server: StalledServer) => Promise<void>
+  ) =>
+    check(name, async () => {
+      const server = await stalledServer()
+      try {
+        await endExecution(server)
+        assert(await server.droppedSoon(), 'The app kept the fetch open')
+      } finally {
+        server.close()
+      }
+    })
+
+  await check('drops the in-flight fetch of a timed-out execution at the timeout', async () => {
+    const server = await stalledServer()
+    try {
+      // Started late in the run, so only dropping it at the timeout, not at its own time
+      // budget, closes it promptly.
+      const result = await execute(
+        `await fetch('http://127.0.0.1:${slowAddress.port}/?ms=800'); await fetch('${server.url}')`,
+        { timeoutMs: 1000 }
+      )
+      assert.equal(result.termination, 'timeout', result.error?.message)
+      assert(await server.droppedSoon(400), 'The timed-out execution kept its fetch open')
+    } finally {
+      server.close()
+    }
+  })
+
+  await checkFetchDropped('drops the in-flight fetch of a cancelled execution', async (server) => {
+    const controller = new AbortController()
+    const pending = execute(`await fetch('${server.url}')`, { signal: controller.signal })
+    await server.reached
+    controller.abort()
+    assert.equal((await pending).termination, 'cancelled')
+  })
+
+  await checkFetchDropped('drops the in-flight fetch when its worker dies', async (server) => {
+    const pending = execute(`await fetch('${server.url}')`)
+    await server.reached
+    for (const pid of workerPidsOf(host.process.pid ?? -1)) process.kill(pid, 'SIGKILL')
+    assert.equal((await pending).error?.isSystemError, true)
+  })
+
+  await checkFetchDropped(
+    'sends a fetch the script left running, then drops it after the time budget',
+    async (server) => {
+      const result = await execute(
+        `fetch('${server.url}', { method: 'POST', body: 'ping' }); return 1`,
+        { timeoutMs: 500 }
+      )
+      assert.equal(result.result, 1, result.error?.message)
+      await server.reached
+    }
+  )
+
   const runTask = (code: string) =>
     execute(code, {
       task: {