Pass whole extensions rather than just context files (#10910)

Co-authored-by: Jake Macdonald <jakemac@google.com>

packages/cli/src/ui/AppContainer.tsx

@@ -90,6 +90,7 @@ import { useSessionStats } from './contexts/SessionContext.js';
 import { useGitBranchName } from './hooks/useGitBranchName.js';
 import { useExtensionUpdates } from './hooks/useExtensionUpdates.js';
 import { ShellFocusContext } from './contexts/ShellFocusContext.js';
+import { ExtensionEnablementManager } from '../config/extensions/extensionEnablement.js';
 
 const CTRL_EXIT_PROMPT_DURATION_MS = 1000;
 const QUEUE_ERROR_DISPLAY_DURATION_MS = 3000;
@@ -159,6 +160,9 @@ export const AppContainer = (props: AppContainerProps) => {
   );
 
   const extensions = config.getExtensions();
+  const [extensionEnablementManager] = useState<ExtensionEnablementManager>(
+    new ExtensionEnablementManager(config.getEnabledExtensions()),
+  );
   const {
     extensionsUpdateState,
     extensionsUpdateStateInternal,
@@ -167,6 +171,7 @@ export const AppContainer = (props: AppContainerProps) => {
     addConfirmUpdateExtensionRequest,
   } = useExtensionUpdates(
     extensions,
+    extensionEnablementManager,
     historyManager.addItem,
     config.getWorkingDir(),
   );
@@ -529,7 +534,7 @@ Logging in with Google... Please restart Gemini CLI to continue.
           config.getDebugMode(),
           config.getFileService(),
           settings.merged,
-          config.getExtensionContextFilePaths(),
+          config.getExtensions(),
           config.isTrustedFolder(),
           settings.merged.context?.importFormat || 'tree', // Use setting or default to 'tree'
           config.getFileFilteringOptions(),

packages/cli/src/ui/commands/directoryCommand.test.tsx

@@ -44,7 +44,6 @@ describe('directoryCommand', () => {
       shouldLoadMemoryFromIncludeDirectories: () => false,
       getDebugMode: () => false,
       getFileService: () => ({}),
-      getExtensionContextFilePaths: () => [],
       getFileFilteringOptions: () => ({ ignore: [], include: [] }),
       setUserMemory: vi.fn(),
       setGeminiMdFileCount: vi.fn(),

packages/cli/src/ui/commands/directoryCommand.tsx

@@ -103,7 +103,7 @@ export const directoryCommand: SlashCommand = {
                 ],
                 config.getDebugMode(),
                 config.getFileService(),
-                config.getExtensionContextFilePaths(),
+                config.getExtensions(),
                 config.getFolderTrust(),
                 context.services.settings.merged.context?.importFormat ||
                   'tree', // Use setting or default to 'tree'

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

@@ -176,7 +176,7 @@ describe('memoryCommand', () => {
         getWorkingDir: () => '/test/dir',
         getDebugMode: () => false,
         getFileService: () => ({}) as FileDiscoveryService,
-        getExtensionContextFilePaths: () => [],
+        getExtensions: () => [],
         shouldLoadMemoryFromIncludeDirectories: () => false,
         getWorkspaceContext: () => ({
           getDirectories: () => [],

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

@@ -91,7 +91,7 @@ export const memoryCommand: SlashCommand = {
                 config.getDebugMode(),
                 config.getFileService(),
                 settings.merged,
-                config.getExtensionContextFilePaths(),
+                config.getExtensions(),
                 config.isTrustedFolder(),
                 settings.merged.context?.importFormat || 'tree',
                 config.getFileFilteringOptions(),

packages/cli/src/ui/components/views/ExtensionsList.test.tsx

@@ -29,7 +29,6 @@ describe('<ExtensionsList />', () => {
   const mockUIState = (
     extensions: unknown[],
     extensionsUpdateState: Map<string, ExtensionUpdateState>,
-    disabledExtensions: string[] = [],
   ) => {
     mockUseUIState.mockReturnValue({
       commandContext: createMockCommandContext({
@@ -37,13 +36,6 @@ describe('<ExtensionsList />', () => {
           config: {
             getExtensions: () => extensions,
           },
-          settings: {
-            merged: {
-              extensions: {
-                disabled: disabledExtensions,
-              },
-            },
-          },
         },
       }),
       extensionsUpdateState,
@@ -58,7 +50,7 @@ describe('<ExtensionsList />', () => {
   });
 
   it('should render a list of extensions with their version and status', () => {
-    mockUIState(mockExtensions, new Map(), ['ext-disabled']);
+    mockUIState(mockExtensions, new Map());
     const { lastFrame } = render(<ExtensionsList />);
     const output = lastFrame();
     expect(output).toContain('ext-one (v1.0.0) - active');

packages/cli/src/ui/components/views/ExtensionsList.tsx

@@ -11,8 +11,6 @@ import { ExtensionUpdateState } from '../../state/extensions.js';
 export const ExtensionsList = () => {
   const { commandContext, extensionsUpdateState } = useUIState();
   const allExtensions = commandContext.services.config!.getExtensions();
-  const settings = commandContext.services.settings;
-  const disabledExtensions = settings.merged.extensions?.disabled ?? [];
 
   if (allExtensions.length === 0) {
     return <Text>No extensions installed.</Text>;
@@ -24,8 +22,9 @@ export const ExtensionsList = () => {
       <Box flexDirection="column" paddingLeft={2}>
         {allExtensions.map((ext) => {
           const state = extensionsUpdateState.get(ext.name);
-          const isActive = !disabledExtensions.includes(ext.name);
+          const isActive = ext.isActive;
           const activeString = isActive ? 'active' : 'disabled';
+          const activeColor = isActive ? 'green' : 'grey';
 
           let stateColor = 'gray';
           const stateText = state || 'unknown state';
@@ -55,7 +54,7 @@ export const ExtensionsList = () => {
             <Box key={ext.name}>
               <Text>
                 <Text color="cyan">{`${ext.name} (v${ext.version})`}</Text>
-                {` - ${activeString}`}
+                <Text color={activeColor}>{` - ${activeString}`}</Text>
                 {<Text color={stateColor}>{` (${stateText})`}</Text>}
               </Text>
             </Box>

packages/cli/src/ui/components/views/McpStatus.tsx

@@ -115,8 +115,8 @@ export const McpStatus: React.FC<McpStatusProps> = ({
         }
 
         let serverDisplayName = serverName;
-        if (server.extensionName) {
-          serverDisplayName += ` (from ${server.extensionName})`;
+        if (server.extension?.name) {
+          serverDisplayName += ` (from ${server.extension?.name})`;
         }
 
         const toolCount = serverTools.length;

packages/cli/src/ui/hooks/useExtensionUpdates.test.ts

@@ -8,21 +8,18 @@ import { vi } from 'vitest';
 import * as fs from 'node:fs';
 import * as os from 'node:os';
 import * as path from 'node:path';
-import {
-  annotateActiveExtensions,
-  loadExtension,
-} from '../../config/extension.js';
+import { loadExtension } from '../../config/extension.js';
 import { createExtension } from '../../test-utils/createExtension.js';
 import { useExtensionUpdates } from './useExtensionUpdates.js';
 import { GEMINI_DIR, type GeminiCLIExtension } from '@google/gemini-cli-core';
 import { renderHook, waitFor } from '@testing-library/react';
 import { MessageType } from '../types.js';
-import { ExtensionEnablementManager } from '../../config/extensions/extensionEnablement.js';
 import {
   checkForAllExtensionUpdates,
   updateExtension,
 } from '../../config/extensions/update.js';
 import { ExtensionUpdateState } from '../state/extensions.js';
+import { ExtensionEnablementManager } from '../../config/extensions/extensionEnablement.js';
 
 vi.mock('os', async (importOriginal) => {
   const mockedOs = await importOriginal<typeof os>();
@@ -76,7 +73,7 @@ describe('useExtensionUpdates', () => {
     const cwd = '/test/cwd';
 
     vi.mocked(checkForAllExtensionUpdates).mockImplementation(
-      async (_extensions, dispatch, _cwd) => {
+      async (_extensions, _extensionEnablementManager, dispatch, _cwd) => {
         dispatch({
           type: 'SET_STATE',
           payload: {
@@ -88,7 +85,12 @@ describe('useExtensionUpdates', () => {
     );
 
     renderHook(() =>
-      useExtensionUpdates(extensions as GeminiCLIExtension[], addItem, cwd),
+      useExtensionUpdates(
+        extensions as GeminiCLIExtension[],
+        new ExtensionEnablementManager(),
+        addItem,
+        cwd,
+      ),
     );
 
     await waitFor(() => {
@@ -113,16 +115,17 @@ describe('useExtensionUpdates', () => {
         autoUpdate: true,
       },
     });
-    const extension = annotateActiveExtensions(
-      [loadExtension({ extensionDir, workspaceDir: tempHomeDir })!],
-      tempHomeDir,
-      new ExtensionEnablementManager(),
-    )[0];
+    const extensionEnablementManager = new ExtensionEnablementManager();
+    const extension = loadExtension({
+      extensionDir,
+      workspaceDir: tempHomeDir,
+      extensionEnablementManager,
+    })!;
 
     const addItem = vi.fn();
 
     vi.mocked(checkForAllExtensionUpdates).mockImplementation(
-      async (_extensions, dispatch, _cwd) => {
+      async (_extensions, _extensionEnablementManager, dispatch, _cwd) => {
         dispatch({
           type: 'SET_STATE',
           payload: {
@@ -139,7 +142,14 @@ describe('useExtensionUpdates', () => {
       name: '',
     });
 
-    renderHook(() => useExtensionUpdates([extension], addItem, tempHomeDir));
+    renderHook(() =>
+      useExtensionUpdates(
+        [extension],
+        extensionEnablementManager,
+        addItem,
+        tempHomeDir,
+      ),
+    );
 
     await waitFor(
       () => {
@@ -177,25 +187,24 @@ describe('useExtensionUpdates', () => {
       },
     });
 
-    const extensions = annotateActiveExtensions(
-      [
-        loadExtension({
-          extensionDir: extensionDir1,
-          workspaceDir: tempHomeDir,
-        })!,
-        loadExtension({
-          extensionDir: extensionDir2,
-          workspaceDir: tempHomeDir,
-        })!,
-      ],
-      tempHomeDir,
-      new ExtensionEnablementManager(),
-    );
+    const extensionEnablementManager = new ExtensionEnablementManager();
+    const extensions = [
+      loadExtension({
+        extensionDir: extensionDir1,
+        workspaceDir: tempHomeDir,
+        extensionEnablementManager,
+      })!,
+      loadExtension({
+        extensionDir: extensionDir2,
+        workspaceDir: tempHomeDir,
+        extensionEnablementManager,
+      })!,
+    ];
 
     const addItem = vi.fn();
 
     vi.mocked(checkForAllExtensionUpdates).mockImplementation(
-      async (_extensions, dispatch, _cwd) => {
+      async (_extensions, _extensionEnablementManager, dispatch, _cwd) => {
         dispatch({
           type: 'SET_STATE',
           payload: {
@@ -225,7 +234,14 @@ describe('useExtensionUpdates', () => {
         name: '',
       });
 
-    renderHook(() => useExtensionUpdates(extensions, addItem, tempHomeDir));
+    renderHook(() =>
+      useExtensionUpdates(
+        extensions,
+        extensionEnablementManager,
+        addItem,
+        tempHomeDir,
+      ),
+    );
 
     await waitFor(
       () => {
@@ -282,7 +298,7 @@ describe('useExtensionUpdates', () => {
     const cwd = '/test/cwd';
 
     vi.mocked(checkForAllExtensionUpdates).mockImplementation(
-      async (_extensions, dispatch, _cwd) => {
+      async (_extensions, _extensionEnablementManager, dispatch, _cwd) => {
         dispatch({ type: 'BATCH_CHECK_START' });
         dispatch({
           type: 'SET_STATE',
@@ -303,8 +319,14 @@ describe('useExtensionUpdates', () => {
       },
     );
 
+    const extensionEnablementManager = new ExtensionEnablementManager();
     renderHook(() =>
-      useExtensionUpdates(extensions as GeminiCLIExtension[], addItem, cwd),
+      useExtensionUpdates(
+        extensions as GeminiCLIExtension[],
+        extensionEnablementManager,
+        addItem,
+        cwd,
+      ),
     );
 
     await waitFor(() => {

packages/cli/src/ui/hooks/useExtensionUpdates.ts

@@ -23,6 +23,7 @@ import {
   type ExtensionUpdateInfo,
 } from '../../config/extension.js';
 import { checkExhaustive } from '../../utils/checks.js';
+import type { ExtensionEnablementManager } from '../../config/extensions/extensionEnablement.js';
 
 type ConfirmationRequestWrapper = {
   prompt: React.ReactNode;
@@ -49,6 +50,7 @@ function confirmationRequestsReducer(
 
 export const useExtensionUpdates = (
   extensions: GeminiCLIExtension[],
+  extensionEnablementManager: ExtensionEnablementManager,
   addItem: UseHistoryManagerReturn['addItem'],
   cwd: string,
 ) => {
@@ -93,11 +95,13 @@ export const useExtensionUpdates = (
     if (extensionsToCheck.length === 0) return;
     checkForAllExtensionUpdates(
       extensionsToCheck,
+      extensionEnablementManager,
       dispatchExtensionStateUpdate,
       cwd,
     );
   }, [
     extensions,
+    extensionEnablementManager,
     extensionsUpdateState.extensionStatuses,
     cwd,
     dispatchExtensionStateUpdate,
@@ -154,6 +158,7 @@ export const useExtensionUpdates = (
       } else {
         const updatePromise = updateExtension(
           extension,
+          extensionEnablementManager,
           cwd,
           (description) =>
             requestConsentInteractive(
@@ -210,6 +215,7 @@ export const useExtensionUpdates = (
     }
   }, [
     extensions,
+    extensionEnablementManager,
     extensionsUpdateState,
     addConfirmUpdateExtensionRequest,
     addItem,