From c93eed6384a73489b46c3015a278c88a508c410d Mon Sep 17 00:00:00 2001
From: Abhi <43648792+abhipatel12@users.noreply.github.com>
Date: Tue, 23 Sep 2025 00:44:07 -0400
Subject: [PATCH] fix(core): Compression was broken when routing enabled
 (#9183)

---
 packages/core/src/core/client.test.ts | 49 ---------------------------
 packages/core/src/core/client.ts      |  3 +-
 2 files changed, 1 insertion(+), 51 deletions(-)

diff --git a/packages/core/src/core/client.test.ts b/packages/core/src/core/client.test.ts
index 6f0034e4b0..db897a6f2b 100644
--- a/packages/core/src/core/client.test.ts
+++ b/packages/core/src/core/client.test.ts
@@ -733,55 +733,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,
-      } as unknown as GeminiChat;
-
-      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', () => {
diff --git a/packages/core/src/core/client.ts b/packages/core/src/core/client.ts
index 450bd42f2f..5d3015195e 100644
--- a/packages/core/src/core/client.ts
+++ b/packages/core/src/core/client.ts
@@ -758,8 +758,7 @@ export class GeminiClient {
 
     const { totalTokens: newTokenCount } =
       await this.getContentGeneratorOrFail().countTokens({
-        // model might change after calling `sendMessage`, so we get the newest value from config
-        model: this.config.getModel(),
+        model,
         contents: chat.getHistory(),
       });
     if (newTokenCount === undefined) {