diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopRepository.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopRepository.kt
index 443e417..85a3126 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopRepository.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopRepository.kt
@@ -30,7 +30,6 @@
 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
@@ -124,7 +123,8 @@
         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 desktop = persistentRepository.readDesktop() ?: return@launch
+
             val maxTasks =
                 DesktopModeStatus.getMaxTaskLimit(context).takeIf { it > 0 }
                     ?: desktop.zOrderedTasksCount
@@ -132,13 +132,11 @@
             desktop.zOrderedTasksList
                 // Reverse it so we initialize the repo from bottom to top.
                 .reversed()
-                .map { taskId ->
-                    desktop.tasksByTaskIdMap.getOrDefault(
-                        taskId,
-                        DesktopTask.getDefaultInstance()
-                    )
+                .mapNotNull { taskId ->
+                    desktop.tasksByTaskIdMap[taskId]?.takeIf {
+                        it.desktopTaskState == DesktopTaskState.VISIBLE
+                    }
                 }
-                .filter { task -> task.desktopTaskState == DesktopTaskState.VISIBLE }
                 .take(maxTasks)
                 .forEach { task ->
                     addOrMoveFreeformTaskToTop(desktop.displayId, task.taskId)
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
index 3f41d7c..2d11e02 100644
--- 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
@@ -73,15 +73,14 @@
      */
     private suspend fun getDesktopRepositoryState(
         userId: Int = DEFAULT_USER_ID
-    ): DesktopRepositoryState =
+    ): DesktopRepositoryState? =
         try {
             dataStoreFlow
                 .first()
-                .desktopRepoByUserMap
-                .getOrDefault(userId, DesktopRepositoryState.getDefaultInstance())
+                .desktopRepoByUserMap[userId]
         } catch (e: Exception) {
             Log.e(TAG, "Unable to read from datastore", e)
-            DesktopRepositoryState.getDefaultInstance()
+            null
         }
 
     /**
@@ -91,13 +90,13 @@
     suspend fun readDesktop(
         userId: Int = DEFAULT_USER_ID,
         desktopId: Int = DEFAULT_DESKTOP_ID,
-    ): Desktop =
+    ): Desktop? =
         try {
             val repository = getDesktopRepositoryState(userId)
-            repository.getDesktopOrThrow(desktopId)
+            repository?.getDesktopOrThrow(desktopId)
         } catch (e: Exception) {
             Log.e(TAG, "Unable to get desktop info from persistent repository", e)
-            Desktop.getDefaultInstance()
+            null
         }
 
     /** Adds or updates a desktop stored in the datastore */
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
index 9b9703f..8495580 100644
--- 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
@@ -115,8 +115,8 @@
                 freeformTasksInZOrder = freeformTasksInZOrder)
 
             val actualDesktop = datastoreRepository.readDesktop(DEFAULT_USER_ID, DEFAULT_DESKTOP_ID)
-            assertThat(actualDesktop.tasksByTaskIdMap).hasSize(2)
-            assertThat(actualDesktop.getZOrderedTasks(0)).isEqualTo(2)
+            assertThat(actualDesktop?.tasksByTaskIdMap).hasSize(2)
+            assertThat(actualDesktop?.getZOrderedTasks(0)).isEqualTo(2)
         }
     }
 
@@ -138,7 +138,7 @@
                 freeformTasksInZOrder = freeformTasksInZOrder)
 
             val actualDesktop = datastoreRepository.readDesktop(DEFAULT_USER_ID, DEFAULT_DESKTOP_ID)
-            assertThat(actualDesktop.tasksByTaskIdMap[task.taskId]?.desktopTaskState)
+            assertThat(actualDesktop?.tasksByTaskIdMap?.get(task.taskId)?.desktopTaskState)
                 .isEqualTo(DesktopTaskState.MINIMIZED)
         }
     }
@@ -161,8 +161,8 @@
                 freeformTasksInZOrder = freeformTasksInZOrder)
 
             val actualDesktop = datastoreRepository.readDesktop(DEFAULT_USER_ID, DEFAULT_DESKTOP_ID)
-            assertThat(actualDesktop.tasksByTaskIdMap).isEmpty()
-            assertThat(actualDesktop.zOrderedTasksList).isEmpty()
+            assertThat(actualDesktop?.tasksByTaskIdMap).isEmpty()
+            assertThat(actualDesktop?.zOrderedTasksList).isEmpty()
         }
     }