improvement(sandbox): return a timed-out workbench run's output and the files it already wrote (#9596)

* improvement(sandbox): return a timed-out workbench run's output and the files it already wrote

A run_code call killed at its own timeout returned an empty stdout and only
"timed out after Nms", so models reran long jobs from scratch. E2B's SDK throws
a stream-less TimeoutError, which dropped the output the stream callbacks had
already observed.

- E2B adapter: a timed-out command returns the stdout/stderr its callbacks observed
- Workbench runs (session sandboxes) that hit their own timeout throw a timeout
  abort carrying the output tail (64 KB) and the files already in the call's
  output directory; every existing timeout check still accepts it
- The route returns that output in stdout with a 408 and says the workbench and
  its files persist, so the next call can rerun only the remaining work
- Output goes through the same session provenance check and secret-name
  reporting as any other result; one-shot sandboxes and billing are unchanged

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

Author: waleedlatif1

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

@@ -613,6 +613,36 @@ describe('E2B session recovery', () => {
     expect(stopCommand).toHaveBeenCalledOnce()
     expect(killSandbox).not.toHaveBeenCalled()
   })
+
+  it('returns what a session command printed before its own timeout stopped it', async () => {
+    runCommand.mockImplementation(
+      async (
+        _command: string,
+        options: { onStdout?: (chunk: string) => void; onStderr?: (chunk: string) => void }
+      ) => {
+        options.onStdout?.('rendered 1 of 3\n')
+        options.onStderr?.('slow frame\n')
+        options.onStdout?.('rendered 2 of 3\n')
+        return {
+          pid: 7,
+          wait: async () => {
+            throw Object.assign(
+              new Error("deadline_exceeded: This error is likely due to exceeding 'timeoutMs'"),
+              { name: 'TimeoutError' }
+            )
+          },
+        }
+      }
+    )
+    const sandbox = await sessionSandbox('reconnected')
+    const result = await sandbox?.runCommand('render', { timeoutMs: 1000 })
+    expect(result).toMatchObject({
+      timedOut: true,
+      exitCode: 124,
+      stdout: 'rendered 1 of 3\nrendered 2 of 3\n',
+      stderr: 'slow frame\n',
+    })
+  })
 })
 
 describe('E2B session lease', () => {

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

@@ -373,6 +373,14 @@ export async function stopE2BSessionProcess(
   if (sandbox) await stopSessionCommand(sandbox, process.id, signal)
 }
 
+/** Joins what the stream callbacks observed with the streams the SDK returned, without repeating either. */
+function mergeObservedOutput(observed: string, returned: string): string {
+  if (!observed) return returned
+  if (!returned || observed.endsWith(returned)) return observed
+  if (returned.startsWith(observed)) return returned
+  return observed + returned
+}
+
 class E2BSandboxHandle implements SandboxHandle {
   private killed = false
   private killPromise: Promise<void> | null = null
@@ -676,12 +684,6 @@ class E2BSandboxHandle implements SandboxHandle {
         // process handle because `background: true` is fixed above.
         result = started
       }
-      const mergeObservedOutput = (observed: string, returned: string): string => {
-        if (!observed) return returned
-        if (!returned || observed.endsWith(returned)) return observed
-        if (returned.startsWith(observed)) return returned
-        return observed + returned
-      }
       const stdout = mergeObservedOutput(observedStdout, result.stdout ?? '')
       const stderr = mergeObservedOutput(observedStderr, result.stderr ?? '')
       assertSandboxProcessOutputWithinLimit([stdout, stderr], options.maxOutputBytes)
@@ -727,16 +729,21 @@ class E2BSandboxHandle implements SandboxHandle {
       )
       const tailIfStreamed = (value: string | undefined, retain: boolean) =>
         retain || value === undefined ? value : tailStreamedSandboxOutput(value)
-      const failureStdout = tailIfStreamed(failure.stdout, retainStdout)
-      const failureStderr = tailIfStreamed(failure.stderr, retainStderr)
       if (isE2BExecutionTimeout(error)) {
+        // The SDK's timeout error carries no streams; the callbacks observed the run's output.
+        const stdout = mergeObservedOutput(observedStdout, failure.stdout ?? '')
+        const stderr = mergeObservedOutput(observedStderr, failure.stderr ?? '')
         return {
-          stdout: failureStdout ?? '',
-          stderr: failureStderr ?? failure.message ?? '',
+          stdout: retainStdout ? stdout : tailStreamedSandboxOutput(stdout),
+          stderr:
+            (retainStderr ? stderr : tailStreamedSandboxOutput(stderr)) ||
+            `Command timed out after ${options.timeoutMs}ms`,
           exitCode: 124,
           timedOut: true,
         }
       }
+      const failureStdout = tailIfStreamed(failure.stdout, retainStdout)
+      const failureStderr = tailIfStreamed(failure.stderr, retainStderr)
       return {
         stdout: failureStdout ?? '',
         stderr: failureStderr ?? failure.message ?? getErrorMessage(error),

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

@@ -9,6 +9,7 @@ import {
   isTimeoutAbortReason,
 } from '@/lib/core/execution-limits'
 import { recordSandboxTeardownFailure } from '@/lib/core/execution-limits/metrics'
+import { withinDeadline } from '@/lib/core/utils/deadline'
 import { buildJavaScriptRuntimeBindingsSource } from '@/lib/execution/code-placeholders/javascript-runtime'
 import { SANDBOX_SYSTEM_PATH } from '@/lib/execution/remote-sandbox/cli-tools.server'
 import { describeProcessExit } from '@/lib/execution/remote-sandbox/code-failure'
@@ -65,6 +66,10 @@ import {
   recordSessionFileInput,
 } from '@/lib/execution/remote-sandbox/session-file-provenance'
 import { withSandboxSessionLock } from '@/lib/execution/remote-sandbox/session-lock'
+import {
+  SandboxExecutionTimeoutError,
+  sandboxTimeoutProgress,
+} from '@/lib/execution/remote-sandbox/timeout'
 import type {
   CreateSandboxOptions,
   SandboxCodeResult,
@@ -258,6 +263,40 @@ function throwIfSandboxTimedOut(result: { timedOut?: boolean }): void {
   if (result.timedOut) throw new DOMException('timeout', 'AbortError')
 }
 
+/** Bounds the output-directory listing a timeout report adds; a slow listing reports no files. */
+const TIMED_OUT_OUTPUT_LISTING_MS = 3_000
+
+/**
+ * A workbench run stopped by its own timeout reports what it produced, because the workbench and
+ * its files outlive the call and the next call can resume from them. A one-shot sandbox is torn
+ * down with everything in it, so it keeps the plain timeout.
+ */
+async function throwIfRunTimedOut(
+  result: { timedOut?: boolean },
+  streams: Array<string | undefined>,
+  run: {
+    sandbox: SandboxHandle
+    lease: SandboxLease
+    outputSandboxDir: string | undefined
+    timeoutMs: number
+    signal: AbortSignal
+  }
+): Promise<void> {
+  if (!result.timedOut) return
+  if (!run.lease.session) throw new DOMException('timeout', 'AbortError')
+  const { outputSandboxDir, sandbox } = run
+  const listed = outputSandboxDir
+    ? await withinDeadline(
+        () => sandbox.listFiles(outputSandboxDir, { depth: SANDBOX_OUTPUT_DIR_MAX_DEPTH }),
+        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))
+}
+
 function bindSandboxAbort(
   sandbox: SandboxHandle,
   provider: SandboxProviderId,
@@ -969,10 +1008,11 @@ async function executeInSandboxWithinBudget(
       Object.keys(privateInputFiles.environment).length > 0
 
     let execution: SandboxCodeResult
+    const runTimeoutMs = codeBudgetMs(req, signal)
     try {
       const endRun = phases.start('run')
       execution = await sandbox.runCode(code, {
-        timeoutMs: codeBudgetMs(req, signal),
+        timeoutMs: runTimeoutMs,
         javascriptPreload: buildJavaScriptRuntimeBindingsSource(req.runtimeBindings ?? []),
         maxOutputBytes: MAX_SANDBOX_PROCESS_OUTPUT_BYTES,
         signal,
@@ -984,7 +1024,13 @@ async function executeInSandboxWithinBudget(
       throw error
     }
     throwIfAborted(signal)
-    throwIfSandboxTimedOut(execution)
+    await throwIfRunTimedOut(execution, [execution.text, execution.stdout, execution.stderr], {
+      sandbox,
+      lease,
+      outputSandboxDir: req.outputSandboxDir,
+      timeoutMs: runTimeoutMs,
+      signal,
+    })
 
     if (execution.error) {
       const errorMessage = `${execution.error.name}: ${execution.error.value}`
@@ -1166,6 +1212,7 @@ async function executeShellInSandboxWithinBudget(
     endPrepare()
 
     let result: SandboxCommandResult
+    const runTimeoutMs = codeBudgetMs(req, signal)
     try {
       const endRun = phases.start('run')
       result = await sandbox.runCommand(code, {
@@ -1177,7 +1224,7 @@ async function executeShellInSandboxWithinBudget(
           ...privateInputFiles.environment,
           ...(req.session && req.outputSandboxDir ? { SIM_OUTPUT_DIR: req.outputSandboxDir } : {}),
         },
-        timeoutMs: codeBudgetMs(req, signal),
+        timeoutMs: runTimeoutMs,
         maxOutputBytes: MAX_SANDBOX_PROCESS_OUTPUT_BYTES,
         signal,
         rootUser: !lease.session,
@@ -1189,7 +1236,13 @@ async function executeShellInSandboxWithinBudget(
       throw error
     }
     throwIfAborted(signal)
-    throwIfSandboxTimedOut(result)
+    await throwIfRunTimedOut(result, [result.stdout, result.stderr], {
+      sandbox,
+      lease,
+      outputSandboxDir: req.outputSandboxDir,
+      timeoutMs: runTimeoutMs,
+      signal,
+    })
 
     const stdout = [result.stdout, result.stderr].filter(Boolean).join('\n')
 

apps/sim/lib/execution/remote-sandbox/output-limits.ts

@@ -48,15 +48,16 @@ const STREAMED_OUTPUT_TRUNCATION_NOTE =
  */
 export function tailStreamedSandboxOutput(
   value: string | undefined,
-  limitBytes = MAX_SANDBOX_STREAMED_OUTPUT_TAIL_BYTES
+  limitBytes = MAX_SANDBOX_STREAMED_OUTPUT_TAIL_BYTES,
+  truncationNote = STREAMED_OUTPUT_TRUNCATION_NOTE
 ): string {
   if (!value) return ''
   const buffer = Buffer.from(value, 'utf8')
   if (buffer.length <= limitBytes) return value
 
   let start = buffer.length - limitBytes
   while (start < buffer.length && (buffer[start] & 0xc0) === 0x80) start += 1
-  return `${STREAMED_OUTPUT_TRUNCATION_NOTE}${buffer.subarray(start).toString('utf8')}`
+  return `${truncationNote}${buffer.subarray(start).toString('utf8')}`
 }
 
 /**