Skip to content

Commit 733f4e9

Browse files
authored
perf(sandbox): grant the run_code session lease in the reconnect instead of extra E2B round trips (#8621)
* perf(sandbox): grant the run_code session lease in the reconnect instead of extra E2B round trips Reused Mothership workbench calls made five E2B control-plane requests: list, connect, getInfo + setTimeout on acquire (the set always fired), and getInfo on release. Connect now asks for max(5 min, remaining, lease), the handle records the deadline it requested, and acquisition/release skip the provider while that lower bound covers the request. Final deadlines are unchanged; unrequested deadlines are still read back. * test(sandbox): model connect as setting the deadline so a dropped preserve is caught
1 parent 1056845 commit 733f4e9

6 files changed

Lines changed: 150 additions & 9 deletions

File tree

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

Lines changed: 75 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -606,3 +606,78 @@ describe('E2B session recovery', () => {
606606
expect(killSandbox).not.toHaveBeenCalled()
607607
})
608608
})
609+
610+
describe('E2B session lease', () => {
611+
const IDLE_MS = 20 * 60_000
612+
const LEASE_MS = IDLE_MS + 60_000
613+
614+
/** Counts control-plane requests; connect and setTimeout both set the deadline, so a caller must preserve it. */
615+
function controlPlane(remainingMs: number) {
616+
const plane = { endAtMs: Date.now() + remainingMs, requests: 0 }
617+
list.mockReturnValue({ nextItems, hasNext: false })
618+
nextItems.mockResolvedValue([{ ...candidate('retained', 10), endAt: new Date(plane.endAtMs) }])
619+
const sandbox = {
620+
sandboxId: 'retained',
621+
getInfo: async () => {
622+
plane.requests++
623+
return { endAt: new Date(plane.endAtMs) }
624+
},
625+
setTimeout: async (timeoutMs: number) => {
626+
plane.requests++
627+
plane.endAtMs = Date.now() + timeoutMs
628+
},
629+
}
630+
connect.mockImplementation(async (_id: string, options: { timeoutMs: number }) => {
631+
plane.requests++
632+
plane.endAtMs = Date.now() + options.timeoutMs
633+
return sandbox
634+
})
635+
create.mockImplementation(async (_template: string, options: { timeoutMs: number }) => {
636+
plane.requests++
637+
plane.endAtMs = Date.now() + options.timeoutMs
638+
return sandbox
639+
})
640+
return plane
641+
}
642+
643+
it('grants a reused lease in the reconnect and keeps the idle window without more requests', async () => {
644+
const plane = controlPlane(5 * 60_000)
645+
const sandbox = await e2bProvider.findSessionSandbox?.('chat', { lifetimeMs: LEASE_MS })
646+
expect(plane.endAtMs).toBeGreaterThanOrEqual(Date.now() + LEASE_MS - 1000)
647+
expect(sandbox?.outlives?.(IDLE_MS)).toBe(true)
648+
await sandbox?.extendLifetime?.(IDLE_MS)
649+
expect(plane.endAtMs).toBeGreaterThanOrEqual(Date.now() + IDLE_MS)
650+
expect(plane.requests).toBe(1)
651+
})
652+
653+
it('grants a created lease at creation without reading it back', async () => {
654+
const plane = controlPlane(0)
655+
const sandbox = await e2bProvider.create('mothership', {
656+
sessionKey: 'chat',
657+
lifetimeMs: LEASE_MS,
658+
})
659+
expect(sandbox.outlives?.(IDLE_MS)).toBe(true)
660+
await sandbox.extendLifetime?.(IDLE_MS)
661+
expect(plane.endAtMs).toBeGreaterThanOrEqual(Date.now() + LEASE_MS - 1000)
662+
expect(plane.requests).toBe(1)
663+
})
664+
665+
it('keeps a later deadline another job already granted', async () => {
666+
const plane = controlPlane(2 * 3_600_000)
667+
const sandbox = await e2bProvider.findSessionSandbox?.('chat', { lifetimeMs: LEASE_MS })
668+
await sandbox?.extendLifetime?.(IDLE_MS)
669+
expect(plane.endAtMs).toBeGreaterThanOrEqual(Date.now() + 2 * 3_600_000 - 1000)
670+
expect(plane.requests).toBe(1)
671+
})
672+
673+
it('reads back and extends a lease it did not grant itself', async () => {
674+
const plane = controlPlane(5 * 60_000)
675+
const sandbox = await e2bProvider.findSessionSandbox?.('chat', {})
676+
expect(sandbox?.outlives?.(IDLE_MS)).toBe(false)
677+
await sandbox?.extendLifetime?.(LEASE_MS)
678+
expect(plane.endAtMs).toBeGreaterThanOrEqual(Date.now() + LEASE_MS - 1000)
679+
expect(plane.requests).toBe(3)
680+
await sandbox?.extendLifetime?.(IDLE_MS)
681+
expect(plane.requests).toBe(3)
682+
})
683+
})

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

Lines changed: 37 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -359,25 +359,41 @@ class E2BSandboxHandle implements SandboxHandle {
359359
private killed = false
360360
private killPromise: Promise<void> | null = null
361361

362+
/**
363+
* @param sessionDeadlineAtMs Earliest time the provider can reap this session sandbox, as set
364+
* by this handle's own create, connect, or timeout request. Session deadlines only ever move
365+
* later — every update path extends and none shortens — so it stays a valid lower bound.
366+
*/
362367
constructor(
363368
private readonly sandbox: E2BSandbox,
364369
private readonly language: CodeLanguage,
365370
private readonly providerLimitAtMs?: number,
366-
private readonly sessionKey?: string
371+
private readonly sessionKey?: string,
372+
private sessionDeadlineAtMs?: number
367373
) {}
368374

369375
get sandboxId(): string {
370376
return this.sandbox.sandboxId
371377
}
372378

379+
outlives(lifetimeMs: number): boolean {
380+
return (
381+
this.sessionDeadlineAtMs !== undefined &&
382+
this.sessionDeadlineAtMs >= Date.now() + e2bTimeoutMs(lifetimeMs)
383+
)
384+
}
385+
373386
async extendLifetime(lifetimeMs: number): Promise<void> {
374387
const timeoutMs = e2bTimeoutMs(lifetimeMs)
375388
if (this.sessionKey !== undefined) {
389+
if (this.outlives(lifetimeMs)) return
376390
/** Session callers serialize updates so a short job cannot shorten another job's lease. */
377391
const info = await this.sandbox.getInfo()
378392
if (info.endAt.getTime() >= Date.now() + timeoutMs) return
379393
}
394+
const requestedAtMs = Date.now()
380395
await this.sandbox.setTimeout(timeoutMs)
396+
if (this.sessionKey !== undefined) this.sessionDeadlineAtMs = requestedAtMs + timeoutMs
381397
}
382398

383399
async runCode(
@@ -1072,13 +1088,16 @@ export const e2bProvider: SandboxProvider = {
10721088
effectiveLifetimeMs === E2B_MAX_SANDBOX_LIFETIME_MS
10731089
? lifetimeStartedAtMs + E2B_MAX_SANDBOX_LIFETIME_MS
10741090
: undefined,
1075-
options?.sessionKey
1091+
options?.sessionKey,
1092+
options?.sessionKey && effectiveLifetimeMs !== undefined
1093+
? lifetimeStartedAtMs + effectiveLifetimeMs
1094+
: undefined
10761095
)
10771096
},
10781097

10791098
async findSessionSandbox(
10801099
key: string,
1081-
options: { language?: CodeLanguage }
1100+
options: { language?: CodeLanguage; lifetimeMs?: number }
10821101
): Promise<SandboxHandle | null> {
10831102
const apiKey = env.E2B_API_KEY
10841103
if (!apiKey) throw new Error('E2B_API_KEY is required when E2B is enabled')
@@ -1097,9 +1116,21 @@ export const e2bProvider: SandboxProvider = {
10971116
'This workbench predates durable execution ownership and requires recovery before reuse'
10981117
)
10991118
}
1100-
// Connect also sets a timeout, including for running sandboxes. Preserve the active deadline.
1101-
const timeoutMs = Math.max(5 * 60_000, candidate.endAt.getTime() - Date.now())
1119+
// Connect also sets a timeout, including for running sandboxes. Preserve the active deadline,
1120+
// and grant the requested lease in the same request instead of a later getInfo + setTimeout.
1121+
const requestedAtMs = Date.now()
1122+
const timeoutMs = Math.max(
1123+
5 * 60_000,
1124+
candidate.endAt.getTime() - requestedAtMs,
1125+
options.lifetimeMs === undefined ? 0 : e2bTimeoutMs(options.lifetimeMs)
1126+
)
11021127
const sandbox = await Sandbox.connect(candidate.sandboxId, { apiKey, timeoutMs })
1103-
return new E2BSandboxHandle(sandbox, options.language ?? CodeLanguage.Python, undefined, key)
1128+
return new E2BSandboxHandle(
1129+
sandbox,
1130+
options.language ?? CodeLanguage.Python,
1131+
undefined,
1132+
key,
1133+
options.lifetimeMs === undefined ? undefined : requestedAtMs + timeoutMs
1134+
)
11041135
},
11051136
}

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -135,6 +135,8 @@ async function leaseSandbox(
135135
created,
136136
session: status,
137137
release: async () => {
138+
// A deadline that already covers the idle window needs no serialized update.
139+
if (created.sandbox.outlives?.(SESSION_SANDBOX_IDLE_MS)) return
138140
// Cleanup failure cannot relabel a completed mutation as a failed execution.
139141
try {
140142
await withSandboxSessionLock(session.key, AbortSignal.timeout(30_000), async () => {

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

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -771,6 +771,29 @@ describe('session sandbox lease', () => {
771771
expect(calls.extendLifetime.length).toBeGreaterThanOrEqual(2)
772772
})
773773

774+
it('trusts a lease granted by a slow reconnect for the whole call, then refreshes nothing', async () => {
775+
const { handle, calls } = fakeSandbox('sb-covered')
776+
let grantedUntilMs = 0
777+
handle.outlives = (lifetimeMs) => grantedUntilMs >= Date.now() + lifetimeMs
778+
mockFindSessionSandbox.mockImplementation(
779+
async (_key: string, options: { lifetimeMs?: number }) => {
780+
grantedUntilMs = Date.now() + (options.lifetimeMs ?? 0)
781+
await sleep(20)
782+
return handle
783+
}
784+
)
785+
786+
const result = await executeInSandbox({
787+
...CODE_REQUEST,
788+
sandboxKind: 'mothership',
789+
session: { key: 'mothership-chat:c3' },
790+
})
791+
792+
expect(result.sandboxSession).toBe('reused')
793+
expect(grantedUntilMs).toBeGreaterThanOrEqual(Date.now() + 20 * 60_000)
794+
expect(calls.extendLifetime).toHaveLength(0)
795+
})
796+
774797
it('does not rewrite an unchanged executable while earlier code can still use it', async () => {
775798
const { handle } = fakeSandbox('unchanged-tooling')
776799
mockFindSessionSandbox.mockResolvedValue(handle)

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

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -35,8 +35,10 @@ export async function ensureSessionSandbox(args: {
3535
signal.throwIfAborted()
3636
if (!provider.findSessionSandbox) throw new Error('This deployment has no persistent workbench')
3737
const lifetimeMs = SESSION_SANDBOX_IDLE_MS + (options.lifetimeMs ?? 0)
38+
const requestedAtMs = Date.now()
3839
const existing = await provider.findSessionSandbox(session.key, {
3940
...(options.language ? { language: options.language } : {}),
41+
lifetimeMs,
4042
})
4143
signal.throwIfAborted()
4244
const created: CreatedSandbox = existing
@@ -61,7 +63,9 @@ export async function ensureSessionSandbox(args: {
6163
providerId: provider.id,
6264
sandboxId: created.sandbox.sandboxId,
6365
})
64-
await created.sandbox.extendLifetime?.(lifetimeMs)
66+
// The budget is anchored before acquisition, so a lease granted by the lookup or create covers it.
67+
if (!created.sandbox.outlives?.(Math.max(0, lifetimeMs - (Date.now() - requestedAtMs))))
68+
await created.sandbox.extendLifetime?.(lifetimeMs)
6569
signal.throwIfAborted()
6670
if (session.cli) {
6771
await ensureSessionCli(created.sandbox, session.cli, signal, args.bootstrapTimeoutMs)

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

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -278,6 +278,11 @@ export interface SandboxHandle {
278278
* served an execution. Absent on providers without session support.
279279
*/
280280
extendLifetime?(lifetimeMs: number): Promise<void>
281+
/**
282+
* True when a deadline this handle already established keeps the sandbox alive for
283+
* `lifetimeMs` from now, so {@link extendLifetime} would make no provider request.
284+
*/
285+
outlives?(lifetimeMs: number): boolean
281286
/** Reads provider metadata without materializing the file contents. */
282287
getFileSize(path: string): Promise<number>
283288
readFile(path: string): Promise<string>
@@ -475,10 +480,11 @@ export interface SandboxProvider {
475480
* {@link CreateSandboxOptions.sessionKey}, or resolves null when none is
476481
* available. Lookup failures must throw rather than masquerade as absence.
477482
* Providers without session support omit this method; callers then
478-
* run every execution in a fresh sandbox.
483+
* run every execution in a fresh sandbox. `lifetimeMs` asks the reconnect to keep the
484+
* sandbox alive at least that long, never shortening a later deadline.
479485
*/
480486
findSessionSandbox?(
481487
key: string,
482-
options: { language?: CodeLanguage }
488+
options: { language?: CodeLanguage; lifetimeMs?: number }
483489
): Promise<SandboxHandle | null>
484490
}

0 commit comments

Comments
 (0)