Merge "Use transient Home launch in enter/cancel transition" into main
diff --git a/core/java/android/app/ActivityOptions.java b/core/java/android/app/ActivityOptions.java
index 26f1c4b..9c279c3 100644
--- a/core/java/android/app/ActivityOptions.java
+++ b/core/java/android/app/ActivityOptions.java
@@ -2563,11 +2563,14 @@
         public static final int TYPE_LOCKSCREEN = 3;
         /** Launched from recents gesture handler. */
         public static final int TYPE_RECENTS_ANIMATION = 4;
+        /** Launched from desktop's transition handler. */
+        public static final int TYPE_DESKTOP_ANIMATION = 5;
 
         @IntDef(prefix = { "TYPE_" }, value = {
                 TYPE_LAUNCHER,
                 TYPE_NOTIFICATION,
                 TYPE_LOCKSCREEN,
+                TYPE_DESKTOP_ANIMATION
         })
         @Retention(RetentionPolicy.SOURCE)
         public @interface SourceType {}
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/RootTaskDisplayAreaOrganizer.java b/libs/WindowManager/Shell/src/com/android/wm/shell/RootTaskDisplayAreaOrganizer.java
index ab61a48..5143d41 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/RootTaskDisplayAreaOrganizer.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/RootTaskDisplayAreaOrganizer.java
@@ -115,6 +115,19 @@
         b.setParent(sc);
     }
 
+    /**
+     * Re-parents the provided surface to the leash of the provided display.
+     *
+     * @param displayId the display area to reparent to.
+     * @param sc the surface to be reparented.
+     * @param t a {@link SurfaceControl.Transaction} in which to reparent.
+     */
+    public void reparentToDisplayArea(int displayId, SurfaceControl sc,
+                                      SurfaceControl.Transaction t) {
+        final SurfaceControl displayAreaLeash = mLeashes.get(displayId);
+        t.reparent(sc, displayAreaLeash);
+    }
+
     public void setPosition(@NonNull SurfaceControl.Transaction tx, int displayId, int x, int y) {
         final SurfaceControl sc = mLeashes.get(displayId);
         if (sc == null) {
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellModule.java b/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellModule.java
index 47769a8..71bf487 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellModule.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellModule.java
@@ -59,6 +59,7 @@
 import com.android.wm.shell.desktopmode.DesktopModeStatus;
 import com.android.wm.shell.desktopmode.DesktopModeTaskRepository;
 import com.android.wm.shell.desktopmode.DesktopTasksController;
+import com.android.wm.shell.desktopmode.DragToDesktopTransitionHandler;
 import com.android.wm.shell.desktopmode.EnterDesktopTaskTransitionHandler;
 import com.android.wm.shell.desktopmode.ExitDesktopTaskTransitionHandler;
 import com.android.wm.shell.desktopmode.ToggleResizeDesktopTaskTransitionHandler;
@@ -498,6 +499,7 @@
             EnterDesktopTaskTransitionHandler enterDesktopTransitionHandler,
             ExitDesktopTaskTransitionHandler exitDesktopTransitionHandler,
             ToggleResizeDesktopTaskTransitionHandler toggleResizeDesktopTaskTransitionHandler,
+            DragToDesktopTransitionHandler dragToDesktopTransitionHandler,
             @DynamicOverride DesktopModeTaskRepository desktopModeTaskRepository,
             LaunchAdjacentController launchAdjacentController,
             RecentsTransitionHandler recentsTransitionHandler,
@@ -506,8 +508,19 @@
         return new DesktopTasksController(context, shellInit, shellCommandHandler, shellController,
                 displayController, shellTaskOrganizer, syncQueue, rootTaskDisplayAreaOrganizer,
                 transitions, enterDesktopTransitionHandler, exitDesktopTransitionHandler,
-                toggleResizeDesktopTaskTransitionHandler, desktopModeTaskRepository,
-                launchAdjacentController, recentsTransitionHandler, mainExecutor);
+                toggleResizeDesktopTaskTransitionHandler, dragToDesktopTransitionHandler,
+                desktopModeTaskRepository, launchAdjacentController, recentsTransitionHandler,
+                mainExecutor);
+    }
+
+    @WMSingleton
+    @Provides
+    static DragToDesktopTransitionHandler provideDragToDesktopTransitionHandler(
+            Context context,
+            Transitions transitions,
+            RootTaskDisplayAreaOrganizer rootTaskDisplayAreaOrganizer) {
+        return new DragToDesktopTransitionHandler(context, transitions,
+                rootTaskDisplayAreaOrganizer);
     }
 
     @WMSingleton
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksController.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksController.kt
index 8e12991..4a9ea6f 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksController.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksController.kt
@@ -59,6 +59,7 @@
 import com.android.wm.shell.common.split.SplitScreenConstants.SPLIT_POSITION_TOP_OR_LEFT
 import com.android.wm.shell.desktopmode.DesktopModeTaskRepository.VisibleTasksListener
 import com.android.wm.shell.desktopmode.DesktopModeVisualIndicator.TO_DESKTOP_INDICATOR
+import com.android.wm.shell.desktopmode.DragToDesktopTransitionHandler.DragToDesktopStateListener
 import com.android.wm.shell.protolog.ShellProtoLogGroup.WM_SHELL_DESKTOP_MODE
 import com.android.wm.shell.recents.RecentsTransitionHandler
 import com.android.wm.shell.recents.RecentsTransitionStateListener
@@ -92,6 +93,7 @@
         private val exitDesktopTaskTransitionHandler: ExitDesktopTaskTransitionHandler,
         private val toggleResizeDesktopTaskTransitionHandler:
         ToggleResizeDesktopTaskTransitionHandler,
+        private val dragToDesktopTransitionHandler: DragToDesktopTransitionHandler,
         private val desktopModeTaskRepository: DesktopModeTaskRepository,
         private val launchAdjacentController: LaunchAdjacentController,
         private val recentsTransitionHandler: RecentsTransitionHandler,
@@ -110,6 +112,20 @@
             launchAdjacentController.launchAdjacentEnabled = !hasVisibleFreeformTasks
         }
     }
+    private val dragToDesktopStateListener = object : DragToDesktopStateListener {
+        override fun onCommitToDesktopAnimationStart(tx: SurfaceControl.Transaction) {
+            removeVisualIndicator(tx)
+        }
+
+        override fun onCancelToDesktopAnimationEnd(tx: SurfaceControl.Transaction) {
+            removeVisualIndicator(tx)
+        }
+
+        private fun removeVisualIndicator(tx: SurfaceControl.Transaction) {
+            visualIndicator?.releaseVisualIndicator(tx)
+            visualIndicator = null
+        }
+    }
 
     private val transitionAreaHeight
         get() = context.resources.getDimensionPixelSize(
@@ -122,9 +138,7 @@
         )
 
     private var recentsAnimationRunning = false
-
-    // This is public to avoid cyclic dependency; it is set by SplitScreenController
-    lateinit var splitScreenController: SplitScreenController
+    private lateinit var splitScreenController: SplitScreenController
 
     init {
         desktopMode = DesktopModeImpl()
@@ -143,7 +157,7 @@
         )
         transitions.addHandler(this)
         desktopModeTaskRepository.addVisibleTasksListener(taskVisibilityListener, mainExecutor)
-
+        dragToDesktopTransitionHandler.setDragToDesktopStateListener(dragToDesktopStateListener)
         recentsTransitionHandler.addTransitionStateListener(
             object : RecentsTransitionStateListener {
                 override fun onAnimationStateChanged(running: Boolean) {
@@ -158,6 +172,12 @@
         )
     }
 
+    /** Setter needed to avoid cyclic dependency. */
+    fun setSplitScreenController(controller: SplitScreenController) {
+        splitScreenController = controller
+        dragToDesktopTransitionHandler.setSplitScreenController(controller)
+    }
+
     /** Show all tasks, that are part of the desktop, on top of launcher */
     fun showDesktopApps(displayId: Int, remoteTransition: RemoteTransition? = null) {
         KtProtoLog.v(WM_SHELL_DESKTOP_MODE, "DesktopTasksController: showDesktopApps")
@@ -248,56 +268,43 @@
     }
 
     /**
-     * The first part of the animated move to desktop transition. Applies the changes to move task
-     * to desktop mode and sets the taskBounds to the passed in bounds, startBounds. This is
-     * followed with a call to {@link finishMoveToDesktop} or {@link cancelMoveToDesktop}.
+     * The first part of the animated drag to desktop transition. This is
+     * followed with a call to [finalizeDragToDesktop] or [cancelDragToDesktop].
      */
-    fun startMoveToDesktop(
+    fun startDragToDesktop(
             taskInfo: RunningTaskInfo,
-            startBounds: Rect,
-            dragToDesktopValueAnimator: MoveToDesktopAnimator
+            dragToDesktopValueAnimator: MoveToDesktopAnimator,
+            windowDecor: DesktopModeWindowDecoration
     ) {
         KtProtoLog.v(
-            WM_SHELL_DESKTOP_MODE,
-            "DesktopTasksController: startMoveToDesktop taskId=%d",
-            taskInfo.taskId
+                WM_SHELL_DESKTOP_MODE,
+                "DesktopTasksController: startDragToDesktop taskId=%d",
+                taskInfo.taskId
+        )
+        dragToDesktopTransitionHandler.startDragToDesktopTransition(
+                taskInfo.taskId,
+                dragToDesktopValueAnimator,
+                windowDecor
+        )
+    }
+
+    /**
+     * The second part of the animated drag to desktop transition, called after
+     * [startDragToDesktop].
+     */
+    private fun finalizeDragToDesktop(taskInfo: RunningTaskInfo, freeformBounds: Rect) {
+        KtProtoLog.v(
+                WM_SHELL_DESKTOP_MODE,
+                "DesktopTasksController: finalizeDragToDesktop taskId=%d",
+                taskInfo.taskId
         )
         val wct = WindowContainerTransaction()
         exitSplitIfApplicable(wct, taskInfo)
         moveHomeTaskToFront(wct)
-        addMoveToDesktopChanges(wct, taskInfo)
-        wct.setBounds(taskInfo.token, startBounds)
-
-        if (Transitions.ENABLE_SHELL_TRANSITIONS) {
-            enterDesktopTaskTransitionHandler.startMoveToDesktop(wct, dragToDesktopValueAnimator,
-                    mOnAnimationFinishedCallback)
-        } else {
-            shellTaskOrganizer.applyTransaction(wct)
-        }
-    }
-
-    /**
-     * The second part of the animated move to desktop transition, called after
-     * {@link startMoveToDesktop}. Brings apps to front and sets freeform task bounds.
-     */
-    private fun finalizeMoveToDesktop(taskInfo: RunningTaskInfo, freeformBounds: Rect) {
-        KtProtoLog.v(
-            WM_SHELL_DESKTOP_MODE,
-            "DesktopTasksController: finalizeMoveToDesktop taskId=%d",
-            taskInfo.taskId
-        )
-        val wct = WindowContainerTransaction()
         bringDesktopAppsToFront(taskInfo.displayId, wct)
         addMoveToDesktopChanges(wct, taskInfo)
         wct.setBounds(taskInfo.token, freeformBounds)
-
-        if (Transitions.ENABLE_SHELL_TRANSITIONS) {
-            enterDesktopTaskTransitionHandler.finalizeMoveToDesktop(wct,
-                    mOnAnimationFinishedCallback)
-        } else {
-            shellTaskOrganizer.applyTransaction(wct)
-            releaseVisualIndicator()
-        }
+        dragToDesktopTransitionHandler.finishDragToDesktopTransition(wct)
     }
 
     /**
@@ -353,40 +360,40 @@
     }
 
     private fun exitSplitIfApplicable(wct: WindowContainerTransaction, taskInfo: RunningTaskInfo) {
-        if (taskInfo.windowingMode == WINDOWING_MODE_MULTI_WINDOW) {
-            splitScreenController.prepareExitSplitScreen(wct,
-                splitScreenController.getStageOfTask(taskInfo.taskId), EXIT_REASON_ENTER_DESKTOP)
+        if (splitScreenController.isTaskInSplitScreen(taskInfo.taskId)) {
+            splitScreenController.prepareExitSplitScreen(
+                    wct,
+                    splitScreenController.getStageOfTask(taskInfo.taskId),
+                    EXIT_REASON_ENTER_DESKTOP
+            )
+            getOtherSplitTask(taskInfo.taskId)?.let { otherTaskInfo ->
+                wct.removeTask(otherTaskInfo.token)
+            }
         }
     }
 
+    private fun getOtherSplitTask(taskId: Int): RunningTaskInfo? {
+        val remainingTaskPosition: Int =
+                if (splitScreenController.getSplitPosition(taskId)
+                        == SPLIT_POSITION_BOTTOM_OR_RIGHT) {
+                    SPLIT_POSITION_TOP_OR_LEFT
+                } else {
+                    SPLIT_POSITION_BOTTOM_OR_RIGHT
+                }
+        return splitScreenController.getTaskInfo(remainingTaskPosition)
+    }
+
     /**
-     * The second part of the animated move to desktop transition, called after
-     * {@link startMoveToDesktop}. Move a task to fullscreen after being dragged from fullscreen
-     * and released back into status bar area.
+     * The second part of the animated drag to desktop transition, called after
+     * [startDragToDesktop].
      */
-    fun cancelMoveToDesktop(task: RunningTaskInfo, moveToDesktopAnimator: MoveToDesktopAnimator) {
+    fun cancelDragToDesktop(task: RunningTaskInfo) {
         KtProtoLog.v(
             WM_SHELL_DESKTOP_MODE,
-            "DesktopTasksController: cancelMoveToDesktop taskId=%d",
+            "DesktopTasksController: cancelDragToDesktop taskId=%d",
             task.taskId
         )
-        val wct = WindowContainerTransaction()
-        wct.setBounds(task.token, Rect())
-
-        if (Transitions.ENABLE_SHELL_TRANSITIONS) {
-            enterDesktopTaskTransitionHandler.startCancelMoveToDesktopMode(wct,
-                    moveToDesktopAnimator) { t ->
-                val callbackWCT = WindowContainerTransaction()
-                visualIndicator?.releaseVisualIndicator(t)
-                visualIndicator = null
-                addMoveToFullscreenChanges(callbackWCT, task)
-                transitions.startTransition(TRANSIT_CHANGE, callbackWCT, null /* handler */)
-            }
-        } else {
-            addMoveToFullscreenChanges(wct, task)
-            shellTaskOrganizer.applyTransaction(wct)
-            releaseVisualIndicator()
-        }
+        dragToDesktopTransitionHandler.cancelDragToDesktopTransition()
     }
 
     private fun moveToFullscreenWithAnimation(task: RunningTaskInfo, position: Point) {
@@ -966,6 +973,11 @@
             visualIndicator = DesktopModeVisualIndicator(syncQueue, taskInfo,
                     displayController, context, taskSurface, shellTaskOrganizer,
                     rootTaskDisplayAreaOrganizer, TO_DESKTOP_INDICATOR)
+            // TODO(b/301106941): don't show the indicator until the drag-to-desktop animation has
+            // started, or it'll be visible too early on top of the task surface, especially in
+            // the cancel-early case. Also because it shouldn't even be shown in the cancel-early
+            // case since its dismissal is tied to the cancel animation end, which doesn't even run
+            // in cancel-early.
             visualIndicator?.createIndicatorWithAnimatedBounds()
         }
         val indicator = visualIndicator ?: return
@@ -988,7 +1000,7 @@
             taskInfo: RunningTaskInfo,
             freeformBounds: Rect
     ) {
-        finalizeMoveToDesktop(taskInfo, freeformBounds)
+        finalizeDragToDesktop(taskInfo, freeformBounds)
     }
 
     private fun getStatusBarHeight(taskInfo: RunningTaskInfo): Int {
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DragToDesktopTransitionHandler.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DragToDesktopTransitionHandler.kt
new file mode 100644
index 0000000..75d27d9
--- /dev/null
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DragToDesktopTransitionHandler.kt
@@ -0,0 +1,543 @@
+package com.android.wm.shell.desktopmode
+
+import android.animation.Animator
+import android.animation.AnimatorListenerAdapter
+import android.animation.RectEvaluator
+import android.animation.ValueAnimator
+import android.app.ActivityOptions
+import android.app.ActivityOptions.SourceInfo
+import android.app.PendingIntent
+import android.app.PendingIntent.FLAG_ALLOW_UNSAFE_IMPLICIT_INTENT
+import android.app.PendingIntent.FLAG_MUTABLE
+import android.app.WindowConfiguration.ACTIVITY_TYPE_HOME
+import android.app.WindowConfiguration.WINDOWING_MODE_FREEFORM
+import android.content.Context
+import android.content.Intent
+import android.content.Intent.FILL_IN_COMPONENT
+import android.graphics.Rect
+import android.os.IBinder
+import android.os.SystemClock
+import android.view.SurfaceControl
+import android.view.WindowManager.TRANSIT_CLOSE
+import android.window.TransitionInfo
+import android.window.TransitionInfo.Change
+import android.window.TransitionRequestInfo
+import android.window.WindowContainerToken
+import android.window.WindowContainerTransaction
+import com.android.wm.shell.RootTaskDisplayAreaOrganizer
+import com.android.wm.shell.splitscreen.SplitScreenController
+import com.android.wm.shell.transition.Transitions
+import com.android.wm.shell.transition.Transitions.TRANSIT_DESKTOP_MODE_CANCEL_DRAG_TO_DESKTOP
+import com.android.wm.shell.transition.Transitions.TRANSIT_DESKTOP_MODE_END_DRAG_TO_DESKTOP
+import com.android.wm.shell.transition.Transitions.TRANSIT_DESKTOP_MODE_START_DRAG_TO_DESKTOP
+import com.android.wm.shell.transition.Transitions.TransitionHandler
+import com.android.wm.shell.util.TransitionUtil
+import com.android.wm.shell.windowdecor.DesktopModeWindowDecoration
+import com.android.wm.shell.windowdecor.MoveToDesktopAnimator
+import com.android.wm.shell.windowdecor.MoveToDesktopAnimator.Companion.DRAG_FREEFORM_SCALE
+import java.util.function.Supplier
+
+/**
+ * Handles the transition to enter desktop from fullscreen by dragging on the handle bar. It also
+ * handles the cancellation case where the task is dragged back to the status bar area in the same
+ * gesture.
+ */
+class DragToDesktopTransitionHandler(
+        private val context: Context,
+        private val transitions: Transitions,
+        private val taskDisplayAreaOrganizer: RootTaskDisplayAreaOrganizer,
+        private val transactionSupplier: Supplier<SurfaceControl.Transaction>
+) : TransitionHandler {
+
+    constructor(
+            context: Context,
+            transitions: Transitions,
+            rootTaskDisplayAreaOrganizer: RootTaskDisplayAreaOrganizer
+    ) : this(
+            context,
+            transitions,
+            rootTaskDisplayAreaOrganizer,
+            Supplier { SurfaceControl.Transaction() }
+    )
+
+    private val rectEvaluator = RectEvaluator(Rect())
+    private val launchHomeIntent = Intent(Intent.ACTION_MAIN)
+            .addCategory(Intent.CATEGORY_HOME)
+
+    private var dragToDesktopStateListener: DragToDesktopStateListener? = null
+    private var splitScreenController: SplitScreenController? = null
+    private var transitionState: TransitionState? = null
+
+    /** Sets a listener to receive callback about events during the transition animation. */
+    fun setDragToDesktopStateListener(listener: DragToDesktopStateListener) {
+        dragToDesktopStateListener = listener
+    }
+
+    /** Setter needed to avoid cyclic dependency. */
+    fun setSplitScreenController(controller: SplitScreenController) {
+        splitScreenController = controller
+    }
+
+    /**
+     * Starts a transition that performs a transient launch of Home so that Home is brought to the
+     * front while still keeping the currently focused task that is being dragged resumed. This
+     * allows the animation handler to reorder the task to the front and to scale it with the
+     * gesture into the desktop area with the Home and wallpaper behind it.
+     *
+     * Note that the transition handler for this transition doesn't call the finish callback until
+     * after one of the "end" or "cancel" transitions is merged into this transition.
+     */
+    fun startDragToDesktopTransition(
+            taskId: Int,
+            dragToDesktopAnimator: MoveToDesktopAnimator,
+            windowDecoration: DesktopModeWindowDecoration
+    ) {
+        if (transitionState != null) {
+            error("A drag to desktop is already in progress")
+        }
+
+        val options = ActivityOptions.makeBasic().apply {
+            setTransientLaunch()
+            setSourceInfo(SourceInfo.TYPE_DESKTOP_ANIMATION, SystemClock.uptimeMillis())
+        }
+        val pendingIntent = PendingIntent.getActivity(
+                context,
+                0 /* requestCode */,
+                launchHomeIntent,
+                FLAG_MUTABLE or FLAG_ALLOW_UNSAFE_IMPLICIT_INTENT or FILL_IN_COMPONENT
+        )
+        val wct = WindowContainerTransaction()
+        wct.sendPendingIntent(pendingIntent, launchHomeIntent, options.toBundle())
+        val startTransitionToken = transitions
+                .startTransition(TRANSIT_DESKTOP_MODE_START_DRAG_TO_DESKTOP, wct, this)
+
+        transitionState = if (isSplitTask(taskId)) {
+            TransitionState.FromSplit(
+                    draggedTaskId = taskId,
+                    dragAnimator = dragToDesktopAnimator,
+                    windowDecoration = windowDecoration,
+                    startTransitionToken = startTransitionToken
+            )
+        } else {
+            TransitionState.FromFullscreen(
+                    draggedTaskId = taskId,
+                    dragAnimator = dragToDesktopAnimator,
+                    windowDecoration = windowDecoration,
+                    startTransitionToken = startTransitionToken
+            )
+        }
+    }
+
+    /**
+     * Starts a transition that "finishes" the drag to desktop gesture. This transition is intended
+     * to merge into the "start" transition and is the one that actually applies the bounds and
+     * windowing mode changes to the dragged task. This is called when the dragged task is released
+     * inside the desktop drop zone.
+     */
+    fun finishDragToDesktopTransition(wct: WindowContainerTransaction) {
+        transitions.startTransition(TRANSIT_DESKTOP_MODE_END_DRAG_TO_DESKTOP, wct, this)
+    }
+
+    /**
+     * Starts a transition that "cancels" the drag to desktop gesture. This transition is intended
+     * to merge into the "start" transition and it restores the transient state that was used to
+     * launch the Home task over the dragged task. This is called when the dragged task is released
+     * outside the desktop drop zone and is instead dropped back into the status bar region that
+     * means the user wants to remain in their current windowing mode.
+     */
+    fun cancelDragToDesktopTransition() {
+        val state = requireTransitionState()
+        state.cancelled = true
+        if (state.draggedTaskChange != null) {
+            // Regular case, transient launch of Home happened as is waiting for the cancel
+            // transient to start and merge. Animate the cancellation (scale back to original
+            // bounds) first before actually starting the cancel transition so that the wallpaper
+            // is visible behind the animating task.
+            startCancelAnimation()
+        } else {
+            // There's no dragged task, this can happen when the "cancel" happened too quickly
+            // before the "start" transition is even ready (like on a fling gesture). The
+            // "shrink" animation didn't even start, so there's no need to animate the "cancel".
+            // We also don't want to start the cancel transition yet since we don't have
+            // enough info to restore the order. We'll check for the cancelled state flag when
+            // the "start" animation is ready and cancel from #startAnimation instead.
+        }
+    }
+
+    override fun startAnimation(
+            transition: IBinder,
+            info: TransitionInfo,
+            startTransaction: SurfaceControl.Transaction,
+            finishTransaction: SurfaceControl.Transaction,
+            finishCallback: Transitions.TransitionFinishCallback
+    ): Boolean {
+        val state = requireTransitionState()
+
+        val isStartDragToDesktop = info.type == TRANSIT_DESKTOP_MODE_START_DRAG_TO_DESKTOP &&
+                transition == state.startTransitionToken
+        if (!isStartDragToDesktop) {
+            return false
+        }
+
+        // Layering: non-wallpaper, non-home tasks excluding the dragged task go at the bottom,
+        // then Home on top of that, wallpaper on top of that and finally the dragged task on top
+        // of everything.
+        val appLayers = info.changes.size
+        val homeLayers = info.changes.size * 2
+        val wallpaperLayers = info.changes.size * 3
+        val dragLayer = wallpaperLayers
+        val leafTaskFilter = TransitionUtil.LeafTaskFilter()
+        info.changes.withIndex().forEach { (i, change) ->
+            if (TransitionUtil.isWallpaper(change)) {
+                val layer = wallpaperLayers - i
+                startTransaction.apply {
+                    setLayer(change.leash, layer)
+                    show(change.leash)
+                }
+            } else if (isHomeChange(change)) {
+                state.homeToken = change.container
+                val layer = homeLayers - i
+                startTransaction.apply {
+                    setLayer(change.leash, layer)
+                    show(change.leash)
+                }
+            } else if (TransitionInfo.isIndependent(change, info)) {
+                // Root.
+                when (state) {
+                    is TransitionState.FromSplit -> {
+                        state.splitRootChange = change
+                        val layer = if (!state.cancelled) {
+                            // Normal case, split root goes to the bottom behind everything else.
+                            appLayers - i
+                        } else {
+                            // Cancel-early case, pretend nothing happened so split root stays top.
+                            dragLayer
+                        }
+                        startTransaction.apply {
+                            setLayer(change.leash, layer)
+                            show(change.leash)
+                        }
+                    }
+                    is TransitionState.FromFullscreen -> {
+                        if (change.taskInfo?.taskId == state.draggedTaskId) {
+                            state.draggedTaskChange = change
+                            val bounds = change.endAbsBounds
+                            startTransaction.apply {
+                                setLayer(change.leash, dragLayer)
+                                setWindowCrop(change.leash, bounds.width(), bounds.height())
+                                show(change.leash)
+                            }
+                        } else {
+                            throw IllegalStateException("Expected root to be dragged task")
+                        }
+                    }
+                }
+            } else if (leafTaskFilter.test(change)) {
+                // When dragging one of the split tasks, the dragged leaf needs to be re-parented
+                // so that it can be layered separately from the rest of the split root/stages.
+                // The split root including the other split side was layered behind the wallpaper
+                // and home while the dragged split needs to be layered in front of them.
+                // Do not do this in the cancel-early case though, since in that case nothing should
+                // happen on screen so the layering will remain the same as if no transition
+                // occurred.
+                if (change.taskInfo?.taskId == state.draggedTaskId && !state.cancelled) {
+                    state.draggedTaskChange = change
+                    taskDisplayAreaOrganizer.reparentToDisplayArea(
+                            change.endDisplayId, change.leash, startTransaction)
+                    val bounds = change.endAbsBounds
+                    startTransaction.apply {
+                        setLayer(change.leash, dragLayer)
+                        setWindowCrop(change.leash, bounds.width(), bounds.height())
+                        show(change.leash)
+                    }
+                }
+            }
+        }
+        state.startTransitionFinishCb = finishCallback
+        state.startTransitionFinishTransaction = finishTransaction
+        startTransaction.apply()
+
+        if (!state.cancelled) {
+            // Normal case, start animation to scale down the dragged task. It'll also be moved to
+            // follow the finger and when released we'll start the next phase/transition.
+            state.dragAnimator.startAnimation()
+        } else {
+            // Cancel-early case, the state was flagged was cancelled already, which means the
+            // gesture ended in the cancel region. This can happen even before the start transition
+            // is ready/animate here when cancelling quickly like with a fling. There's no point
+            // in starting the scale down animation that we would scale up anyway, so just jump
+            // directly into starting the cancel transition to restore WM order. Surfaces should
+            // not move as if no transition happened.
+            startCancelDragToDesktopTransition()
+        }
+        return true
+    }
+
+    override fun mergeAnimation(
+            transition: IBinder,
+            info: TransitionInfo,
+            t: SurfaceControl.Transaction,
+            mergeTarget: IBinder,
+            finishCallback: Transitions.TransitionFinishCallback
+    ) {
+        val state = requireTransitionState()
+        val isCancelTransition = info.type == TRANSIT_DESKTOP_MODE_CANCEL_DRAG_TO_DESKTOP &&
+                transition == state.cancelTransitionToken &&
+                mergeTarget == state.startTransitionToken
+        val isEndTransition = info.type == TRANSIT_DESKTOP_MODE_END_DRAG_TO_DESKTOP &&
+                mergeTarget == state.startTransitionToken
+
+        val startTransactionFinishT = state.startTransitionFinishTransaction
+                ?: error("Start transition expected to be waiting for merge but wasn't")
+        val startTransitionFinishCb = state.startTransitionFinishCb
+                ?: error("Start transition expected to be waiting for merge but wasn't")
+        if (isEndTransition) {
+            info.changes.withIndex().forEach { (i, change) ->
+                if (change.mode == TRANSIT_CLOSE) {
+                    t.hide(change.leash)
+                    startTransactionFinishT.hide(change.leash)
+                } else if (change.taskInfo?.taskId == state.draggedTaskId) {
+                    t.show(change.leash)
+                    startTransactionFinishT.show(change.leash)
+                    state.draggedTaskChange = change
+                } else if (change.taskInfo?.windowingMode == WINDOWING_MODE_FREEFORM) {
+                    // Other freeform tasks that are being restored go behind the dragged task.
+                    val draggedTaskLeash = state.draggedTaskChange?.leash
+                            ?: error("Expected dragged leash to be non-null")
+                    t.setRelativeLayer(change.leash, draggedTaskLeash, -i)
+                    startTransactionFinishT.setRelativeLayer(change.leash, draggedTaskLeash, -i)
+                }
+            }
+
+            val draggedTaskChange = state.draggedTaskChange
+                    ?: throw IllegalStateException("Expected non-null change of dragged task")
+            val draggedTaskLeash = draggedTaskChange.leash
+            val startBounds = draggedTaskChange.startAbsBounds
+            val endBounds = draggedTaskChange.endAbsBounds
+
+            // TODO(b/301106941): Instead of forcing-finishing the animation that scales the
+            //  surface down and then starting another that scales it back up to the final size,
+            //  blend the two animations.
+            state.dragAnimator.endAnimator()
+            // Using [DRAG_FREEFORM_SCALE] to calculate animated width/height is possible because
+            // it is known that the animation scale is finished because the animation was
+            // force-ended above. This won't be true when the two animations are blended.
+            val animStartWidth = (startBounds.width() * DRAG_FREEFORM_SCALE).toInt()
+            val animStartHeight = (startBounds.height() * DRAG_FREEFORM_SCALE).toInt()
+            // Using end bounds here to find the left/top also assumes the center animation has
+            // finished and the surface is placed exactly in the center of the screen which matches
+            // the end/default bounds of the now freeform task.
+            val animStartLeft = endBounds.centerX() - (animStartWidth / 2)
+            val animStartTop = endBounds.centerY() - (animStartHeight / 2)
+            val animStartBounds = Rect(
+                    animStartLeft,
+                    animStartTop,
+                    animStartLeft + animStartWidth,
+                    animStartTop + animStartHeight
+            )
+
+
+            dragToDesktopStateListener?.onCommitToDesktopAnimationStart(t)
+            t.apply {
+                setScale(draggedTaskLeash, 1f, 1f)
+                setPosition(
+                        draggedTaskLeash,
+                        animStartBounds.left.toFloat(),
+                        animStartBounds.top.toFloat()
+                )
+                setWindowCrop(
+                        draggedTaskLeash,
+                        animStartBounds.width(),
+                        animStartBounds.height()
+                )
+            }
+            // Accept the merge by applying the merging transaction (applied by #showResizeVeil)
+            // and finish callback. Show the veil and position the task at the first frame before
+            // starting the final animation.
+            state.windowDecoration.showResizeVeil(t, animStartBounds)
+            finishCallback.onTransitionFinished(null /* wct */)
+
+            // Because the task surface was scaled down during the drag, we must use the animated
+            // bounds instead of the [startAbsBounds].
+            val tx: SurfaceControl.Transaction = transactionSupplier.get()
+            ValueAnimator.ofObject(rectEvaluator, animStartBounds, endBounds)
+                    .setDuration(DRAG_TO_DESKTOP_FINISH_ANIM_DURATION_MS)
+                    .apply {
+                        addUpdateListener { animator ->
+                            val animBounds = animator.animatedValue as Rect
+                            tx.apply {
+                                setScale(draggedTaskLeash, 1f, 1f)
+                                 setPosition(
+                                         draggedTaskLeash,
+                                         animBounds.left.toFloat(),
+                                         animBounds.top.toFloat()
+                                 )
+                                setWindowCrop(
+                                        draggedTaskLeash,
+                                        animBounds.width(),
+                                        animBounds.height()
+                                )
+                            }
+                            state.windowDecoration.updateResizeVeil(tx, animBounds)
+                        }
+                        addListener(object : AnimatorListenerAdapter() {
+                            override fun onAnimationEnd(animation: Animator) {
+                                state.windowDecoration.hideResizeVeil()
+                                startTransitionFinishCb.onTransitionFinished(null /* null */)
+                                clearState()
+                            }
+                        })
+                        start()
+                    }
+        } else if (isCancelTransition) {
+            info.changes.forEach { change ->
+                t.show(change.leash)
+                startTransactionFinishT.show(change.leash)
+            }
+            t.apply()
+            finishCallback.onTransitionFinished(null /* wct */)
+            startTransitionFinishCb.onTransitionFinished(null /* wct */)
+            clearState()
+        }
+    }
+
+    override fun handleRequest(
+            transition: IBinder,
+            request: TransitionRequestInfo
+    ): WindowContainerTransaction? {
+        // Only handle transitions started from shell.
+        return null
+    }
+
+    private fun isHomeChange(change: Change): Boolean {
+        return change.taskInfo?.activityType == ACTIVITY_TYPE_HOME
+    }
+
+    private fun startCancelAnimation() {
+        val state = requireTransitionState()
+        val dragToDesktopAnimator = state.dragAnimator
+
+        val draggedTaskChange = state.draggedTaskChange
+                ?: throw IllegalStateException("Expected non-null task change")
+        val sc = draggedTaskChange.leash
+        // TODO(b/301106941): Don't end the animation and start one to scale it back, merge them
+        //  instead.
+        // End the animation that shrinks the window when task is first dragged from fullscreen
+        dragToDesktopAnimator.endAnimator()
+        // Then animate the scaled window back to its original bounds.
+        val x: Float = dragToDesktopAnimator.position.x
+        val y: Float = dragToDesktopAnimator.position.y
+        val targetX = draggedTaskChange.endAbsBounds.left
+        val targetY = draggedTaskChange.endAbsBounds.top
+        val dx = targetX - x
+        val dy = targetY - y
+        val tx: SurfaceControl.Transaction = transactionSupplier.get()
+        ValueAnimator.ofFloat(DRAG_FREEFORM_SCALE, 1f)
+                .setDuration(DRAG_TO_DESKTOP_FINISH_ANIM_DURATION_MS)
+                .apply {
+                    addUpdateListener { animator ->
+                        val scale = animator.animatedValue as Float
+                        val fraction = animator.animatedFraction
+                        val animX = x + (dx * fraction)
+                        val animY = y + (dy * fraction)
+                        tx.apply {
+                            setPosition(sc, animX, animY)
+                            setScale(sc, scale, scale)
+                            show(sc)
+                            apply()
+                        }
+                    }
+                    addListener(object : AnimatorListenerAdapter() {
+                        override fun onAnimationEnd(animation: Animator) {
+                            dragToDesktopStateListener?.onCancelToDesktopAnimationEnd(tx)
+                            // Start the cancel transition to restore order.
+                            startCancelDragToDesktopTransition()
+                        }
+                    })
+                    start()
+                }
+    }
+
+    private fun startCancelDragToDesktopTransition() {
+        val state = requireTransitionState()
+        val wct = WindowContainerTransaction()
+        when (state) {
+            is TransitionState.FromFullscreen -> {
+                val wc = state.draggedTaskChange?.container
+                        ?: error("Dragged task should be non-null before cancelling")
+                wct.reorder(wc, true /* toTop */)
+            }
+            is TransitionState.FromSplit -> {
+                val wc = state.splitRootChange?.container
+                        ?: error("Split root should be non-null before cancelling")
+                wct.reorder(wc, true /* toTop */)
+            }
+        }
+        val homeWc = state.homeToken ?: error("Home task should be non-null before cancelling")
+        wct.restoreTransientOrder(homeWc)
+
+        state.cancelTransitionToken = transitions.startTransition(
+                TRANSIT_DESKTOP_MODE_CANCEL_DRAG_TO_DESKTOP, wct, this)
+    }
+
+    private fun clearState() {
+        transitionState = null
+    }
+
+    private fun isSplitTask(taskId: Int): Boolean {
+        return splitScreenController?.isTaskInSplitScreen(taskId) ?: false
+    }
+
+    private fun requireTransitionState(): TransitionState {
+        return transitionState ?: error("Expected non-null transition state")
+    }
+
+    interface DragToDesktopStateListener {
+        fun onCommitToDesktopAnimationStart(tx: SurfaceControl.Transaction)
+        fun onCancelToDesktopAnimationEnd(tx: SurfaceControl.Transaction)
+    }
+
+    sealed class TransitionState {
+        abstract val draggedTaskId: Int
+        abstract val dragAnimator: MoveToDesktopAnimator
+        abstract val windowDecoration: DesktopModeWindowDecoration
+        abstract val startTransitionToken: IBinder
+        abstract var startTransitionFinishCb: Transitions.TransitionFinishCallback?
+        abstract var startTransitionFinishTransaction: SurfaceControl.Transaction?
+        abstract var cancelTransitionToken: IBinder?
+        abstract var homeToken: WindowContainerToken?
+        abstract var draggedTaskChange: Change?
+        abstract var cancelled: Boolean
+
+        data class FromFullscreen(
+                override val draggedTaskId: Int,
+                override val dragAnimator: MoveToDesktopAnimator,
+                override val windowDecoration: DesktopModeWindowDecoration,
+                override val startTransitionToken: IBinder,
+                override var startTransitionFinishCb: Transitions.TransitionFinishCallback? = null,
+                override var startTransitionFinishTransaction: SurfaceControl.Transaction? = null,
+                override var cancelTransitionToken: IBinder? = null,
+                override var homeToken: WindowContainerToken? = null,
+                override var draggedTaskChange: Change? = null,
+                override var cancelled: Boolean = false,
+        ) : TransitionState()
+        data class FromSplit(
+                override val draggedTaskId: Int,
+                override val dragAnimator: MoveToDesktopAnimator,
+                override val windowDecoration: DesktopModeWindowDecoration,
+                override val startTransitionToken: IBinder,
+                override var startTransitionFinishCb: Transitions.TransitionFinishCallback? = null,
+                override var startTransitionFinishTransaction: SurfaceControl.Transaction? = null,
+                override var cancelTransitionToken: IBinder? = null,
+                override var homeToken: WindowContainerToken? = null,
+                override var draggedTaskChange: Change? = null,
+                override var cancelled: Boolean = false,
+                var splitRootChange: Change? = null,
+        ) : TransitionState()
+    }
+
+    companion object {
+        /** The duration of the animation to commit or cancel the drag-to-desktop gesture. */
+        private const val DRAG_TO_DESKTOP_FINISH_ANIM_DURATION_MS = 336L
+    }
+}
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/EnterDesktopTaskTransitionHandler.java b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/EnterDesktopTaskTransitionHandler.java
index 024465b..605600f 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/EnterDesktopTaskTransitionHandler.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/EnterDesktopTaskTransitionHandler.java
@@ -18,12 +18,13 @@
 
 import static android.app.WindowConfiguration.WINDOWING_MODE_FREEFORM;
 
+import static com.android.wm.shell.transition.Transitions.TRANSIT_MOVE_TO_DESKTOP;
+
 import android.animation.Animator;
 import android.animation.AnimatorListenerAdapter;
 import android.animation.RectEvaluator;
 import android.animation.ValueAnimator;
 import android.app.ActivityManager;
-import android.graphics.PointF;
 import android.graphics.Rect;
 import android.os.IBinder;
 import android.util.Slog;
@@ -38,11 +39,9 @@
 
 import com.android.wm.shell.transition.Transitions;
 import com.android.wm.shell.windowdecor.DesktopModeWindowDecoration;
-import com.android.wm.shell.windowdecor.MoveToDesktopAnimator;
 
 import java.util.ArrayList;
 import java.util.List;
-import java.util.function.Consumer;
 import java.util.function.Supplier;
 
 /**
@@ -60,8 +59,6 @@
     public static final int FREEFORM_ANIMATION_DURATION = 336;
 
     private final List<IBinder> mPendingTransitionTokens = new ArrayList<>();
-    private Consumer<SurfaceControl.Transaction> mOnAnimationFinishedCallback;
-    private MoveToDesktopAnimator mMoveToDesktopAnimator;
     private DesktopModeWindowDecoration mDesktopModeWindowDecoration;
 
     public EnterDesktopTaskTransitionHandler(
@@ -77,61 +74,6 @@
     }
 
     /**
-     * Starts Transition of a given type
-     * @param type Transition type
-     * @param wct WindowContainerTransaction for transition
-     * @param onAnimationEndCallback to be called after animation
-     */
-    private void startTransition(@WindowManager.TransitionType int type,
-            @NonNull WindowContainerTransaction wct,
-            Consumer<SurfaceControl.Transaction> onAnimationEndCallback) {
-        mOnAnimationFinishedCallback = onAnimationEndCallback;
-        final IBinder token = mTransitions.startTransition(type, wct, this);
-        mPendingTransitionTokens.add(token);
-    }
-
-    /**
-     * Starts Transition of type TRANSIT_START_DRAG_TO_DESKTOP_MODE
-     * @param wct WindowContainerTransaction for transition
-     * @param moveToDesktopAnimator Animator that shrinks and positions task during two part move
-     *                              to desktop animation
-     * @param onAnimationEndCallback to be called after animation
-     */
-    public void startMoveToDesktop(@NonNull WindowContainerTransaction wct,
-            @NonNull MoveToDesktopAnimator moveToDesktopAnimator,
-            Consumer<SurfaceControl.Transaction> onAnimationEndCallback) {
-        mMoveToDesktopAnimator = moveToDesktopAnimator;
-        startTransition(Transitions.TRANSIT_START_DRAG_TO_DESKTOP_MODE, wct,
-                onAnimationEndCallback);
-    }
-
-    /**
-     * Starts Transition of type TRANSIT_FINALIZE_DRAG_TO_DESKTOP_MODE
-     * @param wct WindowContainerTransaction for transition
-     * @param onAnimationEndCallback to be called after animation
-     */
-    public void finalizeMoveToDesktop(@NonNull WindowContainerTransaction wct,
-            Consumer<SurfaceControl.Transaction> onAnimationEndCallback) {
-        startTransition(Transitions.TRANSIT_FINALIZE_DRAG_TO_DESKTOP_MODE, wct,
-                onAnimationEndCallback);
-    }
-
-    /**
-     * Starts Transition of type TRANSIT_CANCEL_ENTERING_DESKTOP_MODE
-     * @param wct WindowContainerTransaction for transition
-     * @param moveToDesktopAnimator Animator that shrinks and positions task during two part move
-     *                              to desktop animation
-     * @param onAnimationEndCallback to be called after animation
-     */
-    public void startCancelMoveToDesktopMode(@NonNull WindowContainerTransaction wct,
-            MoveToDesktopAnimator moveToDesktopAnimator,
-            Consumer<SurfaceControl.Transaction> onAnimationEndCallback) {
-        mMoveToDesktopAnimator = moveToDesktopAnimator;
-        startTransition(Transitions.TRANSIT_CANCEL_DRAG_TO_DESKTOP_MODE, wct,
-                onAnimationEndCallback);
-    }
-
-    /**
      * Starts Transition of type TRANSIT_MOVE_TO_DESKTOP
      * @param wct WindowContainerTransaction for transition
      * @param decor {@link DesktopModeWindowDecoration} of task being animated
@@ -139,8 +81,8 @@
     public void moveToDesktop(@NonNull WindowContainerTransaction wct,
             DesktopModeWindowDecoration decor) {
         mDesktopModeWindowDecoration = decor;
-        startTransition(Transitions.TRANSIT_MOVE_TO_DESKTOP, wct,
-                null /* onAnimationEndCallback */);
+        final IBinder token = mTransitions.startTransition(TRANSIT_MOVE_TO_DESKTOP, wct, this);
+        mPendingTransitionTokens.add(token);
     }
 
     @Override
@@ -182,30 +124,11 @@
         }
 
         final ActivityManager.RunningTaskInfo taskInfo = change.getTaskInfo();
-        if (type == Transitions.TRANSIT_MOVE_TO_DESKTOP
+        if (type == TRANSIT_MOVE_TO_DESKTOP
                 && taskInfo.getWindowingMode() == WINDOWING_MODE_FREEFORM) {
             return animateMoveToDesktop(change, startT, finishCallback);
         }
 
-        if (type == Transitions.TRANSIT_START_DRAG_TO_DESKTOP_MODE
-                && taskInfo.getWindowingMode() == WINDOWING_MODE_FREEFORM) {
-            return animateStartDragToDesktopMode(change, startT, finishT, finishCallback);
-        }
-
-        final Rect endBounds = change.getEndAbsBounds();
-        if (type == Transitions.TRANSIT_FINALIZE_DRAG_TO_DESKTOP_MODE
-                && taskInfo.getWindowingMode() == WINDOWING_MODE_FREEFORM
-                && !endBounds.isEmpty()) {
-            return animateFinalizeDragToDesktopMode(change, startT, finishT, finishCallback,
-                    endBounds);
-        }
-
-        if (type == Transitions.TRANSIT_CANCEL_DRAG_TO_DESKTOP_MODE
-                && taskInfo.getWindowingMode() == WINDOWING_MODE_FREEFORM) {
-            return animateCancelDragToDesktopMode(change, startT, finishT, finishCallback,
-                    endBounds);
-        }
-
         return false;
     }
 
@@ -248,142 +171,6 @@
         return true;
     }
 
-    private boolean animateStartDragToDesktopMode(
-            @NonNull TransitionInfo.Change change,
-            @NonNull SurfaceControl.Transaction startT,
-            @NonNull SurfaceControl.Transaction finishT,
-            @NonNull Transitions.TransitionFinishCallback finishCallback) {
-        // Transitioning to freeform but keeping fullscreen bounds, so the crop is set
-        // to null and we don't require an animation
-        final SurfaceControl sc = change.getLeash();
-        startT.setWindowCrop(sc, null);
-
-        if (mMoveToDesktopAnimator == null
-                || mMoveToDesktopAnimator.getTaskId() != change.getTaskInfo().taskId) {
-            Slog.e(TAG, "No animator available for this transition");
-            return false;
-        }
-
-        // Calculate and set position of the task
-        final PointF position = mMoveToDesktopAnimator.getPosition();
-        startT.setPosition(sc, position.x, position.y);
-        finishT.setPosition(sc, position.x, position.y);
-
-        startT.apply();
-
-        mTransitions.getMainExecutor().execute(() -> finishCallback.onTransitionFinished(null));
-
-        return true;
-    }
-
-    private boolean animateFinalizeDragToDesktopMode(
-            @NonNull TransitionInfo.Change change,
-            @NonNull SurfaceControl.Transaction startT,
-            @NonNull SurfaceControl.Transaction finishT,
-            @NonNull Transitions.TransitionFinishCallback finishCallback,
-            @NonNull Rect endBounds) {
-        // This Transition animates a task to freeform bounds after being dragged into freeform
-        // mode and brings the remaining freeform tasks to front
-        final SurfaceControl sc = change.getLeash();
-        startT.setWindowCrop(sc, endBounds.width(),
-                endBounds.height());
-        startT.apply();
-
-        // End the animation that shrinks the window when task is first dragged from fullscreen
-        if (mMoveToDesktopAnimator != null) {
-            mMoveToDesktopAnimator.endAnimator();
-        }
-
-        // We want to find the scale of the current bounds relative to the end bounds. The
-        // task is currently scaled to DRAG_FREEFORM_SCALE and the final bounds will be
-        // scaled to FINAL_FREEFORM_SCALE. So, it is scaled to
-        // DRAG_FREEFORM_SCALE / FINAL_FREEFORM_SCALE relative to the freeform bounds
-        final ValueAnimator animator =
-                ValueAnimator.ofFloat(
-                        MoveToDesktopAnimator.DRAG_FREEFORM_SCALE / FINAL_FREEFORM_SCALE, 1f);
-        animator.setDuration(FREEFORM_ANIMATION_DURATION);
-        final SurfaceControl.Transaction t = mTransactionSupplier.get();
-        animator.addUpdateListener(animation -> {
-            final float animationValue = (float) animation.getAnimatedValue();
-            t.setScale(sc, animationValue, animationValue);
-
-            final float animationWidth = endBounds.width() * animationValue;
-            final float animationHeight = endBounds.height() * animationValue;
-            final int animationX = endBounds.centerX() - (int) (animationWidth / 2);
-            final int animationY = endBounds.centerY() - (int) (animationHeight / 2);
-
-            t.setPosition(sc, animationX, animationY);
-            t.apply();
-        });
-
-        animator.addListener(new AnimatorListenerAdapter() {
-            @Override
-            public void onAnimationEnd(Animator animation) {
-                if (mOnAnimationFinishedCallback != null) {
-                    mOnAnimationFinishedCallback.accept(finishT);
-                }
-                mTransitions.getMainExecutor().execute(
-                        () -> finishCallback.onTransitionFinished(null));
-            }
-        });
-
-        animator.start();
-        return true;
-    }
-    private boolean animateCancelDragToDesktopMode(
-            @NonNull TransitionInfo.Change change,
-            @NonNull SurfaceControl.Transaction startT,
-            @NonNull SurfaceControl.Transaction finishT,
-            @NonNull Transitions.TransitionFinishCallback finishCallback,
-            @NonNull Rect endBounds) {
-        // This Transition animates a task to fullscreen after being dragged from the status
-        // bar and then released back into the status bar area
-        final SurfaceControl sc = change.getLeash();
-        // Hide the first (fullscreen) frame because the animation will start from the smaller
-        // scale size.
-        startT.hide(sc)
-                .setWindowCrop(sc, endBounds.width(), endBounds.height())
-                .apply();
-
-        if (mMoveToDesktopAnimator == null
-                || mMoveToDesktopAnimator.getTaskId() != change.getTaskInfo().taskId) {
-            Slog.e(TAG, "No animator available for this transition");
-            return false;
-        }
-
-        // End the animation that shrinks the window when task is first dragged from fullscreen
-        mMoveToDesktopAnimator.endAnimator();
-
-        final ValueAnimator animator = new ValueAnimator();
-        animator.setFloatValues(MoveToDesktopAnimator.DRAG_FREEFORM_SCALE, 1f);
-        animator.setDuration(FREEFORM_ANIMATION_DURATION);
-        final SurfaceControl.Transaction t = mTransactionSupplier.get();
-
-        // Get position of the task
-        final float x = mMoveToDesktopAnimator.getPosition().x;
-        final float y = mMoveToDesktopAnimator.getPosition().y;
-
-        animator.addUpdateListener(animation -> {
-            final float scale = (float) animation.getAnimatedValue();
-            t.setPosition(sc, x * (1 - scale), y * (1 - scale))
-                    .setScale(sc, scale, scale)
-                    .show(sc)
-                    .apply();
-        });
-        animator.addListener(new AnimatorListenerAdapter() {
-            @Override
-            public void onAnimationEnd(Animator animation) {
-                if (mOnAnimationFinishedCallback != null) {
-                    mOnAnimationFinishedCallback.accept(finishT);
-                }
-                mTransitions.getMainExecutor().execute(
-                        () -> finishCallback.onTransitionFinished(null));
-            }
-        });
-        animator.start();
-        return true;
-    }
-
     @Nullable
     @Override
     public WindowContainerTransaction handleRequest(@NonNull IBinder transition,
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 ab5c063..41ec33c 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
@@ -150,19 +150,19 @@
     /** Transition type for maximize to freeform transition. */
     public static final int TRANSIT_RESTORE_FROM_MAXIMIZE = WindowManager.TRANSIT_FIRST_CUSTOM + 9;
 
-    /** Transition type for starting the move to desktop mode. */
-    public static final int TRANSIT_START_DRAG_TO_DESKTOP_MODE =
+    /** Transition type for starting the drag to desktop mode. */
+    public static final int TRANSIT_DESKTOP_MODE_START_DRAG_TO_DESKTOP =
             WindowManager.TRANSIT_FIRST_CUSTOM + 10;
 
-    /** Transition type for finalizing the move to desktop mode. */
-    public static final int TRANSIT_FINALIZE_DRAG_TO_DESKTOP_MODE =
+    /** Transition type for finalizing the drag to desktop mode. */
+    public static final int TRANSIT_DESKTOP_MODE_END_DRAG_TO_DESKTOP =
             WindowManager.TRANSIT_FIRST_CUSTOM + 11;
 
     /** Transition type to fullscreen from desktop mode. */
     public static final int TRANSIT_EXIT_DESKTOP_MODE = WindowManager.TRANSIT_FIRST_CUSTOM + 12;
 
-    /** Transition type to animate back to fullscreen when drag to freeform is cancelled. */
-    public static final int TRANSIT_CANCEL_DRAG_TO_DESKTOP_MODE =
+    /** Transition type to cancel the drag to desktop mode. */
+    public static final int TRANSIT_DESKTOP_MODE_CANCEL_DRAG_TO_DESKTOP =
             WindowManager.TRANSIT_FIRST_CUSTOM + 13;
 
     /** Transition type to animate the toggle resize between the max and default desktop sizes. */
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorViewModel.java b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorViewModel.java
index e206039..3add6f4 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorViewModel.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorViewModel.java
@@ -268,13 +268,19 @@
             @NonNull TransitionInfo info,
             @NonNull TransitionInfo.Change change) {
         if (change.getMode() == WindowManager.TRANSIT_CHANGE
-                && (info.getType() == Transitions.TRANSIT_FINALIZE_DRAG_TO_DESKTOP_MODE
-                || info.getType() == Transitions.TRANSIT_CANCEL_DRAG_TO_DESKTOP_MODE
-                || info.getType() == Transitions.TRANSIT_EXIT_DESKTOP_MODE
+                && (info.getType() == Transitions.TRANSIT_EXIT_DESKTOP_MODE
                 || info.getType() == Transitions.TRANSIT_DESKTOP_MODE_TOGGLE_RESIZE
                 || info.getType() == Transitions.TRANSIT_MOVE_TO_DESKTOP)) {
             mWindowDecorByTaskId.get(change.getTaskInfo().taskId)
                     .addTransitionPausingRelayout(transition);
+        } else if (change.getMode() == WindowManager.TRANSIT_TO_BACK
+                && info.getType() == Transitions.TRANSIT_DESKTOP_MODE_START_DRAG_TO_DESKTOP
+                && change.getTaskInfo() != null) {
+            final DesktopModeWindowDecoration decor =
+                    mWindowDecorByTaskId.get(change.getTaskInfo().taskId);
+            if (decor != null) {
+                decor.addTransitionPausingRelayout(transition);
+            }
         }
     }
 
@@ -765,10 +771,8 @@
                         mMoveToDesktopAnimator = null;
                         return;
                     } else if (mMoveToDesktopAnimator != null) {
-                        relevantDecor.incrementRelayoutBlock();
                         mDesktopTasksController.ifPresent(
-                                c -> c.cancelMoveToDesktop(relevantDecor.mTaskInfo,
-                                        mMoveToDesktopAnimator));
+                                c -> c.cancelDragToDesktop(relevantDecor.mTaskInfo));
                         mMoveToDesktopAnimator = null;
                         return;
                     }
@@ -790,15 +794,24 @@
                             relevantDecor.mTaskInfo.displayId);
                     if (ev.getY() > statusBarHeight) {
                         if (mMoveToDesktopAnimator == null) {
-                            closeOtherSplitTask(relevantDecor.mTaskInfo.taskId);
                             mMoveToDesktopAnimator = new MoveToDesktopAnimator(
-                                    mDragToDesktopAnimationStartBounds, relevantDecor.mTaskInfo,
-                                    relevantDecor.mTaskSurface);
+                                    mContext, mDragToDesktopAnimationStartBounds,
+                                    relevantDecor.mTaskInfo, relevantDecor.mTaskSurface);
                             mDesktopTasksController.ifPresent(
-                                    c -> c.startMoveToDesktop(relevantDecor.mTaskInfo,
-                                            mDragToDesktopAnimationStartBounds,
-                                            mMoveToDesktopAnimator));
-                            mMoveToDesktopAnimator.startAnimation();
+                                    c -> {
+                                        final int taskId = relevantDecor.mTaskInfo.taskId;
+                                        relevantDecor.incrementRelayoutBlock();
+                                        if (isTaskInSplitScreen(taskId)) {
+                                            final DesktopModeWindowDecoration otherDecor =
+                                                    mWindowDecorByTaskId.get(
+                                                            getOtherSplitTask(taskId).taskId);
+                                            if (otherDecor != null) {
+                                                otherDecor.incrementRelayoutBlock();
+                                            }
+                                        }
+                                        c.startDragToDesktop(relevantDecor.mTaskInfo,
+                                                mMoveToDesktopAnimator, relevantDecor);
+                                    });
                         }
                     }
                     if (mMoveToDesktopAnimator != null) {
@@ -837,7 +850,6 @@
      */
     private void animateToDesktop(DesktopModeWindowDecoration relevantDecor,
             MotionEvent ev) {
-        relevantDecor.incrementRelayoutBlock();
         centerAndMoveToDesktopWithAnimation(relevantDecor, ev);
     }
 
@@ -853,15 +865,15 @@
         final SurfaceControl sc = relevantDecor.mTaskSurface;
         final Rect endBounds = calculateFreeformBounds(ev.getDisplayId(), DRAG_FREEFORM_SCALE);
         final Transaction t = mTransactionFactory.get();
-        final float diffX = endBounds.centerX() - ev.getX();
-        final float diffY = endBounds.top - ev.getY();
-        final float startingX = ev.getX() - DRAG_FREEFORM_SCALE
+        final float diffX = endBounds.centerX() - ev.getRawX();
+        final float diffY = endBounds.top - ev.getRawY();
+        final float startingX = ev.getRawX() - DRAG_FREEFORM_SCALE
                 * mDragToDesktopAnimationStartBounds.width() / 2;
 
         animator.addUpdateListener(animation -> {
             final float animatorValue = (float) animation.getAnimatedValue();
             final float x = startingX + diffX * animatorValue;
-            final float y = ev.getY() + diffY * animatorValue;
+            final float y = ev.getRawY() + diffY * animatorValue;
             t.setPosition(sc, x, y);
             t.apply();
         });
@@ -869,9 +881,11 @@
             @Override
             public void onAnimationEnd(Animator animation) {
                 mDesktopTasksController.ifPresent(
-                        c -> c.onDragPositioningEndThroughStatusBar(
-                                relevantDecor.mTaskInfo,
-                                calculateFreeformBounds(ev.getDisplayId(), FINAL_FREEFORM_SCALE)));
+                        c -> {
+                            c.onDragPositioningEndThroughStatusBar(relevantDecor.mTaskInfo,
+                                    calculateFreeformBounds(ev.getDisplayId(),
+                                            FINAL_FREEFORM_SCALE));
+                        });
             }
         });
         animator.start();
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/MoveToDesktopAnimator.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/MoveToDesktopAnimator.kt
index b2267dd..af05523 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/MoveToDesktopAnimator.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/MoveToDesktopAnimator.kt
@@ -2,10 +2,12 @@
 
 import android.animation.ValueAnimator
 import android.app.ActivityManager.RunningTaskInfo
+import android.content.Context
 import android.graphics.PointF
 import android.graphics.Rect
 import android.view.MotionEvent
 import android.view.SurfaceControl
+import com.android.internal.policy.ScreenDecorationsUtils
 
 /**
  * Creates an animator to shrink and position task after a user drags a fullscreen task from
@@ -14,6 +16,7 @@
  * accessed by the EnterDesktopTaskTransitionHandler.
  */
 class MoveToDesktopAnimator @JvmOverloads constructor(
+        private val context: Context,
         private val startBounds: Rect,
         private val taskInfo: RunningTaskInfo,
         private val taskSurface: SurfaceControl,
@@ -33,9 +36,11 @@
             .setDuration(ANIMATION_DURATION.toLong())
             .apply {
                 val t = SurfaceControl.Transaction()
+                val cornerRadius = ScreenDecorationsUtils.getWindowCornerRadius(context)
                 addUpdateListener { animation ->
                     val animatorValue = animation.animatedValue as Float
                     t.setScale(taskSurface, animatorValue, animatorValue)
+                            .setCornerRadius(taskSurface, cornerRadius)
                             .apply()
                 }
             }
@@ -44,19 +49,40 @@
     val position: PointF = PointF(0.0f, 0.0f)
 
     /**
+     * Whether motion events from the drag gesture should affect the dragged surface or not. Used
+     * to disallow moving the surface's position prematurely since it should not start moving at
+     * all until the drag-to-desktop transition is ready to animate and the wallpaper/home are
+     * ready to be revealed behind the dragged/scaled task.
+     */
+    private var allowSurfaceChangesOnMove = false
+
+    /**
      * Starts the animation that scales the task down.
      */
     fun startAnimation() {
+        allowSurfaceChangesOnMove = true
         dragToDesktopAnimator.start()
     }
 
     /**
-     * Uses the position of the motion event and the current scale of the task as defined by the
-     * ValueAnimator to update the local position variable and set the task surface's position
+     * Uses the position of the motion event of the drag-to-desktop gesture to update the dragged
+     * task's position on screen to follow the touch point. Note that the position change won't
+     * be applied immediately always, such as near the beginning where it waits until the wallpaper
+     * or home are visible behind it. Once they're visible the surface will catch-up to the most
+     * recent touch position.
      */
     fun updatePosition(ev: MotionEvent) {
-        position.x = ev.x - animatedTaskWidth / 2
-        position.y = ev.y
+        // Using rawX/Y because when dragging a task in split, the local X/Y is relative to the
+        // split stages, but the split task surface is re-parented to the task display area to
+        // allow dragging beyond its stage across any region of the display. Because of that, the
+        // rawX/Y are more true to where the gesture is on screen and where the surface should be
+        // positioned.
+        position.x = ev.rawX - animatedTaskWidth / 2
+        position.y = ev.rawY
+
+        if (!allowSurfaceChangesOnMove) {
+            return
+        }
 
         val t = transactionFactory()
         t.setPosition(taskSurface, position.x, position.y)
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksControllerTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksControllerTest.kt
index ebcb640..fde6acb 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksControllerTest.kt
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksControllerTest.kt
@@ -100,6 +100,7 @@
     @Mock lateinit var enterDesktopTransitionHandler: EnterDesktopTaskTransitionHandler
     @Mock lateinit var mToggleResizeDesktopTaskTransitionHandler:
             ToggleResizeDesktopTaskTransitionHandler
+    @Mock lateinit var dragToDesktopTransitionHandler: DragToDesktopTransitionHandler
     @Mock lateinit var launchAdjacentController: LaunchAdjacentController
     @Mock lateinit var desktopModeWindowDecoration: DesktopModeWindowDecoration
     @Mock lateinit var splitScreenController: SplitScreenController
@@ -127,7 +128,7 @@
         whenever(transitions.startTransition(anyInt(), any(), isNull())).thenAnswer { Binder() }
 
         controller = createController()
-        controller.splitScreenController = splitScreenController
+        controller.setSplitScreenController(splitScreenController)
 
         shellInit.init()
 
@@ -150,6 +151,7 @@
             enterDesktopTransitionHandler,
             exitDesktopTransitionHandler,
             mToggleResizeDesktopTaskTransitionHandler,
+            dragToDesktopTransitionHandler,
             desktopModeTaskRepository,
             launchAdjacentController,
             recentsTransitionHandler,
@@ -757,6 +759,7 @@
 
     private fun setUpSplitScreenTask(displayId: Int = DEFAULT_DISPLAY): RunningTaskInfo {
         val task = createSplitScreenTask(displayId)
+        whenever(splitScreenController.isTaskInSplitScreen(task.taskId)).thenReturn(true)
         whenever(shellTaskOrganizer.getRunningTaskInfo(task.taskId)).thenReturn(task)
         runningTasks.add(task)
         return task
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DragToDesktopTransitionHandlerTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DragToDesktopTransitionHandlerTest.kt
new file mode 100644
index 0000000..a5629c8
--- /dev/null
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DragToDesktopTransitionHandlerTest.kt
@@ -0,0 +1,211 @@
+package com.android.wm.shell.desktopmode
+
+import android.app.ActivityManager.RunningTaskInfo
+import android.app.WindowConfiguration.ACTIVITY_TYPE_HOME
+import android.app.WindowConfiguration.ACTIVITY_TYPE_STANDARD
+import android.app.WindowConfiguration.WINDOWING_MODE_FULLSCREEN
+import android.app.WindowConfiguration.WINDOWING_MODE_MULTI_WINDOW
+import android.app.WindowConfiguration.WindowingMode
+import android.graphics.PointF
+import android.os.IBinder
+import android.testing.AndroidTestingRunner
+import android.testing.TestableLooper.RunWithLooper
+import android.view.SurfaceControl
+import android.window.TransitionInfo
+import android.window.TransitionInfo.FLAG_IS_WALLPAPER
+import androidx.test.filters.SmallTest
+import com.android.server.testutils.any
+import com.android.server.testutils.mock
+import com.android.wm.shell.RootTaskDisplayAreaOrganizer
+import com.android.wm.shell.ShellTestCase
+import com.android.wm.shell.TestRunningTaskInfoBuilder
+import com.android.wm.shell.splitscreen.SplitScreenController
+import com.android.wm.shell.transition.Transitions
+import com.android.wm.shell.transition.Transitions.TRANSIT_DESKTOP_MODE_CANCEL_DRAG_TO_DESKTOP
+import com.android.wm.shell.transition.Transitions.TRANSIT_DESKTOP_MODE_START_DRAG_TO_DESKTOP
+import com.android.wm.shell.windowdecor.MoveToDesktopAnimator
+import java.util.function.Supplier
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.ArgumentMatchers.eq
+import org.mockito.Mock
+import org.mockito.kotlin.never
+import org.mockito.kotlin.verify
+import org.mockito.kotlin.verifyZeroInteractions
+import org.mockito.kotlin.whenever
+
+/** Tests of [DragToDesktopTransitionHandler]. */
+@SmallTest
+@RunWithLooper
+@RunWith(AndroidTestingRunner::class)
+class DragToDesktopTransitionHandlerTest : ShellTestCase() {
+
+    @Mock private lateinit var transitions: Transitions
+    @Mock private lateinit var taskDisplayAreaOrganizer: RootTaskDisplayAreaOrganizer
+    @Mock private lateinit var splitScreenController: SplitScreenController
+
+    private val transactionSupplier = Supplier { mock<SurfaceControl.Transaction>() }
+
+    private lateinit var handler: DragToDesktopTransitionHandler
+
+    @Before
+    fun setUp() {
+        handler =
+            DragToDesktopTransitionHandler(
+                    context,
+                    transitions,
+                    taskDisplayAreaOrganizer,
+                    transactionSupplier
+                )
+                .apply { setSplitScreenController(splitScreenController) }
+    }
+
+    @Test
+    fun startDragToDesktop_animateDragWhenReady() {
+        val task = createTask()
+        val dragAnimator = mock<MoveToDesktopAnimator>()
+        // Simulate transition is started.
+        val transition = startDragToDesktopTransition(task, dragAnimator)
+
+        // Now it's ready to animate.
+        handler.startAnimation(
+            transition = transition,
+            info =
+                createTransitionInfo(
+                    type = TRANSIT_DESKTOP_MODE_START_DRAG_TO_DESKTOP,
+                    draggedTask = task
+                ),
+            startTransaction = mock(),
+            finishTransaction = mock(),
+            finishCallback = {}
+        )
+
+        verify(dragAnimator).startAnimation()
+    }
+
+    @Test
+    fun startDragToDesktop_cancelledBeforeReady_startCancelTransition() {
+        val task = createTask()
+        val dragAnimator = mock<MoveToDesktopAnimator>()
+        // Simulate transition is started and is ready to animate.
+        val transition = startDragToDesktopTransition(task, dragAnimator)
+
+        handler.cancelDragToDesktopTransition()
+
+        handler.startAnimation(
+            transition = transition,
+            info =
+                createTransitionInfo(
+                    type = TRANSIT_DESKTOP_MODE_START_DRAG_TO_DESKTOP,
+                    draggedTask = task
+                ),
+            startTransaction = mock(),
+            finishTransaction = mock(),
+            finishCallback = {}
+        )
+
+        // Don't even animate the "drag" since it was already cancelled.
+        verify(dragAnimator, never()).startAnimation()
+        // Instead, start the cancel transition.
+        verify(transitions)
+            .startTransition(eq(TRANSIT_DESKTOP_MODE_CANCEL_DRAG_TO_DESKTOP), any(), eq(handler))
+    }
+
+    @Test
+    fun cancelDragToDesktop_startWasReady_cancel() {
+        val task = createTask()
+        val dragAnimator = mock<MoveToDesktopAnimator>()
+        whenever(dragAnimator.position).thenReturn(PointF())
+        // Simulate transition is started and is ready to animate.
+        val transition = startDragToDesktopTransition(task, dragAnimator)
+        handler.startAnimation(
+            transition = transition,
+            info =
+                createTransitionInfo(
+                    type = TRANSIT_DESKTOP_MODE_START_DRAG_TO_DESKTOP,
+                    draggedTask = task
+                ),
+            startTransaction = mock(),
+            finishTransaction = mock(),
+            finishCallback = {}
+        )
+
+        // Then user cancelled after it had already started.
+        handler.cancelDragToDesktopTransition()
+
+        // Cancel animation should run since it had already started.
+        verify(dragAnimator).endAnimator()
+    }
+
+    @Test
+    fun cancelDragToDesktop_startWasNotReady_animateCancel() {
+        val task = createTask()
+        val dragAnimator = mock<MoveToDesktopAnimator>()
+        // Simulate transition is started and is ready to animate.
+        startDragToDesktopTransition(task, dragAnimator)
+
+        // Then user cancelled before the transition was ready and animated.
+        handler.cancelDragToDesktopTransition()
+
+        // No need to animate the cancel since the start animation couldn't even start.
+        verifyZeroInteractions(dragAnimator)
+    }
+
+    private fun startDragToDesktopTransition(
+        task: RunningTaskInfo,
+        dragAnimator: MoveToDesktopAnimator
+    ): IBinder {
+        val token = mock<IBinder>()
+        whenever(
+                transitions.startTransition(
+                    eq(TRANSIT_DESKTOP_MODE_START_DRAG_TO_DESKTOP),
+                    any(),
+                    eq(handler)
+                )
+            )
+            .thenReturn(token)
+        handler.startDragToDesktopTransition(task.taskId, dragAnimator, mock())
+        return token
+    }
+
+    private fun createTask(
+        @WindowingMode windowingMode: Int = WINDOWING_MODE_FULLSCREEN,
+        isHome: Boolean = false,
+    ): RunningTaskInfo {
+        return TestRunningTaskInfoBuilder()
+            .setActivityType(if (isHome) ACTIVITY_TYPE_HOME else ACTIVITY_TYPE_STANDARD)
+            .setWindowingMode(windowingMode)
+            .build()
+            .also {
+                whenever(splitScreenController.isTaskInSplitScreen(it.taskId))
+                    .thenReturn(windowingMode == WINDOWING_MODE_MULTI_WINDOW)
+            }
+    }
+
+    private fun createTransitionInfo(type: Int, draggedTask: RunningTaskInfo): TransitionInfo {
+        return TransitionInfo(type, 0 /* flags */).apply {
+            addChange( // Home.
+                TransitionInfo.Change(mock(), mock()).apply {
+                    parent = null
+                    taskInfo =
+                        TestRunningTaskInfoBuilder().setActivityType(ACTIVITY_TYPE_HOME).build()
+                    flags = flags or FLAG_IS_WALLPAPER
+                }
+            )
+            addChange( // Dragged Task.
+                TransitionInfo.Change(mock(), mock()).apply {
+                    parent = null
+                    taskInfo = draggedTask
+                }
+            )
+            addChange( // Wallpaper.
+                TransitionInfo.Change(mock(), mock()).apply {
+                    parent = null
+                    taskInfo = null
+                    flags = flags or FLAG_IS_WALLPAPER
+                }
+            )
+        }
+    }
+}
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/EnterDesktopTaskTransitionHandlerTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/EnterDesktopTaskTransitionHandlerTest.java
deleted file mode 100644
index 772d97d..0000000
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/EnterDesktopTaskTransitionHandlerTest.java
+++ /dev/null
@@ -1,171 +0,0 @@
-/*
- * Copyright (C) 2023 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package com.android.wm.shell.desktopmode;
-
-import static android.app.WindowConfiguration.WINDOWING_MODE_FREEFORM;
-
-import static androidx.test.internal.runner.junit4.statement.UiThreadStatement.runOnUiThread;
-
-import static org.junit.Assert.assertTrue;
-import static org.mockito.Mockito.doReturn;
-import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.verify;
-
-import android.annotation.NonNull;
-import android.app.ActivityManager;
-import android.app.WindowConfiguration;
-import android.graphics.PointF;
-import android.graphics.Rect;
-import android.os.IBinder;
-import android.view.SurfaceControl;
-import android.view.WindowManager;
-import android.window.IWindowContainerToken;
-import android.window.TransitionInfo;
-import android.window.WindowContainerToken;
-import android.window.WindowContainerTransaction;
-
-import androidx.test.filters.SmallTest;
-
-import com.android.wm.shell.common.ShellExecutor;
-import com.android.wm.shell.transition.Transitions;
-import com.android.wm.shell.windowdecor.MoveToDesktopAnimator;
-
-import junit.framework.AssertionFailedError;
-
-import org.junit.Before;
-import org.junit.Test;
-import org.mockito.Mock;
-import org.mockito.MockitoAnnotations;
-
-import java.util.function.Supplier;
-
-/** Tests of {@link com.android.wm.shell.desktopmode.EnterDesktopTaskTransitionHandler} */
-@SmallTest
-public class EnterDesktopTaskTransitionHandlerTest {
-
-    @Mock
-    private Transitions mTransitions;
-    @Mock
-    IBinder mToken;
-    @Mock
-    Supplier<SurfaceControl.Transaction> mTransactionFactory;
-    @Mock
-    SurfaceControl.Transaction mStartT;
-    @Mock
-    SurfaceControl.Transaction mFinishT;
-    @Mock
-    SurfaceControl.Transaction mAnimationT;
-    @Mock
-    Transitions.TransitionFinishCallback mTransitionFinishCallback;
-    @Mock
-    ShellExecutor mExecutor;
-    @Mock
-    SurfaceControl mSurfaceControl;
-    @Mock
-    MoveToDesktopAnimator mMoveToDesktopAnimator;
-    @Mock
-    PointF mPosition;
-
-    private EnterDesktopTaskTransitionHandler mEnterDesktopTaskTransitionHandler;
-
-    @Before
-    public void setUp() {
-        MockitoAnnotations.initMocks(this);
-
-        doReturn(mExecutor).when(mTransitions).getMainExecutor();
-        doReturn(mAnimationT).when(mTransactionFactory).get();
-        doReturn(mPosition).when(mMoveToDesktopAnimator).getPosition();
-
-        mEnterDesktopTaskTransitionHandler = new EnterDesktopTaskTransitionHandler(mTransitions,
-                mTransactionFactory);
-    }
-
-    @Test
-    public void testEnterFreeformAnimation() {
-        final int taskId = 1;
-        WindowContainerTransaction wct = new WindowContainerTransaction();
-        doReturn(mToken).when(mTransitions)
-                .startTransition(Transitions.TRANSIT_START_DRAG_TO_DESKTOP_MODE, wct,
-                        mEnterDesktopTaskTransitionHandler);
-        doReturn(taskId).when(mMoveToDesktopAnimator).getTaskId();
-
-        mEnterDesktopTaskTransitionHandler.startMoveToDesktop(wct,
-                mMoveToDesktopAnimator, null);
-
-        TransitionInfo.Change change =
-                createChange(WindowManager.TRANSIT_CHANGE, taskId, WINDOWING_MODE_FREEFORM);
-        TransitionInfo info = createTransitionInfo(Transitions.TRANSIT_START_DRAG_TO_DESKTOP_MODE,
-                change);
-
-
-        assertTrue(mEnterDesktopTaskTransitionHandler
-                .startAnimation(mToken, info, mStartT, mFinishT, mTransitionFinishCallback));
-
-        verify(mStartT).setWindowCrop(mSurfaceControl, null);
-        verify(mStartT).apply();
-    }
-
-    @Test
-    public void testTransitEnterDesktopModeAnimation() throws Throwable {
-        final int transitionType = Transitions.TRANSIT_FINALIZE_DRAG_TO_DESKTOP_MODE;
-        final int taskId = 1;
-        WindowContainerTransaction wct = new WindowContainerTransaction();
-        doReturn(mToken).when(mTransitions)
-                .startTransition(transitionType, wct, mEnterDesktopTaskTransitionHandler);
-        mEnterDesktopTaskTransitionHandler.finalizeMoveToDesktop(wct, null);
-
-        TransitionInfo.Change change =
-                createChange(WindowManager.TRANSIT_CHANGE, taskId, WINDOWING_MODE_FREEFORM);
-        change.setEndAbsBounds(new Rect(0, 0, 1, 1));
-        TransitionInfo info = createTransitionInfo(
-                Transitions.TRANSIT_FINALIZE_DRAG_TO_DESKTOP_MODE, change);
-
-        runOnUiThread(() -> {
-            try {
-                assertTrue(mEnterDesktopTaskTransitionHandler
-                                .startAnimation(mToken, info, mStartT, mFinishT,
-                                        mTransitionFinishCallback));
-            } catch (Exception e) {
-                throw new AssertionFailedError(e.getMessage());
-            }
-        });
-
-        verify(mStartT).setWindowCrop(mSurfaceControl, change.getEndAbsBounds().width(),
-                change.getEndAbsBounds().height());
-        verify(mStartT).apply();
-    }
-
-    private TransitionInfo.Change createChange(@WindowManager.TransitionType int type, int taskId,
-            @WindowConfiguration.WindowingMode int windowingMode) {
-        final ActivityManager.RunningTaskInfo taskInfo = new ActivityManager.RunningTaskInfo();
-        taskInfo.taskId = taskId;
-        taskInfo.configuration.windowConfiguration.setWindowingMode(windowingMode);
-        final TransitionInfo.Change change = new TransitionInfo.Change(
-                new WindowContainerToken(mock(IWindowContainerToken.class)), mSurfaceControl);
-        change.setMode(type);
-        change.setTaskInfo(taskInfo);
-        return change;
-    }
-
-    private static TransitionInfo createTransitionInfo(
-            @WindowManager.TransitionType int type, @NonNull TransitionInfo.Change change) {
-        TransitionInfo info = new TransitionInfo(type, 0);
-        info.addChange(change);
-        return info;
-    }
-
-}
diff --git a/services/core/java/com/android/server/wm/InputMonitor.java b/services/core/java/com/android/server/wm/InputMonitor.java
index c8fd16d..997b608 100644
--- a/services/core/java/com/android/server/wm/InputMonitor.java
+++ b/services/core/java/com/android/server/wm/InputMonitor.java
@@ -393,8 +393,12 @@
      */
     void setActiveRecents(@Nullable ActivityRecord activity, @Nullable ActivityRecord layer) {
         final boolean clear = activity == null;
+        final boolean wasActive = mActiveRecentsActivity != null && mActiveRecentsLayerRef != null;
         mActiveRecentsActivity = clear ? null : new WeakReference<>(activity);
         mActiveRecentsLayerRef = clear ? null : new WeakReference<>(layer);
+        if (clear && wasActive) {
+            setUpdateInputWindowsNeededLw();
+        }
     }
 
     private static <T> T getWeak(WeakReference<T> ref) {
diff --git a/services/core/java/com/android/server/wm/Transition.java b/services/core/java/com/android/server/wm/Transition.java
index f700944..caa57bb 100644
--- a/services/core/java/com/android/server/wm/Transition.java
+++ b/services/core/java/com/android/server/wm/Transition.java
@@ -1344,6 +1344,7 @@
             final DisplayContent dc =
                     mController.mAtm.mRootWindowContainer.getDisplayContent(mRecentsDisplayId);
             dc.getInputMonitor().setActiveRecents(null /* activity */, null /* layer */);
+            dc.getInputMonitor().updateInputWindowsLw(false /* force */);
         }
         if (mTransientLaunches != null) {
             for (int i = mTransientLaunches.size() - 1; i >= 0; --i) {