fix: call correct handlers in backwards-compat task poll path

Previously, the code only called the underlying task store, and the tests were not complex enough to validate that the handlers were being called, so they missed this.

Author: LucaButBoring

src/server/mcp.ts

@@ -374,27 +374,28 @@ export class McpServer {
         const handler = tool.handler as ToolTaskHandler<ZodRawShapeCompat | undefined>;
         const taskExtra = { ...extra, taskStore: extra.taskStore };
 
-        const createTaskResult: CreateTaskResult = args // undefined only if tool.inputSchema is undefined
-            ? await Promise.resolve((handler as ToolTaskHandler<ZodRawShapeCompat>).createTask(args, taskExtra))
-            : // eslint-disable-next-line @typescript-eslint/no-explicit-any
-              await Promise.resolve(((handler as ToolTaskHandler<undefined>).createTask as any)(taskExtra));
+        const wrappedHandler = toolTaskHandlerByArgs(handler, args);
+
+        const createTaskResult = await wrappedHandler.createTask(taskExtra);
 
         // Poll until completion
         const taskId = createTaskResult.task.taskId;
+        const taskExtraComplete = { ...extra, taskId, taskStore: extra.taskStore };
         let task = createTaskResult.task;
         const pollInterval = task.pollInterval ?? 5000;
 
         while (task.status !== 'completed' && task.status !== 'failed' && task.status !== 'cancelled') {
             await new Promise(resolve => setTimeout(resolve, pollInterval));
-            const updatedTask = await extra.taskStore.getTask(taskId);
+            const getTaskResult = await wrappedHandler.getTask(taskExtraComplete);
+            const updatedTask = getTaskResult;
             if (!updatedTask) {
                 throw new McpError(ErrorCode.InternalError, `Task ${taskId} not found during polling`);
             }
             task = updatedTask;
         }
 
         // Return the final result
-        return (await extra.taskStore.getTaskResult(taskId)) as CallToolResult;
+        return await wrappedHandler.getTaskResult(taskExtraComplete);
     }
 
     private _completionHandlerInitialized = false;
@@ -1540,3 +1541,29 @@ const EMPTY_COMPLETION_RESULT: CompleteResult = {
         hasMore: false
     }
 };
+
+/**
+ * Wraps a tool task handler such that it can be used without checking if it needs to be called in a one-arg manner.
+ * @param handler The task handler to wrap.
+ * @param args The tool arguments.
+ * @returns A wrapped task handler for a tool, which only exposes a no-args interface.
+ */
+function toolTaskHandlerByArgs<Args extends AnySchema | ZodRawShapeCompat | undefined>(
+    handler: ToolTaskHandler<Args>,
+    args: unknown
+): ToolTaskHandler<undefined> {
+    return {
+        createTask: extra =>
+            args // undefined only if tool.inputSchema is undefined
+                ? Promise.resolve((handler as ToolTaskHandler<ZodRawShapeCompat>).createTask(args, extra))
+                : Promise.resolve((handler as ToolTaskHandler<undefined>).createTask(extra)),
+        getTask: extra =>
+            args
+                ? (handler as ToolTaskHandler<ZodRawShapeCompat>).getTask(args, extra)
+                : (handler as ToolTaskHandler<undefined>).getTask(extra),
+        getTaskResult: extra =>
+            args
+                ? (handler as ToolTaskHandler<ZodRawShapeCompat>).getTaskResult(args, extra)
+                : (handler as ToolTaskHandler<undefined>).getTaskResult(extra)
+    };
+}

test/server/mcp.test.ts

@@ -6355,6 +6355,11 @@ describe.each(zodTestMatrix)('$zodVersionLabel', (entry: ZodMatrixEntry) => {
             const taskStore = new InMemoryTaskStore();
             const { releaseLatch, waitForLatch } = createLatch();
 
+            // Spies to verify handler invocations
+            const createTaskSpy = vi.fn();
+            const getTaskSpy = vi.fn();
+            const getTaskResultSpy = vi.fn();
+
             const mcpServer = new McpServer(
                 {
                     name: 'test server',
@@ -6407,6 +6412,7 @@ describe.each(zodTestMatrix)('$zodVersionLabel', (entry: ZodMatrixEntry) => {
                 },
                 {
                     createTask: async ({ value }, extra) => {
+                        createTaskSpy({ value });
                         const task = await extra.taskStore.createTask({ ttl: 60000, pollInterval: 100 });
 
                         // Capture taskStore for use in setTimeout
@@ -6423,13 +6429,15 @@ describe.each(zodTestMatrix)('$zodVersionLabel', (entry: ZodMatrixEntry) => {
                         return { task };
                     },
                     getTask: async (_args, extra) => {
+                        getTaskSpy(extra.taskId);
                         const task = await extra.taskStore.getTask(extra.taskId);
                         if (!task) {
                             throw new Error('Task not found');
                         }
                         return task;
                     },
                     getTaskResult: async (_value, extra) => {
+                        getTaskResultSpy(extra.taskId);
                         const result = await extra.taskStore.getTaskResult(extra.taskId);
                         return result as CallToolResult;
                     }
@@ -6454,6 +6462,12 @@ describe.each(zodTestMatrix)('$zodVersionLabel', (entry: ZodMatrixEntry) => {
             expect(result.content).toEqual([{ type: 'text' as const, text: 'Result: 42' }]);
             expect(result).not.toHaveProperty('task');
 
+            // Verify all three handler methods were called
+            expect(createTaskSpy).toHaveBeenCalledOnce();
+            expect(createTaskSpy).toHaveBeenCalledWith({ value: 21 });
+            expect(getTaskSpy).toHaveBeenCalled(); // Called at least once during polling
+            expect(getTaskResultSpy).toHaveBeenCalledOnce();
+
             // Wait for async operations to complete
             await waitForLatch();
             taskStore.cleanup();