Add parameter to save task snapshot when it is requested

- In legacy recents animation, the snapshot requested is also saved to
  the task snapshot cache, which triggers a subsequent notification to
  Launcher and updates the cached snapshot sent to AiAi for smart
  select. This change just adds an option to the existing call to allow
  the caller to request that the snapshot is saved to trigger the same
  path.

Bug: 276340272
Test: Open Chrome, scroll a bit, swipe up to overview and trigger smart
      select
Change-Id: I5364f42d9ddc5f3ed67def19c9e9944279591e58

7 files changed, 29 insertions, 11 deletions

diff --git a/core/java/android/app/IActivityTaskManager.aidl b/core/java/android/app/IActivityTaskManager.aidl
index 3249b41da934..4e2b6fa56b17 100644
--- a/core/java/android/app/IActivityTaskManager.aidl
+++ b/core/java/android/app/IActivityTaskManager.aidl
@@ -261,6 +261,9 @@ interface IActivityTaskManager {
     void cancelTaskWindowTransition(int taskId);
 
     /**
+     * Fetches the snapshot for the task with the given id, taking a new snapshot if it is not in
+     * the task snapshot cache and it is requested.
+     *
      * @param taskId the id of the task to retrieve the sAutoapshots for
      * @param isLowResolution if set, if the snapshot needs to be loaded from disk, this will load
      *                          a reduced resolution of it, which is much faster
@@ -272,10 +275,14 @@ interface IActivityTaskManager {
             int taskId, boolean isLowResolution, boolean takeSnapshotIfNeeded);
 
     /**
+     * Requests for a new snapshot to be taken for the task with the given id, storing it in the
+     * task snapshot cache only if requested.
+     *
      * @param taskId the id of the task to take a snapshot of
+     * @param updateCache whether to store the new snapshot in the system's task snapshot cache
      * @return a graphic buffer representing a screenshot of a task
      */
-    android.window.TaskSnapshot takeTaskSnapshot(int taskId);
+    android.window.TaskSnapshot takeTaskSnapshot(int taskId, boolean updateCache);
 
     /**
      * Return the user id of last resumed activity.
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/recents/RecentsTransitionHandler.java b/libs/WindowManager/Shell/src/com/android/wm/shell/recents/RecentsTransitionHandler.java
index 387d39056ffc..a9ad3c90949f 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/recents/RecentsTransitionHandler.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/recents/RecentsTransitionHandler.java
@@ -301,7 +301,7 @@ public class RecentsTransitionHandler implements Transitions.TransitionHandler {
                 try {
                     for (int i = 0; i < mPausingTasks.size(); ++i) {
                         snapshots[i] = ActivityTaskManager.getService().takeTaskSnapshot(
-                                mPausingTasks.get(0).mTaskInfo.taskId);
+                                mPausingTasks.get(0).mTaskInfo.taskId, false /* updateCache */);
                     }
                 } catch (RemoteException e) {
                     taskIds = null;
@@ -671,7 +671,8 @@ public class RecentsTransitionHandler implements Transitions.TransitionHandler {
             try {
                 ProtoLog.v(ShellProtoLogGroup.WM_SHELL_RECENTS_TRANSITION,
                         "[%d] RecentsController.screenshotTask: taskId=%d", mInstanceId, taskId);
-                return ActivityTaskManager.getService().takeTaskSnapshot(taskId);
+                return ActivityTaskManager.getService().takeTaskSnapshot(taskId,
+                        true /* updateCache */);
             } catch (RemoteException e) {
                 Slog.e(TAG, "Failed to screenshot task", e);
             }
diff --git a/packages/SystemUI/shared/src/com/android/systemui/shared/system/RemoteTransitionCompat.java b/packages/SystemUI/shared/src/com/android/systemui/shared/system/RemoteTransitionCompat.java
index 1fbf743836a1..74c325dea15c 100644
--- a/packages/SystemUI/shared/src/com/android/systemui/shared/system/RemoteTransitionCompat.java
+++ b/packages/SystemUI/shared/src/com/android/systemui/shared/system/RemoteTransitionCompat.java
@@ -352,7 +352,8 @@ public class RemoteTransitionCompat {
 
         @Override public TaskSnapshot screenshotTask(int taskId) {
             try {
-                return ActivityTaskManager.getService().takeTaskSnapshot(taskId);
+                return ActivityTaskManager.getService().takeTaskSnapshot(taskId,
+                        true /* updateCache */);
             } catch (RemoteException e) {
                 Log.e(TAG, "Failed to screenshot task", e);
             }
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/ImageCaptureImpl.kt b/packages/SystemUI/src/com/android/systemui/screenshot/ImageCaptureImpl.kt
index 67e9a87e4f18..2e47ab6417ae 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/ImageCaptureImpl.kt
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/ImageCaptureImpl.kt
@@ -48,7 +48,9 @@ open class ImageCaptureImpl @Inject constructor(
     }
 
     override suspend fun captureTask(taskId: Int): Bitmap? {
-        val snapshot = withContext(bgContext) { atmService.takeTaskSnapshot(taskId) } ?: return null
+        val snapshot = withContext(bgContext) {
+            atmService.takeTaskSnapshot(taskId, false /* updateCache */)
+        } ?: return null
         return Bitmap.wrapHardwareBuffer(snapshot.hardwareBuffer, snapshot.colorSpace)
     }
 }
diff --git a/services/core/java/com/android/server/wm/ActivityTaskManagerService.java b/services/core/java/com/android/server/wm/ActivityTaskManagerService.java
index cff65547c673..1360a956dc0b 100644
--- a/services/core/java/com/android/server/wm/ActivityTaskManagerService.java
+++ b/services/core/java/com/android/server/wm/ActivityTaskManagerService.java
@@ -3789,7 +3789,7 @@ public class ActivityTaskManagerService extends IActivityTaskManager.Stub {
             TaskSnapshot taskSnapshot = mWindowManager.mTaskSnapshotController.getSnapshot(taskId,
                     task.mUserId, true /* restoreFromDisk */, isLowResolution);
             if (taskSnapshot == null && takeSnapshotIfNeeded) {
-                taskSnapshot = takeTaskSnapshot(taskId);
+                taskSnapshot = takeTaskSnapshot(taskId, false /* updateCache */);
             }
             return taskSnapshot;
         } finally {
@@ -3798,7 +3798,7 @@ public class ActivityTaskManagerService extends IActivityTaskManager.Stub {
     }
 
     @Override
-    public TaskSnapshot takeTaskSnapshot(int taskId) {
+    public TaskSnapshot takeTaskSnapshot(int taskId, boolean updateCache) {
         mAmInternal.enforceCallingPermission(READ_FRAME_BUFFER, "takeTaskSnapshot()");
         final long ident = Binder.clearCallingIdentity();
         try {
@@ -3809,8 +3809,13 @@ public class ActivityTaskManagerService extends IActivityTaskManager.Stub {
                     Slog.w(TAG, "takeTaskSnapshot: taskId=" + taskId + " not found or not visible");
                     return null;
                 }
-                return mWindowManager.mTaskSnapshotController.captureSnapshot(
-                        task, true /* snapshotHome */);
+                if (updateCache) {
+                    return mWindowManager.mTaskSnapshotController.recordSnapshot(task,
+                            true /* snapshotHome */);
+                } else {
+                    return mWindowManager.mTaskSnapshotController.captureSnapshot(task,
+                            true /* snapshotHome */);
+                }
             }
         } finally {
             Binder.restoreCallingIdentity(ident);
diff --git a/services/core/java/com/android/server/wm/InputMonitor.java b/services/core/java/com/android/server/wm/InputMonitor.java
index b67ccd2a7672..cea886ff6e87 100644
--- a/services/core/java/com/android/server/wm/InputMonitor.java
+++ b/services/core/java/com/android/server/wm/InputMonitor.java
@@ -440,7 +440,8 @@ final class InputMonitor {
                         if (app != null) {
                             mDisplayContent.removeImeSurfaceImmediately();
                             if (app.getTask() != null) {
-                                mDisplayContent.mAtmService.takeTaskSnapshot(app.getTask().mTaskId);
+                                mDisplayContent.mAtmService.takeTaskSnapshot(app.getTask().mTaskId,
+                                        true /* updateCache */);
                             }
                         }
                     } else {
diff --git a/services/core/java/com/android/server/wm/TaskSnapshotController.java b/services/core/java/com/android/server/wm/TaskSnapshotController.java
index 4d0bff961e60..20ef0171358b 100644
--- a/services/core/java/com/android/server/wm/TaskSnapshotController.java
+++ b/services/core/java/com/android/server/wm/TaskSnapshotController.java
@@ -171,13 +171,14 @@ class TaskSnapshotController extends AbsAppSnapshotController<Task, TaskSnapshot
         snapshotTasks(tasks, false /* allowSnapshotHome */);
     }
 
-    void recordSnapshot(Task task, boolean allowSnapshotHome) {
+    TaskSnapshot recordSnapshot(Task task, boolean allowSnapshotHome) {
         final boolean snapshotHome = allowSnapshotHome && task.isActivityTypeHome();
         final TaskSnapshot snapshot = recordSnapshotInner(task, allowSnapshotHome);
         if (!snapshotHome && snapshot != null) {
             mPersister.persistSnapshot(task.mTaskId, task.mUserId, snapshot);
             task.onSnapshotChanged(snapshot);
         }
+        return snapshot;
     }
 
     private void snapshotTasks(ArraySet<Task> tasks, boolean allowSnapshotHome) {