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 every client. 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 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, 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 fans an event out to every registered connection without blocking 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 } h.mu.RLock() defer h.mu.RUnlock() for c := range h.conns { 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 }