diff --git a/services/core/java/com/android/server/wm/ActivityRecord.java b/services/core/java/com/android/server/wm/ActivityRecord.java
index 91c849c4110f6..85f333a19aa07 100644
--- a/services/core/java/com/android/server/wm/ActivityRecord.java
+++ b/services/core/java/com/android/server/wm/ActivityRecord.java
@@ -3589,7 +3589,7 @@ final class ActivityRecord extends WindowToken implements WindowManagerService.A
 
     @Override
     ActivityRecord getActivity(Predicate<ActivityRecord> callback, boolean traverseTopToBottom,
-            WindowContainer boundary) {
+            ActivityRecord boundary) {
         return callback.test(this) ? this : null;
     }
 
diff --git a/services/core/java/com/android/server/wm/ActivityStack.java b/services/core/java/com/android/server/wm/ActivityStack.java
index 70488b4b29d91..7dcdbdfc52124 100644
--- a/services/core/java/com/android/server/wm/ActivityStack.java
+++ b/services/core/java/com/android/server/wm/ActivityStack.java
@@ -2689,7 +2689,9 @@ class ActivityStack extends Task {
      */
     @Nullable
     private ActivityRecord getOccludingActivityAbove(ActivityRecord activity) {
-        return getActivity((ar) -> ar.occludesParent(), true /* traverseTopToBottom */, activity);
+        ActivityRecord top = getActivity((ar) -> ar.occludesParent(),
+                true /* traverseTopToBottom */, activity);
+        return top != activity ? top : null;
     }
 
     boolean willActivityBeVisible(IBinder token) {
diff --git a/services/core/java/com/android/server/wm/WindowContainer.java b/services/core/java/com/android/server/wm/WindowContainer.java
index 5d7ec127da464..aa90b76ebb120 100644
--- a/services/core/java/com/android/server/wm/WindowContainer.java
+++ b/services/core/java/com/android/server/wm/WindowContainer.java
@@ -1415,11 +1415,12 @@ class WindowContainer<E extends WindowContainer> extends ConfigurationContainer<
     }
 
     ActivityRecord getActivity(Predicate<ActivityRecord> callback, boolean traverseTopToBottom,
-            WindowContainer boundary) {
+            ActivityRecord boundary) {
         if (traverseTopToBottom) {
             for (int i = mChildren.size() - 1; i >= 0; --i) {
                 final WindowContainer wc = mChildren.get(i);
-                if (wc == boundary) return null;
+                // TODO(b/156986561): Improve the correctness of the boundary check.
+                if (wc == boundary) return boundary;
 
                 final ActivityRecord r = wc.getActivity(callback, traverseTopToBottom, boundary);
                 if (r != null) {
@@ -1430,7 +1431,8 @@ class WindowContainer<E extends WindowContainer> extends ConfigurationContainer<
             final int count = mChildren.size();
             for (int i = 0; i < count; i++) {
                 final WindowContainer wc = mChildren.get(i);
-                if (wc == boundary) return null;
+                // TODO(b/156986561): Improve the correctness of the boundary check.
+                if (wc == boundary) return boundary;
 
                 final ActivityRecord r = wc.getActivity(callback, traverseTopToBottom, boundary);
                 if (r != null) {