@sethdouglasford/claude-flow
Version:
Claude Code Flow - Advanced AI-powered development workflows with SPARC methodology
318 lines • 11.6 kB
JavaScript
/**
* Task scheduler implementation
*/
import { SystemEvents } from "../utils/types.js";
import { TaskError, TaskTimeoutError, TaskDependencyError } from "../utils/errors.js";
/**
* Task scheduler for managing task assignment and execution
*/
export class TaskScheduler {
config;
eventBus;
logger;
tasks = new Map();
agentTasks = new Map(); // agentId -> taskIds
taskDependencies = new Map(); // taskId -> dependent taskIds
completedTasks = new Set();
constructor(config, eventBus, logger) {
this.config = config;
this.eventBus = eventBus;
this.logger = logger;
}
initialize() {
this.logger.info("Initializing task scheduler");
// Set up periodic cleanup
setInterval(() => this.cleanup(), 60000); // Every minute
return Promise.resolve();
}
async shutdown() {
this.logger.info("Shutting down task scheduler");
// Cancel all active tasks
const taskIds = Array.from(this.tasks.keys());
await Promise.all(taskIds.map(id => this.cancelTask(id, "Scheduler shutdown")));
this.tasks.clear();
this.agentTasks.clear();
this.taskDependencies.clear();
this.completedTasks.clear();
}
assignTask(task, agentId) {
this.logger.info("Assigning task", { taskId: task.id, agentId });
// Check dependencies
if (task.dependencies.length > 0) {
const unmetDependencies = task.dependencies.filter(depId => !this.completedTasks.has(depId));
if (unmetDependencies.length > 0) {
throw new TaskDependencyError(task.id, unmetDependencies);
}
}
// Create scheduled task
const scheduledTask = {
task: { ...task, status: "assigned", assignedAgent: agentId },
agentId,
attempts: 0,
};
// Store task
this.tasks.set(task.id, scheduledTask);
// Update agent tasks
let agentTaskSet = this.agentTasks.get(agentId);
if (!agentTaskSet) {
agentTaskSet = new Set();
this.agentTasks.set(agentId, agentTaskSet);
}
agentTaskSet.add(task.id);
// Update dependencies
for (const depId of task.dependencies) {
let depSet = this.taskDependencies.get(depId);
if (!depSet) {
depSet = new Set();
this.taskDependencies.set(depId, depSet);
}
depSet.add(task.id);
}
// Start task execution
this.startTask(task.id);
return Promise.resolve();
}
completeTask(taskId, result) {
const scheduled = this.tasks.get(taskId);
if (!scheduled) {
throw new TaskError(`Task not found: ${taskId}`);
}
this.logger.info("Task completed", { taskId, agentId: scheduled.agentId });
// Update task status
scheduled.task.status = "completed";
scheduled.task.output = result;
scheduled.task.completedAt = new Date();
// Clear timeout
if (scheduled.timeout) {
clearTimeout(scheduled.timeout);
}
// Remove from active tasks
this.tasks.delete(taskId);
this.agentTasks.get(scheduled.agentId)?.delete(taskId);
// Add to completed tasks
this.completedTasks.add(taskId);
// Check and start dependent tasks
const dependents = this.taskDependencies.get(taskId);
if (dependents) {
for (const dependentId of dependents) {
const dependent = this.tasks.get(dependentId);
if (dependent && this.canStartTask(dependent.task)) {
this.startTask(dependentId);
}
}
}
return Promise.resolve();
}
async failTask(taskId, error) {
const scheduled = this.tasks.get(taskId);
if (!scheduled) {
throw new TaskError(`Task not found: ${taskId}`);
}
this.logger.error("Task failed", {
taskId,
agentId: scheduled.agentId,
attempt: scheduled.attempts,
error,
});
// Clear timeout
if (scheduled.timeout) {
clearTimeout(scheduled.timeout);
}
scheduled.attempts++;
scheduled.lastAttempt = new Date();
// Check if we should retry
if (scheduled.attempts < this.config.maxRetries) {
this.logger.info("Retrying task", {
taskId,
attempt: scheduled.attempts,
maxRetries: this.config.maxRetries,
});
// Schedule retry with exponential backoff
const retryDelay = this.config.retryDelay * Math.pow(2, scheduled.attempts - 1);
setTimeout(() => {
this.startTask(taskId);
}, retryDelay);
}
else {
// Max retries exceeded, mark as failed
scheduled.task.status = "failed";
scheduled.task.error = error;
scheduled.task.completedAt = new Date();
// Remove from active tasks
this.tasks.delete(taskId);
this.agentTasks.get(scheduled.agentId)?.delete(taskId);
// Cancel dependent tasks
await this.cancelDependentTasks(taskId, "Parent task failed");
}
}
async cancelTask(taskId, reason) {
const scheduled = this.tasks.get(taskId);
if (!scheduled) {
return; // Already cancelled or completed
}
this.logger.info("Cancelling task", { taskId, reason });
// Clear timeout
if (scheduled.timeout) {
clearTimeout(scheduled.timeout);
}
// Update task status
scheduled.task.status = "cancelled";
scheduled.task.completedAt = new Date();
// Emit cancellation event
this.eventBus.emit(SystemEvents.TASK_CANCELLED, { taskId, reason });
// Remove from active tasks
this.tasks.delete(taskId);
this.agentTasks.get(scheduled.agentId)?.delete(taskId);
// Cancel dependent tasks
await this.cancelDependentTasks(taskId, "Parent task cancelled");
}
async cancelAgentTasks(agentId) {
const taskIds = this.agentTasks.get(agentId);
if (!taskIds) {
return;
}
this.logger.info("Cancelling all tasks for agent", {
agentId,
taskCount: taskIds.size,
});
const promises = Array.from(taskIds).map(taskId => this.cancelTask(taskId, "Agent terminated"));
await Promise.all(promises);
this.agentTasks.delete(agentId);
}
rescheduleAgentTasks(agentId) {
const taskIds = this.agentTasks.get(agentId);
if (!taskIds || taskIds.size === 0) {
return Promise.resolve();
}
this.logger.info("Rescheduling tasks for agent", {
agentId,
taskCount: taskIds.size,
});
for (const taskId of taskIds) {
const scheduled = this.tasks.get(taskId);
if (scheduled && scheduled.task.status === "running") {
// Reset task status
scheduled.task.status = "queued";
scheduled.attempts = 0;
// Re-emit task created event for reassignment
this.eventBus.emit(SystemEvents.TASK_CREATED, {
task: scheduled.task,
});
}
}
return Promise.resolve();
}
getAgentTaskCount(agentId) {
return this.agentTasks.get(agentId)?.size ?? 0;
}
getHealthStatus() {
const activeTasks = this.tasks.size;
const completedTasks = this.completedTasks.size;
const agentsWithTasks = this.agentTasks.size;
const tasksByStatus = {
pending: 0,
queued: 0,
assigned: 0,
running: 0,
completed: completedTasks,
failed: 0,
cancelled: 0,
};
for (const scheduled of this.tasks.values()) {
tasksByStatus[scheduled.task.status]++;
}
return Promise.resolve({
healthy: true,
metrics: {
activeTasks,
completedTasks,
agentsWithTasks,
...tasksByStatus,
},
});
}
getAgentTasks(agentId) {
const taskIds = this.agentTasks.get(agentId);
if (!taskIds) {
return Promise.resolve([]);
}
const tasks = [];
for (const taskId of taskIds) {
const scheduled = this.tasks.get(taskId);
if (scheduled) {
tasks.push(scheduled.task);
}
}
return Promise.resolve(tasks);
}
async performMaintenance() {
this.logger.debug("Performing task scheduler maintenance");
// Cleanup old completed tasks
this.cleanup();
// Check for stuck tasks
const now = new Date();
for (const [taskId, scheduled] of this.tasks) {
if (scheduled.task.status === "running" && scheduled.task.startedAt) {
const runtime = now.getTime() - scheduled.task.startedAt.getTime();
if (runtime > this.config.resourceTimeout * 2) {
this.logger.warn("Found stuck task", {
taskId,
runtime,
agentId: scheduled.agentId,
});
// Force fail the task
await this.failTask(taskId, new TaskTimeoutError(taskId, runtime));
}
}
}
}
startTask(taskId) {
const scheduled = this.tasks.get(taskId);
if (!scheduled) {
return;
}
// Update status
scheduled.task.status = "running";
scheduled.task.startedAt = new Date();
// Emit task started event
this.eventBus.emit(SystemEvents.TASK_STARTED, {
taskId,
agentId: scheduled.agentId,
});
// Set timeout for task execution
const timeoutMs = this.config.resourceTimeout;
scheduled.timeout = setTimeout(() => {
this.failTask(taskId, new TaskTimeoutError(taskId, timeoutMs)).catch(error => {
this.logger.error("Failed to handle task timeout", { taskId, error });
});
}, timeoutMs);
}
canStartTask(task) {
// Check if all dependencies are completed
return task.dependencies.every(depId => this.completedTasks.has(depId));
}
async cancelDependentTasks(taskId, reason) {
const dependents = this.taskDependencies.get(taskId);
if (!dependents) {
return;
}
for (const dependentId of dependents) {
await this.cancelTask(dependentId, reason);
}
}
cleanup() {
// Clean up old completed tasks (keep last 1000)
if (this.completedTasks.size > 1000) {
const toRemove = this.completedTasks.size - 1000;
const iterator = this.completedTasks.values();
for (let i = 0; i < toRemove; i++) {
const result = iterator.next();
if (!result.done && result.value) {
this.completedTasks.delete(result.value);
this.taskDependencies.delete(result.value);
}
}
}
}
}
//# sourceMappingURL=scheduler.js.map