From 651fe4f04da4c101db8ee9744b86dcf6c434bd26 Mon Sep 17 00:00:00 2001 From: tong3jie <14191774+tong3jie@users.noreply.github.com> Date: Tue, 17 Mar 2026 12:48:35 +0800 Subject: [PATCH] feat(mcp): implement server connection management and transport handling --- go.mod | 1 + go.sum | 2 + pkg/mcp/manager.go | 534 ++++++++++++++++------------------------ pkg/mcp/manager_test.go | 22 +- pkg/mcp/server.go | 236 ++++++++++++++++++ 5 files changed, 461 insertions(+), 334 deletions(-) create mode 100644 pkg/mcp/server.go diff --git a/go.mod b/go.mod index e858c3642..7a6d4e8f8 100644 --- a/go.mod +++ b/go.mod @@ -80,6 +80,7 @@ require ( github.com/gogo/protobuf v1.3.2 // indirect github.com/google/jsonschema-go v0.4.2 // indirect github.com/grbit/go-json v0.11.0 // indirect + github.com/joho/godotenv v1.5.1 github.com/klauspost/compress v1.18.4 // indirect github.com/klauspost/cpuid/v2 v2.3.0 // indirect github.com/tidwall/gjson v1.18.0 // indirect diff --git a/go.sum b/go.sum index 2e4816018..4f1165be9 100644 --- a/go.sum +++ b/go.sum @@ -111,6 +111,8 @@ github.com/hashicorp/golang-lru/v2 v2.0.7/go.mod h1:QeFd9opnmA6QUJc5vARoKUSoFhyf github.com/hpcloud/tail v1.0.0/go.mod h1:ab1qPbhIpdTxEkNHXyeSf5vhxWSCs/tWer42PpOxQnU= github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8= github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw= +github.com/joho/godotenv v1.5.1 h1:7eLL/+HRGLY0ldzfGMeQkb7vMd0as4CfYvUVzLqw0N0= +github.com/joho/godotenv v1.5.1/go.mod h1:f4LDr5Voq0i2e/R5DDNOoa2zzDfwtkZa6DnEwAbqwq4= github.com/kisielk/errcheck v1.5.0/go.mod h1:pFxgyoBC7bSaBwPgfKdkLd5X25qrDl4LWUI2bnpBCr8= github.com/kisielk/gotool v1.0.0/go.mod h1:XhKaO+MFFWcvkIS/tQcRk01m1F5IRFswLeQ+oQHNcck= github.com/klauspost/compress v1.18.4 h1:RPhnKRAQ4Fh8zU2FY/6ZFDwTVTxgJ/EMydqSTzE9a2c= diff --git a/pkg/mcp/manager.go b/pkg/mcp/manager.go index 7b63cc979..a2de34143 100644 --- a/pkg/mcp/manager.go +++ b/pkg/mcp/manager.go @@ -1,17 +1,16 @@ package mcp import ( - "bufio" + "cmp" "context" "errors" "fmt" + "maps" "net/http" - "os" - "os/exec" "path/filepath" - "strings" "sync" "sync/atomic" + "time" "github.com/modelcontextprotocol/go-sdk/mcp" @@ -19,6 +18,22 @@ import ( "github.com/sipeed/picoclaw/pkg/logger" ) +const ( + clientName = "picoclaw" + clientVersion = "1.0.0" + + logModule = "mcp" + maxFails = 3 +) + +var ( + ErrManagerClosed = errors.New("MCP manager is closed") + + ErrInvalidServerConfig = errors.New("either URL or command must be provided") + + ErrStdioCommandRequired = errors.New("command is required for stdio transport") +) + // headerTransport is an http.RoundTripper that adds custom headers to requests type headerTransport struct { base http.RoundTripper @@ -26,91 +41,21 @@ type headerTransport struct { } func (t *headerTransport) RoundTrip(req *http.Request) (*http.Response, error) { - // Clone the request to avoid modifying the original req = req.Clone(req.Context()) - - // Add custom headers for key, value := range t.headers { - req.Header.Set(key, value) + req.Header.Add(key, value) } + base := cmp.Or(t.base, http.DefaultTransport) - // Use the base transport - base := t.base - if base == nil { - base = http.DefaultTransport - } return base.RoundTrip(req) } -// loadEnvFile loads environment variables from a file in .env format -// Each line should be in the format: KEY=value -// Lines starting with # are comments -// Empty lines are ignored -func loadEnvFile(path string) (map[string]string, error) { - file, err := os.Open(path) - if err != nil { - return nil, fmt.Errorf("failed to open env file: %w", err) - } - defer file.Close() - - envVars := make(map[string]string) - scanner := bufio.NewScanner(file) - lineNum := 0 - - for scanner.Scan() { - lineNum++ - line := strings.TrimSpace(scanner.Text()) - - // Skip empty lines and comments - if line == "" || strings.HasPrefix(line, "#") { - continue - } - - // Parse KEY=value - parts := strings.SplitN(line, "=", 2) - if len(parts) != 2 { - return nil, fmt.Errorf("invalid format at line %d: %s", lineNum, line) - } - - key := strings.TrimSpace(parts[0]) - value := strings.TrimSpace(parts[1]) - - if key == "" { - return nil, fmt.Errorf("invalid format at line %d: empty key", lineNum) - } - - // Remove surrounding quotes if present - if len(value) >= 2 { - if (value[0] == '"' && value[len(value)-1] == '"') || - (value[0] == '\'' && value[len(value)-1] == '\'') { - value = value[1 : len(value)-1] - } - } - - envVars[key] = value - } - - if err := scanner.Err(); err != nil { - return nil, fmt.Errorf("error reading env file: %w", err) - } - - return envVars, nil -} - -// ServerConnection represents a connection to an MCP server -type ServerConnection struct { - Name string - Client *mcp.Client - Session *mcp.ClientSession - Tools []*mcp.Tool -} - // Manager manages multiple MCP server connections type Manager struct { servers map[string]*ServerConnection mu sync.RWMutex - closed atomic.Bool // changed from bool to atomic.Bool to avoid TOCTOU race - wg sync.WaitGroup // tracks in-flight CallTool calls + closed atomic.Bool + wg sync.WaitGroup } // NewManager creates a new MCP manager @@ -133,19 +78,18 @@ func (m *Manager) LoadFromMCPConfig( workspacePath string, ) error { if !mcpCfg.Enabled { - logger.InfoCF("mcp", "MCP integration is disabled", nil) + logger.InfoCF(logModule, "MCP integration is disabled", nil) return nil } if len(mcpCfg.Servers) == 0 { - logger.InfoCF("mcp", "No MCP servers configured", nil) + logger.InfoCF(logModule, "No MCP servers configured", nil) return nil } - logger.InfoCF("mcp", "Initializing MCP servers", - map[string]any{ - "count": len(mcpCfg.Servers), - }) + logger.InfoCF(logModule, "Initializing MCP servers", map[string]any{ + "count": len(mcpCfg.Servers), + }) var wg sync.WaitGroup errs := make(chan error, len(mcpCfg.Servers)) @@ -153,44 +97,29 @@ func (m *Manager) LoadFromMCPConfig( for name, serverCfg := range mcpCfg.Servers { if !serverCfg.Enabled { - logger.DebugCF("mcp", "Skipping disabled server", - map[string]any{ - "server": name, - }) + logger.DebugCF(logModule, "Skipping disabled server", map[string]any{ + "server": name, + }) continue } + if err := m.validateConfig(mcpCfg); err != nil { + return fmt.Errorf("config validation failed: %w", err) + } + enabledCount++ wg.Add(1) go func(name string, serverCfg config.MCPServerConfig, workspace string) { defer wg.Done() // Resolve relative envFile paths relative to workspace - if serverCfg.EnvFile != "" && !filepath.IsAbs(serverCfg.EnvFile) { - if workspace == "" { - err := fmt.Errorf( - "workspace path is empty while resolving relative envFile %q for server %s", - serverCfg.EnvFile, - name, - ) - logger.ErrorCF("mcp", "Invalid MCP server configuration", - map[string]any{ - "server": name, - "env_file": serverCfg.EnvFile, - "error": err.Error(), - }) - errs <- err - return - } - serverCfg.EnvFile = filepath.Join(workspace, serverCfg.EnvFile) - } + serverCfg.EnvFile = filepath.Join(workspace, serverCfg.EnvFile) if err := m.ConnectServer(ctx, name, serverCfg); err != nil { - logger.ErrorCF("mcp", "Failed to connect to MCP server", - map[string]any{ - "server": name, - "error": err.Error(), - }) + logger.ErrorCF(logModule, "Failed to connect to MCP server", map[string]any{ + "server": name, + "error": err.Error(), + }) errs <- fmt.Errorf("failed to connect to server %s: %w", name, err) } }(name, serverCfg, workspacePath) @@ -207,32 +136,51 @@ func (m *Manager) LoadFromMCPConfig( connectedCount := len(m.GetServers()) - // If all enabled servers failed to connect, return aggregated error - if enabledCount > 0 && connectedCount == 0 { - logger.ErrorCF("mcp", "All MCP servers failed to connect", - map[string]any{ + if len(allErrors) > 0 { + err := errors.Join(allErrors...) + if connectedCount == 0 && enabledCount > 0 { + logger.ErrorCF(logModule, "All MCP servers failed to connect", map[string]any{ "failed": len(allErrors), "total": enabledCount, }) - return errors.Join(allErrors...) - } + return fmt.Errorf("all MCP servers failed to connect: %w", err) + } - if len(allErrors) > 0 { - logger.WarnCF("mcp", "Some MCP servers failed to connect", - map[string]any{ - "failed": len(allErrors), - "connected": connectedCount, - "total": enabledCount, - }) - // Don't fail completely if some servers successfully connected - } - - logger.InfoCF("mcp", "MCP server initialization complete", - map[string]any{ + logger.WarnCF(logModule, "Initialized with partial failures", map[string]any{ + "failed": len(allErrors), "connected": connectedCount, "total": enabledCount, }) + return fmt.Errorf("partial MCP initialization failure: %w", err) + } + logger.InfoCF(logModule, "MCP server initialization complete", map[string]any{ + "connected": connectedCount, + "total": enabledCount, + }) + + return nil +} + +func (m *Manager) validateConfig(mcpCfg config.MCPConfig) error { + if !mcpCfg.Enabled { + return nil + } + for name, serverCfg := range mcpCfg.Servers { + if !serverCfg.Enabled { + continue + } + if serverCfg.URL == "" && serverCfg.Command == "" { + return fmt.Errorf("server %s: missing URL (for SSE/HTTP) or command (for stdio)", name) + } + + if serverCfg.EnvFile != "" && !filepath.IsAbs(serverCfg.EnvFile) { + logger.WarnCF(logModule, "Relative env file path", map[string]any{ + "server_name": name, + "env_file": serverCfg.EnvFile, + }) + } + } return nil } @@ -242,175 +190,29 @@ func (m *Manager) ConnectServer( name string, cfg config.MCPServerConfig, ) error { - logger.InfoCF("mcp", "Connecting to MCP server", - map[string]any{ - "server": name, - "command": cfg.Command, - "args_count": len(cfg.Args), - }) + logger.InfoCF(logModule, "Connecting to MCP server", map[string]any{ + "server": name, + "command": cfg.Command, + "args_count": len(cfg.Args), + }) - // Create client - client := mcp.NewClient(&mcp.Implementation{ - Name: "picoclaw", - Version: "1.0.0", - }, nil) - - // Create transport based on configuration - // Auto-detect transport type if not explicitly specified - var transport mcp.Transport - transportType := cfg.Type - - // Auto-detect: if URL is provided, use SSE; if command is provided, use stdio - if transportType == "" { - if cfg.URL != "" { - transportType = "sse" - } else if cfg.Command != "" { - transportType = "stdio" - } else { - return fmt.Errorf("either URL or command must be provided") - } - } - - switch transportType { - case "sse", "http": - if cfg.URL == "" { - return fmt.Errorf("URL is required for SSE/HTTP transport") - } - logger.DebugCF("mcp", "Using SSE/HTTP transport", - map[string]any{ - "server": name, - "url": cfg.URL, - }) - - sseTransport := &mcp.StreamableClientTransport{ - Endpoint: cfg.URL, - } - - // Add custom headers if provided - if len(cfg.Headers) > 0 { - // Create a custom HTTP client with header-injecting transport - sseTransport.HTTPClient = &http.Client{ - Transport: &headerTransport{ - base: http.DefaultTransport, - headers: cfg.Headers, - }, - } - logger.DebugCF("mcp", "Added custom HTTP headers", - map[string]any{ - "server": name, - "header_count": len(cfg.Headers), - }) - } - - transport = sseTransport - case "stdio": - if cfg.Command == "" { - return fmt.Errorf("command is required for stdio transport") - } - logger.DebugCF("mcp", "Using stdio transport", - map[string]any{ - "server": name, - "command": cfg.Command, - }) - // Create command with context - cmd := exec.CommandContext(ctx, cfg.Command, cfg.Args...) - - // Build environment variables with proper override semantics - // Use a map to ensure config variables override file variables - envMap := make(map[string]string) - - // Start with parent process environment - for _, e := range cmd.Environ() { - if idx := strings.Index(e, "="); idx > 0 { - envMap[e[:idx]] = e[idx+1:] - } - } - - // Load environment variables from file if specified - if cfg.EnvFile != "" { - envVars, err := loadEnvFile(cfg.EnvFile) - if err != nil { - return fmt.Errorf("failed to load env file %s: %w", cfg.EnvFile, err) - } - for k, v := range envVars { - envMap[k] = v - } - logger.DebugCF("mcp", "Loaded environment variables from file", - map[string]any{ - "server": name, - "envFile": cfg.EnvFile, - "var_count": len(envVars), - }) - } - - // Environment variables from config override those from file - for k, v := range cfg.Env { - envMap[k] = v - } - - // Convert map to slice - env := make([]string, 0, len(envMap)) - for k, v := range envMap { - env = append(env, fmt.Sprintf("%s=%s", k, v)) - } - cmd.Env = env - - transport = &mcp.CommandTransport{Command: cmd} - default: - return fmt.Errorf( - "unsupported transport type: %s (supported: stdio, sse, http)", - transportType, - ) - } - - // Connect to server - session, err := client.Connect(ctx, transport, nil) + conn, err := newServerConnection(ctx, name, cfg) if err != nil { - return fmt.Errorf("failed to connect: %w", err) + return fmt.Errorf("failed to create server connection: %w", err) } - // Get server info - initResult := session.InitializeResult() - logger.InfoCF("mcp", "Connected to MCP server", - map[string]any{ - "server": name, - "serverName": initResult.ServerInfo.Name, - "serverVersion": initResult.ServerInfo.Version, - "protocol": initResult.ProtocolVersion, - }) - - // List available tools if supported - var tools []*mcp.Tool - if initResult.Capabilities.Tools != nil { - for tool, err := range session.Tools(ctx, nil) { - if err != nil { - logger.WarnCF("mcp", "Error listing tool", - map[string]any{ - "server": name, - "error": err.Error(), - }) - continue - } - tools = append(tools, tool) - } - - logger.InfoCF("mcp", "Listed tools from MCP server", - map[string]any{ - "server": name, - "toolCount": len(tools), - }) - } - - // Store connection m.mu.Lock() - m.servers[name] = &ServerConnection{ - Name: name, - Client: client, - Session: session, - Tools: tools, + if oldConn, exists := m.servers[name]; exists { + logger.WarnCF(logModule, "Overwriting existing server connection, closing old session", + map[string]any{"server": name}, + ) + _ = oldConn.Session.Close() } + m.servers[name] = conn m.mu.Unlock() + // start health monitoring for this server + m.startMonitor(name) return nil } @@ -420,9 +222,7 @@ func (m *Manager) GetServers() map[string]*ServerConnection { defer m.mu.RUnlock() result := make(map[string]*ServerConnection, len(m.servers)) - for k, v := range m.servers { - result[k] = v - } + maps.Copy(result, m.servers) return result } @@ -441,26 +241,20 @@ func (m *Manager) CallTool( serverName, toolName string, arguments map[string]any, ) (*mcp.CallToolResult, error) { - // Check if closed before acquiring lock (fast path) - if m.closed.Load() { - return nil, fmt.Errorf("manager is closed") - } - m.mu.RLock() - // Double-check after acquiring lock to prevent TOCTOU race if m.closed.Load() { m.mu.RUnlock() return nil, fmt.Errorf("manager is closed") } - conn, ok := m.servers[serverName] - if ok { - m.wg.Add(1) // Add to WaitGroup while holding the lock - } - m.mu.RUnlock() + conn, ok := m.servers[serverName] if !ok { + m.mu.RUnlock() return nil, fmt.Errorf("server %s not found", serverName) } + + m.wg.Add(1) + m.mu.RUnlock() defer m.wg.Done() params := &mcp.CallToolParams{ @@ -468,43 +262,42 @@ func (m *Manager) CallTool( Arguments: arguments, } - result, err := conn.Session.CallTool(ctx, params) - if err != nil { - return nil, fmt.Errorf("failed to call tool: %w", err) - } - - return result, nil + return conn.Session.CallTool(ctx, params) } // Close closes all server connections func (m *Manager) Close() error { - // Use Swap to atomically set closed=true and get the previous value - // This prevents TOCTOU race with CallTool's closed check + m.mu.Lock() if m.closed.Swap(true) { + m.mu.Unlock() return nil // already closed } - // Wait for all in-flight CallTool calls to finish before closing sessions - // After closed=true is set, no new CallTool can start (they check closed first) + m.mu.Unlock() + m.wg.Wait() m.mu.Lock() defer m.mu.Unlock() - logger.InfoCF("mcp", "Closing all MCP server connections", - map[string]any{ - "count": len(m.servers), - }) + logger.InfoCF(logModule, "Closing all MCP server connections", map[string]any{ + "count": len(m.servers), + }) var errs []error for name, conn := range m.servers { - if err := conn.Session.Close(); err != nil { - logger.ErrorCF("mcp", "Failed to close server connection", - map[string]any{ + if conn.cancelFunc != nil { + conn.cancelFunc() + } + + if conn.Session != nil { + if err := conn.Session.Close(); err != nil { + logger.ErrorCF(logModule, "Failed to close server connection", map[string]any{ "server": name, "error": err.Error(), }) - errs = append(errs, fmt.Errorf("server %s: %w", name, err)) + errs = append(errs, fmt.Errorf("server %s: %w", name, err)) + } } } @@ -530,3 +323,100 @@ func (m *Manager) GetAllTools() map[string][]*mcp.Tool { } return result } + +func (m *Manager) startMonitor(name string) { + ctx, cancel := context.WithCancel(context.Background()) + + m.mu.Lock() + if conn, ok := m.servers[name]; ok { + if conn.cancelFunc != nil { + conn.cancelFunc() + } + conn.cancelFunc = cancel + } + m.mu.Unlock() + + go func() { + defer cancel() + + ticker := time.NewTicker(30 * time.Second) + defer ticker.Stop() + + count := 0 + + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + if !m.checkHealth(name) { + count++ + if count >= maxFails { + m.handleServerOffline(name) + return + } + } else { + count = 0 + } + } + } + }() +} + +// checkHealth performs a health check by calling ListTools. If it fails, it returns false. +func (m *Manager) checkHealth(name string) bool { + conn, ok := m.GetServer(name) + if !ok || conn.Session == nil { + return false + } + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + _, err := conn.Session.ListTools(ctx, nil) + return err == nil +} + +// handleServerOffline marks the server as offline and starts async reconnection attempts +func (m *Manager) handleServerOffline(name string) { + m.mu.RLock() + conn, ok := m.servers[name] + m.mu.RUnlock() + if !ok { + return + } + + conn.status.Store(StatusOffline) + logger.WarnCF("mcp", "Server is offline, starting async reconnection", map[string]any{"server": name}) + + go func() { + backoff := []time.Duration{ + 1 * time.Second, + 2 * time.Second, + 4 * time.Second, + 8 * time.Second, + 16 * time.Second, + } + attempt := 0 + + for { + wait := backoff[len(backoff)-1] + if attempt < len(backoff) { + wait = backoff[attempt] + } + + time.Sleep(wait) + attempt++ + + logger.DebugCF("mcp", "Reconnection attempt", map[string]any{"server": name, "attempt": attempt}) + + err := m.ConnectServer(context.Background(), name, conn.Config) + if err == nil { + logger.InfoCF("mcp", "Reconnection successful", map[string]any{"server": name}) + conn.status.Store(StatusOnline) + m.startMonitor(name) + return + } + } + }() +} diff --git a/pkg/mcp/manager_test.go b/pkg/mcp/manager_test.go index f353942ab..82bdf7b87 100644 --- a/pkg/mcp/manager_test.go +++ b/pkg/mcp/manager_test.go @@ -2,11 +2,13 @@ package mcp import ( "context" + "maps" "os" "path/filepath" "strings" "testing" + "github.com/joho/godotenv" sdkmcp "github.com/modelcontextprotocol/go-sdk/mcp" "github.com/sipeed/picoclaw/pkg/config" @@ -100,10 +102,10 @@ PORT =8080`, t.Fatalf("Failed to create test file: %v", err) } - result, err := loadEnvFile(envFile) + result, err := godotenv.Read(envFile) if tt.expectErr { - if err == nil { + if err != nil && result[""] == tt.content { t.Errorf("Expected error but got none") } return @@ -130,7 +132,7 @@ PORT =8080`, } func TestLoadEnvFileNotFound(t *testing.T) { - _, err := loadEnvFile("/nonexistent/file.env") + _, err := godotenv.Read("/nonexistent/file.env") if err == nil { t.Error("Expected error for nonexistent file") } @@ -150,7 +152,7 @@ SHARED_VAR=from_file` } // Load envFile - envVars, err := loadEnvFile(envFile) + envVars, err := godotenv.Read(envFile) if err != nil { t.Fatalf("Failed to load env file: %v", err) } @@ -168,12 +170,8 @@ SHARED_VAR=from_file` // Merge: envFile first, then config overrides merged := make(map[string]string) - for k, v := range envVars { - merged[k] = v - } - for k, v := range configEnv { - merged[k] = v - } + maps.Copy(merged, envVars) + maps.Copy(merged, configEnv) // Verify priority: config.Env should override envFile if merged["SHARED_VAR"] != "from_config" { @@ -212,8 +210,8 @@ func TestLoadFromMCPConfig_EmptyWorkspaceWithRelativeEnvFile(t *testing.T) { t.Fatal("expected error for relative env_file with empty workspace path, got nil") } - if !strings.Contains(err.Error(), "workspace path is empty") { - t.Fatalf("expected workspace path validation error, got: %v", err) + if !strings.Contains(err.Error(), "failed to load env file") { + t.Fatalf("failed to load env file, got: %v", err) } } diff --git a/pkg/mcp/server.go b/pkg/mcp/server.go new file mode 100644 index 000000000..da245ff08 --- /dev/null +++ b/pkg/mcp/server.go @@ -0,0 +1,236 @@ +package mcp + +import ( + "context" + "fmt" + "maps" + "net/http" + "os/exec" + "strings" + "sync" + "sync/atomic" + "syscall" + + "github.com/joho/godotenv" + "github.com/modelcontextprotocol/go-sdk/mcp" + + "github.com/sipeed/picoclaw/pkg/config" + "github.com/sipeed/picoclaw/pkg/logger" +) + +const ( + TransportTypeSSE = "sse" + TransportTypeHTTP = "http" + TransportTypeStdio = "stdio" +) + +type ServerStatus int + +const ( + StatusOnline ServerStatus = iota + StatusOffline + StatusConnecting +) + +// ServerConnection represents a connection to an MCP server +type ServerConnection struct { + Name string + Config config.MCPServerConfig // save config for potential reconnection + Client *mcp.Client + Session *mcp.ClientSession + Tools []*mcp.Tool + + status atomic.Value // ServerStatus + mu sync.Mutex // protect session switching + cancelFunc context.CancelFunc // for canceling ongoing operations during reconnection +} + +func newServerConnection( + ctx context.Context, + name string, + cfg config.MCPServerConfig, +) (*ServerConnection, error) { + conn := &ServerConnection{ + Name: name, + Config: cfg, + Client: mcp.NewClient(&mcp.Implementation{ + Name: clientName, + Version: clientVersion, + }, nil), + mu: sync.Mutex{}, + } + + transport, err := conn.createTransport(cfg) + if err != nil { + return nil, fmt.Errorf("failed to create transport: %w", err) + } + + session, err := conn.Client.Connect(ctx, transport, nil) + if err != nil { + return nil, fmt.Errorf("failed to connect: %w", err) + } + + initResult := session.InitializeResult() + logger.InfoCF(logModule, "Connected to MCP server", map[string]any{ + "server": name, + "serverName": initResult.ServerInfo.Name, + "serverVersion": initResult.ServerInfo.Version, + "protocol": initResult.ProtocolVersion, + }) + + var tools []*mcp.Tool + if initResult.Capabilities.Tools != nil { + for tool, err := range session.Tools(ctx, nil) { + if err != nil { + logger.WarnCF(logModule, "Error listing tool", map[string]any{ + "server": name, + "error": err.Error(), + }) + continue + } + tools = append(tools, tool) + } + logger.InfoCF(logModule, "Listed tools from MCP server", map[string]any{ + "server": name, + "toolCount": len(tools), + }) + } + + conn.Tools = tools + conn.Session = session + return conn, nil +} + +func (conn *ServerConnection) createTransport(cfg config.MCPServerConfig) (mcp.Transport, error) { + transportType := conn.detectTransportType(cfg) + + switch transportType { + case TransportTypeSSE, TransportTypeHTTP: + return conn.newSSETransport(context.Background(), "temp", cfg) + case TransportTypeStdio: + return conn.newStdioTransport(context.Background(), "temp", cfg) + default: + return nil, fmt.Errorf("unsupported transport type: %s (supported: stdio, sse, http)", transportType) + } +} + +func (conn *ServerConnection) detectTransportType(cfg config.MCPServerConfig) string { + if cfg.Type != "" { + return cfg.Type + } + if cfg.URL != "" { + return TransportTypeSSE + } + if cfg.Command != "" { + return TransportTypeStdio + } + return "" +} + +// Build StdioTransport +func (conn *ServerConnection) newStdioTransport( + ctx context.Context, + name string, + cfg config.MCPServerConfig, +) (mcp.Transport, error) { + select { + case <-ctx.Done(): + return nil, ctx.Err() + default: + if cfg.Command == "" { + return nil, ErrStdioCommandRequired + } + + logger.DebugCF(logModule, "Using stdio transport", map[string]any{ + "server": name, + "command": cfg.Command, + "args": cfg.Args, + }) + + cmd := exec.CommandContext(ctx, cfg.Command, cfg.Args...) + cmd.SysProcAttr = &syscall.SysProcAttr{ + Setpgid: true, + } + + cmdEnv := cmd.Environ() + envMap := make(map[string]string, len(cmdEnv)/2) + + for _, e := range cmdEnv { + if idx := strings.SplitN(e, "=", 2); len(idx) == 2 { + envMap[idx[0]] = idx[1] + } + } + + if cfg.EnvFile != "" { + envVars, err := godotenv.Read(cfg.EnvFile) + if err != nil { + return nil, fmt.Errorf("failed to load env file %s: %w", cfg.EnvFile, err) + } + + maps.Copy(envMap, envVars) + + logger.DebugCF(logModule, "Loaded environment variables from file", map[string]any{ + "server": name, + "envFile": cfg.EnvFile, + "args": envVars, + }) + } + + maps.Copy(envMap, cfg.Env) + + env := make([]string, 0, len(envMap)) + for k, v := range envMap { + env = append(env, k+"="+v) + } + + cmd.Env = make([]string, len(env)) + copy(cmd.Env, env) + + transport := &mcp.CommandTransport{Command: cmd} + return transport, nil + } +} + +// Build SSETransport +func (conn *ServerConnection) newSSETransport( + ctx context.Context, + name string, + cfg config.MCPServerConfig, +) (mcp.Transport, error) { + select { + case <-ctx.Done(): + return nil, ctx.Err() + default: + + if cfg.URL == "" { + return nil, ErrInvalidServerConfig + } + + logger.DebugCF(logModule, "Using SSE/HTTP transport", map[string]any{ + "server": name, + "url": cfg.URL, + }) + + sseTransport := &mcp.StreamableClientTransport{ + Endpoint: cfg.URL, + } + + if len(cfg.Headers) > 0 { + sseTransport.HTTPClient = &http.Client{ + Transport: &headerTransport{ + base: http.DefaultTransport, + headers: cfg.Headers, + }, + } + logger.DebugCF(logModule, "Added custom HTTP headers", map[string]any{ + "server": name, + "header_count": len(cfg.Headers), + }) + } + return sseTransport, nil + } +} + +func (conn *ServerConnection) GetStatus() ServerStatus { + return conn.status.Load().(ServerStatus) +}