Cleanup extension update logic (#10514)

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
2026-04-09 21:00:56 -07:00 · 2025-10-03 21:06:26 -07:00
parent 1a06282061
commit 7f8537a130
10 changed files with 310 additions and 292 deletions
--- a/packages/cli/src/ui/commands/extensionsCommand.test.ts
+++ b/packages/cli/src/ui/commands/extensionsCommand.test.ts
@@ -5,43 +5,23 @@
 */

 import type { GeminiCLIExtension } from '@google/gemini-cli-core';
-import {
-  updateAllUpdatableExtensions,
-  updateExtension,
-} from '../../config/extensions/update.js';
 import { createMockCommandContext } from '../../test-utils/mockCommandContext.js';
 import { MessageType } from '../types.js';
 import { extensionsCommand } from './extensionsCommand.js';
 import { type CommandContext } from './types.js';
-import {
-  describe,
-  it,
-  expect,
-  vi,
-  beforeEach,
-  type MockedFunction,
-} from 'vitest';
-import { ExtensionUpdateState } from '../state/extensions.js';
+import { describe, it, expect, vi, beforeEach } from 'vitest';
+import { type ExtensionUpdateAction } from '../state/extensions.js';

 vi.mock('../../config/extensions/update.js', () => ({
  updateExtension: vi.fn(),
-  updateAllUpdatableExtensions: vi.fn(),
  checkForAllExtensionUpdates: vi.fn(),
 }));

-const mockUpdateExtension = updateExtension as MockedFunction<
-  typeof updateExtension
->;
-
-const mockUpdateAllUpdatableExtensions =
-  updateAllUpdatableExtensions as MockedFunction<
-    typeof updateAllUpdatableExtensions
-  >;
-
 const mockGetExtensions = vi.fn();

 describe('extensionsCommand', () => {
  let mockContext: CommandContext;
+  const mockDispatchExtensionState = vi.fn();

  beforeEach(() => {
    vi.resetAllMocks();
@@ -53,7 +33,7 @@ describe('extensionsCommand', () => {
        },
      },
      ui: {
-        dispatchExtensionStateUpdate: vi.fn(),
+        dispatchExtensionStateUpdate: mockDispatchExtensionState,
      },
    });
  });
@@ -93,7 +73,14 @@ describe('extensionsCommand', () => {
    });

    it('should inform user if there are no extensions to update with --all', async () => {
-      mockUpdateAllUpdatableExtensions.mockResolvedValue([]);
+      mockDispatchExtensionState.mockImplementationOnce(
+        (action: ExtensionUpdateAction) => {
+          if (action.type === 'SCHEDULE_UPDATE') {
+            action.payload.onComplete([]);
+          }
+        },
+      );
+
      await updateAction(mockContext, '--all');
      expect(mockContext.ui.addItem).toHaveBeenCalledWith(
        {
@@ -105,18 +92,24 @@ describe('extensionsCommand', () => {
    });

    it('should call setPendingItem and addItem in a finally block on success', async () => {
-      mockUpdateAllUpdatableExtensions.mockResolvedValue([
-        {
-          name: 'ext-one',
-          originalVersion: '1.0.0',
-          updatedVersion: '1.0.1',
+      mockDispatchExtensionState.mockImplementationOnce(
+        (action: ExtensionUpdateAction) => {
+          if (action.type === 'SCHEDULE_UPDATE') {
+            action.payload.onComplete([
+              {
+                name: 'ext-one',
+                originalVersion: '1.0.0',
+                updatedVersion: '1.0.1',
+              },
+              {
+                name: 'ext-two',
+                originalVersion: '2.0.0',
+                updatedVersion: '2.0.1',
+              },
+            ]);
+          }
        },
-        {
-          name: 'ext-two',
-          originalVersion: '2.0.0',
-          updatedVersion: '2.0.1',
-        },
-      ]);
+      );
      await updateAction(mockContext, '--all');
      expect(mockContext.ui.setPendingItem).toHaveBeenCalledWith({
        type: MessageType.EXTENSIONS_LIST,
@@ -131,9 +124,9 @@ describe('extensionsCommand', () => {
    });

    it('should call setPendingItem and addItem in a finally block on failure', async () => {
-      mockUpdateAllUpdatableExtensions.mockRejectedValue(
-        new Error('Something went wrong'),
-      );
+      mockDispatchExtensionState.mockImplementationOnce((_) => {
+        throw new Error('Something went wrong');
+      });
      await updateAction(mockContext, '--all');
      expect(mockContext.ui.setPendingItem).toHaveBeenCalledWith({
        type: MessageType.EXTENSIONS_LIST,
@@ -155,95 +148,58 @@ describe('extensionsCommand', () => {
    });

    it('should update a single extension by name', async () => {
-      const extension: GeminiCLIExtension = {
-        name: 'ext-one',
-        version: '1.0.0',
-        isActive: true,
-        path: '/test/dir/ext-one',
-        installMetadata: {
-          type: 'git',
-          autoUpdate: false,
-          source: 'https://github.com/some/extension.git',
+      mockDispatchExtensionState.mockImplementationOnce(
+        (action: ExtensionUpdateAction) => {
+          if (action.type === 'SCHEDULE_UPDATE') {
+            action.payload.onComplete([
+              {
+                name: 'ext-one',
+                originalVersion: '1.0.0',
+                updatedVersion: '1.0.1',
+              },
+            ]);
+          }
        },
-      };
-      mockUpdateExtension.mockResolvedValue({
-        name: extension.name,
-        originalVersion: extension.version,
-        updatedVersion: '1.0.1',
-      });
-      mockGetExtensions.mockReturnValue([extension]);
-      mockContext.ui.extensionsUpdateState.set(extension.name, {
-        status: ExtensionUpdateState.UPDATE_AVAILABLE,
-        processed: false,
-      });
-      await updateAction(mockContext, 'ext-one');
-      expect(mockUpdateExtension).toHaveBeenCalledWith(
-        extension,
-        '/test/dir',
-        expect.any(Function),
-        ExtensionUpdateState.UPDATE_AVAILABLE,
-        expect.any(Function),
      );
-    });
-
-    it('should handle errors when updating a single extension', async () => {
-      mockUpdateExtension.mockRejectedValue(new Error('Extension not found'));
-      mockGetExtensions.mockReturnValue([]);
      await updateAction(mockContext, 'ext-one');
-      expect(mockContext.ui.addItem).toHaveBeenCalledWith(
-        {
-          type: MessageType.ERROR,
-          text: 'Extension ext-one not found.',
+      expect(mockDispatchExtensionState).toHaveBeenCalledWith({
+        type: 'SCHEDULE_UPDATE',
+        payload: {
+          all: false,
+          names: ['ext-one'],
+          onComplete: expect.any(Function),
        },
-        expect.any(Number),
-      );
+      });
    });

    it('should update multiple extensions by name', async () => {
-      const extensionOne: GeminiCLIExtension = {
-        name: 'ext-one',
-        version: '1.0.0',
-        isActive: true,
-        path: '/test/dir/ext-one',
-        installMetadata: {
-          type: 'git',
-          autoUpdate: false,
-          source: 'https://github.com/some/extension.git',
+      mockDispatchExtensionState.mockImplementationOnce(
+        (action: ExtensionUpdateAction) => {
+          if (action.type === 'SCHEDULE_UPDATE') {
+            action.payload.onComplete([
+              {
+                name: 'ext-one',
+                originalVersion: '1.0.0',
+                updatedVersion: '1.0.1',
+              },
+              {
+                name: 'ext-two',
+                originalVersion: '1.0.0',
+                updatedVersion: '1.0.1',
+              },
+            ]);
+          }
        },
-      };
-      const extensionTwo: GeminiCLIExtension = {
-        name: 'ext-two',
-        version: '1.0.0',
-        isActive: true,
-        path: '/test/dir/ext-two',
-        installMetadata: {
-          type: 'git',
-          autoUpdate: false,
-          source: 'https://github.com/some/extension.git',
-        },
-      };
-      mockGetExtensions.mockReturnValue([extensionOne, extensionTwo]);
-      mockContext.ui.extensionsUpdateState.set(
-        extensionOne.name,
-        ExtensionUpdateState.UPDATE_AVAILABLE,
      );
-      mockContext.ui.extensionsUpdateState.set(
-        extensionTwo.name,
-        ExtensionUpdateState.UPDATE_AVAILABLE,
-      );
-      mockUpdateExtension
-        .mockResolvedValueOnce({
-          name: 'ext-one',
-          originalVersion: '1.0.0',
-          updatedVersion: '1.0.1',
-        })
-        .mockResolvedValueOnce({
-          name: 'ext-two',
-          originalVersion: '2.0.0',
-          updatedVersion: '2.0.1',
-        });
      await updateAction(mockContext, 'ext-one ext-two');
-      expect(mockUpdateExtension).toHaveBeenCalledTimes(2);
+      expect(mockDispatchExtensionState).toHaveBeenCalledWith({
+        type: 'SCHEDULE_UPDATE',
+        payload: {
+          all: false,
+          names: ['ext-one', 'ext-two'],
+          onComplete: expect.any(Function),
+        },
+      });
      expect(mockContext.ui.setPendingItem).toHaveBeenCalledWith({
        type: MessageType.EXTENSIONS_LIST,
      });
--- a/packages/cli/src/ui/commands/extensionsCommand.ts
+++ b/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 = {