diff --git a/js/lb-direct-add.js b/js/lb-direct-add.js
index d6502ea..defe210 100644
--- a/js/lb-direct-add.js
+++ b/js/lb-direct-add.js
@@ -28,21 +28,10 @@
    */
   Drupal.behaviors.lbDirectAdd = {
     attach: function (context, settings) {
-      var $directAddWidgets = $(once('lb_direct_add', '.layout-builder__add-block', context));
-
-      if ($directAddWidgets.length) {
-        var $body = $(once('lb_direct_add-click', 'body', context));
-
-        if ($body.length) {
-          $body.on('click', '.layout-builder__direct-add__toggle', directAddClickHandler);
-        }
-
-        var il = $directAddWidgets.length;
-
-        for (var i = 0; i < il; i++) {
-          DirectAddWidget.widgets.push(new DirectAddWidget($directAddWidgets[i], settings.lbDirectAdd));
-        }
-      }
+      $(once('bind-click', '.layout-builder__add-block', context)).each(function() {
+        $(this).on('click', directAddClickHandler);
+        DirectAddWidget.widgets.push(new DirectAddWidget(this, settings.lbDirectAdd));
+      });
     }
   };
   $.extend(DirectAddWidget, {