Mark media notification controller and filter command buttons

The `MediaNotificationManager` registers an internal controller
to each session. This change marks this controller through its
connection hints and provides an API for apps to hide
implementation details of the marking.

Issue: androidx/media#389
PiperOrigin-RevId: 549712768

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

@@ -298,6 +298,16 @@ public class DefaultMediaNotificationProvider implements MediaNotification.Provi
       Callback onNotificationChangedCallback) {
     ensureNotificationChannel();
 
+    ImmutableList.Builder<CommandButton> customLayoutWithEnabledCommandButtonsOnly =
+        new ImmutableList.Builder<>();
+    for (int i = 0; i < customLayout.size(); i++) {
+      CommandButton button = customLayout.get(i);
+      if (button.sessionCommand != null
+          && button.sessionCommand.commandCode == SessionCommand.COMMAND_CODE_CUSTOM
+          && button.isEnabled) {
+        customLayoutWithEnabledCommandButtonsOnly.add(customLayout.get(i));
+      }
+    }
     Player player = mediaSession.getPlayer();
     NotificationCompat.Builder builder = new NotificationCompat.Builder(context, channelId);
     int notificationId = notificationIdProvider.getNotificationId(mediaSession);
@@ -309,7 +319,7 @@ public class DefaultMediaNotificationProvider implements MediaNotification.Provi
             getMediaButtons(
                 mediaSession,
                 player.getAvailableCommands(),
-                customLayout,
+                customLayoutWithEnabledCommandButtonsOnly.build(),
                 /* showPauseButton= */ player.getPlayWhenReady()
                     && player.getPlaybackState() != STATE_ENDED),
             builder,

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

@@ -61,7 +61,6 @@ import java.util.concurrent.Future;
 /* package */ class MediaLibrarySessionImpl extends MediaSessionImpl {
 
   private static final String RECENT_LIBRARY_ROOT_MEDIA_ID = "androidx.media3.session.recent.root";
-  private static final String SYSTEM_UI_PACKAGE_NAME = "com.android.systemui";
 
   private final MediaLibrarySession instance;
   private final MediaLibrarySession.Callback callback;
@@ -145,9 +144,7 @@ import java.util.concurrent.Future;
 
   public ListenableFuture<LibraryResult<MediaItem>> onGetLibraryRootOnHandler(
       ControllerInfo browser, @Nullable LibraryParams params) {
-    if (params != null
+    if (params != null && params.isRecent && isSystemUiController(browser)) {
-        && params.isRecent
-        && Objects.equals(browser.getPackageName(), SYSTEM_UI_PACKAGE_NAME)) {
       // Advertise support for playback resumption, if enabled.
       return !canResumePlaybackOnStart()
           ? Futures.immediateFuture(LibraryResult.ofError(RESULT_ERROR_NOT_SUPPORTED))

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

@@ -17,7 +17,6 @@ package androidx.media3.session;
 
 import static android.app.Service.STOP_FOREGROUND_DETACH;
 import static android.app.Service.STOP_FOREGROUND_REMOVE;
-import static androidx.media3.common.util.Assertions.checkStateNotNull;
 
 import android.annotation.SuppressLint;
 import android.app.Notification;
@@ -56,6 +55,8 @@ import java.util.concurrent.TimeoutException;
  */
 /* package */ final class MediaNotificationManager {
 
+  /* package */ static final String KEY_MEDIA_NOTIFICATION_MANAGER =
+      "androidx.media3.session.MediaNotificationManager";
   private static final String TAG = "MediaNtfMng";
 
   private final MediaSessionService mediaSessionService;
@@ -65,7 +66,6 @@ import java.util.concurrent.TimeoutException;
   private final Executor mainExecutor;
   private final Intent startSelfIntent;
   private final Map<MediaSession, ListenableFuture<MediaController>> controllerMap;
-  private final Map<MediaSession, ImmutableList<CommandButton>> customLayoutMap;
 
   private int totalNotificationCount;
   @Nullable private MediaNotification mediaNotification;
@@ -83,7 +83,6 @@ import java.util.concurrent.TimeoutException;
     mainExecutor = (runnable) -> Util.postOrRun(mainHandler, runnable);
     startSelfIntent = new Intent(mediaSessionService, mediaSessionService.getClass());
     controllerMap = new HashMap<>();
-    customLayoutMap = new HashMap<>();
     startedInForeground = false;
   }
 
@@ -91,11 +90,12 @@ import java.util.concurrent.TimeoutException;
     if (controllerMap.containsKey(session)) {
       return;
     }
-    customLayoutMap.put(session, ImmutableList.of());
+    MediaControllerListener listener = new MediaControllerListener(mediaSessionService, session);
-    MediaControllerListener listener =
+    Bundle connectionHints = new Bundle();
-        new MediaControllerListener(mediaSessionService, session, customLayoutMap);
+    connectionHints.putBoolean(KEY_MEDIA_NOTIFICATION_MANAGER, true);
     ListenableFuture<MediaController> controllerFuture =
         new MediaController.Builder(mediaSessionService, session.getToken())
+            .setConnectionHints(connectionHints)
             .setListener(listener)
             .setApplicationLooper(Looper.getMainLooper())
             .buildAsync();
@@ -118,7 +118,6 @@ import java.util.concurrent.TimeoutException;
   }
 
   public void removeSession(MediaSession session) {
-    customLayoutMap.remove(session);
     @Nullable ListenableFuture<MediaController> controllerFuture = controllerMap.remove(session);
     if (controllerFuture != null) {
       MediaController.releaseFuture(controllerFuture);
@@ -154,7 +153,19 @@ import java.util.concurrent.TimeoutException;
     }
 
     int notificationSequence = ++totalNotificationCount;
-    ImmutableList<CommandButton> customLayout = checkStateNotNull(customLayoutMap.get(session));
+    MediaController mediaNotificationController = null;
+    @Nullable ListenableFuture<MediaController> controllerFuture = controllerMap.get(session);
+    if (controllerFuture != null && controllerFuture.isDone()) {
+      try {
+        mediaNotificationController = Futures.getDone(controllerFuture);
+      } catch (ExecutionException e) {
+        // Ignore.
+      }
+    }
+    ImmutableList<CommandButton> customLayout =
+        mediaNotificationController != null
+            ? mediaNotificationController.getCustomLayout()
+            : ImmutableList.of();
     MediaNotification.Provider.Callback callback =
         notification ->
             mainExecutor.execute(
@@ -297,15 +308,10 @@ import java.util.concurrent.TimeoutException;
       implements MediaController.Listener, Player.Listener {
     private final MediaSessionService mediaSessionService;
     private final MediaSession session;
-    private final Map<MediaSession, ImmutableList<CommandButton>> customLayoutMap;
 
-    public MediaControllerListener(
+    public MediaControllerListener(MediaSessionService mediaSessionService, MediaSession session) {
-        MediaSessionService mediaSessionService,
-        MediaSession session,
-        Map<MediaSession, ImmutableList<CommandButton>> customLayoutMap) {
       this.mediaSessionService = mediaSessionService;
       this.session = session;
-      this.customLayoutMap = customLayoutMap;
     }
 
     public void onConnected(boolean shouldShowNotification) {
@@ -316,12 +322,16 @@ import java.util.concurrent.TimeoutException;
     }
 
     @Override
-    public ListenableFuture<SessionResult> onSetCustomLayout(
+    public void onCustomLayoutChanged(MediaController controller, List<CommandButton> layout) {
-        MediaController controller, List<CommandButton> layout) {
+      mediaSessionService.onUpdateNotificationInternal(
-      customLayoutMap.put(session, ImmutableList.copyOf(layout));
+          session, /* startInForegroundWhenPaused= */ false);
+    }
+
+    @Override
+    public void onAvailableSessionCommandsChanged(
+        MediaController controller, SessionCommands commands) {
       mediaSessionService.onUpdateNotificationInternal(
           session, /* startInForegroundWhenPaused= */ false);
-      return Futures.immediateFuture(new SessionResult(SessionResult.RESULT_SUCCESS));
     }
 
     @Override

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

@@ -751,6 +751,35 @@ public class MediaSession {
     return impl.getControllerForCurrentRequest();
   }
 
+  /**
+   * 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
+   * 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.
+   */
+  @UnstableApi
+  public boolean isMediaNotificationController(ControllerInfo controllerInfo) {
+    return impl.isMediaNotificationController(controllerInfo);
+  }
+
+  /**
+   * Returns the {@link ControllerInfo} of the media notification controller.
+   *
+   * <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
+   * notification.
+   */
+  @UnstableApi
+  @Nullable
+  public ControllerInfo getMediaNotificationControllerInfo() {
+    return impl.getMediaNotificationControllerInfo();
+  }
+
   /**
    * Sets the custom layout for the given Media3 controller.
    *
@@ -775,11 +804,12 @@ public class MediaSession {
    * @param controller The controller for which to set the custom layout.
    * @param layout The ordered list of {@linkplain CommandButton command buttons}.
    */
+  @CanIgnoreReturnValue
   public final ListenableFuture<SessionResult> setCustomLayout(
       ControllerInfo controller, List<CommandButton> layout) {
     checkNotNull(controller, "controller must not be null");
     checkNotNull(layout, "layout must not be null");
-    return impl.setCustomLayout(controller, layout);
+    return impl.setCustomLayout(controller, ImmutableList.copyOf(layout));
   }
 
   /**
@@ -811,7 +841,7 @@ public class MediaSession {
    */
   public final void setCustomLayout(List<CommandButton> layout) {
     checkNotNull(layout, "layout must not be null");
-    impl.setCustomLayout(layout);
+    impl.setCustomLayout(ImmutableList.copyOf(layout));
   }
 
   /**

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

@@ -86,6 +86,7 @@ import org.checkerframework.checker.nullness.qual.MonotonicNonNull;
 
 /* package */ class MediaSessionImpl {
 
+  private static final String SYSTEM_UI_PACKAGE_NAME = "com.android.systemui";
   private static final String WRONG_THREAD_ERROR_MESSAGE =
       "Player callback method is called from a wrong thread. "
           + "See javadoc of MediaSession for details.";
@@ -305,6 +306,68 @@ import org.checkerframework.checker.nullness.qual.MonotonicNonNull;
         || sessionLegacyStub.getConnectedControllersManager().isConnected(controller);
   }
 
+  /**
+   * Returns whether the given {@link ControllerInfo} belongs to the the System UI controller.
+   *
+   * @param controllerInfo The controller info.
+   * @return Whether the controller info belongs to the System UI controller.
+   */
+  protected boolean isSystemUiController(@Nullable MediaSession.ControllerInfo controllerInfo) {
+    return controllerInfo != null
+        && controllerInfo.getControllerVersion() == ControllerInfo.LEGACY_CONTROLLER_VERSION
+        && Objects.equals(controllerInfo.getPackageName(), SYSTEM_UI_PACKAGE_NAME);
+  }
+
+  /**
+   * Returns whether the given {@link ControllerInfo} belongs to the media notification controller.
+   *
+   * @param controllerInfo The controller info.
+   * @return Whether the given controller info belongs to the media notification controller.
+   */
+  public boolean isMediaNotificationController(MediaSession.ControllerInfo controllerInfo) {
+    return Objects.equals(controllerInfo.getPackageName(), context.getPackageName())
+        && controllerInfo.getControllerVersion()
+            != ControllerInfo.LEGACY_CONTROLLER_INTERFACE_VERSION
+        && controllerInfo
+            .getConnectionHints()
+            .getBoolean(
+                MediaNotificationManager.KEY_MEDIA_NOTIFICATION_MANAGER, /* defaultValue= */ false);
+  }
+
+  /**
+   * Returns the {@link ControllerInfo} of the system UI notification controller, or {@code null} if
+   * the System UI controller is not connected.
+   */
+  @Nullable
+  protected ControllerInfo getSystemUiControllerInfo() {
+    ImmutableList<ControllerInfo> connectedControllers =
+        sessionLegacyStub.getConnectedControllersManager().getConnectedControllers();
+    for (int i = 0; i < connectedControllers.size(); i++) {
+      ControllerInfo controllerInfo = connectedControllers.get(i);
+      if (isSystemUiController(controllerInfo)) {
+        return controllerInfo;
+      }
+    }
+    return null;
+  }
+
+  /**
+   * Returns the {@link ControllerInfo} of the media notification controller, or {@code null} if the
+   * media notification controller is not connected.
+   */
+  @Nullable
+  public ControllerInfo getMediaNotificationControllerInfo() {
+    ImmutableList<ControllerInfo> connectedControllers =
+        sessionStub.getConnectedControllersManager().getConnectedControllers();
+    for (int i = 0; i < connectedControllers.size(); i++) {
+      ControllerInfo controllerInfo = connectedControllers.get(i);
+      if (isMediaNotificationController(controllerInfo)) {
+        return controllerInfo;
+      }
+    }
+    return null;
+  }
+
   public ListenableFuture<SessionResult> setCustomLayout(
       ControllerInfo controller, List<CommandButton> layout) {
     return dispatchRemoteControllerTask(

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

@@ -49,6 +49,7 @@ import androidx.test.core.app.ApplicationProvider;
 import androidx.test.ext.junit.runners.AndroidJUnit4;
 import com.google.common.collect.ImmutableList;
 import com.google.common.util.concurrent.SettableFuture;
+import java.util.ArrayList;
 import java.util.List;
 import org.junit.Before;
 import org.junit.Test;
@@ -606,6 +607,62 @@ public class DefaultMediaNotificationProviderTest {
     verifyNoInteractions(mockOnNotificationChangedCallback1);
   }
 
+  @Test
+  public void createNotification_invalidButtons_enabledSessionCommandsOnlyForGetMediaButtons() {
+    DefaultActionFactory defaultActionFactory =
+        new DefaultActionFactory(Robolectric.setupService(TestService.class));
+    List<CommandButton> filteredEnabledLayout = new ArrayList<>();
+    DefaultMediaNotificationProvider defaultMediaNotificationProvider =
+        new DefaultMediaNotificationProvider(ApplicationProvider.getApplicationContext()) {
+          @Override
+          protected ImmutableList<CommandButton> getMediaButtons(
+              MediaSession session,
+              Commands playerCommands,
+              ImmutableList<CommandButton> customLayout,
+              boolean showPauseButton) {
+            filteredEnabledLayout.addAll(customLayout);
+            return super.getMediaButtons(session, playerCommands, customLayout, showPauseButton);
+          }
+        };
+    MediaSession mediaSession =
+        new MediaSession.Builder(
+                ApplicationProvider.getApplicationContext(),
+                new TestExoPlayerBuilder(context).build())
+            .build();
+    CommandButton button1 =
+        new CommandButton.Builder()
+            .setDisplayName("button1")
+            .setIconResId(R.drawable.media3_notification_small_icon)
+            .setPlayerCommand(Player.COMMAND_SEEK_TO_PREVIOUS)
+            .build();
+    CommandButton button2 =
+        new CommandButton.Builder()
+            .setDisplayName("button2")
+            .setIconResId(R.drawable.media3_notification_small_icon)
+            .setSessionCommand(new SessionCommand("command2", Bundle.EMPTY))
+            .build()
+            .copyWithIsEnabled(true);
+    CommandButton button3 =
+        new CommandButton.Builder()
+            .setDisplayName("button3")
+            .setIconResId(R.drawable.media3_notification_small_icon)
+            .setPlayerCommand(Player.COMMAND_PLAY_PAUSE)
+            .build()
+            .copyWithIsEnabled(true);
+
+    defaultMediaNotificationProvider.createNotification(
+        mediaSession,
+        /* customLayout= */ ImmutableList.of(button1, button2, button3),
+        defaultActionFactory,
+        notification -> {
+          /* Do nothing. */
+        });
+
+    assertThat(filteredEnabledLayout).containsExactly(button2);
+    mediaSession.getPlayer().release();
+    mediaSession.release();
+  }
+
   @Test
   public void provider_idsNotSpecified_usesDefaultIds() {
     Context context = ApplicationProvider.getApplicationContext();

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

@@ -16,11 +16,11 @@
 package androidx.media3.session;
 
 import static androidx.media3.test.utils.robolectric.RobolectricUtil.runMainLooperUntil;
-import static com.google.common.truth.Truth8.assertThat;
+import static com.google.common.truth.Truth.assertThat;
-import static java.util.Arrays.stream;
 
 import android.app.NotificationManager;
 import android.content.Context;
+import android.os.Bundle;
 import android.os.Handler;
 import android.os.HandlerThread;
 import android.service.notification.StatusBarNotification;
@@ -30,6 +30,10 @@ import androidx.media3.exoplayer.ExoPlayer;
 import androidx.media3.test.utils.TestExoPlayerBuilder;
 import androidx.test.core.app.ApplicationProvider;
 import androidx.test.ext.junit.runners.AndroidJUnit4;
+import com.google.common.collect.ImmutableList;
+import java.util.concurrent.TimeoutException;
+import org.junit.After;
+import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.robolectric.Robolectric;
@@ -39,14 +43,29 @@ import org.robolectric.shadows.ShadowLooper;
 @RunWith(AndroidJUnit4.class)
 public class MediaSessionServiceTest {
 
+  private Context context;
+  private NotificationManager notificationManager;
+  private ServiceController<TestService> serviceController;
+
+  @Before
+  public void setUp() {
+    context = ApplicationProvider.getApplicationContext();
+    notificationManager =
+        (NotificationManager) context.getSystemService(Context.NOTIFICATION_SERVICE);
+    serviceController = Robolectric.buildService(TestService.class);
+  }
+
+  @After
+  public void tearDown() {
+    serviceController.destroy();
+  }
+
   @Test
   public void service_multipleSessionsOnMainThread_createsNotificationForEachSession() {
-    Context context = ApplicationProvider.getApplicationContext();
     ExoPlayer player1 = new TestExoPlayerBuilder(context).build();
     ExoPlayer player2 = new TestExoPlayerBuilder(context).build();
     MediaSession session1 = new MediaSession.Builder(context, player1).setId("1").build();
     MediaSession session2 = new MediaSession.Builder(context, player2).setId("2").build();
-    ServiceController<TestService> serviceController = Robolectric.buildService(TestService.class);
     TestService service = serviceController.create().get();
     service.setMediaNotificationProvider(
         new DefaultMediaNotificationProvider(
@@ -66,13 +85,9 @@ public class MediaSessionServiceTest {
     player2.play();
     ShadowLooper.idleMainLooper();
 
-    NotificationManager notificationService =
+    assertThat(getStatusBarNotification(2001)).isNotNull();
-        (NotificationManager) context.getSystemService(Context.NOTIFICATION_SERVICE);
+    assertThat(getStatusBarNotification(2002)).isNotNull();
-    assertThat(
-            stream(notificationService.getActiveNotifications()).map(StatusBarNotification::getId))
-        .containsExactly(2001, 2002);
 
-    serviceController.destroy();
     session1.release();
     session2.release();
     player1.release();
@@ -82,7 +97,6 @@ public class MediaSessionServiceTest {
   @Test
   public void service_multipleSessionsOnDifferentThreads_createsNotificationForEachSession()
       throws Exception {
-    Context context = ApplicationProvider.getApplicationContext();
     HandlerThread thread1 = new HandlerThread("player1");
     HandlerThread thread2 = new HandlerThread("player2");
     thread1.start();
@@ -91,7 +105,6 @@ public class MediaSessionServiceTest {
     ExoPlayer player2 = new TestExoPlayerBuilder(context).setLooper(thread2.getLooper()).build();
     MediaSession session1 = new MediaSession.Builder(context, player1).setId("1").build();
     MediaSession session2 = new MediaSession.Builder(context, player2).setId("2").build();
-    ServiceController<TestService> serviceController = Robolectric.buildService(TestService.class);
     TestService service = serviceController.create().get();
     service.setMediaNotificationProvider(
         new DefaultMediaNotificationProvider(
@@ -99,8 +112,6 @@ public class MediaSessionServiceTest {
             session -> 2000 + Integer.parseInt(session.getId()),
             DefaultMediaNotificationProvider.DEFAULT_CHANNEL_ID,
             DefaultMediaNotificationProvider.DEFAULT_CHANNEL_NAME_RESOURCE_ID));
-    NotificationManager notificationService =
-        (NotificationManager) context.getSystemService(Context.NOTIFICATION_SERVICE);
 
     service.addSession(session1);
     service.addSession(session2);
@@ -119,13 +130,11 @@ public class MediaSessionServiceTest {
               player2.prepare();
               player2.play();
             });
-    runMainLooperUntil(() -> notificationService.getActiveNotifications().length == 2);
+    runMainLooperUntil(() -> notificationManager.getActiveNotifications().length == 2);
 
-    assertThat(
+    assertThat(getStatusBarNotification(2001)).isNotNull();
-            stream(notificationService.getActiveNotifications()).map(StatusBarNotification::getId))
+    assertThat(getStatusBarNotification(2002)).isNotNull();
-        .containsExactly(2001, 2002);
 
-    serviceController.destroy();
     session1.release();
     session2.release();
     new Handler(thread1.getLooper()).post(player1::release);
@@ -134,6 +143,192 @@ public class MediaSessionServiceTest {
     thread2.quit();
   }
 
+  @Test
+  public void mediaNotificationController_setCustomLayout_correctNotificationActions()
+      throws TimeoutException {
+    SessionCommand command1 = new SessionCommand("command1", Bundle.EMPTY);
+    SessionCommand command2 = new SessionCommand("command2", Bundle.EMPTY);
+    CommandButton button1 =
+        new CommandButton.Builder()
+            .setDisplayName("customAction1")
+            .setIconResId(R.drawable.media3_notification_small_icon)
+            .setSessionCommand(command1)
+            .build();
+    CommandButton button2 =
+        new CommandButton.Builder()
+            .setDisplayName("customAction2")
+            .setIconResId(R.drawable.media3_notification_small_icon)
+            .setSessionCommand(command2)
+            .build();
+    ExoPlayer player = new TestExoPlayerBuilder(context).build();
+    MediaSession session =
+        new MediaSession.Builder(context, player)
+            .setCustomLayout(ImmutableList.of(button1, button2))
+            .setCallback(
+                new MediaSession.Callback() {
+                  @Override
+                  public MediaSession.ConnectionResult onConnect(
+                      MediaSession session, MediaSession.ControllerInfo controller) {
+                    if (session.isMediaNotificationController(controller)) {
+                      return new MediaSession.ConnectionResult.AcceptedResultBuilder(session)
+                          .setAvailableSessionCommands(
+                              MediaSession.ConnectionResult.DEFAULT_SESSION_COMMANDS
+                                  .buildUpon()
+                                  .add(command1)
+                                  .add(command2)
+                                  .build())
+                          .build();
+                    }
+                    return new MediaSession.ConnectionResult.AcceptedResultBuilder(session).build();
+                  }
+                })
+            .build();
+    TestService service = serviceController.create().get();
+    service.setMediaNotificationProvider(
+        new DefaultMediaNotificationProvider(
+            service,
+            mediaSession -> 2000,
+            DefaultMediaNotificationProvider.DEFAULT_CHANNEL_ID,
+            DefaultMediaNotificationProvider.DEFAULT_CHANNEL_NAME_RESOURCE_ID));
+    service.addSession(session);
+    // Play media to create a notification.
+    player.setMediaItems(
+        ImmutableList.of(
+            MediaItem.fromUri("asset:///media/mp4/sample.mp4"),
+            MediaItem.fromUri("asset:///media/mp4/sample.mp4")));
+    player.prepare();
+    player.play();
+    runMainLooperUntil(() -> notificationManager.getActiveNotifications().length == 1);
+
+    StatusBarNotification mediaNotification = getStatusBarNotification(2000);
+
+    assertThat(mediaNotification.getNotification().actions).hasLength(5);
+    assertThat(mediaNotification.getNotification().actions[0].title.toString())
+        .isEqualTo("Seek to previous item");
+    assertThat(mediaNotification.getNotification().actions[1].title.toString()).isEqualTo("Pause");
+    assertThat(mediaNotification.getNotification().actions[2].title.toString())
+        .isEqualTo("Seek to next item");
+    assertThat(mediaNotification.getNotification().actions[3].title.toString())
+        .isEqualTo("customAction1");
+    assertThat(mediaNotification.getNotification().actions[4].title.toString())
+        .isEqualTo("customAction2");
+
+    player.pause();
+    session.setCustomLayout(
+        session.getMediaNotificationControllerInfo(), ImmutableList.of(button2));
+    ShadowLooper.idleMainLooper();
+    mediaNotification = getStatusBarNotification(2000);
+
+    assertThat(mediaNotification.getNotification().actions).hasLength(4);
+    assertThat(mediaNotification.getNotification().actions[0].title.toString())
+        .isEqualTo("Seek to previous item");
+    assertThat(mediaNotification.getNotification().actions[1].title.toString()).isEqualTo("Play");
+    assertThat(mediaNotification.getNotification().actions[2].title.toString())
+        .isEqualTo("Seek to next item");
+    assertThat(mediaNotification.getNotification().actions[3].title.toString())
+        .isEqualTo("customAction2");
+    session.release();
+    player.release();
+  }
+
+  @Test
+  public void mediaNotificationController_setAvailableCommands_correctNotificationActions()
+      throws TimeoutException {
+    SessionCommand command1 = new SessionCommand("command1", Bundle.EMPTY);
+    SessionCommand command2 = new SessionCommand("command2", Bundle.EMPTY);
+    CommandButton button1 =
+        new CommandButton.Builder()
+            .setDisplayName("customAction1")
+            .setIconResId(R.drawable.media3_notification_small_icon)
+            .setSessionCommand(command1)
+            .build();
+    CommandButton button2 =
+        new CommandButton.Builder()
+            .setDisplayName("customAction2")
+            .setIconResId(R.drawable.media3_notification_small_icon)
+            .setSessionCommand(command2)
+            .build();
+    Context context = ApplicationProvider.getApplicationContext();
+    ExoPlayer player = new TestExoPlayerBuilder(context).build();
+    MediaSession session =
+        new MediaSession.Builder(context, player)
+            .setId("1")
+            .setCustomLayout(ImmutableList.of(button1, button2))
+            .setCallback(
+                new MediaSession.Callback() {
+                  @Override
+                  public MediaSession.ConnectionResult onConnect(
+                      MediaSession session, MediaSession.ControllerInfo controller) {
+                    if (session.isMediaNotificationController(controller)) {
+                      return new MediaSession.ConnectionResult.AcceptedResultBuilder(session)
+                          .setAvailableSessionCommands(
+                              MediaSession.ConnectionResult.DEFAULT_SESSION_COMMANDS
+                                  .buildUpon()
+                                  .add(command2)
+                                  .build())
+                          .build();
+                    }
+                    return new MediaSession.ConnectionResult.AcceptedResultBuilder(session).build();
+                  }
+                })
+            .build();
+    TestService service = serviceController.create().get();
+    service.setMediaNotificationProvider(
+        new DefaultMediaNotificationProvider(
+            service,
+            mediaSession -> 2000,
+            DefaultMediaNotificationProvider.DEFAULT_CHANNEL_ID,
+            DefaultMediaNotificationProvider.DEFAULT_CHANNEL_NAME_RESOURCE_ID));
+    service.addSession(session);
+    // Start the players so that we also create notifications for them.
+    player.setMediaItem(MediaItem.fromUri("asset:///media/mp4/sample.mp4"));
+    player.prepare();
+    player.play();
+    runMainLooperUntil(() -> notificationManager.getActiveNotifications().length == 1);
+
+    StatusBarNotification mediaNotification = getStatusBarNotification(2000);
+
+    assertThat(mediaNotification.getNotification().actions[0].title.toString())
+        .isEqualTo("Seek to previous item");
+    assertThat(mediaNotification.getNotification().actions[1].title.toString()).isEqualTo("Pause");
+    assertThat(mediaNotification.getNotification().actions[2].title.toString())
+        .isEqualTo("customAction2");
+
+    player.pause();
+    session.setAvailableCommands(
+        session.getMediaNotificationControllerInfo(),
+        MediaSession.ConnectionResult.DEFAULT_SESSION_COMMANDS
+            .buildUpon()
+            .add(command1)
+            .add(command2)
+            .build(),
+        MediaSession.ConnectionResult.DEFAULT_PLAYER_COMMANDS);
+    ShadowLooper.idleMainLooper();
+    mediaNotification = getStatusBarNotification(2000);
+
+    assertThat(mediaNotification.getNotification().actions).hasLength(4);
+    assertThat(mediaNotification.getNotification().actions[0].title.toString())
+        .isEqualTo("Seek to previous item");
+    assertThat(mediaNotification.getNotification().actions[1].title.toString()).isEqualTo("Play");
+    assertThat(mediaNotification.getNotification().actions[2].title.toString())
+        .isEqualTo("customAction1");
+    assertThat(mediaNotification.getNotification().actions[3].title.toString())
+        .isEqualTo("customAction2");
+
+    session.release();
+    player.release();
+  }
+
+  @Nullable
+  private StatusBarNotification getStatusBarNotification(int notificationId) {
+    for (StatusBarNotification notification : notificationManager.getActiveNotifications()) {
+      if (notification.getId() == notificationId) {
+        return notification;
+      }
+    }
+    return null;
+  }
+
   private static final class TestService extends MediaSessionService {
     @Nullable
     @Override