Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions packages/core/src/code_assist/experiments/flagNames.ts
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ export const ExperimentFlags = {
PRO_MODEL_NO_ACCESS: 45768879,
GEMINI_3_1_FLASH_LITE_LAUNCHED: 45771641,
DEFAULT_REQUEST_TIMEOUT: 45773134,
MAX_ATTEMPTS: 45774515,
} as const;

export type ExperimentFlagName =
Expand Down
47 changes: 47 additions & 0 deletions packages/core/src/config/config.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -305,6 +305,53 @@ describe('Server Config (config.ts)', () => {
});
expect(config.getMaxAttempts()).toBe(DEFAULT_MAX_ATTEMPTS);
});

it('should use experiment flag if present and valid', () => {
const config = new Config({
...baseParams,
experiments: {
flags: {
[ExperimentFlags.MAX_ATTEMPTS]: {
intValue: '15',
},
},
experimentIds: [],
},
});
expect(config.getMaxAttempts()).toBe(15);
});

it('should fallback to maxAttempts if experiment flag is invalid', () => {
const config = new Config({
...baseParams,
maxAttempts: 5,
experiments: {
flags: {
[ExperimentFlags.MAX_ATTEMPTS]: {
intValue: 'abc',
},
},
experimentIds: [],
},
});
expect(config.getMaxAttempts()).toBe(5);
});

it('should fallback to maxAttempts if experiment flag is non-positive', () => {
const config = new Config({
...baseParams,
maxAttempts: 5,
experiments: {
flags: {
[ExperimentFlags.MAX_ATTEMPTS]: {
intValue: '0',
},
},
experimentIds: [],
},
});
expect(config.getMaxAttempts()).toBe(5);
});
});

beforeEach(() => {
Expand Down
8 changes: 8 additions & 0 deletions packages/core/src/config/config.ts
Original file line number Diff line number Diff line change
Expand Up @@ -3318,6 +3318,14 @@ export class Config implements McpContext, AgentLoopContext {
}

getMaxAttempts(): number {
const flagVal =
this.experiments?.flags?.[ExperimentFlags.MAX_ATTEMPTS]?.intValue;
if (flagVal !== undefined) {
const parsed = parseInt(flagVal, 10);
if (!isNaN(parsed) && parsed > 0) {
return parsed;
}
}
return this.maxAttempts;
}

Expand Down
2 changes: 1 addition & 1 deletion packages/core/src/core/baseLlmClient.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -252,7 +252,7 @@ describe('BaseLlmClient', () => {
expect(retryWithBackoff).toHaveBeenCalledWith(
expect.any(Function),
expect.objectContaining({
maxAttempts: 5,
maxAttempts: 3,
}),
);
});
Expand Down
10 changes: 6 additions & 4 deletions packages/core/src/core/baseLlmClient.ts
Original file line number Diff line number Diff line change
Expand Up @@ -36,8 +36,6 @@ import {
createAvailabilityContextProvider,
} from '../availability/policyHelpers.js';

const DEFAULT_MAX_ATTEMPTS = 5;

/**
* Options for the generateJson utility function.
*/
Expand Down Expand Up @@ -328,7 +326,9 @@ export class BaseLlmClient {
return await retryWithBackoff(apiCall, {
shouldRetryOnContent,
maxAttempts:
availabilityMaxAttempts ?? maxAttempts ?? DEFAULT_MAX_ATTEMPTS,
availabilityMaxAttempts ??
maxAttempts ??
this.config.getMaxAttempts(),
getAvailabilityContext,
onPersistent429: this.config.isInteractive()
? (authType, error) =>
Expand All @@ -339,7 +339,9 @@ export class BaseLlmClient {
retryFetchErrors: this.config.getRetryFetchErrors(),
onRetry: (attempt, error, delayMs) => {
const actualMaxAttempts =
availabilityMaxAttempts ?? maxAttempts ?? DEFAULT_MAX_ATTEMPTS;
availabilityMaxAttempts ??
maxAttempts ??
this.config.getMaxAttempts();
const modelName = getDisplayString(currentModel);
const errorType = getRetryErrorType(error);

Expand Down
2 changes: 1 addition & 1 deletion packages/core/src/core/client.ts
Original file line number Diff line number Diff line change
Expand Up @@ -1133,7 +1133,7 @@ export class GeminiClient {
onPersistent429: onPersistent429Callback,
onValidationRequired: onValidationRequiredCallback,
authType: this.config.getContentGeneratorConfig()?.authType,
maxAttempts: availabilityMaxAttempts,
maxAttempts: availabilityMaxAttempts ?? this.config.getMaxAttempts(),
retryFetchErrors: this.config.getRetryFetchErrors(),
getAvailabilityContext,
onRetry: (attempt, error, delayMs) => {
Expand Down
2 changes: 1 addition & 1 deletion packages/core/src/core/geminiChat.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -176,7 +176,7 @@ describe('GeminiChat', () => {
},
getContentGenerator: vi.fn().mockReturnValue(mockContentGenerator),
getRetryFetchErrors: vi.fn().mockReturnValue(false),
getMaxAttempts: vi.fn().mockReturnValue(10),
getMaxAttempts: vi.fn().mockReturnValue(4),
getUserTier: vi.fn().mockReturnValue(undefined),
modelConfigService: {
getResolvedConfig: vi.fn().mockImplementation((modelConfigKey) => {
Expand Down
9 changes: 2 additions & 7 deletions packages/core/src/core/geminiChat.ts
Original file line number Diff line number Diff line change
Expand Up @@ -78,16 +78,13 @@ export type StreamEvent =
* Options for retrying mid-stream errors (e.g. invalid content or API disconnects).
*/
interface MidStreamRetryOptions {
/** Total number of attempts to make (1 initial + N retries). */
maxAttempts: number;
/** The base delay in milliseconds for backoff. */
initialDelayMs: number;
/** Whether to use exponential backoff instead of linear. */
useExponentialBackoff: boolean;
}

const MID_STREAM_RETRY_OPTIONS: MidStreamRetryOptions = {
maxAttempts: 4, // 1 initial call + 3 retries mid-stream
initialDelayMs: 1000,
useExponentialBackoff: true,
};
Expand Down Expand Up @@ -420,10 +417,8 @@ export class GeminiChat {
: getRetryErrorType(error);

if (isContentError || (isRetryable && !signal.aborted)) {
// The issue requests exactly 3 retries (4 attempts) for API errors during stream iteration.
// Regardless of the global maxAttempts (e.g. 10), we only want to retry these mid-stream API errors
// up to 3 times before finally throwing the error to the user.
const maxMidStreamAttempts = MID_STREAM_RETRY_OPTIONS.maxAttempts;
// We retry mid-stream API errors up to maxAttempts times before finally throwing the error to the user.
const maxMidStreamAttempts = this.context.config.getMaxAttempts();

if (
attempt < maxAttempts - 1 &&
Expand Down
2 changes: 2 additions & 0 deletions packages/core/src/tools/web-fetch.ts
Original file line number Diff line number Diff line change
Expand Up @@ -309,6 +309,7 @@ class WebFetchToolInvocation extends BaseToolInvocation<
return res;
},
{
maxAttempts: this.context.config.getMaxAttempts(),
retryFetchErrors: this.context.config.getRetryFetchErrors(),
onRetry: (attempt, error, delayMs) =>
this.handleRetry(attempt, error, delayMs),
Expand Down Expand Up @@ -643,6 +644,7 @@ ${aggregatedContent}
return res;
},
{
maxAttempts: this.context.config.getMaxAttempts(),
retryFetchErrors: this.context.config.getRetryFetchErrors(),
onRetry: (attempt, error, delayMs) =>
this.handleRetry(attempt, error, delayMs),
Expand Down
112 changes: 111 additions & 1 deletion packages/core/src/utils/retry.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -9,12 +9,13 @@ import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
import { ApiError } from '@google/genai';
import { AuthType } from '../core/contentGenerator.js';
import { type HttpError, ModelNotFoundError } from './httpErrors.js';
import { retryWithBackoff } from './retry.js';
import { retryWithBackoff, isRetryableError } from './retry.js';
import { setSimulate429 } from './testUtils.js';
import { debugLogger } from './debugLogger.js';
import {
TerminalQuotaError,
RetryableQuotaError,
ValidationRequiredError,
} from './googleQuotaErrors.js';
import { PREVIEW_GEMINI_MODEL } from '../config/models.js';
import type { ModelPolicy } from '../availability/modelPolicy.js';
Expand Down Expand Up @@ -332,6 +333,81 @@ describe('retryWithBackoff', () => {
});
});

it('should call onRetry callback on each retry', async () => {
const mockFn = createFailingFunction(2);
const onRetry = vi.fn();
const promise = retryWithBackoff(mockFn, {
maxAttempts: 3,
initialDelayMs: 10,
onRetry,
});

await vi.runAllTimersAsync();

await promise;
expect(onRetry).toHaveBeenCalledTimes(2);
expect(onRetry).toHaveBeenCalledWith(
1,
expect.any(Error),
expect.any(Number),
);
expect(onRetry).toHaveBeenCalledWith(
2,
expect.any(Error),
expect.any(Number),
);
});

it('should handle ValidationRequiredError using onValidationRequired', async () => {
const error = new ValidationRequiredError('Validation required', {} as any);
let validationCalled = false;
const mockFn = vi.fn().mockImplementation(async () => {
if (!validationCalled) {
throw error;
}
return 'success';
});

const onValidationRequired = vi.fn().mockImplementation(async () => {
validationCalled = true;
return 'verify';
});

const promise = retryWithBackoff(mockFn, {
maxAttempts: 3,
initialDelayMs: 10,
onValidationRequired,
});

await vi.runAllTimersAsync();

const result = await promise;
expect(result).toBe('success');
expect(onValidationRequired).toHaveBeenCalledWith(error);
expect(mockFn).toHaveBeenCalledTimes(2);
});

it('should throw ValidationRequiredError if onValidationRequired returns cancel', async () => {
const error = new ValidationRequiredError('Validation required', {} as any);
const mockFn = vi.fn().mockImplementation(async () => {
throw error;
});

const onValidationRequired = vi.fn().mockResolvedValue('cancel');

const promise = retryWithBackoff(mockFn, {
maxAttempts: 3,
initialDelayMs: 10,
onValidationRequired,
});

await expect(promise).rejects.toThrow('Validation required');
await vi.runAllTimersAsync();

expect(error.userHandled).toBe(true);
expect(mockFn).toHaveBeenCalledTimes(1);
});

describe('Fetch error retries', () => {
it("should retry on 'fetch failed' when retryFetchErrors is true", async () => {
const mockFn = vi.fn();
Expand Down Expand Up @@ -886,3 +962,37 @@ describe('retryWithBackoff', () => {
});
});
});

describe('isRetryableError', () => {
it('should return true for 429 errors', () => {
const error = new ApiError({ message: 'Quota exceeded', status: 429 });
expect(isRetryableError(error)).toBe(true);
});

it('should return true for 499 errors', () => {
const error = new ApiError({
message: 'Client closed request',
status: 499,
});
expect(isRetryableError(error)).toBe(true);
});

it('should return true for 500 errors', () => {
const error = new ApiError({
message: 'Internal Server Error',
status: 500,
});
expect(isRetryableError(error)).toBe(true);
});

it('should return false for 400 errors', () => {
const error = new ApiError({ message: 'Bad Request', status: 400 });
expect(isRetryableError(error)).toBe(false);
});

it('should return true for network error codes like ECONNRESET', () => {
const error = new Error('ECONNRESET');
(error as any).code = 'ECONNRESET';
expect(isRetryableError(error)).toBe(true);
});
});
9 changes: 9 additions & 0 deletions packages/core/src/utils/retry.ts
Original file line number Diff line number Diff line change
Expand Up @@ -17,9 +17,18 @@ import { getErrorStatus, ModelNotFoundError } from './httpErrors.js';
import type { RetryAvailabilityContext } from '../availability/modelPolicy.js';

export type { RetryAvailabilityContext };

/**
* Global fallback for maximum retry attempts when not explicitly provided.
* Most callers should use config.getMaxAttempts() instead.
*/
export const DEFAULT_MAX_ATTEMPTS = 10;

export interface RetryOptions {
/**
* Total number of attempts (1 initial + N retries).
* Defaults to DEFAULT_MAX_ATTEMPTS (10) if not specified.
*/
maxAttempts: number;
initialDelayMs: number;
maxDelayMs: number;
Expand Down
Loading