diff --git a/go/quickstep/src/com/android/quickstep/TaskOverlayFactoryGo.java b/go/quickstep/src/com/android/quickstep/TaskOverlayFactoryGo.java
index 872f168..350e0d1 100644
--- a/go/quickstep/src/com/android/quickstep/TaskOverlayFactoryGo.java
+++ b/go/quickstep/src/com/android/quickstep/TaskOverlayFactoryGo.java
@@ -20,25 +20,17 @@
 import static com.android.quickstep.views.OverviewActionsView.DISABLED_ROTATED;
 
 import android.annotation.SuppressLint;
-import android.app.ActivityTaskManager;
-import android.app.IAssistDataReceiver;
 import android.app.assist.AssistContent;
 import android.content.Context;
 import android.content.Intent;
-import android.graphics.Bitmap;
 import android.graphics.Matrix;
-import android.net.Uri;
-import android.os.Bundle;
-import android.os.Handler;
-import android.os.Looper;
-import android.os.RemoteException;
 import android.os.SystemClock;
 import android.text.TextUtils;
-import android.util.Log;
 
 import androidx.annotation.VisibleForTesting;
 
 import com.android.launcher3.R;
+import com.android.quickstep.util.AssistContentRequester;
 import com.android.quickstep.views.OverviewActionsView;
 import com.android.quickstep.views.TaskThumbnailView;
 import com.android.systemui.shared.recents.model.Task;
@@ -53,7 +45,6 @@
     public static final String ACTION_SEARCH = "com.android.quickstep.ACTION_SEARCH";
     public static final String ELAPSED_NANOS = "niu_actions_elapsed_realtime_nanos";
     public static final String ACTIONS_URL = "niu_actions_app_url";
-    private static final String ASSIST_KEY_CONTENT = "content";
     private static final String TAG = "TaskOverlayFactoryGo";
 
     // Empty constructor required for ResourceBasedOverride
@@ -72,13 +63,10 @@
      */
     public static final class TaskOverlayGo<T extends OverviewActionsView> extends TaskOverlay {
         private String mNIUPackageName;
-        private int mTaskId;
-        private Bundle mAssistData;
-        private final Handler mMainThreadHandler;
+        private String mWebUrl;
 
         private TaskOverlayGo(TaskThumbnailView taskThumbnailView) {
             super(taskThumbnailView);
-            mMainThreadHandler = new Handler(Looper.getMainLooper());
         }
 
         /**
@@ -99,28 +87,23 @@
             boolean isAllowedByPolicy = mThumbnailView.isRealSnapshot();
             getActionsView().setCallbacks(new OverlayUICallbacksGoImpl(isAllowedByPolicy, task));
 
-            mTaskId = task.key.id;
-            AssistDataReceiverImpl receiver = new AssistDataReceiverImpl();
-            receiver.setOverlay(this);
-
-            try {
-                ActivityTaskManager.getService().requestAssistDataForTask(receiver, mTaskId,
-                        mApplicationContext.getPackageName());
-            } catch (RemoteException e) {
-                Log.e(TAG, "Unable to request AssistData");
-            }
+            int taskId = task.key.id;
+            AssistContentRequester contentRequester =
+                    new AssistContentRequester(mApplicationContext);
+            contentRequester.requestAssistContent(taskId, this::onAssistContentReceived);
         }
 
-        /**
-         * Called when AssistDataReceiverImpl receives data from ActivityTaskManagerService's
-         * AssistDataRequester
-         */
-        public void onAssistDataReceived(Bundle data) {
-            mMainThreadHandler.post(() -> {
-                if (data != null) {
-                    mAssistData = data;
-                }
-            });
+        /** Provide Assist Content to the overlay. */
+        @VisibleForTesting
+        public void onAssistContentReceived(AssistContent assistContent) {
+            mWebUrl = assistContent.getWebUri() != null
+                    ? assistContent.getWebUri().toString() : null;
+        }
+
+        @Override
+        public void reset() {
+            super.reset();
+            mWebUrl = null;
         }
 
         /**
@@ -139,12 +122,8 @@
                     .setPackage(mNIUPackageName)
                     .putExtra(ELAPSED_NANOS, SystemClock.elapsedRealtimeNanos());
 
-            if (mAssistData != null) {
-                final AssistContent content = mAssistData.getParcelable(ASSIST_KEY_CONTENT);
-                Uri webUri = (content == null) ? null : content.getWebUri();
-                if (webUri != null) {
-                    intent.putExtra(ACTIONS_URL, webUri.toString());
-                }
+            if (mWebUrl != null) {
+                intent.putExtra(ACTIONS_URL, mWebUrl);
             }
 
             return intent;
@@ -191,26 +170,6 @@
     }
 
     /**
-     * Basic AssistDataReceiver. This is passed to ActivityTaskManagerService, which then requests
-     * the data.
-     */
-    private static final class AssistDataReceiverImpl extends IAssistDataReceiver.Stub {
-        private TaskOverlayGo mOverlay;
-
-        public void setOverlay(TaskOverlayGo overlay) {
-            mOverlay = overlay;
-        }
-
-        @Override
-        public void onHandleAssistData(Bundle data) {
-            mOverlay.onAssistDataReceived(data);
-        }
-
-        @Override
-        public void onHandleAssistScreenshot(Bitmap screenshot) {}
-    }
-
-    /**
      * Callbacks the Ui can generate. This is the only way for a Ui to call methods on the
      * controller.
      */
diff --git a/quickstep/src/com/android/quickstep/util/AssistContentRequester.java b/quickstep/src/com/android/quickstep/util/AssistContentRequester.java
new file mode 100644
index 0000000..3730284
--- /dev/null
+++ b/quickstep/src/com/android/quickstep/util/AssistContentRequester.java
@@ -0,0 +1,103 @@
+/*
+ * 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.
+ */
+
+package com.android.quickstep.util;
+
+import android.app.ActivityTaskManager;
+import android.app.IActivityTaskManager;
+import android.app.IAssistDataReceiver;
+import android.app.assist.AssistContent;
+import android.content.Context;
+import android.graphics.Bitmap;
+import android.os.Bundle;
+import android.os.RemoteException;
+import android.util.Log;
+
+import com.android.launcher3.util.Executors;
+
+import java.util.concurrent.Executor;
+
+/**
+ * Can be used to request the AssistContent from a provided task id, useful for getting the web uri
+ * if provided from the task.
+ */
+public class AssistContentRequester {
+    private static final String TAG = "AssistContentRequester";
+    private static final String ASSIST_KEY_CONTENT = "content";
+
+    /** For receiving content, called on the main thread. */
+    public interface Callback {
+        /**
+         * Called when the {@link android.app.assist.AssistContent} of the requested task is
+         * available.
+         **/
+        void onAssistContentAvailable(AssistContent assistContent);
+    }
+
+    private final IActivityTaskManager mActivityTaskManager;
+    private final String mPackageName;
+    private final Executor mCallbackExecutor;
+
+    public AssistContentRequester(Context context) {
+        mActivityTaskManager = ActivityTaskManager.getService();
+        mPackageName = context.getApplicationContext().getPackageName();
+        mCallbackExecutor = Executors.MAIN_EXECUTOR;
+    }
+
+    /**
+     * Request the {@link AssistContent} from the task with the provided id.
+     *
+     * @param taskId to query for the content.
+     * @param callback to call when the content is available, called on the main thread.
+     */
+    public void requestAssistContent(int taskId, Callback callback) {
+        try {
+            mActivityTaskManager.requestAssistDataForTask(
+                    new AssistDataReceiver(callback, mCallbackExecutor), taskId, mPackageName);
+        } catch (RemoteException e) {
+            Log.e(TAG, "Requesting assist content failed for task: " + taskId, e);
+        }
+    }
+
+    private static final class AssistDataReceiver extends IAssistDataReceiver.Stub {
+
+        private final Executor mExecutor;
+        private final Callback mCallback;
+
+        AssistDataReceiver(Callback callback, Executor callbackExecutor) {
+            mCallback = callback;
+            mExecutor = callbackExecutor;
+        }
+
+        @Override
+        public void onHandleAssistData(Bundle data) {
+            if (data == null) {
+                return;
+            }
+
+            final AssistContent content = data.getParcelable(ASSIST_KEY_CONTENT);
+            if (content == null) {
+                Log.e(TAG, "Received AssistData, but no AssistContent found");
+                return;
+            }
+
+            mExecutor.execute(() -> mCallback.onAssistContentAvailable(content));
+        }
+
+        @Override
+        public void onHandleAssistScreenshot(Bitmap screenshot) {}
+    }
+}