feat: add custom message for 429 errors (#1366)

This commit is contained in:
Abhi 2025-06-23 23:43:00 -04:00 committed by GitHub
parent f7caca5f94
commit d3f13c71ae
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
6 changed files with 133 additions and 25 deletions

View File

@ -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);

View File

@ -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()) {

View File

@ -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,
);
});
});
});
});

View File

@ -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,
],
);

View File

@ -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', () => {

View File

@ -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;
}