apps/backend/src/mcp/mcpWebsocketServer.ts

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
    });
  }
}