diff --git a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleStackView.java b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleStackView.java
index a4a0fe18f5cb268216c065bfb3f3d324c01f7330..00ed4f1dd505c35c7094de56ad972c02c253a0c5 100644
--- a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleStackView.java
+++ b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleStackView.java
@@ -717,12 +717,10 @@ public class BubbleStackView extends FrameLayout {
 
     /** Moves the bubbles out of the way if they're going to be over the keyboard. */
     public void onImeVisibilityChanged(boolean visible, int height) {
+        mStackAnimationController.setImeHeight(height + mImeOffset);
+
         if (!mIsExpanded) {
-            if (visible) {
-                mStackAnimationController.updateBoundsForVisibleImeAndAnimate(height + mImeOffset);
-            } else {
-                mStackAnimationController.updateBoundsForInvisibleImeAndAnimate();
-            }
+            mStackAnimationController.animateForImeVisibility(visible);
         }
     }
 
diff --git a/packages/SystemUI/src/com/android/systemui/bubbles/animation/StackAnimationController.java b/packages/SystemUI/src/com/android/systemui/bubbles/animation/StackAnimationController.java
index 74a6b60054500a63781247d1c4565c6695121b28..e0912bf31ec9a1194b4dec61eb9c7591e8ca23a8 100644
--- a/packages/SystemUI/src/com/android/systemui/bubbles/animation/StackAnimationController.java
+++ b/packages/SystemUI/src/com/android/systemui/bubbles/animation/StackAnimationController.java
@@ -334,41 +334,38 @@ public class StackAnimationController extends
         mLayout.removeEndActionForProperty(DynamicAnimation.TRANSLATION_Y);
     }
 
-    /**
-     * Save the IME height so that the allowable stack bounds reflect the now-visible IME, and
-     * animate the stack out of the way if necessary.
-     */
-    public void updateBoundsForVisibleImeAndAnimate(int imeHeight) {
+    /** Save the current IME height so that we know where the stack bounds should be. */
+    public void setImeHeight(int imeHeight) {
         mImeHeight = imeHeight;
-
-        final float maxBubbleY = getAllowableStackPositionRegion().bottom;
-        if (mStackPosition.y > maxBubbleY && mPreImeY == Float.MIN_VALUE) {
-            mPreImeY = mStackPosition.y;
-
-            springFirstBubbleWithStackFollowing(
-                    DynamicAnimation.TRANSLATION_Y,
-                    getSpringForce(DynamicAnimation.TRANSLATION_Y, /* view */ null)
-                            .setStiffness(SpringForce.STIFFNESS_LOW),
-                    /* startVel */ 0f,
-                    maxBubbleY);
-        }
     }
 
     /**
-     * Clear the IME height from the bounds and animate the stack back to its original position,
-     * assuming it wasn't moved in the meantime.
+     * Animates the stack either away from the newly visible IME, or back to its original position
+     * due to the IME going away.
      */
-    public void updateBoundsForInvisibleImeAndAnimate() {
-        mImeHeight = 0;
+    public void animateForImeVisibility(boolean imeVisible) {
+        final float maxBubbleY = getAllowableStackPositionRegion().bottom;
+        float destinationY = Float.MIN_VALUE;
 
-        if (mPreImeY > Float.MIN_VALUE) {
+        if (imeVisible) {
+            if (mStackPosition.y > maxBubbleY && mPreImeY == Float.MIN_VALUE) {
+                mPreImeY = mStackPosition.y;
+                destinationY = maxBubbleY;
+            }
+        } else {
+            if (mPreImeY > Float.MIN_VALUE) {
+                destinationY = mPreImeY;
+                mPreImeY = Float.MIN_VALUE;
+            }
+        }
+
+        if (destinationY > Float.MIN_VALUE) {
             springFirstBubbleWithStackFollowing(
                     DynamicAnimation.TRANSLATION_Y,
                     getSpringForce(DynamicAnimation.TRANSLATION_Y, /* view */ null)
-                        .setStiffness(SpringForce.STIFFNESS_LOW),
+                            .setStiffness(SpringForce.STIFFNESS_LOW),
                     /* startVel */ 0f,
-                    mPreImeY);
-            mPreImeY = Float.MIN_VALUE;
+                    destinationY);
         }
     }