diff --git a/libs/WindowManager/Shell/res/layout/pip2_menu_action.xml b/libs/WindowManager/Shell/res/layout/pip2_menu_action.xml
new file mode 100644
index 0000000..04ece31
--- /dev/null
+++ b/libs/WindowManager/Shell/res/layout/pip2_menu_action.xml
@@ -0,0 +1,39 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!--
+    Copyright (C) 2024 The Android Open Source Project
+
+    Licensed under the Apache License, Version 2.0 (the "License");
+    you may not use this file except in compliance with the License.
+    You may obtain a copy of the License at
+
+         http://www.apache.org/licenses/LICENSE-2.0
+
+    Unless required by applicable law or agreed to in writing, software
+    distributed under the License is distributed on an "AS IS" BASIS,
+    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+    See the License for the specific language governing permissions and
+    limitations under the License.
+-->
+<com.android.wm.shell.pip2.phone.PipMenuActionView
+    xmlns:android="http://schemas.android.com/apk/res/android"
+    android:layout_width="@dimen/pip_action_size"
+    android:layout_height="@dimen/pip_action_size"
+    android:background="?android:selectableItemBackgroundBorderless"
+    android:forceHasOverlappingRendering="false">
+
+    <ImageView
+        android:id="@+id/custom_close_bg"
+        android:layout_width="@dimen/pip_custom_close_bg_size"
+        android:layout_height="@dimen/pip_custom_close_bg_size"
+        android:layout_gravity="center"
+        android:src="@drawable/pip_custom_close_bg"
+        android:visibility="gone"/>
+
+    <ImageView
+        android:id="@+id/image"
+        android:layout_width="@dimen/pip_action_inner_size"
+        android:layout_height="@dimen/pip_action_inner_size"
+        android:layout_gravity="center"
+        android:scaleType="fitXY"/>
+
+</com.android.wm.shell.pip2.phone.PipMenuActionView>
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/pip/Pip2Module.java b/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/pip/Pip2Module.java
index 702552e..3508ece 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/pip/Pip2Module.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/pip/Pip2Module.java
@@ -127,11 +127,9 @@
     @Provides
     static PipScheduler providePipScheduler(Context context,
             PipBoundsState pipBoundsState,
-            PhonePipMenuController pipMenuController,
             @ShellMainThread ShellExecutor mainExecutor,
             PipTransitionState pipTransitionState) {
-        return new PipScheduler(context, pipBoundsState, pipMenuController,
-                mainExecutor, pipTransitionState);
+        return new PipScheduler(context, pipBoundsState, mainExecutor, pipTransitionState);
     }
 
     @WMSingleton
@@ -140,10 +138,13 @@
             PipBoundsState pipBoundsState, PipMediaController pipMediaController,
             SystemWindows systemWindows,
             PipUiEventLogger pipUiEventLogger,
+            PipTaskListener pipTaskListener,
+            @NonNull PipTransitionState pipTransitionState,
             @ShellMainThread ShellExecutor mainExecutor,
             @ShellMainThread Handler mainHandler) {
         return new PhonePipMenuController(context, pipBoundsState, pipMediaController,
-                systemWindows, pipUiEventLogger, mainExecutor, mainHandler);
+                systemWindows, pipUiEventLogger, pipTaskListener, pipTransitionState, mainExecutor,
+                mainHandler);
     }
 
 
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/phone/PhonePipMenuController.java b/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/phone/PhonePipMenuController.java
index 9cfe162..8c1e5e6 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/phone/PhonePipMenuController.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/phone/PhonePipMenuController.java
@@ -18,10 +18,12 @@
 
 import static android.view.WindowManager.SHELL_ROOT_LAYER_PIP;
 
+import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.app.RemoteAction;
 import android.content.Context;
 import android.graphics.Rect;
+import android.os.Bundle;
 import android.os.Debug;
 import android.os.Handler;
 import android.os.RemoteException;
@@ -52,7 +54,8 @@
  * The current media session provides actions whenever there are no valid actions provided by the
  * current PiP activity. Otherwise, those actions always take precedence.
  */
-public class PhonePipMenuController implements PipMenuController {
+public class PhonePipMenuController implements PipMenuController,
+        PipTransitionState.PipTransitionStateChangedListener {
 
     private static final String TAG = "PhonePipMenuController";
     private static final boolean DEBUG = false;
@@ -113,6 +116,11 @@
 
     private PipMenuView mPipMenuView;
 
+    private final PipTaskListener mPipTaskListener;
+
+    @NonNull
+    private final PipTransitionState mPipTransitionState;
+
     private SurfaceControl mLeash;
 
     private ActionListener mMediaActionListener = new ActionListener() {
@@ -125,15 +133,27 @@
 
     public PhonePipMenuController(Context context, PipBoundsState pipBoundsState,
             PipMediaController mediaController, SystemWindows systemWindows,
-            PipUiEventLogger pipUiEventLogger,
-            ShellExecutor mainExecutor, Handler mainHandler) {
+            PipUiEventLogger pipUiEventLogger, PipTaskListener pipTaskListener,
+            @NonNull PipTransitionState pipTransitionState, ShellExecutor mainExecutor,
+            Handler mainHandler) {
         mContext = context;
         mPipBoundsState = pipBoundsState;
         mMediaController = mediaController;
         mSystemWindows = systemWindows;
+        mPipTaskListener = pipTaskListener;
+        mPipTransitionState = pipTransitionState;
         mMainExecutor = mainExecutor;
         mMainHandler = mainHandler;
         mPipUiEventLogger = pipUiEventLogger;
+
+        mPipTransitionState.addPipTransitionStateChangedListener(this);
+
+        mPipTaskListener.addParamsChangedListener(new PipTaskListener.PipParamsChangedCallback() {
+            @Override
+            public void onActionsChanged(List<RemoteAction> actions, RemoteAction closeAction) {
+                setAppActions(actions, closeAction);
+            }
+        });
     }
 
     public boolean isMenuVisible() {
@@ -438,8 +458,7 @@
      * Sets the menu actions to the actions provided by the current PiP menu.
      */
     @Override
-    public void setAppActions(List<RemoteAction> appActions,
-            RemoteAction closeAction) {
+    public void setAppActions(List<RemoteAction> appActions, RemoteAction closeAction) {
         mAppActions = appActions;
         mCloseAction = closeAction;
         updateMenuActions();
@@ -468,8 +487,8 @@
      */
     private void updateMenuActions() {
         if (mPipMenuView != null) {
-            mPipMenuView.setActions(mPipBoundsState.getBounds(),
-                    resolveMenuActions(), mCloseAction);
+            mPipMenuView.setActions(mPipBoundsState.getBounds(), resolveMenuActions(),
+                    mCloseAction);
         }
     }
 
@@ -567,6 +586,29 @@
         }
     }
 
+    @Override
+    public void onPipTransitionStateChanged(@PipTransitionState.TransitionState int oldState,
+            @PipTransitionState.TransitionState int newState, Bundle extra) {
+        switch (newState) {
+            case PipTransitionState.ENTERED_PIP:
+                attach(mPipTransitionState.mPinnedTaskLeash);
+                break;
+            case PipTransitionState.EXITED_PIP:
+                detach();
+                break;
+            case PipTransitionState.CHANGED_PIP_BOUNDS:
+                updateMenuLayout(mPipBoundsState.getBounds());
+                hideMenu();
+                break;
+            case PipTransitionState.CHANGING_PIP_BOUNDS:
+                hideMenu();
+                break;
+            case PipTransitionState.SCHEDULED_BOUNDS_CHANGE:
+                hideMenu();
+                break;
+        }
+    }
+
     void dump(PrintWriter pw, String prefix) {
         final String innerPrefix = prefix + "  ";
         pw.println(prefix + TAG);
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/phone/PipMenuView.java b/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/phone/PipMenuView.java
index a29104c..0910919 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/phone/PipMenuView.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/phone/PipMenuView.java
@@ -447,7 +447,7 @@
                 final LayoutInflater inflater = LayoutInflater.from(mContext);
                 while (mActionsGroup.getChildCount() < mActions.size()) {
                     final PipMenuActionView actionView = (PipMenuActionView) inflater.inflate(
-                            R.layout.pip_menu_action, mActionsGroup, false);
+                            R.layout.pip2_menu_action, mActionsGroup, false);
                     mActionsGroup.addView(actionView);
                 }
 
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/phone/PipMotionHelper.java b/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/phone/PipMotionHelper.java
index 0324fdb..268c3a2 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/phone/PipMotionHelper.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/phone/PipMotionHelper.java
@@ -336,7 +336,7 @@
         }
         cancelPhysicsAnimation();
         mMenuController.hideMenu(ANIM_TYPE_NONE, false /* resize */);
-        // mPipTaskOrganizer.exitPip(skipAnimation ? 0 : LEAVE_PIP_DURATION, enterSplit);
+        mPipScheduler.scheduleExitPipViaExpand();
     }
 
     /**
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/phone/PipScheduler.java b/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/phone/PipScheduler.java
index f4defdc..d4f190e 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/phone/PipScheduler.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/phone/PipScheduler.java
@@ -52,7 +52,6 @@
 
     private final Context mContext;
     private final PipBoundsState mPipBoundsState;
-    private final PhonePipMenuController mPipMenuController;
     private final ShellExecutor mMainExecutor;
     private final PipTransitionState mPipTransitionState;
     private PipSchedulerReceiver mSchedulerReceiver;
@@ -97,12 +96,10 @@
 
     public PipScheduler(Context context,
             PipBoundsState pipBoundsState,
-            PhonePipMenuController pipMenuController,
             ShellExecutor mainExecutor,
             PipTransitionState pipTransitionState) {
         mContext = context;
         mPipBoundsState = pipBoundsState;
-        mPipMenuController = pipMenuController;
         mMainExecutor = mainExecutor;
         mPipTransitionState = pipTransitionState;
 
@@ -263,7 +260,6 @@
             return;
         }
         mPipBoundsState.setBounds(newBounds);
-        mPipMenuController.updateMenuLayout(newBounds);
         maybeUpdateMovementBounds();
     }
 }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/phone/PipTaskListener.java b/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/phone/PipTaskListener.java
index 262c14d..c58de2c 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/phone/PipTaskListener.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/phone/PipTaskListener.java
@@ -20,6 +20,7 @@
 
 import android.app.ActivityManager;
 import android.app.PictureInPictureParams;
+import android.app.RemoteAction;
 import android.content.Context;
 import android.graphics.Rect;
 import android.os.Bundle;
@@ -37,6 +38,9 @@
 import com.android.wm.shell.pip2.animation.PipResizeAnimator;
 import com.android.wm.shell.shared.annotations.ShellMainThread;
 
+import java.util.ArrayList;
+import java.util.List;
+
 /**
  * A Task Listener implementation used only for CUJs and trigger paths that cannot be initiated via
  * Transitions framework directly.
@@ -57,6 +61,7 @@
             new PictureInPictureParams.Builder().build();
 
     private boolean mWaitingForAspectRatioChange = false;
+    private final List<PipParamsChangedCallback> mPipParamsChangedListeners = new ArrayList<>();
 
     public PipTaskListener(Context context,
             ShellTaskOrganizer shellTaskOrganizer,
@@ -85,10 +90,25 @@
         if (mPictureInPictureParams.equals(params)) {
             return;
         }
+        if (PipUtils.remoteActionsChanged(params.getActions(), mPictureInPictureParams.getActions())
+                || !PipUtils.remoteActionsMatch(params.getCloseAction(),
+                mPictureInPictureParams.getCloseAction())) {
+            for (PipParamsChangedCallback listener : mPipParamsChangedListeners) {
+                listener.onActionsChanged(params.getActions(), params.getCloseAction());
+            }
+        }
         mPictureInPictureParams.copyOnlySet(params != null ? params
                 : new PictureInPictureParams.Builder().build());
     }
 
+    /** Add a PipParamsChangedCallback listener. */
+    public void addParamsChangedListener(PipParamsChangedCallback listener) {
+        if (mPipParamsChangedListeners.contains(listener)) {
+            return;
+        }
+        mPipParamsChangedListeners.add(listener);
+    }
+
     @NonNull
     public PictureInPictureParams getPictureInPictureParams() {
         return mPictureInPictureParams;
@@ -164,4 +184,12 @@
                 break;
         }
     }
+
+    public interface PipParamsChangedCallback {
+        /**
+         * Called if either the actions or the close action changed.
+         */
+        default void onActionsChanged(List<RemoteAction> actions, RemoteAction closeAction) {
+        }
+    }
 }