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

3 files changed, 24 insertions, 66 deletions

diff --git a/packages/core/src/config/config.test.ts b/packages/core/src/config/config.test.ts
index 3f0b3db5..f2169790 100644
--- a/packages/core/src/config/config.test.ts
+++ b/packages/core/src/config/config.test.ts
@@ -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);
     });
   });
 
diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts
index 7ccfdbc8..ee9067c6 100644
--- a/packages/core/src/config/config.ts
+++ b/packages/core/src/config/config.ts
@@ -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 {
diff --git a/packages/core/src/config/flashFallback.test.ts b/packages/core/src/config/flashFallback.test.ts
index 64f0f6fd..cd78cd34 100644
--- a/packages/core/src/config/flashFallback.test.ts
+++ b/packages/core/src/config/flashFallback.test.ts
@@ -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);
     });
   });
 });