diff --git a/js/dom.js b/js/dom.js
index 091f58be..d3b09859 100644
--- a/js/dom.js
+++ b/js/dom.js
@@ -169,7 +169,7 @@ function animateElement(
   element, {
     className = 'highlight',
     removeExtraClasses = [],
-    remove = false,
+    onComplete,
   } = {}) {
   return element && new Promise(resolve => {
     element.addEventListener('animationend', function _() {
@@ -180,9 +180,9 @@ function animateElement(
         // This is helpful to clean-up on the same frame
         ...removeExtraClasses
       );
-      // TODO: investigate why animation restarts if the elements is removed in .then()
-      if (remove) {
-        element.remove();
+      // TODO: investigate why animation restarts for 'display' modification in .then()
+      if (typeof onComplete === 'function') {
+        onComplete.call(element);
       }
       resolve();
     });
diff --git a/msgbox/msgbox.js b/msgbox/msgbox.js
index 4e418832..468ecb96 100644
--- a/msgbox/msgbox.js
+++ b/msgbox/msgbox.js
@@ -45,8 +45,10 @@ function messageBox({
   function resolveWith(value) {
     unbindGlobalListeners();
     setTimeout(messageBox.resolve, 0, value);
-    animateElement(messageBox.element, {className: 'fadeout', remove: true})
-      .then(removeSelf);
+    animateElement(messageBox.element, {
+      className: 'fadeout',
+      onComplete: removeSelf,
+    });
   }
 
   function createElement() {
diff --git a/popup/popup.js b/popup/popup.js
index 4b4f2ad1..40b4a2f6 100644
--- a/popup/popup.js
+++ b/popup/popup.js
@@ -312,8 +312,10 @@ Object.assign(handleEvent, {
     };
     function confirm(ok) {
       window.onkeydown = null;
-      animateElement(box, {className: 'lights-on'})
-        .then(() => (box.dataset.display = false));
+      animateElement(box, {
+        className: 'lights-on',
+        onComplete: () => (box.dataset.display = false),
+      });
       if (ok) {
         deleteStyleSafe({id}).then(() => {
           // update view with 'No styles installed for this site' message