fix(cli): wait for background auto-update before relaunching (#20904)

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);

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<void> {
+  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})`,
     });

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);
   });

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<void> {
+  await waitForUpdateCompletion();
   await runExitCleanup();
   process.exit(RELAUNCH_EXIT_CODE);
 }