Add new fallback state as prefactor for routing (#5065)
This commit is contained in:
parent
b6c2c64f9b
commit
b08679c906
|
@ -398,6 +398,7 @@ const App = ({ config, settings, startupWarnings = [], version }: AppProps) => {
|
|||
|
||||
// Switch model for future use but return false to stop current retry
|
||||
config.setModel(fallbackModel);
|
||||
config.setFallbackMode(true);
|
||||
logFlashFallback(
|
||||
config,
|
||||
new FlashFallbackEvent(config.getContentGeneratorConfig().authType!),
|
||||
|
|
|
@ -152,6 +152,10 @@ describe('Server Config (config.ts)', () => {
|
|||
|
||||
(createContentGeneratorConfig as Mock).mockReturnValue(mockContentConfig);
|
||||
|
||||
// Set fallback mode to true to ensure it gets reset
|
||||
config.setFallbackMode(true);
|
||||
expect(config.isInFallbackMode()).toBe(true);
|
||||
|
||||
await config.refreshAuth(authType);
|
||||
|
||||
expect(createContentGeneratorConfig).toHaveBeenCalledWith(
|
||||
|
@ -163,6 +167,8 @@ describe('Server Config (config.ts)', () => {
|
|||
expect(config.getContentGeneratorConfig().model).toBe(newModel);
|
||||
expect(config.getModel()).toBe(newModel); // getModel() should return the updated model
|
||||
expect(GeminiClient).toHaveBeenCalledWith(config);
|
||||
// Verify that fallback mode is reset
|
||||
expect(config.isInFallbackMode()).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
|
|
|
@ -226,7 +226,7 @@ export class Config {
|
|||
private readonly noBrowser: boolean;
|
||||
private readonly ideMode: boolean;
|
||||
private readonly ideClient: IdeClient | undefined;
|
||||
private modelSwitchedDuringSession: boolean = false;
|
||||
private inFallbackMode = false;
|
||||
private readonly maxSessionTurns: number;
|
||||
private readonly listExtensions: boolean;
|
||||
private readonly _extensions: GeminiCLIExtension[];
|
||||
|
@ -330,7 +330,7 @@ export class Config {
|
|||
await this.geminiClient.initialize(this.contentGeneratorConfig);
|
||||
|
||||
// Reset the session flag since we're explicitly changing auth and using default model
|
||||
this.modelSwitchedDuringSession = false;
|
||||
this.inFallbackMode = false;
|
||||
}
|
||||
|
||||
getSessionId(): string {
|
||||
|
@ -348,19 +348,15 @@ export class Config {
|
|||
setModel(newModel: string): void {
|
||||
if (this.contentGeneratorConfig) {
|
||||
this.contentGeneratorConfig.model = newModel;
|
||||
this.modelSwitchedDuringSession = true;
|
||||
}
|
||||
}
|
||||
|
||||
isModelSwitchedDuringSession(): boolean {
|
||||
return this.modelSwitchedDuringSession;
|
||||
isInFallbackMode(): boolean {
|
||||
return this.inFallbackMode;
|
||||
}
|
||||
|
||||
resetModelToDefault(): void {
|
||||
if (this.contentGeneratorConfig) {
|
||||
this.contentGeneratorConfig.model = this.model; // Reset to the original default model
|
||||
this.modelSwitchedDuringSession = false;
|
||||
}
|
||||
setFallbackMode(active: boolean): void {
|
||||
this.inFallbackMode = active;
|
||||
}
|
||||
|
||||
setFlashFallbackHandler(handler: FlashFallbackHandler): void {
|
||||
|
|
|
@ -29,26 +29,11 @@ describe('Flash Model Fallback Configuration', () => {
|
|||
};
|
||||
});
|
||||
|
||||
// These tests do not actually test fallback. isInFallbackMode() only returns true,
|
||||
// when setFallbackMode is marked as true. This is to decouple setting a model
|
||||
// with the fallback mechanism. This will be necessary we introduce more
|
||||
// intelligent model routing.
|
||||
describe('setModel', () => {
|
||||
it('should update the model and mark as switched during session', () => {
|
||||
expect(config.getModel()).toBe(DEFAULT_GEMINI_MODEL);
|
||||
expect(config.isModelSwitchedDuringSession()).toBe(false);
|
||||
|
||||
config.setModel(DEFAULT_GEMINI_FLASH_MODEL);
|
||||
|
||||
expect(config.getModel()).toBe(DEFAULT_GEMINI_FLASH_MODEL);
|
||||
expect(config.isModelSwitchedDuringSession()).toBe(true);
|
||||
});
|
||||
|
||||
it('should handle multiple model switches during session', () => {
|
||||
config.setModel(DEFAULT_GEMINI_FLASH_MODEL);
|
||||
expect(config.isModelSwitchedDuringSession()).toBe(true);
|
||||
|
||||
config.setModel('gemini-1.5-pro');
|
||||
expect(config.getModel()).toBe('gemini-1.5-pro');
|
||||
expect(config.isModelSwitchedDuringSession()).toBe(true);
|
||||
});
|
||||
|
||||
it('should only mark as switched if contentGeneratorConfig exists', () => {
|
||||
// Create config without initializing contentGeneratorConfig
|
||||
const newConfig = new Config({
|
||||
|
@ -61,7 +46,7 @@ describe('Flash Model Fallback Configuration', () => {
|
|||
|
||||
// Should not crash when contentGeneratorConfig is undefined
|
||||
newConfig.setModel(DEFAULT_GEMINI_FLASH_MODEL);
|
||||
expect(newConfig.isModelSwitchedDuringSession()).toBe(false);
|
||||
expect(newConfig.isInFallbackMode()).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
|
@ -86,54 +71,25 @@ describe('Flash Model Fallback Configuration', () => {
|
|||
});
|
||||
});
|
||||
|
||||
describe('isModelSwitchedDuringSession', () => {
|
||||
describe('isInFallbackMode', () => {
|
||||
it('should start as false for new session', () => {
|
||||
expect(config.isModelSwitchedDuringSession()).toBe(false);
|
||||
expect(config.isInFallbackMode()).toBe(false);
|
||||
});
|
||||
|
||||
it('should remain false if no model switch occurs', () => {
|
||||
// Perform other operations that don't involve model switching
|
||||
expect(config.isModelSwitchedDuringSession()).toBe(false);
|
||||
expect(config.isInFallbackMode()).toBe(false);
|
||||
});
|
||||
|
||||
it('should persist switched state throughout session', () => {
|
||||
config.setModel(DEFAULT_GEMINI_FLASH_MODEL);
|
||||
expect(config.isModelSwitchedDuringSession()).toBe(true);
|
||||
// Setting state for fallback mode as is expected of clients
|
||||
config.setFallbackMode(true);
|
||||
expect(config.isInFallbackMode()).toBe(true);
|
||||
|
||||
// Should remain true even after getting model
|
||||
config.getModel();
|
||||
expect(config.isModelSwitchedDuringSession()).toBe(true);
|
||||
});
|
||||
});
|
||||
|
||||
describe('resetModelToDefault', () => {
|
||||
it('should reset model to default and clear session switch flag', () => {
|
||||
// Switch to Flash first
|
||||
config.setModel(DEFAULT_GEMINI_FLASH_MODEL);
|
||||
expect(config.getModel()).toBe(DEFAULT_GEMINI_FLASH_MODEL);
|
||||
expect(config.isModelSwitchedDuringSession()).toBe(true);
|
||||
|
||||
// Reset to default
|
||||
config.resetModelToDefault();
|
||||
|
||||
// Should be back to default with flag cleared
|
||||
expect(config.getModel()).toBe(DEFAULT_GEMINI_MODEL);
|
||||
expect(config.isModelSwitchedDuringSession()).toBe(false);
|
||||
});
|
||||
|
||||
it('should handle case where contentGeneratorConfig is not initialized', () => {
|
||||
// Create config without initializing contentGeneratorConfig
|
||||
const newConfig = new Config({
|
||||
sessionId: 'test-session-2',
|
||||
targetDir: '/test',
|
||||
debugMode: false,
|
||||
cwd: '/test',
|
||||
model: DEFAULT_GEMINI_MODEL,
|
||||
});
|
||||
|
||||
// Should not crash when contentGeneratorConfig is undefined
|
||||
expect(() => newConfig.resetModelToDefault()).not.toThrow();
|
||||
expect(newConfig.isModelSwitchedDuringSession()).toBe(false);
|
||||
expect(config.isInFallbackMode()).toBe(true);
|
||||
});
|
||||
});
|
||||
});
|
||||
|
|
|
@ -201,6 +201,7 @@ describe('Gemini Client (client.ts)', () => {
|
|||
getUsageStatisticsEnabled: vi.fn().mockReturnValue(true),
|
||||
getIdeMode: vi.fn().mockReturnValue(false),
|
||||
getGeminiClient: vi.fn(),
|
||||
setFallbackMode: vi.fn(),
|
||||
};
|
||||
const MockedConfig = vi.mocked(Config, true);
|
||||
MockedConfig.mockImplementation(
|
||||
|
@ -1262,7 +1263,8 @@ Here are some files the user has open, with the most recent at the top:
|
|||
|
||||
// mock config been changed
|
||||
const currentModel = initialModel + '-changed';
|
||||
vi.spyOn(client['config'], 'getModel').mockReturnValueOnce(currentModel);
|
||||
const getModelSpy = vi.spyOn(client['config'], 'getModel');
|
||||
getModelSpy.mockReturnValue(currentModel);
|
||||
|
||||
const mockFallbackHandler = vi.fn().mockResolvedValue(true);
|
||||
client['config'].flashFallbackHandler = mockFallbackHandler;
|
||||
|
|
|
@ -717,6 +717,7 @@ export class GeminiClient {
|
|||
);
|
||||
if (accepted !== false && accepted !== null) {
|
||||
this.config.setModel(fallbackModel);
|
||||
this.config.setFallbackMode(true);
|
||||
return fallbackModel;
|
||||
}
|
||||
// Check if the model was switched manually in the handler
|
||||
|
|
|
@ -225,6 +225,7 @@ export class GeminiChat {
|
|||
);
|
||||
if (accepted !== false && accepted !== null) {
|
||||
this.config.setModel(fallbackModel);
|
||||
this.config.setFallbackMode(true);
|
||||
return fallbackModel;
|
||||
}
|
||||
// Check if the model was switched manually in the handler
|
||||
|
|
Loading…
Reference in New Issue