fix(mcp): Improve connection health checks

michiosw · michiosw · commit e4d3dfa33eb7 · 2025-04-12T12:35:26.000-07:00
- Verify connection health before reuse with ListTools
- Use existing session context for health checks
- Add debug logging for health check failures
diff --git a/internal/mcp/server.go b/internal/mcp/server.go
@@ -297,7 +297,13 @@ func connectToRemoteServer(ctx context.Context, logger log.Logger, remote Remote
 	if existing, exists := remoteConnections[remote.URL]; exists {
 		connectionsMutex.RUnlock()
 		if existing.ctx != nil && existing.ctx.Err() == nil {
-			return existing, 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
+				}
+			}
 		}
 	} else {
 		connectionsMutex.RUnlock()
@@ -522,11 +528,21 @@ func (cm *ConnectionManager) checkConnections() {
 
 // isConnectionHealthy checks if a connection is working properly
 func (cm *ConnectionManager) isConnectionHealthy(conn *MCPConnection) bool {
+	// Use the existing connection context which has the session ID
 	ctx, cancel := context.WithTimeout(conn.ctx, 5*time.Second)
 	defer cancel()
 
-	_, err := conn.client.ListTools(ctx, mcp.ListToolsRequest{})
-	return err == nil
+	// Use the same ListTools request as during initial connection
+	toolsRequest := mcp.ListToolsRequest{}
+	_, err := conn.client.ListTools(ctx, toolsRequest)
+	if err != nil {
+		cm.logger.Debug().
+			Err(err).
+			Str("url", conn.url).
+			Msg("Health check failed")
+		return false
+	}
+	return true
 }
 
 // connectWithRetry attempts to connect to a remote server with exponential backoff
