go-llm/v2/mcp.go

package llm

import (
	"context"
	"encoding/json"
	"fmt"
	"os"
	"os/exec"
	"sync"

	"github.com/modelcontextprotocol/go-sdk/mcp"
)

// MCPTransport specifies how to connect to an MCP server.
type MCPTransport string

const (
	MCPStdio MCPTransport = "stdio"
	MCPSSE   MCPTransport = "sse"
	MCPHTTP  MCPTransport = "http"
)

// MCPServer represents a connection to an MCP server.
type MCPServer struct {
	name      string
	transport MCPTransport

	// stdio fields
	command string
	args    []string
	env     []string

	// network fields
	url string

	// internal
	client  *mcp.Client
	session *mcp.ClientSession
	tools   map[string]*mcp.Tool
	mu      sync.RWMutex
}

// MCPOption configures an MCP server.
type MCPOption func(*MCPServer)

// WithMCPEnv adds environment variables for the subprocess.
func WithMCPEnv(env ...string) MCPOption {
	return func(s *MCPServer) { s.env = env }
}

// WithMCPName sets a friendly name for logging.
func WithMCPName(name string) MCPOption {
	return func(s *MCPServer) { s.name = name }
}

// MCPStdioServer creates and connects to an MCP server via stdio transport.
//
// Example:
//
//	server, err := llm.MCPStdioServer(ctx, "npx", "-y", "@modelcontextprotocol/server-filesystem", "/tmp")
func MCPStdioServer(ctx context.Context, command string, args ...string) (*MCPServer, error) {
	s := &MCPServer{
		name:      command,
		transport: MCPStdio,
		command:   command,
		args:      args,
	}
	if err := s.connect(ctx); err != nil {
		return nil, err
	}
	return s, nil
}

// MCPHTTPServer creates and connects to an MCP server via streamable HTTP transport.
//
// Example:
//
//	server, err := llm.MCPHTTPServer(ctx, "https://mcp.example.com")
func MCPHTTPServer(ctx context.Context, url string, opts ...MCPOption) (*MCPServer, error) {
	s := &MCPServer{
		name:      url,
		transport: MCPHTTP,
		url:       url,
	}
	for _, opt := range opts {
		opt(s)
	}
	if err := s.connect(ctx); err != nil {
		return nil, err
	}
	return s, nil
}

// MCPSSEServer creates and connects to an MCP server via SSE transport.
func MCPSSEServer(ctx context.Context, url string, opts ...MCPOption) (*MCPServer, error) {
	s := &MCPServer{
		name:      url,
		transport: MCPSSE,
		url:       url,
	}
	for _, opt := range opts {
		opt(s)
	}
	if err := s.connect(ctx); err != nil {
		return nil, err
	}
	return s, nil
}

func (s *MCPServer) connect(ctx context.Context) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	if s.session != nil {
		return nil
	}

	s.client = mcp.NewClient(&mcp.Implementation{
		Name:    "go-llm-v2",
		Version: "2.0.0",
	}, nil)

	var transport mcp.Transport

	switch s.transport {
	case MCPSSE:
		transport = &mcp.SSEClientTransport{
			Endpoint: s.url,
		}
	case MCPHTTP:
		transport = &mcp.StreamableClientTransport{
			Endpoint: s.url,
		}
	default: // stdio
		cmd := exec.Command(s.command, s.args...)
		cmd.Env = append(os.Environ(), s.env...)
		transport = &mcp.CommandTransport{
			Command: cmd,
		}
	}

	session, err := s.client.Connect(ctx, transport, nil)
	if err != nil {
		return fmt.Errorf("failed to connect to MCP server %s: %w", s.name, err)
	}

	s.session = session

	// Load tools
	s.tools = make(map[string]*mcp.Tool)
	for tool, err := range session.Tools(ctx, nil) {
		if err != nil {
			s.session.Close()
			s.session = nil
			return fmt.Errorf("failed to list tools from %s: %w", s.name, err)
		}
		s.tools[tool.Name] = tool
	}

	return nil
}

// Close closes the connection to the MCP server.
func (s *MCPServer) Close() error {
	s.mu.Lock()
	defer s.mu.Unlock()

	if s.session == nil {
		return nil
	}

	err := s.session.Close()
	s.session = nil
	s.tools = nil
	return err
}

// IsConnected returns true if the server is connected.
func (s *MCPServer) IsConnected() bool {
	s.mu.RLock()
	defer s.mu.RUnlock()
	return s.session != nil
}

// ListTools returns Tool definitions for all tools this server provides.
func (s *MCPServer) ListTools() []Tool {
	s.mu.RLock()
	defer s.mu.RUnlock()

	var tools []Tool
	for _, t := range s.tools {
		tools = append(tools, s.toTool(t))
	}
	return tools
}

// CallTool invokes a tool on the server.
func (s *MCPServer) CallTool(ctx context.Context, name string, arguments map[string]any) (string, error) {
	s.mu.RLock()
	session := s.session
	s.mu.RUnlock()

	if session == nil {
		return "", fmt.Errorf("%w: %s", ErrNotConnected, s.name)
	}

	result, err := session.CallTool(ctx, &mcp.CallToolParams{
		Name:      name,
		Arguments: arguments,
	})
	if err != nil {
		return "", err
	}

	if len(result.Content) == 0 {
		return "", nil
	}

	return contentToString(result.Content), nil
}

func (s *MCPServer) toTool(t *mcp.Tool) Tool {
	var inputSchema map[string]any
	if t.InputSchema != nil {
		data, err := json.Marshal(t.InputSchema)
		if err == nil {
			_ = json.Unmarshal(data, &inputSchema)
		}
	}

	if inputSchema == nil {
		inputSchema = map[string]any{
			"type":       "object",
			"properties": map[string]any{},
		}
	}

	return Tool{
		Name:        t.Name,
		Description: t.Description,
		Schema:      inputSchema,
		isMCP:       true,
		mcpServer:   s,
	}
}

func contentToString(content []mcp.Content) string {
	var parts []string
	for _, c := range content {
		switch tc := c.(type) {
		case *mcp.TextContent:
			parts = append(parts, tc.Text)
		default:
			if data, err := json.Marshal(c); err == nil {
				parts = append(parts, string(data))
			}
		}
	}
	if len(parts) == 1 {
		return parts[0]
	}
	data, _ := json.Marshal(parts)
	return string(data)
}