Skip to content

Commit a535c19

Browse files
sehoon38kunal-10-cloud
sehoon38
authored and
kunal-10-cloud
committed
Guard pro model usage (google-gemini#22665)
1 parent 4eafdc6 commit a535c19

7 files changed

Lines changed: 252 additions & 9 deletions

File tree

‎packages/cli/src/ui/components/ModelDialog.test.tsx‎

Lines changed: 112 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,9 @@ import {
1919
PREVIEW_GEMINI_3_1_MODEL,
2020
PREVIEW_GEMINI_3_1_CUSTOM_TOOLS_MODEL,
2121
PREVIEW_GEMINI_FLASH_MODEL,
22+
PREVIEW_GEMINI_3_1_FLASH_LITE_MODEL,
2223
AuthType,
24+
UserTierId,
2325
} from '@google/gemini-cli-core';
2426
import type { Config, ModelSlashCommandEvent } from '@google/gemini-cli-core';
2527

@@ -28,8 +30,9 @@ const mockGetDisplayString = vi.fn();
2830
const mockLogModelSlashCommand = vi.fn();
2931
const mockModelSlashCommandEvent = vi.fn();
3032

31-
vi.mock('@google/gemini-cli-core', async () => {
32-
const actual = await vi.importActual('@google/gemini-cli-core');
33+
vi.mock('@google/gemini-cli-core', async (importOriginal) => {
34+
const actual =
35+
await importOriginal<typeof import('@google/gemini-cli-core')>();
3336
return {
3437
...actual,
3538
getDisplayString: (val: string) => mockGetDisplayString(val),
@@ -40,6 +43,7 @@ vi.mock('@google/gemini-cli-core', async () => {
4043
mockModelSlashCommandEvent(model);
4144
}
4245
},
46+
PREVIEW_GEMINI_3_1_FLASH_LITE_MODEL: 'gemini-3.1-flash-lite-preview',
4347
};
4448
});
4549

@@ -49,13 +53,19 @@ describe('<ModelDialog />', () => {
4953
const mockOnClose = vi.fn();
5054
const mockGetHasAccessToPreviewModel = vi.fn();
5155
const mockGetGemini31LaunchedSync = vi.fn();
56+
const mockGetProModelNoAccess = vi.fn();
57+
const mockGetProModelNoAccessSync = vi.fn();
58+
const mockGetUserTier = vi.fn();
5259

5360
interface MockConfig extends Partial<Config> {
5461
setModel: (model: string, isTemporary?: boolean) => void;
5562
getModel: () => string;
5663
getHasAccessToPreviewModel: () => boolean;
5764
getIdeMode: () => boolean;
5865
getGemini31LaunchedSync: () => boolean;
66+
getProModelNoAccess: () => Promise<boolean>;
67+
getProModelNoAccessSync: () => boolean;
68+
getUserTier: () => UserTierId | undefined;
5969
}
6070

6171
const mockConfig: MockConfig = {
@@ -64,13 +74,19 @@ describe('<ModelDialog />', () => {
6474
getHasAccessToPreviewModel: mockGetHasAccessToPreviewModel,
6575
getIdeMode: () => false,
6676
getGemini31LaunchedSync: mockGetGemini31LaunchedSync,
77+
getProModelNoAccess: mockGetProModelNoAccess,
78+
getProModelNoAccessSync: mockGetProModelNoAccessSync,
79+
getUserTier: mockGetUserTier,
6780
};
6881

6982
beforeEach(() => {
7083
vi.resetAllMocks();
7184
mockGetModel.mockReturnValue(DEFAULT_GEMINI_MODEL_AUTO);
7285
mockGetHasAccessToPreviewModel.mockReturnValue(false);
7386
mockGetGemini31LaunchedSync.mockReturnValue(false);
87+
mockGetProModelNoAccess.mockResolvedValue(false);
88+
mockGetProModelNoAccessSync.mockReturnValue(false);
89+
mockGetUserTier.mockReturnValue(UserTierId.STANDARD);
7490

7591
// Default implementation for getDisplayString
7692
mockGetDisplayString.mockImplementation((val: string) => {
@@ -109,6 +125,55 @@ describe('<ModelDialog />', () => {
109125
unmount();
110126
});
111127

128+
it('renders the "manual" view initially for users with no pro access and filters Pro models with correct order', async () => {
129+
mockGetProModelNoAccessSync.mockReturnValue(true);
130+
mockGetProModelNoAccess.mockResolvedValue(true);
131+
mockGetHasAccessToPreviewModel.mockReturnValue(true);
132+
mockGetUserTier.mockReturnValue(UserTierId.FREE);
133+
mockGetDisplayString.mockImplementation((val: string) => val);
134+
135+
const { lastFrame, unmount } = await renderComponent();
136+
137+
const output = lastFrame();
138+
expect(output).toContain('Select Model');
139+
expect(output).not.toContain(DEFAULT_GEMINI_MODEL);
140+
expect(output).not.toContain(PREVIEW_GEMINI_MODEL);
141+
142+
// Verify order: Flash Preview -> Flash Lite Preview -> Flash -> Flash Lite
143+
const flashPreviewIdx = output.indexOf(PREVIEW_GEMINI_FLASH_MODEL);
144+
const flashLitePreviewIdx = output.indexOf(
145+
PREVIEW_GEMINI_3_1_FLASH_LITE_MODEL,
146+
);
147+
const flashIdx = output.indexOf(DEFAULT_GEMINI_FLASH_MODEL);
148+
const flashLiteIdx = output.indexOf(DEFAULT_GEMINI_FLASH_LITE_MODEL);
149+
150+
expect(flashPreviewIdx).toBeLessThan(flashLitePreviewIdx);
151+
expect(flashLitePreviewIdx).toBeLessThan(flashIdx);
152+
expect(flashIdx).toBeLessThan(flashLiteIdx);
153+
154+
expect(output).not.toContain('Auto');
155+
unmount();
156+
});
157+
158+
it('closes dialog on escape in "manual" view for users with no pro access', async () => {
159+
mockGetProModelNoAccessSync.mockReturnValue(true);
160+
mockGetProModelNoAccess.mockResolvedValue(true);
161+
const { stdin, waitUntilReady, unmount } = await renderComponent();
162+
163+
// Already in manual view
164+
await act(async () => {
165+
stdin.write('\u001B'); // Escape
166+
});
167+
await act(async () => {
168+
await waitUntilReady();
169+
});
170+
171+
await waitFor(() => {
172+
expect(mockOnClose).toHaveBeenCalled();
173+
});
174+
unmount();
175+
});
176+
112177
it('switches to "manual" view when "Manual" is selected and uses getDisplayString for models', async () => {
113178
mockGetDisplayString.mockImplementation((val: string) => {
114179
if (val === DEFAULT_GEMINI_MODEL) return 'Formatted Pro Model';
@@ -369,5 +434,50 @@ describe('<ModelDialog />', () => {
369434
});
370435
unmount();
371436
});
437+
438+
it('hides Flash Lite Preview model for users with pro access', async () => {
439+
mockGetProModelNoAccessSync.mockReturnValue(false);
440+
mockGetProModelNoAccess.mockResolvedValue(false);
441+
mockGetHasAccessToPreviewModel.mockReturnValue(true);
442+
const { lastFrame, stdin, waitUntilReady, unmount } =
443+
await renderComponent();
444+
445+
// Go to manual view
446+
await act(async () => {
447+
stdin.write('\u001B[B'); // Manual
448+
});
449+
await waitUntilReady();
450+
await act(async () => {
451+
stdin.write('\r');
452+
});
453+
await waitUntilReady();
454+
455+
const output = lastFrame();
456+
expect(output).not.toContain(PREVIEW_GEMINI_3_1_FLASH_LITE_MODEL);
457+
unmount();
458+
});
459+
460+
it('shows Flash Lite Preview model for free tier users', async () => {
461+
mockGetProModelNoAccessSync.mockReturnValue(false);
462+
mockGetProModelNoAccess.mockResolvedValue(false);
463+
mockGetHasAccessToPreviewModel.mockReturnValue(true);
464+
mockGetUserTier.mockReturnValue(UserTierId.FREE);
465+
const { lastFrame, stdin, waitUntilReady, unmount } =
466+
await renderComponent();
467+
468+
// Go to manual view
469+
await act(async () => {
470+
stdin.write('\u001B[B'); // Manual
471+
});
472+
await waitUntilReady();
473+
await act(async () => {
474+
stdin.write('\r');
475+
});
476+
await waitUntilReady();
477+
478+
const output = lastFrame();
479+
expect(output).toContain(PREVIEW_GEMINI_3_1_FLASH_LITE_MODEL);
480+
unmount();
481+
});
372482
});
373483
});

‎packages/cli/src/ui/components/ModelDialog.tsx‎

Lines changed: 49 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -5,12 +5,13 @@
55
*/
66

77
import type React from 'react';
8-
import { useCallback, useContext, useMemo, useState } from 'react';
8+
import { useCallback, useContext, useMemo, useState, useEffect } from 'react';
99
import { Box, Text } from 'ink';
1010
import {
1111
PREVIEW_GEMINI_MODEL,
1212
PREVIEW_GEMINI_3_1_MODEL,
1313
PREVIEW_GEMINI_FLASH_MODEL,
14+
PREVIEW_GEMINI_3_1_FLASH_LITE_MODEL,
1415
PREVIEW_GEMINI_MODEL_AUTO,
1516
DEFAULT_GEMINI_MODEL,
1617
DEFAULT_GEMINI_FLASH_MODEL,
@@ -21,6 +22,8 @@ import {
2122
getDisplayString,
2223
AuthType,
2324
PREVIEW_GEMINI_3_1_CUSTOM_TOOLS_MODEL,
25+
isProModel,
26+
UserTierId,
2427
} from '@google/gemini-cli-core';
2528
import { useKeypress } from '../hooks/useKeypress.js';
2629
import { theme } from '../semantic-colors.js';
@@ -35,9 +38,26 @@ interface ModelDialogProps {
3538
export function ModelDialog({ onClose }: ModelDialogProps): React.JSX.Element {
3639
const config = useContext(ConfigContext);
3740
const settings = useSettings();
38-
const [view, setView] = useState<'main' | 'manual'>('main');
41+
const [hasAccessToProModel, setHasAccessToProModel] = useState<boolean>(
42+
() => !(config?.getProModelNoAccessSync() ?? false),
43+
);
44+
const [view, setView] = useState<'main' | 'manual'>(() =>
45+
config?.getProModelNoAccessSync() ? 'manual' : 'main',
46+
);
3947
const [persistMode, setPersistMode] = useState(false);
4048

49+
useEffect(() => {
50+
async function checkAccess() {
51+
if (!config) return;
52+
const noAccess = await config.getProModelNoAccess();
53+
setHasAccessToProModel(!noAccess);
54+
if (noAccess) {
55+
setView('manual');
56+
}
57+
}
58+
void checkAccess();
59+
}, [config]);
60+
4161
// Determine the Preferred Model (read once when the dialog opens).
4262
const preferredModel = config?.getModel() || DEFAULT_GEMINI_MODEL_AUTO;
4363

@@ -66,7 +86,7 @@ export function ModelDialog({ onClose }: ModelDialogProps): React.JSX.Element {
6686
useKeypress(
6787
(key) => {
6888
if (key.name === 'escape') {
69-
if (view === 'manual') {
89+
if (view === 'manual' && hasAccessToProModel) {
7090
setView('main');
7191
} else {
7292
onClose();
@@ -115,6 +135,7 @@ export function ModelDialog({ onClose }: ModelDialogProps): React.JSX.Element {
115135
}, [shouldShowPreviewModels, manualModelSelected, useGemini31]);
116136

117137
const manualOptions = useMemo(() => {
138+
const isFreeTier = config?.getUserTier() === UserTierId.FREE;
118139
const list = [
119140
{
120141
value: DEFAULT_GEMINI_MODEL,
@@ -142,7 +163,7 @@ export function ModelDialog({ onClose }: ModelDialogProps): React.JSX.Element {
142163
? PREVIEW_GEMINI_3_1_CUSTOM_TOOLS_MODEL
143164
: previewProModel;
144165

145-
list.unshift(
166+
const previewOptions = [
146167
{
147168
value: previewProValue,
148169
title: getDisplayString(previewProModel),
@@ -153,10 +174,32 @@ export function ModelDialog({ onClose }: ModelDialogProps): React.JSX.Element {
153174
title: getDisplayString(PREVIEW_GEMINI_FLASH_MODEL),
154175
key: PREVIEW_GEMINI_FLASH_MODEL,
155176
},
156-
);
177+
];
178+
179+
if (isFreeTier) {
180+
previewOptions.push({
181+
value: PREVIEW_GEMINI_3_1_FLASH_LITE_MODEL,
182+
title: getDisplayString(PREVIEW_GEMINI_3_1_FLASH_LITE_MODEL),
183+
key: PREVIEW_GEMINI_3_1_FLASH_LITE_MODEL,
184+
});
185+
}
186+
187+
list.unshift(...previewOptions);
157188
}
189+
190+
if (!hasAccessToProModel) {
191+
// Filter out all Pro models for free tier
192+
return list.filter((option) => !isProModel(option.value));
193+
}
194+
158195
return list;
159-
}, [shouldShowPreviewModels, useGemini31, useCustomToolModel]);
196+
}, [
197+
shouldShowPreviewModels,
198+
useGemini31,
199+
useCustomToolModel,
200+
hasAccessToProModel,
201+
config,
202+
]);
160203

161204
const options = view === 'main' ? mainOptions : manualOptions;
162205

‎packages/core/src/code_assist/experiments/flagNames.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ export const ExperimentFlags = {
1717
MASKING_PRUNABLE_THRESHOLD: 45758818,
1818
MASKING_PROTECT_LATEST_TURN: 45758819,
1919
GEMINI_3_1_PRO_LAUNCHED: 45760185,
20+
PRO_MODEL_NO_ACCESS: 45768879,
2021
} as const;
2122

2223
export type ExperimentFlagName =

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

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -65,6 +65,8 @@ import {
6565
DEFAULT_GEMINI_MODEL,
6666
PREVIEW_GEMINI_3_1_MODEL,
6767
DEFAULT_GEMINI_MODEL_AUTO,
68+
PREVIEW_GEMINI_MODEL_AUTO,
69+
PREVIEW_GEMINI_FLASH_MODEL,
6870
} from './models.js';
6971
import { Storage } from './storage.js';
7072
import type { AgentLoopContext } from './agent-loop-context.js';
@@ -687,6 +689,46 @@ describe('Server Config (config.ts)', () => {
687689
loopContext.geminiClient.stripThoughtsFromHistory,
688690
).not.toHaveBeenCalledWith();
689691
});
692+
693+
it('should switch to flash model if user has no Pro access and model is auto', async () => {
694+
vi.mocked(getExperiments).mockResolvedValue({
695+
experimentIds: [],
696+
flags: {
697+
[ExperimentFlags.PRO_MODEL_NO_ACCESS]: {
698+
boolValue: true,
699+
},
700+
},
701+
});
702+
703+
const config = new Config({
704+
...baseParams,
705+
model: PREVIEW_GEMINI_MODEL_AUTO,
706+
});
707+
708+
await config.refreshAuth(AuthType.LOGIN_WITH_GOOGLE);
709+
710+
expect(config.getModel()).toBe(PREVIEW_GEMINI_FLASH_MODEL);
711+
});
712+
713+
it('should NOT switch to flash model if user has Pro access and model is auto', async () => {
714+
vi.mocked(getExperiments).mockResolvedValue({
715+
experimentIds: [],
716+
flags: {
717+
[ExperimentFlags.PRO_MODEL_NO_ACCESS]: {
718+
boolValue: false,
719+
},
720+
},
721+
});
722+
723+
const config = new Config({
724+
...baseParams,
725+
model: PREVIEW_GEMINI_MODEL_AUTO,
726+
});
727+
728+
await config.refreshAuth(AuthType.LOGIN_WITH_GOOGLE);
729+
730+
expect(config.getModel()).toBe(PREVIEW_GEMINI_MODEL_AUTO);
731+
});
690732
});
691733

692734
it('Config constructor should store userMemory correctly', () => {

‎packages/core/src/config/config.ts‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1395,6 +1395,10 @@ export class Config implements McpContext, AgentLoopContext {
13951395
},
13961396
);
13971397
this.setRemoteAdminSettings(adminControls);
1398+
1399+
if ((await this.getProModelNoAccess()) && isAutoModel(this.model)) {
1400+
this.setModel(PREVIEW_GEMINI_FLASH_MODEL);
1401+
}
13981402
}
13991403

14001404
async getExperimentsAsync(): Promise<Experiments | undefined> {
@@ -2690,6 +2694,30 @@ export class Config implements McpContext, AgentLoopContext {
26902694
);
26912695
}
26922696

2697+
/**
2698+
* Returns whether the user has access to Pro models.
2699+
* This is determined by the PRO_MODEL_NO_ACCESS experiment flag.
2700+
*/
2701+
async getProModelNoAccess(): Promise<boolean> {
2702+
await this.ensureExperimentsLoaded();
2703+
return this.getProModelNoAccessSync();
2704+
}
2705+
2706+
/**
2707+
* Returns whether the user has access to Pro models synchronously.
2708+
*
2709+
* Note: This method should only be called after startup, once experiments have been loaded.
2710+
*/
2711+
getProModelNoAccessSync(): boolean {
2712+
if (this.contentGeneratorConfig?.authType !== AuthType.LOGIN_WITH_GOOGLE) {
2713+
return false;
2714+
}
2715+
return (
2716+
this.experiments?.flags[ExperimentFlags.PRO_MODEL_NO_ACCESS]?.boolValue ??
2717+
false
2718+
);
2719+
}
2720+
26932721
/**
26942722
* Returns whether Gemini 3.1 has been launched.
26952723
* This method is async and ensures that experiments are loaded before returning the result.

0 commit comments

Comments
 (0)