4 files changed, 21 insertions, 12 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);