Merge "Add TransitionCookie and ephemeral return functionality." into main

7 files changed, 435 insertions, 19 deletions

diff --git a/core/java/android/window/TransitionFilter.java b/core/java/android/window/TransitionFilter.java
index 64fe66e36bc3..ec4e3e9163fa 100644
--- a/core/java/android/window/TransitionFilter.java
+++ b/core/java/android/window/TransitionFilter.java
@@ -25,6 +25,7 @@ import android.annotation.Nullable;
 import android.app.ActivityManager;
 import android.app.WindowConfiguration;
 import android.content.ComponentName;
+import android.os.IBinder;
 import android.os.Parcel;
 import android.os.Parcelable;
 import android.view.WindowManager;
@@ -180,6 +181,7 @@ public final class TransitionFilter implements Parcelable {
 
         public @ContainerOrder int mOrder = CONTAINER_ORDER_ANY;
         public ComponentName mTopActivity;
+        public IBinder mLaunchCookie;
 
         public Requirement() {
         }
@@ -193,6 +195,7 @@ public final class TransitionFilter implements Parcelable {
             mMustBeTask = in.readBoolean();
             mOrder = in.readInt();
             mTopActivity = in.readTypedObject(ComponentName.CREATOR);
+            mLaunchCookie = in.readStrongBinder();
         }
 
         /** Go through changes and find if at-least one change matches this filter */
@@ -231,6 +234,9 @@ public final class TransitionFilter implements Parcelable {
                 if (mMustBeTask && change.getTaskInfo() == null) {
                     continue;
                 }
+                if (!matchesCookie(change.getTaskInfo())) {
+                    continue;
+                }
                 return true;
             }
             return false;
@@ -247,13 +253,25 @@ public final class TransitionFilter implements Parcelable {
             return false;
         }
 
+        private boolean matchesCookie(ActivityManager.RunningTaskInfo info) {
+            if (mLaunchCookie == null) return true;
+            if (info == null) return false;
+            for (IBinder cookie : info.launchCookies) {
+                if (mLaunchCookie.equals(cookie)) {
+                    return true;
+                }
+            }
+            return false;
+        }
+
         /** Check if the request matches this filter. It may generate false positives */
         boolean matches(@NonNull TransitionRequestInfo request) {
             // Can't check modes/order since the transition hasn't been built at this point.
             if (mActivityType == ACTIVITY_TYPE_UNDEFINED) return true;
             return request.getTriggerTask() != null
                     && request.getTriggerTask().getActivityType() == mActivityType
-                    && matchesTopActivity(request.getTriggerTask(), null /* activityCmp */);
+                    && matchesTopActivity(request.getTriggerTask(), null /* activityCmp */)
+                    && matchesCookie(request.getTriggerTask());
         }
 
         @Override
@@ -267,6 +285,7 @@ public final class TransitionFilter implements Parcelable {
             dest.writeBoolean(mMustBeTask);
             dest.writeInt(mOrder);
             dest.writeTypedObject(mTopActivity, flags);
+            dest.writeStrongBinder(mLaunchCookie);
         }
 
         @NonNull
@@ -307,6 +326,7 @@ public final class TransitionFilter implements Parcelable {
             out.append(" mustBeTask=" + mMustBeTask);
             out.append(" order=" + containerOrderToString(mOrder));
             out.append(" topActivity=").append(mTopActivity);
+            out.append(" launchCookie=").append(mLaunchCookie);
             out.append("}");
             return out.toString();
         }
diff --git a/packages/SystemUI/animation/src/com/android/systemui/animation/ActivityTransitionAnimator.kt b/packages/SystemUI/animation/src/com/android/systemui/animation/ActivityTransitionAnimator.kt
index d4660fa74f0a..23df26fdb246 100644
--- a/packages/SystemUI/animation/src/com/android/systemui/animation/ActivityTransitionAnimator.kt
+++ b/packages/SystemUI/animation/src/com/android/systemui/animation/ActivityTransitionAnimator.kt
@@ -23,6 +23,7 @@ import android.app.TaskInfo
 import android.graphics.Matrix
 import android.graphics.Rect
 import android.graphics.RectF
+import android.os.Binder
 import android.os.Build
 import android.os.Handler
 import android.os.Looper
@@ -36,7 +37,11 @@ import android.view.SyncRtSurfaceTransactionApplier
 import android.view.View
 import android.view.ViewGroup
 import android.view.WindowManager
+import android.view.WindowManager.TRANSIT_CLOSE
+import android.view.WindowManager.TRANSIT_TO_BACK
 import android.view.animation.PathInterpolator
+import android.window.RemoteTransition
+import android.window.TransitionFilter
 import androidx.annotation.AnyThread
 import androidx.annotation.BinderThread
 import androidx.annotation.UiThread
@@ -44,6 +49,9 @@ import com.android.app.animation.Interpolators
 import com.android.internal.annotations.VisibleForTesting
 import com.android.internal.policy.ScreenDecorationsUtils
 import com.android.systemui.Flags.activityTransitionUseLargestWindow
+import com.android.systemui.shared.Flags.returnAnimationFrameworkLibrary
+import com.android.wm.shell.shared.IShellTransitions
+import com.android.wm.shell.shared.ShellTransitions
 import java.util.concurrent.Executor
 import kotlin.math.roundToInt
 
@@ -59,6 +67,9 @@ constructor(
     /** The executor that runs on the main thread. */
     private val mainExecutor: Executor,
 
+    /** The object used to register ephemeral returns and long-lived transitions. */
+    private val transitionRegister: TransitionRegister? = null,
+
     /** The animator used when animating a View into an app. */
     private val transitionAnimator: TransitionAnimator = defaultTransitionAnimator(mainExecutor),
 
@@ -74,6 +85,36 @@ constructor(
     // TODO(b/301385865): Remove this flag.
     private val disableWmTimeout: Boolean = false,
 ) {
+    @JvmOverloads
+    constructor(
+        mainExecutor: Executor,
+        shellTransitions: ShellTransitions,
+        transitionAnimator: TransitionAnimator = defaultTransitionAnimator(mainExecutor),
+        dialogToAppAnimator: TransitionAnimator = defaultDialogToAppAnimator(mainExecutor),
+        disableWmTimeout: Boolean = false,
+    ) : this(
+        mainExecutor,
+        TransitionRegister.fromShellTransitions(shellTransitions),
+        transitionAnimator,
+        dialogToAppAnimator,
+        disableWmTimeout,
+    )
+
+    @JvmOverloads
+    constructor(
+        mainExecutor: Executor,
+        iShellTransitions: IShellTransitions,
+        transitionAnimator: TransitionAnimator = defaultTransitionAnimator(mainExecutor),
+        dialogToAppAnimator: TransitionAnimator = defaultDialogToAppAnimator(mainExecutor),
+        disableWmTimeout: Boolean = false,
+    ) : this(
+        mainExecutor,
+        TransitionRegister.fromIShellTransitions(iShellTransitions),
+        transitionAnimator,
+        dialogToAppAnimator,
+        disableWmTimeout,
+    )
+
     companion object {
         /** The timings when animating a View into an app. */
         @JvmField
@@ -233,6 +274,10 @@ constructor(
             }
         }
 
+        if (animationAdapter != null && controller.transitionCookie != null) {
+            registerEphemeralReturnAnimation(controller, transitionRegister)
+        }
+
         val launchResult = intentStarter(animationAdapter)
 
         // Only animate if the app is not already on top and will be opened, unless we are on the
@@ -302,6 +347,66 @@ constructor(
         }
     }
 
+    /**
+     * Uses [transitionRegister] to set up the return animation for the given [launchController].
+     *
+     * De-registration is set up automatically once the return animation is run.
+     *
+     * TODO(b/339194555): automatically de-register when the launchable is detached.
+     */
+    private fun registerEphemeralReturnAnimation(
+        launchController: Controller,
+        transitionRegister: TransitionRegister?
+    ) {
+        if (!returnAnimationFrameworkLibrary()) return
+
+        var cleanUpRunnable: Runnable? = null
+        val returnRunner =
+            createRunner(
+                object : DelegateTransitionAnimatorController(launchController) {
+                    override val isLaunching = false
+
+                    override fun onTransitionAnimationCancelled(
+                        newKeyguardOccludedState: Boolean?
+                    ) {
+                        super.onTransitionAnimationCancelled(newKeyguardOccludedState)
+                        cleanUp()
+                    }
+
+                    override fun onTransitionAnimationEnd(isExpandingFullyAbove: Boolean) {
+                        super.onTransitionAnimationEnd(isExpandingFullyAbove)
+                        cleanUp()
+                    }
+
+                    private fun cleanUp() {
+                        cleanUpRunnable?.run()
+                    }
+                }
+            )
+
+        // mTypeSet and mModes match back signals only, and not home. This is on purpose, because
+        // we only want ephemeral return animations triggered in these scenarios.
+        val filter =
+            TransitionFilter().apply {
+                mTypeSet = intArrayOf(TRANSIT_CLOSE, TRANSIT_TO_BACK)
+                mRequirements =
+                    arrayOf(
+                        TransitionFilter.Requirement().apply {
+                            mLaunchCookie = launchController.transitionCookie
+                            mModes = intArrayOf(TRANSIT_CLOSE, TRANSIT_TO_BACK)
+                        }
+                    )
+            }
+        val transition =
+            RemoteTransition(
+                RemoteAnimationRunnerCompat.wrap(returnRunner),
+                "${launchController.transitionCookie}_returnTransition"
+            )
+
+        transitionRegister?.register(filter, transition)
+        cleanUpRunnable = Runnable { transitionRegister?.unregister(transition) }
+    }
+
     /** Add a [Listener] that can listen to transition animations. */
     fun addListener(listener: Listener) {
         listeners.add(listener)
@@ -386,8 +491,14 @@ constructor(
              * Note: The background of [view] should be a (rounded) rectangle so that it can be
              * properly animated.
              */
+            @JvmOverloads
             @JvmStatic
-            fun fromView(view: View, cujType: Int? = null): Controller? {
+            fun fromView(
+                view: View,
+                cujType: Int? = null,
+                cookie: TransitionCookie? = null,
+                returnCujType: Int? = null
+            ): Controller? {
                 // Make sure the View we launch from implements LaunchableView to avoid visibility
                 // issues.
                 if (view !is LaunchableView) {
@@ -408,7 +519,7 @@ constructor(
                     return null
                 }
 
-                return GhostedViewTransitionAnimatorController(view, cujType)
+                return GhostedViewTransitionAnimatorController(view, cujType, cookie, returnCujType)
             }
         }
 
@@ -432,6 +543,17 @@ constructor(
             get() = false
 
         /**
+         * The cookie associated with the transition controlled by this [Controller].
+         *
+         * This should be defined for all return [Controller] (when [isLaunching] is false) and for
+         * their associated launch [Controller]s.
+         *
+         * For the recommended format, see [TransitionCookie].
+         */
+        val transitionCookie: TransitionCookie?
+            get() = null
+
+        /**
          * The intent was started. If [willAnimate] is false, nothing else will happen and the
          * animation will not be started.
          */
@@ -652,7 +774,7 @@ constructor(
                 return
             }
 
-            val window = findRootTaskIfPossible(apps)
+            val window = findTargetWindowIfPossible(apps)
             if (window == null) {
                 Log.i(TAG, "Aborting the animation as no window is opening")
                 callback?.invoke()
@@ -676,7 +798,7 @@ constructor(
             startAnimation(window, navigationBar, callback)
         }
 
-        private fun findRootTaskIfPossible(
+        private fun findTargetWindowIfPossible(
             apps: Array<out RemoteAnimationTarget>?
         ): RemoteAnimationTarget? {
             if (apps == null) {
@@ -694,6 +816,19 @@ constructor(
             for (it in apps) {
                 if (it.mode == targetMode) {
                     if (activityTransitionUseLargestWindow()) {
+                        if (returnAnimationFrameworkLibrary()) {
+                            // If the controller contains a cookie, _only_ match if the candidate
+                            // contains the matching cookie.
+                            if (
+                                controller.transitionCookie != null &&
+                                    it.taskInfo
+                                        ?.launchCookies
+                                        ?.contains(controller.transitionCookie) != true
+                            ) {
+                                continue
+                            }
+                        }
+
                         if (
                             candidate == null ||
                                 !it.hasAnimatingParent && candidate.hasAnimatingParent
@@ -806,11 +941,7 @@ constructor(
                         progress: Float,
                         linearProgress: Float
                     ) {
-                        // Apply the state to the window only if it is visible, i.e. when the
-                        // expanding view is *not* visible.
-                        if (!state.visible) {
-                            applyStateToWindow(window, state, linearProgress)
-                        }
+                        applyStateToWindow(window, state, linearProgress)
                         navigationBar?.let { applyStateToNavigationBar(it, state, linearProgress) }
 
                         listener?.onTransitionAnimationProgress(linearProgress)
@@ -1048,4 +1179,72 @@ constructor(
             return (this.width() * this.height()) > (other.width() * other.height())
         }
     }
+
+    /**
+     * Wraps one of the two methods we have to register remote transitions with WM Shell:
+     * - for in-process registrations (e.g. System UI) we use [ShellTransitions]
+     * - for cross-process registrations (e.g. Launcher) we use [IShellTransitions]
+     *
+     * Important: each instance of this class must wrap exactly one of the two.
+     */
+    class TransitionRegister
+    private constructor(
+        private val shellTransitions: ShellTransitions? = null,
+        private val iShellTransitions: IShellTransitions? = null,
+    ) {
+        init {
+            assert((shellTransitions != null).xor(iShellTransitions != null))
+        }
+
+        companion object {
+            /** Provides a [TransitionRegister] instance wrapping [ShellTransitions]. */
+            fun fromShellTransitions(shellTransitions: ShellTransitions): TransitionRegister {
+                return TransitionRegister(shellTransitions = shellTransitions)
+            }
+
+            /** Provides a [TransitionRegister] instance wrapping [IShellTransitions]. */
+            fun fromIShellTransitions(iShellTransitions: IShellTransitions): TransitionRegister {
+                return TransitionRegister(iShellTransitions = iShellTransitions)
+            }
+        }
+
+        /** Register [remoteTransition] with WM Shell using the given [filter]. */
+        internal fun register(
+            filter: TransitionFilter,
+            remoteTransition: RemoteTransition,
+        ) {
+            shellTransitions?.registerRemote(filter, remoteTransition)
+            iShellTransitions?.registerRemote(filter, remoteTransition)
+        }
+
+        /** Unregister [remoteTransition] from WM Shell. */
+        internal fun unregister(remoteTransition: RemoteTransition) {
+            shellTransitions?.unregisterRemote(remoteTransition)
+            iShellTransitions?.unregisterRemote(remoteTransition)
+        }
+    }
+
+    /**
+     * A cookie used to uniquely identify a task launched using an
+     * [ActivityTransitionAnimator.Controller].
+     *
+     * The [String] encapsulated by this class should be formatted in such a way to be unique across
+     * the system, but reliably constant for the same associated launchable.
+     *
+     * Recommended naming scheme:
+     * - DO use the fully qualified name of the class that owns the instance of the launchable,
+     *   along with a concise and precise description of the purpose of the launchable in question.
+     * - DO NOT introduce uniqueness through the use of timestamps or other runtime variables that
+     *   will change if the instance is destroyed and re-created.
+     *
+     * Example: "com.not.the.real.class.name.ShadeController_openSettingsButton"
+     *
+     * Note that sometimes (e.g. in recycler views) there could be multiple instances of the same
+     * launchable, and no static knowledge to adequately differentiate between them using a single
+     * description. In this case, the recommendation is to append a unique identifier related to the
+     * contents of the launchable.
+     *
+     * Example: “com.not.the.real.class.name.ToastWebResult_launchAga_id143256”
+     */
+    data class TransitionCookie(private val cookie: String) : Binder()
 }
diff --git a/packages/SystemUI/animation/src/com/android/systemui/animation/Expandable.kt b/packages/SystemUI/animation/src/com/android/systemui/animation/Expandable.kt
index e4bb2adbefb4..21557b8bb402 100644
--- a/packages/SystemUI/animation/src/com/android/systemui/animation/Expandable.kt
+++ b/packages/SystemUI/animation/src/com/android/systemui/animation/Expandable.kt
@@ -25,10 +25,30 @@ interface Expandable {
      * [Expandable] into an Activity, or return `null` if this [Expandable] should not be animated
      * (e.g. if it is currently not attached or visible).
      *
-     * @param cujType the CUJ type from the [com.android.internal.jank.InteractionJankMonitor]
+     * @param launchCujType The CUJ type from the [com.android.internal.jank.InteractionJankMonitor]
      *   associated to the launch that will use this controller.
+     * @param cookie The unique cookie associated with the launch that will use this controller.
+     *   This is required iff the a return animation should be included.
+     * @param returnCujType The CUJ type from the [com.android.internal.jank.InteractionJankMonitor]
+     *   associated to the return animation that will use this controller.
      */
-    fun activityTransitionController(cujType: Int? = null): ActivityTransitionAnimator.Controller?
+    fun activityTransitionController(
+        launchCujType: Int? = null,
+        cookie: ActivityTransitionAnimator.TransitionCookie? = null,
+        returnCujType: Int? = null
+    ): ActivityTransitionAnimator.Controller?
+
+    /**
+     * See [activityTransitionController] above.
+     *
+     * Interfaces don't support [JvmOverloads], so this is a useful overload for Java usages that
+     * don't use the return-related parameters.
+     */
+    fun activityTransitionController(
+        launchCujType: Int? = null
+    ): ActivityTransitionAnimator.Controller? {
+        return activityTransitionController(launchCujType, cookie = null, returnCujType = null)
+    }
 
     /**
      * Create a [DialogTransitionAnimator.Controller] that can be used to expand this [Expandable]
@@ -48,9 +68,16 @@ interface Expandable {
         fun fromView(view: View): Expandable {
             return object : Expandable {
                 override fun activityTransitionController(
-                    cujType: Int?,
+                    launchCujType: Int?,
+                    cookie: ActivityTransitionAnimator.TransitionCookie?,
+                    returnCujType: Int?
                 ): ActivityTransitionAnimator.Controller? {
-                    return ActivityTransitionAnimator.Controller.fromView(view, cujType)
+                    return ActivityTransitionAnimator.Controller.fromView(
+                        view,
+                        launchCujType,
+                        cookie,
+                        returnCujType
+                    )
                 }
 
                 override fun dialogTransitionController(
diff --git a/packages/SystemUI/animation/src/com/android/systemui/animation/GhostedViewTransitionAnimatorController.kt b/packages/SystemUI/animation/src/com/android/systemui/animation/GhostedViewTransitionAnimatorController.kt
index fd79f62debce..9d4507337e51 100644
--- a/packages/SystemUI/animation/src/com/android/systemui/animation/GhostedViewTransitionAnimatorController.kt
+++ b/packages/SystemUI/animation/src/com/android/systemui/animation/GhostedViewTransitionAnimatorController.kt
@@ -59,8 +59,12 @@ constructor(
     /** The view that will be ghosted and from which the background will be extracted. */
     private val ghostedView: View,
 
-    /** The [CujType] associated to this animation. */
-    private val cujType: Int? = null,
+    /** The [CujType] associated to this launch animation. */
+    private val launchCujType: Int? = null,
+    override val transitionCookie: ActivityTransitionAnimator.TransitionCookie? = null,
+
+    /** The [CujType] associated to this return animation. */
+    private val returnCujType: Int? = null,
     private var interactionJankMonitor: InteractionJankMonitor =
         InteractionJankMonitor.getInstance(),
 ) : ActivityTransitionAnimator.Controller {
@@ -104,6 +108,15 @@ constructor(
      */
     private val background: Drawable?
 
+    /** CUJ identifier accounting for whether this controller is for a launch or a return. */
+    private val cujType: Int?
+        get() =
+            if (isLaunching) {
+                launchCujType
+            } else {
+                returnCujType
+            }
+
     init {
         // Make sure the View we launch from implements LaunchableView to avoid visibility issues.
         if (ghostedView !is LaunchableView) {
diff --git a/packages/SystemUI/compose/core/src/com/android/compose/animation/ExpandableController.kt b/packages/SystemUI/compose/core/src/com/android/compose/animation/ExpandableController.kt
index c7f0a965206e..17a606171a9e 100644
--- a/packages/SystemUI/compose/core/src/com/android/compose/animation/ExpandableController.kt
+++ b/packages/SystemUI/compose/core/src/com/android/compose/animation/ExpandableController.kt
@@ -134,13 +134,15 @@ internal class ExpandableControllerImpl(
     override val expandable: Expandable =
         object : Expandable {
             override fun activityTransitionController(
-                cujType: Int?,
+                launchCujType: Int?,
+                cookie: ActivityTransitionAnimator.TransitionCookie?,
+                returnCujType: Int?
             ): ActivityTransitionAnimator.Controller? {
                 if (!isComposed.value) {
                     return null
                 }
 
-                return activityController(cujType)
+                return activityController(launchCujType, cookie, returnCujType)
             }
 
             override fun dialogTransitionController(
@@ -262,10 +264,27 @@ internal class ExpandableControllerImpl(
     }
 
     /** Create an [ActivityTransitionAnimator.Controller] that can be used to animate activities. */
-    private fun activityController(cujType: Int?): ActivityTransitionAnimator.Controller {
+    private fun activityController(
+        launchCujType: Int?,
+        cookie: ActivityTransitionAnimator.TransitionCookie?,
+        returnCujType: Int?
+    ): ActivityTransitionAnimator.Controller {
         val delegate = transitionController()
         return object :
             ActivityTransitionAnimator.Controller, TransitionAnimator.Controller by delegate {
+            /**
+             * CUJ identifier accounting for whether this controller is for a launch or a return.
+             */
+            private val cujType: Int?
+                get() =
+                    if (isLaunching) {
+                        launchCujType
+                    } else {
+                        returnCujType
+                    }
+
+            override val transitionCookie = cookie
+
             override fun onTransitionAnimationStart(isExpandingFullyAbove: Boolean) {
                 delegate.onTransitionAnimationStart(isExpandingFullyAbove)
                 overlay.value = composeViewRoot.rootView.overlay as ViewGroupOverlay
diff --git a/packages/SystemUI/tests/src/com/android/systemui/animation/ActivityTransitionAnimatorTest.kt b/packages/SystemUI/tests/src/com/android/systemui/animation/ActivityTransitionAnimatorTest.kt
index 8e4c155593e2..fd37cad72371 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/animation/ActivityTransitionAnimatorTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/animation/ActivityTransitionAnimatorTest.kt
@@ -8,6 +8,7 @@ import android.content.pm.ApplicationInfo
 import android.graphics.Point
 import android.graphics.Rect
 import android.os.Looper
+import android.platform.test.flag.junit.SetFlagsRule
 import android.testing.AndroidTestingRunner
 import android.testing.TestableLooper.RunWithLooper
 import android.view.IRemoteAnimationFinishedCallback
@@ -17,15 +18,20 @@ import android.view.SurfaceControl
 import android.view.ViewGroup
 import android.widget.FrameLayout
 import android.widget.LinearLayout
+import android.window.RemoteTransition
+import android.window.TransitionFilter
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
+import com.android.systemui.shared.Flags
 import com.android.systemui.util.mockito.any
+import com.android.wm.shell.shared.ShellTransitions
 import junit.framework.Assert.assertFalse
 import junit.framework.Assert.assertNotNull
 import junit.framework.Assert.assertNull
 import junit.framework.Assert.assertTrue
 import junit.framework.AssertionFailedError
 import kotlin.concurrent.thread
+import kotlin.test.assertEquals
 import org.junit.After
 import org.junit.Assert.assertThrows
 import org.junit.Before
@@ -48,6 +54,7 @@ class ActivityTransitionAnimatorTest : SysuiTestCase() {
     private val transitionContainer = LinearLayout(mContext)
     private val mainExecutor = context.mainExecutor
     private val testTransitionAnimator = fakeTransitionAnimator(mainExecutor)
+    private val testShellTransitions = FakeShellTransitions()
     @Mock lateinit var callback: ActivityTransitionAnimator.Callback
     @Mock lateinit var listener: ActivityTransitionAnimator.Listener
     @Spy private val controller = TestTransitionAnimatorController(transitionContainer)
@@ -55,12 +62,16 @@ class ActivityTransitionAnimatorTest : SysuiTestCase() {
 
     private lateinit var activityTransitionAnimator: ActivityTransitionAnimator
     @get:Rule val rule = MockitoJUnit.rule()
+    @get:Rule val setFlagsRule = SetFlagsRule()
 
     @Before
     fun setup() {
         activityTransitionAnimator =
             ActivityTransitionAnimator(
                 mainExecutor,
+                ActivityTransitionAnimator.TransitionRegister.fromShellTransitions(
+                    testShellTransitions
+                ),
                 testTransitionAnimator,
                 testTransitionAnimator,
                 disableWmTimeout = true,
@@ -164,6 +175,34 @@ class ActivityTransitionAnimatorTest : SysuiTestCase() {
     }
 
     @Test
+    fun registersReturnIffCookieIsPresent() {
+        setFlagsRule.enableFlags(Flags.FLAG_RETURN_ANIMATION_FRAMEWORK_LIBRARY)
+        `when`(callback.isOnKeyguard()).thenReturn(false)
+
+        startIntentWithAnimation(activityTransitionAnimator, controller) { _ ->
+            ActivityManager.START_DELIVERED_TO_TOP
+        }
+
+        waitForIdleSync()
+        assertTrue(testShellTransitions.remotes.isEmpty())
+        assertTrue(testShellTransitions.remotesForTakeover.isEmpty())
+
+        val controller =
+            object : DelegateTransitionAnimatorController(controller) {
+                override val transitionCookie
+                    get() = ActivityTransitionAnimator.TransitionCookie("testCookie")
+            }
+
+        startIntentWithAnimation(activityTransitionAnimator, controller) { _ ->
+            ActivityManager.START_DELIVERED_TO_TOP
+        }
+
+        waitForIdleSync()
+        assertEquals(1, testShellTransitions.remotes.size)
+        assertTrue(testShellTransitions.remotesForTakeover.isEmpty())
+    }
+
+    @Test
     fun doesNotStartIfAnimationIsCancelled() {
         val runner = activityTransitionAnimator.createRunner(controller)
         runner.onAnimationCancelled()
@@ -243,6 +282,35 @@ class ActivityTransitionAnimatorTest : SysuiTestCase() {
 }
 
 /**
+ * A fake implementation of [ShellTransitions] which saves filter-transition pairs locally and
+ * allows inspection.
+ */
+private class FakeShellTransitions : ShellTransitions {
+    val remotes = mutableMapOf<TransitionFilter, RemoteTransition>()
+    val remotesForTakeover = mutableMapOf<TransitionFilter, RemoteTransition>()
+
+    override fun registerRemote(filter: TransitionFilter, remoteTransition: RemoteTransition) {
+        remotes[filter] = remoteTransition
+    }
+
+    override fun registerRemoteForTakeover(
+        filter: TransitionFilter,
+        remoteTransition: RemoteTransition
+    ) {
+        remotesForTakeover[filter] = remoteTransition
+    }
+
+    override fun unregisterRemote(remoteTransition: RemoteTransition) {
+        while (remotes.containsValue(remoteTransition)) {
+            remotes.values.remove(remoteTransition)
+        }
+        while (remotesForTakeover.containsValue(remoteTransition)) {
+            remotesForTakeover.values.remove(remoteTransition)
+        }
+    }
+}
+
+/**
  * A simple implementation of [ActivityTransitionAnimator.Controller] which throws if it is called
  * outside of the main thread.
  */
diff --git a/packages/SystemUI/tests/src/com/android/systemui/animation/GhostedViewTransitionAnimatorControllerTest.kt b/packages/SystemUI/tests/src/com/android/systemui/animation/GhostedViewTransitionAnimatorControllerTest.kt
index b31fe21f8e91..42fcd547408a 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/animation/GhostedViewTransitionAnimatorControllerTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/animation/GhostedViewTransitionAnimatorControllerTest.kt
@@ -16,12 +16,16 @@
 
 package com.android.systemui.animation
 
+import android.os.HandlerThread
 import android.testing.AndroidTestingRunner
 import android.testing.TestableLooper
+import android.view.View
 import android.widget.FrameLayout
 import androidx.test.filters.SmallTest
+import com.android.internal.jank.InteractionJankMonitor
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.animation.view.LaunchableFrameLayout
+import com.google.common.truth.Truth.assertThat
 import org.junit.Assert.assertThrows
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -30,6 +34,13 @@ import org.junit.runner.RunWith
 @RunWith(AndroidTestingRunner::class)
 @TestableLooper.RunWithLooper
 class GhostedViewTransitionAnimatorControllerTest : SysuiTestCase() {
+    companion object {
+        private const val LAUNCH_CUJ = 0
+        private const val RETURN_CUJ = 1
+    }
+
+    private val interactionJankMonitor = FakeInteractionJankMonitor()
+
     @Test
     fun animatingOrphanViewDoesNotCrash() {
         val state = TransitionAnimator.State(top = 0, bottom = 0, left = 0, right = 0)
@@ -47,4 +58,63 @@ class GhostedViewTransitionAnimatorControllerTest : SysuiTestCase() {
             GhostedViewTransitionAnimatorController(FrameLayout(mContext))
         }
     }
+
+    @Test
+    fun cujsAreLoggedCorrectly() {
+        val parent = FrameLayout(mContext)
+
+        val launchView = LaunchableFrameLayout(mContext)
+        parent.addView((launchView))
+        val launchController =
+            GhostedViewTransitionAnimatorController(
+                    launchView,
+                launchCujType = LAUNCH_CUJ,
+                returnCujType = RETURN_CUJ,
+                interactionJankMonitor = interactionJankMonitor
+            )
+        launchController.onTransitionAnimationStart(isExpandingFullyAbove = true)
+        assertThat(interactionJankMonitor.ongoing).containsExactly(LAUNCH_CUJ)
+        launchController.onTransitionAnimationEnd(isExpandingFullyAbove = true)
+        assertThat(interactionJankMonitor.ongoing).isEmpty()
+        assertThat(interactionJankMonitor.finished).containsExactly(LAUNCH_CUJ)
+
+        val returnView = LaunchableFrameLayout(mContext)
+        parent.addView((returnView))
+        val returnController =
+            object : GhostedViewTransitionAnimatorController(
+                returnView,
+                launchCujType = LAUNCH_CUJ,
+                returnCujType = RETURN_CUJ,
+                interactionJankMonitor = interactionJankMonitor
+            ) {
+                override val isLaunching = false
+            }
+        returnController.onTransitionAnimationStart(isExpandingFullyAbove = true)
+        assertThat(interactionJankMonitor.ongoing).containsExactly(RETURN_CUJ)
+        returnController.onTransitionAnimationEnd(isExpandingFullyAbove = true)
+        assertThat(interactionJankMonitor.ongoing).isEmpty()
+        assertThat(interactionJankMonitor.finished).containsExactly(LAUNCH_CUJ, RETURN_CUJ)
+    }
+
+    /**
+     * A fake implementation of [InteractionJankMonitor] which stores ongoing and finished CUJs and
+     * allows inspection.
+     */
+    private class FakeInteractionJankMonitor : InteractionJankMonitor(
+        HandlerThread("testThread")
+    ) {
+        val ongoing: MutableSet<Int> = mutableSetOf()
+        val finished: MutableSet<Int> = mutableSetOf()
+
+        override fun begin(v: View?, cujType: Int): Boolean {
+            ongoing.add(cujType)
+            return true
+        }
+
+        override fun end(cujType: Int): Boolean {
+            ongoing.remove(cujType)
+            finished.add(cujType)
+            return true
+        }
+    }
 }