fix(core): Discard thought signature when switching from Gemini API Key to OAuth (#6090)

Co-authored-by: Jacob Richman <[email protected]>

2 files changed, 97 insertions, 1 deletions

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