gochat/backend/internal/ws/heartbeat.go

// Package ws provides heartbeat and presence lifecycle management for
// WebSocket connections. Handles periodic ping/pong, presence refresh,
// and presence broadcast on connect/disconnect.
package ws

import (
	"context"
	"time"

	"github.com/gochat/gochat/pkg/logger"
)

// HeartbeatConfig holds configuration for WebSocket heartbeat and presence intervals.
type HeartbeatConfig struct {
	// PingInterval is how often the server sends ping frames to clients.
	// Default: 30 seconds (mirrors Chatwoot's ActionCable ping interval).
	PingInterval time.Duration

	// WriteTimeout is the maximum time to wait for a pong response after ping.
	// Default: 10 seconds. Connections that don't respond are disconnected.
	WriteTimeout time.Duration

	// PresenceRefreshInterval is how often agent presence is refreshed in Redis.
	// Default: 10 seconds. Must be < PresenceDurationAgentSec (20s) to stay online.
	PresenceRefreshInterval time.Duration

	// PresenceCleanupInterval is how often expired presence entries are cleaned.
	// Default: 30 seconds.
	PresenceCleanupInterval time.Duration
}

// DefaultHeartbeatConfig returns sensible defaults mirroring Chatwoot.
func DefaultHeartbeatConfig() *HeartbeatConfig {
	return &HeartbeatConfig{
		PingInterval:             30 * time.Second,
		WriteTimeout:             10 * time.Second,
		PresenceRefreshInterval:  10 * time.Second,
		PresenceCleanupInterval:  30 * time.Second,
	}
}

// PresenceManager coordinates presence lifecycle for WebSocket connections.
// When a client connects, it marks the user as online; when disconnecting,
// it marks them as offline. Between connect/disconnect, it periodically
// refreshes the agent's presence in Redis to prevent expiry-based cleanup.
//
// Reference: Chatwoot's OnlineStatusTracker is refreshed by ActionCable
// ping/pong — each ping from the server triggers a presence refresh
// for the connected agent.
type PresenceManager struct {
	presence *PresenceTracker
	config   *HeartbeatConfig
}

// NewPresenceManager creates a presence lifecycle manager.
func NewPresenceManager(presence *PresenceTracker, config *HeartbeatConfig) *PresenceManager {
	return &PresenceManager{
		presence: presence,
		config:   config,
	}
}

// OnAgentConnect marks an agent as online when they establish a WebSocket connection.
// Starts a background goroutine that periodically refreshes their presence.
//
// Returns a context.CancelFunc that should be called when the agent disconnects
// to stop the refresh loop and mark them as offline.
func (pm *PresenceManager) OnAgentConnect(ctx context.Context, agentID, accountID uint) context.CancelFunc {
	// Mark online immediately
	if err := pm.presence.SetAgentOnline(ctx, agentID, accountID); err != nil {
		logger.L().Warnf("ws heartbeat: failed to mark agent %d online: %v", agentID, err)
	}

	// Start periodic presence refresh
	refreshCtx, cancel := context.WithCancel(ctx)
	go pm.agentPresenceRefreshLoop(refreshCtx, agentID, accountID)

	return cancel
}

// OnAgentDisconnect marks an agent as offline when their WebSocket connection closes.
// The cancel function from OnAgentConnect should also be called to stop refresh loops.
func (pm *PresenceManager) OnAgentDisconnect(ctx context.Context, agentID, accountID uint) {
	if err := pm.presence.SetAgentOffline(ctx, agentID, accountID); err != nil {
		logger.L().Warnf("ws heartbeat: failed to mark agent %d offline: %v", agentID, err)
	}
}

// OnContactConnect marks a contact as online when they establish a WebSocket connection.
func (pm *PresenceManager) OnContactConnect(ctx context.Context, contactID, accountID uint) {
	if err := pm.presence.SetContactOnline(ctx, contactID, accountID); err != nil {
		logger.L().Warnf("ws heartbeat: failed to mark contact %d online: %v", contactID, err)
	}
}

// OnContactDisconnect marks a contact as offline when their WebSocket connection closes.
func (pm *PresenceManager) OnContactDisconnect(ctx context.Context, contactID, accountID uint) {
	if err := pm.presence.SetContactOffline(ctx, contactID, accountID); err != nil {
		logger.L().Warnf("ws heartbeat: failed to mark contact %d offline: %v", contactID, err)
	}
}

// StartPresenceCleanup starts a background goroutine that periodically
// removes expired presence entries from Redis. Returns a cancel function
// to stop the cleanup loop.
func (pm *PresenceManager) StartPresenceCleanup(ctx context.Context) context.CancelFunc {
	cleanupCtx, cancel := context.WithCancel(ctx)
	go pm.presenceCleanupLoop(cleanupCtx)
	return cancel
}

// agentPresenceRefreshLoop periodically refreshes an agent's presence in Redis.
// Stops when the context is cancelled (i.e., when the agent disconnects).
func (pm *PresenceManager) agentPresenceRefreshLoop(ctx context.Context, agentID, accountID uint) {
	ticker := time.NewTicker(pm.config.PresenceRefreshInterval)
	defer ticker.Stop()

	for {
		select {
		case <-ctx.Done():
			logger.L().Debugf("ws heartbeat: presence refresh stopped for agent %d", agentID)
			return
		case <-ticker.C:
			if err := pm.presence.RefreshAgentPresence(ctx, agentID, accountID); err != nil {
				logger.L().Warnf("ws heartbeat: failed to refresh agent %d presence: %v", agentID, err)
			}
		}
	}
}

// presenceCleanupLoop periodically removes expired presence entries from Redis.
func (pm *PresenceManager) presenceCleanupLoop(ctx context.Context) {
	ticker := time.NewTicker(pm.config.PresenceCleanupInterval)
	defer ticker.Stop()

	for {
		select {
		case <-ctx.Done():
			logger.L().Info("ws heartbeat: presence cleanup loop stopped")
			return
		case <-ticker.C:
			if err := pm.presence.CleanupExpired(ctx); err != nil {
				logger.L().Warnf("ws heartbeat: presence cleanup failed: %v", err)
			}
		}
	}
}