2 files changed, 102 insertions, 34 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;