summaryrefslogtreecommitdiff
path: root/packages/core/src
diff options
context:
space:
mode:
Diffstat (limited to 'packages/core/src')
-rw-r--r--packages/core/src/config/config.test.ts88
-rw-r--r--packages/core/src/config/config.ts10
-rw-r--r--packages/core/src/core/client.test.ts69
-rw-r--r--packages/core/src/core/client.ts28
4 files changed, 192 insertions, 3 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
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;
}