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

1 files changed, 64 insertions, 21 deletions

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;