fix(sandbox,db): keep a short workbench run's output on timeout; require dev push for the Project bootstrap switch

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

@@ -93,7 +93,7 @@ describe('E2B session recovery', () => {
     const sandbox = {
       sandboxId: 'retained',
       getInfo,
-      files: { write: writeFile, read: readFile },
+      files: { write: writeFile, read: readFile, remove: vi.fn().mockResolvedValue(undefined) },
       setTimeout,
       commands: {
         run: (command: string, options: { background?: boolean }) =>
@@ -644,6 +644,71 @@ describe('E2B session recovery', () => {
       stderr: 'slow frame\n',
     })
   })
+  describe('when an abort lands before E2B replies', () => {
+    /** Streams one chunk, then the abort lands while the command waits on E2B. */
+    function abortWhileWaiting(controller: AbortController, reason: Error): void {
+      runCommand.mockImplementationOnce(
+        async (_command: string, options: { onStdout?: (chunk: string) => void }) => {
+          options.onStdout?.('rendered 1 of 3\n')
+          return {
+            pid: 8,
+            wait: () =>
+              new Promise((_resolve, reject) => {
+                controller.signal.addEventListener(
+                  'abort',
+                  () => reject(new Error('stream aborted')),
+                  { once: true }
+                )
+                controller.abort(reason)
+              }),
+          }
+        }
+      )
+    }
+
+    async function run(kind: 'code' | 'command', signal: AbortSignal, oneShot = false) {
+      const sandbox = oneShot
+        ? await e2bProvider.create('mothership', {})
+        : await sessionSandbox('reconnected')
+      if (!sandbox) throw new Error('Missing sandbox')
+      return kind === 'code'
+        ? sandbox.runCode('render()', { timeoutMs: 10_000, signal })
+        : sandbox.runCommand('render', { timeoutMs: 10_000, signal })
+    }
+
+    it.each(['code', 'command'] as const)(
+      'returns what the %s printed when its own deadline stops it',
+      async (kind) => {
+        const controller = new AbortController()
+        abortWhileWaiting(controller, new DOMException('timeout', 'AbortError'))
+        const result = await run(kind, controller.signal)
+        expect(result).toMatchObject({ timedOut: true, stdout: 'rendered 1 of 3\n' })
+        expect(stopCommand).toHaveBeenCalledOnce()
+      }
+    )
+
+    it('keeps a Stop of running code a cancellation', async () => {
+      const controller = new AbortController()
+      abortWhileWaiting(controller, new Error('User Stop'))
+      await expect(run('code', controller.signal)).rejects.toThrow('stream aborted')
+      expect(stopCommand).toHaveBeenCalledOnce()
+    })
+
+    it('leaves a one-shot sandbox’s deadline to the caller', async () => {
+      const controller = new AbortController()
+      abortWhileWaiting(controller, new DOMException('timeout', 'AbortError'))
+      await expect(run('code', controller.signal, true)).rejects.toThrow('stream aborted')
+    })
+
+    it('never reports a Stop of a command as its timeout', async () => {
+      const controller = new AbortController()
+      abortWhileWaiting(controller, new Error('User Stop'))
+      const result = await run('command', controller.signal)
+      expect(result.timedOut).toBeUndefined()
+      expect(stopCommand).toHaveBeenCalledOnce()
+    })
+  })
+
   it.each([
     ['stale', SessionCliStaleError],
     ['failed', Error],

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

@@ -17,6 +17,7 @@ import {
   SANDBOX_RELEASE_GENERATION_ERROR,
 } from '@sim/utils/sandbox-references'
 import { env } from '@/lib/core/config/env'
+import { isTimeoutAbortReason } from '@/lib/core/execution-limits'
 import { recordSandboxProviderLimit } from '@/lib/core/execution-limits/metrics'
 import {
   isPayloadSizeLimitError,
@@ -749,6 +750,11 @@ class E2BSandboxHandle implements SandboxHandle {
           providerFailure: 'provider_limit',
         }
       }
+      // A short workbench run hands the command its whole budget, so the caller's deadline can
+      // stop it before E2B's own timeout replies; a Stop stays a cancellation.
+      const timedOut =
+        isE2BExecutionTimeout(error) ||
+        (this.sessionKey !== undefined && isTimeoutAbortReason(options.signal?.reason))
       // The SDK throws on non-zero exit; callers want the streams, not a throw.
       const failure = error as {
         stdout?: string
@@ -761,7 +767,7 @@ class E2BSandboxHandle implements SandboxHandle {
         failure.stdout === undefined &&
         failure.stderr === undefined &&
         failure.exitCode === undefined &&
-        !isE2BExecutionTimeout(error)
+        !timedOut
       ) {
         throw error
       }
@@ -771,7 +777,7 @@ class E2BSandboxHandle implements SandboxHandle {
       )
       const tailIfStreamed = (value: string | undefined, retain: boolean) =>
         retain || value === undefined ? value : tailStreamedSandboxOutput(value)
-      if (isE2BExecutionTimeout(error)) {
+      if (timedOut) {
         // 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 ?? '')

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

@@ -190,6 +190,11 @@ function throwIfAborted(signal?: AbortSignal): void {
   throw signal.reason instanceof Error ? signal.reason : new Error('Execution cancelled')
 }
 
+/** Throws for a cancellation; the run's own deadline passes so a timed-out run can report. */
+function throwIfCancelled(signal: AbortSignal): void {
+  if (!isTimeoutAbortReason(signal.reason)) throwIfAborted(signal)
+}
+
 function abortReason(signal: AbortSignal): Error {
   return signal.reason instanceof Error
     ? signal.reason
@@ -325,13 +330,15 @@ async function throwIfRunTimedOut(
   }
 ): Promise<void> {
   if (!result.timedOut) return
+  throwIfCancelled(run.signal)
   if (!run.lease.session) throw new DOMException('timeout', 'AbortError')
   const { outputSandboxDir, sandbox } = run
   // Half of what the call has left, so the report lands before the call's deadline.
   const listingMs = Math.min(
     TIMED_OUT_OUTPUT_LISTING_MS,
     (getRemainingExecutionMs(run.signal) ?? Number.POSITIVE_INFINITY) / 2
   )
+  // withinDeadline refuses an aborted signal, so a run its own deadline stopped lists no files.
   const listed = outputSandboxDir
     ? await withinDeadline(
         () => sandbox.listFiles(outputSandboxDir, { depth: SANDBOX_OUTPUT_DIR_MAX_DEPTH }),
@@ -340,7 +347,7 @@ async function throwIfRunTimedOut(
       ).catch(() => [])
     : []
   // A cancellation that arrives while listing stays a cancellation.
-  throwIfAborted(run.signal)
+  throwIfCancelled(run.signal)
   throw new SandboxExecutionTimeoutError(sandboxTimeoutProgress(streams, run.timeoutMs, listed))
 }
 
@@ -1092,14 +1099,14 @@ async function executeInSandboxWithinBudget(
       throwIfAborted(signal)
       throw error
     }
-    throwIfAborted(signal)
     await throwIfRunTimedOut(execution, [execution.text, execution.stdout, execution.stderr], {
       sandbox,
       lease,
       outputSandboxDir: req.outputSandboxDir,
       timeoutMs: runTimeoutMs,
       signal,
     })
+    throwIfAborted(signal)
 
     if (execution.error) {
       const errorMessage = `${execution.error.name}: ${execution.error.value}`
@@ -1316,14 +1323,14 @@ async function executeShellInSandboxWithinBudget(
       throwIfAborted(signal)
       throw error
     }
-    throwIfAborted(signal)
     await throwIfRunTimedOut(result, [result.stdout, result.stderr], {
       sandbox,
       lease,
       outputSandboxDir: req.outputSandboxDir,
       timeoutMs: runTimeoutMs,
       signal,
     })
+    throwIfAborted(signal)
 
     const stdout = [result.stdout, result.stderr].filter(Boolean).join('\n')
 

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

@@ -1046,20 +1046,36 @@ describe('session sandbox lease', () => {
       kind: 'code' | 'shell',
       streams: { stdout: string; stderr: string },
       listFiles: SandboxHandle['listFiles'],
-      acceptOutputProvenance?: SandboxSessionRequest['acceptOutputProvenance']
+      options: {
+        acceptOutputProvenance?: SandboxSessionRequest['acceptOutputProvenance']
+        /** Holds the provider's timed-out reply until it settles. */
+        reply?: (signal: AbortSignal | undefined) => Promise<void>
+        signal?: AbortSignal
+        timeoutMs?: number
+      } = {}
     ): Promise<unknown> {
+      const { acceptOutputProvenance, reply, signal, timeoutMs } = options
       const { handle } = fakeSandbox(`timeout-${kind}`)
       mockFindSessionSandbox.mockResolvedValue(handle)
-      handle.runCode = async () => ({ text: '', ...streams, timedOut: true })
-      handle.runCommand = async () => ({ ...streams, exitCode: 124, timedOut: true })
+      mockCreate.mockResolvedValue(handle)
+      handle.runCode = async (_code, runOptions) => {
+        await reply?.(runOptions.signal)
+        return { text: '', ...streams, timedOut: true }
+      }
+      handle.runCommand = async (_command, runOptions) => {
+        await reply?.(runOptions.signal)
+        return { ...streams, exitCode: 124, timedOut: true }
+      }
       handle.listFiles = listFiles
       const request = {
         ...CODE_REQUEST,
+        ...(timeoutMs ? { timeoutMs } : {}),
         session: {
           key: `timeout-${kind}`,
           ...(acceptOutputProvenance ? { acceptOutputProvenance } : {}),
         },
         outputSandboxDir: OUTPUT_DIR,
+        ...(signal ? { signal } : {}),
       }
       return (
         kind === 'code'
@@ -1195,13 +1211,87 @@ describe('session sandbox lease', () => {
       expect(isTimeoutAbortReason(error)).toBe(false)
     })
 
+    describe('when an abort lands before the provider replies', () => {
+      /**
+       * The provider never replies within the command's budget, while any abort makes the run
+       * yield what already streamed as timed out, so only the caller tells its own deadline apart
+       * from a Stop.
+       */
+      function shortRun(
+        kind: 'code' | 'shell',
+        options: { onLaunch?: () => void; signal?: AbortSignal } = {}
+      ): Promise<unknown> {
+        return timedOutRun(kind, { stdout: 'rendered 1 of 3', stderr: '' }, async () => [], {
+          timeoutMs: 10_000,
+          signal: options.signal,
+          reply: (signal) =>
+            new Promise<void>((resolve) => {
+              signal?.addEventListener('abort', () => resolve(), { once: true })
+              options.onLaunch?.()
+            }),
+        })
+      }
+
+      it.each(['code', 'shell'] as const)(
+        'reports the %s output that already streamed when its own deadline fires',
+        async (kind) => {
+          vi.useFakeTimers({ toFake: ['setTimeout', 'clearTimeout'] })
+          try {
+            const pending = shortRun(kind)
+            await vi.advanceTimersByTimeAsync(10_000)
+            const error = await pending
+            if (!(error instanceof SandboxExecutionTimeoutError)) throw error
+            expect(error.progress.output).toBe('rendered 1 of 3')
+            expect(error.progress.outputFiles).toEqual([])
+          } finally {
+            vi.useRealTimers()
+          }
+        }
+      )
+
+      it.each(['code', 'shell'] as const)('keeps a %s Stop a cancellation', async (kind) => {
+        const controller = new AbortController()
+        const error = await shortRun(kind, {
+          signal: controller.signal,
+          onLaunch: () => controller.abort(new Error('User Stop')),
+        })
+        expect(error).not.toBeInstanceOf(SandboxExecutionTimeoutError)
+        expect(getErrorMessage(error)).toBe('User Stop')
+      })
+
+      it('keeps a Stop a cancellation on a provider without workbenches', async () => {
+        const { mockResolveProvider } = remoteSandboxProviderMockFns
+        const workbenchProvider = mockResolveProvider.getMockImplementation()
+        if (!workbenchProvider) throw new Error('Missing provider mock')
+        mockResolveProvider.mockImplementation(() => ({
+          ...workbenchProvider(),
+          findSessionSandbox: undefined,
+        }))
+        try {
+          const controller = new AbortController()
+          const error = await shortRun('shell', {
+            signal: controller.signal,
+            onLaunch: () => controller.abort(new Error('User Stop')),
+          })
+          expect(isTimeoutAbortReason(error)).toBe(false)
+          expect(getErrorMessage(error)).toBe('User Stop')
+        } finally {
+          mockResolveProvider.mockImplementation(workbenchProvider)
+        }
+      })
+    })
+
     it('withholds the output when the workbench secret provenance is unavailable', async () => {
       const error = await timedOutRun(
         'shell',
         { stdout: 'token=secret-value', stderr: '' },
         async () => [],
-        async () => {
-          throw new Error('Workbench output withheld because its secret provenance is unavailable')
+        {
+          acceptOutputProvenance: async () => {
+            throw new Error(
+              'Workbench output withheld because its secret provenance is unavailable'
+            )
+          },
         }
       )
       expect(error).not.toBeInstanceOf(SandboxExecutionTimeoutError)

packages/db/dev-push.test.ts

@@ -44,3 +44,23 @@ describe('dev push index options', () => {
     expect(statements[1]).toContain('WITH (m=16,ef_construction=64)')
   })
 })
+
+describe('dev push Project bootstrap switch', () => {
+  async function importWorkspaceProjectId() {
+    vi.resetModules()
+    const { workspace } = await import('@sim/db/schema')
+    return workspace.projectId
+  }
+
+  it('keeps workspace.project_id required when only the bootstrap flag is set', async () => {
+    vi.stubEnv('SIM_DEV_DB_PUSH', undefined)
+    vi.stubEnv('SIM_DB_PUSH_PROJECT_BOOTSTRAP', '1')
+    expect((await importWorkspaceProjectId()).notNull).toBe(true)
+  })
+
+  it('pushes workspace.project_id nullable only during a development push bootstrap', async () => {
+    vi.stubEnv('SIM_DEV_DB_PUSH', '1')
+    vi.stubEnv('SIM_DB_PUSH_PROJECT_BOOTSTRAP', '1')
+    expect((await importWorkspaceProjectId()).notNull).toBe(false)
+  })
+})

packages/db/schema.ts

@@ -2174,13 +2174,14 @@ export type WorkspaceMode = (typeof workspaceModeEnum.enumValues)[number]
 /**
  * A development push that bootstraps Projects first pushes this column nullable, so the columns
  * `0032_project_membership` reads exist before it fills the column and enforces NOT NULL; the
- * push that follows uses the required column. Every other push and generation reads it required.
+ * push that follows uses the required column. Every other push and generation reads it required,
+ * including one that inherits a stray bootstrap flag outside a development push.
  */
 function workspaceProjectIdColumn() {
   const column = text('project_id').references((): AnyPgColumn => project.id, {
     onDelete: 'restrict',
   })
-  return process.env.SIM_DB_PUSH_PROJECT_BOOTSTRAP === '1'
+  return process.env.SIM_DEV_DB_PUSH === '1' && process.env.SIM_DB_PUSH_PROJECT_BOOTSTRAP === '1'
     ? (column as ReturnType<typeof column.notNull>)
     : column.notNull()
 }