fix: address review feedback on TaskManager extraction

felixweinberger · felixweinberger · commit b783fcdf7c02 · 2026-02-03T10:06:13.000Z
- Move assertTaskCapability/assertTaskHandlerCapability from mutable properties on TaskManager to TaskManagerOptions constructor params, eliminating fragile post-construction wiring - Add _requireHost getter with clear error message instead of _host! non-null assertions throughout TaskManager - Extract duplicated relatedTask metadata injection in Protocol.notification() into _buildJsonRpcNotification() helper - Add @internal JSDoc to TaskManagerHost (will be removed from public exports in PR #1447) - Add changeset for the breaking changes - Revert unrelated pnpm-workspace.yaml change
diff --git a/.changeset/extract-task-manager.md b/.changeset/extract-task-manager.md
@@ -0,0 +1,14 @@
+---
+"@modelcontextprotocol/core": minor
+"@modelcontextprotocol/client": minor
+"@modelcontextprotocol/server": minor
+---
+
+refactor: extract task orchestration from Protocol into TaskManager
+
+**Breaking changes:**
+- `extra.taskId` → `extra.task?.taskId`
+- `extra.taskStore` → `extra.task?.taskStore`
+- `extra.taskRequestedTtl` → `extra.task?.requestedTtl`
+- `ProtocolOptions` no longer accepts `taskStore`/`taskMessageQueue` — pass via `TaskManagerOptions` in `ClientOptions`/`ServerOptions`
+- Abstract methods `assertTaskCapability`/`assertTaskHandlerCapability` removed from Protocol
diff --git a/packages/client/src/client/client.ts b/packages/client/src/client/client.ts
@@ -265,23 +265,21 @@ export class Client<
             taskStore: options?.taskStore,
             taskMessageQueue: options?.taskMessageQueue,
             defaultTaskPollInterval: options?.defaultTaskPollInterval,
-            maxTaskQueueSize: options?.maxTaskQueueSize
+            maxTaskQueueSize: options?.maxTaskQueueSize,
+            assertTaskCapability: (method: string) => {
+                assertToolsCallTaskCapability(this._serverCapabilities?.tasks?.requests, method, 'Server');
+            },
+            assertTaskHandlerCapability: (method: string) => {
+                if (!this._capabilities) return;
+                assertClientRequestTaskCapability(this._capabilities.tasks?.requests, method, 'Client');
+            }
         });
 
         super(options, taskManager);
         this._capabilities = options?.capabilities ?? {};
         this._jsonSchemaValidator = options?.jsonSchemaValidator ?? new AjvJsonSchemaValidator();
         this._enforceStrictCapabilities = options?.enforceStrictCapabilities ?? false;
 
-        // Wire task capability assertions
-        taskManager.assertTaskCapability = (method: string) => {
-            assertToolsCallTaskCapability(this._serverCapabilities?.tasks?.requests, method, 'Server');
-        };
-        taskManager.assertTaskHandlerCapability = (method: string) => {
-            if (!this._capabilities) return;
-            assertClientRequestTaskCapability(this._capabilities.tasks?.requests, method, 'Client');
-        };
-
         // Store list changed config for setup after connection (when we know server capabilities)
         if (options?.listChanged) {
             this._pendingListChangedConfig = options.listChanged;
diff --git a/packages/core/src/shared/protocol.ts b/packages/core/src/shared/protocol.ts
@@ -837,49 +837,37 @@ export abstract class Protocol<SendRequestT extends Request, SendNotificationT e
                     return;
                 }
 
-                let jsonrpcNotification: JSONRPCNotification = {
-                    ...notification,
-                    jsonrpc: '2.0'
-                };
-
-                if (options?.relatedTask) {
-                    jsonrpcNotification = {
-                        ...jsonrpcNotification,
-                        params: {
-                            ...jsonrpcNotification.params,
-                            _meta: {
-                                ...jsonrpcNotification.params?._meta,
-                                [RELATED_TASK_META_KEY]: options.relatedTask
-                            }
-                        }
-                    };
-                }
-
+                const jsonrpcNotification = this._buildJsonRpcNotification(notification, options);
                 this._transport?.send(jsonrpcNotification, options).catch(error => this._onerror(error));
             });
 
             return;
         }
 
-        let jsonrpcNotification: JSONRPCNotification = {
+        const jsonrpcNotification = this._buildJsonRpcNotification(notification, options);
+        await this._transport.send(jsonrpcNotification, options);
+    }
+
+    private _buildJsonRpcNotification(notification: SendNotificationT, options?: NotificationOptions): JSONRPCNotification {
+        const jsonrpcNotification: JSONRPCNotification = {
             ...notification,
             jsonrpc: '2.0'
         };
 
-        if (options?.relatedTask) {
-            jsonrpcNotification = {
-                ...jsonrpcNotification,
-                params: {
-                    ...jsonrpcNotification.params,
-                    _meta: {
-                        ...jsonrpcNotification.params?._meta,
-                        [RELATED_TASK_META_KEY]: options.relatedTask
-                    }
-                }
-            };
+        if (!options?.relatedTask) {
+            return jsonrpcNotification;
         }
 
-        await this._transport.send(jsonrpcNotification, options);
+        return {
+            ...jsonrpcNotification,
+            params: {
+                ...jsonrpcNotification.params,
+                _meta: {
+                    ...jsonrpcNotification.params?._meta,
+                    [RELATED_TASK_META_KEY]: options.relatedTask
+                }
+            }
+        };
     }
 
     /**
diff --git a/packages/core/src/shared/taskManager.ts b/packages/core/src/shared/taskManager.ts
@@ -63,6 +63,16 @@ export type TaskManagerOptions = {
      * If undefined, the queue size is unbounded.
      */
     maxTaskQueueSize?: number;
+    /**
+     * Checks if task creation is supported for a given request method.
+     * Wired by Client/Server to enforce capability negotiation.
+     */
+    assertTaskCapability?: (method: string) => void;
+    /**
+     * Checks if task handling is supported for a given request method.
+     * Wired by Client/Server to enforce capability negotiation.
+     */
+    assertTaskHandlerCapability?: (method: string) => void;
 };
 
 /**
@@ -94,6 +104,8 @@ export interface TaskContext {
 /**
  * Internal interface that TaskManager needs from the Protocol host
  * for sending outbound requests and notifications.
+ *
+ * @internal Not intended for external use. Will be removed from public exports.
  */
 export interface TaskManagerHost {
     request<T extends AnySchema>(request: Request, resultSchema: T, options?: RequestOptions): Promise<SchemaOutput<T>>;
@@ -136,6 +148,13 @@ export class TaskManager {
         this._host = host;
     }
 
+    private get _requireHost(): TaskManagerHost {
+        if (!this._host) {
+            throw new McpError(ErrorCode.InternalError, 'TaskManager is not bound to a Protocol host — call bind() first');
+        }
+        return this._host;
+    }
+
     get taskStore(): TaskStore | undefined {
         return this._taskStore;
     }
@@ -151,17 +170,13 @@ export class TaskManager {
         return this._taskMessageQueue;
     }
 
-    /**
-     * Checks if task creation is supported for a given request method.
-     * Set by Client/Server to wire capability checking.
-     */
-    assertTaskCapability?: (method: string) => void;
+    get assertTaskCapability(): ((method: string) => void) | undefined {
+        return this._options.assertTaskCapability;
+    }
 
-    /**
-     * Checks if task handling is supported for a given request method.
-     * Set by Client/Server to wire capability checking.
-     */
-    assertTaskHandlerCapability?: (method: string) => void;
+    get assertTaskHandlerCapability(): ((method: string) => void) | undefined {
+        return this._options.assertTaskHandlerCapability;
+    }
 
     // =========================================================================
     // Public API methods (previously protected on Protocol)
@@ -176,7 +191,7 @@ export class TaskManager {
         resultSchema: T,
         options?: RequestOptions
     ): AsyncGenerator<ResponseMessage<SchemaOutput<T>>, void, void> {
-        const host = this._host!;
+        const host = this._requireHost;
         const { task } = options ?? {};
 
         // For non-task requests, just yield the result
@@ -247,23 +262,23 @@ export class TaskManager {
     }
 
     async getTask(params: GetTaskRequest['params'], options?: RequestOptions): Promise<GetTaskResult> {
-        return this._host!.request({ method: 'tasks/get', params }, GetTaskResultSchema, options);
+        return this._requireHost.request({ method: 'tasks/get', params }, GetTaskResultSchema, options);
     }
 
     async getTaskResult<T extends AnySchema>(
         params: GetTaskPayloadRequest['params'],
         resultSchema: T,
         options?: RequestOptions
     ): Promise<SchemaOutput<T>> {
-        return this._host!.request({ method: 'tasks/result', params }, resultSchema, options);
+        return this._requireHost.request({ method: 'tasks/result', params }, resultSchema, options);
     }
 
     async listTasks(params?: { cursor?: string }, options?: RequestOptions): Promise<SchemaOutput<typeof ListTasksResultSchema>> {
-        return this._host!.request({ method: 'tasks/list', params }, ListTasksResultSchema, options);
+        return this._requireHost.request({ method: 'tasks/list', params }, ListTasksResultSchema, options);
     }
 
     async cancelTask(params: { taskId: string }, options?: RequestOptions): Promise<SchemaOutput<typeof CancelTaskResultSchema>> {
-        return this._host!.request({ method: 'tasks/cancel', params }, CancelTaskResultSchema, options);
+        return this._requireHost.request({ method: 'tasks/cancel', params }, CancelTaskResultSchema, options);
     }
 
     // =========================================================================
diff --git a/packages/server/src/server/server.ts b/packages/server/src/server/server.ts
@@ -159,23 +159,21 @@ export class Server<
             taskStore: options?.taskStore,
             taskMessageQueue: options?.taskMessageQueue,
             defaultTaskPollInterval: options?.defaultTaskPollInterval,
-            maxTaskQueueSize: options?.maxTaskQueueSize
+            maxTaskQueueSize: options?.maxTaskQueueSize,
+            assertTaskCapability: (method: string) => {
+                assertClientRequestTaskCapability(this._clientCapabilities?.tasks?.requests, method, 'Client');
+            },
+            assertTaskHandlerCapability: (method: string) => {
+                if (!this._capabilities) return;
+                assertToolsCallTaskCapability(this._capabilities.tasks?.requests, method, 'Server');
+            }
         });
 
         super(options, taskManager);
         this._capabilities = options?.capabilities ?? {};
         this._instructions = options?.instructions;
         this._jsonSchemaValidator = options?.jsonSchemaValidator ?? new AjvJsonSchemaValidator();
 
-        // Wire task capability assertions
-        taskManager.assertTaskCapability = (method: string) => {
-            assertClientRequestTaskCapability(this._clientCapabilities?.tasks?.requests, method, 'Client');
-        };
-        taskManager.assertTaskHandlerCapability = (method: string) => {
-            if (!this._capabilities) return;
-            assertToolsCallTaskCapability(this._capabilities.tasks?.requests, method, 'Server');
-        };
-
         this.setRequestHandler('initialize', request => this._oninitialize(request));
         this.setNotificationHandler('notifications/initialized', () => this.oninitialized?.());
 
