From 1e42fdf6c266796e69d6e965653dd02187b7646e Mon Sep 17 00:00:00 2001 From: Adam Weidman <65992621+adamfweidman@users.noreply.github.com> Date: Thu, 6 Nov 2025 17:13:38 -0500 Subject: [PATCH] fix(cli): handle flash model errors gracefully (#12667) --- .../src/ui/hooks/useQuotaAndFallback.test.ts | 65 ++++++++++++++++--- .../cli/src/ui/hooks/useQuotaAndFallback.ts | 34 ++++++---- packages/core/src/fallback/handler.test.ts | 7 +- packages/core/src/fallback/handler.ts | 2 - 4 files changed, 81 insertions(+), 27 deletions(-) diff --git a/packages/cli/src/ui/hooks/useQuotaAndFallback.test.ts b/packages/cli/src/ui/hooks/useQuotaAndFallback.test.ts index 1ab226e351..968e41e0e2 100644 --- a/packages/cli/src/ui/hooks/useQuotaAndFallback.test.ts +++ b/packages/cli/src/ui/hooks/useQuotaAndFallback.test.ts @@ -103,15 +103,6 @@ describe('useQuotaAndFallback', () => { 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({ @@ -125,6 +116,62 @@ describe('useQuotaAndFallback', () => { expect(mockHistoryManager.addItem).not.toHaveBeenCalled(); }); + describe('Flash Model Fallback', () => { + it('should show a terminal quota message and stop, without offering a fallback', async () => { + const handler = getRegisteredHandler(); + const result = await handler( + 'gemini-2.5-flash', + 'gemini-2.5-flash', + new TerminalQuotaError('flash quota', mockGoogleApiError), + ); + + expect(result).toBe('stop'); + expect(mockHistoryManager.addItem).toHaveBeenCalledTimes(1); + const message = (mockHistoryManager.addItem as Mock).mock.calls[0][0] + .text; + expect(message).toContain( + 'You have reached your daily gemini-2.5-flash', + ); + expect(message).not.toContain('continue with the fallback model'); + }); + + it('should show a capacity message and stop', async () => { + const handler = getRegisteredHandler(); + // let result: FallbackIntent | null = null; + const result = await handler( + 'gemini-2.5-flash', + 'gemini-2.5-flash', + new Error('capacity'), + ); + + expect(result).toBe('stop'); + expect(mockHistoryManager.addItem).toHaveBeenCalledTimes(1); + const message = (mockHistoryManager.addItem as Mock).mock.calls[0][0] + .text; + expect(message).toContain( + 'Pardon Our Congestion! It looks like gemini-2.5-flash is very popular', + ); + }); + + it('should show a capacity message and stop, even when already in fallback mode', async () => { + vi.spyOn(mockConfig, 'isInFallbackMode').mockReturnValue(true); + const handler = getRegisteredHandler(); + const result = await handler( + 'gemini-2.5-flash', + 'gemini-2.5-flash', + new Error('capacity'), + ); + + expect(result).toBe('stop'); + expect(mockHistoryManager.addItem).toHaveBeenCalledTimes(1); + const message = (mockHistoryManager.addItem as Mock).mock.calls[0][0] + .text; + expect(message).toContain( + 'Pardon Our Congestion! It looks like gemini-2.5-flash is very popular', + ); + }); + }); + describe('Interactive Fallback', () => { // Pro Quota Errors it('should set an interactive request and wait for user choice', async () => { diff --git a/packages/cli/src/ui/hooks/useQuotaAndFallback.ts b/packages/cli/src/ui/hooks/useQuotaAndFallback.ts index 38fd990a46..63f9fe8d29 100644 --- a/packages/cli/src/ui/hooks/useQuotaAndFallback.ts +++ b/packages/cli/src/ui/hooks/useQuotaAndFallback.ts @@ -11,6 +11,7 @@ import { type FallbackIntent, TerminalQuotaError, UserTierId, + DEFAULT_GEMINI_FLASH_MODEL, } from '@google/gemini-cli-core'; import { useCallback, useEffect, useRef, useState } from 'react'; import { type UseHistoryManagerReturn } from './useHistoryManager.js'; @@ -41,10 +42,6 @@ export function useQuotaAndFallback({ fallbackModel, error, ): Promise => { - if (config.isInFallbackMode()) { - return null; - } - // Fallbacks are currently only handled for OAuth users. const contentGeneratorConfig = config.getContentGeneratorConfig(); if ( @@ -58,28 +55,35 @@ export function useQuotaAndFallback({ const isPaidTier = userTier === UserTierId.LEGACY || userTier === UserTierId.STANDARD; + const isFallbackModel = failedModel === DEFAULT_GEMINI_FLASH_MODEL; let message: string; if (error instanceof TerminalQuotaError) { - // Pro Quota specific messages (Interactive) + // Common part of the message for both tiers + const messageLines = [ + `⚡ You have reached your daily ${failedModel} quota limit.`, + `⚡ You can choose to authenticate with a paid API key${ + isFallbackModel ? '.' : ' or continue with the fallback model.' + }`, + ]; + + // Tier-specific part 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.`, + messageLines.push( `⚡ Increase your limits by using a Gemini API Key. See: https://goo.gle/gemini-cli-docs-auth#gemini-api-key`, `⚡ You can switch authentication methods by typing /auth`, - ].join('\n'); + ); } 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.`, + messageLines.push( `⚡ Increase your limits by `, `⚡ - signing up for a plan with higher limits at https://goo.gle/set-up-gemini-code-assist`, `⚡ - or using a Gemini API Key. See: https://goo.gle/gemini-cli-docs-auth#gemini-api-key`, `⚡ You can switch authentication methods by typing /auth`, - ].join('\n'); + ); } + message = messageLines.join('\n'); } else { + // Capacity error message = [ `🚦Pardon Our Congestion! It looks like ${failedModel} is very popular at the moment.`, `Please retry again later.`, @@ -95,6 +99,10 @@ export function useQuotaAndFallback({ Date.now(), ); + if (isFallbackModel) { + return 'stop'; + } + setModelSwitchedFromQuotaError(true); config.setQuotaErrorOccurred(true); diff --git a/packages/core/src/fallback/handler.test.ts b/packages/core/src/fallback/handler.test.ts index a2a2626037..eb97a7f5e0 100644 --- a/packages/core/src/fallback/handler.test.ts +++ b/packages/core/src/fallback/handler.test.ts @@ -73,14 +73,15 @@ describe('handleFallback', () => { expect(mockConfig.setFallbackMode).not.toHaveBeenCalled(); }); - it('should return null if the failed model is already the fallback model', async () => { + it('should still consult the handler if the failed model is the fallback model', async () => { + mockHandler.mockResolvedValue('stop'); const result = await handleFallback( mockConfig, FALLBACK_MODEL, // Failed model is Flash AUTH_OAUTH, ); - expect(result).toBeNull(); - expect(mockHandler).not.toHaveBeenCalled(); + expect(result).toBe(false); + expect(mockHandler).toHaveBeenCalled(); }); it('should return null if no fallbackHandler is injected in config', async () => { diff --git a/packages/core/src/fallback/handler.ts b/packages/core/src/fallback/handler.ts index 8490e8393e..77bff07798 100644 --- a/packages/core/src/fallback/handler.ts +++ b/packages/core/src/fallback/handler.ts @@ -21,8 +21,6 @@ export async function handleFallback( const fallbackModel = DEFAULT_GEMINI_FLASH_MODEL; - if (failedModel === fallbackModel) return null; - // Consult UI Handler for Intent const fallbackModelHandler = config.fallbackModelHandler; if (typeof fallbackModelHandler !== 'function') return null;