Cleanup extension update logic (#10514)

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

packages/cli/src/config/extensions/github.test.ts

@@ -132,11 +132,7 @@ describe('git extension helpers', () => {
           source: '',
         },
       };
-      let result: ExtensionUpdateState | undefined = undefined;
-      await checkForExtensionUpdate(
-        extension,
-        (newState) => (result = newState),
-      );
+      const result = await checkForExtensionUpdate(extension);
       expect(result).toBe(ExtensionUpdateState.NOT_UPDATABLE);
     });
 
@@ -152,11 +148,7 @@ describe('git extension helpers', () => {
         },
       };
       mockGit.getRemotes.mockResolvedValue([]);
-      let result: ExtensionUpdateState | undefined = undefined;
-      await checkForExtensionUpdate(
-        extension,
-        (newState) => (result = newState),
-      );
+      const result = await checkForExtensionUpdate(extension);
       expect(result).toBe(ExtensionUpdateState.ERROR);
     });
 
@@ -177,11 +169,7 @@ describe('git extension helpers', () => {
       mockGit.listRemote.mockResolvedValue('remote-hash\tHEAD');
       mockGit.revparse.mockResolvedValue('local-hash');
 
-      let result: ExtensionUpdateState | undefined = undefined;
-      await checkForExtensionUpdate(
-        extension,
-        (newState) => (result = newState),
-      );
+      const result = await checkForExtensionUpdate(extension);
       expect(result).toBe(ExtensionUpdateState.UPDATE_AVAILABLE);
     });
 
@@ -202,11 +190,7 @@ describe('git extension helpers', () => {
       mockGit.listRemote.mockResolvedValue('same-hash\tHEAD');
       mockGit.revparse.mockResolvedValue('same-hash');
 
-      let result: ExtensionUpdateState | undefined = undefined;
-      await checkForExtensionUpdate(
-        extension,
-        (newState) => (result = newState),
-      );
+      const result = await checkForExtensionUpdate(extension);
       expect(result).toBe(ExtensionUpdateState.UP_TO_DATE);
     });
 
@@ -223,11 +207,7 @@ describe('git extension helpers', () => {
       };
       mockGit.getRemotes.mockRejectedValue(new Error('git error'));
 
-      let result: ExtensionUpdateState | undefined = undefined;
-      await checkForExtensionUpdate(
-        extension,
-        (newState) => (result = newState),
-      );
+      const result = await checkForExtensionUpdate(extension);
       expect(result).toBe(ExtensionUpdateState.ERROR);
     });
   });

packages/cli/src/config/extensions/github.ts

@@ -119,10 +119,8 @@ async function fetchReleaseFromGithub(
 
 export async function checkForExtensionUpdate(
   extension: GeminiCLIExtension,
-  setExtensionUpdateState: (updateState: ExtensionUpdateState) => void,
   cwd: string = process.cwd(),
-): Promise<void> {
-  setExtensionUpdateState(ExtensionUpdateState.CHECKING_FOR_UPDATES);
+): Promise<ExtensionUpdateState> {
   const installMetadata = extension.installMetadata;
   if (installMetadata?.type === 'local') {
     const newExtension = loadExtension({
@@ -133,23 +131,19 @@ export async function checkForExtensionUpdate(
       console.error(
         `Failed to check for update for local extension "${extension.name}". Could not load extension from source path: ${installMetadata.source}`,
       );
-      setExtensionUpdateState(ExtensionUpdateState.ERROR);
-      return;
+      return ExtensionUpdateState.ERROR;
     }
     if (newExtension.config.version !== extension.version) {
-      setExtensionUpdateState(ExtensionUpdateState.UPDATE_AVAILABLE);
-      return;
+      return ExtensionUpdateState.UPDATE_AVAILABLE;
     }
-    setExtensionUpdateState(ExtensionUpdateState.UP_TO_DATE);
-    return;
+    return ExtensionUpdateState.UP_TO_DATE;
   }
   if (
     !installMetadata ||
     (installMetadata.type !== 'git' &&
       installMetadata.type !== 'github-release')
   ) {
-    setExtensionUpdateState(ExtensionUpdateState.NOT_UPDATABLE);
-    return;
+    return ExtensionUpdateState.NOT_UPDATABLE;
   }
   try {
     if (installMetadata.type === 'git') {
@@ -157,14 +151,12 @@ export async function checkForExtensionUpdate(
       const remotes = await git.getRemotes(true);
       if (remotes.length === 0) {
         console.error('No git remotes found.');
-        setExtensionUpdateState(ExtensionUpdateState.ERROR);
-        return;
+        return ExtensionUpdateState.ERROR;
       }
       const remoteUrl = remotes[0].refs.fetch;
       if (!remoteUrl) {
         console.error(`No fetch URL found for git remote ${remotes[0].name}.`);
-        setExtensionUpdateState(ExtensionUpdateState.ERROR);
-        return;
+        return ExtensionUpdateState.ERROR;
       }
 
       // Determine the ref to check on the remote.
@@ -174,8 +166,7 @@ export async function checkForExtensionUpdate(
 
       if (typeof lsRemoteOutput !== 'string' || lsRemoteOutput.trim() === '') {
         console.error(`Git ref ${refToCheck} not found.`);
-        setExtensionUpdateState(ExtensionUpdateState.ERROR);
-        return;
+        return ExtensionUpdateState.ERROR;
       }
 
       const remoteHash = lsRemoteOutput.split('\t')[0];
@@ -185,21 +176,17 @@ export async function checkForExtensionUpdate(
         console.error(
           `Unable to parse hash from git ls-remote output "${lsRemoteOutput}"`,
         );
-        setExtensionUpdateState(ExtensionUpdateState.ERROR);
-        return;
+        return ExtensionUpdateState.ERROR;
       }
       if (remoteHash === localHash) {
-        setExtensionUpdateState(ExtensionUpdateState.UP_TO_DATE);
-        return;
+        return ExtensionUpdateState.UP_TO_DATE;
       }
-      setExtensionUpdateState(ExtensionUpdateState.UPDATE_AVAILABLE);
-      return;
+      return ExtensionUpdateState.UPDATE_AVAILABLE;
     } else {
       const { source, releaseTag } = installMetadata;
       if (!source) {
         console.error(`No "source" provided for extension.`);
-        setExtensionUpdateState(ExtensionUpdateState.ERROR);
-        return;
+        return ExtensionUpdateState.ERROR;
       }
       const { owner, repo } = parseGitHubRepoForReleases(source);
 
@@ -209,18 +196,15 @@ export async function checkForExtensionUpdate(
         installMetadata.ref,
       );
       if (releaseData.tag_name !== releaseTag) {
-        setExtensionUpdateState(ExtensionUpdateState.UPDATE_AVAILABLE);
-        return;
+        return ExtensionUpdateState.UPDATE_AVAILABLE;
       }
-      setExtensionUpdateState(ExtensionUpdateState.UP_TO_DATE);
-      return;
+      return ExtensionUpdateState.UP_TO_DATE;
     }
   } catch (error) {
     console.error(
       `Failed to check for updates for extension "${installMetadata.source}": ${getErrorMessage(error)}`,
     );
-    setExtensionUpdateState(ExtensionUpdateState.ERROR);
-    return;
+    return ExtensionUpdateState.ERROR;
   }
 }
 export interface GitHubDownloadResult {

packages/cli/src/config/extensions/update.test.ts

@@ -302,7 +302,11 @@ describe('update tests', () => {
       mockGit.revparse.mockResolvedValue('localHash');
 
       const dispatch = vi.fn();
-      await checkForAllExtensionUpdates([extension], dispatch);
+      await checkForAllExtensionUpdates(
+        [extension],
+        dispatch,
+        tempWorkspaceDir,
+      );
       expect(dispatch).toHaveBeenCalledWith({
         type: 'SET_STATE',
         payload: {
@@ -340,7 +344,11 @@ describe('update tests', () => {
       mockGit.revparse.mockResolvedValue('sameHash');
 
       const dispatch = vi.fn();
-      await checkForAllExtensionUpdates([extension], dispatch);
+      await checkForAllExtensionUpdates(
+        [extension],
+        dispatch,
+        tempWorkspaceDir,
+      );
       expect(dispatch).toHaveBeenCalledWith({
         type: 'SET_STATE',
         payload: {
@@ -375,7 +383,11 @@ describe('update tests', () => {
         new ExtensionEnablementManager(ExtensionStorage.getUserExtensionsDir()),
       )[0];
       const dispatch = vi.fn();
-      await checkForAllExtensionUpdates([extension], dispatch);
+      await checkForAllExtensionUpdates(
+        [extension],
+        dispatch,
+        tempWorkspaceDir,
+      );
       expect(dispatch).toHaveBeenCalledWith({
         type: 'SET_STATE',
         payload: {
@@ -410,7 +422,11 @@ describe('update tests', () => {
         new ExtensionEnablementManager(ExtensionStorage.getUserExtensionsDir()),
       )[0];
       const dispatch = vi.fn();
-      await checkForAllExtensionUpdates([extension], dispatch);
+      await checkForAllExtensionUpdates(
+        [extension],
+        dispatch,
+        tempWorkspaceDir,
+      );
       expect(dispatch).toHaveBeenCalledWith({
         type: 'SET_STATE',
         payload: {
@@ -444,7 +460,11 @@ describe('update tests', () => {
       mockGit.getRemotes.mockRejectedValue(new Error('Git error'));
 
       const dispatch = vi.fn();
-      await checkForAllExtensionUpdates([extension], dispatch);
+      await checkForAllExtensionUpdates(
+        [extension],
+        dispatch,
+        tempWorkspaceDir,
+      );
       expect(dispatch).toHaveBeenCalledWith({
         type: 'SET_STATE',
         payload: {

packages/cli/src/config/extensions/update.ts

@@ -154,6 +154,7 @@ export interface ExtensionUpdateCheckResult {
 export async function checkForAllExtensionUpdates(
   extensions: GeminiCLIExtension[],
   dispatch: (action: ExtensionUpdateAction) => void,
+  cwd: string = process.cwd(),
 ): Promise<void> {
   dispatch({ type: 'BATCH_CHECK_START' });
   const promises: Array<Promise<void>> = [];
@@ -168,13 +169,20 @@ export async function checkForAllExtensionUpdates(
       });
       continue;
     }
+    dispatch({
+      type: 'SET_STATE',
+      payload: {
+        name: extension.name,
+        state: ExtensionUpdateState.CHECKING_FOR_UPDATES,
+      },
+    });
     promises.push(
-      checkForExtensionUpdate(extension, (updatedState) => {
+      checkForExtensionUpdate(extension, cwd).then((state) =>
         dispatch({
           type: 'SET_STATE',
-          payload: { name: extension.name, state: updatedState },
-        });
-      }),
+          payload: { name: extension.name, state },
+        }),
+      ),
     );
   }
   await Promise.all(promises);