diff --git a/quickstep/res/layout/taskbar.xml b/quickstep/res/layout/taskbar.xml
index b4c168c..83ad9f3 100644
--- a/quickstep/res/layout/taskbar.xml
+++ b/quickstep/res/layout/taskbar.xml
@@ -52,6 +52,7 @@
             android:orientation="horizontal"
             android:paddingLeft="@dimen/taskbar_nav_buttons_spacing"
             android:paddingRight="@dimen/taskbar_nav_buttons_spacing"
+            android:layout_marginEnd="@dimen/taskbar_contextual_button_margin"
             android:gravity="center_vertical"
             android:layout_gravity="end"/>
 
diff --git a/quickstep/res/layout/taskbar_contextual_button.xml b/quickstep/res/layout/taskbar_contextual_button.xml
new file mode 100644
index 0000000..cbbbfab
--- /dev/null
+++ b/quickstep/res/layout/taskbar_contextual_button.xml
@@ -0,0 +1,23 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!-- Copyright (C) 2021 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.
+-->
+<ImageView
+    xmlns:android="http://schemas.android.com/apk/res/android"
+    android:layout_width="@dimen/taskbar_contextual_buttons_size"
+    android:layout_height="@dimen/taskbar_contextual_buttons_size"
+    android:layout_marginTop="@dimen/taskbar_contextual_button_margin"
+    android:paddingStart="@dimen/taskbar_nav_buttons_spacing"
+    android:background="@drawable/taskbar_icon_click_feedback_roundrect"
+    android:scaleType="center"/>
\ No newline at end of file
diff --git a/quickstep/res/layout/taskbar_nav_button.xml b/quickstep/res/layout/taskbar_nav_button.xml
index 985f928..4ffb8d8 100644
--- a/quickstep/res/layout/taskbar_nav_button.xml
+++ b/quickstep/res/layout/taskbar_nav_button.xml
@@ -1,4 +1,18 @@
 <?xml version="1.0" encoding="utf-8"?>
+<!-- Copyright (C) 2021 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.
+-->
 <ImageView
     xmlns:android="http://schemas.android.com/apk/res/android"
     android:layout_width="@dimen/taskbar_nav_buttons_size"
diff --git a/quickstep/res/values/dimens.xml b/quickstep/res/values/dimens.xml
index e7de0ac..c649082 100644
--- a/quickstep/res/values/dimens.xml
+++ b/quickstep/res/values/dimens.xml
@@ -163,6 +163,8 @@
     <dimen name="taskbar_folder_margin">16dp</dimen>
     <dimen name="taskbar_nav_buttons_spacing">16dp</dimen>
     <dimen name="taskbar_nav_buttons_size">48dp</dimen>
+    <dimen name="taskbar_contextual_button_margin">16dp</dimen>
+    <dimen name="taskbar_contextual_buttons_size">35dp</dimen>
     <dimen name="taskbar_stashed_size">24dp</dimen>
     <dimen name="taskbar_stashed_handle_width">220dp</dimen>
     <dimen name="taskbar_stashed_handle_height">6dp</dimen>
diff --git a/quickstep/robolectric_tests/src/com/android/quickstep/util/TaskViewSimulatorTest.java b/quickstep/robolectric_tests/src/com/android/quickstep/util/TaskViewSimulatorTest.java
deleted file mode 100644
index fd93d98..0000000
--- a/quickstep/robolectric_tests/src/com/android/quickstep/util/TaskViewSimulatorTest.java
+++ /dev/null
@@ -1,206 +0,0 @@
-/*
- * Copyright (C) 2020 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package com.android.quickstep.util;
-
-import static android.view.Display.DEFAULT_DISPLAY;
-
-import static org.mockito.Mockito.mock;
-
-import android.content.Context;
-import android.graphics.Rect;
-import android.graphics.RectF;
-import android.hardware.display.DisplayManager;
-import android.view.Surface;
-import android.view.SurfaceControl;
-
-import com.android.launcher3.DeviceProfile;
-import com.android.launcher3.InvariantDeviceProfile;
-import com.android.launcher3.shadows.LShadowDisplay;
-import com.android.launcher3.util.DisplayController;
-import com.android.quickstep.LauncherActivityInterface;
-import com.android.systemui.shared.system.RemoteAnimationTargetCompat;
-import com.android.systemui.shared.system.SyncRtSurfaceTransactionApplierCompat.SurfaceParams;
-
-import org.hamcrest.Description;
-import org.hamcrest.TypeSafeMatcher;
-import org.junit.Assert;
-import org.junit.Test;
-import org.junit.runner.RunWith;
-import org.robolectric.RobolectricTestRunner;
-import org.robolectric.RuntimeEnvironment;
-import org.robolectric.annotation.LooperMode;
-import org.robolectric.annotation.LooperMode.Mode;
-import org.robolectric.shadow.api.Shadow;
-import org.robolectric.shadows.ShadowDisplayManager;
-
-@RunWith(RobolectricTestRunner.class)
-@LooperMode(Mode.PAUSED)
-public class TaskViewSimulatorTest {
-
-    @Test
-    public void taskProperlyScaled_portrait_noRotation_sameInsets1() {
-        new TaskMatrixVerifier()
-                .withLauncherSize(1200, 2450)
-                .withInsets(new Rect(0, 80, 0, 120))
-                .verifyNoTransforms();
-    }
-
-    @Test
-    public void taskProperlyScaled_portrait_noRotation_sameInsets2() {
-        new TaskMatrixVerifier()
-                .withLauncherSize(1200, 2450)
-                .withInsets(new Rect(55, 80, 55, 120))
-                .verifyNoTransforms();
-    }
-
-    @Test
-    public void taskProperlyScaled_landscape_noRotation_sameInsets1() {
-        new TaskMatrixVerifier()
-                .withLauncherSize(2450, 1250)
-                .withInsets(new Rect(0, 80, 0, 40))
-                .verifyNoTransforms();
-    }
-
-    @Test
-    public void taskProperlyScaled_landscape_noRotation_sameInsets2() {
-        new TaskMatrixVerifier()
-                .withLauncherSize(2450, 1250)
-                .withInsets(new Rect(0, 80, 120, 0))
-                .verifyNoTransforms();
-    }
-
-    @Test
-    public void taskProperlyScaled_landscape_noRotation_sameInsets3() {
-        new TaskMatrixVerifier()
-                .withLauncherSize(2450, 1250)
-                .withInsets(new Rect(55, 80, 55, 120))
-                .verifyNoTransforms();
-    }
-
-    @Test
-    public void taskProperlyScaled_landscape_rotated() {
-        new TaskMatrixVerifier()
-                .withLauncherSize(1200, 2450)
-                .withInsets(new Rect(0, 80, 0, 120))
-                .withAppBounds(
-                        new Rect(0, 0, 2450, 1200),
-                        new Rect(0, 80, 0, 120),
-                        Surface.ROTATION_90)
-                .verifyNoTransforms();
-    }
-
-    private static class TaskMatrixVerifier extends TransformParams {
-
-        private final Context mContext = RuntimeEnvironment.application;
-
-        private Rect mAppBounds = new Rect();
-        private Rect mLauncherInsets = new Rect();
-
-        private Rect mAppInsets;
-
-        private int mAppRotation = -1;
-        private DeviceProfile mDeviceProfile;
-
-        TaskMatrixVerifier withLauncherSize(int width, int height) {
-            ShadowDisplayManager.changeDisplay(DEFAULT_DISPLAY,
-                    String.format("w%sdp-h%sdp-mdpi", width, height));
-            if (mAppBounds.isEmpty()) {
-                mAppBounds.set(0, 0, width, height);
-            }
-            return this;
-        }
-
-        TaskMatrixVerifier withInsets(Rect insets) {
-            LShadowDisplay shadowDisplay = Shadow.extract(
-                    mContext.getSystemService(DisplayManager.class).getDisplay(DEFAULT_DISPLAY));
-            shadowDisplay.setInsets(insets);
-            mLauncherInsets.set(insets);
-            return this;
-        }
-
-        TaskMatrixVerifier withAppBounds(Rect bounds, Rect insets, int appRotation) {
-            mAppBounds.set(bounds);
-            mAppInsets = insets;
-            mAppRotation = appRotation;
-            return this;
-        }
-
-        void verifyNoTransforms() {
-            mDeviceProfile = InvariantDeviceProfile.INSTANCE.get(mContext)
-                    .getDeviceProfile(mContext);
-            mDeviceProfile.updateInsets(mLauncherInsets);
-
-            TaskViewSimulator tvs = new TaskViewSimulator(mContext,
-                    LauncherActivityInterface.INSTANCE);
-            tvs.setDp(mDeviceProfile);
-
-            int launcherRotation = DisplayController.INSTANCE.get(mContext).getInfo().rotation;
-            if (mAppRotation < 0) {
-                mAppRotation = launcherRotation;
-            }
-            tvs.getOrientationState().update(launcherRotation, mAppRotation);
-            if (mAppInsets == null) {
-                mAppInsets = new Rect(mLauncherInsets);
-            }
-            tvs.setPreviewBounds(mAppBounds, mAppInsets);
-
-            tvs.fullScreenProgress.value = 1;
-            tvs.recentsViewScale.value = tvs.getFullScreenScale();
-            tvs.apply(this);
-        }
-
-        @Override
-        public SurfaceParams[] createSurfaceParams(BuilderProxy proxy) {
-            SurfaceParams.Builder builder = new SurfaceParams.Builder((SurfaceControl) null);
-            proxy.onBuildTargetParams(builder, mock(RemoteAnimationTargetCompat.class), this);
-            return new SurfaceParams[] {builder.build()};
-        }
-
-        @Override
-        public void applySurfaceParams(SurfaceParams[] params) {
-            // Verify that the task position remains the same
-            RectF newAppBounds = new RectF(mAppBounds);
-            params[0].matrix.mapRect(newAppBounds);
-            Assert.assertThat(newAppBounds, new AlmostSame(mAppBounds));
-
-            System.err.println("Bounds mapped: " + mAppBounds + " => " + newAppBounds);
-        }
-    }
-
-    private static class AlmostSame extends TypeSafeMatcher<RectF>  {
-
-        // Allow 1px error margin to account for float to int conversions
-        private final float mError = 1f;
-        private final Rect mExpected;
-
-        AlmostSame(Rect expected) {
-            mExpected = expected;
-        }
-
-        @Override
-        protected boolean matchesSafely(RectF item) {
-            return Math.abs(item.left - mExpected.left) < mError
-                    && Math.abs(item.top - mExpected.top) < mError
-                    && Math.abs(item.right - mExpected.right) < mError
-                    && Math.abs(item.bottom - mExpected.bottom) < mError;
-        }
-
-        @Override
-        public void describeTo(Description description) {
-            description.appendValue(mExpected);
-        }
-    }
-}
diff --git a/quickstep/src/com/android/launcher3/taskbar/LauncherTaskbarUIController.java b/quickstep/src/com/android/launcher3/taskbar/LauncherTaskbarUIController.java
index b172095..764b0d3 100644
--- a/quickstep/src/com/android/launcher3/taskbar/LauncherTaskbarUIController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/LauncherTaskbarUIController.java
@@ -33,6 +33,7 @@
 import com.android.launcher3.LauncherState;
 import com.android.launcher3.QuickstepTransitionManager;
 import com.android.launcher3.R;
+import com.android.launcher3.Utilities;
 import com.android.launcher3.anim.AnimatorListeners;
 import com.android.launcher3.config.FeatureFlags;
 import com.android.launcher3.model.data.ItemInfoWithIcon;
@@ -285,6 +286,7 @@
      */
     public void showEdu() {
         if (!FeatureFlags.ENABLE_TASKBAR_EDU.get()
+                || Utilities.IS_RUNNING_IN_TEST_HARNESS
                 || mLauncher.getOnboardingPrefs().getBoolean(OnboardingPrefs.TASKBAR_EDU_SEEN)) {
             return;
         }
diff --git a/quickstep/src/com/android/launcher3/taskbar/NavbarButtonsViewController.java b/quickstep/src/com/android/launcher3/taskbar/NavbarButtonsViewController.java
index 046ee6f..79bea03 100644
--- a/quickstep/src/com/android/launcher3/taskbar/NavbarButtonsViewController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/NavbarButtonsViewController.java
@@ -34,6 +34,7 @@
 import android.animation.ObjectAnimator;
 import android.annotation.DrawableRes;
 import android.annotation.IdRes;
+import android.annotation.LayoutRes;
 import android.graphics.Rect;
 import android.graphics.Region;
 import android.graphics.Region.Op;
@@ -123,21 +124,16 @@
                         .getProperty(ALPHA_INDEX_IME),
                 flags -> (flags & FLAG_IME_VISIBLE) == 0, MultiValueAlpha.VALUE, 1, 0));
 
+        boolean isThreeButtonNav = mContext.isThreeButtonNav();
         // IME switcher
         View imeSwitcherButton = addButton(R.drawable.ic_ime_switcher, BUTTON_IME_SWITCH,
-                mEndContextualContainer, mControllers.navButtonController, R.id.ime_switcher);
+                isThreeButtonNav ? mStartContextualContainer : mEndContextualContainer,
+                mControllers.navButtonController, R.id.ime_switcher);
         mPropertyHolders.add(new StatePropertyHolder(imeSwitcherButton,
                 flags -> ((flags & MASK_IME_SWITCHER_VISIBLE) == MASK_IME_SWITCHER_VISIBLE)
                         && ((flags & FLAG_ROTATION_BUTTON_VISIBLE) == 0)
                         && ((flags & FLAG_A11Y_VISIBLE) == 0)));
 
-        View imeDownButton = addButton(R.drawable.ic_sysbar_back, BUTTON_BACK,
-                mStartContextualContainer, mControllers.navButtonController, R.id.back);
-        imeDownButton.setRotation(Utilities.isRtl(mContext.getResources()) ? 90 : -90);
-        // Rotate when Ime visible
-        mPropertyHolders.add(new StatePropertyHolder(imeDownButton,
-                flags -> (flags & FLAG_IME_VISIBLE) != 0));
-
         mPropertyHolders.add(new StatePropertyHolder(
                 mControllers.taskbarViewController.getTaskbarIconAlpha()
                         .getProperty(ALPHA_INDEX_KEYGUARD),
@@ -150,7 +146,7 @@
         // Force nav buttons (specifically back button) to be visible during setup wizard.
         boolean areButtonsForcedVisible = !SettingsCache.INSTANCE.get(mContext).getValue(
                 Settings.Secure.getUriFor(Settings.Secure.USER_SETUP_COMPLETE), 0);
-        if (mContext.isThreeButtonNav() || areButtonsForcedVisible) {
+        if (isThreeButtonNav || areButtonsForcedVisible) {
             initButtons(mNavButtonContainer, mEndContextualContainer,
                     mControllers.navButtonController);
 
@@ -163,11 +159,18 @@
 
             // Rotation button
             RotationButton rotationButton = new RotationButtonImpl(
-                    addButton(mEndContextualContainer, R.id.rotate_suggestion));
+                    addButton(mEndContextualContainer, R.id.rotate_suggestion,
+                            R.layout.taskbar_contextual_button));
             rotationButton.hide();
             mControllers.rotationButtonController.setRotationButton(rotationButton);
         } else {
             mControllers.rotationButtonController.setRotationButton(new RotationButton() {});
+            View imeDownButton = addButton(R.drawable.ic_sysbar_back, BUTTON_BACK,
+                    mStartContextualContainer, mControllers.navButtonController, R.id.back);
+            imeDownButton.setRotation(Utilities.isRtl(mContext.getResources()) ? 90 : -90);
+            // Rotate when Ime visible
+            mPropertyHolders.add(new StatePropertyHolder(imeDownButton,
+                    flags -> (flags & FLAG_IME_VISIBLE) != 0));
         }
 
         applyState();
@@ -180,8 +183,11 @@
         mBackButton = addButton(R.drawable.ic_sysbar_back, BUTTON_BACK,
                 mNavButtonContainer, mControllers.navButtonController, R.id.back);
         mPropertyHolders.add(new StatePropertyHolder(mBackButton,
-                flags -> (flags & FLAG_IME_VISIBLE) == 0 &&
-                        (flags & FLAG_DISABLE_BACK) == 0));
+                flags -> (flags & FLAG_DISABLE_BACK) == 0));
+        boolean isRtl = Utilities.isRtl(mContext.getResources());
+        mPropertyHolders.add(new StatePropertyHolder(
+                mBackButton, flags -> (flags & FLAG_IME_VISIBLE) != 0, View.ROTATION,
+                isRtl ? 90 : -90, 0));
         // Hide when keyguard is showing, show when bouncer is showing
         mPropertyHolders.add(new StatePropertyHolder(mBackButton,
                 flags -> (flags & FLAG_KEYGUARD_VISIBLE) == 0 ||
@@ -191,19 +197,18 @@
         View homeButton = addButton(R.drawable.ic_sysbar_home, BUTTON_HOME, navContainer,
                 navButtonController, R.id.home);
         mPropertyHolders.add(new StatePropertyHolder(homeButton,
-                flags -> (flags & FLAG_IME_VISIBLE) == 0 &&
-                        (flags & FLAG_KEYGUARD_VISIBLE) == 0 &&
+                flags -> (flags & FLAG_KEYGUARD_VISIBLE) == 0 &&
                         (flags & FLAG_DISABLE_HOME) == 0));
         View recentsButton = addButton(R.drawable.ic_sysbar_recent, BUTTON_RECENTS,
                 navContainer, navButtonController, R.id.recent_apps);
         mPropertyHolders.add(new StatePropertyHolder(recentsButton,
-                flags -> (flags & FLAG_IME_VISIBLE) == 0 &&
-                        (flags & FLAG_KEYGUARD_VISIBLE) == 0 &&
+                flags -> (flags & FLAG_KEYGUARD_VISIBLE) == 0 &&
                         (flags & FLAG_DISABLE_RECENTS) == 0));
 
         // A11y button
         mA11yButton = addButton(R.drawable.ic_sysbar_accessibility_button, BUTTON_A11Y,
-                endContainer, navButtonController, R.id.accessibility_button);
+                endContainer, navButtonController, R.id.accessibility_button,
+                R.layout.taskbar_contextual_button);
         mPropertyHolders.add(new StatePropertyHolder(mA11yButton,
                 flags -> (flags & FLAG_A11Y_VISIBLE) != 0
                         && (flags & FLAG_ROTATION_BUTTON_VISIBLE) == 0));
@@ -306,15 +311,22 @@
 
     private ImageView addButton(@DrawableRes int drawableId, @TaskbarButton int buttonType,
             ViewGroup parent, TaskbarNavButtonController navButtonController, @IdRes int id) {
-        ImageView buttonView = addButton(parent, id);
+        return addButton(drawableId, buttonType, parent, navButtonController, id,
+                R.layout.taskbar_nav_button);
+    }
+
+    private ImageView addButton(@DrawableRes int drawableId, @TaskbarButton int buttonType,
+            ViewGroup parent, TaskbarNavButtonController navButtonController, @IdRes int id,
+            @LayoutRes int layoutId) {
+        ImageView buttonView = addButton(parent, id, layoutId);
         buttonView.setImageResource(drawableId);
         buttonView.setOnClickListener(view -> navButtonController.onButtonClick(buttonType));
         return buttonView;
     }
 
-    private ImageView addButton(ViewGroup parent, int id) {
+    private ImageView addButton(ViewGroup parent, @IdRes int id, @LayoutRes int layoutId) {
         ImageView buttonView = (ImageView) mContext.getLayoutInflater()
-                .inflate(R.layout.taskbar_nav_button, parent, false);
+                .inflate(layoutId, parent, false);
         buttonView.setId(id);
         parent.addView(buttonView);
         mAllButtons.add(buttonView);
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java
index ec98bbf..453bf1c 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java
@@ -22,7 +22,10 @@
 import static com.android.launcher3.util.DisplayController.CHANGE_DENSITY;
 import static com.android.launcher3.util.DisplayController.CHANGE_SUPPORTED_BOUNDS;
 
+import android.content.ComponentCallbacks;
 import android.content.Context;
+import android.content.pm.ActivityInfo;
+import android.content.res.Configuration;
 import android.hardware.display.DisplayManager;
 import android.net.Uri;
 import android.provider.Settings;
@@ -58,6 +61,7 @@
     private final SysUINavigationMode mSysUINavigationMode;
     private final TaskbarNavButtonController mNavButtonController;
     private final SettingsCache.OnChangeListener mUserSetupCompleteListener;
+    private final ComponentCallbacks mComponentCallbacks;
 
     // The source for this provider is set when Launcher is available
     private final ScopedUnfoldTransitionProgressProvider mUnfoldProgressProvider =
@@ -84,11 +88,27 @@
         mContext = service.createWindowContext(display, TYPE_APPLICATION_OVERLAY, null);
         mNavButtonController = new TaskbarNavButtonController(service);
         mUserSetupCompleteListener = isUserSetupComplete -> recreateTaskbar();
+        mComponentCallbacks = new ComponentCallbacks() {
+            private Configuration mOldConfig = mContext.getResources().getConfiguration();
+
+            @Override
+            public void onConfigurationChanged(Configuration newConfig) {
+                if ((mOldConfig.diff(newConfig) & ActivityInfo.CONFIG_ASSETS_PATHS) != 0) {
+                    // Color has changed, recreate taskbar to reload background color & icons.
+                    recreateTaskbar();
+                }
+                mOldConfig = newConfig;
+            }
+
+            @Override
+            public void onLowMemory() { }
+        };
 
         mDisplayController.addChangeListener(this);
         mSysUINavigationMode.addModeChangeListener(this);
         SettingsCache.INSTANCE.get(mContext).register(USER_SETUP_COMPLETE_URI,
                 mUserSetupCompleteListener);
+        mContext.registerComponentCallbacks(mComponentCallbacks);
 
         recreateTaskbar();
     }
@@ -210,6 +230,7 @@
         mSysUINavigationMode.removeModeChangeListener(this);
         SettingsCache.INSTANCE.get(mContext).unregister(USER_SETUP_COMPLETE_URI,
                 mUserSetupCompleteListener);
+        mContext.unregisterComponentCallbacks(mComponentCallbacks);
     }
 
     public @Nullable TaskbarActivityContext getCurrentActivityContext() {
diff --git a/quickstep/src/com/android/launcher3/uioverrides/ApiWrapper.java b/quickstep/src/com/android/launcher3/uioverrides/ApiWrapper.java
index 8181a84..c9909cc 100644
--- a/quickstep/src/com/android/launcher3/uioverrides/ApiWrapper.java
+++ b/quickstep/src/com/android/launcher3/uioverrides/ApiWrapper.java
@@ -56,8 +56,14 @@
     public static int getHotseatEndOffset(Context context) {
         if (SysUINavigationMode.INSTANCE.get(context).getMode() == Mode.THREE_BUTTONS) {
             Resources res = context.getResources();
+            /*
+            * 2 (left + right) x Padding +
+            * 3 nav buttons +
+            * Little space at the end for contextual buttons
+            */
             return 2 * res.getDimensionPixelSize(R.dimen.taskbar_nav_buttons_spacing)
-                    + 3 * res.getDimensionPixelSize(R.dimen.taskbar_nav_buttons_size);
+                    + 3 * res.getDimensionPixelSize(R.dimen.taskbar_nav_buttons_size)
+                    + res.getDimensionPixelSize(R.dimen.taskbar_contextual_button_margin);
         } else {
             return 0;
         }
diff --git a/quickstep/src/com/android/quickstep/util/RecentsOrientedState.java b/quickstep/src/com/android/quickstep/util/RecentsOrientedState.java
index 841e578..a4db596 100644
--- a/quickstep/src/com/android/quickstep/util/RecentsOrientedState.java
+++ b/quickstep/src/com/android/quickstep/util/RecentsOrientedState.java
@@ -600,17 +600,7 @@
             width = Math.min(currentSize.x, currentSize.y);
             height = Math.max(currentSize.x, currentSize.y);
         }
-
-        DeviceProfile bestMatch = idp.supportedProfiles.get(0);
-        float minDiff = Float.MAX_VALUE;
-        for (DeviceProfile profile : idp.supportedProfiles) {
-            float diff = Math.abs(profile.widthPx - width) + Math.abs(profile.heightPx - height);
-            if (diff < minDiff) {
-                minDiff = diff;
-                bestMatch = profile;
-            }
-        }
-        return bestMatch;
+        return idp.getBestMatch(width, height);
     }
 
     private static String nameAndAddress(Object obj) {
diff --git a/quickstep/src/com/android/quickstep/views/RecentsView.java b/quickstep/src/com/android/quickstep/views/RecentsView.java
index 3b962d8..b690982 100644
--- a/quickstep/src/com/android/quickstep/views/RecentsView.java
+++ b/quickstep/src/com/android/quickstep/views/RecentsView.java
@@ -1866,7 +1866,7 @@
         setCurrentTask(-1);
         mIgnoreResetTaskId = -1;
         mTaskListChangeId = -1;
-        mFocusedTaskViewId = getTaskViewCount() > 0 ? getTaskViewAt(0).getTaskViewId() : -1;
+        mFocusedTaskViewId = -1;
 
         if (mRecentsAnimationController != null) {
             if (ENABLE_QUICKSTEP_LIVE_TILE.get() && mEnableDrawingLiveTile) {
diff --git a/quickstep/robolectric_tests/src/com/android/quickstep/util/RecentsOrientedStateTest.java b/quickstep/tests/src/com/android/quickstep/util/RecentsOrientedStateTest.java
similarity index 89%
rename from quickstep/robolectric_tests/src/com/android/quickstep/util/RecentsOrientedStateTest.java
rename to quickstep/tests/src/com/android/quickstep/util/RecentsOrientedStateTest.java
index 656379f..47ef13b 100644
--- a/quickstep/robolectric_tests/src/com/android/quickstep/util/RecentsOrientedStateTest.java
+++ b/quickstep/tests/src/com/android/quickstep/util/RecentsOrientedStateTest.java
@@ -19,33 +19,34 @@
 import static android.view.Surface.ROTATION_180;
 import static android.view.Surface.ROTATION_90;
 
+import static androidx.test.core.app.ApplicationProvider.getApplicationContext;
+
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotEquals;
 
 import android.content.Context;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+import androidx.test.filters.SmallTest;
+
 import com.android.quickstep.FallbackActivityInterface;
 
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
-import org.robolectric.RobolectricTestRunner;
-import org.robolectric.RuntimeEnvironment;
-import org.robolectric.annotation.LooperMode;
-import org.robolectric.annotation.LooperMode.Mode;
 
 /**
  * Tests for {@link RecentsOrientedState}
  */
-@RunWith(RobolectricTestRunner.class)
-@LooperMode(Mode.PAUSED)
+@SmallTest
+@RunWith(AndroidJUnit4.class)
 public class RecentsOrientedStateTest {
 
     private RecentsOrientedState mR1, mR2;
 
     @Before
     public void setup() {
-        Context context = RuntimeEnvironment.application;
+        Context context = getApplicationContext();
         mR1 = new RecentsOrientedState(context, FallbackActivityInterface.INSTANCE, i -> { });
         mR2 = new RecentsOrientedState(context, FallbackActivityInterface.INSTANCE, i -> { });
         assertEquals(mR1.getStateId(), mR2.getStateId());
diff --git a/quickstep/tests/src/com/android/quickstep/util/TaskViewSimulatorTest.java b/quickstep/tests/src/com/android/quickstep/util/TaskViewSimulatorTest.java
new file mode 100644
index 0000000..9c5cfcd
--- /dev/null
+++ b/quickstep/tests/src/com/android/quickstep/util/TaskViewSimulatorTest.java
@@ -0,0 +1,246 @@
+/*
+ * Copyright (C) 2020 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.android.quickstep.util;
+
+import static android.view.Display.DEFAULT_DISPLAY;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.mock;
+
+import android.graphics.Point;
+import android.graphics.Rect;
+import android.graphics.RectF;
+import android.view.Display;
+import android.view.Surface;
+import android.view.SurfaceControl;
+
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+import androidx.test.filters.SmallTest;
+
+import com.android.launcher3.DeviceProfile;
+import com.android.launcher3.InvariantDeviceProfile;
+import com.android.launcher3.util.DisplayController;
+import com.android.launcher3.util.DisplayController.Info;
+import com.android.launcher3.util.LauncherModelHelper;
+import com.android.launcher3.util.ReflectionHelpers;
+import com.android.quickstep.FallbackActivityInterface;
+import com.android.quickstep.SysUINavigationMode;
+import com.android.quickstep.SystemUiProxy;
+import com.android.systemui.shared.system.RemoteAnimationTargetCompat;
+import com.android.systemui.shared.system.SyncRtSurfaceTransactionApplierCompat.SurfaceParams;
+
+import org.hamcrest.Description;
+import org.hamcrest.TypeSafeMatcher;
+import org.junit.Assert;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+@SmallTest
+@RunWith(AndroidJUnit4.class)
+public class TaskViewSimulatorTest {
+
+    @Test
+    public void taskProperlyScaled_portrait_noRotation_sameInsets1() {
+        new TaskMatrixVerifier()
+                .withLauncherSize(1200, 2450)
+                .withInsets(new Rect(0, 80, 0, 120))
+                .verifyNoTransforms();
+    }
+
+    @Test
+    public void taskProperlyScaled_portrait_noRotation_sameInsets2() {
+        new TaskMatrixVerifier()
+                .withLauncherSize(1200, 2450)
+                .withInsets(new Rect(55, 80, 55, 120))
+                .verifyNoTransforms();
+    }
+
+    @Test
+    public void taskProperlyScaled_landscape_noRotation_sameInsets1() {
+        new TaskMatrixVerifier()
+                .withLauncherSize(2450, 1250)
+                .withInsets(new Rect(0, 80, 0, 40))
+                .verifyNoTransforms();
+    }
+
+    @Test
+    public void taskProperlyScaled_landscape_noRotation_sameInsets2() {
+        new TaskMatrixVerifier()
+                .withLauncherSize(2450, 1250)
+                .withInsets(new Rect(0, 80, 120, 0))
+                .verifyNoTransforms();
+    }
+
+    @Test
+    public void taskProperlyScaled_landscape_noRotation_sameInsets3() {
+        new TaskMatrixVerifier()
+                .withLauncherSize(2450, 1250)
+                .withInsets(new Rect(55, 80, 55, 120))
+                .verifyNoTransforms();
+    }
+
+    @Test
+    public void taskProperlyScaled_landscape_rotated() {
+        new TaskMatrixVerifier()
+                .withLauncherSize(1200, 2450)
+                .withInsets(new Rect(0, 80, 0, 120))
+                .withAppBounds(
+                        new Rect(0, 0, 2450, 1200),
+                        new Rect(0, 80, 0, 120),
+                        Surface.ROTATION_90)
+                .verifyNoTransforms();
+    }
+
+    private static class TaskMatrixVerifier extends TransformParams {
+
+        private Point mDisplaySize = new Point();
+        private Rect mDisplayInsets = new Rect();
+        private Rect mAppBounds = new Rect();
+        private Rect mLauncherInsets = new Rect();
+
+        private Rect mAppInsets;
+
+        private int mAppRotation = -1;
+        private DeviceProfile mDeviceProfile;
+
+        TaskMatrixVerifier withLauncherSize(int width, int height) {
+            mDisplaySize.set(width, height);
+            if (mAppBounds.isEmpty()) {
+                mAppBounds.set(0, 0, width, height);
+            }
+            return this;
+        }
+
+        TaskMatrixVerifier withInsets(Rect insets) {
+            mDisplayInsets.set(insets);
+            mLauncherInsets.set(insets);
+            return this;
+        }
+
+        TaskMatrixVerifier withAppBounds(Rect bounds, Rect insets, int appRotation) {
+            mAppBounds.set(bounds);
+            mAppInsets = insets;
+            mAppRotation = appRotation;
+            return this;
+        }
+
+        void verifyNoTransforms() {
+            LauncherModelHelper helper = new LauncherModelHelper();
+            try {
+                helper.sandboxContext.allow(SystemUiProxy.INSTANCE);
+                helper.sandboxContext.allow(SysUINavigationMode.INSTANCE);
+
+                Display display = mock(Display.class);
+                doReturn(DEFAULT_DISPLAY).when(display).getDisplayId();
+                doReturn(mDisplaySize.x > mDisplaySize.y ? Surface.ROTATION_90 : Surface.ROTATION_0)
+                        .when(display).getRotation();
+                doAnswer(i -> {
+                    ((Point) i.getArgument(0)).set(mDisplaySize.x, mDisplaySize.y);
+                    return null;
+                }).when(display).getRealSize(any());
+                doAnswer(i -> {
+                    Point smallestSize = i.getArgument(0);
+                    Point largestSize = i.getArgument(1);
+                    smallestSize.x = smallestSize.y = Math.min(mDisplaySize.x, mDisplaySize.y);
+                    largestSize.x = largestSize.y = Math.max(mDisplaySize.x, mDisplaySize.y);
+
+                    smallestSize.x -= mDisplayInsets.left + mDisplayInsets.right;
+                    largestSize.x -= mDisplayInsets.left + mDisplayInsets.right;
+
+                    smallestSize.y -= mDisplayInsets.top + mDisplayInsets.bottom;
+                    largestSize.y -= mDisplayInsets.top + mDisplayInsets.bottom;
+                    return null;
+                }).when(display).getCurrentSizeRange(any(), any());
+                DisplayController.Info mockInfo = new Info(helper.sandboxContext, display);
+
+                DisplayController controller =
+                        DisplayController.INSTANCE.get(helper.sandboxContext);
+                controller.close();
+                ReflectionHelpers.setField(controller, "mInfo", mockInfo);
+
+                mDeviceProfile = InvariantDeviceProfile.INSTANCE.get(helper.sandboxContext)
+                        .getBestMatch(mAppBounds.width(), mAppBounds.height());
+                mDeviceProfile.updateInsets(mLauncherInsets);
+
+                TaskViewSimulator tvs = new TaskViewSimulator(helper.sandboxContext,
+                        FallbackActivityInterface.INSTANCE);
+                tvs.setDp(mDeviceProfile);
+
+                int launcherRotation = mockInfo.rotation;
+                if (mAppRotation < 0) {
+                    mAppRotation = launcherRotation;
+                }
+
+                tvs.getOrientationState().update(launcherRotation, mAppRotation);
+                if (mAppInsets == null) {
+                    mAppInsets = new Rect(mLauncherInsets);
+                }
+                tvs.setPreviewBounds(mAppBounds, mAppInsets);
+
+                tvs.fullScreenProgress.value = 1;
+                tvs.recentsViewScale.value = tvs.getFullScreenScale();
+                tvs.apply(this);
+            } finally {
+                helper.destroy();
+            }
+        }
+
+        @Override
+        public SurfaceParams[] createSurfaceParams(BuilderProxy proxy) {
+            SurfaceParams.Builder builder = new SurfaceParams.Builder((SurfaceControl) null);
+            proxy.onBuildTargetParams(builder, mock(RemoteAnimationTargetCompat.class), this);
+            return new SurfaceParams[] {builder.build()};
+        }
+
+        @Override
+        public void applySurfaceParams(SurfaceParams[] params) {
+            // Verify that the task position remains the same
+            RectF newAppBounds = new RectF(mAppBounds);
+            params[0].matrix.mapRect(newAppBounds);
+            Assert.assertThat(newAppBounds, new AlmostSame(mAppBounds));
+
+            System.err.println("Bounds mapped: " + mAppBounds + " => " + newAppBounds);
+        }
+    }
+
+    private static class AlmostSame extends TypeSafeMatcher<RectF>  {
+
+        // Allow .1% error margin to account for float to int conversions
+        private final float mErrorFactor = .001f;
+        private final Rect mExpected;
+
+        AlmostSame(Rect expected) {
+            mExpected = expected;
+        }
+
+        @Override
+        protected boolean matchesSafely(RectF item) {
+            float errorWidth = mErrorFactor * mExpected.width();
+            float errorHeight = mErrorFactor * mExpected.height();
+            return Math.abs(item.left - mExpected.left) < errorWidth
+                    && Math.abs(item.top - mExpected.top) < errorHeight
+                    && Math.abs(item.right - mExpected.right) < errorWidth
+                    && Math.abs(item.bottom - mExpected.bottom) < errorHeight;
+        }
+
+        @Override
+        public void describeTo(Description description) {
+            description.appendValue(mExpected);
+        }
+    }
+}
diff --git a/res/layout/qsb_preview.xml b/res/layout/qsb_preview.xml
new file mode 100644
index 0000000..801fb04
--- /dev/null
+++ b/res/layout/qsb_preview.xml
@@ -0,0 +1,30 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!--
+     Copyright (C) 2021 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.
+-->
+<com.android.launcher3.qsb.QsbContainerView
+    xmlns:android="http://schemas.android.com/apk/res/android"
+    android:orientation="vertical"
+    android:layout_width="match_parent"
+    android:layout_height="0dp"
+    android:id="@id/search_container_workspace"
+    android:padding="0dp" >
+
+    <fragment
+        android:name="com.android.launcher3.qsb.QsbContainerView$QsbFragment"
+        android:layout_width="match_parent"
+        android:tag="qsb_view"
+        android:layout_height="match_parent"/>
+</com.android.launcher3.qsb.QsbContainerView>
\ No newline at end of file
diff --git a/src/com/android/launcher3/InvariantDeviceProfile.java b/src/com/android/launcher3/InvariantDeviceProfile.java
index a2189c9..244cb59 100644
--- a/src/com/android/launcher3/InvariantDeviceProfile.java
+++ b/src/com/android/launcher3/InvariantDeviceProfile.java
@@ -548,15 +548,18 @@
         Resources res = context.getResources();
         Configuration config = context.getResources().getConfiguration();
 
-        float availableWidth = config.screenWidthDp * res.getDisplayMetrics().density;
-        float availableHeight = config.screenHeightDp * res.getDisplayMetrics().density;
+        float screenWidth = config.screenWidthDp * res.getDisplayMetrics().density;
+        float screenHeight = config.screenHeightDp * res.getDisplayMetrics().density;
+        return getBestMatch(screenWidth, screenHeight);
+    }
 
+    public DeviceProfile getBestMatch(float screenWidth, float screenHeight) {
         DeviceProfile bestMatch = supportedProfiles.get(0);
         float minDiff = Float.MAX_VALUE;
 
         for (DeviceProfile profile : supportedProfiles) {
-            float diff = Math.abs(profile.availableWidthPx - availableWidth)
-                    + Math.abs(profile.availableHeightPx - availableHeight);
+            float diff = Math.abs(profile.widthPx - screenWidth)
+                    + Math.abs(profile.heightPx - screenHeight);
             if (diff < minDiff) {
                 minDiff = diff;
                 bestMatch = profile;
diff --git a/src/com/android/launcher3/Launcher.java b/src/com/android/launcher3/Launcher.java
index a9b44d3..ed9f044 100644
--- a/src/com/android/launcher3/Launcher.java
+++ b/src/com/android/launcher3/Launcher.java
@@ -600,9 +600,11 @@
 
     private void initDeviceProfile(InvariantDeviceProfile idp) {
         // Load configuration-specific DeviceProfile
-        mDeviceProfile = isInMultiWindowMode()
-                ? mDeviceProfile.getMultiWindowProfile(this, getMultiWindowDisplaySize())
-                : idp.getDeviceProfile(this);
+        mDeviceProfile = idp.getDeviceProfile(this);
+        if (isInMultiWindowMode()) {
+            mDeviceProfile = mDeviceProfile.getMultiWindowProfile(
+                    this, getMultiWindowDisplaySize());
+        }
 
         onDeviceProfileInitiated();
         mModelWriter = mModel.getWriter(getDeviceProfile().isVerticalBarLayout(), true, this);
diff --git a/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java b/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java
index 43ac8f9..94fc708 100644
--- a/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java
+++ b/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java
@@ -324,7 +324,9 @@
     }
 
     private void inflateAndAddFolder(FolderInfo info) {
-        CellLayout screen = mWorkspaceScreens.get(info.screenId);
+        CellLayout screen = info.container == Favorites.CONTAINER_DESKTOP
+                ? mWorkspaceScreens.get(info.screenId)
+                : mHotseat;
         FolderIcon folderIcon = FolderIcon.inflateIcon(R.layout.folder_icon, this, screen,
                 info);
         addInScreenFromBind(folderIcon, info);
@@ -471,8 +473,8 @@
         // Add first page QSB
         if (FeatureFlags.QSB_ON_FIRST_SCREEN) {
             CellLayout firstScreen = mWorkspaceScreens.get(FIRST_SCREEN_ID);
-            View qsb = mHomeElementInflater.inflate(
-                    R.layout.search_container_workspace, firstScreen, false);
+            View qsb = mHomeElementInflater.inflate(R.layout.qsb_preview, firstScreen,
+                    false);
             CellLayout.LayoutParams lp =
                     new CellLayout.LayoutParams(0, 0, firstScreen.getCountX(), 1);
             lp.canReorder = false;
diff --git a/src/com/android/launcher3/pm/InstallSessionHelper.java b/src/com/android/launcher3/pm/InstallSessionHelper.java
index ab35bd6..4b86f65 100644
--- a/src/com/android/launcher3/pm/InstallSessionHelper.java
+++ b/src/com/android/launcher3/pm/InstallSessionHelper.java
@@ -17,7 +17,6 @@
 package com.android.launcher3.pm;
 
 import static com.android.launcher3.Utilities.getPrefs;
-import static com.android.launcher3.util.Executors.MODEL_EXECUTOR;
 
 import android.content.Context;
 import android.content.pm.ApplicationInfo;
@@ -25,7 +24,6 @@
 import android.content.pm.PackageInstaller;
 import android.content.pm.PackageInstaller.SessionInfo;
 import android.content.pm.PackageManager;
-import android.os.Build;
 import android.os.Process;
 import android.os.UserHandle;
 import android.text.TextUtils;
@@ -238,24 +236,12 @@
     }
 
     public InstallSessionTracker registerInstallTracker(InstallSessionTracker.Callback callback) {
-        InstallSessionTracker tracker = new InstallSessionTracker(this, callback);
-
-        if (Build.VERSION.SDK_INT < Build.VERSION_CODES.Q) {
-            mInstaller.registerSessionCallback(tracker, MODEL_EXECUTOR.getHandler());
-        } else {
-            mLauncherApps.registerPackageInstallerSessionCallback(MODEL_EXECUTOR, tracker);
-        }
+        InstallSessionTracker tracker = new InstallSessionTracker(
+                this, callback, mInstaller, mLauncherApps);
+        tracker.register();
         return tracker;
     }
 
-    void unregister(InstallSessionTracker tracker) {
-        if (Build.VERSION.SDK_INT < Build.VERSION_CODES.Q) {
-            mInstaller.unregisterSessionCallback(tracker);
-        } else {
-            mLauncherApps.unregisterPackageInstallerSessionCallback(tracker);
-        }
-    }
-
     public static UserHandle getUserHandle(SessionInfo info) {
         return Utilities.ATLEAST_Q ? info.getUser() : Process.myUserHandle();
     }
diff --git a/src/com/android/launcher3/pm/InstallSessionTracker.java b/src/com/android/launcher3/pm/InstallSessionTracker.java
index b0b907a..e1b3c1a 100644
--- a/src/com/android/launcher3/pm/InstallSessionTracker.java
+++ b/src/com/android/launcher3/pm/InstallSessionTracker.java
@@ -18,9 +18,12 @@
 import static com.android.launcher3.pm.InstallSessionHelper.getUserHandle;
 import static com.android.launcher3.pm.PackageInstallInfo.STATUS_FAILED;
 import static com.android.launcher3.pm.PackageInstallInfo.STATUS_INSTALLED;
+import static com.android.launcher3.util.Executors.MODEL_EXECUTOR;
 
+import android.content.pm.LauncherApps;
 import android.content.pm.PackageInstaller;
 import android.content.pm.PackageInstaller.SessionInfo;
+import android.os.Build;
 import android.os.UserHandle;
 import android.util.SparseArray;
 
@@ -28,35 +31,53 @@
 
 import com.android.launcher3.util.PackageUserKey;
 
+import java.lang.ref.WeakReference;
+
 @WorkerThread
 public class InstallSessionTracker extends PackageInstaller.SessionCallback {
 
     // Lazily initialized
     private SparseArray<PackageUserKey> mActiveSessions = null;
 
-    private final InstallSessionHelper mInstallerCompat;
-    private final Callback mCallback;
+    private final WeakReference<InstallSessionHelper> mWeakHelper;
+    private final WeakReference<Callback> mWeakCallback;
+    private final PackageInstaller mInstaller;
+    private final LauncherApps mLauncherApps;
 
-    InstallSessionTracker(InstallSessionHelper installerCompat, Callback callback) {
-        mInstallerCompat = installerCompat;
-        mCallback = callback;
+
+    InstallSessionTracker(InstallSessionHelper installerCompat, Callback callback,
+            PackageInstaller installer, LauncherApps launcherApps) {
+        mWeakHelper = new WeakReference<>(installerCompat);
+        mWeakCallback = new WeakReference<>(callback);
+        mInstaller = installer;
+        mLauncherApps = launcherApps;
     }
 
     @Override
     public void onCreated(int sessionId) {
-        SessionInfo sessionInfo = pushSessionDisplayToLauncher(sessionId);
+        InstallSessionHelper helper = mWeakHelper.get();
+        Callback callback = mWeakCallback.get();
+        if (callback == null || helper == null) {
+            return;
+        }
+        SessionInfo sessionInfo = pushSessionDisplayToLauncher(sessionId, helper, callback);
         if (sessionInfo != null) {
-            mCallback.onInstallSessionCreated(PackageInstallInfo.fromInstallingState(sessionInfo));
+            callback.onInstallSessionCreated(PackageInstallInfo.fromInstallingState(sessionInfo));
         }
 
-        mInstallerCompat.tryQueuePromiseAppIcon(sessionInfo);
+        helper.tryQueuePromiseAppIcon(sessionInfo);
     }
 
     @Override
     public void onFinished(int sessionId, boolean success) {
+        InstallSessionHelper helper = mWeakHelper.get();
+        Callback callback = mWeakCallback.get();
+        if (callback == null || helper == null) {
+            return;
+        }
         // For a finished session, we can't get the session info. So use the
         // packageName from our local cache.
-        SparseArray<PackageUserKey> activeSessions = getActiveSessionMap();
+        SparseArray<PackageUserKey> activeSessions = getActiveSessionMap(helper);
         PackageUserKey key = activeSessions.get(sessionId);
         activeSessions.remove(sessionId);
 
@@ -65,21 +86,26 @@
             PackageInstallInfo info = PackageInstallInfo.fromState(
                     success ? STATUS_INSTALLED : STATUS_FAILED,
                     packageName, key.mUser);
-            mCallback.onPackageStateChanged(info);
+            callback.onPackageStateChanged(info);
 
-            if (!success && mInstallerCompat.promiseIconAddedForId(sessionId)) {
-                mCallback.onSessionFailure(packageName, key.mUser);
+            if (!success && helper.promiseIconAddedForId(sessionId)) {
+                callback.onSessionFailure(packageName, key.mUser);
                 // If it is successful, the id is removed in the the package added flow.
-                mInstallerCompat.removePromiseIconId(sessionId);
+                helper.removePromiseIconId(sessionId);
             }
         }
     }
 
     @Override
     public void onProgressChanged(int sessionId, float progress) {
-        SessionInfo session = mInstallerCompat.getVerifiedSessionInfo(sessionId);
+        InstallSessionHelper helper = mWeakHelper.get();
+        Callback callback = mWeakCallback.get();
+        if (callback == null || helper == null) {
+            return;
+        }
+        SessionInfo session = helper.getVerifiedSessionInfo(sessionId);
         if (session != null && session.getAppPackageName() != null) {
-            mCallback.onPackageStateChanged(PackageInstallInfo.fromInstallingState(session));
+            callback.onPackageStateChanged(PackageInstallInfo.fromInstallingState(session));
         }
     }
 
@@ -88,35 +114,53 @@
 
     @Override
     public void onBadgingChanged(int sessionId) {
-        SessionInfo sessionInfo = pushSessionDisplayToLauncher(sessionId);
+        InstallSessionHelper helper = mWeakHelper.get();
+        Callback callback = mWeakCallback.get();
+        if (callback == null || helper == null) {
+            return;
+        }
+        SessionInfo sessionInfo = pushSessionDisplayToLauncher(sessionId, helper, callback);
         if (sessionInfo != null) {
-            mInstallerCompat.tryQueuePromiseAppIcon(sessionInfo);
+            helper.tryQueuePromiseAppIcon(sessionInfo);
         }
     }
 
-    private SessionInfo pushSessionDisplayToLauncher(int sessionId) {
-        SessionInfo session = mInstallerCompat.getVerifiedSessionInfo(sessionId);
+    private SessionInfo pushSessionDisplayToLauncher(
+            int sessionId, InstallSessionHelper helper, Callback callback) {
+        SessionInfo session = helper.getVerifiedSessionInfo(sessionId);
         if (session != null && session.getAppPackageName() != null) {
             PackageUserKey key =
                     new PackageUserKey(session.getAppPackageName(), getUserHandle(session));
-            getActiveSessionMap().put(session.getSessionId(), key);
-            mCallback.onUpdateSessionDisplay(key, session);
+            getActiveSessionMap(helper).put(session.getSessionId(), key);
+            callback.onUpdateSessionDisplay(key, session);
             return session;
         }
         return null;
     }
 
-    private SparseArray<PackageUserKey> getActiveSessionMap() {
+    private SparseArray<PackageUserKey> getActiveSessionMap(InstallSessionHelper helper) {
         if (mActiveSessions == null) {
             mActiveSessions = new SparseArray<>();
-            mInstallerCompat.getActiveSessions().forEach(
+            helper.getActiveSessions().forEach(
                     (key, si) -> mActiveSessions.put(si.getSessionId(), key));
         }
         return mActiveSessions;
     }
 
+    void register() {
+        if (Build.VERSION.SDK_INT < Build.VERSION_CODES.Q) {
+            mInstaller.registerSessionCallback(this, MODEL_EXECUTOR.getHandler());
+        } else {
+            mLauncherApps.registerPackageInstallerSessionCallback(MODEL_EXECUTOR, this);
+        }
+    }
+
     public void unregister() {
-        mInstallerCompat.unregister(this);
+        if (Build.VERSION.SDK_INT < Build.VERSION_CODES.Q) {
+            mInstaller.unregisterSessionCallback(this);
+        } else {
+            mLauncherApps.unregisterPackageInstallerSessionCallback(this);
+        }
     }
 
     public interface Callback {
diff --git a/src/com/android/launcher3/settings/SettingsActivity.java b/src/com/android/launcher3/settings/SettingsActivity.java
index 915e140..f348a33 100644
--- a/src/com/android/launcher3/settings/SettingsActivity.java
+++ b/src/com/android/launcher3/settings/SettingsActivity.java
@@ -62,8 +62,9 @@
         SharedPreferences.OnSharedPreferenceChangeListener{
 
     /** List of fragments that can be hosted by this activity. */
-    private static final List<String> VALID_PREFERENCE_FRAGMENTS = Collections.singletonList(
-            DeveloperOptionsFragment.class.getName());
+    private static final List<String> VALID_PREFERENCE_FRAGMENTS =
+            !Utilities.IS_DEBUG_DEVICE ? Collections.emptyList()
+                    : Collections.singletonList(DeveloperOptionsFragment.class.getName());
 
     private static final String DEVELOPER_OPTIONS_KEY = "pref_developer_options";
     private static final String FLAGS_PREFERENCE_KEY = "flag_toggler";
diff --git a/src/com/android/launcher3/util/ItemInfoMatcher.java b/src/com/android/launcher3/util/ItemInfoMatcher.java
index e8ba28f..ab3083d 100644
--- a/src/com/android/launcher3/util/ItemInfoMatcher.java
+++ b/src/com/android/launcher3/util/ItemInfoMatcher.java
@@ -32,6 +32,11 @@
  */
 public interface ItemInfoMatcher {
 
+    /**
+     * Empty component used for match testing
+     */
+    ComponentName EMPTY_COMPONENT = new ComponentName("", "");
+
     boolean matches(ItemInfo info, ComponentName cn);
 
     /**
@@ -40,7 +45,7 @@
     default boolean matchesInfo(ItemInfo info) {
         if (info != null) {
             ComponentName cn = info.getTargetComponent();
-            return cn != null && matches(info, cn);
+            return matches(info, cn != null ? cn : EMPTY_COMPONENT);
         } else {
             return false;
         }
diff --git a/tests/src/com/android/launcher3/util/LauncherModelHelper.java b/tests/src/com/android/launcher3/util/LauncherModelHelper.java
index c9b63ae..59966ee 100644
--- a/tests/src/com/android/launcher3/util/LauncherModelHelper.java
+++ b/tests/src/com/android/launcher3/util/LauncherModelHelper.java
@@ -66,7 +66,6 @@
 import com.android.launcher3.pm.UserCache;
 import com.android.launcher3.testing.TestInformationProvider;
 import com.android.launcher3.uioverrides.plugins.PluginManagerWrapper;
-import com.android.launcher3.util.MainThreadInitializedObject.ObjectProvider;
 import com.android.launcher3.util.MainThreadInitializedObject.SandboxContext;
 import com.android.launcher3.widget.custom.CustomWidgetManager;
 
@@ -533,12 +532,6 @@
             super.onDestroy();
         }
 
-
-        @Override
-        protected <T> T createObject(ObjectProvider<T> provider) {
-            return spy(provider.get(this));
-        }
-
         @Override
         public PackageManager getPackageManager() {
             return mPm;