diff --git a/packages/cli/src/ui/hooks/useQuotaAndFallback.test.ts b/packages/cli/src/ui/hooks/useQuotaAndFallback.test.ts index 60064957e8..5ea4e2a84e 100644 --- a/packages/cli/src/ui/hooks/useQuotaAndFallback.test.ts +++ b/packages/cli/src/ui/hooks/useQuotaAndFallback.test.ts @@ -70,6 +70,7 @@ describe('useQuotaAndFallback', () => { setFallbackHandlerSpy = vi.spyOn(mockConfig, 'setFallbackModelHandler'); vi.spyOn(mockConfig, 'setQuotaErrorOccurred'); vi.spyOn(mockConfig, 'setModel'); + vi.spyOn(mockConfig, 'setActiveModel'); }); afterEach(() => { @@ -164,8 +165,8 @@ describe('useQuotaAndFallback', () => { const intent = await promise!; expect(intent).toBe('retry_always'); - // Verify setModel was called with isFallbackModel=true - expect(mockConfig.setModel).toHaveBeenCalledWith('gemini-flash', true); + // Verify setActiveModel was called + expect(mockConfig.setActiveModel).toHaveBeenCalledWith('gemini-flash'); // The pending request should be cleared from the state expect(result.current.proQuotaRequest).toBeNull(); @@ -278,8 +279,8 @@ describe('useQuotaAndFallback', () => { const intent = await promise!; expect(intent).toBe('retry_always'); - // Verify setModel was called with isFallbackModel=true - expect(mockConfig.setModel).toHaveBeenCalledWith('model-B', true); + // Verify setActiveModel was called + expect(mockConfig.setActiveModel).toHaveBeenCalledWith('model-B'); // The pending request should be cleared from the state expect(result.current.proQuotaRequest).toBeNull(); @@ -336,10 +337,9 @@ To disable gemini-3-pro-preview, disable "Preview features" in /settings.`, const intent = await promise!; expect(intent).toBe('retry_always'); - // Verify setModel was called with isFallbackModel=true - expect(mockConfig.setModel).toHaveBeenCalledWith( + // Verify setActiveModel was called + expect(mockConfig.setActiveModel).toHaveBeenCalledWith( 'gemini-2.5-pro', - true, ); expect(result.current.proQuotaRequest).toBeNull(); @@ -425,8 +425,12 @@ To disable gemini-3-pro-preview, disable "Preview features" in /settings.`, expect(intent).toBe('retry_always'); expect(result.current.proQuotaRequest).toBeNull(); - // Verify setModel was called with isFallbackModel=true - expect(mockConfig.setModel).toHaveBeenCalledWith('gemini-flash', true); + // Verify setActiveModel was called + expect(mockConfig.setActiveModel).toHaveBeenCalledWith('gemini-flash'); + + // Verify quota error flags are reset + expect(mockSetModelSwitchedFromQuotaError).toHaveBeenCalledWith(false); + expect(mockConfig.setQuotaErrorOccurred).toHaveBeenCalledWith(false); // Check for the "Switched to fallback model" message expect(mockHistoryManager.addItem).toHaveBeenCalledTimes(1); diff --git a/packages/cli/src/ui/hooks/useQuotaAndFallback.ts b/packages/cli/src/ui/hooks/useQuotaAndFallback.ts index eb53513fdc..d83b082950 100644 --- a/packages/cli/src/ui/hooks/useQuotaAndFallback.ts +++ b/packages/cli/src/ui/hooks/useQuotaAndFallback.ts @@ -129,21 +129,27 @@ export function useQuotaAndFallback({ setProQuotaRequest(null); isDialogPending.current = false; // Reset the flag here - if (choice === 'retry_always') { - // Set the model to the fallback model for the current session. - // This ensures the Footer updates and future turns use this model. - // The change is not persisted, so the original model is restored on restart. - config.activateFallbackMode(proQuotaRequest.fallbackModel); - historyManager.addItem( - { - type: MessageType.INFO, - text: `Switched to fallback model ${proQuotaRequest.fallbackModel}`, - }, - Date.now(), - ); + if (choice === 'retry_always' || choice === 'retry_once') { + // Reset quota error flags to allow the agent loop to continue. + setModelSwitchedFromQuotaError(false); + config.setQuotaErrorOccurred(false); + + if (choice === 'retry_always') { + // Set the model to the fallback model for the current session. + // This ensures the Footer updates and future turns use this model. + // The change is not persisted, so the original model is restored on restart. + config.activateFallbackMode(proQuotaRequest.fallbackModel); + historyManager.addItem( + { + type: MessageType.INFO, + text: `Switched to fallback model ${proQuotaRequest.fallbackModel}`, + }, + Date.now(), + ); + } } }, - [proQuotaRequest, historyManager, config], + [proQuotaRequest, historyManager, config, setModelSwitchedFromQuotaError], ); return { diff --git a/packages/core/src/availability/fallbackIntegration.test.ts b/packages/core/src/availability/fallbackIntegration.test.ts new file mode 100644 index 0000000000..39cbe2e0b4 --- /dev/null +++ b/packages/core/src/availability/fallbackIntegration.test.ts @@ -0,0 +1,78 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, it, expect, vi, beforeEach } from 'vitest'; +import { applyModelSelection } from './policyHelpers.js'; +import type { Config } from '../config/config.js'; +import { + PREVIEW_GEMINI_MODEL, + PREVIEW_GEMINI_FLASH_MODEL, + PREVIEW_GEMINI_MODEL_AUTO, +} from '../config/models.js'; +import { ModelAvailabilityService } from './modelAvailabilityService.js'; +import { ModelConfigService } from '../services/modelConfigService.js'; +import { DEFAULT_MODEL_CONFIGS } from '../config/defaultModelConfigs.js'; + +describe('Fallback Integration', () => { + let config: Config; + let availabilityService: ModelAvailabilityService; + let modelConfigService: ModelConfigService; + + beforeEach(() => { + // Mocking Config because it has many dependencies + config = { + getModel: () => PREVIEW_GEMINI_MODEL_AUTO, + getActiveModel: () => PREVIEW_GEMINI_MODEL_AUTO, + setActiveModel: vi.fn(), + getPreviewFeatures: () => true, // Preview enabled for Gemini 3 + getUserTier: () => undefined, + getModelAvailabilityService: () => availabilityService, + modelConfigService: undefined as unknown as ModelConfigService, + } as unknown as Config; + + availabilityService = new ModelAvailabilityService(); + modelConfigService = new ModelConfigService(DEFAULT_MODEL_CONFIGS); + // eslint-disable-next-line @typescript-eslint/no-explicit-any + (config as any).modelConfigService = modelConfigService; + }); + + it('should select fallback model when primary model is terminal and config is in AUTO mode', () => { + // 1. Simulate "Pro" failing with a terminal quota error + // The policy chain for PREVIEW_GEMINI_MODEL_AUTO is [PREVIEW_GEMINI_MODEL, PREVIEW_GEMINI_FLASH_MODEL] + availabilityService.markTerminal(PREVIEW_GEMINI_MODEL, 'quota'); + + // 2. Request "Pro" explicitly (as Agent would) + const requestedModel = PREVIEW_GEMINI_MODEL; + + // 3. Apply model selection + const result = applyModelSelection(config, { model: requestedModel }); + + // 4. Expect fallback to Flash + expect(result.model).toBe(PREVIEW_GEMINI_FLASH_MODEL); + + // 5. Expect active model to be updated + expect(config.setActiveModel).toHaveBeenCalledWith( + PREVIEW_GEMINI_FLASH_MODEL, + ); + }); + + it('should NOT fallback if config is NOT in AUTO mode', () => { + // 1. Config is explicitly set to Pro, not Auto + vi.spyOn(config, 'getModel').mockReturnValue(PREVIEW_GEMINI_MODEL); + + // 2. Simulate "Pro" failing + availabilityService.markTerminal(PREVIEW_GEMINI_MODEL, 'quota'); + + // 3. Request "Pro" + const requestedModel = PREVIEW_GEMINI_MODEL; + + // 4. Apply model selection + const result = applyModelSelection(config, { model: requestedModel }); + + // 5. Expect it to stay on Pro (because single model chain) + expect(result.model).toBe(PREVIEW_GEMINI_MODEL); + }); +}); diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index f877a2e797..6fb07754ec 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -939,7 +939,8 @@ export class Config { } activateFallbackMode(model: string): void { - this.setModel(model, true); + this.setActiveModel(model); + coreEvents.emitModelChanged(model); const authType = this.getContentGeneratorConfig()?.authType; if (authType) { logFlashFallback(this, new FlashFallbackEvent(authType)); diff --git a/packages/core/src/config/flashFallback.test.ts b/packages/core/src/config/flashFallback.test.ts index 320d69c565..96adf37655 100644 --- a/packages/core/src/config/flashFallback.test.ts +++ b/packages/core/src/config/flashFallback.test.ts @@ -65,9 +65,12 @@ describe('Flash Model Fallback Configuration', () => { }); describe('activateFallbackMode', () => { - it('should set model to fallback and log event', () => { + it('should set active model to fallback and log event', () => { config.activateFallbackMode(DEFAULT_GEMINI_FLASH_MODEL); - expect(config.getModel()).toBe(DEFAULT_GEMINI_FLASH_MODEL); + expect(config.getActiveModel()).toBe(DEFAULT_GEMINI_FLASH_MODEL); + // Ensure the persisted model setting is NOT changed (to preserve AUTO behavior) + expect(config.getModel()).toBe(DEFAULT_GEMINI_MODEL); + expect(logFlashFallback).toHaveBeenCalledWith( config, expect.any(FlashFallbackEvent), diff --git a/packages/core/src/core/loggingContentGenerator.ts b/packages/core/src/core/loggingContentGenerator.ts index b8cf49a091..e559846366 100644 --- a/packages/core/src/core/loggingContentGenerator.ts +++ b/packages/core/src/core/loggingContentGenerator.ts @@ -195,6 +195,7 @@ export class LoggingContentGenerator implements ContentGenerator { req.config, serverDetails, ); + try { const response = await this.wrapped.generateContent( req,