sketch: add MCP support
Lets you initialize extra tools via MCP. This is additive, so it's
harmless enough.
Here are some examples of the kind of things you can pass to the -mcp
flag.
{"name": "context7", "type": "http", "url": "https://mcp.context7.com/mcp"}
{"name": "context7-http", "type": "http", "url": "https://mcp.context7.com/mcp"}
{"name": "context7-stdio", "type": "stdio", "command": "npx", "args": ["-y", "@upstash/context7-mcp"]}
{"name": "context7-sse", "type": "sse", "url": "https://mcp.context7.com/sse"}
{"name": "local-tool", "type": "stdio", "command": "my_tool", "args": ["--option", "value"], "env": {"TOKEN": "secret"}}
{ "name": "playwright", "command": "npx", "args": [ "@playwright/mcp@latest" ]}
Co-Authored-By: sketch <hello@sketch.dev>
Change-ID: s259a35d11e7bd660k
diff --git a/cmd/sketch/main.go b/cmd/sketch/main.go
index f0b66b7..ef06e49 100644
--- a/cmd/sketch/main.go
+++ b/cmd/sketch/main.go
@@ -218,6 +218,7 @@
branchPrefix string
sshConnectionString string
subtraceToken string
+ mcpServers StringSliceFlag
}
// parseCLIFlags parses all command-line flags and returns a CLIFlags struct
@@ -250,6 +251,7 @@
userFlags.Var(&flags.mounts, "mount", "volume to mount in the container in format /path/on/host:/path/in/container (can be repeated)")
userFlags.BoolVar(&flags.termUI, "termui", true, "enable terminal UI")
userFlags.StringVar(&flags.branchPrefix, "branch-prefix", "sketch/", "prefix for git branches created by sketch")
+ userFlags.Var(&flags.mcpServers, "mcp", "MCP server configuration as JSON (can be repeated). Schema: {\"name\": \"server-name\", \"type\": \"stdio|http|sse\", \"url\": \"...\", \"command\": \"...\", \"args\": [...], \"env\": {...}, \"headers\": {...}}")
// Internal flags (for sketch developers or internal use)
// Args to sketch innie:
@@ -419,6 +421,7 @@
BranchPrefix: flags.branchPrefix,
LinkToGitHub: flags.linkToGitHub,
SubtraceToken: flags.subtraceToken,
+ MCPServers: flags.mcpServers,
}
if err := dockerimg.LaunchContainer(ctx, config); err != nil {
@@ -538,6 +541,7 @@
BranchPrefix: flags.branchPrefix,
LinkToGitHub: flags.linkToGitHub,
SSHConnectionString: flags.sshConnectionString,
+ MCPServers: flags.mcpServers,
}
// Create SkabandClient if skaband address is provided
diff --git a/dockerimg/dockerimg.go b/dockerimg/dockerimg.go
index e5ad648..c17093a 100644
--- a/dockerimg/dockerimg.go
+++ b/dockerimg/dockerimg.go
@@ -129,6 +129,9 @@
// SubtraceToken enables running sketch under subtrace.dev (development only)
SubtraceToken string
+
+ // MCPServers contains MCP server configurations
+ MCPServers []string
}
// LaunchContainer creates a docker container for a project, installs sketch and opens a connection to it.
@@ -624,6 +627,10 @@
// TODO: select and forward the relevant API key based on the model
cmdArgs = append(cmdArgs, "-llm-api-key="+os.Getenv("ANTHROPIC_API_KEY"))
}
+ // Add MCP server configurations
+ for _, mcpServer := range config.MCPServers {
+ cmdArgs = append(cmdArgs, "-mcp", mcpServer)
+ }
// Add additional docker arguments if provided
if config.DockerArgs != "" {
diff --git a/go.mod b/go.mod
index 363a739..66ab7e1 100644
--- a/go.mod
+++ b/go.mod
@@ -13,6 +13,7 @@
github.com/google/go-cmp v0.7.0
github.com/google/uuid v1.6.0
github.com/kevinburke/ssh_config v1.2.0
+ github.com/mark3labs/mcp-go v0.32.0
github.com/oklog/ulid/v2 v2.1.0
github.com/pkg/sftp v1.13.9
github.com/richardlehane/crock32 v1.0.1
@@ -36,6 +37,8 @@
github.com/kr/fs v0.1.0 // indirect
github.com/mattn/go-colorable v0.1.13 // indirect
github.com/mattn/go-isatty v0.0.20 // indirect
+ github.com/spf13/cast v1.7.1 // indirect
+ github.com/yosida95/uritemplate/v3 v3.0.2 // indirect
golang.org/x/mod v0.24.0 // indirect
golang.org/x/sys v0.33.0 // indirect
golang.org/x/text v0.24.0 // indirect
diff --git a/go.sum b/go.sum
index 431391b..9cbb509 100644
--- a/go.sum
+++ b/go.sum
@@ -17,6 +17,8 @@
github.com/evanw/esbuild v0.25.2/go.mod h1:D2vIQZqV/vIf/VRHtViaUtViZmG7o+kKmlBfVQuRi48=
github.com/fatih/color v1.18.0 h1:S8gINlzdQ840/4pfAwic/ZE0djQEH3wM94VfqLTZcOM=
github.com/fatih/color v1.18.0/go.mod h1:4FelSpRwEGDpQ12mAdzqdOukCy4u8WUtOY6lkT/6HfU=
+github.com/frankban/quicktest v1.14.6 h1:7Xjx+VpznH+oBnejlPUj8oUpdxnVs4f8XU8WnHkI4W8=
+github.com/frankban/quicktest v1.14.6/go.mod h1:4ptaffx2x8+WTWXmUCuVU6aPUX1/Mz7zb5vbUoiM6w0=
github.com/gliderlabs/ssh v0.3.8 h1:a4YXD1V7xMF9g5nTkdfnja3Sxy1PVDCj1Zg4Wb8vY6c=
github.com/gliderlabs/ssh v0.3.8/go.mod h1:xYoytBv1sV0aL3CavoDuJIQNURXkkfPA/wxQ1pL1fAU=
github.com/go-json-experiment/json v0.0.0-20250211171154-1ae217ad3535 h1:yE7argOs92u+sSCRgqqe6eF+cDaVhSPlioy1UkA0p/w=
@@ -44,6 +46,8 @@
github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE=
github.com/ledongthuc/pdf v0.0.0-20220302134840-0c2507a12d80 h1:6Yzfa6GP0rIo/kULo2bwGEkFvCePZ3qHDDTC3/J9Swo=
github.com/ledongthuc/pdf v0.0.0-20220302134840-0c2507a12d80/go.mod h1:imJHygn/1yfhB7XSJJKlFZKl/J+dCPAknuiaGOshXAs=
+github.com/mark3labs/mcp-go v0.32.0 h1:fgwmbfL2gbd67obg57OfV2Dnrhs1HtSdlY/i5fn7MU8=
+github.com/mark3labs/mcp-go v0.32.0/go.mod h1:rXqOudj/djTORU/ThxYx8fqEVj/5pvTuuebQ2RC7uk4=
github.com/mattn/go-colorable v0.1.13 h1:fFA4WZxdEF4tXPZVKMLwD8oUnCTTo08duU7wxecdEvA=
github.com/mattn/go-colorable v0.1.13/go.mod h1:7S9/ev0klgBDR4GtXTXX8a3vIGJpMovkB8vQcUbaXHg=
github.com/mattn/go-isatty v0.0.16/go.mod h1:kYGgaQfpe5nmfYZH+SKPsOc2e4SrIfOl2e/yFXSvRLM=
@@ -64,12 +68,16 @@
github.com/rogpeppe/go-internal v1.14.1/go.mod h1:MaRKkUm5W0goXpeCfT7UZI6fk/L7L7so1lCWt35ZSgc=
github.com/sashabaranov/go-openai v1.38.2 h1:akrssjj+6DY3lWuDwHv6cBvJ8Z+FZDM9XEaaYFt0Auo=
github.com/sashabaranov/go-openai v1.38.2/go.mod h1:lj5b/K+zjTSFxVLijLSTDZuP7adOgerWeFyZLUhAKRg=
+github.com/spf13/cast v1.7.1 h1:cuNEagBQEHWN1FnbGEjCXL2szYEXqfJPbP2HNUaca9Y=
+github.com/spf13/cast v1.7.1/go.mod h1:ancEpBxwJDODSW/UG4rDrAqiKolqNNh2DX3mk86cAdo=
github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw=
github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU=
github.com/stretchr/testify v1.10.0 h1:Xv5erBjTwe/5IxqUQTdXv5kgmIvbHo3QQyRwhJsOfJA=
github.com/stretchr/testify v1.10.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
+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.4.13/go.mod h1:6yULJ656Px+3vBD8DxQVa3kxgyrAnzto9xy5taEt/CY=
go.skia.org/infra v0.0.0-20250421160028-59e18403fd4a h1:XqDi+8oE4eakFiXZXmQlsPaZTTdsPOy54jP3my6lIcU=
go.skia.org/infra v0.0.0-20250421160028-59e18403fd4a/go.mod h1:itQeLiwIYtXPJJEqdxRpOlS77LNv/quHjkyy+SaXrkw=
@@ -121,8 +129,6 @@
golang.org/x/sys v0.17.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
golang.org/x/sys v0.20.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
golang.org/x/sys v0.28.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
-golang.org/x/sys v0.32.0 h1:s77OFDvIQeibCmezSnk/q6iAfkdiQaJi4VzroCFrN20=
-golang.org/x/sys v0.32.0/go.mod h1:BJP2sWEmIv4KK5OTEluFJCKSidICx8ciO85XgH3Ak8k=
golang.org/x/sys v0.33.0 h1:q3i8TbbEz+JRD9ywIRlyRAQbM0qF7hu24q3teo2hbuw=
golang.org/x/sys v0.33.0/go.mod h1:BJP2sWEmIv4KK5OTEluFJCKSidICx8ciO85XgH3Ak8k=
golang.org/x/telemetry v0.0.0-20240228155512-f48c80bd79b2/go.mod h1:TeRTkGYfJXctD9OcfyVLyj2J3IxLnKwHJR8f4D8a3YE=
@@ -134,8 +140,6 @@
golang.org/x/term v0.17.0/go.mod h1:lLRBjIVuehSbZlaOtGMbcMncT+aqLLLmKrsjNrUguwk=
golang.org/x/term v0.20.0/go.mod h1:8UkIAJTvZgivsXaD6/pH6U9ecQzZ45awqEOzuCvwpFY=
golang.org/x/term v0.27.0/go.mod h1:iMsnZpn0cago0GOrHO2+Y7u7JPn5AylBrcoWkElMTSM=
-golang.org/x/term v0.31.0 h1:erwDkOK1Msy6offm1mOgvspSkslFnIGsFnxOKoufg3o=
-golang.org/x/term v0.31.0/go.mod h1:R4BeIy7D95HzImkxGkTW1UQTtP54tio2RyHz7PwK0aw=
golang.org/x/term v0.32.0 h1:DR4lr0TjUs3epypdhTOkMmuF5CDFJ/8pOnbzMZPQ7bg=
golang.org/x/term v0.32.0/go.mod h1:uZG1FhGx848Sqfsq4/DlJr3xGGsYMu/L5GW4abiaEPQ=
golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
@@ -162,7 +166,5 @@
gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
-mvdan.cc/sh/v3 v3.11.0 h1:q5h+XMDRfUGUedCqFFsjoFjrhwf2Mvtt1rkMvVz0blw=
-mvdan.cc/sh/v3 v3.11.0/go.mod h1:LRM+1NjoYCzuq/WZ6y44x14YNAI0NK7FLPeQSaFagGg=
mvdan.cc/sh/v3 v3.11.1-0.20250530001257-46bb4f2b309f h1:T7SkxUwIOTm9iowqyQuUMY9oGEgZy5fE+TWNWgOj+yU=
mvdan.cc/sh/v3 v3.11.1-0.20250530001257-46bb4f2b309f/go.mod h1:Se6Cj17eYSn+sNooLZiEUnNNmNxg0imoYlTu4CyaGyg=
diff --git a/loop/agent.go b/loop/agent.go
index 5c36bcf..ee6d8d7 100644
--- a/loop/agent.go
+++ b/loop/agent.go
@@ -28,6 +28,7 @@
"sketch.dev/llm"
"sketch.dev/llm/ant"
"sketch.dev/llm/conversation"
+ "sketch.dev/mcp"
"sketch.dev/skabandclient"
)
@@ -425,6 +426,8 @@
outsideWorkingDir string
// URL of the git remote 'origin' if it exists
gitOrigin string
+ // MCP manager for handling MCP server connections
+ mcpManager *mcp.MCPManager
// Time when the current turn started (reset at the beginning of InnerLoop)
startOfTurn time.Time
@@ -1028,6 +1031,8 @@
SSHConnectionString string
// Skaband client for session history (optional)
SkabandClient *skabandclient.SkabandClient
+ // MCP server configurations
+ MCPServers []string
}
// NewAgent creates a new Agent.
@@ -1059,6 +1064,7 @@
workingDir: config.WorkingDir,
outsideHTTP: config.OutsideHTTP,
portMonitor: NewPortMonitor(),
+ mcpManager: mcp.NewMCPManager(),
}
return agent
}
@@ -1273,6 +1279,37 @@
convo.Tools = append(convo.Tools, sessionHistoryTools...)
}
+ // Add MCP tools if configured
+ if len(a.config.MCPServers) > 0 {
+ slog.InfoContext(ctx, "Initializing MCP connections", "servers", len(a.config.MCPServers))
+ mcpConnections, mcpErrors := a.mcpManager.ConnectToServers(ctx, a.config.MCPServers, 10*time.Second)
+
+ if len(mcpErrors) > 0 {
+ for _, err := range mcpErrors {
+ slog.ErrorContext(ctx, "MCP connection error", "error", err)
+ // Send agent message about MCP connection failures
+ a.pushToOutbox(ctx, AgentMessage{
+ Type: ErrorMessageType,
+ Content: fmt.Sprintf("MCP server connection failed: %v", err),
+ })
+ }
+ }
+
+ if len(mcpConnections) > 0 {
+ // Add tools from all successful connections
+ totalTools := 0
+ for _, connection := range mcpConnections {
+ convo.Tools = append(convo.Tools, connection.Tools...)
+ totalTools += len(connection.Tools)
+ // Log tools per server using structured data
+ slog.InfoContext(ctx, "Added MCP tools from server", "server", connection.ServerName, "count", len(connection.Tools), "tools", connection.ToolNames)
+ }
+ slog.InfoContext(ctx, "Total MCP tools added", "count", totalTools)
+ } else {
+ slog.InfoContext(ctx, "No MCP tools available after connection attempts")
+ }
+ }
+
convo.Listener = a
return convo
}
@@ -1448,6 +1485,13 @@
a.portMonitor.Start(ctxOuter)
}
+ // Set up cleanup when context is done
+ defer func() {
+ if a.mcpManager != nil {
+ a.mcpManager.Close()
+ }
+ }()
+
for {
select {
case <-ctxOuter.Done():
diff --git a/mcp/client.go b/mcp/client.go
new file mode 100644
index 0000000..e87a465
--- /dev/null
+++ b/mcp/client.go
@@ -0,0 +1,344 @@
+package mcp
+
+import (
+ "context"
+ "encoding/json"
+ "fmt"
+ "log/slog"
+ "strings"
+ "sync"
+ "time"
+
+ "github.com/mark3labs/mcp-go/client"
+ "github.com/mark3labs/mcp-go/client/transport"
+ "github.com/mark3labs/mcp-go/mcp"
+ "sketch.dev/llm"
+)
+
+// ServerConfig represents the configuration for an MCP server
+type ServerConfig struct {
+ Name string `json:"name,omitempty"`
+ Type string `json:"type,omitempty"` // "stdio", "http", "sse"
+ URL string `json:"url,omitempty"` // for http/sse
+ Command string `json:"command,omitempty"` // for stdio
+ Args []string `json:"args,omitempty"` // for stdio
+ Env map[string]string `json:"env,omitempty"` // for stdio
+ Headers map[string]string `json:"headers,omitempty"` // for http/sse
+}
+
+// MCPManager manages multiple MCP server connections
+type MCPManager struct {
+ mu sync.RWMutex
+ clients map[string]*MCPClientWrapper
+}
+
+// MCPClientWrapper wraps an MCP client connection
+type MCPClientWrapper struct {
+ name string
+ config ServerConfig
+ client *client.Client
+ tools []*llm.Tool
+}
+
+// MCPServerConnection represents a successful MCP server connection with its tools
+type MCPServerConnection struct {
+ ServerName string
+ Tools []*llm.Tool
+ ToolNames []string // Original tool names without server prefix
+}
+
+// NewMCPManager creates a new MCP manager
+func NewMCPManager() *MCPManager {
+ return &MCPManager{
+ clients: make(map[string]*MCPClientWrapper),
+ }
+}
+
+// ParseServerConfigs parses JSON configuration strings into ServerConfig structs
+func ParseServerConfigs(ctx context.Context, configs []string) ([]ServerConfig, []error) {
+ if len(configs) == 0 {
+ return nil, nil
+ }
+
+ var serverConfigs []ServerConfig
+ var errors []error
+
+ for i, configStr := range configs {
+ var config ServerConfig
+ if err := json.Unmarshal([]byte(configStr), &config); err != nil {
+ slog.ErrorContext(ctx, "Failed to parse MCP server config", "config", configStr, "error", err)
+ errors = append(errors, fmt.Errorf("config %d: %w", i, err))
+ continue
+ }
+ // Require a name
+ if config.Name == "" {
+ errors = append(errors, fmt.Errorf("config %d: name is required", i))
+ continue
+ }
+ serverConfigs = append(serverConfigs, config)
+ }
+
+ return serverConfigs, errors
+}
+
+// ConnectToServers connects to multiple MCP servers in parallel
+func (m *MCPManager) ConnectToServers(ctx context.Context, configs []string, timeout time.Duration) ([]MCPServerConnection, []error) {
+ serverConfigs, parseErrors := ParseServerConfigs(ctx, configs)
+ if len(serverConfigs) == 0 {
+ if len(parseErrors) > 0 {
+ return nil, parseErrors
+ }
+ return nil, nil
+ }
+ return m.ConnectToServerConfigs(ctx, serverConfigs, timeout, parseErrors)
+} // ConnectToServerConfigs connects to multiple parsed MCP server configs in parallel
+func (m *MCPManager) ConnectToServerConfigs(ctx context.Context, serverConfigs []ServerConfig, timeout time.Duration, existingErrors []error) ([]MCPServerConnection, []error) {
+ if len(serverConfigs) == 0 {
+ return nil, existingErrors
+ }
+
+ slog.InfoContext(ctx, "Connecting to MCP servers", "count", len(serverConfigs), "timeout", timeout)
+
+ // Connect to servers in parallel using sync.WaitGroup
+ type result struct {
+ tools []*llm.Tool
+ err error
+ serverName string
+ originalTools []string // Original tool names without server prefix
+ }
+
+ results := make(chan result, len(serverConfigs))
+ ctxWithTimeout, cancel := context.WithTimeout(ctx, timeout)
+ defer cancel()
+
+ for _, config := range serverConfigs {
+ go func(cfg ServerConfig) {
+ slog.InfoContext(ctx, "Connecting to MCP server", "server", cfg.Name, "type", cfg.Type, "url", cfg.URL, "command", cfg.Command)
+ tools, originalToolNames, err := m.connectToServerWithNames(ctxWithTimeout, cfg)
+ results <- result{
+ tools: tools,
+ err: err,
+ serverName: cfg.Name,
+ originalTools: originalToolNames,
+ }
+ }(config)
+ }
+
+ // Collect results
+ var connections []MCPServerConnection
+ errors := make([]error, 0, len(existingErrors))
+ errors = append(errors, existingErrors...)
+
+ for range len(serverConfigs) {
+ select {
+ case res := <-results:
+ if res.err != nil {
+ slog.ErrorContext(ctx, "Failed to connect to MCP server", "server", res.serverName, "error", res.err)
+ errors = append(errors, fmt.Errorf("MCP server %q: %w", res.serverName, res.err))
+ } else {
+ connection := MCPServerConnection{
+ ServerName: res.serverName,
+ Tools: res.tools,
+ ToolNames: res.originalTools,
+ }
+ connections = append(connections, connection)
+ slog.InfoContext(ctx, "Successfully connected to MCP server", "server", res.serverName, "tools", len(res.tools), "tool_names", res.originalTools)
+ }
+ case <-ctxWithTimeout.Done():
+ errors = append(errors, fmt.Errorf("timeout connecting to MCP servers"))
+ break
+ }
+ }
+
+ return connections, errors
+}
+
+// connectToServerWithNames connects to a single MCP server and returns tools with original names
+func (m *MCPManager) connectToServerWithNames(ctx context.Context, config ServerConfig) ([]*llm.Tool, []string, error) {
+ tools, err := m.connectToServer(ctx, config)
+ if err != nil {
+ return nil, nil, err
+ }
+
+ // Extract original tool names (remove server prefix)
+ originalNames := make([]string, len(tools))
+ for i, tool := range tools {
+ // Tool names are in format "servername_toolname"
+ parts := strings.SplitN(tool.Name, "_", 2)
+ if len(parts) == 2 {
+ originalNames[i] = parts[1]
+ } else {
+ originalNames[i] = tool.Name // fallback if no prefix
+ }
+ }
+
+ return tools, originalNames, nil
+}
+
+// connectToServer connects to a single MCP server
+func (m *MCPManager) connectToServer(ctx context.Context, config ServerConfig) ([]*llm.Tool, error) {
+ var mcpClient *client.Client
+ var err error
+
+ // Convert environment variables to []string format
+ var envVars []string
+ for k, v := range config.Env {
+ envVars = append(envVars, k+"="+v)
+ }
+
+ switch config.Type {
+ case "stdio", "":
+ if config.Command == "" {
+ return nil, fmt.Errorf("command is required for stdio transport")
+ }
+ mcpClient, err = client.NewStdioMCPClient(config.Command, envVars, config.Args...)
+ case "http":
+ if config.URL == "" {
+ return nil, fmt.Errorf("URL is required for HTTP transport")
+ }
+ // Use streamable HTTP client for HTTP transport
+ var httpOptions []transport.StreamableHTTPCOption
+ if len(config.Headers) > 0 {
+ httpOptions = append(httpOptions, transport.WithHTTPHeaders(config.Headers))
+ }
+ mcpClient, err = client.NewStreamableHttpClient(config.URL, httpOptions...)
+ case "sse":
+ if config.URL == "" {
+ return nil, fmt.Errorf("URL is required for SSE transport")
+ }
+ var sseOptions []transport.ClientOption
+ if len(config.Headers) > 0 {
+ sseOptions = append(sseOptions, transport.WithHeaders(config.Headers))
+ }
+ mcpClient, err = client.NewSSEMCPClient(config.URL, sseOptions...)
+ default:
+ return nil, fmt.Errorf("unsupported MCP transport type: %s", config.Type)
+ }
+
+ if err != nil {
+ return nil, fmt.Errorf("failed to create MCP client: %w", err)
+ }
+
+ // Start the client first
+ if err := mcpClient.Start(ctx); err != nil {
+ return nil, fmt.Errorf("failed to start MCP client: %w", err)
+ }
+
+ // Initialize the client
+ initReq := mcp.InitializeRequest{
+ Params: mcp.InitializeParams{
+ ProtocolVersion: mcp.LATEST_PROTOCOL_VERSION,
+ Capabilities: mcp.ClientCapabilities{},
+ ClientInfo: mcp.Implementation{
+ Name: "sketch",
+ Version: "1.0.0",
+ },
+ },
+ }
+ if _, err := mcpClient.Initialize(ctx, initReq); err != nil {
+ return nil, fmt.Errorf("failed to initialize MCP client: %w", err)
+ }
+
+ // Get available tools
+ toolsReq := mcp.ListToolsRequest{}
+ toolsResp, err := mcpClient.ListTools(ctx, toolsReq)
+ if err != nil {
+ return nil, fmt.Errorf("failed to list tools: %w", err)
+ }
+
+ // Convert MCP tools to llm.Tool
+ llmTools, err := m.convertMCPTools(config.Name, mcpClient, toolsResp.Tools)
+ if err != nil {
+ return nil, fmt.Errorf("failed to convert tools: %w", err)
+ }
+
+ // Store the client
+ clientWrapper := &MCPClientWrapper{
+ name: config.Name,
+ config: config,
+ client: mcpClient,
+ tools: llmTools,
+ }
+
+ m.mu.Lock()
+ m.clients[config.Name] = clientWrapper
+ m.mu.Unlock()
+
+ return llmTools, nil
+}
+
+// convertMCPTools converts MCP tools to llm.Tool format
+func (m *MCPManager) convertMCPTools(serverName string, mcpClient *client.Client, mcpTools []mcp.Tool) ([]*llm.Tool, error) {
+ var llmTools []*llm.Tool
+
+ for _, mcpTool := range mcpTools {
+ // Convert the input schema
+ schemaBytes, err := json.Marshal(mcpTool.InputSchema)
+ if err != nil {
+ return nil, fmt.Errorf("failed to marshal input schema for tool %s: %w", mcpTool.Name, err)
+ }
+
+ llmTool := &llm.Tool{
+ Name: fmt.Sprintf("%s_%s", serverName, mcpTool.Name),
+ Description: mcpTool.Description,
+ InputSchema: json.RawMessage(schemaBytes),
+ Run: func(toolName string, client *client.Client) func(ctx context.Context, input json.RawMessage) ([]llm.Content, error) {
+ return func(ctx context.Context, input json.RawMessage) ([]llm.Content, error) {
+ result, err := m.executeMCPTool(ctx, client, toolName, input)
+ if err != nil {
+ return nil, err
+ }
+ // Convert result to llm.Content
+ return []llm.Content{llm.StringContent(fmt.Sprintf("%v", result))}, nil
+ }
+ }(mcpTool.Name, mcpClient),
+ }
+
+ llmTools = append(llmTools, llmTool)
+ }
+
+ return llmTools, nil
+}
+
+// executeMCPTool executes an MCP tool call
+func (m *MCPManager) executeMCPTool(ctx context.Context, mcpClient *client.Client, toolName string, input json.RawMessage) (any, error) {
+ // Add timeout for tool execution
+ ctxWithTimeout, cancel := context.WithTimeout(ctx, 30*time.Second)
+ defer cancel()
+
+ // Parse input arguments
+ var args map[string]any
+ if len(input) > 0 {
+ if err := json.Unmarshal(input, &args); err != nil {
+ return nil, fmt.Errorf("failed to parse tool arguments: %w", err)
+ }
+ }
+
+ // Call the MCP tool
+ req := mcp.CallToolRequest{
+ Params: mcp.CallToolParams{
+ Name: toolName,
+ Arguments: args,
+ },
+ }
+ resp, err := mcpClient.CallTool(ctxWithTimeout, req)
+ if err != nil {
+ return nil, fmt.Errorf("MCP tool call failed: %w", err)
+ }
+
+ // Return the content from the response
+ return resp.Content, nil
+}
+
+// Close closes all MCP client connections
+func (m *MCPManager) Close() {
+ m.mu.Lock()
+ defer m.mu.Unlock()
+
+ for _, clientWrapper := range m.clients {
+ if clientWrapper.client != nil {
+ clientWrapper.client.Close()
+ }
+ }
+ m.clients = make(map[string]*MCPClientWrapper)
+}