Skip to content

Commit 2ec4e89

Browse files
authored
fix(sandbox,db): keep a short workbench run's output on timeout; require dev push for the Project bootstrap switch (#9617)
- 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.
1 parent 32af8bc commit 2ec4e89

6 files changed

Lines changed: 202 additions & 13 deletions

File tree

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

Lines changed: 66 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -93,7 +93,7 @@ describe('E2B session recovery', () => {
9393
const sandbox = {
9494
sandboxId: 'retained',
9595
getInfo,
96-
files: { write: writeFile, read: readFile },
96+
files: { write: writeFile, read: readFile, remove: vi.fn().mockResolvedValue(undefined) },
9797
setTimeout,
9898
commands: {
9999
run: (command: string, options: { background?: boolean }) =>
@@ -644,6 +644,71 @@ describe('E2B session recovery', () => {
644644
stderr: 'slow frame\n',
645645
})
646646
})
647+
describe('when an abort lands before E2B replies', () => {
648+
/** Streams one chunk, then the abort lands while the command waits on E2B. */
649+
function abortWhileWaiting(controller: AbortController, reason: Error): void {
650+
runCommand.mockImplementationOnce(
651+
async (_command: string, options: { onStdout?: (chunk: string) => void }) => {
652+
options.onStdout?.('rendered 1 of 3\n')
653+
return {
654+
pid: 8,
655+
wait: () =>
656+
new Promise((_resolve, reject) => {
657+
controller.signal.addEventListener(
658+
'abort',
659+
() => reject(new Error('stream aborted')),
660+
{ once: true }
661+
)
662+
controller.abort(reason)
663+
}),
664+
}
665+
}
666+
)
667+
}
668+
669+
async function run(kind: 'code' | 'command', signal: AbortSignal, oneShot = false) {
670+
const sandbox = oneShot
671+
? await e2bProvider.create('mothership', {})
672+
: await sessionSandbox('reconnected')
673+
if (!sandbox) throw new Error('Missing sandbox')
674+
return kind === 'code'
675+
? sandbox.runCode('render()', { timeoutMs: 10_000, signal })
676+
: sandbox.runCommand('render', { timeoutMs: 10_000, signal })
677+
}
678+
679+
it.each(['code', 'command'] as const)(
680+
'returns what the %s printed when its own deadline stops it',
681+
async (kind) => {
682+
const controller = new AbortController()
683+
abortWhileWaiting(controller, new DOMException('timeout', 'AbortError'))
684+
const result = await run(kind, controller.signal)
685+
expect(result).toMatchObject({ timedOut: true, stdout: 'rendered 1 of 3\n' })
686+
expect(stopCommand).toHaveBeenCalledOnce()
687+
}
688+
)
689+
690+
it('keeps a Stop of running code a cancellation', async () => {
691+
const controller = new AbortController()
692+
abortWhileWaiting(controller, new Error('User Stop'))
693+
await expect(run('code', controller.signal)).rejects.toThrow('stream aborted')
694+
expect(stopCommand).toHaveBeenCalledOnce()
695+
})
696+
697+
it('leaves a one-shot sandbox’s deadline to the caller', async () => {
698+
const controller = new AbortController()
699+
abortWhileWaiting(controller, new DOMException('timeout', 'AbortError'))
700+
await expect(run('code', controller.signal, true)).rejects.toThrow('stream aborted')
701+
})
702+
703+
it('never reports a Stop of a command as its timeout', async () => {
704+
const controller = new AbortController()
705+
abortWhileWaiting(controller, new Error('User Stop'))
706+
const result = await run('command', controller.signal)
707+
expect(result.timedOut).toBeUndefined()
708+
expect(stopCommand).toHaveBeenCalledOnce()
709+
})
710+
})
711+
647712
it.each([
648713
['stale', SessionCliStaleError],
649714
['failed', Error],

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

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ import {
1717
SANDBOX_RELEASE_GENERATION_ERROR,
1818
} from '@sim/utils/sandbox-references'
1919
import { env } from '@/lib/core/config/env'
20+
import { isTimeoutAbortReason } from '@/lib/core/execution-limits'
2021
import { recordSandboxProviderLimit } from '@/lib/core/execution-limits/metrics'
2122
import {
2223
isPayloadSizeLimitError,
@@ -749,6 +750,11 @@ class E2BSandboxHandle implements SandboxHandle {
749750
providerFailure: 'provider_limit',
750751
}
751752
}
753+
// A short workbench run hands the command its whole budget, so the caller's deadline can
754+
// stop it before E2B's own timeout replies; a Stop stays a cancellation.
755+
const timedOut =
756+
isE2BExecutionTimeout(error) ||
757+
(this.sessionKey !== undefined && isTimeoutAbortReason(options.signal?.reason))
752758
// The SDK throws on non-zero exit; callers want the streams, not a throw.
753759
const failure = error as {
754760
stdout?: string
@@ -761,7 +767,7 @@ class E2BSandboxHandle implements SandboxHandle {
761767
failure.stdout === undefined &&
762768
failure.stderr === undefined &&
763769
failure.exitCode === undefined &&
764-
!isE2BExecutionTimeout(error)
770+
!timedOut
765771
) {
766772
throw error
767773
}
@@ -771,7 +777,7 @@ class E2BSandboxHandle implements SandboxHandle {
771777
)
772778
const tailIfStreamed = (value: string | undefined, retain: boolean) =>
773779
retain || value === undefined ? value : tailStreamedSandboxOutput(value)
774-
if (isE2BExecutionTimeout(error)) {
780+
if (timedOut) {
775781
// The SDK's timeout error carries no streams; the callbacks observed the run's output.
776782
const stdout = mergeObservedOutput(observedStdout, failure.stdout ?? '')
777783
const stderr = mergeObservedOutput(observedStderr, failure.stderr ?? '')

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

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -190,6 +190,11 @@ function throwIfAborted(signal?: AbortSignal): void {
190190
throw signal.reason instanceof Error ? signal.reason : new Error('Execution cancelled')
191191
}
192192

193+
/** Throws for a cancellation; the run's own deadline passes so a timed-out run can report. */
194+
function throwIfCancelled(signal: AbortSignal): void {
195+
if (!isTimeoutAbortReason(signal.reason)) throwIfAborted(signal)
196+
}
197+
193198
function abortReason(signal: AbortSignal): Error {
194199
return signal.reason instanceof Error
195200
? signal.reason
@@ -325,13 +330,15 @@ async function throwIfRunTimedOut(
325330
}
326331
): Promise<void> {
327332
if (!result.timedOut) return
333+
throwIfCancelled(run.signal)
328334
if (!run.lease.session) throw new DOMException('timeout', 'AbortError')
329335
const { outputSandboxDir, sandbox } = run
330336
// Half of what the call has left, so the report lands before the call's deadline.
331337
const listingMs = Math.min(
332338
TIMED_OUT_OUTPUT_LISTING_MS,
333339
(getRemainingExecutionMs(run.signal) ?? Number.POSITIVE_INFINITY) / 2
334340
)
341+
// withinDeadline refuses an aborted signal, so a run its own deadline stopped lists no files.
335342
const listed = outputSandboxDir
336343
? await withinDeadline(
337344
() => sandbox.listFiles(outputSandboxDir, { depth: SANDBOX_OUTPUT_DIR_MAX_DEPTH }),
@@ -340,7 +347,7 @@ async function throwIfRunTimedOut(
340347
).catch(() => [])
341348
: []
342349
// A cancellation that arrives while listing stays a cancellation.
343-
throwIfAborted(run.signal)
350+
throwIfCancelled(run.signal)
344351
throw new SandboxExecutionTimeoutError(sandboxTimeoutProgress(streams, run.timeoutMs, listed))
345352
}
346353

@@ -1092,14 +1099,14 @@ async function executeInSandboxWithinBudget(
10921099
throwIfAborted(signal)
10931100
throw error
10941101
}
1095-
throwIfAborted(signal)
10961102
await throwIfRunTimedOut(execution, [execution.text, execution.stdout, execution.stderr], {
10971103
sandbox,
10981104
lease,
10991105
outputSandboxDir: req.outputSandboxDir,
11001106
timeoutMs: runTimeoutMs,
11011107
signal,
11021108
})
1109+
throwIfAborted(signal)
11031110

11041111
if (execution.error) {
11051112
const errorMessage = `${execution.error.name}: ${execution.error.value}`
@@ -1316,14 +1323,14 @@ async function executeShellInSandboxWithinBudget(
13161323
throwIfAborted(signal)
13171324
throw error
13181325
}
1319-
throwIfAborted(signal)
13201326
await throwIfRunTimedOut(result, [result.stdout, result.stderr], {
13211327
sandbox,
13221328
lease,
13231329
outputSandboxDir: req.outputSandboxDir,
13241330
timeoutMs: runTimeoutMs,
13251331
signal,
13261332
})
1333+
throwIfAborted(signal)
13271334

13281335
const stdout = [result.stdout, result.stderr].filter(Boolean).join('\n')
13291336

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

Lines changed: 95 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1046,20 +1046,36 @@ describe('session sandbox lease', () => {
10461046
kind: 'code' | 'shell',
10471047
streams: { stdout: string; stderr: string },
10481048
listFiles: SandboxHandle['listFiles'],
1049-
acceptOutputProvenance?: SandboxSessionRequest['acceptOutputProvenance']
1049+
options: {
1050+
acceptOutputProvenance?: SandboxSessionRequest['acceptOutputProvenance']
1051+
/** Holds the provider's timed-out reply until it settles. */
1052+
reply?: (signal: AbortSignal | undefined) => Promise<void>
1053+
signal?: AbortSignal
1054+
timeoutMs?: number
1055+
} = {}
10501056
): Promise<unknown> {
1057+
const { acceptOutputProvenance, reply, signal, timeoutMs } = options
10511058
const { handle } = fakeSandbox(`timeout-${kind}`)
10521059
mockFindSessionSandbox.mockResolvedValue(handle)
1053-
handle.runCode = async () => ({ text: '', ...streams, timedOut: true })
1054-
handle.runCommand = async () => ({ ...streams, exitCode: 124, timedOut: true })
1060+
mockCreate.mockResolvedValue(handle)
1061+
handle.runCode = async (_code, runOptions) => {
1062+
await reply?.(runOptions.signal)
1063+
return { text: '', ...streams, timedOut: true }
1064+
}
1065+
handle.runCommand = async (_command, runOptions) => {
1066+
await reply?.(runOptions.signal)
1067+
return { ...streams, exitCode: 124, timedOut: true }
1068+
}
10551069
handle.listFiles = listFiles
10561070
const request = {
10571071
...CODE_REQUEST,
1072+
...(timeoutMs ? { timeoutMs } : {}),
10581073
session: {
10591074
key: `timeout-${kind}`,
10601075
...(acceptOutputProvenance ? { acceptOutputProvenance } : {}),
10611076
},
10621077
outputSandboxDir: OUTPUT_DIR,
1078+
...(signal ? { signal } : {}),
10631079
}
10641080
return (
10651081
kind === 'code'
@@ -1195,13 +1211,87 @@ describe('session sandbox lease', () => {
11951211
expect(isTimeoutAbortReason(error)).toBe(false)
11961212
})
11971213

1214+
describe('when an abort lands before the provider replies', () => {
1215+
/**
1216+
* The provider never replies within the command's budget, while any abort makes the run
1217+
* yield what already streamed as timed out, so only the caller tells its own deadline apart
1218+
* from a Stop.
1219+
*/
1220+
function shortRun(
1221+
kind: 'code' | 'shell',
1222+
options: { onLaunch?: () => void; signal?: AbortSignal } = {}
1223+
): Promise<unknown> {
1224+
return timedOutRun(kind, { stdout: 'rendered 1 of 3', stderr: '' }, async () => [], {
1225+
timeoutMs: 10_000,
1226+
signal: options.signal,
1227+
reply: (signal) =>
1228+
new Promise<void>((resolve) => {
1229+
signal?.addEventListener('abort', () => resolve(), { once: true })
1230+
options.onLaunch?.()
1231+
}),
1232+
})
1233+
}
1234+
1235+
it.each(['code', 'shell'] as const)(
1236+
'reports the %s output that already streamed when its own deadline fires',
1237+
async (kind) => {
1238+
vi.useFakeTimers({ toFake: ['setTimeout', 'clearTimeout'] })
1239+
try {
1240+
const pending = shortRun(kind)
1241+
await vi.advanceTimersByTimeAsync(10_000)
1242+
const error = await pending
1243+
if (!(error instanceof SandboxExecutionTimeoutError)) throw error
1244+
expect(error.progress.output).toBe('rendered 1 of 3')
1245+
expect(error.progress.outputFiles).toEqual([])
1246+
} finally {
1247+
vi.useRealTimers()
1248+
}
1249+
}
1250+
)
1251+
1252+
it.each(['code', 'shell'] as const)('keeps a %s Stop a cancellation', async (kind) => {
1253+
const controller = new AbortController()
1254+
const error = await shortRun(kind, {
1255+
signal: controller.signal,
1256+
onLaunch: () => controller.abort(new Error('User Stop')),
1257+
})
1258+
expect(error).not.toBeInstanceOf(SandboxExecutionTimeoutError)
1259+
expect(getErrorMessage(error)).toBe('User Stop')
1260+
})
1261+
1262+
it('keeps a Stop a cancellation on a provider without workbenches', async () => {
1263+
const { mockResolveProvider } = remoteSandboxProviderMockFns
1264+
const workbenchProvider = mockResolveProvider.getMockImplementation()
1265+
if (!workbenchProvider) throw new Error('Missing provider mock')
1266+
mockResolveProvider.mockImplementation(() => ({
1267+
...workbenchProvider(),
1268+
findSessionSandbox: undefined,
1269+
}))
1270+
try {
1271+
const controller = new AbortController()
1272+
const error = await shortRun('shell', {
1273+
signal: controller.signal,
1274+
onLaunch: () => controller.abort(new Error('User Stop')),
1275+
})
1276+
expect(isTimeoutAbortReason(error)).toBe(false)
1277+
expect(getErrorMessage(error)).toBe('User Stop')
1278+
} finally {
1279+
mockResolveProvider.mockImplementation(workbenchProvider)
1280+
}
1281+
})
1282+
})
1283+
11981284
it('withholds the output when the workbench secret provenance is unavailable', async () => {
11991285
const error = await timedOutRun(
12001286
'shell',
12011287
{ stdout: 'token=secret-value', stderr: '' },
12021288
async () => [],
1203-
async () => {
1204-
throw new Error('Workbench output withheld because its secret provenance is unavailable')
1289+
{
1290+
acceptOutputProvenance: async () => {
1291+
throw new Error(
1292+
'Workbench output withheld because its secret provenance is unavailable'
1293+
)
1294+
},
12051295
}
12061296
)
12071297
expect(error).not.toBeInstanceOf(SandboxExecutionTimeoutError)

‎packages/db/dev-push.test.ts‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -44,3 +44,23 @@ describe('dev push index options', () => {
4444
expect(statements[1]).toContain('WITH (m=16,ef_construction=64)')
4545
})
4646
})
47+
48+
describe('dev push Project bootstrap switch', () => {
49+
async function importWorkspaceProjectId() {
50+
vi.resetModules()
51+
const { workspace } = await import('@sim/db/schema')
52+
return workspace.projectId
53+
}
54+
55+
it('keeps workspace.project_id required when only the bootstrap flag is set', async () => {
56+
vi.stubEnv('SIM_DEV_DB_PUSH', undefined)
57+
vi.stubEnv('SIM_DB_PUSH_PROJECT_BOOTSTRAP', '1')
58+
expect((await importWorkspaceProjectId()).notNull).toBe(true)
59+
})
60+
61+
it('pushes workspace.project_id nullable only during a development push bootstrap', async () => {
62+
vi.stubEnv('SIM_DEV_DB_PUSH', '1')
63+
vi.stubEnv('SIM_DB_PUSH_PROJECT_BOOTSTRAP', '1')
64+
expect((await importWorkspaceProjectId()).notNull).toBe(false)
65+
})
66+
})

‎packages/db/schema.ts‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2174,13 +2174,14 @@ export type WorkspaceMode = (typeof workspaceModeEnum.enumValues)[number]
21742174
/**
21752175
* A development push that bootstraps Projects first pushes this column nullable, so the columns
21762176
* `0032_project_membership` reads exist before it fills the column and enforces NOT NULL; the
2177-
* push that follows uses the required column. Every other push and generation reads it required.
2177+
* push that follows uses the required column. Every other push and generation reads it required,
2178+
* including one that inherits a stray bootstrap flag outside a development push.
21782179
*/
21792180
function workspaceProjectIdColumn() {
21802181
const column = text('project_id').references((): AnyPgColumn => project.id, {
21812182
onDelete: 'restrict',
21822183
})
2183-
return process.env.SIM_DB_PUSH_PROJECT_BOOTSTRAP === '1'
2184+
return process.env.SIM_DEV_DB_PUSH === '1' && process.env.SIM_DB_PUSH_PROJECT_BOOTSTRAP === '1'
21842185
? (column as ReturnType<typeof column.notNull>)
21852186
: column.notNull()
21862187
}

0 commit comments

Comments
 (0)