feat: add Pro Quota Dialog (#7094)

JayadityaGit · gemini-code-assist[bot] · jacob314 · acoliver · commit b11fab6c003e · 2025-09-17T15:21:17.000-03:00
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: Jacob Richman <jacob314@gmail.com> (cherry picked from commit a63e678)
diff --git a/packages/cli/src/ui/App.tsx b/packages/cli/src/ui/App.tsx
@@ -82,6 +82,7 @@ import {
   isGenericQuotaExceededError,
   UserTierId,
   uiTelemetryService,
+  DEFAULT_GEMINI_FLASH_MODEL,
 } from '@vybestack/llxprt-code-core';
 import {
   IdeIntegrationNudge,
@@ -122,6 +123,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 { getProviderManager } from '../providers/providerManagerInstance.js';
@@ -377,13 +379,22 @@ const App = (props: AppInternalProps) => {
   const [showEscapePrompt, setShowEscapePrompt] = useState(false);
   const [isProcessing, setIsProcessing] = useState<boolean>(false);
   const [providerModels, setProviderModels] = useState<IModel[]>([]);
+
   const {
     showWorkspaceMigrationDialog,
     workspaceExtensions,
     onWorkspaceMigrationDialogOpen,
     onWorkspaceMigrationDialogClose,
   } = useWorkspaceMigration(settings);
 
+  const [isProQuotaDialogOpen, setIsProQuotaDialogOpen] = useState(false);
+  const [proQuotaDialogResolver, setProQuotaDialogResolver] = useState<
+    ((value: boolean) => void) | null
+  >(null);
+  const [proQuotaDialogResolver, setProQuotaDialogResolver] = useState<
+    ((value: boolean) => void) | null
+  >(null);
+
   useEffect(() => {
     const unsubscribe = ideContext.subscribeToIdeContext(setIdeContextState);
     // Set the initial value
@@ -752,6 +763,12 @@ const App = (props: AppInternalProps) => {
       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;
 
       const contentGenConfig = config.getContentGeneratorConfig();
@@ -764,14 +781,15 @@ const App = (props: AppInternalProps) => {
         // Check if this is a Pro quota exceeded error
         if (error && isProQuotaExceededError(error)) {
           if (isPaidTier) {
-            message = `You have reached your daily ${currentModel} quota limit.
-To continue using ${currentModel}, you can use /auth to switch to using a paid API key from AI Studio at https://aistudio.google.com/apikey
-Or you can switch to a different model using the /model command`;
+            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.
-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 or switch to a different model using /model`;
+            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) {
@@ -812,6 +830,40 @@ You can switch authentication methods by typing /auth or switch to a different m
           );
         }
 
+        // 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
@@ -1223,7 +1275,11 @@ You can switch authentication methods by typing /auth or switch to a different m
   }, [history, logger]);
 
   const isInputActive =
-    streamingState === StreamingState.Idle && !initError && !isProcessing;
+    (streamingState === StreamingState.Idle ||
+      streamingState === StreamingState.Responding) &&
+    !initError &&
+    !isProcessing &&
+    !isProQuotaDialogOpen;
 
   const handleClearScreen = useCallback(() => {
     clearItems();
@@ -1469,6 +1525,31 @@ You can switch authentication methods by typing /auth or switch to a different m
               ide={currentIDE}
               onComplete={handleIdePromptComplete}
             />
+          ) : isProQuotaDialogOpen ? (
+            <ProQuotaDialog
+              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') {
+                  openAuthDialog();
+                } else {
+                  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 ? (
             <FolderTrustDialog
               onSelect={handleFolderTrustSelect}
diff --git a/packages/cli/src/ui/components/ProQuotaDialog.test.tsx 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(
+      <ProQuotaDialog
+        currentModel="gemini-2.5-pro"
+        fallbackModel="gemini-2.5-flash"
+        onChoice={() => {}}
+      />,
+    );
+
+    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(
+      <ProQuotaDialog
+        currentModel="gemini-2.5-pro"
+        fallbackModel="gemini-2.5-flash"
+        onChoice={mockOnChoice}
+      />,
+    );
+
+    // 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(
+      <ProQuotaDialog
+        currentModel="gemini-2.5-pro"
+        fallbackModel="gemini-2.5-flash"
+        onChoice={mockOnChoice}
+      />,
+    );
+
+    // 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');
+  });
+});
diff --git a/packages/cli/src/ui/components/ProQuotaDialog.tsx b/packages/cli/src/ui/components/ProQuotaDialog.tsx
@@ -0,0 +1,52 @@
+/**
+ * @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';
+import { Colors } from '../colors.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 (
+    <Box borderStyle="round" flexDirection="column" paddingX={1}>
+      <Text bold color={Colors.AccentYellow}>
+        Pro quota limit reached for {currentModel}.
+      </Text>
+      <Box marginTop={1}>
+        <RadioButtonSelect
+          items={items}
+          initialIndex={1}
+          onSelect={handleSelect}
+        />
+      </Box>
+    </Box>
+  );
+}
