PSS: Make sure recent task thumbnails are up to date

Task thumbnails are usually only updated when opening them, closing
them, or when going to overview.
This was causing PSS app selector to show an out of date thumbnail
for the last foreground task.
The solution is to make sure to take a new thumbnail when the app
selector starts.

Fixes: 304442979
Test: Manual
Test: MediaProjectionAppSelectorControllerTest
Test: ShellRecentTaskListProviderTest
Test: MediaProjectionAppSelectorScreenshotTest
Change-Id: Id82c0cf47a9be7a610095e6c669fc5aa634d2327

8 files changed, 260 insertions, 31 deletions

diff --git a/packages/SystemUI/shared/src/com/android/systemui/shared/system/ActivityManagerWrapper.java b/packages/SystemUI/shared/src/com/android/systemui/shared/system/ActivityManagerWrapper.java
index 80040a384b9d..631423e4b7fc 100644
--- a/packages/SystemUI/shared/src/com/android/systemui/shared/system/ActivityManagerWrapper.java
+++ b/packages/SystemUI/shared/src/com/android/systemui/shared/system/ActivityManagerWrapper.java
@@ -155,6 +155,26 @@ public class ActivityManagerWrapper {
         }
     }
 
+
+    /**
+     * Requests for a new snapshot to be taken for the given task, stores it in the cache, and
+     * returns a {@link ThumbnailData} with the result.
+     */
+    @NonNull
+    public ThumbnailData takeTaskThumbnail(int taskId) {
+        TaskSnapshot snapshot = null;
+        try {
+            snapshot = getService().takeTaskSnapshot(taskId, /* updateCache= */ true);
+        } catch (RemoteException e) {
+            Log.w(TAG, "Failed to take task snapshot", e);
+        }
+        if (snapshot != null) {
+            return new ThumbnailData(snapshot);
+        } else {
+            return new ThumbnailData();
+        }
+    }
+
     /**
      * Removes the outdated snapshot of home task.
      *
diff --git a/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/MediaProjectionAppSelectorController.kt b/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/MediaProjectionAppSelectorController.kt
index e61650fbb163..fced117a8132 100644
--- a/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/MediaProjectionAppSelectorController.kt
+++ b/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/MediaProjectionAppSelectorController.kt
@@ -20,10 +20,15 @@ import android.content.ComponentName
 import android.os.UserHandle
 import com.android.systemui.mediaprojection.appselector.data.RecentTask
 import com.android.systemui.mediaprojection.appselector.data.RecentTaskListProvider
+import com.android.systemui.mediaprojection.appselector.data.RecentTaskThumbnailLoader
 import com.android.systemui.mediaprojection.devicepolicy.ScreenCaptureDevicePolicyResolver
+import com.android.systemui.shared.recents.model.ThumbnailData
 import javax.inject.Inject
 import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.Deferred
+import kotlinx.coroutines.async
 import kotlinx.coroutines.cancel
+import kotlinx.coroutines.coroutineScope
 import kotlinx.coroutines.launch
 
 @MediaProjectionAppSelectorScope
@@ -36,7 +41,8 @@ constructor(
     @HostUserHandle private val hostUserHandle: UserHandle,
     @MediaProjectionAppSelector private val scope: CoroutineScope,
     @MediaProjectionAppSelector private val appSelectorComponentName: ComponentName,
-    @MediaProjectionAppSelector private val callerPackageName: String?
+    @MediaProjectionAppSelector private val callerPackageName: String?,
+    private val thumbnailLoader: RecentTaskThumbnailLoader,
 ) {
 
     fun init() {
@@ -46,6 +52,11 @@ constructor(
             val tasks =
                 recentTasks.filterDevicePolicyRestrictedTasks().filterAppSelector().sortedTasks()
 
+            // Thumbnails are not fresh for the foreground task(s). They are only refreshed at
+            // launch, going to home, or going to overview.
+            // For this reason, we need to refresh them here.
+            refreshForegroundTaskThumbnails(tasks)
+
             view.bind(tasks)
         }
     }
@@ -54,6 +65,16 @@ constructor(
         scope.cancel()
     }
 
+    private suspend fun refreshForegroundTaskThumbnails(tasks: List<RecentTask>) {
+        coroutineScope {
+            val thumbnails: List<Deferred<ThumbnailData?>> =
+                tasks
+                    .filter { it.isForegroundTask }
+                    .map { async { thumbnailLoader.captureThumbnail(it.taskId) } }
+            thumbnails.forEach { thumbnail -> thumbnail.await() }
+        }
+    }
+
     /** Removes all recent tasks that should be blocked according to the policy */
     private fun List<RecentTask>.filterDevicePolicyRestrictedTasks(): List<RecentTask> = filter {
         devicePolicyResolver.isScreenCaptureAllowed(
diff --git a/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/data/RecentTask.kt b/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/data/RecentTask.kt
index 41e22860d0ad..a9e6c53b3bcd 100644
--- a/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/data/RecentTask.kt
+++ b/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/data/RecentTask.kt
@@ -25,5 +25,6 @@ data class RecentTask(
     @UserIdInt val userId: Int,
     val topActivityComponent: ComponentName?,
     val baseIntentComponent: ComponentName?,
-    @ColorInt val colorBackground: Int?
+    @ColorInt val colorBackground: Int?,
+    val isForegroundTask: Boolean,
 )
diff --git a/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/data/RecentTaskListProvider.kt b/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/data/RecentTaskListProvider.kt
index 01398cf81314..aa4c4e55c718 100644
--- a/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/data/RecentTaskListProvider.kt
+++ b/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/data/RecentTaskListProvider.kt
@@ -48,9 +48,14 @@ constructor(
 
     override suspend fun loadRecentTasks(): List<RecentTask> =
         withContext(coroutineDispatcher) {
-            val rawRecentTasks: List<GroupedRecentTaskInfo> = recents?.getTasks() ?: emptyList()
-
-            rawRecentTasks
+            val groupedTasks: List<GroupedRecentTaskInfo> = recents?.getTasks() ?: emptyList()
+            // Note: the returned task list is from the most-recent to least-recent order.
+            // The last foreground task is at index 1, because at index 0 will be our app selector.
+            val foregroundGroup = groupedTasks.elementAtOrNull(1)
+            val foregroundTaskId1 = foregroundGroup?.taskInfo1?.taskId
+            val foregroundTaskId2 = foregroundGroup?.taskInfo2?.taskId
+            val foregroundTaskIds = listOfNotNull(foregroundTaskId1, foregroundTaskId2)
+            groupedTasks
                 .flatMap { listOfNotNull(it.taskInfo1, it.taskInfo2) }
                 .map {
                     RecentTask(
@@ -58,7 +63,8 @@ constructor(
                         it.userId,
                         it.topActivity,
                         it.baseIntent?.component,
-                        it.taskDescription?.backgroundColor
+                        it.taskDescription?.backgroundColor,
+                        isForegroundTask = it.taskId in foregroundTaskIds
                     )
                 }
         }
diff --git a/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/data/RecentTaskThumbnailLoader.kt b/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/data/RecentTaskThumbnailLoader.kt
index 47faaed10302..ccf272cbd3c2 100644
--- a/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/data/RecentTaskThumbnailLoader.kt
+++ b/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/data/RecentTaskThumbnailLoader.kt
@@ -25,6 +25,8 @@ import kotlinx.coroutines.withContext
 
 interface RecentTaskThumbnailLoader {
     suspend fun loadThumbnail(taskId: Int): ThumbnailData?
+
+    suspend fun captureThumbnail(taskId: Int): ThumbnailData?
 }
 
 class ActivityTaskManagerThumbnailLoader
@@ -36,8 +38,13 @@ constructor(
 
     override suspend fun loadThumbnail(taskId: Int): ThumbnailData? =
         withContext(coroutineDispatcher) {
-            val thumbnailData =
-                activityManager.getTaskThumbnail(taskId, /* isLowResolution= */ false)
-            if (thumbnailData.thumbnail == null) null else thumbnailData
+            activityManager.getTaskThumbnail(taskId, /* isLowResolution= */ false).takeIf {
+                it.thumbnail != null
+            }
+        }
+
+    override suspend fun captureThumbnail(taskId: Int): ThumbnailData? =
+        withContext(coroutineDispatcher) {
+            activityManager.takeTaskThumbnail(taskId).takeIf { it.thumbnail != null }
         }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/mediaprojection/appselector/MediaProjectionAppSelectorControllerTest.kt b/packages/SystemUI/tests/src/com/android/systemui/mediaprojection/appselector/MediaProjectionAppSelectorControllerTest.kt
index f25cd24dfcb0..34360d2ddd5c 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/mediaprojection/appselector/MediaProjectionAppSelectorControllerTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/mediaprojection/appselector/MediaProjectionAppSelectorControllerTest.kt
@@ -7,12 +7,16 @@ import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.mediaprojection.appselector.data.RecentTask
 import com.android.systemui.mediaprojection.appselector.data.RecentTaskListProvider
+import com.android.systemui.mediaprojection.appselector.data.RecentTaskThumbnailLoader
 import com.android.systemui.mediaprojection.devicepolicy.ScreenCaptureDevicePolicyResolver
+import com.android.systemui.shared.recents.model.ThumbnailData
 import com.android.systemui.util.mockito.any
 import com.android.systemui.util.mockito.mock
 import com.android.systemui.util.mockito.whenever
+import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.test.runTest
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -34,6 +38,8 @@ class MediaProjectionAppSelectorControllerTest : SysuiTestCase() {
     private val view: MediaProjectionAppSelectorView = mock()
     private val policyResolver: ScreenCaptureDevicePolicyResolver = mock()
 
+    private val thumbnailLoader = FakeThumbnailLoader()
+
     private val controller =
         MediaProjectionAppSelectorController(
             taskListProvider,
@@ -42,7 +48,8 @@ class MediaProjectionAppSelectorControllerTest : SysuiTestCase() {
             personalUserHandle,
             scope,
             appSelectorComponentName,
-            callerPackageName
+            callerPackageName,
+            thumbnailLoader,
         )
 
     @Before
@@ -69,6 +76,22 @@ class MediaProjectionAppSelectorControllerTest : SysuiTestCase() {
     }
 
     @Test
+    fun init_refreshesThumbnailsOfForegroundTasks() = runTest {
+        val tasks =
+            listOf(
+                createRecentTask(taskId = 1, isForegroundTask = false),
+                createRecentTask(taskId = 2, isForegroundTask = true),
+                createRecentTask(taskId = 3, isForegroundTask = true),
+                createRecentTask(taskId = 4, isForegroundTask = false),
+            )
+        taskListProvider.tasks = tasks
+
+        controller.init()
+
+        assertThat(thumbnailLoader.capturedTaskIds).containsExactly(2, 3)
+    }
+
+    @Test
     fun initMultipleRecentTasksWithoutAppSelectorTask_bindsListInTheSameOrder() {
         val tasks =
             listOf(
@@ -188,14 +211,16 @@ class MediaProjectionAppSelectorControllerTest : SysuiTestCase() {
     private fun createRecentTask(
         taskId: Int,
         topActivityComponent: ComponentName? = null,
-        userId: Int = personalUserHandle.identifier
+        userId: Int = personalUserHandle.identifier,
+        isForegroundTask: Boolean = false
     ): RecentTask {
         return RecentTask(
             taskId = taskId,
             topActivityComponent = topActivityComponent,
             baseIntentComponent = ComponentName("com", "Test"),
             userId = userId,
-            colorBackground = 0
+            colorBackground = 0,
+            isForegroundTask = isForegroundTask,
         )
     }
 
@@ -205,4 +230,18 @@ class MediaProjectionAppSelectorControllerTest : SysuiTestCase() {
 
         override suspend fun loadRecentTasks(): List<RecentTask> = tasks
     }
+
+    private class FakeThumbnailLoader : RecentTaskThumbnailLoader {
+
+        val capturedTaskIds = mutableListOf<Int>()
+
+        override suspend fun loadThumbnail(taskId: Int): ThumbnailData? {
+            return null
+        }
+
+        override suspend fun captureThumbnail(taskId: Int): ThumbnailData? {
+            capturedTaskIds += taskId
+            return null
+        }
+    }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/mediaprojection/appselector/data/ActivityTaskManagerThumbnailLoaderTest.kt b/packages/SystemUI/tests/src/com/android/systemui/mediaprojection/appselector/data/ActivityTaskManagerThumbnailLoaderTest.kt
new file mode 100644
index 000000000000..db275ec190ac
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/mediaprojection/appselector/data/ActivityTaskManagerThumbnailLoaderTest.kt
@@ -0,0 +1,109 @@
+package com.android.systemui.mediaprojection.appselector.data
+
+import android.app.WindowConfiguration
+import android.content.ComponentName
+import android.content.res.Configuration
+import android.graphics.ColorSpace
+import android.graphics.Point
+import android.graphics.Rect
+import android.hardware.HardwareBuffer
+import android.testing.AndroidTestingRunner
+import android.view.Surface
+import android.window.TaskSnapshot
+import androidx.test.filters.SmallTest
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.shared.recents.model.ThumbnailData
+import com.android.systemui.shared.system.ActivityManagerWrapper
+import com.android.systemui.util.mockito.mock
+import com.android.systemui.util.mockito.whenever
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.test.TestScope
+import kotlinx.coroutines.test.UnconfinedTestDispatcher
+import kotlinx.coroutines.test.runTest
+import org.junit.Test
+import org.junit.runner.RunWith
+
+@RunWith(AndroidTestingRunner::class)
+@SmallTest
+@OptIn(ExperimentalCoroutinesApi::class)
+class ActivityTaskManagerThumbnailLoaderTest : SysuiTestCase() {
+
+    private val dispatcher = UnconfinedTestDispatcher()
+    private val testScope = TestScope(dispatcher)
+    private val activityManager = mock<ActivityManagerWrapper>()
+    private val loader = ActivityTaskManagerThumbnailLoader(dispatcher, activityManager)
+
+    @Test
+    fun loadThumbnail_emptyThumbnail_returnsNull() =
+        testScope.runTest {
+            val taskId = 123
+            val isLowResolution = false
+            val thumbnailData = ThumbnailData()
+            whenever(activityManager.getTaskThumbnail(taskId, isLowResolution))
+                .thenReturn(thumbnailData)
+
+            assertThat(loader.loadThumbnail(taskId)).isNull()
+        }
+
+    @Test
+    fun loadThumbnail_thumbnailAvailable_returnsThumbnailData() =
+        testScope.runTest {
+            val taskId = 123
+            val isLowResolution = false
+            val snapshot = createTaskSnapshot()
+            val thumbnailData = ThumbnailData(snapshot)
+            whenever(activityManager.getTaskThumbnail(taskId, isLowResolution))
+                .thenReturn(thumbnailData)
+
+            assertThat(loader.loadThumbnail(taskId)).isEqualTo(thumbnailData)
+        }
+
+    @Test
+    fun captureThumbnail_emptyThumbnail_returnsNull() =
+        testScope.runTest {
+            val taskId = 321
+            val emptyThumbnailData = ThumbnailData()
+
+            whenever(activityManager.takeTaskThumbnail(taskId)).thenReturn(emptyThumbnailData)
+
+            assertThat(loader.captureThumbnail(taskId)).isNull()
+        }
+
+    @Test
+    fun captureThumbnail_thumbnailAvailable_returnsThumbnailData() =
+        testScope.runTest {
+            val taskId = 321
+            val thumbnailData = ThumbnailData(createTaskSnapshot())
+
+            whenever(activityManager.takeTaskThumbnail(taskId)).thenReturn(thumbnailData)
+
+            assertThat(loader.captureThumbnail(taskId)).isEqualTo(thumbnailData)
+        }
+
+    private fun createTaskSnapshot() =
+        TaskSnapshot(
+            /* id= */ 123,
+            /* captureTime= */ 0,
+            /* topActivityComponent= */ ComponentName("package", "class"),
+            /* snapshot= */ HardwareBuffer.create(
+                /* width= */ 100,
+                /* height= */ 100,
+                HardwareBuffer.RGBA_8888,
+                /* layers= */ 1,
+                /* usage= */ HardwareBuffer.USAGE_CPU_READ_OFTEN
+            ),
+            ColorSpace.get(ColorSpace.Named.SRGB),
+            Configuration.ORIENTATION_PORTRAIT,
+            Surface.ROTATION_0,
+            /* taskSize= */ Point(100, 100),
+            /* contentInsets= */ Rect(),
+            /* letterboxInsets= */ Rect(),
+            /* isLowResolution= */ false,
+            /* isRealSnapshot= */ true,
+            WindowConfiguration.WINDOWING_MODE_FULLSCREEN,
+            /* appearance= */ 0,
+            /* isTranslucent= */ false,
+            /* hasImeSurface= */ false
+        )
+}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/mediaprojection/appselector/data/ShellRecentTaskListProviderTest.kt b/packages/SystemUI/tests/src/com/android/systemui/mediaprojection/appselector/data/ShellRecentTaskListProviderTest.kt
index d35a21236ae8..2c7ee56e9408 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/mediaprojection/appselector/data/ShellRecentTaskListProviderTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/mediaprojection/appselector/data/ShellRecentTaskListProviderTest.kt
@@ -11,7 +11,7 @@ import com.android.systemui.util.mockito.whenever
 import com.android.wm.shell.recents.RecentTasks
 import com.android.wm.shell.util.GroupedRecentTaskInfo
 import com.google.common.truth.Truth.assertThat
-import java.util.*
+import java.util.Optional
 import java.util.function.Consumer
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.runBlocking
@@ -52,12 +52,7 @@ class ShellRecentTaskListProviderTest : SysuiTestCase() {
 
         val result = runBlocking { recentTaskListProvider.loadRecentTasks() }
 
-        assertThat(result)
-            .containsExactly(
-                createRecentTask(taskId = 1),
-                createRecentTask(taskId = 2),
-                createRecentTask(taskId = 3),
-            )
+        assertThat(result.map { it.taskId }).containsExactly(1, 2, 3).inOrder()
     }
 
     @Test
@@ -66,8 +61,7 @@ class ShellRecentTaskListProviderTest : SysuiTestCase() {
 
         val result = runBlocking { recentTaskListProvider.loadRecentTasks() }
 
-        assertThat(result)
-            .containsExactly(createRecentTask(taskId = 1), createRecentTask(taskId = 2))
+        assertThat(result.map { it.taskId }).containsExactly(1, 2).inOrder()
     }
 
     @Test
@@ -81,15 +75,46 @@ class ShellRecentTaskListProviderTest : SysuiTestCase() {
 
         val result = runBlocking { recentTaskListProvider.loadRecentTasks() }
 
-        assertThat(result)
-            .containsExactly(
-                createRecentTask(taskId = 1),
-                createRecentTask(taskId = 2),
-                createRecentTask(taskId = 3),
-                createRecentTask(taskId = 4),
-                createRecentTask(taskId = 5),
-                createRecentTask(taskId = 6),
-            )
+        assertThat(result.map { it.taskId }).containsExactly(1, 2, 3, 4, 5, 6).inOrder()
+    }
+
+    @Test
+    fun loadRecentTasks_singleTask_returnsTaskAsNotForeground() {
+        givenRecentTasks(
+            createSingleTask(taskId = 1),
+        )
+
+        val result = runBlocking { recentTaskListProvider.loadRecentTasks() }
+
+        assertThat(result[0].isForegroundTask).isFalse()
+    }
+
+    @Test
+    fun loadRecentTasks_multipleTasks_returnsSecondTaskAsForegroundTask() {
+        givenRecentTasks(
+            createSingleTask(taskId = 1),
+            createSingleTask(taskId = 2),
+            createSingleTask(taskId = 3),
+        )
+
+        val result = runBlocking { recentTaskListProvider.loadRecentTasks() }
+
+        assertThat(result.map { it.isForegroundTask }).containsExactly(false, true, false).inOrder()
+    }
+
+    @Test
+    fun loadRecentTasks_secondTaskIsGrouped_marksBothGroupedTasksAsForeground() {
+        givenRecentTasks(
+            createSingleTask(taskId = 1),
+            createTaskPair(taskId1 = 2, taskId2 = 3),
+            createSingleTask(taskId = 4),
+        )
+
+        val result = runBlocking { recentTaskListProvider.loadRecentTasks() }
+
+        assertThat(result.map { it.isForegroundTask })
+            .containsExactly(false, true, true, false)
+            .inOrder()
     }
 
     @Suppress("UNCHECKED_CAST")
@@ -106,7 +131,8 @@ class ShellRecentTaskListProviderTest : SysuiTestCase() {
             userId = 0,
             topActivityComponent = null,
             baseIntentComponent = null,
-            colorBackground = null
+            colorBackground = null,
+            isForegroundTask = false,
         )
 
     private fun createSingleTask(taskId: Int): GroupedRecentTaskInfo =