mirror of
https://github.com/google-gemini/gemini-cli.git
synced 2026-03-16 17:11:04 -07:00
refactor: Centralize and improve model fallback handling (#7634)
This commit is contained in:
@@ -4,31 +4,55 @@
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
|
||||
import {
|
||||
describe,
|
||||
it,
|
||||
expect,
|
||||
vi,
|
||||
beforeEach,
|
||||
afterEach,
|
||||
type Mock,
|
||||
} from 'vitest';
|
||||
import { render, cleanup } from 'ink-testing-library';
|
||||
import { AppContainer } from './AppContainer.js';
|
||||
import { type Config, makeFakeConfig } from '@google/gemini-cli-core';
|
||||
import type { LoadedSettings } from '../config/settings.js';
|
||||
import type { InitializationResult } from '../core/initializer.js';
|
||||
import { useQuotaAndFallback } from './hooks/useQuotaAndFallback.js';
|
||||
import { UIStateContext, type UIState } from './contexts/UIStateContext.js';
|
||||
import {
|
||||
UIActionsContext,
|
||||
type UIActions,
|
||||
} from './contexts/UIActionsContext.js';
|
||||
import { useContext } from 'react';
|
||||
|
||||
// Helper component will read the context values provided by AppContainer
|
||||
// so we can assert against them in our tests.
|
||||
let capturedUIState: UIState;
|
||||
let capturedUIActions: UIActions;
|
||||
function TestContextConsumer() {
|
||||
capturedUIState = useContext(UIStateContext)!;
|
||||
capturedUIActions = useContext(UIActionsContext)!;
|
||||
return null;
|
||||
}
|
||||
|
||||
// Mock App component to isolate AppContainer testing
|
||||
vi.mock('./App.js', () => ({
|
||||
App: () => 'App Component',
|
||||
App: TestContextConsumer,
|
||||
}));
|
||||
|
||||
// Mock all the hooks and utilities
|
||||
vi.mock('./hooks/useHistory.js');
|
||||
vi.mock('./hooks/useQuotaAndFallback.js');
|
||||
vi.mock('./hooks/useHistoryManager.js');
|
||||
vi.mock('./hooks/useThemeCommand.js');
|
||||
vi.mock('./hooks/useAuthCommand.js');
|
||||
vi.mock('./auth/useAuth.js');
|
||||
vi.mock('./hooks/useEditorSettings.js');
|
||||
vi.mock('./hooks/useSettingsCommand.js');
|
||||
vi.mock('./hooks/useSlashCommandProcessor.js');
|
||||
vi.mock('./hooks/slashCommandProcessor.js');
|
||||
vi.mock('./hooks/useConsoleMessages.js');
|
||||
vi.mock('./hooks/useTerminalSize.js', () => ({
|
||||
useTerminalSize: vi.fn(() => ({ columns: 80, rows: 24 })),
|
||||
}));
|
||||
vi.mock('./hooks/useGeminiStream.js');
|
||||
vi.mock('./hooks/useVim.js');
|
||||
vi.mock('./hooks/vim.js');
|
||||
vi.mock('./hooks/useFocus.js');
|
||||
vi.mock('./hooks/useBracketedPaste.js');
|
||||
vi.mock('./hooks/useKeypress.js');
|
||||
@@ -40,7 +64,7 @@ vi.mock('./hooks/useWorkspaceMigration.js');
|
||||
vi.mock('./hooks/useGitBranchName.js');
|
||||
vi.mock('./contexts/VimModeContext.js');
|
||||
vi.mock('./contexts/SessionContext.js');
|
||||
vi.mock('./hooks/useTextBuffer.js');
|
||||
vi.mock('./components/shared/text-buffer.js');
|
||||
vi.mock('./hooks/useLogger.js');
|
||||
|
||||
// Mock external utilities
|
||||
@@ -49,14 +73,153 @@ vi.mock('../utils/handleAutoUpdate.js');
|
||||
vi.mock('./utils/ConsolePatcher.js');
|
||||
vi.mock('../utils/cleanup.js');
|
||||
|
||||
import { useHistory } from './hooks/useHistoryManager.js';
|
||||
import { useThemeCommand } from './hooks/useThemeCommand.js';
|
||||
import { useAuthCommand } from './auth/useAuth.js';
|
||||
import { useEditorSettings } from './hooks/useEditorSettings.js';
|
||||
import { useSettingsCommand } from './hooks/useSettingsCommand.js';
|
||||
import { useSlashCommandProcessor } from './hooks/slashCommandProcessor.js';
|
||||
import { useConsoleMessages } from './hooks/useConsoleMessages.js';
|
||||
import { useGeminiStream } from './hooks/useGeminiStream.js';
|
||||
import { useVim } from './hooks/vim.js';
|
||||
import { useFolderTrust } from './hooks/useFolderTrust.js';
|
||||
import { useMessageQueue } from './hooks/useMessageQueue.js';
|
||||
import { useAutoAcceptIndicator } from './hooks/useAutoAcceptIndicator.js';
|
||||
import { useWorkspaceMigration } from './hooks/useWorkspaceMigration.js';
|
||||
import { useGitBranchName } from './hooks/useGitBranchName.js';
|
||||
import { useVimMode } from './contexts/VimModeContext.js';
|
||||
import { useSessionStats } from './contexts/SessionContext.js';
|
||||
import { useTextBuffer } from './components/shared/text-buffer.js';
|
||||
import { useLogger } from './hooks/useLogger.js';
|
||||
import { useLoadingIndicator } from './hooks/useLoadingIndicator.js';
|
||||
|
||||
describe('AppContainer State Management', () => {
|
||||
let mockConfig: Config;
|
||||
let mockSettings: LoadedSettings;
|
||||
let mockInitResult: InitializationResult;
|
||||
|
||||
// Create typed mocks for all hooks
|
||||
const mockedUseQuotaAndFallback = useQuotaAndFallback as Mock;
|
||||
const mockedUseHistory = useHistory as Mock;
|
||||
const mockedUseThemeCommand = useThemeCommand as Mock;
|
||||
const mockedUseAuthCommand = useAuthCommand as Mock;
|
||||
const mockedUseEditorSettings = useEditorSettings as Mock;
|
||||
const mockedUseSettingsCommand = useSettingsCommand as Mock;
|
||||
const mockedUseSlashCommandProcessor = useSlashCommandProcessor as Mock;
|
||||
const mockedUseConsoleMessages = useConsoleMessages as Mock;
|
||||
const mockedUseGeminiStream = useGeminiStream as Mock;
|
||||
const mockedUseVim = useVim as Mock;
|
||||
const mockedUseFolderTrust = useFolderTrust as Mock;
|
||||
const mockedUseMessageQueue = useMessageQueue as Mock;
|
||||
const mockedUseAutoAcceptIndicator = useAutoAcceptIndicator as Mock;
|
||||
const mockedUseWorkspaceMigration = useWorkspaceMigration as Mock;
|
||||
const mockedUseGitBranchName = useGitBranchName as Mock;
|
||||
const mockedUseVimMode = useVimMode as Mock;
|
||||
const mockedUseSessionStats = useSessionStats as Mock;
|
||||
const mockedUseTextBuffer = useTextBuffer as Mock;
|
||||
const mockedUseLogger = useLogger as Mock;
|
||||
const mockedUseLoadingIndicator = useLoadingIndicator as Mock;
|
||||
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
|
||||
capturedUIState = null!;
|
||||
capturedUIActions = null!;
|
||||
|
||||
// **Provide a default return value for EVERY mocked hook.**
|
||||
mockedUseQuotaAndFallback.mockReturnValue({
|
||||
proQuotaRequest: null,
|
||||
handleProQuotaChoice: vi.fn(),
|
||||
});
|
||||
mockedUseHistory.mockReturnValue({
|
||||
history: [],
|
||||
addItem: vi.fn(),
|
||||
updateItem: vi.fn(),
|
||||
clearItems: vi.fn(),
|
||||
loadHistory: vi.fn(),
|
||||
});
|
||||
mockedUseThemeCommand.mockReturnValue({
|
||||
isThemeDialogOpen: false,
|
||||
openThemeDialog: vi.fn(),
|
||||
handleThemeSelect: vi.fn(),
|
||||
handleThemeHighlight: vi.fn(),
|
||||
});
|
||||
mockedUseAuthCommand.mockReturnValue({
|
||||
authState: 'authenticated',
|
||||
setAuthState: vi.fn(),
|
||||
authError: null,
|
||||
onAuthError: vi.fn(),
|
||||
});
|
||||
mockedUseEditorSettings.mockReturnValue({
|
||||
isEditorDialogOpen: false,
|
||||
openEditorDialog: vi.fn(),
|
||||
handleEditorSelect: vi.fn(),
|
||||
exitEditorDialog: vi.fn(),
|
||||
});
|
||||
mockedUseSettingsCommand.mockReturnValue({
|
||||
isSettingsDialogOpen: false,
|
||||
openSettingsDialog: vi.fn(),
|
||||
closeSettingsDialog: vi.fn(),
|
||||
});
|
||||
mockedUseSlashCommandProcessor.mockReturnValue({
|
||||
handleSlashCommand: vi.fn(),
|
||||
slashCommands: [],
|
||||
pendingHistoryItems: [],
|
||||
commandContext: {},
|
||||
shellConfirmationRequest: null,
|
||||
confirmationRequest: null,
|
||||
});
|
||||
mockedUseConsoleMessages.mockReturnValue({
|
||||
consoleMessages: [],
|
||||
handleNewMessage: vi.fn(),
|
||||
clearConsoleMessages: vi.fn(),
|
||||
});
|
||||
mockedUseGeminiStream.mockReturnValue({
|
||||
streamingState: 'idle',
|
||||
submitQuery: vi.fn(),
|
||||
initError: null,
|
||||
pendingHistoryItems: [],
|
||||
thought: null,
|
||||
cancelOngoingRequest: vi.fn(),
|
||||
});
|
||||
mockedUseVim.mockReturnValue({ handleInput: vi.fn() });
|
||||
mockedUseFolderTrust.mockReturnValue({
|
||||
isFolderTrustDialogOpen: false,
|
||||
handleFolderTrustSelect: vi.fn(),
|
||||
isRestarting: false,
|
||||
});
|
||||
mockedUseMessageQueue.mockReturnValue({
|
||||
messageQueue: [],
|
||||
addMessage: vi.fn(),
|
||||
clearQueue: vi.fn(),
|
||||
getQueuedMessagesText: vi.fn().mockReturnValue(''),
|
||||
});
|
||||
mockedUseAutoAcceptIndicator.mockReturnValue(false);
|
||||
mockedUseWorkspaceMigration.mockReturnValue({
|
||||
showWorkspaceMigrationDialog: false,
|
||||
workspaceExtensions: [],
|
||||
onWorkspaceMigrationDialogOpen: vi.fn(),
|
||||
onWorkspaceMigrationDialogClose: vi.fn(),
|
||||
});
|
||||
mockedUseGitBranchName.mockReturnValue('main');
|
||||
mockedUseVimMode.mockReturnValue({
|
||||
isVimEnabled: false,
|
||||
toggleVimEnabled: vi.fn(),
|
||||
});
|
||||
mockedUseSessionStats.mockReturnValue({ stats: {} });
|
||||
mockedUseTextBuffer.mockReturnValue({
|
||||
text: '',
|
||||
setText: vi.fn(),
|
||||
// Add other properties if AppContainer uses them
|
||||
});
|
||||
mockedUseLogger.mockReturnValue({
|
||||
getPreviousUserMessages: vi.fn().mockResolvedValue([]),
|
||||
});
|
||||
mockedUseLoadingIndicator.mockReturnValue({
|
||||
elapsedTime: '0.0s',
|
||||
currentLoadingPhrase: '',
|
||||
});
|
||||
|
||||
// Mock Config
|
||||
mockConfig = makeFakeConfig();
|
||||
|
||||
@@ -325,7 +488,73 @@ describe('AppContainer State Management', () => {
|
||||
expect(() => unmount()).not.toThrow();
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
// TODO: Add comprehensive integration test once all hook mocks are complete
|
||||
// For now, the 14 passing unit tests provide good coverage of AppContainer functionality
|
||||
describe('Quota and Fallback Integration', () => {
|
||||
it('passes a null proQuotaRequest to UIStateContext by default', () => {
|
||||
// The default mock from beforeEach already sets proQuotaRequest to null
|
||||
render(
|
||||
<AppContainer
|
||||
config={mockConfig}
|
||||
settings={mockSettings}
|
||||
version="1.0.0"
|
||||
initializationResult={mockInitResult}
|
||||
/>,
|
||||
);
|
||||
|
||||
// Assert that the context value is as expected
|
||||
expect(capturedUIState.proQuotaRequest).toBeNull();
|
||||
});
|
||||
|
||||
it('passes a valid proQuotaRequest to UIStateContext when provided by the hook', () => {
|
||||
// Arrange: Create a mock request object that a UI dialog would receive
|
||||
const mockRequest = {
|
||||
failedModel: 'gemini-pro',
|
||||
fallbackModel: 'gemini-flash',
|
||||
resolve: vi.fn(),
|
||||
};
|
||||
mockedUseQuotaAndFallback.mockReturnValue({
|
||||
proQuotaRequest: mockRequest,
|
||||
handleProQuotaChoice: vi.fn(),
|
||||
});
|
||||
|
||||
// Act: Render the container
|
||||
render(
|
||||
<AppContainer
|
||||
config={mockConfig}
|
||||
settings={mockSettings}
|
||||
version="1.0.0"
|
||||
initializationResult={mockInitResult}
|
||||
/>,
|
||||
);
|
||||
|
||||
// Assert: The mock request is correctly passed through the context
|
||||
expect(capturedUIState.proQuotaRequest).toEqual(mockRequest);
|
||||
});
|
||||
|
||||
it('passes the handleProQuotaChoice function to UIActionsContext', () => {
|
||||
// Arrange: Create a mock handler function
|
||||
const mockHandler = vi.fn();
|
||||
mockedUseQuotaAndFallback.mockReturnValue({
|
||||
proQuotaRequest: null,
|
||||
handleProQuotaChoice: mockHandler,
|
||||
});
|
||||
|
||||
// Act: Render the container
|
||||
render(
|
||||
<AppContainer
|
||||
config={mockConfig}
|
||||
settings={mockSettings}
|
||||
version="1.0.0"
|
||||
initializationResult={mockInitResult}
|
||||
/>,
|
||||
);
|
||||
|
||||
// Assert: The action in the context is the mock handler we provided
|
||||
expect(capturedUIActions.handleProQuotaChoice).toBe(mockHandler);
|
||||
|
||||
// You can even verify that the plumbed function is callable
|
||||
capturedUIActions.handleProQuotaChoice('auth');
|
||||
expect(mockHandler).toHaveBeenCalledWith('auth');
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
@@ -24,18 +24,15 @@ import { MessageType, StreamingState } from './types.js';
|
||||
import {
|
||||
type EditorType,
|
||||
type Config,
|
||||
IdeClient,
|
||||
type DetectedIde,
|
||||
ideContext,
|
||||
type IdeContext,
|
||||
type UserTierId,
|
||||
DEFAULT_GEMINI_FLASH_MODEL,
|
||||
IdeClient,
|
||||
ideContext,
|
||||
getErrorMessage,
|
||||
getAllGeminiMdFilenames,
|
||||
UserTierId,
|
||||
AuthType,
|
||||
isProQuotaExceededError,
|
||||
isGenericQuotaExceededError,
|
||||
logFlashFallback,
|
||||
FlashFallbackEvent,
|
||||
clearCachedCredentialFile,
|
||||
} from '@google/gemini-cli-core';
|
||||
import { validateAuthMethod } from '../config/auth.js';
|
||||
@@ -44,6 +41,7 @@ import process from 'node:process';
|
||||
import { useHistory } from './hooks/useHistoryManager.js';
|
||||
import { useThemeCommand } from './hooks/useThemeCommand.js';
|
||||
import { useAuthCommand } from './auth/useAuth.js';
|
||||
import { useQuotaAndFallback } from './hooks/useQuotaAndFallback.js';
|
||||
import { useEditorSettings } from './hooks/useEditorSettings.js';
|
||||
import { useSettingsCommand } from './hooks/useSettingsCommand.js';
|
||||
import { useSlashCommandProcessor } from './hooks/slashCommandProcessor.js';
|
||||
@@ -123,12 +121,18 @@ export const AppContainer = (props: AppContainerProps) => {
|
||||
const [isTrustedFolder, setIsTrustedFolder] = useState<boolean | undefined>(
|
||||
config.isTrustedFolder(),
|
||||
);
|
||||
const [currentModel, setCurrentModel] = useState(config.getModel());
|
||||
|
||||
// Helper to determine the effective model, considering the fallback state.
|
||||
const getEffectiveModel = useCallback(() => {
|
||||
if (config.isInFallbackMode()) {
|
||||
return DEFAULT_GEMINI_FLASH_MODEL;
|
||||
}
|
||||
return config.getModel();
|
||||
}, [config]);
|
||||
|
||||
const [currentModel, setCurrentModel] = useState(getEffectiveModel());
|
||||
|
||||
const [userTier, setUserTier] = useState<UserTierId | undefined>(undefined);
|
||||
const [isProQuotaDialogOpen, setIsProQuotaDialogOpen] = useState(false);
|
||||
const [proQuotaDialogResolver, setProQuotaDialogResolver] = useState<
|
||||
((value: boolean) => void) | null
|
||||
>(null);
|
||||
|
||||
// Auto-accept indicator
|
||||
const showAutoAcceptIndicator = useAutoAcceptIndicator({
|
||||
@@ -167,18 +171,17 @@ export const AppContainer = (props: AppContainerProps) => {
|
||||
// Watch for model changes (e.g., from Flash fallback)
|
||||
useEffect(() => {
|
||||
const checkModelChange = () => {
|
||||
const configModel = config.getModel();
|
||||
if (configModel !== currentModel) {
|
||||
setCurrentModel(configModel);
|
||||
const effectiveModel = getEffectiveModel();
|
||||
if (effectiveModel !== currentModel) {
|
||||
setCurrentModel(effectiveModel);
|
||||
}
|
||||
};
|
||||
|
||||
// Check immediately and then periodically
|
||||
checkModelChange();
|
||||
const interval = setInterval(checkModelChange, 1000); // Check every second
|
||||
|
||||
return () => clearInterval(interval);
|
||||
}, [config, currentModel]);
|
||||
}, [config, currentModel, getEffectiveModel]);
|
||||
|
||||
const {
|
||||
consoleMessages,
|
||||
@@ -273,6 +276,14 @@ export const AppContainer = (props: AppContainerProps) => {
|
||||
config,
|
||||
);
|
||||
|
||||
const { proQuotaRequest, handleProQuotaChoice } = useQuotaAndFallback({
|
||||
config,
|
||||
historyManager,
|
||||
userTier,
|
||||
setAuthState,
|
||||
setModelSwitchedFromQuotaError,
|
||||
});
|
||||
|
||||
// Derive auth state variables for backward compatibility with UIStateContext
|
||||
const isAuthDialogOpen = authState === AuthState.Updating;
|
||||
const isAuthenticating = authState === AuthState.Unauthenticated;
|
||||
@@ -477,132 +488,6 @@ Logging in with Google... Please restart Gemini CLI to continue.
|
||||
}
|
||||
}, [config, historyManager, settings.merged]);
|
||||
|
||||
// Set up Flash fallback handler
|
||||
useEffect(() => {
|
||||
const flashFallbackHandler = async (
|
||||
currentModel: string,
|
||||
fallbackModel: string,
|
||||
error?: unknown,
|
||||
): Promise<boolean> => {
|
||||
// Check if we've already switched to the fallback model
|
||||
if (config.isInFallbackMode()) {
|
||||
// If we're already in fallback mode, don't show the dialog again
|
||||
return false;
|
||||
}
|
||||
|
||||
let message: string;
|
||||
|
||||
if (
|
||||
config.getContentGeneratorConfig().authType ===
|
||||
AuthType.LOGIN_WITH_GOOGLE
|
||||
) {
|
||||
// Use actual user tier if available; otherwise, default to FREE tier behavior (safe default)
|
||||
const isPaidTier =
|
||||
userTier === UserTierId.LEGACY || userTier === UserTierId.STANDARD;
|
||||
|
||||
// Check if this is a Pro quota exceeded error
|
||||
if (error && isProQuotaExceededError(error)) {
|
||||
if (isPaidTier) {
|
||||
message = `⚡ You have reached your daily ${currentModel} quota limit.
|
||||
⚡ You can choose to authenticate with a paid API key or continue with the fallback model.
|
||||
⚡ To continue accessing the ${currentModel} model today, consider using /auth to switch to using a paid API key from AI Studio at https://aistudio.google.com/apikey`;
|
||||
} else {
|
||||
message = `⚡ You have reached your daily ${currentModel} quota limit.
|
||||
⚡ 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)) {
|
||||
if (isPaidTier) {
|
||||
message = `⚡ You have reached your daily quota limit.
|
||||
⚡ Automatically switching from ${currentModel} to ${fallbackModel} for the remainder of this session.
|
||||
⚡ To continue accessing the ${currentModel} model today, consider using /auth to switch to using a paid API key from AI Studio at https://aistudio.google.com/apikey`;
|
||||
} else {
|
||||
message = `⚡ You have reached your daily quota limit.
|
||||
⚡ Automatically switching from ${currentModel} to ${fallbackModel} for the remainder of this session.
|
||||
⚡ 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 (isPaidTier) {
|
||||
// Default fallback message for other cases (like consecutive 429s)
|
||||
message = `⚡ Automatically switching from ${currentModel} to ${fallbackModel} for faster responses for the remainder of this session.
|
||||
⚡ Possible reasons for this are that you have received multiple consecutive capacity errors or you have reached your daily ${currentModel} quota limit
|
||||
⚡ To continue accessing the ${currentModel} model today, consider using /auth to switch to using a paid API key from AI Studio at https://aistudio.google.com/apikey`;
|
||||
} else {
|
||||
// Default fallback message for other cases (like consecutive 429s)
|
||||
message = `⚡ Automatically switching from ${currentModel} to ${fallbackModel} for faster responses for the remainder of this session.
|
||||
⚡ Possible reasons for this are that you have received multiple consecutive capacity errors or you have reached your daily ${currentModel} quota limit
|
||||
⚡ 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`;
|
||||
}
|
||||
}
|
||||
|
||||
// Add message to UI history
|
||||
historyManager.addItem(
|
||||
{
|
||||
type: MessageType.INFO,
|
||||
text: message,
|
||||
},
|
||||
Date.now(),
|
||||
);
|
||||
|
||||
// For Pro quota errors, show the dialog and wait for user's choice
|
||||
if (error && isProQuotaExceededError(error)) {
|
||||
// Set the flag to prevent tool continuation
|
||||
setModelSwitchedFromQuotaError(true);
|
||||
// Set global quota error flag to prevent Flash model calls
|
||||
config.setQuotaErrorOccurred(true);
|
||||
|
||||
// Show the ProQuotaDialog and wait for user's choice
|
||||
const shouldContinueWithFallback = await new Promise<boolean>(
|
||||
(resolve) => {
|
||||
setIsProQuotaDialogOpen(true);
|
||||
setProQuotaDialogResolver(() => resolve);
|
||||
},
|
||||
);
|
||||
|
||||
// If user chose to continue with fallback, we don't need to stop the current prompt
|
||||
if (shouldContinueWithFallback) {
|
||||
// Switch to fallback model for future use
|
||||
config.setModel(fallbackModel);
|
||||
config.setFallbackMode(true);
|
||||
logFlashFallback(
|
||||
config,
|
||||
new FlashFallbackEvent(
|
||||
config.getContentGeneratorConfig().authType!,
|
||||
),
|
||||
);
|
||||
return true; // Continue with current prompt using fallback model
|
||||
}
|
||||
|
||||
// If user chose to authenticate, stop current prompt
|
||||
return false;
|
||||
}
|
||||
|
||||
// For other quota errors, automatically switch to fallback model
|
||||
// Set the flag to prevent tool continuation
|
||||
setModelSwitchedFromQuotaError(true);
|
||||
// Set global quota error flag to prevent Flash model calls
|
||||
config.setQuotaErrorOccurred(true);
|
||||
}
|
||||
|
||||
// Switch model for future use but return false to stop current retry
|
||||
config.setModel(fallbackModel);
|
||||
config.setFallbackMode(true);
|
||||
logFlashFallback(
|
||||
config,
|
||||
new FlashFallbackEvent(config.getContentGeneratorConfig().authType!),
|
||||
);
|
||||
return false; // Don't continue with current prompt
|
||||
};
|
||||
|
||||
config.setFlashFallbackHandler(flashFallbackHandler);
|
||||
}, [config, historyManager, userTier]);
|
||||
|
||||
const cancelHandlerRef = useRef<() => void>(() => {});
|
||||
|
||||
const {
|
||||
@@ -681,22 +566,6 @@ Logging in with Google... Please restart Gemini CLI to continue.
|
||||
refreshStatic();
|
||||
}, [historyManager, clearConsoleMessagesState, refreshStatic]);
|
||||
|
||||
const handleProQuotaChoice = useCallback(
|
||||
(choice: 'auth' | 'continue') => {
|
||||
setIsProQuotaDialogOpen(false);
|
||||
if (proQuotaDialogResolver) {
|
||||
if (choice === 'auth') {
|
||||
proQuotaDialogResolver(false); // Don't continue with fallback, show auth dialog
|
||||
setAuthState(AuthState.Updating);
|
||||
} else {
|
||||
proQuotaDialogResolver(true); // Continue with fallback model
|
||||
}
|
||||
setProQuotaDialogResolver(null);
|
||||
}
|
||||
},
|
||||
[proQuotaDialogResolver, setAuthState],
|
||||
);
|
||||
|
||||
const { handleInput: vimHandleInput } = useVim(buffer, handleFinalSubmit);
|
||||
|
||||
/**
|
||||
@@ -712,7 +581,7 @@ Logging in with Google... Please restart Gemini CLI to continue.
|
||||
!isProcessing &&
|
||||
(streamingState === StreamingState.Idle ||
|
||||
streamingState === StreamingState.Responding) &&
|
||||
!isProQuotaDialogOpen;
|
||||
!proQuotaRequest;
|
||||
|
||||
// Compute available terminal height based on controls measurement
|
||||
const availableTerminalHeight = useMemo(() => {
|
||||
@@ -1029,7 +898,7 @@ Logging in with Google... Please restart Gemini CLI to continue.
|
||||
isAuthDialogOpen ||
|
||||
isEditorDialogOpen ||
|
||||
showPrivacyNotice ||
|
||||
isProQuotaDialogOpen,
|
||||
!!proQuotaRequest,
|
||||
[
|
||||
showWorkspaceMigrationDialog,
|
||||
shouldShowIdePrompt,
|
||||
@@ -1042,7 +911,7 @@ Logging in with Google... Please restart Gemini CLI to continue.
|
||||
isAuthDialogOpen,
|
||||
isEditorDialogOpen,
|
||||
showPrivacyNotice,
|
||||
isProQuotaDialogOpen,
|
||||
proQuotaRequest,
|
||||
],
|
||||
);
|
||||
|
||||
@@ -1101,11 +970,9 @@ Logging in with Google... Please restart Gemini CLI to continue.
|
||||
showAutoAcceptIndicator,
|
||||
showWorkspaceMigrationDialog,
|
||||
workspaceExtensions,
|
||||
// Use current state values instead of config.getModel()
|
||||
currentModel,
|
||||
userTier,
|
||||
isProQuotaDialogOpen,
|
||||
// New fields
|
||||
proQuotaRequest,
|
||||
contextFileNames,
|
||||
errorCount,
|
||||
availableTerminalHeight,
|
||||
@@ -1174,10 +1041,8 @@ Logging in with Google... Please restart Gemini CLI to continue.
|
||||
showAutoAcceptIndicator,
|
||||
showWorkspaceMigrationDialog,
|
||||
workspaceExtensions,
|
||||
// Quota-related state dependencies
|
||||
userTier,
|
||||
isProQuotaDialogOpen,
|
||||
// New fields dependencies
|
||||
proQuotaRequest,
|
||||
contextFileNames,
|
||||
errorCount,
|
||||
availableTerminalHeight,
|
||||
@@ -1196,7 +1061,6 @@ Logging in with Google... Please restart Gemini CLI to continue.
|
||||
updateInfo,
|
||||
showIdeRestartPrompt,
|
||||
isRestarting,
|
||||
// Quota-related dependencies
|
||||
currentModel,
|
||||
],
|
||||
);
|
||||
|
||||
@@ -22,7 +22,6 @@ import { useUIState } from '../contexts/UIStateContext.js';
|
||||
import { useUIActions } from '../contexts/UIActionsContext.js';
|
||||
import { useConfig } from '../contexts/ConfigContext.js';
|
||||
import { useSettings } from '../contexts/SettingsContext.js';
|
||||
import { DEFAULT_GEMINI_FLASH_MODEL } from '@google/gemini-cli-core';
|
||||
import process from 'node:process';
|
||||
|
||||
// Props for DialogManager
|
||||
@@ -54,11 +53,11 @@ export const DialogManager = () => {
|
||||
/>
|
||||
);
|
||||
}
|
||||
if (uiState.isProQuotaDialogOpen) {
|
||||
if (uiState.proQuotaRequest) {
|
||||
return (
|
||||
<ProQuotaDialog
|
||||
currentModel={uiState.currentModel}
|
||||
fallbackModel={DEFAULT_GEMINI_FLASH_MODEL}
|
||||
failedModel={uiState.proQuotaRequest.failedModel}
|
||||
fallbackModel={uiState.proQuotaRequest.fallbackModel}
|
||||
onChoice={uiActions.handleProQuotaChoice}
|
||||
/>
|
||||
);
|
||||
|
||||
@@ -22,7 +22,7 @@ describe('ProQuotaDialog', () => {
|
||||
it('should render with correct title and options', () => {
|
||||
const { lastFrame } = render(
|
||||
<ProQuotaDialog
|
||||
currentModel="gemini-2.5-pro"
|
||||
failedModel="gemini-2.5-pro"
|
||||
fallbackModel="gemini-2.5-flash"
|
||||
onChoice={() => {}}
|
||||
/>,
|
||||
@@ -53,7 +53,7 @@ describe('ProQuotaDialog', () => {
|
||||
const mockOnChoice = vi.fn();
|
||||
render(
|
||||
<ProQuotaDialog
|
||||
currentModel="gemini-2.5-pro"
|
||||
failedModel="gemini-2.5-pro"
|
||||
fallbackModel="gemini-2.5-flash"
|
||||
onChoice={mockOnChoice}
|
||||
/>,
|
||||
@@ -72,7 +72,7 @@ describe('ProQuotaDialog', () => {
|
||||
const mockOnChoice = vi.fn();
|
||||
render(
|
||||
<ProQuotaDialog
|
||||
currentModel="gemini-2.5-pro"
|
||||
failedModel="gemini-2.5-pro"
|
||||
fallbackModel="gemini-2.5-flash"
|
||||
onChoice={mockOnChoice}
|
||||
/>,
|
||||
|
||||
@@ -10,13 +10,13 @@ import { RadioButtonSelect } from './shared/RadioButtonSelect.js';
|
||||
import { Colors } from '../colors.js';
|
||||
|
||||
interface ProQuotaDialogProps {
|
||||
currentModel: string;
|
||||
failedModel: string;
|
||||
fallbackModel: string;
|
||||
onChoice: (choice: 'auth' | 'continue') => void;
|
||||
}
|
||||
|
||||
export function ProQuotaDialog({
|
||||
currentModel,
|
||||
failedModel,
|
||||
fallbackModel,
|
||||
onChoice,
|
||||
}: ProQuotaDialogProps): React.JSX.Element {
|
||||
@@ -38,7 +38,7 @@ export function ProQuotaDialog({
|
||||
return (
|
||||
<Box borderStyle="round" flexDirection="column" paddingX={1}>
|
||||
<Text bold color={Colors.AccentYellow}>
|
||||
Pro quota limit reached for {currentModel}.
|
||||
Pro quota limit reached for {failedModel}.
|
||||
</Text>
|
||||
<Box marginTop={1}>
|
||||
<RadioButtonSelect
|
||||
|
||||
@@ -21,11 +21,18 @@ import type {
|
||||
ApprovalMode,
|
||||
UserTierId,
|
||||
DetectedIde,
|
||||
FallbackIntent,
|
||||
} from '@google/gemini-cli-core';
|
||||
import type { DOMElement } from 'ink';
|
||||
import type { SessionStatsState } from '../contexts/SessionContext.js';
|
||||
import type { UpdateObject } from '../utils/updateCheck.js';
|
||||
|
||||
export interface ProQuotaDialogRequest {
|
||||
failedModel: string;
|
||||
fallbackModel: string;
|
||||
resolve: (intent: FallbackIntent) => void;
|
||||
}
|
||||
|
||||
export interface UIState {
|
||||
history: HistoryItem[];
|
||||
isThemeDialogOpen: boolean;
|
||||
@@ -78,9 +85,8 @@ export interface UIState {
|
||||
workspaceExtensions: any[]; // Extension[]
|
||||
// Quota-related state
|
||||
userTier: UserTierId | undefined;
|
||||
isProQuotaDialogOpen: boolean;
|
||||
proQuotaRequest: ProQuotaDialogRequest | null;
|
||||
currentModel: string;
|
||||
// New fields for complete state management
|
||||
contextFileNames: string[];
|
||||
errorCount: number;
|
||||
availableTerminalHeight: number | undefined;
|
||||
|
||||
391
packages/cli/src/ui/hooks/useQuotaAndFallback.test.ts
Normal file
391
packages/cli/src/ui/hooks/useQuotaAndFallback.test.ts
Normal file
@@ -0,0 +1,391 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2025 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import {
|
||||
vi,
|
||||
describe,
|
||||
it,
|
||||
expect,
|
||||
beforeEach,
|
||||
afterEach,
|
||||
type Mock,
|
||||
} from 'vitest';
|
||||
import { act, renderHook } from '@testing-library/react';
|
||||
import {
|
||||
type Config,
|
||||
type FallbackModelHandler,
|
||||
UserTierId,
|
||||
AuthType,
|
||||
isGenericQuotaExceededError,
|
||||
isProQuotaExceededError,
|
||||
makeFakeConfig,
|
||||
} 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>;
|
||||
|
||||
describe('useQuotaAndFallback', () => {
|
||||
let mockConfig: Config;
|
||||
let mockHistoryManager: UseHistoryManagerReturn;
|
||||
let mockSetAuthState: Mock;
|
||||
let mockSetModelSwitchedFromQuotaError: Mock;
|
||||
let setFallbackHandlerSpy: SpyInstance;
|
||||
|
||||
const mockedIsGenericQuotaExceededError = isGenericQuotaExceededError as Mock;
|
||||
const mockedIsProQuotaExceededError = isProQuotaExceededError as Mock;
|
||||
|
||||
beforeEach(() => {
|
||||
mockConfig = makeFakeConfig();
|
||||
|
||||
// Spy on the method that requires the private field and mock its return.
|
||||
// This is cleaner than modifying the config class for tests.
|
||||
vi.spyOn(mockConfig, 'getContentGeneratorConfig').mockReturnValue({
|
||||
model: 'gemini-pro',
|
||||
authType: AuthType.LOGIN_WITH_GOOGLE,
|
||||
});
|
||||
|
||||
mockHistoryManager = {
|
||||
addItem: vi.fn(),
|
||||
history: [],
|
||||
updateItem: vi.fn(),
|
||||
clearItems: vi.fn(),
|
||||
loadHistory: vi.fn(),
|
||||
};
|
||||
mockSetAuthState = vi.fn();
|
||||
mockSetModelSwitchedFromQuotaError = vi.fn();
|
||||
|
||||
setFallbackHandlerSpy = vi.spyOn(mockConfig, 'setFallbackModelHandler');
|
||||
vi.spyOn(mockConfig, 'setQuotaErrorOccurred');
|
||||
|
||||
mockedIsGenericQuotaExceededError.mockReturnValue(false);
|
||||
mockedIsProQuotaExceededError.mockReturnValue(false);
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
vi.clearAllMocks();
|
||||
});
|
||||
|
||||
it('should register a fallback handler on initialization', () => {
|
||||
renderHook(() =>
|
||||
useQuotaAndFallback({
|
||||
config: mockConfig,
|
||||
historyManager: mockHistoryManager,
|
||||
userTier: UserTierId.FREE,
|
||||
setAuthState: mockSetAuthState,
|
||||
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
|
||||
}),
|
||||
);
|
||||
|
||||
expect(setFallbackHandlerSpy).toHaveBeenCalledTimes(1);
|
||||
expect(setFallbackHandlerSpy.mock.calls[0][0]).toBeInstanceOf(Function);
|
||||
});
|
||||
|
||||
describe('Fallback Handler Logic', () => {
|
||||
// Helper function to render the hook and extract the registered handler
|
||||
const getRegisteredHandler = (
|
||||
userTier: UserTierId = UserTierId.FREE,
|
||||
): FallbackModelHandler => {
|
||||
renderHook(
|
||||
(props) =>
|
||||
useQuotaAndFallback({
|
||||
config: mockConfig,
|
||||
historyManager: mockHistoryManager,
|
||||
userTier: props.userTier,
|
||||
setAuthState: mockSetAuthState,
|
||||
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
|
||||
}),
|
||||
{ initialProps: { userTier } },
|
||||
);
|
||||
return setFallbackHandlerSpy.mock.calls[0][0] as FallbackModelHandler;
|
||||
};
|
||||
|
||||
it('should return null and take no action if already in fallback mode', async () => {
|
||||
vi.spyOn(mockConfig, 'isInFallbackMode').mockReturnValue(true);
|
||||
const handler = getRegisteredHandler();
|
||||
const result = await handler('gemini-pro', 'gemini-flash', new Error());
|
||||
|
||||
expect(result).toBeNull();
|
||||
expect(mockHistoryManager.addItem).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
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({
|
||||
model: 'gemini-pro',
|
||||
authType: AuthType.USE_GEMINI,
|
||||
});
|
||||
|
||||
const handler = getRegisteredHandler();
|
||||
const result = await handler('gemini-pro', 'gemini-flash', new Error());
|
||||
|
||||
expect(result).toBeNull();
|
||||
expect(mockHistoryManager.addItem).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
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,
|
||||
expectedMessageSnippets: [
|
||||
'Automatically switching from model-A to model-B for faster responses',
|
||||
'upgrade to a Gemini Code Assist Standard or Enterprise plan',
|
||||
],
|
||||
},
|
||||
{
|
||||
errorType: 'other',
|
||||
tier: UserTierId.LEGACY, // Paid tier
|
||||
expectedMessageSnippets: [
|
||||
'Automatically switching from model-A to model-B for faster responses',
|
||||
'switch to using a paid API key from AI Studio',
|
||||
],
|
||||
},
|
||||
];
|
||||
|
||||
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('quota exceeded'),
|
||||
);
|
||||
|
||||
// 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)', () => {
|
||||
beforeEach(() => {
|
||||
mockedIsProQuotaExceededError.mockReturnValue(true);
|
||||
});
|
||||
|
||||
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,
|
||||
}),
|
||||
);
|
||||
|
||||
const handler = setFallbackHandlerSpy.mock
|
||||
.calls[0][0] as FallbackModelHandler;
|
||||
|
||||
// Call the handler but do not await it, to check the intermediate state
|
||||
const promise = handler(
|
||||
'gemini-pro',
|
||||
'gemini-flash',
|
||||
new Error('pro quota'),
|
||||
);
|
||||
|
||||
await act(async () => {});
|
||||
|
||||
// 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('gemini-pro');
|
||||
|
||||
// Simulate the user choosing to continue with the fallback model
|
||||
act(() => {
|
||||
result.current.handleProQuotaChoice('continue');
|
||||
});
|
||||
|
||||
// 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();
|
||||
});
|
||||
|
||||
it('should handle race conditions by stopping subsequent requests', async () => {
|
||||
const { result } = renderHook(() =>
|
||||
useQuotaAndFallback({
|
||||
config: mockConfig,
|
||||
historyManager: mockHistoryManager,
|
||||
userTier: UserTierId.FREE,
|
||||
setAuthState: mockSetAuthState,
|
||||
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
|
||||
}),
|
||||
);
|
||||
|
||||
const handler = setFallbackHandlerSpy.mock
|
||||
.calls[0][0] as FallbackModelHandler;
|
||||
|
||||
const promise1 = handler(
|
||||
'gemini-pro',
|
||||
'gemini-flash',
|
||||
new Error('pro quota 1'),
|
||||
);
|
||||
await act(async () => {});
|
||||
|
||||
const firstRequest = result.current.proQuotaRequest;
|
||||
expect(firstRequest).not.toBeNull();
|
||||
|
||||
const result2 = await handler(
|
||||
'gemini-pro',
|
||||
'gemini-flash',
|
||||
new Error('pro quota 2'),
|
||||
);
|
||||
|
||||
// The lock should have stopped the second request
|
||||
expect(result2).toBe('stop');
|
||||
expect(result.current.proQuotaRequest).toBe(firstRequest);
|
||||
|
||||
act(() => {
|
||||
result.current.handleProQuotaChoice('continue');
|
||||
});
|
||||
|
||||
const intent1 = await promise1;
|
||||
expect(intent1).toBe('retry');
|
||||
expect(result.current.proQuotaRequest).toBeNull();
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe('handleProQuotaChoice', () => {
|
||||
beforeEach(() => {
|
||||
mockedIsProQuotaExceededError.mockReturnValue(true);
|
||||
});
|
||||
|
||||
it('should do nothing if there is no pending pro quota request', () => {
|
||||
const { result } = renderHook(() =>
|
||||
useQuotaAndFallback({
|
||||
config: mockConfig,
|
||||
historyManager: mockHistoryManager,
|
||||
userTier: UserTierId.FREE,
|
||||
setAuthState: mockSetAuthState,
|
||||
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
|
||||
}),
|
||||
);
|
||||
|
||||
act(() => {
|
||||
result.current.handleProQuotaChoice('auth');
|
||||
});
|
||||
|
||||
expect(mockSetAuthState).not.toHaveBeenCalled();
|
||||
expect(mockHistoryManager.addItem).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should resolve intent to "auth" and trigger auth state update', async () => {
|
||||
const { result } = renderHook(() =>
|
||||
useQuotaAndFallback({
|
||||
config: mockConfig,
|
||||
historyManager: mockHistoryManager,
|
||||
userTier: UserTierId.FREE,
|
||||
setAuthState: mockSetAuthState,
|
||||
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
|
||||
}),
|
||||
);
|
||||
|
||||
const handler = setFallbackHandlerSpy.mock
|
||||
.calls[0][0] as FallbackModelHandler;
|
||||
const promise = handler(
|
||||
'gemini-pro',
|
||||
'gemini-flash',
|
||||
new Error('pro quota'),
|
||||
);
|
||||
await act(async () => {}); // Allow state to update
|
||||
|
||||
act(() => {
|
||||
result.current.handleProQuotaChoice('auth');
|
||||
});
|
||||
|
||||
const intent = await promise;
|
||||
expect(intent).toBe('auth');
|
||||
expect(mockSetAuthState).toHaveBeenCalledWith(AuthState.Updating);
|
||||
expect(result.current.proQuotaRequest).toBeNull();
|
||||
});
|
||||
|
||||
it('should resolve intent to "retry" and add info message on continue', async () => {
|
||||
const { result } = renderHook(() =>
|
||||
useQuotaAndFallback({
|
||||
config: mockConfig,
|
||||
historyManager: mockHistoryManager,
|
||||
userTier: UserTierId.FREE,
|
||||
setAuthState: mockSetAuthState,
|
||||
setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError,
|
||||
}),
|
||||
);
|
||||
|
||||
const handler = setFallbackHandlerSpy.mock
|
||||
.calls[0][0] as FallbackModelHandler;
|
||||
// The first `addItem` call is for the initial quota error message
|
||||
const promise = handler(
|
||||
'gemini-pro',
|
||||
'gemini-flash',
|
||||
new Error('pro quota'),
|
||||
);
|
||||
await act(async () => {}); // Allow state to update
|
||||
|
||||
act(() => {
|
||||
result.current.handleProQuotaChoice('continue');
|
||||
});
|
||||
|
||||
const intent = await promise;
|
||||
expect(intent).toBe('retry');
|
||||
expect(result.current.proQuotaRequest).toBeNull();
|
||||
|
||||
// Check for the second "Switched to fallback model" message
|
||||
expect(mockHistoryManager.addItem).toHaveBeenCalledTimes(2);
|
||||
const lastCall = (mockHistoryManager.addItem as Mock).mock.calls[1][0];
|
||||
expect(lastCall.type).toBe(MessageType.INFO);
|
||||
expect(lastCall.text).toContain('Switched to fallback model.');
|
||||
});
|
||||
});
|
||||
});
|
||||
175
packages/cli/src/ui/hooks/useQuotaAndFallback.ts
Normal file
175
packages/cli/src/ui/hooks/useQuotaAndFallback.ts
Normal file
@@ -0,0 +1,175 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2025 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import {
|
||||
AuthType,
|
||||
type Config,
|
||||
type FallbackModelHandler,
|
||||
type FallbackIntent,
|
||||
isGenericQuotaExceededError,
|
||||
isProQuotaExceededError,
|
||||
UserTierId,
|
||||
} from '@google/gemini-cli-core';
|
||||
import { useCallback, useEffect, useRef, useState } from 'react';
|
||||
import { type UseHistoryManagerReturn } from './useHistoryManager.js';
|
||||
import { AuthState, MessageType } from '../types.js';
|
||||
import { type ProQuotaDialogRequest } from '../contexts/UIStateContext.js';
|
||||
|
||||
interface UseQuotaAndFallbackArgs {
|
||||
config: Config;
|
||||
historyManager: UseHistoryManagerReturn;
|
||||
userTier: UserTierId | undefined;
|
||||
setAuthState: (state: AuthState) => void;
|
||||
setModelSwitchedFromQuotaError: (value: boolean) => void;
|
||||
}
|
||||
|
||||
export function useQuotaAndFallback({
|
||||
config,
|
||||
historyManager,
|
||||
userTier,
|
||||
setAuthState,
|
||||
setModelSwitchedFromQuotaError,
|
||||
}: UseQuotaAndFallbackArgs) {
|
||||
const [proQuotaRequest, setProQuotaRequest] =
|
||||
useState<ProQuotaDialogRequest | null>(null);
|
||||
const isDialogPending = useRef(false);
|
||||
|
||||
// Set up Flash fallback handler
|
||||
useEffect(() => {
|
||||
const fallbackHandler: FallbackModelHandler = async (
|
||||
failedModel,
|
||||
fallbackModel,
|
||||
error,
|
||||
): Promise<FallbackIntent | null> => {
|
||||
if (config.isInFallbackMode()) {
|
||||
return null;
|
||||
}
|
||||
|
||||
// Fallbacks are currently only handled for OAuth users.
|
||||
const contentGeneratorConfig = config.getContentGeneratorConfig();
|
||||
if (
|
||||
!contentGeneratorConfig ||
|
||||
contentGeneratorConfig.authType !== AuthType.LOGIN_WITH_GOOGLE
|
||||
) {
|
||||
return null;
|
||||
}
|
||||
|
||||
// Use actual user tier if available; otherwise, default to FREE tier behavior (safe default)
|
||||
const isPaidTier =
|
||||
userTier === UserTierId.LEGACY || userTier === UserTierId.STANDARD;
|
||||
|
||||
let message: string;
|
||||
|
||||
if (error && isProQuotaExceededError(error)) {
|
||||
// Pro Quota specific messages (Interactive)
|
||||
if (isPaidTier) {
|
||||
message = `⚡ You have reached your daily ${failedModel} quota limit.
|
||||
⚡ You can choose to authenticate with a paid API key or continue with the fallback model.
|
||||
⚡ 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 = `⚡ You have reached your daily ${failedModel} quota limit.
|
||||
⚡ 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 {
|
||||
// Consecutive 429s or other errors (Automatic fallback)
|
||||
const actionMessage = `⚡ Automatically switching from ${failedModel} to ${fallbackModel} for faster responses for the remainder of this session.`;
|
||||
|
||||
if (isPaidTier) {
|
||||
message = `${actionMessage}
|
||||
⚡ Possible reasons for this are that you have received multiple consecutive capacity errors or you have reached your daily ${failedModel} quota limit
|
||||
⚡ 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}
|
||||
⚡ Possible reasons for this are that you have received multiple consecutive capacity errors or you have reached your daily ${failedModel} quota limit
|
||||
⚡ 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`;
|
||||
}
|
||||
}
|
||||
|
||||
// Add message to UI history
|
||||
historyManager.addItem(
|
||||
{
|
||||
type: MessageType.INFO,
|
||||
text: message,
|
||||
},
|
||||
Date.now(),
|
||||
);
|
||||
|
||||
setModelSwitchedFromQuotaError(true);
|
||||
config.setQuotaErrorOccurred(true);
|
||||
|
||||
// Interactive Fallback for Pro quota
|
||||
if (error && isProQuotaExceededError(error)) {
|
||||
if (isDialogPending.current) {
|
||||
return 'stop'; // A dialog is already active, so just stop this request.
|
||||
}
|
||||
isDialogPending.current = true;
|
||||
|
||||
const intent: FallbackIntent = await new Promise<FallbackIntent>(
|
||||
(resolve) => {
|
||||
setProQuotaRequest({
|
||||
failedModel,
|
||||
fallbackModel,
|
||||
resolve,
|
||||
});
|
||||
},
|
||||
);
|
||||
|
||||
return intent;
|
||||
}
|
||||
|
||||
return 'stop';
|
||||
};
|
||||
|
||||
config.setFallbackModelHandler(fallbackHandler);
|
||||
}, [config, historyManager, userTier, setModelSwitchedFromQuotaError]);
|
||||
|
||||
const handleProQuotaChoice = useCallback(
|
||||
(choice: 'auth' | 'continue') => {
|
||||
if (!proQuotaRequest) return;
|
||||
|
||||
const intent: FallbackIntent = choice === 'auth' ? 'auth' : 'retry';
|
||||
proQuotaRequest.resolve(intent);
|
||||
setProQuotaRequest(null);
|
||||
isDialogPending.current = false; // Reset the flag here
|
||||
|
||||
if (choice === 'auth') {
|
||||
setAuthState(AuthState.Updating);
|
||||
} else {
|
||||
historyManager.addItem(
|
||||
{
|
||||
type: MessageType.INFO,
|
||||
text: 'Switched to fallback model. Tip: Press Ctrl+P (or Up Arrow) to recall your previous prompt and submit it again if you wish.',
|
||||
},
|
||||
Date.now(),
|
||||
);
|
||||
}
|
||||
},
|
||||
[proQuotaRequest, setAuthState, historyManager],
|
||||
);
|
||||
|
||||
return {
|
||||
proQuotaRequest,
|
||||
handleProQuotaChoice,
|
||||
};
|
||||
}
|
||||
Reference in New Issue
Block a user