From b5ee70fdd5efadb39d4deb62c11132b6283c6d08 Mon Sep 17 00:00:00 2001 From: JayadityaGit Date: Tue, 26 Aug 2025 12:52:22 +0530 Subject: [PATCH 1/6] feat: add Pro Quota Dialog --- packages/cli/src/ui/App.tsx | 95 ++++++++++++++++++- .../cli/src/ui/components/InputPrompt.tsx | 4 +- .../cli/src/ui/components/ProQuotaDialog.tsx | 57 +++++++++++ .../ui/components/messages/InfoMessage.tsx | 2 +- 4 files changed, 150 insertions(+), 8 deletions(-) create mode 100644 packages/cli/src/ui/components/ProQuotaDialog.tsx diff --git a/packages/cli/src/ui/App.tsx b/packages/cli/src/ui/App.tsx index 27883cebf01..04e4b8a19cb 100644 --- a/packages/cli/src/ui/App.tsx +++ b/packages/cli/src/ui/App.tsx @@ -70,6 +70,7 @@ import { isProQuotaExceededError, isGenericQuotaExceededError, UserTierId, + DEFAULT_GEMINI_FLASH_MODEL, } from '@google/gemini-cli-core'; import type { IdeIntegrationNudgeResult } from './IdeIntegrationNudge.js'; import { IdeIntegrationNudge } from './IdeIntegrationNudge.js'; @@ -100,6 +101,7 @@ import { ShowMoreLines } from './components/ShowMoreLines.js'; import { PrivacyNotice } from './privacy/PrivacyNotice.js'; import { useSettingsCommand } from './hooks/useSettingsCommand.js'; import { SettingsDialog } from './components/SettingsDialog.js'; +import { ProQuotaDialog } from './components/ProQuotaDialog.js'; import { setUpdateHandler } from '../utils/handleAutoUpdate.js'; import { appEvents, AppEvent } from '../utils/events.js'; import { isNarrowWidth } from './utils/isNarrowWidth.js'; @@ -225,6 +227,7 @@ const App = ({ config, settings, startupWarnings = [], version }: AppProps) => { >(); const [showEscapePrompt, setShowEscapePrompt] = useState(false); const [isProcessing, setIsProcessing] = useState(false); + const { showWorkspaceMigrationDialog, workspaceExtensions, @@ -232,6 +235,12 @@ const App = ({ config, settings, startupWarnings = [], version }: AppProps) => { onWorkspaceMigrationDialogClose, } = useWorkspaceMigration(settings); + const [isProQuotaDialogOpen, setIsProQuotaDialogOpen] = useState(false); + const [proQuotaDialogResolver, setProQuotaDialogResolver] = useState< + ((value: boolean) => void) | null + >(null); + + useEffect(() => { const unsubscribe = ideContext.subscribeToIdeContext(setIdeContextState); // Set the initial value @@ -408,6 +417,12 @@ const App = ({ config, settings, startupWarnings = [], version }: AppProps) => { fallbackModel: string, error?: unknown, ): Promise => { + // 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 ( @@ -422,11 +437,11 @@ const App = ({ config, settings, startupWarnings = [], version }: AppProps) => { if (error && isProQuotaExceededError(error)) { if (isPaidTier) { message = `⚡ You have reached your daily ${currentModel} quota limit. -⚡ Automatically switching from ${currentModel} to ${fallbackModel} for the remainder of this session. +⚡ 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. -⚡ Automatically switching from ${currentModel} to ${fallbackModel} for the remainder of this session. +⚡ 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`; @@ -468,6 +483,40 @@ const App = ({ config, settings, startupWarnings = [], version }: AppProps) => { 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( + (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 @@ -821,10 +870,10 @@ const App = ({ config, settings, startupWarnings = [], version }: AppProps) => { }, [history, logger]); const isInputActive = - (streamingState === StreamingState.Idle || - streamingState === StreamingState.Responding) && + streamingState === StreamingState.Idle && !initError && - !isProcessing; + !isProcessing && + !isProQuotaDialogOpen; const handleClearScreen = useCallback(() => { clearItems(); @@ -1037,6 +1086,42 @@ const App = ({ config, settings, startupWarnings = [], version }: AppProps) => { ide={currentIDE} onComplete={handleIdePromptComplete} /> + ) : isProQuotaDialogOpen ? ( + { + if (choice === 'auth') { + setIsProQuotaDialogOpen(false); + if (proQuotaDialogResolver) { + proQuotaDialogResolver(false); + setProQuotaDialogResolver(null); + } + openAuthDialog(); + } else { + setIsProQuotaDialogOpen(false); + // For the 'continue' option, we want to: + // 1. Resolve the promise to allow the current request to finish + // 2. Switch to fallback mode + // 3. Show a tip about using Ctrl+P to get the previous prompt + if (proQuotaDialogResolver) { + // Resolve with true to indicate we should continue with fallback + proQuotaDialogResolver(true); + setProQuotaDialogResolver(null); + } + // Switch to fallback model immediately + config.setModel(DEFAULT_GEMINI_FLASH_MODEL); + config.setFallbackMode(true); + addItem( + { + type: MessageType.INFO, + text: 'Switched to fallback model. Tip: Press Ctrl+P to recall your previous prompt and submit it again if you wish.', + }, + Date.now(), + ); + } + }} + /> ) : isFolderTrustDialogOpen ? ( = ({ if (shellModeActive) { shellHistory.addCommandToHistory(submittedValue); } - // Clear the buffer *before* calling onSubmit to prevent potential re-submission + onSubmit(submittedValue); + // Clear the buffer *after* calling onSubmit to prevent potential re-submission // if onSubmit triggers a re-render while the buffer still holds the old value. buffer.setText(''); - onSubmit(submittedValue); resetCompletionState(); resetReverseSearchCompletionState(); }, diff --git a/packages/cli/src/ui/components/ProQuotaDialog.tsx b/packages/cli/src/ui/components/ProQuotaDialog.tsx new file mode 100644 index 00000000000..7e8eda935fc --- /dev/null +++ b/packages/cli/src/ui/components/ProQuotaDialog.tsx @@ -0,0 +1,57 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import type React from 'react'; +import { Box, Text } from 'ink'; +import { RadioButtonSelect } from './shared/RadioButtonSelect.js'; + +interface ProQuotaDialogProps { + currentModel: string; + fallbackModel: string; + onChoice: (choice: 'auth' | 'continue') => void; +} + +export function ProQuotaDialog({ + currentModel, + fallbackModel, + onChoice, +}: ProQuotaDialogProps): React.JSX.Element { + const items = [ + { + label: 'Change auth (executes the /auth command)', + value: 'auth' as const, + }, + { + label: `Continue with ${fallbackModel}`, + value: 'continue' as const, + }, + ]; + + const handleSelect = (choice: 'auth' | 'continue') => { + onChoice(choice); + }; + + return ( + + + Pro quota limit reached for {currentModel}. + + + + + + ); +} diff --git a/packages/cli/src/ui/components/messages/InfoMessage.tsx b/packages/cli/src/ui/components/messages/InfoMessage.tsx index 3d7866bec08..c0babd1d6fd 100644 --- a/packages/cli/src/ui/components/messages/InfoMessage.tsx +++ b/packages/cli/src/ui/components/messages/InfoMessage.tsx @@ -14,7 +14,7 @@ interface InfoMessageProps { } export const InfoMessage: React.FC = ({ text }) => { - const prefix = 'ℹ '; + const prefix = ''; const prefixWidth = prefix.length; return ( From e5ce796c3434755184b9809eda19dbc304a08131 Mon Sep 17 00:00:00 2001 From: JAYADITYA <96861162+JayadityaGit@users.noreply.github.com> Date: Tue, 26 Aug 2025 15:29:15 +0530 Subject: [PATCH 2/6] Update packages/cli/src/ui/App.tsx Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- packages/cli/src/ui/App.tsx | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/packages/cli/src/ui/App.tsx b/packages/cli/src/ui/App.tsx index 04e4b8a19cb..dce39c8ee0c 100644 --- a/packages/cli/src/ui/App.tsx +++ b/packages/cli/src/ui/App.tsx @@ -870,7 +870,8 @@ const App = ({ config, settings, startupWarnings = [], version }: AppProps) => { }, [history, logger]); const isInputActive = - streamingState === StreamingState.Idle && + (streamingState === StreamingState.Idle || + streamingState === StreamingState.Responding) && !initError && !isProcessing && !isProQuotaDialogOpen; From c713d92dc54dca00d13ba89831bd70f54d9d16de Mon Sep 17 00:00:00 2001 From: JayadityaGit Date: Tue, 26 Aug 2025 15:37:17 +0530 Subject: [PATCH 3/6] refactor(ui): simplify onChoice handler in ProQuotaDialog --- packages/cli/src/ui/App.tsx | 25 +++++++------------------ 1 file changed, 7 insertions(+), 18 deletions(-) diff --git a/packages/cli/src/ui/App.tsx b/packages/cli/src/ui/App.tsx index dce39c8ee0c..961d788f932 100644 --- a/packages/cli/src/ui/App.tsx +++ b/packages/cli/src/ui/App.tsx @@ -1092,27 +1092,16 @@ const App = ({ config, settings, startupWarnings = [], version }: AppProps) => { currentModel={config.getModel()} fallbackModel={DEFAULT_GEMINI_FLASH_MODEL} onChoice={(choice) => { + setIsProQuotaDialogOpen(false); + if (!proQuotaDialogResolver) return; + + const resolveValue = choice !== 'auth'; + proQuotaDialogResolver(resolveValue); + setProQuotaDialogResolver(null); + if (choice === 'auth') { - setIsProQuotaDialogOpen(false); - if (proQuotaDialogResolver) { - proQuotaDialogResolver(false); - setProQuotaDialogResolver(null); - } openAuthDialog(); } else { - setIsProQuotaDialogOpen(false); - // For the 'continue' option, we want to: - // 1. Resolve the promise to allow the current request to finish - // 2. Switch to fallback mode - // 3. Show a tip about using Ctrl+P to get the previous prompt - if (proQuotaDialogResolver) { - // Resolve with true to indicate we should continue with fallback - proQuotaDialogResolver(true); - setProQuotaDialogResolver(null); - } - // Switch to fallback model immediately - config.setModel(DEFAULT_GEMINI_FLASH_MODEL); - config.setFallbackMode(true); addItem( { type: MessageType.INFO, From 7c7cb5adc443b51eacd9bb6e275dad414f53a841 Mon Sep 17 00:00:00 2001 From: JayadityaGit Date: Wed, 27 Aug 2025 11:34:26 +0530 Subject: [PATCH 4/6] refinements --- packages/cli/src/ui/components/InputPrompt.tsx | 6 +++--- packages/cli/src/ui/components/ProQuotaDialog.tsx | 7 +++---- packages/cli/src/ui/components/messages/InfoMessage.tsx | 4 ++-- 3 files changed, 8 insertions(+), 9 deletions(-) diff --git a/packages/cli/src/ui/components/InputPrompt.tsx b/packages/cli/src/ui/components/InputPrompt.tsx index 980f3038a68..11760ac97d3 100644 --- a/packages/cli/src/ui/components/InputPrompt.tsx +++ b/packages/cli/src/ui/components/InputPrompt.tsx @@ -138,10 +138,10 @@ export const InputPrompt: React.FC = ({ if (shellModeActive) { shellHistory.addCommandToHistory(submittedValue); } - onSubmit(submittedValue); - // Clear the buffer *after* calling onSubmit to prevent potential re-submission + // Clear the buffer *before* calling onSubmit to prevent potential re-submission // if onSubmit triggers a re-render while the buffer still holds the old value. buffer.setText(''); + onSubmit(submittedValue); resetCompletionState(); resetReverseSearchCompletionState(); }, @@ -829,4 +829,4 @@ export const InputPrompt: React.FC = ({ )} ); -}; +}; \ No newline at end of file diff --git a/packages/cli/src/ui/components/ProQuotaDialog.tsx b/packages/cli/src/ui/components/ProQuotaDialog.tsx index 7e8eda935fc..df62dcfac14 100644 --- a/packages/cli/src/ui/components/ProQuotaDialog.tsx +++ b/packages/cli/src/ui/components/ProQuotaDialog.tsx @@ -7,6 +7,7 @@ import type React from 'react'; import { Box, Text } from 'ink'; import { RadioButtonSelect } from './shared/RadioButtonSelect.js'; +import { Colors } from '../colors.js'; interface ProQuotaDialogProps { currentModel: string; @@ -38,11 +39,9 @@ export function ProQuotaDialog({ - + Pro quota limit reached for {currentModel}. diff --git a/packages/cli/src/ui/components/messages/InfoMessage.tsx b/packages/cli/src/ui/components/messages/InfoMessage.tsx index c0babd1d6fd..10965f19428 100644 --- a/packages/cli/src/ui/components/messages/InfoMessage.tsx +++ b/packages/cli/src/ui/components/messages/InfoMessage.tsx @@ -14,7 +14,7 @@ interface InfoMessageProps { } export const InfoMessage: React.FC = ({ text }) => { - const prefix = ''; + const prefix = 'ℹ '; const prefixWidth = prefix.length; return ( @@ -29,4 +29,4 @@ export const InfoMessage: React.FC = ({ text }) => { ); -}; +}; \ No newline at end of file From ebcea0e1ea9e06dfa659b83a75dbab82297bad6b Mon Sep 17 00:00:00 2001 From: JayadityaGit Date: Wed, 27 Aug 2025 11:41:16 +0530 Subject: [PATCH 5/6] refinements --- packages/cli/src/ui/components/InputPrompt.tsx | 2 +- packages/cli/src/ui/components/ProQuotaDialog.tsx | 6 +----- packages/cli/src/ui/components/messages/InfoMessage.tsx | 2 +- 3 files changed, 3 insertions(+), 7 deletions(-) diff --git a/packages/cli/src/ui/components/InputPrompt.tsx b/packages/cli/src/ui/components/InputPrompt.tsx index 11760ac97d3..09897885a29 100644 --- a/packages/cli/src/ui/components/InputPrompt.tsx +++ b/packages/cli/src/ui/components/InputPrompt.tsx @@ -829,4 +829,4 @@ export const InputPrompt: React.FC = ({ )} ); -}; \ No newline at end of file +}; diff --git a/packages/cli/src/ui/components/ProQuotaDialog.tsx b/packages/cli/src/ui/components/ProQuotaDialog.tsx index df62dcfac14..d94d0698571 100644 --- a/packages/cli/src/ui/components/ProQuotaDialog.tsx +++ b/packages/cli/src/ui/components/ProQuotaDialog.tsx @@ -36,11 +36,7 @@ export function ProQuotaDialog({ }; return ( - + Pro quota limit reached for {currentModel}. diff --git a/packages/cli/src/ui/components/messages/InfoMessage.tsx b/packages/cli/src/ui/components/messages/InfoMessage.tsx index 10965f19428..3d7866bec08 100644 --- a/packages/cli/src/ui/components/messages/InfoMessage.tsx +++ b/packages/cli/src/ui/components/messages/InfoMessage.tsx @@ -29,4 +29,4 @@ export const InfoMessage: React.FC = ({ text }) => { ); -}; \ No newline at end of file +}; From 61450db300bbeececacc13913854d8a0d785cbfa Mon Sep 17 00:00:00 2001 From: JayadityaGit Date: Thu, 28 Aug 2025 12:48:17 +0530 Subject: [PATCH 6/6] feat(cli): Add tests for ProQuotaDialog Adds tests for the `ProQuotaDialog` component. The tests cover the following scenarios: - The dialog renders with the correct title and options. - The `onChoice` callback is called with the correct value when the user selects an option. --- packages/cli/src/ui/App.tsx | 3 +- .../src/ui/components/ProQuotaDialog.test.tsx | 89 +++++++++++++++++++ 2 files changed, 90 insertions(+), 2 deletions(-) create mode 100644 packages/cli/src/ui/components/ProQuotaDialog.test.tsx diff --git a/packages/cli/src/ui/App.tsx b/packages/cli/src/ui/App.tsx index 6ff37f2aaf4..ff10918a77d 100644 --- a/packages/cli/src/ui/App.tsx +++ b/packages/cli/src/ui/App.tsx @@ -229,7 +229,7 @@ const App = ({ config, settings, startupWarnings = [], version }: AppProps) => { >(); const [showEscapePrompt, setShowEscapePrompt] = useState(false); const [isProcessing, setIsProcessing] = useState(false); - + const { showWorkspaceMigrationDialog, workspaceExtensions, @@ -242,7 +242,6 @@ const App = ({ config, settings, startupWarnings = [], version }: AppProps) => { ((value: boolean) => void) | null >(null); - useEffect(() => { const unsubscribe = ideContext.subscribeToIdeContext(setIdeContextState); // Set the initial value diff --git a/packages/cli/src/ui/components/ProQuotaDialog.test.tsx b/packages/cli/src/ui/components/ProQuotaDialog.test.tsx new file mode 100644 index 00000000000..31bb4f03f67 --- /dev/null +++ b/packages/cli/src/ui/components/ProQuotaDialog.test.tsx @@ -0,0 +1,89 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { render } from 'ink-testing-library'; +import { describe, it, expect, vi, beforeEach, type Mock } from 'vitest'; +import { ProQuotaDialog } from './ProQuotaDialog.js'; +import { RadioButtonSelect } from './shared/RadioButtonSelect.js'; + +// Mock the child component to make it easier to test the parent +vi.mock('./shared/RadioButtonSelect.js', () => ({ + RadioButtonSelect: vi.fn(), +})); + +describe('ProQuotaDialog', () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it('should render with correct title and options', () => { + const { lastFrame } = render( + {}} + />, + ); + + const output = lastFrame(); + expect(output).toContain('Pro quota limit reached for gemini-2.5-pro.'); + + // Check that RadioButtonSelect was called with the correct items + expect(RadioButtonSelect).toHaveBeenCalledWith( + expect.objectContaining({ + items: [ + { + label: 'Change auth (executes the /auth command)', + value: 'auth', + }, + { + label: `Continue with gemini-2.5-flash`, + value: 'continue', + }, + ], + }), + undefined, + ); + }); + + it('should call onChoice with "auth" when "Change auth" is selected', () => { + const mockOnChoice = vi.fn(); + render( + , + ); + + // Get the onSelect function passed to RadioButtonSelect + const onSelect = (RadioButtonSelect as Mock).mock.calls[0][0].onSelect; + + // Simulate the selection + onSelect('auth'); + + expect(mockOnChoice).toHaveBeenCalledWith('auth'); + }); + + it('should call onChoice with "continue" when "Continue with flash" is selected', () => { + const mockOnChoice = vi.fn(); + render( + , + ); + + // Get the onSelect function passed to RadioButtonSelect + const onSelect = (RadioButtonSelect as Mock).mock.calls[0][0].onSelect; + + // Simulate the selection + onSelect('continue'); + + expect(mockOnChoice).toHaveBeenCalledWith('continue'); + }); +});