diff --git a/src/main/agent-auth-restart-preservation.ts b/src/main/agent-auth-restart-preservation.ts index a41be9a0183..6721b16dbf0 100644 --- a/src/main/agent-auth-restart-preservation.ts +++ b/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 & +export type CodexRuntimeAuthSync = Pick & Partial> type ShutdownStore = Pick diff --git a/src/main/startup/main-process-runtime-launch-activation.test.ts b/src/main/startup/main-process-runtime-launch-activation.test.ts index 3b40deb6021..80bd815606e 100644 --- a/src/main/startup/main-process-runtime-launch-activation.test.ts +++ b/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 + // 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 + } + }) }) diff --git a/src/main/startup/main-process-runtime-launch.ts b/src/main/startup/main-process-runtime-launch.ts index d9429a61d11..2789aca517c 100644 --- a/src/main/startup/main-process-runtime-launch.ts +++ b/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) }) diff --git a/src/main/startup/main-window-core-services.ts b/src/main/startup/main-window-core-services.ts index 7185b743e89..c350a4e573f 100644 --- a/src/main/startup/main-window-core-services.ts +++ b/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 diff --git a/src/main/update-restart-preservation.ts b/src/main/update-restart-preservation.ts new file mode 100644 index 00000000000..205b5e8b48e --- /dev/null +++ b/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 +}): Promise { + await preserveAgentAuthBeforeRestart({ codexRuntimeHome, store }) + await store.flushPendingOrThrowAsync({ fullCheckpoint: true }) +}