Merge "Implement persistent repository." into main
diff --git a/libs/WindowManager/Shell/Android.bp b/libs/WindowManager/Shell/Android.bp
index 94809f2..f857429 100644
--- a/libs/WindowManager/Shell/Android.bp
+++ b/libs/WindowManager/Shell/Android.bp
@@ -147,8 +147,10 @@
 java_library {
     name: "WindowManager-Shell-lite-proto",
 
-    srcs: ["src/com/android/wm/shell/desktopmode/education/data/proto/**/*.proto"],
-
+    srcs: [
+        "src/com/android/wm/shell/desktopmode/education/data/proto/**/*.proto",
+        "src/com/android/wm/shell/desktopmode/persistence/*.proto",
+    ],
     proto: {
         type: "lite",
     },
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 80a9b67..308bd0b 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
@@ -79,6 +79,7 @@
 import com.android.wm.shell.desktopmode.education.AppHandleEducationController;
 import com.android.wm.shell.desktopmode.education.AppHandleEducationFilter;
 import com.android.wm.shell.desktopmode.education.data.AppHandleEducationDatastoreRepository;
+import com.android.wm.shell.desktopmode.persistence.DesktopPersistentRepository;
 import com.android.wm.shell.draganddrop.DragAndDropController;
 import com.android.wm.shell.draganddrop.GlobalDragListener;
 import com.android.wm.shell.freeform.FreeformComponents;
@@ -712,8 +713,14 @@
     @WMSingleton
     @Provides
     @DynamicOverride
-    static DesktopModeTaskRepository provideDesktopModeTaskRepository() {
-        return new DesktopModeTaskRepository();
+    static DesktopModeTaskRepository provideDesktopModeTaskRepository(
+            Context context,
+            ShellInit shellInit,
+            DesktopPersistentRepository desktopPersistentRepository,
+            @ShellMainThread CoroutineScope mainScope
+    ) {
+        return new DesktopModeTaskRepository(context, shellInit, desktopPersistentRepository,
+                mainScope);
     }
 
     @WMSingleton
@@ -798,6 +805,14 @@
                 shellTaskOrganizer, appHandleEducationDatastoreRepository, applicationScope);
     }
 
+    @WMSingleton
+    @Provides
+    static DesktopPersistentRepository provideDesktopPersistentRepository(
+            Context context,
+            @ShellBackgroundThread CoroutineScope bgScope) {
+        return new DesktopPersistentRepository(context, bgScope);
+    }
+
     //
     // Drag and drop
     //
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 9d04169..759ed03 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
@@ -16,6 +16,7 @@
 
 package com.android.wm.shell.desktopmode
 
+import android.content.Context
 import android.graphics.Rect
 import android.graphics.Region
 import android.util.ArrayMap
@@ -27,13 +28,27 @@
 import androidx.core.util.keyIterator
 import androidx.core.util.valueIterator
 import com.android.internal.protolog.ProtoLog
+import com.android.window.flags.Flags
+import com.android.wm.shell.desktopmode.persistence.DesktopPersistentRepository
+import com.android.wm.shell.desktopmode.persistence.DesktopTask
+import com.android.wm.shell.desktopmode.persistence.DesktopTaskState
 import com.android.wm.shell.protolog.ShellProtoLogGroup.WM_SHELL_DESKTOP_MODE
+import com.android.wm.shell.shared.annotations.ShellMainThread
+import com.android.wm.shell.shared.desktopmode.DesktopModeStatus
+import com.android.wm.shell.sysui.ShellInit
 import java.io.PrintWriter
 import java.util.concurrent.Executor
 import java.util.function.Consumer
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.launch
 
 /** Tracks task data for Desktop Mode. */
-class DesktopModeTaskRepository {
+class DesktopModeTaskRepository (
+    private val context: Context,
+    shellInit: ShellInit,
+    private val persistentRepository: DesktopPersistentRepository,
+    @ShellMainThread private val mainCoroutineScope: CoroutineScope,
+){
 
     /**
      * Task data tracked per desktop.
@@ -54,7 +69,15 @@
         // TODO(b/332682201): Remove when the repository state is updated via TransitionObserver
         val closingTasks: ArraySet<Int> = ArraySet(),
         val freeformTasksInZOrder: ArrayList<Int> = ArrayList(),
-    )
+    ) {
+        fun deepCopy(): DesktopTaskData = DesktopTaskData(
+            activeTasks = ArraySet(activeTasks),
+            visibleTasks = ArraySet(visibleTasks),
+            minimizedTasks = ArraySet(minimizedTasks),
+            closingTasks = ArraySet(closingTasks),
+            freeformTasksInZOrder = ArrayList(freeformTasksInZOrder)
+        )
+    }
 
     /* Current wallpaper activity token to remove wallpaper activity when last task is removed. */
     var wallpaperActivityToken: WindowContainerToken? = null
@@ -77,6 +100,40 @@
             this[displayId] ?: DesktopTaskData().also { this[displayId] = it }
     }
 
+    init {
+        if (DesktopModeStatus.canEnterDesktopMode(context)) {
+            shellInit.addInitCallback(::initRepoFromPersistentStorage, this)
+        }
+    }
+
+    private fun initRepoFromPersistentStorage() {
+        if (!Flags.enableDesktopWindowingPersistence()) return
+        //  TODO: b/365962554 - Handle the case that user moves to desktop before it's initialized
+        mainCoroutineScope.launch {
+            val desktop = persistentRepository.readDesktop()
+            val maxTasks =
+                DesktopModeStatus.getMaxTaskLimit(context).takeIf { it > 0 }
+                    ?: desktop.zOrderedTasksCount
+
+            desktop.zOrderedTasksList
+                // Reverse it so we initialize the repo from bottom to top.
+                .reversed()
+                .map { taskId ->
+                    desktop.tasksByTaskIdMap.getOrDefault(
+                        taskId,
+                        DesktopTask.getDefaultInstance()
+                    )
+                }
+                .filter { task -> task.desktopTaskState == DesktopTaskState.VISIBLE }
+                .take(maxTasks)
+                .forEach { task ->
+                    addOrMoveFreeformTaskToTop(desktop.displayId, task.taskId)
+                    addActiveTask(desktop.displayId, task.taskId)
+                    updateTaskVisibility(desktop.displayId, task.taskId, visible = false)
+                }
+        }
+    }
+
     /** Adds [activeTasksListener] to be notified of updates to active tasks. */
     fun addActiveTaskListener(activeTasksListener: ActiveTasksListener) {
         activeTasksListeners.add(activeTasksListener)
@@ -266,12 +323,18 @@
         desktopTaskDataByDisplayId.getOrCreate(displayId).freeformTasksInZOrder.add(0, taskId)
         // Unminimize the task if it is minimized.
         unminimizeTask(displayId, taskId)
+        if (Flags.enableDesktopWindowingPersistence()) {
+            updatePersistentRepository(displayId)
+        }
     }
 
     /** Minimizes the task for [taskId] and [displayId] */
     fun minimizeTask(displayId: Int, taskId: Int) {
         logD("Minimize Task: display=%d, task=%d", displayId, taskId)
         desktopTaskDataByDisplayId.getOrCreate(displayId).minimizedTasks.add(taskId)
+        if (Flags.enableDesktopWindowingPersistence()) {
+            updatePersistentRepository(displayId)
+        }
     }
 
     /** Unminimizes the task for [taskId] and [displayId] */
@@ -315,7 +378,10 @@
         // Remove task from unminimized task if it is minimized.
         unminimizeTask(displayId, taskId)
         removeActiveTask(taskId)
-        updateTaskVisibility(displayId, taskId, visible = false);
+        updateTaskVisibility(displayId, taskId, visible = false)
+        if (Flags.enableDesktopWindowingPersistence()) {
+            updatePersistentRepository(displayId)
+        }
     }
 
     /**
@@ -352,6 +418,27 @@
     fun saveBoundsBeforeMaximize(taskId: Int, bounds: Rect) =
         boundsBeforeMaximizeByTaskId.set(taskId, Rect(bounds))
 
+    private fun updatePersistentRepository(displayId: Int) {
+        // Create a deep copy of the data
+        desktopTaskDataByDisplayId[displayId]?.deepCopy()?.let { desktopTaskDataByDisplayIdCopy ->
+            mainCoroutineScope.launch {
+                try {
+                    persistentRepository.addOrUpdateDesktop(
+                        visibleTasks = desktopTaskDataByDisplayIdCopy.visibleTasks,
+                        minimizedTasks = desktopTaskDataByDisplayIdCopy.minimizedTasks,
+                        freeformTasksInZOrder = desktopTaskDataByDisplayIdCopy.freeformTasksInZOrder
+                    )
+                } catch (exception: Exception) {
+                    logE(
+                        "An exception occurred while updating the persistent repository \n%s",
+                        exception.stackTrace
+                    )
+                }
+            }
+        }
+    }
+
+
     internal fun dump(pw: PrintWriter, prefix: String) {
         val innerPrefix = "$prefix  "
         pw.println("${prefix}DesktopModeTaskRepository")
@@ -390,6 +477,10 @@
         ProtoLog.w(WM_SHELL_DESKTOP_MODE, "%s: $msg", TAG, *arguments)
     }
 
+    private fun logE(msg: String, vararg arguments: Any?) {
+        ProtoLog.e(WM_SHELL_DESKTOP_MODE, "%s: $msg", TAG, *arguments)
+    }
+
     companion object {
         private const val TAG = "DesktopModeTaskRepository"
     }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksController.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksController.kt
index f3ae3ed..968f40c 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksController.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksController.kt
@@ -58,6 +58,7 @@
 import com.android.internal.jank.InteractionJankMonitor
 import com.android.internal.policy.ScreenDecorationsUtils
 import com.android.internal.protolog.ProtoLog
+import com.android.window.flags.Flags
 import com.android.wm.shell.RootTaskDisplayAreaOrganizer
 import com.android.wm.shell.ShellTaskOrganizer
 import com.android.wm.shell.common.DisplayController
@@ -80,8 +81,8 @@
 import com.android.wm.shell.recents.RecentTasksController
 import com.android.wm.shell.recents.RecentsTransitionHandler
 import com.android.wm.shell.recents.RecentsTransitionStateListener
-import com.android.wm.shell.shared.TransitionUtil
 import com.android.wm.shell.shared.ShellSharedConstants
+import com.android.wm.shell.shared.TransitionUtil
 import com.android.wm.shell.shared.annotations.ExternalThread
 import com.android.wm.shell.shared.annotations.ShellMainThread
 import android.window.flags.DesktopModeFlags
@@ -728,7 +729,7 @@
                 // exclude current task since maximize/restore transition has not taken place yet.
                 .filterNot { taskId -> taskId == excludeTaskId }
                 .any { taskId ->
-                    val taskInfo = shellTaskOrganizer.getRunningTaskInfo(taskId)!!
+                    val taskInfo = shellTaskOrganizer.getRunningTaskInfo(taskId) ?: return false
                     val displayLayout = displayController.getDisplayLayout(taskInfo.displayId)
                     val stableBounds = Rect().apply { displayLayout?.getStableBounds(this) }
                     logD("taskInfo = %s", taskInfo)
@@ -896,6 +897,7 @@
         val nonMinimizedTasksOrderedFrontToBack =
             taskRepository.getActiveNonMinimizedOrderedTasks(displayId)
         // If we're adding a new Task we might need to minimize an old one
+        // TODO(b/365725441): Handle non running task minimization
         val taskToMinimize: RunningTaskInfo? =
             if (newTaskIdInFront != null && desktopTasksLimiter.isPresent) {
                 desktopTasksLimiter
@@ -907,12 +909,26 @@
             } else {
                 null
             }
+
         nonMinimizedTasksOrderedFrontToBack
             // If there is a Task to minimize, let it stay behind the Home Task
             .filter { taskId -> taskId != taskToMinimize?.taskId }
-            .mapNotNull { taskId -> shellTaskOrganizer.getRunningTaskInfo(taskId) }
             .reversed() // Start from the back so the front task is brought forward last
-            .forEach { task -> wct.reorder(task.token, /* onTop= */ true) }
+            .forEach { taskId ->
+                val runningTaskInfo = shellTaskOrganizer.getRunningTaskInfo(taskId)
+                if (runningTaskInfo != null) {
+                    // Task is already running, reorder it to the front
+                    wct.reorder(runningTaskInfo.token, /* onTop= */ true)
+                } else if (Flags.enableDesktopWindowingPersistence()) {
+                    // Task is not running, start it
+                    wct.startTask(
+                        taskId,
+                        ActivityOptions.makeBasic().apply {
+                            launchWindowingMode = WINDOWING_MODE_FREEFORM
+                        }.toBundle(),
+                    )
+                }
+            }
 
         taskbarDesktopTaskListener?.
             onTaskbarCornerRoundingUpdate(doesAnyTaskRequireTaskbarRounding(displayId))
@@ -1211,6 +1227,7 @@
             wct.reorder(task.token, true)
             return wct
         }
+        // TODO(b/365723620): Handle non running tasks that were launched after reboot.
         // If task is already visible, it must have been handled already and added to desktop mode.
         // Cascade task only if it's not visible yet.
         if (DesktopModeFlags.ENABLE_CASCADING_WINDOWS.isTrue()
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/persistence/DesktopPersistentRepository.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/persistence/DesktopPersistentRepository.kt
new file mode 100644
index 0000000..3f41d7c
--- /dev/null
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/persistence/DesktopPersistentRepository.kt
@@ -0,0 +1,201 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.wm.shell.desktopmode.persistence
+
+import android.content.Context
+import android.util.ArraySet
+import android.util.Log
+import android.view.Display.DEFAULT_DISPLAY
+import androidx.datastore.core.CorruptionException
+import androidx.datastore.core.DataStore
+import androidx.datastore.core.DataStoreFactory
+import androidx.datastore.core.Serializer
+import androidx.datastore.dataStoreFile
+import com.android.framework.protobuf.InvalidProtocolBufferException
+import com.android.wm.shell.shared.annotations.ShellBackgroundThread
+import java.io.IOException
+import java.io.InputStream
+import java.io.OutputStream
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.catch
+import kotlinx.coroutines.flow.first
+
+/**
+ * Persistent repository for storing desktop mode related data.
+ *
+ * The main constructor is public only for testing purposes.
+ */
+class DesktopPersistentRepository(
+    private val dataStore: DataStore<DesktopPersistentRepositories>,
+) {
+    constructor(
+        context: Context,
+        @ShellBackgroundThread bgCoroutineScope: CoroutineScope,
+    ) : this(
+        DataStoreFactory.create(
+            serializer = DesktopPersistentRepositoriesSerializer,
+            produceFile = { context.dataStoreFile(DESKTOP_REPOSITORIES_DATASTORE_FILE) },
+            scope = bgCoroutineScope))
+
+    /** Provides `dataStore.data` flow and handles exceptions thrown during collection */
+    private val dataStoreFlow: Flow<DesktopPersistentRepositories> =
+        dataStore.data.catch { exception ->
+            // dataStore.data throws an IOException when an error is encountered when reading data
+            if (exception is IOException) {
+                Log.e(
+                    TAG,
+                    "Error in reading desktop mode related data from datastore, data is " +
+                        "stored in a file named $DESKTOP_REPOSITORIES_DATASTORE_FILE",
+                    exception)
+            } else {
+                throw exception
+            }
+        }
+
+    /**
+     * Reads and returns the [DesktopRepositoryState] proto object from the DataStore for a user. If
+     * the DataStore is empty or there's an error reading, it returns the default value of Proto.
+     */
+    private suspend fun getDesktopRepositoryState(
+        userId: Int = DEFAULT_USER_ID
+    ): DesktopRepositoryState =
+        try {
+            dataStoreFlow
+                .first()
+                .desktopRepoByUserMap
+                .getOrDefault(userId, DesktopRepositoryState.getDefaultInstance())
+        } catch (e: Exception) {
+            Log.e(TAG, "Unable to read from datastore", e)
+            DesktopRepositoryState.getDefaultInstance()
+        }
+
+    /**
+     * Reads the [Desktop] of a desktop filtering by the [userId] and [desktopId]. Executes the
+     * [callback] using the [mainCoroutineScope].
+     */
+    suspend fun readDesktop(
+        userId: Int = DEFAULT_USER_ID,
+        desktopId: Int = DEFAULT_DESKTOP_ID,
+    ): Desktop =
+        try {
+            val repository = getDesktopRepositoryState(userId)
+            repository.getDesktopOrThrow(desktopId)
+        } catch (e: Exception) {
+            Log.e(TAG, "Unable to get desktop info from persistent repository", e)
+            Desktop.getDefaultInstance()
+        }
+
+    /** Adds or updates a desktop stored in the datastore */
+    suspend fun addOrUpdateDesktop(
+        userId: Int = DEFAULT_USER_ID,
+        desktopId: Int = 0,
+        visibleTasks: ArraySet<Int> = ArraySet(),
+        minimizedTasks: ArraySet<Int> = ArraySet(),
+        freeformTasksInZOrder: ArrayList<Int> = ArrayList(),
+    ) {
+        // TODO: b/367609270 - Improve the API to support multi-user
+        try {
+            dataStore.updateData { desktopPersistentRepositories: DesktopPersistentRepositories ->
+                val currentRepository =
+                    desktopPersistentRepositories.getDesktopRepoByUserOrDefault(
+                        userId, DesktopRepositoryState.getDefaultInstance())
+                val desktop =
+                    getDesktop(currentRepository, desktopId)
+                        .toBuilder()
+                        .updateTaskStates(visibleTasks, minimizedTasks)
+                        .updateZOrder(freeformTasksInZOrder)
+
+                desktopPersistentRepositories
+                    .toBuilder()
+                    .putDesktopRepoByUser(
+                        userId,
+                        currentRepository
+                            .toBuilder()
+                            .putDesktop(desktopId, desktop.build())
+                            .build())
+                    .build()
+            }
+        } catch (exception: IOException) {
+            Log.e(
+                TAG,
+                "Error in updating desktop mode related data, data is " +
+                    "stored in a file named $DESKTOP_REPOSITORIES_DATASTORE_FILE",
+                exception)
+        }
+    }
+
+    private fun getDesktop(currentRepository: DesktopRepositoryState, desktopId: Int): Desktop =
+        // If there are no desktops set up, create one on the default display
+        currentRepository.getDesktopOrDefault(
+            desktopId,
+            Desktop.newBuilder().setDesktopId(desktopId).setDisplayId(DEFAULT_DISPLAY).build())
+
+    companion object {
+        private const val TAG = "DesktopPersistenceRepo"
+        private const val DESKTOP_REPOSITORIES_DATASTORE_FILE = "desktop_persistent_repositories.pb"
+
+        private const val DEFAULT_USER_ID = 1000
+        private const val DEFAULT_DESKTOP_ID = 0
+
+        object DesktopPersistentRepositoriesSerializer : Serializer<DesktopPersistentRepositories> {
+
+            override val defaultValue: DesktopPersistentRepositories =
+                DesktopPersistentRepositories.getDefaultInstance()
+
+            override suspend fun readFrom(input: InputStream): DesktopPersistentRepositories =
+                try {
+                    DesktopPersistentRepositories.parseFrom(input)
+                } catch (exception: InvalidProtocolBufferException) {
+                    throw CorruptionException("Cannot read proto.", exception)
+                }
+
+            override suspend fun writeTo(t: DesktopPersistentRepositories, output: OutputStream) =
+                t.writeTo(output)
+        }
+
+        private fun Desktop.Builder.updateTaskStates(
+            visibleTasks: ArraySet<Int>,
+            minimizedTasks: ArraySet<Int>
+        ): Desktop.Builder {
+            clearTasksByTaskId()
+            putAllTasksByTaskId(
+                visibleTasks.associateWith {
+                    createDesktopTask(it, state = DesktopTaskState.VISIBLE)
+                })
+            putAllTasksByTaskId(
+                minimizedTasks.associateWith {
+                    createDesktopTask(it, state = DesktopTaskState.MINIMIZED)
+                })
+            return this
+        }
+
+        private fun Desktop.Builder.updateZOrder(
+            freeformTasksInZOrder: ArrayList<Int>
+        ): Desktop.Builder {
+            clearZOrderedTasks()
+            addAllZOrderedTasks(freeformTasksInZOrder)
+            return this
+        }
+
+        private fun createDesktopTask(
+            taskId: Int,
+            state: DesktopTaskState = DesktopTaskState.VISIBLE
+        ): DesktopTask =
+            DesktopTask.newBuilder().setTaskId(taskId).setDesktopTaskState(state).build()
+    }
+}
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/persistence/persistent_desktop_repositories.proto b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/persistence/persistent_desktop_repositories.proto
new file mode 100644
index 0000000..0105231
--- /dev/null
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/persistence/persistent_desktop_repositories.proto
@@ -0,0 +1,33 @@
+syntax = "proto2";
+
+option java_package = "com.android.wm.shell.desktopmode.persistence";
+option java_multiple_files = true;
+
+// Represents the state of a task in desktop.
+enum DesktopTaskState {
+  VISIBLE = 0;
+  MINIMIZED = 1;
+}
+
+message DesktopTask {
+  optional int32 task_id = 1;
+  optional DesktopTaskState desktop_task_state= 2;
+}
+
+message Desktop {
+  optional int32 display_id = 1;
+  optional int32 desktop_id = 2;
+  // Stores a mapping between task id and the tasks. The key is the task id.
+  map<int32, DesktopTask> tasks_by_task_id = 3;
+  repeated int32 z_ordered_tasks = 4;
+}
+
+message DesktopRepositoryState {
+  // Stores a mapping between a repository and the desktops in it. The key is the desktop id.
+  map<int32, Desktop> desktop = 1;
+}
+
+message DesktopPersistentRepositories {
+  // Stores a mapping between a user and their desktop repository. The key is the user id.
+  map<int32, DesktopRepositoryState> desktop_repo_by_user = 1;
+}
\ No newline at end of file
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopActivityOrientationChangeHandlerTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopActivityOrientationChangeHandlerTest.kt
index b14f163..628c9cdd 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopActivityOrientationChangeHandlerTest.kt
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopActivityOrientationChangeHandlerTest.kt
@@ -41,12 +41,22 @@
 import com.android.wm.shell.common.TaskStackListenerImpl
 import com.android.wm.shell.desktopmode.DesktopTestHelpers.Companion.createFreeformTask
 import com.android.wm.shell.desktopmode.DesktopTestHelpers.Companion.createFullscreenTask
+import com.android.wm.shell.desktopmode.persistence.Desktop
+import com.android.wm.shell.desktopmode.persistence.DesktopPersistentRepository
 import com.android.wm.shell.shared.desktopmode.DesktopModeStatus
 import com.android.wm.shell.sysui.ShellInit
 import com.android.wm.shell.transition.Transitions
 import junit.framework.Assert.assertEquals
 import junit.framework.Assert.assertTrue
 import kotlin.test.assertNotNull
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.cancel
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.SupervisorJob
+import kotlinx.coroutines.runBlocking
+import kotlinx.coroutines.test.StandardTestDispatcher
+import kotlinx.coroutines.test.setMain
 import org.junit.After
 import org.junit.Before
 import org.junit.Rule
@@ -73,6 +83,7 @@
  */
 @SmallTest
 @RunWith(AndroidTestingRunner::class)
+@ExperimentalCoroutinesApi
 @EnableFlags(FLAG_ENABLE_DESKTOP_WINDOWING_MODE, FLAG_RESPECT_ORIENTATION_CHANGE_FOR_UNRESIZEABLE)
 class DesktopActivityOrientationChangeHandlerTest : ShellTestCase() {
     @JvmField @Rule val setFlagsRule = SetFlagsRule()
@@ -82,16 +93,19 @@
     @Mock lateinit var transitions: Transitions
     @Mock lateinit var resizeTransitionHandler: ToggleResizeDesktopTaskTransitionHandler
     @Mock lateinit var taskStackListener: TaskStackListenerImpl
+    @Mock lateinit var persistentRepository: DesktopPersistentRepository
 
     private lateinit var mockitoSession: StaticMockitoSession
     private lateinit var handler: DesktopActivityOrientationChangeHandler
     private lateinit var shellInit: ShellInit
     private lateinit var taskRepository: DesktopModeTaskRepository
+    private lateinit var testScope: CoroutineScope
     // 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() {
+        Dispatchers.setMain(StandardTestDispatcher())
         mockitoSession =
             mockitoSession()
                 .strictness(Strictness.LENIENT)
@@ -99,10 +113,15 @@
                 .startMocking()
         doReturn(true).`when` { DesktopModeStatus.isDesktopModeSupported(any()) }
 
+        testScope = CoroutineScope(Dispatchers.Unconfined + SupervisorJob())
         shellInit = spy(ShellInit(testExecutor))
-        taskRepository = DesktopModeTaskRepository()
+        taskRepository =
+            DesktopModeTaskRepository(context, shellInit, persistentRepository, testScope)
         whenever(shellTaskOrganizer.getRunningTasks(anyInt())).thenAnswer { runningTasks }
         whenever(transitions.startTransition(anyInt(), any(), isNull())).thenAnswer { Binder() }
+        whenever(runBlocking { persistentRepository.readDesktop(any(), any()) }).thenReturn(
+            Desktop.getDefaultInstance()
+        )
 
         handler = DesktopActivityOrientationChangeHandler(context, shellInit, shellTaskOrganizer,
             taskStackListener, resizeTransitionHandler, taskRepository)
@@ -115,6 +134,7 @@
         mockitoSession.finishMocking()
 
         runningTasks.clear()
+        testScope.cancel()
     }
 
     @Test
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 d3404f7..bc40d89 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
@@ -17,27 +17,70 @@
 package com.android.wm.shell.desktopmode
 
 import android.graphics.Rect
+import android.platform.test.annotations.EnableFlags
 import android.testing.AndroidTestingRunner
+import android.util.ArraySet
 import android.view.Display.DEFAULT_DISPLAY
 import android.view.Display.INVALID_DISPLAY
 import androidx.test.filters.SmallTest
+import com.android.window.flags.Flags.FLAG_ENABLE_DESKTOP_WINDOWING_PERSISTENCE
 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.persistence.Desktop
+import com.android.wm.shell.desktopmode.persistence.DesktopPersistentRepository
+import com.android.wm.shell.sysui.ShellInit
 import com.google.common.truth.Truth.assertThat
 import junit.framework.Assert.fail
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.SupervisorJob
+import kotlinx.coroutines.cancel
+import kotlinx.coroutines.runBlocking
+import kotlinx.coroutines.test.StandardTestDispatcher
+import kotlinx.coroutines.test.runTest
+import kotlinx.coroutines.test.setMain
+import org.junit.After
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
+import org.mockito.Mock
+import org.mockito.Mockito.inOrder
+import org.mockito.Mockito.spy
+import org.mockito.kotlin.any
+import org.mockito.kotlin.never
+import org.mockito.kotlin.verify
+import org.mockito.kotlin.whenever
 
 @SmallTest
 @RunWith(AndroidTestingRunner::class)
+@ExperimentalCoroutinesApi
 class DesktopModeTaskRepositoryTest : ShellTestCase() {
 
     private lateinit var repo: DesktopModeTaskRepository
+    private lateinit var shellInit: ShellInit
+    private lateinit var datastoreScope: CoroutineScope
+
+    @Mock private lateinit var testExecutor: ShellExecutor
+    @Mock private lateinit var persistentRepository: DesktopPersistentRepository
 
     @Before
     fun setUp() {
-        repo = DesktopModeTaskRepository()
+        Dispatchers.setMain(StandardTestDispatcher())
+        datastoreScope = CoroutineScope(Dispatchers.Unconfined + SupervisorJob())
+        shellInit = spy(ShellInit(testExecutor))
+
+        repo = DesktopModeTaskRepository(context, shellInit, persistentRepository, datastoreScope)
+        whenever(runBlocking { persistentRepository.readDesktop(any(), any()) }).thenReturn(
+            Desktop.getDefaultInstance()
+        )
+        shellInit.init()
+    }
+
+    @After
+    fun tearDown() {
+        datastoreScope.cancel()
     }
 
     @Test
@@ -455,6 +498,44 @@
     }
 
     @Test
+    @EnableFlags(FLAG_ENABLE_DESKTOP_WINDOWING_PERSISTENCE)
+    fun addOrMoveFreeformTaskToTop_noTaskExists_persistenceEnabled_addsToTop() =
+        runTest(StandardTestDispatcher()) {
+            repo.addOrMoveFreeformTaskToTop(DEFAULT_DISPLAY, 5)
+            repo.addOrMoveFreeformTaskToTop(DEFAULT_DISPLAY, 6)
+            repo.addOrMoveFreeformTaskToTop(DEFAULT_DISPLAY, 7)
+
+            val tasks = repo.getFreeformTasksInZOrder(DEFAULT_DISPLAY)
+            assertThat(tasks).containsExactly(7, 6, 5).inOrder()
+            inOrder(persistentRepository).run {
+                verify(persistentRepository)
+                    .addOrUpdateDesktop(
+                        DEFAULT_USER_ID,
+                        DEFAULT_DESKTOP_ID,
+                        visibleTasks = ArraySet(),
+                        minimizedTasks = ArraySet(),
+                        freeformTasksInZOrder = arrayListOf(5)
+                    )
+                verify(persistentRepository)
+                    .addOrUpdateDesktop(
+                        DEFAULT_USER_ID,
+                        DEFAULT_DESKTOP_ID,
+                        visibleTasks = ArraySet(),
+                        minimizedTasks = ArraySet(),
+                        freeformTasksInZOrder = arrayListOf(6, 5)
+                    )
+                verify(persistentRepository)
+                    .addOrUpdateDesktop(
+                        DEFAULT_USER_ID,
+                        DEFAULT_DESKTOP_ID,
+                        visibleTasks = ArraySet(),
+                        minimizedTasks = ArraySet(),
+                        freeformTasksInZOrder = arrayListOf(7, 6, 5)
+                    )
+            }
+    }
+
+    @Test
     fun addOrMoveFreeformTaskToTop_alreadyExists_movesToTop() {
         repo.addOrMoveFreeformTaskToTop(DEFAULT_DISPLAY, 5)
         repo.addOrMoveFreeformTaskToTop(DEFAULT_DISPLAY, 6)
@@ -480,6 +561,55 @@
     }
 
     @Test
+    @EnableFlags(FLAG_ENABLE_DESKTOP_WINDOWING_PERSISTENCE)
+    fun minimizeTask_persistenceEnabled_taskIsPersistedAsMinimized() =
+        runTest(StandardTestDispatcher()) {
+            repo.addOrMoveFreeformTaskToTop(DEFAULT_DISPLAY, 5)
+            repo.addOrMoveFreeformTaskToTop(DEFAULT_DISPLAY, 6)
+            repo.addOrMoveFreeformTaskToTop(DEFAULT_DISPLAY, 7)
+
+            repo.minimizeTask(displayId = 0, taskId = 6)
+
+            val tasks = repo.getFreeformTasksInZOrder(DEFAULT_DISPLAY)
+            assertThat(tasks).containsExactly(7, 6, 5).inOrder()
+            assertThat(repo.isMinimizedTask(taskId = 6)).isTrue()
+            inOrder(persistentRepository).run {
+                verify(persistentRepository)
+                    .addOrUpdateDesktop(
+                        DEFAULT_USER_ID,
+                        DEFAULT_DESKTOP_ID,
+                        visibleTasks = ArraySet(),
+                        minimizedTasks = ArraySet(),
+                        freeformTasksInZOrder = arrayListOf(5)
+                    )
+                verify(persistentRepository)
+                    .addOrUpdateDesktop(
+                        DEFAULT_USER_ID,
+                        DEFAULT_DESKTOP_ID,
+                        visibleTasks = ArraySet(),
+                        minimizedTasks = ArraySet(),
+                        freeformTasksInZOrder = arrayListOf(6, 5)
+                    )
+                verify(persistentRepository)
+                    .addOrUpdateDesktop(
+                        DEFAULT_USER_ID,
+                        DEFAULT_DESKTOP_ID,
+                        visibleTasks = ArraySet(),
+                        minimizedTasks = ArraySet(),
+                        freeformTasksInZOrder = arrayListOf(7, 6, 5)
+                    )
+                verify(persistentRepository)
+                    .addOrUpdateDesktop(
+                        DEFAULT_USER_ID,
+                        DEFAULT_DESKTOP_ID,
+                        visibleTasks = ArraySet(),
+                        minimizedTasks = ArraySet(arrayOf(6)),
+                        freeformTasksInZOrder = arrayListOf(7, 6, 5)
+                    )
+            }
+    }
+
+    @Test
     fun addOrMoveFreeformTaskToTop_taskIsUnminimized_noop() {
         repo.addOrMoveFreeformTaskToTop(DEFAULT_DISPLAY, 5)
         repo.addOrMoveFreeformTaskToTop(DEFAULT_DISPLAY, 6)
@@ -503,6 +633,33 @@
     }
 
     @Test
+    @EnableFlags(FLAG_ENABLE_DESKTOP_WINDOWING_PERSISTENCE)
+    fun removeFreeformTask_invalidDisplay_persistenceEnabled_removesTaskFromFreeformTasks() {
+        runTest(StandardTestDispatcher()) {
+            repo.addOrMoveFreeformTaskToTop(DEFAULT_DISPLAY, taskId = 1)
+
+            repo.removeFreeformTask(INVALID_DISPLAY, taskId = 1)
+
+            verify(persistentRepository)
+                .addOrUpdateDesktop(
+                    DEFAULT_USER_ID,
+                    DEFAULT_DESKTOP_ID,
+                    visibleTasks = ArraySet(),
+                    minimizedTasks = ArraySet(),
+                    freeformTasksInZOrder = arrayListOf(1)
+                )
+            verify(persistentRepository)
+                .addOrUpdateDesktop(
+                    DEFAULT_USER_ID,
+                    DEFAULT_DESKTOP_ID,
+                    visibleTasks = ArraySet(),
+                    minimizedTasks = ArraySet(),
+                    freeformTasksInZOrder = ArrayList()
+                )
+        }
+    }
+
+    @Test
     fun removeFreeformTask_validDisplay_removesTaskFromFreeformTasks() {
         repo.addOrMoveFreeformTaskToTop(DEFAULT_DISPLAY, taskId = 1)
 
@@ -513,6 +670,33 @@
     }
 
     @Test
+    @EnableFlags(FLAG_ENABLE_DESKTOP_WINDOWING_PERSISTENCE)
+    fun removeFreeformTask_validDisplay_persistenceEnabled_removesTaskFromFreeformTasks() {
+        runTest(StandardTestDispatcher()) {
+            repo.addOrMoveFreeformTaskToTop(DEFAULT_DISPLAY, taskId = 1)
+
+            repo.removeFreeformTask(DEFAULT_DISPLAY, taskId = 1)
+
+            verify(persistentRepository)
+                .addOrUpdateDesktop(
+                    DEFAULT_USER_ID,
+                    DEFAULT_DESKTOP_ID,
+                    visibleTasks = ArraySet(),
+                    minimizedTasks = ArraySet(),
+                    freeformTasksInZOrder = arrayListOf(1)
+                )
+            verify(persistentRepository)
+                .addOrUpdateDesktop(
+                    DEFAULT_USER_ID,
+                    DEFAULT_DESKTOP_ID,
+                    visibleTasks = ArraySet(),
+                    minimizedTasks = ArraySet(),
+                    freeformTasksInZOrder = ArrayList()
+                )
+        }
+    }
+
+    @Test
     fun removeFreeformTask_validDisplay_differentDisplay_doesNotRemovesTask() {
         repo.addOrMoveFreeformTaskToTop(DEFAULT_DISPLAY, taskId = 1)
 
@@ -523,6 +707,33 @@
     }
 
     @Test
+    @EnableFlags(FLAG_ENABLE_DESKTOP_WINDOWING_PERSISTENCE)
+    fun removeFreeformTask_validDisplayButDifferentDisplay_persistenceEnabled_doesNotRemoveTask() {
+        runTest(StandardTestDispatcher()) {
+            repo.addOrMoveFreeformTaskToTop(DEFAULT_DISPLAY, taskId = 1)
+
+            repo.removeFreeformTask(SECOND_DISPLAY, taskId = 1)
+
+            verify(persistentRepository)
+                .addOrUpdateDesktop(
+                    DEFAULT_USER_ID,
+                    DEFAULT_DESKTOP_ID,
+                    visibleTasks = ArraySet(),
+                    minimizedTasks = ArraySet(),
+                    freeformTasksInZOrder = arrayListOf(1)
+                )
+            verify(persistentRepository, never())
+                .addOrUpdateDesktop(
+                    DEFAULT_USER_ID,
+                    DEFAULT_DESKTOP_ID,
+                    visibleTasks = ArraySet(),
+                    minimizedTasks = ArraySet(),
+                    freeformTasksInZOrder = ArrayList()
+                )
+        }
+    }
+
+    @Test
     fun removeFreeformTask_removesTaskBoundsBeforeMaximize() {
         val taskId = 1
         repo.addActiveTask(THIRD_DISPLAY, taskId)
@@ -709,5 +920,7 @@
     companion object {
         const val SECOND_DISPLAY = 1
         const val THIRD_DISPLAY = 345
+        private const val DEFAULT_USER_ID = 1000
+        private const val DEFAULT_DESKTOP_ID = 0
     }
 }
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksControllerTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksControllerTest.kt
index 8f20841..ee54520 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksControllerTest.kt
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksControllerTest.kt
@@ -93,6 +93,8 @@
 import com.android.wm.shell.desktopmode.DesktopTestHelpers.Companion.createFullscreenTask
 import com.android.wm.shell.desktopmode.DesktopTestHelpers.Companion.createHomeTask
 import com.android.wm.shell.desktopmode.DesktopTestHelpers.Companion.createSplitScreenTask
+import com.android.wm.shell.desktopmode.persistence.Desktop
+import com.android.wm.shell.desktopmode.persistence.DesktopPersistentRepository
 import com.android.wm.shell.draganddrop.DragAndDropController
 import com.android.wm.shell.recents.RecentTasksController
 import com.android.wm.shell.recents.RecentsTransitionHandler
@@ -117,6 +119,14 @@
 import junit.framework.Assert.assertTrue
 import kotlin.test.assertNotNull
 import kotlin.test.assertNull
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.SupervisorJob
+import kotlinx.coroutines.cancel
+import kotlinx.coroutines.runBlocking
+import kotlinx.coroutines.test.StandardTestDispatcher
+import kotlinx.coroutines.test.setMain
 import org.junit.After
 import org.junit.Assume.assumeTrue
 import org.junit.Before
@@ -148,6 +158,7 @@
  */
 @SmallTest
 @RunWith(AndroidTestingRunner::class)
+@ExperimentalCoroutinesApi
 @EnableFlags(FLAG_ENABLE_DESKTOP_WINDOWING_MODE)
 class DesktopTasksControllerTest : ShellTestCase() {
 
@@ -183,6 +194,7 @@
   @Mock private lateinit var mockSurface: SurfaceControl
   @Mock private lateinit var taskbarDesktopTaskListener: TaskbarDesktopTaskListener
   @Mock private lateinit var mockHandler: Handler
+  @Mock lateinit var persistentRepository: DesktopPersistentRepository
 
   private lateinit var mockitoSession: StaticMockitoSession
   private lateinit var controller: DesktopTasksController
@@ -190,6 +202,7 @@
   private lateinit var taskRepository: DesktopModeTaskRepository
   private lateinit var desktopTasksLimiter: DesktopTasksLimiter
   private lateinit var recentsTransitionStateListener: RecentsTransitionStateListener
+  private lateinit var testScope: CoroutineScope
 
   private val shellExecutor = TestShellExecutor()
 
@@ -207,6 +220,7 @@
 
   @Before
   fun setUp() {
+    Dispatchers.setMain(StandardTestDispatcher())
     mockitoSession =
         mockitoSession()
             .strictness(Strictness.LENIENT)
@@ -214,8 +228,9 @@
             .startMocking()
     doReturn(true).`when` { DesktopModeStatus.isDesktopModeSupported(any()) }
 
+    testScope = CoroutineScope(Dispatchers.Unconfined + SupervisorJob())
     shellInit = spy(ShellInit(testExecutor))
-    taskRepository = DesktopModeTaskRepository()
+    taskRepository = DesktopModeTaskRepository(context, shellInit, persistentRepository, testScope)
     desktopTasksLimiter =
         DesktopTasksLimiter(
             transitions,
@@ -233,6 +248,9 @@
     whenever(displayLayout.getStableBounds(any())).thenAnswer { i ->
       (i.arguments.first() as Rect).set(STABLE_BOUNDS)
     }
+    whenever(runBlocking { persistentRepository.readDesktop(any(), any()) }).thenReturn(
+      Desktop.getDefaultInstance()
+    )
 
     val tda = DisplayAreaInfo(MockToken().token(), DEFAULT_DISPLAY, 0)
     tda.configuration.windowConfiguration.windowingMode = WINDOWING_MODE_FULLSCREEN
@@ -287,6 +305,7 @@
     mockitoSession.finishMocking()
 
     runningTasks.clear()
+    testScope.cancel()
   }
 
   @Test
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksLimiterTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksLimiterTest.kt
index 61d03ca..045e077 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksLimiterTest.kt
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksLimiterTest.kt
@@ -35,13 +35,23 @@
 import com.android.internal.jank.InteractionJankMonitor
 import com.android.wm.shell.ShellTaskOrganizer
 import com.android.wm.shell.ShellTestCase
+import com.android.wm.shell.common.ShellExecutor
 import com.android.wm.shell.desktopmode.DesktopTestHelpers.Companion.createFreeformTask
+import com.android.wm.shell.desktopmode.persistence.DesktopPersistentRepository
 import com.android.wm.shell.shared.desktopmode.DesktopModeStatus
+import com.android.wm.shell.sysui.ShellInit
 import com.android.wm.shell.transition.TransitionInfoBuilder
 import com.android.wm.shell.transition.Transitions
 import com.android.wm.shell.util.StubTransaction
 import com.google.common.truth.Truth.assertThat
 import kotlin.test.assertFailsWith
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.SupervisorJob
+import kotlinx.coroutines.cancel
+import kotlinx.coroutines.test.StandardTestDispatcher
+import kotlinx.coroutines.test.setMain
 import org.junit.After
 import org.junit.Before
 import org.junit.Rule
@@ -49,6 +59,7 @@
 import org.junit.runner.RunWith
 import org.mockito.Mock
 import org.mockito.Mockito.any
+import org.mockito.Mockito.spy
 import org.mockito.Mockito.`when`
 import org.mockito.kotlin.eq
 import org.mockito.kotlin.verify
@@ -62,6 +73,7 @@
  */
 @SmallTest
 @RunWith(AndroidTestingRunner::class)
+@ExperimentalCoroutinesApi
 class DesktopTasksLimiterTest : ShellTestCase() {
 
     @JvmField
@@ -72,19 +84,26 @@
     @Mock lateinit var transitions: Transitions
     @Mock lateinit var interactionJankMonitor: InteractionJankMonitor
     @Mock lateinit var handler: Handler
+    @Mock lateinit var testExecutor: ShellExecutor
+    @Mock lateinit var persistentRepository: DesktopPersistentRepository
 
     private lateinit var mockitoSession: StaticMockitoSession
     private lateinit var desktopTasksLimiter: DesktopTasksLimiter
     private lateinit var desktopTaskRepo: DesktopModeTaskRepository
+    private lateinit var shellInit: ShellInit
+    private lateinit var testScope: CoroutineScope
 
     @Before
     fun setUp() {
         mockitoSession = ExtendedMockito.mockitoSession().strictness(Strictness.LENIENT)
                 .spyStatic(DesktopModeStatus::class.java).startMocking()
         doReturn(true).`when`{ DesktopModeStatus.canEnterDesktopMode(any()) }
+        shellInit = spy(ShellInit(testExecutor))
+        Dispatchers.setMain(StandardTestDispatcher())
+        testScope = CoroutineScope(Dispatchers.Unconfined + SupervisorJob())
 
-        desktopTaskRepo = DesktopModeTaskRepository()
-
+        desktopTaskRepo =
+            DesktopModeTaskRepository(context, shellInit, persistentRepository, testScope)
         desktopTasksLimiter =
             DesktopTasksLimiter(transitions, desktopTaskRepo, shellTaskOrganizer, MAX_TASK_LIMIT,
                 interactionJankMonitor, mContext, handler)
@@ -93,6 +112,7 @@
     @After
     fun tearDown() {
         mockitoSession.finishMocking()
+        testScope.cancel()
     }
 
     @Test
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/persistence/DesktopPersistentRepositoryTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/persistence/DesktopPersistentRepositoryTest.kt
new file mode 100644
index 0000000..9b9703f
--- /dev/null
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/persistence/DesktopPersistentRepositoryTest.kt
@@ -0,0 +1,198 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.wm.shell.desktopmode.persistence
+
+import android.content.Context
+import android.platform.test.annotations.EnableFlags
+import android.testing.AndroidTestingRunner
+import android.util.ArraySet
+import android.view.Display.DEFAULT_DISPLAY
+import androidx.datastore.core.DataStore
+import androidx.datastore.core.DataStoreFactory
+import androidx.datastore.dataStoreFile
+import androidx.test.core.app.ApplicationProvider
+import androidx.test.filters.SmallTest
+import androidx.test.platform.app.InstrumentationRegistry
+import com.android.window.flags.Flags.FLAG_ENABLE_DESKTOP_WINDOWING_MODE
+import com.android.window.flags.Flags.FLAG_ENABLE_DESKTOP_WINDOWING_PERSISTENCE
+import com.android.wm.shell.ShellTestCase
+import com.google.common.truth.Truth.assertThat
+import java.io.File
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.SupervisorJob
+import kotlinx.coroutines.cancel
+import kotlinx.coroutines.test.StandardTestDispatcher
+import kotlinx.coroutines.test.runTest
+import kotlinx.coroutines.test.setMain
+import org.junit.After
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+
+@SmallTest
+@RunWith(AndroidTestingRunner::class)
+@ExperimentalCoroutinesApi
+@EnableFlags(FLAG_ENABLE_DESKTOP_WINDOWING_MODE, FLAG_ENABLE_DESKTOP_WINDOWING_PERSISTENCE)
+class DesktopPersistentRepositoryTest : ShellTestCase() {
+    private val testContext: Context = InstrumentationRegistry.getInstrumentation().targetContext
+    private lateinit var testDatastore: DataStore<DesktopPersistentRepositories>
+    private lateinit var datastoreRepository: DesktopPersistentRepository
+    private lateinit var datastoreScope: CoroutineScope
+
+    @Before
+    fun setUp() {
+        Dispatchers.setMain(StandardTestDispatcher())
+        datastoreScope = CoroutineScope(Dispatchers.Unconfined + SupervisorJob())
+        testDatastore =
+            DataStoreFactory.create(
+                serializer =
+                    DesktopPersistentRepository.Companion.DesktopPersistentRepositoriesSerializer,
+                scope = datastoreScope) {
+                    testContext.dataStoreFile(DESKTOP_REPOSITORY_STATES_DATASTORE_TEST_FILE)
+                }
+        datastoreRepository = DesktopPersistentRepository(testDatastore)
+    }
+
+    @After
+    fun tearDown() {
+        File(ApplicationProvider.getApplicationContext<Context>().filesDir, "datastore")
+            .deleteRecursively()
+
+        datastoreScope.cancel()
+    }
+
+    @Test
+    fun readRepository_returnsCorrectDesktop() {
+        runTest(StandardTestDispatcher()) {
+            val task = createDesktopTask(1)
+            val desk = createDesktop(task)
+            val repositoryState =
+                DesktopRepositoryState.newBuilder().putDesktop(DEFAULT_DESKTOP_ID, desk)
+            val DesktopPersistentRepositories =
+                DesktopPersistentRepositories.newBuilder()
+                    .putDesktopRepoByUser(DEFAULT_USER_ID, repositoryState.build())
+                    .build()
+            testDatastore.updateData { DesktopPersistentRepositories }
+
+            val actualDesktop = datastoreRepository.readDesktop(DEFAULT_USER_ID, DEFAULT_DESKTOP_ID)
+
+            assertThat(actualDesktop).isEqualTo(desk)
+        }
+    }
+
+    @Test
+    fun addOrUpdateTask_addNewTaskToDesktop() {
+        runTest(StandardTestDispatcher()) {
+            // Create a basic repository state
+            val task = createDesktopTask(1)
+            val DesktopPersistentRepositories = createRepositoryWithOneDesk(task)
+            testDatastore.updateData { DesktopPersistentRepositories }
+            // Create a new state to be initialized
+            val visibleTasks = ArraySet(listOf(1, 2))
+            val minimizedTasks = ArraySet<Int>()
+            val freeformTasksInZOrder = ArrayList(listOf(2, 1))
+
+            // Update with new state
+            datastoreRepository.addOrUpdateDesktop(
+                visibleTasks = visibleTasks,
+                minimizedTasks = minimizedTasks,
+                freeformTasksInZOrder = freeformTasksInZOrder)
+
+            val actualDesktop = datastoreRepository.readDesktop(DEFAULT_USER_ID, DEFAULT_DESKTOP_ID)
+            assertThat(actualDesktop.tasksByTaskIdMap).hasSize(2)
+            assertThat(actualDesktop.getZOrderedTasks(0)).isEqualTo(2)
+        }
+    }
+
+    @Test
+    fun addOrUpdateTask_changeTaskStateToMinimize_taskStateIsMinimized() {
+        runTest(StandardTestDispatcher()) {
+            val task = createDesktopTask(1)
+            val DesktopPersistentRepositories = createRepositoryWithOneDesk(task)
+            testDatastore.updateData { DesktopPersistentRepositories }
+            // Create a new state to be initialized
+            val visibleTasks = ArraySet(listOf(1))
+            val minimizedTasks = ArraySet(listOf(1))
+            val freeformTasksInZOrder = ArrayList(listOf(1))
+
+            // Update with new state
+            datastoreRepository.addOrUpdateDesktop(
+                visibleTasks = visibleTasks,
+                minimizedTasks = minimizedTasks,
+                freeformTasksInZOrder = freeformTasksInZOrder)
+
+            val actualDesktop = datastoreRepository.readDesktop(DEFAULT_USER_ID, DEFAULT_DESKTOP_ID)
+            assertThat(actualDesktop.tasksByTaskIdMap[task.taskId]?.desktopTaskState)
+                .isEqualTo(DesktopTaskState.MINIMIZED)
+        }
+    }
+
+    @Test
+    fun removeTask_previouslyAddedTaskIsRemoved() {
+        runTest(StandardTestDispatcher()) {
+            val task = createDesktopTask(1)
+            val DesktopPersistentRepositories = createRepositoryWithOneDesk(task)
+            testDatastore.updateData { DesktopPersistentRepositories }
+            // Create a new state to be initialized
+            val visibleTasks = ArraySet<Int>()
+            val minimizedTasks = ArraySet<Int>()
+            val freeformTasksInZOrder = ArrayList<Int>()
+
+            // Update with new state
+            datastoreRepository.addOrUpdateDesktop(
+                visibleTasks = visibleTasks,
+                minimizedTasks = minimizedTasks,
+                freeformTasksInZOrder = freeformTasksInZOrder)
+
+            val actualDesktop = datastoreRepository.readDesktop(DEFAULT_USER_ID, DEFAULT_DESKTOP_ID)
+            assertThat(actualDesktop.tasksByTaskIdMap).isEmpty()
+            assertThat(actualDesktop.zOrderedTasksList).isEmpty()
+        }
+    }
+
+    private companion object {
+        const val DESKTOP_REPOSITORY_STATES_DATASTORE_TEST_FILE = "desktop_repo_test.pb"
+        const val DEFAULT_USER_ID = 1000
+        const val DEFAULT_DESKTOP_ID = 0
+
+        fun createRepositoryWithOneDesk(task: DesktopTask): DesktopPersistentRepositories {
+            val desk = createDesktop(task)
+            val repositoryState =
+                DesktopRepositoryState.newBuilder().putDesktop(DEFAULT_DESKTOP_ID, desk)
+            val DesktopPersistentRepositories =
+                DesktopPersistentRepositories.newBuilder()
+                    .putDesktopRepoByUser(DEFAULT_USER_ID, repositoryState.build())
+                    .build()
+            return DesktopPersistentRepositories
+        }
+
+        fun createDesktop(task: DesktopTask): Desktop? =
+            Desktop.newBuilder()
+                .setDisplayId(DEFAULT_DISPLAY)
+                .addZOrderedTasks(task.taskId)
+                .putTasksByTaskId(task.taskId, task)
+                .build()
+
+        fun createDesktopTask(
+            taskId: Int,
+            state: DesktopTaskState = DesktopTaskState.VISIBLE
+        ): DesktopTask =
+            DesktopTask.newBuilder().setTaskId(taskId).setDesktopTaskState(state).build()
+    }
+}