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

Refactor progress notification timeout handling in protocol

Browse source
This commit is contained in:
Tuomas 2025-02-19 12:23:16 -06:00
commit 28415d3917
2 changed files with 140 additions and 160 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -63,180 +63,160 @@ describe("protocol tests", () => {
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(),
describe("progress notification timeout behavior", () => {
beforeEach(() => {
jest.useFakeTimers();
});
afterEach(() => {
jest.useRealTimers();
});
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,
},
test("should reset timeout when progress notification is received", async () => {
await protocol.connect(transport);
const request = { method: "example", params: {} };
const mockSchema: ZodType<{ result: string }> = z.object({
result: z.string(),
});
}
// 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" },
const onProgressMock = jest.fn();
const requestPromise = protocol.request(request, mockSchema, {
timeout: 1000,
resetTimeoutOnProgress: true,
onprogress: onProgressMock,
});
}
// 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,
progress: 50,
total: 100,
},
});
}
// Verify progress handler was called
await Promise.resolve();
expect(onProgressMock).toHaveBeenNthCalledWith(i, {
progress: i * 25,
expect(onProgressMock).toHaveBeenCalledWith({
progress: 50,
total: 100,
});
}
jest.advanceTimersByTime(800);
if (transport.onmessage) {
transport.onmessage({
jsonrpc: "2.0",
id: 0,
result: { result: "success" },
});
}
await Promise.resolve();
await expect(requestPromise).resolves.toEqual({ result: "success" });
});
// Send success response
if (transport.onmessage) {
transport.onmessage({
jsonrpc: "2.0",
id: 0,
result: { result: "success" },
test("should respect maxTotalTimeout", async () => {
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: 150,
resetTimeoutOnProgress: true,
onprogress: onProgressMock,
});
}
await Promise.resolve();
await expect(requestPromise).resolves.toEqual({ result: "success" });
// First progress notification should work
jest.advanceTimersByTime(80);
if (transport.onmessage) {
transport.onmessage({
jsonrpc: "2.0",
method: "notifications/progress",
params: {
progressToken: 0,
progress: 50,
total: 100,
},
});
}
await Promise.resolve();
expect(onProgressMock).toHaveBeenCalledWith({
progress: 50,
total: 100,
});
jest.advanceTimersByTime(80);
if (transport.onmessage) {
transport.onmessage({
jsonrpc: "2.0",
method: "notifications/progress",
params: {
progressToken: 0,
progress: 75,
total: 100,
},
});
}
await expect(requestPromise).rejects.toThrow("Maximum total timeout exceeded");
expect(onProgressMock).toHaveBeenCalledTimes(1);
});
jest.useRealTimers();
test("should timeout if no progress received within timeout period", async () => {
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,
});
jest.advanceTimersByTime(101);
await expect(requestPromise).rejects.toThrow("Request timed out");
});
test("should handle multiple progress notifications correctly", async () => {
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++) {
jest.advanceTimersByTime(800);
if (transport.onmessage) {
transport.onmessage({
jsonrpc: "2.0",
method: "notifications/progress",
params: {
progressToken: 0,
progress: i * 25,
total: 100,
},
});
}
await Promise.resolve();
expect(onProgressMock).toHaveBeenNthCalledWith(i, {
progress: i * 25,
total: 100,
});
}
if (transport.onmessage) {
transport.onmessage({
jsonrpc: "2.0",
id: 0,
result: { result: "success" },
});
}
await Promise.resolve();
await expect(requestPromise).resolves.toEqual({ result: "success" });
});
});
});

24
src/shared/protocol.ts
View file

@ -73,8 +73,8 @@ export type RequestOptions = {
resetTimeoutOnProgress?: boolean;
/**
* Maximum total time (in milliseconds) to wait for a response, even if progress notifications are received.
* Only used when resetTimeoutOnProgress is true.
* Maximum total time (in milliseconds) to wait for a response.
* If exceeded, an McpError with code `RequestTimeout` will be raised, regardless of progress notifications.
* If not specified, there is no maximum total timeout.
*/
maxTotalTimeout?: number;
@ -190,19 +190,18 @@ export abstract class Protocol<
});
}
private _resetTimeout(messageId: number, cancel: (reason: unknown) => void): boolean {
private _resetTimeout(messageId: number): 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(
throw new McpError(
ErrorCode.RequestTimeout,
"Maximum total timeout exceeded",
{ maxTotalTimeout: info.maxTotalTimeout, totalElapsed }
));
return false;
);
}
clearTimeout(info.timeoutId);
@ -360,7 +359,12 @@ export abstract class Protocol<
const responseHandler = this._responseHandlers.get(messageId);
if (this._timeoutInfo.has(messageId) && responseHandler) {
if (!this._resetTimeout(messageId, (reason) => responseHandler(reason as Error))) {
try {
if (!this._resetTimeout(messageId)) {
return;
}
} catch (error) {
responseHandler(error as Error);
return;
}
}
@ -518,11 +522,7 @@ export abstract class Protocol<
{ timeout }
));
if (options?.resetTimeoutOnProgress) {
this._setupTimeout(messageId, timeout, options.maxTotalTimeout, timeoutHandler);
} else {
this._setupTimeout(messageId, timeout, undefined, timeoutHandler);
}
this._setupTimeout(messageId, timeout, options?.maxTotalTimeout, timeoutHandler);
this._transport.send(jsonrpcRequest).catch((error) => {
this._cleanupTimeout(messageId);