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 67963a3..d94e8e4 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 @@
                 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 @@
     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 @@
         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 @@
         } 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 @@
         } 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 57ba6bb..2d7fd7e 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 @@
             // 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 2b93682..8a21039 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 @@
         // 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 @@
                 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 @@
         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 @@
             @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 @@
             @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 @@
         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 38ac719..bc082b1 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 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 @@
     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 @@
     /**
      * 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 @@
     /**
      * 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 @@
         }
     }
 
-    @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 @@
 
     /**
      * Updates the last requested windowing mode.
+     * @see WindowContainerTransaction#setWindowingMode
      */
     void setLastRequestedWindowingMode(@WindowingMode int windowingModes) {
         mLastRequestedWindowingMode = windowingModes;
@@ -610,6 +631,7 @@
 
     /**
      * 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 @@
 
     /**
      * 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 a41e63f..96dd86a 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 @@
     }
 
     @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 e8aa2c8..9333d65 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.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 @@
                     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 eb06b91..294e90b 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 @@
     void removeImmediately() {
         mIsRemovalRequested = false;
         resetAdjacentTaskFragment();
-        cleanUp();
+        cleanUpEmbeddedTaskFragment();
         final boolean shouldExecuteAppTransition =
                 mClearedTaskFragmentForPip && isTaskVisibleRequested();
         super.removeImmediately();
@@ -2830,10 +2830,20 @@
     }
 
     /** 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 5099869..dab842c 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.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 @@
 
     @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 @@
     @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 e5efe05..3e8d259 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.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 @@
     }
 
     /**
-     * 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 @@
                 .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. */