diff --git a/packages/core/src/core/geminiChat.test.ts b/packages/core/src/core/geminiChat.test.ts index 6e997ec9b1..067f0d160c 100644 --- a/packages/core/src/core/geminiChat.test.ts +++ b/packages/core/src/core/geminiChat.test.ts @@ -10,6 +10,7 @@ import type { GenerateContentConfig, GenerateContentResponse, } from '@google/genai'; +import { ApiError } from '@google/genai'; import type { ContentGenerator } from '../core/contentGenerator.js'; import { GeminiChat, @@ -902,6 +903,168 @@ describe('GeminiChat', () => { const history = chat.getHistory(); expect(history.length).toBe(0); }); + + describe('API error retry behavior', () => { + beforeEach(() => { + // Use a more direct mock for retry testing + mockRetryWithBackoff.mockImplementation(async (apiCall, options) => { + try { + return await apiCall(); + } catch (error) { + if (options?.shouldRetry && options.shouldRetry(error)) { + // Try again + return await apiCall(); + } + throw error; + } + }); + }); + + it('should not retry on 400 Bad Request errors', async () => { + const error400 = new ApiError({ message: 'Bad Request', status: 400 }); + + vi.mocked(mockContentGenerator.generateContentStream).mockRejectedValue( + error400, + ); + + const stream = await chat.sendMessageStream( + 'test-model', + { message: 'test' }, + 'prompt-id-400', + ); + + await expect( + (async () => { + for await (const _ of stream) { + /* consume stream */ + } + })(), + ).rejects.toThrow(error400); + + // Should only be called once (no retry) + expect( + mockContentGenerator.generateContentStream, + ).toHaveBeenCalledTimes(1); + }); + + it('should retry on 429 Rate Limit errors', async () => { + const error429 = new ApiError({ message: 'Rate Limited', status: 429 }); + + vi.mocked(mockContentGenerator.generateContentStream) + .mockRejectedValueOnce(error429) + .mockResolvedValueOnce( + (async function* () { + yield { + candidates: [ + { + content: { parts: [{ text: 'Success after retry' }] }, + finishReason: 'STOP', + }, + ], + } as unknown as GenerateContentResponse; + })(), + ); + + const stream = await chat.sendMessageStream( + 'test-model', + { message: 'test' }, + 'prompt-id-429-retry', + ); + + const events: StreamEvent[] = []; + for await (const event of stream) { + events.push(event); + } + + // Should be called twice (initial + retry) + expect( + mockContentGenerator.generateContentStream, + ).toHaveBeenCalledTimes(2); + + // Should have successful content + expect( + events.some( + (e) => + e.type === StreamEventType.CHUNK && + e.value.candidates?.[0]?.content?.parts?.[0]?.text === + 'Success after retry', + ), + ).toBe(true); + }); + + it('should not retry on schema depth errors', async () => { + const schemaError = new ApiError({ + message: 'Request failed: maximum schema depth exceeded', + status: 500, + }); + + vi.mocked(mockContentGenerator.generateContentStream).mockRejectedValue( + schemaError, + ); + + const stream = await chat.sendMessageStream( + 'test-model', + { message: 'test' }, + 'prompt-id-schema', + ); + + await expect( + (async () => { + for await (const _ of stream) { + /* consume stream */ + } + })(), + ).rejects.toThrow(schemaError); + + // Should only be called once (no retry) + expect( + mockContentGenerator.generateContentStream, + ).toHaveBeenCalledTimes(1); + }); + + it('should retry on 5xx server errors', async () => { + const error500 = new ApiError({ + message: 'Internal Server Error 500', + status: 500, + }); + + vi.mocked(mockContentGenerator.generateContentStream) + .mockRejectedValueOnce(error500) + .mockResolvedValueOnce( + (async function* () { + yield { + candidates: [ + { + content: { parts: [{ text: 'Recovered from 500' }] }, + finishReason: 'STOP', + }, + ], + } as unknown as GenerateContentResponse; + })(), + ); + + const stream = await chat.sendMessageStream( + 'test-model', + { message: 'test' }, + 'prompt-id-500-retry', + ); + + const events: StreamEvent[] = []; + for await (const event of stream) { + events.push(event); + } + + // Should be called twice (initial + retry) + expect( + mockContentGenerator.generateContentStream, + ).toHaveBeenCalledTimes(2); + }); + + afterEach(() => { + // Reset to default behavior + mockRetryWithBackoff.mockImplementation(async (apiCall) => apiCall()); + }); + }); }); it('should correctly retry and append to an existing history mid-conversation', async () => { // 1. Setup @@ -1427,7 +1590,8 @@ describe('GeminiChat', () => { }); describe('Fallback Integration (Retries)', () => { - const error429 = Object.assign(new Error('API Error 429: Quota exceeded'), { + const error429 = new ApiError({ + message: 'API Error 429: Quota exceeded', status: 429, }); diff --git a/packages/core/src/core/geminiChat.ts b/packages/core/src/core/geminiChat.ts index 947242e04b..7a9d083dad 100644 --- a/packages/core/src/core/geminiChat.ts +++ b/packages/core/src/core/geminiChat.ts @@ -15,6 +15,7 @@ import { type Part, type Tool, FinishReason, + ApiError, } from '@google/genai'; import { toParts } from '../code_assist/converter.js'; import { createUserContent } from '@google/genai'; @@ -376,10 +377,11 @@ export class GeminiChat { const streamResponse = await retryWithBackoff(apiCall, { shouldRetry: (error: unknown) => { - if (error instanceof Error && error.message) { + if (error instanceof ApiError && error.message) { + if (error.status === 400) return false; if (isSchemaDepthError(error.message)) return false; - if (error.message.includes('429')) return true; - if (error.message.match(/5\d{2}/)) return true; + if (error.status === 429) return true; + if (error.status >= 500 && error.status < 600) return true; } return false; },