agents_and_robots/shell/mcp/manager.go

package mcp

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

	"github.com/enmanuel/agents/internal/config"
)

// Manager manages multiple MCP client connections.
type Manager struct {
	clients map[string]*Client // server name → client
	logger  *slog.Logger
}

// NewManager creates a new MCP manager and initializes all configured servers.
func NewManager(ctx context.Context, servers []config.MCPServerCfg, logger *slog.Logger) (*Manager, error) {
	if logger == nil {
		logger = slog.Default()
	}

	m := &Manager{
		clients: make(map[string]*Client),
		logger:  logger,
	}

	for _, serverCfg := range servers {
		if err := m.addServer(ctx, serverCfg); err != nil {
			// Close any already-created clients before returning error
			m.Close()
			return nil, fmt.Errorf("failed to initialize MCP server %q: %w", serverCfg.Name, err)
		}
	}

	logger.Info("MCP manager initialized", "servers", len(m.clients))
	return m, nil
}

// addServer creates and adds a single MCP client to the manager.
func (m *Manager) addServer(ctx context.Context, cfg config.MCPServerCfg) error {
	if cfg.Name == "" {
		return fmt.Errorf("MCP server must have a name")
	}

	// Auto-detect transport if not specified
	transport := cfg.Transport
	if transport == "" {
		if cfg.Command != "" {
			transport = "stdio"
		} else if cfg.URL != "" {
			transport = "sse"
		} else {
			return fmt.Errorf("MCP server %q must specify either command (stdio) or url (sse)", cfg.Name)
		}
	}

	var client *Client
	var err error

	switch transport {
	case "stdio":
		if cfg.Command == "" {
			return fmt.Errorf("MCP server %q with stdio transport must have a command", cfg.Name)
		}
		// Expand environment variables in command and args
		command := os.ExpandEnv(cfg.Command)
		args := make([]string, len(cfg.Args))
		for i, arg := range cfg.Args {
			args[i] = os.ExpandEnv(arg)
		}
		// Expand env vars in environment map
		env := make(map[string]string, len(cfg.Env))
		for k, v := range cfg.Env {
			env[k] = os.ExpandEnv(v)
		}
		client, err = NewStdioClient(ctx, cfg.Name, command, args, env, m.logger)

	case "sse":
		if cfg.URL == "" {
			return fmt.Errorf("MCP server %q with sse transport must have a url", cfg.Name)
		}
		url := os.ExpandEnv(cfg.URL)
		headers := make(map[string]string, len(cfg.Headers))
		for k, v := range cfg.Headers {
			headers[k] = os.ExpandEnv(v)
		}
		client, err = NewSSEClient(ctx, cfg.Name, url, headers, m.logger)

	default:
		return fmt.Errorf("unknown transport %q for MCP server %q (must be stdio or sse)", transport, cfg.Name)
	}

	if err != nil {
		return err
	}

	m.clients[cfg.Name] = client
	m.logger.Info("MCP server connected", "name", cfg.Name, "transport", transport, "tools", len(client.Tools()))
	return nil
}

// GetClient returns an MCP client by name, or nil if not found.
func (m *Manager) GetClient(name string) *Client {
	return m.clients[name]
}

// AllClients returns all MCP clients managed by this manager.
func (m *Manager) AllClients() map[string]*Client {
	return m.clients
}

// Close closes all MCP client connections.
func (m *Manager) Close() error {
	var errs []string
	for name, client := range m.clients {
		if err := client.Close(); err != nil {
			errs = append(errs, fmt.Sprintf("%s: %v", name, err))
		}
	}

	if len(errs) > 0 {
		return fmt.Errorf("errors closing MCP clients: %s", strings.Join(errs, "; "))
	}

	m.logger.Info("MCP manager closed", "servers", len(m.clients))
	return nil
}