@machina.ai/cell-cli-core 1.6.1-rc2 → 1.10.0-rc1

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

@@ -4,7 +4,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 import { describe, it, expect, vi, beforeEach, afterEach, } from 'vitest';
-import { findIndexAfterFraction, isThinkingDefault, isThinkingSupported, GeminiClient, } from './client.js';
+import { findCompressSplitPoint, isThinkingDefault, isThinkingSupported, GeminiClient, } from './client.js';
 import { AuthType, } from './contentGenerator.js';
 import {} from './geminiChat.js';
 import { CompressionStatus, GeminiEventType, Turn, } from './turn.js';
@@ -15,6 +15,7 @@ import { setSimulate429 } from '../utils/testUtils.js';
 import { tokenLimit } from './tokenLimits.js';
 import { ideContextStore } from '../ide/ideContext.js';
 import { ClearcutLogger } from '../telemetry/clearcut-logger/clearcut-logger.js';
+import { uiTelemetryService } from '../telemetry/uiTelemetry.js';
 // Mock fs module to prevent actual file system operations during tests
 const mockFileSystem = new Map();
 vi.mock('node:fs', () => {
@@ -76,6 +77,12 @@ vi.mock('../telemetry/index.js', () => ({
     logApiError: vi.fn(),
 }));
 vi.mock('../ide/ideContext.js');
+vi.mock('../telemetry/uiTelemetry.js', () => ({
+    uiTelemetryService: {
+        setLastPromptTokenCount: vi.fn(),
+        getLastPromptTokenCount: vi.fn(),
+    },
+}));
 /**
  * Array.fromAsync ponyfill, which will be available in es 2024.
  *
@@ -88,42 +95,60 @@ async function fromAsync(promise) {
     }
     return results;
 }
-describe('findIndexAfterFraction', () => {
-    const history = [
-        { role: 'user', parts: [{ text: 'This is the first message.' }] }, // JSON length: 66
-        { role: 'model', parts: [{ text: 'This is the second message.' }] }, // JSON length: 68
-        { role: 'user', parts: [{ text: 'This is the third message.' }] }, // JSON length: 66
-        { role: 'model', parts: [{ text: 'This is the fourth message.' }] }, // JSON length: 68
-        { role: 'user', parts: [{ text: 'This is the fifth message.' }] }, // JSON length: 65
-    ];
-    // Total length: 333
+describe('findCompressSplitPoint', () => {
     it('should throw an error for non-positive numbers', () => {
-        expect(() => findIndexAfterFraction(history, 0)).toThrow('Fraction must be between 0 and 1');
+        expect(() => findCompressSplitPoint([], 0)).toThrow('Fraction must be between 0 and 1');
     });
     it('should throw an error for a fraction greater than or equal to 1', () => {
-        expect(() => findIndexAfterFraction(history, 1)).toThrow('Fraction must be between 0 and 1');
+        expect(() => findCompressSplitPoint([], 1)).toThrow('Fraction must be between 0 and 1');
+    });
+    it('should handle an empty history', () => {
+        expect(findCompressSplitPoint([], 0.5)).toBe(0);
     });
     it('should handle a fraction in the middle', () => {
-        // 333 * 0.5 = 166.5
-        // 0: 66
-        // 1: 66 + 68 = 134
-        // 2: 134 + 66 = 200
-        // 200 >= 166.5, so index is 3
-        expect(findIndexAfterFraction(history, 0.5)).toBe(3);
+        const history = [
+            { role: 'user', parts: [{ text: 'This is the first message.' }] }, // JSON length: 66 (19%)
+            { role: 'model', parts: [{ text: 'This is the second message.' }] }, // JSON length: 68 (40%)
+            { role: 'user', parts: [{ text: 'This is the third message.' }] }, // JSON length: 66 (60%)
+            { 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(4);
+    });
+    it('should handle a fraction of last index', () => {
+        const history = [
+            { role: 'user', parts: [{ text: 'This is the first message.' }] }, // JSON length: 66 (19%)
+            { role: 'model', parts: [{ text: 'This is the second message.' }] }, // JSON length: 68 (40%)
+            { role: 'user', parts: [{ text: 'This is the third message.' }] }, // JSON length: 66 (60%)
+            { 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.9)).toBe(4);
     });
-    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 5
-        expect(findIndexAfterFraction(history, 0.9)).toBe(5);
+    it('should handle a fraction of after last index', () => {
+        const history = [
+            { role: 'user', parts: [{ text: 'This is the first message.' }] }, // JSON length: 66 (24%%)
+            { role: 'model', parts: [{ text: 'This is the second message.' }] }, // JSON length: 68 (50%)
+            { role: 'user', parts: [{ text: 'This is the third message.' }] }, // JSON length: 66 (74%)
+            { role: 'model', parts: [{ text: 'This is the fourth message.' }] }, // JSON length: 68 (100%)
+        ];
+        expect(findCompressSplitPoint(history, 0.8)).toBe(4);
     });
-    it('should handle an empty history', () => {
-        expect(findIndexAfterFraction([], 0.5)).toBe(0);
+    it('should return earlier splitpoint if no valid ones are after threshhold', () => {
+        const history = [
+            { role: 'user', parts: [{ text: 'This is the first message.' }] },
+            { role: 'model', parts: [{ text: 'This is the second message.' }] },
+            { role: 'user', parts: [{ text: 'This is the third message.' }] },
+            { role: 'model', parts: [{ functionCall: {} }] },
+        ];
+        // Can't return 4 because the previous item has a function call.
+        expect(findCompressSplitPoint(history, 0.99)).toBe(2);
     });
     it('should handle a history with only one item', () => {
-        expect(findIndexAfterFraction(history.slice(0, 1), 0.5)).toBe(1);
+        const historyWithEmptyParts = [
+            { role: 'user', parts: [{ text: 'Message 1' }] },
+        ];
+        expect(findCompressSplitPoint(historyWithEmptyParts, 0.5)).toBe(0);
     });
     it('should handle history with weird parts', () => {
         const historyWithEmptyParts = [
@@ -131,7 +156,7 @@ describe('findIndexAfterFraction', () => {
             { role: 'model', parts: [{ fileData: { fileUri: 'derp' } }] },
             { role: 'user', parts: [{ text: 'Message 2' }] },
         ];
-        expect(findIndexAfterFraction(historyWithEmptyParts, 0.5)).toBe(2);
+        expect(findCompressSplitPoint(historyWithEmptyParts, 0.5)).toBe(2);
     });
 });
 describe('isThinkingSupported', () => {
@@ -168,6 +193,7 @@ describe('Gemini Client (client.ts)', () => {
     let mockGenerateContentFn;
     beforeEach(async () => {
         vi.resetAllMocks();
+        vi.mocked(uiTelemetryService.setLastPromptTokenCount).mockClear();
         mockGenerateContentFn = vi.fn().mockResolvedValue({
             candidates: [{ content: { parts: [{ text: '{"key": "value"}' }] } }],
         });
@@ -176,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)
@@ -229,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'),
@@ -299,72 +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,
                 });
+                // 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,
@@ -378,43 +460,80 @@ 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);
         });
+        it('should return NOOP if history is too short to compress', async () => {
+            const { client } = setup({
+                chatHistory: [{ role: 'user', parts: [{ text: 'hi' }] }],
+                originalTokenCount: 50,
+            });
+            const result = await client.tryCompressChat('prompt-id-noop', false);
+            expect(result).toEqual({
+                compressionStatus: CompressionStatus.NOOP,
+                originalTokenCount: 50,
+                newTokenCount: 50,
+            });
+            expect(mockGenerateContentFn).not.toHaveBeenCalled();
+        });
         it('logs a telemetry event when compressing', async () => {
             vi.spyOn(ClearcutLogger.prototype, 'logChatCompressionEvent');
             const MOCKED_TOKEN_LIMIT = 1000;
             const MOCKED_CONTEXT_PERCENTAGE_THRESHOLD = 0.5;
-            vi.mocked(tokenLimit).mockReturnValue(MOCKED_TOKEN_LIMIT);
             vi.spyOn(client['config'], 'getChatCompression').mockReturnValue({
                 contextPercentageThreshold: MOCKED_CONTEXT_PERCENTAGE_THRESHOLD,
             });
-            mockGetHistory.mockReturnValue([
+            const history = [
                 { role: 'user', parts: [{ text: '...history...' }] },
-            ]);
+                { role: 'model', parts: [{ text: '...history...' }] },
+                { role: 'user', parts: [{ text: '...history...' }] },
+                { role: 'model', parts: [{ text: '...history...' }] },
+                { role: 'user', parts: [{ text: '...history...' }] },
+                { role: 'model', 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 }],
                         },
                     },
                 ],
@@ -424,6 +543,8 @@ describe('Gemini Client (client.ts)', () => {
                 tokens_before: originalTokenCount,
                 tokens_after: newTokenCount,
             }));
+            expect(uiTelemetryService.setLastPromptTokenCount).toHaveBeenCalledWith(newTokenCount);
+            expect(uiTelemetryService.setLastPromptTokenCount).toHaveBeenCalledTimes(1);
         });
         it('should trigger summarization if token count is at threshold with contextPercentageThreshold setting', async () => {
             const MOCKED_TOKEN_LIMIT = 1000;
@@ -432,21 +553,46 @@ describe('Gemini Client (client.ts)', () => {
             vi.spyOn(client['config'], 'getChatCompression').mockReturnValue({
                 contextPercentageThreshold: MOCKED_CONTEXT_PERCENTAGE_THRESHOLD,
             });
-            mockGetHistory.mockReturnValue([
+            const history = [
                 { role: 'user', parts: [{ text: '...history...' }] },
-            ]);
+                { role: 'model', parts: [{ text: '...history...' }] },
+                { role: 'user', parts: [{ text: '...history...' }] },
+                { role: 'model', parts: [{ text: '...history...' }] },
+                { role: 'user', parts: [{ text: '...history...' }] },
+                { role: 'model', 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 }],
                         },
                     },
                 ],
@@ -468,7 +614,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...' }] },
@@ -485,19 +631,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 }],
                         },
                     },
                 ],
@@ -515,35 +685,60 @@ 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([
+            const history = [
                 { role: 'user', parts: [{ text: '...history...' }] },
-            ]);
-            const originalTokenCount = 10; // Well below threshold
-            const newTokenCount = 5;
-            vi.mocked(mockContentGenerator.countTokens)
-                .mockResolvedValueOnce({ totalTokens: originalTokenCount })
-                .mockResolvedValueOnce({ totalTokens: newTokenCount });
+                { role: 'model', parts: [{ text: '...history...' }] },
+                { role: 'user', parts: [{ text: '...history...' }] },
+                { role: 'model', parts: [{ text: '...history...' }] },
+                { role: 'user', parts: [{ text: '...history...' }] },
+                { role: 'model', 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({
@@ -581,9 +776,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* () {
@@ -841,6 +1033,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);
@@ -996,6 +1189,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(() => {
@@ -1081,6 +1350,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' };
@@ -1649,7 +2001,7 @@ ${JSON.stringify({
                 model: DEFAULT_GEMINI_FLASH_MODEL,
                 config: {
                     abortSignal,
-                    systemInstruction: getCoreSystemPrompt(''),
+                    systemInstruction: getCoreSystemPrompt({}, ''),
                     temperature: 0.5,
                     topP: 1,
                 },