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(() => {
vi.useFakeTimers();
vi.spyOn(Config.prototype, 'getHasAccessToPreviewModel').mockReturnValue(
true,
);
// Mock fs to avoid real file system access
vi.mocked(fs.existsSync).mockReturnValue(true);
@@ -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;
@@ -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({
@@ -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,
+4 -4
View File
@@ -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({
+7 -6
View File
@@ -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,7 +2286,6 @@ export class Config implements McpContext, AgentLoopContext {
});
}
}
this.emitQuotaChangedEvent();
}
const hasAccess =
@@ -2297,6 +2293,11 @@ export class Config implements McpContext, AgentLoopContext {
(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);
@@ -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,
+9 -1
View File
@@ -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;
@@ -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: {