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

1 files changed, 12 insertions, 56 deletions

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