@machina.ai/cell-cli-core 1.0.21-rc4 → 1.4.0-rc1

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

@@ -3,12 +3,13 @@
  * Copyright 2025 Google LLC
  * SPDX-License-Identifier: Apache-2.0
  */
-import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
-import { GoogleGenAI, } from '@google/genai';
-import { findIndexAfterFraction, GeminiClient } from './client.js';
-import { AuthType } from './contentGenerator.js';
+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 { GeminiEventType, Turn } from './turn.js';
+import { CompressionStatus, GeminiEventType, Turn, } from './turn.js';
 import { getCoreSystemPrompt } from './prompts.js';
 import { DEFAULT_GEMINI_FLASH_MODEL } from '../config/models.js';
 import { FileDiscoveryService } from '../services/fileDiscoveryService.js';
@@ -16,13 +17,37 @@ import { setSimulate429 } from '../utils/testUtils.js';
 import { tokenLimit } from './tokenLimits.js';
 import { ideContext } from '../ide/ideContext.js';
 import { ClearcutLogger } from '../telemetry/clearcut-logger/clearcut-logger.js';
+// Mock fs module to prevent actual file system operations during tests
+const mockFileSystem = new Map();
+vi.mock('node:fs', () => {
+    const fsModule = {
+        mkdirSync: vi.fn(),
+        writeFileSync: vi.fn((path, data) => {
+            mockFileSystem.set(path, data);
+        }),
+        readFileSync: vi.fn((path) => {
+            if (mockFileSystem.has(path)) {
+                return mockFileSystem.get(path);
+            }
+            throw Object.assign(new Error('ENOENT: no such file or directory'), {
+                code: 'ENOENT',
+            });
+        }),
+        existsSync: vi.fn((path) => mockFileSystem.has(path)),
+    };
+    return {
+        default: fsModule,
+        ...fsModule,
+    };
+});
 // --- Mocks ---
 const mockChatCreateFn = vi.fn();
 const mockGenerateContentFn = vi.fn();
 const mockEmbedContentFn = vi.fn();
 const mockTurnRunFn = vi.fn();
 vi.mock('@google/genai');
-vi.mock('./turn', () => {
+vi.mock('./turn', async (importOriginal) => {
+    const actual = await importOriginal();
     // Define a mock class that has the same shape as the real Turn
     class MockTurn {
         pendingToolCalls = [];
@@ -34,11 +59,8 @@ vi.mock('./turn', () => {
     }
     // Export the mock class as 'Turn'
     return {
+        ...actual,
         Turn: MockTurn,
-        GeminiEventType: {
-            MaxSessionTurns: 'MaxSessionTurns',
-            ChatCompressed: 'ChatCompressed',
-        },
     };
 });
 vi.mock('../config/config.js');
@@ -60,6 +82,18 @@ vi.mock('../telemetry/index.js', () => ({
     logApiError: vi.fn(),
 }));
 vi.mock('../ide/ideContext.js');
+/**
+ * Array.fromAsync ponyfill, which will be available in es 2024.
+ *
+ * Buffers an async generator into an array and returns the result.
+ */
+async function fromAsync(promise) {
+    const results = [];
+    for await (const result of promise) {
+        results.push(result);
+    }
+    return results;
+}
 describe('findIndexAfterFraction', () => {
     const history = [
         { role: 'user', parts: [{ text: 'This is the first message.' }] }, // JSON length: 66
@@ -106,6 +140,33 @@ describe('findIndexAfterFraction', () => {
         expect(findIndexAfterFraction(historyWithEmptyParts, 0.5)).toBe(1);
     });
 });
+describe('isThinkingSupported', () => {
+    it('should return true for gemini-2.5', () => {
+        expect(isThinkingSupported('gemini-2.5')).toBe(true);
+    });
+    it('should return true for gemini-2.5-pro', () => {
+        expect(isThinkingSupported('gemini-2.5-pro')).toBe(true);
+    });
+    it('should return false for other models', () => {
+        expect(isThinkingSupported('gemini-1.5-flash')).toBe(false);
+        expect(isThinkingSupported('some-other-model')).toBe(false);
+    });
+});
+describe('isThinkingDefault', () => {
+    it('should return false for gemini-2.5-flash-lite', () => {
+        expect(isThinkingDefault('gemini-2.5-flash-lite')).toBe(false);
+    });
+    it('should return true for gemini-2.5', () => {
+        expect(isThinkingDefault('gemini-2.5')).toBe(true);
+    });
+    it('should return true for gemini-2.5-pro', () => {
+        expect(isThinkingDefault('gemini-2.5-pro')).toBe(true);
+    });
+    it('should return false for other models', () => {
+        expect(isThinkingDefault('gemini-1.5-flash')).toBe(false);
+        expect(isThinkingDefault('some-other-model')).toBe(false);
+    });
+});
 describe('Gemini Client (client.ts)', () => {
     let client;
     beforeEach(async () => {
@@ -152,7 +213,7 @@ describe('Gemini Client (client.ts)', () => {
             getContentGeneratorConfig: vi
                 .fn()
                 .mockReturnValue(contentGeneratorConfig),
-            getToolRegistry: vi.fn().mockResolvedValue(mockToolRegistry),
+            getToolRegistry: vi.fn().mockReturnValue(mockToolRegistry),
             getModel: vi.fn().mockReturnValue('test-model'),
             getEmbeddingModel: vi.fn().mockReturnValue('test-embedding-model'),
             getApiKey: vi.fn().mockReturnValue('test-key'),
@@ -178,6 +239,12 @@ describe('Gemini Client (client.ts)', () => {
             getGeminiClient: vi.fn(),
             setFallbackMode: vi.fn(),
             getChatCompression: vi.fn().mockReturnValue(undefined),
+            getSkipNextSpeakerCheck: vi.fn().mockReturnValue(false),
+            getUseSmartEdit: vi.fn().mockReturnValue(false),
+            getProjectRoot: vi.fn().mockReturnValue('/test/project/root'),
+            storage: {
+                getProjectTempDir: vi.fn().mockReturnValue('/test/temp'),
+            },
         };
         const MockedConfig = vi.mocked(Config, true);
         MockedConfig.mockImplementation(() => mockConfigObject);
@@ -274,30 +341,6 @@ describe('Gemini Client (client.ts)', () => {
             await expect(client.generateEmbedding(texts)).rejects.toThrow('API Failure');
         });
     });
-    describe('generateContent', () => {
-        it('should call generateContent with the correct parameters', async () => {
-            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);
-            expect(mockGenerateContentFn).toHaveBeenCalledWith({
-                model: 'test-model',
-                config: {
-                    abortSignal,
-                    systemInstruction: getCoreSystemPrompt(''),
-                    temperature: 0.5,
-                    topP: 1,
-                },
-                contents,
-            }, 'test-session-id');
-        });
-    });
     describe('generateJson', () => {
         it('should call generateContent with the correct parameters', async () => {
             const contents = [{ role: 'user', parts: [{ text: 'hello' }] }];
@@ -309,9 +352,9 @@ describe('Gemini Client (client.ts)', () => {
                 generateContent: mockGenerateContentFn,
             };
             client['contentGenerator'] = mockGenerator;
-            await client.generateJson(contents, schema, abortSignal);
+            await client.generateJson(contents, schema, abortSignal, DEFAULT_GEMINI_FLASH_MODEL);
             expect(mockGenerateContentFn).toHaveBeenCalledWith({
-                model: 'test-model', // Should use current model from config
+                model: DEFAULT_GEMINI_FLASH_MODEL,
                 config: {
                     abortSignal,
                     systemInstruction: getCoreSystemPrompt(''),
@@ -324,7 +367,9 @@ describe('Gemini Client (client.ts)', () => {
             }, 'test-session-id');
         });
         it('should allow overriding model and config', async () => {
-            const contents = [{ role: 'user', parts: [{ text: 'hello' }] }];
+            const contents = [
+                { role: 'user', parts: [{ text: 'hello' }] },
+            ];
             const schema = { type: 'string' };
             const abortSignal = new AbortController().signal;
             const customModel = 'custom-json-model';
@@ -404,6 +449,106 @@ describe('Gemini Client (client.ts)', () => {
                 sendMessage: mockSendMessage,
             };
         });
+        function setup({ chatHistory = [
+            { role: 'user', parts: [{ text: 'Long conversation' }] },
+            { role: 'model', parts: [{ text: 'Long response' }] },
+        ], } = {}) {
+            const mockChat = {
+                getHistory: vi.fn().mockReturnValue(chatHistory),
+                setHistory: vi.fn(),
+                sendMessage: vi.fn().mockResolvedValue({ text: 'Summary' }),
+            };
+            const mockCountTokens = vi
+                .fn()
+                .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 };
+        }
+        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({
+                    totalTokens: 1000,
+                });
+                await client.tryCompressChat('prompt-id-4'); // Fails
+                const result = await client.tryCompressChat('prompt-id-4', true);
+                expect(result).toEqual({
+                    compressionStatus: CompressionStatus.COMPRESSED,
+                    newTokenCount: 1000,
+                    originalTokenCount: 1000,
+                });
+            });
+            it('yields the result even if the compression inflated the tokens', async () => {
+                const { client } = setup();
+                const result = await client.tryCompressChat('prompt-id-4', true);
+                expect(result).toEqual({
+                    compressionStatus: CompressionStatus.COMPRESSION_FAILED_INFLATED_TOKEN_COUNT,
+                    newTokenCount: 5000,
+                    originalTokenCount: 1000,
+                });
+            });
+            it('does not manipulate the source chat', async () => {
+                const { client, mockChat } = setup();
+                await client.tryCompressChat('prompt-id-4', true);
+                expect(client['chat']).toBe(mockChat); // a new chat session was not created
+            });
+            it('restores the history back to the original', async () => {
+                vi.mocked(tokenLimit).mockReturnValue(1000);
+                mockCountTokens.mockResolvedValue({
+                    totalTokens: 999,
+                });
+                const originalHistory = [
+                    { role: 'user', parts: [{ text: 'what is your wisdom?' }] },
+                    { role: 'model', parts: [{ text: 'some wisdom' }] },
+                    { role: 'user', parts: [{ text: 'ahh that is a good a wisdom' }] },
+                ];
+                const { client } = setup({
+                    chatHistory: originalHistory,
+                });
+                const { compressionStatus } = await client.tryCompressChat('prompt-id-4');
+                expect(compressionStatus).toBe(CompressionStatus.COMPRESSION_FAILED_INFLATED_TOKEN_COUNT);
+                expect(client['chat']?.setHistory).toHaveBeenCalledWith(originalHistory);
+            });
+            it('will not attempt to compress context after a failure', async () => {
+                const { client, mockGenerator } = 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(result).toEqual({
+                    compressionStatus: CompressionStatus.NOOP,
+                    newTokenCount: 0,
+                    originalTokenCount: 0,
+                });
+            });
+        });
+        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);
@@ -417,7 +562,11 @@ describe('Gemini Client (client.ts)', () => {
             const result = await client.tryCompressChat('prompt-id-2');
             const newChat = client.getChat();
             expect(tokenLimit).toHaveBeenCalled();
-            expect(result).toBeNull();
+            expect(result).toEqual({
+                compressionStatus: CompressionStatus.NOOP,
+                newTokenCount: 699,
+                originalTokenCount: 699,
+            });
             expect(newChat).toBe(initialChat);
         });
         it('logs a telemetry event when compressing', async () => {
@@ -474,6 +623,7 @@ describe('Gemini Client (client.ts)', () => {
             expect(mockSendMessage).toHaveBeenCalled();
             // Assert that summarization happened and returned the correct stats
             expect(result).toEqual({
+                compressionStatus: CompressionStatus.COMPRESSED,
                 originalTokenCount,
                 newTokenCount,
             });
@@ -518,6 +668,7 @@ describe('Gemini Client (client.ts)', () => {
             expect(mockSendMessage).toHaveBeenCalled();
             // Assert that summarization happened and returned the correct stats
             expect(result).toEqual({
+                compressionStatus: CompressionStatus.COMPRESSED,
                 originalTokenCount,
                 newTokenCount,
             });
@@ -549,14 +700,129 @@ describe('Gemini Client (client.ts)', () => {
             const newChat = client.getChat();
             expect(mockSendMessage).toHaveBeenCalled();
             expect(result).toEqual({
+                compressionStatus: CompressionStatus.COMPRESSED,
                 originalTokenCount,
                 newTokenCount,
             });
             // Assert that the chat was reset
             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()
+                .mockResolvedValueOnce({ totalTokens: 100000 })
+                .mockResolvedValueOnce({ totalTokens: 5000 });
+            const mockSendMessage = vi.fn().mockResolvedValue({ text: 'Summary' });
+            const mockChatHistory = [
+                { role: 'user', parts: [{ text: 'Long conversation' }] },
+                { role: 'model', parts: [{ text: 'Long response' }] },
+            ];
+            const mockChat = {
+                getHistory: vi.fn().mockReturnValue(mockChatHistory),
+                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, {
+                model: firstCurrentModel,
+                contents: mockChatHistory,
+            });
+            expect(mockCountTokens).toHaveBeenNthCalledWith(2, {
+                model: secondCurrentModel,
+                contents: expect.any(Array),
+            });
+            expect(result).toEqual({
+                compressionStatus: CompressionStatus.COMPRESSED,
+                originalTokenCount: 100000,
+                newTokenCount: 5000,
+            });
+        });
     });
     describe('sendMessageStream', () => {
+        it('emits a compression event when the context was automatically compressed', async () => {
+            // Arrange
+            const mockStream = (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,
+                newTokenCount: 500,
+            };
+            vi.spyOn(client, 'tryCompressChat').mockResolvedValueOnce(compressionInfo);
+            // Act
+            const stream = client.sendMessageStream([{ text: 'Hi' }], new AbortController().signal, 'prompt-id-1');
+            const events = await fromAsync(stream);
+            // Assert
+            expect(events).toContainEqual({
+                type: GeminiEventType.ChatCompressed,
+                value: compressionInfo,
+            });
+        });
+        it.each([
+            {
+                compressionStatus: CompressionStatus.COMPRESSION_FAILED_INFLATED_TOKEN_COUNT,
+            },
+            { compressionStatus: CompressionStatus.NOOP },
+            {
+                compressionStatus: CompressionStatus.COMPRESSION_FAILED_TOKEN_COUNT_ERROR,
+            },
+        ])('does not emit a compression event when the status is $compressionStatus', async ({ compressionStatus }) => {
+            // Arrange
+            const mockStream = (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,
+                originalTokenCount: 1000,
+                newTokenCount: 500,
+            };
+            vi.spyOn(client, 'tryCompressChat').mockResolvedValueOnce(compressionInfo);
+            // Act
+            const stream = client.sendMessageStream([{ text: 'Hi' }], new AbortController().signal, 'prompt-id-1');
+            const events = await fromAsync(stream);
+            // Assert
+            expect(events).not.toContainEqual({
+                type: GeminiEventType.ChatCompressed,
+                value: expect.anything(),
+            });
+        });
         it('should include editor context when ideMode is enabled', async () => {
             // Arrange
             vi.mocked(ideContext.getIdeContext).mockReturnValue({
@@ -976,7 +1242,11 @@ ${JSON.stringify({
             })();
             beforeEach(() => {
                 client['forceFullIdeContext'] = false; // Reset before each delta test
-                vi.spyOn(client, 'tryCompressChat').mockResolvedValue(null);
+                vi.spyOn(client, 'tryCompressChat').mockResolvedValue({
+                    originalTokenCount: 0,
+                    newTokenCount: 0,
+                    compressionStatus: CompressionStatus.COMPRESSED,
+                });
                 vi.spyOn(client['config'], 'getIdeMode').mockReturnValue(true);
                 mockTurnRunFn.mockReturnValue(mockStream);
                 const mockChat = {
@@ -1189,7 +1459,11 @@ ${JSON.stringify({
         describe('IDE context with pending tool calls', () => {
             let mockChat;
             beforeEach(() => {
-                vi.spyOn(client, 'tryCompressChat').mockResolvedValue(null);
+                vi.spyOn(client, 'tryCompressChat').mockResolvedValue({
+                    originalTokenCount: 0,
+                    newTokenCount: 0,
+                    compressionStatus: CompressionStatus.COMPRESSED,
+                });
                 const mockStream = (async function* () {
                     yield { type: 'content', value: 'response' };
                 })();
@@ -1452,8 +1726,89 @@ ${JSON.stringify({
                 expect(JSON.stringify(finalCall)).toContain('fileC.ts');
             });
         });
+        it('should not call checkNextSpeaker when turn.run() yields an error', async () => {
+            // Arrange
+            const { checkNextSpeaker } = await import('../utils/nextSpeakerChecker.js');
+            const mockCheckNextSpeaker = vi.mocked(checkNextSpeaker);
+            const mockStream = (async function* () {
+                yield {
+                    type: GeminiEventType.Error,
+                    value: { error: { message: 'test error' } },
+                };
+            })();
+            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;
+            // Act
+            const stream = client.sendMessageStream([{ text: 'Hi' }], new AbortController().signal, 'prompt-id-error');
+            for await (const _ of stream) {
+                // consume stream
+            }
+            // Assert
+            expect(mockCheckNextSpeaker).not.toHaveBeenCalled();
+        });
+        it('should not call checkNextSpeaker when turn.run() yields a value then an error', async () => {
+            // Arrange
+            const { checkNextSpeaker } = await import('../utils/nextSpeakerChecker.js');
+            const mockCheckNextSpeaker = vi.mocked(checkNextSpeaker);
+            const mockStream = (async function* () {
+                yield { type: GeminiEventType.Content, value: 'some content' };
+                yield {
+                    type: GeminiEventType.Error,
+                    value: { error: { message: 'test error' } },
+                };
+            })();
+            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;
+            // Act
+            const stream = client.sendMessageStream([{ text: 'Hi' }], new AbortController().signal, 'prompt-id-error');
+            for await (const _ of stream) {
+                // consume stream
+            }
+            // Assert
+            expect(mockCheckNextSpeaker).not.toHaveBeenCalled();
+        });
     });
     describe('generateContent', () => {
+        it('should call generateContent with the correct parameters', async () => {
+            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({
+                model: DEFAULT_GEMINI_FLASH_MODEL,
+                config: {
+                    abortSignal,
+                    systemInstruction: getCoreSystemPrompt(''),
+                    temperature: 0.5,
+                    topP: 1,
+                },
+                contents,
+            }, 'test-session-id');
+        });
         it('should use current model from config for content generation', async () => {
             const initialModel = client['config'].getModel();
             const contents = [{ role: 'user', parts: [{ text: 'test' }] }];
@@ -1464,64 +1819,19 @@ ${JSON.stringify({
                 generateContent: mockGenerateContentFn,
             };
             client['contentGenerator'] = mockGenerator;
-            await client.generateContent(contents, {}, new AbortController().signal);
+            await client.generateContent(contents, {}, new AbortController().signal, DEFAULT_GEMINI_FLASH_MODEL);
             expect(mockGenerateContentFn).not.toHaveBeenCalledWith({
                 model: initialModel,
                 config: expect.any(Object),
                 contents,
             });
             expect(mockGenerateContentFn).toHaveBeenCalledWith({
-                model: currentModel,
+                model: DEFAULT_GEMINI_FLASH_MODEL,
                 config: expect.any(Object),
                 contents,
             }, 'test-session-id');
         });
     });
-    describe('tryCompressChat', () => {
-        it('should use current model from config for token counting after sendMessage', async () => {
-            const initialModel = client['config'].getModel();
-            const mockCountTokens = vi
-                .fn()
-                .mockResolvedValueOnce({ totalTokens: 100000 })
-                .mockResolvedValueOnce({ totalTokens: 5000 });
-            const mockSendMessage = vi.fn().mockResolvedValue({ text: 'Summary' });
-            const mockChatHistory = [
-                { role: 'user', parts: [{ text: 'Long conversation' }] },
-                { role: 'model', parts: [{ text: 'Long response' }] },
-            ];
-            const mockChat = {
-                getHistory: vi.fn().mockReturnValue(mockChatHistory),
-                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, {
-                model: firstCurrentModel,
-                contents: mockChatHistory,
-            });
-            expect(mockCountTokens).toHaveBeenNthCalledWith(2, {
-                model: secondCurrentModel,
-                contents: expect.any(Array),
-            });
-            expect(result).toEqual({
-                originalTokenCount: 100000,
-                newTokenCount: 5000,
-            });
-        });
-    });
     describe('handleFlashFallback', () => {
         it('should use current model from config when checking for fallback', async () => {
             const initialModel = client['config'].getModel();