diff --git a/ext_tests/src/com/android/launcher3/testing/DebugTestInformationHandler.java b/ext_tests/src/com/android/launcher3/testing/DebugTestInformationHandler.java
index dacd8a2..0f61d14 100644
--- a/ext_tests/src/com/android/launcher3/testing/DebugTestInformationHandler.java
+++ b/ext_tests/src/com/android/launcher3/testing/DebugTestInformationHandler.java
@@ -18,6 +18,8 @@
 
 import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
 
+import android.app.Activity;
+import android.app.Application;
 import android.content.Context;
 import android.os.Binder;
 import android.os.Bundle;
@@ -31,7 +33,10 @@
 
 import java.util.ArrayList;
 import java.util.Collection;
+import java.util.Collections;
 import java.util.LinkedList;
+import java.util.Map;
+import java.util.WeakHashMap;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
 
@@ -41,9 +46,48 @@
 public class DebugTestInformationHandler extends TestInformationHandler {
     private static LinkedList sLeaks;
     private static Collection<String> sEvents;
+    private static Application.ActivityLifecycleCallbacks sActivityLifecycleCallbacks;
+    private static final Map<Activity, Boolean> sActivities =
+            Collections.synchronizedMap(new WeakHashMap<>());
+    private static int sActivitiesCreatedCount = 0;
 
     public DebugTestInformationHandler(Context context) {
         init(context);
+        if (sActivityLifecycleCallbacks == null) {
+            sActivityLifecycleCallbacks = new Application.ActivityLifecycleCallbacks() {
+                @Override
+                public void onActivityCreated(Activity activity, Bundle bundle) {
+                    sActivities.put(activity, true);
+                    ++sActivitiesCreatedCount;
+                }
+
+                @Override
+                public void onActivityStarted(Activity activity) {
+                }
+
+                @Override
+                public void onActivityResumed(Activity activity) {
+                }
+
+                @Override
+                public void onActivityPaused(Activity activity) {
+                }
+
+                @Override
+                public void onActivityStopped(Activity activity) {
+                }
+
+                @Override
+                public void onActivitySaveInstanceState(Activity activity, Bundle bundle) {
+                }
+
+                @Override
+                public void onActivityDestroyed(Activity activity) {
+                }
+            };
+            ((Application) context.getApplicationContext())
+                    .registerActivityLifecycleCallbacks(sActivityLifecycleCallbacks);
+        }
     }
 
     private static void runGcAndFinalizersSync() {
@@ -160,6 +204,20 @@
                 }
             }
 
+            case TestProtocol.REQUEST_GET_ACTIVITIES_CREATED_COUNT: {
+                response.putInt(TestProtocol.TEST_INFO_RESPONSE_FIELD, sActivitiesCreatedCount);
+                return response;
+            }
+
+            case TestProtocol.REQUEST_GET_ACTIVITIES: {
+                response.putStringArray(TestProtocol.TEST_INFO_RESPONSE_FIELD,
+                        sActivities.keySet().stream().map(
+                                a -> a.getClass().getSimpleName() + " ("
+                                        + (a.isDestroyed() ? "destroyed" : "current") + ")")
+                                .toArray(String[]::new));
+                return response;
+            }
+
             default:
                 return super.call(method, arg);
         }
diff --git a/quickstep/src/com/android/quickstep/views/TaskThumbnailView.java b/quickstep/src/com/android/quickstep/views/TaskThumbnailView.java
index 28044e4..a9db400 100644
--- a/quickstep/src/com/android/quickstep/views/TaskThumbnailView.java
+++ b/quickstep/src/com/android/quickstep/views/TaskThumbnailView.java
@@ -47,20 +47,17 @@
 import com.android.launcher3.BaseActivity;
 import com.android.launcher3.DeviceProfile;
 import com.android.launcher3.Utilities;
-import com.android.launcher3.uioverrides.plugins.PluginManagerWrapper;
 import com.android.launcher3.util.MainThreadInitializedObject;
 import com.android.launcher3.util.SystemUiController;
 import com.android.quickstep.TaskOverlayFactory.TaskOverlay;
 import com.android.quickstep.views.TaskView.FullscreenDrawParams;
-import com.android.systemui.plugins.OverviewScreenshotActions;
-import com.android.systemui.plugins.PluginListener;
 import com.android.systemui.shared.recents.model.Task;
 import com.android.systemui.shared.recents.model.ThumbnailData;
 
 /**
  * A task in the Recents view.
  */
-public class TaskThumbnailView extends View implements PluginListener<OverviewScreenshotActions> {
+public class TaskThumbnailView extends View {
     private static final MainThreadInitializedObject<FullscreenDrawParams> TEMP_PARAMS =
             new MainThreadInitializedObject<>(FullscreenDrawParams::new);
 
@@ -97,7 +94,6 @@
     private float mDimAlpha = 0f;
 
     private boolean mOverlayEnabled;
-    private OverviewScreenshotActions mOverviewScreenshotActionsPlugin;
 
     public TaskThumbnailView(Context context) {
         this(context, null);
@@ -171,9 +167,6 @@
             mPaint.setShader(null);
             getTaskOverlay().reset();
         }
-        if (mOverviewScreenshotActionsPlugin != null) {
-            mOverviewScreenshotActionsPlugin.setupActions(getTaskView(), getThumbnail(), mActivity);
-        }
         updateThumbnailPaintFilter();
     }
 
@@ -265,33 +258,6 @@
         canvas.restore();
     }
 
-    @Override
-    public void onPluginConnected(OverviewScreenshotActions overviewScreenshotActions,
-            Context context) {
-        mOverviewScreenshotActionsPlugin = overviewScreenshotActions;
-        mOverviewScreenshotActionsPlugin.setupActions(getTaskView(), getThumbnail(), mActivity);
-    }
-
-    @Override
-    public void onPluginDisconnected(OverviewScreenshotActions plugin) {
-        if (mOverviewScreenshotActionsPlugin != null) {
-            mOverviewScreenshotActionsPlugin = null;
-        }
-    }
-
-    @Override
-    protected void onAttachedToWindow() {
-        super.onAttachedToWindow();
-        PluginManagerWrapper.INSTANCE.get(getContext())
-            .addPluginListener(this, OverviewScreenshotActions.class);
-    }
-
-    @Override
-    protected void onDetachedFromWindow() {
-        super.onDetachedFromWindow();
-        PluginManagerWrapper.INSTANCE.get(getContext()).removePluginListener(this);
-    }
-
     public PreviewPositionHelper getPreviewPositionHelper() {
         return mPreviewPositionHelper;
     }
diff --git a/src/com/android/launcher3/Launcher.java b/src/com/android/launcher3/Launcher.java
index 8a35185..409ee83 100644
--- a/src/com/android/launcher3/Launcher.java
+++ b/src/com/android/launcher3/Launcher.java
@@ -23,9 +23,11 @@
 import static android.view.accessibility.AccessibilityEvent.TYPE_WINDOW_STATE_CHANGED;
 
 import static com.android.launcher3.AbstractFloatingView.TYPE_ALL;
+import static com.android.launcher3.AbstractFloatingView.TYPE_FOLDER;
 import static com.android.launcher3.AbstractFloatingView.TYPE_ICON_SURFACE;
 import static com.android.launcher3.AbstractFloatingView.TYPE_REBIND_SAFE;
 import static com.android.launcher3.AbstractFloatingView.TYPE_SNACKBAR;
+import static com.android.launcher3.AbstractFloatingView.getTopOpenViewWithType;
 import static com.android.launcher3.LauncherAnimUtils.SPRING_LOADED_EXIT_DELAY;
 import static com.android.launcher3.LauncherState.ALL_APPS;
 import static com.android.launcher3.LauncherState.FLAG_CLOSE_POPUPS;
@@ -2901,13 +2903,26 @@
      * Shows the default options popup
      */
     public void showDefaultOptions(float x, float y) {
+        OptionsPopupView.show(this, getPopupTarget(x, y), OptionsPopupView.getOptions(this),
+                false);
+    }
+
+    /**
+     * Returns target rectangle for anchoring a popup menu.
+     */
+    protected RectF getPopupTarget(float x, float y) {
         float halfSize = getResources().getDimension(R.dimen.options_menu_thumb_size) / 2;
         if (x < 0 || y < 0) {
             x = mDragLayer.getWidth() / 2;
             y = mDragLayer.getHeight() / 2;
         }
-        RectF target = new RectF(x - halfSize, y - halfSize, x + halfSize, y + halfSize);
-        OptionsPopupView.show(this, target, OptionsPopupView.getOptions(this), false);
+        return new RectF(x - halfSize, y - halfSize, x + halfSize, y + halfSize);
+    }
+
+    @Override
+    public boolean shouldUseColorExtractionForPopup() {
+        return getTopOpenViewWithType(this, TYPE_FOLDER) == null
+                && getStateManager().getState() != LauncherState.ALL_APPS;
     }
 
     @Override
diff --git a/src/com/android/launcher3/popup/ArrowPopup.java b/src/com/android/launcher3/popup/ArrowPopup.java
index 112a24e..117ae42 100644
--- a/src/com/android/launcher3/popup/ArrowPopup.java
+++ b/src/com/android/launcher3/popup/ArrowPopup.java
@@ -51,21 +51,20 @@
 import androidx.annotation.NonNull;
 
 import com.android.launcher3.AbstractFloatingView;
-import com.android.launcher3.BaseDraggingActivity;
 import com.android.launcher3.InsettableFrameLayout;
 import com.android.launcher3.Launcher;
-import com.android.launcher3.LauncherState;
 import com.android.launcher3.R;
 import com.android.launcher3.Utilities;
 import com.android.launcher3.Workspace;
 import com.android.launcher3.dragndrop.DragLayer;
 import com.android.launcher3.shortcuts.DeepShortcutView;
-import com.android.launcher3.statemanager.StatefulActivity;
 import com.android.launcher3.util.Themes;
+import com.android.launcher3.views.ActivityContext;
 import com.android.launcher3.views.BaseDragLayer;
 import com.android.launcher3.widget.LocalColorExtractor;
 
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
 
@@ -74,7 +73,7 @@
  *
  * @param <T> The activity on with the popup shows
  */
-public abstract class ArrowPopup<T extends StatefulActivity<LauncherState>>
+public abstract class ArrowPopup<T extends Context & ActivityContext>
         extends AbstractFloatingView {
 
     // Duration values (ms) for popup open and close animations.
@@ -98,7 +97,7 @@
 
     protected final LayoutInflater mInflater;
     private final float mOutlineRadius;
-    protected final T mLauncher;
+    protected final T mActivityContext;
     protected final boolean mIsRtl;
 
     private final int mArrowOffsetVertical;
@@ -131,13 +130,13 @@
 
     private final String mIterateChildrenTag;
 
-    private final int[] mColors;
+    private final int[] mColorIds;
 
     public ArrowPopup(Context context, AttributeSet attrs, int defStyleAttr) {
         super(context, attrs, defStyleAttr);
         mInflater = LayoutInflater.from(context);
         mOutlineRadius = Themes.getDialogCornerRadius(context);
-        mLauncher = BaseDraggingActivity.fromContext(context);
+        mActivityContext = ActivityContext.lookupContext(context);
         mIsRtl = Utilities.isRtl(getResources());
 
         mBackgroundColor = Themes.getAttrColor(context, R.attr.popupColorPrimary);
@@ -169,22 +168,18 @@
 
         mIterateChildrenTag = getContext().getString(R.string.popup_container_iterate_children);
 
-        boolean isAboveAnotherSurface = getTopOpenViewWithType(mLauncher, TYPE_FOLDER) != null
-                || mLauncher.getStateManager().getState() == LauncherState.ALL_APPS;
-        if (!isAboveAnotherSurface && Utilities.ATLEAST_S && ENABLE_LOCAL_COLOR_POPUPS.get()) {
+        boolean shouldUseColorExtraction = mActivityContext.shouldUseColorExtractionForPopup();
+        if (shouldUseColorExtraction && Utilities.ATLEAST_S && ENABLE_LOCAL_COLOR_POPUPS.get()) {
             mColorExtractors = new ArrayList<>();
         } else {
             mColorExtractors = null;
         }
 
-        if (isAboveAnotherSurface) {
-            mColors = new int[] {
-                    getColorStateList(context, R.color.popup_shade_first).getDefaultColor()};
+        if (shouldUseColorExtraction) {
+            mColorIds = new int[]{R.color.popup_shade_first, R.color.popup_shade_second,
+                    R.color.popup_shade_third};
         } else {
-            mColors = new int[] {
-                    getColorStateList(context, R.color.popup_shade_first).getDefaultColor(),
-                    getColorStateList(context, R.color.popup_shade_second).getDefaultColor(),
-                    getColorStateList(context, R.color.popup_shade_third).getDefaultColor()};
+            mColorIds = new int[]{R.color.popup_shade_first};
         }
     }
 
@@ -236,17 +231,22 @@
     }
 
     /**
-     * @param backgroundColor When Color.TRANSPARENT, we get color from {@link #mColors}.
+     * @param backgroundColor When Color.TRANSPARENT, we get color from {@link #mColorIds}.
      *                        Otherwise, we will use this color for all child views.
      */
     private void assignMarginsAndBackgrounds(ViewGroup viewGroup, int backgroundColor) {
-        final boolean getColorFromColorArray = backgroundColor == Color.TRANSPARENT;
+        int[] colors = null;
+        if (backgroundColor == Color.TRANSPARENT) {
+            // Lazily get the colors so they match the current wallpaper colors.
+            colors = Arrays.stream(mColorIds).map(
+                    r -> getColorStateList(getContext(), r).getDefaultColor()).toArray();
+        }
 
         int count = viewGroup.getChildCount();
         int totalVisibleShortcuts = 0;
         for (int i = 0; i < count; i++) {
             View view = viewGroup.getChildAt(i);
-            if (view.getVisibility() == VISIBLE && view instanceof DeepShortcutView) {
+            if (view.getVisibility() == VISIBLE && isShortcutOrWrapper(view)) {
                 totalVisibleShortcuts++;
             }
         }
@@ -266,9 +266,8 @@
                 MarginLayoutParams mlp = (MarginLayoutParams) lastView.getLayoutParams();
                 mlp.bottomMargin = 0;
 
-
-                if (getColorFromColorArray) {
-                    backgroundColor = mColors[numVisibleChild % mColors.length];
+                if (colors != null) {
+                    backgroundColor = colors[numVisibleChild % colors.length];
                 }
 
                 if (view instanceof ViewGroup && mIterateChildrenTag.equals(view.getTag())) {
@@ -277,7 +276,7 @@
                     continue;
                 }
 
-                if (view instanceof DeepShortcutView) {
+                if (isShortcutOrWrapper(view)) {
                     if (totalVisibleShortcuts == 1) {
                         view.setBackgroundResource(R.drawable.single_item_primary);
                     } else if (totalVisibleShortcuts > 1) {
@@ -310,6 +309,12 @@
         measure(MeasureSpec.UNSPECIFIED, MeasureSpec.UNSPECIFIED);
     }
 
+    /**
+     * Returns {@code true} if the child is a shortcut or wraps a shortcut.
+     */
+    protected boolean isShortcutOrWrapper(View view) {
+        return view instanceof DeepShortcutView;
+    }
 
     @TargetApi(Build.VERSION_CODES.S)
     private int getExtractedColor(SparseIntArray colors) {
@@ -427,7 +432,7 @@
     /**
      * Shows the popup at the desired location.
      */
-    protected void show() {
+    public void show() {
         setupForDisplay();
         onInflationComplete(false);
         assignMarginsAndBackgrounds(this);
@@ -807,6 +812,6 @@
     }
 
     protected BaseDragLayer getPopupContainer() {
-        return mLauncher.getDragLayer();
+        return mActivityContext.getDragLayer();
     }
 }
diff --git a/src/com/android/launcher3/popup/PopupContainerWithArrow.java b/src/com/android/launcher3/popup/PopupContainerWithArrow.java
index 454dc6e..e340b21 100644
--- a/src/com/android/launcher3/popup/PopupContainerWithArrow.java
+++ b/src/com/android/launcher3/popup/PopupContainerWithArrow.java
@@ -151,7 +151,7 @@
 
     public OnClickListener getItemClickListener() {
         return (view) -> {
-            mLauncher.getItemOnClickListener().onClick(view);
+            mActivityContext.getItemOnClickListener().onClick(view);
             close(true);
         };
     }
@@ -326,7 +326,7 @@
 
         // Load the shortcuts on a background thread and update the container as it animates.
         MODEL_EXECUTOR.getHandler().postAtFrontOfQueue(PopupPopulator.createUpdateRunnable(
-                mLauncher, originalItemInfo, new Handler(Looper.getMainLooper()),
+                mActivityContext, originalItemInfo, new Handler(Looper.getMainLooper()),
                 this, mShortcuts, notificationKeys));
     }
 
@@ -439,7 +439,7 @@
 
     private void updateNotificationHeader() {
         ItemInfoWithIcon itemInfo = (ItemInfoWithIcon) mOriginalIcon.getTag();
-        DotInfo dotInfo = mLauncher.getDotInfoForItem(itemInfo);
+        DotInfo dotInfo = mActivityContext.getDotInfoForItem(itemInfo);
         if (mNotificationContainer != null && dotInfo != null) {
             mNotificationContainer.updateHeader(dotInfo.getNotificationCount());
         }
@@ -481,7 +481,7 @@
     @Override
     protected void closeComplete() {
         super.closeComplete();
-        PopupContainerWithArrow openPopup = getOpen(mLauncher);
+        PopupContainerWithArrow openPopup = getOpen(mActivityContext);
         if (openPopup == null || openPopup.mOriginalIcon != mOriginalIcon) {
             mOriginalIcon.setTextVisibility(mOriginalIcon.shouldTextBeVisible());
             mOriginalIcon.setForceHideDot(false);
diff --git a/src/com/android/launcher3/testing/TestProtocol.java b/src/com/android/launcher3/testing/TestProtocol.java
index 3aecaa5..9e52dac 100644
--- a/src/com/android/launcher3/testing/TestProtocol.java
+++ b/src/com/android/launcher3/testing/TestProtocol.java
@@ -99,6 +99,9 @@
     public static final String REQUEST_CLEAR_DATA = "clear-data";
     public static final String REQUEST_IS_TABLET = "is-tablet";
     public static final String REQUEST_IS_TWO_PANELS = "is-two-panel";
+    public static final String REQUEST_GET_ACTIVITIES_CREATED_COUNT =
+            "get-activities-created-count";
+    public static final String REQUEST_GET_ACTIVITIES = "get-activities";
 
     public static Long sForcePauseTimeout;
     public static final String REQUEST_SET_FORCE_PAUSE_TIMEOUT = "set-force-pause-timeout";
diff --git a/src/com/android/launcher3/views/ActivityContext.java b/src/com/android/launcher3/views/ActivityContext.java
index ebcd379..e07d71e 100644
--- a/src/com/android/launcher3/views/ActivityContext.java
+++ b/src/com/android/launcher3/views/ActivityContext.java
@@ -116,6 +116,13 @@
     }
 
     /**
+     * Returns {@code true} if popups should use color extraction.
+     */
+    default boolean shouldUseColorExtractionForPopup() {
+        return true;
+    }
+
+    /**
      * Returns whether we can show the IME for elements hosted by this ActivityContext.
      */
     default boolean supportsIme() {
diff --git a/src/com/android/launcher3/views/OptionsPopupView.java b/src/com/android/launcher3/views/OptionsPopupView.java
index 257b18f..33ab0d2 100644
--- a/src/com/android/launcher3/views/OptionsPopupView.java
+++ b/src/com/android/launcher3/views/OptionsPopupView.java
@@ -59,7 +59,7 @@
 /**
  * Popup shown on long pressing an empty space in launcher
  */
-public class OptionsPopupView extends ArrowPopup
+public class OptionsPopupView extends ArrowPopup<Launcher>
         implements OnClickListener, OnLongClickListener {
 
     private final ArrayMap<View, OptionItem> mItemMap = new ArrayMap<>();
@@ -74,6 +74,10 @@
         super(context, attrs, defStyleAttr);
     }
 
+    public void setTargetRect(RectF targetRect) {
+        mTargetRect = targetRect;
+    }
+
     @Override
     public void onClick(View view) {
         handleViewClick(view);
@@ -90,7 +94,7 @@
             return false;
         }
         if (item.eventId.getId() > 0) {
-            mLauncher.getStatsLogManager().logger().log(item.eventId);
+            mActivityContext.getStatsLogManager().logger().log(item.eventId);
         }
         if (item.clickListener.onLongClick(view)) {
             close(true);
diff --git a/src/com/android/launcher3/widget/util/WidgetSizes.java b/src/com/android/launcher3/widget/util/WidgetSizes.java
index 451ed6e..b211f4c 100644
--- a/src/com/android/launcher3/widget/util/WidgetSizes.java
+++ b/src/com/android/launcher3/widget/util/WidgetSizes.java
@@ -17,7 +17,6 @@
 
 import static android.appwidget.AppWidgetHostView.getDefaultPaddingForWidget;
 
-
 import android.appwidget.AppWidgetHostView;
 import android.appwidget.AppWidgetManager;
 import android.content.ComponentName;
@@ -87,7 +86,13 @@
     }
 
     /**
-     * Returns the size of a WidgetItem.
+     * Returns the size of a {@link WidgetItem}.
+     *
+     * <p>This size is used by the widget picker. It should NEVER be shared with app widgets.
+     *
+     * <p>For sizes shared with app widgets, please refer to
+     * {@link #getWidgetPaddedSizes(Context, ComponentName, int, int)} &
+     * {@link #getWidgetPaddedSizePx(Context, ComponentName, DeviceProfile, int, int)}.
      */
     public static Size getWidgetItemSizePx(Context context, DeviceProfile profile,
             WidgetItem widgetItem) {
@@ -96,8 +101,15 @@
                     .getDimensionPixelSize(R.dimen.widget_preview_shortcut_padding);
             return new Size(dimension, dimension);
         }
-        return getWidgetPaddedSizePx(context, widgetItem.componentName, profile, widgetItem.spanX,
-                widgetItem.spanY);
+        Size widgetItemSize = getWidgetSizePx(profile, widgetItem.spanX,
+                widgetItem.spanY, /* recycledCellSize= */ null);
+        if (profile.shouldInsetWidgets()) {
+            Rect inset = new Rect();
+            AppWidgetHostView.getDefaultPaddingForWidget(context, widgetItem.componentName, inset);
+            return new Size(widgetItemSize.getWidth() + inset.left + inset.right,
+                    widgetItemSize.getHeight() + inset.top + inset.bottom);
+        }
+        return widgetItemSize;
     }
 
     private static Size getWidgetSizePx(DeviceProfile profile, int spanX, int spanY,
diff --git a/src_plugins/com/android/systemui/plugins/OverviewScreenshotActions.java b/src_plugins/com/android/systemui/plugins/OverviewScreenshotActions.java
deleted file mode 100644
index 8d9c0f4..0000000
--- a/src_plugins/com/android/systemui/plugins/OverviewScreenshotActions.java
+++ /dev/null
@@ -1,41 +0,0 @@
-/*
- * Copyright (C) 2019 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package com.android.systemui.plugins;
-
-import android.app.Activity;
-import android.graphics.Bitmap;
-import android.view.ViewGroup;
-
-import com.android.systemui.plugins.annotations.ProvidesInterface;
-
-/**
- * Implement this interface to add action buttons for overview screenshots, e.g. share, edit etc.
- */
-@ProvidesInterface(
-        action = OverviewScreenshotActions.ACTION, version = OverviewScreenshotActions.VERSION)
-public interface OverviewScreenshotActions extends Plugin {
-    String ACTION = "com.android.systemui.action.PLUGIN_OVERVIEW_SCREENSHOT_ACTIONS";
-    int VERSION = 1;
-
-    /**
-     * Setup the actions for the screenshot, including edit, save, etc.
-     * @param parent The parent view to add buttons on.
-     * @param screenshot The screenshot we will do actions on.
-     * @param activity THe host activity.
-     */
-    void setupActions(ViewGroup parent, Bitmap screenshot, Activity activity);
-}
diff --git a/tests/Android.bp b/tests/Android.bp
index aeddc4c..3670c37 100644
--- a/tests/Android.bp
+++ b/tests/Android.bp
@@ -31,7 +31,6 @@
     name: "launcher-oop-tests-src",
     srcs: [
       "src/com/android/launcher3/ui/AbstractLauncherUiTest.java",
-      "src/com/android/launcher3/ui/ActivityLeakTracker.java",
       "src/com/android/launcher3/ui/PortraitLandscapeRunner.java",
       "src/com/android/launcher3/util/Wait.java",
       "src/com/android/launcher3/util/WidgetUtils.java",
diff --git a/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java b/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
index 1a6ce8c..b6b6cdd 100644
--- a/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
+++ b/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
@@ -36,7 +36,6 @@
 import android.os.Debug;
 import android.os.Process;
 import android.os.RemoteException;
-import android.os.StrictMode;
 import android.os.UserHandle;
 import android.os.UserManager;
 import android.util.Log;
@@ -99,11 +98,9 @@
     public static final long DEFAULT_UI_TIMEOUT = 10000;
     private static final String TAG = "AbstractLauncherUiTest";
 
-    private static String sStrictmodeDetectedActivityLeak;
     private static boolean sDumpWasGenerated = false;
-    private static boolean sActivityLeakReported;
+    private static boolean sActivityLeakReported = false;
     private static final String SYSTEMUI_PACKAGE = "com.android.systemui";
-    protected static final ActivityLeakTracker ACTIVITY_LEAK_TRACKER = new ActivityLeakTracker();
 
     protected LooperExecutor mMainThreadExecutor = MAIN_EXECUTOR;
     protected final UiDevice mDevice = UiDevice.getInstance(getInstrumentation());
@@ -112,45 +109,25 @@
     protected String mTargetPackage;
     private int mLauncherPid;
 
-    static {
-        if (TestHelpers.isInLauncherProcess()) {
-            StrictMode.VmPolicy.Builder builder =
-                    new StrictMode.VmPolicy.Builder()
-                            .penaltyLog()
-                            .penaltyListener(Runnable::run, violation -> {
-                                if (sStrictmodeDetectedActivityLeak == null) {
-                                    sStrictmodeDetectedActivityLeak = violation.toString() + ", "
-                                            + dumpHprofData() + ".";
-                                }
-                            });
-            StrictMode.setVmPolicy(builder.build());
-        }
-    }
-
     public static void checkDetectedLeaks(LauncherInstrumentation launcher) {
         if (sActivityLeakReported) return;
 
-        if (sStrictmodeDetectedActivityLeak != null) {
-            // Report from the test thread strictmode violations detected in the main thread.
-            sActivityLeakReported = true;
-            Assert.fail(sStrictmodeDetectedActivityLeak);
-        }
-
         // Check whether activity leak detector has found leaked activities.
-        Wait.atMost(AbstractLauncherUiTest::getActivityLeakErrorMessage,
+        Wait.atMost(() -> getActivityLeakErrorMessage(launcher),
                 () -> {
                     launcher.forceGc();
                     return MAIN_EXECUTOR.submit(
-                            () -> ACTIVITY_LEAK_TRACKER.noLeakedActivities()).get();
+                            () -> launcher.noLeakedActivities()).get();
                 }, DEFAULT_UI_TIMEOUT, launcher);
     }
 
-    private static String getActivityLeakErrorMessage() {
+    private static String getActivityLeakErrorMessage(LauncherInstrumentation launcher) {
         sActivityLeakReported = true;
-        return "Activity leak detector has found leaked activities, " + dumpHprofData() + ".";
+        return "Activity leak detector has found leaked activities, "
+                + dumpHprofData(launcher) + ".";
     }
 
-    public static String dumpHprofData() {
+    public static String dumpHprofData(LauncherInstrumentation launcher) {
         String result;
         if (sDumpWasGenerated) {
             Log.d("b/195319692", "dump has already been generated by another test",
@@ -176,8 +153,7 @@
                 result = "failed to save memory dump";
             }
         }
-        return result
-                + ". Full list of activities: " + ACTIVITY_LEAK_TRACKER.getActivitiesList();
+        return result + ". Full list of activities: " + launcher.getRootedActivitiesList();
     }
 
     protected AbstractLauncherUiTest() {
diff --git a/tests/src/com/android/launcher3/ui/ActivityLeakTracker.java b/tests/src/com/android/launcher3/ui/ActivityLeakTracker.java
deleted file mode 100644
index 2db7472..0000000
--- a/tests/src/com/android/launcher3/ui/ActivityLeakTracker.java
+++ /dev/null
@@ -1,90 +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.launcher3.ui;
-
-import android.app.Activity;
-import android.app.Application;
-import android.os.Bundle;
-
-import androidx.test.InstrumentationRegistry;
-
-import com.android.launcher3.tapl.TestHelpers;
-
-import java.util.WeakHashMap;
-import java.util.stream.Collectors;
-
-public class ActivityLeakTracker implements Application.ActivityLifecycleCallbacks {
-    private final WeakHashMap<Activity, Boolean> mActivities = new WeakHashMap<>();
-
-    private int mActivitiesCreated;
-
-    ActivityLeakTracker() {
-        if (!TestHelpers.isInLauncherProcess()) return;
-        final Application app =
-                (Application) InstrumentationRegistry.getTargetContext().getApplicationContext();
-        app.registerActivityLifecycleCallbacks(this);
-    }
-
-    public int getActivitiesCreated() {
-        return mActivitiesCreated;
-    }
-
-    @Override
-    public void onActivityCreated(Activity activity, Bundle bundle) {
-        mActivities.put(activity, true);
-        ++mActivitiesCreated;
-    }
-
-    @Override
-    public void onActivityStarted(Activity activity) {
-    }
-
-    @Override
-    public void onActivityResumed(Activity activity) {
-    }
-
-    @Override
-    public void onActivityPaused(Activity activity) {
-    }
-
-    @Override
-    public void onActivityStopped(Activity activity) {
-    }
-
-    @Override
-    public void onActivitySaveInstanceState(Activity activity, Bundle bundle) {
-    }
-
-    @Override
-    public void onActivityDestroyed(Activity activity) {
-    }
-
-    public boolean noLeakedActivities() {
-        for (Activity activity : mActivities.keySet()) {
-            if (activity.isDestroyed()) {
-                return false;
-            }
-        }
-
-        return mActivities.size() <= 2;
-    }
-
-    public String getActivitiesList() {
-        return mActivities.keySet().stream().map(a -> a.getClass().getSimpleName())
-                .collect(Collectors.joining(","));
-    }
-}
diff --git a/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java b/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
index f7c6044..0360a7d 100644
--- a/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
+++ b/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
@@ -1502,6 +1502,30 @@
         getTestInfo(TestProtocol.REQUEST_CLEAR_DATA);
     }
 
+    private String[] getActivities() {
+        return getTestInfo(TestProtocol.REQUEST_GET_ACTIVITIES)
+                .getStringArray(TestProtocol.TEST_INFO_RESPONSE_FIELD);
+    }
+
+    public String getRootedActivitiesList() {
+        return String.join(", ", getActivities());
+    }
+
+    public boolean noLeakedActivities() {
+        final String[] activities = getActivities();
+        for (String activity : activities) {
+            if (activity.contains("(destroyed)")) {
+                return false;
+            }
+        }
+        return activities.length <= 2;
+    }
+
+    public int getActivitiesCreated() {
+        return getTestInfo(TestProtocol.REQUEST_GET_ACTIVITIES_CREATED_COUNT)
+                .getInt(TestProtocol.TEST_INFO_RESPONSE_FIELD);
+    }
+
     public Closable eventsCheck() {
         Assert.assertTrue("Nested event checking", mEventChecker == null);
         disableSensorRotation();