Add groups to mcpbox

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
Signed-off-by: mudler <mudler@localai.io>

pkg/stdio/server.go

@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"io"
 	"net/http"
+	"os"
 	"os/exec"
 	"sync"
 	"time"
@@ -16,6 +17,7 @@ import (
 // Process represents a running process with its stdio streams
 type Process struct {
 	ID        string
+	GroupID   string
 	Cmd       *exec.Cmd
 	Stdin     io.WriteCloser
 	Stdout    io.ReadCloser
@@ -26,6 +28,7 @@ type Process struct {
 // Server handles process management and stdio streaming
 type Server struct {
 	processes map[string]*Process
+	groups    map[string][]string // maps group ID to process IDs
 	mu        sync.RWMutex
 	upgrader  websocket.Upgrader
 }
@@ -34,6 +37,7 @@ type Server struct {
 func NewServer() *Server {
 	return &Server{
 		processes: make(map[string]*Process),
+		groups:    make(map[string][]string),
 		upgrader: websocket.Upgrader{
 			CheckOrigin: func(r *http.Request) bool { return true },
 		},
@@ -41,9 +45,13 @@ func NewServer() *Server {
 }
 
 // StartProcess starts a new process and returns its ID
-func (s *Server) StartProcess(ctx context.Context, command string, args []string) (string, error) {
+func (s *Server) StartProcess(ctx context.Context, command string, args []string, env []string, groupID string) (string, error) {
 	cmd := exec.CommandContext(ctx, command, args...)
 
+	if len(env) > 0 {
+		cmd.Env = append(os.Environ(), env...)
+	}
+
 	stdin, err := cmd.StdinPipe()
 	if err != nil {
 		return "", fmt.Errorf("failed to create stdin pipe: %w", err)
@@ -65,6 +73,7 @@ func (s *Server) StartProcess(ctx context.Context, command string, args []string
 
 	process := &Process{
 		ID:        fmt.Sprintf("%d", time.Now().UnixNano()),
+		GroupID:   groupID,
 		Cmd:       cmd,
 		Stdin:     stdin,
 		Stdout:    stdout,
@@ -74,6 +83,9 @@ func (s *Server) StartProcess(ctx context.Context, command string, args []string
 
 	s.mu.Lock()
 	s.processes[process.ID] = process
+	if groupID != "" {
+		s.groups[groupID] = append(s.groups[groupID], process.ID)
+	}
 	s.mu.Unlock()
 
 	return process.ID, nil
@@ -87,6 +99,21 @@ func (s *Server) StopProcess(id string) error {
 		s.mu.Unlock()
 		return fmt.Errorf("process not found: %s", id)
 	}
+
+	// Remove from group if it exists
+	if process.GroupID != "" {
+		groupProcesses := s.groups[process.GroupID]
+		for i, pid := range groupProcesses {
+			if pid == id {
+				s.groups[process.GroupID] = append(groupProcesses[:i], groupProcesses[i+1:]...)
+				break
+			}
+		}
+		if len(s.groups[process.GroupID]) == 0 {
+			delete(s.groups, process.GroupID)
+		}
+	}
+
 	delete(s.processes, id)
 	s.mu.Unlock()
 
@@ -97,6 +124,57 @@ func (s *Server) StopProcess(id string) error {
 	return nil
 }
 
+// StopGroup stops all processes in a group
+func (s *Server) StopGroup(groupID string) error {
+	s.mu.Lock()
+	processIDs, exists := s.groups[groupID]
+	if !exists {
+		s.mu.Unlock()
+		return fmt.Errorf("group not found: %s", groupID)
+	}
+	s.mu.Unlock()
+
+	for _, pid := range processIDs {
+		if err := s.StopProcess(pid); err != nil {
+			return fmt.Errorf("failed to stop process %s in group %s: %w", pid, groupID, err)
+		}
+	}
+
+	return nil
+}
+
+// GetGroupProcesses returns all processes in a group
+func (s *Server) GetGroupProcesses(groupID string) ([]*Process, error) {
+	s.mu.RLock()
+	processIDs, exists := s.groups[groupID]
+	if !exists {
+		s.mu.RUnlock()
+		return nil, fmt.Errorf("group not found: %s", groupID)
+	}
+
+	processes := make([]*Process, 0, len(processIDs))
+	for _, pid := range processIDs {
+		if process, exists := s.processes[pid]; exists {
+			processes = append(processes, process)
+		}
+	}
+	s.mu.RUnlock()
+
+	return processes, nil
+}
+
+// ListGroups returns all group IDs
+func (s *Server) ListGroups() []string {
+	s.mu.RLock()
+	defer s.mu.RUnlock()
+
+	groups := make([]string, 0, len(s.groups))
+	for groupID := range s.groups {
+		groups = append(groups, groupID)
+	}
+	return groups
+}
+
 // GetProcess returns a process by ID
 func (s *Server) GetProcess(id string) (*Process, error) {
 	s.mu.RLock()
@@ -128,6 +206,8 @@ func (s *Server) Start(addr string) error {
 	http.HandleFunc("/processes", s.handleProcesses)
 	http.HandleFunc("/processes/", s.handleProcess)
 	http.HandleFunc("/ws/", s.handleWebSocket)
+	http.HandleFunc("/groups", s.handleGroups)
+	http.HandleFunc("/groups/", s.handleGroup)
 
 	return http.ListenAndServe(addr, nil)
 }
@@ -141,13 +221,15 @@ func (s *Server) handleProcesses(w http.ResponseWriter, r *http.Request) {
 		var req struct {
 			Command string   `json:"command"`
 			Args    []string `json:"args"`
+			Env     []string `json:"env"`
+			GroupID string   `json:"group_id"`
 		}
 		if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
 			http.Error(w, err.Error(), http.StatusBadRequest)
 			return
 		}
 
-		id, err := s.StartProcess(r.Context(), req.Command, req.Args)
+		id, err := s.StartProcess(r.Context(), req.Command, req.Args, req.Env, req.GroupID)
 		if err != nil {
 			http.Error(w, err.Error(), http.StatusInternalServerError)
 			return
@@ -212,6 +294,7 @@ func (s *Server) handleWebSocket(w http.ResponseWriter, r *http.Request) {
 	go func() {
 		buf := make([]byte, 1024)
 		for {
+
 			n, err := process.Stdout.Read(buf)
 			if err != nil {
 				return
@@ -235,3 +318,36 @@ func (s *Server) handleWebSocket(w http.ResponseWriter, r *http.Request) {
 	// Wait for process to exit
 	process.Cmd.Wait()
 }
+
+// Add new handlers for group management
+func (s *Server) handleGroups(w http.ResponseWriter, r *http.Request) {
+	switch r.Method {
+	case http.MethodGet:
+		groups := s.ListGroups()
+		json.NewEncoder(w).Encode(groups)
+	default:
+		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
+	}
+}
+
+func (s *Server) handleGroup(w http.ResponseWriter, r *http.Request) {
+	groupID := r.URL.Path[len("/groups/"):]
+
+	switch r.Method {
+	case http.MethodGet:
+		processes, err := s.GetGroupProcesses(groupID)
+		if err != nil {
+			http.Error(w, err.Error(), http.StatusNotFound)
+			return
+		}
+		json.NewEncoder(w).Encode(processes)
+	case http.MethodDelete:
+		if err := s.StopGroup(groupID); err != nil {
+			http.Error(w, err.Error(), http.StatusNotFound)
+			return
+		}
+		w.WriteHeader(http.StatusNoContent)
+	default:
+		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
+	}
+}