Spaces:
Paused
Paused
File size: 4,137 Bytes
34367da | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 |
import { WebSocketServer, WebSocket } from 'ws';
import { Server } from 'http';
import { MCPMessage } from '@widget-tdc/mcp-types';
import { mcpRegistry } from './mcpRegistry.js';
import rateLimit from 'express-rate-limit';
import { z } from 'zod';
// Define Zod schema for incoming MCP messages
const mcpMessageSchema = z.object({
id: z.string(),
// Add other properties of MCPMessage here
});
export class MCPWebSocketServer {
private wss: WebSocketServer;
private clients: Map<string, WebSocket> = new Map();
constructor(server: Server) {
this.wss = new WebSocketServer({ server, path: '/mcp/ws' });
this.setupWebSocketServer();
}
private setupWebSocketServer(): void {
const limiter = rateLimit({
windowMs: 15 * 60 * 1000, // 15 minutes
max: 100, // Limit each IP to 100 requests per windowMs
message: 'Too many requests from this IP, please try again after 15 minutes'
});
this.wss.on('connection', (ws: WebSocket, req) => {
limiter(req as any, {} as any, () => {});
const clientId = Math.random().toString(36).substring(7);
this.clients.set(clientId, ws);
console.log(`MCP WebSocket client connected: ${clientId}`);
ws.on('message', async (data: Buffer) => {
try {
const rawMessage = data.toString();
const message = JSON.parse(rawMessage);
// Validate the message against the Zod schema
mcpMessageSchema.parse(message);
// Route the message
const result = await mcpRegistry.route(message as MCPMessage);
// Send response back to client
ws.send(JSON.stringify({
success: true,
messageId: message.id,
result,
}));
// Broadcast to other clients if needed
this.broadcast(message, clientId);
} catch (error: any) {
ws.send(JSON.stringify({
success: false,
error: error.message,
}));
}
});
ws.on('close', () => {
this.clients.delete(clientId);
console.log(`MCP WebSocket client disconnected: ${clientId}`);
});
// Send welcome message
ws.send(JSON.stringify({
type: 'welcome',
clientId,
availableTools: mcpRegistry.getRegisteredTools(),
}));
});
}
private broadcast(message: MCPMessage, excludeClientId?: string): void {
const data = JSON.stringify({
type: 'broadcast',
message,
});
this.clients.forEach((client, clientId) => {
if (clientId !== excludeClientId && client.readyState === WebSocket.OPEN) {
client.send(data);
}
});
}
public sendToAll(message: any): void {
const data = JSON.stringify(message);
this.clients.forEach((client) => {
if (client.readyState === WebSocket.OPEN) {
client.send(data);
}
});
}
/**
* Emit autonomous decision event to all connected clients
*/
public emitAutonomousDecision(decision: {
queryId: string;
selectedSource: string;
confidence: number;
alternatives: string[];
reasoning: string;
latency: number;
}): void {
this.sendToAll({
type: 'autonomous:decision',
timestamp: new Date().toISOString(),
data: decision
});
}
/**
* Emit source health update
*/
public emitSourceHealth(sourceName: string, health: {
healthy: boolean;
score: number;
latency?: number;
}): void {
this.sendToAll({
type: 'autonomous:health',
timestamp: new Date().toISOString(),
data: {
source: sourceName,
...health
}
});
}
/**
* Emit learning progress update
*/
public emitLearningProgress(progress: {
patternsLearned: number;
decisionsMade: number;
averageConfidence: number;
}): void {
this.sendToAll({
type: 'autonomous:learning',
timestamp: new Date().toISOString(),
data: progress
});
}
}
|