Add MCP tool support via official Go SDK
This commit is contained in:
Spark
parent
d83fb6e081
commit
bc4b3e2048
9 changed files with 767 additions and 16 deletions
58
README.md
58
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"
|
||||
}
|
||||
]
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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": {
|
||||
|
|
|
|||
|
|
@ -77,6 +77,10 @@
|
|||
"api_key": "YOUR_BRAVE_API_KEY",
|
||||
"max_results": 5
|
||||
}
|
||||
},
|
||||
"mcp": {
|
||||
"enabled": false,
|
||||
"servers": []
|
||||
}
|
||||
},
|
||||
"gateway": {
|
||||
|
|
|
|||
3
go.mod
3
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
|
||||
|
|
|
|||
12
go.sum
12
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=
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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{},
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
|
|
|||
446
pkg/tools/mcp.go
Normal file
446
pkg/tools/mcp.go
Normal file
|
|
@ -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
|
||||
}
|
||||
176
pkg/tools/mcp_test.go
Normal file
176
pkg/tools/mcp_test.go
Normal file
|
|
@ -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
|
||||
}
|
||||
Loading…
Add table
Reference in a new issue