fix(codex): fence detached restart ownership races

src/renderer/src/App.tsx

@@ -124,7 +124,7 @@ import {
 } from './lib/workspace-session'
 import { createSessionWriteSubscriber } from './lib/session-write-subscriber'
 import { sweepRestoredCodexPanesForStaleAccounts } from './lib/codex-stale-pane-sweep'
-import { installCodexDetachedPaneRestartExecutor } from '@/components/terminal-pane/codex-detached-pane-restart'
+import { installCodexDetachedPaneRestartExecutor } from '@/components/terminal-pane/codex-detached-pane-restart-scheduler'
 import { buildActiveViewUnloadPatch } from './lib/active-view-persist'
 import {
   buildWorkspaceSessionHostSnapshots,

src/renderer/src/components/terminal-pane/codex-detached-pane-restart-scheduler.ts

@@ -0,0 +1,74 @@
+import { useAppStore } from '@/store'
+import {
+  resetDetachedCodexPaneRestartClaimsForTests,
+  sweepUnclaimedCodexPaneRestarts
+} from './codex-detached-pane-restart'
+
+let executorInstalled = false
+let executorGeneration = 0
+let sweepQueued = false
+let sweepRunning = false
+let sweepRequestedAfterRun = false
+
+/** Installed once at app startup; returns the uninstaller (tests). */
+export function installCodexDetachedPaneRestartExecutor(): () => void {
+  executorInstalled = true
+  const generation = ++executorGeneration
+  const unsubscribe = useAppStore.subscribe((state, previousState) => {
+    const addedPendingId = Object.keys(state.pendingCodexPaneRestartIds).some(
+      (ptyId) => !previousState.pendingCodexPaneRestartIds[ptyId]
+    )
+    if (addedPendingId) {
+      scheduleClaimSweep()
+    }
+  })
+  scheduleClaimSweep()
+  return () => {
+    unsubscribe()
+    if (executorGeneration === generation) {
+      executorInstalled = false
+      executorGeneration += 1
+      sweepQueued = false
+      sweepRequestedAfterRun = false
+    }
+  }
+}
+
+export function resetCodexDetachedPaneRestartExecutorForTests(): void {
+  executorInstalled = false
+  executorGeneration += 1
+  sweepQueued = false
+  sweepRunning = false
+  sweepRequestedAfterRun = false
+  resetDetachedCodexPaneRestartClaimsForTests()
+}
+
+function scheduleClaimSweep(): void {
+  if (!executorInstalled) {
+    return
+  }
+  if (sweepRunning) {
+    sweepRequestedAfterRun = true
+    return
+  }
+  if (sweepQueued) {
+    return
+  }
+  sweepQueued = true
+  const generation = executorGeneration
+  // Exact mounted-owner checks fence the claim; a microtask only exits the store write.
+  queueMicrotask(() => {
+    sweepQueued = false
+    if (!executorInstalled || executorGeneration !== generation) {
+      return
+    }
+    sweepRunning = true
+    void sweepUnclaimedCodexPaneRestarts().finally(() => {
+      sweepRunning = false
+      if (executorInstalled && sweepRequestedAfterRun) {
+        sweepRequestedAfterRun = false
+        scheduleClaimSweep()
+      }
+    })
+  })
+}

src/renderer/src/components/terminal-pane/codex-detached-pane-restart.test.ts

@@ -3,11 +3,11 @@ import { useAppStore } from '@/store'
 import { registerRuntimeTerminalTab } from '@/runtime/sync-runtime-graph'
 import { awaitsCodexRestartAnswer, blocksCodexPaneInput } from '../codex-restart-notice-state'
 import { ptyDataHandlers } from './pty-dispatcher'
+import { sweepUnclaimedCodexPaneRestarts } from './codex-detached-pane-restart'
 import {
   installCodexDetachedPaneRestartExecutor,
-  resetCodexDetachedPaneRestartExecutorForTests,
-  sweepUnclaimedCodexPaneRestarts
-} from './codex-detached-pane-restart'
+  resetCodexDetachedPaneRestartExecutorForTests
+} from './codex-detached-pane-restart-scheduler'
 
 const ACCOUNT_A = 'a@example.com'
 const ACCOUNT_B = 'b@example.com'
@@ -15,6 +15,11 @@ const LEAF_ID = '11111111-1111-4111-8111-111111111111'
 const OLD_PTY = 'wt1@@old'
 const NEW_PTY = 'wt1@@new'
 
+function deferred<T>(): { promise: Promise<T>; resolve: (value: T) => void } {
+  let resolve!: (value: T) => void
+  return { promise: new Promise<T>((done) => (resolve = done)), resolve }
+}
+
 function seedQueuedRestart(
   opts: { leafId?: string | null; layoutRoot?: 'leaf' | 'none' } = {}
 ): void {
@@ -124,22 +129,22 @@ describe('codex detached pane restart executor', () => {
     expect(state.tabsByWorktree.wt1?.[0]?.ptyId).toBe(NEW_PTY)
     expect(state.terminalLayoutsByTabId['tab-1']?.ptyIdsByLeafId).toEqual({ [LEAF_ID]: NEW_PTY })
     expect(state.pendingCodexPaneRestartIds).toEqual({})
+    expect(state.suppressedPtyExitIds[OLD_PTY]).toBeUndefined()
     // The whole point: the restart completed, so nothing may block input anymore.
     expect(blocksCodexPaneInput(state.codexRestartNoticeByPtyId[OLD_PTY])).toBe(false)
     expect(blocksCodexPaneInput(state.codexRestartNoticeByPtyId[NEW_PTY])).toBe(false)
   })
 
-  it('executes via the store subscription once the claim grace elapses', async () => {
-    vi.useFakeTimers()
+  it('executes via the store subscription without a lifecycle timeout', async () => {
     const uninstall = installCodexDetachedPaneRestartExecutor()
     try {
       seedQueuedRestart()
       expect(window.api.pty.spawn).not.toHaveBeenCalled()
 
-      await vi.advanceTimersByTimeAsync(400)
+      await vi.waitFor(() => expect(window.api.pty.spawn).toHaveBeenCalledTimes(1))
+      await vi.waitFor(() => expect(window.api.pty.kill).toHaveBeenCalledExactlyOnceWith(OLD_PTY))
 
-      expect(window.api.pty.spawn).toHaveBeenCalledTimes(1)
-      expect(window.api.pty.kill).toHaveBeenCalledExactlyOnceWith(OLD_PTY)
+      expect(useAppStore.getState().pendingCodexPaneRestartIds).toEqual({})
     } finally {
       uninstall()
     }
@@ -251,6 +256,65 @@ describe('codex detached pane restart executor', () => {
     }
   })
 
+  it('reaps a detached spawn and requeues when a pane mounts during the spawn', async () => {
+    seedQueuedRestart()
+    const pendingSpawn = deferred<{ id: string }>()
+    vi.mocked(window.api.pty.spawn).mockReturnValue(pendingSpawn.promise)
+
+    const restart = sweepUnclaimedCodexPaneRestarts()
+    await vi.waitFor(() => expect(window.api.pty.spawn).toHaveBeenCalledTimes(1))
+    const unregister = registerRuntimeTerminalTab({
+      tabId: 'tab-1',
+      worktreeId: 'wt1',
+      getManager: () => null,
+      getContainer: () => null,
+      getPtyIdForPane: () => OLD_PTY
+    })
+    try {
+      pendingSpawn.resolve({ id: NEW_PTY })
+      await restart
+
+      expect(window.api.pty.kill).toHaveBeenCalledExactlyOnceWith(NEW_PTY)
+      expect(useAppStore.getState().ptyIdsByTabId['tab-1']).toEqual([OLD_PTY])
+      expect(useAppStore.getState().pendingCodexPaneRestartIds).toEqual({ [OLD_PTY]: true })
+    } finally {
+      unregister()
+    }
+  })
+
+  it('rejects a detached spawn after the tab generation and leaf owner change', async () => {
+    seedQueuedRestart()
+    const pendingSpawn = deferred<{ id: string }>()
+    vi.mocked(window.api.pty.spawn).mockReturnValue(pendingSpawn.promise)
+
+    const restart = sweepUnclaimedCodexPaneRestarts()
+    await vi.waitFor(() => expect(window.api.pty.spawn).toHaveBeenCalledTimes(1))
+    const state = useAppStore.getState()
+    const notice = state.codexRestartNoticeByPtyId[OLD_PTY]
+    useAppStore.setState({
+      tabsByWorktree: {
+        wt1: [{ ...state.tabsByWorktree.wt1![0]!, generation: 1, ptyId: 'wt1@@successor' }]
+      },
+      ptyIdsByTabId: { 'tab-1': ['wt1@@successor'] },
+      terminalLayoutsByTabId: {
+        'tab-1': {
+          ...state.terminalLayoutsByTabId['tab-1']!,
+          ptyIdsByLeafId: { [LEAF_ID]: 'wt1@@successor' }
+        }
+      },
+      codexRestartNoticeByPtyId: { 'wt1@@successor': notice! }
+    })
+
+    pendingSpawn.resolve({ id: NEW_PTY })
+    await restart
+
+    const after = useAppStore.getState()
+    expect(window.api.pty.kill).toHaveBeenCalledExactlyOnceWith(NEW_PTY)
+    expect(after.ptyIdsByTabId['tab-1']).toEqual(['wt1@@successor'])
+    expect(after.terminalLayoutsByTabId['tab-1']?.ptyIdsByLeafId?.[LEAF_ID]).toBe('wt1@@successor')
+    expect(awaitsCodexRestartAnswer(after.codexRestartNoticeByPtyId['wt1@@successor'])).toBe(true)
+  })
+
   it('leaves a sleep-retained pending id alone so wake can migrate it', async () => {
     // Why: hibernation unbinds a pane's PTY but keeps its pending restart, and
     // wake moves that entry onto the respawned PTY. No notice means nothing is
@@ -296,4 +360,21 @@ describe('codex detached pane restart executor', () => {
     )
     expect(blocksCodexPaneInput(state.codexRestartNoticeByPtyId[OLD_PTY])).toBe(false)
   })
+
+  it('publishes a rootless replacement startup before waiting for old PTY teardown', async () => {
+    seedQueuedRestart({ leafId: null })
+    const pendingKill = deferred<void>()
+    vi.mocked(window.api.pty.kill).mockReturnValue(pendingKill.promise)
+
+    const restart = sweepUnclaimedCodexPaneRestarts()
+    await vi.waitFor(() => expect(window.api.pty.kill).toHaveBeenCalledExactlyOnceWith(OLD_PTY))
+
+    expect(useAppStore.getState().ptyIdsByTabId['tab-1']).toEqual([])
+    expect(useAppStore.getState().pendingStartupByTabId['tab-1']).toMatchObject({
+      command: 'codex',
+      launchAgent: 'codex'
+    })
+    pendingKill.resolve()
+    await restart
+  })
 })

src/renderer/src/components/terminal-pane/codex-detached-pane-restart.ts

@@ -27,49 +27,12 @@ import { ptyDataHandlers, unregisterPtyDataHandlers } from './pty-dispatcher'
 import { discardPreHandlerPtyState } from './pty-pre-handler-buffer'
 import { disposeParkedTerminalWatchersForPtyIds } from './terminal-parked-watcher-registry'
 
-// Why this long: a mounted pane claims its pending restart from a React effect
-// in the same commit that observed the queue write, so anything still pending
-// after a few frames has no pane coming for it.
-const CLAIM_GRACE_MS = 400
-
-let sweepTimer: ReturnType<typeof setTimeout> | null = null
 const inFlightPtyIds = new Set<string>()
 
-/** Installed once at app startup; returns the uninstaller (tests). */
-export function installCodexDetachedPaneRestartExecutor(): () => void {
-  const unsubscribe = useAppStore.subscribe((state, previousState) => {
-    if (state.pendingCodexPaneRestartIds === previousState.pendingCodexPaneRestartIds) {
-      return
-    }
-    scheduleClaimSweep()
-  })
-  return () => {
-    unsubscribe()
-    if (sweepTimer !== null) {
-      clearTimeout(sweepTimer)
-      sweepTimer = null
-    }
-  }
-}
-
-export function resetCodexDetachedPaneRestartExecutorForTests(): void {
-  if (sweepTimer !== null) {
-    clearTimeout(sweepTimer)
-    sweepTimer = null
-  }
+export function resetDetachedCodexPaneRestartClaimsForTests(): void {
   inFlightPtyIds.clear()
 }
 
-function scheduleClaimSweep(): void {
-  if (sweepTimer !== null) {
-    return
-  }
-  sweepTimer = setTimeout(() => {
-    sweepTimer = null
-    void sweepUnclaimedCodexPaneRestarts()
-  }, CLAIM_GRACE_MS)
-}
-
 export async function sweepUnclaimedCodexPaneRestarts(): Promise<void> {
   for (const ptyId of Object.keys(useAppStore.getState().pendingCodexPaneRestartIds)) {
     // Why: remote-runtime spawns need that machine's transport assembly, which
@@ -112,7 +75,7 @@ export async function sweepUnclaimedCodexPaneRestarts(): Promise<void> {
       console.warn('[codex-restart] detached pane restart failed:', err)
       // Why: the answered-but-unexecuted state is the one this module exists to
       // prevent — put the question back rather than leave a silent input block.
-      useAppStore.getState().reopenCodexRestartPrompt(ptyId)
+      reopenCurrentCodexRestartPrompt(located, ptyId)
     } finally {
       inFlightPtyIds.delete(ptyId)
     }
@@ -123,6 +86,7 @@ type LocatedCodexPane = {
   worktreeId: string
   tab: TerminalTab
   leafId: string | null
+  generation: number
 }
 
 function locateCodexPane(state: AppState, ptyId: string): LocatedCodexPane | null {
@@ -140,7 +104,8 @@ function locateCodexPane(state: AppState, ptyId: string): LocatedCodexPane | nul
       return {
         worktreeId,
         tab,
-        leafId: leafId !== null && isTerminalLeafId(leafId) ? leafId : null
+        leafId: leafId !== null && isTerminalLeafId(leafId) ? leafId : null,
+        generation: tab.generation ?? 0
       }
     }
   }
@@ -191,11 +156,16 @@ async function executeDetachedCodexPaneRestart(
   if (!located.leafId) {
     // Why: without a usable layout leaf the replacement cannot be bound in
     // place, so kill now and let the tab's next mount run the Codex startup.
-    await killReplacedCodexPanePty(ptyId)
+    if (!isLocatedCodexPaneCurrent(state, located, ptyId)) {
+      reopenCurrentCodexRestartPrompt(located, ptyId)
+      return
+    }
+    const killPromise = killReplacedCodexPanePty(ptyId)
     const store = useAppStore.getState()
     store.clearTabPtyId(located.tab.id, ptyId)
     store.queueTabStartupCommand(located.tab.id, { ...CODEX_ACCOUNT_RESTART_STARTUP })
     store.clearCodexRestartNotice(ptyId)
+    await killPromise
     return
   }
   const { worktreeId, tab, leafId } = located
@@ -213,6 +183,16 @@ async function executeDetachedCodexPaneRestart(
     availableWslDistros: capabilities?.wslDistros ?? null
   })
 
+  const currentState = useAppStore.getState()
+  if (!isLocatedCodexPaneCurrent(currentState, located, ptyId)) {
+    reopenCurrentCodexRestartPrompt(located, ptyId)
+    return
+  }
+  if (hasRegisteredRuntimeTerminalTab(tab.id) || ptyDataHandlers.has(ptyId)) {
+    currentState.queueCodexPaneRestarts([ptyId])
+    return
+  }
+
   const spawned = await window.api.pty.spawn({
     cols: size?.cols ?? 80,
     rows: size?.rows ?? 24,
@@ -231,6 +211,16 @@ async function executeDetachedCodexPaneRestart(
   })
 
   const store = useAppStore.getState()
+  if (!isLocatedCodexPaneCurrent(store, located, ptyId)) {
+    await window.api.pty.kill(spawned.id).catch(() => {})
+    reopenCurrentCodexRestartPrompt(located, ptyId)
+    return
+  }
+  if (hasRegisteredRuntimeTerminalTab(tab.id) || ptyDataHandlers.has(ptyId)) {
+    await window.api.pty.kill(spawned.id).catch(() => {})
+    store.queueCodexPaneRestarts([ptyId])
+    return
+  }
   store.updateTabPtyId(tab.id, spawned.id, ptyId)
   if (!useAppStore.getState().ptyIdsByTabId[tab.id]?.includes(spawned.id)) {
     // Why: the tab was retired while the spawn was in flight; without a binding
@@ -248,6 +238,44 @@ async function executeDetachedCodexPaneRestart(
   await killReplacedCodexPanePty(ptyId)
 }
 
+function isLocatedCodexPaneCurrent(
+  state: AppState,
+  located: LocatedCodexPane,
+  ptyId: string
+): boolean {
+  const currentTab = state.tabsByWorktree[located.worktreeId]?.find(
+    (candidate) => candidate.id === located.tab.id
+  )
+  if (
+    !currentTab ||
+    currentTab.worktreeId !== located.worktreeId ||
+    (currentTab.generation ?? 0) !== located.generation ||
+    !(state.ptyIdsByTabId[located.tab.id] ?? []).includes(ptyId)
+  ) {
+    return false
+  }
+  return (
+    located.leafId === null ||
+    state.terminalLayoutsByTabId[located.tab.id]?.ptyIdsByLeafId?.[located.leafId] === ptyId
+  )
+}
+
+function reopenCurrentCodexRestartPrompt(located: LocatedCodexPane, replacedPtyId: string): void {
+  const state = useAppStore.getState()
+  const currentTab = state.tabsByWorktree[located.worktreeId]?.find(
+    (candidate) => candidate.id === located.tab.id
+  )
+  const currentPtyId = located.leafId
+    ? state.terminalLayoutsByTabId[located.tab.id]?.ptyIdsByLeafId?.[located.leafId]
+    : currentTab?.ptyId
+  for (const candidate of [currentPtyId, replacedPtyId]) {
+    if (candidate && state.codexRestartNoticeByPtyId[candidate]?.restartRequested) {
+      state.reopenCodexRestartPrompt(candidate)
+      return
+    }
+  }
+}
+
 function layoutRootContainsLeaf(
   node: TerminalPaneLayoutNode | null | undefined,
   leafId: string
@@ -294,4 +322,6 @@ async function killReplacedCodexPanePty(ptyId: string): Promise<void> {
     console.warn('[codex-restart] failed to kill replaced Codex pane PTY:', err)
   }
   discardPreHandlerPtyState(ptyId)
+  // No exit consumer remains after the watcher/handler disposal above.
+  useAppStore.getState().consumeSuppressedPtyExit(ptyId)
 }