diff --git a/packages/cli/src/utils/handleAutoUpdate.test.ts b/packages/cli/src/utils/handleAutoUpdate.test.ts index 0af2de37b1..5317bf00e4 100644 --- a/packages/cli/src/utils/handleAutoUpdate.test.ts +++ b/packages/cli/src/utils/handleAutoUpdate.test.ts @@ -12,7 +12,13 @@ import type { UpdateObject } from '../ui/utils/updateCheck.js'; import type { LoadedSettings } from '../config/settings.js'; import EventEmitter from 'node:events'; import type { ChildProcess } from 'node:child_process'; -import { handleAutoUpdate, setUpdateHandler } from './handleAutoUpdate.js'; +import { + handleAutoUpdate, + setUpdateHandler, + isUpdateInProgress, + waitForUpdateCompletion, + _setUpdateStateForTesting, +} from './handleAutoUpdate.js'; import { MessageType } from '../ui/types.js'; vi.mock('./installationInfo.js', async () => { @@ -79,6 +85,7 @@ describe('handleAutoUpdate', () => { afterEach(() => { vi.unstubAllEnvs(); vi.clearAllMocks(); + _setUpdateStateForTesting(false); }); it('should do nothing if update info is null', () => { @@ -88,6 +95,80 @@ describe('handleAutoUpdate', () => { expect(mockSpawn).not.toHaveBeenCalled(); }); + it('should track update progress state', async () => { + mockGetInstallationInfo.mockReturnValue({ + updateCommand: 'npm i -g @google/gemini-cli@latest', + updateMessage: 'This is an additional message.', + isGlobal: false, + packageManager: PackageManager.NPM, + }); + + expect(isUpdateInProgress()).toBe(false); + + handleAutoUpdate(mockUpdateInfo, mockSettings, '/root', mockSpawn); + + expect(isUpdateInProgress()).toBe(true); + + mockChildProcess.emit('close', 0); + + expect(isUpdateInProgress()).toBe(false); + }); + + it('should track update progress state on error', async () => { + mockGetInstallationInfo.mockReturnValue({ + updateCommand: 'npm i -g @google/gemini-cli@latest', + updateMessage: 'This is an additional message.', + isGlobal: false, + packageManager: PackageManager.NPM, + }); + + handleAutoUpdate(mockUpdateInfo, mockSettings, '/root', mockSpawn); + + expect(isUpdateInProgress()).toBe(true); + + mockChildProcess.emit('error', new Error('fail')); + + expect(isUpdateInProgress()).toBe(false); + }); + + it('should resolve waitForUpdateCompletion when update succeeds', async () => { + _setUpdateStateForTesting(true); + + const waitPromise = waitForUpdateCompletion(); + updateEventEmitter.emit('update-success', {}); + + await expect(waitPromise).resolves.toBeUndefined(); + }); + + it('should resolve waitForUpdateCompletion when update fails', async () => { + _setUpdateStateForTesting(true); + + const waitPromise = waitForUpdateCompletion(); + updateEventEmitter.emit('update-failed', {}); + + await expect(waitPromise).resolves.toBeUndefined(); + }); + + it('should resolve waitForUpdateCompletion immediately if not in progress', async () => { + _setUpdateStateForTesting(false); + + const waitPromise = waitForUpdateCompletion(); + + await expect(waitPromise).resolves.toBeUndefined(); + }); + + it('should timeout waitForUpdateCompletion', async () => { + vi.useFakeTimers(); + _setUpdateStateForTesting(true); + + const waitPromise = waitForUpdateCompletion(1000); + + vi.advanceTimersByTime(1001); + + await expect(waitPromise).resolves.toBeUndefined(); + vi.useRealTimers(); + }); + it('should do nothing if update prompts are disabled', () => { mockSettings.merged.general.enableAutoUpdateNotification = false; handleAutoUpdate(mockUpdateInfo, mockSettings, '/root', mockSpawn); diff --git a/packages/cli/src/utils/handleAutoUpdate.ts b/packages/cli/src/utils/handleAutoUpdate.ts index a6d0cdc574..8a7b6f3925 100644 --- a/packages/cli/src/utils/handleAutoUpdate.ts +++ b/packages/cli/src/utils/handleAutoUpdate.ts @@ -12,6 +12,54 @@ import type { HistoryItem } from '../ui/types.js'; import { MessageType } from '../ui/types.js'; import { spawnWrapper } from './spawnWrapper.js'; import type { spawn } from 'node:child_process'; +import { debugLogger } from '@google/gemini-cli-core'; + +let _updateInProgress = false; + +/** @internal */ +export function _setUpdateStateForTesting(value: boolean) { + _updateInProgress = value; +} + +export function isUpdateInProgress() { + return _updateInProgress; +} + +/** + * Returns a promise that resolves when the update process completes or times out. + */ +export async function waitForUpdateCompletion( + timeoutMs = 30000, +): Promise { + if (!_updateInProgress) { + return; + } + + debugLogger.log( + '\nGemini CLI is waiting for a background update to complete before restarting...', + ); + + return new Promise((resolve) => { + // Re-check the condition inside the promise executor to avoid a race condition. + // If the update finished between the initial check and now, resolve immediately. + if (!_updateInProgress) { + resolve(); + return; + } + + const timer = setTimeout(cleanup, timeoutMs); + + function cleanup() { + clearTimeout(timer); + updateEventEmitter.off('update-success', cleanup); + updateEventEmitter.off('update-failed', cleanup); + resolve(); + } + + updateEventEmitter.once('update-success', cleanup); + updateEventEmitter.once('update-failed', cleanup); + }); +} export function handleAutoUpdate( info: UpdateObject | null, @@ -62,6 +110,11 @@ export function handleAutoUpdate( ) { return; } + + if (_updateInProgress) { + return; + } + const isNightly = info.update.latest.includes('nightly'); const updateCommand = installationInfo.updateCommand.replace( @@ -73,10 +126,14 @@ export function handleAutoUpdate( shell: true, detached: true, }); + + _updateInProgress = true; + // Un-reference the child process to allow the parent to exit independently. updateProcess.unref(); updateProcess.on('close', (code) => { + _updateInProgress = false; if (code === 0) { updateEventEmitter.emit('update-success', { message: @@ -90,6 +147,7 @@ export function handleAutoUpdate( }); updateProcess.on('error', (err) => { + _updateInProgress = false; updateEventEmitter.emit('update-failed', { message: `Automatic update failed. Please try updating manually. (error: ${err.message})`, }); diff --git a/packages/cli/src/utils/processUtils.test.ts b/packages/cli/src/utils/processUtils.test.ts index be85a4dbad..009c17a9d4 100644 --- a/packages/cli/src/utils/processUtils.test.ts +++ b/packages/cli/src/utils/processUtils.test.ts @@ -7,6 +7,11 @@ import { vi } from 'vitest'; import { RELAUNCH_EXIT_CODE, relaunchApp } from './processUtils.js'; import * as cleanup from './cleanup.js'; +import * as handleAutoUpdate from './handleAutoUpdate.js'; + +vi.mock('./handleAutoUpdate.js', () => ({ + waitForUpdateCompletion: vi.fn().mockResolvedValue(undefined), +})); describe('processUtils', () => { const processExit = vi @@ -14,8 +19,11 @@ describe('processUtils', () => { .mockReturnValue(undefined as never); const runExitCleanup = vi.spyOn(cleanup, 'runExitCleanup'); - it('should run cleanup and exit with the relaunch code', async () => { + afterEach(() => vi.clearAllMocks()); + + it('should wait for updates, run cleanup, and exit with the relaunch code', async () => { await relaunchApp(); + expect(handleAutoUpdate.waitForUpdateCompletion).toHaveBeenCalledTimes(1); expect(runExitCleanup).toHaveBeenCalledTimes(1); expect(processExit).toHaveBeenCalledWith(RELAUNCH_EXIT_CODE); }); diff --git a/packages/cli/src/utils/processUtils.ts b/packages/cli/src/utils/processUtils.ts index 1122a2b0dc..c55caf023b 100644 --- a/packages/cli/src/utils/processUtils.ts +++ b/packages/cli/src/utils/processUtils.ts @@ -5,6 +5,7 @@ */ import { runExitCleanup } from './cleanup.js'; +import { waitForUpdateCompletion } from './handleAutoUpdate.js'; /** * Exit code used to signal that the CLI should be relaunched. @@ -15,6 +16,7 @@ export const RELAUNCH_EXIT_CODE = 199; * Exits the process with a special code to signal that the parent process should relaunch it. */ export async function relaunchApp(): Promise { + await waitForUpdateCompletion(); await runExitCleanup(); process.exit(RELAUNCH_EXIT_CODE); }