From d3f13c71ae78494b69c36a6d62dd9b3852bdfec2 Mon Sep 17 00:00:00 2001 From: Abhi <43648792+abhipatel12@users.noreply.github.com> Date: Mon, 23 Jun 2025 23:43:00 -0400 Subject: [PATCH] feat: add custom message for 429 errors (#1366) --- packages/cli/src/nonInteractiveCli.test.ts | 1 + packages/cli/src/nonInteractiveCli.ts | 7 ++- .../cli/src/ui/hooks/useGeminiStream.test.tsx | 58 ++++++++++++++++++- packages/cli/src/ui/hooks/useGeminiStream.ts | 13 +++-- .../cli/src/ui/utils/errorParsing.test.ts | 46 +++++++++++---- packages/cli/src/ui/utils/errorParsing.ts | 33 +++++++++-- 6 files changed, 133 insertions(+), 25 deletions(-) diff --git a/packages/cli/src/nonInteractiveCli.test.ts b/packages/cli/src/nonInteractiveCli.test.ts index 2bee1f24..3b6f1f2d 100644 --- a/packages/cli/src/nonInteractiveCli.test.ts +++ b/packages/cli/src/nonInteractiveCli.test.ts @@ -53,6 +53,7 @@ describe('runNonInteractive', () => { mockConfig = { getToolRegistry: vi.fn().mockReturnValue(mockToolRegistry), getGeminiClient: vi.fn().mockReturnValue(mockGeminiClient), + getContentGeneratorConfig: vi.fn().mockReturnValue({}), } as unknown as Config; mockProcessStdoutWrite = vi.fn().mockImplementation(() => true); diff --git a/packages/cli/src/nonInteractiveCli.ts b/packages/cli/src/nonInteractiveCli.ts index 02bbfd52..28db73f1 100644 --- a/packages/cli/src/nonInteractiveCli.ts +++ b/packages/cli/src/nonInteractiveCli.ts @@ -128,7 +128,12 @@ export async function runNonInteractive( } } } catch (error) { - console.error(parseAndFormatApiError(error)); + console.error( + parseAndFormatApiError( + error, + config.getContentGeneratorConfig().authType, + ), + ); process.exit(1); } finally { if (isTelemetrySdkInitialized()) { diff --git a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx index 29871e8a..ed1eea4d 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx +++ b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx @@ -16,7 +16,7 @@ import { TrackedExecutingToolCall, TrackedCancelledToolCall, } from './useReactToolScheduler.js'; -import { Config, EditorType } from '@gemini-cli/core'; +import { Config, EditorType, AuthType } from '@gemini-cli/core'; import { Part, PartListUnion } from '@google/genai'; import { UseHistoryManagerReturn } from './useHistoryManager.js'; import { HistoryItem, MessageType, StreamingState } from '../types.js'; @@ -117,6 +117,11 @@ vi.mock('./slashCommandProcessor.js', () => ({ handleSlashCommand: vi.fn().mockReturnValue(false), })); +const mockParseAndFormatApiError = vi.hoisted(() => vi.fn()); +vi.mock('../utils/errorParsing.js', () => ({ + parseAndFormatApiError: mockParseAndFormatApiError, +})); + // --- END MOCKS --- describe('mergePartListUnions', () => { @@ -1033,4 +1038,55 @@ describe('useGeminiStream', () => { }); }); }); + + describe('Error Handling', () => { + it('should call parseAndFormatApiError with the correct authType on stream initialization failure', async () => { + // 1. Setup + const mockError = new Error('Rate limit exceeded'); + const mockAuthType = AuthType.LOGIN_WITH_GOOGLE_ENTERPRISE; + mockParseAndFormatApiError.mockClear(); + mockSendMessageStream.mockReturnValue( + (async function* () { + yield { type: 'content', value: '' }; + throw mockError; + })(), + ); + + const testConfig = { + ...mockConfig, + getContentGeneratorConfig: vi.fn(() => ({ + authType: mockAuthType, + })), + } as unknown as Config; + + const { result } = renderHook(() => + useGeminiStream( + new MockedGeminiClientClass(testConfig), + [], + mockAddItem, + mockSetShowHelp, + testConfig, + mockOnDebugMessage, + mockHandleSlashCommand, + false, + () => 'vscode' as EditorType, + () => {}, + () => Promise.resolve(), + ), + ); + + // 2. Action + await act(async () => { + await result.current.submitQuery('test query'); + }); + + // 3. Assertion + await waitFor(() => { + expect(mockParseAndFormatApiError).toHaveBeenCalledWith( + 'Rate limit exceeded', + mockAuthType, + ); + }); + }); + }); }); diff --git a/packages/cli/src/ui/hooks/useGeminiStream.ts b/packages/cli/src/ui/hooks/useGeminiStream.ts index 86540b68..a8816f98 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.ts +++ b/packages/cli/src/ui/hooks/useGeminiStream.ts @@ -398,12 +398,15 @@ export const useGeminiStream = ( addItem( { type: MessageType.ERROR, - text: parseAndFormatApiError(eventValue.error), + text: parseAndFormatApiError( + eventValue.error, + config.getContentGeneratorConfig().authType, + ), }, userMessageTimestamp, ); }, - [addItem, pendingHistoryItemRef, setPendingHistoryItem], + [addItem, pendingHistoryItemRef, setPendingHistoryItem, config], ); const handleChatCompressionEvent = useCallback( @@ -533,10 +536,6 @@ export const useGeminiStream = ( setPendingHistoryItem(null); } } catch (error: unknown) { - console.log( - 'GEMINI_DEBUG: Caught error in useGeminiStream.ts:', - JSON.stringify(error), - ); if (error instanceof UnauthorizedError) { onAuthError(); } else if (!isNodeError(error) || error.name !== 'AbortError') { @@ -545,6 +544,7 @@ export const useGeminiStream = ( type: MessageType.ERROR, text: parseAndFormatApiError( getErrorMessage(error) || 'Unknown error', + config.getContentGeneratorConfig().authType, ), }, userMessageTimestamp, @@ -566,6 +566,7 @@ export const useGeminiStream = ( geminiClient, startNewTurn, onAuthError, + config, ], ); diff --git a/packages/cli/src/ui/utils/errorParsing.test.ts b/packages/cli/src/ui/utils/errorParsing.test.ts index 0dbd75c8..ccc2522e 100644 --- a/packages/cli/src/ui/utils/errorParsing.test.ts +++ b/packages/cli/src/ui/utils/errorParsing.test.ts @@ -6,11 +6,12 @@ import { describe, it, expect } from 'vitest'; import { parseAndFormatApiError } from './errorParsing.js'; -import { StructuredError } from '@gemini-cli/core'; +import { AuthType, StructuredError } from '@gemini-cli/core'; describe('parseAndFormatApiError', () => { - const rateLimitMessage = - 'Please wait and try again later. To increase your limits, upgrade to a plan with higher limits, or use /auth to switch to using a paid API key from AI Studio at https://aistudio.google.com/apikey'; + const enterpriseMessage = 'upgrade to a plan with higher limits'; + const vertexMessage = 'request a quota increase through Vertex'; + const geminiMessage = 'request a quota increase through AI Studio'; it('should format a valid API error JSON', () => { const errorMessage = @@ -20,11 +21,31 @@ describe('parseAndFormatApiError', () => { expect(parseAndFormatApiError(errorMessage)).toBe(expected); }); - it('should format a 429 API error JSON with the custom message', () => { + it('should format a 429 API error with the default message', () => { const errorMessage = 'got status: 429 Too Many Requests. {"error":{"code":429,"message":"Rate limit exceeded","status":"RESOURCE_EXHAUSTED"}}'; - const expected = `[API Error: Rate limit exceeded (Status: RESOURCE_EXHAUSTED)]\n${rateLimitMessage}`; - expect(parseAndFormatApiError(errorMessage)).toBe(expected); + const result = parseAndFormatApiError(errorMessage); + expect(result).toContain('[API Error: Rate limit exceeded'); + expect(result).toContain('Your request has been rate limited'); + }); + + it('should format a 429 API error with the enterprise message', () => { + const errorMessage = + 'got status: 429 Too Many Requests. {"error":{"code":429,"message":"Rate limit exceeded","status":"RESOURCE_EXHAUSTED"}}'; + const result = parseAndFormatApiError( + errorMessage, + AuthType.LOGIN_WITH_GOOGLE_ENTERPRISE, + ); + expect(result).toContain('[API Error: Rate limit exceeded'); + expect(result).toContain(enterpriseMessage); + }); + + it('should format a 429 API error with the vertex message', () => { + const errorMessage = + 'got status: 429 Too Many Requests. {"error":{"code":429,"message":"Rate limit exceeded","status":"RESOURCE_EXHAUSTED"}}'; + const result = parseAndFormatApiError(errorMessage, AuthType.USE_VERTEX_AI); + expect(result).toContain('[API Error: Rate limit exceeded'); + expect(result).toContain(vertexMessage); }); it('should return the original message if it is not a JSON error', () => { @@ -66,9 +87,9 @@ describe('parseAndFormatApiError', () => { }, }); - const expected = `[API Error: Gemini 2.5 Pro Preview doesn't have a free quota tier. For more information on this error, head to: https://ai.google.dev/gemini-api/docs/rate-limits. (Status: Too Many Requests)]\n${rateLimitMessage}`; - - expect(parseAndFormatApiError(errorMessage)).toBe(expected); + const result = parseAndFormatApiError(errorMessage, AuthType.USE_GEMINI); + expect(result).toContain('Gemini 2.5 Pro Preview'); + expect(result).toContain(geminiMessage); }); it('should format a StructuredError', () => { @@ -80,13 +101,14 @@ describe('parseAndFormatApiError', () => { expect(parseAndFormatApiError(error)).toBe(expected); }); - it('should format a 429 StructuredError with the custom message', () => { + it('should format a 429 StructuredError with the vertex message', () => { const error: StructuredError = { message: 'Rate limit exceeded', status: 429, }; - const expected = `[API Error: Rate limit exceeded]\n${rateLimitMessage}`; - expect(parseAndFormatApiError(error)).toBe(expected); + const result = parseAndFormatApiError(error, AuthType.USE_VERTEX_AI); + expect(result).toContain('[API Error: Rate limit exceeded]'); + expect(result).toContain(vertexMessage); }); it('should handle an unknown error type', () => { diff --git a/packages/cli/src/ui/utils/errorParsing.ts b/packages/cli/src/ui/utils/errorParsing.ts index 1aca15ae..b0f2d00a 100644 --- a/packages/cli/src/ui/utils/errorParsing.ts +++ b/packages/cli/src/ui/utils/errorParsing.ts @@ -4,10 +4,16 @@ * SPDX-License-Identifier: Apache-2.0 */ -import { StructuredError } from '@gemini-cli/core'; +import { AuthType, StructuredError } from '@gemini-cli/core'; -const RATE_LIMIT_ERROR_MESSAGE = +const RATE_LIMIT_ERROR_MESSAGE_GOOGLE = '\nPlease wait and try again later. To increase your limits, upgrade to a plan with higher limits, or use /auth to switch to using a paid API key from AI Studio at https://aistudio.google.com/apikey'; +const RATE_LIMIT_ERROR_MESSAGE_USE_GEMINI = + '\nPlease wait and try again later. To increase your limits, request a quota increase through AI Studio, or switch to another /auth method'; +const RATE_LIMIT_ERROR_MESSAGE_VERTEX = + '\nPlease wait and try again later. To increase your limits, request a quota increase through Vertex, or switch to another /auth method'; +const RATE_LIMIT_ERROR_MESSAGE_DEFAULT = + 'Your request has been rate limited. Please wait and try again later.'; export interface ApiError { error: { @@ -37,11 +43,28 @@ function isStructuredError(error: unknown): error is StructuredError { ); } -export function parseAndFormatApiError(error: unknown): string { +function getRateLimitMessage(authType?: AuthType): string { + switch (authType) { + case AuthType.LOGIN_WITH_GOOGLE_PERSONAL: + case AuthType.LOGIN_WITH_GOOGLE_ENTERPRISE: + return RATE_LIMIT_ERROR_MESSAGE_GOOGLE; + case AuthType.USE_GEMINI: + return RATE_LIMIT_ERROR_MESSAGE_USE_GEMINI; + case AuthType.USE_VERTEX_AI: + return RATE_LIMIT_ERROR_MESSAGE_VERTEX; + default: + return RATE_LIMIT_ERROR_MESSAGE_DEFAULT; + } +} + +export function parseAndFormatApiError( + error: unknown, + authType?: AuthType, +): string { if (isStructuredError(error)) { let text = `[API Error: ${error.message}]`; if (error.status === 429) { - text += RATE_LIMIT_ERROR_MESSAGE; + text += getRateLimitMessage(authType); } return text; } @@ -70,7 +93,7 @@ export function parseAndFormatApiError(error: unknown): string { } let text = `[API Error: ${finalMessage} (Status: ${parsedError.error.status})]`; if (parsedError.error.code === 429) { - text += RATE_LIMIT_ERROR_MESSAGE; + text += getRateLimitMessage(authType); } return text; }