refactor: Add Effective Timeout for MCP Fetch (#9585)

This commit is contained in:
Danny Avila 2025-09-11 19:09:13 -04:00 committed by GitHub
parent 51f2d43fed
commit d75fb76338
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -65,6 +65,7 @@ function isStreamableHTTPOptions(options: t.MCPOptions): options is t.Streamable
} }
const FIVE_MINUTES = 5 * 60 * 1000; const FIVE_MINUTES = 5 * 60 * 1000;
const DEFAULT_TIMEOUT = 60000;
interface MCPConnectionParams { interface MCPConnectionParams {
serverName: string; serverName: string;
@ -145,19 +146,22 @@ export class MCPConnection extends EventEmitter {
* This helps prevent memory leaks by only passing necessary dependencies. * This helps prevent memory leaks by only passing necessary dependencies.
* *
* @param getHeaders Function to retrieve request headers * @param getHeaders Function to retrieve request headers
* @param timeout Timeout value for the agent (in milliseconds)
* @returns A fetch function that merges headers appropriately * @returns A fetch function that merges headers appropriately
*/ */
private createFetchFunction( private createFetchFunction(
getHeaders: () => Record<string, string> | null | undefined, getHeaders: () => Record<string, string> | null | undefined,
timeout?: number,
): (input: UndiciRequestInfo, init?: UndiciRequestInit) => Promise<UndiciResponse> { ): (input: UndiciRequestInfo, init?: UndiciRequestInit) => Promise<UndiciResponse> {
return function customFetch( return function customFetch(
input: UndiciRequestInfo, input: UndiciRequestInfo,
init?: UndiciRequestInit, init?: UndiciRequestInit,
): Promise<UndiciResponse> { ): Promise<UndiciResponse> {
const requestHeaders = getHeaders(); const requestHeaders = getHeaders();
const effectiveTimeout = timeout || DEFAULT_TIMEOUT;
const agent = new Agent({ const agent = new Agent({
bodyTimeout: 0, bodyTimeout: effectiveTimeout,
headersTimeout: 0, headersTimeout: effectiveTimeout,
}); });
if (!requestHeaders) { if (!requestHeaders) {
return undiciFetch(input, { ...init, dispatcher: agent }); return undiciFetch(input, { ...init, dispatcher: agent });
@ -243,6 +247,7 @@ export class MCPConnection extends EventEmitter {
headers['Authorization'] = `Bearer ${this.oauthTokens.access_token}`; headers['Authorization'] = `Bearer ${this.oauthTokens.access_token}`;
} }
const timeoutValue = this.timeout || DEFAULT_TIMEOUT;
const transport = new SSEClientTransport(url, { const transport = new SSEClientTransport(url, {
requestInit: { requestInit: {
headers, headers,
@ -252,8 +257,8 @@ export class MCPConnection extends EventEmitter {
fetch: (url, init) => { fetch: (url, init) => {
const fetchHeaders = new Headers(Object.assign({}, init?.headers, headers)); const fetchHeaders = new Headers(Object.assign({}, init?.headers, headers));
const agent = new Agent({ const agent = new Agent({
bodyTimeout: 0, bodyTimeout: timeoutValue,
headersTimeout: 0, headersTimeout: timeoutValue,
}); });
return undiciFetch(url, { return undiciFetch(url, {
...init, ...init,
@ -264,6 +269,7 @@ export class MCPConnection extends EventEmitter {
}, },
fetch: this.createFetchFunction( fetch: this.createFetchFunction(
this.getRequestHeaders.bind(this), this.getRequestHeaders.bind(this),
this.timeout,
) as unknown as FetchLike, ) as unknown as FetchLike,
}); });
@ -304,6 +310,7 @@ export class MCPConnection extends EventEmitter {
}, },
fetch: this.createFetchFunction( fetch: this.createFetchFunction(
this.getRequestHeaders.bind(this), this.getRequestHeaders.bind(this),
this.timeout,
) as unknown as FetchLike, ) as unknown as FetchLike,
}); });