diff --git a/internal/agent/loop_mcp_user.go b/internal/agent/loop_mcp_user.go index ce6ba393..11a536f1 100644 --- a/internal/agent/loop_mcp_user.go +++ b/internal/agent/loop_mcp_user.go @@ -86,7 +86,7 @@ func (l *Loop) getUserMCPTools(ctx context.Context, userID string) []tools.Tool // shared tool registry so ExecuteWithContext can resolve them by name. reg, _ := l.tools.(*tools.Registry) for _, mcpTool := range entry.MCPTools() { - bt := mcpbridge.NewBridgeTool(srv.Name, mcpTool, entry.Client(), srv.ToolPrefix, srv.TimeoutSec, entry.Connected()) + bt := mcpbridge.NewBridgeTool(srv.Name, mcpTool, entry.ClientPtr(), srv.ToolPrefix, srv.TimeoutSec, entry.Connected()) // Register in registry so ExecuteWithContext can find them. // Skip if already registered (another user loaded this server with same tool names). if reg != nil { diff --git a/internal/mcp/bridge_tool.go b/internal/mcp/bridge_tool.go index 5d3215cc..69026a86 100644 --- a/internal/mcp/bridge_tool.go +++ b/internal/mcp/bridge_tool.go @@ -15,6 +15,8 @@ import ( // BridgeTool adapts an MCP tool into the tools.Tool interface. // It delegates Execute calls to the MCP server via the client. +// The client pointer is loaded atomically from clientPtr to support +// safe reconnection without data races. type BridgeTool struct { serverName string toolName string // original MCP tool name @@ -22,7 +24,7 @@ type BridgeTool struct { description string inputSchema map[string]any // JSON Schema for parameters requiredSet map[string]bool - client *mcpclient.Client + clientPtr *atomic.Pointer[mcpclient.Client] // shared with serverState for atomic swap on reconnect timeoutSec int connected *atomic.Bool } @@ -30,7 +32,9 @@ type BridgeTool struct { // NewBridgeTool creates a BridgeTool from an MCP Tool definition. // The tool name is always prefixed with "mcp_" to distinguish MCP tools from native tools. // If prefix is empty, it is auto-derived from the server name. -func NewBridgeTool(serverName string, mcpTool mcpgo.Tool, client *mcpclient.Client, prefix string, timeoutSec int, connected *atomic.Bool) *BridgeTool { +// clientPtr is a shared atomic pointer from serverState — reconnection swaps it +// atomically, and all BridgeTools see the new client without explicit notification. +func NewBridgeTool(serverName string, mcpTool mcpgo.Tool, clientPtr *atomic.Pointer[mcpclient.Client], prefix string, timeoutSec int, connected *atomic.Bool) *BridgeTool { name := mcpTool.Name effectivePrefix := ensureMCPPrefix(prefix, serverName) registered := effectivePrefix + "__" + name @@ -53,7 +57,7 @@ func NewBridgeTool(serverName string, mcpTool mcpgo.Tool, client *mcpclient.Clie description: mcpTool.Description, inputSchema: schema, requiredSet: reqSet, - client: client, + clientPtr: clientPtr, timeoutSec: timeoutSec, connected: connected, } @@ -99,6 +103,11 @@ func (t *BridgeTool) Execute(ctx context.Context, args map[string]any) *tools.Re return tools.ErrorResult(fmt.Sprintf("MCP server %q is disconnected", t.serverName)) } + client := t.clientPtr.Load() // atomic load — safe during concurrent reconnect + if client == nil { + return tools.ErrorResult(fmt.Sprintf("MCP server %q has no active client", t.serverName)) + } + callCtx, cancel := context.WithTimeout(ctx, time.Duration(t.timeoutSec)*time.Second) defer cancel() @@ -111,7 +120,7 @@ func (t *BridgeTool) Execute(ctx context.Context, args map[string]any) *tools.Re req.Params.Name = t.toolName req.Params.Arguments = cleanedArgs - result, err := t.client.CallTool(callCtx, req) + result, err := client.CallTool(callCtx, req) if err != nil { if errors.Is(callCtx.Err(), context.DeadlineExceeded) { return tools.ErrorResult(fmt.Sprintf("MCP tool %q timeout after %ds", t.registeredName, t.timeoutSec)) diff --git a/internal/mcp/manager.go b/internal/mcp/manager.go index c5ebab06..df63d1d4 100644 --- a/internal/mcp/manager.go +++ b/internal/mcp/manager.go @@ -25,6 +25,7 @@ const ( initialBackoff = 2 * time.Second maxBackoff = 60 * time.Second maxReconnectAttempts = 10 + reconnectCooldown = 5 * time.Minute // wait after exhausting reconnect attempts before retrying // mcpToolInlineMaxCount is the threshold above which MCP tools switch // to search mode (deferred loading via mcp_tool_search) instead of @@ -41,15 +42,35 @@ type ServerStatus struct { Error string `json:"error,omitempty"` } +// connParams stores connection parameters needed to re-establish a dead connection. +// Populated during initial connectAndDiscover and used by tryReconnect. +type connParams struct { + command string + args []string + env map[string]string + url string + headers map[string]string +} + // serverState tracks a single MCP server connection. +// +// Dual-pointer design for the MCP client: +// - client: direct pointer used by healthLoop (single goroutine, no contention). +// - clientPtr: atomic pointer shared with all BridgeTools via NewBridgeTool. +// BridgeTools call clientPtr.Load() in Execute for race-safe access. +// +// On reconnect, fullReconnect() updates BOTH: ss.client for healthLoop and +// ss.clientPtr.Store() for BridgeTools. The old client is closed AFTER the swap. type serverState struct { name string transport string - client *mcpclient.Client + client *mcpclient.Client // direct ref for health checks (single-goroutine access) + clientPtr atomic.Pointer[mcpclient.Client] // shared atomic ref for BridgeTools (multi-goroutine safe) connected atomic.Bool toolNames []string // registered tool names in the registry timeoutSec int cancel context.CancelFunc + conn connParams // connection params for reconnect mu sync.Mutex reconnAttempts int diff --git a/internal/mcp/manager_connect.go b/internal/mcp/manager_connect.go index 9bb6972a..a71c3ecb 100644 --- a/internal/mcp/manager_connect.go +++ b/internal/mcp/manager_connect.go @@ -86,7 +86,15 @@ func connectAndDiscover(ctx context.Context, name, transportType, command string transport: transportType, client: client, timeoutSec: timeoutSec, + conn: connParams{ + command: command, + args: args, + env: env, + url: url, + headers: headers, + }, } + ss.clientPtr.Store(client) ss.connected.Store(true) return ss, toolsResult.Tools, nil @@ -133,7 +141,7 @@ func (m *Manager) connectServer(ctx context.Context, name, transportType, comman func (m *Manager) registerBridgeTools(ss *serverState, mcpTools []mcpgo.Tool, serverName, toolPrefix string, timeoutSec int) []string { var registeredNames []string for _, mcpTool := range mcpTools { - bt := NewBridgeTool(serverName, mcpTool, ss.client, toolPrefix, timeoutSec, &ss.connected) + bt := NewBridgeTool(serverName, mcpTool, &ss.clientPtr, toolPrefix, timeoutSec, &ss.connected) if _, exists := m.registry.Get(bt.Name()); exists { slog.Warn("mcp.tool.name_collision", @@ -199,7 +207,7 @@ func (m *Manager) connectViaPool(ctx context.Context, tenantID uuid.UUID, name, func (m *Manager) registerPoolBridgeTools(entry *poolEntry, serverName, toolPrefix string, timeoutSec int) []string { var registeredNames []string for _, mcpTool := range entry.tools { - bt := NewBridgeTool(serverName, mcpTool, entry.state.client, toolPrefix, timeoutSec, &entry.state.connected) + bt := NewBridgeTool(serverName, mcpTool, &entry.state.clientPtr, toolPrefix, timeoutSec, &entry.state.connected) if _, exists := m.registry.Get(bt.Name()); exists { slog.Warn("mcp.tool.name_collision", @@ -301,13 +309,26 @@ func (m *Manager) healthLoop(ctx context.Context, ss *serverState) { } // tryReconnect attempts to reconnect with exponential backoff. +// For SSE/HTTP transports, a server-side restart invalidates the session, +// so pinging the old client will never succeed. In that case, we create a +// fresh client and atomically swap it via serverState.clientPtr — all +// BridgeTools see the new client on their next Execute call. func (m *Manager) tryReconnect(ctx context.Context, ss *serverState) { ss.mu.Lock() if ss.reconnAttempts >= maxReconnectAttempts { - ss.lastErr = fmt.Sprintf("max reconnect attempts (%d) reached", maxReconnectAttempts) + ss.lastErr = fmt.Sprintf("max reconnect attempts (%d) reached, entering cooldown", maxReconnectAttempts) ss.mu.Unlock() - slog.Error("mcp.server.reconnect_exhausted", "server", ss.name) - return + slog.Warn("mcp.server.reconnect_cooldown", "server", ss.name, "cooldown", reconnectCooldown) + // Cooldown: wait before resetting attempts so the server has time to recover. + select { + case <-ctx.Done(): + return + case <-time.After(reconnectCooldown): + } + ss.mu.Lock() + ss.reconnAttempts = 0 + ss.mu.Unlock() + return // will retry on next health tick } ss.reconnAttempts++ attempt := ss.reconnAttempts @@ -327,13 +348,71 @@ func (m *Manager) tryReconnect(ctx context.Context, ss *serverState) { case <-time.After(backoff): } - // Try to ping again — transport may have auto-reconnected + // Fast path: ping existing client — works for transient network blips + // where the server-side session is still alive. if err := ss.client.Ping(ctx); err == nil { ss.connected.Store(true) ss.mu.Lock() ss.reconnAttempts = 0 + ss.healthFailures = 0 ss.lastErr = "" ss.mu.Unlock() slog.Info("mcp.server.reconnected", "server", ss.name) + return + } + + // Slow path: server-side session is dead (container restart, OOM, etc.). + if fullReconnect(ctx, ss) { + slog.Info("mcp.server.reconnected", "server", ss.name, "method", "full_reconnect") } } + +// fullReconnect creates a fresh MCP client, atomically swaps it into serverState, +// and closes the old one. Returns true on success. Shared by Manager.tryReconnect +// and poolTryReconnect to avoid duplication. +// +// The new client is created and validated FIRST, then swapped via clientPtr.Store() +// so BridgeTools see the new client immediately. The old client is closed AFTER +// the swap to avoid a window where ss.client points to a closed client. +func fullReconnect(ctx context.Context, ss *serverState) bool { + slog.Info("mcp.full_reconnect", "server", ss.name, "transport", ss.transport) + + newClient, err := createClient(ss.transport, ss.conn.command, ss.conn.args, ss.conn.env, ss.conn.url, ss.conn.headers) + if err != nil { + slog.Warn("mcp.reconnect_create_failed", "server", ss.name, "error", err) + return false + } + + if ss.transport != "stdio" { + if err := newClient.Start(ctx); err != nil { + _ = newClient.Close() + slog.Warn("mcp.reconnect_start_failed", "server", ss.name, "error", err) + return false + } + } + + initReq := mcpgo.InitializeRequest{} + initReq.Params.ProtocolVersion = mcpgo.LATEST_PROTOCOL_VERSION + initReq.Params.ClientInfo = mcpgo.Implementation{Name: "goclaw", Version: "1.0.0"} + + if _, err := newClient.Initialize(ctx, initReq); err != nil { + _ = newClient.Close() + slog.Warn("mcp.reconnect_init_failed", "server", ss.name, "error", err) + return false + } + + // Swap atomically: store new client, then close old. + // BridgeTools use clientPtr.Load() so they see the new client immediately. + oldClient := ss.client + ss.client = newClient + ss.clientPtr.Store(newClient) + ss.connected.Store(true) + ss.mu.Lock() + ss.reconnAttempts = 0 + ss.healthFailures = 0 + ss.lastErr = "" + ss.mu.Unlock() + + _ = oldClient.Close() + return true +} diff --git a/internal/mcp/pool.go b/internal/mcp/pool.go index 2d19a332..df86675b 100644 --- a/internal/mcp/pool.go +++ b/internal/mcp/pool.go @@ -591,6 +591,10 @@ func (p *Pool) evictOldestIdleLocked() bool { // Client returns the MCP client for this pool entry. func (e *poolEntry) Client() *mcpclient.Client { return e.state.client } +// ClientPtr returns the atomic client pointer for this pool entry. +// Used by BridgeTools to atomically load the current client during reconnect. +func (e *poolEntry) ClientPtr() *atomic.Pointer[mcpclient.Client] { return &e.state.clientPtr } + // Connected returns a pointer to the connected flag for this pool entry. func (e *poolEntry) Connected() *atomic.Bool { return &e.state.connected } @@ -598,6 +602,8 @@ func (e *poolEntry) Connected() *atomic.Bool { return &e.state.connected } func (e *poolEntry) MCPTools() []mcpgo.Tool { return e.tools } // poolHealthLoop is a standalone health loop for pool-managed connections. +// After consecutive ping failures, it attempts a full reconnect by creating +// a fresh client, mirroring the Manager.tryReconnect slow path. func poolHealthLoop(ctx context.Context, ss *serverState) { ticker := newHealthTicker() defer ticker.Stop() @@ -625,6 +631,7 @@ func poolHealthLoop(ctx context.Context, ss *serverState) { if failures >= healthFailThreshold { ss.connected.Store(false) + poolTryReconnect(ctx, ss) } } else { ss.connected.Store(true) @@ -637,3 +644,53 @@ func poolHealthLoop(ctx context.Context, ss *serverState) { } } } + +// poolTryReconnect attempts a full reconnect for a pool-managed connection. +// Creates a fresh client, atomically swaps ss.clientPtr so all BridgeTools +// see the new client, then closes the old one. +func poolTryReconnect(ctx context.Context, ss *serverState) { + ss.mu.Lock() + if ss.reconnAttempts >= maxReconnectAttempts { + ss.lastErr = fmt.Sprintf("max reconnect attempts (%d) reached, entering cooldown", maxReconnectAttempts) + ss.mu.Unlock() + slog.Warn("mcp.pool.reconnect_cooldown", "server", ss.name, "cooldown", reconnectCooldown) + select { + case <-ctx.Done(): + return + case <-time.After(reconnectCooldown): + } + ss.mu.Lock() + ss.reconnAttempts = 0 + ss.mu.Unlock() + return + } + ss.reconnAttempts++ + attempt := ss.reconnAttempts + ss.mu.Unlock() + + backoff := min(initialBackoff*time.Duration(1<<(attempt-1)), maxBackoff) + slog.Info("mcp.pool.reconnecting", "server", ss.name, "attempt", attempt, "backoff", backoff) + + select { + case <-ctx.Done(): + return + case <-time.After(backoff): + } + + // Fast path: transport may have auto-reconnected. + if err := ss.client.Ping(ctx); err == nil { + ss.connected.Store(true) + ss.mu.Lock() + ss.reconnAttempts = 0 + ss.healthFailures = 0 + ss.lastErr = "" + ss.mu.Unlock() + slog.Info("mcp.pool.reconnected", "server", ss.name) + return + } + + // Slow path: create fresh client, swap atomically, close old. + if fullReconnect(ctx, ss) { + slog.Info("mcp.pool.reconnected", "server", ss.name, "method", "full_reconnect") + } +}