[JUNE 25] Permanent failover to Flash model for OAuth users after persistent 429 errors (#1376)

Co-authored-by: Scott Densmore <[email protected]>

2 files changed, 196 insertions, 2 deletions

diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts
index a92dd7ba..b266512c 100644
--- a/packages/core/src/config/config.ts
+++ b/packages/core/src/config/config.ts
@@ -35,7 +35,10 @@ import {
   TelemetryTarget,
   StartSessionEvent,
 } from '../telemetry/index.js';
-import { DEFAULT_GEMINI_EMBEDDING_MODEL } from './models.js';
+import {
+  DEFAULT_GEMINI_EMBEDDING_MODEL,
+  DEFAULT_GEMINI_FLASH_MODEL,
+} from './models.js';
 import { ClearcutLogger } from '../telemetry/clearcut-logger/clearcut-logger.js';
 
 export enum ApprovalMode {
@@ -85,6 +88,11 @@ export interface SandboxConfig {
   image: string;
 }
 
+export type FlashFallbackHandler = (
+  currentModel: string,
+  fallbackModel: string,
+) => Promise<boolean>;
+
 export interface ConfigParameters {
   sessionId: string;
   embeddingModel?: string;
@@ -156,6 +164,8 @@ export class Config {
   private readonly bugCommand: BugCommandSettings | undefined;
   private readonly model: string;
   private readonly extensionContextFilePaths: string[];
+  private modelSwitchedDuringSession: boolean = false;
+  flashFallbackHandler?: FlashFallbackHandler;
 
   constructor(params: ConfigParameters) {
     this.sessionId = params.sessionId;
@@ -216,9 +226,24 @@ export class Config {
   }
 
   async refreshAuth(authMethod: AuthType) {
+    // Check if this is actually a switch to a different auth method
+    const previousAuthType = this.contentGeneratorConfig?.authType;
+    const _isAuthMethodSwitch =
+      previousAuthType && previousAuthType !== authMethod;
+
+    // Always use the original default model when switching auth methods
+    // This ensures users don't stay on Flash after switching between auth types
+    // and allows API key users to get proper fallback behavior from getEffectiveModel
+    const modelToUse = this.model; // Use the original default model
+
+    // Temporarily clear contentGeneratorConfig to prevent getModel() from returning
+    // the previous session's model (which might be Flash)
+    this.contentGeneratorConfig = undefined!;
+
     const contentConfig = await createContentGeneratorConfig(
-      this.getModel(),
+      modelToUse,
       authMethod,
+      this,
     );
 
     const gc = new GeminiClient(this);
@@ -226,6 +251,11 @@ export class Config {
     this.toolRegistry = await createToolRegistry(this);
     await gc.initialize(contentConfig);
     this.contentGeneratorConfig = contentConfig;
+
+    // Reset the session flag since we're explicitly changing auth and using default model
+    this.modelSwitchedDuringSession = false;
+
+    // Note: In the future, we may want to reset any cached state when switching auth methods
   }
 
   getSessionId(): string {
@@ -240,6 +270,28 @@ export class Config {
     return this.contentGeneratorConfig?.model || this.model;
   }
 
+  setModel(newModel: string): void {
+    if (this.contentGeneratorConfig) {
+      this.contentGeneratorConfig.model = newModel;
+      this.modelSwitchedDuringSession = true;
+    }
+  }
+
+  isModelSwitchedDuringSession(): boolean {
+    return this.modelSwitchedDuringSession;
+  }
+
+  resetModelToDefault(): void {
+    if (this.contentGeneratorConfig) {
+      this.contentGeneratorConfig.model = this.model; // Reset to the original default model
+      this.modelSwitchedDuringSession = false;
+    }
+  }
+
+  setFlashFallbackHandler(handler: FlashFallbackHandler): void {
+    this.flashFallbackHandler = handler;
+  }
+
   getEmbeddingModel(): string {
     return this.embeddingModel;
   }
@@ -445,3 +497,6 @@ export function createToolRegistry(config: Config): Promise<ToolRegistry> {
     return registry;
   })();
 }
+
+// Export model constants for use in CLI
+export { DEFAULT_GEMINI_FLASH_MODEL };
diff --git a/packages/core/src/config/flashFallback.test.ts b/packages/core/src/config/flashFallback.test.ts
new file mode 100644
index 00000000..325cc064
--- /dev/null
+++ b/packages/core/src/config/flashFallback.test.ts
@@ -0,0 +1,139 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { describe, it, expect, beforeEach } from 'vitest';
+import { Config } from './config.js';
+import { DEFAULT_GEMINI_MODEL, DEFAULT_GEMINI_FLASH_MODEL } from './models.js';
+
+describe('Flash Model Fallback Configuration', () => {
+  let config: Config;
+
+  beforeEach(() => {
+    config = new Config({
+      sessionId: 'test-session',
+      targetDir: '/test',
+      debugMode: false,
+      cwd: '/test',
+      model: DEFAULT_GEMINI_MODEL,
+    });
+
+    // Initialize contentGeneratorConfig for testing
+    (
+      config as unknown as { contentGeneratorConfig: unknown }
+    ).contentGeneratorConfig = {
+      model: DEFAULT_GEMINI_MODEL,
+      authType: 'oauth-personal',
+    };
+  });
+
+  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({
+        sessionId: 'test-session-2',
+        targetDir: '/test',
+        debugMode: false,
+        cwd: '/test',
+        model: DEFAULT_GEMINI_MODEL,
+      });
+
+      // Should not crash when contentGeneratorConfig is undefined
+      newConfig.setModel(DEFAULT_GEMINI_FLASH_MODEL);
+      expect(newConfig.isModelSwitchedDuringSession()).toBe(false);
+    });
+  });
+
+  describe('getModel', () => {
+    it('should return contentGeneratorConfig model if available', () => {
+      // Simulate initialized content generator config
+      config.setModel(DEFAULT_GEMINI_FLASH_MODEL);
+      expect(config.getModel()).toBe(DEFAULT_GEMINI_FLASH_MODEL);
+    });
+
+    it('should fallback to initial model if contentGeneratorConfig is not available', () => {
+      // Test with fresh config where contentGeneratorConfig might not be set
+      const newConfig = new Config({
+        sessionId: 'test-session-2',
+        targetDir: '/test',
+        debugMode: false,
+        cwd: '/test',
+        model: 'custom-model',
+      });
+
+      expect(newConfig.getModel()).toBe('custom-model');
+    });
+  });
+
+  describe('isModelSwitchedDuringSession', () => {
+    it('should start as false for new session', () => {
+      expect(config.isModelSwitchedDuringSession()).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);
+    });
+
+    it('should persist switched state throughout session', () => {
+      config.setModel(DEFAULT_GEMINI_FLASH_MODEL);
+      expect(config.isModelSwitchedDuringSession()).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);
+    });
+  });
+});