Do not assume a valid queue in 3rd party sessions

This change fixes an issue that can be reproduced when
a controller `onConnect` creates a `QueueTimeline` out
of the state of a legacy session and then `prepare` is called.

`activeQueueItemId`, `metadata` and the `queue` of the legacy
session are used when a `QueueTimeline` is created. The change
adds unit tests to cover the different combinatoric cases these
properties being set or unset.

PiperOrigin-RevId: 505731288

RELEASENOTES.md

@@ -60,6 +60,8 @@
         onto Player ([#156](https://github.com/androidx/media/issues/156)).
     *   Avoid double tap detection for non-Bluetooth media button events
         ([#233](https://github.com/androidx/media/issues/233)).
+    *   Make `QueueTimeline` more robust in case of a shady legacy session state
+        ([#241](https://github.com/androidx/media/issues/241)).
 *   Metadata:
     *   Parse multiple null-separated values from ID3 frames, as permitted by
         ID3 v2.4.

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

@@ -1829,6 +1829,8 @@ import org.checkerframework.checker.nullness.compatqual.NullableType;
                   + " MediaItem.");
           MediaItem fakeMediaItem =
               MediaUtils.convertToMediaItem(newLegacyPlayerInfo.mediaMetadataCompat, ratingType);
+          // Ad a tag to make sure the fake media item can't have an equal instance by accident.
+          fakeMediaItem = fakeMediaItem.buildUpon().setTag(new Object()).build();
           currentTimeline = currentTimeline.copyWithFakeMediaItem(fakeMediaItem);
           currentMediaItemIndex = currentTimeline.getWindowCount() - 1;
         } else {
@@ -1843,7 +1845,7 @@ import org.checkerframework.checker.nullness.compatqual.NullableType;
         if (hasMediaMetadataCompat) {
           MediaItem mediaItem =
               MediaUtils.convertToMediaItem(
-                  currentTimeline.getMediaItemAt(currentMediaItemIndex).mediaId,
+                  checkNotNull(currentTimeline.getMediaItemAt(currentMediaItemIndex)).mediaId,
                   newLegacyPlayerInfo.mediaMetadataCompat,
                   ratingType);
           currentTimeline =
@@ -2000,7 +2002,7 @@ import org.checkerframework.checker.nullness.compatqual.NullableType;
       MediaItem oldCurrentMediaItem =
           checkStateNotNull(oldControllerInfo.playerInfo.getCurrentMediaItem());
       int oldCurrentMediaItemIndexInNewTimeline =
-          ((QueueTimeline) newControllerInfo.playerInfo.timeline).findIndexOf(oldCurrentMediaItem);
+          ((QueueTimeline) newControllerInfo.playerInfo.timeline).indexOf(oldCurrentMediaItem);
       if (oldCurrentMediaItemIndexInNewTimeline == C.INDEX_UNSET) {
         // Old current item is removed.
         discontinuityReason = Player.DISCONTINUITY_REASON_REMOVE;

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

@@ -15,6 +15,10 @@
  */
 package androidx.media3.session;
 
+import static androidx.media3.common.util.Assertions.checkArgument;
+import static androidx.media3.common.util.Assertions.checkNotNull;
+
+import android.support.v4.media.MediaMetadataCompat;
 import android.support.v4.media.session.MediaSessionCompat.QueueItem;
 import androidx.annotation.Nullable;
 import androidx.media3.common.C;
@@ -25,20 +29,18 @@ import com.google.common.base.Objects;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import java.util.ArrayList;
-import java.util.Collections;
+import java.util.HashMap;
-import java.util.IdentityHashMap;
 import java.util.List;
 import java.util.Map;
 
 /**
- * An immutable class to represent the current {@link Timeline} backed by {@link QueueItem}.
+ * An immutable class to represent the current {@link Timeline} backed by {@linkplain QueueItem
+ * queue items}.
  *
- * <p>This supports the fake item that represents the removed but currently playing media item. In
+ * <p>This timeline supports the case in which the current {@link MediaMetadataCompat} is not
- * that case, a fake item would be inserted at the end of the {@link MediaItem media item list}
+ * included in the queue of the session. In such a case a fake media item is inserted at the end of
- * converted from {@link QueueItem queue item list}. Without the fake item support, the timeline
+ * the timeline and the size of the timeline is by one larger than the size of the corresponding
- * should be always recreated to handle the case when the fake item is no longer necessary and
+ * queue in the session.
- * timeline change isn't precisely detected. Queue item doesn't support equals(), so it's better not
- * to use equals() on the converted MediaItem.
  */
 /* package */ final class QueueTimeline extends Timeline {
 
@@ -48,66 +50,29 @@ import java.util.Map;
   private static final Object FAKE_WINDOW_UID = new Object();
 
   private final ImmutableList<MediaItem> mediaItems;
-  private final Map<MediaItem, Long> unmodifiableMediaItemToQueueIdMap;
+  private final ImmutableMap<MediaItem, Long> mediaItemToQueueIdMap;
   @Nullable private final MediaItem fakeMediaItem;
 
-  private QueueTimeline(
+  /** Creates a new instance. */
-      ImmutableList<MediaItem> mediaItems,
-      Map<MediaItem, Long> unmodifiableMediaItemToQueueIdMap,
-      @Nullable MediaItem fakeMediaItem) {
-    this.mediaItems = mediaItems;
-    this.unmodifiableMediaItemToQueueIdMap = unmodifiableMediaItemToQueueIdMap;
-    this.fakeMediaItem = fakeMediaItem;
-  }
-
   public QueueTimeline(QueueTimeline queueTimeline) {
     this.mediaItems = queueTimeline.mediaItems;
-    this.unmodifiableMediaItemToQueueIdMap = queueTimeline.unmodifiableMediaItemToQueueIdMap;
+    this.mediaItemToQueueIdMap = queueTimeline.mediaItemToQueueIdMap;
     this.fakeMediaItem = queueTimeline.fakeMediaItem;
   }
 
-  public QueueTimeline copyWithFakeMediaItem(@Nullable MediaItem fakeMediaItem) {
+  private QueueTimeline(
-    return new QueueTimeline(mediaItems, unmodifiableMediaItemToQueueIdMap, fakeMediaItem);
+      ImmutableList<MediaItem> mediaItems,
-  }
+      ImmutableMap<MediaItem, Long> mediaItemToQueueIdMap,
-
+      @Nullable MediaItem fakeMediaItem) {
-  public QueueTimeline copyWithNewMediaItem(int replaceIndex, MediaItem newMediaItem) {
+    this.mediaItems = mediaItems;
-    ImmutableList.Builder<MediaItem> newMediaItemsBuilder = new ImmutableList.Builder<>();
+    this.mediaItemToQueueIdMap = mediaItemToQueueIdMap;
-    newMediaItemsBuilder.addAll(mediaItems.subList(0, replaceIndex));
+    this.fakeMediaItem = fakeMediaItem;
-    newMediaItemsBuilder.add(newMediaItem);
-    newMediaItemsBuilder.addAll(mediaItems.subList(replaceIndex + 1, mediaItems.size()));
-    return new QueueTimeline(
-        newMediaItemsBuilder.build(), unmodifiableMediaItemToQueueIdMap, fakeMediaItem);
-  }
-
-  public QueueTimeline copyWithNewMediaItems(int index, List<MediaItem> newMediaItems) {
-    ImmutableList.Builder<MediaItem> newMediaItemsBuilder = new ImmutableList.Builder<>();
-    newMediaItemsBuilder.addAll(mediaItems.subList(0, index));
-    newMediaItemsBuilder.addAll(newMediaItems);
-    newMediaItemsBuilder.addAll(mediaItems.subList(index, mediaItems.size()));
-    return new QueueTimeline(
-        newMediaItemsBuilder.build(), unmodifiableMediaItemToQueueIdMap, fakeMediaItem);
-  }
-
-  public QueueTimeline copyWithRemovedMediaItems(int fromIndex, int toIndex) {
-    ImmutableList.Builder<MediaItem> newMediaItemsBuilder = new ImmutableList.Builder<>();
-    newMediaItemsBuilder.addAll(mediaItems.subList(0, fromIndex));
-    newMediaItemsBuilder.addAll(mediaItems.subList(toIndex, mediaItems.size()));
-    return new QueueTimeline(
-        newMediaItemsBuilder.build(), unmodifiableMediaItemToQueueIdMap, fakeMediaItem);
-  }
-
-  public QueueTimeline copyWithMovedMediaItems(int fromIndex, int toIndex, int newIndex) {
-    List<MediaItem> list = new ArrayList<>(mediaItems);
-    Util.moveItems(list, fromIndex, toIndex, newIndex);
-    return new QueueTimeline(
-        new ImmutableList.Builder<MediaItem>().addAll(list).build(),
-        unmodifiableMediaItemToQueueIdMap,
-        fakeMediaItem);
   }
 
+  /** Creates a {@link QueueTimeline} from a list of {@linkplain QueueItem queue items}. */
   public static QueueTimeline create(List<QueueItem> queue) {
     ImmutableList.Builder<MediaItem> mediaItemsBuilder = new ImmutableList.Builder<>();
-    IdentityHashMap<MediaItem, Long> mediaItemToQueueIdMap = new IdentityHashMap<>();
+    ImmutableMap.Builder<MediaItem, Long> mediaItemToQueueIdMap = new ImmutableMap.Builder<>();
     for (int i = 0; i < queue.size(); i++) {
       QueueItem queueItem = queue.get(i);
       MediaItem mediaItem = MediaUtils.convertToMediaItem(queueItem);
@@ -115,20 +80,122 @@ import java.util.Map;
       mediaItemToQueueIdMap.put(mediaItem, queueItem.getQueueId());
     }
     return new QueueTimeline(
-        mediaItemsBuilder.build(),
+        mediaItemsBuilder.build(), mediaItemToQueueIdMap.buildOrThrow(), /* fakeMediaItem= */ null);
-        Collections.unmodifiableMap(mediaItemToQueueIdMap),
-        /* fakeMediaItem= */ null);
   }
 
+  /**
+   * Gets the queue ID of the media item at the given index or {@link QueueItem#UNKNOWN_ID} if not
+   * known.
+   *
+   * @param mediaItemIndex The media item index.
+   * @return The corresponding queue ID or {@link QueueItem#UNKNOWN_ID} if not known.
+   */
   public long getQueueId(int mediaItemIndex) {
-    @Nullable MediaItem mediaItem = mediaItems.get(mediaItemIndex);
+    MediaItem mediaItem = getMediaItemAt(mediaItemIndex);
-    if (mediaItem == null) {
+    @Nullable Long queueId = mediaItemToQueueIdMap.get(mediaItem);
-      return QueueItem.UNKNOWN_ID;
-    }
-    Long queueId = unmodifiableMediaItemToQueueIdMap.get(mediaItem);
     return queueId == null ? QueueItem.UNKNOWN_ID : queueId;
   }
 
+  /**
+   * Copies the timeline with the given fake media item.
+   *
+   * @param fakeMediaItem The fake media item.
+   * @return A new {@link QueueTimeline} reflecting the update.
+   */
+  public QueueTimeline copyWithFakeMediaItem(@Nullable MediaItem fakeMediaItem) {
+    return new QueueTimeline(mediaItems, mediaItemToQueueIdMap, fakeMediaItem);
+  }
+
+  /**
+   * Replaces the media item at {@code replaceIndex} with the new media item.
+   *
+   * @param replaceIndex The index at which to replace the media item.
+   * @param newMediaItem The new media item that replaces the old one.
+   * @return A new {@link QueueTimeline} reflecting the update.
+   */
+  public QueueTimeline copyWithNewMediaItem(int replaceIndex, MediaItem newMediaItem) {
+    checkArgument(
+        replaceIndex < mediaItems.size()
+            || (replaceIndex == mediaItems.size() && fakeMediaItem != null));
+    if (replaceIndex == mediaItems.size()) {
+      return new QueueTimeline(mediaItems, mediaItemToQueueIdMap, newMediaItem);
+    }
+    MediaItem oldMediaItem = mediaItems.get(replaceIndex);
+    // Create the new play list.
+    ImmutableList.Builder<MediaItem> newMediaItemsBuilder = new ImmutableList.Builder<>();
+    newMediaItemsBuilder.addAll(mediaItems.subList(0, replaceIndex));
+    newMediaItemsBuilder.add(newMediaItem);
+    newMediaItemsBuilder.addAll(mediaItems.subList(replaceIndex + 1, mediaItems.size()));
+    // Update the map of items to queue IDs accordingly.
+    Map<MediaItem, Long> newMediaItemToQueueIdMap = new HashMap<>(mediaItemToQueueIdMap);
+    Long queueId = checkNotNull(newMediaItemToQueueIdMap.remove(oldMediaItem));
+    newMediaItemToQueueIdMap.put(newMediaItem, queueId);
+    return new QueueTimeline(
+        newMediaItemsBuilder.build(), ImmutableMap.copyOf(newMediaItemToQueueIdMap), fakeMediaItem);
+  }
+
+  /**
+   * Replaces the media item at the given index with a list of new media items. The timeline grows
+   * by one less than the size of the new list of items.
+   *
+   * @param index The index of the media item to be replaced.
+   * @param newMediaItems The list of new {@linkplain MediaItem media items} to insert.
+   * @return A new {@link QueueTimeline} reflecting the update.
+   */
+  public QueueTimeline copyWithNewMediaItems(int index, List<MediaItem> newMediaItems) {
+    ImmutableList.Builder<MediaItem> newMediaItemsBuilder = new ImmutableList.Builder<>();
+    newMediaItemsBuilder.addAll(mediaItems.subList(0, index));
+    newMediaItemsBuilder.addAll(newMediaItems);
+    newMediaItemsBuilder.addAll(mediaItems.subList(index, mediaItems.size()));
+    return new QueueTimeline(newMediaItemsBuilder.build(), mediaItemToQueueIdMap, fakeMediaItem);
+  }
+
+  /**
+   * Removes the range of media items in the current timeline.
+   *
+   * @param fromIndex The index to start removing items from.
+   * @param toIndex The index up to which to remove items (exclusive).
+   * @return A new {@link QueueTimeline} reflecting the update.
+   */
+  public QueueTimeline copyWithRemovedMediaItems(int fromIndex, int toIndex) {
+    ImmutableList.Builder<MediaItem> newMediaItemsBuilder = new ImmutableList.Builder<>();
+    newMediaItemsBuilder.addAll(mediaItems.subList(0, fromIndex));
+    newMediaItemsBuilder.addAll(mediaItems.subList(toIndex, mediaItems.size()));
+    return new QueueTimeline(newMediaItemsBuilder.build(), mediaItemToQueueIdMap, fakeMediaItem);
+  }
+
+  /**
+   * Moves the defined range of media items to a new position.
+   *
+   * @param fromIndex The start index of the range to be moved.
+   * @param toIndex The (exclusive) end index of the range to be moved.
+   * @param newIndex The new index to move the first item of the range to.
+   * @return A new {@link QueueTimeline} reflecting the update.
+   */
+  public QueueTimeline copyWithMovedMediaItems(int fromIndex, int toIndex, int newIndex) {
+    List<MediaItem> list = new ArrayList<>(mediaItems);
+    Util.moveItems(list, fromIndex, toIndex, newIndex);
+    return new QueueTimeline(
+        new ImmutableList.Builder<MediaItem>().addAll(list).build(),
+        mediaItemToQueueIdMap,
+        fakeMediaItem);
+  }
+
+  /**
+   * Returns the media item index of the given media item in the timeline, or {@link C#INDEX_UNSET}
+   * if the item is not part of this timeline.
+   *
+   * @param mediaItem The media item of interest.
+   * @return The index of the item or {@link C#INDEX_UNSET} if the item is not part of the timeline.
+   */
+  public int indexOf(MediaItem mediaItem) {
+    if (mediaItem.equals(fakeMediaItem)) {
+      return mediaItems.size();
+    }
+    int mediaItemIndex = mediaItems.indexOf(mediaItem);
+    return mediaItemIndex == -1 ? C.INDEX_UNSET : mediaItemIndex;
+  }
+
   @Nullable
   public MediaItem getMediaItemAt(int mediaItemIndex) {
     if (mediaItemIndex >= 0 && mediaItemIndex < mediaItems.size()) {
@@ -137,14 +204,6 @@ import java.util.Map;
     return (mediaItemIndex == mediaItems.size()) ? fakeMediaItem : null;
   }
 
-  public int findIndexOf(MediaItem mediaItem) {
-    if (mediaItem == fakeMediaItem) {
-      return mediaItems.size();
-    }
-    int mediaItemIndex = mediaItems.indexOf(mediaItem);
-    return mediaItemIndex == -1 ? C.INDEX_UNSET : mediaItemIndex;
-  }
-
   @Override
   public int getWindowCount() {
     return mediaItems.size() + ((fakeMediaItem == null) ? 0 : 1);
@@ -198,14 +257,14 @@ import java.util.Map;
       return false;
     }
     QueueTimeline other = (QueueTimeline) obj;
-    return mediaItems == other.mediaItems
+    return Objects.equal(mediaItems, other.mediaItems)
-        && unmodifiableMediaItemToQueueIdMap == other.unmodifiableMediaItemToQueueIdMap
+        && Objects.equal(mediaItemToQueueIdMap, other.mediaItemToQueueIdMap)
-        && fakeMediaItem == other.fakeMediaItem;
+        && Objects.equal(fakeMediaItem, other.fakeMediaItem);
   }
 
   @Override
   public int hashCode() {
-    return Objects.hashCode(mediaItems, unmodifiableMediaItemToQueueIdMap, fakeMediaItem);
+    return Objects.hashCode(mediaItems, mediaItemToQueueIdMap, fakeMediaItem);
   }
 
   private static Window getWindow(Window window, MediaItem mediaItem, int windowIndex) {

libraries/test_session_common/src/main/aidl/androidx/media3/test/session/common/IRemoteMediaSessionCompat.aidl

@@ -42,4 +42,5 @@ interface IRemoteMediaSessionCompat {
   void sendSessionEvent(String sessionTag, String event, in Bundle extras);
   void setCaptioningEnabled(String sessionTag, boolean enabled);
   void setSessionExtras(String sessionTag, in Bundle extras);
+  int getCallbackMethodCount(String sessionTag, String methodName);
 }

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

@@ -1780,6 +1780,269 @@ public class MediaControllerWithMediaSessionCompatTest {
     assertThat(totalBufferedDurationMs).isEqualTo(testTotalBufferedDurationMs);
   }
 
+  @Test
+  public void prepare_empty_correctInitializationState() throws Exception {
+    session.setPlaybackState(
+        new PlaybackStateCompat.Builder()
+            .setState(PlaybackStateCompat.STATE_NONE, /* position= */ 0, /* playbackSpeed= */ 0.0f)
+            .build());
+    MediaController controller = controllerTestRule.createController(session.getSessionToken());
+
+    // Assert the constructed timeline and start index after connecting to an empty session.
+    int mediaItemCount = threadTestRule.getHandler().postAndSync(controller::getMediaItemCount);
+    int currentMediaItemIndex =
+        threadTestRule.getHandler().postAndSync(controller::getCurrentMediaItemIndex);
+    assertThat(mediaItemCount).isEqualTo(0);
+    assertThat(currentMediaItemIndex).isEqualTo(0);
+  }
+
+  @Test
+  public void prepare_withMetadata_callsPrepareFromMediaId() throws Exception {
+    session.setPlaybackState(
+        new PlaybackStateCompat.Builder()
+            .setState(PlaybackStateCompat.STATE_NONE, /* position= */ 0, /* playbackSpeed= */ 0.0f)
+            .build());
+    session.setMetadata(
+        new MediaMetadataCompat.Builder()
+            .putString(MediaMetadataCompat.METADATA_KEY_MEDIA_ID, "mediaItem_2")
+            .putString(MediaMetadataCompat.METADATA_KEY_DISPLAY_TITLE, "Title")
+            .putString(MediaMetadataCompat.METADATA_KEY_DISPLAY_SUBTITLE, "Subtitle")
+            .putString(MediaMetadataCompat.METADATA_KEY_ARTIST, "Artist")
+            .build());
+    MediaController controller = controllerTestRule.createController(session.getSessionToken());
+    CountDownLatch countDownLatch = new CountDownLatch(1);
+    controller.addListener(
+        new Player.Listener() {
+          @Override
+          public void onEvents(Player player, Player.Events events) {
+            if (events.contains(Player.EVENT_MEDIA_METADATA_CHANGED)) {
+              countDownLatch.countDown();
+            }
+          }
+        });
+
+    // Assert the constructed timeline and start index for preparation.
+    int mediaItemCount = threadTestRule.getHandler().postAndSync(controller::getMediaItemCount);
+    int currentMediaItemIndex =
+        threadTestRule.getHandler().postAndSync(controller::getCurrentMediaItemIndex);
+    assertThat(mediaItemCount).isEqualTo(1);
+    assertThat(currentMediaItemIndex).isEqualTo(0);
+
+    threadTestRule.getHandler().postAndSync(controller::prepare);
+
+    // Assert whether the correct preparation method has been called and received by the session.
+    assertThat(countDownLatch.await(TIMEOUT_MS, MILLISECONDS)).isTrue();
+    int callbackMethodCount =
+        session.getCallbackMethodCount(
+            MediaSessionCompatProviderService.METHOD_ON_PREPARE_FROM_MEDIA_ID);
+    assertThat(callbackMethodCount).isEqualTo(1);
+  }
+
+  @Test
+  public void prepare_withMetadataAndActiveQueueItemId_callsPrepareFromMediaId() throws Exception {
+    session.setPlaybackState(
+        new PlaybackStateCompat.Builder()
+            .setActiveQueueItemId(4)
+            .setState(PlaybackStateCompat.STATE_NONE, /* position= */ 0, /* playbackSpeed= */ 0.0f)
+            .build());
+    session.setMetadata(
+        new MediaMetadataCompat.Builder()
+            .putString(MediaMetadataCompat.METADATA_KEY_MEDIA_ID, "mediaItem_2")
+            .putString(MediaMetadataCompat.METADATA_KEY_DISPLAY_TITLE, "Title")
+            .putString(MediaMetadataCompat.METADATA_KEY_DISPLAY_SUBTITLE, "Subtitle")
+            .putString(MediaMetadataCompat.METADATA_KEY_ARTIST, "Artist")
+            .build());
+    MediaController controller = controllerTestRule.createController(session.getSessionToken());
+    CountDownLatch countDownLatch = new CountDownLatch(1);
+    controller.addListener(
+        new Player.Listener() {
+          @Override
+          public void onEvents(Player player, Player.Events events) {
+            if (events.contains(Player.EVENT_MEDIA_METADATA_CHANGED)) {
+              countDownLatch.countDown();
+            }
+          }
+        });
+
+    // Assert the constructed timeline and start index for preparation.
+    int mediaItemCount = threadTestRule.getHandler().postAndSync(controller::getMediaItemCount);
+    int currentMediaItemIndex =
+        threadTestRule.getHandler().postAndSync(controller::getCurrentMediaItemIndex);
+    assertThat(mediaItemCount).isEqualTo(1);
+    assertThat(currentMediaItemIndex).isEqualTo(0);
+
+    threadTestRule.getHandler().postAndSync(controller::prepare);
+
+    // Assert whether the correct preparation method has been called and received by the session.
+    assertThat(countDownLatch.await(TIMEOUT_MS, MILLISECONDS)).isTrue();
+    int callbackMethodCount =
+        session.getCallbackMethodCount(
+            MediaSessionCompatProviderService.METHOD_ON_PREPARE_FROM_MEDIA_ID);
+    assertThat(callbackMethodCount).isEqualTo(1);
+  }
+
+  @Test
+  public void prepare_withQueue_callsPrepare() throws Exception {
+    List<MediaItem> testMediaItems = MediaTestUtils.createMediaItems(10);
+    List<QueueItem> testQueue = MediaTestUtils.convertToQueueItemsWithoutBitmap(testMediaItems);
+    session.setPlaybackState(
+        new PlaybackStateCompat.Builder()
+            .setState(PlaybackStateCompat.STATE_NONE, /* position= */ 0, /* playbackSpeed= */ 0.0f)
+            .build());
+    session.setQueue(testQueue);
+    MediaController controller = controllerTestRule.createController(session.getSessionToken());
+    CountDownLatch countDownLatch = new CountDownLatch(1);
+    controller.addListener(
+        new Player.Listener() {
+          @Override
+          public void onEvents(Player player, Player.Events events) {
+            if (events.contains(Player.EVENT_MEDIA_METADATA_CHANGED)) {
+              countDownLatch.countDown();
+            }
+          }
+        });
+
+    // Assert the constructed timeline and start index for preparation.
+    int mediaItemCount = threadTestRule.getHandler().postAndSync(controller::getMediaItemCount);
+    int currentMediaItemIndex =
+        threadTestRule.getHandler().postAndSync(controller::getCurrentMediaItemIndex);
+    assertThat(mediaItemCount).isEqualTo(10);
+    assertThat(currentMediaItemIndex).isEqualTo(0);
+
+    threadTestRule.getHandler().postAndSync(controller::prepare);
+
+    // Assert whether the correct preparation method has been called and received by the session.
+    assertThat(countDownLatch.await(TIMEOUT_MS, MILLISECONDS)).isTrue();
+    int callbackMethodCount =
+        session.getCallbackMethodCount(MediaSessionCompatProviderService.METHOD_ON_PREPARE);
+    assertThat(callbackMethodCount).isEqualTo(1);
+  }
+
+  @Test
+  public void prepare_withQueueAndActiveQueueItemId_callsPrepare() throws Exception {
+    List<MediaItem> testMediaItems = MediaTestUtils.createMediaItems(10);
+    List<QueueItem> testQueue = MediaTestUtils.convertToQueueItemsWithoutBitmap(testMediaItems);
+    session.setPlaybackState(
+        new PlaybackStateCompat.Builder()
+            .setActiveQueueItemId(5)
+            .setState(PlaybackStateCompat.STATE_NONE, /* position= */ 0, /* playbackSpeed= */ 0.0f)
+            .build());
+    session.setQueue(testQueue);
+    MediaController controller = controllerTestRule.createController(session.getSessionToken());
+    CountDownLatch countDownLatch = new CountDownLatch(1);
+    controller.addListener(
+        new Player.Listener() {
+          @Override
+          public void onEvents(Player player, Player.Events events) {
+            if (events.contains(Player.EVENT_MEDIA_METADATA_CHANGED)) {
+              countDownLatch.countDown();
+            }
+          }
+        });
+
+    // Assert the constructed timeline and start index for preparation.
+    int mediaItemCount = threadTestRule.getHandler().postAndSync(controller::getMediaItemCount);
+    int currentMediaItemIndex =
+        threadTestRule.getHandler().postAndSync(controller::getCurrentMediaItemIndex);
+    assertThat(mediaItemCount).isEqualTo(10);
+    assertThat(currentMediaItemIndex).isEqualTo(5);
+
+    threadTestRule.getHandler().postAndSync(controller::prepare);
+
+    // Assert whether the correct preparation method has been called and received by the session.
+    assertThat(countDownLatch.await(TIMEOUT_MS, MILLISECONDS)).isTrue();
+    int callbackMethodCount =
+        session.getCallbackMethodCount(MediaSessionCompatProviderService.METHOD_ON_PREPARE);
+    assertThat(callbackMethodCount).isEqualTo(1);
+  }
+
+  @Test
+  public void prepare_withQueueAndMetadata_callsPrepareFromMediaId() throws Exception {
+    List<MediaItem> testMediaItems = MediaTestUtils.createMediaItems(10);
+    List<QueueItem> testQueue = MediaTestUtils.convertToQueueItemsWithoutBitmap(testMediaItems);
+    session.setPlaybackState(
+        new PlaybackStateCompat.Builder()
+            .setState(PlaybackStateCompat.STATE_NONE, /* position= */ 0, /* playbackSpeed= */ 0.0f)
+            .build());
+    session.setMetadata(
+        new MediaMetadataCompat.Builder()
+            .putString(MediaMetadataCompat.METADATA_KEY_MEDIA_ID, "mediaItem_2")
+            .putString(MediaMetadataCompat.METADATA_KEY_DISPLAY_TITLE, "Title")
+            .putString(MediaMetadataCompat.METADATA_KEY_DISPLAY_SUBTITLE, "Subtitle")
+            .putString(MediaMetadataCompat.METADATA_KEY_ARTIST, "Artist")
+            .build());
+    session.setQueue(testQueue);
+    MediaController controller = controllerTestRule.createController(session.getSessionToken());
+    CountDownLatch countDownLatch = new CountDownLatch(1);
+    controller.addListener(
+        new Player.Listener() {
+          @Override
+          public void onEvents(Player player, Player.Events events) {
+            if (events.contains(Player.EVENT_MEDIA_METADATA_CHANGED)) {
+              countDownLatch.countDown();
+            }
+          }
+        });
+
+    // Assert the constructed timeline and start index for preparation.
+    int mediaItemCount = threadTestRule.getHandler().postAndSync(controller::getMediaItemCount);
+    int currentMediaItemIndex =
+        threadTestRule.getHandler().postAndSync(controller::getCurrentMediaItemIndex);
+    assertThat(mediaItemCount).isEqualTo(11);
+    assertThat(currentMediaItemIndex).isEqualTo(10);
+
+    threadTestRule.getHandler().postAndSync(controller::prepare);
+
+    // Assert whether the correct preparation method has been called and received by the session.
+    assertThat(countDownLatch.await(TIMEOUT_MS, MILLISECONDS)).isTrue();
+    int callbackMethodCount =
+        session.getCallbackMethodCount(
+            MediaSessionCompatProviderService.METHOD_ON_PREPARE_FROM_MEDIA_ID);
+    assertThat(callbackMethodCount).isEqualTo(1);
+  }
+
+  @Test
+  public void prepare_withQueueAndMetadataAndActiveQueueItemId_callsPrepare() throws Exception {
+    List<MediaItem> testMediaItems = MediaTestUtils.createMediaItems(10);
+    List<QueueItem> testQueue = MediaTestUtils.convertToQueueItemsWithoutBitmap(testMediaItems);
+    session.setPlaybackState(
+        new PlaybackStateCompat.Builder()
+            .setActiveQueueItemId(4)
+            .setState(PlaybackStateCompat.STATE_NONE, /* position= */ 0, /* playbackSpeed= */ 0.0f)
+            .build());
+    session.setMetadata(
+        new MediaMetadataCompat.Builder()
+            .putString(MediaMetadataCompat.METADATA_KEY_MEDIA_ID, "mediaItem_5")
+            .build());
+    session.setQueue(testQueue);
+    MediaController controller = controllerTestRule.createController(session.getSessionToken());
+    CountDownLatch countDownLatch = new CountDownLatch(1);
+    controller.addListener(
+        new Player.Listener() {
+          @Override
+          public void onEvents(Player player, Player.Events events) {
+            if (events.contains(Player.EVENT_MEDIA_METADATA_CHANGED)) {
+              countDownLatch.countDown();
+            }
+          }
+        });
+
+    // Assert the constructed timeline and start index for preparation.
+    int mediaItemCount = threadTestRule.getHandler().postAndSync(controller::getMediaItemCount);
+    int currentMediaItemIndex =
+        threadTestRule.getHandler().postAndSync(controller::getCurrentMediaItemIndex);
+    assertThat(mediaItemCount).isEqualTo(10);
+    assertThat(currentMediaItemIndex).isEqualTo(4);
+
+    threadTestRule.getHandler().postAndSync(controller::prepare);
+
+    // Assert whether the correct preparation method has been called and received by the session.
+    assertThat(countDownLatch.await(TIMEOUT_MS, MILLISECONDS)).isTrue();
+    int callbackMethodCount =
+        session.getCallbackMethodCount(MediaSessionCompatProviderService.METHOD_ON_PREPARE);
+    assertThat(callbackMethodCount).isEqualTo(1);
+  }
+
   @Nullable
   private Bitmap getBitmapFromMetadata(MediaMetadata metadata) throws Exception {
     @Nullable Bitmap bitmap = null;

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

@@ -49,9 +49,13 @@ import java.util.concurrent.Executor;
 @UnstableApi
 public class MediaSessionCompatProviderService extends Service {
 
+  public static final String METHOD_ON_PREPARE_FROM_MEDIA_ID = "onPrepareFromMediaId";
+  public static final String METHOD_ON_PREPARE = "onPrepare";
+
   private static final String TAG = "MSCProviderService";
 
   Map<String, MediaSessionCompat> sessionMap = new HashMap<>();
+  Map<String, CallCountingCallback> callbackMap = new HashMap<>();
   RemoteMediaSessionCompatStub sessionBinder;
 
   TestHandler handler;
@@ -88,7 +92,10 @@ public class MediaSessionCompatProviderService extends Service {
             () -> {
               MediaSessionCompat session =
                   new MediaSessionCompat(MediaSessionCompatProviderService.this, sessionTag);
+              CallCountingCallback callback = new CallCountingCallback(sessionTag);
+              session.setCallback(callback);
               sessionMap.put(sessionTag, session);
+              callbackMap.put(sessionTag, callback);
             });
       } catch (Exception e) {
         Log.e(TAG, "Exception occurred while creating MediaSessionCompat", e);
@@ -212,15 +219,61 @@ public class MediaSessionCompatProviderService extends Service {
     }
 
     @Override
-    public void setCaptioningEnabled(String sessionTag, boolean enabled) throws RemoteException {
+    public void setCaptioningEnabled(String sessionTag, boolean enabled) {
       MediaSessionCompat session = sessionMap.get(sessionTag);
       session.setCaptioningEnabled(enabled);
     }
 
     @Override
-    public void setSessionExtras(String sessionTag, Bundle extras) throws RemoteException {
+    public void setSessionExtras(String sessionTag, Bundle extras) {
       MediaSessionCompat session = sessionMap.get(sessionTag);
       session.setExtras(extras);
     }
+
+    @Override
+    public int getCallbackMethodCount(String sessionTag, String methodName) {
+      CallCountingCallback callCountingCallback = callbackMap.get(sessionTag);
+      if (callCountingCallback != null) {
+        Integer count = callCountingCallback.callbackCallCounters.get(methodName);
+        return count != null ? count : 0;
+      }
+      return 0;
+    }
+  }
+
+  private class CallCountingCallback extends MediaSessionCompat.Callback {
+
+    private final String sessionTag;
+    private final Map<String, Integer> callbackCallCounters;
+
+    public CallCountingCallback(String sessionTag) {
+      this.sessionTag = sessionTag;
+      callbackCallCounters = new HashMap<>();
+    }
+
+    @Override
+    public void onPrepareFromMediaId(String mediaId, Bundle extras) {
+      countCallbackCall(METHOD_ON_PREPARE_FROM_MEDIA_ID);
+      sessionMap
+          .get(sessionTag)
+          .setMetadata(
+              new MediaMetadataCompat.Builder()
+                  .putString(MediaMetadataCompat.METADATA_KEY_MEDIA_ID, mediaId)
+                  .build());
+    }
+
+    @Override
+    public void onPrepare() {
+      countCallbackCall(METHOD_ON_PREPARE);
+      sessionMap.get(sessionTag).setMetadata(new MediaMetadataCompat.Builder().build());
+    }
+
+    private void countCallbackCall(String callbackName) {
+      int count = 0;
+      if (callbackCallCounters.containsKey(callbackName)) {
+        count = callbackCallCounters.get(callbackName);
+      }
+      callbackCallCounters.put(callbackName, ++count);
+    }
   }
 }

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

@@ -111,6 +111,10 @@ public class RemoteMediaSessionCompat {
     binder.setPlaybackToLocal(sessionTag, stream);
   }
 
+  public int getCallbackMethodCount(String callbackMethodName) throws RemoteException {
+    return binder.getCallbackMethodCount(sessionTag, callbackMethodName);
+  }
+
   /**
    * Since we cannot pass VolumeProviderCompat directly, we pass volumeControl, maxVolume,
    * currentVolume instead.