fix(cli): support quota error fallbacks for all authentication types (#20475)

Co-authored-by: Gal Zahavi <38544478+galz10@users.noreply.github.com>
This commit is contained in:
Sehoon Shon
2026-02-26 17:39:25 -05:00
committed by GitHub
parent 10c5bd8ce9
commit edb1fdea30
7 changed files with 89 additions and 36 deletions

View File

@@ -135,6 +135,7 @@ export const DialogManager = ({
isModelNotFoundError={
!!uiState.quota.proQuotaRequest.isModelNotFoundError
}
authType={uiState.quota.proQuotaRequest.authType}
onChoice={uiActions.handleProQuotaChoice}
/>
);

View File

@@ -13,6 +13,7 @@ import { RadioButtonSelect } from './shared/RadioButtonSelect.js';
import {
PREVIEW_GEMINI_MODEL,
DEFAULT_GEMINI_FLASH_MODEL,
AuthType,
} from '@google/gemini-cli-core';
// Mock the child component to make it easier to test the parent
@@ -62,7 +63,7 @@ describe('ProQuotaDialog', () => {
describe('for non-flash model failures', () => {
describe('when it is a terminal quota error', () => {
it('should render switch, upgrade, and stop options for paid tiers', () => {
it('should render switch, upgrade, and stop options for LOGIN_WITH_GOOGLE', () => {
const { unmount } = render(
<ProQuotaDialog
failedModel="gemini-2.5-pro"
@@ -70,6 +71,7 @@ describe('ProQuotaDialog', () => {
message="paid tier quota error"
isTerminalQuotaError={true}
isModelNotFoundError={false}
authType={AuthType.LOGIN_WITH_GOOGLE}
onChoice={mockOnChoice}
/>,
);
@@ -99,6 +101,39 @@ describe('ProQuotaDialog', () => {
unmount();
});
it('should NOT render upgrade option for USE_GEMINI', () => {
const { unmount } = render(
<ProQuotaDialog
failedModel="gemini-2.5-pro"
fallbackModel="gemini-2.5-flash"
message="paid tier quota error"
isTerminalQuotaError={true}
isModelNotFoundError={false}
authType={AuthType.USE_GEMINI}
onChoice={mockOnChoice}
/>,
);
expect(RadioButtonSelect).toHaveBeenCalledWith(
expect.objectContaining({
items: [
{
label: 'Switch to gemini-2.5-flash',
value: 'retry_always',
key: 'retry_always',
},
{
label: 'Stop',
value: 'retry_later',
key: 'retry_later',
},
],
}),
undefined,
);
unmount();
});
it('should render "Keep trying" and "Stop" options when failed model and fallback model are the same', () => {
const { unmount } = render(
<ProQuotaDialog
@@ -130,7 +165,7 @@ describe('ProQuotaDialog', () => {
unmount();
});
it('should render switch, upgrade, and stop options for free tier', () => {
it('should render switch, upgrade, and stop options for LOGIN_WITH_GOOGLE (free tier)', () => {
const { unmount } = render(
<ProQuotaDialog
failedModel="gemini-2.5-pro"
@@ -138,6 +173,7 @@ describe('ProQuotaDialog', () => {
message="free tier quota error"
isTerminalQuotaError={true}
isModelNotFoundError={false}
authType={AuthType.LOGIN_WITH_GOOGLE}
onChoice={mockOnChoice}
/>,
);
@@ -204,7 +240,7 @@ describe('ProQuotaDialog', () => {
});
describe('when it is a model not found error', () => {
it('should render switch and stop options regardless of tier', () => {
it('should render switch, upgrade, and stop options for LOGIN_WITH_GOOGLE', () => {
const { unmount } = render(
<ProQuotaDialog
failedModel="gemini-3-pro-preview"
@@ -212,6 +248,7 @@ describe('ProQuotaDialog', () => {
message="You don't have access to gemini-3-pro-preview yet."
isTerminalQuotaError={false}
isModelNotFoundError={true}
authType={AuthType.LOGIN_WITH_GOOGLE}
onChoice={mockOnChoice}
/>,
);
@@ -241,7 +278,7 @@ describe('ProQuotaDialog', () => {
unmount();
});
it('should render switch and stop options for paid tier as well', () => {
it('should NOT render upgrade option for USE_GEMINI', () => {
const { unmount } = render(
<ProQuotaDialog
failedModel="gemini-3-pro-preview"
@@ -249,6 +286,7 @@ describe('ProQuotaDialog', () => {
message="You don't have access to gemini-3-pro-preview yet."
isTerminalQuotaError={false}
isModelNotFoundError={true}
authType={AuthType.USE_GEMINI}
onChoice={mockOnChoice}
/>,
);
@@ -261,11 +299,6 @@ describe('ProQuotaDialog', () => {
value: 'retry_always',
key: 'retry_always',
},
{
label: 'Upgrade for higher limits',
value: 'upgrade',
key: 'upgrade',
},
{
label: 'Stop',
value: 'retry_later',

View File

@@ -8,6 +8,7 @@ import type React from 'react';
import { Box, Text } from 'ink';
import { RadioButtonSelect } from './shared/RadioButtonSelect.js';
import { theme } from '../semantic-colors.js';
import { AuthType } from '@google/gemini-cli-core';
interface ProQuotaDialogProps {
failedModel: string;
@@ -15,6 +16,7 @@ interface ProQuotaDialogProps {
message: string;
isTerminalQuotaError: boolean;
isModelNotFoundError?: boolean;
authType?: AuthType;
onChoice: (
choice: 'retry_later' | 'retry_once' | 'retry_always' | 'upgrade',
) => void;
@@ -26,6 +28,7 @@ export function ProQuotaDialog({
message,
isTerminalQuotaError,
isModelNotFoundError,
authType,
onChoice,
}: ProQuotaDialogProps): React.JSX.Element {
let items;
@@ -51,11 +54,15 @@ export function ProQuotaDialog({
value: 'retry_always' as const,
key: 'retry_always',
},
{
label: 'Upgrade for higher limits',
value: 'upgrade' as const,
key: 'upgrade',
},
...(authType === AuthType.LOGIN_WITH_GOOGLE
? [
{
label: 'Upgrade for higher limits',
value: 'upgrade' as const,
key: 'upgrade',
},
]
: []),
{
label: `Stop`,
value: 'retry_later' as const,