MediaMetz/gemini-cli
1
0
Fork 0
mirror of https://github.com/google-gemini/gemini-cli.git synced 2026-07-16 13:00:39 -07:00
Code Issues Packages Projects Releases Wiki Activity

fix: handle request retries and model fallback correctly (#11624)

Browse Source
This commit is contained in:
Gaurav
2025-10-24 11:09:06 -07:00
committed by GitHub
parent 24b2c411c9
commit 95b804ee37
14 changed files with 1357 additions and 804 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/core/src/utils/retry.test.ts
+46 -135
View File
@@ -7,10 +7,15 @@
/* eslint-disable @typescript-eslint/no-explicit-any */
import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
import { ApiError } from '@google/genai';
import { AuthType } from '../core/contentGenerator.js';
import type { HttpError } from './retry.js';
import { retryWithBackoff } from './retry.js';
import { setSimulate429 } from './testUtils.js';
import { debugLogger } from './debugLogger.js';
import {
TerminalQuotaError,
RetryableQuotaError,
} from './googleQuotaErrors.js';
// Helper to create a mock function that fails a certain number of times
const createFailingFunction = (
@@ -100,26 +105,26 @@ describe('retryWithBackoff', () => {
// Expect it to fail with the error from the 5th attempt.
await Promise.all([
expect(promise).rejects.toThrow('Simulated error attempt 5'),
expect(promise).rejects.toThrow('Simulated error attempt 3'),
vi.runAllTimersAsync(),
]);
expect(mockFn).toHaveBeenCalledTimes(5);
expect(mockFn).toHaveBeenCalledTimes(3);
});
it('should default to 5 maxAttempts if options.maxAttempts is undefined', async () => {
// This function will fail more than 5 times to ensure all retries are used.
it('should default to 3 maxAttempts if options.maxAttempts is undefined', async () => {
// This function will fail more than 3 times to ensure all retries are used.
const mockFn = createFailingFunction(10);
const promise = retryWithBackoff(mockFn, { maxAttempts: undefined });
// Expect it to fail with the error from the 5th attempt.
await Promise.all([
expect(promise).rejects.toThrow('Simulated error attempt 5'),
expect(promise).rejects.toThrow('Simulated error attempt 3'),
vi.runAllTimersAsync(),
]);
expect(mockFn).toHaveBeenCalledTimes(5);
expect(mockFn).toHaveBeenCalledTimes(3);
});
it('should not retry if shouldRetry returns false', async () => {
@@ -336,15 +341,13 @@ describe('retryWithBackoff', () => {
});
describe('Flash model fallback for OAuth users', () => {
it('should trigger fallback for OAuth personal users after persistent 429 errors', async () => {
it('should trigger fallback for OAuth personal users on TerminalQuotaError', async () => {
const fallbackCallback = vi.fn().mockResolvedValue('gemini-2.5-flash');
let fallbackOccurred = false;
const mockFn = vi.fn().mockImplementation(async () => {
if (!fallbackOccurred) {
const error: HttpError = new Error('Rate limit exceeded');
error.status = 429;
throw error;
throw new TerminalQuotaError('Daily limit reached', {} as any);
}
return 'success';
});
@@ -352,154 +355,62 @@ describe('retryWithBackoff', () => {
const promise = retryWithBackoff(mockFn, {
maxAttempts: 3,
initialDelayMs: 100,
onPersistent429: async (authType?: string) => {
onPersistent429: async (authType?: string, error?: unknown) => {
fallbackOccurred = true;
return await fallbackCallback(authType);
return await fallbackCallback(authType, error);
},
authType: 'oauth-personal',
});
// Advance all timers to complete retries
await vi.runAllTimersAsync();
// Should succeed after fallback
await expect(promise).resolves.toBe('success');
// Verify callback was called with correct auth type
expect(fallbackCallback).toHaveBeenCalledWith('oauth-personal');
// Should retry again after fallback
expect(mockFn).toHaveBeenCalledTimes(3); // 2 initial attempts + 1 after fallback
});
it('should NOT trigger fallback for API key users', async () => {
const fallbackCallback = vi.fn();
const mockFn = vi.fn(async () => {
const error: HttpError = new Error('Rate limit exceeded');
error.status = 429;
throw error;
});
const promise = retryWithBackoff(mockFn, {
maxAttempts: 3,
initialDelayMs: 100,
onPersistent429: fallbackCallback,
authType: 'gemini-api-key',
});
// Handle the promise properly to avoid unhandled rejections
const resultPromise = promise.catch((error) => error);
await vi.runAllTimersAsync();
const result = await resultPromise;
// Should fail after all retries without fallback
expect(result).toBeInstanceOf(Error);
expect(result.message).toBe('Rate limit exceeded');
// Callback should not be called for API key users
expect(fallbackCallback).not.toHaveBeenCalled();
});
it('should reset attempt counter and continue after successful fallback', async () => {
let fallbackCalled = false;
const fallbackCallback = vi.fn().mockImplementation(async () => {
fallbackCalled = true;
return 'gemini-2.5-flash';
});
const mockFn = vi.fn().mockImplementation(async () => {
if (!fallbackCalled) {
const error: HttpError = new Error('Rate limit exceeded');
error.status = 429;
throw error;
}
return 'success';
});
const promise = retryWithBackoff(mockFn, {
maxAttempts: 3,
initialDelayMs: 100,
onPersistent429: fallbackCallback,
authType: 'oauth-personal',
});
await vi.runAllTimersAsync();
await expect(promise).resolves.toBe('success');
expect(fallbackCallback).toHaveBeenCalledOnce();
});
it('should continue with original error if fallback is rejected', async () => {
const fallbackCallback = vi.fn().mockResolvedValue(null); // User rejected fallback
const mockFn = vi.fn(async () => {
const error: HttpError = new Error('Rate limit exceeded');
error.status = 429;
throw error;
});
const promise = retryWithBackoff(mockFn, {
maxAttempts: 3,
initialDelayMs: 100,
onPersistent429: fallbackCallback,
authType: 'oauth-personal',
});
// Handle the promise properly to avoid unhandled rejections
const resultPromise = promise.catch((error) => error);
await vi.runAllTimersAsync();
const result = await resultPromise;
// Should fail with original error when fallback is rejected
expect(result).toBeInstanceOf(Error);
expect(result.message).toBe('Rate limit exceeded');
expect(fallbackCallback).toHaveBeenCalledWith(
'oauth-personal',
expect.any(Error),
expect.any(TerminalQuotaError),
);
expect(mockFn).toHaveBeenCalledTimes(2);
});
it('should handle mixed error types (only count consecutive 429s)', async () => {
const fallbackCallback = vi.fn().mockResolvedValue('gemini-2.5-flash');
let attempts = 0;
let fallbackOccurred = false;
it('should use retryDelayMs from RetryableQuotaError', async () => {
const setTimeoutSpy = vi.spyOn(global, 'setTimeout');
const mockFn = vi.fn().mockImplementation(async () => {
attempts++;
if (fallbackOccurred) {
return 'success';
}
if (attempts === 1) {
// First attempt: 500 error (resets consecutive count)
const error: HttpError = new Error('Server error');
error.status = 500;
throw error;
} else {
// Remaining attempts: 429 errors
const error: HttpError = new Error('Rate limit exceeded');
error.status = 429;
throw error;
}
throw new RetryableQuotaError('Per-minute limit', {} as any, 12.345);
});
const promise = retryWithBackoff(mockFn, {
maxAttempts: 5,
maxAttempts: 2,
initialDelayMs: 100,
onPersistent429: async (authType?: string) => {
fallbackOccurred = true;
return await fallbackCallback(authType);
},
authType: 'oauth-personal',
});
// Attach the rejection expectation *before* running timers
// eslint-disable-next-line vitest/valid-expect
const assertionPromise = expect(promise).rejects.toThrow();
await vi.runAllTimersAsync();
await assertionPromise;
await expect(promise).resolves.toBe('success');
// Should trigger fallback after 2 consecutive 429s (attempts 2-3)
expect(fallbackCallback).toHaveBeenCalledWith('oauth-personal');
expect(setTimeoutSpy).toHaveBeenCalledWith(expect.any(Function), 12345);
});
it.each([[AuthType.USE_GEMINI], [AuthType.USE_VERTEX_AI], [undefined]])(
'should not trigger fallback for non-Google auth users (authType: %s) on TerminalQuotaError',
async (authType) => {
const fallbackCallback = vi.fn();
const mockFn = vi.fn().mockImplementation(async () => {
throw new TerminalQuotaError('Daily limit reached', {} as any);
});
const promise = retryWithBackoff(mockFn, {
maxAttempts: 3,
onPersistent429: fallbackCallback,
authType,
});
await expect(promise).rejects.toThrow('Daily limit reached');
expect(fallbackCallback).not.toHaveBeenCalled();
expect(mockFn).toHaveBeenCalledTimes(1);
},
);
});
it('should abort the retry loop when the signal is aborted', async () => {
const abortController = new AbortController();