refactor: update experimental wrappers and fix task context extraction

- Rewrite ExperimentalClientTasks to delegate to client.tasks (no more unsafe casts)
- Rewrite ExperimentalServerTasks to delegate to server.tasks (no more unsafe casts)
- Make taskStore optional in TaskManagerOptions (clients need TaskManager for outbound API)
- Always create TaskManager in Client/Server (outbound API always available)
- Fix extractInboundTaskContext to create context for task creation requests too
- Update integration test references (extra.task!.taskStore etc.)

Core: 430/430 pass. Integration: 475/487 pass (12 failures in message queuing/auto-poll).

Author: felixweinberger

packages/client/src/client/client.ts

@@ -264,31 +264,28 @@ export class Client<
         private _clientInfo: Implementation,
         options?: ClientOptions
     ) {
-        // Create TaskManager if taskStore is provided
-        const taskManager = options?.taskStore
-            ? new TaskManager({
-                taskStore: options.taskStore,
-                taskMessageQueue: options.taskMessageQueue,
-                defaultTaskPollInterval: options.defaultTaskPollInterval,
-                maxTaskQueueSize: options.maxTaskQueueSize
-            })
-            : undefined;
+        // Always create a TaskManager for outbound task API (getTask, requestStream, etc.)
+        // If taskStore is provided, it also enables handling incoming task requests
+        const taskManager = new TaskManager({
+            taskStore: options?.taskStore,
+            taskMessageQueue: options?.taskMessageQueue,
+            defaultTaskPollInterval: options?.defaultTaskPollInterval,
+            maxTaskQueueSize: options?.maxTaskQueueSize
+        });
 
         super(options, taskManager);
         this._capabilities = options?.capabilities ?? {};
         this._jsonSchemaValidator = options?.jsonSchemaValidator ?? new AjvJsonSchemaValidator();
         this._enforceStrictCapabilities = options?.enforceStrictCapabilities ?? false;
 
         // Wire task capability assertions
-        if (taskManager) {
-            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');
-            };
-        }
+        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) {

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

@@ -6,7 +6,6 @@
  */
 
 import type {
-    AnyObjectSchema,
     CallToolRequest,
     CancelTaskResult,
     ClientRequest,
@@ -21,19 +20,15 @@ import type {
     SchemaOutput
 } from '@modelcontextprotocol/core';
 import { CallToolResultSchema, ErrorCode, McpError } from '@modelcontextprotocol/core';
+import type { AnyObjectSchema } from '@modelcontextprotocol/core';
 
 import type { Client } from '../../client/client.js';
 
 /**
  * Internal interface for accessing Client's private methods.
  * @internal
  */
-interface ClientInternal<RequestT extends Request> {
-    requestStream<T extends AnyObjectSchema>(
-        request: ClientRequest | RequestT,
-        resultSchema: T,
-        options?: RequestOptions
-    ): AsyncGenerator<ResponseMessage<SchemaOutput<T>>, void, void>;
+interface ClientInternal {
     isToolTask(toolName: string): boolean;
     getToolOutputValidator(toolName: string): ((data: unknown) => { valid: boolean; errorMessage?: string }) | undefined;
 }
@@ -58,57 +53,33 @@ export class ExperimentalClientTasks<
 
     /**
      * Calls a tool and returns an AsyncGenerator that yields response messages.
-     * The generator is guaranteed to end with either a 'result' or 'error' message.
-     *
-     * This method provides streaming access to tool execution, allowing you to
-     * observe intermediate task status updates for long-running tool calls.
      * Automatically validates structured output if the tool has an outputSchema.
      *
-     * @example
-     * ```typescript
-     * const stream = client.experimental.tasks.callToolStream({ name: 'myTool', arguments: {} });
-     * for await (const message of stream) {
-     *   switch (message.type) {
-     *     case 'taskCreated':
-     *       console.log('Tool execution started:', message.task.taskId);
-     *       break;
-     *     case 'taskStatus':
-     *       console.log('Tool status:', message.task.status);
-     *       break;
-     *     case 'result':
-     *       console.log('Tool result:', message.result);
-     *       break;
-     *     case 'error':
-     *       console.error('Tool error:', message.error);
-     *       break;
-     *   }
-     * }
-     * ```
-     *
-     * @param params - Tool call parameters (name and arguments)
-     * @param resultSchema - Zod schema for validating the result (defaults to CallToolResultSchema)
-     * @param options - Optional request options (timeout, signal, task creation params, etc.)
-     * @returns AsyncGenerator that yields ResponseMessage objects
-     *
      * @experimental
      */
     async *callToolStream<T extends typeof CallToolResultSchema | typeof CompatibilityCallToolResultSchema>(
         params: CallToolRequest['params'],
         resultSchema: T = CallToolResultSchema as T,
         options?: RequestOptions
     ): AsyncGenerator<ResponseMessage<SchemaOutput<T>>, void, void> {
-        // Access Client's internal methods
-        const clientInternal = this._client as unknown as ClientInternal<RequestT>;
+        const tasks = this._client.tasks;
+        if (!tasks) {
+            throw new McpError(ErrorCode.InternalError, 'Task support is not enabled on this client');
+        }
+
+        const clientInternal = this._client as unknown as ClientInternal;
 
         // Add task creation parameters if server supports it and not explicitly provided
         const optionsWithTask = {
             ...options,
-            // We check if the tool is known to be a task during auto-configuration, but assume
-            // the caller knows what they're doing if they pass this explicitly
             task: options?.task ?? (clientInternal.isToolTask(params.name) ? {} : undefined)
         };
 
-        const stream = clientInternal.requestStream({ method: 'tools/call', params }, resultSchema, optionsWithTask);
+        const stream = tasks.requestStream(
+            { method: 'tools/call', params } as ClientRequest | RequestT,
+            resultSchema,
+            optionsWithTask
+        );
 
         // Get the validator for this tool (if it has an output schema)
         const validator = clientInternal.getToolOutputValidator(params.name);
@@ -134,7 +105,6 @@ export class ExperimentalClientTasks<
                 // Only validate structured content if present (not when there's an error)
                 if (result.structuredContent) {
                     try {
-                        // Validate the structured content against the schema
                         const validationResult = validator(result.structuredContent);
 
                         if (!validationResult.valid) {
@@ -171,104 +141,55 @@ export class ExperimentalClientTasks<
 
     /**
      * Gets the current status of a task.
-     *
-     * @param taskId - The task identifier
-     * @param options - Optional request options
-     * @returns The task status
-     *
      * @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);
+        const tasks = this._client.tasks;
+        if (!tasks) throw new McpError(ErrorCode.InternalError, 'Task support is not enabled');
+        return tasks.getTask({ taskId }, options);
     }
 
     /**
      * Retrieves the result of a completed task.
-     *
-     * @param taskId - The task identifier
-     * @param resultSchema - Zod schema for validating the result
-     * @param options - Optional request options
-     * @returns The task result
-     *
      * @experimental
      */
     async getTaskResult<T extends AnyObjectSchema>(taskId: string, resultSchema?: T, options?: RequestOptions): Promise<SchemaOutput<T>> {
-        // Delegate to the client's underlying Protocol method
-        return (
-            this._client as unknown as {
-                getTaskResult: <U extends AnyObjectSchema>(
-                    params: { taskId: string },
-                    resultSchema?: U,
-                    options?: RequestOptions
-                ) => Promise<SchemaOutput<U>>;
-            }
-        ).getTaskResult({ taskId }, resultSchema, options);
+        const tasks = this._client.tasks;
+        if (!tasks) throw new McpError(ErrorCode.InternalError, 'Task support is not enabled');
+        return tasks.getTaskResult({ taskId }, resultSchema!, options);
     }
 
     /**
      * Lists tasks with optional pagination.
-     *
-     * @param cursor - Optional pagination cursor
-     * @param options - Optional request options
-     * @returns List of tasks with optional next cursor
-     *
      * @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);
+        const tasks = this._client.tasks;
+        if (!tasks) throw new McpError(ErrorCode.InternalError, 'Task support is not enabled');
+        return tasks.listTasks(cursor ? { cursor } : undefined, options);
     }
 
     /**
      * Cancels a running task.
-     *
-     * @param taskId - The task identifier
-     * @param options - Optional request options
-     *
      * @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);
+        const tasks = this._client.tasks;
+        if (!tasks) throw new McpError(ErrorCode.InternalError, 'Task support is not enabled');
+        return tasks.cancelTask({ taskId }, options);
     }
 
     /**
      * Sends a request and returns an AsyncGenerator that yields response messages.
-     * The generator is guaranteed to end with either a 'result' or 'error' message.
-     *
-     * This method provides streaming access to request processing, allowing you to
-     * observe intermediate task status updates for task-augmented requests.
-     *
-     * @param request - The request to send
-     * @param resultSchema - Zod schema for validating the result
-     * @param options - Optional request options (timeout, signal, task creation params, etc.)
-     * @returns AsyncGenerator that yields ResponseMessage objects
-     *
      * @experimental
      */
     requestStream<T extends AnyObjectSchema>(
         request: ClientRequest | RequestT,
         resultSchema: T,
         options?: RequestOptions
     ): AsyncGenerator<ResponseMessage<SchemaOutput<T>>, void, void> {
-        // Delegate to the client's underlying Protocol method
-        type ClientWithRequestStream = {
-            requestStream<U extends AnyObjectSchema>(
-                request: ClientRequest | RequestT,
-                resultSchema: U,
-                options?: RequestOptions
-            ): AsyncGenerator<ResponseMessage<SchemaOutput<U>>, void, void>;
-        };
-        return (this._client as unknown as ClientWithRequestStream).requestStream(request, resultSchema, options);
+        const tasks = this._client.tasks;
+        if (!tasks) throw new McpError(ErrorCode.InternalError, 'Task support is not enabled');
+        return tasks.requestStream(request, resultSchema, options);
     }
 }

packages/core/src/shared/protocol.ts

@@ -291,7 +291,10 @@ export abstract class Protocol<SendRequestT extends Request, SendNotificationT e
                 reportError: (error) => this._onerror(error),
                 removeProgressHandler: (token) => this._progressHandlers.delete(token)
             });
-            this._registerTaskHandlers();
+            // Only register task handlers if a task store is configured
+            if (taskManager.taskStore) {
+                this._registerTaskHandlers();
+            }
         }
     }