Merge "Update formatting of VirtualCameraController dumpsys" into main
diff --git a/core/java/android/hardware/camera2/CameraManager.java b/core/java/android/hardware/camera2/CameraManager.java
index 2191fd5..a019612 100644
--- a/core/java/android/hardware/camera2/CameraManager.java
+++ b/core/java/android/hardware/camera2/CameraManager.java
@@ -69,6 +69,7 @@
 import android.util.ArrayMap;
 import android.util.ArraySet;
 import android.util.Log;
+import android.util.Pair;
 import android.util.Size;
 import android.view.Display;
 
@@ -379,7 +380,8 @@
      */
     @NonNull
     public Set<Set<String>> getConcurrentCameraIds() throws CameraAccessException {
-        return CameraManagerGlobal.get().getConcurrentCameraIds();
+        return CameraManagerGlobal.get().getConcurrentCameraIds(mContext.getDeviceId(),
+                getDevicePolicyFromContext(mContext));
     }
 
     /**
@@ -418,7 +420,8 @@
             @NonNull Map<String, SessionConfiguration> cameraIdAndSessionConfig)
             throws CameraAccessException {
         return CameraManagerGlobal.get().isConcurrentSessionConfigurationSupported(
-                cameraIdAndSessionConfig, mContext.getApplicationInfo().targetSdkVersion);
+                cameraIdAndSessionConfig, mContext.getApplicationInfo().targetSdkVersion,
+                mContext.getDeviceId(), getDevicePolicyFromContext(mContext));
     }
 
     /**
@@ -794,7 +797,7 @@
 
             boolean hasConcurrentStreams =
                     CameraManagerGlobal.get().cameraIdHasConcurrentStreamsLocked(cameraId,
-                            mContext.getDeviceId());
+                            mContext.getDeviceId(), getDevicePolicyFromContext(mContext));
             metadata.setHasMandatoryConcurrentStreams(hasConcurrentStreams);
 
             Size displaySize = getDisplaySize();
@@ -2097,7 +2100,7 @@
         // Opened Camera ID -> apk name map
         private final ArrayMap<DeviceCameraInfo, String> mOpenedDevices = new ArrayMap<>();
 
-        private final Set<Set<String>> mConcurrentCameraIdCombinations = new ArraySet<>();
+        private final Set<Set<DeviceCameraInfo>> mConcurrentCameraIdCombinations = new ArraySet<>();
 
         // Registered availability callbacks and their executors
         private final ArrayMap<AvailabilityCallback, Executor> mCallbackMap = new ArrayMap<>();
@@ -2256,7 +2259,13 @@
                 ConcurrentCameraIdCombination[] cameraIdCombinations =
                         cameraService.getConcurrentCameraIds();
                 for (ConcurrentCameraIdCombination comb : cameraIdCombinations) {
-                    mConcurrentCameraIdCombinations.add(comb.getConcurrentCameraIdCombination());
+                    Set<Pair<String, Integer>> combination =
+                            comb.getConcurrentCameraIdCombination();
+                    Set<DeviceCameraInfo> deviceCameraInfoSet = new ArraySet<>();
+                    for (Pair<String, Integer> entry : combination) {
+                        deviceCameraInfoSet.add(new DeviceCameraInfo(entry.first, entry.second));
+                    }
+                    mConcurrentCameraIdCombinations.add(deviceCameraInfoSet);
                 }
             } catch (ServiceSpecificException e) {
                 // Unexpected failure
@@ -2341,13 +2350,12 @@
             return cameraIds.toArray(new String[0]);
         }
 
-        private Set<Set<String>> extractConcurrentCameraIdListLocked() {
+        private Set<Set<String>> extractConcurrentCameraIdListLocked(int deviceId,
+                int devicePolicy) {
             Set<Set<String>> concurrentCameraIds = new ArraySet<>();
-            for (Set<String> cameraIds : mConcurrentCameraIdCombinations) {
+            for (Set<DeviceCameraInfo> deviceCameraInfos : mConcurrentCameraIdCombinations) {
                 Set<String> extractedCameraIds = new ArraySet<>();
-                for (String cameraId : cameraIds) {
-                    // TODO(b/291736219): This to be made device-aware.
-                    DeviceCameraInfo info = new DeviceCameraInfo(cameraId, DEVICE_ID_DEFAULT);
+                for (DeviceCameraInfo info : deviceCameraInfos) {
                     // if the camera id status is NOT_PRESENT or ENUMERATING; skip the device.
                     // TODO: Would a device status NOT_PRESENT ever be in the map ? it gets removed
                     // in the callback anyway.
@@ -2360,9 +2368,14 @@
                             || status == ICameraServiceListener.STATUS_NOT_PRESENT) {
                         continue;
                     }
-                    extractedCameraIds.add(cameraId);
+                    if (shouldHideCamera(deviceId, devicePolicy, info)) {
+                        continue;
+                    }
+                    extractedCameraIds.add(info.mCameraId);
                 }
-                concurrentCameraIds.add(extractedCameraIds);
+                if (!extractedCameraIds.isEmpty()) {
+                    concurrentCameraIds.add(extractedCameraIds);
+                }
             }
             return concurrentCameraIds;
         }
@@ -2523,12 +2536,13 @@
             return cameraIds;
         }
 
-        public @NonNull Set<Set<String>> getConcurrentCameraIds() {
+        public @NonNull Set<Set<String>> getConcurrentCameraIds(int deviceId, int devicePolicy) {
             Set<Set<String>> concurrentStreamingCameraIds;
             synchronized (mLock) {
                 // Try to make sure we have an up-to-date list of concurrent camera devices.
                 connectCameraServiceLocked();
-                concurrentStreamingCameraIds = extractConcurrentCameraIdListLocked();
+                concurrentStreamingCameraIds = extractConcurrentCameraIdListLocked(deviceId,
+                        devicePolicy);
             }
             // TODO: Some sort of sorting  ?
             return concurrentStreamingCameraIds;
@@ -2536,13 +2550,12 @@
 
         public boolean isConcurrentSessionConfigurationSupported(
                 @NonNull Map<String, SessionConfiguration> cameraIdsAndSessionConfigurations,
-                int targetSdkVersion) throws CameraAccessException {
+                int targetSdkVersion, int deviceId, int devicePolicy)
+                throws CameraAccessException {
             if (cameraIdsAndSessionConfigurations == null) {
                 throw new IllegalArgumentException("cameraIdsAndSessionConfigurations was null");
             }
 
-            // TODO(b/291736219): Check if this API needs to be made device-aware.
-
             int size = cameraIdsAndSessionConfigurations.size();
             if (size == 0) {
                 throw new IllegalArgumentException("camera id and session combination is empty");
@@ -2552,14 +2565,20 @@
                 // Go through all the elements and check if the camera ids are valid at least /
                 // belong to one of the combinations returned by getConcurrentCameraIds()
                 boolean subsetFound = false;
-                for (Set<String> combination : mConcurrentCameraIdCombinations) {
-                    if (combination.containsAll(cameraIdsAndSessionConfigurations.keySet())) {
+                for (Set<DeviceCameraInfo> combination : mConcurrentCameraIdCombinations) {
+                    Set<DeviceCameraInfo> infos = new ArraySet<>();
+                    for (String cameraId : cameraIdsAndSessionConfigurations.keySet()) {
+                        infos.add(new DeviceCameraInfo(cameraId,
+                                devicePolicy == DEVICE_POLICY_DEFAULT
+                                        ? DEVICE_ID_DEFAULT : deviceId));
+                    }
+                    if (combination.containsAll(infos)) {
                         subsetFound = true;
                     }
                 }
                 if (!subsetFound) {
                     Log.v(TAG, "isConcurrentSessionConfigurationSupported called with a subset of"
-                            + "camera ids not returned by getConcurrentCameraIds");
+                            + " camera ids not returned by getConcurrentCameraIds");
                     return false;
                 }
                 CameraIdAndSessionConfiguration [] cameraIdsAndConfigs =
@@ -2573,7 +2592,7 @@
                 }
                 try {
                     return mCameraService.isConcurrentSessionConfigurationSupported(
-                            cameraIdsAndConfigs, targetSdkVersion);
+                            cameraIdsAndConfigs, targetSdkVersion, deviceId, devicePolicy);
                 } catch (ServiceSpecificException e) {
                     throw ExceptionUtils.throwAsPublicException(e);
                 } catch (RemoteException e) {
@@ -2592,8 +2611,10 @@
          * @return Whether the camera device was found in the set of combinations returned by
          *         getConcurrentCameraIds
          */
-        public boolean cameraIdHasConcurrentStreamsLocked(String cameraId, int deviceId) {
-            DeviceCameraInfo info = new DeviceCameraInfo(cameraId, deviceId);
+        public boolean cameraIdHasConcurrentStreamsLocked(String cameraId, int deviceId,
+                int devicePolicy) {
+            DeviceCameraInfo info = new DeviceCameraInfo(cameraId,
+                    devicePolicy == DEVICE_POLICY_DEFAULT ? DEVICE_ID_DEFAULT : deviceId);
             if (!mDeviceStatus.containsKey(info)) {
                 // physical camera ids aren't advertised in concurrent camera id combinations.
                 if (DEBUG) {
@@ -2602,8 +2623,8 @@
                 }
                 return false;
             }
-            for (Set<String> comb : mConcurrentCameraIdCombinations) {
-                if (comb.contains(cameraId)) {
+            for (Set<DeviceCameraInfo> comb : mConcurrentCameraIdCombinations) {
+                if (comb.contains(info)) {
                     return true;
                 }
             }
diff --git a/core/java/android/hardware/camera2/utils/ConcurrentCameraIdCombination.java b/core/java/android/hardware/camera2/utils/ConcurrentCameraIdCombination.java
index 8f4d636..1a8bf1d 100644
--- a/core/java/android/hardware/camera2/utils/ConcurrentCameraIdCombination.java
+++ b/core/java/android/hardware/camera2/utils/ConcurrentCameraIdCombination.java
@@ -13,13 +13,15 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
+
 package android.hardware.camera2.utils;
 
 import android.annotation.NonNull;
 import android.os.Parcel;
 import android.os.Parcelable;
+import android.util.ArraySet;
+import android.util.Pair;
 
-import java.util.HashSet;
 import java.util.Set;
 
 /**
@@ -30,21 +32,21 @@
  */
 public class ConcurrentCameraIdCombination implements Parcelable {
 
-    private Set<String> mConcurrentCameraIds = new HashSet<>();
+    private final Set<Pair<String, Integer>> mConcurrentCameraIdDeviceIdPairs = new ArraySet<>();
 
     public static final @NonNull
             Parcelable.Creator<ConcurrentCameraIdCombination> CREATOR =
-            new Parcelable.Creator<ConcurrentCameraIdCombination>() {
-        @Override
-        public ConcurrentCameraIdCombination createFromParcel(Parcel in) {
-            return new ConcurrentCameraIdCombination(in);
-        }
+            new Parcelable.Creator<>() {
+                @Override
+                public ConcurrentCameraIdCombination createFromParcel(Parcel in) {
+                    return new ConcurrentCameraIdCombination(in);
+                }
 
-        @Override
-        public ConcurrentCameraIdCombination[] newArray(int size) {
-            return new ConcurrentCameraIdCombination[size];
-        }
-    };
+                @Override
+                public ConcurrentCameraIdCombination[] newArray(int size) {
+                    return new ConcurrentCameraIdCombination[size];
+                }
+            };
 
     private ConcurrentCameraIdCombination(Parcel in) {
         readFromParcel(in);
@@ -57,9 +59,10 @@
 
     @Override
     public void writeToParcel(Parcel dest, int flags) {
-        dest.writeInt(mConcurrentCameraIds.size());
-        for (String cameraId : mConcurrentCameraIds) {
-            dest.writeString(cameraId);
+        dest.writeInt(mConcurrentCameraIdDeviceIdPairs.size());
+        for (Pair<String, Integer> cameraIdDeviceIdPair : mConcurrentCameraIdDeviceIdPairs) {
+            dest.writeString(cameraIdDeviceIdPair.first);
+            dest.writeInt(cameraIdDeviceIdPair.second);
         }
     }
 
@@ -67,7 +70,7 @@
      * helper for CREATOR
      */
     public void readFromParcel(Parcel in) {
-        mConcurrentCameraIds.clear();
+        mConcurrentCameraIdDeviceIdPairs.clear();
         int cameraCombinationSize = in.readInt();
         if (cameraCombinationSize < 0) {
             throw new RuntimeException("cameraCombinationSize " + cameraCombinationSize
@@ -78,14 +81,15 @@
             if (cameraId == null) {
                 throw new RuntimeException("Failed to read camera id from Parcel");
             }
-            mConcurrentCameraIds.add(cameraId);
+            int deviceId = in.readInt();
+            mConcurrentCameraIdDeviceIdPairs.add(new Pair<>(cameraId, deviceId));
         }
     }
 
     /**
      * Get this concurrent camera id combination.
      */
-    public Set<String> getConcurrentCameraIdCombination() {
-        return mConcurrentCameraIds;
+    public Set<Pair<String, Integer>> getConcurrentCameraIdCombination() {
+        return mConcurrentCameraIdDeviceIdPairs;
     }
 }
diff --git a/core/java/android/tracing/flags.aconfig b/core/java/android/tracing/flags.aconfig
index 74428aa..d7389ba 100644
--- a/core/java/android/tracing/flags.aconfig
+++ b/core/java/android/tracing/flags.aconfig
@@ -30,3 +30,11 @@
     is_fixed_read_only: true
     bug: "276433199"
 }
+
+flag {
+    name: "perfetto_view_capture_tracing"
+    namespace: "windowing_tools"
+    description: "Migrate ViewCapture tracing to Perfetto"
+    is_fixed_read_only: true
+    bug: "323166383"
+}
diff --git a/core/tests/coretests/src/android/window/WindowOnBackInvokedDispatcherTest.java b/core/tests/coretests/src/android/window/WindowOnBackInvokedDispatcherTest.java
index d54b862..50d7f59 100644
--- a/core/tests/coretests/src/android/window/WindowOnBackInvokedDispatcherTest.java
+++ b/core/tests/coretests/src/android/window/WindowOnBackInvokedDispatcherTest.java
@@ -384,11 +384,7 @@
         clearInvocations(mCallback1);
 
         callbackInfo.getCallback().onBackCancelled();
-
         waitForIdle();
-        // verify onBackCancelled not yet called (since BackProgressAnimator animates
-        // progress to 0 first)
-        verify(mCallback1, never()).onBackCancelled();
 
         // simulate start of new gesture while cancel animation is still running
         callbackInfo.getCallback().onBackStarted(mBackEvent);
diff --git a/packages/SystemUI/aconfig/systemui.aconfig b/packages/SystemUI/aconfig/systemui.aconfig
index 24d6413..fd90bd9 100644
--- a/packages/SystemUI/aconfig/systemui.aconfig
+++ b/packages/SystemUI/aconfig/systemui.aconfig
@@ -803,6 +803,16 @@
 }
 
 flag {
+    name: "shade_collapse_activity_launch_fix"
+    namespace: "systemui"
+    description: "Avoid collapsing the shade on activity launch if it is already collapsed, as this causes a flicker."
+    bug: "331591373"
+    metadata {
+      purpose: PURPOSE_BUGFIX
+    }
+}
+
+flag {
     name: "slice_broadcast_relay_in_background"
     namespace: "systemui"
     description: "Move handling of slice broadcast relay broadcasts to background threads"
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/shade/ui/composable/OverlayShade.kt b/packages/SystemUI/compose/features/src/com/android/systemui/shade/ui/composable/OverlayShade.kt
new file mode 100644
index 0000000..d528736
--- /dev/null
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/shade/ui/composable/OverlayShade.kt
@@ -0,0 +1,147 @@
+/*
+ * 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.shade.ui.composable
+
+import androidx.compose.foundation.background
+import androidx.compose.foundation.clickable
+import androidx.compose.foundation.layout.Arrangement
+import androidx.compose.foundation.layout.Box
+import androidx.compose.foundation.layout.Row
+import androidx.compose.foundation.layout.Spacer
+import androidx.compose.foundation.layout.fillMaxSize
+import androidx.compose.foundation.layout.padding
+import androidx.compose.foundation.layout.width
+import androidx.compose.foundation.shape.RoundedCornerShape
+import androidx.compose.material3.MaterialTheme
+import androidx.compose.material3.Text
+import androidx.compose.runtime.Composable
+import androidx.compose.runtime.ReadOnlyComposable
+import androidx.compose.runtime.collectAsState
+import androidx.compose.runtime.getValue
+import androidx.compose.ui.Alignment
+import androidx.compose.ui.Modifier
+import androidx.compose.ui.draw.clip
+import androidx.compose.ui.graphics.Color
+import androidx.compose.ui.unit.dp
+import com.android.compose.animation.scene.ElementKey
+import com.android.compose.animation.scene.LowestZIndexScenePicker
+import com.android.compose.animation.scene.SceneScope
+import com.android.systemui.scene.shared.model.Scenes
+import com.android.systemui.shade.ui.viewmodel.OverlayShadeViewModel
+
+/** The overlay shade renders a lightweight shade UI container on top of a background scene. */
+@Composable
+fun SceneScope.OverlayShade(
+    viewModel: OverlayShadeViewModel,
+    horizontalArrangement: Arrangement.Horizontal,
+    modifier: Modifier = Modifier,
+    content: @Composable () -> Unit,
+) {
+    val backgroundScene by viewModel.backgroundScene.collectAsState()
+
+    Box(modifier) {
+        if (backgroundScene == Scenes.Lockscreen) {
+            Lockscreen()
+        }
+
+        Scrim(onClicked = viewModel::onScrimClicked)
+
+        Row(
+            modifier = Modifier.fillMaxSize().padding(OverlayShade.Dimensions.ScrimContentPadding),
+            horizontalArrangement = horizontalArrangement,
+        ) {
+            Panel(content = content)
+        }
+    }
+}
+
+@Composable
+private fun Lockscreen(
+    modifier: Modifier = Modifier,
+) {
+    // TODO(b/338025605): This is a placeholder, replace with the actual lockscreen.
+    Box(modifier = modifier.fillMaxSize().background(Color.LightGray)) {
+        Text(text = "Lockscreen", modifier = Modifier.align(Alignment.Center))
+    }
+}
+
+@Composable
+private fun SceneScope.Scrim(
+    onClicked: () -> Unit,
+    modifier: Modifier = Modifier,
+) {
+    Spacer(
+        modifier =
+            modifier
+                .element(OverlayShade.Elements.Scrim)
+                .fillMaxSize()
+                .background(OverlayShade.Colors.ScrimBackground)
+                .clickable(onClick = onClicked, interactionSource = null, indication = null)
+    )
+}
+
+@Composable
+private fun SceneScope.Panel(
+    modifier: Modifier = Modifier,
+    content: @Composable () -> Unit,
+) {
+    Box(
+        modifier =
+            modifier
+                .width(OverlayShade.Dimensions.PanelWidth)
+                .clip(OverlayShade.Shapes.RoundedCornerPanel)
+    ) {
+        Spacer(
+            modifier =
+                Modifier.element(OverlayShade.Elements.PanelBackground)
+                    .matchParentSize()
+                    .background(
+                        color = OverlayShade.Colors.PanelBackground,
+                        shape = OverlayShade.Shapes.RoundedCornerPanel,
+                    ),
+        )
+
+        // This content is intentionally rendered as a separate element from the background in order
+        // to allow for more flexibility when defining transitions.
+        content()
+    }
+}
+
+object OverlayShade {
+    object Elements {
+        val Scrim = ElementKey("OverlayShadeScrim", scenePicker = LowestZIndexScenePicker)
+        val PanelBackground =
+            ElementKey("OverlayShadePanelBackground", scenePicker = LowestZIndexScenePicker)
+    }
+
+    object Colors {
+        val ScrimBackground = Color(0, 0, 0, alpha = 255 / 3)
+        val PanelBackground: Color
+            @Composable @ReadOnlyComposable get() = MaterialTheme.colorScheme.surfaceContainer
+    }
+
+    object Dimensions {
+        val ScrimContentPadding = 16.dp
+        val PanelCornerRadius = 46.dp
+        // TODO(b/338033836): This width should not be fixed.
+        val PanelWidth = 390.dp
+    }
+
+    object Shapes {
+        val RoundedCornerPanel = RoundedCornerShape(Dimensions.PanelCornerRadius)
+    }
+}
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractorTest.kt
index 12f8918..323087c 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractorTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractorTest.kt
@@ -24,23 +24,23 @@
 import com.android.compose.animation.scene.ObservableTransitionState
 import com.android.systemui.Flags.FLAG_KEYGUARD_WM_STATE_REFACTOR
 import com.android.systemui.SysuiTestCase
-import com.android.systemui.bouncer.data.repository.FakeKeyguardBouncerRepository
-import com.android.systemui.common.ui.data.repository.FakeConfigurationRepository
-import com.android.systemui.common.ui.domain.interactor.ConfigurationInteractor
+import com.android.systemui.bouncer.data.repository.keyguardBouncerRepository
+import com.android.systemui.common.ui.data.repository.fakeConfigurationRepository
 import com.android.systemui.coroutines.collectLastValue
 import com.android.systemui.flags.EnableSceneContainer
-import com.android.systemui.keyguard.data.repository.FakeCommandQueue
+import com.android.systemui.keyguard.data.repository.fakeCommandQueue
 import com.android.systemui.keyguard.data.repository.fakeKeyguardRepository
 import com.android.systemui.keyguard.data.repository.fakeKeyguardTransitionRepository
 import com.android.systemui.keyguard.shared.model.CameraLaunchSourceModel
 import com.android.systemui.keyguard.shared.model.KeyguardState
 import com.android.systemui.keyguard.shared.model.StatusBarState
+import com.android.systemui.keyguard.shared.model.TransitionState
+import com.android.systemui.keyguard.shared.model.TransitionStep
 import com.android.systemui.kosmos.testScope
-import com.android.systemui.power.domain.interactor.PowerInteractorFactory
+import com.android.systemui.res.R
 import com.android.systemui.scene.domain.interactor.sceneInteractor
 import com.android.systemui.scene.shared.model.Scenes
-import com.android.systemui.shade.data.repository.FakeShadeRepository
-import com.android.systemui.statusbar.notification.stack.domain.interactor.sharedNotificationContainerInteractor
+import com.android.systemui.shade.data.repository.shadeRepository
 import com.android.systemui.testKosmos
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.ExperimentalCoroutinesApi
@@ -60,33 +60,17 @@
 
     private val kosmos = testKosmos()
     private val testScope = kosmos.testScope
-    private val repository by lazy { kosmos.fakeKeyguardRepository }
-    private val sceneInteractor by lazy { kosmos.sceneInteractor }
-    private val fromGoneTransitionInteractor by lazy { kosmos.fromGoneTransitionInteractor }
-    private val commandQueue by lazy { FakeCommandQueue() }
-    private val bouncerRepository = FakeKeyguardBouncerRepository()
-    private val shadeRepository = FakeShadeRepository()
+    private val repository = kosmos.fakeKeyguardRepository
+    private val sceneInteractor = kosmos.sceneInteractor
+    private val fromGoneTransitionInteractor = kosmos.fromGoneTransitionInteractor
+    private val commandQueue = kosmos.fakeCommandQueue
+    private val configRepository = kosmos.fakeConfigurationRepository
+    private val bouncerRepository = kosmos.keyguardBouncerRepository
+    private val shadeRepository = kosmos.shadeRepository
     private val keyguardTransitionRepository = kosmos.fakeKeyguardTransitionRepository
     private val transitionState: MutableStateFlow<ObservableTransitionState> =
         MutableStateFlow(ObservableTransitionState.Idle(Scenes.Gone))
-
-    private val underTest by lazy {
-        KeyguardInteractor(
-            repository = repository,
-            commandQueue = commandQueue,
-            powerInteractor = PowerInteractorFactory.create().powerInteractor,
-            bouncerRepository = bouncerRepository,
-            configurationInteractor = ConfigurationInteractor(FakeConfigurationRepository()),
-            shadeRepository = shadeRepository,
-            keyguardTransitionInteractor = kosmos.keyguardTransitionInteractor,
-            sceneInteractorProvider = { sceneInteractor },
-            fromGoneTransitionInteractor = { fromGoneTransitionInteractor },
-            sharedNotificationContainerInteractor = {
-                kosmos.sharedNotificationContainerInteractor
-            },
-            applicationScope = testScope,
-        )
-    }
+    private val underTest = kosmos.keyguardInteractor
 
     @Before
     fun setUp() {
@@ -247,6 +231,84 @@
         }
 
     @Test
+    fun keyguardTranslationY_whenGoneEmitsZero() =
+        testScope.runTest {
+            val keyguardTranslationY by collectLastValue(underTest.keyguardTranslationY)
+
+            configRepository.setDimensionPixelSize(
+                R.dimen.keyguard_translate_distance_on_swipe_up,
+                100
+            )
+            configRepository.onAnyConfigurationChange()
+
+            shadeRepository.setLegacyShadeExpansion(0f)
+
+            keyguardTransitionRepository.sendTransitionSteps(
+                from = KeyguardState.AOD,
+                to = KeyguardState.GONE,
+                testScope,
+            )
+
+            assertThat(keyguardTranslationY).isEqualTo(0f)
+        }
+
+    @Test
+    fun keyguardTranslationY_whenNotGoneAndShadeIsFullyCollapsedEmitsZero() =
+        testScope.runTest {
+            val keyguardTranslationY by collectLastValue(underTest.keyguardTranslationY)
+
+            configRepository.setDimensionPixelSize(
+                R.dimen.keyguard_translate_distance_on_swipe_up,
+                100
+            )
+            configRepository.onAnyConfigurationChange()
+
+            shadeRepository.setLegacyShadeExpansion(0f)
+
+            keyguardTransitionRepository.sendTransitionSteps(
+                from = KeyguardState.AOD,
+                to = KeyguardState.LOCKSCREEN,
+                testScope,
+            )
+
+            assertThat(keyguardTranslationY).isEqualTo(0f)
+        }
+
+    @Test
+    fun keyguardTranslationY_whenTransitioningToGoneAndShadeIsExpandingEmitsNonZero() =
+        testScope.runTest {
+            val keyguardTranslationY by collectLastValue(underTest.keyguardTranslationY)
+
+            configRepository.setDimensionPixelSize(
+                R.dimen.keyguard_translate_distance_on_swipe_up,
+                100
+            )
+            configRepository.onAnyConfigurationChange()
+
+            shadeRepository.setLegacyShadeExpansion(0.5f)
+
+            keyguardTransitionRepository.sendTransitionSteps(
+                listOf(
+                    TransitionStep(
+                        from = KeyguardState.AOD,
+                        to = KeyguardState.GONE,
+                        value = 0f,
+                        transitionState = TransitionState.STARTED,
+                    ),
+                    TransitionStep(
+                        from = KeyguardState.AOD,
+                        to = KeyguardState.GONE,
+                        value = 0.1f,
+                        transitionState = TransitionState.RUNNING,
+                    ),
+                ),
+                testScope,
+            )
+
+            assertThat(keyguardTranslationY).isGreaterThan(0f)
+        }
+
+    @Test
     @EnableSceneContainer
     fun animationDozingTransitions() =
         testScope.runTest {
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/shade/ui/viewmodel/OverlayShadeViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/shade/ui/viewmodel/OverlayShadeViewModelTest.kt
new file mode 100644
index 0000000..0ffabd8
--- /dev/null
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/shade/ui/viewmodel/OverlayShadeViewModelTest.kt
@@ -0,0 +1,163 @@
+/*
+ * 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.shade.ui.viewmodel
+
+import android.testing.TestableLooper
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.authentication.data.repository.fakeAuthenticationRepository
+import com.android.systemui.authentication.shared.model.AuthenticationMethodModel
+import com.android.systemui.coroutines.collectLastValue
+import com.android.systemui.deviceentry.data.repository.fakeDeviceEntryRepository
+import com.android.systemui.deviceentry.domain.interactor.deviceUnlockedInteractor
+import com.android.systemui.flags.EnableSceneContainer
+import com.android.systemui.keyguard.data.repository.fakeDeviceEntryFingerprintAuthRepository
+import com.android.systemui.keyguard.shared.model.SuccessFingerprintAuthenticationStatus
+import com.android.systemui.kosmos.testScope
+import com.android.systemui.scene.domain.interactor.sceneInteractor
+import com.android.systemui.scene.shared.model.Scenes
+import com.android.systemui.testKosmos
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.test.TestScope
+import kotlinx.coroutines.test.runCurrent
+import kotlinx.coroutines.test.runTest
+import org.junit.Test
+import org.junit.runner.RunWith
+
+@OptIn(ExperimentalCoroutinesApi::class)
+@SmallTest
+@RunWith(AndroidJUnit4::class)
+@TestableLooper.RunWithLooper
+@EnableSceneContainer
+class OverlayShadeViewModelTest : SysuiTestCase() {
+
+    private val kosmos = testKosmos()
+    private val testScope = kosmos.testScope
+    private val sceneInteractor = kosmos.sceneInteractor
+    private val deviceUnlockedInteractor by lazy { kosmos.deviceUnlockedInteractor }
+
+    private val underTest = kosmos.overlayShadeViewModel
+
+    @Test
+    fun backgroundScene_deviceLocked_lockscreen() =
+        testScope.runTest {
+            val backgroundScene by collectLastValue(underTest.backgroundScene)
+
+            lockDevice()
+
+            assertThat(backgroundScene).isEqualTo(Scenes.Lockscreen)
+        }
+
+    @Test
+    fun backgroundScene_deviceUnlocked_gone() =
+        testScope.runTest {
+            val backgroundScene by collectLastValue(underTest.backgroundScene)
+
+            lockDevice()
+            unlockDevice()
+
+            assertThat(backgroundScene).isEqualTo(Scenes.Gone)
+        }
+
+    @Test
+    fun backgroundScene_authMethodSwipe_lockscreenNotDismissed_goesToLockscreen() =
+        testScope.runTest {
+            val backgroundScene by collectLastValue(underTest.backgroundScene)
+            val deviceUnlockStatus by collectLastValue(deviceUnlockedInteractor.deviceUnlockStatus)
+
+            kosmos.fakeDeviceEntryRepository.setLockscreenEnabled(true)
+            kosmos.fakeAuthenticationRepository.setAuthenticationMethod(
+                AuthenticationMethodModel.None
+            )
+            assertThat(deviceUnlockStatus?.isUnlocked).isTrue()
+            sceneInteractor.changeScene(Scenes.Lockscreen, "reason")
+            runCurrent()
+
+            assertThat(backgroundScene).isEqualTo(Scenes.Lockscreen)
+        }
+
+    @Test
+    fun backgroundScene_authMethodSwipe_lockscreenDismissed_goesToGone() =
+        testScope.runTest {
+            val backgroundScene by collectLastValue(underTest.backgroundScene)
+            val deviceUnlockStatus by collectLastValue(deviceUnlockedInteractor.deviceUnlockStatus)
+
+            kosmos.fakeDeviceEntryRepository.setLockscreenEnabled(true)
+            kosmos.fakeAuthenticationRepository.setAuthenticationMethod(
+                AuthenticationMethodModel.None
+            )
+            assertThat(deviceUnlockStatus?.isUnlocked).isTrue()
+            sceneInteractor.changeScene(Scenes.Gone, "reason")
+            runCurrent()
+
+            assertThat(backgroundScene).isEqualTo(Scenes.Gone)
+        }
+
+    @Test
+    fun onScrimClicked_onLockscreen_goesToLockscreen() =
+        testScope.runTest {
+            val currentScene by collectLastValue(sceneInteractor.currentScene)
+            lockDevice()
+            sceneInteractor.changeScene(Scenes.Bouncer, "reason")
+            runCurrent()
+            assertThat(currentScene).isNotEqualTo(Scenes.Lockscreen)
+
+            underTest.onScrimClicked()
+
+            assertThat(currentScene).isEqualTo(Scenes.Lockscreen)
+        }
+
+    @Test
+    fun onScrimClicked_deviceWasEntered_goesToGone() =
+        testScope.runTest {
+            val currentScene by collectLastValue(sceneInteractor.currentScene)
+            val backgroundScene by collectLastValue(underTest.backgroundScene)
+
+            lockDevice()
+            unlockDevice()
+            sceneInteractor.changeScene(Scenes.QuickSettings, "reason")
+            runCurrent()
+            assertThat(backgroundScene).isEqualTo(Scenes.Gone)
+            assertThat(currentScene).isNotEqualTo(Scenes.Gone)
+
+            underTest.onScrimClicked()
+
+            assertThat(currentScene).isEqualTo(Scenes.Gone)
+        }
+
+    private fun TestScope.lockDevice() {
+        val deviceUnlockStatus by collectLastValue(deviceUnlockedInteractor.deviceUnlockStatus)
+
+        kosmos.fakeAuthenticationRepository.setAuthenticationMethod(AuthenticationMethodModel.Pin)
+        assertThat(deviceUnlockStatus?.isUnlocked).isFalse()
+        sceneInteractor.changeScene(Scenes.Lockscreen, "reason")
+        runCurrent()
+    }
+
+    private fun TestScope.unlockDevice() {
+        val deviceUnlockStatus by collectLastValue(deviceUnlockedInteractor.deviceUnlockStatus)
+
+        kosmos.fakeDeviceEntryFingerprintAuthRepository.setAuthenticationStatus(
+            SuccessFingerprintAuthenticationStatus(0, true)
+        )
+        assertThat(deviceUnlockStatus?.isUnlocked).isTrue()
+        sceneInteractor.changeScene(Scenes.Gone, "reason")
+        runCurrent()
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardViewMediator.java b/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardViewMediator.java
index 5d31d1e..6ef39f3 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardViewMediator.java
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardViewMediator.java
@@ -3214,6 +3214,10 @@
         mHideAnimationRun = false;
         adjustStatusBarLocked();
         sendUserPresentBroadcast();
+
+        if (!KeyguardWmStateRefactor.isEnabled()) {
+            mKeyguardInteractor.dismissKeyguard();
+        }
     }
 
     private Configuration.Builder createInteractionJankMonitorConf(int cuj) {
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractor.kt b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractor.kt
index d191768..8d107ab 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractor.kt
@@ -90,6 +90,7 @@
     keyguardTransitionInteractor: KeyguardTransitionInteractor,
     sceneInteractorProvider: Provider<SceneInteractor>,
     private val fromGoneTransitionInteractor: Provider<FromGoneTransitionInteractor>,
+    private val fromLockscreenTransitionInteractor: Provider<FromLockscreenTransitionInteractor>,
     sharedNotificationContainerInteractor: Provider<SharedNotificationContainerInteractor>,
     @Application applicationScope: CoroutineScope,
 ) {
@@ -307,19 +308,27 @@
         configurationInteractor
             .dimensionPixelSize(R.dimen.keyguard_translate_distance_on_swipe_up)
             .flatMapLatest { translationDistance ->
-                combine(
+                combineTransform(
                     shadeRepository.legacyShadeExpansion.onStart { emit(0f) },
                     keyguardTransitionInteractor.transitionValue(GONE).onStart { emit(0f) },
                 ) { legacyShadeExpansion, goneValue ->
-                    if (goneValue == 1f || legacyShadeExpansion == 0f) {
+                    if (goneValue == 1f || (goneValue == 0f && legacyShadeExpansion == 0f)) {
                         // Reset the translation value
-                        0f
-                    } else {
-                        // On swipe up, translate the keyguard to reveal the bouncer
-                        MathUtils.lerp(
-                            translationDistance,
-                            0,
-                            Interpolators.FAST_OUT_LINEAR_IN.getInterpolation(legacyShadeExpansion)
+                        emit(0f)
+                    } else if (legacyShadeExpansion > 0f && legacyShadeExpansion < 1f) {
+                        // On swipe up, translate the keyguard to reveal the bouncer, OR a GONE
+                        // transition is running, which means this is a swipe to dismiss. Values of
+                        // 0f and 1f need to be ignored in the legacy shade expansion. These can
+                        // flip arbitrarily as the legacy shade is reset, and would cause the
+                        // translation value to jump around unexpectedly.
+                        emit(
+                            MathUtils.lerp(
+                                translationDistance,
+                                0,
+                                Interpolators.FAST_OUT_LINEAR_IN.getInterpolation(
+                                    legacyShadeExpansion
+                                ),
+                            )
                         )
                     }
                 }
@@ -417,6 +426,11 @@
         fromGoneTransitionInteractor.get().showKeyguard()
     }
 
+    /** Temporary shim, until [KeyguardWmStateRefactor] is enabled */
+    fun dismissKeyguard() {
+        fromLockscreenTransitionInteractor.get().dismissKeyguard()
+    }
+
     companion object {
         private const val TAG = "KeyguardInteractor"
     }
diff --git a/packages/SystemUI/src/com/android/systemui/qs/panels/dagger/PanelsModule.kt b/packages/SystemUI/src/com/android/systemui/qs/panels/dagger/PanelsModule.kt
index ee4eeb8..e3ba36f 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/panels/dagger/PanelsModule.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/panels/dagger/PanelsModule.kt
@@ -18,20 +18,31 @@
 
 import com.android.systemui.qs.panels.data.repository.IconTilesRepository
 import com.android.systemui.qs.panels.data.repository.IconTilesRepositoryImpl
-import com.android.systemui.qs.panels.shared.model.GridLayoutTypeKey
+import com.android.systemui.qs.panels.shared.model.GridLayoutType
 import com.android.systemui.qs.panels.shared.model.InfiniteGridLayoutType
 import com.android.systemui.qs.panels.ui.compose.GridLayout
 import com.android.systemui.qs.panels.ui.compose.InfiniteGridLayout
 import dagger.Binds
 import dagger.Module
-import dagger.multibindings.IntoMap
+import dagger.Provides
+import dagger.multibindings.IntoSet
 
 @Module
 interface PanelsModule {
     @Binds fun bindIconTilesRepository(impl: IconTilesRepositoryImpl): IconTilesRepository
 
-    @Binds
-    @IntoMap
-    @GridLayoutTypeKey(InfiniteGridLayoutType::class)
-    fun bindGridLayout(impl: InfiniteGridLayout): GridLayout
+    companion object {
+        @Provides
+        @IntoSet
+        fun provideGridLayout(gridLayout: InfiniteGridLayout): Pair<GridLayoutType, GridLayout> {
+            return Pair(InfiniteGridLayoutType, gridLayout)
+        }
+
+        @Provides
+        fun provideGridLayoutMap(
+            entries: Set<@JvmSuppressWildcards Pair<GridLayoutType, GridLayout>>
+        ): Map<GridLayoutType, GridLayout> {
+            return entries.toMap()
+        }
+    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/GridLayout.kt b/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/GridLayout.kt
index 920cbe7..68ce5d8 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/GridLayout.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/GridLayout.kt
@@ -25,6 +25,5 @@
     fun TileGrid(
         tiles: List<TileViewModel>,
         modifier: Modifier,
-        tile: @Composable (TileViewModel) -> Unit,
     )
 }
diff --git a/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/InfiniteGridLayout.kt b/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/InfiniteGridLayout.kt
index 4d0089e7..e2143e0 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/InfiniteGridLayout.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/InfiniteGridLayout.kt
@@ -16,32 +16,77 @@
 
 package com.android.systemui.qs.panels.ui.compose
 
+import android.graphics.drawable.Animatable
+import android.text.TextUtils
+import androidx.appcompat.content.res.AppCompatResources
+import androidx.compose.animation.graphics.ExperimentalAnimationGraphicsApi
+import androidx.compose.animation.graphics.res.animatedVectorResource
+import androidx.compose.animation.graphics.res.rememberAnimatedVectorPainter
+import androidx.compose.animation.graphics.vector.AnimatedImageVector
+import androidx.compose.foundation.Image
+import androidx.compose.foundation.background
+import androidx.compose.foundation.basicMarquee
+import androidx.compose.foundation.clickable
 import androidx.compose.foundation.layout.Arrangement
+import androidx.compose.foundation.layout.Column
+import androidx.compose.foundation.layout.Row
+import androidx.compose.foundation.layout.fillMaxHeight
+import androidx.compose.foundation.layout.fillMaxWidth
+import androidx.compose.foundation.layout.height
+import androidx.compose.foundation.layout.padding
+import androidx.compose.foundation.layout.size
 import androidx.compose.foundation.lazy.grid.GridCells
 import androidx.compose.foundation.lazy.grid.GridItemSpan
 import androidx.compose.foundation.lazy.grid.LazyVerticalGrid
+import androidx.compose.foundation.shape.RoundedCornerShape
+import androidx.compose.material3.Text
 import androidx.compose.runtime.Composable
 import androidx.compose.runtime.DisposableEffect
+import androidx.compose.runtime.LaunchedEffect
+import androidx.compose.runtime.collectAsState
+import androidx.compose.runtime.getValue
+import androidx.compose.runtime.mutableStateOf
+import androidx.compose.runtime.remember
+import androidx.compose.runtime.setValue
+import androidx.compose.ui.Alignment
 import androidx.compose.ui.Modifier
+import androidx.compose.ui.draw.clip
+import androidx.compose.ui.graphics.Color
+import androidx.compose.ui.graphics.ColorFilter
+import androidx.compose.ui.platform.LocalContext
 import androidx.compose.ui.res.dimensionResource
 import androidx.compose.ui.res.integerResource
+import com.android.compose.theme.colorAttr
+import com.android.systemui.common.shared.model.Icon
+import com.android.systemui.common.ui.compose.Icon
+import com.android.systemui.dagger.SysUISingleton
+import com.android.systemui.qs.panels.domain.interactor.IconTilesInteractor
+import com.android.systemui.qs.panels.ui.viewmodel.TileUiState
 import com.android.systemui.qs.panels.ui.viewmodel.TileViewModel
+import com.android.systemui.qs.panels.ui.viewmodel.toUiState
+import com.android.systemui.qs.tileimpl.QSTileImpl
 import com.android.systemui.res.R
 import javax.inject.Inject
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.delay
+import kotlinx.coroutines.flow.mapLatest
 
-class InfiniteGridLayout @Inject constructor() : GridLayout {
+@SysUISingleton
+class InfiniteGridLayout @Inject constructor(private val iconTilesInteractor: IconTilesInteractor) :
+    GridLayout {
 
     @Composable
     override fun TileGrid(
         tiles: List<TileViewModel>,
         modifier: Modifier,
-        tile: @Composable (TileViewModel) -> Unit
     ) {
         DisposableEffect(tiles) {
             val token = Any()
             tiles.forEach { it.startListening(token) }
             onDispose { tiles.forEach { it.stopListening(token) } }
         }
+        val iconTilesSpecs by
+            iconTilesInteractor.iconTilesSpecs.collectAsState(initial = emptySet())
 
         LazyVerticalGrid(
             columns =
@@ -54,14 +99,128 @@
                 Arrangement.spacedBy(dimensionResource(R.dimen.qs_tile_margin_horizontal)),
             modifier = modifier
         ) {
-            tiles.forEach { item(span = { it.span() }) { tile(it) } }
+            items(
+                tiles.size,
+                span = { index ->
+                    val iconOnly = iconTilesSpecs.contains(tiles[index].spec)
+                    if (iconOnly) {
+                        GridItemSpan(1)
+                    } else {
+                        GridItemSpan(2)
+                    }
+                }
+            ) { index ->
+                Tile(
+                    tiles[index],
+                    iconTilesSpecs.contains(tiles[index].spec),
+                    Modifier.height(dimensionResource(id = R.dimen.qs_tile_height))
+                )
+            }
         }
     }
 
-    private fun TileViewModel.span(): GridItemSpan =
-        if (iconOnly) {
-            GridItemSpan(1)
-        } else {
-            GridItemSpan(2)
+    @OptIn(ExperimentalCoroutinesApi::class)
+    @Composable
+    private fun Tile(
+        tile: TileViewModel,
+        iconOnly: Boolean,
+        modifier: Modifier,
+    ) {
+        val state: TileUiState by
+            tile.state
+                .mapLatest { it.toUiState() }
+                .collectAsState(initial = tile.currentState.toUiState())
+        val context = LocalContext.current
+        val horizontalAlignment =
+            if (iconOnly) {
+                Alignment.CenterHorizontally
+            } else {
+                Alignment.Start
+            }
+
+        Row(
+            modifier =
+                modifier
+                    .fillMaxWidth()
+                    .clip(RoundedCornerShape(dimensionResource(R.dimen.qs_corner_radius)))
+                    .clickable { tile.onClick(null) }
+                    .background(colorAttr(state.colors.background))
+                    .padding(
+                        horizontal = dimensionResource(id = R.dimen.qs_label_container_margin)
+                    ),
+            verticalAlignment = Alignment.CenterVertically,
+            horizontalArrangement =
+                Arrangement.spacedBy(
+                    space = dimensionResource(id = R.dimen.qs_label_container_margin),
+                    alignment = horizontalAlignment
+                )
+        ) {
+            val icon =
+                remember(state.icon) {
+                    state.icon.get().let {
+                        if (it is QSTileImpl.ResourceIcon) {
+                            Icon.Resource(it.resId, null)
+                        } else {
+                            Icon.Loaded(it.getDrawable(context), null)
+                        }
+                    }
+                }
+            TileIcon(icon, colorAttr(state.colors.icon))
+
+            if (!iconOnly) {
+                Column(
+                    verticalArrangement = Arrangement.Center,
+                    modifier = Modifier.fillMaxHeight()
+                ) {
+                    Text(
+                        state.label.toString(),
+                        color = colorAttr(state.colors.label),
+                        modifier = Modifier.basicMarquee(),
+                    )
+                    if (!TextUtils.isEmpty(state.secondaryLabel)) {
+                        Text(
+                            state.secondaryLabel.toString(),
+                            color = colorAttr(state.colors.secondaryLabel),
+                            modifier = Modifier.basicMarquee(),
+                        )
+                    }
+                }
+            }
         }
+    }
+
+    @OptIn(ExperimentalAnimationGraphicsApi::class)
+    @Composable
+    private fun TileIcon(icon: Icon, color: Color) {
+        val modifier = Modifier.size(dimensionResource(id = R.dimen.qs_icon_size))
+        val context = LocalContext.current
+        val loadedDrawable =
+            remember(icon, context) {
+                when (icon) {
+                    is Icon.Loaded -> icon.drawable
+                    is Icon.Resource -> AppCompatResources.getDrawable(context, icon.res)
+                }
+            }
+        if (loadedDrawable !is Animatable) {
+            Icon(
+                icon = icon,
+                tint = color,
+                modifier = modifier,
+            )
+        } else if (icon is Icon.Resource) {
+            val image = AnimatedImageVector.animatedVectorResource(id = icon.res)
+            var atEnd by remember(icon.res) { mutableStateOf(false) }
+            LaunchedEffect(key1 = icon.res) {
+                delay(350)
+                atEnd = true
+            }
+            val painter = rememberAnimatedVectorPainter(animatedImageVector = image, atEnd = atEnd)
+            Image(
+                painter = painter,
+                contentDescription = null,
+                colorFilter = ColorFilter.tint(color = color),
+                modifier = modifier
+            )
+        }
+    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/Tile.kt b/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/Tile.kt
deleted file mode 100644
index 35f2970..0000000
--- a/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/Tile.kt
+++ /dev/null
@@ -1,155 +0,0 @@
-/*
- * 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.qs.panels.ui.compose
-
-import android.graphics.drawable.Animatable
-import android.text.TextUtils
-import androidx.appcompat.content.res.AppCompatResources
-import androidx.compose.animation.graphics.ExperimentalAnimationGraphicsApi
-import androidx.compose.animation.graphics.res.animatedVectorResource
-import androidx.compose.animation.graphics.res.rememberAnimatedVectorPainter
-import androidx.compose.animation.graphics.vector.AnimatedImageVector
-import androidx.compose.foundation.Image
-import androidx.compose.foundation.background
-import androidx.compose.foundation.basicMarquee
-import androidx.compose.foundation.clickable
-import androidx.compose.foundation.layout.Arrangement
-import androidx.compose.foundation.layout.Column
-import androidx.compose.foundation.layout.Row
-import androidx.compose.foundation.layout.fillMaxHeight
-import androidx.compose.foundation.layout.fillMaxWidth
-import androidx.compose.foundation.layout.padding
-import androidx.compose.foundation.layout.size
-import androidx.compose.foundation.shape.RoundedCornerShape
-import androidx.compose.material3.Text
-import androidx.compose.runtime.Composable
-import androidx.compose.runtime.LaunchedEffect
-import androidx.compose.runtime.collectAsState
-import androidx.compose.runtime.getValue
-import androidx.compose.runtime.mutableStateOf
-import androidx.compose.runtime.remember
-import androidx.compose.runtime.setValue
-import androidx.compose.ui.Alignment
-import androidx.compose.ui.Modifier
-import androidx.compose.ui.draw.clip
-import androidx.compose.ui.graphics.Color
-import androidx.compose.ui.graphics.ColorFilter
-import androidx.compose.ui.platform.LocalContext
-import androidx.compose.ui.res.dimensionResource
-import com.android.compose.theme.colorAttr
-import com.android.systemui.common.shared.model.Icon as IconModel
-import com.android.systemui.common.ui.compose.Icon
-import com.android.systemui.qs.panels.ui.viewmodel.TileUiState
-import com.android.systemui.qs.panels.ui.viewmodel.TileViewModel
-import com.android.systemui.qs.tileimpl.QSTileImpl
-import com.android.systemui.res.R
-import kotlinx.coroutines.delay
-
-@Composable
-fun Tile(
-    tileViewModel: TileViewModel,
-    modifier: Modifier = Modifier,
-) {
-    val state: TileUiState by tileViewModel.state.collectAsState(tileViewModel.currentState)
-    val context = LocalContext.current
-    val horizontalAlignment =
-        if (state.iconOnly) {
-            Alignment.CenterHorizontally
-        } else {
-            Alignment.Start
-        }
-
-    Row(
-        modifier =
-            modifier
-                .fillMaxWidth()
-                .clip(RoundedCornerShape(dimensionResource(R.dimen.qs_corner_radius)))
-                .clickable { tileViewModel.onClick(null) }
-                .background(colorAttr(state.colors.background))
-                .padding(horizontal = dimensionResource(id = R.dimen.qs_label_container_margin)),
-        verticalAlignment = Alignment.CenterVertically,
-        horizontalArrangement =
-            Arrangement.spacedBy(
-                space = dimensionResource(id = R.dimen.qs_label_container_margin),
-                alignment = horizontalAlignment
-            )
-    ) {
-        val icon =
-            remember(state.icon) {
-                state.icon.get().let {
-                    if (it is QSTileImpl.ResourceIcon) {
-                        IconModel.Resource(it.resId, null)
-                    } else {
-                        IconModel.Loaded(it.getDrawable(context), null)
-                    }
-                }
-            }
-        TileIcon(icon, colorAttr(state.colors.icon))
-
-        if (!state.iconOnly) {
-            Column(verticalArrangement = Arrangement.Center, modifier = Modifier.fillMaxHeight()) {
-                Text(
-                    state.label.toString(),
-                    color = colorAttr(state.colors.label),
-                    modifier = Modifier.basicMarquee(),
-                )
-                if (!TextUtils.isEmpty(state.secondaryLabel)) {
-                    Text(
-                        state.secondaryLabel.toString(),
-                        color = colorAttr(state.colors.secondaryLabel),
-                        modifier = Modifier.basicMarquee(),
-                    )
-                }
-            }
-        }
-    }
-}
-
-@OptIn(ExperimentalAnimationGraphicsApi::class)
-@Composable
-private fun TileIcon(icon: IconModel, color: Color) {
-    val modifier = Modifier.size(dimensionResource(id = R.dimen.qs_icon_size))
-    val context = LocalContext.current
-    val loadedDrawable =
-        remember(icon, context) {
-            when (icon) {
-                is IconModel.Loaded -> icon.drawable
-                is IconModel.Resource -> AppCompatResources.getDrawable(context, icon.res)
-            }
-        }
-    if (loadedDrawable !is Animatable) {
-        Icon(
-            icon = icon,
-            tint = color,
-            modifier = modifier,
-        )
-    } else if (icon is IconModel.Resource) {
-        val image = AnimatedImageVector.animatedVectorResource(id = icon.res)
-        var atEnd by remember(icon.res) { mutableStateOf(false) }
-        LaunchedEffect(key1 = icon.res) {
-            delay(350)
-            atEnd = true
-        }
-        val painter = rememberAnimatedVectorPainter(animatedImageVector = image, atEnd = atEnd)
-        Image(
-            painter = painter,
-            contentDescription = null,
-            colorFilter = ColorFilter.tint(color = color),
-            modifier = modifier
-        )
-    }
-}
diff --git a/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/TileGrid.kt b/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/TileGrid.kt
index a528eed..2f32d72 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/TileGrid.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/TileGrid.kt
@@ -16,21 +16,16 @@
 
 package com.android.systemui.qs.panels.ui.compose
 
-import androidx.compose.foundation.layout.height
 import androidx.compose.runtime.Composable
 import androidx.compose.runtime.collectAsState
 import androidx.compose.runtime.getValue
 import androidx.compose.ui.Modifier
-import androidx.compose.ui.res.dimensionResource
 import com.android.systemui.qs.panels.ui.viewmodel.TileGridViewModel
-import com.android.systemui.res.R
 
 @Composable
 fun TileGrid(viewModel: TileGridViewModel, modifier: Modifier = Modifier) {
-    val gridLayout by viewModel.gridLayout.collectAsState(InfiniteGridLayout())
+    val gridLayout by viewModel.gridLayout.collectAsState()
     val tiles by viewModel.tileViewModels.collectAsState(emptyList())
 
-    gridLayout.TileGrid(tiles, modifier) {
-        Tile(it, modifier = Modifier.height(dimensionResource(id = R.dimen.qs_tile_height)))
-    }
+    gridLayout.TileGrid(tiles, modifier)
 }
diff --git a/packages/SystemUI/src/com/android/systemui/qs/panels/ui/viewmodel/TileGridViewModel.kt b/packages/SystemUI/src/com/android/systemui/qs/panels/ui/viewmodel/TileGridViewModel.kt
index fc13460..5eee691 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/panels/ui/viewmodel/TileGridViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/panels/ui/viewmodel/TileGridViewModel.kt
@@ -17,34 +17,39 @@
 package com.android.systemui.qs.panels.ui.viewmodel
 
 import com.android.systemui.dagger.SysUISingleton
+import com.android.systemui.dagger.qualifiers.Application
 import com.android.systemui.qs.panels.domain.interactor.GridLayoutTypeInteractor
-import com.android.systemui.qs.panels.domain.interactor.IconTilesInteractor
 import com.android.systemui.qs.panels.shared.model.GridLayoutType
 import com.android.systemui.qs.panels.ui.compose.GridLayout
 import com.android.systemui.qs.panels.ui.compose.InfiniteGridLayout
 import com.android.systemui.qs.pipeline.domain.interactor.CurrentTilesInteractor
 import javax.inject.Inject
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.flow.Flow
-import kotlinx.coroutines.flow.combine
+import kotlinx.coroutines.flow.SharingStarted
+import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.flow.mapLatest
+import kotlinx.coroutines.flow.stateIn
 
+@OptIn(ExperimentalCoroutinesApi::class)
 @SysUISingleton
 class TileGridViewModel
 @Inject
 constructor(
     gridLayoutTypeInteractor: GridLayoutTypeInteractor,
-    gridLayoutMap: Map<Class<out GridLayoutType>, @JvmSuppressWildcards GridLayout>,
+    gridLayoutMap: Map<GridLayoutType, @JvmSuppressWildcards GridLayout>,
     tilesInteractor: CurrentTilesInteractor,
-    iconTilesInteractor: IconTilesInteractor,
+    defaultGridLayout: InfiniteGridLayout,
+    @Application private val applicationScope: CoroutineScope
 ) {
-    val gridLayout: Flow<GridLayout> =
-        gridLayoutTypeInteractor.layout.map {
-            gridLayoutMap[it::class.java] ?: InfiniteGridLayout()
-        }
+    val gridLayout: StateFlow<GridLayout> =
+        gridLayoutTypeInteractor.layout
+            .map { gridLayoutMap[it] ?: defaultGridLayout }
+            .stateIn(applicationScope, SharingStarted.Eagerly, defaultGridLayout)
     val tileViewModels: Flow<List<TileViewModel>> =
-        combine(tilesInteractor.currentTiles, iconTilesInteractor.iconTilesSpecs) {
-            tiles,
-            iconTilesSpecs ->
-            tiles.map { TileViewModel(it.tile, iconTilesSpecs.contains(it.spec)) }
+        tilesInteractor.currentTiles.mapLatest { tiles ->
+            tiles.map { TileViewModel(it.tile, it.spec) }
         }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/qs/panels/ui/viewmodel/TileUiState.kt b/packages/SystemUI/src/com/android/systemui/qs/panels/ui/viewmodel/TileUiState.kt
index f4b7255..58d07c3 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/panels/ui/viewmodel/TileUiState.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/panels/ui/viewmodel/TileUiState.kt
@@ -24,15 +24,13 @@
     val secondaryLabel: CharSequence,
     val colors: TileColorAttributes,
     val icon: Supplier<QSTile.Icon>,
-    val iconOnly: Boolean,
 )
 
-fun QSTile.State.toUiState(iconOnly: Boolean): TileUiState {
+fun QSTile.State.toUiState(): TileUiState {
     return TileUiState(
         label ?: "",
         secondaryLabel ?: "",
         colors(),
         icon?.let { Supplier { icon } } ?: iconSupplier,
-        iconOnly,
     )
 }
diff --git a/packages/SystemUI/src/com/android/systemui/qs/panels/ui/viewmodel/TileViewModel.kt b/packages/SystemUI/src/com/android/systemui/qs/panels/ui/viewmodel/TileViewModel.kt
index 08e9119..8fc66d3 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/panels/ui/viewmodel/TileViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/panels/ui/viewmodel/TileViewModel.kt
@@ -19,16 +19,16 @@
 import android.view.View
 import android.view.View.OnLongClickListener
 import com.android.systemui.plugins.qs.QSTile
+import com.android.systemui.qs.pipeline.shared.TileSpec
 import com.android.systemui.utils.coroutines.flow.conflatedCallbackFlow
 import kotlinx.coroutines.channels.awaitClose
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.distinctUntilChanged
-import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.flow.onStart
 
-class TileViewModel(private val tile: QSTile, val iconOnly: Boolean = false) :
+class TileViewModel(private val tile: QSTile, val spec: TileSpec) :
     OnLongClickListener, View.OnClickListener {
-    val state: Flow<TileUiState> =
+    val state: Flow<QSTile.State> =
         conflatedCallbackFlow {
                 val callback = QSTile.Callback { trySend(it.copy()) }
 
@@ -37,11 +37,10 @@
                 awaitClose { tile.removeCallback(callback) }
             }
             .onStart { emit(tile.state) }
-            .map { it.toUiState(iconOnly) }
             .distinctUntilChanged()
 
-    val currentState: TileUiState
-        get() = tile.state.toUiState(iconOnly)
+    val currentState: QSTile.State
+        get() = tile.state
 
     override fun onClick(view: View?) {
         tile.click(view)
diff --git a/packages/SystemUI/src/com/android/systemui/shade/NotificationPanelViewController.java b/packages/SystemUI/src/com/android/systemui/shade/NotificationPanelViewController.java
index 4d4d177..9f1b423 100644
--- a/packages/SystemUI/src/com/android/systemui/shade/NotificationPanelViewController.java
+++ b/packages/SystemUI/src/com/android/systemui/shade/NotificationPanelViewController.java
@@ -25,6 +25,7 @@
 import static com.android.keyguard.KeyguardClockSwitch.LARGE;
 import static com.android.keyguard.KeyguardClockSwitch.SMALL;
 import static com.android.systemui.Flags.predictiveBackAnimateShade;
+import static com.android.systemui.Flags.shadeCollapseActivityLaunchFix;
 import static com.android.systemui.Flags.smartspaceRelocateToBottom;
 import static com.android.systemui.classifier.Classifier.BOUNCER_UNLOCK;
 import static com.android.systemui.classifier.Classifier.GENERIC;
@@ -4102,7 +4103,11 @@
 
     @Override
     public boolean canBeCollapsed() {
-        return !isFullyCollapsed() && !isTracking() && !isClosing();
+        return !isFullyCollapsed() && !isTracking() && !isClosing()
+                // Don't try to collapse if on keyguard, as the expansion fraction is 1 in this
+                // case.
+                && !(shadeCollapseActivityLaunchFix() && mExpandedFraction == 1f
+                && mBarState == KEYGUARD);
     }
 
     public void instantCollapse() {
diff --git a/packages/SystemUI/src/com/android/systemui/shade/ui/viewmodel/OverlayShadeViewModel.kt b/packages/SystemUI/src/com/android/systemui/shade/ui/viewmodel/OverlayShadeViewModel.kt
new file mode 100644
index 0000000..b8dd628
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/shade/ui/viewmodel/OverlayShadeViewModel.kt
@@ -0,0 +1,67 @@
+/*
+ * 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.shade.ui.viewmodel
+
+import com.android.compose.animation.scene.SceneKey
+import com.android.systemui.dagger.SysUISingleton
+import com.android.systemui.dagger.qualifiers.Application
+import com.android.systemui.deviceentry.domain.interactor.DeviceEntryInteractor
+import com.android.systemui.scene.domain.interactor.SceneInteractor
+import com.android.systemui.scene.shared.model.Scenes
+import javax.inject.Inject
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.flow.SharingStarted
+import kotlinx.coroutines.flow.StateFlow
+import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.flow.stateIn
+
+/**
+ * Models UI state and handles user input for the overlay shade UI, which shows a shade as an
+ * overlay on top of another scene UI.
+ */
+@OptIn(ExperimentalCoroutinesApi::class)
+@SysUISingleton
+class OverlayShadeViewModel
+@Inject
+constructor(
+    @Application private val applicationScope: CoroutineScope,
+    private val sceneInteractor: SceneInteractor,
+    deviceEntryInteractor: DeviceEntryInteractor,
+) {
+    /** The scene to show in the background when the overlay shade is open. */
+    val backgroundScene: StateFlow<SceneKey> =
+        deviceEntryInteractor.isDeviceEntered
+            .map(::backgroundScene)
+            .stateIn(
+                scope = applicationScope,
+                started = SharingStarted.WhileSubscribed(),
+                initialValue = backgroundScene(deviceEntryInteractor.isDeviceEntered.value)
+            )
+
+    /** Notifies that the user has clicked the semi-transparent background scrim. */
+    fun onScrimClicked() {
+        sceneInteractor.changeScene(
+            toScene = backgroundScene.value,
+            loggingReason = "Shade scrim clicked",
+        )
+    }
+
+    private fun backgroundScene(isDeviceEntered: Boolean): SceneKey {
+        return if (isDeviceEntered) Scenes.Gone else Scenes.Lockscreen
+    }
+}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/qs/panels/ui/viewmodel/TileGridViewModelTest.kt b/packages/SystemUI/tests/src/com/android/systemui/qs/panels/ui/viewmodel/TileGridViewModelTest.kt
deleted file mode 100644
index e8c5fd9..0000000
--- a/packages/SystemUI/tests/src/com/android/systemui/qs/panels/ui/viewmodel/TileGridViewModelTest.kt
+++ /dev/null
@@ -1,84 +0,0 @@
-/*
- * 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.qs.panels.ui.viewmodel
-
-import android.testing.AndroidTestingRunner
-import androidx.test.filters.SmallTest
-import com.android.systemui.SysuiTestCase
-import com.android.systemui.coroutines.collectLastValue
-import com.android.systemui.kosmos.testScope
-import com.android.systemui.plugins.qs.QSTile
-import com.android.systemui.qs.FakeQSFactory
-import com.android.systemui.qs.pipeline.data.repository.tileSpecRepository
-import com.android.systemui.qs.pipeline.domain.interactor.FakeQSTile
-import com.android.systemui.qs.pipeline.shared.TileSpec
-import com.android.systemui.qs.qsTileFactory
-import com.android.systemui.testKosmos
-import kotlinx.coroutines.test.runTest
-import org.junit.Assert
-import org.junit.Test
-import org.junit.runner.RunWith
-
-@SmallTest
-@RunWith(AndroidTestingRunner::class)
-class TileGridViewModelTest : SysuiTestCase() {
-
-    private val kosmos = testKosmos().apply { qsTileFactory = FakeQSFactory(::tileCreator) }
-    private val underTest = with(kosmos) { tileGridViewModel }
-
-    @Test
-    fun noIconTiles() =
-        with(kosmos) {
-            testScope.runTest {
-                val latest by collectLastValue(underTest.tileViewModels)
-
-                tileSpecRepository.setTiles(
-                    0,
-                    listOf(
-                        TileSpec.create("bluetooth"),
-                        TileSpec.create("internet"),
-                        TileSpec.create("alarm")
-                    )
-                )
-
-                latest!!.forEach { Assert.assertFalse(it.iconOnly) }
-            }
-        }
-
-    @Test
-    fun withIconTiles() =
-        with(kosmos) {
-            testScope.runTest {
-                val latest by collectLastValue(underTest.tileViewModels)
-
-                tileSpecRepository.setTiles(
-                    0,
-                    listOf(
-                        TileSpec.create("airplane"),
-                        TileSpec.create("flashlight"),
-                        TileSpec.create("rotation")
-                    )
-                )
-
-                latest!!.forEach { Assert.assertTrue(it.iconOnly) }
-            }
-        }
-
-    private fun tileCreator(spec: String): QSTile {
-        return FakeQSTile(0).apply { tileSpec = spec }
-    }
-}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/shade/NotificationPanelViewControllerTest.java b/packages/SystemUI/tests/src/com/android/systemui/shade/NotificationPanelViewControllerTest.java
index 81e20c1..6b57f6e 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/shade/NotificationPanelViewControllerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/shade/NotificationPanelViewControllerTest.java
@@ -18,6 +18,7 @@
 
 import static com.android.keyguard.KeyguardClockSwitch.LARGE;
 import static com.android.keyguard.KeyguardClockSwitch.SMALL;
+import static com.android.systemui.Flags.FLAG_SHADE_COLLAPSE_ACTIVITY_LAUNCH_FIX;
 import static com.android.systemui.shade.ShadeExpansionStateManagerKt.STATE_CLOSED;
 import static com.android.systemui.shade.ShadeExpansionStateManagerKt.STATE_OPEN;
 import static com.android.systemui.shade.ShadeExpansionStateManagerKt.STATE_OPENING;
@@ -47,6 +48,7 @@
 import android.graphics.Point;
 import android.os.PowerManager;
 import android.platform.test.annotations.DisableFlags;
+import android.platform.test.annotations.EnableFlags;
 import android.testing.AndroidTestingRunner;
 import android.testing.TestableLooper;
 import android.view.MotionEvent;
@@ -675,6 +677,32 @@
     }
 
     @Test
+    @EnableFlags(FLAG_SHADE_COLLAPSE_ACTIVITY_LAUNCH_FIX)
+    public void testCanBeCollapsed_expandedInKeyguard() {
+        mStatusBarStateController.setState(KEYGUARD);
+        mNotificationPanelViewController.setExpandedFraction(1f);
+
+        assertThat(mNotificationPanelViewController.canBeCollapsed()).isFalse();
+    }
+
+    @Test
+    @EnableFlags(FLAG_SHADE_COLLAPSE_ACTIVITY_LAUNCH_FIX)
+    public void testCanBeCollapsed_expandedInShade() {
+        mStatusBarStateController.setState(SHADE);
+        mNotificationPanelViewController.setExpandedFraction(1f);
+        assertThat(mNotificationPanelViewController.canBeCollapsed()).isTrue();
+    }
+
+    @Test
+    @DisableFlags(FLAG_SHADE_COLLAPSE_ACTIVITY_LAUNCH_FIX)
+    public void testCanBeCollapsed_expandedInKeyguard_flagDisabled() {
+        mStatusBarStateController.setState(KEYGUARD);
+        mNotificationPanelViewController.setExpandedFraction(1f);
+
+        assertThat(mNotificationPanelViewController.canBeCollapsed()).isTrue();
+    }
+
+    @Test
     public void testSwipeWhileLocked_notifiesKeyguardState() {
         mStatusBarStateController.setState(KEYGUARD);
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/shade/QuickSettingsControllerImplBaseTest.java b/packages/SystemUI/tests/src/com/android/systemui/shade/QuickSettingsControllerImplBaseTest.java
index 317e35c..04fa590 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/shade/QuickSettingsControllerImplBaseTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/shade/QuickSettingsControllerImplBaseTest.java
@@ -216,6 +216,7 @@
                 keyguardTransitionInteractor,
                 () -> sceneInteractor,
                 () -> mKosmos.getFromGoneTransitionInteractor(),
+                () -> mKosmos.getFromLockscreenTransitionInteractor(),
                 () -> mKosmos.getSharedNotificationContainerInteractor(),
                 mTestScope);
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/KeyguardStatusBarViewControllerTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/KeyguardStatusBarViewControllerTest.java
index 8e9840a..dfee737 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/KeyguardStatusBarViewControllerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/KeyguardStatusBarViewControllerTest.java
@@ -178,6 +178,7 @@
                 keyguardTransitionInteractor,
                 () -> mKosmos.getSceneInteractor(),
                 () -> mKosmos.getFromGoneTransitionInteractor(),
+                () -> mKosmos.getFromLockscreenTransitionInteractor(),
                 () -> mKosmos.getSharedNotificationContainerInteractor(),
                 mTestScope);
         mViewModel =
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractorFactory.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractorFactory.kt
index 2e751cc..02842cc 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractorFactory.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractorFactory.kt
@@ -53,6 +53,7 @@
         shadeRepository: FakeShadeRepository = FakeShadeRepository(),
         sceneInteractor: SceneInteractor = mock(),
         fromGoneTransitionInteractor: FromGoneTransitionInteractor = mock(),
+        fromLockscreenTransitionInteractor: FromLockscreenTransitionInteractor = mock(),
         sharedNotificationContainerInteractor: SharedNotificationContainerInteractor? = null,
         powerInteractor: PowerInteractor = PowerInteractorFactory.create().powerInteractor,
         testScope: CoroutineScope = TestScope(),
@@ -98,6 +99,7 @@
                 keyguardTransitionInteractor = keyguardTransitionInteractor,
                 sceneInteractorProvider = { sceneInteractor },
                 fromGoneTransitionInteractor = { fromGoneTransitionInteractor },
+                fromLockscreenTransitionInteractor = { fromLockscreenTransitionInteractor },
                 sharedNotificationContainerInteractor = { sncInteractor },
                 applicationScope = testScope,
             ),
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractorKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractorKosmos.kt
index 9426718..81d8f0b 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractorKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractorKosmos.kt
@@ -18,20 +18,20 @@
 
 import com.android.systemui.bouncer.data.repository.keyguardBouncerRepository
 import com.android.systemui.common.ui.domain.interactor.configurationInteractor
+import com.android.systemui.keyguard.data.repository.fakeCommandQueue
 import com.android.systemui.keyguard.data.repository.keyguardRepository
 import com.android.systemui.kosmos.Kosmos
 import com.android.systemui.kosmos.testScope
 import com.android.systemui.power.domain.interactor.powerInteractor
 import com.android.systemui.scene.domain.interactor.sceneInteractor
 import com.android.systemui.shade.data.repository.shadeRepository
-import com.android.systemui.statusbar.commandQueue
 import com.android.systemui.statusbar.notification.stack.domain.interactor.sharedNotificationContainerInteractor
 
 val Kosmos.keyguardInteractor: KeyguardInteractor by
     Kosmos.Fixture {
         KeyguardInteractor(
             repository = keyguardRepository,
-            commandQueue = commandQueue,
+            commandQueue = fakeCommandQueue,
             powerInteractor = powerInteractor,
             bouncerRepository = keyguardBouncerRepository,
             configurationInteractor = configurationInteractor,
@@ -39,6 +39,7 @@
             keyguardTransitionInteractor = keyguardTransitionInteractor,
             sceneInteractorProvider = { sceneInteractor },
             fromGoneTransitionInteractor = { fromGoneTransitionInteractor },
+            fromLockscreenTransitionInteractor = { fromLockscreenTransitionInteractor },
             sharedNotificationContainerInteractor = { sharedNotificationContainerInteractor },
             applicationScope = testScope.backgroundScope,
         )
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/domain/interactor/GridLayoutTypeInteractorKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/domain/interactor/GridLayoutTypeInteractorKosmos.kt
index e44fa07..c951642 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/domain/interactor/GridLayoutTypeInteractorKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/domain/interactor/GridLayoutTypeInteractorKosmos.kt
@@ -18,6 +18,12 @@
 
 import com.android.systemui.kosmos.Kosmos
 import com.android.systemui.qs.panels.data.repository.gridLayoutTypeRepository
+import com.android.systemui.qs.panels.shared.model.GridLayoutType
+import com.android.systemui.qs.panels.shared.model.InfiniteGridLayoutType
+import com.android.systemui.qs.panels.ui.compose.GridLayout
 
 val Kosmos.gridLayoutTypeInteractor by
     Kosmos.Fixture { GridLayoutTypeInteractor(gridLayoutTypeRepository) }
+
+val Kosmos.gridLayoutMap: Map<GridLayoutType, GridLayout> by
+    Kosmos.Fixture { mapOf(Pair(InfiniteGridLayoutType, infiniteGridLayout)) }
diff --git a/packages/SystemUI/src/com/android/systemui/qs/panels/shared/model/GridLayoutTypeKey.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/domain/interactor/InfiniteGridLayoutKosmos.kt
similarity index 63%
rename from packages/SystemUI/src/com/android/systemui/qs/panels/shared/model/GridLayoutTypeKey.kt
rename to packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/domain/interactor/InfiniteGridLayoutKosmos.kt
index 0dbaaba..1893c30 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/panels/shared/model/GridLayoutTypeKey.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/domain/interactor/InfiniteGridLayoutKosmos.kt
@@ -14,15 +14,9 @@
  * limitations under the License.
  */
 
-package com.android.systemui.qs.panels.shared.model
+package com.android.systemui.qs.panels.domain.interactor
 
-import dagger.MapKey
-import kotlin.reflect.KClass
+import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.qs.panels.ui.compose.InfiniteGridLayout
 
-/**
- * Dagger map key to associate a [GridLayoutType] with its
- * [com.android.systemui.qs.panels.ui.compose.GridLayout].
- */
-@Retention(AnnotationRetention.RUNTIME)
-@MapKey
-annotation class GridLayoutTypeKey(val value: KClass<out GridLayoutType>)
+val Kosmos.infiniteGridLayout by Kosmos.Fixture { InfiniteGridLayout(iconTilesInteractor) }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/ui/viewmodel/TileGridViewModelKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/ui/viewmodel/TileGridViewModelKosmos.kt
index 9851f04..5fd8762 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/ui/viewmodel/TileGridViewModelKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/ui/viewmodel/TileGridViewModelKosmos.kt
@@ -17,9 +17,10 @@
 package com.android.systemui.qs.panels.ui.viewmodel
 
 import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.kosmos.applicationCoroutineScope
+import com.android.systemui.qs.panels.domain.interactor.gridLayoutMap
 import com.android.systemui.qs.panels.domain.interactor.gridLayoutTypeInteractor
 import com.android.systemui.qs.panels.domain.interactor.iconTilesInteractor
-import com.android.systemui.qs.panels.shared.model.InfiniteGridLayoutType
 import com.android.systemui.qs.panels.ui.compose.InfiniteGridLayout
 import com.android.systemui.qs.pipeline.domain.interactor.currentTilesInteractor
 
@@ -27,8 +28,9 @@
     Kosmos.Fixture {
         TileGridViewModel(
             gridLayoutTypeInteractor,
-            mapOf(Pair(InfiniteGridLayoutType::class.java, InfiniteGridLayout())),
+            gridLayoutMap,
             currentTilesInteractor,
-            iconTilesInteractor
+            InfiniteGridLayout(iconTilesInteractor),
+            applicationCoroutineScope,
         )
     }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/shade/ui/viewmodel/OverlayShadeViewModelKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/shade/ui/viewmodel/OverlayShadeViewModelKosmos.kt
new file mode 100644
index 0000000..45ec032
--- /dev/null
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/shade/ui/viewmodel/OverlayShadeViewModelKosmos.kt
@@ -0,0 +1,34 @@
+/*
+ * 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.
+ */
+
+@file:OptIn(ExperimentalCoroutinesApi::class)
+
+package com.android.systemui.shade.ui.viewmodel
+
+import com.android.systemui.deviceentry.domain.interactor.deviceEntryInteractor
+import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.kosmos.applicationCoroutineScope
+import com.android.systemui.scene.domain.interactor.sceneInteractor
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+
+val Kosmos.overlayShadeViewModel: OverlayShadeViewModel by
+    Kosmos.Fixture {
+        OverlayShadeViewModel(
+            applicationScope = applicationCoroutineScope,
+            sceneInteractor = sceneInteractor,
+            deviceEntryInteractor = deviceEntryInteractor,
+        )
+    }
diff --git a/tests/FlickerTests/IME/OWNERS b/tests/FlickerTests/IME/OWNERS
index 301fafa..ae1098d 100644
--- a/tests/FlickerTests/IME/OWNERS
+++ b/tests/FlickerTests/IME/OWNERS
@@ -1,2 +1,3 @@
 # ime
 # Bug component: 34867
+include /services/core/java/com/android/server/inputmethod/OWNERS
diff --git a/tests/FlickerTests/IME/src/com/android/server/wm/flicker/ime/ShowImeOnUnlockScreenTest.kt b/tests/FlickerTests/IME/src/com/android/server/wm/flicker/ime/ShowImeOnUnlockScreenTest.kt
index 31506b5..d22bdcf 100644
--- a/tests/FlickerTests/IME/src/com/android/server/wm/flicker/ime/ShowImeOnUnlockScreenTest.kt
+++ b/tests/FlickerTests/IME/src/com/android/server/wm/flicker/ime/ShowImeOnUnlockScreenTest.kt
@@ -24,6 +24,7 @@
 import android.tools.flicker.legacy.LegacyFlickerTest
 import android.tools.flicker.legacy.LegacyFlickerTestFactory
 import android.tools.traces.component.ComponentNameMatcher
+import androidx.test.filters.FlakyTest
 import com.android.server.wm.flicker.BaseTest
 import com.android.server.wm.flicker.helpers.ImeAppHelper
 import org.junit.FixMethodOrder
@@ -90,6 +91,11 @@
     @Ignore("Not applicable to this CUJ. Display turns off during transition")
     override fun taskBarWindowIsAlwaysVisible() {}
 
+    @FlakyTest(bugId = 338178020)
+    @Test
+    override fun visibleWindowsShownMoreThanOneConsecutiveEntry() =
+        super.visibleWindowsShownMoreThanOneConsecutiveEntry()
+
     companion object {
         @Parameterized.Parameters(name = "{0}")
         @JvmStatic