progress manager split

KKonstantinov · KKonstantinov · commit be7ab37b8532 · 2026-02-03T12:24:37.000+02:00
diff --git a/packages/core/src/index.ts b/packages/core/src/index.ts
@@ -2,6 +2,7 @@ export * from './auth/errors.js';
 export * from './shared/auth.js';
 export * from './shared/authUtils.js';
 export * from './shared/metadataUtils.js';
+export * from './shared/progressManager.js';
 export * from './shared/protocol.js';
 export * from './shared/responseMessage.js';
 export * from './shared/stdio.js';
diff --git a/packages/core/src/shared/progressManager.ts b/packages/core/src/shared/progressManager.ts
@@ -0,0 +1,172 @@
+/**
+ * Progress Manager
+ *
+ * Manages progress tracking for the Protocol class.
+ * Extracted from Protocol to follow Single Responsibility Principle.
+ */
+
+import type { Progress, ProgressNotification } from '../types/types.js';
+
+/**
+ * Callback for progress notifications.
+ */
+export type ProgressCallback = (progress: Progress) => void;
+
+/**
+ * Manages progress tracking for requests.
+ *
+ * This class handles registration, lookup, and invocation of progress callbacks,
+ * as well as task-to-progress-token associations for long-running task operations.
+ *
+ * @example
+ * ```typescript
+ * const progressManager = new ProgressManager();
+ *
+ * // Register a progress handler for a request
+ * progressManager.registerHandler(messageId, (progress) => {
+ *   console.log(`Progress: ${progress.progress}/${progress.total}`);
+ * });
+ *
+ * // Handle incoming progress notification
+ * progressManager.handleProgress(notification);
+ *
+ * // Clean up when done
+ * progressManager.removeHandler(messageId);
+ * ```
+ */
+export class ProgressManager {
+    /**
+     * Maps message IDs to progress callbacks.
+     */
+    #progressHandlers: Map<number, ProgressCallback> = new Map();
+
+    /**
+     * Maps task IDs to progress tokens to keep handlers alive after CreateTaskResult.
+     */
+    #taskProgressTokens: Map<string, number> = new Map();
+
+    /**
+     * Registers a progress callback for a message.
+     *
+     * @param messageId - The message ID (used as progress token)
+     * @param callback - The callback to invoke when progress is received
+     */
+    registerHandler(messageId: number, callback: ProgressCallback): void {
+        this.#progressHandlers.set(messageId, callback);
+    }
+
+    /**
+     * Gets the progress callback for a message.
+     *
+     * @param messageId - The message ID
+     * @returns The progress callback or undefined if not registered
+     */
+    getHandler(messageId: number): ProgressCallback | undefined {
+        return this.#progressHandlers.get(messageId);
+    }
+
+    /**
+     * Removes the progress callback for a message.
+     *
+     * @param messageId - The message ID
+     */
+    removeHandler(messageId: number): void {
+        this.#progressHandlers.delete(messageId);
+    }
+
+    /**
+     * Checks if a progress handler exists for the given message ID.
+     *
+     * @param messageId - The message ID
+     * @returns true if a handler is registered, false otherwise
+     */
+    hasHandler(messageId: number): boolean {
+        return this.#progressHandlers.has(messageId);
+    }
+
+    /**
+     * Handles an incoming progress notification by invoking the registered callback.
+     * Returns true if the progress was handled, false if no handler was found.
+     *
+     * @param notification - The progress notification
+     * @returns true if handled, false otherwise
+     */
+    handleProgress(notification: ProgressNotification): boolean {
+        const token = notification.params.progressToken;
+        if (typeof token !== 'number') {
+            // Token must be a number for our internal tracking
+            return false;
+        }
+
+        const callback = this.#progressHandlers.get(token);
+        if (callback) {
+            callback({
+                progress: notification.params.progress,
+                total: notification.params.total,
+                message: notification.params.message
+            });
+            return true;
+        }
+
+        return false;
+    }
+
+    /**
+     * Links a task ID to a progress token.
+     * This keeps the progress handler alive after CreateTaskResult is returned,
+     * allowing progress notifications to continue for long-running tasks.
+     *
+     * @param taskId - The task identifier
+     * @param progressToken - The progress token (message ID)
+     */
+    linkTaskToProgressToken(taskId: string, progressToken: number): void {
+        this.#taskProgressTokens.set(taskId, progressToken);
+    }
+
+    /**
+     * Gets the progress token associated with a task.
+     *
+     * @param taskId - The task identifier
+     * @returns The progress token or undefined if not linked
+     */
+    getTaskProgressToken(taskId: string): number | undefined {
+        return this.#taskProgressTokens.get(taskId);
+    }
+
+    /**
+     * Cleans up the progress handler associated with a task.
+     * Should be called when a task reaches a terminal status.
+     *
+     * @param taskId - The task identifier
+     */
+    cleanupTaskProgressHandler(taskId: string): void {
+        const progressToken = this.#taskProgressTokens.get(taskId);
+        if (progressToken !== undefined) {
+            this.#progressHandlers.delete(progressToken);
+            this.#taskProgressTokens.delete(taskId);
+        }
+    }
+
+    /**
+     * Clears all progress handlers and task progress tokens.
+     * Typically called when the connection is closed.
+     */
+    clear(): void {
+        this.#progressHandlers.clear();
+        this.#taskProgressTokens.clear();
+    }
+
+    /**
+     * Gets the number of active progress handlers.
+     */
+    get handlerCount(): number {
+        return this.#progressHandlers.size;
+    }
+
+    /**
+     * Gets the number of active task-to-progress-token links.
+     */
+    get taskTokenCount(): number {
+        return this.#taskProgressTokens.size;
+    }
+}
diff --git a/packages/core/src/shared/protocol.ts b/packages/core/src/shared/protocol.ts
@@ -16,7 +16,6 @@ import type {
     Notification,
     NotificationMethod,
     NotificationTypeMap,
-    Progress,
     ProgressNotification,
     RelatedTaskMetadata,
     Request,
@@ -51,14 +50,11 @@ import {
 import type { AnySchema, SchemaOutput } from '../util/zodCompat.js';
 import { safeParse } from '../util/zodCompat.js';
 import { parseWithCompat } from '../util/zodJsonSchemaCompat.js';
+import type { ProgressCallback } from './progressManager.js';
+import { ProgressManager } from './progressManager.js';
 import type { ResponseMessage } from './responseMessage.js';
 import type { Transport, TransportSendOptions } from './transport.js';
 
-/**
- * Callback for progress notifications.
- */
-export type ProgressCallback = (progress: Progress) => void;
-
 /**
  * Additional initialization options.
  */
@@ -330,13 +326,10 @@ export abstract class Protocol<SendRequestT extends Request, SendNotificationT e
     private _requestHandlerAbortControllers: Map<RequestId, AbortController> = new Map();
     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 _progressManager: ProgressManager = new ProgressManager();
     private _timeoutInfo: Map<number, TimeoutInfo> = new Map();
     private _pendingDebouncedNotifications = new Set<string>();
 
-    // Maps task IDs to progress tokens to keep handlers alive after CreateTaskResult
-    private _taskProgressTokens: Map<string, number> = new Map();
-
     private _taskStore?: TaskStore;
     private _taskMessageQueue?: TaskMessageQueue;
 
@@ -639,8 +632,7 @@ export abstract class Protocol<SendRequestT extends Request, SendNotificationT e
     private _onclose(): void {
         const responseHandlers = this._responseHandlers;
         this._responseHandlers = new Map();
-        this._progressHandlers.clear();
-        this._taskProgressTokens.clear();
+        this._progressManager.clear();
         this._pendingDebouncedNotifications.clear();
 
         const error = McpError.fromError(ErrorCode.ConnectionClosed, 'Connection closed');
@@ -826,11 +818,10 @@ export abstract class Protocol<SendRequestT extends Request, SendNotificationT e
     }
 
     private _onprogress(notification: ProgressNotification): void {
-        const { progressToken, ...params } = notification.params;
+        const progressToken = notification.params.progressToken;
         const messageId = Number(progressToken);
 
-        const handler = this._progressHandlers.get(messageId);
-        if (!handler) {
+        if (!this._progressManager.hasHandler(messageId)) {
             this._onerror(new Error(`Received a progress notification for an unknown token: ${JSON.stringify(notification)}`));
             return;
         }
@@ -844,14 +835,14 @@ 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._progressManager.removeHandler(messageId);
                 this._cleanupTimeout(messageId);
                 responseHandler(error as Error);
                 return;
             }
         }
 
-        handler(params);
+        this._progressManager.handleProgress(notification);
     }
 
     private _onresponse(response: JSONRPCResponse | JSONRPCErrorResponse): void {
@@ -887,13 +878,13 @@ export abstract class Protocol<SendRequestT extends Request, SendNotificationT e
                 const task = result.task as Record<string, unknown>;
                 if (typeof task.taskId === 'string') {
                     isTaskResponse = true;
-                    this._taskProgressTokens.set(task.taskId, messageId);
+                    this._progressManager.linkTaskToProgressToken(task.taskId, messageId);
                 }
             }
         }
 
         if (!isTaskResponse) {
-            this._progressHandlers.delete(messageId);
+            this._progressManager.removeHandler(messageId);
         }
 
         if (isJSONRPCResultResponse(response)) {
@@ -1116,7 +1107,7 @@ export abstract class Protocol<SendRequestT extends Request, SendNotificationT e
             };
 
             if (options?.onprogress) {
-                this._progressHandlers.set(messageId, options.onprogress);
+                this._progressManager.registerHandler(messageId, options.onprogress);
                 jsonrpcRequest.params = {
                     ...request.params,
                     _meta: {
@@ -1147,7 +1138,7 @@ export abstract class Protocol<SendRequestT extends Request, SendNotificationT e
 
             const cancel = (reason: unknown) => {
                 this._responseHandlers.delete(messageId);
-                this._progressHandlers.delete(messageId);
+                this._progressManager.removeHandler(messageId);
                 this._cleanupTimeout(messageId);
 
                 this._transport
@@ -1459,11 +1450,7 @@ export abstract class Protocol<SendRequestT extends Request, SendNotificationT e
      * This should be called when a task reaches a terminal status.
      */
     private _cleanupTaskProgressHandler(taskId: string): void {
-        const progressToken = this._taskProgressTokens.get(taskId);
-        if (progressToken !== undefined) {
-            this._progressHandlers.delete(progressToken);
-            this._taskProgressTokens.delete(taskId);
-        }
+        this._progressManager.cleanupTaskProgressHandler(taskId);
     }
 
     /**
diff --git a/packages/core/test/shared/protocol.test.ts b/packages/core/test/shared/protocol.test.ts
@@ -36,7 +36,10 @@ interface TestProtocol {
     _taskMessageQueue?: TaskMessageQueue;
     _requestResolvers: Map<RequestId, (response: JSONRPCResultResponse | Error) => void>;
     _responseHandlers: Map<RequestId, (response: JSONRPCResultResponse | Error) => void>;
-    _taskProgressTokens: Map<string, number>;
+    _progressManager: {
+        getTaskProgressToken(taskId: string): number | undefined;
+        cleanupTaskProgressHandler(taskId: string): void;
+    };
     _clearTaskQueue: (taskId: string, sessionId?: string) => Promise<void>;
     requestTaskStore: (request: Request, authInfo: unknown) => TaskStore;
     // Protected task methods (exposed for testing)
@@ -2564,9 +2567,8 @@ describe('Progress notification support for tasks', () => {
         expect(progressCallback).toHaveBeenCalledTimes(1);
 
         // Verify the task-progress association was created
-        const taskProgressTokens = (protocol as unknown as TestProtocol)._taskProgressTokens as Map<string, number>;
-        expect(taskProgressTokens.has(taskId)).toBe(true);
-        expect(taskProgressTokens.get(taskId)).toBe(progressToken);
+        const progressManager = (protocol as unknown as TestProtocol)._progressManager;
+        expect(progressManager.getTaskProgressToken(taskId)).toBe(progressToken);
 
         // Simulate task completion by calling through the protocol's task store
         // This will trigger the cleanup logic
@@ -2578,7 +2580,7 @@ describe('Progress notification support for tasks', () => {
         await new Promise(resolve => setTimeout(resolve, 50));
 
         // Verify the association was cleaned up
-        expect(taskProgressTokens.has(taskId)).toBe(false);
+        expect(progressManager.getTaskProgressToken(taskId)).toBeUndefined();
 
         // Try to send progress notification after task completion - should be ignored
         progressCallback.mockClear();
