-
Notifications
You must be signed in to change notification settings - Fork 1.7k
[v1.x] Fix registerToolTask's getTask and getTaskResult handlers not being invoked #1335
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: v1.x
Are you sure you want to change the base?
Changes from all commits
7503353
fbe5df4
8ae7217
fd14b33
a4e24fb
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -82,9 +82,43 @@ export class McpServer { | |
| 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. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. 🟡 Extended reasoning...What the bug isThe How it manifestsEvery task creation adds a Step-by-step proof
Why existing code doesn't prevent itSearching for all references to ImpactEach 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 fixThe simplest approach would be to add a lazy cleanup in There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. The cleanup added in the The The original suggestion of lazy cleanup in |
||
|
|
||
| constructor(serverInfo: Implementation, options?: ServerOptions) { | ||
| this.server = new Server(serverInfo, options); | ||
| const taskHandlerHooks = { | ||
| getTask: async (taskId: string, extra: RequestHandlerExtra<ServerRequest, ServerNotification>) => { | ||
| // taskStore is guaranteed to exist here because Protocol only calls hooks when taskStore is configured | ||
| const taskStore = extra.taskStore!; | ||
| const handler = this._getTaskHandler(taskId); | ||
| if (handler) { | ||
| return await handler.getTask({ ...extra, taskId, taskStore }); | ||
| } | ||
| return await taskStore.getTask(taskId); | ||
| }, | ||
| getTaskResult: async (taskId: string, extra: RequestHandlerExtra<ServerRequest, ServerNotification>) => { | ||
| const taskStore = extra.taskStore!; | ||
| const handler = this._getTaskHandler(taskId); | ||
| try { | ||
| if (handler) { | ||
| return await handler.getTaskResult({ ...extra, taskId, taskStore }); | ||
| } | ||
| return await taskStore.getTaskResult(taskId); | ||
| } finally { | ||
| // Once the result has been retrieved the task is complete; | ||
| // drop the taskId → toolName mapping to avoid unbounded growth. | ||
| this._taskToolMap.delete(taskId); | ||
| } | ||
| } | ||
| }; | ||
| this.server = new Server(serverInfo, { ...options, taskHandlerHooks }); | ||
|
Comment on lines
+85
to
+113
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. 🔴 The Extended reasoning...Gap 1 — Premature deletion on error (bug_003)In the Step-by-step proof:
The fix is straightforward: use a success flag (or move the Gap 2 — Cancel path never cleans up (bug_001)The Step-by-step proof:
Individual entries are small (two short strings), but on long-running servers handling many task cancellations, the map grows without bound. The fix requires either adding a Gap 3 — Auto-polling path never populates the map (bug_004)In Step-by-step proof:
Addressing the refutation: The refuter argues this scenario is implausible because clients using the |
||
| } | ||
|
|
||
| private _getTaskHandler(taskId: string): ToolTaskHandler<ZodRawShapeCompat | undefined> | null { | ||
| const toolName = this._taskToolMap.get(taskId); | ||
| if (!toolName) return null; | ||
| const tool = this._registeredTools[toolName]; | ||
| if (!tool || !('createTask' in (tool.handler as AnyToolHandler<ZodRawShapeCompat>))) return null; | ||
| return tool.handler as ToolTaskHandler<ZodRawShapeCompat | undefined>; | ||
| } | ||
|
|
||
| /** | ||
|
|
@@ -215,6 +249,10 @@ export class McpServer { | |
|
|
||
| // Return CreateTaskResult immediately for task requests | ||
| if (isTaskRequest) { | ||
| const taskResult = result as CreateTaskResult; | ||
| if (taskResult.task?.taskId) { | ||
| this._taskToolMap.set(taskResult.task.taskId, request.params.name); | ||
| } | ||
| return result; | ||
| } | ||
|
|
||
|
|
@@ -374,27 +412,28 @@ export class McpServer { | |
| const handler = tool.handler as ToolTaskHandler<ZodRawShapeCompat | undefined>; | ||
| const taskExtra = { ...extra, taskStore: extra.taskStore }; | ||
|
|
||
| const createTaskResult: CreateTaskResult = args // undefined only if tool.inputSchema is undefined | ||
| ? await Promise.resolve((handler as ToolTaskHandler<ZodRawShapeCompat>).createTask(args, taskExtra)) | ||
| : // eslint-disable-next-line @typescript-eslint/no-explicit-any | ||
| await Promise.resolve(((handler as ToolTaskHandler<undefined>).createTask as any)(taskExtra)); | ||
| const wrappedHandler = toolTaskHandlerByArgs(handler, args); | ||
|
|
||
| const createTaskResult = await wrappedHandler.createTask(taskExtra); | ||
|
|
||
| // Poll until completion | ||
| const taskId = createTaskResult.task.taskId; | ||
| const taskExtraComplete = { ...extra, taskId, taskStore: extra.taskStore }; | ||
| let task = createTaskResult.task; | ||
| const pollInterval = task.pollInterval ?? 5000; | ||
|
|
||
| while (task.status !== 'completed' && task.status !== 'failed' && task.status !== 'cancelled') { | ||
| await new Promise(resolve => setTimeout(resolve, pollInterval)); | ||
| const updatedTask = await extra.taskStore.getTask(taskId); | ||
| const getTaskResult = await wrappedHandler.getTask(taskExtraComplete); | ||
| const updatedTask = getTaskResult; | ||
| if (!updatedTask) { | ||
| throw new McpError(ErrorCode.InternalError, `Task ${taskId} not found during polling`); | ||
| } | ||
| task = updatedTask; | ||
| } | ||
|
|
||
| // Return the final result | ||
| return (await extra.taskStore.getTaskResult(taskId)) as CallToolResult; | ||
| return await wrappedHandler.getTaskResult(taskExtraComplete); | ||
| } | ||
|
|
||
| private _completionHandlerInitialized = false; | ||
|
|
@@ -1545,3 +1584,24 @@ const EMPTY_COMPLETION_RESULT: CompleteResult = { | |
| hasMore: false | ||
| } | ||
| }; | ||
|
|
||
| /** | ||
| * Wraps a tool task handler's createTask to handle args uniformly. | ||
| * getTask and getTaskResult don't take args, so they're passed through directly. | ||
| * @param handler The task handler to wrap. | ||
| * @param args The tool arguments. | ||
| * @returns A wrapped task handler for a tool, which only exposes a no-args interface for createTask. | ||
| */ | ||
| function toolTaskHandlerByArgs<Args extends AnySchema | ZodRawShapeCompat | undefined>( | ||
| handler: ToolTaskHandler<Args>, | ||
| args: unknown | ||
| ): ToolTaskHandler<undefined> { | ||
| return { | ||
| createTask: extra => | ||
| args // undefined only if tool.inputSchema is undefined | ||
| ? Promise.resolve((handler as ToolTaskHandler<ZodRawShapeCompat>).createTask(args, extra)) | ||
| : Promise.resolve((handler as ToolTaskHandler<undefined>).createTask(extra)), | ||
| getTask: handler.getTask, | ||
| getTaskResult: handler.getTaskResult | ||
| }; | ||
| } | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -98,6 +98,20 @@ export type ProtocolOptions = { | |
| * appropriately (e.g., by failing the task, dropping messages, etc.). | ||
| */ | ||
| maxTaskQueueSize?: number; | ||
| /** | ||
| * Optional hooks for customizing task request handling. | ||
| * If a hook is provided, it fully owns the behavior (no fallback to TaskStore). | ||
| */ | ||
| taskHandlerHooks?: { | ||
| /** | ||
| * Called when tasks/get is received. If provided, must return the task. | ||
| */ | ||
| getTask?: (taskId: string, extra: RequestHandlerExtra<Request, Notification>) => Promise<GetTaskResult>; | ||
| /** | ||
| * Called when tasks/payload needs to retrieve the final result. If provided, must return the result. | ||
| */ | ||
| getTaskResult?: (taskId: string, extra: RequestHandlerExtra<Request, Notification>) => Promise<Result>; | ||
|
Comment on lines
+110
to
+113
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. 🟡 The JSDoc comment for Extended reasoning...What the bug isIn /**
* Called when tasks/payload needs to retrieve the final result. If provided, must return the result.
*/
getTaskResult?: (taskId: string, extra: RequestHandlerExtra<Request, Notification>) => Promise<Result>;The method name The specific code pathThe Why existing code does not prevent itThis is a purely documentation error in a JSDoc comment. TypeScript does not validate comment text, so no compilation or test failure would catch this. The comment was introduced as a new ImpactThis is a documentation-only issue with no runtime impact. However, any developer reading the How to fix itChange the comment from /**
* Called when tasks/result is received to retrieve the final result. If provided, must return the result.
*/Step-by-step proof
|
||
| }; | ||
| }; | ||
|
|
||
| /** | ||
|
|
@@ -383,6 +397,16 @@ export abstract class Protocol<SendRequestT extends Request, SendNotificationT e | |
| this._taskMessageQueue = _options?.taskMessageQueue; | ||
| if (this._taskStore) { | ||
| this.setRequestHandler(GetTaskRequestSchema, async (request, extra) => { | ||
| // Use hook if provided, otherwise fall back to TaskStore | ||
| if (_options?.taskHandlerHooks?.getTask) { | ||
| const hookResult = await _options.taskHandlerHooks.getTask( | ||
| request.params.taskId, | ||
| extra as unknown as RequestHandlerExtra<Request, Notification> | ||
| ); | ||
| // @ts-expect-error SendResultT cannot contain GetTaskResult | ||
| return hookResult as SendResultT; | ||
| } | ||
|
|
||
| const task = await this._taskStore!.getTask(request.params.taskId, extra.sessionId); | ||
| if (!task) { | ||
| throw new McpError(ErrorCode.InvalidParams, 'Failed to retrieve task: Task not found'); | ||
|
|
@@ -462,7 +486,13 @@ export abstract class Protocol<SendRequestT extends Request, SendNotificationT e | |
|
|
||
| // If task is terminal, return the result | ||
| if (isTerminal(task.status)) { | ||
| const result = await this._taskStore!.getTaskResult(taskId, extra.sessionId); | ||
| // Use hook if provided, otherwise fall back to TaskStore | ||
| const result = this._options?.taskHandlerHooks?.getTaskResult | ||
| ? await this._options.taskHandlerHooks.getTaskResult( | ||
| taskId, | ||
| extra as unknown as RequestHandlerExtra<Request, Notification> | ||
| ) | ||
| : await this._taskStore!.getTaskResult(taskId, extra.sessionId); | ||
|
|
||
| this._clearTaskQueue(taskId); | ||
|
|
||
|
|
||
Uh oh!
There was an error while loading. Please reload this page.