[flexiglass] Add SceneTransitionLayout (STL) overlays to the framework.

This CL integrates overlays and defines a basic API, based off of
ag/28338922. Overlay transition state isn't included in this initial
version.

Bug: 359173565
Flag: com.android.systemui.scene_container
Test: Added a few unit tests, existing ones still pass. More thorough
unit tests will be added in a follow-up CL, once we've defined a couple
of overlays.

Change-Id: I0fb11d094ce2d7e14e3860b2e0f0f032e320b8d8
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/Overlay.kt b/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/Overlay.kt
new file mode 100644
index 0000000..d62befd
--- /dev/null
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/Overlay.kt
@@ -0,0 +1,37 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.scene.ui.composable
+
+import androidx.compose.runtime.Composable
+import androidx.compose.ui.Modifier
+import com.android.compose.animation.scene.ContentScope
+import com.android.compose.animation.scene.OverlayKey
+import com.android.systemui.lifecycle.Activatable
+
+/**
+ * Defines interface for classes that can describe an "overlay".
+ *
+ * In the scene framework, there can be multiple overlays in a single scene "container". The
+ * container takes care of rendering any current overlays and allowing overlays to be shown, hidden,
+ * or replaced based on a user action.
+ */
+interface Overlay : Activatable {
+    /** Uniquely-identifying key for this overlay. The key must be unique within its container. */
+    val key: OverlayKey
+
+    @Composable fun ContentScope.Content(modifier: Modifier)
+}
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 e17cb31..f9723d9 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
@@ -31,7 +31,9 @@
 import androidx.compose.ui.Modifier
 import androidx.compose.ui.graphics.Color
 import androidx.compose.ui.input.pointer.pointerInput
+import com.android.compose.animation.scene.ContentKey
 import com.android.compose.animation.scene.MutableSceneTransitionLayoutState
+import com.android.compose.animation.scene.OverlayKey
 import com.android.compose.animation.scene.SceneKey
 import com.android.compose.animation.scene.SceneTransitionLayout
 import com.android.compose.animation.scene.UserAction
@@ -57,12 +59,18 @@
  *   and only the scenes on this container. In other words: (a) there should be no scene in this map
  *   that is not in the configuration for this container and (b) all scenes in the configuration
  *   must have entries in this map.
+ * @param overlayByKey Mapping of [Overlay] by [OverlayKey], ordered by z-order such that the last
+ *   overlay is rendered on top of all other overlays. It's critical that this map contains exactly
+ *   and only the overlays on this container. In other words: (a) there should be no overlay in this
+ *   map that is not in the configuration for this container and (b) all overlays in the
+ *   configuration must have entries in this map.
  * @param modifier A modifier.
  */
 @Composable
 fun SceneContainer(
     viewModel: SceneContainerViewModel,
     sceneByKey: Map<SceneKey, ComposableScene>,
+    overlayByKey: Map<OverlayKey, Overlay>,
     initialSceneKey: SceneKey,
     dataSourceDelegator: SceneDataSourceDelegator,
     modifier: Modifier = Modifier,
@@ -89,16 +97,19 @@
         onDispose { viewModel.setTransitionState(null) }
     }
 
-    val userActionsBySceneKey: MutableMap<SceneKey, Map<UserAction, UserActionResult>> = remember {
-        mutableStateMapOf()
-    }
+    val userActionsByContentKey: MutableMap<ContentKey, Map<UserAction, UserActionResult>> =
+        remember {
+            mutableStateMapOf()
+        }
+    // TODO(b/359173565): Add overlay user actions when the API is final.
     LaunchedEffect(currentSceneKey) {
         try {
             sceneByKey[currentSceneKey]?.destinationScenes?.collectLatest { userActions ->
-                userActionsBySceneKey[currentSceneKey] = viewModel.resolveSceneFamilies(userActions)
+                userActionsByContentKey[currentSceneKey] =
+                    viewModel.resolveSceneFamilies(userActions)
             }
         } finally {
-            userActionsBySceneKey[currentSceneKey] = emptyMap()
+            userActionsByContentKey[currentSceneKey] = emptyMap()
         }
     }
 
@@ -115,7 +126,7 @@
             sceneByKey.forEach { (sceneKey, composableScene) ->
                 scene(
                     key = sceneKey,
-                    userActions = userActionsBySceneKey.getOrDefault(sceneKey, emptyMap())
+                    userActions = userActionsByContentKey.getOrDefault(sceneKey, emptyMap())
                 ) {
                     // Activate the scene.
                     LaunchedEffect(composableScene) { composableScene.activate() }
@@ -128,6 +139,15 @@
                     }
                 }
             }
+            overlayByKey.forEach { (overlayKey, composableOverlay) ->
+                overlay(
+                    key = overlayKey,
+                    userActions = userActionsByContentKey.getOrDefault(overlayKey, emptyMap())
+                ) {
+                    // Render the overlay.
+                    with(composableOverlay) { this@overlay.Content(Modifier) }
+                }
+            }
         }
 
         BottomRightCornerRibbon(
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/SceneTransitionLayoutDataSource.kt b/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/SceneTransitionLayoutDataSource.kt
index 4b4b7ed..e12a8bd 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/SceneTransitionLayoutDataSource.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/SceneTransitionLayoutDataSource.kt
@@ -18,7 +18,9 @@
 
 package com.android.systemui.scene.ui.composable
 
+import androidx.compose.runtime.snapshotFlow
 import com.android.compose.animation.scene.MutableSceneTransitionLayoutState
+import com.android.compose.animation.scene.OverlayKey
 import com.android.compose.animation.scene.SceneKey
 import com.android.compose.animation.scene.TransitionKey
 import com.android.compose.animation.scene.observableTransitionState
@@ -52,6 +54,14 @@
                 initialValue = state.transitionState.currentScene,
             )
 
+    override val currentOverlays: StateFlow<Set<OverlayKey>> =
+        snapshotFlow { state.currentOverlays }
+            .stateIn(
+                scope = coroutineScope,
+                started = SharingStarted.WhileSubscribed(),
+                initialValue = emptySet(),
+            )
+
     override fun changeScene(
         toScene: SceneKey,
         transitionKey: TransitionKey?,
@@ -68,4 +78,29 @@
             scene = toScene,
         )
     }
+
+    override fun showOverlay(overlay: OverlayKey, transitionKey: TransitionKey?) {
+        state.showOverlay(
+            overlay = overlay,
+            animationScope = coroutineScope,
+            transitionKey = transitionKey,
+        )
+    }
+
+    override fun hideOverlay(overlay: OverlayKey, transitionKey: TransitionKey?) {
+        state.hideOverlay(
+            overlay = overlay,
+            animationScope = coroutineScope,
+            transitionKey = transitionKey,
+        )
+    }
+
+    override fun replaceOverlay(from: OverlayKey, to: OverlayKey, transitionKey: TransitionKey?) {
+        state.replaceOverlay(
+            from = from,
+            to = to,
+            animationScope = coroutineScope,
+            transitionKey = transitionKey,
+        )
+    }
 }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/data/repository/SceneContainerRepositoryTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/data/repository/SceneContainerRepositoryTest.kt
index df30c4b..227b3a6 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/data/repository/SceneContainerRepositoryTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/data/repository/SceneContainerRepositoryTest.kt
@@ -25,6 +25,7 @@
 import com.android.systemui.coroutines.collectLastValue
 import com.android.systemui.flags.EnableSceneContainer
 import com.android.systemui.kosmos.testScope
+import com.android.systemui.scene.overlayKeys
 import com.android.systemui.scene.sceneContainerConfig
 import com.android.systemui.scene.sceneKeys
 import com.android.systemui.scene.shared.model.Scenes
@@ -46,19 +47,9 @@
     private val testScope = kosmos.testScope
 
     @Test
-    fun allSceneKeys() {
+    fun allContentKeys() {
         val underTest = kosmos.sceneContainerRepository
-        assertThat(underTest.allSceneKeys())
-            .isEqualTo(
-                listOf(
-                    Scenes.QuickSettings,
-                    Scenes.Shade,
-                    Scenes.Lockscreen,
-                    Scenes.Bouncer,
-                    Scenes.Gone,
-                    Scenes.Communal,
-                )
-            )
+        assertThat(underTest.allContentKeys).isEqualTo(kosmos.sceneKeys + kosmos.overlayKeys)
     }
 
     @Test
@@ -75,6 +66,18 @@
             assertThat(currentScene).isEqualTo(Scenes.QuickSettings)
         }
 
+    // TODO(b/356596436): Add tests for showing, hiding, and replacing overlays after we've defined
+    //  them.
+    @Test
+    fun currentOverlays() =
+        testScope.runTest {
+            val underTest = kosmos.sceneContainerRepository
+            val currentOverlays by collectLastValue(underTest.currentOverlays)
+            assertThat(currentOverlays).isEmpty()
+
+            // TODO(b/356596436): When we have a first overlay, add it here and assert contains.
+        }
+
     @Test(expected = IllegalStateException::class)
     fun changeScene_noSuchSceneInContainer_throws() {
         kosmos.sceneKeys = listOf(Scenes.QuickSettings, Scenes.Lockscreen)
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/domain/interactor/SceneInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/domain/interactor/SceneInteractorTest.kt
index 35cefa6..4a7d8b0 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/domain/interactor/SceneInteractorTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/domain/interactor/SceneInteractorTest.kt
@@ -35,6 +35,7 @@
 import com.android.systemui.scene.data.repository.sceneContainerRepository
 import com.android.systemui.scene.data.repository.setSceneTransition
 import com.android.systemui.scene.domain.resolver.homeSceneFamilyResolver
+import com.android.systemui.scene.overlayKeys
 import com.android.systemui.scene.sceneContainerConfig
 import com.android.systemui.scene.sceneKeys
 import com.android.systemui.scene.shared.model.SceneFamilies
@@ -72,9 +73,11 @@
         kosmos.keyguardEnabledInteractor
     }
 
+    // TODO(b/356596436): Add tests for showing, hiding, and replacing overlays after we've defined
+    //  them.
     @Test
-    fun allSceneKeys() {
-        assertThat(underTest.allSceneKeys()).isEqualTo(kosmos.sceneKeys)
+    fun allContentKeys() {
+        assertThat(underTest.allContentKeys).isEqualTo(kosmos.sceneKeys + kosmos.overlayKeys)
     }
 
     @Test
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/shared/model/SceneDataSourceDelegatorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/shared/model/SceneDataSourceDelegatorTest.kt
index 32c0172..2720c57 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/shared/model/SceneDataSourceDelegatorTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/shared/model/SceneDataSourceDelegatorTest.kt
@@ -37,6 +37,8 @@
     private val initialSceneKey = kosmos.initialSceneKey
     private val fakeSceneDataSource = kosmos.fakeSceneDataSource
 
+    // TODO(b/356596436): Add tests for showing, hiding, and replacing overlays after we've defined
+    //  them.
     private val underTest = kosmos.sceneDataSourceDelegator
 
     @Test
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/ui/viewmodel/SceneContainerViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/ui/viewmodel/SceneContainerViewModelTest.kt
index 832e7b1..3558f17 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/ui/viewmodel/SceneContainerViewModelTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/ui/viewmodel/SceneContainerViewModelTest.kt
@@ -32,7 +32,6 @@
 import com.android.systemui.power.domain.interactor.powerInteractor
 import com.android.systemui.scene.domain.interactor.sceneInteractor
 import com.android.systemui.scene.sceneContainerConfig
-import com.android.systemui.scene.sceneKeys
 import com.android.systemui.scene.shared.logger.sceneLogger
 import com.android.systemui.scene.shared.model.Scenes
 import com.android.systemui.scene.shared.model.fakeSceneDataSource
@@ -49,6 +48,7 @@
 import org.junit.Test
 import org.junit.runner.RunWith
 
+@OptIn(ExperimentalCoroutinesApi::class)
 @SmallTest
 @RunWith(AndroidJUnit4::class)
 @EnableSceneContainer
@@ -113,11 +113,6 @@
         }
 
     @Test
-    fun allSceneKeys() {
-        assertThat(underTest.allSceneKeys).isEqualTo(kosmos.sceneKeys)
-    }
-
-    @Test
     fun sceneTransition() =
         testScope.runTest {
             val currentScene by collectLastValue(underTest.currentScene)
diff --git a/packages/SystemUI/src/com/android/systemui/scene/EmptySceneModule.kt b/packages/SystemUI/src/com/android/systemui/scene/EmptySceneModule.kt
index efb9375..4c730a0 100644
--- a/packages/SystemUI/src/com/android/systemui/scene/EmptySceneModule.kt
+++ b/packages/SystemUI/src/com/android/systemui/scene/EmptySceneModule.kt
@@ -17,6 +17,7 @@
 package com.android.systemui.scene
 
 import com.android.systemui.scene.shared.model.Scene
+import com.android.systemui.scene.ui.composable.Overlay
 import dagger.Module
 import dagger.Provides
 import dagger.multibindings.ElementsIntoSet
@@ -29,4 +30,10 @@
     fun emptySceneSet(): Set<Scene> {
         return emptySet()
     }
+
+    @Provides
+    @ElementsIntoSet
+    fun emptyOverlaySet(): Set<Overlay> {
+        return emptySet()
+    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/scene/ShadelessSceneContainerFrameworkModule.kt b/packages/SystemUI/src/com/android/systemui/scene/ShadelessSceneContainerFrameworkModule.kt
index 9a7eef8..16ed59f4 100644
--- a/packages/SystemUI/src/com/android/systemui/scene/ShadelessSceneContainerFrameworkModule.kt
+++ b/packages/SystemUI/src/com/android/systemui/scene/ShadelessSceneContainerFrameworkModule.kt
@@ -53,6 +53,7 @@
                     Scenes.Bouncer,
                 ),
             initialSceneKey = Scenes.Lockscreen,
+            overlayKeys = emptyList(),
             navigationDistances =
                 mapOf(
                     Scenes.Gone to 0,
diff --git a/packages/SystemUI/src/com/android/systemui/scene/data/repository/SceneContainerRepository.kt b/packages/SystemUI/src/com/android/systemui/scene/data/repository/SceneContainerRepository.kt
index beb6816..d60f05e 100644
--- a/packages/SystemUI/src/com/android/systemui/scene/data/repository/SceneContainerRepository.kt
+++ b/packages/SystemUI/src/com/android/systemui/scene/data/repository/SceneContainerRepository.kt
@@ -18,7 +18,9 @@
 
 package com.android.systemui.scene.data.repository
 
+import com.android.compose.animation.scene.ContentKey
 import com.android.compose.animation.scene.ObservableTransitionState
+import com.android.compose.animation.scene.OverlayKey
 import com.android.compose.animation.scene.SceneKey
 import com.android.compose.animation.scene.TransitionKey
 import com.android.systemui.dagger.SysUISingleton
@@ -43,11 +45,27 @@
 @Inject
 constructor(
     @Application applicationScope: CoroutineScope,
-    private val config: SceneContainerConfig,
+    config: SceneContainerConfig,
     private val dataSource: SceneDataSource,
 ) {
+    /**
+     * The keys of all scenes and overlays in the container.
+     *
+     * They will be sorted in z-order such that the last one is the one that should be rendered on
+     * top of all previous ones.
+     */
+    val allContentKeys: List<ContentKey> = config.sceneKeys + config.overlayKeys
+
     val currentScene: StateFlow<SceneKey> = dataSource.currentScene
 
+    /**
+     * The current set of overlays to be shown (may be empty).
+     *
+     * Note that during a transition between overlays, a different set of overlays may be rendered -
+     * but only the ones in this set are considered the current overlays.
+     */
+    val currentOverlays: StateFlow<Set<OverlayKey>> = dataSource.currentOverlays
+
     private val _isVisible = MutableStateFlow(true)
     val isVisible: StateFlow<Boolean> = _isVisible.asStateFlow()
 
@@ -72,16 +90,6 @@
                 initialValue = defaultTransitionState,
             )
 
-    /**
-     * Returns the keys to all scenes in the container.
-     *
-     * The scenes will be sorted in z-order such that the last one is the one that should be
-     * rendered on top of all previous ones.
-     */
-    fun allSceneKeys(): List<SceneKey> {
-        return config.sceneKeys
-    }
-
     fun changeScene(
         toScene: SceneKey,
         transitionKey: TransitionKey? = null,
@@ -100,6 +108,48 @@
         )
     }
 
+    /**
+     * Request to show [overlay] so that it animates in from [currentScene] and ends up being
+     * visible on screen.
+     *
+     * After this returns, this overlay will be included in [currentOverlays]. This does nothing if
+     * [overlay] is already shown.
+     */
+    fun showOverlay(overlay: OverlayKey, transitionKey: TransitionKey? = null) {
+        dataSource.showOverlay(
+            overlay = overlay,
+            transitionKey = transitionKey,
+        )
+    }
+
+    /**
+     * Request to hide [overlay] so that it animates out to [currentScene] and ends up *not* being
+     * visible on screen.
+     *
+     * After this returns, this overlay will not be included in [currentOverlays]. This does nothing
+     * if [overlay] is already hidden.
+     */
+    fun hideOverlay(overlay: OverlayKey, transitionKey: TransitionKey? = null) {
+        dataSource.hideOverlay(
+            overlay = overlay,
+            transitionKey = transitionKey,
+        )
+    }
+
+    /**
+     * Replace [from] by [to] so that [from] ends up not being visible on screen and [to] ends up
+     * being visible.
+     *
+     * This throws if [from] is not currently shown or if [to] is already shown.
+     */
+    fun replaceOverlay(from: OverlayKey, to: OverlayKey, transitionKey: TransitionKey? = null) {
+        dataSource.replaceOverlay(
+            from = from,
+            to = to,
+            transitionKey = transitionKey,
+        )
+    }
+
     /** Sets whether the container is visible. */
     fun setVisible(isVisible: Boolean) {
         _isVisible.value = isVisible
diff --git a/packages/SystemUI/src/com/android/systemui/scene/domain/interactor/SceneInteractor.kt b/packages/SystemUI/src/com/android/systemui/scene/domain/interactor/SceneInteractor.kt
index 4c404e2..bdb148a 100644
--- a/packages/SystemUI/src/com/android/systemui/scene/domain/interactor/SceneInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/scene/domain/interactor/SceneInteractor.kt
@@ -16,7 +16,9 @@
 
 package com.android.systemui.scene.domain.interactor
 
+import com.android.compose.animation.scene.ContentKey
 import com.android.compose.animation.scene.ObservableTransitionState
+import com.android.compose.animation.scene.OverlayKey
 import com.android.compose.animation.scene.SceneKey
 import com.android.compose.animation.scene.TransitionKey
 import com.android.systemui.dagger.SysUISingleton
@@ -51,6 +53,7 @@
  * other feature modules should depend on and call into this class when their parts of the
  * application state change.
  */
+@OptIn(ExperimentalCoroutinesApi::class)
 @SysUISingleton
 class SceneInteractor
 @Inject
@@ -76,6 +79,14 @@
     private val onSceneAboutToChangeListener = mutableSetOf<OnSceneAboutToChangeListener>()
 
     /**
+     * The keys of all scenes and overlays in the container.
+     *
+     * They will be sorted in z-order such that the last one is the one that should be rendered on
+     * top of all previous ones.
+     */
+    val allContentKeys: List<ContentKey> = repository.allContentKeys
+
+    /**
      * The current scene.
      *
      * Note that during a transition between scenes, more than one scene might be rendered but only
@@ -84,6 +95,14 @@
     val currentScene: StateFlow<SceneKey> = repository.currentScene
 
     /**
+     * The current set of overlays to be shown (may be empty).
+     *
+     * Note that during a transition between overlays, a different set of overlays may be rendered -
+     * but only the ones in this set are considered the current overlays.
+     */
+    val currentOverlays: StateFlow<Set<OverlayKey>> = repository.currentOverlays
+
+    /**
      * The current state of the transition.
      *
      * Consumers should use this state to know:
@@ -192,16 +211,6 @@
         }
     }
 
-    /**
-     * Returns the keys of all scenes in the container.
-     *
-     * The scenes will be sorted in z-order such that the last one is the one that should be
-     * rendered on top of all previous ones.
-     */
-    fun allSceneKeys(): List<SceneKey> {
-        return repository.allSceneKeys()
-    }
-
     fun registerSceneStateProcessor(processor: OnSceneAboutToChangeListener) {
         onSceneAboutToChangeListener.add(processor)
     }
@@ -284,6 +293,105 @@
     }
 
     /**
+     * Request to show [overlay] so that it animates in from [currentScene] and ends up being
+     * visible on screen.
+     *
+     * After this returns, this overlay will be included in [currentOverlays]. This does nothing if
+     * [overlay] is already shown.
+     *
+     * @param overlay The overlay to be shown
+     * @param loggingReason The reason why the transition is requested, for logging purposes
+     * @param transitionKey The transition key for this animated transition
+     */
+    @JvmOverloads
+    fun showOverlay(
+        overlay: OverlayKey,
+        loggingReason: String,
+        transitionKey: TransitionKey? = null,
+    ) {
+        if (!validateOverlayChange(to = overlay, loggingReason = loggingReason)) {
+            return
+        }
+
+        logger.logOverlayChangeRequested(
+            to = overlay,
+            reason = loggingReason,
+        )
+
+        repository.showOverlay(
+            overlay = overlay,
+            transitionKey = transitionKey,
+        )
+    }
+
+    /**
+     * Request to hide [overlay] so that it animates out to [currentScene] and ends up *not* being
+     * visible on screen.
+     *
+     * After this returns, this overlay will not be included in [currentOverlays]. This does nothing
+     * if [overlay] is already hidden.
+     *
+     * @param overlay The overlay to be hidden
+     * @param loggingReason The reason why the transition is requested, for logging purposes
+     * @param transitionKey The transition key for this animated transition
+     */
+    @JvmOverloads
+    fun hideOverlay(
+        overlay: OverlayKey,
+        loggingReason: String,
+        transitionKey: TransitionKey? = null,
+    ) {
+        if (!validateOverlayChange(from = overlay, loggingReason = loggingReason)) {
+            return
+        }
+
+        logger.logOverlayChangeRequested(
+            from = overlay,
+            reason = loggingReason,
+        )
+
+        repository.hideOverlay(
+            overlay = overlay,
+            transitionKey = transitionKey,
+        )
+    }
+
+    /**
+     * Replace [from] by [to] so that [from] ends up not being visible on screen and [to] ends up
+     * being visible.
+     *
+     * This throws if [from] is not currently shown or if [to] is already shown.
+     *
+     * @param from The overlay to be hidden, if any
+     * @param to The overlay to be shown, if any
+     * @param loggingReason The reason why the transition is requested, for logging purposes
+     * @param transitionKey The transition key for this animated transition
+     */
+    @JvmOverloads
+    fun replaceOverlay(
+        from: OverlayKey,
+        to: OverlayKey,
+        loggingReason: String,
+        transitionKey: TransitionKey? = null,
+    ) {
+        if (!validateOverlayChange(from = from, to = to, loggingReason = loggingReason)) {
+            return
+        }
+
+        logger.logOverlayChangeRequested(
+            from = from,
+            to = to,
+            reason = loggingReason,
+        )
+
+        repository.replaceOverlay(
+            from = from,
+            to = to,
+            transitionKey = transitionKey,
+        )
+    }
+
+    /**
      * Sets the visibility of the container.
      *
      * Please do not call this from outside of the scene framework. If you are trying to force the
@@ -388,7 +496,7 @@
         to: SceneKey,
         loggingReason: String,
     ): Boolean {
-        if (!repository.allSceneKeys().contains(to)) {
+        if (to !in repository.allContentKeys) {
             return false
         }
 
@@ -409,6 +517,34 @@
         return from != to
     }
 
+    /**
+     * Validates that the given overlay change is allowed.
+     *
+     * Will throw a runtime exception for illegal states.
+     *
+     * @param from The overlay to be hidden, if any
+     * @param to The overlay to be shown, if any
+     * @param loggingReason The reason why the transition is requested, for logging purposes
+     * @return `true` if the scene change is valid; `false` if it shouldn't happen
+     */
+    private fun validateOverlayChange(
+        from: OverlayKey? = null,
+        to: OverlayKey? = null,
+        loggingReason: String,
+    ): Boolean {
+        check(from != null || to != null) {
+            "No overlay key provided for requested change." +
+                " Current transition state is ${transitionState.value}." +
+                " Logging reason for overlay change was: $loggingReason"
+        }
+
+        val isFromValid = (from == null) || (from in currentOverlays.value)
+        val isToValid =
+            (to == null) || (to !in currentOverlays.value && to in repository.allContentKeys)
+
+        return isFromValid && isToValid && from != to
+    }
+
     /** Returns a flow indicating if the currently visible scene can be resolved from [family]. */
     fun isCurrentSceneInFamily(family: SceneKey): Flow<Boolean> =
         currentScene.map { currentScene -> isSceneInFamily(currentScene, family) }
diff --git a/packages/SystemUI/src/com/android/systemui/scene/shared/logger/SceneLogger.kt b/packages/SystemUI/src/com/android/systemui/scene/shared/logger/SceneLogger.kt
index 045a887..aa418e6 100644
--- a/packages/SystemUI/src/com/android/systemui/scene/shared/logger/SceneLogger.kt
+++ b/packages/SystemUI/src/com/android/systemui/scene/shared/logger/SceneLogger.kt
@@ -17,6 +17,7 @@
 package com.android.systemui.scene.shared.logger
 
 import com.android.compose.animation.scene.ObservableTransitionState
+import com.android.compose.animation.scene.OverlayKey
 import com.android.compose.animation.scene.SceneKey
 import com.android.systemui.log.LogBuffer
 import com.android.systemui.log.core.LogLevel
@@ -94,6 +95,34 @@
         }
     }
 
+    fun logOverlayChangeRequested(
+        from: OverlayKey? = null,
+        to: OverlayKey? = null,
+        reason: String,
+    ) {
+        logBuffer.log(
+            tag = TAG,
+            level = LogLevel.INFO,
+            messageInitializer = {
+                str1 = from?.toString()
+                str2 = to?.toString()
+                str3 = reason
+            },
+            messagePrinter = {
+                buildString {
+                    append("Overlay change requested: ")
+                    if (str1 != null) {
+                        append(str1)
+                        append(if (str2 == null) " (hidden)" else " → $str2")
+                    } else {
+                        append("$str2 (shown)")
+                    }
+                    append(", reason: $str3")
+                }
+            },
+        )
+    }
+
     fun logVisibilityChange(
         from: Boolean,
         to: Boolean,
diff --git a/packages/SystemUI/src/com/android/systemui/scene/shared/model/SceneContainerConfig.kt b/packages/SystemUI/src/com/android/systemui/scene/shared/model/SceneContainerConfig.kt
index 0a30c31..2311e47 100644
--- a/packages/SystemUI/src/com/android/systemui/scene/shared/model/SceneContainerConfig.kt
+++ b/packages/SystemUI/src/com/android/systemui/scene/shared/model/SceneContainerConfig.kt
@@ -16,6 +16,7 @@
 
 package com.android.systemui.scene.shared.model
 
+import com.android.compose.animation.scene.OverlayKey
 import com.android.compose.animation.scene.SceneKey
 
 /** Models the configuration of the scene container. */
@@ -38,6 +39,13 @@
     val initialSceneKey: SceneKey,
 
     /**
+     * The keys to all overlays in the container, sorted by z-order such that the last one renders
+     * on top of all previous ones. Overlay keys within the same container must not repeat but it's
+     * okay to have the same overlay keys in different containers.
+     */
+    val overlayKeys: List<OverlayKey> = emptyList(),
+
+    /**
      * Navigation distance of each scene.
      *
      * The navigation distance is a measure of how many non-back user action "steps" away from the
diff --git a/packages/SystemUI/src/com/android/systemui/scene/shared/model/SceneDataSource.kt b/packages/SystemUI/src/com/android/systemui/scene/shared/model/SceneDataSource.kt
index 034da25..4538d1c 100644
--- a/packages/SystemUI/src/com/android/systemui/scene/shared/model/SceneDataSource.kt
+++ b/packages/SystemUI/src/com/android/systemui/scene/shared/model/SceneDataSource.kt
@@ -16,6 +16,7 @@
 
 package com.android.systemui.scene.shared.model
 
+import com.android.compose.animation.scene.OverlayKey
 import com.android.compose.animation.scene.SceneKey
 import com.android.compose.animation.scene.TransitionKey
 import kotlinx.coroutines.flow.StateFlow
@@ -33,6 +34,14 @@
     val currentScene: StateFlow<SceneKey>
 
     /**
+     * The current set of overlays to be shown (may be empty).
+     *
+     * Note that during a transition between overlays, a different set of overlays may be rendered -
+     * but only the ones in this set are considered the current overlays.
+     */
+    val currentOverlays: StateFlow<Set<OverlayKey>>
+
+    /**
      * Asks for an asynchronous scene switch to [toScene], which will use the corresponding
      * installed transition or the one specified by [transitionKey], if provided.
      */
@@ -47,4 +56,40 @@
     fun snapToScene(
         toScene: SceneKey,
     )
+
+    /**
+     * Request to show [overlay] so that it animates in from [currentScene] and ends up being
+     * visible on screen.
+     *
+     * After this returns, this overlay will be included in [currentOverlays]. This does nothing if
+     * [overlay] is already shown.
+     */
+    fun showOverlay(
+        overlay: OverlayKey,
+        transitionKey: TransitionKey? = null,
+    )
+
+    /**
+     * Request to hide [overlay] so that it animates out to [currentScene] and ends up *not* being
+     * visible on screen.
+     *
+     * After this returns, this overlay will not be included in [currentOverlays]. This does nothing
+     * if [overlay] is already hidden.
+     */
+    fun hideOverlay(
+        overlay: OverlayKey,
+        transitionKey: TransitionKey? = null,
+    )
+
+    /**
+     * Replace [from] by [to] so that [from] ends up not being visible on screen and [to] ends up
+     * being visible.
+     *
+     * This throws if [from] is not currently shown or if [to] is already shown.
+     */
+    fun replaceOverlay(
+        from: OverlayKey,
+        to: OverlayKey,
+        transitionKey: TransitionKey? = null,
+    )
 }
diff --git a/packages/SystemUI/src/com/android/systemui/scene/shared/model/SceneDataSourceDelegator.kt b/packages/SystemUI/src/com/android/systemui/scene/shared/model/SceneDataSourceDelegator.kt
index 43c3635..eb4c0f2 100644
--- a/packages/SystemUI/src/com/android/systemui/scene/shared/model/SceneDataSourceDelegator.kt
+++ b/packages/SystemUI/src/com/android/systemui/scene/shared/model/SceneDataSourceDelegator.kt
@@ -18,6 +18,7 @@
 
 package com.android.systemui.scene.shared.model
 
+import com.android.compose.animation.scene.OverlayKey
 import com.android.compose.animation.scene.SceneKey
 import com.android.compose.animation.scene.TransitionKey
 import kotlinx.coroutines.CoroutineScope
@@ -49,6 +50,15 @@
                 initialValue = config.initialSceneKey,
             )
 
+    override val currentOverlays: StateFlow<Set<OverlayKey>> =
+        delegateMutable
+            .flatMapLatest { delegate -> delegate.currentOverlays }
+            .stateIn(
+                scope = applicationScope,
+                started = SharingStarted.WhileSubscribed(),
+                initialValue = emptySet(),
+            )
+
     override fun changeScene(toScene: SceneKey, transitionKey: TransitionKey?) {
         delegateMutable.value.changeScene(
             toScene = toScene,
@@ -62,6 +72,28 @@
         )
     }
 
+    override fun showOverlay(overlay: OverlayKey, transitionKey: TransitionKey?) {
+        delegateMutable.value.showOverlay(
+            overlay = overlay,
+            transitionKey = transitionKey,
+        )
+    }
+
+    override fun hideOverlay(overlay: OverlayKey, transitionKey: TransitionKey?) {
+        delegateMutable.value.hideOverlay(
+            overlay = overlay,
+            transitionKey = transitionKey,
+        )
+    }
+
+    override fun replaceOverlay(from: OverlayKey, to: OverlayKey, transitionKey: TransitionKey?) {
+        delegateMutable.value.replaceOverlay(
+            from = from,
+            to = to,
+            transitionKey = transitionKey,
+        )
+    }
+
     /**
      * Binds the current, dependency injection provided [SceneDataSource] to the given object.
      *
@@ -82,8 +114,21 @@
         override val currentScene: StateFlow<SceneKey> =
             MutableStateFlow(initialSceneKey).asStateFlow()
 
+        override val currentOverlays: StateFlow<Set<OverlayKey>> =
+            MutableStateFlow(emptySet<OverlayKey>()).asStateFlow()
+
         override fun changeScene(toScene: SceneKey, transitionKey: TransitionKey?) = Unit
 
         override fun snapToScene(toScene: SceneKey) = Unit
+
+        override fun showOverlay(overlay: OverlayKey, transitionKey: TransitionKey?) = Unit
+
+        override fun hideOverlay(overlay: OverlayKey, transitionKey: TransitionKey?) = Unit
+
+        override fun replaceOverlay(
+            from: OverlayKey,
+            to: OverlayKey,
+            transitionKey: TransitionKey?
+        ) = Unit
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/scene/ui/view/SceneWindowRootView.kt b/packages/SystemUI/src/com/android/systemui/scene/ui/view/SceneWindowRootView.kt
index 8aa601f..c1bb6fb 100644
--- a/packages/SystemUI/src/com/android/systemui/scene/ui/view/SceneWindowRootView.kt
+++ b/packages/SystemUI/src/com/android/systemui/scene/ui/view/SceneWindowRootView.kt
@@ -9,6 +9,7 @@
 import com.android.systemui.scene.shared.model.Scene
 import com.android.systemui.scene.shared.model.SceneContainerConfig
 import com.android.systemui.scene.shared.model.SceneDataSourceDelegator
+import com.android.systemui.scene.ui.composable.Overlay
 import com.android.systemui.scene.ui.viewmodel.SceneContainerViewModel
 import com.android.systemui.shade.TouchLogger
 import com.android.systemui.statusbar.notification.stack.ui.view.SharedNotificationContainer
@@ -35,6 +36,7 @@
         containerConfig: SceneContainerConfig,
         sharedNotificationContainer: SharedNotificationContainer,
         scenes: Set<Scene>,
+        overlays: Set<Overlay>,
         layoutInsetController: LayoutInsetsController,
         sceneDataSourceDelegator: SceneDataSourceDelegator,
         alternateBouncerDependencies: AlternateBouncerDependencies,
@@ -50,6 +52,7 @@
             containerConfig = containerConfig,
             sharedNotificationContainer = sharedNotificationContainer,
             scenes = scenes,
+            overlays = overlays,
             onVisibilityChangedInternal = { isVisible ->
                 super.setVisibility(if (isVisible) View.VISIBLE else View.INVISIBLE)
             },
diff --git a/packages/SystemUI/src/com/android/systemui/scene/ui/view/SceneWindowRootViewBinder.kt b/packages/SystemUI/src/com/android/systemui/scene/ui/view/SceneWindowRootViewBinder.kt
index 0f05af6..b870a4e 100644
--- a/packages/SystemUI/src/com/android/systemui/scene/ui/view/SceneWindowRootViewBinder.kt
+++ b/packages/SystemUI/src/com/android/systemui/scene/ui/view/SceneWindowRootViewBinder.kt
@@ -29,6 +29,7 @@
 import androidx.compose.ui.unit.dp
 import androidx.core.view.isVisible
 import androidx.lifecycle.Lifecycle
+import com.android.compose.animation.scene.OverlayKey
 import com.android.compose.animation.scene.SceneKey
 import com.android.compose.theme.PlatformTheme
 import com.android.internal.policy.ScreenDecorationsUtils
@@ -47,6 +48,7 @@
 import com.android.systemui.scene.shared.model.SceneContainerConfig
 import com.android.systemui.scene.shared.model.SceneDataSourceDelegator
 import com.android.systemui.scene.ui.composable.ComposableScene
+import com.android.systemui.scene.ui.composable.Overlay
 import com.android.systemui.scene.ui.composable.SceneContainer
 import com.android.systemui.scene.ui.viewmodel.SceneContainerViewModel
 import com.android.systemui.statusbar.notification.stack.ui.view.SharedNotificationContainer
@@ -70,6 +72,7 @@
         containerConfig: SceneContainerConfig,
         sharedNotificationContainer: SharedNotificationContainer,
         scenes: Set<Scene>,
+        overlays: Set<Overlay>,
         onVisibilityChangedInternal: (isVisible: Boolean) -> Unit,
         dataSourceDelegator: SceneDataSourceDelegator,
         alternateBouncerDependencies: AlternateBouncerDependencies,
@@ -86,6 +89,19 @@
             }
         }
 
+        val unsortedOverlayByKey: Map<OverlayKey, Overlay> =
+            overlays.associateBy { overlay -> overlay.key }
+        val sortedOverlayByKey: Map<OverlayKey, Overlay> = buildMap {
+            containerConfig.overlayKeys.forEach { overlayKey ->
+                val overlay =
+                    checkNotNull(unsortedOverlayByKey[overlayKey]) {
+                        "Overlay not found for key \"$overlayKey\"!"
+                    }
+
+                put(overlayKey, overlay)
+            }
+        }
+
         view.repeatWhenAttached {
             view.viewModel(
                 minWindowLifecycleState = WindowLifecycleState.ATTACHED,
@@ -112,6 +128,7 @@
                                 viewModel = viewModel,
                                 windowInsets = windowInsets,
                                 sceneByKey = sortedSceneByKey,
+                                overlayByKey = sortedOverlayByKey,
                                 dataSourceDelegator = dataSourceDelegator,
                                 containerConfig = containerConfig,
                             )
@@ -156,6 +173,7 @@
         viewModel: SceneContainerViewModel,
         windowInsets: StateFlow<WindowInsets?>,
         sceneByKey: Map<SceneKey, Scene>,
+        overlayByKey: Map<OverlayKey, Overlay>,
         dataSourceDelegator: SceneDataSourceDelegator,
         containerConfig: SceneContainerConfig,
     ): View {
@@ -170,6 +188,7 @@
                             viewModel = viewModel,
                             sceneByKey =
                                 sceneByKey.mapValues { (_, scene) -> scene as ComposableScene },
+                            overlayByKey = overlayByKey,
                             initialSceneKey = containerConfig.initialSceneKey,
                             dataSourceDelegator = dataSourceDelegator,
                         )
diff --git a/packages/SystemUI/src/com/android/systemui/scene/ui/viewmodel/SceneContainerViewModel.kt b/packages/SystemUI/src/com/android/systemui/scene/ui/viewmodel/SceneContainerViewModel.kt
index c4be26a..e2947d3 100644
--- a/packages/SystemUI/src/com/android/systemui/scene/ui/viewmodel/SceneContainerViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/scene/ui/viewmodel/SceneContainerViewModel.kt
@@ -48,14 +48,6 @@
     private val logger: SceneLogger,
     @Assisted private val motionEventHandlerReceiver: (MotionEventHandler?) -> Unit,
 ) : SysUiViewModel, ExclusiveActivatable() {
-    /**
-     * Keys of all scenes in the container.
-     *
-     * The scenes will be sorted in z-order such that the last one is the one that should be
-     * rendered on top of all previous ones.
-     */
-    val allSceneKeys: List<SceneKey> = sceneInteractor.allSceneKeys()
-
     /** The scene that should be rendered. */
     val currentScene: StateFlow<SceneKey> = sceneInteractor.currentScene
 
diff --git a/packages/SystemUI/src/com/android/systemui/shade/ShadeViewProviderModule.kt b/packages/SystemUI/src/com/android/systemui/shade/ShadeViewProviderModule.kt
index 606fef0..018144b 100644
--- a/packages/SystemUI/src/com/android/systemui/shade/ShadeViewProviderModule.kt
+++ b/packages/SystemUI/src/com/android/systemui/shade/ShadeViewProviderModule.kt
@@ -14,6 +14,8 @@
  * limitations under the License.
  */
 
+@file:OptIn(ExperimentalCoroutinesApi::class)
+
 package com.android.systemui.shade
 
 import android.annotation.SuppressLint
@@ -38,6 +40,7 @@
 import com.android.systemui.scene.shared.model.Scene
 import com.android.systemui.scene.shared.model.SceneContainerConfig
 import com.android.systemui.scene.shared.model.SceneDataSourceDelegator
+import com.android.systemui.scene.ui.composable.Overlay
 import com.android.systemui.scene.ui.view.SceneWindowRootView
 import com.android.systemui.scene.ui.view.WindowRootView
 import com.android.systemui.scene.ui.viewmodel.SceneContainerViewModel
@@ -59,6 +62,7 @@
 import dagger.Provides
 import javax.inject.Named
 import javax.inject.Provider
+import kotlinx.coroutines.ExperimentalCoroutinesApi
 
 /** Module for providing views related to the shade. */
 @Module
@@ -82,6 +86,7 @@
             viewModelFactory: SceneContainerViewModel.Factory,
             containerConfigProvider: Provider<SceneContainerConfig>,
             scenesProvider: Provider<Set<@JvmSuppressWildcards Scene>>,
+            overlaysProvider: Provider<Set<@JvmSuppressWildcards Overlay>>,
             layoutInsetController: NotificationInsetsController,
             sceneDataSourceDelegator: Provider<SceneDataSourceDelegator>,
             alternateBouncerDependencies: Provider<AlternateBouncerDependencies>,
@@ -96,6 +101,7 @@
                     sharedNotificationContainer =
                         sceneWindowRootView.requireViewById(R.id.shared_notification_container),
                     scenes = scenesProvider.get(),
+                    overlays = overlaysProvider.get(),
                     layoutInsetController = layoutInsetController,
                     sceneDataSourceDelegator = sceneDataSourceDelegator.get(),
                     alternateBouncerDependencies = alternateBouncerDependencies.get(),
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/scene/SceneKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/scene/SceneKosmos.kt
index dd93141..7dfe802 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/scene/SceneKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/scene/SceneKosmos.kt
@@ -1,10 +1,12 @@
 package com.android.systemui.scene
 
+import com.android.compose.animation.scene.OverlayKey
 import com.android.systemui.kosmos.Kosmos
 import com.android.systemui.kosmos.Kosmos.Fixture
 import com.android.systemui.scene.shared.model.FakeScene
 import com.android.systemui.scene.shared.model.SceneContainerConfig
 import com.android.systemui.scene.shared.model.Scenes
+import com.android.systemui.scene.ui.FakeOverlay
 
 var Kosmos.sceneKeys by Fixture {
     listOf(
@@ -22,6 +24,17 @@
 val Kosmos.scenes by Fixture { fakeScenes }
 
 val Kosmos.initialSceneKey by Fixture { Scenes.Lockscreen }
+
+var Kosmos.overlayKeys by Fixture {
+    listOf<OverlayKey>(
+        // TODO(b/356596436): Add overlays here when we have them.
+    )
+}
+
+val Kosmos.fakeOverlays by Fixture { overlayKeys.map { key -> FakeOverlay(key) }.toSet() }
+
+val Kosmos.overlays by Fixture { fakeOverlays }
+
 var Kosmos.sceneContainerConfig by Fixture {
     val navigationDistances =
         mapOf(
@@ -32,5 +45,11 @@
             Scenes.QuickSettings to 3,
             Scenes.Bouncer to 4,
         )
-    SceneContainerConfig(sceneKeys, initialSceneKey, navigationDistances)
+
+    SceneContainerConfig(
+        sceneKeys = sceneKeys,
+        initialSceneKey = initialSceneKey,
+        overlayKeys = overlayKeys,
+        navigationDistances = navigationDistances,
+    )
 }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/scene/shared/model/FakeSceneDataSource.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/scene/shared/model/FakeSceneDataSource.kt
index 957a60f..f52572a 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/scene/shared/model/FakeSceneDataSource.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/scene/shared/model/FakeSceneDataSource.kt
@@ -16,6 +16,7 @@
 
 package com.android.systemui.scene.shared.model
 
+import com.android.compose.animation.scene.OverlayKey
 import com.android.compose.animation.scene.SceneKey
 import com.android.compose.animation.scene.TransitionKey
 import kotlinx.coroutines.flow.MutableStateFlow
@@ -29,11 +30,18 @@
     private val _currentScene = MutableStateFlow(initialSceneKey)
     override val currentScene: StateFlow<SceneKey> = _currentScene.asStateFlow()
 
+    private val _currentOverlays = MutableStateFlow<Set<OverlayKey>>(emptySet())
+    override val currentOverlays: StateFlow<Set<OverlayKey>> = _currentOverlays.asStateFlow()
+
     var isPaused = false
         private set
+
     var pendingScene: SceneKey? = null
         private set
 
+    var pendingOverlays: Set<OverlayKey>? = null
+        private set
+
     override fun changeScene(toScene: SceneKey, transitionKey: TransitionKey?) {
         if (isPaused) {
             pendingScene = toScene
@@ -46,10 +54,32 @@
         changeScene(toScene)
     }
 
+    override fun showOverlay(overlay: OverlayKey, transitionKey: TransitionKey?) {
+        if (isPaused) {
+            pendingOverlays = (pendingOverlays ?: currentOverlays.value) + overlay
+        } else {
+            _currentOverlays.value += overlay
+        }
+    }
+
+    override fun hideOverlay(overlay: OverlayKey, transitionKey: TransitionKey?) {
+        if (isPaused) {
+            pendingOverlays = (pendingOverlays ?: currentOverlays.value) - overlay
+        } else {
+            _currentOverlays.value -= overlay
+        }
+    }
+
+    override fun replaceOverlay(from: OverlayKey, to: OverlayKey, transitionKey: TransitionKey?) {
+        hideOverlay(from, transitionKey)
+        showOverlay(to, transitionKey)
+    }
+
     /**
-     * Pauses scene changes.
+     * Pauses scene and overlay changes.
      *
-     * Any following calls to [changeScene] will be conflated and the last one will be remembered.
+     * Any following calls to [changeScene] or overlay changing functions will be conflated and the
+     * last one will be remembered.
      */
     fun pause() {
         check(!isPaused) { "Can't pause what's already paused!" }
@@ -58,11 +88,14 @@
     }
 
     /**
-     * Unpauses scene changes.
+     * Unpauses scene and overlay changes.
      *
      * If there were any calls to [changeScene] since [pause] was called, the latest of the bunch
      * will be replayed.
      *
+     * If there were any calls to show, hide or replace overlays since [pause] was called, they will
+     * all be applied at once.
+     *
      * If [force] is `true`, there will be no check that [isPaused] is true.
      *
      * If [expectedScene] is provided, will assert that it's indeed the latest called.
@@ -76,6 +109,8 @@
         isPaused = false
         pendingScene?.let { _currentScene.value = it }
         pendingScene = null
+        pendingOverlays?.let { _currentOverlays.value = it }
+        pendingOverlays = null
 
         check(expectedScene == null || currentScene.value == expectedScene) {
             """
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/scene/ui/FakeOverlay.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/scene/ui/FakeOverlay.kt
new file mode 100644
index 0000000..f4f30cd
--- /dev/null
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/scene/ui/FakeOverlay.kt
@@ -0,0 +1,36 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.scene.ui
+
+import androidx.compose.runtime.Composable
+import androidx.compose.ui.Modifier
+import com.android.compose.animation.scene.ContentScope
+import com.android.compose.animation.scene.OverlayKey
+import com.android.systemui.lifecycle.ExclusiveActivatable
+import com.android.systemui.scene.ui.composable.Overlay
+import kotlinx.coroutines.awaitCancellation
+
+class FakeOverlay(
+    override val key: OverlayKey,
+) : ExclusiveActivatable(), Overlay {
+
+    @Composable override fun ContentScope.Content(modifier: Modifier) = Unit
+
+    override suspend fun onActivated(): Nothing {
+        awaitCancellation()
+    }
+}