Exit desktop windowing transition

- Added `FreeformTaskTransitionStarterInitializer` to extract transition
  starter setup logic from `FreeformTaskTransitionHandler`.
- Added `DesktopMixedTransitionHandler` to coordinate desktop task close
  transition animation between Launcher (via `dispatchTransition`),
  `CloseDesktopTaskTransitionHandler` and fallback `FreeformTaskTransitionHandler`.
- Added `CloseDesktopTaskTransitionHandler` to animate close desktop
  window transition, when not leaving desktop mode.

Bug: 331165070
Test: atest WMShellUnitTests:DesktopMixedTransitionHandlerTest WMShellUnitTests:CloseDesktopTaskTransitionHandlerTest
Flag: com.android.window.flags.enable_desktop_windowing_exit_transitions
Change-Id: Ic5bf3673a5f5e8ae2fa22a9d3dbe77cefd244cb9
diff --git a/libs/WindowManager/Shell/Android.bp b/libs/WindowManager/Shell/Android.bp
index f857429..a796ecc 100644
--- a/libs/WindowManager/Shell/Android.bp
+++ b/libs/WindowManager/Shell/Android.bp
@@ -220,6 +220,7 @@
         "//frameworks/libs/systemui:com_android_systemui_shared_flags_lib",
         "//frameworks/libs/systemui:iconloader_base",
         "com_android_wm_shell_flags_lib",
+        "PlatformAnimationLib",
         "WindowManager-Shell-proto",
         "WindowManager-Shell-lite-proto",
         "WindowManager-Shell-shared",
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 96a0775..888fc62e 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
@@ -16,6 +16,7 @@
 
 package com.android.wm.shell.dagger;
 
+import static android.window.flags.DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_EXIT_TRANSITIONS;
 import static android.window.flags.DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_TASK_LIMIT;
 
 import android.annotation.Nullable;
@@ -59,8 +60,10 @@
 import com.android.wm.shell.common.TaskStackListenerImpl;
 import com.android.wm.shell.dagger.back.ShellBackAnimationModule;
 import com.android.wm.shell.dagger.pip.PipModule;
+import com.android.wm.shell.desktopmode.CloseDesktopTaskTransitionHandler;
 import com.android.wm.shell.desktopmode.DefaultDragToDesktopTransitionHandler;
 import com.android.wm.shell.desktopmode.DesktopActivityOrientationChangeHandler;
+import com.android.wm.shell.desktopmode.DesktopMixedTransitionHandler;
 import com.android.wm.shell.desktopmode.DesktopModeDragAndDropTransitionHandler;
 import com.android.wm.shell.desktopmode.DesktopModeEventLogger;
 import com.android.wm.shell.desktopmode.DesktopModeLoggerTransitionObserver;
@@ -85,6 +88,8 @@
 import com.android.wm.shell.freeform.FreeformTaskListener;
 import com.android.wm.shell.freeform.FreeformTaskTransitionHandler;
 import com.android.wm.shell.freeform.FreeformTaskTransitionObserver;
+import com.android.wm.shell.freeform.FreeformTaskTransitionStarter;
+import com.android.wm.shell.freeform.FreeformTaskTransitionStarterInitializer;
 import com.android.wm.shell.keyguard.KeyguardTransitionHandler;
 import com.android.wm.shell.onehanded.OneHandedController;
 import com.android.wm.shell.pip.PipTransitionController;
@@ -317,9 +322,13 @@
     static FreeformComponents provideFreeformComponents(
             FreeformTaskListener taskListener,
             FreeformTaskTransitionHandler transitionHandler,
-            FreeformTaskTransitionObserver transitionObserver) {
+            FreeformTaskTransitionObserver transitionObserver,
+            FreeformTaskTransitionStarterInitializer transitionStarterInitializer) {
         return new FreeformComponents(
-                taskListener, Optional.of(transitionHandler), Optional.of(transitionObserver));
+                taskListener,
+                Optional.of(transitionHandler),
+                Optional.of(transitionObserver),
+                Optional.of(transitionStarterInitializer));
     }
 
     @WMSingleton
@@ -343,27 +352,15 @@
     @WMSingleton
     @Provides
     static FreeformTaskTransitionHandler provideFreeformTaskTransitionHandler(
-            ShellInit shellInit,
             Transitions transitions,
-            Context context,
-            WindowDecorViewModel windowDecorViewModel,
             DisplayController displayController,
             @ShellMainThread ShellExecutor mainExecutor,
-            @ShellAnimationThread ShellExecutor animExecutor,
-            @DynamicOverride DesktopModeTaskRepository desktopModeTaskRepository,
-            InteractionJankMonitor interactionJankMonitor,
-            @ShellMainThread Handler handler) {
+            @ShellAnimationThread ShellExecutor animExecutor) {
         return new FreeformTaskTransitionHandler(
-                shellInit,
                 transitions,
-                context,
-                windowDecorViewModel,
                 displayController,
                 mainExecutor,
-                animExecutor,
-                desktopModeTaskRepository,
-                interactionJankMonitor,
-                handler);
+                animExecutor);
     }
 
     @WMSingleton
@@ -377,6 +374,23 @@
                 context, shellInit, transitions, windowDecorViewModel);
     }
 
+    @WMSingleton
+    @Provides
+    static FreeformTaskTransitionStarterInitializer provideFreeformTaskTransitionStarterInitializer(
+            ShellInit shellInit,
+            WindowDecorViewModel windowDecorViewModel,
+            FreeformTaskTransitionHandler freeformTaskTransitionHandler,
+            Optional<DesktopMixedTransitionHandler> desktopMixedTransitionHandler) {
+        FreeformTaskTransitionStarter transitionStarter;
+        if (desktopMixedTransitionHandler.isPresent()) {
+            transitionStarter = desktopMixedTransitionHandler.get();
+        } else {
+            transitionStarter = freeformTaskTransitionHandler;
+        }
+        return new FreeformTaskTransitionStarterInitializer(shellInit, windowDecorViewModel,
+                transitionStarter);
+    }
+
     //
     // One handed mode
     //
@@ -686,7 +700,17 @@
             InteractionJankMonitor interactionJankMonitor,
             @ShellMainThread Handler handler) {
         return new ExitDesktopTaskTransitionHandler(
-            transitions, context, interactionJankMonitor, handler);
+                transitions, context, interactionJankMonitor, handler);
+    }
+
+    @WMSingleton
+    @Provides
+    static CloseDesktopTaskTransitionHandler provideCloseDesktopTaskTransitionHandler(
+            Context context,
+            @ShellMainThread ShellExecutor mainExecutor,
+            @ShellAnimationThread ShellExecutor animExecutor
+    ) {
+        return new CloseDesktopTaskTransitionHandler(context, mainExecutor, animExecutor);
     }
 
     @WMSingleton
@@ -745,6 +769,32 @@
 
     @WMSingleton
     @Provides
+    static Optional<DesktopMixedTransitionHandler> provideDesktopMixedTransitionHandler(
+            Context context,
+            Transitions transitions,
+            @DynamicOverride DesktopModeTaskRepository desktopModeTaskRepository,
+            FreeformTaskTransitionHandler freeformTaskTransitionHandler,
+            CloseDesktopTaskTransitionHandler closeDesktopTaskTransitionHandler,
+            InteractionJankMonitor interactionJankMonitor,
+            @ShellMainThread Handler handler
+    ) {
+        if (!DesktopModeStatus.canEnterDesktopMode(context)
+                || !ENABLE_DESKTOP_WINDOWING_EXIT_TRANSITIONS.isTrue()) {
+            return Optional.empty();
+        }
+        return Optional.of(
+                new DesktopMixedTransitionHandler(
+                        context,
+                        transitions,
+                        desktopModeTaskRepository,
+                        freeformTaskTransitionHandler,
+                        closeDesktopTaskTransitionHandler,
+                        interactionJankMonitor,
+                        handler));
+    }
+
+    @WMSingleton
+    @Provides
     static DesktopModeLoggerTransitionObserver provideDesktopModeLoggerTransitionObserver(
             Context context,
             ShellInit shellInit,
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/CloseDesktopTaskTransitionHandler.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/CloseDesktopTaskTransitionHandler.kt
new file mode 100644
index 0000000..a16c15df
--- /dev/null
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/CloseDesktopTaskTransitionHandler.kt
@@ -0,0 +1,153 @@
+/*
+ * Copyright (C) 2024 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 android.animation.Animator
+import android.animation.AnimatorSet
+import android.animation.RectEvaluator
+import android.animation.ValueAnimator
+import android.app.WindowConfiguration.WINDOWING_MODE_FREEFORM
+import android.content.Context
+import android.graphics.Rect
+import android.os.IBinder
+import android.util.TypedValue
+import android.view.SurfaceControl.Transaction
+import android.view.WindowManager
+import android.window.TransitionInfo
+import android.window.TransitionRequestInfo
+import android.window.WindowContainerTransaction
+import androidx.core.animation.addListener
+import com.android.app.animation.Interpolators
+import com.android.wm.shell.common.ShellExecutor
+import com.android.wm.shell.transition.Transitions
+import java.util.function.Supplier
+
+/** The [Transitions.TransitionHandler] that handles transitions for closing desktop mode tasks. */
+class CloseDesktopTaskTransitionHandler
+@JvmOverloads
+constructor(
+    private val context: Context,
+    private val mainExecutor: ShellExecutor,
+    private val animExecutor: ShellExecutor,
+    private val transactionSupplier: Supplier<Transaction> = Supplier { Transaction() },
+) : Transitions.TransitionHandler {
+
+    private val runningAnimations = mutableMapOf<IBinder, List<Animator>>()
+
+    /** Returns null, as it only handles transitions started from Shell. */
+    override fun handleRequest(
+        transition: IBinder,
+        request: TransitionRequestInfo,
+    ): WindowContainerTransaction? = null
+
+    override fun startAnimation(
+        transition: IBinder,
+        info: TransitionInfo,
+        startTransaction: Transaction,
+        finishTransaction: Transaction,
+        finishCallback: Transitions.TransitionFinishCallback,
+    ): Boolean {
+        if (info.type != WindowManager.TRANSIT_CLOSE) return false
+        val animations = mutableListOf<Animator>()
+        val onAnimFinish: (Animator) -> Unit = { animator ->
+            mainExecutor.execute {
+                // Animation completed
+                animations.remove(animator)
+                if (animations.isEmpty()) {
+                    // All animations completed, finish the transition
+                    runningAnimations.remove(transition)
+                    finishCallback.onTransitionFinished(/* wct= */ null)
+                }
+            }
+        }
+        animations +=
+            info.changes
+                .filter {
+                    it.mode == WindowManager.TRANSIT_CLOSE &&
+                        it.taskInfo?.windowingMode == WINDOWING_MODE_FREEFORM
+                }
+                .map { createCloseAnimation(it, finishTransaction, onAnimFinish) }
+        if (animations.isEmpty()) return false
+        runningAnimations[transition] = animations
+        animExecutor.execute { animations.forEach(Animator::start) }
+        return true
+    }
+
+    private fun createCloseAnimation(
+        change: TransitionInfo.Change,
+        finishTransaction: Transaction,
+        onAnimFinish: (Animator) -> Unit,
+    ): Animator {
+        finishTransaction.hide(change.leash)
+        return AnimatorSet().apply {
+            playTogether(createBoundsCloseAnimation(change), createAlphaCloseAnimation(change))
+            addListener(onEnd = onAnimFinish)
+        }
+    }
+
+    private fun createBoundsCloseAnimation(change: TransitionInfo.Change): Animator {
+        val startBounds = change.startAbsBounds
+        val endBounds =
+            Rect(startBounds).apply {
+                // Scale the end bounds of the window down with an anchor in the center
+                inset(
+                    (startBounds.width().toFloat() * (1 - CLOSE_ANIM_SCALE) / 2).toInt(),
+                    (startBounds.height().toFloat() * (1 - CLOSE_ANIM_SCALE) / 2).toInt()
+                )
+                val offsetY =
+                    TypedValue.applyDimension(
+                            TypedValue.COMPLEX_UNIT_DIP,
+                            CLOSE_ANIM_OFFSET_Y,
+                            context.resources.displayMetrics
+                        )
+                        .toInt()
+                offset(/* dx= */ 0, offsetY)
+            }
+        return ValueAnimator.ofObject(RectEvaluator(), startBounds, endBounds).apply {
+            duration = CLOSE_ANIM_DURATION_BOUNDS
+            interpolator = Interpolators.STANDARD_ACCELERATE
+            addUpdateListener { animation ->
+                val animBounds = animation.animatedValue as Rect
+                val animScale = 1 - (1 - CLOSE_ANIM_SCALE) * animation.animatedFraction
+                transactionSupplier
+                    .get()
+                    .setPosition(change.leash, animBounds.left.toFloat(), animBounds.top.toFloat())
+                    .setScale(change.leash, animScale, animScale)
+                    .apply()
+            }
+        }
+    }
+
+    private fun createAlphaCloseAnimation(change: TransitionInfo.Change): Animator =
+        ValueAnimator.ofFloat(1f, 0f).apply {
+            duration = CLOSE_ANIM_DURATION_ALPHA
+            interpolator = Interpolators.LINEAR
+            addUpdateListener { animation ->
+                transactionSupplier
+                    .get()
+                    .setAlpha(change.leash, animation.animatedValue as Float)
+                    .apply()
+            }
+        }
+
+    private companion object {
+        const val CLOSE_ANIM_DURATION_BOUNDS = 200L
+        const val CLOSE_ANIM_DURATION_ALPHA = 100L
+        const val CLOSE_ANIM_SCALE = 0.95f
+        const val CLOSE_ANIM_OFFSET_Y = 36.0f
+    }
+}
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopMixedTransitionHandler.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopMixedTransitionHandler.kt
new file mode 100644
index 0000000..ec3f8c5
--- /dev/null
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopMixedTransitionHandler.kt
@@ -0,0 +1,157 @@
+/*
+ * Copyright (C) 2024 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 android.app.ActivityTaskManager.INVALID_TASK_ID
+import android.app.WindowConfiguration.WINDOWING_MODE_FREEFORM
+import android.content.Context
+import android.os.Handler
+import android.os.IBinder
+import android.view.SurfaceControl
+import android.view.WindowManager
+import android.window.TransitionInfo
+import android.window.TransitionRequestInfo
+import android.window.WindowContainerTransaction
+import com.android.internal.jank.Cuj.CUJ_DESKTOP_MODE_EXIT_MODE_ON_LAST_WINDOW_CLOSE
+import com.android.internal.jank.InteractionJankMonitor
+import com.android.internal.protolog.ProtoLog
+import com.android.wm.shell.freeform.FreeformTaskTransitionHandler
+import com.android.wm.shell.freeform.FreeformTaskTransitionStarter
+import com.android.wm.shell.protolog.ShellProtoLogGroup.WM_SHELL_DESKTOP_MODE
+import com.android.wm.shell.shared.annotations.ShellMainThread
+import com.android.wm.shell.transition.MixedTransitionHandler
+import com.android.wm.shell.transition.Transitions
+
+/** The [Transitions.TransitionHandler] coordinates transition handlers in desktop windowing. */
+class DesktopMixedTransitionHandler(
+    private val context: Context,
+    private val transitions: Transitions,
+    private val desktopTaskRepository: DesktopModeTaskRepository,
+    private val freeformTaskTransitionHandler: FreeformTaskTransitionHandler,
+    private val closeDesktopTaskTransitionHandler: CloseDesktopTaskTransitionHandler,
+    private val interactionJankMonitor: InteractionJankMonitor,
+    @ShellMainThread private val handler: Handler,
+) : MixedTransitionHandler, FreeformTaskTransitionStarter {
+
+    /** Delegates starting transition to [FreeformTaskTransitionHandler]. */
+    override fun startWindowingModeTransition(
+        targetWindowingMode: Int,
+        wct: WindowContainerTransaction?,
+    ) = freeformTaskTransitionHandler.startWindowingModeTransition(targetWindowingMode, wct)
+
+    /** Delegates starting minimized mode transition to [FreeformTaskTransitionHandler]. */
+    override fun startMinimizedModeTransition(wct: WindowContainerTransaction?): IBinder =
+        freeformTaskTransitionHandler.startMinimizedModeTransition(wct)
+
+    /** Starts close transition and handles or delegates desktop task close animation. */
+    override fun startRemoveTransition(wct: WindowContainerTransaction?) {
+        requireNotNull(wct)
+        transitions.startTransition(WindowManager.TRANSIT_CLOSE, wct, /* handler= */ this)
+    }
+
+    /** Returns null, as it only handles transitions started from Shell. */
+    override fun handleRequest(
+        transition: IBinder,
+        request: TransitionRequestInfo,
+    ): WindowContainerTransaction? = null
+
+    override fun startAnimation(
+        transition: IBinder,
+        info: TransitionInfo,
+        startTransaction: SurfaceControl.Transaction,
+        finishTransaction: SurfaceControl.Transaction,
+        finishCallback: Transitions.TransitionFinishCallback,
+    ): Boolean {
+        val closeChange = findCloseDesktopTaskChange(info)
+        if (closeChange == null) {
+            ProtoLog.w(WM_SHELL_DESKTOP_MODE, "%s: Should have closing desktop task", TAG)
+            return false
+        }
+        if (isLastDesktopTask(closeChange)) {
+            // Dispatch close desktop task animation to the default transition handlers.
+            return dispatchCloseLastDesktopTaskAnimation(
+                transition,
+                info,
+                closeChange,
+                startTransaction,
+                finishTransaction,
+                finishCallback,
+            )
+        }
+        // Animate close desktop task transition with [CloseDesktopTaskTransitionHandler].
+        return closeDesktopTaskTransitionHandler.startAnimation(
+            transition,
+            info,
+            startTransaction,
+            finishTransaction,
+            finishCallback,
+        )
+    }
+
+    /**
+     * Dispatch close desktop task animation to the default transition handlers. Allows delegating
+     * it to Launcher to animate in sync with show Home transition.
+     */
+    private fun dispatchCloseLastDesktopTaskAnimation(
+        transition: IBinder,
+        info: TransitionInfo,
+        change: TransitionInfo.Change,
+        startTransaction: SurfaceControl.Transaction,
+        finishTransaction: SurfaceControl.Transaction,
+        finishCallback: Transitions.TransitionFinishCallback,
+    ): Boolean {
+        // Starting the jank trace if closing the last window in desktop mode.
+        interactionJankMonitor.begin(
+            change.leash,
+            context,
+            handler,
+            CUJ_DESKTOP_MODE_EXIT_MODE_ON_LAST_WINDOW_CLOSE,
+        )
+        // Dispatch the last desktop task closing animation.
+        return transitions.dispatchTransition(
+            transition,
+            info,
+            startTransaction,
+            finishTransaction,
+            { wct ->
+                // Finish the jank trace when closing the last window in desktop mode.
+                interactionJankMonitor.end(CUJ_DESKTOP_MODE_EXIT_MODE_ON_LAST_WINDOW_CLOSE)
+                finishCallback.onTransitionFinished(wct)
+            },
+            /* skip= */ this
+        ) != null
+    }
+
+    private fun isLastDesktopTask(change: TransitionInfo.Change): Boolean =
+        change.taskInfo?.let {
+            desktopTaskRepository.getActiveNonMinimizedTaskCount(it.displayId) == 1
+        } ?: false
+
+    private fun findCloseDesktopTaskChange(info: TransitionInfo): TransitionInfo.Change? {
+        if (info.type != WindowManager.TRANSIT_CLOSE) return null
+        return info.changes.firstOrNull { change ->
+            change.mode == WindowManager.TRANSIT_CLOSE &&
+                !change.hasFlags(TransitionInfo.FLAG_IS_WALLPAPER) &&
+                change.taskInfo?.taskId != INVALID_TASK_ID &&
+                change.taskInfo?.windowingMode == WINDOWING_MODE_FREEFORM
+        }
+    }
+
+    companion object {
+        private const val TAG = "DesktopMixedTransitionHandler"
+    }
+}
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/freeform/FreeformComponents.java b/libs/WindowManager/Shell/src/com/android/wm/shell/freeform/FreeformComponents.java
index eee5aae..3379ff2 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/freeform/FreeformComponents.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/freeform/FreeformComponents.java
@@ -35,6 +35,7 @@
     public final ShellTaskOrganizer.TaskListener mTaskListener;
     public final Optional<Transitions.TransitionHandler> mTransitionHandler;
     public final Optional<Transitions.TransitionObserver> mTransitionObserver;
+    public final Optional<FreeformTaskTransitionStarterInitializer> mTransitionStarterInitializer;
 
     /**
      * Creates an instance with the given components.
@@ -42,10 +43,12 @@
     public FreeformComponents(
             ShellTaskOrganizer.TaskListener taskListener,
             Optional<Transitions.TransitionHandler> transitionHandler,
-            Optional<Transitions.TransitionObserver> transitionObserver) {
+            Optional<Transitions.TransitionObserver> transitionObserver,
+            Optional<FreeformTaskTransitionStarterInitializer> transitionStarterInitializer) {
         mTaskListener = taskListener;
         mTransitionHandler = transitionHandler;
         mTransitionObserver = transitionObserver;
+        mTransitionStarterInitializer = transitionStarterInitializer;
     }
 
     /**
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/freeform/FreeformTaskTransitionHandler.java b/libs/WindowManager/Shell/src/com/android/wm/shell/freeform/FreeformTaskTransitionHandler.java
index 517e209..6aaf001 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/freeform/FreeformTaskTransitionHandler.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/freeform/FreeformTaskTransitionHandler.java
@@ -19,16 +19,12 @@
 import static android.app.WindowConfiguration.WINDOWING_MODE_FREEFORM;
 import static android.app.WindowConfiguration.WINDOWING_MODE_FULLSCREEN;
 
-import static com.android.internal.jank.Cuj.CUJ_DESKTOP_MODE_EXIT_MODE_ON_LAST_WINDOW_CLOSE;
-
 import android.animation.Animator;
 import android.animation.AnimatorListenerAdapter;
 import android.animation.ValueAnimator;
 import android.app.ActivityManager;
 import android.app.WindowConfiguration;
-import android.content.Context;
 import android.graphics.Rect;
-import android.os.Handler;
 import android.os.IBinder;
 import android.util.ArrayMap;
 import android.view.SurfaceControl;
@@ -40,14 +36,9 @@
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
 
-import com.android.internal.jank.InteractionJankMonitor;
 import com.android.wm.shell.common.DisplayController;
 import com.android.wm.shell.common.ShellExecutor;
-import com.android.wm.shell.desktopmode.DesktopModeTaskRepository;
-import com.android.wm.shell.shared.annotations.ShellMainThread;
-import com.android.wm.shell.sysui.ShellInit;
 import com.android.wm.shell.transition.Transitions;
-import com.android.wm.shell.windowdecor.WindowDecorViewModel;
 
 import java.util.ArrayList;
 import java.util.List;
@@ -59,48 +50,24 @@
 public class FreeformTaskTransitionHandler
         implements Transitions.TransitionHandler, FreeformTaskTransitionStarter {
     private static final int CLOSE_ANIM_DURATION = 400;
-    private final Context mContext;
     private final Transitions mTransitions;
-    private final WindowDecorViewModel mWindowDecorViewModel;
-    private final DesktopModeTaskRepository mDesktopModeTaskRepository;
     private final DisplayController mDisplayController;
-    private final InteractionJankMonitor mInteractionJankMonitor;
     private final ShellExecutor mMainExecutor;
     private final ShellExecutor mAnimExecutor;
-    @ShellMainThread
-    private final Handler mHandler;
 
     private final List<IBinder> mPendingTransitionTokens = new ArrayList<>();
 
     private final ArrayMap<IBinder, ArrayList<Animator>> mAnimations = new ArrayMap<>();
 
     public FreeformTaskTransitionHandler(
-            ShellInit shellInit,
             Transitions transitions,
-            Context context,
-            WindowDecorViewModel windowDecorViewModel,
             DisplayController displayController,
             ShellExecutor mainExecutor,
-            ShellExecutor animExecutor,
-            DesktopModeTaskRepository desktopModeTaskRepository,
-            InteractionJankMonitor interactionJankMonitor,
-            @ShellMainThread Handler handler) {
+            ShellExecutor animExecutor) {
         mTransitions = transitions;
-        mContext = context;
-        mWindowDecorViewModel = windowDecorViewModel;
-        mDesktopModeTaskRepository = desktopModeTaskRepository;
         mDisplayController = displayController;
-        mInteractionJankMonitor = interactionJankMonitor;
         mMainExecutor = mainExecutor;
         mAnimExecutor = animExecutor;
-        mHandler = handler;
-        if (Transitions.ENABLE_SHELL_TRANSITIONS) {
-            shellInit.addInitCallback(this::onInit, this);
-        }
-    }
-
-    private void onInit() {
-        mWindowDecorViewModel.setFreeformTaskTransitionStarter(this);
     }
 
     @Override
@@ -269,20 +236,12 @@
                     startBounds.top + (animation.getAnimatedFraction() * screenHeight));
             t.apply();
         });
-        if (mDesktopModeTaskRepository.getActiveNonMinimizedTaskCount(
-                        change.getTaskInfo().displayId) == 1) {
-            // Starting the jank trace if closing the last window in desktop mode.
-            mInteractionJankMonitor.begin(
-                    sc, mContext, mHandler, CUJ_DESKTOP_MODE_EXIT_MODE_ON_LAST_WINDOW_CLOSE);
-        }
         animator.addListener(
                 new AnimatorListenerAdapter() {
                     @Override
                     public void onAnimationEnd(Animator animation) {
                         animations.remove(animator);
                         onAnimFinish.run();
-                        mInteractionJankMonitor.end(
-                                CUJ_DESKTOP_MODE_EXIT_MODE_ON_LAST_WINDOW_CLOSE);
                     }
                 });
         animations.add(animator);
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/freeform/FreeformTaskTransitionStarterInitializer.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/freeform/FreeformTaskTransitionStarterInitializer.kt
new file mode 100644
index 0000000..98bdf05
--- /dev/null
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/freeform/FreeformTaskTransitionStarterInitializer.kt
@@ -0,0 +1,41 @@
+/*
+ * Copyright (C) 2024 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.freeform
+
+import com.android.wm.shell.sysui.ShellInit
+import com.android.wm.shell.windowdecor.WindowDecorViewModel
+
+/**
+ * Sets up [FreeformTaskTransitionStarter] for [WindowDecorViewModel] when shell finishes
+ * initializing.
+ *
+ * Used to extract the setup logic from the starter implementation.
+ */
+class FreeformTaskTransitionStarterInitializer(
+    shellInit: ShellInit,
+    private val windowDecorViewModel: WindowDecorViewModel,
+    private val freeformTaskTransitionStarter: FreeformTaskTransitionStarter
+) {
+    init {
+        shellInit.addInitCallback(::onShellInit, this)
+    }
+
+    /** Sets up [WindowDecorViewModel] transition starter with [FreeformTaskTransitionStarter] */
+    private fun onShellInit() {
+        windowDecorViewModel.setFreeformTaskTransitionStarter(freeformTaskTransitionStarter)
+    }
+}
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 d03832d..70aaac4 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
@@ -1026,9 +1026,14 @@
      * Gives every handler (in order) a chance to animate until one consumes the transition.
      * @return the handler which consumed the transition.
      */
-    TransitionHandler dispatchTransition(@NonNull IBinder transition, @NonNull TransitionInfo info,
-            @NonNull SurfaceControl.Transaction startT, @NonNull SurfaceControl.Transaction finishT,
-            @NonNull TransitionFinishCallback finishCB, @Nullable TransitionHandler skip) {
+    public TransitionHandler dispatchTransition(
+            @NonNull IBinder transition,
+            @NonNull TransitionInfo info,
+            @NonNull SurfaceControl.Transaction startT,
+            @NonNull SurfaceControl.Transaction finishT,
+            @NonNull TransitionFinishCallback finishCB,
+            @Nullable TransitionHandler skip
+    ) {
         for (int i = mHandlers.size() - 1; i >= 0; --i) {
             if (mHandlers.get(i) == skip) continue;
             ProtoLog.v(ShellProtoLogGroup.WM_SHELL_TRANSITIONS, " try handler %s",
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/CloseDesktopTaskTransitionHandlerTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/CloseDesktopTaskTransitionHandlerTest.kt
new file mode 100644
index 0000000..9b4cc17
--- /dev/null
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/CloseDesktopTaskTransitionHandlerTest.kt
@@ -0,0 +1,160 @@
+/*
+ * Copyright (C) 2024 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 android.app.ActivityManager.RunningTaskInfo
+import android.app.WindowConfiguration.ACTIVITY_TYPE_STANDARD
+import android.app.WindowConfiguration.WINDOWING_MODE_FREEFORM
+import android.app.WindowConfiguration.WINDOWING_MODE_FULLSCREEN
+import android.app.WindowConfiguration.WindowingMode
+import android.testing.AndroidTestingRunner
+import android.testing.TestableLooper.RunWithLooper
+import android.view.SurfaceControl
+import android.view.WindowManager
+import android.window.TransitionInfo
+import androidx.test.filters.SmallTest
+import com.android.wm.shell.ShellTestCase
+import com.android.wm.shell.TestRunningTaskInfoBuilder
+import com.android.wm.shell.common.ShellExecutor
+import java.util.function.Supplier
+import org.junit.Assert.assertFalse
+import org.junit.Assert.assertNull
+import org.junit.Assert.assertTrue
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.Mock
+import org.mockito.kotlin.mock
+
+/**
+ * Test class for [CloseDesktopTaskTransitionHandler]
+ *
+ * Usage: atest WMShellUnitTests:CloseDesktopTaskTransitionHandlerTest
+ */
+@SmallTest
+@RunWithLooper
+@RunWith(AndroidTestingRunner::class)
+class CloseDesktopTaskTransitionHandlerTest : ShellTestCase() {
+
+    @Mock lateinit var testExecutor: ShellExecutor
+    @Mock lateinit var closingTaskLeash: SurfaceControl
+
+    private val transactionSupplier = Supplier { mock<SurfaceControl.Transaction>() }
+
+    private lateinit var handler: CloseDesktopTaskTransitionHandler
+
+    @Before
+    fun setUp() {
+        handler =
+            CloseDesktopTaskTransitionHandler(
+                context,
+                testExecutor,
+                testExecutor,
+                transactionSupplier
+            )
+    }
+
+    @Test
+    fun handleRequest_returnsNull() {
+        assertNull(handler.handleRequest(mock(), mock()))
+    }
+
+    @Test
+    fun startAnimation_openTransition_returnsFalse() {
+        val animates =
+            handler.startAnimation(
+                transition = mock(),
+                info =
+                    createTransitionInfo(
+                        type = WindowManager.TRANSIT_OPEN,
+                        task = createTask(WINDOWING_MODE_FREEFORM)
+                    ),
+                startTransaction = mock(),
+                finishTransaction = mock(),
+                finishCallback = {}
+            )
+
+        assertFalse("Should not animate open transition", animates)
+    }
+
+    @Test
+    fun startAnimation_closeTransitionFullscreenTask_returnsFalse() {
+        val animates =
+            handler.startAnimation(
+                transition = mock(),
+                info = createTransitionInfo(task = createTask(WINDOWING_MODE_FULLSCREEN)),
+                startTransaction = mock(),
+                finishTransaction = mock(),
+                finishCallback = {}
+            )
+
+        assertFalse("Should not animate fullscreen task close transition", animates)
+    }
+
+    @Test
+    fun startAnimation_closeTransitionOpeningFreeformTask_returnsFalse() {
+        val animates =
+            handler.startAnimation(
+                transition = mock(),
+                info =
+                    createTransitionInfo(
+                        changeMode = WindowManager.TRANSIT_OPEN,
+                        task = createTask(WINDOWING_MODE_FREEFORM)
+                    ),
+                startTransaction = mock(),
+                finishTransaction = mock(),
+                finishCallback = {}
+            )
+
+        assertFalse("Should not animate opening freeform task close transition", animates)
+    }
+
+    @Test
+    fun startAnimation_closeTransitionClosingFreeformTask_returnsTrue() {
+        val animates =
+            handler.startAnimation(
+                transition = mock(),
+                info = createTransitionInfo(task = createTask(WINDOWING_MODE_FREEFORM)),
+                startTransaction = mock(),
+                finishTransaction = mock(),
+                finishCallback = {}
+            )
+
+        assertTrue("Should animate closing freeform task close transition", animates)
+    }
+
+    private fun createTransitionInfo(
+        type: Int = WindowManager.TRANSIT_CLOSE,
+        changeMode: Int = WindowManager.TRANSIT_CLOSE,
+        task: RunningTaskInfo
+    ): TransitionInfo =
+        TransitionInfo(type, 0 /* flags */).apply {
+            addChange(
+                TransitionInfo.Change(mock(), closingTaskLeash).apply {
+                    mode = changeMode
+                    parent = null
+                    taskInfo = task
+                }
+            )
+        }
+
+    private fun createTask(@WindowingMode windowingMode: Int): RunningTaskInfo =
+        TestRunningTaskInfoBuilder()
+            .setActivityType(ACTIVITY_TYPE_STANDARD)
+            .setWindowingMode(windowingMode)
+            .build()
+}
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopMixedTransitionHandlerTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopMixedTransitionHandlerTest.kt
new file mode 100644
index 0000000..2b60200
--- /dev/null
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopMixedTransitionHandlerTest.kt
@@ -0,0 +1,220 @@
+/*
+ * Copyright (C) 2024 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 android.app.ActivityManager.RunningTaskInfo
+import android.app.WindowConfiguration.ACTIVITY_TYPE_STANDARD
+import android.app.WindowConfiguration.WINDOWING_MODE_FREEFORM
+import android.app.WindowConfiguration.WINDOWING_MODE_FULLSCREEN
+import android.app.WindowConfiguration.WindowingMode
+import android.os.Handler
+import android.os.IBinder
+import android.testing.AndroidTestingRunner
+import android.testing.TestableLooper.RunWithLooper
+import android.view.SurfaceControl
+import android.view.WindowManager
+import android.window.TransitionInfo
+import android.window.WindowContainerTransaction
+import androidx.test.filters.SmallTest
+import com.android.internal.jank.Cuj.CUJ_DESKTOP_MODE_EXIT_MODE_ON_LAST_WINDOW_CLOSE
+import com.android.internal.jank.InteractionJankMonitor
+import com.android.wm.shell.ShellTestCase
+import com.android.wm.shell.TestRunningTaskInfoBuilder
+import com.android.wm.shell.freeform.FreeformTaskTransitionHandler
+import com.android.wm.shell.transition.Transitions
+import org.junit.Assert.assertFalse
+import org.junit.Assert.assertNull
+import org.junit.Assert.assertTrue
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.Mock
+import org.mockito.Mockito.verify
+import org.mockito.kotlin.any
+import org.mockito.kotlin.eq
+import org.mockito.kotlin.mock
+import org.mockito.kotlin.whenever
+
+/**
+ * Test class for [DesktopMixedTransitionHandler]
+ *
+ * Usage: atest WMShellUnitTests:DesktopMixedTransitionHandlerTest
+ */
+@SmallTest
+@RunWithLooper
+@RunWith(AndroidTestingRunner::class)
+class DesktopMixedTransitionHandlerTest : ShellTestCase() {
+
+    @Mock lateinit var transitions: Transitions
+    @Mock lateinit var desktopTaskRepository: DesktopModeTaskRepository
+    @Mock lateinit var freeformTaskTransitionHandler: FreeformTaskTransitionHandler
+    @Mock lateinit var closeDesktopTaskTransitionHandler: CloseDesktopTaskTransitionHandler
+    @Mock lateinit var interactionJankMonitor: InteractionJankMonitor
+    @Mock lateinit var mockHandler: Handler
+    @Mock lateinit var closingTaskLeash: SurfaceControl
+
+    private lateinit var mixedHandler: DesktopMixedTransitionHandler
+
+    @Before
+    fun setUp() {
+        mixedHandler =
+            DesktopMixedTransitionHandler(
+                context,
+                transitions,
+                desktopTaskRepository,
+                freeformTaskTransitionHandler,
+                closeDesktopTaskTransitionHandler,
+                interactionJankMonitor,
+                mockHandler
+            )
+    }
+
+    @Test
+    fun startWindowingModeTransition_callsFreeformTaskTransitionHandler() {
+        val windowingMode = WINDOWING_MODE_FULLSCREEN
+        val wct = WindowContainerTransaction()
+
+        mixedHandler.startWindowingModeTransition(windowingMode, wct)
+
+        verify(freeformTaskTransitionHandler).startWindowingModeTransition(windowingMode, wct)
+    }
+
+    @Test
+    fun startMinimizedModeTransition_callsFreeformTaskTransitionHandler() {
+        val wct = WindowContainerTransaction()
+        whenever(freeformTaskTransitionHandler.startMinimizedModeTransition(any()))
+            .thenReturn(mock())
+
+        mixedHandler.startMinimizedModeTransition(wct)
+
+        verify(freeformTaskTransitionHandler).startMinimizedModeTransition(wct)
+    }
+
+    @Test
+    fun startRemoveTransition_startsCloseTransition() {
+        val wct = WindowContainerTransaction()
+
+        mixedHandler.startRemoveTransition(wct)
+
+        verify(transitions).startTransition(WindowManager.TRANSIT_CLOSE, wct, mixedHandler)
+    }
+
+    @Test
+    fun handleRequest_returnsNull() {
+        assertNull(mixedHandler.handleRequest(mock(), mock()))
+    }
+
+    @Test
+    fun startAnimation_withoutClosingDesktopTask_returnsFalse() {
+        val transition = mock<IBinder>()
+        val transitionInfo =
+            createTransitionInfo(
+                changeMode = WindowManager.TRANSIT_OPEN,
+                task = createTask(WINDOWING_MODE_FREEFORM)
+            )
+        whenever(freeformTaskTransitionHandler.startAnimation(any(), any(), any(), any(), any()))
+            .thenReturn(true)
+
+        val started = mixedHandler.startAnimation(
+            transition = transition,
+            info = transitionInfo,
+            startTransaction = mock(),
+            finishTransaction = mock(),
+            finishCallback = {}
+        )
+
+        assertFalse("Should not start animation without closing desktop task", started)
+    }
+
+    @Test
+    fun startAnimation_withClosingDesktopTask_callsCloseTaskHandler() {
+        val transition = mock<IBinder>()
+        val transitionInfo = createTransitionInfo(task = createTask(WINDOWING_MODE_FREEFORM))
+        whenever(desktopTaskRepository.getActiveNonMinimizedTaskCount(any())).thenReturn(2)
+        whenever(
+                closeDesktopTaskTransitionHandler.startAnimation(any(), any(), any(), any(), any())
+            )
+            .thenReturn(true)
+
+        val started = mixedHandler.startAnimation(
+            transition = transition,
+            info = transitionInfo,
+            startTransaction = mock(),
+            finishTransaction = mock(),
+            finishCallback = {}
+        )
+
+        assertTrue("Should delegate animation to close transition handler", started)
+        verify(closeDesktopTaskTransitionHandler)
+            .startAnimation(eq(transition), eq(transitionInfo), any(), any(), any())
+    }
+
+    @Test
+    fun startAnimation_withClosingLastDesktopTask_dispatchesTransition() {
+        val transition = mock<IBinder>()
+        val transitionInfo = createTransitionInfo(task = createTask(WINDOWING_MODE_FREEFORM))
+        whenever(desktopTaskRepository.getActiveNonMinimizedTaskCount(any())).thenReturn(1)
+        whenever(transitions.dispatchTransition(any(), any(), any(), any(), any(), any()))
+            .thenReturn(mock())
+
+        mixedHandler.startAnimation(
+            transition = transition,
+            info = transitionInfo,
+            startTransaction = mock(),
+            finishTransaction = mock(),
+            finishCallback = {}
+        )
+
+        verify(transitions)
+            .dispatchTransition(
+                eq(transition),
+                eq(transitionInfo),
+                any(),
+                any(),
+                any(),
+                eq(mixedHandler)
+            )
+        verify(interactionJankMonitor)
+            .begin(
+                closingTaskLeash,
+                context,
+                mockHandler,
+                CUJ_DESKTOP_MODE_EXIT_MODE_ON_LAST_WINDOW_CLOSE
+            )
+    }
+
+    private fun createTransitionInfo(
+        type: Int = WindowManager.TRANSIT_CLOSE,
+        changeMode: Int = WindowManager.TRANSIT_CLOSE,
+        task: RunningTaskInfo
+    ): TransitionInfo =
+        TransitionInfo(type, 0 /* flags */).apply {
+            addChange(
+                TransitionInfo.Change(mock(), closingTaskLeash).apply {
+                    mode = changeMode
+                    parent = null
+                    taskInfo = task
+                }
+            )
+        }
+
+    private fun createTask(@WindowingMode windowingMode: Int): RunningTaskInfo =
+        TestRunningTaskInfoBuilder()
+            .setActivityType(ACTIVITY_TYPE_STANDARD)
+            .setWindowingMode(windowingMode)
+            .build()
+}