Implement takeover functionality for the Home gesture on the Shell side.

This CL does a few things:
* Add a new registry for remotes so that they can be differentiated
  between "regular" and takeover remotes (requires registering twice for
  both functionalities).
* Add an API to Transitions to find a suitable handler to takeover a
  transition, and to TransitionHandler to handler the requests.
* Add an API to IRecentsAnimationController to hand off the recents
  animation, and to RecentsTransitionHandler to route it to the takeover
  remote correctly.

Bug: 323863002
Flag: ACONFIG com.android.systemui.shared.return_animation_framework_library DISABLED
Test: updated unit tests and tested manually
Change-Id: I58e2cccffa0da33260ae1e4b0d8a29f5965fa39a
diff --git a/libs/WindowManager/Shell/Android.bp b/libs/WindowManager/Shell/Android.bp
index 5c978e2..89781fd 100644
--- a/libs/WindowManager/Shell/Android.bp
+++ b/libs/WindowManager/Shell/Android.bp
@@ -210,6 +210,7 @@
         "androidx.recyclerview_recyclerview",
         "kotlinx-coroutines-android",
         "kotlinx-coroutines-core",
+        "//frameworks/libs/systemui:com_android_systemui_shared_flags_lib",
         "//frameworks/libs/systemui:iconloader_base",
         "com_android_wm_shell_flags_lib",
         "com.android.window.flags.window-aconfig-java",
diff --git a/libs/WindowManager/Shell/shared/src/com/android/wm/shell/shared/IShellTransitions.aidl b/libs/WindowManager/Shell/shared/src/com/android/wm/shell/shared/IShellTransitions.aidl
index 526407e..3256abf 100644
--- a/libs/WindowManager/Shell/shared/src/com/android/wm/shell/shared/IShellTransitions.aidl
+++ b/libs/WindowManager/Shell/shared/src/com/android/wm/shell/shared/IShellTransitions.aidl
@@ -28,13 +28,14 @@
 interface IShellTransitions {
 
     /**
-     * Registers a remote transition handler.
+     * Registers a remote transition handler for all operations excluding takeovers (see
+     * registerRemoteForTakeover()).
      */
     oneway void registerRemote(in TransitionFilter filter,
             in RemoteTransition remoteTransition) = 1;
 
     /**
-     * Unregisters a remote transition handler.
+     * Unregisters a remote transition handler for all operations.
      */
     oneway void unregisterRemote(in RemoteTransition remoteTransition) = 2;
 
@@ -52,4 +53,10 @@
      * Returns a container surface for the home root task.
      */
     SurfaceControl getHomeTaskOverlayContainer() = 5;
+
+    /**
+     * Registers a remote transition for takeover operations only.
+     */
+    oneway void registerRemoteForTakeover(in TransitionFilter filter,
+            in RemoteTransition remoteTransition) = 6;
 }
diff --git a/libs/WindowManager/Shell/shared/src/com/android/wm/shell/shared/ShellTransitions.java b/libs/WindowManager/Shell/shared/src/com/android/wm/shell/shared/ShellTransitions.java
index 5e49f55..6d4ab4c 100644
--- a/libs/WindowManager/Shell/shared/src/com/android/wm/shell/shared/ShellTransitions.java
+++ b/libs/WindowManager/Shell/shared/src/com/android/wm/shell/shared/ShellTransitions.java
@@ -28,13 +28,20 @@
 @ExternalThread
 public interface ShellTransitions {
     /**
-     * Registers a remote transition.
+     * Registers a remote transition for all operations excluding takeovers (see
+     * {@link ShellTransitions#registerRemoteForTakeover(TransitionFilter, RemoteTransition)}).
      */
     default void registerRemote(@NonNull TransitionFilter filter,
             @NonNull RemoteTransition remoteTransition) {}
 
     /**
-     * Unregisters a remote transition.
+     * Registers a remote transition for takeover operations only.
+     */
+    default void registerRemoteForTakeover(@NonNull TransitionFilter filter,
+            @NonNull RemoteTransition remoteTransition) {}
+
+    /**
+     * Unregisters a remote transition for all operations.
      */
     default void unregisterRemote(@NonNull RemoteTransition remoteTransition) {}
 }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/recents/RecentsTransitionHandler.java b/libs/WindowManager/Shell/src/com/android/wm/shell/recents/RecentsTransitionHandler.java
index c625b69..a7829c9 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/recents/RecentsTransitionHandler.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/recents/RecentsTransitionHandler.java
@@ -27,6 +27,7 @@
 import static android.view.WindowManager.TRANSIT_TO_FRONT;
 import static android.window.TransitionInfo.FLAG_TRANSLUCENT;
 
+import static com.android.wm.shell.sysui.ShellSharedConstants.KEY_EXTRA_SHELL_CAN_HAND_OFF_ANIMATION;
 import static com.android.wm.shell.util.SplitBounds.KEY_EXTRA_SPLIT_BOUNDS;
 
 import android.annotation.Nullable;
@@ -54,6 +55,7 @@
 import android.window.TaskSnapshot;
 import android.window.TransitionInfo;
 import android.window.TransitionRequestInfo;
+import android.window.WindowAnimationState;
 import android.window.WindowContainerToken;
 import android.window.WindowContainerTransaction;
 
@@ -283,6 +285,7 @@
         private IBinder mTransition = null;
         private boolean mKeyguardLocked = false;
         private boolean mWillFinishToHome = false;
+        private Transitions.TransitionHandler mTakeoverHandler = null;
 
         /** The animation is idle, waiting for the user to choose a task to switch to. */
         private static final int STATE_NORMAL = 0;
@@ -576,9 +579,13 @@
             ProtoLog.v(ShellProtoLogGroup.WM_SHELL_RECENTS_TRANSITION,
                     "Applying transaction=%d", t.getId());
             t.apply();
-            Bundle b = new Bundle(1 /*capacity*/);
+
+            mTakeoverHandler = mTransitions.getHandlerForTakeover(mTransition, info);
+
+            Bundle b = new Bundle(2 /*capacity*/);
             b.putParcelable(KEY_EXTRA_SPLIT_BOUNDS,
                     mRecentTasksController.getSplitBoundsForTaskId(closingSplitTaskId));
+            b.putBoolean(KEY_EXTRA_SHELL_CAN_HAND_OFF_ANIMATION, mTakeoverHandler != null);
             try {
                 ProtoLog.v(ShellProtoLogGroup.WM_SHELL_RECENTS_TRANSITION,
                         "[%d] RecentsController.start: calling onAnimationStart with %d apps",
@@ -597,6 +604,63 @@
             return true;
         }
 
+        @Override
+        public void handOffAnimation(
+                RemoteAnimationTarget[] targets, WindowAnimationState[] states) {
+            mExecutor.execute(() -> {
+                ProtoLog.v(ShellProtoLogGroup.WM_SHELL_RECENTS_TRANSITION,
+                        "[%d] RecentsController.handOffAnimation", mInstanceId);
+
+                if (mTakeoverHandler == null) {
+                    Slog.e(TAG, "Tried to hand off an animation without a valid takeover "
+                            + "handler.");
+                    return;
+                }
+
+                if (targets.length != states.length) {
+                    Slog.e(TAG, "Tried to hand off an animation, but the number of targets "
+                            + "(" + targets.length + ") doesn't match the number of states "
+                            + "(" + states.length + ")");
+                    return;
+                }
+
+                ProtoLog.v(ShellProtoLogGroup.WM_SHELL_RECENTS_TRANSITION,
+                        "[%d] RecentsController.handOffAnimation: got %d states for %d "
+                                + "changes", mInstanceId, states.length, mInfo.getChanges().size());
+                WindowAnimationState[] updatedStates =
+                        new WindowAnimationState[mInfo.getChanges().size()];
+
+                // Ensure that the ordering of animation states is the same as that of  matching
+                // changes in mInfo. prefixOrderIndex is set up in reverse order to that of the
+                // changes, so that's what we use to get to the correct ordering.
+                for (int i = 0; i < targets.length; i++) {
+                    RemoteAnimationTarget target = targets[i];
+                    updatedStates[updatedStates.length - target.prefixOrderIndex] = states[i];
+                }
+
+                Transitions.TransitionFinishCallback finishCB = mFinishCB;
+                // Reset the callback here, so any stray calls that aren't coming from the new
+                // handler are ignored.
+                mFinishCB = null;
+
+                ProtoLog.v(ShellProtoLogGroup.WM_SHELL_RECENTS_TRANSITION,
+                        "[%d] RecentsController.handOffAnimation: calling "
+                                + "takeOverAnimation with %d states", mInstanceId,
+                        updatedStates.length);
+                mTakeoverHandler.takeOverAnimation(
+                        mTransition, mInfo, new SurfaceControl.Transaction(),
+                        wct -> {
+                            ProtoLog.v(ShellProtoLogGroup.WM_SHELL_RECENTS_TRANSITION,
+                                    "[%d] RecentsController.handOffAnimation: finish "
+                                            + "callback", mInstanceId);
+                            // Set the callback once again so we can finish correctly.
+                            mFinishCB = finishCB;
+                            finishInner(true /* toHome */, false /* userLeave */,
+                                    null /* finishCb */);
+                        }, updatedStates);
+            });
+        }
+
         /**
          * Updates this controller when a new transition is requested mid-recents transition.
          */
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/sysui/ShellSharedConstants.java b/libs/WindowManager/Shell/src/com/android/wm/shell/sysui/ShellSharedConstants.java
index 56c0d0e..c886cc9 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/sysui/ShellSharedConstants.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/sysui/ShellSharedConstants.java
@@ -42,4 +42,7 @@
     public static final String KEY_EXTRA_SHELL_DESKTOP_MODE = "extra_shell_desktop_mode";
     // See IDragAndDrop.aidl
     public static final String KEY_EXTRA_SHELL_DRAG_AND_DROP = "extra_shell_drag_and_drop";
+    // See IRecentsAnimationController.aidl
+    public static final String KEY_EXTRA_SHELL_CAN_HAND_OFF_ANIMATION =
+            "extra_shell_can_hand_off_animation";
 }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/transition/OneShotRemoteHandler.java b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/OneShotRemoteHandler.java
index 94519a0..69c4167 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/transition/OneShotRemoteHandler.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/OneShotRemoteHandler.java
@@ -27,6 +27,7 @@
 import android.window.RemoteTransition;
 import android.window.TransitionInfo;
 import android.window.TransitionRequestInfo;
+import android.window.WindowAnimationState;
 import android.window.WindowContainerTransaction;
 
 import com.android.internal.protolog.common.ProtoLog;
@@ -65,30 +66,9 @@
         ProtoLog.v(ShellProtoLogGroup.WM_SHELL_TRANSITIONS, "Using registered One-shot remote"
                 + " transition %s for (#%d).", mRemote, info.getDebugId());
 
-        final IBinder.DeathRecipient remoteDied = () -> {
-            Log.e(Transitions.TAG, "Remote transition died, finishing");
-            mMainExecutor.execute(
-                    () -> finishCallback.onTransitionFinished(null /* wct */));
-        };
-        IRemoteTransitionFinishedCallback cb = new IRemoteTransitionFinishedCallback.Stub() {
-            @Override
-            public void onTransitionFinished(WindowContainerTransaction wct,
-                    SurfaceControl.Transaction sct) {
-                ProtoLog.v(ShellProtoLogGroup.WM_SHELL_TRANSITIONS,
-                        "Finished one-shot remote transition %s for (#%d).", mRemote,
-                        info.getDebugId());
-                if (mRemote.asBinder() != null) {
-                    mRemote.asBinder().unlinkToDeath(remoteDied, 0 /* flags */);
-                }
-                if (sct != null) {
-                    finishTransaction.merge(sct);
-                }
-                mMainExecutor.execute(() -> {
-                    finishCallback.onTransitionFinished(wct);
-                    mRemote = null;
-                });
-            }
-        };
+        final IBinder.DeathRecipient remoteDied = createDeathRecipient(finishCallback);
+        IRemoteTransitionFinishedCallback cb =
+                createFinishedCallback(info, finishTransaction, finishCallback, remoteDied);
         Transitions.setRunningRemoteTransitionDelegate(mRemote.getAppThread());
         try {
             if (mRemote.asBinder() != null) {
@@ -152,6 +132,51 @@
     }
 
     @Override
+    public boolean takeOverAnimation(
+            @NonNull IBinder transition, @NonNull TransitionInfo info,
+            @NonNull SurfaceControl.Transaction transaction,
+            @NonNull Transitions.TransitionFinishCallback finishCallback,
+            @NonNull WindowAnimationState[] states) {
+        if (mTransition != transition) return false;
+        ProtoLog.v(ShellProtoLogGroup.WM_SHELL_RECENTS_TRANSITION, "Using registered One-shot "
+                + "remote transition %s to take over (#%d).", mRemote, info.getDebugId());
+
+        final IBinder.DeathRecipient remoteDied = createDeathRecipient(finishCallback);
+        IRemoteTransitionFinishedCallback cb = createFinishedCallback(
+                info, null /* finishTransaction */, finishCallback, remoteDied);
+
+        Transitions.setRunningRemoteTransitionDelegate(mRemote.getAppThread());
+
+        try {
+            if (mRemote.asBinder() != null) {
+                mRemote.asBinder().linkToDeath(remoteDied, 0 /* flags */);
+            }
+
+            // If the remote is actually in the same process, then make a copy of parameters since
+            // remote impls assume that they have to clean-up native references.
+            final SurfaceControl.Transaction remoteStartT =
+                    RemoteTransitionHandler.copyIfLocal(transaction, mRemote.getRemoteTransition());
+            final TransitionInfo remoteInfo =
+                    remoteStartT == transaction ? info : info.localRemoteCopy();
+            mRemote.getRemoteTransition().takeOverAnimation(
+                    transition, remoteInfo, remoteStartT, cb, states);
+
+            // Assume that remote will apply the transaction.
+            transaction.clear();
+            return true;
+        } catch (RemoteException e) {
+            Log.e(Transitions.TAG, "Error running remote transition takeover.", e);
+            if (mRemote.asBinder() != null) {
+                mRemote.asBinder().unlinkToDeath(remoteDied, 0 /* flags */);
+            }
+            finishCallback.onTransitionFinished(null /* wct */);
+            mRemote = null;
+        }
+
+        return false;
+    }
+
+    @Override
     @Nullable
     public WindowContainerTransaction handleRequest(@NonNull IBinder transition,
             @Nullable TransitionRequestInfo request) {
@@ -174,6 +199,41 @@
         }
     }
 
+    private IBinder.DeathRecipient createDeathRecipient(
+            Transitions.TransitionFinishCallback finishCallback) {
+        return () -> {
+            Log.e(Transitions.TAG, "Remote transition died, finishing");
+            mMainExecutor.execute(
+                    () -> finishCallback.onTransitionFinished(null /* wct */));
+        };
+    }
+
+    private IRemoteTransitionFinishedCallback createFinishedCallback(
+            @NonNull TransitionInfo info,
+            @Nullable SurfaceControl.Transaction finishTransaction,
+            @NonNull Transitions.TransitionFinishCallback finishCallback,
+            @NonNull IBinder.DeathRecipient remoteDied) {
+        return new IRemoteTransitionFinishedCallback.Stub() {
+            @Override
+            public void onTransitionFinished(WindowContainerTransaction wct,
+                    SurfaceControl.Transaction sct) {
+                ProtoLog.v(ShellProtoLogGroup.WM_SHELL_TRANSITIONS,
+                        "Finished one-shot remote transition %s for (#%d).", mRemote,
+                        info.getDebugId());
+                if (mRemote.asBinder() != null) {
+                    mRemote.asBinder().unlinkToDeath(remoteDied, 0 /* flags */);
+                }
+                if (finishTransaction != null && sct != null) {
+                    finishTransaction.merge(sct);
+                }
+                mMainExecutor.execute(() -> {
+                    finishCallback.onTransitionFinished(wct);
+                    mRemote = null;
+                });
+            }
+        };
+    }
+
     @Override
     public String toString() {
         return "OneShotRemoteHandler:" + mRemote.getDebugName() + ":"
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/transition/RemoteTransitionHandler.java b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/RemoteTransitionHandler.java
index 4c4c580..d686046 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/transition/RemoteTransitionHandler.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/RemoteTransitionHandler.java
@@ -16,6 +16,8 @@
 
 package com.android.wm.shell.transition;
 
+import static com.android.systemui.shared.Flags.returnAnimationFrameworkLibrary;
+
 import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.os.IBinder;
@@ -32,6 +34,7 @@
 import android.window.TransitionFilter;
 import android.window.TransitionInfo;
 import android.window.TransitionRequestInfo;
+import android.window.WindowAnimationState;
 import android.window.WindowContainerTransaction;
 
 import androidx.annotation.BinderThread;
@@ -41,7 +44,9 @@
 import com.android.wm.shell.protolog.ShellProtoLogGroup;
 import com.android.wm.shell.shared.TransitionUtil;
 
+import java.io.PrintWriter;
 import java.util.ArrayList;
+import java.util.Arrays;
 
 /**
  * Handler that deals with RemoteTransitions. It will only request to handle a transition
@@ -58,6 +63,8 @@
     /** Ordered by specificity. Last filters will be checked first */
     private final ArrayList<Pair<TransitionFilter, RemoteTransition>> mFilters =
             new ArrayList<>();
+    private final ArrayList<Pair<TransitionFilter, RemoteTransition>> mTakeoverFilters =
+            new ArrayList<>();
 
     private final ArrayMap<IBinder, RemoteDeathHandler> mDeathHandlers = new ArrayMap<>();
 
@@ -70,14 +77,23 @@
         mFilters.add(new Pair<>(filter, remote));
     }
 
+    void addFilteredForTakeover(TransitionFilter filter, RemoteTransition remote) {
+        handleDeath(remote.asBinder(), null /* finishCallback */);
+        mTakeoverFilters.add(new Pair<>(filter, remote));
+    }
+
     void removeFiltered(RemoteTransition remote) {
         boolean removed = false;
-        for (int i = mFilters.size() - 1; i >= 0; --i) {
-            if (mFilters.get(i).second.asBinder().equals(remote.asBinder())) {
-                mFilters.remove(i);
-                removed = true;
+        for (ArrayList<Pair<TransitionFilter, RemoteTransition>> filters
+                : Arrays.asList(mFilters, mTakeoverFilters)) {
+            for (int i = filters.size() - 1; i >= 0; --i) {
+                if (filters.get(i).second.asBinder().equals(remote.asBinder())) {
+                    filters.remove(i);
+                    removed = true;
+                }
             }
         }
+
         if (removed) {
             unhandleDeath(remote.asBinder(), null /* finishCallback */);
         }
@@ -237,6 +253,47 @@
         }
     }
 
+    @Nullable
+    @Override
+    public Transitions.TransitionHandler getHandlerForTakeover(
+            @NonNull IBinder transition, @NonNull TransitionInfo info) {
+        if (!returnAnimationFrameworkLibrary()) {
+            return null;
+        }
+
+        for (Pair<TransitionFilter, RemoteTransition> registered : mTakeoverFilters) {
+            if (registered.first.matches(info)) {
+                ProtoLog.v(ShellProtoLogGroup.WM_SHELL_RECENTS_TRANSITION,
+                        "Found matching remote to takeover (#%d)", info.getDebugId());
+
+                OneShotRemoteHandler oneShot =
+                        new OneShotRemoteHandler(mMainExecutor, registered.second);
+                oneShot.setTransition(transition);
+                return oneShot;
+            }
+        }
+
+        ProtoLog.v(ShellProtoLogGroup.WM_SHELL_RECENTS_TRANSITION,
+                "No matching remote found to takeover (#%d)", info.getDebugId());
+        return null;
+    }
+
+    @Override
+    public boolean takeOverAnimation(
+            @NonNull IBinder transition, @NonNull TransitionInfo info,
+            @NonNull SurfaceControl.Transaction transaction,
+            @NonNull Transitions.TransitionFinishCallback finishCallback,
+            @NonNull WindowAnimationState[] states) {
+        Transitions.TransitionHandler handler = getHandlerForTakeover(transition, info);
+        if (handler == null) {
+            ProtoLog.v(ShellProtoLogGroup.WM_SHELL_RECENTS_TRANSITION,
+                    "Take over request failed: no matching remote for (#%d)", info.getDebugId());
+            return false;
+        }
+        ((OneShotRemoteHandler) handler).setTransition(transition);
+        return handler.takeOverAnimation(transition, info, transaction, finishCallback, states);
+    }
+
     @Override
     @Nullable
     public WindowContainerTransaction handleRequest(@NonNull IBinder transition,
@@ -284,6 +341,34 @@
         }
     }
 
+    void dump(@NonNull PrintWriter pw, String prefix) {
+        final String innerPrefix = prefix + "  ";
+
+        pw.println(prefix + "Registered Remotes:");
+        if (mFilters.isEmpty()) {
+            pw.println(innerPrefix + "none");
+        } else {
+            for (Pair<TransitionFilter, RemoteTransition> entry : mFilters) {
+                dumpRemote(pw, innerPrefix, entry.second);
+            }
+        }
+
+        pw.println(prefix + "Registered Takeover Remotes:");
+        if (mTakeoverFilters.isEmpty()) {
+            pw.println(innerPrefix + "none");
+        } else {
+            for (Pair<TransitionFilter, RemoteTransition> entry : mTakeoverFilters) {
+                dumpRemote(pw, innerPrefix, entry.second);
+            }
+        }
+    }
+
+    private void dumpRemote(@NonNull PrintWriter pw, String prefix, RemoteTransition remote) {
+        pw.print(prefix);
+        pw.print(remote.getDebugName());
+        pw.println(" (" + Integer.toHexString(System.identityHashCode(remote)) + ")");
+    }
+
     /** NOTE: binder deaths can alter the filter order */
     private class RemoteDeathHandler implements IBinder.DeathRecipient {
         private final IBinder mRemote;
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/transition/Transitions.java b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/Transitions.java
index 437a00e..c6b127d 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/transition/Transitions.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/Transitions.java
@@ -35,6 +35,7 @@
 import static android.window.TransitionInfo.FLAG_NO_ANIMATION;
 import static android.window.TransitionInfo.FLAG_STARTING_WINDOW_TRANSFER_RECIPIENT;
 
+import static com.android.systemui.shared.Flags.returnAnimationFrameworkLibrary;
 import static com.android.wm.shell.common.ExecutorUtils.executeRemoteCallWithTaskPermission;
 import static com.android.wm.shell.shared.TransitionUtil.isClosingType;
 import static com.android.wm.shell.shared.TransitionUtil.isOpeningType;
@@ -62,6 +63,7 @@
 import android.window.TransitionInfo;
 import android.window.TransitionMetrics;
 import android.window.TransitionRequestInfo;
+import android.window.WindowAnimationState;
 import android.window.WindowContainerTransaction;
 
 import androidx.annotation.BinderThread;
@@ -419,12 +421,24 @@
         mHandlers.set(0, handler);
     }
 
-    /** Register a remote transition to be used when `filter` matches an incoming transition */
+    /**
+     * Register a remote transition to be used for all operations except takeovers when `filter`
+     * matches an incoming transition.
+     */
     public void registerRemote(@NonNull TransitionFilter filter,
             @NonNull RemoteTransition remoteTransition) {
         mRemoteTransitionHandler.addFiltered(filter, remoteTransition);
     }
 
+    /**
+     * Register a remote transition to be used for all operations except takeovers when `filter`
+     * matches an incoming transition.
+     */
+    public void registerRemoteForTakeover(@NonNull TransitionFilter filter,
+            @NonNull RemoteTransition remoteTransition) {
+        mRemoteTransitionHandler.addFilteredForTakeover(filter, remoteTransition);
+    }
+
     /** Unregisters a remote transition and all associated filters */
     public void unregisterRemote(@NonNull RemoteTransition remoteTransition) {
         mRemoteTransitionHandler.removeFiltered(remoteTransition);
@@ -1187,6 +1201,29 @@
     }
 
     /**
+     * Checks whether a handler exists capable of taking over the given transition, and returns it.
+     * Otherwise it returns null.
+     */
+    @Nullable
+    public TransitionHandler getHandlerForTakeover(
+            @NonNull IBinder transition, @NonNull TransitionInfo info) {
+        if (!returnAnimationFrameworkLibrary()) {
+            ProtoLog.v(ShellProtoLogGroup.WM_SHELL_RECENTS_TRANSITION,
+                    "Trying to get a handler for takeover but the flag is disabled");
+            return null;
+        }
+
+        for (TransitionHandler handler : mHandlers) {
+            TransitionHandler candidate = handler.getHandlerForTakeover(transition, info);
+            if (candidate != null) {
+                return candidate;
+            }
+        }
+
+        return null;
+    }
+
+    /**
      * Finish running animations (almost) immediately when a SLEEP transition comes in. We use this
      * as both a way to reduce unnecessary work (animations not visible while screen off) and as a
      * failsafe to unblock "stuck" animations (in particular remote animations).
@@ -1328,6 +1365,49 @@
                 @NonNull TransitionFinishCallback finishCallback) { }
 
         /**
+         * Checks whether this handler is capable of taking over a transition matching `info`.
+         * {@link TransitionHandler#takeOverAnimation(IBinder, TransitionInfo,
+         * SurfaceControl.Transaction, TransitionFinishCallback, WindowAnimationState[])} is
+         * guaranteed to succeed if called on the handler returned by this method.
+         *
+         * Note that the handler returned by this method can either be itself, or a different one
+         * selected by this handler to take care of the transition on its behalf.
+         *
+         * @param transition The transition that should be taken over.
+         * @param info Information about the transition to be taken over.
+         * @return A handler capable of taking over a matching transition, or null.
+         */
+        @Nullable
+        default TransitionHandler getHandlerForTakeover(
+                @NonNull IBinder transition, @NonNull TransitionInfo info) {
+            return null;
+        }
+
+        /**
+         * Attempt to take over a running transition. This must succeed if this handler was returned
+         * by {@link TransitionHandler#getHandlerForTakeover(IBinder, TransitionInfo)}.
+         *
+         * @param transition The transition that should be taken over.
+         * @param info Information about the what is changing in the transition.
+         * @param transaction Contains surface changes that resulted from the transition. Any
+         *                    additional changes should be added to this transaction and committed
+         *                    inside this method.
+         * @param finishCallback Call this at the end of the animation, if the take-over succeeds.
+         *                       Note that this will be called instead of the callback originally
+         *                       passed to startAnimation(), so the caller should make sure all
+         *                       necessary cleanup happens here. This MUST be called on main thread.
+         * @param states The animation states of the transition's window at the time this method was
+         *               called.
+         * @return true if the transition was taken over, false if not.
+         */
+        default boolean takeOverAnimation(@NonNull IBinder transition, @NonNull TransitionInfo info,
+                @NonNull SurfaceControl.Transaction transaction,
+                @NonNull TransitionFinishCallback finishCallback,
+                @NonNull WindowAnimationState[] states) {
+            return false;
+        }
+
+        /**
          * Potentially handles a startTransition request.
          *
          * @param transition The transition whose start is being requested.
@@ -1432,16 +1512,21 @@
         @Override
         public void registerRemote(@NonNull TransitionFilter filter,
                 @NonNull RemoteTransition remoteTransition) {
-            mMainExecutor.execute(() -> {
-                mRemoteTransitionHandler.addFiltered(filter, remoteTransition);
-            });
+            mMainExecutor.execute(
+                    () -> mRemoteTransitionHandler.addFiltered(filter, remoteTransition));
+        }
+
+        @Override
+        public void registerRemoteForTakeover(@NonNull TransitionFilter filter,
+                @NonNull RemoteTransition remoteTransition) {
+            mMainExecutor.execute(() -> mRemoteTransitionHandler.addFilteredForTakeover(
+                    filter, remoteTransition));
         }
 
         @Override
         public void unregisterRemote(@NonNull RemoteTransition remoteTransition) {
-            mMainExecutor.execute(() -> {
-                mRemoteTransitionHandler.removeFiltered(remoteTransition);
-            });
+            mMainExecutor.execute(
+                    () -> mRemoteTransitionHandler.removeFiltered(remoteTransition));
         }
     }
 
@@ -1470,17 +1555,23 @@
         public void registerRemote(@NonNull TransitionFilter filter,
                 @NonNull RemoteTransition remoteTransition) {
             executeRemoteCallWithTaskPermission(mTransitions, "registerRemote",
-                    (transitions) -> {
-                        transitions.mRemoteTransitionHandler.addFiltered(filter, remoteTransition);
-                    });
+                    (transitions) -> transitions.mRemoteTransitionHandler.addFiltered(
+                            filter, remoteTransition));
+        }
+
+        @Override
+        public void registerRemoteForTakeover(@NonNull TransitionFilter filter,
+                @NonNull RemoteTransition remoteTransition) {
+            executeRemoteCallWithTaskPermission(mTransitions, "registerRemoteForTakeover",
+                    (transitions) -> transitions.mRemoteTransitionHandler.addFilteredForTakeover(
+                            filter, remoteTransition));
         }
 
         @Override
         public void unregisterRemote(@NonNull RemoteTransition remoteTransition) {
             executeRemoteCallWithTaskPermission(mTransitions, "unregisterRemote",
-                    (transitions) -> {
-                        transitions.mRemoteTransitionHandler.removeFiltered(remoteTransition);
-                    });
+                    (transitions) ->
+                            transitions.mRemoteTransitionHandler.removeFiltered(remoteTransition));
         }
 
         @Override
@@ -1565,6 +1656,8 @@
             pw.println(" (" + Integer.toHexString(System.identityHashCode(handler)) + ")");
         }
 
+        mRemoteTransitionHandler.dump(pw, prefix);
+
         pw.println(prefix + "Observers:");
         for (TransitionObserver observer : mObservers) {
             pw.print(innerPrefix);
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/transition/ShellTransitionTests.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/transition/ShellTransitionTests.java
index 2366917..964d86e 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/transition/ShellTransitionTests.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/transition/ShellTransitionTests.java
@@ -73,6 +73,7 @@
 import android.os.IBinder;
 import android.os.Looper;
 import android.os.RemoteException;
+import android.platform.test.flag.junit.SetFlagsRule;
 import android.util.ArraySet;
 import android.util.Pair;
 import android.view.IRecentsAnimationRunner;
@@ -87,6 +88,7 @@
 import android.window.TransitionFilter;
 import android.window.TransitionInfo;
 import android.window.TransitionRequestInfo;
+import android.window.WindowAnimationState;
 import android.window.WindowContainerToken;
 import android.window.WindowContainerTransaction;
 
@@ -98,6 +100,7 @@
 
 import com.android.internal.R;
 import com.android.internal.policy.TransitionAnimation;
+import com.android.systemui.shared.Flags;
 import com.android.wm.shell.RootTaskDisplayAreaOrganizer;
 import com.android.wm.shell.ShellTaskOrganizer;
 import com.android.wm.shell.ShellTestCase;
@@ -114,6 +117,7 @@
 import com.android.wm.shell.util.StubTransaction;
 
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.Answers;
@@ -141,6 +145,9 @@
     private final TestTransitionHandler mDefaultHandler = new TestTransitionHandler();
     private final Handler mMainHandler = new Handler(Looper.getMainLooper());
 
+    @Rule
+    public final SetFlagsRule setFlagsRule = new SetFlagsRule();
+
     @Before
     public void setUp() {
         doAnswer(invocation -> new Binder())
@@ -475,11 +482,97 @@
     }
 
     @Test
+    public void testRegisteredRemoteTransitionTakeover() {
+        Transitions transitions = createTestTransitions();
+        transitions.replaceDefaultHandlerForTest(mDefaultHandler);
+
+        IRemoteTransition testRemote = new RemoteTransitionStub() {
+            @Override
+            public void startAnimation(IBinder token, TransitionInfo info,
+                    SurfaceControl.Transaction t,
+                    IRemoteTransitionFinishedCallback finishCallback) throws RemoteException {
+                final Transitions.TransitionHandler takeoverHandler =
+                        transitions.getHandlerForTakeover(token, info);
+
+                if (takeoverHandler == null) {
+                    finishCallback.onTransitionFinished(null /* wct */, null /* sct */);
+                    return;
+                }
+
+                takeoverHandler.takeOverAnimation(token, info, new SurfaceControl.Transaction(),
+                        wct -> {
+                            try {
+                                finishCallback.onTransitionFinished(wct, null /* sct */);
+                            } catch (RemoteException e) {
+                                // Fail
+                            }
+                        }, new WindowAnimationState[info.getChanges().size()]);
+            }
+        };
+        final boolean[] takeoverRemoteCalled = new boolean[]{false};
+        IRemoteTransition testTakeoverRemote = new RemoteTransitionStub() {
+            @Override
+            public void startAnimation(IBinder token, TransitionInfo info,
+                    SurfaceControl.Transaction t,
+                    IRemoteTransitionFinishedCallback finishCallback) {}
+
+            @Override
+            public void takeOverAnimation(IBinder transition, TransitionInfo info,
+                    SurfaceControl.Transaction startTransaction,
+                    IRemoteTransitionFinishedCallback finishCallback, WindowAnimationState[] states)
+                    throws RemoteException {
+                takeoverRemoteCalled[0] = true;
+                finishCallback.onTransitionFinished(null /* wct */, null /* sct */);
+            }
+        };
+
+        TransitionFilter filter = new TransitionFilter();
+        filter.mRequirements =
+                new TransitionFilter.Requirement[]{new TransitionFilter.Requirement()};
+        filter.mRequirements[0].mModes = new int[]{TRANSIT_OPEN, TRANSIT_TO_FRONT};
+
+        transitions.registerRemote(filter, new RemoteTransition(testRemote, "Test"));
+        transitions.registerRemoteForTakeover(
+                filter, new RemoteTransition(testTakeoverRemote, "Test"));
+        mMainExecutor.flushAll();
+
+        // Takeover shouldn't happen when the flag is disabled.
+        setFlagsRule.disableFlags(Flags.FLAG_RETURN_ANIMATION_FRAMEWORK_LIBRARY);
+        IBinder transitToken = new Binder();
+        transitions.requestStartTransition(transitToken,
+                new TransitionRequestInfo(TRANSIT_OPEN, null /* trigger */, null /* remote */));
+        TransitionInfo info = new TransitionInfoBuilder(TRANSIT_OPEN)
+                .addChange(TRANSIT_OPEN).addChange(TRANSIT_CLOSE).build();
+        transitions.onTransitionReady(transitToken, info, new StubTransaction(),
+                new StubTransaction());
+        assertEquals(0, mDefaultHandler.activeCount());
+        assertFalse(takeoverRemoteCalled[0]);
+        mDefaultHandler.finishAll();
+        mMainExecutor.flushAll();
+        verify(mOrganizer, times(1)).finishTransition(eq(transitToken), any());
+
+        // Takeover should happen when the flag is enabled.
+        setFlagsRule.enableFlags(Flags.FLAG_RETURN_ANIMATION_FRAMEWORK_LIBRARY);
+        transitions.requestStartTransition(transitToken,
+                new TransitionRequestInfo(TRANSIT_OPEN, null /* trigger */, null /* remote */));
+        info = new TransitionInfoBuilder(TRANSIT_OPEN)
+                .addChange(TRANSIT_OPEN).addChange(TRANSIT_CLOSE).build();
+        transitions.onTransitionReady(transitToken, info, new StubTransaction(),
+                new StubTransaction());
+        assertEquals(0, mDefaultHandler.activeCount());
+        assertTrue(takeoverRemoteCalled[0]);
+        mDefaultHandler.finishAll();
+        mMainExecutor.flushAll();
+        verify(mOrganizer, times(2)).finishTransition(eq(transitToken), any());
+    }
+
+    @Test
     public void testOneShotRemoteHandler() {
         Transitions transitions = createTestTransitions();
         transitions.replaceDefaultHandlerForTest(mDefaultHandler);
 
         final boolean[] remoteCalled = new boolean[]{false};
+        final boolean[] takeoverRemoteCalled = new boolean[]{false};
         final WindowContainerTransaction remoteFinishWCT = new WindowContainerTransaction();
         IRemoteTransition testRemote = new RemoteTransitionStub() {
             @Override
@@ -489,12 +582,22 @@
                 remoteCalled[0] = true;
                 finishCallback.onTransitionFinished(remoteFinishWCT, null /* sct */);
             }
+
+            @Override
+            public void takeOverAnimation(IBinder transition, TransitionInfo info,
+                    SurfaceControl.Transaction startTransaction,
+                    IRemoteTransitionFinishedCallback finishCallback, WindowAnimationState[] states)
+                    throws RemoteException {
+                takeoverRemoteCalled[0] = true;
+                finishCallback.onTransitionFinished(remoteFinishWCT, null /* sct */);
+            }
         };
 
         final int transitType = TRANSIT_FIRST_CUSTOM + 1;
 
         OneShotRemoteHandler oneShot = new OneShotRemoteHandler(mMainExecutor,
                 new RemoteTransition(testRemote, "Test"));
+
         // Verify that it responds to the remote but not other things.
         IBinder transitToken = new Binder();
         assertNotNull(oneShot.handleRequest(transitToken,
@@ -505,6 +608,7 @@
 
         Transitions.TransitionFinishCallback testFinish =
                 mock(Transitions.TransitionFinishCallback.class);
+
         // Verify that it responds to animation properly
         oneShot.setTransition(transitToken);
         IBinder anotherToken = new Binder();
@@ -514,6 +618,16 @@
         assertTrue(oneShot.startAnimation(transitToken, new TransitionInfo(transitType, 0),
                 new StubTransaction(), new StubTransaction(),
                 testFinish));
+        assertTrue(remoteCalled[0]);
+
+        // Verify that it handles takeovers properly
+        IBinder newToken = new Binder();
+        oneShot.setTransition(newToken);
+        assertFalse(oneShot.takeOverAnimation(transitToken, new TransitionInfo(transitType, 0),
+                new StubTransaction(), testFinish, new WindowAnimationState[0]));
+        assertTrue(oneShot.takeOverAnimation(newToken, new TransitionInfo(transitType, 0),
+                new StubTransaction(), testFinish, new WindowAnimationState[0]));
+        assertTrue(takeoverRemoteCalled[0]);
     }
 
     @Test