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,

View File

@@ -24,6 +24,7 @@ import type {
ApprovalMode,
UserTierId,
IdeInfo,
AuthType,
FallbackIntent,
ValidationIntent,
AgentDefinition,
@@ -42,6 +43,7 @@ export interface ProQuotaDialogRequest {
message: string;
isTerminalQuotaError: boolean;
isModelNotFoundError?: boolean;
authType?: AuthType;
resolve: (intent: FallbackIntent) => void;
}

View File

@@ -96,9 +96,13 @@ describe('useQuotaAndFallback', () => {
});
describe('Fallback Handler Logic', () => {
// Helper function to render the hook and extract the registered handler
const getRegisteredHandler = (): FallbackModelHandler => {
renderHook(() =>
it('should show fallback dialog but omit switch to API key message if authType is not LOGIN_WITH_GOOGLE', async () => {
// Override the default mock from beforeEach for this specific test
vi.spyOn(mockConfig, 'getContentGeneratorConfig').mockReturnValue({
authType: AuthType.USE_GEMINI,
});
const { result } = renderHook(() =>
useQuotaAndFallback({
config: mockConfig,
historyManager: mockHistoryManager,
@@ -107,20 +111,24 @@ describe('useQuotaAndFallback', () => {
onShowAuthSelection: mockOnShowAuthSelection,
}),
);
return setFallbackHandlerSpy.mock.calls[0][0] as FallbackModelHandler;
};
it('should return null and take no action if authType is not LOGIN_WITH_GOOGLE', async () => {
// Override the default mock from beforeEach for this specific test
vi.spyOn(mockConfig, 'getContentGeneratorConfig').mockReturnValue({
authType: AuthType.USE_GEMINI,
const handler = setFallbackHandlerSpy.mock
.calls[0][0] as FallbackModelHandler;
const error = new TerminalQuotaError(
'pro quota',
mockGoogleApiError,
1000 * 60 * 5,
);
act(() => {
void handler('gemini-pro', 'gemini-flash', error);
});
const handler = getRegisteredHandler();
const result = await handler('gemini-pro', 'gemini-flash', new Error());
expect(result).toBeNull();
expect(mockHistoryManager.addItem).not.toHaveBeenCalled();
expect(result.current.proQuotaRequest).not.toBeNull();
expect(result.current.proQuotaRequest?.message).not.toContain(
'/auth to switch to API key.',
);
});
describe('Interactive Fallback', () => {

View File

@@ -55,14 +55,7 @@ export function useQuotaAndFallback({
fallbackModel,
error,
): Promise<FallbackIntent | null> => {
// Fallbacks are currently only handled for OAuth users.
const contentGeneratorConfig = config.getContentGeneratorConfig();
if (
!contentGeneratorConfig ||
contentGeneratorConfig.authType !== AuthType.LOGIN_WITH_GOOGLE
) {
return null;
}
let message: string;
let isTerminalQuotaError = false;
@@ -78,7 +71,9 @@ export function useQuotaAndFallback({
error.retryDelayMs ? getResetTimeMessage(error.retryDelayMs) : null,
`/stats model for usage details`,
`/model to switch models.`,
`/auth to switch to API key.`,
contentGeneratorConfig?.authType === AuthType.LOGIN_WITH_GOOGLE
? `/auth to switch to API key.`
: null,
].filter(Boolean);
message = messageLines.join('\n');
} else if (error instanceof ModelNotFoundError) {
@@ -122,6 +117,7 @@ export function useQuotaAndFallback({
message,
isTerminalQuotaError,
isModelNotFoundError,
authType: contentGeneratorConfig?.authType,
});
},
);

View File

@@ -77,6 +77,12 @@ export function getAuthTypeFromEnv(): AuthType | undefined {
if (process.env['GEMINI_API_KEY']) {
return AuthType.USE_GEMINI;
}
if (
process.env['CLOUD_SHELL'] === 'true' ||
process.env['GEMINI_CLI_USE_COMPUTE_ADC'] === 'true'
) {
return AuthType.COMPUTE_ADC;
}
return undefined;
}