improvement(sandbox): exempt caller-consumed streams from the output retention budget

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

@@ -119,6 +119,7 @@ import { E2B_MAX_SANDBOX_LIFETIME_MS, e2bProvider } from '@/lib/execution/remote
 import {
   MAX_SANDBOX_OUTPUT_BYTES,
   MAX_SANDBOX_PROCESS_OUTPUT_BYTES,
+  MAX_SANDBOX_STREAMED_OUTPUT_TAIL_BYTES,
 } from '@/lib/execution/remote-sandbox/output-limits'
 import {
   PI_SANDBOX_MIN_LIFETIME_MS,
@@ -293,6 +294,9 @@ beforeEach(() => {
   mockGetSessionCommand.mockResolvedValue({ exitCode: 0 })
 })
 
+/** Headroom for the note `tailStreamedSandboxOutput` prepends when it truncates. */
+const TRUNCATION_NOTE_ALLOWANCE_BYTES = 256
+
 describe.each(PROVIDERS)('sandbox conformance [%s]', (provider) => {
   beforeEach(() => useProvider(provider))
 
@@ -373,6 +377,101 @@ describe.each(PROVIDERS)('sandbox conformance [%s]', (provider) => {
     expect(provider === 'e2b' ? mockE2BKill : mockDelete).toHaveBeenCalledTimes(1)
   })
 
+  it('exempts a caller-consumed stream from the budget and keeps a diagnostic tail', async () => {
+    // A Pi agent turn streams one JSONL event per step and routinely passes the retention budget
+    // while producing no oversized result — the caller parses every chunk and keeps none of it.
+    const oversized = 'x'.repeat(MAX_SANDBOX_PROCESS_OUTPUT_BYTES + 1024)
+    if (provider === 'e2b') {
+      mockE2BCommandsRun.mockImplementationOnce(async (_cmd, options) => {
+        options.onStdout(`${oversized}TAIL_MARKER`)
+        return { stdout: `${oversized}TAIL_MARKER`, stderr: '', exitCode: 0 }
+      })
+    } else {
+      mockGetSessionCommandLogs.mockImplementationOnce(
+        async (_sessionId: string, _commandId: string, onStdout: (chunk: string) => void) => {
+          onStdout(`${oversized}TAIL_MARKER`)
+        }
+      )
+      mockGetSessionCommand.mockResolvedValue({ exitCode: 0 })
+    }
+
+    let streamedBytes = 0
+    const result = await withPiSandbox({}, (runner) =>
+      runner.run('pi run', {
+        timeoutMs: 1000,
+        onStdout: (chunk) => {
+          streamedBytes += chunk.length
+        },
+      })
+    )
+
+    // Delivered in full to the caller...
+    expect(streamedBytes).toBeGreaterThan(MAX_SANDBOX_PROCESS_OUTPUT_BYTES)
+    expect(result.exitCode).toBe(0)
+    // ...but only the tail is retained. Reaching exitCode 0 at all is the point: before the
+    // exemption this threw `sandbox_output_limit_exceeded` and killed the sandbox mid-run.
+    expect(result.stdout).toContain('TAIL_MARKER')
+    // The tail plus its truncation note, NOT a multiple of it. A looser bound here passes on both
+    // providers even when one returns twice as much as the other, which is the divergence this
+    // pair exists to prevent.
+    expect(Buffer.byteLength(result.stdout)).toBeLessThanOrEqual(
+      MAX_SANDBOX_STREAMED_OUTPUT_TAIL_BYTES + TRUNCATION_NOTE_ALLOWANCE_BYTES
+    )
+  })
+
+  it('cuts the retained tail identically when a stream ends between one and two tails', async () => {
+    // The Daytona appender only collapses once the accumulator passes twice the tail, so a stream
+    // finishing inside that band is the case where the two providers can disagree.
+    const midBand = 'y'.repeat(Math.floor(MAX_SANDBOX_STREAMED_OUTPUT_TAIL_BYTES * 1.5))
+    if (provider === 'e2b') {
+      mockE2BCommandsRun.mockImplementationOnce(async (_cmd, options) => {
+        options.onStdout?.(midBand)
+        return { stdout: midBand, stderr: '', exitCode: 0 }
+      })
+    } else {
+      mockGetSessionCommandLogs.mockImplementationOnce(
+        async (_sessionId: string, _commandId: string, onStdout: (chunk: string) => void) => {
+          onStdout(midBand)
+        }
+      )
+      mockGetSessionCommand.mockResolvedValue({ exitCode: 0 })
+    }
+
+    const result = await withPiSandbox({}, (runner) =>
+      runner.run('pi run', { timeoutMs: 1000, onStdout: () => {} })
+    )
+
+    expect(Buffer.byteLength(result.stdout)).toBeLessThanOrEqual(
+      MAX_SANDBOX_STREAMED_OUTPUT_TAIL_BYTES + TRUNCATION_NOTE_ALLOWANCE_BYTES
+    )
+    expect(Buffer.byteLength(result.stdout)).toBeLessThan(Buffer.byteLength(midBand))
+  })
+
+  it('still bounds a stream the caller does not consume', async () => {
+    const oversized = 'x'.repeat(MAX_SANDBOX_PROCESS_OUTPUT_BYTES + 1)
+    if (provider === 'e2b') {
+      mockE2BCommandsRun.mockImplementationOnce(async (_cmd, options) => {
+        options.onStderr(oversized)
+      })
+    } else {
+      mockGetSessionCommandLogs.mockImplementationOnce(
+        async (
+          _sessionId: string,
+          _commandId: string,
+          _onStdout: (chunk: string) => void,
+          onStderr: (chunk: string) => void
+        ) => {
+          onStderr(oversized)
+        }
+      )
+    }
+
+    // stdout is streamed, stderr is not — the exemption is per stream, not per command.
+    await expect(
+      withPiSandbox({}, (runner) => runner.run('pi run', { timeoutMs: 1000, onStdout: () => {} }))
+    ).rejects.toMatchObject({ code: 'sandbox_output_limit_exceeded', outputKind: 'process' })
+  })
+
   it('normalizes execution errors to the same shape', async () => {
     stubShellCommand(provider, '', 'Traceback...\nValueError: boom', 1)
 

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

@@ -13,12 +13,14 @@ import {
 } from '@/lib/core/utils/stream-limits'
 import { CodeLanguage } from '@/lib/execution/languages'
 import {
+  appendStreamedSandboxOutput,
   isSandboxOutputLimitError,
   MAX_SANDBOX_OUTPUT_BYTES,
   MAX_SANDBOX_PROCESS_OUTPUT_BYTES,
   SandboxOutputFileError,
   SandboxOutputLimitError,
   SandboxProcessOutputBudget,
+  tailStreamedSandboxOutput,
 } from '@/lib/execution/remote-sandbox/output-limits'
 import type {
   CreateSandboxOptions,
@@ -246,6 +248,18 @@ class DaytonaSandboxHandle implements SandboxHandle {
     const outputBudget = new SandboxProcessOutputBudget(
       options.maxOutputBytes ?? MAX_SANDBOX_PROCESS_OUTPUT_BYTES
     )
+    // Matches the E2B adapter: the budget bounds what Sim retains, so a stream the caller consumes
+    // itself is exempt and only a diagnostic tail is kept. Per stream, so a caller that streams
+    // stdout but not stderr still has stderr fully bounded. The failover must not change behavior.
+    const retainStdout = options.onStdout === undefined
+    const retainStderr = options.onStderr === undefined
+    // The appender keeps the accumulator under twice the tail so it is not re-cut on every chunk,
+    // which leaves it anywhere in that band when the stream ends. E2B tails the value it returns,
+    // so the final cut has to happen here too or a stream finishing between one and two tails comes
+    // back longer on Daytona than on E2B — a failover divergence, which is what this adapter pair
+    // must never have.
+    const finalStdout = () => (retainStdout ? stdout : tailStreamedSandboxOutput(stdout))
+    const finalStderr = () => (retainStderr ? stderr : tailStreamedSandboxOutput(stderr))
     try {
       await this.sandbox.process.createSession(sessionId)
       sessionCreated = true
@@ -282,14 +296,17 @@ class DaytonaSandboxHandle implements SandboxHandle {
       const appendOutput = (
         chunk: string,
         append: (value: string) => void,
+        retain: boolean,
         callback?: (value: string) => void
       ) => {
-        try {
-          outputBudget.add(chunk)
-        } catch {
-          void this.kill().catch(() => {})
-          resolveOutputLimit()
-          return
+        if (retain) {
+          try {
+            outputBudget.add(chunk)
+          } catch {
+            void this.kill().catch(() => {})
+            resolveOutputLimit()
+            return
+          }
         }
         append(chunk)
         callback?.(chunk)
@@ -302,17 +319,19 @@ class DaytonaSandboxHandle implements SandboxHandle {
             appendOutput(
               chunk,
               (value) => {
-                stdout += value
+                stdout = retainStdout ? stdout + value : appendStreamedSandboxOutput(stdout, value)
               },
+              retainStdout,
               options.onStdout
             )
           },
           (chunk: string) => {
             appendOutput(
               chunk,
               (value) => {
-                stderr += value
+                stderr = retainStderr ? stderr + value : appendStreamedSandboxOutput(stderr, value)
               },
+              retainStderr,
               options.onStderr
             )
           }
@@ -355,8 +374,8 @@ class DaytonaSandboxHandle implements SandboxHandle {
       if (outcome === 'output-limit' || outputBudget.error) throw outputBudget.error
       if (outcome === 'timeout') {
         return {
-          stdout,
-          stderr: stderr || `Command timed out after ${options.timeoutMs}ms`,
+          stdout: finalStdout(),
+          stderr: finalStderr() || `Command timed out after ${options.timeoutMs}ms`,
           exitCode: 124,
           timedOut: true,
         }
@@ -365,12 +384,17 @@ class DaytonaSandboxHandle implements SandboxHandle {
         if (outputBudget.error) throw outputBudget.error
         const timedOut = isDaytonaExecutionTimeout(streamError)
         if (!timedOut) throw streamError
-        return { stdout, stderr: stderr || getErrorMessage(streamError), exitCode: 124, timedOut }
+        return {
+          stdout: finalStdout(),
+          stderr: finalStderr() || getErrorMessage(streamError),
+          exitCode: 124,
+          timedOut,
+        }
       }
 
       const finished = await this.sandbox.process.getSessionCommand(sessionId, commandId)
       const exitCode = finished.exitCode ?? 0
-      return { stdout, stderr, exitCode }
+      return { stdout: finalStdout(), stderr: finalStderr(), exitCode }
     } catch (error) {
       if (isSandboxOutputLimitError(error)) {
         void this.kill().catch(() => {})
@@ -382,10 +406,15 @@ class DaytonaSandboxHandle implements SandboxHandle {
           : new DOMException('Execution cancelled', 'AbortError')
       }
       if (isDaytonaExecutionTimeout(error)) {
-        return { stdout, stderr: stderr || getErrorMessage(error), exitCode: 124, timedOut: true }
+        return {
+          stdout: finalStdout(),
+          stderr: finalStderr() || getErrorMessage(error),
+          exitCode: 124,
+          timedOut: true,
+        }
       }
       if (operation === 'code') throw error
-      return { stdout, stderr: stderr || getErrorMessage(error), exitCode: 1 }
+      return { stdout: finalStdout(), stderr: finalStderr() || getErrorMessage(error), exitCode: 1 }
     } finally {
       if (sessionCreated) {
         try {

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

@@ -36,6 +36,7 @@ import {
   SandboxOutputFileError,
   SandboxOutputLimitError,
   SandboxProcessOutputBudget,
+  tailStreamedSandboxOutput,
 } from '@/lib/execution/remote-sandbox/output-limits'
 import {
   quoteDependency,
@@ -389,12 +390,19 @@ class E2BSandboxHandle implements SandboxHandle {
     const outputBudget = new SandboxProcessOutputBudget(
       options.maxOutputBytes ?? MAX_SANDBOX_PROCESS_OUTPUT_BYTES
     )
-    const guardOutput = (value: string, callback?: (chunk: string) => void) => {
-      try {
-        outputBudget.add(value)
-      } catch (error) {
-        void this.kill().catch(() => {})
-        throw error
+    // The budget bounds what Sim RETAINS, so a stream the caller consumes itself is exempt: it has
+    // already been delivered chunk by chunk, and only a diagnostic tail is kept. Per stream, not
+    // per command — a caller that streams stdout but not stderr still has stderr fully bounded.
+    const retainStdout = options.onStdout === undefined
+    const retainStderr = options.onStderr === undefined
+    const guardOutput = (value: string, retain: boolean, callback?: (chunk: string) => void) => {
+      if (retain) {
+        try {
+          outputBudget.add(value)
+        } catch (error) {
+          void this.kill().catch(() => {})
+          throw error
+        }
       }
       callback?.(value)
     }
@@ -405,11 +413,18 @@ class E2BSandboxHandle implements SandboxHandle {
         timeoutMs: e2bTimeoutMs(options.timeoutMs),
         ...(options.signal ? { signal: options.signal } : {}),
         ...(options.rootUser ? { user: 'root' as const } : {}),
-        onStdout: (chunk) => guardOutput(chunk, options.onStdout),
-        onStderr: (chunk) => guardOutput(chunk, options.onStderr),
+        onStdout: (chunk) => guardOutput(chunk, retainStdout, options.onStdout),
+        onStderr: (chunk) => guardOutput(chunk, retainStderr, options.onStderr),
       })
-      assertSandboxProcessOutputWithinLimit([result.stdout, result.stderr], options.maxOutputBytes)
-      return { stdout: result.stdout, stderr: result.stderr, exitCode: result.exitCode }
+      assertSandboxProcessOutputWithinLimit(
+        [retainStdout ? result.stdout : undefined, retainStderr ? result.stderr : undefined],
+        options.maxOutputBytes
+      )
+      return {
+        stdout: retainStdout ? result.stdout : tailStreamedSandboxOutput(result.stdout),
+        stderr: retainStderr ? result.stderr : tailStreamedSandboxOutput(result.stderr),
+        exitCode: result.exitCode,
+      }
     } catch (error) {
       if (outputBudget.error) throw outputBudget.error
       if (reachedE2BProviderLimit(error, this.providerLimitAtMs, options.signal)) {
@@ -432,21 +447,33 @@ class E2BSandboxHandle implements SandboxHandle {
       ) {
         throw error
       }
+      // The SDK throws on a non-zero exit, so this is the ordinary path for a failing streamed
+      // command — the same retention exemption has to apply here or a failing Pi turn still trips
+      // the budget on output the caller already consumed. `message` never streams, so it is always
+      // retained and billed.
       assertSandboxProcessOutputWithinLimit(
-        [failure.stdout, failure.stderr, failure.message],
+        [
+          retainStdout ? failure.stdout : undefined,
+          retainStderr ? failure.stderr : undefined,
+          failure.message,
+        ],
         options.maxOutputBytes
       )
+      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)) {
         return {
-          stdout: failure.stdout ?? '',
-          stderr: failure.stderr ?? failure.message ?? '',
+          stdout: failureStdout ?? '',
+          stderr: failureStderr ?? failure.message ?? '',
           exitCode: 124,
           timedOut: true,
         }
       }
       return {
-        stdout: failure.stdout ?? '',
-        stderr: failure.stderr ?? failure.message ?? getErrorMessage(error),
+        stdout: failureStdout ?? '',
+        stderr: failureStderr ?? failure.message ?? getErrorMessage(error),
         exitCode: failure.exitCode ?? 1,
       }
     }

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

@@ -7,6 +7,51 @@ export const MAX_SANDBOX_OUTPUT_BYTES = 50 * 1024 * 1024
  */
 export const MAX_SANDBOX_PROCESS_OUTPUT_BYTES = 10 * 1024 * 1024
 
+/**
+ * Diagnostic tail kept from a stream the caller consumed itself.
+ *
+ * A caller that passes `onStdout`/`onStderr` takes delivery of every chunk as it arrives, so the
+ * adapter's accumulated copy is never the result — it is only ever read back to explain a failure.
+ * Billing that copy to the retention budget kills runs whose live stream is legitimately long while
+ * producing no oversized result: a Pi agent turn emits one JSONL event per step and passes 10 MB on
+ * an ordinary session, even though the caller has already parsed every event and keeps none of it.
+ */
+export const MAX_SANDBOX_STREAMED_OUTPUT_TAIL_BYTES = 64 * 1024
+
+const STREAMED_OUTPUT_TRUNCATION_NOTE =
+  '[earlier output truncated — it was streamed to the caller]\n'
+
+/**
+ * Keeps the last {@link MAX_SANDBOX_STREAMED_OUTPUT_TAIL_BYTES} of a streamed output. The cut is
+ * advanced past any UTF-8 continuation bytes so the tail starts on a code-point boundary rather
+ * than decoding to replacement characters.
+ */
+export function tailStreamedSandboxOutput(
+  value: string | undefined,
+  limitBytes = MAX_SANDBOX_STREAMED_OUTPUT_TAIL_BYTES
+): 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')}`
+}
+
+/**
+ * Appends to a streamed-output accumulator, collapsing it back to the diagnostic tail once it grows
+ * past twice that tail. Truncating on every chunk would be quadratic over a long stream. The
+ * threshold compares UTF-16 length rather than bytes because it only decides *when* to collapse —
+ * {@link tailStreamedSandboxOutput} does the byte-exact cut.
+ */
+export function appendStreamedSandboxOutput(current: string, chunk: string): string {
+  const next = current + chunk
+  return next.length > MAX_SANDBOX_STREAMED_OUTPUT_TAIL_BYTES * 2
+    ? tailStreamedSandboxOutput(next)
+    : next
+}
+
 export const SANDBOX_OUTPUT_LIMIT_CODE = 'sandbox_output_limit_exceeded' as const
 export const SANDBOX_OUTPUT_FILE_INVALID_CODE = 'sandbox_output_file_invalid' as const