vibe-coder-mcp
Version:
Production-ready MCP server with complete agent integration, multi-transport support, and comprehensive development automation tools for AI-assisted workflows.
352 lines (351 loc) • 14 kB
JavaScript
import { WebSocketServer, WebSocket } from 'ws';
import { createServer } from 'http';
import logger from '../../logger.js';
import { AgentRegistry } from '../../tools/agent-registry/index.js';
class WebSocketServerManager {
static instance;
server;
httpServer;
connections = new Map();
agentConnections = new Map();
port = 8080;
heartbeatInterval;
static getInstance() {
if (!WebSocketServerManager.instance) {
WebSocketServerManager.instance = new WebSocketServerManager();
}
return WebSocketServerManager.instance;
}
async start(port) {
try {
if (!port || port <= 0 || port > 65535) {
throw new Error(`Invalid port provided: ${port}. Port should be pre-allocated by Transport Manager.`);
}
this.port = port;
logger.debug({ port }, 'Starting WebSocket server with pre-allocated port');
this.httpServer = createServer();
this.server = new WebSocketServer({
server: this.httpServer,
path: '/agent-ws'
});
this.server.on('connection', this.handleConnection.bind(this));
this.server.on('error', this.handleServerError.bind(this));
await new Promise((resolve, reject) => {
this.httpServer.listen(port, (err) => {
if (err) {
if (err.message.includes('EADDRINUSE')) {
const enhancedError = new Error(`Port ${port} is already in use. This should not happen with pre-allocated ports. ` +
`Transport Manager port allocation may have failed.`);
enhancedError.name = 'PortAllocationError';
reject(enhancedError);
}
else {
reject(err);
}
}
else {
resolve();
}
});
});
this.startHeartbeatMonitoring();
logger.info({
port,
path: '/agent-ws',
note: 'Using pre-allocated port from Transport Manager'
}, 'WebSocket server started successfully');
}
catch (error) {
logger.error({
err: error,
port,
context: 'WebSocket server startup with pre-allocated port'
}, 'Failed to start WebSocket server');
if (error instanceof Error) {
error.message = `WebSocket server startup failed on pre-allocated port ${port}: ${error.message}`;
}
throw error;
}
}
async stop() {
try {
if (this.heartbeatInterval) {
clearInterval(this.heartbeatInterval);
}
for (const [sessionId, connection] of this.connections.entries()) {
connection.ws.close(1000, 'Server shutdown');
this.connections.delete(sessionId);
}
if (this.server) {
await new Promise((resolve) => {
this.server.close(() => resolve());
});
}
if (this.httpServer) {
await new Promise((resolve) => {
this.httpServer.close(() => resolve());
});
}
logger.info('WebSocket server stopped');
}
catch (error) {
logger.error({ err: error }, 'Error stopping WebSocket server');
throw error;
}
}
handleConnection(ws, request) {
const sessionId = this.generateSessionId();
const connection = {
ws,
sessionId,
lastSeen: Date.now(),
authenticated: false
};
this.connections.set(sessionId, connection);
logger.info({ sessionId, remoteAddress: request.socket.remoteAddress }, 'New WebSocket connection');
ws.on('message', (data) => this.handleMessage(sessionId, data));
ws.on('close', (code, reason) => this.handleDisconnection(sessionId, code, reason));
ws.on('error', (error) => this.handleConnectionError(sessionId, error));
ws.on('pong', () => this.handlePong(sessionId));
this.sendMessage(sessionId, {
type: 'register',
sessionId,
data: {
message: 'WebSocket connection established. Please register your agent.',
timestamp: Date.now()
}
});
}
async handleMessage(sessionId, rawData) {
try {
const connection = this.connections.get(sessionId);
if (!connection) {
logger.warn({ sessionId }, 'Message received for unknown connection');
return;
}
connection.lastSeen = Date.now();
let message;
try {
if (typeof rawData !== 'string' && !Buffer.isBuffer(rawData)) {
this.sendError(sessionId, 'Invalid message data type');
return;
}
message = JSON.parse(rawData.toString());
}
catch {
this.sendError(sessionId, 'Invalid JSON message format');
return;
}
if (!message.timestamp) {
message.timestamp = Date.now();
}
logger.debug({ sessionId, messageType: message.type }, 'WebSocket message received');
switch (message.type) {
case 'register':
await this.handleAgentRegistration(sessionId, message);
break;
case 'task_response':
await this.handleTaskResponse(sessionId, message);
break;
case 'heartbeat':
await this.handleHeartbeat(sessionId, message);
break;
default:
this.sendError(sessionId, `Unknown message type: ${message.type}`);
}
}
catch (error) {
logger.error({ err: error, sessionId }, 'Error handling WebSocket message');
this.sendError(sessionId, 'Internal server error processing message');
}
}
async handleAgentRegistration(sessionId, message) {
try {
const data = message.data;
const { agentId, capabilities, maxConcurrentTasks } = data || {};
if (!agentId || !capabilities) {
this.sendError(sessionId, 'Agent registration requires agentId and capabilities');
return;
}
const connection = this.connections.get(sessionId);
if (!connection) {
this.sendError(sessionId, 'Connection not found');
return;
}
connection.agentId = agentId;
connection.authenticated = true;
const agentRegistry = AgentRegistry.getInstance();
await agentRegistry.registerAgent({
agentId,
capabilities,
transportType: 'websocket',
sessionId,
maxConcurrentTasks: maxConcurrentTasks || 1
});
this.agentConnections.set(agentId, sessionId);
this.sendMessage(sessionId, {
type: 'register',
agentId,
data: {
agentId,
capabilities: capabilities || [],
transportType: 'websocket',
maxConcurrentTasks: maxConcurrentTasks || 1
}
});
logger.info({ sessionId, agentId }, 'Agent registered via WebSocket');
}
catch (error) {
logger.error({ err: error, sessionId }, 'Failed to register agent via WebSocket');
this.sendError(sessionId, `Registration failed: ${error instanceof Error ? error.message : 'Unknown error'}`);
}
}
async handleTaskResponse(sessionId, message) {
try {
const connection = this.connections.get(sessionId);
if (!connection?.agentId) {
this.sendError(sessionId, 'Agent must be registered before submitting task responses');
return;
}
const { AgentResponseProcessor } = await import('../../tools/agent-response/index.js');
const responseProcessor = AgentResponseProcessor.getInstance();
const responseData = message.data;
await responseProcessor.processResponse({
agentId: connection.agentId,
taskId: responseData.taskId,
status: responseData.status === 'IN_PROGRESS' ? 'PARTIAL' : responseData.status,
response: String(responseData.response || ''),
completionDetails: responseData.error ? { errorDetails: responseData.error } : undefined,
receivedAt: Date.now()
});
this.sendMessage(sessionId, {
type: 'task_response',
agentId: connection.agentId,
data: {
taskId: responseData.taskId,
status: 'DONE',
response: 'acknowledged'
}
});
logger.info({
sessionId,
agentId: connection.agentId,
taskId: responseData.taskId
}, 'Task response received via WebSocket');
}
catch (error) {
logger.error({ err: error, sessionId }, 'Failed to process task response via WebSocket');
this.sendError(sessionId, `Task response processing failed: ${error instanceof Error ? error.message : 'Unknown error'}`);
}
}
async handleHeartbeat(sessionId, _message) {
const connection = this.connections.get(sessionId);
if (connection) {
connection.lastSeen = Date.now();
this.sendMessage(sessionId, {
type: 'heartbeat',
agentId: connection.agentId,
data: {
timestamp: Date.now(),
status: 'alive'
}
});
}
}
handleDisconnection(sessionId, code, reason) {
const connection = this.connections.get(sessionId);
if (connection?.agentId) {
this.agentConnections.delete(connection.agentId);
const agentRegistry = AgentRegistry.getInstance();
agentRegistry.updateAgentStatus(connection.agentId, 'offline').catch(error => {
logger.error({ err: error, agentId: connection.agentId }, 'Failed to update agent status on disconnect');
});
}
this.connections.delete(sessionId);
logger.info({
sessionId,
agentId: connection?.agentId,
code,
reason: reason.toString()
}, 'WebSocket connection closed');
}
handleConnectionError(sessionId, error) {
logger.error({ err: error, sessionId }, 'WebSocket connection error');
const connection = this.connections.get(sessionId);
if (connection) {
connection.ws.close(1011, 'Connection error');
}
}
handlePong(sessionId) {
const connection = this.connections.get(sessionId);
if (connection) {
connection.lastSeen = Date.now();
}
}
handleServerError(error) {
logger.error({ err: error }, 'WebSocket server error');
}
async sendTaskToAgent(agentId, taskPayload) {
try {
const sessionId = this.agentConnections.get(agentId);
if (!sessionId) {
logger.warn({ agentId }, 'No WebSocket connection found for agent');
return false;
}
this.sendMessage(sessionId, {
type: 'task_assignment',
agentId,
data: taskPayload
});
return true;
}
catch (error) {
logger.error({ err: error, agentId }, 'Failed to send task via WebSocket');
return false;
}
}
sendMessage(sessionId, message) {
const connection = this.connections.get(sessionId);
if (connection && connection.ws.readyState === WebSocket.OPEN) {
connection.ws.send(JSON.stringify(message));
}
}
sendError(sessionId, errorMessage) {
this.sendMessage(sessionId, {
type: 'error',
data: {
error: errorMessage,
timestamp: Date.now()
}
});
}
generateSessionId() {
return `ws-${Date.now()}-${Math.random().toString(36).substr(2, 9)}`;
}
startHeartbeatMonitoring() {
this.heartbeatInterval = setInterval(() => {
const now = Date.now();
const timeout = 60000;
for (const [sessionId, connection] of this.connections.entries()) {
if (now - connection.lastSeen > timeout) {
logger.warn({ sessionId, agentId: connection.agentId }, 'WebSocket connection timed out');
connection.ws.close(1000, 'Connection timeout');
}
else if (connection.ws.readyState === WebSocket.OPEN) {
connection.ws.ping();
}
}
}, 30000);
}
getConnectionCount() {
return this.connections.size;
}
getConnectedAgents() {
return Array.from(this.agentConnections.keys());
}
isAgentConnected(agentId) {
return this.agentConnections.has(agentId);
}
}
export const websocketServer = WebSocketServerManager.getInstance();
export { WebSocketServerManager };