Introduce TransitionKey

This CL introduces TransitionKey, which can be used to specify which
transition should be used when transitioning between scenes. It can be
specified both for user gestures through UserActionResult or for
triggered animations using MutableSceneTransitionLayoutState. Note that
there is no API for triggered animations backed by a hoisted
SceneTransitionLayoutState.

Bug: 315921512
Test: SceneTransitionLayoutStateTest
Test: SwipeToSceneTest
Flag: N/A
Change-Id: I2cdc2d78da41c8bb1b3234e09096faa15c3f265c

12 files changed, 230 insertions, 57 deletions

diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/AnimateToScene.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/AnimateToScene.kt
index 705754594903..b9b472f5ac6d 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/AnimateToScene.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/AnimateToScene.kt
@@ -30,6 +30,7 @@ import kotlinx.coroutines.launch
 internal fun CoroutineScope.animateToScene(
     layoutState: BaseSceneTransitionLayoutState,
     target: SceneKey,
+    transitionKey: TransitionKey?,
 ): TransitionState.Transition? {
     val transitionState = layoutState.transitionState
     if (transitionState.currentScene == target) {
@@ -45,7 +46,7 @@ internal fun CoroutineScope.animateToScene(
     }
 
     return when (transitionState) {
-        is TransitionState.Idle -> animate(layoutState, target)
+        is TransitionState.Idle -> animate(layoutState, target, transitionKey)
         is TransitionState.Transition -> {
             // A transition is currently running: first check whether `transition.toScene` or
             // `transition.fromScene` is the same as our target scene, in which case the transition
@@ -67,7 +68,7 @@ internal fun CoroutineScope.animateToScene(
                     // The transition is in progress: start the canned animation at the same
                     // progress as it was in.
                     // TODO(b/290184746): Also take the current velocity into account.
-                    animate(layoutState, target, startProgress = progress)
+                    animate(layoutState, target, transitionKey, startProgress = progress)
                 }
             } else if (transitionState.fromScene == target) {
                 // There is a transition from [target] to another scene: simply animate the same
@@ -82,12 +83,18 @@ internal fun CoroutineScope.animateToScene(
                     null
                 } else {
                     // TODO(b/290184746): Also take the current velocity into account.
-                    animate(layoutState, target, startProgress = progress, reversed = true)
+                    animate(
+                        layoutState,
+                        target,
+                        transitionKey,
+                        startProgress = progress,
+                        reversed = true,
+                    )
                 }
             } else {
                 // Generic interruption; the current transition is neither from or to [target].
                 // TODO(b/290930950): Better handle interruptions here.
-                animate(layoutState, target)
+                animate(layoutState, target, transitionKey)
             }
         }
     }
@@ -96,6 +103,7 @@ internal fun CoroutineScope.animateToScene(
 private fun CoroutineScope.animate(
     layoutState: BaseSceneTransitionLayoutState,
     target: SceneKey,
+    transitionKey: TransitionKey?,
     startProgress: Float = 0f,
     reversed: Boolean = false,
 ): TransitionState.Transition {
@@ -127,7 +135,7 @@ private fun CoroutineScope.animate(
     // Change the current layout state to start this new transition. This will compute the
     // TransformationSpec associated to this transition, which we need to initialize the Animatable
     // that will actually animate it.
-    layoutState.startTransition(transition)
+    layoutState.startTransition(transition, transitionKey)
 
     // The transformation now contains the spec that we should use to instantiate the Animatable.
     val animationSpec = layoutState.transformationSpec.progressSpec
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/Key.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/Key.kt
index c68519547c7e..2596d4a34b6f 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/Key.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/Key.kt
@@ -55,6 +55,7 @@ class SceneKey(
 
     // Implementation of [UserActionResult].
     override val toScene: SceneKey = this
+    override val transitionKey: TransitionKey? = null
     override val distance: UserActionDistance? = null
 
     override fun toString(): String {
@@ -104,3 +105,13 @@ class ValueKey(debugName: String, identity: Any = Object()) : Key(debugName, ide
         return "ValueKey(debugName=$debugName)"
     }
 }
+
+/**
+ * Key for a transition. This can be used to specify which transition spec should be used when
+ * starting the transition between two scenes.
+ */
+class TransitionKey(debugName: String, identity: Any = Object()) : Key(debugName, identity) {
+    override fun toString(): String {
+        return "TransitionKey(debugName=$debugName)"
+    }
+}
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneGestureHandler.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneGestureHandler.kt
index 35754d624ccd..23af5acd0ee8 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneGestureHandler.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneGestureHandler.kt
@@ -54,7 +54,7 @@ internal class SceneGestureHandler(
 
     private fun updateTransition(newTransition: SwipeTransition, force: Boolean = false) {
         if (isDrivingTransition || force) {
-            layoutState.startTransition(newTransition)
+            layoutState.startTransition(newTransition, newTransition.key)
 
             // Initialize SwipeTransition.swipeSpec. Note that this must be called right after
             // layoutState.startTransition() is called, because it computes the
@@ -237,7 +237,11 @@ internal class SceneGestureHandler(
                 }
         swipeTransition.dragOffset += acceleratedOffset
 
-        if (isNewFromScene || result.toScene != swipeTransition.toScene) {
+        if (
+            isNewFromScene ||
+                result.toScene != swipeTransition.toScene ||
+                result.transitionKey != swipeTransition.key
+        ) {
             updateTransition(
                 SwipeTransition(fromScene, result).apply {
                     this.dragOffset = swipeTransition.dragOffset
@@ -484,6 +488,7 @@ internal class SceneGestureHandler(
 
     private fun SwipeTransition(fromScene: Scene, result: UserActionResult): SwipeTransition {
         return SwipeTransition(
+            result.transitionKey,
             fromScene,
             layoutImpl.scene(result.toScene),
             computeAbsoluteDistance(fromScene, result),
@@ -491,6 +496,7 @@ internal class SceneGestureHandler(
     }
 
     internal class SwipeTransition(
+        val key: TransitionKey?,
         val _fromScene: Scene,
         val _toScene: Scene,
         /**
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitionLayout.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitionLayout.kt
index 7e0aa9c3e2b1..5258078e0e9a 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitionLayout.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitionLayout.kt
@@ -388,8 +388,9 @@ interface SwipeSourceDetector {
 /**
  * The result of performing a [UserAction].
  *
- * Note: [UserActionResult] is implemented by [SceneKey], and you can also use [withDistance] to
- * easily create a [UserActionResult] with a fixed distance:
+ * Note: [UserActionResult] is implemented by [SceneKey], so you can also use scene keys directly
+ * when defining your [UserActionResult]s.
+ *
  * ```
  * SceneTransitionLayout(...) {
  *     scene(
@@ -397,7 +398,7 @@ interface SwipeSourceDetector {
  *         userActions =
  *             mapOf(
  *                 Swipe.Right to Scene.Bar,
- *                 Swipe.Down to Scene.Doe withDistance 100.dp,
+ *                 Swipe.Down to Scene.Doe,
  *             )
  *         )
  *     ) { ... }
@@ -408,6 +409,9 @@ interface UserActionResult {
     /** The scene we should be transitioning to during the [UserAction]. */
     val toScene: SceneKey
 
+    /** The key of the transition that should be used. */
+    val transitionKey: TransitionKey?
+
     /**
      * The distance the action takes to animate from 0% to 100%.
      *
@@ -416,6 +420,32 @@ interface UserActionResult {
     val distance: UserActionDistance?
 }
 
+/** Create a [UserActionResult] to [toScene] with the given [distance] and [transitionKey]. */
+fun UserActionResult(
+    toScene: SceneKey,
+    distance: UserActionDistance? = null,
+    transitionKey: TransitionKey? = null,
+): UserActionResult {
+    return object : UserActionResult {
+        override val toScene: SceneKey = toScene
+        override val transitionKey: TransitionKey? = transitionKey
+        override val distance: UserActionDistance? = distance
+    }
+}
+
+/** Create a [UserActionResult] to [toScene] with the given fixed [distance] and [transitionKey]. */
+fun UserActionResult(
+    toScene: SceneKey,
+    distance: Dp,
+    transitionKey: TransitionKey? = null,
+): UserActionResult {
+    return UserActionResult(
+        toScene = toScene,
+        distance = FixedDistance(distance),
+        transitionKey = transitionKey,
+    )
+}
+
 interface UserActionDistance {
     /**
      * Return the **absolute** distance of the user action given the size of the scene we are
@@ -424,22 +454,6 @@ interface UserActionDistance {
     fun Density.absoluteDistance(fromSceneSize: IntSize, orientation: Orientation): Float
 }
 
-/**
- * A utility function to make it possible to define user actions with a distance using the syntax
- * `Swipe.Up to Scene.foo withDistance 100.dp`
- */
-infix fun Pair<UserAction, SceneKey>.withDistance(
-    distance: Dp
-): Pair<UserAction, UserActionResult> {
-    val scene = second
-    val distance = FixedDistance(distance)
-    return first to
-        object : UserActionResult {
-            override val toScene: SceneKey = scene
-            override val distance: UserActionDistance = distance
-        }
-}
-
 /** The user action has a fixed [absoluteDistance]. */
 private class FixedDistance(private val distance: Dp) : UserActionDistance {
     override fun Density.absoluteDistance(fromSceneSize: IntSize, orientation: Orientation): Float {
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitionLayoutState.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitionLayoutState.kt
index 956e326dc03e..aee6f9e15620 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitionLayoutState.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitionLayoutState.kt
@@ -92,6 +92,7 @@ sealed interface MutableSceneTransitionLayoutState : SceneTransitionLayoutState
     fun setTargetScene(
         targetScene: SceneKey,
         coroutineScope: CoroutineScope,
+        transitionKey: TransitionKey? = null,
     ): TransitionState.Transition?
 }
 
@@ -213,11 +214,14 @@ internal abstract class BaseSceneTransitionLayoutState(initialScene: SceneKey) :
     }
 
     /** Start a new [transition], instantly interrupting any ongoing transition if there was one. */
-    internal fun startTransition(transition: TransitionState.Transition) {
+    internal fun startTransition(
+        transition: TransitionState.Transition,
+        transitionKey: TransitionKey?,
+    ) {
         // Compute the [TransformationSpec] when the transition starts.
         transformationSpec =
             transitions
-                .transitionSpec(transition.fromScene, transition.toScene)
+                .transitionSpec(transition.fromScene, transition.toScene, key = transitionKey)
                 .transformationSpec()
 
         transitionState = transition
@@ -265,7 +269,11 @@ internal class HoistedSceneTransitionLayoutScene(
                 // Inspired by AnimateAsState.kt: let's poll the last value to avoid being one frame
                 // late.
                 val newKey = targetSceneChannel.tryReceive().getOrNull() ?: newKey
-                animateToScene(layoutState = this@HoistedSceneTransitionLayoutScene, newKey)
+                animateToScene(
+                    layoutState = this@HoistedSceneTransitionLayoutScene,
+                    target = newKey,
+                    transitionKey = null,
+                )
             }
         }
     }
@@ -278,14 +286,14 @@ internal class MutableSceneTransitionLayoutStateImpl(
 ) : MutableSceneTransitionLayoutState, BaseSceneTransitionLayoutState(initialScene) {
     override fun setTargetScene(
         targetScene: SceneKey,
-        coroutineScope: CoroutineScope
+        coroutineScope: CoroutineScope,
+        transitionKey: TransitionKey?,
     ): TransitionState.Transition? {
-        return with(this) {
-            coroutineScope.animateToScene(
-                layoutState = this@MutableSceneTransitionLayoutStateImpl,
-                target = targetScene,
-            )
-        }
+        return coroutineScope.animateToScene(
+            layoutState = this@MutableSceneTransitionLayoutStateImpl,
+            target = targetScene,
+            transitionKey = transitionKey,
+        )
     }
 
     override fun CoroutineScope.onChangeScene(scene: SceneKey) {
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitions.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitions.kt
index ac423b70fd50..b8f9359463de 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitions.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitions.kt
@@ -42,32 +42,42 @@ internal constructor(
     internal val defaultSwipeSpec: SpringSpec<Float>,
     internal val transitionSpecs: List<TransitionSpecImpl>,
 ) {
-    private val cache = mutableMapOf<SceneKey, MutableMap<SceneKey, TransitionSpecImpl>>()
+    private val cache =
+        mutableMapOf<
+            SceneKey, MutableMap<SceneKey, MutableMap<TransitionKey?, TransitionSpecImpl>>
+        >()
 
-    internal fun transitionSpec(from: SceneKey, to: SceneKey): TransitionSpecImpl {
-        return cache.getOrPut(from) { mutableMapOf() }.getOrPut(to) { findSpec(from, to) }
+    internal fun transitionSpec(
+        from: SceneKey,
+        to: SceneKey,
+        key: TransitionKey?,
+    ): TransitionSpecImpl {
+        return cache
+            .getOrPut(from) { mutableMapOf() }
+            .getOrPut(to) { mutableMapOf() }
+            .getOrPut(key) { findSpec(from, to, key) }
     }
 
-    private fun findSpec(from: SceneKey, to: SceneKey): TransitionSpecImpl {
-        val spec = transition(from, to) { it.from == from && it.to == to }
+    private fun findSpec(from: SceneKey, to: SceneKey, key: TransitionKey?): TransitionSpecImpl {
+        val spec = transition(from, to, key) { it.from == from && it.to == to }
         if (spec != null) {
             return spec
         }
 
-        val reversed = transition(from, to) { it.from == to && it.to == from }
+        val reversed = transition(from, to, key) { it.from == to && it.to == from }
         if (reversed != null) {
             return reversed.reversed()
         }
 
         val relaxedSpec =
-            transition(from, to) {
+            transition(from, to, key) {
                 (it.from == from && it.to == null) || (it.to == to && it.from == null)
             }
         if (relaxedSpec != null) {
             return relaxedSpec
         }
 
-        return transition(from, to) {
+        return transition(from, to, key) {
                 (it.from == to && it.to == null) || (it.to == from && it.from == null)
             }
             ?.reversed()
@@ -77,11 +87,12 @@ internal constructor(
     private fun transition(
         from: SceneKey,
         to: SceneKey,
+        key: TransitionKey?,
         filter: (TransitionSpecImpl) -> Boolean,
     ): TransitionSpecImpl? {
         var match: TransitionSpecImpl? = null
         transitionSpecs.fastForEach { spec ->
-            if (filter(spec)) {
+            if (spec.key == key && filter(spec)) {
                 if (match != null) {
                     error("Found multiple transition specs for transition $from => $to")
                 }
@@ -92,7 +103,7 @@ internal constructor(
     }
 
     private fun defaultTransition(from: SceneKey, to: SceneKey) =
-        TransitionSpecImpl(from, to, TransformationSpec.EmptyProvider)
+        TransitionSpecImpl(key = null, from, to, TransformationSpec.EmptyProvider)
 
     companion object {
         internal val DefaultSwipeSpec =
@@ -107,6 +118,9 @@ internal constructor(
 
 /** The definition of a transition between [from] and [to]. */
 interface TransitionSpec {
+    /** The key of this [TransitionSpec]. */
+    val key: TransitionKey?
+
     /**
      * The scene we are transitioning from. If `null`, this spec can be used to animate from any
      * scene.
@@ -164,12 +178,14 @@ interface TransformationSpec {
 }
 
 internal class TransitionSpecImpl(
+    override val key: TransitionKey?,
     override val from: SceneKey?,
     override val to: SceneKey?,
     private val transformationSpec: () -> TransformationSpecImpl,
 ) : TransitionSpec {
     override fun reversed(): TransitionSpecImpl {
         return TransitionSpecImpl(
+            key = key,
             from = to,
             to = from,
             transformationSpec = {
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/TransitionDsl.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/TransitionDsl.kt
index 04e093721824..d93911d2de42 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/TransitionDsl.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/TransitionDsl.kt
@@ -42,10 +42,14 @@ interface SceneTransitionsBuilder {
      * any scene. For the animation specification to apply only when transitioning between two
      * specific scenes, use [from] instead.
      *
+     * If [key] is not `null`, then this transition will only be used if the same key is specified
+     * when triggering the transition.
+     *
      * @see from
      */
     fun to(
         to: SceneKey,
+        key: TransitionKey? = null,
         builder: TransitionBuilder.() -> Unit = {},
     ): TransitionSpec
 
@@ -55,7 +59,8 @@ interface SceneTransitionsBuilder {
      * the destination scene via the [to] argument.
      *
      * When looking up which transition should be used when animating from scene A to scene B, we
-     * pick the single transition matching one of these predicates (in order of importance):
+     * pick the single transition with the given [key] and matching one of these predicates (in
+     * order of importance):
      * 1. from == A && to == B
      * 2. to == A && from == B, which is then treated in reverse.
      * 3. (from == A && to == null) || (from == null && to == B)
@@ -64,6 +69,7 @@ interface SceneTransitionsBuilder {
     fun from(
         from: SceneKey,
         to: SceneKey? = null,
+        key: TransitionKey? = null,
         builder: TransitionBuilder.() -> Unit = {},
     ): TransitionSpec
 }
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/TransitionDslImpl.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/TransitionDslImpl.kt
index df186a15181e..9b16d46bfcc8 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/TransitionDslImpl.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/TransitionDslImpl.kt
@@ -49,21 +49,27 @@ private class SceneTransitionsBuilderImpl : SceneTransitionsBuilder {
 
     val transitionSpecs = mutableListOf<TransitionSpecImpl>()
 
-    override fun to(to: SceneKey, builder: TransitionBuilder.() -> Unit): TransitionSpec {
-        return transition(from = null, to = to, builder)
+    override fun to(
+        to: SceneKey,
+        key: TransitionKey?,
+        builder: TransitionBuilder.() -> Unit
+    ): TransitionSpec {
+        return transition(from = null, to = to, key = key, builder)
     }
 
     override fun from(
         from: SceneKey,
         to: SceneKey?,
+        key: TransitionKey?,
         builder: TransitionBuilder.() -> Unit
     ): TransitionSpec {
-        return transition(from = from, to = to, builder)
+        return transition(from = from, to = to, key = key, builder)
     }
 
     private fun transition(
         from: SceneKey?,
         to: SceneKey?,
+        key: TransitionKey?,
         builder: TransitionBuilder.() -> Unit,
     ): TransitionSpec {
         fun transformationSpec(): TransformationSpecImpl {
@@ -75,7 +81,7 @@ private class SceneTransitionsBuilderImpl : SceneTransitionsBuilder {
             )
         }
 
-        val spec = TransitionSpecImpl(from, to, ::transformationSpec)
+        val spec = TransitionSpecImpl(key, from, to, ::transformationSpec)
         transitionSpecs.add(spec)
         return spec
     }
diff --git a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/SceneTransitionLayoutStateTest.kt b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/SceneTransitionLayoutStateTest.kt
index 9a25d43cb62d..c61917dbfb17 100644
--- a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/SceneTransitionLayoutStateTest.kt
+++ b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/SceneTransitionLayoutStateTest.kt
@@ -45,7 +45,10 @@ class SceneTransitionLayoutStateTest {
     @Test
     fun isTransitioningTo_transition() {
         val state = MutableSceneTransitionLayoutStateImpl(TestScenes.SceneA, SceneTransitions.Empty)
-        state.startTransition(transition(from = TestScenes.SceneA, to = TestScenes.SceneB))
+        state.startTransition(
+            transition(from = TestScenes.SceneA, to = TestScenes.SceneB),
+            transitionKey = null
+        )
 
         assertThat(state.isTransitioning()).isTrue()
         assertThat(state.isTransitioning(from = TestScenes.SceneA)).isTrue()
@@ -116,4 +119,43 @@ class SceneTransitionLayoutStateTest {
         testScheduler.advanceUntilIdle()
         assertThat(state.transitionState).isEqualTo(TransitionState.Idle(TestScenes.SceneB))
     }
+
+    @Test
+    fun setTargetScene_withTransitionKey() = runMonotonicClockTest {
+        val transitionkey = TransitionKey(debugName = "foo")
+        val state =
+            MutableSceneTransitionLayoutState(
+                TestScenes.SceneA,
+                transitions =
+                    transitions {
+                        from(TestScenes.SceneA, to = TestScenes.SceneB) { fade(TestElements.Foo) }
+                        from(TestScenes.SceneA, to = TestScenes.SceneB, key = transitionkey) {
+                            fade(TestElements.Foo)
+                            fade(TestElements.Bar)
+                        }
+                    },
+            )
+                as MutableSceneTransitionLayoutStateImpl
+
+        // Default transition from A to B.
+        assertThat(state.setTargetScene(TestScenes.SceneB, coroutineScope = this)).isNotNull()
+        assertThat(state.transformationSpec.transformations).hasSize(1)
+
+        // Go back to A.
+        state.setTargetScene(TestScenes.SceneA, coroutineScope = this)
+        testScheduler.advanceUntilIdle()
+        assertThat(state.currentTransition).isNull()
+        assertThat(state.transitionState.currentScene).isEqualTo(TestScenes.SceneA)
+
+        // Specific transition from A to B.
+        assertThat(
+                state.setTargetScene(
+                    TestScenes.SceneB,
+                    coroutineScope = this,
+                    transitionKey = transitionkey,
+                )
+            )
+            .isNotNull()
+        assertThat(state.transformationSpec.transformations).hasSize(2)
+    }
 }
diff --git a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/SwipeToSceneTest.kt b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/SwipeToSceneTest.kt
index d81631aad13a..18b4af38e310 100644
--- a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/SwipeToSceneTest.kt
+++ b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/SwipeToSceneTest.kt
@@ -384,7 +384,13 @@ class SwipeToSceneTest {
                 scene(
                     TestScenes.SceneA,
                     userActions =
-                        mapOf(Swipe.Down to TestScenes.SceneB withDistance verticalSwipeDistance),
+                        mapOf(
+                            Swipe.Down to
+                                UserActionResult(
+                                    toScene = TestScenes.SceneB,
+                                    distance = verticalSwipeDistance,
+                                )
+                        ),
                 ) {
                     Spacer(Modifier.fillMaxSize())
                 }
@@ -475,4 +481,54 @@ class SwipeToSceneTest {
         }
         assertThat(layoutState.currentTransition).isNotNull()
     }
+
+    @Test
+    fun transitionKey() {
+        val transitionkey = TransitionKey(debugName = "foo")
+        val state =
+            MutableSceneTransitionLayoutState(
+                TestScenes.SceneA,
+                transitions {
+                    from(TestScenes.SceneA, to = TestScenes.SceneB) { fade(TestElements.Foo) }
+                    from(TestScenes.SceneA, to = TestScenes.SceneB, key = transitionkey) {
+                        fade(TestElements.Foo)
+                        fade(TestElements.Bar)
+                    }
+                }
+            )
+                as MutableSceneTransitionLayoutStateImpl
+
+        var touchSlop = 0f
+        rule.setContent {
+            touchSlop = LocalViewConfiguration.current.touchSlop
+            SceneTransitionLayout(state, Modifier.size(LayoutWidth, LayoutHeight)) {
+                scene(
+                    TestScenes.SceneA,
+                    userActions =
+                        mapOf(
+                            Swipe.Down to TestScenes.SceneB,
+                            Swipe.Up to
+                                UserActionResult(TestScenes.SceneB, transitionKey = transitionkey)
+                        )
+                ) {
+                    Box(Modifier.fillMaxSize())
+                }
+                scene(TestScenes.SceneB) { Box(Modifier.fillMaxSize()) }
+            }
+        }
+
+        // Swipe down for the default transition from A to B.
+        rule.onRoot().performTouchInput {
+            down(middle)
+            moveBy(Offset(0f, touchSlop), delayMillis = 1_000)
+        }
+
+        assertThat(state.isTransitioning(from = TestScenes.SceneA, to = TestScenes.SceneB)).isTrue()
+        assertThat(state.transformationSpec.transformations).hasSize(1)
+
+        // Move the pointer up to swipe to scene B using the new transition.
+        rule.onRoot().performTouchInput { moveBy(Offset(0f, -1.dp.toPx()), delayMillis = 1_000) }
+        assertThat(state.isTransitioning(from = TestScenes.SceneA, to = TestScenes.SceneB)).isTrue()
+        assertThat(state.transformationSpec.transformations).hasSize(2)
+    }
 }
diff --git a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/TransitionDslTest.kt b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/TransitionDslTest.kt
index 140bacaec738..1beafcc2cbb2 100644
--- a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/TransitionDslTest.kt
+++ b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/TransitionDslTest.kt
@@ -177,7 +177,7 @@ class TransitionDslTest {
         // to B we defined.
         val transformations =
             transitions
-                .transitionSpec(from = TestScenes.SceneB, to = TestScenes.SceneA)
+                .transitionSpec(from = TestScenes.SceneB, to = TestScenes.SceneA, key = null)
                 .transformationSpec()
                 .transformations
 
@@ -207,7 +207,7 @@ class TransitionDslTest {
         // A => B does not have a custom spec.
         assertThat(
                 transitions
-                    .transitionSpec(from = TestScenes.SceneA, to = TestScenes.SceneB)
+                    .transitionSpec(from = TestScenes.SceneA, to = TestScenes.SceneB, key = null)
                     .transformationSpec()
                     .swipeSpec
             )
@@ -216,7 +216,7 @@ class TransitionDslTest {
         // A => C has a custom swipe spec.
         assertThat(
                 transitions
-                    .transitionSpec(from = TestScenes.SceneA, to = TestScenes.SceneC)
+                    .transitionSpec(from = TestScenes.SceneA, to = TestScenes.SceneC, key = null)
                     .transformationSpec()
                     .swipeSpec
             )
diff --git a/packages/SystemUI/compose/scene/tests/utils/src/com/android/compose/animation/scene/TestTransition.kt b/packages/SystemUI/compose/scene/tests/utils/src/com/android/compose/animation/scene/TestTransition.kt
index 6d6d57506a3a..2d71a6e50ac2 100644
--- a/packages/SystemUI/compose/scene/tests/utils/src/com/android/compose/animation/scene/TestTransition.kt
+++ b/packages/SystemUI/compose/scene/tests/utils/src/com/android/compose/animation/scene/TestTransition.kt
@@ -89,7 +89,7 @@ fun ComposeContentTestRule.testTransition(
             SceneTransitionLayout(
                 currentScene,
                 onChangeScene,
-                transitions { from(fromScene, to = toScene, transition) },
+                transitions { from(fromScene, to = toScene, builder = transition) },
                 layoutModifier,
             ) {
                 scene(fromScene, content = fromSceneContent)