Merge "Enable view anchoring in ActivityStarter." into main

8 files changed, 289 insertions, 50 deletions

diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/ActivityStarterImplTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/ActivityStarterImplTest.kt
index c3c5a48dc197..b0b80a9419e2 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/ActivityStarterImplTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/ActivityStarterImplTest.kt
@@ -18,9 +18,14 @@ package com.android.systemui.statusbar.phone
 
 import android.app.PendingIntent
 import android.content.Intent
+import android.platform.test.annotations.DisableFlags
+import android.platform.test.annotations.EnableFlags
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
+import com.android.systemui.animation.ActivityTransitionAnimator
+import com.android.systemui.flags.EnableSceneContainer
+import com.android.systemui.shared.Flags as SharedFlags
 import com.android.systemui.statusbar.SysuiStatusBarStateController
 import com.android.systemui.util.concurrency.FakeExecutor
 import com.android.systemui.util.time.FakeSystemClock
@@ -32,6 +37,9 @@ import org.mockito.Mock
 import org.mockito.Mockito.mock
 import org.mockito.Mockito.verify
 import org.mockito.MockitoAnnotations
+import org.mockito.kotlin.any
+import org.mockito.kotlin.eq
+import org.mockito.kotlin.never
 
 @SmallTest
 @RunWith(AndroidJUnit4::class)
@@ -54,6 +62,62 @@ class ActivityStarterImplTest : SysuiTestCase() {
             )
     }
 
+    @EnableFlags(
+        SharedFlags.FLAG_RETURN_ANIMATION_FRAMEWORK_LIBRARY,
+        SharedFlags.FLAG_RETURN_ANIMATION_FRAMEWORK_LONG_LIVED,
+    )
+    @EnableSceneContainer
+    @Test
+    fun registerTransition_forwardsTheRequest() {
+        val cookie = mock(ActivityTransitionAnimator.TransitionCookie::class.java)
+        val controllerFactory = mock(ActivityTransitionAnimator.ControllerFactory::class.java)
+
+        underTest.registerTransition(cookie, controllerFactory)
+
+        verify(activityStarterInternal).registerTransition(eq(cookie), eq(controllerFactory))
+    }
+
+    @DisableFlags(
+        SharedFlags.FLAG_RETURN_ANIMATION_FRAMEWORK_LIBRARY,
+        SharedFlags.FLAG_RETURN_ANIMATION_FRAMEWORK_LONG_LIVED,
+    )
+    @Test
+    fun registerTransition_doesNotForwardTheRequest_whenFlaggedOff() {
+        val cookie = mock(ActivityTransitionAnimator.TransitionCookie::class.java)
+        val controllerFactory = mock(ActivityTransitionAnimator.ControllerFactory::class.java)
+
+        underTest.registerTransition(cookie, controllerFactory)
+
+        verify(activityStarterInternal, never()).registerTransition(any(), any())
+    }
+
+    @EnableFlags(
+        SharedFlags.FLAG_RETURN_ANIMATION_FRAMEWORK_LIBRARY,
+        SharedFlags.FLAG_RETURN_ANIMATION_FRAMEWORK_LONG_LIVED,
+    )
+    @EnableSceneContainer
+    @Test
+    fun unregisterTransition_forwardsTheRequest() {
+        val cookie = mock(ActivityTransitionAnimator.TransitionCookie::class.java)
+
+        underTest.unregisterTransition(cookie)
+
+        verify(activityStarterInternal).unregisterTransition(eq(cookie))
+    }
+
+    @DisableFlags(
+        SharedFlags.FLAG_RETURN_ANIMATION_FRAMEWORK_LIBRARY,
+        SharedFlags.FLAG_RETURN_ANIMATION_FRAMEWORK_LONG_LIVED,
+    )
+    @Test
+    fun unregisterTransition_doesNotForwardTheRequest_whenFlaggedOff() {
+        val cookie = mock(ActivityTransitionAnimator.TransitionCookie::class.java)
+
+        underTest.unregisterTransition(cookie)
+
+        verify(activityStarterInternal, never()).unregisterTransition(any())
+    }
+
     @Test
     fun postStartActivityDismissingKeyguard_pendingIntent_postsOnMain() {
         val intent = mock(PendingIntent::class.java)
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/LegacyActivityStarterInternalImplTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/LegacyActivityStarterInternalImplTest.kt
index bac79a9cc520..5406acf694ff 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/LegacyActivityStarterInternalImplTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/LegacyActivityStarterInternalImplTest.kt
@@ -48,6 +48,7 @@ import com.android.systemui.shade.ShadeController
 import com.android.systemui.shade.data.repository.FakeShadeRepository
 import com.android.systemui.shade.data.repository.ShadeAnimationRepository
 import com.android.systemui.shade.domain.interactor.ShadeAnimationInteractorLegacyImpl
+import com.android.systemui.shared.Flags as SharedFlags
 import com.android.systemui.statusbar.CommandQueue
 import com.android.systemui.statusbar.NotificationLockscreenUserManager
 import com.android.systemui.statusbar.NotificationShadeWindowController
@@ -63,6 +64,7 @@ import com.google.common.truth.Truth.assertThat
 import java.util.Optional
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.flow.MutableStateFlow
+import org.junit.Assert.assertThrows
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -149,6 +151,63 @@ class LegacyActivityStarterInternalImplTest : SysuiTestCase() {
         `when`(communalSceneInteractor.isLaunchingWidget).thenReturn(MutableStateFlow(false))
     }
 
+    @EnableFlags(
+        SharedFlags.FLAG_RETURN_ANIMATION_FRAMEWORK_LIBRARY,
+        SharedFlags.FLAG_RETURN_ANIMATION_FRAMEWORK_LONG_LIVED,
+    )
+    @Test
+    fun registerTransition_registers() {
+        val cookie = mock(ActivityTransitionAnimator.TransitionCookie::class.java)
+        val controllerFactory = mock(ActivityTransitionAnimator.ControllerFactory::class.java)
+        `when`(controllerFactory.cookie).thenReturn(cookie)
+
+        underTest.registerTransition(cookie, controllerFactory)
+
+        verify(activityTransitionAnimator).register(eq(cookie), any())
+    }
+
+    @DisableFlags(
+        SharedFlags.FLAG_RETURN_ANIMATION_FRAMEWORK_LIBRARY,
+        SharedFlags.FLAG_RETURN_ANIMATION_FRAMEWORK_LONG_LIVED,
+    )
+    @Test
+    fun registerTransition_throws_whenFlagsAreDisabled() {
+        val cookie = mock(ActivityTransitionAnimator.TransitionCookie::class.java)
+        val controllerFactory = mock(ActivityTransitionAnimator.ControllerFactory::class.java)
+
+        assertThrows(IllegalStateException::class.java) {
+            underTest.registerTransition(cookie, controllerFactory)
+        }
+
+        verify(activityTransitionAnimator, never()).register(any(), any())
+    }
+
+    @EnableFlags(
+        SharedFlags.FLAG_RETURN_ANIMATION_FRAMEWORK_LIBRARY,
+        SharedFlags.FLAG_RETURN_ANIMATION_FRAMEWORK_LONG_LIVED,
+    )
+    @Test
+    fun unregisterTransition_unregisters() {
+        val cookie = mock(ActivityTransitionAnimator.TransitionCookie::class.java)
+
+        underTest.unregisterTransition(cookie)
+
+        verify(activityTransitionAnimator).unregister(eq(cookie))
+    }
+
+    @DisableFlags(
+        SharedFlags.FLAG_RETURN_ANIMATION_FRAMEWORK_LIBRARY,
+        SharedFlags.FLAG_RETURN_ANIMATION_FRAMEWORK_LONG_LIVED,
+    )
+    @Test
+    fun unregisterTransition_throws_whenFlagsAreDisabled() {
+        val cookie = mock(ActivityTransitionAnimator.TransitionCookie::class.java)
+
+        assertThrows(IllegalStateException::class.java) { underTest.unregisterTransition(cookie) }
+
+        verify(activityTransitionAnimator, never()).unregister(eq(cookie))
+    }
+
     @Test
     fun startActivityDismissingKeyguard_dismissShadeWhenOccluded_runAfterKeyguardGone() {
         val intent = mock(Intent::class.java)
@@ -216,7 +275,7 @@ class LegacyActivityStarterInternalImplTest : SysuiTestCase() {
         underTest.startPendingIntentDismissingKeyguard(
             intent = pendingIntent,
             dismissShade = true,
-            customMessage = customMessage
+            customMessage = customMessage,
         )
         mainExecutor.runAllReady()
 
@@ -296,7 +355,7 @@ class LegacyActivityStarterInternalImplTest : SysuiTestCase() {
                 nullable(PendingIntent.OnFinished::class.java),
                 nullable(Handler::class.java),
                 nullable(String::class.java),
-                bundleCaptor.capture()
+                bundleCaptor.capture(),
             )
         val options = ActivityOptions.fromBundle(bundleCaptor.firstValue)
         assertThat(options.getPendingIntentBackgroundActivityStartMode())
@@ -339,7 +398,7 @@ class LegacyActivityStarterInternalImplTest : SysuiTestCase() {
             dismissShade = true,
             animationController = controller,
             showOverLockscreen = true,
-            skipLockscreenChecks = true
+            skipLockscreenChecks = true,
         )
         mainExecutor.runAllReady()
 
@@ -373,7 +432,7 @@ class LegacyActivityStarterInternalImplTest : SysuiTestCase() {
             dismissShade = true,
             animationController = controller,
             showOverLockscreen = true,
-            skipLockscreenChecks = true
+            skipLockscreenChecks = true,
         )
         mainExecutor.runAllReady()
 
@@ -413,7 +472,7 @@ class LegacyActivityStarterInternalImplTest : SysuiTestCase() {
             dismissShade = false,
             animationController = controller,
             showOverLockscreen = true,
-            skipLockscreenChecks = false
+            skipLockscreenChecks = false,
         )
         mainExecutor.runAllReady()
 
@@ -458,7 +517,7 @@ class LegacyActivityStarterInternalImplTest : SysuiTestCase() {
             dismissShade = false,
             animationController = controller,
             showOverLockscreen = true,
-            skipLockscreenChecks = false
+            skipLockscreenChecks = false,
         )
         mainExecutor.runAllReady()
 
@@ -583,7 +642,7 @@ class LegacyActivityStarterInternalImplTest : SysuiTestCase() {
             },
             {},
             false,
-            customMessage
+            customMessage,
         )
 
         verify(centralSurfaces).awakenDreams()
@@ -602,7 +661,7 @@ class LegacyActivityStarterInternalImplTest : SysuiTestCase() {
             cancelAction = null,
             dismissShade = false,
             afterKeyguardGone = false,
-            deferred = false
+            deferred = false,
         )
 
         verify(centralSurfaces, times(1)).awakenDreams()
@@ -620,7 +679,7 @@ class LegacyActivityStarterInternalImplTest : SysuiTestCase() {
             cancelAction = null,
             dismissShade = false,
             afterKeyguardGone = false,
-            deferred = false
+            deferred = false,
         )
 
         verify(centralSurfaces, never()).awakenDreams()
diff --git a/packages/SystemUI/plugin/src/com/android/systemui/plugins/ActivityStarter.java b/packages/SystemUI/plugin/src/com/android/systemui/plugins/ActivityStarter.java
index abb721ab0dd8..55be9f79e598 100644
--- a/packages/SystemUI/plugin/src/com/android/systemui/plugins/ActivityStarter.java
+++ b/packages/SystemUI/plugin/src/com/android/systemui/plugins/ActivityStarter.java
@@ -16,6 +16,7 @@ package com.android.systemui.plugins;
 
 import android.annotation.Nullable;
 import android.app.PendingIntent;
+import android.content.ComponentName;
 import android.content.Intent;
 import android.os.Bundle;
 import android.os.UserHandle;
@@ -33,6 +34,22 @@ import com.android.systemui.plugins.annotations.ProvidesInterface;
 public interface ActivityStarter {
     int VERSION = 2;
 
+    /**
+     * Registers the given {@link ActivityTransitionAnimator.ControllerFactory} for launching and
+     * closing transitions matching the {@link ActivityTransitionAnimator.TransitionCookie} and the
+     * {@link ComponentName} that it contains.
+     */
+    void registerTransition(
+            ActivityTransitionAnimator.TransitionCookie cookie,
+            ActivityTransitionAnimator.ControllerFactory controllerFactory);
+
+    /**
+     * Unregisters the {@link ActivityTransitionAnimator.ControllerFactory} previously registered
+     * containing the given {@link ActivityTransitionAnimator.TransitionCookie}. If no such
+     * registration exists, this is a no-op.
+     */
+    void unregisterTransition(ActivityTransitionAnimator.TransitionCookie cookie);
+
     void startPendingIntentDismissingKeyguard(PendingIntent intent);
 
     /**
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/dagger/CentralSurfacesDependenciesModule.java b/packages/SystemUI/src/com/android/systemui/statusbar/dagger/CentralSurfacesDependenciesModule.java
index 7df7ef187e26..254b792f8152 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/dagger/CentralSurfacesDependenciesModule.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/dagger/CentralSurfacesDependenciesModule.java
@@ -63,6 +63,7 @@ import com.android.systemui.statusbar.phone.ui.StatusBarIconController;
 import com.android.systemui.statusbar.phone.ui.StatusBarIconControllerImpl;
 import com.android.systemui.statusbar.phone.ui.StatusBarIconList;
 import com.android.systemui.statusbar.policy.KeyguardStateController;
+import com.android.wm.shell.shared.ShellTransitions;
 
 import dagger.Binds;
 import dagger.Lazy;
@@ -214,8 +215,8 @@ public interface CentralSurfacesDependenciesModule {
     @Provides
     @SysUISingleton
     static ActivityTransitionAnimator provideActivityTransitionAnimator(
-            @Main Executor mainExecutor) {
-        return new ActivityTransitionAnimator(mainExecutor);
+            @Main Executor mainExecutor, ShellTransitions shellTransitions) {
+        return new ActivityTransitionAnimator(mainExecutor, shellTransitions);
     }
 
     /** */
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/ActivityStarterImpl.kt b/packages/SystemUI/src/com/android/systemui/statusbar/phone/ActivityStarterImpl.kt
index 86c7c6bd1e86..4751293a16cc 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/ActivityStarterImpl.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/ActivityStarterImpl.kt
@@ -20,6 +20,7 @@ import android.os.Bundle
 import android.os.UserHandle
 import android.view.View
 import com.android.systemui.animation.ActivityTransitionAnimator
+import com.android.systemui.animation.TransitionAnimator
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.dagger.qualifiers.Main
 import com.android.systemui.plugins.ActivityStarter
@@ -38,7 +39,7 @@ constructor(
     private val statusBarStateController: SysuiStatusBarStateController,
     @Main private val mainExecutor: DelayableExecutor,
     activityStarterInternal: Lazy<ActivityStarterInternalImpl>,
-    legacyActivityStarter: Lazy<LegacyActivityStarterInternalImpl>
+    legacyActivityStarter: Lazy<LegacyActivityStarterInternalImpl>,
 ) : ActivityStarter {
 
     private val activityStarterInternal: ActivityStarterInternal =
@@ -48,10 +49,23 @@ constructor(
             legacyActivityStarter.get()
         }
 
+    override fun registerTransition(
+        cookie: ActivityTransitionAnimator.TransitionCookie,
+        controllerFactory: ActivityTransitionAnimator.ControllerFactory,
+    ) {
+        if (!TransitionAnimator.longLivedReturnAnimationsEnabled()) return
+        activityStarterInternal.registerTransition(cookie, controllerFactory)
+    }
+
+    override fun unregisterTransition(cookie: ActivityTransitionAnimator.TransitionCookie) {
+        if (!TransitionAnimator.longLivedReturnAnimationsEnabled()) return
+        activityStarterInternal.unregisterTransition(cookie)
+    }
+
     override fun startPendingIntentDismissingKeyguard(intent: PendingIntent) {
         activityStarterInternal.startPendingIntentDismissingKeyguard(
             intent = intent,
-            dismissShade = true
+            dismissShade = true,
         )
     }
 
@@ -98,7 +112,7 @@ constructor(
         intentSentUiThreadCallback: Runnable?,
         animationController: ActivityTransitionAnimator.Controller?,
         fillInIntent: Intent?,
-        extraOptions: Bundle?
+        extraOptions: Bundle?,
     ) {
         activityStarterInternal.startPendingIntentDismissingKeyguard(
             intent = intent,
@@ -115,7 +129,7 @@ constructor(
     override fun startPendingIntentMaybeDismissingKeyguard(
         intent: PendingIntent,
         intentSentUiThreadCallback: Runnable?,
-        animationController: ActivityTransitionAnimator.Controller?
+        animationController: ActivityTransitionAnimator.Controller?,
     ) {
         activityStarterInternal.startPendingIntentDismissingKeyguard(
             intent = intent,
@@ -245,7 +259,7 @@ constructor(
 
     override fun postStartActivityDismissingKeyguard(
         intent: PendingIntent,
-        animationController: ActivityTransitionAnimator.Controller?
+        animationController: ActivityTransitionAnimator.Controller?,
     ) {
         postOnUiThread {
             activityStarterInternal.startPendingIntentDismissingKeyguard(
@@ -381,7 +395,7 @@ constructor(
         postOnUiThread {
             statusBarStateController.setLeaveOpenOnKeyguardHide(true)
             activityStarterInternal.executeRunnableDismissingKeyguard(
-                runnable = { runnable?.let { postOnUiThread(runnable = it) } },
+                runnable = { runnable?.let { postOnUiThread(runnable = it) } }
             )
         }
     }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/ActivityStarterInternal.kt b/packages/SystemUI/src/com/android/systemui/statusbar/phone/ActivityStarterInternal.kt
index 93ce6e8561ac..5e427fbf1f7e 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/ActivityStarterInternal.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/ActivityStarterInternal.kt
@@ -17,6 +17,7 @@
 package com.android.systemui.statusbar.phone
 
 import android.app.PendingIntent
+import android.content.ComponentName
 import android.content.Intent
 import android.os.Bundle
 import android.os.UserHandle
@@ -27,6 +28,21 @@ import com.android.systemui.plugins.ActivityStarter
 
 interface ActivityStarterInternal {
     /**
+     * Registers the given [controllerFactory] for launching and closing transitions matching the
+     * [cookie] and the [ComponentName] that it contains.
+     */
+    fun registerTransition(
+        cookie: ActivityTransitionAnimator.TransitionCookie,
+        controllerFactory: ActivityTransitionAnimator.ControllerFactory,
+    )
+
+    /**
+     * Unregisters the [ActivityTransitionAnimator.Controller] previously registered containing the
+     * given [cookie]. If no such registration exists, this is a no-op.
+     */
+    fun unregisterTransition(cookie: ActivityTransitionAnimator.TransitionCookie)
+
+    /**
      * Starts a pending intent after dismissing keyguard.
      *
      * This can be called in a background thread (to prevent calls in [ActivityIntentHelper] in the
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/ActivityStarterInternalImpl.kt b/packages/SystemUI/src/com/android/systemui/statusbar/phone/ActivityStarterInternalImpl.kt
index f2ef2f0ab48f..33e4fed4073c 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/ActivityStarterInternalImpl.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/ActivityStarterInternalImpl.kt
@@ -36,6 +36,7 @@ import com.android.systemui.ActivityIntentHelper
 import com.android.systemui.Flags
 import com.android.systemui.animation.ActivityTransitionAnimator
 import com.android.systemui.animation.DelegateTransitionAnimatorController
+import com.android.systemui.animation.TransitionAnimator
 import com.android.systemui.assist.AssistManager
 import com.android.systemui.camera.CameraIntents
 import com.android.systemui.communal.domain.interactor.CommunalSceneInteractor
@@ -103,6 +104,44 @@ constructor(
     private val centralSurfaces: CentralSurfaces?
         get() = centralSurfacesOptLazy.get().getOrNull()
 
+    override fun registerTransition(
+        cookie: ActivityTransitionAnimator.TransitionCookie,
+        controllerFactory: ActivityTransitionAnimator.ControllerFactory,
+    ) {
+        check(TransitionAnimator.longLivedReturnAnimationsEnabled())
+
+        val factory =
+            object :
+                ActivityTransitionAnimator.ControllerFactory(
+                    controllerFactory.cookie,
+                    controllerFactory.component,
+                    controllerFactory.launchCujType,
+                    controllerFactory.returnCujType,
+                ) {
+                override fun createController(
+                    forLaunch: Boolean
+                ): ActivityTransitionAnimator.Controller {
+                    val baseController = controllerFactory.createController(forLaunch)
+                    val rootView = baseController.transitionContainer.rootView
+                    val controllerFromStatusBar: Optional<ActivityTransitionAnimator.Controller> =
+                        statusBarWindowControllerStore.defaultDisplay
+                            .wrapAnimationControllerIfInStatusBar(rootView, baseController)
+                    return if (controllerFromStatusBar.isPresent) {
+                        controllerFromStatusBar.get()
+                    } else {
+                        baseController
+                    }
+                }
+            }
+
+        activityTransitionAnimator.register(cookie, factory)
+    }
+
+    override fun unregisterTransition(cookie: ActivityTransitionAnimator.TransitionCookie) {
+        check(TransitionAnimator.longLivedReturnAnimationsEnabled())
+        activityTransitionAnimator.unregister(cookie)
+    }
+
     override fun startPendingIntentDismissingKeyguard(
         intent: PendingIntent,
         dismissShade: Boolean,
@@ -134,7 +173,7 @@ constructor(
                 (skipLockscreenChecks ||
                     activityIntentHelper.wouldPendingShowOverLockscreen(
                         intent,
-                        lockScreenUserManager.currentUserId
+                        lockScreenUserManager.currentUserId,
                     ))
 
         val animate =
@@ -190,7 +229,7 @@ constructor(
                                 null,
                                 null,
                                 null,
-                                options.toBundle()
+                                options.toBundle(),
                             )
                         }
                     },
@@ -239,7 +278,7 @@ constructor(
         animationController: ActivityTransitionAnimator.Controller?,
         customMessage: String?,
         disallowEnterPictureInPictureWhileLaunching: Boolean,
-        userHandle: UserHandle?
+        userHandle: UserHandle?,
     ) {
         if (SceneContainerFlag.isUnexpectedlyInLegacyMode()) return
         val userHandle: UserHandle = userHandle ?: getActivityUserHandle(intent)
@@ -280,7 +319,7 @@ constructor(
             activityTransitionAnimator.startIntentWithAnimation(
                 animController,
                 animate,
-                intent.getPackage()
+                intent.getPackage(),
             ) { adapter: RemoteAnimationAdapter? ->
                 val options =
                     ActivityOptions(CentralSurfaces.getActivityOptions(displayId, adapter))
@@ -359,7 +398,7 @@ constructor(
         dismissShade: Boolean,
         animationController: ActivityTransitionAnimator.Controller?,
         showOverLockscreenWhenLocked: Boolean,
-        userHandle: UserHandle?
+        userHandle: UserHandle?,
     ) {
         if (SceneContainerFlag.isUnexpectedlyInLegacyMode()) return
         val userHandle = userHandle ?: getActivityUserHandle(intent)
@@ -383,7 +422,7 @@ constructor(
             animationController != null &&
                 shouldAnimateLaunch(
                     isActivityIntent = true,
-                    showOverLockscreen = showOverLockscreenWhenLocked
+                    showOverLockscreen = showOverLockscreenWhenLocked,
                 )
 
         var controller: ActivityTransitionAnimator.Controller? = null
@@ -413,7 +452,7 @@ constructor(
             controller,
             animate,
             intent.getPackage(),
-            showOverLockscreenWhenLocked
+            showOverLockscreenWhenLocked,
         ) { adapter: RemoteAnimationAdapter? ->
             TaskStackBuilder.create(context)
                 .addNextIntent(intent)
@@ -425,7 +464,7 @@ constructor(
         action: ActivityStarter.OnDismissAction,
         cancel: Runnable?,
         afterKeyguardGone: Boolean,
-        customMessage: String?
+        customMessage: String?,
     ) {
         if (SceneContainerFlag.isUnexpectedlyInLegacyMode()) return
         Log.i(TAG, "Invoking dismissKeyguardThenExecute, afterKeyguardGone: $afterKeyguardGone")
@@ -453,7 +492,7 @@ constructor(
         afterKeyguardGone: Boolean,
         deferred: Boolean,
         willAnimateOnKeyguard: Boolean,
-        customMessage: String?
+        customMessage: String?,
     ) {
         if (SceneContainerFlag.isUnexpectedlyInLegacyMode()) return
         val onDismissAction: ActivityStarter.OnDismissAction =
@@ -482,12 +521,7 @@ constructor(
                     return willAnimateOnKeyguard
                 }
             }
-        dismissKeyguardThenExecute(
-            onDismissAction,
-            cancelAction,
-            afterKeyguardGone,
-            customMessage,
-        )
+        dismissKeyguardThenExecute(onDismissAction, cancelAction, afterKeyguardGone, customMessage)
     }
 
     override fun shouldAnimateLaunch(isActivityIntent: Boolean): Boolean {
@@ -565,7 +599,7 @@ constructor(
         val controllerFromStatusBar: Optional<ActivityTransitionAnimator.Controller> =
             statusBarWindowControllerStore.defaultDisplay.wrapAnimationControllerIfInStatusBar(
                 rootView,
-                animationController
+                animationController,
             )
         if (controllerFromStatusBar.isPresent) {
             return controllerFromStatusBar.get()
@@ -582,7 +616,7 @@ constructor(
                     notifShadeWindowControllerLazy.get(),
                     commandQueue,
                     displayId,
-                    isLaunchForActivity
+                    isLaunchForActivity,
                 )
             }
         }
@@ -596,7 +630,7 @@ constructor(
      */
     private fun wrapAnimationControllerForLockscreen(
         dismissShade: Boolean,
-        animationController: ActivityTransitionAnimator.Controller?
+        animationController: ActivityTransitionAnimator.Controller?,
     ): ActivityTransitionAnimator.Controller? {
         return animationController?.let {
             object : DelegateTransitionAnimatorController(it) {
@@ -613,7 +647,7 @@ constructor(
                         communalSceneInteractor.snapToScene(
                             newScene = CommunalScenes.Blank,
                             loggingReason = "ActivityStarterInternalImpl",
-                            delayMillis = ActivityTransitionAnimator.TIMINGS.totalDuration
+                            delayMillis = ActivityTransitionAnimator.TIMINGS.totalDuration,
                         )
                     }
                 }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/LegacyActivityStarterInternalImpl.kt b/packages/SystemUI/src/com/android/systemui/statusbar/phone/LegacyActivityStarterInternalImpl.kt
index d7cc65d22663..d7a29c36f2ce 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/LegacyActivityStarterInternalImpl.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/LegacyActivityStarterInternalImpl.kt
@@ -36,6 +36,7 @@ import com.android.systemui.ActivityIntentHelper
 import com.android.systemui.Flags.mediaLockscreenLaunchAnimation
 import com.android.systemui.animation.ActivityTransitionAnimator
 import com.android.systemui.animation.DelegateTransitionAnimatorController
+import com.android.systemui.animation.TransitionAnimator
 import com.android.systemui.assist.AssistManager
 import com.android.systemui.camera.CameraIntents
 import com.android.systemui.communal.domain.interactor.CommunalSceneInteractor
@@ -98,6 +99,44 @@ constructor(
     private val centralSurfaces: CentralSurfaces?
         get() = centralSurfacesOptLazy.get().getOrNull()
 
+    override fun registerTransition(
+        cookie: ActivityTransitionAnimator.TransitionCookie,
+        controllerFactory: ActivityTransitionAnimator.ControllerFactory,
+    ) {
+        check(TransitionAnimator.longLivedReturnAnimationsEnabled())
+
+        val factory =
+            object :
+                ActivityTransitionAnimator.ControllerFactory(
+                    controllerFactory.cookie,
+                    controllerFactory.component,
+                    controllerFactory.launchCujType,
+                    controllerFactory.returnCujType,
+                ) {
+                override fun createController(
+                    forLaunch: Boolean
+                ): ActivityTransitionAnimator.Controller {
+                    val baseController = controllerFactory.createController(forLaunch)
+                    val rootView = baseController.transitionContainer.rootView
+                    val controllerFromStatusBar: Optional<ActivityTransitionAnimator.Controller> =
+                        statusBarWindowControllerStore.defaultDisplay
+                            .wrapAnimationControllerIfInStatusBar(rootView, baseController)
+                    return if (controllerFromStatusBar.isPresent) {
+                        controllerFromStatusBar.get()
+                    } else {
+                        baseController
+                    }
+                }
+            }
+
+        activityTransitionAnimator.register(cookie, factory)
+    }
+
+    override fun unregisterTransition(cookie: ActivityTransitionAnimator.TransitionCookie) {
+        check(TransitionAnimator.longLivedReturnAnimationsEnabled())
+        activityTransitionAnimator.unregister(cookie)
+    }
+
     override fun startActivityDismissingKeyguard(
         intent: Intent,
         dismissShade: Boolean,
@@ -116,7 +155,7 @@ constructor(
         val willLaunchResolverActivity: Boolean =
             activityIntentHelper.wouldLaunchResolverActivity(
                 intent,
-                lockScreenUserManager.currentUserId
+                lockScreenUserManager.currentUserId,
             )
 
         val animate =
@@ -147,7 +186,7 @@ constructor(
             activityTransitionAnimator.startIntentWithAnimation(
                 animController,
                 animate,
-                intent.getPackage()
+                intent.getPackage(),
             ) { adapter: RemoteAnimationAdapter? ->
                 val options =
                     ActivityOptions(CentralSurfaces.getActivityOptions(displayId, adapter))
@@ -259,7 +298,7 @@ constructor(
                 (skipLockscreenChecks ||
                     activityIntentHelper.wouldPendingShowOverLockscreen(
                         intent,
-                        lockScreenUserManager.currentUserId
+                        lockScreenUserManager.currentUserId,
                     ))
 
         val animate =
@@ -317,7 +356,7 @@ constructor(
                                 null,
                                 null,
                                 null,
-                                options.toBundle()
+                                options.toBundle(),
                             )
                         }
                     },
@@ -409,7 +448,7 @@ constructor(
             controller,
             animate,
             intent.getPackage(),
-            showOverLockscreenWhenLocked
+            showOverLockscreenWhenLocked,
         ) { adapter: RemoteAnimationAdapter? ->
             TaskStackBuilder.create(context)
                 .addNextIntent(intent)
@@ -495,12 +534,7 @@ constructor(
                     return willAnimateOnKeyguard
                 }
             }
-        dismissKeyguardThenExecute(
-            onDismissAction,
-            cancelAction,
-            afterKeyguardGone,
-            customMessage,
-        )
+        dismissKeyguardThenExecute(onDismissAction, cancelAction, afterKeyguardGone, customMessage)
     }
 
     /**
@@ -528,7 +562,7 @@ constructor(
         val controllerFromStatusBar: Optional<ActivityTransitionAnimator.Controller> =
             statusBarWindowControllerStore.defaultDisplay.wrapAnimationControllerIfInStatusBar(
                 rootView,
-                animationController
+                animationController,
             )
         if (controllerFromStatusBar.isPresent) {
             return controllerFromStatusBar.get()
@@ -545,7 +579,7 @@ constructor(
                     notifShadeWindowControllerLazy.get(),
                     commandQueue,
                     displayId,
-                    isLaunchForActivity
+                    isLaunchForActivity,
                 )
             }
         }
@@ -559,7 +593,7 @@ constructor(
      */
     private fun wrapAnimationControllerForLockscreen(
         dismissShade: Boolean,
-        animationController: ActivityTransitionAnimator.Controller?
+        animationController: ActivityTransitionAnimator.Controller?,
     ): ActivityTransitionAnimator.Controller? {
         return animationController?.let {
             object : DelegateTransitionAnimatorController(it) {