diff --git a/apct-tests/perftests/core/src/android/conscrypt/conscrypt/ClientSocketPerfTest.java b/apct-tests/perftests/core/src/android/conscrypt/conscrypt/ClientSocketPerfTest.java
index 9e45c4a..bcc0a3b 100644
--- a/apct-tests/perftests/core/src/android/conscrypt/conscrypt/ClientSocketPerfTest.java
+++ b/apct-tests/perftests/core/src/android/conscrypt/conscrypt/ClientSocketPerfTest.java
@@ -112,36 +112,20 @@
         for (EndpointFactory endpointFactory : EndpointFactory.values()) {
             for (ChannelType channelType : ChannelType.values()) {
                 for (PerfTestProtocol protocol : PerfTestProtocol.values()) {
-                    params.add(
-                            new Object[] {
-                                new Config(
-                                        endpointFactory,
-                                        endpointFactory,
-                                        64,
-                                        "TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256",
-                                        channelType,
-                                        protocol)
-                            });
-                    params.add(
-                            new Object[] {
-                                new Config(
-                                        endpointFactory,
-                                        endpointFactory,
-                                        512,
-                                        "TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256",
-                                        channelType,
-                                        protocol)
-                            });
-                    params.add(
-                            new Object[] {
-                                new Config(
-                                        endpointFactory,
-                                        endpointFactory,
-                                        4096,
-                                        "TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256",
-                                        channelType,
-                                        protocol)
-                            });
+                    for (int messageSize : ConscryptParams.messageSizes) {
+                        for (String cipher : ConscryptParams.ciphers) {
+                            params.add(
+                                    new Object[] {
+                                            new Config(
+                                                    endpointFactory,
+                                                    endpointFactory,
+                                                    messageSize,
+                                                    cipher,
+                                                    channelType,
+                                                    protocol)
+                                    });
+                        }
+                    }
                 }
             }
         }
diff --git a/apct-tests/perftests/core/src/android/conscrypt/conscrypt/ConscryptParams.java b/apct-tests/perftests/core/src/android/conscrypt/conscrypt/ConscryptParams.java
new file mode 100644
index 0000000..e5131b8
--- /dev/null
+++ b/apct-tests/perftests/core/src/android/conscrypt/conscrypt/ConscryptParams.java
@@ -0,0 +1,35 @@
+/*
+ * Copyright 2025 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 android.conscrypt;
+
+import java.util.List;
+
+public class ConscryptParams {
+    public static final List<String> ciphers = List.of(
+        "TLS_RSA_WITH_AES_128_GCM_SHA256",
+        "TLS_RSA_WITH_AES_256_GCM_SHA384",
+        "TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256",
+        "TLS_ECDHE_RSA_WITH_AES_256_GCM_SHA384",
+        "TLS_ECDHE_RSA_WITH_CHACHA20_POLY1305_SHA256"
+    );
+
+    public static final List<Integer> messageSizes = List.of(
+        64,
+        512,
+        4096
+    );
+}
\ No newline at end of file
diff --git a/apct-tests/perftests/core/src/android/conscrypt/conscrypt/EngineHandshakePerfTest.java b/apct-tests/perftests/core/src/android/conscrypt/conscrypt/EngineHandshakePerfTest.java
index cd0ac96..341d8e6 100644
--- a/apct-tests/perftests/core/src/android/conscrypt/conscrypt/EngineHandshakePerfTest.java
+++ b/apct-tests/perftests/core/src/android/conscrypt/conscrypt/EngineHandshakePerfTest.java
@@ -87,11 +87,13 @@
         }
     }
 
+
     public Collection getParams() {
         final List<Object[]> params = new ArrayList<>();
         for (BufferType bufferType : BufferType.values()) {
-            params.add(new Object[] {new Config(bufferType,
-                "TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256", 100)});
+            for (String cipher : ConscryptParams.ciphers) {
+                params.add(new Object[] {new Config(bufferType, cipher, 100)});
+            }
         }
         return params;
     }
diff --git a/apct-tests/perftests/core/src/android/conscrypt/conscrypt/EngineWrapPerfTest.java b/apct-tests/perftests/core/src/android/conscrypt/conscrypt/EngineWrapPerfTest.java
index 1fee218..23b642e 100644
--- a/apct-tests/perftests/core/src/android/conscrypt/conscrypt/EngineWrapPerfTest.java
+++ b/apct-tests/perftests/core/src/android/conscrypt/conscrypt/EngineWrapPerfTest.java
@@ -37,10 +37,10 @@
 import static org.junit.Assert.assertEquals;
 
 import java.nio.ByteBuffer;
-import java.util.Locale;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.List;
+import java.util.Locale;
 import javax.net.ssl.SSLEngine;
 import javax.net.ssl.SSLEngineResult;
 import javax.net.ssl.SSLException;
@@ -94,12 +94,11 @@
     public Collection getParams() {
         final List<Object[]> params = new ArrayList<>();
         for (BufferType bufferType : BufferType.values()) {
-            params.add(new Object[] {new Config(bufferType, 64,
-                                    "TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256")});
-            params.add(new Object[] {new Config(bufferType, 512,
-                                    "TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256")});
-            params.add(new Object[] {new Config(bufferType, 4096,
-                                    "TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256")});
+            for (int messageSize : ConscryptParams.messageSizes) {
+                for (String cipher : ConscryptParams.ciphers) {
+                    params.add(new Object[] {new Config(bufferType, messageSize, cipher)});
+                }
+            }
         }
         return params;
     }
diff --git a/apct-tests/perftests/core/src/android/conscrypt/conscrypt/ServerSocketPerfTest.java b/apct-tests/perftests/core/src/android/conscrypt/conscrypt/ServerSocketPerfTest.java
index 90a87ce..343bb12 100644
--- a/apct-tests/perftests/core/src/android/conscrypt/conscrypt/ServerSocketPerfTest.java
+++ b/apct-tests/perftests/core/src/android/conscrypt/conscrypt/ServerSocketPerfTest.java
@@ -102,15 +102,12 @@
         final List<Object[]> params = new ArrayList<>();
         for (EndpointFactory endpointFactory : EndpointFactory.values()) {
             for (ChannelType channelType : ChannelType.values()) {
-                params.add(new Object[] {new Config(endpointFactory,
-                    endpointFactory, 64,
-                    "TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256", channelType)});
-                params.add(new Object[] {new Config(endpointFactory,
-                    endpointFactory, 512,
-                    "TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256", channelType)});
-                params.add(new Object[] {new Config(endpointFactory,
-                    endpointFactory, 4096,
-                    "TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256", channelType)});
+                for (int messageSize : ConscryptParams.messageSizes) {
+                    for (String cipher : ConscryptParams.ciphers) {
+                        params.add(new Object[] {new Config(endpointFactory,
+                            endpointFactory, messageSize, cipher, channelType)});
+                    }
+                }
             }
         }
         return params;
diff --git a/core/java/android/view/InsetsController.java b/core/java/android/view/InsetsController.java
index 3659e78..0d82acd 100644
--- a/core/java/android/view/InsetsController.java
+++ b/core/java/android/view/InsetsController.java
@@ -1578,11 +1578,7 @@
             Trace.asyncTraceBegin(TRACE_TAG_VIEW, "IC.pendingAnim", 0);
         }
 
-        if (Flags.refactorInsetsController()) {
-            onAnimationStateChanged(typesReady, true /* running */);
-        } else {
-            onAnimationStateChanged(types, true /* running */);
-        }
+        onAnimationStateChanged(types, true /* running */);
 
         if (fromIme) {
             switch (animationType) {
diff --git a/libs/WindowManager/Shell/res/layout/caption_window_decor.xml b/libs/WindowManager/Shell/res/layout/caption_window_decor.xml
index f3d2198..819d4ab 100644
--- a/libs/WindowManager/Shell/res/layout/caption_window_decor.xml
+++ b/libs/WindowManager/Shell/res/layout/caption_window_decor.xml
@@ -37,20 +37,17 @@
         style="@style/CaptionButtonStyle"
         android:id="@+id/minimize_window"
         android:layout_gravity="center_vertical|end"
-        android:contentDescription="@string/minimize_button_text"
         android:background="@drawable/decor_minimize_button_dark"
         android:duplicateParentState="true"/>
     <Button
         style="@style/CaptionButtonStyle"
         android:id="@+id/maximize_window"
         android:layout_gravity="center_vertical|end"
-        android:contentDescription="@string/maximize_button_text"
         android:background="@drawable/decor_maximize_button_dark"
         android:duplicateParentState="true"/>
     <Button
         style="@style/CaptionButtonStyle"
         android:id="@+id/close_window"
-        android:contentDescription="@string/close_button_text"
         android:background="@drawable/decor_close_button_dark"
         android:duplicateParentState="true"/>
 </com.android.wm.shell.windowdecor.WindowDecorLinearLayout>
\ No newline at end of file
diff --git a/libs/WindowManager/Shell/res/values/strings.xml b/libs/WindowManager/Shell/res/values/strings.xml
index f576549..2179128 100644
--- a/libs/WindowManager/Shell/res/values/strings.xml
+++ b/libs/WindowManager/Shell/res/values/strings.xml
@@ -277,11 +277,13 @@
 
     <!-- Freeform window caption strings -->
     <!-- Accessibility text for the maximize window button [CHAR LIMIT=NONE] -->
-    <string name="maximize_button_text">Maximize</string>
+    <string name="maximize_button_text">Maximize <xliff:g id="app_name" example="Chrome">%1$s</xliff:g></string>
+    <!-- Accessibility text for the restore window button [CHAR LIMIT=NONE] -->
+    <string name="restore_button_text">Restore <xliff:g id="app_name" example="Chrome">%1$s</xliff:g></string>
      <!-- Accessibility text for the minimize window button [CHAR LIMIT=NONE] -->
-     <string name="minimize_button_text">Minimize</string>
+     <string name="minimize_button_text">Minimize <xliff:g id="app_name" example="Chrome">%1$s</xliff:g></string>
     <!-- Accessibility text for the close window button [CHAR LIMIT=NONE] -->
-    <string name="close_button_text">Close</string>
+    <string name="close_button_text">Close <xliff:g id="app_name" example="Chrome">%1$s</xliff:g></string>
     <!-- Accessibility text for the caption back button [CHAR LIMIT=NONE] -->
     <string name="back_button_text">Back</string>
     <!-- Accessibility text for the caption handle [CHAR LIMIT=NONE] -->
@@ -353,10 +355,14 @@
     <string name="maximize_menu_talkback_action_snap_right_text">Resize window to right</string>
     <!-- Accessibility action replacement for maximize menu enter maximize/restore button [CHAR LIMIT=NONE] -->
     <string name="maximize_menu_talkback_action_maximize_restore_text">Maximize or restore window size</string>
-    <!-- Accessibility action replacement for app header maximize/restore button [CHAR LIMIT=NONE] -->
-    <string name="maximize_button_talkback_action_maximize_restore_text">Maximize or restore window size</string>
+    <!-- Accessibility action replacement for app header maximize button [CHAR LIMIT=NONE] -->
+    <string name="app_header_talkback_action_maximize_button_text">Maximize app window size</string>
+    <!-- Accessibility action replacement for app header restore button [CHAR LIMIT=NONE] -->
+    <string name="app_header_talkback_action_restore_button_text">Restore window size</string>
     <!-- Accessibility action replacement for app header minimize button [CHAR LIMIT=NONE] -->
-    <string name="minimize_button_talkback_action_maximize_restore_text">Minimize app window</string>
+    <string name="app_header_talkback_action_minimize_button_text">Minimize app window</string>
+    <!-- Accessibility action replacement for app header close button [CHAR LIMIT=NONE] -->
+    <string name="app_header_talkback_action_close_button_text">Close app window</string>
 
     <!-- Accessibility text for open by default settings button [CHAR LIMIT=NONE] -->
     <string name="open_by_default_settings_text">Open by default settings</string>
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopModeVisualIndicator.java b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopModeVisualIndicator.java
index 99f05283..56de48d 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopModeVisualIndicator.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopModeVisualIndicator.java
@@ -49,6 +49,7 @@
 import com.android.wm.shell.shared.bubbles.BubbleAnythingFlagHelper;
 import com.android.wm.shell.shared.bubbles.BubbleDropTargetBoundsProvider;
 import com.android.wm.shell.shared.desktopmode.DesktopModeStatus;
+import com.android.wm.shell.windowdecor.tiling.SnapEventHandler;
 
 /**
  * Animated visual indicator for Desktop Mode windowing transitions.
@@ -98,7 +99,9 @@
                 return FROM_SPLIT;
             } else if (taskInfo.isFreeform()) {
                 return FROM_FREEFORM;
-            } else return null;
+            } else {
+                return null;
+            }
         }
     }
 
@@ -110,6 +113,7 @@
 
     private IndicatorType mCurrentType;
     private final DragStartState mDragStartState;
+    private final SnapEventHandler mSnapEventHandler;
 
     public DesktopModeVisualIndicator(@ShellDesktopThread ShellExecutor desktopExecutor,
             @ShellMainThread ShellExecutor mainExecutor,
@@ -118,18 +122,20 @@
             Context context, SurfaceControl taskSurface,
             RootTaskDisplayAreaOrganizer taskDisplayAreaOrganizer,
             DragStartState dragStartState,
-            @Nullable BubbleDropTargetBoundsProvider bubbleBoundsProvider) {
+            @Nullable BubbleDropTargetBoundsProvider bubbleBoundsProvider,
+            SnapEventHandler snapEventHandler) {
         SurfaceControl.Builder builder = new SurfaceControl.Builder();
         taskDisplayAreaOrganizer.attachToDisplayArea(taskInfo.displayId, builder);
         mVisualIndicatorViewContainer = new VisualIndicatorViewContainer(
                 DesktopModeFlags.ENABLE_DESKTOP_INDICATOR_IN_SEPARATE_THREAD_BUGFIX.isTrue()
                         ? desktopExecutor : mainExecutor,
-                mainExecutor, builder, syncQueue, bubbleBoundsProvider);
+                mainExecutor, builder, syncQueue, bubbleBoundsProvider, snapEventHandler);
         mTaskInfo = taskInfo;
         mDisplayController = displayController;
         mContext = context;
         mCurrentType = NO_INDICATOR;
         mDragStartState = dragStartState;
+        mSnapEventHandler = snapEventHandler;
         mVisualIndicatorViewContainer.createView(
                 mContext,
                 mDisplayController.getDisplay(mTaskInfo.displayId),
@@ -143,7 +149,8 @@
     public void fadeOutIndicator(
             @NonNull Runnable callback) {
         mVisualIndicatorViewContainer.fadeOutIndicator(
-                mDisplayController.getDisplayLayout(mTaskInfo.displayId), mCurrentType, callback
+                mDisplayController.getDisplayLayout(mTaskInfo.displayId), mCurrentType, callback,
+                mTaskInfo.displayId, mSnapEventHandler
         );
     }
 
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksController.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksController.kt
index 1c88056..2d9aea0 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksController.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksController.kt
@@ -3051,6 +3051,7 @@
                     rootTaskDisplayAreaOrganizer,
                     dragStartState,
                     bubbleController.getOrNull()?.bubbleDropTargetBoundsProvider,
+                    snapEventHandler,
                 )
         if (visualIndicator == null) visualIndicator = indicator
         return indicator.updateIndicatorType(PointF(inputX, taskTop))
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/VisualIndicatorViewContainer.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/VisualIndicatorViewContainer.kt
index 2317274..919e816 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/VisualIndicatorViewContainer.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/VisualIndicatorViewContainer.kt
@@ -44,6 +44,7 @@
 import com.android.wm.shell.shared.annotations.ShellMainThread
 import com.android.wm.shell.shared.bubbles.BubbleDropTargetBoundsProvider
 import com.android.wm.shell.windowdecor.WindowDecoration.SurfaceControlViewHostFactory
+import com.android.wm.shell.windowdecor.tiling.SnapEventHandler
 
 /**
  * Container for the view / viewhost of the indicator, ensuring it is created / animated off the
@@ -60,6 +61,7 @@
     private val surfaceControlViewHostFactory: SurfaceControlViewHostFactory =
         object : SurfaceControlViewHostFactory {},
     private val bubbleBoundsProvider: BubbleDropTargetBoundsProvider?,
+    private val snapEventHandler: SnapEventHandler,
 ) {
     @VisibleForTesting var indicatorView: View? = null
     private var indicatorViewHost: SurfaceControlViewHost? = null
@@ -164,9 +166,15 @@
                 displayController.getDisplayLayout(taskInfo.displayId)
                     ?: error("Expected to find DisplayLayout for taskId${taskInfo.taskId}.")
             if (currentType == IndicatorType.NO_INDICATOR) {
-                fadeInIndicator(layout, newType)
+                fadeInIndicator(layout, newType, taskInfo.displayId, snapEventHandler)
             } else if (newType == IndicatorType.NO_INDICATOR) {
-                fadeOutIndicator(layout, currentType, /* finishCallback= */ null)
+                fadeOutIndicator(
+                    layout,
+                    currentType,
+                    /* finishCallback= */ null,
+                    taskInfo.displayId,
+                    snapEventHandler,
+                )
             } else {
                 val animStartType = IndicatorType.valueOf(currentType.name)
                 val animator =
@@ -177,6 +185,8 @@
                             animStartType,
                             newType,
                             bubbleBoundsProvider,
+                            taskInfo.displayId,
+                            snapEventHandler,
                         )
                     } ?: return@execute
                 animator.start()
@@ -188,12 +198,24 @@
      * Fade indicator in as provided type. Animator fades it in while expanding the bounds outwards.
      */
     @VisibleForTesting
-    fun fadeInIndicator(layout: DisplayLayout, type: IndicatorType) {
+    fun fadeInIndicator(
+        layout: DisplayLayout,
+        type: IndicatorType,
+        displayId: Int,
+        snapEventHandler: SnapEventHandler,
+    ) {
         desktopExecutor.assertCurrentThread()
         indicatorView?.let {
             it.setBackgroundResource(R.drawable.desktop_windowing_transition_background)
             val animator =
-                VisualIndicatorAnimator.fadeBoundsIn(it, type, layout, bubbleBoundsProvider)
+                VisualIndicatorAnimator.fadeBoundsIn(
+                    it,
+                    type,
+                    layout,
+                    bubbleBoundsProvider,
+                    displayId,
+                    snapEventHandler,
+                )
             animator.start()
         }
     }
@@ -207,6 +229,8 @@
         layout: DisplayLayout,
         currentType: IndicatorType,
         finishCallback: Runnable?,
+        displayId: Int,
+        snapEventHandler: SnapEventHandler,
     ) {
         if (currentType == IndicatorType.NO_INDICATOR) {
             // In rare cases, fade out can be requested before the indicator has determined its
@@ -223,6 +247,8 @@
                         animStartType,
                         layout,
                         bubbleBoundsProvider,
+                        displayId,
+                        snapEventHandler,
                     )
                 animator.addListener(
                     object : AnimatorListenerAdapter() {
@@ -328,8 +354,17 @@
                 type: IndicatorType,
                 displayLayout: DisplayLayout,
                 bubbleBoundsProvider: BubbleDropTargetBoundsProvider?,
+                displayId: Int,
+                snapEventHandler: SnapEventHandler,
             ): VisualIndicatorAnimator {
-                val endBounds = getIndicatorBounds(displayLayout, type, bubbleBoundsProvider)
+                val endBounds =
+                    getIndicatorBounds(
+                        displayLayout,
+                        type,
+                        bubbleBoundsProvider,
+                        displayId,
+                        snapEventHandler,
+                    )
                 val startBounds = getMinBounds(endBounds)
                 view.background.bounds = startBounds
 
@@ -345,11 +380,19 @@
                 type: IndicatorType,
                 displayLayout: DisplayLayout,
                 bubbleBoundsProvider: BubbleDropTargetBoundsProvider?,
+                displayId: Int,
+                snapEventHandler: SnapEventHandler,
             ): VisualIndicatorAnimator {
-                val startBounds = getIndicatorBounds(displayLayout, type, bubbleBoundsProvider)
+                val startBounds =
+                    getIndicatorBounds(
+                        displayLayout,
+                        type,
+                        bubbleBoundsProvider,
+                        displayId,
+                        snapEventHandler,
+                    )
                 val endBounds = getMinBounds(startBounds)
                 view.background.bounds = startBounds
-
                 val animator = VisualIndicatorAnimator(view, startBounds, endBounds)
                 animator.interpolator = DecelerateInterpolator()
                 setupIndicatorAnimation(animator, AlphaAnimType.ALPHA_FADE_OUT_ANIM)
@@ -375,9 +418,25 @@
                 origType: IndicatorType,
                 newType: IndicatorType,
                 bubbleBoundsProvider: BubbleDropTargetBoundsProvider?,
+                displayId: Int,
+                snapEventHandler: SnapEventHandler,
             ): VisualIndicatorAnimator {
-                val startBounds = getIndicatorBounds(displayLayout, origType, bubbleBoundsProvider)
-                val endBounds = getIndicatorBounds(displayLayout, newType, bubbleBoundsProvider)
+                val startBounds =
+                    getIndicatorBounds(
+                        displayLayout,
+                        origType,
+                        bubbleBoundsProvider,
+                        displayId,
+                        snapEventHandler,
+                    )
+                val endBounds =
+                    getIndicatorBounds(
+                        displayLayout,
+                        newType,
+                        bubbleBoundsProvider,
+                        displayId,
+                        snapEventHandler,
+                    )
                 val animator = VisualIndicatorAnimator(view, startBounds, endBounds)
                 animator.interpolator = DecelerateInterpolator()
                 setupIndicatorAnimation(animator, AlphaAnimType.ALPHA_NO_CHANGE_ANIM)
@@ -389,6 +448,8 @@
                 layout: DisplayLayout,
                 type: IndicatorType,
                 bubbleBoundsProvider: BubbleDropTargetBoundsProvider?,
+                displayId: Int,
+                snapEventHandler: SnapEventHandler,
             ): Rect {
                 val desktopStableBounds = Rect()
                 layout.getStableBounds(desktopStableBounds)
@@ -417,21 +478,25 @@
                         )
                     }
 
-                    IndicatorType.TO_SPLIT_LEFT_INDICATOR ->
+                    IndicatorType.TO_SPLIT_LEFT_INDICATOR -> {
+                        val currentLeftBounds = snapEventHandler.getLeftSnapBoundsIfTiled(displayId)
                         return Rect(
                             padding,
                             padding,
-                            desktopStableBounds.width() / 2 - padding,
+                            currentLeftBounds.right - padding,
                             desktopStableBounds.height(),
                         )
-
-                    IndicatorType.TO_SPLIT_RIGHT_INDICATOR ->
+                    }
+                    IndicatorType.TO_SPLIT_RIGHT_INDICATOR -> {
+                        val currentRightBounds =
+                            snapEventHandler.getRightSnapBoundsIfTiled(displayId)
                         return Rect(
-                            desktopStableBounds.width() / 2 + padding,
+                            currentRightBounds.left + padding,
                             padding,
                             desktopStableBounds.width() - padding,
                             desktopStableBounds.height(),
                         )
+                    }
                     IndicatorType.TO_BUBBLE_LEFT_INDICATOR ->
                         return bubbleBoundsProvider?.getBubbleBarExpandedViewDropTargetBounds(
                             /* onLeft= */ true
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorViewModel.java b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorViewModel.java
index d9afd15..0d773ec 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorViewModel.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorViewModel.java
@@ -159,6 +159,8 @@
 import kotlin.Unit;
 import kotlin.jvm.functions.Function1;
 
+import org.jetbrains.annotations.NotNull;
+
 import kotlinx.coroutines.CoroutineScope;
 import kotlinx.coroutines.ExperimentalCoroutinesApi;
 import kotlinx.coroutines.MainCoroutineDispatcher;
@@ -935,6 +937,18 @@
         return mDesktopTilingDecorViewModel.moveTaskToFrontIfTiled(taskInfo);
     }
 
+    @Override
+    @NotNull
+    public Rect getLeftSnapBoundsIfTiled(int displayId) {
+        return mDesktopTilingDecorViewModel.getLeftSnapBoundsIfTiled(displayId);
+    }
+
+    @Override
+    @NotNull
+    public Rect getRightSnapBoundsIfTiled(int displayId) {
+        return mDesktopTilingDecorViewModel.getRightSnapBoundsIfTiled(displayId);
+    }
+
     private class DesktopModeTouchEventListener extends GestureDetector.SimpleOnGestureListener
             implements View.OnClickListener, View.OnTouchListener, View.OnLongClickListener,
             View.OnGenericMotionListener, DragDetector.MotionEventHandler {
@@ -974,7 +988,7 @@
             final int touchSlop = ViewConfiguration.get(mContext).getScaledTouchSlop();
             final long appHandleHoldToDragDuration =
                     DesktopModeFlags.ENABLE_HOLD_TO_DRAG_APP_HANDLE.isTrue()
-                    ? APP_HANDLE_HOLD_TO_DRAG_DURATION_MS : 0;
+                            ? APP_HANDLE_HOLD_TO_DRAG_DURATION_MS : 0;
             mHandleDragDetector = new DragDetector(this, appHandleHoldToDragDuration,
                     touchSlop);
             mHeaderDragDetector = new DragDetector(this, APP_HEADER_HOLD_TO_DRAG_DURATION_MS,
@@ -1027,7 +1041,7 @@
                         decoration.mTaskInfo.userId);
                 if (DesktopModeFlags.ENABLE_FULLY_IMMERSIVE_IN_DESKTOP.isTrue()
                         && desktopRepository.isTaskInFullImmersiveState(
-                                decoration.mTaskInfo.taskId)) {
+                        decoration.mTaskInfo.taskId)) {
                     // Task is in immersive and should exit.
                     onEnterOrExitImmersive(decoration.mTaskInfo);
                 } else {
@@ -1321,6 +1335,7 @@
         /**
          * Perform a task size toggle on release of the double-tap, assuming no drag event
          * was handled during the double-tap.
+         *
          * @param e The motion event that occurred during the double-tap gesture.
          * @return true if the event should be consumed, false if not
          */
@@ -1346,6 +1361,7 @@
     class EventReceiver extends InputEventReceiver {
         private InputMonitor mInputMonitor;
         private int mTasksOnDisplay;
+
         EventReceiver(InputMonitor inputMonitor, InputChannel channel, Looper looper) {
             super(channel, looper);
             mInputMonitor = inputMonitor;
@@ -1397,6 +1413,7 @@
     /**
      * Check if an EventReceiver exists on a particular display.
      * If it does, increment its task count. Otherwise, create one for that display.
+     *
      * @param displayId the display to check against
      */
     private void incrementEventReceiverTasks(int displayId) {
@@ -1902,7 +1919,7 @@
         @Override
         public void onAnimationStart(int taskId, Transaction t, Rect bounds) {
             final DesktopModeWindowDecoration decoration = mWindowDecorByTaskId.get(taskId);
-            if (decoration == null)  {
+            if (decoration == null) {
                 t.apply();
                 return;
             }
@@ -1986,15 +2003,15 @@
             return activityTaskManager.getRecentTasks(Integer.MAX_VALUE,
                     ActivityManager.RECENT_WITH_EXCLUDED,
                     info.userId).getList().stream().filter(
-                            recentTaskInfo -> {
-                                if (recentTaskInfo.taskId == info.taskId) {
-                                    return false;
-                                }
-                                final String recentTaskPackageName =
-                                        ComponentUtils.getPackageName(recentTaskInfo);
-                                return packageName != null
-                                        && packageName.equals(recentTaskPackageName);
-                            }
+                    recentTaskInfo -> {
+                        if (recentTaskInfo.taskId == info.taskId) {
+                            return false;
+                        }
+                        final String recentTaskPackageName =
+                                ComponentUtils.getPackageName(recentTaskInfo);
+                        return packageName != null
+                                && packageName.equals(recentTaskPackageName);
+                    }
             ).toList().size();
         } catch (RemoteException e) {
             throw new RuntimeException(e);
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/MaximizeButtonView.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/MaximizeButtonView.kt
index 581d186..bdde096d 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/MaximizeButtonView.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/MaximizeButtonView.kt
@@ -39,10 +39,7 @@
 private const val OPEN_MAXIMIZE_MENU_DELAY_ON_HOVER_MS = 350
 private const val MAX_DRAWABLE_ALPHA = 255
 
-class MaximizeButtonView(
-        context: Context,
-        attrs: AttributeSet
-) : FrameLayout(context, attrs) {
+class MaximizeButtonView(context: Context, attrs: AttributeSet) : FrameLayout(context, attrs) {
     lateinit var onHoverAnimationFinishedListener: () -> Unit
     private val hoverProgressAnimatorSet = AnimatorSet()
     var hoverDisabled = false
@@ -53,10 +50,6 @@
         (stubProgressBarContainer.inflate() as FrameLayout)
             .requireViewById(R.id.progress_bar)
     }
-    private val maximizeButtonText =
-        context.resources.getString(R.string.desktop_mode_maximize_menu_maximize_button_text)
-    private val restoreButtonText =
-        context.resources.getString(R.string.desktop_mode_maximize_menu_restore_button_text)
 
     init {
         LayoutInflater.from(context).inflate(R.layout.maximize_menu_button, this, true)
@@ -158,12 +151,6 @@
     /** Set the drawable resource to use for the maximize button. */
     fun setIcon(@DrawableRes icon: Int) {
         maximizeWindow.setImageResource(icon)
-        when (icon) {
-            R.drawable.decor_desktop_mode_immersive_or_maximize_exit_button_dark ->
-                maximizeWindow.contentDescription = restoreButtonText
-            R.drawable.decor_desktop_mode_maximize_button_dark ->
-                maximizeWindow.contentDescription = maximizeButtonText
-        }
     }
 
     companion object {
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/tiling/DesktopTilingDecorViewModel.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/tiling/DesktopTilingDecorViewModel.kt
index 8747f63..ee5d0e8 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/tiling/DesktopTilingDecorViewModel.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/tiling/DesktopTilingDecorViewModel.kt
@@ -25,6 +25,7 @@
 import android.window.WindowContainerTransaction
 import androidx.core.util.valueIterator
 import com.android.internal.annotations.VisibleForTesting
+import com.android.wm.shell.R
 import com.android.wm.shell.RootTaskDisplayAreaOrganizer
 import com.android.wm.shell.ShellTaskOrganizer
 import com.android.wm.shell.common.DisplayChangeController
@@ -148,4 +149,45 @@
         if ((fromRotation % 2 == toRotation % 2)) return
         tilingTransitionHandlerByDisplayId.get(displayId)?.resetTilingSession()
     }
+
+    fun getRightSnapBoundsIfTiled(displayId: Int): Rect {
+        val tilingBounds =
+            tilingTransitionHandlerByDisplayId.get(displayId)?.getRightSnapBoundsIfTiled()
+        if (tilingBounds != null) {
+            return tilingBounds
+        }
+        val displayLayout = displayController.getDisplayLayout(displayId)
+        val stableBounds = Rect()
+        displayLayout?.getStableBounds(stableBounds)
+        val snapBounds =
+            Rect(
+                stableBounds.left +
+                    stableBounds.width() / 2 +
+                    context.resources.getDimensionPixelSize(R.dimen.split_divider_bar_width) / 2,
+                stableBounds.top,
+                stableBounds.right,
+                stableBounds.bottom,
+            )
+        return snapBounds
+    }
+
+    fun getLeftSnapBoundsIfTiled(displayId: Int): Rect {
+        val tilingBounds =
+            tilingTransitionHandlerByDisplayId.get(displayId)?.getLeftSnapBoundsIfTiled()
+        if (tilingBounds != null) {
+            return tilingBounds
+        }
+        val displayLayout = displayController.getDisplayLayout(displayId)
+        val stableBounds = Rect()
+        displayLayout?.getStableBounds(stableBounds)
+        val snapBounds =
+            Rect(
+                stableBounds.left,
+                stableBounds.top,
+                stableBounds.left + stableBounds.width() / 2 -
+                    context.resources.getDimensionPixelSize(R.dimen.split_divider_bar_width) / 2,
+                stableBounds.bottom,
+            )
+        return snapBounds
+    }
 }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/tiling/DesktopTilingWindowDecoration.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/tiling/DesktopTilingWindowDecoration.kt
index 666d4bd..9833325 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/tiling/DesktopTilingWindowDecoration.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/tiling/DesktopTilingWindowDecoration.kt
@@ -36,7 +36,6 @@
 import com.android.internal.annotations.VisibleForTesting
 import com.android.launcher3.icons.BaseIconFactory
 import com.android.window.flags.Flags
-import com.android.wm.shell.shared.FocusTransitionListener
 import com.android.wm.shell.R
 import com.android.wm.shell.RootTaskDisplayAreaOrganizer
 import com.android.wm.shell.ShellTaskOrganizer
@@ -50,6 +49,7 @@
 import com.android.wm.shell.desktopmode.DesktopUserRepositories
 import com.android.wm.shell.desktopmode.ReturnToDragStartAnimator
 import com.android.wm.shell.desktopmode.ToggleResizeDesktopTaskTransitionHandler
+import com.android.wm.shell.shared.FocusTransitionListener
 import com.android.wm.shell.shared.annotations.ShellBackgroundThread
 import com.android.wm.shell.shared.annotations.ShellMainThread
 import com.android.wm.shell.transition.FocusTransitionObserver
@@ -112,7 +112,7 @@
         position: SnapPosition,
         currentBounds: Rect,
     ): Boolean {
-        val destinationBounds = getSnapBounds(taskInfo, position)
+        val destinationBounds = getSnapBounds(position)
         val resizeMetadata =
             AppResizingHelper(
                 taskInfo,
@@ -502,9 +502,11 @@
     }
 
     // Overriding FocusTransitionListener
-    override fun onFocusedTaskChanged(taskId: Int,
-            isFocusedOnDisplay: Boolean,
-            isFocusedGlobally: Boolean) {
+    override fun onFocusedTaskChanged(
+        taskId: Int,
+        isFocusedOnDisplay: Boolean,
+        isFocusedGlobally: Boolean,
+    ) {
         if (!Flags.enableDisplayFocusInShellTransitions()) return
         moveTiledPairToFront(taskId, isFocusedOnDisplay)
     }
@@ -512,7 +514,7 @@
     // Only called if [taskInfo] relates to a focused task
     private fun isTilingRefocused(taskId: Int): Boolean {
         return taskId == leftTaskResizingHelper?.taskInfo?.taskId ||
-                taskId == rightTaskResizingHelper?.taskInfo?.taskId
+            taskId == rightTaskResizingHelper?.taskInfo?.taskId
     }
 
     private fun buildTiledTasksMoveToFront(leftOnTop: Boolean): WindowContainerTransaction {
@@ -623,26 +625,24 @@
         val t = transactionSupplier.get()
         if (!Flags.enableDisplayFocusInShellTransitions()) isTilingFocused = true
         if (taskId == leftTaskResizingHelper?.taskInfo?.taskId) {
-          desktopTilingDividerWindowManager?.onRelativeLeashChanged(
-              leftTiledTask.getLeash(),
-              t,
-          )
+            desktopTilingDividerWindowManager?.onRelativeLeashChanged(leftTiledTask.getLeash(), t)
         }
         if (taskId == rightTaskResizingHelper?.taskInfo?.taskId) {
-          desktopTilingDividerWindowManager?.onRelativeLeashChanged(
-              rightTiledTask.getLeash(),
-              t,
-          )
+            desktopTilingDividerWindowManager?.onRelativeLeashChanged(rightTiledTask.getLeash(), t)
         }
-        transitions.startTransition(
-            TRANSIT_TO_FRONT,
-            buildTiledTasksMoveToFront(isLeftOnTop),
-            null,
-        )
+        transitions.startTransition(TRANSIT_TO_FRONT, buildTiledTasksMoveToFront(isLeftOnTop), null)
         t.apply()
         return true
     }
 
+    fun getRightSnapBoundsIfTiled(): Rect {
+        return getSnapBounds(SnapPosition.RIGHT)
+    }
+
+    fun getLeftSnapBoundsIfTiled(): Rect {
+        return getSnapBounds(SnapPosition.LEFT)
+    }
+
     private fun allTiledTasksVisible(): Boolean {
         val leftTiledTask = leftTaskResizingHelper ?: return false
         val rightTiledTask = rightTaskResizingHelper ?: return false
@@ -674,8 +674,8 @@
             )
     }
 
-    private fun getSnapBounds(taskInfo: RunningTaskInfo, position: SnapPosition): Rect {
-        val displayLayout = displayController.getDisplayLayout(taskInfo.displayId) ?: return Rect()
+    private fun getSnapBounds(position: SnapPosition): Rect {
+        val displayLayout = displayController.getDisplayLayout(displayId) ?: return Rect()
 
         val stableBounds = Rect()
         displayLayout.getStableBounds(stableBounds)
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/tiling/SnapEventHandler.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/tiling/SnapEventHandler.kt
index 52e24d6..b9d6741 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/tiling/SnapEventHandler.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/tiling/SnapEventHandler.kt
@@ -40,4 +40,16 @@
 
     /** If a task is tiled, delegate moving to front to tiling infrastructure. */
     fun moveTaskToFrontIfTiled(taskInfo: RunningTaskInfo): Boolean
+
+    /**
+     * Returns the bounds of a task tiled on the left on the specified display, defaults to default
+     * snapping bounds if no task is tiled.
+     */
+    fun getLeftSnapBoundsIfTiled(displayId: Int): Rect
+
+    /**
+     * Returns the bounds of a task tiled on the right on the specified display, defaults to default
+     * snapping bounds if no task is tiled.
+     */
+    fun getRightSnapBoundsIfTiled(displayId: Int): Rect
 }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/viewholder/AppHeaderViewHolder.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/viewholder/AppHeaderViewHolder.kt
index 90c865e..870c894 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/viewholder/AppHeaderViewHolder.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/viewholder/AppHeaderViewHolder.kt
@@ -37,10 +37,13 @@
 import android.widget.ImageButton
 import android.widget.ImageView
 import android.widget.TextView
+import android.window.DesktopModeFlags
 import androidx.compose.material3.dynamicDarkColorScheme
 import androidx.compose.material3.dynamicLightColorScheme
 import androidx.compose.ui.graphics.toArgb
 import androidx.core.content.withStyledAttributes
+import androidx.core.view.ViewCompat
+import androidx.core.view.accessibility.AccessibilityNodeInfoCompat.AccessibilityActionCompat
 import androidx.core.view.isGone
 import androidx.core.view.isVisible
 import com.android.internal.R.color.materialColorOnSecondaryContainer
@@ -50,9 +53,6 @@
 import com.android.internal.R.color.materialColorSurfaceContainerLow
 import com.android.internal.R.color.materialColorSurfaceDim
 import com.android.wm.shell.R
-import android.window.DesktopModeFlags
-import androidx.core.view.ViewCompat
-import androidx.core.view.accessibility.AccessibilityNodeInfoCompat.AccessibilityActionCompat
 import com.android.wm.shell.windowdecor.MaximizeButtonView
 import com.android.wm.shell.windowdecor.common.DecorThemeUtil
 import com.android.wm.shell.windowdecor.common.OPACITY_100
@@ -145,6 +145,15 @@
     val appNameTextWidth: Int
         get() = appNameTextView.width
 
+    private val a11yAnnounceTextMaximize: String =
+        context.getString(R.string.app_header_talkback_action_maximize_button_text)
+    private val a11yAnnounceTextRestore: String =
+        context.getString(R.string.app_header_talkback_action_restore_button_text)
+
+    private lateinit var sizeToggleDirection: SizeToggleDirection
+    private lateinit var a11yTextMaximize: String
+    private lateinit var a11yTextRestore: String
+
     init {
         captionView.setOnTouchListener(onCaptionTouchListener)
         captionHandle.setOnTouchListener(onCaptionTouchListener)
@@ -163,15 +172,15 @@
 
         val a11yActionSnapLeft = AccessibilityAction(
             R.id.action_snap_left,
-            context.resources.getString(R.string.desktop_mode_a11y_action_snap_left)
+            context.getString(R.string.desktop_mode_a11y_action_snap_left)
         )
         val a11yActionSnapRight = AccessibilityAction(
             R.id.action_snap_right,
-            context.resources.getString(R.string.desktop_mode_a11y_action_snap_right)
+            context.getString(R.string.desktop_mode_a11y_action_snap_right)
         )
         val a11yActionMaximizeRestore = AccessibilityAction(
             R.id.action_maximize_restore,
-            context.resources.getString(R.string.desktop_mode_a11y_action_maximize_restore)
+            context.getString(R.string.desktop_mode_a11y_action_maximize_restore)
         )
 
         captionHandle.accessibilityDelegate = object : View.AccessibilityDelegate() {
@@ -236,19 +245,19 @@
             null
         )
 
-        // Update a11y announcement to say "double tap to maximize or restore window size"
-        ViewCompat.replaceAccessibilityAction(
-            maximizeWindowButton,
-            AccessibilityActionCompat.ACTION_CLICK,
-            context.getString(R.string.maximize_button_talkback_action_maximize_restore_text),
-            null
-        )
-
-        // Update a11y announcement out to say "double tap to minimize app window"
+        // Update a11y announcement to say "double tap to minimize app window"
         ViewCompat.replaceAccessibilityAction(
             minimizeWindowButton,
             AccessibilityActionCompat.ACTION_CLICK,
-            context.getString(R.string.minimize_button_talkback_action_maximize_restore_text),
+            context.getString(R.string.app_header_talkback_action_minimize_button_text),
+            null
+        )
+
+        // Update a11y announcement to say "double tap to close app window"
+        ViewCompat.replaceAccessibilityAction(
+            closeWindowButton,
+            AccessibilityActionCompat.ACTION_CLICK,
+            context.getString(R.string.app_header_talkback_action_close_button_text),
             null
         )
     }
@@ -268,6 +277,26 @@
         appNameTextView.text = name
         openMenuButton.contentDescription =
             context.getString(R.string.desktop_mode_app_header_chip_text, name)
+
+        closeWindowButton.contentDescription = context.getString(R.string.close_button_text, name)
+        minimizeWindowButton.contentDescription =
+            context.getString(R.string.minimize_button_text, name)
+
+        a11yTextMaximize = context.getString(R.string.maximize_button_text, name)
+        a11yTextRestore = context.getString(R.string.restore_button_text, name)
+
+        updateMaximizeButtonContentDescription()
+    }
+
+    private fun updateMaximizeButtonContentDescription() {
+        if (this::a11yTextRestore.isInitialized &&
+            this::a11yTextMaximize.isInitialized &&
+            this::sizeToggleDirection.isInitialized) {
+            maximizeWindowButton.contentDescription = when (sizeToggleDirection) {
+                SizeToggleDirection.MAXIMIZE -> a11yTextMaximize
+                SizeToggleDirection.RESTORE -> a11yTextRestore
+            }
+        }
     }
 
     /** Sets the app's icon in the header. */
@@ -388,7 +417,34 @@
                     drawableInsets = maximizeDrawableInsets
                 )
             )
-            setIcon(getMaximizeButtonIcon(isTaskMaximized, inFullImmersiveState))
+            val icon = getMaximizeButtonIcon(isTaskMaximized, inFullImmersiveState)
+            setIcon(icon)
+
+            when (icon) {
+                R.drawable.decor_desktop_mode_immersive_or_maximize_exit_button_dark -> {
+                    sizeToggleDirection = SizeToggleDirection.RESTORE
+
+                    // Update a11y announcement to say "double tap to maximize app window size"
+                    ViewCompat.replaceAccessibilityAction(
+                        maximizeWindowButton,
+                        AccessibilityActionCompat.ACTION_CLICK,
+                        a11yAnnounceTextRestore,
+                        null
+                    )
+                }
+                R.drawable.decor_desktop_mode_maximize_button_dark -> {
+                    sizeToggleDirection = SizeToggleDirection.MAXIMIZE
+
+                    // Update a11y announcement to say "double tap to restore app window size"
+                    ViewCompat.replaceAccessibilityAction(
+                        maximizeWindowButton,
+                        AccessibilityActionCompat.ACTION_CLICK,
+                        a11yAnnounceTextMaximize,
+                        null
+                    )
+                }
+            }
+            updateMaximizeButtonContentDescription()
         }
         // Close button.
         closeWindowButton.apply {
@@ -625,6 +681,10 @@
         )
     }
 
+    private enum class SizeToggleDirection {
+        MAXIMIZE, RESTORE
+    }
+
     private data class DrawableInsets(val l: Int, val t: Int, val r: Int, val b: Int) {
         constructor(vertical: Int = 0, horizontal: Int = 0) :
                 this(horizontal, vertical, horizontal, vertical)
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopModeVisualIndicatorTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopModeVisualIndicatorTest.kt
index 20d50aa..dcc9e24 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopModeVisualIndicatorTest.kt
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopModeVisualIndicatorTest.kt
@@ -36,6 +36,7 @@
 import com.android.wm.shell.common.DisplayLayout
 import com.android.wm.shell.common.SyncTransactionQueue
 import com.android.wm.shell.shared.bubbles.BubbleDropTargetBoundsProvider
+import com.android.wm.shell.windowdecor.tiling.SnapEventHandler
 import com.google.common.truth.Truth.assertThat
 import org.junit.Before
 import org.junit.Rule
@@ -67,6 +68,7 @@
     @Mock private lateinit var taskDisplayAreaOrganizer: RootTaskDisplayAreaOrganizer
     @Mock private lateinit var displayLayout: DisplayLayout
     @Mock private lateinit var bubbleBoundsProvider: BubbleDropTargetBoundsProvider
+    @Mock private lateinit var snapEventHandler: SnapEventHandler
 
     private lateinit var visualIndicator: DesktopModeVisualIndicator
     private val desktopExecutor = TestShellExecutor()
@@ -356,6 +358,7 @@
                 taskDisplayAreaOrganizer,
                 dragStartState,
                 bubbleBoundsProvider,
+                snapEventHandler,
             )
     }
 
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/VisualIndicatorViewContainerTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/VisualIndicatorViewContainerTest.kt
index 79b0f1c..4c8cb38 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/VisualIndicatorViewContainerTest.kt
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/VisualIndicatorViewContainerTest.kt
@@ -38,6 +38,7 @@
 import com.android.wm.shell.common.SyncTransactionQueue
 import com.android.wm.shell.shared.bubbles.BubbleDropTargetBoundsProvider
 import com.android.wm.shell.windowdecor.WindowDecoration.SurfaceControlViewHostFactory
+import com.android.wm.shell.windowdecor.tiling.SnapEventHandler
 import com.google.common.truth.Truth.assertThat
 import kotlin.test.Test
 import org.junit.Before
@@ -71,6 +72,7 @@
     @Mock private lateinit var mockSurfaceControlViewHostFactory: SurfaceControlViewHostFactory
     @Mock private lateinit var mockBackground: LayerDrawable
     @Mock private lateinit var bubbleDropTargetBoundsProvider: BubbleDropTargetBoundsProvider
+    @Mock private lateinit var snapEventHandler: SnapEventHandler
     private val taskInfo: RunningTaskInfo = createTaskInfo()
     private val mainExecutor = TestShellExecutor()
     private val desktopExecutor = TestShellExecutor()
@@ -81,6 +83,8 @@
         whenever(displayLayout.getStableBounds(any())).thenAnswer { i ->
             (i.arguments.first() as Rect).set(DISPLAY_BOUNDS)
         }
+        whenever(snapEventHandler.getRightSnapBoundsIfTiled(any())).thenReturn(Rect(1, 2, 3, 4))
+        whenever(snapEventHandler.getLeftSnapBoundsIfTiled(any())).thenReturn(Rect(5, 6, 7, 8))
         whenever(mockSurfaceControlViewHostFactory.create(any(), any(), any()))
             .thenReturn(mock(SurfaceControlViewHost::class.java))
     }
@@ -117,7 +121,7 @@
             DesktopModeVisualIndicator.IndicatorType.TO_FULLSCREEN_INDICATOR,
         )
         desktopExecutor.flushAll()
-        verify(spyViewContainer).fadeInIndicator(any(), any())
+        verify(spyViewContainer).fadeInIndicator(any(), any(), any(), any())
     }
 
     @Test
@@ -135,6 +139,8 @@
                 any(),
                 eq(DesktopModeVisualIndicator.IndicatorType.TO_FULLSCREEN_INDICATOR),
                 anyOrNull(),
+                eq(taskInfo.displayId),
+                eq(snapEventHandler),
             )
     }
 
@@ -167,6 +173,8 @@
                     DesktopModeVisualIndicator.IndicatorType.TO_FULLSCREEN_INDICATOR,
                     displayLayout,
                     bubbleDropTargetBoundsProvider,
+                    taskInfo.displayId,
+                    snapEventHandler,
                 )
             }
         assertThat(animator?.indicatorStartBounds).isEqualTo(Rect(15, 15, 985, 985))
@@ -174,6 +182,46 @@
     }
 
     @Test
+    fun testFadeInBoundsCalculationForLeftSnap() {
+        val spyIndicator = setupSpyViewContainer()
+        val animator =
+            spyIndicator.indicatorView?.let {
+                VisualIndicatorViewContainer.VisualIndicatorAnimator.fadeBoundsIn(
+                    it,
+                    DesktopModeVisualIndicator.IndicatorType.TO_SPLIT_LEFT_INDICATOR,
+                    displayLayout,
+                    bubbleDropTargetBoundsProvider,
+                    taskInfo.displayId,
+                    snapEventHandler,
+                )
+            }
+
+        // Right bound is the same as whatever right bound snapEventHandler returned minus padding,
+        // in this case, the right bound for the left app is 7.
+        assertThat(animator?.indicatorEndBounds).isEqualTo(Rect(0, 0, 7, 1000))
+    }
+
+    @Test
+    fun testFadeInBoundsCalculationForRightSnap() {
+        val spyIndicator = setupSpyViewContainer()
+        val animator =
+            spyIndicator.indicatorView?.let {
+                VisualIndicatorViewContainer.VisualIndicatorAnimator.fadeBoundsIn(
+                    it,
+                    DesktopModeVisualIndicator.IndicatorType.TO_SPLIT_RIGHT_INDICATOR,
+                    displayLayout,
+                    bubbleDropTargetBoundsProvider,
+                    taskInfo.displayId,
+                    snapEventHandler,
+                )
+            }
+
+        // Left bound is the same as whatever left bound snapEventHandler returned plus padding
+        // in this case, the left bound of the right app is 1.
+        assertThat(animator?.indicatorEndBounds).isEqualTo(Rect(1, 0, 1000, 1000))
+    }
+
+    @Test
     fun testFadeOutBoundsCalculation() {
         val spyIndicator = setupSpyViewContainer()
         val animator =
@@ -183,6 +231,8 @@
                     DesktopModeVisualIndicator.IndicatorType.TO_FULLSCREEN_INDICATOR,
                     displayLayout,
                     bubbleDropTargetBoundsProvider,
+                    taskInfo.displayId,
+                    snapEventHandler,
                 )
             }
         assertThat(animator?.indicatorStartBounds).isEqualTo(Rect(0, 0, 1000, 1000))
@@ -199,6 +249,8 @@
                 DesktopModeVisualIndicator.IndicatorType.TO_FULLSCREEN_INDICATOR,
                 DesktopModeVisualIndicator.IndicatorType.TO_SPLIT_LEFT_INDICATOR,
                 bubbleDropTargetBoundsProvider,
+                taskInfo.displayId,
+                snapEventHandler,
             )
         // Test desktop to split-right bounds.
         animator =
@@ -208,6 +260,8 @@
                 DesktopModeVisualIndicator.IndicatorType.TO_DESKTOP_INDICATOR,
                 DesktopModeVisualIndicator.IndicatorType.TO_SPLIT_RIGHT_INDICATOR,
                 bubbleDropTargetBoundsProvider,
+                taskInfo.displayId,
+                snapEventHandler,
             )
     }
 
@@ -220,6 +274,7 @@
                 syncQueue,
                 mockSurfaceControlViewHostFactory,
                 bubbleDropTargetBoundsProvider,
+                snapEventHandler,
             )
         viewContainer.createView(
             context,
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/tiling/DesktopTilingDecorViewModelTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/tiling/DesktopTilingDecorViewModelTest.kt
index 2cabb9a..646ec21 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/tiling/DesktopTilingDecorViewModelTest.kt
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/tiling/DesktopTilingDecorViewModelTest.kt
@@ -16,6 +16,7 @@
 package com.android.wm.shell.windowdecor.tiling
 
 import android.content.Context
+import android.content.res.Resources
 import android.graphics.Rect
 import android.testing.AndroidTestingRunner
 import androidx.test.filters.SmallTest
@@ -23,12 +24,13 @@
 import com.android.wm.shell.ShellTaskOrganizer
 import com.android.wm.shell.ShellTestCase
 import com.android.wm.shell.common.DisplayController
+import com.android.wm.shell.common.DisplayLayout
 import com.android.wm.shell.common.ShellExecutor
 import com.android.wm.shell.common.SyncTransactionQueue
 import com.android.wm.shell.desktopmode.DesktopModeEventLogger
-import com.android.wm.shell.desktopmode.DesktopUserRepositories
 import com.android.wm.shell.desktopmode.DesktopTasksController
 import com.android.wm.shell.desktopmode.DesktopTestHelpers.createFreeformTask
+import com.android.wm.shell.desktopmode.DesktopUserRepositories
 import com.android.wm.shell.desktopmode.ReturnToDragStartAnimator
 import com.android.wm.shell.desktopmode.ToggleResizeDesktopTaskTransitionHandler
 import com.android.wm.shell.transition.FocusTransitionObserver
@@ -52,6 +54,7 @@
 @RunWith(AndroidTestingRunner::class)
 class DesktopTilingDecorViewModelTest : ShellTestCase() {
     private val contextMock: Context = mock()
+    private val resourcesMock: Resources = mock()
     private val mainDispatcher: MainCoroutineDispatcher = mock()
     private val bgScope: CoroutineScope = mock()
     private val displayControllerMock: DisplayController = mock()
@@ -70,6 +73,7 @@
     private val desktopTilingDecoration: DesktopTilingWindowDecoration = mock()
     private val taskResourceLoader: WindowDecorTaskResourceLoader = mock()
     private val focusTransitionObserver: FocusTransitionObserver = mock()
+    private val displayLayout: DisplayLayout = mock()
     private val mainExecutor: ShellExecutor = mock()
     private lateinit var desktopTilingDecorViewModel: DesktopTilingDecorViewModel
 
@@ -91,9 +95,16 @@
                 desktopModeEventLogger,
                 taskResourceLoader,
                 focusTransitionObserver,
-                mainExecutor
+                mainExecutor,
             )
         whenever(contextMock.createContextAsUser(any(), any())).thenReturn(contextMock)
+        whenever(displayControllerMock.getDisplayLayout(any())).thenReturn(displayLayout)
+        whenever(displayLayout.getStableBounds(any())).thenAnswer { i ->
+            (i.arguments.first() as Rect).set(STABLE_BOUNDS)
+        }
+        whenever(contextMock.createContextAsUser(any(), any())).thenReturn(context)
+        whenever(contextMock.resources).thenReturn(resourcesMock)
+        whenever(resourcesMock.getDimensionPixelSize(any())).thenReturn(10)
     }
 
     @Test
@@ -202,7 +213,21 @@
         verify(desktopTilingDecoration, times(1)).resetTilingSession()
     }
 
+    @Test
+    fun getTiledAppBounds_NoTilingTransitionHandlerObject() {
+        // Right bound of the left app here represents default 8 / 2 - 2 ( {Right bound} / 2 -
+        // {divider pixel size})
+        assertThat(desktopTilingDecorViewModel.getLeftSnapBoundsIfTiled(1))
+            .isEqualTo(Rect(6, 7, 2, 9))
+
+        // Left bound of the right app here represents default 8 / 2 + 6 + 2 ( {Left bound} +
+        // {width}/ 2 + {divider pixel size})
+        assertThat(desktopTilingDecorViewModel.getRightSnapBoundsIfTiled(1))
+            .isEqualTo(Rect(12, 7, 8, 9))
+    }
+
     companion object {
         private val BOUNDS = Rect(1, 2, 3, 4)
+        private val STABLE_BOUNDS = Rect(6, 7, 8, 9)
     }
 }
diff --git a/packages/SystemUI/compose/core/src/com/android/compose/theme/AndroidColorScheme.kt b/packages/SystemUI/compose/core/src/com/android/compose/theme/AndroidColorScheme.kt
index 1256641..ed144bd 100644
--- a/packages/SystemUI/compose/core/src/com/android/compose/theme/AndroidColorScheme.kt
+++ b/packages/SystemUI/compose/core/src/com/android/compose/theme/AndroidColorScheme.kt
@@ -41,6 +41,7 @@
  */
 @Immutable
 class AndroidColorScheme(
+    // fixed tokens
     val primaryFixed: Color,
     val primaryFixedDim: Color,
     val onPrimaryFixed: Color,
@@ -53,6 +54,30 @@
     val tertiaryFixedDim: Color,
     val onTertiaryFixed: Color,
     val onTertiaryFixedVariant: Color,
+
+    // custom tokens
+    val brandA: Color,
+    val brandB: Color,
+    val brandC: Color,
+    val brandD: Color,
+    val clockHour: Color,
+    val clockMinute: Color,
+    val clockSecond: Color,
+    val onShadeActive: Color,
+    val onShadeActiveVariant: Color,
+    val onShadeInactive: Color,
+    val onShadeInactiveVariant: Color,
+    val onThemeApp: Color,
+    val overviewBackground: Color,
+    val shadeActive: Color,
+    val shadeDisabled: Color,
+    val shadeInactive: Color,
+    val themeApp: Color,
+    val themeAppRing: Color,
+    val themeNotif: Color,
+    val underSurface: Color,
+    val weatherTemp: Color,
+    val widgetBackground: Color,
 ) {
     companion object {
         internal fun color(context: Context, @ColorRes id: Int): Color {
@@ -61,6 +86,7 @@
 
         operator fun invoke(context: Context): AndroidColorScheme {
             return AndroidColorScheme(
+                // Fixed tokens.
                 primaryFixed = color(context, R.color.system_primary_fixed),
                 primaryFixedDim = color(context, R.color.system_primary_fixed_dim),
                 onPrimaryFixed = color(context, R.color.system_on_primary_fixed),
@@ -73,6 +99,30 @@
                 tertiaryFixedDim = color(context, R.color.system_tertiary_fixed_dim),
                 onTertiaryFixed = color(context, R.color.system_on_tertiary_fixed),
                 onTertiaryFixedVariant = color(context, R.color.system_on_tertiary_fixed_variant),
+
+                // Custom tokens.
+                brandA = color(context, R.color.customColorBrandA),
+                brandB = color(context, R.color.customColorBrandB),
+                brandC = color(context, R.color.customColorBrandC),
+                brandD = color(context, R.color.customColorBrandD),
+                clockHour = color(context, R.color.customColorClockHour),
+                clockMinute = color(context, R.color.customColorClockMinute),
+                clockSecond = color(context, R.color.customColorClockSecond),
+                onShadeActive = color(context, R.color.customColorOnShadeActive),
+                onShadeActiveVariant = color(context, R.color.customColorOnShadeActiveVariant),
+                onShadeInactive = color(context, R.color.customColorOnShadeInactive),
+                onShadeInactiveVariant = color(context, R.color.customColorOnShadeInactiveVariant),
+                onThemeApp = color(context, R.color.customColorOnThemeApp),
+                overviewBackground = color(context, R.color.customColorOverviewBackground),
+                shadeActive = color(context, R.color.customColorShadeActive),
+                shadeDisabled = color(context, R.color.customColorShadeDisabled),
+                shadeInactive = color(context, R.color.customColorShadeInactive),
+                themeApp = color(context, R.color.customColorThemeApp),
+                themeAppRing = color(context, R.color.customColorThemeAppRing),
+                themeNotif = color(context, R.color.customColorThemeNotif),
+                underSurface = color(context, R.color.customColorUnderSurface),
+                weatherTemp = color(context, R.color.customColorWeatherTemp),
+                widgetBackground = color(context, R.color.customColorWidgetBackground),
             )
         }
     }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/BundleEntryTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/BundleEntryTest.kt
index 83e26c4..5d8b68e 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/BundleEntryTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/BundleEntryTest.kt
@@ -65,8 +65,8 @@
 
     @Test
     @EnableFlags(NotificationBundleUi.FLAG_NAME)
-    fun getGroupRoot_adapter() {
-        assertThat(underTest.entryAdapter.groupRoot).isEqualTo(underTest.entryAdapter)
+    fun isGroupRoot_adapter() {
+        assertThat(underTest.entryAdapter.isGroupRoot).isTrue()
     }
 
     @Test
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/NotificationEntryTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/NotificationEntryTest.java
index 1f5c672..34dff24 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/NotificationEntryTest.java
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/NotificationEntryTest.java
@@ -542,7 +542,7 @@
 
     @Test
     @EnableFlags(NotificationBundleUi.FLAG_NAME)
-    public void getGroupRoot_adapter_groupSummary() {
+    public void isGroupRoot_adapter_groupSummary() {
         ExpandableNotificationRow row = mock(ExpandableNotificationRow.class);
         Notification notification = new Notification.Builder(mContext, "")
                 .setSmallIcon(R.drawable.ic_person)
@@ -562,12 +562,12 @@
                 .build();
         entry.setRow(row);
 
-        assertThat(entry.getEntryAdapter().getGroupRoot()).isNull();
+        assertThat(entry.getEntryAdapter().isGroupRoot()).isFalse();
     }
 
     @Test
     @EnableFlags(NotificationBundleUi.FLAG_NAME)
-    public void getGroupRoot_adapter_groupChild() {
+    public void isGroupRoot_adapter_groupChild() {
         Notification notification = new Notification.Builder(mContext, "")
                 .setSmallIcon(R.drawable.ic_person)
                 .setGroupSummary(true)
@@ -591,7 +591,7 @@
                 .setParent(groupEntry.build())
                 .build();
 
-        assertThat(entry.getEntryAdapter().getGroupRoot()).isEqualTo(parent.getEntryAdapter());
+        assertThat(entry.getEntryAdapter().isGroupRoot()).isFalse();
     }
 
     @Test
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/coordinator/VisualStabilityCoordinatorTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/coordinator/VisualStabilityCoordinatorTest.java
index 2aa1efa..2aa405a 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/coordinator/VisualStabilityCoordinatorTest.java
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/coordinator/VisualStabilityCoordinatorTest.java
@@ -22,9 +22,10 @@
 
 import static junit.framework.Assert.assertFalse;
 
+import static kotlinx.coroutines.flow.StateFlowKt.MutableStateFlow;
 
-import static org.junit.Assume.assumeFalse;
 import static org.junit.Assert.assertTrue;
+import static org.junit.Assume.assumeFalse;
 import static org.junit.Assume.assumeTrue;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.eq;
@@ -36,8 +37,6 @@
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
-import static kotlinx.coroutines.flow.StateFlowKt.MutableStateFlow;
-
 import android.platform.test.annotations.EnableFlags;
 import android.platform.test.flag.junit.FlagsParameterization;
 import android.testing.TestableLooper;
@@ -74,13 +73,16 @@
 import com.android.systemui.statusbar.notification.collection.listbuilder.pluggable.NotifStabilityManager;
 import com.android.systemui.statusbar.notification.collection.listbuilder.pluggable.Pluggable;
 import com.android.systemui.statusbar.notification.collection.provider.VisualStabilityProvider;
+import com.android.systemui.statusbar.notification.data.repository.HeadsUpRepository;
 import com.android.systemui.statusbar.notification.domain.interactor.SeenNotificationsInteractor;
-import com.android.systemui.statusbar.notification.headsup.HeadsUpManager;
 import com.android.systemui.statusbar.policy.KeyguardStateController;
 import com.android.systemui.util.concurrency.FakeExecutor;
 import com.android.systemui.util.kotlin.JavaAdapter;
 import com.android.systemui.util.time.FakeSystemClock;
 
+import kotlinx.coroutines.flow.MutableStateFlow;
+import kotlinx.coroutines.test.TestScope;
+
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -90,15 +92,12 @@
 import org.mockito.MockitoAnnotations;
 import org.mockito.verification.VerificationMode;
 
-import java.util.List;
-import java.util.Set;
-
-import kotlinx.coroutines.flow.MutableStateFlow;
-import kotlinx.coroutines.test.TestScope;
-
 import platform.test.runner.parameterized.ParameterizedAndroidJunit4;
 import platform.test.runner.parameterized.Parameters;
 
+import java.util.List;
+import java.util.Set;
+
 @SmallTest
 @RunWith(ParameterizedAndroidJunit4.class)
 @TestableLooper.RunWithLooper
@@ -118,7 +117,7 @@
     @Mock private StatusBarStateController mStatusBarStateController;
     @Mock private Pluggable.PluggableListener<NotifStabilityManager> mInvalidateListener;
     @Mock private SeenNotificationsInteractor mSeenNotificationsInteractor;
-    @Mock private HeadsUpManager mHeadsUpManager;
+    @Mock private HeadsUpRepository mHeadsUpRepository;
     @Mock private VisibilityLocationProvider mVisibilityLocationProvider;
     @Mock private VisualStabilityProvider mVisualStabilityProvider;
     @Mock private VisualStabilityCoordinatorLogger mLogger;
@@ -159,7 +158,7 @@
                 mFakeBackgroundExecutor,
                 mFakeMainExecutor,
                 mDumpManager,
-                mHeadsUpManager,
+                mHeadsUpRepository,
                 mShadeAnimationInteractor,
                 mJavaAdapter,
                 mSeenNotificationsInteractor,
@@ -172,6 +171,8 @@
                 mKosmos.getKeyguardTransitionInteractor(),
                 mKeyguardStateController,
                 mLogger);
+
+        when(mHeadsUpRepository.isTrackingHeadsUp()).thenReturn(MutableStateFlow(false));
         mCoordinator.attach(mNotifPipeline);
         mTestScope.getTestScheduler().runCurrent();
 
@@ -194,7 +195,7 @@
                 .setSummary(mEntry)
                 .build();
 
-        when(mHeadsUpManager.isHeadsUpEntry(mEntry.getKey())).thenReturn(false);
+        when(mHeadsUpRepository.isHeadsUpEntry(mEntry.getKey())).thenReturn(false);
 
         // Whenever we invalidate, the pipeline runs again, so we invalidate the state
         doAnswer(i -> {
@@ -461,7 +462,7 @@
         setSleepy(false);
 
         // WHEN a notification is alerting and visible
-        when(mHeadsUpManager.isHeadsUpEntry(mEntry.getKey())).thenReturn(true);
+        when(mHeadsUpRepository.isHeadsUpEntry(mEntry.getKey())).thenReturn(true);
         when(mVisibilityLocationProvider.isInVisibleLocation(any(NotificationEntry.class)))
                 .thenReturn(true);
 
@@ -477,7 +478,7 @@
         setSleepy(false);
 
         // WHEN a notification is alerting but not visible
-        when(mHeadsUpManager.isHeadsUpEntry(mEntry.getKey())).thenReturn(true);
+        when(mHeadsUpRepository.isHeadsUpEntry(mEntry.getKey())).thenReturn(true);
         when(mVisibilityLocationProvider.isInVisibleLocation(any(NotificationEntry.class)))
                 .thenReturn(false);
 
@@ -701,7 +702,7 @@
         assertFalse(mNotifStabilityManager.isSectionChangeAllowed(mEntry));
 
         // GIVEN mEntry is a HUN
-        when(mHeadsUpManager.isHeadsUpEntry(mEntry.getKey())).thenReturn(true);
+        when(mHeadsUpRepository.isHeadsUpEntry(mEntry.getKey())).thenReturn(true);
 
         // THEN group + section changes are allowed
         assertTrue(mNotifStabilityManager.isGroupChangeAllowed(mEntry));
@@ -768,7 +769,7 @@
         //  GIVEN - there is a group heads-up.
         String headsUpGroupKey = "heads_up_group_key";
         mCoordinator.setHeadsUpGroupKeys(Set.of(headsUpGroupKey));
-        when(mHeadsUpManager.isHeadsUpEntry(headsUpGroupKey)).thenReturn(true);
+        when(mHeadsUpRepository.isHeadsUpEntry(headsUpGroupKey)).thenReturn(true);
 
         // GIVEN - HUN Group Summary
         final NotificationEntry nonHeadsUpGroupSummary = mock(NotificationEntry.class);
@@ -793,7 +794,7 @@
         //  GIVEN - there is a group heads-up.
         final String headsUpGroupKey = "heads_up_group_key";
         mCoordinator.setHeadsUpGroupKeys(Set.of(headsUpGroupKey));
-        when(mHeadsUpManager.isHeadsUpEntry(headsUpGroupKey)).thenReturn(true);
+        when(mHeadsUpRepository.isHeadsUpEntry(headsUpGroupKey)).thenReturn(true);
 
         // GIVEN - HUN Group
         final NotificationEntry headsUpGroupSummary = mock(NotificationEntry.class);
@@ -825,7 +826,7 @@
         //  GIVEN - there is a group heads-up.
         final String headsUpGroupKey = "heads_up_group_key";
         mCoordinator.setHeadsUpGroupKeys(Set.of(headsUpGroupKey));
-        when(mHeadsUpManager.isHeadsUpEntry(headsUpGroupKey)).thenReturn(true);
+        when(mHeadsUpRepository.isHeadsUpEntry(headsUpGroupKey)).thenReturn(true);
 
         // GIVEN - HUN Group
         final NotificationEntry headsUpGroupSummary = mock(NotificationEntry.class);
@@ -858,7 +859,7 @@
         //  GIVEN - there is a group heads-up.
         String headsUpGroupKey = "heads_up_group_key";
         mCoordinator.setHeadsUpGroupKeys(Set.of(headsUpGroupKey));
-        when(mHeadsUpManager.isHeadsUpEntry(headsUpGroupKey)).thenReturn(true);
+        when(mHeadsUpRepository.isHeadsUpEntry(headsUpGroupKey)).thenReturn(true);
 
         // GIVEN - non HUN parent Group Summary
         final NotificationEntry groupSummary = mock(NotificationEntry.class);
@@ -891,7 +892,7 @@
         // GIVEN - there is a group heads-up.
         final String headsUpGroupKey = "heads_up_group_key";
         mCoordinator.setHeadsUpGroupKeys(Set.of(headsUpGroupKey));
-        when(mHeadsUpManager.isHeadsUpEntry(headsUpGroupKey)).thenReturn(true);
+        when(mHeadsUpRepository.isHeadsUpEntry(headsUpGroupKey)).thenReturn(true);
 
         // GIVEN - HUN Group Summary
         final NotificationEntry headsUpGroupSummary = mock(NotificationEntry.class);
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/render/GroupExpansionManagerTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/render/GroupExpansionManagerTest.kt
index 3dd0982..8e6aedca 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/render/GroupExpansionManagerTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/render/GroupExpansionManagerTest.kt
@@ -17,6 +17,7 @@
 package com.android.systemui.statusbar.notification.collection.render
 
 import android.os.Build
+import android.os.UserHandle
 import android.platform.test.annotations.DisableFlags
 import android.platform.test.annotations.EnableFlags
 import android.platform.test.flag.junit.SetFlagsRule
@@ -29,9 +30,12 @@
 import com.android.systemui.statusbar.notification.collection.GroupEntryBuilder
 import com.android.systemui.statusbar.notification.collection.ListEntry
 import com.android.systemui.statusbar.notification.collection.NotifPipeline
+import com.android.systemui.statusbar.notification.collection.NotificationEntry
 import com.android.systemui.statusbar.notification.collection.NotificationEntryBuilder
 import com.android.systemui.statusbar.notification.collection.listbuilder.OnBeforeRenderListListener
 import com.android.systemui.statusbar.notification.collection.render.GroupExpansionManager.OnGroupExpansionChangeListener
+import com.android.systemui.statusbar.notification.row.ExpandableNotificationRow
+import com.android.systemui.statusbar.notification.row.NotificationTestHelper
 import com.android.systemui.statusbar.notification.shared.NotificationBundleUi
 import com.android.systemui.util.mockito.any
 import com.android.systemui.util.mockito.mock
@@ -55,56 +59,58 @@
 
     private lateinit var underTest: GroupExpansionManagerImpl
 
+    private lateinit var testHelper: NotificationTestHelper
     private val dumpManager: DumpManager = mock()
     private val groupMembershipManager: GroupMembershipManager = mock()
 
     private val pipeline: NotifPipeline = mock()
     private lateinit var beforeRenderListListener: OnBeforeRenderListListener
 
-    private val summary1 = notificationSummaryEntry("foo", 1)
-    private val summary2 = notificationSummaryEntry("bar", 1)
-    private val entries =
-        listOf<ListEntry>(
-            GroupEntryBuilder()
-                .setSummary(summary1)
-                .setChildren(
-                    listOf(
-                        notificationEntry("foo", 2),
-                        notificationEntry("foo", 3),
-                        notificationEntry("foo", 4)
-                    )
-                )
-                .build(),
-            GroupEntryBuilder()
-                .setSummary(summary2)
-                .setChildren(
-                    listOf(
-                        notificationEntry("bar", 2),
-                        notificationEntry("bar", 3),
-                        notificationEntry("bar", 4)
-                    )
-                )
-                .build(),
-            notificationEntry("baz", 1)
-        )
+    private lateinit var summary1: NotificationEntry
+    private lateinit var summary2: NotificationEntry
+    private lateinit var entries: List<ListEntry>
 
-    private fun notificationEntry(pkg: String, id: Int) =
-        NotificationEntryBuilder().setPkg(pkg).setId(id).build().apply { row = mock() }
-
-    private fun notificationSummaryEntry(pkg: String, id: Int) =
-        NotificationEntryBuilder().setPkg(pkg).setId(id).setParent(GroupEntry.ROOT_ENTRY).build()
-            .apply { row = mock() }
+    private fun notificationEntry(pkg: String, id: Int, parent: ExpandableNotificationRow?) =
+        NotificationEntryBuilder().setPkg(pkg).setId(id).build().apply {
+            row = testHelper.createRow().apply {
+                setIsChildInGroup(true, parent)
+            }
+        }
 
     @Before
     fun setUp() {
+        testHelper = NotificationTestHelper(mContext, mDependency)
+
+        summary1 = testHelper.createRow().entry
+        summary2 = testHelper.createRow().entry
+        entries =
+            listOf<ListEntry>(
+                GroupEntryBuilder()
+                    .setSummary(summary1)
+                    .setChildren(
+                        listOf(
+                            notificationEntry("foo", 2, summary1.row),
+                            notificationEntry("foo", 3, summary1.row),
+                            notificationEntry("foo", 4, summary1.row)
+                        )
+                    )
+                    .build(),
+                GroupEntryBuilder()
+                    .setSummary(summary2)
+                    .setChildren(
+                        listOf(
+                            notificationEntry("bar", 2, summary2.row),
+                            notificationEntry("bar", 3, summary2.row),
+                            notificationEntry("bar", 4, summary2.row)
+                        )
+                    )
+                    .build(),
+                notificationEntry("baz", 1, null)
+            )
+
         whenever(groupMembershipManager.getGroupSummary(summary1)).thenReturn(summary1)
         whenever(groupMembershipManager.getGroupSummary(summary2)).thenReturn(summary2)
 
-        whenever(groupMembershipManager.getGroupRoot(summary1.entryAdapter))
-            .thenReturn(summary1.entryAdapter)
-        whenever(groupMembershipManager.getGroupRoot(summary2.entryAdapter))
-            .thenReturn(summary2.entryAdapter)
-
         underTest = GroupExpansionManagerImpl(dumpManager, groupMembershipManager)
     }
 
@@ -221,4 +227,15 @@
         verify(listener).onGroupExpansionChange(summary1.row, false)
         verifyNoMoreInteractions(listener)
     }
+
+    @Test
+    @EnableFlags(NotificationBundleUi.FLAG_NAME)
+    fun isGroupExpanded() {
+        underTest.setGroupExpanded(summary1.entryAdapter, true)
+
+        assertThat(underTest.isGroupExpanded(summary1.entryAdapter)).isTrue();
+        assertThat(underTest.isGroupExpanded(
+            (entries[0] as? GroupEntry)?.getChildren()?.get(0)?.entryAdapter))
+            .isTrue();
+    }
 }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/render/GroupMembershipManagerTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/render/GroupMembershipManagerTest.kt
index dcbf44e..2bbf094 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/render/GroupMembershipManagerTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/render/GroupMembershipManagerTest.kt
@@ -170,6 +170,21 @@
 
     @Test
     @EnableFlags(NotificationBundleUi.FLAG_NAME)
+    fun isChildEntryAdapterInGroup_child() {
+        val groupKey = "group"
+        val summary =
+            NotificationEntryBuilder()
+                .setGroup(mContext, groupKey)
+                .setGroupSummary(mContext, true)
+                .build()
+        val entry = NotificationEntryBuilder().setGroup(mContext, groupKey).build()
+        GroupEntryBuilder().setKey(groupKey).setSummary(summary).addChild(entry).build()
+
+        assertThat(underTest.isChildInGroup(entry.entryAdapter)).isTrue()
+    }
+
+    @Test
+    @EnableFlags(NotificationBundleUi.FLAG_NAME)
     fun isGroupRoot_topLevelEntry() {
         val entry = NotificationEntryBuilder().setParent(GroupEntry.ROOT_ENTRY).build()
         assertThat(underTest.isGroupRoot(entry.entryAdapter)).isFalse()
@@ -203,40 +218,4 @@
 
         assertThat(underTest.isGroupRoot(entry.entryAdapter)).isFalse()
     }
-
-    @Test
-    @EnableFlags(NotificationBundleUi.FLAG_NAME)
-    fun getGroupRoot_topLevelEntry() {
-        val entry = NotificationEntryBuilder().setParent(GroupEntry.ROOT_ENTRY).build()
-        assertThat(underTest.getGroupRoot(entry.entryAdapter)).isNull()
-    }
-
-    @Test
-    @EnableFlags(NotificationBundleUi.FLAG_NAME)
-    fun getGroupRoot_summary() {
-        val groupKey = "group"
-        val summary =
-            NotificationEntryBuilder()
-                .setGroup(mContext, groupKey)
-                .setGroupSummary(mContext, true)
-                .build()
-        GroupEntryBuilder().setKey(groupKey).setSummary(summary).build()
-
-        assertThat(underTest.getGroupRoot(summary.entryAdapter)).isEqualTo(summary.entryAdapter)
-    }
-
-    @Test
-    @EnableFlags(NotificationBundleUi.FLAG_NAME)
-    fun getGroupRoot_child() {
-        val groupKey = "group"
-        val summary =
-            NotificationEntryBuilder()
-                .setGroup(mContext, groupKey)
-                .setGroupSummary(mContext, true)
-                .build()
-        val entry = NotificationEntryBuilder().setGroup(mContext, groupKey).build()
-        GroupEntryBuilder().setKey(groupKey).setSummary(summary).addChild(entry).build()
-
-        assertThat(underTest.getGroupRoot(entry.entryAdapter)).isEqualTo(summary.entryAdapter)
-    }
 }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/row/NotificationTestHelper.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/row/NotificationTestHelper.java
index f2131da..92ceb60 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/row/NotificationTestHelper.java
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/row/NotificationTestHelper.java
@@ -101,6 +101,7 @@
 
 import kotlin.coroutines.CoroutineContext;
 
+import kotlinx.coroutines.flow.StateFlowKt;
 import kotlinx.coroutines.test.TestScope;
 
 import org.mockito.ArgumentCaptor;
@@ -166,11 +167,21 @@
         this(context, dependency, testLooper, new FakeFeatureFlagsClassic());
     }
 
+
     public NotificationTestHelper(
             Context context,
             TestableDependency dependency,
             @Nullable TestableLooper testLooper,
             @NonNull FakeFeatureFlagsClassic featureFlags) {
+        this(context, dependency, testLooper, featureFlags, mockHeadsUpManager());
+    }
+
+    public NotificationTestHelper(
+            Context context,
+            TestableDependency dependency,
+            @Nullable TestableLooper testLooper,
+            @NonNull FakeFeatureFlagsClassic featureFlags,
+            @NonNull HeadsUpManager headsUpManager) {
         mContext = context;
         mFeatureFlags = Objects.requireNonNull(featureFlags);
         dependency.injectTestDependency(FeatureFlagsClassic.class, mFeatureFlags);
@@ -182,7 +193,7 @@
         mKeyguardBypassController = mock(KeyguardBypassController.class);
         mGroupMembershipManager = mock(GroupMembershipManager.class);
         mGroupExpansionManager = mock(GroupExpansionManager.class);
-        mHeadsUpManager = mock(HeadsUpManager.class);
+        mHeadsUpManager = headsUpManager;
         mIconManager = new IconManager(
                 mock(CommonNotifCollection.class),
                 mock(LauncherApps.class),
@@ -689,6 +700,12 @@
                 .build();
     }
 
+    private static HeadsUpManager mockHeadsUpManager() {
+        HeadsUpManager mock = mock(HeadsUpManager.class);
+        when(mock.isTrackingHeadsUp()).thenReturn(StateFlowKt.MutableStateFlow(false));
+        return mock;
+    }
+
     private static class MockSmartReplyInflater implements SmartReplyStateInflater {
         @Override
         public InflatedSmartReplyState inflateSmartReplyState(NotificationEntry entry) {
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/stack/MagneticNotificationRowManagerImplTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/stack/MagneticNotificationRowManagerImplTest.kt
index 885e71e..ccc8be7 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/stack/MagneticNotificationRowManagerImplTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/stack/MagneticNotificationRowManagerImplTest.kt
@@ -80,7 +80,7 @@
             // THEN the magnetic and roundable targets are defined and the state is TARGETS_SET
             assertThat(underTest.currentState).isEqualTo(State.TARGETS_SET)
             assertThat(underTest.currentMagneticListeners.isNotEmpty()).isTrue()
-            assertThat(underTest.currentRoundableTargets).isNotNull()
+            assertThat(underTest.isSwipedViewRoundableSet).isTrue()
         }
 
     @Test
@@ -281,6 +281,19 @@
             assertThat(magneticAnimationsCancelled[neighborIndex]).isTrue()
         }
 
+    @Test
+    fun onResetRoundness_swipedRoundableGetsCleared() =
+        kosmos.testScope.runTest {
+            // GIVEN targets are set
+            setTargets()
+
+            // WHEN we reset the roundness
+            underTest.resetRoundness()
+
+            // THEN the swiped roundable gets cleared
+            assertThat(underTest.isSwipedViewRoundableSet).isFalse()
+        }
+
     @After
     fun tearDown() {
         // We reset the manager so that all MagneticRowListener can cancel all animations
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/StatusBarNotificationPresenterTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/StatusBarNotificationPresenterTest.kt
index baea1a1..47967b3 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/StatusBarNotificationPresenterTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/StatusBarNotificationPresenterTest.kt
@@ -54,6 +54,7 @@
 import com.android.systemui.statusbar.notification.interruption.VisualInterruptionRefactor
 import com.android.systemui.statusbar.notification.interruption.VisualInterruptionType
 import com.android.systemui.statusbar.notification.row.ExpandableNotificationRow
+import com.android.systemui.statusbar.notification.shared.NotificationBundleUi
 import com.android.systemui.statusbar.notification.stack.notificationStackScrollLayoutController
 import com.android.systemui.statusbar.notification.visualInterruptionDecisionProvider
 import com.android.systemui.statusbar.notificationLockscreenUserManager
@@ -293,6 +294,7 @@
 
     @Test
     @EnableSceneContainer
+    @DisableFlags(NotificationBundleUi.FLAG_NAME)
     fun testExpandSensitiveNotification_onLockScreen_opensShade() =
         kosmos.runTest {
             // Given we are on the keyguard
@@ -303,11 +305,10 @@
             )
 
             // When the user expands a sensitive Notification
-            val row = createRow()
-            val entry =
-                row.entry.apply { setSensitive(/* sensitive= */ true, /* deviceSensitive= */ true) }
+            val row = createRow(createNotificationEntry())
+            row.entry.apply { setSensitive(/* sensitive= */ true, /* deviceSensitive= */ true) }
 
-            underTest.onExpandClicked(entry, mock(), /* nowExpanded= */ true)
+            underTest.onExpandClicked(row.entry, mock(), /* nowExpanded= */ true)
 
             // Then we open the locked shade
             verify(kosmos.lockscreenShadeTransitionController)
@@ -317,6 +318,32 @@
 
     @Test
     @EnableSceneContainer
+    @EnableFlags(NotificationBundleUi.FLAG_NAME)
+    fun testExpandSensitiveNotification_onLockScreen_opensShade_entryAdapter() =
+        kosmos.runTest {
+            // Given we are on the keyguard
+            kosmos.sysuiStatusBarStateController.state = StatusBarState.KEYGUARD
+            // And the device is locked
+            kosmos.fakeAuthenticationRepository.setAuthenticationMethod(
+                AuthenticationMethodModel.Pin
+            )
+
+            // When the user expands a sensitive Notification
+            val entry = createNotificationEntry()
+            val row = createRow(entry)
+            entry.setSensitive(/* sensitive= */ true, /* deviceSensitive= */ true)
+
+            underTest.onExpandClicked(row, row.entryAdapter, /* nowExpanded= */ true)
+
+            // Then we open the locked shade
+            verify(kosmos.lockscreenShadeTransitionController)
+                // Explicit parameters to avoid issues with Kotlin default arguments in Mockito
+                .goToLockedShade(row, true)
+        }
+
+    @Test
+    @EnableSceneContainer
+    @DisableFlags(NotificationBundleUi.FLAG_NAME)
     fun testExpandSensitiveNotification_onLockedShade_showsBouncer() =
         kosmos.runTest {
             // Given we are on the locked shade
@@ -328,7 +355,7 @@
 
             // When the user expands a sensitive Notification
             val entry =
-                createRow().entry.apply {
+                createRow(createNotificationEntry()).entry.apply {
                     setSensitive(/* sensitive= */ true, /* deviceSensitive= */ true)
                 }
             underTest.onExpandClicked(entry, mock(), /* nowExpanded= */ true)
@@ -337,6 +364,29 @@
             verify(kosmos.activityStarter).dismissKeyguardThenExecute(any(), eq(null), eq(false))
         }
 
+    @Test
+    @EnableSceneContainer
+    @EnableFlags(NotificationBundleUi.FLAG_NAME)
+    fun testExpandSensitiveNotification_onLockedShade_showsBouncer_entryAdapter() =
+        kosmos.runTest {
+            // Given we are on the locked shade
+            kosmos.sysuiStatusBarStateController.state = StatusBarState.SHADE_LOCKED
+            // And the device is locked
+            kosmos.fakeAuthenticationRepository.setAuthenticationMethod(
+                AuthenticationMethodModel.Pin
+            )
+
+            // When the user expands a sensitive Notification
+            val entry = createNotificationEntry()
+            val row = createRow(entry)
+            entry.setSensitive(/* sensitive= */ true, /* deviceSensitive= */ true)
+
+            underTest.onExpandClicked(row, row.entryAdapter, /* nowExpanded= */ true)
+
+            // Then we show the bouncer
+            verify(kosmos.activityStarter).dismissKeyguardThenExecute(any(), eq(null), eq(false))
+        }
+
     private fun createPresenter(): StatusBarNotificationPresenter {
         val initController: InitController = InitController()
         return StatusBarNotificationPresenter(
@@ -398,10 +448,13 @@
         interruptSuppressor = suppressorCaptor.lastValue
     }
 
-    private fun createRow(): ExpandableNotificationRow {
+    private fun createRow(entry: NotificationEntry): ExpandableNotificationRow {
         val row: ExpandableNotificationRow = mock()
-        val entry: NotificationEntry = createNotificationEntry()
-        whenever(row.entry).thenReturn(entry)
+        if (NotificationBundleUi.isEnabled) {
+            whenever(row.entryAdapter).thenReturn(entry.entryAdapter)
+        } else {
+            whenever(row.entry).thenReturn(entry)
+        }
         entry.row = row
         return row
     }
diff --git a/packages/SystemUI/src/com/android/systemui/shade/NotificationPanelViewController.java b/packages/SystemUI/src/com/android/systemui/shade/NotificationPanelViewController.java
index 28d619a..09a04a7 100644
--- a/packages/SystemUI/src/com/android/systemui/shade/NotificationPanelViewController.java
+++ b/packages/SystemUI/src/com/android/systemui/shade/NotificationPanelViewController.java
@@ -1788,7 +1788,8 @@
         // height - which means user is swiping down. Otherwise shade QS will either not show at all
         // with HUN movement or it will blink when touching HUN initially
         boolean qsShouldExpandWithHeadsUp = !mSplitShadeEnabled
-                || (!mHeadsUpManager.isTrackingHeadsUp() || expandedHeight > mHeadsUpStartHeight);
+                || (!mHeadsUpManager.isTrackingHeadsUp().getValue()
+                || expandedHeight > mHeadsUpStartHeight);
         if (goingBetweenClosedShadeAndExpandedQs && qsShouldExpandWithHeadsUp) {
             float qsExpansionFraction;
             if (mSplitShadeEnabled) {
@@ -2048,7 +2049,7 @@
         // motion has the expected speed. We also only want this on non-lockscreen for now.
         if (mSplitShadeEnabled && mBarState == SHADE) {
             boolean transitionFromHeadsUp = (mHeadsUpManager != null
-                    && mHeadsUpManager.isTrackingHeadsUp()) || mExpandingFromHeadsUp;
+                    && mHeadsUpManager.isTrackingHeadsUp().getValue()) || mExpandingFromHeadsUp;
             // heads-up starting height is too close to mSplitShadeFullTransitionDistance and
             // when dragging HUN transition is already 90% complete. It makes shade become
             // immediately visible when starting to drag. We want to set distance so that
diff --git a/packages/SystemUI/src/com/android/systemui/shade/QuickSettingsControllerImpl.java b/packages/SystemUI/src/com/android/systemui/shade/QuickSettingsControllerImpl.java
index d058372..42c63f9 100644
--- a/packages/SystemUI/src/com/android/systemui/shade/QuickSettingsControllerImpl.java
+++ b/packages/SystemUI/src/com/android/systemui/shade/QuickSettingsControllerImpl.java
@@ -2158,6 +2158,8 @@
 
     /** */
     public final class QsFragmentListener implements FragmentHostManager.FragmentListener {
+        private boolean mPreviouslyVisibleMedia = false;
+
         /** */
         @Override
         public void onFragmentViewCreated(String tag, Fragment fragment) {
@@ -2183,7 +2185,12 @@
                     setAnimateNextNotificationBounds(
                             StackStateAnimator.ANIMATION_DURATION_STANDARD, 0);
                     mNotificationStackScrollLayoutController.animateNextTopPaddingChange();
+                    if (QSComposeFragment.isEnabled() && mPreviouslyVisibleMedia && !visible) {
+                        updateHeightsOnShadeLayoutChange();
+                        mPanelViewControllerLazy.get().positionClockAndNotifications();
+                    }
                 }
+                mPreviouslyVisibleMedia = visible;
             });
             mLockscreenShadeTransitionController.setQS(mQs);
             if (QSComposeFragment.isEnabled()) {
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/NotificationGroupingUtil.java b/packages/SystemUI/src/com/android/systemui/statusbar/NotificationGroupingUtil.java
index 6ebe024..4bb12e5 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/NotificationGroupingUtil.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/NotificationGroupingUtil.java
@@ -57,12 +57,6 @@
     private static final VisibilityApplicator VISIBILITY_APPLICATOR = new VisibilityApplicator();
     private static final VisibilityApplicator APP_NAME_APPLICATOR = new AppNameApplicator();
     private static final ResultApplicator LEFT_ICON_APPLICATOR = new LeftIconApplicator();
-    private static final DataExtractor ICON_EXTRACTOR = new DataExtractor() {
-        @Override
-        public Object extractData(ExpandableNotificationRow row) {
-            return row.getEntry().getSbn().getNotification();
-        }
-    };
 
     private final ExpandableNotificationRow mRow;
     private final ArrayList<Processor> mProcessors = new ArrayList<>();
@@ -109,31 +103,26 @@
         // To hide the icons if they are the same and the color is the same
         mProcessors.add(new Processor(mRow,
                 com.android.internal.R.id.icon,
-                ICON_EXTRACTOR,
                 iconVisibilityComparator,
                 VISIBILITY_APPLICATOR));
         // To grey out the icons when they are not the same, or they have the same color
         mProcessors.add(new Processor(mRow,
                 com.android.internal.R.id.status_bar_latest_event_content,
-                ICON_EXTRACTOR,
                 greyComparator,
                 greyApplicator));
         // To show the large icon on the left side instead if all the small icons are the same
         mProcessors.add(new Processor(mRow,
                 com.android.internal.R.id.status_bar_latest_event_content,
-                ICON_EXTRACTOR,
                 iconVisibilityComparator,
                 LEFT_ICON_APPLICATOR));
         // To only show the work profile icon in the group header
         mProcessors.add(new Processor(mRow,
                 com.android.internal.R.id.profile_badge,
-                null /* Extractor */,
                 BADGE_COMPARATOR,
                 VISIBILITY_APPLICATOR));
         // To hide the app name in group children
         mProcessors.add(new Processor(mRow,
                 com.android.internal.R.id.app_name_text,
-                null,
                 APP_NAME_COMPARATOR,
                 APP_NAME_APPLICATOR));
         // To hide the header text if it's the same
@@ -253,23 +242,20 @@
 
     private static class Processor {
         private final int mId;
-        private final DataExtractor mExtractor;
         private final ViewComparator mComparator;
         private final ResultApplicator mApplicator;
         private final ExpandableNotificationRow mParentRow;
         private boolean mApply;
         private View mParentView;
-        private Object mParentData;
 
         public static Processor forTextView(ExpandableNotificationRow row, int id) {
-            return new Processor(row, id, null, TEXT_VIEW_COMPARATOR, VISIBILITY_APPLICATOR);
+            return new Processor(row, id, TEXT_VIEW_COMPARATOR, VISIBILITY_APPLICATOR);
         }
 
-        Processor(ExpandableNotificationRow row, int id, DataExtractor extractor,
+        Processor(ExpandableNotificationRow row, int id,
                 ViewComparator comparator,
                 ResultApplicator applicator) {
             mId = id;
-            mExtractor = extractor;
             mApplicator = applicator;
             mComparator = comparator;
             mParentRow = row;
@@ -279,7 +265,6 @@
             NotificationViewWrapper wrapper = mParentRow.getNotificationViewWrapper();
             View header = wrapper == null ? null : wrapper.getNotificationHeader();
             mParentView = header == null ? null : header.findViewById(mId);
-            mParentData = mExtractor == null ? null : mExtractor.extractData(mParentRow);
             mApply = !mComparator.isEmpty(mParentView);
         }
 
@@ -297,9 +282,7 @@
                 // when for example showing an undo notification
                 return;
             }
-            Object childData = mExtractor == null ? null : mExtractor.extractData(row);
-            mApply = mComparator.compare(mParentView, ownView,
-                    mParentData, childData);
+            mApply = mComparator.compare(mParentView, ownView);
         }
 
         public void apply(ExpandableNotificationRow row) {
@@ -331,11 +314,9 @@
         /**
          * @param parent     the view with the given id in the group header
          * @param child      the view with the given id in the child notification
-         * @param parentData optional data for the parent
-         * @param childData  optional data for the child
          * @return whether to views are the same
          */
-        boolean compare(View parent, View child, Object parentData, Object childData);
+        boolean compare(View parent, View child);
 
         boolean isEmpty(View view);
     }
@@ -346,7 +327,7 @@
 
     private static class BadgeComparator implements ViewComparator {
         @Override
-        public boolean compare(View parent, View child, Object parentData, Object childData) {
+        public boolean compare(View parent, View child) {
             return parent.getVisibility() != View.GONE;
         }
 
@@ -364,7 +345,7 @@
 
     private static class TextViewComparator implements ViewComparator {
         @Override
-        public boolean compare(View parent, View child, Object parentData, Object childData) {
+        public boolean compare(View parent, View child) {
             TextView parentView = (TextView) parent;
             CharSequence parentText = parentView == null ? "" : parentView.getText();
             TextView childView = (TextView) child;
@@ -380,7 +361,7 @@
 
     private abstract static class IconComparator implements ViewComparator {
         @Override
-        public boolean compare(View parent, View child, Object parentData, Object childData) {
+        public boolean compare(View parent, View child) {
             return false;
         }
 
@@ -440,14 +421,14 @@
 
     private static class AppNameComparator extends TextViewComparator {
         @Override
-        public boolean compare(View parent, View child, Object parentData, Object childData) {
+        public boolean compare(View parent, View child) {
             if (isEmpty(child)) {
                 // In headerless notifications the AppName view exists but is usually GONE (and not
                 // populated).  We need to treat this case as equal to the header in order to
                 // deduplicate the view.
                 return true;
             }
-            return super.compare(parent, child, parentData, childData);
+            return super.compare(parent, child);
         }
     }
 
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/NotificationShelf.java b/packages/SystemUI/src/com/android/systemui/statusbar/NotificationShelf.java
index 31fdec6..f88c618 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/NotificationShelf.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/NotificationShelf.java
@@ -50,6 +50,7 @@
 import com.android.systemui.statusbar.notification.row.ActivatableNotificationView;
 import com.android.systemui.statusbar.notification.row.ExpandableNotificationRow;
 import com.android.systemui.statusbar.notification.row.ExpandableView;
+import com.android.systemui.statusbar.notification.shared.NotificationBundleUi;
 import com.android.systemui.statusbar.notification.shared.NotificationMinimalism;
 import com.android.systemui.statusbar.notification.shelf.NotificationShelfBackgroundView;
 import com.android.systemui.statusbar.notification.shelf.NotificationShelfIconContainer;
@@ -672,7 +673,9 @@
             // if the shelf is clipped, lets make sure we also clip the icon
             maxTop = Math.max(maxTop, getTranslationY() + getClipTopAmount());
         }
-        StatusBarIconView icon = row.getEntry().getIcons().getShelfIcon();
+        StatusBarIconView icon = NotificationBundleUi.isEnabled()
+                ? row.getEntryAdapter().getIcons().getShelfIcon()
+                : row.getEntry().getIcons().getShelfIcon();
         float shelfIconPosition = getTranslationY() + icon.getTop() + icon.getTranslationY();
         if (shelfIconPosition < maxTop && !mAmbientState.isFullyHidden()) {
             int top = (int) (maxTop - shelfIconPosition);
@@ -684,7 +687,9 @@
     }
 
     private void updateContinuousClipping(final ExpandableNotificationRow row) {
-        StatusBarIconView icon = row.getEntry().getIcons().getShelfIcon();
+        StatusBarIconView icon = NotificationBundleUi.isEnabled()
+                ? row.getEntryAdapter().getIcons().getShelfIcon()
+                : row.getEntry().getIcons().getShelfIcon();
         boolean needsContinuousClipping = ViewState.isAnimatingY(icon) && !mAmbientState.isDozing();
         boolean isContinuousClipping = icon.getTag(TAG_CONTINUOUS_CLIPPING) != null;
         if (needsContinuousClipping && !isContinuousClipping) {
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/BundleEntry.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/BundleEntry.java
index 35a2828..c79cae7 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/BundleEntry.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/BundleEntry.java
@@ -25,6 +25,7 @@
 import android.content.Context;
 import android.os.Build;
 
+import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
 import androidx.annotation.VisibleForTesting;
 
@@ -34,6 +35,10 @@
 
 import java.util.List;
 
+import kotlinx.coroutines.flow.MutableStateFlow;
+import kotlinx.coroutines.flow.StateFlow;
+import kotlinx.coroutines.flow.StateFlowKt;
+
 /**
  * Class to represent notifications bundled by classification.
  */
@@ -41,6 +46,9 @@
 
     private final BundleEntryAdapter mEntryAdapter;
 
+    // TODO(b/394483200): move NotificationEntry's implementation to PipelineEntry?
+    private final MutableStateFlow<Boolean> mSensitive = StateFlowKt.MutableStateFlow(false);
+
     // TODO (b/389839319): implement the row
     private ExpandableNotificationRow mRow;
 
@@ -97,20 +105,26 @@
             return true;
         }
 
+        @NonNull
         @Override
         public String getKey() {
             return mKey;
         }
 
         @Override
+        @Nullable
         public ExpandableNotificationRow getRow() {
             return mRow;
         }
 
-        @Nullable
         @Override
-        public EntryAdapter getGroupRoot() {
-            return this;
+        public boolean isGroupRoot() {
+            return true;
+        }
+
+        @Override
+        public StateFlow<Boolean> isSensitive() {
+            return BundleEntry.this.mSensitive;
         }
 
         @Override
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/EntryAdapter.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/EntryAdapter.java
index 6431cac..109ebe6 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/EntryAdapter.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/EntryAdapter.java
@@ -24,6 +24,8 @@
 import com.android.systemui.statusbar.notification.icon.IconPack;
 import com.android.systemui.statusbar.notification.row.ExpandableNotificationRow;
 
+import kotlinx.coroutines.flow.StateFlow;
+
 /**
  * Adapter interface for UI to get relevant info.
  */
@@ -51,15 +53,10 @@
     ExpandableNotificationRow getRow();
 
     /**
-     * Gets the EntryAdapter that is the nearest root of the collection of rows the given entry
-     * belongs to. If the given entry is a BundleEntry or an isolated child of a BundleEntry, the
-     * BundleEntry will be returned. If the given notification is a group summary NotificationEntry,
-     * or a child of a group summary, the summary NotificationEntry will be returned, even if that
-     * summary belongs to a BundleEntry. If the entry is a notification that does not belong to any
-     * group or bundle grouping, null will be returned.
+     * Whether this entry is the root of its collapsable 'group' - either a BundleEntry or a
+     * notification group summary
      */
-    @Nullable
-    EntryAdapter getGroupRoot();
+    boolean isGroupRoot();
 
     /**
      * @return whether the row can be removed with the 'Clear All' action
@@ -107,4 +104,9 @@
      * Retrieves the pack of icons associated with this entry
      */
     IconPack getIcons();
+
+    /**
+     * Returns whether the content of this entry is sensitive
+     */
+    StateFlow<Boolean> isSensitive();
 }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/NotificationEntry.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/NotificationEntry.java
index e5b72d4..fb2a66c 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/NotificationEntry.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/NotificationEntry.java
@@ -297,20 +297,20 @@
             return NotificationEntry.this.getRow();
         }
 
-        @Nullable
         @Override
-        public EntryAdapter getGroupRoot() {
-            // TODO (b/395857098): for backwards compatibility this will return null if called
-            // on a group summary that's not in a bundles, but it should return itself.
+        public boolean isGroupRoot() {
             if (isTopLevelEntry() || getParent() == null) {
-                return null;
+                return false;
             }
             if (NotificationEntry.this.getParent() instanceof GroupEntry parentGroupEntry) {
-                if (parentGroupEntry.getSummary() != null) {
-                    return parentGroupEntry.getSummary().mEntryAdapter;
-                }
+                return parentGroupEntry.getSummary() == NotificationEntry.this;
             }
-            return null;
+            return false;
+        }
+
+        @Override
+        public StateFlow<Boolean> isSensitive() {
+            return NotificationEntry.this.isSensitive();
         }
 
         @Override
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/coordinator/VisualStabilityCoordinator.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/coordinator/VisualStabilityCoordinator.java
index 3e5655a..bdbdc53 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/coordinator/VisualStabilityCoordinator.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/coordinator/VisualStabilityCoordinator.java
@@ -45,8 +45,8 @@
 import com.android.systemui.statusbar.notification.collection.listbuilder.OnBeforeRenderListListener;
 import com.android.systemui.statusbar.notification.collection.listbuilder.pluggable.NotifStabilityManager;
 import com.android.systemui.statusbar.notification.collection.provider.VisualStabilityProvider;
+import com.android.systemui.statusbar.notification.data.repository.HeadsUpRepository;
 import com.android.systemui.statusbar.notification.domain.interactor.SeenNotificationsInteractor;
-import com.android.systemui.statusbar.notification.headsup.HeadsUpManager;
 import com.android.systemui.statusbar.notification.shared.NotificationMinimalism;
 import com.android.systemui.statusbar.policy.KeyguardStateController;
 import com.android.systemui.util.concurrency.DelayableExecutor;
@@ -72,7 +72,7 @@
 public class VisualStabilityCoordinator implements Coordinator, Dumpable {
     private final DelayableExecutor mDelayableExecutor;
     private final DelayableExecutor mMainExecutor;
-    private final HeadsUpManager mHeadsUpManager;
+    private final HeadsUpRepository mHeadsUpRepository;
     private final SeenNotificationsInteractor mSeenNotificationsInteractor;
     private final ShadeAnimationInteractor mShadeAnimationInteractor;
     private final StatusBarStateController mStatusBarStateController;
@@ -94,6 +94,7 @@
     private boolean mNotifPanelLaunchingActivity;
     private boolean mCommunalShowing = false;
     private boolean mLockscreenShowing = false;
+    private boolean mTrackingHeadsUp = false;
     private boolean mLockscreenInGoneTransition = false;
     private Set<String> mHeadsUpGroupKeys = new HashSet<>();
 
@@ -117,7 +118,7 @@
             @Background DelayableExecutor delayableExecutor,
             @Main DelayableExecutor mainExecutor,
             DumpManager dumpManager,
-            HeadsUpManager headsUpManager,
+            HeadsUpRepository headsUpRepository,
             ShadeAnimationInteractor shadeAnimationInteractor,
             JavaAdapter javaAdapter,
             SeenNotificationsInteractor seenNotificationsInteractor,
@@ -130,7 +131,7 @@
             KeyguardTransitionInteractor keyguardTransitionInteractor,
             KeyguardStateController keyguardStateController,
             VisualStabilityCoordinatorLogger logger) {
-        mHeadsUpManager = headsUpManager;
+        mHeadsUpRepository = headsUpRepository;
         mShadeAnimationInteractor = shadeAnimationInteractor;
         mJavaAdapter = javaAdapter;
         mSeenNotificationsInteractor = seenNotificationsInteractor;
@@ -177,6 +178,8 @@
             mJavaAdapter.alwaysCollectFlow(mKeyguardTransitionInteractor.transitionValue(
                             KeyguardState.LOCKSCREEN),
                     this::onLockscreenKeyguardStateTransitionValueChanged);
+            mJavaAdapter.alwaysCollectFlow(mHeadsUpRepository.isTrackingHeadsUp(),
+                    this::onTrackingHeadsUpModeChanged);
         }
 
         if (Flags.checkLockscreenGoneTransition()) {
@@ -239,7 +242,7 @@
                     boolean isTopUnseen = NotificationMinimalism.isEnabled()
                             && (mSeenNotificationsInteractor.isTopUnseenNotification(entry)
                                 || mSeenNotificationsInteractor.isTopOngoingNotification(entry));
-                    if (isTopUnseen || mHeadsUpManager.isHeadsUpEntry(entry.getKey())) {
+                    if (isTopUnseen || mHeadsUpRepository.isHeadsUpEntry(entry.getKey())) {
                         return !mVisibilityLocationProvider.isInVisibleLocation(entry);
                     }
                     return false;
@@ -275,7 +278,7 @@
                         return false;
                     }
 
-                    return mHeadsUpManager.isHeadsUpEntry(summary.getKey());
+                    return mHeadsUpRepository.isHeadsUpEntry(summary.getKey());
                 }
                 /**
                  * When reordering is enabled, non-heads-up groups can be pruned.
@@ -415,7 +418,7 @@
                             if (summary == null) continue;
 
                             final String summaryKey = summary.getKey();
-                            if (mHeadsUpManager.isHeadsUpEntry(summaryKey)) {
+                            if (mHeadsUpRepository.isHeadsUpEntry(summaryKey)) {
                                 currentHeadsUpKeys.add(summaryKey);
                             }
                         }
@@ -475,11 +478,19 @@
 
     private boolean isReorderingAllowed() {
         final boolean sleepyAndDozed = mFullyDozed && mSleepy;
-        final boolean stackShowing = mPanelExpanded
-                || (SceneContainerFlag.isEnabled() && mLockscreenShowing);
+        final boolean stackShowing = isStackShowing();
         return (sleepyAndDozed || !stackShowing || mCommunalShowing) && !mPulsing;
     }
 
+    /** @return true when the notification stack is visible to the user */
+    private boolean isStackShowing() {
+        if (SceneContainerFlag.isEnabled()) {
+            return mPanelExpanded || mLockscreenShowing || mTrackingHeadsUp;
+        } else {
+            return mPanelExpanded;
+        }
+    }
+
     /**
      * Allows this notification entry to be re-ordered in the notification list temporarily until
      * the timeout has passed.
@@ -610,6 +621,11 @@
         updateAllowedStates("lockscreenShowing", isShowing);
     }
 
+    private void onTrackingHeadsUpModeChanged(boolean isTrackingHeadsUp) {
+        mTrackingHeadsUp = isTrackingHeadsUp;
+        updateAllowedStates("trackingHeadsUp", isTrackingHeadsUp);
+    }
+
     private void onLockscreenInGoneTransitionChanged(boolean inGoneTransition) {
         if (!Flags.checkLockscreenGoneTransition()) {
             return;
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/render/GroupExpansionManagerImpl.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/render/GroupExpansionManagerImpl.java
index b179a69..c5ae875 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/render/GroupExpansionManagerImpl.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/render/GroupExpansionManagerImpl.java
@@ -29,6 +29,7 @@
 import com.android.systemui.statusbar.notification.collection.NotifPipeline;
 import com.android.systemui.statusbar.notification.collection.NotificationEntry;
 import com.android.systemui.statusbar.notification.collection.listbuilder.OnBeforeRenderListListener;
+import com.android.systemui.statusbar.notification.row.ExpandableNotificationRow;
 import com.android.systemui.statusbar.notification.shared.NotificationBundleUi;
 
 import java.io.PrintWriter;
@@ -162,41 +163,38 @@
     @Override
     public boolean isGroupExpanded(EntryAdapter entry) {
         NotificationBundleUi.assertInNewMode();
-        return mExpandedCollections.contains(mGroupMembershipManager.getGroupRoot(entry));
+        ExpandableNotificationRow parent = entry.getRow().getNotificationParent();
+        return mExpandedCollections.contains(entry)
+                || (parent != null && mExpandedCollections.contains(parent.getEntryAdapter()));
     }
 
     @Override
-    public void setGroupExpanded(EntryAdapter entry, boolean expanded) {
+    public void setGroupExpanded(EntryAdapter groupRoot, boolean expanded) {
         NotificationBundleUi.assertInNewMode();
-        EntryAdapter groupParent = mGroupMembershipManager.getGroupRoot(entry);
-        if (!entry.isAttached()) {
+        if (!groupRoot.isAttached()) {
             if (expanded) {
                 Log.wtf(TAG, "Cannot expand group that is not attached");
-            } else {
-                // The entry is no longer attached, but we still want to make sure we don't have
-                // a stale expansion state.
-                groupParent = entry;
             }
         }
 
         boolean changed;
         if (expanded) {
-            changed = mExpandedCollections.add(groupParent);
+            changed = mExpandedCollections.add(groupRoot);
         } else {
-            changed = mExpandedCollections.remove(groupParent);
+            changed = mExpandedCollections.remove(groupRoot);
         }
 
         // Only notify listeners if something changed.
         if (changed) {
-            sendOnGroupExpandedChange(entry, expanded);
+            sendOnGroupExpandedChange(groupRoot, expanded);
         }
     }
 
     @Override
-    public boolean toggleGroupExpansion(EntryAdapter entry) {
+    public boolean toggleGroupExpansion(EntryAdapter groupRoot) {
         NotificationBundleUi.assertInNewMode();
-        setGroupExpanded(entry, !isGroupExpanded(entry));
-        return isGroupExpanded(entry);
+        setGroupExpanded(groupRoot, !isGroupExpanded(groupRoot));
+        return isGroupExpanded(groupRoot);
     }
 
     @Override
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/render/GroupMembershipManager.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/render/GroupMembershipManager.java
index 3edbfaf..86aa4a3 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/render/GroupMembershipManager.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/render/GroupMembershipManager.java
@@ -51,17 +51,6 @@
     NotificationEntry getGroupSummary(@NonNull NotificationEntry entry);
 
     /**
-     * Gets the EntryAdapter that is the nearest root of the collection of rows the given entry
-     * belongs to. If the given entry is a BundleEntry or an isolated child of a BundleEntry, the
-     * BundleEntry will be returned. If the given notification is a group summary NotificationEntry,
-     * or a child of a group summary, the summary NotificationEntry will be returned, even if that
-     * summary belongs to a BundleEntry. If the entry is a notification that does not belong to any
-     * group or bundle grouping, null will be returned.
-     */
-    @Nullable
-    EntryAdapter getGroupRoot(@NonNull EntryAdapter entry);
-
-    /**
      * @return whether a given notification is a child in a group
      */
     boolean isChildInGroup(@NonNull NotificationEntry entry);
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/render/GroupMembershipManagerImpl.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/render/GroupMembershipManagerImpl.java
index a1a23e3..aec0d70 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/render/GroupMembershipManagerImpl.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/render/GroupMembershipManagerImpl.java
@@ -60,7 +60,7 @@
     @Override
     public boolean isGroupRoot(@NonNull EntryAdapter entry) {
         NotificationBundleUi.assertInNewMode();
-        return entry == entry.getGroupRoot();
+        return entry.isGroupRoot();
     }
 
     @Nullable
@@ -76,13 +76,6 @@
         return null;
     }
 
-    @Nullable
-    @Override
-    public EntryAdapter getGroupRoot(@NonNull EntryAdapter entry) {
-        NotificationBundleUi.assertInNewMode();
-        return entry.getGroupRoot();
-    }
-
     @Override
     public boolean isChildInGroup(@NonNull NotificationEntry entry) {
         NotificationBundleUi.assertInLegacyMode();
@@ -94,7 +87,7 @@
     public boolean isChildInGroup(@NonNull EntryAdapter entry) {
         NotificationBundleUi.assertInNewMode();
         // An entry is a child if it's not a group root or top level entry, but it is attached.
-        return entry.isAttached() && entry != getGroupRoot(entry) && !entry.isTopLevelEntry();
+        return entry.isAttached() && !entry.isGroupRoot() && !entry.isTopLevelEntry();
     }
 
     @Nullable
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/data/repository/HeadsUpRepository.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/data/repository/HeadsUpRepository.kt
index 28e3995..222f94b 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/data/repository/HeadsUpRepository.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/data/repository/HeadsUpRepository.kt
@@ -40,6 +40,16 @@
     /** Set of currently active top-level heads up rows to be displayed. */
     val activeHeadsUpRows: Flow<Set<HeadsUpRowRepository>>
 
+    /** Whether the user is swiping on a heads up row */
+    val isTrackingHeadsUp: Flow<Boolean>
+
+    /** @return true if the actively managed heads up notifications contain an entry for this key */
+    fun isHeadsUpEntry(key: String): Boolean
+
+    /**
+     * set whether a HUN is currently animation out, to keep its view container visible during the
+     * animation
+     */
     fun setHeadsUpAnimatingAway(animatingAway: Boolean)
 
     /** Snooze the currently pinned HUN. */
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/headsup/HeadsUpManager.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/headsup/HeadsUpManager.kt
index 95234da..9728fcf 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/headsup/HeadsUpManager.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/headsup/HeadsUpManager.kt
@@ -19,12 +19,16 @@
 import android.graphics.Region
 import com.android.systemui.Dumpable
 import com.android.systemui.dagger.SysUISingleton
+import com.android.systemui.statusbar.notification.collection.EntryAdapter
 import com.android.systemui.statusbar.notification.collection.NotificationEntry
+import com.android.systemui.statusbar.notification.row.ExpandableNotificationRow
 import dagger.Binds
 import dagger.Module
 import java.io.PrintWriter
 import java.util.stream.Stream
 import javax.inject.Inject
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.StateFlow
 
 /**
  * A manager which handles heads up notifications which is a special mode where they simply peek
@@ -110,7 +114,7 @@
      * Returns the value of the tracking-heads-up flag. See the doc of {@code setTrackingHeadsUp} as
      * well.
      */
-    fun isTrackingHeadsUp(): Boolean
+    fun isTrackingHeadsUp(): StateFlow<Boolean>
 
     fun onExpandingFinished()
 
@@ -151,6 +155,12 @@
     fun setAnimationStateHandler(handler: AnimationStateHandler)
 
     /**
+    * Set an entry to be expanded and therefore stick in the heads up area if it's pinned until
+    * it's collapsed again.
+    */
+    fun setExpanded(key: String, row: ExpandableNotificationRow, expanded: Boolean)
+
+    /**
      * Set an entry to be expanded and therefore stick in the heads up area if it's pinned until
      * it's collapsed again.
      */
@@ -284,12 +294,12 @@
 
     override fun isHeadsUpAnimatingAwayValue() = false
 
+    override fun isTrackingHeadsUp(): StateFlow<Boolean> = MutableStateFlow(false)
+
     override fun isSnoozed(packageName: String) = false
 
     override fun isSticky(key: String?) = false
 
-    override fun isTrackingHeadsUp() = false
-
     override fun onExpandingFinished() {}
 
     override fun releaseAllImmediately() {}
@@ -308,6 +318,8 @@
 
     override fun setAnimationStateHandler(handler: AnimationStateHandler) {}
 
+    override fun setExpanded(key: String, row: ExpandableNotificationRow, expanded: Boolean) {}
+
     override fun setExpanded(entry: NotificationEntry, expanded: Boolean) {}
 
     override fun setGutsShown(entry: NotificationEntry, gutsShown: Boolean) {}
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/headsup/HeadsUpManagerImpl.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/headsup/HeadsUpManagerImpl.java
index 010f080..7c5f3b5 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/headsup/HeadsUpManagerImpl.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/headsup/HeadsUpManagerImpl.java
@@ -135,6 +135,8 @@
             StateFlowKt.MutableStateFlow(new HashSet<>());
     private final MutableStateFlow<Boolean> mHeadsUpAnimatingAway =
             StateFlowKt.MutableStateFlow(false);
+    private final MutableStateFlow<Boolean> mTrackingHeadsUp =
+            StateFlowKt.MutableStateFlow(false);
     private final HashSet<String> mSwipedOutKeys = new HashSet<>();
     private final HashSet<NotificationEntry> mEntriesToRemoveAfterExpand = new HashSet<>();
     @VisibleForTesting
@@ -142,7 +144,6 @@
             = new ArraySet<>();
 
     private boolean mReleaseOnExpandFinish;
-    private boolean mTrackingHeadsUp;
     private boolean mIsShadeOrQsExpanded;
     private boolean mIsQsExpanded;
     private int mStatusBarState;
@@ -417,8 +418,8 @@
     }
 
     @Override
-    public void setTrackingHeadsUp(boolean trackingHeadsUp) {
-        mTrackingHeadsUp = trackingHeadsUp;
+    public void setTrackingHeadsUp(boolean isTrackingHeadsUp) {
+        mTrackingHeadsUp.setValue(isTrackingHeadsUp);
     }
 
     @Override
@@ -510,9 +511,7 @@
                 || !mAvalancheController.getWaitingEntryList().isEmpty();
     }
 
-    /**
-     * @return true if the notification is managed by this manager
-     */
+    @Override
     public boolean isHeadsUpEntry(@NonNull String key) {
         return mHeadsUpEntryMap.containsKey(key) || mAvalancheController.isWaiting(key);
     }
@@ -879,10 +878,8 @@
             ExpandableNotificationRow topRow = topEntry.getRow();
             if (topEntry.rowIsChildInGroup()) {
                 if (NotificationBundleUi.isEnabled()) {
-                    final EntryAdapter adapter = mGroupMembershipManager.getGroupRoot(
-                            topRow.getEntryAdapter());
-                    if (adapter != null) {
-                        topRow = adapter.getRow();
+                    if (topRow.getNotificationParent() != null) {
+                        topRow = topRow.getNotificationParent();
                     }
                 } else {
                     final NotificationEntry groupSummary =
@@ -1066,8 +1063,9 @@
         }
     }
 
+    @NonNull
     @Override
-    public boolean isTrackingHeadsUp() {
+    public StateFlow<Boolean> isTrackingHeadsUp() {
         return mTrackingHeadsUp;
     }
 
@@ -1093,7 +1091,23 @@
      * Set an entry to be expanded and therefore stick in the heads up area if it's pinned
      * until it's collapsed again.
      */
+    @Override
+    public void setExpanded(@NonNull String entryKey, @NonNull ExpandableNotificationRow row,
+            boolean expanded) {
+        NotificationBundleUi.assertInNewMode();
+        HeadsUpEntry headsUpEntry = getHeadsUpEntry(entryKey);
+        if (headsUpEntry != null && row.getPinnedStatus().isPinned()) {
+            headsUpEntry.setExpanded(expanded);
+        }
+    }
+
+    /**
+     * Set an entry to be expanded and therefore stick in the heads up area if it's pinned
+     * until it's collapsed again.
+     */
+    @Override
     public void setExpanded(@NonNull NotificationEntry entry, boolean expanded) {
+        NotificationBundleUi.assertInLegacyMode();
         HeadsUpEntry headsUpEntry = getHeadsUpEntry(entry.getKey());
         if (headsUpEntry != null && entry.isRowPinned()) {
             headsUpEntry.setExpanded(expanded);
@@ -1660,7 +1674,7 @@
                     mEntriesToRemoveWhenReorderingAllowed.add(entry);
                     mVisualStabilityProvider.addTemporaryReorderingAllowedListener(
                             mOnReorderingAllowedListener);
-                } else if (mTrackingHeadsUp) {
+                } else if (mTrackingHeadsUp.getValue()) {
                     mEntriesToRemoveAfterExpand.add(entry);
                     mLogger.logRemoveEntryAfterExpand(entry);
                 } else if (mVisualStabilityProvider.isReorderingAllowed()
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRow.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRow.java
index 4c74408..5e34b3b 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRow.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRow.java
@@ -404,17 +404,25 @@
                 : mGroupMembershipManager.isGroupSummary(mEntry);
         if (!shouldShowPublic() && (!mIsMinimized || isExpanded()) && isGroupRoot) {
             mGroupExpansionChanging = true;
-            final boolean wasExpanded = NotificationBundleUi.isEnabled()
-                    ? mGroupExpansionManager.isGroupExpanded(mEntryAdapter)
-                    : mGroupExpansionManager.isGroupExpanded(mEntry);
-            boolean nowExpanded = NotificationBundleUi.isEnabled()
-                    ? mGroupExpansionManager.toggleGroupExpansion(mEntryAdapter)
-                    : mGroupExpansionManager.toggleGroupExpansion(mEntry);
-            mOnExpandClickListener.onExpandClicked(mEntry, v, nowExpanded);
-            if (shouldLogExpandClickMetric) {
-                mMetricsLogger.action(MetricsEvent.ACTION_NOTIFICATION_GROUP_EXPANDER, nowExpanded);
+            if (NotificationBundleUi.isEnabled()) {
+                final boolean wasExpanded =  mGroupExpansionManager.isGroupExpanded(mEntryAdapter);
+                boolean nowExpanded = mGroupExpansionManager.toggleGroupExpansion(mEntryAdapter);
+                mOnExpandClickListener.onExpandClicked(this, mEntryAdapter, nowExpanded);
+                if (shouldLogExpandClickMetric) {
+                    mMetricsLogger.action(
+                            MetricsEvent.ACTION_NOTIFICATION_GROUP_EXPANDER, nowExpanded);
+                }
+                onExpansionChanged(true /* userAction */, wasExpanded);
+            } else {
+                final boolean wasExpanded = mGroupExpansionManager.isGroupExpanded(mEntry);
+                boolean nowExpanded = mGroupExpansionManager.toggleGroupExpansion(mEntry);
+                mOnExpandClickListener.onExpandClicked(mEntry, v, nowExpanded);
+                if (shouldLogExpandClickMetric) {
+                    mMetricsLogger.action(
+                            MetricsEvent.ACTION_NOTIFICATION_GROUP_EXPANDER, nowExpanded);
+                }
+                onExpansionChanged(true /* userAction */, wasExpanded);
             }
-            onExpansionChanged(true /* userAction */, wasExpanded);
         } else if (mEnableNonGroupedNotificationExpand) {
             if (v.isAccessibilityFocused()) {
                 mPrivateLayout.setFocusOnVisibilityChange();
@@ -435,7 +443,11 @@
             }
 
             notifyHeightChanged(/* needsAnimation= */ true);
-            mOnExpandClickListener.onExpandClicked(mEntry, v, nowExpanded);
+            if (NotificationBundleUi.isEnabled()) {
+                mOnExpandClickListener.onExpandClicked(this, mEntryAdapter, nowExpanded);
+            } else {
+                mOnExpandClickListener.onExpandClicked(mEntry, v, nowExpanded);
+            }
             if (shouldLogExpandClickMetric) {
                 mMetricsLogger.action(MetricsEvent.ACTION_NOTIFICATION_EXPANDER, nowExpanded);
             }
@@ -2946,7 +2958,9 @@
                 && !mChildrenContainer.showingAsLowPriority()) {
             final boolean wasExpanded = isGroupExpanded();
             if (NotificationBundleUi.isEnabled()) {
-                mGroupExpansionManager.setGroupExpanded(mEntryAdapter, userExpanded);
+                if (mEntryAdapter.isGroupRoot()) {
+                    mGroupExpansionManager.setGroupExpanded(mEntryAdapter, userExpanded);
+                }
             } else {
                 mGroupExpansionManager.setGroupExpanded(mEntry, userExpanded);
             }
@@ -3440,9 +3454,9 @@
     public void makeActionsVisibile() {
         setUserExpanded(true, true);
         if (isChildInGroup()) {
-            if (NotificationBundleUi.isEnabled()) {
-                mGroupExpansionManager.setGroupExpanded(mEntryAdapter, true);
-            } else {
+            if (!NotificationBundleUi.isEnabled()) {
+                // this is only called if row.getParent() instanceof NotificationStackScrollLayout,
+                // so there is never a group to expand
                 mGroupExpansionManager.setGroupExpanded(mEntry, true);
             }
         }
@@ -3728,7 +3742,7 @@
         if (!ignoreTemporaryStates && mGuts != null && mGuts.isExposed()) {
             return mGuts.getIntrinsicHeight();
         } else if (!ignoreTemporaryStates && canShowHeadsUp() && mIsHeadsUp
-                && mHeadsUpManager.isTrackingHeadsUp()) {
+                && mHeadsUpManager.isTrackingHeadsUp().getValue()) {
             return getPinnedHeadsUpHeight(false /* atLeastMinHeight */);
         } else if (mIsSummaryWithChildren && !isGroupExpanded() && !shouldShowPublic()) {
             return mChildrenContainer.getMinHeight();
@@ -4023,6 +4037,9 @@
 
     public interface OnExpandClickListener {
         void onExpandClicked(NotificationEntry clickedEntry, View clickedView, boolean nowExpanded);
+
+        void onExpandClicked(ExpandableNotificationRow row, EntryAdapter clickedEntry,
+                boolean nowExpanded);
     }
 
     @Override
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/wrapper/NotificationViewWrapper.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/wrapper/NotificationViewWrapper.java
index 752a8ab..3987ca6 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/wrapper/NotificationViewWrapper.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/wrapper/NotificationViewWrapper.java
@@ -44,6 +44,7 @@
 import com.android.systemui.statusbar.notification.NotificationFadeAware;
 import com.android.systemui.statusbar.notification.TransformState;
 import com.android.systemui.statusbar.notification.row.ExpandableNotificationRow;
+import com.android.systemui.statusbar.notification.shared.NotificationBundleUi;
 
 /**
  * Wraps the actual notification content view; used to implement behaviors which are different for
@@ -135,7 +136,10 @@
         }
 
         // Apps targeting Q should fix their dark mode bugs.
-        if (mRow.getEntry().targetSdk >= Build.VERSION_CODES.Q) {
+        int targetSdk = NotificationBundleUi.isEnabled()
+                ? mRow.getEntryAdapter().getTargetSdk()
+                : mRow.getEntry().targetSdk;
+        if (targetSdk >= Build.VERSION_CODES.Q) {
             return false;
         }
 
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/MagneticNotificationRowManager.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/MagneticNotificationRowManager.kt
index 02336e4..aa69517 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/MagneticNotificationRowManager.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/MagneticNotificationRowManager.kt
@@ -87,6 +87,9 @@
      */
     fun onMagneticInteractionEnd(row: ExpandableNotificationRow, velocity: Float? = null)
 
+    /* Reset any roundness that magnetic targets may have */
+    fun resetRoundness()
+
     /**
      * Reset any magnetic and roundable targets set, as well as any internal state.
      *
@@ -124,6 +127,8 @@
                         velocity: Float?,
                     ) {}
 
+                    override fun resetRoundness() {}
+
                     override fun reset() {}
                 }
     }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/MagneticNotificationRowManagerImpl.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/MagneticNotificationRowManagerImpl.kt
index de4af37..da98858 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/MagneticNotificationRowManagerImpl.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/MagneticNotificationRowManagerImpl.kt
@@ -27,6 +27,7 @@
 import javax.inject.Inject
 import kotlin.math.abs
 import kotlin.math.pow
+import org.jetbrains.annotations.TestOnly
 
 @SysUISingleton
 class MagneticNotificationRowManagerImpl
@@ -41,15 +42,16 @@
     var currentState = State.IDLE
         private set
 
-    // Magnetic and roundable targets
+    // Magnetic targets
     var currentMagneticListeners = listOf<MagneticRowListener?>()
         private set
 
-    var currentRoundableTargets: RoundableTargets? = null
-        private set
-
     private var magneticDetachThreshold = Float.POSITIVE_INFINITY
 
+    // Has the roundable target been set for the magnetic view that is being swiped.
+    val isSwipedViewRoundableSet: Boolean
+        @TestOnly get() = notificationRoundnessManager.isSwipedViewSet
+
     // Animation spring forces
     private val detachForce =
         SpringForce().setStiffness(DETACH_STIFFNESS).setDampingRatio(DETACH_DAMPING_RATIO)
@@ -83,12 +85,14 @@
         sectionsManager: NotificationSectionsManager,
     ) {
         // Update roundable targets
-        currentRoundableTargets =
+        notificationRoundnessManager.clear()
+        val currentRoundableTargets =
             notificationTargetsHelper.findRoundableTargets(
                 expandableNotificationRow,
                 stackScrollLayout,
                 sectionsManager,
             )
+        notificationRoundnessManager.setRoundableTargets(currentRoundableTargets)
 
         // Update magnetic targets
         val newListeners =
@@ -127,6 +131,7 @@
                 currentState = State.PULLING
             }
             State.PULLING -> {
+                updateRoundness(translation)
                 if (canTargetBeDismissed) {
                     pullDismissibleRow(translation)
                 } else {
@@ -141,6 +146,14 @@
         return true
     }
 
+    private fun updateRoundness(translation: Float) {
+        val normalizedTranslation = abs(swipedRowMultiplier * translation) / magneticDetachThreshold
+        notificationRoundnessManager.setRoundnessForAffectedViews(
+            /* roundness */ normalizedTranslation.coerceIn(0f, MAX_PRE_DETACH_ROUNDNESS),
+            /* animate */ false,
+        )
+    }
+
     private fun pullDismissibleRow(translation: Float) {
         val targetTranslation = swipedRowMultiplier * translation
         val crossedThreshold = abs(targetTranslation) >= magneticDetachThreshold
@@ -203,9 +216,10 @@
     private fun detach(listener: MagneticRowListener, toPosition: Float) {
         listener.cancelMagneticAnimations()
         listener.triggerMagneticForce(toPosition, detachForce)
-        currentRoundableTargets?.let {
-            notificationRoundnessManager.setViewsAffectedBySwipe(it.before, it.swiped, it.after)
-        }
+        notificationRoundnessManager.setRoundnessForAffectedViews(
+            /* roundness */ 1f,
+            /* animate */ true,
+        )
         msdlPlayer.playToken(MSDLToken.SWIPE_THRESHOLD_INDICATOR)
     }
 
@@ -240,6 +254,8 @@
         }
     }
 
+    override fun resetRoundness() = notificationRoundnessManager.clear()
+
     override fun reset() {
         currentMagneticListeners.forEach {
             it?.cancelMagneticAnimations()
@@ -247,7 +263,7 @@
         }
         currentState = State.IDLE
         currentMagneticListeners = listOf()
-        currentRoundableTargets = null
+        notificationRoundnessManager.clear()
     }
 
     private fun List<MagneticRowListener?>.swipedListener(): MagneticRowListener? =
@@ -256,6 +272,11 @@
     private fun ExpandableNotificationRow.isSwipedTarget(): Boolean =
         magneticRowListener == currentMagneticListeners.swipedListener()
 
+    private fun NotificationRoundnessManager.clear() = setViewsAffectedBySwipe(null, null, null)
+
+    private fun NotificationRoundnessManager.setRoundableTargets(targets: RoundableTargets) =
+        setViewsAffectedBySwipe(targets.before, targets.swiped, targets.after)
+
     enum class State {
         IDLE,
         TARGETS_SET,
@@ -280,6 +301,9 @@
         private const val SNAP_BACK_STIFFNESS = 550f
         private const val SNAP_BACK_DAMPING_RATIO = 0.52f
 
+        // Maximum value of corner roundness that gets applied during the pre-detach dragging
+        private const val MAX_PRE_DETACH_ROUNDNESS = 0.8f
+
         private val VIBRATION_ATTRIBUTES_PIPELINING =
             VibrationAttributes.Builder()
                 .setUsage(VibrationAttributes.USAGE_TOUCH)
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationRoundnessManager.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationRoundnessManager.java
index fa1843e..a53e837 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationRoundnessManager.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationRoundnessManager.java
@@ -71,10 +71,8 @@
             Roundable viewBefore,
             ExpandableView viewSwiped,
             Roundable viewAfter) {
-        // This method requires you to change the roundness of the current View targets and reset
-        // the roundness of the old View targets (if any) to 0f.
-        // To avoid conflicts, it generates a set of old Views and removes the current Views
-        // from this set.
+        // This method caches a new set of current View targets and reset the roundness of the old
+        // View targets (if any) to 0f.
         HashSet<Roundable> oldViews = new HashSet<>();
         if (mViewBeforeSwipedView != null) oldViews.add(mViewBeforeSwipedView);
         if (mSwipedView != null) oldViews.add(mSwipedView);
@@ -83,19 +81,16 @@
         mViewBeforeSwipedView = viewBefore;
         if (viewBefore != null) {
             oldViews.remove(viewBefore);
-            viewBefore.requestRoundness(/* top = */ 0f, /* bottom = */ 1f, DISMISS_ANIMATION);
         }
 
         mSwipedView = viewSwiped;
         if (viewSwiped != null) {
             oldViews.remove(viewSwiped);
-            viewSwiped.requestRoundness(/* top = */ 1f, /* bottom = */ 1f, DISMISS_ANIMATION);
         }
 
         mViewAfterSwipedView = viewAfter;
         if (viewAfter != null) {
             oldViews.remove(viewAfter);
-            viewAfter.requestRoundness(/* top = */ 1f, /* bottom = */ 0f, DISMISS_ANIMATION);
         }
 
         // After setting the current Views, reset the views that are still present in the set.
@@ -104,6 +99,34 @@
         }
     }
 
+    void setRoundnessForAffectedViews(float roundness) {
+        if (mViewBeforeSwipedView != null) {
+            mViewBeforeSwipedView.requestBottomRoundness(roundness, DISMISS_ANIMATION);
+        }
+
+        if (mSwipedView != null) {
+            mSwipedView.requestRoundness(roundness, roundness, DISMISS_ANIMATION);
+        }
+
+        if (mViewAfterSwipedView != null) {
+            mViewAfterSwipedView.requestTopRoundness(roundness, DISMISS_ANIMATION);
+        }
+    }
+
+    void setRoundnessForAffectedViews(float roundness, boolean animate) {
+        if (mViewBeforeSwipedView != null) {
+            mViewBeforeSwipedView.requestBottomRoundness(roundness, DISMISS_ANIMATION, animate);
+        }
+
+        if (mSwipedView != null) {
+            mSwipedView.requestRoundness(roundness, roundness, DISMISS_ANIMATION, animate);
+        }
+
+        if (mViewAfterSwipedView != null) {
+            mViewAfterSwipedView.requestTopRoundness(roundness, DISMISS_ANIMATION, animate);
+        }
+    }
+
     void setClearAllInProgress(boolean isClearingAll) {
         mIsClearAllInProgress = isClearingAll;
     }
@@ -138,4 +161,8 @@
     public void setShouldRoundPulsingViews(boolean shouldRoundPulsingViews) {
         mRoundForPulsingViews = shouldRoundPulsingViews;
     }
+
+    public boolean isSwipedViewSet() {
+        return mSwipedView != null;
+    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayout.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayout.java
index 3ff18ef..89ede09 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayout.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayout.java
@@ -1810,16 +1810,22 @@
 
     private ExpandableNotificationRow getTopHeadsUpRow() {
         ExpandableNotificationRow row = mTopHeadsUpRow;
-        if (row.isChildInGroup()) {
-            final NotificationEntry groupSummary =
-                    mGroupMembershipManager.getGroupSummary(row.getEntry());
-            if (groupSummary != null) {
-                row = groupSummary.getRow();
+        if (NotificationBundleUi.isEnabled()) {
+            if (mGroupMembershipManager.isChildInGroup(row.getEntryAdapter())
+                    && row.isChildInGroup()) {
+                row = row.getNotificationParent();
+            }
+        } else {
+            if (row.isChildInGroup()) {
+                final NotificationEntry groupSummary =
+                        mGroupMembershipManager.getGroupSummary(row.getEntry());
+                if (groupSummary != null) {
+                    row = groupSummary.getRow();
+                }
             }
         }
         return row;
     }
-
     /**
      * @return the position from where the appear transition ends when expanding.
      * Measured in absolute height.
@@ -1966,10 +1972,19 @@
                     && touchY >= top && touchY <= bottom && touchX >= left && touchX <= right) {
                 if (slidingChild instanceof ExpandableNotificationRow row) {
                     NotificationEntry entry = row.getEntry();
+                    boolean isEntrySummaryForTopHun;
+                    if (NotificationBundleUi.isEnabled()) {
+                        isEntrySummaryForTopHun = Objects.equals(
+                                ((ExpandableNotificationRow) slidingChild).getNotificationParent(),
+                                mTopHeadsUpRow);
+                    } else {
+                        isEntrySummaryForTopHun =
+                                mGroupMembershipManager.getGroupSummary(mTopHeadsUpRow.getEntry())
+                                == entry;
+                    }
                     if (!mIsExpanded && row.isHeadsUp() && row.isPinned()
                             && mTopHeadsUpRow != row
-                            && mGroupMembershipManager.getGroupSummary(mTopHeadsUpRow.getEntry())
-                            != entry) {
+                            && !isEntrySummaryForTopHun) {
                         continue;
                     }
                     return row.getViewAtPosition(touchY - childTop);
@@ -5825,7 +5840,8 @@
                             targets.getBefore(),
                             targets.getSwiped(),
                             targets.getAfter());
-
+            mController.getNotificationRoundnessManager()
+                    .setRoundnessForAffectedViews(/* roundness */ 1f);
         }
 
         updateFirstAndLastBackgroundViews();
@@ -5836,8 +5852,10 @@
 
     void onSwipeEnd() {
         updateFirstAndLastBackgroundViews();
-        mController.getNotificationRoundnessManager()
-                .setViewsAffectedBySwipe(null, null, null);
+        if (!magneticNotificationSwipes()) {
+            mController.getNotificationRoundnessManager()
+                    .setViewsAffectedBySwipe(null, null, null);
+        }
         // Round bottom corners for notification right before shelf.
         mShelf.updateAppearance();
     }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayoutController.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayoutController.java
index 888c8cc..01ef90a 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayoutController.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayoutController.java
@@ -545,6 +545,7 @@
 
                 public void handleChildViewDismissed(View view) {
                     // The View needs to clean up the Swipe states, e.g. roundness.
+                    mMagneticNotificationRowManager.resetRoundness();
                     mView.onSwipeEnd();
                     if (mView.getClearAllInProgress()) {
                         return;
@@ -629,7 +630,7 @@
                 @Override
                 public void onChildSnapBackOvershoots() {
                     if (Flags.magneticNotificationSwipes()) {
-                        mNotificationRoundnessManager.setViewsAffectedBySwipe(null, null, null);
+                        mMagneticNotificationRowManager.resetRoundness();
                     }
                 }
 
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBarNotificationPresenter.java b/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBarNotificationPresenter.java
index 4d1d64e..74b1c3b 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBarNotificationPresenter.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBarNotificationPresenter.java
@@ -58,6 +58,7 @@
 import com.android.systemui.statusbar.SysuiStatusBarStateController;
 import com.android.systemui.statusbar.notification.AboveShelfObserver;
 import com.android.systemui.statusbar.notification.DynamicPrivacyController;
+import com.android.systemui.statusbar.notification.collection.EntryAdapter;
 import com.android.systemui.statusbar.notification.collection.NotificationEntry;
 import com.android.systemui.statusbar.notification.collection.render.NotifShadeEventSource;
 import com.android.systemui.statusbar.notification.domain.interactor.NotificationAlertsInteractor;
@@ -262,6 +263,23 @@
         }
     }
 
+    @Override
+    public void onExpandClicked(ExpandableNotificationRow row, EntryAdapter clickedEntry,
+            boolean nowExpanded) {
+        mHeadsUpManager.setExpanded(clickedEntry.getKey(), row, nowExpanded);
+        mPowerInteractor.wakeUpIfDozing("NOTIFICATION_CLICK", PowerManager.WAKE_REASON_GESTURE);
+        if (nowExpanded) {
+            if (mStatusBarStateController.getState() == StatusBarState.KEYGUARD) {
+                mShadeTransitionController.goToLockedShade(row, /* needsQSAnimation = */ true);
+            } else if (clickedEntry.isSensitive().getValue() && isInLockedDownShade()) {
+                mStatusBarStateController.setLeaveOpenOnKeyguardHide(true);
+                // launch the bouncer if the device is locked
+                mActivityStarter.dismissKeyguardThenExecute(() -> false /* dismissAction */
+                        , null /* cancelRunnable */, false /* afterKeyguardGone */);
+            }
+        }
+    }
+
     /** @return true if the Shade is shown over the Lockscreen, and the device is locked */
     private boolean isInLockedDownShade() {
         if (SceneContainerFlag.isEnabled()) {
diff --git a/packages/SystemUI/tests/src/com/android/systemui/stylus/StylusManagerTest.kt b/packages/SystemUI/tests/src/com/android/systemui/stylus/StylusManagerTest.kt
index 516541d..242865b 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/stylus/StylusManagerTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/stylus/StylusManagerTest.kt
@@ -28,6 +28,7 @@
 import com.android.dx.mockito.inline.extended.ExtendedMockito.never
 import com.android.dx.mockito.inline.extended.ExtendedMockito.times
 import com.android.dx.mockito.inline.extended.ExtendedMockito.verify
+import com.android.dx.mockito.inline.extended.MockedVoidMethod
 import com.android.dx.mockito.inline.extended.StaticMockitoSession
 import com.android.internal.logging.InstanceId
 import com.android.internal.logging.UiEventLogger
@@ -245,7 +246,7 @@
     @Test
     fun onInputDeviceAdded_btStylus_firstUsed_setsFlag() {
         stylusManager.onInputDeviceAdded(BT_STYLUS_DEVICE_ID)
-        verify({ InputSettings.setStylusEverUsed(mContext, true) }, times(1))
+        verify(MockedVoidMethod { InputSettings.setStylusEverUsed(mContext, true) }, times(1))
     }
 
     @Test
@@ -511,7 +512,7 @@
 
         stylusManager.onBatteryStateChanged(STYLUS_DEVICE_ID, 1, batteryState)
 
-        verify({ InputSettings.setStylusEverUsed(mContext, true) }, times(1))
+        verify(MockedVoidMethod { InputSettings.setStylusEverUsed(mContext, true) }, times(1))
     }
 
     @Test
@@ -612,7 +613,7 @@
 
         stylusManager.onBatteryStateChanged(STYLUS_DEVICE_ID, 1, batteryState)
 
-        verify({ InputSettings.setStylusEverUsed(mContext, true) }, never())
+        verify(MockedVoidMethod { InputSettings.setStylusEverUsed(mContext, true) }, never())
     }
 
     @Test
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRowBuilder.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRowBuilder.kt
index 09c632c..771e1a5 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRowBuilder.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRowBuilder.kt
@@ -91,16 +91,17 @@
 import com.android.systemui.util.Assert.runWithCurrentThreadAsMainThread
 import com.android.systemui.util.DeviceConfigProxyFake
 import com.android.systemui.util.concurrency.FakeExecutor
-import com.android.systemui.util.mockito.whenever
 import com.android.systemui.util.time.FakeSystemClock
 import java.util.concurrent.CountDownLatch
 import java.util.concurrent.Executor
 import java.util.concurrent.TimeUnit
+import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.test.TestScope
 import org.junit.Assert.assertTrue
 import org.mockito.ArgumentCaptor
 import org.mockito.ArgumentMatchers
 import org.mockito.Mockito
+import org.mockito.kotlin.whenever
 
 class ExpandableNotificationRowBuilder(
     private val context: Context,
@@ -149,7 +150,10 @@
 
         mGroupExpansionManager = GroupExpansionManagerImpl(mDumpManager, mGroupMembershipManager)
         mUserManager = Mockito.mock(UserManager::class.java, STUB_ONLY)
-        mHeadsUpManager = Mockito.mock(HeadsUpManager::class.java, STUB_ONLY)
+        mHeadsUpManager =
+            Mockito.mock(HeadsUpManager::class.java, STUB_ONLY).apply {
+                whenever(isTrackingHeadsUp()).thenReturn(MutableStateFlow(false))
+            }
         mIconManager =
             IconManager(
                 Mockito.mock(CommonNotifCollection::class.java, STUB_ONLY),
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/notification/stack/data/repository/HeadsUpNotificationRepositoryKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/notification/stack/data/repository/HeadsUpNotificationRepositoryKosmos.kt
index 1fa6236..3406d81 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/notification/stack/data/repository/HeadsUpNotificationRepositoryKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/notification/stack/data/repository/HeadsUpNotificationRepositoryKosmos.kt
@@ -35,6 +35,10 @@
         orderedHeadsUpRows.map { it.firstOrNull() }.distinctUntilChanged()
     override val activeHeadsUpRows: Flow<Set<HeadsUpRowRepository>> =
         orderedHeadsUpRows.map { it.toSet() }.distinctUntilChanged()
+    override val isTrackingHeadsUp: MutableStateFlow<Boolean> = MutableStateFlow(false)
+
+    override fun isHeadsUpEntry(key: String): Boolean =
+        orderedHeadsUpRows.value.any { it.key == key }
 
     override fun setHeadsUpAnimatingAway(animatingAway: Boolean) {
         isHeadsUpAnimatingAway.value = animatingAway
diff --git a/services/core/java/com/android/server/integrity/OWNERS b/services/core/java/com/android/server/integrity/OWNERS
index 33561fd..352724a 100644
--- a/services/core/java/com/android/server/integrity/OWNERS
+++ b/services/core/java/com/android/server/integrity/OWNERS
@@ -1,5 +1,4 @@
 omernebil@google.com
 khelmy@google.com
 mdchurchill@google.com
-sturla@google.com
 
diff --git a/services/core/java/com/android/server/wm/DisplayContent.java b/services/core/java/com/android/server/wm/DisplayContent.java
index da28378..a874ef6 100644
--- a/services/core/java/com/android/server/wm/DisplayContent.java
+++ b/services/core/java/com/android/server/wm/DisplayContent.java
@@ -46,6 +46,7 @@
 import static android.view.Display.FLAG_SHOULD_SHOW_SYSTEM_DECORATIONS;
 import static android.view.Display.INVALID_DISPLAY;
 import static android.view.Display.STATE_UNKNOWN;
+import static android.view.Display.TYPE_EXTERNAL;
 import static android.view.Display.isSuspendedState;
 import static android.view.InsetsSource.ID_IME;
 import static android.view.Surface.ROTATION_0;
@@ -155,6 +156,7 @@
 import static com.android.server.wm.utils.DisplayInfoOverrides.copyDisplayInfoFields;
 import static com.android.server.wm.utils.RegionUtils.forEachRectReverse;
 import static com.android.server.wm.utils.RegionUtils.rectListToRegion;
+import static com.android.window.flags.Flags.enablePersistingDensityScaleForConnectedDisplays;
 
 import android.annotation.IntDef;
 import android.annotation.NonNull;
@@ -426,6 +428,12 @@
      * @see WindowManagerService#setForcedDisplayDensityForUser(int, int, int)
      */
     int mBaseDisplayDensity = 0;
+
+    /**
+     * Ratio between overridden display density for current user and the initial display density,
+     * used only for external displays.
+     */
+    float mExternalDisplayForcedDensityRatio = 0.0f;
     boolean mIsDensityForced = false;
 
     /**
@@ -3065,6 +3073,17 @@
                 mDisplayPolicy.physicalDisplayChanged();
             }
 
+            // Real display metrics changed, so we should also update initial values.
+            mInitialDisplayWidth = newWidth;
+            mInitialDisplayHeight = newHeight;
+            mInitialDisplayDensity = newDensity;
+            mInitialPhysicalXDpi = newXDpi;
+            mInitialPhysicalYDpi = newYDpi;
+            mInitialDisplayCutout = newCutout;
+            mInitialRoundedCorners = newRoundedCorners;
+            mInitialDisplayShape = newDisplayShape;
+            mCurrentUniqueDisplayId = newUniqueId;
+
             // If there is an override set for base values - use it, otherwise use new values.
             updateBaseDisplayMetrics(mIsSizeForced ? mBaseDisplayWidth : newWidth,
                     mIsSizeForced ? mBaseDisplayHeight : newHeight,
@@ -3081,16 +3100,6 @@
                 mWmService.mDisplayWindowSettings.applyRotationSettingsToDisplayLocked(this);
             }
 
-            // Real display metrics changed, so we should also update initial values.
-            mInitialDisplayWidth = newWidth;
-            mInitialDisplayHeight = newHeight;
-            mInitialDisplayDensity = newDensity;
-            mInitialPhysicalXDpi = newXDpi;
-            mInitialPhysicalYDpi = newYDpi;
-            mInitialDisplayCutout = newCutout;
-            mInitialRoundedCorners = newRoundedCorners;
-            mInitialDisplayShape = newDisplayShape;
-            mCurrentUniqueDisplayId = newUniqueId;
             reconfigureDisplayLocked();
 
             if (physicalDisplayChanged) {
@@ -3143,6 +3152,12 @@
                         + mBaseDisplayHeight + " on display:" + getDisplayId());
             }
         }
+        // Update the base density if there is a forced density ratio.
+        if (enablePersistingDensityScaleForConnectedDisplays()
+                && mIsDensityForced && mExternalDisplayForcedDensityRatio != 0.0f) {
+            mBaseDisplayDensity = (int)
+                    (mInitialDisplayDensity * mExternalDisplayForcedDensityRatio + 0.5);
+        }
         if (mDisplayReady && !mDisplayPolicy.shouldKeepCurrentDecorInsets()) {
             mDisplayPolicy.mDecorInsets.invalidate();
         }
@@ -3172,6 +3187,14 @@
         if (density == getInitialDisplayDensity()) {
             density = 0;
         }
+        // Save the new density ratio to settings for external displays.
+        if (enablePersistingDensityScaleForConnectedDisplays()
+                && mDisplayInfo.type == TYPE_EXTERNAL) {
+            mExternalDisplayForcedDensityRatio = (float)
+                    mBaseDisplayDensity / getInitialDisplayDensity();
+            mWmService.mDisplayWindowSettings.setForcedDensityRatio(getDisplayInfo(),
+                    mExternalDisplayForcedDensityRatio);
+        }
         mWmService.mDisplayWindowSettings.setForcedDensity(getDisplayInfo(), density, userId);
     }
 
diff --git a/services/core/java/com/android/server/wm/DisplayPolicy.java b/services/core/java/com/android/server/wm/DisplayPolicy.java
index a03ecf5..4908df0 100644
--- a/services/core/java/com/android/server/wm/DisplayPolicy.java
+++ b/services/core/java/com/android/server/wm/DisplayPolicy.java
@@ -1875,18 +1875,40 @@
     }
 
     void notifyDisplayAddSystemDecorations() {
-        mHandler.post(() -> {
+        if (enableDisplayContentModeManagement()) {
             final int displayId = getDisplayId();
-            StatusBarManagerInternal statusBar = getStatusBarManagerInternal();
-            if (statusBar != null) {
-                statusBar.onDisplayAddSystemDecorations(displayId);
-            }
-            final WallpaperManagerInternal wpMgr = LocalServices
-                    .getService(WallpaperManagerInternal.class);
-            if (wpMgr != null) {
-                wpMgr.onDisplayAddSystemDecorations(displayId);
-            }
-        });
+            final boolean isSystemDecorationsSupported =
+                    mDisplayContent.isSystemDecorationsSupported();
+            final boolean isHomeSupported = mDisplayContent.isHomeSupported();
+            mHandler.post(() -> {
+                if (isSystemDecorationsSupported) {
+                    StatusBarManagerInternal statusBar = getStatusBarManagerInternal();
+                    if (statusBar != null) {
+                        statusBar.onDisplayAddSystemDecorations(displayId);
+                    }
+                }
+                if (isHomeSupported) {
+                    final WallpaperManagerInternal wpMgr =
+                            LocalServices.getService(WallpaperManagerInternal.class);
+                    if (wpMgr != null) {
+                        wpMgr.onDisplayAddSystemDecorations(displayId);
+                    }
+                }
+            });
+        } else {
+            mHandler.post(() -> {
+                final int displayId = getDisplayId();
+                StatusBarManagerInternal statusBar = getStatusBarManagerInternal();
+                if (statusBar != null) {
+                    statusBar.onDisplayAddSystemDecorations(displayId);
+                }
+                final WallpaperManagerInternal wpMgr = LocalServices
+                        .getService(WallpaperManagerInternal.class);
+                if (wpMgr != null) {
+                    wpMgr.onDisplayAddSystemDecorations(displayId);
+                }
+            });
+        }
     }
 
     void notifyDisplayRemoveSystemDecorations() {
diff --git a/services/core/java/com/android/server/wm/DisplayWindowSettings.java b/services/core/java/com/android/server/wm/DisplayWindowSettings.java
index 1173875..c6892e9 100644
--- a/services/core/java/com/android/server/wm/DisplayWindowSettings.java
+++ b/services/core/java/com/android/server/wm/DisplayWindowSettings.java
@@ -23,11 +23,10 @@
 import static android.view.WindowManager.REMOVE_CONTENT_MODE_MOVE_TO_PRIMARY;
 import static android.view.WindowManager.REMOVE_CONTENT_MODE_UNDEFINED;
 
+import static com.android.server.display.feature.flags.Flags.enableDisplayContentModeManagement;
 import static com.android.server.wm.DisplayContent.FORCE_SCALING_MODE_AUTO;
 import static com.android.server.wm.DisplayContent.FORCE_SCALING_MODE_DISABLED;
 
-import static com.android.server.display.feature.flags.Flags.enableDisplayContentModeManagement;
-
 import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.app.WindowConfiguration;
@@ -100,6 +99,13 @@
         mSettingsProvider.updateOverrideSettings(info, overrideSettings);
     }
 
+    void setForcedDensityRatio(@NonNull DisplayInfo info, float ratio) {
+        final SettingsProvider.SettingsEntry overrideSettings =
+                mSettingsProvider.getOverrideSettings(info);
+        overrideSettings.mForcedDensityRatio = ratio;
+        mSettingsProvider.updateOverrideSettings(info, overrideSettings);
+    }
+
     void setForcedScalingMode(@NonNull DisplayContent displayContent, @ForceScalingMode int mode) {
         if (displayContent.isDefaultDisplay) {
             Settings.Global.putInt(mService.mContext.getContentResolver(),
@@ -367,6 +373,7 @@
                 mFixedToUserRotation);
 
         final boolean hasDensityOverride = settings.mForcedDensity != 0;
+        final boolean hasDensityOverrideRatio = settings.mForcedDensityRatio != 0.0f;
         final boolean hasSizeOverride = settings.mForcedWidth != 0 && settings.mForcedHeight != 0;
         dc.mIsDensityForced = hasDensityOverride;
         dc.mIsSizeForced = hasSizeOverride;
@@ -378,6 +385,10 @@
         final int height = hasSizeOverride ? settings.mForcedHeight : dc.mInitialDisplayHeight;
         final int density = hasDensityOverride ? settings.mForcedDensity
                 : dc.getInitialDisplayDensity();
+        if (hasDensityOverrideRatio) {
+            dc.mExternalDisplayForcedDensityRatio = settings.mForcedDensityRatio;
+        }
+
         dc.updateBaseDisplayMetrics(width, height, density, dc.mBaseDisplayPhysicalXDpi,
                 dc.mBaseDisplayPhysicalYDpi);
 
@@ -496,6 +507,13 @@
             int mForcedWidth;
             int mForcedHeight;
             int mForcedDensity;
+            /**
+             * The ratio of the forced density to the initial density of the display. This is only
+             * saved for external displays, and used to make sure ratio between forced density and
+             * initial density persist when a resolution change happens. Ratio is updated when
+             * mForcedDensity is changed.
+             */
+            float mForcedDensityRatio;
             @Nullable
             @ForceScalingMode
             Integer mForcedScalingMode;
@@ -561,6 +579,10 @@
                     mForcedDensity = other.mForcedDensity;
                     changed = true;
                 }
+                if (other.mForcedDensityRatio != mForcedDensityRatio) {
+                    mForcedDensityRatio = other.mForcedDensityRatio;
+                    changed = true;
+                }
                 if (!Objects.equals(other.mForcedScalingMode, mForcedScalingMode)) {
                     mForcedScalingMode = other.mForcedScalingMode;
                     changed = true;
@@ -649,6 +671,11 @@
                     mForcedDensity = delta.mForcedDensity;
                     changed = true;
                 }
+                if (delta.mForcedDensityRatio != 0
+                        && delta.mForcedDensityRatio != mForcedDensityRatio) {
+                    mForcedDensityRatio = delta.mForcedDensityRatio;
+                    changed = true;
+                }
                 if (delta.mForcedScalingMode != null
                         && !Objects.equals(delta.mForcedScalingMode, mForcedScalingMode)) {
                     mForcedScalingMode = delta.mForcedScalingMode;
@@ -713,6 +740,7 @@
                         && mUserRotationMode == null
                         && mUserRotation == null
                         && mForcedWidth == 0 && mForcedHeight == 0 && mForcedDensity == 0
+                        && mForcedDensityRatio == 0.0f
                         && mForcedScalingMode == null
                         && mRemoveContentMode == REMOVE_CONTENT_MODE_UNDEFINED
                         && mShouldShowWithInsecureKeyguard == null
@@ -736,6 +764,7 @@
                         && mForcedHeight == that.mForcedHeight
                         && mForcedDensity == that.mForcedDensity
                         && mRemoveContentMode == that.mRemoveContentMode
+                        && mForcedDensityRatio == that.mForcedDensityRatio
                         && Objects.equals(mUserRotationMode, that.mUserRotationMode)
                         && Objects.equals(mUserRotation, that.mUserRotation)
                         && Objects.equals(mForcedScalingMode, that.mForcedScalingMode)
@@ -755,10 +784,11 @@
             @Override
             public int hashCode() {
                 return Objects.hash(mWindowingMode, mUserRotationMode, mUserRotation, mForcedWidth,
-                        mForcedHeight, mForcedDensity, mForcedScalingMode, mRemoveContentMode,
-                        mShouldShowWithInsecureKeyguard, mShouldShowSystemDecors, mIsHomeSupported,
-                        mImePolicy, mFixedToUserRotation, mIgnoreOrientationRequest,
-                        mIgnoreDisplayCutout, mDontMoveToTop, mIgnoreActivitySizeRestrictions);
+                        mForcedHeight, mForcedDensity, mForcedDensityRatio, mForcedScalingMode,
+                        mRemoveContentMode, mShouldShowWithInsecureKeyguard,
+                        mShouldShowSystemDecors, mIsHomeSupported, mImePolicy, mFixedToUserRotation,
+                        mIgnoreOrientationRequest, mIgnoreDisplayCutout, mDontMoveToTop,
+                        mIgnoreActivitySizeRestrictions);
             }
 
             @Override
@@ -770,6 +800,7 @@
                         + ", mForcedWidth=" + mForcedWidth
                         + ", mForcedHeight=" + mForcedHeight
                         + ", mForcedDensity=" + mForcedDensity
+                        + ", mForcedDensityRatio=" + mForcedDensityRatio
                         + ", mForcedScalingMode=" + mForcedScalingMode
                         + ", mRemoveContentMode=" + mRemoveContentMode
                         + ", mShouldShowWithInsecureKeyguard=" + mShouldShowWithInsecureKeyguard
diff --git a/services/core/java/com/android/server/wm/DisplayWindowSettingsProvider.java b/services/core/java/com/android/server/wm/DisplayWindowSettingsProvider.java
index 7135c3b..e7a1fdd 100644
--- a/services/core/java/com/android/server/wm/DisplayWindowSettingsProvider.java
+++ b/services/core/java/com/android/server/wm/DisplayWindowSettingsProvider.java
@@ -511,6 +511,8 @@
                     0 /* defaultValue */);
             settingsEntry.mForcedDensity = getIntAttribute(parser, "forcedDensity",
                     0 /* defaultValue */);
+            settingsEntry.mForcedDensityRatio = parser.getAttributeFloat(null, "forcedDensityRatio",
+                    0.0f /* defaultValue */);
             settingsEntry.mForcedScalingMode = getIntegerAttribute(parser, "forcedScalingMode",
                     null /* defaultValue */);
             settingsEntry.mRemoveContentMode = getIntAttribute(parser, "removeContentMode",
@@ -599,6 +601,10 @@
                 if (settingsEntry.mForcedDensity != 0) {
                     out.attributeInt(null, "forcedDensity", settingsEntry.mForcedDensity);
                 }
+                if (settingsEntry.mForcedDensityRatio != 0.0f) {
+                    out.attributeFloat(null, "forcedDensityRatio",
+                            settingsEntry.mForcedDensityRatio);
+                }
                 if (settingsEntry.mForcedScalingMode != null) {
                     out.attributeInt(null, "forcedScalingMode",
                             settingsEntry.mForcedScalingMode);
diff --git a/services/core/java/com/android/server/wm/RootWindowContainer.java b/services/core/java/com/android/server/wm/RootWindowContainer.java
index f309372..e864ecf 100644
--- a/services/core/java/com/android/server/wm/RootWindowContainer.java
+++ b/services/core/java/com/android/server/wm/RootWindowContainer.java
@@ -2794,13 +2794,7 @@
         }
 
         startHomeOnDisplay(mCurrentUser, reason, displayContent.getDisplayId());
-        if (enableDisplayContentModeManagement()) {
-            if (displayContent.isSystemDecorationsSupported()) {
-                displayContent.getDisplayPolicy().notifyDisplayAddSystemDecorations();
-            }
-        } else {
-            displayContent.getDisplayPolicy().notifyDisplayAddSystemDecorations();
-        }
+        displayContent.getDisplayPolicy().notifyDisplayAddSystemDecorations();
     }
 
     @Override
diff --git a/services/tests/mockingservicestests/src/com/android/server/am/CachedAppOptimizerTest.java b/services/tests/mockingservicestests/src/com/android/server/am/CachedAppOptimizerTest.java
index 4b53f13..46bc70e 100644
--- a/services/tests/mockingservicestests/src/com/android/server/am/CachedAppOptimizerTest.java
+++ b/services/tests/mockingservicestests/src/com/android/server/am/CachedAppOptimizerTest.java
@@ -383,6 +383,12 @@
         // When we override new reasonable throttle values after init...
         mCountDown = new CountDownLatch(8);
         DeviceConfig.setProperty(DeviceConfig.NAMESPACE_ACTIVITY_MANAGER,
+                CachedAppOptimizer.KEY_COMPACT_THROTTLE_MIN_OOM_ADJ,
+                Long.toString(CachedAppOptimizer.DEFAULT_COMPACT_THROTTLE_MIN_OOM_ADJ + 1), false);
+        DeviceConfig.setProperty(DeviceConfig.NAMESPACE_ACTIVITY_MANAGER,
+                CachedAppOptimizer.KEY_COMPACT_THROTTLE_MAX_OOM_ADJ,
+                Long.toString(CachedAppOptimizer.DEFAULT_COMPACT_THROTTLE_MAX_OOM_ADJ - 1), false);
+        DeviceConfig.setProperty(DeviceConfig.NAMESPACE_ACTIVITY_MANAGER,
                 CachedAppOptimizer.KEY_COMPACT_THROTTLE_1,
                 Long.toString(CachedAppOptimizer.DEFAULT_COMPACT_THROTTLE_1 + 1), false);
         DeviceConfig.setProperty(DeviceConfig.NAMESPACE_ACTIVITY_MANAGER,
@@ -400,12 +406,6 @@
         DeviceConfig.setProperty(DeviceConfig.NAMESPACE_ACTIVITY_MANAGER,
                 CachedAppOptimizer.KEY_COMPACT_THROTTLE_6,
                 Long.toString(CachedAppOptimizer.DEFAULT_COMPACT_THROTTLE_6 + 1), false);
-        DeviceConfig.setProperty(DeviceConfig.NAMESPACE_ACTIVITY_MANAGER,
-                CachedAppOptimizer.KEY_COMPACT_THROTTLE_MIN_OOM_ADJ,
-                Long.toString(CachedAppOptimizer.DEFAULT_COMPACT_THROTTLE_MIN_OOM_ADJ + 1), false);
-        DeviceConfig.setProperty(DeviceConfig.NAMESPACE_ACTIVITY_MANAGER,
-                CachedAppOptimizer.KEY_COMPACT_THROTTLE_MAX_OOM_ADJ,
-                Long.toString(CachedAppOptimizer.DEFAULT_COMPACT_THROTTLE_MAX_OOM_ADJ - 1), false);
         assertThat(mCountDown.await(7, TimeUnit.SECONDS)).isTrue();
 
         // Then those flags values are reflected in the compactor.
diff --git a/services/tests/servicestests/Android.bp b/services/tests/servicestests/Android.bp
index 009ce88..d702cae 100644
--- a/services/tests/servicestests/Android.bp
+++ b/services/tests/servicestests/Android.bp
@@ -33,9 +33,6 @@
         "test-apps/DisplayManagerTestApp/src/**/*.java",
     ],
 
-    kotlincflags: [
-        "-Werror",
-    ],
     static_libs: [
         "a11ychecker",
         "aatf",
diff --git a/services/tests/wmtests/src/com/android/server/wm/DisplayContentTests.java b/services/tests/wmtests/src/com/android/server/wm/DisplayContentTests.java
index 02ed67b..cfd501a 100644
--- a/services/tests/wmtests/src/com/android/server/wm/DisplayContentTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/DisplayContentTests.java
@@ -83,6 +83,7 @@
 import static com.android.window.flags.Flags.FLAG_ENABLE_CAMERA_COMPAT_FOR_DESKTOP_WINDOWING;
 import static com.android.window.flags.Flags.FLAG_ENABLE_DESKTOP_WINDOWING_MODE;
 import static com.android.server.display.feature.flags.Flags.FLAG_ENABLE_DISPLAY_CONTENT_MODE_MANAGEMENT;
+import static com.android.window.flags.Flags.FLAG_ENABLE_PERSISTING_DENSITY_SCALE_FOR_CONNECTED_DISPLAYS;
 
 import static com.google.common.truth.Truth.assertThat;
 
@@ -2872,6 +2873,74 @@
         assertFalse(dc.mWmService.mDisplayWindowSettings.shouldShowSystemDecorsLocked(dc));
     }
 
+    @EnableFlags(FLAG_ENABLE_PERSISTING_DENSITY_SCALE_FOR_CONNECTED_DISPLAYS)
+    @Test
+    public void testForcedDensityRatioSetForExternalDisplays_persistDensityScaleFlagEnabled() {
+        final DisplayInfo displayInfo = new DisplayInfo(mDisplayInfo);
+        displayInfo.displayId = DEFAULT_DISPLAY + 1;
+        displayInfo.type = Display.TYPE_EXTERNAL;
+        final DisplayContent displayContent = createNewDisplay(displayInfo);
+        final int baseWidth = 1280;
+        final int baseHeight = 720;
+        final int baseDensity = 320;
+        final float baseXDpi = 60;
+        final float baseYDpi = 60;
+
+        displayContent.mInitialDisplayWidth = baseWidth;
+        displayContent.mInitialDisplayHeight = baseHeight;
+        displayContent.mInitialDisplayDensity = baseDensity;
+        displayContent.updateBaseDisplayMetrics(baseWidth, baseHeight, baseDensity, baseXDpi,
+                baseYDpi);
+
+        final int forcedDensity = 640;
+
+        // Verify that forcing the density is honored and the size doesn't change.
+        displayContent.setForcedDensity(forcedDensity, 0 /* userId */);
+        verifySizes(displayContent, baseWidth, baseHeight, forcedDensity);
+
+        // Verify that density ratio is set correctly.
+        assertEquals((float) forcedDensity / baseDensity,
+                displayContent.mExternalDisplayForcedDensityRatio, 0.01);
+    }
+
+    @EnableFlags(FLAG_ENABLE_PERSISTING_DENSITY_SCALE_FOR_CONNECTED_DISPLAYS)
+    @Test
+    public void testForcedDensityUpdateForExternalDisplays_persistDensityScaleFlagEnabled() {
+        final DisplayInfo displayInfo = new DisplayInfo(mDisplayInfo);
+        displayInfo.displayId = DEFAULT_DISPLAY + 1;
+        displayInfo.type = Display.TYPE_EXTERNAL;
+        final DisplayContent displayContent = createNewDisplay(displayInfo);
+        final int baseWidth = 1280;
+        final int baseHeight = 720;
+        final int baseDensity = 320;
+        final float baseXDpi = 60;
+        final float baseYDpi = 60;
+
+        displayContent.mInitialDisplayWidth = baseWidth;
+        displayContent.mInitialDisplayHeight = baseHeight;
+        displayContent.mInitialDisplayDensity = baseDensity;
+        displayContent.updateBaseDisplayMetrics(baseWidth, baseHeight, baseDensity, baseXDpi,
+                baseYDpi);
+
+        final int forcedDensity = 640;
+
+        // Verify that forcing the density is honored and the size doesn't change.
+        displayContent.setForcedDensity(forcedDensity, 0 /* userId */);
+        verifySizes(displayContent, baseWidth, baseHeight, forcedDensity);
+
+        // Verify that density ratio is set correctly.
+        assertEquals((float) 2.0f,
+                displayContent.mExternalDisplayForcedDensityRatio, 0.001);
+
+
+        displayContent.mInitialDisplayDensity = 160;
+        displayContent.updateBaseDisplayMetrics(baseWidth, baseHeight, baseDensity, baseXDpi,
+                baseYDpi);
+
+        // Verify that forced density is updated based on the ratio.
+        assertEquals(320, displayContent.mBaseDisplayDensity);
+    }
+
     @EnableFlags(FLAG_ENABLE_DISPLAY_CONTENT_MODE_MANAGEMENT)
     @Test
     public void testSetShouldShowSystemDecorations_nonDefaultNonPrivateDisplay() {
diff --git a/services/tests/wmtests/src/com/android/server/wm/DisplayWindowSettingsTests.java b/services/tests/wmtests/src/com/android/server/wm/DisplayWindowSettingsTests.java
index b1cad51..a57577a96 100644
--- a/services/tests/wmtests/src/com/android/server/wm/DisplayWindowSettingsTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/DisplayWindowSettingsTests.java
@@ -42,6 +42,7 @@
 import android.annotation.NonNull;
 import android.app.WindowConfiguration;
 import android.content.ContentResolver;
+import android.platform.test.annotations.EnableFlags;
 import android.platform.test.annotations.Presubmit;
 import android.provider.Settings;
 import android.view.Display;
@@ -53,6 +54,7 @@
 import com.android.server.LocalServices;
 import com.android.server.policy.WindowManagerPolicy;
 import com.android.server.wm.DisplayWindowSettings.SettingsProvider.SettingsEntry;
+import com.android.window.flags.Flags;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -272,6 +274,23 @@
                 mSecondaryDisplay.mBaseDisplayDensity);
     }
 
+    @EnableFlags(Flags.FLAG_ENABLE_PERSISTING_DENSITY_SCALE_FOR_CONNECTED_DISPLAYS)
+    @Test
+    public void testSetForcedDensityRatio() {
+        mDisplayWindowSettings.setForcedDensity(mSecondaryDisplay.getDisplayInfo(),
+                300 /* density */, 0 /* userId */);
+        mDisplayWindowSettings.setForcedDensityRatio(mSecondaryDisplay.getDisplayInfo(),
+                2.0f /* ratio */);
+        mDisplayWindowSettings.applySettingsToDisplayLocked(mSecondaryDisplay);
+
+        assertEquals(mSecondaryDisplay.mInitialDisplayDensity * 2.0f,
+                mSecondaryDisplay.mBaseDisplayDensity, 0.01);
+
+        mWm.clearForcedDisplayDensityForUser(mSecondaryDisplay.getDisplayId(), 0 /* userId */);
+        assertEquals(mSecondaryDisplay.mInitialDisplayDensity,
+                mSecondaryDisplay.mBaseDisplayDensity);
+    }
+
     @Test
     public void testSetForcedScalingMode() {
         mDisplayWindowSettings.setForcedScalingMode(mSecondaryDisplay,
diff --git a/tests/Codegen/OWNERS b/tests/Codegen/OWNERS
index da723b3..e69de29 100644
--- a/tests/Codegen/OWNERS
+++ b/tests/Codegen/OWNERS
@@ -1 +0,0 @@
-eugenesusla@google.com
\ No newline at end of file
diff --git a/tests/Input/Android.bp b/tests/Input/Android.bp
index 168141b..1f0bd61 100644
--- a/tests/Input/Android.bp
+++ b/tests/Input/Android.bp
@@ -19,9 +19,6 @@
         "src/**/*.kt",
     ],
     asset_dirs: ["assets"],
-    kotlincflags: [
-        "-Werror",
-    ],
     platform_apis: true,
     certificate: "platform",
     static_libs: [