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 tuning constants. The ping/pong cadence keeps intermediary proxies
// from dropping idle connections.
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 is the outbound (server→client) frame. TurnID/Seq are optional and
// used by the streaming path so a client can correlate a burst of streamed
// events to one assistant turn and keep them in order.
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 is the client→server frame. Data carries the per-type payload.
type wsInbound struct {
	Type  string          `json:"type"`
	ReqID string          `json:"reqId,omitempty"`
	Data  json.RawMessage `json:"data,omitempty"`
}

// hub owns the connection set and relays bus events to clients, each frame
// gated by its type (see 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 resolves each turn's role (design/22 §4.5); it mirrors the
	// server's resolver so the WS path and the REST path never diverge.
	envResolver turnEnvResolver
	log         *slog.Logger
	upgrader    websocket.Upgrader
	origins     map[string]bool
	allowAll    bool
	cancel      context.CancelFunc
	streamSeq   atomic.Uint64
}

// newHub builds a hub bound to the server's dependencies.
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 subscribes the hub to the bus event stream. Publishing to ws.> arrives
// here and is fanned out to clients; the producer bridge patterns (iot/task/
// automation/reminder/subagent-progress) are also subscribed and normalized into
// ws.* event types. Nothing happens when no bus is wired.
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 cancels the bus subscription, interrupts every orchestrator session (so
// no background turn outlives the hub), and shuts every live connection down.
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 adds a connection to the set.
func (h *hub) register(c *wsConn) {
	h.mu.Lock()
	h.conns[c] = struct{}{}
	h.mu.Unlock()
}

// unregister removes a connection from the set.
func (h *hub) unregister(c *wsConn) {
	h.mu.Lock()
	if _, ok := h.conns[c]; ok {
		delete(h.conns, c)
	}
	h.mu.Unlock()
}

// connCount reports the number of live connections (observability/tests).
func (h *hub) connCount() int {
	h.mu.RLock()
	defer h.mu.RUnlock()
	return len(h.conns)
}

// broadcast is the single delivery entry for outbound frames, and the design/11
// §7.1 role gate itself: it looks the frame's type up (frameReachesAll) and fans
// a registered 「状态机信号」 out to every registered connection, while every other
// type — content, tool, persisted-message, graph, private-state, and any type not
// classified yet — reaches owner + admin only. Registration happens only after the
// JWT is verified (ws_hub_conn.go), so the connection set is the authenticated set.
//
// Gating here rather than at the call sites is deliberate: there is exactly one
// fan-out loop, so a frame cannot be sent to the wrong audience by a call site that
// forgot to pick the restricted entry. It never blocks on any single (possibly
// slow) client: a full send buffer drops the frame.
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 maps a bus subject to its outbound event type, dropping the ws. prefix.
func eventType(subject string) string {
	return strings.TrimPrefix(subject, wsPrefix)
}

// rawData converts a possibly-empty raw payload into a value safe to publish.
func rawData(d json.RawMessage) any {
	if len(d) == 0 {
		return nil
	}
	return d
}