fix: respect explicit model selection after Flash quota exhaustion (#26759) (#26872)

packages/core/src/availability/modelAvailabilityService.test.ts

@@ -168,4 +168,46 @@ describe('ModelAvailabilityService', () => {
       reason: 'quota',
     });
   });
+
+  describe('prefix normalization', () => {
+    it('treats prefixed and non-prefixed models as identical when marking terminal', () => {
+      service.markTerminal('models/gemini-3.1-pro-preview', 'quota');
+
+      // Checking the non-prefixed version should show it as unavailable
+      expect(service.snapshot('gemini-3.1-pro-preview')).toEqual({
+        available: false,
+        reason: 'quota',
+      });
+
+      // Checking the prefixed version should also show it as unavailable
+      expect(service.snapshot('models/gemini-3.1-pro-preview')).toEqual({
+        available: false,
+        reason: 'quota',
+      });
+    });
+
+    it('treats prefixed and non-prefixed models as identical when selecting', () => {
+      service.markTerminal('gemini-3-flash-preview', 'quota');
+
+      // Attempting to select the prefixed version should skip it because the base is exhausted
+      const result = service.selectFirstAvailable([
+        'models/gemini-3-flash-preview',
+        'gemini-3.1-pro-preview',
+      ]);
+
+      expect(result.selectedModel).toBe('gemini-3.1-pro-preview');
+      expect(result.skipped).toEqual([
+        { model: 'gemini-3-flash-preview', reason: 'quota' },
+      ]);
+    });
+
+    it('treats prefixed and non-prefixed models as identical when marking healthy', () => {
+      service.markTerminal('gemini-3-flash-preview', 'quota');
+      service.markHealthy('models/gemini-3-flash-preview');
+
+      expect(service.snapshot('gemini-3-flash-preview')).toEqual({
+        available: true,
+      });
+    });
+  });
 });

packages/core/src/availability/modelAvailabilityService.ts

@@ -39,21 +39,26 @@ export interface ModelSelectionResult {
   }>;
 }
 
+import { normalizeModelId } from '../utils/modelUtils.js';
+
 export class ModelAvailabilityService {
   private readonly health = new Map<ModelId, HealthState>();
 
-  markTerminal(model: ModelId, reason: TerminalUnavailabilityReason) {
+  markTerminal(modelId: ModelId, reason: TerminalUnavailabilityReason) {
+    const model = normalizeModelId(modelId);
     this.setState(model, {
       status: 'terminal',
       reason,
     });
   }
 
-  markHealthy(model: ModelId) {
+  markHealthy(modelId: ModelId) {
+    const model = normalizeModelId(modelId);
     this.clearState(model);
   }
 
-  markRetryOncePerTurn(model: ModelId, attempts: number = 1) {
+  markRetryOncePerTurn(modelId: ModelId, attempts: number = 1) {
+    const model = normalizeModelId(modelId);
     const currentState = this.health.get(model);
     // Do not override a terminal failure with a transient one.
     if (currentState?.status === 'terminal') {
@@ -75,14 +80,16 @@ export class ModelAvailabilityService {
     });
   }
 
-  consumeStickyAttempt(model: ModelId) {
+  consumeStickyAttempt(modelId: ModelId) {
+    const model = normalizeModelId(modelId);
     const state = this.health.get(model);
     if (state?.status === 'sticky_retry') {
       this.setState(model, { ...state, consumed: true });
     }
   }
 
-  snapshot(model: ModelId): ModelAvailabilitySnapshot {
+  snapshot(modelId: ModelId): ModelAvailabilitySnapshot {
+    const model = normalizeModelId(modelId);
     const state = this.health.get(model);
 
     if (!state) {
@@ -100,10 +107,11 @@ export class ModelAvailabilityService {
     return { available: true };
   }
 
-  selectFirstAvailable(models: ModelId[]): ModelSelectionResult {
+  selectFirstAvailable(modelIds: ModelId[]): ModelSelectionResult {
     const skipped: ModelSelectionResult['skipped'] = [];
 
-    for (const model of models) {
+    for (const modelId of modelIds) {
+      const model = normalizeModelId(modelId);
       const snapshot = this.snapshot(model);
       if (snapshot.available) {
         const state = this.health.get(model);

packages/core/src/availability/policyHelpers.test.ts

@@ -96,10 +96,11 @@ describe('policyHelpers', () => {
 
     it('starts chain from preferredModel when model is "auto"', () => {
       const config = createMockConfig({
-        getModel: () => DEFAULT_GEMINI_MODEL_AUTO,
+        getModel: () => 'auto',
       });
       const chain = resolvePolicyChain(config, 'gemini-2.5-flash');
-      expect(chain).toHaveLength(1);
+      // Due to Gemini 2.x wrapsAround, the chain will contain both flash and pro
+      expect(chain.length).toBeGreaterThanOrEqual(1);
       expect(chain[0]?.model).toBe('gemini-2.5-flash');
     });
 

packages/core/src/availability/policyHelpers.ts

@@ -28,6 +28,7 @@ import {
   isGemini3Model,
   resolveModel,
 } from '../config/models.js';
+import { normalizeModelId } from '../utils/modelUtils.js';
 import type { ModelSelectionResult } from './modelAvailabilityService.js';
 import type { ModelConfigKey } from '../services/modelConfigService.js';
 import { ApprovalMode } from '../policy/types.js';
@@ -41,9 +42,13 @@ export function resolvePolicyChain(
   preferredModel?: string,
   wrapsAround: boolean = false,
 ): ModelPolicyChain {
-  const modelFromConfig =
-    preferredModel ?? config.getActiveModel?.() ?? config.getModel();
-  const configuredModel = config.getModel();
+  const normalizedPreferredModel = preferredModel
+    ? normalizeModelId(preferredModel)
+    : undefined;
+  const modelFromConfig = normalizeModelId(
+    normalizedPreferredModel ?? config.getActiveModel?.() ?? config.getModel(),
+  );
+  const configuredModel = normalizeModelId(config.getModel());
 
   let chain: ModelPolicyChain | undefined;
   const useGemini31 = config.getGemini31LaunchedSync?.() ?? false;
@@ -52,19 +57,29 @@ export function resolvePolicyChain(
   const useCustomToolModel = config.getUseCustomToolModelSync?.() ?? false;
   const hasAccessToPreview = config.getHasAccessToPreviewModel?.() ?? true;
 
-  const resolvedModel = resolveModel(
-    modelFromConfig,
-    useGemini31,
-    useGemini31FlashLite,
-    useCustomToolModel,
-    hasAccessToPreview,
-    config,
+  const resolvedModel = normalizeModelId(
+    resolveModel(
+      modelFromConfig,
+      useGemini31,
+      useGemini31FlashLite,
+      useCustomToolModel,
+      hasAccessToPreview,
+      config,
+    ),
   );
-  const isAutoPreferred = preferredModel
-    ? isAutoModel(preferredModel, config)
+  const isAutoPreferred = normalizedPreferredModel
+    ? isAutoModel(normalizedPreferredModel, config)
     : false;
   const isAutoConfigured = isAutoModel(configuredModel, config);
 
+  // We always wrap around for Gemini 3 chains to ensure maximum availability
+  // between models in the same family (e.g. fallback to Pro if Flash is exhausted).
+  const effectiveWrapsAround =
+    wrapsAround ||
+    isAutoPreferred ||
+    isAutoConfigured ||
+    isGemini3Model(resolvedModel, config);
+
   // --- DYNAMIC PATH ---
   if (config.getExperimentalDynamicModelConfiguration?.() === true) {
     const context = {
@@ -76,7 +91,7 @@ export function resolvePolicyChain(
     if (resolvedModel === DEFAULT_GEMINI_FLASH_LITE_MODEL) {
       chain = config.modelConfigService.resolveChain('lite', context);
     } else if (
-      isGemini3Model(resolvedModel, config) ||
+      isGemini3Model(normalizeModelId(resolvedModel), config) ||
       isAutoPreferred ||
       isAutoConfigured
     ) {
@@ -96,7 +111,7 @@ export function resolvePolicyChain(
         const previewEnabled =
           hasAccessToPreview &&
           (isGemini3Model(resolvedModel, config) ||
-            preferredModel === PREVIEW_GEMINI_MODEL_AUTO ||
+            normalizedPreferredModel === PREVIEW_GEMINI_MODEL_AUTO ||
             configuredModel === PREVIEW_GEMINI_MODEL_AUTO);
         const autoPrefix = isAutoSelection ? 'auto-' : '';
         const chainKey = previewEnabled ? 'preview' : 'default';
@@ -110,7 +125,7 @@ export function resolvePolicyChain(
       // No matching modelChains found, default to single model chain
       chain = createSingleModelChain(modelFromConfig);
     }
-    chain = applyDynamicSlicing(chain, resolvedModel, wrapsAround);
+    chain = applyDynamicSlicing(chain, resolvedModel, effectiveWrapsAround);
   } else {
     // --- LEGACY PATH ---
 
@@ -125,7 +140,7 @@ export function resolvePolicyChain(
       if (hasAccessToPreview) {
         const previewEnabled =
           isGemini3Model(resolvedModel, config) ||
-          preferredModel === PREVIEW_GEMINI_MODEL_AUTO ||
+          normalizedPreferredModel === PREVIEW_GEMINI_MODEL_AUTO ||
           configuredModel === PREVIEW_GEMINI_MODEL_AUTO;
         chain = getModelPolicyChain({
           previewEnabled,
@@ -150,7 +165,7 @@ export function resolvePolicyChain(
     } else {
       chain = createSingleModelChain(modelFromConfig);
     }
-    chain = applyDynamicSlicing(chain, resolvedModel, wrapsAround);
+    chain = applyDynamicSlicing(chain, resolvedModel, effectiveWrapsAround);
   }
   // Apply Unified Silent Injection for Plan Mode with defensive checks
   if (config?.getApprovalMode?.() === ApprovalMode.PLAN) {
@@ -171,8 +186,9 @@ function applyDynamicSlicing(
   resolvedModel: string,
   wrapsAround: boolean,
 ): ModelPolicyChain {
+  const normalizedResolved = normalizeModelId(resolvedModel);
   const activeIndex = chain.findIndex(
-    (policy) => policy.model === resolvedModel,
+    (policy) => normalizeModelId(policy.model) === normalizedResolved,
   );
   if (activeIndex !== -1) {
     return wrapsAround
@@ -200,7 +216,10 @@ export function buildFallbackPolicyContext(
   failedPolicy?: ModelPolicy;
   candidates: ModelPolicy[];
 } {
-  const index = chain.findIndex((policy) => policy.model === failedModel);
+  const normalizedFailed = normalizeModelId(failedModel);
+  const index = chain.findIndex(
+    (policy) => normalizeModelId(policy.model) === normalizedFailed,
+  );
   if (index === -1) {
     return { failedPolicy: undefined, candidates: chain };
   }