Guard pro model usage (#22665)

This commit is contained in:
Sehoon Shon
2026-03-16 13:44:25 -04:00
committed by GitHub
parent ef5627eece
commit 48130ebd25
7 changed files with 252 additions and 9 deletions

View File

@@ -5,12 +5,13 @@
*/
import type React from 'react';
import { useCallback, useContext, useMemo, useState } from 'react';
import { useCallback, useContext, useMemo, useState, useEffect } from 'react';
import { Box, Text } from 'ink';
import {
PREVIEW_GEMINI_MODEL,
PREVIEW_GEMINI_3_1_MODEL,
PREVIEW_GEMINI_FLASH_MODEL,
PREVIEW_GEMINI_3_1_FLASH_LITE_MODEL,
PREVIEW_GEMINI_MODEL_AUTO,
DEFAULT_GEMINI_MODEL,
DEFAULT_GEMINI_FLASH_MODEL,
@@ -21,6 +22,8 @@ import {
getDisplayString,
AuthType,
PREVIEW_GEMINI_3_1_CUSTOM_TOOLS_MODEL,
isProModel,
UserTierId,
} from '@google/gemini-cli-core';
import { useKeypress } from '../hooks/useKeypress.js';
import { theme } from '../semantic-colors.js';
@@ -35,9 +38,26 @@ interface ModelDialogProps {
export function ModelDialog({ onClose }: ModelDialogProps): React.JSX.Element {
const config = useContext(ConfigContext);
const settings = useSettings();
const [view, setView] = useState<'main' | 'manual'>('main');
const [hasAccessToProModel, setHasAccessToProModel] = useState<boolean>(
() => !(config?.getProModelNoAccessSync() ?? false),
);
const [view, setView] = useState<'main' | 'manual'>(() =>
config?.getProModelNoAccessSync() ? 'manual' : 'main',
);
const [persistMode, setPersistMode] = useState(false);
useEffect(() => {
async function checkAccess() {
if (!config) return;
const noAccess = await config.getProModelNoAccess();
setHasAccessToProModel(!noAccess);
if (noAccess) {
setView('manual');
}
}
void checkAccess();
}, [config]);
// Determine the Preferred Model (read once when the dialog opens).
const preferredModel = config?.getModel() || DEFAULT_GEMINI_MODEL_AUTO;
@@ -66,7 +86,7 @@ export function ModelDialog({ onClose }: ModelDialogProps): React.JSX.Element {
useKeypress(
(key) => {
if (key.name === 'escape') {
if (view === 'manual') {
if (view === 'manual' && hasAccessToProModel) {
setView('main');
} else {
onClose();
@@ -115,6 +135,7 @@ export function ModelDialog({ onClose }: ModelDialogProps): React.JSX.Element {
}, [shouldShowPreviewModels, manualModelSelected, useGemini31]);
const manualOptions = useMemo(() => {
const isFreeTier = config?.getUserTier() === UserTierId.FREE;
const list = [
{
value: DEFAULT_GEMINI_MODEL,
@@ -142,7 +163,7 @@ export function ModelDialog({ onClose }: ModelDialogProps): React.JSX.Element {
? PREVIEW_GEMINI_3_1_CUSTOM_TOOLS_MODEL
: previewProModel;
list.unshift(
const previewOptions = [
{
value: previewProValue,
title: getDisplayString(previewProModel),
@@ -153,10 +174,32 @@ export function ModelDialog({ onClose }: ModelDialogProps): React.JSX.Element {
title: getDisplayString(PREVIEW_GEMINI_FLASH_MODEL),
key: PREVIEW_GEMINI_FLASH_MODEL,
},
);
];
if (isFreeTier) {
previewOptions.push({
value: PREVIEW_GEMINI_3_1_FLASH_LITE_MODEL,
title: getDisplayString(PREVIEW_GEMINI_3_1_FLASH_LITE_MODEL),
key: PREVIEW_GEMINI_3_1_FLASH_LITE_MODEL,
});
}
list.unshift(...previewOptions);
}
if (!hasAccessToProModel) {
// Filter out all Pro models for free tier
return list.filter((option) => !isProModel(option.value));
}
return list;
}, [shouldShowPreviewModels, useGemini31, useCustomToolModel]);
}, [
shouldShowPreviewModels,
useGemini31,
useCustomToolModel,
hasAccessToProModel,
config,
]);
const options = view === 'main' ? mainOptions : manualOptions;