package sse import ( "encoding/json" "fmt" "net/http" "sync" "time" "github.com/google/uuid" ) // eventRing 是单用户事件的有界环形缓冲区。 // // EventSource 断线重连时自带 Last-Event-ID 头:服务端据此回放断线期间的事件。 // 没有它,重连后永远看不到断线期间收到的邮件 —— 而这正是实时协作的体验核心。 // // 缓冲区大小 500 条:一条事件约 200B(typical),500 条 ≈ 100KB/用户。 // 20 个在线用户 ≈ 2MB,远低于 OOM 风险。 type eventRing struct { mu sync.Mutex events []StoredEvent cap int head int // 下一次写入的位置 full bool } // StoredEvent 是缓冲区中的单条事件。 type StoredEvent struct { ID string // 自增序列号,EventSource 的 Last-Event-ID 值 EventType string Data []byte Timestamp time.Time } func newEventRing(cap int) *eventRing { return &eventRing{events: make([]StoredEvent, cap), cap: cap} } // push 追加一条事件到缓冲区。满了就覆盖最旧的。 func (r *eventRing) push(evt StoredEvent) { r.mu.Lock() defer r.mu.Unlock() r.events[r.head] = evt r.head = (r.head + 1) % r.cap if r.head == 0 && !r.full { r.full = true } } // replay 从 afterID 之后的所有事件回放给 ResponseWriter。 // afterID 为空时:缓冲区未满不回放(首次连接无历史);满了也不回放 // (首次连接的 EventSource 不传 Last-Event-ID)。 // afterID 非空时:找到该 ID 的位置,从下一条开始回放。 func (r *eventRing) replay(afterID string, flush http.Flusher, res http.ResponseWriter) { r.mu.Lock() defer r.mu.Unlock() if afterID == "" { return // 首次连接,不回放 } start := -1 total := r.cap for i := 0; i < r.cap; i++ { idx := (r.head + i) % r.cap if r.events[idx].ID == afterID { start = (idx + 1) % r.cap break } } if start == -1 { // afterID 不在缓冲区里(已被覆盖或从未存在), // 回放缓冲区里所有事件 —— 宁可重复也不丢失 start = 0 if !r.full { total = r.head } } else { // 从 start 开始到 head 结束 total = r.head - start if total < 0 { total += r.cap } } for i := 0; i < total; i++ { idx := (start + i) % r.cap evt := &r.events[idx] if evt.ID == "" { continue } fmt.Fprintf(res, "id: %s\nevent: %s\ndata: %s\n\n", evt.ID, evt.EventType, evt.Data) } flush.Flush() } // 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 // eventBuffer:per-user/agent 的事件环形缓冲区,供 Last-Event-ID 回放。 // Key 是 userName(人类)或 agentName(Agent),二者共享一个 map。 // 不是连接级别的 —— 同一用户断线重连后仍能从同一个缓冲区拿到断线期间的事件。 eventBuffer map[string]*eventRing bufMu sync.RWMutex seqCounter uint64 // 全局递增序列号,用作事件 ID seqMu sync.Mutex } // Default 是全局 SSE 管理器 var Default = &Manager{ clients: make(map[string]*Client), eventBuffer: make(map[string]*eventRing), } const eventBufferCap = 500 // 每用户最多保留 500 条事件 // nextEventID 生成下一个全局递增的事件 ID func (m *Manager) nextEventID() string { m.seqMu.Lock() defer m.seqMu.Unlock() m.seqCounter++ return fmt.Sprintf("%d", m.seqCounter) } // getOrCreateRing 获取或创建用户的环形缓冲区 func (m *Manager) getOrCreateRing(key string) *eventRing { if key == "" { return nil } m.bufMu.RLock() ring, ok := m.eventBuffer[key] m.bufMu.RUnlock() if ok { return ring } m.bufMu.Lock() defer m.bufMu.Unlock() // double-check if ring, ok = m.eventBuffer[key]; ok { return ring } ring = newEventRing(eventBufferCap) m.eventBuffer[key] = ring return ring } // AddClient 注册一个新 SSE 客户端(agentName 与 userName 二者恰其一) func (m *Manager) AddClient(res http.ResponseWriter, r *http.Request, 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") // Last-Event-ID 回放:EventSource 断线重连时自带这个头, // 服务端据此把断线期间的事件补上 —— 否则重连后永远看不到那段时间的邮件。 lastID := r.Header.Get("Last-Event-ID") key := m.bufferKey(userName, agentName) if ring := m.getOrCreateRing(key); ring != nil && lastID != "" { ring.replay(lastID, flusher, res) } m.mu.Lock() m.clients[id] = client m.mu.Unlock() // 发送连接确认(带 id 让客户端知道自己的 ID) evtID := m.nextEventID() client.SendWithID(evtID, "connected", map[string]string{"id": id}) // 启动心跳 go m.heartbeat(client) fmt.Printf("[SSE] Client connected: %s (agent=%q user=%q) lastID=%q\n", id, agentName, userName, lastID) return client } // bufferKey 返回缓冲区 key:优先 userName(人类),其次 agentName(Agent) func (m *Manager) bufferKey(userName, agentName string) string { if userName != "" { return "u:" + userName } if agentName != "" { return "a:" + agentName } return "" } // 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 } // 写入缓冲区 evtID := m.nextEventID() raw, _ := json.Marshal(data) if ring := m.getOrCreateRing(m.bufferKey("", agentName)); ring != nil { ring.push(StoredEvent{ID: evtID, EventType: eventType, Data: raw, Timestamp: time.Now()}) } m.mu.RLock() defer m.mu.RUnlock() for _, c := range m.clients { if c.AgentName == agentName { c.SendWithID(evtID, eventType, data) } } } // SendToUser 向指定人类用户的所有前端连接推送事件 func (m *Manager) SendToUser(userName, eventType string, data interface{}) { if userName == "" { return } evtID := m.nextEventID() raw, _ := json.Marshal(data) if ring := m.getOrCreateRing(m.bufferKey(userName, "")); ring != nil { ring.push(StoredEvent{ID: evtID, EventType: eventType, Data: raw, Timestamp: time.Now()}) } m.mu.RLock() defer m.mu.RUnlock() for _, c := range m.clients { if c.UserName == userName { c.SendWithID(evtID, eventType, data) } } } // SendToRecipient 根据收件人名同时尝试 Agent 通道与人类用户通道 func (m *Manager) SendToRecipient(name, eventType string, data interface{}) { if name == "" { return } evtID := m.nextEventID() raw, _ := json.Marshal(data) // 同时写两个缓冲区(人类或 Agent,或两者都有) if ring := m.getOrCreateRing(m.bufferKey(name, "")); ring != nil { ring.push(StoredEvent{ID: evtID, EventType: eventType, Data: raw, Timestamp: time.Now()}) } if ring := m.getOrCreateRing(m.bufferKey("", name)); ring != nil { ring.push(StoredEvent{ID: evtID, EventType: eventType, Data: raw, Timestamp: time.Now()}) } m.mu.RLock() defer m.mu.RUnlock() for _, c := range m.clients { if c.AgentName == name || c.UserName == name { c.SendWithID(evtID, eventType, data) } } } // Broadcast 向所有客户端广播事件(心跳、系统通知等) func (m *Manager) Broadcast(eventType string, data interface{}) { evtID := m.nextEventID() raw, _ := json.Marshal(data) // 广播写入所有用户的缓冲区(确保任何用户重连都能回放) m.bufMu.RLock() for key, ring := range m.eventBuffer { ring.push(StoredEvent{ID: evtID, EventType: eventType, Data: raw, Timestamp: time.Now()}) _ = key // key 仅用于日志,此处不需 } m.bufMu.RUnlock() m.mu.RLock() defer m.mu.RUnlock() for _, c := range m.clients { c.SendWithID(evtID, eventType, data) } } // ClientCount 返回当前连接数 func (m *Manager) ClientCount() int { m.mu.RLock() defer m.mu.RUnlock() return len(m.clients) } // Send 向单个客户端发送事件(无 ID) func (c *Client) Send(eventType string, data interface{}) { defer func() { recover() }() jsonData, err := json.Marshal(data) if err != nil { return } fmt.Fprintf(c.Res, "event: %s\ndata: %s\n\n", eventType, jsonData) c.Flusher.Flush() } // SendWithID 向单个客户端发送带 ID 的事件 func (c *Client) SendWithID(id, eventType string, data interface{}) { defer func() { recover() }() jsonData, err := json.Marshal(data) if err != nil { return } fmt.Fprintf(c.Res, "id: %s\nevent: %s\ndata: %s\n\n", id, 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: defer func() { recover() }() fmt.Fprintf(client.Res, ": heartbeat\n\n") client.Flusher.Flush() } } }