wait_for_task
Monitor task completion in MCP Task server by waiting for specific tasks or batches to finish, fail, or cancel with configurable timeout options.
Instructions
Wait for a task or any task in a batch to complete, fail, or be cancelled. Only waits for tasks that complete AFTER this call is made - ignores tasks that were already completed.
Input Schema
TableJSON Schema
| Name | Required | Description | Default |
|---|---|---|---|
| task_id | No | Wait for this specific task to complete (required if batch_id not provided) | |
| batch_id | No | Wait for any task in this batch to complete (required if task_id not provided) | |
| timeout_seconds | No | Maximum seconds to wait before timing out (default: 300, max: 600) | |
| return_all | No | For batch_id: return all completed tasks instead of just the first one (default: false) |
Implementation Reference
- src/serve.ts:821-1074 (handler)The core handler implementation for the 'wait_for_task' tool. It polls the TaskManager either for a specific task_id or any new completions in a batch_id, with configurable timeout and option to wait for all or first completion. Ignores tasks already complete at call time. Returns JSON status and output.
case 'wait_for_task': { // Validate that at least one parameter is provided if (!args.task_id && !args.batch_id) { throw new Error('Either task_id or batch_id is required'); } if (args.task_id && args.batch_id) { throw new Error( 'Provide either task_id or batch_id, not both' ); } const timeoutSeconds = Math.min( args.timeout_seconds || 300, 600 ); const returnAll = args.return_all || false; const startTime = Date.now(); const timeoutMs = timeoutSeconds * 1000; const pollIntervalMs = 1000; // Poll every second // Helper function to check if a task is complete const isTaskComplete = (task: any) => { return ( task.status === 'completed' || task.status === 'failed' || task.status === 'cancelled' ); }; // Wait for single task if (args.task_id) { while (Date.now() - startTime < timeoutMs) { const task = taskManager.getTask(args.task_id); if (!task) { throw new Error(`Task ${args.task_id} not found`); } if (isTaskComplete(task)) { return { content: [ { type: 'text', text: JSON.stringify( { task_id: task.id, status: task.status, output: task.output, completed_at: task.completedAt, wait_time_seconds: Math.round( (Date.now() - startTime) / 1000 ), }, null, 2 ), }, ], }; } // Wait before next check await new Promise(resolve => setTimeout(resolve, pollIntervalMs) ); } // Timeout reached const task = taskManager.getTask(args.task_id); return { content: [ { type: 'text', text: JSON.stringify( { task_id: args.task_id, status: task?.status || 'unknown', timeout: true, message: `Timeout after ${timeoutSeconds} seconds waiting for task ${args.task_id}`, }, null, 2 ), }, ], }; } // Wait for batch tasks if (args.batch_id) { // Track which tasks were already completed when wait started const initialTasks = taskManager.getAllTasks(); const initialBatchTasks = initialTasks.filter( t => t.batchId === args.batch_id ); if (initialBatchTasks.length === 0) { throw new Error( `No tasks found with batch_id: ${args.batch_id}` ); } // Record which tasks were already completed at the start const alreadyCompletedIds = new Set( initialBatchTasks.filter(isTaskComplete).map(t => t.id) ); const completedTasks: any[] = []; while (Date.now() - startTime < timeoutMs) { const allTasks = taskManager.getAllTasks(); const batchTasks = allTasks.filter( t => t.batchId === args.batch_id ); if (batchTasks.length === 0) { throw new Error( `No tasks found with batch_id: ${args.batch_id}` ); } // Check for newly completed tasks (not already completed when wait started) const newlyCompleted = batchTasks.filter( t => isTaskComplete(t) && !alreadyCompletedIds.has(t.id) && !completedTasks.find(ct => ct.id === t.id) ); completedTasks.push(...newlyCompleted); // If we want all tasks, check if all are complete if (returnAll) { const allComplete = batchTasks.every(isTaskComplete); if (allComplete) { return { content: [ { type: 'text', text: JSON.stringify( { batch_id: args.batch_id, all_complete: true, tasks: batchTasks.map( t => ({ id: t.id, status: t.status, output: t.output, completed_at: t.completedAt, }) ), wait_time_seconds: Math.round( (Date.now() - startTime) / 1000 ), }, null, 2 ), }, ], }; } } else { // Return first newly completed task if (completedTasks.length > 0) { const firstCompleted = completedTasks[0]; return { content: [ { type: 'text', text: JSON.stringify( { batch_id: args.batch_id, newly_completed: true, task_id: firstCompleted.id, status: firstCompleted.status, output: firstCompleted.output, completed_at: firstCompleted.completedAt, remaining_tasks: batchTasks.filter( t => !isTaskComplete( t ) ).length, previously_completed: alreadyCompletedIds.size, wait_time_seconds: Math.round( (Date.now() - startTime) / 1000 ), }, null, 2 ), }, ], }; } } // Wait before next check await new Promise(resolve => setTimeout(resolve, pollIntervalMs) ); } // Timeout reached const allTasks = taskManager.getAllTasks(); const batchTasks = allTasks.filter( t => t.batchId === args.batch_id ); const runningCount = batchTasks.filter( t => t.status === 'running' ).length; const pendingCount = batchTasks.filter( t => t.status === 'pending' ).length; return { content: [ { type: 'text', text: JSON.stringify( { batch_id: args.batch_id, timeout: true, message: `Timeout after ${timeoutSeconds} seconds`, completed_tasks: completedTasks.length, running_tasks: runningCount, pending_tasks: pendingCount, total_tasks: batchTasks.length, }, null, 2 ), }, ], }; } // Should never reach here throw new Error('Invalid wait_for_task parameters'); } - src/serve.ts:244-284 (schema)The Tool object definition including name, description, annotations, and detailed inputSchema for the wait_for_task tool.
const WAIT_FOR_TASK_TOOL: Tool = { name: 'wait_for_task', description: 'Wait for a task or any task in a batch to complete, fail, or be cancelled. Only waits for tasks that complete AFTER this call is made - ignores tasks that were already completed.', annotations: { title: 'Wait For Task Completion', readOnlyHint: true, // Only reads task status destructiveHint: false, // Doesn't modify or destroy data idempotentHint: true, // Waiting multiple times is safe openWorldHint: false, // Only queries local task state }, inputSchema: { type: 'object', properties: { task_id: { type: 'string', description: 'Wait for this specific task to complete (required if batch_id not provided)', }, batch_id: { type: 'string', description: 'Wait for any task in this batch to complete (required if task_id not provided)', }, timeout_seconds: { type: 'number', description: 'Maximum seconds to wait before timing out (default: 300, max: 600)', default: 300, maximum: 600, }, return_all: { type: 'boolean', description: 'For batch_id: return all completed tasks instead of just the first one (default: false)', default: false, }, }, required: [], // At least one must be provided, validated in handler }, }; - src/serve.ts:558-579 (registration)Registration of the wait_for_task tool as part of the list of available tools served in response to ListToolsRequestSchema.
server.setRequestHandler(ListToolsRequestSchema, async () => { if (process.env.MCP_MODE !== 'true') { logger.debug('Received ListTools request'); } const response = { tools: [ RUN_TASK_TOOL, CHECK_TASK_STATUS_TOOL, GET_TASK_RESULT_TOOL, CANCEL_TASK_TOOL, WAIT_FOR_TASK_TOOL, LIST_TASKS_TOOL, ], }; if (process.env.MCP_MODE !== 'true') { logger.debug( 'Returning tools:', response.tools.map(t => t.name) ); } return response; });