From 357546a2aac918702f6ebfa4a97bd95ccd614e5d Mon Sep 17 00:00:00 2001
From: Tommaso Sciortino <sciortino@gmail.com>
Date: Mon, 7 Jul 2025 15:01:59 -0700
Subject: Initialize MCP tools once at start up instead of every time we auth.
 (#3483)

---
 packages/core/src/config/config.ts | 131 ++++++++++++++++++-------------------
 1 file changed, 65 insertions(+), 66 deletions(-)

(limited to 'packages/core/src/config/config.ts')

diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts
index fd96af91..ca0714f0 100644
--- a/packages/core/src/config/config.ts
+++ b/packages/core/src/config/config.ts
@@ -232,32 +232,30 @@ export class Config {
     }
   }
 
-  async refreshAuth(authMethod: AuthType) {
-    // Always use the original default model when switching auth methods
-    // This ensures users don't stay on Flash after switching between auth types
-    // 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!;
+  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();
+  }
 
-    const contentConfig = await createContentGeneratorConfig(
-      modelToUse,
+  async refreshAuth(authMethod: AuthType) {
+    this.contentGeneratorConfig = await createContentGeneratorConfig(
+      this.model,
       authMethod,
-      this,
     );
 
-    const gc = new GeminiClient(this);
-    this.geminiClient = gc;
-    this.toolRegistry = await createToolRegistry(this);
-    await gc.initialize(contentConfig);
-    this.contentGeneratorConfig = contentConfig;
+    this.geminiClient = new GeminiClient(this);
+    await this.geminiClient.initialize(this.contentGeneratorConfig);
 
     // 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,58 +467,59 @@ export class Config {
 
     return { memoryContent, fileCount };
   }
-}
 
-export function createToolRegistry(config: Config): Promise<ToolRegistry> {
-  const registry = new ToolRegistry(config);
-  const targetDir = config.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 excludeTools = config.getExcludeTools();
-
-    let isEnabled = false;
-    if (coreTools === undefined) {
-      isEnabled = true;
-    } else {
-      isEnabled = coreTools.some(
-        (tool) =>
-          tool === className ||
-          tool === toolName ||
-          tool.startsWith(`${className}(`) ||
-          tool.startsWith(`${toolName}(`),
-      );
-    }
-
-    if (excludeTools?.includes(className) || excludeTools?.includes(toolName)) {
-      isEnabled = false;
-    }
+  async createToolRegistry(): Promise<ToolRegistry> {
+    const registry = new ToolRegistry(this);
+    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 = this.getCoreTools();
+      const excludeTools = this.getExcludeTools();
+
+      let isEnabled = false;
+      if (coreTools === undefined) {
+        isEnabled = true;
+      } else {
+        isEnabled = coreTools.some(
+          (tool) =>
+            tool === className ||
+            tool === toolName ||
+            tool.startsWith(`${className}(`) ||
+            tool.startsWith(`${toolName}(`),
+        );
+      }
+
+      if (
+        excludeTools?.includes(className) ||
+        excludeTools?.includes(toolName)
+      ) {
+        isEnabled = false;
+      }
+
+      if (isEnabled) {
+        registry.registerTool(new ToolClass(...args));
+      }
+    };
 
-    if (isEnabled) {
-      registry.registerTool(new ToolClass(...args));
-    }
-  };
+    registerCoreTool(LSTool, targetDir, this);
+    registerCoreTool(ReadFileTool, targetDir, this);
+    registerCoreTool(GrepTool, targetDir);
+    registerCoreTool(GlobTool, targetDir, this);
+    registerCoreTool(EditTool, this);
+    registerCoreTool(WriteFileTool, this);
+    registerCoreTool(WebFetchTool, this);
+    registerCoreTool(ReadManyFilesTool, targetDir, this);
+    registerCoreTool(ShellTool, this);
+    registerCoreTool(MemoryTool);
+    registerCoreTool(WebSearchTool, this);
 
-  registerCoreTool(LSTool, targetDir, config);
-  registerCoreTool(ReadFileTool, targetDir, config);
-  registerCoreTool(GrepTool, targetDir);
-  registerCoreTool(GlobTool, targetDir, config);
-  registerCoreTool(EditTool, config);
-  registerCoreTool(WriteFileTool, config);
-  registerCoreTool(WebFetchTool, config);
-  registerCoreTool(ReadManyFilesTool, targetDir, config);
-  registerCoreTool(ShellTool, config);
-  registerCoreTool(MemoryTool);
-  registerCoreTool(WebSearchTool, config);
-  return (async () => {
     await registry.discoverTools();
     return registry;
-  })();
+  }
 }
-
 // Export model constants for use in CLI
 export { DEFAULT_GEMINI_FLASH_MODEL };
-- 
cgit v1.2.3