diff --git a/apct-tests/perftests/core/src/android/wm/RecentsAnimationPerfTest.java b/apct-tests/perftests/core/src/android/wm/RecentsAnimationPerfTest.java
index 4e2b281..2b6f455 100644
--- a/apct-tests/perftests/core/src/android/wm/RecentsAnimationPerfTest.java
+++ b/apct-tests/perftests/core/src/android/wm/RecentsAnimationPerfTest.java
@@ -26,6 +26,7 @@
 import static org.hamcrest.core.Is.is;
 
 import android.app.Activity;
+import android.app.ActivityManager.TaskSnapshot;
 import android.app.ActivityTaskManager;
 import android.app.IActivityTaskManager;
 import android.content.ComponentName;
@@ -216,7 +217,7 @@
             }
 
             @Override
-            public void onAnimationCanceled(boolean deferredWithScreenshot) throws RemoteException {
+            public void onAnimationCanceled(TaskSnapshot taskSnapshot) throws RemoteException {
                 Assume.assumeNoException(
                         new AssertionError("onAnimationCanceled should not be called"));
             }
diff --git a/core/java/android/view/IRecentsAnimationRunner.aidl b/core/java/android/view/IRecentsAnimationRunner.aidl
index 6cda60c..724d96e 100644
--- a/core/java/android/view/IRecentsAnimationRunner.aidl
+++ b/core/java/android/view/IRecentsAnimationRunner.aidl
@@ -16,6 +16,7 @@
 
 package android.view;
 
+import android.app.ActivityManager;
 import android.graphics.Rect;
 import android.view.RemoteAnimationTarget;
 import android.view.IRecentsAnimationController;
@@ -33,16 +34,16 @@
      * wallpaper not drawing in time, or the handler not finishing the animation within a predefined
      * amount of time.
      *
-     * @param deferredWithScreenshot If set to {@code true}, the contents of the task will be
-     *                               replaced with a screenshot, such that the runner's leash is
-     *                               still active. As soon as the runner doesn't need the leash
-     *                               anymore, it must call
-     *                               {@link IRecentsAnimationController#cleanupScreenshot).
+     * @param taskSnapshot If the snapshot is null, the animation will be cancelled and the leash
+     *                     will be inactive immediately. Otherwise, the contents of the task will be
+     *                     replaced with {@param taskSnapshot}, such that the runner's leash is
+     *                     still active. As soon as the runner doesn't need the leash anymore, it
+     *                     must call {@link IRecentsAnimationController#cleanupScreenshot).
      *
      * @see {@link RecentsAnimationController#cleanupScreenshot}
      */
     @UnsupportedAppUsage
-    void onAnimationCanceled(boolean deferredWithScreenshot) = 1;
+    void onAnimationCanceled(in @nullable ActivityManager.TaskSnapshot taskSnapshot) = 1;
 
     /**
      * Called when the system is ready for the handler to start animating all the visible tasks.
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 13fc702..82287873 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
@@ -31,6 +31,7 @@
 import android.app.ActivityManager;
 import android.app.ActivityManager.RecentTaskInfo;
 import android.app.ActivityManager.RunningTaskInfo;
+import android.app.ActivityManager.TaskSnapshot;
 import android.app.ActivityOptions;
 import android.app.ActivityTaskManager;
 import android.app.AppGlobals;
@@ -235,9 +236,9 @@
                     }
 
                     @Override
-                    public void onAnimationCanceled(boolean deferredWithScreenshot) {
+                    public void onAnimationCanceled(TaskSnapshot taskSnapshot) {
                         animationHandler.onAnimationCanceled(
-                                deferredWithScreenshot ? new ThumbnailData() : null);
+                                taskSnapshot != null ? new ThumbnailData(taskSnapshot) : null);
                     }
                 };
             }
diff --git a/services/core/java/com/android/server/wm/RecentsAnimation.java b/services/core/java/com/android/server/wm/RecentsAnimation.java
index 036bef7..4f0332c 100644
--- a/services/core/java/com/android/server/wm/RecentsAnimation.java
+++ b/services/core/java/com/android/server/wm/RecentsAnimation.java
@@ -460,7 +460,7 @@
      */
     static void notifyAnimationCancelBeforeStart(IRecentsAnimationRunner recentsAnimationRunner) {
         try {
-            recentsAnimationRunner.onAnimationCanceled(false /* deferredWithScreenshot */);
+            recentsAnimationRunner.onAnimationCanceled(null /* taskSnapshot */);
         } catch (RemoteException e) {
             Slog.e(TAG, "Failed to cancel recents animation before start", e);
         }
diff --git a/services/core/java/com/android/server/wm/RecentsAnimationController.java b/services/core/java/com/android/server/wm/RecentsAnimationController.java
index 6ea4d58..724a72e 100644
--- a/services/core/java/com/android/server/wm/RecentsAnimationController.java
+++ b/services/core/java/com/android/server/wm/RecentsAnimationController.java
@@ -23,8 +23,8 @@
 import static android.view.RemoteAnimationTarget.MODE_OPENING;
 import static android.view.WindowManager.DOCKED_INVALID;
 import static android.view.WindowManager.INPUT_CONSUMER_RECENTS_ANIMATION;
-
 import static android.view.WindowManager.LayoutParams.TYPE_BASE_APPLICATION;
+
 import static com.android.server.policy.WindowManagerPolicy.FINISH_LAYOUT_REDO_WALLPAPER;
 import static com.android.server.wm.ActivityTaskManagerInternal.APP_TRANSITION_RECENTS_ANIM;
 import static com.android.server.wm.AnimationAdapterProto.REMOTE;
@@ -525,18 +525,23 @@
                 // Screen shot previous task when next task starts transition and notify the runner.
                 // We will actually finish the animation once the runner calls cleanUpScreenshot().
                 final Task task = mPendingAnimations.get(0).mTask;
-                screenshotRecentTask(task, reorderMode, runSynchronously);
+                final TaskSnapshot taskSnapshot = screenshotRecentTask(task, reorderMode,
+                        runSynchronously);
                 try {
-                    mRunner.onAnimationCanceled(true /* deferredWithScreenshot */);
+                    mRunner.onAnimationCanceled(taskSnapshot);
                 } catch (RemoteException e) {
                     Slog.e(TAG, "Failed to cancel recents animation", e);
                 }
+                if (taskSnapshot == null) {
+                    mCallbacks.onAnimationFinished(reorderMode, runSynchronously,
+                            false /* sendUserLeaveHint */);
+                }
             } else {
                 // Otherwise, notify the runner and clean up the animation immediately
                 // Note: In the fallback case, this can trigger multiple onAnimationCancel() calls
                 // to the runner if we this actually triggers cancel twice on the caller
                 try {
-                    mRunner.onAnimationCanceled(false /* deferredWithScreenshot */);
+                    mRunner.onAnimationCanceled(null /* taskSnapshot */);
                 } catch (RemoteException e) {
                     Slog.e(TAG, "Failed to cancel recents animation", e);
                 }
@@ -587,20 +592,32 @@
         return mRequestDeferCancelUntilNextTransition && mCancelDeferredWithScreenshot;
     }
 
-    void screenshotRecentTask(Task task, @ReorderMode int reorderMode, boolean runSynchronously) {
-        final TaskScreenshotAnimatable animatable = TaskScreenshotAnimatable.create(task);
-        if (animatable != null) {
-            mRecentScreenshotAnimator = new SurfaceAnimator(
-                    animatable,
-                    () -> {
-                        if (DEBUG_RECENTS_ANIMATIONS) {
-                            Slog.d(TAG, "mRecentScreenshotAnimator finish");
-                        }
-                        mCallbacks.onAnimationFinished(reorderMode, runSynchronously,
-                                false /* sendUserLeaveHint */);
-                    }, mService);
-            mRecentScreenshotAnimator.transferAnimation(task.mSurfaceAnimator);
+    TaskSnapshot screenshotRecentTask(Task task, @ReorderMode int reorderMode,
+            boolean runSynchronously) {
+        final TaskSnapshotController snapshotController = mService.mTaskSnapshotController;
+        final ArraySet<Task> tasks = Sets.newArraySet(task);
+        snapshotController.snapshotTasks(tasks);
+        snapshotController.addSkipClosingAppSnapshotTasks(tasks);
+        final TaskSnapshot taskSnapshot = snapshotController.getSnapshot(task.mTaskId,
+                task.mUserId, false /* restoreFromDisk */, false /* reducedResolution */);
+        if (taskSnapshot == null) {
+            return null;
         }
+
+        final TaskScreenshotAnimatable animatable = new TaskScreenshotAnimatable(task,
+                new SurfaceControl.ScreenshotGraphicBuffer(taskSnapshot.getSnapshot(),
+                        taskSnapshot.getColorSpace(), false /* containsSecureLayers */));
+        mRecentScreenshotAnimator = new SurfaceAnimator(
+                animatable,
+                () -> {
+                    if (DEBUG_RECENTS_ANIMATIONS) {
+                        Slog.d(TAG, "mRecentScreenshotAnimator finish");
+                    }
+                    mCallbacks.onAnimationFinished(reorderMode, runSynchronously,
+                            false /* sendUserLeaveHint */);
+                }, mService);
+        mRecentScreenshotAnimator.transferAnimation(task.mSurfaceAnimator);
+        return taskSnapshot;
     }
 
     void cleanupAnimation(@ReorderMode int reorderMode) {
diff --git a/services/core/java/com/android/server/wm/TaskScreenshotAnimatable.java b/services/core/java/com/android/server/wm/TaskScreenshotAnimatable.java
index 143543e..d6c2f66 100644
--- a/services/core/java/com/android/server/wm/TaskScreenshotAnimatable.java
+++ b/services/core/java/com/android/server/wm/TaskScreenshotAnimatable.java
@@ -18,7 +18,6 @@
 import static com.android.server.wm.WindowManagerDebugConfig.DEBUG_RECENTS_ANIMATIONS;
 
 import android.graphics.GraphicBuffer;
-import android.graphics.Rect;
 import android.util.Slog;
 import android.view.Surface;
 import android.view.SurfaceControl;
@@ -37,20 +36,7 @@
     private int mWidth;
     private int mHeight;
 
-    public static TaskScreenshotAnimatable create(Task task) {
-        return new TaskScreenshotAnimatable(task, getBufferFromTask(task));
-    }
-
-    private static SurfaceControl.ScreenshotGraphicBuffer getBufferFromTask(Task task) {
-        if (task == null) {
-            return null;
-        }
-        final Rect tmpRect = task.getBounds();
-        tmpRect.offset(0, 0);
-        return SurfaceControl.captureLayers(task.getSurfaceControl(), tmpRect, 1f);
-    }
-
-    private TaskScreenshotAnimatable(Task task,
+    TaskScreenshotAnimatable(Task task,
             SurfaceControl.ScreenshotGraphicBuffer screenshotBuffer) {
         GraphicBuffer buffer = screenshotBuffer == null
                 ? null : screenshotBuffer.getGraphicBuffer();
diff --git a/services/tests/wmtests/src/com/android/server/wm/RecentsAnimationControllerTest.java b/services/tests/wmtests/src/com/android/server/wm/RecentsAnimationControllerTest.java
index cd292b2..f5a1d75 100644
--- a/services/tests/wmtests/src/com/android/server/wm/RecentsAnimationControllerTest.java
+++ b/services/tests/wmtests/src/com/android/server/wm/RecentsAnimationControllerTest.java
@@ -44,6 +44,7 @@
 import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.ArgumentMatchers.eq;
 
+import android.app.ActivityManager.TaskSnapshot;
 import android.os.Binder;
 import android.os.IInterface;
 import android.platform.test.annotations.Presubmit;
@@ -73,6 +74,7 @@
     @Mock OnAnimationFinishedCallback mFinishedCallback;
     @Mock IRecentsAnimationRunner mMockRunner;
     @Mock RecentsAnimationController.RecentsAnimationCallbacks mAnimationCallbacks;
+    @Mock TaskSnapshot mMockTaskSnapshot;
     private RecentsAnimationController mController;
 
     @Before
@@ -104,7 +106,7 @@
         // Verify that the finish callback to reparent the leash is called
         verify(mFinishedCallback).onAnimationFinished(eq(adapter));
         // Verify the animation canceled callback to the app was made
-        verify(mMockRunner).onAnimationCanceled(false);
+        verify(mMockRunner).onAnimationCanceled(null /* taskSnapshot */);
         verifyNoMoreInteractionsExceptAsBinder(mMockRunner);
     }
 
@@ -162,7 +164,7 @@
 
         mController.setDeferredCancel(true /* deferred */, false /* screenshot */);
         mController.cancelAnimationWithScreenshot(false /* screenshot */);
-        verify(mMockRunner).onAnimationCanceled(false /* deferredWithScreenshot */);
+        verify(mMockRunner).onAnimationCanceled(null /* taskSnapshot */);
         assertNull(mController.mRecentScreenshotAnimator);
 
         // Simulate the app transition finishing
@@ -183,9 +185,14 @@
         mController.addAnimation(appWindow.getTask(), false /* isRecentTaskInvisible */);
         assertTrue(mController.isAnimatingTask(appWindow.getTask()));
 
+        spyOn(mWm.mTaskSnapshotController);
+        doNothing().when(mWm.mTaskSnapshotController).notifyAppVisibilityChanged(any(),
+                anyBoolean());
+        doReturn(mMockTaskSnapshot).when(mWm.mTaskSnapshotController).getSnapshot(anyInt(),
+                anyInt(), eq(false) /* restoreFromDisk */, eq(false) /* reducedResolution */);
         mController.setDeferredCancel(true /* deferred */, true /* screenshot */);
         mController.cancelAnimationWithScreenshot(true /* screenshot */);
-        verify(mMockRunner).onAnimationCanceled(true /* deferredWithScreenshot */);
+        verify(mMockRunner).onAnimationCanceled(mMockTaskSnapshot /* taskSnapshot */);
         assertNotNull(mController.mRecentScreenshotAnimator);
         assertTrue(mController.mRecentScreenshotAnimator.isAnimating());