fix: call correct handlers in task-required call path

Author: LucaButBoring

src/server/mcp.ts

@@ -82,9 +82,37 @@ export class McpServer {
     private _registeredTools: { [name: string]: RegisteredTool } = {};
     private _registeredPrompts: { [name: string]: RegisteredPrompt } = {};
     private _experimental?: { tasks: ExperimentalMcpServerTasks };
+    private _taskToolMap: Map<string, string> = new Map();
 
     constructor(serverInfo: Implementation, options?: ServerOptions) {
-        this.server = new Server(serverInfo, options);
+        const taskHandlerHooks = {
+            getTask: async (taskId: string, extra: RequestHandlerExtra<ServerRequest, ServerNotification>) => {
+                // taskStore is guaranteed to exist here because Protocol only calls hooks when taskStore is configured
+                const taskStore = extra.taskStore!;
+                const handler = this._getTaskHandler(taskId);
+                if (handler) {
+                    return await handler.getTask({ ...extra, taskId, taskStore });
+                }
+                return await taskStore.getTask(taskId);
+            },
+            getTaskResult: async (taskId: string, extra: RequestHandlerExtra<ServerRequest, ServerNotification>) => {
+                const taskStore = extra.taskStore!;
+                const handler = this._getTaskHandler(taskId);
+                if (handler) {
+                    return await handler.getTaskResult({ ...extra, taskId, taskStore });
+                }
+                return await taskStore.getTaskResult(taskId);
+            }
+        };
+        this.server = new Server(serverInfo, { ...options, taskHandlerHooks });
+    }
+
+    private _getTaskHandler(taskId: string): ToolTaskHandler<ZodRawShapeCompat | undefined> | null {
+        const toolName = this._taskToolMap.get(taskId);
+        if (!toolName) return null;
+        const tool = this._registeredTools[toolName];
+        if (!tool || !('createTask' in (tool.handler as AnyToolHandler<ZodRawShapeCompat>))) return null;
+        return tool.handler as ToolTaskHandler<ZodRawShapeCompat | undefined>;
     }
 
     /**
@@ -215,6 +243,10 @@ export class McpServer {
 
                 // Return CreateTaskResult immediately for task requests
                 if (isTaskRequest) {
+                    const taskResult = result as CreateTaskResult;
+                    if (taskResult.task?.taskId) {
+                        this._taskToolMap.set(taskResult.task.taskId, request.params.name);
+                    }
                     return result;
                 }
 

src/shared/protocol.ts

@@ -98,6 +98,20 @@ export type ProtocolOptions = {
      * appropriately (e.g., by failing the task, dropping messages, etc.).
      */
     maxTaskQueueSize?: number;
+    /**
+     * Optional hooks for customizing task request handling.
+     * If a hook is provided, it fully owns the behavior (no fallback to TaskStore).
+     */
+    taskHandlerHooks?: {
+        /**
+         * Called when tasks/get is received. If provided, must return the task.
+         */
+        getTask?: (taskId: string, extra: RequestHandlerExtra<Request, Notification>) => Promise<GetTaskResult>;
+        /**
+         * Called when tasks/payload needs to retrieve the final result. If provided, must return the result.
+         */
+        getTaskResult?: (taskId: string, extra: RequestHandlerExtra<Request, Notification>) => Promise<Result>;
+    };
 };
 
 /**
@@ -383,6 +397,16 @@ export abstract class Protocol<SendRequestT extends Request, SendNotificationT e
         this._taskMessageQueue = _options?.taskMessageQueue;
         if (this._taskStore) {
             this.setRequestHandler(GetTaskRequestSchema, async (request, extra) => {
+                // Use hook if provided, otherwise fall back to TaskStore
+                if (_options?.taskHandlerHooks?.getTask) {
+                    const hookResult = await _options.taskHandlerHooks.getTask(
+                        request.params.taskId,
+                        extra as unknown as RequestHandlerExtra<Request, Notification>
+                    );
+                    // @ts-expect-error SendResultT cannot contain GetTaskResult
+                    return hookResult as SendResultT;
+                }
+
                 const task = await this._taskStore!.getTask(request.params.taskId, extra.sessionId);
                 if (!task) {
                     throw new McpError(ErrorCode.InvalidParams, 'Failed to retrieve task: Task not found');
@@ -462,7 +486,13 @@ export abstract class Protocol<SendRequestT extends Request, SendNotificationT e
 
                     // If task is terminal, return the result
                     if (isTerminal(task.status)) {
-                        const result = await this._taskStore!.getTaskResult(taskId, extra.sessionId);
+                        // Use hook if provided, otherwise fall back to TaskStore
+                        const result = this._options?.taskHandlerHooks?.getTaskResult
+                            ? await this._options.taskHandlerHooks.getTaskResult(
+                                  taskId,
+                                  extra as unknown as RequestHandlerExtra<Request, Notification>
+                              )
+                            : await this._taskStore!.getTaskResult(taskId, extra.sessionId);
 
                         this._clearTaskQueue(taskId);
 

test/server/mcp.test.ts

@@ -6,8 +6,10 @@ import {
     CallToolResultSchema,
     type CallToolResult,
     CompleteResultSchema,
+    CreateTaskResultSchema,
     ElicitRequestSchema,
     GetPromptResultSchema,
+    GetTaskResultSchema,
     ListPromptsResultSchema,
     ListResourcesResultSchema,
     ListResourceTemplatesResultSchema,
@@ -6890,5 +6892,113 @@ describe.each(zodTestMatrix)('$zodVersionLabel', (entry: ZodMatrixEntry) => {
 
             taskStore.cleanup();
         });
+
+        test('should call custom getTask and getTaskResult handlers when client polls task directly', async () => {
+            const taskStore = new InMemoryTaskStore();
+
+            const getTaskSpy = vi.fn();
+            const getTaskResultSpy = vi.fn();
+            let taskCreatedAt: number;
+
+            const mcpServer = new McpServer(
+                {
+                    name: 'test server',
+                    version: '1.0'
+                },
+                {
+                    capabilities: {
+                        tools: {},
+                        tasks: {
+                            requests: {
+                                tools: {
+                                    call: {}
+                                }
+                            }
+                        }
+                    },
+                    taskStore
+                }
+            );
+
+            const client = new Client(
+                {
+                    name: 'test client',
+                    version: '1.0'
+                },
+                {
+                    capabilities: {
+                        tasks: {
+                            requests: {
+                                tools: {
+                                    call: {}
+                                }
+                            }
+                        }
+                    }
+                }
+            );
+
+            mcpServer.experimental.tasks.registerToolTask(
+                'task-tool',
+                {
+                    description: 'A task tool',
+                    inputSchema: { data: z.string() },
+                    execution: { taskSupport: 'required' }
+                },
+                {
+                    createTask: async (_args, extra) => {
+                        const task = await extra.taskStore.createTask({ ttl: 60000, pollInterval: 100 });
+                        taskCreatedAt = Date.now();
+                        return { task };
+                    },
+                    getTask: async extra => {
+                        getTaskSpy(extra.taskId);
+                        // Complete the task after 50ms - only occurs if getTask is actually called
+                        if (Date.now() - taskCreatedAt >= 50) {
+                            await extra.taskStore.storeTaskResult(extra.taskId, 'completed', {
+                                content: [{ type: 'text' as const, text: 'Done' }]
+                            });
+                        }
+                        return await extra.taskStore.getTask(extra.taskId);
+                    },
+                    getTaskResult: async extra => {
+                        getTaskResultSpy(extra.taskId);
+                        return (await extra.taskStore.getTaskResult(extra.taskId)) as CallToolResult;
+                    }
+                }
+            );
+
+            const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair();
+            await Promise.all([client.connect(clientTransport), mcpServer.connect(serverTransport)]);
+
+            // Create task
+            const createResult = await client.request(
+                {
+                    method: 'tools/call',
+                    params: {
+                        name: 'task-tool',
+                        arguments: { data: 'test' },
+                        task: { ttl: 60000 }
+                    }
+                },
+                CreateTaskResultSchema
+            );
+            const taskId = createResult.task.taskId;
+
+            // Wait for task to be ready to complete
+            await new Promise(resolve => setTimeout(resolve, 60));
+
+            // Client directly calls tasks/get - should invoke custom handler which completes the task
+            const getResult = await client.request({ method: 'tasks/get', params: { taskId } }, GetTaskResultSchema);
+            expect(getResult.status).toBe('completed');
+            expect(getTaskSpy).toHaveBeenCalledWith(taskId);
+
+            // Client directly calls tasks/result - should invoke custom handler
+            const payloadResult = await client.request({ method: 'tasks/result', params: { taskId } }, CallToolResultSchema);
+            expect(payloadResult.content).toBeDefined();
+            expect(getTaskResultSpy).toHaveBeenCalledWith(taskId);
+
+            taskStore.cleanup();
+        });
     });
 });