task-engine-ai-core
Version:
Revolutionary AI-driven task management system with complete transformation trilogy: Frontend v0.1.0, Backend v0.2.0, CLI v0.3.0 - Enterprise-grade performance with 95% improvements
637 lines (554 loc) • 21.5 kB
JavaScript
/**
* WebSocket-to-Stdio Bridge Server
* Bridges WebSocket connections from IDEs to stdio-based MCP communication
*/
import WebSocket, { WebSocketServer } from 'ws';
import { spawn } from 'child_process';
import { EventEmitter } from 'events';
import path from 'path';
import { fileURLToPath } from 'url';
import logger from '../../mcp-server/src/logger.js';
const __filename = fileURLToPath(import.meta.url);
const __dirname = path.dirname(__filename);
export class WebSocketBridge extends EventEmitter {
constructor(options = {}) {
super();
this.port = options.port || 8765;
this.host = options.host || 'localhost';
this.mcpServerPath = options.mcpServerPath || path.join(__dirname, '../../mcp-server/server.js');
this.maxConnections = options.maxConnections || 10;
this.timeout = options.timeout || 30000;
this.wss = null;
this.activeConnections = new Map();
this.mcpProcesses = new Map();
this.heartbeatInterval = null;
this.connectionCounter = 0;
this.messageBuffer = new Map(); // Buffer for incomplete messages
}
/**
* Start the WebSocket bridge server
*/
async start() {
try {
// Check if port is available
await this.checkPortAvailability();
this.wss = new WebSocketServer({
port: this.port,
host: this.host,
perMessageDeflate: false,
maxPayload: 1024 * 1024 * 10 // 10MB max payload
});
this.setupServerEventHandlers();
this.startHeartbeat();
logger.info(`WebSocket bridge started on ${this.host}:${this.port}`);
logger.info(`Max connections: ${this.maxConnections}`);
logger.info(`MCP server path: ${this.mcpServerPath}`);
this.emit('started', { port: this.port, host: this.host });
return this;
} catch (error) {
logger.error('Failed to start WebSocket bridge:', error);
throw error;
}
}
/**
* Check if port is available
*/
async checkPortAvailability() {
return new Promise((resolve, reject) => {
const testServer = new WebSocketServer({ port: this.port, host: this.host });
testServer.on('listening', () => {
testServer.close(() => resolve());
});
testServer.on('error', (error) => {
if (error.code === 'EADDRINUSE') {
reject(new Error(`Port ${this.port} is already in use`));
} else {
reject(error);
}
});
});
}
/**
* Setup WebSocket server event handlers
*/
setupServerEventHandlers() {
this.wss.on('connection', (ws, request) => {
this.handleConnection(ws, request);
});
this.wss.on('error', (error) => {
logger.error('WebSocket server error:', error);
this.emit('error', error);
});
this.wss.on('close', () => {
logger.info('WebSocket server closed');
this.emit('closed');
});
}
/**
* Handle new WebSocket connection from IDE
*/
handleConnection(ws, request) {
// Check connection limit
if (this.activeConnections.size >= this.maxConnections) {
logger.warn('Connection limit reached, rejecting new connection');
ws.close(1013, 'Server overloaded');
return;
}
const connectionId = this.generateConnectionId();
const clientIP = request.socket.remoteAddress;
logger.info(`New IDE connection: ${connectionId} from ${clientIP}`);
// Store connection metadata
const connectionInfo = {
ws,
connected: true,
lastActivity: Date.now(),
clientIP,
userAgent: request.headers['user-agent'] || 'unknown',
messageCount: 0,
bytesReceived: 0,
bytesSent: 0
};
this.activeConnections.set(connectionId, connectionInfo);
this.messageBuffer.set(connectionId, '');
// Spawn MCP server process for this connection
try {
const mcpProcess = this.spawnMCPServer(connectionId);
this.mcpProcesses.set(connectionId, mcpProcess);
// Set up message routing
this.setupMessageRouting(connectionId, ws, mcpProcess);
// Set up connection event handlers
this.setupConnectionEventHandlers(connectionId, ws);
// Send welcome message
this.sendWelcomeMessage(connectionId, ws);
this.emit('connection', connectionId, connectionInfo);
} catch (error) {
logger.error(`Failed to setup connection ${connectionId}:`, error);
this.cleanupConnection(connectionId);
}
}
/**
* Generate unique connection ID
*/
generateConnectionId() {
this.connectionCounter++;
return `ide_${Date.now()}_${this.connectionCounter.toString().padStart(4, '0')}`;
}
/**
* Spawn MCP server process for connection
*/
spawnMCPServer(connectionId) {
const env = {
...process.env,
CONNECTION_ID: connectionId,
BRIDGE_MODE: 'true',
NODE_ENV: process.env.NODE_ENV || 'development'
};
logger.debug(`Spawning MCP server for ${connectionId}`);
const mcpProcess = spawn('node', [this.mcpServerPath], {
stdio: ['pipe', 'pipe', 'pipe'],
env,
cwd: process.cwd()
});
// Set up process event handlers
mcpProcess.on('error', (error) => {
logger.error(`MCP process error for ${connectionId}:`, error);
this.cleanupConnection(connectionId);
});
mcpProcess.on('exit', (code, signal) => {
logger.info(`MCP process exited for ${connectionId} with code ${code}, signal ${signal}`);
this.cleanupConnection(connectionId);
});
mcpProcess.on('spawn', () => {
logger.debug(`MCP process spawned successfully for ${connectionId}`);
});
// Handle stderr
mcpProcess.stderr.on('data', (data) => {
const errorMsg = data.toString().trim();
if (errorMsg) {
logger.error(`MCP stderr for ${connectionId}: ${errorMsg}`);
}
});
return mcpProcess;
}
/**
* Send welcome message to new connection
*/
sendWelcomeMessage(connectionId, ws) {
const welcomeMessage = {
type: 'bridge-welcome',
connectionId,
timestamp: new Date().toISOString(),
bridgeVersion: '1.0.0',
capabilities: ['mcp-bridge', 'ide-integration', 'real-time-communication']
};
try {
ws.send(JSON.stringify(welcomeMessage));
logger.debug(`Welcome message sent to ${connectionId}`);
} catch (error) {
logger.error(`Failed to send welcome message to ${connectionId}:`, error);
}
}
/**
* Set up connection event handlers
*/
setupConnectionEventHandlers(connectionId, ws) {
// Handle connection close
ws.on('close', (code, reason) => {
logger.info(`Connection ${connectionId} closed with code ${code}: ${reason}`);
this.cleanupConnection(connectionId);
});
// Handle connection errors
ws.on('error', (error) => {
logger.error(`WebSocket error for ${connectionId}:`, error);
this.cleanupConnection(connectionId);
});
// Handle pong responses (for heartbeat)
ws.on('pong', () => {
const connection = this.activeConnections.get(connectionId);
if (connection) {
connection.lastActivity = Date.now();
}
});
}
/**
* Start heartbeat mechanism
*/
startHeartbeat() {
this.heartbeatInterval = setInterval(() => {
const now = Date.now();
const timeoutThreshold = now - this.timeout;
for (const [connectionId, connection] of this.activeConnections) {
if (connection.lastActivity < timeoutThreshold) {
logger.warn(`Connection ${connectionId} timed out`);
this.cleanupConnection(connectionId);
} else if (connection.ws.readyState === WebSocket.OPEN) {
// Send ping
try {
connection.ws.ping();
} catch (error) {
logger.error(`Failed to ping ${connectionId}:`, error);
this.cleanupConnection(connectionId);
}
}
}
}, 10000); // Check every 10 seconds
}
/**
* Set up bidirectional message routing between WebSocket and MCP process
*/
setupMessageRouting(connectionId, ws, mcpProcess) {
// WebSocket -> MCP Server (stdin)
ws.on('message', (data) => {
try {
this.handleWebSocketMessage(connectionId, data, mcpProcess);
} catch (error) {
logger.error(`Error handling WebSocket message from ${connectionId}:`, error);
this.sendErrorMessage(connectionId, ws, 'Message processing error', error.message);
}
});
// MCP Server -> WebSocket (stdout)
mcpProcess.stdout.on('data', (data) => {
try {
this.handleMCPResponse(connectionId, data, ws);
} catch (error) {
logger.error(`Error handling MCP response for ${connectionId}:`, error);
}
});
}
/**
* Handle incoming WebSocket message
*/
handleWebSocketMessage(connectionId, data, mcpProcess) {
const connection = this.activeConnections.get(connectionId);
if (!connection) {
logger.warn(`Received message for unknown connection: ${connectionId}`);
return;
}
// Update connection stats
connection.lastActivity = Date.now();
connection.messageCount++;
connection.bytesReceived += data.length;
let message;
try {
const messageText = data.toString();
message = JSON.parse(messageText);
} catch (error) {
logger.error(`Invalid JSON from IDE ${connectionId}:`, error);
this.sendErrorMessage(connectionId, connection.ws, 'Invalid JSON', 'Message must be valid JSON');
return;
}
// Add bridge metadata
const enrichedMessage = {
...message,
_bridge: {
connectionId,
timestamp: Date.now(),
messageId: `${connectionId}_${connection.messageCount}`
}
};
// Route to MCP server
this.routeToMCP(connectionId, enrichedMessage, mcpProcess);
}
/**
* Handle MCP server response
*/
handleMCPResponse(connectionId, data, ws) {
const connection = this.activeConnections.get(connectionId);
if (!connection) {
logger.warn(`Received MCP response for unknown connection: ${connectionId}`);
return;
}
// Buffer incomplete messages
let buffer = this.messageBuffer.get(connectionId) + data.toString();
const lines = buffer.split('\n');
// Keep the last incomplete line in buffer
this.messageBuffer.set(connectionId, lines.pop() || '');
// Process complete lines
for (const line of lines) {
if (line.trim()) {
try {
const message = JSON.parse(line);
this.routeToIDE(connectionId, message, ws);
} catch (error) {
logger.error(`Invalid JSON from MCP server for ${connectionId}:`, error);
}
}
}
}
/**
* Route message from IDE to MCP server
*/
routeToMCP(connectionId, message, mcpProcess) {
try {
const jsonMessage = JSON.stringify(message) + '\n';
if (mcpProcess.stdin && mcpProcess.stdin.writable) {
mcpProcess.stdin.write(jsonMessage);
logger.debug(`Routed to MCP ${connectionId}:`, message.method || message.type || 'unknown');
} else {
logger.error(`MCP process stdin not writable for ${connectionId}`);
throw new Error('MCP process not ready');
}
} catch (error) {
logger.error(`Failed to route message to MCP for ${connectionId}:`, error);
throw error;
}
}
/**
* Route message from MCP server to IDE
*/
routeToIDE(connectionId, message, ws) {
const connection = this.activeConnections.get(connectionId);
if (!connection) {
return;
}
try {
// Remove bridge metadata before sending to IDE
if (message._bridge) {
delete message._bridge;
}
const messageText = JSON.stringify(message);
if (ws.readyState === WebSocket.OPEN) {
ws.send(messageText);
connection.bytesSent += messageText.length;
logger.debug(`Routed to IDE ${connectionId}:`, message.method || message.type || 'response');
} else {
logger.warn(`WebSocket not open for ${connectionId}, state: ${ws.readyState}`);
}
} catch (error) {
logger.error(`Error routing to IDE ${connectionId}:`, error);
}
}
/**
* Send error message to IDE
*/
sendErrorMessage(connectionId, ws, errorType, errorMessage) {
const errorResponse = {
type: 'bridge-error',
error: {
type: errorType,
message: errorMessage,
connectionId,
timestamp: new Date().toISOString()
}
};
try {
if (ws.readyState === WebSocket.OPEN) {
ws.send(JSON.stringify(errorResponse));
}
} catch (error) {
logger.error(`Failed to send error message to ${connectionId}:`, error);
}
}
/**
* Clean up connection and associated resources
*/
cleanupConnection(connectionId) {
logger.debug(`Cleaning up connection: ${connectionId}`);
// Close WebSocket
const connection = this.activeConnections.get(connectionId);
if (connection && connection.ws.readyState === WebSocket.OPEN) {
try {
connection.ws.close(1000, 'Connection cleanup');
} catch (error) {
logger.debug(`Error closing WebSocket for ${connectionId}:`, error);
}
}
this.activeConnections.delete(connectionId);
// Terminate MCP process
const mcpProcess = this.mcpProcesses.get(connectionId);
if (mcpProcess && !mcpProcess.killed) {
try {
mcpProcess.stdin.end();
mcpProcess.kill('SIGTERM');
// Force kill after timeout
setTimeout(() => {
if (!mcpProcess.killed) {
mcpProcess.kill('SIGKILL');
}
}, 5000);
} catch (error) {
logger.debug(`Error terminating MCP process for ${connectionId}:`, error);
}
}
this.mcpProcesses.delete(connectionId);
// Clean up message buffer
this.messageBuffer.delete(connectionId);
logger.info(`Cleaned up connection: ${connectionId}`);
this.emit('disconnection', connectionId);
}
/**
* Stop the bridge server
*/
async stop() {
logger.info('Stopping WebSocket bridge server...');
// Stop heartbeat
if (this.heartbeatInterval) {
clearInterval(this.heartbeatInterval);
this.heartbeatInterval = null;
}
// Close all connections
const connectionIds = Array.from(this.activeConnections.keys());
for (const connectionId of connectionIds) {
this.cleanupConnection(connectionId);
}
// Close WebSocket server
if (this.wss) {
return new Promise((resolve) => {
this.wss.close(() => {
logger.info('WebSocket bridge server stopped');
this.emit('stopped');
resolve();
});
});
}
}
/**
* Get bridge status and statistics
*/
getStatus() {
const connections = Array.from(this.activeConnections.entries()).map(([id, conn]) => ({
id,
clientIP: conn.clientIP,
userAgent: conn.userAgent,
connected: conn.connected,
lastActivity: new Date(conn.lastActivity).toISOString(),
messageCount: conn.messageCount,
bytesReceived: conn.bytesReceived,
bytesSent: conn.bytesSent,
uptime: Date.now() - (conn.connectedAt || Date.now())
}));
return {
running: !!this.wss,
port: this.port,
host: this.host,
activeConnections: this.activeConnections.size,
maxConnections: this.maxConnections,
mcpProcesses: this.mcpProcesses.size,
uptime: process.uptime(),
connections,
stats: {
totalConnections: this.connectionCounter,
totalMessages: connections.reduce((sum, conn) => sum + conn.messageCount, 0),
totalBytesReceived: connections.reduce((sum, conn) => sum + conn.bytesReceived, 0),
totalBytesSent: connections.reduce((sum, conn) => sum + conn.bytesSent, 0)
}
};
}
/**
* Get connection by ID
*/
getConnection(connectionId) {
const connection = this.activeConnections.get(connectionId);
if (!connection) {
return null;
}
return {
id: connectionId,
clientIP: connection.clientIP,
userAgent: connection.userAgent,
connected: connection.connected,
lastActivity: new Date(connection.lastActivity).toISOString(),
messageCount: connection.messageCount,
bytesReceived: connection.bytesReceived,
bytesSent: connection.bytesSent
};
}
/**
* Broadcast message to all connected IDEs
*/
broadcast(message) {
const messageText = JSON.stringify(message);
let sentCount = 0;
for (const [connectionId, connection] of this.activeConnections) {
if (connection.ws.readyState === WebSocket.OPEN) {
try {
connection.ws.send(messageText);
sentCount++;
} catch (error) {
logger.error(`Failed to broadcast to ${connectionId}:`, error);
}
}
}
logger.debug(`Broadcast message sent to ${sentCount} connections`);
return sentCount;
}
/**
* Send message to specific connection
*/
sendToConnection(connectionId, message) {
const connection = this.activeConnections.get(connectionId);
if (!connection) {
throw new Error(`Connection ${connectionId} not found`);
}
if (connection.ws.readyState !== WebSocket.OPEN) {
throw new Error(`Connection ${connectionId} is not open`);
}
try {
connection.ws.send(JSON.stringify(message));
connection.bytesSent += JSON.stringify(message).length;
return true;
} catch (error) {
logger.error(`Failed to send message to ${connectionId}:`, error);
throw error;
}
}
/**
* Health check for the bridge
*/
healthCheck() {
const status = this.getStatus();
const health = {
healthy: true,
timestamp: new Date().toISOString(),
checks: {
server: !!this.wss,
connections: status.activeConnections >= 0,
processes: status.mcpProcesses >= 0,
heartbeat: !!this.heartbeatInterval
}
};
health.healthy = Object.values(health.checks).every(check => check);
return health;
}
}
export default WebSocketBridge;