From e9bef2967dcd6c399637482f8cb1739cf90bf6e0 Mon Sep 17 00:00:00 2001
From: Tianyi Cui <53024+tianyicui@users.noreply.github.com>
Date: Wed, 29 Jul 2026 07:07:31 +0800
Subject: [PATCH] fix(e2b): preserve signals after output drain

---
 packages/e2b/subprocess-e2b/src/process.ts    |  4 +++-
 .../subprocess-e2b/tests/subprocess.spec.ts   | 22 ++++++++++++++++++-
 2 files changed, 24 insertions(+), 2 deletions(-)

diff --git a/packages/e2b/subprocess-e2b/src/process.ts b/packages/e2b/subprocess-e2b/src/process.ts
index 3a8047eb54..473c1a0af4 100644
--- a/packages/e2b/subprocess-e2b/src/process.ts
+++ b/packages/e2b/subprocess-e2b/src/process.ts
@@ -512,7 +512,9 @@ export class E2BSubprocessHandle implements SubprocessHandle {
         this.stdoutReader?.invalidateSpill()
         this.stderrReader?.invalidateSpill()
         await handle.disconnect()
-        return { exitCode, signal: null }
+        return this.terminationSignal === null
+          ? { exitCode, signal: null }
+          : { exitCode: null, signal: this.terminationSignal }
       }
       const completed = await Promise.race([settlement, waitTick().then(() => undefined)])
       if (completed !== undefined) return this.commandOutcome(completed)
diff --git a/packages/e2b/subprocess-e2b/tests/subprocess.spec.ts b/packages/e2b/subprocess-e2b/tests/subprocess.spec.ts
index d81b49a600..457e68f24d 100644
--- a/packages/e2b/subprocess-e2b/tests/subprocess.spec.ts
+++ b/packages/e2b/subprocess-e2b/tests/subprocess.spec.ts
@@ -103,6 +103,7 @@ class FakeSandbox {
   readonly signalErrors: unknown[] = []
   trapsTerm = false
   delaysKill = false
+  delaysKillCompletion = false
   sdkKillStops = true
   alive = true
   ambient = 'PATH=/ambient/bin\0KEEP=safe\0NPM_TOKEN=secret\0DSH_STALE=old\0BROKEN\0=bad\0'
@@ -247,7 +248,7 @@ class FakeSandbox {
             throw error
           }
           if (!this.delaysKill) this.alive = false
-          this.handle.fail(137)
+          if (!this.delaysKillCompletion) this.handle.fail(137)
           return { exitCode: 0, stdout: '', stderr: '' }
         }
         if ((options as StartOptions | undefined)?.background === true) {
@@ -473,6 +474,25 @@ describe('E2BSubprocessHandle', () => {
     expect(fake.handle.disconnects).toBe(0)
   })
 
+  it('preserves a requested signal when output draining expires', async () => {
+    const fake = new FakeSandbox()
+    fake.trapsTerm = true
+    fake.delaysKill = true
+    fake.delaysKillCompletion = true
+    fake.sdkKillStops = false
+    const handle = new E2BSubprocessHandle(runtime(fake), spec({ graceMs: 5 }), '/runtime/drain-signal')
+    await flush()
+
+    handle.terminate()
+    await vi.waitFor(() => { expect(fake.commandsSeen).toContain('kill -KILL -- -4242') })
+    fake.exitStatus = '143\n'
+
+    await expect(handle.done).resolves.toEqual({ exitCode: null, signal: 'SIGKILL' })
+    expect(fake.handle.disconnects).toBe(1)
+    fake.alive = false
+    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')