Fix an issue where the agent stops prematurely (#16269)

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

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 {

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

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

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),

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,