Cleanup extension update logic (#10514)

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

packages/cli/src/ui/commands/extensionsCommand.ts

@@ -4,15 +4,8 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { requestConsentInteractive } from '../../config/extension.js';
-import {
-  updateAllUpdatableExtensions,
-  type ExtensionUpdateInfo,
-  updateExtension,
-  checkForAllExtensionUpdates,
-} from '../../config/extensions/update.js';
+import type { ExtensionUpdateInfo } from '../../config/extension.js';
 import { getErrorMessage } from '../../utils/errors.js';
-import { ExtensionUpdateState } from '../state/extensions.js';
 import { MessageType } from '../types.js';
 import {
   type CommandContext,
@@ -29,11 +22,10 @@ async function listAction(context: CommandContext) {
   );
 }
 
-async function updateAction(context: CommandContext, args: string) {
+function updateAction(context: CommandContext, args: string): Promise<void> {
   const updateArgs = args.split(' ').filter((value) => value.length > 0);
   const all = updateArgs.length === 1 && updateArgs[0] === '--all';
-  const names = all ? undefined : updateArgs;
-  let updateInfos: ExtensionUpdateInfo[] = [];
+  const names = all ? null : updateArgs;
 
   if (!all && names?.length === 0) {
     context.ui.addItem(
@@ -43,32 +35,48 @@ async function updateAction(context: CommandContext, args: string) {
       },
       Date.now(),
     );
-    return;
+    return Promise.resolve();
   }
 
-  try {
-    await checkForAllExtensionUpdates(
-      context.services.config!.getExtensions(),
-      context.ui.dispatchExtensionStateUpdate,
+  let resolveUpdateComplete: (updateInfo: ExtensionUpdateInfo[]) => void;
+  const updateComplete = new Promise<ExtensionUpdateInfo[]>(
+    (resolve) => (resolveUpdateComplete = resolve),
+  );
+  updateComplete.then((updateInfos) => {
+    if (updateInfos.length === 0) {
+      context.ui.addItem(
+        {
+          type: MessageType.INFO,
+          text: 'No extensions to update.',
+        },
+        Date.now(),
+      );
+    }
+    context.ui.addItem(
+      {
+        type: MessageType.EXTENSIONS_LIST,
+      },
+      Date.now(),
     );
+    context.ui.setPendingItem(null);
+  });
+
+  try {
     context.ui.setPendingItem({
       type: MessageType.EXTENSIONS_LIST,
     });
-    if (all) {
-      updateInfos = await updateAllUpdatableExtensions(
-        context.services.config!.getWorkingDir(),
-        // We don't have the ability to prompt for consent yet in this flow.
-        (description) =>
-          requestConsentInteractive(
-            description,
-            context.ui.addConfirmUpdateExtensionRequest,
-          ),
-        context.services.config!.getExtensions(),
-        context.ui.extensionsUpdateState,
-        context.ui.dispatchExtensionStateUpdate,
-      );
-    } else if (names?.length) {
-      const workingDir = context.services.config!.getWorkingDir();
+
+    context.ui.dispatchExtensionStateUpdate({
+      type: 'SCHEDULE_UPDATE',
+      payload: {
+        all,
+        names,
+        onComplete: (updateInfos) => {
+          resolveUpdateComplete(updateInfos);
+        },
+      },
+    });
+    if (names?.length) {
       const extensions = context.services.config!.getExtensions();
       for (const name of names) {
         const extension = extensions.find(
@@ -84,33 +92,10 @@ async function updateAction(context: CommandContext, args: string) {
           );
           continue;
         }
-        const updateInfo = await updateExtension(
-          extension,
-          workingDir,
-          (description) =>
-            requestConsentInteractive(
-              description,
-              context.ui.addConfirmUpdateExtensionRequest,
-            ),
-          context.ui.extensionsUpdateState.get(extension.name)?.status ??
-            ExtensionUpdateState.UNKNOWN,
-          context.ui.dispatchExtensionStateUpdate,
-        );
-        if (updateInfo) updateInfos.push(updateInfo);
       }
     }
-
-    if (updateInfos.length === 0) {
-      context.ui.addItem(
-        {
-          type: MessageType.INFO,
-          text: 'No extensions to update.',
-        },
-        Date.now(),
-      );
-      return;
-    }
   } catch (error) {
+    resolveUpdateComplete!([]);
     context.ui.addItem(
       {
         type: MessageType.ERROR,
@@ -118,15 +103,8 @@ async function updateAction(context: CommandContext, args: string) {
       },
       Date.now(),
     );
-  } finally {
-    context.ui.addItem(
-      {
-        type: MessageType.EXTENSIONS_LIST,
-      },
-      Date.now(),
-    );
-    context.ui.setPendingItem(null);
   }
+  return updateComplete.then((_) => {});
 }
 
 const listExtensionsCommand: SlashCommand = {