fix: Prevent reconnect loop and decouple tool refresh

michiosw · michiosw · commit 6e0ae8ade2e8 · 2025-04-14T17:14:12.000-07:00
- Performs cleanup (cancel/close) of old MCP connections asynchronously in a separate goroutine within `connectToRemoteServer` to prevent interference with newly established connections. This resolves the immediate context cancellation and rapid reconnect loop observed previously.
- Defines `ToolsProcessedInDBEvent` triggered by `BackendService` after successfully processing tools from a `ToolsUpdatedEvent`.
- `ConnectionManager` now subscribes to `ToolsProcessedInDBEvent` and calls `refreshMCPServerTools` only upon receiving this event, decoupling the internal MCP server tool update from the database update process.
- Removes the direct call to `refreshMCPServerTools` from `updateServerTools` in `ConnectionManager`.
diff --git a/internal/backend/service.go b/internal/backend/service.go
@@ -330,6 +330,14 @@ func (s *service) HandleToolsUpdated(event events.Event) {
 		Int("updated", updatedCount).
 		Msg("Finished processing ToolsUpdatedEvent")
 
+	// Publish event to signal completion of DB processing for this server's tools
+	if !hadError {
+		s.logger.Info().Int64("serverID", serverID).Str("url", toolsEvent.ServerURL).Msg("Publishing ToolsProcessedInDBEvent")
+		s.bus.Publish(events.NewToolsProcessedInDBEvent(serverID, toolsEvent.ServerURL))
+	} else {
+		s.logger.Warn().Int64("serverID", serverID).Msg("Skipping ToolsProcessedInDBEvent due to errors during tool upsert")
+	}
+
 	// TODO: Potentially publish another event like 'ToolsProcessed' if other components
 	// need to know the DB update is complete (e.g., to trigger MCP server tool refresh).
 }
diff --git a/internal/events/events.go b/internal/events/events.go
@@ -20,6 +20,8 @@ const (
 	ToolsUpdated EventType = "tools.updated"
 	// ServerStatusChanged indicates the connection status of a server has changed
 	ServerStatusChanged EventType = "server.status.changed"
+	// ToolsProcessedInDB indicates that the backend service has finished processing tools from an update event.
+	ToolsProcessedInDB EventType = "tools.processed.db"
 	// Add more event types as needed
 )
 
@@ -117,3 +119,22 @@ func NewServerStatusChangedEvent(serverID int64, serverURL string, newState mode
 		LastError: lastError,
 	}
 }
+
+// --- NEW Event ---
+
+// ToolsProcessedInDBEvent is published by the BackendService after it finishes processing tools from a ToolsUpdatedEvent.
+// This signals the ConnectionManager that it's safe to update the internal MCP server's tool list.
+type ToolsProcessedInDBEvent struct {
+	baseEvent
+	ServerID  int64
+	ServerURL string // Include URL as ID might not be known to all potential future listeners
+}
+
+// NewToolsProcessedInDBEvent creates a new ToolsProcessedInDBEvent.
+func NewToolsProcessedInDBEvent(serverID int64, serverURL string) *ToolsProcessedInDBEvent {
+	return &ToolsProcessedInDBEvent{
+		baseEvent: newBaseEvent(ToolsProcessedInDB),
+		ServerID:  serverID,
+		ServerURL: serverURL,
+	}
+}
diff --git a/internal/mcp/server.go b/internal/mcp/server.go
@@ -275,81 +275,131 @@ var connectionsMutex sync.RWMutex
 
 // connectToRemoteServer establishes a connection to a remote MCP server
 func connectToRemoteServer(ctx context.Context, logger log.Logger, remote RemoteMCPServer) (*MCPConnection, error) {
-	// Check if we already have a valid connection
+	// --- Step 1: Check for existing connection (read lock) ---
 	connectionsMutex.RLock()
-	if existing, exists := remoteConnections[remote.URL]; exists {
-		connectionsMutex.RUnlock()
-		if existing.ctx != nil && existing.ctx.Err() == nil {
-			// Verify connection is actually healthy before reusing
-			if existing.client != nil {
-				toolsRequest := mcp.ListToolsRequest{}
-				if _, err := existing.client.ListTools(existing.ctx, toolsRequest); err == nil {
-					return existing, nil
-				}
+	existingConn, alreadyConnected := remoteConnections[remote.URL]
+	connectionsMutex.RUnlock()
+
+	if alreadyConnected && existingConn.ctx != nil && existingConn.ctx.Err() == nil {
+		// Verify connection is actually healthy before reusing
+		if existingConn.client != nil {
+			// Use a short timeout for the health check
+			checkCtx, checkCancel := context.WithTimeout(existingConn.ctx, 3*time.Second)
+			toolsRequest := mcp.ListToolsRequest{}
+			_, err := existingConn.client.ListTools(checkCtx, toolsRequest)
+			checkCancel()
+			if err == nil {
+				logger.Debug().Str("url", remote.URL).Msg("Reusing existing healthy connection")
+				return existingConn, nil
 			}
+			logger.Warn().Err(err).Str("url", remote.URL).Msg("Existing connection found but failed health check, proceeding to reconnect.")
 		}
-	} else {
-		connectionsMutex.RUnlock()
 	}
 
+	// --- Step 2: Establish NEW connection (no lock held) ---
+	logger.Debug().Str("url", remote.URL).Msg("Creating new MCP client")
 	mcpClient, err := mcpclient.NewSSEMCPClient(remote.URL)
 	if err != nil {
-		return nil, fmt.Errorf("failed to create client for %s: %v", remote.URL, err)
+		return nil, fmt.Errorf("failed to create client for %s: %w", remote.URL, err)
 	}
 
-	// Create new context
-	connCtx, cancel := context.WithCancel(ctx)
+	// Create new context derived from the parent context (e.g., cm.ctx)
+	connCtx, baseCancel := context.WithCancel(ctx) // Get base cancel func
+	// Create a logging wrapper for the cancel function
+	cancel := func() {
+		logger.Warn().Str("url", remote.URL).Msgf("!!! Cancel function for context %p called !!!", connCtx)
+		// Optional: Print stack trace for debugging
+		// debug.PrintStack()
+		baseCancel() // Call the original cancel func
+	}
 
-	// Start the client
+	logger.Debug().Str("url", remote.URL).Msg("Starting new MCP client")
 	if err := mcpClient.Start(connCtx); err != nil {
-		cancel()
-		return nil, fmt.Errorf("failed to start client: %v", err)
+		cancel() // Clean up the new context if start fails
+		return nil, fmt.Errorf("failed to start client: %w", err)
 	}
 
 	// Initialize the client
+	logger.Debug().Str("url", remote.URL).Msg("Initializing new MCP client")
 	initRequest := mcp.InitializeRequest{}
 	initRequest.Params.ProtocolVersion = mcp.LATEST_PROTOCOL_VERSION
 	initRequest.Params.ClientInfo = mcp.Implementation{
 		Name:    "cobrowser-agent",
 		Version: "1.0.0",
 	}
-
 	if _, err := mcpClient.Initialize(connCtx, initRequest); err != nil {
-		cancel()
-		return nil, fmt.Errorf("failed to initialize client: %v", err)
+		cancel()          // Clean up the new context
+		mcpClient.Close() // Close the client
+		return nil, fmt.Errorf("failed to initialize client: %w", err)
 	}
 
 	// Get initial tools list
+	logger.Debug().Str("url", remote.URL).Msg("Listing tools from new MCP client")
 	toolsRequest := mcp.ListToolsRequest{}
 	toolsResult, err := mcpClient.ListTools(connCtx, toolsRequest)
 	if err != nil {
-		cancel()
-		return nil, fmt.Errorf("failed to list tools: %v", err)
+		cancel()          // Clean up the new context
+		mcpClient.Close() // Close the client
+		return nil, fmt.Errorf("failed to list tools: %w", err)
 	}
 
-	conn := &MCPConnection{
+	// Create the new connection object using the wrapped cancel func
+	newConn := &MCPConnection{
 		client: mcpClient,
 		url:    remote.URL,
 		ctx:    connCtx,
-		cancel: cancel,
+		cancel: cancel, // Use the logging wrapper cancel func
 		tools:  toolsResult.Tools,
 	}
 
-	// Store the connection in the map
+	// --- Step 3: Atomically swap connections (write lock) ---
 	connectionsMutex.Lock()
-	if existing, exists := remoteConnections[remote.URL]; exists && existing.cancel != nil {
-		existing.cancel() // Clean up existing connection
-		existing.client.Close()
+	// Check again for existing connection that might have appeared
+	// or use the one we found in Step 1 (existingConn).
+	connToReplace, existsNow := remoteConnections[remote.URL]
+	var oldConnToCleanup *MCPConnection // Store the connection needing cleanup
+	if existsNow && connToReplace != nil {
+		if connToReplace != newConn { // Ensure we don't cleanup the one we just created
+			logger.Warn().Str("url", remote.URL).Msg("Marking existing connection (found during final swap) for cleanup.")
+			oldConnToCleanup = connToReplace
+		} else {
+			logger.Warn().Str("url", remote.URL).Msg("New connection object was already in map during final swap? Not cleaning up.")
+		}
+	} else if alreadyConnected && existingConn != nil { // Use the one from Step 1 if none appeared
+		if existingConn != newConn {
+			logger.Info().Str("url", remote.URL).Msg("Marking previous connection (from step 1) for cleanup.")
+			oldConnToCleanup = existingConn
+		}
+	}
+
+	// Store the NEW connection in the map
+	remoteConnections[remote.URL] = newConn
+	connectionsMutex.Unlock() // Release lock BEFORE cleanup
+
+	// --- Step 4: Cleanup old connection asynchronously (if necessary) ---
+	if oldConnToCleanup != nil {
+		go func(connToClose *MCPConnection) {
+			logger.Info().Str("url", connToClose.url).Msg("Starting asynchronous cleanup of old connection...")
+			// Add a small delay just in case, although ideally not needed now
+			time.Sleep(50 * time.Millisecond)
+			if connToClose.cancel != nil {
+				logger.Info().Str("url", connToClose.url).Msg("Asynchronously cancelling old connection context.")
+				connToClose.cancel()
+			}
+			if connToClose.client != nil {
+				logger.Info().Str("url", connToClose.url).Msg("Asynchronously closing old connection client.")
+				connToClose.client.Close()
+			}
+			logger.Info().Str("url", connToClose.url).Msg("Asynchronous cleanup of old connection finished.")
+		}(oldConnToCleanup) // Pass the connection to the goroutine
 	}
-	remoteConnections[remote.URL] = conn
-	connectionsMutex.Unlock()
 
 	logger.Info().
 		Str("url", remote.URL).
-		Int("tool_count", len(toolsResult.Tools)).
-		Msg("Successfully connected to MCP server")
+		Int("tool_count", len(newConn.tools)).
+		Msg("Successfully established and stored new MCP connection")
 
-	return conn, nil
+	return newConn, nil
 }
 
 // findRemoteServerForTool finds the remote server config for a tool
@@ -441,6 +491,7 @@ func RegisterEventSubscribers(bus events.Bus, cm *ConnectionManager, logger log.
 	logger.Info().Msg("Registering ConnectionManager event subscribers...")
 	bus.Subscribe(events.ServerAdded, cm.handleServerAdded)
 	bus.Subscribe(events.ServerRemoved, cm.handleServerRemoved)
+	bus.Subscribe(events.ToolsProcessedInDB, cm.handleToolsProcessed)
 }
 
 // Stop gracefully shuts down all connections
@@ -655,7 +706,11 @@ func (cm *ConnectionManager) connectWithRetry(remote RemoteMCPServer) {
 			cm.updateServerTools(remote.URL, conn.getTools())
 
 			// Monitor connection
+			cm.logger.Debug().Str("url", remote.URL).Msg("Connection successful, monitoring context...")
 			<-conn.ctx.Done()
+			// Log why the context was done
+			ctxErr := conn.ctx.Err()
+			cm.logger.Warn().Err(ctxErr).Str("url", remote.URL).Msg("Connection context finished. Reason logged.")
 
 			// --- Check again if permanently removed *after* disconnect ---
 			cm.mu.RLock()
@@ -793,79 +848,28 @@ func (cm *ConnectionManager) updateServerTools(serverURL string, fetchedTools []
 	// Publish event instead of calling API
 	cm.eventBus.Publish(events.NewToolsUpdatedEvent(serverID, serverURL, modelTools))
 
-	// --- Remove API Client Call ---
-	/*
-		start := time.Now()
-		defer func() {
-			mcpToolSyncLatency.WithLabelValues("update_tools").Observe(time.Since(start).Seconds())
-		}()
-
-		cm.mu.Lock()
-		defer cm.mu.Unlock()
-
-		// Update runtime state
-		if tools == nil {
-			delete(cm.toolsByServer, serverURL)
-			mcpToolsTotal.WithLabelValues(serverURL).Set(0)
-		} else {
-			cm.toolsByServer[serverURL] = tools
-			mcpToolsTotal.WithLabelValues(serverURL).Set(float64(len(tools)))
-		}
-
-		// Update persistent state through API
-		servers, err := cm.apiClient.ListMCPServers(cm.ctx)
-		if err != nil {
-			cm.logger.Error().Err(err).Msg("Failed to list servers for tool sync")
-			mcpToolSyncErrors.WithLabelValues("list_servers").Inc()
-			return
-		}
-
-		var serverID int64
-		for _, s := range servers {
-			if s.URL == serverURL {
-				serverID = s.ID
-				break
-			}
-		}
-
-		if serverID != 0 && tools != nil {
-			// Convert tools to FetchedTool format
-			var fetchedTools []models.FetchedTool
-			for _, tool := range tools {
-				fetchedTools = append(fetchedTools, models.FetchedTool{
-					ID:          tool.Name,
-					Name:        tool.Name,
-					Description: tool.Description,
-				})
-			}
-
-			// Sync through API
-			if err := cm.apiClient.ProcessFetchedTools(cm.ctx, serverID, fetchedTools); err != nil {
-				cm.logger.Error().
-					Err(err).
-					Str("server", serverURL).
-					Int64("serverID", serverID).
-					Msg("Failed to sync tools with API")
-				mcpToolSyncErrors.WithLabelValues("process_tools").Inc()
-			} else {
-				cm.logger.Info().
-					Str("server", serverURL).
-					Int64("serverID", serverID).
-					Int("toolCount", len(fetchedTools)).
-					Msg("Synced tools with API")
-			}
-		}
-
-		// Update MCP server tools
-		cm.refreshMCPServerTools()
-	*/
 	// TODO: Decide when/how cm.refreshMCPServerTools() should be called now.
 	// Maybe the backend service publishes another event after DB update?
 	// Or maybe ConnectionManager subscribes to its own published events? Less ideal.
 	// For now, let's call it directly after publishing, but this is not quite right.
-	cm.mu.Lock()
-	cm.refreshMCPServerTools()
-	cm.mu.Unlock()
+
+	/*
+		Calling refreshMCPServerTools() immediately after publishing the ToolsUpdatedEvent
+		means the internal MCP server's tool list is updated before the BackendService
+		has necessarily finished processing the event and updating the database.
+		This could lead to inconsistent states or potentially trigger other unintended
+		actions that result in repeated connection/update cycles.
+
+		Recommendation:
+		The refreshMCPServerTools() call should be decoupled and triggered only after
+		the BackendService confirms the database update for the tools is complete.
+		This could be done by:
+		1. The BackendService.HandleToolsUpdated method publishing a new event (e.g., ToolsProcessedInDBEvent).
+		2. The ConnectionManager subscribing to this new event and calling refreshMCPServerTools() in its handler.
+	*/
+	// cm.mu.Lock() // REMOVE Call to refreshMCPServerTools
+	// cm.refreshMCPServerTools() // REMOVE Call
+	// cm.mu.Unlock() // REMOVE Call
 }
 
 // createToolHandler creates a new handler function for a tool
@@ -1061,6 +1065,26 @@ func (cm *ConnectionManager) handleServerRemoved(event events.Event) {
 	}
 }
 
+// handleToolsProcessed is called when the BackendService confirms tools have been processed in the DB.
+// This is the correct time to refresh the tools exposed by this agent's MCP server.
+func (cm *ConnectionManager) handleToolsProcessed(event events.Event) {
+	processedEvent, ok := event.(*events.ToolsProcessedInDBEvent)
+	if !ok {
+		cm.logger.Error().Str("eventType", string(event.Type())).Msg("Received event of unexpected type in handleToolsProcessed")
+		return
+	}
+
+	cm.logger.Info().
+		Int64("serverID", processedEvent.ServerID).
+		Str("url", processedEvent.ServerURL).
+		Msg("Received ToolsProcessedInDBEvent, refreshing MCP server tools.")
+
+	// Now we refresh the tools served by *our* MCP server
+	cm.mu.Lock()
+	cm.refreshMCPServerTools()
+	cm.mu.Unlock()
+}
+
 // refreshMCPServerTools rebuilds the runtime tool state.
 // IMPORTANT: This function assumes the caller holds the cm.mu lock.
 func (cm *ConnectionManager) refreshMCPServerTools() {
