diff --git a/services/core/java/com/android/server/wm/AppWindowToken.java b/services/core/java/com/android/server/wm/AppWindowToken.java
index 1f71b8f..8540feb 100644
--- a/services/core/java/com/android/server/wm/AppWindowToken.java
+++ b/services/core/java/com/android/server/wm/AppWindowToken.java
@@ -252,6 +252,7 @@
     private final Point mTmpPoint = new Point();
     private final Rect mTmpRect = new Rect();
     private RemoteAnimationDefinition mRemoteAnimationDefinition;
+    private AnimatingAppWindowTokenRegistry mAnimatingAppWindowTokenRegistry;
 
     AppWindowToken(WindowManagerService service, IApplicationToken token, boolean voiceInteraction,
             DisplayContent dc, long inputDispatchingTimeoutNanos, boolean fullscreen,
@@ -781,6 +782,16 @@
                 task.mStack.mExitingAppTokens.remove(this);
             }
         }
+        final TaskStack stack = getStack();
+
+        // If we reparent, make sure to remove ourselves from the old animation registry.
+        if (mAnimatingAppWindowTokenRegistry != null) {
+            mAnimatingAppWindowTokenRegistry.notifyFinished(this);
+        }
+        mAnimatingAppWindowTokenRegistry = stack != null
+                ? stack.getAnimatingAppWindowTokenRegistry()
+                : null;
+
         mLastParent = task;
     }
 
@@ -1784,6 +1795,21 @@
     }
 
     @Override
+    public boolean shouldDeferAnimationFinish(Runnable endDeferFinishCallback) {
+        return mAnimatingAppWindowTokenRegistry != null
+                && mAnimatingAppWindowTokenRegistry.notifyAboutToFinish(
+                        this, endDeferFinishCallback);
+    }
+
+    @Override
+    public void onAnimationLeashDestroyed(Transaction t) {
+        super.onAnimationLeashDestroyed(t);
+        if (mAnimatingAppWindowTokenRegistry != null) {
+            mAnimatingAppWindowTokenRegistry.notifyFinished(this);
+        }
+    }
+
+    @Override
     protected void setLayer(Transaction t, int layer) {
         if (!mSurfaceAnimator.hasLeash()) {
             t.setLayer(mSurfaceControl, layer);
@@ -1825,6 +1851,9 @@
 
         final DisplayContent dc = getDisplayContent();
         dc.assignStackOrdering(t);
+        if (mAnimatingAppWindowTokenRegistry != null) {
+            mAnimatingAppWindowTokenRegistry.notifyStarting(this);
+        }
     }
 
     /**