MediaMetz/gemini-cli
1
0
Fork 0
mirror of https://github.com/google-gemini/gemini-cli.git synced 2026-04-21 18:44:30 -07:00
Code Issues Packages Projects Releases Wiki Activity

Let users know when falling back to flash, and update the error messa… (#12640)

Browse Source
This commit is contained in:
Sehoon Shon
2025-11-06 13:43:21 -05:00
committed by GitHub
parent c585470a71
commit 31b34b11ab
11 changed files with 174 additions and 225 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/cli/src/ui/hooks/useQuotaAndFallback.test.ts
+117 -103
View File
@@ -28,7 +28,7 @@ import {
} from '@google/gemini-cli-core';
import { useQuotaAndFallback } from './useQuotaAndFallback.js';
import type { UseHistoryManagerReturn } from './useHistoryManager.js';
import { AuthState, MessageType } from '../types.js';
import { MessageType } from '../types.js';
// Use a type alias for SpyInstance as it's not directly exported
type SpyInstance = ReturnType<typeof vi.spyOn>;
@@ -36,7 +36,6 @@ type SpyInstance = ReturnType<typeof vi.spyOn>;
describe('useQuotaAndFallback', () => {
let mockConfig: Config;
let mockHistoryManager: UseHistoryManagerReturn;
let mockSetAuthState: Mock;
let mockSetModelSwitchedFromQuotaError: Mock;
let setFallbackHandlerSpy: SpyInstance;
let mockGoogleApiError: GoogleApiError;
@@ -62,7 +61,6 @@ describe('useQuotaAndFallback', () => {
clearItems: vi.fn(),
loadHistory: vi.fn(),
};
mockSetAuthState = vi.fn();
mockSetModelSwitchedFromQuotaError = vi.fn();
setFallbackHandlerSpy = vi.spyOn(mockConfig, 'setFallbackModelHandler');
@@ -79,7 +77,6 @@ describe('useQuotaAndFallback', () => {
config: mockConfig,
historyManager: mockHistoryManager,
userTier: UserTierId.FREE,
setAuthState: mockSetAuthState,
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
}),
);
@@ -99,7 +96,6 @@ describe('useQuotaAndFallback', () => {
config: mockConfig,
historyManager: mockHistoryManager,
userTier: props.userTier,
setAuthState: mockSetAuthState,
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
}),
{ initialProps: { userTier } },
@@ -129,97 +125,14 @@ describe('useQuotaAndFallback', () => {
expect(mockHistoryManager.addItem).not.toHaveBeenCalled();
});
describe('Automatic Fallback Scenarios', () => {
const testCases = [
{
description: 'other error for FREE tier',
tier: UserTierId.FREE,
error: new Error('some error'),
expectedMessageSnippets: [
'Automatically switching from model-A to model-B for faster responses',
'upgrade to a Gemini Code Assist Standard or Enterprise plan',
],
},
{
description: 'other error for LEGACY tier',
tier: UserTierId.LEGACY, // Paid tier
error: new Error('some error'),
expectedMessageSnippets: [
'Automatically switching from model-A to model-B for faster responses',
'switch to using a paid API key from AI Studio',
],
},
{
description: 'retryable quota error for FREE tier',
tier: UserTierId.FREE,
error: new RetryableQuotaError(
'retryable quota',
mockGoogleApiError,
5,
),
expectedMessageSnippets: [
'Your requests are being throttled right now due to server being at capacity for model-A',
'Automatically switching from model-A to model-B',
'upgrading to a Gemini Code Assist Standard or Enterprise plan',
],
},
{
description: 'retryable quota error for LEGACY tier',
tier: UserTierId.LEGACY, // Paid tier
error: new RetryableQuotaError(
'retryable quota',
mockGoogleApiError,
5,
),
expectedMessageSnippets: [
'Your requests are being throttled right now due to server being at capacity for model-A',
'Automatically switching from model-A to model-B',
'switch to using a paid API key from AI Studio',
],
},
];
for (const {
description,
tier,
error,
expectedMessageSnippets,
} of testCases) {
it(`should handle ${description} correctly`, async () => {
const handler = getRegisteredHandler(tier);
let result: FallbackIntent | null;
await act(async () => {
result = await handler('model-A', 'model-B', error);
});
// Automatic fallbacks should return 'stop'
expect(result!).toBe('stop');
expect(mockHistoryManager.addItem).toHaveBeenCalledWith(
expect.objectContaining({ type: MessageType.INFO }),
expect.any(Number),
);
const message = (mockHistoryManager.addItem as Mock).mock.calls[0][0]
.text;
for (const snippet of expectedMessageSnippets) {
expect(message).toContain(snippet);
}
expect(mockSetModelSwitchedFromQuotaError).toHaveBeenCalledWith(true);
expect(mockConfig.setQuotaErrorOccurred).toHaveBeenCalledWith(true);
});
}
});
describe('Interactive Fallback (Pro Quota Error)', () => {
describe('Interactive Fallback', () => {
// Pro Quota Errors
it('should set an interactive request and wait for user choice', async () => {
const { result } = renderHook(() =>
useQuotaAndFallback({
config: mockConfig,
historyManager: mockHistoryManager,
userTier: UserTierId.FREE,
setAuthState: mockSetAuthState,
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
}),
);
@@ -243,7 +156,7 @@ describe('useQuotaAndFallback', () => {
// Simulate the user choosing to continue with the fallback model
await act(() => {
result.current.handleProQuotaChoice('continue');
result.current.handleProQuotaChoice('retry');
});
// The original promise from the handler should now resolve
@@ -260,7 +173,6 @@ describe('useQuotaAndFallback', () => {
config: mockConfig,
historyManager: mockHistoryManager,
userTier: UserTierId.FREE,
setAuthState: mockSetAuthState,
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
}),
);
@@ -294,13 +206,120 @@ describe('useQuotaAndFallback', () => {
expect(result.current.proQuotaRequest).toBe(firstRequest);
await act(() => {
result.current.handleProQuotaChoice('continue');
result.current.handleProQuotaChoice('retry');
});
const intent1 = await promise1!;
expect(intent1).toBe('retry');
expect(result.current.proQuotaRequest).toBeNull();
});
// Non-Quota error test cases
const testCases = [
{
description: 'other error for FREE tier',
tier: UserTierId.FREE,
error: new Error('some error'),
expectedMessageSnippets: [
'🚦Pardon Our Congestion! It looks like model-A is very popular at the moment.',
'Please retry again later.',
],
},
{
description: 'other error for LEGACY tier',
tier: UserTierId.LEGACY, // Paid tier
error: new Error('some error'),
expectedMessageSnippets: [
'🚦Pardon Our Congestion! It looks like model-A is very popular at the moment.',
'Please retry again later.',
],
},
{
description: 'retryable quota error for FREE tier',
tier: UserTierId.FREE,
error: new RetryableQuotaError(
'retryable quota',
mockGoogleApiError,
5,
),
expectedMessageSnippets: [
'🚦Pardon Our Congestion! It looks like model-A is very popular at the moment.',
'Please retry again later.',
],
},
{
description: 'retryable quota error for LEGACY tier',
tier: UserTierId.LEGACY, // Paid tier
error: new RetryableQuotaError(
'retryable quota',
mockGoogleApiError,
5,
),
expectedMessageSnippets: [
'🚦Pardon Our Congestion! It looks like model-A is very popular at the moment.',
'Please retry again later.',
],
},
];
for (const {
description,
tier,
error,
expectedMessageSnippets,
} of testCases) {
it(`should handle ${description} correctly`, async () => {
const { result } = renderHook(
(props) =>
useQuotaAndFallback({
config: mockConfig,
historyManager: mockHistoryManager,
userTier: props.tier,
setModelSwitchedFromQuotaError:
mockSetModelSwitchedFromQuotaError,
}),
{ initialProps: { tier } },
);
const handler = setFallbackHandlerSpy.mock
.calls[0][0] as FallbackModelHandler;
// Call the handler but do not await it, to check the intermediate state
let promise: Promise<FallbackIntent | null>;
await act(() => {
promise = handler('model-A', 'model-B', error);
});
// The hook should now have a pending request for the UI to handle
expect(result.current.proQuotaRequest).not.toBeNull();
expect(result.current.proQuotaRequest?.failedModel).toBe('model-A');
// Check that the correct initial message was added
expect(mockHistoryManager.addItem).toHaveBeenCalledWith(
expect.objectContaining({ type: MessageType.INFO }),
expect.any(Number),
);
const message = (mockHistoryManager.addItem as Mock).mock.calls[0][0]
.text;
for (const snippet of expectedMessageSnippets) {
expect(message).toContain(snippet);
}
// Simulate the user choosing to continue with the fallback model
await act(() => {
result.current.handleProQuotaChoice('retry');
});
expect(mockSetModelSwitchedFromQuotaError).toHaveBeenCalledWith(true);
// The original promise from the handler should now resolve
const intent = await promise!;
expect(intent).toBe('retry');
// The pending request should be cleared from the state
expect(result.current.proQuotaRequest).toBeNull();
expect(mockConfig.setQuotaErrorOccurred).toHaveBeenCalledWith(true);
});
}
});
});
@@ -311,26 +330,23 @@ describe('useQuotaAndFallback', () => {
config: mockConfig,
historyManager: mockHistoryManager,
userTier: UserTierId.FREE,
setAuthState: mockSetAuthState,
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
}),
);
act(() => {
result.current.handleProQuotaChoice('auth');
result.current.handleProQuotaChoice('retry_later');
});
expect(mockSetAuthState).not.toHaveBeenCalled();
expect(mockHistoryManager.addItem).not.toHaveBeenCalled();
});
it('should resolve intent to "auth" and trigger auth state update', async () => {
it('should resolve intent to "retry_later"', async () => {
const { result } = renderHook(() =>
useQuotaAndFallback({
config: mockConfig,
historyManager: mockHistoryManager,
userTier: UserTierId.FREE,
setAuthState: mockSetAuthState,
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
}),
);
@@ -347,12 +363,11 @@ describe('useQuotaAndFallback', () => {
});
await act(() => {
result.current.handleProQuotaChoice('auth');
result.current.handleProQuotaChoice('retry_later');
});
const intent = await promise!;
expect(intent).toBe('auth');
expect(mockSetAuthState).toHaveBeenCalledWith(AuthState.Updating);
expect(intent).toBe('retry_later');
expect(result.current.proQuotaRequest).toBeNull();
});
@@ -362,7 +377,6 @@ describe('useQuotaAndFallback', () => {
config: mockConfig,
historyManager: mockHistoryManager,
userTier: UserTierId.FREE,
setAuthState: mockSetAuthState,
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
}),
);
@@ -380,7 +394,7 @@ describe('useQuotaAndFallback', () => {
});
await act(() => {
result.current.handleProQuotaChoice('continue');
result.current.handleProQuotaChoice('retry');
});
const intent = await promise!;