diff --git a/packages/core/src/availability/autoRoutingFallback.integration.test.ts b/packages/core/src/availability/autoRoutingFallback.integration.test.ts index f4e157503b..9ea062e1ab 100644 --- a/packages/core/src/availability/autoRoutingFallback.integration.test.ts +++ b/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); diff --git a/packages/core/src/availability/fallbackIntegration.test.ts b/packages/core/src/availability/fallbackIntegration.test.ts index 6c49938ed9..1ee53da37e 100644 --- a/packages/core/src/availability/fallbackIntegration.test.ts +++ b/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; diff --git a/packages/core/src/availability/policyHelpers.ts b/packages/core/src/availability/policyHelpers.ts index 530bb43dce..e818ab52a1 100644 --- a/packages/core/src/availability/policyHelpers.ts +++ b/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({ diff --git a/packages/core/src/code_assist/setup.test.ts b/packages/core/src/code_assist/setup.test.ts index 6779143b9a..a76525470f 100644 --- a/packages/core/src/code_assist/setup.test.ts +++ b/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, diff --git a/packages/core/src/config/config.test.ts b/packages/core/src/config/config.test.ts index c0256281a5..ad5f839515 100644 --- a/packages/core/src/config/config.test.ts +++ b/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({ diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index 9206a01f15..f1077058e9 100644 --- a/packages/core/src/config/config.ts +++ b/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,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); diff --git a/packages/core/src/fallback/handler.test.ts b/packages/core/src/fallback/handler.test.ts index 698a5d7cfb..0bc3096f70 100644 --- a/packages/core/src/fallback/handler.test.ts +++ b/packages/core/src/fallback/handler.test.ts @@ -77,6 +77,7 @@ const createMockConfig = (overrides: Partial = {}): 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, diff --git a/packages/core/src/fallback/handler.ts b/packages/core/src/fallback/handler.ts index f216f9216c..3e8dcc5267 100644 --- a/packages/core/src/fallback/handler.ts +++ b/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 { + 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; diff --git a/packages/core/src/policy/core-tools-mapping.test.ts b/packages/core/src/policy/core-tools-mapping.test.ts index 95877c6ac4..8ef042d6bb 100644 --- a/packages/core/src/policy/core-tools-mapping.test.ts +++ b/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: {