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

This commit is contained in:
Gal Zahavi
2026-05-13 14:46:41 -07:00
committed by GitHub
parent 1814c7f358
commit 77078b3e8a
9 changed files with 50 additions and 28 deletions
@@ -29,6 +29,9 @@ describe('Auto Routing Fallback Integration', () => {
beforeEach(() => { beforeEach(() => {
vi.useFakeTimers(); vi.useFakeTimers();
vi.spyOn(Config.prototype, 'getHasAccessToPreviewModel').mockReturnValue(
true,
);
// Mock fs to avoid real file system access // Mock fs to avoid real file system access
vi.mocked(fs.existsSync).mockReturnValue(true); vi.mocked(fs.existsSync).mockReturnValue(true);
@@ -28,6 +28,7 @@ describe('Fallback Integration', () => {
getActiveModel: () => PREVIEW_GEMINI_MODEL_AUTO, getActiveModel: () => PREVIEW_GEMINI_MODEL_AUTO,
setActiveModel: vi.fn(), setActiveModel: vi.fn(),
getUserTier: () => undefined, getUserTier: () => undefined,
getHasAccessToPreviewModel: () => true,
getModelAvailabilityService: () => availabilityService, getModelAvailabilityService: () => availabilityService,
modelConfigService: undefined as unknown as ModelConfigService, modelConfigService: undefined as unknown as ModelConfigService,
} as unknown as Config; } as unknown as Config;
@@ -55,7 +55,10 @@ export function resolvePolicyChain(
const useGemini31FlashLite = const useGemini31FlashLite =
config.getGemini31FlashLiteLaunchedSync?.() ?? false; config.getGemini31FlashLiteLaunchedSync?.() ?? false;
const useCustomToolModel = config.getUseCustomToolModelSync?.() ?? 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( const resolvedModel = normalizeModelId(
resolveModel( resolveModel(
@@ -75,10 +78,7 @@ export function resolvePolicyChain(
// We always wrap around for Gemini 3 chains to ensure maximum availability // 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). // between models in the same family (e.g. fallback to Pro if Flash is exhausted).
const effectiveWrapsAround = const effectiveWrapsAround =
wrapsAround || wrapsAround || isAutoPreferred || isAutoConfigured || isOriginallyGemini3;
isAutoPreferred ||
isAutoConfigured ||
isGemini3Model(resolvedModel, config);
// --- DYNAMIC PATH --- // --- DYNAMIC PATH ---
if (config.getExperimentalDynamicModelConfiguration?.() === true) { if (config.getExperimentalDynamicModelConfiguration?.() === true) {
@@ -91,11 +91,7 @@ export function resolvePolicyChain(
if (resolvedModel === DEFAULT_GEMINI_FLASH_LITE_MODEL) { if (resolvedModel === DEFAULT_GEMINI_FLASH_LITE_MODEL) {
chain = config.modelConfigService.resolveChain('lite', context); chain = config.modelConfigService.resolveChain('lite', context);
} else if ( } else if (isOriginallyGemini3 || isAutoPreferred || isAutoConfigured) {
isGemini3Model(normalizeModelId(resolvedModel), config) ||
isAutoPreferred ||
isAutoConfigured
) {
// 1. Try to find a chain specifically for the current configured alias // 1. Try to find a chain specifically for the current configured alias
if ( if (
isAutoConfigured && isAutoConfigured &&
@@ -132,15 +128,11 @@ export function resolvePolicyChain(
if (resolvedModel === DEFAULT_GEMINI_FLASH_LITE_MODEL) { if (resolvedModel === DEFAULT_GEMINI_FLASH_LITE_MODEL) {
chain = getFlashLitePolicyChain(); chain = getFlashLitePolicyChain();
} else if ( } else if (isOriginallyGemini3 || isAutoPreferred || isAutoConfigured) {
isGemini3Model(resolvedModel, config) ||
isAutoPreferred ||
isAutoConfigured
) {
const isAutoSelection = isAutoPreferred || isAutoConfigured; const isAutoSelection = isAutoPreferred || isAutoConfigured;
if (hasAccessToPreview) { if (hasAccessToPreview) {
const previewEnabled = const previewEnabled =
isGemini3Model(resolvedModel, config) || isOriginallyGemini3 ||
normalizedPreferredModel === PREVIEW_GEMINI_MODEL_AUTO || normalizedPreferredModel === PREVIEW_GEMINI_MODEL_AUTO ||
configuredModel === PREVIEW_GEMINI_MODEL_AUTO; configuredModel === PREVIEW_GEMINI_MODEL_AUTO;
chain = getModelPolicyChain({ chain = getModelPolicyChain({
@@ -228,6 +228,7 @@ describe('setupUser', () => {
}); });
it('should throw InvalidNumericProjectIdError when GOOGLE_CLOUD_PROJECT_ID is numeric', async () => { it('should throw InvalidNumericProjectIdError when GOOGLE_CLOUD_PROJECT_ID is numeric', async () => {
vi.stubEnv('GOOGLE_CLOUD_PROJECT', '');
vi.stubEnv('GOOGLE_CLOUD_PROJECT_ID', '1234567890'); vi.stubEnv('GOOGLE_CLOUD_PROJECT_ID', '1234567890');
await expect(setupUser({} as OAuth2Client, mockConfig)).rejects.toThrow( await expect(setupUser({} as OAuth2Client, mockConfig)).rejects.toThrow(
InvalidNumericProjectIdError, InvalidNumericProjectIdError,
+4 -4
View File
@@ -3251,8 +3251,8 @@ describe('Config Quota & Preview Model Access', () => {
vi.mocked(getCodeAssistServer).mockReturnValue(undefined); vi.mocked(getCodeAssistServer).mockReturnValue(undefined);
const result = await config.refreshUserQuota(); const result = await config.refreshUserQuota();
expect(result).toBeUndefined(); expect(result).toBeUndefined();
// Never set => stays null (unknown); getter returns true so UI shows preview // Never set => stays null (unknown); getter returns false by default
expect(config.getHasAccessToPreviewModel()).toBe(true); expect(config.getHasAccessToPreviewModel()).toBe(false);
}); });
it('should return undefined if retrieveUserQuota fails', async () => { it('should return undefined if retrieveUserQuota fails', async () => {
@@ -3261,8 +3261,8 @@ describe('Config Quota & Preview Model Access', () => {
); );
const result = await config.refreshUserQuota(); const result = await config.refreshUserQuota();
expect(result).toBeUndefined(); expect(result).toBeUndefined();
// Never set => stays null (unknown); getter returns true so UI shows preview // Never set => stays null (unknown); getter returns false by default
expect(config.getHasAccessToPreviewModel()).toBe(true); expect(config.getHasAccessToPreviewModel()).toBe(false);
}); });
it('should derive quota from remainingFraction when remainingAmount is missing', async () => { it('should derive quota from remainingFraction when remainingAmount is missing', async () => {
mockCodeAssistServer.retrieveUserQuota.mockResolvedValue({ mockCodeAssistServer.retrieveUserQuota.mockResolvedValue({
+7 -6
View File
@@ -1622,10 +1622,7 @@ export class Config implements McpContext, AgentLoopContext {
this.baseLlmClient = new BaseLlmClient(this.contentGenerator, this); this.baseLlmClient = new BaseLlmClient(this.contentGenerator, this);
const authType = this.contentGeneratorConfig.authType; const authType = this.contentGeneratorConfig.authType;
if ( if (authType === AuthType.USE_GEMINI) {
authType === AuthType.USE_GEMINI ||
authType === AuthType.USE_VERTEX_AI
) {
this.setHasAccessToPreviewModel(true); this.setHasAccessToPreviewModel(true);
} }
@@ -2227,7 +2224,7 @@ export class Config implements McpContext, AgentLoopContext {
} }
getHasAccessToPreviewModel(): boolean { getHasAccessToPreviewModel(): boolean {
return this.hasAccessToPreviewModel !== false; return this.hasAccessToPreviewModel ?? false;
} }
setHasAccessToPreviewModel(hasAccess: boolean | null): void { setHasAccessToPreviewModel(hasAccess: boolean | null): void {
@@ -2289,7 +2286,6 @@ export class Config implements McpContext, AgentLoopContext {
}); });
} }
} }
this.emitQuotaChangedEvent();
} }
const hasAccess = const hasAccess =
@@ -2297,6 +2293,11 @@ export class Config implements McpContext, AgentLoopContext {
(b) => b.modelId && isPreviewModel(b.modelId, this), (b) => b.modelId && isPreviewModel(b.modelId, this),
) ?? false; ) ?? false;
this.setHasAccessToPreviewModel(hasAccess); this.setHasAccessToPreviewModel(hasAccess);
if (quota.buckets) {
this.emitQuotaChangedEvent();
}
return quota; return quota;
} catch (e) { } catch (e) {
debugLogger.debug('Failed to retrieve user quota', e); debugLogger.debug('Failed to retrieve user quota', e);
@@ -77,6 +77,7 @@ const createMockConfig = (overrides: Partial<Config> = {}): Config =>
getModel: vi.fn(() => MOCK_PRO_MODEL), getModel: vi.fn(() => MOCK_PRO_MODEL),
getUserTier: vi.fn(() => undefined), getUserTier: vi.fn(() => undefined),
isInteractive: vi.fn(() => false), isInteractive: vi.fn(() => false),
getHasAccessToPreviewModel: vi.fn(() => false),
...overrides, ...overrides,
}) as unknown as Config; }) as unknown as Config;
@@ -234,6 +235,7 @@ describe('handleFallback', () => {
vi.mocked(policyConfig.getModel).mockReturnValue( vi.mocked(policyConfig.getModel).mockReturnValue(
PREVIEW_GEMINI_MODEL_AUTO, PREVIEW_GEMINI_MODEL_AUTO,
); );
vi.mocked(policyConfig.getHasAccessToPreviewModel).mockReturnValue(true);
const result = await handleFallback( const result = await handleFallback(
policyConfig, policyConfig,
+9 -1
View File
@@ -20,6 +20,8 @@ import {
applyAvailabilityTransition, applyAvailabilityTransition,
} from '../availability/policyHelpers.js'; } from '../availability/policyHelpers.js';
import { isPreviewModel } from '../config/models.js';
export const UPGRADE_URL_PAGE = 'https://goo.gle/set-up-gemini-code-assist'; export const UPGRADE_URL_PAGE = 'https://goo.gle/set-up-gemini-code-assist';
export async function handleFallback( export async function handleFallback(
@@ -28,13 +30,19 @@ export async function handleFallback(
authType?: string, authType?: string,
error?: unknown, error?: unknown,
): Promise<string | boolean | null> { ): 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 chain = resolvePolicyChain(config);
const { failedPolicy, candidates } = buildFallbackPolicyContext( const { failedPolicy, candidates } = buildFallbackPolicyContext(
chain, chain,
failedModel, failedModel,
); );
const failureKind = classifyFailureKind(error);
const availability = config.getModelAvailabilityService(); const availability = config.getModelAvailabilityService();
const getAvailabilityContext = () => { const getAvailabilityContext = () => {
if (!failedPolicy) return undefined; if (!failedPolicy) return undefined;
@@ -4,12 +4,26 @@
* SPDX-License-Identifier: Apache-2.0 * 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 { createPolicyEngineConfig } from './config.js';
import { PolicyEngine } from './policy-engine.js'; import { PolicyEngine } from './policy-engine.js';
import { PolicyDecision, ApprovalMode } from './types.js'; import { PolicyDecision, ApprovalMode } from './types.js';
import { Storage } from '../config/storage.js';
describe('PolicyEngine - Core Tools Mapping', () => { 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 () => { it('should allow tools explicitly listed in settings.tools.core', async () => {
const settings = { const settings = {
tools: { tools: {