diff --git a/quickstep/src/com/android/launcher3/statehandlers/DesktopVisibilityController.kt b/quickstep/src/com/android/launcher3/statehandlers/DesktopVisibilityController.kt
index 03f5d96..62543e4 100644
--- a/quickstep/src/com/android/launcher3/statehandlers/DesktopVisibilityController.kt
+++ b/quickstep/src/com/android/launcher3/statehandlers/DesktopVisibilityController.kt
@@ -62,14 +62,14 @@
      * @property canCreateDesks true if it's possible to create new desks on the display represented
      *   by this object.
      * @property activeDeskId The ID of the active desk on the associated display (if any). It has a
-     *   value of `-1` if there are no active desks. Note that there can only be at most one active
-     *   desk on each display.
+     *   value of `INACTIVE_DESK_ID` (-1) if there are no active desks. Note that there can only be
+     *   at most one active desk on each display.
      * @property deskIds a set containing the IDs of the desks on the associated display.
      */
     private data class DisplayDeskConfig(
         val displayId: Int,
         var canCreateDesks: Boolean,
-        var activeDeskId: Int = -1,
+        var activeDeskId: Int = INACTIVE_DESK_ID,
         val deskIds: MutableSet<Int>,
     )
 
@@ -79,6 +79,8 @@
     private val desktopVisibilityListeners: MutableSet<DesktopVisibilityListener> = HashSet()
     private val taskbarDesktopModeListeners: MutableSet<TaskbarDesktopModeListener> = HashSet()
 
+    // TODO: b/394387739 - Deprecate this and replace it with something that tracks the count per
+    //  desk.
     /** Number of visible desktop windows in desktop mode. */
     var visibleDesktopTasksCount: Int = 0
         /**
@@ -142,8 +144,36 @@
         }
     }
 
+    /** Returns whether a desk is currently active on the display with the given [displayId]. */
+    fun isInDesktopMode(displayId: Int): Boolean {
+        if (!DesktopModeStatus.enableMultipleDesktops(context)) {
+            return areDesktopTasksVisible()
+        }
+
+        val isInDesktopMode = displaysDesksConfigsMap[displayId].activeDeskId != INACTIVE_DESK_ID
+        if (DEBUG) {
+            Log.d(TAG, "isInDesktopMode: $isInDesktopMode")
+        }
+        return isInDesktopMode
+    }
+
+    /**
+     * Returns whether a desk is currently active on the display with the given [displayId] and
+     * Overview is not active.
+     */
+    fun isInDesktopModeAndNotInOverview(displayId: Int): Boolean {
+        if (!DesktopModeStatus.enableMultipleDesktops(context)) {
+            return areDesktopTasksVisibleAndNotInOverview()
+        }
+
+        if (DEBUG) {
+            Log.d(TAG, "isInDesktopModeAndNotInOverview: overview=$inOverviewState")
+        }
+        return isInDesktopMode(displayId) && !inOverviewState
+    }
+
     /** Whether desktop tasks are visible in desktop mode. */
-    fun areDesktopTasksVisible(): Boolean {
+    private fun areDesktopTasksVisible(): Boolean {
         val desktopTasksVisible: Boolean = visibleDesktopTasksCount > 0
         if (DEBUG) {
             Log.d(TAG, "areDesktopTasksVisible: desktopVisible=$desktopTasksVisible")
@@ -152,7 +182,7 @@
     }
 
     /** Whether desktop tasks are visible in desktop mode. */
-    fun areDesktopTasksVisibleAndNotInOverview(): Boolean {
+    private fun areDesktopTasksVisibleAndNotInOverview(): Boolean {
         val desktopTasksVisible: Boolean = visibleDesktopTasksCount > 0
         if (DEBUG) {
             Log.d(
@@ -250,6 +280,13 @@
     }
 
     private fun notifyDesktopVisibilityListeners(areDesktopTasksVisible: Boolean) {
+        if (DesktopModeStatus.enableMultipleDesktops(context)) {
+            // This is triggered using the legacy way of determining the visibility of the desktop
+            // mode and should not be used when the multi-desks feature is enabled.
+            // (See b/394685645 for details).
+            return
+        }
+
         if (DEBUG) {
             Log.d(TAG, "notifyDesktopVisibilityListeners: visible=$areDesktopTasksVisible")
         }
@@ -391,7 +428,7 @@
                 "Removing non-existing desk Id: $deskId on display: $displayId"
             }
             if (it.activeDeskId == deskId) {
-                it.activeDeskId = -1
+                it.activeDeskId = INACTIVE_DESK_ID
             }
         }
     }
@@ -542,5 +579,7 @@
 
         private const val TAG = "DesktopVisController"
         private const val DEBUG = false
+
+        private const val INACTIVE_DESK_ID = -1
     }
 }
diff --git a/quickstep/src/com/android/launcher3/taskbar/LauncherTaskbarUIController.java b/quickstep/src/com/android/launcher3/taskbar/LauncherTaskbarUIController.java
index 4143157..3f3700b 100644
--- a/quickstep/src/com/android/launcher3/taskbar/LauncherTaskbarUIController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/LauncherTaskbarUIController.java
@@ -246,7 +246,7 @@
 
         if (!ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY.isTrue()
                 && mControllers.taskbarDesktopModeController
-                    .getAreDesktopTasksVisibleAndNotInOverview()) {
+                    .isInDesktopModeAndNotInOverview(mLauncher.getDisplayId())) {
             // TODO: b/333533253 - Remove after flag rollout
             isVisible = false;
         }
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
index 8880abd..b9d1275 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
@@ -1291,7 +1291,7 @@
 
     boolean areDesktopTasksVisible() {
         return mControllers != null
-                && mControllers.taskbarDesktopModeController.getAreDesktopTasksVisible();
+                && mControllers.taskbarDesktopModeController.isInDesktopMode(getDisplayId());
     }
 
     protected void onTaskbarIconClicked(View view) {
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarControllers.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarControllers.java
index af60f10..b244be9 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarControllers.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarControllers.java
@@ -194,7 +194,8 @@
                 voiceInteractionWindowController
         };
 
-        if (taskbarDesktopModeController.getAreDesktopTasksVisibleAndNotInOverview()) {
+        if (taskbarDesktopModeController.isInDesktopModeAndNotInOverview(
+                taskbarActivityContext.getDisplayId())) {
             mCornerRoundness.value = taskbarDesktopModeController.getTaskbarCornerRoundness(
                     mSharedState.showCornerRadiusInDesktopMode);
         } else {
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarDesktopModeController.kt b/quickstep/src/com/android/launcher3/taskbar/TaskbarDesktopModeController.kt
index cb399e8..f71dea9 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarDesktopModeController.kt
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarDesktopModeController.kt
@@ -30,18 +30,17 @@
     private lateinit var taskbarControllers: TaskbarControllers
     private lateinit var taskbarSharedState: TaskbarSharedState
 
-    val areDesktopTasksVisibleAndNotInOverview: Boolean
-        get() = desktopVisibilityController.areDesktopTasksVisibleAndNotInOverview()
-
-    val areDesktopTasksVisible: Boolean
-        get() = desktopVisibilityController.areDesktopTasksVisible()
-
     fun init(controllers: TaskbarControllers, sharedState: TaskbarSharedState) {
         taskbarControllers = controllers
         taskbarSharedState = sharedState
         desktopVisibilityController.registerTaskbarDesktopModeListener(this)
     }
 
+    fun isInDesktopMode(displayId: Int) = desktopVisibilityController.isInDesktopMode(displayId)
+
+    fun isInDesktopModeAndNotInOverview(displayId: Int) =
+        desktopVisibilityController.isInDesktopModeAndNotInOverview(displayId)
+
     override fun onTaskbarCornerRoundingUpdate(doesAnyTaskRequireTaskbarRounding: Boolean) {
         taskbarSharedState.showCornerRadiusInDesktopMode = doesAnyTaskRequireTaskbarRounding
         val cornerRadius = getTaskbarCornerRoundness(doesAnyTaskRequireTaskbarRounding)
@@ -49,7 +48,7 @@
     }
 
     fun shouldShowDesktopTasksInTaskbar(): Boolean {
-        return desktopVisibilityController.areDesktopTasksVisible() ||
+        return isInDesktopMode(context.displayId) ||
             DisplayController.showDesktopTaskbarForFreeformDisplay(context) ||
             (DisplayController.showLockedTaskbarOnHome(context) &&
                 taskbarControllers.taskbarStashController.isOnHome)
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarDragController.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarDragController.java
index f36c481..b91f512 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarDragController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarDragController.java
@@ -347,7 +347,7 @@
         // Pre-drag has ended, start the global system drag.
         if (mDisallowGlobalDrag
                 || mControllers.taskbarDesktopModeController
-                    .getAreDesktopTasksVisibleAndNotInOverview()) {
+                    .isInDesktopModeAndNotInOverview(mActivity.getDisplayId())) {
             AbstractFloatingView.closeAllOpenViewsExcept(mActivity, TYPE_TASKBAR_ALL_APPS);
             return;
         }
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarLauncherStateController.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarLauncherStateController.java
index cada5a3..dd9f61e 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarLauncherStateController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarLauncherStateController.java
@@ -631,7 +631,8 @@
 
         float cornerRoundness = isInLauncher ? 0 : 1;
 
-        if (mControllers.taskbarDesktopModeController.getAreDesktopTasksVisibleAndNotInOverview()
+        if (mControllers.taskbarDesktopModeController.isInDesktopModeAndNotInOverview(
+                mControllers.taskbarActivityContext.getDisplayId())
                 && mControllers.getSharedState() != null) {
             cornerRoundness =
                     mControllers.taskbarDesktopModeController.getTaskbarCornerRoundness(
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarPinningController.kt b/quickstep/src/com/android/launcher3/taskbar/TaskbarPinningController.kt
index 23c5070..7141bb8 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarPinningController.kt
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarPinningController.kt
@@ -58,8 +58,9 @@
                 }
                 val shouldPinTaskbar =
                     if (
-                        controllers.taskbarDesktopModeController
-                            .areDesktopTasksVisibleAndNotInOverview
+                        controllers.taskbarDesktopModeController.isInDesktopModeAndNotInOverview(
+                            context.displayId
+                        )
                     ) {
                         !launcherPrefs.get(TASKBAR_PINNING_IN_DESKTOP_MODE)
                     } else {
@@ -140,7 +141,11 @@
     @VisibleForTesting
     fun recreateTaskbarAndUpdatePinningValue() {
         updateIsAnimatingTaskbarPinningAndNotifyTaskbarDragLayer(false)
-        if (controllers.taskbarDesktopModeController.areDesktopTasksVisibleAndNotInOverview) {
+        if (
+            controllers.taskbarDesktopModeController.isInDesktopModeAndNotInOverview(
+                context.displayId
+            )
+        ) {
             launcherPrefs.put(
                 TASKBAR_PINNING_IN_DESKTOP_MODE,
                 !launcherPrefs.get(TASKBAR_PINNING_IN_DESKTOP_MODE),
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarPopupController.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarPopupController.java
index feb9b33..5d8b821 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarPopupController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarPopupController.java
@@ -217,7 +217,7 @@
         }
         shortcuts.add(APP_INFO);
         if (!mControllers.taskbarDesktopModeController
-                .getAreDesktopTasksVisibleAndNotInOverview()) {
+                .isInDesktopModeAndNotInOverview(mContext.getDisplayId())) {
             shortcuts.addAll(mControllers.uiController.getSplitMenuOptions().toList());
         }
         if (BubbleAnythingFlagHelper.enableCreateAnyBubble()) {
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarStashController.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarStashController.java
index 2ded1bf..6016394 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarStashController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarStashController.java
@@ -1189,7 +1189,7 @@
         if (mActivity.isHardwareKeyboard()
                 && mActivity.isThreeButtonNav()
                 && mControllers.taskbarDesktopModeController
-                    .getAreDesktopTasksVisibleAndNotInOverview()) {
+                    .isInDesktopModeAndNotInOverview(mActivity.getDisplayId())) {
             return false;
         }
 
diff --git a/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java b/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java
index 23f4f67..bb57d6e 100644
--- a/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java
+++ b/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java
@@ -1047,7 +1047,7 @@
         DesktopVisibilityController desktopVisibilityController =
                 DesktopVisibilityController.INSTANCE.get(this);
         if (!ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY.isTrue()
-                && desktopVisibilityController.areDesktopTasksVisibleAndNotInOverview()
+                && desktopVisibilityController.isInDesktopModeAndNotInOverview(getDisplayId())
                 && !desktopVisibilityController.isRecentsGestureInProgress()) {
             // Return early to skip setting activity to appear as resumed
             // TODO: b/333533253 - Remove after flag rollout
@@ -1369,7 +1369,7 @@
     @Override
     public boolean areDesktopTasksVisible() {
         return DesktopVisibilityController.INSTANCE.get(this)
-                .areDesktopTasksVisibleAndNotInOverview();
+                .isInDesktopModeAndNotInOverview(getDisplayId());
     }
 
     @Override
diff --git a/quickstep/src/com/android/quickstep/BaseActivityInterface.java b/quickstep/src/com/android/quickstep/BaseActivityInterface.java
index 7cab751..549c2f8 100644
--- a/quickstep/src/com/android/quickstep/BaseActivityInterface.java
+++ b/quickstep/src/com/android/quickstep/BaseActivityInterface.java
@@ -87,7 +87,7 @@
             // We were on our way to this state when we got canceled, end there instead.
             startState = stateFromGestureEndTarget(endTarget);
             if (DesktopVisibilityController.INSTANCE.get(activity)
-                    .areDesktopTasksVisibleAndNotInOverview()
+                    .isInDesktopModeAndNotInOverview(activity.getDisplayId())
                     && endTarget == LAST_TASK) {
                 // When we are cancelling the transition and going back to last task, move to
                 // rest state instead when desktop tasks are visible.
diff --git a/quickstep/src/com/android/quickstep/BaseContainerInterface.java b/quickstep/src/com/android/quickstep/BaseContainerInterface.java
index 6d588d9..c64067a 100644
--- a/quickstep/src/com/android/quickstep/BaseContainerInterface.java
+++ b/quickstep/src/com/android/quickstep/BaseContainerInterface.java
@@ -233,8 +233,10 @@
         if (endTarget != null) {
             // We were on our way to this state when we got canceled, end there instead.
             startState = stateFromGestureEndTarget(endTarget);
-            if (DesktopVisibilityController.INSTANCE.get(recentsView.getContext())
-                    .areDesktopTasksVisibleAndNotInOverview() && endTarget == LAST_TASK) {
+            final var context = recentsView.getContext();
+            if (DesktopVisibilityController.INSTANCE.get(context)
+                    .isInDesktopModeAndNotInOverview(context.getDisplayId())
+                    && endTarget == LAST_TASK) {
                 // When we are cancelling the transition and going back to last task, move to
                 // rest state instead when desktop tasks are visible.
                 // If a fullscreen task is visible, launcher goes to normal state when the
diff --git a/quickstep/src/com/android/quickstep/LauncherSwipeHandlerV2.java b/quickstep/src/com/android/quickstep/LauncherSwipeHandlerV2.java
index e1e9c99..0a77688 100644
--- a/quickstep/src/com/android/quickstep/LauncherSwipeHandlerV2.java
+++ b/quickstep/src/com/android/quickstep/LauncherSwipeHandlerV2.java
@@ -107,7 +107,7 @@
                 && workspaceView.isAttachedToWindow()
                 && workspaceView.getHeight() > 0
                 && !DesktopVisibilityController.INSTANCE.get(mContainer)
-                        .areDesktopTasksVisibleAndNotInOverview();
+                        .isInDesktopModeAndNotInOverview(mContainer.getDisplayId());
 
         mContainer.getRootView().setForceHideBackArrow(true);
 
diff --git a/quickstep/src/com/android/quickstep/util/SystemWindowManagerProxy.java b/quickstep/src/com/android/quickstep/util/SystemWindowManagerProxy.java
index 10ae7a3..d92cc86 100644
--- a/quickstep/src/com/android/quickstep/util/SystemWindowManagerProxy.java
+++ b/quickstep/src/com/android/quickstep/util/SystemWindowManagerProxy.java
@@ -73,8 +73,8 @@
     }
 
     @Override
-    public boolean isInDesktopMode() {
-        return mDesktopVisibilityController.areDesktopTasksVisible();
+    public boolean isInDesktopMode(int displayId) {
+        return mDesktopVisibilityController.isInDesktopMode(displayId);
     }
 
     @Override
diff --git a/quickstep/src/com/android/quickstep/views/LauncherRecentsView.java b/quickstep/src/com/android/quickstep/views/LauncherRecentsView.java
index c6bd677..0f1c294 100644
--- a/quickstep/src/com/android/quickstep/views/LauncherRecentsView.java
+++ b/quickstep/src/com/android/quickstep/views/LauncherRecentsView.java
@@ -275,7 +275,8 @@
         boolean showDesktopApps = false;
         GestureState.GestureEndTarget endTarget = mCurrentGestureEndTarget;
         if (endTarget == GestureState.GestureEndTarget.LAST_TASK
-                && desktopVisibilityController.areDesktopTasksVisibleAndNotInOverview()) {
+                && desktopVisibilityController.isInDesktopModeAndNotInOverview(
+                        mContainer.getDisplayId())) {
             // Recents gesture was cancelled and we are returning to the previous task.
             // After super class has handled clean up, show desktop apps on top again
             showDesktopApps = true;
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/taskbar/controllers/TaskbarPinningControllerTest.kt b/quickstep/tests/multivalentTests/src/com/android/quickstep/taskbar/controllers/TaskbarPinningControllerTest.kt
index 5b42d6c..26418d8 100644
--- a/quickstep/tests/multivalentTests/src/com/android/quickstep/taskbar/controllers/TaskbarPinningControllerTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/taskbar/controllers/TaskbarPinningControllerTest.kt
@@ -71,8 +71,9 @@
         whenever(taskbarActivityContext.dragLayer).thenReturn(taskbarDragLayer)
         whenever(taskbarActivityContext.statsLogManager).thenReturn(statsLogManager)
         whenever(
-                taskbarControllers.taskbarDesktopModeController
-                    .areDesktopTasksVisibleAndNotInOverview
+                taskbarControllers.taskbarDesktopModeController.isInDesktopModeAndNotInOverview(
+                    taskbarActivityContext.displayId
+                )
             )
             .thenAnswer { _ -> isInDesktopMode }
         pinningController = spy(TaskbarPinningController(taskbarActivityContext))
diff --git a/src/com/android/launcher3/util/DisplayController.java b/src/com/android/launcher3/util/DisplayController.java
index ee1af81..8008687 100644
--- a/src/com/android/launcher3/util/DisplayController.java
+++ b/src/com/android/launcher3/util/DisplayController.java
@@ -448,7 +448,7 @@
             mIsTaskbarPinned = LauncherPrefs.get(displayInfoContext).get(TASKBAR_PINNING);
             mIsTaskbarPinnedInDesktopMode = LauncherPrefs.get(displayInfoContext).get(
                     TASKBAR_PINNING_IN_DESKTOP_MODE);
-            mIsInDesktopMode = wmProxy.isInDesktopMode();
+            mIsInDesktopMode = wmProxy.isInDesktopMode(DEFAULT_DISPLAY);
             mShowLockedTaskbarOnHome = wmProxy.showLockedTaskbarOnHome(displayInfoContext);
             mShowDesktopTaskbarForFreeformDisplay = wmProxy.showDesktopTaskbarForFreeformDisplay(
                     displayInfoContext);
diff --git a/src/com/android/launcher3/util/window/WindowManagerProxy.java b/src/com/android/launcher3/util/window/WindowManagerProxy.java
index f511ef2..bbb9e28 100644
--- a/src/com/android/launcher3/util/window/WindowManagerProxy.java
+++ b/src/com/android/launcher3/util/window/WindowManagerProxy.java
@@ -109,7 +109,7 @@
     /**
      * Returns if we are in desktop mode or not.
      */
-    public boolean isInDesktopMode() {
+    public boolean isInDesktopMode(int displayId) {
         return false;
     }
 
diff --git a/tests/multivalentTests/src/com/android/launcher3/util/DisplayControllerTest.kt b/tests/multivalentTests/src/com/android/launcher3/util/DisplayControllerTest.kt
index 588a668..aa1451b 100644
--- a/tests/multivalentTests/src/com/android/launcher3/util/DisplayControllerTest.kt
+++ b/tests/multivalentTests/src/com/android/launcher3/util/DisplayControllerTest.kt
@@ -203,7 +203,7 @@
     fun testTaskbarPinningChangeInLockedTaskbarChange() {
         whenever(windowManagerProxy.showLockedTaskbarOnHome(any())).thenReturn(true)
         whenever(windowManagerProxy.isHomeVisible(any())).thenReturn(true)
-        whenever(windowManagerProxy.isInDesktopMode()).thenReturn(false)
+        whenever(windowManagerProxy.isInDesktopMode(any())).thenReturn(false)
         whenever(launcherPrefs.get(TASKBAR_PINNING)).thenReturn(false)
         DisplayController.enableTaskbarModePreferenceForTests(true)
 
@@ -219,7 +219,7 @@
     fun testLockedTaskbarChangeOnConfigurationChanged() {
         whenever(windowManagerProxy.showLockedTaskbarOnHome(any())).thenReturn(true)
         whenever(windowManagerProxy.isHomeVisible(any())).thenReturn(true)
-        whenever(windowManagerProxy.isInDesktopMode()).thenReturn(false)
+        whenever(windowManagerProxy.isInDesktopMode(any())).thenReturn(false)
         whenever(launcherPrefs.get(TASKBAR_PINNING)).thenReturn(false)
         DisplayController.enableTaskbarModePreferenceForTests(true)
         assertTrue(displayController.getInfo().isTransientTaskbar())
@@ -237,7 +237,7 @@
         whenever(windowManagerProxy.showDesktopTaskbarForFreeformDisplay(any())).thenReturn(true)
         whenever(launcherPrefs.get(TASKBAR_PINNING)).thenReturn(false)
         whenever(launcherPrefs.get(TASKBAR_PINNING_IN_DESKTOP_MODE)).thenReturn(false)
-        whenever(windowManagerProxy.isInDesktopMode()).thenReturn(true)
+        whenever(windowManagerProxy.isInDesktopMode(any())).thenReturn(true)
         whenever(windowManagerProxy.isHomeVisible(any())).thenReturn(false)
         DisplayController.enableTaskbarModePreferenceForTests(true)
 
@@ -256,7 +256,7 @@
         whenever(windowManagerProxy.showDesktopTaskbarForFreeformDisplay(any())).thenReturn(true)
         whenever(launcherPrefs.get(TASKBAR_PINNING)).thenReturn(false)
         whenever(launcherPrefs.get(TASKBAR_PINNING_IN_DESKTOP_MODE)).thenReturn(false)
-        whenever(windowManagerProxy.isInDesktopMode()).thenReturn(false)
+        whenever(windowManagerProxy.isInDesktopMode(any())).thenReturn(false)
         whenever(windowManagerProxy.isHomeVisible(any())).thenReturn(false)
         DisplayController.enableTaskbarModePreferenceForTests(true)
 
@@ -275,7 +275,7 @@
         whenever(windowManagerProxy.showDesktopTaskbarForFreeformDisplay(any())).thenReturn(true)
         whenever(launcherPrefs.get(TASKBAR_PINNING)).thenReturn(false)
         whenever(launcherPrefs.get(TASKBAR_PINNING_IN_DESKTOP_MODE)).thenReturn(false)
-        whenever(windowManagerProxy.isInDesktopMode()).thenReturn(false)
+        whenever(windowManagerProxy.isInDesktopMode(any())).thenReturn(false)
         whenever(windowManagerProxy.isHomeVisible(any())).thenReturn(true)
         DisplayController.enableTaskbarModePreferenceForTests(true)