@google/gemini-cli-core 0.5.0-preview.1 → 0.7.0-nightly.20250912.68035591

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

@@ -13,7 +13,7 @@ import { DEFAULT_GEMINI_FLASH_MODEL } from '../config/models.js';
 import { FileDiscoveryService } from '../services/fileDiscoveryService.js';
 import { setSimulate429 } from '../utils/testUtils.js';
 import { tokenLimit } from './tokenLimits.js';
-import { ideContext } from '../ide/ideContext.js';
+import { ideContextStore } 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();
@@ -108,22 +108,22 @@ describe('findIndexAfterFraction', () => {
         // 0: 66
         // 1: 66 + 68 = 134
         // 2: 134 + 66 = 200
-        // 200 >= 166.5, so index is 2
-        expect(findIndexAfterFraction(history, 0.5)).toBe(2);
+        // 200 >= 166.5, so index is 3
+        expect(findIndexAfterFraction(history, 0.5)).toBe(3);
     });
     it('should handle a fraction that results in the last index', () => {
         // 333 * 0.9 = 299.7
         // ...
         // 3: 200 + 68 = 268
         // 4: 268 + 65 = 333
-        // 333 >= 299.7, so index is 4
-        expect(findIndexAfterFraction(history, 0.9)).toBe(4);
+        // 333 >= 299.7, so index is 5
+        expect(findIndexAfterFraction(history, 0.9)).toBe(5);
     });
     it('should handle an empty history', () => {
         expect(findIndexAfterFraction([], 0.5)).toBe(0);
     });
     it('should handle a history with only one item', () => {
-        expect(findIndexAfterFraction(history.slice(0, 1), 0.5)).toBe(0);
+        expect(findIndexAfterFraction(history.slice(0, 1), 0.5)).toBe(1);
     });
     it('should handle history with weird parts', () => {
         const historyWithEmptyParts = [
@@ -131,7 +131,7 @@ describe('findIndexAfterFraction', () => {
             { role: 'model', parts: [{ fileData: { fileUri: 'derp' } }] },
             { role: 'user', parts: [{ text: 'Message 2' }] },
         ];
-        expect(findIndexAfterFraction(historyWithEmptyParts, 0.5)).toBe(1);
+        expect(findIndexAfterFraction(historyWithEmptyParts, 0.5)).toBe(2);
     });
 });
 describe('isThinkingSupported', () => {
@@ -176,7 +176,7 @@ describe('Gemini Client (client.ts)', () => {
         mockContentGenerator = {
             generateContent: mockGenerateContentFn,
             generateContentStream: vi.fn(),
-            countTokens: vi.fn(),
+            countTokens: vi.fn().mockResolvedValue({ totalTokens: 100 }),
             embedContent: vi.fn(),
             batchEmbedContents: vi.fn(),
         };
@@ -189,7 +189,6 @@ describe('Gemini Client (client.ts)', () => {
         };
         const fileService = new FileDiscoveryService('/test/dir');
         const contentGeneratorConfig = {
-            model: 'test-model',
             apiKey: 'test-key',
             vertexai: false,
             authType: AuthType.USE_GEMINI,
@@ -222,6 +221,9 @@ describe('Gemini Client (client.ts)', () => {
                 getDirectories: vi.fn().mockReturnValue(['/test/dir']),
             }),
             getGeminiClient: vi.fn(),
+            getModelRouterService: vi.fn().mockReturnValue({
+                route: vi.fn().mockResolvedValue({ model: 'default-routed-model' }),
+            }),
             isInFallbackMode: vi.fn().mockReturnValue(false),
             setFallbackMode: vi.fn(),
             getChatCompression: vi.fn().mockReturnValue(undefined),
@@ -400,7 +402,6 @@ describe('Gemini Client (client.ts)', () => {
         });
     });
     describe('tryCompressChat', () => {
-        const mockSendMessage = vi.fn();
         const mockGetHistory = vi.fn();
         beforeEach(() => {
             vi.mock('./tokenLimits', () => ({
@@ -410,7 +411,6 @@ describe('Gemini Client (client.ts)', () => {
                 getHistory: mockGetHistory,
                 addHistory: vi.fn(),
                 setHistory: vi.fn(),
-                sendMessage: mockSendMessage,
             };
         });
         function setup({ chatHistory = [
@@ -420,7 +420,6 @@ describe('Gemini Client (client.ts)', () => {
             const mockChat = {
                 getHistory: vi.fn().mockReturnValue(chatHistory),
                 setHistory: vi.fn(),
-                sendMessage: vi.fn().mockResolvedValue({ text: 'Summary' }),
             };
             vi.mocked(mockContentGenerator.countTokens)
                 .mockResolvedValueOnce({ totalTokens: 1000 })
@@ -527,9 +526,15 @@ describe('Gemini Client (client.ts)', () => {
                 .mockResolvedValueOnce({ totalTokens: originalTokenCount }) // First call for the check
                 .mockResolvedValueOnce({ totalTokens: newTokenCount }); // Second call for the new history
             // Mock the summary response from the chat
-            mockSendMessage.mockResolvedValue({
-                role: 'model',
-                parts: [{ text: 'This is a summary.' }],
+            mockGenerateContentFn.mockResolvedValue({
+                candidates: [
+                    {
+                        content: {
+                            role: 'model',
+                            parts: [{ text: 'This is a summary.' }],
+                        },
+                    },
+                ],
             });
             await client.tryCompressChat('prompt-id-3');
             expect(ClearcutLogger.prototype.logChatCompressionEvent).toHaveBeenCalledWith(expect.objectContaining({
@@ -553,15 +558,21 @@ describe('Gemini Client (client.ts)', () => {
                 .mockResolvedValueOnce({ totalTokens: originalTokenCount }) // First call for the check
                 .mockResolvedValueOnce({ totalTokens: newTokenCount }); // Second call for the new history
             // Mock the summary response from the chat
-            mockSendMessage.mockResolvedValue({
-                role: 'model',
-                parts: [{ text: 'This is a summary.' }],
+            mockGenerateContentFn.mockResolvedValue({
+                candidates: [
+                    {
+                        content: {
+                            role: 'model',
+                            parts: [{ text: 'This is a summary.' }],
+                        },
+                    },
+                ],
             });
             const initialChat = client.getChat();
             const result = await client.tryCompressChat('prompt-id-3');
             const newChat = client.getChat();
             expect(tokenLimit).toHaveBeenCalled();
-            expect(mockSendMessage).toHaveBeenCalled();
+            expect(mockGenerateContentFn).toHaveBeenCalled();
             // Assert that summarization happened and returned the correct stats
             expect(result).toEqual({
                 compressionStatus: CompressionStatus.COMPRESSED,
@@ -598,15 +609,21 @@ describe('Gemini Client (client.ts)', () => {
                 .mockResolvedValueOnce({ totalTokens: originalTokenCount }) // First call for the check
                 .mockResolvedValueOnce({ totalTokens: newTokenCount }); // Second call for the new history
             // Mock the summary response from the chat
-            mockSendMessage.mockResolvedValue({
-                role: 'model',
-                parts: [{ text: 'This is a summary.' }],
+            mockGenerateContentFn.mockResolvedValue({
+                candidates: [
+                    {
+                        content: {
+                            role: 'model',
+                            parts: [{ text: 'This is a summary.' }],
+                        },
+                    },
+                ],
             });
             const initialChat = client.getChat();
             const result = await client.tryCompressChat('prompt-id-3');
             const newChat = client.getChat();
             expect(tokenLimit).toHaveBeenCalled();
-            expect(mockSendMessage).toHaveBeenCalled();
+            expect(mockGenerateContentFn).toHaveBeenCalled();
             // Assert that summarization happened and returned the correct stats
             expect(result).toEqual({
                 compressionStatus: CompressionStatus.COMPRESSED,
@@ -632,14 +649,20 @@ describe('Gemini Client (client.ts)', () => {
                 .mockResolvedValueOnce({ totalTokens: originalTokenCount })
                 .mockResolvedValueOnce({ totalTokens: newTokenCount });
             // Mock the summary response from the chat
-            mockSendMessage.mockResolvedValue({
-                role: 'model',
-                parts: [{ text: 'This is a summary.' }],
+            mockGenerateContentFn.mockResolvedValue({
+                candidates: [
+                    {
+                        content: {
+                            role: 'model',
+                            parts: [{ text: 'This is a summary.' }],
+                        },
+                    },
+                ],
             });
             const initialChat = client.getChat();
             const result = await client.tryCompressChat('prompt-id-1', true); // force = true
             const newChat = client.getChat();
-            expect(mockSendMessage).toHaveBeenCalled();
+            expect(mockGenerateContentFn).toHaveBeenCalled();
             expect(result).toEqual({
                 compressionStatus: CompressionStatus.COMPRESSED,
                 originalTokenCount,
@@ -740,7 +763,7 @@ describe('Gemini Client (client.ts)', () => {
         });
         it('should include editor context when ideMode is enabled', async () => {
             // Arrange
-            vi.mocked(ideContext.getIdeContext).mockReturnValue({
+            vi.mocked(ideContextStore.get).mockReturnValue({
                 workspaceState: {
                     openFiles: [
                         {
@@ -777,7 +800,7 @@ describe('Gemini Client (client.ts)', () => {
                 // consume stream
             }
             // Assert
-            expect(ideContext.getIdeContext).toHaveBeenCalled();
+            expect(ideContextStore.get).toHaveBeenCalled();
             const expectedContext = `
 Here is the user's editor context as a JSON object. This is for your information only.
 \`\`\`json
@@ -802,7 +825,7 @@ ${JSON.stringify({
         });
         it('should not add context if ideMode is enabled but no open files', async () => {
             // Arrange
-            vi.mocked(ideContext.getIdeContext).mockReturnValue({
+            vi.mocked(ideContextStore.get).mockReturnValue({
                 workspaceState: {
                     openFiles: [],
                 },
@@ -824,12 +847,16 @@ ${JSON.stringify({
                 // consume stream
             }
             // Assert
-            expect(ideContext.getIdeContext).toHaveBeenCalled();
-            expect(mockTurnRunFn).toHaveBeenCalledWith(initialRequest, expect.any(Object));
+            expect(ideContextStore.get).toHaveBeenCalled();
+            // The `turn.run` method is now called with the model name as the first
+            // argument. We use `expect.any(String)` because this test is
+            // concerned with the IDE context logic, not the model routing,
+            // which is tested in its own dedicated suite.
+            expect(mockTurnRunFn).toHaveBeenCalledWith(expect.any(String), initialRequest, expect.any(Object));
         });
         it('should add context if ideMode is enabled and there is one active file', async () => {
             // Arrange
-            vi.mocked(ideContext.getIdeContext).mockReturnValue({
+            vi.mocked(ideContextStore.get).mockReturnValue({
                 workspaceState: {
                     openFiles: [
                         {
@@ -859,7 +886,7 @@ ${JSON.stringify({
                 // consume stream
             }
             // Assert
-            expect(ideContext.getIdeContext).toHaveBeenCalled();
+            expect(ideContextStore.get).toHaveBeenCalled();
             const expectedContext = `
 Here is the user's editor context as a JSON object. This is for your information only.
 \`\`\`json
@@ -883,7 +910,7 @@ ${JSON.stringify({
         });
         it('should add context if ideMode is enabled and there are open files but no active file', async () => {
             // Arrange
-            vi.mocked(ideContext.getIdeContext).mockReturnValue({
+            vi.mocked(ideContextStore.get).mockReturnValue({
                 workspaceState: {
                     openFiles: [
                         {
@@ -914,7 +941,7 @@ ${JSON.stringify({
                 // consume stream
             }
             // Assert
-            expect(ideContext.getIdeContext).toHaveBeenCalled();
+            expect(ideContextStore.get).toHaveBeenCalled();
             const expectedContext = `
 Here is the user's editor context as a JSON object. This is for your information only.
 \`\`\`json
@@ -1110,6 +1137,91 @@ ${JSON.stringify({
             console.log(`Infinite loop protection working: checkNextSpeaker called ${callCount} times, ` +
                 `${eventCount} events generated (properly bounded by MAX_TURNS)`);
         });
+        describe('Model Routing', () => {
+            let mockRouterService;
+            beforeEach(() => {
+                mockRouterService = {
+                    route: vi
+                        .fn()
+                        .mockResolvedValue({ model: 'routed-model', reason: 'test' }),
+                };
+                vi.mocked(mockConfig.getModelRouterService).mockReturnValue(mockRouterService);
+                mockTurnRunFn.mockReturnValue((async function* () {
+                    yield { type: 'content', value: 'Hello' };
+                })());
+            });
+            it('should use the model router service to select a model on the first turn', async () => {
+                const stream = client.sendMessageStream([{ text: 'Hi' }], new AbortController().signal, 'prompt-1');
+                await fromAsync(stream); // consume stream
+                expect(mockConfig.getModelRouterService).toHaveBeenCalled();
+                expect(mockRouterService.route).toHaveBeenCalled();
+                expect(mockTurnRunFn).toHaveBeenCalledWith('routed-model', // The model from the router
+                [{ text: 'Hi' }], expect.any(Object));
+            });
+            it('should use the same model for subsequent turns in the same prompt (stickiness)', async () => {
+                // First turn
+                let stream = client.sendMessageStream([{ text: 'Hi' }], new AbortController().signal, 'prompt-1');
+                await fromAsync(stream);
+                expect(mockRouterService.route).toHaveBeenCalledTimes(1);
+                expect(mockTurnRunFn).toHaveBeenCalledWith('routed-model', [{ text: 'Hi' }], expect.any(Object));
+                // Second turn
+                stream = client.sendMessageStream([{ text: 'Continue' }], new AbortController().signal, 'prompt-1');
+                await fromAsync(stream);
+                // Router should not be called again
+                expect(mockRouterService.route).toHaveBeenCalledTimes(1);
+                // Should stick to the first model
+                expect(mockTurnRunFn).toHaveBeenCalledWith('routed-model', [{ text: 'Continue' }], expect.any(Object));
+            });
+            it('should reset the sticky model and re-route when the prompt_id changes', async () => {
+                // First prompt
+                let stream = client.sendMessageStream([{ text: 'Hi' }], new AbortController().signal, 'prompt-1');
+                await fromAsync(stream);
+                expect(mockRouterService.route).toHaveBeenCalledTimes(1);
+                expect(mockTurnRunFn).toHaveBeenCalledWith('routed-model', [{ text: 'Hi' }], expect.any(Object));
+                // New prompt
+                mockRouterService.route.mockResolvedValue({
+                    model: 'new-routed-model',
+                    reason: 'test',
+                });
+                stream = client.sendMessageStream([{ text: 'A new topic' }], new AbortController().signal, 'prompt-2');
+                await fromAsync(stream);
+                // Router should be called again for the new prompt
+                expect(mockRouterService.route).toHaveBeenCalledTimes(2);
+                // Should use the newly routed model
+                expect(mockTurnRunFn).toHaveBeenCalledWith('new-routed-model', [{ text: 'A new topic' }], expect.any(Object));
+            });
+            it('should use the fallback model and bypass routing when in fallback mode', async () => {
+                vi.mocked(mockConfig.isInFallbackMode).mockReturnValue(true);
+                mockRouterService.route.mockResolvedValue({
+                    model: DEFAULT_GEMINI_FLASH_MODEL,
+                    reason: 'fallback',
+                });
+                const stream = client.sendMessageStream([{ text: 'Hi' }], new AbortController().signal, 'prompt-1');
+                await fromAsync(stream);
+                expect(mockTurnRunFn).toHaveBeenCalledWith(DEFAULT_GEMINI_FLASH_MODEL, [{ text: 'Hi' }], expect.any(Object));
+            });
+            it('should stick to the fallback model for the entire sequence even if fallback mode ends', async () => {
+                // Start the sequence in fallback mode
+                vi.mocked(mockConfig.isInFallbackMode).mockReturnValue(true);
+                mockRouterService.route.mockResolvedValue({
+                    model: DEFAULT_GEMINI_FLASH_MODEL,
+                    reason: 'fallback',
+                });
+                let stream = client.sendMessageStream([{ text: 'Hi' }], new AbortController().signal, 'prompt-fallback-stickiness');
+                await fromAsync(stream);
+                // First call should use fallback model
+                expect(mockTurnRunFn).toHaveBeenCalledWith(DEFAULT_GEMINI_FLASH_MODEL, [{ text: 'Hi' }], expect.any(Object));
+                // End fallback mode
+                vi.mocked(mockConfig.isInFallbackMode).mockReturnValue(false);
+                // Second call in the same sequence
+                stream = client.sendMessageStream([{ text: 'Continue' }], new AbortController().signal, 'prompt-fallback-stickiness');
+                await fromAsync(stream);
+                // Router should still not be called, and it should stick to the fallback model
+                expect(mockTurnRunFn).toHaveBeenCalledTimes(2); // Ensure it was called again
+                expect(mockTurnRunFn).toHaveBeenLastCalledWith(DEFAULT_GEMINI_FLASH_MODEL, // Still the fallback model
+                [{ text: 'Continue' }], expect.any(Object));
+            });
+        });
         describe('Editor context delta', () => {
             const mockStream = (async function* () {
                 yield { type: 'content', value: 'Hello' };
@@ -1126,7 +1238,6 @@ ${JSON.stringify({
                 const mockChat = {
                     addHistory: vi.fn(),
                     setHistory: vi.fn(),
-                    sendMessage: vi.fn().mockResolvedValue({ text: 'summary' }),
                     // Assume history is not empty for delta checks
                     getHistory: vi
                         .fn()
@@ -1250,7 +1361,7 @@ ${JSON.stringify({
                     },
                 };
                 // Setup current context
-                vi.mocked(ideContext.getIdeContext).mockReturnValue({
+                vi.mocked(ideContextStore.get).mockReturnValue({
                     workspaceState: {
                         openFiles: [
                             { ...currentActiveFile, isActive: true, timestamp: Date.now() },
@@ -1296,7 +1407,7 @@ ${JSON.stringify({
                     },
                 };
                 // Setup current context (same as previous)
-                vi.mocked(ideContext.getIdeContext).mockReturnValue({
+                vi.mocked(ideContextStore.get).mockReturnValue({
                     workspaceState: {
                         openFiles: [
                             { ...activeFile, isActive: true, timestamp: Date.now() },
@@ -1341,11 +1452,10 @@ ${JSON.stringify({
                     addHistory: vi.fn(),
                     getHistory: vi.fn().mockReturnValue([]), // Default empty history
                     setHistory: vi.fn(),
-                    sendMessage: vi.fn().mockResolvedValue({ text: 'summary' }),
                 };
                 client['chat'] = mockChat;
                 vi.spyOn(client['config'], 'getIdeMode').mockReturnValue(true);
-                vi.mocked(ideContext.getIdeContext).mockReturnValue({
+                vi.mocked(ideContextStore.get).mockReturnValue({
                     workspaceState: {
                         openFiles: [{ path: '/path/to/file.ts', timestamp: Date.now() }],
                     },
@@ -1421,7 +1531,7 @@ ${JSON.stringify({
                         openFiles: [{ path: '/path/to/fileA.ts', timestamp: Date.now() }],
                     },
                 };
-                vi.mocked(ideContext.getIdeContext).mockReturnValue(initialIdeContext);
+                vi.mocked(ideContextStore.get).mockReturnValue(initialIdeContext);
                 // Act: Send the tool response
                 let stream = client.sendMessageStream([
                     {
@@ -1467,7 +1577,7 @@ ${JSON.stringify({
                         openFiles: [{ path: '/path/to/fileB.ts', timestamp: Date.now() }],
                     },
                 };
-                vi.mocked(ideContext.getIdeContext).mockReturnValue(newIdeContext);
+                vi.mocked(ideContextStore.get).mockReturnValue(newIdeContext);
                 // Act: Send a new, regular user message
                 stream = client.sendMessageStream([{ text: 'Thanks!' }], new AbortController().signal, 'prompt-id-final');
                 for await (const _ of stream) {
@@ -1497,7 +1607,7 @@ ${JSON.stringify({
                         ],
                     },
                 };
-                vi.mocked(ideContext.getIdeContext).mockReturnValue(contextA);
+                vi.mocked(ideContextStore.get).mockReturnValue(contextA);
                 // Act: Send a regular message to establish the initial context
                 let stream = client.sendMessageStream([{ text: 'Initial message' }], new AbortController().signal, 'prompt-id-initial');
                 for await (const _ of stream) {
@@ -1530,7 +1640,7 @@ ${JSON.stringify({
                         ],
                     },
                 };
-                vi.mocked(ideContext.getIdeContext).mockReturnValue(contextB);
+                vi.mocked(ideContextStore.get).mockReturnValue(contextB);
                 // Act: Send the tool response
                 stream = client.sendMessageStream([
                     {
@@ -1575,7 +1685,7 @@ ${JSON.stringify({
                         ],
                     },
                 };
-                vi.mocked(ideContext.getIdeContext).mockReturnValue(contextC);
+                vi.mocked(ideContextStore.get).mockReturnValue(contextC);
                 // Act: Send a new, regular user message
                 stream = client.sendMessageStream([{ text: 'Thanks!' }], new AbortController().signal, 'prompt-id-final');
                 for await (const _ of stream) {