Introduce shim task interface for tools

Author: LucaButBoring

src/client/index.test.ts

@@ -2495,14 +2495,16 @@ test('should respect server task capabilities', async () => {
 
     // This should throw because server doesn't support task creation for tools/list
     await expect(
-        client.beginRequest(
-            {
-                method: 'tools/list',
-                params: {}
-            },
-            z.object({ tools: z.array(z.any()) }),
-            { task: { taskId: 'test-task-2', keepAlive: 60000 } }
-        ).result()
+        client
+            .beginRequest(
+                {
+                    method: 'tools/list',
+                    params: {}
+                },
+                z.object({ tools: z.array(z.any()) }),
+                { task: { taskId: 'test-task-2', keepAlive: 60000 } }
+            )
+            .result()
     ).rejects.toThrow('Server does not support task creation for tools/list');
 
     serverTaskStore.cleanup();

src/examples/server/simpleStreamableHttp.ts

@@ -454,7 +454,7 @@ const getServer = () => {
     );
 
     // Register a long-running tool that demonstrates task execution
-    server.registerTool(
+    server.registerToolTask(
         'delay',
         {
             title: 'Delay',
@@ -463,16 +463,37 @@ const getServer = () => {
                 duration: z.number().describe('Duration in milliseconds').default(5000)
             }
         },
-        async ({ duration }): Promise<CallToolResult> => {
-            await new Promise(resolve => setTimeout(resolve, duration));
-            return {
-                content: [
-                    {
-                        type: 'text',
-                        text: `Completed ${duration}ms delay`
-                    }
-                ]
-            };
+        {
+            async createTask({ duration }, { taskId, taskStore, taskRequestedKeepAlive }) {
+                // Simulate out-of-band work
+                (async () => {
+                    await new Promise(resolve => setTimeout(resolve, duration));
+                    await taskStore.storeTaskResult(taskId, {
+                        content: [
+                            {
+                                type: 'text',
+                                text: `Completed ${duration}ms delay`
+                            }
+                        ]
+                    });
+                })();
+                return await taskStore.createTask({
+                    taskId,
+                    keepAlive: taskRequestedKeepAlive
+                });
+            },
+            async getTask(_args, { taskId, taskStore }) {
+                const task = await taskStore.getTask(taskId);
+                if (!task) {
+                    throw new Error(`Task ${taskId} not found`);
+                }
+
+                return task;
+            },
+            async getTaskResult(_args, { taskId, taskStore }) {
+                const result = await taskStore.getTaskResult(taskId);
+                return result as CallToolResult;
+            }
         }
     );
 

src/examples/shared/inMemoryTaskStore.ts

@@ -21,7 +21,7 @@ export class InMemoryTaskStore implements TaskStore {
     private tasks = new Map<string, StoredTask>();
     private cleanupTimers = new Map<string, ReturnType<typeof setTimeout>>();
 
-    async createTask(metadata: TaskMetadata, requestId: RequestId, request: Request, _sessionId?: string): Promise<void> {
+    async createTask(metadata: TaskMetadata, requestId: RequestId, request: Request, _sessionId?: string): Promise<Task> {
         const taskId = metadata.taskId;
 
         if (this.tasks.has(taskId)) {
@@ -50,6 +50,8 @@ export class InMemoryTaskStore implements TaskStore {
 
             this.cleanupTimers.set(taskId, timer);
         }
+
+        return task;
     }
 
     async getTask(taskId: string, _sessionId?: string): Promise<Task | null> {

src/server/mcp.ts

@@ -31,15 +31,20 @@ import {
     ServerNotification,
     ToolAnnotations,
     LoggingMessageNotification,
+    CreateTaskResult,
+    GetTaskResult,
+    Result,
+    TASK_META_KEY,
     CompleteRequestPrompt,
     CompleteRequestResourceTemplate,
     assertCompleteRequestPrompt,
     assertCompleteRequestResourceTemplate
 } from '../types.js';
 import { Completable, CompletableDef } from './completable.js';
 import { UriTemplate, Variables } from '../shared/uriTemplate.js';
-import { RequestHandlerExtra } from '../shared/protocol.js';
+import { RequestHandlerExtra, RequestTaskStore } from '../shared/protocol.js';
 import { Transport } from '../shared/transport.js';
+import { isTerminal } from '../shared/task.js';
 
 /**
  * High-level MCP server that provides a simpler API for working with resources, tools, and prompts.
@@ -836,6 +841,51 @@ export class McpServer {
         );
     }
 
+    /**
+     * Registers a task-based tool with a config object and callback.
+     */
+    registerToolTask<InputArgs extends ZodRawShape, OutputArgs extends ZodRawShape>(
+        name: string,
+        config: {
+            title?: string;
+            description?: string;
+            inputSchema?: InputArgs;
+            outputSchema?: OutputArgs;
+            annotations?: ToolAnnotations;
+            _meta?: Record<string, unknown>;
+        },
+        handler: ToolTaskHandler<InputArgs>
+    ): RegisteredTool {
+        // TODO: Attach to individual request handlers and remove this wrapper
+        const cb: ToolCallback<InputArgs> = (async (...args) => {
+            const [inputArgs, extra] = args;
+
+            const taskStore = extra.taskStore;
+            if (!taskStore) {
+                throw new Error('Task store is not available');
+            }
+
+            const taskMetadata = extra._meta?.[TASK_META_KEY];
+            const taskId = taskMetadata?.taskId;
+            if (!taskId) {
+                throw new Error('No task ID provided');
+            }
+
+            // Internal polling to allow using this interface before internals are hooked up
+            const taskExtra = { ...extra, taskId, taskStore };
+            let task = await handler.createTask(inputArgs, taskExtra);
+            do {
+                await new Promise(resolve => setTimeout(resolve, task.pollInterval ?? 5000));
+                task = await handler.getTask(inputArgs, taskExtra);
+            } while (!isTerminal(task.status));
+
+            const result: CallToolResult = await handler.getTaskResult(inputArgs, taskExtra);
+            return result;
+        }) as ToolCallback<InputArgs>;
+
+        return this.registerTool(name, { ...config, annotations: { ...config.annotations, taskHint: true } }, cb);
+    }
+
     /**
      * Registers a zero-argument prompt `name`, which will run the given function when the client calls it.
      * @deprecated Use `registerPrompt` instead.
@@ -1044,6 +1094,21 @@ export type ToolCallback<Args extends undefined | ZodRawShape | ZodType<object>
       ? (args: T, extra: RequestHandlerExtra<ServerRequest, ServerNotification>) => CallToolResult | Promise<CallToolResult>
       : (extra: RequestHandlerExtra<ServerRequest, ServerNotification>) => CallToolResult | Promise<CallToolResult>;
 
+export interface TaskRequestHandlerExtra extends RequestHandlerExtra<ServerRequest, ServerNotification> {
+    taskId: string;
+    taskStore: RequestTaskStore;
+}
+
+export type TaskRequestHandler<SendResultT extends Result, Args extends undefined | ZodRawShape = undefined> = Args extends ZodRawShape
+    ? (args: z.objectOutputType<Args, ZodTypeAny>, extra: TaskRequestHandlerExtra) => SendResultT | Promise<SendResultT>
+    : (extra: TaskRequestHandlerExtra) => SendResultT | Promise<SendResultT>;
+
+export interface ToolTaskHandler<Args extends undefined | ZodRawShape = undefined> {
+    createTask: TaskRequestHandler<CreateTaskResult, Args>;
+    getTask: TaskRequestHandler<GetTaskResult, Args>;
+    getTaskResult: TaskRequestHandler<CallToolResult, Args>;
+}
+
 export type RegisteredTool = {
     title?: string;
     description?: string;

src/shared/protocol.test.ts

@@ -37,14 +37,14 @@ function createMockTaskStore(options?: {
     const tasks: Record<string, Task & { result?: Result }> = {};
     return {
         createTask: jest.fn((taskMetadata: TaskMetadata, _1: RequestId, _2: Request) => {
-            tasks[taskMetadata.taskId] = {
+            const task = (tasks[taskMetadata.taskId] = {
                 taskId: taskMetadata.taskId,
                 status: (taskMetadata.status as Task['status'] | undefined) ?? 'submitted',
                 keepAlive: taskMetadata.keepAlive ?? null,
                 pollInterval: (taskMetadata.pollInterval as Task['pollInterval'] | undefined) ?? 1000
-            };
+            });
             options?.onStatus?.('submitted');
-            return Promise.resolve();
+            return Promise.resolve(task);
         }),
         getTask: jest.fn((taskId: string) => {
             return Promise.resolve(tasks[taskId] ?? null);