diff --git a/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java b/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
index f60b225..1b1cb85 100644
--- a/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
+++ b/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
@@ -151,6 +151,7 @@
 
     // Null if the recents animation hasn't started yet or has been canceled or finished.
     protected @Nullable RecentsAnimationController mRecentsAnimationController;
+    protected @Nullable RecentsAnimationController mDeferredCleanupRecentsAnimationController;
     protected RecentsAnimationTargets mRecentsAnimationTargets;
     protected T mActivity;
     protected Q mRecentsView;
@@ -435,6 +436,9 @@
                     mRecentsView.switchToScreenshot(snapshots, () -> {
                         if (mRecentsAnimationController != null) {
                             mRecentsAnimationController.cleanupScreenshot();
+                        } else if (mDeferredCleanupRecentsAnimationController != null) {
+                            mDeferredCleanupRecentsAnimationController.cleanupScreenshot();
+                            mDeferredCleanupRecentsAnimationController = null;
                         }
                     });
                     mRecentsView.onRecentsAnimationComplete();
@@ -839,6 +843,9 @@
     public void onRecentsAnimationCanceled(HashMap<Integer, ThumbnailData> thumbnailDatas) {
         ActiveGestureLog.INSTANCE.addLog("cancelRecentsAnimation");
         mActivityInitListener.unregister();
+        // Cache the recents animation controller so we can defer its cleanup to after having
+        // properly cleaned up the screenshot without accidentally using it.
+        mDeferredCleanupRecentsAnimationController = mRecentsAnimationController;
         mStateCallback.setStateOnUiThread(STATE_GESTURE_CANCELLED | STATE_HANDLER_INVALIDATED);
 
         if (mRecentsAnimationTargets != null) {
diff --git a/quickstep/src/com/android/quickstep/util/SurfaceTransactionApplier.java b/quickstep/src/com/android/quickstep/util/SurfaceTransactionApplier.java
index 3b1c150..4757d4b 100644
--- a/quickstep/src/com/android/quickstep/util/SurfaceTransactionApplier.java
+++ b/quickstep/src/com/android/quickstep/util/SurfaceTransactionApplier.java
@@ -53,6 +53,7 @@
         mTargetViewRootImpl = targetView.getViewRootImpl();
         mBarrierSurfaceControl = mTargetViewRootImpl.getSurfaceControl();
         mApplyHandler = new Handler(this::onApplyMessage);
+        setCanRelease(true);
     }
 
     protected boolean onApplyMessage(Message msg) {
diff --git a/src/com/android/launcher3/anim/AnimationSuccessListener.java b/src/com/android/launcher3/anim/AnimationSuccessListener.java
index 6196df2..a312070 100644
--- a/src/com/android/launcher3/anim/AnimationSuccessListener.java
+++ b/src/com/android/launcher3/anim/AnimationSuccessListener.java
@@ -19,8 +19,6 @@
 import android.animation.Animator;
 import android.animation.AnimatorListenerAdapter;
 
-import androidx.annotation.CallSuper;
-
 /**
  * Extension of {@link AnimatorListenerAdapter} for listening for non-cancelled animations
  */
@@ -29,7 +27,6 @@
     protected boolean mCancelled = false;
 
     @Override
-    @CallSuper
     public void onAnimationCancel(Animator animation) {
         mCancelled = true;
     }
diff --git a/src/com/android/launcher3/statemanager/StateManager.java b/src/com/android/launcher3/statemanager/StateManager.java
index 9f50ff9..2aa9dde 100644
--- a/src/com/android/launcher3/statemanager/StateManager.java
+++ b/src/com/android/launcher3/statemanager/StateManager.java
@@ -335,13 +335,7 @@
             @Override
             public void onAnimationStart(Animator animation) {
                 // Change the internal state only when the transition actually starts
-                onStateTransitionStart(mCancelled ? mCurrentStableState : state);
-            }
-
-            @Override
-            public void onAnimationCancel(Animator animation) {
-                super.onAnimationCancel(animation);
-                mState = mCurrentStableState;
+                onStateTransitionStart(state);
             }
 
             @Override