v2 - RFC Extract Tasks out of protocol.ts into TaskManager (#1673)

Co-authored-by: Felix Weinberger <fweinberger@anthropic.com>

Author: KKonstantinov

.changeset/extract-task-manager.md

@@ -0,0 +1,10 @@
+---
+"@modelcontextprotocol/core": minor
+"@modelcontextprotocol/client": minor
+"@modelcontextprotocol/server": minor
+---
+
+refactor: extract task orchestration from Protocol into TaskManager
+
+**Breaking changes:**
+- `taskStore`, `taskMessageQueue`, `defaultTaskPollInterval`, and `maxTaskQueueSize` moved from `ProtocolOptions` to `capabilities.tasks` on `ClientOptions`/`ServerOptions`

examples/client/src/simpleStreamableHttp.ts

@@ -265,14 +265,14 @@ async function connect(url?: string): Promise<void> {
                         form: {}
                     },
                     tasks: {
+                        taskStore: clientTaskStore,
                         requests: {
                             elicitation: {
                                 create: {}
                             }
                         }
                     }
-                },
-                taskStore: clientTaskStore
+                }
             }
         );
         client.onerror = error => {

examples/server/src/simpleStreamableHttp.ts

@@ -41,9 +41,14 @@ const getServer = () => {
             websiteUrl: 'https://github.com/modelcontextprotocol/typescript-sdk'
         },
         {
-            capabilities: { logging: {}, tasks: { requests: { tools: { call: {} } } } },
-            taskStore, // Enable task support
-            taskMessageQueue: new InMemoryTaskMessageQueue()
+            capabilities: {
+                logging: {},
+                tasks: {
+                    requests: { tools: { call: {} } },
+                    taskStore,
+                    taskMessageQueue: new InMemoryTaskMessageQueue()
+                }
+            }
         }
     );
 

packages/client/src/client/client.ts

@@ -30,6 +30,7 @@ import type {
     ResultTypeMap,
     ServerCapabilities,
     SubscribeRequest,
+    TaskManagerOptions,
     Tool,
     Transport,
     UnsubscribeRequest
@@ -46,6 +47,7 @@ import {
     ElicitRequestSchema,
     ElicitResultSchema,
     EmptyResultSchema,
+    extractTaskManagerOptions,
     GetPromptResultSchema,
     InitializeResultSchema,
     LATEST_PROTOCOL_VERSION,
@@ -140,11 +142,19 @@ export function getSupportedElicitationModes(capabilities: ClientCapabilities['e
     return { supportsFormMode, supportsUrlMode };
 }
 
+/**
+ * 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']> & TaskManagerOptions;
+
 export type ClientOptions = ProtocolOptions & {
     /**
      * Capabilities to advertise as being supported by this client.
      */
-    capabilities?: ClientCapabilities;
+    capabilities?: Omit<ClientCapabilities, 'tasks'> & {
+        tasks?: ClientTasksCapabilityWithRuntime;
+    };
 
     /**
      * JSON Schema validator for tool output validation.
@@ -213,11 +223,22 @@ export class Client extends Protocol<ClientContext> {
         private _clientInfo: Implementation,
         options?: ClientOptions
     ) {
-        super(options);
-        this._capabilities = options?.capabilities ?? {};
+        super({
+            ...options,
+            tasks: extractTaskManagerOptions(options?.capabilities?.tasks)
+        });
+        this._capabilities = options?.capabilities ? { ...options.capabilities } : {};
         this._jsonSchemaValidator = options?.jsonSchemaValidator ?? new DefaultJsonSchemaValidator();
         this._enforceStrictCapabilities = options?.enforceStrictCapabilities ?? false;
 
+        // 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;
+            this._capabilities.tasks = wireCapabilities;
+        }
+
         // Store list changed config for setup after connection (when we know server capabilities)
         if (options?.listChanged) {
             this._pendingListChangedConfig = options.listChanged;
@@ -650,12 +671,6 @@ export class Client extends Protocol<ClientContext> {
     }
 
     protected assertRequestHandlerCapability(method: string): void {
-        // Task handlers are registered in Protocol constructor before _capabilities is initialized
-        // Skip capability check for task methods during initialization
-        if (!this._capabilities) {
-            return;
-        }
-
         switch (method) {
             case 'sampling/createMessage': {
                 if (!this._capabilities.sampling) {
@@ -687,19 +702,6 @@ export class Client extends Protocol<ClientContext> {
                 break;
             }
 
-            case 'tasks/get':
-            case 'tasks/list':
-            case 'tasks/result':
-            case 'tasks/cancel': {
-                if (!this._capabilities.tasks) {
-                    throw new SdkError(
-                        SdkErrorCode.CapabilityNotSupported,
-                        `Client does not support tasks capability (required for ${method})`
-                    );
-                }
-                break;
-            }
-
             case 'ping': {
                 // No specific capability required for ping
                 break;
@@ -712,16 +714,9 @@ export class Client extends Protocol<ClientContext> {
     }
 
     protected assertTaskHandlerCapability(method: string): void {
-        // Task handlers are registered in Protocol constructor before _capabilities is initialized
-        // Skip capability check for task methods during initialization
-        if (!this._capabilities) {
-            return;
-        }
-
-        assertClientRequestTaskCapability(this._capabilities.tasks?.requests, method, 'Client');
+        assertClientRequestTaskCapability(this._capabilities?.tasks?.requests, method, 'Client');
     }
 
-    /** Sends a ping to the server to check connectivity. */
     async ping(options?: RequestOptions) {
         return this._requestWithSchema({ method: 'ping' }, EmptyResultSchema, options);
     }

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

@@ -6,19 +6,27 @@
  */
 
 import type {
+    AnyObjectSchema,
     CallToolRequest,
     CallToolResult,
     CancelTaskResult,
     CreateTaskResult,
     GetTaskPayloadResult,
     GetTaskResult,
     ListTasksResult,
+    Request,
     RequestMethod,
     RequestOptions,
     ResponseMessage,
     ResultTypeMap
 } from '@modelcontextprotocol/core';
-import { GetTaskPayloadResultSchema, ProtocolError, ProtocolErrorCode } from '@modelcontextprotocol/core';
+import {
+    CallToolResultSchema,
+    getResultSchema,
+    GetTaskPayloadResultSchema,
+    ProtocolError,
+    ProtocolErrorCode
+} from '@modelcontextprotocol/core';
 
 import type { Client } from '../../client/client.js';
 
@@ -27,10 +35,6 @@ import type { Client } from '../../client/client.js';
  * @internal
  */
 interface ClientInternal {
-    requestStream<M extends RequestMethod>(
-        request: { method: M; params?: Record<string, unknown> },
-        options?: RequestOptions
-    ): AsyncGenerator<ResponseMessage<ResultTypeMap[M]>, void, void>;
     isToolTask(toolName: string): boolean;
     getToolOutputValidator(toolName: string): ((data: unknown) => { valid: boolean; errorMessage?: string }) | undefined;
 }
@@ -49,6 +53,10 @@ interface ClientInternal {
 export class ExperimentalClientTasks {
     constructor(private readonly _client: Client) {}
 
+    private get _module() {
+        return this._client.taskManager;
+    }
+
     /**
      * Calls a tool and returns an AsyncGenerator that yields response messages.
      * The generator is guaranteed to end with either a `'result'` or `'error'` message.
@@ -103,7 +111,7 @@ export class ExperimentalClientTasks {
             task: options?.task ?? (clientInternal.isToolTask(params.name) ? {} : undefined)
         };
 
-        const stream = clientInternal.requestStream({ method: 'tools/call', params }, optionsWithTask);
+        const stream = this._module.requestStream({ method: 'tools/call', params }, CallToolResultSchema, optionsWithTask);
 
         // Get the validator for this tool (if it has an output schema)
         const validator = clientInternal.getToolOutputValidator(params.name);
@@ -175,9 +183,7 @@ export class ExperimentalClientTasks {
      * @experimental
      */
     async getTask(taskId: string, options?: RequestOptions): Promise<GetTaskResult> {
-        // Delegate to the client's underlying Protocol method
-        type ClientWithGetTask = { getTask(params: { taskId: string }, options?: RequestOptions): Promise<GetTaskResult> };
-        return (this._client as unknown as ClientWithGetTask).getTask({ taskId }, options);
+        return this._module.getTask({ taskId }, options);
     }
 
     /**
@@ -191,15 +197,7 @@ export class ExperimentalClientTasks {
      * @experimental
      */
     async getTaskResult(taskId: string, options?: RequestOptions): Promise<GetTaskPayloadResult> {
-        return (
-            this._client as unknown as {
-                getTaskResult: (
-                    params: { taskId: string },
-                    resultSchema: typeof GetTaskPayloadResultSchema,
-                    options?: RequestOptions
-                ) => Promise<GetTaskPayloadResult>;
-            }
-        ).getTaskResult({ taskId }, GetTaskPayloadResultSchema, options);
+        return this._module.getTaskResult({ taskId }, GetTaskPayloadResultSchema, options);
     }
 
     /**
@@ -212,12 +210,7 @@ export class ExperimentalClientTasks {
      * @experimental
      */
     async listTasks(cursor?: string, options?: RequestOptions): Promise<ListTasksResult> {
-        // Delegate to the client's underlying Protocol method
-        return (
-            this._client as unknown as {
-                listTasks: (params?: { cursor?: string }, options?: RequestOptions) => Promise<ListTasksResult>;
-            }
-        ).listTasks(cursor ? { cursor } : undefined, options);
+        return this._module.listTasks(cursor ? { cursor } : undefined, options);
     }
 
     /**
@@ -229,12 +222,7 @@ export class ExperimentalClientTasks {
      * @experimental
      */
     async cancelTask(taskId: string, options?: RequestOptions): Promise<CancelTaskResult> {
-        // Delegate to the client's underlying Protocol method
-        return (
-            this._client as unknown as {
-                cancelTask: (params: { taskId: string }, options?: RequestOptions) => Promise<CancelTaskResult>;
-            }
-        ).cancelTask({ taskId }, options);
+        return this._module.cancelTask({ taskId }, options);
     }
 
     /**
@@ -279,7 +267,11 @@ export class ExperimentalClientTasks {
         request: { method: M; params?: Record<string, unknown> },
         options?: RequestOptions
     ): AsyncGenerator<ResponseMessage<ResultTypeMap[M]>, void, void> {
-        // Delegate to the client's underlying Protocol method
-        return (this._client as unknown as ClientInternal).requestStream(request, options);
+        const resultSchema = getResultSchema(request.method) as unknown as AnyObjectSchema;
+        return this._module.requestStream(request as Request, resultSchema, options) as AsyncGenerator<
+            ResponseMessage<ResultTypeMap[M]>,
+            void,
+            void
+        >;
     }
 }

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

@@ -17,7 +17,7 @@ interface TaskRequestsCapability {
 
 /**
  * Asserts that task creation is supported for `tools/call`.
- * Used by {@linkcode @modelcontextprotocol/client!client/client.Client.assertTaskCapability | Client.assertTaskCapability} and {@linkcode @modelcontextprotocol/server!server/server.Server.assertTaskHandlerCapability | Server.assertTaskHandlerCapability}.
+ * 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 by {@linkcode @modelcontextprotocol/server!server/server.Server.assertTaskCapability | Server.assertTaskCapability} and {@linkcode @modelcontextprotocol/client!client/client.Client.assertTaskHandlerCapability | Client.assertTaskHandlerCapability}.
+ * 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/experimental/tasks/interfaces.ts

@@ -3,7 +3,8 @@
  * WARNING: These APIs are experimental and may change without notice.
  */
 
-import type { RequestTaskStore, ServerContext } from '../../shared/protocol.js';
+import type { ServerContext } from '../../shared/protocol.js';
+import type { RequestTaskStore } from '../../shared/taskManager.js';
 import type {
     JSONRPCErrorResponse,
     JSONRPCNotification,

packages/core/src/index.ts

@@ -6,6 +6,8 @@ export * from './shared/metadataUtils.js';
 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 { extractTaskManagerOptions, NullTaskManager, TaskManager } from './shared/taskManager.js';
 export * from './shared/toolNameValidation.js';
 export * from './shared/transport.js';
 export * from './shared/uriTemplate.js';