diff --git a/packages/cli/src/ui/components/DialogManager.tsx b/packages/cli/src/ui/components/DialogManager.tsx index acd2f3472f5..1b5c00c3723 100644 --- a/packages/cli/src/ui/components/DialogManager.tsx +++ b/packages/cli/src/ui/components/DialogManager.tsx @@ -86,6 +86,7 @@ export const DialogManager = ({ message={quotaState.proQuotaRequest.message} isTerminalQuotaError={quotaState.proQuotaRequest.isTerminalQuotaError} isModelNotFoundError={!!quotaState.proQuotaRequest.isModelNotFoundError} + isCapacityExceeded={!!quotaState.proQuotaRequest.isCapacityExceeded} authType={quotaState.proQuotaRequest.authType} tierName={config?.getUserTierName()} onChoice={uiActions.handleProQuotaChoice} diff --git a/packages/cli/src/ui/components/ProQuotaDialog.test.tsx b/packages/cli/src/ui/components/ProQuotaDialog.test.tsx index 1f1ece6ca67..0d8d3fc5f93 100644 --- a/packages/cli/src/ui/components/ProQuotaDialog.test.tsx +++ b/packages/cli/src/ui/components/ProQuotaDialog.test.tsx @@ -271,6 +271,40 @@ describe('ProQuotaDialog', () => { ); unmount(); }); + + it('should render keep trying, switch, and stop options even if isTerminalQuotaError is true when isCapacityExceeded is true', async () => { + const { unmount } = await render( + , + ); + + expect(RadioButtonSelect).toHaveBeenCalledWith( + expect.objectContaining({ + items: [ + { + label: 'Keep trying', + value: 'retry_once', + key: 'retry_once', + }, + { + label: 'Switch to gemini-2.5-flash', + value: 'retry_always', + key: 'retry_always', + }, + { label: 'Stop', value: 'retry_later', key: 'retry_later' }, + ], + }), + undefined, + ); + unmount(); + }); }); describe('when it is a model not found error', () => { diff --git a/packages/cli/src/ui/components/ProQuotaDialog.tsx b/packages/cli/src/ui/components/ProQuotaDialog.tsx index e9e869edb0d..4e0ecee0b4e 100644 --- a/packages/cli/src/ui/components/ProQuotaDialog.tsx +++ b/packages/cli/src/ui/components/ProQuotaDialog.tsx @@ -17,6 +17,7 @@ interface ProQuotaDialogProps { message: string; isTerminalQuotaError: boolean; isModelNotFoundError?: boolean; + isCapacityExceeded?: boolean; authType?: AuthType; tierName?: string; onChoice: ( @@ -30,6 +31,7 @@ export function ProQuotaDialog({ message, isTerminalQuotaError, isModelNotFoundError, + isCapacityExceeded, authType, tierName, onChoice, @@ -49,6 +51,24 @@ export function ProQuotaDialog({ key: 'retry_later', }, ]; + } else if (isCapacityExceeded) { + items = [ + { + label: 'Keep trying', + value: 'retry_once' as const, + key: 'retry_once', + }, + { + label: `Switch to ${fallbackModel}`, + value: 'retry_always' as const, + key: 'retry_always', + }, + { + label: 'Stop', + value: 'retry_later' as const, + key: 'retry_later', + }, + ]; } else if (isModelNotFoundError || isTerminalQuotaError) { const isUltra = isUltraTier(tierName); @@ -75,7 +95,7 @@ export function ProQuotaDialog({ }, ]; } else { - // capacity error + // capacity error or generic fallback items = [ { label: 'Keep trying', diff --git a/packages/cli/src/ui/contexts/UIStateContext.tsx b/packages/cli/src/ui/contexts/UIStateContext.tsx index eb998a9de05..77412a0f48e 100644 --- a/packages/cli/src/ui/contexts/UIStateContext.tsx +++ b/packages/cli/src/ui/contexts/UIStateContext.tsx @@ -40,6 +40,7 @@ export interface ProQuotaDialogRequest { message: string; isTerminalQuotaError: boolean; isModelNotFoundError?: boolean; + isCapacityExceeded?: boolean; authType?: AuthType; resolve: (intent: FallbackIntent) => void; } diff --git a/packages/cli/src/ui/hooks/useQuotaAndFallback.test.ts b/packages/cli/src/ui/hooks/useQuotaAndFallback.test.ts index b032dc28ecd..d2955b23590 100644 --- a/packages/cli/src/ui/hooks/useQuotaAndFallback.test.ts +++ b/packages/cli/src/ui/hooks/useQuotaAndFallback.test.ts @@ -222,6 +222,126 @@ describe('useQuotaAndFallback', () => { await promise!; }); + it('should auto-retry terminal quota capacity failures in low verbosity mode', async () => { + const { result } = await renderHook(() => + useQuotaAndFallback({ + config: mockConfig, + historyManager: mockHistoryManager, + userTier: UserTierId.FREE, + setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError, + onShowAuthSelection: mockOnShowAuthSelection, + paidTier: null, + settings: mockSettings, + errorVerbosity: 'low', + }), + ); + + const handler = setFallbackHandlerSpy.mock + .calls[0][0] as FallbackModelHandler; + const intent = await handler( + 'gemini-pro', + 'gemini-flash', + new TerminalQuotaError( + 'pro capacity exhausted', + mockGoogleApiError, + undefined, + 'MODEL_CAPACITY_EXHAUSTED', + ), + ); + + expect(intent).toBe('retry_once'); + expect(result.current.proQuotaRequest).toBeNull(); + }); + + it('should auto-retry capacity failures matched by regex on message in low verbosity mode', async () => { + const { result } = await renderHook(() => + useQuotaAndFallback({ + config: mockConfig, + historyManager: mockHistoryManager, + userTier: UserTierId.FREE, + setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError, + onShowAuthSelection: mockOnShowAuthSelection, + paidTier: null, + settings: mockSettings, + errorVerbosity: 'low', + }), + ); + + const handler = setFallbackHandlerSpy.mock + .calls[0][0] as FallbackModelHandler; + const intent = await handler( + 'gemini-pro', + 'gemini-flash', + new Error('you have exhausted your capacity limit'), + ); + + expect(intent).toBe('retry_once'); + expect(result.current.proQuotaRequest).toBeNull(); + }); + + it('should auto-retry capacity failures thrown as raw string error in low verbosity mode', async () => { + const { result } = await renderHook(() => + useQuotaAndFallback({ + config: mockConfig, + historyManager: mockHistoryManager, + userTier: UserTierId.FREE, + setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError, + onShowAuthSelection: mockOnShowAuthSelection, + paidTier: null, + settings: mockSettings, + errorVerbosity: 'low', + }), + ); + + const handler = setFallbackHandlerSpy.mock + .calls[0][0] as FallbackModelHandler; + const intent = await handler( + 'gemini-pro', + 'gemini-flash', + 'MODEL_CAPACITY_EXHAUSTED', + ); + + expect(intent).toBe('retry_once'); + expect(result.current.proQuotaRequest).toBeNull(); + }); + + it('should show high demand message for MODEL_CAPACITY_EXHAUSTED', async () => { + const { result } = await renderHook(() => + useQuotaAndFallback({ + config: mockConfig, + historyManager: mockHistoryManager, + userTier: UserTierId.FREE, + setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError, + onShowAuthSelection: mockOnShowAuthSelection, + paidTier: null, + settings: mockSettings, + }), + ); + + const handler = setFallbackHandlerSpy.mock + .calls[0][0] as FallbackModelHandler; + + const error = new TerminalQuotaError( + 'pro capacity exhausted', + mockGoogleApiError, + undefined, + 'MODEL_CAPACITY_EXHAUSTED', + ); + + act(() => { + void handler('gemini-pro', 'gemini-flash', error); + }); + + expect(result.current.proQuotaRequest).not.toBeNull(); + expect(result.current.proQuotaRequest?.isCapacityExceeded).toBe(true); + expect(result.current.proQuotaRequest?.message).toContain( + 'We are currently experiencing high demand', + ); + expect(result.current.proQuotaRequest?.message).not.toContain( + 'Usage limit reached', + ); + }); + describe('Interactive Fallback', () => { it('should set an interactive request for a terminal quota error', async () => { const { result } = await renderHook(() => @@ -1040,7 +1160,7 @@ Your admin might have disabled the access. Contact them to enable the Preview Re ); }); - it('should show a special message when falling back from the preview model, but do not show periodical check message for flash model fallback', async () => { + it('should show a special message when falling back from the preview model, but not show the periodical check message for flash model fallbacks', async () => { const { result } = await renderHook(() => useQuotaAndFallback({ config: mockConfig, diff --git a/packages/cli/src/ui/hooks/useQuotaAndFallback.ts b/packages/cli/src/ui/hooks/useQuotaAndFallback.ts index a8e757cca49..4a91b19da83 100644 --- a/packages/cli/src/ui/hooks/useQuotaAndFallback.ts +++ b/packages/cli/src/ui/hooks/useQuotaAndFallback.ts @@ -45,6 +45,9 @@ interface UseQuotaAndFallbackArgs { errorVerbosity?: 'low' | 'full'; } +const isObject = (val: unknown): val is Record => + typeof val === 'object' && val !== null; + export function useQuotaAndFallback({ config, historyManager, @@ -79,6 +82,28 @@ export function useQuotaAndFallback({ let message: string; let isTerminalQuotaError = false; let isModelNotFoundError = false; + + const errorObj = isObject(error) ? error : null; + + const errorReasonValue = errorObj?.['reason']; + const errorReason = + typeof errorReasonValue === 'string' ? errorReasonValue : undefined; + + const errorMessageValue = errorObj?.['message']; + const errorMessage = + typeof errorMessageValue === 'string' ? errorMessageValue : undefined; + + const isCapacityExceeded = + errorReason === 'MODEL_CAPACITY_EXHAUSTED' || + errorReason === 'MODEL_CAPACITY_EXCEEDED' || + (typeof errorMessage === 'string' && + /exhausted your capacity|capacity exceeded|MODEL_CAPACITY_EXHAUSTED/i.test( + errorMessage, + )) || + (typeof error === 'string' && + /exhausted your capacity|capacity exceeded|MODEL_CAPACITY_EXHAUSTED/i.test( + error, + )); const usageLimitReachedModel = isProModel(failedModel) ? 'all Pro models' : failedModel; @@ -121,18 +146,30 @@ export function useQuotaAndFallback({ } // Default: Show existing ProQuotaDialog (for overageStrategy: 'never' or non-G1 users) - const messageLines = [ - `Usage limit reached for ${usageLimitReachedModel}.`, - error.retryDelayMs - ? `Access resets at ${getResetTimeMessage(error.retryDelayMs)}.` - : null, - `/stats model for usage details`, - `/model to switch models.`, - contentGeneratorConfig?.authType === AuthType.LOGIN_WITH_GOOGLE - ? `/auth to switch to API key.` - : null, - ].filter(Boolean); - message = messageLines.join('\n'); + if (isCapacityExceeded) { + const messageLines = [ + `We are currently experiencing high demand for ${usageLimitReachedModel}.`, + 'We apologize and appreciate your patience.', + error.retryDelayMs + ? `Access resets at ${getResetTimeMessage(error.retryDelayMs)}.` + : null, + `/model to switch models.`, + ].filter(Boolean); + message = messageLines.join('\n'); + } else { + const messageLines = [ + `Usage limit reached for ${usageLimitReachedModel}.`, + error.retryDelayMs + ? `Access resets at ${getResetTimeMessage(error.retryDelayMs)}.` + : null, + `/stats model for usage details`, + `/model to switch models.`, + contentGeneratorConfig?.authType === AuthType.LOGIN_WITH_GOOGLE + ? `/auth to switch to API key.` + : null, + ].filter(Boolean); + message = messageLines.join('\n'); + } } else if (error instanceof ModelNotFoundError) { isModelNotFoundError = true; if ( @@ -174,7 +211,7 @@ export function useQuotaAndFallback({ // without interrupting with a dialog. if ( errorVerbosity === 'low' && - !isTerminalQuotaError && + (!isTerminalQuotaError || isCapacityExceeded) && !isModelNotFoundError ) { return 'retry_once'; @@ -197,6 +234,7 @@ export function useQuotaAndFallback({ message, isTerminalQuotaError, isModelNotFoundError, + isCapacityExceeded, authType: contentGeneratorConfig?.authType, }); }, diff --git a/packages/core/src/config/config.test.ts b/packages/core/src/config/config.test.ts index 69a1dc5c292..7863f4f46e9 100644 --- a/packages/core/src/config/config.test.ts +++ b/packages/core/src/config/config.test.ts @@ -3198,6 +3198,24 @@ describe('Config Quota & Preview Model Access', () => { expect(config.getHasAccessToPreviewModel()).toBe(false); }); + it('should reverse-map gemini-3-flash back to gemini-3.5-flash in modelQuotas', async () => { + mockCodeAssistServer.retrieveUserQuota.mockResolvedValue({ + buckets: [ + { + modelId: 'gemini-3-flash', + remainingAmount: '90', + remainingFraction: 0.9, + }, + ], + }); + + config.setModel('gemini-3.5-flash'); + await config.refreshUserQuota(); + + expect(config.getQuotaRemaining()).toBe(90); + expect(config.getQuotaLimit()).toBe(100); + }); + it('should calculate pooled quota correctly for auto models', async () => { mockCodeAssistServer.retrieveUserQuota.mockResolvedValue({ buckets: [ @@ -4134,7 +4152,9 @@ describe('Plans Directory Initialization', () => { const plansDir = config.storage.getPlansDir(); // Should NOT create the directory eagerly - expect(fs.promises.mkdir).not.toHaveBeenCalled(); + expect(fs.promises.mkdir).not.toHaveBeenCalledWith(plansDir, { + recursive: true, + }); // Should check if it exists expect(fs.promises.access).toHaveBeenCalledWith(plansDir); @@ -4152,7 +4172,9 @@ describe('Plans Directory Initialization', () => { await config.initialize(); const plansDir = config.storage.getPlansDir(); - expect(fs.promises.mkdir).not.toHaveBeenCalled(); + expect(fs.promises.mkdir).not.toHaveBeenCalledWith(plansDir, { + recursive: true, + }); expect(fs.promises.access).toHaveBeenCalledWith(plansDir); const context = config.getWorkspaceContext(); diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index 59e7852b5c9..a685886e83d 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -87,6 +87,8 @@ import { PREVIEW_GEMINI_FLASH_MODEL, resolveModel, setFlashModels, + DEFAULT_GEMINI_3_5_FLASH_MODEL, + SECONDARY_GEMINI_3_5_FLASH_MODEL, } from './models.js'; import { shouldAttemptBrowserLaunch } from '../utils/browser.js'; import type { MCPOAuthConfig } from '../mcp/oauth-provider.js'; @@ -2320,6 +2322,11 @@ export class Config implements McpContext, AgentLoopContext { continue; } + let modelId = bucket.modelId; + if (modelId === SECONDARY_GEMINI_3_5_FLASH_MODEL) { + modelId = DEFAULT_GEMINI_3_5_FLASH_MODEL; + } + let remaining: number; let limit: number; @@ -2328,7 +2335,7 @@ export class Config implements McpContext, AgentLoopContext { limit = bucket.remainingFraction > 0 ? Math.round(remaining / bucket.remainingFraction) - : (this.modelQuotas.get(bucket.modelId)?.limit ?? 0); + : (this.modelQuotas.get(modelId)?.limit ?? 0); } else { // Server only sent remainingFraction — use a normalized scale. limit = 100; @@ -2336,7 +2343,7 @@ export class Config implements McpContext, AgentLoopContext { } if (!isNaN(remaining) && Number.isFinite(limit) && limit > 0) { - this.modelQuotas.set(bucket.modelId, { + this.modelQuotas.set(modelId, { remaining, limit, resetTime: bucket.resetTime,