Merge "Add bug fix flag for inset update seq" into main
diff --git a/core/java/android/view/KeyboardShortcutInfo.java b/core/java/android/view/KeyboardShortcutInfo.java
index 118b03c..3f6fd64 100644
--- a/core/java/android/view/KeyboardShortcutInfo.java
+++ b/core/java/android/view/KeyboardShortcutInfo.java
@@ -28,8 +28,8 @@
  * Information about a Keyboard Shortcut.
  */
 public final class KeyboardShortcutInfo implements Parcelable {
-    private final CharSequence mLabel;
-    private final Icon mIcon;
+    @Nullable private final CharSequence mLabel;
+    @Nullable private Icon mIcon;
     private final char mBaseCharacter;
     private final int mKeycode;
     private final int mModifiers;
@@ -116,6 +116,15 @@
     }
 
     /**
+     * Removes an icon that was previously set.
+     *
+     * @hide
+     */
+    public void clearIcon() {
+        mIcon = null;
+    }
+
+    /**
      * Returns the base keycode that, combined with the modifiers, triggers this shortcut. If the
      * base character was set instead, returns {@link KeyEvent#KEYCODE_UNKNOWN}. Valid keycodes are
      * defined as constants in {@link KeyEvent}.
@@ -165,4 +174,4 @@
             return new KeyboardShortcutInfo[size];
         }
     };
-}
\ No newline at end of file
+}
diff --git a/core/java/android/view/ViewRootImpl.java b/core/java/android/view/ViewRootImpl.java
index 5335b1a..bdada11 100644
--- a/core/java/android/view/ViewRootImpl.java
+++ b/core/java/android/view/ViewRootImpl.java
@@ -7273,7 +7273,8 @@
             if (dispatcher.isBackGestureInProgress()) {
                 return FINISH_NOT_HANDLED;
             }
-            if (topCallback instanceof OnBackAnimationCallback) {
+            if (topCallback instanceof OnBackAnimationCallback
+                    && !(topCallback instanceof ImeBackAnimationController)) {
                 final OnBackAnimationCallback animationCallback =
                         (OnBackAnimationCallback) topCallback;
                 switch (keyEvent.getAction()) {
diff --git a/core/java/android/window/SnapshotDrawerUtils.java b/core/java/android/window/SnapshotDrawerUtils.java
index 29bb32e..f928f50 100644
--- a/core/java/android/window/SnapshotDrawerUtils.java
+++ b/core/java/android/window/SnapshotDrawerUtils.java
@@ -50,7 +50,6 @@
 import android.app.ActivityThread;
 import android.content.Context;
 import android.graphics.Canvas;
-import android.graphics.Color;
 import android.graphics.GraphicBuffer;
 import android.graphics.Paint;
 import android.graphics.PixelFormat;
@@ -68,6 +67,7 @@
 import com.android.internal.R;
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.policy.DecorView;
+import com.android.window.flags.Flags;
 
 /**
  * Utils class to help draw a snapshot on a surface.
@@ -181,7 +181,8 @@
 
             // We consider nearly matched dimensions as there can be rounding errors and the user
             // won't notice very minute differences from scaling one dimension more than the other
-            boolean aspectRatioMismatch = !isAspectRatioMatch(mFrame, mSnapshotW, mSnapshotH);
+            boolean aspectRatioMismatch = !isAspectRatioMatch(mFrame, mSnapshotW, mSnapshotH)
+                    && !Flags.drawSnapshotAspectRatioMatch();
 
             // Keep a reference to it such that it doesn't get destroyed when finalized.
             SurfaceControl childSurfaceControl = new SurfaceControl.Builder(session)
@@ -382,8 +383,8 @@
         }
         final SnapshotSurface drawSurface = new SnapshotSurface(
                 rootSurface, snapshot, lp.getTitle());
-
-        final WindowManager.LayoutParams attrs = info.topOpaqueWindowLayoutParams;
+        final WindowManager.LayoutParams attrs = Flags.drawSnapshotAspectRatioMatch()
+                ? info.mainWindowLayoutParams : info.topOpaqueWindowLayoutParams;
         final ActivityManager.RunningTaskInfo runningTaskInfo = info.taskInfo;
         final ActivityManager.TaskDescription taskDescription =
                 getOrCreateTaskDescription(runningTaskInfo);
@@ -400,7 +401,8 @@
     public static WindowManager.LayoutParams createLayoutParameters(StartingWindowInfo info,
             CharSequence title, @WindowManager.LayoutParams.WindowType int windowType,
             int pixelFormat, IBinder token) {
-        final WindowManager.LayoutParams attrs = info.topOpaqueWindowLayoutParams;
+        final WindowManager.LayoutParams attrs = Flags.drawSnapshotAspectRatioMatch()
+                ? info.mainWindowLayoutParams : info.topOpaqueWindowLayoutParams;
         final WindowManager.LayoutParams mainWindowParams = info.mainWindowLayoutParams;
         final InsetsState topWindowInsetsState = info.topOpaqueWindowInsetsState;
         if (attrs == null || mainWindowParams == null || topWindowInsetsState == null) {
@@ -527,7 +529,7 @@
 
         void drawStatusBarBackground(Canvas c, @Nullable Rect alreadyDrawnFrame,
                 int statusBarHeight) {
-            if (statusBarHeight > 0 && Color.alpha(mStatusBarColor) != 0
+            if (statusBarHeight > 0 && alpha(mStatusBarColor) != 0
                     && (alreadyDrawnFrame == null || c.getWidth() > alreadyDrawnFrame.right)) {
                 final int rightInset = (int) (mSystemBarInsets.right * mScale);
                 final int left = alreadyDrawnFrame != null ? alreadyDrawnFrame.right : 0;
@@ -541,7 +543,7 @@
             getNavigationBarRect(c.getWidth(), c.getHeight(), mSystemBarInsets, navigationBarRect,
                     mScale);
             final boolean visible = isNavigationBarColorViewVisible();
-            if (visible && Color.alpha(mNavigationBarColor) != 0
+            if (visible && alpha(mNavigationBarColor) != 0
                     && !navigationBarRect.isEmpty()) {
                 c.drawRect(navigationBarRect, mNavigationBarPaint);
             }
diff --git a/core/java/android/window/flags/windowing_frontend.aconfig b/core/java/android/window/flags/windowing_frontend.aconfig
index f08f5b8..d6f65f8 100644
--- a/core/java/android/window/flags/windowing_frontend.aconfig
+++ b/core/java/android/window/flags/windowing_frontend.aconfig
@@ -171,4 +171,15 @@
     description: "Actively release task snapshot memory"
     bug: "238206323"
     is_fixed_read_only: true
+}
+
+flag {
+  name: "draw_snapshot_aspect_ratio_match"
+  namespace: "windowing_frontend"
+  description: "The aspect ratio should always match when drawing snapshot"
+  bug: "341020277"
+  is_fixed_read_only: true
+  metadata {
+    purpose: PURPOSE_BUGFIX
+  }
 }
\ No newline at end of file
diff --git a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitController.java b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitController.java
index 9aa12aa..f78e2b5 100644
--- a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitController.java
+++ b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitController.java
@@ -1322,7 +1322,9 @@
             mPresenter.expandTaskFragment(wct, container);
         } else {
             // Put activity into a new expanded container.
-            final TaskFragmentContainer newContainer = newContainer(activity, getTaskId(activity));
+            final TaskFragmentContainer newContainer =
+                    new TaskFragmentContainer.Builder(this, getTaskId(activity), activity)
+                            .setPendingAppearedActivity(activity).build();
             mPresenter.expandActivity(wct, newContainer.getTaskFragmentToken(), activity);
         }
     }
@@ -1738,9 +1740,13 @@
             // Can't find any activity in the Task that we can use as the owner activity.
             return null;
         }
-        final TaskFragmentContainer container = newContainer(null /* pendingAppearedActivity */,
-                intent, activityInTask, taskId, null /* pairedPrimaryContainer*/, overlayTag,
-                launchOptions, associateLaunchingActivity);
+        final TaskFragmentContainer container =
+                new TaskFragmentContainer.Builder(this, taskId, activityInTask)
+                        .setPendingAppearedIntent(intent)
+                        .setOverlayTag(overlayTag)
+                        .setLaunchOptions(launchOptions)
+                        .setAssociatedActivity(associateLaunchingActivity ? activityInTask : null)
+                        .build();
         final IBinder taskFragmentToken = container.getTaskFragmentToken();
         // Note that taskContainer will not exist before calling #newContainer if the container
         // is the first embedded TF in the task.
@@ -1818,74 +1824,6 @@
         return null;
     }
 
-    @GuardedBy("mLock")
-    TaskFragmentContainer newContainer(@NonNull Activity pendingAppearedActivity, int taskId) {
-        return newContainer(pendingAppearedActivity, pendingAppearedActivity, taskId);
-    }
-
-    @GuardedBy("mLock")
-    TaskFragmentContainer newContainer(@NonNull Activity pendingAppearedActivity,
-            @NonNull Activity activityInTask, int taskId) {
-        return newContainer(pendingAppearedActivity, null /* pendingAppearedIntent */,
-                activityInTask, taskId, null /* pairedPrimaryContainer */, null /* tag */,
-                null /* launchOptions */, false /* associateLaunchingActivity */);
-    }
-
-    @GuardedBy("mLock")
-    TaskFragmentContainer newContainer(@NonNull Intent pendingAppearedIntent,
-            @NonNull Activity activityInTask, int taskId) {
-        return newContainer(null /* pendingAppearedActivity */, pendingAppearedIntent,
-                activityInTask, taskId, null /* pairedPrimaryContainer */, null /* tag */,
-                null /* launchOptions */, false /* associateLaunchingActivity */);
-    }
-
-    @GuardedBy("mLock")
-    TaskFragmentContainer newContainer(@NonNull Intent pendingAppearedIntent,
-            @NonNull Activity activityInTask, int taskId,
-            @NonNull TaskFragmentContainer pairedPrimaryContainer) {
-        return newContainer(null /* pendingAppearedActivity */, pendingAppearedIntent,
-                activityInTask, taskId, pairedPrimaryContainer, null /* tag */,
-                null /* launchOptions */, false /* associateLaunchingActivity */);
-    }
-
-    /**
-     * Creates and registers a new organized container with an optional activity that will be
-     * re-parented to it in a WCT.
-     *
-     * @param pendingAppearedActivity the activity that will be reparented to the TaskFragment.
-     * @param pendingAppearedIntent   the Intent that will be started in the TaskFragment.
-     * @param activityInTask          activity in the same Task so that we can get the Task bounds
-     *                                if needed.
-     * @param taskId                  parent Task of the new TaskFragment.
-     * @param pairedContainer  the paired primary {@link TaskFragmentContainer}. When it is
-     *                                set, the new container will be added right above it.
-     * @param overlayTag              The tag for the new created overlay container. It must be
-     *                                needed if {@code isOverlay} is {@code true}. Otherwise,
-     *                                it should be {@code null}.
-     * @param launchOptions           The launch options bundle to create a container. Must be
-     *                                specified for overlay container.
-     * @param associateLaunchingActivity {@code true} to indicate this overlay container
-     *                                   should associate with launching activity.
-     */
-    @GuardedBy("mLock")
-    TaskFragmentContainer newContainer(@Nullable Activity pendingAppearedActivity,
-            @Nullable Intent pendingAppearedIntent, @NonNull Activity activityInTask, int taskId,
-            @Nullable TaskFragmentContainer pairedContainer, @Nullable String overlayTag,
-            @Nullable Bundle launchOptions, boolean associateLaunchingActivity) {
-        if (activityInTask == null) {
-            throw new IllegalArgumentException("activityInTask must not be null,");
-        }
-        if (!mTaskContainers.contains(taskId)) {
-            mTaskContainers.put(taskId, new TaskContainer(taskId, activityInTask));
-            mDividerPresenters.put(taskId, new DividerPresenter(taskId, this, mExecutor));
-        }
-        final TaskContainer taskContainer = mTaskContainers.get(taskId);
-        final TaskFragmentContainer container = new TaskFragmentContainer(pendingAppearedActivity,
-                pendingAppearedIntent, taskContainer, this, pairedContainer, overlayTag,
-                launchOptions, associateLaunchingActivity ? activityInTask : null);
-        return container;
-    }
-
     /**
      * Creates and registers a new split with the provided containers and configuration. Finishes
      * existing secondary containers if found for the given primary container.
@@ -2581,6 +2519,12 @@
         return mTaskContainers.get(taskId);
     }
 
+    @GuardedBy("mLock")
+    void addTaskContainer(int taskId, TaskContainer taskContainer) {
+        mTaskContainers.put(taskId, taskContainer);
+        mDividerPresenters.put(taskId, new DividerPresenter(taskId, this, mExecutor));
+    }
+
     Handler getHandler() {
         return mHandler;
     }
diff --git a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitPresenter.java b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitPresenter.java
index 1cb410e..eade86e 100644
--- a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitPresenter.java
+++ b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitPresenter.java
@@ -186,8 +186,9 @@
 
         // Create new empty task fragment
         final int taskId = primaryContainer.getTaskId();
-        final TaskFragmentContainer secondaryContainer = mController.newContainer(
-                secondaryIntent, primaryActivity, taskId);
+        final TaskFragmentContainer secondaryContainer =
+                new TaskFragmentContainer.Builder(mController, taskId, primaryActivity)
+                        .setPendingAppearedIntent(secondaryIntent).build();
         final Rect secondaryRelBounds = getRelBoundsForPosition(POSITION_END, taskProperties,
                 splitAttributes);
         final int windowingMode = mController.getTaskContainer(taskId)
@@ -261,7 +262,8 @@
         TaskFragmentContainer container = mController.getContainerWithActivity(activity);
         final int taskId = container != null ? container.getTaskId() : activity.getTaskId();
         if (container == null || container == containerToAvoid) {
-            container = mController.newContainer(activity, taskId);
+            container = new TaskFragmentContainer.Builder(mController, taskId, activity)
+                    .setPendingAppearedActivity(activity).build();
             final int windowingMode = mController.getTaskContainer(taskId)
                     .getWindowingModeForTaskFragment(relBounds);
             final IBinder reparentActivityToken = activity.getActivityToken();
@@ -304,15 +306,19 @@
         TaskFragmentContainer primaryContainer = mController.getContainerWithActivity(
                 launchingActivity);
         if (primaryContainer == null) {
-            primaryContainer = mController.newContainer(launchingActivity,
-                    launchingActivity.getTaskId());
+            primaryContainer = new TaskFragmentContainer.Builder(mController,
+                    launchingActivity.getTaskId(), launchingActivity)
+                    .setPendingAppearedActivity(launchingActivity).build();
         }
 
         final int taskId = primaryContainer.getTaskId();
-        final TaskFragmentContainer secondaryContainer = mController.newContainer(activityIntent,
-                launchingActivity, taskId,
-                // Pass in the primary container to make sure it is added right above the primary.
-                primaryContainer);
+        final TaskFragmentContainer secondaryContainer =
+                new TaskFragmentContainer.Builder(mController, taskId, launchingActivity)
+                        .setPendingAppearedIntent(activityIntent)
+                        // Pass in the primary container to make sure it is added right above the
+                        // primary.
+                        .setPairedPrimaryContainer(primaryContainer)
+                        .build();
         final TaskContainer taskContainer = mController.getTaskContainer(taskId);
         final int windowingMode = taskContainer.getWindowingModeForTaskFragment(
                 primaryRelBounds);
diff --git a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/TaskFragmentContainer.java b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/TaskFragmentContainer.java
index d0b6a01..7173b0c 100644
--- a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/TaskFragmentContainer.java
+++ b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/TaskFragmentContainer.java
@@ -195,20 +195,6 @@
     private boolean mLastDimOnTask;
 
     /**
-     * @see #TaskFragmentContainer(Activity, Intent, TaskContainer, SplitController,
-     * TaskFragmentContainer, String, Bundle, Activity)
-     */
-    TaskFragmentContainer(@Nullable Activity pendingAppearedActivity,
-                          @Nullable Intent pendingAppearedIntent,
-                          @NonNull TaskContainer taskContainer,
-                          @NonNull SplitController controller,
-                          @Nullable TaskFragmentContainer pairedPrimaryContainer) {
-        this(pendingAppearedActivity, pendingAppearedIntent, taskContainer,
-                controller, pairedPrimaryContainer, null /* overlayTag */,
-                null /* launchOptions */, null /* associatedActivity */);
-    }
-
-    /**
      * Creates a container with an existing activity that will be re-parented to it in a window
      * container transaction.
      * @param pairedPrimaryContainer    when it is set, the new container will be add right above it
@@ -218,7 +204,7 @@
      * @param associatedActivity        the associated activity of the overlay container. Must be
      *                                  {@code null} for a non-overlay container.
      */
-    TaskFragmentContainer(@Nullable Activity pendingAppearedActivity,
+    private TaskFragmentContainer(@Nullable Activity pendingAppearedActivity,
             @Nullable Intent pendingAppearedIntent, @NonNull TaskContainer taskContainer,
             @NonNull SplitController controller,
             @Nullable TaskFragmentContainer pairedPrimaryContainer, @Nullable String overlayTag,
@@ -232,12 +218,6 @@
         mToken = new Binder("TaskFragmentContainer");
         mTaskContainer = taskContainer;
         mOverlayTag = overlayTag;
-        if (overlayTag != null) {
-            Objects.requireNonNull(launchOptions);
-        } else if (associatedActivity != null) {
-            throw new IllegalArgumentException("Associated activity must be null for "
-                    + "non-overlay activity.");
-        }
         mAssociatedActivityToken = associatedActivity != null
                 ? associatedActivity.getActivityToken() : null;
 
@@ -1116,6 +1096,117 @@
         return sb.append("]").toString();
     }
 
+    static final class Builder {
+        @NonNull
+        private final SplitController mSplitController;
+
+        // The parent Task id of the new TaskFragment.
+        private final int mTaskId;
+
+        // The activity in the same Task so that we can get the Task bounds if needed.
+        @NonNull
+        private final Activity mActivityInTask;
+
+        // The activity that will be reparented to the TaskFragment.
+        @Nullable
+        private Activity mPendingAppearedActivity;
+
+        // The Intent that will be started in the TaskFragment.
+        @Nullable
+        private Intent mPendingAppearedIntent;
+
+        // The paired primary {@link TaskFragmentContainer}. When it is set, the new container
+        // will be added right above it.
+        @Nullable
+        private TaskFragmentContainer mPairedPrimaryContainer;
+
+        // The launch options bundle to create a container. Must be specified for overlay container.
+        @Nullable
+        private Bundle mLaunchOptions;
+
+        // The tag for the new created overlay container. This is required when creating an
+        // overlay container.
+        @Nullable
+        private String mOverlayTag;
+
+        // The associated activity of the overlay container. Must be {@code null} for a
+        // non-overlay container.
+        @Nullable
+        private Activity mAssociatedActivity;
+
+        Builder(@NonNull SplitController splitController, int taskId,
+                @Nullable Activity activityInTask) {
+            if (taskId <= 0) {
+                throw new IllegalArgumentException("taskId is invalid, " + taskId);
+            }
+
+            mSplitController = splitController;
+            mTaskId = taskId;
+            mActivityInTask = activityInTask;
+        }
+
+        @NonNull
+        Builder setPendingAppearedActivity(@Nullable Activity pendingAppearedActivity) {
+            mPendingAppearedActivity = pendingAppearedActivity;
+            return this;
+        }
+
+        @NonNull
+        Builder setPendingAppearedIntent(@Nullable Intent pendingAppearedIntent) {
+            mPendingAppearedIntent = pendingAppearedIntent;
+            return this;
+        }
+
+        @NonNull
+        Builder setPairedPrimaryContainer(@Nullable TaskFragmentContainer pairedPrimaryContainer) {
+            mPairedPrimaryContainer = pairedPrimaryContainer;
+            return this;
+        }
+
+        @NonNull
+        Builder setLaunchOptions(@Nullable Bundle launchOptions) {
+            mLaunchOptions = launchOptions;
+            return this;
+        }
+
+        @NonNull
+        Builder setOverlayTag(@Nullable String overlayTag) {
+            mOverlayTag = overlayTag;
+            return this;
+        }
+
+        @NonNull
+        Builder setAssociatedActivity(@Nullable Activity associatedActivity) {
+            mAssociatedActivity = associatedActivity;
+            return this;
+        }
+
+        @NonNull
+        TaskFragmentContainer build() {
+            if (mOverlayTag != null) {
+                Objects.requireNonNull(mLaunchOptions);
+            } else if (mAssociatedActivity != null) {
+                throw new IllegalArgumentException("Associated activity must be null for "
+                        + "non-overlay activity.");
+            }
+
+            TaskContainer taskContainer = mSplitController.getTaskContainer(mTaskId);
+            if (taskContainer == null && mActivityInTask == null) {
+                throw new IllegalArgumentException("mActivityInTask must be set.");
+            }
+
+            if (taskContainer == null) {
+                // Adding a TaskContainer if no existed one.
+                taskContainer = new TaskContainer(mTaskId, mActivityInTask);
+                mSplitController.addTaskContainer(mTaskId, taskContainer);
+            }
+
+            return new TaskFragmentContainer(mPendingAppearedActivity, mPendingAppearedIntent,
+                    taskContainer, mSplitController, mPairedPrimaryContainer, mOverlayTag,
+                    mLaunchOptions, mAssociatedActivity);
+        }
+    }
+
     static class OverlayContainerRestoreParams {
         /** The token of the overlay container */
         @NonNull
diff --git a/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/EmbeddingTestUtils.java b/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/EmbeddingTestUtils.java
index a069ac7..d649c6d 100644
--- a/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/EmbeddingTestUtils.java
+++ b/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/EmbeddingTestUtils.java
@@ -248,4 +248,17 @@
         return new SplitPlaceholderRule.Builder(placeholderIntent, activityPredicate,
                 intentPredicate, windowMetricsPredicate);
     }
+
+    @NonNull
+    static TaskFragmentContainer createTfContainer(
+            @NonNull SplitController splitController, @NonNull Activity activity) {
+        return createTfContainer(splitController, TASK_ID, activity);
+    }
+
+    @NonNull
+    static TaskFragmentContainer createTfContainer(
+            @NonNull SplitController splitController, int taskId, @NonNull Activity activity) {
+        return new TaskFragmentContainer.Builder(splitController, taskId, activity)
+                .setPendingAppearedActivity(activity).build();
+    }
 }
diff --git a/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/JetpackTaskFragmentOrganizerTest.java b/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/JetpackTaskFragmentOrganizerTest.java
index 76e6a0f..7b473b0 100644
--- a/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/JetpackTaskFragmentOrganizerTest.java
+++ b/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/JetpackTaskFragmentOrganizerTest.java
@@ -25,6 +25,7 @@
 
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
+import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.mock;
 
@@ -105,8 +106,11 @@
     @Test
     public void testExpandTaskFragment() {
         final TaskContainer taskContainer = createTestTaskContainer();
-        final TaskFragmentContainer container = new TaskFragmentContainer(null /* activity */,
-                new Intent(), taskContainer, mSplitController, null /* pairedPrimaryContainer */);
+        doReturn(taskContainer).when(mSplitController).getTaskContainer(anyInt());
+        final TaskFragmentContainer container =  new TaskFragmentContainer.Builder(mSplitController,
+                taskContainer.getTaskId(), null /* activityInTask */)
+                .setPendingAppearedIntent(new Intent())
+                .build();
         final TaskFragmentInfo info = createMockInfo(container);
         mOrganizer.mFragmentInfos.put(container.getTaskFragmentToken(), info);
         container.setInfo(mTransaction, info);
diff --git a/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/OverlayPresentationTest.java b/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/OverlayPresentationTest.java
index 86b7e88..0972d40 100644
--- a/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/OverlayPresentationTest.java
+++ b/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/OverlayPresentationTest.java
@@ -29,6 +29,7 @@
 import static androidx.window.extensions.embedding.EmbeddingTestUtils.createSplitPairRuleBuilder;
 import static androidx.window.extensions.embedding.EmbeddingTestUtils.createSplitPlaceholderRuleBuilder;
 import static androidx.window.extensions.embedding.EmbeddingTestUtils.createSplitRule;
+import static androidx.window.extensions.embedding.EmbeddingTestUtils.createTfContainer;
 import static androidx.window.extensions.embedding.SplitPresenter.sanitizeBounds;
 import static androidx.window.extensions.embedding.WindowAttributes.DIM_AREA_ON_TASK;
 
@@ -530,8 +531,8 @@
 
     @Test
     public void testUpdateActivityStackAttributes_nullContainer_earlyReturn() {
-        final TaskFragmentContainer container = mSplitController.newContainer(mActivity,
-                mActivity.getTaskId());
+        final TaskFragmentContainer container = createTfContainer(mSplitController,
+                mActivity.getTaskId(), mActivity);
         mSplitController.updateActivityStackAttributes(
                 ActivityStack.Token.createFromBinder(container.getTaskFragmentToken()),
                 new ActivityStackAttributes.Builder().build());
@@ -837,8 +838,9 @@
         final Intent intent = new Intent();
         final IBinder fillTaskActivityToken = new Binder();
         final IBinder lastOverlayToken = new Binder();
-        final TaskFragmentContainer overlayContainer = mSplitController.newContainer(intent,
-                mActivity, TASK_ID);
+        final TaskFragmentContainer overlayContainer =
+                new TaskFragmentContainer.Builder(mSplitController, TASK_ID, mActivity)
+                        .setPendingAppearedIntent(intent).build();
         final TaskFragmentContainer.OverlayContainerRestoreParams params = mock(
                 TaskFragmentContainer.OverlayContainerRestoreParams.class);
         doReturn(params).when(mSplitController).getOverlayContainerRestoreParams(any(), any());
@@ -884,8 +886,8 @@
     @NonNull
     private TaskFragmentContainer createMockTaskFragmentContainer(
             @NonNull Activity activity, boolean isVisible) {
-        final TaskFragmentContainer container = mSplitController.newContainer(activity,
-                activity.getTaskId());
+        final TaskFragmentContainer container = createTfContainer(mSplitController,
+                activity.getTaskId(), activity);
         setupTaskFragmentInfo(container, activity, isVisible);
         return container;
     }
@@ -918,10 +920,13 @@
             @Nullable Activity launchingActivity) {
         final Activity activity = launchingActivity != null
                 ? launchingActivity : createMockActivity();
-        TaskFragmentContainer overlayContainer = mSplitController.newContainer(
-                null /* pendingAppearedActivity */, mIntent, activity, taskId,
-                null /* pairedPrimaryContainer */, tag, Bundle.EMPTY,
-                associateLaunchingActivity);
+        TaskFragmentContainer overlayContainer =
+                new TaskFragmentContainer.Builder(mSplitController, taskId, activity)
+                        .setPendingAppearedIntent(mIntent)
+                        .setOverlayTag(tag)
+                        .setLaunchOptions(Bundle.EMPTY)
+                        .setAssociatedActivity(associateLaunchingActivity ? activity : null)
+                        .build();
         setupTaskFragmentInfo(overlayContainer, createMockActivity(), isVisible);
         return overlayContainer;
     }
diff --git a/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/SplitControllerTest.java b/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/SplitControllerTest.java
index 35353db..640b1fc 100644
--- a/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/SplitControllerTest.java
+++ b/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/SplitControllerTest.java
@@ -41,6 +41,7 @@
 import static androidx.window.extensions.embedding.EmbeddingTestUtils.createSplitPlaceholderRuleBuilder;
 import static androidx.window.extensions.embedding.EmbeddingTestUtils.createSplitRule;
 import static androidx.window.extensions.embedding.EmbeddingTestUtils.createTestTaskContainer;
+import static androidx.window.extensions.embedding.EmbeddingTestUtils.createTfContainer;
 import static androidx.window.extensions.embedding.EmbeddingTestUtils.getSplitBounds;
 import static androidx.window.extensions.embedding.SplitRule.FINISH_ALWAYS;
 
@@ -59,7 +60,6 @@
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyBoolean;
 import static org.mockito.ArgumentMatchers.anyInt;
-import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.ArgumentMatchers.isNull;
 import static org.mockito.Mockito.clearInvocations;
@@ -198,7 +198,7 @@
 
     @Test
     public void testOnTaskFragmentVanished() {
-        final TaskFragmentContainer tf = mSplitController.newContainer(mActivity, TASK_ID);
+        final TaskFragmentContainer tf = createTfContainer(mSplitController, mActivity);
         doReturn(tf.getTaskFragmentToken()).when(mInfo).getFragmentToken();
 
         // The TaskFragment has been removed in the server, we only need to cleanup the reference.
@@ -213,7 +213,7 @@
     public void testOnTaskFragmentAppearEmptyTimeout() {
         // Setup to make sure a transaction record is started.
         mTransactionManager.startNewTransaction();
-        final TaskFragmentContainer tf = mSplitController.newContainer(mActivity, TASK_ID);
+        final TaskFragmentContainer tf = createTfContainer(mSplitController, mActivity);
         doCallRealMethod().when(mSplitController).onTaskFragmentAppearEmptyTimeout(any(), any());
         mSplitController.onTaskFragmentAppearEmptyTimeout(mTransaction, tf);
 
@@ -224,7 +224,7 @@
     @Test
     public void testOnActivityDestroyed() {
         doReturn(new Binder()).when(mActivity).getActivityToken();
-        final TaskFragmentContainer tf = mSplitController.newContainer(mActivity, TASK_ID);
+        final TaskFragmentContainer tf = createTfContainer(mSplitController, mActivity);
 
         assertTrue(tf.hasActivity(mActivity.getActivityToken()));
 
@@ -245,12 +245,9 @@
     public void testNewContainer() {
         // Must pass in a valid activity.
         assertThrows(IllegalArgumentException.class, () ->
-                mSplitController.newContainer(null /* activity */, TASK_ID));
-        assertThrows(IllegalArgumentException.class, () ->
-                mSplitController.newContainer(mActivity, null /* launchingActivity */, TASK_ID));
+                createTfContainer(mSplitController, null /* activity */));
 
-        final TaskFragmentContainer tf = mSplitController.newContainer(mActivity, mActivity,
-                TASK_ID);
+        final TaskFragmentContainer tf = createTfContainer(mSplitController, mActivity);
         final TaskContainer taskContainer = mSplitController.getTaskContainer(TASK_ID);
 
         assertNotNull(tf);
@@ -263,7 +260,7 @@
     public void testUpdateContainer() {
         // Make SplitController#launchPlaceholderIfNecessary(TaskFragmentContainer) return true
         // and verify if shouldContainerBeExpanded() not called.
-        final TaskFragmentContainer tf = mSplitController.newContainer(mActivity, TASK_ID);
+        final TaskFragmentContainer tf = createTfContainer(mSplitController, mActivity);
         spyOn(tf);
         doReturn(mActivity).when(tf).getTopNonFinishingActivity();
         doReturn(true).when(tf).isEmpty();
@@ -369,8 +366,12 @@
     public void testOnStartActivityResultError() {
         final Intent intent = new Intent();
         final TaskContainer taskContainer = createTestTaskContainer();
-        final TaskFragmentContainer container = new TaskFragmentContainer(null /* activity */,
-                intent, taskContainer, mSplitController, null /* pairedPrimaryContainer */);
+        final int taskId = taskContainer.getTaskId();
+        mSplitController.addTaskContainer(taskId, taskContainer);
+        final TaskFragmentContainer container = new TaskFragmentContainer.Builder(mSplitController,
+                taskId, null /* activityInTask */)
+                .setPendingAppearedIntent(intent)
+                .build();
         final SplitController.ActivityStartMonitor monitor =
                 mSplitController.getActivityStartMonitor();
 
@@ -410,7 +411,8 @@
     @Test
     public void testOnActivityReparentedToTask_diffProcess() {
         // Create an empty TaskFragment to initialize for the Task.
-        mSplitController.newContainer(new Intent(), mActivity, TASK_ID);
+        new TaskFragmentContainer.Builder(mSplitController, TASK_ID, mActivity)
+                .setPendingAppearedIntent(new Intent()).build();
         final IBinder activityToken = new Binder();
         final Intent intent = new Intent();
 
@@ -595,8 +597,9 @@
         verify(mTransaction, never()).reparentActivityToTaskFragment(any(), any());
 
         // Place in the top container if there is no other rule matched.
-        final TaskFragmentContainer topContainer = mSplitController
-                .newContainer(new Intent(), mActivity, TASK_ID);
+        final TaskFragmentContainer topContainer =
+                new TaskFragmentContainer.Builder(mSplitController, TASK_ID, mActivity)
+                        .setPendingAppearedIntent(new Intent()).build();
         mSplitController.placeActivityInTopContainer(mTransaction, mActivity);
 
         verify(mTransaction).reparentActivityToTaskFragment(topContainer.getTaskFragmentToken(),
@@ -604,7 +607,7 @@
 
         // Not reparent if activity is in a TaskFragment.
         clearInvocations(mTransaction);
-        mSplitController.newContainer(mActivity, TASK_ID);
+        createTfContainer(mSplitController, mActivity);
         mSplitController.placeActivityInTopContainer(mTransaction, mActivity);
 
         verify(mTransaction, never()).reparentActivityToTaskFragment(any(), any());
@@ -616,8 +619,7 @@
                 false /* isOnReparent */);
 
         assertFalse(result);
-        verify(mSplitController, never()).newContainer(any(), any(), any(), anyInt(), any(),
-                anyString(), any(), anyBoolean());
+        verify(mSplitController, never()).addTaskContainer(anyInt(), any());
     }
 
     @Test
@@ -632,7 +634,6 @@
 
         assertTrue(result);
         assertNotNull(container);
-        verify(mSplitController).newContainer(mActivity, TASK_ID);
         verify(mSplitPresenter).expandActivity(mTransaction, container.getTaskFragmentToken(),
                 mActivity);
     }
@@ -642,7 +643,7 @@
         setupExpandRule(mActivity);
 
         // When the activity is not in any TaskFragment, create a new expanded TaskFragment for it.
-        final TaskFragmentContainer container = mSplitController.newContainer(mActivity, TASK_ID);
+        final TaskFragmentContainer container = createTfContainer(mSplitController, mActivity);
         final boolean result = mSplitController.resolveActivityToContainer(mTransaction, mActivity,
                 false /* isOnReparent */);
 
@@ -692,8 +693,8 @@
 
         // Don't launch placeholder if the activity is not in the topmost active TaskFragment.
         final Activity activity = createMockActivity();
-        mSplitController.newContainer(mActivity, TASK_ID);
-        mSplitController.newContainer(activity, TASK_ID);
+        createTfContainer(mSplitController, mActivity);
+        createTfContainer(mSplitController, activity);
         final boolean result = mSplitController.resolveActivityToContainer(mTransaction, mActivity,
                 false /* isOnReparent */);
 
@@ -711,7 +712,7 @@
                 (SplitPlaceholderRule) mSplitController.getSplitRules().get(0);
 
         // Launch placeholder if the activity is in the topmost expanded TaskFragment.
-        mSplitController.newContainer(mActivity, TASK_ID);
+        createTfContainer(mSplitController, mActivity);
         final boolean result = mSplitController.resolveActivityToContainer(mTransaction, mActivity,
                 false /* isOnReparent */);
 
@@ -763,10 +764,11 @@
         final SplitPairRule splitRule = (SplitPairRule) mSplitController.getSplitRules().get(0);
 
         // Activity is already in primary split, no need to create new split.
-        final TaskFragmentContainer primaryContainer = mSplitController.newContainer(mActivity,
-                TASK_ID);
-        final TaskFragmentContainer secondaryContainer = mSplitController.newContainer(
-                secondaryIntent, mActivity, TASK_ID);
+        final TaskFragmentContainer primaryContainer =
+                createTfContainer(mSplitController, mActivity);
+        final TaskFragmentContainer secondaryContainer =
+                new TaskFragmentContainer.Builder(mSplitController, TASK_ID, mActivity)
+                        .setPendingAppearedIntent(secondaryIntent).build();
         mSplitController.registerSplit(
                 mTransaction,
                 primaryContainer,
@@ -779,8 +781,6 @@
                 false /* isOnReparent */);
 
         assertTrue(result);
-        verify(mSplitController, never()).newContainer(any(), any(), any(), anyInt(), any(),
-                anyString(), any(), anyBoolean());
         verify(mSplitController, never()).registerSplit(any(), any(), any(), any(), any(), any());
     }
 
@@ -792,10 +792,11 @@
 
         // The new launched activity is in primary split, but there is no rule for it to split with
         // the secondary, so return false.
-        final TaskFragmentContainer primaryContainer = mSplitController.newContainer(mActivity,
-                TASK_ID);
-        final TaskFragmentContainer secondaryContainer = mSplitController.newContainer(
-                secondaryIntent, mActivity, TASK_ID);
+        final TaskFragmentContainer primaryContainer =
+                createTfContainer(mSplitController, mActivity);
+        final TaskFragmentContainer secondaryContainer =
+                new TaskFragmentContainer.Builder(mSplitController, TASK_ID, mActivity)
+                        .setPendingAppearedIntent(secondaryIntent).build();
         mSplitController.registerSplit(
                 mTransaction,
                 primaryContainer,
@@ -822,8 +823,6 @@
                 false /* isOnReparent */);
 
         assertTrue(result);
-        verify(mSplitController, never()).newContainer(any(), any(), any(), anyInt(), any(),
-                anyString(), any(), anyBoolean());
         verify(mSplitController, never()).registerSplit(any(), any(), any(), any(), any(), any());
     }
 
@@ -852,10 +851,10 @@
         doReturn(PLACEHOLDER_INTENT).when(mActivity).getIntent();
 
         // Activity is a placeholder.
-        final TaskFragmentContainer primaryContainer = mSplitController.newContainer(
-                primaryActivity, TASK_ID);
-        final TaskFragmentContainer secondaryContainer = mSplitController.newContainer(mActivity,
-                TASK_ID);
+        final TaskFragmentContainer primaryContainer =
+                createTfContainer(mSplitController, primaryActivity);
+        final TaskFragmentContainer secondaryContainer =
+                createTfContainer(mSplitController, mActivity);
         mSplitController.registerSplit(
                 mTransaction,
                 primaryContainer,
@@ -874,8 +873,7 @@
         final Activity activityBelow = createMockActivity();
         setupSplitRule(activityBelow, mActivity);
 
-        final TaskFragmentContainer container = mSplitController.newContainer(activityBelow,
-                TASK_ID);
+        final TaskFragmentContainer container = createTfContainer(mSplitController, activityBelow);
         container.addPendingAppearedActivity(mActivity);
         final boolean result = mSplitController.resolveActivityToContainer(mTransaction, mActivity,
                 false /* isOnReparent */);
@@ -890,8 +888,7 @@
         setupSplitRule(mActivity, activityBelow);
 
         // Disallow to split as primary.
-        final TaskFragmentContainer container = mSplitController.newContainer(activityBelow,
-                TASK_ID);
+        final TaskFragmentContainer container = createTfContainer(mSplitController, activityBelow);
         container.addPendingAppearedActivity(mActivity);
         boolean result = mSplitController.resolveActivityToContainer(mTransaction, mActivity,
                 false /* isOnReparent */);
@@ -961,8 +958,7 @@
 
         doReturn(createActivityInfoWithMinDimensions()).when(mActivity).getActivityInfo();
 
-        final TaskFragmentContainer container = mSplitController.newContainer(activityBelow,
-                TASK_ID);
+        final TaskFragmentContainer container = createTfContainer(mSplitController, activityBelow);
         container.addPendingAppearedActivity(mActivity);
 
         // Allow to split as primary.
@@ -980,8 +976,7 @@
 
         doReturn(createActivityInfoWithMinDimensions()).when(mActivity).getActivityInfo();
 
-        final TaskFragmentContainer container = mSplitController.newContainer(activityBelow,
-                TASK_ID);
+        final TaskFragmentContainer container = createTfContainer(mSplitController, activityBelow);
         container.addPendingAppearedActivity(mActivity);
 
         boolean result = mSplitController.resolveActivityToContainer(mTransaction, mActivity,
@@ -1044,8 +1039,8 @@
     public void testResolveActivityToContainer_skipIfNonTopOrPinned() {
         final TaskFragmentContainer container = createMockTaskFragmentContainer(mActivity);
         final Activity pinnedActivity = createMockActivity();
-        final TaskFragmentContainer topContainer = mSplitController.newContainer(pinnedActivity,
-                TASK_ID);
+        final TaskFragmentContainer topContainer =
+                createTfContainer(mSplitController, pinnedActivity);
         final TaskContainer taskContainer = container.getTaskContainer();
         spyOn(taskContainer);
         doReturn(container).when(taskContainer).getTopNonFinishingTaskFragmentContainer(false);
@@ -1351,7 +1346,7 @@
         // Launch placeholder for activity in top TaskFragment.
         setupPlaceholderRule(mActivity);
         mTransactionManager.startNewTransaction();
-        final TaskFragmentContainer container = mSplitController.newContainer(mActivity, TASK_ID);
+        final TaskFragmentContainer container = createTfContainer(mSplitController, mActivity);
         mSplitController.launchPlaceholderIfNecessary(mTransaction, mActivity,
                 true /* isOnCreated */);
 
@@ -1365,9 +1360,10 @@
         // Do not launch placeholder for invisible activity below the top TaskFragment.
         setupPlaceholderRule(mActivity);
         mTransactionManager.startNewTransaction();
-        final TaskFragmentContainer bottomTf = mSplitController.newContainer(mActivity, TASK_ID);
-        final TaskFragmentContainer topTf = mSplitController.newContainer(new Intent(), mActivity,
-                TASK_ID);
+        final TaskFragmentContainer bottomTf = createTfContainer(mSplitController, mActivity);
+        final TaskFragmentContainer topTf =
+                new TaskFragmentContainer.Builder(mSplitController, TASK_ID, mActivity)
+                        .setPendingAppearedIntent(new Intent()).build();
         bottomTf.setInfo(mTransaction, createMockTaskFragmentInfo(bottomTf, mActivity,
                 false /* isVisible */));
         topTf.setInfo(mTransaction, createMockTaskFragmentInfo(topTf, createMockActivity()));
@@ -1383,9 +1379,10 @@
         // Launch placeholder for visible activity below the top TaskFragment.
         setupPlaceholderRule(mActivity);
         mTransactionManager.startNewTransaction();
-        final TaskFragmentContainer bottomTf = mSplitController.newContainer(mActivity, TASK_ID);
-        final TaskFragmentContainer topTf = mSplitController.newContainer(new Intent(), mActivity,
-                TASK_ID);
+        final TaskFragmentContainer bottomTf = createTfContainer(mSplitController, mActivity);
+        final TaskFragmentContainer topTf =
+                new TaskFragmentContainer.Builder(mSplitController, TASK_ID, mActivity)
+                        .setPendingAppearedIntent(new Intent()).build();
         bottomTf.setInfo(mTransaction, createMockTaskFragmentInfo(bottomTf, mActivity,
                 true /* isVisible */));
         topTf.setInfo(mTransaction, createMockTaskFragmentInfo(topTf, createMockActivity()));
@@ -1412,7 +1409,7 @@
 
     @Test
     public void testFinishActivityStacks_finishSingleActivityStack() {
-        TaskFragmentContainer tf = mSplitController.newContainer(mActivity, TASK_ID);
+        TaskFragmentContainer tf = createTfContainer(mSplitController, mActivity);
         tf.setInfo(mTransaction, createMockTaskFragmentInfo(tf, mActivity));
 
         final TaskContainer taskContainer = mSplitController.mTaskContainers.get(TASK_ID);
@@ -1426,8 +1423,8 @@
 
     @Test
     public void testFinishActivityStacks_finishActivityStacksInOrder() {
-        TaskFragmentContainer bottomTf = mSplitController.newContainer(mActivity, TASK_ID);
-        TaskFragmentContainer topTf = mSplitController.newContainer(mActivity, TASK_ID);
+        TaskFragmentContainer bottomTf = createTfContainer(mSplitController, mActivity);
+        TaskFragmentContainer topTf = createTfContainer(mSplitController, mActivity);
         bottomTf.setInfo(mTransaction, createMockTaskFragmentInfo(bottomTf, mActivity));
         topTf.setInfo(mTransaction, createMockTaskFragmentInfo(topTf, createMockActivity()));
 
@@ -1687,8 +1684,8 @@
 
     /** Creates a mock TaskFragment that has been registered and appeared in the organizer. */
     private TaskFragmentContainer createMockTaskFragmentContainer(@NonNull Activity activity) {
-        final TaskFragmentContainer container = mSplitController.newContainer(activity,
-                activity.getTaskId());
+        final TaskFragmentContainer container = createTfContainer(mSplitController,
+                activity.getTaskId(), activity);
         setupTaskFragmentInfo(container, activity);
         return container;
     }
diff --git a/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/SplitPresenterTest.java b/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/SplitPresenterTest.java
index 3fbce9ec..816e2da 100644
--- a/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/SplitPresenterTest.java
+++ b/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/SplitPresenterTest.java
@@ -31,6 +31,7 @@
 import static androidx.window.extensions.embedding.EmbeddingTestUtils.createMockTaskFragmentInfo;
 import static androidx.window.extensions.embedding.EmbeddingTestUtils.createSplitPairRuleBuilder;
 import static androidx.window.extensions.embedding.EmbeddingTestUtils.createSplitRule;
+import static androidx.window.extensions.embedding.EmbeddingTestUtils.createTfContainer;
 import static androidx.window.extensions.embedding.EmbeddingTestUtils.createWindowLayoutInfo;
 import static androidx.window.extensions.embedding.EmbeddingTestUtils.getSplitBounds;
 import static androidx.window.extensions.embedding.SplitPresenter.EXPAND_CONTAINERS_ATTRIBUTES;
@@ -139,7 +140,7 @@
 
     @Test
     public void testCreateTaskFragment() {
-        final TaskFragmentContainer container = mController.newContainer(mActivity, TASK_ID);
+        final TaskFragmentContainer container = createTfContainer(mController, mActivity);
         mPresenter.createTaskFragment(mTransaction, container.getTaskFragmentToken(),
                 mActivity.getActivityToken(), TASK_BOUNDS, WINDOWING_MODE_MULTI_WINDOW);
 
@@ -150,7 +151,7 @@
 
     @Test
     public void testResizeTaskFragment() {
-        final TaskFragmentContainer container = mController.newContainer(mActivity, TASK_ID);
+        final TaskFragmentContainer container = createTfContainer(mController, mActivity);
         mPresenter.mFragmentInfos.put(container.getTaskFragmentToken(), mTaskFragmentInfo);
         mPresenter.resizeTaskFragment(mTransaction, container.getTaskFragmentToken(), TASK_BOUNDS);
 
@@ -166,7 +167,7 @@
 
     @Test
     public void testUpdateWindowingMode() {
-        final TaskFragmentContainer container = mController.newContainer(mActivity, TASK_ID);
+        final TaskFragmentContainer container = createTfContainer(mController, mActivity);
         mPresenter.mFragmentInfos.put(container.getTaskFragmentToken(), mTaskFragmentInfo);
         mPresenter.updateWindowingMode(mTransaction, container.getTaskFragmentToken(),
                 WINDOWING_MODE_MULTI_WINDOW);
@@ -184,8 +185,8 @@
 
     @Test
     public void testSetAdjacentTaskFragments() {
-        final TaskFragmentContainer container0 = mController.newContainer(mActivity, TASK_ID);
-        final TaskFragmentContainer container1 = mController.newContainer(mActivity, TASK_ID);
+        final TaskFragmentContainer container0 = createTfContainer(mController, mActivity);
+        final TaskFragmentContainer container1 = createTfContainer(mController, mActivity);
 
         mPresenter.setAdjacentTaskFragments(mTransaction, container0.getTaskFragmentToken(),
                 container1.getTaskFragmentToken(), null /* adjacentParams */);
@@ -202,8 +203,8 @@
 
     @Test
     public void testClearAdjacentTaskFragments() {
-        final TaskFragmentContainer container0 = mController.newContainer(mActivity, TASK_ID);
-        final TaskFragmentContainer container1 = mController.newContainer(mActivity, TASK_ID);
+        final TaskFragmentContainer container0 = createTfContainer(mController, mActivity);
+        final TaskFragmentContainer container1 = createTfContainer(mController, mActivity);
 
         // No request to clear as it is not set by default.
         mPresenter.clearAdjacentTaskFragments(mTransaction, container0.getTaskFragmentToken());
@@ -224,8 +225,8 @@
 
     @Test
     public void testSetCompanionTaskFragment() {
-        final TaskFragmentContainer container0 = mController.newContainer(mActivity, TASK_ID);
-        final TaskFragmentContainer container1 = mController.newContainer(mActivity, TASK_ID);
+        final TaskFragmentContainer container0 = createTfContainer(mController, mActivity);
+        final TaskFragmentContainer container1 = createTfContainer(mController, mActivity);
 
         mPresenter.setCompanionTaskFragment(mTransaction, container0.getTaskFragmentToken(),
                 container1.getTaskFragmentToken());
@@ -242,7 +243,7 @@
 
     @Test
     public void testSetTaskFragmentDimOnTask() {
-        final TaskFragmentContainer container = mController.newContainer(mActivity, TASK_ID);
+        final TaskFragmentContainer container = createTfContainer(mController, mActivity);
 
         mPresenter.setTaskFragmentDimOnTask(mTransaction, container.getTaskFragmentToken(), true);
         verify(mTransaction).addTaskFragmentOperation(eq(container.getTaskFragmentToken()), any());
@@ -255,7 +256,7 @@
 
     @Test
     public void testUpdateAnimationParams() {
-        final TaskFragmentContainer container = mController.newContainer(mActivity, TASK_ID);
+        final TaskFragmentContainer container = createTfContainer(mController, mActivity);
 
         // Verify the default.
         assertTrue(container.areLastRequestedAnimationParamsEqual(
@@ -287,7 +288,7 @@
 
     @Test
     public void testSetTaskFragmentPinned() {
-        final TaskFragmentContainer container = mController.newContainer(mActivity, TASK_ID);
+        final TaskFragmentContainer container = createTfContainer(mController, mActivity);
 
         // Verify the default.
         assertFalse(container.isPinned());
@@ -667,8 +668,8 @@
     public void testExpandSplitContainerIfNeeded() {
         Activity secondaryActivity = createMockActivity();
         SplitRule splitRule = createSplitRule(mActivity, secondaryActivity);
-        TaskFragmentContainer primaryTf = mController.newContainer(mActivity, TASK_ID);
-        TaskFragmentContainer secondaryTf = mController.newContainer(secondaryActivity, TASK_ID);
+        TaskFragmentContainer primaryTf = createTfContainer(mController, mActivity);
+        TaskFragmentContainer secondaryTf = createTfContainer(mController, secondaryActivity);
         SplitContainer splitContainer = new SplitContainer(primaryTf, secondaryActivity,
                 secondaryTf, splitRule, SPLIT_ATTRIBUTES);
 
@@ -710,8 +711,8 @@
     @Test
     public void testCreateNewSplitContainer_secondaryAbovePrimary() {
         final Activity secondaryActivity = createMockActivity();
-        final TaskFragmentContainer bottomTf = mController.newContainer(secondaryActivity, TASK_ID);
-        final TaskFragmentContainer primaryTf = mController.newContainer(mActivity, TASK_ID);
+        final TaskFragmentContainer bottomTf = createTfContainer(mController, secondaryActivity);
+        final TaskFragmentContainer primaryTf = createTfContainer(mController, mActivity);
         final SplitPairRule rule = createSplitPairRuleBuilder(pair ->
                 pair.first == mActivity && pair.second == secondaryActivity, pair -> false,
                 metrics -> true)
diff --git a/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/TaskContainerTest.java b/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/TaskContainerTest.java
index 8913b22..2847232 100644
--- a/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/TaskContainerTest.java
+++ b/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/TaskContainerTest.java
@@ -29,6 +29,7 @@
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
+import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.mock;
 
@@ -57,6 +58,9 @@
  * Build/Install/Run:
  *  atest WMJetpackUnitTests:TaskContainerTest
  */
+
+// Suppress GuardedBy warning on unit tests
+@SuppressWarnings("GuardedBy")
 @Presubmit
 @SmallTest
 @RunWith(AndroidJUnit4.class)
@@ -126,8 +130,11 @@
 
         assertTrue(taskContainer.isEmpty());
 
-        final TaskFragmentContainer tf = new TaskFragmentContainer(null /* activity */,
-                new Intent(), taskContainer, mController, null /* pairedPrimaryContainer */);
+        doReturn(taskContainer).when(mController).getTaskContainer(anyInt());
+        final TaskFragmentContainer tf = new TaskFragmentContainer.Builder(mController,
+                taskContainer.getTaskId(), null /* activityInTask */)
+                .setPendingAppearedIntent(new Intent())
+                .build();
 
         assertFalse(taskContainer.isEmpty());
 
@@ -142,12 +149,17 @@
         final TaskContainer taskContainer = createTestTaskContainer();
         assertNull(taskContainer.getTopNonFinishingTaskFragmentContainer());
 
-        final TaskFragmentContainer tf0 = new TaskFragmentContainer(null /* activity */,
-                new Intent(), taskContainer, mController, null /* pairedPrimaryContainer */);
+        doReturn(taskContainer).when(mController).getTaskContainer(anyInt());
+        final TaskFragmentContainer tf0 = new TaskFragmentContainer.Builder(mController,
+                taskContainer.getTaskId(), null /* activityInTask */)
+                        .setPendingAppearedIntent(new Intent())
+                        .build();
         assertEquals(tf0, taskContainer.getTopNonFinishingTaskFragmentContainer());
 
-        final TaskFragmentContainer tf1 = new TaskFragmentContainer(null /* activity */,
-                new Intent(), taskContainer, mController, null /* pairedPrimaryContainer */);
+        final TaskFragmentContainer tf1 = new TaskFragmentContainer.Builder(mController,
+                taskContainer.getTaskId(), null /* activityInTask */)
+                        .setPendingAppearedIntent(new Intent())
+                        .build();
         assertEquals(tf1, taskContainer.getTopNonFinishingTaskFragmentContainer());
     }
 
diff --git a/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/TaskFragmentContainerTest.java b/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/TaskFragmentContainerTest.java
index 44ab2c4..7fab371 100644
--- a/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/TaskFragmentContainerTest.java
+++ b/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/TaskFragmentContainerTest.java
@@ -100,24 +100,27 @@
     @Test
     public void testNewContainer() {
         final TaskContainer taskContainer = createTestTaskContainer();
+        mController.addTaskContainer(taskContainer.getTaskId(), taskContainer);
 
         // One of the activity and the intent must be non-null
         assertThrows(IllegalArgumentException.class,
-                () -> new TaskFragmentContainer(null, null, taskContainer, mController,
-                        null /* pairedPrimaryContainer */));
+                () -> new TaskFragmentContainer.Builder(mController, taskContainer.getTaskId(),
+                        null /* activityInTask */).build());
 
         // One of the activity and the intent must be null.
         assertThrows(IllegalArgumentException.class,
-                () -> new TaskFragmentContainer(mActivity, mIntent, taskContainer, mController,
-                        null /* pairedPrimaryContainer */));
+                () -> new TaskFragmentContainer.Builder(mController, taskContainer.getTaskId(),
+                        null /* activityInTask */)
+                        .setPendingAppearedActivity(createMockActivity())
+                        .setPendingAppearedIntent(mIntent)
+                        .build());
     }
 
     @Test
     public void testFinish() {
         final TaskContainer taskContainer = createTestTaskContainer();
-        final TaskFragmentContainer container = new TaskFragmentContainer(mActivity,
-                null /* pendingAppearedIntent */, taskContainer, mController,
-                null /* pairedPrimaryContainer */);
+        final TaskFragmentContainer container = createTaskFragmentContainer(taskContainer,
+                mActivity, null /* pendingAppearedIntent */);
         doReturn(container).when(mController).getContainerWithActivity(mActivity);
 
         // Only remove the activity, but not clear the reference until appeared.
@@ -148,15 +151,13 @@
     @Test
     public void testFinish_notFinishActivityThatIsReparenting() {
         final TaskContainer taskContainer = createTestTaskContainer();
-        final TaskFragmentContainer container0 = new TaskFragmentContainer(mActivity,
-                null /* pendingAppearedIntent */, taskContainer, mController,
-                null /* pairedPrimaryContainer */);
+        final TaskFragmentContainer container0 = createTaskFragmentContainer(taskContainer,
+                mActivity, null /* pendingAppearedIntent */);
         final TaskFragmentInfo info = createMockTaskFragmentInfo(container0, mActivity);
         container0.setInfo(mTransaction, info);
         // Request to reparent the activity to a new TaskFragment.
-        final TaskFragmentContainer container1 = new TaskFragmentContainer(mActivity,
-                null /* pendingAppearedIntent */, taskContainer, mController,
-                null /* pairedPrimaryContainer */);
+        final TaskFragmentContainer container1 = createTaskFragmentContainer(taskContainer,
+                mActivity, null /* pendingAppearedIntent */);
         doReturn(container1).when(mController).getContainerWithActivity(mActivity);
 
         // The activity is requested to be reparented, so don't finish it.
@@ -171,15 +172,13 @@
     public void testFinish_alwaysFinishPlaceholder() {
         // Register container1 as a placeholder
         final TaskContainer taskContainer = createTestTaskContainer();
-        final TaskFragmentContainer container0 = new TaskFragmentContainer(mActivity,
-                null /* pendingAppearedIntent */, taskContainer, mController,
-                null /* pairedPrimaryContainer */);
+        final TaskFragmentContainer container0 = createTaskFragmentContainer(taskContainer,
+                mActivity, null /* pendingAppearedIntent */);
         final TaskFragmentInfo info0 = createMockTaskFragmentInfo(container0, mActivity);
         container0.setInfo(mTransaction, info0);
         final Activity placeholderActivity = createMockActivity();
-        final TaskFragmentContainer container1 = new TaskFragmentContainer(placeholderActivity,
-                null /* pendingAppearedIntent */, taskContainer, mController,
-                null /* pairedPrimaryContainer */);
+        final TaskFragmentContainer container1 = createTaskFragmentContainer(taskContainer,
+                placeholderActivity, null /* pendingAppearedIntent */);
         final TaskFragmentInfo info1 = createMockTaskFragmentInfo(container1, placeholderActivity);
         container1.setInfo(mTransaction, info1);
         final SplitAttributes splitAttributes = new SplitAttributes.Builder().build();
@@ -207,9 +206,8 @@
     public void testSetInfo() {
         final TaskContainer taskContainer = createTestTaskContainer();
         // Pending activity should be cleared when it has appeared on server side.
-        final TaskFragmentContainer pendingActivityContainer = new TaskFragmentContainer(mActivity,
-                null /* pendingAppearedIntent */, taskContainer, mController,
-                null /* pairedPrimaryContainer */);
+        final TaskFragmentContainer pendingActivityContainer = createTaskFragmentContainer(
+                taskContainer, mActivity, null /* pendingAppearedIntent */);
 
         assertTrue(pendingActivityContainer.mPendingAppearedActivities.contains(
                 mActivity.getActivityToken()));
@@ -221,9 +219,8 @@
         assertTrue(pendingActivityContainer.mPendingAppearedActivities.isEmpty());
 
         // Pending intent should be cleared when the container becomes non-empty.
-        final TaskFragmentContainer pendingIntentContainer = new TaskFragmentContainer(
-                null /* pendingAppearedActivity */, mIntent, taskContainer, mController,
-                null /* pairedPrimaryContainer */);
+        final TaskFragmentContainer pendingIntentContainer = createTaskFragmentContainer(
+                taskContainer, null /* pendingAppearedActivity */, mIntent);
 
         assertEquals(mIntent, pendingIntentContainer.getPendingAppearedIntent());
 
@@ -237,8 +234,8 @@
     @Test
     public void testIsWaitingActivityAppear() {
         final TaskContainer taskContainer = createTestTaskContainer();
-        final TaskFragmentContainer container = new TaskFragmentContainer(null /* activity */,
-                mIntent, taskContainer, mController, null /* pairedPrimaryContainer */);
+        final TaskFragmentContainer container = createTaskFragmentContainer(
+                taskContainer, null  /* pendingAppearedActivity */, mIntent);
 
         assertTrue(container.isWaitingActivityAppear());
 
@@ -259,8 +256,8 @@
     public void testAppearEmptyTimeout() {
         doNothing().when(mController).onTaskFragmentAppearEmptyTimeout(any(), any());
         final TaskContainer taskContainer = createTestTaskContainer();
-        final TaskFragmentContainer container = new TaskFragmentContainer(null /* activity */,
-                mIntent, taskContainer, mController, null /* pairedPrimaryContainer */);
+        final TaskFragmentContainer container = createTaskFragmentContainer(
+                taskContainer, null /* pendingAppearedActivity */, mIntent);
 
         assertNull(container.mAppearEmptyTimeout);
 
@@ -299,8 +296,8 @@
     @Test
     public void testCollectNonFinishingActivities() {
         final TaskContainer taskContainer = createTestTaskContainer();
-        final TaskFragmentContainer container = new TaskFragmentContainer(null /* activity */,
-                mIntent, taskContainer, mController, null /* pairedPrimaryContainer */);
+        final TaskFragmentContainer container = createTaskFragmentContainer(
+                taskContainer, null /* pendingAppearedActivity */, mIntent);
         List<Activity> activities = container.collectNonFinishingActivities();
 
         assertTrue(activities.isEmpty());
@@ -327,8 +324,8 @@
     @Test
     public void testCollectNonFinishingActivities_checkIfStable() {
         final TaskContainer taskContainer = createTestTaskContainer();
-        final TaskFragmentContainer container = new TaskFragmentContainer(null /* activity */,
-                mIntent, taskContainer, mController, null /* pairedPrimaryContainer */);
+        final TaskFragmentContainer container = createTaskFragmentContainer(
+                taskContainer, null /* pendingAppearedActivity */, mIntent);
 
         // In case mInfo is null, collectNonFinishingActivities(true) should return null.
         List<Activity> activities =
@@ -353,8 +350,8 @@
     @Test
     public void testAddPendingActivity() {
         final TaskContainer taskContainer = createTestTaskContainer();
-        final TaskFragmentContainer container = new TaskFragmentContainer(null /* activity */,
-                mIntent, taskContainer, mController, null /* pairedPrimaryContainer */);
+        final TaskFragmentContainer container = createTaskFragmentContainer(
+                taskContainer, null /* pendingAppearedActivity */, mIntent);
         container.addPendingAppearedActivity(mActivity);
 
         assertEquals(1, container.collectNonFinishingActivities().size());
@@ -367,10 +364,10 @@
     @Test
     public void testIsAbove() {
         final TaskContainer taskContainer = createTestTaskContainer();
-        final TaskFragmentContainer container0 = new TaskFragmentContainer(null /* activity */,
-                mIntent, taskContainer, mController, null /* pairedPrimaryContainer */);
-        final TaskFragmentContainer container1 = new TaskFragmentContainer(null /* activity */,
-                mIntent, taskContainer, mController, null /* pairedPrimaryContainer */);
+        final TaskFragmentContainer container0 = createTaskFragmentContainer(
+                taskContainer, null /* pendingAppearedActivity */, mIntent);
+        final TaskFragmentContainer container1 = createTaskFragmentContainer(
+                taskContainer, null /* pendingAppearedActivity */, mIntent);
 
         assertTrue(container1.isAbove(container0));
         assertFalse(container0.isAbove(container1));
@@ -379,8 +376,8 @@
     @Test
     public void testGetBottomMostActivity() {
         final TaskContainer taskContainer = createTestTaskContainer();
-        final TaskFragmentContainer container = new TaskFragmentContainer(null /* activity */,
-                mIntent, taskContainer, mController, null /* pairedPrimaryContainer */);
+        final TaskFragmentContainer container = createTaskFragmentContainer(
+                taskContainer, null /* pendingAppearedActivity */, mIntent);
         container.addPendingAppearedActivity(mActivity);
 
         assertEquals(mActivity, container.getBottomMostActivity());
@@ -396,8 +393,8 @@
     @Test
     public void testOnActivityDestroyed() {
         final TaskContainer taskContainer = createTestTaskContainer(mController);
-        final TaskFragmentContainer container = new TaskFragmentContainer(null /* activity */,
-                mIntent, taskContainer, mController, null /* pairedPrimaryContainer */);
+        final TaskFragmentContainer container = createTaskFragmentContainer(
+                taskContainer, null /* pendingAppearedActivity */, mIntent);
         container.addPendingAppearedActivity(mActivity);
         final List<IBinder> activities = new ArrayList<>();
         activities.add(mActivity.getActivityToken());
@@ -416,8 +413,8 @@
     public void testIsInIntermediateState() {
         // True if no info set.
         final TaskContainer taskContainer = createTestTaskContainer();
-        final TaskFragmentContainer container = new TaskFragmentContainer(null /* activity */,
-                mIntent, taskContainer, mController, null /* pairedPrimaryContainer */);
+        final TaskFragmentContainer container = createTaskFragmentContainer(
+                taskContainer, null /* pendingAppearedActivity */, mIntent);
         spyOn(taskContainer);
         doReturn(true).when(taskContainer).isVisible();
 
@@ -479,8 +476,8 @@
     @Test
     public void testHasAppearedActivity() {
         final TaskContainer taskContainer = createTestTaskContainer();
-        final TaskFragmentContainer container = new TaskFragmentContainer(null /* activity */,
-                mIntent, taskContainer, mController, null /* pairedPrimaryContainer */);
+        final TaskFragmentContainer container = createTaskFragmentContainer(
+                taskContainer, null /* pendingAppearedActivity */, mIntent);
         container.addPendingAppearedActivity(mActivity);
 
         assertFalse(container.hasAppearedActivity(mActivity.getActivityToken()));
@@ -496,8 +493,8 @@
     @Test
     public void testHasPendingAppearedActivity() {
         final TaskContainer taskContainer = createTestTaskContainer();
-        final TaskFragmentContainer container = new TaskFragmentContainer(null /* activity */,
-                mIntent, taskContainer, mController, null /* pairedPrimaryContainer */);
+        final TaskFragmentContainer container = createTaskFragmentContainer(
+                taskContainer, null /* pendingAppearedActivity */, mIntent);
         container.addPendingAppearedActivity(mActivity);
 
         assertTrue(container.hasPendingAppearedActivity(mActivity.getActivityToken()));
@@ -513,10 +510,10 @@
     @Test
     public void testHasActivity() {
         final TaskContainer taskContainer = createTestTaskContainer(mController);
-        final TaskFragmentContainer container1 = new TaskFragmentContainer(null /* activity */,
-                mIntent, taskContainer, mController, null /* pairedPrimaryContainer */);
-        final TaskFragmentContainer container2 = new TaskFragmentContainer(null /* activity */,
-                mIntent, taskContainer, mController, null /* pairedPrimaryContainer */);
+        final TaskFragmentContainer container1 = createTaskFragmentContainer(
+                taskContainer, null /* pendingAppearedActivity */, mIntent);
+        final TaskFragmentContainer container2 = createTaskFragmentContainer(
+                taskContainer, null /* pendingAppearedActivity */, mIntent);
 
         // Activity is pending appeared on container2.
         container2.addPendingAppearedActivity(mActivity);
@@ -550,17 +547,19 @@
     @Test
     public void testNewContainerWithPairedPrimaryContainer() {
         final TaskContainer taskContainer = createTestTaskContainer();
-        final TaskFragmentContainer tf0 = new TaskFragmentContainer(
-                null /* pendingAppearedActivity */, new Intent(), taskContainer, mController,
-                null /* pairedPrimaryTaskFragment */);
-        final TaskFragmentContainer tf1 = new TaskFragmentContainer(
-                null /* pendingAppearedActivity */, new Intent(), taskContainer, mController,
-                null /* pairedPrimaryTaskFragment */);
+        mController.addTaskContainer(taskContainer.getTaskId(), taskContainer);
+        final TaskFragmentContainer tf0 = createTaskFragmentContainer(
+                taskContainer, null /* pendingAppearedActivity */, new Intent());
+        final TaskFragmentContainer tf1 = createTaskFragmentContainer(
+                taskContainer, null /* pendingAppearedActivity */, new Intent());
 
         // When tf2 is created with using tf0 as pairedPrimaryContainer, tf2 should be inserted
         // right above tf0.
-        final TaskFragmentContainer tf2 = new TaskFragmentContainer(
-                null /* pendingAppearedActivity */, new Intent(), taskContainer, mController, tf0);
+        final TaskFragmentContainer tf2 = new TaskFragmentContainer.Builder(mController,
+                taskContainer.getTaskId(), null /* activityInTask */)
+                .setPendingAppearedIntent(new Intent())
+                .setPairedPrimaryContainer(tf0)
+                .build();
         assertEquals(0, taskContainer.indexOf(tf0));
         assertEquals(1, taskContainer.indexOf(tf2));
         assertEquals(2, taskContainer.indexOf(tf1));
@@ -569,18 +568,15 @@
     @Test
     public void testNewContainerWithPairedPendingAppearedActivity() {
         final TaskContainer taskContainer = createTestTaskContainer();
-        final TaskFragmentContainer tf0 = new TaskFragmentContainer(
-                createMockActivity(), null /* pendingAppearedIntent */, taskContainer, mController,
-                null /* pairedPrimaryTaskFragment */);
-        final TaskFragmentContainer tf1 = new TaskFragmentContainer(
-                null /* pendingAppearedActivity */, new Intent(), taskContainer, mController,
-                null /* pairedPrimaryTaskFragment */);
+        final TaskFragmentContainer tf0 = createTaskFragmentContainer(taskContainer,
+                createMockActivity(), null /* pendingAppearedIntent */);
+        final TaskFragmentContainer tf1 = createTaskFragmentContainer(taskContainer,
+                null /* pendingAppearedActivity */, new Intent());
 
         // When tf2 is created with pendingAppearedActivity, tf2 should be inserted below any
         // TaskFragment without any Activity.
-        final TaskFragmentContainer tf2 = new TaskFragmentContainer(
-                createMockActivity(), null /* pendingAppearedIntent */, taskContainer, mController,
-                null /* pairedPrimaryTaskFragment */);
+        final TaskFragmentContainer tf2 = createTaskFragmentContainer(taskContainer,
+                createMockActivity(), null /* pendingAppearedIntent */);
         assertEquals(0, taskContainer.indexOf(tf0));
         assertEquals(1, taskContainer.indexOf(tf2));
         assertEquals(2, taskContainer.indexOf(tf1));
@@ -589,9 +585,8 @@
     @Test
     public void testIsVisible() {
         final TaskContainer taskContainer = createTestTaskContainer();
-        final TaskFragmentContainer container = new TaskFragmentContainer(
-                null /* pendingAppearedActivity */, new Intent(), taskContainer, mController,
-                null /* pairedPrimaryTaskFragment */);
+        final TaskFragmentContainer container = createTaskFragmentContainer(taskContainer,
+                null /* pendingAppearedActivity */, new Intent());
 
         // Not visible when there is not appeared.
         assertFalse(container.isVisible());
@@ -617,4 +612,14 @@
         doReturn(activity).when(mController).getActivity(activityToken);
         return activity;
     }
+
+    private TaskFragmentContainer createTaskFragmentContainer(TaskContainer taskContainer,
+            Activity pendingAppearedActivity, Intent pendingAppearedIntent) {
+        final int taskId = taskContainer.getTaskId();
+        mController.addTaskContainer(taskId, taskContainer);
+        return new TaskFragmentContainer.Builder(mController, taskId, pendingAppearedActivity)
+                .setPendingAppearedActivity(pendingAppearedActivity)
+                .setPendingAppearedIntent(pendingAppearedIntent)
+                .build();
+    }
 }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellModule.java b/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellModule.java
index fb0a1ab..12bbd51 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellModule.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellModule.java
@@ -522,14 +522,16 @@
             RecentsTransitionHandler recentsTransitionHandler,
             MultiInstanceHelper multiInstanceHelper,
             @ShellMainThread ShellExecutor mainExecutor,
-            Optional<DesktopTasksLimiter> desktopTasksLimiter) {
+            Optional<DesktopTasksLimiter> desktopTasksLimiter,
+            Optional<RecentTasksController> recentTasksController) {
         return new DesktopTasksController(context, shellInit, shellCommandHandler, shellController,
                 displayController, shellTaskOrganizer, syncQueue, rootTaskDisplayAreaOrganizer,
                 dragAndDropController, transitions, enterDesktopTransitionHandler,
                 exitDesktopTransitionHandler, toggleResizeDesktopTaskTransitionHandler,
                 dragToDesktopTransitionHandler, desktopModeTaskRepository,
                 desktopModeLoggerTransitionObserver, launchAdjacentController,
-                recentsTransitionHandler, multiInstanceHelper, mainExecutor, desktopTasksLimiter);
+                recentsTransitionHandler, multiInstanceHelper,
+                mainExecutor, desktopTasksLimiter, recentTasksController.orElse(null));
     }
 
     @WMSingleton
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 6e45397..ef384c7 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
@@ -70,6 +70,7 @@
 import com.android.wm.shell.desktopmode.DragToDesktopTransitionHandler.DragToDesktopStateListener
 import com.android.wm.shell.draganddrop.DragAndDropController
 import com.android.wm.shell.protolog.ShellProtoLogGroup.WM_SHELL_DESKTOP_MODE
+import com.android.wm.shell.recents.RecentTasksController
 import com.android.wm.shell.recents.RecentsTransitionHandler
 import com.android.wm.shell.recents.RecentsTransitionStateListener
 import com.android.wm.shell.shared.DesktopModeStatus
@@ -118,6 +119,7 @@
     private val multiInstanceHelper: MultiInstanceHelper,
     @ShellMainThread private val mainExecutor: ShellExecutor,
     private val desktopTasksLimiter: Optional<DesktopTasksLimiter>,
+    private val recentTasksController: RecentTasksController?
 ) :
     RemoteCallable<DesktopTasksController>,
     Transitions.TransitionHandler,
@@ -293,24 +295,49 @@
         taskId: Int,
         wct: WindowContainerTransaction = WindowContainerTransaction()
     ): Boolean {
-        shellTaskOrganizer.getRunningTaskInfo(taskId)?.let { task -> moveToDesktop(task, wct) }
-            ?: return false
+        shellTaskOrganizer.getRunningTaskInfo(taskId)?.let {
+            moveToDesktop(it, wct)
+        } ?: moveToDesktopFromNonRunningTask(taskId, wct)
         return true
     }
 
-    /** Move a task to desktop */
+    private fun moveToDesktopFromNonRunningTask(
+        taskId: Int,
+        wct: WindowContainerTransaction
+    ): Boolean {
+        recentTasksController?.findTaskInBackground(taskId)?.let {
+            KtProtoLog.v(
+                WM_SHELL_DESKTOP_MODE,
+                "DesktopTasksController: moveToDesktopFromNonRunningTask taskId=%d",
+                taskId
+            )
+            // TODO(342378842): Instead of using default display, support multiple displays
+            val taskToMinimize =
+                bringDesktopAppsToFrontBeforeShowingNewTask(DEFAULT_DISPLAY, wct, taskId)
+            addMoveToDesktopChangesNonRunningTask(wct, taskId)
+            // TODO(343149901): Add DPI changes for task launch
+            val transition = enterDesktopTaskTransitionHandler.moveToDesktop(wct)
+            addPendingMinimizeTransition(transition, taskToMinimize)
+            return true
+        } ?: return false
+    }
+
+    private fun addMoveToDesktopChangesNonRunningTask(
+        wct: WindowContainerTransaction,
+        taskId: Int
+    ) {
+        val options = ActivityOptions.makeBasic()
+        options.launchWindowingMode = WINDOWING_MODE_FREEFORM
+        wct.startTask(taskId, options.toBundle())
+    }
+
+    /**
+     * Move a task to desktop
+     */
     fun moveToDesktop(
         task: RunningTaskInfo,
         wct: WindowContainerTransaction = WindowContainerTransaction()
     ) {
-        if (!DesktopModeStatus.canEnterDesktopMode(context)) {
-            KtProtoLog.w(
-                WM_SHELL_DESKTOP_MODE,
-                "DesktopTasksController: Cannot enter desktop, " +
-                    "display does not meet minimum size requirements"
-            )
-            return
-        }
         if (Flags.enableDesktopWindowingModalsPolicy() && isSingleTopActivityTranslucent(task)) {
             KtProtoLog.w(
                 WM_SHELL_DESKTOP_MODE,
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/recents/RecentTasksController.java b/libs/WindowManager/Shell/src/com/android/wm/shell/recents/RecentTasksController.java
index d8f2c02..863202d 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/recents/RecentTasksController.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/recents/RecentTasksController.java
@@ -446,6 +446,25 @@
         return null;
     }
 
+    /**
+     * Find the background task that match the given taskId.
+     */
+    @Nullable
+    public ActivityManager.RecentTaskInfo findTaskInBackground(int taskId) {
+        List<ActivityManager.RecentTaskInfo> tasks = mActivityTaskManager.getRecentTasks(
+                Integer.MAX_VALUE, ActivityManager.RECENT_IGNORE_UNAVAILABLE,
+                ActivityManager.getCurrentUser());
+        for (int i = 0; i < tasks.size(); i++) {
+            final ActivityManager.RecentTaskInfo task = tasks.get(i);
+            if (task.isVisible) {
+                continue;
+            }
+            if (taskId == task.taskId) {
+                return task;
+            }
+        }
+        return null;
+    }
     public void dump(@NonNull PrintWriter pw, String prefix) {
         final String innerPrefix = prefix + "  ";
         pw.println(prefix + TAG);
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksControllerTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksControllerTest.kt
index ac67bd1..cf6cea2 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksControllerTest.kt
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksControllerTest.kt
@@ -16,6 +16,7 @@
 
 package com.android.wm.shell.desktopmode
 
+import android.app.ActivityManager.RecentTaskInfo
 import android.app.ActivityManager.RunningTaskInfo
 import android.app.WindowConfiguration.ACTIVITY_TYPE_HOME
 import android.app.WindowConfiguration.ACTIVITY_TYPE_STANDARD
@@ -47,13 +48,16 @@
 import android.view.WindowManager.TRANSIT_TO_BACK
 import android.view.WindowManager.TRANSIT_TO_FRONT
 import android.window.DisplayAreaInfo
+import android.window.IWindowContainerToken
 import android.window.RemoteTransition
 import android.window.TransitionRequestInfo
 import android.window.WindowContainerToken
 import android.window.WindowContainerTransaction
+import android.window.WindowContainerTransaction.HierarchyOp.HIERARCHY_OP_TYPE_LAUNCH_TASK
 import android.window.WindowContainerTransaction.HierarchyOp.HIERARCHY_OP_TYPE_PENDING_INTENT
 import android.window.WindowContainerTransaction.HierarchyOp.HIERARCHY_OP_TYPE_REMOVE_TASK
 import android.window.WindowContainerTransaction.HierarchyOp.HIERARCHY_OP_TYPE_REORDER
+import android.window.WindowContainerTransaction.HierarchyOp.LAUNCH_KEY_TASK_ID
 import androidx.test.filters.SmallTest
 import com.android.dx.mockito.inline.extended.ExtendedMockito.doReturn
 import com.android.dx.mockito.inline.extended.ExtendedMockito.mockitoSession
@@ -78,6 +82,7 @@
 import com.android.wm.shell.desktopmode.DesktopTestHelpers.Companion.createHomeTask
 import com.android.wm.shell.desktopmode.DesktopTestHelpers.Companion.createSplitScreenTask
 import com.android.wm.shell.draganddrop.DragAndDropController
+import com.android.wm.shell.recents.RecentTasksController
 import com.android.wm.shell.recents.RecentsTransitionHandler
 import com.android.wm.shell.recents.RecentsTransitionStateListener
 import com.android.wm.shell.shared.DesktopModeStatus
@@ -93,6 +98,7 @@
 import com.android.wm.shell.transition.Transitions.TransitionHandler
 import com.google.common.truth.Truth.assertThat
 import com.google.common.truth.Truth.assertWithMessage
+import java.util.Optional
 import org.junit.After
 import org.junit.Assume.assumeTrue
 import org.junit.Before
@@ -115,7 +121,6 @@
 import org.mockito.kotlin.atLeastOnce
 import org.mockito.kotlin.capture
 import org.mockito.quality.Strictness
-import java.util.Optional
 import junit.framework.Assert.assertFalse
 import junit.framework.Assert.assertTrue
 import org.mockito.Mockito.`when` as whenever
@@ -154,6 +159,7 @@
     @Mock lateinit var multiInstanceHelper: MultiInstanceHelper
     @Mock lateinit var desktopModeLoggerTransitionObserver: DesktopModeLoggerTransitionObserver
     @Mock lateinit var desktopModeVisualIndicator: DesktopModeVisualIndicator
+    @Mock lateinit var recentTasksController: RecentTasksController
 
     private lateinit var mockitoSession: StaticMockitoSession
     private lateinit var controller: DesktopTasksController
@@ -233,6 +239,7 @@
             multiInstanceHelper,
             shellExecutor,
             Optional.of(desktopTasksLimiter),
+            recentTasksController
         )
     }
 
@@ -622,7 +629,7 @@
         controller.moveToDesktop(task)
         val wct = getLatestMoveToDesktopWct()
         assertThat(wct.changes[task.token.asBinder()]?.windowingMode)
-            .isEqualTo(WINDOWING_MODE_FREEFORM)
+                .isEqualTo(WINDOWING_MODE_FREEFORM)
     }
 
     @Test
@@ -643,14 +650,17 @@
     }
 
     @Test
-    fun moveToDesktop_deviceNotSupported_doesNothing() {
-        val task = setUpFullscreenTask()
+    fun moveToDesktop_nonRunningTask_launchesInFreeform() {
+        whenever(shellTaskOrganizer.getRunningTaskInfo(anyInt())).thenReturn(null)
 
-        // Simulate non compatible device
-        doReturn(false).`when` { DesktopModeStatus.isDesktopModeSupported(any()) }
+        val task = createTaskInfo(1)
 
-        controller.moveToDesktop(task)
-        verifyWCTNotExecuted()
+        whenever(recentTasksController.findTaskInBackground(anyInt())).thenReturn(task)
+
+        controller.moveToDesktop(task.taskId)
+        with(getLatestMoveToDesktopWct()){
+            assertLaunchTaskAt(0, task.taskId, WINDOWING_MODE_FREEFORM)
+        }
     }
 
     @Test
@@ -666,6 +676,17 @@
     }
 
     @Test
+    fun moveToDesktop_deviceNotSupported_doesNothing() {
+        val task = setUpFullscreenTask()
+
+        // Simulate non compatible device
+        doReturn(false).`when` { DesktopModeStatus.isDesktopModeSupported(any()) }
+
+        controller.moveToDesktop(task)
+        verifyWCTNotExecuted()
+    }
+
+    @Test
     fun moveToDesktop_deviceNotSupported_deviceRestrictionsOverridden_taskIsMovedToDesktop() {
         val task = setUpFullscreenTask()
 
@@ -1834,6 +1855,20 @@
     assertThat(op.pendingIntent?.intent?.component).isEqualTo(intent.component)
 }
 
+private fun WindowContainerTransaction.assertLaunchTaskAt(
+    index: Int,
+    taskId: Int,
+    windowingMode: Int
+) {
+    val keyLaunchWindowingMode = "android.activity.windowingMode"
+
+    assertIndexInBounds(index)
+    val op = hierarchyOps[index]
+    assertThat(op.type).isEqualTo(HIERARCHY_OP_TYPE_LAUNCH_TASK)
+    assertThat(op.launchOptions?.getInt(LAUNCH_KEY_TASK_ID)).isEqualTo(taskId)
+    assertThat(op.launchOptions?.getInt(keyLaunchWindowingMode, WINDOWING_MODE_UNDEFINED))
+            .isEqualTo(windowingMode)
+}
 private fun WindowContainerTransaction?.anyDensityConfigChange(
     token: WindowContainerToken
 ): Boolean {
@@ -1841,3 +1876,7 @@
         change.key == token.asBinder() && ((change.value.configSetMask and CONFIG_DENSITY) != 0)
     } ?: false
 }
+private fun createTaskInfo(id: Int) = RecentTaskInfo().apply {
+    taskId = id
+    token = WindowContainerToken(mock(IWindowContainerToken::class.java))
+}
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/notifications/ui/composable/NotificationHeadsUpHeight.kt b/packages/SystemUI/compose/features/src/com/android/systemui/notifications/ui/composable/NotificationHeadsUpHeight.kt
new file mode 100644
index 0000000..75a565b
--- /dev/null
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/notifications/ui/composable/NotificationHeadsUpHeight.kt
@@ -0,0 +1,87 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.notifications.ui.composable
+
+import androidx.compose.ui.Modifier
+import androidx.compose.ui.layout.Measurable
+import androidx.compose.ui.layout.MeasureResult
+import androidx.compose.ui.layout.MeasureScope
+import androidx.compose.ui.node.LayoutModifierNode
+import androidx.compose.ui.node.ModifierNodeElement
+import androidx.compose.ui.node.invalidateMeasurement
+import androidx.compose.ui.unit.Constraints
+import androidx.compose.ui.unit.IntOffset
+import com.android.systemui.statusbar.notification.stack.ui.view.NotificationScrollView
+
+/**
+ * Modify element, which updates the height to the height of current top heads up notification, or
+ * to 0 if there is none.
+ *
+ * @param view Notification stack scroll view
+ */
+fun Modifier.notificationHeadsUpHeight(view: NotificationScrollView) =
+    this then HeadsUpLayoutElement(view)
+
+private data class HeadsUpLayoutElement(
+    val view: NotificationScrollView,
+) : ModifierNodeElement<HeadsUpLayoutNode>() {
+
+    override fun create(): HeadsUpLayoutNode = HeadsUpLayoutNode(view)
+
+    override fun update(node: HeadsUpLayoutNode) {
+        check(view == node.view) { "Trying to reuse the node with a new View." }
+    }
+}
+
+private class HeadsUpLayoutNode(val view: NotificationScrollView) :
+    LayoutModifierNode, Modifier.Node() {
+
+    private val headsUpHeightChangedListener = Runnable { invalidateMeasureIfAttached() }
+
+    override fun onAttach() {
+        super.onAttach()
+        view.addHeadsUpHeightChangedListener(headsUpHeightChangedListener)
+    }
+
+    override fun onDetach() {
+        super.onDetach()
+        view.removeHeadsUpHeightChangedListener(headsUpHeightChangedListener)
+    }
+
+    override fun MeasureScope.measure(
+        measurable: Measurable,
+        constraints: Constraints
+    ): MeasureResult {
+        // TODO(b/339181697) make sure, that the row is already measured.
+        val contentHeight = view.topHeadsUpHeight
+        val placeable =
+            measurable.measure(
+                constraints.copy(minHeight = contentHeight, maxHeight = contentHeight)
+            )
+        return layout(placeable.width, placeable.height) { placeable.place(IntOffset.Zero) }
+    }
+
+    override fun toString(): String {
+        return "HeadsUpLayoutNode(view=$view)"
+    }
+
+    fun invalidateMeasureIfAttached() {
+        if (isAttached) {
+            this.invalidateMeasurement()
+        }
+    }
+}
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/notifications/ui/composable/Notifications.kt b/packages/SystemUI/compose/features/src/com/android/systemui/notifications/ui/composable/Notifications.kt
index e6132c6..c26259f 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/notifications/ui/composable/Notifications.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/notifications/ui/composable/Notifications.kt
@@ -67,7 +67,6 @@
 import com.android.compose.animation.scene.ElementKey
 import com.android.compose.animation.scene.NestedScrollBehavior
 import com.android.compose.animation.scene.SceneScope
-import com.android.compose.modifiers.height
 import com.android.compose.modifiers.thenIf
 import com.android.systemui.common.ui.compose.windowinsets.LocalRawScreenHeight
 import com.android.systemui.common.ui.compose.windowinsets.LocalScreenCornerRadius
@@ -108,18 +107,17 @@
  */
 @Composable
 fun SceneScope.HeadsUpNotificationSpace(
+    stackScrollView: NotificationScrollView,
     viewModel: NotificationsPlaceholderViewModel,
     modifier: Modifier = Modifier,
     isPeekFromBottom: Boolean = false,
 ) {
-    val headsUpHeight = viewModel.headsUpHeight.collectAsStateWithLifecycle()
-
     Element(
         Notifications.Elements.HeadsUpNotificationPlaceholder,
         modifier =
             modifier
-                .height { headsUpHeight.value.roundToInt() }
                 .fillMaxWidth()
+                .notificationHeadsUpHeight(stackScrollView)
                 .debugBackground(viewModel, DEBUG_HUN_COLOR)
                 .onGloballyPositioned { coordinates: LayoutCoordinates ->
                     val boundsInWindow = coordinates.boundsInWindow()
@@ -152,6 +150,7 @@
             modifier = Modifier.fillMaxSize(),
         )
         HeadsUpNotificationSpace(
+            stackScrollView = stackScrollView,
             viewModel = viewModel,
             modifier = Modifier.align(Alignment.TopCenter),
         )
@@ -358,7 +357,7 @@
                         .onSizeChanged { size -> stackHeight.intValue = size.height },
             )
         }
-        HeadsUpNotificationSpace(viewModel = viewModel)
+        HeadsUpNotificationSpace(stackScrollView = stackScrollView, viewModel = viewModel)
     }
 }
 
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/qs/ui/composable/QuickSettingsScene.kt b/packages/SystemUI/compose/features/src/com/android/systemui/qs/ui/composable/QuickSettingsScene.kt
index d76b19f..0ee485c 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/qs/ui/composable/QuickSettingsScene.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/qs/ui/composable/QuickSettingsScene.kt
@@ -42,7 +42,6 @@
 import androidx.compose.foundation.layout.fillMaxWidth
 import androidx.compose.foundation.layout.height
 import androidx.compose.foundation.layout.navigationBars
-import androidx.compose.foundation.layout.offset
 import androidx.compose.foundation.layout.padding
 import androidx.compose.foundation.layout.wrapContentHeight
 import androidx.compose.foundation.rememberScrollState
@@ -60,7 +59,6 @@
 import androidx.compose.ui.platform.LocalDensity
 import androidx.compose.ui.platform.LocalLifecycleOwner
 import androidx.compose.ui.res.colorResource
-import androidx.compose.ui.unit.IntOffset
 import androidx.compose.ui.unit.dp
 import androidx.lifecycle.compose.collectAsStateWithLifecycle
 import com.android.compose.animation.scene.SceneScope
@@ -79,14 +77,13 @@
 import com.android.systemui.media.controls.ui.controller.MediaCarouselController
 import com.android.systemui.media.controls.ui.view.MediaHost
 import com.android.systemui.media.dagger.MediaModule
-import com.android.systemui.notifications.ui.composable.NotificationScrollingStack
+import com.android.systemui.notifications.ui.composable.HeadsUpNotificationSpace
 import com.android.systemui.qs.footer.ui.compose.FooterActionsWithAnimatedVisibility
 import com.android.systemui.qs.ui.viewmodel.QuickSettingsSceneViewModel
 import com.android.systemui.res.R
 import com.android.systemui.scene.session.ui.composable.SaveableSession
 import com.android.systemui.scene.shared.model.Scenes
 import com.android.systemui.scene.ui.composable.ComposableScene
-import com.android.systemui.shade.shared.model.ShadeMode
 import com.android.systemui.shade.ui.composable.CollapsedShadeHeader
 import com.android.systemui.shade.ui.composable.ExpandedShadeHeader
 import com.android.systemui.shade.ui.composable.Shade
@@ -99,7 +96,6 @@
 import dagger.Lazy
 import javax.inject.Inject
 import javax.inject.Named
-import kotlin.math.roundToInt
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.stateIn
@@ -368,15 +364,11 @@
                     Modifier.align(Alignment.CenterHorizontally).sysuiResTag("qs_footer_actions"),
             )
         }
-        NotificationScrollingStack(
+        HeadsUpNotificationSpace(
             stackScrollView = notificationStackScrollView,
             viewModel = notificationsPlaceholderViewModel,
-            shadeSession = shadeSession,
-            maxScrimTop = { screenHeight },
-            shouldPunchHoleBehindScrim = shouldPunchHoleBehindScrim,
-            shadeMode = ShadeMode.Single,
-            modifier =
-                Modifier.fillMaxWidth().offset { IntOffset(x = 0, y = screenHeight.roundToInt()) },
+            modifier = Modifier.align(Alignment.BottomCenter),
+            isPeekFromBottom = true,
         )
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/KeyboardShortcutListSearch.java b/packages/SystemUI/src/com/android/systemui/statusbar/KeyboardShortcutListSearch.java
index 0d8030f..5268009 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/KeyboardShortcutListSearch.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/KeyboardShortcutListSearch.java
@@ -19,9 +19,14 @@
 import static android.view.View.IMPORTANT_FOR_ACCESSIBILITY_YES;
 import static android.view.WindowManager.LayoutParams.TYPE_SYSTEM_DIALOG;
 
+import static com.android.systemui.Flags.validateKeyboardShortcutHelperIconUri;
+
 import android.annotation.NonNull;
 import android.annotation.Nullable;
+import android.app.ActivityManager;
 import android.app.AppGlobals;
+import android.app.SynchronousUserSwitchObserver;
+import android.app.UserSwitchObserver;
 import android.content.ComponentName;
 import android.content.Context;
 import android.content.Intent;
@@ -37,6 +42,7 @@
 import android.graphics.drawable.Icon;
 import android.hardware.input.InputManagerGlobal;
 import android.os.Handler;
+import android.os.HandlerThread;
 import android.os.Looper;
 import android.os.RemoteException;
 import android.text.Editable;
@@ -136,6 +142,8 @@
     };
 
     private final Handler mHandler = new Handler(Looper.getMainLooper());
+    private final HandlerThread mHandlerThread = new HandlerThread("KeyboardShortcutHelper");
+    @VisibleForTesting Handler mBackgroundHandler;
     @VisibleForTesting public Context mContext;
     private final IPackageManager mPackageManager;
 
@@ -143,6 +151,13 @@
     private KeyCharacterMap mKeyCharacterMap;
     private KeyCharacterMap mBackupKeyCharacterMap;
 
+    private final UserSwitchObserver mUserSwitchObserver = new SynchronousUserSwitchObserver() {
+            @Override
+            public void onUserSwitching(int newUserId) throws RemoteException {
+                dismiss();
+            }
+    };
+
     @VisibleForTesting
     KeyboardShortcutListSearch(Context context, WindowManager windowManager) {
         this.mContext = new ContextThemeWrapper(
@@ -413,36 +428,75 @@
     private boolean mAppShortcutsReceived;
     private boolean mImeShortcutsReceived;
 
-    @VisibleForTesting
-    public void showKeyboardShortcuts(int deviceId) {
-        retrieveKeyCharacterMap(deviceId);
-        mAppShortcutsReceived = false;
-        mImeShortcutsReceived = false;
-        mWindowManager.requestAppKeyboardShortcuts(result -> {
-            // Add specific app shortcuts
+    private void onAppSpecificShortcutsReceived(List<KeyboardShortcutGroup> result) {
+        // Add specific app shortcuts
+        if (result != null) {
             if (result.isEmpty()) {
                 mCurrentAppPackageName = null;
                 mKeySearchResultMap.put(SHORTCUT_SPECIFICAPP_INDEX, false);
             } else {
                 mCurrentAppPackageName = result.get(0).getPackageName();
-                mSpecificAppGroup.addAll(reMapToKeyboardShortcutMultiMappingGroup(result));
+                if (validateKeyboardShortcutHelperIconUri()) {
+                    KeyboardShortcuts.sanitiseShortcuts(result);
+                }
+                mSpecificAppGroup.addAll(
+                        reMapToKeyboardShortcutMultiMappingGroup(result));
                 mKeySearchResultMap.put(SHORTCUT_SPECIFICAPP_INDEX, true);
             }
-            mAppShortcutsReceived = true;
-            if (mImeShortcutsReceived) {
-                mergeAndShowKeyboardShortcutsGroups();
-            }
-        }, deviceId);
-        mWindowManager.requestImeKeyboardShortcuts(result -> {
-            // Add specific Ime shortcuts
+        }
+        mAppShortcutsReceived = true;
+        if (mImeShortcutsReceived) {
+            mergeAndShowKeyboardShortcutsGroups();
+        }
+    }
+
+    private void onImeSpecificShortcutsReceived(List<KeyboardShortcutGroup> result) {
+        // Add specific Ime shortcuts
+        if (result != null) {
             if (!result.isEmpty()) {
-                mInputGroup.addAll(reMapToKeyboardShortcutMultiMappingGroup(result));
+                if (validateKeyboardShortcutHelperIconUri()) {
+                    KeyboardShortcuts.sanitiseShortcuts(result);
+                }
+                mInputGroup.addAll(
+                        reMapToKeyboardShortcutMultiMappingGroup(result));
             }
-            mImeShortcutsReceived = true;
-            if (mAppShortcutsReceived) {
-                mergeAndShowKeyboardShortcutsGroups();
+        }
+        mImeShortcutsReceived = true;
+        if (mAppShortcutsReceived) {
+            mergeAndShowKeyboardShortcutsGroups();
+        }
+    }
+
+    @VisibleForTesting
+    public void showKeyboardShortcuts(int deviceId) {
+        if (mBackgroundHandler == null) {
+            mHandlerThread.start();
+            mBackgroundHandler = new Handler(mHandlerThread.getLooper());
+        }
+
+        if (validateKeyboardShortcutHelperIconUri()) {
+            try {
+                ActivityManager.getService().registerUserSwitchObserver(mUserSwitchObserver, TAG);
+            } catch (RemoteException e) {
+                Log.e(TAG, "could not register user switch observer", e);
             }
-        }, deviceId);
+        }
+
+        retrieveKeyCharacterMap(deviceId);
+        mAppShortcutsReceived = false;
+        mImeShortcutsReceived = false;
+        mWindowManager.requestAppKeyboardShortcuts(
+                result -> {
+                    mBackgroundHandler.post(() -> {
+                        onAppSpecificShortcutsReceived(result);
+                    });
+                }, deviceId);
+        mWindowManager.requestImeKeyboardShortcuts(
+                result -> {
+                    mBackgroundHandler.post(() -> {
+                        onImeSpecificShortcutsReceived(result);
+                    });
+                }, deviceId);
     }
 
     private void mergeAndShowKeyboardShortcutsGroups() {
@@ -508,6 +562,14 @@
             mKeyboardShortcutsBottomSheetDialog.dismiss();
             mKeyboardShortcutsBottomSheetDialog = null;
         }
+        mHandlerThread.quit();
+        if (validateKeyboardShortcutHelperIconUri()) {
+            try {
+                ActivityManager.getService().unregisterUserSwitchObserver(mUserSwitchObserver);
+            } catch (RemoteException e) {
+                Log.e(TAG, "Could not unregister user switch observer", e);
+            }
+        }
     }
 
     private KeyboardShortcutMultiMappingGroup getMultiMappingSystemShortcuts(Context context) {
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/KeyboardShortcuts.java b/packages/SystemUI/src/com/android/systemui/statusbar/KeyboardShortcuts.java
index 21f608e..d00916a 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/KeyboardShortcuts.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/KeyboardShortcuts.java
@@ -20,11 +20,16 @@
 import static android.view.View.IMPORTANT_FOR_ACCESSIBILITY_YES;
 import static android.view.WindowManager.LayoutParams.TYPE_SYSTEM_DIALOG;
 
+import static com.android.systemui.Flags.validateKeyboardShortcutHelperIconUri;
+
 import android.annotation.NonNull;
 import android.annotation.Nullable;
+import android.app.ActivityManager;
 import android.app.AlertDialog;
 import android.app.AppGlobals;
 import android.app.Dialog;
+import android.app.SynchronousUserSwitchObserver;
+import android.app.UserSwitchObserver;
 import android.content.ComponentName;
 import android.content.Context;
 import android.content.DialogInterface;
@@ -39,6 +44,7 @@
 import android.graphics.drawable.Icon;
 import android.hardware.input.InputManager;
 import android.os.Handler;
+import android.os.HandlerThread;
 import android.os.Looper;
 import android.os.RemoteException;
 import android.util.Log;
@@ -93,6 +99,8 @@
     };
 
     private final Handler mHandler = new Handler(Looper.getMainLooper());
+    private final HandlerThread mHandlerThread = new HandlerThread("KeyboardShortcutHelper");
+    @VisibleForTesting Handler mBackgroundHandler;
     @VisibleForTesting public Context mContext;
     private final IPackageManager mPackageManager;
     private final OnClickListener mDialogCloseListener = new DialogInterface.OnClickListener() {
@@ -129,6 +137,13 @@
     @Nullable private List<KeyboardShortcutGroup> mReceivedAppShortcutGroups = null;
     @Nullable private List<KeyboardShortcutGroup> mReceivedImeShortcutGroups = null;
 
+    private final UserSwitchObserver mUserSwitchObserver = new SynchronousUserSwitchObserver() {
+            @Override
+            public void onUserSwitching(int newUserId) throws RemoteException {
+                dismiss();
+            }
+    };
+
     @VisibleForTesting
     KeyboardShortcuts(Context context, WindowManager windowManager) {
         this.mContext = new ContextThemeWrapper(
@@ -374,21 +389,68 @@
 
     @VisibleForTesting
     public void showKeyboardShortcuts(int deviceId) {
+        if (mBackgroundHandler == null) {
+            mHandlerThread.start();
+            mBackgroundHandler = new Handler(mHandlerThread.getLooper());
+        }
+
+        if (validateKeyboardShortcutHelperIconUri()) {
+            try {
+                ActivityManager.getService().registerUserSwitchObserver(mUserSwitchObserver, TAG);
+            } catch (RemoteException e) {
+                Log.e(TAG, "could not register user switch observer", e);
+            }
+        }
+
         retrieveKeyCharacterMap(deviceId);
+
         mReceivedAppShortcutGroups = null;
         mReceivedImeShortcutGroups = null;
+
         mWindowManager.requestAppKeyboardShortcuts(
                 result -> {
-                    mReceivedAppShortcutGroups = result;
-                    maybeMergeAndShowKeyboardShortcuts();
+                    mBackgroundHandler.post(() -> {
+                        onAppSpecificShortcutsReceived(result);
+                    });
                 }, deviceId);
         mWindowManager.requestImeKeyboardShortcuts(
                 result -> {
-                    mReceivedImeShortcutGroups = result;
-                    maybeMergeAndShowKeyboardShortcuts();
+                    mBackgroundHandler.post(() -> {
+                        onImeSpecificShortcutsReceived(result);
+                    });
                 }, deviceId);
     }
 
+    private void onAppSpecificShortcutsReceived(List<KeyboardShortcutGroup> result) {
+        mReceivedAppShortcutGroups =
+                result == null ? Collections.emptyList() : result;
+
+        if (validateKeyboardShortcutHelperIconUri()) {
+            sanitiseShortcuts(mReceivedAppShortcutGroups);
+        }
+
+        maybeMergeAndShowKeyboardShortcuts();
+    }
+
+    private void onImeSpecificShortcutsReceived(List<KeyboardShortcutGroup> result) {
+        mReceivedImeShortcutGroups =
+                result == null ? Collections.emptyList() : result;
+
+        if (validateKeyboardShortcutHelperIconUri()) {
+            sanitiseShortcuts(mReceivedImeShortcutGroups);
+        }
+
+        maybeMergeAndShowKeyboardShortcuts();
+    }
+
+    static void sanitiseShortcuts(List<KeyboardShortcutGroup> shortcutGroups) {
+        for (KeyboardShortcutGroup group : shortcutGroups) {
+            for (KeyboardShortcutInfo info : group.getItems()) {
+                info.clearIcon();
+            }
+        }
+    }
+
     private void maybeMergeAndShowKeyboardShortcuts() {
         if (mReceivedAppShortcutGroups == null || mReceivedImeShortcutGroups == null) {
             return;
@@ -413,6 +475,14 @@
             mKeyboardShortcutsDialog.dismiss();
             mKeyboardShortcutsDialog = null;
         }
+        mHandlerThread.quit();
+        if (validateKeyboardShortcutHelperIconUri()) {
+            try {
+                ActivityManager.getService().unregisterUserSwitchObserver(mUserSwitchObserver);
+            } catch (RemoteException e) {
+                Log.e(TAG, "Could not unregister user switch observer", e);
+            }
+        }
     }
 
     private KeyboardShortcutGroup getSystemShortcuts() {
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 edd2961..99ce454 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
@@ -2843,14 +2843,10 @@
         mIsSystemChildExpanded = expanded;
     }
 
-    public void setLayoutListener(LayoutListener listener) {
+    public void setLayoutListener(@Nullable LayoutListener listener) {
         mLayoutListener = listener;
     }
 
-    public void removeListener() {
-        mLayoutListener = null;
-    }
-
     @Override
     protected void onLayout(boolean changed, int left, int top, int right, int bottom) {
         Trace.beginSection(appendTraceStyleTag("ExpNotRow#onLayout"));
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/NotificationMenuRow.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/NotificationMenuRow.java
index c10c09c..bdfbc4b 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/NotificationMenuRow.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/NotificationMenuRow.java
@@ -242,7 +242,7 @@
     public void onLayout() {
         mIconsPlaced = false; // Force icons to be re-placed
         setMenuLocation();
-        mParent.removeListener();
+        mParent.setLayoutListener(null);
     }
 
     private void createMenuViews(boolean resetState) {
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 a9d7cc0..84b7478 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
@@ -255,7 +255,8 @@
      * The raw amount of the overScroll on the bottom, which is not rubber-banded.
      */
     private float mOverScrolledBottomPixels;
-    private ListenerSet<Runnable> mStackHeightChangedListeners = new ListenerSet<>();
+    private final ListenerSet<Runnable> mStackHeightChangedListeners = new ListenerSet<>();
+    private final ListenerSet<Runnable> mHeadsUpHeightChangedListeners = new ListenerSet<>();
     private NotificationLogger.OnChildLocationsChangedListener mListener;
     private OnNotificationLocationsChangedListener mLocationsChangedListener;
     private OnOverscrollTopChangedListener mOverscrollTopChangedListener;
@@ -1114,6 +1115,28 @@
         mStackHeightChangedListeners.remove(runnable);
     }
 
+    private void notifyHeadsUpHeightChangedForView(View view) {
+        if (mTopHeadsUpRow == view) {
+            notifyHeadsUpHeightChangedListeners();
+        }
+    }
+
+    private void notifyHeadsUpHeightChangedListeners() {
+        for (Runnable listener : mHeadsUpHeightChangedListeners) {
+            listener.run();
+        }
+    }
+
+    @Override
+    public void addHeadsUpHeightChangedListener(@NonNull Runnable runnable) {
+        mHeadsUpHeightChangedListeners.addIfAbsent(runnable);
+    }
+
+    @Override
+    public void removeHeadsUpHeightChangedListener(@NonNull Runnable runnable) {
+        mHeadsUpHeightChangedListeners.remove(runnable);
+    }
+
     @Override
     protected void onLayout(boolean changed, int l, int t, int r, int b) {
         if (!mSuppressChildrenMeasureAndLayout) {
@@ -2444,6 +2467,11 @@
         return mScrollViewFields.getIntrinsicStackHeight();
     }
 
+    @Override
+    public int getTopHeadsUpHeight() {
+        return getTopHeadsUpPinnedHeight();
+    }
+
     /**
      * Calculate the gap height between two different views
      *
@@ -4193,12 +4221,14 @@
             requestAnimationOnViewResize(row);
         }
         requestChildrenUpdate();
+        notifyHeadsUpHeightChangedForView(view);
         mAnimateStackYForContentHeightChange = previouslyNeededAnimation;
     }
 
     void onChildHeightReset(ExpandableView view) {
         updateAnimationState(view);
         updateChronometerForChild(view);
+        notifyHeadsUpHeightChangedForView(view);
     }
 
     private void updateScrollPositionOnExpandInBottom(ExpandableView view) {
@@ -5573,6 +5603,7 @@
      */
     public void setTopHeadsUpRow(@Nullable ExpandableNotificationRow topHeadsUpRow) {
         mTopHeadsUpRow = topHeadsUpRow;
+        notifyHeadsUpHeightChangedListeners();
     }
 
     public boolean getIsExpanded() {
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/data/repository/NotificationViewHeightRepository.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/data/repository/NotificationViewHeightRepository.kt
index 463c631..f6d9351 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/data/repository/NotificationViewHeightRepository.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/data/repository/NotificationViewHeightRepository.kt
@@ -27,9 +27,6 @@
 @SysUISingleton
 class NotificationViewHeightRepository @Inject constructor() {
 
-    /** The height in px of the current heads up notification. */
-    val headsUpHeight = MutableStateFlow(0f)
-
     /**
      * The amount in px that the notification stack should scroll due to internal expansion. This
      * should only happen when a notification expansion hits the bottom of the screen, so it is
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/domain/interactor/NotificationStackAppearanceInteractor.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/domain/interactor/NotificationStackAppearanceInteractor.kt
index e7acbe3..afcf3ae 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/domain/interactor/NotificationStackAppearanceInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/domain/interactor/NotificationStackAppearanceInteractor.kt
@@ -65,9 +65,6 @@
             }
             .distinctUntilChanged()
 
-    /** The height in px of the contents of the HUN. */
-    val headsUpHeight: StateFlow<Float> = viewHeightRepository.headsUpHeight.asStateFlow()
-
     /** The alpha of the Notification Stack for the brightness mirror */
     val alphaForBrightnessMirror: StateFlow<Float> =
         placeholderRepository.alphaForBrightnessMirror.asStateFlow()
@@ -110,11 +107,6 @@
         placeholderRepository.shadeScrimBounds.value = bounds
     }
 
-    /** Sets the height of heads up notification. */
-    fun setHeadsUpHeight(height: Float) {
-        viewHeightRepository.headsUpHeight.value = height
-    }
-
     /** Sets whether the notification stack is scrolled to the top. */
     fun setScrolledToTop(scrolledToTop: Boolean) {
         placeholderRepository.scrolledToTop.value = scrolledToTop
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/view/NotificationScrollView.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/view/NotificationScrollView.kt
index 14b882f..eaaa9a1 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/view/NotificationScrollView.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/view/NotificationScrollView.kt
@@ -32,6 +32,9 @@
      */
     val intrinsicStackHeight: Int
 
+    /** Height in pixels required to display the top HeadsUp Notification. */
+    val topHeadsUpHeight: Int
+
     /**
      * Since this is an interface rather than a literal View, this provides cast-like access to the
      * underlying view.
@@ -72,9 +75,18 @@
     /** Sets whether the view is displayed in doze mode. */
     fun setDozing(dozing: Boolean)
 
-    /** Sets a listener to be notified, when the stack height might have changed. */
+    /** Adds a listener to be notified, when the stack height might have changed. */
     fun addStackHeightChangedListener(runnable: Runnable)
 
     /** @see addStackHeightChangedListener */
     fun removeStackHeightChangedListener(runnable: Runnable)
+
+    /**
+     * Adds a listener to be notified, when the height of the top heads up notification might have
+     * changed.
+     */
+    fun addHeadsUpHeightChangedListener(runnable: Runnable)
+
+    /** @see addHeadsUpHeightChangedListener */
+    fun removeHeadsUpHeightChangedListener(runnable: Runnable)
 }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewbinder/NotificationScrollViewBinder.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewbinder/NotificationScrollViewBinder.kt
index 622d8e7..fd08e89 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewbinder/NotificationScrollViewBinder.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewbinder/NotificationScrollViewBinder.kt
@@ -80,7 +80,6 @@
 
         launch { viewModel.maxAlpha.collect { view.setMaxAlpha(it) } }
         launch { viewModel.scrolledToTop.collect { view.setScrolledToTop(it) } }
-        launch { viewModel.headsUpTop.collect { view.setHeadsUpTop(it) } }
         launch { viewModel.expandFraction.collect { view.setExpandFraction(it.coerceIn(0f, 1f)) } }
         launch { viewModel.isScrollable.collect { view.setScrollingEnabled(it) } }
         launch { viewModel.isDozing.collect { isDozing -> view.setDozing(isDozing) } }
@@ -88,11 +87,9 @@
         launchAndDispose {
             view.setSyntheticScrollConsumer(viewModel.syntheticScrollConsumer)
             view.setCurrentGestureOverscrollConsumer(viewModel.currentGestureOverscrollConsumer)
-            view.setHeadsUpHeightConsumer(viewModel.headsUpHeightConsumer)
             DisposableHandle {
                 view.setSyntheticScrollConsumer(null)
                 view.setCurrentGestureOverscrollConsumer(null)
-                view.setHeadsUpHeightConsumer(null)
             }
         }
     }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/NotificationScrollViewModel.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/NotificationScrollViewModel.kt
index c2ce114..a99fbfc 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/NotificationScrollViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/NotificationScrollViewModel.kt
@@ -150,8 +150,6 @@
      */
     val currentGestureOverscrollConsumer: (Boolean) -> Unit =
         stackAppearanceInteractor::setCurrentGestureOverscroll
-    /** Receives the height of the heads up notification. */
-    val headsUpHeightConsumer: (Float) -> Unit = stackAppearanceInteractor::setHeadsUpHeight
 
     /** Whether the notification stack is scrollable or not. */
     val isScrollable: Flow<Boolean> =
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/NotificationsPlaceholderViewModel.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/NotificationsPlaceholderViewModel.kt
index 97b86e3..ea33be0 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/NotificationsPlaceholderViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/NotificationsPlaceholderViewModel.kt
@@ -29,7 +29,6 @@
 import com.android.systemui.util.kotlin.FlowDumperImpl
 import javax.inject.Inject
 import kotlinx.coroutines.flow.Flow
-import kotlinx.coroutines.flow.StateFlow
 
 /**
  * ViewModel used by the Notification placeholders inside the scene container to update the
@@ -74,9 +73,6 @@
     val shadeScrimRounding: Flow<ShadeScrimRounding> =
         interactor.shadeScrimRounding.dumpWhileCollecting("shadeScrimRounding")
 
-    /** The height in px of the contents of the HUN. */
-    val headsUpHeight: StateFlow<Float> = interactor.headsUpHeight.dumpValue("headsUpHeight")
-
     /**
      * The amount [0-1] that the shade or quick settings has been opened. At 0, the shade is closed;
      * at 1, either the shade or quick settings is open.
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/KeyboardShortcutListSearchTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/KeyboardShortcutListSearchTest.java
index 22c9e45..6985a27 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/KeyboardShortcutListSearchTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/KeyboardShortcutListSearchTest.java
@@ -20,14 +20,21 @@
 
 import static org.mockito.Mockito.any;
 import static org.mockito.Mockito.anyInt;
+import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
+import android.graphics.drawable.Icon;
+import android.os.Handler;
+import android.platform.test.annotations.EnableFlags;
+import android.view.KeyboardShortcutGroup;
+import android.view.KeyboardShortcutInfo;
 import android.view.WindowManager;
 
 import androidx.test.filters.SmallTest;
 import androidx.test.runner.AndroidJUnit4;
 
+import com.android.systemui.Flags;
 import com.android.systemui.SysuiTestCase;
 
 import com.google.android.material.bottomsheet.BottomSheetDialog;
@@ -36,10 +43,14 @@
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.mockito.ArgumentCaptor;
 import org.mockito.Mock;
 import org.mockito.junit.MockitoJUnit;
 import org.mockito.junit.MockitoRule;
 
+import java.util.Arrays;
+import java.util.Collections;
+
 @SmallTest
 @RunWith(AndroidJUnit4.class)
 public class KeyboardShortcutListSearchTest extends SysuiTestCase {
@@ -51,6 +62,7 @@
 
     @Mock private BottomSheetDialog mBottomSheetDialog;
     @Mock WindowManager mWindowManager;
+    @Mock Handler mHandler;
 
     @Before
     public void setUp() {
@@ -58,6 +70,7 @@
         mKeyboardShortcutListSearch.sInstance = mKeyboardShortcutListSearch;
         mKeyboardShortcutListSearch.mKeyboardShortcutsBottomSheetDialog = mBottomSheetDialog;
         mKeyboardShortcutListSearch.mContext = mContext;
+        mKeyboardShortcutListSearch.mBackgroundHandler = mHandler;
     }
 
     @Test
@@ -78,4 +91,59 @@
         verify(mWindowManager).requestAppKeyboardShortcuts(any(), anyInt());
         verify(mWindowManager).requestImeKeyboardShortcuts(any(), anyInt());
     }
+
+    @Test
+    @EnableFlags(Flags.FLAG_VALIDATE_KEYBOARD_SHORTCUT_HELPER_ICON_URI)
+    public void requestAppKeyboardShortcuts_callback_sanitisesIcons() {
+        KeyboardShortcutGroup group = createKeyboardShortcutGroupForIconTests();
+
+        mKeyboardShortcutListSearch.toggle(mContext, DEVICE_ID);
+
+        ArgumentCaptor<WindowManager.KeyboardShortcutsReceiver> callbackCaptor =
+                ArgumentCaptor.forClass(WindowManager.KeyboardShortcutsReceiver.class);
+        ArgumentCaptor<Runnable> handlerRunnableCaptor = ArgumentCaptor.forClass(Runnable.class);
+        verify(mWindowManager).requestAppKeyboardShortcuts(callbackCaptor.capture(), anyInt());
+        callbackCaptor.getValue().onKeyboardShortcutsReceived(Collections.singletonList(group));
+        verify(mHandler).post(handlerRunnableCaptor.capture());
+        handlerRunnableCaptor.getValue().run();
+
+        verify(group.getItems().get(0)).clearIcon();
+        verify(group.getItems().get(1)).clearIcon();
+    }
+
+    @Test
+    @EnableFlags(Flags.FLAG_VALIDATE_KEYBOARD_SHORTCUT_HELPER_ICON_URI)
+    public void requestImeKeyboardShortcuts_callback_sanitisesIcons() {
+        KeyboardShortcutGroup group = createKeyboardShortcutGroupForIconTests();
+
+        mKeyboardShortcutListSearch.toggle(mContext, DEVICE_ID);
+
+        ArgumentCaptor<WindowManager.KeyboardShortcutsReceiver> callbackCaptor =
+                ArgumentCaptor.forClass(WindowManager.KeyboardShortcutsReceiver.class);
+        ArgumentCaptor<Runnable> handlerRunnableCaptor = ArgumentCaptor.forClass(Runnable.class);
+        verify(mWindowManager).requestImeKeyboardShortcuts(callbackCaptor.capture(), anyInt());
+        callbackCaptor.getValue().onKeyboardShortcutsReceived(Collections.singletonList(group));
+        verify(mHandler).post(handlerRunnableCaptor.capture());
+        handlerRunnableCaptor.getValue().run();
+
+        verify(group.getItems().get(0)).clearIcon();
+        verify(group.getItems().get(1)).clearIcon();
+
+    }
+
+    private KeyboardShortcutGroup createKeyboardShortcutGroupForIconTests() {
+        Icon icon = mock(Icon.class);
+
+        KeyboardShortcutInfo info1 = mock(KeyboardShortcutInfo.class);
+        KeyboardShortcutInfo info2 = mock(KeyboardShortcutInfo.class);
+        when(info1.getIcon()).thenReturn(icon);
+        when(info2.getIcon()).thenReturn(icon);
+        when(info1.getLabel()).thenReturn("label");
+        when(info2.getLabel()).thenReturn("label");
+
+        KeyboardShortcutGroup group = new KeyboardShortcutGroup("label",
+                Arrays.asList(new KeyboardShortcutInfo[]{ info1, info2}));
+        group.setPackageName("com.example");
+        return group;
+    }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/KeyboardShortcutsTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/KeyboardShortcutsTest.java
index a3ecde0..2b3f139 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/KeyboardShortcutsTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/KeyboardShortcutsTest.java
@@ -20,25 +20,36 @@
 
 import static org.mockito.Mockito.any;
 import static org.mockito.Mockito.anyInt;
+import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
 import android.app.Dialog;
+import android.graphics.drawable.Icon;
+import android.os.Handler;
+import android.platform.test.annotations.EnableFlags;
+import android.view.KeyboardShortcutGroup;
+import android.view.KeyboardShortcutInfo;
 import android.view.WindowManager;
 
 import androidx.test.filters.SmallTest;
 import androidx.test.runner.AndroidJUnit4;
 
+import com.android.systemui.Flags;
 import com.android.systemui.SysuiTestCase;
 
 import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.mockito.ArgumentCaptor;
 import org.mockito.Mock;
 import org.mockito.junit.MockitoJUnit;
 import org.mockito.junit.MockitoRule;
 
+import java.util.Arrays;
+import java.util.Collections;
+
 @SmallTest
 @RunWith(AndroidJUnit4.class)
 public class KeyboardShortcutsTest extends SysuiTestCase {
@@ -50,6 +61,7 @@
 
     @Mock private Dialog mDialog;
     @Mock WindowManager mWindowManager;
+    @Mock Handler mHandler;
 
     @Before
     public void setUp() {
@@ -57,6 +69,7 @@
         mKeyboardShortcuts.sInstance = mKeyboardShortcuts;
         mKeyboardShortcuts.mKeyboardShortcutsDialog = mDialog;
         mKeyboardShortcuts.mContext = mContext;
+        mKeyboardShortcuts.mBackgroundHandler = mHandler;
     }
 
     @Test
@@ -77,4 +90,78 @@
         verify(mWindowManager).requestAppKeyboardShortcuts(any(), anyInt());
         verify(mWindowManager).requestImeKeyboardShortcuts(any(), anyInt());
     }
+
+    @Test
+    public void sanitiseShortcuts_clearsIcons() {
+        KeyboardShortcutGroup group = createKeyboardShortcutGroupForIconTests();
+
+        KeyboardShortcuts.sanitiseShortcuts(Collections.singletonList(group));
+
+        verify(group.getItems().get(0)).clearIcon();
+        verify(group.getItems().get(1)).clearIcon();
+    }
+
+    @Test
+    public void sanitiseShortcuts_nullPackage_clearsIcons() {
+        KeyboardShortcutGroup group = createKeyboardShortcutGroupForIconTests();
+        group.setPackageName(null);
+
+        KeyboardShortcuts.sanitiseShortcuts(Collections.singletonList(group));
+
+        verify(group.getItems().get(0)).clearIcon();
+        verify(group.getItems().get(1)).clearIcon();
+    }
+
+    @Test
+    @EnableFlags(Flags.FLAG_VALIDATE_KEYBOARD_SHORTCUT_HELPER_ICON_URI)
+    public void requestAppKeyboardShortcuts_callback_sanitisesIcons() {
+        KeyboardShortcutGroup group = createKeyboardShortcutGroupForIconTests();
+
+        mKeyboardShortcuts.toggle(mContext, DEVICE_ID);
+
+        ArgumentCaptor<WindowManager.KeyboardShortcutsReceiver> callbackCaptor =
+                ArgumentCaptor.forClass(WindowManager.KeyboardShortcutsReceiver.class);
+        ArgumentCaptor<Runnable> handlerRunnableCaptor = ArgumentCaptor.forClass(Runnable.class);
+        verify(mWindowManager).requestAppKeyboardShortcuts(callbackCaptor.capture(), anyInt());
+        callbackCaptor.getValue().onKeyboardShortcutsReceived(Collections.singletonList(group));
+        verify(mHandler).post(handlerRunnableCaptor.capture());
+        handlerRunnableCaptor.getValue().run();
+
+        verify(group.getItems().get(0)).clearIcon();
+        verify(group.getItems().get(1)).clearIcon();
+    }
+
+    @Test
+    @EnableFlags(Flags.FLAG_VALIDATE_KEYBOARD_SHORTCUT_HELPER_ICON_URI)
+    public void requestImeKeyboardShortcuts_callback_sanitisesIcons() {
+        KeyboardShortcutGroup group = createKeyboardShortcutGroupForIconTests();
+
+        mKeyboardShortcuts.toggle(mContext, DEVICE_ID);
+
+        ArgumentCaptor<WindowManager.KeyboardShortcutsReceiver> callbackCaptor =
+                ArgumentCaptor.forClass(WindowManager.KeyboardShortcutsReceiver.class);
+        ArgumentCaptor<Runnable> handlerRunnableCaptor = ArgumentCaptor.forClass(Runnable.class);
+        verify(mWindowManager).requestImeKeyboardShortcuts(callbackCaptor.capture(), anyInt());
+        callbackCaptor.getValue().onKeyboardShortcutsReceived(Collections.singletonList(group));
+        verify(mHandler).post(handlerRunnableCaptor.capture());
+        handlerRunnableCaptor.getValue().run();
+
+        verify(group.getItems().get(0)).clearIcon();
+        verify(group.getItems().get(1)).clearIcon();
+
+    }
+
+    private KeyboardShortcutGroup createKeyboardShortcutGroupForIconTests() {
+        Icon icon = mock(Icon.class);
+
+        KeyboardShortcutInfo info1 = mock(KeyboardShortcutInfo.class);
+        KeyboardShortcutInfo info2 = mock(KeyboardShortcutInfo.class);
+        when(info1.getIcon()).thenReturn(icon);
+        when(info2.getIcon()).thenReturn(icon);
+
+        KeyboardShortcutGroup group = new KeyboardShortcutGroup("label",
+                Arrays.asList(new KeyboardShortcutInfo[]{ info1, info2}));
+        group.setPackageName("com.example");
+        return group;
+    }
 }
diff --git a/services/core/java/com/android/server/locales/LocaleManagerBackupHelper.java b/services/core/java/com/android/server/locales/LocaleManagerBackupHelper.java
index d932bd4..563f93e 100644
--- a/services/core/java/com/android/server/locales/LocaleManagerBackupHelper.java
+++ b/services/core/java/com/android/server/locales/LocaleManagerBackupHelper.java
@@ -434,7 +434,7 @@
                         ATTR_PACKAGE_NAME);
                 String languageTags = parser.getAttributeValue(/* namespace= */ null, ATTR_LOCALES);
                 boolean delegateSelector = parser.getAttributeBoolean(/* namespace= */ null,
-                        ATTR_DELEGATE_SELECTOR);
+                        ATTR_DELEGATE_SELECTOR, false);
 
                 if (!TextUtils.isEmpty(packageName) && !TextUtils.isEmpty(languageTags)) {
                     LocalesInfo localesInfo = new LocalesInfo(languageTags, delegateSelector);
diff --git a/services/core/java/com/android/server/wm/Task.java b/services/core/java/com/android/server/wm/Task.java
index e802527..cd22591 100644
--- a/services/core/java/com/android/server/wm/Task.java
+++ b/services/core/java/com/android/server/wm/Task.java
@@ -3624,13 +3624,15 @@
         // If the developer has persist a different configuration, we need to override it to the
         // starting window because persisted configuration does not effect to Task.
         info.taskInfo.configuration.setTo(activity.getConfiguration());
-        final ActivityRecord topFullscreenActivity = getTopFullscreenActivity();
-        if (topFullscreenActivity != null) {
-            final WindowState mainWindow = topFullscreenActivity.findMainWindow(false);
-            if (mainWindow != null) {
-                info.topOpaqueWindowInsetsState =
-                        mainWindow.getInsetsStateWithVisibilityOverride();
-                info.topOpaqueWindowLayoutParams = mainWindow.getAttrs();
+        if (!Flags.drawSnapshotAspectRatioMatch()) {
+            final ActivityRecord topFullscreenActivity = getTopFullscreenActivity();
+            if (topFullscreenActivity != null) {
+                final WindowState mainWindow = topFullscreenActivity.findMainWindow(false);
+                if (mainWindow != null) {
+                    info.topOpaqueWindowInsetsState =
+                            mainWindow.getInsetsStateWithVisibilityOverride();
+                    info.topOpaqueWindowLayoutParams = mainWindow.getAttrs();
+                }
             }
         }
         return info;