diff --git a/services/core/java/com/android/server/wm/RootWindowContainer.java b/services/core/java/com/android/server/wm/RootWindowContainer.java
index ddad1db..eaf7693 100644
--- a/services/core/java/com/android/server/wm/RootWindowContainer.java
+++ b/services/core/java/com/android/server/wm/RootWindowContainer.java
@@ -3372,7 +3372,7 @@
     }
 
     /**
-     * Find all task stacks containing {@param userId} and intercept them with an activity
+     * Find all visible task stacks containing {@param userId} and intercept them with an activity
      * to block out the contents and possibly start a credential-confirming intent.
      *
      * @param userId user handle for the locked managed profile.
@@ -3380,18 +3380,42 @@
     void lockAllProfileTasks(@UserIdInt int userId) {
         mService.deferWindowLayout();
         try {
-            forAllLeafTasks(task -> {
-                if (task.getActivity(activity -> !activity.finishing && activity.mUserId == userId)
-                        != null) {
-                    mService.getTaskChangeNotificationController().notifyTaskProfileLocked(
-                            task.mTaskId, userId);
-                }
-            }, true /* traverseTopToBottom */);
+            final PooledConsumer c = PooledLambda.obtainConsumer(
+                    RootWindowContainer::taskTopActivityIsUser, this, PooledLambda.__(Task.class),
+                    userId);
+            forAllLeafTasks(c, true /* traverseTopToBottom */);
+            c.recycle();
         } finally {
             mService.continueWindowLayout();
         }
     }
 
+    /**
+     * Detects whether we should show a lock screen in front of this task for a locked user.
+     * <p>
+     * We'll do this if either of the following holds:
+     * <ul>
+     *   <li>The top activity explicitly belongs to {@param userId}.</li>
+     *   <li>The top activity returns a result to an activity belonging to {@param userId}.</li>
+     * </ul>
+     *
+     * @return {@code true} if the top activity looks like it belongs to {@param userId}.
+     */
+    private void taskTopActivityIsUser(Task task, @UserIdInt int userId) {
+        // To handle the case that work app is in the task but just is not the top one.
+        final ActivityRecord activityRecord = task.getTopNonFinishingActivity();
+        final ActivityRecord resultTo = (activityRecord != null ? activityRecord.resultTo : null);
+
+        // Check the task for a top activity belonging to userId, or returning a
+        // result to an activity belonging to userId. Example case: a document
+        // picker for personal files, opened by a work app, should still get locked.
+        if ((activityRecord != null && activityRecord.mUserId == userId)
+                || (resultTo != null && resultTo.mUserId == userId)) {
+            mService.getTaskChangeNotificationController().notifyTaskProfileLocked(
+                    task.mTaskId, userId);
+        }
+    }
+
     void cancelInitializingActivities() {
         for (int displayNdx = getChildCount() - 1; displayNdx >= 0; --displayNdx) {
             final DisplayContent display = getChildAt(displayNdx);
diff --git a/services/tests/wmtests/src/com/android/server/wm/RootWindowContainerTests.java b/services/tests/wmtests/src/com/android/server/wm/RootWindowContainerTests.java
index 1aff8a7..35d1b17 100644
--- a/services/tests/wmtests/src/com/android/server/wm/RootWindowContainerTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/RootWindowContainerTests.java
@@ -25,7 +25,6 @@
 import static android.view.WindowManager.LayoutParams.TYPE_STATUS_BAR;
 import static android.view.WindowManager.LayoutParams.TYPE_TOAST;
 
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.spyOn;
 import static com.android.server.wm.ActivityStack.ActivityState.FINISHING;
 import static com.android.server.wm.ActivityStack.ActivityState.PAUSED;
 import static com.android.server.wm.ActivityStack.ActivityState.PAUSING;
@@ -37,13 +36,10 @@
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
-import static org.mockito.ArgumentMatchers.eq;
-import static org.mockito.Mockito.verify;
 
 import android.app.WindowConfiguration;
 import android.content.ComponentName;
 import android.content.pm.ActivityInfo;
-import android.os.UserHandle;
 import android.platform.test.annotations.Presubmit;
 
 import androidx.test.filters.SmallTest;
@@ -173,34 +169,5 @@
         activity.setState(FINISHING, "test FINISHING");
         assertThat(mWm.mRoot.allPausedActivitiesComplete()).isTrue();
     }
-
-    @Test
-    public void testLockAllProfileTasks() {
-        // Make an activity visible with the user id set to 0
-        DisplayContent displayContent = mWm.mRoot.getDisplayContent(DEFAULT_DISPLAY);
-        TaskDisplayArea taskDisplayArea = displayContent.getTaskDisplayAreaAt(0);
-        final ActivityStack stack = createTaskStackOnDisplay(WINDOWING_MODE_FULLSCREEN,
-                ACTIVITY_TYPE_STANDARD, displayContent);
-        final ActivityRecord activity = new ActivityTestsBase.ActivityBuilder(stack.mAtmService)
-                .setStack(stack)
-                .setUid(0)
-                .setCreateTask(true)
-                .build();
-
-        // Create another activity on top and the user id is 1
-        Task task = activity.getTask();
-        final ActivityRecord topActivity = new ActivityTestsBase.ActivityBuilder(mWm.mAtmService)
-                .setStack(stack)
-                .setUid(UserHandle.PER_USER_RANGE + 1)
-                .setTask(task)
-                .build();
-
-        // Make sure the listeners will be notified for putting the task to locked state
-        TaskChangeNotificationController controller =
-                mWm.mAtmService.getTaskChangeNotificationController();
-        spyOn(controller);
-        mWm.mRoot.lockAllProfileTasks(0);
-        verify(controller).notifyTaskProfileLocked(eq(task.mTaskId), eq(0));
-    }
 }