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', () => { diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index 7c61f239..5c11667b 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -379,13 +379,21 @@ export class Config { const newGeminiClient = new GeminiClient(this); await newGeminiClient.initialize(newContentGeneratorConfig); + // Vertex and Genai have incompatible encryption and sending history with + // throughtSignature from Genai to Vertex will fail, we need to strip them + const fromGenaiToVertex = + this.contentGeneratorConfig?.authType === AuthType.USE_GEMINI && + authMethod === AuthType.LOGIN_WITH_GOOGLE; + // Only assign to instance properties after successful initialization this.contentGeneratorConfig = newContentGeneratorConfig; this.geminiClient = newGeminiClient; // Restore the conversation history to the new client if (existingHistory.length > 0) { - this.geminiClient.setHistory(existingHistory); + this.geminiClient.setHistory(existingHistory, { + stripThoughts: fromGenaiToVertex, + }); } // Reset the session flag since we're explicitly changing auth and using default model diff --git a/packages/core/src/core/client.test.ts b/packages/core/src/core/client.test.ts index 5e68cfb6..9f6dcbe9 100644 --- a/packages/core/src/core/client.test.ts +++ b/packages/core/src/core/client.test.ts @@ -1596,4 +1596,73 @@ ${JSON.stringify( ); }); }); + + describe('setHistory', () => { + it('should strip thought signatures when stripThoughts is true', () => { + const mockChat = { + setHistory: vi.fn(), + }; + client['chat'] = mockChat as unknown as GeminiChat; + + const historyWithThoughts: Content[] = [ + { + role: 'user', + parts: [{ text: 'hello' }], + }, + { + role: 'model', + parts: [ + { text: 'thinking...', thoughtSignature: 'thought-123' }, + { + functionCall: { name: 'test', args: {} }, + thoughtSignature: 'thought-456', + }, + ], + }, + ]; + + client.setHistory(historyWithThoughts, { stripThoughts: true }); + + const expectedHistory: Content[] = [ + { + role: 'user', + parts: [{ text: 'hello' }], + }, + { + role: 'model', + parts: [ + { text: 'thinking...' }, + { functionCall: { name: 'test', args: {} } }, + ], + }, + ]; + + expect(mockChat.setHistory).toHaveBeenCalledWith(expectedHistory); + }); + + it('should not strip thought signatures when stripThoughts is false', () => { + const mockChat = { + setHistory: vi.fn(), + }; + client['chat'] = mockChat as unknown as GeminiChat; + + const historyWithThoughts: Content[] = [ + { + role: 'user', + parts: [{ text: 'hello' }], + }, + { + role: 'model', + parts: [ + { text: 'thinking...', thoughtSignature: 'thought-123' }, + { text: 'ok', thoughtSignature: 'thought-456' }, + ], + }, + ]; + + client.setHistory(historyWithThoughts, { stripThoughts: false }); + + expect(mockChat.setHistory).toHaveBeenCalledWith(historyWithThoughts); + }); + }); }); diff --git a/packages/core/src/core/client.ts b/packages/core/src/core/client.ts index 96be4111..93de190d 100644 --- a/packages/core/src/core/client.ts +++ b/packages/core/src/core/client.ts @@ -164,8 +164,32 @@ export class GeminiClient { return this.getChat().getHistory(); } - setHistory(history: Content[]) { - this.getChat().setHistory(history); + setHistory( + history: Content[], + { stripThoughts = false }: { stripThoughts?: boolean } = {}, + ) { + const historyToSet = stripThoughts + ? history.map((content) => { + const newContent = { ...content }; + if (newContent.parts) { + newContent.parts = newContent.parts.map((part) => { + if ( + part && + typeof part === 'object' && + 'thoughtSignature' in part + ) { + const newPart = { ...part }; + delete (newPart as { thoughtSignature?: string }) + .thoughtSignature; + return newPart; + } + return part; + }); + } + return newContent; + }) + : history; + this.getChat().setHistory(historyToSet); this.forceFullIdeContext = true; }