Switch to a reducer for tracking update state fixing flicker issues due to continuous renders (#10280)

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

@@ -26,6 +26,7 @@ import { ExtensionUpdateState } from '../state/extensions.js';
 vi.mock('../../config/extensions/update.js', () => ({
   updateExtension: vi.fn(),
   updateAllUpdatableExtensions: vi.fn(),
+  checkForAllExtensionUpdates: vi.fn(),
 }));
 
 const mockUpdateExtension = updateExtension as MockedFunction<
@@ -51,6 +52,9 @@ describe('extensionsCommand', () => {
           getWorkingDir: () => '/test/dir',
         },
       },
+      ui: {
+        dispatchExtensionStateUpdate: vi.fn(),
+      },
     });
   });
 
@@ -168,10 +172,10 @@ describe('extensionsCommand', () => {
         updatedVersion: '1.0.1',
       });
       mockGetExtensions.mockReturnValue([extension]);
-      mockContext.ui.extensionsUpdateState.set(
-        extension.name,
-        ExtensionUpdateState.UPDATE_AVAILABLE,
-      );
+      mockContext.ui.extensionsUpdateState.set(extension.name, {
+        status: ExtensionUpdateState.UPDATE_AVAILABLE,
+        processed: false,
+      });
       await updateAction(mockContext, 'ext-one');
       expect(mockUpdateExtension).toHaveBeenCalledWith(
         extension,

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

@@ -9,6 +9,7 @@ import {
   updateAllUpdatableExtensions,
   type ExtensionUpdateInfo,
   updateExtension,
+  checkForAllExtensionUpdates,
 } from '../../config/extensions/update.js';
 import { getErrorMessage } from '../../utils/errors.js';
 import { ExtensionUpdateState } from '../state/extensions.js';
@@ -46,6 +47,10 @@ async function updateAction(context: CommandContext, args: string) {
   }
 
   try {
+    await checkForAllExtensionUpdates(
+      context.services.config!.getExtensions(),
+      context.ui.dispatchExtensionStateUpdate,
+    );
     context.ui.setPendingItem({
       type: MessageType.EXTENSIONS_LIST,
     });
@@ -60,7 +65,7 @@ async function updateAction(context: CommandContext, args: string) {
           ),
         context.services.config!.getExtensions(),
         context.ui.extensionsUpdateState,
-        context.ui.setExtensionsUpdateState,
+        context.ui.dispatchExtensionStateUpdate,
       );
     } else if (names?.length) {
       const workingDir = context.services.config!.getWorkingDir();
@@ -87,15 +92,9 @@ async function updateAction(context: CommandContext, args: string) {
               description,
               context.ui.addConfirmUpdateExtensionRequest,
             ),
-          context.ui.extensionsUpdateState.get(extension.name) ??
+          context.ui.extensionsUpdateState.get(extension.name)?.status ??
             ExtensionUpdateState.UNKNOWN,
-          (updateState) => {
-            context.ui.setExtensionsUpdateState((prev) => {
-              const newState = new Map(prev);
-              newState.set(name, updateState);
-              return newState;
-            });
-          },
+          context.ui.dispatchExtensionStateUpdate,
         );
         if (updateInfo) updateInfos.push(updateInfo);
       }

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

@@ -4,7 +4,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import type { Dispatch, ReactNode, SetStateAction } from 'react';
+import type { ReactNode } from 'react';
 import type { Content, PartListUnion } from '@google/genai';
 import type {
   HistoryItemWithoutId,
@@ -15,7 +15,10 @@ import type { Config, GitService, Logger } from '@google/gemini-cli-core';
 import type { LoadedSettings } from '../../config/settings.js';
 import type { UseHistoryManagerReturn } from '../hooks/useHistoryManager.js';
 import type { SessionStatsState } from '../contexts/SessionContext.js';
-import type { ExtensionUpdateState } from '../state/extensions.js';
+import type {
+  ExtensionUpdateAction,
+  ExtensionUpdateStatus,
+} from '../state/extensions.js';
 
 // Grouped dependencies for clarity and easier mocking
 export interface CommandContext {
@@ -66,10 +69,8 @@ export interface CommandContext {
     toggleVimEnabled: () => Promise<boolean>;
     setGeminiMdFileCount: (count: number) => void;
     reloadCommands: () => void;
-    extensionsUpdateState: Map<string, ExtensionUpdateState>;
-    setExtensionsUpdateState: Dispatch<
-      SetStateAction<Map<string, ExtensionUpdateState>>
-    >;
+    extensionsUpdateState: Map<string, ExtensionUpdateStatus>;
+    dispatchExtensionStateUpdate: (action: ExtensionUpdateAction) => void;
     addConfirmUpdateExtensionRequest: (value: ConfirmationRequest) => void;
   };
   // Session-specific data