Avoid duplicated WindowContainerTransaction for TaskFragment

Before, when there is Task bounds change without triggering TaskFragment
bounds change, we may still trigger a WCT to update adjacentTF that
isn't different from the previous request.

Record the last request value to avoid duplicated request.

Also make sure the requested changes are reset when TaskFragment enters
PiP.

Bug: 265271880
Test: atest WmTests:TaskFragmentTest
Test: atest WMJetpackUnitTests:SplitPresenterTest
Change-Id: Ifd9bba2de3e5f2caaa3efd1a6af66f3bcebff7d1

9 files changed, 299 insertions, 59 deletions

diff --git a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/JetpackTaskFragmentOrganizer.java b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/JetpackTaskFragmentOrganizer.java
index 67963a36713c..d94e8e426c4b 100644
--- a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/JetpackTaskFragmentOrganizer.java
+++ b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/JetpackTaskFragmentOrganizer.java
@@ -154,7 +154,7 @@ class JetpackTaskFragmentOrganizer extends TaskFragmentOrganizer {
                 activityOptions);
 
         // Set adjacent to each other so that the containers below will be invisible.
-        setAdjacentTaskFragments(wct, launchingFragmentToken, secondaryFragmentToken, rule);
+        setAdjacentTaskFragmentsWithRule(wct, launchingFragmentToken, secondaryFragmentToken, rule);
         setCompanionTaskFragment(wct, launchingFragmentToken, secondaryFragmentToken, rule,
                 false /* isStacked */);
     }
@@ -167,7 +167,7 @@ class JetpackTaskFragmentOrganizer extends TaskFragmentOrganizer {
     void expandTaskFragment(@NonNull WindowContainerTransaction wct,
             @NonNull IBinder fragmentToken) {
         resizeTaskFragment(wct, fragmentToken, new Rect());
-        setAdjacentTaskFragments(wct, fragmentToken, null /* secondary */, null /* splitRule */);
+        clearAdjacentTaskFragments(wct, fragmentToken);
         updateWindowingMode(wct, fragmentToken, WINDOWING_MODE_UNDEFINED);
         updateAnimationParams(wct, fragmentToken, TaskFragmentAnimationParams.DEFAULT);
     }
@@ -238,26 +238,37 @@ class JetpackTaskFragmentOrganizer extends TaskFragmentOrganizer {
         wct.reparentActivityToTaskFragment(fragmentToken, reparentActivityToken);
     }
 
-    void setAdjacentTaskFragments(@NonNull WindowContainerTransaction wct,
-            @NonNull IBinder primary, @Nullable IBinder secondary, @Nullable SplitRule splitRule) {
-        if (secondary == null) {
-            wct.clearAdjacentTaskFragments(primary);
-            return;
-        }
-
+    /**
+     * Sets the two given TaskFragments as adjacent to each other with respecting the given
+     * {@link SplitRule} for {@link WindowContainerTransaction.TaskFragmentAdjacentParams}.
+     */
+    void setAdjacentTaskFragmentsWithRule(@NonNull WindowContainerTransaction wct,
+            @NonNull IBinder primary, @NonNull IBinder secondary, @NonNull SplitRule splitRule) {
         WindowContainerTransaction.TaskFragmentAdjacentParams adjacentParams = null;
         final boolean finishSecondaryWithPrimary =
-                splitRule != null && SplitContainer.shouldFinishSecondaryWithPrimary(splitRule);
+                SplitContainer.shouldFinishSecondaryWithPrimary(splitRule);
         final boolean finishPrimaryWithSecondary =
-                splitRule != null && SplitContainer.shouldFinishPrimaryWithSecondary(splitRule);
+                SplitContainer.shouldFinishPrimaryWithSecondary(splitRule);
         if (finishSecondaryWithPrimary || finishPrimaryWithSecondary) {
             adjacentParams = new WindowContainerTransaction.TaskFragmentAdjacentParams();
             adjacentParams.setShouldDelayPrimaryLastActivityRemoval(finishSecondaryWithPrimary);
             adjacentParams.setShouldDelaySecondaryLastActivityRemoval(finishPrimaryWithSecondary);
         }
+        setAdjacentTaskFragments(wct, primary, secondary, adjacentParams);
+    }
+
+    void setAdjacentTaskFragments(@NonNull WindowContainerTransaction wct,
+            @NonNull IBinder primary, @NonNull IBinder secondary,
+            @Nullable WindowContainerTransaction.TaskFragmentAdjacentParams adjacentParams) {
         wct.setAdjacentTaskFragments(primary, secondary, adjacentParams);
     }
 
+    void clearAdjacentTaskFragments(@NonNull WindowContainerTransaction wct,
+            @NonNull IBinder fragmentToken) {
+        // Clear primary will also clear secondary.
+        wct.clearAdjacentTaskFragments(fragmentToken);
+    }
+
     void setCompanionTaskFragment(@NonNull WindowContainerTransaction wct,
             @NonNull IBinder primary, @NonNull IBinder secondary, @NonNull SplitRule splitRule,
             boolean isStacked) {
@@ -268,7 +279,7 @@ class JetpackTaskFragmentOrganizer extends TaskFragmentOrganizer {
         } else {
             finishPrimaryWithSecondary = shouldFinishPrimaryWithSecondary(splitRule);
         }
-        wct.setCompanionTaskFragment(primary, finishPrimaryWithSecondary ? secondary : null);
+        setCompanionTaskFragment(wct, primary, finishPrimaryWithSecondary ? secondary : null);
 
         final boolean finishSecondaryWithPrimary;
         if (isStacked) {
@@ -277,7 +288,12 @@ class JetpackTaskFragmentOrganizer extends TaskFragmentOrganizer {
         } else {
             finishSecondaryWithPrimary = shouldFinishSecondaryWithPrimary(splitRule);
         }
-        wct.setCompanionTaskFragment(secondary, finishSecondaryWithPrimary ? primary : null);
+        setCompanionTaskFragment(wct, secondary, finishSecondaryWithPrimary ? primary : null);
+    }
+
+    void setCompanionTaskFragment(@NonNull WindowContainerTransaction wct, @NonNull IBinder primary,
+            @Nullable IBinder secondary) {
+        wct.setCompanionTaskFragment(primary, secondary);
     }
 
     void resizeTaskFragment(@NonNull WindowContainerTransaction wct, @NonNull IBinder fragmentToken,
diff --git a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitController.java b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitController.java
index 57ba6bbb04f0..2d7fd7e6a085 100644
--- a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitController.java
+++ b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitController.java
@@ -424,6 +424,9 @@ public class SplitController implements JetpackTaskFragmentOrganizer.TaskFragmen
             // All overrides will be cleanup.
             container.setLastRequestedBounds(null /* bounds */);
             container.setLastRequestedWindowingMode(WINDOWING_MODE_UNDEFINED);
+            container.clearLastAdjacentTaskFragment();
+            container.setLastCompanionTaskFragment(null /* fragmentToken */);
+            container.setLastRequestAnimationParams(TaskFragmentAnimationParams.DEFAULT);
             cleanupForEnterPip(wct, container);
         } else if (wasInPip) {
             // Exit PIP.
diff --git a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitPresenter.java b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitPresenter.java
index 2b93682d791f..8a21039ffe4e 100644
--- a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitPresenter.java
+++ b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitPresenter.java
@@ -387,10 +387,9 @@ class SplitPresenter extends JetpackTaskFragmentOrganizer {
         // secondaryContainer could not be finished.
         boolean isStacked = !shouldShowSplit(splitAttributes);
         if (isStacked) {
-            setAdjacentTaskFragments(wct, primaryContainer.getTaskFragmentToken(),
-                    null /* secondary */, null /* splitRule */);
+            clearAdjacentTaskFragments(wct, primaryContainer.getTaskFragmentToken());
         } else {
-            setAdjacentTaskFragments(wct, primaryContainer.getTaskFragmentToken(),
+            setAdjacentTaskFragmentsWithRule(wct, primaryContainer.getTaskFragmentToken(),
                     secondaryContainer.getTaskFragmentToken(), splitRule);
         }
         setCompanionTaskFragment(wct, primaryContainer.getTaskFragmentToken(),
@@ -427,7 +426,7 @@ class SplitPresenter extends JetpackTaskFragmentOrganizer {
                 fragmentOptions.getFragmentToken());
         if (container == null) {
             throw new IllegalStateException(
-                    "Creating a task fragment that is not registered with controller.");
+                    "Creating a TaskFragment that is not registered with controller.");
         }
 
         container.setLastRequestedBounds(fragmentOptions.getInitialRelativeBounds());
@@ -441,7 +440,7 @@ class SplitPresenter extends JetpackTaskFragmentOrganizer {
         TaskFragmentContainer container = mController.getContainer(fragmentToken);
         if (container == null) {
             throw new IllegalStateException(
-                    "Resizing a task fragment that is not registered with controller.");
+                    "Resizing a TaskFragment that is not registered with controller.");
         }
 
         if (container.areLastRequestedBoundsEqual(relBounds)) {
@@ -458,7 +457,7 @@ class SplitPresenter extends JetpackTaskFragmentOrganizer {
             @NonNull IBinder fragmentToken, @WindowingMode int windowingMode) {
         final TaskFragmentContainer container = mController.getContainer(fragmentToken);
         if (container == null) {
-            throw new IllegalStateException("Setting windowing mode for a task fragment that is"
+            throw new IllegalStateException("Setting windowing mode for a TaskFragment that is"
                     + " not registered with controller.");
         }
 
@@ -476,7 +475,7 @@ class SplitPresenter extends JetpackTaskFragmentOrganizer {
             @NonNull IBinder fragmentToken, @NonNull TaskFragmentAnimationParams animationParams) {
         final TaskFragmentContainer container = mController.getContainer(fragmentToken);
         if (container == null) {
-            throw new IllegalStateException("Setting animation params for a task fragment that is"
+            throw new IllegalStateException("Setting animation params for a TaskFragment that is"
                     + " not registered with controller.");
         }
 
@@ -489,6 +488,64 @@ class SplitPresenter extends JetpackTaskFragmentOrganizer {
         super.updateAnimationParams(wct, fragmentToken, animationParams);
     }
 
+    @Override
+    void setAdjacentTaskFragments(@NonNull WindowContainerTransaction wct,
+            @NonNull IBinder primary, @NonNull IBinder secondary,
+            @Nullable WindowContainerTransaction.TaskFragmentAdjacentParams adjacentParams) {
+        final TaskFragmentContainer primaryContainer = mController.getContainer(primary);
+        final TaskFragmentContainer secondaryContainer = mController.getContainer(secondary);
+        if (primaryContainer == null || secondaryContainer == null) {
+            throw new IllegalStateException("setAdjacentTaskFragments on TaskFragment that is"
+                    + " not registered with controller.");
+        }
+
+        if (primaryContainer.isLastAdjacentTaskFragmentEqual(secondary, adjacentParams)
+                && secondaryContainer.isLastAdjacentTaskFragmentEqual(primary, adjacentParams)) {
+            // Return early if the same adjacent TaskFragments were already requested
+            return;
+        }
+
+        primaryContainer.setLastAdjacentTaskFragment(secondary, adjacentParams);
+        secondaryContainer.setLastAdjacentTaskFragment(primary, adjacentParams);
+        super.setAdjacentTaskFragments(wct, primary, secondary, adjacentParams);
+    }
+
+    @Override
+    void clearAdjacentTaskFragments(@NonNull WindowContainerTransaction wct,
+            @NonNull IBinder fragmentToken) {
+        final TaskFragmentContainer container = mController.getContainer(fragmentToken);
+        if (container == null) {
+            throw new IllegalStateException("clearAdjacentTaskFragments on TaskFragment that is"
+                    + " not registered with controller.");
+        }
+
+        if (container.isLastAdjacentTaskFragmentEqual(null /* fragmentToken*/, null /* params */)) {
+            // Return early if no adjacent TaskFragment was yet requested
+            return;
+        }
+
+        container.clearLastAdjacentTaskFragment();
+        super.clearAdjacentTaskFragments(wct, fragmentToken);
+    }
+
+    @Override
+    void setCompanionTaskFragment(@NonNull WindowContainerTransaction wct, @NonNull IBinder primary,
+            @Nullable IBinder secondary) {
+        final TaskFragmentContainer container = mController.getContainer(primary);
+        if (container == null) {
+            throw new IllegalStateException("setCompanionTaskFragment on TaskFragment that is"
+                    + " not registered with controller.");
+        }
+
+        if (container.isLastCompanionTaskFragmentEqual(secondary)) {
+            // Return early if the same companion TaskFragment was already requested
+            return;
+        }
+
+        container.setLastCompanionTaskFragment(secondary);
+        super.setCompanionTaskFragment(wct, primary, secondary);
+    }
+
     /**
      * Expands the split container if the current split bounds are smaller than the Activity or
      * Intent that is added to the container.
diff --git a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/TaskFragmentContainer.java b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/TaskFragmentContainer.java
index 38ac719ebbe2..bc082b134acf 100644
--- a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/TaskFragmentContainer.java
+++ b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/TaskFragmentContainer.java
@@ -39,6 +39,7 @@ import com.android.internal.annotations.VisibleForTesting;
 import java.util.ArrayList;
 import java.util.Iterator;
 import java.util.List;
+import java.util.Objects;
 
 /**
  * Client-side container for a stack of activities. Corresponds to an instance of TaskFragment
@@ -116,6 +117,27 @@ class TaskFragmentContainer {
     private TaskFragmentAnimationParams mLastAnimationParams = TaskFragmentAnimationParams.DEFAULT;
 
     /**
+     * TaskFragment token that was requested last via
+     * {@link android.window.TaskFragmentOperation#OP_TYPE_SET_ADJACENT_TASK_FRAGMENTS}.
+     */
+    @Nullable
+    private IBinder mLastAdjacentTaskFragment;
+
+    /**
+     * {@link WindowContainerTransaction.TaskFragmentAdjacentParams} token that was requested last
+     * via {@link android.window.TaskFragmentOperation#OP_TYPE_SET_ADJACENT_TASK_FRAGMENTS}.
+     */
+    @Nullable
+    private WindowContainerTransaction.TaskFragmentAdjacentParams mLastAdjacentParams;
+
+    /**
+     * TaskFragment token that was requested last via
+     * {@link android.window.TaskFragmentOperation#OP_TYPE_SET_COMPANION_TASK_FRAGMENT}.
+     */
+    @Nullable
+    private IBinder mLastCompanionTaskFragment;
+
+    /**
      * When the TaskFragment has appeared in server, but is empty, we should remove the TaskFragment
      * if it is still empty after the timeout.
      */
@@ -571,6 +593,7 @@ class TaskFragmentContainer {
     /**
      * Checks if last requested bounds are equal to the provided value.
      * The requested bounds are relative bounds in parent coordinate.
+     * @see WindowContainerTransaction#setRelativeBounds
      */
     boolean areLastRequestedBoundsEqual(@Nullable Rect relBounds) {
         return (relBounds == null && mLastRequestedBounds.isEmpty())
@@ -580,6 +603,7 @@ class TaskFragmentContainer {
     /**
      * Updates the last requested bounds.
      * The requested bounds are relative bounds in parent coordinate.
+     * @see WindowContainerTransaction#setRelativeBounds
      */
     void setLastRequestedBounds(@Nullable Rect relBounds) {
         if (relBounds == null) {
@@ -589,13 +613,9 @@ class TaskFragmentContainer {
         }
     }
 
-    @NonNull
-    Rect getLastRequestedBounds() {
-        return mLastRequestedBounds;
-    }
-
     /**
      * Checks if last requested windowing mode is equal to the provided value.
+     * @see WindowContainerTransaction#setWindowingMode
      */
     boolean isLastRequestedWindowingModeEqual(@WindowingMode int windowingMode) {
         return mLastRequestedWindowingMode == windowingMode;
@@ -603,6 +623,7 @@ class TaskFragmentContainer {
 
     /**
      * Updates the last requested windowing mode.
+     * @see WindowContainerTransaction#setWindowingMode
      */
     void setLastRequestedWindowingMode(@WindowingMode int windowingModes) {
         mLastRequestedWindowingMode = windowingModes;
@@ -610,6 +631,7 @@ class TaskFragmentContainer {
 
     /**
      * Checks if last requested {@link TaskFragmentAnimationParams} are equal to the provided value.
+     * @see android.window.TaskFragmentOperation#OP_TYPE_SET_ANIMATION_PARAMS
      */
     boolean areLastRequestedAnimationParamsEqual(
             @NonNull TaskFragmentAnimationParams animationParams) {
@@ -618,11 +640,66 @@ class TaskFragmentContainer {
 
     /**
      * Updates the last requested {@link TaskFragmentAnimationParams}.
+     * @see android.window.TaskFragmentOperation#OP_TYPE_SET_ANIMATION_PARAMS
      */
     void setLastRequestAnimationParams(@NonNull TaskFragmentAnimationParams animationParams) {
         mLastAnimationParams = animationParams;
     }
 
+    /**
+     * Checks if last requested adjacent TaskFragment token and params are equal to the provided
+     * values.
+     * @see android.window.TaskFragmentOperation#OP_TYPE_SET_ADJACENT_TASK_FRAGMENTS
+     * @see android.window.TaskFragmentOperation#OP_TYPE_CLEAR_ADJACENT_TASK_FRAGMENTS
+     */
+    boolean isLastAdjacentTaskFragmentEqual(@Nullable IBinder fragmentToken,
+            @Nullable WindowContainerTransaction.TaskFragmentAdjacentParams params) {
+        return Objects.equals(mLastAdjacentTaskFragment, fragmentToken)
+                && Objects.equals(mLastAdjacentParams, params);
+    }
+
+    /**
+     * Updates the last requested adjacent TaskFragment token and params.
+     * @see android.window.TaskFragmentOperation#OP_TYPE_SET_ADJACENT_TASK_FRAGMENTS
+     */
+    void setLastAdjacentTaskFragment(@NonNull IBinder fragmentToken,
+            @NonNull WindowContainerTransaction.TaskFragmentAdjacentParams params) {
+        mLastAdjacentTaskFragment = fragmentToken;
+        mLastAdjacentParams = params;
+    }
+
+    /**
+     * Clears the last requested adjacent TaskFragment token and params.
+     * @see android.window.TaskFragmentOperation#OP_TYPE_CLEAR_ADJACENT_TASK_FRAGMENTS
+     */
+    void clearLastAdjacentTaskFragment() {
+        final TaskFragmentContainer lastAdjacentTaskFragment = mLastAdjacentTaskFragment != null
+                ? mController.getContainer(mLastAdjacentTaskFragment)
+                : null;
+        mLastAdjacentTaskFragment = null;
+        mLastAdjacentParams = null;
+        if (lastAdjacentTaskFragment != null) {
+            // Clear the previous adjacent TaskFragment as well.
+            lastAdjacentTaskFragment.clearLastAdjacentTaskFragment();
+        }
+    }
+
+    /**
+     * Checks if last requested companion TaskFragment token is equal to the provided value.
+     * @see android.window.TaskFragmentOperation#OP_TYPE_SET_COMPANION_TASK_FRAGMENT
+     */
+    boolean isLastCompanionTaskFragmentEqual(@Nullable IBinder fragmentToken) {
+        return Objects.equals(mLastCompanionTaskFragment, fragmentToken);
+    }
+
+    /**
+     * Updates the last requested companion TaskFragment token.
+     * @see android.window.TaskFragmentOperation#OP_TYPE_SET_COMPANION_TASK_FRAGMENT
+     */
+    void setLastCompanionTaskFragment(@Nullable IBinder fragmentToken) {
+        mLastCompanionTaskFragment = fragmentToken;
+    }
+
     /** Gets the parent leaf Task id. */
     int getTaskId() {
         return mTaskContainer.getTaskId();
diff --git a/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/SplitPresenterTest.java b/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/SplitPresenterTest.java
index a41e63f46f55..96dd86aff7c0 100644
--- a/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/SplitPresenterTest.java
+++ b/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/SplitPresenterTest.java
@@ -177,6 +177,64 @@ public class SplitPresenterTest {
     }
 
     @Test
+    public void testSetAdjacentTaskFragments() {
+        final TaskFragmentContainer container0 = mController.newContainer(mActivity, TASK_ID);
+        final TaskFragmentContainer container1 = mController.newContainer(mActivity, TASK_ID);
+
+        mPresenter.setAdjacentTaskFragments(mTransaction, container0.getTaskFragmentToken(),
+                container1.getTaskFragmentToken(), null /* adjacentParams */);
+        verify(mTransaction).setAdjacentTaskFragments(container0.getTaskFragmentToken(),
+                container1.getTaskFragmentToken(), null /* adjacentParams */);
+
+        // No request to set the same adjacent TaskFragments.
+        clearInvocations(mTransaction);
+        mPresenter.setAdjacentTaskFragments(mTransaction, container0.getTaskFragmentToken(),
+                container1.getTaskFragmentToken(), null /* adjacentParams */);
+
+        verify(mTransaction, never()).setAdjacentTaskFragments(any(), any(), any());
+    }
+
+    @Test
+    public void testClearAdjacentTaskFragments() {
+        final TaskFragmentContainer container0 = mController.newContainer(mActivity, TASK_ID);
+        final TaskFragmentContainer container1 = mController.newContainer(mActivity, TASK_ID);
+
+        // No request to clear as it is not set by default.
+        mPresenter.clearAdjacentTaskFragments(mTransaction, container0.getTaskFragmentToken());
+        verify(mTransaction, never()).clearAdjacentTaskFragments(any());
+
+        mPresenter.setAdjacentTaskFragments(mTransaction, container0.getTaskFragmentToken(),
+                container1.getTaskFragmentToken(), null /* adjacentParams */);
+        mPresenter.clearAdjacentTaskFragments(mTransaction, container0.getTaskFragmentToken());
+        verify(mTransaction).clearAdjacentTaskFragments(container0.getTaskFragmentToken());
+
+        // No request to clear on either of the previous cleared TasKFragments.
+        clearInvocations(mTransaction);
+        mPresenter.clearAdjacentTaskFragments(mTransaction, container0.getTaskFragmentToken());
+        mPresenter.clearAdjacentTaskFragments(mTransaction, container1.getTaskFragmentToken());
+
+        verify(mTransaction, never()).clearAdjacentTaskFragments(any());
+    }
+
+    @Test
+    public void testSetCompanionTaskFragment() {
+        final TaskFragmentContainer container0 = mController.newContainer(mActivity, TASK_ID);
+        final TaskFragmentContainer container1 = mController.newContainer(mActivity, TASK_ID);
+
+        mPresenter.setCompanionTaskFragment(mTransaction, container0.getTaskFragmentToken(),
+                container1.getTaskFragmentToken());
+        verify(mTransaction).setCompanionTaskFragment(container0.getTaskFragmentToken(),
+                container1.getTaskFragmentToken());
+
+        // No request to set the same adjacent TaskFragments.
+        clearInvocations(mTransaction);
+        mPresenter.setCompanionTaskFragment(mTransaction, container0.getTaskFragmentToken(),
+                container1.getTaskFragmentToken());
+
+        verify(mTransaction, never()).setCompanionTaskFragment(any(), any());
+    }
+
+    @Test
     public void testUpdateAnimationParams() {
         final TaskFragmentContainer container = mController.newContainer(mActivity, TASK_ID);
 
diff --git a/services/core/java/com/android/server/wm/RootWindowContainer.java b/services/core/java/com/android/server/wm/RootWindowContainer.java
index e8aa2c8d4e60..9333d65c7148 100644
--- a/services/core/java/com/android/server/wm/RootWindowContainer.java
+++ b/services/core/java/com/android/server/wm/RootWindowContainer.java
@@ -136,6 +136,7 @@ import android.view.Display;
 import android.view.DisplayInfo;
 import android.view.SurfaceControl;
 import android.view.WindowManager;
+import android.window.TaskFragmentAnimationParams;
 import android.window.WindowContainerToken;
 
 import com.android.internal.annotations.VisibleForTesting;
@@ -2082,6 +2083,8 @@ class RootWindowContainer extends WindowContainer<DisplayContent>
                     return;
                 }
                 tf.resetAdjacentTaskFragment();
+                tf.setCompanionTaskFragment(null /* companionTaskFragment */);
+                tf.setAnimationParams(TaskFragmentAnimationParams.DEFAULT);
                 if (tf.getTopNonFinishingActivity() != null) {
                     // When the Task is entering picture-in-picture, we should clear all override
                     // from the client organizer, so the PIP activity can get the correct config
diff --git a/services/core/java/com/android/server/wm/TaskFragment.java b/services/core/java/com/android/server/wm/TaskFragment.java
index eb06b911e2d8..294e90b858a4 100644
--- a/services/core/java/com/android/server/wm/TaskFragment.java
+++ b/services/core/java/com/android/server/wm/TaskFragment.java
@@ -2813,7 +2813,7 @@ class TaskFragment extends WindowContainer<WindowContainer> {
     void removeImmediately() {
         mIsRemovalRequested = false;
         resetAdjacentTaskFragment();
-        cleanUp();
+        cleanUpEmbeddedTaskFragment();
         final boolean shouldExecuteAppTransition =
                 mClearedTaskFragmentForPip && isTaskVisibleRequested();
         super.removeImmediately();
@@ -2830,10 +2830,20 @@ class TaskFragment extends WindowContainer<WindowContainer> {
     }
 
     /** Called on remove to cleanup. */
-    private void cleanUp() {
-        if (mIsEmbedded) {
-            mAtmService.mWindowOrganizerController.cleanUpEmbeddedTaskFragment(this);
+    private void cleanUpEmbeddedTaskFragment() {
+        if (!mIsEmbedded) {
+            return;
         }
+        mAtmService.mWindowOrganizerController.cleanUpEmbeddedTaskFragment(this);
+        final Task task = getTask();
+        if (task == null) {
+            return;
+        }
+        task.forAllLeafTaskFragments(taskFragment -> {
+            if (taskFragment.getCompanionTaskFragment() == this) {
+                taskFragment.setCompanionTaskFragment(null /* companionTaskFragment */);
+            }
+        }, false /* traverseTopToBottom */);
     }
 
     @Override
diff --git a/services/tests/wmtests/src/com/android/server/wm/TaskFragmentTest.java b/services/tests/wmtests/src/com/android/server/wm/TaskFragmentTest.java
index 50998693e5ed..dab842c4aa5a 100644
--- a/services/tests/wmtests/src/com/android/server/wm/TaskFragmentTest.java
+++ b/services/tests/wmtests/src/com/android/server/wm/TaskFragmentTest.java
@@ -45,11 +45,13 @@ import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.Mockito.clearInvocations;
 
 import android.content.res.Configuration;
+import android.graphics.Color;
 import android.graphics.Rect;
 import android.os.Binder;
 import android.platform.test.annotations.Presubmit;
 import android.view.SurfaceControl;
 import android.window.ITaskFragmentOrganizer;
+import android.window.TaskFragmentAnimationParams;
 import android.window.TaskFragmentInfo;
 import android.window.TaskFragmentOrganizer;
 
@@ -299,35 +301,44 @@ public class TaskFragmentTest extends WindowTestsBase {
 
     @Test
     public void testEmbeddedTaskFragmentEnterPip_singleActivity_resetOrganizerOverrideConfig() {
-        final TaskFragment taskFragment = new TaskFragmentBuilder(mAtm)
-                .setOrganizer(mOrganizer)
-                .setFragmentToken(new Binder())
-                .setCreateParentTask()
-                .createActivityCount(1)
+        final Task task = createTask(mDisplayContent);
+        final TaskFragment taskFragment0 = createTaskFragmentWithEmbeddedActivity(task, mOrganizer);
+        final TaskFragment taskFragment1 = new TaskFragmentBuilder(mAtm)
+                .setParentTask(task)
                 .build();
-        final Task task = taskFragment.getTask();
-        final ActivityRecord activity = taskFragment.getTopMostActivity();
+        final ActivityRecord activity = taskFragment0.getTopMostActivity();
         final Rect taskFragmentBounds = new Rect(0, 0, 300, 1000);
         task.setWindowingMode(WINDOWING_MODE_FULLSCREEN);
-        taskFragment.setWindowingMode(WINDOWING_MODE_MULTI_WINDOW);
-        taskFragment.setBounds(taskFragmentBounds);
+        taskFragment0.setWindowingMode(WINDOWING_MODE_MULTI_WINDOW);
+        taskFragment0.setBounds(taskFragmentBounds);
+        taskFragment0.setAdjacentTaskFragment(taskFragment1);
+        taskFragment0.setCompanionTaskFragment(taskFragment1);
+        taskFragment0.setAnimationParams(new TaskFragmentAnimationParams.Builder()
+                .setAnimationBackgroundColor(Color.GREEN)
+                .build());
 
         assertEquals(taskFragmentBounds, activity.getBounds());
         assertEquals(WINDOWING_MODE_MULTI_WINDOW, activity.getWindowingMode());
+        assertEquals(taskFragment1, taskFragment0.getAdjacentTaskFragment());
+        assertEquals(taskFragment1, taskFragment0.getCompanionTaskFragment());
+        assertNotEquals(TaskFragmentAnimationParams.DEFAULT, taskFragment0.getAnimationParams());
 
         // Move activity to pinned root task.
         mRootWindowContainer.moveActivityToPinnedRootTask(activity,
                 null /* launchIntoPipHostActivity */, "test");
 
         // Ensure taskFragment requested config is reset.
-        assertEquals(taskFragment, activity.getOrganizedTaskFragment());
+        assertEquals(taskFragment0, activity.getOrganizedTaskFragment());
         assertEquals(task, activity.getTask());
         assertTrue(task.inPinnedWindowingMode());
-        assertTrue(taskFragment.inPinnedWindowingMode());
+        assertTrue(taskFragment0.inPinnedWindowingMode());
         final Rect taskBounds = task.getBounds();
-        assertEquals(taskBounds, taskFragment.getBounds());
+        assertEquals(taskBounds, taskFragment0.getBounds());
         assertEquals(taskBounds, activity.getBounds());
-        assertEquals(Configuration.EMPTY, taskFragment.getRequestedOverrideConfiguration());
+        assertEquals(Configuration.EMPTY, taskFragment0.getRequestedOverrideConfiguration());
+        assertNull(taskFragment0.getAdjacentTaskFragment());
+        assertNull(taskFragment0.getCompanionTaskFragment());
+        assertEquals(TaskFragmentAnimationParams.DEFAULT, taskFragment0.getAnimationParams());
         // Because the whole Task is entering PiP, no need to record for future reparent.
         assertNull(activity.mLastTaskFragmentOrganizerBeforePip);
     }
@@ -335,18 +346,8 @@ public class TaskFragmentTest extends WindowTestsBase {
     @Test
     public void testEmbeddedTaskFragmentEnterPip_multiActivities_notifyOrganizer() {
         final Task task = createTask(mDisplayContent);
-        final TaskFragment taskFragment0 = new TaskFragmentBuilder(mAtm)
-                .setParentTask(task)
-                .setOrganizer(mOrganizer)
-                .setFragmentToken(new Binder())
-                .createActivityCount(1)
-                .build();
-        final TaskFragment taskFragment1 = new TaskFragmentBuilder(mAtm)
-                .setParentTask(task)
-                .setOrganizer(mOrganizer)
-                .setFragmentToken(new Binder())
-                .createActivityCount(1)
-                .build();
+        final TaskFragment taskFragment0 = createTaskFragmentWithEmbeddedActivity(task, mOrganizer);
+        final TaskFragment taskFragment1 = createTaskFragmentWithEmbeddedActivity(task, mOrganizer);
         final ActivityRecord activity0 = taskFragment0.getTopMostActivity();
         final ActivityRecord activity1 = taskFragment1.getTopMostActivity();
         activity0.setVisibility(true /* visible */, false /* deferHidingClient */);
diff --git a/services/tests/wmtests/src/com/android/server/wm/WindowTestsBase.java b/services/tests/wmtests/src/com/android/server/wm/WindowTestsBase.java
index e5efe05b8b32..3e8d259e26ec 100644
--- a/services/tests/wmtests/src/com/android/server/wm/WindowTestsBase.java
+++ b/services/tests/wmtests/src/com/android/server/wm/WindowTestsBase.java
@@ -79,6 +79,7 @@ import android.graphics.Insets;
 import android.graphics.Rect;
 import android.hardware.HardwareBuffer;
 import android.hardware.display.DisplayManager;
+import android.os.Binder;
 import android.os.Build;
 import android.os.Bundle;
 import android.os.IBinder;
@@ -750,11 +751,11 @@ class WindowTestsBase extends SystemServiceTestsBase {
     }
 
     /**
-     * Creates a {@link TaskFragment} with {@link ActivityRecord} and attach it to the
+     * Creates a {@link TaskFragment} with {@link ActivityRecord}, and attaches it to the
      * {@code parentTask}.
      *
-     * @param parentTask the {@link Task} this TaskFragment is going to be attached
-     * @return the created TaskFragment
+     * @param parentTask the {@link Task} this {@link TaskFragment} is going to be attached.
+     * @return the created {@link TaskFragment}
      */
     static TaskFragment createTaskFragmentWithActivity(@NonNull Task parentTask) {
         return new TaskFragmentBuilder(parentTask.mAtmService)
@@ -763,13 +764,27 @@ class WindowTestsBase extends SystemServiceTestsBase {
                 .build();
     }
 
+    /**
+     * Creates an embedded {@link TaskFragment} organized by {@code organizer} with
+     * {@link ActivityRecord}, and attaches it to the {@code parentTask}.
+     *
+     * @param parentTask the {@link Task} this {@link TaskFragment} is going to be attached.
+     * @param organizer  the {@link TaskFragmentOrganizer} this {@link TaskFragment} is going to be
+     *                   organized by.
+     * @return the created {@link TaskFragment}
+     */
     static TaskFragment createTaskFragmentWithEmbeddedActivity(@NonNull Task parentTask,
-            TaskFragmentOrganizer organizer) {
-        return new TaskFragmentBuilder(parentTask.mAtmService)
+            @NonNull TaskFragmentOrganizer organizer) {
+        final IBinder fragmentToken = new Binder();
+        final TaskFragment taskFragment = new TaskFragmentBuilder(parentTask.mAtmService)
                 .setParentTask(parentTask)
                 .createActivityCount(1)
                 .setOrganizer(organizer)
+                .setFragmentToken(fragmentToken)
                 .build();
+        parentTask.mAtmService.mWindowOrganizerController.mLaunchTaskFragments
+                .put(fragmentToken, taskFragment);
+        return taskFragment;
     }
 
     /** Creates a {@link DisplayContent} that supports IME and adds it to the system. */