[v1.x] Fix registerToolTask's getTask and getTaskResult handlers not being invoked#1335
[v1.x] Fix registerToolTask's getTask and getTaskResult handlers not being invoked#1335LucaButBoring wants to merge 5 commits intomodelcontextprotocol:v1.xfrom
Conversation
|
commit: |
|
@claude review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| async getTask({ taskId, taskStore }) { | ||
| return await taskStore.getTask(taskId); | ||
| }, | ||
| async getTaskResult(_args, { taskId, taskStore }) { | ||
| async getTaskResult({ taskId, taskStore }) { |
There was a problem hiding this comment.
🔴 Bug: collect-user-info-task at lines 382-388 still uses the old two-argument signature (_args, { taskId, taskStore }) for getTask and getTaskResult, while the PR updated TaskRequestHandler to accept a single extra argument. This will cause a TypeError: Cannot destructure property 'taskId' of 'undefined' at runtime when these handlers are invoked, since the entire extra object is passed as the first argument and undefined as the second. The delay tool at lines 608-614 was correctly updated but this tool was missed.
Extended reasoning...
What the bug is
The PR changes TaskRequestHandler from a multi-argument type (BaseToolCallback<SendResultT, TaskRequestHandlerExtra, Args>) to a single-argument function type ((extra: TaskRequestHandlerExtra) => SendResultT | Promise<SendResultT>). This means getTask and getTaskResult handlers now receive a single object containing { taskId, taskStore, ...extra } instead of the old (args, extra) two-argument pattern.
The specific code path
In src/server/mcp.ts, the McpServer constructor sets up taskHandlerHooks that call handler.getTask({ ...extra, taskId, taskStore }) and handler.getTaskResult({ ...extra, taskId, taskStore }) with a single argument. When collect-user-info-task's handler at line 382 declares async getTask(_args, { taskId, taskStore: getTaskStore }), JavaScript passes the entire { ...extra, taskId, taskStore } object as _args and undefined as the second parameter.
Why existing code doesn't prevent it
The delay tool was correctly updated in this PR (lines 608-614 in the diff show the fix from (_args, { taskId, taskStore }) to ({ taskId, taskStore })), but the collect-user-info-task tool at lines 382-388 was overlooked. TypeScript should catch this as a compilation error since a 2-parameter function is not assignable to the new 1-parameter TaskRequestHandler type, but example files may not be compiled as part of CI.
Step-by-step proof
- Client sends
tools/callwithname: "collect-user-info-task"andtask: { ttl: 60000 } - Server creates the task via
createTask, returnsCreateTaskResult - Client later calls
tasks/getwith thetaskId Protocoldispatches totaskHandlerHooks.getTask(taskId, extra)McpServer's hook callshandler.getTask({ ...extra, taskId, taskStore })— one argument- The old handler signature
async getTask(_args, { taskId, taskStore: getTaskStore })receives:_args = { signal, sessionId, taskId, taskStore, ... }and the second parameter =undefined - Destructuring
{ taskId, taskStore: getTaskStore }fromundefinedthrowsTypeError: Cannot destructure property 'taskId' of 'undefined' - The server crashes with an unhandled exception
Impact
This is a runtime crash in the example server. Any client that creates a task using the collect-user-info-task tool and then polls its status will crash the server. The delay tool works correctly since it was updated.
Fix
Update lines 382-388 to use the new single-argument signature, matching the delay tool:
async getTask({ taskId, taskStore: getTaskStore }) {
return await getTaskStore.getTask(taskId);
},
async getTaskResult({ taskId, taskStore: getResultTaskStore }) {
const result = await getResultTaskStore.getTaskResult(taskId);
return result as CallToolResult;
}| private _registeredTools: { [name: string]: RegisteredTool } = {}; | ||
| private _registeredPrompts: { [name: string]: RegisteredPrompt } = {}; | ||
| private _experimental?: { tasks: ExperimentalMcpServerTasks }; | ||
| private _taskToolMap: Map<string, string> = new Map(); |
There was a problem hiding this comment.
🟡 _taskToolMap entries (taskId → toolName) are added at line 248 but never removed — there is no .delete() or .clear() call anywhere. For long-running servers processing many tasks, this map grows unboundedly even after tasks reach terminal states or expire via TTL. Consider adding cleanup when a task completes/fails/is cancelled, or lazily when _getTaskHandler finds a task no longer exists in the store.
Extended reasoning...
What the bug is
The _taskToolMap field (declared at line 85 as Map<string, string>) stores a mapping from taskId to the tool name that created it. Entries are added at line 248 via this._taskToolMap.set(taskResult.task.taskId, request.params.name) whenever a task-augmented tool call returns a CreateTaskResult. However, there is no corresponding .delete() call anywhere in the codebase — entries persist for the lifetime of the McpServer instance.
How it manifests
Every task creation adds a string → string entry to the map. When a task completes, fails, is cancelled, or expires via its TTL and gets cleaned up from the TaskStore, the corresponding _taskToolMap entry remains. Over time, for a server that processes many tasks, this map grows monotonically.
Step-by-step proof
- A client calls
tools/callwithtask: { ttl: 60000 }for a registered tool task. - The
CallToolRequestSchemahandler at line 243-249 executes:this._taskToolMap.set(taskResult.task.taskId, request.params.name). - The task completes —
TaskStore.storeTaskResult()is called, the task enters a terminal state. - The task's TTL expires and
InMemoryTaskStorecleans it up internally. - The
_taskToolMapstill holds thetaskId → toolNameentry. There is no code path that removes it. - Repeat steps 1-5 thousands of times — the map now holds thousands of stale entries.
Why existing code doesn't prevent it
Searching for all references to _taskToolMap reveals exactly three: the declaration (line 85), a .get() call (line 111), and the .set() call (line 248). No .delete(), .clear(), or any other cleanup mechanism exists.
Impact
Each entry is two short strings (taskId + toolName), so individual entries are small. For typical short-lived MCP server instances or low task throughput, this is unlikely to cause issues. However, for long-running servers processing a high volume of tasks (e.g., a persistent production server), memory usage will grow linearly and unboundedly over time.
Suggested fix
The simplest approach would be to add a lazy cleanup in _getTaskHandler: if the taskId is found in _taskToolMap but the task no longer exists in the TaskStore, delete the entry. Alternatively, cleanup could be added when a task reaches a terminal state (in the taskHandlerHooks or after storeTaskResult calls).
Previously, the code only called the underlying task store, and the tests were not complex enough to validate that the handlers were being called, so they missed this.
They weren't being populated correctly, and can't be without changing the TaskStore interface to require restoring the original request when retrieving a Task.
This removes the setTimeout logic we had in tests, which was masking an issue where the getTask handlers weren't being called. The appropriate logic has been moved into the getTask handlers themselves.
- Update collect-user-info-task example to use single-arg getTask/getTaskResult signature matching the updated TaskRequestHandler type - Clean up _taskToolMap entries after getTaskResult to prevent unbounded growth
felixweinberger
force-pushed
the
fix/registerToolTask-v1
branch
from
30f2a51 to
a4e24fb
Compare
|
Rebased onto v1.x and pushed a couple of fixes from the review:
Build and tests pass locally now. @LucaButBoring — when you get a chance, mind taking a look to make sure this still matches your intent? Since #1332 was superseded by #1764 on main with a different approach, want to confirm we're happy keeping this simpler fix on v1.x rather than backporting the larger refactor. |
3 participants
Note: This is the v1 backport of #1332.
This PR fixes a bug where custom
getTaskandgetTaskResulthandlers registered viaregisterToolTaskwere never invoked. TheProtocolclass's task handlers bypassed them entirely and usedTaskStoredirectly. This was a refactoring oversight that was missed due to (1) the existing tests not explicitly checking if those handlers were called, and (2)setTimeoutbeing used increateTaskin many tests inadvertently masking the issue.This also removes the argument-forwarding to
getTaskandgetTaskResult, as that was originally built before the currentTaskStoredesign was finalized, which broke the assumption that the original request would reliably be stored by the implementor. The currentTaskStoredesign allows theRequestto be saved, but does notrequirethat, and also exposes no way to directly retrieve it ingetTaskorgetTaskResult(it was possible but no longer intended at the time of the rewrite).getTaskandgetTaskResultnow only have theextraargument.Motivation and Context
When using
registerToolTask, developers could provide customgetTaskandgetTaskResulthandlers:These handlers were never invoked because:
Protocolclass'stasks/getandtasks/resulthandlers directly calledTaskStoreinstead of forwarding to the custom handlers.McpServer's backwards-compat polling wrapper also bypassed the custom handlerssetTimeoutto complete tasks and did not explicitly assert on the handlers being called, inadvertently masking the issue since tasks completed regardless of whether handlers were invokedHow Has This Been Tested?
Updated unit tests with stricter/more robust assertions.
Breaking Changes
Yes, due to
argsno longer being passed togetTaskorgetTaskResult. We could defer this part of the PR to v2.Types of changes
Checklist
Additional context