fix(serve): save Codex auth before a headless serve update restart (#28001)

* fix(serve): save Codex auth before a headless serve update restart

* refactor(updater): keep the update-restart step in its own module so auth mocks still apply

---------

Co-authored-by: m4air <m4air@Mac.localdomain>

src/main/agent-auth-restart-preservation.ts

@@ -3,7 +3,7 @@ import type { Store } from './persistence'
 
 const AUTH_PRESERVATION_TIMEOUT_MS = 2_000
 
-type CodexRuntimeAuthSync = Pick<CodexRuntimeHomeService, 'syncForCurrentSelection'> &
+export type CodexRuntimeAuthSync = Pick<CodexRuntimeHomeService, 'syncForCurrentSelection'> &
   Partial<Pick<CodexRuntimeHomeService, 'syncActiveWslSelectionsBeforeRestart'>>
 type ShutdownStore = Pick<Store, 'flushPendingOrThrowAsync'>
 

src/main/startup/main-process-runtime-launch-activation.test.ts

@@ -90,6 +90,8 @@ vi.mock('../ipc/startup-notification-registration', () => ({
 }))
 vi.mock('./main-process-push-startup', () => ({ startDesktopPushService: vi.fn() }))
 vi.mock('./startup-diagnostics', () => ({ logStartupMilestone: vi.fn() }))
+const setupHeadlessServeAutoUpdater = vi.hoisted(() => vi.fn())
+vi.mock('../window/main-window-updater', () => ({ setupHeadlessServeAutoUpdater }))
 vi.mock('../server/serve-stdout-boundary', () => ({ emitServeBrowserIdentityActionLine: vi.fn() }))
 vi.mock('../browser/browser-identity-mode-store', () => ({
   getBrowserIdentityModeStatus: vi.fn()
@@ -309,4 +311,59 @@ describe('desktop startup activation', () => {
       })
     )
   })
+  it('saves Codex auth before a headless serve update restart, like the desktop does', async () => {
+    setupHeadlessServeAutoUpdater.mockClear()
+    state.isServeMode = true
+    vi.mocked(getServeOptions).mockReturnValueOnce({
+      json: false,
+      pairingAddress: null,
+      noPairing: true,
+      mobilePairing: false,
+      grantDesktopControl: false,
+      recipeJson: false,
+      projectRoot: null
+    })
+    Object.assign(state.runtime!, {
+      refreshRestoredOrchestrationAuthority: vi.fn(async () => {}),
+      reconcileLegacyWorkerTerminals: vi.fn(async () => {})
+    })
+    const order: string[] = []
+    const store = {
+      flushPendingOrThrowAsync: vi.fn(async (options?: { fullCheckpoint?: boolean }) => {
+        order.push(options?.fullCheckpoint ? 'store-checkpoint' : 'store-flush')
+      })
+    }
+    const codexRuntimeHome = {
+      syncForCurrentSelection: vi.fn(() => {
+        order.push('codex-auth')
+      }),
+      syncActiveWslSelectionsBeforeRestart: vi.fn(async () => {})
+    }
+    // oxlint-disable-next-line typescript/consistent-type-assertions -- SAFETY: the update restart hook only flushes the store.
+    state.store = store as unknown as NonNullable<typeof state.store>
+    // oxlint-disable-next-line typescript/consistent-type-assertions -- SAFETY: the update restart hook only calls the two sync methods.
+    state.codexRuntimeHome = codexRuntimeHome as unknown as NonNullable<
+      typeof state.codexRuntimeHome
+    >
+    try {
+      await initializeMainProcessReady({
+        openMainWindow: () => {
+          throw new Error('serve must not open a window')
+        },
+        handleMacAppActivation: vi.fn()
+      })
+
+      expect(setupHeadlessServeAutoUpdater).toHaveBeenCalledTimes(1)
+      const hooks = setupHeadlessServeAutoUpdater.mock.calls[0][1]
+      expect(hooks.onBeforeUpdateQuitFailure).toBe('abort')
+      await hooks.onBeforeUpdateQuit()
+      expect(codexRuntimeHome.syncForCurrentSelection).toHaveBeenCalledTimes(1)
+      expect(codexRuntimeHome.syncActiveWslSelectionsBeforeRestart).toHaveBeenCalledTimes(1)
+      expect(order.at(0)).toBe('codex-auth')
+      expect(order.at(-1)).toBe('store-checkpoint')
+    } finally {
+      state.store = null
+      state.codexRuntimeHome = null
+    }
+  })
 })

src/main/startup/main-process-runtime-launch.ts

@@ -41,6 +41,7 @@ import { startDesktopPushService } from './main-process-push-startup'
 import { mainProcessState as state } from './main-process-state'
 import { logStartupMilestone } from './startup-diagnostics'
 import { setupHeadlessServeAutoUpdater } from '../window/main-window-updater'
+import { preserveStateBeforeUpdateRestart } from '../update-restart-preservation'
 import { resolveUpdateInstallMode } from '../updater'
 import { scheduleAgentLaunchRecordWarmup } from './agent-launch-record-warmup'
 import { emitServeBrowserIdentityActionLine } from '../server/serve-stdout-boundary'
@@ -219,7 +220,8 @@ async function launchServeMode(
   await printServeReady(serveOptions)
   const store = state.store!
   setupHeadlessServeAutoUpdater(store, {
-    onBeforeUpdateQuit: () => store.flushPendingOrThrowAsync({ fullCheckpoint: true }),
+    onBeforeUpdateQuit: () =>
+      preserveStateBeforeUpdateRestart({ codexRuntimeHome: state.codexRuntimeHome, store }),
     onBeforeUpdateQuitFailure: 'abort',
     updateInstallMode: resolveUpdateInstallMode(true)
   })

src/main/startup/main-window-core-services.ts

@@ -7,6 +7,7 @@ import { mainProcessState as state } from './main-process-state'
 import { prepareCodexAiVaultSessionResume } from '../codex/codex-ai-vault-session-resume'
 import { resolveHostCodexSessionSourceHome } from '../codex/codex-session-source-home'
 import { preserveAgentAuthBeforeRestart } from '../agent-auth-restart-preservation'
+import { preserveStateBeforeUpdateRestart } from '../update-restart-preservation'
 import {
   emitPluginWorktreeLifecycle,
   handleCodexHomePtySpawned,
@@ -134,10 +135,7 @@ export function attachMainWindowCoreServices(
       isRecoveryReloadInFlight,
       onCodexHomePtySpawned: handleCodexHomePtySpawned,
       onPtyExit: handlePtyExit,
-      onBeforeUpdateQuit: async () => {
-        await preserveAgentAuthBeforeRestart({ codexRuntimeHome, store })
-        await store.flushPendingOrThrowAsync({ fullCheckpoint: true })
-      },
+      onBeforeUpdateQuit: () => preserveStateBeforeUpdateRestart({ codexRuntimeHome, store }),
       onBeforeUpdateQuitFailure: 'abort',
       updateInstallMode: resolveUpdateInstallMode(state.isServeMode),
       onWorktreeLifecycle: emitPluginWorktreeLifecycle

src/main/update-restart-preservation.ts

@@ -0,0 +1,17 @@
+import {
+  preserveAgentAuthBeforeRestart,
+  type CodexRuntimeAuthSync
+} from './agent-auth-restart-preservation'
+import type { Store } from './persistence'
+
+/** Before any update restart, desktop or headless serve: save agent auth, then a full store checkpoint. */
+export async function preserveStateBeforeUpdateRestart({
+  codexRuntimeHome,
+  store
+}: {
+  codexRuntimeHome: CodexRuntimeAuthSync | null | undefined
+  store: Pick<Store, 'flushPendingOrThrowAsync'>
+}): Promise<void> {
+  await preserveAgentAuthBeforeRestart({ codexRuntimeHome, store })
+  await store.flushPendingOrThrowAsync({ fullCheckpoint: true })
+}