Make it working, expose MCP prepare script to UI

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

pkg/stdio/client.go

@@ -58,8 +58,7 @@ func (c *Client) CreateProcess(ctx context.Context, command string, args []strin
 
 	resp, err := http.Post(url, "application/json", bytes.NewReader(reqBody))
 	if err != nil {
-		body, _ := io.ReadAll(resp.Body)
-		return nil, fmt.Errorf("failed to start process: %w. body: %s", err, string(body))
+		return nil, fmt.Errorf("failed to start process: %w", err)
 	}
 	defer resp.Body.Close()
 
@@ -259,9 +258,10 @@ type websocketWriter struct {
 }
 
 func (w *websocketWriter) Write(p []byte) (n int, err error) {
-	err = w.conn.WriteMessage(websocket.TextMessage, p)
+	// Use BinaryMessage type for better compatibility
+	err = w.conn.WriteMessage(websocket.BinaryMessage, p)
 	if err != nil {
-		return 0, err
+		return 0, fmt.Errorf("failed to write WebSocket message: %w", err)
 	}
 	return len(p), nil
 }

pkg/stdio/client_suite_test.go

@@ -1,6 +1,7 @@
 package stdio
 
 import (
+	"os"
 	"testing"
 
 	. "github.com/onsi/ginkgo/v2"
@@ -11,3 +12,17 @@ func TestSTDIOTransport(t *testing.T) {
 	RegisterFailHandler(Fail)
 	RunSpecs(t, "STDIOTransport test suite")
 }
+
+var baseURL string
+
+func init() {
+	baseURL = os.Getenv("STDIO_SERVER_URL")
+	if baseURL == "" {
+		baseURL = "http://localhost:8080"
+	}
+}
+
+var _ = AfterSuite(func() {
+	client := NewClient(baseURL)
+	client.StopGroup("test-group")
+})

pkg/stdio/client_test.go

@@ -2,24 +2,21 @@ package stdio
 
 import (
 	"context"
-	"os"
 	"time"
 
+	mcp "github.com/metoro-io/mcp-golang"
+	"github.com/metoro-io/mcp-golang/transport/stdio"
+	"github.com/mudler/LocalAGI/pkg/xlog"
 	. "github.com/onsi/ginkgo/v2"
 	. "github.com/onsi/gomega"
 )
 
 var _ = Describe("Client", func() {
 	var (
-		client  *Client
-		baseURL string
+		client *Client
 	)
 
 	BeforeEach(func() {
-		baseURL = os.Getenv("STDIO_SERVER_URL")
-		if baseURL == "" {
-			baseURL = "http://localhost:8080"
-		}
 		client = NewClient(baseURL)
 	})
 
@@ -183,5 +180,56 @@ var _ = Describe("Client", func() {
 			err = client.StopProcess(process.ID)
 			Expect(err).NotTo(HaveOccurred())
 		})
+
+		It("MCP", func() {
+			ctx := context.Background()
+			process, err := client.CreateProcess(ctx,
+				"docker", []string{"run", "-i", "--rm", "-e", "GITHUB_PERSONAL_ACCESS_TOKEN", "ghcr.io/github/github-mcp-server"},
+				[]string{"GITHUB_PERSONAL_ACCESS_TOKEN=test"}, "test-group")
+			Expect(err).NotTo(HaveOccurred())
+			Expect(process).NotTo(BeNil())
+			Expect(process.ID).NotTo(BeEmpty())
+
+			defer client.StopProcess(process.ID)
+
+			// MCP client
+
+			read, writer, err := client.GetProcessIO(process.ID)
+			Expect(err).NotTo(HaveOccurred())
+			Expect(read).NotTo(BeNil())
+			Expect(writer).NotTo(BeNil())
+
+			transport := stdio.NewStdioServerTransportWithIO(read, writer)
+
+			// Create a new client
+			mcpClient := mcp.NewClient(transport)
+			// Initialize the client
+			response, e := mcpClient.Initialize(ctx)
+			Expect(e).NotTo(HaveOccurred())
+			Expect(response).NotTo(BeNil())
+
+			Expect(mcpClient.Ping(ctx)).To(Succeed())
+
+			xlog.Debug("Client initialized: %v", response.Instructions)
+
+			alltools := []mcp.ToolRetType{}
+			var cursor *string
+			for {
+				tools, err := mcpClient.ListTools(ctx, cursor)
+				Expect(err).NotTo(HaveOccurred())
+				Expect(tools).NotTo(BeNil())
+				Expect(tools.Tools).NotTo(BeEmpty())
+				alltools = append(alltools, tools.Tools...)
+
+				if tools.NextCursor == nil {
+					break // No more pages
+				}
+				cursor = tools.NextCursor
+			}
+
+			for _, tool := range alltools {
+				xlog.Debug("Tool: %v", tool)
+			}
+		})
 	})
 })

pkg/stdio/server.go

@@ -1,7 +1,6 @@
 package stdio
 
 import (
-	"bufio"
 	"context"
 	"encoding/json"
 	"fmt"
@@ -41,9 +40,7 @@ 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 },
-		},
+		upgrader:  websocket.Upgrader{},
 	}
 }
 
@@ -354,7 +351,7 @@ func (s *Server) handleWebSocket(w http.ResponseWriter, r *http.Request) {
 		}
 	}()
 
-	// Handle stdout and stderr using bufio.Scanner
+	// Handle stdout and stderr
 	go func() {
 		defer func() {
 			select {
@@ -365,46 +362,29 @@ func (s *Server) handleWebSocket(w http.ResponseWriter, r *http.Request) {
 			}
 		}()
 
-		// Create a scanner that reads from both stdout and stderr
-		scanner := bufio.NewScanner(io.MultiReader(process.Stdout, process.Stderr))
-		// Set a larger buffer size for JSON-RPC messages (10MB)
-		scanner.Buffer(make([]byte, 10*1024*1024), 10*1024*1024)
-		// Use a custom split function to handle JSON-RPC messages
-		scanner.Split(func(data []byte, atEOF bool) (advance int, token []byte, err error) {
-			if atEOF && len(data) == 0 {
-				return 0, nil, nil
-			}
+		// Create a buffer for reading
+		buf := make([]byte, 4096)
+		reader := io.MultiReader(process.Stdout, process.Stderr)
 
-			// Look for the end of a JSON-RPC message
-			for i := 0; i < len(data); i++ {
-				if data[i] == '\n' {
-					return i + 1, data[:i], nil
-				}
-			}
-
-			// If we're at EOF, return the remaining data
-			if atEOF {
-				return len(data), data, nil
-			}
-
-			// Request more data
-			return 0, nil, nil
-		})
-
-		for scanner.Scan() {
-			line := scanner.Text()
-			xlog.Debug("Sending message", "processID", id, "message", line)
-			if err := conn.WriteMessage(websocket.TextMessage, []byte(line)); err != nil {
-				if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseNormalClosure) {
-					xlog.Debug("WebSocket output write error", "processID", id, "error", err)
+		for {
+			n, err := reader.Read(buf)
+			if err != nil {
+				if err != io.EOF {
+					xlog.Debug("Read error", "processID", id, "error", err)
 				}
 				return
 			}
-			xlog.Debug("Message sent to client", "processID", id, "message", line)
-		}
 
-		if err := scanner.Err(); err != nil {
-			xlog.Debug("Scanner error", "processID", id, "error", err)
+			if n > 0 {
+				xlog.Debug("Sending message", "processID", id, "size", n)
+				if err := conn.WriteMessage(websocket.BinaryMessage, buf[:n]); err != nil {
+					if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseNormalClosure) {
+						xlog.Debug("WebSocket output write error", "processID", id, "error", err)
+					}
+					return
+				}
+				xlog.Debug("Message sent to client", "processID", id, "size", n)
+			}
 		}
 	}()