feat: local MCP server support (#61)

* wip Signed-off-by: Ettore Di Giacinto <mudler@localai.io> Signed-off-by: mudler <mudler@localai.io> * Add groups to mcpbox Signed-off-by: Ettore Di Giacinto <mudler@localai.io> Signed-off-by: mudler <mudler@localai.io> * Add mcpbox dockerfile and entrypoint Signed-off-by: mudler <mudler@localai.io> * Attach mcp stdio box to agent Signed-off-by: Ettore Di Giacinto <mudler@localai.io> Signed-off-by: mudler <mudler@localai.io> * Add to dockerfile Signed-off-by: mudler <mudler@localai.io> * Attach to config Signed-off-by: mudler <mudler@localai.io> * Attach to ui Signed-off-by: mudler <mudler@localai.io> * Revert "Attach to ui" This reverts commit 088d0c47e87ee8f84297e47d178fb7384bbe6d45. Signed-off-by: mudler <mudler@localai.io> * add one-time process, attach to UI the mcp server json configuration Signed-off-by: mudler <mudler@localai.io> * quality of life improvements Signed-off-by: mudler <mudler@localai.io> * fixes Signed-off-by: mudler <mudler@localai.io> * Make it working, expose MCP prepare script to UI Signed-off-by: mudler <mudler@localai.io> * Add container image to CI builds * Wire mcpbox to tests * Improve setup' * Not needed anymore, using tests Signed-off-by: mudler <mudler@localai.io> * fix: do not override actions Signed-off-by: mudler <mudler@localai.io> * chore(tests): fix env var Signed-off-by: mudler <mudler@localai.io> --------- Signed-off-by: Ettore Di Giacinto <mudler@localai.io> Signed-off-by: mudler <mudler@localai.io>
2025-04-24 16:39:20 +02:00
parent ce997d2425
commit eb8663ada1
15 changed files with 1575 additions and 65 deletions
--- a/core/agent/agent.go
+++ b/core/agent/agent.go
@@ -241,6 +241,7 @@ func (a *Agent) Stop() {
 	a.Lock()
 	defer a.Unlock()
 	xlog.Debug("Stopping agent", "agent", a.Character.Name)
+	a.closeMCPSTDIOServers()
 	a.context.Cancel()
 }

--- a/core/agent/mcp.go
+++ b/core/agent/mcp.go
@@ -3,12 +3,14 @@ package agent
 import (
 	"context"
 	"encoding/json"
-	"errors"

 	mcp "github.com/metoro-io/mcp-golang"
 	"github.com/metoro-io/mcp-golang/transport/http"
+	stdioTransport "github.com/metoro-io/mcp-golang/transport/stdio"
 	"github.com/mudler/LocalAGI/core/types"
+	"github.com/mudler/LocalAGI/pkg/stdio"
 	"github.com/mudler/LocalAGI/pkg/xlog"
+
 	"github.com/sashabaranov/go-openai/jsonschema"
 )

@@ -19,6 +21,12 @@ type MCPServer struct {
 	Token string `json:"token"`
 }

+type MCPSTDIOServer struct {
+	Args []string `json:"args"`
+	Env  []string `json:"env"`
+	Cmd  string   `json:"cmd"`
+}
+
 type mcpAction struct {
 	mcpClient       *mcp.Client
 	inputSchema     ToolInputSchema
@@ -79,6 +87,68 @@ type ToolInputSchema struct {
 	Required   []string               `json:"required,omitempty"`
 }

+func (a *Agent) addTools(client *mcp.Client) (types.Actions, error) {
+
+	var generatedActions types.Actions
+	xlog.Debug("Initializing client")
+	// Initialize the client
+	response, e := client.Initialize(a.context)
+	if e != nil {
+		xlog.Error("Failed to initialize client", "error", e.Error())
+		return nil, e
+	}
+
+	xlog.Debug("Client initialized: %v", response.Instructions)
+
+	var cursor *string
+	for {
+		tools, err := client.ListTools(a.context, cursor)
+		if err != nil {
+			xlog.Error("Failed to list tools", "error", err.Error())
+			return nil, err
+		}
+
+		for _, t := range tools.Tools {
+			desc := ""
+			if t.Description != nil {
+				desc = *t.Description
+			}
+
+			xlog.Debug("Tool", "name", t.Name, "description", desc)
+
+			dat, err := json.Marshal(t.InputSchema)
+			if err != nil {
+				xlog.Error("Failed to marshal input schema", "error", err.Error())
+			}
+
+			xlog.Debug("Input schema", "tool", t.Name, "schema", string(dat))
+
+			// XXX: This is a wild guess, to verify (data types might be incompatible)
+			var inputSchema ToolInputSchema
+			err = json.Unmarshal(dat, &inputSchema)
+			if err != nil {
+				xlog.Error("Failed to unmarshal input schema", "error", err.Error())
+			}
+
+			// Create a new action with Client + tool
+			generatedActions = append(generatedActions, &mcpAction{
+				mcpClient:       client,
+				toolName:        t.Name,
+				inputSchema:     inputSchema,
+				toolDescription: desc,
+			})
+		}
+
+		if tools.NextCursor == nil {
+			break // No more pages
+		}
+		cursor = tools.NextCursor
+	}
+
+	return generatedActions, nil
+
+}
+
 func (a *Agent) initMCPActions() error {

 	a.mcpActions = nil
@@ -86,6 +156,7 @@ func (a *Agent) initMCPActions() error {

 	generatedActions := types.Actions{}

+	// MCP HTTP Servers
 	for _, mcpServer := range a.options.mcpServers {
 		transport := http.NewHTTPClientTransport("/mcp")
 		transport.WithBaseURL(mcpServer.URL)
@@ -95,70 +166,60 @@ func (a *Agent) initMCPActions() error {

 		// Create a new client
 		client := mcp.NewClient(transport)
+		xlog.Debug("Adding tools for MCP server", "server", mcpServer)
+		actions, err := a.addTools(client)
+		if err != nil {
+			xlog.Error("Failed to add tools for MCP server", "server", mcpServer, "error", err.Error())
+		}
+		generatedActions = append(generatedActions, actions...)
+	}

-		xlog.Debug("Initializing client", "server", mcpServer.URL)
-		// Initialize the client
-		response, e := client.Initialize(a.context)
-		if e != nil {
-			xlog.Error("Failed to initialize client", "error", e.Error(), "server", mcpServer)
-			if err == nil {
-				err = e
-			} else {
-				err = errors.Join(err, e)
-			}
+	// MCP STDIO Servers
+
+	a.closeMCPSTDIOServers() // Make sure we stop all previous servers if any is active
+
+	if a.options.mcpPrepareScript != "" {
+		xlog.Debug("Preparing MCP box", "script", a.options.mcpPrepareScript)
+		client := stdio.NewClient(a.options.mcpBoxURL)
+		client.RunProcess(a.context, "/bin/bash", []string{"-c", a.options.mcpPrepareScript}, []string{})
+	}
+
+	for _, mcpStdioServer := range a.options.mcpStdioServers {
+		client := stdio.NewClient(a.options.mcpBoxURL)
+		p, err := client.CreateProcess(a.context,
+			mcpStdioServer.Cmd,
+			mcpStdioServer.Args,
+			mcpStdioServer.Env,
+			a.Character.Name)
+		if err != nil {
+			xlog.Error("Failed to create process", "error", err.Error())
+			continue
+		}
+		read, writer, err := client.GetProcessIO(p.ID)
+		if err != nil {
+			xlog.Error("Failed to get process IO", "error", err.Error())
 			continue
 		}

-		xlog.Debug("Client initialized: %v", response.Instructions)
+		transport := stdioTransport.NewStdioServerTransportWithIO(read, writer)

-		var cursor *string
-		for {
-			tools, err := client.ListTools(a.context, cursor)
-			if err != nil {
-				xlog.Error("Failed to list tools", "error", err.Error())
-				return err
-			}
+		// Create a new client
+		mcpClient := mcp.NewClient(transport)

-			for _, t := range tools.Tools {
-				desc := ""
-				if t.Description != nil {
-					desc = *t.Description
-				}
-
-				xlog.Debug("Tool", "mcpServer", mcpServer, "name", t.Name, "description", desc)
-
-				dat, err := json.Marshal(t.InputSchema)
-				if err != nil {
-					xlog.Error("Failed to marshal input schema", "error", err.Error())
-				}
-
-				xlog.Debug("Input schema", "mcpServer", mcpServer, "tool", t.Name, "schema", string(dat))
-
-				// XXX: This is a wild guess, to verify (data types might be incompatible)
-				var inputSchema ToolInputSchema
-				err = json.Unmarshal(dat, &inputSchema)
-				if err != nil {
-					xlog.Error("Failed to unmarshal input schema", "error", err.Error())
-				}
-
-				// Create a new action with Client + tool
-				generatedActions = append(generatedActions, &mcpAction{
-					mcpClient:       client,
-					toolName:        t.Name,
-					inputSchema:     inputSchema,
-					toolDescription: desc,
-				})
-			}
-
-			if tools.NextCursor == nil {
-				break // No more pages
-			}
-			cursor = tools.NextCursor
+		xlog.Debug("Adding tools for MCP server (stdio)", "server", mcpStdioServer)
+		actions, err := a.addTools(mcpClient)
+		if err != nil {
+			xlog.Error("Failed to add tools for MCP server", "server", mcpStdioServer, "error", err.Error())
 		}
-
+		generatedActions = append(generatedActions, actions...)
 	}

 	a.mcpActions = generatedActions

 	return err
 }
+
+func (a *Agent) closeMCPSTDIOServers() {
+	client := stdio.NewClient(a.options.mcpBoxURL)
+	client.StopGroup(a.Character.Name)
+}
--- a/core/agent/options.go
+++ b/core/agent/options.go
@@ -50,8 +50,10 @@ type options struct {

 	conversationsPath string

-	mcpServers []MCPServer
-
+	mcpServers                  []MCPServer
+	mcpStdioServers             []MCPSTDIOServer
+	mcpBoxURL                   string
+	mcpPrepareScript            string
 	newConversationsSubscribers []func(openai.ChatCompletionMessage)

 	observer     Observer
@@ -207,6 +209,27 @@ func WithMCPServers(servers ...MCPServer) Option {
 	}
 }

+func WithMCPSTDIOServers(servers ...MCPSTDIOServer) Option {
+	return func(o *options) error {
+		o.mcpStdioServers = servers
+		return nil
+	}
+}
+
+func WithMCPBoxURL(url string) Option {
+	return func(o *options) error {
+		o.mcpBoxURL = url
+		return nil
+	}
+}
+
+func WithMCPPrepareScript(script string) Option {
+	return func(o *options) error {
+		o.mcpPrepareScript = script
+		return nil
+	}
+}
+
 func WithLLMAPIURL(url string) Option {
 	return func(o *options) error {
 		o.LLMAPI.APIURL = url
--- a/core/state/config.go
+++ b/core/state/config.go
@@ -2,6 +2,8 @@ package state

 import (
 	"encoding/json"
+	"fmt"
+	"strings"

 	"github.com/mudler/LocalAGI/core/agent"
 	"github.com/mudler/LocalAGI/core/types"
@@ -30,10 +32,13 @@ func (d DynamicPromptsConfig) ToMap() map[string]string {
 }

 type AgentConfig struct {
-	Connector      []ConnectorConfig      `json:"connectors" form:"connectors" `
-	Actions        []ActionsConfig        `json:"actions" form:"actions"`
-	DynamicPrompts []DynamicPromptsConfig `json:"dynamic_prompts" form:"dynamic_prompts"`
-	MCPServers     []agent.MCPServer      `json:"mcp_servers" form:"mcp_servers"`
+	Connector        []ConnectorConfig      `json:"connectors" form:"connectors" `
+	Actions          []ActionsConfig        `json:"actions" form:"actions"`
+	DynamicPrompts   []DynamicPromptsConfig `json:"dynamic_prompts" form:"dynamic_prompts"`
+	MCPServers       []agent.MCPServer      `json:"mcp_servers" form:"mcp_servers"`
+	MCPSTDIOServers  []agent.MCPSTDIOServer `json:"mcp_stdio_servers" form:"mcp_stdio_servers"`
+	MCPPrepareScript string                 `json:"mcp_prepare_script" form:"mcp_prepare_script"`
+	MCPBoxURL        string                 `json:"mcp_box_url" form:"mcp_box_url"`

 	Description string `json:"description" form:"description"`

@@ -271,6 +276,22 @@ func NewAgentConfigMeta(
 				HelpText:     "Number of concurrent tasks that can run in parallel",
 				Tags:         config.Tags{Section: "AdvancedSettings"},
 			},
+			{
+				Name:         "mcp_stdio_servers",
+				Label:        "MCP STDIO Servers",
+				Type:         "textarea",
+				DefaultValue: "",
+				HelpText:     "JSON configuration for MCP STDIO servers",
+				Tags:         config.Tags{Section: "AdvancedSettings"},
+			},
+			{
+				Name:         "mcp_prepare_script",
+				Label:        "MCP Prepare Script",
+				Type:         "textarea",
+				DefaultValue: "",
+				HelpText:     "Script to prepare the MCP box",
+				Tags:         config.Tags{Section: "AdvancedSettings"},
+			},
 		},
 		MCPServers: []config.Field{
 			{
@@ -297,3 +318,148 @@ type Connector interface {
 	AgentReasoningCallback() func(state types.ActionCurrentState) bool
 	Start(a *agent.Agent)
 }
+
+// UnmarshalJSON implements json.Unmarshaler for AgentConfig
+func (a *AgentConfig) UnmarshalJSON(data []byte) error {
+	// Create a temporary type to avoid infinite recursion
+	type Alias AgentConfig
+	aux := &struct {
+		*Alias
+		MCPSTDIOServersConfig interface{} `json:"mcp_stdio_servers"`
+	}{
+		Alias: (*Alias)(a),
+	}
+
+	if err := json.Unmarshal(data, &aux); err != nil {
+		return err
+	}
+
+	// Handle MCP STDIO servers configuration
+	if aux.MCPSTDIOServersConfig != nil {
+		switch v := aux.MCPSTDIOServersConfig.(type) {
+		case string:
+			// Parse string configuration
+			var mcpConfig struct {
+				MCPServers map[string]struct {
+					Command string            `json:"command"`
+					Args    []string          `json:"args"`
+					Env     map[string]string `json:"env"`
+				} `json:"mcpServers"`
+			}
+
+			if err := json.Unmarshal([]byte(v), &mcpConfig); err != nil {
+				return fmt.Errorf("failed to parse MCP STDIO servers configuration: %w", err)
+			}
+
+			a.MCPSTDIOServers = make([]agent.MCPSTDIOServer, 0, len(mcpConfig.MCPServers))
+			for _, server := range mcpConfig.MCPServers {
+				// Convert env map to slice of "KEY=VALUE" strings
+				envSlice := make([]string, 0, len(server.Env))
+				for k, v := range server.Env {
+					envSlice = append(envSlice, fmt.Sprintf("%s=%s", k, v))
+				}
+
+				a.MCPSTDIOServers = append(a.MCPSTDIOServers, agent.MCPSTDIOServer{
+					Cmd:  server.Command,
+					Args: server.Args,
+					Env:  envSlice,
+				})
+			}
+		case []interface{}:
+			// Parse array configuration
+			a.MCPSTDIOServers = make([]agent.MCPSTDIOServer, 0, len(v))
+			for _, server := range v {
+				serverMap, ok := server.(map[string]interface{})
+				if !ok {
+					return fmt.Errorf("invalid server configuration format")
+				}
+
+				cmd, _ := serverMap["cmd"].(string)
+				args := make([]string, 0)
+				if argsInterface, ok := serverMap["args"].([]interface{}); ok {
+					for _, arg := range argsInterface {
+						if argStr, ok := arg.(string); ok {
+							args = append(args, argStr)
+						}
+					}
+				}
+
+				env := make([]string, 0)
+				if envInterface, ok := serverMap["env"].([]interface{}); ok {
+					for _, e := range envInterface {
+						if envStr, ok := e.(string); ok {
+							env = append(env, envStr)
+						}
+					}
+				}
+
+				a.MCPSTDIOServers = append(a.MCPSTDIOServers, agent.MCPSTDIOServer{
+					Cmd:  cmd,
+					Args: args,
+					Env:  env,
+				})
+			}
+		}
+	}
+
+	return nil
+}
+
+// MarshalJSON implements json.Marshaler for AgentConfig
+func (a *AgentConfig) MarshalJSON() ([]byte, error) {
+	// Create a temporary type to avoid infinite recursion
+	type Alias AgentConfig
+	aux := &struct {
+		*Alias
+		MCPSTDIOServersConfig string `json:"mcp_stdio_servers,omitempty"`
+	}{
+		Alias: (*Alias)(a),
+	}
+
+	// Convert MCPSTDIOServers back to the expected JSON format
+	if len(a.MCPSTDIOServers) > 0 {
+		mcpConfig := struct {
+			MCPServers map[string]struct {
+				Command string            `json:"command"`
+				Args    []string          `json:"args"`
+				Env     map[string]string `json:"env"`
+			} `json:"mcpServers"`
+		}{
+			MCPServers: make(map[string]struct {
+				Command string            `json:"command"`
+				Args    []string          `json:"args"`
+				Env     map[string]string `json:"env"`
+			}),
+		}
+
+		// Convert each MCPSTDIOServer to the expected format
+		for i, server := range a.MCPSTDIOServers {
+			// Convert env slice back to map
+			envMap := make(map[string]string)
+			for _, env := range server.Env {
+				if parts := strings.SplitN(env, "=", 2); len(parts) == 2 {
+					envMap[parts[0]] = parts[1]
+				}
+			}
+
+			mcpConfig.MCPServers[fmt.Sprintf("server%d", i)] = struct {
+				Command string            `json:"command"`
+				Args    []string          `json:"args"`
+				Env     map[string]string `json:"env"`
+			}{
+				Command: server.Cmd,
+				Args:    server.Args,
+				Env:     envMap,
+			}
+		}
+
+		// Marshal the MCP config to JSON string
+		mcpConfigJSON, err := json.Marshal(mcpConfig)
+		if err != nil {
+			return nil, fmt.Errorf("failed to marshal MCP STDIO servers configuration: %w", err)
+		}
+		aux.MCPSTDIOServersConfig = string(mcpConfigJSON)
+	}
+
+	return json.Marshal(aux)
+}
--- a/core/state/pool.go
+++ b/core/state/pool.go
@@ -33,6 +33,7 @@ type AgentPool struct {
 	managers                                     map[string]sse.Manager
 	agentStatus                                  map[string]*Status
 	apiURL, defaultModel, defaultMultimodalModel string
+	mcpBoxURL                                    string
 	imageModel, localRAGAPI, localRAGKey, apiKey string
 	availableActions                             func(*AgentConfig) func(ctx context.Context, pool *AgentPool) []types.Action
 	connectors                                   func(*AgentConfig) []Connector
@@ -72,7 +73,7 @@ func loadPoolFromFile(path string) (*AgentPoolData, error) {
 }

 func NewAgentPool(
-	defaultModel, defaultMultimodalModel, imageModel, apiURL, apiKey, directory string,
+	defaultModel, defaultMultimodalModel, imageModel, apiURL, apiKey, directory, mcpBoxURL string,
 	LocalRAGAPI string,
 	availableActions func(*AgentConfig) func(ctx context.Context, pool *AgentPool) []types.Action,
 	connectors func(*AgentConfig) []Connector,
@@ -98,6 +99,7 @@ func NewAgentPool(
 			apiURL:                 apiURL,
 			defaultModel:           defaultModel,
 			defaultMultimodalModel: defaultMultimodalModel,
+			mcpBoxURL:              mcpBoxURL,
 			imageModel:             imageModel,
 			localRAGAPI:            LocalRAGAPI,
 			apiKey:                 apiKey,
@@ -123,6 +125,7 @@ func NewAgentPool(
 		pooldir:                directory,
 		defaultModel:           defaultModel,
 		defaultMultimodalModel: defaultMultimodalModel,
+		mcpBoxURL:              mcpBoxURL,
 		imageModel:             imageModel,
 		apiKey:                 apiKey,
 		agents:                 make(map[string]*Agent),
@@ -336,6 +339,10 @@ func (a *AgentPool) startAgentWithConfig(name string, config *AgentConfig, obs O
 		model = config.Model
 	}

+	if config.MCPBoxURL != "" {
+		a.mcpBoxURL = config.MCPBoxURL
+	}
+
 	if config.PeriodicRuns == "" {
 		config.PeriodicRuns = "10m"
 	}
@@ -396,7 +403,10 @@ func (a *AgentPool) startAgentWithConfig(name string, config *AgentConfig, obs O
 		WithMCPServers(config.MCPServers...),
 		WithPeriodicRuns(config.PeriodicRuns),
 		WithPermanentGoal(config.PermanentGoal),
+		WithMCPSTDIOServers(config.MCPSTDIOServers...),
+		WithMCPBoxURL(a.mcpBoxURL),
 		WithPrompts(promptBlocks...),
+		WithMCPPrepareScript(config.MCPPrepareScript),
 		//	WithDynamicPrompts(dynamicPrompts...),
 		WithCharacter(Character{
 			Name: name,