diff --git a/quickstep/src/com/android/quickstep/OtherActivityTouchConsumer.java b/quickstep/src/com/android/quickstep/OtherActivityTouchConsumer.java
index db50a00..bd02034 100644
--- a/quickstep/src/com/android/quickstep/OtherActivityTouchConsumer.java
+++ b/quickstep/src/com/android/quickstep/OtherActivityTouchConsumer.java
@@ -334,7 +334,7 @@
                 new AssistDataReceiver() {
                     @Override
                     public void onHandleAssistData(Bundle bundle) {
-                        // Pass to AIAI
+                        mRecentsModel.preloadAssistData(mRunningTask.id, bundle);
                     }
                 },
                 new RecentsAnimationListener() {
diff --git a/quickstep/src/com/android/quickstep/RecentsModel.java b/quickstep/src/com/android/quickstep/RecentsModel.java
index 182803e..7fe7751 100644
--- a/quickstep/src/com/android/quickstep/RecentsModel.java
+++ b/quickstep/src/com/android/quickstep/RecentsModel.java
@@ -22,12 +22,16 @@
 import android.content.res.Resources;
 import android.graphics.drawable.Drawable;
 import android.os.Build;
+import android.os.Bundle;
 import android.os.Looper;
 import android.os.UserHandle;
+import android.support.annotation.WorkerThread;
 import android.util.LruCache;
+import android.util.SparseArray;
 
 import com.android.launcher3.MainThreadExecutor;
 import com.android.launcher3.R;
+import com.android.launcher3.util.Preconditions;
 import com.android.systemui.shared.recents.ISystemUiProxy;
 import com.android.systemui.shared.recents.model.IconLoader;
 import com.android.systemui.shared.recents.model.RecentsTaskLoadPlan;
@@ -38,6 +42,7 @@
 import com.android.systemui.shared.system.BackgroundExecutor;
 import com.android.systemui.shared.system.TaskStackChangeListener;
 
+import java.util.ArrayList;
 import java.util.concurrent.ExecutionException;
 import java.util.function.Consumer;
 
@@ -66,6 +71,9 @@
         return INSTANCE;
     }
 
+    private final SparseArray<Bundle> mCachedAssistData = new SparseArray<>(1);
+    private final ArrayList<AssistDataListener> mAssistDataListeners = new ArrayList<>();
+
     private final Context mContext;
     private final RecentsTaskLoader mRecentsTaskLoader;
     private final MainThreadExecutor mMainThreadExecutor;
@@ -74,6 +82,7 @@
     private int mLastLoadPlanId;
     private int mTaskChangeId;
     private ISystemUiProxy mSystemUiProxy;
+    private boolean mClearAssistCacheOnStackChange = true;
 
     private RecentsModel(Context context) {
         mContext = context;
@@ -144,6 +153,13 @@
     @Override
     public void onTaskStackChanged() {
         mTaskChangeId++;
+
+        Preconditions.assertUIThread();
+        if (mClearAssistCacheOnStackChange) {
+            mCachedAssistData.clear();
+        } else {
+            mClearAssistCacheOnStackChange = true;
+        }
     }
 
     public boolean isLoadPlanValid(int resultId) {
@@ -161,4 +177,40 @@
     public ISystemUiProxy getSystemUiProxy() {
         return mSystemUiProxy;
     }
+
+    @WorkerThread
+    public void preloadAssistData(int taskId, Bundle data) {
+        mMainThreadExecutor.execute(() -> {
+            mCachedAssistData.put(taskId, data);
+            // We expect a stack change callback after the assist data is set. So ignore the
+            // very next stack change callback.
+            mClearAssistCacheOnStackChange = false;
+
+            int count = mAssistDataListeners.size();
+            for (int i = 0; i < count; i++) {
+                mAssistDataListeners.get(i).onAssistDataReceived(taskId);
+            }
+        });
+    }
+
+    public Bundle getAssistData(int taskId) {
+        Preconditions.assertUIThread();
+        return mCachedAssistData.get(taskId);
+    }
+
+    public void addAssistDataListener(AssistDataListener listener) {
+        mAssistDataListeners.add(listener);
+    }
+
+    public void removeAssistDataListener(AssistDataListener listener) {
+        mAssistDataListeners.remove(listener);
+    }
+
+    /**
+     * Callback for receiving assist data
+     */
+    public interface AssistDataListener {
+
+        void onAssistDataReceived(int taskId);
+    }
 }
diff --git a/quickstep/src/com/android/quickstep/TaskOverlayFactory.java b/quickstep/src/com/android/quickstep/TaskOverlayFactory.java
index c2fb7be..66969c6 100644
--- a/quickstep/src/com/android/quickstep/TaskOverlayFactory.java
+++ b/quickstep/src/com/android/quickstep/TaskOverlayFactory.java
@@ -23,6 +23,7 @@
 import com.android.launcher3.R;
 import com.android.launcher3.Utilities;
 import com.android.launcher3.util.Preconditions;
+import com.android.systemui.shared.recents.model.Task;
 import com.android.systemui.shared.recents.model.ThumbnailData;
 
 /**
@@ -47,7 +48,7 @@
 
     public static class TaskOverlay {
 
-        public void setTaskInfo(ThumbnailData thumbnail, Matrix matrix) { }
+        public void setTaskInfo(Task task, ThumbnailData thumbnail, Matrix matrix) { }
 
         public void reset() { }
 
diff --git a/quickstep/src/com/android/quickstep/TaskThumbnailView.java b/quickstep/src/com/android/quickstep/TaskThumbnailView.java
index 4f93b1c..9b9c6f2 100644
--- a/quickstep/src/com/android/quickstep/TaskThumbnailView.java
+++ b/quickstep/src/com/android/quickstep/TaskThumbnailView.java
@@ -55,6 +55,7 @@
 
     private final Matrix mMatrix = new Matrix();
 
+    private Task mTask;
     private ThumbnailData mThumbnailData;
     protected BitmapShader mBitmapShader;
 
@@ -83,6 +84,7 @@
      * Updates this thumbnail.
      */
     public void setThumbnail(Task task, ThumbnailData thumbnailData) {
+        mTask = task;
         mPaint.setColor(task == null ? Color.BLACK : task.colorBackground | 0xFF000000);
 
         if (thumbnailData != null && thumbnailData.thumbnail != null) {
@@ -189,7 +191,7 @@
             mPaint.setShader(shader);
         }
 
-        mOverlay.setTaskInfo(mThumbnailData, mMatrix);
+        mOverlay.setTaskInfo(mTask, mThumbnailData, mMatrix);
         invalidate();
     }