Cyrene_v2/internal/services/gateway/ws_hub.go

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
	// documents 解析回合携带的文档引用（design/24 §6.1）；同样与 server 的
	// 解析器同源，使 WS 帧与 REST 的文档校验口径一致。为 nil 时带附件的回合
	// 被明确拒绝（见 resolveDocuments），绝不静默丢弃。
	documents documentResolver
	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,
		documents:   s.resolveChatDocuments,
		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
}