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