Cyrene_v2/internal/services/gateway/ws_hub.go

package gateway

import (
	"context"
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"log/slog"
	"net/http"
	"strings"
	"sync"
	"sync/atomic"
	"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"
	wsStreamSubject = "ws.stream"
	wsVoiceInput    = "ws.voice_input"
	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"`
	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
	auditor     SubagentAuditor
	taskOutputs TaskOutputStore
	streaming   bool
	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,
		streaming:   s.streaming,
		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 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 and the producer bridge
// patterns. The ws.chat subject is a streaming trigger (run the loop and emit
// chunks); a ws.* subject becomes an outbound event of its design type; a
// bridged producer subject (device_update, task output/done, automation result,
// reminder, subagent progress) is normalized into the matching ws.* event type.
// Events with no WS-visible mapping are dropped so internal agent-to-agent
// traffic never leaks to clients.
func (h *hub) relay(ctx context.Context, e bus.Event) error {
	if e.Subject == wsChatSubject {
		go h.streamChat(ctx, e.Payload)
		return nil
	}
	if e.Subject == wsStreamSubject {
		// OnStream→bus→relay path: the loop already published a structured
		// model.StreamEvent; map it to the typed downstream event.
		h.broadcastStream(e.Payload)
		return nil
	}
	ev, ok := normalizeEvent(e)
	if !ok {
		return nil
	}
	h.broadcast(ev)
	return nil
}

// streamChat runs one agent turn for a chat trigger and broadcasts the turn
// terminus (stream_end) with a fresh turnId; on the non-streaming fallback path
// it also emits a single stream_chunk. When the loop is streaming, the per-event
// deltas arrive via the OnStream→bus→relay path as structured message.*/
// thinking.*/tool_call.* events. 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)
	}
	turnID := newTurnID()
	if h.runner == nil || strings.TrimSpace(req.Text) == "" {
		h.broadcast(wsEvent{Type: "stream_end", TurnID: turnID, Seq: int(h.streamSeq.Add(1)), 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", TurnID: turnID, Seq: int(h.streamSeq.Add(1)), Data: map[string]any{"stopReason": "error", "error": err.Error()}})
		return
	}
	// When the loop is streaming (OnStream wired), the structured deltas already
	// fan out as message.*/thinking.*/tool_call.* events through the
	// OnStream→bus→relay path; a second whole-text chunk would duplicate the
	// tail. The non-streaming path (or a stub runner) still emits one chunk so
	// the stream contract holds in tests.
	if !h.streaming {
		h.broadcast(wsEvent{Type: "stream_chunk", TurnID: turnID, Seq: int(h.streamSeq.Add(1)), Data: map[string]any{"delta": res.Text}})
	}
	h.broadcast(wsEvent{Type: "stream_end", TurnID: turnID, Seq: int(h.streamSeq.Add(1)), Data: map[string]any{"stopReason": res.StopReason, "steps": res.Steps}})
}

// broadcastStream maps a streamed model.StreamEvent to its downstream WS event
// type and fans it out to every client. It runs on the bus subscription worker;
// seq is assigned atomically so a burst stays ordered even if the loop emits
// events from more than one goroutine. TurnId is left empty here because the
// OnStream bus payload does not carry the turn identity (future work would
// thread a turn id through loop → streamChunk → payload).
func (h *hub) broadcastStream(payload any) {
	ev, ok := payload.(model.StreamEvent)
	if !ok {
		return
	}
	typ, data := streamEventData(ev)
	h.broadcast(wsEvent{Type: typ, Seq: int(h.streamSeq.Add(1)), Data: data})
}

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

// 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. It tolerates a nil *websocket.Conn
// so tests that build a fake connection can rely on a safe teardown.
func (c *wsConn) shutdown() {
	c.once.Do(func() {
		close(c.done)
		if c.conn != nil {
			_ = 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)
}

// newTurnID returns a fresh random hex identifier used to correlate a burst of
// streamed WS events to one assistant turn. It burns the crypto/rand source and
// falls back to a nanosecond timestamp if the source is unavailable, so a id is
// always produced (a blank turnId would break client correlation silently).
func newTurnID() string {
	var b [8]byte
	if _, err := rand.Read(b[:]); err == nil {
		return "turn-" + hex.EncodeToString(b[:])
	}
	return "turn-" + time.Now().UTC().Format("20060102150405.000000000")
}

// 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, same-host, and whitelisted WebSocket origins,
// mirroring the CORS allow list. A "*" allows all.
//
// Same-host relaxation: when cyrene serves its own embedded frontend (production
// mode), the page origin is exactly the gateway's Host (e.g. http://localhost:10942),
// which is not in the dev-phase ALLOWED_ORIGINS list. We must let that through so a
// user who changes PORT does not need to edit .env for the WS to connect. Only truly
// cross-origin clients (external IM / dashboards on another origin) hit the allow-list.
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
		}
		if allowed[origin] {
			return true
		}
		// Same-host: the Origin matches this request's Host (scheme://host:port == host[:port]).
		return sameOriginAsHost(origin, r)
	}
}

// sameOriginAsHost reports whether origin (a full URL origin string such as
// "http://localhost:10942") corresponds to this request's Host header. It compares
// the host[:port] portion with loopback equivalence: "localhost" and "127.0.0.1"
// (and the IPv6 form) are treated as the same host, so serving the embedded frontend
// on 127.0.0.1 while a client connects via localhost (or vice versa) still passes.
func sameOriginAsHost(origin string, r *http.Request) bool {
	if origin == "" || r == nil {
		return false
	}
	after := origin
	if i := strings.Index(after, "://"); i >= 0 {
		after = after[i+3:]
	}
	if i := strings.IndexByte(after, '/'); i >= 0 {
		after = after[:i]
	}
	return hostEquivalent(after, r.Host)
}

// hostEquivalent compares two host[:port] strings, treating empty-port variants,
// case, and loopback aliases as equal. It handles "localhost" vs "127.0.0.1" and
// bracketed IPv6 loopback "[::1]".
func hostEquivalent(a, b string) bool {
	na := normalizeHost(a)
	nb := normalizeHost(b)
	if na == nb {
		return true
	}
	// Loopback aliases: localhost, 127.0.0.1, [::1]
	loop := map[string]bool{"localhost": true, "127.0.0.1": true, "[::1]": true, "::1": true, "127.0.0.2": true}
	return loop[na] && loop[nb]
}

// normalizeHost strips the port (if present), lowercases, and trims brackets.
func normalizeHost(h string) string {
	h = strings.ToLower(strings.TrimSpace(h))
	if i := strings.LastIndex(h, ":"); i >= 0 && !strings.Contains(h[i:], "]") {
		h = h[:i]
	}
	return strings.TrimPrefix(strings.TrimSuffix(h, "]"), "[")
}