Snap for 13258483 from 88b2f23a83252572a53cae7dbf4d27de54592553 to 25Q2-release

Change-Id: Iedd6e4899c8772ccfac46029ff90fb71bf04b65a
diff --git a/quickstep/src/com/android/launcher3/statehandlers/DepthController.java b/quickstep/src/com/android/launcher3/statehandlers/DepthController.java
index 360210b..d980830 100644
--- a/quickstep/src/com/android/launcher3/statehandlers/DepthController.java
+++ b/quickstep/src/com/android/launcher3/statehandlers/DepthController.java
@@ -30,6 +30,8 @@
 import android.view.ViewRootImpl;
 import android.view.ViewTreeObserver;
 
+import androidx.annotation.VisibleForTesting;
+
 import com.android.launcher3.BaseActivity;
 import com.android.launcher3.Launcher;
 import com.android.launcher3.LauncherState;
@@ -46,8 +48,8 @@
  */
 public class DepthController extends BaseDepthController implements StateHandler<LauncherState>,
         BaseActivity.MultiWindowModeChangedListener {
-
-    private final ViewTreeObserver.OnDrawListener mOnDrawListener = this::onLauncherDraw;
+    @VisibleForTesting
+    final ViewTreeObserver.OnDrawListener mOnDrawListener = this::onLauncherDraw;
 
     private final Consumer<Boolean> mCrossWindowBlurListener = this::setCrossWindowBlursEnabled;
 
@@ -58,6 +60,10 @@
 
     private View.OnAttachStateChangeListener mOnAttachListener;
 
+    // Ensure {@link mOnDrawListener} is added only once to avoid spamming DragLayer's mRunQueue
+    // via {@link View#post(Runnable)}
+    private boolean mIsOnDrawListenerAdded = false;
+
     public DepthController(Launcher l) {
         super(l);
     }
@@ -66,33 +72,37 @@
         View view = mLauncher.getDragLayer();
         ViewRootImpl viewRootImpl = view.getViewRootImpl();
         setBaseSurface(viewRootImpl != null ? viewRootImpl.getSurfaceControl() : null);
-        view.post(() -> view.getViewTreeObserver().removeOnDrawListener(mOnDrawListener));
+        view.post(this::removeOnDrawListener);
     }
 
     private void ensureDependencies() {
-        if (mLauncher.getRootView() != null && mOnAttachListener == null) {
-            View rootView = mLauncher.getRootView();
-            mOnAttachListener = new View.OnAttachStateChangeListener() {
-                @Override
-                public void onViewAttachedToWindow(View view) {
-                    UI_HELPER_EXECUTOR.execute(() ->
-                            CrossWindowBlurListeners.getInstance().addListener(
-                                    mLauncher.getMainExecutor(), mCrossWindowBlurListener));
-                    mLauncher.getScrimView().addOpaquenessListener(mOpaquenessListener);
+        View rootView = mLauncher.getRootView();
+        if (rootView == null) {
+            return;
+        }
+        if (mOnAttachListener != null) {
+            return;
+        }
+        mOnAttachListener = new View.OnAttachStateChangeListener() {
+            @Override
+            public void onViewAttachedToWindow(View view) {
+                UI_HELPER_EXECUTOR.execute(() ->
+                        CrossWindowBlurListeners.getInstance().addListener(
+                                mLauncher.getMainExecutor(), mCrossWindowBlurListener));
+                mLauncher.getScrimView().addOpaquenessListener(mOpaquenessListener);
 
-                    // To handle the case where window token is invalid during last setDepth call.
-                    applyDepthAndBlur();
-                }
-
-                @Override
-                public void onViewDetachedFromWindow(View view) {
-                    removeSecondaryListeners();
-                }
-            };
-            rootView.addOnAttachStateChangeListener(mOnAttachListener);
-            if (rootView.isAttachedToWindow()) {
-                mOnAttachListener.onViewAttachedToWindow(rootView);
+                // To handle the case where window token is invalid during last setDepth call.
+                applyDepthAndBlur();
             }
+
+            @Override
+            public void onViewDetachedFromWindow(View view) {
+                removeSecondaryListeners();
+            }
+        };
+        rootView.addOnAttachStateChangeListener(mOnAttachListener);
+        if (rootView.isAttachedToWindow()) {
+            mOnAttachListener.onViewAttachedToWindow(rootView);
         }
     }
 
@@ -109,11 +119,9 @@
     }
 
     private void removeSecondaryListeners() {
-        if (mCrossWindowBlurListener != null) {
-            UI_HELPER_EXECUTOR.execute(() ->
-                    CrossWindowBlurListeners.getInstance()
-                            .removeListener(mCrossWindowBlurListener));
-        }
+        UI_HELPER_EXECUTOR.execute(() ->
+                CrossWindowBlurListeners.getInstance()
+                        .removeListener(mCrossWindowBlurListener));
         if (mOpaquenessListener != null) {
             mLauncher.getScrimView().removeOpaquenessListener(mOpaquenessListener);
         }
@@ -124,9 +132,9 @@
      */
     public void setActivityStarted(boolean isStarted) {
         if (isStarted) {
-            mLauncher.getDragLayer().getViewTreeObserver().addOnDrawListener(mOnDrawListener);
+            addOnDrawListener();
         } else {
-            mLauncher.getDragLayer().getViewTreeObserver().removeOnDrawListener(mOnDrawListener);
+            removeOnDrawListener();
             setBaseSurface(null);
         }
     }
@@ -139,7 +147,7 @@
 
         stateDepth.setValue(toState.getDepth(mLauncher));
         if (toState == LauncherState.BACKGROUND_APP) {
-            mLauncher.getDragLayer().getViewTreeObserver().addOnDrawListener(mOnDrawListener);
+            addOnDrawListener();
         }
     }
 
@@ -165,7 +173,23 @@
     @Override
     protected void onInvalidSurface() {
         // Lets wait for surface to become valid again
+        addOnDrawListener();
+    }
+
+    private void addOnDrawListener() {
+        if (mIsOnDrawListenerAdded) {
+            return;
+        }
         mLauncher.getDragLayer().getViewTreeObserver().addOnDrawListener(mOnDrawListener);
+        mIsOnDrawListenerAdded = true;
+    }
+
+    private void removeOnDrawListener() {
+        if (!mIsOnDrawListenerAdded) {
+            return;
+        }
+        mLauncher.getDragLayer().getViewTreeObserver().removeOnDrawListener(mOnDrawListener);
+        mIsOnDrawListenerAdded = false;
     }
 
     @Override
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarViewController.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarViewController.java
index f15ef8c..ac9d2ba 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarViewController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarViewController.java
@@ -496,7 +496,7 @@
 
         float finalScale;
         TaskbarSharedState sharedState = mControllers.getSharedState();
-        if (sharedState != null || mControllers.getSharedState().startTaskbarVariantIsTransient) {
+        if (sharedState != null && sharedState.startTaskbarVariantIsTransient) {
             finalScale = mapRange(scale, 1f, ((float) mPersistentIconSize / mTransientIconSize));
         } else {
             finalScale = mapRange(scale, ((float) mTransientIconSize / mPersistentIconSize), 1f);
diff --git a/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java b/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
index 47aa1de..48630b1 100644
--- a/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
+++ b/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
@@ -161,7 +161,6 @@
 import com.android.quickstep.views.TaskViewType;
 import com.android.systemui.animation.TransitionAnimator;
 import com.android.systemui.contextualeducation.GestureType;
-import com.android.systemui.shared.recents.model.Task;
 import com.android.systemui.shared.recents.model.ThumbnailData;
 import com.android.systemui.shared.system.ActivityManagerWrapper;
 import com.android.systemui.shared.system.InputConsumerController;
@@ -170,6 +169,7 @@
 import com.android.systemui.shared.system.TaskStackChangeListener;
 import com.android.systemui.shared.system.TaskStackChangeListeners;
 import com.android.wm.shell.Flags;
+import com.android.wm.shell.shared.GroupedTaskInfo;
 import com.android.wm.shell.shared.TransactionPool;
 import com.android.wm.shell.shared.desktopmode.DesktopModeStatus;
 import com.android.wm.shell.shared.startingsurface.SplashScreenExitAnimationUtils;
@@ -182,7 +182,6 @@
 import java.util.Arrays;
 import java.util.HashMap;
 import java.util.List;
-import java.util.Objects;
 import java.util.Optional;
 import java.util.OptionalInt;
 import java.util.function.Consumer;
@@ -693,26 +692,23 @@
     }
 
     protected void notifyGestureAnimationStartToRecents() {
-        Task[] runningTasks;
-        TopTaskTracker.CachedTaskInfo cachedTaskInfo = mGestureState.getRunningTask();
-        if (mIsSwipeForSplit) {
-            int[] splitTaskIds = TopTaskTracker.INSTANCE.get(mContext).getRunningSplitTaskIds();
-            runningTasks = cachedTaskInfo.getSplitPlaceholderTasks(splitTaskIds);
-        } else {
-            runningTasks = cachedTaskInfo.getPlaceholderTasks();
-        }
+        int[] splitTaskIds = mIsSwipeForSplit
+                ? TopTaskTracker.INSTANCE.get(mContext).getRunningSplitTaskIds()
+                : null;
+        GroupedTaskInfo groupedTaskInfo =
+                mGestureState.getRunningTask().getPlaceholderGroupedTaskInfo(splitTaskIds);
 
         // Safeguard against any null tasks being sent to recents view, happens when quickswitching
         // very quickly w/ split tasks because TopTaskTracker provides stale information compared to
         // actual running tasks in the recents animation.
         // TODO(b/236226779), Proper fix (ag/22237143)
-        if (Arrays.stream(runningTasks).anyMatch(Objects::isNull)) {
+        if (groupedTaskInfo == null) {
             return;
         }
         if (mRecentsView == null) {
             return;
         }
-        mRecentsView.onGestureAnimationStart(runningTasks);
+        mRecentsView.onGestureAnimationStart(groupedTaskInfo);
         TaskView currentPageTaskView = mRecentsView.getCurrentPageTaskView();
         if (currentPageTaskView != null) {
             mPreviousTaskViewType = currentPageTaskView.getType();
diff --git a/quickstep/src/com/android/quickstep/FallbackSwipeHandler.java b/quickstep/src/com/android/quickstep/FallbackSwipeHandler.java
index 7d8a53d..9365383 100644
--- a/quickstep/src/com/android/quickstep/FallbackSwipeHandler.java
+++ b/quickstep/src/com/android/quickstep/FallbackSwipeHandler.java
@@ -218,7 +218,8 @@
         if (mRunningOverHome) {
             if (DisplayController.getNavigationMode(mContext).hasGestures) {
                 mRecentsView.onGestureAnimationStartOnHome(
-                        mGestureState.getRunningTask().getPlaceholderTasks());
+                        mGestureState.getRunningTask().getPlaceholderGroupedTaskInfo(
+                                /* splitTaskIds = */ null));
             }
         } else {
             super.notifyGestureAnimationStartToRecents();
diff --git a/quickstep/src/com/android/quickstep/TopTaskTracker.java b/quickstep/src/com/android/quickstep/TopTaskTracker.java
index 7e773e3..43e8ff9 100644
--- a/quickstep/src/com/android/quickstep/TopTaskTracker.java
+++ b/quickstep/src/com/android/quickstep/TopTaskTracker.java
@@ -29,9 +29,13 @@
 import static com.android.wm.shell.Flags.enableShellTopTaskTracking;
 import static com.android.wm.shell.Flags.enableFlexibleSplit;
 import static com.android.wm.shell.shared.GroupedTaskInfo.TYPE_SPLIT;
+import static com.android.launcher3.statehandlers.DesktopVisibilityController.INACTIVE_DESK_ID;
+import static com.android.wm.shell.shared.desktopmode.DesktopModeStatus.canEnterDesktopMode;
 
 import android.app.ActivityManager.RunningTaskInfo;
 import android.app.TaskInfo;
+import android.app.WindowConfiguration;
+import android.content.Context;
 import android.util.ArrayMap;
 import android.util.Log;
 
@@ -39,6 +43,7 @@
 import androidx.annotation.Nullable;
 import androidx.annotation.UiThread;
 
+import com.android.launcher3.dagger.ApplicationContext;
 import com.android.launcher3.dagger.LauncherAppSingleton;
 import com.android.launcher3.util.DaggerSingletonObject;
 import com.android.launcher3.util.DaggerSingletonTracker;
@@ -50,8 +55,6 @@
 import com.android.quickstep.dagger.QuickstepBaseAppComponent;
 import com.android.quickstep.util.DesksUtils;
 import com.android.quickstep.util.ExternalDisplaysKt;
-import com.android.systemui.shared.recents.model.Task;
-import com.android.systemui.shared.recents.model.Task.TaskKey;
 import com.android.systemui.shared.system.ActivityManagerWrapper;
 import com.android.systemui.shared.system.TaskStackChangeListener;
 import com.android.systemui.shared.system.TaskStackChangeListeners;
@@ -92,8 +95,11 @@
     // most.
     private ArrayMap<Integer, GroupedTaskInfo> mVisibleTasks = new ArrayMap<>();
 
+    private final boolean mCanEnterDesktopMode;
+
     @Inject
-    public TopTaskTracker(DaggerSingletonTracker tracker, SystemUiProxy systemUiProxy) {
+    public TopTaskTracker(@ApplicationContext Context context, DaggerSingletonTracker tracker,
+            SystemUiProxy systemUiProxy) {
         if (!enableShellTopTaskTracking()) {
             mMainStagePosition.stageType = SplitConfigurationOptions.STAGE_TYPE_MAIN;
             mSideStagePosition.stageType = SplitConfigurationOptions.STAGE_TYPE_SIDE;
@@ -110,6 +116,8 @@
             TaskStackChangeListeners.getInstance().unregisterTaskStackListener(this);
             systemUiProxy.unregisterSplitScreenListener(this);
         });
+
+        mCanEnterDesktopMode = canEnterDesktopMode(context);
     }
 
     @Override
@@ -324,7 +332,7 @@
             // TODO(346588978): Currently ignore filterOnlyVisibleRecents, but perhaps make this an
             //  explicit filter For things to ignore (ie. PIP/Bubbles/Assistant/etc/so that this is
             //  explicit)
-            return new CachedTaskInfo(mVisibleTasks.get(displayId));
+            return new CachedTaskInfo(mVisibleTasks.get(displayId), mCanEnterDesktopMode);
         } else {
             if (filterOnlyVisibleRecents) {
                 // Since we only know about the top most task, any filtering may not be applied on
@@ -335,9 +343,9 @@
                 if (enableOverviewOnConnectedDisplays()) {
                     return new CachedTaskInfo(Arrays.stream(tasks).filter(
                             info -> ExternalDisplaysKt.getSafeDisplayId(info)
-                                    == displayId).toList());
+                                    == displayId).toList(), mCanEnterDesktopMode);
                 } else {
-                    return new CachedTaskInfo(Arrays.asList(tasks));
+                    return new CachedTaskInfo(Arrays.asList(tasks), mCanEnterDesktopMode);
                 }
             }
 
@@ -354,9 +362,10 @@
                     ||  DesksUtils.isDesktopWallpaperTask(t));
             if (enableOverviewOnConnectedDisplays()) {
                 return new CachedTaskInfo(tasks.stream().filter(
-                        info -> ExternalDisplaysKt.getSafeDisplayId(info) == displayId).toList());
+                        info -> ExternalDisplaysKt.getSafeDisplayId(info) == displayId).toList(),
+                        mCanEnterDesktopMode);
             } else {
-                return new CachedTaskInfo(tasks);
+                return new CachedTaskInfo(tasks, mCanEnterDesktopMode);
             }
         }
     }
@@ -376,6 +385,8 @@
      * during the lifecycle of the task.
      */
     public static class CachedTaskInfo {
+        // TODO: b/402362465 - Provide a valid value while tracking top task per display.
+        private final int mDisplayId = DEFAULT_DISPLAY;
         // Only used when enableShellTopTaskTracking() is disabled
         @Nullable
         private final TaskInfo mTopTask;
@@ -386,20 +397,22 @@
         @Nullable
         private final GroupedTaskInfo mVisibleTasks;
 
+        private final boolean mCanEnterDesktopMode;
 
         // Only used when enableShellTopTaskTracking() is enabled
-        CachedTaskInfo(@Nullable GroupedTaskInfo visibleTasks) {
+        CachedTaskInfo(@Nullable GroupedTaskInfo visibleTasks, boolean canEnterDesktopMode) {
             mAllCachedTasks = null;
             mTopTask = null;
             mVisibleTasks = visibleTasks;
-
+            mCanEnterDesktopMode = canEnterDesktopMode;
         }
 
         // Only used when enableShellTopTaskTracking() is disabled
-        CachedTaskInfo(@NonNull List<TaskInfo> allCachedTasks) {
+        CachedTaskInfo(@NonNull List<TaskInfo> allCachedTasks, boolean canEnterDesktopMode) {
             mVisibleTasks = null;
             mAllCachedTasks = allCachedTasks;
             mTopTask = allCachedTasks.isEmpty() ? null : allCachedTasks.get(0);
+            mCanEnterDesktopMode = canEnterDesktopMode;
         }
 
         /**
@@ -509,58 +522,73 @@
                             && t.getActivityType() != ACTIVITY_TYPE_RECENTS)
                     .toList();
             return visibleNonExcludedTasks.isEmpty() ? null
-                    : new CachedTaskInfo(visibleNonExcludedTasks);
+                    : new CachedTaskInfo(visibleNonExcludedTasks, mCanEnterDesktopMode);
         }
 
         /**
-         * Returns {@link Task} array which can be used as a placeholder until the true object
-         * is loaded by the model
+         * Returns {@link TaskInfo} array corresponding to the provided task ids which can be
+         * used as a placeholder until the true object is loaded by the model. Only used when
+         * enableShellTopTaskTracking() is disabled.
          */
-        public Task[] getPlaceholderTasks() {
-            final TaskInfo baseTask = getLegacyBaseTask();
-            // TODO(346588978): Update this to return more than a single task once the callers
-            //  are refactored
-            return baseTask == null
-                    ? new Task[0]
-                    : new Task[]{Task.from(new TaskKey(baseTask), baseTask, false)};
-        }
-
-        /**
-         * Returns {@link Task} array corresponding to the provided task ids which can be used as a
-         * placeholder until the true object is loaded by the model
-         */
-        public Task[] getSplitPlaceholderTasks(int[] taskIds) {
-            if (enableShellTopTaskTracking()) {
-                if (mVisibleTasks == null || !mVisibleTasks.isBaseType(TYPE_SPLIT)) {
-                    return new Task[0];
-                }
-
-                GroupedTaskInfo splitTask = mVisibleTasks.getBaseGroupedTask();
-                Task[] result = new Task[taskIds.length];
-                for (int i = 0; i < taskIds.length; i++) {
-                    TaskInfo info = splitTask.getTaskById(taskIds[i]);
-                    if (info == null) {
-                        Log.w(TAG, "Requested task (" + taskIds[i] + ") not found");
-                        return new Task[0];
+        private TaskInfo[] getSplitPlaceholderTasksInfo(int[] splitTaskIds) {
+            if (mTopTask == null) {
+                return new TaskInfo[0];
+            }
+            TaskInfo[] result = new TaskInfo[splitTaskIds.length];
+            for (int i = 0; i < splitTaskIds.length; i++) {
+                final int index = i;
+                int taskId = splitTaskIds[i];
+                mAllCachedTasks.forEach(rti -> {
+                    if (rti.taskId == taskId) {
+                        result[index] = rti;
                     }
-                    result[i] = Task.from(new TaskKey(info), info, false);
+                });
+            }
+            return result;
+        }
+
+        private boolean isDesktopTask(TaskInfo taskInfo) {
+            return mCanEnterDesktopMode
+                    && taskInfo.configuration.windowConfiguration.getWindowingMode()
+                    == WindowConfiguration.WINDOWING_MODE_FREEFORM;
+        }
+
+        // TODO(346588978): Update this to return more than a single task once the callers
+        //  are refactored.
+
+        /**
+         * Returns a {@link GroupedTaskInfo} which can be used as a placeholder until the true
+         * object is loaded by the model.
+         *
+         * @param splitTaskIds provide if it is for split, which represents the task ids of the
+         *                     paired tasks. Otherwise, provide null.
+         */
+        public GroupedTaskInfo getPlaceholderGroupedTaskInfo(@Nullable int[] splitTaskIds) {
+            if (enableShellTopTaskTracking()) {
+                if (mVisibleTasks == null) {
+                    return null;
                 }
-                return result;
+                return mVisibleTasks.getBaseGroupedTask();
             } else {
-                if (mTopTask == null) {
-                    return new Task[0];
+                final TaskInfo baseTaskInfo = getLegacyBaseTask();
+                if (baseTaskInfo == null) {
+                    return null;
                 }
-                Task[] result = new Task[taskIds.length];
-                for (int i = 0; i < taskIds.length; i++) {
-                    final int index = i;
-                    int taskId = taskIds[i];
-                    mAllCachedTasks.forEach(rti -> {
-                        if (rti.taskId == taskId) {
-                            result[index] = Task.from(new TaskKey(rti), rti, false);
-                        }
-                    });
+                if (splitTaskIds != null && splitTaskIds.length >= 2) {
+                    TaskInfo[] splitTasksInfo = getSplitPlaceholderTasksInfo(splitTaskIds);
+                    if (splitTasksInfo[0] == null || splitTasksInfo[1] == null) {
+                        return null;
+                    }
+                    return GroupedTaskInfo.forSplitTasks(splitTasksInfo[0],
+                            splitTasksInfo[1], /* splitBounds = */ null);
+                } else if (isDesktopTask(baseTaskInfo)) {
+                    return GroupedTaskInfo.forDeskTasks(INACTIVE_DESK_ID, mDisplayId,
+                            Collections.singletonList(
+                                    baseTaskInfo), /* minimizedFreeformTaskIds = */
+                            Collections.emptySet());
+                } else {
+                    return GroupedTaskInfo.forFullscreenTasks(baseTaskInfo);
                 }
-                return result;
             }
         }
 
diff --git a/quickstep/src/com/android/quickstep/fallback/FallbackRecentsView.java b/quickstep/src/com/android/quickstep/fallback/FallbackRecentsView.java
index 695c77c..dc1cdde 100644
--- a/quickstep/src/com/android/quickstep/fallback/FallbackRecentsView.java
+++ b/quickstep/src/com/android/quickstep/fallback/FallbackRecentsView.java
@@ -56,6 +56,7 @@
 import com.android.quickstep.views.TaskContainer;
 import com.android.quickstep.views.TaskView;
 import com.android.systemui.shared.recents.model.Task;
+import com.android.wm.shell.shared.GroupedTaskInfo;
 
 import java.util.ArrayList;
 import java.util.Arrays;
@@ -116,11 +117,13 @@
      * to the home task. This allows us to handle quick-switch similarly to a quick-switching
      * from a foreground task.
      */
-    public void onGestureAnimationStartOnHome(Task[] homeTask) {
+    public void onGestureAnimationStartOnHome(GroupedTaskInfo homeTaskInfo) {
         // TODO(b/195607777) General fallback love, but this might be correct
         //  Home task should be defined as the front-most task info I think?
-        mHomeTask = homeTask.length > 0 ? homeTask[0] : null;
-        onGestureAnimationStart(homeTask);
+        if (homeTaskInfo != null) {
+            mHomeTask = Task.from(homeTaskInfo.getTaskInfo1());
+        }
+        onGestureAnimationStart(homeTaskInfo);
     }
 
     /**
@@ -175,13 +178,13 @@
     }
 
     @Override
-    protected boolean shouldAddStubTaskView(Task[] runningTasks) {
-        if (runningTasks.length > 1) {
+    protected boolean shouldAddStubTaskView(GroupedTaskInfo groupedTaskInfo) {
+        if (!groupedTaskInfo.isBaseType(GroupedTaskInfo.TYPE_FULLSCREEN)) {
             // can't be in split screen w/ home task
-            return super.shouldAddStubTaskView(runningTasks);
+            return super.shouldAddStubTaskView(groupedTaskInfo);
         }
 
-        Task runningTask = runningTasks[0];
+        Task runningTask = Task.from(groupedTaskInfo.getTaskInfo1());
         if (mHomeTask != null && runningTask != null
                 && mHomeTask.key.id == runningTask.key.id
                 && !hasTaskViews() && mLoadPlanEverApplied) {
@@ -190,7 +193,7 @@
             // Ignore empty task signal if applyLoadPlan has never run.
             return false;
         }
-        return super.shouldAddStubTaskView(runningTasks);
+        return super.shouldAddStubTaskView(groupedTaskInfo);
     }
 
     @Override
diff --git a/quickstep/src/com/android/quickstep/fallback/window/RecentsWindowSwipeHandler.java b/quickstep/src/com/android/quickstep/fallback/window/RecentsWindowSwipeHandler.java
index b365ddf..0d139b4 100644
--- a/quickstep/src/com/android/quickstep/fallback/window/RecentsWindowSwipeHandler.java
+++ b/quickstep/src/com/android/quickstep/fallback/window/RecentsWindowSwipeHandler.java
@@ -257,7 +257,8 @@
         if (mRunningOverHome) {
             if (DisplayController.getNavigationMode(mContext).hasGestures) {
                 mRecentsView.onGestureAnimationStartOnHome(
-                        mGestureState.getRunningTask().getPlaceholderTasks());
+                        mGestureState.getRunningTask().getPlaceholderGroupedTaskInfo(
+                                /* splitTaskIds = */ null));
             }
         } else {
             super.notifyGestureAnimationStartToRecents();
@@ -283,7 +284,7 @@
         private final AnimatedFloat mHomeAlpha = new AnimatedFloat(this::updateAppTransforms);
         private final AnimatedFloat mVerticalShiftForScale =
                 new AnimatedFloat(this::updateAppTransforms);
-        private final AnimatedFloat mRecentsAlpha = new AnimatedFloat(this:: updateAppTransforms);
+        private final AnimatedFloat mRecentsAlpha = new AnimatedFloat(this::updateAppTransforms);
 
         private final RectF mTargetRect = new RectF();
         private SurfaceControl mSurfaceControl;
diff --git a/quickstep/src/com/android/quickstep/views/LauncherRecentsView.java b/quickstep/src/com/android/quickstep/views/LauncherRecentsView.java
index c1282b9..b86d0c8 100644
--- a/quickstep/src/com/android/quickstep/views/LauncherRecentsView.java
+++ b/quickstep/src/com/android/quickstep/views/LauncherRecentsView.java
@@ -54,7 +54,7 @@
 import com.android.quickstep.SystemUiProxy;
 import com.android.quickstep.util.AnimUtils;
 import com.android.quickstep.util.SplitSelectStateController;
-import com.android.systemui.shared.recents.model.Task;
+import com.android.wm.shell.shared.GroupedTaskInfo;
 
 import kotlin.Unit;
 
@@ -275,8 +275,8 @@
     }
 
     @Override
-    public void onGestureAnimationStart(Task[] runningTasks) {
-        super.onGestureAnimationStart(runningTasks);
+    public void onGestureAnimationStart(GroupedTaskInfo groupedTaskInfo) {
+        super.onGestureAnimationStart(groupedTaskInfo);
         if (!ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY.isTrue()) {
             // TODO: b/333533253 - Remove after flag rollout
             DesktopVisibilityController.INSTANCE.get(mContainer).setRecentsGestureStart();
diff --git a/quickstep/src/com/android/quickstep/views/RecentsView.java b/quickstep/src/com/android/quickstep/views/RecentsView.java
index 8f7dc80..ea95206 100644
--- a/quickstep/src/com/android/quickstep/views/RecentsView.java
+++ b/quickstep/src/com/android/quickstep/views/RecentsView.java
@@ -87,7 +87,6 @@
 import android.animation.PropertyValuesHolder;
 import android.animation.ValueAnimator;
 import android.annotation.SuppressLint;
-import android.app.WindowConfiguration;
 import android.content.Context;
 import android.content.Intent;
 import android.content.LocusId;
@@ -234,6 +233,7 @@
 import com.android.quickstep.util.VibrationConstants;
 import com.android.systemui.plugins.ResourceProvider;
 import com.android.systemui.shared.recents.model.Task;
+import com.android.systemui.shared.recents.model.Task.TaskKey;
 import com.android.systemui.shared.recents.model.ThumbnailData;
 import com.android.systemui.shared.system.ActivityManagerWrapper;
 import com.android.systemui.shared.system.InteractionJankMonitorWrapper;
@@ -241,6 +241,7 @@
 import com.android.systemui.shared.system.TaskStackChangeListener;
 import com.android.systemui.shared.system.TaskStackChangeListeners;
 import com.android.wm.shell.common.pip.IPipAnimationListener;
+import com.android.wm.shell.shared.GroupedTaskInfo;
 import com.android.wm.shell.shared.desktopmode.DesktopModeStatus;
 import com.android.wm.shell.shared.desktopmode.DesktopModeTransitionSource;
 
@@ -567,10 +568,10 @@
     private DesktopVisibilityController mDesktopVisibilityController = null;
 
     /**
-     * Reflects if Recents is currently in the middle of a gesture, and if so, which tasks are
-     * running. If a gesture is not in progress, this will be null.
+     * Reflects if Recents is currently in the middle of a gesture, and if so, which related
+     * [GroupedTaskInfo] is running. If a gesture is not in progress, this will be null.
      */
-    private @Nullable Task[] mActiveGestureRunningTasks;
+    private @Nullable GroupedTaskInfo mActiveGestureGroupedTaskInfo;
 
     // Keeps track of the previously known visible tasks for purposes of loading/unloading task data
     private final SparseBooleanArray mHasVisibleTaskData = new SparseBooleanArray();
@@ -582,7 +583,7 @@
     private final ViewPool<GroupedTaskView> mGroupedTaskViewPool;
     private final ViewPool<DesktopTaskView> mDesktopTaskViewPool;
 
-    private final TaskOverlayFactory mTaskOverlayFactory;
+    protected final TaskOverlayFactory mTaskOverlayFactory;
 
     protected boolean mDisallowScrollToClearAll;
     // True if it is not allowed to scroll to [AddDesktopButton].
@@ -671,7 +672,7 @@
                 return;
             }
             Log.d(TAG, "onTaskRemoved: " + taskId);
-            Task.TaskKey taskKey = taskContainer.getTask().key;
+            TaskKey taskKey = taskContainer.getTask().key;
             UI_HELPER_EXECUTOR.execute(new CancellableTask<>(
                     () -> PackageManagerWrapper.getInstance()
                             .getActivityInfo(taskKey.getComponent(), taskKey.userId) == null,
@@ -2100,12 +2101,12 @@
         if (newRunningTaskView != null) {
             setRunningTaskViewId(newRunningTaskView.getTaskViewId());
         } else {
-            if (mActiveGestureRunningTasks != null) {
+            if (mActiveGestureGroupedTaskInfo != null) {
                 // This will update mRunningTaskViewId and create a stub view if necessary.
                 // We try to avoid this because it can cause a scroll jump, but it is needed
                 // for cases where the running task isn't included in this load plan (e.g. if
                 // the current running task is excludedFromRecents.)
-                showCurrentTask(mActiveGestureRunningTasks, "applyLoadPlan");
+                showCurrentTask(mActiveGestureGroupedTaskInfo, "applyLoadPlan");
                 newRunningTaskView = getRunningTaskView();
             } else {
                 setRunningTaskViewId(INVALID_TASK_ID);
@@ -2826,7 +2827,7 @@
      * Handle the edge case where Recents could increment task count very high over long
      * period of device usage. Probably will never happen, but meh.
      */
-    private TaskView getTaskViewFromPool(TaskViewType type) {
+    protected TaskView getTaskViewFromPool(TaskViewType type) {
         TaskView taskView;
         switch (type) {
             case GROUPED:
@@ -2880,10 +2881,9 @@
      */
     // TODO: b/401582344 - Implement a way to exclude the `DesktopWallpaperActivity` from being
     //  considered in Overview.
-    public void onGestureAnimationStart(Task[] runningTasks) {
-        Log.d(TAG, "onGestureAnimationStart - runningTasks: " + Arrays.toString(runningTasks));
-        mActiveGestureRunningTasks = runningTasks;
-
+    public void onGestureAnimationStart(GroupedTaskInfo groupedTaskInfo) {
+        Log.d(TAG, "onGestureAnimationStart - groupedTaskInfo: " + groupedTaskInfo);
+        mActiveGestureGroupedTaskInfo = groupedTaskInfo;
 
         // This needs to be called before the other states are set since it can create the task view
         if (mOrientationState.setGestureActive(true)) {
@@ -2893,7 +2893,7 @@
             updateSizeAndPadding();
         }
 
-        showCurrentTask(mActiveGestureRunningTasks, "onGestureAnimationStart");
+        showCurrentTask(groupedTaskInfo, "onGestureAnimationStart");
         setEnableFreeScroll(false);
         setEnableDrawingLiveTile(false);
         setRunningTaskHidden(true);
@@ -2908,7 +2908,7 @@
     }
 
     private boolean isGestureActive() {
-        return mActiveGestureRunningTasks != null;
+        return mActiveGestureGroupedTaskInfo != null;
     }
 
     /**
@@ -3046,7 +3046,7 @@
      * Called when a gesture from an app has finished, and the animation to the target has ended.
      */
     public void onGestureAnimationEnd() {
-        mActiveGestureRunningTasks = null;
+        mActiveGestureGroupedTaskInfo = null;
         if (mOrientationState.setGestureActive(false)) {
             updateOrientationHandler(/* forceRecreateDragLayerControllers = */ false);
         }
@@ -3085,10 +3085,17 @@
     /**
      * Returns true if we should add a stub taskView for the running task id
      */
-    protected boolean shouldAddStubTaskView(Task[] runningTasks) {
-        int[] runningTaskIds = Arrays.stream(runningTasks).mapToInt(task -> task.key.id).toArray();
+    protected boolean shouldAddStubTaskView(GroupedTaskInfo groupedTaskInfo) {
+        int[] runningTaskIds;
+        if (groupedTaskInfo != null) {
+            runningTaskIds = groupedTaskInfo.getTaskInfoList().stream().mapToInt(
+                    taskInfo -> taskInfo.taskId).toArray();
+        } else {
+            runningTaskIds = new int[0];
+        }
         TaskView matchingTaskView = null;
-        if (hasDesktopTask(runningTasks) && runningTaskIds.length == 1) {
+        if (groupedTaskInfo != null && groupedTaskInfo.isBaseType(GroupedTaskInfo.TYPE_DESK)
+                && runningTaskIds.length == 1) {
             // TODO(b/342635213): Unsure if it's expected, desktop runningTasks only have a single
             // taskId, therefore we match any DesktopTaskView that contains the runningTaskId.
             TaskView taskview = getTaskViewByTaskId(runningTaskIds[0]);
@@ -3102,53 +3109,36 @@
     }
 
     /**
-     * Creates a `DesktopTaskView` for the currently active desk on this display, which contains the
-     * gievn `runningTasks`.
-     */
-    private DesktopTaskView createDesktopTaskViewForActiveDesk(Task[] runningTasks) {
-        final int activeDeskId = mUtils.getActiveDeskIdOnThisDisplay();
-        final var desktopTaskView = (DesktopTaskView) getTaskViewFromPool(TaskViewType.DESKTOP);
-
-        // TODO: b/401582344 - Implement a way to exclude the `DesktopWallpaperActivity`.
-        desktopTaskView.bind(
-                new DesktopTask(activeDeskId, mContainer.getDisplayId(),
-                        Arrays.asList(runningTasks)),
-                mOrientationState, mTaskOverlayFactory);
-        return desktopTaskView;
-    }
-
-    /**
-     * Creates a task view (if necessary) to represent the task with the {@param runningTaskId}.
+     * Creates a task view (if necessary) to represent the tasks with the {@param groupedTaskInfo}.
      *
      * All subsequent calls to reload will keep the task as the first item until {@link #reset()}
      * is called.  Also scrolls the view to this task.
      */
-    private void showCurrentTask(Task[] runningTasks, String caller) {
-        Log.d(TAG, "showCurrentTask(" + caller + ") - runningTasks: "
-                + Arrays.toString(runningTasks));
-        if (runningTasks.length == 0) {
+    private void showCurrentTask(GroupedTaskInfo groupedTaskInfo, String caller) {
+        Log.d(TAG, "showCurrentTask(" + caller + ") - groupedTaskInfo: " + groupedTaskInfo);
+        if (groupedTaskInfo == null) {
             return;
         }
 
         int runningTaskViewId = -1;
-        if (shouldAddStubTaskView(runningTasks)) {
+        if (shouldAddStubTaskView(groupedTaskInfo)) {
             boolean wasEmpty = getChildCount() == 0;
             // Add an empty view for now until the task plan is loaded and applied
             final TaskView taskView;
-            final boolean needGroupTaskView = runningTasks.length > 1;
-            final boolean needDesktopTask = hasDesktopTask(runningTasks);
-            if (needDesktopTask) {
-                taskView = createDesktopTaskViewForActiveDesk(runningTasks);
-            } else if (needGroupTaskView) {
+            if (groupedTaskInfo.isBaseType(GroupedTaskInfo.TYPE_DESK)) {
+                taskView = mUtils.createDesktopTaskViewForActiveDesk(groupedTaskInfo);
+            } else if (groupedTaskInfo.isBaseType(GroupedTaskInfo.TYPE_SPLIT)) {
                 taskView = getTaskViewFromPool(TaskViewType.GROUPED);
                 // When we create a placeholder task view mSplitBoundsConfig will be null, but with
                 // the actual app running we won't need to show the thumbnail until all the tasks
                 // load later anyways
-                ((GroupedTaskView) taskView).bind(runningTasks[0], runningTasks[1],
-                        mOrientationState, mTaskOverlayFactory, mSplitBoundsConfig);
+                ((GroupedTaskView) taskView).bind(Task.from(groupedTaskInfo.getTaskInfo1()),
+                        Task.from(groupedTaskInfo.getTaskInfo2()), mOrientationState,
+                        mTaskOverlayFactory, mSplitBoundsConfig);
             } else {
                 taskView = getTaskViewFromPool(TaskViewType.SINGLE);
-                taskView.bind(runningTasks[0], mOrientationState, mTaskOverlayFactory);
+                taskView.bind(Task.from(groupedTaskInfo.getTaskInfo1()), mOrientationState,
+                        mTaskOverlayFactory);
             }
             if (mAddDesktopButton != null && wasEmpty) {
                 addView(mAddDesktopButton);
@@ -3165,7 +3155,7 @@
                     makeMeasureSpec(getMeasuredHeight(), EXACTLY));
             layout(getLeft(), getTop(), getRight(), getBottom());
         } else {
-            var runningTaskView = getTaskViewByTaskId(runningTasks[0].key.id);
+            var runningTaskView = getTaskViewByTaskId(groupedTaskInfo.getTaskInfo1().taskId);
             if (runningTaskView != null) {
                 runningTaskViewId = runningTaskView.getTaskViewId();
             }
@@ -3198,22 +3188,6 @@
         reloadIfNeeded();
     }
 
-    private boolean hasDesktopTask(Task[] runningTasks) {
-        if (!DesktopModeStatus.canEnterDesktopMode(mContext)) {
-            return false;
-        }
-        for (Task task : runningTasks) {
-            if (task.key.windowingMode == WindowConfiguration.WINDOWING_MODE_FREEFORM) {
-                return true;
-            }
-        }
-
-        // A running empty desk will have a single running app for the `DesktopWallpaperActivity`.
-        // TODO: b/401582344 - Implement a way to exclude the `DesktopWallpaperActivity`.
-
-        return false;
-    }
-
     /**
      * Sets the running task id, cleaning up the old running task if necessary.
      */
diff --git a/quickstep/src/com/android/quickstep/views/RecentsViewUtils.kt b/quickstep/src/com/android/quickstep/views/RecentsViewUtils.kt
index 50941fe..b265b13 100644
--- a/quickstep/src/com/android/quickstep/views/RecentsViewUtils.kt
+++ b/quickstep/src/com/android/quickstep/views/RecentsViewUtils.kt
@@ -39,7 +39,9 @@
 import com.android.quickstep.util.GroupTask
 import com.android.quickstep.util.isExternalDisplay
 import com.android.quickstep.views.RecentsView.RUNNING_TASK_ATTACH_ALPHA
+import com.android.systemui.shared.recents.model.Task
 import com.android.systemui.shared.recents.model.ThumbnailData
+import com.android.wm.shell.shared.GroupedTaskInfo
 import java.util.function.BiConsumer
 import kotlin.math.min
 import kotlin.reflect.KMutableProperty1
@@ -459,6 +461,22 @@
         }
     }
 
+    /**
+     * Creates a [DesktopTaskView] for the currently active desk on this display, which contains the
+     * tasks with the given [groupedTaskInfo].
+     */
+    fun createDesktopTaskViewForActiveDesk(groupedTaskInfo: GroupedTaskInfo): DesktopTaskView {
+        val desktopTaskView =
+            recentsView.getTaskViewFromPool(TaskViewType.DESKTOP) as DesktopTaskView
+        val tasks: List<Task> = groupedTaskInfo.taskInfoList.map { taskInfo -> Task.from(taskInfo) }
+        desktopTaskView.bind(
+            DesktopTask(groupedTaskInfo.deskId, groupedTaskInfo.deskDisplayId, tasks),
+            recentsView.mOrientationState,
+            recentsView.mTaskOverlayFactory,
+        )
+        return desktopTaskView
+    }
+
     companion object {
         class RecentsViewFloatProperty(
             private val utilsProperty: KMutableProperty1<RecentsViewUtils, Float>
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/AbsSwipeUpHandlerTestCase.java b/quickstep/tests/multivalentTests/src/com/android/quickstep/AbsSwipeUpHandlerTestCase.java
index 2f5f160..568da2d 100644
--- a/quickstep/tests/multivalentTests/src/com/android/quickstep/AbsSwipeUpHandlerTestCase.java
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/AbsSwipeUpHandlerTestCase.java
@@ -101,7 +101,8 @@
     protected final ActivityManager.RunningTaskInfo mRunningTaskInfo =
             new ActivityManager.RunningTaskInfo();
     protected final TopTaskTracker.CachedTaskInfo mCachedTaskInfo =
-            new TopTaskTracker.CachedTaskInfo(Collections.singletonList(mRunningTaskInfo));
+            new TopTaskTracker.CachedTaskInfo(
+                    Collections.singletonList(mRunningTaskInfo), /* canEnterDesktopMode = */ false);
     protected final RemoteAnimationTarget mRemoteAnimationTarget = new RemoteAnimationTarget(
             /* taskId= */ 0,
             /* mode= */ RemoteAnimationTarget.MODE_CLOSING,
diff --git a/quickstep/tests/src/com/android/launcher3/statehandlers/DepthControllerTest.kt b/quickstep/tests/src/com/android/launcher3/statehandlers/DepthControllerTest.kt
new file mode 100644
index 0000000..17cca0b
--- /dev/null
+++ b/quickstep/tests/src/com/android/launcher3/statehandlers/DepthControllerTest.kt
@@ -0,0 +1,105 @@
+/*
+ * 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.statehandlers
+
+import android.content.res.Resources
+import android.view.ViewTreeObserver
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import com.android.launcher3.Launcher
+import com.android.launcher3.R
+import com.android.launcher3.dragndrop.DragLayer
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.Mock
+import org.mockito.Mockito.reset
+import org.mockito.Mockito.same
+import org.mockito.Mockito.verify
+import org.mockito.Mockito.verifyNoMoreInteractions
+import org.mockito.Mockito.`when`
+import org.mockito.MockitoAnnotations
+
+@SmallTest
+@RunWith(AndroidJUnit4::class)
+class DepthControllerTest {
+
+    private lateinit var underTest: DepthController
+    @Mock private lateinit var launcher: Launcher
+    @Mock private lateinit var resource: Resources
+    @Mock private lateinit var dragLayer: DragLayer
+    @Mock private lateinit var viewTreeObserver: ViewTreeObserver
+
+    @Before
+    fun setUp() {
+        MockitoAnnotations.initMocks(this)
+        `when`(launcher.resources).thenReturn(resource)
+        `when`(resource.getInteger(R.integer.max_depth_blur_radius)).thenReturn(30)
+        `when`(launcher.dragLayer).thenReturn(dragLayer)
+        `when`(dragLayer.viewTreeObserver).thenReturn(viewTreeObserver)
+
+        underTest = DepthController(launcher)
+    }
+
+    @Test
+    fun setActivityStarted_add_onDrawListener() {
+        underTest.setActivityStarted(true)
+
+        verify(viewTreeObserver).addOnDrawListener(same(underTest.mOnDrawListener))
+    }
+
+    @Test
+    fun setActivityStopped_not_remove_onDrawListener() {
+        underTest.setActivityStarted(false)
+
+        // Because underTest.mOnDrawListener is never added
+        verifyNoMoreInteractions(viewTreeObserver)
+    }
+
+    @Test
+    fun setActivityStared_then_stopped_remove_onDrawListener() {
+        underTest.setActivityStarted(true)
+        reset(viewTreeObserver)
+
+        underTest.setActivityStarted(false)
+
+        verify(viewTreeObserver).removeOnDrawListener(same(underTest.mOnDrawListener))
+    }
+
+    @Test
+    fun setActivityStared_then_stopped_multiple_times_remove_onDrawListener_once() {
+        underTest.setActivityStarted(true)
+        reset(viewTreeObserver)
+
+        underTest.setActivityStarted(false)
+        underTest.setActivityStarted(false)
+        underTest.setActivityStarted(false)
+
+        // Should just remove mOnDrawListener once
+        verify(viewTreeObserver).removeOnDrawListener(same(underTest.mOnDrawListener))
+    }
+
+    @Test
+    fun test_onInvalidSurface_multiple_times_add_onDrawListener_once() {
+        underTest.onInvalidSurface()
+        underTest.onInvalidSurface()
+        underTest.onInvalidSurface()
+
+        // We should only call addOnDrawListener 1 time
+        verify(viewTreeObserver).addOnDrawListener(same(underTest.mOnDrawListener))
+    }
+}
diff --git a/src/com/android/launcher3/graphics/GridCustomizationsProxy.java b/src/com/android/launcher3/graphics/GridCustomizationsProxy.java
index b40e099..7bd9a28 100644
--- a/src/com/android/launcher3/graphics/GridCustomizationsProxy.java
+++ b/src/com/android/launcher3/graphics/GridCustomizationsProxy.java
@@ -27,6 +27,7 @@
 import android.database.Cursor;
 import android.database.MatrixCursor;
 import android.net.Uri;
+import android.os.Binder;
 import android.os.Bundle;
 import android.os.Handler;
 import android.os.IBinder.DeathRecipient;
@@ -313,7 +314,7 @@
         RunnableList lifeCycleTracker = new RunnableList();
         try {
             PreviewSurfaceRenderer renderer = new PreviewSurfaceRenderer(
-                    mContext, lifeCycleTracker, request);
+                    mContext, lifeCycleTracker, request, Binder.getCallingPid());
             PreviewLifecycleObserver observer =
                     new PreviewLifecycleObserver(lifeCycleTracker, renderer);
 
diff --git a/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java b/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java
index e4e4b90..a062a54 100644
--- a/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java
+++ b/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java
@@ -23,6 +23,7 @@
 import static com.android.launcher3.BubbleTextView.DISPLAY_TASKBAR;
 import static com.android.launcher3.BubbleTextView.DISPLAY_WORKSPACE;
 import static com.android.launcher3.DeviceProfile.DEFAULT_SCALE;
+import static com.android.launcher3.Flags.extendibleThemeManager;
 import static com.android.launcher3.Hotseat.ALPHA_CHANNEL_PREVIEW_RENDERER;
 import static com.android.launcher3.LauncherPrefs.GRID_NAME;
 import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_HOTSEAT;
@@ -30,6 +31,7 @@
 import static com.android.launcher3.Utilities.SHOULD_SHOW_FIRST_PAGE_WIDGET;
 import static com.android.launcher3.graphics.ThemeManager.PREF_ICON_SHAPE;
 import static com.android.launcher3.model.ModelUtils.currentScreenContentFilter;
+import static com.android.launcher3.widget.LauncherWidgetHolder.APPWIDGET_HOST_ID;
 
 import android.app.Fragment;
 import android.app.WallpaperColors;
@@ -44,6 +46,7 @@
 import android.graphics.Rect;
 import android.os.Handler;
 import android.os.Looper;
+import android.text.TextUtils;
 import android.util.AttributeSet;
 import android.util.Size;
 import android.util.SparseArray;
@@ -76,13 +79,20 @@
 import com.android.launcher3.celllayout.CellLayoutLayoutParams;
 import com.android.launcher3.celllayout.CellPosMapper;
 import com.android.launcher3.config.FeatureFlags;
+import com.android.launcher3.dagger.ApiWrapperModule;
+import com.android.launcher3.dagger.AppModule;
 import com.android.launcher3.dagger.LauncherAppComponent;
-import com.android.launcher3.dagger.LauncherAppModule;
 import com.android.launcher3.dagger.LauncherAppSingleton;
+import com.android.launcher3.dagger.LauncherComponentProvider;
+import com.android.launcher3.dagger.PluginManagerWrapperModule;
+import com.android.launcher3.dagger.StaticObjectModule;
+import com.android.launcher3.dagger.WindowManagerProxyModule;
 import com.android.launcher3.folder.FolderIcon;
 import com.android.launcher3.model.BaseLauncherBinder.BaseLauncherBinderFactory;
 import com.android.launcher3.model.BgDataModel;
 import com.android.launcher3.model.BgDataModel.FixedContainerItems;
+import com.android.launcher3.model.LayoutParserFactory;
+import com.android.launcher3.model.LayoutParserFactory.XmlLayoutParserFactory;
 import com.android.launcher3.model.LoaderTask.LoaderTaskFactory;
 import com.android.launcher3.model.data.AppPairInfo;
 import com.android.launcher3.model.data.CollectionInfo;
@@ -104,6 +114,7 @@
 import com.android.launcher3.widget.BaseLauncherAppWidgetHostView;
 import com.android.launcher3.widget.LauncherAppWidgetProviderInfo;
 import com.android.launcher3.widget.LauncherWidgetHolder;
+import com.android.launcher3.widget.LauncherWidgetHolder.WidgetHolderFactory;
 import com.android.launcher3.widget.LocalColorExtractor;
 import com.android.launcher3.widget.util.WidgetSizes;
 import com.android.systemui.shared.Flags;
@@ -111,6 +122,8 @@
 import dagger.BindsInstance;
 import dagger.Component;
 
+import java.io.File;
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
@@ -139,21 +152,63 @@
 
         private final String mPrefName;
 
+        private final File mDbDir;
+
         public PreviewContext(Context base, String gridName, String shapeKey) {
+            this(base, gridName, shapeKey, APPWIDGET_HOST_ID, null);
+        }
+
+        public PreviewContext(Context base, String gridName, String shapeKey,
+                int widgetHostId, @Nullable String layoutXml) {
             super(base);
-            mPrefName = "preview-" + UUID.randomUUID().toString();
+            String randomUid = UUID.randomUUID().toString();
+            mPrefName = "preview-" + randomUid;
             LauncherPrefs prefs =
                     new ProxyPrefs(this, getSharedPreferences(mPrefName, MODE_PRIVATE));
             prefs.put(GRID_NAME, gridName);
             prefs.put(PREF_ICON_SHAPE, shapeKey);
-            initDaggerComponent(
-                    DaggerLauncherPreviewRenderer_PreviewAppComponent.builder().bindPrefs(prefs));
+
+            PreviewAppComponent.Builder builder =
+                    DaggerLauncherPreviewRenderer_PreviewAppComponent.builder().bindPrefs(prefs);
+            if (TextUtils.isEmpty(layoutXml) || !extendibleThemeManager()) {
+                mDbDir = null;
+                builder.bindParserFactory(new LayoutParserFactory(this))
+                        .bindWidgetsFactory(
+                                LauncherComponentProvider.get(base).getWidgetHolderFactory());
+            } else {
+                mDbDir = new File(base.getFilesDir(), randomUid);
+                emptyDbDir();
+                mDbDir.mkdirs();
+                builder.bindParserFactory(new XmlLayoutParserFactory(this, layoutXml))
+                        .bindWidgetsFactory(c -> new LauncherWidgetHolder(c, widgetHostId));
+            }
+            initDaggerComponent(builder);
+
+            if (!TextUtils.isEmpty(layoutXml)) {
+                // Use null the DB file so that we use a new in-memory DB
+                InvariantDeviceProfile.INSTANCE.get(this).dbFile = null;
+            }
+        }
+
+        private void emptyDbDir() {
+            if (mDbDir != null && mDbDir.exists()) {
+                Arrays.stream(mDbDir.listFiles()).forEach(File::delete);
+            }
         }
 
         @Override
         protected void cleanUpObjects() {
             super.cleanUpObjects();
             deleteSharedPreferences(mPrefName);
+            if (mDbDir != null) {
+                emptyDbDir();
+                mDbDir.delete();
+            }
+        }
+
+        @Override
+        public File getDatabasePath(String name) {
+            return mDbDir != null ? new File(mDbDir, name) :  super.getDatabasePath(name);
         }
     }
 
@@ -173,13 +228,15 @@
 
     public LauncherPreviewRenderer(Context context,
             InvariantDeviceProfile idp,
+            int widgetHostId,
             WallpaperColors wallpaperColorsOverride,
             @Nullable final SparseArray<Size> launcherWidgetSpanInfo) {
-        this(context, idp, null, wallpaperColorsOverride, launcherWidgetSpanInfo);
+        this(context, idp, widgetHostId, null, wallpaperColorsOverride, launcherWidgetSpanInfo);
     }
 
     public LauncherPreviewRenderer(Context context,
             InvariantDeviceProfile idp,
+            int widgetHostId,
             SparseIntArray previewColorOverride,
             WallpaperColors wallpaperColorsOverride,
             @Nullable final SparseArray<Size> launcherWidgetSpanInfo) {
@@ -262,9 +319,13 @@
                     wallpaperColors)
                     : null;
         }
-        mAppWidgetHost = new LauncherPreviewAppWidgetHost(context);
+        mAppWidgetHost = new LauncherPreviewAppWidgetHost(context, widgetHostId);
 
         onViewCreated();
+
+        if (widgetHostId != APPWIDGET_HOST_ID) {
+            mAppWidgetHost.stopListening();
+        }
     }
 
     @Override
@@ -567,8 +628,8 @@
 
     private class LauncherPreviewAppWidgetHost extends AppWidgetHost {
 
-        private LauncherPreviewAppWidgetHost(Context context) {
-            super(context, LauncherWidgetHolder.APPWIDGET_HOST_ID);
+        private LauncherPreviewAppWidgetHost(Context context, int hostId) {
+            super(context, hostId);
         }
 
         @Override
@@ -604,7 +665,12 @@
     }
 
     @LauncherAppSingleton
-    @Component(modules = LauncherAppModule.class)
+    // Exclude widget module since we bind widget holder separately
+    @Component(modules = {WindowManagerProxyModule.class,
+            ApiWrapperModule.class,
+            PluginManagerWrapperModule.class,
+            StaticObjectModule.class,
+            AppModule.class})
     public interface PreviewAppComponent extends LauncherAppComponent {
 
         LoaderTaskFactory getLoaderTaskFactory();
@@ -615,6 +681,8 @@
         @Component.Builder
         interface Builder extends LauncherAppComponent.Builder {
             @BindsInstance Builder bindPrefs(LauncherPrefs prefs);
+            @BindsInstance Builder bindParserFactory(LayoutParserFactory parserFactory);
+            @BindsInstance Builder bindWidgetsFactory(WidgetHolderFactory holderFactory);
             PreviewAppComponent build();
         }
     }
diff --git a/src/com/android/launcher3/graphics/PreviewSurfaceRenderer.java b/src/com/android/launcher3/graphics/PreviewSurfaceRenderer.java
index 457d12e..93e5f32 100644
--- a/src/com/android/launcher3/graphics/PreviewSurfaceRenderer.java
+++ b/src/com/android/launcher3/graphics/PreviewSurfaceRenderer.java
@@ -20,12 +20,18 @@
 import static android.content.res.Configuration.UI_MODE_NIGHT_YES;
 import static android.view.Display.DEFAULT_DISPLAY;
 
+import static com.android.launcher3.Flags.extendibleThemeManager;
 import static com.android.launcher3.LauncherPrefs.GRID_NAME;
+import static com.android.launcher3.WorkspaceLayoutManager.FIRST_SCREEN_ID;
+import static com.android.launcher3.WorkspaceLayoutManager.SECOND_SCREEN_ID;
 import static com.android.launcher3.graphics.ThemeManager.PREF_ICON_SHAPE;
+import static com.android.launcher3.provider.LauncherDbUtils.selectionForWorkspaceScreen;
 import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
 import static com.android.launcher3.util.Executors.MODEL_EXECUTOR;
+import static com.android.launcher3.widget.LauncherWidgetHolder.APPWIDGET_HOST_ID;
 
 import android.app.WallpaperColors;
+import android.appwidget.AppWidgetHost;
 import android.appwidget.AppWidgetProviderInfo;
 import android.content.Context;
 import android.content.res.Configuration;
@@ -33,6 +39,7 @@
 import android.hardware.display.DisplayManager;
 import android.os.Bundle;
 import android.os.IBinder;
+import android.text.TextUtils;
 import android.util.Log;
 import android.util.Size;
 import android.util.SparseArray;
@@ -56,7 +63,6 @@
 import com.android.launcher3.LauncherAppState;
 import com.android.launcher3.LauncherPrefs;
 import com.android.launcher3.LauncherSettings;
-import com.android.launcher3.Workspace;
 import com.android.launcher3.dagger.LauncherComponentProvider;
 import com.android.launcher3.graphics.LauncherPreviewRenderer.PreviewAppComponent;
 import com.android.launcher3.graphics.LauncherPreviewRenderer.PreviewContext;
@@ -89,18 +95,22 @@
     private static final String KEY_COLOR_RESOURCE_IDS = "color_resource_ids";
     private static final String KEY_COLOR_VALUES = "color_values";
     private static final String KEY_DARK_MODE = "use_dark_mode";
+    private static final String KEY_LAYOUT_XML = "layout_xml";
     public static final String KEY_SKIP_ANIMATIONS = "skip_animations";
 
     private final Context mContext;
     private SparseIntArray mPreviewColorOverride;
     private String mGridName;
     private String mShapeKey;
+    private String mLayoutXml;
 
     @Nullable private Boolean mDarkMode;
     private boolean mDestroyed = false;
     private boolean mHideQsb;
     @Nullable private FrameLayout mViewRoot = null;
+    private boolean mDeletingHostOnExit = false;
 
+    private final int mCallingPid;
     private final IBinder mHostToken;
     private final int mWidth;
     private final int mHeight;
@@ -111,11 +121,11 @@
     private final RunnableList mLifeCycleTracker;
     private final SurfaceControlViewHost mSurfaceControlViewHost;
 
-
-    public PreviewSurfaceRenderer(
-            Context context, RunnableList lifecycleTracker, Bundle bundle) throws Exception {
+    public PreviewSurfaceRenderer(Context context, RunnableList lifecycleTracker, Bundle bundle,
+            int callingPid) throws Exception {
         mContext = context;
         mLifeCycleTracker = lifecycleTracker;
+        mCallingPid = callingPid;
         mGridName = bundle.getString("name");
         bundle.remove("name");
         if (mGridName == null) {
@@ -135,6 +145,7 @@
         mDisplayId = bundle.getInt(KEY_DISPLAY_ID);
         mDisplay = context.getSystemService(DisplayManager.class)
                 .getDisplay(mDisplayId);
+        mLayoutXml = bundle.getString(KEY_LAYOUT_XML);
         if (mDisplay == null) {
             throw new IllegalArgumentException("Display ID does not match any displays.");
         }
@@ -335,41 +346,52 @@
     private void loadModelData() {
         final Context inflationContext = getPreviewContext();
         if (!mGridName.equals(LauncherPrefs.INSTANCE.get(mContext).get(GRID_NAME))
-                || !mShapeKey.equals(LauncherPrefs.INSTANCE.get(mContext).get(PREF_ICON_SHAPE))) {
+                || !mShapeKey.equals(LauncherPrefs.INSTANCE.get(mContext).get(PREF_ICON_SHAPE))
+                || !TextUtils.isEmpty(mLayoutXml)) {
+
+            boolean isCustomLayout = extendibleThemeManager() &&  !TextUtils.isEmpty(mLayoutXml);
+            int widgetHostId = isCustomLayout ? APPWIDGET_HOST_ID + mCallingPid : APPWIDGET_HOST_ID;
+
             // Start the migration
-            PreviewContext previewContext =
-                    new PreviewContext(inflationContext, mGridName, mShapeKey);
+            PreviewContext previewContext = new PreviewContext(
+                    inflationContext, mGridName, mShapeKey, widgetHostId, mLayoutXml);
             PreviewAppComponent appComponent =
                     (PreviewAppComponent) LauncherComponentProvider.get(previewContext);
 
+            if (extendibleThemeManager() && isCustomLayout && !mDeletingHostOnExit) {
+                mDeletingHostOnExit = true;
+                mLifeCycleTracker.add(() -> {
+                    AppWidgetHost host = new AppWidgetHost(mContext, widgetHostId);
+                    // Start listening here, so that any previous active host is disabled
+                    host.startListening();
+                    host.stopListening();
+                    host.deleteHost();
+                });
+            }
+
             LoaderTask task = appComponent.getLoaderTaskFactory().newLoaderTask(
                     appComponent.getBaseLauncherBinderFactory().createBinder(new Callbacks[0]),
                     new UserManagerState());
 
             InvariantDeviceProfile idp = appComponent.getIDP();
             DeviceProfile deviceProfile = idp.getDeviceProfile(previewContext);
-            String query =
-                    LauncherSettings.Favorites.SCREEN + " = " + Workspace.FIRST_SCREEN_ID
-                            + " or " + LauncherSettings.Favorites.CONTAINER + " = "
-                            + LauncherSettings.Favorites.CONTAINER_HOTSEAT;
-            if (deviceProfile.isTwoPanels) {
-                query += " or " + LauncherSettings.Favorites.SCREEN + " = "
-                        + Workspace.SECOND_SCREEN_ID;
-            }
-
+            String query = deviceProfile.isTwoPanels
+                    ? selectionForWorkspaceScreen(FIRST_SCREEN_ID, SECOND_SCREEN_ID)
+                    : selectionForWorkspaceScreen(FIRST_SCREEN_ID);
             Map<ComponentKey, AppWidgetProviderInfo> widgetProviderInfoMap = new HashMap<>();
             task.loadWorkspaceForPreview(query, widgetProviderInfoMap);
             final SparseArray<Size> spanInfo = getLoadedLauncherWidgetInfo();
             MAIN_EXECUTOR.execute(() -> {
-                renderView(previewContext, appComponent.getDataModel(), widgetProviderInfoMap,
-                        spanInfo, idp);
+                renderView(previewContext, appComponent.getDataModel(), widgetHostId,
+                        widgetProviderInfoMap, spanInfo, idp);
                 mLifeCycleTracker.add(previewContext::onDestroy);
             });
         } else {
             LauncherAppState.getInstance(inflationContext).getModel().loadAsync(dataModel -> {
                 if (dataModel != null) {
-                    MAIN_EXECUTOR.execute(() -> renderView(inflationContext, dataModel, null,
-                            null, LauncherAppState.getIDP(inflationContext)));
+                    MAIN_EXECUTOR.execute(() -> renderView(inflationContext, dataModel,
+                            APPWIDGET_HOST_ID, null, null,
+                            LauncherAppState.getIDP(inflationContext)));
                 } else {
                     Log.e(TAG, "Model loading failed");
                 }
@@ -378,7 +400,7 @@
     }
 
     @UiThread
-    private void renderView(Context inflationContext, BgDataModel dataModel,
+    private void renderView(Context inflationContext, BgDataModel dataModel, int widgetHostId,
             Map<ComponentKey, AppWidgetProviderInfo> widgetProviderInfoMap,
             @Nullable final SparseArray<Size> launcherWidgetSpanInfo, InvariantDeviceProfile idp) {
         if (mDestroyed) {
@@ -386,10 +408,10 @@
         }
         LauncherPreviewRenderer renderer;
         if (Flags.newCustomizationPickerUi()) {
-            renderer = new LauncherPreviewRenderer(inflationContext, idp, mPreviewColorOverride,
-                    mWallpaperColors, launcherWidgetSpanInfo);
+            renderer = new LauncherPreviewRenderer(inflationContext, idp, widgetHostId,
+                    mPreviewColorOverride, mWallpaperColors, launcherWidgetSpanInfo);
         } else {
-            renderer = new LauncherPreviewRenderer(inflationContext, idp,
+            renderer = new LauncherPreviewRenderer(inflationContext, idp, widgetHostId,
                     mWallpaperColors, launcherWidgetSpanInfo);
         }
         renderer.hideBottomRow(mHideQsb);
diff --git a/src/com/android/launcher3/model/LayoutParserFactory.kt b/src/com/android/launcher3/model/LayoutParserFactory.kt
new file mode 100644
index 0000000..32aa58a
--- /dev/null
+++ b/src/com/android/launcher3/model/LayoutParserFactory.kt
@@ -0,0 +1,181 @@
+/*
+ * 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.app.blob.BlobHandle
+import android.app.blob.BlobStoreManager
+import android.content.Context
+import android.os.ParcelFileDescriptor.AutoCloseInputStream
+import android.provider.Settings.Secure
+import android.text.TextUtils
+import android.util.Base64
+import android.util.Log
+import android.util.Xml
+import com.android.launcher3.AutoInstallsLayout
+import com.android.launcher3.AutoInstallsLayout.SourceResources
+import com.android.launcher3.DefaultLayoutParser
+import com.android.launcher3.DefaultLayoutParser.RES_PARTNER_DEFAULT_LAYOUT
+import com.android.launcher3.LauncherSettings.Settings
+import com.android.launcher3.dagger.ApplicationContext
+import com.android.launcher3.util.IOUtils
+import com.android.launcher3.util.Partner
+import com.android.launcher3.widget.LauncherWidgetHolder
+import java.io.StringReader
+import javax.inject.Inject
+
+private const val TAG = "LayoutParserFactory"
+
+/** Utility class for providing default layout parsers */
+open class LayoutParserFactory
+@Inject
+constructor(@ApplicationContext private val context: Context) {
+
+    open fun createExternalLayoutParser(
+        widgetHolder: LauncherWidgetHolder,
+        openHelper: DatabaseHelper,
+    ): AutoInstallsLayout? {
+
+        createWorkspaceLoaderFromAppRestriction(widgetHolder, openHelper)?.let {
+            return it
+        }
+        AutoInstallsLayout.get(context, widgetHolder, openHelper)?.let {
+            return it
+        }
+
+        val partner = Partner.get(context.packageManager)
+        if (partner != null) {
+            val workspaceResId = partner.getXmlResId(RES_PARTNER_DEFAULT_LAYOUT)
+            if (workspaceResId != 0) {
+                return DefaultLayoutParser(
+                    context,
+                    widgetHolder,
+                    openHelper,
+                    partner.resources,
+                    workspaceResId,
+                )
+            }
+        }
+        return null
+    }
+
+    /**
+     * Creates workspace loader from an XML resource listed in the app restrictions.
+     *
+     * @return the loader if the restrictions are set and the resource exists; null otherwise.
+     */
+    private fun createWorkspaceLoaderFromAppRestriction(
+        widgetHolder: LauncherWidgetHolder,
+        openHelper: DatabaseHelper,
+    ): AutoInstallsLayout? {
+        val systemLayoutProvider =
+            Secure.getString(context.contentResolver, Settings.LAYOUT_PROVIDER_KEY)
+        if (TextUtils.isEmpty(systemLayoutProvider)) {
+            return null
+        }
+
+        // Try the blob store first
+        val blobManager = context.getSystemService(BlobStoreManager::class.java)
+        if (systemLayoutProvider.startsWith(Settings.BLOB_KEY_PREFIX) && blobManager != null) {
+            val blobHandlerDigest = systemLayoutProvider.substring(Settings.BLOB_KEY_PREFIX.length)
+            try {
+                AutoCloseInputStream(
+                        blobManager.openBlob(
+                            BlobHandle.createWithSha256(
+                                Base64.decode(
+                                    blobHandlerDigest,
+                                    Base64.NO_WRAP or Base64.NO_PADDING,
+                                ),
+                                Settings.LAYOUT_DIGEST_LABEL,
+                                0,
+                                Settings.LAYOUT_DIGEST_TAG,
+                            )
+                        )
+                    )
+                    .use {
+                        return getAutoInstallsLayoutFromIS(
+                            widgetHolder,
+                            openHelper,
+                            String(IOUtils.toByteArray(it)),
+                        )
+                    }
+            } catch (e: Exception) {
+                Log.e(TAG, "Error getting layout from blob handle", e)
+                return null
+            }
+        }
+
+        // Try contentProvider based provider
+        val pm = context.packageManager
+        val pi = pm.resolveContentProvider(systemLayoutProvider, 0)
+        if (pi == null) {
+            Log.e(TAG, "No provider found for authority $systemLayoutProvider")
+            return null
+        }
+        val uri = ModelDbController.getLayoutUri(systemLayoutProvider, context)
+        try {
+            context.contentResolver.openInputStream(uri)?.use {
+                Log.d(TAG, "Loading layout from $systemLayoutProvider")
+                val res = pm.getResourcesForApplication(pi.applicationInfo)
+                return getAutoInstallsLayoutFromIS(
+                    widgetHolder,
+                    openHelper,
+                    String(IOUtils.toByteArray(it)),
+                    SourceResources.wrap(res),
+                )
+            }
+        } catch (e: Exception) {
+            Log.e(TAG, "Error getting layout stream from: $systemLayoutProvider", e)
+        }
+        return null
+    }
+
+    @Throws(Exception::class)
+    protected fun getAutoInstallsLayoutFromIS(
+        widgetHolder: LauncherWidgetHolder,
+        openHelper: DatabaseHelper,
+        xml: String,
+        res: SourceResources = object : SourceResources {},
+    ): AutoInstallsLayout {
+        val parser = Xml.newPullParser()
+        parser.setInput(StringReader(xml))
+
+        return AutoInstallsLayout(
+            context,
+            widgetHolder,
+            openHelper,
+            res,
+            { parser },
+            AutoInstallsLayout.TAG_WORKSPACE,
+        )
+    }
+
+    /** Layout parser factory with fixed xml */
+    class XmlLayoutParserFactory(ctx: Context, private val xml: String) : LayoutParserFactory(ctx) {
+
+        override fun createExternalLayoutParser(
+            widgetHolder: LauncherWidgetHolder,
+            openHelper: DatabaseHelper,
+        ): AutoInstallsLayout? {
+            try {
+                return getAutoInstallsLayoutFromIS(widgetHolder, openHelper, xml)
+            } catch (e: Exception) {
+                Log.e(TAG, "Error getting layout from provided xml", e)
+                return super.createExternalLayoutParser(widgetHolder, openHelper)
+            }
+        }
+    }
+}
diff --git a/src/com/android/launcher3/model/LoaderCursor.java b/src/com/android/launcher3/model/LoaderCursor.java
index 8f116bb..fd8e2f7 100644
--- a/src/com/android/launcher3/model/LoaderCursor.java
+++ b/src/com/android/launcher3/model/LoaderCursor.java
@@ -45,6 +45,7 @@
 import com.android.launcher3.Flags;
 import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.LauncherModel;
+import com.android.launcher3.LauncherPrefs;
 import com.android.launcher3.LauncherSettings.Favorites;
 import com.android.launcher3.Utilities;
 import com.android.launcher3.Workspace;
@@ -421,7 +422,8 @@
     ) {
         boolean isPreArchived = Flags.enableSupportForArchiving()
                 && Flags.restoreArchivedAppIconsFromDb()
-                && info.isInactiveArchive();
+                && info.isInactiveArchive()
+                && LauncherPrefs.get(mContext).get(LauncherPrefs.IS_FIRST_LOAD_AFTER_RESTORE);
         boolean preArchivedIconNotFound = isPreArchived && !loadIconFromDb(info);
         if (preArchivedIconNotFound) {
             Log.d(TAG, "loadIconFromDb failed for pre-archived icon, loading from cache."
diff --git a/src/com/android/launcher3/model/ModelDbController.java b/src/com/android/launcher3/model/ModelDbController.java
index 64b9c1c..f48d848 100644
--- a/src/com/android/launcher3/model/ModelDbController.java
+++ b/src/com/android/launcher3/model/ModelDbController.java
@@ -16,10 +16,7 @@
 package com.android.launcher3.model;
 
 import static android.provider.BaseColumns._ID;
-import static android.util.Base64.NO_PADDING;
-import static android.util.Base64.NO_WRAP;
 
-import static com.android.launcher3.DefaultLayoutParser.RES_PARTNER_DEFAULT_LAYOUT;
 import static com.android.launcher3.LauncherPrefs.DB_FILE;
 import static com.android.launcher3.LauncherPrefs.NO_DB_FILES_RESTORED;
 import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER;
@@ -27,40 +24,25 @@
 import static com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_APP_PAIR;
 import static com.android.launcher3.LauncherSettings.Favorites.TABLE_NAME;
 import static com.android.launcher3.LauncherSettings.Favorites.addTableToDb;
-import static com.android.launcher3.LauncherSettings.Settings.BLOB_KEY_PREFIX;
-import static com.android.launcher3.LauncherSettings.Settings.LAYOUT_DIGEST_LABEL;
-import static com.android.launcher3.LauncherSettings.Settings.LAYOUT_DIGEST_TAG;
-import static com.android.launcher3.LauncherSettings.Settings.LAYOUT_PROVIDER_KEY;
 import static com.android.launcher3.provider.LauncherDbUtils.tableExists;
 
-import android.app.blob.BlobHandle;
-import android.app.blob.BlobStoreManager;
-import android.content.ContentResolver;
 import android.content.ContentValues;
 import android.content.Context;
-import android.content.pm.PackageManager;
-import android.content.pm.ProviderInfo;
-import android.content.res.Resources;
 import android.database.Cursor;
 import android.database.SQLException;
 import android.database.sqlite.SQLiteDatabase;
 import android.net.Uri;
 import android.os.Bundle;
-import android.os.ParcelFileDescriptor;
 import android.os.Process;
 import android.os.UserHandle;
 import android.os.UserManager;
-import android.provider.Settings;
 import android.text.TextUtils;
-import android.util.Base64;
 import android.util.Log;
-import android.util.Xml;
 
 import androidx.annotation.Nullable;
 import androidx.annotation.WorkerThread;
 
 import com.android.launcher3.AutoInstallsLayout;
-import com.android.launcher3.AutoInstallsLayout.SourceResources;
 import com.android.launcher3.ConstantItem;
 import com.android.launcher3.DefaultLayoutParser;
 import com.android.launcher3.EncryptionType;
@@ -80,16 +62,10 @@
 import com.android.launcher3.provider.LauncherDbUtils;
 import com.android.launcher3.provider.LauncherDbUtils.SQLiteTransaction;
 import com.android.launcher3.provider.RestoreDbTask;
-import com.android.launcher3.util.IOUtils;
 import com.android.launcher3.util.IntArray;
-import com.android.launcher3.util.Partner;
 import com.android.launcher3.widget.LauncherWidgetHolder;
 
-import org.xmlpull.v1.XmlPullParser;
-
 import java.io.File;
-import java.io.InputStream;
-import java.io.StringReader;
 import java.util.List;
 import java.util.stream.Collectors;
 
@@ -113,17 +89,20 @@
     private final InvariantDeviceProfile mIdp;
     private final LauncherPrefs mPrefs;
     private final UserCache mUserCache;
+    private final LayoutParserFactory mLayoutParserFactory;
 
     @Inject
     ModelDbController(
             @ApplicationContext Context context,
             InvariantDeviceProfile idp,
             LauncherPrefs prefs,
-            UserCache userCache) {
+            UserCache userCache,
+            LayoutParserFactory layoutParserFactory) {
         mContext = context;
         mIdp = idp;
         mPrefs = prefs;
         mUserCache = userCache;
+        mLayoutParserFactory = layoutParserFactory;
     }
 
     private void printDBs(String prefix) {
@@ -368,7 +347,6 @@
     public void attemptMigrateDb(LauncherRestoreEventLogger restoreEventLogger,
             ModelDelegate modelDelegate) throws Exception {
         createDbIfNotExists();
-
         if (shouldResetDb()) {
             resetLauncherDb(restoreEventLogger);
             return;
@@ -634,20 +612,8 @@
 
             LauncherWidgetHolder widgetHolder = mOpenHelper.newLauncherWidgetHolder();
             try {
-                AutoInstallsLayout loader = createWorkspaceLoaderFromAppRestriction(widgetHolder);
-                if (loader == null) {
-                    loader = AutoInstallsLayout.get(mContext, widgetHolder, mOpenHelper);
-                }
-                if (loader == null) {
-                    final Partner partner = Partner.get(mContext.getPackageManager());
-                    if (partner != null) {
-                        int workspaceResId = partner.getXmlResId(RES_PARTNER_DEFAULT_LAYOUT);
-                        if (workspaceResId != 0) {
-                            loader = new DefaultLayoutParser(mContext, widgetHolder,
-                                    mOpenHelper, partner.getResources(), workspaceResId);
-                        }
-                    }
-                }
+                AutoInstallsLayout loader =
+                        mLayoutParserFactory.createExternalLayoutParser(widgetHolder, mOpenHelper);
 
                 final boolean usingExternallyProvidedLayout = loader != null;
                 if (loader == null) {
@@ -672,64 +638,6 @@
         }
     }
 
-    /**
-     * Creates workspace loader from an XML resource listed in the app restrictions.
-     *
-     * @return the loader if the restrictions are set and the resource exists; null otherwise.
-     */
-    private AutoInstallsLayout createWorkspaceLoaderFromAppRestriction(
-            LauncherWidgetHolder widgetHolder) {
-        ContentResolver cr = mContext.getContentResolver();
-        String systemLayoutProvider = Settings.Secure.getString(cr, LAYOUT_PROVIDER_KEY);
-        if (TextUtils.isEmpty(systemLayoutProvider)) {
-            return null;
-        }
-
-        // Try the blob store first
-        if (systemLayoutProvider.startsWith(BLOB_KEY_PREFIX)) {
-            BlobStoreManager blobManager = mContext.getSystemService(BlobStoreManager.class);
-            String blobHandlerDigest = systemLayoutProvider.substring(BLOB_KEY_PREFIX.length());
-            try (InputStream in = new ParcelFileDescriptor.AutoCloseInputStream(
-                    blobManager.openBlob(BlobHandle.createWithSha256(
-                            Base64.decode(blobHandlerDigest, NO_WRAP | NO_PADDING),
-                            LAYOUT_DIGEST_LABEL, 0, LAYOUT_DIGEST_TAG)))) {
-                return getAutoInstallsLayoutFromIS(in, widgetHolder, new SourceResources() { });
-            } catch (Exception e) {
-                Log.e(TAG, "Error getting layout from blob handle" , e);
-                return null;
-            }
-        }
-
-        // Try contentProvider based provider
-        PackageManager pm = mContext.getPackageManager();
-        ProviderInfo pi = pm.resolveContentProvider(systemLayoutProvider, 0);
-        if (pi == null) {
-            Log.e(TAG, "No provider found for authority " + systemLayoutProvider);
-            return null;
-        }
-        Uri uri = getLayoutUri(systemLayoutProvider, mContext);
-        try (InputStream in = cr.openInputStream(uri)) {
-            Log.d(TAG, "Loading layout from " + systemLayoutProvider);
-
-            Resources res = pm.getResourcesForApplication(pi.applicationInfo);
-            return getAutoInstallsLayoutFromIS(in, widgetHolder, SourceResources.wrap(res));
-        } catch (Exception e) {
-            Log.e(TAG, "Error getting layout stream from: " + systemLayoutProvider , e);
-            return null;
-        }
-    }
-
-    private AutoInstallsLayout getAutoInstallsLayoutFromIS(InputStream in,
-            LauncherWidgetHolder widgetHolder, SourceResources res) throws Exception {
-        // Read the full xml so that we fail early in case of any IO error.
-        String layout = new String(IOUtils.toByteArray(in));
-        XmlPullParser parser = Xml.newPullParser();
-        parser.setInput(new StringReader(layout));
-
-        return new AutoInstallsLayout(mContext, widgetHolder, mOpenHelper, res,
-                () -> parser, AutoInstallsLayout.TAG_WORKSPACE);
-    }
-
     public static Uri getLayoutUri(String authority, Context ctx) {
         InvariantDeviceProfile grid = LauncherAppState.getIDP(ctx);
         return new Uri.Builder().scheme("content").authority(authority).path("launcher_layout")
diff --git a/src/com/android/launcher3/provider/LauncherDbUtils.kt b/src/com/android/launcher3/provider/LauncherDbUtils.kt
index 3641371..7ee4988 100644
--- a/src/com/android/launcher3/provider/LauncherDbUtils.kt
+++ b/src/com/android/launcher3/provider/LauncherDbUtils.kt
@@ -27,7 +27,13 @@
 import android.os.UserManager
 import android.text.TextUtils
 import com.android.launcher3.LauncherSettings
+import com.android.launcher3.LauncherSettings.Favorites
+import com.android.launcher3.LauncherSettings.Favorites.CONTAINER
 import com.android.launcher3.LauncherSettings.Favorites.CONTAINER_DESKTOP
+import com.android.launcher3.LauncherSettings.Favorites.CONTAINER_HOTSEAT
+import com.android.launcher3.LauncherSettings.Favorites.SCREEN
+import com.android.launcher3.LauncherSettings.Favorites.TABLE_NAME
+import com.android.launcher3.LauncherSettings.Favorites._ID
 import com.android.launcher3.Utilities
 import com.android.launcher3.dagger.LauncherComponentProvider.appComponent
 import com.android.launcher3.icons.IconCache
@@ -45,6 +51,14 @@
      */
     @JvmStatic fun itemIdMatch(itemId: Int): String = "_id=$itemId"
 
+    /**
+     * Returns a string which can be used as a where clause for DB query to match the given
+     * workspace screens or hotseat or a collection in workspace screens or hotseat
+     */
+    @JvmStatic
+    fun selectionForWorkspaceScreen(vararg screens: Int) =
+        "$SCREEN in (${screens.joinToString()}) or $CONTAINER = $CONTAINER_HOTSEAT or $CONTAINER in (select $_ID from $TABLE_NAME where $SCREEN in (${screens.joinToString()}) or $CONTAINER = $CONTAINER_HOTSEAT)"
+
     @JvmStatic
     fun queryIntArray(
         distinct: Boolean,
diff --git a/src/com/android/launcher3/widget/LauncherWidgetHolder.java b/src/com/android/launcher3/widget/LauncherWidgetHolder.java
index 642f35a..f7c34f6 100644
--- a/src/com/android/launcher3/widget/LauncherWidgetHolder.java
+++ b/src/com/android/launcher3/widget/LauncherWidgetHolder.java
@@ -107,7 +107,11 @@
 
     @AssistedInject
     protected LauncherWidgetHolder(@Assisted("UI_CONTEXT") @NonNull Context context) {
-        this(context, new LauncherAppWidgetHost(context, APPWIDGET_HOST_ID));
+        this(context, APPWIDGET_HOST_ID);
+    }
+
+    public LauncherWidgetHolder(@NonNull Context context, int hostId) {
+        this(context, new LauncherAppWidgetHost(context, hostId));
     }
 
     protected LauncherWidgetHolder(
diff --git a/tests/multivalentTests/src/com/android/launcher3/model/LoaderCursorTest.java b/tests/multivalentTests/src/com/android/launcher3/model/LoaderCursorTest.java
index b848d27..d1292cf 100644
--- a/tests/multivalentTests/src/com/android/launcher3/model/LoaderCursorTest.java
+++ b/tests/multivalentTests/src/com/android/launcher3/model/LoaderCursorTest.java
@@ -21,6 +21,7 @@
 
 import static androidx.test.InstrumentationRegistry.getContext;
 
+import static com.android.launcher3.LauncherPrefs.IS_FIRST_LOAD_AFTER_RESTORE;
 import static com.android.launcher3.LauncherSettings.Favorites.APPWIDGET_ID;
 import static com.android.launcher3.LauncherSettings.Favorites.APPWIDGET_PROVIDER;
 import static com.android.launcher3.LauncherSettings.Favorites.APPWIDGET_SOURCE;
@@ -69,6 +70,7 @@
 import com.android.launcher3.Flags;
 import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.LauncherAppState;
+import com.android.launcher3.LauncherPrefs;
 import com.android.launcher3.icons.BitmapInfo;
 import com.android.launcher3.icons.LauncherIcons;
 import com.android.launcher3.model.data.ItemInfo;
@@ -96,6 +98,7 @@
 
     private LauncherModelHelper mModelHelper;
     private LauncherAppState mApp;
+    private LauncherPrefs mPrefs;
 
     private MatrixCursor mCursor;
     private InvariantDeviceProfile mIDP;
@@ -113,6 +116,7 @@
     public void setup() {
         mModelHelper = new LauncherModelHelper();
         mContext = mModelHelper.sandboxContext;
+        mPrefs = LauncherPrefs.get(mContext);
         mIDP = InvariantDeviceProfile.INSTANCE.get(mContext);
         mApp = LauncherAppState.getInstance(mContext);
 
@@ -131,6 +135,7 @@
 
     @After
     public void tearDown() {
+        mPrefs.putSync(IS_FIRST_LOAD_AFTER_RESTORE.to(false));
         mCursor.close();
         mModelHelper.destroy();
     }
@@ -253,8 +258,9 @@
 
     @Test
     @EnableFlags(Flags.FLAG_RESTORE_ARCHIVED_APP_ICONS_FROM_DB)
-    public void ifArchivedWithFlag_whenloadWorkspaceTitleAndIcon_thenLoadIconFromDb() {
+    public void ifArchivedWithFlagAndRestore_whenloadWorkspaceTitleAndIcon_thenLoadIconFromDb() {
         // Given
+        mPrefs.putSync(IS_FIRST_LOAD_AFTER_RESTORE.to(true));
         initCursor(ITEM_TYPE_APPLICATION, "title");
         assertTrue(mLoaderCursor.moveToNext());
         WorkspaceItemInfo itemInfo = new WorkspaceItemInfo();
@@ -262,6 +268,7 @@
         Bitmap expectedBitmap = LauncherIcons.obtain(mContext)
                 .createIconBitmap(decodeByteArray(sTestBlob, 0, sTestBlob.length))
                 .icon;
+
         // When
         mLoaderCursor.loadWorkspaceTitleAndIcon(false, true, itemInfo);
         // Then
@@ -271,6 +278,23 @@
 
     @Test
     @EnableFlags(Flags.FLAG_RESTORE_ARCHIVED_APP_ICONS_FROM_DB)
+    public void ifArchivedWithFlagAndNotRestore_whenloadWorkspaceTitleAndIcon_thenLoadIconFromDb() {
+        // Given
+        mPrefs.putSync(IS_FIRST_LOAD_AFTER_RESTORE.to(false));
+        initCursor(ITEM_TYPE_APPLICATION, "title");
+        assertTrue(mLoaderCursor.moveToNext());
+        WorkspaceItemInfo itemInfo = new WorkspaceItemInfo();
+        BitmapInfo original = itemInfo.bitmap;
+        itemInfo.runtimeStatusFlags |= FLAG_ARCHIVED;
+
+        // When
+        mLoaderCursor.loadWorkspaceTitleAndIcon(false, true, itemInfo);
+        // Then
+        assertThat(itemInfo.bitmap).isEqualTo(original);
+    }
+
+    @Test
+    @EnableFlags(Flags.FLAG_RESTORE_ARCHIVED_APP_ICONS_FROM_DB)
     public void ifArchivedWithFlag_whenLoadIconFromDb_thenLoadIconFromBlob() {
         // Given
         initCursor(ITEM_TYPE_APPLICATION, "title");