modelcontextprotocol/typescript-sdk
Watch
1
0
Fork
You've already forked typescript-sdk
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 1

Add timeout reset on progress notifications

Browse source
This commit is contained in:
Tuomas 2025-02-17 15:57:02 -06:00
commit 7f6c0465f5
2 changed files with 292 additions and 50 deletions
Download patch file Download diff file Expand all files Collapse all files

176
src/shared/protocol.test.ts
View file

@ -62,6 +62,182 @@ describe("protocol tests", () => {
await transport.close();
expect(oncloseMock).toHaveBeenCalled();
});
test("should reset timeout when progress notification is received", async () => {
jest.useFakeTimers();
await protocol.connect(transport);
const request = { method: "example", params: {} };
const mockSchema: ZodType<{ result: string }> = z.object({
result: z.string(),
});
const onProgressMock = jest.fn();
const requestPromise = protocol.request(request, mockSchema, {
timeout: 1000, // Increased timeout for more reliable testing
resetTimeoutOnProgress: true,
onprogress: onProgressMock,
});
// Advance time close to timeout
jest.advanceTimersByTime(800);
// Send progress notification
if (transport.onmessage) {
transport.onmessage({
jsonrpc: "2.0",
method: "notifications/progress",
params: {
progressToken: 0,
progress: 50,
total: 100,
},
});
}
// Run all pending promises to ensure progress handler is called
await Promise.resolve();
// Verify progress handler was called
expect(onProgressMock).toHaveBeenCalledWith({
progress: 50,
total: 100,
});
// Send success response
if (transport.onmessage) {
transport.onmessage({
jsonrpc: "2.0",
id: 0,
result: { result: "success" },
});
}
// Run all pending promises
await Promise.resolve();
await expect(requestPromise).resolves.toEqual({ result: "success" });
jest.useRealTimers();
});
test("should respect maxTotalTimeout", async () => {
jest.useFakeTimers();
await protocol.connect(transport);
const request = { method: "example", params: {} };
const mockSchema: ZodType<{ result: string }> = z.object({
result: z.string(),
});
const onProgressMock = jest.fn();
const requestPromise = protocol.request(request, mockSchema, {
timeout: 1000,
maxTotalTimeout: 100,
resetTimeoutOnProgress: true,
onprogress: onProgressMock,
});
// Advance time beyond maxTotalTimeout
jest.advanceTimersByTime(150);
// Send progress notification after maxTotalTimeout
if (transport.onmessage) {
transport.onmessage({
jsonrpc: "2.0",
method: "notifications/progress",
params: {
progressToken: 0,
progress: 50,
total: 100,
},
});
}
await expect(requestPromise).rejects.toThrow("Maximum total timeout exceeded");
expect(onProgressMock).not.toHaveBeenCalled();
jest.useRealTimers();
});
test("should timeout if no progress received within timeout period", async () => {
jest.useFakeTimers();
await protocol.connect(transport);
const request = { method: "example", params: {} };
const mockSchema: ZodType<{ result: string }> = z.object({
result: z.string(),
});
const requestPromise = protocol.request(request, mockSchema, {
timeout: 100,
resetTimeoutOnProgress: true,
});
// Advance time beyond timeout
jest.advanceTimersByTime(101);
await expect(requestPromise).rejects.toThrow("Request timed out");
jest.useRealTimers();
});
test("should handle multiple progress notifications correctly", async () => {
jest.useFakeTimers();
await protocol.connect(transport);
const request = { method: "example", params: {} };
const mockSchema: ZodType<{ result: string }> = z.object({
result: z.string(),
});
const onProgressMock = jest.fn();
const requestPromise = protocol.request(request, mockSchema, {
timeout: 1000,
resetTimeoutOnProgress: true,
onprogress: onProgressMock,
});
// Simulate multiple progress updates
for (let i = 1; i <= 3; i++) {
// Advance close to timeout
jest.advanceTimersByTime(800);
// Send progress notification
if (transport.onmessage) {
transport.onmessage({
jsonrpc: "2.0",
method: "notifications/progress",
params: {
progressToken: 0,
progress: i * 25,
total: 100,
},
});
}
// Verify progress handler was called
await Promise.resolve();
expect(onProgressMock).toHaveBeenNthCalledWith(i, {
progress: i * 25,
total: 100,
});
}
// Send success response
if (transport.onmessage) {
transport.onmessage({
jsonrpc: "2.0",
id: 0,
result: { result: "success" },
});
}
await Promise.resolve();
await expect(requestPromise).resolves.toEqual({ result: "success" });
jest.useRealTimers();
});
});
describe("mergeCapabilities", () => {

166
src/shared/protocol.ts
View file

@ -64,6 +64,20 @@ export type RequestOptions = {
* If not specified, `DEFAULT_REQUEST_TIMEOUT_MSEC` will be used as the timeout.
*/
timeout?: number;
/**
* If true, receiving a progress notification will reset the request timeout.
* This is useful for long-running operations that send periodic progress updates.
* Default: false
*/
resetTimeoutOnProgress?: boolean;
/**
* Maximum total time (in milliseconds) to wait for a response, even if progress notifications are received.
* Only used when resetTimeoutOnProgress is true.
* If not specified, there is no maximum total timeout.
*/
maxTotalTimeout?: number;
};
/**
@ -76,6 +90,17 @@ export type RequestHandlerExtra = {
signal: AbortSignal;
};
/**
* Information about a request's timeout state
*/
type TimeoutInfo = {
timeoutId: ReturnType<typeof setTimeout>;
startTime: number;
timeout: number;
maxTotalTimeout?: number;
onTimeout: () => void;
};
/**
* Implements MCP protocol framing on top of a pluggable transport, including
* features like request/response linking, notifications, and progress.
@ -105,6 +130,7 @@ export abstract class Protocol<
(response: JSONRPCResponse | Error) => void
> = new Map();
private _progressHandlers: Map<number, ProgressCallback> = new Map();
private _timeoutInfo: Map<number, TimeoutInfo> = new Map();
/**
* Callback for when the connection is closed for any reason.
@ -149,6 +175,49 @@ export abstract class Protocol<
);
}
private _setupTimeout(
messageId: number,
timeout: number,
maxTotalTimeout: number | undefined,
onTimeout: () => void
) {
this._timeoutInfo.set(messageId, {
timeoutId: setTimeout(onTimeout, timeout),
startTime: Date.now(),
timeout,
maxTotalTimeout,
onTimeout
});
}
private _resetTimeout(messageId: number, cancel: (reason: unknown) => void): boolean {
const info = this._timeoutInfo.get(messageId);
if (!info) return false;
const totalElapsed = Date.now() - info.startTime;
if (info.maxTotalTimeout && totalElapsed >= info.maxTotalTimeout) {
this._timeoutInfo.delete(messageId);
cancel(new McpError(
ErrorCode.RequestTimeout,
"Maximum total timeout exceeded",
{ maxTotalTimeout: info.maxTotalTimeout, totalElapsed }
));
return false;
}
clearTimeout(info.timeoutId);
info.timeoutId = setTimeout(info.onTimeout, info.timeout);
return true;
}
private _cleanupTimeout(messageId: number) {
const info = this._timeoutInfo.get(messageId);
if (info) {
clearTimeout(info.timeoutId);
this._timeoutInfo.delete(messageId);
}
}
/**
* Attaches to the given transport, starts it, and starts listening for messages.
*
@ -281,22 +350,27 @@ export abstract class Protocol<
private _onprogress(notification: ProgressNotification): void {
const { progressToken, ...params } = notification.params;
const handler = this._progressHandlers.get(Number(progressToken));
if (handler === undefined) {
this._onerror(
new Error(
`Received a progress notification for an unknown token: ${JSON.stringify(notification)}`,
),
);
const messageId = Number(progressToken);
const handler = this._progressHandlers.get(messageId);
if (!handler) {
this._onerror(new Error(`Received a progress notification for an unknown token: ${JSON.stringify(notification)}`));
return;
}
const responseHandler = this._responseHandlers.get(messageId);
if (this._timeoutInfo.has(messageId) && responseHandler) {
if (!this._resetTimeout(messageId, (reason) => responseHandler(reason as Error))) {
return;
}
}
handler(params);
}
private _onresponse(response: JSONRPCResponse | JSONRPCError): void {
const messageId = response.id;
const handler = this._responseHandlers.get(Number(messageId));
const messageId = Number(response.id);
const handler = this._responseHandlers.get(messageId);
if (handler === undefined) {
this._onerror(
new Error(
@ -306,8 +380,10 @@ export abstract class Protocol<
return;
}
this._responseHandlers.delete(Number(messageId));
this._progressHandlers.delete(Number(messageId));
this._responseHandlers.delete(messageId);
this._progressHandlers.delete(messageId);
this._cleanupTimeout(messageId);
if ("result" in response) {
handler(response);
} else {
@ -393,32 +469,10 @@ export abstract class Protocol<
};
}
let timeoutId: ReturnType<typeof setTimeout> | undefined = undefined;
this._responseHandlers.set(messageId, (response) => {
if (timeoutId !== undefined) {
clearTimeout(timeoutId);
}
if (options?.signal?.aborted) {
return;
}
if (response instanceof Error) {
return reject(response);
}
try {
const result = resultSchema.parse(response.result);
resolve(result);
} catch (error) {
reject(error);
}
});
const cancel = (reason: unknown) => {
this._responseHandlers.delete(messageId);
this._progressHandlers.delete(messageId);
this._cleanupTimeout(messageId);
this._transport
?.send({
@ -436,30 +490,42 @@ export abstract class Protocol<
reject(reason);
};
options?.signal?.addEventListener("abort", () => {
if (timeoutId !== undefined) {
clearTimeout(timeoutId);
this._responseHandlers.set(messageId, (response) => {
if (options?.signal?.aborted) {
return;
}
if (response instanceof Error) {
return reject(response);
}
try {
const result = resultSchema.parse(response.result);
resolve(result);
} catch (error) {
reject(error);
}
});
options?.signal?.addEventListener("abort", () => {
cancel(options?.signal?.reason);
});
const timeout = options?.timeout ?? DEFAULT_REQUEST_TIMEOUT_MSEC;
timeoutId = setTimeout(
() =>
cancel(
new McpError(ErrorCode.RequestTimeout, "Request timed out", {
timeout,
}),
),
timeout,
);
const timeoutHandler = () => cancel(new McpError(
ErrorCode.RequestTimeout,
"Request timed out",
{ timeout }
));
if (options?.resetTimeoutOnProgress) {
this._setupTimeout(messageId, timeout, options.maxTotalTimeout, timeoutHandler);
} else {
this._setupTimeout(messageId, timeout, undefined, timeoutHandler);
}
this._transport.send(jsonrpcRequest).catch((error) => {
if (timeoutId !== undefined) {
clearTimeout(timeoutId);
}
this._cleanupTimeout(messageId);
reject(error);
});
});