[v2] Make ToolTaskHandler.getTask/getTaskResult optional and actually invoke them

.changeset/remove-tooltask-get-handlers.md

@@ -0,0 +1,15 @@
+---
+"@modelcontextprotocol/core": minor
+"@modelcontextprotocol/server": minor
+---
+
+Make `ToolTaskHandler.getTask`/`getTaskResult` optional and actually invoke them
+
+**Bug fix:** `getTask` and `getTaskResult` handlers registered via `registerToolTask` were never invoked — `tasks/get` and `tasks/result` requests always hit `TaskStore` directly.
+
+**Breaking changes (experimental API):**
+
+- `getTask` and `getTaskResult` are now **optional** on `ToolTaskHandler`. When omitted, `TaskStore` handles the requests (previous de-facto behavior).
+- `TaskRequestHandler` signature changed: handlers receive only `(ctx: TaskServerContext)`, not the tool's input arguments.
+
+**Migration:** If your handlers just delegated to `ctx.task.store`, delete them. If you're proxying an external job system (Step Functions, CI/CD pipelines), keep them and drop the `args` parameter.

docs/migration.md

@@ -500,8 +500,39 @@
 
 ```typescript
 import { JSONRPCErrorResponse, ResourceTemplateReference, isJSONRPCErrorResponse } from '@modelcontextprotocol/server';
 ```
 
+### `ToolTaskHandler.getTask` and `getTaskResult` are now optional (experimental)
+
+`getTask` and `getTaskResult` are now optional on `ToolTaskHandler`. When omitted, `tasks/get` and `tasks/result` are served directly from the configured `TaskStore`. Their signature has also changed — they no longer receive the tool's input arguments (which aren't available at `tasks/get`/`tasks/result` time).
+
+If your handlers just delegated to the store, delete them:
+
+**Before:**
+
+```typescript
+server.experimental.tasks.registerToolTask('long-task', config, {
+    createTask: async (args, ctx) => { /* ... */ },
+    getTask: async (args, ctx) => ctx.task.store.getTask(ctx.task.id),
+    getTaskResult: async (args, ctx) => ctx.task.store.getTaskResult(ctx.task.id)
+});
+```
+
+**After:**
+
+```typescript
+server.experimental.tasks.registerToolTask('long-task', config, {
+    createTask: async (args, ctx) => { /* ... */ }
+});
+```
+
+Keep them if you're proxying an external job system (AWS Step Functions, CI/CD pipelines, etc.) — the new signature takes only `ctx`:
+
+```typescript
+getTask: async (ctx) => describeStepFunctionExecution(ctx.task.id),
+getTaskResult: async (ctx) => getStepFunctionOutput(ctx.task.id)
+```
+
 ### Request handler context types
 
 The `RequestHandlerExtra` type has been replaced with a structured context type hierarchy using nested groups:

examples/server/src/simpleStreamableHttp.ts

@@ -480,13 +480,6 @@ const getServer = () => {
                 return {
                     task
                 };
-            },
-            async getTask(_args, ctx) {
-                return await ctx.task.store.getTask(ctx.task.id);
-            },
-            async getTaskResult(_args, ctx) {
-                const result = await ctx.task.store.getTaskResult(ctx.task.id);
-                return result as CallToolResult;
             }
         }
     );
@@ -588,13 +581,6 @@ const getServer = () => {
                 })();
 
                 return { task };
-            },
-            async getTask(_args, ctx) {
-                return await ctx.task.store.getTask(ctx.task.id);
-            },
-            async getTaskResult(_args, ctx) {
-                const result = await ctx.task.store.getTaskResult(ctx.task.id);
-                return result as CallToolResult;
             }
         }
     );

packages/core/src/shared/taskManager.ts

@@ -154,6 +154,16 @@ export type TaskContext = {
     requestedTtl?: number;
 };
 
+/**
+ * Overrides for `tasks/get` and `tasks/result` lookups. Consulted before
+ * the configured {@linkcode TaskStore}; return `undefined` to fall through.
+ * @internal
+ */
+export type TaskLookupOverrides = {
+    getTask?: (taskId: string, ctx: BaseContext) => Promise<Task | undefined>;
+    getTaskResult?: (taskId: string, ctx: BaseContext) => Promise<Result | undefined>;
+};
+
 export type TaskManagerOptions = {
     /**
      * Task storage implementation. Required for handling incoming task requests (server-side).
@@ -199,20 +209,30 @@ export class TaskManager {
     private _requestResolvers: Map<RequestId, (response: JSONRPCResultResponse | Error) => void> = new Map();
     private _options: TaskManagerOptions;
     private _host?: TaskManagerHost;
+    private _overrides?: TaskLookupOverrides;
 
     constructor(options: TaskManagerOptions) {
         this._options = options;
         this._taskStore = options.taskStore;
         this._taskMessageQueue = options.taskMessageQueue;
     }
 
+    /**
+     * Installs per-task lookup overrides consulted before the {@linkcode TaskStore}.
+     * Used by McpServer to dispatch to per-tool `getTask`/`getTaskResult` handlers.
+     * @internal
+     */
+    setTaskOverrides(overrides: TaskLookupOverrides): void {
+        this._overrides = overrides;
+    }
+
     bind(host: TaskManagerHost): void {
         this._host = host;
 
         if (this._taskStore) {
             host.registerHandler('tasks/get', async (request, ctx) => {
                 const params = request.params as { taskId: string };
-                const task = await this.handleGetTask(params.taskId, ctx.sessionId);
+                const task = await this.handleGetTask(params.taskId, ctx);
                 // Per spec: tasks/get responses SHALL NOT include related-task metadata
                 // as the taskId parameter is the source of truth
                 return {
@@ -222,7 +242,7 @@ export class TaskManager {
 
             host.registerHandler('tasks/result', async (request, ctx) => {
                 const params = request.params as { taskId: string };
-                return await this.handleGetTaskPayload(params.taskId, ctx.sessionId, ctx.mcpReq.signal, async message => {
+                return await this.handleGetTaskPayload(params.taskId, ctx, async message => {
                     // Send the message on the response stream by passing the relatedRequestId
                     // This tells the transport to write the message to the tasks/result response stream
                     await host.sendOnResponseStream(message, ctx.mcpReq.id);
@@ -362,8 +382,11 @@ export class TaskManager {
 
     // -- Handler bodies (delegated from Protocol's registered handlers) --
 
-    private async handleGetTask(taskId: string, sessionId?: string): Promise<Task> {
-        const task = await this._requireTaskStore.getTask(taskId, sessionId);
+    private async handleGetTask(taskId: string, ctx: BaseContext): Promise<Task> {
+        const override = await this._overrides?.getTask?.(taskId, ctx);
+        if (override !== undefined) return override;
+
+        const task = await this._requireTaskStore.getTask(taskId, ctx.sessionId);
         if (!task) {
             throw new ProtocolError(ProtocolErrorCode.InvalidParams, 'Failed to retrieve task: Task not found');
         }
@@ -372,10 +395,12 @@ export class TaskManager {
 
     private async handleGetTaskPayload(
         taskId: string,
-        sessionId: string | undefined,
-        signal: AbortSignal,
+        ctx: BaseContext,
         sendOnResponseStream: (message: JSONRPCNotification | JSONRPCRequest) => Promise<void>
     ): Promise<Result> {
+        const sessionId = ctx.sessionId;
+        const signal = ctx.mcpReq.signal;
+
         const handleTaskResult = async (): Promise<Result> => {
             if (this._taskMessageQueue) {
                 let queuedMessage: QueuedMessage | undefined;
@@ -404,17 +429,15 @@ export class TaskManager {
                 }
             }
 
-            const task = await this._requireTaskStore.getTask(taskId, sessionId);
-            if (!task) {
-                throw new ProtocolError(ProtocolErrorCode.InvalidParams, `Task not found: ${taskId}`);
-            }
+            const task = await this.handleGetTask(taskId, ctx);
 
             if (!isTerminal(task.status)) {
                 await this._waitForTaskUpdate(task.pollInterval, signal);
                 return await handleTaskResult();
             }
 
-            const result = await this._requireTaskStore.getTaskResult(taskId, sessionId);
+            const override = await this._overrides?.getTaskResult?.(taskId, ctx);
+            const result = override ?? (await this._requireTaskStore.getTaskResult(taskId, sessionId));
             await this._clearTaskQueue(taskId);
 
             return {

packages/server/src/experimental/tasks/interfaces.ts

@@ -13,7 +13,7 @@ import type {
     TaskServerContext
 } from '@modelcontextprotocol/core';
 
-import type { BaseToolCallback } from '../../server/mcp.js';
+import type { AnyToolHandler, BaseToolCallback } from '../../server/mcp.js';
 
 // ============================================================================
 // Task Handler Types (for registerToolTask)
@@ -30,19 +30,27 @@ export type CreateTaskRequestHandler<
 
 /**
  * Handler for task operations (`get`, `getResult`).
+ *
+ * Receives only the context (no tool arguments — they are not available at
+ * `tasks/get` or `tasks/result` time). Access the task ID via `ctx.task.id`.
+ *
  * @experimental
  */
-export type TaskRequestHandler<SendResultT extends Result, Args extends StandardSchemaWithJSON | undefined = undefined> = BaseToolCallback<
-    SendResultT,
-    TaskServerContext,
-    Args
->;
+export type TaskRequestHandler<SendResultT extends Result> = (ctx: TaskServerContext) => SendResultT | Promise<SendResultT>;
 
 /**
  * Interface for task-based tool handlers.
  *
- * Task-based tools split a long-running operation into three phases:
- * `createTask`, `getTask`, and `getTaskResult`.
+ * Task-based tools create a task on `tools/call` and by default let the SDK's
+ * `TaskStore` handle subsequent `tasks/get` and `tasks/result` requests.
+ *
+ * Provide `getTask` and `getTaskResult` to override the default lookups — useful
+ * when proxying an external job system (e.g., AWS Step Functions, CI/CD pipelines)
+ * where the external system is the source of truth for task state.
+ *
+ * **Note:** the taskId → tool mapping used to dispatch `getTask`/`getTaskResult`
+ * is held in-memory and does not survive server restarts or span multiple
+ * instances. In those scenarios, requests fall through to the `TaskStore`.
  *
  * @see {@linkcode @modelcontextprotocol/server!experimental/tasks/mcpServer.ExperimentalMcpServerTasks#registerToolTask | registerToolTask} for registration.
  * @experimental
@@ -56,11 +64,23 @@ export interface ToolTaskHandler<Args extends StandardSchemaWithJSON | undefined
      */
     createTask: CreateTaskRequestHandler<CreateTaskResult, Args>;
     /**
-     * Handler for `tasks/get` requests.
+     * Optional handler for `tasks/get` requests. When omitted, the configured
+     * `TaskStore` is consulted directly.
      */
-    getTask: TaskRequestHandler<GetTaskResult, Args>;
+    getTask?: TaskRequestHandler<GetTaskResult>;
     /**
-     * Handler for `tasks/result` requests.
+     * Optional handler for `tasks/result` requests. When omitted, the configured
+     * `TaskStore` is consulted directly.
      */
-    getTaskResult: TaskRequestHandler<CallToolResult, Args>;
+    getTaskResult?: TaskRequestHandler<CallToolResult>;
+}
+
+/**
+ * Type guard for {@linkcode ToolTaskHandler}.
+ * @experimental
+ */
+export function isToolTaskHandler(
+    handler: AnyToolHandler<StandardSchemaWithJSON | undefined>
+): handler is ToolTaskHandler<StandardSchemaWithJSON | undefined> {
+    return 'createTask' in handler;
 }

packages/server/src/experimental/tasks/mcpServer.ts

@@ -5,13 +5,25 @@
  * @experimental
  */
 
-import type { StandardSchemaWithJSON, TaskToolExecution, ToolAnnotations, ToolExecution } from '@modelcontextprotocol/core';
+import type {
+    BaseContext,
+    CallToolResult,
+    GetTaskResult,
+    ServerContext,
+    StandardSchemaWithJSON,
+    TaskManager,
+    TaskServerContext,
+    TaskToolExecution,
+    ToolAnnotations,
+    ToolExecution
+} from '@modelcontextprotocol/core';
 
 import type { AnyToolHandler, McpServer, RegisteredTool } from '../../server/mcp.js';
 import type { ToolTaskHandler } from './interfaces.js';
+import { isToolTaskHandler } from './interfaces.js';
 
 /**
- * Internal interface for accessing {@linkcode McpServer}'s private _createRegisteredTool method.
+ * Internal interface for accessing {@linkcode McpServer}'s private members.
  * @internal
  */
 interface McpServerInternal {
@@ -26,6 +38,7 @@
         _meta: Record<string, unknown> | undefined,
         handler: AnyToolHandler<StandardSchemaWithJSON | undefined>
     ): RegisteredTool;
+    _registeredTools: { [name: string]: RegisteredTool };
 }
 
 /**
@@ -39,14 +52,74 @@
  * @experimental
  */
 export class ExperimentalMcpServerTasks {
+    /**
+     * Maps taskId → toolName for tasks whose handlers define custom
+     * `getTask` or `getTaskResult`. In-memory only; after a server restart
+     * or on a different instance, lookups fall through to the TaskStore.
+     */
+    private _taskToTool = new Map<string, string>();
+
     constructor(private readonly _mcpServer: McpServer) {}
 
+    /** @internal */
+    _installOverrides(taskManager: TaskManager): void {
+        taskManager.setTaskOverrides({
+            getTask: (taskId, ctx) => this._dispatch(taskId, ctx, 'getTask'),
+            getTaskResult: (taskId, ctx) => this._dispatch(taskId, ctx, 'getTaskResult')
+        });
+    }
+
+    /** @internal */
+    _recordTask(taskId: string, toolName: string): void {
+        const tool = (this._mcpServer as unknown as McpServerInternal)._registeredTools[toolName];
+        if (tool && isToolTaskHandler(tool.handler) && (tool.handler.getTask || tool.handler.getTaskResult)) {
+            this._taskToTool.set(taskId, toolName);
+        }
+    }
+
+    /** @internal */
+    onClose(): void {
+        this._taskToTool.clear();
+    }
+
+    private async _dispatch<M extends 'getTask' | 'getTaskResult'>(
+        taskId: string,
+        ctx: BaseContext,
+        method: M
+    ): Promise<(M extends 'getTask' ? GetTaskResult : CallToolResult) | undefined> {
+        const toolName = this._taskToTool.get(taskId);
+        if (!toolName) return undefined;
+
+        const tool = (this._mcpServer as unknown as McpServerInternal)._registeredTools[toolName];
+        if (!tool || !isToolTaskHandler(tool.handler)) return undefined;
+
+        const handler = tool.handler[method];
+        if (!handler) return undefined;
+
+        const serverCtx = ctx as ServerContext;
+        if (!serverCtx.task?.store) return undefined;
+
+        const taskCtx: TaskServerContext = {
+            ...serverCtx,
+            task: { ...serverCtx.task, id: taskId, store: serverCtx.task.store }
+        };
+
+        const result = (await handler(taskCtx)) as M extends 'getTask' ? GetTaskResult : CallToolResult;
+        // getTaskResult is terminal — drop the mapping only after the handler resolves
+        // so a transient throw doesn't orphan the task on retry.
+        if (method === 'getTaskResult') {
+            this._taskToTool.delete(taskId);
+        }
+        return result;
+    }
+
     /**
      * Registers a task-based tool with a config object and handler.
      *
      * Task-based tools support long-running operations that can be polled for status
-     * and results. The handler must implement {@linkcode ToolTaskHandler.createTask | createTask}, {@linkcode ToolTaskHandler.getTask | getTask}, and {@linkcode ToolTaskHandler.getTaskResult | getTaskResult}
-     * methods.
+     * and results. The handler implements {@linkcode ToolTaskHandler.createTask | createTask}
+     * to start the task; subsequent `tasks/get` and `tasks/result` requests are served
+     * from the configured `TaskStore`.
      *
      * @example
      * ```typescript
@@ -59,19 +132,13 @@
      *     const task = await ctx.task.store.createTask({ ttl: 300000 });
      *     startBackgroundWork(task.taskId, args);
      *     return { task };
-     *   },
-     *   getTask: async (args, ctx) => {
-     *     return ctx.task.store.getTask(ctx.task.id);
-     *   },
-     *   getTaskResult: async (args, ctx) => {
-     *     return ctx.task.store.getTaskResult(ctx.task.id);
      *   }
      * });
      * ```
      *
      * @param name - The tool name
      * @param config - Tool configuration (description, schemas, etc.)
-     * @param handler - Task handler with {@linkcode ToolTaskHandler.createTask | createTask}, {@linkcode ToolTaskHandler.getTask | getTask}, {@linkcode ToolTaskHandler.getTaskResult | getTaskResult} methods
+     * @param handler - Task handler with {@linkcode ToolTaskHandler.createTask | createTask}
      * @returns {@linkcode server/mcp.RegisteredTool | RegisteredTool} for managing the tool's lifecycle
      *
      * @experimental