diff --git a/cmd/web-ui/main.go b/cmd/web-ui/main.go index 2ecfdd1..c0ab8d4 100644 --- a/cmd/web-ui/main.go +++ b/cmd/web-ui/main.go @@ -10,11 +10,17 @@ import ( "net/url" "os" "strings" + + "github.com/gorilla/websocket" ) //go:embed static var staticFiles embed.FS +var wsUpgrader = websocket.Upgrader{ + CheckOrigin: func(r *http.Request) bool { return true }, +} + func main() { addr := envDefault("AGENTMON_UI_ADDR", ":8082") queryAPIBase := envDefault("AGENTMON_QUERY_BASE", "http://query-api") @@ -36,6 +42,44 @@ func main() { proxy.ServeHTTP(w, r) }) + // WebSocket proxy to query-api + mux.HandleFunc("/api/v1/ws", func(w http.ResponseWriter, r *http.Request) { + queryWSURL := "ws://" + queryURL.Host + "/v1/ws" + conn, _, err := websocket.DefaultDialer.Dial(queryWSURL, nil) + if err != nil { + http.Error(w, "failed to connect to upstream WebSocket", http.StatusBadGateway) + return + } + defer conn.Close() + + uiConn, err := wsUpgrader.Upgrade(w, r, nil) + if err != nil { + return + } + defer uiConn.Close() + + // Bidirectional copy + done := make(chan struct{}) + go func() { + for { + _, msg, err := conn.ReadMessage() + if err != nil { + break + } + uiConn.WriteMessage(websocket.TextMessage, msg) + } + close(done) + }() + for { + _, msg, err := uiConn.ReadMessage() + if err != nil { + break + } + conn.WriteMessage(websocket.TextMessage, msg) + } + <-done + }) + // Static files staticFS, _ := fs.Sub(staticFiles, "static") fileServer := http.FileServer(http.FS(staticFS)) @@ -48,7 +92,7 @@ func main() { // SPA catch-all: serve index.html for all other routes mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) { // Serve index.html for SPA routes - if r.URL.Path == "/" || strings.HasPrefix(r.URL.Path, "/sessions") || strings.HasPrefix(r.URL.Path, "/runs") || strings.HasPrefix(r.URL.Path, "/openclaw") || strings.HasPrefix(r.URL.Path, "/agents") { + if r.URL.Path == "/" || strings.HasPrefix(r.URL.Path, "/sessions") || strings.HasPrefix(r.URL.Path, "/runs") || strings.HasPrefix(r.URL.Path, "/infrastructure") || strings.HasPrefix(r.URL.Path, "/agents") { f, err := staticFiles.Open("static/index.html") if err != nil { http.Error(w, "index.html not found", http.StatusInternalServerError)