diff --git a/quickstep/src/com/android/launcher3/BaseQuickstepLauncher.java b/quickstep/src/com/android/launcher3/BaseQuickstepLauncher.java
index 4204597..3643d6d 100644
--- a/quickstep/src/com/android/launcher3/BaseQuickstepLauncher.java
+++ b/quickstep/src/com/android/launcher3/BaseQuickstepLauncher.java
@@ -316,12 +316,4 @@
     public void setHintUserWillBeActive() {
         addActivityFlags(ACTIVITY_STATE_USER_WILL_BE_ACTIVE);
     }
-
-    @Override
-    public void onAttachedToWindow() {
-        super.onAttachedToWindow();
-        if (Utilities.ATLEAST_R) {
-            InteractionJankMonitorWrapper.init(getWindow().getDecorView());
-        }
-    }
 }
diff --git a/quickstep/src/com/android/launcher3/QuickstepAppTransitionManagerImpl.java b/quickstep/src/com/android/launcher3/QuickstepAppTransitionManagerImpl.java
index 108006b..8e6a5b6 100644
--- a/quickstep/src/com/android/launcher3/QuickstepAppTransitionManagerImpl.java
+++ b/quickstep/src/com/android/launcher3/QuickstepAppTransitionManagerImpl.java
@@ -763,7 +763,7 @@
         anim.addListener(new AnimationSuccessListener() {
             @Override
             public void onAnimationStart(Animator animation) {
-                InteractionJankMonitorWrapper.begin(cuj);
+                InteractionJankMonitorWrapper.begin(mDragLayer, cuj);
                 super.onAnimationStart(animation);
             }
 
diff --git a/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java b/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
index 5f196fe..a7e4741 100644
--- a/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
+++ b/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
@@ -721,9 +721,9 @@
 
     @UiThread
     public void onGestureStarted(boolean isLikelyToStartNewTask) {
-        InteractionJankMonitorWrapper.begin(
+        InteractionJankMonitorWrapper.begin(mRecentsView,
                 InteractionJankMonitorWrapper.CUJ_QUICK_SWITCH, 2000 /* ms timeout */);
-        InteractionJankMonitorWrapper.begin(
+        InteractionJankMonitorWrapper.begin(mRecentsView,
                 InteractionJankMonitorWrapper.CUJ_APP_CLOSE_TO_HOME);
         notifyGestureStartedAsync();
         setIsLikelyToStartNewTask(isLikelyToStartNewTask, false /* animate */);
@@ -805,7 +805,8 @@
         maybeUpdateRecentsAttachedState(false);
         final GestureEndTarget endTarget = mGestureState.getEndTarget();
         if (endTarget != NEW_TASK) {
-            InteractionJankMonitorWrapper.cancel(InteractionJankMonitorWrapper.CUJ_QUICK_SWITCH);
+            InteractionJankMonitorWrapper.cancel(
+                    InteractionJankMonitorWrapper.CUJ_QUICK_SWITCH);
         }
         if (endTarget != HOME) {
             InteractionJankMonitorWrapper.cancel(
@@ -1165,7 +1166,8 @@
                 TaskInfoCompat.getPipSourceRectHint(runningTaskTarget.pictureInPictureParams),
                 TaskInfoCompat.getWindowConfigurationBounds(taskInfo),
                 startBounds,
-                destinationBounds);
+                destinationBounds,
+                mRecentsView);
         // We would assume home and app window always in the same rotation While homeRotation
         // is not ROTATION_0 (which implies the rotation is turned on in launcher settings).
         if (homeRotation == ROTATION_0
diff --git a/quickstep/src/com/android/quickstep/OverviewCommandHelper.java b/quickstep/src/com/android/quickstep/OverviewCommandHelper.java
index fd9c315..ec6e303 100644
--- a/quickstep/src/com/android/quickstep/OverviewCommandHelper.java
+++ b/quickstep/src/com/android/quickstep/OverviewCommandHelper.java
@@ -175,7 +175,9 @@
                 return;
             }
 
-            InteractionJankMonitorWrapper.begin(InteractionJankMonitorWrapper.CUJ_QUICK_SWITCH);
+            InteractionJankMonitorWrapper.begin(
+                    mActivityInterface.getCreatedActivity().getRootView(),
+                    InteractionJankMonitorWrapper.CUJ_QUICK_SWITCH);
 
             // Otherwise, start overview.
             mListener = mActivityInterface.createActivityInitListener(this::onActivityReady);
diff --git a/quickstep/src/com/android/quickstep/util/SwipePipToHomeAnimator.java b/quickstep/src/com/android/quickstep/util/SwipePipToHomeAnimator.java
index 8fbd645..378f25b 100644
--- a/quickstep/src/com/android/quickstep/util/SwipePipToHomeAnimator.java
+++ b/quickstep/src/com/android/quickstep/util/SwipePipToHomeAnimator.java
@@ -28,6 +28,7 @@
 import android.util.Log;
 import android.view.Surface;
 import android.view.SurfaceControl;
+import android.view.View;
 
 import androidx.annotation.NonNull;
 
@@ -91,7 +92,8 @@
             @NonNull Rect sourceRectHint,
             @NonNull Rect appBounds,
             @NonNull Rect startBounds,
-            @NonNull Rect destinationBounds) {
+            @NonNull Rect destinationBounds,
+            @NonNull View view) {
         mTaskId = taskId;
         mComponentName = componentName;
         mLeash = leash;
@@ -110,7 +112,7 @@
         addListener(new AnimationSuccessListener() {
             @Override
             public void onAnimationStart(Animator animation) {
-                InteractionJankMonitorWrapper.begin(CUJ_APP_CLOSE_TO_PIP);
+                InteractionJankMonitorWrapper.begin(view, CUJ_APP_CLOSE_TO_PIP);
                 super.onAnimationStart(animation);
             }
 
diff --git a/src/com/android/launcher3/allapps/search/AllAppsSearchBarController.java b/src/com/android/launcher3/allapps/search/AllAppsSearchBarController.java
index f09f1ca..522f1d4 100644
--- a/src/com/android/launcher3/allapps/search/AllAppsSearchBarController.java
+++ b/src/com/android/launcher3/allapps/search/AllAppsSearchBarController.java
@@ -15,7 +15,6 @@
  */
 package com.android.launcher3.allapps.search;
 
-import android.os.Bundle;
 import android.text.Editable;
 import android.text.TextUtils;
 import android.text.TextWatcher;
@@ -36,7 +35,6 @@
 import com.android.systemui.plugins.AllAppsSearchPlugin;
 
 import java.util.ArrayList;
-import java.util.List;
 import java.util.function.Consumer;
 
 /**
@@ -62,7 +60,7 @@
      */
     public final void initialize(
             SearchAlgorithm searchAlgorithm, ExtendedEditText input,
-            BaseDraggingActivity launcher, Callbacks cb, Consumer<List<Bundle>> secondaryCb) {
+            BaseDraggingActivity launcher, Callbacks cb) {
         mCb = cb;
         mLauncher = launcher;
 
diff --git a/src/com/android/launcher3/allapps/search/AppsSearchContainerLayout.java b/src/com/android/launcher3/allapps/search/AppsSearchContainerLayout.java
index bd2f04d..4f79fb8 100644
--- a/src/com/android/launcher3/allapps/search/AppsSearchContainerLayout.java
+++ b/src/com/android/launcher3/allapps/search/AppsSearchContainerLayout.java
@@ -24,7 +24,6 @@
 
 import android.content.Context;
 import android.graphics.Rect;
-import android.os.Bundle;
 import android.text.Selection;
 import android.text.SpannableStringBuilder;
 import android.text.method.TextKeyListener;
@@ -50,15 +49,13 @@
 import com.android.launcher3.config.FeatureFlags;
 
 import java.util.ArrayList;
-import java.util.List;
-import java.util.function.Consumer;
 
 /**
  * Layout to contain the All-apps search UI.
  */
 public class AppsSearchContainerLayout extends ExtendedEditText
         implements SearchUiManager, AllAppsSearchBarController.Callbacks,
-        AllAppsStore.OnUpdateListener, Insettable, Consumer<List<Bundle>> {
+        AllAppsStore.OnUpdateListener, Insettable {
 
     private final BaseDraggingActivity mLauncher;
     private final AllAppsSearchBarController mSearchBarController;
@@ -141,7 +138,7 @@
         mAppsView = appsView;
         mSearchBarController.initialize(
                 new DefaultAppSearchAlgorithm(mLauncher, LauncherAppState.getInstance(mLauncher)),
-                this, mLauncher, this, this);
+                this, mLauncher, this);
     }
 
     @Override
@@ -234,9 +231,4 @@
     public EditText getEditText() {
         return this;
     }
-
-    @Override
-    public void accept(List<Bundle> bundles) {
-        // TODO: Render the result on mAppsView object
-    }
 }
diff --git a/tests/src/com/android/launcher3/ui/WorkTabTest.java b/tests/src/com/android/launcher3/ui/WorkTabTest.java
index ac0d355..aef26ae 100644
--- a/tests/src/com/android/launcher3/ui/WorkTabTest.java
+++ b/tests/src/com/android/launcher3/ui/WorkTabTest.java
@@ -150,11 +150,15 @@
         waitForLauncherCondition("Launcher did not show the next edu screen", l -> {
             Log.d(TestProtocol.WORK_PROFILE_REMOVED,
                     "running test attempt" + attempt.getAndIncrement());
+            if (!(l.getAppsView().getContentView() instanceof AllAppsPagedView)) {
+                Log.d(TestProtocol.WORK_PROFILE_REMOVED, "Work tab not setup. Skipping test");
+                return false;
+            }
             return ((AllAppsPagedView) l.getAppsView().getContentView()).getCurrentPage()
                     == WORK_PAGE && ((TextView) workEduView.findViewById(
                     R.id.content_text)).getText().equals(
                     l.getResources().getString(R.string.work_profile_edu_work_apps));
-        });
+        }, 60000);
     }
 
     @Test