Merge "[flexiglass] Clear session storage if we are idle on Gone scene" into main
diff --git a/packages/SystemUI/src/com/android/systemui/scene/domain/startable/SceneContainerStartable.kt b/packages/SystemUI/src/com/android/systemui/scene/domain/startable/SceneContainerStartable.kt
index 3fca84e..5b50133 100644
--- a/packages/SystemUI/src/com/android/systemui/scene/domain/startable/SceneContainerStartable.kt
+++ b/packages/SystemUI/src/com/android/systemui/scene/domain/startable/SceneContainerStartable.kt
@@ -172,12 +172,28 @@
 
     private fun resetShadeSessions() {
         applicationScope.launch {
-            sceneBackInteractor.backStack
-                // We are in a session if either Shade or QuickSettings is on the back stack
-                .map { backStack ->
-                    backStack.asIterable().any { it == Scenes.Shade || it == Scenes.QuickSettings }
+            combine(
+                    sceneBackInteractor.backStack
+                        // We are in a session if either Shade or QuickSettings is on the back stack
+                        .map { backStack ->
+                            backStack.asIterable().any {
+                                it == Scenes.Shade || it == Scenes.QuickSettings
+                            }
+                        }
+                        .distinctUntilChanged(),
+                    sceneInteractor.transitionState
+                        .mapNotNull { state ->
+                            // We are also in a session if either Shade or QuickSettings is the
+                            // current scene
+                            when (state) {
+                                is ObservableTransitionState.Idle -> state.currentScene
+                                is ObservableTransitionState.Transition -> state.fromScene
+                            }.let { it == Scenes.Shade || it == Scenes.QuickSettings }
+                        }
+                        .distinctUntilChanged()
+                ) { inBackStack, isCurrentScene ->
+                    inBackStack || isCurrentScene
                 }
-                .distinctUntilChanged()
                 // Once a session has ended, clear the session storage.
                 .filter { inSession -> !inSession }
                 .collect { shadeSessionStorage.clear() }