From 94d61e4ed3d5b8f78092c672a43f886b36b76482 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Mon, 2 Nov 2026 22:25:01 -0800 Subject: [PATCH] fix(sandbox,db): keep a short workbench run's output on timeout; require dev push for the Project bootstrap switch - A workbench run with a budget of 20 s or less hands the command its whole budget, so the run's own deadline can fire before the provider's timeout reply. The E2B session adapter now treats that timeout-reason abort like the provider's own timeout and returns the output it already streamed, and throwIfRunTimedOut lets only the run's deadline through, so a Stop stays a cancellation. No budget is taken from short runs. - The nullable workspace.project_id switch now also requires SIM_DEV_DB_PUSH=1, like hnswIndexOptions, so a stray bootstrap flag cannot make an ordinary push or generation drop NOT NULL. --- .../remote-sandbox/e2b-session.test.ts | 67 +++++++++++- apps/sim/lib/execution/remote-sandbox/e2b.ts | 10 +- .../sim/lib/execution/remote-sandbox/index.ts | 13 ++- .../remote-sandbox/session-sandbox.test.ts | 100 +++++++++++++++++- packages/db/dev-push.test.ts | 20 ++++ packages/db/schema.ts | 5 +- 6 files changed, 202 insertions(+), 13 deletions(-) diff --git a/apps/sim/lib/execution/remote-sandbox/e2b-session.test.ts b/apps/sim/lib/execution/remote-sandbox/e2b-session.test.ts index 62e4c016f3a..cdfbad46c1f 100644 --- a/apps/sim/lib/execution/remote-sandbox/e2b-session.test.ts +++ b/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], diff --git a/apps/sim/lib/execution/remote-sandbox/e2b.ts b/apps/sim/lib/execution/remote-sandbox/e2b.ts index 264308044ff..6d0059ccd55 100644 --- a/apps/sim/lib/execution/remote-sandbox/e2b.ts +++ b/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 ?? '') diff --git a/apps/sim/lib/execution/remote-sandbox/index.ts b/apps/sim/lib/execution/remote-sandbox/index.ts index 3b9c35f5d3d..f682f3621f1 100644 --- a/apps/sim/lib/execution/remote-sandbox/index.ts +++ b/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,6 +330,7 @@ async function throwIfRunTimedOut( } ): Promise { 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. @@ -332,6 +338,7 @@ async function throwIfRunTimedOut( 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,7 +1099,6 @@ async function executeInSandboxWithinBudget( throwIfAborted(signal) throw error } - throwIfAborted(signal) await throwIfRunTimedOut(execution, [execution.text, execution.stdout, execution.stderr], { sandbox, lease, @@ -1100,6 +1106,7 @@ async function executeInSandboxWithinBudget( timeoutMs: runTimeoutMs, signal, }) + throwIfAborted(signal) if (execution.error) { const errorMessage = `${execution.error.name}: ${execution.error.value}` @@ -1316,7 +1323,6 @@ async function executeShellInSandboxWithinBudget( throwIfAborted(signal) throw error } - throwIfAborted(signal) await throwIfRunTimedOut(result, [result.stdout, result.stderr], { sandbox, lease, @@ -1324,6 +1330,7 @@ async function executeShellInSandboxWithinBudget( timeoutMs: runTimeoutMs, signal, }) + throwIfAborted(signal) const stdout = [result.stdout, result.stderr].filter(Boolean).join('\n') diff --git a/apps/sim/lib/execution/remote-sandbox/session-sandbox.test.ts b/apps/sim/lib/execution/remote-sandbox/session-sandbox.test.ts index 0cb4b375993..4be2c186461 100644 --- a/apps/sim/lib/execution/remote-sandbox/session-sandbox.test.ts +++ b/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 + signal?: AbortSignal + timeoutMs?: number + } = {} ): Promise { + 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 { + return timedOutRun(kind, { stdout: 'rendered 1 of 3', stderr: '' }, async () => [], { + timeoutMs: 10_000, + signal: options.signal, + reply: (signal) => + new Promise((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) diff --git a/packages/db/dev-push.test.ts b/packages/db/dev-push.test.ts index cfef616d22e..8901148f33e 100644 --- a/packages/db/dev-push.test.ts +++ b/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) + }) +}) diff --git a/packages/db/schema.ts b/packages/db/schema.ts index 7a78be87da7..e2d5582dc81 100644 --- a/packages/db/schema.ts +++ b/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) : column.notNull() }