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

Co-authored-by: Jacob Richman <jacob314@gmail.com>
This commit is contained in:
Sandy Tao
2025-08-13 13:33:04 -07:00
committed by GitHub
parent a90aeb3d8f
commit 61047173a8
4 changed files with 192 additions and 3 deletions

View File

@@ -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', () => {

View File

@@ -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