6 files changed, 29 insertions, 67 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);
     });
   });
 });
diff --git a/packages/core/src/core/client.test.ts b/packages/core/src/core/client.test.ts
index 8c46d7f5..1da355f4 100644
--- a/packages/core/src/core/client.test.ts
+++ b/packages/core/src/core/client.test.ts
@@ -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;
diff --git a/packages/core/src/core/client.ts b/packages/core/src/core/client.ts
index e58e7040..6337bac2 100644
--- a/packages/core/src/core/client.ts
+++ b/packages/core/src/core/client.ts
@@ -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
diff --git a/packages/core/src/core/geminiChat.ts b/packages/core/src/core/geminiChat.ts
index 4c3cd4c8..d3b2e060 100644
--- a/packages/core/src/core/geminiChat.ts
+++ b/packages/core/src/core/geminiChat.ts
@@ -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