mirror of
https://github.com/google-gemini/gemini-cli.git
synced 2026-03-10 14:10:37 -07:00
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:
@@ -135,6 +135,7 @@ export const DialogManager = ({
|
||||
isModelNotFoundError={
|
||||
!!uiState.quota.proQuotaRequest.isModelNotFoundError
|
||||
}
|
||||
authType={uiState.quota.proQuotaRequest.authType}
|
||||
onChoice={uiActions.handleProQuotaChoice}
|
||||
/>
|
||||
);
|
||||
|
||||
@@ -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',
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
|
||||
|
||||
@@ -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', () => {
|
||||
|
||||
@@ -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,
|
||||
});
|
||||
},
|
||||
);
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user