diff --git a/libraries/ui/src/main/java/androidx/media3/ui/PlayerNotificationManager.java b/libraries/ui/src/main/java/androidx/media3/ui/PlayerNotificationManager.java
index 6821b54480..768274cbab 100644
--- a/libraries/ui/src/main/java/androidx/media3/ui/PlayerNotificationManager.java
+++ b/libraries/ui/src/main/java/androidx/media3/ui/PlayerNotificationManager.java
@@ -46,6 +46,7 @@ import android.content.Intent;
 import android.content.IntentFilter;
 import android.graphics.Bitmap;
 import android.graphics.Color;
+import android.media.session.MediaSession;
 import android.os.Handler;
 import android.os.Looper;
 import android.os.Message;
@@ -54,9 +55,10 @@ import androidx.annotation.DrawableRes;
 import androidx.annotation.IntDef;
 import androidx.annotation.IntRange;
 import androidx.annotation.Nullable;
+import androidx.annotation.RequiresApi;
+import androidx.core.app.NotificationBuilderWithBuilderAccessor;
 import androidx.core.app.NotificationCompat;
 import androidx.core.app.NotificationManagerCompat;
-import androidx.media.app.NotificationCompat.MediaStyle;
 import androidx.media3.common.C;
 import androidx.media3.common.Player;
 import androidx.media3.common.util.NotificationUtil;
@@ -174,7 +176,7 @@ import java.util.Map;
  * Builder#Builder(Context, int, String, MediaDescriptionAdapter)}.
  *
  * <p>Note: This class would require {@link android.Manifest.permission#POST_NOTIFICATIONS}
- * permission if used without a {@linkplain #setMediaSessionToken(MediaSessionCompat.Token) media
+ * permission if used without a {@linkplain #setMediaSessionToken(MediaSession.Token) media
  * session}.
  */
 @UnstableApi
@@ -707,7 +709,7 @@ public class PlayerNotificationManager {
   @Nullable private Player player;
   private boolean isNotificationStarted;
   private int currentNotificationTag;
-  @Nullable private MediaSessionCompat.Token mediaSessionToken;
+  @Nullable private MediaSession.Token mediaSessionToken;
   private boolean usePreviousAction;
   private boolean useNextAction;
   private boolean usePreviousActionInCompactView;
@@ -1008,11 +1010,26 @@ public class PlayerNotificationManager {
   }
 
   /**
-   * Sets the {@link MediaSessionCompat.Token}.
-   *
-   * @param token The {@link MediaSessionCompat.Token}.
+   * @deprecated Use {@link #setMediaSessionToken(MediaSession.Token)} and pass in {@code
+   *     (MediaSession.Token) compatToken.getToken()}.
    */
-  public final void setMediaSessionToken(MediaSessionCompat.Token token) {
+  @Deprecated
+  public final void setMediaSessionToken(MediaSessionCompat.Token compatToken) {
+    if (Util.SDK_INT >= 21) {
+      setMediaSessionToken((MediaSession.Token) compatToken.getToken());
+    }
+  }
+
+  /**
+   * Sets the {@link MediaSession.Token}.
+   *
+   * <p>When using {@code MediaSessionCompat}, this token can be obtained with {@code
+   * (MediaSession.Token) compatToken.getToken()}.
+   *
+   * @param token The {@link MediaSession.Token}.
+   */
+  @RequiresApi(21)
+  public final void setMediaSessionToken(MediaSession.Token token) {
     if (!Util.areEqual(this.mediaSessionToken, token)) {
       mediaSessionToken = token;
       invalidate();
@@ -1273,15 +1290,18 @@ public class PlayerNotificationManager {
       }
     }
 
-    MediaStyle mediaStyle = new MediaStyle();
-    if (mediaSessionToken != null) {
-      mediaStyle.setMediaSession(mediaSessionToken);
+    int[] actionIndicesForCompactView = getActionIndicesForCompactView(actionNames, player);
+    if (Util.SDK_INT >= 21) {
+      builder.setStyle(new MediaStyle(mediaSessionToken, actionIndicesForCompactView));
+    } else {
+      // TODO: b/333355694 - Remove dependency on androidx.media once this logic is gone.
+      androidx.media.app.NotificationCompat.MediaStyle mediaStyle =
+          new androidx.media.app.NotificationCompat.MediaStyle();
+      mediaStyle.setShowActionsInCompactView(actionIndicesForCompactView);
+      mediaStyle.setShowCancelButton(!ongoing);
+      mediaStyle.setCancelButtonIntent(dismissPendingIntent);
+      builder.setStyle(mediaStyle);
     }
-    mediaStyle.setShowActionsInCompactView(getActionIndicesForCompactView(actionNames, player));
-    // Configure dismiss action prior to API 21 ('x' button).
-    mediaStyle.setShowCancelButton(!ongoing);
-    mediaStyle.setCancelButtonIntent(dismissPendingIntent);
-    builder.setStyle(mediaStyle);
 
     // Set intent which is sent if the user selects 'clear all'
     builder.setDeleteIntent(dismissPendingIntent);
@@ -1562,7 +1582,6 @@ public class PlayerNotificationManager {
 
   private class NotificationBroadcastReceiver extends BroadcastReceiver {
 
-    @SuppressWarnings("deprecation")
     @Override
     public void onReceive(Context context, Intent intent) {
       Player player = PlayerNotificationManager.this.player;
@@ -1608,4 +1627,26 @@ public class PlayerNotificationManager {
       }
     }
   }
+
+  @RequiresApi(21)
+  private static final class MediaStyle extends androidx.core.app.NotificationCompat.Style {
+
+    private final int[] actionsToShowInCompact;
+    @Nullable private final MediaSession.Token token;
+
+    public MediaStyle(@Nullable MediaSession.Token token, int[] actionsToShowInCompact) {
+      this.token = token;
+      this.actionsToShowInCompact = actionsToShowInCompact;
+    }
+
+    @Override
+    public void apply(NotificationBuilderWithBuilderAccessor builder) {
+      Notification.MediaStyle style = new Notification.MediaStyle();
+      style.setShowActionsInCompactView(actionsToShowInCompact);
+      if (token != null) {
+        style.setMediaSession(token);
+      }
+      builder.getBuilder().setStyle(style);
+    }
+  }
 }