From 608d7aba856db4277e2b6aeb9728350fda006f92 Mon Sep 17 00:00:00 2001
From: mudler <mudler@localai.io>
Date: Wed, 23 Apr 2025 18:39:35 +0200
Subject: [PATCH] quality of life improvements

Signed-off-by: mudler <mudler@localai.io>
---
 Makefile            |  9 ++++-
 pkg/stdio/server.go | 84 +++++++++++++++++++++++++++++++++++++++------
 2 files changed, 82 insertions(+), 11 deletions(-)

diff --git a/Makefile b/Makefile
index 6edd743..e9eca2f 100644
--- a/Makefile
+++ b/Makefile
@@ -1,5 +1,6 @@
 GOCMD?=go
 IMAGE_NAME?=webui
+MCPBOX_IMAGE_NAME?=mcpbox
 ROOT_DIR:=$(shell dirname $(realpath $(lastword $(MAKEFILE_LIST))))
 
 prepare-tests:
@@ -23,10 +24,16 @@ build: webui/react-ui/dist
 
 .PHONY: run
 run: webui/react-ui/dist
-	$(GOCMD) run ./
+	LOCALAGI_MCPBOX_URL="http://localhost:9090" $(GOCMD) run ./
 
 build-image:
 	docker build -t $(IMAGE_NAME) -f Dockerfile.webui .
 
 image-push:
 	docker push $(IMAGE_NAME)
+
+build-mcpbox:
+	docker build -t $(MCPBOX_IMAGE_NAME) -f Dockerfile.mcpbox .
+
+run-mcpbox:
+	docker run -v /var/run/docker.sock:/var/run/docker.sock --privileged -p 9090:8080 -ti mcpbox
\ No newline at end of file
diff --git a/pkg/stdio/server.go b/pkg/stdio/server.go
index bafd198..7b3bf54 100644
--- a/pkg/stdio/server.go
+++ b/pkg/stdio/server.go
@@ -1,6 +1,7 @@
 package stdio
 
 import (
+	"bytes"
 	"context"
 	"encoding/json"
 	"fmt"
@@ -306,47 +307,110 @@ func (s *Server) handleWebSocket(w http.ResponseWriter, r *http.Request) {
 
 	log.Printf("WebSocket connection established for process: %s", id)
 
+	// Create a done channel to signal process completion
+	done := make(chan struct{})
+
+	// Create buffers to capture output
+	var stdoutBuf, stderrBuf bytes.Buffer
+	stdoutTee := io.TeeReader(process.Stdout, &stdoutBuf)
+	stderrTee := io.TeeReader(process.Stderr, &stderrBuf)
+
 	// Handle stdin
 	go func() {
+		defer func() {
+			select {
+			case <-done:
+				// Process already done, this is expected
+			default:
+				log.Printf("WebSocket stdin connection closed for process %s", id)
+			}
+		}()
+
 		for {
 			_, message, err := conn.ReadMessage()
 			if err != nil {
-				log.Printf("WebSocket stdin read error for process %s: %v", id, err)
+				if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseNormalClosure) {
+					log.Printf("WebSocket stdin unexpected error for process %s: %v", id, err)
+				}
+				return
+			}
+			if _, err := process.Stdin.Write(message); err != nil {
+				if err != io.EOF {
+					log.Printf("WebSocket stdin write error for process %s: %v", id, err)
+				}
 				return
 			}
-			process.Stdin.Write(message)
 		}
 	}()
 
 	// Handle stdout
 	go func() {
+		defer func() {
+			select {
+			case <-done:
+				// Process already done, this is expected
+			default:
+				log.Printf("WebSocket stdout connection closed for process %s", id)
+			}
+		}()
+
 		buf := make([]byte, 1024)
 		for {
-			n, err := process.Stdout.Read(buf)
+			n, err := stdoutTee.Read(buf)
 			if err != nil {
-				log.Printf("WebSocket stdout read error for process %s: %v", id, err)
+				if err != io.EOF {
+					log.Printf("WebSocket stdout read error for process %s: %v", id, err)
+				}
+				return
+			}
+			if err := conn.WriteMessage(websocket.TextMessage, buf[:n]); err != nil {
+				if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseNormalClosure) {
+					log.Printf("WebSocket stdout write error for process %s: %v", id, err)
+				}
 				return
 			}
-			conn.WriteMessage(websocket.TextMessage, buf[:n])
 		}
 	}()
 
 	// Handle stderr
 	go func() {
+		defer func() {
+			select {
+			case <-done:
+				// Process already done, this is expected
+			default:
+				log.Printf("WebSocket stderr connection closed for process %s", id)
+			}
+		}()
+
 		buf := make([]byte, 1024)
 		for {
-			n, err := process.Stderr.Read(buf)
+			n, err := stderrTee.Read(buf)
 			if err != nil {
-				log.Printf("WebSocket stderr read error for process %s: %v", id, err)
+				if err != io.EOF {
+					log.Printf("WebSocket stderr read error for process %s: %v", id, err)
+				}
+				return
+			}
+			if err := conn.WriteMessage(websocket.TextMessage, buf[:n]); err != nil {
+				if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseNormalClosure) {
+					log.Printf("WebSocket stderr write error for process %s: %v", id, err)
+				}
 				return
 			}
-			conn.WriteMessage(websocket.TextMessage, buf[:n])
 		}
 	}()
 
 	// Wait for process to exit
-	process.Cmd.Wait()
-	log.Printf("Process %s exited", id)
+	err = process.Cmd.Wait()
+	close(done) // Signal that the process is done
+
+	if err != nil {
+		log.Printf("Process %s exited with error: %v\nstdout: %s\nstderr: %s",
+			id, err, stdoutBuf.String(), stderrBuf.String())
+	} else {
+		log.Printf("Process %s exited successfully", id)
+	}
 }
 
 // Add new handlers for group management