Fix: Preserve conversation history when changing auth methods via /auth (#5216)
Co-authored-by: Jacob Richman <jacob314@gmail.com>
This commit is contained in:
parent
acd48a1259
commit
a8984a9b30
|
@ -184,6 +184,87 @@ describe('Server Config (config.ts)', () => {
|
|||
// Verify that fallback mode is reset
|
||||
expect(config.isInFallbackMode()).toBe(false);
|
||||
});
|
||||
|
||||
it('should preserve conversation history when refreshing auth', async () => {
|
||||
const config = new Config(baseParams);
|
||||
const authType = AuthType.USE_GEMINI;
|
||||
const mockContentConfig = {
|
||||
model: 'gemini-pro',
|
||||
apiKey: 'test-key',
|
||||
};
|
||||
|
||||
(createContentGeneratorConfig as Mock).mockReturnValue(mockContentConfig);
|
||||
|
||||
// Mock the existing client with some history
|
||||
const mockExistingHistory = [
|
||||
{ role: 'user', parts: [{ text: 'Hello' }] },
|
||||
{ role: 'model', parts: [{ text: 'Hi there!' }] },
|
||||
{ role: 'user', parts: [{ text: 'How are you?' }] },
|
||||
];
|
||||
|
||||
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),
|
||||
};
|
||||
|
||||
// Set the existing client
|
||||
(
|
||||
config as unknown as { geminiClient: typeof mockExistingClient }
|
||||
).geminiClient = mockExistingClient;
|
||||
(GeminiClient as Mock).mockImplementation(() => mockNewClient);
|
||||
|
||||
await config.refreshAuth(authType);
|
||||
|
||||
// Verify that existing history was retrieved
|
||||
expect(mockExistingClient.getHistory).toHaveBeenCalled();
|
||||
|
||||
// Verify that new client was created and initialized
|
||||
expect(GeminiClient).toHaveBeenCalledWith(config);
|
||||
expect(mockNewClient.initialize).toHaveBeenCalledWith(mockContentConfig);
|
||||
|
||||
// Verify that history was restored to the new client
|
||||
expect(mockNewClient.setHistory).toHaveBeenCalledWith(
|
||||
mockExistingHistory,
|
||||
);
|
||||
});
|
||||
|
||||
it('should handle case when no existing client is initialized', async () => {
|
||||
const config = new Config(baseParams);
|
||||
const authType = AuthType.USE_GEMINI;
|
||||
const mockContentConfig = {
|
||||
model: 'gemini-pro',
|
||||
apiKey: 'test-key',
|
||||
};
|
||||
|
||||
(createContentGeneratorConfig as Mock).mockReturnValue(mockContentConfig);
|
||||
|
||||
const mockNewClient = {
|
||||
isInitialized: vi.fn().mockReturnValue(true),
|
||||
getHistory: vi.fn().mockReturnValue([]),
|
||||
setHistory: vi.fn(),
|
||||
initialize: vi.fn().mockResolvedValue(undefined),
|
||||
};
|
||||
|
||||
// No existing client
|
||||
(config as unknown as { geminiClient: null }).geminiClient = null;
|
||||
(GeminiClient as Mock).mockImplementation(() => mockNewClient);
|
||||
|
||||
await config.refreshAuth(authType);
|
||||
|
||||
// Verify that new client was created and initialized
|
||||
expect(GeminiClient).toHaveBeenCalledWith(config);
|
||||
expect(mockNewClient.initialize).toHaveBeenCalledWith(mockContentConfig);
|
||||
|
||||
// Verify that setHistory was not called since there was no existing history
|
||||
expect(mockNewClient.setHistory).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
|
||||
it('Config constructor should store userMemory correctly', () => {
|
||||
|
|
|
@ -47,6 +47,7 @@ import { ClearcutLogger } from '../telemetry/clearcut-logger/clearcut-logger.js'
|
|||
import { shouldAttemptBrowserLaunch } from '../utils/browser.js';
|
||||
import { MCPOAuthConfig } from '../mcp/oauth-provider.js';
|
||||
import { IdeClient } from '../ide/ide-client.js';
|
||||
import type { Content } from '@google/genai';
|
||||
|
||||
// Re-export OAuth config type
|
||||
export type { MCPOAuthConfig };
|
||||
|
@ -332,13 +333,30 @@ export class Config {
|
|||
}
|
||||
|
||||
async refreshAuth(authMethod: AuthType) {
|
||||
this.contentGeneratorConfig = createContentGeneratorConfig(
|
||||
// Save the current conversation history before creating a new client
|
||||
let existingHistory: Content[] = [];
|
||||
if (this.geminiClient && this.geminiClient.isInitialized()) {
|
||||
existingHistory = this.geminiClient.getHistory();
|
||||
}
|
||||
|
||||
// Create new content generator config
|
||||
const newContentGeneratorConfig = createContentGeneratorConfig(
|
||||
this,
|
||||
authMethod,
|
||||
);
|
||||
|
||||
this.geminiClient = new GeminiClient(this);
|
||||
await this.geminiClient.initialize(this.contentGeneratorConfig);
|
||||
// Create and initialize new client in local variable first
|
||||
const newGeminiClient = new GeminiClient(this);
|
||||
await newGeminiClient.initialize(newContentGeneratorConfig);
|
||||
|
||||
// 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);
|
||||
}
|
||||
|
||||
// Reset the session flag since we're explicitly changing auth and using default model
|
||||
this.inFallbackMode = false;
|
||||
|
|
Loading…
Reference in New Issue