MediaMetz/gemini-cli
1
0
Fork 0
mirror of https://github.com/google-gemini/gemini-cli.git synced 2026-03-24 13:01:29 -07:00
Code Issues Packages Projects Releases Wiki Activity

fix(patch): cherry-pick 996c9f5 to release/v0.9.0-preview.4-pr-11164 [CONFLICTS] (#11166)

Browse Source 
Co-authored-by: Gaurav Ghosh <gaghosh@google.com>
This commit is contained in:
gemini-cli-robot
2025-10-14 17:38:02 -07:00
committed by GitHub
parent 78acfa4416
commit 3b6d90cfc6
13 changed files with 821 additions and 1013 deletions
Download Patch File Download Diff File Expand all files Collapse all files

67
packages/cli/src/ui/hooks/useQuotaAndFallback.test.ts
View File

@@ -19,14 +19,25 @@ import {
type FallbackModelHandler,
UserTierId,
AuthType,
TerminalQuotaError,
isGenericQuotaExceededError,
isProQuotaExceededError,
makeFakeConfig,
type GoogleApiError,
} from '@google/gemini-cli-core';
import { useQuotaAndFallback } from './useQuotaAndFallback.js';
import type { UseHistoryManagerReturn } from './useHistoryManager.js';
import { AuthState, MessageType } from '../types.js';
// Mock the error checking functions from the core package to control test scenarios
vi.mock('@google/gemini-cli-core', async (importOriginal) => {
const original =
await importOriginal<typeof import('@google/gemini-cli-core')>();
return {
...original,
isGenericQuotaExceededError: vi.fn(),
isProQuotaExceededError: vi.fn(),
};
});
// Use a type alias for SpyInstance as it's not directly exported
type SpyInstance = ReturnType<typeof vi.spyOn>;
@@ -36,15 +47,12 @@ describe('useQuotaAndFallback', () => {
let mockSetAuthState: Mock;
let mockSetModelSwitchedFromQuotaError: Mock;
let setFallbackHandlerSpy: SpyInstance;
let mockGoogleApiError: GoogleApiError;
const mockedIsGenericQuotaExceededError = isGenericQuotaExceededError as Mock;
const mockedIsProQuotaExceededError = isProQuotaExceededError as Mock;
beforeEach(() => {
mockConfig = makeFakeConfig();
mockGoogleApiError = {
code: 429,
message: 'mock error',
details: [],
};
// Spy on the method that requires the private field and mock its return.
// This is cleaner than modifying the config class for tests.
@@ -64,6 +72,9 @@ describe('useQuotaAndFallback', () => {
setFallbackHandlerSpy = vi.spyOn(mockConfig, 'setFallbackModelHandler');
vi.spyOn(mockConfig, 'setQuotaErrorOccurred');
mockedIsGenericQuotaExceededError.mockReturnValue(false);
mockedIsProQuotaExceededError.mockReturnValue(false);
});
afterEach(() => {
@@ -128,6 +139,22 @@ describe('useQuotaAndFallback', () => {
describe('Automatic Fallback Scenarios', () => {
const testCases = [
{
errorType: 'generic',
tier: UserTierId.FREE,
expectedMessageSnippets: [
'Automatically switching from model-A to model-B',
'upgrade to a Gemini Code Assist Standard or Enterprise plan',
],
},
{
errorType: 'generic',
tier: UserTierId.STANDARD, // Paid tier
expectedMessageSnippets: [
'Automatically switching from model-A to model-B',
'switch to using a paid API key from AI Studio',
],
},
{
errorType: 'other',
tier: UserTierId.FREE,
@@ -148,11 +175,15 @@ describe('useQuotaAndFallback', () => {
for (const { errorType, tier, expectedMessageSnippets } of testCases) {
it(`should handle ${errorType} error for ${tier} tier correctly`, async () => {
mockedIsGenericQuotaExceededError.mockReturnValue(
errorType === 'generic',
);
const handler = getRegisteredHandler(tier);
const result = await handler(
'model-A',
'model-B',
new Error('some error'),
new Error('quota exceeded'),
);
// Automatic fallbacks should return 'stop'
@@ -176,6 +207,10 @@ describe('useQuotaAndFallback', () => {
});
describe('Interactive Fallback (Pro Quota Error)', () => {
beforeEach(() => {
mockedIsProQuotaExceededError.mockReturnValue(true);
});
it('should set an interactive request and wait for user choice', async () => {
const { result } = renderHook(() =>
useQuotaAndFallback({
@@ -194,7 +229,7 @@ describe('useQuotaAndFallback', () => {
const promise = handler(
'gemini-pro',
'gemini-flash',
new TerminalQuotaError('pro quota', mockGoogleApiError),
new Error('pro quota'),
);
await act(async () => {});
@@ -233,7 +268,7 @@ describe('useQuotaAndFallback', () => {
const promise1 = handler(
'gemini-pro',
'gemini-flash',
new TerminalQuotaError('pro quota 1', mockGoogleApiError),
new Error('pro quota 1'),
);
await act(async () => {});
@@ -243,7 +278,7 @@ describe('useQuotaAndFallback', () => {
const result2 = await handler(
'gemini-pro',
'gemini-flash',
new TerminalQuotaError('pro quota 2', mockGoogleApiError),
new Error('pro quota 2'),
);
// The lock should have stopped the second request
@@ -262,6 +297,10 @@ describe('useQuotaAndFallback', () => {
});
describe('handleProQuotaChoice', () => {
beforeEach(() => {
mockedIsProQuotaExceededError.mockReturnValue(true);
});
it('should do nothing if there is no pending pro quota request', () => {
const { result } = renderHook(() =>
useQuotaAndFallback({
@@ -297,7 +336,7 @@ describe('useQuotaAndFallback', () => {
const promise = handler(
'gemini-pro',
'gemini-flash',
new TerminalQuotaError('pro quota', mockGoogleApiError),
new Error('pro quota'),
);
await act(async () => {}); // Allow state to update
@@ -328,7 +367,7 @@ describe('useQuotaAndFallback', () => {
const promise = handler(
'gemini-pro',
'gemini-flash',
new TerminalQuotaError('pro quota', mockGoogleApiError),
new Error('pro quota'),
);
await act(async () => {}); // Allow state to update

20
packages/cli/src/ui/hooks/useQuotaAndFallback.ts
View File

@@ -9,7 +9,8 @@ import {
type Config,
type FallbackModelHandler,
type FallbackIntent,
TerminalQuotaError,
isGenericQuotaExceededError,
isProQuotaExceededError,
UserTierId,
} from '@google/gemini-cli-core';
import { useCallback, useEffect, useRef, useState } from 'react';
@@ -62,7 +63,7 @@ export function useQuotaAndFallback({
let message: string;
if (error instanceof TerminalQuotaError) {
if (error && isProQuotaExceededError(error)) {
// Pro Quota specific messages (Interactive)
if (isPaidTier) {
message = `⚡ You have reached your daily ${failedModel} quota limit.
@@ -73,6 +74,19 @@ export function useQuotaAndFallback({
⚡ You can choose to authenticate with a paid API key or continue with the fallback model.
⚡ To increase your limits, upgrade to a Gemini Code Assist Standard or Enterprise plan with higher limits at https://goo.gle/set-up-gemini-code-assist
⚡ Or you can utilize a Gemini API Key. See: https://goo.gle/gemini-cli-docs-auth#gemini-api-key
⚡ You can switch authentication methods by typing /auth`;
}
} else if (error && isGenericQuotaExceededError(error)) {
// Generic Quota (Automatic fallback)
const actionMessage = `⚡ You have reached your daily quota limit.\n⚡ Automatically switching from ${failedModel} to ${fallbackModel} for the remainder of this session.`;
if (isPaidTier) {
message = `${actionMessage}
⚡ To continue accessing the ${failedModel} model today, consider using /auth to switch to using a paid API key from AI Studio at https://aistudio.google.com/apikey`;
} else {
message = `${actionMessage}
⚡ To increase your limits, upgrade to a Gemini Code Assist Standard or Enterprise plan with higher limits at https://goo.gle/set-up-gemini-code-assist
⚡ Or you can utilize a Gemini API Key. See: https://goo.gle/gemini-cli-docs-auth#gemini-api-key
⚡ You can switch authentication methods by typing /auth`;
}
} else {
@@ -105,7 +119,7 @@ export function useQuotaAndFallback({
config.setQuotaErrorOccurred(true);
// Interactive Fallback for Pro quota
if (error instanceof TerminalQuotaError) {
if (error && isProQuotaExceededError(error)) {
if (isDialogPending.current) {
return 'stop'; // A dialog is already active, so just stop this request.
}