@machina.ai/cell-cli-core 1.6.1-rc1 → 1.8.2-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 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 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 an empty history', () => {
-        expect(findIndexAfterFraction([], 0.5)).toBe(0);
+    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 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)
@@ -299,72 +324,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,17 +459,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);
         });
@@ -400,21 +480,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 }],
                         },
                     },
                 ],
@@ -424,6 +523,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 +533,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 }],
                         },
                     },
                 ],
@@ -468,7 +587,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 +604,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 +658,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({
@@ -581,9 +742,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* () {