agents_and_robots/shell/mcp/client.go

// Package mcp provides MCP client and server implementations.
package mcp

import (
	"context"
	"fmt"
	"log/slog"
	"os"
	"time"

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

// Client wraps an MCP client (stdio or SSE) and exposes discovered tools.
type Client struct {
	name      string
	transport string // "stdio" | "sse"
	mcpClient *client.Client
	tools     []mcp.Tool
	logger    *slog.Logger
}

// NewStdioClient creates an MCP client that connects to a stdio-based MCP server.
func NewStdioClient(ctx context.Context, name, command string, args []string, env map[string]string, logger *slog.Logger) (*Client, error) {
	logger.Info("creating stdio MCP client", "name", name, "command", command, "args", args)

	// Prepare environment
	envSlice := os.Environ()
	for k, v := range env {
		envSlice = append(envSlice, fmt.Sprintf("%s=%s", k, v))
	}

	// Create stdio client
	mcpClient, err := client.NewStdioMCPClient(command, envSlice, args...)
	if err != nil {
		return nil, fmt.Errorf("failed to create stdio client: %w", err)
	}

	// Initialize
	initReq := mcp.InitializeRequest{
		Params: mcp.InitializeParams{
			ProtocolVersion: mcp.LATEST_PROTOCOL_VERSION,
			ClientInfo: mcp.Implementation{
				Name:    "agents-mcp-client",
				Version: "1.0.0",
			},
			Capabilities: mcp.ClientCapabilities{},
		},
	}

	_, err = mcpClient.Initialize(ctx, initReq)
	if err != nil {
		mcpClient.Close()
		return nil, fmt.Errorf("failed to initialize MCP client: %w", err)
	}

	// Discover tools
	toolsResp, err := mcpClient.ListTools(ctx, mcp.ListToolsRequest{})
	if err != nil {
		mcpClient.Close()
		return nil, fmt.Errorf("failed to list tools: %w", err)
	}

	logger.Info("discovered MCP tools", "name", name, "count", len(toolsResp.Tools))

	return &Client{
		name:      name,
		transport: "stdio",
		mcpClient: mcpClient,
		tools:     toolsResp.Tools,
		logger:    logger,
	}, nil
}

// NewSSEClient creates an MCP client that connects to an SSE/HTTP-based MCP server.
func NewSSEClient(ctx context.Context, name, url string, headers map[string]string, logger *slog.Logger) (*Client, error) {
	logger.Info("creating SSE MCP client", "name", name, "url", url)

	// Create SSE client (no custom headers support in basic API, would need transport options)
	mcpClient, err := client.NewSSEMCPClient(url)
	if err != nil {
		return nil, fmt.Errorf("failed to create SSE client: %w", err)
	}

	// Initialize
	initReq := mcp.InitializeRequest{
		Params: mcp.InitializeParams{
			ProtocolVersion: mcp.LATEST_PROTOCOL_VERSION,
			ClientInfo: mcp.Implementation{
				Name:    "agents-mcp-client",
				Version: "1.0.0",
			},
			Capabilities: mcp.ClientCapabilities{},
		},
	}

	_, err = mcpClient.Initialize(ctx, initReq)
	if err != nil {
		mcpClient.Close()
		return nil, fmt.Errorf("failed to initialize MCP client: %w", err)
	}

	// Discover tools
	toolsResp, err := mcpClient.ListTools(ctx, mcp.ListToolsRequest{})
	if err != nil {
		mcpClient.Close()
		return nil, fmt.Errorf("failed to list tools: %w", err)
	}

	logger.Info("discovered MCP tools", "name", name, "count", len(toolsResp.Tools))

	return &Client{
		name:      name,
		transport: "sse",
		mcpClient: mcpClient,
		tools:     toolsResp.Tools,
		logger:    logger,
	}, nil
}

// Tools returns the discovered MCP tools.
func (c *Client) Tools() []mcp.Tool {
	return c.tools
}

// Name returns the client name.
func (c *Client) Name() string {
	return c.name
}

// CallTool invokes an MCP tool by name with the given arguments.
func (c *Client) CallTool(ctx context.Context, name string, args map[string]any, timeout time.Duration) (*mcp.CallToolResult, error) {
	if timeout > 0 {
		var cancel context.CancelFunc
		ctx, cancel = context.WithTimeout(ctx, timeout)
		defer cancel()
	}

	req := mcp.CallToolRequest{
		Params: mcp.CallToolParams{
			Name:      name,
			Arguments: args,
		},
	}

	result, err := c.mcpClient.CallTool(ctx, req)
	if err != nil {
		return nil, fmt.Errorf("tool call failed: %w", err)
	}

	return result, nil
}

// Close closes the MCP client connection.
func (c *Client) Close() error {
	c.logger.Info("closing MCP client", "name", c.name, "transport", c.transport)
	return c.mcpClient.Close()
}