Add new fallback state as prefactor for routing (#5065)

This commit is contained in:
Abhi 2025-07-28 15:55:50 -04:00 committed by GitHub
parent b6c2c64f9b
commit b08679c906
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
7 changed files with 30 additions and 67 deletions

View File

@ -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!),

View File

@ -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);
});
});

View File

@ -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 {

View File

@ -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);
});
});
});

View File

@ -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;

View File

@ -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

View File

@ -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