fix: fix auto-poll task store access and sendRequest wrapping

- Fix handleAutomaticTaskPolling to create task store from server.tasks when
  extra.task is not set (auto-poll creates tasks internally)
- Fix sendRequest wrapping to activate on taskContext (not just relatedTaskId)
  so task creation handlers can properly set input_required status

Integration: 482/487 pass (5 failures in HTTP message queuing delivery).

Author: felixweinberger

packages/core/src/shared/protocol.ts

@@ -509,13 +509,14 @@ export abstract class Protocol<SendRequestT extends Request, SendNotificationT e
         }
 
         // Build sendRequest closure, potentially wrapped with task metadata
+        // Wrap when there's a task context (either relatedTaskId or task creation params)
         let sendRequest = async <U extends AnySchema>(r: SendRequestT, resultSchema: U, options?: RequestOptions) => {
             return await this.request(r, resultSchema, { ...options, relatedRequestId: request.id });
         };
-        if (relatedTaskId && this._taskManager) {
-            const taskStore = taskInfo?.taskContext?.taskStore;
+        if (taskInfo?.taskContext && this._taskManager) {
+            const taskStore = taskInfo.taskContext.taskStore;
             sendRequest = this._taskManager.wrapSendRequest(
-                relatedTaskId,
+                relatedTaskId ?? '',
                 taskStore,
                 (r: SendRequestT, resultSchema: AnySchema, options?: RequestOptions) =>
                     this.request(r, resultSchema, { ...options, relatedRequestId: request.id })

packages/server/src/server/mcp.ts

@@ -9,6 +9,7 @@ import type {
     CompleteResult,
     CreateTaskRequestHandlerExtra,
     CreateTaskResult,
+    JSONRPCRequest,
     GetPromptResult,
     Implementation,
     ListPromptsResult,
@@ -366,14 +367,27 @@ export class McpServer {
         request: RequestT,
         extra: RequestHandlerExtra<ServerRequest, ServerNotification>
     ): Promise<CallToolResult> {
-        if (!extra.task?.taskStore) {
+        // For automatic polling, we need a task store. Use the one from extra.task if available,
+        // otherwise create one from the server's TaskManager (auto-poll creates tasks internally).
+        const taskStore = extra.task?.taskStore ?? this.server.tasks?.createRequestTaskStore(
+            { jsonrpc: '2.0', id: extra.requestId, method: request.method, params: request.params } as JSONRPCRequest,
+            extra.sessionId
+        );
+        if (!taskStore) {
             throw new Error('No task store provided for task-capable tool.');
         }
 
+        // Build a task context for the handler
+        const taskContext: CreateTaskRequestHandlerExtra['task'] = extra.task ?? {
+            taskId: '',
+            taskStore,
+            requestedTtl: undefined
+        };
+        const taskExtra = { ...extra, task: taskContext } as CreateTaskRequestHandlerExtra;
+
         // Validate input and create task
         const args = await this.validateToolInput(tool, request.params.arguments, request.params.name);
         const handler = tool.handler as ToolTaskHandler<ZodRawShapeCompat | undefined>;
-        const taskExtra = { ...extra, task: extra.task } as CreateTaskRequestHandlerExtra;
 
         const createTaskResult: CreateTaskResult = args // undefined only if tool.inputSchema is undefined
             ? await Promise.resolve((handler as ToolTaskHandler<ZodRawShapeCompat>).createTask(args, taskExtra))
@@ -387,15 +401,15 @@ export class McpServer {
 
         while (task.status !== 'completed' && task.status !== 'failed' && task.status !== 'cancelled') {
             await new Promise(resolve => setTimeout(resolve, pollInterval));
-            const updatedTask = await extra.task.taskStore.getTask(taskId);
+            const updatedTask = await taskStore.getTask(taskId);
             if (!updatedTask) {
                 throw new McpError(ErrorCode.InternalError, `Task ${taskId} not found during polling`);
             }
             task = updatedTask;
         }
 
         // Return the final result
-        return (await extra.task.taskStore.getTaskResult(taskId)) as CallToolResult;
+        return (await taskStore.getTaskResult(taskId)) as CallToolResult;
     }
 
     private _completionHandlerInitialized = false;