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.

411 lines (410 loc) 18.4 kB
import { sseNotifier } from '../../services/sse-notifier/index.js'; import { registerTool } from '../../services/routing/toolRegistry.js'; import { transportManager } from '../../services/transport-manager/index.js'; import { InitializationMonitor } from '../../utils/initialization-monitor.js'; import { dependencyContainer } from '../../services/dependency-container.js'; import { z } from 'zod'; class AgentRegistry { static instance; static isInitializing = false; agents = new Map(); sessionToAgent = new Map(); integrationBridge = null; isBridgeRegistration = false; static getInstance() { if (AgentRegistry.isInitializing) { console.warn('Circular initialization detected in AgentRegistry, using safe fallback'); return AgentRegistry.createSafeFallback(); } if (!AgentRegistry.instance) { const monitor = InitializationMonitor.getInstance(); monitor.startServiceInitialization('AgentRegistry', [ 'SSENotifier', 'TransportManager' ]); AgentRegistry.isInitializing = true; try { monitor.startPhase('AgentRegistry', 'constructor'); AgentRegistry.instance = new AgentRegistry(); monitor.endPhase('AgentRegistry', 'constructor'); monitor.endServiceInitialization('AgentRegistry'); } catch (error) { monitor.endPhase('AgentRegistry', 'constructor', error); monitor.endServiceInitialization('AgentRegistry', error); throw error; } finally { AgentRegistry.isInitializing = false; } } return AgentRegistry.instance; } static createSafeFallback() { const fallback = Object.create(AgentRegistry.prototype); fallback.agents = new Map(); fallback.sessionToAgent = new Map(); fallback.integrationBridge = null; fallback.isBridgeRegistration = false; fallback.registerAgent = async () => { console.warn('AgentRegistry fallback: registerAgent called during initialization'); return { success: false, message: 'Registry initializing' }; }; fallback.getAgent = async () => { console.warn('AgentRegistry fallback: getAgent called during initialization'); return null; }; fallback.getOnlineAgents = async () => { console.warn('AgentRegistry fallback: getOnlineAgents called during initialization'); return []; }; return fallback; } async initializeIntegrationBridge() { if (!this.integrationBridge) { this.integrationBridge = await dependencyContainer.getAgentIntegrationBridge(); if (!this.integrationBridge) { console.warn('Integration bridge not available, using fallback'); } } } async registerAgent(registration) { this.validateRegistration(registration); const existingAgent = this.agents.get(registration.agentId); if (existingAgent) { await this.updateAgent(registration); } else { await this.createAgent(registration); } this.sessionToAgent.set(registration.sessionId, registration.agentId); if (!this.isBridgeRegistration) { await this.initializeIntegrationBridge(); if (this.integrationBridge) { try { await this.integrationBridge.registerAgent({ id: registration.agentId, capabilities: registration.capabilities, status: registration.status || 'online', maxConcurrentTasks: registration.maxConcurrentTasks, currentTasks: registration.currentTasks || [], transportType: registration.transportType, sessionId: registration.sessionId, pollingInterval: registration.pollingInterval, registeredAt: registration.registeredAt || Date.now(), lastSeen: registration.lastSeen || Date.now(), lastHeartbeat: new Date(registration.lastSeen || Date.now()), performance: { tasksCompleted: 0, averageCompletionTime: 0, successRate: 1.0 }, httpEndpoint: registration.httpEndpoint, httpAuthToken: registration.httpAuthToken, websocketConnection: registration.websocketConnection, metadata: { version: '1.0.0', supportedProtocols: [registration.transportType], preferences: { transportType: registration.transportType, sessionId: registration.sessionId, pollingInterval: registration.pollingInterval, httpEndpoint: registration.httpEndpoint, httpAuthToken: registration.httpAuthToken } } }); console.log(`Agent ${registration.agentId} registered in both registry and orchestrator via integration bridge`); } catch (bridgeError) { console.warn(`Integration bridge registration failed for agent ${registration.agentId}:`, bridgeError); } } } if (registration.transportType === 'sse') { await this.notifyAgentRegistered(registration); } } validateRegistration(registration) { if (!registration.agentId || registration.agentId.trim() === '') { throw new Error('Agent ID is required'); } if (!registration.capabilities || registration.capabilities.length === 0) { throw new Error('Agent capabilities are required'); } if (!['stdio', 'sse'].includes(registration.transportType)) { throw new Error('Transport type must be stdio or sse'); } if (!registration.sessionId || registration.sessionId.trim() === '') { throw new Error('Session ID is required'); } if (registration.maxConcurrentTasks < 1 || registration.maxConcurrentTasks > 10) { throw new Error('Max concurrent tasks must be between 1 and 10'); } } async createAgent(registration) { const agentData = { ...registration, status: 'online', registeredAt: Date.now(), lastSeen: Date.now(), currentTasks: [] }; this.agents.set(registration.agentId, agentData); console.log(`Agent ${registration.agentId} registered with ${registration.transportType} transport`); } async updateAgent(registration) { const existing = this.agents.get(registration.agentId); const updated = { ...existing, ...registration, lastSeen: Date.now(), status: 'online' }; this.agents.set(registration.agentId, updated); console.log(`Agent ${registration.agentId} updated registration`); } async notifyAgentRegistered(registration) { try { await sseNotifier.broadcastEvent('agentRegistered', { agentId: registration.agentId, capabilities: registration.capabilities, status: 'online', transportType: registration.transportType, registeredAt: Date.now() }); console.log(`SSE notification sent for agent ${registration.agentId} registration`); } catch (error) { console.error('Failed to send SSE notification for agent registration:', error); } } async getAgent(agentId) { return this.agents.get(agentId); } async getAgentBySession(sessionId) { const agentId = this.sessionToAgent.get(sessionId); return agentId ? this.agents.get(agentId) : undefined; } async getAllAgents() { return Array.from(this.agents.values()); } async getOnlineAgents() { return Array.from(this.agents.values()).filter(agent => agent.status === 'online'); } async updateAgentStatus(agentId, status) { const agent = this.agents.get(agentId); if (agent) { agent.status = status; agent.lastSeen = Date.now(); if (agent.transportType === 'sse') { await this.notifyAgentStatusUpdate(agent); } } } async notifyAgentStatusUpdate(agent) { try { await sseNotifier.broadcastEvent('agentStatusUpdate', { agentId: agent.agentId, status: agent.status, lastSeen: agent.lastSeen, currentTasks: agent.currentTasks?.length || 0 }); } catch (error) { console.error('Failed to send SSE notification for agent status update:', error); } } async unregisterAgent(agentId) { const agent = this.agents.get(agentId); if (agent) { this.agents.delete(agentId); this.sessionToAgent.delete(agent.sessionId); if (agent.transportType === 'sse') { await this.notifyAgentUnregistered(agent); } console.log(`Agent ${agentId} unregistered`); } } async notifyAgentUnregistered(agent) { try { await sseNotifier.broadcastEvent('agentUnregistered', { agentId: agent.agentId, unregisteredAt: Date.now() }); } catch (error) { console.error('Failed to send SSE notification for agent unregistration:', error); } } getTransportEndpoints() { const allocatedPorts = transportManager.getAllocatedPorts(); const endpoints = {}; if (allocatedPorts.websocket) { endpoints.websocket = `ws://localhost:${allocatedPorts.websocket}/agent-ws`; } if (allocatedPorts.http) { endpoints.http = `http://localhost:${allocatedPorts.http}`; } if (allocatedPorts.sse) { endpoints.sse = `http://localhost:${allocatedPorts.sse}/events`; } return endpoints; } getTransportInstructions(registration) { const endpoints = this.getTransportEndpoints(); switch (registration.transportType) { case 'stdio': return `Poll for tasks using 'get-agent-tasks' every ${registration.pollingInterval}ms`; case 'sse': { const sseEndpoint = endpoints.sse || 'http://localhost:3000/events'; return `Connect to SSE endpoint: ${sseEndpoint}/{sessionId} for real-time task notifications`; } case 'websocket': { const wsEndpoint = endpoints.websocket || 'ws://localhost:8080/agent-ws'; return `Connect to WebSocket endpoint: ${wsEndpoint} for real-time task notifications`; } case 'http': { const httpEndpoint = endpoints.http || 'http://localhost:3001'; return `Register with HTTP API: ${httpEndpoint}/agents/register. ` + `Tasks will be sent to your endpoint: ${registration.httpEndpoint}. ` + `Poll for additional tasks at: ${httpEndpoint}/agents/${registration.agentId}/tasks every ${registration.pollingInterval}ms`; } default: return 'Transport-specific instructions not available'; } } async performHealthCheck() { const now = Date.now(); const timeoutMs = 5 * 60 * 1000; for (const [agentId, agent] of this.agents.entries()) { if (agent.status === 'online' && (now - (agent.lastSeen || 0)) > timeoutMs) { await this.updateAgentStatus(agentId, 'offline'); console.log(`Agent ${agentId} marked as offline due to inactivity`); } } } } export const registerAgentTool = { name: 'register-agent', description: 'Register an AI agent with the task management system', inputSchema: { type: 'object', properties: { agentId: { type: 'string', description: 'Unique agent identifier (e.g., claude-agent-001)' }, capabilities: { type: 'array', items: { type: 'string' }, description: 'List of agent capabilities (e.g., code_generation, testing, debugging)' }, transportType: { type: 'string', enum: ['stdio', 'sse', 'websocket', 'http'], description: 'Communication transport type' }, sessionId: { type: 'string', description: 'MCP session identifier' }, maxConcurrentTasks: { type: 'number', default: 1, minimum: 1, maximum: 10, description: 'Maximum number of concurrent tasks this agent can handle' }, pollingInterval: { type: 'number', default: 5000, minimum: 1000, maximum: 30000, description: 'Polling interval in milliseconds (stdio and http transports only)' }, httpEndpoint: { type: 'string', description: 'HTTP callback endpoint URL (required for http transport)' }, httpAuthToken: { type: 'string', description: 'Authentication token for HTTP callbacks (optional for http transport)' } }, required: ['agentId', 'capabilities', 'transportType', 'sessionId'] } }; export async function handleRegisterAgent(args) { try { const registry = AgentRegistry.getInstance(); const registration = { agentId: args.agentId, capabilities: args.capabilities, transportType: args.transportType, sessionId: args.sessionId, maxConcurrentTasks: args.maxConcurrentTasks || 1, pollingInterval: args.pollingInterval || 5000, httpEndpoint: args.httpEndpoint, httpAuthToken: args.httpAuthToken }; if (registration.transportType === 'http' && !registration.httpEndpoint) { throw new Error('HTTP endpoint is required for http transport'); } await registry.registerAgent(registration); const transportInstructions = registry.getTransportInstructions(registration); const endpoints = registry.getTransportEndpoints(); const endpointInfo = Object.entries(endpoints) .filter(([_, url]) => url) .map(([transport, url]) => `${transport.toUpperCase()}: ${url}`) .join('\n'); return { content: [{ type: 'text', text: `✅ Agent Registration Successful\n\n` + `Agent ID: ${registration.agentId}\n` + `Transport: ${registration.transportType}\n` + `Capabilities: ${registration.capabilities.join(', ')}\n` + `Max Concurrent Tasks: ${registration.maxConcurrentTasks}\n` + `Session: ${registration.sessionId}\n\n` + `🌐 Available Endpoints (Dynamic Port Allocation):\n${endpointInfo || 'No endpoints available yet'}\n\n` + `📋 Next Steps:\n${transportInstructions}\n\n` + `🔧 Available Commands:\n` + `- get-agent-tasks: Poll for new task assignments\n` + `- submit-task-response: Submit completed task results` }] }; } catch (error) { console.error('Agent registration failed:', error); return { content: [{ type: 'text', text: `❌ Agent Registration Failed\n\nError: ${error instanceof Error ? error.message : 'Unknown error'}\n\n` + `Please check your registration parameters and try again.` }], isError: true }; } } export { AgentRegistry }; setInterval(() => { AgentRegistry.getInstance().performHealthCheck(); }, 60000); const registerAgentInputSchemaShape = { agentId: z.string().min(1, { message: "Agent ID is required" }).describe("Unique agent identifier (e.g., claude-agent-001)"), capabilities: z.array(z.string()).min(1, { message: "At least one capability is required" }).describe("List of agent capabilities (e.g., code_generation, testing, debugging)"), transportType: z.enum(['stdio', 'sse', 'websocket', 'http']).describe("Communication transport type"), sessionId: z.string().min(1, { message: "Session ID is required" }).describe("MCP session identifier"), maxConcurrentTasks: z.number().min(1).max(10).default(1).describe("Maximum number of concurrent tasks this agent can handle"), pollingInterval: z.number().min(1000).max(30000).default(5000).describe("Polling interval in milliseconds (stdio and http transports only)"), httpEndpoint: z.string().url().optional().describe("HTTP callback endpoint URL (required for http transport)"), httpAuthToken: z.string().optional().describe("Authentication token for HTTP callbacks (optional for http transport)") }; const registerAgentToolDefinition = { name: "register-agent", description: "Register an AI agent with the task management system. Supports both stdio and SSE transports for universal agent communication.", inputSchema: registerAgentInputSchemaShape, executor: handleRegisterAgent }; registerTool(registerAgentToolDefinition);