From 0f240533e1cd4e41bf4e15c5489f6eb51e75f446 Mon Sep 17 00:00:00 2001 From: Brennan Benson <79079362+brennanb2025@users.noreply.github.com> Date: Wed, 17 Jun 2026 16:20:37 -0700 Subject: [PATCH] Keep Codex and Claude signed in across Orca restarts (#5639) --- .../agent-auth-restart-preservation.test.ts | 172 ++++++++++++++++++ src/main/agent-auth-restart-preservation.ts | 123 +++++++++++++ .../runtime-home-service.test.ts | 70 +++++++ .../codex-accounts/runtime-home-service.ts | 61 +++++++ src/main/index.ts | 9 +- src/main/ipc/app.test.ts | 67 ++++++- src/main/ipc/app.ts | 25 ++- src/main/ipc/register-core-handlers.ts | 2 +- src/main/updater-events.ts | 2 +- src/main/updater-mac-install.ts | 11 +- src/main/updater.mac-install.test.ts | 88 ++++++++- src/main/updater.test.ts | 48 +++++ src/main/updater.ts | 50 ++++- .../attach-main-window-services.test.ts | 35 +++- .../window/attach-main-window-services.ts | 9 +- 15 files changed, 744 insertions(+), 28 deletions(-) create mode 100644 src/main/agent-auth-restart-preservation.test.ts create mode 100644 src/main/agent-auth-restart-preservation.ts diff --git a/src/main/agent-auth-restart-preservation.test.ts b/src/main/agent-auth-restart-preservation.test.ts new file mode 100644 index 000000000..057e3d876 --- /dev/null +++ b/src/main/agent-auth-restart-preservation.test.ts @@ -0,0 +1,172 @@ +import { afterEach, describe, expect, it, vi } from 'vitest' + +import { preserveAgentAuthBeforeRestart } from './agent-auth-restart-preservation' + +describe('preserveAgentAuthBeforeRestart', () => { + afterEach(() => { + vi.useRealTimers() + vi.restoreAllMocks() + }) + + it('syncs Codex then Claude before flushing the store', async () => { + const calls: string[] = [] + + await preserveAgentAuthBeforeRestart({ + codexRuntimeHome: { + syncForCurrentSelection: vi.fn(() => { + calls.push('codex') + }), + syncActiveWslSelectionsBeforeRestart: vi.fn() + }, + claudeRuntimeAuth: { + syncForCurrentSelection: vi.fn(async () => { + calls.push('claude') + }) + }, + store: { + flush: vi.fn(() => { + calls.push('flush') + }) + } + }) + + expect(calls).toEqual(['codex', 'claude', 'flush']) + }) + + it('runs WSL Codex preservation through the runtime service', async () => { + const syncForCurrentSelection = vi.fn() + const syncActiveWslSelectionsBeforeRestart = vi.fn() + + await preserveAgentAuthBeforeRestart({ + codexRuntimeHome: { + syncForCurrentSelection, + syncActiveWslSelectionsBeforeRestart + }, + store: { + flush: vi.fn() + } + }) + + expect(syncForCurrentSelection).toHaveBeenCalledTimes(1) + expect(syncForCurrentSelection).toHaveBeenNthCalledWith(1) + expect(syncActiveWslSelectionsBeforeRestart).toHaveBeenCalledTimes(1) + }) + + it('runs Claude preservation before WSL Codex preservation', async () => { + const calls: string[] = [] + + await preserveAgentAuthBeforeRestart({ + codexRuntimeHome: { + syncForCurrentSelection: vi.fn(() => { + calls.push('codex-host') + }), + syncActiveWslSelectionsBeforeRestart: vi.fn(() => { + calls.push('codex-wsl') + }) + }, + claudeRuntimeAuth: { + syncForCurrentSelection: vi.fn(async () => { + calls.push('claude') + }) + }, + store: { + flush: vi.fn(() => { + calls.push('flush') + }) + } + }) + + expect(calls).toEqual(['codex-host', 'claude', 'codex-wsl', 'flush']) + }) + + it('continues after WSL Codex preservation fails', async () => { + const warn = vi.spyOn(console, 'warn').mockImplementation(() => {}) + const flush = vi.fn() + + await preserveAgentAuthBeforeRestart({ + codexRuntimeHome: { + syncForCurrentSelection: vi.fn(), + syncActiveWslSelectionsBeforeRestart: vi.fn(() => { + throw new Error('wsl-token-secret') + }) + }, + store: { + flush + } + }) + + expect(flush).toHaveBeenCalledTimes(1) + expect(JSON.stringify(warn.mock.calls)).not.toContain('token-secret') + }) + + it('flushes the store when auth services are missing', async () => { + const flush = vi.fn() + + await preserveAgentAuthBeforeRestart({ store: { flush } }) + + expect(flush).toHaveBeenCalledTimes(1) + }) + + it('logs secret-free warnings and does not throw when sync fails', async () => { + const warn = vi.spyOn(console, 'warn').mockImplementation(() => {}) + const flush = vi.fn() + + await expect( + preserveAgentAuthBeforeRestart({ + codexRuntimeHome: { + syncForCurrentSelection: vi.fn(() => { + throw new Error('codex-token-secret') + }), + syncActiveWslSelectionsBeforeRestart: vi.fn() + }, + claudeRuntimeAuth: { + syncForCurrentSelection: vi.fn(async () => { + throw new Error('claude-token-secret') + }) + }, + store: { flush } + }) + ).resolves.toBeUndefined() + + expect(flush).toHaveBeenCalledTimes(1) + expect(warn).toHaveBeenCalledTimes(2) + expect(JSON.stringify(warn.mock.calls)).not.toContain('token-secret') + }) + + it('releases the lifecycle path on timeout without canceling in-flight sync', async () => { + vi.useFakeTimers() + const warn = vi.spyOn(console, 'warn').mockImplementation(() => {}) + const calls: string[] = [] + let finishClaude!: () => void + + const preservation = preserveAgentAuthBeforeRestart({ + claudeRuntimeAuth: { + syncForCurrentSelection: vi.fn(async () => { + calls.push('claude-start') + await new Promise((resolve) => { + finishClaude = resolve + }) + calls.push('claude-finish') + }) + }, + store: { + flush: vi.fn(() => { + calls.push('flush') + }) + } + }) + + await vi.advanceTimersByTimeAsync(2_000) + await preservation + + expect(calls).toEqual(['claude-start', 'flush']) + expect(warn).toHaveBeenCalledWith( + '[agent-auth-restart] Claude auth preservation exceeded 2000ms; continuing restart/update' + ) + + finishClaude() + await Promise.resolve() + + expect(calls).toEqual(['claude-start', 'flush', 'claude-finish']) + }) +}) diff --git a/src/main/agent-auth-restart-preservation.ts b/src/main/agent-auth-restart-preservation.ts new file mode 100644 index 000000000..cb3fe9492 --- /dev/null +++ b/src/main/agent-auth-restart-preservation.ts @@ -0,0 +1,123 @@ +import type { ClaudeRuntimeAuthService } from './claude-accounts/runtime-auth-service' +import type { CodexRuntimeHomeService } from './codex-accounts/runtime-home-service' +import type { Store } from './persistence' + +const AUTH_PRESERVATION_TIMEOUT_MS = 2_000 + +type CodexRuntimeAuthSync = Pick< + CodexRuntimeHomeService, + 'syncForCurrentSelection' | 'syncActiveWslSelectionsBeforeRestart' +> +type ClaudeRuntimeAuthSync = Pick +type ShutdownStore = Pick + +type AuthPreservationStep = 'Codex auth preservation' | 'Claude auth preservation' + +export type AgentAuthRestartPreservationOptions = { + codexRuntimeHome?: CodexRuntimeAuthSync | null + claudeRuntimeAuth?: ClaudeRuntimeAuthSync | null + store?: ShutdownStore | null +} + +export async function preserveAgentAuthBeforeRestart({ + codexRuntimeHome, + claudeRuntimeAuth, + store +}: AgentAuthRestartPreservationOptions): Promise { + const startedAt = Date.now() + + runCodexPreservationStep(codexRuntimeHome) + + const remainingMs = Math.max(0, AUTH_PRESERVATION_TIMEOUT_MS - (Date.now() - startedAt)) + if (claudeRuntimeAuth && remainingMs > 0) { + await runWithinLifecycleTimeout( + 'Claude auth preservation', + () => claudeRuntimeAuth.syncForCurrentSelection(), + remainingMs + ) + } else if (claudeRuntimeAuth) { + logStepTimeout('Claude auth preservation', 0) + } + + if (codexRuntimeHome && Date.now() - startedAt < AUTH_PRESERVATION_TIMEOUT_MS) { + runWslCodexPreservationStep(codexRuntimeHome) + } else if (codexRuntimeHome) { + logStepTimeout('Codex auth preservation', 0) + } + + try { + store?.flush() + } catch (error) { + logStoreFlushFailure(error) + } +} + +function runCodexPreservationStep(codexRuntimeHome: CodexRuntimeAuthSync | null | undefined): void { + try { + codexRuntimeHome?.syncForCurrentSelection() + } catch (error) { + logStepFailure('Codex auth preservation', error) + } +} + +function runWslCodexPreservationStep( + codexRuntimeHome: CodexRuntimeAuthSync | null | undefined +): void { + try { + codexRuntimeHome?.syncActiveWslSelectionsBeforeRestart() + } catch (error) { + logStepFailure('Codex auth preservation', error) + } +} + +async function runWithinLifecycleTimeout( + step: AuthPreservationStep, + run: () => Promise, + timeoutMs: number +): Promise { + let timeout: ReturnType | null = null + const operation = Promise.resolve() + .then(run) + .catch((error) => { + logStepFailure(step, error) + }) + + // Why: this timeout only releases the restart/update path. It does not + // cancel a sync that already started, and Codex sync is synchronous today. + const timeoutResult = new Promise<'timeout'>((resolve) => { + timeout = setTimeout(() => resolve('timeout'), timeoutMs) + }) + + const result = await Promise.race([operation.then(() => 'done' as const), timeoutResult]) + if (result === 'timeout') { + logStepTimeout(step, timeoutMs) + return + } + + if (timeout) { + clearTimeout(timeout) + } +} + +function logStepFailure(step: AuthPreservationStep, error: unknown): void { + console.warn( + `[agent-auth-restart] ${step} failed (${describeErrorKind(error)}); continuing restart/update` + ) +} + +function logStepTimeout(step: AuthPreservationStep, timeoutMs: number): void { + console.warn(`[agent-auth-restart] ${step} exceeded ${timeoutMs}ms; continuing restart/update`) +} + +function logStoreFlushFailure(error: unknown): void { + console.warn( + `[agent-auth-restart] Store flush failed (${describeErrorKind(error)}); continuing restart/update` + ) +} + +function describeErrorKind(error: unknown): string { + if (error instanceof Error) { + return error.name || 'Error' + } + return typeof error +} diff --git a/src/main/codex-accounts/runtime-home-service.test.ts b/src/main/codex-accounts/runtime-home-service.test.ts index 0c9e2635f..1431ca67a 100644 --- a/src/main/codex-accounts/runtime-home-service.test.ts +++ b/src/main/codex-accounts/runtime-home-service.test.ts @@ -1552,6 +1552,76 @@ describe('CodexRuntimeHomeService', () => { } }) + it('reads active WSL token refreshes back before restart using the selected distro', async () => { + const originalPlatform = Object.getOwnPropertyDescriptor(process, 'platform') + Object.defineProperty(process, 'platform', { configurable: true, value: 'win32' }) + const wslHome = join(testState.userDataDir, 'wsl-home') + vi.doMock('../wsl', () => ({ + getDefaultWslDistro: () => 'Ubuntu', + getWslHome: () => wslHome + })) + const managedAuth = createCodexAuthJson('wsl@example.com', 'acct-wsl', 'managed', 1_000) + const refreshedAuth = createCodexAuthJson( + 'wsl@example.com', + 'acct-wsl', + 'runtime-refreshed', + 2_000 + ) + const managedHomePath = createManagedAuth(testState.userDataDir, 'account-1', managedAuth) + const managedAuthPath = join(managedHomePath, 'auth.json') + const store = createStore( + createSettings({ + codexManagedAccounts: [ + { + id: 'account-1', + email: 'wsl@example.com', + managedHomePath, + managedHomeRuntime: 'wsl', + wslDistro: null, + wslLinuxHomePath: '/home/alice/.local/share/orca/codex-accounts/account-1/home', + providerAccountId: 'acct-wsl', + workspaceLabel: null, + workspaceAccountId: 'acct-wsl', + createdAt: 1, + updatedAt: 1, + lastAuthenticatedAt: 1 + } + ], + activeCodexManagedAccountIdsByRuntime: { + host: null, + wsl: { Ubuntu: 'account-1' } + } + }) + ) + + try { + const { CodexRuntimeHomeService } = await import('./runtime-home-service') + const service = new CodexRuntimeHomeService(store as never) + const target = { runtime: 'wsl' as const, wslDistro: 'Ubuntu' } + const wslRuntimeHomePath = join( + wslHome, + '.local', + 'share', + 'orca', + 'codex-runtime-home', + 'home' + ) + const runtimeAuthPath = join(wslRuntimeHomePath, 'auth.json') + + expect(service.prepareForCodexLaunch(target)).toBe(wslRuntimeHomePath) + writeFileSync(runtimeAuthPath, refreshedAuth, 'utf-8') + + service.syncActiveWslSelectionsBeforeRestart() + + expect(readFileSync(managedAuthPath, 'utf-8')).toBe(refreshedAuth) + expect(readFileSync(runtimeAuthPath, 'utf-8')).toBe(refreshedAuth) + } finally { + if (originalPlatform) { + Object.defineProperty(process, 'platform', originalPlatform) + } + } + }) + it('uses the stable WSL runtime home for WSL system-default rate-limit fetches', async () => { const originalPlatform = Object.getOwnPropertyDescriptor(process, 'platform') Object.defineProperty(process, 'platform', { configurable: true, value: 'win32' }) diff --git a/src/main/codex-accounts/runtime-home-service.ts b/src/main/codex-accounts/runtime-home-service.ts index 4c5723f95..7024b8fa9 100644 --- a/src/main/codex-accounts/runtime-home-service.ts +++ b/src/main/codex-accounts/runtime-home-service.ts @@ -43,6 +43,7 @@ import { syncSystemCodexSessionsIntoManagedHome } from '../codex/codex-session-b import { syncSystemConfigIntoManagedCodexHome } from '../codex/codex-config-mirror' import { parseWslUncPath } from '../../shared/wsl-paths' import { + getWslSelectionKey, getSelectedCodexAccountIdForTarget, normalizeCodexRuntimeSelection, setSelectedCodexAccountIdForTarget, @@ -97,6 +98,7 @@ export class CodexRuntimeHomeService { // newer than managed storage. private readonly lastWrittenWslAuthJsonByDistro = new Map() private readonly lastSyncedWslAccountIdByDistro = new Map() + private readonly wslRuntimeHomePathByDistro = new Map() private skipNextReadBackForAccountId: string | null = null constructor(private readonly store: Store) { @@ -139,6 +141,26 @@ export class CodexRuntimeHomeService { return this.getRuntimeHomePath() } + syncActiveWslSelectionsBeforeRestart(): void { + if (process.platform !== 'win32') { + return + } + + const settings = this.store.getSettings() + for (const [selectedDistroKey, accountId] of Object.entries( + normalizeCodexRuntimeSelection(settings).wsl + )) { + if (!accountId) { + continue + } + const account = this.getActiveAccount(settings.codexManagedAccounts, accountId) + if (!account || account.managedHomeRuntime !== 'wsl') { + continue + } + this.safeReadBackActiveWslAccountBeforeRestart(account, selectedDistroKey) + } + } + private getWslSystemCodexHomePath(target: CodexAccountSelectionTarget): string | null { if (process.platform !== 'win32') { return null @@ -452,6 +474,7 @@ export class CodexRuntimeHomeService { if (!runtimeHomePath) { return null } + this.wslRuntimeHomePathByDistro.set(distro, runtimeHomePath) mkdirSync(runtimeHomePath, { recursive: true }) this.safeMigrateLegacyWslActiveHomePointer(distro, runtimeHomePath) @@ -544,6 +567,44 @@ export class CodexRuntimeHomeService { : null } + private safeReadBackActiveWslAccountBeforeRestart( + account: CodexManagedAccount, + selectedDistroKey: string + ): void { + try { + this.readBackActiveWslAccountBeforeRestart(account, selectedDistroKey) + } catch (error) { + console.warn('[codex-runtime-home] Failed to preserve WSL Codex auth before restart:', error) + } + } + + private readBackActiveWslAccountBeforeRestart( + account: CodexManagedAccount, + selectedDistroKey: string + ): void { + const distro = + selectedDistroKey === getWslSelectionKey(null) + ? account.wslDistro?.trim() + : selectedDistroKey.trim() || account.wslDistro?.trim() + if (!distro) { + return + } + + const runtimeHomePath = this.wslRuntimeHomePathByDistro.get(distro) + if (!runtimeHomePath) { + return + } + + this.readBackRefreshedTokensFromPath(join(runtimeHomePath, 'auth.json'), { + updateLastWrittenAuthJson: true, + lastWrittenAuthJson: this.lastWrittenWslAuthJsonByDistro.get(distro) ?? null, + setLastWrittenAuthJson: (contents) => { + this.lastWrittenWslAuthJsonByDistro.set(distro, contents) + }, + expectedAccountId: account.id + }) + } + private safeMigrateLegacyWslActiveHomePointer(distro: string, runtimeHomePath: string): void { try { this.migrateLegacyWslActiveHomePointer(distro, runtimeHomePath) diff --git a/src/main/index.ts b/src/main/index.ts index 892b74a6f..a96483bc8 100644 --- a/src/main/index.ts +++ b/src/main/index.ts @@ -143,6 +143,7 @@ import { import type { AgentStatusState } from '../shared/agent-status-types' import { KeybindingService } from './keybindings/keybinding-service' import { applyElectronProxySettings } from './network/proxy-settings' +import { preserveAgentAuthBeforeRestart } from './agent-auth-restart-preservation' let mainWindow: BrowserWindow | null = null /** Whether a manual app.quit() (Cmd+Q, etc.) is in progress. Shared with the @@ -712,9 +713,9 @@ function openMainWindow(): BrowserWindow { { getAdditionalAiVaultCodexHomePaths: () => codexRuntimeHome ? [codexRuntimeHome.getHostRuntimeHomePath()] : [], - onBeforeRelaunch: () => { + onBeforeRelaunch: async () => { isQuitting = true - store?.flush() + await preserveAgentAuthBeforeRestart({ codexRuntimeHome, claudeRuntimeAuth, store }) } } ) @@ -733,7 +734,9 @@ function openMainWindow(): BrowserWindow { markExpectedRendererReload(webContentsId) } recordCrashBreadcrumb('renderer_reload_requested', { ignoreCache }) - } + }, + onBeforeUpdateQuit: () => + preserveAgentAuthBeforeRestart({ codexRuntimeHome, claudeRuntimeAuth, store }) } ) rateLimits.attach(window) diff --git a/src/main/ipc/app.test.ts b/src/main/ipc/app.test.ts index 125658aca..76cc5fb69 100644 --- a/src/main/ipc/app.test.ts +++ b/src/main/ipc/app.test.ts @@ -57,34 +57,89 @@ describe('registerAppHandlers', () => { Object.defineProperty(process, 'platform', { value: originalPlatform, configurable: true }) }) - it('marks relaunch as expected shutdown before exiting', () => { + it('marks relaunch as expected shutdown before exiting', async () => { const onBeforeRelaunch = vi.fn() registerAppHandlers({} as never, { onBeforeRelaunch }) - handlers.get('app:relaunch')?.(null) + const relaunchPromise = Promise.resolve(handlers.get('app:relaunch')?.(null)) expect(onBeforeRelaunch).toHaveBeenCalledTimes(1) expect(appRelaunchMock).not.toHaveBeenCalled() expect(appExitMock).not.toHaveBeenCalled() - vi.advanceTimersByTime(150) + await relaunchPromise + await vi.advanceTimersByTimeAsync(150) expect(appRelaunchMock).toHaveBeenCalledTimes(1) expect(appExitMock).toHaveBeenCalledWith(0) }) - it('marks restart as expected shutdown before quitting through the normal pipeline', () => { + it('waits for pre-relaunch cleanup before exiting', async () => { + let finishCleanup!: () => void + const onBeforeRelaunch = vi.fn( + () => + new Promise((resolve) => { + finishCleanup = resolve + }) + ) + registerAppHandlers({} as never, { onBeforeRelaunch }) + + const relaunchPromise = Promise.resolve(handlers.get('app:relaunch')?.(null)) + + expect(onBeforeRelaunch).toHaveBeenCalledTimes(1) + + await vi.advanceTimersByTimeAsync(150) + expect(appRelaunchMock).not.toHaveBeenCalled() + expect(appExitMock).not.toHaveBeenCalled() + + finishCleanup() + await relaunchPromise + await vi.advanceTimersByTimeAsync(150) + + expect(appRelaunchMock).toHaveBeenCalledTimes(1) + expect(appExitMock).toHaveBeenCalledWith(0) + }) + + it('marks restart as expected shutdown before quitting through the normal pipeline', async () => { const onBeforeRelaunch = vi.fn() registerAppHandlers({} as never, { onBeforeRelaunch }) - handlers.get('app:restart')?.(null) + const restartPromise = Promise.resolve(handlers.get('app:restart')?.(null)) expect(onBeforeRelaunch).toHaveBeenCalledTimes(1) expect(appRelaunchMock).not.toHaveBeenCalled() expect(appQuitMock).not.toHaveBeenCalled() expect(appExitMock).not.toHaveBeenCalled() - vi.advanceTimersByTime(150) + await restartPromise + await vi.advanceTimersByTimeAsync(150) + + expect(appRelaunchMock).toHaveBeenCalledTimes(1) + expect(appQuitMock).toHaveBeenCalledTimes(1) + expect(appExitMock).not.toHaveBeenCalled() + }) + + it('waits for pre-relaunch cleanup before restarting through the normal pipeline', async () => { + let finishCleanup!: () => void + const onBeforeRelaunch = vi.fn( + () => + new Promise((resolve) => { + finishCleanup = resolve + }) + ) + registerAppHandlers({} as never, { onBeforeRelaunch }) + + const restartPromise = Promise.resolve(handlers.get('app:restart')?.(null)) + + expect(onBeforeRelaunch).toHaveBeenCalledTimes(1) + + await vi.advanceTimersByTimeAsync(150) + expect(appRelaunchMock).not.toHaveBeenCalled() + expect(appQuitMock).not.toHaveBeenCalled() + + finishCleanup() + await restartPromise + await vi.advanceTimersByTimeAsync(150) expect(appRelaunchMock).toHaveBeenCalledTimes(1) expect(appQuitMock).toHaveBeenCalledTimes(1) diff --git a/src/main/ipc/app.ts b/src/main/ipc/app.ts index c34e4b5c0..7424c4db8 100644 --- a/src/main/ipc/app.ts +++ b/src/main/ipc/app.ts @@ -23,7 +23,7 @@ import { isMarkdownDocumentName, markdownDocumentFromFilePath } from './markdown const KEYBOARD_INPUT_SOURCE_TIMEOUT_MS = 500 type RegisterAppHandlersOptions = { - onBeforeRelaunch?: () => void + onBeforeRelaunch?: () => void | Promise } async function pickFloatingMarkdownDocument( @@ -206,24 +206,24 @@ export function registerAppHandlers(store: Store, options: RegisterAppHandlersOp } }) - ipcMain.handle('app:relaunch', () => { + ipcMain.handle('app:relaunch', async () => { // Why: small delay lets the renderer finish painting any "Restarting…" // UI state before the window tears down. `app.relaunch()` schedules a // spawn; `app.exit(0)` triggers the actual quit without invoking // before-quit handlers that could block on confirmation dialogs. // Mark shutdown first because app.exit() can bypass the usual quit latch. - options.onBeforeRelaunch?.() + await runBeforeRelaunchCleanup(options.onBeforeRelaunch) setTimeout(() => { app.relaunch() app.exit(0) }, 150) }) - ipcMain.handle('app:restart', () => { + ipcMain.handle('app:restart', async () => { // Why: the hidden admin restart should mirror the update relaunch path: // schedule a new Orca process, then use the normal quit pipeline so daemon // checkpoints, runtime metadata, and telemetry flush before exit. - options.onBeforeRelaunch?.() + await runBeforeRelaunchCleanup(options.onBeforeRelaunch) setTimeout(() => { app.relaunch() app.quit() @@ -246,3 +246,18 @@ export function registerAppHandlers(store: Store, options: RegisterAppHandlersOp pickFloatingWorkspaceDirectory(event, store) ) } + +async function runBeforeRelaunchCleanup( + onBeforeRelaunch?: () => void | Promise +): Promise { + try { + await onBeforeRelaunch?.() + } catch (error) { + // Why: restart/relaunch must not get trapped if best-effort shutdown + // cleanup fails; the cleanup path logs without exposing secret contents. + console.warn( + '[app] Pre-relaunch cleanup failed; continuing relaunch:', + error instanceof Error ? error.name : typeof error + ) + } +} diff --git a/src/main/ipc/register-core-handlers.ts b/src/main/ipc/register-core-handlers.ts index d92e7a338..b496659e6 100644 --- a/src/main/ipc/register-core-handlers.ts +++ b/src/main/ipc/register-core-handlers.ts @@ -65,7 +65,7 @@ import type { KeybindingService } from '../keybindings/keybinding-service' let registered = false type CoreHandlerLifecycleOptions = { - onBeforeRelaunch?: () => void + onBeforeRelaunch?: () => void | Promise getAdditionalAiVaultCodexHomePaths?: () => readonly string[] } diff --git a/src/main/updater-events.ts b/src/main/updater-events.ts index 5b979b3bb..50d973621 100644 --- a/src/main/updater-events.ts +++ b/src/main/updater-events.ts @@ -28,7 +28,7 @@ type UpdaterHandlerContext = { getUserInitiatedCheck: () => boolean hasNewerDownloadedVersion: () => boolean markMissingManifestPrereleaseFallbackChecking: () => void - performQuitAndInstall: () => void + performQuitAndInstall: () => void | Promise recordCompletedUpdateCheck: () => void sendCheckFailureStatus: ( message: string, diff --git a/src/main/updater-mac-install.ts b/src/main/updater-mac-install.ts index 370741241..086e7d5fc 100644 --- a/src/main/updater-mac-install.ts +++ b/src/main/updater-mac-install.ts @@ -112,14 +112,21 @@ export function deferMacQuitUntilInstallerReady( export function handleMacInstallerReady( hasNewerDownloadedVersion: boolean, - onReadyToInstall: () => void, + onReadyToInstall: () => void | Promise, onReadyToReportDownloaded: () => void ): void { squirrelReady = true clearPendingInstallTimeout() if (installRequestedAfterSquirrelReady && hasNewerDownloadedVersion) { - onReadyToInstall() + void Promise.resolve() + .then(() => onReadyToInstall()) + .catch((error) => { + console.warn( + '[updater] Deferred macOS install handoff failed:', + error instanceof Error ? error.name : typeof error + ) + }) return } diff --git a/src/main/updater.mac-install.test.ts b/src/main/updater.mac-install.test.ts index c4baf8cec..b47b3181a 100644 --- a/src/main/updater.mac-install.test.ts +++ b/src/main/updater.mac-install.test.ts @@ -164,7 +164,9 @@ describe('updater mac install handoff', () => { nativeDownloadedHandler?.() - expect(autoUpdaterMock.quitAndInstall).toHaveBeenCalledWith(false, true) + await vi.waitFor(() => { + expect(autoUpdaterMock.quitAndInstall).toHaveBeenCalledWith(false, true) + }) expect(sendMock).toHaveBeenCalledWith('updater:status', { state: 'downloading', percent: 100, @@ -173,6 +175,90 @@ describe('updater mac install handoff', () => { } ) + it.runIf(process.platform === 'darwin')( + 'ignores duplicate quit requests while deferred mac install cleanup is running', + async () => { + vi.useFakeTimers() + + let finishCleanup!: () => void + const onBeforeQuit = vi.fn( + () => + new Promise((resolve) => { + finishCleanup = resolve + }) + ) + const mainWindow = { webContents: { send: vi.fn() } } + + autoUpdaterMock.checkForUpdates.mockResolvedValue(undefined) + const { setupAutoUpdater, quitAndInstall } = await import('./updater') + + setupAutoUpdater(mainWindow as never, { onBeforeQuit }) + autoUpdaterMock.emit('update-available', { version: '1.0.61' }) + await vi.advanceTimersByTimeAsync(0) + autoUpdaterMock.emit('update-downloaded', { version: '1.0.61' }) + + const preventDefault = vi.fn() + appMock.emit('before-quit', { preventDefault }) + expect(preventDefault).toHaveBeenCalledTimes(1) + + const nativeDownloadedHandler = nativeUpdaterMock.on.mock.calls.find( + ([eventName]) => eventName === 'update-downloaded' + )?.[1] as (() => void) | undefined + expect(nativeDownloadedHandler).toBeTypeOf('function') + + nativeDownloadedHandler?.() + await vi.advanceTimersByTimeAsync(0) + + expect(onBeforeQuit).toHaveBeenCalledTimes(1) + expect(autoUpdaterMock.quitAndInstall).not.toHaveBeenCalled() + + quitAndInstall() + finishCleanup() + await vi.advanceTimersByTimeAsync(0) + + expect(onBeforeQuit).toHaveBeenCalledTimes(1) + expect(killAllPtyMock).toHaveBeenCalledTimes(1) + expect(autoUpdaterMock.quitAndInstall).toHaveBeenCalledTimes(1) + } + ) + + it.runIf(process.platform === 'darwin')( + 'logs rejected deferred mac install handoffs without unhandled rejection', + async () => { + const warn = vi.spyOn(console, 'warn').mockImplementation(() => {}) + const reportDownloaded = vi.fn() + const { deferMacQuitUntilInstallerReady, handleMacInstallerReady } = + await import('./updater-mac-install') + + expect( + deferMacQuitUntilInstallerReady( + { state: 'downloading', percent: 100, version: '1.0.61' }, + true, + () => '1.0.61', + vi.fn() + ) + ).toBe(true) + + handleMacInstallerReady( + true, + async () => { + throw new Error('handoff-secret') + }, + reportDownloaded + ) + await Promise.resolve() + + expect(reportDownloaded).not.toHaveBeenCalled() + await vi.waitFor(() => { + expect(warn).toHaveBeenCalledWith( + '[updater] Deferred macOS install handoff failed:', + 'Error' + ) + }) + expect(JSON.stringify(warn.mock.calls)).not.toContain('handoff-secret') + } + ) + it.runIf(process.platform === 'darwin')( 'falls back to a normal quit if Squirrel.Mac never becomes ready', async () => { diff --git a/src/main/updater.test.ts b/src/main/updater.test.ts index 0d7aae865..a9ff22a5b 100644 --- a/src/main/updater.test.ts +++ b/src/main/updater.test.ts @@ -559,6 +559,25 @@ describe('updater', () => { expect(autoUpdaterMock.quitAndInstall).toHaveBeenCalledWith(false, true) }) + it('runs pre-quit cleanup before killing PTYs during update install', async () => { + vi.useFakeTimers() + + const onBeforeQuit = vi.fn() + const mainWindow = { webContents: { send: vi.fn() } } + const { setupAutoUpdater, quitAndInstall } = await import('./updater') + + setupAutoUpdater(mainWindow as never, { onBeforeQuit }) + quitAndInstall() + + await vi.advanceTimersByTimeAsync(100) + + expect(onBeforeQuit).toHaveBeenCalledTimes(1) + expect(killAllPtyMock).toHaveBeenCalledTimes(1) + expect(onBeforeQuit.mock.invocationCallOrder[0]).toBeLessThan( + killAllPtyMock.mock.invocationCallOrder[0] + ) + }) + it('ignores duplicate quitAndInstall requests while the shared delay is pending', async () => { vi.useFakeTimers() @@ -574,6 +593,35 @@ describe('updater', () => { expect(autoUpdaterMock.quitAndInstall).toHaveBeenCalledTimes(1) }) + it('ignores duplicate quitAndInstall requests while async pre-quit cleanup is running', async () => { + vi.useFakeTimers() + + let finishCleanup!: () => void + const onBeforeQuit = vi.fn( + () => + new Promise((resolve) => { + finishCleanup = resolve + }) + ) + const mainWindow = { webContents: { send: vi.fn() } } + const { setupAutoUpdater, quitAndInstall } = await import('./updater') + + setupAutoUpdater(mainWindow as never, { onBeforeQuit }) + quitAndInstall() + + await vi.advanceTimersByTimeAsync(100) + + expect(onBeforeQuit).toHaveBeenCalledTimes(1) + expect(autoUpdaterMock.quitAndInstall).not.toHaveBeenCalled() + + quitAndInstall() + finishCleanup() + await vi.advanceTimersByTimeAsync(0) + + expect(onBeforeQuit).toHaveBeenCalledTimes(1) + expect(autoUpdaterMock.quitAndInstall).toHaveBeenCalledTimes(1) + }) + it('runs a startup check immediately when the last background check is stale', async () => { const mainWindow = { webContents: { send: vi.fn() } } const setLastUpdateCheckAt = vi.fn() diff --git a/src/main/updater.ts b/src/main/updater.ts index 7669e422d..ae5cbca76 100644 --- a/src/main/updater.ts +++ b/src/main/updater.ts @@ -35,11 +35,12 @@ const AUTO_UPDATE_RETRY_INTERVAL_MS = 60 * 60 * 1000 const NUDGE_POLL_INTERVAL_MS = 30 * 60 * 1000 const NUDGE_ACTIVATION_COOLDOWN_MS = 5 * 60 * 1000 const QUIT_AND_INSTALL_DELAY_MS = 100 +const PRE_QUIT_CLEANUP_TIMEOUT_MS = 2_500 let mainWindowRef: BrowserWindow | null = null let currentStatus: UpdateStatus = { state: 'idle' } let userInitiatedCheck = false -let onBeforeQuitCleanup: (() => void) | null = null +let onBeforeQuitCleanup: (() => void | Promise) | null = null let autoUpdaterInitialized = false // Why: Shift-clicking "Check for Updates" opts the user into the RC release // channel for the rest of this process. The generic feed still gets pinned to @@ -52,6 +53,7 @@ let pendingCheckFailurePromise: Promise | null = null let autoUpdateCheckTimer: ReturnType | null = null let nudgeCheckTimer: ReturnType | null = null let pendingQuitAndInstallTimer: ReturnType | null = null +let quitAndInstallInProgress = false let persistLastUpdateCheckAt: ((timestamp: number) => void) | null = null let _getLastUpdateCheckAt: (() => number | null) | null = null let backgroundCheckLaunchPending = false @@ -283,7 +285,12 @@ function clearPrereleaseFallbackContextIfSettled(): void { } } -function performQuitAndInstall(): void { +async function performQuitAndInstall(): Promise { + if (quitAndInstallInProgress) { + return + } + quitAndInstallInProgress = true + if (pendingQuitAndInstallTimer) { clearTimeout(pendingQuitAndInstallTimer) pendingQuitAndInstallTimer = null @@ -299,8 +306,8 @@ function performQuitAndInstall(): void { // either can't replace it or the user ends up on the old version. quittingForUpdate = true + await runBeforeUpdateQuitCleanup() killAllPty() - onBeforeQuitCleanup?.() for (const win of BrowserWindow.getAllWindows()) { win.removeAllListeners('close') @@ -309,6 +316,37 @@ function performQuitAndInstall(): void { getAutoUpdater().quitAndInstall(false, true) } +async function runBeforeUpdateQuitCleanup(): Promise { + if (!onBeforeQuitCleanup) { + return + } + + let timeout: ReturnType | null = null + const cleanup = Promise.resolve() + .then(() => onBeforeQuitCleanup?.()) + .catch((error) => { + console.warn( + '[updater] Pre-quit cleanup failed; continuing update install:', + error instanceof Error ? error.name : typeof error + ) + }) + const timeoutResult = new Promise<'timeout'>((resolve) => { + timeout = setTimeout(() => resolve('timeout'), PRE_QUIT_CLEANUP_TIMEOUT_MS) + }) + + const result = await Promise.race([cleanup.then(() => 'done' as const), timeoutResult]) + if (result === 'timeout') { + console.warn( + `[updater] Pre-quit cleanup exceeded ${PRE_QUIT_CLEANUP_TIMEOUT_MS}ms; continuing update install` + ) + return + } + + if (timeout) { + clearTimeout(timeout) + } +} + async function sendCheckFailureStatus( message: string, userInitiated?: boolean, @@ -742,7 +780,7 @@ export function isQuittingForUpdate(): boolean { } export function quitAndInstall(): void { - if (pendingQuitAndInstallTimer) { + if (pendingQuitAndInstallTimer || quitAndInstallInProgress) { return } @@ -762,7 +800,7 @@ export function quitAndInstall(): void { // a moment to flush dismissals/state updates before windows start closing, // and centralizing it avoids drift between the toast flow and settings UI. pendingQuitAndInstallTimer = setTimeout(() => { - performQuitAndInstall() + void performQuitAndInstall() }, QUIT_AND_INSTALL_DELAY_MS) } @@ -835,7 +873,7 @@ export function setupAutoUpdater( mainWindow: BrowserWindow, opts?: { getLastUpdateCheckAt?: () => number | null - onBeforeQuit?: () => void + onBeforeQuit?: () => void | Promise setLastUpdateCheckAt?: (timestamp: number) => void getPendingUpdateNudgeId?: () => string | null getDismissedUpdateNudgeId?: () => string | null diff --git a/src/main/window/attach-main-window-services.test.ts b/src/main/window/attach-main-window-services.test.ts index 7be18282a..e6d04fe6b 100644 --- a/src/main/window/attach-main-window-services.test.ts +++ b/src/main/window/attach-main-window-services.test.ts @@ -1,5 +1,6 @@ /* eslint-disable max-lines -- Why: attachMainWindowServices centralizes main-window IPC wiring; keeping its integration-style mocks together avoids brittle cross-file setup. */ import { beforeEach, describe, expect, it, vi } from 'vitest' +import type { Store } from '../persistence' const { onMock, @@ -133,8 +134,8 @@ function createMainWindow(extraWebContents: { on?: MockFn; send?: MockFn } = {}) } } -function createStore(): never { - return { flush: vi.fn() } as never +function createStore(): Store & { flush: MockFn } { + return { flush: vi.fn() } as Store & { flush: MockFn } } function createRuntime(): RuntimeStub { @@ -231,6 +232,36 @@ describe('attachMainWindowServices', () => { expect(hydrateLocalPtyRegistryAtBootMock).toHaveBeenLastCalledWith(store) }) + it('passes injected update quit cleanup to the auto-updater', async () => { + const onBeforeUpdateQuit = vi.fn() + const store = createStore() + + attachMainWindowServices( + createMainWindow() as never, + store, + createRuntime() as never, + undefined, + undefined, + { onBeforeUpdateQuit } + ) + + expect(setupAutoUpdaterMock).toHaveBeenCalledTimes(1) + await setupAutoUpdaterMock.mock.calls[0][1].onBeforeQuit() + + expect(onBeforeUpdateQuit).toHaveBeenCalledTimes(1) + expect(store.flush).toHaveBeenCalledTimes(1) + }) + + it('flushes the store before update quit when no cleanup is injected', async () => { + const store = createStore() + + attachMainWindowServices(createMainWindow() as never, store, createRuntime() as never) + + await setupAutoUpdaterMock.mock.calls[0][1].onBeforeQuit() + + expect(store.flush).toHaveBeenCalledTimes(1) + }) + it('ignores app reload requests from non-main webContents', async () => { const onBeforeRendererReload = vi.fn() const mainWindow = createMainWindow() diff --git a/src/main/window/attach-main-window-services.ts b/src/main/window/attach-main-window-services.ts index 16df825d8..8548e9f8f 100644 --- a/src/main/window/attach-main-window-services.ts +++ b/src/main/window/attach-main-window-services.ts @@ -53,6 +53,7 @@ export function attachMainWindowServices( options?: { awaitLocalPtyStartup?: () => Promise onBeforeRendererReload?: (args: { webContentsId: number; ignoreCache: boolean }) => void + onBeforeUpdateQuit?: () => void | Promise } ): void { registerAppReloadHandler(mainWindow, options?.onBeforeRendererReload) @@ -109,7 +110,13 @@ export function attachMainWindowServices( registerFileDropRelay(mainWindow) setupAutoUpdater(mainWindow, { getLastUpdateCheckAt: () => store.getUI().lastUpdateCheckAt, - onBeforeQuit: () => store.flush(), + onBeforeQuit: async () => { + try { + await options?.onBeforeUpdateQuit?.() + } finally { + store.flush() + } + }, setLastUpdateCheckAt: (timestamp) => { store.updateUI({ lastUpdateCheckAt: timestamp }) },