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/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
-    const getRegisteredHandler = (): FallbackModelHandler => {
-      renderHook(() =>
+    it('should show fallback dialog but omit switch to API key message if authType is not LOGIN_WITH_GOOGLE', async () => {
+      // Override the default mock from beforeEach for this specific test
+      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 () => {
-      // Override the default mock from beforeEach for this specific test
-      vi.spyOn(mockConfig, 'getContentGeneratorConfig').mockReturnValue({
-        authType: AuthType.USE_GEMINI,
+      const handler = setFallbackHandlerSpy.mock
+        .calls[0][0] as FallbackModelHandler;
+
+      const error = new TerminalQuotaError(
+        'pro quota',
+        mockGoogleApiError,
+        1000 * 60 * 5,
+      );
+
+      act(() => {
+        void handler('gemini-pro', 'gemini-flash', error);
       });
 
-      const handler = getRegisteredHandler();
-      const result = await handler('gemini-pro', 'gemini-flash', new Error());
-
-      expect(result).toBeNull();
-      expect(mockHistoryManager.addItem).not.toHaveBeenCalled();
+      expect(result.current.proQuotaRequest).not.toBeNull();
+      expect(result.current.proQuotaRequest?.message).not.toContain(
+        '/auth to switch to API key.',
+      );
     });
 
     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,
           });
         },
       );