MailUI4Agents/gateway/internal/sse/manager.go

package sse

import (
	"encoding/json"
	"fmt"
	"net/http"
	"sync"
	"time"

	"github.com/google/uuid"
)

// Client 是一个 SSE 连接客户端
type Client struct {
	ID        string
	AgentName string // 非空 = Agent 侧连接
	UserName  string // 非空 = 已登录人类用户的前端连接
	Res       http.ResponseWriter
	Flusher   http.Flusher
	done      chan struct{}
}

// Manager 管理所有 SSE 客户端连接
type Manager struct {
	mu      sync.RWMutex
	clients map[string]*Client
}

var Default = &Manager{
	clients: make(map[string]*Client),
}

// AddClient 注册一个新 SSE 客户端（agentName 与 userName 二者恰其一）
func (m *Manager) AddClient(res http.ResponseWriter, agentName, userName string) *Client {
	flusher, ok := res.(http.Flusher)
	if !ok {
		return nil
	}

	id := uuid.New().String()[:8]
	client := &Client{
		ID:        id,
		AgentName: agentName,
		UserName:  userName,
		Res:       res,
		Flusher:   flusher,
		done:      make(chan struct{}),
	}

	// 设置 SSE 响应头
	res.Header().Set("Content-Type", "text/event-stream")
	res.Header().Set("Cache-Control", "no-cache")
	res.Header().Set("Connection", "keep-alive")
	res.Header().Set("X-Accel-Buffering", "no")

	m.mu.Lock()
	m.clients[id] = client
	m.mu.Unlock()

	// 发送连接确认
	client.Send("connected", map[string]string{"id": id})

	// 启动心跳
	go m.heartbeat(client)

	fmt.Printf("[SSE] Client connected: %s (agent=%q user=%q)\n", id, agentName, userName)
	return client
}

// RemoveClient 移除一个客户端
func (m *Manager) RemoveClient(id string) {
	m.mu.Lock()
	if c, ok := m.clients[id]; ok {
		close(c.done)
		delete(m.clients, id)
		fmt.Printf("[SSE] Client disconnected: %s\n", id)
	}
	m.mu.Unlock()
}

// SendToAgent 向指定 Agent 名的所有客户端推送事件
func (m *Manager) SendToAgent(agentName, eventType string, data interface{}) {
	if agentName == "" {
		return
	}
	m.mu.RLock()
	defer m.mu.RUnlock()

	for _, c := range m.clients {
		if c.AgentName == agentName {
			c.Send(eventType, data)
		}
	}
}

// SendToUser 向指定人类用户的所有前端连接推送事件
func (m *Manager) SendToUser(userName, eventType string, data interface{}) {
	if userName == "" {
		return
	}
	m.mu.RLock()
	defer m.mu.RUnlock()

	for _, c := range m.clients {
		if c.UserName == userName {
			c.Send(eventType, data)
		}
	}
}

// SendToRecipient 根据收件人名同时尝试 Agent 通道与人类用户通道
// （三维地址的 name 位共享命名空间，投递时不必先判断对方是人还是 Agent）
func (m *Manager) SendToRecipient(name, eventType string, data interface{}) {
	if name == "" {
		return
	}
	m.mu.RLock()
	defer m.mu.RUnlock()

	for _, c := range m.clients {
		if c.AgentName == name || c.UserName == name {
			c.Send(eventType, data)
		}
	}
}

// Broadcast 向所有客户端广播事件
func (m *Manager) Broadcast(eventType string, data interface{}) {
	m.mu.RLock()
	defer m.mu.RUnlock()

	for _, c := range m.clients {
		c.Send(eventType, data)
	}
}

// ClientCount 返回当前连接数
func (m *Manager) ClientCount() int {
	m.mu.RLock()
	defer m.mu.RUnlock()
	return len(m.clients)
}

// Send 向单个客户端发送事件
func (c *Client) Send(eventType string, data interface{}) {
	defer func() { recover() }() // 防止向已关闭的连接写入 panic

	jsonData, err := json.Marshal(data)
	if err != nil {
		return
	}

	fmt.Fprintf(c.Res, "event: %s\ndata: %s\n\n", eventType, jsonData)
	c.Flusher.Flush()
}

// heartbeat 定期发送心跳保活
func (m *Manager) heartbeat(client *Client) {
	ticker := time.NewTicker(30 * time.Second)
	defer ticker.Stop()

	for {
		select {
		case <-client.done:
			return
		case <-ticker.C:
			// SSE 注释行作为心跳
			defer func() { recover() }()
			fmt.Fprintf(client.Res, ": heartbeat\n\n")
			client.Flusher.Flush()
		}
	}
}