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

4 files changed, 8 insertions, 21 deletions

diff --git a/packages/core/src/tools/edit.test.ts b/packages/core/src/tools/edit.test.ts
index ab42450a..84ad1daf 100644
--- a/packages/core/src/tools/edit.test.ts
+++ b/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
-    const mockClientInstanceWithGenerateJson = {
+    geminiClient = {
       generateJson: mockGenerateJson, // mockGenerateJson is already defined and hoisted
     };
 
     mockConfig = {
-      getGeminiClient: vi
-        .fn()
-        .mockReturnValue(mockClientInstanceWithGenerateJson),
+      getGeminiClient: vi.fn().mockReturnValue(geminiClient),
       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,
diff --git a/packages/core/src/tools/edit.ts b/packages/core/src/tools/edit.ts
index 2df01a22..f388b9f5 100644
--- a/packages/core/src/tools/edit.ts
+++ b/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;
diff --git a/packages/core/src/tools/tool-registry.ts b/packages/core/src/tools/tool-registry.ts
index bc628f03..1778c6d6 100644
--- a/packages/core/src/tools/tool-registry.ts
+++ b/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
       }
     }
 
diff --git a/packages/core/src/tools/write-file.ts b/packages/core/src/tools/write-file.ts
index 37a1ba78..ab30891b 100644
--- a/packages/core/src/tools/write-file.ts
+++ b/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,
       );
     }