fix(cli): support quota error fallbacks for all authentication types (#20475)

Co-authored-by: Gal Zahavi <38544478+galz10@users.noreply.github.com>

packages/cli/src/ui/components/DialogManager.tsx

@@ -135,6 +135,7 @@ export const DialogManager = ({
         isModelNotFoundError={
           !!uiState.quota.proQuotaRequest.isModelNotFoundError
         }
+        authType={uiState.quota.proQuotaRequest.authType}
         onChoice={uiActions.handleProQuotaChoice}
       />
     );

packages/cli/src/ui/components/ProQuotaDialog.test.tsx

@@ -13,6 +13,7 @@ import { RadioButtonSelect } from './shared/RadioButtonSelect.js';
 import {
   PREVIEW_GEMINI_MODEL,
   DEFAULT_GEMINI_FLASH_MODEL,
+  AuthType,
 } from '@google/gemini-cli-core';
 
 // Mock the child component to make it easier to test the parent
@@ -62,7 +63,7 @@ describe('ProQuotaDialog', () => {
 
   describe('for non-flash model failures', () => {
     describe('when it is a terminal quota error', () => {
-      it('should render switch, upgrade, and stop options for paid tiers', () => {
+      it('should render switch, upgrade, and stop options for LOGIN_WITH_GOOGLE', () => {
         const { unmount } = render(
           <ProQuotaDialog
             failedModel="gemini-2.5-pro"
@@ -70,6 +71,7 @@ describe('ProQuotaDialog', () => {
             message="paid tier quota error"
             isTerminalQuotaError={true}
             isModelNotFoundError={false}
+            authType={AuthType.LOGIN_WITH_GOOGLE}
             onChoice={mockOnChoice}
           />,
         );
@@ -99,6 +101,39 @@ describe('ProQuotaDialog', () => {
         unmount();
       });
 
+      it('should NOT render upgrade option for USE_GEMINI', () => {
+        const { unmount } = render(
+          <ProQuotaDialog
+            failedModel="gemini-2.5-pro"
+            fallbackModel="gemini-2.5-flash"
+            message="paid tier quota error"
+            isTerminalQuotaError={true}
+            isModelNotFoundError={false}
+            authType={AuthType.USE_GEMINI}
+            onChoice={mockOnChoice}
+          />,
+        );
+
+        expect(RadioButtonSelect).toHaveBeenCalledWith(
+          expect.objectContaining({
+            items: [
+              {
+                label: 'Switch to gemini-2.5-flash',
+                value: 'retry_always',
+                key: 'retry_always',
+              },
+              {
+                label: 'Stop',
+                value: 'retry_later',
+                key: 'retry_later',
+              },
+            ],
+          }),
+          undefined,
+        );
+        unmount();
+      });
+
       it('should render "Keep trying" and "Stop" options when failed model and fallback model are the same', () => {
         const { unmount } = render(
           <ProQuotaDialog
@@ -130,7 +165,7 @@ describe('ProQuotaDialog', () => {
         unmount();
       });
 
-      it('should render switch, upgrade, and stop options for free tier', () => {
+      it('should render switch, upgrade, and stop options for LOGIN_WITH_GOOGLE (free tier)', () => {
         const { unmount } = render(
           <ProQuotaDialog
             failedModel="gemini-2.5-pro"
@@ -138,6 +173,7 @@ describe('ProQuotaDialog', () => {
             message="free tier quota error"
             isTerminalQuotaError={true}
             isModelNotFoundError={false}
+            authType={AuthType.LOGIN_WITH_GOOGLE}
             onChoice={mockOnChoice}
           />,
         );
@@ -204,7 +240,7 @@ describe('ProQuotaDialog', () => {
     });
 
     describe('when it is a model not found error', () => {
-      it('should render switch and stop options regardless of tier', () => {
+      it('should render switch, upgrade, and stop options for LOGIN_WITH_GOOGLE', () => {
         const { unmount } = render(
           <ProQuotaDialog
             failedModel="gemini-3-pro-preview"
@@ -212,6 +248,7 @@ describe('ProQuotaDialog', () => {
             message="You don't have access to gemini-3-pro-preview yet."
             isTerminalQuotaError={false}
             isModelNotFoundError={true}
+            authType={AuthType.LOGIN_WITH_GOOGLE}
             onChoice={mockOnChoice}
           />,
         );
@@ -241,7 +278,7 @@ describe('ProQuotaDialog', () => {
         unmount();
       });
 
-      it('should render switch and stop options for paid tier as well', () => {
+      it('should NOT render upgrade option for USE_GEMINI', () => {
         const { unmount } = render(
           <ProQuotaDialog
             failedModel="gemini-3-pro-preview"
@@ -249,6 +286,7 @@ describe('ProQuotaDialog', () => {
             message="You don't have access to gemini-3-pro-preview yet."
             isTerminalQuotaError={false}
             isModelNotFoundError={true}
+            authType={AuthType.USE_GEMINI}
             onChoice={mockOnChoice}
           />,
         );
@@ -261,11 +299,6 @@ describe('ProQuotaDialog', () => {
                 value: 'retry_always',
                 key: 'retry_always',
               },
-              {
-                label: 'Upgrade for higher limits',
-                value: 'upgrade',
-                key: 'upgrade',
-              },
               {
                 label: 'Stop',
                 value: 'retry_later',

packages/cli/src/ui/components/ProQuotaDialog.tsx

@@ -8,6 +8,7 @@ import type React from 'react';
 import { Box, Text } from 'ink';
 import { RadioButtonSelect } from './shared/RadioButtonSelect.js';
 import { theme } from '../semantic-colors.js';
+import { AuthType } from '@google/gemini-cli-core';
 
 interface ProQuotaDialogProps {
   failedModel: string;
@@ -15,6 +16,7 @@ interface ProQuotaDialogProps {
   message: string;
   isTerminalQuotaError: boolean;
   isModelNotFoundError?: boolean;
+  authType?: AuthType;
   onChoice: (
     choice: 'retry_later' | 'retry_once' | 'retry_always' | 'upgrade',
   ) => void;
@@ -26,6 +28,7 @@ export function ProQuotaDialog({
   message,
   isTerminalQuotaError,
   isModelNotFoundError,
+  authType,
   onChoice,
 }: ProQuotaDialogProps): React.JSX.Element {
   let items;
@@ -51,11 +54,15 @@ export function ProQuotaDialog({
         value: 'retry_always' as const,
         key: 'retry_always',
       },
+      ...(authType === AuthType.LOGIN_WITH_GOOGLE
+        ? [
             {
               label: 'Upgrade for higher limits',
               value: 'upgrade' as const,
               key: 'upgrade',
             },
+          ]
+        : []),
       {
         label: `Stop`,
         value: 'retry_later' as const,

packages/cli/src/ui/contexts/UIStateContext.tsx

@@ -24,6 +24,7 @@ import type {
   ApprovalMode,
   UserTierId,
   IdeInfo,
+  AuthType,
   FallbackIntent,
   ValidationIntent,
   AgentDefinition,
@@ -42,6 +43,7 @@ export interface ProQuotaDialogRequest {
   message: string;
   isTerminalQuotaError: boolean;
   isModelNotFoundError?: boolean;
+  authType?: AuthType;
   resolve: (intent: FallbackIntent) => void;
 }
 

packages/cli/src/ui/hooks/useQuotaAndFallback.test.ts

@@ -96,9 +96,13 @@ describe('useQuotaAndFallback', () => {
   });
 
   describe('Fallback Handler Logic', () => {
-    // Helper function to render the hook and extract the registered handler
+    it('should show fallback dialog but omit switch to API key message if authType is not LOGIN_WITH_GOOGLE', async () => {
-    const getRegisteredHandler = (): FallbackModelHandler => {
+      // Override the default mock from beforeEach for this specific test
-      renderHook(() =>
+      vi.spyOn(mockConfig, 'getContentGeneratorConfig').mockReturnValue({
+        authType: AuthType.USE_GEMINI,
+      });
+
+      const { result } = renderHook(() =>
         useQuotaAndFallback({
           config: mockConfig,
           historyManager: mockHistoryManager,
@@ -107,20 +111,24 @@ describe('useQuotaAndFallback', () => {
           onShowAuthSelection: mockOnShowAuthSelection,
         }),
       );
-      return setFallbackHandlerSpy.mock.calls[0][0] as FallbackModelHandler;
-    };
 
-    it('should return null and take no action if authType is not LOGIN_WITH_GOOGLE', async () => {
+      const handler = setFallbackHandlerSpy.mock
-      // Override the default mock from beforeEach for this specific test
+        .calls[0][0] as FallbackModelHandler;
-      vi.spyOn(mockConfig, 'getContentGeneratorConfig').mockReturnValue({
+
-        authType: AuthType.USE_GEMINI,
+      const error = new TerminalQuotaError(
+        'pro quota',
+        mockGoogleApiError,
+        1000 * 60 * 5,
+      );
+
+      act(() => {
+        void handler('gemini-pro', 'gemini-flash', error);
       });
 
-      const handler = getRegisteredHandler();
+      expect(result.current.proQuotaRequest).not.toBeNull();
-      const result = await handler('gemini-pro', 'gemini-flash', new Error());
+      expect(result.current.proQuotaRequest?.message).not.toContain(
-
+        '/auth to switch to API key.',
-      expect(result).toBeNull();
+      );
-      expect(mockHistoryManager.addItem).not.toHaveBeenCalled();
     });
 
     describe('Interactive Fallback', () => {

packages/cli/src/ui/hooks/useQuotaAndFallback.ts

@@ -55,14 +55,7 @@ export function useQuotaAndFallback({
       fallbackModel,
       error,
     ): Promise<FallbackIntent | null> => {
-      // Fallbacks are currently only handled for OAuth users.
       const contentGeneratorConfig = config.getContentGeneratorConfig();
-      if (
-        !contentGeneratorConfig ||
-        contentGeneratorConfig.authType !== AuthType.LOGIN_WITH_GOOGLE
-      ) {
-        return null;
-      }
 
       let message: string;
       let isTerminalQuotaError = false;
@@ -78,7 +71,9 @@ export function useQuotaAndFallback({
           error.retryDelayMs ? getResetTimeMessage(error.retryDelayMs) : null,
           `/stats model for usage details`,
           `/model to switch models.`,
-          `/auth to switch to API key.`,
+          contentGeneratorConfig?.authType === AuthType.LOGIN_WITH_GOOGLE
+            ? `/auth to switch to API key.`
+            : null,
         ].filter(Boolean);
         message = messageLines.join('\n');
       } else if (error instanceof ModelNotFoundError) {
@@ -122,6 +117,7 @@ export function useQuotaAndFallback({
             message,
             isTerminalQuotaError,
             isModelNotFoundError,
+            authType: contentGeneratorConfig?.authType,
           });
         },
       );

packages/core/src/core/contentGenerator.ts

@@ -77,6 +77,12 @@ export function getAuthTypeFromEnv(): AuthType | undefined {
   if (process.env['GEMINI_API_KEY']) {
     return AuthType.USE_GEMINI;
   }
+  if (
+    process.env['CLOUD_SHELL'] === 'true' ||
+    process.env['GEMINI_CLI_USE_COMPUTE_ADC'] === 'true'
+  ) {
+    return AuthType.COMPUTE_ADC;
+  }
   return undefined;
 }