fix(protocol): ignore late progress for recently completed tokens

.changeset/ignore-late-progress-notifications.md

@@ -0,0 +1,5 @@
+---
+'@modelcontextprotocol/sdk': patch
+---
+
+Ignore late progress notifications for recently finished requests while preserving errors for progress tokens that were never registered.

src/shared/protocol.ts

@@ -104,6 +104,7 @@ export type ProtocolOptions = {
  * The default request timeout, in miliseconds.
  */
 export const DEFAULT_REQUEST_TIMEOUT_MSEC = 60000;
+const MAX_RECENTLY_COMPLETED_PROGRESS_TOKENS = 64;
 
 /**
  * Options that can be given per request.
@@ -328,6 +329,7 @@ export abstract class Protocol<SendRequestT extends Request, SendNotificationT e
     private _notificationHandlers: Map<string, (notification: JSONRPCNotification) => Promise<void>> = new Map();
     private _responseHandlers: Map<number, (response: JSONRPCResultResponse | Error) => void> = new Map();
     private _progressHandlers: Map<number, ProgressCallback> = new Map();
+    private _recentlyCompletedProgressTokens = new Set<number>();
     private _timeoutInfo: Map<number, TimeoutInfo> = new Map();
     private _pendingDebouncedNotifications = new Set<string>();
 
@@ -645,6 +647,7 @@ export abstract class Protocol<SendRequestT extends Request, SendNotificationT e
         const responseHandlers = this._responseHandlers;
         this._responseHandlers = new Map();
         this._progressHandlers.clear();
+        this._recentlyCompletedProgressTokens.clear();
         this._taskProgressTokens.clear();
         this._pendingDebouncedNotifications.clear();
 
@@ -859,6 +862,12 @@ export abstract class Protocol<SendRequestT extends Request, SendNotificationT e
 
         const handler = this._progressHandlers.get(messageId);
         if (!handler) {
+            // A request can finish before progress notifications that were already in transit
+            // are dispatched. Drop progress for a recently finished request, while preserving
+            // the error for tokens that were never registered.
+            if (this._recentlyCompletedProgressTokens.has(messageId)) {
+                return;
+            }
             this._onerror(new Error(`Received a progress notification for an unknown token: ${JSON.stringify(notification)}`));
             return;
         }
@@ -872,7 +881,7 @@ export abstract class Protocol<SendRequestT extends Request, SendNotificationT e
             } catch (error) {
                 // Clean up if maxTotalTimeout was exceeded
                 this._responseHandlers.delete(messageId);
-                this._progressHandlers.delete(messageId);
+                this._cleanupProgressHandler(messageId);
                 this._cleanupTimeout(messageId);
                 responseHandler(error as Error);
                 return;
@@ -882,6 +891,23 @@ export abstract class Protocol<SendRequestT extends Request, SendNotificationT e
         handler(params);
     }
 
+    private _cleanupProgressHandler(progressToken: number): void {
+        if (this._progressHandlers.delete(progressToken)) {
+            this._rememberCompletedProgressToken(progressToken);
+        }
+    }
+
+    private _rememberCompletedProgressToken(progressToken: number): void {
+        this._recentlyCompletedProgressTokens.add(progressToken);
+
+        if (this._recentlyCompletedProgressTokens.size > MAX_RECENTLY_COMPLETED_PROGRESS_TOKENS) {
+            const oldestProgressToken = this._recentlyCompletedProgressTokens.values().next().value;
+            if (oldestProgressToken !== undefined) {
+                this._recentlyCompletedProgressTokens.delete(oldestProgressToken);
+            }
+        }
+    }
+
     private _onresponse(response: JSONRPCResponse | JSONRPCErrorResponse): void {
         const messageId = Number(response.id);
 
@@ -921,7 +947,7 @@ export abstract class Protocol<SendRequestT extends Request, SendNotificationT e
         }
 
         if (!isTaskResponse) {
-            this._progressHandlers.delete(messageId);
+            this._cleanupProgressHandler(messageId);
         }
 
         if (isJSONRPCResultResponse(response)) {
@@ -1164,7 +1190,7 @@ export abstract class Protocol<SendRequestT extends Request, SendNotificationT e
 
             const cancel = (reason: unknown) => {
                 this._responseHandlers.delete(messageId);
-                this._progressHandlers.delete(messageId);
+                this._cleanupProgressHandler(messageId);
                 this._cleanupTimeout(messageId);
 
                 this._transport
@@ -1477,7 +1503,7 @@ export abstract class Protocol<SendRequestT extends Request, SendNotificationT e
     private _cleanupTaskProgressHandler(taskId: string): void {
         const progressToken = this._taskProgressTokens.get(taskId);
         if (progressToken !== undefined) {
-            this._progressHandlers.delete(progressToken);
+            this._cleanupProgressHandler(progressToken);
             this._taskProgressTokens.delete(taskId);
         }
     }

test/shared/protocol.test.ts

@@ -14,7 +14,7 @@ import {
     type Notification,
     type Result
 } from '../../src/types.js';
-import { Protocol, mergeCapabilities } from '../../src/shared/protocol.js';
+import { Protocol, mergeCapabilities, type ProgressCallback, type RequestOptions } from '../../src/shared/protocol.js';
 import { Transport, TransportSendOptions } from '../../src/shared/transport.js';
 import { TaskStore, TaskMessageQueue, QueuedMessage, QueuedNotification, QueuedRequest } from '../../src/experimental/tasks/interfaces.js';
 import { MockInstance, vi } from 'vitest';
@@ -351,6 +351,151 @@ describe('protocol tests', () => {
         });
     });
 
+    describe('late progress notifications', () => {
+        const resultSchema = z.object({
+            result: z.string()
+        });
+
+        const startProgressRequest = (onprogress: ProgressCallback = vi.fn(), options: Omit<RequestOptions, 'onprogress'> = {}) => {
+            const requestPromise = protocol.request({ method: 'example', params: {} }, resultSchema, { ...options, onprogress });
+            const sentRequest = sendSpy.mock.calls.at(-1)![0] as {
+                id: number;
+                params: { _meta: { progressToken: number } };
+            };
+            return { requestPromise, sentRequest };
+        };
+
+        const completeRequest = async (requestPromise: Promise<{ result: string }>, requestId: number) => {
+            transport.onmessage?.({
+                jsonrpc: '2.0',
+                id: requestId,
+                result: { result: 'success' }
+            });
+            await expect(requestPromise).resolves.toEqual({ result: 'success' });
+        };
+
+        const sendProgress = async (progressToken: number, progress: number) => {
+            transport.onmessage?.({
+                jsonrpc: '2.0',
+                method: 'notifications/progress',
+                params: {
+                    progressToken,
+                    progress,
+                    total: 100
+                }
+            });
+            await Promise.resolve();
+        };
+
+        test('should deliver progress received before the response', async () => {
+            const onProgressMock = vi.fn();
+            const onErrorMock = vi.fn();
+            protocol.onerror = onErrorMock;
+            await protocol.connect(transport);
+
+            const { requestPromise, sentRequest } = startProgressRequest(onProgressMock);
+            const progressToken = sentRequest.params._meta.progressToken;
+
+            await sendProgress(progressToken, 50);
+            await completeRequest(requestPromise, sentRequest.id);
+            expect(onProgressMock).toHaveBeenCalledWith({
+                progress: 50,
+                total: 100
+            });
+            expect(onErrorMock).not.toHaveBeenCalled();
+        });
+
+        test('should ignore progress received after the response for the same token', async () => {
+            const onProgressMock = vi.fn();
+            const onErrorMock = vi.fn();
+            protocol.onerror = onErrorMock;
+            await protocol.connect(transport);
+
+            const { requestPromise, sentRequest } = startProgressRequest(onProgressMock);
+            const progressToken = sentRequest.params._meta.progressToken;
+
+            await completeRequest(requestPromise, sentRequest.id);
+            await sendProgress(progressToken, 100);
+
+            expect(onProgressMock).not.toHaveBeenCalled();
+            expect(onErrorMock).not.toHaveBeenCalled();
+        });
+
+        test('should ignore progress received after caller abort', async () => {
+            const onErrorMock = vi.fn();
+            protocol.onerror = onErrorMock;
+            await protocol.connect(transport);
+
+            const abortController = new AbortController();
+            const { requestPromise, sentRequest } = startProgressRequest(vi.fn(), { signal: abortController.signal });
+            const requestRejection = requestPromise.catch(() => undefined);
+
+            abortController.abort(new Error('caller aborted'));
+            await requestRejection;
+            await sendProgress(sentRequest.params._meta.progressToken, 50);
+
+            expect(onErrorMock).not.toHaveBeenCalled();
+        });
+
+        test('should ignore progress received after request timeout', async () => {
+            vi.useFakeTimers();
+            try {
+                const onErrorMock = vi.fn();
+                protocol.onerror = onErrorMock;
+                await protocol.connect(transport);
+
+                const { requestPromise, sentRequest } = startProgressRequest(vi.fn(), { timeout: 10 });
+                const requestRejection = requestPromise.catch(() => undefined);
+
+                await vi.advanceTimersByTimeAsync(20);
+                await requestRejection;
+                await sendProgress(sentRequest.params._meta.progressToken, 50);
+
+                expect(onErrorMock).not.toHaveBeenCalled();
+            } finally {
+                vi.useRealTimers();
+            }
+        });
+
+        test('should report progress for a token that was never registered', async () => {
+            const onErrorMock = vi.fn();
+            protocol.onerror = onErrorMock;
+            await protocol.connect(transport);
+
+            const requestPromise = protocol.request({ method: 'example', params: {} }, resultSchema);
+            const sentRequest = sendSpy.mock.calls.at(-1)![0] as { id: number };
+
+            await completeRequest(requestPromise, sentRequest.id);
+            await sendProgress(sentRequest.id, 50);
+            expect(onErrorMock).toHaveBeenCalledOnce();
+            const error = onErrorMock.mock.calls[0][0] as Error;
+            expect(error).toBeInstanceOf(Error);
+            expect(error.message).toContain('Received a progress notification for an unknown token');
+            expect(error.message).toContain(`"progressToken":${sentRequest.id}`);
+        });
+
+        test('should evict the oldest recently completed progress token', async () => {
+            const onErrorMock = vi.fn();
+            protocol.onerror = onErrorMock;
+            await protocol.connect(transport);
+
+            const retainedTokenCount = 64;
+            const completedProgressTokens: number[] = [];
+
+            for (let i = 0; i <= retainedTokenCount; i++) {
+                const { requestPromise, sentRequest } = startProgressRequest();
+                completedProgressTokens.push(sentRequest.params._meta.progressToken);
+                await completeRequest(requestPromise, sentRequest.id);
+            }
+
+            await sendProgress(completedProgressTokens[0], 100);
+            expect(onErrorMock).toHaveBeenCalledOnce();
+
+            await sendProgress(completedProgressTokens.at(-1)!, 100);
+            expect(onErrorMock).toHaveBeenCalledOnce();
+        });
+    });
+
     describe('progress notification timeout behavior', () => {
         beforeEach(() => {
             vi.useFakeTimers();
@@ -445,6 +590,8 @@ describe('protocol tests', () => {
                 result: z.string()
             });
             const onProgressMock = vi.fn();
+            const onErrorMock = vi.fn();
+            protocol.onerror = onErrorMock;
             const requestPromise = protocol.request(request, mockSchema, {
                 timeout: 1000,
                 maxTotalTimeout: 150,
@@ -484,6 +631,22 @@ describe('protocol tests', () => {
             }
             await expect(requestPromise).rejects.toThrow('Maximum total timeout exceeded');
             expect(onProgressMock).toHaveBeenCalledTimes(1);
+
+            // A progress notification already in flight after maxTotalTimeout cleanup
+            // should be ignored rather than reported as an unknown token.
+            if (transport.onmessage) {
+                transport.onmessage({
+                    jsonrpc: '2.0',
+                    method: 'notifications/progress',
+                    params: {
+                        progressToken: 0,
+                        progress: 100,
+                        total: 100
+                    }
+                });
+            }
+            await Promise.resolve();
+            expect(onErrorMock).not.toHaveBeenCalled();
         });
 
         test('should timeout if no progress received within timeout period', async () => {
@@ -2476,6 +2639,8 @@ describe('Progress notification support for tasks', () => {
 
         const transport = new MockTransport();
         const sendSpy = vi.spyOn(transport, 'send');
+        const onErrorMock = vi.fn();
+        protocol.onerror = onErrorMock;
         await protocol.connect(transport);
 
         // Set up a request handler that will complete the task
@@ -2597,6 +2762,7 @@ describe('Progress notification support for tasks', () => {
 
         // Progress callback should NOT be invoked after task completion
         expect(progressCallback).not.toHaveBeenCalled();
+        expect(onErrorMock).not.toHaveBeenCalled();
     });
 
     it('should stop progress notifications when task reaches terminal status (failed)', async () => {