Introduce STLState.currentTransitions: List<TransitionState>

This CL adds the foundation for interruptions in STL: multiple
transitions can now run in parallel. This means that the state is now
either Idle or List<Transition>. However, for backward compatibility,
the old API is preserved and this CL exposes a new currentTransitions
list that can be used by consumers that want to handle interruptions.
For the current callers, STLState.currentTransition still represents the
last/current transition and the behavior is unchanged.

This is mostly a pure refactoring that should not have any impact on
current usages, but there is an important change: transitions now *have
to* call STLState.finishTransition() once they are finished, even if
they are not the last/current transition. This change only impacts
library STL code and not consumer code given that user code is not
supposed to create custom Transitions (yet).

This CL introduces a flag to disable interruptions. This flag is enabled
by default so that the tests cover the new code, but it is explicitly
disabled in current production usages (Flexiglass, Bouncer, Lockscreen,
Communal). It is enabled by default in the STL demo app.

See b/290930950#comment5 for more information about the upcoming
interruptions support.

Bug: 290930950
Test: atest SceneTransitionLayoutStateTest
Test: Ran all STL tests with both values for the interruptions flag
Flag: N/A

Change-Id: Ia6669cc6d305b2f17ee9a805082f0b7bda56f06e

12 files changed, 365 insertions, 39 deletions

diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/bouncer/ui/composable/BouncerContent.kt b/packages/SystemUI/compose/features/src/com/android/systemui/bouncer/ui/composable/BouncerContent.kt
index 621ddf796f58..85f03c95bc65 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/bouncer/ui/composable/BouncerContent.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/bouncer/ui/composable/BouncerContent.kt
@@ -480,6 +480,7 @@ private fun FoldAware(
         onChangeScene = {},
         transitions = SceneTransitions,
         modifier = modifier,
+        enableInterruptions = false,
     ) {
         scene(SceneKeys.ContiguousSceneKey) {
             FoldableScene(
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/communal/ui/compose/CommunalContainer.kt b/packages/SystemUI/compose/features/src/com/android/systemui/communal/ui/compose/CommunalContainer.kt
index d0c498475d0b..a1d8c29c2a39 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/communal/ui/compose/CommunalContainer.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/communal/ui/compose/CommunalContainer.kt
@@ -71,6 +71,7 @@ fun CommunalContainer(
             currentScene,
             onChangeScene = { viewModel.onSceneChanged(it) },
             transitions = sceneTransitions,
+            enableInterruptions = false,
         )
     val touchesAllowed by viewModel.touchesAllowed.collectAsState(initial = false)
 
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/keyguard/ui/composable/LockscreenContent.kt b/packages/SystemUI/compose/features/src/com/android/systemui/keyguard/ui/composable/LockscreenContent.kt
index bc4e55505579..1178cc843d60 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/keyguard/ui/composable/LockscreenContent.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/keyguard/ui/composable/LockscreenContent.kt
@@ -74,6 +74,7 @@ constructor(
             transitions =
                 transitions { sceneKeyByBlueprintId.values.forEach { sceneKey -> to(sceneKey) } },
             modifier = modifier,
+            enableInterruptions = false,
         ) {
             sceneKeyByBlueprint.entries.forEach { (blueprint, sceneKey) ->
                 scene(sceneKey) { with(blueprint) { Content(Modifier.fillMaxSize()) } }
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/keyguard/ui/composable/section/TopAreaSection.kt b/packages/SystemUI/compose/features/src/com/android/systemui/keyguard/ui/composable/section/TopAreaSection.kt
index 763584182c97..d72d5cad31b4 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/keyguard/ui/composable/section/TopAreaSection.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/keyguard/ui/composable/section/TopAreaSection.kt
@@ -92,6 +92,7 @@ constructor(
             currentScene = currentScene,
             onChangeScene = {},
             transitions = ClockTransition.defaultClockTransitions,
+            enableInterruptions = false,
         ) {
             scene(ClockScenes.splitShadeLargeClockScene) {
                 Row(
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/SceneContainer.kt b/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/SceneContainer.kt
index 0fdaabe75306..fe6701cc8d89 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/SceneContainer.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/SceneContainer.kt
@@ -79,6 +79,7 @@ fun SceneContainer(
             initialScene = currentSceneKey,
             canChangeScene = { toScene -> viewModel.canChangeScene(toScene) },
             transitions = SceneContainerTransitions,
+            enableInterruptions = false,
         )
     }
 
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/DraggableHandler.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/DraggableHandler.kt
index 54249445d223..1b0627576af7 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/DraggableHandler.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/DraggableHandler.kt
@@ -18,7 +18,6 @@
 
 package com.android.compose.animation.scene
 
-import android.util.Log
 import androidx.compose.animation.core.Animatable
 import androidx.compose.animation.core.AnimationVector1D
 import androidx.compose.foundation.gestures.Orientation
@@ -144,16 +143,6 @@ internal class DraggableHandlerImpl(
         }
 
         val transitionState = layoutImpl.state.transitionState
-        if (transitionState is TransitionState.Transition) {
-            // TODO(b/290184746): Better handle interruptions here if state != idle.
-            Log.w(
-                TAG,
-                "start from TransitionState.Transition is not fully supported: from" +
-                    " ${transitionState.fromScene} to ${transitionState.toScene} " +
-                    "(progress ${transitionState.progress})"
-            )
-        }
-
         val fromScene = layoutImpl.scene(transitionState.currentScene)
         val swipes = computeSwipes(fromScene, startedPosition, pointersDown)
         val result =
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 b7e2dd13f321..c13eda22e470 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
@@ -96,9 +96,17 @@ fun SceneTransitionLayout(
     modifier: Modifier = Modifier,
     swipeSourceDetector: SwipeSourceDetector = DefaultEdgeDetector,
     @FloatRange(from = 0.0, to = 0.5) transitionInterceptionThreshold: Float = 0f,
+    enableInterruptions: Boolean = DEFAULT_INTERRUPTIONS_ENABLED,
     scenes: SceneTransitionLayoutScope.() -> Unit,
 ) {
-    val state = updateSceneTransitionLayoutState(currentScene, onChangeScene, transitions)
+    val state =
+        updateSceneTransitionLayoutState(
+            currentScene,
+            onChangeScene,
+            transitions,
+            enableInterruptions = enableInterruptions,
+        )
+
     SceneTransitionLayout(
         state,
         modifier,
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 ac2d82e024d4..f13c016e9d68 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
@@ -16,15 +16,16 @@
 
 package com.android.compose.animation.scene
 
+import android.util.Log
+import androidx.annotation.VisibleForTesting
 import androidx.compose.foundation.gestures.Orientation
 import androidx.compose.runtime.Composable
 import androidx.compose.runtime.LaunchedEffect
 import androidx.compose.runtime.SideEffect
 import androidx.compose.runtime.Stable
-import androidx.compose.runtime.getValue
-import androidx.compose.runtime.mutableStateOf
 import androidx.compose.runtime.remember
-import androidx.compose.runtime.setValue
+import androidx.compose.runtime.snapshots.SnapshotStateList
+import androidx.compose.ui.util.fastAll
 import androidx.compose.ui.util.fastFilter
 import androidx.compose.ui.util.fastForEach
 import com.android.compose.animation.scene.transition.link.LinkedTransition
@@ -50,10 +51,21 @@ sealed interface SceneTransitionLayoutState {
      */
     val transitionState: TransitionState
 
-    /** The current transition, or `null` if we are idle. */
+    /**
+     * The current transition, or `null` if we are idle.
+     *
+     * Note: If you need to handle interruptions and multiple transitions running in parallel, use
+     * [currentTransitions] instead.
+     */
     val currentTransition: TransitionState.Transition?
         get() = transitionState as? TransitionState.Transition
 
+    /**
+     * The list of [TransitionState.Transition] currently running. This will be the empty list if we
+     * are idle.
+     */
+    val currentTransitions: List<TransitionState.Transition>
+
     /** The [SceneTransitions] used when animating this state. */
     val transitions: SceneTransitions
 
@@ -120,12 +132,14 @@ fun MutableSceneTransitionLayoutState(
     transitions: SceneTransitions = SceneTransitions.Empty,
     canChangeScene: (SceneKey) -> Boolean = { true },
     stateLinks: List<StateLink> = emptyList(),
+    enableInterruptions: Boolean = DEFAULT_INTERRUPTIONS_ENABLED,
 ): MutableSceneTransitionLayoutState {
     return MutableSceneTransitionLayoutStateImpl(
         initialScene,
         transitions,
         canChangeScene,
         stateLinks,
+        enableInterruptions,
     )
 }
 
@@ -154,6 +168,7 @@ fun updateSceneTransitionLayoutState(
     transitions: SceneTransitions = SceneTransitions.Empty,
     canChangeScene: (SceneKey) -> Boolean = { true },
     stateLinks: List<StateLink> = emptyList(),
+    enableInterruptions: Boolean = DEFAULT_INTERRUPTIONS_ENABLED,
 ): SceneTransitionLayoutState {
     return remember {
             HoistedSceneTransitionLayoutState(
@@ -162,9 +177,19 @@ fun updateSceneTransitionLayoutState(
                 onChangeScene,
                 canChangeScene,
                 stateLinks,
+                enableInterruptions,
+            )
+        }
+        .apply {
+            update(
+                currentScene,
+                onChangeScene,
+                canChangeScene,
+                transitions,
+                stateLinks,
+                enableInterruptions,
             )
         }
-        .apply { update(currentScene, onChangeScene, canChangeScene, transitions, stateLinks) }
 }
 
 @Stable
@@ -302,13 +327,42 @@ sealed interface TransitionState {
 internal abstract class BaseSceneTransitionLayoutState(
     initialScene: SceneKey,
     protected var stateLinks: List<StateLink>,
+
+    // TODO(b/290930950): Remove this flag.
+    internal var enableInterruptions: Boolean,
 ) : SceneTransitionLayoutState {
-    override var transitionState: TransitionState by
-        mutableStateOf(TransitionState.Idle(initialScene))
-        protected set
+    /**
+     * The current [TransitionState]. This list will either be:
+     * 1. A list with a single [TransitionState.Idle] element, when we are idle.
+     * 2. A list with one or more [TransitionState.Transition], when we are transitioning.
+     */
+    @VisibleForTesting
+    internal val transitionStates: MutableList<TransitionState> =
+        SnapshotStateList<TransitionState>().apply { add(TransitionState.Idle(initialScene)) }
+
+    override val transitionState: TransitionState
+        get() = transitionStates.last()
 
     private val activeTransitionLinks = mutableMapOf<StateLink, LinkedTransition>()
 
+    override val currentTransitions: List<TransitionState.Transition>
+        get() {
+            if (transitionStates.last() is TransitionState.Idle) {
+                check(transitionStates.size == 1)
+                return emptyList()
+            } else {
+                @Suppress("UNCHECKED_CAST")
+                return transitionStates as List<TransitionState.Transition>
+            }
+        }
+
+    /**
+     * The mapping of transitions that are finished, i.e. for which [finishTransition] was called,
+     * to their idle scene.
+     */
+    @VisibleForTesting
+    internal val finishedTransitions = mutableMapOf<TransitionState.Transition, SceneKey>()
+
     /** Whether we can transition to the given [scene]. */
     internal abstract fun canChangeScene(scene: SceneKey): Boolean
 
@@ -330,7 +384,11 @@ internal abstract class BaseSceneTransitionLayoutState(
         return transition.isTransitioningBetween(scene, other)
     }
 
-    /** Start a new [transition], instantly interrupting any ongoing transition if there was one. */
+    /**
+     * Start a new [transition], instantly interrupting any ongoing transition if there was one.
+     *
+     * Important: you *must* call [finishTransition] once the transition is finished.
+     */
     internal fun startTransition(
         transition: TransitionState.Transition,
         transitionKey: TransitionKey?,
@@ -356,8 +414,64 @@ internal abstract class BaseSceneTransitionLayoutState(
         cancelActiveTransitionLinks()
         setupTransitionLinks(transition)
 
-        // Set the current transition.
-        transitionState = transition
+        if (!enableInterruptions) {
+            // Set the current transition.
+            check(transitionStates.size == 1)
+            transitionStates[0] = transition
+            return
+        }
+
+        when (val currentState = transitionStates.last()) {
+            is TransitionState.Idle -> {
+                // Replace [Idle] by [transition].
+                check(transitionStates.size == 1)
+                transitionStates[0] = transition
+            }
+            is TransitionState.Transition -> {
+                // Force the current transition to finish to currentScene.
+                currentState.finish().invokeOnCompletion {
+                    // Make sure [finishTransition] is called at the end of the transition.
+                    finishTransition(currentState, currentState.currentScene)
+                }
+
+                // Check that we don't have too many concurrent transitions.
+                if (transitionStates.size >= MAX_CONCURRENT_TRANSITIONS) {
+                    Log.wtf(
+                        TAG,
+                        buildString {
+                            appendLine("Potential leak detected in SceneTransitionLayoutState!")
+                            appendLine(
+                                "  Some transition(s) never called STLState.finishTransition()."
+                            )
+                            appendLine("  Transitions (size=${transitionStates.size}):")
+                            transitionStates.fastForEach { state ->
+                                val transition = state as TransitionState.Transition
+                                val from = transition.fromScene
+                                val to = transition.toScene
+                                val indicator =
+                                    if (finishedTransitions.contains(transition)) "x" else " "
+                                appendLine("  [$indicator] $from => $to ($transition)")
+                            }
+                        }
+                    )
+
+                    // Force finish all transitions.
+                    while (currentTransitions.isNotEmpty()) {
+                        val transition = transitionStates[0] as TransitionState.Transition
+                        finishTransition(transition, transition.currentScene)
+                    }
+
+                    // We finished all transitions, so we are now idle. We remove this state so that
+                    // we end up only with the new transition after appending it.
+                    check(transitionStates.size == 1)
+                    check(transitionStates[0] is TransitionState.Idle)
+                    transitionStates.clear()
+                }
+
+                // Append the new transition.
+                transitionStates.add(transition)
+            }
+        }
     }
 
     private fun cancelActiveTransitionLinks() {
@@ -397,13 +511,54 @@ internal abstract class BaseSceneTransitionLayoutState(
      * nothing if [transition] was interrupted since it was started.
      */
     internal fun finishTransition(transition: TransitionState.Transition, idleScene: SceneKey) {
-        resolveActiveTransitionLinks(idleScene)
-        if (transitionState == transition) {
-            transitionState = TransitionState.Idle(idleScene)
+        val existingIdleScene = finishedTransitions[transition]
+        if (existingIdleScene != null) {
+            // This transition was already finished.
+            check(idleScene == existingIdleScene) {
+                "Transition $transition was finished multiple times with different " +
+                    "idleScene ($existingIdleScene != $idleScene)"
+            }
+            return
+        }
+
+        if (!transitionStates.contains(transition)) {
+            // This transition was already removed from transitionStates.
+            return
+        }
+
+        check(transitionStates.fastAll { it is TransitionState.Transition })
+
+        // Mark this transition as finished and save the scene it is settling at.
+        finishedTransitions[transition] = idleScene
+
+        // Finish all linked transitions.
+        finishActiveTransitionLinks(idleScene)
+
+        // Keep a reference to the idle scene of the last removed transition, in case we remove all
+        // transitions and should settle to Idle.
+        var lastRemovedIdleScene: SceneKey? = null
+
+        // Remove all first n finished transitions.
+        while (transitionStates.isNotEmpty()) {
+            val firstTransition = transitionStates[0]
+            if (!finishedTransitions.contains(firstTransition)) {
+                // Stop here.
+                break
+            }
+
+            // Remove the transition from the list and from the set of finished transitions.
+            transitionStates.removeAt(0)
+            lastRemovedIdleScene = finishedTransitions.remove(firstTransition)
+        }
+
+        // If all transitions are finished, we are idle.
+        if (transitionStates.isEmpty()) {
+            check(finishedTransitions.isEmpty())
+            transitionStates.add(TransitionState.Idle(checkNotNull(lastRemovedIdleScene)))
         }
     }
 
-    private fun resolveActiveTransitionLinks(idleScene: SceneKey) {
+    private fun finishActiveTransitionLinks(idleScene: SceneKey) {
         val previousTransition = this.transitionState as? TransitionState.Transition ?: return
         for ((link, linkedTransition) in activeTransitionLinks) {
             if (previousTransition.fromScene == idleScene) {
@@ -424,20 +579,39 @@ internal abstract class BaseSceneTransitionLayoutState(
      * Check if a transition is in progress. If the progress value is near 0 or 1, immediately snap
      * to the closest scene.
      *
+     * Important: Snapping to the closest scene will instantly finish *all* ongoing transitions,
+     * only the progress of the last transition will be checked.
+     *
      * @return true if snapped to the closest scene.
      */
     internal fun snapToIdleIfClose(threshold: Float): Boolean {
         val transition = currentTransition ?: return false
         val progress = transition.progress
+
         fun isProgressCloseTo(value: Float) = (progress - value).absoluteValue <= threshold
 
+        fun finishAllTransitions(lastTransitionIdleScene: SceneKey) {
+            // Force finish all transitions.
+            while (currentTransitions.isNotEmpty()) {
+                val transition = transitionStates[0] as TransitionState.Transition
+                val idleScene =
+                    if (transitionStates.size == 1) {
+                        lastTransitionIdleScene
+                    } else {
+                        transition.currentScene
+                    }
+
+                finishTransition(transition, idleScene)
+            }
+        }
+
         return when {
             isProgressCloseTo(0f) -> {
-                finishTransition(transition, transition.fromScene)
+                finishAllTransitions(transition.fromScene)
                 true
             }
             isProgressCloseTo(1f) -> {
-                finishTransition(transition, transition.toScene)
+                finishAllTransitions(transition.toScene)
                 true
             }
             else -> false
@@ -455,7 +629,8 @@ internal class HoistedSceneTransitionLayoutState(
     private var changeScene: (SceneKey) -> Unit,
     private var canChangeScene: (SceneKey) -> Boolean,
     stateLinks: List<StateLink> = emptyList(),
-) : BaseSceneTransitionLayoutState(initialScene, stateLinks) {
+    enableInterruptions: Boolean = DEFAULT_INTERRUPTIONS_ENABLED,
+) : BaseSceneTransitionLayoutState(initialScene, stateLinks, enableInterruptions) {
     private val targetSceneChannel = Channel<SceneKey>(Channel.CONFLATED)
 
     override fun canChangeScene(scene: SceneKey): Boolean = canChangeScene.invoke(scene)
@@ -469,12 +644,14 @@ internal class HoistedSceneTransitionLayoutState(
         canChangeScene: (SceneKey) -> Boolean,
         transitions: SceneTransitions,
         stateLinks: List<StateLink>,
+        enableInterruptions: Boolean,
     ) {
         SideEffect {
             this.changeScene = onChangeScene
             this.canChangeScene = canChangeScene
             this.transitions = transitions
             this.stateLinks = stateLinks
+            this.enableInterruptions = enableInterruptions
 
             targetSceneChannel.trySend(currentScene)
         }
@@ -500,7 +677,10 @@ internal class MutableSceneTransitionLayoutStateImpl(
     override var transitions: SceneTransitions,
     private val canChangeScene: (SceneKey) -> Boolean = { true },
     stateLinks: List<StateLink> = emptyList(),
-) : MutableSceneTransitionLayoutState, BaseSceneTransitionLayoutState(initialScene, stateLinks) {
+    enableInterruptions: Boolean = DEFAULT_INTERRUPTIONS_ENABLED,
+) :
+    MutableSceneTransitionLayoutState,
+    BaseSceneTransitionLayoutState(initialScene, stateLinks, enableInterruptions) {
     override fun setTargetScene(
         targetScene: SceneKey,
         coroutineScope: CoroutineScope,
@@ -519,3 +699,15 @@ internal class MutableSceneTransitionLayoutStateImpl(
         setTargetScene(scene, coroutineScope = this)
     }
 }
+
+private const val TAG = "SceneTransitionLayoutState"
+
+/** Whether support for interruptions in enabled by default. */
+internal const val DEFAULT_INTERRUPTIONS_ENABLED = true
+
+/**
+ * The max number of concurrent transitions. If the number of transitions goes past this number,
+ * this probably means that there is a leak and we will Log.wtf before clearing the list of
+ * transitions.
+ */
+private const val MAX_CONCURRENT_TRANSITIONS = 100
diff --git a/packages/SystemUI/compose/scene/tests/Android.bp b/packages/SystemUI/compose/scene/tests/Android.bp
index 59cc63aa5eef..af1389680bd2 100644
--- a/packages/SystemUI/compose/scene/tests/Android.bp
+++ b/packages/SystemUI/compose/scene/tests/Android.bp
@@ -26,7 +26,6 @@ android_test {
     name: "PlatformComposeSceneTransitionLayoutTests",
     manifest: "AndroidManifest.xml",
     test_suites: ["device-tests"],
-    sdk_version: "current",
     certificate: "platform",
 
     srcs: [
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 0cce99c53d46..93e94f8f95a2 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
@@ -16,6 +16,7 @@
 
 package com.android.compose.animation.scene
 
+import android.util.Log
 import androidx.compose.foundation.gestures.Orientation
 import androidx.compose.runtime.mutableStateOf
 import androidx.compose.ui.test.junit4.createComposeRule
@@ -28,8 +29,12 @@ import com.android.compose.animation.scene.transition.link.StateLink
 import com.android.compose.test.runMonotonicClockTest
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.CoroutineStart
+import kotlinx.coroutines.Job
 import kotlinx.coroutines.cancelAndJoin
 import kotlinx.coroutines.launch
+import kotlinx.coroutines.sync.Mutex
+import kotlinx.coroutines.sync.withLock
+import kotlinx.coroutines.test.runTest
 import org.junit.Rule
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -270,11 +275,21 @@ class SceneTransitionLayoutStateTest {
     }
 
     @Test
-    fun linkedTransition_startsLinkButLinkedStateIsTakenOver() {
+    fun linkedTransition_startsLinkButLinkedStateIsTakenOver() = runTest {
         val (parentState, childState) = setupLinkedStates()
 
-        val childTransition = transition(SceneA, SceneB)
-        val parentTransition = transition(SceneC, SceneA)
+        val childTransition =
+            transition(
+                SceneA,
+                SceneB,
+                onFinish = { launch { /* Do nothing. */} },
+            )
+        val parentTransition =
+            transition(
+                SceneC,
+                SceneA,
+                onFinish = { launch { /* Do nothing. */} },
+            )
         childState.startTransition(childTransition, null)
         parentState.startTransition(parentTransition, null)
 
@@ -326,7 +341,7 @@ class SceneTransitionLayoutStateTest {
     fun snapToIdleIfClose_snapToStart() = runMonotonicClockTest {
         val state = MutableSceneTransitionLayoutStateImpl(SceneA, SceneTransitions.Empty)
         state.startTransition(
-            transition(from = SceneA, to = TestScenes.SceneB, progress = { 0.2f }),
+            transition(from = SceneA, to = SceneB, progress = { 0.2f }),
             transitionKey = null
         )
         assertThat(state.isTransitioning()).isTrue()
@@ -345,7 +360,7 @@ class SceneTransitionLayoutStateTest {
     fun snapToIdleIfClose_snapToEnd() = runMonotonicClockTest {
         val state = MutableSceneTransitionLayoutStateImpl(SceneA, SceneTransitions.Empty)
         state.startTransition(
-            transition(from = SceneA, to = TestScenes.SceneB, progress = { 0.8f }),
+            transition(from = SceneA, to = SceneB, progress = { 0.8f }),
             transitionKey = null
         )
         assertThat(state.isTransitioning()).isTrue()
@@ -357,7 +372,35 @@ class SceneTransitionLayoutStateTest {
         // Go to the final scene if it is close to 1.
         assertThat(state.snapToIdleIfClose(threshold = 0.2f)).isTrue()
         assertThat(state.isTransitioning()).isFalse()
-        assertThat(state.transitionState).isEqualTo(TransitionState.Idle(TestScenes.SceneB))
+        assertThat(state.transitionState).isEqualTo(TransitionState.Idle(SceneB))
+    }
+
+    @Test
+    fun snapToIdleIfClose_multipleTransitions() = runMonotonicClockTest {
+        val state = MutableSceneTransitionLayoutStateImpl(SceneA, SceneTransitions.Empty)
+
+        val aToB =
+            transition(
+                from = SceneA,
+                to = SceneB,
+                progress = { 0.5f },
+                onFinish = { launch { /* do nothing */} },
+            )
+        state.startTransition(aToB, transitionKey = null)
+        assertThat(state.currentTransitions).containsExactly(aToB).inOrder()
+
+        val bToC = transition(from = SceneB, to = SceneC, progress = { 0.8f })
+        state.startTransition(bToC, transitionKey = null)
+        assertThat(state.currentTransitions).containsExactly(aToB, bToC).inOrder()
+
+        // Ignore the request if the progress is not close to 0 or 1, using the threshold.
+        assertThat(state.snapToIdleIfClose(threshold = 0.1f)).isFalse()
+        assertThat(state.currentTransitions).containsExactly(aToB, bToC).inOrder()
+
+        // Go to the final scene if it is close to 1.
+        assertThat(state.snapToIdleIfClose(threshold = 0.2f)).isTrue()
+        assertThat(state.transitionState).isEqualTo(TransitionState.Idle(SceneC))
+        assertThat(state.currentTransitions).isEmpty()
     }
 
     @Test
@@ -508,4 +551,82 @@ class SceneTransitionLayoutStateTest {
         progress.value = 1.1f
         assertThat(state.currentTransition?.currentOverscrollSpec).isNull()
     }
+
+    @Test
+    fun multipleTransitions() = runTest {
+        val finishingTransitions = mutableSetOf<TransitionState.Transition>()
+        fun onFinish(transition: TransitionState.Transition): Job {
+            // Instead of letting the transition finish, we put the transition in the
+            // finishingTransitions set so that we can verify that finish() is called when expected
+            // and then we call state STLState.finishTransition() ourselves.
+            finishingTransitions.add(transition)
+
+            return backgroundScope.launch {
+                // Try to acquire a locked mutex so that this code never completes.
+                Mutex(locked = true).withLock {}
+            }
+        }
+
+        val state = MutableSceneTransitionLayoutStateImpl(SceneA, EmptyTestTransitions)
+        val aToB = transition(SceneA, SceneB, onFinish = ::onFinish)
+        val bToC = transition(SceneB, SceneC, onFinish = ::onFinish)
+        val cToA = transition(SceneC, SceneA, onFinish = ::onFinish)
+
+        // Starting state.
+        assertThat(finishingTransitions).isEmpty()
+        assertThat(state.currentTransitions).isEmpty()
+
+        // A => B.
+        state.startTransition(aToB, transitionKey = null)
+        assertThat(finishingTransitions).isEmpty()
+        assertThat(state.finishedTransitions).isEmpty()
+        assertThat(state.currentTransitions).containsExactly(aToB).inOrder()
+
+        // B => C. This should automatically call finish() on aToB.
+        state.startTransition(bToC, transitionKey = null)
+        assertThat(finishingTransitions).containsExactly(aToB)
+        assertThat(state.finishedTransitions).isEmpty()
+        assertThat(state.currentTransitions).containsExactly(aToB, bToC).inOrder()
+
+        // C => A. This should automatically call finish() on bToC.
+        state.startTransition(cToA, transitionKey = null)
+        assertThat(finishingTransitions).containsExactly(aToB, bToC)
+        assertThat(state.finishedTransitions).isEmpty()
+        assertThat(state.currentTransitions).containsExactly(aToB, bToC, cToA).inOrder()
+
+        // Mark bToC as finished. The list of current transitions does not change because aToB is
+        // still not marked as finished.
+        state.finishTransition(bToC, idleScene = bToC.currentScene)
+        assertThat(state.finishedTransitions).containsExactly(bToC, bToC.currentScene)
+        assertThat(state.currentTransitions).containsExactly(aToB, bToC, cToA).inOrder()
+
+        // Mark aToB as finished. This will remove both aToB and bToC from the list of transitions.
+        state.finishTransition(aToB, idleScene = aToB.currentScene)
+        assertThat(state.finishedTransitions).isEmpty()
+        assertThat(state.currentTransitions).containsExactly(cToA).inOrder()
+    }
+
+    @Test
+    fun tooManyTransitionsLogsWtfAndClearsTransitions() = runTest {
+        val state = MutableSceneTransitionLayoutStateImpl(SceneA, EmptyTestTransitions)
+
+        fun startTransition() {
+            val transition = transition(SceneA, SceneB, onFinish = { launch { /* do nothing */} })
+            state.startTransition(transition, transitionKey = null)
+        }
+
+        var hasLoggedWtf = false
+        val originalHandler = Log.setWtfHandler { _, _, _ -> hasLoggedWtf = true }
+        try {
+            repeat(100) { startTransition() }
+            assertThat(hasLoggedWtf).isFalse()
+            assertThat(state.currentTransitions).hasSize(100)
+
+            startTransition()
+            assertThat(hasLoggedWtf).isTrue()
+            assertThat(state.currentTransitions).hasSize(1)
+        } finally {
+            Log.setWtfHandler(originalHandler)
+        }
+    }
 }
diff --git a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/SceneTransitionLayoutTest.kt b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/SceneTransitionLayoutTest.kt
index efaea71f8d2c..723a1825f205 100644
--- a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/SceneTransitionLayoutTest.kt
+++ b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/SceneTransitionLayoutTest.kt
@@ -299,6 +299,11 @@ class SceneTransitionLayoutTest {
             .isWithin(DpOffsetSubject.DefaultTolerance)
             .of(DpOffset(expectedOffset, expectedOffset))
 
+        // Wait for the transition to C to finish.
+        rule.mainClock.advanceTimeBy(TestTransitionDuration)
+        assertThat(layoutState.transitionState).isInstanceOf(TransitionState.Idle::class.java)
+        assertThat(layoutState.transitionState.currentScene).isEqualTo(TestScenes.SceneC)
+
         // Go back to scene A. This should happen instantly (once the animation started, i.e. after
         // 2 frames) given that we use a snap() animation spec.
         currentScene = TestScenes.SceneA
diff --git a/packages/SystemUI/compose/scene/tests/utils/src/com/android/compose/animation/scene/Transition.kt b/packages/SystemUI/compose/scene/tests/utils/src/com/android/compose/animation/scene/Transition.kt
index a32fe2273804..767057b585b8 100644
--- a/packages/SystemUI/compose/scene/tests/utils/src/com/android/compose/animation/scene/Transition.kt
+++ b/packages/SystemUI/compose/scene/tests/utils/src/com/android/compose/animation/scene/Transition.kt
@@ -29,6 +29,7 @@ fun transition(
     isUpOrLeft: Boolean = false,
     bouncingScene: SceneKey? = null,
     orientation: Orientation = Orientation.Horizontal,
+    onFinish: ((TransitionState.Transition) -> Job)? = null,
 ): TransitionState.Transition {
     return object : TransitionState.Transition(from, to), TransitionState.HasOverscrollProperties {
         override val currentScene: SceneKey = from
@@ -46,7 +47,13 @@ fun transition(
             }
 
         override fun finish(): Job {
-            error("finish() is not supported in test transitions")
+            val onFinish =
+                onFinish
+                    ?: error(
+                        "onFinish() must be provided if finish() is called on test transitions"
+                    )
+
+            return onFinish(this)
         }
     }
 }