Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
24 changes: 8 additions & 16 deletions packages/core/src/availability/policyHelpers.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand All @@ -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) {
Expand All @@ -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 &&
Expand Down Expand Up @@ -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({
Expand Down
1 change: 1 addition & 0 deletions packages/core/src/code_assist/setup.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
8 changes: 4 additions & 4 deletions packages/core/src/config/config.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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 () => {
Expand All @@ -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({
Expand Down
13 changes: 7 additions & 6 deletions packages/core/src/config/config.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}

Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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);
Expand Down
2 changes: 2 additions & 0 deletions packages/core/src/fallback/handler.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down Expand Up @@ -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,
Expand Down
10 changes: 9 additions & 1 deletion packages/core/src/fallback/handler.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand All @@ -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;
Expand Down
16 changes: 15 additions & 1 deletion packages/core/src/policy/core-tools-mapping.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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: {
Expand Down