[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-SKILL.md

@@ -96,6 +96,8 @@ Notes:
 | `ErrorCode.RequestTimeout`               | `SdkErrorCode.RequestTimeout`                            |
 | `ErrorCode.ConnectionClosed`             | `SdkErrorCode.ConnectionClosed`                          |
 | `StreamableHTTPError`                    | REMOVED (use `SdkError` with `SdkErrorCode.ClientHttp*`) |
+| `ToolTaskHandler.getTask(args, ctx)`     | `ToolTaskHandler.getTask?(ctx)` — now optional, no args  |
+| `ToolTaskHandler.getTaskResult(args, ctx)` | `ToolTaskHandler.getTaskResult?(ctx)` — now optional, no args |
 
 All other symbols from `@modelcontextprotocol/sdk/types.js` retain their original names (e.g., `CallToolResultSchema`, `ListToolsResultSchema`, etc.).
 

docs/migration.md

@@ -478,6 +478,37 @@ import { JSONRPCError, ResourceReference, isJSONRPCError } from '@modelcontextpr
 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 | null;
 };
 
+/**
+ * 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;
 }