diff --git a/core/java/android/view/InsetsController.java b/core/java/android/view/InsetsController.java
index 72ddaca..43c7bed 100644
--- a/core/java/android/view/InsetsController.java
+++ b/core/java/android/view/InsetsController.java
@@ -18,6 +18,7 @@
 
 import static android.view.InsetsState.ITYPE_CAPTION_BAR;
 import static android.view.InsetsState.ITYPE_IME;
+import static android.view.InsetsState.toInternalType;
 import static android.view.InsetsState.toPublicType;
 import static android.view.WindowInsets.Type.all;
 import static android.view.WindowInsets.Type.ime;
@@ -471,7 +472,7 @@
         if (!localStateChanged && mLastDispachedState.equals(state)) {
             return false;
         }
-        mState.set(state);
+        updateState(state);
         mLastDispachedState.set(state, true /* copySources */);
         applyLocalVisibilityOverride();
         if (localStateChanged) {
@@ -480,11 +481,25 @@
         if (!mState.equals(mLastDispachedState, true /* excludingCaptionInsets */)) {
             sendStateToWindowManager();
         }
+        return true;
+    }
+
+    private void updateState(InsetsState newState) {
+        mState.setDisplayFrame(newState.getDisplayFrame());
+        for (int i = newState.getSourcesCount() - 1; i >= 0; i--) {
+            InsetsSource source = newState.sourceAt(i);
+            getSourceConsumer(source.getType()).updateSource(source);
+        }
+        for (int i = mState.getSourcesCount() - 1; i >= 0; i--) {
+            InsetsSource source = mState.sourceAt(i);
+            if (newState.peekSource(source.getType()) == null) {
+                mState.removeSource(source.getType());
+            }
+        }
         if (mCaptionInsetsHeight != 0) {
             mState.getSource(ITYPE_CAPTION_BAR).setFrame(new Rect(mFrame.left, mFrame.top,
                     mFrame.right, mFrame.top + mCaptionInsetsHeight));
         }
-        return true;
     }
 
     private boolean captionInsetsUnchanged() {
@@ -879,8 +894,15 @@
             control.cancel();
         }
         for (int i = mRunningAnimations.size() - 1; i >= 0; i--) {
-            if (mRunningAnimations.get(i).runner == control) {
+            RunningAnimation runningAnimation = mRunningAnimations.get(i);
+            if (runningAnimation.runner == control) {
                 mRunningAnimations.remove(i);
+                ArraySet<Integer> types = toInternalType(control.getTypes());
+                for (int j = types.size() - 1; j >= 0; j--) {
+                    if (getSourceConsumer(types.valueAt(j)).notifyAnimationFinished()) {
+                        mViewRoot.notifyInsetsChanged();
+                    }
+                }
                 break;
             }
         }
diff --git a/core/java/android/view/InsetsSourceConsumer.java b/core/java/android/view/InsetsSourceConsumer.java
index f74221d..69bab4d 100644
--- a/core/java/android/view/InsetsSourceConsumer.java
+++ b/core/java/android/view/InsetsSourceConsumer.java
@@ -16,11 +16,13 @@
 
 package android.view;
 
+import static android.view.InsetsController.ANIMATION_TYPE_NONE;
 import static android.view.InsetsController.AnimationType;
 import static android.view.InsetsState.toPublicType;
 
 import android.annotation.IntDef;
 import android.annotation.Nullable;
+import android.graphics.Rect;
 import android.view.InsetsState.InternalInsetsType;
 import android.view.SurfaceControl.Transaction;
 import android.view.WindowInsets.Type.InsetsType;
@@ -64,6 +66,8 @@
     private final Supplier<Transaction> mTransactionSupplier;
     private @Nullable InsetsSourceControl mSourceControl;
     private boolean mHasWindowFocus;
+    private Rect mPendingFrame;
+    private Rect mPendingVisibleFrame;
 
     public InsetsSourceConsumer(@InternalInsetsType int type, InsetsState state,
             Supplier<Transaction> transactionSupplier, InsetsController controller) {
@@ -215,6 +219,38 @@
         // no-op for types that always return ShowResult#SHOW_IMMEDIATELY.
     }
 
+    void updateSource(InsetsSource newSource) {
+        InsetsSource source = mState.peekSource(mType);
+        if (source == null || mController.getAnimationType(mType) == ANIMATION_TYPE_NONE
+                || source.getFrame().equals(newSource.getFrame())) {
+            mState.addSource(newSource);
+            return;
+        }
+
+        // Frame is changing while animating. Keep note of the new frame but keep existing frame
+        // until animaition is finished.
+        newSource = new InsetsSource(newSource);
+        mPendingFrame = new Rect(newSource.getFrame());
+        mPendingVisibleFrame = newSource.getVisibleFrame() != null
+                ? new Rect(newSource.getVisibleFrame())
+                : null;
+        newSource.setFrame(source.getFrame());
+        newSource.setVisibleFrame(source.getVisibleFrame());
+        mState.addSource(newSource);
+    }
+
+    boolean notifyAnimationFinished() {
+        if (mPendingFrame != null) {
+            InsetsSource source = mState.getSource(mType);
+            source.setFrame(mPendingFrame);
+            source.setVisibleFrame(mPendingVisibleFrame);
+            mPendingFrame = null;
+            mPendingVisibleFrame = null;
+            return true;
+        }
+        return false;
+    }
+
     /**
      * Sets requested visibility from the client, regardless of whether we are able to control it at
      * the moment.
diff --git a/core/tests/coretests/src/android/view/InsetsControllerTest.java b/core/tests/coretests/src/android/view/InsetsControllerTest.java
index b9c71a2..cbb379b 100644
--- a/core/tests/coretests/src/android/view/InsetsControllerTest.java
+++ b/core/tests/coretests/src/android/view/InsetsControllerTest.java
@@ -32,6 +32,7 @@
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotEquals;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
@@ -638,7 +639,31 @@
         });
     }
 
+    @Test
+    public void testFrameUpdateDuringAnimation() {
+        InstrumentationRegistry.getInstrumentation().runOnMainSync(() -> {
 
+            mController.onControlsChanged(createSingletonControl(ITYPE_IME));
+
+            // Pretend IME is calling
+            mController.show(ime(), true /* fromIme */);
+
+            InsetsState copy = new InsetsState(mController.getState(), true /* copySources */);
+            copy.getSource(ITYPE_IME).setFrame(0, 1, 2, 3);
+            copy.getSource(ITYPE_IME).setVisibleFrame(new Rect(4, 5, 6, 7));
+            mController.onStateChanged(copy);
+            assertNotEquals(new Rect(0, 1, 2, 3),
+                    mController.getState().getSource(ITYPE_IME).getFrame());
+            assertNotEquals(new Rect(4, 5, 6, 7),
+                    mController.getState().getSource(ITYPE_IME).getVisibleFrame());
+            mController.cancelExistingAnimation();
+            assertEquals(new Rect(0, 1, 2, 3),
+                    mController.getState().getSource(ITYPE_IME).getFrame());
+            assertEquals(new Rect(4, 5, 6, 7),
+                    mController.getState().getSource(ITYPE_IME).getVisibleFrame());
+        });
+        InstrumentationRegistry.getInstrumentation().waitForIdleSync();
+    }
 
     @Test
     public void testCaptionInsetsStateAssemble() {
diff --git a/tests/WindowInsetsTests/src/com/google/android/test/windowinsetstests/WindowInsetsActivity.java b/tests/WindowInsetsTests/src/com/google/android/test/windowinsetstests/WindowInsetsActivity.java
index 548af0c..498cb7c 100644
--- a/tests/WindowInsetsTests/src/com/google/android/test/windowinsetstests/WindowInsetsActivity.java
+++ b/tests/WindowInsetsTests/src/com/google/android/test/windowinsetstests/WindowInsetsActivity.java
@@ -18,7 +18,6 @@
 
 import static android.view.WindowInsets.Type.ime;
 import static android.view.WindowInsetsAnimation.Callback.DISPATCH_MODE_STOP;
-
 import static java.lang.Math.max;
 import static java.lang.Math.min;
 
@@ -31,6 +30,7 @@
 import android.graphics.Insets;
 import android.os.Bundle;
 import android.util.AttributeSet;
+import android.util.Log;
 import android.view.MotionEvent;
 import android.view.View;
 import android.view.ViewConfiguration;
@@ -44,11 +44,11 @@
 import android.view.animation.LinearInterpolator;
 import android.widget.LinearLayout;
 
-import androidx.appcompat.app.AppCompatActivity;
-
 import java.util.ArrayList;
 import java.util.List;
 
+import androidx.appcompat.app.AppCompatActivity;
+
 public class WindowInsetsActivity extends AppCompatActivity {
 
     private View mRoot;
@@ -191,6 +191,40 @@
                 mTransitions.forEach(it -> it.onFinish(animation));
             }
         });
+
+        findViewById(R.id.floating_action_button).setOnClickListener(
+                v -> v.getWindowInsetsController().controlWindowInsetsAnimation(ime(), -1,
+                new LinearInterpolator(), null /* cancellationSignal */,
+                new WindowInsetsAnimationControlListener() {
+                    @Override
+                    public void onReady(
+                            WindowInsetsAnimationController controller,
+                            int types) {
+                        ValueAnimator anim = ValueAnimator.ofFloat(0f, 1f);
+                        anim.setDuration(1500);
+                        anim.addUpdateListener(animation
+                                -> controller.setInsetsAndAlpha(
+                                controller.getShownStateInsets(),
+                                (float) animation.getAnimatedValue(),
+                                anim.getAnimatedFraction()));
+                        anim.addListener(new AnimatorListenerAdapter() {
+                            @Override
+                            public void onAnimationEnd(Animator animation) {
+                                super.onAnimationEnd(animation);
+                                controller.finish(true);
+                            }
+                        });
+                        anim.start();
+                    }
+
+                    @Override
+                    public void onCancelled(WindowInsetsAnimationController controller) {
+                    }
+
+                    @Override
+                    public void onFinished(WindowInsetsAnimationController controller) {
+                    }
+                }));
     }
 
     @Override
@@ -200,57 +234,6 @@
         getWindow().getDecorView().post(() -> getWindow().setDecorFitsSystemWindows(false));
     }
 
-    @Override
-    public void onAttachedToWindow() {
-        super.onAttachedToWindow();
-        getWindow().getInsetsController().addOnControllableInsetsChangedListener(
-                new OnControllableInsetsChangedListener() {
-
-                    boolean hasControl = false;
-                    @Override
-                    public void onControllableInsetsChanged(WindowInsetsController controller,
-                            int types) {
-                        if ((types & ime()) != 0 && !hasControl) {
-                            hasControl = true;
-                            controller.controlWindowInsetsAnimation(ime(), -1,
-                                    new LinearInterpolator(), null /* cancellationSignal */,
-                                    new WindowInsetsAnimationControlListener() {
-                                        @Override
-                                        public void onReady(
-                                                WindowInsetsAnimationController controller,
-                                                int types) {
-                                            ValueAnimator anim = ValueAnimator.ofFloat(0f, 1f);
-                                            anim.setDuration(1500);
-                                            anim.addUpdateListener(animation
-                                                    -> controller.setInsetsAndAlpha(
-                                                    controller.getShownStateInsets(),
-                                                    (float) animation.getAnimatedValue(),
-                                                    anim.getAnimatedFraction()));
-                                            anim.addListener(new AnimatorListenerAdapter() {
-                                                @Override
-                                                public void onAnimationEnd(Animator animation) {
-                                                    super.onAnimationEnd(animation);
-                                                    controller.finish(true);
-                                                }
-                                            });
-                                            anim.start();
-                                        }
-
-                                        @Override
-                                        public void onFinished(
-                                                WindowInsetsAnimationController controller) {
-                                        }
-
-                                        @Override
-                                        public void onCancelled(
-                                                WindowInsetsAnimationController controller) {
-                                        }
-                                    });
-                        }
-                    }
-                });
-    }
-
     static class Transition {
         private int mEndBottom;
         private int mStartBottom;