9 files changed, 39 insertions, 16 deletions
diff --git a/packages/cli/src/nonInteractiveCli.test.ts b/packages/cli/src/nonInteractiveCli.test.ts
index c071dedc..97e4416a 100644
--- a/packages/cli/src/nonInteractiveCli.test.ts
+++ b/packages/cli/src/nonInteractiveCli.test.ts
@@ -51,6 +51,7 @@ describe('runNonInteractive', () => {
 
     mockConfig = {
       getToolRegistry: vi.fn().mockReturnValue(mockToolRegistry),
+      getGeminiClient: vi.fn().mockReturnValue(mockGeminiClient),
     } as unknown as Config;
 
     mockProcessStdoutWrite = vi.fn().mockImplementation(() => true);
diff --git a/packages/cli/src/nonInteractiveCli.ts b/packages/cli/src/nonInteractiveCli.ts
index 7505c736..45c9ce2d 100644
--- a/packages/cli/src/nonInteractiveCli.ts
+++ b/packages/cli/src/nonInteractiveCli.ts
@@ -6,7 +6,6 @@
 
 import {
   Config,
-  GeminiClient,
   ToolCallRequestInfo,
   executeToolCall,
   ToolRegistry,
@@ -39,7 +38,7 @@ export async function runNonInteractive(
   config: Config,
   input: string,
 ): Promise<void> {
-  const geminiClient = new GeminiClient(config);
+  const geminiClient = config.getGeminiClient();
   const toolRegistry: ToolRegistry = await config.getToolRegistry();
 
   const chat = await geminiClient.getChat();
diff --git a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx
index d46fab9e..4dfea5e0 100644
--- a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx
+++ b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx
@@ -25,20 +25,20 @@ const mockSendMessageStream = vi
   .mockReturnValue((async function* () {})());
 const mockStartChat = vi.fn();
 
-vi.mock('@gemini-code/core', async (importOriginal) => {
-  const actualCoreModule = (await importOriginal()) as any;
-  const MockedGeminiClientClass = vi.fn().mockImplementation(function (
-    this: any,
-    _config: any,
-  ) {
+const MockedGeminiClientClass = vi.hoisted(() =>
+  vi.fn().mockImplementation(function (this: any, _config: any) {
     // _config
     this.startChat = mockStartChat;
     this.sendMessageStream = mockSendMessageStream;
-  });
+  }),
+);
+
+vi.mock('@gemini-code/core', async (importOriginal) => {
+  const actualCoreModule = (await importOriginal()) as any;
   return {
     ...(actualCoreModule || {}),
-    GeminiClient: MockedGeminiClientClass,
-    // GeminiChat will be from actualCoreModule if it exists, otherwise undefined
+    GeminiClient: MockedGeminiClientClass, // Export the class for type checking or other direct uses
+    Config: actualCoreModule.Config, // Ensure Config is passed through
   };
 });
 
@@ -235,6 +235,14 @@ describe('useGeminiStream', () => {
 
     mockAddItem = vi.fn();
     mockSetShowHelp = vi.fn();
+    // Define the mock for getGeminiClient
+    const mockGetGeminiClient = vi.fn().mockImplementation(() => {
+      // MockedGeminiClientClass is defined in the module scope by the previous change.
+      // It will use the mockStartChat and mockSendMessageStream that are managed within beforeEach.
+      const clientInstance = new MockedGeminiClientClass(mockConfig);
+      return clientInstance;
+    });
+
     mockConfig = {
       apiKey: 'test-api-key',
       model: 'gemini-pro',
@@ -258,6 +266,7 @@ describe('useGeminiStream', () => {
       getToolRegistry: vi.fn(
         () => ({ getToolSchemaList: vi.fn(() => []) }) as any,
       ),
+      getGeminiClient: mockGetGeminiClient,
     } as unknown as Config;
     mockOnDebugMessage = vi.fn();
     mockHandleSlashCommand = vi.fn().mockReturnValue(false);
diff --git a/packages/cli/src/ui/hooks/useGeminiStream.ts b/packages/cli/src/ui/hooks/useGeminiStream.ts
index 284709cf..c50d4c43 100644
--- a/packages/cli/src/ui/hooks/useGeminiStream.ts
+++ b/packages/cli/src/ui/hooks/useGeminiStream.ts
@@ -145,7 +145,7 @@ export const useGeminiStream = (
     setInitError(null);
     if (!geminiClientRef.current) {
       try {
-        geminiClientRef.current = new GeminiClient(config);
+        geminiClientRef.current = config.getGeminiClient();
       } catch (error: unknown) {
         const errorMsg = `Failed to initialize client: ${getErrorMessage(error) || 'Unknown error'}`;
         setInitError(errorMsg);
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 {