UNPKG

vibe-coder-mcp

Version:

Production-ready MCP server with complete agent integration, multi-transport support, and comprehensive development automation tools for AI-assisted workflows.

346 lines (345 loc) 12.9 kB
import { AgentOrchestrator } from './agent-orchestrator.js'; import { TaskScheduler } from './task-scheduler.js'; import { AppError, ValidationError } from '../../../utils/errors.js'; import { MemoryManager } from '../../code-map-generator/cache/memoryManager.js'; import logger from '../../../logger.js'; export class TaskStreamer { static instance = null; taskQueue = []; claims = new Map(); agentOrchestrator; taskScheduler; memoryManager; config; streamTimer; status; eventListeners = new Map(); constructor(config) { this.config = { batchSize: 5, streamInterval: 5000, maxQueueSize: 1000, priorityThreshold: 2, enableRealTimeStreaming: true, loadBalancingEnabled: true, ...config }; this.agentOrchestrator = AgentOrchestrator.getInstance(); this.taskScheduler = new TaskScheduler(); this.memoryManager = new MemoryManager(); this.status = { isActive: false, queuedTasks: 0, streamedTasks: 0, failedTasks: 0, averageStreamTime: 0 }; logger.info({ config: this.config }, 'Task streamer initialized'); } static getInstance(config) { if (!TaskStreamer.instance) { TaskStreamer.instance = new TaskStreamer(config); } return TaskStreamer.instance; } async startStreaming() { try { if (this.status.isActive) { logger.warn('Task streaming already active'); return; } this.status.isActive = true; this.status.lastStreamAt = new Date(); if (this.config.enableRealTimeStreaming) { this.streamTimer = setInterval(() => { this.processTaskQueue().catch(error => { logger.error({ err: error }, 'Error in task streaming interval'); }); }, this.config.streamInterval); } this.emitEvent('stream_started', {}); logger.info('Task streaming started'); } catch (error) { logger.error({ err: error }, 'Failed to start task streaming'); throw new AppError('Task streaming startup failed', { cause: error }); } } async stopStreaming() { try { this.status.isActive = false; if (this.streamTimer) { clearInterval(this.streamTimer); this.streamTimer = undefined; } await this.releaseAllClaims(); this.emitEvent('stream_stopped', {}); logger.info('Task streaming stopped'); } catch (error) { logger.error({ err: error }, 'Failed to stop task streaming'); throw new AppError('Task streaming shutdown failed', { cause: error }); } } async queueTask(task, context, epicTitle) { try { if (this.taskQueue.length >= this.config.maxQueueSize) { this.emitEvent('queue_full', { taskId: task.id }); throw new ValidationError('Task queue is full'); } const priority = this.calculateTaskPriority(task); const queueItem = { task, context, epicTitle, priority }; this.insertTaskByPriority(queueItem); this.status.queuedTasks = this.taskQueue.length; this.emitEvent('task_queued', { taskId: task.id }); logger.debug({ taskId: task.id, priority, queueSize: this.taskQueue.length }, 'Task queued for streaming'); if (priority <= this.config.priorityThreshold && this.status.isActive) { await this.processTaskQueue(); } this.memoryManager.getMemoryStats(); } catch (error) { logger.error({ err: error, taskId: task.id }, 'Failed to queue task'); throw new AppError('Task queuing failed', { cause: error }); } } async claimTask(taskId, agentId) { try { const queueIndex = this.taskQueue.findIndex(item => item.task.id === taskId); if (queueIndex === -1) { logger.warn({ taskId, agentId }, 'Attempted to claim non-existent task'); return null; } const existingClaim = this.claims.get(taskId); if (existingClaim && existingClaim.status === 'claimed') { logger.warn({ taskId, agentId, existingAgent: existingClaim.agentId }, 'Task already claimed'); return null; } const claim = { taskId, agentId, claimedAt: new Date(), expiresAt: new Date(Date.now() + 30 * 60 * 1000), status: 'claimed' }; this.claims.set(taskId, claim); this.emitEvent('task_claimed', { taskId, agentId }); logger.info({ taskId, agentId }, 'Task claimed by agent'); return claim; } catch (error) { logger.error({ err: error, taskId, agentId }, 'Failed to claim task'); throw new AppError('Task claim failed', { cause: error }); } } async releaseTask(taskId, agentId) { try { const claim = this.claims.get(taskId); if (!claim) { logger.warn({ taskId, agentId }, 'Attempted to release non-existent claim'); return; } if (claim.agentId !== agentId) { logger.warn({ taskId, agentId, claimOwner: claim.agentId }, 'Attempted to release claim owned by different agent'); return; } claim.status = 'released'; this.claims.delete(taskId); this.emitEvent('task_released', { taskId, agentId }); logger.info({ taskId, agentId }, 'Task claim released'); } catch (error) { logger.error({ err: error, taskId, agentId }, 'Failed to release task'); throw new AppError('Task release failed', { cause: error }); } } async getReadyTasks(limit) { try { const batchSize = limit || this.config.batchSize; const readyTasks = []; for (const queueItem of this.taskQueue.slice(0, batchSize * 2)) { if (readyTasks.length >= batchSize) break; const task = queueItem.task; const claim = this.claims.get(task.id); if (claim && claim.status === 'claimed') continue; const dependenciesMet = await this.checkDependencies(task); if (!dependenciesMet) continue; readyTasks.push(task); } return readyTasks; } catch (error) { logger.error({ err: error }, 'Failed to get ready tasks'); throw new AppError('Ready tasks retrieval failed', { cause: error }); } } getStatus() { return { ...this.status, queuedTasks: this.taskQueue.length }; } getQueueInfo() { const now = new Date(); const highPriorityTasks = this.taskQueue.filter(item => item.priority <= this.config.priorityThreshold).length; const claimedTasks = Array.from(this.claims.values()).filter(claim => claim.status === 'claimed').length; let oldestTaskAge = 0; if (this.taskQueue.length > 0) { const oldestTask = this.taskQueue[this.taskQueue.length - 1]; oldestTaskAge = now.getTime() - new Date(oldestTask.task.createdAt).getTime(); } return { totalTasks: this.taskQueue.length, highPriorityTasks, claimedTasks, oldestTaskAge }; } addEventListener(event, listener) { if (!this.eventListeners.has(event)) { this.eventListeners.set(event, []); } this.eventListeners.get(event).push(listener); } removeEventListener(event, listener) { const listeners = this.eventListeners.get(event); if (listeners) { const index = listeners.indexOf(listener); if (index > -1) { listeners.splice(index, 1); } } } async processTaskQueue() { try { if (!this.status.isActive || this.taskQueue.length === 0) { return; } const startTime = Date.now(); const readyTasks = await this.getReadyTasks(); if (readyTasks.length === 0) { return; } let streamedCount = 0; for (const task of readyTasks) { const queueItem = this.taskQueue.find(item => item.task.id === task.id); if (!queueItem) continue; try { const assignment = await this.agentOrchestrator.assignTask(task, queueItem.context, queueItem.epicTitle); if (assignment) { this.removeTaskFromQueue(task.id); streamedCount++; this.emitEvent('task_streamed', { taskId: task.id, agentId: assignment.agentId }); logger.debug({ taskId: task.id, agentId: assignment.agentId }, 'Task streamed to agent'); } else { this.emitEvent('agent_unavailable', { taskId: task.id }); } } catch (error) { logger.error({ err: error, taskId: task.id }, 'Failed to stream task'); this.status.failedTasks++; } } if (streamedCount > 0) { this.status.streamedTasks += streamedCount; this.status.lastStreamAt = new Date(); const streamTime = Date.now() - startTime; this.status.averageStreamTime = (this.status.averageStreamTime + streamTime) / 2; } this.status.queuedTasks = this.taskQueue.length; } catch (error) { logger.error({ err: error }, 'Error processing task queue'); } } calculateTaskPriority(task) { const priorityMap = { 'critical': 1, 'high': 2, 'medium': 3, 'low': 4 }; return priorityMap[task.priority.toLowerCase()] || 3; } insertTaskByPriority(queueItem) { let insertIndex = this.taskQueue.length; for (let i = 0; i < this.taskQueue.length; i++) { if (queueItem.priority < this.taskQueue[i].priority) { insertIndex = i; break; } } this.taskQueue.splice(insertIndex, 0, queueItem); } removeTaskFromQueue(taskId) { const index = this.taskQueue.findIndex(item => item.task.id === taskId); if (index > -1) { this.taskQueue.splice(index, 1); } } async checkDependencies(task) { if (!task.dependencies || task.dependencies.length === 0) { return true; } return true; } async releaseAllClaims() { for (const [taskId, claim] of this.claims.entries()) { if (claim.status === 'claimed') { claim.status = 'released'; this.emitEvent('task_released', { taskId, agentId: claim.agentId }); } } this.claims.clear(); } emitEvent(event, data) { const eventData = { event, timestamp: new Date(), ...data }; const listeners = this.eventListeners.get(event); if (listeners) { listeners.forEach(listener => { try { listener(eventData); } catch (error) { logger.error({ err: error, event }, 'Error in stream event listener'); } }); } } destroy() { this.stopStreaming().catch(error => { logger.error({ err: error }, 'Error stopping streaming during destroy'); }); this.taskQueue = []; this.claims.clear(); this.eventListeners.clear(); TaskStreamer.instance = null; logger.info('Task streamer destroyed'); } }