diff --git a/core/java/android/app/contextualsearch/CallbackToken.java b/core/java/android/app/contextualsearch/CallbackToken.java
index 378193f..94cdc73 100644
--- a/core/java/android/app/contextualsearch/CallbackToken.java
+++ b/core/java/android/app/contextualsearch/CallbackToken.java
@@ -68,6 +68,8 @@
      * invocations of this method will result in {@link OutcomeReceiver#onError} being called with
      * an {@link IllegalAccessException}.
      *
+     * Note that the callback could be invoked multiple times, e.g. in the case of split screen.
+     *
      * @param executor The executor which will be used to invoke the callback.
      * @param callback The callback which will be used to return {@link ContextualSearchState}
      *                 if/when it is available via {@link OutcomeReceiver#onResult}. It will also be
diff --git a/services/contextualsearch/java/com/android/server/contextualsearch/ContextualSearchManagerService.java b/services/contextualsearch/java/com/android/server/contextualsearch/ContextualSearchManagerService.java
index fdd1025..7f64786 100644
--- a/services/contextualsearch/java/com/android/server/contextualsearch/ContextualSearchManagerService.java
+++ b/services/contextualsearch/java/com/android/server/contextualsearch/ContextualSearchManagerService.java
@@ -17,6 +17,8 @@
 package com.android.server.contextualsearch;
 
 import static android.Manifest.permission.ACCESS_CONTEXTUAL_SEARCH;
+import static android.app.AppOpsManager.OP_ASSIST_SCREENSHOT;
+import static android.app.AppOpsManager.OP_ASSIST_STRUCTURE;
 import static android.content.Context.CONTEXTUAL_SEARCH_SERVICE;
 import static android.content.Intent.FLAG_ACTIVITY_NEW_TASK;
 import static android.content.Intent.FLAG_ACTIVITY_NO_ANIMATION;
@@ -29,11 +31,17 @@
 
 import static com.android.server.contextualsearch.flags.Flags.enableExcludePersistentUi;
 
+import static com.android.server.wm.ActivityTaskManagerInternal.ASSIST_KEY_CONTENT;
+import static com.android.server.wm.ActivityTaskManagerInternal.ASSIST_KEY_STRUCTURE;
+
 import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.annotation.RequiresPermission;
 import android.app.ActivityOptions;
+import android.app.AppOpsManager;
 import android.app.admin.DevicePolicyManagerInternal;
+import android.app.assist.AssistContent;
+import android.app.assist.AssistStructure;
 import android.app.contextualsearch.CallbackToken;
 import android.app.contextualsearch.ContextualSearchManager;
 import android.app.contextualsearch.ContextualSearchState;
@@ -42,6 +50,7 @@
 import android.content.ComponentName;
 import android.content.Context;
 import android.content.Intent;
+import android.content.pm.PackageManagerInternal;
 import android.content.pm.ResolveInfo;
 import android.graphics.Bitmap;
 import android.os.Binder;
@@ -54,15 +63,19 @@
 import android.os.Process;
 import android.os.RemoteException;
 import android.os.ResultReceiver;
+import android.os.ServiceManager;
 import android.os.ShellCallback;
 import android.util.Log;
 import android.util.Slog;
+import android.view.IWindowManager;
 import android.window.ScreenCapture;
 
 import com.android.internal.R;
 import com.android.internal.annotations.GuardedBy;
 import com.android.server.LocalServices;
 import com.android.server.SystemService;
+import com.android.server.am.AssistDataRequester;
+import com.android.server.am.AssistDataRequester.AssistDataRequesterCallbacks;
 import com.android.server.wm.ActivityAssistInfo;
 import com.android.server.wm.ActivityTaskManagerInternal;
 import com.android.server.wm.WindowManagerInternal;
@@ -74,19 +87,66 @@
 import java.util.Set;
 
 public class ContextualSearchManagerService extends SystemService {
-
+    private static final String TAG = ContextualSearchManagerService.class.getSimpleName();
     private static final int MSG_RESET_TEMPORARY_PACKAGE = 0;
     private static final int MAX_TEMP_PACKAGE_DURATION_MS = 1_000 * 60 * 2; // 2 minutes
+
     private final Context mContext;
     private final ActivityTaskManagerInternal mAtmInternal;
+    private final PackageManagerInternal mPackageManager;
     private final WindowManagerInternal mWmInternal;
     private final DevicePolicyManagerInternal mDpmInternal;
+    private final Object mLock = new Object();
+    private final AssistDataRequester mAssistDataRequester;
 
-    private Handler mTemporaryHandler;
+    private final AssistDataRequesterCallbacks mAssistDataCallbacks =
+            new AssistDataRequesterCallbacks() {
+                @Override
+                public boolean canHandleReceivedAssistDataLocked() {
+                    synchronized (mLock) {
+                        return mStateCallback != null;
+                    }
+                }
+
+                @Override
+                public void onAssistDataReceivedLocked(
+                        final Bundle data,
+                        final int activityIndex,
+                        final int activityCount) {
+                    final IContextualSearchCallback callback;
+                    synchronized (mLock) {
+                        callback = mStateCallback;
+                    }
+
+                    if (callback != null) {
+                        try {
+                            callback.onResult(new ContextualSearchState(
+                                    data.getParcelable(ASSIST_KEY_STRUCTURE, AssistStructure.class),
+                                    data.getParcelable(ASSIST_KEY_CONTENT, AssistContent.class),
+                                    data));
+                        } catch (RemoteException e) {
+                            Log.e(TAG, "Error invoking ContextualSearchCallback", e);
+                        }
+                    } else {
+                        Log.w(TAG, "Callback went away!");
+                    }
+                }
+
+                @Override
+                public void onAssistRequestCompleted() {
+                    synchronized (mLock) {
+                        mStateCallback = null;
+                    }
+                }
+            };
 
     @GuardedBy("this")
+    private Handler mTemporaryHandler;
+    @GuardedBy("this")
     private String mTemporaryPackage = null;
-    private static final String TAG = ContextualSearchManagerService.class.getSimpleName();
+
+    @GuardedBy("mLock")
+    private IContextualSearchCallback mStateCallback;
 
     public ContextualSearchManagerService(@NonNull Context context) {
         super(context);
@@ -94,8 +154,14 @@
         mContext = context;
         mAtmInternal = Objects.requireNonNull(
                 LocalServices.getService(ActivityTaskManagerInternal.class));
+        mPackageManager = LocalServices.getService(PackageManagerInternal.class);
         mWmInternal = Objects.requireNonNull(LocalServices.getService(WindowManagerInternal.class));
         mDpmInternal = LocalServices.getService(DevicePolicyManagerInternal.class);
+        mAssistDataRequester = new AssistDataRequester(
+                mContext,
+                IWindowManager.Stub.asInterface(ServiceManager.getService(Context.WINDOW_SERVICE)),
+                mContext.getSystemService(AppOpsManager.class),
+                mAssistDataCallbacks, mLock, OP_ASSIST_STRUCTURE, OP_ASSIST_SCREENSHOT);
     }
 
     @Override
@@ -184,18 +250,37 @@
         launchIntent.putExtra(ContextualSearchManager.EXTRA_TOKEN, mToken);
         boolean isAssistDataAllowed = mAtmInternal.isAssistDataAllowed();
         final List<ActivityAssistInfo> records = mAtmInternal.getTopVisibleActivities();
+        final List<IBinder> activityTokens = new ArrayList<>(records.size());
         ArrayList<String> visiblePackageNames = new ArrayList<>();
         boolean isManagedProfileVisible = false;
         for (ActivityAssistInfo record : records) {
             // Add the package name to the list only if assist data is allowed.
             if (isAssistDataAllowed) {
                 visiblePackageNames.add(record.getComponentName().getPackageName());
+                activityTokens.add(record.getActivityToken());
             }
             if (mDpmInternal != null
                     && mDpmInternal.isUserOrganizationManaged(record.getUserId())) {
                 isManagedProfileVisible = true;
             }
         }
+        if (isAssistDataAllowed) {
+            try {
+                final String csPackage = Objects.requireNonNull(launchIntent.getPackage());
+                final int csUid = mPackageManager.getPackageUid(csPackage, 0, 0);
+                mAssistDataRequester.requestAssistData(
+                        activityTokens,
+                        /* fetchData */ true,
+                        /* fetchScreenshot */ false,
+                        /* allowFetchData */ true,
+                        /* allowFetchScreenshot */ false,
+                        csUid,
+                        csPackage,
+                        null);
+            } catch (Exception e) {
+                Log.e(TAG, "Could not request assist data", e);
+            }
+        }
         final ScreenCapture.ScreenshotHardwareBuffer shb;
         if (mWmInternal != null) {
             if (enableExcludePersistentUi()) {
@@ -276,6 +361,7 @@
             synchronized (this) {
                 if (DEBUG_USER) Log.d(TAG, "startContextualSearch");
                 enforcePermission("startContextualSearch");
+                mAssistDataRequester.cancel();
                 mToken = new CallbackToken();
                 // We get the launch intent with the system server's identity because the system
                 // server has READ_FRAME_BUFFER permission to get the screenshot and because only
@@ -310,17 +396,19 @@
                 return;
             }
             mToken = null;
-            // Process data request
-            try {
-                callback.onResult(new ContextualSearchState(null, null, Bundle.EMPTY));
-            } catch (RemoteException e) {
-                Log.e(TAG, "Could not invoke onResult callback", e);
+            synchronized (mLock) {
+                mStateCallback = callback;
             }
+            mAssistDataRequester.processPendingAssistData();
         }
 
-        public void onShellCommand(@Nullable FileDescriptor in, @Nullable FileDescriptor out,
-                @Nullable FileDescriptor err, @NonNull String[] args,
-                @Nullable ShellCallback callback, @NonNull ResultReceiver resultReceiver) {
+        public void onShellCommand(
+                @Nullable FileDescriptor in,
+                @Nullable FileDescriptor out,
+                @Nullable FileDescriptor err,
+                @NonNull String[] args,
+                @Nullable ShellCallback callback,
+                @NonNull ResultReceiver resultReceiver) {
             new ContextualSearchManagerShellCommand(ContextualSearchManagerService.this)
                     .exec(this, in, out, err, args, callback, resultReceiver);
         }