fix(core): ensure stable fallback for restricted preview models (google-gemini#26999)

Author: galdawave

packages/core/src/availability/autoRoutingFallback.integration.test.ts

@@ -29,6 +29,9 @@ describe('Auto Routing Fallback Integration', () => {
 
   beforeEach(() => {
     vi.useFakeTimers();
+    vi.spyOn(Config.prototype, 'getHasAccessToPreviewModel').mockReturnValue(
+      true,
+    );
 
     // Mock fs to avoid real file system access
     vi.mocked(fs.existsSync).mockReturnValue(true);

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

@@ -28,6 +28,7 @@ describe('Fallback Integration', () => {
       getActiveModel: () => PREVIEW_GEMINI_MODEL_AUTO,
       setActiveModel: vi.fn(),
       getUserTier: () => undefined,
+      getHasAccessToPreviewModel: () => true,
       getModelAvailabilityService: () => availabilityService,
       modelConfigService: undefined as unknown as ModelConfigService,
     } as unknown as Config;

packages/core/src/availability/policyHelpers.ts

@@ -55,7 +55,10 @@ export function resolvePolicyChain(
   const useGemini31FlashLite =
     config.getGemini31FlashLiteLaunchedSync?.() ?? false;
   const useCustomToolModel = config.getUseCustomToolModelSync?.() ?? false;
-  const hasAccessToPreview = config.getHasAccessToPreviewModel?.() ?? true;
+  const hasAccessToPreview = config.getHasAccessToPreviewModel?.() ?? false;
+
+  // Capture the original family intent before any normalization or early downgrade.
+  const isOriginallyGemini3 = isGemini3Model(modelFromConfig, config);
 
   const resolvedModel = normalizeModelId(
     resolveModel(
@@ -75,10 +78,7 @@ export function resolvePolicyChain(
   // 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);
+    wrapsAround || isAutoPreferred || isAutoConfigured || isOriginallyGemini3;
 
   // --- DYNAMIC PATH ---
   if (config.getExperimentalDynamicModelConfiguration?.() === true) {
@@ -91,11 +91,7 @@ export function resolvePolicyChain(
 
     if (resolvedModel === DEFAULT_GEMINI_FLASH_LITE_MODEL) {
       chain = config.modelConfigService.resolveChain('lite', context);
-    } else if (
-      isGemini3Model(normalizeModelId(resolvedModel), config) ||
-      isAutoPreferred ||
-      isAutoConfigured
-    ) {
+    } else if (isOriginallyGemini3 || isAutoPreferred || isAutoConfigured) {
       // 1. Try to find a chain specifically for the current configured alias
       if (
         isAutoConfigured &&
@@ -132,15 +128,11 @@ export function resolvePolicyChain(
 
     if (resolvedModel === DEFAULT_GEMINI_FLASH_LITE_MODEL) {
       chain = getFlashLitePolicyChain();
-    } else if (
-      isGemini3Model(resolvedModel, config) ||
-      isAutoPreferred ||
-      isAutoConfigured
-    ) {
+    } else if (isOriginallyGemini3 || isAutoPreferred || isAutoConfigured) {
       const isAutoSelection = isAutoPreferred || isAutoConfigured;
       if (hasAccessToPreview) {
         const previewEnabled =
-          isGemini3Model(resolvedModel, config) ||
+          isOriginallyGemini3 ||
           normalizedPreferredModel === PREVIEW_GEMINI_MODEL_AUTO ||
           configuredModel === PREVIEW_GEMINI_MODEL_AUTO;
         chain = getModelPolicyChain({

packages/core/src/code_assist/setup.test.ts

@@ -228,6 +228,7 @@ describe('setupUser', () => {
     });
 
     it('should throw InvalidNumericProjectIdError when GOOGLE_CLOUD_PROJECT_ID is numeric', async () => {
+      vi.stubEnv('GOOGLE_CLOUD_PROJECT', '');
       vi.stubEnv('GOOGLE_CLOUD_PROJECT_ID', '1234567890');
       await expect(setupUser({} as OAuth2Client, mockConfig)).rejects.toThrow(
         InvalidNumericProjectIdError,

packages/core/src/config/config.test.ts

@@ -3251,8 +3251,8 @@ describe('Config Quota & Preview Model Access', () => {
       vi.mocked(getCodeAssistServer).mockReturnValue(undefined);
       const result = await config.refreshUserQuota();
       expect(result).toBeUndefined();
-      // Never set => stays null (unknown); getter returns true so UI shows preview
-      expect(config.getHasAccessToPreviewModel()).toBe(true);
+      // Never set => stays null (unknown); getter returns false by default
+      expect(config.getHasAccessToPreviewModel()).toBe(false);
     });
 
     it('should return undefined if retrieveUserQuota fails', async () => {
@@ -3261,8 +3261,8 @@ describe('Config Quota & Preview Model Access', () => {
       );
       const result = await config.refreshUserQuota();
       expect(result).toBeUndefined();
-      // Never set => stays null (unknown); getter returns true so UI shows preview
-      expect(config.getHasAccessToPreviewModel()).toBe(true);
+      // Never set => stays null (unknown); getter returns false by default
+      expect(config.getHasAccessToPreviewModel()).toBe(false);
     });
     it('should derive quota from remainingFraction when remainingAmount is missing', async () => {
       mockCodeAssistServer.retrieveUserQuota.mockResolvedValue({

packages/core/src/config/config.ts

@@ -1622,10 +1622,7 @@ export class Config implements McpContext, AgentLoopContext {
     this.baseLlmClient = new BaseLlmClient(this.contentGenerator, this);
 
     const authType = this.contentGeneratorConfig.authType;
-    if (
-      authType === AuthType.USE_GEMINI ||
-      authType === AuthType.USE_VERTEX_AI
-    ) {
+    if (authType === AuthType.USE_GEMINI) {
       this.setHasAccessToPreviewModel(true);
     }
 
@@ -2227,7 +2224,7 @@ export class Config implements McpContext, AgentLoopContext {
   }
 
   getHasAccessToPreviewModel(): boolean {
-    return this.hasAccessToPreviewModel !== false;
+    return this.hasAccessToPreviewModel ?? false;
   }
 
   setHasAccessToPreviewModel(hasAccess: boolean | null): void {
@@ -2289,14 +2286,18 @@ export class Config implements McpContext, AgentLoopContext {
             });
           }
         }
-        this.emitQuotaChangedEvent();
       }
 
       const hasAccess =
         quota.buckets?.some(
           (b) => b.modelId && isPreviewModel(b.modelId, this),
         ) ?? false;
       this.setHasAccessToPreviewModel(hasAccess);
+
+      if (quota.buckets) {
+        this.emitQuotaChangedEvent();
+      }
+
       return quota;
     } catch (e) {
       debugLogger.debug('Failed to retrieve user quota', e);

packages/core/src/fallback/handler.test.ts

@@ -77,6 +77,7 @@ const createMockConfig = (overrides: Partial<Config> = {}): Config =>
     getModel: vi.fn(() => MOCK_PRO_MODEL),
     getUserTier: vi.fn(() => undefined),
     isInteractive: vi.fn(() => false),
+    getHasAccessToPreviewModel: vi.fn(() => false),
     ...overrides,
   }) as unknown as Config;
 
@@ -234,6 +235,7 @@ describe('handleFallback', () => {
       vi.mocked(policyConfig.getModel).mockReturnValue(
         PREVIEW_GEMINI_MODEL_AUTO,
       );
+      vi.mocked(policyConfig.getHasAccessToPreviewModel).mockReturnValue(true);
 
       const result = await handleFallback(
         policyConfig,

packages/core/src/fallback/handler.ts

@@ -20,6 +20,8 @@ import {
   applyAvailabilityTransition,
 } from '../availability/policyHelpers.js';
 
+import { isPreviewModel } from '../config/models.js';
+
 export const UPGRADE_URL_PAGE = 'https://goo.gle/set-up-gemini-code-assist';
 
 export async function handleFallback(
@@ -28,13 +30,19 @@ export async function handleFallback(
   authType?: string,
   error?: unknown,
 ): Promise<string | boolean | null> {
+  const failureKind = classifyFailureKind(error);
+
+  // If a preview model is not found, record that the user lacks preview access.
+  if (failureKind === 'not_found' && isPreviewModel(failedModel, config)) {
+    config.setHasAccessToPreviewModel?.(false);
+  }
+
   const chain = resolvePolicyChain(config);
   const { failedPolicy, candidates } = buildFallbackPolicyContext(
     chain,
     failedModel,
   );
 
-  const failureKind = classifyFailureKind(error);
   const availability = config.getModelAvailabilityService();
   const getAvailabilityContext = () => {
     if (!failedPolicy) return undefined;

packages/core/src/policy/core-tools-mapping.test.ts

@@ -4,12 +4,26 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { describe, it, expect } from 'vitest';
+import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
 import { createPolicyEngineConfig } from './config.js';
 import { PolicyEngine } from './policy-engine.js';
 import { PolicyDecision, ApprovalMode } from './types.js';
+import { Storage } from '../config/storage.js';
 
 describe('PolicyEngine - Core Tools Mapping', () => {
+  beforeEach(() => {
+    vi.spyOn(Storage, 'getUserPoliciesDir').mockReturnValue(
+      '/mock/user/policies',
+    );
+    vi.spyOn(Storage, 'getSystemPoliciesDir').mockReturnValue(
+      '/mock/system/policies',
+    );
+  });
+
+  afterEach(() => {
+    vi.restoreAllMocks();
+  });
+
   it('should allow tools explicitly listed in settings.tools.core', async () => {
     const settings = {
       tools: {