Merge "PSS: Make sure recent task thumbnails are up to date" into main

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 =