diff --git a/README.md b/README.md index 381998211..220dcfb94 100644 --- a/README.md +++ b/README.md @@ -198,6 +198,19 @@ picoclaw onboard "api_key": "YOUR_BRAVE_API_KEY", "max_results": 5 } + }, + "mcp": { + "enabled": true, + "servers": [ + { + "name": "filesystem", + "enabled": true, + "transport": "command", + "command": "npx", + "args": ["-y", "@modelcontextprotocol/server-filesystem", "/tmp"], + "tool_prefix": "mcp_fs" + } + ] } } } @@ -218,6 +231,38 @@ picoclaw agent -m "What is 2+2?" That's it! You have a working AI assistant in 2 minutes. +### 🔌 MCP Servers (Official Go SDK) + +PicoClaw supports MCP clients via the official SDK (`github.com/modelcontextprotocol/go-sdk`). + +- `transport: "command"` for stdio subprocess MCP servers +- `transport: "streamable_http"` for streamable HTTP servers +- `transport: "sse"` for legacy SSE servers + +Example: + +```json +{ + "tools": { + "mcp": { + "enabled": true, + "servers": [ + { + "name": "filesystem", + "enabled": true, + "transport": "command", + "command": "npx", + "args": ["-y", "@modelcontextprotocol/server-filesystem", "/tmp"], + "tool_prefix": "mcp_fs", + "startup_timeout_ms": 8000, + "call_timeout_ms": 30000 + } + ] + } + } +} +``` + --- ## 💬 Chat Apps @@ -512,6 +557,19 @@ picoclaw agent -m "Hello" "search": { "api_key": "BSA..." } + }, + "mcp": { + "enabled": true, + "servers": [ + { + "name": "filesystem", + "enabled": true, + "transport": "command", + "command": "npx", + "args": ["-y", "@modelcontextprotocol/server-filesystem", "/tmp"], + "tool_prefix": "mcp_fs" + } + ] } } } diff --git a/config/config.example.json b/config/config.example.json index ed5cb7048..8d7360e33 100644 --- a/config/config.example.json +++ b/config/config.example.json @@ -98,6 +98,21 @@ "api_key": "YOUR_BRAVE_API_KEY", "max_results": 5 } + }, + "mcp": { + "enabled": false, + "servers": [ + { + "name": "filesystem", + "enabled": false, + "transport": "command", + "command": "npx", + "args": ["-y", "@modelcontextprotocol/server-filesystem", "/tmp"], + "tool_prefix": "mcp_fs", + "startup_timeout_ms": 8000, + "call_timeout_ms": 30000 + } + ] } }, "gateway": { diff --git a/config/config.openrouter.json b/config/config.openrouter.json index 4aca883d4..4eacb2cfa 100644 --- a/config/config.openrouter.json +++ b/config/config.openrouter.json @@ -77,10 +77,14 @@ "api_key": "YOUR_BRAVE_API_KEY", "max_results": 5 } + }, + "mcp": { + "enabled": false, + "servers": [] } }, "gateway": { "host": "0.0.0.0", "port": 18790 } -} \ No newline at end of file +} diff --git a/go.mod b/go.mod index f4c233ea8..34c069d30 100644 --- a/go.mod +++ b/go.mod @@ -11,6 +11,7 @@ require ( github.com/google/uuid v1.6.0 github.com/gorilla/websocket v1.5.3 github.com/larksuite/oapi-sdk-go/v3 v3.5.3 + github.com/modelcontextprotocol/go-sdk v1.2.0 github.com/mymmrac/telego v1.6.0 github.com/open-dingtalk/dingtalk-stream-sdk-go v0.9.1 github.com/openai/openai-go/v3 v3.21.0 @@ -27,6 +28,7 @@ require ( github.com/cloudwego/base64x v0.1.6 // indirect github.com/go-resty/resty/v2 v2.17.1 // indirect github.com/gogo/protobuf v1.3.2 // indirect + github.com/google/jsonschema-go v0.3.0 // indirect github.com/grbit/go-json v0.11.0 // indirect github.com/klauspost/compress v1.18.4 // indirect github.com/klauspost/cpuid/v2 v2.3.0 // indirect @@ -38,6 +40,7 @@ require ( github.com/valyala/bytebufferpool v1.0.0 // indirect github.com/valyala/fasthttp v1.69.0 // indirect github.com/valyala/fastjson v1.6.7 // indirect + github.com/yosida95/uritemplate/v3 v3.0.2 // indirect golang.org/x/arch v0.24.0 // indirect golang.org/x/crypto v0.48.0 // indirect golang.org/x/net v0.50.0 // indirect diff --git a/go.sum b/go.sum index 9174d2889..b2873c117 100644 --- a/go.sum +++ b/go.sum @@ -41,6 +41,8 @@ github.com/go-test/deep v1.1.1 h1:0r/53hagsehfO4bzD2Pgr/+RgHqhmf+k1Bpse2cTu1U= github.com/go-test/deep v1.1.1/go.mod h1:5C2ZWiW0ErCdrYzpqxLbTX7MG14M9iiw8DgHncVwcsE= github.com/gogo/protobuf v1.3.2 h1:Ov1cvc58UF3b5XjBnZv7+opcTcQFZebYjWzi34vdm4Q= github.com/gogo/protobuf v1.3.2/go.mod h1:P1XiOD3dCwIKUDQYPy72D8LYyHL2YPYrpS2s69NZV8Q= +github.com/golang-jwt/jwt/v5 v5.2.2 h1:Rl4B7itRWVtYIHFrSNd7vhTiz9UpLdi6gZhZ3wEeDy8= +github.com/golang-jwt/jwt/v5 v5.2.2/go.mod h1:pqrtFR0X4osieyHYxtmOUWsAWrfe1Q5UVIyoH402zdk= github.com/golang/protobuf v1.2.0/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U= github.com/golang/protobuf v1.4.0-rc.1/go.mod h1:ceaxUfeHdC40wWswd/P6IGgMaK3YpKi5j83Wpe3EHw8= github.com/golang/protobuf v1.4.0-rc.1.0.20200221234624-67d41d38c208/go.mod h1:xKAWHe0F5eneWXFV3EuXVDTCmh+JuBKY0li0aMyXATA= @@ -56,6 +58,10 @@ github.com/google/go-cmp v0.4.0/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/ github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= github.com/google/go-cmp v0.5.6/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= github.com/google/go-cmp v0.5.9/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY= +github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8= +github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU= +github.com/google/jsonschema-go v0.3.0 h1:6AH2TxVNtk3IlvkkhjrtbUc4S8AvO0Xii0DxIygDg+Q= +github.com/google/jsonschema-go v0.3.0/go.mod h1:r5quNTdLOYEz95Ru18zA0ydNbBuYoo9tgaYcxEYhJVE= github.com/google/uuid v1.3.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= @@ -80,6 +86,8 @@ github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI= github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE= github.com/larksuite/oapi-sdk-go/v3 v3.5.3 h1:xvf8Dv29kBXC5/DNDCLhHkAFW8l/0LlQJimO5Zn+JUk= github.com/larksuite/oapi-sdk-go/v3 v3.5.3/go.mod h1:ZEplY+kwuIrj/nqw5uSCINNATcH3KdxSN7y+UxYY5fI= +github.com/modelcontextprotocol/go-sdk v1.2.0 h1:Y23co09300CEk8iZ/tMxIX1dVmKZkzoSBZOpJwUnc/s= +github.com/modelcontextprotocol/go-sdk v1.2.0/go.mod h1:6fM3LCm3yV7pAs8isnKLn07oKtB0MP9LHd3DfAcKw10= github.com/mymmrac/telego v1.6.0 h1:Zc8rgyHozvd/7ZgyrigyHdAF9koHYMfilYfyB6wlFC0= github.com/mymmrac/telego v1.6.0/go.mod h1:xt6ZWA8zi8KmuzryE1ImEdl9JSwjHNpM4yhC7D8hU4Y= github.com/nxadm/tail v1.4.4/go.mod h1:kenIhsEOeOJmVchQTgglprH7qJGnHDVpk1VPCcaMI8A= @@ -137,6 +145,8 @@ github.com/valyala/fastjson v1.6.7 h1:ZE4tRy0CIkh+qDc5McjatheGX2czdn8slQjomexVpB github.com/valyala/fastjson v1.6.7/go.mod h1:CLCAqky6SMuOcxStkYQvblddUtoRxhYMGLrsQns1aXY= github.com/xyproto/randomstring v1.0.5 h1:YtlWPoRdgMu3NZtP45drfy1GKoojuR7hmRcnhZqKjWU= github.com/xyproto/randomstring v1.0.5/go.mod h1:rgmS5DeNXLivK7YprL0pY+lTuhNQW3iGxZ18UQApw/E= +github.com/yosida95/uritemplate/v3 v3.0.2 h1:Ed3Oyj9yrmi9087+NczuL5BwkIc4wvTb5zIM+UJPGz4= +github.com/yosida95/uritemplate/v3 v3.0.2/go.mod h1:ILOh0sOhIJR3+L/8afwt/kE++YT040gmv5BQTMR2HP4= github.com/yuin/goldmark v1.1.27/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74= github.com/yuin/goldmark v1.2.1/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74= github.com/yuin/goldmark v1.4.13/go.mod h1:6yULJ656Px+3vBD8DxQVa3kxgyrAnzto9xy5taEt/CY= @@ -224,6 +234,8 @@ golang.org/x/tools v0.0.0-20201224043029-2b0845dc783e/go.mod h1:emZCQorbCU4vsT4f golang.org/x/tools v0.0.0-20210106214847-113979e3529a/go.mod h1:emZCQorbCU4vsT4fOWvOPXz4eW1wZW4PmDk9uLelYpA= golang.org/x/tools v0.1.12/go.mod h1:hNGJHUnrk76NpqgfD5Aqm5Crs+Hm0VOH/i9J2+nxYbc= golang.org/x/tools v0.6.0/go.mod h1:Xwgl3UAJ/d3gWutnCtw505GrjyAbvKui8lOU390QaIU= +golang.org/x/tools v0.34.0 h1:qIpSLOxeCYGg9TrcJokLBG4KFA6d795g0xkBkiESGlo= +golang.org/x/tools v0.34.0/go.mod h1:pAP9OwEaY1CAW3HOmg3hLZC5Z0CCmzjAF2UQMSqNARg= golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go index fac2856e9..ac4c659dd 100644 --- a/pkg/agent/loop.go +++ b/pkg/agent/loop.go @@ -31,13 +31,13 @@ type AgentLoop struct { provider providers.LLMProvider workspace string model string - contextWindow int // Maximum context window size in tokens + contextWindow int // Maximum context window size in tokens maxIterations int sessions *session.SessionManager contextBuilder *ContextBuilder tools *tools.ToolRegistry running atomic.Bool - summarizing sync.Map // Tracks which sessions are currently being summarized + summarizing sync.Map // Tracks which sessions are currently being summarized } // processOptions configures how a message is processed @@ -89,6 +89,18 @@ func NewAgentLoop(cfg *config.Config, msgBus *bus.MessageBus, provider providers toolsRegistry.Register(editFileTool) toolsRegistry.Register(tools.NewAppendFileTool(workspace, restrict)) + // Register MCP-discovered tools (best effort; continue on per-server failures) + mcpTools, mcpErr := tools.LoadMCPTools(context.Background(), cfg.Tools.MCP, workspace) + if mcpErr != nil { + logger.WarnCF("agent", "Some MCP servers failed to load", + map[string]interface{}{ + "error": mcpErr.Error(), + }) + } + for _, tool := range mcpTools { + toolsRegistry.Register(tool) + } + sessionsManager := session.NewSessionManager(filepath.Join(workspace, "sessions")) // Create context builder and set tools registry diff --git a/pkg/config/config.go b/pkg/config/config.go index 56f1e1958..e8c5bceb3 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -57,13 +57,13 @@ type AgentsConfig struct { } type AgentDefaults struct { - Workspace string `json:"workspace" env:"PICOCLAW_AGENTS_DEFAULTS_WORKSPACE"` - RestrictToWorkspace bool `json:"restrict_to_workspace" env:"PICOCLAW_AGENTS_DEFAULTS_RESTRICT_TO_WORKSPACE"` - Provider string `json:"provider" env:"PICOCLAW_AGENTS_DEFAULTS_PROVIDER"` - Model string `json:"model" env:"PICOCLAW_AGENTS_DEFAULTS_MODEL"` - MaxTokens int `json:"max_tokens" env:"PICOCLAW_AGENTS_DEFAULTS_MAX_TOKENS"` - Temperature float64 `json:"temperature" env:"PICOCLAW_AGENTS_DEFAULTS_TEMPERATURE"` - MaxToolIterations int `json:"max_tool_iterations" env:"PICOCLAW_AGENTS_DEFAULTS_MAX_TOOL_ITERATIONS"` + Workspace string `json:"workspace" env:"PICOCLAW_AGENTS_DEFAULTS_WORKSPACE"` + RestrictToWorkspace bool `json:"restrict_to_workspace" env:"PICOCLAW_AGENTS_DEFAULTS_RESTRICT_TO_WORKSPACE"` + Provider string `json:"provider" env:"PICOCLAW_AGENTS_DEFAULTS_PROVIDER"` + Model string `json:"model" env:"PICOCLAW_AGENTS_DEFAULTS_MODEL"` + MaxTokens int `json:"max_tokens" env:"PICOCLAW_AGENTS_DEFAULTS_MAX_TOKENS"` + Temperature float64 `json:"temperature" env:"PICOCLAW_AGENTS_DEFAULTS_TEMPERATURE"` + MaxToolIterations int `json:"max_tool_iterations" env:"PICOCLAW_AGENTS_DEFAULTS_MAX_TOOL_ITERATIONS"` } type ChannelsConfig struct { @@ -166,21 +166,42 @@ type WebToolsConfig struct { Search WebSearchConfig `json:"search"` } +type MCPServerConfig struct { + Name string `json:"name"` + Enabled bool `json:"enabled"` + Transport string `json:"transport"` // command|streamable_http|sse + Command string `json:"command,omitempty"` + Args []string `json:"args,omitempty"` + Env map[string]string `json:"env,omitempty"` + WorkingDir string `json:"working_dir,omitempty"` + URL string `json:"url,omitempty"` + ToolPrefix string `json:"tool_prefix,omitempty"` + StartupTimeoutMS int `json:"startup_timeout_ms,omitempty"` + CallTimeoutMS int `json:"call_timeout_ms,omitempty"` + TerminateTimeoutMS int `json:"terminate_timeout_ms,omitempty"` +} + +type MCPToolsConfig struct { + Enabled bool `json:"enabled"` + Servers []MCPServerConfig `json:"servers"` +} + type ToolsConfig struct { Web WebToolsConfig `json:"web"` + MCP MCPToolsConfig `json:"mcp"` } func DefaultConfig() *Config { return &Config{ Agents: AgentsConfig{ Defaults: AgentDefaults{ - Workspace: "~/.picoclaw/workspace", + Workspace: "~/.picoclaw/workspace", RestrictToWorkspace: true, - Provider: "", - Model: "glm-4.7", - MaxTokens: 8192, - Temperature: 0.7, - MaxToolIterations: 20, + Provider: "", + Model: "glm-4.7", + MaxTokens: 8192, + Temperature: 0.7, + MaxToolIterations: 20, }, }, Channels: ChannelsConfig{ @@ -254,6 +275,10 @@ func DefaultConfig() *Config { MaxResults: 5, }, }, + MCP: MCPToolsConfig{ + Enabled: false, + Servers: []MCPServerConfig{}, + }, }, } } diff --git a/pkg/tools/mcp.go b/pkg/tools/mcp.go new file mode 100644 index 000000000..c4c6a0783 --- /dev/null +++ b/pkg/tools/mcp.go @@ -0,0 +1,446 @@ +package tools + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "os" + "os/exec" + "path/filepath" + "regexp" + "sort" + "strings" + "time" + + "github.com/modelcontextprotocol/go-sdk/mcp" + "github.com/sipeed/picoclaw/pkg/config" +) + +const ( + defaultMCPStartupTimeout = 8 * time.Second + defaultMCPCallTimeout = 30 * time.Second + maxToolNameLength = 64 +) + +var toolNameSanitizer = regexp.MustCompile(`[^a-zA-Z0-9_-]+`) + +// LoadMCPTools discovers tools from configured MCP servers and returns them as local tools. +// Discovery is best-effort across servers: individual server failures are aggregated in the returned error. +func LoadMCPTools(ctx context.Context, cfg config.MCPToolsConfig, workspace string) ([]Tool, error) { + if !cfg.Enabled || len(cfg.Servers) == 0 { + return nil, nil + } + + usedNames := make(map[string]int) + loaded := make([]Tool, 0) + errs := make([]error, 0) + + for _, serverCfg := range cfg.Servers { + if !serverCfg.Enabled { + continue + } + + client := newMCPClient(serverCfg, workspace) + startupTimeout := durationFromMS(serverCfg.StartupTimeoutMS, defaultMCPStartupTimeout) + + connectCtx, cancel := context.WithTimeout(ctx, startupTimeout) + remoteTools, err := client.ListTools(connectCtx) + cancel() + if err != nil { + errs = append(errs, fmt.Errorf("mcp server %q discovery failed: %w", serverCfg.Name, err)) + continue + } + + for _, rt := range remoteTools { + if rt == nil || strings.TrimSpace(rt.Name) == "" { + continue + } + + localName := buildLocalToolName(serverCfg, rt.Name, usedNames) + description := buildMCPToolDescription(serverCfg.Name, rt.Name, rt.Description) + parameters := normalizeMCPInputSchema(rt.InputSchema) + + loaded = append(loaded, &MCPTool{ + localName: localName, + remoteName: rt.Name, + description: description, + parameters: parameters, + callTimeout: durationFromMS(serverCfg.CallTimeoutMS, defaultMCPCallTimeout), + client: client, + }) + } + } + + return loaded, errors.Join(errs...) +} + +type MCPTool struct { + localName string + remoteName string + description string + parameters map[string]interface{} + callTimeout time.Duration + client *mcpClient +} + +func (t *MCPTool) Name() string { + return t.localName +} + +func (t *MCPTool) Description() string { + return t.description +} + +func (t *MCPTool) Parameters() map[string]interface{} { + return t.parameters +} + +func (t *MCPTool) Execute(ctx context.Context, args map[string]interface{}) (string, error) { + callCtx := ctx + if t.callTimeout > 0 { + var cancel context.CancelFunc + callCtx, cancel = context.WithTimeout(ctx, t.callTimeout) + defer cancel() + } + return t.client.CallTool(callCtx, t.remoteName, args) +} + +type mcpClient struct { + cfg config.MCPServerConfig + workspace string + client *mcp.Client +} + +func newMCPClient(cfg config.MCPServerConfig, workspace string) *mcpClient { + implName := strings.TrimSpace(cfg.Name) + if implName == "" { + implName = "picoclaw-mcp" + } + return &mcpClient{ + cfg: cfg, + workspace: workspace, + client: mcp.NewClient(&mcp.Implementation{ + Name: "picoclaw-" + sanitizeToolName(implName), + Version: "v0.1.0", + }, nil), + } +} + +func (c *mcpClient) ListTools(ctx context.Context) ([]*mcp.Tool, error) { + session, err := c.connect(ctx) + if err != nil { + return nil, err + } + defer session.Close() + + all := make([]*mcp.Tool, 0) + cursor := "" + for { + params := &mcp.ListToolsParams{} + if cursor != "" { + params.Cursor = cursor + } + res, err := session.ListTools(ctx, params) + if err != nil { + return nil, fmt.Errorf("list tools: %w", err) + } + all = append(all, res.Tools...) + if res.NextCursor == "" { + break + } + cursor = res.NextCursor + } + return all, nil +} + +func (c *mcpClient) CallTool(ctx context.Context, toolName string, args map[string]interface{}) (string, error) { + session, err := c.connect(ctx) + if err != nil { + return "", err + } + defer session.Close() + + result, err := session.CallTool(ctx, &mcp.CallToolParams{ + Name: toolName, + Arguments: args, + }) + if err != nil { + return "", fmt.Errorf("call tool %q: %w", toolName, err) + } + + return formatMCPCallToolResult(result) +} + +func (c *mcpClient) connect(ctx context.Context) (*mcp.ClientSession, error) { + transport, err := c.buildTransport() + if err != nil { + return nil, err + } + session, err := c.client.Connect(ctx, transport, nil) + if err != nil { + return nil, fmt.Errorf("connect mcp server %q: %w", c.cfg.Name, err) + } + return session, nil +} + +func (c *mcpClient) buildTransport() (mcp.Transport, error) { + transport := strings.ToLower(strings.TrimSpace(c.cfg.Transport)) + if transport == "" { + transport = "command" + } + + switch transport { + case "command": + command := strings.TrimSpace(c.cfg.Command) + if command == "" { + return nil, fmt.Errorf("mcp server %q: command is required for command transport", c.cfg.Name) + } + cmd := exec.Command(command, c.cfg.Args...) + + if wd := resolvePath(c.cfg.WorkingDir, c.workspace); wd != "" { + cmd.Dir = wd + } + + if len(c.cfg.Env) > 0 { + cmd.Env = mergeEnv(os.Environ(), c.cfg.Env) + } + cmd.Stderr = os.Stderr + + tr := &mcp.CommandTransport{ + Command: cmd, + } + if c.cfg.TerminateTimeoutMS > 0 { + tr.TerminateDuration = time.Duration(c.cfg.TerminateTimeoutMS) * time.Millisecond + } + return tr, nil + case "streamable_http": + if strings.TrimSpace(c.cfg.URL) == "" { + return nil, fmt.Errorf("mcp server %q: url is required for streamable_http transport", c.cfg.Name) + } + return &mcp.StreamableClientTransport{ + Endpoint: c.cfg.URL, + }, nil + case "sse": + if strings.TrimSpace(c.cfg.URL) == "" { + return nil, fmt.Errorf("mcp server %q: url is required for sse transport", c.cfg.Name) + } + return &mcp.SSEClientTransport{ + Endpoint: c.cfg.URL, + }, nil + default: + return nil, fmt.Errorf("mcp server %q: unsupported transport %q", c.cfg.Name, c.cfg.Transport) + } +} + +func formatMCPCallToolResult(result *mcp.CallToolResult) (string, error) { + if result == nil { + return "", fmt.Errorf("empty MCP response") + } + + textOnly, ok := singleTextResult(result) + if ok && result.StructuredContent == nil { + if result.IsError { + return "MCP tool error: " + textOnly, nil + } + return textOnly, nil + } + + out := map[string]interface{}{ + "is_error": result.IsError, + } + if len(result.Content) > 0 { + out["content"] = result.Content + } + if result.StructuredContent != nil { + out["structured_content"] = result.StructuredContent + } + + if len(out) == 1 && !result.IsError { + return "(empty MCP tool response)", nil + } + + data, err := json.MarshalIndent(out, "", " ") + if err != nil { + return "", fmt.Errorf("marshal MCP tool response: %w", err) + } + return string(data), nil +} + +func singleTextResult(result *mcp.CallToolResult) (string, bool) { + if result == nil || len(result.Content) != 1 { + return "", false + } + tc, ok := result.Content[0].(*mcp.TextContent) + if !ok { + return "", false + } + return tc.Text, true +} + +func buildLocalToolName(serverCfg config.MCPServerConfig, remoteToolName string, used map[string]int) string { + prefix := strings.TrimSpace(serverCfg.ToolPrefix) + if prefix == "" { + baseServer := sanitizeToolName(serverCfg.Name) + if baseServer == "" { + baseServer = "server" + } + prefix = "mcp_" + baseServer + } + + base := sanitizeToolName(prefix + "_" + remoteToolName) + if base == "" { + base = "mcp_tool" + } + + candidate := truncateToolName(base) + if used[candidate] == 0 { + used[candidate] = 1 + return candidate + } + + for i := 2; ; i++ { + suffix := fmt.Sprintf("_%d", i) + candidate = truncateWithSuffix(base, suffix) + if used[candidate] == 0 { + used[candidate] = 1 + return candidate + } + } +} + +func buildMCPToolDescription(serverName, remoteName, rawDescription string) string { + base := strings.TrimSpace(rawDescription) + if base == "" { + base = fmt.Sprintf("Call MCP tool %q.", remoteName) + } + + serverName = strings.TrimSpace(serverName) + if serverName == "" { + return "[MCP] " + base + } + + return fmt.Sprintf("[MCP %s/%s] %s", serverName, remoteName, base) +} + +func normalizeMCPInputSchema(schema any) map[string]interface{} { + fallback := map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{}, + } + + if schema == nil { + return fallback + } + + var out map[string]interface{} + switch v := schema.(type) { + case map[string]interface{}: + out = v + default: + data, err := json.Marshal(v) + if err != nil { + return fallback + } + if err := json.Unmarshal(data, &out); err != nil { + return fallback + } + } + + if out == nil { + return fallback + } + if _, ok := out["type"]; !ok { + out["type"] = "object" + } + if out["type"] == "object" { + if _, ok := out["properties"]; !ok { + out["properties"] = map[string]interface{}{} + } + } + return out +} + +func sanitizeToolName(name string) string { + name = strings.TrimSpace(name) + if name == "" { + return "" + } + name = strings.ReplaceAll(name, " ", "_") + name = toolNameSanitizer.ReplaceAllString(name, "_") + name = strings.Trim(name, "_-") + return name +} + +func truncateToolName(name string) string { + if len(name) <= maxToolNameLength { + return name + } + return name[:maxToolNameLength] +} + +func truncateWithSuffix(base, suffix string) string { + if len(suffix) >= maxToolNameLength { + return suffix[len(suffix)-maxToolNameLength:] + } + maxBase := maxToolNameLength - len(suffix) + if len(base) > maxBase { + base = base[:maxBase] + } + return base + suffix +} + +func durationFromMS(value int, fallback time.Duration) time.Duration { + if value <= 0 { + return fallback + } + return time.Duration(value) * time.Millisecond +} + +func resolvePath(path, workspace string) string { + path = strings.TrimSpace(path) + if path == "" { + return "" + } + path = expandHome(path) + if filepath.IsAbs(path) { + return path + } + if workspace != "" { + return filepath.Join(workspace, path) + } + return path +} + +func expandHome(path string) string { + if path == "" || path[0] != '~' { + return path + } + home, err := os.UserHomeDir() + if err != nil { + return path + } + if len(path) == 1 { + return home + } + if path[1] == '/' { + return home + path[1:] + } + return path +} + +func mergeEnv(base []string, extra map[string]string) []string { + if len(extra) == 0 { + return base + } + merged := append([]string{}, base...) + keys := make([]string, 0, len(extra)) + for k := range extra { + keys = append(keys, k) + } + sort.Strings(keys) + for _, k := range keys { + merged = append(merged, fmt.Sprintf("%s=%s", k, extra[k])) + } + return merged +} diff --git a/pkg/tools/mcp_test.go b/pkg/tools/mcp_test.go new file mode 100644 index 000000000..1662e9cdf --- /dev/null +++ b/pkg/tools/mcp_test.go @@ -0,0 +1,176 @@ +package tools + +import ( + "context" + "os" + "strings" + "testing" + + "github.com/modelcontextprotocol/go-sdk/mcp" + "github.com/sipeed/picoclaw/pkg/config" +) + +const mcpHelperEnv = "PICOCLAW_MCP_TEST_HELPER" + +func TestMain(m *testing.M) { + if os.Getenv(mcpHelperEnv) == "1" { + runMCPHelperServer() + os.Exit(0) + } + os.Exit(m.Run()) +} + +func runMCPHelperServer() { + type GreetInput struct { + Name string `json:"name" jsonschema:"name to greet"` + } + type GreetOutput struct { + Greeting string `json:"greeting"` + } + type SumInput struct { + A int `json:"a" jsonschema:"first number"` + B int `json:"b" jsonschema:"second number"` + } + + server := mcp.NewServer(&mcp.Implementation{Name: "picoclaw-test-server", Version: "v1.0.0"}, nil) + mcp.AddTool(server, &mcp.Tool{Name: "greet", Description: "return a greeting"}, func(_ context.Context, _ *mcp.CallToolRequest, in GreetInput) (*mcp.CallToolResult, GreetOutput, error) { + return &mcp.CallToolResult{ + Content: []mcp.Content{ + &mcp.TextContent{Text: "Hello " + in.Name}, + }, + }, GreetOutput{Greeting: "Hello " + in.Name}, nil + }) + mcp.AddTool(server, &mcp.Tool{Name: "sum", Description: "sum two integers"}, func(_ context.Context, _ *mcp.CallToolRequest, in SumInput) (*mcp.CallToolResult, map[string]int, error) { + return nil, map[string]int{"sum": in.A + in.B}, nil + }) + + if err := server.Run(context.Background(), &mcp.StdioTransport{}); err != nil { + os.Exit(1) + } +} + +func TestLoadMCPTools_CommandTransport(t *testing.T) { + cfg := config.MCPToolsConfig{ + Enabled: true, + Servers: []config.MCPServerConfig{ + { + Name: "helper", + Enabled: true, + Transport: "command", + Command: os.Args[0], + Args: []string{}, + Env: map[string]string{mcpHelperEnv: "1"}, + StartupTimeoutMS: 8000, + CallTimeoutMS: 5000, + ToolPrefix: "mcp_helper", + }, + }, + } + + tools, err := LoadMCPTools(context.Background(), cfg, t.TempDir()) + if err != nil { + t.Fatalf("LoadMCPTools() error: %v", err) + } + if len(tools) != 2 { + t.Fatalf("LoadMCPTools() got %d tools, want 2", len(tools)) + } + + var greetTool Tool + var sumTool Tool + for _, tool := range tools { + switch tool.Name() { + case "mcp_helper_greet": + greetTool = tool + case "mcp_helper_sum": + sumTool = tool + } + } + + if greetTool == nil { + t.Fatalf("missing discovered tool mcp_helper_greet; got names=%v", toolNames(tools)) + } + if sumTool == nil { + t.Fatalf("missing discovered tool mcp_helper_sum; got names=%v", toolNames(tools)) + } + + gotGreeting, err := greetTool.Execute(context.Background(), map[string]interface{}{"name": "Ada"}) + if err != nil { + t.Fatalf("greetTool.Execute() error: %v", err) + } + if !strings.Contains(gotGreeting, "Hello Ada") { + t.Fatalf("greetTool.Execute() missing greeting: %s", gotGreeting) + } + + gotSum, err := sumTool.Execute(context.Background(), map[string]interface{}{"a": 2, "b": 3}) + if err != nil { + t.Fatalf("sumTool.Execute() error: %v", err) + } + if !strings.Contains(gotSum, `"sum": 5`) { + t.Fatalf("sumTool.Execute() output missing sum result: %s", gotSum) + } +} + +func TestBuildLocalToolName_EnsuresUniqueness(t *testing.T) { + used := map[string]int{} + cfg := config.MCPServerConfig{Name: "my server", ToolPrefix: "mcp_my_server"} + + name1 := buildLocalToolName(cfg, "echo", used) + name2 := buildLocalToolName(cfg, "echo", used) + + if name1 == name2 { + t.Fatalf("expected unique names, got both %q", name1) + } + if len(name1) > maxToolNameLength || len(name2) > maxToolNameLength { + t.Fatalf("tool name length exceeded %d: %q / %q", maxToolNameLength, name1, name2) + } +} + +func TestNormalizeMCPInputSchema_DefaultObject(t *testing.T) { + schema := normalizeMCPInputSchema(nil) + if schema["type"] != "object" { + t.Fatalf("schema.type = %v, want object", schema["type"]) + } + if _, ok := schema["properties"]; !ok { + t.Fatalf("schema.properties missing") + } +} + +func TestResolvePath_RelativeUsesWorkspace(t *testing.T) { + got := resolvePath("servers/time", "/tmp/workspace") + if got != "/tmp/workspace/servers/time" { + t.Fatalf("resolvePath() = %q, want %q", got, "/tmp/workspace/servers/time") + } +} + +func TestLoadMCPTools_InvalidServerAggregatesError(t *testing.T) { + cfg := config.MCPToolsConfig{ + Enabled: true, + Servers: []config.MCPServerConfig{ + { + Name: "broken", + Enabled: true, + Transport: "command", + Command: "", + }, + }, + } + + tools, err := LoadMCPTools(context.Background(), cfg, t.TempDir()) + if len(tools) != 0 { + t.Fatalf("expected no tools, got %d", len(tools)) + } + if err == nil { + t.Fatalf("expected discovery error, got nil") + } + if !strings.Contains(err.Error(), "discovery failed") { + t.Fatalf("unexpected error: %v", err) + } +} + +func toolNames(tools []Tool) []string { + out := make([]string, 0, len(tools)) + for _, tool := range tools { + out = append(out, tool.Name()) + } + return out +}