claude-flow
Version:
Ruflo - Enterprise AI agent orchestration for Claude Code. Deploy 60+ specialized agents in coordinated swarms with self-learning, fault-tolerant consensus, vector memory, and MCP integration
485 lines (426 loc) • 12.3 kB
text/typescript
/**
* V3 MCP WebSocket Transport
*
* Standalone WebSocket transport for MCP communication:
* - Native WebSocket server without HTTP dependency
* - Binary message support for efficiency
* - Heartbeat/ping-pong for connection health
* - Automatic reconnection handling
*
* Performance Targets:
* - Message latency: <3ms
* - Connection overhead: <10ms
*/
import { EventEmitter } from 'events';
import { WebSocketServer, WebSocket, RawData } from 'ws';
import { createServer, Server } from 'http';
import {
ITransport,
TransportType,
MCPRequest,
MCPResponse,
MCPNotification,
RequestHandler,
NotificationHandler,
TransportHealthStatus,
ILogger,
AuthConfig,
} from '../types.js';
/**
* WebSocket Transport Configuration
*/
export interface WebSocketTransportConfig {
host: string;
port: number;
path?: string;
maxConnections?: number;
heartbeatInterval?: number;
heartbeatTimeout?: number;
maxMessageSize?: number;
auth?: AuthConfig;
enableBinaryMode?: boolean;
}
/**
* Client connection info
*/
interface ClientConnection {
id: string;
ws: WebSocket;
createdAt: Date;
lastActivity: Date;
messageCount: number;
isAlive: boolean;
isAuthenticated: boolean;
}
/**
* WebSocket Transport Implementation
*/
export class WebSocketTransport extends EventEmitter implements ITransport {
public readonly type: TransportType = 'websocket';
private requestHandler?: RequestHandler;
private notificationHandler?: NotificationHandler;
private server?: Server;
private wss?: WebSocketServer;
private clients: Map<string, ClientConnection> = new Map();
private heartbeatTimer?: NodeJS.Timeout;
private running = false;
private connectionCounter = 0;
// Statistics
private messagesReceived = 0;
private messagesSent = 0;
private errors = 0;
private totalConnections = 0;
constructor(
private readonly logger: ILogger,
private readonly config: WebSocketTransportConfig
) {
super();
}
/**
* Start the transport
*/
async start(): Promise<void> {
if (this.running) {
throw new Error('WebSocket transport already running');
}
this.logger.info('Starting WebSocket transport', {
host: this.config.host,
port: this.config.port,
path: this.config.path || '/ws',
});
// Create HTTP server for WebSocket upgrade
this.server = createServer((req, res) => {
// Simple HTTP response for non-WebSocket requests
res.writeHead(426, { 'Content-Type': 'text/plain' });
res.end('Upgrade Required - WebSocket connection expected');
});
// Create WebSocket server
this.wss = new WebSocketServer({
server: this.server,
path: this.config.path || '/ws',
maxPayload: this.config.maxMessageSize || 10 * 1024 * 1024,
perMessageDeflate: true, // Enable compression
});
this.setupWebSocketHandlers();
this.startHeartbeat();
// Start server
await new Promise<void>((resolve, reject) => {
this.server!.listen(this.config.port, this.config.host, () => {
resolve();
});
this.server!.on('error', reject);
});
this.running = true;
this.logger.info('WebSocket transport started', {
url: `ws://${this.config.host}:${this.config.port}${this.config.path || '/ws'}`,
});
}
/**
* Stop the transport
*/
async stop(): Promise<void> {
if (!this.running) {
return;
}
this.logger.info('Stopping WebSocket transport');
this.running = false;
this.stopHeartbeat();
// Close all client connections
for (const client of this.clients.values()) {
try {
client.ws.close(1000, 'Server shutting down');
} catch {
// Ignore errors
}
}
this.clients.clear();
// Close WebSocket server
if (this.wss) {
this.wss.close();
this.wss = undefined;
}
// Close HTTP server
if (this.server) {
await new Promise<void>((resolve) => {
this.server!.close(() => resolve());
});
this.server = undefined;
}
this.logger.info('WebSocket transport stopped');
}
/**
* Register request handler
*/
onRequest(handler: RequestHandler): void {
this.requestHandler = handler;
}
/**
* Register notification handler
*/
onNotification(handler: NotificationHandler): void {
this.notificationHandler = handler;
}
/**
* Get health status
*/
async getHealthStatus(): Promise<TransportHealthStatus> {
return {
healthy: this.running,
metrics: {
messagesReceived: this.messagesReceived,
messagesSent: this.messagesSent,
errors: this.errors,
activeConnections: this.clients.size,
totalConnections: this.totalConnections,
},
};
}
/**
* Send notification to all connected clients
*/
async sendNotification(notification: MCPNotification): Promise<void> {
const message = this.serializeMessage(notification);
for (const client of this.clients.values()) {
try {
if (client.ws.readyState === WebSocket.OPEN) {
client.ws.send(message);
this.messagesSent++;
}
} catch (error) {
this.logger.error('Failed to send notification', { clientId: client.id, error });
this.errors++;
}
}
}
/**
* Send notification to specific client
*/
async sendToClient(clientId: string, notification: MCPNotification): Promise<boolean> {
const client = this.clients.get(clientId);
if (!client || client.ws.readyState !== WebSocket.OPEN) {
return false;
}
try {
client.ws.send(this.serializeMessage(notification));
this.messagesSent++;
return true;
} catch (error) {
this.logger.error('Failed to send to client', { clientId, error });
this.errors++;
return false;
}
}
/**
* Get connected clients
*/
getClients(): string[] {
return Array.from(this.clients.keys());
}
/**
* Get client info
*/
getClientInfo(clientId: string): ClientConnection | undefined {
return this.clients.get(clientId);
}
/**
* Disconnect specific client
*/
disconnectClient(clientId: string, reason = 'Disconnected by server'): boolean {
const client = this.clients.get(clientId);
if (!client) {
return false;
}
try {
client.ws.close(1000, reason);
return true;
} catch {
return false;
}
}
/**
* Setup WebSocket handlers
*/
private setupWebSocketHandlers(): void {
if (!this.wss) return;
this.wss.on('connection', (ws, req) => {
// Check max connections
if (this.config.maxConnections && this.clients.size >= this.config.maxConnections) {
this.logger.warn('Max connections reached, rejecting client');
ws.close(1013, 'Server at capacity');
return;
}
const clientId = `client-${++this.connectionCounter}`;
const client: ClientConnection = {
id: clientId,
ws,
createdAt: new Date(),
lastActivity: new Date(),
messageCount: 0,
isAlive: true,
isAuthenticated: !this.config.auth?.enabled,
};
this.clients.set(clientId, client);
this.totalConnections++;
this.logger.info('Client connected', {
id: clientId,
total: this.clients.size,
});
// Setup message handler
ws.on('message', async (data) => {
await this.handleMessage(client, data);
});
// Setup pong handler for heartbeat
ws.on('pong', () => {
client.isAlive = true;
});
// Setup close handler
ws.on('close', (code, reason) => {
this.clients.delete(clientId);
this.logger.info('Client disconnected', {
id: clientId,
code,
reason: reason.toString(),
total: this.clients.size,
});
this.emit('client:disconnected', clientId);
});
// Setup error handler
ws.on('error', (error) => {
this.logger.error('Client error', { id: clientId, error });
this.errors++;
this.clients.delete(clientId);
});
this.emit('client:connected', clientId);
});
}
/**
* Handle incoming message
*/
private async handleMessage(client: ClientConnection, data: RawData): Promise<void> {
client.lastActivity = new Date();
client.messageCount++;
this.messagesReceived++;
try {
const message = this.parseMessage(data);
// Check authentication for non-authenticated clients
if (!client.isAuthenticated && this.config.auth?.enabled) {
if (message.method !== 'authenticate') {
client.ws.send(this.serializeMessage({
jsonrpc: '2.0',
id: message.id || null,
error: { code: -32001, message: 'Authentication required' },
} as MCPResponse));
return;
}
}
if (message.jsonrpc !== '2.0') {
client.ws.send(this.serializeMessage({
jsonrpc: '2.0',
id: message.id || null,
error: { code: -32600, message: 'Invalid JSON-RPC version' },
} as MCPResponse));
return;
}
if (message.id === undefined) {
// Notification
if (this.notificationHandler) {
await this.notificationHandler(message as MCPNotification);
}
} else {
// Request
if (!this.requestHandler) {
client.ws.send(this.serializeMessage({
jsonrpc: '2.0',
id: message.id,
error: { code: -32603, message: 'No request handler' },
} as MCPResponse));
return;
}
const startTime = performance.now();
const response = await this.requestHandler(message as MCPRequest);
const duration = performance.now() - startTime;
this.logger.debug('Request processed', {
clientId: client.id,
method: message.method,
duration: `${duration.toFixed(2)}ms`,
});
client.ws.send(this.serializeMessage(response));
this.messagesSent++;
}
} catch (error) {
this.errors++;
this.logger.error('Message handling error', { clientId: client.id, error });
try {
client.ws.send(this.serializeMessage({
jsonrpc: '2.0',
id: null,
error: { code: -32700, message: 'Parse error' },
} as MCPResponse));
} catch {
// Ignore send errors
}
}
}
/**
* Parse incoming message
*/
private parseMessage(data: RawData): any {
if (this.config.enableBinaryMode && Buffer.isBuffer(data)) {
// Could implement binary protocol here
return JSON.parse(data.toString());
}
return JSON.parse(data.toString());
}
/**
* Serialize outgoing message
*/
private serializeMessage(message: MCPResponse | MCPNotification): string | Buffer {
if (this.config.enableBinaryMode) {
// Could implement binary protocol here
return JSON.stringify(message);
}
return JSON.stringify(message);
}
/**
* Start heartbeat interval
*/
private startHeartbeat(): void {
const interval = this.config.heartbeatInterval || 30000; // 30 seconds
const timeout = this.config.heartbeatTimeout || 10000; // 10 seconds
this.heartbeatTimer = setInterval(() => {
for (const client of this.clients.values()) {
if (!client.isAlive) {
// Client didn't respond to last ping
this.logger.warn('Client heartbeat timeout', { id: client.id });
client.ws.terminate();
this.clients.delete(client.id);
continue;
}
client.isAlive = false;
try {
client.ws.ping();
} catch {
// Ignore ping errors
}
}
}, interval);
}
/**
* Stop heartbeat interval
*/
private stopHeartbeat(): void {
if (this.heartbeatTimer) {
clearInterval(this.heartbeatTimer);
this.heartbeatTimer = undefined;
}
}
}
/**
* Create WebSocket transport
*/
export function createWebSocketTransport(
logger: ILogger,
config: WebSocketTransportConfig
): WebSocketTransport {
return new WebSocketTransport(logger, config);
}