Use proxy controller to maintain platform session and notification

With this change, the notification controller that is connected by
`MediaNotificationManager`, is used as a proxy controller of the
System UI controller. An app can use the proxy at connection time
and during the lifetime of the session for configuration of the
platform session and the media notification on all API levels.

This includes using custom layout and available player and session
commands of the proxy to maintain the platform session (actions,
custom actions, session extras) and the `MediaNotification.Provider`.

The legacy System UI controller is hidden from the public API,
instead the app interacts with the Media3 proxy:

- System UI is hidden from `MediaSession.getConnectedControllers()`.
- Calls from System UI to methods of `MediaSession.Callback`/
  `MediaLibrarySession.Callback` are mapped to the `ControllerInfo`
  of the proxy controller.
- When `getControllerForCurrentRequest()` is called during an operation of
  System UI the proxy `ControllerInfo` is returned.

PiperOrigin-RevId: 567606117

RELEASENOTES.md

@@ -67,6 +67,9 @@
         `android.media.session.MediaSession.setMediaButtonBroadcastReceiver()`
         above API 31 to avoid problems with deprecated API on Samsung devices
         ([#167](https://github.com/androidx/media/issues/167)).
+    *   Use the media notification controller as proxy to set available commands
+        and custom layout used to populate the notification and the platform
+        session.
 *   UI:
 *   Downloads:
 *   OkHttp Extension:

demos/session/src/main/java/androidx/media3/demo/session/PlaybackService.kt

@@ -33,7 +33,6 @@ import androidx.media3.common.util.UnstableApi
 import androidx.media3.datasource.DataSourceBitmapLoader
 import androidx.media3.exoplayer.ExoPlayer
 import androidx.media3.session.*
-import androidx.media3.session.LibraryResult.RESULT_ERROR_NOT_SUPPORTED
 import androidx.media3.session.MediaSession.ConnectionResult
 import androidx.media3.session.MediaSession.ControllerInfo
 import com.google.common.collect.ImmutableList
@@ -45,7 +44,7 @@ class PlaybackService : MediaLibraryService() {
 
   private lateinit var player: ExoPlayer
   private lateinit var mediaLibrarySession: MediaLibrarySession
-  private lateinit var customCommands: List<CommandButton>
+  private lateinit var customLayoutCommandButtons: List<CommandButton>
 
   companion object {
     private const val CUSTOM_COMMAND_TOGGLE_SHUFFLE_MODE_ON =
@@ -60,7 +59,7 @@ class PlaybackService : MediaLibraryService() {
   @OptIn(UnstableApi::class) // MediaSessionService.setListener
   override fun onCreate() {
     super.onCreate()
-    customCommands =
+    customLayoutCommandButtons =
       listOf(
         getShuffleCommandButton(
           SessionCommand(CUSTOM_COMMAND_TOGGLE_SHUFFLE_MODE_ON, Bundle.EMPTY)
@@ -100,33 +99,47 @@ class PlaybackService : MediaLibraryService() {
     // ConnectionResult.AcceptedResultBuilder
     @OptIn(UnstableApi::class)
     override fun onConnect(session: MediaSession, controller: ControllerInfo): ConnectionResult {
+      if (session.isMediaNotificationController(controller)) {
+        // Set the required available session commands and the custom layout for the notification
+        // on all API levels.
         val availableSessionCommands =
           ConnectionResult.DEFAULT_SESSION_AND_LIBRARY_COMMANDS.buildUpon()
-      for (commandButton in customCommands) {
+        // Add the session commands of all command buttons.
-        // Add custom command to available session commands.
+        customLayoutCommandButtons.forEach { commandButton ->
           commandButton.sessionCommand?.let { availableSessionCommands.add(it) }
         }
+        // Select the buttons to display.
+        val customLayout =
+          ImmutableList.of(customLayoutCommandButtons[if (player.shuffleModeEnabled) 1 else 0])
         return ConnectionResult.AcceptedResultBuilder(session)
           .setAvailableSessionCommands(availableSessionCommands.build())
+          .setCustomLayout(customLayout)
           .build()
       }
+      // Default commands without custom layout for common controllers.
+      return ConnectionResult.AcceptedResultBuilder(session).build()
+    }
 
+    @OptIn(UnstableApi::class) // MediaSession.isMediaNotificationController
     override fun onCustomCommand(
       session: MediaSession,
       controller: ControllerInfo,
       customCommand: SessionCommand,
       args: Bundle
     ): ListenableFuture<SessionResult> {
+      if (!session.isMediaNotificationController(controller)) {
+        return Futures.immediateFuture(SessionResult(SessionResult.RESULT_ERROR_NOT_SUPPORTED))
+      }
       if (CUSTOM_COMMAND_TOGGLE_SHUFFLE_MODE_ON == customCommand.customAction) {
         // Enable shuffling.
         player.shuffleModeEnabled = true
         // Change the custom layout to contain the `Disable shuffling` command.
-        session.setCustomLayout(ImmutableList.of(customCommands[1]))
+        session.setCustomLayout(controller, ImmutableList.of(customLayoutCommandButtons[1]))
       } else if (CUSTOM_COMMAND_TOGGLE_SHUFFLE_MODE_OFF == customCommand.customAction) {
         // Disable shuffling.
         player.shuffleModeEnabled = false
         // Change the custom layout to contain the `Enable shuffling` command.
-        session.setCustomLayout(ImmutableList.of(customCommands[0]))
+        session.setCustomLayout(controller, ImmutableList.of(customLayoutCommandButtons[0]))
       }
       return Futures.immediateFuture(SessionResult(SessionResult.RESULT_SUCCESS))
     }
@@ -136,12 +149,6 @@ class PlaybackService : MediaLibraryService() {
       browser: ControllerInfo,
       params: LibraryParams?
     ): ListenableFuture<LibraryResult<MediaItem>> {
-      if (params != null && params.isRecent) {
-        // The service currently does not support playback resumption. Tell System UI by returning
-        // an error of type 'RESULT_ERROR_NOT_SUPPORTED' for a `params.isRecent` request. See
-        // https://github.com/androidx/media/issues/355
-        return Futures.immediateFuture(LibraryResult.ofError(RESULT_ERROR_NOT_SUPPORTED))
-      }
       return Futures.immediateFuture(LibraryResult.ofItem(MediaItemTree.getRootItem(), params))
     }
 
@@ -234,7 +241,6 @@ class PlaybackService : MediaLibraryService() {
     mediaLibrarySession =
       MediaLibrarySession.Builder(this, player, librarySessionCallback)
         .setSessionActivity(getSingleTopActivity())
-        .setCustomLayout(ImmutableList.of(customCommands[0]))
         .setBitmapLoader(CacheBitmapLoader(DataSourceBitmapLoader(/* context= */ this)))
         .build()
   }
@@ -270,10 +276,6 @@ class PlaybackService : MediaLibraryService() {
       .build()
   }
 
-  private fun ignoreFuture(customLayout: ListenableFuture<SessionResult>) {
-    /* Do nothing. */
-  }
-
   @OptIn(UnstableApi::class) // MediaSessionService.Listener
   private inner class MediaSessionServiceListener : Listener {
 

libraries/session/src/main/java/androidx/media3/session/CommandButton.java

@@ -215,6 +215,7 @@ public final class CommandButton implements Bundleable {
         sessionCommand, playerCommand, iconResId, displayName, new Bundle(extras), isEnabled);
   }
 
+  /** Checks the given command button for equality while ignoring {@link #extras}. */
   @Override
   public boolean equals(@Nullable Object obj) {
     if (this == obj) {

libraries/session/src/main/java/androidx/media3/session/MediaLibrarySessionImpl.java

@@ -128,6 +128,13 @@ import java.util.concurrent.Future;
 
   public void notifyChildrenChanged(
       ControllerInfo browser, String parentId, int itemCount, @Nullable LibraryParams params) {
+    if (isMediaNotificationControllerConnected() && isMediaNotificationController(browser)) {
+      ControllerInfo systemUiBrowser = getSystemUiControllerInfo();
+      if (systemUiBrowser == null) {
+        return;
+      }
+      browser = systemUiBrowser;
+    }
     dispatchRemoteControllerTaskWithoutReturn(
         browser,
         (callback, seq) -> {
@@ -140,6 +147,13 @@ import java.util.concurrent.Future;
 
   public void notifySearchResultChanged(
       ControllerInfo browser, String query, int itemCount, @Nullable LibraryParams params) {
+    if (isMediaNotificationControllerConnected() && isMediaNotificationController(browser)) {
+      ControllerInfo systemUiBrowser = getSystemUiControllerInfo();
+      if (systemUiBrowser == null) {
+        return;
+      }
+      browser = systemUiBrowser;
+    }
     dispatchRemoteControllerTaskWithoutReturn(
         browser, (callback, seq) -> callback.onSearchResultChanged(seq, query, itemCount, params));
   }
@@ -163,7 +177,7 @@ import java.util.concurrent.Future;
                   params));
     }
     ListenableFuture<LibraryResult<MediaItem>> future =
-        callback.onGetLibraryRoot(instance, browser, params);
+        callback.onGetLibraryRoot(instance, resolveControllerInfoForCallback(browser), params);
     future.addListener(
         () -> {
           @Nullable LibraryResult<MediaItem> result = tryGetFutureResult(future);
@@ -204,7 +218,8 @@ import java.util.concurrent.Future;
                   params));
     }
     ListenableFuture<LibraryResult<ImmutableList<MediaItem>>> future =
-        callback.onGetChildren(instance, browser, parentId, page, pageSize, params);
+        callback.onGetChildren(
+            instance, resolveControllerInfoForCallback(browser), parentId, page, pageSize, params);
     future.addListener(
         () -> {
           @Nullable LibraryResult<ImmutableList<MediaItem>> result = tryGetFutureResult(future);
@@ -220,7 +235,7 @@ import java.util.concurrent.Future;
   public ListenableFuture<LibraryResult<MediaItem>> onGetItemOnHandler(
       ControllerInfo browser, String mediaId) {
     ListenableFuture<LibraryResult<MediaItem>> future =
-        callback.onGetItem(instance, browser, mediaId);
+        callback.onGetItem(instance, resolveControllerInfoForCallback(browser), mediaId);
     future.addListener(
         () -> {
           @Nullable LibraryResult<MediaItem> result = tryGetFutureResult(future);
@@ -251,7 +266,8 @@ import java.util.concurrent.Future;
     // so we explicitly null-check the result to fail early if an app accidentally returns null.
     ListenableFuture<LibraryResult<Void>> future =
         checkNotNull(
-            callback.onSubscribe(instance, browser, parentId, params),
+            callback.onSubscribe(
+                instance, resolveControllerInfoForCallback(browser), parentId, params),
             "onSubscribe must return non-null future");
 
     // When error happens, remove from the subscription list.
@@ -270,7 +286,7 @@ import java.util.concurrent.Future;
   public ListenableFuture<LibraryResult<Void>> onUnsubscribeOnHandler(
       ControllerInfo browser, String parentId) {
     ListenableFuture<LibraryResult<Void>> future =
-        callback.onUnsubscribe(instance, browser, parentId);
+        callback.onUnsubscribe(instance, resolveControllerInfoForCallback(browser), parentId);
 
     future.addListener(
         () -> removeSubscription(checkStateNotNull(browser.getControllerCb()), parentId),
@@ -282,7 +298,7 @@ import java.util.concurrent.Future;
   public ListenableFuture<LibraryResult<Void>> onSearchOnHandler(
       ControllerInfo browser, String query, @Nullable LibraryParams params) {
     ListenableFuture<LibraryResult<Void>> future =
-        callback.onSearch(instance, browser, query, params);
+        callback.onSearch(instance, resolveControllerInfoForCallback(browser), query, params);
     future.addListener(
         () -> {
           @Nullable LibraryResult<Void> result = tryGetFutureResult(future);
@@ -301,7 +317,8 @@ import java.util.concurrent.Future;
       int pageSize,
       @Nullable LibraryParams params) {
     ListenableFuture<LibraryResult<ImmutableList<MediaItem>>> future =
-        callback.onGetSearchResult(instance, browser, query, page, pageSize, params);
+        callback.onGetSearchResult(
+            instance, resolveControllerInfoForCallback(browser), query, page, pageSize, params);
     future.addListener(
         () -> {
           @Nullable LibraryResult<ImmutableList<MediaItem>> result = tryGetFutureResult(future);
@@ -410,6 +427,10 @@ import java.util.concurrent.Future;
           ControllerInfo controller, @Nullable LibraryParams params) {
     SettableFuture<LibraryResult<ImmutableList<MediaItem>>> settableFuture =
         SettableFuture.create();
+    controller =
+        isMediaNotificationControllerConnected()
+            ? checkNotNull(getMediaNotificationControllerInfo())
+            : controller;
     ListenableFuture<MediaSession.MediaItemsWithStartPosition> future =
         callback.onPlaybackResumption(instance, controller);
     Futures.addCallback(

libraries/session/src/main/java/androidx/media3/session/MediaSession.java

@@ -32,7 +32,6 @@ import android.os.Looper;
 import android.os.RemoteException;
 import android.support.v4.media.session.MediaControllerCompat;
 import android.support.v4.media.session.MediaSessionCompat;
-import android.support.v4.media.session.PlaybackStateCompat;
 import android.view.KeyEvent;
 import androidx.annotation.GuardedBy;
 import androidx.annotation.Nullable;
@@ -775,9 +774,7 @@ public class MediaSession {
   /**
    * Returns whether the given media controller info belongs to the media notification controller.
    *
-   * <p>Use this method for instance in {@link Callback#onConnect(MediaSession, ControllerInfo)} to
+   * <p>See {@link #getMediaNotificationControllerInfo()}.
-   * recognize the media notification controller and provide a {@link ConnectionResult} with a
-   * custom layout specific for this controller.
    *
    * @param controllerInfo The controller info.
    * @return Whether the controller info belongs to the media notification controller.
@@ -792,8 +789,28 @@ public class MediaSession {
    *
    * <p>Use this controller info to set {@linkplain #setAvailableCommands(ControllerInfo,
    * SessionCommands, Player.Commands) available commands} and {@linkplain
-   * #setCustomLayout(ControllerInfo, List) custom layout} that are applied to the media
+   * #setCustomLayout(ControllerInfo, List) custom layout} that are consistently applied to the
-   * notification.
+   * media notification on all API levels.
+   *
+   * <p>Available {@linkplain SessionCommands session commands} of the media notification controller
+   * are used to enable or disable buttons of the custom layout before it is passed to the
+   * {@linkplain MediaNotification.Provider#createNotification(MediaSession, ImmutableList,
+   * MediaNotification.ActionFactory, MediaNotification.Provider.Callback) notification provider}.
+   * Disabled command buttons are not converted to notification actions when using {@link
+   * DefaultMediaNotificationProvider}. This affects the media notification displayed by System UI
+   * below API 33.
+   *
+   * <p>The available session commands of the media notification controller are used to maintain
+   * custom actions of the platform session (see {@code PlaybackStateCompat.getCustomActions()}).
+   * Command buttons of the custom layout are disabled or enabled according to the available session
+   * commands. Disabled command buttons are not converted to custom actions of the platform session.
+   * This affects the media notification displayed by System UI <a
+   * href="https://developer.android.com/about/versions/13/behavior-changes-13#playback-controls">starting
+   * with API 33</a>.
+   *
+   * <p>The available {@linkplain Player.Commands player commands} are intersected with the actual
+   * available commands of the underlying player to determine the playback actions of the platform
+   * session (see {@code PlaybackStateCompat.getActions()}).
    */
   @UnstableApi
   @Nullable
@@ -815,7 +832,8 @@ public class MediaSession {
    * <p>On the controller side, {@link
    * MediaController.Listener#onCustomLayoutChanged(MediaController, List)} is only called if the
    * new custom layout is different to the custom layout the {@link
-   * MediaController#getCustomLayout() controller already has available}.
+   * MediaController#getCustomLayout() controller already has available}. Note that this comparison
+   * uses {@link CommandButton#equals} and therefore ignores {@link CommandButton#extras}.
    *
    * <p>It's up to controller's decision how to represent the layout in its own UI.
    *
@@ -836,22 +854,17 @@ public class MediaSession {
   /**
    * Sets the custom layout that can initially be set when building the session.
    *
-   * <p>Calling this method broadcasts the custom layout to all connected Media3 controllers and
+   * <p>Calling this method broadcasts the custom layout to all connected Media3 controllers,
-   * converts the {@linkplain CommandButton command buttons} to {@linkplain
+   * including the {@linkplain #getMediaNotificationControllerInfo() media notification controller}.
-   * PlaybackStateCompat.CustomAction custom actions of the playback state} of the platform media
-   * session (see {@code
-   * PlaybackStateCompat.Builder#addCustomAction(PlaybackStateCompat.CustomAction)}). The {@link
-   * CommandButton#isEnabled} flag is set according to the available commands of the controller and
-   * overrides a value that has been set by the app. The platform media session won't see any
-   * commands that are disabled.
    *
-   * <p>On the controller side, {@link
+   * <p>On the controller side, the {@linkplain CommandButton#isEnabled enabled} flag is set
-   * MediaController.Listener#onCustomLayoutChanged(MediaController, List)} is only called if the
+   * according to the available commands of the controller which overrides a value that has been set
-   * new custom layout is different to the custom layout the {@linkplain
+   * by the session.
-   * MediaController#getCustomLayout() controller already has available}.
    *
-   * <p>When converting, the {@linkplain SessionCommand#customExtras custom extras of the session
+   * <p>{@link MediaController.Listener#onCustomLayoutChanged(MediaController, List)} is only called
-   * command} is used for the extras of the legacy custom action.
+   * if the new custom layout is different to the custom layout the {@linkplain
+   * MediaController#getCustomLayout() controller already has available}. Note that {@link Bundle
+   * extras} are ignored when comparing {@linkplain CommandButton command buttons}.
    *
    * <p>Controllers that connect after calling this method will have the new custom layout available
    * with the initial connection result. A custom layout specific to a controller can be set when

libraries/session/src/main/java/androidx/media3/session/MediaSessionImpl.java

@@ -130,6 +130,7 @@ import org.checkerframework.checker.nullness.qual.MonotonicNonNull;
 
   // Should be only accessed on the application looper
   private long sessionPositionUpdateDelayMs;
+  private boolean isMediaNotificationControllerConnected;
   private ImmutableList<CommandButton> customLayout;
 
   public MediaSessionImpl(
@@ -191,10 +192,19 @@ import org.checkerframework.checker.nullness.qual.MonotonicNonNull;
 
     sessionLegacyStub =
         new MediaSessionLegacyStub(/* session= */ thisRef, sessionUri, applicationHandler);
-
+    // For PlayerWrapper, use the same default commands as the proxy controller gets when the app
-    PlayerWrapper playerWrapper = new PlayerWrapper(player, playIfSuppressed);
+    // doesn't overrides the default commands in `onConnect`. When the default is overridden by the
+    // app in `onConnect`, the default set here will be overridden with these values.
+    MediaSession.ConnectionResult connectionResult =
+        new MediaSession.ConnectionResult.AcceptedResultBuilder(instance).build();
+    PlayerWrapper playerWrapper =
+        new PlayerWrapper(
+            player,
+            playIfSuppressed,
+            customLayout,
+            connectionResult.availableSessionCommands,
+            connectionResult.availablePlayerCommands);
     this.playerWrapper = playerWrapper;
-    this.playerWrapper.setCustomLayout(customLayout);
     postOrRun(
         applicationHandler,
         () ->
@@ -212,13 +222,18 @@ import org.checkerframework.checker.nullness.qual.MonotonicNonNull;
       return;
     }
     setPlayerInternal(
-        /* oldPlayerWrapper= */ playerWrapper, new PlayerWrapper(player, playIfSuppressed));
+        /* oldPlayerWrapper= */ playerWrapper,
+        new PlayerWrapper(
+            player,
+            playIfSuppressed,
+            playerWrapper.getCustomLayout(),
+            playerWrapper.getAvailableSessionCommands(),
+            playerWrapper.getAvailablePlayerCommands()));
   }
 
   private void setPlayerInternal(
       @Nullable PlayerWrapper oldPlayerWrapper, PlayerWrapper newPlayerWrapper) {
     playerWrapper = newPlayerWrapper;
-    playerWrapper.setCustomLayout(customLayout);
     if (oldPlayerWrapper != null) {
       oldPlayerWrapper.removeListener(checkStateNotNull(this.playerListener));
     }
@@ -295,14 +310,27 @@ import org.checkerframework.checker.nullness.qual.MonotonicNonNull;
   public List<ControllerInfo> getConnectedControllers() {
     List<ControllerInfo> controllers = new ArrayList<>();
     controllers.addAll(sessionStub.getConnectedControllersManager().getConnectedControllers());
+    if (isMediaNotificationControllerConnected) {
+      ImmutableList<ControllerInfo> legacyControllers =
+          sessionLegacyStub.getConnectedControllersManager().getConnectedControllers();
+      for (int i = 0; i < legacyControllers.size(); i++) {
+        ControllerInfo legacyController = legacyControllers.get(i);
+        if (!isSystemUiController(legacyController)) {
+          controllers.add(legacyController);
+        }
+      }
+    } else {
       controllers.addAll(
           sessionLegacyStub.getConnectedControllersManager().getConnectedControllers());
+    }
     return controllers;
   }
 
   @Nullable
   public ControllerInfo getControllerForCurrentRequest() {
-    return controllerForCurrentRequest;
+    return controllerForCurrentRequest != null
+        ? resolveControllerInfoForCallback(controllerForCurrentRequest)
+        : null;
   }
 
   public boolean isConnected(ControllerInfo controller) {
@@ -372,17 +400,39 @@ import org.checkerframework.checker.nullness.qual.MonotonicNonNull;
     return null;
   }
 
-  public ListenableFuture<SessionResult> setCustomLayout(
+  /** Returns whether the media notification controller is connected. */
-      ControllerInfo controller, List<CommandButton> layout) {
+  protected boolean isMediaNotificationControllerConnected() {
-    return dispatchRemoteControllerTask(
+    return isMediaNotificationControllerConnected;
-        controller, (controller1, seq) -> controller1.setCustomLayout(seq, layout));
   }
 
-  public void setCustomLayout(List<CommandButton> layout) {
+  /**
-    customLayout = ImmutableList.copyOf(layout);
+   * Sets the custom layout for the given {@link MediaController}.
+   *
+   * @param controller The controller.
+   * @param customLayout The custom layout.
+   * @return The session result from the controller.
+   */
+  public ListenableFuture<SessionResult> setCustomLayout(
+      ControllerInfo controller, ImmutableList<CommandButton> customLayout) {
+    if (isMediaNotificationController(controller)) {
+      playerWrapper.setCustomLayout(customLayout);
+      sessionLegacyStub.updateLegacySessionPlaybackStateCompat();
+    }
+    return dispatchRemoteControllerTask(
+        controller, (controller1, seq) -> controller1.setCustomLayout(seq, customLayout));
+  }
+
+  /** Sets the custom layout of the session and sends the custom layout to all controllers. */
+  public void setCustomLayout(ImmutableList<CommandButton> customLayout) {
+    this.customLayout = customLayout;
     playerWrapper.setCustomLayout(customLayout);
     dispatchRemoteControllerTaskWithoutReturn(
-        (controller, seq) -> controller.setCustomLayout(seq, layout));
+        (controller, seq) -> controller.setCustomLayout(seq, customLayout));
+  }
+
+  /** Returns the custom layout. */
+  public ImmutableList<CommandButton> getCustomLayout() {
+    return customLayout;
   }
 
   public void setSessionExtras(Bundle sessionExtras) {
@@ -408,6 +458,18 @@ import org.checkerframework.checker.nullness.qual.MonotonicNonNull;
   public void setAvailableCommands(
       ControllerInfo controller, SessionCommands sessionCommands, Player.Commands playerCommands) {
     if (sessionStub.getConnectedControllersManager().isConnected(controller)) {
+      if (isMediaNotificationController(controller)) {
+        playerWrapper.setAvailableCommands(sessionCommands, playerCommands);
+        sessionLegacyStub.updateLegacySessionPlaybackStateCompat();
+        ControllerInfo systemUiControllerInfo = getSystemUiControllerInfo();
+        if (systemUiControllerInfo != null) {
+          // Set the available commands of the proxy controller to the ConnectedControllerRecord of
+          // the hidden System UI controller.
+          sessionLegacyStub
+              .getConnectedControllersManager()
+              .updateCommandsFromSession(systemUiControllerInfo, sessionCommands, playerCommands);
+        }
+      }
       sessionStub
           .getConnectedControllersManager()
           .updateCommandsFromSession(controller, sessionCommands, playerCommands);
@@ -482,45 +544,103 @@ import org.checkerframework.checker.nullness.qual.MonotonicNonNull;
   }
 
   public MediaSession.ConnectionResult onConnectOnHandler(ControllerInfo controller) {
-    return checkNotNull(
+    if (isMediaNotificationControllerConnected && isSystemUiController(controller)) {
-        callback.onConnect(instance, controller), "Callback.onConnect must return non-null future");
+      // Hide System UI and provide the connection result from the `PlayerWrapper` state.
+      return new MediaSession.ConnectionResult.AcceptedResultBuilder(instance)
+          .setAvailableSessionCommands(playerWrapper.getAvailableSessionCommands())
+          .setAvailablePlayerCommands(playerWrapper.getAvailablePlayerCommands())
+          .setCustomLayout(playerWrapper.getCustomLayout())
+          .build();
+    }
+    MediaSession.ConnectionResult connectionResult =
+        checkNotNull(
+            callback.onConnect(instance, controller),
+            "Callback.onConnect must return non-null future");
+    if (isMediaNotificationController(controller)) {
+      isMediaNotificationControllerConnected = true;
+      playerWrapper.setAvailableCommands(
+          connectionResult.availableSessionCommands, connectionResult.availablePlayerCommands);
+      playerWrapper.setCustomLayout(
+          connectionResult.customLayout != null
+              ? connectionResult.customLayout
+              : instance.getCustomLayout());
+      sessionLegacyStub.updateLegacySessionPlaybackStateCompat();
+    }
+    return connectionResult;
   }
 
   public void onPostConnectOnHandler(ControllerInfo controller) {
+    if (isMediaNotificationControllerConnected && isSystemUiController(controller)) {
+      // Hide System UI. Apps can use the media notification controller to maintain the platform
+      // session
+      return;
+    }
     callback.onPostConnect(instance, controller);
   }
 
   public void onDisconnectedOnHandler(ControllerInfo controller) {
+    if (isMediaNotificationControllerConnected) {
+      if (isSystemUiController(controller)) {
+        // Hide System UI controller. Apps can use the media notification controller to maintain the
+        // platform session.
+        return;
+      } else if (isMediaNotificationController(controller)) {
+        isMediaNotificationControllerConnected = false;
+      }
+    }
     callback.onDisconnected(instance, controller);
   }
 
   @SuppressWarnings("deprecation") // Calling deprecated callback method.
   public @SessionResult.Code int onPlayerCommandRequestOnHandler(
       ControllerInfo controller, @Player.Command int playerCommand) {
-    return callback.onPlayerCommandRequest(instance, controller, playerCommand);
+    return callback.onPlayerCommandRequest(
+        instance, resolveControllerInfoForCallback(controller), playerCommand);
   }
 
   public ListenableFuture<SessionResult> onSetRatingOnHandler(
       ControllerInfo controller, String mediaId, Rating rating) {
     return checkNotNull(
-        callback.onSetRating(instance, controller, mediaId, rating),
+        callback.onSetRating(
+            instance, resolveControllerInfoForCallback(controller), mediaId, rating),
         "Callback.onSetRating must return non-null future");
   }
 
   public ListenableFuture<SessionResult> onSetRatingOnHandler(
       ControllerInfo controller, Rating rating) {
     return checkNotNull(
-        callback.onSetRating(instance, controller, rating),
+        callback.onSetRating(instance, resolveControllerInfoForCallback(controller), rating),
         "Callback.onSetRating must return non-null future");
   }
 
   public ListenableFuture<SessionResult> onCustomCommandOnHandler(
-      ControllerInfo browser, SessionCommand command, Bundle extras) {
+      ControllerInfo controller, SessionCommand command, Bundle extras) {
     return checkNotNull(
-        callback.onCustomCommand(instance, browser, command, extras),
+        callback.onCustomCommand(
+            instance, resolveControllerInfoForCallback(controller), command, extras),
         "Callback.onCustomCommandOnHandler must return non-null future");
   }
 
+  protected ListenableFuture<List<MediaItem>> onAddMediaItemsOnHandler(
+      ControllerInfo controller, List<MediaItem> mediaItems) {
+    return checkNotNull(
+        callback.onAddMediaItems(
+            instance, resolveControllerInfoForCallback(controller), mediaItems),
+        "Callback.onAddMediaItems must return a non-null future");
+  }
+
+  protected ListenableFuture<MediaItemsWithStartPosition> onSetMediaItemsOnHandler(
+      ControllerInfo controller, List<MediaItem> mediaItems, int startIndex, long startPositionMs) {
+    return checkNotNull(
+        callback.onSetMediaItems(
+            instance,
+            resolveControllerInfoForCallback(controller),
+            mediaItems,
+            startIndex,
+            startPositionMs),
+        "Callback.onSetMediaItems must return a non-null future");
+  }
+
   public void connectFromService(
       IMediaController caller,
       int controllerVersion,
@@ -555,20 +675,6 @@ import org.checkerframework.checker.nullness.qual.MonotonicNonNull;
     return applicationHandler;
   }
 
-  protected ListenableFuture<List<MediaItem>> onAddMediaItemsOnHandler(
-      ControllerInfo controller, List<MediaItem> mediaItems) {
-    return checkNotNull(
-        callback.onAddMediaItems(instance, controller, mediaItems),
-        "Callback.onAddMediaItems must return a non-null future");
-  }
-
-  protected ListenableFuture<MediaItemsWithStartPosition> onSetMediaItemsOnHandler(
-      ControllerInfo controller, List<MediaItem> mediaItems, int startIndex, long startPositionMs) {
-    return checkNotNull(
-        callback.onSetMediaItems(instance, controller, mediaItems, startIndex, startPositionMs),
-        "Callback.onSetMediaItems must return a non-null future");
-  }
-
   protected boolean isReleased() {
     synchronized (lock) {
       return closed;
@@ -580,10 +686,6 @@ import org.checkerframework.checker.nullness.qual.MonotonicNonNull;
     return sessionActivity;
   }
 
-  protected ImmutableList<CommandButton> getCustomLayout() {
-    return customLayout;
-  }
-
   @UnstableApi
   protected void setSessionActivity(PendingIntent sessionActivity) {
     if (Objects.equals(this.sessionActivity, sessionActivity)) {
@@ -603,6 +705,12 @@ import org.checkerframework.checker.nullness.qual.MonotonicNonNull;
     }
   }
 
+  protected ControllerInfo resolveControllerInfoForCallback(ControllerInfo controller) {
+    return isMediaNotificationControllerConnected && isSystemUiController(controller)
+        ? checkNotNull(getMediaNotificationControllerInfo())
+        : controller;
+  }
+
   /**
    * Gets the service binder from the MediaBrowserServiceCompat. Should be only called by the thread
    * with a Looper.
@@ -693,7 +801,7 @@ import org.checkerframework.checker.nullness.qual.MonotonicNonNull;
     @Nullable
     ListenableFuture<MediaItemsWithStartPosition> future =
         checkNotNull(
-            callback.onPlaybackResumption(instance, controller),
+            callback.onPlaybackResumption(instance, resolveControllerInfoForCallback(controller)),
             "Callback.onPlaybackResumption must return a non-null future");
     // Use a direct executor when an immediate future is returned to execute the player setup in the
     // caller's looper event on the application thread.

libraries/session/src/main/java/androidx/media3/session/MediaSessionLegacyStub.java

@@ -823,6 +823,15 @@ import org.checkerframework.checker.initialization.qual.Initialized;
     connectionTimeoutMs = timeoutMs;
   }
 
+  public void updateLegacySessionPlaybackStateCompat() {
+    postOrRun(
+        sessionImpl.getApplicationHandler(),
+        () ->
+            sessionImpl
+                .getSessionCompat()
+                .setPlaybackState(sessionImpl.getPlayerWrapper().createPlaybackStateCompat()));
+  }
+
   private void handleMediaRequest(MediaItem mediaItem, boolean play) {
     dispatchSessionTaskWithPlayerCommand(
         COMMAND_SET_MEDIA_ITEM,
@@ -1081,16 +1090,12 @@ import org.checkerframework.checker.initialization.qual.Initialized;
 
     @Override
     public void onPlayerError(int seq, @Nullable PlaybackException playerError) {
-      sessionImpl
+      updateLegacySessionPlaybackStateCompat();
-          .getSessionCompat()
-          .setPlaybackState(sessionImpl.getPlayerWrapper().createPlaybackStateCompat());
     }
 
     @Override
     public void setCustomLayout(int seq, List<CommandButton> layout) {
-      sessionImpl
+      updateLegacySessionPlaybackStateCompat();
-          .getSessionCompat()
-          .setPlaybackState(sessionImpl.getPlayerWrapper().createPlaybackStateCompat());
     }
 
     @Override

libraries/session/src/main/java/androidx/media3/session/PlayerWrapper.java

@@ -21,6 +21,7 @@ import static androidx.media3.common.util.Util.msToUs;
 import static androidx.media3.common.util.Util.postOrRun;
 import static androidx.media3.session.MediaConstants.EXTRAS_KEY_MEDIA_ID_COMPAT;
 import static androidx.media3.session.MediaConstants.EXTRAS_KEY_PLAYBACK_SPEED_COMPAT;
+import static androidx.media3.session.MediaUtils.intersect;
 
 import android.media.AudioManager;
 import android.os.Bundle;
@@ -70,12 +71,43 @@ import java.util.List;
   @Nullable private String legacyErrorMessage;
   @Nullable private Bundle legacyErrorExtras;
   private ImmutableList<CommandButton> customLayout;
+  private SessionCommands availableSessionCommands;
+  private Commands availablePlayerCommands;
 
-  public PlayerWrapper(Player player, boolean playIfSuppressed) {
+  public PlayerWrapper(
+      Player player,
+      boolean playIfSuppressed,
+      ImmutableList<CommandButton> customLayout,
+      SessionCommands availableSessionCommands,
+      Commands availablePlayerCommands) {
     super(player);
     this.playIfSuppressed = playIfSuppressed;
+    this.customLayout = customLayout;
+    this.availableSessionCommands = availableSessionCommands;
+    this.availablePlayerCommands = availablePlayerCommands;
     legacyStatusCode = STATUS_CODE_SUCCESS_COMPAT;
-    customLayout = ImmutableList.of();
+  }
+
+  public void setAvailableCommands(
+      SessionCommands availableSessionCommands, Commands availablePlayerCommands) {
+    this.availableSessionCommands = availableSessionCommands;
+    this.availablePlayerCommands = availablePlayerCommands;
+  }
+
+  public SessionCommands getAvailableSessionCommands() {
+    return availableSessionCommands;
+  }
+
+  public Commands getAvailablePlayerCommands() {
+    return availablePlayerCommands;
+  }
+
+  public void setCustomLayout(ImmutableList<CommandButton> customLayout) {
+    this.customLayout = customLayout;
+  }
+
+  /* package */ ImmutableList<CommandButton> getCustomLayout() {
+    return customLayout;
   }
 
   /**
@@ -104,11 +136,6 @@ import java.util.List;
     return legacyStatusCode;
   }
 
-  /** Sets the custom layout. */
-  public void setCustomLayout(ImmutableList<CommandButton> customLayout) {
-    this.customLayout = customLayout;
-  }
-
   /** Clears the legacy error status. */
   public void clearLegacyErrorStatus() {
     legacyStatusCode = STATUS_CODE_SUCCESS_COMPAT;
@@ -974,7 +1001,7 @@ import java.util.List;
     int state = MediaUtils.convertToPlaybackStateCompatState(/* player= */ this, playIfSuppressed);
     // Always advertise ACTION_SET_RATING.
     long actions = PlaybackStateCompat.ACTION_SET_RATING;
-    Commands availableCommands = getAvailableCommands();
+    Commands availableCommands = intersect(availablePlayerCommands, getAvailableCommands());
     for (int i = 0; i < availableCommands.size(); i++) {
       actions |= convertCommandToPlaybackStateActions(availableCommands.get(i));
     }
@@ -1006,7 +1033,9 @@ import java.util.List;
       CommandButton commandButton = customLayout.get(i);
       if (commandButton.sessionCommand != null) {
         SessionCommand sessionCommand = commandButton.sessionCommand;
-        if (sessionCommand.commandCode == SessionCommand.COMMAND_CODE_CUSTOM) {
+        if (sessionCommand.commandCode == SessionCommand.COMMAND_CODE_CUSTOM
+            && CommandButton.isEnabled(
+                commandButton, availableSessionCommands, availablePlayerCommands)) {
           builder.addCustomAction(
               new PlaybackStateCompat.CustomAction.Builder(
                       sessionCommand.customAction,

libraries/session/src/main/java/androidx/media3/session/SessionCommand.java

@@ -157,6 +157,7 @@ public final class SessionCommand implements Bundleable {
     customExtras = new Bundle(checkNotNull(extras));
   }
 
+  /** Checks the given session command for equality while ignoring extras. */
   @Override
   public boolean equals(@Nullable Object obj) {
     if (!(obj instanceof SessionCommand)) {

libraries/session/src/test/java/androidx/media3/session/PlayerWrapperTest.java

@@ -22,6 +22,7 @@ import static org.mockito.Mockito.when;
 import android.os.Looper;
 import androidx.media3.common.Player;
 import androidx.test.ext.junit.runners.AndroidJUnit4;
+import com.google.common.collect.ImmutableList;
 import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
@@ -42,7 +43,13 @@ public class PlayerWrapperTest {
 
   @Before
   public void setUp() {
-    playerWrapper = new PlayerWrapper(player, /* playIfSuppressed= */ true);
+    playerWrapper =
+        new PlayerWrapper(
+            player,
+            /* playIfSuppressed= */ true,
+            ImmutableList.of(),
+            SessionCommands.EMPTY,
+            Player.Commands.EMPTY);
     when(player.isCommandAvailable(anyInt())).thenReturn(true);
     when(player.getApplicationLooper()).thenReturn(Looper.myLooper());
   }

libraries/test_session_common/src/main/java/androidx/media3/test/session/common/MediaSessionConstants.java

@@ -30,7 +30,8 @@ public class MediaSessionConstants {
       "onTracksChanged_videoToAudioTransition";
   public static final String TEST_SET_SHOW_PLAY_BUTTON_IF_SUPPRESSED_TO_FALSE =
       "testSetShowPlayButtonIfSuppressedToFalse";
-
+  public static final String TEST_MEDIA_CONTROLLER_COMPAT_CALLBACK_WITH_MEDIA_SESSION_TEST =
+      "MediaControllerCompatCallbackWithMediaSessionTest";
   // Bundle keys
   public static final String KEY_AVAILABLE_SESSION_COMMANDS = "availableSessionCommands";
   public static final String KEY_CONTROLLER = "controllerKey";

libraries/test_session_current/build.gradle

@@ -45,6 +45,7 @@ dependencies {
     implementation 'androidx.test:core:' + androidxTestCoreVersion
     implementation project(modulePrefix + 'test-data')
     androidTestImplementation project(modulePrefix + 'lib-exoplayer')
+    androidTestImplementation project(modulePrefix + 'test-utils')
     androidTestImplementation 'androidx.test.ext:junit:' + androidxTestJUnitVersion
     androidTestImplementation 'androidx.test.ext:truth:' + androidxTestTruthVersion
     androidTestImplementation 'androidx.test:core:' + androidxTestCoreVersion

libraries/test_session_current/src/androidTest/java/androidx/media3/session/MediaControllerCompatCallbackWithMediaSessionTest.java

@@ -20,6 +20,7 @@ import static android.support.v4.media.MediaMetadataCompat.METADATA_KEY_MEDIA_ID
 import static android.support.v4.media.MediaMetadataCompat.METADATA_KEY_USER_RATING;
 import static androidx.media3.common.Player.STATE_ENDED;
 import static androidx.media3.common.Player.STATE_READY;
+import static androidx.media3.test.session.common.MediaSessionConstants.TEST_MEDIA_CONTROLLER_COMPAT_CALLBACK_WITH_MEDIA_SESSION_TEST;
 import static androidx.media3.test.session.common.MediaSessionConstants.TEST_SET_SHOW_PLAY_BUTTON_IF_SUPPRESSED_TO_FALSE;
 import static androidx.media3.test.session.common.TestUtils.LONG_TIMEOUT_MS;
 import static androidx.media3.test.session.common.TestUtils.TIMEOUT_MS;
@@ -78,10 +79,10 @@ import org.junit.runner.RunWith;
 @LargeTest
 public class MediaControllerCompatCallbackWithMediaSessionTest {
 
-  private static final String TAG = "MCCCallbackTestWithMS2";
+  private static final String SESSION_ID =
-  private static final float EPSILON = 1e-6f;
+      TEST_MEDIA_CONTROLLER_COMPAT_CALLBACK_WITH_MEDIA_SESSION_TEST;
 
-  @Rule public final HandlerThreadTestRule threadTestRule = new HandlerThreadTestRule(TAG);
+  @Rule public final HandlerThreadTestRule threadTestRule = new HandlerThreadTestRule(SESSION_ID);
 
   private Context context;
   private TestHandler handler;
@@ -92,7 +93,10 @@ public class MediaControllerCompatCallbackWithMediaSessionTest {
   public void setUp() throws Exception {
     context = ApplicationProvider.getApplicationContext();
     handler = threadTestRule.getHandler();
-    session = new RemoteMediaSession(TAG, context, null);
+    Bundle tokenExtras = new Bundle();
+    tokenExtras.putBoolean(
+        MediaSessionProviderService.KEY_ENABLE_FAKE_MEDIA_NOTIFICATION_MANAGER_CONTROLLER, true);
+    session = new RemoteMediaSession(SESSION_ID, context, tokenExtras);
     controllerCompat = new MediaControllerCompat(context, session.getCompatToken());
   }
 
@@ -181,7 +185,7 @@ public class MediaControllerCompatCallbackWithMediaSessionTest {
   public void getError_withPlayerErrorAfterConnected_returnsError() throws Exception {
     PlaybackException testPlayerError =
         new PlaybackException(
-            /* messaage= */ "testremote",
+            /* message= */ "testremote",
             /* cause= */ null,
             PlaybackException.ERROR_CODE_REMOTE_ERROR);
     Bundle playerConfig =
@@ -207,7 +211,7 @@ public class MediaControllerCompatCallbackWithMediaSessionTest {
   public void playerError_notified() throws Exception {
     PlaybackException testPlayerError =
         new PlaybackException(
-            /* messaage= */ "player error",
+            /* message= */ "player error",
             /* cause= */ null,
             PlaybackException.ERROR_CODE_UNSPECIFIED);
 
@@ -937,53 +941,46 @@ public class MediaControllerCompatCallbackWithMediaSessionTest {
 
   @Test
   public void setCustomLayout_onPlaybackStateCompatChangedCalled() throws Exception {
-    List<CommandButton> buttons = new ArrayList<>();
     Bundle extras1 = new Bundle();
     extras1.putString("key", "value-1");
-    CommandButton button1 =
+    SessionCommand command1 = new SessionCommand("command1", extras1);
-        new CommandButton.Builder()
-            .setSessionCommand(new SessionCommand("action1", extras1))
-            .setDisplayName("actionName1")
-            .setIconResId(1)
-            .build();
     Bundle extras2 = new Bundle();
     extras2.putString("key", "value-2");
-    CommandButton button2 =
+    SessionCommand command2 = new SessionCommand("command2", extras2);
+    ImmutableList<CommandButton> customLayout =
+        ImmutableList.of(
             new CommandButton.Builder()
-            .setSessionCommand(new SessionCommand("action2", extras2))
+                .setSessionCommand(command1)
-            .setDisplayName("actionName2")
+                .setDisplayName("command1")
+                .setIconResId(1)
+                .build()
+                .copyWithIsEnabled(true),
+            new CommandButton.Builder()
+                .setSessionCommand(command2)
+                .setDisplayName("command2")
                 .setIconResId(2)
-            .build();
+                .build()
-    buttons.add(button1);
+                .copyWithIsEnabled(true));
-    buttons.add(button2);
+    List<ImmutableList<CommandButton>> reportedCustomLayouts = new ArrayList<>();
-    List<String> receivedActions = new ArrayList<>();
+    CountDownLatch latch1 = new CountDownLatch(2);
-    List<String> receivedDisplayNames = new ArrayList<>();
-    List<String> receivedBundleValues = new ArrayList<>();
-    List<Integer> receivedIconResIds = new ArrayList<>();
-    CountDownLatch latch = new CountDownLatch(1);
     MediaControllerCompat.Callback callback =
         new MediaControllerCompat.Callback() {
           @Override
           public void onPlaybackStateChanged(PlaybackStateCompat state) {
-            List<PlaybackStateCompat.CustomAction> layout = state.getCustomActions();
+            reportedCustomLayouts.add(MediaUtils.convertToCustomLayout(state));
-            for (PlaybackStateCompat.CustomAction action : layout) {
+            latch1.countDown();
-              receivedActions.add(action.getAction());
-              receivedDisplayNames.add(String.valueOf(action.getName()));
-              receivedBundleValues.add(action.getExtras().getString("key"));
-              receivedIconResIds.add(action.getIcon());
-            }
-            latch.countDown();
           }
         };
     controllerCompat.registerCallback(callback, handler);
 
-    session.setCustomLayout(buttons);
+    session.setCustomLayout(customLayout);
+    session.setAvailableCommands(
+        SessionCommands.EMPTY.buildUpon().add(command1).add(command2).build(),
+        Player.Commands.EMPTY);
 
-    assertThat(latch.await(TIMEOUT_MS, MILLISECONDS)).isTrue();
+    assertThat(latch1.await(TIMEOUT_MS, MILLISECONDS)).isTrue();
-    assertThat(receivedActions).containsExactly("action1", "action2").inOrder();
+    assertThat(reportedCustomLayouts.get(0)).containsExactly(customLayout.get(0));
-    assertThat(receivedDisplayNames).containsExactly("actionName1", "actionName2").inOrder();
+    assertThat(reportedCustomLayouts.get(1)).isEqualTo(customLayout);
-    assertThat(receivedIconResIds).containsExactly(1, 2).inOrder();
-    assertThat(receivedBundleValues).containsExactly("value-1", "value-2").inOrder();
   }
 
   @Test
@@ -1403,6 +1400,6 @@ public class MediaControllerCompatCallbackWithMediaSessionTest {
       PlaybackStateCompat state, PlaybackException playerError) {
     assertThat(state.getState()).isEqualTo(PlaybackStateCompat.STATE_ERROR);
     assertThat(state.getErrorCode()).isEqualTo(PlaybackStateCompat.ERROR_CODE_UNKNOWN_ERROR);
-    assertThat(state.getErrorMessage()).isEqualTo(playerError.getMessage());
+    assertThat(state.getErrorMessage().toString()).isEqualTo(playerError.getMessage());
   }
 }

libraries/test_session_current/src/androidTest/java/androidx/media3/session/MediaControllerCompatPlaybackStateCompatActionsWithMediaSessionTest.java

@@ -43,6 +43,8 @@ import androidx.media3.common.Timeline;
 import androidx.media3.common.util.ConditionVariable;
 import androidx.media3.common.util.Consumer;
 import androidx.media3.exoplayer.ExoPlayer;
+import androidx.media3.session.MediaSession.ConnectionResult;
+import androidx.media3.session.MediaSession.ConnectionResult.AcceptedResultBuilder;
 import androidx.media3.test.session.R;
 import androidx.media3.test.session.common.HandlerThreadTestRule;
 import androidx.test.core.app.ApplicationProvider;
@@ -51,6 +53,7 @@ import androidx.test.filters.LargeTest;
 import com.google.common.collect.ImmutableList;
 import com.google.common.util.concurrent.Futures;
 import com.google.common.util.concurrent.ListenableFuture;
+import com.google.common.util.concurrent.MoreExecutors;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.concurrent.CountDownLatch;
@@ -1459,18 +1462,23 @@ public class MediaControllerCompatPlaybackStateCompatActionsWithMediaSessionTest
                 .setIconResId(R.drawable.media3_notification_pause)
                 .setSessionCommand(command2)
                 .build());
-    MediaSession mediaSession = createMediaSession(player, /* callback= */ null, customLayout);
+    MediaSession.Callback callback =
+        new MediaSession.Callback() {
+          @Override
+          public ConnectionResult onConnect(
+              MediaSession session, MediaSession.ControllerInfo controller) {
+            return new AcceptedResultBuilder(session)
+                .setAvailableSessionCommands(
+                    ConnectionResult.DEFAULT_SESSION_COMMANDS.buildUpon().add(command1).build())
+                .build();
+          }
+        };
+    MediaSession mediaSession = createMediaSession(player, callback, customLayout);
+    connectMediaNotificationController(mediaSession);
     MediaControllerCompat controllerCompat = createMediaControllerCompat(mediaSession);
 
-    // Wait until a playback state is sent to the controller.
+    assertThat(MediaUtils.convertToCustomLayout(controllerCompat.getPlaybackState()))
-    PlaybackStateCompat firstPlaybackState =
+        .containsExactly(customLayout.get(0).copyWithIsEnabled(true));
-        getFirstPlaybackState(controllerCompat, threadTestRule.getHandler());
-
-    assertThat(MediaUtils.convertToCustomLayout(firstPlaybackState))
-        .containsExactly(
-            customLayout.get(0).copyWithIsEnabled(true),
-            customLayout.get(1).copyWithIsEnabled(true))
-        .inOrder();
     mediaSession.release();
     releasePlayer(player);
   }
@@ -1497,11 +1505,23 @@ public class MediaControllerCompatPlaybackStateCompatActionsWithMediaSessionTest
                 .setIconResId(R.drawable.media3_notification_pause)
                 .setSessionCommand(command2)
                 .build());
-    MediaSession mediaSession = createMediaSession(player);
+    MediaSession.Callback callback =
+        new MediaSession.Callback() {
+          @Override
+          public ConnectionResult onConnect(
+              MediaSession session, MediaSession.ControllerInfo controller) {
+            return new ConnectionResult.AcceptedResultBuilder(session)
+                .setAvailableSessionCommands(
+                    ConnectionResult.DEFAULT_SESSION_COMMANDS.buildUpon().add(command1).build())
+                .build();
+          }
+        };
+    MediaSession mediaSession = createMediaSession(player, callback);
+    connectMediaNotificationController(mediaSession);
     MediaControllerCompat controllerCompat = createMediaControllerCompat(mediaSession);
+    ImmutableList<CommandButton> initialCustomLayout =
+        MediaUtils.convertToCustomLayout(controllerCompat.getPlaybackState());
     AtomicReference<List<CommandButton>> reportedCustomLayout = new AtomicReference<>();
-    // Wait until a playback state is sent to the controller.
-    getFirstPlaybackState(controllerCompat, threadTestRule.getHandler());
     CountDownLatch latch = new CountDownLatch(1);
     controllerCompat.registerCallback(
         new MediaControllerCompat.Callback() {
@@ -1516,14 +1536,97 @@ public class MediaControllerCompatPlaybackStateCompatActionsWithMediaSessionTest
     getInstrumentation().runOnMainSync(() -> mediaSession.setCustomLayout(customLayout));
 
     assertThat(latch.await(TIMEOUT_MS, MILLISECONDS)).isTrue();
+    assertThat(initialCustomLayout).isEmpty();
     assertThat(reportedCustomLayout.get())
-        .containsExactly(
+        .containsExactly(customLayout.get(0).copyWithIsEnabled(true));
-            customLayout.get(0).copyWithIsEnabled(true),
-            customLayout.get(1).copyWithIsEnabled(true));
     mediaSession.release();
     releasePlayer(player);
   }
 
+  @Test
+  public void
+      playerWithCustomLayout_setCustomLayoutForMediaNotificationController_playbackStateChangedWithCustomActionsChanged()
+          throws Exception {
+    Player player = createDefaultPlayer();
+    Bundle extras1 = new Bundle();
+    extras1.putString("key1", "value1");
+    Bundle extras2 = new Bundle();
+    extras1.putString("key2", "value2");
+    SessionCommand command1 = new SessionCommand("command1", extras1);
+    SessionCommand command2 = new SessionCommand("command2", extras2);
+    ImmutableList<CommandButton> customLayout =
+        ImmutableList.of(
+            new CommandButton.Builder()
+                .setDisplayName("button1")
+                .setIconResId(R.drawable.media3_notification_play)
+                .setSessionCommand(command1)
+                .build(),
+            new CommandButton.Builder()
+                .setDisplayName("button2")
+                .setIconResId(R.drawable.media3_notification_pause)
+                .setSessionCommand(command2)
+                .build());
+    MediaSession.Callback callback =
+        new MediaSession.Callback() {
+          @Override
+          public ConnectionResult onConnect(
+              MediaSession session, MediaSession.ControllerInfo controller) {
+            return new ConnectionResult.AcceptedResultBuilder(session)
+                .setAvailableSessionCommands(
+                    ConnectionResult.DEFAULT_SESSION_COMMANDS.buildUpon().add(command1).build())
+                .build();
+          }
+        };
+    MediaSession mediaSession = createMediaSession(player, callback);
+    connectMediaNotificationController(mediaSession);
+    MediaControllerCompat controllerCompat = createMediaControllerCompat(mediaSession);
+    ImmutableList<CommandButton> initialCustomLayout =
+        MediaUtils.convertToCustomLayout(controllerCompat.getPlaybackState());
+    AtomicReference<List<CommandButton>> reportedCustomLayout = new AtomicReference<>();
+    CountDownLatch latch = new CountDownLatch(1);
+    controllerCompat.registerCallback(
+        new MediaControllerCompat.Callback() {
+          @Override
+          public void onPlaybackStateChanged(PlaybackStateCompat state) {
+            reportedCustomLayout.set(MediaUtils.convertToCustomLayout(state));
+            latch.countDown();
+          }
+        },
+        threadTestRule.getHandler());
+
+    getInstrumentation()
+        .runOnMainSync(
+            () ->
+                mediaSession.setCustomLayout(
+                    mediaSession.getMediaNotificationControllerInfo(), customLayout));
+
+    assertThat(latch.await(TIMEOUT_MS, MILLISECONDS)).isTrue();
+    assertThat(initialCustomLayout).isEmpty();
+    assertThat(reportedCustomLayout.get())
+        .containsExactly(customLayout.get(0).copyWithIsEnabled(true));
+    mediaSession.release();
+    releasePlayer(player);
+  }
+
+  /**
+   * Connect a controller that mimics the media notification controller that is connected by {@link
+   * MediaNotificationManager} when the session is running in the service.
+   */
+  private void connectMediaNotificationController(MediaSession mediaSession)
+      throws InterruptedException {
+    CountDownLatch connectionLatch = new CountDownLatch(1);
+    Bundle connectionHints = new Bundle();
+    connectionHints.putBoolean(MediaNotificationManager.KEY_MEDIA_NOTIFICATION_MANAGER, true);
+    ListenableFuture<MediaController> mediaNotificationControllerFuture =
+        new MediaController.Builder(
+                ApplicationProvider.getApplicationContext(), mediaSession.getToken())
+            .setConnectionHints(connectionHints)
+            .buildAsync();
+    mediaNotificationControllerFuture.addListener(
+        connectionLatch::countDown, MoreExecutors.directExecutor());
+    assertThat(connectionLatch.await(TIMEOUT_MS, MILLISECONDS)).isTrue();
+  }
+
   private PlaybackStateCompat getFirstPlaybackState(
       MediaControllerCompat mediaControllerCompat, Handler handler) throws InterruptedException {
     LinkedBlockingDeque<PlaybackStateCompat> playbackStateCompats = new LinkedBlockingDeque<>();

libraries/test_session_current/src/androidTest/java/androidx/media3/session/MediaSessionServiceTest.java

@@ -27,18 +27,24 @@ import android.content.Context;
 import android.os.Bundle;
 import android.os.Handler;
 import android.os.Looper;
+import android.support.v4.media.session.MediaControllerCompat;
+import android.support.v4.media.session.PlaybackStateCompat;
 import androidx.media3.common.MediaItem;
 import androidx.media3.common.Player;
 import androidx.media3.common.util.ConditionVariable;
 import androidx.media3.exoplayer.ExoPlayer;
 import androidx.media3.session.MediaSession.ControllerInfo;
+import androidx.media3.test.session.R;
 import androidx.media3.test.session.common.HandlerThreadTestRule;
 import androidx.media3.test.session.common.MainLooperTestRule;
+import androidx.media3.test.session.common.TestHandler;
 import androidx.media3.test.session.common.TestUtils;
+import androidx.media3.test.utils.TestExoPlayerBuilder;
 import androidx.test.core.app.ApplicationProvider;
 import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.MediumTest;
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.Lists;
 import com.google.common.util.concurrent.ListenableFuture;
 import com.google.common.util.concurrent.SettableFuture;
 import java.util.ArrayList;
@@ -145,6 +151,100 @@ public class MediaSessionServiceTest {
     service.blockUntilAllControllersUnbind(TIMEOUT_MS);
   }
 
+  @Test
+  public void onCreate_mediaNotificationManagerController_correctSessionStateFromOnConnect()
+      throws Exception {
+    SessionCommand command1 = new SessionCommand("command1", Bundle.EMPTY);
+    SessionCommand command2 = new SessionCommand("command2", Bundle.EMPTY);
+    SessionCommand command3 = new SessionCommand("command3", Bundle.EMPTY);
+    CommandButton button1 =
+        new CommandButton.Builder()
+            .setDisplayName("button1")
+            .setIconResId(R.drawable.media3_notification_small_icon)
+            .setSessionCommand(command1)
+            .build();
+    CommandButton button2 =
+        new CommandButton.Builder()
+            .setDisplayName("button2")
+            .setIconResId(R.drawable.media3_notification_small_icon)
+            .setSessionCommand(command2)
+            .build();
+    CommandButton button3 =
+        new CommandButton.Builder()
+            .setDisplayName("button3")
+            .setIconResId(R.drawable.media3_notification_small_icon)
+            .setSessionCommand(command3)
+            .build();
+    Bundle testHints = new Bundle();
+    testHints.putString("test_key", "test_value");
+    List<ControllerInfo> controllerInfoList = new ArrayList<>();
+    CountDownLatch latch = new CountDownLatch(2);
+    TestHandler handler = new TestHandler(Looper.getMainLooper());
+    ExoPlayer player =
+        handler.postAndSync(
+            () -> {
+              ExoPlayer exoPlayer = new TestExoPlayerBuilder(context).build();
+              exoPlayer.setMediaItem(MediaItem.fromUri("asset:///media/mp4/sample.mp4"));
+              exoPlayer.prepare();
+              return exoPlayer;
+            });
+    MediaSession mediaSession =
+        new MediaSession.Builder(ApplicationProvider.getApplicationContext(), player)
+            .setCustomLayout(Lists.newArrayList(button1, button2))
+            .setCallback(
+                new MediaSession.Callback() {
+                  @Override
+                  public MediaSession.ConnectionResult onConnect(
+                      MediaSession session, ControllerInfo controller) {
+                    controllerInfoList.add(controller);
+                    if (session.isMediaNotificationController(controller)) {
+                      latch.countDown();
+                      return new MediaSession.ConnectionResult.AcceptedResultBuilder(session)
+                          .setAvailableSessionCommands(
+                              SessionCommands.EMPTY.buildUpon().add(command1).add(command3).build())
+                          .setAvailablePlayerCommands(Player.Commands.EMPTY)
+                          .setCustomLayout(ImmutableList.of(button1, button3))
+                          .build();
+                    }
+                    latch.countDown();
+                    return new MediaSession.ConnectionResult.AcceptedResultBuilder(session).build();
+                  }
+                })
+            .build();
+    TestServiceRegistry.getInstance().setOnGetSessionHandler(controllerInfo -> mediaSession);
+    MediaControllerCompat mediaControllerCompat =
+        new MediaControllerCompat(
+            ApplicationProvider.getApplicationContext(), mediaSession.getSessionCompat());
+    ImmutableList<CommandButton> initialCustomLayoutInControllerCompat =
+        MediaUtils.convertToCustomLayout(mediaControllerCompat.getPlaybackState());
+
+    // Start the service by creating a remote controller.
+    RemoteMediaController remoteController =
+        controllerTestRule.createRemoteController(token, /* waitForConnection= */ true, testHints);
+
+    assertThat(latch.await(TIMEOUT_MS, MILLISECONDS)).isTrue();
+    assertThat(
+            controllerInfoList
+                .get(0)
+                .getConnectionHints()
+                .getBoolean(
+                    MediaNotificationManager.KEY_MEDIA_NOTIFICATION_MANAGER,
+                    /* defaultValue= */ false))
+        .isTrue();
+    assertThat(TestUtils.equals(controllerInfoList.get(1).getConnectionHints(), testHints))
+        .isTrue();
+    assertThat(mediaControllerCompat.getPlaybackState().getActions())
+        .isEqualTo(PlaybackStateCompat.ACTION_SET_RATING);
+    assertThat(remoteController.getCustomLayout()).containsExactly(button1, button2).inOrder();
+    assertThat(initialCustomLayoutInControllerCompat).isEmpty();
+    assertThat(MediaUtils.convertToCustomLayout(mediaControllerCompat.getPlaybackState()))
+        .containsExactly(button1.copyWithIsEnabled(true), button3.copyWithIsEnabled(true))
+        .inOrder();
+    mediaSession.release();
+    ((MockMediaSessionService) TestServiceRegistry.getInstance().getServiceInstance())
+        .blockUntilAllControllersUnbind(TIMEOUT_MS);
+  }
+
   /**
    * Tests whether {@link MediaSessionService#onGetSession(ControllerInfo)} is called when
    * controller tries to connect, with the proper arguments.

libraries/test_session_current/src/main/java/androidx/media3/session/MediaSessionProviderService.java

@@ -16,6 +16,7 @@
 package androidx.media3.session;
 
 import static androidx.media3.common.Player.COMMAND_GET_TRACKS;
+import static androidx.media3.session.MediaSession.ConnectionResult.accept;
 import static androidx.media3.test.session.common.CommonConstants.ACTION_MEDIA3_SESSION;
 import static androidx.media3.test.session.common.CommonConstants.KEY_AUDIO_ATTRIBUTES;
 import static androidx.media3.test.session.common.CommonConstants.KEY_AVAILABLE_COMMANDS;
@@ -63,6 +64,7 @@ import static androidx.media3.test.session.common.MediaSessionConstants.TEST_CON
 import static androidx.media3.test.session.common.MediaSessionConstants.TEST_GET_CUSTOM_LAYOUT;
 import static androidx.media3.test.session.common.MediaSessionConstants.TEST_GET_SESSION_ACTIVITY;
 import static androidx.media3.test.session.common.MediaSessionConstants.TEST_IS_SESSION_COMMAND_AVAILABLE;
+import static androidx.media3.test.session.common.MediaSessionConstants.TEST_MEDIA_CONTROLLER_COMPAT_CALLBACK_WITH_MEDIA_SESSION_TEST;
 import static androidx.media3.test.session.common.MediaSessionConstants.TEST_ON_TRACKS_CHANGED_VIDEO_TO_AUDIO_TRANSITION;
 import static androidx.media3.test.session.common.MediaSessionConstants.TEST_ON_VIDEO_SIZE_CHANGED;
 import static androidx.media3.test.session.common.MediaSessionConstants.TEST_SET_SHOW_PLAY_BUTTON_IF_SUPPRESSED_TO_FALSE;
@@ -101,6 +103,7 @@ import androidx.media3.test.session.common.TestHandler;
 import androidx.media3.test.session.common.TestHandler.TestRunnable;
 import androidx.media3.test.session.common.TestUtils;
 import com.google.common.collect.ImmutableList;
+import com.google.common.util.concurrent.ListenableFuture;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
@@ -113,6 +116,8 @@ import java.util.concurrent.Callable;
  */
 public class MediaSessionProviderService extends Service {
 
+  public static final String KEY_ENABLE_FAKE_MEDIA_NOTIFICATION_MANAGER_CONTROLLER =
+      "key_enable_fake_media_notification_manager_controller";
   private static final String TAG = "MSProviderService";
 
   private Map<String, MediaSession> sessionMap = new HashMap<>();
@@ -164,15 +169,18 @@ public class MediaSessionProviderService extends Service {
 
     @Override
     public void create(String sessionId, Bundle tokenExtras) throws RemoteException {
+      if (tokenExtras == null) {
+        tokenExtras = Bundle.EMPTY;
+      }
+      boolean useFakeMediaNotificationManagerController =
+          tokenExtras.getBoolean(
+              KEY_ENABLE_FAKE_MEDIA_NOTIFICATION_MANAGER_CONTROLLER, /* defaultValue= */ false);
       MockPlayer mockPlayer =
           new MockPlayer.Builder().setApplicationLooper(handler.getLooper()).build();
       MediaSession.Builder builder =
           new MediaSession.Builder(MediaSessionProviderService.this, mockPlayer).setId(sessionId);
 
-      if (tokenExtras != null) {
       builder.setExtras(tokenExtras);
-      }
-
       switch (sessionId) {
         case TEST_GET_SESSION_ACTIVITY:
           {
@@ -194,7 +202,7 @@ public class MediaSessionProviderService extends Service {
                   @Override
                   public MediaSession.ConnectionResult onConnect(
                       MediaSession session, ControllerInfo controller) {
-                    return MediaSession.ConnectionResult.accept(
+                    return accept(
                         new SessionCommands.Builder()
                             .add(new SessionCommand("command1", Bundle.EMPTY))
                             .add(new SessionCommand("command2", Bundle.EMPTY))
@@ -216,8 +224,7 @@ public class MediaSessionProviderService extends Service {
                   @Override
                   public MediaSession.ConnectionResult onConnect(
                       MediaSession session, ControllerInfo controller) {
-                    return MediaSession.ConnectionResult.accept(
+                    return accept(availableSessionCommands, Player.Commands.EMPTY);
-                        availableSessionCommands, Player.Commands.EMPTY);
                   }
                 });
             break;
@@ -244,8 +251,7 @@ public class MediaSessionProviderService extends Service {
                   @Override
                   public MediaSession.ConnectionResult onConnect(
                       MediaSession session, ControllerInfo controller) {
-                    return MediaSession.ConnectionResult.accept(
+                    return accept(availableSessionCommands, Player.Commands.EMPTY);
-                        availableSessionCommands, Player.Commands.EMPTY);
                   }
                 });
             break;
@@ -272,8 +278,7 @@ public class MediaSessionProviderService extends Service {
                         .getBoolean(KEY_COMMAND_GET_TASKS_UNAVAILABLE, /* defaultValue= */ false)) {
                       commandBuilder.remove(COMMAND_GET_TRACKS);
                     }
-                    return MediaSession.ConnectionResult.accept(
+                    return accept(SessionCommands.EMPTY, commandBuilder.build());
-                        SessionCommands.EMPTY, commandBuilder.build());
                   }
                 });
             break;
@@ -290,6 +295,31 @@ public class MediaSessionProviderService extends Service {
             builder.setShowPlayButtonIfPlaybackIsSuppressed(false);
             break;
           }
+        case TEST_MEDIA_CONTROLLER_COMPAT_CALLBACK_WITH_MEDIA_SESSION_TEST:
+          {
+            builder.setCallback(
+                new MediaSession.Callback() {
+                  @Override
+                  public MediaSession.ConnectionResult onConnect(
+                      MediaSession session, ControllerInfo controller) {
+                    MediaSession.ConnectionResult connectionResult =
+                        MediaSession.Callback.super.onConnect(session, controller);
+                    SessionCommands availableSessionCommands =
+                        connectionResult.availableSessionCommands;
+                    if (session.isMediaNotificationController(controller)) {
+                      availableSessionCommands =
+                          connectionResult
+                              .availableSessionCommands
+                              .buildUpon()
+                              .add(new SessionCommand("command1", Bundle.EMPTY))
+                              .build();
+                    }
+                    return accept(
+                        availableSessionCommands, connectionResult.availablePlayerCommands);
+                  }
+                });
+            break;
+          }
         default: // fall out
       }
 
@@ -297,6 +327,22 @@ public class MediaSessionProviderService extends Service {
           () -> {
             MediaSession session = builder.build();
             session.setSessionPositionUpdateDelayMs(0L);
+            if (useFakeMediaNotificationManagerController) {
+              Bundle connectionHints = new Bundle();
+              connectionHints.putBoolean("androidx.media3.session.MediaNotificationManager", true);
+              //noinspection unused
+              ListenableFuture<MediaController> unusedFuture =
+                  new MediaController.Builder(getApplicationContext(), session.getToken())
+                      .setListener(
+                          new MediaController.Listener() {
+                            @Override
+                            public void onDisconnected(MediaController controller) {
+                              controller.release();
+                            }
+                          })
+                      .setConnectionHints(connectionHints)
+                      .buildAsync();
+            }
             sessionMap.put(sessionId, session);
           });
     }