diff --git a/services/core/java/com/android/server/wm/DisplayRotationCompatPolicy.java b/services/core/java/com/android/server/wm/DisplayRotationCompatPolicy.java
index 63fe94c..e50a089 100644
--- a/services/core/java/com/android/server/wm/DisplayRotationCompatPolicy.java
+++ b/services/core/java/com/android/server/wm/DisplayRotationCompatPolicy.java
@@ -44,6 +44,7 @@
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.protolog.ProtoLog;
 import com.android.server.UiThread;
+import com.android.window.flags.Flags;
 
 /**
  * Controls camera compatibility treatment that handles orientation mismatch between camera
@@ -69,6 +70,9 @@
     @NonNull
     private final ActivityRefresher mActivityRefresher;
 
+    @Nullable
+    private Task mCameraTask;
+
     @ScreenOrientation
     private int mLastReportedOrientation = SCREEN_ORIENTATION_UNSET;
 
@@ -104,7 +108,7 @@
      * guaranteed to match, the rotation can cause letterboxing.
      *
      * <p>If treatment isn't applicable returns {@link SCREEN_ORIENTATION_UNSPECIFIED}. See {@link
-     * #shouldComputeCameraCompatOrientation} for conditions enabling the treatment.
+     * #isTreatmentEnabledForDisplay} for conditions enabling the treatment.
      */
     @ScreenOrientation
     int getOrientation() {
@@ -136,9 +140,9 @@
         // are aligned when they compute orientation of the preview.
         // This means that even for a landscape-only activity and a device with landscape natural
         // orientation this would return SCREEN_ORIENTATION_PORTRAIT because an assumption that
-        // natural orientation = portrait window = portait camera is the main wrong assumption
+        // natural orientation = portrait window = portrait camera is the main wrong assumption
         // that apps make when they implement camera previews so landscape windows need be
-        // rotated in the orientation oposite to the natural one even if it's portrait.
+        // rotated in the orientation opposite to the natural one even if it's portrait.
         // TODO(b/261475895): Consider allowing more rotations for "sensor" and "user" versions
         // of the portrait and landscape orientation requests.
         final int orientation = (isPortraitActivity && isNaturalDisplayOrientationPortrait)
@@ -296,6 +300,7 @@
     @Override
     public boolean onCameraOpened(@NonNull ActivityRecord cameraActivity,
             @NonNull String cameraId) {
+        mCameraTask = cameraActivity.getTask();
         // Checking whether an activity in fullscreen rather than the task as this camera
         // compat treatment doesn't cover activity embedding.
         if (cameraActivity.getWindowingMode() == WINDOWING_MODE_FULLSCREEN) {
@@ -305,7 +310,7 @@
         }
         // Checking that the whole app is in multi-window mode as we shouldn't show toast
         // for the activity embedding case.
-        if (cameraActivity.getTask().getWindowingMode() == WINDOWING_MODE_MULTI_WINDOW
+        if (mCameraTask != null && mCameraTask.getWindowingMode() == WINDOWING_MODE_MULTI_WINDOW
                 && isTreatmentEnabledForActivity(
                 cameraActivity, /* mustBeFullscreen */ false)) {
             final PackageManager packageManager = mWmService.mContext.getPackageManager();
@@ -343,10 +348,15 @@
 
     @Override
     public boolean onCameraClosed(@NonNull String cameraId) {
-        // Top activity in the same task as the camera activity, or `null` if the task is
-        // closed.
-        final ActivityRecord topActivity = mDisplayContent.topRunningActivity(
-                /* considerKeyguardState= */ true);
+        final ActivityRecord topActivity;
+        if (Flags.cameraCompatFullscreenPickSameTaskActivity()) {
+            topActivity = mCameraTask != null ? mCameraTask.getTopActivity(
+                    /* includeFinishing= */ true, /* includeOverlays= */ false) : null;
+        } else {
+            topActivity = mDisplayContent.topRunningActivity(/* considerKeyguardState= */ true);
+        }
+
+        mCameraTask = null;
         if (topActivity == null) {
             return true;
         }
@@ -368,8 +378,6 @@
                 mDisplayContent.mDisplayId);
         // Checking whether an activity in fullscreen rather than the task as this camera compat
         // treatment doesn't cover activity embedding.
-        // TODO(b/350495350): Consider checking whether this activity is the camera activity, or
-        // whether the top activity has the same task as the one which opened camera.
         if (topActivity.getWindowingMode() != WINDOWING_MODE_FULLSCREEN) {
             return true;
         }