@google/gemini-cli-core 0.0.8999999 → 0.0.77777773

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

@@ -80,6 +80,7 @@ vi.mock('../ide/ideContext.js');
 vi.mock('../telemetry/uiTelemetry.js', () => ({
     uiTelemetryService: {
         setLastPromptTokenCount: vi.fn(),
+        getLastPromptTokenCount: vi.fn(),
     },
 }));
 /**
@@ -94,7 +95,7 @@ async function fromAsync(promise) {
     }
     return results;
 }
-describe('findIndexAfterFraction', () => {
+describe('findCompressSplitPoint', () => {
     it('should throw an error for non-positive numbers', () => {
         expect(() => findCompressSplitPoint([], 0)).toThrow('Fraction must be between 0 and 1');
     });
@@ -112,7 +113,7 @@ describe('findIndexAfterFraction', () => {
             { role: 'model', parts: [{ text: 'This is the fourth message.' }] }, // JSON length: 68 (80%)
             { role: 'user', parts: [{ text: 'This is the fifth message.' }] }, // JSON length: 65 (100%)
         ];
-        expect(findCompressSplitPoint(history, 0.5)).toBe(2);
+        expect(findCompressSplitPoint(history, 0.5)).toBe(4);
     });
     it('should handle a fraction of last index', () => {
         const history = [
@@ -201,7 +202,6 @@ describe('Gemini Client (client.ts)', () => {
         mockContentGenerator = {
             generateContent: mockGenerateContentFn,
             generateContentStream: vi.fn(),
-            countTokens: vi.fn().mockResolvedValue({ totalTokens: 100 }),
             batchEmbedContents: vi.fn(),
         };
         // Because the GeminiClient constructor kicks off an async process (startChat)
@@ -254,6 +254,7 @@ describe('Gemini Client (client.ts)', () => {
             getSkipNextSpeakerCheck: vi.fn().mockReturnValue(false),
             getUseSmartEdit: vi.fn().mockReturnValue(false),
             getUseModelRouter: vi.fn().mockReturnValue(false),
+            getContinueOnFailedApiCall: vi.fn(),
             getProjectRoot: vi.fn().mockReturnValue('/test/project/root'),
             storage: {
                 getProjectTempDir: vi.fn().mockReturnValue('/test/temp'),
@@ -324,74 +325,128 @@ describe('Gemini Client (client.ts)', () => {
         function setup({ chatHistory = [
             { role: 'user', parts: [{ text: 'Long conversation' }] },
             { role: 'model', parts: [{ text: 'Long response' }] },
-        ], } = {}) {
-            const mockChat = {
-                getHistory: vi.fn().mockReturnValue(chatHistory),
+        ], originalTokenCount = 1000, summaryText = 'This is a summary.', } = {}) {
+            const mockOriginalChat = {
+                getHistory: vi.fn((_curated) => chatHistory),
                 setHistory: vi.fn(),
             };
-            vi.mocked(mockContentGenerator.countTokens)
-                .mockResolvedValueOnce({ totalTokens: 1000 })
-                .mockResolvedValueOnce({ totalTokens: 5000 });
-            client['chat'] = mockChat;
-            client['startChat'] = vi.fn().mockResolvedValue({ ...mockChat });
-            return { client, mockChat };
+            client['chat'] = mockOriginalChat;
+            vi.mocked(uiTelemetryService.getLastPromptTokenCount).mockReturnValue(originalTokenCount);
+            mockGenerateContentFn.mockResolvedValue({
+                candidates: [
+                    {
+                        content: {
+                            role: 'model',
+                            parts: [{ text: summaryText }],
+                        },
+                    },
+                ],
+            });
+            // Calculate what the new history will be
+            const splitPoint = findCompressSplitPoint(chatHistory, 0.7); // 1 - 0.3
+            const historyToKeep = chatHistory.slice(splitPoint);
+            // This is the history that the new chat will have.
+            // It includes the default startChat history + the extra history from tryCompressChat
+            const newCompressedHistory = [
+                // Mocked envParts + canned response from startChat
+                {
+                    role: 'user',
+                    parts: [{ text: 'Mocked env context' }],
+                },
+                {
+                    role: 'model',
+                    parts: [{ text: 'Got it. Thanks for the context!' }],
+                },
+                // extraHistory from tryCompressChat
+                {
+                    role: 'user',
+                    parts: [{ text: summaryText }],
+                },
+                {
+                    role: 'model',
+                    parts: [{ text: 'Got it. Thanks for the additional context!' }],
+                },
+                ...historyToKeep,
+            ];
+            const mockNewChat = {
+                getHistory: vi.fn().mockReturnValue(newCompressedHistory),
+                setHistory: vi.fn(),
+            };
+            client['startChat'] = vi
+                .fn()
+                .mockResolvedValue(mockNewChat);
+            const totalChars = newCompressedHistory.reduce((total, content) => total + JSON.stringify(content).length, 0);
+            const estimatedNewTokenCount = Math.floor(totalChars / 4);
+            return {
+                client,
+                mockOriginalChat,
+                mockNewChat,
+                estimatedNewTokenCount,
+            };
         }
         describe('when compression inflates the token count', () => {
             it('allows compression to be forced/manual after a failure', async () => {
-                const { client } = setup();
-                vi.mocked(mockContentGenerator.countTokens).mockResolvedValue({
-                    totalTokens: 1000,
+                // Call 1 (Fails): Setup with a long summary to inflate tokens
+                const longSummary = 'long summary '.repeat(100);
+                const { client, estimatedNewTokenCount: inflatedTokenCount } = setup({
+                    originalTokenCount: 100,
+                    summaryText: longSummary,
                 });
+                expect(inflatedTokenCount).toBeGreaterThan(100); // Ensure setup is correct
                 await client.tryCompressChat('prompt-id-4', false); // Fails
-                const result = await client.tryCompressChat('prompt-id-4', true);
+                // Call 2 (Forced): Re-setup with a short summary
+                const shortSummary = 'short';
+                const { estimatedNewTokenCount: compressedTokenCount } = setup({
+                    originalTokenCount: 100,
+                    summaryText: shortSummary,
+                });
+                expect(compressedTokenCount).toBeLessThanOrEqual(100); // Ensure setup is correct
+                const result = await client.tryCompressChat('prompt-id-4', true); // Forced
                 expect(result).toEqual({
                     compressionStatus: CompressionStatus.COMPRESSED,
-                    newTokenCount: 1000,
-                    originalTokenCount: 1000,
+                    newTokenCount: compressedTokenCount,
+                    originalTokenCount: 100,
                 });
             });
             it('yields the result even if the compression inflated the tokens', async () => {
-                const { client } = setup();
-                vi.mocked(mockContentGenerator.countTokens).mockResolvedValue({
-                    totalTokens: 1000,
+                const longSummary = 'long summary '.repeat(100);
+                const { client, estimatedNewTokenCount } = setup({
+                    originalTokenCount: 100,
+                    summaryText: longSummary,
                 });
+                expect(estimatedNewTokenCount).toBeGreaterThan(100); // Ensure setup is correct
                 const result = await client.tryCompressChat('prompt-id-4', false);
                 expect(result).toEqual({
                     compressionStatus: CompressionStatus.COMPRESSION_FAILED_INFLATED_TOKEN_COUNT,
-                    newTokenCount: 5000,
-                    originalTokenCount: 1000,
+                    newTokenCount: estimatedNewTokenCount,
+                    originalTokenCount: 100,
                 });
-                expect(uiTelemetryService.setLastPromptTokenCount).toHaveBeenCalledWith(5000);
-                expect(uiTelemetryService.setLastPromptTokenCount).toHaveBeenCalledTimes(1);
+                // IMPORTANT: The change in client.ts means setLastPromptTokenCount is NOT called on failure
+                expect(uiTelemetryService.setLastPromptTokenCount).not.toHaveBeenCalled();
             });
             it('does not manipulate the source chat', async () => {
-                const { client, mockChat } = setup();
-                await client.tryCompressChat('prompt-id-4', false);
-                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);
-                vi.mocked(mockContentGenerator.countTokens).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 longSummary = 'long summary '.repeat(100);
+                const { client, mockOriginalChat, estimatedNewTokenCount } = setup({
+                    originalTokenCount: 100,
+                    summaryText: longSummary,
                 });
-                const { compressionStatus } = await client.tryCompressChat('prompt-id-4', false);
-                expect(compressionStatus).toBe(CompressionStatus.COMPRESSION_FAILED_INFLATED_TOKEN_COUNT);
-                expect(client['chat']?.setHistory).toHaveBeenCalledWith(originalHistory);
+                expect(estimatedNewTokenCount).toBeGreaterThan(100); // Ensure setup is correct
+                await client.tryCompressChat('prompt-id-4', false);
+                // On failure, the chat should NOT be replaced
+                expect(client['chat']).toBe(mockOriginalChat);
             });
             it('will not attempt to compress context after a failure', async () => {
-                const { client } = setup();
-                await client.tryCompressChat('prompt-id-4', false);
+                const longSummary = 'long summary '.repeat(100);
+                const { client, estimatedNewTokenCount } = setup({
+                    originalTokenCount: 100,
+                    summaryText: longSummary,
+                });
+                expect(estimatedNewTokenCount).toBeGreaterThan(100); // Ensure setup is correct
+                await client.tryCompressChat('prompt-id-4', false); // This fails and sets hasFailedCompressionAttempt = true
+                // This call should now be a NOOP
                 const result = await client.tryCompressChat('prompt-id-5', false);
-                // it counts tokens for {original, compressed} and then never again
-                expect(mockContentGenerator.countTokens).toHaveBeenCalledTimes(2);
+                // generateContent (for summary) should only have been called once
+                expect(mockGenerateContentFn).toHaveBeenCalledTimes(1);
                 expect(result).toEqual({
                     compressionStatus: CompressionStatus.NOOP,
                     newTokenCount: 0,
@@ -405,17 +460,16 @@ describe('Gemini Client (client.ts)', () => {
             mockGetHistory.mockReturnValue([
                 { role: 'user', parts: [{ text: '...history...' }] },
             ]);
-            vi.mocked(mockContentGenerator.countTokens).mockResolvedValue({
-                totalTokens: MOCKED_TOKEN_LIMIT * 0.699, // TOKEN_THRESHOLD_FOR_SUMMARIZATION = 0.7
-            });
+            const originalTokenCount = MOCKED_TOKEN_LIMIT * 0.699;
+            vi.mocked(uiTelemetryService.getLastPromptTokenCount).mockReturnValue(originalTokenCount);
             const initialChat = client.getChat();
             const result = await client.tryCompressChat('prompt-id-2', false);
             const newChat = client.getChat();
             expect(tokenLimit).toHaveBeenCalled();
             expect(result).toEqual({
                 compressionStatus: CompressionStatus.NOOP,
-                newTokenCount: 699,
-                originalTokenCount: 699,
+                newTokenCount: originalTokenCount,
+                originalTokenCount,
             });
             expect(newChat).toBe(initialChat);
         });
@@ -427,21 +481,40 @@ describe('Gemini Client (client.ts)', () => {
             vi.spyOn(client['config'], 'getChatCompression').mockReturnValue({
                 contextPercentageThreshold: MOCKED_CONTEXT_PERCENTAGE_THRESHOLD,
             });
-            mockGetHistory.mockReturnValue([
-                { role: 'user', parts: [{ text: '...history...' }] },
-            ]);
+            const history = [{ role: 'user', parts: [{ text: '...history...' }] }];
+            mockGetHistory.mockReturnValue(history);
             const originalTokenCount = MOCKED_TOKEN_LIMIT * MOCKED_CONTEXT_PERCENTAGE_THRESHOLD;
-            const newTokenCount = 100;
-            vi.mocked(mockContentGenerator.countTokens)
-                .mockResolvedValueOnce({ totalTokens: originalTokenCount }) // First call for the check
-                .mockResolvedValueOnce({ totalTokens: newTokenCount }); // Second call for the new history
+            vi.mocked(uiTelemetryService.getLastPromptTokenCount).mockReturnValue(originalTokenCount);
+            // We need to control the estimated new token count.
+            // We mock startChat to return a chat with a known history.
+            const summaryText = 'This is a summary.';
+            const splitPoint = findCompressSplitPoint(history, 0.7);
+            const historyToKeep = history.slice(splitPoint);
+            const newCompressedHistory = [
+                { role: 'user', parts: [{ text: 'Mocked env context' }] },
+                { role: 'model', parts: [{ text: 'Got it. Thanks for the context!' }] },
+                { role: 'user', parts: [{ text: summaryText }] },
+                {
+                    role: 'model',
+                    parts: [{ text: 'Got it. Thanks for the additional context!' }],
+                },
+                ...historyToKeep,
+            ];
+            const mockNewChat = {
+                getHistory: vi.fn().mockReturnValue(newCompressedHistory),
+            };
+            client['startChat'] = vi
+                .fn()
+                .mockResolvedValue(mockNewChat);
+            const totalChars = newCompressedHistory.reduce((total, content) => total + JSON.stringify(content).length, 0);
+            const newTokenCount = Math.floor(totalChars / 4);
             // Mock the summary response from the chat
             mockGenerateContentFn.mockResolvedValue({
                 candidates: [
                     {
                         content: {
                             role: 'model',
-                            parts: [{ text: 'This is a summary.' }],
+                            parts: [{ text: summaryText }],
                         },
                     },
                 ],
@@ -461,21 +534,39 @@ describe('Gemini Client (client.ts)', () => {
             vi.spyOn(client['config'], 'getChatCompression').mockReturnValue({
                 contextPercentageThreshold: MOCKED_CONTEXT_PERCENTAGE_THRESHOLD,
             });
-            mockGetHistory.mockReturnValue([
-                { role: 'user', parts: [{ text: '...history...' }] },
-            ]);
+            const history = [{ role: 'user', parts: [{ text: '...history...' }] }];
+            mockGetHistory.mockReturnValue(history);
             const originalTokenCount = MOCKED_TOKEN_LIMIT * MOCKED_CONTEXT_PERCENTAGE_THRESHOLD;
-            const newTokenCount = 100;
-            vi.mocked(mockContentGenerator.countTokens)
-                .mockResolvedValueOnce({ totalTokens: originalTokenCount }) // First call for the check
-                .mockResolvedValueOnce({ totalTokens: newTokenCount }); // Second call for the new history
+            vi.mocked(uiTelemetryService.getLastPromptTokenCount).mockReturnValue(originalTokenCount);
+            // Mock summary and new chat
+            const summaryText = 'This is a summary.';
+            const splitPoint = findCompressSplitPoint(history, 0.7);
+            const historyToKeep = history.slice(splitPoint);
+            const newCompressedHistory = [
+                { role: 'user', parts: [{ text: 'Mocked env context' }] },
+                { role: 'model', parts: [{ text: 'Got it. Thanks for the context!' }] },
+                { role: 'user', parts: [{ text: summaryText }] },
+                {
+                    role: 'model',
+                    parts: [{ text: 'Got it. Thanks for the additional context!' }],
+                },
+                ...historyToKeep,
+            ];
+            const mockNewChat = {
+                getHistory: vi.fn().mockReturnValue(newCompressedHistory),
+            };
+            client['startChat'] = vi
+                .fn()
+                .mockResolvedValue(mockNewChat);
+            const totalChars = newCompressedHistory.reduce((total, content) => total + JSON.stringify(content).length, 0);
+            const newTokenCount = Math.floor(totalChars / 4);
             // Mock the summary response from the chat
             mockGenerateContentFn.mockResolvedValue({
                 candidates: [
                     {
                         content: {
                             role: 'model',
-                            parts: [{ text: 'This is a summary.' }],
+                            parts: [{ text: summaryText }],
                         },
                     },
                 ],
@@ -497,7 +588,7 @@ describe('Gemini Client (client.ts)', () => {
         it('should not compress across a function call response', async () => {
             const MOCKED_TOKEN_LIMIT = 1000;
             vi.mocked(tokenLimit).mockReturnValue(MOCKED_TOKEN_LIMIT);
-            mockGetHistory.mockReturnValue([
+            const history = [
                 { role: 'user', parts: [{ text: '...history 1...' }] },
                 { role: 'model', parts: [{ text: '...history 2...' }] },
                 { role: 'user', parts: [{ text: '...history 3...' }] },
@@ -514,19 +605,43 @@ describe('Gemini Client (client.ts)', () => {
                 { role: 'model', parts: [{ text: '...history 10...' }] },
                 // Instead we will break here.
                 { role: 'user', parts: [{ text: '...history 10...' }] },
-            ]);
+            ];
+            mockGetHistory.mockReturnValue(history);
             const originalTokenCount = 1000 * 0.7;
-            const newTokenCount = 100;
-            vi.mocked(mockContentGenerator.countTokens)
-                .mockResolvedValueOnce({ totalTokens: originalTokenCount }) // First call for the check
-                .mockResolvedValueOnce({ totalTokens: newTokenCount }); // Second call for the new history
+            vi.mocked(uiTelemetryService.getLastPromptTokenCount).mockReturnValue(originalTokenCount);
+            // Mock summary and new chat
+            const summaryText = 'This is a summary.';
+            const splitPoint = findCompressSplitPoint(history, 0.7); // This should be 10
+            expect(splitPoint).toBe(10); // Verify split point logic
+            const historyToKeep = history.slice(splitPoint); // Should keep last user message
+            expect(historyToKeep).toEqual([
+                { role: 'user', parts: [{ text: '...history 10...' }] },
+            ]);
+            const newCompressedHistory = [
+                { role: 'user', parts: [{ text: 'Mocked env context' }] },
+                { role: 'model', parts: [{ text: 'Got it. Thanks for the context!' }] },
+                { role: 'user', parts: [{ text: summaryText }] },
+                {
+                    role: 'model',
+                    parts: [{ text: 'Got it. Thanks for the additional context!' }],
+                },
+                ...historyToKeep,
+            ];
+            const mockNewChat = {
+                getHistory: vi.fn().mockReturnValue(newCompressedHistory),
+            };
+            client['startChat'] = vi
+                .fn()
+                .mockResolvedValue(mockNewChat);
+            const totalChars = newCompressedHistory.reduce((total, content) => total + JSON.stringify(content).length, 0);
+            const newTokenCount = Math.floor(totalChars / 4);
             // Mock the summary response from the chat
             mockGenerateContentFn.mockResolvedValue({
                 candidates: [
                     {
                         content: {
                             role: 'model',
-                            parts: [{ text: 'This is a summary.' }],
+                            parts: [{ text: summaryText }],
                         },
                     },
                 ],
@@ -544,35 +659,53 @@ describe('Gemini Client (client.ts)', () => {
             });
             // Assert that the chat was reset
             expect(newChat).not.toBe(initialChat);
-            // 1. standard start context message
-            // 2. standard canned user start message
-            // 3. compressed summary message
-            // 4. standard canned user summary message
-            // 5. The last user message (not the last 3 because that would start with a function response)
+            // 1. standard start context message (env)
+            // 2. standard canned model response
+            // 3. compressed summary message (user)
+            // 4. standard canned model response
+            // 5. The last user message (historyToKeep)
             expect(newChat.getHistory().length).toEqual(5);
         });
         it('should always trigger summarization when force is true, regardless of token count', async () => {
-            mockGetHistory.mockReturnValue([
-                { role: 'user', parts: [{ text: '...history...' }] },
-            ]);
-            const originalTokenCount = 10; // Well below threshold
-            const newTokenCount = 5;
-            vi.mocked(mockContentGenerator.countTokens)
-                .mockResolvedValueOnce({ totalTokens: originalTokenCount })
-                .mockResolvedValueOnce({ totalTokens: newTokenCount });
+            const history = [{ role: 'user', parts: [{ text: '...history...' }] }];
+            mockGetHistory.mockReturnValue(history);
+            const originalTokenCount = 100; // Well below threshold, but > estimated new count
+            vi.mocked(uiTelemetryService.getLastPromptTokenCount).mockReturnValue(originalTokenCount);
+            // Mock summary and new chat
+            const summaryText = 'This is a summary.';
+            const splitPoint = findCompressSplitPoint(history, 0.7);
+            const historyToKeep = history.slice(splitPoint);
+            const newCompressedHistory = [
+                { role: 'user', parts: [{ text: 'Mocked env context' }] },
+                { role: 'model', parts: [{ text: 'Got it. Thanks for the context!' }] },
+                { role: 'user', parts: [{ text: summaryText }] },
+                {
+                    role: 'model',
+                    parts: [{ text: 'Got it. Thanks for the additional context!' }],
+                },
+                ...historyToKeep,
+            ];
+            const mockNewChat = {
+                getHistory: vi.fn().mockReturnValue(newCompressedHistory),
+            };
+            client['startChat'] = vi
+                .fn()
+                .mockResolvedValue(mockNewChat);
+            const totalChars = newCompressedHistory.reduce((total, content) => total + JSON.stringify(content).length, 0);
+            const newTokenCount = Math.floor(totalChars / 4);
             // Mock the summary response from the chat
             mockGenerateContentFn.mockResolvedValue({
                 candidates: [
                     {
                         content: {
                             role: 'model',
-                            parts: [{ text: 'This is a summary.' }],
+                            parts: [{ text: summaryText }],
                         },
                     },
                 ],
             });
             const initialChat = client.getChat();
-            const result = await client.tryCompressChat('prompt-id-1', false); // force = true
+            const result = await client.tryCompressChat('prompt-id-1', true); // force = true
             const newChat = client.getChat();
             expect(mockGenerateContentFn).toHaveBeenCalled();
             expect(result).toEqual({
@@ -583,45 +716,6 @@ describe('Gemini Client (client.ts)', () => {
             // 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 = 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' });
-            const mockChatHistory = [
-                { role: 'user', parts: [{ text: 'Long conversation' }] },
-                { role: 'model', parts: [{ text: 'Long response' }] },
-            ];
-            const mockChat = {
-                getHistory: vi.fn().mockImplementation(() => [...mockChatHistory]),
-                setHistory: vi.fn(),
-                sendMessage: mockSendMessage,
-            };
-            client['chat'] = mockChat;
-            client['startChat'] = vi.fn().mockResolvedValue(mockChat);
-            const result = await client.tryCompressChat('prompt-id-4', false);
-            expect(mockContentGenerator.countTokens).toHaveBeenCalledTimes(2);
-            expect(mockContentGenerator.countTokens).toHaveBeenNthCalledWith(1, {
-                model: firstCurrentModel,
-                contents: [...mockChatHistory],
-            });
-            expect(mockContentGenerator.countTokens).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 () => {
@@ -649,9 +743,6 @@ describe('Gemini Client (client.ts)', () => {
                 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* () {
@@ -909,6 +1000,7 @@ ${JSON.stringify({
             expect(finalResult).toBeInstanceOf(Turn);
         });
         it('should stop infinite loop after MAX_TURNS when nextSpeaker always returns model', async () => {
+            vi.spyOn(client['config'], 'getContinueOnFailedApiCall').mockReturnValue(true);
             // Get the mocked checkNextSpeaker function and configure it to trigger infinite loop
             const { checkNextSpeaker } = await import('../utils/nextSpeakerChecker.js');
             const mockCheckNextSpeaker = vi.mocked(checkNextSpeaker);
@@ -1064,6 +1156,82 @@ ${JSON.stringify({
             console.log(`Infinite loop protection working: checkNextSpeaker called ${callCount} times, ` +
                 `${eventCount} events generated (properly bounded by MAX_TURNS)`);
         });
+        it('should yield ContextWindowWillOverflow when the context window is about to overflow', async () => {
+            // Arrange
+            const MOCKED_TOKEN_LIMIT = 1000;
+            vi.mocked(tokenLimit).mockReturnValue(MOCKED_TOKEN_LIMIT);
+            // Set last prompt token count
+            const lastPromptTokenCount = 900;
+            vi.mocked(uiTelemetryService.getLastPromptTokenCount).mockReturnValue(lastPromptTokenCount);
+            // Remaining = 100. Threshold (95%) = 95.
+            // We need a request > 95 tokens.
+            // A string of length 400 is roughly 100 tokens.
+            const longText = 'a'.repeat(400);
+            const request = [{ text: longText }];
+            const estimatedRequestTokenCount = Math.floor(JSON.stringify(request).length / 4);
+            const remainingTokenCount = MOCKED_TOKEN_LIMIT - lastPromptTokenCount;
+            // Mock tryCompressChat to not compress
+            vi.spyOn(client, 'tryCompressChat').mockResolvedValue({
+                originalTokenCount: lastPromptTokenCount,
+                newTokenCount: lastPromptTokenCount,
+                compressionStatus: CompressionStatus.NOOP,
+            });
+            // Act
+            const stream = client.sendMessageStream(request, new AbortController().signal, 'prompt-id-overflow');
+            const events = await fromAsync(stream);
+            // Assert
+            expect(events).toContainEqual({
+                type: GeminiEventType.ContextWindowWillOverflow,
+                value: {
+                    estimatedRequestTokenCount,
+                    remainingTokenCount,
+                },
+            });
+            // Ensure turn.run is not called
+            expect(mockTurnRunFn).not.toHaveBeenCalled();
+        });
+        it("should use the sticky model's token limit for the overflow check", async () => {
+            // Arrange
+            const STICKY_MODEL = 'gemini-1.5-flash';
+            const STICKY_MODEL_LIMIT = 1000;
+            const CONFIG_MODEL_LIMIT = 2000;
+            // Set up token limits
+            vi.mocked(tokenLimit).mockImplementation((model) => {
+                if (model === STICKY_MODEL)
+                    return STICKY_MODEL_LIMIT;
+                return CONFIG_MODEL_LIMIT;
+            });
+            // Set the sticky model
+            client['currentSequenceModel'] = STICKY_MODEL;
+            // Set token count
+            const lastPromptTokenCount = 900;
+            vi.mocked(uiTelemetryService.getLastPromptTokenCount).mockReturnValue(lastPromptTokenCount);
+            // Remaining (sticky) = 100. Threshold (95%) = 95.
+            // We need a request > 95 tokens.
+            const longText = 'a'.repeat(400);
+            const request = [{ text: longText }];
+            const estimatedRequestTokenCount = Math.floor(JSON.stringify(request).length / 4);
+            const remainingTokenCount = STICKY_MODEL_LIMIT - lastPromptTokenCount;
+            vi.spyOn(client, 'tryCompressChat').mockResolvedValue({
+                originalTokenCount: lastPromptTokenCount,
+                newTokenCount: lastPromptTokenCount,
+                compressionStatus: CompressionStatus.NOOP,
+            });
+            // Act
+            const stream = client.sendMessageStream(request, new AbortController().signal, 'test-session-id');
+            const events = await fromAsync(stream);
+            // Assert
+            // Should overflow based on the sticky model's limit
+            expect(events).toContainEqual({
+                type: GeminiEventType.ContextWindowWillOverflow,
+                value: {
+                    estimatedRequestTokenCount,
+                    remainingTokenCount,
+                },
+            });
+            expect(tokenLimit).toHaveBeenCalledWith(STICKY_MODEL);
+            expect(mockTurnRunFn).not.toHaveBeenCalled();
+        });
         describe('Model Routing', () => {
             let mockRouterService;
             beforeEach(() => {
@@ -1149,6 +1317,89 @@ ${JSON.stringify({
                 [{ text: 'Continue' }], expect.any(Object));
             });
         });
+        it('should recursively call sendMessageStream with "Please continue." when InvalidStream event is received', async () => {
+            vi.spyOn(client['config'], 'getContinueOnFailedApiCall').mockReturnValue(true);
+            // Arrange
+            const mockStream1 = (async function* () {
+                yield { type: GeminiEventType.InvalidStream };
+            })();
+            const mockStream2 = (async function* () {
+                yield { type: GeminiEventType.Content, value: 'Continued content' };
+            })();
+            mockTurnRunFn
+                .mockReturnValueOnce(mockStream1)
+                .mockReturnValueOnce(mockStream2);
+            const mockChat = {
+                addHistory: vi.fn(),
+                getHistory: vi.fn().mockReturnValue([]),
+            };
+            client['chat'] = mockChat;
+            const initialRequest = [{ text: 'Hi' }];
+            const promptId = 'prompt-id-invalid-stream';
+            const signal = new AbortController().signal;
+            // Act
+            const stream = client.sendMessageStream(initialRequest, signal, promptId);
+            const events = await fromAsync(stream);
+            // Assert
+            expect(events).toEqual([
+                { type: GeminiEventType.InvalidStream },
+                { type: GeminiEventType.Content, value: 'Continued content' },
+            ]);
+            // Verify that turn.run was called twice
+            expect(mockTurnRunFn).toHaveBeenCalledTimes(2);
+            // First call with original request
+            expect(mockTurnRunFn).toHaveBeenNthCalledWith(1, expect.any(String), initialRequest, expect.any(Object));
+            // Second call with "Please continue."
+            expect(mockTurnRunFn).toHaveBeenNthCalledWith(2, expect.any(String), [{ text: 'System: Please continue.' }], expect.any(Object));
+        });
+        it('should not recursively call sendMessageStream with "Please continue." when InvalidStream event is received and flag is false', async () => {
+            vi.spyOn(client['config'], 'getContinueOnFailedApiCall').mockReturnValue(false);
+            // Arrange
+            const mockStream1 = (async function* () {
+                yield { type: GeminiEventType.InvalidStream };
+            })();
+            mockTurnRunFn.mockReturnValueOnce(mockStream1);
+            const mockChat = {
+                addHistory: vi.fn(),
+                getHistory: vi.fn().mockReturnValue([]),
+            };
+            client['chat'] = mockChat;
+            const initialRequest = [{ text: 'Hi' }];
+            const promptId = 'prompt-id-invalid-stream';
+            const signal = new AbortController().signal;
+            // Act
+            const stream = client.sendMessageStream(initialRequest, signal, promptId);
+            const events = await fromAsync(stream);
+            // Assert
+            expect(events).toEqual([{ type: GeminiEventType.InvalidStream }]);
+            // Verify that turn.run was called only once
+            expect(mockTurnRunFn).toHaveBeenCalledTimes(1);
+        });
+        it('should stop recursing after one retry when InvalidStream events are repeatedly received', async () => {
+            vi.spyOn(client['config'], 'getContinueOnFailedApiCall').mockReturnValue(true);
+            // Arrange
+            // Always return a new invalid stream
+            mockTurnRunFn.mockImplementation(() => (async function* () {
+                yield { type: GeminiEventType.InvalidStream };
+            })());
+            const mockChat = {
+                addHistory: vi.fn(),
+                getHistory: vi.fn().mockReturnValue([]),
+            };
+            client['chat'] = mockChat;
+            const initialRequest = [{ text: 'Hi' }];
+            const promptId = 'prompt-id-infinite-invalid-stream';
+            const signal = new AbortController().signal;
+            // Act
+            const stream = client.sendMessageStream(initialRequest, signal, promptId);
+            const events = await fromAsync(stream);
+            // Assert
+            // We expect 2 InvalidStream events (original + 1 retry)
+            expect(events.length).toBe(2);
+            expect(events.every((e) => e.type === GeminiEventType.InvalidStream)).toBe(true);
+            // Verify that turn.run was called twice
+            expect(mockTurnRunFn).toHaveBeenCalledTimes(2);
+        });
         describe('Editor context delta', () => {
             const mockStream = (async function* () {
                 yield { type: 'content', value: 'Hello' };
@@ -1717,7 +1968,7 @@ ${JSON.stringify({
                 model: DEFAULT_GEMINI_FLASH_MODEL,
                 config: {
                     abortSignal,
-                    systemInstruction: getCoreSystemPrompt(''),
+                    systemInstruction: getCoreSystemPrompt({}, ''),
                     temperature: 0.5,
                     topP: 1,
                 },