mirror of
https://github.com/YspCoder/clawgo.git
synced 2026-05-18 11:14:36 +08:00
feat: add MCP tool and web UI management
This commit is contained in:
lpf
354
pkg/tools/mcp_test.go
Normal file
354
pkg/tools/mcp_test.go
Normal file
@@ -0,0 +1,354 @@
|
||||
package tools
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"os"
|
||||
"strconv"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"clawgo/pkg/config"
|
||||
)
|
||||
|
||||
func TestMCPToolListServers(t *testing.T) {
|
||||
tool := NewMCPTool("/tmp/workspace", config.MCPToolsConfig{
|
||||
Enabled: true,
|
||||
RequestTimeoutSec: 5,
|
||||
Servers: map[string]config.MCPServerConfig{
|
||||
"demo": {
|
||||
Enabled: true,
|
||||
Transport: "stdio",
|
||||
Command: "demo-server",
|
||||
Description: "demo",
|
||||
},
|
||||
"disabled": {
|
||||
Enabled: false,
|
||||
Transport: "stdio",
|
||||
Command: "nope",
|
||||
},
|
||||
},
|
||||
})
|
||||
out, err := tool.Execute(context.Background(), map[string]interface{}{"action": "list_servers"})
|
||||
if err != nil {
|
||||
t.Fatalf("list_servers returned error: %v", err)
|
||||
}
|
||||
if !strings.Contains(out, `"name": "demo"`) {
|
||||
t.Fatalf("expected enabled server in output, got: %s", out)
|
||||
}
|
||||
if strings.Contains(out, "disabled") {
|
||||
t.Fatalf("did not expect disabled server in output, got: %s", out)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMCPToolCallTool(t *testing.T) {
|
||||
tool := NewMCPTool(t.TempDir(), config.MCPToolsConfig{
|
||||
Enabled: true,
|
||||
RequestTimeoutSec: 5,
|
||||
Servers: map[string]config.MCPServerConfig{
|
||||
"helper": {
|
||||
Enabled: true,
|
||||
Transport: "stdio",
|
||||
Command: os.Args[0],
|
||||
Args: []string{"-test.run=TestMCPHelperProcess", "--"},
|
||||
Env: map[string]string{
|
||||
"GO_WANT_HELPER_PROCESS": "1",
|
||||
},
|
||||
},
|
||||
},
|
||||
})
|
||||
out, err := tool.Execute(context.Background(), map[string]interface{}{
|
||||
"action": "call_tool",
|
||||
"server": "helper",
|
||||
"tool": "echo",
|
||||
"arguments": map[string]interface{}{"text": "hello"},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("call_tool returned error: %v", err)
|
||||
}
|
||||
if !strings.Contains(out, "echo:hello") {
|
||||
t.Fatalf("expected echo output, got: %s", out)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMCPToolDiscoverTools(t *testing.T) {
|
||||
tool := NewMCPTool(t.TempDir(), config.MCPToolsConfig{
|
||||
Enabled: true,
|
||||
RequestTimeoutSec: 5,
|
||||
Servers: map[string]config.MCPServerConfig{
|
||||
"helper": {
|
||||
Enabled: true,
|
||||
Transport: "stdio",
|
||||
Command: os.Args[0],
|
||||
Args: []string{"-test.run=TestMCPHelperProcess", "--"},
|
||||
Env: map[string]string{
|
||||
"GO_WANT_HELPER_PROCESS": "1",
|
||||
},
|
||||
},
|
||||
},
|
||||
})
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
discovered := tool.DiscoverTools(ctx)
|
||||
if len(discovered) != 1 {
|
||||
t.Fatalf("expected 1 discovered tool, got %d", len(discovered))
|
||||
}
|
||||
if got := discovered[0].Name(); got != "mcp__helper__echo" {
|
||||
t.Fatalf("unexpected discovered tool name: %s", got)
|
||||
}
|
||||
out, err := discovered[0].Execute(ctx, map[string]interface{}{"text": "world"})
|
||||
if err != nil {
|
||||
t.Fatalf("discovered tool execute returned error: %v", err)
|
||||
}
|
||||
if strings.TrimSpace(out) != "echo:world" {
|
||||
t.Fatalf("unexpected discovered tool output: %q", out)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMCPHelperProcess(t *testing.T) {
|
||||
if os.Getenv("GO_WANT_HELPER_PROCESS") != "1" {
|
||||
return
|
||||
}
|
||||
runMCPHelper()
|
||||
os.Exit(0)
|
||||
}
|
||||
|
||||
func runMCPHelper() {
|
||||
reader := bufio.NewReader(os.Stdin)
|
||||
writer := bufio.NewWriter(os.Stdout)
|
||||
for {
|
||||
msg, err := readHelperFrame(reader)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
method, _ := msg["method"].(string)
|
||||
id, hasID := msg["id"]
|
||||
switch method {
|
||||
case "initialize":
|
||||
writeHelperFrame(writer, map[string]interface{}{
|
||||
"jsonrpc": "2.0",
|
||||
"id": id,
|
||||
"result": map[string]interface{}{
|
||||
"protocolVersion": mcpProtocolVersion,
|
||||
"capabilities": map[string]interface{}{
|
||||
"tools": map[string]interface{}{},
|
||||
"resources": map[string]interface{}{},
|
||||
"prompts": map[string]interface{}{},
|
||||
},
|
||||
"serverInfo": map[string]interface{}{
|
||||
"name": "helper",
|
||||
"version": "1.0.0",
|
||||
},
|
||||
},
|
||||
})
|
||||
case "notifications/initialized":
|
||||
continue
|
||||
case "tools/list":
|
||||
writeHelperFrame(writer, map[string]interface{}{
|
||||
"jsonrpc": "2.0",
|
||||
"id": id,
|
||||
"result": map[string]interface{}{
|
||||
"tools": []map[string]interface{}{
|
||||
{
|
||||
"name": "echo",
|
||||
"description": "Echo the provided text",
|
||||
"inputSchema": map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"text": map[string]interface{}{
|
||||
"type": "string",
|
||||
},
|
||||
},
|
||||
"required": []string{"text"},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
})
|
||||
case "tools/call":
|
||||
params, _ := msg["params"].(map[string]interface{})
|
||||
args, _ := params["arguments"].(map[string]interface{})
|
||||
text, _ := args["text"].(string)
|
||||
writeHelperFrame(writer, map[string]interface{}{
|
||||
"jsonrpc": "2.0",
|
||||
"id": id,
|
||||
"result": map[string]interface{}{
|
||||
"content": []map[string]interface{}{
|
||||
{
|
||||
"type": "text",
|
||||
"text": "echo:" + text,
|
||||
},
|
||||
},
|
||||
},
|
||||
})
|
||||
case "resources/list":
|
||||
writeHelperFrame(writer, map[string]interface{}{
|
||||
"jsonrpc": "2.0",
|
||||
"id": id,
|
||||
"result": map[string]interface{}{
|
||||
"resources": []map[string]interface{}{
|
||||
{"uri": "file:///tmp/demo.txt", "name": "demo"},
|
||||
},
|
||||
},
|
||||
})
|
||||
case "resources/read":
|
||||
writeHelperFrame(writer, map[string]interface{}{
|
||||
"jsonrpc": "2.0",
|
||||
"id": id,
|
||||
"result": map[string]interface{}{
|
||||
"contents": []map[string]interface{}{
|
||||
{"uri": "file:///tmp/demo.txt", "mimeType": "text/plain", "text": "demo content"},
|
||||
},
|
||||
},
|
||||
})
|
||||
case "prompts/list":
|
||||
writeHelperFrame(writer, map[string]interface{}{
|
||||
"jsonrpc": "2.0",
|
||||
"id": id,
|
||||
"result": map[string]interface{}{
|
||||
"prompts": []map[string]interface{}{
|
||||
{"name": "greeter", "description": "Greets"},
|
||||
},
|
||||
},
|
||||
})
|
||||
case "prompts/get":
|
||||
params, _ := msg["params"].(map[string]interface{})
|
||||
args, _ := params["arguments"].(map[string]interface{})
|
||||
name, _ := args["name"].(string)
|
||||
writeHelperFrame(writer, map[string]interface{}{
|
||||
"jsonrpc": "2.0",
|
||||
"id": id,
|
||||
"result": map[string]interface{}{
|
||||
"description": "Greets",
|
||||
"messages": []map[string]interface{}{
|
||||
{
|
||||
"role": "user",
|
||||
"content": map[string]interface{}{
|
||||
"type": "text",
|
||||
"text": "hello " + name,
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
})
|
||||
default:
|
||||
if hasID {
|
||||
writeHelperFrame(writer, map[string]interface{}{
|
||||
"jsonrpc": "2.0",
|
||||
"id": id,
|
||||
"error": map[string]interface{}{
|
||||
"code": -32601,
|
||||
"message": "method not found",
|
||||
},
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func readHelperFrame(r *bufio.Reader) (map[string]interface{}, error) {
|
||||
length := 0
|
||||
for {
|
||||
line, err := r.ReadString('\n')
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
line = strings.TrimRight(line, "\r\n")
|
||||
if line == "" {
|
||||
break
|
||||
}
|
||||
parts := strings.SplitN(line, ":", 2)
|
||||
if len(parts) == 2 && strings.EqualFold(strings.TrimSpace(parts[0]), "Content-Length") {
|
||||
length, _ = strconv.Atoi(strings.TrimSpace(parts[1]))
|
||||
}
|
||||
}
|
||||
body := make([]byte, length)
|
||||
if _, err := io.ReadFull(r, body); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var msg map[string]interface{}
|
||||
if err := json.Unmarshal(body, &msg); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return msg, nil
|
||||
}
|
||||
|
||||
func writeHelperFrame(w *bufio.Writer, payload map[string]interface{}) {
|
||||
data, _ := json.Marshal(payload)
|
||||
_, _ = fmt.Fprintf(w, "Content-Length: %d\r\n\r\n%s", len(data), data)
|
||||
_ = w.Flush()
|
||||
}
|
||||
|
||||
func TestValidateMCPTools(t *testing.T) {
|
||||
cfg := config.DefaultConfig()
|
||||
cfg.Tools.MCP.Enabled = true
|
||||
cfg.Tools.MCP.RequestTimeoutSec = 0
|
||||
cfg.Tools.MCP.Servers = map[string]config.MCPServerConfig{
|
||||
"bad": {
|
||||
Enabled: true,
|
||||
Transport: "http",
|
||||
Command: "",
|
||||
WorkingDir: "relative",
|
||||
},
|
||||
}
|
||||
errs := config.Validate(cfg)
|
||||
if len(errs) == 0 {
|
||||
t.Fatal("expected validation errors")
|
||||
}
|
||||
got := make([]string, 0, len(errs))
|
||||
for _, err := range errs {
|
||||
got = append(got, err.Error())
|
||||
}
|
||||
joined := strings.Join(got, "\n")
|
||||
for _, want := range []string{
|
||||
"tools.mcp.request_timeout_sec must be > 0 when tools.mcp.enabled=true",
|
||||
"tools.mcp.servers.bad.transport must be 'stdio'",
|
||||
"tools.mcp.servers.bad.command is required when enabled=true",
|
||||
"tools.mcp.servers.bad.working_dir must be an absolute path",
|
||||
} {
|
||||
if !strings.Contains(joined, want) {
|
||||
t.Fatalf("expected validation error %q in:\n%s", want, joined)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestMCPToolListTools(t *testing.T) {
|
||||
tool := NewMCPTool(t.TempDir(), config.MCPToolsConfig{
|
||||
Enabled: true,
|
||||
RequestTimeoutSec: 5,
|
||||
Servers: map[string]config.MCPServerConfig{
|
||||
"helper": {
|
||||
Enabled: true,
|
||||
Transport: "stdio",
|
||||
Command: os.Args[0],
|
||||
Args: []string{"-test.run=TestMCPHelperProcess", "--"},
|
||||
Env: map[string]string{
|
||||
"GO_WANT_HELPER_PROCESS": "1",
|
||||
},
|
||||
},
|
||||
},
|
||||
})
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
out, err := tool.Execute(ctx, map[string]interface{}{
|
||||
"action": "list_tools",
|
||||
"server": "helper",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("list_tools returned error: %v", err)
|
||||
}
|
||||
if !strings.Contains(out, `"name": "echo"`) {
|
||||
t.Fatalf("expected tool listing, got: %s", out)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildMCPDynamicToolName(t *testing.T) {
|
||||
got := buildMCPDynamicToolName("Context7 Server", "resolve-library.id")
|
||||
if got != "mcp__context7_server__resolve_library_id" {
|
||||
t.Fatalf("unexpected tool name: %s", got)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user