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

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.

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);
             }

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);
             }

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)
     }
 }

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);

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 {

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) {