Fix the IndexOutOfBoundsException when creating output media item list

The current implementation of building the output media item list
keeps all media items in a single list. They try to mark the position
of the original divider items and insert them back to the original
position. However, the list of media devices might change due to
various reasons. So inserting the original divider items back to the
same position as the original media item list could cause the
IndexOutOfBoundsException. To fix this issue, we create three separated
media item lists to store selected media items, suggested media items,
and other speakers and displays media items correspondingly. Then we
insert the divider items when creating the output media item list.

Flag: com.android.media.flags.fix_output_media_item_list_index_out_of_bounds_exception
Bug: b/400241830
Test: manual test and presubmit
Change-Id: I05026adaaae974e8271cb6d74d026da8c3593a4f

5 files changed, 701 insertions, 26 deletions

diff --git a/media/java/android/media/flags/media_better_together.aconfig b/media/java/android/media/flags/media_better_together.aconfig
index e39a0aa8717e..48e2f4e15238 100644
--- a/media/java/android/media/flags/media_better_together.aconfig
+++ b/media/java/android/media/flags/media_better_together.aconfig
@@ -242,3 +242,13 @@ flag {
     description: "Fallbacks to the default handling for volume adjustment when media session has fixed volume handling and its app is in the foreground and setting a media controller."
     bug: "293743975"
 }
+
+flag {
+    name: "fix_output_media_item_list_index_out_of_bounds_exception"
+    namespace: "media_better_together"
+    description: "Fixes a bug of causing IndexOutOfBoundsException when building media item list."
+    bug: "398246089"
+    metadata {
+        purpose: PURPOSE_BUGFIX
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/media/dialog/MediaSwitchingController.java b/packages/SystemUI/src/com/android/systemui/media/dialog/MediaSwitchingController.java
index f79693138e24..d0c6a3e6a3ef 100644
--- a/packages/SystemUI/src/com/android/systemui/media/dialog/MediaSwitchingController.java
+++ b/packages/SystemUI/src/com/android/systemui/media/dialog/MediaSwitchingController.java
@@ -66,6 +66,7 @@ import androidx.annotation.VisibleForTesting;
 import androidx.core.graphics.drawable.IconCompat;
 
 import com.android.internal.annotations.GuardedBy;
+import com.android.media.flags.Flags;
 import com.android.settingslib.RestrictedLockUtilsInternal;
 import com.android.settingslib.Utils;
 import com.android.settingslib.bluetooth.BluetoothUtils;
@@ -78,7 +79,6 @@ import com.android.settingslib.media.InputMediaDevice;
 import com.android.settingslib.media.InputRouteManager;
 import com.android.settingslib.media.LocalMediaManager;
 import com.android.settingslib.media.MediaDevice;
-import com.android.settingslib.media.flags.Flags;
 import com.android.settingslib.utils.ThreadUtils;
 import com.android.systemui.animation.ActivityTransitionAnimator;
 import com.android.systemui.animation.DialogTransitionAnimator;
@@ -226,7 +226,7 @@ public class MediaSwitchingController
                 InfoMediaManager.createInstance(mContext, packageName, userHandle, lbm, token);
         mLocalMediaManager = new LocalMediaManager(mContext, lbm, imm, packageName);
         mMetricLogger = new MediaOutputMetricLogger(mContext, mPackageName);
-        mOutputMediaItemListProxy = new OutputMediaItemListProxy();
+        mOutputMediaItemListProxy = new OutputMediaItemListProxy(context);
         mDialogTransitionAnimator = dialogTransitionAnimator;
         mNearbyMediaDevicesManager = nearbyMediaDevicesManager;
         mMediaOutputColorSchemeLegacy = MediaOutputColorSchemeLegacy.fromSystemColors(mContext);
@@ -308,7 +308,8 @@ public class MediaSwitchingController
     }
 
     private MediaController getMediaController() {
-        if (mToken != null && Flags.usePlaybackInfoForRoutingControls()) {
+        if (mToken != null
+                && com.android.settingslib.media.flags.Flags.usePlaybackInfoForRoutingControls()) {
             return new MediaController(mContext, mToken);
         } else {
             for (NotificationEntry entry : mNotifCollection.getAllNotifs()) {
@@ -577,19 +578,35 @@ public class MediaSwitchingController
 
     private void buildMediaItems(List<MediaDevice> devices) {
         synchronized (mMediaDevicesLock) {
-            List<MediaItem> updatedMediaItems =
-                    buildMediaItems(mOutputMediaItemListProxy.getOutputMediaItemList(), devices);
-            mOutputMediaItemListProxy.clearAndAddAll(updatedMediaItems);
+            if (!mLocalMediaManager.isPreferenceRouteListingExist()) {
+                attachRangeInfo(devices);
+                Collections.sort(devices, Comparator.naturalOrder());
+            }
+            if (Flags.fixOutputMediaItemListIndexOutOfBoundsException()) {
+                // For the first time building list, to make sure the top device is the connected
+                // device.
+                boolean needToHandleMutingExpectedDevice =
+                        hasMutingExpectedDevice() && !isCurrentConnectedDeviceRemote();
+                final MediaDevice connectedMediaDevice =
+                        needToHandleMutingExpectedDevice ? null : getCurrentConnectedMediaDevice();
+                mOutputMediaItemListProxy.updateMediaDevices(
+                        devices,
+                        getSelectedMediaDevice(),
+                        connectedMediaDevice,
+                        needToHandleMutingExpectedDevice,
+                        getConnectNewDeviceItem());
+            } else {
+                List<MediaItem> updatedMediaItems =
+                        buildMediaItems(
+                                mOutputMediaItemListProxy.getOutputMediaItemList(), devices);
+                mOutputMediaItemListProxy.clearAndAddAll(updatedMediaItems);
+            }
         }
     }
 
     protected List<MediaItem> buildMediaItems(
             List<MediaItem> oldMediaItems, List<MediaDevice> devices) {
         synchronized (mMediaDevicesLock) {
-            if (!mLocalMediaManager.isPreferenceRouteListingExist()) {
-                attachRangeInfo(devices);
-                Collections.sort(devices, Comparator.naturalOrder());
-            }
             // For the first time building list, to make sure the top device is the connected
             // device.
             boolean needToHandleMutingExpectedDevice =
@@ -648,8 +665,7 @@ public class MediaSwitchingController
                     .map(MediaItem::createDeviceMediaItem)
                     .collect(Collectors.toList());
 
-            boolean shouldAddFirstSeenSelectedDevice =
-                    com.android.media.flags.Flags.enableOutputSwitcherDeviceGrouping();
+            boolean shouldAddFirstSeenSelectedDevice = Flags.enableOutputSwitcherDeviceGrouping();
 
             if (shouldAddFirstSeenSelectedDevice) {
                 finalMediaItems.clear();
@@ -675,7 +691,7 @@ public class MediaSwitchingController
     }
 
     private boolean enableInputRouting() {
-        return com.android.media.flags.Flags.enableAudioInputDeviceRoutingAndVolumeControl();
+        return Flags.enableAudioInputDeviceRoutingAndVolumeControl();
     }
 
     private void buildInputMediaItems(List<MediaDevice> devices) {
@@ -703,8 +719,7 @@ public class MediaSwitchingController
         if (connectedMediaDevice != null) {
             selectedDevicesIds.add(connectedMediaDevice.getId());
         }
-        boolean groupSelectedDevices =
-                com.android.media.flags.Flags.enableOutputSwitcherDeviceGrouping();
+        boolean groupSelectedDevices = Flags.enableOutputSwitcherDeviceGrouping();
         int nextSelectedItemIndex = 0;
         boolean suggestedDeviceAdded = false;
         boolean displayGroupAdded = false;
@@ -879,6 +894,11 @@ public class MediaSwitchingController
         return mLocalMediaManager.getCurrentConnectedDevice();
     }
 
+    @VisibleForTesting
+    void clearMediaItemList() {
+        mOutputMediaItemListProxy.clear();
+    }
+
     boolean addDeviceToPlayMedia(MediaDevice device) {
         mMetricLogger.logInteractionExpansion(device);
         return mLocalMediaManager.addDeviceToPlayMedia(device);
diff --git a/packages/SystemUI/src/com/android/systemui/media/dialog/OutputMediaItemListProxy.java b/packages/SystemUI/src/com/android/systemui/media/dialog/OutputMediaItemListProxy.java
index 1c9c0b102cb7..45ca2c6ee8e5 100644
--- a/packages/SystemUI/src/com/android/systemui/media/dialog/OutputMediaItemListProxy.java
+++ b/packages/SystemUI/src/com/android/systemui/media/dialog/OutputMediaItemListProxy.java
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2020 The Android Open Source Project
+ * Copyright (C) 2025 The Android Open Source Project
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -16,22 +16,175 @@
 
 package com.android.systemui.media.dialog;
 
+import android.content.Context;
+
+import androidx.annotation.Nullable;
+
+import com.android.media.flags.Flags;
+import com.android.settingslib.media.MediaDevice;
+import com.android.systemui.res.R;
+
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.HashSet;
 import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.Set;
 import java.util.concurrent.CopyOnWriteArrayList;
+import java.util.stream.Collectors;
 
 /** A proxy of holding the list of Output Switcher's output media items. */
 public class OutputMediaItemListProxy {
+    private final Context mContext;
     private final List<MediaItem> mOutputMediaItemList;
 
-    public OutputMediaItemListProxy() {
+    // Use separated lists to hold different media items and create the list of output media items
+    // by using those separated lists and group dividers.
+    private final List<MediaItem> mSelectedMediaItems;
+    private final List<MediaItem> mSuggestedMediaItems;
+    private final List<MediaItem> mSpeakersAndDisplaysMediaItems;
+    @Nullable private MediaItem mConnectNewDeviceMediaItem;
+
+    public OutputMediaItemListProxy(Context context) {
+        mContext = context;
         mOutputMediaItemList = new CopyOnWriteArrayList<>();
+        mSelectedMediaItems = new CopyOnWriteArrayList<>();
+        mSuggestedMediaItems = new CopyOnWriteArrayList<>();
+        mSpeakersAndDisplaysMediaItems = new CopyOnWriteArrayList<>();
     }
 
     /** Returns the list of output media items. */
     public List<MediaItem> getOutputMediaItemList() {
+        if (Flags.fixOutputMediaItemListIndexOutOfBoundsException()) {
+            if (isEmpty() && !mOutputMediaItemList.isEmpty()) {
+                // Ensures mOutputMediaItemList is empty when all individual media item lists are
+                // empty, preventing unexpected state issues.
+                mOutputMediaItemList.clear();
+            } else if (!isEmpty() && mOutputMediaItemList.isEmpty()) {
+                // When any individual media item list is modified, the cached mOutputMediaItemList
+                // is emptied. On the next request for the output media item list, a fresh list is
+                // created and stored in the cache.
+                mOutputMediaItemList.addAll(createOutputMediaItemList());
+            }
+        }
         return mOutputMediaItemList;
     }
 
+    private List<MediaItem> createOutputMediaItemList() {
+        List<MediaItem> finalMediaItems = new CopyOnWriteArrayList<>();
+        finalMediaItems.addAll(mSelectedMediaItems);
+        if (!mSuggestedMediaItems.isEmpty()) {
+            finalMediaItems.add(
+                    MediaItem.createGroupDividerMediaItem(
+                            mContext.getString(
+                                    R.string.media_output_group_title_suggested_device)));
+            finalMediaItems.addAll(mSuggestedMediaItems);
+        }
+        if (!mSpeakersAndDisplaysMediaItems.isEmpty()) {
+            finalMediaItems.add(
+                    MediaItem.createGroupDividerMediaItem(
+                            mContext.getString(
+                                    R.string.media_output_group_title_speakers_and_displays)));
+            finalMediaItems.addAll(mSpeakersAndDisplaysMediaItems);
+        }
+        if (mConnectNewDeviceMediaItem != null) {
+            finalMediaItems.add(mConnectNewDeviceMediaItem);
+        }
+        return finalMediaItems;
+    }
+
+    /** Updates the list of output media items with a given list of media devices. */
+    public void updateMediaDevices(
+            List<MediaDevice> devices,
+            List<MediaDevice> selectedDevices,
+            @Nullable MediaDevice connectedMediaDevice,
+            boolean needToHandleMutingExpectedDevice,
+            @Nullable MediaItem connectNewDeviceMediaItem) {
+        Set<String> selectedOrConnectedMediaDeviceIds =
+                selectedDevices.stream().map(MediaDevice::getId).collect(Collectors.toSet());
+        if (connectedMediaDevice != null) {
+            selectedOrConnectedMediaDeviceIds.add(connectedMediaDevice.getId());
+        }
+
+        List<MediaItem> selectedMediaItems = new ArrayList<>();
+        List<MediaItem> suggestedMediaItems = new ArrayList<>();
+        List<MediaItem> speakersAndDisplaysMediaItems = new ArrayList<>();
+        Map<String, MediaItem> deviceIdToMediaItemMap = new HashMap<>();
+        buildMediaItems(
+                devices,
+                selectedOrConnectedMediaDeviceIds,
+                needToHandleMutingExpectedDevice,
+                selectedMediaItems,
+                suggestedMediaItems,
+                speakersAndDisplaysMediaItems,
+                deviceIdToMediaItemMap);
+
+        List<MediaItem> updatedSelectedMediaItems = new CopyOnWriteArrayList<>();
+        List<MediaItem> updatedSuggestedMediaItems = new CopyOnWriteArrayList<>();
+        List<MediaItem> updatedSpeakersAndDisplaysMediaItems = new CopyOnWriteArrayList<>();
+        if (isEmpty()) {
+            updatedSelectedMediaItems.addAll(selectedMediaItems);
+            updatedSuggestedMediaItems.addAll(suggestedMediaItems);
+            updatedSpeakersAndDisplaysMediaItems.addAll(speakersAndDisplaysMediaItems);
+        } else {
+            Set<String> updatedDeviceIds = new HashSet<>();
+            // Preserve the existing media item order while updating with the latest device
+            // information. Some items may retain their original group (suggested, speakers and
+            // displays) to maintain this order.
+            updateMediaItems(
+                    mSelectedMediaItems,
+                    updatedSelectedMediaItems,
+                    deviceIdToMediaItemMap,
+                    updatedDeviceIds);
+            updateMediaItems(
+                    mSuggestedMediaItems,
+                    updatedSuggestedMediaItems,
+                    deviceIdToMediaItemMap,
+                    updatedDeviceIds);
+            updateMediaItems(
+                    mSpeakersAndDisplaysMediaItems,
+                    updatedSpeakersAndDisplaysMediaItems,
+                    deviceIdToMediaItemMap,
+                    updatedDeviceIds);
+
+            // Append new media items that are not already in the existing lists to the output list.
+            List<MediaItem> remainingMediaItems = new ArrayList<>();
+            remainingMediaItems.addAll(
+                    getRemainingMediaItems(selectedMediaItems, updatedDeviceIds));
+            remainingMediaItems.addAll(
+                    getRemainingMediaItems(suggestedMediaItems, updatedDeviceIds));
+            remainingMediaItems.addAll(
+                    getRemainingMediaItems(speakersAndDisplaysMediaItems, updatedDeviceIds));
+            updatedSpeakersAndDisplaysMediaItems.addAll(remainingMediaItems);
+        }
+
+        if (Flags.enableOutputSwitcherDeviceGrouping() && !updatedSelectedMediaItems.isEmpty()) {
+            MediaItem selectedMediaItem = updatedSelectedMediaItems.get(0);
+            Optional<MediaDevice> mediaDeviceOptional = selectedMediaItem.getMediaDevice();
+            if (mediaDeviceOptional.isPresent()) {
+                MediaItem updatedMediaItem =
+                        MediaItem.createDeviceMediaItem(
+                                mediaDeviceOptional.get(), /* isFirstDeviceInGroup= */ true);
+                updatedSelectedMediaItems.remove(0);
+                updatedSelectedMediaItems.add(0, updatedMediaItem);
+            }
+        }
+
+        mSelectedMediaItems.clear();
+        mSelectedMediaItems.addAll(updatedSelectedMediaItems);
+        mSuggestedMediaItems.clear();
+        mSuggestedMediaItems.addAll(updatedSuggestedMediaItems);
+        mSpeakersAndDisplaysMediaItems.clear();
+        mSpeakersAndDisplaysMediaItems.addAll(updatedSpeakersAndDisplaysMediaItems);
+        mConnectNewDeviceMediaItem = connectNewDeviceMediaItem;
+
+        // The cached mOutputMediaItemList is cleared upon any update to individual media item
+        // lists. This ensures getOutputMediaItemList() computes and caches a fresh list on the next
+        // invocation.
+        mOutputMediaItemList.clear();
+    }
+
     /** Updates the list of output media items with the given list. */
     public void clearAndAddAll(List<MediaItem> updatedMediaItems) {
         mOutputMediaItemList.clear();
@@ -40,16 +193,112 @@ public class OutputMediaItemListProxy {
 
     /** Removes the media items with muting expected devices. */
     public void removeMutingExpectedDevices() {
+        if (Flags.fixOutputMediaItemListIndexOutOfBoundsException()) {
+            mSelectedMediaItems.removeIf((MediaItem::isMutingExpectedDevice));
+            mSuggestedMediaItems.removeIf((MediaItem::isMutingExpectedDevice));
+            mSpeakersAndDisplaysMediaItems.removeIf((MediaItem::isMutingExpectedDevice));
+            if (mConnectNewDeviceMediaItem != null
+                    && mConnectNewDeviceMediaItem.isMutingExpectedDevice()) {
+                mConnectNewDeviceMediaItem = null;
+            }
+        }
         mOutputMediaItemList.removeIf((MediaItem::isMutingExpectedDevice));
     }
 
     /** Clears the output media item list. */
     public void clear() {
+        if (Flags.fixOutputMediaItemListIndexOutOfBoundsException()) {
+            mSelectedMediaItems.clear();
+            mSuggestedMediaItems.clear();
+            mSpeakersAndDisplaysMediaItems.clear();
+            mConnectNewDeviceMediaItem = null;
+        }
         mOutputMediaItemList.clear();
     }
 
     /** Returns whether the output media item list is empty. */
     public boolean isEmpty() {
-        return mOutputMediaItemList.isEmpty();
+        if (Flags.fixOutputMediaItemListIndexOutOfBoundsException()) {
+            return mSelectedMediaItems.isEmpty()
+                    && mSuggestedMediaItems.isEmpty()
+                    && mSpeakersAndDisplaysMediaItems.isEmpty()
+                    && (mConnectNewDeviceMediaItem == null);
+        } else {
+            return mOutputMediaItemList.isEmpty();
+        }
+    }
+
+    private void buildMediaItems(
+            List<MediaDevice> devices,
+            Set<String> selectedOrConnectedMediaDeviceIds,
+            boolean needToHandleMutingExpectedDevice,
+            List<MediaItem> selectedMediaItems,
+            List<MediaItem> suggestedMediaItems,
+            List<MediaItem> speakersAndDisplaysMediaItems,
+            Map<String, MediaItem> deviceIdToMediaItemMap) {
+        for (MediaDevice device : devices) {
+            String deviceId = device.getId();
+            MediaItem mediaItem = MediaItem.createDeviceMediaItem(device);
+            if (needToHandleMutingExpectedDevice && device.isMutingExpectedDevice()) {
+                selectedMediaItems.add(0, mediaItem);
+            } else if (!needToHandleMutingExpectedDevice
+                    && selectedOrConnectedMediaDeviceIds.contains(device.getId())) {
+                if (Flags.enableOutputSwitcherDeviceGrouping()) {
+                    selectedMediaItems.add(mediaItem);
+                } else {
+                    selectedMediaItems.add(0, mediaItem);
+                }
+            } else if (device.isSuggestedDevice()) {
+                suggestedMediaItems.add(mediaItem);
+            } else {
+                speakersAndDisplaysMediaItems.add(mediaItem);
+            }
+            deviceIdToMediaItemMap.put(deviceId, mediaItem);
+        }
+    }
+
+    /** Returns a list of media items that remains the same order as the existing media items. */
+    private void updateMediaItems(
+            List<MediaItem> existingMediaItems,
+            List<MediaItem> updatedMediaItems,
+            Map<String, MediaItem> deviceIdToMediaItemMap,
+            Set<String> updatedDeviceIds) {
+        List<String> existingDeviceIds = getDeviceIds(existingMediaItems);
+        for (String deviceId : existingDeviceIds) {
+            MediaItem mediaItem = deviceIdToMediaItemMap.get(deviceId);
+            if (mediaItem != null) {
+                updatedMediaItems.add(mediaItem);
+                updatedDeviceIds.add(deviceId);
+            }
+        }
+    }
+
+    /**
+     * Returns media items from the input list that are not associated with the given device IDs.
+     */
+    private List<MediaItem> getRemainingMediaItems(
+            List<MediaItem> mediaItems, Set<String> deviceIds) {
+        List<MediaItem> remainingMediaItems = new ArrayList<>();
+        for (MediaItem item : mediaItems) {
+            Optional<MediaDevice> mediaDeviceOptional = item.getMediaDevice();
+            if (mediaDeviceOptional.isPresent()) {
+                String deviceId = mediaDeviceOptional.get().getId();
+                if (!deviceIds.contains(deviceId)) {
+                    remainingMediaItems.add(item);
+                }
+            }
+        }
+        return remainingMediaItems;
+    }
+
+    /** Returns a list of media device IDs for the given list of media items. */
+    private List<String> getDeviceIds(List<MediaItem> mediaItems) {
+        List<String> deviceIds = new ArrayList<>();
+        for (MediaItem item : mediaItems) {
+            if (item != null && item.getMediaDevice().isPresent()) {
+                deviceIds.add(item.getMediaDevice().get().getId());
+            }
+        }
+        return deviceIds;
     }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/media/dialog/MediaSwitchingControllerTest.java b/packages/SystemUI/tests/src/com/android/systemui/media/dialog/MediaSwitchingControllerTest.java
index 5c26dac5eb30..798aa428e73e 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/media/dialog/MediaSwitchingControllerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/media/dialog/MediaSwitchingControllerTest.java
@@ -60,13 +60,13 @@ import android.os.RemoteException;
 import android.os.UserHandle;
 import android.platform.test.annotations.DisableFlags;
 import android.platform.test.annotations.EnableFlags;
+import android.platform.test.flag.junit.FlagsParameterization;
 import android.service.notification.StatusBarNotification;
 import android.testing.TestableLooper;
 import android.text.TextUtils;
 import android.view.View;
 
 import androidx.core.graphics.drawable.IconCompat;
-import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
 
 import com.android.media.flags.Flags;
@@ -101,6 +101,9 @@ import org.mockito.Mock;
 import org.mockito.Mockito;
 import org.mockito.MockitoAnnotations;
 
+import platform.test.runner.parameterized.ParameterizedAndroidJunit4;
+import platform.test.runner.parameterized.Parameters;
+
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
@@ -108,7 +111,7 @@ import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
 
 @SmallTest
-@RunWith(AndroidJUnit4.class)
+@RunWith(ParameterizedAndroidJunit4.class)
 @TestableLooper.RunWithLooper(setAsMainLooper = true)
 public class MediaSwitchingControllerTest extends SysuiTestCase {
     private static final String TEST_DEVICE_1_ID = "test_device_1_id";
@@ -201,6 +204,17 @@ public class MediaSwitchingControllerTest extends SysuiTestCase {
     private MediaDescription mMediaDescription;
     private List<RoutingSessionInfo> mRoutingSessionInfos = new ArrayList<>();
 
+    @Parameters(name = "{0}")
+    public static List<FlagsParameterization> getParams() {
+        return FlagsParameterization.allCombinationsOf(
+                Flags.FLAG_FIX_OUTPUT_MEDIA_ITEM_LIST_INDEX_OUT_OF_BOUNDS_EXCEPTION,
+                Flags.FLAG_ENABLE_OUTPUT_SWITCHER_DEVICE_GROUPING);
+    }
+
+    public MediaSwitchingControllerTest(FlagsParameterization flags) {
+        mSetFlagsRule.setFlagsParameterization(flags);
+    }
+
     @Before
     public void setUp() {
         mPackageName = mContext.getPackageName();
@@ -260,7 +274,6 @@ public class MediaSwitchingControllerTest extends SysuiTestCase {
         mMediaDevices.add(mMediaDevice1);
         mMediaDevices.add(mMediaDevice2);
 
-
         when(mNearbyDevice1.getMediaRoute2Id()).thenReturn(TEST_DEVICE_1_ID);
         when(mNearbyDevice1.getRangeZone()).thenReturn(NearbyDevice.RANGE_FAR);
         when(mNearbyDevice2.getMediaRoute2Id()).thenReturn(TEST_DEVICE_2_ID);
@@ -689,7 +702,7 @@ public class MediaSwitchingControllerTest extends SysuiTestCase {
 
         mMediaSwitchingController.start(mCb);
         reset(mCb);
-        mMediaSwitchingController.getMediaItemList().clear();
+        mMediaSwitchingController.clearMediaItemList();
         mMediaSwitchingController.onDeviceListUpdate(mMediaDevices);
         final List<MediaDevice> devices = new ArrayList<>();
         int dividerSize = 0;
@@ -1528,7 +1541,7 @@ public class MediaSwitchingControllerTest extends SysuiTestCase {
                 .getSelectedMediaDevice();
         mMediaSwitchingController.start(mCb);
         reset(mCb);
-        mMediaSwitchingController.getMediaItemList().clear();
+        mMediaSwitchingController.clearMediaItemList();
 
         mMediaSwitchingController.onDeviceListUpdate(mMediaDevices);
 
@@ -1546,7 +1559,7 @@ public class MediaSwitchingControllerTest extends SysuiTestCase {
                 .getSelectedMediaDevice();
         mMediaSwitchingController.start(mCb);
         reset(mCb);
-        mMediaSwitchingController.getMediaItemList().clear();
+        mMediaSwitchingController.clearMediaItemList();
 
         mMediaSwitchingController.onDeviceListUpdate(mMediaDevices);
 
@@ -1564,7 +1577,7 @@ public class MediaSwitchingControllerTest extends SysuiTestCase {
                 .getSelectedMediaDevice();
         mMediaSwitchingController.start(mCb);
         reset(mCb);
-        mMediaSwitchingController.getMediaItemList().clear();
+        mMediaSwitchingController.clearMediaItemList();
 
         mMediaSwitchingController.onDeviceListUpdate(mMediaDevices);
 
@@ -1582,7 +1595,7 @@ public class MediaSwitchingControllerTest extends SysuiTestCase {
                 .getSelectedMediaDevice();
         mMediaSwitchingController.start(mCb);
         reset(mCb);
-        mMediaSwitchingController.getMediaItemList().clear();
+        mMediaSwitchingController.clearMediaItemList();
         mMediaSwitchingController.onDeviceListUpdate(mMediaDevices);
         mMediaDevices.clear();
         mMediaDevices.add(mMediaDevice2);
diff --git a/packages/SystemUI/tests/src/com/android/systemui/media/dialog/OutputMediaItemListProxyTest.java b/packages/SystemUI/tests/src/com/android/systemui/media/dialog/OutputMediaItemListProxyTest.java
new file mode 100644
index 000000000000..f6edd49f142f
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/media/dialog/OutputMediaItemListProxyTest.java
@@ -0,0 +1,383 @@
+/*
+ * Copyright (C) 2025 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.media.dialog;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import static org.mockito.Mockito.when;
+
+import android.platform.test.annotations.DisableFlags;
+import android.platform.test.annotations.EnableFlags;
+import android.platform.test.flag.junit.FlagsParameterization;
+import android.testing.TestableLooper;
+
+import androidx.test.filters.SmallTest;
+
+import com.android.media.flags.Flags;
+import com.android.settingslib.media.MediaDevice;
+import com.android.systemui.SysuiTestCase;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+
+import platform.test.runner.parameterized.ParameterizedAndroidJunit4;
+import platform.test.runner.parameterized.Parameters;
+
+import java.util.List;
+import java.util.stream.Collectors;
+
+@SmallTest
+@RunWith(ParameterizedAndroidJunit4.class)
+@TestableLooper.RunWithLooper
+public class OutputMediaItemListProxyTest extends SysuiTestCase {
+    private static final String DEVICE_ID_1 = "device_id_1";
+    private static final String DEVICE_ID_2 = "device_id_2";
+    private static final String DEVICE_ID_3 = "device_id_3";
+    private static final String DEVICE_ID_4 = "device_id_4";
+    @Mock private MediaDevice mMediaDevice1;
+    @Mock private MediaDevice mMediaDevice2;
+    @Mock private MediaDevice mMediaDevice3;
+    @Mock private MediaDevice mMediaDevice4;
+
+    private MediaItem mMediaItem1;
+    private MediaItem mMediaItem2;
+    private MediaItem mConnectNewDeviceMediaItem;
+    private OutputMediaItemListProxy mOutputMediaItemListProxy;
+
+    @Parameters(name = "{0}")
+    public static List<FlagsParameterization> getParams() {
+        return FlagsParameterization.allCombinationsOf(
+                Flags.FLAG_FIX_OUTPUT_MEDIA_ITEM_LIST_INDEX_OUT_OF_BOUNDS_EXCEPTION,
+                Flags.FLAG_ENABLE_OUTPUT_SWITCHER_DEVICE_GROUPING);
+    }
+
+    public OutputMediaItemListProxyTest(FlagsParameterization flags) {
+        mSetFlagsRule.setFlagsParameterization(flags);
+    }
+
+    @Before
+    public void setUp() {
+        MockitoAnnotations.initMocks(this);
+
+        when(mMediaDevice1.getId()).thenReturn(DEVICE_ID_1);
+        when(mMediaDevice2.getId()).thenReturn(DEVICE_ID_2);
+        when(mMediaDevice2.isSuggestedDevice()).thenReturn(true);
+        when(mMediaDevice3.getId()).thenReturn(DEVICE_ID_3);
+        when(mMediaDevice4.getId()).thenReturn(DEVICE_ID_4);
+        mMediaItem1 = MediaItem.createDeviceMediaItem(mMediaDevice1);
+        mMediaItem2 = MediaItem.createDeviceMediaItem(mMediaDevice2);
+        mConnectNewDeviceMediaItem = MediaItem.createPairNewDeviceMediaItem();
+
+        mOutputMediaItemListProxy = new OutputMediaItemListProxy(mContext);
+    }
+
+    @EnableFlags(Flags.FLAG_FIX_OUTPUT_MEDIA_ITEM_LIST_INDEX_OUT_OF_BOUNDS_EXCEPTION)
+    @Test
+    public void updateMediaDevices_shouldUpdateMediaItemList() {
+        assertThat(mOutputMediaItemListProxy.isEmpty()).isTrue();
+
+        // Create the initial output media item list with mMediaDevice2 and mMediaDevice3.
+        mOutputMediaItemListProxy.updateMediaDevices(
+                /* devices= */ List.of(mMediaDevice2, mMediaDevice3),
+                /* selectedDevices */ List.of(mMediaDevice3),
+                /* connectedMediaDevice= */ null,
+                /* needToHandleMutingExpectedDevice= */ false,
+                /* connectNewDeviceMediaItem= */ null);
+
+        // Check the output media items to be
+        //     * a media item with the selected mMediaDevice3
+        //     * a group divider for suggested devices
+        //     * a media item with the mMediaDevice2
+        assertThat(getMediaDevices(mOutputMediaItemListProxy.getOutputMediaItemList()))
+                .containsExactly(mMediaDevice3, null, mMediaDevice2);
+        assertThat(mOutputMediaItemListProxy.getOutputMediaItemList().get(0).isFirstDeviceInGroup())
+                .isEqualTo(Flags.enableOutputSwitcherDeviceGrouping());
+
+        // Update the output media item list with more media devices.
+        mOutputMediaItemListProxy.updateMediaDevices(
+                /* devices= */ List.of(mMediaDevice4, mMediaDevice1, mMediaDevice3, mMediaDevice2),
+                /* selectedDevices */ List.of(mMediaDevice3),
+                /* connectedMediaDevice= */ null,
+                /* needToHandleMutingExpectedDevice= */ false,
+                /* connectNewDeviceMediaItem= */ null);
+
+        // Check the output media items to be
+        //     * a media item with the selected route mMediaDevice3
+        //     * a group divider for suggested devices
+        //     * a media item with the route mMediaDevice2
+        //     * a group divider for speakers and displays
+        //     * a media item with the route mMediaDevice4
+        //     * a media item with the route mMediaDevice1
+        assertThat(getMediaDevices(mOutputMediaItemListProxy.getOutputMediaItemList()))
+                .containsExactly(
+                        mMediaDevice3, null, mMediaDevice2, null, mMediaDevice4, mMediaDevice1);
+        assertThat(mOutputMediaItemListProxy.getOutputMediaItemList().get(0).isFirstDeviceInGroup())
+                .isEqualTo(Flags.enableOutputSwitcherDeviceGrouping());
+
+        // Update the output media item list where mMediaDevice4 is offline and new selected device.
+        mOutputMediaItemListProxy.updateMediaDevices(
+                /* devices= */ List.of(mMediaDevice1, mMediaDevice3, mMediaDevice2),
+                /* selectedDevices */ List.of(mMediaDevice1, mMediaDevice3),
+                /* connectedMediaDevice= */ null,
+                /* needToHandleMutingExpectedDevice= */ false,
+                /* connectNewDeviceMediaItem= */ null);
+
+        // Check the output media items to be
+        //     * a media item with the selected route mMediaDevice3
+        //     * a group divider for suggested devices
+        //     * a media item with the route mMediaDevice2
+        //     * a group divider for speakers and displays
+        //     * a media item with the route mMediaDevice1
+        assertThat(getMediaDevices(mOutputMediaItemListProxy.getOutputMediaItemList()))
+                .containsExactly(mMediaDevice3, null, mMediaDevice2, null, mMediaDevice1);
+        assertThat(mOutputMediaItemListProxy.getOutputMediaItemList().get(0).isFirstDeviceInGroup())
+                .isEqualTo(Flags.enableOutputSwitcherDeviceGrouping());
+    }
+
+    @EnableFlags(Flags.FLAG_FIX_OUTPUT_MEDIA_ITEM_LIST_INDEX_OUT_OF_BOUNDS_EXCEPTION)
+    @Test
+    public void updateMediaDevices_multipleSelectedDevices_shouldHaveCorrectDeviceOrdering() {
+        assertThat(mOutputMediaItemListProxy.isEmpty()).isTrue();
+
+        // Create the initial output media item list with mMediaDevice2 and mMediaDevice3.
+        mOutputMediaItemListProxy.updateMediaDevices(
+                /* devices= */ List.of(mMediaDevice2, mMediaDevice4, mMediaDevice3, mMediaDevice1),
+                /* selectedDevices */ List.of(mMediaDevice1, mMediaDevice2, mMediaDevice3),
+                /* connectedMediaDevice= */ null,
+                /* needToHandleMutingExpectedDevice= */ false,
+                /* connectNewDeviceMediaItem= */ null);
+
+        if (Flags.enableOutputSwitcherDeviceGrouping()) {
+            // When the device grouping is enabled, the order of selected devices are preserved:
+            //     * a media item with the selected mMediaDevice2
+            //     * a media item with the selected mMediaDevice3
+            //     * a media item with the selected mMediaDevice1
+            //     * a group divider for speakers and displays
+            //     * a media item with the mMediaDevice4
+            assertThat(getMediaDevices(mOutputMediaItemListProxy.getOutputMediaItemList()))
+                    .containsExactly(
+                            mMediaDevice2, mMediaDevice3, mMediaDevice1, null, mMediaDevice4);
+            assertThat(
+                            mOutputMediaItemListProxy
+                                    .getOutputMediaItemList()
+                                    .get(0)
+                                    .isFirstDeviceInGroup())
+                    .isTrue();
+        } else {
+            // When the device grouping is disabled, the order of selected devices are reverted:
+            //     * a media item with the selected mMediaDevice1
+            //     * a media item with the selected mMediaDevice3
+            //     * a media item with the selected mMediaDevice2
+            //     * a group divider for speakers and displays
+            //     * a media item with the mMediaDevice4
+            assertThat(getMediaDevices(mOutputMediaItemListProxy.getOutputMediaItemList()))
+                    .containsExactly(
+                            mMediaDevice1, mMediaDevice3, mMediaDevice2, null, mMediaDevice4);
+        }
+
+        // Update the output media item list with a selected device being deselected.
+        mOutputMediaItemListProxy.updateMediaDevices(
+                /* devices= */ List.of(mMediaDevice4, mMediaDevice1, mMediaDevice3, mMediaDevice2),
+                /* selectedDevices */ List.of(mMediaDevice2, mMediaDevice3),
+                /* connectedMediaDevice= */ null,
+                /* needToHandleMutingExpectedDevice= */ false,
+                /* connectNewDeviceMediaItem= */ null);
+
+        if (Flags.enableOutputSwitcherDeviceGrouping()) {
+            // When the device grouping is enabled, the order of selected devices are preserved:
+            //     * a media item with the selected mMediaDevice2
+            //     * a media item with the selected mMediaDevice3
+            //     * a media item with the selected mMediaDevice1
+            //     * a group divider for speakers and displays
+            //     * a media item with the mMediaDevice4
+            assertThat(getMediaDevices(mOutputMediaItemListProxy.getOutputMediaItemList()))
+                    .containsExactly(
+                            mMediaDevice2, mMediaDevice3, mMediaDevice1, null, mMediaDevice4);
+            assertThat(
+                            mOutputMediaItemListProxy
+                                    .getOutputMediaItemList()
+                                    .get(0)
+                                    .isFirstDeviceInGroup())
+                    .isTrue();
+        } else {
+            // When the device grouping is disabled, the order of selected devices are reverted:
+            //     * a media item with the selected mMediaDevice1
+            //     * a media item with the selected mMediaDevice3
+            //     * a media item with the selected mMediaDevice2
+            //     * a group divider for speakers and displays
+            //     * a media item with the mMediaDevice4
+            assertThat(getMediaDevices(mOutputMediaItemListProxy.getOutputMediaItemList()))
+                    .containsExactly(
+                            mMediaDevice1, mMediaDevice3, mMediaDevice2, null, mMediaDevice4);
+        }
+
+        // Update the output media item list with a selected device is missing.
+        mOutputMediaItemListProxy.updateMediaDevices(
+                /* devices= */ List.of(mMediaDevice1, mMediaDevice3, mMediaDevice4),
+                /* selectedDevices */ List.of(mMediaDevice3),
+                /* connectedMediaDevice= */ null,
+                /* needToHandleMutingExpectedDevice= */ false,
+                /* connectNewDeviceMediaItem= */ null);
+
+        if (Flags.enableOutputSwitcherDeviceGrouping()) {
+            // When the device grouping is enabled, the order of selected devices are preserved:
+            //     * a media item with the selected mMediaDevice3
+            //     * a media item with the selected mMediaDevice1
+            //     * a group divider for speakers and displays
+            //     * a media item with the mMediaDevice4
+            assertThat(getMediaDevices(mOutputMediaItemListProxy.getOutputMediaItemList()))
+                    .containsExactly(mMediaDevice3, mMediaDevice1, null, mMediaDevice4);
+            assertThat(
+                            mOutputMediaItemListProxy
+                                    .getOutputMediaItemList()
+                                    .get(0)
+                                    .isFirstDeviceInGroup())
+                    .isTrue();
+        } else {
+            // When the device grouping is disabled, the order of selected devices are reverted:
+            //     * a media item with the selected mMediaDevice1
+            //     * a media item with the selected mMediaDevice3
+            //     * a group divider for speakers and displays
+            //     * a media item with the mMediaDevice4
+            assertThat(getMediaDevices(mOutputMediaItemListProxy.getOutputMediaItemList()))
+                    .containsExactly(mMediaDevice1, mMediaDevice3, null, mMediaDevice4);
+        }
+    }
+
+    @EnableFlags(Flags.FLAG_FIX_OUTPUT_MEDIA_ITEM_LIST_INDEX_OUT_OF_BOUNDS_EXCEPTION)
+    @Test
+    public void updateMediaDevices_withConnectNewDeviceMediaItem_shouldUpdateMediaItemList() {
+        assertThat(mOutputMediaItemListProxy.isEmpty()).isTrue();
+
+        // Create the initial output media item list with a connect new device media item.
+        mOutputMediaItemListProxy.updateMediaDevices(
+                /* devices= */ List.of(mMediaDevice2, mMediaDevice3),
+                /* selectedDevices */ List.of(mMediaDevice3),
+                /* connectedMediaDevice= */ null,
+                /* needToHandleMutingExpectedDevice= */ false,
+                mConnectNewDeviceMediaItem);
+
+        // Check the output media items to be
+        //     * a media item with the selected mMediaDevice3
+        //     * a group divider for suggested devices
+        //     * a media item with the mMediaDevice2
+        //     * a connect new device media item
+        assertThat(mOutputMediaItemListProxy.getOutputMediaItemList())
+                .contains(mConnectNewDeviceMediaItem);
+        assertThat(getMediaDevices(mOutputMediaItemListProxy.getOutputMediaItemList()))
+                .containsExactly(mMediaDevice3, null, mMediaDevice2, null);
+
+        // Update the output media item list without a connect new device media item.
+        mOutputMediaItemListProxy.updateMediaDevices(
+                /* devices= */ List.of(mMediaDevice2, mMediaDevice3),
+                /* selectedDevices */ List.of(mMediaDevice3),
+                /* connectedMediaDevice= */ null,
+                /* needToHandleMutingExpectedDevice= */ false,
+                /* connectNewDeviceMediaItem= */ null);
+
+        // Check the output media items to be
+        //     * a media item with the selected mMediaDevice3
+        //     * a group divider for suggested devices
+        //     * a media item with the mMediaDevice2
+        assertThat(mOutputMediaItemListProxy.getOutputMediaItemList())
+                .doesNotContain(mConnectNewDeviceMediaItem);
+        assertThat(getMediaDevices(mOutputMediaItemListProxy.getOutputMediaItemList()))
+                .containsExactly(mMediaDevice3, null, mMediaDevice2);
+    }
+
+    @DisableFlags(Flags.FLAG_FIX_OUTPUT_MEDIA_ITEM_LIST_INDEX_OUT_OF_BOUNDS_EXCEPTION)
+    @Test
+    public void clearAndAddAll_shouldUpdateMediaItemList() {
+        assertThat(mOutputMediaItemListProxy.isEmpty()).isTrue();
+
+        mOutputMediaItemListProxy.clearAndAddAll(List.of(mMediaItem1));
+        assertThat(mOutputMediaItemListProxy.getOutputMediaItemList()).containsExactly(mMediaItem1);
+        assertThat(mOutputMediaItemListProxy.isEmpty()).isFalse();
+
+        mOutputMediaItemListProxy.clearAndAddAll(List.of(mMediaItem2));
+        assertThat(mOutputMediaItemListProxy.getOutputMediaItemList()).containsExactly(mMediaItem2);
+        assertThat(mOutputMediaItemListProxy.isEmpty()).isFalse();
+    }
+
+    @EnableFlags(Flags.FLAG_FIX_OUTPUT_MEDIA_ITEM_LIST_INDEX_OUT_OF_BOUNDS_EXCEPTION)
+    @Test
+    public void clear_flagOn_shouldClearMediaItemList() {
+        assertThat(mOutputMediaItemListProxy.isEmpty()).isTrue();
+
+        mOutputMediaItemListProxy.updateMediaDevices(
+                /* devices= */ List.of(mMediaDevice1),
+                /* selectedDevices */ List.of(),
+                /* connectedMediaDevice= */ null,
+                /* needToHandleMutingExpectedDevice= */ false,
+                /* connectNewDeviceMediaItem= */ null);
+        assertThat(mOutputMediaItemListProxy.isEmpty()).isFalse();
+
+        mOutputMediaItemListProxy.clear();
+        assertThat(mOutputMediaItemListProxy.isEmpty()).isTrue();
+    }
+
+    @DisableFlags(Flags.FLAG_FIX_OUTPUT_MEDIA_ITEM_LIST_INDEX_OUT_OF_BOUNDS_EXCEPTION)
+    @Test
+    public void clear_flagOff_shouldClearMediaItemList() {
+        assertThat(mOutputMediaItemListProxy.isEmpty()).isTrue();
+
+        mOutputMediaItemListProxy.clearAndAddAll(List.of(mMediaItem1));
+        assertThat(mOutputMediaItemListProxy.isEmpty()).isFalse();
+
+        mOutputMediaItemListProxy.clear();
+        assertThat(mOutputMediaItemListProxy.isEmpty()).isTrue();
+    }
+
+    @EnableFlags(Flags.FLAG_FIX_OUTPUT_MEDIA_ITEM_LIST_INDEX_OUT_OF_BOUNDS_EXCEPTION)
+    @Test
+    public void removeMutingExpectedDevices_flagOn_shouldClearMediaItemList() {
+        assertThat(mOutputMediaItemListProxy.isEmpty()).isTrue();
+
+        mOutputMediaItemListProxy.updateMediaDevices(
+                /* devices= */ List.of(mMediaDevice1),
+                /* selectedDevices */ List.of(),
+                /* connectedMediaDevice= */ null,
+                /* needToHandleMutingExpectedDevice= */ false,
+                /* connectNewDeviceMediaItem= */ null);
+        assertThat(mOutputMediaItemListProxy.isEmpty()).isFalse();
+
+        mOutputMediaItemListProxy.removeMutingExpectedDevices();
+        assertThat(mOutputMediaItemListProxy.isEmpty()).isFalse();
+    }
+
+    @DisableFlags(Flags.FLAG_FIX_OUTPUT_MEDIA_ITEM_LIST_INDEX_OUT_OF_BOUNDS_EXCEPTION)
+    @Test
+    public void removeMutingExpectedDevices_flagOff_shouldClearMediaItemList() {
+        assertThat(mOutputMediaItemListProxy.isEmpty()).isTrue();
+
+        mOutputMediaItemListProxy.clearAndAddAll(List.of(mMediaItem1));
+        assertThat(mOutputMediaItemListProxy.isEmpty()).isFalse();
+
+        mOutputMediaItemListProxy.removeMutingExpectedDevices();
+        assertThat(mOutputMediaItemListProxy.getOutputMediaItemList()).containsExactly(mMediaItem1);
+        assertThat(mOutputMediaItemListProxy.isEmpty()).isFalse();
+    }
+
+    private List<MediaDevice> getMediaDevices(List<MediaItem> mediaItems) {
+        return mediaItems.stream()
+                .map(item -> item.getMediaDevice().orElse(null))
+                .collect(Collectors.toList());
+    }
+}