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
|
// Verify that fallback mode is reset
|
||||||
expect(config.isInFallbackMode()).toBe(false);
|
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', () => {
|
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 { shouldAttemptBrowserLaunch } from '../utils/browser.js';
|
||||||
import { MCPOAuthConfig } from '../mcp/oauth-provider.js';
|
import { MCPOAuthConfig } from '../mcp/oauth-provider.js';
|
||||||
import { IdeClient } from '../ide/ide-client.js';
|
import { IdeClient } from '../ide/ide-client.js';
|
||||||
|
import type { Content } from '@google/genai';
|
||||||
|
|
||||||
// Re-export OAuth config type
|
// Re-export OAuth config type
|
||||||
export type { MCPOAuthConfig };
|
export type { MCPOAuthConfig };
|
||||||
|
@ -332,13 +333,30 @@ export class Config {
|
||||||
}
|
}
|
||||||
|
|
||||||
async refreshAuth(authMethod: AuthType) {
|
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,
|
this,
|
||||||
authMethod,
|
authMethod,
|
||||||
);
|
);
|
||||||
|
|
||||||
this.geminiClient = new GeminiClient(this);
|
// Create and initialize new client in local variable first
|
||||||
await this.geminiClient.initialize(this.contentGeneratorConfig);
|
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
|
// Reset the session flag since we're explicitly changing auth and using default model
|
||||||
this.inFallbackMode = false;
|
this.inFallbackMode = false;
|
||||||
|
|
Loading…
Reference in New Issue