Only delete _taskToolMap entry on success; switch test to fake timers; fix JSDoc

Author: felixweinberger

.changeset/fix-registertooltask-handlers.md

@@ -0,0 +1,5 @@
+---
+'@modelcontextprotocol/sdk': patch
+---
+
+Fix registerToolTask custom getTask/getTaskResult handlers not being invoked

src/server/mcp.ts

@@ -98,16 +98,15 @@ export class McpServer {
             getTaskResult: async (taskId: string, extra: RequestHandlerExtra<ServerRequest, ServerNotification>) => {
                 const taskStore = extra.taskStore!;
                 const handler = this._getTaskHandler(taskId);
-                try {
-                    if (handler) {
-                        return await handler.getTaskResult({ ...extra, taskId, taskStore });
-                    }
-                    return await taskStore.getTaskResult(taskId);
-                } finally {
-                    // Once the result has been retrieved the task is complete;
-                    // drop the taskId → toolName mapping to avoid unbounded growth.
-                    this._taskToolMap.delete(taskId);
-                }
+                const result = handler
+                    ? await handler.getTaskResult({ ...extra, taskId, taskStore })
+                    : await taskStore.getTaskResult(taskId);
+                // Once the result has been retrieved the task is complete;
+                // drop the taskId → toolName mapping to avoid unbounded growth.
+                // Cleanup on tasks/cancel would require a protocol-level hook and is
+                // intentionally left out here.
+                this._taskToolMap.delete(taskId);
+                return result;
             }
         };
         this.server = new Server(serverInfo, { ...options, taskHandlerHooks });

src/shared/protocol.ts

@@ -108,7 +108,7 @@ export type ProtocolOptions = {
          */
         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.
+         * Called when tasks/result needs to retrieve the final result. If provided, must return the result.
          */
         getTaskResult?: (taskId: string, extra: RequestHandlerExtra<Request, Notification>) => Promise<Result>;
     };

test/server/mcp.test.ts

@@ -6818,6 +6818,7 @@ describe.each(zodTestMatrix)('$zodVersionLabel', (entry: ZodMatrixEntry) => {
         });
 
         test('should call custom getTask and getTaskResult handlers when client polls task directly', async () => {
+            vi.useFakeTimers();
             const taskStore = new InMemoryTaskStore();
 
             const getTaskSpy = vi.fn();
@@ -6910,7 +6911,7 @@ describe.each(zodTestMatrix)('$zodVersionLabel', (entry: ZodMatrixEntry) => {
             const taskId = createResult.task.taskId;
 
             // Wait for task to be ready to complete
-            await new Promise(resolve => setTimeout(resolve, 60));
+            vi.advanceTimersByTime(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);
@@ -6922,6 +6923,7 @@ describe.each(zodTestMatrix)('$zodVersionLabel', (entry: ZodMatrixEntry) => {
             expect(payloadResult.content).toBeDefined();
             expect(getTaskResultSpy).toHaveBeenCalledWith(taskId);
 
+            vi.useRealTimers();
             taskStore.cleanup();
         });
     });