[v2] Make ToolTaskHandler.getTask/getTaskResult optional and actually invoke them#1764
[v2] Make ToolTaskHandler.getTask/getTaskResult optional and actually invoke them#1764felixweinberger wants to merge 2 commits intomainfrom
Conversation
🦋 Changeset detectedLatest commit: 607e0db The changes in this PR will be included in the next version bump. This PR includes changesets to release 5 packages
Not sure what this means? Click here to learn what changesets are. Click here if you're a maintainer who wants to add another changeset to this PR |
@modelcontextprotocol/client
@modelcontextprotocol/server
@modelcontextprotocol/express
@modelcontextprotocol/hono
@modelcontextprotocol/node
commit: |
…actually invoke them These handlers were defined on the interface but never invoked — three code paths bypassed them and called TaskStore directly: - TaskManager.handleGetTask - TaskManager.handleGetTaskPayload - McpServer.handleAutomaticTaskPolling The handlers exist to support proxying external job systems (AWS Step Functions, CI/CD pipelines) where the external system is the source of truth for task state. But every test/example implementation was pure boilerplate delegation to the store. This change makes the handlers optional: when omitted, TaskStore handles requests (zero boilerplate, previous de-facto behavior). When provided, they're invoked for tasks/get, tasks/result, and automatic polling. Dispatch logic is isolated in ExperimentalMcpServerTasks. Core gains a single setTaskOverrides() method on TaskManager — no changes to public TaskManagerOptions type. McpServer itself gains ~5 lines. Also drops the Args parameter from TaskRequestHandler since tool input arguments aren't available at tasks/get/tasks/result time. BREAKING CHANGE: TaskRequestHandler signature changed from (args, ctx) to (ctx). ToolTaskHandler.getTask and getTaskResult are now optional. Closes #1332 Co-authored-by: Luca Chang <lucalc@amazon.com>
felixweinberger
force-pushed
the
fweinberger/tooltask-handlers
branch
from
b106062 to
0f7ba7c
Compare
felixweinberger
changed the title
TypeDoc can't resolve {@linkcode TaskStore} from the server package
since core isn't in the processed package set.
|
@claude review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
There was a problem hiding this comment.
This PR introduces an override mechanism into core TaskManager, has breaking changes to the experimental task API, and touches 13 files across core/server/tests — worth a human review for the design decisions involved.
Extended reasoning...
Overview
This PR makes getTask and getTaskResult optional on ToolTaskHandler and actually wires them into the tasks/get and tasks/result request paths, fixing a bug where these handlers were defined but never invoked. The changes span 13 files: core TaskManager gains a new TaskLookupOverrides type and setTaskOverrides() method; ExperimentalMcpServerTasks gains ~50 lines of dispatch logic (_taskToTool map, _recordTask, _installOverrides, _dispatch); McpServer is modified to eagerly construct the experimental module and install overrides at construction time; tests remove ~180 lines of boilerplate handlers; and migration docs are updated.
Security risks
No direct security risks. The override mechanism is internal (@internal annotation) and only invoked by McpServer. The _dispatch method validates that the tool exists and has the expected handler shape before calling it. No auth, crypto, or permissions code is affected.
Level of scrutiny
This warrants careful human review for several reasons:
- Core infrastructure change:
TaskManagerin@modelcontextprotocol/coregains a new internal method that alters the lookup flow fortasks/getandtasks/result— this affects all servers using tasks. - Breaking API changes:
TaskRequestHandlersignature changes from(args, ctx)to(ctx), andgetTask/getTaskResultbecome optional. Even though the API is experimental, downstream users may be affected. - Design decisions: The override pattern (McpServer installing overrides into core TaskManager at construction time) is an architectural choice that should be validated by a maintainer.
- Major version bumps: The changeset triggers Major bumps for
node,express, andhonopackages.
Other factors
The PR has good test coverage with two new tests (handlers invoked when provided, TaskStore consulted when omitted) and all existing tests updated. The in-memory _taskToTool map limitation is documented. The code is well-structured with clear separation of concerns. The net LOC change is negative (removing boilerplate), which is positive. However, the scope and the core changes justify human sign-off.
| 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); | ||
| } |
There was a problem hiding this comment.
🟡 _taskToTool Map grows without bound: entries are added via _recordTask() but never deleted when tasks reach terminal status or when the server closes. For long-running servers using custom getTask/getTaskResult handlers, this is an unbounded memory leak. Consider deleting entries after returning a terminal result in _dispatch, and clearing the map on close (similar to how TaskManager._taskProgressTokens has _cleanupTaskProgressHandler() and onClose() cleanup).
Extended reasoning...
What the bug is
The _taskToTool Map (line 60) in ExperimentalMcpServerTasks records taskId -> toolName mappings whenever a task is created for a tool that has custom getTask or getTaskResult handlers. Entries are added in _recordTask() but there are zero calls to _taskToTool.delete() or _taskToTool.clear() anywhere in the class or codebase.
Code path
The flow is:
- A client calls
tools/callwith task augmentation.McpServercallsthis._experimental.tasks._recordTask(taskId, toolName)(mcp.ts lines 200 and 322). _recordTask()checks if the tool has customgetTask/getTaskResulthandlers and, if so, stores the mapping:this._taskToTool.set(taskId, toolName)(mcpServer.ts line 77).- When
tasks/getortasks/resultarrives,_dispatch()reads from the map to route to the correct handler. - No code ever removes the entry -- not after a terminal result is returned, not when the task is cancelled, and not when the server shuts down.
Why existing code does not prevent it
The TaskManager class in core has cleanup patterns for its own maps: _cleanupTaskProgressHandler() deletes entries from _taskProgressTokens when a task reaches terminal status, and onClose() clears the entire map. The ExperimentalMcpServerTasks class follows none of these patterns -- it has no onClose() hook and no per-task cleanup.
Step-by-step proof
- Server starts,
_taskToToolis empty (size 0). - Client creates task T1 for tool
proxy-task(which has customgetTask/getTaskResult)._recordTask("T1", "proxy-task")runs. Map size = 1. - T1 completes.
_dispatchreturns the result. Map still has entry"T1" -> "proxy-task"-- size = 1. - Client creates task T2 for the same tool. Map size = 2.
- T2 completes. Map size = 2.
- After N tasks, map size = N. The entries for T1, T2, etc. are never removed despite those tasks being long finished.
Impact
Each entry is just two short strings (~100 bytes), so the leak is slow. For the common case (no custom handlers), _recordTask skips the .set() call entirely, so there is no leak. The leak only manifests for the uncommon proxy use case (custom getTask/getTaskResult handlers) on a long-running server that processes many tasks. The PR description already acknowledges the in-memory-only nature of this map, but the concern there was about cross-instance behavior -- within a single process the map grows monotonically.
Suggested fix
Delete the entry in _dispatch after a terminal result is returned (when the method is getTaskResult, since that is the final lookup). Also add an onClose() method that calls this._taskToTool.clear() and wire it into the server shutdown path. This mirrors the _taskProgressTokens cleanup pattern in TaskManager.
1 participant
Makes
getTaskandgetTaskResultoptional onToolTaskHandlerand wires them into thetasks/getandtasks/resultrequest paths. Supersedes #1332.Motivation and Context
These handlers were defined on the interface but never invoked — three code paths bypassed them and hit
TaskStoredirectly:TaskManager.handleGetTask→_requireTaskStore.getTaskTaskManager.handleGetTaskPayload→_requireTaskStore.getTask/getTaskResultMcpServer.handleAutomaticTaskPolling→ctx.task.storeThe handlers exist to support proxying external job systems (AWS Step Functions, CI/CD pipelines, etc.) where the external system is the source of truth for task state. But every test and example implementation was pure boilerplate delegation to
ctx.task.store, which made the handlers look vestigial.This PR makes them optional:
TaskStorehandles everything. Zero boilerplate.tasks/get,tasks/result, and automatic polling. Useful for external-system proxies.Isolation
All dispatch logic lives in
ExperimentalMcpServerTasks(the experimental module that already ownsregisterToolTask):_taskToToolmap,_recordTask(),_installOverrides(),_dispatch()— ~50 lines, all contained thereTaskManager.setTaskOverrides()+ two lookup checks. No changes to the publicTaskManagerOptionstype.How Has This Been Tested?
tasks/getandtasks/resultTaskStoreis consulted directly when handlers are omittedBreaking Changes
Yes (experimental API):
TaskRequestHandlersignature changed from(args, ctx)to(ctx)— tool arguments aren't available attasks/get/tasks/resulttimegetTaskandgetTaskResultare now optional onToolTaskHandlerSee
docs/migration.md.Types of changes
Checklist
Additional context
Known limitation: the taskId → tool mapping is in-memory and doesn't survive restarts or span multiple server instances. In those scenarios, requests fall through to
TaskStore. Documented on theToolTaskHandlerinterface.Also adds an
isToolTaskHandlertype guard replacing inline'createTask' in handlerchecks.Closes #1332