feat(autoupdate): Improve update check and refactor for testability (#5389)

5 files changed, 258 insertions, 68 deletions

diff --git a/packages/cli/src/ui/utils/updateCheck.test.ts b/packages/cli/src/ui/utils/updateCheck.test.ts
index fa6f342e..c2b56a03 100644
--- a/packages/cli/src/ui/utils/updateCheck.test.ts
+++ b/packages/cli/src/ui/utils/updateCheck.test.ts
@@ -5,7 +5,7 @@
  */
 
 import { vi, describe, it, expect, beforeEach } from 'vitest';
-import { checkForUpdates, FETCH_TIMEOUT_MS } from './updateCheck.js';
+import { checkForUpdates } from './updateCheck.js';
 
 const getPackageJson = vi.hoisted(() => vi.fn());
 vi.mock('../../utils/package.js', () => ({
@@ -109,24 +109,16 @@ describe('checkForUpdates', () => {
     expect(result).toBeNull();
   });
 
-  it('should return null if fetchInfo times out', async () => {
+  it('should return null if fetchInfo rejects', async () => {
     getPackageJson.mockResolvedValue({
       name: 'test-package',
       version: '1.0.0',
     });
     updateNotifier.mockReturnValue({
-      fetchInfo: vi.fn(
-        async () =>
-          new Promise((resolve) => {
-            setTimeout(() => {
-              resolve({ current: '1.0.0', latest: '1.1.0' });
-            }, FETCH_TIMEOUT_MS + 1);
-          }),
-      ),
+      fetchInfo: vi.fn().mockRejectedValue(new Error('Timeout')),
     });
-    const promise = checkForUpdates();
-    await vi.advanceTimersByTimeAsync(FETCH_TIMEOUT_MS);
-    const result = await promise;
+
+    const result = await checkForUpdates();
     expect(result).toBeNull();
   });
 
@@ -135,4 +127,37 @@ describe('checkForUpdates', () => {
     const result = await checkForUpdates();
     expect(result).toBeNull();
   });
+
+  describe('nightly updates', () => {
+    it('should notify for a newer nightly version when current is nightly', async () => {
+      getPackageJson.mockResolvedValue({
+        name: 'test-package',
+        version: '1.2.3-nightly.1',
+      });
+
+      const fetchInfoMock = vi.fn().mockImplementation(({ distTag }) => {
+        if (distTag === 'nightly') {
+          return Promise.resolve({
+            latest: '1.2.3-nightly.2',
+            current: '1.2.3-nightly.1',
+          });
+        }
+        if (distTag === 'latest') {
+          return Promise.resolve({
+            latest: '1.2.3',
+            current: '1.2.3-nightly.1',
+          });
+        }
+        return Promise.resolve(null);
+      });
+
+      updateNotifier.mockImplementation(({ pkg, distTag }) => ({
+        fetchInfo: () => fetchInfoMock({ pkg, distTag }),
+      }));
+
+      const result = await checkForUpdates();
+      expect(result?.message).toContain('1.2.3-nightly.1 → 1.2.3-nightly.2');
+      expect(result?.update.latest).toBe('1.2.3-nightly.2');
+    });
+  });
 });
diff --git a/packages/cli/src/ui/utils/updateCheck.ts b/packages/cli/src/ui/utils/updateCheck.ts
index 2fe5df39..f4c18586 100644
--- a/packages/cli/src/ui/utils/updateCheck.ts
+++ b/packages/cli/src/ui/utils/updateCheck.ts
@@ -15,38 +15,81 @@ export interface UpdateObject {
   update: UpdateInfo;
 }
 
+/**
+ * From a nightly and stable update, determines which is the "best" one to offer.
+ * The rule is to always prefer nightly if the base versions are the same.
+ */
+function getBestAvailableUpdate(
+  nightly?: UpdateInfo,
+  stable?: UpdateInfo,
+): UpdateInfo | null {
+  if (!nightly) return stable || null;
+  if (!stable) return nightly || null;
+
+  const nightlyVer = nightly.latest;
+  const stableVer = stable.latest;
+
+  if (
+    semver.coerce(stableVer)?.version === semver.coerce(nightlyVer)?.version
+  ) {
+    return nightly;
+  }
+
+  return semver.gt(stableVer, nightlyVer) ? stable : nightly;
+}
+
 export async function checkForUpdates(): Promise<UpdateObject | null> {
   try {
     // Skip update check when running from source (development mode)
     if (process.env.DEV === 'true') {
       return null;
     }
-
     const packageJson = await getPackageJson();
     if (!packageJson || !packageJson.name || !packageJson.version) {
       return null;
     }
-    const notifier = updateNotifier({
-      pkg: {
-        name: packageJson.name,
-        version: packageJson.version,
-      },
-      // check every time
-      updateCheckInterval: 0,
-      // allow notifier to run in scripts
-      shouldNotifyInNpmScript: true,
-    });
-    // avoid blocking by waiting at most FETCH_TIMEOUT_MS for fetchInfo to resolve
-    const timeout = new Promise<null>((resolve) =>
-      setTimeout(resolve, FETCH_TIMEOUT_MS, null),
-    );
-    const updateInfo = await Promise.race([notifier.fetchInfo(), timeout]);
 
-    if (updateInfo && semver.gt(updateInfo.latest, updateInfo.current)) {
-      return {
-        message: `Gemini CLI update available! ${updateInfo.current} → ${updateInfo.latest}`,
-        update: updateInfo,
-      };
+    const { name, version: currentVersion } = packageJson;
+    const isNightly = currentVersion.includes('nightly');
+    const createNotifier = (distTag: 'latest' | 'nightly') =>
+      updateNotifier({
+        pkg: {
+          name,
+          version: currentVersion,
+        },
+        updateCheckInterval: 0,
+        shouldNotifyInNpmScript: true,
+        distTag,
+      });
+
+    if (isNightly) {
+      const [nightlyUpdateInfo, latestUpdateInfo] = await Promise.all([
+        createNotifier('nightly').fetchInfo(),
+        createNotifier('latest').fetchInfo(),
+      ]);
+
+      const bestUpdate = getBestAvailableUpdate(
+        nightlyUpdateInfo,
+        latestUpdateInfo,
+      );
+
+      if (bestUpdate && semver.gt(bestUpdate.latest, currentVersion)) {
+        const message = `A new version of Gemini CLI is available! ${currentVersion} → ${bestUpdate.latest}`;
+        return {
+          message,
+          update: { ...bestUpdate, current: currentVersion },
+        };
+      }
+    } else {
+      const updateInfo = await createNotifier('latest').fetchInfo();
+
+      if (updateInfo && semver.gt(updateInfo.latest, currentVersion)) {
+        const message = `Gemini CLI update available! ${currentVersion} → ${updateInfo.latest}`;
+        return {
+          message,
+          update: { ...updateInfo, current: currentVersion },
+        };
+      }
     }
 
     return null;
diff --git a/packages/cli/src/utils/handleAutoUpdate.test.ts b/packages/cli/src/utils/handleAutoUpdate.test.ts
index adaed932..f292d0c2 100644
--- a/packages/cli/src/utils/handleAutoUpdate.test.ts
+++ b/packages/cli/src/utils/handleAutoUpdate.test.ts
@@ -4,22 +4,13 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
-import { ChildProcess, spawn } from 'node:child_process';
-import { handleAutoUpdate } from './handleAutoUpdate.js';
+import { describe, it, expect, vi, beforeEach, afterEach, Mock } from 'vitest';
 import { getInstallationInfo, PackageManager } from './installationInfo.js';
 import { updateEventEmitter } from './updateEventEmitter.js';
 import { UpdateObject } from '../ui/utils/updateCheck.js';
 import { LoadedSettings } from '../config/settings.js';
-
-// Mock dependencies
-vi.mock('node:child_process', async () => {
-  const actual = await vi.importActual('node:child_process');
-  return {
-    ...actual,
-    spawn: vi.fn(),
-  };
-});
+import EventEmitter from 'node:events';
+import { handleAutoUpdate } from './handleAutoUpdate.js';
 
 vi.mock('./installationInfo.js', async () => {
   const actual = await vi.importActual('./installationInfo.js');
@@ -40,21 +31,26 @@ vi.mock('./updateEventEmitter.js', async () => {
   };
 });
 
-const mockSpawn = vi.mocked(spawn);
+interface MockChildProcess extends EventEmitter {
+  stdin: EventEmitter & {
+    write: Mock;
+    end: Mock;
+  };
+  stderr: EventEmitter;
+}
+
 const mockGetInstallationInfo = vi.mocked(getInstallationInfo);
 const mockUpdateEventEmitter = vi.mocked(updateEventEmitter);
 
 describe('handleAutoUpdate', () => {
+  let mockSpawn: Mock;
   let mockUpdateInfo: UpdateObject;
   let mockSettings: LoadedSettings;
-  let mockChildProcess: {
-    stderr: { on: ReturnType<typeof vi.fn> };
-    stdout: { on: ReturnType<typeof vi.fn> };
-    on: ReturnType<typeof vi.fn>;
-    unref: ReturnType<typeof vi.fn>;
-  };
+  let mockChildProcess: MockChildProcess;
 
   beforeEach(() => {
+    mockSpawn = vi.fn();
+    vi.clearAllMocks();
     mockUpdateInfo = {
       update: {
         latest: '2.0.0',
@@ -71,13 +67,17 @@ describe('handleAutoUpdate', () => {
       },
     } as LoadedSettings;
 
-    mockChildProcess = {
-      stdout: { on: vi.fn() },
-      stderr: { on: vi.fn() },
-      on: vi.fn(),
-      unref: vi.fn(),
-    };
-    mockSpawn.mockReturnValue(mockChildProcess as unknown as ChildProcess);
+    mockChildProcess = Object.assign(new EventEmitter(), {
+      stdin: Object.assign(new EventEmitter(), {
+        write: vi.fn(),
+        end: vi.fn(),
+      }),
+      stderr: new EventEmitter(),
+    }) as MockChildProcess;
+
+    mockSpawn.mockReturnValue(
+      mockChildProcess as unknown as ReturnType<typeof mockSpawn>,
+    );
   });
 
   afterEach(() => {
@@ -85,7 +85,7 @@ describe('handleAutoUpdate', () => {
   });
 
   it('should do nothing if update info is null', () => {
-    handleAutoUpdate(null, mockSettings, '/root');
+    handleAutoUpdate(null, mockSettings, '/root', mockSpawn);
     expect(mockGetInstallationInfo).not.toHaveBeenCalled();
     expect(mockUpdateEventEmitter.emit).not.toHaveBeenCalled();
     expect(mockSpawn).not.toHaveBeenCalled();
@@ -100,7 +100,7 @@ describe('handleAutoUpdate', () => {
       packageManager: PackageManager.NPM,
     });
 
-    handleAutoUpdate(mockUpdateInfo, mockSettings, '/root');
+    handleAutoUpdate(mockUpdateInfo, mockSettings, '/root', mockSpawn);
 
     expect(mockUpdateEventEmitter.emit).toHaveBeenCalledTimes(1);
     expect(mockUpdateEventEmitter.emit).toHaveBeenCalledWith(
@@ -120,7 +120,7 @@ describe('handleAutoUpdate', () => {
       packageManager: PackageManager.NPM,
     });
 
-    handleAutoUpdate(mockUpdateInfo, mockSettings, '/root');
+    handleAutoUpdate(mockUpdateInfo, mockSettings, '/root', mockSpawn);
 
     expect(mockUpdateEventEmitter.emit).toHaveBeenCalledTimes(1);
     expect(mockUpdateEventEmitter.emit).toHaveBeenCalledWith(
@@ -140,7 +140,7 @@ describe('handleAutoUpdate', () => {
       packageManager: PackageManager.NPM,
     });
 
-    handleAutoUpdate(mockUpdateInfo, mockSettings, '/root');
+    handleAutoUpdate(mockUpdateInfo, mockSettings, '/root', mockSpawn);
 
     expect(mockUpdateEventEmitter.emit).toHaveBeenCalledTimes(1);
     expect(mockUpdateEventEmitter.emit).toHaveBeenCalledWith(
@@ -150,4 +150,115 @@ describe('handleAutoUpdate', () => {
       },
     );
   });
+
+  it('should attempt to perform an update when conditions are met', async () => {
+    mockGetInstallationInfo.mockReturnValue({
+      updateCommand: 'npm i -g @google/gemini-cli@latest',
+      updateMessage: 'This is an additional message.',
+      isGlobal: false,
+      packageManager: PackageManager.NPM,
+    });
+
+    // Simulate successful execution
+    setTimeout(() => {
+      mockChildProcess.emit('close', 0);
+    }, 0);
+
+    handleAutoUpdate(mockUpdateInfo, mockSettings, '/root', mockSpawn);
+
+    expect(mockSpawn).toHaveBeenCalledOnce();
+  });
+
+  it('should emit "update-failed" when the update process fails', async () => {
+    await new Promise<void>((resolve) => {
+      mockGetInstallationInfo.mockReturnValue({
+        updateCommand: 'npm i -g @google/gemini-cli@latest',
+        updateMessage: 'This is an additional message.',
+        isGlobal: false,
+        packageManager: PackageManager.NPM,
+      });
+
+      // Simulate failed execution
+      setTimeout(() => {
+        mockChildProcess.stderr.emit('data', 'An error occurred');
+        mockChildProcess.emit('close', 1);
+        resolve();
+      }, 0);
+
+      handleAutoUpdate(mockUpdateInfo, mockSettings, '/root', mockSpawn);
+    });
+
+    expect(mockUpdateEventEmitter.emit).toHaveBeenCalledWith('update-failed', {
+      message:
+        'Automatic update failed. Please try updating manually. (command: npm i -g @google/[email protected], stderr: An error occurred)',
+    });
+  });
+
+  it('should emit "update-failed" when the spawn function throws an error', async () => {
+    await new Promise<void>((resolve) => {
+      mockGetInstallationInfo.mockReturnValue({
+        updateCommand: 'npm i -g @google/gemini-cli@latest',
+        updateMessage: 'This is an additional message.',
+        isGlobal: false,
+        packageManager: PackageManager.NPM,
+      });
+
+      // Simulate an error event
+      setTimeout(() => {
+        mockChildProcess.emit('error', new Error('Spawn error'));
+        resolve();
+      }, 0);
+
+      handleAutoUpdate(mockUpdateInfo, mockSettings, '/root', mockSpawn);
+    });
+
+    expect(mockUpdateEventEmitter.emit).toHaveBeenCalledWith('update-failed', {
+      message:
+        'Automatic update failed. Please try updating manually. (error: Spawn error)',
+    });
+  });
+
+  it('should use the "@nightly" tag for nightly updates', async () => {
+    mockUpdateInfo.update.latest = '2.0.0-nightly';
+    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(mockSpawn).toHaveBeenCalledWith(
+      'npm i -g @google/gemini-cli@nightly',
+      {
+        shell: true,
+        stdio: 'pipe',
+      },
+    );
+  });
+
+  it('should emit "update-success" when the update process succeeds', async () => {
+    await new Promise<void>((resolve) => {
+      mockGetInstallationInfo.mockReturnValue({
+        updateCommand: 'npm i -g @google/gemini-cli@latest',
+        updateMessage: 'This is an additional message.',
+        isGlobal: false,
+        packageManager: PackageManager.NPM,
+      });
+
+      // Simulate successful execution
+      setTimeout(() => {
+        mockChildProcess.emit('close', 0);
+        resolve();
+      }, 0);
+
+      handleAutoUpdate(mockUpdateInfo, mockSettings, '/root', mockSpawn);
+    });
+
+    expect(mockUpdateEventEmitter.emit).toHaveBeenCalledWith('update-success', {
+      message:
+        'Update successful! The new version will be used on your next run.',
+    });
+  });
 });
diff --git a/packages/cli/src/utils/handleAutoUpdate.ts b/packages/cli/src/utils/handleAutoUpdate.ts
index 1ef2d475..03a6a8d6 100644
--- a/packages/cli/src/utils/handleAutoUpdate.ts
+++ b/packages/cli/src/utils/handleAutoUpdate.ts
@@ -4,17 +4,19 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { spawn } from 'node:child_process';
 import { UpdateObject } from '../ui/utils/updateCheck.js';
 import { LoadedSettings } from '../config/settings.js';
 import { getInstallationInfo } from './installationInfo.js';
 import { updateEventEmitter } from './updateEventEmitter.js';
 import { HistoryItem, MessageType } from '../ui/types.js';
+import { spawnWrapper } from './spawnWrapper.js';
+import { spawn } from 'child_process';
 
 export function handleAutoUpdate(
   info: UpdateObject | null,
   settings: LoadedSettings,
   projectRoot: string,
+  spawnFn: typeof spawn = spawnWrapper,
 ) {
   if (!info) {
     return;
@@ -37,13 +39,13 @@ export function handleAutoUpdate(
   if (!installationInfo.updateCommand || settings.merged.disableAutoUpdate) {
     return;
   }
+  const isNightly = info.update.latest.includes('nightly');
 
   const updateCommand = installationInfo.updateCommand.replace(
     '@latest',
-    `@${info.update.latest}`,
+    isNightly ? '@nightly' : `@${info.update.latest}`,
   );
-
-  const updateProcess = spawn(updateCommand, { stdio: 'pipe', shell: true });
+  const updateProcess = spawnFn(updateCommand, { stdio: 'pipe', shell: true });
   let errorOutput = '';
   updateProcess.stderr.on('data', (data) => {
     errorOutput += data.toString();
diff --git a/packages/cli/src/utils/spawnWrapper.ts b/packages/cli/src/utils/spawnWrapper.ts
new file mode 100644
index 00000000..3f3cca94
--- /dev/null
+++ b/packages/cli/src/utils/spawnWrapper.ts
@@ -0,0 +1,9 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { spawn } from 'child_process';
+
+export const spawnWrapper = spawn;