vibe-coder-mcp
Version:
Production-ready MCP server with complete agent integration, multi-transport support, and comprehensive development automation tools for AI-assisted workflows.
289 lines (288 loc) • 11.9 kB
JavaScript
import { sseNotifier } from '../../services/sse-notifier/index.js';
import { registerTool } from '../../services/routing/toolRegistry.js';
import { dependencyContainer } from '../../services/dependency-container.js';
import { z } from 'zod';
class AgentTaskQueue {
static instance;
static isInitializing = false;
queues = new Map();
taskHistory = new Map();
assignmentCounter = 0;
agentRegistryCache = null;
static getInstance() {
if (AgentTaskQueue.isInitializing) {
console.warn('Circular initialization detected in AgentTaskQueue, using safe fallback');
return AgentTaskQueue.createSafeFallback();
}
if (!AgentTaskQueue.instance) {
AgentTaskQueue.isInitializing = true;
try {
AgentTaskQueue.instance = new AgentTaskQueue();
}
finally {
AgentTaskQueue.isInitializing = false;
}
}
return AgentTaskQueue.instance;
}
static createSafeFallback() {
const fallback = Object.create(AgentTaskQueue.prototype);
fallback.queues = new Map();
fallback.taskHistory = new Map();
fallback.assignmentCounter = 0;
fallback.assignTask = async () => {
console.warn('AgentTaskQueue fallback: assignTask called during initialization');
return null;
};
fallback.getTasks = async () => {
console.warn('AgentTaskQueue fallback: getTasks called during initialization');
return [];
};
fallback.getQueueLength = async () => {
console.warn('AgentTaskQueue fallback: getQueueLength called during initialization');
return 0;
};
return fallback;
}
async getAgentRegistry() {
if (!this.agentRegistryCache) {
this.agentRegistryCache = await dependencyContainer.getAgentRegistry();
}
return this.agentRegistryCache;
}
async addTask(agentId, task) {
const taskId = this.generateTaskId();
const taskAssignment = {
...task,
taskId,
agentId,
assignedAt: Date.now()
};
if (!this.queues.has(agentId)) {
this.queues.set(agentId, []);
}
this.queues.get(agentId).push(taskAssignment);
this.taskHistory.set(taskId, taskAssignment);
await this.updateAgentTaskCount(agentId);
const agentRegistry = await this.getAgentRegistry();
const agent = agentRegistry ? await agentRegistry.getAgent(agentId) : null;
if (agent?.transportType === 'sse') {
await this.sendSSETaskNotification(agentId, taskAssignment);
}
console.log(`Task ${taskId} assigned to agent ${agentId} (${agent?.transportType || 'unknown'} transport)`);
return taskId;
}
generateTaskId() {
this.assignmentCounter++;
return `task-${Date.now()}-${this.assignmentCounter.toString().padStart(4, '0')}`;
}
async getTasks(agentId, maxTasks = 1) {
const queue = this.queues.get(agentId) || [];
const tasks = queue.splice(0, Math.min(maxTasks, queue.length));
const agentRegistry = await this.getAgentRegistry();
const agent = agentRegistry ? await agentRegistry.getAgent(agentId) : null;
if (agent) {
agent.lastSeen = Date.now();
await this.updateAgentTaskCount(agentId);
}
return tasks;
}
async getQueueLength(agentId) {
return (this.queues.get(agentId) || []).length;
}
async getAllQueueLengths() {
const lengths = {};
for (const [agentId, queue] of this.queues.entries()) {
lengths[agentId] = queue.length;
}
return lengths;
}
async removeTask(taskId) {
const task = this.taskHistory.get(taskId);
if (!task)
return false;
const queue = this.queues.get(task.agentId);
if (queue) {
const index = queue.findIndex(t => t.taskId === taskId);
if (index !== -1) {
queue.splice(index, 1);
await this.updateAgentTaskCount(task.agentId);
return true;
}
}
return false;
}
async getTask(taskId) {
return this.taskHistory.get(taskId);
}
async updateAgentTaskCount(agentId) {
const agentRegistry = await this.getAgentRegistry();
const agent = agentRegistry ? await agentRegistry.getAgent(agentId) : null;
if (agent && agentRegistry) {
const queueLength = await this.getQueueLength(agentId);
agent.currentTasks = Array.from({ length: queueLength }, (_, i) => `pending-${i + 1}`);
const maxTasks = agent.maxConcurrentTasks || 1;
if (queueLength >= maxTasks) {
await agentRegistry.updateAgentStatus(agentId, 'busy');
}
else if (agent.status === 'busy' && queueLength < maxTasks) {
await agentRegistry.updateAgentStatus(agentId, 'online');
}
}
}
async sendSSETaskNotification(agentId, task) {
try {
const agentRegistry = await this.getAgentRegistry();
const agent = agentRegistry ? await agentRegistry.getAgent(agentId) : null;
if (agent?.sessionId) {
await sseNotifier.sendEvent(agent.sessionId, 'taskAssigned', {
agentId,
taskId: task.taskId,
taskPayload: task.sentinelPayload,
priority: task.priority,
estimatedDuration: task.estimatedDuration,
deadline: task.deadline,
assignedAt: task.assignedAt,
metadata: task.metadata
});
console.log(`SSE task notification sent to agent ${agentId} for task ${task.taskId}`);
}
await sseNotifier.broadcastEvent('taskAssignmentUpdate', {
agentId,
taskId: task.taskId,
priority: task.priority,
assignedAt: task.assignedAt,
queueLength: await this.getQueueLength(agentId)
});
}
catch (error) {
console.error('Failed to send SSE task notification:', error);
}
}
async findBestAgent(requiredCapabilities) {
const agentRegistry = await this.getAgentRegistry();
const onlineAgents = agentRegistry ? await agentRegistry.getOnlineAgents() : [];
const capableAgents = onlineAgents.filter((agent) => requiredCapabilities.every(cap => agent.capabilities?.includes(cap)));
if (capableAgents.length === 0) {
return null;
}
const agentsWithLoad = await Promise.all(capableAgents.map(async (agent) => ({
agent,
queueLength: await this.getQueueLength(agent.agentId)
})));
agentsWithLoad.sort((a, b) => a.queueLength - b.queueLength);
for (const { agent, queueLength } of agentsWithLoad) {
if (queueLength < (agent.maxConcurrentTasks || 1)) {
return agent.agentId || null;
}
}
return agentsWithLoad[0].agent.agentId || null;
}
async clearAgentTasks(agentId) {
const queue = this.queues.get(agentId) || [];
const clearedCount = queue.length;
for (const task of queue) {
this.taskHistory.delete(task.taskId);
}
this.queues.set(agentId, []);
await this.updateAgentTaskCount(agentId);
console.log(`Cleared ${clearedCount} tasks for agent ${agentId}`);
return clearedCount;
}
}
export const getAgentTasksTool = {
name: 'get-agent-tasks',
description: 'Get pending tasks for an agent (stdio polling)',
inputSchema: {
type: 'object',
properties: {
agentId: {
type: 'string',
description: 'Agent identifier'
},
maxTasks: {
type: 'number',
default: 1,
minimum: 1,
maximum: 5,
description: 'Maximum number of tasks to retrieve'
}
},
required: ['agentId']
}
};
export async function handleGetAgentTasks(args) {
try {
const { agentId, maxTasks = 1 } = args;
const agentRegistry = await dependencyContainer.getAgentRegistry();
const agent = agentRegistry ? await agentRegistry.getAgent(agentId) : null;
if (!agent) {
return {
content: [{
type: 'text',
text: `❌ Agent Not Found\n\nAgent ${agentId} is not registered.\n\n` +
`Please register first using the 'register-agent' tool.`
}],
isError: true
};
}
const taskQueue = AgentTaskQueue.getInstance();
const tasks = await taskQueue.getTasks(agentId, maxTasks);
if (tasks.length === 0) {
return {
content: [{
type: 'text',
text: `📭 No Tasks Available\n\n` +
`Agent: ${agentId}\n` +
`Status: ${agent.status}\n` +
`Queue Length: 0\n\n` +
`Continue polling for new task assignments.`
}]
};
}
const taskDetails = tasks.map(task => `📋 Task ID: ${task.taskId}\n` +
`Priority: ${task.priority.toUpperCase()}\n` +
`Assigned: ${new Date(task.assignedAt).toISOString()}\n` +
`${task.estimatedDuration ? `Estimated Duration: ${task.estimatedDuration}ms\n` : ''}` +
`${task.deadline ? `Deadline: ${new Date(task.deadline).toISOString()}\n` : ''}` +
`\n--- Sentinel Protocol Payload ---\n${task.sentinelPayload}\n--- End Payload ---`).join('\n\n');
const remainingTasks = await taskQueue.getQueueLength(agentId);
return {
content: [{
type: 'text',
text: `✅ Task Assignment Retrieved\n\n` +
`Agent: ${agentId}\n` +
`Tasks Retrieved: ${tasks.length}\n` +
`Remaining in Queue: ${remainingTasks}\n\n` +
`${taskDetails}\n\n` +
`🔧 Next Steps:\n` +
`1. Process the task(s) according to the Sentinel Protocol\n` +
`2. Submit results using 'submit-task-response' tool\n` +
`3. Continue polling for additional tasks`
}]
};
}
catch (error) {
console.error('Get agent tasks failed:', error);
return {
content: [{
type: 'text',
text: `❌ Task Retrieval Failed\n\nError: ${error instanceof Error ? error.message : 'Unknown error'}\n\n` +
`Please try again or contact support if the issue persists.`
}],
isError: true
};
}
}
export { AgentTaskQueue };
const getAgentTasksInputSchemaShape = {
agentId: z.string().min(1, { message: "Agent ID is required" }).describe("Agent identifier"),
maxTasks: z.number().min(1).max(5).default(1).describe("Maximum number of tasks to retrieve")
};
const getAgentTasksToolDefinition = {
name: "get-agent-tasks",
description: "Get pending tasks for an agent (stdio polling). Supports both stdio and SSE transports for universal task assignment.",
inputSchema: getAgentTasksInputSchemaShape,
executor: handleGetAgentTasks
};
registerTool(getAgentTasksToolDefinition);