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

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

packages/core/src/config/config.ts

@@ -189,7 +189,7 @@ export class MCPServerConfig {
     readonly description?: string,
     readonly includeTools?: string[],
     readonly excludeTools?: string[],
-    readonly extensionName?: string,
+    readonly extension?: GeminiCLIExtension,
     // OAuth configuration
     readonly oauth?: MCPOAuthConfig,
     readonly authProviderType?: AuthProviderType,
@@ -249,11 +249,11 @@ export interface ConfigParameters {
   includeDirectories?: string[];
   bugCommand?: BugCommandSettings;
   model: string;
-  extensionContextFilePaths?: string[];
   maxSessionTurns?: number;
   experimentalZedIntegration?: boolean;
   listExtensions?: boolean;
   extensions?: GeminiCLIExtension[];
+  enabledExtensions?: string[];
   blockedMcpServers?: Array<{ name: string; extensionName: string }>;
   noBrowser?: boolean;
   summarizeToolOutput?: Record<string, SummarizeToolOutputSettings>;
@@ -332,7 +332,6 @@ export class Config {
   private readonly cwd: string;
   private readonly bugCommand: BugCommandSettings | undefined;
   private model: string;
-  private readonly extensionContextFilePaths: string[];
   private readonly noBrowser: boolean;
   private readonly folderTrust: boolean;
   private ideMode: boolean;
@@ -341,6 +340,7 @@ export class Config {
   private readonly maxSessionTurns: number;
   private readonly listExtensions: boolean;
   private readonly _extensions: GeminiCLIExtension[];
+  private readonly _enabledExtensions: string[];
   private readonly _blockedMcpServers: Array<{
     name: string;
     extensionName: string;
@@ -436,12 +436,12 @@ export class Config {
     this.fileDiscoveryService = params.fileDiscoveryService ?? null;
     this.bugCommand = params.bugCommand;
     this.model = params.model;
-    this.extensionContextFilePaths = params.extensionContextFilePaths ?? [];
     this.maxSessionTurns = params.maxSessionTurns ?? -1;
     this.experimentalZedIntegration =
       params.experimentalZedIntegration ?? false;
     this.listExtensions = params.listExtensions ?? false;
     this._extensions = params.extensions ?? [];
+    this._enabledExtensions = params.enabledExtensions ?? [];
     this._blockedMcpServers = params.blockedMcpServers ?? [];
     this.noBrowser = params.noBrowser ?? false;
     this.summarizeToolOutput = params.summarizeToolOutput;
@@ -542,7 +542,7 @@ export class Config {
 
   async refreshAuth(authMethod: AuthType) {
     // Vertex and Genai have incompatible encryption and sending history with
-    // throughtSignature from Genai to Vertex will fail, we need to strip them
+    // thoughtSignature from Genai to Vertex will fail, we need to strip them
     if (
       this.contentGeneratorConfig?.authType === AuthType.USE_GEMINI &&
       authMethod === AuthType.LOGIN_WITH_GOOGLE
@@ -869,10 +869,6 @@ export class Config {
     return this.usageStatisticsEnabled;
   }
 
-  getExtensionContextFilePaths(): string[] {
-    return this.extensionContextFilePaths;
-  }
-
   getExperimentalZedIntegration(): boolean {
     return this.experimentalZedIntegration;
   }
@@ -889,6 +885,12 @@ export class Config {
     return this._extensions;
   }
 
+  // The list of explicitly enabled extensions, if any were given, may contain
+  // the string "none".
+  getEnabledExtensions(): string[] {
+    return this._enabledExtensions;
+  }
+
   getBlockedMcpServers(): Array<{ name: string; extensionName: string }> {
     return this._blockedMcpServers;
   }

packages/core/src/tools/mcp-client-manager.test.ts

@@ -8,8 +8,6 @@ import { afterEach, describe, expect, it, vi } from 'vitest';
 import { McpClientManager } from './mcp-client-manager.js';
 import { McpClient } from './mcp-client.js';
 import type { ToolRegistry } from './tool-registry.js';
-import type { PromptRegistry } from '../prompts/prompt-registry.js';
-import type { WorkspaceContext } from '../utils/workspaceContext.js';
 import type { Config } from '../config/config.js';
 
 vi.mock('./mcp-client.js', async () => {
@@ -38,18 +36,16 @@ describe('McpClientManager', () => {
     vi.mocked(McpClient).mockReturnValue(
       mockedMcpClient as unknown as McpClient,
     );
-    const manager = new McpClientManager(
-      {
-        'test-server': {},
-      },
-      '',
-      {} as ToolRegistry,
-      {} as PromptRegistry,
-      false,
-      {} as WorkspaceContext,
-    );
+    const manager = new McpClientManager({} as ToolRegistry);
     await manager.discoverAllMcpTools({
       isTrustedFolder: () => true,
+      getMcpServers: () => ({
+        'test-server': {},
+      }),
+      getMcpServerCommand: () => '',
+      getPromptRegistry: () => {},
+      getDebugMode: () => false,
+      getWorkspaceContext: () => {},
     } as unknown as Config);
     expect(mockedMcpClient.connect).toHaveBeenCalledOnce();
     expect(mockedMcpClient.discover).toHaveBeenCalledOnce();
@@ -65,18 +61,16 @@ describe('McpClientManager', () => {
     vi.mocked(McpClient).mockReturnValue(
       mockedMcpClient as unknown as McpClient,
     );
-    const manager = new McpClientManager(
-      {
-        'test-server': {},
-      },
-      '',
-      {} as ToolRegistry,
-      {} as PromptRegistry,
-      false,
-      {} as WorkspaceContext,
-    );
+    const manager = new McpClientManager({} as ToolRegistry);
     await manager.discoverAllMcpTools({
       isTrustedFolder: () => false,
+      getMcpServers: () => ({
+        'test-server': {},
+      }),
+      getMcpServerCommand: () => '',
+      getPromptRegistry: () => {},
+      getDebugMode: () => false,
+      getWorkspaceContext: () => {},
     } as unknown as Config);
     expect(mockedMcpClient.connect).not.toHaveBeenCalled();
     expect(mockedMcpClient.discover).not.toHaveBeenCalled();

packages/core/src/tools/mcp-client-manager.ts

@@ -4,9 +4,8 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import type { Config, MCPServerConfig } from '../config/config.js';
+import type { Config } from '../config/config.js';
 import type { ToolRegistry } from './tool-registry.js';
-import type { PromptRegistry } from '../prompts/prompt-registry.js';
 import {
   McpClient,
   MCPDiscoveryState,
@@ -14,7 +13,6 @@ import {
 } from './mcp-client.js';
 import { getErrorMessage } from '../utils/errors.js';
 import type { EventEmitter } from 'node:events';
-import type { WorkspaceContext } from '../utils/workspaceContext.js';
 
 /**
  * Manages the lifecycle of multiple MCP clients, including local child processes.
@@ -23,30 +21,12 @@ import type { WorkspaceContext } from '../utils/workspaceContext.js';
  */
 export class McpClientManager {
   private clients: Map<string, McpClient> = new Map();
-  private readonly mcpServers: Record<string, MCPServerConfig>;
-  private readonly mcpServerCommand: string | undefined;
   private readonly toolRegistry: ToolRegistry;
-  private readonly promptRegistry: PromptRegistry;
-  private readonly debugMode: boolean;
-  private readonly workspaceContext: WorkspaceContext;
   private discoveryState: MCPDiscoveryState = MCPDiscoveryState.NOT_STARTED;
   private readonly eventEmitter?: EventEmitter;
 
-  constructor(
-    mcpServers: Record<string, MCPServerConfig>,
-    mcpServerCommand: string | undefined,
-    toolRegistry: ToolRegistry,
-    promptRegistry: PromptRegistry,
-    debugMode: boolean,
-    workspaceContext: WorkspaceContext,
-    eventEmitter?: EventEmitter,
-  ) {
-    this.mcpServers = mcpServers;
-    this.mcpServerCommand = mcpServerCommand;
+  constructor(toolRegistry: ToolRegistry, eventEmitter?: EventEmitter) {
     this.toolRegistry = toolRegistry;
-    this.promptRegistry = promptRegistry;
-    this.debugMode = debugMode;
-    this.workspaceContext = workspaceContext;
     this.eventEmitter = eventEmitter;
   }
 
@@ -62,22 +42,23 @@ export class McpClientManager {
     await this.stop();
 
     const servers = populateMcpServerCommand(
-      this.mcpServers,
-      this.mcpServerCommand,
+      cliConfig.getMcpServers() || {},
+      cliConfig.getMcpServerCommand(),
     );
 
     this.discoveryState = MCPDiscoveryState.IN_PROGRESS;
 
     this.eventEmitter?.emit('mcp-client-update', this.clients);
-    const discoveryPromises = Object.entries(servers).map(
-      async ([name, config]) => {
+    const discoveryPromises = Object.entries(servers)
+      .filter(([_, config]) => !config.extension || config.extension.isActive)
+      .map(async ([name, config]) => {
         const client = new McpClient(
           name,
           config,
           this.toolRegistry,
-          this.promptRegistry,
-          this.workspaceContext,
-          this.debugMode,
+          cliConfig.getPromptRegistry(),
+          cliConfig.getWorkspaceContext(),
+          cliConfig.getDebugMode(),
         );
         this.clients.set(name, client);
 
@@ -95,8 +76,7 @@ export class McpClientManager {
             )}`,
           );
         }
-      },
-    );
+      });
 
     await Promise.all(discoveryPromises);
     this.discoveryState = MCPDiscoveryState.COMPLETED;

packages/core/src/tools/tool-registry.ts

@@ -174,15 +174,7 @@ export class ToolRegistry {
 
   constructor(config: Config, eventEmitter?: EventEmitter) {
     this.config = config;
-    this.mcpClientManager = new McpClientManager(
-      this.config.getMcpServers() ?? {},
-      this.config.getMcpServerCommand(),
-      this,
-      this.config.getPromptRegistry(),
-      this.config.getDebugMode(),
-      this.config.getWorkspaceContext(),
-      eventEmitter,
-    );
+    this.mcpClientManager = new McpClientManager(this, eventEmitter);
   }
 
   /**

packages/core/src/utils/memoryDiscovery.test.ts

@@ -15,6 +15,7 @@ import {
 } from '../tools/memoryTool.js';
 import { FileDiscoveryService } from '../services/fileDiscoveryService.js';
 import { GEMINI_DIR } from './paths.js';
+import type { GeminiCLIExtension } from '../config/config.js';
 
 vi.mock('os', async (importOriginal) => {
   const actualOs = await importOriginal<typeof os>();
@@ -87,7 +88,7 @@ describe('loadServerHierarchicalMemory', () => {
         [],
         false,
         new FileDiscoveryService(projectRoot),
-        [],
+        [], // extensions
         false, // untrusted
       );
 
@@ -116,7 +117,7 @@ describe('loadServerHierarchicalMemory', () => {
           [],
           false,
           new FileDiscoveryService(projectRoot),
-          [],
+          [], // extensions
           false, // untrusted
         );
 
@@ -132,7 +133,7 @@ describe('loadServerHierarchicalMemory', () => {
       [],
       false,
       new FileDiscoveryService(projectRoot),
-      [],
+      [], // extensions
       DEFAULT_FOLDER_TRUST,
     );
 
@@ -154,7 +155,7 @@ describe('loadServerHierarchicalMemory', () => {
       [],
       false,
       new FileDiscoveryService(projectRoot),
-      [],
+      [], // extensions
       DEFAULT_FOLDER_TRUST,
     );
 
@@ -181,7 +182,7 @@ default context content
       [],
       false,
       new FileDiscoveryService(projectRoot),
-      [],
+      [], // extensions
       DEFAULT_FOLDER_TRUST,
     );
 
@@ -212,7 +213,7 @@ custom context content
       [],
       false,
       new FileDiscoveryService(projectRoot),
-      [],
+      [], // extensions
       DEFAULT_FOLDER_TRUST,
     );
 
@@ -247,7 +248,7 @@ cwd context content
       [],
       false,
       new FileDiscoveryService(projectRoot),
-      [],
+      [], // extensions
       DEFAULT_FOLDER_TRUST,
     );
 
@@ -279,7 +280,7 @@ Subdir custom memory
       [],
       false,
       new FileDiscoveryService(projectRoot),
-      [],
+      [], // extensions
       DEFAULT_FOLDER_TRUST,
     );
 
@@ -311,7 +312,7 @@ Src directory memory
       [],
       false,
       new FileDiscoveryService(projectRoot),
-      [],
+      [], // extensions
       DEFAULT_FOLDER_TRUST,
     );
 
@@ -355,7 +356,7 @@ Subdir memory
       [],
       false,
       new FileDiscoveryService(projectRoot),
-      [],
+      [], // extensions
       DEFAULT_FOLDER_TRUST,
     );
 
@@ -408,7 +409,7 @@ Subdir memory
       [],
       false,
       new FileDiscoveryService(projectRoot),
-      [],
+      [], // extensions
       DEFAULT_FOLDER_TRUST,
       'tree',
       {
@@ -444,7 +445,7 @@ My code memory
       [],
       true,
       new FileDiscoveryService(projectRoot),
-      [],
+      [], // extensions
       DEFAULT_FOLDER_TRUST,
       'tree', // importFormat
       {
@@ -466,7 +467,7 @@ My code memory
       [],
       false,
       new FileDiscoveryService(projectRoot),
-      [],
+      [], // extensions
       DEFAULT_FOLDER_TRUST,
     );
 
@@ -488,7 +489,12 @@ My code memory
       [],
       false,
       new FileDiscoveryService(projectRoot),
-      [extensionFilePath],
+      [
+        {
+          contextFiles: [extensionFilePath],
+          isActive: true,
+        } as GeminiCLIExtension,
+      ], // extensions
       DEFAULT_FOLDER_TRUST,
     );
 
@@ -515,7 +521,7 @@ Extension memory content
       [includedDir],
       false,
       new FileDiscoveryService(projectRoot),
-      [],
+      [], // extensions
       DEFAULT_FOLDER_TRUST,
     );
 
@@ -550,7 +556,7 @@ included directory memory
       createdFiles.map((f) => path.dirname(f)),
       false,
       new FileDiscoveryService(projectRoot),
-      [],
+      [], // extensions
       DEFAULT_FOLDER_TRUST,
     );
 
@@ -585,7 +591,7 @@ included directory memory
       [childDir, parentDir], // Deliberately include duplicates
       false,
       new FileDiscoveryService(projectRoot),
-      [],
+      [], // extensions
       DEFAULT_FOLDER_TRUST,
     );
 

packages/core/src/utils/memoryDiscovery.ts

@@ -15,6 +15,7 @@ import { processImports } from './memoryImportProcessor.js';
 import type { FileFilteringOptions } from '../config/constants.js';
 import { DEFAULT_MEMORY_FILE_FILTERING_OPTIONS } from '../config/constants.js';
 import { GEMINI_DIR } from './paths.js';
+import type { GeminiCLIExtension } from '../config/config.js';
 
 // Simple console logger, similar to the one previously in CLI's config.ts
 // TODO: Integrate with a more robust server-side logger if available/appropriate.
@@ -84,7 +85,6 @@ async function getGeminiMdFilePathsInternal(
   userHomePath: string,
   debugMode: boolean,
   fileService: FileDiscoveryService,
-  extensionContextFilePaths: string[] = [],
   folderTrust: boolean,
   fileFilteringOptions: FileFilteringOptions,
   maxDirs: number,
@@ -107,7 +107,6 @@ async function getGeminiMdFilePathsInternal(
         userHomePath,
         debugMode,
         fileService,
-        extensionContextFilePaths,
         folderTrust,
         fileFilteringOptions,
         maxDirs,
@@ -137,7 +136,6 @@ async function getGeminiMdFilePathsInternalForEachDir(
   userHomePath: string,
   debugMode: boolean,
   fileService: FileDiscoveryService,
-  extensionContextFilePaths: string[] = [],
   folderTrust: boolean,
   fileFilteringOptions: FileFilteringOptions,
   maxDirs: number,
@@ -226,11 +224,6 @@ async function getGeminiMdFilePathsInternalForEachDir(
     }
   }
 
-  // Add extension context file paths.
-  for (const extensionPath of extensionContextFilePaths) {
-    allPaths.add(extensionPath);
-  }
-
   const finalPaths = Array.from(allPaths);
 
   if (debugMode)
@@ -343,7 +336,7 @@ export async function loadServerHierarchicalMemory(
   includeDirectoriesToReadGemini: readonly string[],
   debugMode: boolean,
   fileService: FileDiscoveryService,
-  extensionContextFilePaths: string[] = [],
+  extensions: GeminiCLIExtension[],
   folderTrust: boolean,
   importFormat: 'flat' | 'tree' = 'tree',
   fileFilteringOptions?: FileFilteringOptions,
@@ -363,11 +356,18 @@ export async function loadServerHierarchicalMemory(
     userHomePath,
     debugMode,
     fileService,
-    extensionContextFilePaths,
     folderTrust,
     fileFilteringOptions || DEFAULT_MEMORY_FILE_FILTERING_OPTIONS,
     maxDirs,
   );
+
+  // Add extension file paths separately since they may be conditionally enabled.
+  filePaths.push(
+    ...extensions
+      .filter((ext) => ext.isActive)
+      .flatMap((ext) => ext.contextFiles),
+  );
+
   if (filePaths.length === 0) {
     if (debugMode)
       logger.debug('No GEMINI.md files found in hierarchy of the workspace.');