Merge "When doing the special case migration where cells shift up, only modify workspace items" into main
diff --git a/PREUPLOAD.cfg b/PREUPLOAD.cfg
index 9051ca8..6948133 100644
--- a/PREUPLOAD.cfg
+++ b/PREUPLOAD.cfg
@@ -10,4 +10,3 @@
 [Hook Scripts]
 checkstyle_hook = ${REPO_ROOT}/prebuilts/checkstyle/checkstyle.py --config_xml tools/checkstyle.xml --sha ${PREUPLOAD_COMMIT}
 
-flag_hook = ${REPO_ROOT}/frameworks/base/packages/SystemUI/flag_check.py --msg=${PREUPLOAD_COMMIT_MESSAGE} --files=${PREUPLOAD_FILES} --project=${REPO_PATH}
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java
index ebd4ee5..17fb959 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java
@@ -91,6 +91,7 @@
     private static final boolean DEBUG = false;
     // TODO(b/382378283) remove all logs with this tag
     public static final String NULL_TASKBAR_ROOT_LAYOUT_TAG = "b/382378283";
+    public static final String ILLEGAL_ARGUMENT_WM_ADD_VIEW = "b/391653300";
 
     /**
      * All the configurations which do not initiate taskbar recreation.
@@ -306,6 +307,7 @@
                 int configDiff = mOldConfig.diff(newConfig) & ~SKIP_RECREATE_CONFIG_CHANGES;
 
                 if ((configDiff & ActivityInfo.CONFIG_UI_MODE) != 0) {
+                    Log.d(ILLEGAL_ARGUMENT_WM_ADD_VIEW, "onConfigurationChanged: theme changed");
                     // Only recreate for theme changes, not other UI mode changes such as docking.
                     int oldUiNightMode = (mOldConfig.uiMode & Configuration.UI_MODE_NIGHT_MASK);
                     int newUiNightMode = (newConfig.uiMode & Configuration.UI_MODE_NIGHT_MASK);
@@ -360,6 +362,7 @@
     }
 
     private void destroyTaskbarForDisplay(int displayId) {
+        Log.d(ILLEGAL_ARGUMENT_WM_ADD_VIEW, "destroyTaskbarForDisplay: " + displayId);
         TaskbarActivityContext taskbar = getTaskbarForDisplay(displayId);
         debugWhyTaskbarNotDestroyed("destroyTaskbarForDisplay: " + taskbar, displayId);
         if (taskbar != null) {
@@ -522,6 +525,7 @@
     private void recreateTaskbarForDisplay(int displayId) {
         Trace.beginSection("recreateTaskbar");
         try {
+            Log.d(ILLEGAL_ARGUMENT_WM_ADD_VIEW, "recreateTaskbarForDisplay: " + displayId);
             // TODO: make this code display specific
             DeviceProfile dp = mUserUnlocked ?
                     LauncherAppState.getIDP(mWindowContext).getDeviceProfile(mWindowContext) : null;
@@ -711,9 +715,9 @@
 
     /**
      * Signal from SysUI indicating that a non-mirroring display was just connected to the
-     * primary device.
+     * primary device or a previously mirroring display is switched to extended mode.
      */
-    public void onDisplayReady(int displayId) {
+    public void onDisplayAddSystemDecorations(int displayId) {
     }
 
     /**
@@ -794,9 +798,15 @@
         }
 
         if (!isTaskbarRootLayoutAddedForDisplay(displayId)) {
-            mWindowManager.addView(getTaskbarRootLayoutForDisplay(displayId),
-                    taskbar.getWindowLayoutParams());
-            mAddedRootLayouts.put(displayId, true);
+            FrameLayout rootLayout = getTaskbarRootLayoutForDisplay(displayId);
+            if (rootLayout != null) {
+                mWindowManager.addView(getTaskbarRootLayoutForDisplay(displayId),
+                        taskbar.getWindowLayoutParams());
+                mAddedRootLayouts.put(displayId, true);
+            } else {
+                Log.d(ILLEGAL_ARGUMENT_WM_ADD_VIEW,
+                        "addTaskbarRootViewToWindow - root layout null | displayId=" + displayId);
+            }
         } else {
             Log.d(NULL_TASKBAR_ROOT_LAYOUT_TAG,
                     "addTaskbarRootViewToWindow - root layout already added | displayId="
@@ -805,6 +815,7 @@
     }
 
     private void removeTaskbarRootViewFromWindow(int displayId) {
+        Log.d(ILLEGAL_ARGUMENT_WM_ADD_VIEW, "removeTaskbarRootViewFromWindow: " + displayId);
         FrameLayout rootLayout = getTaskbarRootLayoutForDisplay(displayId);
         if (!enableTaskbarNoRecreate() || rootLayout == null) {
             return;
@@ -880,6 +891,7 @@
      * @param displayId The ID of the display for which to create the taskbar root layout.
      */
     private void createTaskbarRootLayout(int displayId) {
+        Log.d(ILLEGAL_ARGUMENT_WM_ADD_VIEW, "createTaskbarRootLayout: " + displayId);
         FrameLayout newTaskbarRootLayout = new FrameLayout(mWindowContext) {
             @Override
             public boolean dispatchTouchEvent(MotionEvent ev) {
@@ -907,6 +919,7 @@
      * @return The taskbar root layout {@link FrameLayout} for a given display or {@code null}.
      */
     private FrameLayout getTaskbarRootLayoutForDisplay(int displayId) {
+        Log.d(ILLEGAL_ARGUMENT_WM_ADD_VIEW, "getTaskbarRootLayoutForDisplay: " + displayId);
         FrameLayout frameLayout = mRootLayouts.get(displayId);
         if (frameLayout != null) {
             return frameLayout;
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarModelCallbacks.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarModelCallbacks.java
index f905c5f..6815f97 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarModelCallbacks.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarModelCallbacks.java
@@ -25,7 +25,6 @@
 import com.android.launcher3.model.BgDataModel.FixedContainerItems;
 import com.android.launcher3.model.data.AppInfo;
 import com.android.launcher3.model.data.ItemInfo;
-import com.android.launcher3.model.data.WorkspaceItemInfo;
 import com.android.launcher3.util.ComponentKey;
 import com.android.launcher3.util.IntArray;
 import com.android.launcher3.util.IntSet;
@@ -39,9 +38,9 @@
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
-import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import java.util.function.Predicate;
 
 /**
@@ -114,15 +113,9 @@
         return modified;
     }
 
-
     @Override
-    public void bindWorkspaceItemsChanged(List<WorkspaceItemInfo> updated) {
-        updateWorkspaceItems(updated, mContext);
-    }
-
-    @Override
-    public void bindRestoreItemsChange(HashSet<ItemInfo> updates) {
-        updateRestoreItems(updates, mContext);
+    public void bindItemsUpdated(Set<ItemInfo> updates) {
+        updateContainerItems(updates, mContext);
     }
 
     @Override
diff --git a/quickstep/src/com/android/launcher3/uioverrides/PredictedAppIcon.java b/quickstep/src/com/android/launcher3/uioverrides/PredictedAppIcon.java
index a85e5e0..36a4865 100644
--- a/quickstep/src/com/android/launcher3/uioverrides/PredictedAppIcon.java
+++ b/quickstep/src/com/android/launcher3/uioverrides/PredictedAppIcon.java
@@ -142,7 +142,7 @@
         int shadowSize = context.getResources().getDimensionPixelSize(
                 R.dimen.blur_size_thin_outline);
         mShadowFilter = new BlurMaskFilter(shadowSize, BlurMaskFilter.Blur.OUTER);
-        mShapePath = IconShape.INSTANCE.get(context).getShapeOverridePath(mNormalizedIconSize);
+        mShapePath = IconShape.INSTANCE.get(context).getShape().getPath(mNormalizedIconSize);
     }
 
     @Override
@@ -214,7 +214,7 @@
         boolean animate = shouldAnimateIconChange(info);
         Drawable oldIcon = getIcon();
         int oldPlateColor = mPlateColor.currentColor;
-        applyFromWorkspaceItem(info, null);
+        applyFromWorkspaceItem(info);
 
         setContentDescription(
                 mIsPinned ? info.contentDescription :
diff --git a/quickstep/src/com/android/launcher3/uioverrides/SystemApiWrapper.kt b/quickstep/src/com/android/launcher3/uioverrides/SystemApiWrapper.kt
index 2e2d7cc..d097dba 100644
--- a/quickstep/src/com/android/launcher3/uioverrides/SystemApiWrapper.kt
+++ b/quickstep/src/com/android/launcher3/uioverrides/SystemApiWrapper.kt
@@ -192,4 +192,6 @@
 
     override fun getApplicationInfoHash(appInfo: ApplicationInfo): String =
         (appInfo.sourceDir?.hashCode() ?: 0).toString() + " " + appInfo.longVersionCode
+
+    override fun getRoundIconRes(appInfo: ApplicationInfo) = appInfo.roundIconRes
 }
diff --git a/quickstep/src/com/android/launcher3/uioverrides/touchcontrollers/QuickSwitchTouchController.java b/quickstep/src/com/android/launcher3/uioverrides/touchcontrollers/QuickSwitchTouchController.java
index d673720..2b296c8 100644
--- a/quickstep/src/com/android/launcher3/uioverrides/touchcontrollers/QuickSwitchTouchController.java
+++ b/quickstep/src/com/android/launcher3/uioverrides/touchcontrollers/QuickSwitchTouchController.java
@@ -151,7 +151,7 @@
             int sysuiFlags = 0;
             TaskView tv = mOverviewPanel.getFirstTaskView();
             if (tv != null) {
-                sysuiFlags = tv.getTaskContainers().getFirst().getSysUiStatusNavFlags();
+                sysuiFlags = tv.getSysUiStatusNavFlags();
             }
             mLauncher.getSystemUiController().updateUiState(UI_STATE_FULLSCREEN_TASK, sysuiFlags);
         } else {
diff --git a/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java b/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
index 8b76ce9..78e543c 100644
--- a/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
+++ b/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
@@ -932,7 +932,7 @@
             TaskView runningTask = mRecentsView.getRunningTaskView();
             TaskView centermostTask = mRecentsView.getTaskViewNearestToCenterOfScreen();
             int centermostTaskFlags = centermostTask == null ? 0
-                    : centermostTask.getTaskContainers().getFirst().getSysUiStatusNavFlags();
+                    : centermostTask.getSysUiStatusNavFlags();
             boolean swipeUpThresholdPassed = windowProgress > 1 - UPDATE_SYSUI_FLAGS_THRESHOLD;
             boolean quickswitchThresholdPassed = centermostTask != runningTask;
 
@@ -955,7 +955,7 @@
 
     @Override
     public void onRecentsAnimationStart(RecentsAnimationController controller,
-            RecentsAnimationTargets targets, TransitionInfo transitionInfo) {
+            RecentsAnimationTargets targets, @Nullable TransitionInfo transitionInfo) {
         super.onRecentsAnimationStart(controller, targets, transitionInfo);
         if (targets.hasDesktopTasks(mContext)) {
             mRemoteTargetHandles = mTargetGluer.assignTargetsForDesktop(targets, transitionInfo);
diff --git a/quickstep/src/com/android/quickstep/RecentsAnimationCallbacks.java b/quickstep/src/com/android/quickstep/RecentsAnimationCallbacks.java
index 87bf81c..c6b858b 100644
--- a/quickstep/src/com/android/quickstep/RecentsAnimationCallbacks.java
+++ b/quickstep/src/com/android/quickstep/RecentsAnimationCallbacks.java
@@ -22,6 +22,7 @@
 
 import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
 
+import android.annotation.Nullable;
 import android.graphics.Rect;
 import android.os.Bundle;
 import android.util.ArraySet;
@@ -103,7 +104,7 @@
             RemoteAnimationTarget[] appTargets,
             RemoteAnimationTarget[] wallpaperTargets,
             Rect homeContentInsets, Rect minimizedHomeBounds, Bundle extras,
-            TransitionInfo transitionInfo) {
+            @Nullable TransitionInfo transitionInfo) {
         long appCount = Arrays.stream(appTargets)
                 .filter(app -> app.mode == MODE_CLOSING)
                 .count();
@@ -207,7 +208,7 @@
      */
     public interface RecentsAnimationListener {
         default void onRecentsAnimationStart(RecentsAnimationController controller,
-                RecentsAnimationTargets targets, TransitionInfo transitionInfo) {}
+                RecentsAnimationTargets targets, @Nullable TransitionInfo transitionInfo) {}
 
         /**
          * Callback from the system when the recents animation is canceled. {@param thumbnailData}
diff --git a/quickstep/src/com/android/quickstep/RemoteTargetGluer.java b/quickstep/src/com/android/quickstep/RemoteTargetGluer.java
index ef63b9b..f96bbcb 100644
--- a/quickstep/src/com/android/quickstep/RemoteTargetGluer.java
+++ b/quickstep/src/com/android/quickstep/RemoteTargetGluer.java
@@ -217,7 +217,7 @@
      * transform params per app in {@code targets.apps} list.
      */
     public RemoteTargetHandle[] assignTargetsForDesktop(
-            RemoteAnimationTargets targets, TransitionInfo transitionInfo) {
+            RemoteAnimationTargets targets, @Nullable TransitionInfo transitionInfo) {
         resizeRemoteTargetHandles(targets);
 
         for (int i = 0; i < mRemoteTargetHandles.length; i++) {
diff --git a/quickstep/src/com/android/quickstep/SimpleOrientationTouchTransformer.java b/quickstep/src/com/android/quickstep/SimpleOrientationTouchTransformer.java
index 5264643..d2a491d 100644
--- a/quickstep/src/com/android/quickstep/SimpleOrientationTouchTransformer.java
+++ b/quickstep/src/com/android/quickstep/SimpleOrientationTouchTransformer.java
@@ -22,35 +22,35 @@
 import android.content.Context;
 import android.view.MotionEvent;
 
+import com.android.launcher3.dagger.ApplicationContext;
+import com.android.launcher3.dagger.LauncherAppSingleton;
+import com.android.launcher3.util.DaggerSingletonObject;
+import com.android.launcher3.util.DaggerSingletonTracker;
 import com.android.launcher3.util.DisplayController;
-import com.android.launcher3.util.MainThreadInitializedObject;
-import com.android.launcher3.util.SafeCloseable;
+import com.android.quickstep.dagger.QuickstepBaseAppComponent;
 
+import javax.inject.Inject;
+
+@LauncherAppSingleton
 public class SimpleOrientationTouchTransformer implements
-        DisplayController.DisplayInfoChangeListener, SafeCloseable {
+        DisplayController.DisplayInfoChangeListener {
 
-    public static final MainThreadInitializedObject<SimpleOrientationTouchTransformer> INSTANCE =
-            new MainThreadInitializedObject<>(SimpleOrientationTouchTransformer::new);
+    public static final DaggerSingletonObject<SimpleOrientationTouchTransformer> INSTANCE =
+            new DaggerSingletonObject<>(
+                    QuickstepBaseAppComponent::getSimpleOrientationTouchTransformer);
 
-    private final Context mContext;
     private OrientationRectF mOrientationRectF;
     private OrientationRectF mTouchingOrientationRectF;
     private int mViewRotation;
 
-    public SimpleOrientationTouchTransformer(Context context) {
-        this(context, DisplayController.INSTANCE.get(context));
-    }
-
-    @androidx.annotation.VisibleForTesting
-    public SimpleOrientationTouchTransformer(Context context, DisplayController displayController) {
-        mContext = context;
+    @Inject
+    public SimpleOrientationTouchTransformer(@ApplicationContext Context context,
+            DisplayController displayController,
+            DaggerSingletonTracker tracker) {
         displayController.addChangeListener(this);
-        onDisplayInfoChanged(context, displayController.getInfo(), CHANGE_ALL);
-    }
+        tracker.addCloseable(() -> displayController.removeChangeListener(this));
 
-    @Override
-    public void close() {
-        DisplayController.INSTANCE.get(mContext).removeChangeListener(this);
+        onDisplayInfoChanged(context, displayController.getInfo(), CHANGE_ALL);
     }
 
     @Override
diff --git a/quickstep/src/com/android/quickstep/TaskAnimationManager.java b/quickstep/src/com/android/quickstep/TaskAnimationManager.java
index 1fd7211..63b8aaf 100644
--- a/quickstep/src/com/android/quickstep/TaskAnimationManager.java
+++ b/quickstep/src/com/android/quickstep/TaskAnimationManager.java
@@ -156,7 +156,7 @@
         mCallbacks.addListener(new RecentsAnimationCallbacks.RecentsAnimationListener() {
             @Override
             public void onRecentsAnimationStart(RecentsAnimationController controller,
-                    RecentsAnimationTargets targets, TransitionInfo transitionInfo) {
+                    RecentsAnimationTargets targets, @Nullable TransitionInfo transitionInfo) {
                 if (enableHandleDelayedGestureCallbacks() && mRecentsAnimationStartPending) {
                     ActiveGestureProtoLogProxy.logStartRecentsAnimationCallback(
                             "onRecentsAnimationStart");
diff --git a/quickstep/src/com/android/quickstep/TouchInteractionService.java b/quickstep/src/com/android/quickstep/TouchInteractionService.java
index 516b24c..51e59ff 100644
--- a/quickstep/src/com/android/quickstep/TouchInteractionService.java
+++ b/quickstep/src/com/android/quickstep/TouchInteractionService.java
@@ -301,10 +301,10 @@
 
         @BinderThread
         @Override
-        public void onDisplayReady(int displayId) {
+        public void onDisplayAddSystemDecorations(int displayId) {
             MAIN_EXECUTOR.execute(() -> executeForTouchInteractionService(
-                    tis -> executeForTaskbarManager(
-                            taskbarManager -> taskbarManager.onDisplayReady(displayId))));
+                    tis -> executeForTaskbarManager(taskbarManager ->
+                            taskbarManager.onDisplayAddSystemDecorations(displayId))));
         }
 
         @BinderThread
diff --git a/quickstep/src/com/android/quickstep/dagger/QuickstepBaseAppComponent.java b/quickstep/src/com/android/quickstep/dagger/QuickstepBaseAppComponent.java
index c875251..adc45ae 100644
--- a/quickstep/src/com/android/quickstep/dagger/QuickstepBaseAppComponent.java
+++ b/quickstep/src/com/android/quickstep/dagger/QuickstepBaseAppComponent.java
@@ -23,6 +23,7 @@
 import com.android.quickstep.OverviewComponentObserver;
 import com.android.quickstep.RecentsAnimationDeviceState;
 import com.android.quickstep.RotationTouchHelper;
+import com.android.quickstep.SimpleOrientationTouchTransformer;
 import com.android.quickstep.SystemUiProxy;
 import com.android.quickstep.TopTaskTracker;
 import com.android.quickstep.fallback.window.RecentsDisplayModel;
@@ -61,4 +62,7 @@
     RecentsAnimationDeviceState getRecentsAnimationDeviceState();
 
     SettingsChangeLogger getSettingsChangeLogger();
+
+    SimpleOrientationTouchTransformer getSimpleOrientationTouchTransformer();
+
 }
diff --git a/quickstep/src/com/android/quickstep/fallback/window/RecentsWindowSwipeHandler.java b/quickstep/src/com/android/quickstep/fallback/window/RecentsWindowSwipeHandler.java
index 973fb2f..5adc960 100644
--- a/quickstep/src/com/android/quickstep/fallback/window/RecentsWindowSwipeHandler.java
+++ b/quickstep/src/com/android/quickstep/fallback/window/RecentsWindowSwipeHandler.java
@@ -125,7 +125,7 @@
 
     @Override
     public void onRecentsAnimationStart(RecentsAnimationController controller,
-            RecentsAnimationTargets targets, TransitionInfo transitionInfo) {
+            RecentsAnimationTargets targets, @Nullable TransitionInfo transitionInfo) {
         super.onRecentsAnimationStart(controller, targets, transitionInfo);
         initTransformParams();
     }
diff --git a/quickstep/src/com/android/quickstep/recents/di/RecentsDependencies.kt b/quickstep/src/com/android/quickstep/recents/di/RecentsDependencies.kt
index b4648c4..02f48e6 100644
--- a/quickstep/src/com/android/quickstep/recents/di/RecentsDependencies.kt
+++ b/quickstep/src/com/android/quickstep/recents/di/RecentsDependencies.kt
@@ -26,10 +26,10 @@
 import com.android.quickstep.recents.data.TaskVisualsChangedDelegate
 import com.android.quickstep.recents.data.TaskVisualsChangedDelegateImpl
 import com.android.quickstep.recents.data.TasksRepository
+import com.android.quickstep.recents.domain.usecase.GetSysUiStatusNavFlagsUseCase
 import com.android.quickstep.recents.domain.usecase.GetTaskUseCase
 import com.android.quickstep.recents.usecase.GetThumbnailPositionUseCase
 import com.android.quickstep.recents.usecase.GetThumbnailUseCase
-import com.android.quickstep.recents.usecase.SysUiStatusNavFlagsUseCase
 import com.android.quickstep.recents.viewmodel.RecentsViewData
 import com.android.quickstep.task.thumbnail.SplashAlphaUseCase
 import com.android.quickstep.task.thumbnail.TaskThumbnailViewData
@@ -194,8 +194,7 @@
                 }
                 GetTaskUseCase::class.java -> GetTaskUseCase(repository = inject())
                 GetThumbnailUseCase::class.java -> GetThumbnailUseCase(taskRepository = inject())
-                SysUiStatusNavFlagsUseCase::class.java ->
-                    SysUiStatusNavFlagsUseCase(taskRepository = inject())
+                GetSysUiStatusNavFlagsUseCase::class.java -> GetSysUiStatusNavFlagsUseCase()
                 GetThumbnailPositionUseCase::class.java ->
                     GetThumbnailPositionUseCase(
                         deviceProfileRepository = inject(),
diff --git a/quickstep/src/com/android/quickstep/recents/domain/usecase/GetSysUiStatusNavFlagsUseCase.kt b/quickstep/src/com/android/quickstep/recents/domain/usecase/GetSysUiStatusNavFlagsUseCase.kt
new file mode 100644
index 0000000..da0e2d1
--- /dev/null
+++ b/quickstep/src/com/android/quickstep/recents/domain/usecase/GetSysUiStatusNavFlagsUseCase.kt
@@ -0,0 +1,43 @@
+/*
+ * Copyright (C) 2025 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.quickstep.recents.domain.usecase
+
+import android.view.WindowInsetsController.APPEARANCE_LIGHT_NAVIGATION_BARS
+import android.view.WindowInsetsController.APPEARANCE_LIGHT_STATUS_BARS
+import com.android.launcher3.util.SystemUiController.FLAG_DARK_NAV
+import com.android.launcher3.util.SystemUiController.FLAG_DARK_STATUS
+import com.android.launcher3.util.SystemUiController.FLAG_LIGHT_NAV
+import com.android.launcher3.util.SystemUiController.FLAG_LIGHT_STATUS
+import com.android.systemui.shared.recents.model.ThumbnailData
+
+/** UseCase to calculate flags for status bar and navigation bar */
+class GetSysUiStatusNavFlagsUseCase {
+    operator fun invoke(thumbnailData: ThumbnailData?): Int {
+        if (thumbnailData == null) return 0
+        val thumbnailAppearance = thumbnailData.appearance
+        var flags = 0
+        flags =
+            flags or
+                if (thumbnailAppearance and APPEARANCE_LIGHT_STATUS_BARS != 0) FLAG_LIGHT_STATUS
+                else FLAG_DARK_STATUS
+        flags =
+            flags or
+                if (thumbnailAppearance and APPEARANCE_LIGHT_NAVIGATION_BARS != 0) FLAG_LIGHT_NAV
+                else FLAG_DARK_NAV
+        return flags
+    }
+}
diff --git a/quickstep/src/com/android/quickstep/recents/ui/viewmodel/TaskTileUiState.kt b/quickstep/src/com/android/quickstep/recents/ui/viewmodel/TaskTileUiState.kt
index 54b2389..118a931 100644
--- a/quickstep/src/com/android/quickstep/recents/ui/viewmodel/TaskTileUiState.kt
+++ b/quickstep/src/com/android/quickstep/recents/ui/viewmodel/TaskTileUiState.kt
@@ -29,11 +29,13 @@
  * @property tasks The list of [TaskTileUiState] objects representing the individual tasks.
  * @property isLiveTile Indicates whether this data is intended for a live tile. If `true`, the
  *   running app will be displayed instead of the thumbnail.
+ * @property sysUiStatusNavFlags Flags for status bar and navigation bar
  */
 data class TaskTileUiState(
     val tasks: List<TaskData>,
     val isLiveTile: Boolean,
     val hasHeader: Boolean,
+    val sysUiStatusNavFlags: Int,
 )
 
 sealed class TaskData {
diff --git a/quickstep/src/com/android/quickstep/recents/ui/viewmodel/TaskViewModel.kt b/quickstep/src/com/android/quickstep/recents/ui/viewmodel/TaskViewModel.kt
index ed8197f..4936e30 100644
--- a/quickstep/src/com/android/quickstep/recents/ui/viewmodel/TaskViewModel.kt
+++ b/quickstep/src/com/android/quickstep/recents/ui/viewmodel/TaskViewModel.kt
@@ -22,6 +22,7 @@
 import com.android.launcher3.util.coroutines.DispatcherProvider
 import com.android.quickstep.recents.domain.model.TaskId
 import com.android.quickstep.recents.domain.model.TaskModel
+import com.android.quickstep.recents.domain.usecase.GetSysUiStatusNavFlagsUseCase
 import com.android.quickstep.recents.domain.usecase.GetTaskUseCase
 import com.android.quickstep.recents.viewmodel.RecentsViewData
 import com.android.quickstep.views.TaskViewType
@@ -43,6 +44,7 @@
     private val taskViewType: TaskViewType,
     recentsViewData: RecentsViewData,
     private val getTaskUseCase: GetTaskUseCase,
+    private val getSysUiStatusNavFlagsUseCase: GetSysUiStatusNavFlagsUseCase,
     dispatcherProvider: DispatcherProvider,
 ) {
     private var taskIds = MutableStateFlow(emptySet<Int>())
@@ -57,24 +59,19 @@
             }
             .distinctUntilChanged()
 
+    private val taskData =
+        taskIds.flatMapLatest { ids ->
+            // Combine Tasks requests
+            combine(
+                ids.map { id -> getTaskUseCase(id).map { taskModel -> id to taskModel } },
+                ::mapToTaskData,
+            )
+        }
+
     val tintAmount: Flow<Float> = recentsViewData.tintAmount
 
     val state: Flow<TaskTileUiState> =
-        taskIds
-            .flatMapLatest { ids ->
-                // Combine Tasks requests
-                combine(
-                    ids.map { id -> getTaskUseCase(id).map { taskModel -> id to taskModel } },
-                    ::mapToUiState,
-                )
-            }
-            .combine(isLiveTile) { tasks, isLiveTile ->
-                TaskTileUiState(
-                    tasks = tasks,
-                    isLiveTile = isLiveTile,
-                    hasHeader = taskViewType == TaskViewType.DESKTOP,
-                )
-            }
+        combine(taskData, isLiveTile) { tasks, isLiveTile -> mapToTaskTile(tasks, isLiveTile) }
             .distinctUntilChanged()
             .flowOn(dispatcherProvider.background)
 
@@ -83,10 +80,20 @@
         taskIds.value = taskId.toSet()
     }
 
-    private fun mapToUiState(result: Array<Pair<TaskId, TaskModel?>>): List<TaskData> =
-        result.map { mapToUiState(it.first, it.second) }
+    private fun mapToTaskTile(tasks: List<TaskData>, isLiveTile: Boolean): TaskTileUiState {
+        val firstThumbnailData = (tasks.firstOrNull() as? TaskData.Data)?.thumbnailData
+        return TaskTileUiState(
+            tasks = tasks,
+            isLiveTile = isLiveTile,
+            hasHeader = taskViewType == TaskViewType.DESKTOP,
+            sysUiStatusNavFlags = getSysUiStatusNavFlagsUseCase(firstThumbnailData),
+        )
+    }
 
-    private fun mapToUiState(taskId: TaskId, result: TaskModel?): TaskData =
+    private fun mapToTaskData(result: Array<Pair<TaskId, TaskModel?>>): List<TaskData> =
+        result.map { mapToTaskData(it.first, it.second) }
+
+    private fun mapToTaskData(taskId: TaskId, result: TaskModel?): TaskData =
         result?.let {
             TaskData.Data(
                 taskId = taskId,
diff --git a/quickstep/src/com/android/quickstep/recents/usecase/SysUiStatusNavFlagsUseCase.kt b/quickstep/src/com/android/quickstep/recents/usecase/SysUiStatusNavFlagsUseCase.kt
deleted file mode 100644
index 5be5f4a..0000000
--- a/quickstep/src/com/android/quickstep/recents/usecase/SysUiStatusNavFlagsUseCase.kt
+++ /dev/null
@@ -1,50 +0,0 @@
-/*
- * Copyright (C) 2024 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package com.android.quickstep.recents.usecase
-
-import android.view.WindowInsetsController
-import com.android.launcher3.util.SystemUiController.FLAG_DARK_NAV
-import com.android.launcher3.util.SystemUiController.FLAG_DARK_STATUS
-import com.android.launcher3.util.SystemUiController.FLAG_LIGHT_NAV
-import com.android.launcher3.util.SystemUiController.FLAG_LIGHT_STATUS
-import com.android.quickstep.recents.data.RecentTasksRepository
-
-/** UseCase to calculate flags for status bar and navigation bar */
-class SysUiStatusNavFlagsUseCase(private val taskRepository: RecentTasksRepository) {
-    fun getSysUiStatusNavFlags(taskId: Int): Int {
-        val thumbnailData = taskRepository.getCurrentThumbnailById(taskId) ?: return 0
-
-        val thumbnailAppearance = thumbnailData.appearance
-        var flags = 0
-        flags =
-            flags or
-                if (
-                    thumbnailAppearance and WindowInsetsController.APPEARANCE_LIGHT_STATUS_BARS != 0
-                )
-                    FLAG_LIGHT_STATUS
-                else FLAG_DARK_STATUS
-        flags =
-            flags or
-                if (
-                    thumbnailAppearance and
-                        WindowInsetsController.APPEARANCE_LIGHT_NAVIGATION_BARS != 0
-                )
-                    FLAG_LIGHT_NAV
-                else FLAG_DARK_NAV
-        return flags
-    }
-}
diff --git a/quickstep/src/com/android/quickstep/recents/viewmodel/TaskContainerViewModel.kt b/quickstep/src/com/android/quickstep/recents/viewmodel/TaskContainerViewModel.kt
index 168c1e0..723df55 100644
--- a/quickstep/src/com/android/quickstep/recents/viewmodel/TaskContainerViewModel.kt
+++ b/quickstep/src/com/android/quickstep/recents/viewmodel/TaskContainerViewModel.kt
@@ -16,23 +16,11 @@
 
 package com.android.quickstep.recents.viewmodel
 
-import android.graphics.Bitmap
-import com.android.quickstep.recents.usecase.GetThumbnailUseCase
-import com.android.quickstep.recents.usecase.SysUiStatusNavFlagsUseCase
 import com.android.quickstep.task.thumbnail.SplashAlphaUseCase
 import kotlinx.coroutines.flow.firstOrNull
 import kotlinx.coroutines.runBlocking
 
-class TaskContainerViewModel(
-    private val sysUiStatusNavFlagsUseCase: SysUiStatusNavFlagsUseCase,
-    private val getThumbnailUseCase: GetThumbnailUseCase,
-    private val splashAlphaUseCase: SplashAlphaUseCase,
-) {
-    fun getThumbnail(taskId: Int): Bitmap? = getThumbnailUseCase.run(taskId)
-
-    fun getSysUiStatusNavFlags(taskId: Int) =
-        sysUiStatusNavFlagsUseCase.getSysUiStatusNavFlags(taskId)
-
+class TaskContainerViewModel(private val splashAlphaUseCase: SplashAlphaUseCase) {
     fun shouldShowThumbnailSplash(taskId: Int): Boolean =
         (runBlocking { splashAlphaUseCase.execute(taskId).firstOrNull() } ?: 0f) > 0f
 }
diff --git a/quickstep/src/com/android/quickstep/views/RecentsView.java b/quickstep/src/com/android/quickstep/views/RecentsView.java
index 0eb9dbc..1c8299d 100644
--- a/quickstep/src/com/android/quickstep/views/RecentsView.java
+++ b/quickstep/src/com/android/quickstep/views/RecentsView.java
@@ -5709,7 +5709,7 @@
         updateGridProperties();
         updateScrollSynchronously();
 
-        int targetSysUiFlags = taskView.getTaskContainers().getFirst().getSysUiStatusNavFlags();
+        int targetSysUiFlags = taskView.getSysUiStatusNavFlags();
         final boolean[] passedOverviewThreshold = new boolean[]{false};
         AnimatorSet anim = createAdjacentPageAnimForTaskLaunch(taskView);
         anim.play(new AnimatedFloat(v -> {
diff --git a/quickstep/src/com/android/quickstep/views/TaskContainer.kt b/quickstep/src/com/android/quickstep/views/TaskContainer.kt
index 8cd5601..a9e84ef 100644
--- a/quickstep/src/com/android/quickstep/views/TaskContainer.kt
+++ b/quickstep/src/com/android/quickstep/views/TaskContainer.kt
@@ -63,11 +63,7 @@
 
     // TODO(b/335649589): Ideally create and obtain this from DI.
     private val taskContainerViewModel: TaskContainerViewModel by lazy {
-        TaskContainerViewModel(
-            sysUiStatusNavFlagsUseCase = RecentsDependencies.get(),
-            getThumbnailUseCase = RecentsDependencies.get(),
-            splashAlphaUseCase = RecentsDependencies.get(),
-        )
+        TaskContainerViewModel(splashAlphaUseCase = RecentsDependencies.get())
     }
 
     init {
@@ -86,13 +82,9 @@
         }
     }
 
-    val splitAnimationThumbnail: Bitmap?
-        get() =
-            if (enableRefactorTaskThumbnail()) {
-                taskContainerViewModel.getThumbnail(task.key.id)
-            } else {
-                thumbnailViewDeprecated.thumbnail
-            }
+    var splitAnimationThumbnail: Bitmap? = null
+        get() = if (enableRefactorTaskThumbnail()) field else thumbnailViewDeprecated.thumbnail
+        private set
 
     val thumbnailView: TaskThumbnailView
         get() {
@@ -112,12 +104,6 @@
                 taskContainerViewModel.shouldShowThumbnailSplash(task.key.id)
             else thumbnailViewDeprecated.shouldShowSplashView()
 
-    val sysUiStatusNavFlags: Int
-        get() =
-            if (enableRefactorTaskThumbnail())
-                taskContainerViewModel.getSysUiStatusNavFlags(task.key.id)
-            else thumbnailViewDeprecated.sysUiStatusNavFlags
-
     /** Builds proto for logging */
     val itemInfo: TaskViewItemInfo
         get() = TaskViewItemInfo(this)
@@ -168,6 +154,8 @@
 
     fun setState(state: TaskData?, liveTile: Boolean, hasHeader: Boolean) {
         thumbnailView.setState(TaskUiStateMapper.toTaskThumbnailUiState(state, liveTile, hasHeader))
+        splitAnimationThumbnail =
+            if (state is TaskData.Data) state.thumbnailData?.thumbnail else null
     }
 
     fun updateTintAmount(tintAmount: Float) {
diff --git a/quickstep/src/com/android/quickstep/views/TaskView.kt b/quickstep/src/com/android/quickstep/views/TaskView.kt
index 0eb8cb0..20a385f 100644
--- a/quickstep/src/com/android/quickstep/views/TaskView.kt
+++ b/quickstep/src/com/android/quickstep/views/TaskView.kt
@@ -255,6 +255,11 @@
 
     var taskViewId = UNBOUND_TASK_VIEW_ID
     var isEndQuickSwitchCuj = false
+    var sysUiStatusNavFlags: Int = 0
+        get() =
+            if (enableRefactorTaskThumbnail()) field
+            else taskContainers.first().thumbnailViewDeprecated.sysUiStatusNavFlags
+        private set
 
     // Various animation progress variables.
     // progress: 0 = show icon and no insets; 1 = don't show icon and show full insets.
@@ -756,9 +761,7 @@
                 viewModel!!.tintAmount.onEach(::updateTintAmount).launchIn(coroutineScope)
 
             coroutineJobs +=
-                coroutineScope.launch {
-                    viewModel!!.state.collectLatest(::updateTaskContainerState)
-                }
+                coroutineScope.launch { viewModel!!.state.collectLatest(::updateTaskViewState) }
         }
     }
 
@@ -766,7 +769,10 @@
         taskContainers.forEach { it.updateTintAmount(amount) }
     }
 
-    private fun updateTaskContainerState(state: TaskTileUiState) {
+    private fun updateTaskViewState(state: TaskTileUiState) {
+        sysUiStatusNavFlags = state.sysUiStatusNavFlags
+
+        // Updating containers
         val mapOfTasks = state.tasks.associateBy { it.taskId }
         taskContainers.forEach { container ->
             container.setState(
@@ -820,6 +826,7 @@
                         taskViewType = type,
                         recentsViewData = RecentsDependencies.get(),
                         getTaskUseCase = RecentsDependencies.get(),
+                        getSysUiStatusNavFlagsUseCase = RecentsDependencies.get(),
                         dispatcherProvider = RecentsDependencies.get(),
                     )
                     .apply { bind(*taskIds) }
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/domain/usecase/GetSysUiStatusNavFlagsUseCaseTest.kt b/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/domain/usecase/GetSysUiStatusNavFlagsUseCaseTest.kt
new file mode 100644
index 0000000..d384256
--- /dev/null
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/domain/usecase/GetSysUiStatusNavFlagsUseCaseTest.kt
@@ -0,0 +1,118 @@
+/*
+ * Copyright (C) 2025 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.quickstep.recents.domain.usecase
+
+import android.view.WindowInsetsController.APPEARANCE_LIGHT_CAPTION_BARS
+import android.view.WindowInsetsController.APPEARANCE_LIGHT_NAVIGATION_BARS
+import android.view.WindowInsetsController.APPEARANCE_LIGHT_STATUS_BARS
+import com.android.launcher3.util.SystemUiController.FLAG_DARK_NAV
+import com.android.launcher3.util.SystemUiController.FLAG_DARK_STATUS
+import com.android.launcher3.util.SystemUiController.FLAG_LIGHT_NAV
+import com.android.launcher3.util.SystemUiController.FLAG_LIGHT_STATUS
+import com.android.systemui.shared.recents.model.ThumbnailData
+import com.google.common.truth.Truth.assertThat
+import org.junit.Test
+
+class GetSysUiStatusNavFlagsUseCaseTest {
+    private val sut: GetSysUiStatusNavFlagsUseCase = GetSysUiStatusNavFlagsUseCase()
+
+    @Test
+    fun onLightStatusBarAppearance_returns_LightTheme() {
+        val thumbnailData = ThumbnailData(appearance = APPEARANCE_LIGHT_STATUS_BARS)
+        val flag = sut.invoke(thumbnailData) // 6
+        flag.assertContainsFlag(FLAG_LIGHT_STATUS)
+        flag.assertContainsFlag(FLAG_DARK_NAV)
+        flag.assertDoesNotContainsFlag(FLAG_DARK_STATUS)
+        flag.assertDoesNotContainsFlag(FLAG_LIGHT_NAV)
+    }
+
+    @Test
+    fun onLightNavBarsAppearance_returns_LightTheme() {
+        val thumbnailData = ThumbnailData(appearance = APPEARANCE_LIGHT_NAVIGATION_BARS)
+        val flag = sut.invoke(thumbnailData)
+        flag.assertContainsFlag(FLAG_DARK_STATUS)
+        flag.assertContainsFlag(FLAG_LIGHT_NAV)
+        flag.assertDoesNotContainsFlag(FLAG_LIGHT_STATUS)
+        flag.assertDoesNotContainsFlag(FLAG_DARK_NAV)
+    }
+
+    @Test
+    fun onLightStatusBarAndNavBarAppearance_returns_LightTheme() {
+        val thumbnailData =
+            ThumbnailData(
+                appearance = APPEARANCE_LIGHT_STATUS_BARS or APPEARANCE_LIGHT_NAVIGATION_BARS
+            )
+        val flag = sut.invoke(thumbnailData)
+        flag.assertContainsFlag(FLAG_LIGHT_NAV)
+        flag.assertContainsFlag(FLAG_LIGHT_STATUS)
+        flag.assertDoesNotContainsFlag(FLAG_DARK_STATUS)
+        flag.assertDoesNotContainsFlag(FLAG_DARK_NAV)
+    }
+
+    @Test
+    fun onLightAppearance_returns_LightTheme() {
+        val thumbnailData =
+            ThumbnailData(
+                appearance =
+                    APPEARANCE_LIGHT_CAPTION_BARS or
+                        APPEARANCE_LIGHT_STATUS_BARS or
+                        APPEARANCE_LIGHT_NAVIGATION_BARS
+            )
+        val flag = sut.invoke(thumbnailData)
+        flag.assertContainsFlag(FLAG_LIGHT_NAV)
+        flag.assertContainsFlag(FLAG_LIGHT_STATUS)
+        flag.assertDoesNotContainsFlag(FLAG_DARK_STATUS)
+        flag.assertDoesNotContainsFlag(FLAG_DARK_NAV)
+    }
+
+    @Test
+    fun onDarkAppearance_returns_DarkTheme() {
+        val thumbnailData = ThumbnailData(appearance = 0)
+        val flag = sut.invoke(thumbnailData)
+        flag.assertContainsFlag(FLAG_DARK_STATUS)
+        flag.assertContainsFlag(FLAG_DARK_NAV)
+        flag.assertDoesNotContainsFlag(FLAG_LIGHT_NAV)
+        flag.assertDoesNotContainsFlag(FLAG_LIGHT_STATUS)
+    }
+
+    @Test
+    fun onUnrelatedDarkAppearance_returns_DarkTheme() {
+        val thumbnailData = ThumbnailData(appearance = 1)
+        val flag = sut.invoke(thumbnailData)
+        flag.assertContainsFlag(FLAG_DARK_STATUS)
+        flag.assertContainsFlag(FLAG_DARK_NAV)
+        flag.assertDoesNotContainsFlag(FLAG_LIGHT_NAV)
+        flag.assertDoesNotContainsFlag(FLAG_LIGHT_STATUS)
+    }
+
+    @Test
+    fun whenThumbnailIsNull_returns_default() {
+        val flag = sut.invoke(null)
+        flag.assertDoesNotContainsFlag(FLAG_DARK_STATUS)
+        flag.assertDoesNotContainsFlag(FLAG_LIGHT_NAV)
+        flag.assertDoesNotContainsFlag(FLAG_LIGHT_STATUS)
+        flag.assertDoesNotContainsFlag(FLAG_DARK_NAV)
+    }
+
+    private fun Int.assertContainsFlag(flag: Int) {
+        assertThat(this and flag).isNotEqualTo(0)
+    }
+
+    private fun Int.assertDoesNotContainsFlag(flag: Int) {
+        assertThat(this and flag).isEqualTo(0)
+    }
+}
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/ui/viewmodel/TaskViewModelTest.kt b/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/ui/viewmodel/TaskViewModelTest.kt
index 7a4b5f2..c031150 100644
--- a/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/ui/viewmodel/TaskViewModelTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/ui/viewmodel/TaskViewModelTest.kt
@@ -18,9 +18,15 @@
 
 import android.graphics.Color
 import android.graphics.drawable.ShapeDrawable
+import android.view.WindowInsetsController.APPEARANCE_LIGHT_CAPTION_BARS
+import android.view.WindowInsetsController.APPEARANCE_LIGHT_NAVIGATION_BARS
+import android.view.WindowInsetsController.APPEARANCE_LIGHT_STATUS_BARS
 import androidx.test.ext.junit.runners.AndroidJUnit4
+import com.android.launcher3.util.SystemUiController.FLAG_LIGHT_NAV
+import com.android.launcher3.util.SystemUiController.FLAG_LIGHT_STATUS
 import com.android.launcher3.util.TestDispatcherProvider
 import com.android.quickstep.recents.domain.model.TaskModel
+import com.android.quickstep.recents.domain.usecase.GetSysUiStatusNavFlagsUseCase
 import com.android.quickstep.recents.domain.usecase.GetTaskUseCase
 import com.android.quickstep.recents.viewmodel.RecentsViewData
 import com.android.quickstep.views.TaskViewType
@@ -55,6 +61,7 @@
                 taskViewType = TaskViewType.SINGLE,
                 recentsViewData = recentsViewData,
                 getTaskUseCase = getTaskUseCase,
+                getSysUiStatusNavFlagsUseCase = GetSysUiStatusNavFlagsUseCase(),
                 dispatcherProvider = TestDispatcherProvider(unconfinedTestDispatcher),
             )
         whenever(getTaskUseCase.invoke(TASK_MODEL_1.id)).thenReturn(flow { emit(TASK_MODEL_1) })
@@ -73,6 +80,7 @@
                     tasks = listOf(TASK_MODEL_1.toUiState()),
                     isLiveTile = false,
                     hasHeader = false,
+                    sysUiStatusNavFlags = FLAGS_APPEARANCE_LIGHT_THEME,
                 )
             assertThat(sut.state.first()).isEqualTo(expectedResult)
         }
@@ -93,6 +101,7 @@
                         taskViewType = type,
                         recentsViewData = recentsViewData,
                         getTaskUseCase = getTaskUseCase,
+                        getSysUiStatusNavFlagsUseCase = GetSysUiStatusNavFlagsUseCase(),
                         dispatcherProvider = TestDispatcherProvider(unconfinedTestDispatcher),
                     )
                 sut.bind(TASK_MODEL_1.id)
@@ -115,6 +124,7 @@
                         ),
                     isLiveTile = false,
                     hasHeader = false,
+                    sysUiStatusNavFlags = FLAGS_APPEARANCE_LIGHT_THEME,
                 )
             assertThat(sut.state.first()).isEqualTo(expectedResult)
         }
@@ -136,6 +146,7 @@
                         ),
                     isLiveTile = true,
                     hasHeader = false,
+                    sysUiStatusNavFlags = FLAGS_APPEARANCE_LIGHT_THEME,
                 )
             assertThat(sut.state.first()).isEqualTo(expectedResult)
         }
@@ -157,6 +168,7 @@
                         ),
                     isLiveTile = false,
                     hasHeader = false,
+                    sysUiStatusNavFlags = FLAGS_APPEARANCE_LIGHT_THEME,
                 )
             assertThat(sut.state.first()).isEqualTo(expectedResult)
         }
@@ -177,6 +189,7 @@
                         ),
                     isLiveTile = false,
                     hasHeader = false,
+                    sysUiStatusNavFlags = FLAGS_APPEARANCE_LIGHT_THEME,
                 )
             assertThat(sut.state.first()).isEqualTo(expectedResult)
         }
@@ -193,6 +206,7 @@
                     tasks = listOf(TASK_MODEL_1.toUiState(), TASK_MODEL_2.toUiState()),
                     isLiveTile = false,
                     hasHeader = false,
+                    sysUiStatusNavFlags = FLAGS_APPEARANCE_LIGHT_THEME,
                 )
             assertThat(sut.state.first()).isEqualTo(expectedResult)
         }
@@ -206,6 +220,7 @@
                     listOf(TaskData.NoData(INVALID_TASK_ID)),
                     isLiveTile = false,
                     hasHeader = false,
+                    sysUiStatusNavFlags = FLAGS_APPEARANCE_DEFAULT,
                 )
             assertThat(sut.state.first()).isEqualTo(expectedResult)
         }
@@ -221,15 +236,22 @@
             isLocked = isLocked,
         )
 
-    companion object {
+    private companion object {
         const val INVALID_TASK_ID = -1
+        const val FLAGS_APPEARANCE_LIGHT_THEME = FLAG_LIGHT_STATUS or FLAG_LIGHT_NAV
+        const val FLAGS_APPEARANCE_DEFAULT = 0
+        const val APPEARANCE_LIGHT_THEME =
+            APPEARANCE_LIGHT_CAPTION_BARS or
+                APPEARANCE_LIGHT_STATUS_BARS or
+                APPEARANCE_LIGHT_NAVIGATION_BARS
+
         val TASK_MODEL_1 =
             TaskModel(
                 1,
                 "Title 1",
                 "Content Description 1",
                 ShapeDrawable(),
-                ThumbnailData(),
+                ThumbnailData(appearance = APPEARANCE_LIGHT_THEME),
                 Color.BLACK,
                 false,
             )
@@ -239,7 +261,7 @@
                 "Title 2",
                 "Content Description 2",
                 ShapeDrawable(),
-                ThumbnailData(),
+                ThumbnailData(appearance = APPEARANCE_LIGHT_THEME),
                 Color.RED,
                 true,
             )
@@ -249,7 +271,7 @@
                 "Title 3",
                 "Content Description 3",
                 ShapeDrawable(),
-                ThumbnailData(),
+                ThumbnailData(appearance = APPEARANCE_LIGHT_THEME),
                 Color.BLUE,
                 false,
             )
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/usecase/SysUiStatusNavFlagsUseCaseTest.kt b/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/usecase/SysUiStatusNavFlagsUseCaseTest.kt
deleted file mode 100644
index 92f2efd..0000000
--- a/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/usecase/SysUiStatusNavFlagsUseCaseTest.kt
+++ /dev/null
@@ -1,110 +0,0 @@
-/*
- * Copyright (C) 2024 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package com.android.quickstep.recents.usecase
-
-import android.content.ComponentName
-import android.content.Intent
-import android.graphics.Bitmap
-import android.graphics.Color
-import com.android.quickstep.recents.data.FakeTasksRepository
-import com.android.systemui.shared.recents.model.Task
-import com.android.systemui.shared.recents.model.ThumbnailData
-import com.google.common.truth.Truth.assertThat
-import org.junit.Before
-import org.junit.Test
-import org.mockito.kotlin.mock
-import org.mockito.kotlin.whenever
-
-/** Test for [SysUiStatusNavFlagsUseCase] */
-class SysUiStatusNavFlagsUseCaseTest {
-    private lateinit var tasksRepository: FakeTasksRepository
-    private lateinit var sysUiStatusNavFlagsUseCase: SysUiStatusNavFlagsUseCase
-
-    @Before
-    fun setup() {
-        tasksRepository = FakeTasksRepository()
-        sysUiStatusNavFlagsUseCase = SysUiStatusNavFlagsUseCase(tasksRepository)
-        initTaskRepository()
-    }
-
-    @Test
-    fun onLightAppearanceReturnExpectedFlags() {
-        assertThat(sysUiStatusNavFlagsUseCase.getSysUiStatusNavFlags(FIRST_TASK_ID))
-            .isEqualTo(FLAGS_APPEARANCE_LIGHT_THEME)
-    }
-
-    @Test
-    fun onDarkAppearanceReturnExpectedFlags() {
-        assertThat(sysUiStatusNavFlagsUseCase.getSysUiStatusNavFlags(SECOND_TASK_ID))
-            .isEqualTo(FLAGS_APPEARANCE_DARK_THEME)
-    }
-
-    @Test
-    fun whenThumbnailIsNullReturnDefault() {
-        assertThat(sysUiStatusNavFlagsUseCase.getSysUiStatusNavFlags(UNKNOWN_TASK_ID))
-            .isEqualTo(FLAGS_DEFAULT)
-    }
-
-    private fun initTaskRepository() {
-        val firstTask =
-            Task(Task.TaskKey(FIRST_TASK_ID, 0, Intent(), ComponentName("", ""), 0, 2000)).apply {
-                colorBackground = Color.BLACK
-            }
-        val firstThumbnailData =
-            ThumbnailData(
-                thumbnail =
-                    mock<Bitmap>().apply {
-                        whenever(width).thenReturn(THUMBNAIL_WIDTH)
-                        whenever(height).thenReturn(THUMBNAIL_HEIGHT)
-                    },
-                appearance = APPEARANCE_LIGHT_THEME,
-            )
-
-        val secondTask =
-            Task(Task.TaskKey(SECOND_TASK_ID, 0, Intent(), ComponentName("", ""), 0, 2005)).apply {
-                colorBackground = Color.BLACK
-            }
-        val secondThumbnailData =
-            ThumbnailData(
-                thumbnail =
-                    mock<Bitmap>().apply {
-                        whenever(width).thenReturn(THUMBNAIL_WIDTH)
-                        whenever(height).thenReturn(THUMBNAIL_HEIGHT)
-                    },
-                appearance = APPEARANCE_DARK_THEME,
-            )
-
-        tasksRepository.seedTasks(listOf(firstTask, secondTask))
-        tasksRepository.seedThumbnailData(
-            mapOf(FIRST_TASK_ID to firstThumbnailData, SECOND_TASK_ID to secondThumbnailData)
-        )
-        tasksRepository.setVisibleTasks(setOf(FIRST_TASK_ID, SECOND_TASK_ID))
-    }
-
-    companion object {
-        const val FIRST_TASK_ID = 0
-        const val SECOND_TASK_ID = 100
-        const val UNKNOWN_TASK_ID = 404
-        const val THUMBNAIL_WIDTH = 100
-        const val THUMBNAIL_HEIGHT = 200
-        const val APPEARANCE_LIGHT_THEME = 24
-        const val FLAGS_APPEARANCE_LIGHT_THEME = 5
-        const val APPEARANCE_DARK_THEME = 0
-        const val FLAGS_APPEARANCE_DARK_THEME = 10
-        const val FLAGS_DEFAULT = 0
-    }
-}
diff --git a/quickstep/tests/src/com/android/quickstep/DesktopSystemShortcutTest.kt b/quickstep/tests/src/com/android/quickstep/DesktopSystemShortcutTest.kt
index 52bd2ea..76aab39 100644
--- a/quickstep/tests/src/com/android/quickstep/DesktopSystemShortcutTest.kt
+++ b/quickstep/tests/src/com/android/quickstep/DesktopSystemShortcutTest.kt
@@ -19,6 +19,7 @@
 import android.content.ComponentName
 import android.content.Context
 import android.content.Intent
+import android.content.pm.PackageManager
 import android.platform.test.annotations.EnableFlags
 import android.view.Display.DEFAULT_DISPLAY
 import androidx.test.platform.app.InstrumentationRegistry
@@ -72,7 +73,7 @@
     private val overlayFactory: TaskOverlayFactory = mock()
     private val factory: TaskShortcutFactory =
         DesktopSystemShortcut.createFactory(abstractFloatingViewHelper)
-    private val context: Context = InstrumentationRegistry.getInstrumentation().targetContext
+    private val context: Context = spy(InstrumentationRegistry.getInstrumentation().targetContext)
 
     private lateinit var mockitoSession: StaticMockitoSession
 
@@ -151,6 +152,32 @@
     }
 
     @Test
+    @EnableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_MODALS_POLICY)
+    fun createDesktopTaskShortcutFactory_defaultHomeTask() {
+        val packageManager: PackageManager = mock()
+        val homeActivities = ComponentName("defaultHomePackage", /* class */ "")
+        whenever(context.packageManager).thenReturn(packageManager)
+        whenever(packageManager.getHomeActivities(any())).thenReturn(homeActivities)
+        val taskKey =
+            TaskKey(
+                /* id */ 1,
+                /* windowingMode */ 0,
+                Intent(),
+                homeActivities,
+                /* userId */ 0,
+                /* lastActiveTime */ 2000,
+                DEFAULT_DISPLAY,
+                homeActivities,
+                /* numActivities */ 1,
+                /* isTopActivityNoDisplay */ false,
+                /* isActivityStackTransparent */ false,
+            )
+        val taskContainer = createTaskContainer(Task(taskKey).apply { isDockable = true })
+        val shortcuts = factory.getShortcuts(launcher, taskContainer)
+        assertThat(shortcuts).isNull()
+    }
+
+    @Test
     fun createDesktopTaskShortcutFactory_undockable() {
         val unDockableTask = createTask().apply { isDockable = false }
         val taskContainer = createTaskContainer(unDockableTask)
diff --git a/quickstep/tests/src/com/android/quickstep/ExternalDisplaySystemShortcutTest.kt b/quickstep/tests/src/com/android/quickstep/ExternalDisplaySystemShortcutTest.kt
index 4111dec..818841a 100644
--- a/quickstep/tests/src/com/android/quickstep/ExternalDisplaySystemShortcutTest.kt
+++ b/quickstep/tests/src/com/android/quickstep/ExternalDisplaySystemShortcutTest.kt
@@ -19,6 +19,7 @@
 import android.content.ComponentName
 import android.content.Context
 import android.content.Intent
+import android.content.pm.PackageManager
 import android.platform.test.annotations.EnableFlags
 import android.platform.test.flag.junit.SetFlagsRule
 import android.view.Display.DEFAULT_DISPLAY
@@ -75,7 +76,7 @@
     private val overlayFactory: TaskOverlayFactory = mock()
     private val factory: TaskShortcutFactory =
         ExternalDisplaySystemShortcut.createFactory(abstractFloatingViewHelper)
-    private val context: Context = InstrumentationRegistry.getInstrumentation().targetContext
+    private val context: Context = spy(InstrumentationRegistry.getInstrumentation().targetContext)
 
     private lateinit var mockitoSession: StaticMockitoSession
 
@@ -161,6 +162,35 @@
     }
 
     @Test
+    @EnableFlags(
+        Flags.FLAG_MOVE_TO_EXTERNAL_DISPLAY_SHORTCUT,
+        Flags.FLAG_ENABLE_DESKTOP_WINDOWING_MODALS_POLICY,
+    )
+    fun createExternalDisplayTaskShortcut_defaultHomeTask() {
+        val packageManager: PackageManager = mock()
+        val homeActivities = ComponentName("defaultHomePackage", /* class */ "")
+        whenever(context.packageManager).thenReturn(packageManager)
+        whenever(packageManager.getHomeActivities(any())).thenReturn(homeActivities)
+        val taskKey =
+            TaskKey(
+                /* id */ 1,
+                /* windowingMode */ 0,
+                Intent(),
+                homeActivities,
+                /* userId */ 0,
+                /* lastActiveTime */ 2000,
+                DEFAULT_DISPLAY,
+                homeActivities,
+                /* numActivities */ 1,
+                /* isTopActivityNoDisplay */ false,
+                /* isActivityStackTransparent */ false,
+            )
+        val taskContainer = createTaskContainer(Task(taskKey).apply { isDockable = true })
+        val shortcuts = factory.getShortcuts(launcher, taskContainer)
+        assertThat(shortcuts).isNull()
+    }
+
+    @Test
     @EnableFlags(Flags.FLAG_MOVE_TO_EXTERNAL_DISPLAY_SHORTCUT)
     fun externalDisplaySystemShortcutClicked() {
         val task = createTask()
diff --git a/quickstep/tests/src/com/android/quickstep/OrientationTouchTransformerTest.java b/quickstep/tests/src/com/android/quickstep/OrientationTouchTransformerTest.java
index ff0ad53..a0ec635 100644
--- a/quickstep/tests/src/com/android/quickstep/OrientationTouchTransformerTest.java
+++ b/quickstep/tests/src/com/android/quickstep/OrientationTouchTransformerTest.java
@@ -43,6 +43,7 @@
 import androidx.test.filters.SmallTest;
 
 import com.android.launcher3.testing.shared.ResourceUtils;
+import com.android.launcher3.util.DaggerSingletonTracker;
 import com.android.launcher3.util.DisplayController;
 import com.android.launcher3.util.RotationUtils;
 import com.android.launcher3.util.WindowBounds;
@@ -301,7 +302,8 @@
         final DisplayController displayController = mock(DisplayController.class);
         doReturn(mInfo).when(displayController).getInfo();
         final SimpleOrientationTouchTransformer transformer =
-                new SimpleOrientationTouchTransformer(getApplicationContext(), displayController);
+                new SimpleOrientationTouchTransformer(getApplicationContext(), displayController,
+                        mock(DaggerSingletonTracker.class));
         final MotionEvent move1 = generateMotionEvent(MotionEvent.ACTION_MOVE, 100, 10);
         transformer.transform(move1, Surface.ROTATION_90);
         // The position is transformed to 90 degree.
diff --git a/src/com/android/launcher3/BubbleTextView.java b/src/com/android/launcher3/BubbleTextView.java
index afa18a5..d3684b2 100644
--- a/src/com/android/launcher3/BubbleTextView.java
+++ b/src/com/android/launcher3/BubbleTextView.java
@@ -29,6 +29,7 @@
 import static com.android.launcher3.icons.GraphicsUtils.setColorAlphaBound;
 import static com.android.launcher3.model.data.ItemInfoWithIcon.FLAG_INCREMENTAL_DOWNLOAD_ACTIVE;
 import static com.android.launcher3.model.data.ItemInfoWithIcon.FLAG_INSTALL_SESSION_ACTIVE;
+import static com.android.launcher3.model.data.ItemInfoWithIcon.FLAG_SHOW_DOWNLOAD_PROGRESS_MASK;
 
 import android.animation.Animator;
 import android.animation.AnimatorListenerAdapter;
@@ -63,6 +64,7 @@
 import android.widget.TextView;
 
 import androidx.annotation.DrawableRes;
+import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
 import androidx.annotation.UiThread;
 import androidx.annotation.VisibleForTesting;
@@ -366,11 +368,6 @@
         mDotScaleAnim.start();
     }
 
-    @UiThread
-    public void applyFromWorkspaceItem(WorkspaceItemInfo info) {
-        applyFromWorkspaceItem(info, null);
-    }
-
     @Override
     public void setAccessibilityDelegate(AccessibilityDelegate delegate) {
         if (delegate instanceof BaseAccessibilityDelegate) {
@@ -384,10 +381,10 @@
     }
 
     @UiThread
-    public void applyFromWorkspaceItem(WorkspaceItemInfo info, PreloadIconDrawable icon) {
+    public void applyFromWorkspaceItem(WorkspaceItemInfo info) {
         applyIconAndLabel(info);
         setItemInfo(info);
-        applyLoadingState(icon);
+
         applyDotState(info, false /* animate */);
         setDownloadStateContentDescription(info, info.getProgressLevel());
     }
@@ -395,17 +392,11 @@
     @UiThread
     public void applyFromApplicationInfo(AppInfo info) {
         applyIconAndLabel(info);
-
-        // We don't need to check the info since it's not a WorkspaceItemInfo
         setItemInfo(info);
 
-
         // Verify high res immediately
         verifyHighRes();
 
-        if ((info.runtimeStatusFlags & ItemInfoWithIcon.FLAG_SHOW_DOWNLOAD_PROGRESS_MASK) != 0) {
-            applyProgressLevel();
-        }
         applyDotState(info, false /* animate */);
         setDownloadStateContentDescription(info, info.getProgressLevel());
     }
@@ -449,6 +440,50 @@
     @VisibleForTesting
     @UiThread
     public void applyIconAndLabel(ItemInfoWithIcon info) {
+        FastBitmapDrawable oldIcon = mIcon;
+        if (!canReuseIcon(info)) {
+            setNonPendingIcon(info);
+        }
+        applyLabel(info);
+        maybeApplyProgressLevel(info, oldIcon);
+    }
+
+    /**
+     * Check if we can reuse icon so that any animation is preserved
+     */
+    private boolean canReuseIcon(ItemInfoWithIcon info) {
+        return mIcon instanceof PreloadIconDrawable p
+                && p.hasNotCompleted() && p.isSameInfo(info.bitmap);
+    }
+
+    /**
+     * Apply progress level to the icon if necessary
+     */
+    private void maybeApplyProgressLevel(ItemInfoWithIcon info, FastBitmapDrawable oldIcon) {
+        if (!shouldApplyProgressLevel(info, oldIcon)) {
+            return;
+        }
+        PreloadIconDrawable pendingIcon = applyProgressLevel(info);
+        boolean isNoLongerPending = info instanceof WorkspaceItemInfo wii
+                ? !wii.hasPromiseIconUi() : !info.isArchived();
+        if (isNoLongerPending && info.getProgressLevel() == 100 && pendingIcon != null) {
+            pendingIcon.maybePerformFinishedAnimation(
+                    (oldIcon instanceof PreloadIconDrawable p) ? p : pendingIcon,
+                    () -> setNonPendingIcon(
+                            (getTag() instanceof ItemInfoWithIcon iiwi) ? iiwi : info));
+        }
+    }
+
+    /**
+     * Check if progress level should be applied to the icon
+     */
+    private boolean shouldApplyProgressLevel(ItemInfoWithIcon info, FastBitmapDrawable oldIcon) {
+        return (info.runtimeStatusFlags & FLAG_SHOW_DOWNLOAD_PROGRESS_MASK) != 0
+                || (info instanceof WorkspaceItemInfo wii && wii.hasPromiseIconUi())
+                || (oldIcon instanceof PreloadIconDrawable p && p.hasNotCompleted());
+    }
+
+    private void setNonPendingIcon(ItemInfoWithIcon info) {
         ThemeManager themeManager = ThemeManager.INSTANCE.get(getContext());
         int flags = (shouldUseTheme()
                 && themeManager.isMonoThemeEnabled()) ? FLAG_THEMED : 0;
@@ -463,7 +498,6 @@
         mDotParams.appColor = iconDrawable.getIconColor();
         mDotParams.dotColor = Themes.getAttrColor(getContext(), R.attr.notificationDotColor);
         setIcon(iconDrawable);
-        applyLabel(info);
     }
 
     protected boolean shouldUseTheme() {
@@ -1070,38 +1104,10 @@
         mLongPressHelper.cancelLongPress();
     }
 
-    /**
-     * Applies the loading progress value to the progress bar.
-     *
-     * If this app is installing, the progress bar will be updated with the installation progress.
-     * If this app is installed and downloading incrementally, the progress bar will be updated
-     * with the total download progress.
-     */
-    public void applyLoadingState(PreloadIconDrawable icon) {
-        if (getTag() instanceof ItemInfoWithIcon) {
-            WorkspaceItemInfo info = (WorkspaceItemInfo) getTag();
-            if ((info.runtimeStatusFlags & FLAG_INCREMENTAL_DOWNLOAD_ACTIVE) != 0
-                    || info.hasPromiseIconUi()
-                    || (info.runtimeStatusFlags & FLAG_INSTALL_SESSION_ACTIVE) != 0
-                    || (icon != null)) {
-                updateProgressBarUi(info.getProgressLevel() == 100 ? icon : null);
-            }
-        }
-    }
-
-    private void updateProgressBarUi(PreloadIconDrawable oldIcon) {
-        FastBitmapDrawable originalIcon = mIcon;
-        PreloadIconDrawable preloadDrawable = applyProgressLevel();
-        if (preloadDrawable != null && oldIcon != null) {
-            preloadDrawable.maybePerformFinishedAnimation(oldIcon, () -> setIcon(originalIcon));
-        }
-    }
-
     /** Applies the given progress level to the this icon's progress bar. */
     @Nullable
-    public PreloadIconDrawable applyProgressLevel() {
-        if (!(getTag() instanceof ItemInfoWithIcon info)
-                || ((ItemInfoWithIcon) getTag()).isInactiveArchive()) {
+    private PreloadIconDrawable applyProgressLevel(ItemInfoWithIcon info) {
+        if (info.isInactiveArchive()) {
             return null;
         }
 
@@ -1115,23 +1121,16 @@
             setContentDescription(getContext()
                     .getString(R.string.app_waiting_download_title, info.title));
         }
-        if (mIcon != null) {
-            PreloadIconDrawable preloadIconDrawable;
-            if (mIcon instanceof PreloadIconDrawable) {
-                preloadIconDrawable = (PreloadIconDrawable) mIcon;
-                preloadIconDrawable.setLevel(progressLevel);
-                preloadIconDrawable.setIsDisabled(isIconDisabled(info));
-            } else {
-                preloadIconDrawable = makePreloadIcon();
-                setIcon(preloadIconDrawable);
-                if (info.isArchived() && Flags.useNewIconForArchivedApps()) {
-                    // reapply text without cloud icon as soon as unarchiving is triggered
-                    applyLabel(info);
-                }
-            }
-            return preloadIconDrawable;
+        PreloadIconDrawable pid;
+        if (mIcon instanceof PreloadIconDrawable p) {
+            pid = p;
+            pid.setLevel(progressLevel);
+            pid.setIsDisabled(isIconDisabled(info));
+        } else {
+            pid = makePreloadIcon(info);
+            setIcon(pid);
         }
-        return null;
+        return pid;
     }
 
     /**
@@ -1140,11 +1139,11 @@
      */
     @Nullable
     public PreloadIconDrawable makePreloadIcon() {
-        if (!(getTag() instanceof ItemInfoWithIcon)) {
-            return null;
-        }
+        return getTag() instanceof ItemInfoWithIcon info ? makePreloadIcon(info) : null;
+    }
 
-        ItemInfoWithIcon info = (ItemInfoWithIcon) getTag();
+    @NonNull
+    private PreloadIconDrawable makePreloadIcon(ItemInfoWithIcon info) {
         int progressLevel = info.getProgressLevel();
         final PreloadIconDrawable preloadDrawable = newPendingIcon(getContext(), info);
 
@@ -1212,7 +1211,7 @@
                 setContentDescription(getContext().getString(
                         R.string.app_archived_title, info.title));
             }
-        } else if ((info.runtimeStatusFlags & ItemInfoWithIcon.FLAG_SHOW_DOWNLOAD_PROGRESS_MASK)
+        } else if ((info.runtimeStatusFlags & FLAG_SHOW_DOWNLOAD_PROGRESS_MASK)
                 != 0) {
             String percentageString = NumberFormat.getPercentInstance()
                     .format(progressLevel * 0.01);
diff --git a/src/com/android/launcher3/DeviceProfile.java b/src/com/android/launcher3/DeviceProfile.java
index 9f47da7..b4a24f1 100644
--- a/src/com/android/launcher3/DeviceProfile.java
+++ b/src/com/android/launcher3/DeviceProfile.java
@@ -873,7 +873,7 @@
         if (renderer == null) {
             renderer = new DotRenderer(
                     size,
-                    IconShape.INSTANCE.get(context).getShapeOverridePath(DEFAULT_DOT_SIZE),
+                    IconShape.INSTANCE.get(context).getShape().getPath(DEFAULT_DOT_SIZE),
                     DEFAULT_DOT_SIZE);
             cache.put(size, renderer);
         }
diff --git a/src/com/android/launcher3/Launcher.java b/src/com/android/launcher3/Launcher.java
index 7df4014..647d2ad 100644
--- a/src/com/android/launcher3/Launcher.java
+++ b/src/com/android/launcher3/Launcher.java
@@ -277,11 +277,11 @@
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
-import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
 import java.util.Optional;
+import java.util.Set;
 import java.util.function.Predicate;
 import java.util.function.Supplier;
 import java.util.stream.Stream;
@@ -2598,25 +2598,12 @@
         mModelCallbacks.bindIncrementalDownloadProgressUpdated(app);
     }
 
-    @Override
-    public void bindWidgetsRestored(ArrayList<LauncherAppWidgetInfo> widgets) {
-        mModelCallbacks.bindWidgetsRestored(widgets);
-    }
-
     /**
      * See {@code LauncherBindingDelegate}
      */
     @Override
-    public void bindWorkspaceItemsChanged(List<WorkspaceItemInfo> updated) {
-        mModelCallbacks.bindWorkspaceItemsChanged(updated);
-    }
-
-    /**
-     * See {@code LauncherBindingDelegate}
-     */
-    @Override
-    public void bindRestoreItemsChange(HashSet<ItemInfo> updates) {
-        mModelCallbacks.bindRestoreItemsChange(updates);
+    public void bindItemsUpdated(Set<ItemInfo> updates) {
+        mModelCallbacks.bindItemsUpdated(updates);
     }
 
     /**
diff --git a/src/com/android/launcher3/LauncherPrefs.kt b/src/com/android/launcher3/LauncherPrefs.kt
index d8bb84e..484cef4 100644
--- a/src/com/android/launcher3/LauncherPrefs.kt
+++ b/src/com/android/launcher3/LauncherPrefs.kt
@@ -58,11 +58,11 @@
             else encryptedContext.getSharedPreferences(sharedPrefFile, MODE_PRIVATE)
 
     /** Returns the value with type [T] for [item]. */
-    fun <T> get(item: ContextualItem<T>): T =
+    open fun <T> get(item: ContextualItem<T>): T =
         getInner(item, item.defaultValueFromContext(encryptedContext))
 
     /** Returns the value with type [T] for [item]. */
-    fun <T> get(item: ConstantItem<T>): T = getInner(item, item.defaultValue)
+    open fun <T> get(item: ConstantItem<T>): T = getInner(item, item.defaultValue)
 
     /**
      * Retrieves the value for an [Item] from [SharedPreferences]. It handles method typing via the
@@ -406,3 +406,20 @@
     ENCRYPTED,
     DEVICE_PROTECTED,
 }
+
+/**
+ * LauncherPrefs which delegates all lookup to [prefs] but uses the real prefs for initial values
+ */
+class ProxyPrefs(context: Context, private val prefs: SharedPreferences) : LauncherPrefs(context) {
+
+    private val realPrefs = LauncherPrefs(context)
+
+    override val Item.sharedPrefs: SharedPreferences
+        get() = prefs
+
+    override fun <T> get(item: ConstantItem<T>) =
+        super.get(backedUpItem(item.sharedPrefKey, realPrefs.get(item)))
+
+    override fun <T> get(item: ContextualItem<T>) =
+        super.get(backedUpItem(item.sharedPrefKey, realPrefs.get(item)))
+}
diff --git a/src/com/android/launcher3/ModelCallbacks.kt b/src/com/android/launcher3/ModelCallbacks.kt
index 5d32525..5338fb4 100644
--- a/src/com/android/launcher3/ModelCallbacks.kt
+++ b/src/com/android/launcher3/ModelCallbacks.kt
@@ -17,8 +17,6 @@
 import com.android.launcher3.model.StringCache
 import com.android.launcher3.model.data.AppInfo
 import com.android.launcher3.model.data.ItemInfo
-import com.android.launcher3.model.data.LauncherAppWidgetInfo
-import com.android.launcher3.model.data.WorkspaceItemInfo
 import com.android.launcher3.popup.PopupContainerWithArrow
 import com.android.launcher3.util.ComponentKey
 import com.android.launcher3.util.IntArray as LIntArray
@@ -215,29 +213,13 @@
         launcher.appsView.appsStore.updateProgressBar(app)
     }
 
-    override fun bindWidgetsRestored(widgets: ArrayList<LauncherAppWidgetInfo?>?) {
-        launcher.workspace.widgetsRestored(widgets)
-    }
-
-    /**
-     * Some shortcuts were updated in the background. Implementation of the method from
-     * LauncherModel.Callbacks.
-     *
-     * @param updated list of shortcuts which have changed.
-     */
-    override fun bindWorkspaceItemsChanged(updated: List<WorkspaceItemInfo?>) {
-        if (updated.isNotEmpty()) {
-            launcher.workspace.updateWorkspaceItems(updated, launcher)
-            PopupContainerWithArrow.dismissInvalidPopup(launcher)
-        }
-    }
-
     /**
      * Update the state of a package, typically related to install state. Implementation of the
      * method from LauncherModel.Callbacks.
      */
-    override fun bindRestoreItemsChange(updates: HashSet<ItemInfo?>?) {
-        launcher.workspace.updateRestoreItems(updates, launcher)
+    override fun bindItemsUpdated(updates: Set<ItemInfo>) {
+        launcher.workspace.updateContainerItems(updates, launcher)
+        PopupContainerWithArrow.dismissInvalidPopup(launcher)
     }
 
     /**
diff --git a/src/com/android/launcher3/Workspace.java b/src/com/android/launcher3/Workspace.java
index 97c0d9a..b41a425 100644
--- a/src/com/android/launcher3/Workspace.java
+++ b/src/com/android/launcher3/Workspace.java
@@ -53,8 +53,6 @@
 import android.graphics.PointF;
 import android.graphics.Rect;
 import android.graphics.drawable.Drawable;
-import android.os.Handler;
-import android.os.Message;
 import android.os.Parcelable;
 import android.util.AttributeSet;
 import android.util.Log;
@@ -125,13 +123,9 @@
 import com.android.launcher3.util.Thunk;
 import com.android.launcher3.util.WallpaperOffsetInterpolator;
 import com.android.launcher3.widget.LauncherAppWidgetHostView;
-import com.android.launcher3.widget.LauncherWidgetHolder;
-import com.android.launcher3.widget.LauncherWidgetHolder.ProviderChangedListener;
 import com.android.launcher3.widget.NavigableAppWidgetHostView;
 import com.android.launcher3.widget.PendingAddShortcutInfo;
 import com.android.launcher3.widget.PendingAddWidgetInfo;
-import com.android.launcher3.widget.PendingAppWidgetHostView;
-import com.android.launcher3.widget.WidgetManagerHelper;
 import com.android.launcher3.widget.util.WidgetSizes;
 import com.android.systemui.plugins.shared.LauncherOverlayManager.LauncherOverlayCallbacks;
 import com.android.systemui.plugins.shared.LauncherOverlayManager.LauncherOverlayTouchProxy;
@@ -664,9 +658,6 @@
             bindAndInitFirstWorkspaceScreen();
         }
 
-        // Remove any deferred refresh callbacks
-        mLauncher.mHandler.removeCallbacksAndMessages(DeferredWidgetRefresh.class);
-
         // Re-enable the layout transitions
         enableLayoutTransitions();
     }
@@ -3465,43 +3456,6 @@
         removeItemsByMatcher(matcher);
     }
 
-    public void widgetsRestored(final ArrayList<LauncherAppWidgetInfo> changedInfo) {
-        if (!changedInfo.isEmpty()) {
-            DeferredWidgetRefresh widgetRefresh = new DeferredWidgetRefresh(changedInfo,
-                    mLauncher.getAppWidgetHolder());
-
-            LauncherAppWidgetInfo item = changedInfo.get(0);
-            final AppWidgetProviderInfo widgetInfo;
-            WidgetManagerHelper widgetHelper = new WidgetManagerHelper(getContext());
-            if (item.hasRestoreFlag(LauncherAppWidgetInfo.FLAG_ID_NOT_VALID)) {
-                widgetInfo = widgetHelper.findProvider(item.providerName, item.user);
-            } else {
-                widgetInfo = widgetHelper.getLauncherAppWidgetInfo(item.appWidgetId,
-                        item.getTargetComponent());
-            }
-
-            if (widgetInfo != null) {
-                // Re-inflate the widgets which have changed status
-                widgetRefresh.run();
-            } else {
-                // widgetRefresh will automatically run when the packages are updated.
-                // For now just update the progress bars
-                mapOverItems(new ItemOperator() {
-                    @Override
-                    public boolean evaluate(ItemInfo info, View view) {
-                        if (view instanceof PendingAppWidgetHostView
-                                && changedInfo.contains(info)) {
-                            ((LauncherAppWidgetInfo) info).installProgress = 100;
-                            ((PendingAppWidgetHostView) view).applyState();
-                        }
-                        // process all the shortcuts
-                        return false;
-                    }
-                });
-            }
-        }
-    }
-
     public boolean isOverlayShown() {
         return mOverlayShown;
     }
@@ -3608,62 +3562,6 @@
         return mLauncher.getCellPosMapper();
     }
 
-    /**
-     * Used as a workaround to ensure that the AppWidgetService receives the
-     * PACKAGE_ADDED broadcast before updating widgets.
-     */
-    private class DeferredWidgetRefresh implements Runnable, ProviderChangedListener {
-        private final ArrayList<LauncherAppWidgetInfo> mInfos;
-        private final LauncherWidgetHolder mWidgetHolder;
-        private final Handler mHandler;
-
-        private boolean mRefreshPending;
-
-        DeferredWidgetRefresh(ArrayList<LauncherAppWidgetInfo> infos,
-                              LauncherWidgetHolder holder) {
-            mInfos = infos;
-            mWidgetHolder = holder;
-            mHandler = mLauncher.mHandler;
-            mRefreshPending = true;
-
-            mWidgetHolder.addProviderChangeListener(this);
-            // Force refresh after 10 seconds, if we don't get the provider changed event.
-            // This could happen when the provider is no longer available in the app.
-            Message msg = Message.obtain(mHandler, this);
-            msg.obj = DeferredWidgetRefresh.class;
-            mHandler.sendMessageDelayed(msg, 10000);
-        }
-
-        @Override
-        public void run() {
-            mWidgetHolder.removeProviderChangeListener(this);
-            mHandler.removeCallbacks(this);
-
-            if (!mRefreshPending) {
-                return;
-            }
-
-            mRefreshPending = false;
-
-            ArrayList<PendingAppWidgetHostView> views = new ArrayList<>(mInfos.size());
-            mapOverItems((info, view) -> {
-                if (view instanceof PendingAppWidgetHostView && mInfos.contains(info)) {
-                    views.add((PendingAppWidgetHostView) view);
-                }
-                // process all children
-                return false;
-            });
-            for (PendingAppWidgetHostView view : views) {
-                view.reInflate();
-            }
-        }
-
-        @Override
-        public void notifyWidgetProvidersChanged() {
-            run();
-        }
-    }
-
     private class StateTransitionListener extends AnimatorListenerAdapter
             implements AnimatorUpdateListener {
 
diff --git a/src/com/android/launcher3/allapps/AllAppsStore.java b/src/com/android/launcher3/allapps/AllAppsStore.java
index 9afe06c..d5a4022 100644
--- a/src/com/android/launcher3/allapps/AllAppsStore.java
+++ b/src/com/android/launcher3/allapps/AllAppsStore.java
@@ -17,7 +17,6 @@
 
 import static com.android.launcher3.model.data.AppInfo.COMPONENT_KEY_COMPARATOR;
 import static com.android.launcher3.model.data.AppInfo.EMPTY_ARRAY;
-import static com.android.launcher3.model.data.ItemInfoWithIcon.FLAG_SHOW_DOWNLOAD_PROGRESS_MASK;
 
 import android.content.Context;
 import android.os.UserHandle;
@@ -229,11 +228,7 @@
     public void updateProgressBar(AppInfo app) {
         updateAllIcons((child) -> {
             if (child.getTag() == app) {
-                if ((app.runtimeStatusFlags & FLAG_SHOW_DOWNLOAD_PROGRESS_MASK) == 0) {
-                    child.applyFromApplicationInfo(app);
-                } else {
-                    child.applyProgressLevel();
-                }
+                child.applyFromApplicationInfo(app);
             }
         });
     }
diff --git a/src/com/android/launcher3/dragndrop/DragView.java b/src/com/android/launcher3/dragndrop/DragView.java
index 9c82748..072673d 100644
--- a/src/com/android/launcher3/dragndrop/DragView.java
+++ b/src/com/android/launcher3/dragndrop/DragView.java
@@ -273,7 +273,11 @@
                 Rect shrunkBounds = new Rect(bounds);
                 Utilities.scaleRectAboutCenter(shrunkBounds, 0.98f);
                 adaptiveIcon.setBounds(shrunkBounds);
-                final Path mask = IconShape.INSTANCE.get(getContext()).getShapeOverridePath(w);
+
+                IconShape iconShape = IconShape.INSTANCE.get(getContext());
+                final Path mask = (adaptiveIcon instanceof FolderAdaptiveIcon
+                        ? iconShape.getFolderShape() : iconShape.getShape())
+                        .getPath(shrunkBounds);
 
                 mTranslateX = new SpringFloatValue(DragView.this,
                         w * AdaptiveIconDrawable.getExtraInsetFraction());
diff --git a/src/com/android/launcher3/dragndrop/FolderAdaptiveIcon.java b/src/com/android/launcher3/dragndrop/FolderAdaptiveIcon.java
index 6a43b24..929e52e 100644
--- a/src/com/android/launcher3/dragndrop/FolderAdaptiveIcon.java
+++ b/src/com/android/launcher3/dragndrop/FolderAdaptiveIcon.java
@@ -165,7 +165,7 @@
         Paint backgroundPaint = new Paint(Paint.ANTI_ALIAS_FLAG);
         backgroundPaint.setColor(bg.getBgColor());
         bg.drawShadow(backgroundCanvas);
-        backgroundCanvas.drawCircle(size / 2f, size / 2f, bg.getRadius(), backgroundPaint);
+        backgroundCanvas.drawPaint(backgroundPaint);
         bg.drawBackgroundStroke(backgroundCanvas);
     }
 
diff --git a/src/com/android/launcher3/folder/FolderAnimationManager.java b/src/com/android/launcher3/folder/FolderAnimationManager.java
index 588a6db..d2ff2cb 100644
--- a/src/com/android/launcher3/folder/FolderAnimationManager.java
+++ b/src/com/android/launcher3/folder/FolderAnimationManager.java
@@ -233,7 +233,7 @@
         }
         play(a, getAnimator(mFolder.mFooter, ALPHA, 0, 1f), footerStartDelay, footerAlphaDuration);
 
-        ShapeDelegate shapeDelegate = IconShape.INSTANCE.get(mContext).getShape();
+        ShapeDelegate shapeDelegate = IconShape.INSTANCE.get(mContext).getFolderShape();
         // Create reveal animator for the folder background
         play(a, shapeDelegate.createRevealAnimator(
                 mFolder, startRect, endRect, finalRadius, !mIsOpening));
diff --git a/src/com/android/launcher3/folder/FolderIcon.java b/src/com/android/launcher3/folder/FolderIcon.java
index 2481a1a..b7b378e 100644
--- a/src/com/android/launcher3/folder/FolderIcon.java
+++ b/src/com/android/launcher3/folder/FolderIcon.java
@@ -177,12 +177,16 @@
         FolderIcon icon = inflateIcon(resId, activityContext, group, folderInfo);
         folder.setFolderIcon(icon);
         folder.bind(folderInfo);
+
         icon.setFolder(folder);
+        folderInfo.addListener(icon);
         return icon;
     }
 
     /**
-     * Builds a FolderIcon to be added to the Launcher
+     * Builds a FolderIcon to be added to the activity.
+     * This method doesn't add any listeners to the FolderInfo, and hence any changes to the info
+     * will not be reflected in the folder.
      */
     public static FolderIcon inflateIcon(int resId, ActivityContext activity,
             @Nullable ViewGroup group, FolderInfo folderInfo) {
@@ -228,8 +232,6 @@
         icon.mPreviewVerifier.setFolderInfo(folderInfo);
         icon.updatePreviewItems(false);
 
-        folderInfo.addListener(icon);
-
         return icon;
     }
 
diff --git a/src/com/android/launcher3/folder/PreviewBackground.java b/src/com/android/launcher3/folder/PreviewBackground.java
index df41d47..d9c60db 100644
--- a/src/com/android/launcher3/folder/PreviewBackground.java
+++ b/src/com/android/launcher3/folder/PreviewBackground.java
@@ -261,7 +261,7 @@
     }
 
     private ShapeDelegate getShape() {
-        return IconShape.INSTANCE.get(mContext).getShape();
+        return IconShape.INSTANCE.get(mContext).getFolderShape();
     }
 
     public void drawShadow(Canvas canvas) {
diff --git a/src/com/android/launcher3/graphics/GridCustomizationsProvider.java b/src/com/android/launcher3/graphics/GridCustomizationsProvider.java
index ad176dc..5a2864d 100644
--- a/src/com/android/launcher3/graphics/GridCustomizationsProvider.java
+++ b/src/com/android/launcher3/graphics/GridCustomizationsProvider.java
@@ -16,10 +16,12 @@
 package com.android.launcher3.graphics;
 
 
-import static com.android.launcher3.graphics.IconShape.PREF_ICON_SHAPE;
+import static com.android.launcher3.graphics.ThemeManager.PREF_ICON_SHAPE;
 import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
 import static com.android.launcher3.util.Executors.UI_HELPER_EXECUTOR;
 
+import static java.util.Objects.requireNonNullElse;
+
 import android.content.ContentProvider;
 import android.content.ContentValues;
 import android.content.Context;
@@ -51,10 +53,10 @@
 import com.android.launcher3.util.RunnableList;
 import com.android.systemui.shared.Flags;
 
+import java.lang.ref.WeakReference;
 import java.util.Collections;
-import java.util.List;
 import java.util.Set;
-import java.util.WeakHashMap;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ExecutionException;
 
 /**
@@ -120,7 +122,7 @@
 
     // Set of all active previews used to track duplicate memory allocations
     private final Set<PreviewLifecycleObserver> mActivePreviews =
-            Collections.newSetFromMap(new WeakHashMap<>());
+            Collections.newSetFromMap(new ConcurrentHashMap<>());
 
     @Override
     public boolean onCreate() {
@@ -141,23 +143,15 @@
                 if (Flags.newCustomizationPickerUi()) {
                     MatrixCursor cursor = new MatrixCursor(new String[]{
                             KEY_SHAPE_KEY, KEY_SHAPE_TITLE, KEY_PATH, KEY_IS_DEFAULT});
-                    List<IconShapeModel> shapes = IconShapesProvider.INSTANCE.getShapes()
-                            .values()
-                            .stream()
-                            .toList();
-                    String currentPath = LauncherPrefs.get(context).get(PREF_ICON_SHAPE);
-                    IconShapeModel currentShape = shapes.stream()
-                            .filter(shape -> currentPath.equals(shape.getPathString()))
-                            .findFirst()
-                            .orElse(IconShapesProvider.INSTANCE.getShapes().get("circle"));
-
-                    for (int i = 0; i < shapes.size(); i++) {
-                        IconShapeModel shape = shapes.get(i);
+                    String currentShapePath =
+                            ThemeManager.INSTANCE.get(context).getIconState().getIconMask();
+                    for (IconShapeModel shape : IconShapesProvider.INSTANCE.getShapes().values()) {
                         cursor.newRow()
                                 .add(KEY_SHAPE_KEY, shape.getKey())
                                 .add(KEY_SHAPE_TITLE, shape.getTitle())
                                 .add(KEY_PATH, shape.getPathString())
-                                .add(KEY_IS_DEFAULT, shape.equals(currentShape));
+                                .add(KEY_IS_DEFAULT,
+                                        shape.getPathString().equals(currentShapePath));
                     }
                     return cursor;
                 } else  {
@@ -218,9 +212,8 @@
         switch (path) {
             case KEY_DEFAULT_GRID: {
                 if (Flags.newCustomizationPickerUi()) {
-                    String shapeKey = values.getAsString(KEY_SHAPE_KEY);
-                    IconShapeModel shape = IconShapesProvider.INSTANCE.getShapes().get(shapeKey);
-                    IconShape.INSTANCE.get(context).setShapeOverride(shape);
+                    LauncherPrefs.INSTANCE.get(context).put(PREF_ICON_SHAPE,
+                            requireNonNullElse(values.getAsString(KEY_SHAPE_KEY), ""));
                 }
                 String gridName = values.getAsString(KEY_NAME);
                 InvariantDeviceProfile idp = InvariantDeviceProfile.INSTANCE.get(context);
@@ -321,8 +314,15 @@
             Bundle result = new Bundle();
             result.putParcelable(KEY_SURFACE_PACKAGE, renderer.getSurfacePackage());
 
-            Messenger messenger =
-                    new Messenger(new Handler(UI_HELPER_EXECUTOR.getLooper(), observer));
+            mActivePreviews.add(observer);
+            lifeCycleTracker.add(() -> mActivePreviews.remove(observer));
+
+            // Wrap the callback in a weak reference. This ensures that the callback is not kept
+            // alive due to the Messenger's IBinder
+            Messenger messenger = new Messenger(new Handler(
+                    UI_HELPER_EXECUTOR.getLooper(),
+                    new WeakCallbackWrapper(observer)));
+
             Message msg = Message.obtain();
             msg.replyTo = messenger;
             result.putParcelable(KEY_CALLBACK, msg);
@@ -362,9 +362,7 @@
                     if (Flags.newCustomizationPickerUi()
                             && com.android.launcher3.Flags.enableLauncherIconShapes()) {
                         String shapeKey = message.getData().getString(KEY_SHAPE_KEY);
-                        IconShapeModel shape =
-                                IconShapesProvider.INSTANCE.getShapes().get(shapeKey);
-                        renderer.updateShape(shape);
+                        renderer.updateShape(shapeKey);
                     }
                     break;
                 case MESSAGE_ID_UPDATE_GRID:
@@ -402,4 +400,34 @@
                     && plo.renderer.getDisplayId() == renderer.getDisplayId();
         }
     }
+
+    /**
+     * A WeakReference wrapper around Handler.Callback to avoid passing hard-reference over IPC
+     * when using a Messenger
+     */
+    private static class WeakCallbackWrapper implements Handler.Callback {
+
+        private final WeakReference<Handler.Callback> mActual;
+        private final Message mCleanupMessage;
+
+        WeakCallbackWrapper(Handler.Callback actual) {
+            mActual = new WeakReference<>(actual);
+            mCleanupMessage = new Message();
+        }
+
+        @Override
+        public boolean handleMessage(Message message) {
+            Handler.Callback actual = mActual.get();
+            return actual != null && actual.handleMessage(message);
+        }
+
+        @Override
+        protected void finalize() throws Throwable {
+            super.finalize();
+            Handler.Callback actual = mActual.get();
+            if (actual != null) {
+                actual.handleMessage(mCleanupMessage);
+            }
+        }
+    }
 }
diff --git a/src/com/android/launcher3/graphics/IconShape.kt b/src/com/android/launcher3/graphics/IconShape.kt
index 2c4d8e4..e62936c 100644
--- a/src/com/android/launcher3/graphics/IconShape.kt
+++ b/src/com/android/launcher3/graphics/IconShape.kt
@@ -18,7 +18,6 @@
 import android.animation.Animator
 import android.animation.AnimatorListenerAdapter
 import android.animation.ValueAnimator
-import android.content.Context
 import android.graphics.Canvas
 import android.graphics.Color
 import android.graphics.Matrix
@@ -41,19 +40,12 @@
 import androidx.graphics.shapes.SvgPathParser
 import androidx.graphics.shapes.toPath
 import androidx.graphics.shapes.transformed
-import com.android.launcher3.EncryptionType
-import com.android.launcher3.LauncherPrefs
-import com.android.launcher3.LauncherPrefs.Companion.backedUpItem
 import com.android.launcher3.anim.RoundedRectRevealOutlineProvider
-import com.android.launcher3.dagger.ApplicationContext
 import com.android.launcher3.dagger.LauncherAppComponent
 import com.android.launcher3.dagger.LauncherAppSingleton
-import com.android.launcher3.graphics.LauncherPreviewRenderer.PreviewContext
 import com.android.launcher3.graphics.ThemeManager.ThemeChangeListener
 import com.android.launcher3.icons.GraphicsUtils
-import com.android.launcher3.icons.IconNormalizer
-import com.android.launcher3.shapes.IconShapeModel
-import com.android.launcher3.shapes.IconShapesProvider
+import com.android.launcher3.icons.IconNormalizer.normalizeAdaptiveIcon
 import com.android.launcher3.util.DaggerSingletonObject
 import com.android.launcher3.util.DaggerSingletonTracker
 import com.android.launcher3.views.ClipPathView
@@ -63,56 +55,52 @@
 @LauncherAppSingleton
 class IconShape
 @Inject
-constructor(
-    @ApplicationContext private val context: Context,
-    private val prefs: LauncherPrefs,
-    private val themeManager: ThemeManager,
-    lifeCycle: DaggerSingletonTracker,
-) {
+constructor(private val themeManager: ThemeManager, lifeCycle: DaggerSingletonTracker) {
 
-    var shapeOverride: IconShapeModel? = getShapeFromPathString(prefs.get(PREF_ICON_SHAPE))
-        set(value) {
-            field = value
-            if (context !is PreviewContext) {
-                value?.let { prefs.put(PREF_ICON_SHAPE, value.pathString) }
-            }
-        }
+    val normalizationScale =
+        normalizeAdaptiveIcon(
+            AdaptiveIconDrawable(null, ColorDrawable(Color.BLACK)),
+            AREA_CALC_SIZE,
+        )
 
-    var normalizationScale: Float = IconNormalizer.ICON_VISIBLE_AREA_FACTOR
+    var shape: ShapeDelegate = pickBestShape(themeManager.iconState.iconMask)
         private set
 
-    var shape: ShapeDelegate = pickBestShape(themeManager)
+    var folderShape: ShapeDelegate =
+        themeManager.iconState.run {
+            if (folderShapeMask == iconMask || folderShapeMask.isEmpty()) shape
+            else pickBestShape(folderShapeMask)
+        }
         private set
 
     init {
-        val changeListener = ThemeChangeListener { shape = pickBestShape(themeManager) }
+        val changeListener = ThemeChangeListener {
+            shape = pickBestShape(themeManager.iconState.iconMask)
+            folderShape =
+                themeManager.iconState.run {
+                    if (folderShapeMask == iconMask || folderShapeMask.isEmpty()) shape
+                    else pickBestShape(folderShapeMask)
+                }
+        }
         themeManager.addChangeListener(changeListener)
         lifeCycle.addCloseable { themeManager.removeChangeListener(changeListener) }
     }
 
-    fun getShapeOverridePath(pathSize: Float = DEFAULT_PATH_SIZE): Path {
-        val path = PathParser.createPathFromPathData(themeManager.iconState.iconMask)
-        if (pathSize != DEFAULT_PATH_SIZE) {
-            val matrix = Matrix()
-            val scale: Float = pathSize / DEFAULT_PATH_SIZE
-            matrix.setScale(scale, scale)
-            path.transform(matrix)
-        }
-        return path
-    }
+    interface ShapeDelegate {
+        fun getPath(pathSize: Float = DEFAULT_PATH_SIZE) =
+            Path().apply { addToPath(this, 0f, 0f, pathSize / 2) }
 
-    /** Initializes the shape which is closest to the [AdaptiveIconDrawable] */
-    private fun pickBestShape(themeManager: ThemeManager): ShapeDelegate {
-        val drawable =
-            AdaptiveIconDrawable(null, ColorDrawable(Color.BLACK)).apply {
-                setBounds(0, 0, AREA_CALC_SIZE, AREA_CALC_SIZE)
+        fun getPath(bounds: Rect) =
+            Path().apply {
+                addToPath(
+                    this,
+                    bounds.left.toFloat(),
+                    bounds.top.toFloat(),
+                    // Radius is half of the average size of the icon
+                    (bounds.width() + bounds.height()) / 4f,
+                )
             }
 
-        normalizationScale = IconNormalizer.normalizeAdaptiveIcon(drawable, AREA_CALC_SIZE)
-        return pickBestShape(drawable.iconMask, themeManager.iconState.iconMask)
-    }
-
-    interface ShapeDelegate {
         fun drawShape(canvas: Canvas, offsetX: Float, offsetY: Float, radius: Float, paint: Paint)
 
         fun addToPath(path: Path, offsetX: Float, offsetY: Float, radius: Float)
@@ -126,7 +114,6 @@
         ): ValueAnimator where T : View, T : ClipPathView
     }
 
-    @VisibleForTesting
     class Circle : RoundedSquare(1f) {
 
         override fun drawShape(
@@ -217,14 +204,24 @@
             paint: Paint,
         ) {
             tmpPath.reset()
-            addToPath(tmpPath, offsetX, offsetY, radius)
+            addToPath(tmpPath, offsetX, offsetY, radius, tmpMatrix)
             canvas.drawPath(tmpPath, paint)
         }
 
         override fun addToPath(path: Path, offsetX: Float, offsetY: Float, radius: Float) {
-            tmpMatrix.setScale(radius / 50, radius / 50)
-            tmpMatrix.postTranslate(offsetX, offsetY)
-            basePath.transform(tmpMatrix, path)
+            addToPath(path, offsetX, offsetY, radius, Matrix())
+        }
+
+        private fun addToPath(
+            path: Path,
+            offsetX: Float,
+            offsetY: Float,
+            radius: Float,
+            matrix: Matrix,
+        ) {
+            matrix.setScale(radius / 50, radius / 50)
+            matrix.postTranslate(offsetX, offsetY)
+            basePath.transform(matrix, path)
         }
 
         override fun <T> createRevealAnimator(
@@ -292,20 +289,11 @@
         @JvmField var INSTANCE = DaggerSingletonObject(LauncherAppComponent::getIconShape)
 
         const val TAG = "IconShape"
-        const val KEY_ICON_SHAPE = "icon_shape"
         const val DEFAULT_PATH_SIZE = 100f
         const val AREA_CALC_SIZE = 1000
         // .1% error margin
         const val AREA_DIFF_THRESHOLD = AREA_CALC_SIZE * AREA_CALC_SIZE / 1000
 
-        @JvmField val PREF_ICON_SHAPE = backedUpItem(KEY_ICON_SHAPE, "", EncryptionType.ENCRYPTED)
-
-        private fun getShapeFromPathString(pathString: String): IconShapeModel? {
-            return IconShapesProvider.shapes.values.firstOrNull { shape: IconShapeModel ->
-                shape.pathString == pathString
-            }
-        }
-
         /** Returns a function to calculate area diff from [base] */
         @VisibleForTesting
         fun areaDiffCalculator(base: Path): (ShapeDelegate) -> Int {
@@ -324,6 +312,26 @@
         }
 
         @VisibleForTesting
+        fun pickBestShape(shapeStr: String): ShapeDelegate {
+            val baseShape =
+                if (shapeStr.isNotEmpty()) {
+                    PathParser.createPathFromPathData(shapeStr).apply {
+                        transform(
+                            Matrix().apply {
+                                setScale(AREA_CALC_SIZE / 100f, AREA_CALC_SIZE / 100f)
+                            }
+                        )
+                    }
+                } else {
+                    AdaptiveIconDrawable(null, ColorDrawable(Color.BLACK)).let {
+                        it.setBounds(0, 0, AREA_CALC_SIZE, AREA_CALC_SIZE)
+                        it.iconMask
+                    }
+                }
+            return pickBestShape(baseShape, shapeStr)
+        }
+
+        @VisibleForTesting
         fun pickBestShape(baseShape: Path, shapeStr: String): ShapeDelegate {
             val calcAreaDiff = areaDiffCalculator(baseShape)
 
diff --git a/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java b/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java
index a0b73ae..911064c 100644
--- a/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java
+++ b/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java
@@ -67,7 +67,9 @@
 import com.android.launcher3.InsettableFrameLayout;
 import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.LauncherAppState;
+import com.android.launcher3.LauncherPrefs;
 import com.android.launcher3.LauncherSettings.Favorites;
+import com.android.launcher3.ProxyPrefs;
 import com.android.launcher3.R;
 import com.android.launcher3.Workspace;
 import com.android.launcher3.WorkspaceLayoutManager;
@@ -75,6 +77,9 @@
 import com.android.launcher3.celllayout.CellLayoutLayoutParams;
 import com.android.launcher3.celllayout.CellPosMapper;
 import com.android.launcher3.config.FeatureFlags;
+import com.android.launcher3.dagger.LauncherAppComponent;
+import com.android.launcher3.dagger.LauncherAppModule;
+import com.android.launcher3.dagger.LauncherAppSingleton;
 import com.android.launcher3.folder.FolderIcon;
 import com.android.launcher3.model.BgDataModel;
 import com.android.launcher3.model.BgDataModel.FixedContainerItems;
@@ -100,12 +105,16 @@
 import com.android.launcher3.widget.util.WidgetSizes;
 import com.android.systemui.shared.Flags;
 
+import dagger.BindsInstance;
+import dagger.Component;
+
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Map.Entry;
+import java.util.UUID;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
 
@@ -126,12 +135,25 @@
      */
     public static class PreviewContext extends SandboxContext {
 
+        private final String mPrefName;
+
         public PreviewContext(Context base, InvariantDeviceProfile idp) {
             super(base);
+            mPrefName = "preview-" + UUID.randomUUID().toString();
+            initDaggerComponent(DaggerLauncherPreviewRenderer_PreviewAppComponent.builder()
+                    .bindPrefs(new ProxyPrefs(
+                            this, getSharedPreferences(mPrefName, MODE_PRIVATE))));
+
             putObject(InvariantDeviceProfile.INSTANCE, idp);
             putObject(LauncherAppState.INSTANCE,
                     new LauncherAppState(this, null /* iconCacheFileName */));
         }
+
+        @Override
+        protected void cleanUpObjects() {
+            super.cleanUpObjects();
+            deleteSharedPreferences(mPrefName);
+        }
     }
 
     private final List<OnDeviceProfileChangeListener> mDpChangeListeners = new ArrayList<>();
@@ -578,4 +600,16 @@
             return true;
         }
     }
+
+    @LauncherAppSingleton
+    @Component(modules = LauncherAppModule.class)
+    public interface PreviewAppComponent extends LauncherAppComponent {
+
+        /** Builder for NexusLauncherAppComponent. */
+        @Component.Builder
+        interface Builder extends LauncherAppComponent.Builder {
+            @BindsInstance Builder bindPrefs(LauncherPrefs prefs);
+            PreviewAppComponent build();
+        }
+    }
 }
diff --git a/src/com/android/launcher3/graphics/PreloadIconDrawable.java b/src/com/android/launcher3/graphics/PreloadIconDrawable.java
index 3464e9b..50d6d1c 100644
--- a/src/com/android/launcher3/graphics/PreloadIconDrawable.java
+++ b/src/com/android/launcher3/graphics/PreloadIconDrawable.java
@@ -24,7 +24,6 @@
 import android.animation.AnimatorListenerAdapter;
 import android.animation.ObjectAnimator;
 import android.content.Context;
-import android.graphics.Bitmap;
 import android.graphics.Canvas;
 import android.graphics.Matrix;
 import android.graphics.Paint;
@@ -33,12 +32,14 @@
 import android.graphics.Rect;
 import android.util.Property;
 
+import androidx.annotation.VisibleForTesting;
 import androidx.core.graphics.ColorUtils;
 
 import com.android.launcher3.R;
 import com.android.launcher3.Utilities;
 import com.android.launcher3.anim.AnimatedFloat;
 import com.android.launcher3.anim.AnimatorListeners;
+import com.android.launcher3.icons.BitmapInfo;
 import com.android.launcher3.icons.FastBitmapDrawable;
 import com.android.launcher3.model.data.ItemInfoWithIcon;
 import com.android.launcher3.util.Themes;
@@ -63,8 +64,6 @@
 
     private static final int DEFAULT_PATH_SIZE = 100;
     private static final int MAX_PAINT_ALPHA = 255;
-    private static final int TRACK_ALPHA = (int) (0.27f * MAX_PAINT_ALPHA);
-    private static final int DISABLED_ICON_ALPHA = (int) (0.6f * MAX_PAINT_ALPHA);
 
     private static final long DURATION_SCALE = 500;
     private static final long SCALE_AND_ALPHA_ANIM_DURATION = 500;
@@ -120,7 +119,7 @@
                 IconPalette.getPreloadProgressColor(context, info.bitmap.color),
                 getPreloadColors(context),
                 Utilities.isDarkTheme(context),
-                IconShape.INSTANCE.get(context).getShapeOverridePath(DEFAULT_PATH_SIZE)
+                IconShape.INSTANCE.get(context).getShape().getPath(DEFAULT_PATH_SIZE)
         );
     }
 
@@ -284,20 +283,25 @@
                     (long) ((finalProgress - mInternalStateProgress) * DURATION_SCALE));
             mCurrentAnim.setInterpolator(LINEAR);
             if (isFinish) {
-                if (onFinishCallback != null) {
-                    mCurrentAnim.addListener(AnimatorListeners.forEndCallback(onFinishCallback));
-                }
                 mCurrentAnim.addListener(new AnimatorListenerAdapter() {
                     @Override
                     public void onAnimationEnd(Animator animation) {
                         mRanFinishAnimation = true;
                     }
                 });
+                if (onFinishCallback != null) {
+                    mCurrentAnim.addListener(AnimatorListeners.forEndCallback(onFinishCallback));
+                }
             }
             mCurrentAnim.start();
         }
     }
 
+    @VisibleForTesting
+    public ObjectAnimator getActiveAnimation() {
+        return mCurrentAnim;
+    }
+
     /**
      * Sets the internal progress and updates the UI accordingly
      *   for progress <= 0:
@@ -358,8 +362,7 @@
     @Override
     public FastBitmapConstantState newConstantState() {
         return new PreloadIconConstantState(
-                mBitmap,
-                mIconColor,
+                mBitmapInfo,
                 mItem,
                 mIndicatorColor,
                 new int[] {mSystemAccentColor, mSystemBackgroundColor},
@@ -377,14 +380,13 @@
         private final Path mShapePath;
 
         public PreloadIconConstantState(
-                Bitmap bitmap,
-                int iconColor,
+                BitmapInfo bitmapInfo,
                 ItemInfoWithIcon info,
                 int indicatorColor,
                 int[] preloadColors,
                 boolean isDarkMode,
                 Path shapePath) {
-            super(bitmap, iconColor);
+            super(bitmapInfo);
             mInfo = info;
             mIndicatorColor = indicatorColor;
             mPreloadColors = preloadColors;
diff --git a/src/com/android/launcher3/graphics/PreviewSurfaceRenderer.java b/src/com/android/launcher3/graphics/PreviewSurfaceRenderer.java
index 6afac71..7a60814 100644
--- a/src/com/android/launcher3/graphics/PreviewSurfaceRenderer.java
+++ b/src/com/android/launcher3/graphics/PreviewSurfaceRenderer.java
@@ -21,6 +21,7 @@
 import static android.view.Display.DEFAULT_DISPLAY;
 
 import static com.android.launcher3.LauncherSettings.Favorites.TABLE_NAME;
+import static com.android.launcher3.graphics.ThemeManager.PREF_ICON_SHAPE;
 import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
 import static com.android.launcher3.util.Executors.MODEL_EXECUTOR;
 
@@ -52,6 +53,7 @@
 import com.android.launcher3.DeviceProfile;
 import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.LauncherAppState;
+import com.android.launcher3.LauncherPrefs;
 import com.android.launcher3.LauncherSettings;
 import com.android.launcher3.Workspace;
 import com.android.launcher3.graphics.LauncherPreviewRenderer.PreviewContext;
@@ -62,7 +64,6 @@
 import com.android.launcher3.model.LoaderTask;
 import com.android.launcher3.model.ModelDbController;
 import com.android.launcher3.provider.LauncherDbUtils;
-import com.android.launcher3.shapes.IconShapeModel;
 import com.android.launcher3.util.ComponentKey;
 import com.android.launcher3.util.RunnableList;
 import com.android.launcher3.util.Themes;
@@ -71,6 +72,7 @@
 
 import java.util.ArrayList;
 import java.util.Map;
+import java.util.Objects;
 import java.util.concurrent.TimeUnit;
 
 /** Render preview using surface view. */
@@ -91,10 +93,10 @@
     private Context mContext;
     private SparseIntArray mPreviewColorOverride;
     private String mGridName;
-    private IconShapeModel mShape;
+    private String mShapeKey;
+
     @Nullable private Boolean mDarkMode;
     private boolean mDestroyed = false;
-    private LauncherPreviewRenderer mRenderer;
     private boolean mHideQsb;
     @Nullable private FrameLayout mViewRoot = null;
 
@@ -220,14 +222,14 @@
     /**
      * Update the shapes of the launcher preview
      *
-     * @param shape path for shapes
+     * @param shapeKey key for the IconShape model
      */
-    public void updateShape(@NonNull IconShapeModel shape) {
-        if (shape.equals(mShape)) {
-            Log.w(TAG, "Preview shape already set, skipping. shape=" + shape);
+    public void updateShape(@Nullable String shapeKey) {
+        if (Objects.equals(mShapeKey, shapeKey)) {
+            Log.w(TAG, "Preview shape already set, skipping. shape=" + mShapeKey);
             return;
         }
-        mShape = shape;
+        mShapeKey = shapeKey;
         loadAsync();
     }
 
@@ -237,9 +239,8 @@
      * @param hide True to hide and false to show.
      */
     public void hideBottomRow(boolean hide) {
-        if (mRenderer != null) {
-            mRenderer.hideBottomRow(hide);
-        }
+        mHideQsb = hide;
+        loadAsync();
     }
 
     /**
@@ -317,11 +318,11 @@
         final Context inflationContext = getPreviewContext();
         final InvariantDeviceProfile idp = new InvariantDeviceProfile(inflationContext, mGridName);
         if (GridSizeMigrationDBController.needsToMigrate(inflationContext, idp)
-                || mShape != null) {
+                || mShapeKey != null) {
             // Start the migration
             PreviewContext previewContext = new PreviewContext(inflationContext, idp);
-            if (mShape != null) {
-                IconShape.INSTANCE.get(previewContext).setShapeOverride(mShape);
+            if (mShapeKey != null) {
+                LauncherPrefs.INSTANCE.get(previewContext).put(PREF_ICON_SHAPE, mShapeKey);
             }
             // Copy existing data to preview DB
             LauncherDbUtils.copyTable(LauncherAppState.getInstance(mContext)
@@ -385,15 +386,16 @@
         if (mDestroyed) {
             return;
         }
+        LauncherPreviewRenderer renderer;
         if (Flags.newCustomizationPickerUi()) {
-            mRenderer = new LauncherPreviewRenderer(inflationContext, idp, mPreviewColorOverride,
+            renderer = new LauncherPreviewRenderer(inflationContext, idp, mPreviewColorOverride,
                     mWallpaperColors, launcherWidgetSpanInfo);
         } else {
-            mRenderer = new LauncherPreviewRenderer(inflationContext, idp,
+            renderer = new LauncherPreviewRenderer(inflationContext, idp,
                     mWallpaperColors, launcherWidgetSpanInfo);
         }
-        mRenderer.hideBottomRow(mHideQsb);
-        View view = mRenderer.getRenderedView(dataModel, widgetProviderInfoMap);
+        renderer.hideBottomRow(mHideQsb);
+        View view = renderer.getRenderedView(dataModel, widgetProviderInfoMap);
         // This aspect scales the view to fit in the surface and centers it
         final float scale = Math.min(mWidth / (float) view.getMeasuredWidth(),
                 mHeight / (float) view.getMeasuredHeight());
diff --git a/src/com/android/launcher3/graphics/ThemeManager.kt b/src/com/android/launcher3/graphics/ThemeManager.kt
index 989471f..f24c2ab 100644
--- a/src/com/android/launcher3/graphics/ThemeManager.kt
+++ b/src/com/android/launcher3/graphics/ThemeManager.kt
@@ -25,10 +25,9 @@
 import com.android.launcher3.dagger.ApplicationContext
 import com.android.launcher3.dagger.LauncherAppComponent
 import com.android.launcher3.dagger.LauncherAppSingleton
-import com.android.launcher3.graphics.IconShape.Companion.KEY_ICON_SHAPE
-import com.android.launcher3.graphics.IconShape.Companion.PREF_ICON_SHAPE
 import com.android.launcher3.icons.IconThemeController
 import com.android.launcher3.icons.mono.MonoIconThemeController
+import com.android.launcher3.shapes.IconShapesProvider
 import com.android.launcher3.util.DaggerSingletonObject
 import com.android.launcher3.util.DaggerSingletonTracker
 import com.android.launcher3.util.Executors.MAIN_EXECUTOR
@@ -93,20 +92,26 @@
     fun removeChangeListener(listener: ThemeChangeListener) = listeners.remove(listener)
 
     private fun parseIconState(): IconState {
-        val shapeOverride = prefs.get(PREF_ICON_SHAPE)
+        val shapeModel =
+            prefs.get(PREF_ICON_SHAPE).let { shapeOverride ->
+                IconShapesProvider.shapes.values.firstOrNull { it.key == shapeOverride }
+            }
+        val iconMask =
+            when {
+                shapeModel != null -> shapeModel.pathString
+                CONFIG_ICON_MASK_RES_ID == Resources.ID_NULL -> ""
+                else -> context.resources.getString(CONFIG_ICON_MASK_RES_ID)
+            }
         return IconState(
-            iconMask =
-                when {
-                    shapeOverride.isNotEmpty() -> shapeOverride
-                    CONFIG_ICON_MASK_RES_ID == Resources.ID_NULL -> ""
-                    else -> context.resources.getString(CONFIG_ICON_MASK_RES_ID)
-                },
+            iconMask = iconMask,
+            folderShapeMask = shapeModel?.folderPathString ?: iconMask,
             isMonoTheme = isMonoThemeEnabled,
         )
     }
 
     data class IconState(
         val iconMask: String,
+        val folderShapeMask: String,
         val isMonoTheme: Boolean,
         val themeCode: String = if (isMonoTheme) "with-theme" else "no-theme",
     ) {
@@ -121,9 +126,11 @@
     companion object {
 
         @JvmField val INSTANCE = DaggerSingletonObject(LauncherAppComponent::getThemeManager)
+        const val KEY_ICON_SHAPE = "icon_shape_model"
 
         const val KEY_THEMED_ICONS = "themed_icons"
         @JvmField val THEMED_ICONS = backedUpItem(KEY_THEMED_ICONS, false, EncryptionType.ENCRYPTED)
+        @JvmField val PREF_ICON_SHAPE = backedUpItem(KEY_ICON_SHAPE, "", EncryptionType.ENCRYPTED)
 
         private const val ACTION_OVERLAY_CHANGED = "android.intent.action.OVERLAY_CHANGED"
         private val CONFIG_ICON_MASK_RES_ID: Int =
diff --git a/src/com/android/launcher3/icons/LauncherIconProvider.java b/src/com/android/launcher3/icons/LauncherIconProvider.java
index e40f526..482360c 100644
--- a/src/com/android/launcher3/icons/LauncherIconProvider.java
+++ b/src/com/android/launcher3/icons/LauncherIconProvider.java
@@ -19,14 +19,18 @@
 import android.content.pm.ApplicationInfo;
 import android.content.res.Resources;
 import android.content.res.XmlResourceParser;
+import android.graphics.drawable.AdaptiveIconDrawable;
+import android.graphics.drawable.Drawable;
 import android.text.TextUtils;
 import android.util.ArrayMap;
 import android.util.Log;
 
 import androidx.annotation.NonNull;
+import androidx.annotation.Nullable;
 
 import com.android.launcher3.R;
 import com.android.launcher3.config.FeatureFlags;
+import com.android.launcher3.graphics.IconShape;
 import com.android.launcher3.graphics.ThemeManager;
 import com.android.launcher3.util.ApiWrapper;
 
@@ -49,9 +53,14 @@
 
     private Map<String, ThemeData> mThemedIconMap;
 
+    private final ApiWrapper mApiWrapper;
+    private final IconShape mIconShape;
+
     public LauncherIconProvider(Context context) {
         super(context);
         setIconThemeSupported(ThemeManager.INSTANCE.get(context).isMonoThemeEnabled());
+        mApiWrapper = ApiWrapper.INSTANCE.get(context);
+        mIconShape = IconShape.INSTANCE.get(context);
     }
 
     /**
@@ -75,7 +84,25 @@
 
     @Override
     protected String getApplicationInfoHash(@NonNull ApplicationInfo appInfo) {
-        return ApiWrapper.INSTANCE.get(mContext).getApplicationInfoHash(appInfo);
+        return mApiWrapper.getApplicationInfoHash(appInfo);
+    }
+
+    @Nullable
+    @Override
+    protected Drawable loadAppInfoIcon(ApplicationInfo info, Resources resources, int density) {
+        // Tries to load the round icon res, if the app defines it as an adaptive icon
+        if (mIconShape.getShape() instanceof IconShape.Circle) {
+            int roundIconRes = mApiWrapper.getRoundIconRes(info);
+            if (roundIconRes != 0 && roundIconRes != info.icon) {
+                try {
+                    Drawable d = resources.getDrawableForDensity(roundIconRes, density);
+                    if (d instanceof AdaptiveIconDrawable) {
+                        return d;
+                    }
+                } catch (Resources.NotFoundException exc) { }
+            }
+        }
+        return super.loadAppInfoIcon(info, resources, density);
     }
 
     private Map<String, ThemeData> getThemedIconMap() {
diff --git a/src/com/android/launcher3/icons/LauncherIcons.java b/src/com/android/launcher3/icons/LauncherIcons.java
index 2ffbeb8..59fff62 100644
--- a/src/com/android/launcher3/icons/LauncherIcons.java
+++ b/src/com/android/launcher3/icons/LauncherIcons.java
@@ -17,22 +17,18 @@
 package com.android.launcher3.icons;
 
 import android.content.Context;
-import android.graphics.Matrix;
 import android.graphics.Path;
 import android.graphics.Rect;
-import android.graphics.RectF;
 import android.graphics.drawable.AdaptiveIconDrawable;
 import android.os.UserHandle;
 
 import androidx.annotation.NonNull;
-import androidx.core.graphics.PathParser;
 
 import com.android.launcher3.Flags;
 import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.graphics.IconShape;
 import com.android.launcher3.graphics.ThemeManager;
 import com.android.launcher3.pm.UserCache;
-import com.android.launcher3.shapes.IconShapeModel;
 import com.android.launcher3.util.MainThreadInitializedObject;
 import com.android.launcher3.util.SafeCloseable;
 import com.android.launcher3.util.UserIconInfo;
@@ -87,22 +83,7 @@
     @Override
     public Path getShapePath(AdaptiveIconDrawable drawable, Rect iconBounds) {
         if (!Flags.enableLauncherIconShapes()) return drawable.getIconMask();
-
-        IconShapeModel shapeOverride = IconShape.INSTANCE.get(mContext).getShapeOverride();
-        if (shapeOverride != null) {
-            Path maskPath = PathParser.createPathFromPathData(shapeOverride.getPathString());
-            Matrix matrix = new Matrix();
-            // Assuming Path is in [0, 0, 100, 100] coordinate space.
-            matrix.setRectToRect(
-                    new RectF(0, 0, 100, 100),
-                    new RectF(iconBounds),
-                    Matrix.ScaleToFit.CENTER // Todo: CENTER or FILL?
-            );
-            maskPath.transform(matrix);
-            return maskPath;
-        } else {
-            return drawable.getIconMask();
-        }
+        return IconShape.INSTANCE.get(mContext).getShape().getPath(iconBounds);
     }
 
     @Override
diff --git a/src/com/android/launcher3/model/BgDataModel.java b/src/com/android/launcher3/model/BgDataModel.java
index a04cbfb..ddc775d 100644
--- a/src/com/android/launcher3/model/BgDataModel.java
+++ b/src/com/android/launcher3/model/BgDataModel.java
@@ -49,7 +49,6 @@
 import com.android.launcher3.model.data.AppInfo;
 import com.android.launcher3.model.data.CollectionInfo;
 import com.android.launcher3.model.data.ItemInfo;
-import com.android.launcher3.model.data.LauncherAppWidgetInfo;
 import com.android.launcher3.model.data.WorkspaceItemInfo;
 import com.android.launcher3.pm.UserCache;
 import com.android.launcher3.shortcuts.ShortcutKey;
@@ -70,7 +69,6 @@
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashMap;
-import java.util.HashSet;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
@@ -419,9 +417,9 @@
          * Binds updated incremental download progress
          */
         default void bindIncrementalDownloadProgressUpdated(AppInfo app) { }
-        default void bindWorkspaceItemsChanged(List<WorkspaceItemInfo> updated) { }
-        default void bindWidgetsRestored(ArrayList<LauncherAppWidgetInfo> widgets) { }
-        default void bindRestoreItemsChange(HashSet<ItemInfo> updates) { }
+
+        /** Called when a runtime property of the ItemInfo is updated due to some system event */
+        default void bindItemsUpdated(Set<ItemInfo> updates) { }
         default void bindWorkspaceComponentsRemoved(Predicate<ItemInfo> matcher) { }
 
         /**
diff --git a/src/com/android/launcher3/model/CacheDataUpdatedTask.java b/src/com/android/launcher3/model/CacheDataUpdatedTask.java
index b544b91..48934e2 100644
--- a/src/com/android/launcher3/model/CacheDataUpdatedTask.java
+++ b/src/com/android/launcher3/model/CacheDataUpdatedTask.java
@@ -15,6 +15,9 @@
  */
 package com.android.launcher3.model;
 
+import static com.android.launcher3.icons.cache.CacheLookupFlag.DEFAULT_LOOKUP_FLAG;
+import static com.android.launcher3.model.ModelUtils.WIDGET_FILTER;
+
 import android.content.ComponentName;
 import android.os.UserHandle;
 
@@ -23,6 +26,8 @@
 import com.android.launcher3.LauncherModel.ModelUpdateTask;
 import com.android.launcher3.LauncherSettings;
 import com.android.launcher3.icons.IconCache;
+import com.android.launcher3.model.data.ItemInfo;
+import com.android.launcher3.model.data.LauncherAppWidgetInfo;
 import com.android.launcher3.model.data.WorkspaceItemInfo;
 
 import java.util.ArrayList;
@@ -55,7 +60,7 @@
     public void execute(@NonNull ModelTaskController taskController, @NonNull BgDataModel dataModel,
             @NonNull AllAppsList apps) {
         IconCache iconCache = taskController.getApp().getIconCache();
-        ArrayList<WorkspaceItemInfo> updatedShortcuts = new ArrayList<>();
+        ArrayList<ItemInfo> updatedItems = new ArrayList<>();
 
         synchronized (dataModel) {
             dataModel.forAllWorkspaceItemInfos(mUser, si -> {
@@ -64,12 +69,25 @@
                         && isValidShortcut(si) && cn != null
                         && mPackages.contains(cn.getPackageName())) {
                     iconCache.getTitleAndIcon(si, si.getMatchingLookupFlag());
-                    updatedShortcuts.add(si);
+                    updatedItems.add(si);
                 }
             });
+
+            dataModel.itemsIdMap.stream()
+                    .filter(WIDGET_FILTER)
+                    .filter(item -> mUser.equals(item.user))
+                    .map(item -> (LauncherAppWidgetInfo) item)
+                    .filter(widget -> mPackages.contains(widget.providerName.getPackageName())
+                            && widget.pendingItemInfo != null)
+                    .forEach(widget -> {
+                        iconCache.getTitleAndIconForApp(
+                                widget.pendingItemInfo, DEFAULT_LOOKUP_FLAG);
+                        updatedItems.add(widget);
+                    });
+
             apps.updateIconsAndLabels(mPackages, mUser);
         }
-        taskController.bindUpdatedWorkspaceItems(updatedShortcuts);
+        taskController.bindUpdatedWorkspaceItems(updatedItems);
         taskController.bindApplicationsIfNeeded();
     }
 
diff --git a/src/com/android/launcher3/model/ModelTaskController.kt b/src/com/android/launcher3/model/ModelTaskController.kt
index fc53343..40ea17d 100644
--- a/src/com/android/launcher3/model/ModelTaskController.kt
+++ b/src/com/android/launcher3/model/ModelTaskController.kt
@@ -22,7 +22,6 @@
 import com.android.launcher3.celllayout.CellPosMapper
 import com.android.launcher3.model.BgDataModel.FixedContainerItems
 import com.android.launcher3.model.data.ItemInfo
-import com.android.launcher3.model.data.WorkspaceItemInfo
 import com.android.launcher3.util.PackageUserKey
 import com.android.launcher3.widget.model.WidgetsListBaseEntriesBuilder
 import java.util.Objects
@@ -51,18 +50,17 @@
      */
     fun getModelWriter() = model.getWriter(false /* verifyChanges */, CellPosMapper.DEFAULT, null)
 
-    fun bindUpdatedWorkspaceItems(allUpdates: List<WorkspaceItemInfo>) {
+    fun bindUpdatedWorkspaceItems(allUpdates: Collection<ItemInfo>) {
         // Bind workspace items
-        val workspaceUpdates =
-            allUpdates.stream().filter { info -> info.id != ItemInfo.NO_ID }.toList()
+        val workspaceUpdates = allUpdates.filter { it.id != ItemInfo.NO_ID }.toSet()
         if (workspaceUpdates.isNotEmpty()) {
-            scheduleCallbackTask { it.bindWorkspaceItemsChanged(workspaceUpdates) }
+            scheduleCallbackTask { it.bindItemsUpdated(workspaceUpdates) }
         }
 
         // Bind extra items if any
         allUpdates
             .stream()
-            .mapToInt { info: WorkspaceItemInfo -> info.container }
+            .mapToInt { it.container }
             .distinct()
             .mapToObj { dataModel.extraItems.get(it) }
             .filter { Objects.nonNull(it) }
diff --git a/src/com/android/launcher3/model/PackageInstallStateChangedTask.java b/src/com/android/launcher3/model/PackageInstallStateChangedTask.java
index 4103937..a216042 100644
--- a/src/com/android/launcher3/model/PackageInstallStateChangedTask.java
+++ b/src/com/android/launcher3/model/PackageInstallStateChangedTask.java
@@ -99,8 +99,7 @@
                     });
 
             if (!updates.isEmpty()) {
-                taskController.scheduleCallbackTask(
-                        callbacks -> callbacks.bindRestoreItemsChange(updates));
+                taskController.bindUpdatedWorkspaceItems(updates);
             }
         }
     }
diff --git a/src/com/android/launcher3/model/PackageUpdatedTask.java b/src/com/android/launcher3/model/PackageUpdatedTask.java
index 1153f48..6bef292 100644
--- a/src/com/android/launcher3/model/PackageUpdatedTask.java
+++ b/src/com/android/launcher3/model/PackageUpdatedTask.java
@@ -214,8 +214,7 @@
 
         // Update shortcut infos
         if (mOp == OP_ADD || flagOp != FlagOp.NO_OP) {
-            final ArrayList<WorkspaceItemInfo> updatedWorkspaceItems = new ArrayList<>();
-            final ArrayList<LauncherAppWidgetInfo> widgets = new ArrayList<>();
+            final ArrayList<ItemInfo> updatedWorkspaceItems = new ArrayList<>();
 
             // For system apps, package manager send OP_UPDATE when an app is enabled.
             final boolean isNewApkAvailable = mOp == OP_ADD || mOp == OP_UPDATE;
@@ -364,8 +363,8 @@
                             // if the widget has a config activity. In case there is no config
                             // activity, it will be marked as 'restored' during bind.
                             widgetInfo.restoreStatus |= LauncherAppWidgetInfo.FLAG_UI_NOT_READY;
-
-                            widgets.add(widgetInfo);
+                            widgetInfo.installProgress = 100;
+                            updatedWorkspaceItems.add(widgetInfo);
                             taskController.getModelWriter().updateItemInDatabase(widgetInfo);
                         });
             }
@@ -377,10 +376,6 @@
                         "removing shortcuts with invalid target components."
                                 + " ids=" + removedShortcuts);
             }
-
-            if (!widgets.isEmpty()) {
-                taskController.scheduleCallbackTask(c -> c.bindWidgetsRestored(widgets));
-            }
         }
 
         final HashSet<String> removedPackages = new HashSet<>();
diff --git a/src/com/android/launcher3/model/ShortcutsChangedTask.java b/src/com/android/launcher3/model/ShortcutsChangedTask.java
deleted file mode 100644
index b5a7382..0000000
--- a/src/com/android/launcher3/model/ShortcutsChangedTask.java
+++ /dev/null
@@ -1,136 +0,0 @@
-/*
- * Copyright (C) 2016 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.launcher3.model;
-
-import android.content.Context;
-import android.content.pm.ShortcutInfo;
-import android.os.UserHandle;
-
-import androidx.annotation.NonNull;
-
-import com.android.launcher3.LauncherAppState;
-import com.android.launcher3.LauncherModel.ModelUpdateTask;
-import com.android.launcher3.LauncherSettings;
-import com.android.launcher3.icons.CacheableShortcutInfo;
-import com.android.launcher3.model.data.WorkspaceItemInfo;
-import com.android.launcher3.shortcuts.ShortcutKey;
-import com.android.launcher3.shortcuts.ShortcutRequest;
-import com.android.launcher3.util.ApplicationInfoWrapper;
-import com.android.launcher3.util.ItemInfoMatcher;
-
-import java.util.ArrayList;
-import java.util.HashSet;
-import java.util.List;
-import java.util.Set;
-import java.util.stream.Collectors;
-
-/**
- * Handles changes due to shortcut manager updates (deep shortcut changes)
- */
-public class ShortcutsChangedTask implements ModelUpdateTask {
-
-    @NonNull
-    private final String mPackageName;
-
-    @NonNull
-    private final List<ShortcutInfo> mShortcuts;
-
-    @NonNull
-    private final UserHandle mUser;
-
-    private final boolean mUpdateIdMap;
-
-    public ShortcutsChangedTask(@NonNull final String packageName,
-            @NonNull final List<ShortcutInfo> shortcuts, @NonNull final UserHandle user,
-            final boolean updateIdMap) {
-        mPackageName = packageName;
-        mShortcuts = shortcuts;
-        mUser = user;
-        mUpdateIdMap = updateIdMap;
-    }
-
-    @Override
-    public void execute(@NonNull ModelTaskController taskController, @NonNull BgDataModel dataModel,
-            @NonNull AllAppsList apps) {
-        final LauncherAppState app = taskController.getApp();
-        final Context context = app.getContext();
-        // Find WorkspaceItemInfo's that have changed on the workspace.
-        ArrayList<WorkspaceItemInfo> matchingWorkspaceItems = new ArrayList<>();
-
-        synchronized (dataModel) {
-            dataModel.forAllWorkspaceItemInfos(mUser, si -> {
-                if ((si.itemType == LauncherSettings.Favorites.ITEM_TYPE_DEEP_SHORTCUT)
-                        && mPackageName.equals(si.getIntent().getPackage())) {
-                    matchingWorkspaceItems.add(si);
-                }
-            });
-        }
-
-        if (!matchingWorkspaceItems.isEmpty()) {
-            ApplicationInfoWrapper infoWrapper =
-                    new ApplicationInfoWrapper(context, mPackageName, mUser);
-            if (mShortcuts.isEmpty()) {
-                // Verify that the app is indeed installed.
-                if (!infoWrapper.isInstalled() && !infoWrapper.isArchived()) {
-                    // App is not installed or archived, ignoring package events
-                    return;
-                }
-            }
-            // Update the workspace to reflect the changes to updated shortcuts residing on it.
-            List<String> allLauncherKnownIds = matchingWorkspaceItems.stream()
-                    .map(WorkspaceItemInfo::getDeepShortcutId)
-                    .distinct()
-                    .collect(Collectors.toList());
-            List<ShortcutInfo> shortcuts = new ShortcutRequest(context, mUser)
-                    .forPackage(mPackageName, allLauncherKnownIds)
-                    .query(ShortcutRequest.ALL);
-
-            Set<String> nonPinnedIds = new HashSet<>(allLauncherKnownIds);
-            ArrayList<WorkspaceItemInfo> updatedWorkspaceItemInfos = new ArrayList<>();
-            for (ShortcutInfo fullDetails : shortcuts) {
-                if (!fullDetails.isPinned()) {
-                    continue;
-                }
-                String sid = fullDetails.getId();
-                nonPinnedIds.remove(sid);
-                matchingWorkspaceItems
-                        .stream()
-                        .filter(itemInfo -> sid.equals(itemInfo.getDeepShortcutId()))
-                        .forEach(workspaceItemInfo -> {
-                            workspaceItemInfo.updateFromDeepShortcutInfo(fullDetails, context);
-                            app.getIconCache().getShortcutIcon(workspaceItemInfo,
-                                    new CacheableShortcutInfo(fullDetails, infoWrapper));
-                            updatedWorkspaceItemInfos.add(workspaceItemInfo);
-                        });
-            }
-
-            taskController.bindUpdatedWorkspaceItems(updatedWorkspaceItemInfos);
-            if (!nonPinnedIds.isEmpty()) {
-                taskController.deleteAndBindComponentsRemoved(ItemInfoMatcher.ofShortcutKeys(
-                        nonPinnedIds.stream()
-                                .map(id -> new ShortcutKey(mPackageName, mUser, id))
-                                .collect(Collectors.toSet())),
-                        "removed because the shortcut is no longer available in shortcut service");
-            }
-        }
-
-        if (mUpdateIdMap) {
-            // Update the deep shortcut map if the list of ids has changed for an activity.
-            dataModel.updateDeepShortcutCounts(mPackageName, mUser, mShortcuts);
-            taskController.bindDeepShortcuts(dataModel);
-        }
-    }
-}
diff --git a/src/com/android/launcher3/model/ShortcutsChangedTask.kt b/src/com/android/launcher3/model/ShortcutsChangedTask.kt
new file mode 100644
index 0000000..2e4f75f
--- /dev/null
+++ b/src/com/android/launcher3/model/ShortcutsChangedTask.kt
@@ -0,0 +1,114 @@
+/*
+ * Copyright (C) 2025 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.android.launcher3.model
+
+import android.content.pm.ShortcutInfo
+import android.os.UserHandle
+import com.android.launcher3.LauncherModel.ModelUpdateTask
+import com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_DEEP_SHORTCUT
+import com.android.launcher3.icons.CacheableShortcutInfo
+import com.android.launcher3.model.data.WorkspaceItemInfo
+import com.android.launcher3.shortcuts.ShortcutKey
+import com.android.launcher3.shortcuts.ShortcutRequest
+import com.android.launcher3.util.ApplicationInfoWrapper
+import com.android.launcher3.util.ItemInfoMatcher
+
+/** Handles changes due to shortcut manager updates (deep shortcut changes) */
+class ShortcutsChangedTask(
+    private val packageName: String,
+    private val shortcuts: List<ShortcutInfo>,
+    private val user: UserHandle,
+    private val shouldUpdateIdMap: Boolean,
+) : ModelUpdateTask {
+
+    override fun execute(
+        taskController: ModelTaskController,
+        dataModel: BgDataModel,
+        apps: AllAppsList,
+    ) {
+        val app = taskController.app
+        val context = app.context
+        // Find WorkspaceItemInfo's that have changed on the workspace.
+        val matchingWorkspaceItems = ArrayList<WorkspaceItemInfo>()
+
+        synchronized(dataModel) {
+            dataModel.forAllWorkspaceItemInfos(user) { wai: WorkspaceItemInfo ->
+                if (
+                    (wai.itemType == ITEM_TYPE_DEEP_SHORTCUT) &&
+                        packageName == wai.getIntent().getPackage()
+                ) {
+                    matchingWorkspaceItems.add(wai)
+                }
+            }
+        }
+
+        if (matchingWorkspaceItems.isNotEmpty()) {
+            val infoWrapper = ApplicationInfoWrapper(context, packageName, user)
+            if (shortcuts.isEmpty()) {
+                // Verify that the app is indeed installed.
+                if (!infoWrapper.isInstalled() && !infoWrapper.isArchived()) {
+                    // App is not installed or archived, ignoring package events
+                    return
+                }
+            }
+            // Update the workspace to reflect the changes to updated shortcuts residing on it.
+            val allLauncherKnownIds =
+                matchingWorkspaceItems.map { item -> item.deepShortcutId }.distinct()
+            val shortcuts: List<ShortcutInfo> =
+                ShortcutRequest(context, user)
+                    .forPackage(packageName, allLauncherKnownIds)
+                    .query(ShortcutRequest.ALL)
+
+            val nonPinnedIds: MutableSet<String> = HashSet(allLauncherKnownIds)
+            val updatedWorkspaceItemInfos = ArrayList<WorkspaceItemInfo>()
+            for (fullDetails in shortcuts) {
+                if (!fullDetails.isPinned) {
+                    continue
+                }
+                val shortcutId = fullDetails.id
+                nonPinnedIds.remove(shortcutId)
+                matchingWorkspaceItems
+                    .filter { itemInfo: WorkspaceItemInfo -> shortcutId == itemInfo.deepShortcutId }
+                    .forEach { workspaceItemInfo: WorkspaceItemInfo ->
+                        workspaceItemInfo.updateFromDeepShortcutInfo(fullDetails, context)
+                        app.iconCache.getShortcutIcon(
+                            workspaceItemInfo,
+                            CacheableShortcutInfo(fullDetails, infoWrapper),
+                        )
+                        updatedWorkspaceItemInfos.add(workspaceItemInfo)
+                    }
+            }
+
+            taskController.bindUpdatedWorkspaceItems(updatedWorkspaceItemInfos)
+            if (nonPinnedIds.isNotEmpty()) {
+                taskController.deleteAndBindComponentsRemoved(
+                    ItemInfoMatcher.ofShortcutKeys(
+                        nonPinnedIds
+                            .map { id: String? -> ShortcutKey(packageName, user, id) }
+                            .toSet()
+                    ),
+                    "removed because the shortcut is no longer available in shortcut service",
+                )
+            }
+        }
+
+        if (shouldUpdateIdMap) {
+            // Update the deep shortcut map if the list of ids has changed for an activity.
+            dataModel.updateDeepShortcutCounts(packageName, user, shortcuts)
+            taskController.bindDeepShortcuts(dataModel)
+        }
+    }
+}
diff --git a/src/com/android/launcher3/shapes/IconShapeModel.kt b/src/com/android/launcher3/shapes/IconShapeModel.kt
index 10b7f91..dd6c432 100644
--- a/src/com/android/launcher3/shapes/IconShapeModel.kt
+++ b/src/com/android/launcher3/shapes/IconShapeModel.kt
@@ -16,4 +16,9 @@
 
 package com.android.launcher3.shapes
 
-data class IconShapeModel(val key: String, val title: String, val pathString: String)
+data class IconShapeModel(
+    val key: String,
+    val title: String,
+    val pathString: String,
+    val folderPathString: String = pathString,
+)
diff --git a/src/com/android/launcher3/shapes/IconShapesProvider.kt b/src/com/android/launcher3/shapes/IconShapesProvider.kt
index a190e22..8608437 100644
--- a/src/com/android/launcher3/shapes/IconShapesProvider.kt
+++ b/src/com/android/launcher3/shapes/IconShapesProvider.kt
@@ -28,7 +28,7 @@
                         key = "arch",
                         title = "arch",
                         pathString =
-                            "M50 0C77.614 0 100 22.386 100 50C100 85.471 100 86.476 99.9 87.321 99.116 93.916 93.916 99.116 87.321 99.9 86.476 100 85.471 100 83.46 100H16.54C14.529 100 13.524 100 12.679 99.9 6.084 99.116.884 93.916.1 87.321 0 86.476 0 85.471 0 83.46L0 50C0 22.386 22.386 0 50 0Z",
+                            "M50 0C77.614 0 100 22.386 100 50C100 85.471 100 86.476 99.9 87.321 99.116 93.916 93.916 99.116 87.321 99.9 86.476 100 85.471 100 83.46 100H16.54C14.529 100 13.524 100 12.679 99.9 6.084 99.116 .884 93.916 .1 87.321 0 86.476 0 85.471 0 83.46L0 50C0 22.386 22.386 0 50 0Z",
                     ),
                 "4_sided_cookie" to
                     IconShapeModel(
@@ -42,7 +42,7 @@
                         key = "seven_sided_cookie",
                         title = "7 sided cookie",
                         pathString =
-                            "M35.209 4.878C36.326 3.895 36.884 3.404 37.397 3.006 44.82 -2.742 55.18 -2.742 62.603 3.006 63.116 3.404 63.674 3.895 64.791 4.878 65.164 5.207 65.351 5.371 65.539 5.529 68.167 7.734 71.303 9.248 74.663 9.932 74.902 9.981 75.147 10.025 75.637 10.113 77.1 10.375 77.831 10.506 78.461 10.66 87.573 12.893 94.032 21.011 94.176 30.412 94.186 31.062 94.151 31.805 94.08 33.293 94.057 33.791 94.045 34.04 94.039 34.285 93.958 37.72 94.732 41.121 96.293 44.18 96.404 44.399 96.522 44.618 96.759 45.056 97.467 46.366 97.821 47.021 98.093 47.611 102.032 56.143 99.727 66.266 92.484 72.24 91.983 72.653 91.381 73.089 90.177 73.961 89.774 74.254 89.572 74.4 89.377 74.548 86.647 76.626 84.477 79.353 83.063 82.483 82.962 82.707 82.865 82.936 82.671 83.395 82.091 84.766 81.8 85.451 81.51 86.033 77.31 94.44 67.977 98.945 58.801 96.994 58.166 96.859 57.451 96.659 56.019 96.259 55.54 96.125 55.3 96.058 55.063 95.998 51.74 95.154 48.26 95.154 44.937 95.998 44.699 96.058 44.46 96.125 43.981 96.259 42.549 96.659 41.834 96.859 41.199 96.994 32.023 98.945 22.69 94.44 18.49 86.033 18.2 85.451 17.909 84.766 17.329 83.395 17.135 82.936 17.038 82.707 16.937 82.483 15.523 79.353 13.353 76.626 10.623 74.548 10.428 74.4 10.226 74.254 9.823 73.961 8.619 73.089 8.017 72.653 7.516 72.24.273 66.266 -2.032 56.143 1.907 47.611 2.179 47.021 2.533 46.366 3.241 45.056 3.478 44.618 3.596 44.399 3.707 44.18 5.268 41.121 6.042 37.72 5.961 34.285 5.955 34.04 5.943 33.791 5.92 33.293 5.849 31.805 5.814 31.062 5.824 30.412 5.968 21.011 12.427 12.893 21.539 10.66 22.169 10.506 22.9 10.375 24.363 10.113 24.853 10.025 25.098 9.981 25.337 9.932 28.697 9.248 31.833 7.734 34.461 5.529 34.649 5.371 34.836 5.207 35.209 4.878Z",
+                            "M35.209 4.878C36.326 3.895 36.884 3.404 37.397 3.006 44.82 -2.742 55.18 -2.742 62.603 3.006 63.116 3.404 63.674 3.895 64.791 4.878 65.164 5.207 65.351 5.371 65.539 5.529 68.167 7.734 71.303 9.248 74.663 9.932 74.902 9.981 75.147 10.025 75.637 10.113 77.1 10.375 77.831 10.506 78.461 10.66 87.573 12.893 94.032 21.011 94.176 30.412 94.186 31.062 94.151 31.805 94.08 33.293 94.057 33.791 94.045 34.04 94.039 34.285 93.958 37.72 94.732 41.121 96.293 44.18 96.404 44.399 96.522 44.618 96.759 45.056 97.467 46.366 97.821 47.021 98.093 47.611 102.032 56.143 99.727 66.266 92.484 72.24 91.983 72.653 91.381 73.089 90.177 73.961 89.774 74.254 89.572 74.4 89.377 74.548 86.647 76.626 84.477 79.353 83.063 82.483 82.962 82.707 82.865 82.936 82.671 83.395 82.091 84.766 81.8 85.451 81.51 86.033 77.31 94.44 67.977 98.945 58.801 96.994 58.166 96.859 57.451 96.659 56.019 96.259 55.54 96.125 55.3 96.058 55.063 95.998 51.74 95.154 48.26 95.154 44.937 95.998 44.699 96.058 44.46 96.125 43.981 96.259 42.549 96.659 41.834 96.859 41.199 96.994 32.023 98.945 22.69 94.44 18.49 86.033 18.2 85.451 17.909 84.766 17.329 83.395 17.135 82.936 17.038 82.707 16.937 82.483 15.523 79.353 13.353 76.626 10.623 74.548 10.428 74.4 10.226 74.254 9.823 73.961 8.619 73.089 8.017 72.653 7.516 72.24 .273 66.266 -2.032 56.143 1.907 47.611 2.179 47.021 2.533 46.366 3.241 45.056 3.478 44.618 3.596 44.399 3.707 44.18 5.268 41.121 6.042 37.72 5.961 34.285 5.955 34.04 5.943 33.791 5.92 33.293 5.849 31.805 5.814 31.062 5.824 30.412 5.968 21.011 12.427 12.893 21.539 10.66 22.169 10.506 22.9 10.375 24.363 10.113 24.853 10.025 25.098 9.981 25.337 9.932 28.697 9.248 31.833 7.734 34.461 5.529 34.649 5.371 34.836 5.207 35.209 4.878Z",
                     ),
                 "sunny" to
                     IconShapeModel(
@@ -62,7 +62,7 @@
                         key = "square",
                         title = "square",
                         pathString =
-                            "M53.689 0.82L53.689.82C67.434.82 74.306.82 79.758 2.978 87.649 6.103 93.897 12.351 97.022 20.242 99.18 25.694 99.18 32.566 99.18 46.311V53.689C99.18 67.434 99.18 74.306 97.022 79.758 93.897 87.649 87.649 93.897 79.758 97.022 74.306 99.18 67.434 99.18 53.689 99.18H46.311C32.566 99.18 25.694 99.18 20.242 97.022 12.351 93.897 6.103 87.649 2.978 79.758.82 74.306.82 67.434.82 53.689L.82 46.311C.82 32.566.82 25.694 2.978 20.242 6.103 12.351 12.351 6.103 20.242 2.978 25.694.82 32.566.82 46.311.82Z",
+                            "M53.689 0.82 L53.689 .82 C67.434 .82 74.306 .82 79.758 2.978 87.649 6.103 93.897 12.351 97.022 20.242 99.18 25.694 99.18 32.566 99.18 46.311 V53.689 C99.18 67.434 99.18 74.306 97.022 79.758 93.897 87.649 87.649 93.897 79.758 97.022 74.306 99.18 67.434 99.18 53.689 99.18 H46.311 C32.566 99.18 25.694 99.18 20.242 97.022 12.351 93.897 6.103 87.649 2.978 79.758 .82 74.306 .82 67.434 .82 53.689 L.82 46.311 C.82 32.566 .82 25.694 2.978 20.242 6.103 12.351 12.351 6.103 20.242 2.978 25.694 .82 32.566 .82 46.311 .82Z",
                     ),
             )
         } else {
diff --git a/src/com/android/launcher3/touch/ItemClickHandler.java b/src/com/android/launcher3/touch/ItemClickHandler.java
index 78709b8..381d17a 100644
--- a/src/com/android/launcher3/touch/ItemClickHandler.java
+++ b/src/com/android/launcher3/touch/ItemClickHandler.java
@@ -228,10 +228,9 @@
     private static void onClickPendingAppItem(View v, Launcher launcher, String packageName,
             boolean downloadStarted) {
         ItemInfo item = (ItemInfo) v.getTag();
-        CompletableFuture<SessionInfo> siFuture;
-        siFuture = CompletableFuture.supplyAsync(() ->
-                        InstallSessionHelper.INSTANCE.get(launcher)
-                                .getActiveSessionInfo(item.user, packageName),
+        CompletableFuture<SessionInfo> siFuture = CompletableFuture.supplyAsync(() ->
+                InstallSessionHelper.INSTANCE.get(launcher)
+                        .getActiveSessionInfo(item.user, packageName),
                 UI_HELPER_EXECUTOR);
         Consumer<SessionInfo> marketLaunchAction = sessionInfo -> {
             if (sessionInfo != null) {
@@ -245,8 +244,8 @@
                 }
             }
             // Fallback to using custom market intent.
-            Intent intent = ApiWrapper.INSTANCE.get(launcher).getAppMarketActivityIntent(
-                    packageName, Process.myUserHandle());
+            Intent intent = ApiWrapper.INSTANCE.get(launcher).getMarketSearchIntent(
+                    packageName, item.user);
             launcher.startActivitySafely(v, intent, item);
         };
 
@@ -358,9 +357,7 @@
         // Check for abandoned promise
         if ((v instanceof BubbleTextView) && shortcut.hasPromiseIconUi()
                 && (!Flags.enableSupportForArchiving() || !shortcut.isArchived())) {
-            String packageName = shortcut.getIntent().getComponent() != null
-                    ? shortcut.getIntent().getComponent().getPackageName()
-                    : shortcut.getIntent().getPackage();
+            String packageName = shortcut.getTargetPackage();
             if (!TextUtils.isEmpty(packageName)) {
                 onClickPendingAppItem(
                         v,
diff --git a/src/com/android/launcher3/util/ApiWrapper.java b/src/com/android/launcher3/util/ApiWrapper.java
index 467a7ec..48e033a 100644
--- a/src/com/android/launcher3/util/ApiWrapper.java
+++ b/src/com/android/launcher3/util/ApiWrapper.java
@@ -28,6 +28,7 @@
 import android.content.pm.ShortcutInfo;
 import android.graphics.drawable.ColorDrawable;
 import android.net.Uri;
+import android.os.Process;
 import android.os.UserHandle;
 import android.os.UserManager;
 import android.util.ArrayMap;
@@ -120,6 +121,21 @@
      * Activity).
      */
     public Intent getAppMarketActivityIntent(String packageName, UserHandle user) {
+        return createMarketIntent(packageName);
+    }
+
+    /**
+     * Returns an intent which can be used to start a search for a package on app market
+     */
+    public Intent getMarketSearchIntent(String packageName, UserHandle user) {
+        // If we are search for the current user, just launch the market directly as the
+        // system won't have the installer details either
+        return  (Process.myUserHandle().equals(user))
+                ? createMarketIntent(packageName)
+                : getAppMarketActivityIntent(packageName, user);
+    }
+
+    private static Intent createMarketIntent(String packageName) {
         return new Intent(Intent.ACTION_VIEW)
                 .setData(new Uri.Builder()
                         .scheme("market")
@@ -171,6 +187,13 @@
         return appInfo.sourceDir;
     }
 
+    /**
+     * Returns the round icon resource Id if defined by the app
+     */
+    public int getRoundIconRes(@NonNull ApplicationInfo appInfo) {
+        return 0;
+    }
+
     private static class NoopDrawable extends ColorDrawable {
         @Override
         public int getIntrinsicHeight() {
diff --git a/src/com/android/launcher3/util/DisplayController.java b/src/com/android/launcher3/util/DisplayController.java
index 9472f5f..475dc04 100644
--- a/src/com/android/launcher3/util/DisplayController.java
+++ b/src/com/android/launcher3/util/DisplayController.java
@@ -69,6 +69,7 @@
 import java.util.Objects;
 import java.util.Set;
 import java.util.StringJoiner;
+import java.util.concurrent.CopyOnWriteArrayList;
 
 import javax.inject.Inject;
 
@@ -114,7 +115,8 @@
 
     // The callback in this listener updates DeviceProfile, which other listeners might depend on
     private DisplayInfoChangeListener mPriorityListener;
-    private final ArrayList<DisplayInfoChangeListener> mListeners = new ArrayList<>();
+    private final CopyOnWriteArrayList<DisplayInfoChangeListener> mListeners =
+            new CopyOnWriteArrayList<>();
 
     // We will register broadcast receiver on main thread to ensure not missing changes on
     // TARGET_OVERLAY_PACKAGE and ACTION_OVERLAY_CHANGED.
diff --git a/src/com/android/launcher3/util/LauncherBindableItemsContainer.java b/src/com/android/launcher3/util/LauncherBindableItemsContainer.java
index 02779ce..20e3eaf 100644
--- a/src/com/android/launcher3/util/LauncherBindableItemsContainer.java
+++ b/src/com/android/launcher3/util/LauncherBindableItemsContainer.java
@@ -15,24 +15,20 @@
  */
 package com.android.launcher3.util;
 
-import android.graphics.drawable.Drawable;
 import android.view.View;
 
 import com.android.launcher3.BubbleTextView;
 import com.android.launcher3.apppairs.AppPairIcon;
 import com.android.launcher3.folder.Folder;
 import com.android.launcher3.folder.FolderIcon;
-import com.android.launcher3.graphics.PreloadIconDrawable;
 import com.android.launcher3.model.data.AppPairInfo;
 import com.android.launcher3.model.data.FolderInfo;
 import com.android.launcher3.model.data.ItemInfo;
-import com.android.launcher3.model.data.LauncherAppWidgetInfo;
 import com.android.launcher3.model.data.WorkspaceItemInfo;
 import com.android.launcher3.views.ActivityContext;
 import com.android.launcher3.widget.PendingAppWidgetHostView;
 
-import java.util.HashSet;
-import java.util.List;
+import java.util.Set;
 
 /**
  * Interface representing a container which can bind Launcher items with some utility methods
@@ -41,27 +37,22 @@
 
     /**
      * Called to update workspace items as a result of
-     * {@link com.android.launcher3.model.BgDataModel.Callbacks#bindWorkspaceItemsChanged(List)}
+     * {@link com.android.launcher3.model.BgDataModel.Callbacks#bindItemsUpdated(Set)}
      */
-    default void updateWorkspaceItems(List<WorkspaceItemInfo> shortcuts, ActivityContext context) {
-        final HashSet<WorkspaceItemInfo> updates = new HashSet<>(shortcuts);
+    default void updateContainerItems(Set<ItemInfo> updates, ActivityContext context) {
         ItemOperator op = (info, v) -> {
-            if (v instanceof BubbleTextView && updates.contains(info)) {
-                WorkspaceItemInfo si = (WorkspaceItemInfo) info;
-                BubbleTextView shortcut = (BubbleTextView) v;
-                Drawable oldIcon = shortcut.getIcon();
-                boolean oldPromiseState = (oldIcon instanceof PreloadIconDrawable)
-                        && ((PreloadIconDrawable) oldIcon).hasNotCompleted();
-                shortcut.applyFromWorkspaceItem(
-                        si,
-                        si.isPromise() != oldPromiseState
-                                && oldIcon instanceof PreloadIconDrawable
-                                ? (PreloadIconDrawable) oldIcon
-                                : null);
-            } else if (info instanceof FolderInfo && v instanceof FolderIcon) {
-                ((FolderIcon) v).updatePreviewItems(updates::contains);
+            if (v instanceof BubbleTextView shortcut
+                    && info instanceof WorkspaceItemInfo wii
+                    && updates.contains(info)) {
+                shortcut.applyFromWorkspaceItem(wii);
+            } else if (info instanceof FolderInfo && v instanceof FolderIcon folderIcon) {
+                folderIcon.updatePreviewItems(updates::contains);
             } else if (info instanceof AppPairInfo && v instanceof AppPairIcon appPairIcon) {
                 appPairIcon.maybeRedrawForWorkspaceUpdate(updates::contains);
+            } else if (v instanceof PendingAppWidgetHostView pendingView
+                    && updates.contains(info)) {
+                pendingView.applyState();
+                pendingView.postProviderAvailabilityCheck();
             }
 
             // Iterate all items
@@ -76,35 +67,6 @@
     }
 
     /**
-     * Called to update restored items as a result of
-     * {@link com.android.launcher3.model.BgDataModel.Callbacks#bindRestoreItemsChange(HashSet)}}
-     */
-    default void updateRestoreItems(final HashSet<ItemInfo> updates, ActivityContext context) {
-        ItemOperator op = (info, v) -> {
-            if (info instanceof WorkspaceItemInfo && v instanceof BubbleTextView
-                    && updates.contains(info)) {
-                ((BubbleTextView) v).applyLoadingState(null);
-            } else if (v instanceof PendingAppWidgetHostView
-                    && info instanceof LauncherAppWidgetInfo
-                    && updates.contains(info)) {
-                ((PendingAppWidgetHostView) v).applyState();
-            } else if (v instanceof FolderIcon && info instanceof FolderInfo) {
-                ((FolderIcon) v).updatePreviewItems(updates::contains);
-            } else if (info instanceof AppPairInfo && v instanceof AppPairIcon appPairIcon) {
-                appPairIcon.maybeRedrawForWorkspaceUpdate(updates::contains);
-            }
-            // process all the shortcuts
-            return false;
-        };
-
-        mapOverItems(op);
-        Folder folder = Folder.getOpen(context);
-        if (folder != null) {
-            folder.iterateOverItems(op);
-        }
-    }
-
-    /**
      * Map the operator over the shortcuts and widgets.
      *
      * @param op the operator to map over the shortcuts
diff --git a/src/com/android/launcher3/util/MSDLPlayerWrapper.java b/src/com/android/launcher3/util/MSDLPlayerWrapper.java
index 8a1d923..fc3fa72 100644
--- a/src/com/android/launcher3/util/MSDLPlayerWrapper.java
+++ b/src/com/android/launcher3/util/MSDLPlayerWrapper.java
@@ -16,8 +16,6 @@
 
 package com.android.launcher3.util;
 
-import static com.android.launcher3.util.Executors.UI_HELPER_EXECUTOR;
-
 import android.content.Context;
 import android.os.Vibrator;
 
@@ -50,7 +48,9 @@
     @Inject
     public MSDLPlayerWrapper(@ApplicationContext Context context) {
         Vibrator vibrator = context.getSystemService(Vibrator.class);
-        mMSDLPlayer = MSDLPlayer.Companion.createPlayer(vibrator, UI_HELPER_EXECUTOR, null);
+        mMSDLPlayer = MSDLPlayer.Companion.createPlayer(vibrator,
+                java.util.concurrent.Executors.newSingleThreadExecutor(),
+                null /* useHapticFeedbackForToken */);
     }
 
     /** Perform MSDL feedback for a token with interaction properties */
diff --git a/src/com/android/launcher3/widget/PendingAppWidgetHostView.java b/src/com/android/launcher3/widget/PendingAppWidgetHostView.java
index 9c9b80d..cd8e457 100644
--- a/src/com/android/launcher3/widget/PendingAppWidgetHostView.java
+++ b/src/com/android/launcher3/widget/PendingAppWidgetHostView.java
@@ -21,7 +21,7 @@
 import static android.graphics.Paint.FILTER_BITMAP_FLAG;
 
 import static com.android.launcher3.graphics.PreloadIconDrawable.newPendingIcon;
-import static com.android.launcher3.icons.FastBitmapDrawable.getDisabledColorFilter;
+import static com.android.launcher3.model.data.LauncherAppWidgetInfo.FLAG_PROVIDER_NOT_READY;
 import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
 
 import android.appwidget.AppWidgetProviderInfo;
@@ -37,6 +37,9 @@
 import android.graphics.drawable.ColorDrawable;
 import android.graphics.drawable.Drawable;
 import android.os.Bundle;
+import android.os.Handler;
+import android.os.Looper;
+import android.os.Message;
 import android.text.Layout;
 import android.text.StaticLayout;
 import android.text.TextPaint;
@@ -60,8 +63,10 @@
 import com.android.launcher3.model.data.ItemInfoWithIcon;
 import com.android.launcher3.model.data.LauncherAppWidgetInfo;
 import com.android.launcher3.model.data.PackageItemInfo;
+import com.android.launcher3.util.RunnableList;
 import com.android.launcher3.util.SafeCloseable;
 import com.android.launcher3.util.Themes;
+import com.android.launcher3.widget.LauncherWidgetHolder.ProviderChangedListener;
 
 import java.util.List;
 
@@ -81,6 +86,8 @@
     private final Matrix mMatrix = new Matrix();
     private final RectF mPreviewBitmapRect = new RectF();
     private final RectF mCanvasRect = new RectF();
+    private final Handler mHandler = new Handler(Looper.getMainLooper());
+    private final RunnableList mOnDetachCleanup = new RunnableList();
 
     private final LauncherWidgetHolder mWidgetHolder;
     private final LauncherAppWidgetProviderInfo mAppwidget;
@@ -90,7 +97,6 @@
     private final CharSequence mLabel;
 
     private OnClickListener mClickListener;
-    private SafeCloseable mOnDetachCleanup;
 
     private int mDragFlags;
 
@@ -210,16 +216,15 @@
     protected void onAttachedToWindow() {
         super.onAttachedToWindow();
 
+        mOnDetachCleanup.executeAllAndClear();
         if ((mAppwidget != null)
                 && !mInfo.hasRestoreFlag(LauncherAppWidgetInfo.FLAG_ID_NOT_VALID)
                 && mInfo.restoreStatus != LauncherAppWidgetInfo.RESTORE_COMPLETED) {
             // If the widget is not completely restored, but has a valid ID, then listen of
             // updates from provider app for potential restore complete.
-            if (mOnDetachCleanup != null) {
-                mOnDetachCleanup.close();
-            }
-            mOnDetachCleanup = mWidgetHolder.addOnUpdateListener(
+            SafeCloseable updateCleanup = mWidgetHolder.addOnUpdateListener(
                     mInfo.appWidgetId, mAppwidget, this::checkIfRestored);
+            mOnDetachCleanup.add(updateCleanup::close);
             checkIfRestored();
         }
     }
@@ -227,10 +232,7 @@
     @Override
     protected void onDetachedFromWindow() {
         super.onDetachedFromWindow();
-        if (mOnDetachCleanup != null) {
-            mOnDetachCleanup.close();
-            mOnDetachCleanup = null;
-        }
+        mOnDetachCleanup.executeAllAndClear();
     }
 
     /**
@@ -295,43 +297,30 @@
             mCenterDrawable.setCallback(null);
             mCenterDrawable = null;
         }
-        mDragFlags = 0;
-        if (info.bitmap.icon != null) {
-            mDragFlags = FLAG_DRAW_ICON;
+        mDragFlags = FLAG_DRAW_ICON;
 
-            Drawable widgetCategoryIcon = getWidgetCategoryIcon();
-            // The view displays three modes,
-            //   1) App icon in the center
-            //   2) Preload icon in the center
-            //   3) App icon in the center with a setup icon on the top left corner.
-            if (mDisabledForSafeMode) {
-                if (widgetCategoryIcon == null) {
-                    FastBitmapDrawable disabledIcon = info.newIcon(getContext());
-                    disabledIcon.setIsDisabled(true);
-                    mCenterDrawable = disabledIcon;
-                } else {
-                    widgetCategoryIcon.setColorFilter(getDisabledColorFilter());
-                    mCenterDrawable = widgetCategoryIcon;
-                }
-                mSettingIconDrawable = null;
-            } else if (isReadyForClickSetup()) {
-                mCenterDrawable = widgetCategoryIcon == null
-                        ? info.newIcon(getContext())
-                        : widgetCategoryIcon;
-                mSettingIconDrawable = getResources().getDrawable(R.drawable.ic_setting).mutate();
-                updateSettingColor(info.bitmap.color);
+        // The view displays three modes,
+        //   1) App icon in the center
+        //   2) Preload icon in the center
+        //   3) App icon in the center with a setup icon on the top left corner.
+        if (mDisabledForSafeMode) {
+            FastBitmapDrawable disabledIcon = info.newIcon(getContext());
+            disabledIcon.setIsDisabled(true);
+            mCenterDrawable = disabledIcon;
+            mSettingIconDrawable = null;
+        } else if (isReadyForClickSetup()) {
+            mCenterDrawable = info.newIcon(getContext());
+            mSettingIconDrawable = getResources().getDrawable(R.drawable.ic_setting).mutate();
+            updateSettingColor(info.bitmap.color);
 
-                mDragFlags |= FLAG_DRAW_SETTINGS | FLAG_DRAW_LABEL;
-            } else {
-                mCenterDrawable = widgetCategoryIcon == null
-                        ? newPendingIcon(getContext(), info)
-                        : widgetCategoryIcon;
-                mSettingIconDrawable = null;
-                applyState();
-            }
-            mCenterDrawable.setCallback(this);
-            mDrawableSizeChanged = true;
+            mDragFlags |= FLAG_DRAW_SETTINGS | FLAG_DRAW_LABEL;
+        } else {
+            mCenterDrawable = newPendingIcon(getContext(), info);
+            mSettingIconDrawable = null;
+            applyState();
         }
+        mCenterDrawable.setCallback(this);
+        mDrawableSizeChanged = true;
         invalidate();
     }
 
@@ -350,6 +339,11 @@
     }
 
     public void applyState() {
+        if (mCenterDrawable instanceof FastBitmapDrawable fb
+                && mInfo.pendingItemInfo != null
+                && !fb.isSameInfo(mInfo.pendingItemInfo.bitmap)) {
+            reapplyItemInfo(mInfo.pendingItemInfo);
+        }
         if (mCenterDrawable != null) {
             mCenterDrawable.setLevel(Math.max(mInfo.installProgress, 0));
         }
@@ -486,16 +480,72 @@
     }
 
     /**
-     * Returns the widget category icon for {@link #mInfo}.
-     *
-     * <p>If {@link #mInfo}'s category is {@code PackageItemInfo#NO_CATEGORY} or unknown, returns
-     * {@code null}.
+     * Creates a runnable runnable which tries to refresh the widget if it is restored
      */
-    @Nullable
-    private Drawable getWidgetCategoryIcon() {
-        if (mInfo.pendingItemInfo.widgetCategory == WidgetSections.NO_CATEGORY) {
-            return null;
+    public void postProviderAvailabilityCheck() {
+        if (!mInfo.hasRestoreFlag(FLAG_PROVIDER_NOT_READY) && getAppWidgetInfo() == null) {
+            // If the info state suggests that the provider is ready, but there is no
+            // provider info attached on this pending view, recreate when the provider is available
+            DeferredWidgetRefresh restoreRunnable = new DeferredWidgetRefresh();
+            mOnDetachCleanup.add(restoreRunnable::cleanup);
+            mHandler.post(restoreRunnable::notifyWidgetProvidersChanged);
         }
-        return mInfo.pendingItemInfo.newIcon(getContext());
+    }
+
+    /**
+     * Used as a workaround to ensure that the AppWidgetService receives the
+     * PACKAGE_ADDED broadcast before updating widgets.
+     *
+     * This class will periodically check for the availability of the WidgetProvider as a result
+     * of providerChanged callback from the host. When the provider is available or a timeout of
+     * 10-sec is reached, it reinflates the pending-widget which in-turn goes through the process
+     * of re-evaluating the pending state of the widget,
+     */
+    private class DeferredWidgetRefresh implements Runnable, ProviderChangedListener {
+        private boolean mRefreshPending = true;
+
+        DeferredWidgetRefresh() {
+            mWidgetHolder.addProviderChangeListener(this);
+            // Force refresh after 10 seconds, if we don't get the provider changed event.
+            // This could happen when the provider is no longer available in the app.
+            Message msg = Message.obtain(getHandler(), this);
+            msg.obj = DeferredWidgetRefresh.class;
+            mHandler.sendMessageDelayed(msg, 10000);
+        }
+
+        /**
+         * Reinflate the widget if it is still attached.
+         */
+        @Override
+        public void run() {
+            cleanup();
+            if (mRefreshPending) {
+                reInflate();
+                mRefreshPending = false;
+            }
+        }
+
+        @Override
+        public void notifyWidgetProvidersChanged() {
+            final AppWidgetProviderInfo widgetInfo;
+            WidgetManagerHelper widgetHelper = new WidgetManagerHelper(getContext());
+            if (mInfo.hasRestoreFlag(LauncherAppWidgetInfo.FLAG_ID_NOT_VALID)) {
+                widgetInfo = widgetHelper.findProvider(mInfo.providerName, mInfo.user);
+            } else {
+                widgetInfo = widgetHelper.getLauncherAppWidgetInfo(mInfo.appWidgetId,
+                        mInfo.getTargetComponent());
+            }
+            if (widgetInfo != null) {
+                run();
+            }
+        }
+
+        /**
+         * Removes any scheduled callbacks and change listeners, no-op if nothing is scheduled
+         */
+        public void cleanup() {
+            mWidgetHolder.removeProviderChangeListener(this);
+            mHandler.removeCallbacks(this);
+        }
     }
 }
diff --git a/tests/multivalentTests/src/com/android/launcher3/model/ShortcutsChangedTaskTest.kt b/tests/multivalentTests/src/com/android/launcher3/model/ShortcutsChangedTaskTest.kt
new file mode 100644
index 0000000..fb6d038
--- /dev/null
+++ b/tests/multivalentTests/src/com/android/launcher3/model/ShortcutsChangedTaskTest.kt
@@ -0,0 +1,313 @@
+/*
+ * Copyright (C) 2025 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.launcher3.model
+
+import android.content.ComponentName
+import android.content.Intent
+import android.content.pm.ApplicationInfo
+import android.content.pm.ApplicationInfo.FLAG_INSTALLED
+import android.content.pm.LauncherApps
+import android.content.pm.ShortcutInfo
+import android.os.Process.myUserHandle
+import android.os.UserHandle
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import com.android.launcher3.LauncherAppState
+import com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_DEEP_SHORTCUT
+import com.android.launcher3.icons.BitmapInfo
+import com.android.launcher3.icons.CacheableShortcutInfo
+import com.android.launcher3.icons.IconCache
+import com.android.launcher3.model.data.ItemInfo
+import com.android.launcher3.model.data.WorkspaceItemInfo
+import com.android.launcher3.shortcuts.ShortcutKey
+import com.android.launcher3.util.ComponentKey
+import com.android.launcher3.util.IntSparseArrayMap
+import com.android.launcher3.util.LauncherModelHelper
+import com.android.launcher3.util.LauncherModelHelper.SandboxModelContext
+import com.google.common.truth.Truth.assertThat
+import java.util.function.Predicate
+import org.junit.After
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.kotlin.any
+import org.mockito.kotlin.doReturn
+import org.mockito.kotlin.eq
+import org.mockito.kotlin.mock
+import org.mockito.kotlin.times
+import org.mockito.kotlin.verify
+import org.mockito.kotlin.whenever
+
+@SmallTest
+@RunWith(AndroidJUnit4::class)
+class ShortcutsChangedTaskTest {
+    private lateinit var shortcutsChangedTask: ShortcutsChangedTask
+    private lateinit var modelHelper: LauncherModelHelper
+    private lateinit var context: SandboxModelContext
+    private lateinit var launcherApps: LauncherApps
+    private var shortcuts: List<ShortcutInfo> = emptyList()
+
+    private val expectedPackage: String = "expected"
+    private val expectedShortcutId: String = "shortcut_id"
+    private val user: UserHandle = myUserHandle()
+    private val mockTaskController: ModelTaskController = mock()
+    private val mockAllApps: AllAppsList = mock()
+    private val mockAppState: LauncherAppState = mock()
+    private val mockIconCache: IconCache = mock()
+
+    private val expectedWai =
+        WorkspaceItemInfo().apply {
+            id = 1
+            itemType = ITEM_TYPE_DEEP_SHORTCUT
+            intent =
+                Intent().apply {
+                    `package` = expectedPackage
+                    putExtra(ShortcutKey.EXTRA_SHORTCUT_ID, expectedShortcutId)
+                }
+        }
+
+    @Before
+    fun setup() {
+        modelHelper = LauncherModelHelper()
+        modelHelper.loadModelSync()
+        context = modelHelper.sandboxContext
+        launcherApps = context.spyService(LauncherApps::class.java)
+        whenever(mockTaskController.app).thenReturn(mockAppState)
+        whenever(mockAppState.context).thenReturn(context)
+        whenever(mockAppState.iconCache).thenReturn(mockIconCache)
+        whenever(mockIconCache.getShortcutIcon(eq(expectedWai), any<CacheableShortcutInfo>()))
+            .then { _ -> { expectedWai.bitmap = BitmapInfo.LOW_RES_INFO } }
+        shortcuts = emptyList()
+        shortcutsChangedTask = ShortcutsChangedTask(expectedPackage, shortcuts, user, false)
+    }
+
+    @After
+    fun teardown() {
+        modelHelper.destroy()
+    }
+
+    @Test
+    fun `When installed pinned shortcut is found then keep in workspace`() {
+        // Given
+        shortcuts =
+            listOf(
+                mock<ShortcutInfo>().apply {
+                    whenever(isPinned).thenReturn(true)
+                    whenever(id).thenReturn(expectedShortcutId)
+                }
+            )
+        val items: IntSparseArrayMap<ItemInfo> = modelHelper.bgDataModel.itemsIdMap
+        items.put(expectedWai.id, expectedWai)
+        doReturn(
+                ApplicationInfo().apply {
+                    enabled = true
+                    flags = flags or FLAG_INSTALLED
+                    isArchived = false
+                }
+            )
+            .whenever(launcherApps)
+            .getApplicationInfo(eq(expectedPackage), any(), eq(user))
+        doReturn(shortcuts).whenever(launcherApps).getShortcuts(any(), eq(user))
+        // When
+        shortcutsChangedTask.execute(mockTaskController, modelHelper.bgDataModel, mockAllApps)
+        // Then
+        verify(mockAppState.iconCache)
+            .getShortcutIcon(eq(expectedWai), any<CacheableShortcutInfo>())
+        verify(mockTaskController).bindUpdatedWorkspaceItems(listOf(expectedWai))
+    }
+
+    @Test
+    fun `When installed unpinned shortcut is found then remove from workspace`() {
+        // Given
+        shortcuts =
+            listOf(
+                mock<ShortcutInfo>().apply {
+                    whenever(isPinned).thenReturn(false)
+                    whenever(id).thenReturn(expectedShortcutId)
+                }
+            )
+        val items: IntSparseArrayMap<ItemInfo> = modelHelper.bgDataModel.itemsIdMap
+        items.put(expectedWai.id, expectedWai)
+        doReturn(
+                ApplicationInfo().apply {
+                    enabled = true
+                    flags = flags or FLAG_INSTALLED
+                    isArchived = false
+                }
+            )
+            .whenever(launcherApps)
+            .getApplicationInfo(eq(expectedPackage), any(), eq(user))
+        doReturn(shortcuts).whenever(launcherApps).getShortcuts(any(), eq(user))
+        // When
+        shortcutsChangedTask.execute(mockTaskController, modelHelper.bgDataModel, mockAllApps)
+        // Then
+        verify(mockTaskController)
+            .deleteAndBindComponentsRemoved(
+                any<Predicate<ItemInfo?>>(),
+                eq("removed because the shortcut is no longer available in shortcut service"),
+            )
+    }
+
+    @Test
+    fun `When shortcut app is uninstalled then skip handling`() {
+        // Given
+        shortcuts =
+            listOf(
+                mock<ShortcutInfo>().apply {
+                    whenever(isPinned).thenReturn(true)
+                    whenever(id).thenReturn(expectedShortcutId)
+                }
+            )
+        val items: IntSparseArrayMap<ItemInfo> = modelHelper.bgDataModel.itemsIdMap
+        items.put(expectedWai.id, expectedWai)
+        doReturn(
+                ApplicationInfo().apply {
+                    enabled = true
+                    flags = flags and FLAG_INSTALLED.inv()
+                    isArchived = false
+                }
+            )
+            .whenever(launcherApps)
+            .getApplicationInfo(eq(expectedPackage), any(), eq(user))
+        doReturn(shortcuts).whenever(launcherApps).getShortcuts(any(), eq(user))
+        // When
+        shortcutsChangedTask.execute(mockTaskController, modelHelper.bgDataModel, mockAllApps)
+        // Then
+        verify(mockTaskController, times(0)).deleteAndBindComponentsRemoved(any(), any())
+        verify(mockTaskController, times(0)).bindUpdatedWorkspaceItems(any())
+    }
+
+    @Test
+    fun `When archived pinned shortcut is found then keep in workspace`() {
+        // Given
+        shortcuts =
+            listOf(
+                mock<ShortcutInfo>().apply {
+                    whenever(isPinned).thenReturn(true)
+                    whenever(id).thenReturn(expectedShortcutId)
+                }
+            )
+        val items: IntSparseArrayMap<ItemInfo> = modelHelper.bgDataModel.itemsIdMap
+        items.put(expectedWai.id, expectedWai)
+        doReturn(
+                ApplicationInfo().apply {
+                    enabled = true
+                    flags = flags or FLAG_INSTALLED
+                    isArchived = true
+                }
+            )
+            .whenever(launcherApps)
+            .getApplicationInfo(eq(expectedPackage), any(), eq(user))
+        doReturn(shortcuts).whenever(launcherApps).getShortcuts(any(), eq(user))
+        // When
+        shortcutsChangedTask.execute(mockTaskController, modelHelper.bgDataModel, mockAllApps)
+        // Then
+        verify(mockAppState.iconCache)
+            .getShortcutIcon(eq(expectedWai), any<CacheableShortcutInfo>())
+        verify(mockTaskController).bindUpdatedWorkspaceItems(listOf(expectedWai))
+    }
+
+    @Test
+    fun `When archived unpinned shortcut is found then keep in workspace`() {
+        // Given
+        shortcuts =
+            listOf(
+                mock<ShortcutInfo>().apply {
+                    whenever(isPinned).thenReturn(true)
+                    whenever(id).thenReturn(expectedShortcutId)
+                }
+            )
+        val items: IntSparseArrayMap<ItemInfo> = modelHelper.bgDataModel.itemsIdMap
+        items.put(expectedWai.id, expectedWai)
+        doReturn(
+                ApplicationInfo().apply {
+                    enabled = true
+                    flags = flags or FLAG_INSTALLED
+                    isArchived = true
+                }
+            )
+            .whenever(launcherApps)
+            .getApplicationInfo(eq(expectedPackage), any(), eq(user))
+        doReturn(shortcuts).whenever(launcherApps).getShortcuts(any(), eq(user))
+        // When
+        shortcutsChangedTask.execute(mockTaskController, modelHelper.bgDataModel, mockAllApps)
+        // Then
+        verify(mockAppState.iconCache)
+            .getShortcutIcon(eq(expectedWai), any<CacheableShortcutInfo>())
+        verify(mockTaskController).bindUpdatedWorkspaceItems(listOf(expectedWai))
+    }
+
+    @Test
+    fun `When updateIdMap true then trigger deep shortcut binding`() {
+        // Given
+        val expectedShortcut =
+            mock<ShortcutInfo>().apply {
+                whenever(isEnabled).thenReturn(true)
+                whenever(isDeclaredInManifest).thenReturn(true)
+                whenever(activity).thenReturn(ComponentName(expectedPackage, "expectedClass"))
+                whenever(id).thenReturn(expectedShortcutId)
+                whenever(userHandle).thenReturn(user)
+            }
+        shortcuts = listOf(expectedShortcut)
+        val expectedKey = ComponentKey(expectedShortcut.activity, expectedShortcut.userHandle)
+        doReturn(ApplicationInfo())
+            .whenever(launcherApps)
+            .getApplicationInfo(eq(expectedPackage), any(), eq(user))
+        shortcutsChangedTask =
+            ShortcutsChangedTask(
+                packageName = expectedPackage,
+                shortcuts = shortcuts,
+                user = user,
+                shouldUpdateIdMap = true,
+            )
+        // When
+        shortcutsChangedTask.execute(mockTaskController, modelHelper.bgDataModel, mockAllApps)
+        // Then
+        assertThat(modelHelper.bgDataModel.deepShortcutMap).containsEntry(expectedKey, 1)
+        verify(mockTaskController).bindDeepShortcuts(eq(modelHelper.bgDataModel))
+    }
+
+    @Test
+    fun `When updateIdMap false then do not trigger deep shortcut binding`() {
+        // Given
+        val expectedShortcut =
+            mock<ShortcutInfo>().apply {
+                whenever(isEnabled).thenReturn(true)
+                whenever(isDeclaredInManifest).thenReturn(true)
+                whenever(activity).thenReturn(ComponentName(expectedPackage, "expectedClass"))
+                whenever(id).thenReturn(expectedShortcutId)
+                whenever(userHandle).thenReturn(user)
+            }
+        shortcuts = listOf(expectedShortcut)
+        val expectedKey = ComponentKey(expectedShortcut.activity, expectedShortcut.userHandle)
+        doReturn(ApplicationInfo())
+            .whenever(launcherApps)
+            .getApplicationInfo(eq(expectedPackage), any(), eq(user))
+        shortcutsChangedTask =
+            ShortcutsChangedTask(
+                packageName = expectedPackage,
+                shortcuts = shortcuts,
+                user = user,
+                shouldUpdateIdMap = false,
+            )
+        // When
+        shortcutsChangedTask.execute(mockTaskController, modelHelper.bgDataModel, mockAllApps)
+        // Then
+        assertThat(modelHelper.bgDataModel.deepShortcutMap).doesNotContainKey(expectedKey)
+        verify(mockTaskController, times(0)).bindDeepShortcuts(eq(modelHelper.bgDataModel))
+    }
+}
diff --git a/tests/multivalentTests/src/com/android/launcher3/shapes/IconShapesProviderTest.kt b/tests/multivalentTests/src/com/android/launcher3/shapes/IconShapesProviderTest.kt
new file mode 100644
index 0000000..234e050
--- /dev/null
+++ b/tests/multivalentTests/src/com/android/launcher3/shapes/IconShapesProviderTest.kt
@@ -0,0 +1,90 @@
+/*
+ * Copyright (C) 2025 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.launcher3.shapes
+
+import android.platform.test.annotations.EnableFlags
+import android.platform.test.flag.junit.SetFlagsRule
+import androidx.core.graphics.PathParser
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import com.android.launcher3.Flags.FLAG_ENABLE_LAUNCHER_ICON_SHAPES
+import com.android.launcher3.graphics.IconShape.GenericPathShape
+import com.android.systemui.shared.Flags.FLAG_NEW_CUSTOMIZATION_PICKER_UI
+import org.junit.Rule
+import org.junit.Test
+import org.junit.runner.RunWith
+
+@SmallTest
+@RunWith(AndroidJUnit4::class)
+class IconShapesProviderTest {
+
+    @get:Rule val setFlagsRule: SetFlagsRule = SetFlagsRule()
+
+    @Test
+    @EnableFlags(FLAG_ENABLE_LAUNCHER_ICON_SHAPES, FLAG_NEW_CUSTOMIZATION_PICKER_UI)
+    fun `verify valid path arch`() {
+        IconShapesProvider.shapes["arch"]?.apply {
+            GenericPathShape(pathString)
+            PathParser.createPathFromPathData(pathString)
+        }
+    }
+
+    @Test
+    @EnableFlags(FLAG_ENABLE_LAUNCHER_ICON_SHAPES, FLAG_NEW_CUSTOMIZATION_PICKER_UI)
+    fun `verify valid path 4_sided_cookie`() {
+        IconShapesProvider.shapes["4_sided_cookie"]?.apply {
+            GenericPathShape(pathString)
+            PathParser.createPathFromPathData(pathString)
+        }
+    }
+
+    @Test
+    @EnableFlags(FLAG_ENABLE_LAUNCHER_ICON_SHAPES, FLAG_NEW_CUSTOMIZATION_PICKER_UI)
+    fun `verify valid path seven_sided_cookie`() {
+        IconShapesProvider.shapes["seven_sided_cookie"]?.apply {
+            GenericPathShape(pathString)
+            PathParser.createPathFromPathData(pathString)
+        }
+    }
+
+    @Test
+    @EnableFlags(FLAG_ENABLE_LAUNCHER_ICON_SHAPES, FLAG_NEW_CUSTOMIZATION_PICKER_UI)
+    fun `verify valid path sunny`() {
+        IconShapesProvider.shapes["sunny"]?.apply {
+            GenericPathShape(pathString)
+            PathParser.createPathFromPathData(pathString)
+        }
+    }
+
+    @Test
+    @EnableFlags(FLAG_ENABLE_LAUNCHER_ICON_SHAPES, FLAG_NEW_CUSTOMIZATION_PICKER_UI)
+    fun `verify valid path circle`() {
+        IconShapesProvider.shapes["circle"]?.apply {
+            GenericPathShape(pathString)
+            PathParser.createPathFromPathData(pathString)
+        }
+    }
+
+    @Test
+    @EnableFlags(FLAG_ENABLE_LAUNCHER_ICON_SHAPES, FLAG_NEW_CUSTOMIZATION_PICKER_UI)
+    fun `verify valid path square`() {
+        IconShapesProvider.shapes["square"]?.apply {
+            GenericPathShape(pathString)
+            PathParser.createPathFromPathData(pathString)
+        }
+    }
+}
diff --git a/tests/multivalentTests/src/com/android/launcher3/ui/BubbleTextViewTest.java b/tests/multivalentTests/src/com/android/launcher3/ui/BubbleTextViewTest.java
index f51871b..faf6b91 100644
--- a/tests/multivalentTests/src/com/android/launcher3/ui/BubbleTextViewTest.java
+++ b/tests/multivalentTests/src/com/android/launcher3/ui/BubbleTextViewTest.java
@@ -27,7 +27,9 @@
 import static com.android.launcher3.Flags.FLAG_ENABLE_SUPPORT_FOR_ARCHIVING;
 import static com.android.launcher3.Flags.FLAG_USE_NEW_ICON_FOR_ARCHIVED_APPS;
 import static com.android.launcher3.LauncherPrefs.ENABLE_TWOLINE_ALLAPPS_TOGGLE;
+import static com.android.launcher3.anim.AnimatorListeners.forEndCallback;
 import static com.android.launcher3.model.data.ItemInfoWithIcon.FLAG_ARCHIVED;
+import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
 
 import static com.google.common.truth.Truth.assertThat;
 
@@ -39,6 +41,8 @@
 import android.content.ComponentName;
 import android.content.Context;
 import android.content.Intent;
+import android.graphics.Bitmap;
+import android.graphics.Bitmap.Config;
 import android.graphics.Typeface;
 import android.os.Build;
 import android.os.UserHandle;
@@ -57,13 +61,17 @@
 import com.android.launcher3.Flags;
 import com.android.launcher3.LauncherPrefs;
 import com.android.launcher3.Utilities;
+import com.android.launcher3.graphics.PreloadIconDrawable;
+import com.android.launcher3.icons.BitmapInfo;
 import com.android.launcher3.model.data.AppInfo;
 import com.android.launcher3.model.data.ItemInfoWithIcon;
+import com.android.launcher3.pm.PackageInstallInfo;
 import com.android.launcher3.search.StringMatcherUtility;
 import com.android.launcher3.util.ActivityContextWrapper;
 import com.android.launcher3.util.FlagOp;
 import com.android.launcher3.util.IntArray;
 import com.android.launcher3.util.LauncherModelHelper.SandboxModelContext;
+import com.android.launcher3.util.TestUtil;
 import com.android.launcher3.views.BaseDragLayer;
 
 import org.junit.After;
@@ -73,6 +81,8 @@
 import org.junit.runner.RunWith;
 import org.mockito.MockitoAnnotations;
 
+import java.util.concurrent.CountDownLatch;
+
 /**
  * Unit tests for testing modifyTitleToSupportMultiLine() in BubbleTextView.java
  * This class tests a couple of strings and uses the getMaxLines() to determine if the test passes.
@@ -485,4 +495,40 @@
 
         assertThat(mBubbleTextView.getIcon().hasBadge()).isEqualTo(true);
     }
+
+    @Test
+    public void applyingPendingIcon_preserves_last_icon() throws Exception {
+        mItemInfoWithIcon.bitmap =
+                BitmapInfo.fromBitmap(Bitmap.createBitmap(100, 100, Config.ARGB_8888));
+        mItemInfoWithIcon.setProgressLevel(30, PackageInstallInfo.STATUS_INSTALLING);
+
+        TestUtil.runOnExecutorSync(MAIN_EXECUTOR,
+                () -> mBubbleTextView.applyIconAndLabel(mItemInfoWithIcon));
+        assertThat(mBubbleTextView.getIcon()).isInstanceOf(PreloadIconDrawable.class);
+        assertThat(mBubbleTextView.getIcon().getLevel()).isEqualTo(30);
+        PreloadIconDrawable oldIcon = (PreloadIconDrawable) mBubbleTextView.getIcon();
+
+        // Same icon is used when progress changes
+        mItemInfoWithIcon.setProgressLevel(50, PackageInstallInfo.STATUS_INSTALLING);
+        TestUtil.runOnExecutorSync(MAIN_EXECUTOR,
+                () -> mBubbleTextView.applyIconAndLabel(mItemInfoWithIcon));
+        assertThat(mBubbleTextView.getIcon()).isSameInstanceAs(oldIcon);
+        assertThat(mBubbleTextView.getIcon().getLevel()).isEqualTo(50);
+
+        // Icon is replaced with a non pending icon when download finishes
+        mItemInfoWithIcon.setProgressLevel(100, PackageInstallInfo.STATUS_INSTALLED);
+
+        CountDownLatch animWait = new CountDownLatch(1);
+        TestUtil.runOnExecutorSync(MAIN_EXECUTOR, () -> {
+            mBubbleTextView.applyIconAndLabel(mItemInfoWithIcon);
+            assertThat(mBubbleTextView.getIcon()).isSameInstanceAs(oldIcon);
+            assertThat(oldIcon.getActiveAnimation()).isNotNull();
+            oldIcon.getActiveAnimation().addListener(forEndCallback(animWait::countDown));
+        });
+        animWait.await();
+
+        // Assert that the icon is replaced with a non-pending icon
+        assertThat(mBubbleTextView.getIcon()).isNotInstanceOf(PreloadIconDrawable.class);
+    }
+
 }
diff --git a/tests/multivalentTests/src/com/android/launcher3/util/LauncherBindableItemsContainerTest.kt b/tests/multivalentTests/src/com/android/launcher3/util/LauncherBindableItemsContainerTest.kt
new file mode 100644
index 0000000..93be5f5
--- /dev/null
+++ b/tests/multivalentTests/src/com/android/launcher3/util/LauncherBindableItemsContainerTest.kt
@@ -0,0 +1,137 @@
+/*
+ * Copyright (C) 2025 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.launcher3.util
+
+import android.content.ComponentName
+import android.content.pm.LauncherApps
+import android.graphics.Bitmap
+import android.graphics.Bitmap.Config.ARGB_8888
+import android.os.Process.myUserHandle
+import android.platform.uiautomatorhelpers.DeviceHelpers.context
+import android.view.View
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import com.android.launcher3.BubbleTextView
+import com.android.launcher3.graphics.PreloadIconDrawable
+import com.android.launcher3.icons.BitmapInfo
+import com.android.launcher3.icons.FastBitmapDrawable
+import com.android.launcher3.icons.PlaceHolderIconDrawable
+import com.android.launcher3.model.data.AppInfo
+import com.android.launcher3.model.data.AppInfo.makeLaunchIntent
+import com.android.launcher3.model.data.ItemInfo
+import com.android.launcher3.model.data.WorkspaceItemInfo
+import com.android.launcher3.pm.PackageInstallInfo
+import com.android.launcher3.util.Executors.MAIN_EXECUTOR
+import com.android.launcher3.util.LauncherBindableItemsContainer.ItemOperator
+import com.android.launcher3.util.LauncherModelHelper.TEST_ACTIVITY
+import com.android.launcher3.util.LauncherModelHelper.TEST_ACTIVITY2
+import com.android.launcher3.util.LauncherModelHelper.TEST_ACTIVITY3
+import com.android.launcher3.util.LauncherModelHelper.TEST_PACKAGE
+import com.google.common.truth.Truth.assertThat
+import org.junit.Test
+import org.junit.runner.RunWith
+
+@SmallTest
+@RunWith(AndroidJUnit4::class)
+class LauncherBindableItemsContainerTest {
+
+    private val icon1 by lazy { getLAI(TEST_ACTIVITY) }
+    private val icon2 by lazy { getLAI(TEST_ACTIVITY2) }
+    private val icon3 by lazy { getLAI(TEST_ACTIVITY3) }
+
+    private val container = TestContainer()
+
+    @Test
+    fun `icon bitmap is updated`() {
+        container.addIcon(icon1)
+        container.addIcon(icon2)
+        container.addIcon(icon3)
+
+        assertThat(container.getAppIcon(icon1).icon)
+            .isInstanceOf(PlaceHolderIconDrawable::class.java)
+        assertThat(container.getAppIcon(icon2).icon)
+            .isInstanceOf(PlaceHolderIconDrawable::class.java)
+        assertThat(container.getAppIcon(icon3).icon)
+            .isInstanceOf(PlaceHolderIconDrawable::class.java)
+
+        icon2.bitmap = BitmapInfo.fromBitmap(Bitmap.createBitmap(200, 200, ARGB_8888))
+        TestUtil.runOnExecutorSync(MAIN_EXECUTOR) {
+            container.updateContainerItems(setOf(icon2), container)
+        }
+
+        assertThat(container.getAppIcon(icon1).icon)
+            .isInstanceOf(PlaceHolderIconDrawable::class.java)
+        assertThat(container.getAppIcon(icon3).icon)
+            .isInstanceOf(PlaceHolderIconDrawable::class.java)
+        assertThat(container.getAppIcon(icon2).icon)
+            .isNotInstanceOf(PlaceHolderIconDrawable::class.java)
+        assertThat(container.getAppIcon(icon2).icon).isInstanceOf(FastBitmapDrawable::class.java)
+    }
+
+    @Test
+    fun `icon download progress updated`() {
+        container.addIcon(icon1)
+        container.addIcon(icon2)
+        assertThat(container.getAppIcon(icon1).icon)
+            .isInstanceOf(PlaceHolderIconDrawable::class.java)
+        assertThat(container.getAppIcon(icon2).icon)
+            .isInstanceOf(PlaceHolderIconDrawable::class.java)
+
+        icon1.status = WorkspaceItemInfo.FLAG_RESTORED_ICON
+        icon1.bitmap = BitmapInfo.fromBitmap(Bitmap.createBitmap(200, 200, ARGB_8888))
+        icon1.setProgressLevel(30, PackageInstallInfo.STATUS_INSTALLING)
+        TestUtil.runOnExecutorSync(MAIN_EXECUTOR) {
+            container.updateContainerItems(setOf(icon1), container)
+        }
+
+        assertThat(container.getAppIcon(icon2).icon)
+            .isInstanceOf(PlaceHolderIconDrawable::class.java)
+        assertThat(container.getAppIcon(icon1).icon).isInstanceOf(PreloadIconDrawable::class.java)
+        val oldIcon = container.getAppIcon(icon1).icon as PreloadIconDrawable
+        assertThat(oldIcon.level).isEqualTo(30)
+    }
+
+    private fun getLAI(className: String): WorkspaceItemInfo =
+        AppInfo(
+                context,
+                context
+                    .getSystemService(LauncherApps::class.java)!!
+                    .resolveActivity(
+                        makeLaunchIntent(ComponentName(TEST_PACKAGE, className)),
+                        myUserHandle(),
+                    )!!,
+                myUserHandle(),
+            )
+            .makeWorkspaceItem(context)
+
+    class TestContainer : ActivityContextWrapper(context), LauncherBindableItemsContainer {
+
+        val items = mutableMapOf<ItemInfo, View>()
+
+        override fun mapOverItems(op: ItemOperator) {
+            items.forEach { (item, view) -> if (op.evaluate(item, view)) return@forEach }
+        }
+
+        fun addIcon(info: WorkspaceItemInfo) {
+            val btv = BubbleTextView(this)
+            btv.applyFromWorkspaceItem(info)
+            items[info] = btv
+        }
+
+        fun getAppIcon(info: WorkspaceItemInfo) = items[info] as BubbleTextView
+    }
+}
diff --git a/tests/src/com/android/launcher3/ui/TaplTestsLauncher3Test.java b/tests/src/com/android/launcher3/ui/TaplTestsLauncher3Test.java
index 342eedf..1338e60 100644
--- a/tests/src/com/android/launcher3/ui/TaplTestsLauncher3Test.java
+++ b/tests/src/com/android/launcher3/ui/TaplTestsLauncher3Test.java
@@ -21,7 +21,6 @@
 import androidx.test.filters.LargeTest;
 import androidx.test.runner.AndroidJUnit4;
 
-import com.android.launcher3.util.rule.ScreenRecordRule.ScreenRecord;
 import com.android.launcher3.Launcher;
 
 import org.junit.Test;
@@ -31,9 +30,8 @@
 @RunWith(AndroidJUnit4.class)
 public class TaplTestsLauncher3Test extends AbstractLauncherUiTest<Launcher> {
 
-    @ScreenRecord // b/322823478
     @Test
-    public void testDevicePressMenu() throws Exception {
+    public void testDevicePressMenu() {
         mDevice.pressMenu();
         mDevice.waitForIdle();
         executeOnLauncher(