Cyrene_v2/backend/services/gateway/ws/hub.go

// Package ws provides WebSocket connection management.
package ws

import (
	"context"
	"encoding/json"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/logger"
	"net/http"
	"strings"
	"sync"

	"github.com/gorilla/websocket"

	"git.yeij.top/AskaEth/Cyrene_v2/services/ai-core/model"
)

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	CheckOrigin:     func(r *http.Request) bool { return true },
}

// Orchestrator is the interface for the AI dialogue pipeline.
type Orchestrator interface {
	ProcessInput(ctx context.Context, sessionID, userMessage string) (<-chan model.StreamToken, *model.Intent, error)
	// FinishRound records the assistant reply and triggers post-round work
	// (e.g. memory extraction) after the response stream completes.
	FinishRound(sessionID, userID, assistantReply string)
}

// Hub manages all active WebSocket connections.
type Hub struct {
	mu           sync.RWMutex
	clients      map[*Client]bool
	register     chan *Client
	unregister   chan *Client
	broadcast    chan []byte
	orchestrator Orchestrator
}

// Client represents a single WebSocket connection.
type Client struct {
	hub       *Hub
	conn      *websocket.Conn
	send      chan []byte
	UserID    string
	SessionID string
	closed    bool
	mu        sync.Mutex
}

// WSMessage represents a client-server message.
type WSMessage struct {
	Type      string `json:"type"`
	Content   string `json:"content,omitempty"`
	SessionID string `json:"session_id,omitempty"`
	Timestamp int64  `json:"timestamp,omitempty"`
	Error     string `json:"error,omitempty"`
	Devices   []any  `json:"devices,omitempty"`
}

// NewHub creates a WebSocket hub with an optional orchestrator.
func NewHub(orch Orchestrator) *Hub {
	h := &Hub{
		clients:      make(map[*Client]bool),
		register:     make(chan *Client),
		unregister:   make(chan *Client),
		broadcast:    make(chan []byte, 256),
		orchestrator: orch,
	}
	go h.run()
	return h
}

func (h *Hub) run() {
	for {
		select {
		case c := <-h.register:
			h.mu.Lock()
			h.clients[c] = true
			n := len(h.clients)
			h.mu.Unlock()
			logger.Infof("[ws] client connected (total=%d)", n)

		case c := <-h.unregister:
			h.mu.Lock()
			if _, ok := h.clients[c]; ok {
				delete(h.clients, c)
				close(c.send)
			}
			h.mu.Unlock()

		case msg := <-h.broadcast:
			h.mu.RLock()
			for c := range h.clients {
				select {
				case c.send <- msg:
				default:
				}
			}
			h.mu.RUnlock()
		}
	}
}

// HandleUpgrade upgrades an HTTP connection to WebSocket.
func (h *Hub) HandleUpgrade(w http.ResponseWriter, r *http.Request) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		logger.Infof("[ws] upgrade error: %v", err)
		return
	}

	uid := r.URL.Query().Get("user_id")
	if uid == "" { uid = "anonymous" }

	client := &Client{
		hub:    h,
		conn:   conn,
		send:   make(chan []byte, 128),
		UserID: uid,
	}
	h.register <- client
	go client.writePump()
	go client.readPump()
}

func (c *Client) readPump() {
	defer func() {
		c.hub.unregister <- c
		c.conn.Close()
	}()

	for {
		_, raw, err := c.conn.ReadMessage()
		if err != nil {
			break
		}
		var msg WSMessage
		if err := json.Unmarshal(raw, &msg); err != nil {
			continue
		}

		switch msg.Type {
		case "message":
			c.handleChatMessage(msg)
		default:
			c.sendJSON(WSMessage{Type: "response", Content: "unknown message type: " + msg.Type})
		}
	}
}

func (c *Client) handleChatMessage(msg WSMessage) {
	c.SessionID = msg.SessionID

	// If orchestrator not available, echo
	if c.hub.orchestrator == nil {
		c.sendJSON(WSMessage{Type: "response", Content: "[系统] AI 核心未配置，当前为 echo 模式: " + msg.Content, SessionID: msg.SessionID})
		return
	}

	// Call orchestrator
	tokenCh, intent, err := c.hub.orchestrator.ProcessInput(context.Background(), msg.SessionID, msg.Content)
	if err != nil {
		c.sendJSON(WSMessage{Type: "error", Error: err.Error()})
		return
	}

	// Send intent info
	if intent != nil {
		c.sendJSON(WSMessage{Type: "response", Content: "[Intent: " + intent.Type + "]", SessionID: msg.SessionID})
	}

	// Stream tokens
	var fullContent strings.Builder
	for token := range tokenCh {
		fullContent.WriteString(token.Content)
		if token.Done {
			c.sendJSON(WSMessage{Type: "stream_end", SessionID: msg.SessionID})
			break
		}
		c.sendJSON(WSMessage{Type: "stream_chunk", Content: token.Content, SessionID: msg.SessionID})
	}

	// Record the reply and trigger post-round work (memory extraction etc.)
	c.hub.orchestrator.FinishRound(msg.SessionID, c.UserID, fullContent.String())
}

func (c *Client) writePump() {
	defer c.conn.Close()
	for msg := range c.send {
		c.mu.Lock()
		err := c.conn.WriteMessage(websocket.TextMessage, msg)
		c.mu.Unlock()
		if err != nil {
			return
		}
	}
}

func (c *Client) sendJSON(v any) {
	data, _ := json.Marshal(v)
	select {
	case c.send <- data:
	default:
	}
}

// Broadcast sends a message to all connected clients.
func (h *Hub) Broadcast(msg any) error {
	data, err := json.Marshal(msg)
	if err != nil {
		return err
	}
	h.broadcast <- data
	return nil
}

// Count returns the number of connected clients.
func (h *Hub) Count() int {
	h.mu.RLock()
	defer h.mu.RUnlock()
	return len(h.clients)
}

// UpgradeHandler returns an http.HandlerFunc for WebSocket upgrades.
func (h *Hub) UpgradeHandler() http.HandlerFunc {
	return h.HandleUpgrade
}