Skip to content

Commit 9f8e3f3

Browse files
committed
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.
1 parent 15298b2 commit 9f8e3f3

10 files changed

Lines changed: 138 additions & 14 deletions

File tree

packages/cli/src/test-utils/mockConfig.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -136,6 +136,7 @@ export const createMockConfig = (overrides: Partial<Config> = {}): Config =>
136136
getRetryFetchErrors: vi.fn().mockReturnValue(true),
137137
getEnableShellOutputEfficiency: vi.fn().mockReturnValue(true),
138138
getShellToolInactivityTimeout: vi.fn().mockReturnValue(300000),
139+
getRequestTimeoutMs: vi.fn().mockReturnValue(undefined),
139140
getShellExecutionConfig: vi.fn().mockReturnValue({
140141
sandboxManager: new NoopSandboxManager(),
141142
sanitizationConfig: {

packages/core/src/code_assist/experiments/flagNames.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ export const ExperimentFlags = {
1919
GEMINI_3_1_PRO_LAUNCHED: 45760185,
2020
PRO_MODEL_NO_ACCESS: 45768879,
2121
GEMINI_3_1_FLASH_LITE_LAUNCHED: 45771641,
22+
DEFAULT_REQUEST_TIMEOUT: 45773134,
2223
} as const;
2324

2425
export type ExperimentFlagName =

packages/core/src/config/config.test.ts

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -644,6 +644,58 @@ describe('Server Config (config.ts)', () => {
644644
},
645645
);
646646
});
647+
648+
describe('getRequestTimeoutMs', () => {
649+
it('should return undefined if the flag is not set', () => {
650+
const config = new Config(baseParams);
651+
expect(config.getRequestTimeoutMs()).toBeUndefined();
652+
});
653+
654+
it('should return timeout in milliseconds if flag is set', () => {
655+
const config = new Config({
656+
...baseParams,
657+
experiments: {
658+
flags: {
659+
[ExperimentFlags.DEFAULT_REQUEST_TIMEOUT]: {
660+
intValue: '30',
661+
},
662+
},
663+
experimentIds: [],
664+
},
665+
} as unknown as ConfigParameters);
666+
expect(config.getRequestTimeoutMs()).toBe(30000);
667+
});
668+
669+
it('should return undefined if intValue is not a valid integer', () => {
670+
const config = new Config({
671+
...baseParams,
672+
experiments: {
673+
flags: {
674+
[ExperimentFlags.DEFAULT_REQUEST_TIMEOUT]: {
675+
intValue: 'abc',
676+
},
677+
},
678+
experimentIds: [],
679+
},
680+
} as unknown as ConfigParameters);
681+
expect(config.getRequestTimeoutMs()).toBeUndefined();
682+
});
683+
684+
it('should return undefined if intValue is negative', () => {
685+
const config = new Config({
686+
...baseParams,
687+
experiments: {
688+
flags: {
689+
[ExperimentFlags.DEFAULT_REQUEST_TIMEOUT]: {
690+
intValue: '-10',
691+
},
692+
},
693+
experimentIds: [],
694+
},
695+
} as unknown as ConfigParameters);
696+
expect(config.getRequestTimeoutMs()).toBeUndefined();
697+
});
698+
});
647699
});
648700

649701
describe('refreshAuth', () => {

packages/core/src/config/config.ts

Lines changed: 21 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -160,7 +160,7 @@ import {
160160
} from '../code_assist/experiments/experiments.js';
161161
import { AgentRegistry } from '../agents/registry.js';
162162
import { AcknowledgedAgentsService } from '../agents/acknowledgedAgents.js';
163-
import { setGlobalProxy } from '../utils/fetch.js';
163+
import { setGlobalProxy, updateGlobalFetchTimeouts } from '../utils/fetch.js';
164164
import { SubagentTool } from '../agents/subagent-tool.js';
165165
import { ExperimentFlags } from '../code_assist/experiments/flagNames.js';
166166
import { debugLogger } from '../utils/debugLogger.js';
@@ -1584,6 +1584,11 @@ export class Config implements McpContext, AgentLoopContext {
15841584
// Fetch admin controls
15851585
const experiments = await this.experimentsPromise;
15861586

1587+
const requestTimeoutMs = this.getRequestTimeoutMs();
1588+
if (requestTimeoutMs !== undefined) {
1589+
updateGlobalFetchTimeouts(requestTimeoutMs);
1590+
}
1591+
15871592
const adminControlsEnabled =
15881593
experiments?.flags[ExperimentFlags.ENABLE_ADMIN_CONTROLS]?.boolValue ??
15891594
false;
@@ -3149,6 +3154,21 @@ export class Config implements McpContext, AgentLoopContext {
31493154
);
31503155
}
31513156

3157+
/**
3158+
* Returns the configured default request timeout in milliseconds.
3159+
*/
3160+
getRequestTimeoutMs(): number | undefined {
3161+
const flag =
3162+
this.experiments?.flags?.[ExperimentFlags.DEFAULT_REQUEST_TIMEOUT];
3163+
if (flag?.intValue !== undefined) {
3164+
const seconds = parseInt(flag.intValue, 10);
3165+
if (Number.isInteger(seconds) && seconds >= 0) {
3166+
return seconds * 1000; // Convert seconds to milliseconds
3167+
}
3168+
}
3169+
return undefined;
3170+
}
3171+
31523172
/**
31533173
* Returns whether Gemini 3.1 Flash Lite has been launched.
31543174
*

packages/core/src/core/baseLlmClient.test.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -102,6 +102,7 @@ describe('BaseLlmClient', () => {
102102
);
103103

104104
mockConfig = {
105+
getRequestTimeoutMs: vi.fn().mockReturnValue(undefined),
105106
getSessionId: vi.fn().mockReturnValue('test-session-id'),
106107
getContentGeneratorConfig: vi
107108
.fn()

packages/core/src/core/client.test.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -203,6 +203,7 @@ describe('Gemini Client (client.ts)', () => {
203203
authType: AuthType.USE_GEMINI,
204204
};
205205
mockConfig = {
206+
getRequestTimeoutMs: vi.fn().mockReturnValue(undefined),
206207
getContentGeneratorConfig: vi
207208
.fn()
208209
.mockReturnValue(contentGeneratorConfig),

packages/core/src/core/geminiChat.test.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -142,6 +142,7 @@ describe('GeminiChat', () => {
142142
let currentActiveModel = 'gemini-pro';
143143

144144
mockConfig = {
145+
getRequestTimeoutMs: vi.fn().mockReturnValue(undefined),
145146
get config() {
146147
return this;
147148
},

packages/core/src/core/geminiChat_network_retry.test.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -83,6 +83,7 @@ describe('GeminiChat Network Retries', () => {
8383
const testMessageBus = { publish: vi.fn(), subscribe: vi.fn() };
8484

8585
mockConfig = {
86+
getRequestTimeoutMs: vi.fn().mockReturnValue(undefined),
8687
get config() {
8788
return this;
8889
},

packages/core/src/utils/fetch.test.ts

Lines changed: 37 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -4,21 +4,37 @@
44
* SPDX-License-Identifier: Apache-2.0
55
*/
66

7+
import { updateGlobalFetchTimeouts } from './fetch.js';
78
import { describe, it, expect, vi, beforeEach, afterAll } from 'vitest';
8-
import {
9-
isPrivateIp,
10-
isPrivateIpAsync,
11-
isAddressPrivate,
12-
fetchWithTimeout,
13-
} from './fetch.js';
149
import * as dnsPromises from 'node:dns/promises';
1510
import type { LookupAddress, LookupAllOptions } from 'node:dns';
1611
import ipaddr from 'ipaddr.js';
1712

13+
const { setGlobalDispatcher, Agent, ProxyAgent } = vi.hoisted(() => ({
14+
setGlobalDispatcher: vi.fn(),
15+
Agent: vi.fn(),
16+
ProxyAgent: vi.fn(),
17+
}));
18+
19+
vi.mock('undici', () => ({
20+
setGlobalDispatcher,
21+
Agent,
22+
ProxyAgent,
23+
}));
24+
1825
vi.mock('node:dns/promises', () => ({
1926
lookup: vi.fn(),
2027
}));
2128

29+
// Import after mocks are established
30+
const {
31+
isPrivateIp,
32+
isPrivateIpAsync,
33+
isAddressPrivate,
34+
fetchWithTimeout,
35+
setGlobalProxy,
36+
} = await import('./fetch.js');
37+
2238
// Mock global fetch
2339
const originalFetch = global.fetch;
2440
global.fetch = vi.fn();
@@ -183,4 +199,19 @@ describe('fetch utils', () => {
183199
);
184200
});
185201
});
202+
203+
describe('setGlobalProxy', () => {
204+
it('should configure ProxyAgent with experiment flag timeout', () => {
205+
const proxyUrl = 'http://proxy.example.com';
206+
updateGlobalFetchTimeouts(45773134);
207+
setGlobalProxy(proxyUrl);
208+
209+
expect(ProxyAgent).toHaveBeenCalledWith({
210+
uri: proxyUrl,
211+
headersTimeout: 45773134,
212+
bodyTimeout: 45773134,
213+
});
214+
expect(setGlobalDispatcher).toHaveBeenCalled();
215+
});
216+
});
186217
});

packages/core/src/utils/fetch.ts

Lines changed: 22 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -10,9 +10,6 @@ import { Agent, ProxyAgent, setGlobalDispatcher } from 'undici';
1010
import ipaddr from 'ipaddr.js';
1111
import { lookup } from 'node:dns/promises';
1212

13-
const DEFAULT_HEADERS_TIMEOUT = 300000; // 5 minutes
14-
const DEFAULT_BODY_TIMEOUT = 300000; // 5 minutes
15-
1613
export class FetchError extends Error {
1714
constructor(
1815
message: string,
@@ -31,14 +28,31 @@ export class PrivateIpError extends Error {
3128
}
3229
}
3330

31+
let defaultTimeout = 300000; // 5 minutes
32+
let currentProxy: string | undefined = undefined;
33+
3434
// Configure default global dispatcher with higher timeouts
3535
setGlobalDispatcher(
3636
new Agent({
37-
headersTimeout: DEFAULT_HEADERS_TIMEOUT,
38-
bodyTimeout: DEFAULT_BODY_TIMEOUT,
37+
headersTimeout: defaultTimeout,
38+
bodyTimeout: defaultTimeout,
3939
}),
4040
);
4141

42+
export function updateGlobalFetchTimeouts(timeoutMs: number) {
43+
defaultTimeout = timeoutMs;
44+
if (currentProxy) {
45+
setGlobalProxy(currentProxy);
46+
} else {
47+
setGlobalDispatcher(
48+
new Agent({
49+
headersTimeout: defaultTimeout,
50+
bodyTimeout: defaultTimeout,
51+
}),
52+
);
53+
}
54+
}
55+
4256
/**
4357
* Sanitizes a hostname by stripping IPv6 brackets if present.
4458
*/
@@ -191,11 +205,12 @@ export async function fetchWithTimeout(
191205
}
192206

193207
export function setGlobalProxy(proxy: string) {
208+
currentProxy = proxy;
194209
setGlobalDispatcher(
195210
new ProxyAgent({
196211
uri: proxy,
197-
headersTimeout: DEFAULT_HEADERS_TIMEOUT,
198-
bodyTimeout: DEFAULT_BODY_TIMEOUT,
212+
headersTimeout: defaultTimeout,
213+
bodyTimeout: defaultTimeout,
199214
}),
200215
);
201216
}

0 commit comments

Comments
 (0)