fix: dismiss Codex account prompt and return focus to terminal (#24683)

src/renderer/src/components/CodexRestartChip.tsx

@@ -20,10 +20,12 @@ type RestartNotice = {
 export default function CodexRestartChip({
   isVisible = true,
   ptyId,
+  onReturnFocus,
   shouldFocus = false
 }: {
   isVisible?: boolean
   ptyId: string
+  onReturnFocus: () => void
   shouldFocus?: boolean
 }): React.JSX.Element | null {
   // Why: one O(1) selector per mounted pane stays idle when unrelated PTY maps
@@ -35,10 +37,12 @@ export default function CodexRestartChip({
 
   const handleRestart = (): void => {
     useAppStore.getState().queueCodexPaneRestarts([ptyId])
+    onReturnFocus()
   }
 
   const handleDismiss = (): void => {
     useAppStore.getState().dismissCodexRestartNotices([ptyId])
+    onReturnFocus()
     // Why: notices are renderer-only, so the persisted launch record must be
     // cleared for this pane or the startup sweep re-raises its answered prompt.
     void window.api.codexAccounts.forgetStalePanes({ ptyIds: [ptyId] }).catch((err: unknown) => {
@@ -106,7 +110,19 @@ function LoudRestartOverlay({
       aria-live="assertive"
       aria-labelledby={titleId}
       aria-describedby={bodyId}
-      className="pointer-events-none absolute inset-0 z-50 flex items-center justify-center p-6 outline-none"
+      onMouseDown={(event) => {
+        if (event.button === 0 && event.target === event.currentTarget) {
+          onDismiss()
+        }
+      }}
+      onKeyDown={(event) => {
+        if (event.key === 'Escape') {
+          event.preventDefault()
+          event.stopPropagation()
+          onDismiss()
+        }
+      }}
+      className="absolute inset-0 z-50 flex items-center justify-center p-6 outline-none"
     >
       <div className="pointer-events-auto flex w-full max-w-[30rem] flex-col gap-3 rounded-lg border border-border bg-card p-6 pb-5 text-card-foreground shadow-xs">
         <div className="flex items-start gap-3">

src/renderer/src/components/codex-restart-chip.test.tsx

@@ -4,7 +4,14 @@ import React, { act, Profiler } from 'react'
 import { createRoot, type Root } from 'react-dom/client'
 import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
 import { useAppStore } from '../store'
-import CodexRestartChip from './CodexRestartChip'
+import CodexRestartChipComponent from './CodexRestartChip'
+
+const returnFocus = vi.fn()
+function CodexRestartChip(
+  props: Omit<React.ComponentProps<typeof CodexRestartChipComponent>, 'onReturnFocus'>
+) {
+  return <CodexRestartChipComponent {...props} onReturnFocus={returnFocus} />
+}
 
 globalThis.IS_REACT_ACT_ENVIRONMENT = true
 
@@ -30,6 +37,7 @@ function button(scope: ParentNode, label: string): HTMLButtonElement {
 }
 
 beforeEach(() => {
+  returnFocus.mockReset()
   useAppStore.setState(useAppStore.getInitialState(), true)
   container = document.createElement('div')
   document.body.appendChild(container)
@@ -118,6 +126,7 @@ describe('CodexRestartChip pane ownership', () => {
       button(firstPane, 'Restart').click()
     })
 
+    expect(returnFocus).toHaveBeenCalledOnce()
     const state = useAppStore.getState()
     expect(state.pendingCodexPaneRestartIds).toEqual({ [PTY_ONE]: true })
     expect(state.codexRestartNoticeByPtyId[PTY_ONE]?.restartRequested).toBe(true)
@@ -155,6 +164,7 @@ describe('CodexRestartChip pane ownership', () => {
       button(firstPane, 'Keep old account').click()
     })
 
+    expect(returnFocus).toHaveBeenCalledOnce()
     expect(forgetStalePanes).toHaveBeenCalledExactlyOnceWith({ ptyIds: [PTY_ONE] })
     expect(useAppStore.getState().codexRestartNoticeByPtyId[PTY_ONE]?.dismissed).toBe(true)
     expect(useAppStore.getState().codexRestartNoticeByPtyId[PTY_TWO]?.dismissed).toBeUndefined()
@@ -273,6 +283,67 @@ describe('CodexRestartChip pane focus', () => {
     return terminalInput
   }
 
+  it.each(['outside', 'keep', 'escape', 'restart'])(
+    'returns focus to the terminal after %s',
+    async (action) => {
+      const terminalInput = await renderPane(true)
+      returnFocus.mockImplementation(() => terminalInput.focus())
+      const dialog = container.querySelector('[role="dialog"]')!
+      await act(async () => {
+        if (action === 'outside') {
+          dialog.dispatchEvent(new MouseEvent('mousedown', { bubbles: true, button: 0 }))
+        } else if (action === 'escape') {
+          dialog.dispatchEvent(new KeyboardEvent('keydown', { key: 'Escape', bubbles: true }))
+        } else {
+          button(container, action === 'restart' ? 'Restart' : 'Keep old account').click()
+        }
+      })
+      expect(document.activeElement).toBe(terminalInput)
+      expect(container.querySelector('[role="dialog"]')).toBeNull()
+      expect(returnFocus).toHaveBeenCalledOnce()
+      if (action !== 'restart') {
+        expect(forgetStalePanes).toHaveBeenCalledExactlyOnceWith({ ptyIds: [PTY_ONE] })
+        expect(useAppStore.getState().pendingCodexPaneRestartIds).toEqual({})
+      }
+    }
+  )
+
+  it('does not dismiss when the card body is clicked', async () => {
+    await renderPane(true)
+    await act(async () => {
+      container
+        .querySelector('[role="dialog"] > div')!
+        .dispatchEvent(new MouseEvent('click', { bubbles: true }))
+    })
+    expect(container.querySelector('[role="dialog"]')).not.toBeNull()
+    expect(returnFocus).not.toHaveBeenCalled()
+    expect(forgetStalePanes).not.toHaveBeenCalled()
+  })
+
+  it('does not dismiss when a card text selection ends on the backdrop', async () => {
+    await renderPane(true)
+    const dialog = container.querySelector('[role="dialog"]')!
+    await act(async () => {
+      dialog.querySelector('div')!.dispatchEvent(new MouseEvent('mousedown', { bubbles: true }))
+      dialog.dispatchEvent(new MouseEvent('mouseup', { bubbles: true }))
+      dialog.dispatchEvent(new MouseEvent('click', { bubbles: true }))
+    })
+    expect(container.querySelector('[role="dialog"]')).not.toBeNull()
+    expect(returnFocus).not.toHaveBeenCalled()
+    expect(forgetStalePanes).not.toHaveBeenCalled()
+  })
+
+  it('does not dismiss on a secondary mouse button press', async () => {
+    await renderPane(true)
+    await act(async () => {
+      container
+        .querySelector('[role="dialog"]')!
+        .dispatchEvent(new MouseEvent('mousedown', { bubbles: true, button: 2 }))
+    })
+    expect(container.querySelector('[role="dialog"]')).not.toBeNull()
+    expect(returnFocus).not.toHaveBeenCalled()
+  })
+
   it('focuses the dialog itself for the active stale pane', async () => {
     const terminalInput = await renderPane(true)
 

src/renderer/src/components/terminal-pane/TerminalPaneRuntimePortals.tsx

@@ -1,4 +1,5 @@
 import { createPortal } from 'react-dom'
+import { focusTerminalTabSurface } from '@/lib/focus-terminal-tab-surface'
 import CodexRestartChip from '../CodexRestartChip'
 import { CodexSharedServerBanner } from './CodexSharedServerBanner'
 import { TerminalSshReconnectOverlay } from './TerminalSshReconnectOverlay'
@@ -16,8 +17,16 @@ export function TerminalPaneCodexRestartPortals({
 }: {
   controller: TerminalPaneController
 }): React.JSX.Element {
-  const { activePane, isActive, isVisible, managedPanes, paneTransportsRef, savedLayout, tabId } =
-    controller
+  const {
+    activePane,
+    isActive,
+    isVisible,
+    managedPanes,
+    managerRef,
+    paneTransportsRef,
+    savedLayout,
+    tabId
+  } = controller
   return (
     <>
       {managedPanes.map((pane) => {
@@ -34,6 +43,10 @@ export function TerminalPaneCodexRestartPortals({
               key={`codex-restart-${pane.id}-${ptyId}`}
               isVisible={isVisible}
               ptyId={ptyId}
+              onReturnFocus={() => {
+                managerRef.current?.setActivePane(pane.id, { focus: false })
+                focusTerminalTabSurface(tabId, pane.leafId)
+              }}
               shouldFocus={isActive && isVisible && activePane?.id === pane.id}
             />
             <CodexSharedServerBanner

tests/e2e/codex-account-switch-dismiss.spec.ts

@@ -0,0 +1,76 @@
+import { expect, test } from './helpers/orca-app'
+import {
+  configureGoldenStubAgent,
+  getGoldenStubAgentLaunchEnv,
+  launchGoldenStubAgentFromNewTab
+} from './helpers/golden-stub-agent'
+import { ensureTerminalVisible, waitForActiveWorktree, waitForSessionReady } from './helpers/store'
+import { waitForActivePanePtyId, waitForTerminalOutput } from './helpers/terminal'
+
+test.use({ launchEnv: getGoldenStubAgentLaunchEnv() })
+
+test('account switch actions return keyboard input to the terminal', async ({
+  orcaPage
+}, testInfo) => {
+  await waitForSessionReady(orcaPage)
+  await waitForActiveWorktree(orcaPage)
+  await ensureTerminalVisible(orcaPage)
+  await configureGoldenStubAgent(orcaPage)
+  await launchGoldenStubAgentFromNewTab(orcaPage)
+
+  for (const action of ['outside', 'keep', 'escape', 'restart'] as const) {
+    const ptyId = await waitForActivePanePtyId(orcaPage)
+    await orcaPage.evaluate(
+      ({ ptyId, action }) => {
+        window.__store.getState().markCodexRestartNotices([
+          {
+            ptyId,
+            previousAccountLabel: 'Previous account',
+            nextAccountLabel: `Current account (${action})`
+          }
+        ])
+      },
+      { ptyId, action }
+    )
+    const dialog = orcaPage.getByRole('dialog').filter({ hasText: 'Account switched' })
+    await expect(dialog).toBeVisible()
+    await expect(dialog).toBeFocused()
+    if (action === 'outside') {
+      const cardText = await dialog.getByText(/Restart this session to use/).boundingBox()
+      const backdrop = await dialog.boundingBox()
+      if (!cardText || !backdrop) {
+        throw new Error('Account prompt has no rendered bounds')
+      }
+      await orcaPage.mouse.move(cardText.x + 8, cardText.y + 8)
+      await orcaPage.mouse.down()
+      await orcaPage.mouse.move(backdrop.x + 8, backdrop.y + 8)
+      await orcaPage.mouse.up()
+      await expect(dialog).toBeVisible()
+      await orcaPage.screenshot({ path: testInfo.outputPath('account-switched.png') })
+      await dialog.click({ position: { x: 8, y: 8 } })
+    } else if (action === 'escape') {
+      await orcaPage.keyboard.press('Escape')
+    } else {
+      await dialog
+        .getByRole('button', {
+          name: action === 'keep' ? 'Keep old account' : 'Restart',
+          exact: true
+        })
+        .click()
+    }
+    await expect(dialog).toHaveCount(0)
+    const input = orcaPage.locator('.xterm-helper-textarea:focus')
+    await expect(input).toHaveCount(1)
+    if (action === 'restart') {
+      // This check covers the focus handoff; process replacement is a separate contract.
+      break
+    }
+    // A printed reply proves input reached the process without an extra terminal click.
+    const marker = `ACCOUNT_SWITCH_${action.toUpperCase()}`
+    await orcaPage.keyboard.type(marker)
+    await orcaPage.keyboard.press('Enter')
+    await waitForTerminalOutput(orcaPage, `[GOLDEN_STUB_AGENT_SUBMITTED] ${marker}`, 20_000)
+    await expect(input).toHaveCount(1)
+    await orcaPage.screenshot({ path: testInfo.outputPath(`after-${action}.png`) })
+  }
+})