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