// Package ws (handler/ws) provides the WebSocket Hub that manages connected // clients, room subscriptions, and message routing. Enhanced with Redis // Pub/Sub relay for cross-instance broadcasting. package ws import ( "context" "crypto/rand" "encoding/hex" "encoding/json" "fmt" "sync" "time" "github.com/gorilla/websocket" wspkg "github.com/gochat/gochat/internal/ws" "github.com/gochat/gochat/pkg/logger" ) // Client represents a connected WebSocket client. // Each client has a unique ID, user identity from auth claims, // and a buffered Send channel for outgoing messages. type Client struct { ID string // unique connection ID (uuid) UserID uint // from WSClaims.UserID AccountID uint // from WSClaims.AccountID Role string // from WSClaims.Role IsContact bool // from WSClaims.IsContact PubsubToken string // from WSClaims.PubsubToken ContactID uint // from WSClaims.ContactID (only for contacts) InboxID uint // from WSClaims.InboxID (only for contacts) Conn *websocket.Conn // gorilla/websocket connection Send chan []byte // buffered outgoing message channel (256 capacity) Hub *Hub // reference back to Hub SubscribedRooms map[string]bool // rooms this client is subscribed to CancelPresence context.CancelFunc // cancel presence refresh on disconnect } // NewClient creates a new WebSocket client with the given identity and connection. func NewClient(userID, accountID uint, conn *websocket.Conn, hub *Hub) *Client { return &Client{ ID: generateClientID(), UserID: userID, AccountID: accountID, Conn: conn, Send: make(chan []byte, SendChannelSize), Hub: hub, SubscribedRooms: make(map[string]bool), } } // Subscribe adds the client to a room and registers with the Hub. func (c *Client) Subscribe(room string) { c.Hub.mu.Lock() c.SubscribedRooms[room] = true c.Hub.subscribeClient(c.ID, room) c.Hub.mu.Unlock() } // Unsubscribe removes the client from a room and unregisters with the Hub. func (c *Client) Unsubscribe(room string) { c.Hub.mu.Lock() delete(c.SubscribedRooms, room) c.Hub.unsubscribeClient(c.ID, room) c.Hub.mu.Unlock() } // generateClientID creates a unique client ID using crypto/rand. func generateClientID() string { b := make([]byte, 16) _, _ = rand.Read(b) return hex.EncodeToString(b) } const ( // MaxMessageSize is the maximum size of a message that can be read from a client. MaxMessageSize = 512 * 1024 // 512KB // SendChannelSize is the buffer size for client Send channels. SendChannelSize = 256 // WriteWait is the time allowed to write a message to the peer. WriteWait = 10 * time.Second // PongWait is the time allowed to read the next pong message from the peer. PongWait = 60 * time.Second // PingPeriod is how often pings are sent to peers. Must be < PongWait. PingPeriod = (PongWait * 9) / 10 ) // Hub maintains the set of active clients and broadcasts messages to them. // Enhanced with Redis Pub/Sub relay for cross-instance message delivery, // typing indicator support, and agent/contact presence tracking. type Hub struct { // Registered clients indexed by connection ID clients map[string]*Client // Room subscriptions: room name → set of client IDs in that room rooms map[string]map[string]bool // Account subscriptions: account ID → set of client IDs accounts map[uint]map[string]bool // Broadcast relay for cross-instance message delivery via Redis Pub/Sub relay *wspkg.BroadcastRelay // Typing tracker for typing indicator management typing *wspkg.TypingTracker // Presence tracker for online/offline status presence *wspkg.PresenceTracker // Presence lifecycle manager presenceMgr *wspkg.PresenceManager // Inbound messages from clients (commands) commandChan chan *ClientCommand mu sync.RWMutex } // ClientCommand wraps a command from a client with the client reference. type ClientCommand struct { Client *Client Cmd wspkg.WSCommand } // NewHub creates a new Hub with Redis Pub/Sub relay and presence subsystems. func NewHub(relay *wspkg.BroadcastRelay, typing *wspkg.TypingTracker, presence *wspkg.PresenceTracker, presenceMgr *wspkg.PresenceManager) *Hub { return &Hub{ clients: make(map[string]*Client), rooms: make(map[string]map[string]bool), accounts: make(map[uint]map[string]bool), relay: relay, typing: typing, presence: presence, presenceMgr: presenceMgr, commandChan: make(chan *ClientCommand, 256), } } // NewHubSimple creates a minimal Hub without Redis subsystems. // Used for development/testing when Redis is not available. func NewHubSimple() *Hub { return &Hub{ clients: make(map[string]*Client), rooms: make(map[string]map[string]bool), accounts: make(map[uint]map[string]bool), commandChan: make(chan *ClientCommand, 256), } } // Run starts the Hub's main event loop. Processes client registrations, // unregistrations, and commands. Must be called in a goroutine. func (h *Hub) Run(ctx context.Context) { logger.L().Info("ws hub: starting event loop") for { select { case <-ctx.Done(): logger.L().Info("ws hub: shutting down event loop") h.shutdown() return case cmd := <-h.commandChan: h.processCommand(cmd) } } } // Register adds a client to the Hub and sets up presence. func (h *Hub) Register(c *Client) { h.mu.Lock() defer h.mu.Unlock() h.clients[c.ID] = c // Auto-subscribe to account room on connect roomName := accountRoomName(c.AccountID) h.subscribeClient(c.ID, roomName) c.SubscribedRooms[roomName] = true // Set up presence tracking if h.presenceMgr != nil { if c.IsContact { h.presenceMgr.OnContactConnect(context.Background(), c.ContactID, c.AccountID) // Also auto-subscribe to pubsub_token room (Chatwoot RoomChannel pattern) if c.PubsubToken != "" { tokenRoom := pubsubTokenRoomName(c.PubsubToken) h.subscribeClient(c.ID, tokenRoom) c.SubscribedRooms[tokenRoom] = true } } else { cancelFn := h.presenceMgr.OnAgentConnect(context.Background(), c.UserID, c.AccountID) c.CancelPresence = cancelFn } } logger.L().Infof("ws hub: client registered (id=%s, user_id=%d, account_id=%d, is_contact=%v)", c.ID, c.UserID, c.AccountID, c.IsContact) // Send welcome message welcomeMsg, _ := json.Marshal(wspkg.WSMessage{ Event: wspkg.EventWelcome, Data: map[string]any{"client_id": c.ID}, }) c.Send <- welcomeMsg } // Unregister removes a client from the Hub and cleans up presence. func (h *Hub) Unregister(c *Client) { h.mu.Lock() defer h.mu.Unlock() if _, ok := h.clients[c.ID]; !ok { return // already unregistered } // Remove from all subscribed rooms for room := range c.SubscribedRooms { h.unsubscribeClient(c.ID, room) } // Remove from Hub's client map delete(h.clients, c.ID) // Clean up presence tracking if h.presenceMgr != nil { if c.IsContact { h.presenceMgr.OnContactDisconnect(context.Background(), c.ContactID, c.AccountID) } else { if c.CancelPresence != nil { c.CancelPresence() // stop presence refresh loop } h.presenceMgr.OnAgentDisconnect(context.Background(), c.UserID, c.AccountID) } } // Close Send channel close(c.Send) logger.L().Infof("ws hub: client unregistered (id=%s, user_id=%d)", c.ID, c.UserID) } // SendToAccount sends a message to all clients subscribed to an account room. func (h *Hub) SendToAccount(accountID uint, data []byte) { h.mu.RLock() defer h.mu.RUnlock() roomName := accountRoomName(accountID) if clientIDs, ok := h.rooms[roomName]; ok { for clientID := range clientIDs { if client, ok := h.clients[clientID]; ok { select { case client.Send <- data: default: logger.L().Warnf("ws hub: dropping message for slow client %s", clientID) } } } } } // SendToAccountConversation sends a message to all clients subscribed to // a specific conversation within an account. func (h *Hub) SendToAccountConversation(accountID uint, conversationID uint, data []byte) { h.mu.RLock() defer h.mu.RUnlock() roomName := conversationRoomName(accountID, conversationID) if clientIDs, ok := h.rooms[roomName]; ok { for clientID := range clientIDs { if client, ok := h.clients[clientID]; ok { select { case client.Send <- data: default: logger.L().Warnf("ws hub: dropping message for slow client %s", clientID) } } } } } // SendToRoom sends a message to all clients in a named room. func (h *Hub) SendToRoom(room string, data []byte) { h.mu.RLock() defer h.mu.RUnlock() if clientIDs, ok := h.rooms[room]; ok { for clientID := range clientIDs { if client, ok := h.clients[clientID]; ok { select { case client.Send <- data: default: logger.L().Warnf("ws hub: dropping message for slow client %s", clientID) } } } } } // SendToClient sends a message to a specific client by ID. func (h *Hub) SendToClient(clientID string, data []byte) { h.mu.RLock() defer h.mu.RUnlock() if client, ok := h.clients[clientID]; ok { select { case client.Send <- data: default: logger.L().Warnf("ws hub: dropping message for slow client %s", clientID) } } } // SubmitCommand enqueues a client command for the Hub's event loop to process. func (h *Hub) SubmitCommand(client *Client, cmd wspkg.WSCommand) { h.commandChan <- &ClientCommand{Client: client, Cmd: cmd} } // processCommand handles a client command in the Hub's event loop. func (h *Hub) processCommand(cmd *ClientCommand) { switch cmd.Cmd.Command { case "subscribe": h.handleSubscribe(cmd) case "unsubscribe": h.handleUnsubscribe(cmd) case "ping": h.handlePing(cmd) case "typing_on": h.handleTypingOn(cmd) case "typing_off": h.handleTypingOff(cmd) case "update_presence": h.handleUpdatePresence(cmd) default: logger.L().Warnf("ws hub: unknown command '%s' from client %s", cmd.Cmd.Command, cmd.Client.ID) } } // handleSubscribe processes a subscribe command. // Mirrors Chatwoot's RoomChannel.subscribe flow. func (h *Hub) handleSubscribe(cmd *ClientCommand) { var data wspkg.SubscribeData if err := json.Unmarshal([]byte(cmd.Cmd.Data), &data); err != nil { logger.L().Warnf("ws hub: invalid subscribe data from client %s: %v", cmd.Client.ID, err) return } // Verify account access if data.AccountID != cmd.Client.AccountID { rejectMsg, _ := json.Marshal(wspkg.WSMessage{ Event: wspkg.EventSubscribeReject, Data: map[string]any{"reason": "account_id mismatch"}, }) cmd.Client.Send <- rejectMsg return } var roomName string switch data.Channel { case wspkg.ChannelAccount: roomName = accountRoomName(data.AccountID) case wspkg.ChannelConversation: if data.ConversationID == 0 { rejectMsg, _ := json.Marshal(wspkg.WSMessage{ Event: wspkg.EventSubscribeReject, Data: map[string]any{"reason": "conversation_id required for ConversationChannel"}, }) cmd.Client.Send <- rejectMsg return } roomName = conversationRoomName(data.AccountID, data.ConversationID) default: rejectMsg, _ := json.Marshal(wspkg.WSMessage{ Event: wspkg.EventSubscribeReject, Data: map[string]any{"reason": "unknown channel type"}, }) cmd.Client.Send <- rejectMsg return } h.mu.Lock() h.subscribeClient(cmd.Client.ID, roomName) cmd.Client.SubscribedRooms[roomName] = true h.mu.Unlock() confirmMsg, _ := json.Marshal(wspkg.WSMessage{ Event: wspkg.EventSubscribeConfirm, Data: map[string]any{"room": roomName, "channel": data.Channel}, }) cmd.Client.Send <- confirmMsg logger.L().Infof("ws hub: client %s subscribed to room %s", cmd.Client.ID, roomName) } // handleUnsubscribe processes an unsubscribe command. func (h *Hub) handleUnsubscribe(cmd *ClientCommand) { var data wspkg.SubscribeData if err := json.Unmarshal([]byte(cmd.Cmd.Data), &data); err != nil { logger.L().Warnf("ws hub: invalid unsubscribe data from client %s: %v", cmd.Client.ID, err) return } var roomName string switch data.Channel { case wspkg.ChannelAccount: roomName = accountRoomName(data.AccountID) case wspkg.ChannelConversation: roomName = conversationRoomName(data.AccountID, data.ConversationID) default: return } h.mu.Lock() h.unsubscribeClient(cmd.Client.ID, roomName) delete(cmd.Client.SubscribedRooms, roomName) h.mu.Unlock() confirmMsg, _ := json.Marshal(wspkg.WSMessage{ Event: wspkg.EventUnsubscribeConfirm, Data: map[string]any{"room": roomName}, }) cmd.Client.Send <- confirmMsg logger.L().Infof("ws hub: client %s unsubscribed from room %s", cmd.Client.ID, roomName) } // handlePing responds to a ping command from the client. func (h *Hub) handlePing(cmd *ClientCommand) { pingMsg, _ := json.Marshal(wspkg.WSMessage{ Event: wspkg.EventPingResponse, Data: map[string]any{"timestamp": time.Now().UnixMilli()}, }) cmd.Client.Send <- pingMsg } // handleTypingOn processes a typing_on command. func (h *Hub) handleTypingOn(cmd *ClientCommand) { if h.typing == nil { return } var data wspkg.TypingData if err := json.Unmarshal([]byte(cmd.Cmd.Data), &data); err != nil { logger.L().Warnf("ws hub: invalid typing_on data from client %s: %v", cmd.Client.ID, err) return } performerType := "user" if cmd.Client.IsContact { performerType = "contact" } performer := &wspkg.Performer{ ID: cmd.Client.UserID, Type: performerType, } if err := h.typing.SetTypingOn(context.Background(), data.AccountID, data.ConversationID, performer); err != nil { logger.L().Warnf("ws hub: typing_on failed: %v", err) } } // handleTypingOff processes a typing_off command. func (h *Hub) handleTypingOff(cmd *ClientCommand) { if h.typing == nil { return } var data wspkg.TypingData if err := json.Unmarshal([]byte(cmd.Cmd.Data), &data); err != nil { logger.L().Warnf("ws hub: invalid typing_off data from client %s: %v", cmd.Client.ID, err) return } performerType := "user" if cmd.Client.IsContact { performerType = "contact" } performer := &wspkg.Performer{ ID: cmd.Client.UserID, Type: performerType, } if err := h.typing.SetTypingOff(context.Background(), data.AccountID, data.ConversationID, performer); err != nil { logger.L().Warnf("ws hub: typing_off failed: %v", err) } } // handleUpdatePresence processes an update_presence command. func (h *Hub) handleUpdatePresence(cmd *ClientCommand) { if h.presence == nil { return } var data wspkg.PresenceData if err := json.Unmarshal([]byte(cmd.Cmd.Data), &data); err != nil { logger.L().Warnf("ws hub: invalid update_presence data from client %s: %v", cmd.Client.ID, err) return } ctx := context.Background() switch data.Status { case "online": h.presence.SetAgentOnline(ctx, cmd.Client.UserID, cmd.Client.AccountID) case "busy": h.presence.SetAgentBusy(ctx, cmd.Client.UserID, cmd.Client.AccountID) case "offline": h.presence.SetAgentOffline(ctx, cmd.Client.UserID, cmd.Client.AccountID) default: logger.L().Warnf("ws hub: unknown presence status '%s' from client %s", data.Status, cmd.Client.ID) } } // --- Internal helpers --- // subscribeClient adds a client ID to a room's member set. // Must be called with h.mu held. func (h *Hub) subscribeClient(clientID string, room string) { if _, ok := h.rooms[room]; !ok { h.rooms[room] = make(map[string]bool) } h.rooms[room][clientID] = true } // unsubscribeClient removes a client ID from a room's member set. // Must be called with h.mu held. func (h *Hub) unsubscribeClient(clientID string, room string) { if clients, ok := h.rooms[room]; ok { delete(clients, clientID) if len(clients) == 0 { delete(h.rooms, room) } } } // shutdown disconnects all clients and cleans up. func (h *Hub) shutdown() { h.mu.Lock() defer h.mu.Unlock() for id, client := range h.clients { if client.CancelPresence != nil { client.CancelPresence() } close(client.Send) client.Conn.Close() delete(h.clients, id) } h.rooms = make(map[string]map[string]bool) h.accounts = make(map[uint]map[string]bool) logger.L().Info("ws hub: all clients disconnected, shutdown complete") } // Shutdown gracefully closes the Hub, disconnecting all WebSocket clients. // This is the public API called during application graceful shutdown. // It sends a close message to each client and waits for them to finish, // respecting the context deadline. func (h *Hub) Shutdown(ctx context.Context) { logger.L().Info("ws hub: initiating graceful shutdown") h.shutdown() } // Room name helpers following Chatwoot naming conventions. func accountRoomName(accountID uint) string { return fmt.Sprintf("account_%d", accountID) } func conversationRoomName(accountID uint, conversationID uint) string { return fmt.Sprintf("account_%d_conversation_%d", accountID, conversationID) } func pubsubTokenRoomName(token string) string { return fmt.Sprintf("pubsub_token_%s", token) }