Merge changes I9907038a,Ic5de4d43,I0435e07a into tm-qpr-dev

* changes:
  Implement move to desktop and fullscreen
  Implement showDesktopApps
  Create controller for desktop prototype 2
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellBaseModule.java b/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellBaseModule.java
index c743582..09f5cf1 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellBaseModule.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellBaseModule.java
@@ -61,6 +61,7 @@
 import com.android.wm.shell.desktopmode.DesktopModeController;
 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.displayareahelper.DisplayAreaHelper;
 import com.android.wm.shell.displayareahelper.DisplayAreaHelperController;
 import com.android.wm.shell.draganddrop.DragAndDropController;
@@ -677,7 +678,11 @@
     @WMSingleton
     @Provides
     static Optional<DesktopMode> provideDesktopMode(
-            Optional<DesktopModeController> desktopModeController) {
+            Optional<DesktopModeController> desktopModeController,
+            Optional<DesktopTasksController> desktopTasksController) {
+        if (DesktopModeStatus.isProto2Enabled()) {
+            return desktopTasksController.map(DesktopTasksController::asDesktopMode);
+        }
         return desktopModeController.map(DesktopModeController::asDesktopMode);
     }
 
@@ -700,6 +705,23 @@
 
     @BindsOptionalOf
     @DynamicOverride
+    abstract DesktopTasksController optionalDesktopTasksController();
+
+    @WMSingleton
+    @Provides
+    static Optional<DesktopTasksController> providesDesktopTasksController(
+            @DynamicOverride Optional<Lazy<DesktopTasksController>> desktopTasksController) {
+        // Use optional-of-lazy for the dependency that this provider relies on.
+        // Lazy ensures that this provider will not be the cause the dependency is created
+        // when it will not be returned due to the condition below.
+        if (DesktopModeStatus.isProto2Enabled()) {
+            return desktopTasksController.map(Lazy::get);
+        }
+        return Optional.empty();
+    }
+
+    @BindsOptionalOf
+    @DynamicOverride
     abstract DesktopModeTaskRepository optionalDesktopModeTaskRepository();
 
     @WMSingleton
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 6be8305..701a3a4 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
@@ -50,6 +50,7 @@
 import com.android.wm.shell.common.annotations.ShellMainThread;
 import com.android.wm.shell.desktopmode.DesktopModeController;
 import com.android.wm.shell.desktopmode.DesktopModeTaskRepository;
+import com.android.wm.shell.desktopmode.DesktopTasksController;
 import com.android.wm.shell.draganddrop.DragAndDropController;
 import com.android.wm.shell.freeform.FreeformComponents;
 import com.android.wm.shell.freeform.FreeformTaskListener;
@@ -189,7 +190,8 @@
             ShellTaskOrganizer taskOrganizer,
             DisplayController displayController,
             SyncTransactionQueue syncQueue,
-            Optional<DesktopModeController> desktopModeController) {
+            Optional<DesktopModeController> desktopModeController,
+            Optional<DesktopTasksController> desktopTasksController) {
         return new CaptionWindowDecorViewModel(
                     context,
                     mainHandler,
@@ -197,7 +199,8 @@
                     taskOrganizer,
                     displayController,
                     syncQueue,
-                    desktopModeController);
+                    desktopModeController,
+                    desktopTasksController);
     }
 
     //
@@ -616,6 +619,22 @@
     @WMSingleton
     @Provides
     @DynamicOverride
+    static DesktopTasksController provideDesktopTasksController(
+            Context context,
+            ShellInit shellInit,
+            ShellController shellController,
+            ShellTaskOrganizer shellTaskOrganizer,
+            Transitions transitions,
+            @DynamicOverride DesktopModeTaskRepository desktopModeTaskRepository,
+            @ShellMainThread ShellExecutor mainExecutor
+    ) {
+        return new DesktopTasksController(context, shellInit, shellController, shellTaskOrganizer,
+                transitions, desktopModeTaskRepository, mainExecutor);
+    }
+
+    @WMSingleton
+    @Provides
+    @DynamicOverride
     static DesktopModeTaskRepository provideDesktopModeTaskRepository() {
         return new DesktopModeTaskRepository();
     }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopModeStatus.java b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopModeStatus.java
index 67f4a19..055949f 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopModeStatus.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopModeStatus.java
@@ -70,9 +70,13 @@
      * @return {@code true} if active
      */
     public static boolean isActive(Context context) {
-        if (!IS_SUPPORTED) {
+        if (!isAnyEnabled()) {
             return false;
         }
+        if (isProto2Enabled()) {
+            // Desktop mode is always active in prototype 2
+            return true;
+        }
         try {
             int result = Settings.System.getIntForUser(context.getContentResolver(),
                     Settings.System.DESKTOP_MODE, UserHandle.USER_CURRENT);
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
new file mode 100644
index 0000000..b075b14
--- /dev/null
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksController.kt
@@ -0,0 +1,231 @@
+/*
+ * Copyright (C) 2022 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
+import android.app.WindowConfiguration
+import android.app.WindowConfiguration.ACTIVITY_TYPE_HOME
+import android.app.WindowConfiguration.WINDOWING_MODE_UNDEFINED
+import android.app.WindowConfiguration.WindowingMode
+import android.content.Context
+import android.view.WindowManager
+import android.window.WindowContainerTransaction
+import androidx.annotation.BinderThread
+import com.android.internal.protolog.common.ProtoLog
+import com.android.wm.shell.ShellTaskOrganizer
+import com.android.wm.shell.common.ExecutorUtils
+import com.android.wm.shell.common.ExternalInterfaceBinder
+import com.android.wm.shell.common.RemoteCallable
+import com.android.wm.shell.common.ShellExecutor
+import com.android.wm.shell.common.annotations.ExternalThread
+import com.android.wm.shell.common.annotations.ShellMainThread
+import com.android.wm.shell.desktopmode.DesktopModeTaskRepository.VisibleTasksListener
+import com.android.wm.shell.protolog.ShellProtoLogGroup.WM_SHELL_DESKTOP_MODE
+import com.android.wm.shell.sysui.ShellController
+import com.android.wm.shell.sysui.ShellInit
+import com.android.wm.shell.sysui.ShellSharedConstants
+import com.android.wm.shell.transition.Transitions
+import java.util.concurrent.Executor
+import java.util.function.Consumer
+
+/** Handles moving tasks in and out of desktop */
+class DesktopTasksController(
+    private val context: Context,
+    shellInit: ShellInit,
+    private val shellController: ShellController,
+    private val shellTaskOrganizer: ShellTaskOrganizer,
+    private val transitions: Transitions,
+    private val desktopModeTaskRepository: DesktopModeTaskRepository,
+    @ShellMainThread private val mainExecutor: ShellExecutor
+) : RemoteCallable<DesktopTasksController> {
+
+    private val desktopMode: DesktopModeImpl
+
+    init {
+        desktopMode = DesktopModeImpl()
+        if (DesktopModeStatus.isProto2Enabled()) {
+            shellInit.addInitCallback({ onInit() }, this)
+        }
+    }
+
+    private fun onInit() {
+        ProtoLog.d(WM_SHELL_DESKTOP_MODE, "Initialize DesktopTasksController")
+        shellController.addExternalInterface(
+            ShellSharedConstants.KEY_EXTRA_SHELL_DESKTOP_MODE,
+            { createExternalInterface() },
+            this
+        )
+    }
+
+    /** Show all tasks, that are part of the desktop, on top of launcher */
+    fun showDesktopApps() {
+        ProtoLog.v(WM_SHELL_DESKTOP_MODE, "showDesktopApps")
+        val wct = WindowContainerTransaction()
+
+        bringDesktopAppsToFront(wct)
+
+        // Execute transaction if there are pending operations
+        if (!wct.isEmpty) {
+            if (Transitions.ENABLE_SHELL_TRANSITIONS) {
+                transitions.startTransition(WindowManager.TRANSIT_TO_FRONT, wct, null /* handler */)
+            } else {
+                shellTaskOrganizer.applyTransaction(wct)
+            }
+        }
+    }
+
+    /** Move a task with given `taskId` to desktop */
+    fun moveToDesktop(taskId: Int) {
+        shellTaskOrganizer.getRunningTaskInfo(taskId)?.let { task -> moveToDesktop(task) }
+    }
+
+    /** Move a task to desktop */
+    fun moveToDesktop(task: ActivityManager.RunningTaskInfo) {
+        ProtoLog.v(WM_SHELL_DESKTOP_MODE, "moveToDesktop: %d", task.taskId)
+
+        val wct = WindowContainerTransaction()
+        // Bring other apps to front first
+        bringDesktopAppsToFront(wct)
+
+        wct.setWindowingMode(task.getToken(), WindowConfiguration.WINDOWING_MODE_FREEFORM)
+        wct.reorder(task.getToken(), true /* onTop */)
+
+        if (Transitions.ENABLE_SHELL_TRANSITIONS) {
+            transitions.startTransition(WindowManager.TRANSIT_CHANGE, wct, null /* handler */)
+        } else {
+            shellTaskOrganizer.applyTransaction(wct)
+        }
+    }
+
+    /** Move a task with given `taskId` to fullscreen */
+    fun moveToFullscreen(taskId: Int) {
+        shellTaskOrganizer.getRunningTaskInfo(taskId)?.let { task -> moveToFullscreen(task) }
+    }
+
+    /** Move a task to fullscreen */
+    fun moveToFullscreen(task: ActivityManager.RunningTaskInfo) {
+        ProtoLog.v(WM_SHELL_DESKTOP_MODE, "moveToFullscreen: %d", task.taskId)
+
+        val wct = WindowContainerTransaction()
+        wct.setWindowingMode(task.getToken(), WindowConfiguration.WINDOWING_MODE_FULLSCREEN)
+        wct.setBounds(task.getToken(), null)
+        if (Transitions.ENABLE_SHELL_TRANSITIONS) {
+            transitions.startTransition(WindowManager.TRANSIT_CHANGE, wct, null /* handler */)
+        } else {
+            shellTaskOrganizer.applyTransaction(wct)
+        }
+    }
+
+    /**
+     * Get windowing move for a given `taskId`
+     *
+     * @return [WindowingMode] for the task or [WINDOWING_MODE_UNDEFINED] if task is not found
+     */
+    @WindowingMode
+    fun getTaskWindowingMode(taskId: Int): Int {
+        return shellTaskOrganizer.getRunningTaskInfo(taskId)?.windowingMode
+            ?: WINDOWING_MODE_UNDEFINED
+    }
+
+    private fun bringDesktopAppsToFront(wct: WindowContainerTransaction) {
+        val activeTasks = desktopModeTaskRepository.getActiveTasks()
+
+        // Skip if all tasks are already visible
+        if (activeTasks.isNotEmpty() && activeTasks.all(desktopModeTaskRepository::isVisibleTask)) {
+            ProtoLog.d(
+                WM_SHELL_DESKTOP_MODE,
+                "bringDesktopAppsToFront: active tasks are already in front, skipping."
+            )
+            return
+        }
+
+        ProtoLog.v(WM_SHELL_DESKTOP_MODE, "bringDesktopAppsToFront")
+
+        // First move home to front and then other tasks on top of it
+        moveHomeTaskToFront(wct)
+
+        val allTasksInZOrder = desktopModeTaskRepository.getFreeformTasksInZOrder()
+        activeTasks
+            // Sort descending as the top task is at index 0. It should be ordered to top last
+            .sortedByDescending { taskId -> allTasksInZOrder.indexOf(taskId) }
+            .mapNotNull { taskId -> shellTaskOrganizer.getRunningTaskInfo(taskId) }
+            .forEach { task -> wct.reorder(task.token, true /* onTop */) }
+    }
+
+    private fun moveHomeTaskToFront(wct: WindowContainerTransaction) {
+        shellTaskOrganizer
+            .getRunningTasks(context.displayId)
+            .firstOrNull { task -> task.activityType == ACTIVITY_TYPE_HOME }
+            ?.let { homeTask -> wct.reorder(homeTask.getToken(), true /* onTop */) }
+    }
+
+    override fun getContext(): Context {
+        return context
+    }
+
+    override fun getRemoteCallExecutor(): ShellExecutor {
+        return mainExecutor
+    }
+
+    /** Creates a new instance of the external interface to pass to another process. */
+    private fun createExternalInterface(): ExternalInterfaceBinder {
+        return IDesktopModeImpl(this)
+    }
+
+    /** Get connection interface between sysui and shell */
+    fun asDesktopMode(): DesktopMode {
+        return desktopMode
+    }
+
+    /**
+     * Adds a listener to find out about changes in the visibility of freeform tasks.
+     *
+     * @param listener the listener to add.
+     * @param callbackExecutor the executor to call the listener on.
+     */
+    fun addListener(listener: VisibleTasksListener, callbackExecutor: Executor) {
+        desktopModeTaskRepository.addVisibleTasksListener(listener, callbackExecutor)
+    }
+
+    /** The interface for calls from outside the shell, within the host process. */
+    @ExternalThread
+    private inner class DesktopModeImpl : DesktopMode {
+        override fun addListener(listener: VisibleTasksListener, callbackExecutor: Executor) {
+            mainExecutor.execute {
+                this@DesktopTasksController.addListener(listener, callbackExecutor)
+            }
+        }
+    }
+
+    /** The interface for calls from outside the host process. */
+    @BinderThread
+    private class IDesktopModeImpl(private var controller: DesktopTasksController?) :
+        IDesktopMode.Stub(), ExternalInterfaceBinder {
+        /** Invalidates this instance, preventing future calls from updating the controller. */
+        override fun invalidate() {
+            controller = null
+        }
+
+        override fun showDesktopApps() {
+            ExecutorUtils.executeRemoteCallWithTaskPermission(
+                controller,
+                "showDesktopApps",
+                Consumer(DesktopTasksController::showDesktopApps)
+            )
+        }
+    }
+}
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/CaptionWindowDecorViewModel.java b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/CaptionWindowDecorViewModel.java
index afefd5d..3cb40c5 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/CaptionWindowDecorViewModel.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/CaptionWindowDecorViewModel.java
@@ -52,6 +52,7 @@
 import com.android.wm.shell.common.SyncTransactionQueue;
 import com.android.wm.shell.desktopmode.DesktopModeController;
 import com.android.wm.shell.desktopmode.DesktopModeStatus;
+import com.android.wm.shell.desktopmode.DesktopTasksController;
 import com.android.wm.shell.freeform.FreeformTaskTransitionStarter;
 import com.android.wm.shell.transition.Transitions;
 
@@ -76,6 +77,7 @@
     private final SyncTransactionQueue mSyncQueue;
     private FreeformTaskTransitionStarter mTransitionStarter;
     private Optional<DesktopModeController> mDesktopModeController;
+    private Optional<DesktopTasksController> mDesktopTasksController;
     private boolean mTransitionDragActive;
 
     private SparseArray<EventReceiver> mEventReceiversByDisplay = new SparseArray<>();
@@ -91,7 +93,8 @@
             ShellTaskOrganizer taskOrganizer,
             DisplayController displayController,
             SyncTransactionQueue syncQueue,
-            Optional<DesktopModeController> desktopModeController) {
+            Optional<DesktopModeController> desktopModeController,
+            Optional<DesktopTasksController> desktopTasksController) {
         this(
                 context,
                 mainHandler,
@@ -100,6 +103,7 @@
                 displayController,
                 syncQueue,
                 desktopModeController,
+                desktopTasksController,
                 new CaptionWindowDecoration.Factory(),
                 InputManager::getInstance);
     }
@@ -112,6 +116,7 @@
             DisplayController displayController,
             SyncTransactionQueue syncQueue,
             Optional<DesktopModeController> desktopModeController,
+            Optional<DesktopTasksController> desktopTasksController,
             CaptionWindowDecoration.Factory captionWindowDecorFactory,
             Supplier<InputManager> inputManagerSupplier) {
 
@@ -123,6 +128,7 @@
         mDisplayController = displayController;
         mSyncQueue = syncQueue;
         mDesktopModeController = desktopModeController;
+        mDesktopTasksController = desktopTasksController;
 
         mCaptionWindowDecorFactory = captionWindowDecorFactory;
         mInputManagerSupplier = inputManagerSupplier;
@@ -248,11 +254,13 @@
                 decoration.createHandleMenu();
             } else if (id == R.id.desktop_button) {
                 mDesktopModeController.ifPresent(c -> c.setDesktopModeActive(true));
+                mDesktopTasksController.ifPresent(c -> c.moveToDesktop(mTaskId));
                 decoration.closeHandleMenu();
             } else if (id == R.id.fullscreen_button) {
                 mDesktopModeController.ifPresent(c -> c.setDesktopModeActive(false));
+                mDesktopTasksController.ifPresent(c -> c.moveToFullscreen(mTaskId));
                 decoration.closeHandleMenu();
-                decoration.setButtonVisibility();
+                decoration.setButtonVisibility(false);
             }
         }
 
@@ -305,8 +313,13 @@
          */
         private void handleEventForMove(MotionEvent e) {
             RunningTaskInfo taskInfo = mTaskOrganizer.getRunningTaskInfo(mTaskId);
-            if (mDesktopModeController.isPresent()
-                    && mDesktopModeController.get().getDisplayAreaWindowingMode(taskInfo.displayId)
+            if (DesktopModeStatus.isProto2Enabled()
+                    && taskInfo.getWindowingMode() == WINDOWING_MODE_FULLSCREEN) {
+                return;
+            }
+            if (DesktopModeStatus.isProto1Enabled() && mDesktopModeController.isPresent()
+                    && mDesktopModeController.get().getDisplayAreaWindowingMode(
+                    taskInfo.displayId)
                     == WINDOWING_MODE_FULLSCREEN) {
                 return;
             }
@@ -330,9 +343,20 @@
                             .stableInsets().top;
                     mDragResizeCallback.onDragResizeEnd(
                             e.getRawX(dragPointerIdx), e.getRawY(dragPointerIdx));
-                    if (e.getRawY(dragPointerIdx) <= statusBarHeight
-                            && DesktopModeStatus.isActive(mContext)) {
-                        mDesktopModeController.ifPresent(c -> c.setDesktopModeActive(false));
+                    if (e.getRawY(dragPointerIdx) <= statusBarHeight) {
+                        if (DesktopModeStatus.isProto2Enabled()) {
+                            if (taskInfo.getWindowingMode() == WINDOWING_MODE_FREEFORM) {
+                                // Switch a single task to fullscreen
+                                mDesktopTasksController.ifPresent(
+                                        c -> c.moveToFullscreen(taskInfo));
+                            }
+                        } else if (DesktopModeStatus.isProto1Enabled()) {
+                            if (DesktopModeStatus.isActive(mContext)) {
+                                // Turn off desktop mode
+                                mDesktopModeController.ifPresent(
+                                        c -> c.setDesktopModeActive(false));
+                            }
+                        }
                     }
                     break;
                 }
@@ -420,13 +444,27 @@
      * @param ev the {@link MotionEvent} received by {@link EventReceiver}
      */
     private void handleReceivedMotionEvent(MotionEvent ev, InputMonitor inputMonitor) {
-        if (!DesktopModeStatus.isActive(mContext)) {
-            handleCaptionThroughStatusBar(ev);
+        if (DesktopModeStatus.isProto2Enabled()) {
+            CaptionWindowDecoration focusedDecor = getFocusedDecor();
+            if (focusedDecor == null
+                    || focusedDecor.mTaskInfo.getWindowingMode() != WINDOWING_MODE_FREEFORM) {
+                handleCaptionThroughStatusBar(ev);
+            }
+        } else if (DesktopModeStatus.isProto1Enabled()) {
+            if (!DesktopModeStatus.isActive(mContext)) {
+                handleCaptionThroughStatusBar(ev);
+            }
         }
         handleEventOutsideFocusedCaption(ev);
         // Prevent status bar from reacting to a caption drag.
-        if (mTransitionDragActive && !DesktopModeStatus.isActive(mContext)) {
-            inputMonitor.pilferPointers();
+        if (DesktopModeStatus.isProto2Enabled()) {
+            if (mTransitionDragActive) {
+                inputMonitor.pilferPointers();
+            }
+        } else if (DesktopModeStatus.isProto1Enabled()) {
+            if (mTransitionDragActive && !DesktopModeStatus.isActive(mContext)) {
+                inputMonitor.pilferPointers();
+            }
         }
     }
 
@@ -455,9 +493,20 @@
             case MotionEvent.ACTION_DOWN: {
                 // Begin drag through status bar if applicable.
                 CaptionWindowDecoration focusedDecor = getFocusedDecor();
-                if (focusedDecor != null && !DesktopModeStatus.isActive(mContext)
-                        && focusedDecor.checkTouchEventInHandle(ev)) {
-                    mTransitionDragActive = true;
+                if (focusedDecor != null) {
+                    boolean dragFromStatusBarAllowed = false;
+                    if (DesktopModeStatus.isProto2Enabled()) {
+                        // In proto2 any full screen task can be dragged to freeform
+                        dragFromStatusBarAllowed = focusedDecor.mTaskInfo.getWindowingMode()
+                                == WINDOWING_MODE_FULLSCREEN;
+                    } else if (DesktopModeStatus.isProto1Enabled()) {
+                        // In proto1 task can be dragged to freeform when not in desktop mode
+                        dragFromStatusBarAllowed = !DesktopModeStatus.isActive(mContext);
+                    }
+
+                    if (dragFromStatusBarAllowed && focusedDecor.checkTouchEventInHandle(ev)) {
+                        mTransitionDragActive = true;
+                    }
                 }
                 break;
             }
@@ -472,7 +521,13 @@
                     int statusBarHeight = mDisplayController
                             .getDisplayLayout(focusedDecor.mTaskInfo.displayId).stableInsets().top;
                     if (ev.getY() > statusBarHeight) {
-                        mDesktopModeController.ifPresent(c -> c.setDesktopModeActive(true));
+                        if (DesktopModeStatus.isProto2Enabled()) {
+                            mDesktopTasksController.ifPresent(
+                                    c -> c.moveToDesktop(focusedDecor.mTaskInfo));
+                        } else if (DesktopModeStatus.isProto1Enabled()) {
+                            mDesktopModeController.ifPresent(c -> c.setDesktopModeActive(true));
+                        }
+
                         return;
                     }
                 }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/CaptionWindowDecoration.java b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/CaptionWindowDecoration.java
index 037ca20..f7c7a87 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/CaptionWindowDecoration.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/CaptionWindowDecoration.java
@@ -16,8 +16,9 @@
 
 package com.android.wm.shell.windowdecor;
 
+import static android.app.WindowConfiguration.WINDOWING_MODE_FREEFORM;
+
 import android.app.ActivityManager;
-import android.app.WindowConfiguration;
 import android.content.Context;
 import android.content.res.ColorStateList;
 import android.content.res.Resources;
@@ -117,7 +118,7 @@
                 ? R.dimen.freeform_decor_shadow_focused_thickness
                 : R.dimen.freeform_decor_shadow_unfocused_thickness;
         final boolean isFreeform =
-                taskInfo.getWindowingMode() == WindowConfiguration.WINDOWING_MODE_FREEFORM;
+                taskInfo.getWindowingMode() == WINDOWING_MODE_FREEFORM;
         final boolean isDragResizeable = isFreeform && taskInfo.isResizeable;
 
         WindowDecorLinearLayout oldRootView = mResult.mRootView;
@@ -167,11 +168,17 @@
         // If this task is not focused, do not show caption.
         setCaptionVisibility(mTaskInfo.isFocused);
 
-        // Only handle should show if Desktop Mode is inactive.
-        boolean desktopCurrentStatus = DesktopModeStatus.isActive(mContext);
-        if (mDesktopActive != desktopCurrentStatus && mTaskInfo.isFocused) {
-            mDesktopActive = desktopCurrentStatus;
-            setButtonVisibility();
+        if (mTaskInfo.isFocused) {
+            if (DesktopModeStatus.isProto2Enabled()) {
+                updateButtonVisibility();
+            } else if (DesktopModeStatus.isProto1Enabled()) {
+                // Only handle should show if Desktop Mode is inactive.
+                boolean desktopCurrentStatus = DesktopModeStatus.isActive(mContext);
+                if (mDesktopActive != desktopCurrentStatus) {
+                    mDesktopActive = desktopCurrentStatus;
+                    setButtonVisibility(mDesktopActive);
+                }
+            }
         }
 
         if (!isDragResizeable) {
@@ -214,7 +221,7 @@
         View handle = caption.findViewById(R.id.caption_handle);
         handle.setOnTouchListener(mOnCaptionTouchListener);
         handle.setOnClickListener(mOnCaptionButtonClickListener);
-        setButtonVisibility();
+        updateButtonVisibility();
     }
 
     private void setupHandleMenu() {
@@ -244,14 +251,25 @@
     /**
      * Sets the visibility of buttons and color of caption based on desktop mode status
      */
-    void setButtonVisibility() {
-        mDesktopActive = DesktopModeStatus.isActive(mContext);
-        int v = mDesktopActive ? View.VISIBLE : View.GONE;
+    void updateButtonVisibility() {
+        if (DesktopModeStatus.isProto2Enabled()) {
+            setButtonVisibility(mTaskInfo.getWindowingMode() == WINDOWING_MODE_FREEFORM);
+        } else if (DesktopModeStatus.isProto1Enabled()) {
+            mDesktopActive = DesktopModeStatus.isActive(mContext);
+            setButtonVisibility(mDesktopActive);
+        }
+    }
+
+    /**
+     * Show or hide buttons
+     */
+    void setButtonVisibility(boolean visible) {
+        int visibility = visible ? View.VISIBLE : View.GONE;
         View caption = mResult.mRootView.findViewById(R.id.caption);
         View back = caption.findViewById(R.id.back_button);
         View close = caption.findViewById(R.id.close_window);
-        back.setVisibility(v);
-        close.setVisibility(v);
+        back.setVisibility(visibility);
+        close.setVisibility(visibility);
         int buttonTintColorRes =
                 mDesktopActive ? R.color.decor_button_dark_color
                         : R.color.decor_button_light_color;
@@ -260,7 +278,7 @@
         View handle = caption.findViewById(R.id.caption_handle);
         VectorDrawable handleBackground = (VectorDrawable) handle.getBackground();
         handleBackground.setTintList(buttonTintColor);
-        caption.getBackground().setTint(v == View.VISIBLE ? Color.WHITE : Color.TRANSPARENT);
+        caption.getBackground().setTint(visible ? Color.WHITE : Color.TRANSPARENT);
     }
 
     boolean isHandleMenuActive() {
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopModeControllerTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopModeControllerTest.java
index 707c049..a3ba767 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopModeControllerTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopModeControllerTest.java
@@ -16,8 +16,6 @@
 
 package com.android.wm.shell.desktopmode;
 
-import static android.app.WindowConfiguration.ACTIVITY_TYPE_HOME;
-import static android.app.WindowConfiguration.ACTIVITY_TYPE_STANDARD;
 import static android.app.WindowConfiguration.WINDOWING_MODE_FREEFORM;
 import static android.app.WindowConfiguration.WINDOWING_MODE_FULLSCREEN;
 import static android.app.WindowConfiguration.WINDOWING_MODE_UNDEFINED;
@@ -29,6 +27,9 @@
 import static android.window.WindowContainerTransaction.HierarchyOp.HIERARCHY_OP_TYPE_REORDER;
 
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.mockitoSession;
+import static com.android.wm.shell.desktopmode.DesktopTestHelpers.createFreeformTask;
+import static com.android.wm.shell.desktopmode.DesktopTestHelpers.createFullscreenTask;
+import static com.android.wm.shell.desktopmode.DesktopTestHelpers.createHomeTask;
 
 import static com.google.common.truth.Truth.assertThat;
 
@@ -48,7 +49,6 @@
 import android.testing.AndroidTestingRunner;
 import android.window.DisplayAreaInfo;
 import android.window.TransitionRequestInfo;
-import android.window.WindowContainerToken;
 import android.window.WindowContainerTransaction;
 import android.window.WindowContainerTransaction.Change;
 import android.window.WindowContainerTransaction.HierarchyOp;
@@ -59,7 +59,6 @@
 import com.android.wm.shell.RootTaskDisplayAreaOrganizer;
 import com.android.wm.shell.ShellTaskOrganizer;
 import com.android.wm.shell.ShellTestCase;
-import com.android.wm.shell.TestRunningTaskInfoBuilder;
 import com.android.wm.shell.TestShellExecutor;
 import com.android.wm.shell.common.ShellExecutor;
 import com.android.wm.shell.sysui.ShellController;
@@ -355,7 +354,7 @@
     @Test
     public void testHandleTransitionRequest_taskOpen_returnsWct() {
         RunningTaskInfo trigger = new RunningTaskInfo();
-        trigger.token = new MockToken().mToken;
+        trigger.token = new MockToken().token();
         trigger.configuration.windowConfiguration.setWindowingMode(WINDOWING_MODE_FREEFORM);
         WindowContainerTransaction wct = mController.handleRequest(
                 mock(IBinder.class),
@@ -381,40 +380,13 @@
     }
 
     private DisplayAreaInfo createMockDisplayArea() {
-        DisplayAreaInfo displayAreaInfo = new DisplayAreaInfo(new MockToken().mToken,
+        DisplayAreaInfo displayAreaInfo = new DisplayAreaInfo(new MockToken().token(),
                 mContext.getDisplayId(), 0);
         when(mRootTaskDisplayAreaOrganizer.getDisplayAreaInfo(mContext.getDisplayId()))
                 .thenReturn(displayAreaInfo);
         return displayAreaInfo;
     }
 
-    private RunningTaskInfo createFreeformTask() {
-        return new TestRunningTaskInfoBuilder()
-                .setToken(new MockToken().token())
-                .setActivityType(ACTIVITY_TYPE_STANDARD)
-                .setWindowingMode(WINDOWING_MODE_FREEFORM)
-                .setLastActiveTime(100)
-                .build();
-    }
-
-    private RunningTaskInfo createFullscreenTask() {
-        return new TestRunningTaskInfoBuilder()
-                .setToken(new MockToken().token())
-                .setActivityType(ACTIVITY_TYPE_STANDARD)
-                .setWindowingMode(WINDOWING_MODE_FULLSCREEN)
-                .setLastActiveTime(100)
-                .build();
-    }
-
-    private RunningTaskInfo createHomeTask() {
-        return new TestRunningTaskInfoBuilder()
-                .setToken(new MockToken().token())
-                .setActivityType(ACTIVITY_TYPE_HOME)
-                .setWindowingMode(WINDOWING_MODE_FULLSCREEN)
-                .setLastActiveTime(100)
-                .build();
-    }
-
     private WindowContainerTransaction getDesktopModeSwitchTransaction() {
         ArgumentCaptor<WindowContainerTransaction> arg = ArgumentCaptor.forClass(
                 WindowContainerTransaction.class);
@@ -442,18 +414,4 @@
         assertThat(change.getConfiguration().windowConfiguration.getBounds().isEmpty()).isTrue();
     }
 
-    private static class MockToken {
-        private final WindowContainerToken mToken;
-        private final IBinder mBinder;
-
-        MockToken() {
-            mToken = mock(WindowContainerToken.class);
-            mBinder = mock(IBinder.class);
-            when(mToken.asBinder()).thenReturn(mBinder);
-        }
-
-        WindowContainerToken token() {
-            return mToken;
-        }
-    }
 }
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
new file mode 100644
index 0000000..de2473b
--- /dev/null
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksControllerTest.kt
@@ -0,0 +1,281 @@
+/*
+ * Copyright (C) 2022 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.WINDOWING_MODE_FREEFORM
+import android.app.WindowConfiguration.WINDOWING_MODE_FULLSCREEN
+import android.app.WindowConfiguration.WINDOWING_MODE_UNDEFINED
+import android.testing.AndroidTestingRunner
+import android.window.WindowContainerTransaction
+import android.window.WindowContainerTransaction.HierarchyOp.HIERARCHY_OP_TYPE_REORDER
+import androidx.test.filters.SmallTest
+import com.android.dx.mockito.inline.extended.ExtendedMockito.mockitoSession
+import com.android.dx.mockito.inline.extended.ExtendedMockito.never
+import com.android.dx.mockito.inline.extended.StaticMockitoSession
+import com.android.wm.shell.ShellTaskOrganizer
+import com.android.wm.shell.ShellTestCase
+import com.android.wm.shell.TestShellExecutor
+import com.android.wm.shell.common.ShellExecutor
+import com.android.wm.shell.desktopmode.DesktopTestHelpers.Companion.createFreeformTask
+import com.android.wm.shell.desktopmode.DesktopTestHelpers.Companion.createFullscreenTask
+import com.android.wm.shell.desktopmode.DesktopTestHelpers.Companion.createHomeTask
+import com.android.wm.shell.sysui.ShellController
+import com.android.wm.shell.sysui.ShellInit
+import com.android.wm.shell.transition.Transitions
+import com.google.common.truth.Truth.assertThat
+import com.google.common.truth.Truth.assertWithMessage
+import org.junit.After
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.ArgumentCaptor
+import org.mockito.ArgumentMatchers.isNull
+import org.mockito.Mock
+import org.mockito.Mockito
+import org.mockito.Mockito.any
+import org.mockito.Mockito.anyInt
+import org.mockito.Mockito.clearInvocations
+import org.mockito.Mockito.verify
+import org.mockito.Mockito.`when` as whenever
+
+@SmallTest
+@RunWith(AndroidTestingRunner::class)
+class DesktopTasksControllerTest : ShellTestCase() {
+
+    @Mock lateinit var testExecutor: ShellExecutor
+    @Mock lateinit var shellController: ShellController
+    @Mock lateinit var shellTaskOrganizer: ShellTaskOrganizer
+    @Mock lateinit var transitions: Transitions
+
+    lateinit var mockitoSession: StaticMockitoSession
+    lateinit var controller: DesktopTasksController
+    lateinit var shellInit: ShellInit
+    lateinit var desktopModeTaskRepository: DesktopModeTaskRepository
+
+    // Mock running tasks are registered here so we can get the list from mock shell task organizer
+    private val runningTasks = mutableListOf<RunningTaskInfo>()
+
+    @Before
+    fun setUp() {
+        mockitoSession = mockitoSession().mockStatic(DesktopModeStatus::class.java).startMocking()
+        whenever(DesktopModeStatus.isProto2Enabled()).thenReturn(true)
+
+        shellInit = Mockito.spy(ShellInit(testExecutor))
+        desktopModeTaskRepository = DesktopModeTaskRepository()
+
+        whenever(shellTaskOrganizer.getRunningTasks(anyInt())).thenAnswer { runningTasks }
+
+        controller = createController()
+
+        shellInit.init()
+    }
+
+    private fun createController(): DesktopTasksController {
+        return DesktopTasksController(
+            context,
+            shellInit,
+            shellController,
+            shellTaskOrganizer,
+            transitions,
+            desktopModeTaskRepository,
+            TestShellExecutor()
+        )
+    }
+
+    @After
+    fun tearDown() {
+        mockitoSession.finishMocking()
+
+        runningTasks.clear()
+    }
+
+    @Test
+    fun instantiate_addInitCallback() {
+        verify(shellInit).addInitCallback(any(), any<DesktopTasksController>())
+    }
+
+    @Test
+    fun instantiate_flagOff_doNotAddInitCallback() {
+        whenever(DesktopModeStatus.isProto2Enabled()).thenReturn(false)
+        clearInvocations(shellInit)
+
+        createController()
+
+        verify(shellInit, never()).addInitCallback(any(), any<DesktopTasksController>())
+    }
+
+    @Test
+    fun showDesktopApps_allAppsInvisible_bringsToFront() {
+        val homeTask = setUpHomeTask()
+        val task1 = setUpFreeformTask()
+        val task2 = setUpFreeformTask()
+        markTaskHidden(task1)
+        markTaskHidden(task2)
+
+        controller.showDesktopApps()
+
+        val wct = getLatestWct()
+        assertThat(wct.hierarchyOps).hasSize(3)
+        // Expect order to be from bottom: home, task1, task2
+        wct.assertReorderAt(index = 0, homeTask)
+        wct.assertReorderAt(index = 1, task1)
+        wct.assertReorderAt(index = 2, task2)
+    }
+
+    @Test
+    fun showDesktopApps_appsAlreadyVisible_doesNothing() {
+        setUpHomeTask()
+        val task1 = setUpFreeformTask()
+        val task2 = setUpFreeformTask()
+        markTaskVisible(task1)
+        markTaskVisible(task2)
+
+        controller.showDesktopApps()
+
+        verifyWCTNotExecuted()
+    }
+
+    @Test
+    fun showDesktopApps_someAppsInvisible_reordersAll() {
+        val homeTask = setUpHomeTask()
+        val task1 = setUpFreeformTask()
+        val task2 = setUpFreeformTask()
+        markTaskHidden(task1)
+        markTaskVisible(task2)
+
+        controller.showDesktopApps()
+
+        val wct = getLatestWct()
+        assertThat(wct.hierarchyOps).hasSize(3)
+        // Expect order to be from bottom: home, task1, task2
+        wct.assertReorderAt(index = 0, homeTask)
+        wct.assertReorderAt(index = 1, task1)
+        wct.assertReorderAt(index = 2, task2)
+    }
+
+    @Test
+    fun showDesktopApps_noActiveTasks_reorderHomeToTop() {
+        val homeTask = setUpHomeTask()
+
+        controller.showDesktopApps()
+
+        val wct = getLatestWct()
+        assertThat(wct.hierarchyOps).hasSize(1)
+        wct.assertReorderAt(index = 0, homeTask)
+    }
+
+    @Test
+    fun moveToDesktop() {
+        val task = setUpFullscreenTask()
+        controller.moveToDesktop(task)
+        val wct = getLatestWct()
+        assertThat(wct.changes[task.token.asBinder()]?.windowingMode)
+            .isEqualTo(WINDOWING_MODE_FREEFORM)
+    }
+
+    @Test
+    fun moveToDesktop_nonExistentTask_doesNothing() {
+        controller.moveToDesktop(999)
+        verifyWCTNotExecuted()
+    }
+
+    @Test
+    fun moveToFullscreen() {
+        val task = setUpFreeformTask()
+        controller.moveToFullscreen(task)
+        val wct = getLatestWct()
+        assertThat(wct.changes[task.token.asBinder()]?.windowingMode)
+            .isEqualTo(WINDOWING_MODE_FULLSCREEN)
+    }
+
+    @Test
+    fun moveToFullscreen_nonExistentTask_doesNothing() {
+        controller.moveToFullscreen(999)
+        verifyWCTNotExecuted()
+    }
+
+    @Test
+    fun getTaskWindowingMode() {
+        val fullscreenTask = setUpFullscreenTask()
+        val freeformTask = setUpFreeformTask()
+
+        assertThat(controller.getTaskWindowingMode(fullscreenTask.taskId))
+            .isEqualTo(WINDOWING_MODE_FULLSCREEN)
+        assertThat(controller.getTaskWindowingMode(freeformTask.taskId))
+            .isEqualTo(WINDOWING_MODE_FREEFORM)
+        assertThat(controller.getTaskWindowingMode(999)).isEqualTo(WINDOWING_MODE_UNDEFINED)
+    }
+
+    private fun setUpFreeformTask(): RunningTaskInfo {
+        val task = createFreeformTask()
+        whenever(shellTaskOrganizer.getRunningTaskInfo(task.taskId)).thenReturn(task)
+        desktopModeTaskRepository.addActiveTask(task.taskId)
+        desktopModeTaskRepository.addOrMoveFreeformTaskToTop(task.taskId)
+        runningTasks.add(task)
+        return task
+    }
+
+    private fun setUpHomeTask(): RunningTaskInfo {
+        val task = createHomeTask()
+        whenever(shellTaskOrganizer.getRunningTaskInfo(task.taskId)).thenReturn(task)
+        runningTasks.add(task)
+        return task
+    }
+
+    private fun setUpFullscreenTask(): RunningTaskInfo {
+        val task = createFullscreenTask()
+        whenever(shellTaskOrganizer.getRunningTaskInfo(task.taskId)).thenReturn(task)
+        runningTasks.add(task)
+        return task
+    }
+
+    private fun markTaskVisible(task: RunningTaskInfo) {
+        desktopModeTaskRepository.updateVisibleFreeformTasks(task.taskId, visible = true)
+    }
+
+    private fun markTaskHidden(task: RunningTaskInfo) {
+        desktopModeTaskRepository.updateVisibleFreeformTasks(task.taskId, visible = false)
+    }
+
+    private fun getLatestWct(): WindowContainerTransaction {
+        val arg = ArgumentCaptor.forClass(WindowContainerTransaction::class.java)
+        if (Transitions.ENABLE_SHELL_TRANSITIONS) {
+            verify(transitions).startTransition(anyInt(), arg.capture(), isNull())
+        } else {
+            verify(shellTaskOrganizer).applyTransaction(arg.capture())
+        }
+        return arg.value
+    }
+
+    private fun verifyWCTNotExecuted() {
+        if (Transitions.ENABLE_SHELL_TRANSITIONS) {
+            verify(transitions, never()).startTransition(anyInt(), any(), isNull())
+        } else {
+            verify(shellTaskOrganizer, never()).applyTransaction(any())
+        }
+    }
+}
+
+private fun WindowContainerTransaction.assertReorderAt(index: Int, task: RunningTaskInfo) {
+    assertWithMessage("WCT does not have a hierarchy operation at index $index")
+        .that(hierarchyOps.size)
+        .isGreaterThan(index)
+    val op = hierarchyOps[index]
+    assertThat(op.type).isEqualTo(HIERARCHY_OP_TYPE_REORDER)
+    assertThat(op.container).isEqualTo(task.token.asBinder())
+}
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTestHelpers.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTestHelpers.kt
new file mode 100644
index 0000000..dc91d75
--- /dev/null
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTestHelpers.kt
@@ -0,0 +1,61 @@
+/*
+ * Copyright (C) 2022 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_HOME
+import android.app.WindowConfiguration.ACTIVITY_TYPE_STANDARD
+import android.app.WindowConfiguration.WINDOWING_MODE_FREEFORM
+import android.app.WindowConfiguration.WINDOWING_MODE_FULLSCREEN
+import com.android.wm.shell.TestRunningTaskInfoBuilder
+
+class DesktopTestHelpers {
+    companion object {
+        /** Create a task that has windowing mode set to [WINDOWING_MODE_FREEFORM] */
+        @JvmStatic
+        fun createFreeformTask(): RunningTaskInfo {
+            return TestRunningTaskInfoBuilder()
+                    .setToken(MockToken().token())
+                    .setActivityType(ACTIVITY_TYPE_STANDARD)
+                    .setWindowingMode(WINDOWING_MODE_FREEFORM)
+                    .setLastActiveTime(100)
+                    .build()
+        }
+
+        /** Create a task that has windowing mode set to [WINDOWING_MODE_FULLSCREEN] */
+        @JvmStatic
+        fun createFullscreenTask(): RunningTaskInfo {
+            return TestRunningTaskInfoBuilder()
+                    .setToken(MockToken().token())
+                    .setActivityType(ACTIVITY_TYPE_STANDARD)
+                    .setWindowingMode(WINDOWING_MODE_FULLSCREEN)
+                    .setLastActiveTime(100)
+                    .build()
+        }
+
+        /** Create a new home task */
+        @JvmStatic
+        fun createHomeTask(): RunningTaskInfo {
+            return TestRunningTaskInfoBuilder()
+                    .setToken(MockToken().token())
+                    .setActivityType(ACTIVITY_TYPE_HOME)
+                    .setWindowingMode(WINDOWING_MODE_FULLSCREEN)
+                    .setLastActiveTime(100)
+                    .build()
+        }
+    }
+}
\ No newline at end of file
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/MockToken.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/MockToken.java
new file mode 100644
index 0000000..09d474d
--- /dev/null
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/MockToken.java
@@ -0,0 +1,40 @@
+/*
+ * Copyright (C) 2022 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 org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+import android.os.IBinder;
+import android.window.WindowContainerToken;
+
+/**
+ * {@link WindowContainerToken} wrapper that supports a mock binder
+ */
+class MockToken {
+    private final WindowContainerToken mToken;
+
+    MockToken() {
+        mToken = mock(WindowContainerToken.class);
+        IBinder binder = mock(IBinder.class);
+        when(mToken.asBinder()).thenReturn(binder);
+    }
+
+    WindowContainerToken token() {
+        return mToken;
+    }
+}
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/CaptionWindowDecorViewModelTests.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/CaptionWindowDecorViewModelTests.java
index ad6fced..87f9d21 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/CaptionWindowDecorViewModelTests.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/CaptionWindowDecorViewModelTests.java
@@ -47,6 +47,7 @@
 import com.android.wm.shell.common.DisplayController;
 import com.android.wm.shell.common.SyncTransactionQueue;
 import com.android.wm.shell.desktopmode.DesktopModeController;
+import com.android.wm.shell.desktopmode.DesktopTasksController;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -76,6 +77,8 @@
 
     @Mock private DesktopModeController mDesktopModeController;
 
+    @Mock private DesktopTasksController mDesktopTasksController;
+
     @Mock private InputMonitor mInputMonitor;
 
     @Mock private InputChannel mInputChannel;
@@ -103,6 +106,7 @@
                 mDisplayController,
                 mSyncQueue,
                 Optional.of(mDesktopModeController),
+                Optional.of(mDesktopTasksController),
                 mCaptionWindowDecorFactory,
                 new MockObjectSupplier<>(mMockInputManagers, () -> mock(InputManager.class)));
         mCaptionWindowDecorViewModel.setEventReceiverFactory(mEventReceiverFactory);