clean up

Author: KKonstantinov

packages/client/src/client/client.ts

@@ -38,6 +38,7 @@ import type {
 import {
     assertClientRequestTaskCapability,
     assertToolsCallTaskCapability,
+    extractTaskManagerOptions,
     CallToolResultSchema,
     CompleteResultSchema,
     CreateMessageRequestSchema,
@@ -56,15 +57,13 @@ import {
     ListResourceTemplatesResultSchema,
     ListToolsResultSchema,
     mergeCapabilities,
-    NullTaskManager,
     parseSchema,
     Protocol,
     ProtocolError,
     ProtocolErrorCode,
     ReadResourceResultSchema,
     SdkError,
-    SdkErrorCode,
-    TaskManager
+    SdkErrorCode
 } from '@modelcontextprotocol/core';
 
 import { ExperimentalClientTasks } from '../experimental/tasks/client.js';
@@ -147,8 +146,7 @@ export function getSupportedElicitationModes(capabilities: ClientCapabilities['e
  * Extended tasks capability that includes runtime configuration (store, messageQueue).
  * The runtime-only fields are stripped before advertising capabilities to servers.
  */
-export type ClientTasksCapabilityWithRuntime = NonNullable<ClientCapabilities['tasks']> &
-    Pick<TaskManagerOptions, 'taskStore' | 'taskMessageQueue' | 'defaultTaskPollInterval' | 'maxTaskQueueSize'>;
+export type ClientTasksCapabilityWithRuntime = NonNullable<ClientCapabilities['tasks']> & TaskManagerOptions;
 
 export type ClientOptions = ProtocolOptions & {
     /**
@@ -216,7 +214,6 @@ export class Client extends Protocol<ClientContext> {
     private _listChangedDebounceTimers: Map<string, ReturnType<typeof setTimeout>> = new Map();
     private _pendingListChangedConfig?: ListChangedHandlers;
     private _enforceStrictCapabilities: boolean;
-    private _taskModule: TaskManager;
 
     /**
      * Initializes this client with the given name and version information.
@@ -225,45 +222,28 @@ export class Client extends Protocol<ClientContext> {
         private _clientInfo: Implementation,
         options?: ClientOptions
     ) {
-        super(options);
+        super({
+            ...options,
+            tasks: extractTaskManagerOptions(options?.capabilities?.tasks)
+        });
         this._capabilities = options?.capabilities ? { ...options.capabilities } : {};
         this._jsonSchemaValidator = options?.jsonSchemaValidator ?? new DefaultJsonSchemaValidator();
         this._enforceStrictCapabilities = options?.enforceStrictCapabilities ?? false;
 
-        // Always create TaskManager (NullTaskManager pattern for streaming support)
+        // Strip runtime-only fields from advertised capabilities
         if (options?.capabilities?.tasks) {
+            // eslint-disable-next-line @typescript-eslint/no-unused-vars
             const { taskStore, taskMessageQueue, defaultTaskPollInterval, maxTaskQueueSize, ...wireCapabilities } =
                 options.capabilities.tasks;
-            // Strip runtime-only config from advertised capabilities
             this._capabilities.tasks = wireCapabilities;
-            this._taskModule = new TaskManager({
-                taskStore,
-                taskMessageQueue,
-                defaultTaskPollInterval,
-                maxTaskQueueSize,
-                enforceStrictCapabilities: options?.enforceStrictCapabilities,
-                assertTaskCapability: method => assertToolsCallTaskCapability(this._serverCapabilities?.tasks?.requests, method, 'Server'),
-                assertTaskHandlerCapability: method =>
-                    assertClientRequestTaskCapability(this._capabilities.tasks?.requests, method, 'Client')
-            });
-        } else {
-            this._taskModule = new NullTaskManager();
         }
-        this.setTaskManager(this._taskModule);
 
         // Store list changed config for setup after connection (when we know server capabilities)
         if (options?.listChanged) {
             this._pendingListChangedConfig = options.listChanged;
         }
     }
 
-    /**
-     * Access the task module.
-     */
-    get taskModule(): TaskManager {
-        return this._taskModule;
-    }
-
     protected override buildContext(ctx: BaseContext, _transportInfo?: MessageExtraInfo): ClientContext {
         return ctx;
     }
@@ -714,6 +694,14 @@ export class Client extends Protocol<ClientContext> {
         }
     }
 
+    protected assertTaskCapability(method: string): void {
+        assertToolsCallTaskCapability(this._serverCapabilities?.tasks?.requests, method, 'Server');
+    }
+
+    protected assertTaskHandlerCapability(method: string): void {
+        assertClientRequestTaskCapability(this._capabilities?.tasks?.requests, method, 'Client');
+    }
+
     async ping(options?: RequestOptions) {
         return this._requestWithSchema({ method: 'ping' }, EmptyResultSchema, options);
     }

packages/client/src/experimental/tasks/client.ts

@@ -48,7 +48,7 @@ export class ExperimentalClientTasks {
     constructor(private readonly _client: Client) {}
 
     private get _module() {
-        return this._client.taskModule;
+        return this._client.taskManager;
     }
 
     /**

packages/core/src/experimental/tasks/helpers.ts

@@ -17,7 +17,7 @@ interface TaskRequestsCapability {
 
 /**
  * Asserts that task creation is supported for `tools/call`.
- * Used as the `assertTaskCapability` or `assertTaskHandlerCapability` callback in `TaskManagerOptions`.
+ * Used to implement the `assertTaskCapability` or `assertTaskHandlerCapability` abstract methods on Protocol.
  *
  * @param requests - The task requests capability object
  * @param method - The method being checked
@@ -52,7 +52,7 @@ export function assertToolsCallTaskCapability(
 
 /**
  * Asserts that task creation is supported for `sampling/createMessage` or `elicitation/create`.
- * Used as the `assertTaskCapability` or `assertTaskHandlerCapability` callback in `TaskManagerOptions`.
+ * Used to implement the `assertTaskCapability` or `assertTaskHandlerCapability` abstract methods on Protocol.
  *
  * @param requests - The task requests capability object
  * @param method - The method being checked

packages/core/src/index.ts

@@ -7,7 +7,7 @@ export * from './shared/protocol.js';
 export * from './shared/responseMessage.js';
 export * from './shared/stdio.js';
 export type { RequestTaskStore, TaskContext, TaskManagerOptions, TaskRequestOptions } from './shared/taskManager.js';
-export { NullTaskManager, TaskManager } from './shared/taskManager.js';
+export { extractTaskManagerOptions, NullTaskManager, TaskManager } from './shared/taskManager.js';
 export * from './shared/toolNameValidation.js';
 export * from './shared/transport.js';
 export * from './shared/uriTemplate.js';

packages/core/src/shared/protocol.ts

@@ -47,7 +47,8 @@ import {
 } from '../types/types.js';
 import type { AnySchema, SchemaOutput } from '../util/schema.js';
 import { parseSchema } from '../util/schema.js';
-import type { TaskContext, TaskManager, TaskManagerHost, TaskRequestOptions } from './taskManager.js';
+import type { TaskContext, TaskManagerHost, TaskManagerOptions, TaskRequestOptions } from './taskManager.js';
+import { NullTaskManager, TaskManager } from './taskManager.js';
 import type { Transport, TransportSendOptions } from './transport.js';
 
 /**
@@ -82,6 +83,17 @@ export type ProtocolOptions = {
      * e.g., `['notifications/tools/list_changed']`
      */
     debouncedNotificationMethods?: string[];
+
+    /**
+     * Runtime configuration for task management.
+     * If provided, creates a TaskManager with the given options; otherwise a NullTaskManager is used.
+     *
+     * Capability assertions are wired automatically from the abstract
+     * {@linkcode Protocol.assertTaskCapability | assertTaskCapability()} and
+     * {@linkcode Protocol.assertTaskHandlerCapability | assertTaskHandlerCapability()} methods,
+     * so they should NOT be included here.
+     */
+    tasks?: TaskManagerOptions;
 };
 
 /**
@@ -281,10 +293,6 @@ type TimeoutInfo = {
     onTimeout: () => void;
 };
 
-async function _noopRouteResponse(): Promise<boolean> {
-    return false;
-}
-
 /**
  * Implements MCP protocol framing on top of a pluggable transport, including
  * features like request/response linking, notifications, and progress.
@@ -300,7 +308,7 @@ export abstract class Protocol<ContextT extends BaseContext> {
     private _timeoutInfo: Map<number, TimeoutInfo> = new Map();
     private _pendingDebouncedNotifications = new Set<string>();
 
-    private _taskManager?: TaskManager;
+    private _taskManager: TaskManager;
 
     protected _supportedProtocolVersions: string[];
 
@@ -331,6 +339,10 @@ export abstract class Protocol<ContextT extends BaseContext> {
     constructor(private _options?: ProtocolOptions) {
         this._supportedProtocolVersions = _options?.supportedProtocolVersions ?? SUPPORTED_PROTOCOL_VERSIONS;
 
+        // Create TaskManager from protocol options
+        this._taskManager = _options?.tasks ? new TaskManager(_options.tasks) : new NullTaskManager();
+        this._bindTaskManager();
+
         this.setNotificationHandler('notifications/cancelled', notification => {
             this._oncancel(notification);
         });
@@ -347,11 +359,15 @@ export abstract class Protocol<ContextT extends BaseContext> {
     }
 
     /**
-     * Sets the TaskManager that hooks into the message lifecycle.
-     * The TaskManager is bound to this Protocol and can register handlers, send messages, etc.
+     * Access the TaskManager for task orchestration.
+     * Always available; returns a NullTaskManager when no task store is configured.
      */
-    protected setTaskManager(taskManager: TaskManager): void {
-        this._taskManager = taskManager;
+    get taskManager(): TaskManager {
+        return this._taskManager;
+    }
+
+    private _bindTaskManager(): void {
+        const taskManager = this._taskManager;
         const host: TaskManagerHost = {
             request: (request, resultSchema, options) => this._requestWithSchema(request, resultSchema, options),
             notification: (notification, options) => this.notification(notification, options),
@@ -367,7 +383,10 @@ export abstract class Protocol<ContextT extends BaseContext> {
             },
             sendOnResponseStream: async (message, relatedRequestId) => {
                 await this._transport?.send(message, { relatedRequestId });
-            }
+            },
+            enforceStrictCapabilities: this._options?.enforceStrictCapabilities === true,
+            assertTaskCapability: method => this.assertTaskCapability(method),
+            assertTaskHandlerCapability: method => this.assertTaskHandlerCapability(method)
         };
         taskManager.bind(host);
     }
@@ -473,7 +492,7 @@ export abstract class Protocol<ContextT extends BaseContext> {
         const responseHandlers = this._responseHandlers;
         this._responseHandlers = new Map();
         this._progressHandlers.clear();
-        this._taskManager?.onClose();
+        this._taskManager.onClose();
         this._pendingDebouncedNotifications.clear();
 
         const error = new SdkError(SdkErrorCode.ConnectionClosed, 'Connection closed');
@@ -519,22 +538,14 @@ export abstract class Protocol<ContextT extends BaseContext> {
                 this._requestWithSchema(r, resultSchema, { ...options, relatedRequestId: request.id })
         };
 
-        // Compose results from all modules
-        let sendNotification: (notification: Notification) => Promise<void> = (notification: Notification) =>
-            inboundCtx.sendNotification(notification);
-        let sendRequest = inboundCtx.sendRequest;
-        let routeResponse: (message: JSONRPCResponse | JSONRPCErrorResponse) => Promise<boolean> = _noopRouteResponse;
-        let taskContext: BaseContext['task'] | undefined;
+        // Delegate to TaskManager for task context, wrapped send/notify, and response routing
+        const taskResult = this._taskManager.processInboundRequest(request, inboundCtx);
+        const sendNotification = taskResult.sendNotification;
+        const sendRequest = taskResult.sendRequest;
+        const taskContext = taskResult.taskContext;
+        const routeResponse = taskResult.routeResponse;
         const validators: Array<() => void> = [];
-
-        if (this._taskManager) {
-            const taskResult = this._taskManager.processInboundRequest(request, inboundCtx);
-            sendNotification = taskResult.sendNotification;
-            sendRequest = taskResult.sendRequest;
-            if (taskResult.taskContext !== undefined) taskContext = taskResult.taskContext;
-            if (taskResult.validateInbound) validators.push(taskResult.validateInbound);
-            routeResponse = taskResult.routeResponse;
-        }
+        if (taskResult.validateInbound) validators.push(taskResult.validateInbound);
 
         if (handler === undefined) {
             const errorResponse: JSONRPCErrorResponse = {
@@ -669,12 +680,9 @@ export abstract class Protocol<ContextT extends BaseContext> {
         const messageId = Number(response.id);
 
         // Delegate to TaskManager for task-related response handling
-        let preserveProgress = false;
-        if (this._taskManager) {
-            const taskResult = this._taskManager.processInboundResponse(response, messageId);
-            if (taskResult.consumed) return;
-            preserveProgress = taskResult.preserveProgress;
-        }
+        const taskResult = this._taskManager.processInboundResponse(response, messageId);
+        if (taskResult.consumed) return;
+        const preserveProgress = taskResult.preserveProgress;
 
         const handler = this._responseHandlers.get(messageId);
         if (handler === undefined) {
@@ -730,6 +738,22 @@ export abstract class Protocol<ContextT extends BaseContext> {
      */
     protected abstract assertRequestHandlerCapability(method: string): void;
 
+    /**
+     * A method to check if the remote side supports task creation for the given method.
+     *
+     * Called when sending a task-augmented outbound request (only when enforceStrictCapabilities is true).
+     * This should be implemented by subclasses.
+     */
+    protected abstract assertTaskCapability(method: string): void;
+
+    /**
+     * A method to check if this side supports handling task creation for the given method.
+     *
+     * Called when receiving a task-augmented inbound request.
+     * This should be implemented by subclasses.
+     */
+    protected abstract assertTaskHandlerCapability(method: string): void;
+
     /**
      * Sends a request and waits for a response, resolving the result schema
      * automatically from the method name.
@@ -862,28 +886,20 @@ export abstract class Protocol<ContextT extends BaseContext> {
             };
 
             let outboundQueued = false;
-            if (this._taskManager) {
-                try {
-                    const taskResult = this._taskManager.processOutboundRequest(
-                        jsonrpcRequest,
-                        options,
-                        messageId,
-                        responseHandler,
-                        error => {
-                            this._cleanupTimeout(messageId);
-                            reject(error);
-                        }
-                    );
-                    if (taskResult.queued) {
-                        outboundQueued = true;
-                    }
-                } catch (error) {
-                    this._responseHandlers.delete(messageId);
-                    this._progressHandlers.delete(messageId);
+            try {
+                const taskResult = this._taskManager.processOutboundRequest(jsonrpcRequest, options, messageId, responseHandler, error => {
                     this._cleanupTimeout(messageId);
                     reject(error);
-                    return;
+                });
+                if (taskResult.queued) {
+                    outboundQueued = true;
                 }
+            } catch (error) {
+                this._responseHandlers.delete(messageId);
+                this._progressHandlers.delete(messageId);
+                this._cleanupTimeout(messageId);
+                reject(error);
+                return;
             }
 
             if (!outboundQueued) {
@@ -906,20 +922,10 @@ export abstract class Protocol<ContextT extends BaseContext> {
 
         this.assertNotificationCapability(notification.method as NotificationMethod);
 
-        // Delegate task-related notification routing and JSONRPC building to modules (if registered)
-        let queued = false;
-        let jsonrpcNotification: JSONRPCNotification | undefined;
-
-        if (this._taskManager) {
-            const taskResult = await this._taskManager.processOutboundNotification(notification, options);
-            if (taskResult.queued) {
-                queued = true;
-            } else {
-                jsonrpcNotification = taskResult.jsonrpcNotification;
-            }
-        } else {
-            jsonrpcNotification = { ...notification, jsonrpc: '2.0' };
-        }
+        // Delegate task-related notification routing and JSONRPC building to TaskManager
+        const taskResult = await this._taskManager.processOutboundNotification(notification, options);
+        const queued = taskResult.queued;
+        const jsonrpcNotification = taskResult.queued ? undefined : taskResult.jsonrpcNotification;
 
         if (queued) {
             // Don't send through transport - queued messages are delivered via tasks/result only