Переглянути джерело

fix(e2b): await teardown and honor published status

Tianyi Cui 2 місяців тому
батько
коміт
896bcd64a1

+ 4 - 1
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')
   }
 

+ 7 - 4
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 }

+ 43 - 0
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)