From 61047173a8f5bc279f480c5ab150d74337c0265a Mon Sep 17 00:00:00 2001 From: Sandy Tao Date: Wed, 13 Aug 2025 13:33:04 -0700 Subject: fix(core): Discard thought signature when switching from Gemini API Key to OAuth (#6090) Co-authored-by: Jacob Richman --- packages/core/src/config/config.test.ts | 88 +++++++++++++++++++++++++++++++++ 1 file changed, 88 insertions(+) (limited to 'packages/core/src/config/config.test.ts') diff --git a/packages/core/src/config/config.test.ts b/packages/core/src/config/config.test.ts index 6c57d058..f1d8b965 100644 --- a/packages/core/src/config/config.test.ts +++ b/packages/core/src/config/config.test.ts @@ -15,6 +15,7 @@ import { } from '../telemetry/index.js'; import { AuthType, + ContentGeneratorConfig, createContentGeneratorConfig, } from '../core/contentGenerator.js'; import { GeminiClient } from '../core/client.js'; @@ -249,6 +250,7 @@ describe('Server Config (config.ts)', () => { // Verify that history was restored to the new client expect(mockNewClient.setHistory).toHaveBeenCalledWith( mockExistingHistory, + { stripThoughts: false }, ); }); @@ -282,6 +284,92 @@ describe('Server Config (config.ts)', () => { // Verify that setHistory was not called since there was no existing history expect(mockNewClient.setHistory).not.toHaveBeenCalled(); }); + + it('should strip thoughts when switching from GenAI to Vertex', async () => { + const config = new Config(baseParams); + const mockContentConfig = { + model: 'gemini-pro', + apiKey: 'test-key', + authType: AuthType.USE_GEMINI, + }; + ( + config as unknown as { contentGeneratorConfig: ContentGeneratorConfig } + ).contentGeneratorConfig = mockContentConfig; + + (createContentGeneratorConfig as Mock).mockReturnValue({ + ...mockContentConfig, + authType: AuthType.LOGIN_WITH_GOOGLE, + }); + + const mockExistingHistory = [ + { role: 'user', parts: [{ text: 'Hello' }] }, + ]; + const mockExistingClient = { + isInitialized: vi.fn().mockReturnValue(true), + getHistory: vi.fn().mockReturnValue(mockExistingHistory), + }; + const mockNewClient = { + isInitialized: vi.fn().mockReturnValue(true), + getHistory: vi.fn().mockReturnValue([]), + setHistory: vi.fn(), + initialize: vi.fn().mockResolvedValue(undefined), + }; + + ( + config as unknown as { geminiClient: typeof mockExistingClient } + ).geminiClient = mockExistingClient; + (GeminiClient as Mock).mockImplementation(() => mockNewClient); + + await config.refreshAuth(AuthType.LOGIN_WITH_GOOGLE); + + expect(mockNewClient.setHistory).toHaveBeenCalledWith( + mockExistingHistory, + { stripThoughts: true }, + ); + }); + + it('should not strip thoughts when switching from Vertex to GenAI', async () => { + const config = new Config(baseParams); + const mockContentConfig = { + model: 'gemini-pro', + apiKey: 'test-key', + authType: AuthType.LOGIN_WITH_GOOGLE, + }; + ( + config as unknown as { contentGeneratorConfig: ContentGeneratorConfig } + ).contentGeneratorConfig = mockContentConfig; + + (createContentGeneratorConfig as Mock).mockReturnValue({ + ...mockContentConfig, + authType: AuthType.USE_GEMINI, + }); + + const mockExistingHistory = [ + { role: 'user', parts: [{ text: 'Hello' }] }, + ]; + const mockExistingClient = { + isInitialized: vi.fn().mockReturnValue(true), + getHistory: vi.fn().mockReturnValue(mockExistingHistory), + }; + const mockNewClient = { + isInitialized: vi.fn().mockReturnValue(true), + getHistory: vi.fn().mockReturnValue([]), + setHistory: vi.fn(), + initialize: vi.fn().mockResolvedValue(undefined), + }; + + ( + config as unknown as { geminiClient: typeof mockExistingClient } + ).geminiClient = mockExistingClient; + (GeminiClient as Mock).mockImplementation(() => mockNewClient); + + await config.refreshAuth(AuthType.USE_GEMINI); + + expect(mockNewClient.setHistory).toHaveBeenCalledWith( + mockExistingHistory, + { stripThoughts: false }, + ); + }); }); it('Config constructor should store userMemory correctly', () => { -- cgit v1.2.3