Refactor: Use config.getGeminiClient() for GeminiClient instantiation (#715)

5 files changed, 18 insertions, 4 deletions

diff --git a/packages/core/src/core/client.ts b/packages/core/src/core/client.ts
index fcad1ef0..311f53bf 100644
--- a/packages/core/src/core/client.ts
+++ b/packages/core/src/core/client.ts
@@ -25,6 +25,7 @@ import { checkNextSpeaker } from '../utils/nextSpeakerChecker.js';
 import { reportError } from '../utils/errorReporting.js';
 import { GeminiChat } from './geminiChat.js';
 import { retryWithBackoff } from '../utils/retry.js';
+import { getErrorMessage } from '../utils/errors.js';
 
 export class GeminiClient {
   private chat: Promise<GeminiChat>;
@@ -158,8 +159,7 @@ export class GeminiClient {
         history,
         'startChat',
       );
-      const message = error instanceof Error ? error.message : 'Unknown error.';
-      throw new Error(`Failed to initialize chat: ${message}`);
+      throw new Error(`Failed to initialize chat: ${getErrorMessage(error)}`);
     }
   }
 
diff --git a/packages/core/src/tools/edit.test.ts b/packages/core/src/tools/edit.test.ts
index 3b93708a..87ae3b6f 100644
--- a/packages/core/src/tools/edit.test.ts
+++ b/packages/core/src/tools/edit.test.ts
@@ -39,7 +39,15 @@ describe('EditTool', () => {
     rootDir = path.join(tempDir, 'root');
     fs.mkdirSync(rootDir);
 
+    // The client instance that EditTool will use
+    const mockClientInstanceWithGenerateJson = {
+      generateJson: mockGenerateJson, // mockGenerateJson is already defined and hoisted
+    };
+
     mockConfig = {
+      getGeminiClient: vi
+        .fn()
+        .mockReturnValue(mockClientInstanceWithGenerateJson),
       getTargetDir: () => rootDir,
       getApprovalMode: vi.fn(() => false),
       setApprovalMode: vi.fn(),
diff --git a/packages/core/src/tools/edit.ts b/packages/core/src/tools/edit.ts
index b2f648f8..4a337faa 100644
--- a/packages/core/src/tools/edit.ts
+++ b/packages/core/src/tools/edit.ts
@@ -114,7 +114,7 @@ Expectation for required parameters:
     );
     this.config = config;
     this.rootDirectory = path.resolve(this.config.getTargetDir());
-    this.client = new GeminiClient(this.config);
+    this.client = config.getGeminiClient();
   }
 
   /**
diff --git a/packages/core/src/tools/write-file.test.ts b/packages/core/src/tools/write-file.test.ts
index c94edfd1..4646f30a 100644
--- a/packages/core/src/tools/write-file.test.ts
+++ b/packages/core/src/tools/write-file.test.ts
@@ -53,6 +53,7 @@ const mockConfigInternal = {
   getTargetDir: () => rootDir,
   getApprovalMode: vi.fn(() => ApprovalMode.DEFAULT),
   setApprovalMode: vi.fn(),
+  getGeminiClient: vi.fn(), // Initialize as a plain mock function
   getApiKey: () => 'test-key',
   getModel: () => 'test-model',
   getSandbox: () => false,
@@ -97,6 +98,11 @@ describe('WriteFileTool', () => {
     ) as Mocked<GeminiClient>;
     vi.mocked(GeminiClient).mockImplementation(() => mockGeminiClientInstance);
 
+    // Now that mockGeminiClientInstance is initialized, set the mock implementation for getGeminiClient
+    mockConfigInternal.getGeminiClient.mockReturnValue(
+      mockGeminiClientInstance,
+    );
+
     tool = new WriteFileTool(mockConfig);
 
     // Reset mocks before each test
diff --git a/packages/core/src/tools/write-file.ts b/packages/core/src/tools/write-file.ts
index 2e04a10a..57bfe8ed 100644
--- a/packages/core/src/tools/write-file.ts
+++ b/packages/core/src/tools/write-file.ts
@@ -77,7 +77,7 @@ export class WriteFileTool extends BaseTool<WriteFileToolParams, ToolResult> {
       },
     );
 
-    this.client = new GeminiClient(this.config);
+    this.client = this.config.getGeminiClient();
   }
 
   private isWithinRoot(pathToCheck: string): boolean {