Sidekick/mcp.go

package sidekick

import (
  "context"
  "encoding/json"
  "fmt"
  "io"
  "os"
  "os/exec"
  "strings"
  "time"

  "github.com/mark3labs/mcp-go/client"
  "github.com/mark3labs/mcp-go/client/transport"
  "github.com/mark3labs/mcp-go/mcp"
)

// MCPClientInterface represents a scoped connection to an MCP server
type MCPClientInterface interface {
  CallTool(ctx context.Context, name string, args map[string]interface{}) (string, error)
  ListTools(ctx context.Context) ([]ToolDefinition, error)
  Close() error
}

type mcpGoClientWrapper struct {
  client *client.Client
}

func (w *mcpGoClientWrapper) CallTool(ctx context.Context, name string,
  args map[string]interface{}) (string, error) {
  req := mcp.CallToolRequest{}
  req.Params.Name = name
  req.Params.Arguments = args

  res, err := w.client.CallTool(ctx, req)
  if err != nil {
    return "", err
  }

  if res.IsError {
    return "", fmt.Errorf("tool returned error")
  }

  // Extract text from result content
  var output string
  for _, c := range res.Content {
    if textContent, ok := c.(mcp.TextContent); ok {
      output += textContent.Text + "\n"
    } else {
      output += fmt.Sprintf("%v\n", c)
    }
  }
  return output, nil
}

func (w *mcpGoClientWrapper) ListTools(ctx context.Context) ([]ToolDefinition, error) {
  req := mcp.ListToolsRequest{}
  res, err := w.client.ListTools(ctx, req)
  if err != nil {
    return nil, err
  }

  var tools []ToolDefinition
  for _, t := range res.Tools {
    var params map[string]interface{}
    b, err := json.Marshal(t.InputSchema)
    if err == nil {
      _ = json.Unmarshal(b, &params)
    }

    tools = append(tools, ToolDefinition{
      Name:        t.Name,
      Description: t.Description,
      Parameters:  params,
      Internal:    false,
    })
  }
  return tools, nil
}

func (w *mcpGoClientWrapper) Close() error {
  if w.client != nil {
    return w.client.Close()
  }
  return nil
}

// Global server configurations
var mcpServers map[string]MCPServerConfig

// GlobalLogger is an optional logger for MCP connections
var GlobalLogger func(format string, v ...interface{})

// InitMCPServers sets up the global registry for MCP factories
func InitMCPServers(servers map[string]MCPServerConfig) {
  mcpServers = servers
}

type sidekickLogger struct{}

func (l *sidekickLogger) Infof(format string, v ...any) {
  if GlobalLogger != nil {
    GlobalLogger(format, v...)
  }
}

func (l *sidekickLogger) Errorf(format string, v ...any) {
  if GlobalLogger != nil {
    GlobalLogger("ERROR: "+format, v...)
  }
}

// defaultMCPFactory creates a real MCP client based on the global configuration
var defaultMCPFactory MCPClientFactory = func(toolset string) (MCPClientInterface, error) {
  cfg, ok := mcpServers[toolset]
  if !ok {
    return nil, fmt.Errorf("unknown MCP server: %s", toolset)
  }

  if GlobalLogger != nil {
    GlobalLogger("Initializing MCP connection to %s via %s", toolset, cfg.Transport)
  }

  var c *client.Client
  var err error

  switch cfg.Transport {
  case "stdio":
    if GlobalLogger != nil {
      GlobalLogger("Running stdio MCP: %s %v (env: %d vars)", cfg.Command, cfg.Args, len(cfg.Env))
    }

    command := cfg.Command
    args := cfg.Args

    // If using npx, try to be quiet to avoid stdout pollution
    if command == "npx" {
      quietFound := false
      for _, arg := range args {
        if arg == "--quiet" || arg == "-q" {
          quietFound = true
          break
        }
      }
      if !quietFound {
        // Prepend --quiet to args
        args = append([]string{"--quiet"}, args...)
      }
    }

    env := cfg.Env
    if len(env) > 0 {
      foundPath := false
      for _, e := range env {
        if strings.HasPrefix(strings.ToUpper(e), "PATH=") {
          foundPath = true
          break
        }
      }
      if !foundPath {
        if path := os.Getenv("PATH"); path != "" {
          env = append(env, "PATH="+path)
        }
      }
    } else {
      env = os.Environ()
    }

    // Use WithCommandFunc to configure the command environment
    c, err = client.NewStdioMCPClientWithOptions(command, env, args,
      transport.WithCommandFunc(func(ctx context.Context, command string, env []string, args []string) (*exec.Cmd, error) {
        cmd := exec.CommandContext(ctx, command, args...)
        cmd.Env = env
        return cmd, nil
      }),
    )
  case "sse":
    var opts []transport.ClientOption
    headers := make(map[string]string)
    if cfg.AuthToken != "" {
      headers["Authorization"] = "Bearer " + cfg.AuthToken
    }
    for k, v := range cfg.Headers {
      headers[k] = v
    }
    if len(headers) > 0 {
      opts = append(opts, transport.WithHeaders(headers))
    }
    c, err = client.NewSSEMCPClient(cfg.URL, opts...)
  case "http":
    var opts []transport.StreamableHTTPCOption
    headers := make(map[string]string)
    if cfg.AuthToken != "" {
      headers["Authorization"] = "Bearer " + cfg.AuthToken
    }
    for k, v := range cfg.Headers {
      headers[k] = v
    }
    if len(headers) > 0 {
      opts = append(opts, transport.WithHTTPHeaders(headers))
    }
    c, err = client.NewStreamableHttpClient(cfg.URL, opts...)
  default:
    return nil, fmt.Errorf("unsupported MCP transport: %s", cfg.Transport)
  }

  if err != nil {
    return nil, err
  }

  // Redirect stderr from the stdio transport to os.Stderr
  if cfg.Transport == "stdio" {
    if stderr, ok := client.GetStderr(c); ok {
      go io.Copy(os.Stderr, stderr)
    }
  }

  // Initialize the client session
  initReq := mcp.InitializeRequest{}
  initReq.Params.ProtocolVersion = mcp.LATEST_PROTOCOL_VERSION
  initReq.Params.ClientInfo = mcp.Implementation{
    Name:    "Sidekick",
    Version: "1.0.0",
  }

  initCtx, initCancel := context.WithTimeout(context.Background(), 20*time.Second)
  defer initCancel()
  _, err = c.Initialize(initCtx, initReq)
  if err != nil {
    _ = c.Close()
    return nil, fmt.Errorf("failed to initialize MCP client: %w", err)
  }

  return &mcpGoClientWrapper{client: c}, nil
}

// MCPClientFactory is a function that creates a new MCP client for a specific toolset
type MCPClientFactory func(toolset string) (MCPClientInterface, error)

// SetMCPClientFactory allows overriding the factory for testing
func SetMCPClientFactory(factory MCPClientFactory) {
  defaultMCPFactory = factory
}

// LoadExternalTools connects to all registered MCP servers and fetches their tools
func LoadExternalTools(ctx context.Context) []ToolDefinition {
  var extTools []ToolDefinition
  for toolset := range mcpServers {
    c, err := defaultMCPFactory(toolset)
    if err != nil {
      if GlobalLogger != nil {
        GlobalLogger("Failed to connect to MCP server %s: %v", toolset, err)
      }
      continue
    }

    listCtx, listCancel := context.WithTimeout(ctx, 10*time.Second)
    tools, err := c.ListTools(listCtx)
    listCancel()
    c.Close()
    if err != nil {
      if GlobalLogger != nil {
        GlobalLogger("Failed to list tools from MCP server %s: %v", toolset, err)
      }
      continue
    }

    // Assign toolset to each tool
    for i := range tools {
      tools[i].Toolset = toolset
    }

    if GlobalLogger != nil {
      GlobalLogger("Loaded %d tools from MCP server %s", len(tools), toolset)
    }
    extTools = append(extTools, tools...)
  }
  return extTools
}

// CallExternalTool creates a scoped MCP connection, calls the tool, and cleans up
func CallExternalTool(ctx context.Context, toolset string, toolName string,
  args map[string]interface{}) (string, error) {
  c, err := defaultMCPFactory(toolset)
  if err != nil {
    return "", fmt.Errorf("failed to create MCP client for %s: %w", toolset, err)
  }
  defer c.Close()

  return c.CallTool(ctx, toolName, args)
}