From 896bcd64a10d74f75a674e278bd816738a31062e Mon Sep 17 00:00:00 2001
From: Tianyi Cui <53024+tianyicui@users.noreply.github.com>
Date: Wed, 29 Jul 2026 10:14:06 +0800
Subject: [PATCH] fix(e2b): await teardown and honor published status

---
 packages/e2b/subprocess-e2b/src/index.ts      |  5 ++-
 packages/e2b/subprocess-e2b/src/process.ts    | 11 +++--
 .../subprocess-e2b/tests/subprocess.spec.ts   | 43 +++++++++++++++++++
 3 files changed, 54 insertions(+), 5 deletions(-)

diff --git a/packages/e2b/subprocess-e2b/src/index.ts b/packages/e2b/subprocess-e2b/src/index.ts
index bf7edc41a1..4516bffa84 100644
--- a/packages/e2b/subprocess-e2b/src/index.ts
+++ b/packages/e2b/subprocess-e2b/src/index.ts
@@ -67,7 +67,10 @@ export class E2BSubprocessService extends SubprocessService {
       for (const cleanup of failedTerminalSetupCleanups) {
         pending.push(cleanup().then(() => { this.failedTerminalSetupCleanups.delete(cleanup) }))
       }
-      await Promise.all(pending)
+      const outcomes = await Promise.allSettled(pending)
+      for (const outcome of outcomes) {
+        if (outcome.status === 'rejected') throw outcome.reason
+      }
     }, 'e2b subprocess teardown')
   }
 
diff --git a/packages/e2b/subprocess-e2b/src/process.ts b/packages/e2b/subprocess-e2b/src/process.ts
index 5387966578..478dc96c30 100644
--- a/packages/e2b/subprocess-e2b/src/process.ts
+++ b/packages/e2b/subprocess-e2b/src/process.ts
@@ -524,10 +524,10 @@ export class E2BSubprocessHandle implements SubprocessHandle {
           throw new Error(`subprocess-e2b: remote wrapper published invalid exit code ${JSON.stringify(rawStatus)}`)
         }
         if (this.spec.stdio.stdout === 'pipe' || this.spec.stdio.stderr === 'pipe') {
-          return this.commandOutcome(await settlement)
+          return this.commandOutcome(await settlement, exitCode)
         }
         const completed = await withinMs(settlement, this.spec.graceMs)
-        if (completed !== undefined) return this.commandOutcome(completed)
+        if (completed !== undefined) return this.commandOutcome(completed, exitCode)
         this.outputDrainExpired = true
         this.stdoutReader?.invalidateSpill()
         this.stderrReader?.invalidateSpill()
@@ -539,9 +539,12 @@ export class E2BSubprocessHandle implements SubprocessHandle {
     }
   }
 
-  private commandOutcome(settlement: CommandSettlement): SubprocessOutcome {
-    if (settlement.kind === 'result') return { exitCode: settlement.result.exitCode, signal: null }
+  private commandOutcome(settlement: CommandSettlement, publishedExitCode?: number): SubprocessOutcome {
+    if (settlement.kind === 'result') {
+      return { exitCode: publishedExitCode ?? settlement.result.exitCode, signal: null }
+    }
     if (settlement.error instanceof CommandExitError) {
+      if (publishedExitCode !== undefined) return { exitCode: publishedExitCode, signal: null }
       return this.terminationSignal === null
         ? { exitCode: settlement.error.exitCode, signal: null }
         : { exitCode: null, signal: this.terminationSignal }
diff --git a/packages/e2b/subprocess-e2b/tests/subprocess.spec.ts b/packages/e2b/subprocess-e2b/tests/subprocess.spec.ts
index b7d1642bd6..6d20abd3c1 100644
--- a/packages/e2b/subprocess-e2b/tests/subprocess.spec.ts
+++ b/packages/e2b/subprocess-e2b/tests/subprocess.spec.ts
@@ -547,6 +547,20 @@ describe('E2BSubprocessHandle', () => {
     await expect(handle.waitForExit()).resolves.toBe(true)
   })
 
+  it('preserves a published nonzero exit code when termination settles the SDK inside the drain grace', async () => {
+    const fake = new FakeSandbox()
+    const handle = new E2BSubprocessHandle(runtime(fake), spec({ graceMs: 100 }), '/runtime/drain-signal-settled')
+    await flush()
+    fake.exitStatus = '7\n'
+    fake.afterStatusRead = () => {
+      fake.afterStatusRead = undefined
+      handle.terminate()
+    }
+
+    await expect(handle.done).resolves.toEqual({ exitCode: 7, signal: null })
+    await expect(handle.waitForExit()).resolves.toBe(true)
+  })
+
   it('rejects an invalid direct-command exit status', async () => {
     const fake = new FakeSandbox()
     const handle = new E2BSubprocessHandle(runtime(fake), spec(), '/runtime/invalid-status')
@@ -1336,6 +1350,35 @@ describe('E2BSubprocessService', () => {
     await expect(handle.done).resolves.toEqual({ exitCode: null, signal: 'SIGTERM' })
   })
 
+  it('waits for every owned cleanup before reporting a disposal failure', async () => {
+    const { ctx, fiber } = await service()
+    const failed = {
+      terminate: vi.fn(),
+      waitForExit: vi.fn(async () => { throw new Error('cleanup failed') }),
+      done: Promise.resolve({ exitCode: 0, signal: null }),
+    } as unknown as E2BSubprocessHandle
+    let finishCleanup!: () => void
+    const cleanup = new Promise<boolean>((resolve) => {
+      finishCleanup = () => { resolve(true) }
+    })
+    const draining = {
+      terminate: vi.fn(),
+      waitForExit: vi.fn(() => cleanup),
+      done: Promise.resolve({ exitCode: 0, signal: null }),
+    } as unknown as E2BSubprocessHandle
+    const live = (ctx.subprocess as unknown as { live: Set<E2BSubprocessHandle> }).live
+    live.add(failed)
+    live.add(draining)
+
+    let disposed = false
+    const disposing = fiber.dispose().then(() => { disposed = true })
+    await flush()
+    expect(disposed).toBe(false)
+    finishCleanup()
+    await disposing
+    expect(live).toEqual(new Set([failed]))
+  })
+
   it('releases naturally settled handles before later service disposal', async () => {
     const fake = new FakeSandbox()
     const { ctx, fiber } = await service(fake)