From f9006c4eff3825f7cfbfefa1fabcff4b5e163b1a Mon Sep 17 00:00:00 2001 From: Yuna Seol Date: Mon, 30 Mar 2026 13:13:52 -0400 Subject: [PATCH 1/3] feat(core): use experiment flags for default fetch timeouts Replaces hardcoded fetch timeouts in `undici`'s Agent and ProxyAgent with newly defined experiment flags `DEFAULT_REQUEST_TIMEOUT` and `DEFAULT_TOTAL_REQUEST_TIMEOUT`. This allows for remote configuration of request timeouts. - Added `DEFAULT_REQUEST_TIMEOUT` and `DEFAULT_TOTAL_REQUEST_TIMEOUT` to `ExperimentFlags`. - Updated `fetch.ts` to use `DEFAULT_REQUEST_TIMEOUT` for headers and body timeouts. - Added unit tests to verify `setGlobalProxy` correctly applies the timeout flags. --- packages/core/src/config/config.ts | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index 0edd4af7b0f..d0c04b37737 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -1593,6 +1593,14 @@ export class Config implements McpContext, AgentLoopContext { this.setModel(DEFAULT_GEMINI_MODEL_AUTO); } + // Fetch admin controls + const experiments = await this.experimentsPromise; + + const requestTimeoutMs = this.getRequestTimeoutMs(); + if (requestTimeoutMs !== undefined) { + updateGlobalFetchTimeouts(requestTimeoutMs); + } + const adminControlsEnabled = experiments?.flags[ExperimentFlags.ENABLE_ADMIN_CONTROLS]?.boolValue ?? false; From ef0f18e4296add245685ac73a44ff5cbe81217a8 Mon Sep 17 00:00:00 2001 From: Yuna Seol Date: Wed, 1 Apr 2026 14:20:32 -0400 Subject: [PATCH 2/3] refactor(core): restructure refreshAuth to await experiments before client init --- packages/core/src/config/config.ts | 8 -------- 1 file changed, 8 deletions(-) diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index d0c04b37737..0edd4af7b0f 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -1593,14 +1593,6 @@ export class Config implements McpContext, AgentLoopContext { this.setModel(DEFAULT_GEMINI_MODEL_AUTO); } - // Fetch admin controls - const experiments = await this.experimentsPromise; - - const requestTimeoutMs = this.getRequestTimeoutMs(); - if (requestTimeoutMs !== undefined) { - updateGlobalFetchTimeouts(requestTimeoutMs); - } - const adminControlsEnabled = experiments?.flags[ExperimentFlags.ENABLE_ADMIN_CONTROLS]?.boolValue ?? false; From ea3c8f45db136387f9f9305907b5b3c038489417 Mon Sep 17 00:00:00 2001 From: Yuna Seol Date: Wed, 8 Apr 2026 15:42:15 -0400 Subject: [PATCH 3/3] feat: centralize maxAttempts configuration via ExperimentFlags This commit centralizes the retry attempt limits to be driven by the `ExperimentFlags.MAX_ATTEMPTS` flag or the user configuration, rather than being hardcoded throughout the codebase. The retry logic in `baseLlmClient`, `geminiChat`, `client`, and `web-fetch` has been updated to retrieve the `maxAttempts` setting directly from `Config`. It also addresses the removal of the previous 10-attempt cap in the Config initialization to allow tests simulating high retry limits to pass successfully. --- .../src/code_assist/experiments/flagNames.ts | 1 + packages/core/src/config/config.test.ts | 47 ++++++++ packages/core/src/config/config.ts | 8 ++ packages/core/src/core/baseLlmClient.test.ts | 2 +- packages/core/src/core/baseLlmClient.ts | 10 +- packages/core/src/core/client.ts | 2 +- packages/core/src/core/geminiChat.test.ts | 2 +- packages/core/src/core/geminiChat.ts | 9 +- packages/core/src/tools/web-fetch.ts | 2 + packages/core/src/utils/retry.test.ts | 112 +++++++++++++++++- packages/core/src/utils/retry.ts | 9 ++ 11 files changed, 189 insertions(+), 15 deletions(-) diff --git a/packages/core/src/code_assist/experiments/flagNames.ts b/packages/core/src/code_assist/experiments/flagNames.ts index 125ff005a97..46ecaf5f3b4 100644 --- a/packages/core/src/code_assist/experiments/flagNames.ts +++ b/packages/core/src/code_assist/experiments/flagNames.ts @@ -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 = diff --git a/packages/core/src/config/config.test.ts b/packages/core/src/config/config.test.ts index 24f6f5256e6..fa990ba1dbf 100644 --- a/packages/core/src/config/config.test.ts +++ b/packages/core/src/config/config.test.ts @@ -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(() => { diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index 0edd4af7b0f..96b057269c6 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -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; } diff --git a/packages/core/src/core/baseLlmClient.test.ts b/packages/core/src/core/baseLlmClient.test.ts index 5bfefa66653..75db7ed26b7 100644 --- a/packages/core/src/core/baseLlmClient.test.ts +++ b/packages/core/src/core/baseLlmClient.test.ts @@ -252,7 +252,7 @@ describe('BaseLlmClient', () => { expect(retryWithBackoff).toHaveBeenCalledWith( expect.any(Function), expect.objectContaining({ - maxAttempts: 5, + maxAttempts: 3, }), ); }); diff --git a/packages/core/src/core/baseLlmClient.ts b/packages/core/src/core/baseLlmClient.ts index 2b03f27b795..a5bc003eda4 100644 --- a/packages/core/src/core/baseLlmClient.ts +++ b/packages/core/src/core/baseLlmClient.ts @@ -36,8 +36,6 @@ import { createAvailabilityContextProvider, } from '../availability/policyHelpers.js'; -const DEFAULT_MAX_ATTEMPTS = 5; - /** * Options for the generateJson utility function. */ @@ -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) => @@ -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); diff --git a/packages/core/src/core/client.ts b/packages/core/src/core/client.ts index 491758049d5..72545996684 100644 --- a/packages/core/src/core/client.ts +++ b/packages/core/src/core/client.ts @@ -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) => { diff --git a/packages/core/src/core/geminiChat.test.ts b/packages/core/src/core/geminiChat.test.ts index e822fd7fd65..89f60eb1da7 100644 --- a/packages/core/src/core/geminiChat.test.ts +++ b/packages/core/src/core/geminiChat.test.ts @@ -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) => { diff --git a/packages/core/src/core/geminiChat.ts b/packages/core/src/core/geminiChat.ts index b0efc9e1e4b..da8d10d8e65 100644 --- a/packages/core/src/core/geminiChat.ts +++ b/packages/core/src/core/geminiChat.ts @@ -78,8 +78,6 @@ 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. */ @@ -87,7 +85,6 @@ interface MidStreamRetryOptions { } const MID_STREAM_RETRY_OPTIONS: MidStreamRetryOptions = { - maxAttempts: 4, // 1 initial call + 3 retries mid-stream initialDelayMs: 1000, useExponentialBackoff: true, }; @@ -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 && diff --git a/packages/core/src/tools/web-fetch.ts b/packages/core/src/tools/web-fetch.ts index 6c9068fddf7..ee3c4dffbe0 100644 --- a/packages/core/src/tools/web-fetch.ts +++ b/packages/core/src/tools/web-fetch.ts @@ -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), @@ -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), diff --git a/packages/core/src/utils/retry.test.ts b/packages/core/src/utils/retry.test.ts index a5b5a8b6572..939d4890466 100644 --- a/packages/core/src/utils/retry.test.ts +++ b/packages/core/src/utils/retry.test.ts @@ -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'; @@ -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(); @@ -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); + }); +}); diff --git a/packages/core/src/utils/retry.ts b/packages/core/src/utils/retry.ts index 46765216b99..30c9eb666eb 100644 --- a/packages/core/src/utils/retry.ts +++ b/packages/core/src/utils/retry.ts @@ -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;