Merge "Hook up luma sampling state change from CommandQueue" into main
diff --git a/quickstep/src/com/android/launcher3/uioverrides/QuickstepAppWidgetHost.java b/quickstep/src/com/android/launcher3/uioverrides/QuickstepAppWidgetHost.java
index 6659fa0..45813ce 100644
--- a/quickstep/src/com/android/launcher3/uioverrides/QuickstepAppWidgetHost.java
+++ b/quickstep/src/com/android/launcher3/uioverrides/QuickstepAppWidgetHost.java
@@ -25,6 +25,7 @@
 import androidx.annotation.NonNull;
 
 import com.android.launcher3.LauncherAppState;
+import com.android.launcher3.util.Executors;
 import com.android.launcher3.widget.LauncherAppWidgetProviderInfo;
 import com.android.launcher3.widget.LauncherWidgetHolder;
 
@@ -55,7 +56,10 @@
 
     @Override
     public void onAppWidgetRemoved(int appWidgetId) {
-        mAppWidgetRemovedCallback.accept(appWidgetId);
+        // Route the call via model thread, in case it comes while a loader-bind is in progress
+        Executors.MODEL_EXECUTOR.execute(
+                () -> Executors.MAIN_EXECUTOR.execute(
+                        () -> mAppWidgetRemovedCallback.accept(appWidgetId)));
     }
 
     @Override
diff --git a/quickstep/src/com/android/launcher3/uioverrides/QuickstepWidgetHolder.java b/quickstep/src/com/android/launcher3/uioverrides/QuickstepWidgetHolder.java
index 1a75535..0fb2b17 100644
--- a/quickstep/src/com/android/launcher3/uioverrides/QuickstepWidgetHolder.java
+++ b/quickstep/src/com/android/launcher3/uioverrides/QuickstepWidgetHolder.java
@@ -34,10 +34,10 @@
 import com.android.launcher3.config.FeatureFlags;
 import com.android.launcher3.model.WidgetsModel;
 import com.android.launcher3.util.IntSet;
+import com.android.launcher3.util.SafeCloseable;
 import com.android.launcher3.widget.LauncherAppWidgetHostView;
 import com.android.launcher3.widget.LauncherAppWidgetProviderInfo;
 import com.android.launcher3.widget.LauncherWidgetHolder;
-import com.android.launcher3.widget.custom.CustomWidgetManager;
 
 import java.util.ArrayList;
 import java.util.Collections;
@@ -67,11 +67,11 @@
 
     private static AppWidgetHost sWidgetHost = null;
 
+    private final UpdateHandler mUpdateHandler = this::onWidgetUpdate;
     private final @Nullable RemoteViews.InteractionHandler mInteractionHandler;
 
     private final @NonNull IntConsumer mAppWidgetRemovedCallback;
 
-    private final ArrayList<ProviderChangedListener> mProviderChangedListeners = new ArrayList<>();
     // Map to all pending updated keyed with appWidgetId;
     private final SparseArray<PendingUpdate> mPendingUpdateMap = new SparseArray<>();
 
@@ -175,7 +175,10 @@
     @Override
     public void destroy() {
         try {
-            MAIN_EXECUTOR.submit(() -> sHolders.remove(this)).get();
+            MAIN_EXECUTOR.submit(() -> {
+                clearViews();
+                sHolders.remove(this);
+            }).get();
         } catch (Exception e) {
             Log.e(TAG, "Failed to remove self from holder list", e);
         }
@@ -188,26 +191,6 @@
     }
 
     /**
-     * Add a listener that is triggered when the providers of the widgets are changed
-     * @param listener The listener that notifies when the providers changed
-     */
-    @Override
-    public void addProviderChangeListener(
-            @NonNull LauncherWidgetHolder.ProviderChangedListener listener) {
-        MAIN_EXECUTOR.execute(() -> mProviderChangedListeners.add(listener));
-    }
-
-    /**
-     * Remove the specified listener from the host
-     * @param listener The listener that is to be removed from the host
-     */
-    @Override
-    public void removeProviderChangeListener(
-            LauncherWidgetHolder.ProviderChangedListener listener) {
-        MAIN_EXECUTOR.execute(() -> mProviderChangedListeners.remove(listener));
-    }
-
-    /**
      * Stop the host from updating the widget views
      */
     @Override
@@ -220,44 +203,41 @@
         setListeningFlag(false);
     }
 
-    /**
-     * Create a view for the specified app widget
-     * @param context The activity context for which the view is created
-     * @param appWidgetId The ID of the widget
-     * @param appWidget The {@link LauncherAppWidgetProviderInfo} of the widget
-     * @return A view for the widget
-     */
+    @Override
+    public SafeCloseable addOnUpdateListener(int appWidgetId,
+            LauncherAppWidgetProviderInfo appWidget, Runnable callback) {
+        UpdateHandler handler = new UpdateHandler() {
+            @Override
+            public <T> void onWidgetUpdate(int widgetId, UpdateKey<T> key, T data) {
+                if (KEY_VIEWS_UPDATE == key) {
+                    callback.run();
+                }
+            }
+        };
+        QuickstepWidgetHolderListener holderListener = getHolderListener(appWidgetId);
+        holderListener.addHolder(handler);
+        return () -> holderListener.mListeningHolders.remove(handler);
+    }
+
     @NonNull
     @Override
-    public LauncherAppWidgetHostView createView(@NonNull Context context, int appWidgetId,
-            @NonNull LauncherAppWidgetProviderInfo appWidget) {
-        if (appWidget.isCustomWidget()) {
-            LauncherAppWidgetHostView lahv = new LauncherAppWidgetHostView(context);
-            lahv.setAppWidget(appWidgetId, appWidget);
-            CustomWidgetManager.INSTANCE.get(context).onViewCreated(lahv);
-            return lahv;
-        }
-
-        LauncherAppWidgetHostView widgetView = getPendingView(appWidgetId);
-        if (widgetView != null) {
-            removePendingView(appWidgetId);
-        } else {
-            widgetView = new LauncherAppWidgetHostView(context);
-        }
+    protected LauncherAppWidgetHostView createViewInternal(
+            int appWidgetId, @NonNull LauncherAppWidgetProviderInfo appWidget) {
+        LauncherAppWidgetHostView widgetView = new LauncherAppWidgetHostView(mContext);
         widgetView.setInteractionHandler(mInteractionHandler);
         widgetView.setAppWidget(appWidgetId, appWidget);
-        mViews.put(appWidgetId, widgetView);
+        widgetView.updateAppWidget(getHolderListener(appWidgetId).addHolder(mUpdateHandler));
+        return widgetView;
+    }
 
+    private static QuickstepWidgetHolderListener getHolderListener(int appWidgetId) {
         QuickstepWidgetHolderListener listener = sListeners.get(appWidgetId);
         if (listener == null) {
             listener = new QuickstepWidgetHolderListener(appWidgetId);
             sWidgetHost.setListener(appWidgetId, listener);
             sListeners.put(appWidgetId, listener);
         }
-        RemoteViews remoteViews = listener.addHolder(this);
-        widgetView.updateAppWidget(remoteViews);
-
-        return widgetView;
+        return listener;
     }
 
     /**
@@ -267,7 +247,7 @@
     public void clearViews() {
         mViews.clear();
         for (int i = sListeners.size() - 1; i >= 0; i--) {
-            sListeners.valueAt(i).mListeningHolders.remove(this);
+            sListeners.valueAt(i).mListeningHolders.remove(mUpdateHandler);
         }
     }
 
@@ -275,7 +255,7 @@
             implements AppWidgetHost.AppWidgetHostListener {
 
         // Static listeners should use a set that is backed by WeakHashMap to avoid memory leak
-        private final Set<QuickstepWidgetHolder> mListeningHolders = Collections.newSetFromMap(
+        private final Set<UpdateHandler> mListeningHolders = Collections.newSetFromMap(
                 new WeakHashMap<>());
 
         private final int mWidgetId;
@@ -288,7 +268,7 @@
 
         @UiThread
         @Nullable
-        public RemoteViews addHolder(@NonNull QuickstepWidgetHolder holder) {
+        public RemoteViews addHolder(@NonNull UpdateHandler holder) {
             mListeningHolders.add(holder);
             return mRemoteViews;
         }
@@ -359,11 +339,15 @@
         }
     }
 
+    private interface UpdateKey<T> extends BiConsumer<AppWidgetHostView, T> { }
+
+    private interface UpdateHandler {
+        <T> void onWidgetUpdate(int widgetId, UpdateKey<T> key, T data);
+    }
+
     private static class PendingUpdate {
         public final IntSet changedViews = new IntSet();
         public AppWidgetProviderInfo providerInfo;
         public RemoteViews remoteViews;
     }
-
-    private interface UpdateKey<T> extends BiConsumer<AppWidgetHostView, T> { }
 }
diff --git a/quickstep/src/com/android/launcher3/uioverrides/touchcontrollers/PortraitStatesTouchController.java b/quickstep/src/com/android/launcher3/uioverrides/touchcontrollers/PortraitStatesTouchController.java
index d94cd89..5db5d17 100644
--- a/quickstep/src/com/android/launcher3/uioverrides/touchcontrollers/PortraitStatesTouchController.java
+++ b/quickstep/src/com/android/launcher3/uioverrides/touchcontrollers/PortraitStatesTouchController.java
@@ -106,7 +106,7 @@
     protected StateAnimationConfig getConfigForStates(
             LauncherState fromState, LauncherState toState) {
         final StateAnimationConfig config = new StateAnimationConfig();
-        config.userControlled = true;
+        config.animProps |= StateAnimationConfig.USER_CONTROLLED;
         if (fromState == NORMAL && toState == ALL_APPS) {
             AllAppsSwipeController.applyNormalToAllAppsAnimConfig(mLauncher, config);
         } else if (fromState == ALL_APPS && toState == NORMAL) {
diff --git a/quickstep/src/com/android/quickstep/util/TaskViewSimulator.java b/quickstep/src/com/android/quickstep/util/TaskViewSimulator.java
index baaa062..065a9c5 100644
--- a/quickstep/src/com/android/quickstep/util/TaskViewSimulator.java
+++ b/quickstep/src/com/android/quickstep/util/TaskViewSimulator.java
@@ -74,6 +74,7 @@
     private final boolean mIsRecentsRtl;
 
     private final Rect mTaskRect = new Rect();
+    private final Rect mFullTaskSize = new Rect();
     private final PointF mPivot = new PointF();
     private DeviceProfile mDp;
     @StagePosition
@@ -131,6 +132,47 @@
         mDp = dp;
         mLayoutValid = false;
         mOrientationState.setDeviceProfile(dp);
+        calculateTaskSize();
+    }
+
+    private void calculateTaskSize() {
+        if (mDp == null) {
+            return;
+        }
+
+        if (mIsGridTask) {
+            mSizeStrategy.calculateGridTaskSize(mContext, mDp, mFullTaskSize,
+                    mOrientationState.getOrientationHandler());
+        } else {
+            mSizeStrategy.calculateTaskSize(mContext, mDp, mFullTaskSize,
+                    mOrientationState.getOrientationHandler());
+        }
+
+        if (mSplitBounds != null) {
+            // The task rect changes according to the staged split task sizes, but recents
+            // fullscreen scale and pivot remains the same since the task fits into the existing
+            // sized task space bounds
+            mTaskRect.set(mFullTaskSize);
+            mOrientationState.getOrientationHandler()
+                    .setSplitTaskSwipeRect(mDp, mTaskRect, mSplitBounds, mStagePosition);
+            mTaskRect.offset(mTaskRectTranslationX, mTaskRectTranslationY);
+        } else if (mIsDesktopTask) {
+            // For desktop, tasks can take up only part of the screen size.
+            // Full task size represents the whole screen size, but scaled down to fit in recents.
+            // Task rect will represent the scaled down thumbnail position and is placed inside
+            // full task size as it is on the home screen.
+            PointF fullscreenTaskDimension = new PointF();
+            BaseActivityInterface.getTaskDimension(mContext, mDp, fullscreenTaskDimension);
+            // Calculate the scale down factor used in recents
+            float scale = mFullTaskSize.width() / fullscreenTaskDimension.x;
+            mTaskRect.set(mThumbnailPosition);
+            mTaskRect.scale(scale);
+            // Ensure the task rect is inside the full task rect
+            mTaskRect.offset(mFullTaskSize.left, mFullTaskSize.top);
+        } else {
+            mTaskRect.set(mFullTaskSize);
+            mTaskRect.offset(mTaskRectTranslationX, mTaskRectTranslationY);
+        }
     }
 
     /**
@@ -148,44 +190,11 @@
         if (mDp == null) {
             return 1;
         }
-
-        if (mIsGridTask) {
-            mSizeStrategy.calculateGridTaskSize(mContext, mDp, mTaskRect,
-                    mOrientationState.getOrientationHandler());
-        } else {
-            mSizeStrategy.calculateTaskSize(mContext, mDp, mTaskRect,
-                    mOrientationState.getOrientationHandler());
-        }
-
-        Rect fullTaskSize;
-        if (mSplitBounds != null) {
-            // The task rect changes according to the staged split task sizes, but recents
-            // fullscreen scale and pivot remains the same since the task fits into the existing
-            // sized task space bounds
-            fullTaskSize = new Rect(mTaskRect);
-            mOrientationState.getOrientationHandler()
-                    .setSplitTaskSwipeRect(mDp, mTaskRect, mSplitBounds, mStagePosition);
-            mTaskRect.offset(mTaskRectTranslationX, mTaskRectTranslationY);
-        } else if (mIsDesktopTask) {
-            // For desktop, tasks can take up only part of the screen size.
-            // Full task size represents the whole screen size, but scaled down to fit in recents.
-            // Task rect will represent the scaled down thumbnail position and is placed inside
-            // full task size as it is on the home screen.
-            fullTaskSize = new Rect(mTaskRect);
-            PointF fullscreenTaskDimension = new PointF();
-            BaseActivityInterface.getTaskDimension(mContext, mDp, fullscreenTaskDimension);
-            // Calculate the scale down factor used in recents
-            float scale = fullTaskSize.width() / fullscreenTaskDimension.x;
-            mTaskRect.set(mThumbnailPosition);
-            mTaskRect.scale(scale);
-            // Ensure the task rect is inside the full task rect
-            mTaskRect.offset(fullTaskSize.left, fullTaskSize.top);
-        } else {
-            fullTaskSize = new Rect(mTaskRect);
-            mTaskRect.offset(mTaskRectTranslationX, mTaskRectTranslationY);
-        }
-        fullTaskSize.offset(mTaskRectTranslationX + mPivotOffsetX, mTaskRectTranslationY);
-        return mOrientationState.getFullScreenScaleAndPivot(fullTaskSize, mDp, mPivot);
+        // Copy mFullTaskSize instead of updating it directly so it could be reused next time
+        // without recalculating
+        Rect scaleRect = new Rect(mFullTaskSize);
+        scaleRect.offset(mTaskRectTranslationX + mPivotOffsetX, mTaskRectTranslationY);
+        return mOrientationState.getFullScreenScaleAndPivot(scaleRect, mDp, mPivot);
     }
 
     /**
@@ -209,13 +218,13 @@
         mSplitBounds = splitInfo;
         if (mSplitBounds == null) {
             mStagePosition = STAGE_POSITION_UNDEFINED;
-            return;
+        } else {
+            mStagePosition = mThumbnailPosition.equals(splitInfo.leftTopBounds)
+                    ? STAGE_POSITION_TOP_OR_LEFT : STAGE_POSITION_BOTTOM_OR_RIGHT;
+            mPositionHelper.setSplitBounds(convertLauncherSplitBoundsToShell(mSplitBounds),
+                    mStagePosition);
         }
-        mStagePosition = mThumbnailPosition.equals(splitInfo.leftTopBounds) ?
-                STAGE_POSITION_TOP_OR_LEFT :
-                STAGE_POSITION_BOTTOM_OR_RIGHT;
-        mPositionHelper.setSplitBounds(convertLauncherSplitBoundsToShell(mSplitBounds),
-                mStagePosition);
+        calculateTaskSize();
     }
 
     /**
@@ -261,6 +270,8 @@
     public void setTaskRectTranslation(int taskRectTranslationX, int taskRectTranslationY) {
         mTaskRectTranslationX = taskRectTranslationX;
         mTaskRectTranslationY = taskRectTranslationY;
+        // Re-calculate task size after changing translation
+        calculateTaskSize();
     }
 
     /**
@@ -269,7 +280,7 @@
     public void addAppToOverviewAnim(PendingAnimation pa, TimeInterpolator interpolator) {
         pa.addFloat(fullScreenProgress, AnimatedFloat.VALUE, 1, 0, interpolator);
         if (enableGridOnlyOverview() && mDp.isTablet) {
-            int translationXToMiddle = mDp.widthPx / 2 - mTaskRect.centerX();
+            int translationXToMiddle = mDp.widthPx / 2 - mFullTaskSize.centerX();
             taskPrimaryTranslation.value = translationXToMiddle;
             mPivotOffsetX = translationXToMiddle;
         }
@@ -324,8 +335,8 @@
     public void applyWindowToHomeRotation(Matrix matrix) {
         matrix.postTranslate(mDp.windowX, mDp.windowY);
         postDisplayRotation(deltaRotation(
-                mOrientationState.getRecentsActivityRotation(),
-                mOrientationState.getDisplayRotation()),
+                        mOrientationState.getRecentsActivityRotation(),
+                        mOrientationState.getDisplayRotation()),
                 mDp.widthPx, mDp.heightPx, matrix);
     }
 
diff --git a/src/com/android/launcher3/Launcher.java b/src/com/android/launcher3/Launcher.java
index e0e35a4..5c49b89 100644
--- a/src/com/android/launcher3/Launcher.java
+++ b/src/com/android/launcher3/Launcher.java
@@ -1011,7 +1011,7 @@
         AppWidgetHostView boundWidget = null;
         if (resultCode == RESULT_OK) {
             animationType = Workspace.COMPLETE_TWO_STAGE_WIDGET_DROP_ANIMATION;
-            final AppWidgetHostView layout = mAppWidgetHolder.createView(this, appWidgetId,
+            final AppWidgetHostView layout = mAppWidgetHolder.createView(appWidgetId,
                     requestArgs.getWidgetHandler().getProviderInfo(this));
             boundWidget = layout;
             onCompleteRunnable = () -> {
@@ -1464,7 +1464,7 @@
 
         if (hostView == null) {
             // Perform actual inflation because we're live
-            hostView = mAppWidgetHolder.createView(this, appWidgetId, appWidgetInfo);
+            hostView = mAppWidgetHolder.createView(appWidgetId, appWidgetInfo);
         }
 
         LauncherAppWidgetInfo launcherInfo;
@@ -2319,7 +2319,7 @@
         }
         final AppWidgetHostView view;
         if (mIsSafeModeEnabled) {
-            view = new PendingAppWidgetHostView(this, item, mIconCache, true);
+            view = new PendingAppWidgetHostView(this, item, null, true);
             prepareAppWidget(view, item);
             return view;
         }
@@ -2450,14 +2450,9 @@
 
                 item.minSpanX = appWidgetInfo.minSpanX;
                 item.minSpanY = appWidgetInfo.minSpanY;
-                view = mAppWidgetHolder.createView(this, item.appWidgetId, appWidgetInfo);
-            } else if (!item.hasRestoreFlag(LauncherAppWidgetInfo.FLAG_ID_NOT_VALID)
-                    && appWidgetInfo != null) {
-                mAppWidgetHolder.addPendingView(item.appWidgetId,
-                        new PendingAppWidgetHostView(this, item, mIconCache, false));
-                view = mAppWidgetHolder.createView(this, item.appWidgetId, appWidgetInfo);
+                view = mAppWidgetHolder.createView(item.appWidgetId, appWidgetInfo);
             } else {
-                view = new PendingAppWidgetHostView(this, item, mIconCache, false);
+                view = new PendingAppWidgetHostView(this, item, appWidgetInfo, false);
             }
             prepareAppWidget(view, item);
         } finally {
diff --git a/src/com/android/launcher3/allapps/AllAppsTransitionController.java b/src/com/android/launcher3/allapps/AllAppsTransitionController.java
index 03ac9df..e2c5795 100644
--- a/src/com/android/launcher3/allapps/AllAppsTransitionController.java
+++ b/src/com/android/launcher3/allapps/AllAppsTransitionController.java
@@ -342,7 +342,7 @@
             });
         }
 
-        if(FeatureFlags.ENABLE_PREMIUM_HAPTICS_ALL_APPS.get() && config.userControlled
+        if (FeatureFlags.ENABLE_PREMIUM_HAPTICS_ALL_APPS.get() && config.isUserControlled()
                 && Utilities.ATLEAST_S) {
             if (toState == ALL_APPS) {
                 builder.addOnFrameListener(
@@ -367,7 +367,7 @@
 
         // need to decide depending on the release velocity
         Interpolator verticalProgressInterpolator = config.getInterpolator(ANIM_VERTICAL_PROGRESS,
-                config.userControlled ? LINEAR : DECELERATE_1_7);
+                config.isUserControlled() ? LINEAR : DECELERATE_1_7);
         Animator anim = createSpringAnimation(mProgress, targetProgress);
         anim.setInterpolator(verticalProgressInterpolator);
         builder.add(anim);
diff --git a/src/com/android/launcher3/folder/PreviewItemManager.java b/src/com/android/launcher3/folder/PreviewItemManager.java
index b39e968e..871643e 100644
--- a/src/com/android/launcher3/folder/PreviewItemManager.java
+++ b/src/com/android/launcher3/folder/PreviewItemManager.java
@@ -37,6 +37,7 @@
 import android.view.View;
 
 import androidx.annotation.NonNull;
+import androidx.annotation.VisibleForTesting;
 
 import com.android.launcher3.BubbleTextView;
 import com.android.launcher3.Utilities;
@@ -71,7 +72,8 @@
 
     private final Context mContext;
     private final FolderIcon mIcon;
-    private final int mIconSize;
+    @VisibleForTesting
+    public final int mIconSize;
 
     // These variables are all associated with the drawing of the preview; they are stored
     // as member variables for shared usage and to avoid computation on each frame
@@ -428,7 +430,8 @@
         p.anim = anim;
     }
 
-    private void setDrawable(PreviewItemDrawingParams p, WorkspaceItemInfo item) {
+    @VisibleForTesting
+    public void setDrawable(PreviewItemDrawingParams p, WorkspaceItemInfo item) {
         if (item.hasPromiseIconUi() || (item.runtimeStatusFlags
                     & ItemInfoWithIcon.FLAG_SHOW_DOWNLOAD_PROGRESS_MASK) != 0) {
             PreloadIconDrawable drawable = newPendingIcon(mContext, item);
diff --git a/src/com/android/launcher3/statemanager/StateManager.java b/src/com/android/launcher3/statemanager/StateManager.java
index 360ff7e..e8b3066 100644
--- a/src/com/android/launcher3/statemanager/StateManager.java
+++ b/src/com/android/launcher3/statemanager/StateManager.java
@@ -19,6 +19,7 @@
 import static android.animation.ValueAnimator.areAnimatorsEnabled;
 
 import static com.android.launcher3.anim.AnimatorPlaybackController.callListenerCommandRecursively;
+import static com.android.launcher3.states.StateAnimationConfig.HANDLE_STATE_APPLY;
 import static com.android.launcher3.states.StateAnimationConfig.SKIP_ALL_ANIMATIONS;
 
 import android.animation.Animator;
@@ -36,10 +37,13 @@
 import com.android.launcher3.anim.PendingAnimation;
 import com.android.launcher3.states.StateAnimationConfig;
 import com.android.launcher3.states.StateAnimationConfig.AnimationFlags;
+import com.android.launcher3.states.StateAnimationConfig.AnimationPropertyFlags;
 import com.android.launcher3.testing.shared.TestProtocol;
 
 import java.io.PrintWriter;
 import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.stream.Collectors;
 
 /**
  * Class to manage transitions between different states for a StatefulActivity based on different
@@ -189,7 +193,7 @@
 
     public void reapplyState(boolean cancelCurrentAnimation) {
         boolean wasInAnimation = mConfig.currentAnimation != null;
-        if (cancelCurrentAnimation) {
+        if (cancelCurrentAnimation && (mConfig.animProps & HANDLE_STATE_APPLY) == 0) {
             // Animation canceling can trigger a cleanup routine, causing problems when we are in a
             // launcher state that relies on member variable data. So if we are in one of those
             // states, accelerate the current animation to its end point rather than canceling it
@@ -227,7 +231,15 @@
 
     private void goToState(
             STATE_TYPE state, boolean animated, long delay, AnimatorListener listener) {
-        Log.d(TestProtocol.OVERVIEW_OVER_HOME, "go to state " + state);
+        String stackTrace = Log.getStackTraceString(new Exception("tracing state transition"));
+        String truncatedTrace =
+                Arrays.stream(stackTrace.split("\\n"))
+                    .limit(5)
+                    .skip(1) // Removes the line "java.lang.Exception: tracing state transition"
+                    .filter(traceLine -> !traceLine.contains("StateManager.goToState"))
+                    .collect(Collectors.joining("\n"));
+        Log.d(TestProtocol.OVERVIEW_OVER_HOME,
+                "go to state " + state + " partial trace:\n" + truncatedTrace);
 
         animated &= areAnimatorsEnabled();
         if (mActivity.isInState(state)) {
@@ -237,7 +249,7 @@
                     listener.onAnimationEnd(null);
                 }
                 return;
-            } else if ((!mConfig.userControlled && animated && mConfig.targetState == state)
+            } else if ((!mConfig.isUserControlled() && animated && mConfig.targetState == state)
                     || mState.shouldPreserveDataStateOnReapply()) {
                 // We are running the same animation as requested, and/or target state should not be
                 // reset -- allow the current animation to complete instead of canceling it.
@@ -343,7 +355,7 @@
 
     public AnimatorPlaybackController createAnimationToNewWorkspace(STATE_TYPE state,
             StateAnimationConfig config) {
-        config.userControlled = true;
+        config.animProps |= StateAnimationConfig.USER_CONTROLLED;
         cancelAnimation();
         config.copyTo(mConfig);
         mConfig.playbackController = createAnimationToNewWorkspaceInternal(state)
@@ -418,7 +430,7 @@
     }
 
     public void moveToRestState(boolean isAnimated) {
-        if (mConfig.currentAnimation != null && mConfig.userControlled) {
+        if (mConfig.currentAnimation != null && mConfig.isUserControlled()) {
             // The user is doing something. Lets not mess it up
             return;
         }
@@ -450,10 +462,18 @@
         }
     }
 
+    /**
+     * Sets the provided controller as the current user controlled state animation
+     */
     public void setCurrentUserControlledAnimation(AnimatorPlaybackController controller) {
+        setCurrentAnimation(controller, StateAnimationConfig.USER_CONTROLLED);
+    }
+
+    public void setCurrentAnimation(AnimatorPlaybackController controller,
+            @AnimationPropertyFlags int animationProps) {
         clearCurrentAnimation();
         setCurrentAnimation(controller.getTarget());
-        mConfig.userControlled = true;
+        mConfig.animProps = animationProps;
         mConfig.playbackController = controller;
     }
 
diff --git a/src/com/android/launcher3/statemanager/StatefulActivity.java b/src/com/android/launcher3/statemanager/StatefulActivity.java
index 520f33c..30ba703 100644
--- a/src/com/android/launcher3/statemanager/StatefulActivity.java
+++ b/src/com/android/launcher3/statemanager/StatefulActivity.java
@@ -126,16 +126,8 @@
 
     @Override
     public void reapplyUi() {
-        reapplyUi(true /* cancelCurrentAnimation */);
-    }
-
-    /**
-     * Re-applies if any state transition is not running, optionally cancelling
-     * the transition if requested.
-     */
-    public void reapplyUi(boolean cancelCurrentAnimation) {
         getRootView().dispatchInsets();
-        getStateManager().reapplyState(cancelCurrentAnimation);
+        getStateManager().reapplyState(true /* cancelCurrentAnimation */);
     }
 
     @Override
diff --git a/src/com/android/launcher3/states/StateAnimationConfig.java b/src/com/android/launcher3/states/StateAnimationConfig.java
index 0d9e010..0ca5afd 100644
--- a/src/com/android/launcher3/states/StateAnimationConfig.java
+++ b/src/com/android/launcher3/states/StateAnimationConfig.java
@@ -40,8 +40,19 @@
     public static final int SKIP_DEPTH_CONTROLLER = 1 << 2;
     public static final int SKIP_SCRIM = 1 << 3;
 
+    @IntDef(flag = true, value = {
+            USER_CONTROLLED,
+            HANDLE_STATE_APPLY
+    })
+    @Retention(RetentionPolicy.SOURCE)
+    public @interface AnimationPropertyFlags {}
+    // Indicates that the animation is controlled by the user
+    public static final int USER_CONTROLLED = 1 << 0;
+    // Indicates that he animation can survive state UI resets due to inset or config changes
+    public static final int HANDLE_STATE_APPLY = 1 << 1;
+
     public long duration;
-    public boolean userControlled;
+    public @AnimationPropertyFlags int animProps = 0;
     public @AnimationFlags int animFlags = 0;
 
 
@@ -105,12 +116,16 @@
     public void copyTo(StateAnimationConfig target) {
         target.duration = duration;
         target.animFlags = animFlags;
-        target.userControlled = userControlled;
+        target.animProps = animProps;
         for (int i = 0; i < ANIM_TYPES_COUNT; i++) {
             target.mInterpolators[i] = mInterpolators[i];
         }
     }
 
+    public boolean isUserControlled() {
+        return (animProps & USER_CONTROLLED) != 0;
+    }
+
     /**
      * Returns the interpolator set for animId or fallback if nothing is set
      *
diff --git a/src/com/android/launcher3/touch/AllAppsSwipeController.java b/src/com/android/launcher3/touch/AllAppsSwipeController.java
index 8b9bc19..fe4a83b 100644
--- a/src/com/android/launcher3/touch/AllAppsSwipeController.java
+++ b/src/com/android/launcher3/touch/AllAppsSwipeController.java
@@ -192,7 +192,7 @@
     protected StateAnimationConfig getConfigForStates(LauncherState fromState,
             LauncherState toState) {
         StateAnimationConfig config = super.getConfigForStates(fromState, toState);
-        config.userControlled = true;
+        config.animProps |= StateAnimationConfig.USER_CONTROLLED;
         if (fromState == NORMAL && toState == ALL_APPS) {
             applyNormalToAllAppsAnimConfig(mLauncher, config);
         } else if (fromState == ALL_APPS && toState == NORMAL) {
@@ -209,13 +209,13 @@
             config.setInterpolator(ANIM_SCRIM_FADE,
                     Interpolators.reverse(ALL_APPS_SCRIM_RESPONDER));
             config.setInterpolator(ANIM_ALL_APPS_FADE, FINAL_FRAME);
-            if (!config.userControlled) {
+            if (!config.isUserControlled()) {
                 config.setInterpolator(ANIM_VERTICAL_PROGRESS, EMPHASIZED);
             }
             config.setInterpolator(ANIM_WORKSPACE_SCALE, DECELERATED_EASE);
             config.setInterpolator(ANIM_DEPTH, DECELERATED_EASE);
         } else {
-            if (config.userControlled) {
+            if (config.isUserControlled()) {
                 config.setInterpolator(ANIM_DEPTH, Interpolators.reverse(BLUR_MANUAL));
                 config.setInterpolator(ANIM_WORKSPACE_FADE,
                         Interpolators.reverse(WORKSPACE_FADE_MANUAL));
@@ -250,29 +250,32 @@
         if (launcher.getDeviceProfile().isTablet) {
             config.setInterpolator(ANIM_ALL_APPS_FADE, INSTANT);
             config.setInterpolator(ANIM_SCRIM_FADE, ALL_APPS_SCRIM_RESPONDER);
-            if (!config.userControlled) {
+            if (!config.isUserControlled()) {
                 config.setInterpolator(ANIM_VERTICAL_PROGRESS, EMPHASIZED);
             }
             config.setInterpolator(ANIM_WORKSPACE_SCALE, DECELERATED_EASE);
             config.setInterpolator(ANIM_DEPTH, DECELERATED_EASE);
         } else {
-            config.setInterpolator(ANIM_DEPTH, config.userControlled ? BLUR_MANUAL : BLUR_ATOMIC);
+            config.setInterpolator(ANIM_DEPTH,
+                    config.isUserControlled() ? BLUR_MANUAL : BLUR_ATOMIC);
             config.setInterpolator(ANIM_WORKSPACE_FADE,
-                    config.userControlled ? WORKSPACE_FADE_MANUAL : WORKSPACE_FADE_ATOMIC);
+                    config.isUserControlled() ? WORKSPACE_FADE_MANUAL : WORKSPACE_FADE_ATOMIC);
             config.setInterpolator(ANIM_WORKSPACE_SCALE,
-                    config.userControlled ? WORKSPACE_SCALE_MANUAL : WORKSPACE_SCALE_ATOMIC);
+                    config.isUserControlled() ? WORKSPACE_SCALE_MANUAL : WORKSPACE_SCALE_ATOMIC);
             config.setInterpolator(ANIM_HOTSEAT_FADE,
-                    config.userControlled ? HOTSEAT_FADE_MANUAL : HOTSEAT_FADE_ATOMIC);
+                    config.isUserControlled() ? HOTSEAT_FADE_MANUAL : HOTSEAT_FADE_ATOMIC);
             config.setInterpolator(ANIM_HOTSEAT_SCALE,
-                    config.userControlled ? HOTSEAT_SCALE_MANUAL : HOTSEAT_SCALE_ATOMIC);
+                    config.isUserControlled() ? HOTSEAT_SCALE_MANUAL : HOTSEAT_SCALE_ATOMIC);
             config.setInterpolator(ANIM_HOTSEAT_TRANSLATE,
-                    config.userControlled ? HOTSEAT_TRANSLATE_MANUAL : HOTSEAT_TRANSLATE_ATOMIC);
+                    config.isUserControlled()
+                            ? HOTSEAT_TRANSLATE_MANUAL
+                            : HOTSEAT_TRANSLATE_ATOMIC);
             config.setInterpolator(ANIM_SCRIM_FADE,
-                    config.userControlled ? SCRIM_FADE_MANUAL : SCRIM_FADE_ATOMIC);
+                    config.isUserControlled() ? SCRIM_FADE_MANUAL : SCRIM_FADE_ATOMIC);
             config.setInterpolator(ANIM_ALL_APPS_FADE,
-                    config.userControlled ? ALL_APPS_FADE_MANUAL : ALL_APPS_FADE_ATOMIC);
+                    config.isUserControlled() ? ALL_APPS_FADE_MANUAL : ALL_APPS_FADE_ATOMIC);
             config.setInterpolator(ANIM_VERTICAL_PROGRESS,
-                    config.userControlled
+                    config.isUserControlled()
                             ? ALL_APPS_VERTICAL_PROGRESS_MANUAL
                             : ALL_APPS_VERTICAL_PROGRESS_ATOMIC);
         }
@@ -285,7 +288,7 @@
      */
     public static void applyOverviewToAllAppsAnimConfig(
             DeviceProfile deviceProfile, StateAnimationConfig config, float threshold) {
-        config.userControlled = true;
+        config.animProps |= StateAnimationConfig.USER_CONTROLLED;
         config.animFlags = SKIP_OVERVIEW;
         if (deviceProfile.isTablet) {
             config.setInterpolator(ANIM_ALL_APPS_FADE, INSTANT);
diff --git a/src/com/android/launcher3/util/CannedAnimationCoordinator.kt b/src/com/android/launcher3/util/CannedAnimationCoordinator.kt
index 18f8339..85f81f5 100644
--- a/src/com/android/launcher3/util/CannedAnimationCoordinator.kt
+++ b/src/com/android/launcher3/util/CannedAnimationCoordinator.kt
@@ -26,6 +26,8 @@
 import com.android.launcher3.anim.AnimatorPlaybackController
 import com.android.launcher3.anim.PendingAnimation
 import com.android.launcher3.statemanager.StatefulActivity
+import com.android.launcher3.states.StateAnimationConfig.HANDLE_STATE_APPLY
+import com.android.launcher3.states.StateAnimationConfig.USER_CONTROLLED
 import java.util.function.Consumer
 
 private const val TAG = "CannedAnimCoordinator"
@@ -107,8 +109,15 @@
         }
         // Link this to the state manager so that it auto-cancels when state changes
         recreatePending = false
+        // Animator coordinator takes care of reapplying the animation due to state reset. Set the
+        // flags accordingly
         animationController =
-            controller.apply { activity.stateManager.setCurrentUserControlledAnimation(this) }
+            controller.apply {
+                activity.stateManager.setCurrentAnimation(
+                    this,
+                    USER_CONTROLLED or HANDLE_STATE_APPLY
+                )
+            }
         recreateAnimation(provider)
     }
 
diff --git a/src/com/android/launcher3/widget/LauncherAppWidgetHost.java b/src/com/android/launcher3/widget/LauncherAppWidgetHost.java
index 739e204..b1c477c 100644
--- a/src/com/android/launcher3/widget/LauncherAppWidgetHost.java
+++ b/src/com/android/launcher3/widget/LauncherAppWidgetHost.java
@@ -21,13 +21,22 @@
 import android.appwidget.AppWidgetHost;
 import android.appwidget.AppWidgetProviderInfo;
 import android.content.Context;
+import android.view.accessibility.AccessibilityNodeInfo;
+import android.widget.RemoteViews;
 
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
 
 import com.android.launcher3.LauncherAppState;
+import com.android.launcher3.util.Executors;
+import com.android.launcher3.util.SafeCloseable;
+import com.android.launcher3.widget.LauncherWidgetHolder.ProviderChangedListener;
 
 import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.Set;
+import java.util.WeakHashMap;
 import java.util.function.IntConsumer;
 
 /**
@@ -37,8 +46,7 @@
  */
 class LauncherAppWidgetHost extends AppWidgetHost {
     @NonNull
-    private final ArrayList<LauncherWidgetHolder.ProviderChangedListener>
-            mProviderChangeListeners = new ArrayList<>();
+    private final List<ProviderChangedListener> mProviderChangeListeners;
 
     @NonNull
     private final Context mContext;
@@ -46,33 +54,13 @@
     @Nullable
     private final IntConsumer mAppWidgetRemovedCallback;
 
-    @NonNull
-    private final LauncherWidgetHolder mHolder;
-
     public LauncherAppWidgetHost(@NonNull Context context,
-            @Nullable IntConsumer appWidgetRemovedCallback, @NonNull LauncherWidgetHolder holder) {
+            @Nullable IntConsumer appWidgetRemovedCallback,
+            List<ProviderChangedListener> providerChangeListeners) {
         super(context, APPWIDGET_HOST_ID);
         mContext = context;
         mAppWidgetRemovedCallback = appWidgetRemovedCallback;
-        mHolder = holder;
-    }
-
-    /**
-     * Add a listener that is triggered when the providers of the widgets are changed
-     * @param listener The listener that notifies when the providers changed
-     */
-    public void addProviderChangeListener(
-            @NonNull LauncherWidgetHolder.ProviderChangedListener listener) {
-        mProviderChangeListeners.add(listener);
-    }
-
-    /**
-     * Remove the specified listener from the host
-     * @param listener The listener that is to be removed from the host
-     */
-    public void removeProviderChangeListener(
-            LauncherWidgetHolder.ProviderChangedListener listener) {
-        mProviderChangeListeners.remove(listener);
+        mProviderChangeListeners = providerChangeListeners;
     }
 
     @Override
@@ -89,7 +77,7 @@
     @NonNull
     public LauncherAppWidgetHostView onCreateView(Context context, int appWidgetId,
             AppWidgetProviderInfo appWidget) {
-        return mHolder.onCreateView(context, appWidgetId);
+        return new ListenableHostView(context);
     }
 
     /**
@@ -115,7 +103,10 @@
         if (mAppWidgetRemovedCallback == null) {
             return;
         }
-        mAppWidgetRemovedCallback.accept(appWidgetId);
+        // Route the call via model thread, in case it comes while a loader-bind is in progress
+        Executors.MODEL_EXECUTOR.execute(
+                () -> Executors.MAIN_EXECUTOR.execute(
+                        () -> mAppWidgetRemovedCallback.accept(appWidgetId)));
     }
 
     /**
@@ -126,4 +117,36 @@
         super.clearViews();
     }
 
+    public static class ListenableHostView extends LauncherAppWidgetHostView {
+
+        private Set<Runnable> mUpdateListeners = Collections.EMPTY_SET;
+
+        ListenableHostView(Context context) {
+            super(context);
+        }
+
+        @Override
+        public void updateAppWidget(RemoteViews remoteViews) {
+            super.updateAppWidget(remoteViews);
+            mUpdateListeners.forEach(Runnable::run);
+        }
+
+        @Override
+        public void onInitializeAccessibilityNodeInfo(AccessibilityNodeInfo info) {
+            super.onInitializeAccessibilityNodeInfo(info);
+            info.setClassName(LauncherAppWidgetHostView.class.getName());
+        }
+
+        /**
+         * Adds a callback to be run everytime the provided app widget updates.
+         * @return a closable to remove this callback
+         */
+        public SafeCloseable addUpdateListener(Runnable callback) {
+            if (mUpdateListeners == Collections.EMPTY_SET) {
+                mUpdateListeners = Collections.newSetFromMap(new WeakHashMap<>());
+            }
+            mUpdateListeners.add(callback);
+            return () -> mUpdateListeners.remove(callback);
+        }
+    }
 }
diff --git a/src/com/android/launcher3/widget/LauncherWidgetHolder.java b/src/com/android/launcher3/widget/LauncherWidgetHolder.java
index fbd48cf..23127b3 100644
--- a/src/com/android/launcher3/widget/LauncherWidgetHolder.java
+++ b/src/com/android/launcher3/widget/LauncherWidgetHolder.java
@@ -42,8 +42,12 @@
 import com.android.launcher3.testing.TestLogging;
 import com.android.launcher3.testing.shared.TestProtocol;
 import com.android.launcher3.util.ResourceBasedOverride;
+import com.android.launcher3.util.SafeCloseable;
+import com.android.launcher3.widget.LauncherAppWidgetHost.ListenableHostView;
 import com.android.launcher3.widget.custom.CustomWidgetManager;
 
+import java.util.ArrayList;
+import java.util.List;
 import java.util.function.IntConsumer;
 
 /**
@@ -61,15 +65,14 @@
             FLAG_STATE_IS_NORMAL | FLAG_ACTIVITY_STARTED | FLAG_ACTIVITY_RESUMED;
 
     @NonNull
-    private final Context mContext;
+    protected final Context mContext;
 
     @NonNull
     private final AppWidgetHost mWidgetHost;
 
     @NonNull
     protected final SparseArray<LauncherAppWidgetHostView> mViews = new SparseArray<>();
-    @NonNull
-    private final SparseArray<PendingAppWidgetHostView> mPendingViews = new SparseArray<>();
+    protected final List<ProviderChangedListener> mProviderChangedListeners = new ArrayList<>();
 
     protected int mFlags = FLAG_STATE_IS_NORMAL;
 
@@ -86,7 +89,8 @@
 
     protected AppWidgetHost createHost(
             Context context, @Nullable IntConsumer appWidgetRemovedCallback) {
-        return new LauncherAppWidgetHost(context, appWidgetRemovedCallback, this);
+        return new LauncherAppWidgetHost(
+                context, appWidgetRemovedCallback, mProviderChangedListeners);
     }
 
     /**
@@ -158,28 +162,6 @@
     }
 
     /**
-     * Add the pending view to the host for complete configuration in further steps
-     * @param appWidgetId The ID of the specified app widget
-     * @param view The {@link PendingAppWidgetHostView} of the app widget
-     */
-    public void addPendingView(int appWidgetId, @NonNull PendingAppWidgetHostView view) {
-        mPendingViews.put(appWidgetId, view);
-    }
-
-    /**
-     * @param appWidgetId The app widget id of the specified widget
-     * @return The {@link PendingAppWidgetHostView} of the widget if it exists, null otherwise
-     */
-    @Nullable
-    protected PendingAppWidgetHostView getPendingView(int appWidgetId) {
-        return mPendingViews.get(appWidgetId);
-    }
-
-    protected void removePendingView(int appWidgetId) {
-        mPendingViews.remove(appWidgetId);
-    }
-
-    /**
      * Called when the launcher is destroyed
      */
     public void destroy() {
@@ -201,18 +183,18 @@
      * Add a listener that is triggered when the providers of the widgets are changed
      * @param listener The listener that notifies when the providers changed
      */
-    public void addProviderChangeListener(@NonNull ProviderChangedListener listener) {
-        LauncherAppWidgetHost tempHost = (LauncherAppWidgetHost) mWidgetHost;
-        tempHost.addProviderChangeListener(listener);
+    public void addProviderChangeListener(
+            @NonNull LauncherWidgetHolder.ProviderChangedListener listener) {
+        MAIN_EXECUTOR.execute(() -> mProviderChangedListeners.add(listener));
     }
 
     /**
      * Remove the specified listener from the host
      * @param listener The listener that is to be removed from the host
      */
-    public void removeProviderChangeListener(ProviderChangedListener listener) {
-        LauncherAppWidgetHost tempHost = (LauncherAppWidgetHost) mWidgetHost;
-        tempHost.removeProviderChangeListener(listener);
+    public void removeProviderChangeListener(
+            LauncherWidgetHolder.ProviderChangedListener listener) {
+        MAIN_EXECUTOR.execute(() -> mProviderChangedListeners.remove(listener));
     }
 
     /**
@@ -316,32 +298,51 @@
     }
 
     /**
+     * Adds a callback to be run everytime the provided app widget updates.
+     * @return a closable to remove this callback
+     */
+    public SafeCloseable addOnUpdateListener(
+            int appWidgetId, LauncherAppWidgetProviderInfo appWidget, Runnable callback) {
+        if (createView(appWidgetId, appWidget) instanceof ListenableHostView lhv) {
+            return lhv.addUpdateListener(callback);
+        }
+        return () -> { };
+    }
+
+    /**
      * Create a view for the specified app widget
-     * @param context The activity context for which the view is created
+     *
      * @param appWidgetId The ID of the widget
-     * @param appWidget The {@link LauncherAppWidgetProviderInfo} of the widget
+     * @param appWidget   The {@link LauncherAppWidgetProviderInfo} of the widget
      * @return A view for the widget
      */
     @NonNull
-    public AppWidgetHostView createView(@NonNull Context context, int appWidgetId,
-            @NonNull LauncherAppWidgetProviderInfo appWidget) {
-
+    public AppWidgetHostView createView(
+            int appWidgetId, @NonNull LauncherAppWidgetProviderInfo appWidget) {
         if (appWidget.isCustomWidget()) {
-            LauncherAppWidgetHostView lahv = new LauncherAppWidgetHostView(context);
+            LauncherAppWidgetHostView lahv = new LauncherAppWidgetHostView(mContext);
             lahv.setAppWidget(0, appWidget);
-            CustomWidgetManager.INSTANCE.get(context).onViewCreated(lahv);
+            CustomWidgetManager.INSTANCE.get(mContext).onViewCreated(lahv);
             return lahv;
-        } else if ((mFlags & FLAG_LISTENING) == 0) {
+        }
+
+        LauncherAppWidgetHostView view = createViewInternal(appWidgetId, appWidget);
+        mViews.put(appWidgetId, view);
+        return view;
+    }
+
+    @NonNull
+    protected LauncherAppWidgetHostView createViewInternal(
+            int appWidgetId, @NonNull LauncherAppWidgetProviderInfo appWidget) {
+        if ((mFlags & FLAG_LISTENING) == 0) {
             // Since the launcher hasn't started listening to widget updates, we can't simply call
-            // super.createView here because the later will make a binder call to retrieve
+            // host.createView here because the later will make a binder call to retrieve
             // RemoteViews from system process.
-            LauncherAppWidgetHostView view =
-                    new PendingAppWidgetHostView(context, appWidgetId, appWidget);
-            mViews.put(appWidgetId, view);
-            return view;
+            return new PendingAppWidgetHostView(mContext, appWidgetId, appWidget);
         } else {
             try {
-                return mWidgetHost.createView(context, appWidgetId, appWidget);
+                return (LauncherAppWidgetHostView) mWidgetHost.createView(
+                        mContext, appWidgetId, appWidget);
             } catch (Exception e) {
                 if (!Utilities.isBinderSizeError(e)) {
                     throw new RuntimeException(e);
@@ -352,7 +353,7 @@
                 // will update.
                 LauncherAppWidgetHostView view = mViews.get(appWidgetId);
                 if (view == null) {
-                    view = onCreateView(mContext, appWidgetId);
+                    view = new ListenableHostView(mContext);
                 }
                 view.setAppWidget(appWidgetId, appWidget);
                 view.switchToErrorView();
@@ -372,26 +373,6 @@
     }
 
     /**
-     * Called to return a proper view when creating a view
-     *
-     * @param context     The context for which the widget view is created
-     * @param appWidgetId The ID of the added widget
-     * @return A view for the specified app widget
-     */
-    @NonNull
-    public LauncherAppWidgetHostView onCreateView(Context context, int appWidgetId) {
-        final LauncherAppWidgetHostView view;
-        if (getPendingView(appWidgetId) != null) {
-            view = getPendingView(appWidgetId);
-            removePendingView(appWidgetId);
-        } else {
-            view = new LauncherAppWidgetHostView(context);
-        }
-        mViews.put(appWidgetId, view);
-        return view;
-    }
-
-    /**
      * Clears all the views from the host
      */
     public void clearViews() {
diff --git a/src/com/android/launcher3/widget/PendingAppWidgetHostView.java b/src/com/android/launcher3/widget/PendingAppWidgetHostView.java
index 2bd4c7e..6ad913e 100644
--- a/src/com/android/launcher3/widget/PendingAppWidgetHostView.java
+++ b/src/com/android/launcher3/widget/PendingAppWidgetHostView.java
@@ -18,6 +18,7 @@
 
 import static com.android.launcher3.graphics.PreloadIconDrawable.newPendingIcon;
 import static com.android.launcher3.icons.FastBitmapDrawable.getDisabledColorFilter;
+import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
 
 import android.content.Context;
 import android.graphics.Canvas;
@@ -38,16 +39,18 @@
 import android.view.View.OnClickListener;
 import android.widget.RemoteViews;
 
+import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
 
 import com.android.launcher3.DeviceProfile;
+import com.android.launcher3.LauncherAppState;
 import com.android.launcher3.R;
 import com.android.launcher3.icons.FastBitmapDrawable;
-import com.android.launcher3.icons.IconCache;
 import com.android.launcher3.icons.IconCache.ItemInfoUpdateReceiver;
 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.SafeCloseable;
 import com.android.launcher3.util.Themes;
 
 import java.util.List;
@@ -64,12 +67,16 @@
     private static final int DEFERRED_ALPHA = 0x77;
 
     private final Rect mRect = new Rect();
-    private OnClickListener mClickListener;
+
+    private final LauncherAppWidgetProviderInfo mAppwidget;
     private final LauncherAppWidgetInfo mInfo;
     private final int mStartState;
     private final boolean mDisabledForSafeMode;
     private final CharSequence mLabel;
 
+    private OnClickListener mClickListener;
+    private SafeCloseable mOnDetachCleanup;
+
     private int mDragFlags;
 
     private Drawable mCenterDrawable;
@@ -81,8 +88,8 @@
     private Layout mSetupTextLayout;
 
     public PendingAppWidgetHostView(Context context, LauncherAppWidgetInfo info,
-            IconCache cache, boolean disabledForSafeMode) {
-        this(context, info, disabledForSafeMode,
+            @Nullable LauncherAppWidgetProviderInfo appWidget, boolean disabledForSafeMode) {
+        this(context, info, disabledForSafeMode, appWidget,
                 context.getResources().getText(R.string.gadget_complete_setup_text));
 
         super.updateAppWidget(null);
@@ -91,16 +98,17 @@
         if (info.pendingItemInfo == null) {
             info.pendingItemInfo = new PackageItemInfo(info.providerName.getPackageName(),
                     info.user);
-            cache.updateIconInBackground(this, info.pendingItemInfo);
+            LauncherAppState.getInstance(context).getIconCache()
+                    .updateIconInBackground(this, info.pendingItemInfo);
         } else {
             reapplyItemInfo(info.pendingItemInfo);
         }
     }
 
     public PendingAppWidgetHostView(
-            Context context, int appWidgetId, LauncherAppWidgetProviderInfo appWidget) {
+            Context context, int appWidgetId, @NonNull LauncherAppWidgetProviderInfo appWidget) {
         this(context, new LauncherAppWidgetInfo(appWidgetId, appWidget.provider), false,
-                appWidget.label);
+                appWidget, appWidget.label);
         getBackground().mutate().setAlpha(DEFERRED_ALPHA);
 
         mCenterDrawable = new ColorDrawable(Color.TRANSPARENT);
@@ -109,9 +117,11 @@
     }
 
     private PendingAppWidgetHostView(Context context, LauncherAppWidgetInfo info,
-            boolean disabledForSafeMode, CharSequence label) {
+            boolean disabledForSafeMode, LauncherAppWidgetProviderInfo appwidget,
+            CharSequence label) {
         super(new ContextThemeWrapper(context, R.style.WidgetContainerTheme));
 
+        mAppwidget = appwidget;
         mInfo = info;
         mStartState = info.restoreStatus;
         mDisabledForSafeMode = disabledForSafeMode;
@@ -128,9 +138,40 @@
 
     @Override
     public void updateAppWidget(RemoteViews remoteViews) {
+        checkIfRestored();
+    }
+
+    private void checkIfRestored() {
         WidgetManagerHelper widgetManagerHelper = new WidgetManagerHelper(getContext());
         if (widgetManagerHelper.isAppWidgetRestored(mInfo.appWidgetId)) {
-            reInflate();
+            MAIN_EXECUTOR.getHandler().post(this::reInflate);
+        }
+    }
+
+    @Override
+    protected void onAttachedToWindow() {
+        super.onAttachedToWindow();
+
+        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 = mLauncher.getAppWidgetHolder()
+                    .addOnUpdateListener(mInfo.appWidgetId, mAppwidget, this::checkIfRestored);
+            checkIfRestored();
+        }
+    }
+
+    @Override
+    protected void onDetachedFromWindow() {
+        super.onDetachedFromWindow();
+        if (mOnDetachCleanup != null) {
+            mOnDetachCleanup.close();
+            mOnDetachCleanup = null;
         }
     }
 
diff --git a/src/com/android/launcher3/widget/WidgetHostViewLoader.java b/src/com/android/launcher3/widget/WidgetHostViewLoader.java
index b18cd47..1cc00ef 100644
--- a/src/com/android/launcher3/widget/WidgetHostViewLoader.java
+++ b/src/com/android/launcher3/widget/WidgetHostViewLoader.java
@@ -1,7 +1,6 @@
 package com.android.launcher3.widget;
 
 import android.appwidget.AppWidgetHostView;
-import android.content.Context;
 import android.os.Bundle;
 import android.os.Handler;
 import android.util.Log;
@@ -117,7 +116,7 @@
                     return;
                 }
                 AppWidgetHostView hostView = mLauncher.getAppWidgetHolder().createView(
-                        (Context) mLauncher, mWidgetLoadingId, pInfo);
+                        mWidgetLoadingId, pInfo);
                 mInfo.boundWidget = hostView;
 
                 // We used up the widget Id in binding the above view.
diff --git a/tests/src/com/android/launcher3/folder/PreviewItemManagerTest.kt b/tests/src/com/android/launcher3/folder/PreviewItemManagerTest.kt
new file mode 100644
index 0000000..4ec5b0e
--- /dev/null
+++ b/tests/src/com/android/launcher3/folder/PreviewItemManagerTest.kt
@@ -0,0 +1,217 @@
+/*
+ * Copyright (C) 2023 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.folder
+
+import android.R
+import android.content.Context
+import android.os.Process
+import androidx.test.core.app.ApplicationProvider
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import androidx.test.platform.app.InstrumentationRegistry.getInstrumentation
+import com.android.launcher3.LauncherPrefs
+import com.android.launcher3.LauncherPrefs.Companion.get
+import com.android.launcher3.icons.BaseIconFactory
+import com.android.launcher3.icons.FastBitmapDrawable
+import com.android.launcher3.icons.UserBadgeDrawable
+import com.android.launcher3.model.data.WorkspaceItemInfo
+import com.android.launcher3.util.ActivityContextWrapper
+import com.android.launcher3.util.FlagOp
+import com.android.launcher3.util.LauncherLayoutBuilder
+import com.android.launcher3.util.LauncherModelHelper
+import com.android.launcher3.util.UserIconInfo
+import org.junit.After
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+
+/** Tests for [PreviewItemManager] */
+@SmallTest
+@RunWith(AndroidJUnit4::class)
+class PreviewItemManagerTest {
+
+    private lateinit var previewItemManager: PreviewItemManager
+    private lateinit var context: Context
+    private lateinit var folderItems: ArrayList<WorkspaceItemInfo>
+    private lateinit var modelHelper: LauncherModelHelper
+    private lateinit var folderIcon: FolderIcon
+
+    @Before
+    fun setup() {
+        getInstrumentation().runOnMainSync {
+            folderIcon =
+                FolderIcon(ActivityContextWrapper(ApplicationProvider.getApplicationContext()))
+        }
+        context = getInstrumentation().targetContext
+        previewItemManager = PreviewItemManager(folderIcon)
+        modelHelper = LauncherModelHelper()
+        modelHelper
+            .setupDefaultLayoutProvider(
+                LauncherLayoutBuilder()
+                    .atWorkspace(0, 0, 1)
+                    .putFolder(R.string.copy)
+                    .addApp(LauncherModelHelper.TEST_PACKAGE, LauncherModelHelper.TEST_ACTIVITY)
+                    .addApp(LauncherModelHelper.TEST_PACKAGE, LauncherModelHelper.TEST_ACTIVITY2)
+                    .addApp(LauncherModelHelper.TEST_PACKAGE, LauncherModelHelper.TEST_ACTIVITY3)
+                    .addApp(LauncherModelHelper.TEST_PACKAGE, LauncherModelHelper.TEST_ACTIVITY4)
+                    .build()
+            )
+            .loadModelSync()
+        folderItems = modelHelper.bgDataModel.folders.valueAt(0).contents
+        folderIcon.mInfo = modelHelper.bgDataModel.folders.valueAt(0)
+        folderIcon.mInfo.contents = folderItems
+
+        // Set first icon to be themed.
+        folderItems[0]
+            .bitmap
+            .setMonoIcon(
+                folderItems[0].bitmap.icon,
+                BaseIconFactory(
+                    context,
+                    context.resources.configuration.densityDpi,
+                    previewItemManager.mIconSize
+                )
+            )
+
+        // Set second icon to be non-themed.
+        folderItems[1]
+            .bitmap
+            .setMonoIcon(
+                null,
+                BaseIconFactory(
+                    context,
+                    context.resources.configuration.densityDpi,
+                    previewItemManager.mIconSize
+                )
+            )
+
+        // Set third icon to be themed with badge.
+        folderItems[2]
+            .bitmap
+            .setMonoIcon(
+                folderItems[2].bitmap.icon,
+                BaseIconFactory(
+                    context,
+                    context.resources.configuration.densityDpi,
+                    previewItemManager.mIconSize
+                )
+            )
+        folderItems[2].bitmap =
+            folderItems[2].bitmap.withFlags(profileFlagOp(UserIconInfo.TYPE_WORK))
+
+        // Set fourth icon to be non-themed with badge.
+        folderItems[3].bitmap =
+            folderItems[3].bitmap.withFlags(profileFlagOp(UserIconInfo.TYPE_WORK))
+        folderItems[3]
+            .bitmap
+            .setMonoIcon(
+                null,
+                BaseIconFactory(
+                    context,
+                    context.resources.configuration.densityDpi,
+                    previewItemManager.mIconSize
+                )
+            )
+    }
+    @After
+    @Throws(Exception::class)
+    fun tearDown() {
+        modelHelper.destroy()
+    }
+
+    @Test
+    fun checkThemedIconWithThemingOn_iconShouldBeThemed() {
+        get(context).put(LauncherPrefs.THEMED_ICONS, true)
+        val drawingParams = PreviewItemDrawingParams(0f, 0f, 0f)
+
+        previewItemManager.setDrawable(drawingParams, folderItems[0])
+
+        assert((drawingParams.drawable as FastBitmapDrawable).isThemed)
+    }
+
+    @Test
+    fun checkThemedIconWithThemingOff_iconShouldNotBeThemed() {
+        get(context).put(LauncherPrefs.THEMED_ICONS, false)
+        val drawingParams = PreviewItemDrawingParams(0f, 0f, 0f)
+
+        previewItemManager.setDrawable(drawingParams, folderItems[0])
+
+        assert(!(drawingParams.drawable as FastBitmapDrawable).isThemed)
+    }
+
+    @Test
+    fun checkUnthemedIconWithThemingOn_iconShouldNotBeThemed() {
+        get(context).put(LauncherPrefs.THEMED_ICONS, true)
+        val drawingParams = PreviewItemDrawingParams(0f, 0f, 0f)
+
+        previewItemManager.setDrawable(drawingParams, folderItems[1])
+
+        assert(!(drawingParams.drawable as FastBitmapDrawable).isThemed)
+    }
+
+    @Test
+    fun checkUnthemedIconWithThemingOff_iconShouldNotBeThemed() {
+        get(context).put(LauncherPrefs.THEMED_ICONS, false)
+        val drawingParams = PreviewItemDrawingParams(0f, 0f, 0f)
+
+        previewItemManager.setDrawable(drawingParams, folderItems[1])
+
+        assert(!(drawingParams.drawable as FastBitmapDrawable).isThemed)
+    }
+
+    @Test
+    fun checkThemedIconWithBadgeWithThemingOn_iconAndBadgeShouldBeThemed() {
+        get(context).put(LauncherPrefs.THEMED_ICONS, true)
+        val drawingParams = PreviewItemDrawingParams(0f, 0f, 0f)
+
+        previewItemManager.setDrawable(drawingParams, folderItems[2])
+
+        assert((drawingParams.drawable as FastBitmapDrawable).isThemed)
+        assert(
+            ((drawingParams.drawable as FastBitmapDrawable).badge as UserBadgeDrawable).mIsThemed
+        )
+    }
+
+    @Test
+    fun checkUnthemedIconWithBadgeWithThemingOn_badgeShouldBeThemed() {
+        get(context).put(LauncherPrefs.THEMED_ICONS, true)
+        val drawingParams = PreviewItemDrawingParams(0f, 0f, 0f)
+
+        previewItemManager.setDrawable(drawingParams, folderItems[3])
+
+        assert(!(drawingParams.drawable as FastBitmapDrawable).isThemed)
+        assert(
+            ((drawingParams.drawable as FastBitmapDrawable).badge as UserBadgeDrawable).mIsThemed
+        )
+    }
+
+    @Test
+    fun checkUnthemedIconWithBadgeWithThemingOff_iconAndBadgeShouldNotBeThemed() {
+        get(context).put(LauncherPrefs.THEMED_ICONS, false)
+        val drawingParams = PreviewItemDrawingParams(0f, 0f, 0f)
+
+        previewItemManager.setDrawable(drawingParams, folderItems[3])
+
+        assert(!(drawingParams.drawable as FastBitmapDrawable).isThemed)
+        assert(
+            !((drawingParams.drawable as FastBitmapDrawable).badge as UserBadgeDrawable).mIsThemed
+        )
+    }
+
+    private fun profileFlagOp(type: Int) =
+        UserIconInfo(Process.myUserHandle(), type).applyBitmapInfoFlags(FlagOp.NO_OP)
+}
diff --git a/tests/tapl/com/android/launcher3/tapl/BaseOverview.java b/tests/tapl/com/android/launcher3/tapl/BaseOverview.java
index 2013ce4..1a219a5 100644
--- a/tests/tapl/com/android/launcher3/tapl/BaseOverview.java
+++ b/tests/tapl/com/android/launcher3/tapl/BaseOverview.java
@@ -208,11 +208,15 @@
     public void touchTaskbarBottomCorner(boolean tapRight) {
         try (LauncherInstrumentation.Closable e = mLauncher.eventsCheck()) {
             Taskbar taskbar = new Taskbar(mLauncher);
-            taskbar.touchBottomCorner(tapRight);
             if (mLauncher.isTransientTaskbar()) {
+                mLauncher.runToState(
+                        () -> taskbar.touchBottomCorner(tapRight),
+                        NORMAL_STATE_ORDINAL,
+                        "touching taskbar");
                 // Tapping outside Transient Taskbar returns to Workspace, wait for that state.
                 new Workspace(mLauncher);
             } else {
+                taskbar.touchBottomCorner(tapRight);
                 // Should stay in Overview.
                 verifyActiveContainer();
                 verifyActionsViewVisibility();
diff --git a/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java b/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
index d17f034..68bf298 100644
--- a/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
+++ b/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
@@ -1027,11 +1027,11 @@
             return;
         }
 
-        linearGesture(
-                displaySize.x / 2, displaySize.y - 1,
-                displaySize.x / 2, 0,
-                ZERO_BUTTON_STEPS_FROM_BACKGROUND_TO_HOME,
-                false, GestureScope.EXPECT_PILFER);
+        if (isLauncher3()) {
+            gestureToDismissPopup(displaySize);
+        } else {
+            runToState(() -> gestureToDismissPopup(displaySize), NORMAL_STATE_ORDINAL, "swiping");
+        }
 
         try (LauncherInstrumentation.Closable c1 = addContextLayer(
                 String.format("Swiped up from floating view %s to home", floatingRes.get()))) {
@@ -1040,6 +1040,14 @@
         }
     }
 
+    private void gestureToDismissPopup(Point displaySize) {
+        linearGesture(
+                displaySize.x / 2, displaySize.y - 1,
+                displaySize.x / 2, 0,
+                ZERO_BUTTON_STEPS_FROM_BACKGROUND_TO_HOME,
+                false, GestureScope.EXPECT_PILFER);
+    }
+
     /**
      * @return the Workspace object.
      * @deprecated use goHome().
diff --git a/tests/tapl/com/android/launcher3/tapl/SelectModeButtons.java b/tests/tapl/com/android/launcher3/tapl/SelectModeButtons.java
index 146a67c..e2bc17b 100644
--- a/tests/tapl/com/android/launcher3/tapl/SelectModeButtons.java
+++ b/tests/tapl/com/android/launcher3/tapl/SelectModeButtons.java
@@ -18,6 +18,8 @@
 
 import static android.view.KeyEvent.KEYCODE_ESCAPE;
 
+import static com.android.launcher3.testing.shared.TestProtocol.OVERVIEW_STATE_ORDINAL;
+
 import androidx.annotation.NonNull;
 import androidx.test.uiautomator.UiObject2;
 
@@ -69,7 +71,10 @@
         try (LauncherInstrumentation.Closable e = mLauncher.eventsCheck()) {
             mLauncher.expectEvent(TestProtocol.SEQUENCE_MAIN, EVENT_ALT_ESC_DOWN);
             mLauncher.expectEvent(TestProtocol.SEQUENCE_MAIN, EVENT_ALT_ESC_UP);
-            mLauncher.getDevice().pressKeyCode(KEYCODE_ESCAPE);
+            mLauncher.runToState(
+                    () -> mLauncher.getDevice().pressKeyCode(KEYCODE_ESCAPE),
+                    OVERVIEW_STATE_ORDINAL,
+                    "pressing Esc");
             try (LauncherInstrumentation.Closable c = mLauncher.addContextLayer(
                     "pressed esc key")) {
                 return new Overview(mLauncher);