package gateway import ( "context" "encoding/json" "log/slog" "strings" "sync" "sync/atomic" "time" "github.com/gorilla/websocket" "git.yeij.top/AskaEth/Cyrene_v2/internal/bus" ) // WebSocket 调参常量。ping/pong 节奏用于防止中间代理把空闲连接断开。 const ( wsEventsPattern = "ws.>" wsPrefix = "ws." wsChatSubject = "ws.chat" wsStreamSubject = "ws.stream" wsVoiceInput = "ws.voice_input" wsVoiceOutput = "voice_output" wsReadLimit = 4096 wsWriteWait = 10 * time.Second wsPongWait = 60 * time.Second wsPingPeriod = 54 * time.Second wsSendBuffer = 32 ) // wsEvent 是下行(服务端→客户端)帧。TurnID/Seq 可选,由流式路径使用, // 让客户端能把一串流式事件归到同一个助手回合并保持顺序。 type wsEvent struct { Type string `json:"type"` TurnID string `json:"turnId,omitempty"` Seq int `json:"seq,omitempty"` SessionID string `json:"sessionId,omitempty"` Data any `json:"data,omitempty"` } // wsInbound 是上行(客户端→服务端)帧。Data 承载按类型区分的载荷。 type wsInbound struct { Type string `json:"type"` ReqID string `json:"reqId,omitempty"` Data json.RawMessage `json:"data,omitempty"` } // hub 持有连接集合并把 bus 事件转发给客户端,每个帧按其类型过门 // (见 broadcast / ws_frame_scope.go)。 type hub struct { mu sync.RWMutex conns map[*wsConn]struct{} bus bus.Bus runner Chatter auditor SubagentAuditor taskOutputs TaskOutputStore voice VoiceService streaming bool orchestrator *Orchestrator history HistoryStore // envResolver 解析每个回合的角色(design/22 §4.5);它与 server 的解析器 // 同源,使 WS 路径与 REST 路径不会分叉。 envResolver turnEnvResolver log *slog.Logger upgrader websocket.Upgrader origins map[string]bool allowAll bool cancel context.CancelFunc streamSeq atomic.Uint64 } // newHub 构建一个绑定到 server 依赖的 hub。 func newHub(s *Server) *hub { origins, allowAll := s.allowedOrigins() return &hub{ conns: make(map[*wsConn]struct{}), bus: s.bus, runner: s.runner, auditor: s.auditor, taskOutputs: s.taskOutputs, voice: s.voice, streaming: s.streaming, history: s.history, envResolver: s.envResolver, log: s.log, origins: origins, allowAll: allowAll, upgrader: websocket.Upgrader{ ReadBufferSize: 1024, WriteBufferSize: 1024, CheckOrigin: originChecker(origins, allowAll), }, } } // start 让 hub 订阅 bus 事件流。发往 ws.> 的事件到达此处后广播给客户端; // 生产者桥接模式(iot/task/automation/reminder/subagent-progress)也一并订阅, // 并归一为 ws.* 事件类型。未接入 bus 时什么都不做。 func (h *hub) start() { if h.bus == nil { return } ctx, cancel := context.WithCancel(context.Background()) h.cancel = cancel patterns := append([]string{wsEventsPattern}, bridgePatterns...) for _, p := range patterns { if _, err := h.bus.Subscribe(ctx, p, h.relay); err != nil { h.log.Warn("bus subscribe", "pattern", p, "err", err) } } } // close 取消 bus 订阅、中断所有 orchestrator 会话(使后台回合不会比 hub 活得 // 更久),并关闭每条在线连接。 func (h *hub) close() { if h.cancel != nil { h.cancel() } if h.orchestrator != nil { h.orchestrator.InterruptAll() } h.mu.Lock() conns := make([]*wsConn, 0, len(h.conns)) for c := range h.conns { conns = append(conns, c) } h.conns = make(map[*wsConn]struct{}) h.mu.Unlock() for _, c := range conns { c.shutdown() } } // register 把一条连接加入集合。 func (h *hub) register(c *wsConn) { h.mu.Lock() h.conns[c] = struct{}{} h.mu.Unlock() } // unregister 把一条连接从集合中移除。 func (h *hub) unregister(c *wsConn) { h.mu.Lock() if _, ok := h.conns[c]; ok { delete(h.conns, c) } h.mu.Unlock() } // connCount 返回在线连接数(可观测性/测试)。 func (h *hub) connCount() int { h.mu.RLock() defer h.mu.RUnlock() return len(h.conns) } // broadcast 是下行帧的唯一投递入口,也是 design/11 §7.1 角色门本身:它按帧的 // 类型查表(frameReachesAll),把登记在册的「状态机信号」广播给每条已注册连接, // 而其余类型 —— 内容、工具、持久化消息、图谱、角色私有状态,以及任何尚未分类 // 的类型 —— 只送达 owner + admin。连接只在 JWT 校验通过后才注册 // (ws_hub_conn.go),因此连接集合就是已认证集合。 // // 把门放在这里而非调用点是刻意的:全包只有这一个 fan-out 循环,调用点即使忘了 // 改走受限入口,也不可能把帧投错受众。它也不会被任何单个(可能很慢的)客户端 // 阻塞:发送缓冲满即丢弃该帧。 func (h *hub) broadcast(ev wsEvent) { b, err := json.Marshal(ev) if err != nil { h.log.Warn("marshal broadcast", "err", err) return } adminOnly := !frameReachesAll(ev.Type) h.mu.RLock() defer h.mu.RUnlock() for c := range h.conns { if adminOnly && !ownerAdminClaims(c.claims) { continue } c.enqueue(b) } } // eventType 把 bus subject 映射为下行事件类型,去掉 ws. 前缀。 func eventType(subject string) string { return strings.TrimPrefix(subject, wsPrefix) } // rawData 把可能为空的原始载荷转换为可安全发布的值。 func rawData(d json.RawMessage) any { if len(d) == 0 { return nil } return d }