File size: 4,048 Bytes
1cf0854
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
import { WebSocketServer, WebSocket } from 'ws';
import { IncomingMessage } from 'http';
import { parse } from 'url';

interface ExtendedWebSocket extends WebSocket {
  tenantId?: string;
  sessionToken?: string;
  isAlive?: boolean;
}

export function setupWebSocket(wss: WebSocketServer) {
  const clients = new Map<string, ExtendedWebSocket>();

  wss.on('connection', (ws: ExtendedWebSocket, req: IncomingMessage) => {
    // Parse query parameters
    const query = parse(req.url || '', true).query;
    const tenantId = query.tenantId as string;
    const sessionToken = query.sessionToken as string;

    if (!tenantId || !sessionToken) {
      ws.close(1008, 'Missing tenantId or sessionToken');
      return;
    }

    // Set up connection
    ws.tenantId = tenantId;
    ws.sessionToken = sessionToken;
    ws.isAlive = true;

    // Store client connection
    const clientKey = `${tenantId}:${sessionToken}`;
    clients.set(clientKey, ws);

    console.log(`WebSocket client connected: ${clientKey}`);

    // Handle ping/pong for connection health
    ws.on('pong', () => {
      ws.isAlive = true;
    });

    ws.on('message', (data: Buffer) => {
      try {
        const message = JSON.parse(data.toString());
        handleMessage(ws, message);
      } catch (error) {
        console.error('Invalid WebSocket message:', error);
        ws.send(JSON.stringify({ error: 'Invalid message format' }));
      }
    });

    ws.on('close', () => {
      console.log(`WebSocket client disconnected: ${clientKey}`);
      clients.delete(clientKey);
    });

    ws.on('error', (error) => {
      console.error('WebSocket error:', error);
      clients.delete(clientKey);
    });

    // Send welcome message
    ws.send(JSON.stringify({
      type: 'connection',
      message: 'Connected to chat server',
      timestamp: new Date().toISOString()
    }));
  });

  // Health check - ping clients every 30 seconds
  const interval = setInterval(() => {
    wss.clients.forEach((ws: ExtendedWebSocket) => {
      if (ws.isAlive === false) {
        return ws.terminate();
      }
      
      ws.isAlive = false;
      ws.ping();
    });
  }, 30000);

  wss.on('close', () => {
    clearInterval(interval);
  });

  function handleMessage(ws: ExtendedWebSocket, message: any) {
    switch (message.type) {
      case 'ping':
        ws.send(JSON.stringify({ type: 'pong', timestamp: new Date().toISOString() }));
        break;
        
      case 'typing':
        // Broadcast typing indicator to other participants (if multi-user chat)
        broadcastToSession(ws.tenantId!, ws.sessionToken!, {
          type: 'typing',
          sender: 'user',
          timestamp: new Date().toISOString()
        }, ws);
        break;
        
      case 'stop_typing':
        broadcastToSession(ws.tenantId!, ws.sessionToken!, {
          type: 'stop_typing',
          sender: 'user',
          timestamp: new Date().toISOString()
        }, ws);
        break;
        
      default:
        ws.send(JSON.stringify({ error: 'Unknown message type' }));
    }
  }

  function broadcastToSession(tenantId: string, sessionToken: string, message: any, sender?: WebSocket) {
    const clientKey = `${tenantId}:${sessionToken}`;
    const client = clients.get(clientKey);
    
    if (client && client !== sender && client.readyState === WebSocket.OPEN) {
      client.send(JSON.stringify(message));
    }
  }

  // Utility function to send message to specific session
  function sendToSession(tenantId: string, sessionToken: string, message: any) {
    const clientKey = `${tenantId}:${sessionToken}`;
    const client = clients.get(clientKey);
    
    if (client && client.readyState === WebSocket.OPEN) {
      client.send(JSON.stringify(message));
      return true;
    }
    return false;
  }

  return {
    broadcastToSession,
    sendToSession,
    getConnectedClients: () => clients.size
  };
}