From 97a22a75cf42bf9f48b8d996ac9804153bf922cb Mon Sep 17 00:00:00 2001
From: Brennan Benson <79079362+brennanb2025@users.noreply.github.com>
Date: Mon, 13 Jul 2026 20:39:08 -0700
Subject: [PATCH] fix(pty): clean up failed post-spawn setup

---
 src/main/ipc/pty.test.ts | 36 +++++++++++++++++++++++++++++++++---
 src/main/ipc/pty.ts      | 38 ++++++++++++++++++++++++++++++++++++--
 2 files changed, 69 insertions(+), 5 deletions(-)

diff --git a/src/main/ipc/pty.test.ts b/src/main/ipc/pty.test.ts
index 1c77f7552d3..e46193f2da2 100644
--- a/src/main/ipc/pty.test.ts
+++ b/src/main/ipc/pty.test.ts
@@ -6129,13 +6129,22 @@ describe('registerPtyHandlers', () => {
     // awaits a promise that never settles and the tab hangs forever.
     registerPtyHandlers(mainWindow as never)
     const leafId = '44444444-4444-4444-8444-444444444444'
-    const spawnArgs = { cols: 80, rows: 24, tabId: 'tab-reservation', leafId }
+    const paneKey = makePaneKey('tab-reservation', leafId)
+    const spawnArgs = {
+      cols: 80,
+      rows: 24,
+      tabId: 'tab-reservation',
+      leafId,
+      env: { ORCA_PANE_KEY: paneKey, ORCA_TAB_ID: 'tab-reservation' }
+    }
 
     registerPtyMock.mockImplementationOnce(() => {
       throw new Error('boom: post-spawn registration failed')
     })
 
     await expect(handlers.get('pty:spawn')!(null, spawnArgs)).rejects.toThrow('boom')
+    const orphanedPtyId = (registerPtyMock.mock.calls[0]![0] as { ptyId: string }).ptyId
+    expect(unregisterPtyMock).toHaveBeenCalledWith(orphanedPtyId)
 
     // A second spawn for the same pane must run a fresh spawn rather than await
     // the leaked (never-settled) reservation promise.
@@ -6175,12 +6184,16 @@ describe('registerPtyHandlers', () => {
     }
     let spawnCount = 0
     const providerSpawn = vi.fn(async () => ({ id: `pty-${++spawnCount}` }))
+    const shutdown = vi
+      .fn()
+      .mockRejectedValueOnce(new Error('native shutdown failed'))
+      .mockResolvedValue(undefined)
     setLocalPtyProvider({
       spawn: providerSpawn,
       write: vi.fn(),
       resize: vi.fn(),
       kill: vi.fn(),
-      shutdown: vi.fn(),
+      shutdown,
       sendSignal: vi.fn(),
       getCwd: vi.fn(),
       getInitialCwd: vi.fn(),
@@ -6199,7 +6212,9 @@ describe('registerPtyHandlers', () => {
       getProfiles: vi.fn()
     } as never)
     const store = {
-      persistPtyBinding: vi.fn()
+      persistPtyBinding: vi.fn(),
+      upsertPendingLocalPtyShutdown: vi.fn(),
+      removePendingLocalPtyShutdown: vi.fn()
     }
     let controller: RuntimeSpawnController | null = null
     const runtime = {
@@ -6240,6 +6255,21 @@ describe('registerPtyHandlers', () => {
     }
 
     await expect(spawnController.spawn(spawnArgs)).rejects.toThrow('boom')
+    expect(shutdown).toHaveBeenCalledWith('pty-1', {
+      immediate: true,
+      expectedPaneKey: paneKey,
+      expectedTabId: 'tab-runtime-reservation'
+    })
+    expect(store.upsertPendingLocalPtyShutdown).toHaveBeenCalledWith({
+      ptyId: 'pty-1',
+      expectedPaneKey: paneKey,
+      expectedTabId: 'tab-runtime-reservation',
+      requestedAt: expect.any(Number)
+    })
+    expect(store.removePendingLocalPtyShutdown).not.toHaveBeenCalled()
+    await vi.waitFor(() => expect(shutdown).toHaveBeenCalledTimes(2))
+    expect(store.removePendingLocalPtyShutdown).toHaveBeenCalledWith('pty-1')
+    expect(unregisterPtyMock).toHaveBeenCalledWith('pty-1')
 
     // The reservation must be gone, so a second materialization runs a fresh
     // provider.spawn instead of awaiting the leaked promise.
diff --git a/src/main/ipc/pty.ts b/src/main/ipc/pty.ts
index 5cd761d183e..6fb0c677eec 100644
--- a/src/main/ipc/pty.ts
+++ b/src/main/ipc/pty.ts
@@ -3732,7 +3732,8 @@ export function registerPtyHandlers(
       const paneSpawnReservation = materializedPaneKey
         ? reservePaneSpawn(materializedPaneKey)
         : null
-      let result: PtySpawnResult
+      let result: PtySpawnResult | undefined
+      let failedSpawnCleanupAttempted = false
       try {
         try {
           if (args.preAllocatedHandle) {
@@ -3829,6 +3830,7 @@ export function registerPtyHandlers(
               // writes, but the pre-existing provider session remains owned.
               deletePtyOwnership(result.id)
             } else {
+              failedSpawnCleanupAttempted = true
               if (args.connectionId) {
                 // Why: the binding failed, so the SSH lease is the only durable
                 // identity that can carry cleanup across an app restart.
@@ -3911,6 +3913,21 @@ export function registerPtyHandlers(
         const response = { id: result.id }
         return resolvePaneSpawnReservation(materializedPaneKey, paneSpawnReservation, response)
       } catch (err) {
+        if (result?.isReattach) {
+          // Why: this request published only routing ownership for an existing
+          // provider session; failed setup must not kill that session.
+          deletePtyOwnership(result.id)
+        } else if (result && !failedSpawnCleanupAttempted) {
+          try {
+            await retainFailedSpawnCleanup(result.id, args.connectionId, provider, {
+              immediate: true,
+              ...(metadataPaneKey ? { expectedPaneKey: metadataPaneKey } : {}),
+              ...(typeof args.tabId === 'string' ? { expectedTabId: args.tabId } : {})
+            })
+          } catch (shutdownErr) {
+            console.warn('[pty] failed to clean up PTY after post-spawn failure:', shutdownErr)
+          }
+        }
         // Why: once the reservation is created, any later throw — spawn
         // failure, persist failure, or a post-spawn helper such as
         // registerPty/rememberPaneKeyForPty/track — must settle it. Otherwise
@@ -4642,7 +4659,8 @@ export function registerPtyHandlers(
           if (preSpawnHiddenMarkId !== null) {
             markHiddenRendererPty(preSpawnHiddenMarkId)
           }
-          let result: PtySpawnResult
+          let result: PtySpawnResult | undefined
+          let failedSpawnCleanupAttempted = false
           try {
             try {
               if (preAllocatedHandle) {
@@ -4826,6 +4844,7 @@ export function registerPtyHandlers(
               } catch (err) {
                 console.error('[pty] failed to persist PTY binding after spawn:', err)
                 if (!result.isReattach) {
+                  failedSpawnCleanupAttempted = true
                   try {
                     await retainFailedSpawnCleanup(result.id, args.connectionId, provider, {
                       immediate: true,
@@ -5030,6 +5049,21 @@ export function registerPtyHandlers(
             }
             return resolvePaneSpawnReservation(reservationPaneKey, paneSpawnReservation, response)
           } catch (err) {
+            if (result?.isReattach) {
+              // Why: this request published only routing ownership for an existing
+              // provider session; failed setup must not kill that session.
+              deletePtyOwnership(result.id)
+            } else if (result && !failedSpawnCleanupAttempted) {
+              try {
+                await retainFailedSpawnCleanup(result.id, args.connectionId, provider, {
+                  immediate: true,
+                  ...(validatedPaneKey ? { expectedPaneKey: validatedPaneKey } : {}),
+                  ...(typeof args.tabId === 'string' ? { expectedTabId: args.tabId } : {})
+                })
+              } catch (shutdownErr) {
+                console.warn('[pty] failed to clean up PTY after post-spawn failure:', shutdownErr)
+              }
+            }
             // Why: once the reservation is created, any later throw —
             // spawn failure, persist failure, or a post-spawn helper such as
             // seedHeadlessTerminal/registerPty/track — must settle it. Otherwise