@google/gemini-cli-core 0.5.0-nightly.20250906.968e9389 → 0.5.0-nightly.20250909.2b05cf3b

package/dist/src/core/client.test.js

@@ -4,11 +4,9 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 import { describe, it, expect, vi, beforeEach, afterEach, } from 'vitest';
-import { GoogleGenAI } from '@google/genai';
 import { findIndexAfterFraction, isThinkingDefault, isThinkingSupported, GeminiClient, } from './client.js';
 import { AuthType, } from './contentGenerator.js';
 import {} from './geminiChat.js';
-import { Config } from '../config/config.js';
 import { CompressionStatus, GeminiEventType, Turn, } from './turn.js';
 import { getCoreSystemPrompt } from './prompts.js';
 import { DEFAULT_GEMINI_FLASH_MODEL } from '../config/models.js';
@@ -41,11 +39,7 @@ vi.mock('node:fs', () => {
     };
 });
 // --- Mocks ---
-const mockChatCreateFn = vi.fn();
-const mockGenerateContentFn = vi.fn();
-const mockEmbedContentFn = vi.fn();
 const mockTurnRunFn = vi.fn();
-vi.mock('@google/genai');
 vi.mock('./turn', async (importOriginal) => {
     const actual = await importOriginal();
     // Define a mock class that has the same shape as the real Turn
@@ -168,33 +162,24 @@ describe('isThinkingDefault', () => {
     });
 });
 describe('Gemini Client (client.ts)', () => {
+    let mockContentGenerator;
+    let mockConfig;
     let client;
+    let mockGenerateContentFn;
     beforeEach(async () => {
         vi.resetAllMocks();
+        mockGenerateContentFn = vi.fn().mockResolvedValue({
+            candidates: [{ content: { parts: [{ text: '{"key": "value"}' }] } }],
+        });
         // Disable 429 simulation for tests
         setSimulate429(false);
-        // Set up the mock for GoogleGenAI constructor and its methods
-        const MockedGoogleGenAI = vi.mocked(GoogleGenAI);
-        MockedGoogleGenAI.mockImplementation(() => {
-            const mock = {
-                chats: { create: mockChatCreateFn },
-                models: {
-                    generateContent: mockGenerateContentFn,
-                    embedContent: mockEmbedContentFn,
-                },
-            };
-            return mock;
-        });
-        mockChatCreateFn.mockResolvedValue({});
-        mockGenerateContentFn.mockResolvedValue({
-            candidates: [
-                {
-                    content: {
-                        parts: [{ text: '{"key": "value"}' }],
-                    },
-                },
-            ],
-        });
+        mockContentGenerator = {
+            generateContent: mockGenerateContentFn,
+            generateContentStream: vi.fn(),
+            countTokens: vi.fn(),
+            embedContent: vi.fn(),
+            batchEmbedContents: vi.fn(),
+        };
         // Because the GeminiClient constructor kicks off an async process (startChat)
         // that depends on a fully-formed Config object, we need to mock the
         // entire implementation of Config for these tests.
@@ -209,7 +194,7 @@ describe('Gemini Client (client.ts)', () => {
             vertexai: false,
             authType: AuthType.USE_GEMINI,
         };
-        const mockConfigObject = {
+        mockConfig = {
             getContentGeneratorConfig: vi
                 .fn()
                 .mockReturnValue(contentGeneratorConfig),
@@ -237,6 +222,7 @@ describe('Gemini Client (client.ts)', () => {
                 getDirectories: vi.fn().mockReturnValue(['/test/dir']),
             }),
             getGeminiClient: vi.fn(),
+            isInFallbackMode: vi.fn().mockReturnValue(false),
             setFallbackMode: vi.fn(),
             getChatCompression: vi.fn().mockReturnValue(undefined),
             getSkipNextSpeakerCheck: vi.fn().mockReturnValue(false),
@@ -245,36 +231,15 @@ describe('Gemini Client (client.ts)', () => {
             storage: {
                 getProjectTempDir: vi.fn().mockReturnValue('/test/temp'),
             },
+            getContentGenerator: vi.fn().mockReturnValue(mockContentGenerator),
         };
-        const MockedConfig = vi.mocked(Config, true);
-        MockedConfig.mockImplementation(() => mockConfigObject);
-        // We can instantiate the client here since Config is mocked
-        // and the constructor will use the mocked GoogleGenAI
-        client = new GeminiClient(new Config({ sessionId: 'test-session-id' }));
-        mockConfigObject.getGeminiClient.mockReturnValue(client);
-        await client.initialize(contentGeneratorConfig);
+        client = new GeminiClient(mockConfig);
+        await client.initialize();
+        vi.mocked(mockConfig.getGeminiClient).mockReturnValue(client);
     });
     afterEach(() => {
         vi.restoreAllMocks();
     });
-    // NOTE: The following tests for startChat were removed due to persistent issues with
-    // the @google/genai mock. Specifically, the mockChatCreateFn (representing instance.chats.create)
-    // was not being detected as called by the GeminiClient instance.
-    // This likely points to a subtle issue in how the GoogleGenerativeAI class constructor
-    // and its instance methods are mocked and then used by the class under test.
-    // For future debugging, ensure that the `this.client` in `GeminiClient` (which is an
-    // instance of the mocked GoogleGenerativeAI) correctly has its `chats.create` method
-    // pointing to `mockChatCreateFn`.
-    // it('startChat should call getCoreSystemPrompt with userMemory and pass to chats.create', async () => { ... });
-    // it('startChat should call getCoreSystemPrompt with empty string if userMemory is empty', async () => { ... });
-    // NOTE: The following tests for generateJson were removed due to persistent issues with
-    // the @google/genai mock, similar to the startChat tests. The mockGenerateContentFn
-    // (representing instance.models.generateContent) was not being detected as called, or the mock
-    // was not preventing an actual API call (leading to API key errors).
-    // For future debugging, ensure `this.client.models.generateContent` in `GeminiClient` correctly
-    // uses the `mockGenerateContentFn`.
-    // it('generateJson should call getCoreSystemPrompt with userMemory and pass to generateContent', async () => { ... });
-    // it('generateJson should call getCoreSystemPrompt with empty string if userMemory is empty', async () => { ... });
     describe('generateEmbedding', () => {
         const texts = ['hello world', 'goodbye world'];
         const testEmbeddingModel = 'test-embedding-model';
@@ -283,16 +248,15 @@ describe('Gemini Client (client.ts)', () => {
                 [0.1, 0.2, 0.3],
                 [0.4, 0.5, 0.6],
             ];
-            const mockResponse = {
+            vi.mocked(mockContentGenerator.embedContent).mockResolvedValue({
                 embeddings: [
                     { values: mockEmbeddings[0] },
                     { values: mockEmbeddings[1] },
                 ],
-            };
-            mockEmbedContentFn.mockResolvedValue(mockResponse);
+            });
             const result = await client.generateEmbedding(texts);
-            expect(mockEmbedContentFn).toHaveBeenCalledTimes(1);
-            expect(mockEmbedContentFn).toHaveBeenCalledWith({
+            expect(mockContentGenerator.embedContent).toHaveBeenCalledTimes(1);
+            expect(mockContentGenerator.embedContent).toHaveBeenCalledWith({
                 model: testEmbeddingModel,
                 contents: texts,
             });
@@ -301,43 +265,38 @@ describe('Gemini Client (client.ts)', () => {
         it('should return an empty array if an empty array is passed', async () => {
             const result = await client.generateEmbedding([]);
             expect(result).toEqual([]);
-            expect(mockEmbedContentFn).not.toHaveBeenCalled();
+            expect(mockContentGenerator.embedContent).not.toHaveBeenCalled();
         });
         it('should throw an error if API response has no embeddings array', async () => {
-            mockEmbedContentFn.mockResolvedValue({}); // No `embeddings` key
+            vi.mocked(mockContentGenerator.embedContent).mockResolvedValue({});
             await expect(client.generateEmbedding(texts)).rejects.toThrow('No embeddings found in API response.');
         });
         it('should throw an error if API response has an empty embeddings array', async () => {
-            const mockResponse = {
+            vi.mocked(mockContentGenerator.embedContent).mockResolvedValue({
                 embeddings: [],
-            };
-            mockEmbedContentFn.mockResolvedValue(mockResponse);
+            });
             await expect(client.generateEmbedding(texts)).rejects.toThrow('No embeddings found in API response.');
         });
         it('should throw an error if API returns a mismatched number of embeddings', async () => {
-            const mockResponse = {
+            vi.mocked(mockContentGenerator.embedContent).mockResolvedValue({
                 embeddings: [{ values: [1, 2, 3] }], // Only one for two texts
-            };
-            mockEmbedContentFn.mockResolvedValue(mockResponse);
+            });
             await expect(client.generateEmbedding(texts)).rejects.toThrow('API returned a mismatched number of embeddings. Expected 2, got 1.');
         });
         it('should throw an error if any embedding has nullish values', async () => {
-            const mockResponse = {
+            vi.mocked(mockContentGenerator.embedContent).mockResolvedValue({
                 embeddings: [{ values: [1, 2, 3] }, { values: undefined }], // Second one is bad
-            };
-            mockEmbedContentFn.mockResolvedValue(mockResponse);
+            });
             await expect(client.generateEmbedding(texts)).rejects.toThrow('API returned an empty embedding for input text at index 1: "goodbye world"');
         });
         it('should throw an error if any embedding has an empty values array', async () => {
-            const mockResponse = {
+            vi.mocked(mockContentGenerator.embedContent).mockResolvedValue({
                 embeddings: [{ values: [] }, { values: [1, 2, 3] }], // First one is bad
-            };
-            mockEmbedContentFn.mockResolvedValue(mockResponse);
+            });
             await expect(client.generateEmbedding(texts)).rejects.toThrow('API returned an empty embedding for input text at index 0: "hello world"');
         });
         it('should propagate errors from the API call', async () => {
-            const apiError = new Error('API Failure');
-            mockEmbedContentFn.mockRejectedValue(apiError);
+            vi.mocked(mockContentGenerator.embedContent).mockRejectedValue(new Error('API Failure'));
             await expect(client.generateEmbedding(texts)).rejects.toThrow('API Failure');
         });
     });
@@ -346,14 +305,11 @@ describe('Gemini Client (client.ts)', () => {
             const contents = [{ role: 'user', parts: [{ text: 'hello' }] }];
             const schema = { type: 'string' };
             const abortSignal = new AbortController().signal;
-            // Mock countTokens
-            const mockGenerator = {
-                countTokens: vi.fn().mockResolvedValue({ totalTokens: 1 }),
-                generateContent: mockGenerateContentFn,
-            };
-            client['contentGenerator'] = mockGenerator;
+            vi.mocked(mockContentGenerator.countTokens).mockResolvedValue({
+                totalTokens: 1,
+            });
             await client.generateJson(contents, schema, abortSignal, DEFAULT_GEMINI_FLASH_MODEL);
-            expect(mockGenerateContentFn).toHaveBeenCalledWith({
+            expect(mockContentGenerator.generateContent).toHaveBeenCalledWith({
                 model: DEFAULT_GEMINI_FLASH_MODEL,
                 config: {
                     abortSignal,
@@ -374,13 +330,11 @@ describe('Gemini Client (client.ts)', () => {
             const abortSignal = new AbortController().signal;
             const customModel = 'custom-json-model';
             const customConfig = { temperature: 0.9, topK: 20 };
-            const mockGenerator = {
-                countTokens: vi.fn().mockResolvedValue({ totalTokens: 1 }),
-                generateContent: mockGenerateContentFn,
-            };
-            client['contentGenerator'] = mockGenerator;
+            vi.mocked(mockContentGenerator.countTokens).mockResolvedValue({
+                totalTokens: 1,
+            });
             await client.generateJson(contents, schema, abortSignal, customModel, customConfig);
-            expect(mockGenerateContentFn).toHaveBeenCalledWith({
+            expect(mockContentGenerator.generateContent).toHaveBeenCalledWith({
                 model: customModel,
                 config: {
                     abortSignal,
@@ -394,6 +348,20 @@ describe('Gemini Client (client.ts)', () => {
                 contents,
             }, 'test-session-id');
         });
+        it('should use the Flash model when fallback mode is active', async () => {
+            const contents = [{ role: 'user', parts: [{ text: 'hello' }] }];
+            const schema = { type: 'string' };
+            const abortSignal = new AbortController().signal;
+            const requestedModel = 'gemini-2.5-pro'; // A non-flash model
+            // Mock config to be in fallback mode
+            // We access the mock via the client instance which holds the mocked config
+            vi.spyOn(client['config'], 'isInFallbackMode').mockReturnValue(true);
+            await client.generateJson(contents, schema, abortSignal, requestedModel);
+            // Assert that the Flash model was used, not the requested model
+            expect(mockContentGenerator.generateContent).toHaveBeenCalledWith(expect.objectContaining({
+                model: DEFAULT_GEMINI_FLASH_MODEL,
+            }), 'test-session-id');
+        });
     });
     describe('addHistory', () => {
         it('should call chat.addHistory with the provided content', async () => {
@@ -432,16 +400,12 @@ describe('Gemini Client (client.ts)', () => {
         });
     });
     describe('tryCompressChat', () => {
-        const mockCountTokens = vi.fn();
         const mockSendMessage = vi.fn();
         const mockGetHistory = vi.fn();
         beforeEach(() => {
             vi.mock('./tokenLimits', () => ({
                 tokenLimit: vi.fn(),
             }));
-            client['contentGenerator'] = {
-                countTokens: mockCountTokens,
-            };
             client['chat'] = {
                 getHistory: mockGetHistory,
                 addHistory: vi.fn(),
@@ -458,23 +422,17 @@ describe('Gemini Client (client.ts)', () => {
                 setHistory: vi.fn(),
                 sendMessage: vi.fn().mockResolvedValue({ text: 'Summary' }),
             };
-            const mockCountTokens = vi
-                .fn()
+            vi.mocked(mockContentGenerator.countTokens)
                 .mockResolvedValueOnce({ totalTokens: 1000 })
                 .mockResolvedValueOnce({ totalTokens: 5000 });
-            const mockGenerator = {
-                countTokens: mockCountTokens,
-            };
             client['chat'] = mockChat;
-            client['contentGenerator'] = mockGenerator;
             client['startChat'] = vi.fn().mockResolvedValue({ ...mockChat });
-            return { client, mockChat, mockGenerator };
+            return { client, mockChat };
         }
         describe('when compression inflates the token count', () => {
-            it('uses the truncated history for compression');
             it('allows compression to be forced/manual after a failure', async () => {
-                const { client, mockGenerator } = setup();
-                mockGenerator.countTokens?.mockResolvedValue({
+                const { client } = setup();
+                vi.mocked(mockContentGenerator.countTokens).mockResolvedValue({
                     totalTokens: 1000,
                 });
                 await client.tryCompressChat('prompt-id-4'); // Fails
@@ -487,6 +445,9 @@ describe('Gemini Client (client.ts)', () => {
             });
             it('yields the result even if the compression inflated the tokens', async () => {
                 const { client } = setup();
+                vi.mocked(mockContentGenerator.countTokens).mockResolvedValue({
+                    totalTokens: 1000,
+                });
                 const result = await client.tryCompressChat('prompt-id-4', true);
                 expect(result).toEqual({
                     compressionStatus: CompressionStatus.COMPRESSION_FAILED_INFLATED_TOKEN_COUNT,
@@ -501,7 +462,7 @@ describe('Gemini Client (client.ts)', () => {
             });
             it('restores the history back to the original', async () => {
                 vi.mocked(tokenLimit).mockReturnValue(1000);
-                mockCountTokens.mockResolvedValue({
+                vi.mocked(mockContentGenerator.countTokens).mockResolvedValue({
                     totalTokens: 999,
                 });
                 const originalHistory = [
@@ -517,11 +478,11 @@ describe('Gemini Client (client.ts)', () => {
                 expect(client['chat']?.setHistory).toHaveBeenCalledWith(originalHistory);
             });
             it('will not attempt to compress context after a failure', async () => {
-                const { client, mockGenerator } = setup();
+                const { client } = setup();
                 await client.tryCompressChat('prompt-id-4');
                 const result = await client.tryCompressChat('prompt-id-5');
                 // it counts tokens for {original, compressed} and then never again
-                expect(mockGenerator.countTokens).toHaveBeenCalledTimes(2);
+                expect(mockContentGenerator.countTokens).toHaveBeenCalledTimes(2);
                 expect(result).toEqual({
                     compressionStatus: CompressionStatus.NOOP,
                     newTokenCount: 0,
@@ -529,33 +490,13 @@ describe('Gemini Client (client.ts)', () => {
                 });
             });
         });
-        it('attempts to compress with a maxOutputTokens set to the original token count', async () => {
-            vi.mocked(tokenLimit).mockReturnValue(1000);
-            mockCountTokens.mockResolvedValue({
-                totalTokens: 999,
-            });
-            mockGetHistory.mockReturnValue([
-                { role: 'user', parts: [{ text: '...history...' }] },
-            ]);
-            // Mock the summary response from the chat
-            mockSendMessage.mockResolvedValue({
-                role: 'model',
-                parts: [{ text: 'This is a summary.' }],
-            });
-            await client.tryCompressChat('prompt-id-2', true);
-            expect(mockSendMessage).toHaveBeenCalledWith(expect.objectContaining({
-                config: expect.objectContaining({
-                    maxOutputTokens: 999,
-                }),
-            }), 'prompt-id-2');
-        });
         it('should not trigger summarization if token count is below threshold', async () => {
             const MOCKED_TOKEN_LIMIT = 1000;
             vi.mocked(tokenLimit).mockReturnValue(MOCKED_TOKEN_LIMIT);
             mockGetHistory.mockReturnValue([
                 { role: 'user', parts: [{ text: '...history...' }] },
             ]);
-            mockCountTokens.mockResolvedValue({
+            vi.mocked(mockContentGenerator.countTokens).mockResolvedValue({
                 totalTokens: MOCKED_TOKEN_LIMIT * 0.699, // TOKEN_THRESHOLD_FOR_SUMMARIZATION = 0.7
             });
             const initialChat = client.getChat();
@@ -582,7 +523,7 @@ describe('Gemini Client (client.ts)', () => {
             ]);
             const originalTokenCount = MOCKED_TOKEN_LIMIT * MOCKED_CONTEXT_PERCENTAGE_THRESHOLD;
             const newTokenCount = 100;
-            mockCountTokens
+            vi.mocked(mockContentGenerator.countTokens)
                 .mockResolvedValueOnce({ totalTokens: originalTokenCount }) // First call for the check
                 .mockResolvedValueOnce({ totalTokens: newTokenCount }); // Second call for the new history
             // Mock the summary response from the chat
@@ -608,7 +549,7 @@ describe('Gemini Client (client.ts)', () => {
             ]);
             const originalTokenCount = MOCKED_TOKEN_LIMIT * MOCKED_CONTEXT_PERCENTAGE_THRESHOLD;
             const newTokenCount = 100;
-            mockCountTokens
+            vi.mocked(mockContentGenerator.countTokens)
                 .mockResolvedValueOnce({ totalTokens: originalTokenCount }) // First call for the check
                 .mockResolvedValueOnce({ totalTokens: newTokenCount }); // Second call for the new history
             // Mock the summary response from the chat
@@ -653,7 +594,7 @@ describe('Gemini Client (client.ts)', () => {
             ]);
             const originalTokenCount = 1000 * 0.7;
             const newTokenCount = 100;
-            mockCountTokens
+            vi.mocked(mockContentGenerator.countTokens)
                 .mockResolvedValueOnce({ totalTokens: originalTokenCount }) // First call for the check
                 .mockResolvedValueOnce({ totalTokens: newTokenCount }); // Second call for the new history
             // Mock the summary response from the chat
@@ -687,7 +628,7 @@ describe('Gemini Client (client.ts)', () => {
             ]);
             const originalTokenCount = 10; // Well below threshold
             const newTokenCount = 5;
-            mockCountTokens
+            vi.mocked(mockContentGenerator.countTokens)
                 .mockResolvedValueOnce({ totalTokens: originalTokenCount })
                 .mockResolvedValueOnce({ totalTokens: newTokenCount });
             // Mock the summary response from the chat
@@ -708,9 +649,14 @@ describe('Gemini Client (client.ts)', () => {
             expect(newChat).not.toBe(initialChat);
         });
         it('should use current model from config for token counting after sendMessage', async () => {
-            const initialModel = client['config'].getModel();
-            const mockCountTokens = vi
-                .fn()
+            const initialModel = mockConfig.getModel();
+            // mock the model has been changed between calls of `countTokens`
+            const firstCurrentModel = initialModel + '-changed-1';
+            const secondCurrentModel = initialModel + '-changed-2';
+            vi.mocked(mockConfig.getModel)
+                .mockReturnValueOnce(firstCurrentModel)
+                .mockReturnValueOnce(secondCurrentModel);
+            vi.mocked(mockContentGenerator.countTokens)
                 .mockResolvedValueOnce({ totalTokens: 100000 })
                 .mockResolvedValueOnce({ totalTokens: 5000 });
             const mockSendMessage = vi.fn().mockResolvedValue({ text: 'Summary' });
@@ -723,25 +669,15 @@ describe('Gemini Client (client.ts)', () => {
                 setHistory: vi.fn(),
                 sendMessage: mockSendMessage,
             };
-            const mockGenerator = {
-                countTokens: mockCountTokens,
-            };
-            // mock the model has been changed between calls of `countTokens`
-            const firstCurrentModel = initialModel + '-changed-1';
-            const secondCurrentModel = initialModel + '-changed-2';
-            vi.spyOn(client['config'], 'getModel')
-                .mockReturnValueOnce(firstCurrentModel)
-                .mockReturnValueOnce(secondCurrentModel);
             client['chat'] = mockChat;
-            client['contentGenerator'] = mockGenerator;
             client['startChat'] = vi.fn().mockResolvedValue(mockChat);
             const result = await client.tryCompressChat('prompt-id-4', true);
-            expect(mockCountTokens).toHaveBeenCalledTimes(2);
-            expect(mockCountTokens).toHaveBeenNthCalledWith(1, {
+            expect(mockContentGenerator.countTokens).toHaveBeenCalledTimes(2);
+            expect(mockContentGenerator.countTokens).toHaveBeenNthCalledWith(1, {
                 model: firstCurrentModel,
                 contents: mockChatHistory,
             });
-            expect(mockCountTokens).toHaveBeenNthCalledWith(2, {
+            expect(mockContentGenerator.countTokens).toHaveBeenNthCalledWith(2, {
                 model: secondCurrentModel,
                 contents: expect.any(Array),
             });
@@ -755,20 +691,9 @@ describe('Gemini Client (client.ts)', () => {
     describe('sendMessageStream', () => {
         it('emits a compression event when the context was automatically compressed', async () => {
             // Arrange
-            const mockStream = (async function* () {
+            mockTurnRunFn.mockReturnValue((async function* () {
                 yield { type: 'content', value: 'Hello' };
-            })();
-            mockTurnRunFn.mockReturnValue(mockStream);
-            const mockChat = {
-                addHistory: vi.fn(),
-                getHistory: vi.fn().mockReturnValue([]),
-            };
-            client['chat'] = mockChat;
-            const mockGenerator = {
-                countTokens: vi.fn().mockResolvedValue({ totalTokens: 0 }),
-                generateContent: mockGenerateContentFn,
-            };
-            client['contentGenerator'] = mockGenerator;
+            })());
             const compressionInfo = {
                 compressionStatus: CompressionStatus.COMPRESSED,
                 originalTokenCount: 1000,
@@ -798,16 +723,6 @@ describe('Gemini Client (client.ts)', () => {
                 yield { type: 'content', value: 'Hello' };
             })();
             mockTurnRunFn.mockReturnValue(mockStream);
-            const mockChat = {
-                addHistory: vi.fn(),
-                getHistory: vi.fn().mockReturnValue([]),
-            };
-            client['chat'] = mockChat;
-            const mockGenerator = {
-                countTokens: vi.fn().mockResolvedValue({ totalTokens: 0 }),
-                generateContent: mockGenerateContentFn,
-            };
-            client['contentGenerator'] = mockGenerator;
             const compressionInfo = {
                 compressionStatus,
                 originalTokenCount: 1000,
@@ -846,21 +761,15 @@ describe('Gemini Client (client.ts)', () => {
                     ],
                 },
             });
-            vi.spyOn(client['config'], 'getIdeMode').mockReturnValue(true);
-            const mockStream = (async function* () {
+            vi.mocked(mockConfig.getIdeMode).mockReturnValue(true);
+            mockTurnRunFn.mockReturnValue((async function* () {
                 yield { type: 'content', value: 'Hello' };
-            })();
-            mockTurnRunFn.mockReturnValue(mockStream);
+            })());
             const mockChat = {
                 addHistory: vi.fn(),
                 getHistory: vi.fn().mockReturnValue([]),
             };
             client['chat'] = mockChat;
-            const mockGenerator = {
-                countTokens: vi.fn().mockResolvedValue({ totalTokens: 0 }),
-                generateContent: mockGenerateContentFn,
-            };
-            client['contentGenerator'] = mockGenerator;
             const initialRequest = [{ text: 'Hi' }];
             // Act
             const stream = client.sendMessageStream(initialRequest, new AbortController().signal, 'prompt-id-ide');
@@ -908,11 +817,6 @@ ${JSON.stringify({
                 getHistory: vi.fn().mockReturnValue([]),
             };
             client['chat'] = mockChat;
-            const mockGenerator = {
-                countTokens: vi.fn().mockResolvedValue({ totalTokens: 0 }),
-                generateContent: mockGenerateContentFn,
-            };
-            client['contentGenerator'] = mockGenerator;
             const initialRequest = [{ text: 'Hi' }];
             // Act
             const stream = client.sendMessageStream(initialRequest, new AbortController().signal, 'prompt-id-ide');
@@ -948,11 +852,6 @@ ${JSON.stringify({
                 getHistory: vi.fn().mockReturnValue([]),
             };
             client['chat'] = mockChat;
-            const mockGenerator = {
-                countTokens: vi.fn().mockResolvedValue({ totalTokens: 0 }),
-                generateContent: mockGenerateContentFn,
-            };
-            client['contentGenerator'] = mockGenerator;
             const initialRequest = [{ text: 'Hi' }];
             // Act
             const stream = client.sendMessageStream(initialRequest, new AbortController().signal, 'prompt-id-ide');
@@ -1008,11 +907,6 @@ ${JSON.stringify({
                 getHistory: vi.fn().mockReturnValue([]),
             };
             client['chat'] = mockChat;
-            const mockGenerator = {
-                countTokens: vi.fn().mockResolvedValue({ totalTokens: 0 }),
-                generateContent: mockGenerateContentFn,
-            };
-            client['contentGenerator'] = mockGenerator;
             const initialRequest = [{ text: 'Hi' }];
             // Act
             const stream = client.sendMessageStream(initialRequest, new AbortController().signal, 'prompt-id-ide');
@@ -1046,11 +940,6 @@ ${JSON.stringify({
                 getHistory: vi.fn().mockReturnValue([]),
             };
             client['chat'] = mockChat;
-            const mockGenerator = {
-                countTokens: vi.fn().mockResolvedValue({ totalTokens: 0 }),
-                generateContent: mockGenerateContentFn,
-            };
-            client['contentGenerator'] = mockGenerator;
             // Act
             const stream = client.sendMessageStream([{ text: 'Hi' }], new AbortController().signal, 'prompt-id-1');
             // Consume the stream manually to get the final return value.
@@ -1083,11 +972,6 @@ ${JSON.stringify({
                 getHistory: vi.fn().mockReturnValue([]),
             };
             client['chat'] = mockChat;
-            const mockGenerator = {
-                countTokens: vi.fn().mockResolvedValue({ totalTokens: 0 }),
-                generateContent: mockGenerateContentFn,
-            };
-            client['contentGenerator'] = mockGenerator;
             // Use a signal that never gets aborted
             const abortController = new AbortController();
             const signal = abortController.signal;
@@ -1150,11 +1034,6 @@ ${JSON.stringify({
                 getHistory: vi.fn().mockReturnValue([]),
             };
             client['chat'] = mockChat;
-            const mockGenerator = {
-                countTokens: vi.fn().mockResolvedValue({ totalTokens: 0 }),
-                generateContent: mockGenerateContentFn,
-            };
-            client['contentGenerator'] = mockGenerator;
             // Act & Assert
             // Run up to the limit
             for (let i = 0; i < MAX_SESSION_TURNS; i++) {
@@ -1193,11 +1072,6 @@ ${JSON.stringify({
                 getHistory: vi.fn().mockReturnValue([]),
             };
             client['chat'] = mockChat;
-            const mockGenerator = {
-                countTokens: vi.fn().mockResolvedValue({ totalTokens: 0 }),
-                generateContent: mockGenerateContentFn,
-            };
-            client['contentGenerator'] = mockGenerator;
             // Use a signal that never gets aborted
             const abortController = new AbortController();
             const signal = abortController.signal;
@@ -1261,11 +1135,6 @@ ${JSON.stringify({
                     ]),
                 };
                 client['chat'] = mockChat;
-                const mockGenerator = {
-                    countTokens: vi.fn().mockResolvedValue({ totalTokens: 0 }),
-                    generateContent: mockGenerateContentFn,
-                };
-                client['contentGenerator'] = mockGenerator;
             });
             const testCases = [
                 {
@@ -1475,10 +1344,6 @@ ${JSON.stringify({
                     sendMessage: vi.fn().mockResolvedValue({ text: 'summary' }),
                 };
                 client['chat'] = mockChat;
-                const mockGenerator = {
-                    countTokens: vi.fn().mockResolvedValue({ totalTokens: 0 }),
-                };
-                client['contentGenerator'] = mockGenerator;
                 vi.spyOn(client['config'], 'getIdeMode').mockReturnValue(true);
                 vi.mocked(ideContext.getIdeContext).mockReturnValue({
                     workspaceState: {
@@ -1742,11 +1607,6 @@ ${JSON.stringify({
                 getHistory: vi.fn().mockReturnValue([]),
             };
             client['chat'] = mockChat;
-            const mockGenerator = {
-                countTokens: vi.fn().mockResolvedValue({ totalTokens: 0 }),
-                generateContent: mockGenerateContentFn,
-            };
-            client['contentGenerator'] = mockGenerator;
             // Act
             const stream = client.sendMessageStream([{ text: 'Hi' }], new AbortController().signal, 'prompt-id-error');
             for await (const _ of stream) {
@@ -1772,11 +1632,6 @@ ${JSON.stringify({
                 getHistory: vi.fn().mockReturnValue([]),
             };
             client['chat'] = mockChat;
-            const mockGenerator = {
-                countTokens: vi.fn().mockResolvedValue({ totalTokens: 0 }),
-                generateContent: mockGenerateContentFn,
-            };
-            client['contentGenerator'] = mockGenerator;
             // Act
             const stream = client.sendMessageStream([{ text: 'Hi' }], new AbortController().signal, 'prompt-id-error');
             for await (const _ of stream) {
@@ -1791,14 +1646,8 @@ ${JSON.stringify({
             const contents = [{ role: 'user', parts: [{ text: 'hello' }] }];
             const generationConfig = { temperature: 0.5 };
             const abortSignal = new AbortController().signal;
-            // Mock countTokens
-            const mockGenerator = {
-                countTokens: vi.fn().mockResolvedValue({ totalTokens: 1 }),
-                generateContent: mockGenerateContentFn,
-            };
-            client['contentGenerator'] = mockGenerator;
             await client.generateContent(contents, generationConfig, abortSignal, DEFAULT_GEMINI_FLASH_MODEL);
-            expect(mockGenerateContentFn).toHaveBeenCalledWith({
+            expect(mockContentGenerator.generateContent).toHaveBeenCalledWith({
                 model: DEFAULT_GEMINI_FLASH_MODEL,
                 config: {
                     abortSignal,
@@ -1814,98 +1663,29 @@ ${JSON.stringify({
             const contents = [{ role: 'user', parts: [{ text: 'test' }] }];
             const currentModel = initialModel + '-changed';
             vi.spyOn(client['config'], 'getModel').mockReturnValueOnce(currentModel);
-            const mockGenerator = {
-                countTokens: vi.fn().mockResolvedValue({ totalTokens: 1 }),
-                generateContent: mockGenerateContentFn,
-            };
-            client['contentGenerator'] = mockGenerator;
             await client.generateContent(contents, {}, new AbortController().signal, DEFAULT_GEMINI_FLASH_MODEL);
-            expect(mockGenerateContentFn).not.toHaveBeenCalledWith({
+            expect(mockContentGenerator.generateContent).not.toHaveBeenCalledWith({
                 model: initialModel,
                 config: expect.any(Object),
                 contents,
             });
-            expect(mockGenerateContentFn).toHaveBeenCalledWith({
+            expect(mockContentGenerator.generateContent).toHaveBeenCalledWith({
                 model: DEFAULT_GEMINI_FLASH_MODEL,
                 config: expect.any(Object),
                 contents,
             }, 'test-session-id');
         });
-    });
-    describe('handleFlashFallback', () => {
-        it('should use current model from config when checking for fallback', async () => {
-            const initialModel = client['config'].getModel();
-            const fallbackModel = DEFAULT_GEMINI_FLASH_MODEL;
-            // mock config been changed
-            const currentModel = initialModel + '-changed';
-            const getModelSpy = vi.spyOn(client['config'], 'getModel');
-            getModelSpy.mockReturnValue(currentModel);
-            const mockFallbackHandler = vi.fn().mockResolvedValue(true);
-            client['config'].flashFallbackHandler = mockFallbackHandler;
-            client['config'].setModel = vi.fn();
-            const result = await client['handleFlashFallback'](AuthType.LOGIN_WITH_GOOGLE);
-            expect(result).toBe(fallbackModel);
-            expect(mockFallbackHandler).toHaveBeenCalledWith(currentModel, fallbackModel, undefined);
-        });
-    });
-    describe('setHistory', () => {
-        it('should strip thought signatures when stripThoughts is true', () => {
-            const mockChat = {
-                setHistory: vi.fn(),
-            };
-            client['chat'] = mockChat;
-            const historyWithThoughts = [
-                {
-                    role: 'user',
-                    parts: [{ text: 'hello' }],
-                },
-                {
-                    role: 'model',
-                    parts: [
-                        { text: 'thinking...', thoughtSignature: 'thought-123' },
-                        {
-                            functionCall: { name: 'test', args: {} },
-                            thoughtSignature: 'thought-456',
-                        },
-                    ],
-                },
-            ];
-            client.setHistory(historyWithThoughts, { stripThoughts: true });
-            const expectedHistory = [
-                {
-                    role: 'user',
-                    parts: [{ text: 'hello' }],
-                },
-                {
-                    role: 'model',
-                    parts: [
-                        { text: 'thinking...' },
-                        { functionCall: { name: 'test', args: {} } },
-                    ],
-                },
-            ];
-            expect(mockChat.setHistory).toHaveBeenCalledWith(expectedHistory);
-        });
-        it('should not strip thought signatures when stripThoughts is false', () => {
-            const mockChat = {
-                setHistory: vi.fn(),
-            };
-            client['chat'] = mockChat;
-            const historyWithThoughts = [
-                {
-                    role: 'user',
-                    parts: [{ text: 'hello' }],
-                },
-                {
-                    role: 'model',
-                    parts: [
-                        { text: 'thinking...', thoughtSignature: 'thought-123' },
-                        { text: 'ok', thoughtSignature: 'thought-456' },
-                    ],
-                },
-            ];
-            client.setHistory(historyWithThoughts, { stripThoughts: false });
-            expect(mockChat.setHistory).toHaveBeenCalledWith(historyWithThoughts);
+        it('should use the Flash model when fallback mode is active', async () => {
+            const contents = [{ role: 'user', parts: [{ text: 'hello' }] }];
+            const generationConfig = { temperature: 0.5 };
+            const abortSignal = new AbortController().signal;
+            const requestedModel = 'gemini-2.5-pro'; // A non-flash model
+            // Mock config to be in fallback mode
+            vi.spyOn(client['config'], 'isInFallbackMode').mockReturnValue(true);
+            await client.generateContent(contents, generationConfig, abortSignal, requestedModel);
+            expect(mockGenerateContentFn).toHaveBeenCalledWith(expect.objectContaining({
+                model: DEFAULT_GEMINI_FLASH_MODEL,
+            }), 'test-session-id');
         });
     });
 });