Add immersive state to desktop repository

Changes the app header maximize icon based on whether the task is in
immersive state or not.

Flag: com.android.window.flags.enable_fully_immersive_in_desktop
Bug: 369402633
Test: atest WMShellUnitTests
Change-Id: Ic4f73212d34d07f14057c0475ecab99b3b22820a

11 files changed, 203 insertions, 41 deletions

diff --git a/libs/WindowManager/Shell/res/drawable/decor_desktop_mode_immersive_exit_button_dark.xml b/libs/WindowManager/Shell/res/drawable/decor_desktop_mode_immersive_exit_button_dark.xml
new file mode 100644
index 000000000000..5260450e8a13
--- /dev/null
+++ b/libs/WindowManager/Shell/res/drawable/decor_desktop_mode_immersive_exit_button_dark.xml
@@ -0,0 +1,26 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!--
+  ~ 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.
+  -->
+<vector xmlns:android="http://schemas.android.com/apk/res/android"
+    android:width="24dp"
+    android:height="24dp"
+    android:viewportWidth="960"
+    android:viewportHeight="960"
+    android:tint="?attr/colorControlNormal">
+    <path
+        android:fillColor="@android:color/white"
+        android:pathData="M240,840L240,720L120,720L120,640L320,640L320,840L240,840ZM640,840L640,640L840,640L840,720L720,720L720,840L640,840ZM120,320L120,240L240,240L240,120L320,120L320,320L120,320ZM640,320L640,120L720,120L720,240L840,240L840,320L640,320Z"/>
+</vector>
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 96a07755fea9..60f0eb55ffc8 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
@@ -234,6 +234,7 @@ public abstract class WMShellModule {
             IWindowManager windowManager,
             ShellCommandHandler shellCommandHandler,
             ShellTaskOrganizer taskOrganizer,
+            @DynamicOverride DesktopModeTaskRepository desktopRepository,
             DisplayController displayController,
             ShellController shellController,
             DisplayInsetsController displayInsetsController,
@@ -259,6 +260,7 @@ public abstract class WMShellModule {
                     shellCommandHandler,
                     windowManager,
                     taskOrganizer,
+                    desktopRepository,
                     displayController,
                     shellController,
                     displayInsetsController,
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopModeTaskRepository.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopModeTaskRepository.kt
index 759ed035895e..03f23a2136c2 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopModeTaskRepository.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopModeTaskRepository.kt
@@ -60,6 +60,7 @@ class DesktopModeTaskRepository (
      * @property minimizedTasks task ids for active freeform tasks that are currently minimized.
      * @property closingTasks task ids for tasks that are going to close, but are currently visible.
      * @property freeformTasksInZOrder list of current freeform task ids ordered from top to bottom
+     * @property fullImmersiveTaskId the task id of the desktop task that is in full-immersive mode.
      * (top is at index 0).
      */
     private data class DesktopTaskData(
@@ -69,13 +70,15 @@ class DesktopModeTaskRepository (
         // TODO(b/332682201): Remove when the repository state is updated via TransitionObserver
         val closingTasks: ArraySet<Int> = ArraySet(),
         val freeformTasksInZOrder: ArrayList<Int> = ArrayList(),
+        var fullImmersiveTaskId: Int? = null,
     ) {
         fun deepCopy(): DesktopTaskData = DesktopTaskData(
             activeTasks = ArraySet(activeTasks),
             visibleTasks = ArraySet(visibleTasks),
             minimizedTasks = ArraySet(minimizedTasks),
             closingTasks = ArraySet(closingTasks),
-            freeformTasksInZOrder = ArrayList(freeformTasksInZOrder)
+            freeformTasksInZOrder = ArrayList(freeformTasksInZOrder),
+            fullImmersiveTaskId = fullImmersiveTaskId
         )
     }
 
@@ -300,6 +303,23 @@ class DesktopModeTaskRepository (
         }
     }
 
+    /** Set whether the given task is the full-immersive task in this display. */
+    fun setTaskInFullImmersiveState(displayId: Int, taskId: Int, immersive: Boolean) {
+        val desktopData = desktopTaskDataByDisplayId.getOrCreate(displayId)
+        if (immersive) {
+            desktopData.fullImmersiveTaskId = taskId
+        } else {
+            if (desktopData.fullImmersiveTaskId == taskId) {
+                desktopData.fullImmersiveTaskId = null
+            }
+        }
+    }
+
+    /* Whether the task is in full-immersive state. */
+    fun isTaskInFullImmersiveState(taskId: Int): Boolean {
+        return desktopTaskDataSequence().any { taskId == it.fullImmersiveTaskId }
+    }
+
     private fun notifyVisibleTaskListeners(displayId: Int, visibleTasksCount: Int) {
         visibleTasksListeners.forEach { (listener, executor) ->
             executor.execute { listener.onTasksVisibilityChanged(displayId, visibleTasksCount) }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorViewModel.java b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorViewModel.java
index 272508f46d33..ffccc3570524 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorViewModel.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorViewModel.java
@@ -106,6 +106,7 @@ import com.android.wm.shell.common.MultiInstanceHelper;
 import com.android.wm.shell.common.ShellExecutor;
 import com.android.wm.shell.common.SyncTransactionQueue;
 import com.android.wm.shell.desktopmode.DesktopActivityOrientationChangeHandler;
+import com.android.wm.shell.desktopmode.DesktopModeTaskRepository;
 import com.android.wm.shell.desktopmode.DesktopModeVisualIndicator;
 import com.android.wm.shell.desktopmode.DesktopTasksController;
 import com.android.wm.shell.desktopmode.DesktopTasksController.SnapPosition;
@@ -154,6 +155,7 @@ public class DesktopModeWindowDecorViewModel implements WindowDecorViewModel {
     private final ActivityTaskManager mActivityTaskManager;
     private final ShellCommandHandler mShellCommandHandler;
     private final ShellTaskOrganizer mTaskOrganizer;
+    private final DesktopModeTaskRepository mDesktopRepository;
     private final ShellController mShellController;
     private final Context mContext;
     private final @ShellMainThread Handler mMainHandler;
@@ -224,6 +226,7 @@ public class DesktopModeWindowDecorViewModel implements WindowDecorViewModel {
             ShellCommandHandler shellCommandHandler,
             IWindowManager windowManager,
             ShellTaskOrganizer taskOrganizer,
+            DesktopModeTaskRepository desktopRepository,
             DisplayController displayController,
             ShellController shellController,
             DisplayInsetsController displayInsetsController,
@@ -248,6 +251,7 @@ public class DesktopModeWindowDecorViewModel implements WindowDecorViewModel {
                 shellCommandHandler,
                 windowManager,
                 taskOrganizer,
+                desktopRepository,
                 displayController,
                 shellController,
                 displayInsetsController,
@@ -281,6 +285,7 @@ public class DesktopModeWindowDecorViewModel implements WindowDecorViewModel {
             ShellCommandHandler shellCommandHandler,
             IWindowManager windowManager,
             ShellTaskOrganizer taskOrganizer,
+            DesktopModeTaskRepository desktopRepository,
             DisplayController displayController,
             ShellController shellController,
             DisplayInsetsController displayInsetsController,
@@ -308,6 +313,7 @@ public class DesktopModeWindowDecorViewModel implements WindowDecorViewModel {
         mBgExecutor = bgExecutor;
         mActivityTaskManager = mContext.getSystemService(ActivityTaskManager.class);
         mTaskOrganizer = taskOrganizer;
+        mDesktopRepository = desktopRepository;
         mShellController = shellController;
         mDisplayController = displayController;
         mDisplayInsetsController = displayInsetsController;
@@ -1394,6 +1400,7 @@ public class DesktopModeWindowDecorViewModel implements WindowDecorViewModel {
                         mContext.createContextAsUser(UserHandle.of(taskInfo.userId), 0 /* flags */),
                         mDisplayController,
                         mSplitScreenController,
+                        mDesktopRepository,
                         mTaskOrganizer,
                         taskInfo,
                         taskSurface,
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecoration.java b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecoration.java
index 8a53f5ba4a51..a8890e4ea28e 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecoration.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecoration.java
@@ -89,6 +89,7 @@ import com.android.wm.shell.common.MultiInstanceHelper;
 import com.android.wm.shell.common.ShellExecutor;
 import com.android.wm.shell.common.SyncTransactionQueue;
 import com.android.wm.shell.desktopmode.CaptionState;
+import com.android.wm.shell.desktopmode.DesktopModeTaskRepository;
 import com.android.wm.shell.desktopmode.WindowDecorCaptionHandleRepository;
 import com.android.wm.shell.shared.annotations.ShellBackgroundThread;
 import com.android.wm.shell.shared.desktopmode.DesktopModeStatus;
@@ -188,12 +189,14 @@ public class DesktopModeWindowDecoration extends WindowDecoration<WindowDecorLin
     private final Runnable mCapturedLinkExpiredRunnable = this::onCapturedLinkExpired;
     private final MultiInstanceHelper mMultiInstanceHelper;
     private final WindowDecorCaptionHandleRepository mWindowDecorCaptionHandleRepository;
+    private final DesktopModeTaskRepository mDesktopRepository;
 
     DesktopModeWindowDecoration(
             Context context,
             @NonNull Context userContext,
             DisplayController displayController,
             SplitScreenController splitScreenController,
+            DesktopModeTaskRepository desktopRepository,
             ShellTaskOrganizer taskOrganizer,
             ActivityManager.RunningTaskInfo taskInfo,
             SurfaceControl taskSurface,
@@ -207,8 +210,8 @@ public class DesktopModeWindowDecoration extends WindowDecoration<WindowDecorLin
             AssistContentRequester assistContentRequester,
             MultiInstanceHelper multiInstanceHelper,
             WindowDecorCaptionHandleRepository windowDecorCaptionHandleRepository) {
-        this (context, userContext, displayController, splitScreenController, taskOrganizer,
-                taskInfo, taskSurface, handler, bgExecutor, choreographer, syncQueue,
+        this (context, userContext, displayController, splitScreenController, desktopRepository,
+                taskOrganizer, taskInfo, taskSurface, handler, bgExecutor, choreographer, syncQueue,
                 appHeaderViewHolderFactory, rootTaskDisplayAreaOrganizer, genericLinksParser,
                 assistContentRequester,
                 SurfaceControl.Builder::new, SurfaceControl.Transaction::new,
@@ -225,6 +228,7 @@ public class DesktopModeWindowDecoration extends WindowDecoration<WindowDecorLin
             @NonNull Context userContext,
             DisplayController displayController,
             SplitScreenController splitScreenController,
+            DesktopModeTaskRepository desktopRepository,
             ShellTaskOrganizer taskOrganizer,
             ActivityManager.RunningTaskInfo taskInfo,
             SurfaceControl taskSurface,
@@ -264,6 +268,7 @@ public class DesktopModeWindowDecoration extends WindowDecoration<WindowDecorLin
         mMultiInstanceHelper = multiInstanceHelper;
         mWindowManagerWrapper = windowManagerWrapper;
         mWindowDecorCaptionHandleRepository = windowDecorCaptionHandleRepository;
+        mDesktopRepository = desktopRepository;
     }
 
     /**
@@ -479,11 +484,18 @@ public class DesktopModeWindowDecoration extends WindowDecoration<WindowDecorLin
         if (canEnterDesktopMode(mContext) && Flags.enableDesktopWindowingAppHandleEducation()) {
             notifyCaptionStateChanged();
         }
-        mWindowDecorViewHolder.bindData(mTaskInfo,
-                position,
-                mResult.mCaptionWidth,
-                mResult.mCaptionHeight,
-                isCaptionVisible());
+
+        if (isAppHandle(mWindowDecorViewHolder)) {
+            mWindowDecorViewHolder.bindData(new AppHandleViewHolder.HandleData(
+                    mTaskInfo, position, mResult.mCaptionWidth, mResult.mCaptionHeight,
+                    isCaptionVisible()
+            ));
+        } else {
+            mWindowDecorViewHolder.bindData(new AppHeaderViewHolder.HeaderData(
+                    mTaskInfo, TaskInfoKt.getRequestingImmersive(mTaskInfo),
+                    mDesktopRepository.isTaskInFullImmersiveState(mTaskInfo.taskId)
+            ));
+        }
         Trace.endSection();
 
         if (!mTaskInfo.isFocused) {
@@ -1531,6 +1543,7 @@ public class DesktopModeWindowDecoration extends WindowDecoration<WindowDecorLin
                 @NonNull Context userContext,
                 DisplayController displayController,
                 SplitScreenController splitScreenController,
+                DesktopModeTaskRepository desktopRepository,
                 ShellTaskOrganizer taskOrganizer,
                 ActivityManager.RunningTaskInfo taskInfo,
                 SurfaceControl taskSurface,
@@ -1549,6 +1562,7 @@ public class DesktopModeWindowDecoration extends WindowDecoration<WindowDecorLin
                     userContext,
                     displayController,
                     splitScreenController,
+                    desktopRepository,
                     taskOrganizer,
                     taskInfo,
                     taskSurface,
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/viewholder/AppHandleViewHolder.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/viewholder/AppHandleViewHolder.kt
index 8c102ebfb590..b5700ffb046b 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/viewholder/AppHandleViewHolder.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/viewholder/AppHandleViewHolder.kt
@@ -42,6 +42,7 @@ import com.android.wm.shell.R
 import com.android.wm.shell.shared.animation.Interpolators
 import com.android.wm.shell.windowdecor.WindowManagerWrapper
 import com.android.wm.shell.windowdecor.additionalviewcontainer.AdditionalSystemViewContainer
+import com.android.wm.shell.windowdecor.viewholder.WindowDecorationViewHolder.Data
 
 /**
  * A desktop mode window decoration used when the window is in full "focus" (i.e. fullscreen/split).
@@ -53,11 +54,20 @@ internal class AppHandleViewHolder(
     onCaptionButtonClickListener: OnClickListener,
     private val windowManagerWrapper: WindowManagerWrapper,
     private val handler: Handler
-) : WindowDecorationViewHolder(rootView) {
+) : WindowDecorationViewHolder<AppHandleViewHolder.HandleData>(rootView) {
 
     companion object {
         private const val CAPTION_HANDLE_ANIMATION_DURATION: Long = 100
     }
+
+    data class HandleData(
+        val taskInfo: RunningTaskInfo,
+        val position: Point,
+        val width: Int,
+        val height: Int,
+        val isCaptionVisible: Boolean
+    ) : Data()
+
     private lateinit var taskInfo: RunningTaskInfo
     private val captionView: View = rootView.requireViewById(R.id.desktop_mode_caption)
     private val captionHandle: ImageButton = rootView.requireViewById(R.id.caption_handle)
@@ -89,7 +99,11 @@ internal class AppHandleViewHolder(
         }
     }
 
-    override fun bindData(
+    override fun bindData(data: HandleData) {
+        bindData(data.taskInfo, data.position, data.width, data.height, data.isCaptionVisible)
+    }
+
+    private fun bindData(
         taskInfo: RunningTaskInfo,
         position: Point,
         width: Int,
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/viewholder/AppHeaderViewHolder.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/viewholder/AppHeaderViewHolder.kt
index 306103c9b492..52bf40062cdb 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/viewholder/AppHeaderViewHolder.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/viewholder/AppHeaderViewHolder.kt
@@ -16,12 +16,12 @@
 package com.android.wm.shell.windowdecor.viewholder
 
 import android.annotation.ColorInt
+import android.annotation.DrawableRes
 import android.app.ActivityManager.RunningTaskInfo
 import android.content.res.ColorStateList
 import android.content.res.Configuration
 import android.graphics.Bitmap
 import android.graphics.Color
-import android.graphics.Point
 import android.graphics.Rect
 import android.graphics.drawable.LayerDrawable
 import android.graphics.drawable.RippleDrawable
@@ -60,7 +60,6 @@ import com.android.wm.shell.windowdecor.common.OPACITY_65
 import com.android.wm.shell.windowdecor.common.Theme
 import com.android.wm.shell.windowdecor.extension.isLightCaptionBarAppearance
 import com.android.wm.shell.windowdecor.extension.isTransparentCaptionBarAppearance
-import com.android.wm.shell.windowdecor.extension.requestingImmersive
 
 /**
  * A desktop mode window decoration used when the window is floating (i.e. freeform). It hosts
@@ -76,7 +75,13 @@ class AppHeaderViewHolder(
         appName: CharSequence,
         appIconBitmap: Bitmap,
         onMaximizeHoverAnimationFinishedListener: () -> Unit
-) : WindowDecorationViewHolder(rootView) {
+) : WindowDecorationViewHolder<AppHeaderViewHolder.HeaderData>(rootView) {
+
+    data class HeaderData(
+        val taskInfo: RunningTaskInfo,
+        val isRequestingImmersive: Boolean,
+        val inFullImmersiveState: Boolean,
+    ) : Data()
 
     private val decorThemeUtil = DecorThemeUtil(context)
     private val lightColors = dynamicLightColorScheme(context)
@@ -153,15 +158,17 @@ class AppHeaderViewHolder(
                 onMaximizeHoverAnimationFinishedListener
     }
 
-    override fun bindData(
+    override fun bindData(data: HeaderData) {
+        bindData(data.taskInfo, data.isRequestingImmersive, data.inFullImmersiveState)
+    }
+
+    private fun bindData(
         taskInfo: RunningTaskInfo,
-        position: Point,
-        width: Int,
-        height: Int,
-        isCaptionVisible: Boolean
+        isRequestingImmersive: Boolean,
+        inFullImmersiveState: Boolean,
     ) {
         if (DesktopModeFlags.ENABLE_THEMED_APP_HEADERS.isTrue()) {
-            bindDataWithThemedHeaders(taskInfo)
+            bindDataWithThemedHeaders(taskInfo, isRequestingImmersive, inFullImmersiveState)
         } else {
             bindDataLegacy(taskInfo)
         }
@@ -200,7 +207,11 @@ class AppHeaderViewHolder(
         minimizeWindowButton.isGone = !enableMinimizeButton()
     }
 
-    private fun bindDataWithThemedHeaders(taskInfo: RunningTaskInfo) {
+    private fun bindDataWithThemedHeaders(
+        taskInfo: RunningTaskInfo,
+        requestingImmersive: Boolean,
+        inFullImmersiveState: Boolean
+    ) {
         val header = fillHeaderInfo(taskInfo)
         val headerStyle = getHeaderStyle(header)
 
@@ -254,13 +265,7 @@ class AppHeaderViewHolder(
                     drawableInsets = maximizeDrawableInsets
                 )
             )
-            setIcon(
-                if (taskInfo.requestingImmersive && Flags.enableFullyImmersiveInDesktop()) {
-                    R.drawable.decor_desktop_mode_immersive_button_dark
-                } else {
-                    R.drawable.decor_desktop_mode_maximize_button_dark
-                }
-            )
+            setIcon(getMaximizeButtonIcon(requestingImmersive, inFullImmersiveState))
         }
         // Close button.
         closeWindowButton.apply {
@@ -331,6 +336,32 @@ class AppHeaderViewHolder(
         }
     }
 
+    @DrawableRes
+    private fun getMaximizeButtonIcon(
+        requestingImmersive: Boolean,
+        inFullImmersiveState: Boolean
+    ): Int = when {
+        shouldShowEnterFullImmersiveIcon(requestingImmersive, inFullImmersiveState) -> {
+            R.drawable.decor_desktop_mode_immersive_button_dark
+        }
+        shouldShowExitFullImmersiveIcon(requestingImmersive, inFullImmersiveState) -> {
+            R.drawable.decor_desktop_mode_immersive_exit_button_dark
+        }
+        else -> R.drawable.decor_desktop_mode_maximize_button_dark
+    }
+
+    private fun shouldShowEnterFullImmersiveIcon(
+        requestingImmersive: Boolean,
+        inFullImmersiveState: Boolean
+    ): Boolean = Flags.enableFullyImmersiveInDesktop()
+            && requestingImmersive && !inFullImmersiveState
+
+    private fun shouldShowExitFullImmersiveIcon(
+        requestingImmersive: Boolean,
+        inFullImmersiveState: Boolean
+    ): Boolean = Flags.enableFullyImmersiveInDesktop()
+            && requestingImmersive && inFullImmersiveState
+
     private fun getHeaderStyle(header: Header): HeaderStyle {
         return HeaderStyle(
             background = getHeaderBackground(header),
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/viewholder/WindowDecorationViewHolder.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/viewholder/WindowDecorationViewHolder.kt
index 5ea55b367703..1fe743da966a 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/viewholder/WindowDecorationViewHolder.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/viewholder/WindowDecorationViewHolder.kt
@@ -17,31 +17,28 @@ package com.android.wm.shell.windowdecor.viewholder
 
 import android.app.ActivityManager.RunningTaskInfo
 import android.content.Context
-import android.graphics.Point
 import android.view.View
+import com.android.wm.shell.windowdecor.viewholder.WindowDecorationViewHolder.Data
 
 /**
  * Encapsulates the root [View] of a window decoration and its children to facilitate looking up
  * children (via findViewById) and updating to the latest data from [RunningTaskInfo].
  */
-abstract class WindowDecorationViewHolder(rootView: View) {
+abstract class WindowDecorationViewHolder<T : Data>(rootView: View) {
   val context: Context = rootView.context
 
   /**
    * A signal to the view holder that new data is available and that the views should be updated to
    * reflect it.
    */
-  abstract fun bindData(
-    taskInfo: RunningTaskInfo,
-    position: Point,
-    width: Int,
-    height: Int,
-    isCaptionVisible: Boolean
-  )
+  abstract fun bindData(data: T)
 
   /** Callback when the handle menu is opened. */
   abstract fun onHandleMenuOpened()
 
   /** Callback when the handle menu is closed. */
   abstract fun onHandleMenuClosed()
+
+  /** Data clas that contains the information needed to update the view holder. */
+  abstract class Data
 }
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopModeTaskRepositoryTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopModeTaskRepositoryTest.kt
index bc40d89009bc..794f9d819f4b 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopModeTaskRepositoryTest.kt
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopModeTaskRepositoryTest.kt
@@ -882,6 +882,51 @@ class DesktopModeTaskRepositoryTest : ShellTestCase() {
         assertThat(tasks).containsExactly(1, 3).inOrder()
     }
 
+    @Test
+    fun setTaskInFullImmersiveState_savedAsInImmersiveState() {
+        assertThat(repo.isTaskInFullImmersiveState(taskId = 1)).isFalse()
+
+        repo.setTaskInFullImmersiveState(DEFAULT_DESKTOP_ID, taskId = 1, immersive = true)
+
+        assertThat(repo.isTaskInFullImmersiveState(taskId = 1)).isTrue()
+    }
+
+    @Test
+    fun removeTaskInFullImmersiveState_removedAsInImmersiveState() {
+        repo.setTaskInFullImmersiveState(DEFAULT_DESKTOP_ID, taskId = 1, immersive = true)
+        assertThat(repo.isTaskInFullImmersiveState(taskId = 1)).isTrue()
+
+        repo.setTaskInFullImmersiveState(DEFAULT_DESKTOP_ID, taskId = 1, immersive = false)
+
+        assertThat(repo.isTaskInFullImmersiveState(taskId = 1)).isFalse()
+    }
+
+    @Test
+    fun removeTaskInFullImmersiveState_otherWasImmersive_otherRemainsImmersive() {
+        repo.setTaskInFullImmersiveState(DEFAULT_DESKTOP_ID, taskId = 1, immersive = true)
+
+        repo.setTaskInFullImmersiveState(DEFAULT_DESKTOP_ID, taskId = 2, immersive = false)
+
+        assertThat(repo.isTaskInFullImmersiveState(taskId = 1)).isTrue()
+    }
+
+    @Test
+    fun setTaskInFullImmersiveState_sameDisplay_overridesExistingFullImmersiveTask() {
+        repo.setTaskInFullImmersiveState(DEFAULT_DESKTOP_ID, taskId = 1, immersive = true)
+        repo.setTaskInFullImmersiveState(DEFAULT_DESKTOP_ID, taskId = 2, immersive = true)
+
+        assertThat(repo.isTaskInFullImmersiveState(taskId = 1)).isFalse()
+        assertThat(repo.isTaskInFullImmersiveState(taskId = 2)).isTrue()
+    }
+
+    @Test
+    fun setTaskInFullImmersiveState_differentDisplay_bothAreImmersive() {
+        repo.setTaskInFullImmersiveState(DEFAULT_DESKTOP_ID, taskId = 1, immersive = true)
+        repo.setTaskInFullImmersiveState(DEFAULT_DESKTOP_ID + 1, taskId = 2, immersive = true)
+
+        assertThat(repo.isTaskInFullImmersiveState(taskId = 1)).isTrue()
+        assertThat(repo.isTaskInFullImmersiveState(taskId = 2)).isTrue()
+    }
 
     class TestListener : DesktopModeTaskRepository.ActiveTasksListener {
         var activeChangesOnDefaultDisplay = 0
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorViewModelTests.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorViewModelTests.kt
index 3051714b5ae8..6959187a6837 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorViewModelTests.kt
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorViewModelTests.kt
@@ -85,6 +85,7 @@ import com.android.wm.shell.common.MultiInstanceHelper
 import com.android.wm.shell.common.ShellExecutor
 import com.android.wm.shell.common.SyncTransactionQueue
 import com.android.wm.shell.desktopmode.DesktopActivityOrientationChangeHandler
+import com.android.wm.shell.desktopmode.DesktopModeTaskRepository
 import com.android.wm.shell.desktopmode.DesktopTasksController
 import com.android.wm.shell.desktopmode.DesktopTasksController.SnapPosition
 import com.android.wm.shell.desktopmode.DesktopTasksLimiter
@@ -155,6 +156,7 @@ class DesktopModeWindowDecorViewModelTests : ShellTestCase() {
     @Mock private lateinit var mockTaskOrganizer: ShellTaskOrganizer
     @Mock private lateinit var mockDisplayController: DisplayController
     @Mock private lateinit var mockSplitScreenController: SplitScreenController
+    @Mock private lateinit var mockDesktopRepository: DesktopModeTaskRepository
     @Mock private lateinit var mockDisplayLayout: DisplayLayout
     @Mock private lateinit var displayInsetsController: DisplayInsetsController
     @Mock private lateinit var mockSyncQueue: SyncTransactionQueue
@@ -225,6 +227,7 @@ class DesktopModeWindowDecorViewModelTests : ShellTestCase() {
                 mockShellCommandHandler,
                 mockWindowManager,
                 mockTaskOrganizer,
+                mockDesktopRepository,
                 mockDisplayController,
                 mockShellController,
                 displayInsetsController,
@@ -1213,8 +1216,8 @@ class DesktopModeWindowDecorViewModelTests : ShellTestCase() {
         val decoration = mock(DesktopModeWindowDecoration::class.java)
         whenever(
             mockDesktopModeWindowDecorFactory.create(
-                any(), any(), any(), any(), any(), eq(task), any(), any(), any(), any(), any(),
-                any(), any(), any(), any(), any(), any())
+                any(), any(), any(), any(), any(), any(), eq(task), any(), any(), any(), any(),
+                any(), any(), any(), any(), any(), any(), any())
         ).thenReturn(decoration)
         decoration.mTaskInfo = task
         whenever(decoration.isFocused).thenReturn(task.isFocused)
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorationTests.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorationTests.java
index f007115c6dab..a9713b4f0233 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorationTests.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorationTests.java
@@ -102,6 +102,7 @@ import com.android.wm.shell.common.MultiInstanceHelper;
 import com.android.wm.shell.common.ShellExecutor;
 import com.android.wm.shell.common.SyncTransactionQueue;
 import com.android.wm.shell.desktopmode.CaptionState;
+import com.android.wm.shell.desktopmode.DesktopModeTaskRepository;
 import com.android.wm.shell.desktopmode.WindowDecorCaptionHandleRepository;
 import com.android.wm.shell.shared.desktopmode.DesktopModeStatus;
 import com.android.wm.shell.splitscreen.SplitScreenController;
@@ -157,6 +158,8 @@ public class DesktopModeWindowDecorationTests extends ShellTestCase {
     @Mock
     private ShellTaskOrganizer mMockShellTaskOrganizer;
     @Mock
+    private DesktopModeTaskRepository mMockDesktopRepository;
+    @Mock
     private Choreographer mMockChoreographer;
     @Mock
     private SyncTransactionQueue mMockSyncQueue;
@@ -1086,9 +1089,9 @@ public class DesktopModeWindowDecorationTests extends ShellTestCase {
             boolean relayout) {
         final DesktopModeWindowDecoration windowDecor = new DesktopModeWindowDecoration(mContext,
                 mContext, mMockDisplayController, mMockSplitScreenController,
-                mMockShellTaskOrganizer, taskInfo, mMockSurfaceControl, mMockHandler, mBgExecutor,
-                mMockChoreographer, mMockSyncQueue, mMockAppHeaderViewHolderFactory,
-                mMockRootTaskDisplayAreaOrganizer,
+                mMockDesktopRepository, mMockShellTaskOrganizer, taskInfo, mMockSurfaceControl,
+                mMockHandler, mBgExecutor, mMockChoreographer, mMockSyncQueue,
+                mMockAppHeaderViewHolderFactory, mMockRootTaskDisplayAreaOrganizer,
                 mMockGenericLinksParser, mMockAssistContentRequester, SurfaceControl.Builder::new,
                 mMockTransactionSupplier, WindowContainerTransaction::new, SurfaceControl::new,
                 new WindowManagerWrapper(mMockWindowManager), mMockSurfaceControlViewHostFactory,