Initialize MCP tools once at start up instead of every time we auth. (#3483)

packages/cli/src/gemini.tsx

@@ -115,15 +115,7 @@ export async function main() {
 
   setMaxSizedBoxDebugging(config.getDebugMode());
 
-  // Initialize centralized FileDiscoveryService
+  await config.initialize();
-  config.getFileService();
-  if (config.getCheckpointingEnabled()) {
-    try {
-      await config.getGitService();
-    } catch {
-      // For now swallow the error, later log it.
-    }
-  }
 
   if (settings.merged.theme) {
     if (!themeManager.setActiveTheme(settings.merged.theme)) {
@@ -133,12 +125,11 @@ export async function main() {
     }
   }
 
+  // hop into sandbox if we are outside and sandboxing is enabled
+  if (!process.env.SANDBOX) {
     const memoryArgs = settings.merged.autoConfigureMaxOldSpaceSize
       ? getNodeMemoryArgs(config)
       : [];
-
-  // hop into sandbox if we are outside and sandboxing is enabled
-  if (!process.env.SANDBOX) {
     const sandboxConfig = config.getSandbox();
     if (sandboxConfig) {
       if (settings.merged.selectedAuthType) {

packages/core/src/code_assist/converter.ts

@@ -80,6 +80,7 @@ interface VertexGenerateContentResponse {
   promptFeedback?: GenerateContentResponsePromptFeedback;
   usageMetadata?: GenerateContentResponseUsageMetadata;
 }
+
 export interface CaCountTokenRequest {
   request: VertexCountTokenRequest;
 }

packages/core/src/config/config.ts

@@ -232,32 +232,30 @@ export class Config {
     }
   }
 
+  async initialize(): Promise<void> {
+    // Initialize centralized FileDiscoveryService
+    this.getFileService();
+    if (this.getCheckpointingEnabled()) {
+      try {
+        await this.getGitService();
+      } catch {
+        // For now swallow the error, later log it.
+      }
+    }
+    this.toolRegistry = await this.createToolRegistry();
+  }
+
   async refreshAuth(authMethod: AuthType) {
-    // Always use the original default model when switching auth methods
+    this.contentGeneratorConfig = await createContentGeneratorConfig(
-    // This ensures users don't stay on Flash after switching between auth types
+      this.model,
-    // and allows API key users to get proper fallback behavior from getEffectiveModel
-    const modelToUse = this.model; // Use the original default model
-
-    // Temporarily clear contentGeneratorConfig to prevent getModel() from returning
-    // the previous session's model (which might be Flash)
-    this.contentGeneratorConfig = undefined!;
-
-    const contentConfig = await createContentGeneratorConfig(
-      modelToUse,
       authMethod,
-      this,
     );
 
-    const gc = new GeminiClient(this);
+    this.geminiClient = new GeminiClient(this);
-    this.geminiClient = gc;
+    await this.geminiClient.initialize(this.contentGeneratorConfig);
-    this.toolRegistry = await createToolRegistry(this);
-    await gc.initialize(contentConfig);
-    this.contentGeneratorConfig = contentConfig;
 
     // Reset the session flag since we're explicitly changing auth and using default model
     this.modelSwitchedDuringSession = false;
-
-    // Note: In the future, we may want to reset any cached state when switching auth methods
   }
 
   getSessionId(): string {
@@ -469,19 +467,18 @@ export class Config {
 
     return { memoryContent, fileCount };
   }
-}
 
-export function createToolRegistry(config: Config): Promise<ToolRegistry> {
+  async createToolRegistry(): Promise<ToolRegistry> {
-  const registry = new ToolRegistry(config);
+    const registry = new ToolRegistry(this);
-  const targetDir = config.getTargetDir();
+    const targetDir = this.getTargetDir();
 
     // helper to create & register core tools that are enabled
     // eslint-disable-next-line @typescript-eslint/no-explicit-any
     const registerCoreTool = (ToolClass: any, ...args: unknown[]) => {
       const className = ToolClass.name;
       const toolName = ToolClass.Name || className;
-    const coreTools = config.getCoreTools();
+      const coreTools = this.getCoreTools();
-    const excludeTools = config.getExcludeTools();
+      const excludeTools = this.getExcludeTools();
 
       let isEnabled = false;
       if (coreTools === undefined) {
@@ -496,7 +493,10 @@ export function createToolRegistry(config: Config): Promise<ToolRegistry> {
         );
       }
 
-    if (excludeTools?.includes(className) || excludeTools?.includes(toolName)) {
+      if (
+        excludeTools?.includes(className) ||
+        excludeTools?.includes(toolName)
+      ) {
         isEnabled = false;
       }
 
@@ -505,22 +505,21 @@ export function createToolRegistry(config: Config): Promise<ToolRegistry> {
       }
     };
 
-  registerCoreTool(LSTool, targetDir, config);
+    registerCoreTool(LSTool, targetDir, this);
-  registerCoreTool(ReadFileTool, targetDir, config);
+    registerCoreTool(ReadFileTool, targetDir, this);
     registerCoreTool(GrepTool, targetDir);
-  registerCoreTool(GlobTool, targetDir, config);
+    registerCoreTool(GlobTool, targetDir, this);
-  registerCoreTool(EditTool, config);
+    registerCoreTool(EditTool, this);
-  registerCoreTool(WriteFileTool, config);
+    registerCoreTool(WriteFileTool, this);
-  registerCoreTool(WebFetchTool, config);
+    registerCoreTool(WebFetchTool, this);
-  registerCoreTool(ReadManyFilesTool, targetDir, config);
+    registerCoreTool(ReadManyFilesTool, targetDir, this);
-  registerCoreTool(ShellTool, config);
+    registerCoreTool(ShellTool, this);
     registerCoreTool(MemoryTool);
-  registerCoreTool(WebSearchTool, config);
+    registerCoreTool(WebSearchTool, this);
-  return (async () => {
+
     await registry.discoverTools();
     return registry;
-  })();
   }
-
+}
 // Export model constants for use in CLI
 export { DEFAULT_GEMINI_FLASH_MODEL };

packages/core/src/core/contentGenerator.ts

@@ -50,7 +50,6 @@ export type ContentGeneratorConfig = {
 export async function createContentGeneratorConfig(
   model: string | undefined,
   authType: AuthType | undefined,
-  config?: { getModel?: () => string },
 ): Promise<ContentGeneratorConfig> {
   const geminiApiKey = process.env.GEMINI_API_KEY;
   const googleApiKey = process.env.GOOGLE_API_KEY;
@@ -58,7 +57,7 @@ export async function createContentGeneratorConfig(
   const googleCloudLocation = process.env.GOOGLE_CLOUD_LOCATION;
 
   // Use runtime model from config if available, otherwise fallback to parameter or default
-  const effectiveModel = config?.getModel?.() || model || DEFAULT_GEMINI_MODEL;
+  const effectiveModel = model || DEFAULT_GEMINI_MODEL;
 
   const contentGeneratorConfig: ContentGeneratorConfig = {
     model: effectiveModel,

packages/core/src/tools/edit.test.ts

@@ -38,21 +38,19 @@ describe('EditTool', () => {
   let tempDir: string;
   let rootDir: string;
   let mockConfig: Config;
+  let geminiClient: any;
 
   beforeEach(() => {
     tempDir = fs.mkdtempSync(path.join(os.tmpdir(), 'edit-tool-test-'));
     rootDir = path.join(tempDir, 'root');
     fs.mkdirSync(rootDir);
 
-    // The client instance that EditTool will use
+    geminiClient = {
-    const mockClientInstanceWithGenerateJson = {
       generateJson: mockGenerateJson, // mockGenerateJson is already defined and hoisted
     };
 
     mockConfig = {
-      getGeminiClient: vi
+      getGeminiClient: vi.fn().mockReturnValue(geminiClient),
-        .fn()
-        .mockReturnValue(mockClientInstanceWithGenerateJson),
       getTargetDir: () => rootDir,
       getApprovalMode: vi.fn(),
       setApprovalMode: vi.fn(),
@@ -339,7 +337,7 @@ describe('EditTool', () => {
           mockCalled = true;
           expect(content).toBe(originalContent);
           expect(p).toBe(params);
-          expect(client).toBe((tool as any).client);
+          expect(client).toBe(geminiClient);
           return {
             params: {
               file_path: filePath,

packages/core/src/tools/edit.ts

@@ -18,7 +18,6 @@ import {
 import { SchemaValidator } from '../utils/schemaValidator.js';
 import { makeRelative, shortenPath } from '../utils/paths.js';
 import { isNodeError } from '../utils/errors.js';
-import { GeminiClient } from '../core/client.js';
 import { Config, ApprovalMode } from '../config/config.js';
 import { ensureCorrectEdit } from '../utils/editCorrector.js';
 import { DEFAULT_DIFF_OPTIONS } from './diffOptions.js';
@@ -72,15 +71,13 @@ export class EditTool
   implements ModifiableTool<EditToolParams>
 {
   static readonly Name = 'replace';
-  private readonly config: Config;
   private readonly rootDirectory: string;
-  private readonly client: GeminiClient;
 
   /**
    * Creates a new instance of the EditLogic
    * @param rootDirectory Root directory to ground this tool in.
    */
-  constructor(config: Config) {
+  constructor(private readonly config: Config) {
     super(
       EditTool.Name,
       'Edit',
@@ -123,9 +120,7 @@ Expectation for required parameters:
         type: 'object',
       },
     );
-    this.config = config;
     this.rootDirectory = path.resolve(this.config.getTargetDir());
-    this.client = config.getGeminiClient();
   }
 
   /**
@@ -239,7 +234,7 @@ Expectation for required parameters:
         params.file_path,
         currentContent,
         params,
-        this.client,
+        this.config.getGeminiClient(),
         abortSignal,
       );
       finalOldString = correctedEdit.params.old_string;

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

@@ -154,8 +154,6 @@ export class ToolRegistry {
     for (const tool of this.tools.values()) {
       if (tool instanceof DiscoveredTool || tool instanceof DiscoveredMCPTool) {
         this.tools.delete(tool.name);
-      } else {
-        // Keep manually registered tools
       }
     }
 

packages/core/src/tools/write-file.ts

@@ -23,7 +23,6 @@ import {
   ensureCorrectEdit,
   ensureCorrectFileContent,
 } from '../utils/editCorrector.js';
-import { GeminiClient } from '../core/client.js';
 import { DEFAULT_DIFF_OPTIONS } from './diffOptions.js';
 import { ModifiableTool, ModifyContext } from './modifiable-tool.js';
 import { getSpecificMimeType } from '../utils/fileUtils.js';
@@ -67,7 +66,6 @@ export class WriteFileTool
   implements ModifiableTool<WriteFileToolParams>
 {
   static readonly Name: string = 'write_file';
-  private readonly client: GeminiClient;
 
   constructor(private readonly config: Config) {
     super(
@@ -92,8 +90,6 @@ export class WriteFileTool
         type: 'object',
       },
     );
-
-    this.client = this.config.getGeminiClient();
   }
 
   /**
@@ -374,7 +370,7 @@ export class WriteFileTool
           new_string: proposedContent,
           file_path: filePath,
         },
-        this.client,
+        this.config.getGeminiClient(),
         abortSignal,
       );
       correctedContent = correctedParams.new_string;
@@ -382,7 +378,7 @@ export class WriteFileTool
       // This implies new file (ENOENT)
       correctedContent = await ensureCorrectFileContent(
         proposedContent,
-        this.client,
+        this.config.getGeminiClient(),
         abortSignal,
       );
     }