visionA/local-agent/server/internal/api/ws/hub.go

package ws

import (
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"sync"
	"time"

	"github.com/gorilla/websocket"
)

type Client struct {
	Conn *websocket.Conn
	Send chan []byte
}

type Subscription struct {
	Client *Client
	Room   string
	done   chan struct{} // used by RegisterSync to wait for completion
}

type RoomMessage struct {
	Room    string
	Message []byte
}

// Hub 管理 WebSocket client 訂閱與訊息廣播。
//
// M8-4b：Hub 額外負責「第一個 client 連上時寫 sentinel file」，
// 讓 Wails 端的 StartupPipeline 知道階段 6（Wait for Web UI WebSocket）已完成。
// 詳細設計見 .autoflow/04-architecture/v2/startup-pipeline.md §3。
//
// dataDir 由 main.go 在初始化 Hub 後透過 SetStartupSentinel(dataDir) 注入。
// 若 dataDir 為空，sentinel 寫入會被跳過（單元測試或缺少資料目錄時的安全行為）。
type Hub struct {
	rooms      map[string]map[*Client]bool
	register   chan *Subscription
	unregister chan *Subscription
	broadcast  chan *RoomMessage
	mu         sync.RWMutex

	// M8-4b: 啟動 sentinel file
	sentinelDataDir string    // <dataDir>，由 SetStartupSentinel 設定
	sentinelOnce    sync.Once // 確保只在「第一個」client 連上時寫一次
	bootID          string    // 寫入 sentinel 內容供 debug
}

func NewHub() *Hub {
	return &Hub{
		rooms:      make(map[string]map[*Client]bool),
		register:   make(chan *Subscription, 10),
		unregister: make(chan *Subscription, 10),
		broadcast:  make(chan *RoomMessage, 100),
		bootID:     fmt.Sprintf("boot-%d", time.Now().UnixNano()),
	}
}

// SetStartupSentinel 設定 sentinel file 的根目錄。
// main.go 在 NewHub() 之後、Run() 之前呼叫一次，dataDir 應為完整路徑。
//
// 寫入路徑：<dataDir>/.first-ws-connected
// 內容：boot-id + timestamp（用於 debug，內容對 Wails 端的判斷沒有意義，存在即可）
//
// dataDir 為空字串時 sentinel 機制完全停用。
func (h *Hub) SetStartupSentinel(dataDir string) {
	h.mu.Lock()
	h.sentinelDataDir = dataDir
	h.mu.Unlock()
}

// writeStartupSentinel 在第一個 WebSocket client 連上時呼叫一次。
// 由 sentinelOnce 確保只執行一次；後續連線完全 no-op。
//
// 寫入失敗不會 panic 也不會回 error：sentinel 是 best-effort 機制，
// 若 disk 滿/權限錯，Wails 端會走 hard timeout 路徑進 Error state。
func (h *Hub) writeStartupSentinel() {
	h.sentinelOnce.Do(func() {
		h.mu.RLock()
		dir := h.sentinelDataDir
		bootID := h.bootID
		h.mu.RUnlock()
		if dir == "" {
			return
		}
		path := filepath.Join(dir, ".first-ws-connected")
		// 確保父目錄存在（dataDir 通常已存在，但保險起見）
		_ = os.MkdirAll(dir, 0o755)
		f, err := os.Create(path)
		if err != nil {
			return
		}
		_, _ = fmt.Fprintf(f, "bootId=%s\nts=%d\n", bootID, time.Now().UnixMilli())
		_ = f.Close()
	})
}

func (h *Hub) Run() {
	for {
		select {
		case sub := <-h.register:
			h.mu.Lock()
			if h.rooms[sub.Room] == nil {
				h.rooms[sub.Room] = make(map[*Client]bool)
			}
			h.rooms[sub.Room][sub.Client] = true
			h.mu.Unlock()
			// M8-4b：第一次有 client 加入任何 room → 寫 sentinel file
			// （sync.Once 保證後續呼叫 no-op）
			h.writeStartupSentinel()
			if sub.done != nil {
				close(sub.done)
			}

		case sub := <-h.unregister:
			h.mu.Lock()
			if clients, ok := h.rooms[sub.Room]; ok {
				if _, exists := clients[sub.Client]; exists {
					delete(clients, sub.Client)
					close(sub.Client.Send)
				}
			}
			h.mu.Unlock()

		case msg := <-h.broadcast:
			h.mu.RLock()
			if clients, ok := h.rooms[msg.Room]; ok {
				for client := range clients {
					select {
					case client.Send <- msg.Message:
					default:
						close(client.Send)
						delete(clients, client)
					}
				}
			}
			h.mu.RUnlock()
		}
	}
}

func (h *Hub) Register(sub *Subscription) {
	h.register <- sub
}

// RegisterSync registers a subscription and blocks until the Hub has processed it,
// ensuring the client is in the room before returning.
func (h *Hub) RegisterSync(sub *Subscription) {
	sub.done = make(chan struct{})
	h.register <- sub
	<-sub.done
}

func (h *Hub) Unregister(sub *Subscription) {
	h.unregister <- sub
}

func (h *Hub) BroadcastToRoom(room string, data interface{}) {
	jsonData, err := json.Marshal(data)
	if err != nil {
		return
	}
	h.broadcast <- &RoomMessage{Room: room, Message: jsonData}
}