wait_for_task
Wait for an ongoing MCP Task to complete, fail, or be canceled by specifying a task ID or batch ID. Set timeout and return options to manage monitoring of complex AI workflows effectively.
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 |
|---|---|---|---|
| batch_id | No | Wait for any task in this batch to complete (required if task_id not provided) | |
| return_all | No | For batch_id: return all completed tasks instead of just the first one (default: false) | |
| task_id | No | Wait for this specific task to complete (required if batch_id not provided) | |
| timeout_seconds | No | Maximum seconds to wait before timing out (default: 300, max: 600) |
Implementation Reference
- src/serve.ts:821-1074 (handler)Executes the wait_for_task tool by polling task status until completion, failure, cancellation, or timeout. Supports single task or batch waiting.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)Defines the input schema, description, and annotations for the wait_for_task MCP 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)Registers the wait_for_task tool by including it in the list returned for ListTools requests.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; });
- src/utils/task-manager.ts:261-263 (helper)Retrieves task information used by wait_for_task to poll status.public getTask(taskId: string): TaskInfo | undefined { return this.tasks.get(taskId); }
- src/utils/task-manager.ts:316-318 (helper)Retrieves all tasks used by wait_for_task for batch monitoring.public getAllTasks(): TaskInfo[] { return Array.from(this.tasks.values()); }