Cyrene_v2/internal/services/gateway/ws_hub.go

package gateway

import (
	"context"
	"encoding/json"
	"log/slog"
	"net/http"
	"strings"
	"sync"
	"time"

	"github.com/gorilla/websocket"

	"git.yeij.top/AskaEth/Cyrene_v2/internal/auth"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/bus"
	apperr "git.yeij.top/AskaEth/Cyrene_v2/internal/errors"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/model"
)

// WebSocket tuning constants. The ping/pong cadence keeps intermediary proxies
// from dropping idle connections.
const (
	wsEventsPattern = "ws.>"
	wsPrefix        = "ws."
	wsChatSubject   = "ws.chat"
	wsReadLimit     = 4096
	wsWriteWait     = 10 * time.Second
	wsPongWait      = 60 * time.Second
	wsPingPeriod    = 54 * time.Second
	wsSendBuffer    = 32
)

// wsEvent is the outbound (server→client) frame.
type wsEvent struct {
	Type string `json:"type"`
	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"`
}

// wsConn wraps one WebSocket connection. Writes are serialized through the
// buffered send channel drained by a single writePump, so the gorilla conn is
// only ever touched by one goroutine at a time.
type wsConn struct {
	conn   *websocket.Conn
	send   chan []byte
	done   chan struct{}
	h      *hub
	claims *auth.Claims
	once   sync.Once
}

// 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
	log      *slog.Logger
	upgrader websocket.Upgrader
	origins  map[string]bool
	allowAll bool
	cancel   context.CancelFunc
}

// 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,
		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; 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
	if _, err := h.bus.Subscribe(ctx, wsEventsPattern, h.relay); err != nil {
		h.log.Warn("bus subscribe", "pattern", wsEventsPattern, "err", err)
	}
}

// close cancels the bus subscription and shuts every live connection down.
func (h *hub) close() {
	if h.cancel != nil {
		h.cancel()
	}
	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)
	}
}

// relay is the bus handler for the ws.> pattern. The ws.chat subject is a
// streaming trigger (run the loop and emit chunks); every other subject becomes
// an outbound event of its design type.
func (h *hub) relay(ctx context.Context, e bus.Event) error {
	if e.Subject == wsChatSubject {
		go h.streamChat(ctx, e.Payload)
		return nil
	}
	h.broadcast(wsEvent{Type: eventType(e.Subject), Data: e.Payload})
	return nil
}

// streamChat runs one agent turn for a chat trigger and broadcasts the result as
// stream_chunk/stream_end. It runs on its own goroutine so a slow model call does
// not stall the bus subscription worker.
func (h *hub) streamChat(ctx context.Context, payload any) {
	var req chatStreamRequest
	if b, err := json.Marshal(payload); err == nil {
		_ = json.Unmarshal(b, &req)
	}
	if h.runner == nil || strings.TrimSpace(req.Text) == "" {
		h.broadcast(wsEvent{Type: "stream_end", Data: map[string]any{"stopReason": "idle"}})
		return
	}
	if req.SessionID == "" {
		req.SessionID = "web"
	}
	res, err := h.runner.Run(ctx, req.SessionID, []model.Message{model.TextMessage(model.RoleUser, req.Text)})
	if err != nil {
		h.log.Error("ws chat run", "err", err, "session", req.SessionID)
		h.broadcast(wsEvent{Type: "stream_end", Data: map[string]any{"stopReason": "error", "error": err.Error()}})
		return
	}
	h.broadcast(wsEvent{Type: "stream_chunk", Data: map[string]any{"delta": res.Text}})
	h.broadcast(wsEvent{Type: "stream_end", Data: map[string]any{"stopReason": res.StopReason, "steps": res.Steps}})
}

// chatStreamRequest is the decoded ws.chat payload.
type chatStreamRequest struct {
	SessionID string `json:"sessionId"`
	Text      string `json:"text"`
}

// handleWS upgrades /ws/chat after authenticating the JWT from the query token or
// the Authorization header. A failed upgrade writes a 401 {error:...} before the
// handshake, so a browser/IM client sees a clean reject.
func (s *Server) handleWS(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodGet {
		writeError(w, apperr.Wrap(apperr.ErrBadRequest, "bad_request", "websocket requires GET"))
		return
	}
	claims, err := s.auth.Verify(wsToken(r))
	if err != nil {
		s.log.Debug("ws auth", "err", err)
		writeError(w, apperr.ErrUnauthorized)
		return
	}
	conn, err := s.hub.upgrader.Upgrade(w, r, nil)
	if err != nil {
		s.log.Debug("ws upgrade", "err", err)
		return
	}
	c := &wsConn{
		conn:   conn,
		send:   make(chan []byte, wsSendBuffer),
		done:   make(chan struct{}),
		h:      s.hub,
		claims: claims,
	}
	s.hub.register(c)
	go c.writePump()
	go c.readPump()
}

// readPump reads client frames until the connection closes, routing each by type.
func (c *wsConn) readPump() {
	defer func() {
		c.h.unregister(c)
		c.shutdown()
	}()
	c.conn.SetReadLimit(wsReadLimit)
	_ = c.conn.SetReadDeadline(time.Now().Add(wsPongWait))
	c.conn.SetPongHandler(func(string) error {
		return c.conn.SetReadDeadline(time.Now().Add(wsPongWait))
	})
	for {
		_, msg, err := c.conn.ReadMessage()
		if err != nil {
			if websocket.IsUnexpectedCloseError(err, websocket.CloseNormalClosure, websocket.CloseGoingAway) {
				c.h.log.Debug("ws read", "err", err)
			}
			return
		}
		var in wsInbound
		if err := json.Unmarshal(msg, &in); err != nil {
			c.enqueueJSON(wsEvent{Type: "error", Data: map[string]any{"message": "invalid json"}})
			continue
		}
		c.route(in)
	}
}

// writePump drains the send channel (and sends keepalive pings) until the
// connection closes. It is the only writer on the gorilla conn.
func (c *wsConn) writePump() {
	ticker := time.NewTicker(wsPingPeriod)
	defer func() {
		ticker.Stop()
		_ = c.conn.Close()
	}()
	for {
		select {
		case <-c.done:
			return
		case msg := <-c.send:
			_ = c.conn.SetWriteDeadline(time.Now().Add(wsWriteWait))
			if err := c.conn.WriteMessage(websocket.TextMessage, msg); err != nil {
				return
			}
		case <-ticker.C:
			_ = c.conn.SetWriteDeadline(time.Now().Add(wsWriteWait))
			if err := c.conn.WriteMessage(websocket.PingMessage, nil); err != nil {
				return
			}
		}
	}
}

// route dispatches one inbound frame by its type.
func (c *wsConn) route(in wsInbound) {
	switch in.Type {
	case "mode_switch":
		c.publishPublic("ws.mode_switch", rawData(in.Data))
	case "chat":
		c.publishPublic(wsChatSubject, rawData(in.Data))
	case "subagent_ask":
		c.enqueueJSON(wsEvent{Type: "subagent_event", Data: map[string]any{"state": "pending"}})
	case "task_output_query":
		c.enqueueJSON(wsEvent{Type: "task_output", Data: map[string]any{}})
	default:
		c.enqueueJSON(wsEvent{Type: "error", Data: map[string]any{"message": "unknown event type"}})
	}
}

// publishPublic publishes a bus event without any per-connection reply.
func (c *wsConn) publishPublic(subject string, payload any) {
	if c.h.bus == nil {
		return
	}
	_ = c.h.bus.Publish(context.Background(), subject, payload)
}

// shutdown closes the connection exactly once.
func (c *wsConn) shutdown() {
	c.once.Do(func() {
		close(c.done)
		_ = c.conn.Close()
	})
}

// enqueue queues a serialized frame, dropping it when a slow client's buffer is
// full so the producer never blocks.
func (c *wsConn) enqueue(b []byte) {
	select {
	case c.send <- b:
	default:
	}
}

// enqueueJSON serializes ev and queues it.
func (c *wsConn) enqueueJSON(ev wsEvent) {
	b, err := json.Marshal(ev)
	if err != nil {
		c.h.log.Warn("marshal ws event", "err", err)
		return
	}
	c.enqueue(b)
}

// publishChatStream publishes a ws.chat streaming trigger (REST stream mode).
func (s *Server) publishChatStream(ctx context.Context, sessionID, text string) {
	if s.bus == nil {
		return
	}
	_ = s.bus.Publish(ctx, wsChatSubject, map[string]any{"sessionId": sessionID, "text": text})
}

// wsToken extracts the JWT from the ?token= query or the Authorization header.
func wsToken(r *http.Request) string {
	if t := r.URL.Query().Get("token"); t != "" {
		return t
	}
	const prefix = "Bearer "
	h := r.Header.Get("Authorization")
	if strings.HasPrefix(h, prefix) {
		return strings.TrimSpace(h[len(prefix):])
	}
	return ""
}

// 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
}

// originChecker allows same-origin (empty) and whitelisted WebSocket origins,
// mirroring the CORS allow list. A "*" default allows all.
func originChecker(allowed map[string]bool, allowAll bool) func(*http.Request) bool {
	return func(r *http.Request) bool {
		origin := r.Header.Get("Origin")
		if origin == "" {
			return true
		}
		if allowAll {
			return true
		}
		return allowed[origin]
	}
}