Merge "Update Backlinks to capture data from all tasks" into main
diff --git a/packages/SystemUI/res/layout/app_clips_backlinks_drop_down_entry.xml b/packages/SystemUI/res/layout/app_clips_backlinks_drop_down_entry.xml
new file mode 100644
index 0000000..7eab340
--- /dev/null
+++ b/packages/SystemUI/res/layout/app_clips_backlinks_drop_down_entry.xml
@@ -0,0 +1,24 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!--
+  ~ Copyright (C) 2024 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.
+  -->
+<TextView xmlns:android="http://schemas.android.com/apk/res/android"
+    android:layout_width="match_parent"
+    android:layout_height="48dp"
+    android:drawablePadding="4dp"
+    android:ellipsize="end"
+    android:gravity="center_vertical"
+    android:paddingHorizontal="8dp"
+    android:textColor="?android:textColorSecondary" />
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/appclips/AppClipsActivity.java b/packages/SystemUI/src/com/android/systemui/screenshot/appclips/AppClipsActivity.java
index 8feefa4..9db1f24 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/appclips/AppClipsActivity.java
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/appclips/AppClipsActivity.java
@@ -46,13 +46,17 @@
 import android.os.ResultReceiver;
 import android.util.Log;
 import android.view.View;
+import android.view.ViewGroup;
+import android.widget.ArrayAdapter;
 import android.widget.Button;
 import android.widget.CheckBox;
 import android.widget.ImageView;
+import android.widget.ListPopupWindow;
 import android.widget.TextView;
 
 import androidx.activity.ComponentActivity;
 import androidx.annotation.Nullable;
+import androidx.appcompat.content.res.AppCompatResources;
 import androidx.core.graphics.Insets;
 import androidx.core.view.ViewCompat;
 import androidx.core.view.WindowInsetsCompat;
@@ -67,6 +71,7 @@
 import com.android.systemui.screenshot.scroll.CropView;
 import com.android.systemui.settings.UserTracker;
 
+import java.util.List;
 import java.util.Set;
 
 import javax.inject.Inject;
@@ -92,6 +97,7 @@
 
     private static final String TAG = AppClipsActivity.class.getSimpleName();
     private static final ApplicationInfoFlags APPLICATION_INFO_FLAGS = ApplicationInfoFlags.of(0);
+    private static final int DRAWABLE_END = 2;
 
     private final AppClipsViewModel.Factory mViewModelFactory;
     private final PackageManager mPackageManager;
@@ -192,6 +198,7 @@
         mViewModel.getResultLiveData().observe(this, this::setResultThenFinish);
         mViewModel.getErrorLiveData().observe(this, this::setErrorThenFinish);
         mViewModel.getBacklinksLiveData().observe(this, this::setBacklinksData);
+        mViewModel.mSelectedBacklinksLiveData.observe(this, this::updateBacklinksTextView);
 
         if (savedInstanceState == null) {
             int displayId = getDisplayId();
@@ -305,8 +312,8 @@
 
         if (mBacklinksIncludeDataCheckBox.getVisibility() == View.VISIBLE
                 && mBacklinksIncludeDataCheckBox.isChecked()
-                && mViewModel.getBacklinksLiveData().getValue() != null) {
-            ClipData backlinksData = mViewModel.getBacklinksLiveData().getValue().getClipData();
+                && mViewModel.mSelectedBacklinksLiveData.getValue() != null) {
+            ClipData backlinksData = mViewModel.mSelectedBacklinksLiveData.getValue().getClipData();
             data.putParcelable(EXTRA_CLIP_DATA, backlinksData);
 
             DebugLogger.INSTANCE.logcatMessage(this,
@@ -330,18 +337,80 @@
         finish();
     }
 
-    private void setBacklinksData(InternalBacklinksData backlinksData) {
+    private void setBacklinksData(List<InternalBacklinksData> backlinksData) {
         mBacklinksIncludeDataCheckBox.setVisibility(View.VISIBLE);
         mBacklinksDataTextView.setVisibility(
                 mBacklinksIncludeDataCheckBox.isChecked() ? View.VISIBLE : View.GONE);
 
-        mBacklinksDataTextView.setText(backlinksData.getClipData().getDescription().getLabel());
+        // Set up the dropdown when multiple backlinks are available.
+        if (backlinksData.size() > 1) {
+            setUpListPopupWindow(backlinksData, mBacklinksDataTextView);
+        }
+    }
 
+    private void setUpListPopupWindow(List<InternalBacklinksData> backlinksData, View anchor) {
+        ListPopupWindow listPopupWindow = new ListPopupWindow(this);
+        listPopupWindow.setAnchorView(anchor);
+        listPopupWindow.setOverlapAnchor(true);
+        listPopupWindow.setBackgroundDrawable(
+                AppCompatResources.getDrawable(this, R.drawable.backlinks_rounded_rectangle));
+        listPopupWindow.setOnItemClickListener((parent, view, position, id) -> {
+            mViewModel.mSelectedBacklinksLiveData.setValue(backlinksData.get(position));
+            listPopupWindow.dismiss();
+        });
+
+        ArrayAdapter<InternalBacklinksData> adapter = new ArrayAdapter<>(this,
+                R.layout.app_clips_backlinks_drop_down_entry) {
+            @Override
+            public View getView(int position, @Nullable View convertView, ViewGroup parent) {
+                TextView itemView = (TextView) super.getView(position, convertView, parent);
+                InternalBacklinksData data = backlinksData.get(position);
+                itemView.setText(data.getClipData().getDescription().getLabel());
+
+                Drawable icon = data.getAppIcon();
+                icon.setBounds(createBacklinksTextViewDrawableBounds());
+                itemView.setCompoundDrawablesRelative(/* start= */ icon, /* top= */ null,
+                        /* end= */ null, /* bottom= */ null);
+
+                return itemView;
+            }
+        };
+        adapter.addAll(backlinksData);
+        listPopupWindow.setAdapter(adapter);
+
+        mBacklinksDataTextView.setOnClickListener(unused -> listPopupWindow.show());
+    }
+
+    /**
+     * Updates the {@link #mBacklinksDataTextView} with the currently selected
+     * {@link InternalBacklinksData}. The {@link AppClipsViewModel#getBacklinksLiveData()} is
+     * expected to be already set when this method is called.
+     */
+    private void updateBacklinksTextView(InternalBacklinksData backlinksData) {
+        mBacklinksDataTextView.setText(backlinksData.getClipData().getDescription().getLabel());
         Drawable appIcon = backlinksData.getAppIcon();
-        int size = getResources().getDimensionPixelSize(R.dimen.appclips_backlinks_icon_size);
-        appIcon.setBounds(/* left= */ 0, /* top= */ 0, /* right= */ size, /* bottom= */ size);
+        Rect compoundDrawableBounds = createBacklinksTextViewDrawableBounds();
+        appIcon.setBounds(compoundDrawableBounds);
+
+        // Try to reuse the dropdown down arrow icon if available, will be null if never set.
+        Drawable dropDownIcon = mBacklinksDataTextView.getCompoundDrawablesRelative()[DRAWABLE_END];
+        if (mViewModel.getBacklinksLiveData().getValue().size() > 1 && dropDownIcon == null) {
+            // Set up the dropdown down arrow drawable only if it is required.
+            dropDownIcon = AppCompatResources.getDrawable(this, R.drawable.arrow_pointing_down);
+            dropDownIcon.setBounds(compoundDrawableBounds);
+            dropDownIcon.setTint(Utils.getColorAttr(this,
+                    android.R.attr.textColorSecondary).getDefaultColor());
+        }
+
         mBacklinksDataTextView.setCompoundDrawablesRelative(/* start= */ appIcon, /* top= */
-                null, /* end= */ null, /* bottom= */ null);
+                null, /* end= */ dropDownIcon, /* bottom= */ null);
+    }
+
+    private Rect createBacklinksTextViewDrawableBounds() {
+        int size = getResources().getDimensionPixelSize(R.dimen.appclips_backlinks_icon_size);
+        Rect bounds = new Rect();
+        bounds.set(/* left= */ 0, /* top= */ 0, /* right= */ size, /* bottom= */ size);
+        return bounds;
     }
 
     private void setError(int errorCode) {
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/appclips/AppClipsViewModel.java b/packages/SystemUI/src/com/android/systemui/screenshot/appclips/AppClipsViewModel.java
index bd9e295..3530b3f 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/appclips/AppClipsViewModel.java
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/appclips/AppClipsViewModel.java
@@ -21,12 +21,11 @@
 import static android.content.Intent.CAPTURE_CONTENT_FOR_NOTE_FAILED;
 import static android.content.Intent.CATEGORY_LAUNCHER;
 
-import static com.google.common.util.concurrent.Futures.withTimeout;
-
 import static java.util.concurrent.Executors.newSingleThreadScheduledExecutor;
 
 import android.app.ActivityTaskManager.RootTaskInfo;
 import android.app.IActivityTaskManager;
+import android.app.TaskInfo;
 import android.app.WindowConfiguration;
 import android.app.assist.AssistContent;
 import android.content.ClipData;
@@ -41,7 +40,6 @@
 import android.graphics.RenderNode;
 import android.graphics.drawable.Drawable;
 import android.net.Uri;
-import android.os.RemoteException;
 import android.os.UserHandle;
 import android.util.Log;
 import android.view.Display;
@@ -94,7 +92,8 @@
     private final MutableLiveData<Bitmap> mScreenshotLiveData;
     private final MutableLiveData<Uri> mResultLiveData;
     private final MutableLiveData<Integer> mErrorLiveData;
-    private final MutableLiveData<InternalBacklinksData> mBacklinksLiveData;
+    private final MutableLiveData<List<InternalBacklinksData>> mBacklinksLiveData;
+    final MutableLiveData<InternalBacklinksData> mSelectedBacklinksLiveData;
 
     private AppClipsViewModel(AppClipsCrossProcessHelper appClipsCrossProcessHelper,
             ImageExporter imageExporter, IActivityTaskManager atmService,
@@ -112,6 +111,7 @@
         mResultLiveData = new MutableLiveData<>();
         mErrorLiveData = new MutableLiveData<>();
         mBacklinksLiveData = new MutableLiveData<>();
+        mSelectedBacklinksLiveData = new MutableLiveData<>();
     }
 
     /**
@@ -135,10 +135,11 @@
     /**
      * Triggers the Backlinks flow which:
      * <ul>
-     *     <li>Evaluates the task to query.
-     *     <li>Requests {@link AssistContent} from that task.
-     *     <li>Transforms the {@link AssistContent} into {@link ClipData} for Backlinks.
-     *     <li>The {@link ClipData} is reported to activity via {@link #getBacklinksLiveData()}.
+     *     <li>Evaluates the tasks to query.
+     *     <li>Requests {@link AssistContent} from all valid tasks.
+     *     <li>Transforms {@link AssistContent} into {@link InternalBacklinksData} for Backlinks.
+     *     <li>The {@link InternalBacklinksData}s are reported to activity via
+     *     {@link #getBacklinksLiveData()}.
      * </ul>
      *
      * @param taskIdsToIgnore id of the tasks to ignore when querying for {@link AssistContent}
@@ -146,24 +147,24 @@
      */
     void triggerBacklinks(Set<Integer> taskIdsToIgnore, int displayId) {
         DebugLogger.INSTANCE.logcatMessage(this, () -> "Backlinks triggered");
-        mBgExecutor.execute(() -> {
-            ListenableFuture<InternalBacklinksData> backlinksData = getBacklinksData(
-                    taskIdsToIgnore, displayId);
-            Futures.addCallback(backlinksData, new FutureCallback<>() {
-                @Override
-                public void onSuccess(@Nullable InternalBacklinksData result) {
-                    if (result != null) {
-                        mBacklinksLiveData.setValue(result);
-                    }
+        ListenableFuture<List<InternalBacklinksData>> backlinksData = getAllAvailableBacklinks(
+                taskIdsToIgnore, displayId);
+        Futures.addCallback(backlinksData, new FutureCallback<>() {
+            @Override
+            public void onSuccess(@Nullable List<InternalBacklinksData> result) {
+                if (result != null && !result.isEmpty()) {
+                    // Set the list of backlinks before setting the selected backlink as this is
+                    // required when updating the backlink data text view.
+                    mBacklinksLiveData.setValue(result);
+                    mSelectedBacklinksLiveData.setValue(result.get(0));
                 }
+            }
 
-                @Override
-                public void onFailure(Throwable t) {
-                    Log.e(TAG, "Error querying for Backlinks data", t);
-                }
-            }, mMainExecutor);
-
-        });
+            @Override
+            public void onFailure(Throwable t) {
+                Log.e(TAG, "Error querying for Backlinks data", t);
+            }
+        }, mMainExecutor);
     }
 
     /** Returns a {@link LiveData} that holds the captured screenshot. */
@@ -184,8 +185,11 @@
         return mErrorLiveData;
     }
 
-    /** Returns a {@link LiveData} that holds Backlinks data in {@link InternalBacklinksData}. */
-    LiveData<InternalBacklinksData> getBacklinksLiveData() {
+    /**
+     * Returns a {@link LiveData} that holds all the available Backlinks data and the currently
+     * selected index for displaying the Backlinks in the UI.
+     */
+    LiveData<List<InternalBacklinksData>> getBacklinksLiveData() {
         return mBacklinksLiveData;
     }
 
@@ -230,26 +234,58 @@
         return HardwareRenderer.createHardwareBitmap(output, bounds.width(), bounds.height());
     }
 
-    private ListenableFuture<InternalBacklinksData> getBacklinksData(Set<Integer> taskIdsToIgnore,
-            int displayId) {
-        return getAllRootTaskInfosOnDisplay(displayId)
-                .stream()
-                .filter(taskInfo -> shouldIncludeTask(taskInfo, taskIdsToIgnore))
-                .findFirst()
-                .map(this::getBacklinksDataForTaskId)
-                .orElse(Futures.immediateFuture(null));
+    private ListenableFuture<List<InternalBacklinksData>> getAllAvailableBacklinks(
+            Set<Integer> taskIdsToIgnore, int displayId) {
+        ListenableFuture<List<TaskInfo>> allTasksOnDisplayFuture = getAllTasksOnDisplay(displayId);
+
+        ListenableFuture<List<ListenableFuture<InternalBacklinksData>>> backlinksNestedListFuture =
+                Futures.transform(allTasksOnDisplayFuture, allTasksOnDisplay ->
+                        allTasksOnDisplay
+                                .stream()
+                                .filter(taskInfo -> shouldIncludeTask(taskInfo, taskIdsToIgnore))
+                                .map(this::getBacklinksDataForTaskInfo)
+                                .toList(),
+                        mBgExecutor);
+
+        return Futures.transformAsync(backlinksNestedListFuture, Futures::allAsList, mBgExecutor);
     }
 
-    private List<RootTaskInfo> getAllRootTaskInfosOnDisplay(int displayId) {
-        try {
-            return mAtmService.getAllRootTaskInfosOnDisplay(displayId);
-        } catch (RemoteException e) {
-            Log.e(TAG, String.format("Error while querying for tasks on display %d", displayId), e);
-            return Collections.emptyList();
-        }
+    /**
+     * Returns all tasks on a given display after querying {@link IActivityTaskManager} from the
+     * {@link #mBgExecutor}.
+     */
+    private ListenableFuture<List<TaskInfo>> getAllTasksOnDisplay(int displayId) {
+        SettableFuture<List<TaskInfo>> recentTasksFuture = SettableFuture.create();
+        mBgExecutor.execute(() -> {
+            try {
+                // Directly call into ActivityTaskManagerService instead of going through WMShell
+                // because WMShell is only available in the main SysUI process and App Clips runs
+                // in its own separate process as it deals with bitmaps.
+                List<TaskInfo> allTasksOnDisplay = mAtmService.getTasks(
+                                /* maxNum= */ Integer.MAX_VALUE,
+                                // PIP tasks are not visible in recents. So _not_ filtering for
+                                // tasks that are only visible in recents.
+                                /* filterOnlyVisibleRecents= */ false,
+                                /* keepIntentExtra= */ false,
+                                displayId)
+                        .stream()
+                        .map(runningTaskInfo -> (TaskInfo) runningTaskInfo)
+                        .toList();
+                recentTasksFuture.set(allTasksOnDisplay);
+            } catch (Exception e) {
+                Log.e(TAG, String.format("Error getting all tasks on displayId %d", displayId), e);
+                recentTasksFuture.set(Collections.emptyList());
+            }
+        });
+
+        return withTimeout(recentTasksFuture);
     }
 
-    private boolean shouldIncludeTask(RootTaskInfo taskInfo, Set<Integer> taskIdsToIgnore) {
+    /**
+     * Returns whether the app represented by the provided {@link TaskInfo} should be included for
+     * querying for {@link AssistContent}.
+     */
+    private boolean shouldIncludeTask(TaskInfo taskInfo, Set<Integer> taskIdsToIgnore) {
         DebugLogger.INSTANCE.logcatMessage(this,
                 () -> String.format("shouldIncludeTask taskId %d; topActivity %s", taskInfo.taskId,
                         taskInfo.topActivity));
@@ -262,11 +298,14 @@
                 && taskInfo.numActivities > 0
                 && taskInfo.topActivity != null
                 && taskInfo.topActivityInfo != null
-                && taskInfo.childTaskIds.length > 0
                 && taskInfo.getActivityType() == WindowConfiguration.ACTIVITY_TYPE_STANDARD
                 && canAppStartThroughLauncher(taskInfo.topActivity.getPackageName());
     }
 
+    /**
+     * Returns whether the app represented by the provided {@code packageName} can be launched
+     * through the all apps tray by a user.
+     */
     private boolean canAppStartThroughLauncher(String packageName) {
         // Use Intent.resolveActivity API to check if the intent resolves as that is what Android
         // uses internally when apps use Context.startActivity.
@@ -274,8 +313,12 @@
                 != null;
     }
 
-    private ListenableFuture<InternalBacklinksData> getBacklinksDataForTaskId(
-            RootTaskInfo taskInfo) {
+    /**
+     * Returns an {@link InternalBacklinksData} that represents the Backlink data internally, which
+     * is captured by querying the system using {@link TaskInfo#taskId}.
+     */
+    private ListenableFuture<InternalBacklinksData> getBacklinksDataForTaskInfo(
+            TaskInfo taskInfo) {
         DebugLogger.INSTANCE.logcatMessage(this,
                 () -> String.format("getBacklinksDataForTaskId for taskId %d; topActivity %s",
                         taskInfo.taskId, taskInfo.topActivity));
@@ -284,7 +327,13 @@
         int taskId = taskInfo.taskId;
         mAssistContentRequester.requestAssistContent(taskId, assistContent ->
                 backlinksData.set(getBacklinksDataFromAssistContent(taskInfo, assistContent)));
-        return withTimeout(backlinksData, 5L, TimeUnit.SECONDS, newSingleThreadScheduledExecutor());
+        return withTimeout(backlinksData);
+    }
+
+    /** Returns the same {@link ListenableFuture} but with a 5 {@link TimeUnit#SECONDS} timeout. */
+    private static <V> ListenableFuture<V> withTimeout(ListenableFuture<V> future) {
+        return Futures.withTimeout(future, 5L, TimeUnit.SECONDS,
+                newSingleThreadScheduledExecutor());
     }
 
     /**
@@ -306,7 +355,7 @@
      * @param content the {@link AssistContent} to map into Backlinks {@link ClipData}.
      * @return {@link InternalBacklinksData} that represents the Backlinks data along with app icon.
      */
-    private InternalBacklinksData getBacklinksDataFromAssistContent(RootTaskInfo taskInfo,
+    private InternalBacklinksData getBacklinksDataFromAssistContent(TaskInfo taskInfo,
             @Nullable AssistContent content) {
         DebugLogger.INSTANCE.logcatMessage(this,
                 () -> String.format("getBacklinksDataFromAssistContent taskId %d; topActivity %s",
@@ -365,7 +414,7 @@
         return resolvedComponent.getPackageName().equals(requiredPackageName);
     }
 
-    private String getAppNameOfTask(RootTaskInfo taskInfo) {
+    private String getAppNameOfTask(TaskInfo taskInfo) {
         return taskInfo.topActivityInfo.loadLabel(mPackageManager).toString();
     }
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/screenshot/appclips/AppClipsActivityTest.java b/packages/SystemUI/tests/src/com/android/systemui/screenshot/appclips/AppClipsActivityTest.java
index 9986205..a8d5008 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/screenshot/appclips/AppClipsActivityTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/screenshot/appclips/AppClipsActivityTest.java
@@ -17,6 +17,7 @@
 package com.android.systemui.screenshot.appclips;
 
 import static android.app.Activity.RESULT_OK;
+import static android.app.ActivityManager.RunningTaskInfo;
 import static android.app.WindowConfiguration.ACTIVITY_TYPE_STANDARD;
 
 import static com.android.systemui.screenshot.appclips.AppClipsEvent.SCREENSHOT_FOR_NOTE_ACCEPTED;
@@ -32,7 +33,6 @@
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
-import android.app.ActivityTaskManager.RootTaskInfo;
 import android.app.IActivityTaskManager;
 import android.app.assist.AssistContent;
 import android.content.ComponentName;
@@ -103,7 +103,7 @@
     private static final String BACKLINKS_TASK_APP_NAME = "Backlinks app";
     private static final String BACKLINKS_TASK_PACKAGE_NAME = "backlinksTaskPackageName";
 
-    private static final RootTaskInfo TASK_THAT_SUPPORTS_BACKLINKS =
+    private static final RunningTaskInfo TASK_THAT_SUPPORTS_BACKLINKS =
             createTaskInfoForBacklinksTask();
     private static final AssistContent ASSIST_CONTENT_FOR_BACKLINKS_TASK =
             createAssistContentForBacklinksTask();
@@ -233,6 +233,10 @@
         assertThat(backlinksData.getText().toString()).isEqualTo(BACKLINKS_TASK_APP_NAME);
         assertThat(backlinksData.getCompoundDrawablesRelative()[0]).isEqualTo(FAKE_DRAWABLE);
 
+        // Verify dropdown icon is not shown and there are no click listeners on text view.
+        assertThat(backlinksData.getCompoundDrawablesRelative()[2]).isNull();
+        assertThat(backlinksData.hasOnClickListeners()).isFalse();
+
         CheckBox backlinksIncludeData = mActivity.findViewById(R.id.backlinks_include_data);
         assertThat(backlinksIncludeData.getVisibility()).isEqualTo(View.VISIBLE);
         assertThat(backlinksIncludeData.getText().toString())
@@ -258,20 +262,71 @@
         assertThat(backlinksData.getVisibility()).isEqualTo(View.GONE);
     }
 
+    @Test
+    @EnableFlags(Flags.FLAG_APP_CLIPS_BACKLINKS)
+    public void appClipsLaunched_backlinks_multipleBacklinksAvailable_defaultShown()
+            throws RemoteException {
+        // Set up mocking for multiple backlinks.
+        ResolveInfo resolveInfo1 = createBacklinksTaskResolveInfo();
+
+        int taskId2 = BACKLINKS_TASK_ID + 2;
+        String package2 = BACKLINKS_TASK_PACKAGE_NAME + 2;
+        String appName2 = BACKLINKS_TASK_APP_NAME + 2;
+
+        ResolveInfo resolveInfo2 = createBacklinksTaskResolveInfo();
+        ActivityInfo activityInfo2 = resolveInfo2.activityInfo;
+        activityInfo2.name = appName2;
+        activityInfo2.packageName = package2;
+        activityInfo2.applicationInfo.packageName = package2;
+        RunningTaskInfo runningTaskInfo2 = createTaskInfoForBacklinksTask();
+        runningTaskInfo2.taskId = taskId2;
+        runningTaskInfo2.topActivity = new ComponentName(package2, "backlinksClass");
+        runningTaskInfo2.topActivityInfo = resolveInfo2.activityInfo;
+        runningTaskInfo2.baseIntent = new Intent().setComponent(runningTaskInfo2.topActivity);
+
+        when(mAtmService.getTasks(eq(Integer.MAX_VALUE), eq(false), eq(false),
+                mDisplayIdCaptor.capture()))
+                .thenReturn(List.of(TASK_THAT_SUPPORTS_BACKLINKS, runningTaskInfo2));
+        when(mPackageManager.resolveActivity(any(Intent.class), anyInt())).thenReturn(resolveInfo1,
+                resolveInfo1, resolveInfo1, resolveInfo2, resolveInfo2, resolveInfo2);
+        when(mPackageManager.loadItemIcon(any(), any())).thenReturn(FAKE_DRAWABLE);
+
+        // Using same AssistContent data for both tasks.
+        mockForAssistContent(ASSIST_CONTENT_FOR_BACKLINKS_TASK, BACKLINKS_TASK_ID);
+        mockForAssistContent(ASSIST_CONTENT_FOR_BACKLINKS_TASK, taskId2);
+
+        // Mocking complete, trigger backlinks.
+        launchActivity();
+        waitForIdleSync();
+
+        // Verify default backlink shown to user and text view has on click listener.
+        TextView backlinksData = mActivity.findViewById(R.id.backlinks_data);
+        assertThat(backlinksData.getText().toString()).isEqualTo(BACKLINKS_TASK_APP_NAME);
+        assertThat(backlinksData.hasOnClickListeners()).isTrue();
+
+        // Verify dropdown icon is not null.
+        assertThat(backlinksData.getCompoundDrawablesRelative()[2]).isNotNull();
+    }
+
     private void setUpMocksForBacklinks() throws RemoteException {
-        when(mAtmService.getAllRootTaskInfosOnDisplay(mDisplayIdCaptor.capture()))
+        when(mAtmService.getTasks(eq(Integer.MAX_VALUE), eq(false), eq(false),
+                mDisplayIdCaptor.capture()))
                 .thenReturn(List.of(TASK_THAT_SUPPORTS_BACKLINKS));
-        doAnswer(invocation -> {
-            AssistContentRequester.Callback callback = invocation.getArgument(1);
-            callback.onAssistContentAvailable(ASSIST_CONTENT_FOR_BACKLINKS_TASK);
-            return null;
-        }).when(mAssistContentRequester).requestAssistContent(eq(BACKLINKS_TASK_ID), any());
+        mockForAssistContent(ASSIST_CONTENT_FOR_BACKLINKS_TASK, BACKLINKS_TASK_ID);
         when(mPackageManager
                 .resolveActivity(any(Intent.class), anyInt()))
                 .thenReturn(createBacklinksTaskResolveInfo());
         when(mPackageManager.loadItemIcon(any(), any())).thenReturn(FAKE_DRAWABLE);
     }
 
+    private void mockForAssistContent(AssistContent expected, int taskId) {
+        doAnswer(invocation -> {
+            AssistContentRequester.Callback callback = invocation.getArgument(1);
+            callback.onAssistContentAvailable(expected);
+            return null;
+        }).when(mAssistContentRequester).requestAssistContent(eq(taskId), any());
+    }
+
     private void launchActivity() {
         launchActivity(createResultReceiver(FAKE_CONSUMER));
     }
@@ -319,8 +374,8 @@
         return resolveInfo;
     }
 
-    private static RootTaskInfo createTaskInfoForBacklinksTask() {
-        RootTaskInfo taskInfo = new RootTaskInfo();
+    private static RunningTaskInfo createTaskInfoForBacklinksTask() {
+        RunningTaskInfo taskInfo = new RunningTaskInfo();
         taskInfo.taskId = BACKLINKS_TASK_ID;
         taskInfo.isVisible = true;
         taskInfo.isRunning = true;
@@ -328,7 +383,6 @@
         taskInfo.topActivity = new ComponentName(BACKLINKS_TASK_PACKAGE_NAME, "backlinksClass");
         taskInfo.topActivityInfo = createBacklinksTaskResolveInfo().activityInfo;
         taskInfo.baseIntent = new Intent().setComponent(taskInfo.topActivity);
-        taskInfo.childTaskIds = new int[]{BACKLINKS_TASK_ID + 1};
         taskInfo.configuration.windowConfiguration.setActivityType(ACTIVITY_TYPE_STANDARD);
         return taskInfo;
     }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/screenshot/appclips/AppClipsViewModelTest.java b/packages/SystemUI/tests/src/com/android/systemui/screenshot/appclips/AppClipsViewModelTest.java
index 193d29c..178547e 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/screenshot/appclips/AppClipsViewModelTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/screenshot/appclips/AppClipsViewModelTest.java
@@ -37,7 +37,7 @@
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
-import android.app.ActivityTaskManager.RootTaskInfo;
+import android.app.ActivityManager.RunningTaskInfo;
 import android.app.IActivityTaskManager;
 import android.app.assist.AssistContent;
 import android.content.ClipData;
@@ -107,7 +107,7 @@
         mPackageManagerIntentCaptor = ArgumentCaptor.forClass(Intent.class);
 
         // Set up mocking for backlinks.
-        when(mAtmService.getAllRootTaskInfosOnDisplay(DEFAULT_DISPLAY))
+        when(mAtmService.getTasks(Integer.MAX_VALUE, false, false, DEFAULT_DISPLAY))
                 .thenReturn(List.of(createTaskInfoForBacklinksTask()));
         when(mPackageManager.resolveActivity(mPackageManagerIntentCaptor.capture(), anyInt()))
                 .thenReturn(createBacklinksTaskResolveInfo());
@@ -190,11 +190,7 @@
         Uri expectedUri = Uri.parse("https://developers.android.com");
         AssistContent contentWithUri = new AssistContent();
         contentWithUri.setWebUri(expectedUri);
-        doAnswer(invocation -> {
-            AssistContentRequester.Callback callback = invocation.getArgument(1);
-            callback.onAssistContentAvailable(contentWithUri);
-            return null;
-        }).when(mAssistContentRequester).requestAssistContent(eq(BACKLINKS_TASK_ID), any());
+        mockForAssistContent(contentWithUri, BACKLINKS_TASK_ID);
 
         mViewModel.triggerBacklinks(Collections.emptySet(), DEFAULT_DISPLAY);
         waitForIdleSync();
@@ -203,7 +199,7 @@
         assertThat(queriedIntent.getData()).isEqualTo(expectedUri);
         assertThat(queriedIntent.getAction()).isEqualTo(ACTION_VIEW);
 
-        InternalBacklinksData result = mViewModel.getBacklinksLiveData().getValue();
+        InternalBacklinksData result = mViewModel.mSelectedBacklinksLiveData.getValue();
         assertThat(result.getAppIcon()).isEqualTo(FAKE_DRAWABLE);
         ClipData clipData = result.getClipData();
         ClipDescription resultDescription = clipData.getDescription();
@@ -211,6 +207,8 @@
         assertThat(resultDescription.getMimeType(0)).isEqualTo(MIMETYPE_TEXT_URILIST);
         assertThat(clipData.getItemCount()).isEqualTo(1);
         assertThat(clipData.getItemAt(0).getUri()).isEqualTo(expectedUri);
+
+        assertThat(result).isEqualTo(mViewModel.getBacklinksLiveData().getValue().get(0));
     }
 
     @Test
@@ -218,12 +216,8 @@
         Uri expectedUri = Uri.parse("https://developers.android.com");
         AssistContent contentWithUri = new AssistContent();
         contentWithUri.setWebUri(expectedUri);
+        mockForAssistContent(contentWithUri, BACKLINKS_TASK_ID);
         resetPackageManagerMockingForUsingFallbackBacklinks();
-        doAnswer(invocation -> {
-            AssistContentRequester.Callback callback = invocation.getArgument(1);
-            callback.onAssistContentAvailable(contentWithUri);
-            return null;
-        }).when(mAssistContentRequester).requestAssistContent(eq(BACKLINKS_TASK_ID), any());
 
         mViewModel.triggerBacklinks(Collections.emptySet(), DEFAULT_DISPLAY);
         waitForIdleSync();
@@ -236,11 +230,7 @@
         Intent expectedIntent = new Intent().setPackage(BACKLINKS_TASK_PACKAGE_NAME);
         AssistContent contentWithAppProvidedIntent = new AssistContent();
         contentWithAppProvidedIntent.setIntent(expectedIntent);
-        doAnswer(invocation -> {
-            AssistContentRequester.Callback callback = invocation.getArgument(1);
-            callback.onAssistContentAvailable(contentWithAppProvidedIntent);
-            return null;
-        }).when(mAssistContentRequester).requestAssistContent(eq(BACKLINKS_TASK_ID), any());
+        mockForAssistContent(contentWithAppProvidedIntent, BACKLINKS_TASK_ID);
 
         mViewModel.triggerBacklinks(Collections.emptySet(), DEFAULT_DISPLAY);
         waitForIdleSync();
@@ -248,7 +238,7 @@
         Intent queriedIntent = mPackageManagerIntentCaptor.getValue();
         assertThat(queriedIntent.getPackage()).isEqualTo(expectedIntent.getPackage());
 
-        InternalBacklinksData result = mViewModel.getBacklinksLiveData().getValue();
+        InternalBacklinksData result = mViewModel.mSelectedBacklinksLiveData.getValue();
         assertThat(result.getAppIcon()).isEqualTo(FAKE_DRAWABLE);
         ClipData clipData = result.getClipData();
         ClipDescription resultDescription = clipData.getDescription();
@@ -263,12 +253,8 @@
         Intent expectedIntent = new Intent().setPackage(BACKLINKS_TASK_PACKAGE_NAME);
         AssistContent contentWithAppProvidedIntent = new AssistContent();
         contentWithAppProvidedIntent.setIntent(expectedIntent);
+        mockForAssistContent(contentWithAppProvidedIntent, BACKLINKS_TASK_ID);
         resetPackageManagerMockingForUsingFallbackBacklinks();
-        doAnswer(invocation -> {
-            AssistContentRequester.Callback callback = invocation.getArgument(1);
-            callback.onAssistContentAvailable(contentWithAppProvidedIntent);
-            return null;
-        }).when(mAssistContentRequester).requestAssistContent(eq(BACKLINKS_TASK_ID), any());
 
         mViewModel.triggerBacklinks(Collections.emptySet(), DEFAULT_DISPLAY);
         waitForIdleSync();
@@ -278,11 +264,7 @@
 
     @Test
     public void triggerBacklinks_shouldUpdateBacklinks_withMainLauncherIntent() {
-        doAnswer(invocation -> {
-            AssistContentRequester.Callback callback = invocation.getArgument(1);
-            callback.onAssistContentAvailable(EMPTY_ASSIST_CONTENT);
-            return null;
-        }).when(mAssistContentRequester).requestAssistContent(eq(BACKLINKS_TASK_ID), any());
+        mockForAssistContent(EMPTY_ASSIST_CONTENT, BACKLINKS_TASK_ID);
 
         mViewModel.triggerBacklinks(Collections.emptySet(), DEFAULT_DISPLAY);
         waitForIdleSync();
@@ -298,15 +280,12 @@
     @Test
     public void triggerBacklinks_withNonResolvableMainLauncherIntent_noBacklinksAvailable() {
         reset(mPackageManager);
-        doAnswer(invocation -> {
-            AssistContentRequester.Callback callback = invocation.getArgument(1);
-            callback.onAssistContentAvailable(EMPTY_ASSIST_CONTENT);
-            return null;
-        }).when(mAssistContentRequester).requestAssistContent(eq(BACKLINKS_TASK_ID), any());
+        mockForAssistContent(EMPTY_ASSIST_CONTENT, BACKLINKS_TASK_ID);
 
         mViewModel.triggerBacklinks(Collections.emptySet(), DEFAULT_DISPLAY);
         waitForIdleSync();
 
+        assertThat(mViewModel.mSelectedBacklinksLiveData.getValue()).isNull();
         assertThat(mViewModel.getBacklinksLiveData().getValue()).isNull();
     }
 
@@ -314,14 +293,15 @@
     public void triggerBacklinks_nonStandardActivityIgnored_noBacklinkAvailable()
             throws RemoteException {
         reset(mAtmService);
-        RootTaskInfo taskInfo = createTaskInfoForBacklinksTask();
+        RunningTaskInfo taskInfo = createTaskInfoForBacklinksTask();
         taskInfo.configuration.windowConfiguration.setActivityType(ACTIVITY_TYPE_HOME);
-        when(mAtmService.getAllRootTaskInfosOnDisplay(DEFAULT_DISPLAY))
+        when(mAtmService.getTasks(Integer.MAX_VALUE, false, false, DEFAULT_DISPLAY))
                 .thenReturn(List.of(taskInfo));
 
         mViewModel.triggerBacklinks(Collections.emptySet(), DEFAULT_DISPLAY);
         waitForIdleSync();
 
+        assertThat(mViewModel.mSelectedBacklinksLiveData.getValue()).isNull();
         assertThat(mViewModel.getBacklinksLiveData().getValue()).isNull();
     }
 
@@ -330,9 +310,68 @@
         mViewModel.triggerBacklinks(Set.of(BACKLINKS_TASK_ID), DEFAULT_DISPLAY);
         waitForIdleSync();
 
+        assertThat(mViewModel.mSelectedBacklinksLiveData.getValue()).isNull();
         assertThat(mViewModel.getBacklinksLiveData().getValue()).isNull();
     }
 
+    @Test
+    public void triggerBacklinks_multipleAppsOnScreen_multipleBacklinksAvailable()
+            throws RemoteException {
+        // Set up mocking for multiple backlinks.
+        reset(mAtmService, mPackageManager);
+        RunningTaskInfo runningTaskInfo1 = createTaskInfoForBacklinksTask();
+        ResolveInfo resolveInfo1 = createBacklinksTaskResolveInfo();
+
+        int taskId2 = BACKLINKS_TASK_ID + 2;
+        String package2 = BACKLINKS_TASK_PACKAGE_NAME + 2;
+        String appName2 = BACKLINKS_TASK_APP_NAME + 2;
+
+        ResolveInfo resolveInfo2 = createBacklinksTaskResolveInfo();
+        ActivityInfo activityInfo2 = resolveInfo2.activityInfo;
+        activityInfo2.name = appName2;
+        activityInfo2.packageName = package2;
+        activityInfo2.applicationInfo.packageName = package2;
+        RunningTaskInfo runningTaskInfo2 = createTaskInfoForBacklinksTask();
+        runningTaskInfo2.taskId = taskId2;
+        runningTaskInfo2.topActivity = new ComponentName(package2, "backlinksClass");
+        runningTaskInfo2.topActivityInfo = resolveInfo2.activityInfo;
+        runningTaskInfo2.baseIntent = new Intent().setComponent(runningTaskInfo2.topActivity);
+
+        // For each task, the logic queries PM 3 times, twice for verifying if an app can be
+        // launched via launcher and once with the data provided in backlink intent.
+        when(mPackageManager.resolveActivity(any(), anyInt())).thenReturn(resolveInfo1,
+                resolveInfo1, resolveInfo1, resolveInfo2, resolveInfo2, resolveInfo2);
+        when(mPackageManager.loadItemIcon(any(), any())).thenReturn(FAKE_DRAWABLE);
+        when(mAtmService.getTasks(Integer.MAX_VALUE, false, false, DEFAULT_DISPLAY))
+                .thenReturn(List.of(runningTaskInfo1, runningTaskInfo2));
+
+        // Using app provided web uri for the first backlink.
+        Uri expectedUri = Uri.parse("https://developers.android.com");
+        AssistContent contentWithUri = new AssistContent();
+        contentWithUri.setWebUri(expectedUri);
+        mockForAssistContent(contentWithUri, BACKLINKS_TASK_ID);
+
+        // Using app provided intent for the second backlink.
+        Intent expectedIntent = new Intent().setPackage(package2);
+        AssistContent contentWithAppProvidedIntent = new AssistContent();
+        contentWithAppProvidedIntent.setIntent(expectedIntent);
+        mockForAssistContent(contentWithAppProvidedIntent, taskId2);
+
+        // Set up complete, trigger the backlinks action.
+        mViewModel.triggerBacklinks(Collections.emptySet(), DEFAULT_DISPLAY);
+        waitForIdleSync();
+
+        // Verify two backlinks are received and the first backlink is set as default selected.
+        assertThat(mViewModel.mSelectedBacklinksLiveData.getValue().getClipData().getItemAt(
+                0).getUri()).isEqualTo(expectedUri);
+        List<InternalBacklinksData> actualBacklinks = mViewModel.getBacklinksLiveData().getValue();
+        assertThat(actualBacklinks).hasSize(2);
+        assertThat(actualBacklinks.get(0).getClipData().getItemAt(0).getUri())
+                .isEqualTo(expectedUri);
+        assertThat(actualBacklinks.get(1).getClipData().getItemAt(0).getIntent())
+                .isEqualTo(expectedIntent);
+    }
+
     private void resetPackageManagerMockingForUsingFallbackBacklinks() {
         ResolveInfo backlinksTaskResolveInfo = createBacklinksTaskResolveInfo();
         reset(mPackageManager);
@@ -350,7 +389,7 @@
     }
 
     private void verifyMainLauncherBacklinksIntent() {
-        InternalBacklinksData result = mViewModel.getBacklinksLiveData().getValue();
+        InternalBacklinksData result = mViewModel.mSelectedBacklinksLiveData.getValue();
         assertThat(result.getAppIcon()).isEqualTo(FAKE_DRAWABLE);
 
         ClipData clipData = result.getClipData();
@@ -368,6 +407,14 @@
                 new ComponentName(BACKLINKS_TASK_PACKAGE_NAME, BACKLINKS_TASK_APP_NAME));
     }
 
+    private void mockForAssistContent(AssistContent expected, int taskId) {
+        doAnswer(invocation -> {
+            AssistContentRequester.Callback callback = invocation.getArgument(1);
+            callback.onAssistContentAvailable(expected);
+            return null;
+        }).when(mAssistContentRequester).requestAssistContent(eq(taskId), any());
+    }
+
     private static ResolveInfo createBacklinksTaskResolveInfo() {
         ActivityInfo activityInfo = new ActivityInfo();
         activityInfo.applicationInfo = new ApplicationInfo();
@@ -379,8 +426,8 @@
         return resolveInfo;
     }
 
-    private static RootTaskInfo createTaskInfoForBacklinksTask() {
-        RootTaskInfo taskInfo = new RootTaskInfo();
+    private static RunningTaskInfo createTaskInfoForBacklinksTask() {
+        RunningTaskInfo taskInfo = new RunningTaskInfo();
         taskInfo.taskId = BACKLINKS_TASK_ID;
         taskInfo.isVisible = true;
         taskInfo.isRunning = true;
@@ -388,7 +435,6 @@
         taskInfo.topActivity = new ComponentName(BACKLINKS_TASK_PACKAGE_NAME, "backlinksClass");
         taskInfo.topActivityInfo = createBacklinksTaskResolveInfo().activityInfo;
         taskInfo.baseIntent = new Intent().setComponent(taskInfo.topActivity);
-        taskInfo.childTaskIds = new int[]{BACKLINKS_TASK_ID + 1};
         taskInfo.configuration.windowConfiguration.setActivityType(ACTIVITY_TYPE_STANDARD);
         return taskInfo;
     }