Remove deduplication from Shareousel

Test: atest com.android.intentresolver
BUG: 351911089
Flag: EXEMPT bugfix
Change-Id: I5eb5d93b61891e8f91928361f375b1ca6562f724

12 files changed, 291 insertions, 107 deletions

diff --git a/java/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/interactor/CursorPreviewsInteractor.kt b/java/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/interactor/CursorPreviewsInteractor.kt
index 7d658209..0e198f43 100644
--- a/java/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/interactor/CursorPreviewsInteractor.kt
+++ b/java/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/interactor/CursorPreviewsInteractor.kt
@@ -30,6 +30,7 @@ import com.android.intentresolver.contentpreview.payloadtoggle.domain.model.expa
 import com.android.intentresolver.contentpreview.payloadtoggle.domain.model.numLoadedPages
 import com.android.intentresolver.contentpreview.payloadtoggle.domain.model.shiftWindowLeft
 import com.android.intentresolver.contentpreview.payloadtoggle.domain.model.shiftWindowRight
+import com.android.intentresolver.contentpreview.payloadtoggle.shared.model.PreviewKey
 import com.android.intentresolver.contentpreview.payloadtoggle.shared.model.PreviewModel
 import com.android.intentresolver.inject.FocusedItemIndex
 import com.android.intentresolver.util.cursor.CursorView
@@ -82,16 +83,19 @@ constructor(
                 .toMap(ConcurrentHashMap())
         val pagedCursor: PagedCursor<CursorRow?> = uriCursor.paged(pageSize)
         val startPosition = uriCursor.extras?.getInt(POSITION, 0) ?: 0
+
         val state =
             loadToMaxPages(
-                initialState = readInitialState(pagedCursor, startPosition, unclaimedRecords),
+                startPosition = startPosition,
+                initialState = readInitialState(startPosition, pagedCursor, unclaimedRecords),
                 pagedCursor = pagedCursor,
                 unclaimedRecords = unclaimedRecords,
             )
-        processLoadRequests(state, pagedCursor, unclaimedRecords)
+        processLoadRequests(startPosition, state, pagedCursor, unclaimedRecords)
     }
 
     private suspend fun loadToMaxPages(
+        startPosition: Int,
         initialState: CursorWindow,
         pagedCursor: PagedCursor<CursorRow?>,
         unclaimedRecords: MutableUnclaimedMap,
@@ -113,9 +117,10 @@ constructor(
             state =
                 when {
                     state.hasMoreLeft && loadedLeft < loadedRight ->
-                        state.loadMoreLeft(pagedCursor, unclaimedRecords)
-                    state.hasMoreRight -> state.loadMoreRight(pagedCursor, unclaimedRecords)
-                    else -> state.loadMoreLeft(pagedCursor, unclaimedRecords)
+                        state.loadMoreLeft(startPosition, pagedCursor, unclaimedRecords)
+                    state.hasMoreRight ->
+                        state.loadMoreRight(startPosition, pagedCursor, unclaimedRecords)
+                    else -> state.loadMoreLeft(startPosition, pagedCursor, unclaimedRecords)
                 }
         }
         return state
@@ -123,6 +128,7 @@ constructor(
 
     /** Loop forever, processing any loading requests from the UI and updating local cache. */
     private suspend fun processLoadRequests(
+        startPosition: Int,
         initialState: CursorWindow,
         pagedCursor: PagedCursor<CursorRow?>,
         unclaimedRecords: MutableUnclaimedMap,
@@ -144,7 +150,13 @@ constructor(
                     leftTriggerIndex = leftTriggerIndex,
                     rightTriggerIndex = rightTriggerIndex,
                 )
-            state = loadingState.handleOneLoadRequest(state, pagedCursor, unclaimedRecords)
+            state =
+                loadingState.handleOneLoadRequest(
+                    startPosition,
+                    state,
+                    pagedCursor,
+                    unclaimedRecords,
+                )
         }
     }
 
@@ -153,6 +165,7 @@ constructor(
      * with the loaded data incorporated.
      */
     private suspend fun Flow<LoadDirection?>.handleOneLoadRequest(
+        startPosition: Int,
         state: CursorWindow,
         pagedCursor: PagedCursor<CursorRow?>,
         unclaimedRecords: MutableUnclaimedMap,
@@ -160,8 +173,10 @@ constructor(
         mapLatest { loadDirection ->
                 loadDirection?.let {
                     when (loadDirection) {
-                        LoadDirection.Left -> state.loadMoreLeft(pagedCursor, unclaimedRecords)
-                        LoadDirection.Right -> state.loadMoreRight(pagedCursor, unclaimedRecords)
+                        LoadDirection.Left ->
+                            state.loadMoreLeft(startPosition, pagedCursor, unclaimedRecords)
+                        LoadDirection.Right ->
+                            state.loadMoreRight(startPosition, pagedCursor, unclaimedRecords)
                     }
                 }
             }
@@ -169,12 +184,12 @@ constructor(
             .first()
 
     /**
-     * Returns the initial [CursorWindow], with a single page loaded that contains the given
+     * Returns the initial [CursorWindow], with a single page loaded that contains the
      * [startPosition].
      */
     private suspend fun readInitialState(
-        cursor: PagedCursor<CursorRow?>,
         startPosition: Int,
+        cursor: PagedCursor<CursorRow?>,
         unclaimedRecords: MutableUnclaimedMap,
     ): CursorWindow {
         val startPageIdx = startPosition / pageSize
@@ -184,13 +199,15 @@ constructor(
             if (!hasMoreLeft) {
                 // First read the initial page; this might claim some unclaimed Uris
                 val page =
-                    cursor.getPageRows(startPageIdx)?.toPage(mutableMapOf(), unclaimedRecords)
+                    cursor
+                        .getPageRows(startPageIdx)
+                        ?.toPage(startPosition, mutableMapOf(), unclaimedRecords)
                 // Now that unclaimed Uris are up-to-date, add them first.
                 putAllUnclaimedLeft(unclaimedRecords)
                 // Then add the loaded page
                 page?.let(::putAll)
             } else {
-                cursor.getPageRows(startPageIdx)?.toPage(this, unclaimedRecords)
+                cursor.getPageRows(startPageIdx)?.toPage(startPosition, this, unclaimedRecords)
             }
             // Finally, add the remainder of the unclaimed Uris.
             if (!hasMoreRight) {
@@ -208,13 +225,14 @@ constructor(
     }
 
     private suspend fun CursorWindow.loadMoreRight(
+        startPosition: Int,
         cursor: PagedCursor<CursorRow?>,
         unclaimedRecords: MutableUnclaimedMap,
     ): CursorWindow {
         val pageNum = lastLoadedPageNum + 1
         val hasMoreRight = pageNum < cursor.count - 1
         val newPage: PreviewMap = buildMap {
-            readAndPutPage(this@loadMoreRight, cursor, pageNum, unclaimedRecords)
+            readAndPutPage(startPosition, this@loadMoreRight, cursor, pageNum, unclaimedRecords)
             if (!hasMoreRight) {
                 putAllUnclaimedRight(unclaimedRecords)
             }
@@ -227,6 +245,7 @@ constructor(
     }
 
     private suspend fun CursorWindow.loadMoreLeft(
+        startPosition: Int,
         cursor: PagedCursor<CursorRow?>,
         unclaimedRecords: MutableUnclaimedMap,
     ): CursorWindow {
@@ -235,13 +254,14 @@ constructor(
         val newPage: PreviewMap = buildMap {
             if (!hasMoreLeft) {
                 // First read the page; this might claim some unclaimed Uris
-                val page = readPage(this@loadMoreLeft, cursor, pageNum, unclaimedRecords)
+                val page =
+                    readPage(startPosition, this@loadMoreLeft, cursor, pageNum, unclaimedRecords)
                 // Now that unclaimed URIs are up-to-date, add them first
                 putAllUnclaimedLeft(unclaimedRecords)
                 // Then add the loaded page
                 putAll(page)
             } else {
-                readAndPutPage(this@loadMoreLeft, cursor, pageNum, unclaimedRecords)
+                readAndPutPage(startPosition, this@loadMoreLeft, cursor, pageNum, unclaimedRecords)
             }
         }
         return if (numLoadedPages < maxLoadedPages) {
@@ -259,15 +279,17 @@ constructor(
     }
 
     private suspend fun readPage(
+        startPosition: Int,
         state: CursorWindow,
         pagedCursor: PagedCursor<CursorRow?>,
         pageNum: Int,
         unclaimedRecords: MutableUnclaimedMap,
     ): PreviewMap =
-        mutableMapOf<Uri, PreviewModel>()
-            .readAndPutPage(state, pagedCursor, pageNum, unclaimedRecords)
+        mutableMapOf<PreviewKey, PreviewModel>()
+            .readAndPutPage(startPosition, state, pagedCursor, pageNum, unclaimedRecords)
 
     private suspend fun <M : MutablePreviewMap> M.readAndPutPage(
+        startPosition: Int,
         state: CursorWindow,
         pagedCursor: PagedCursor<CursorRow?>,
         pageNum: Int,
@@ -275,19 +297,23 @@ constructor(
     ): M =
         pagedCursor
             .getPageRows(pageNum) // TODO: what do we do if the load fails?
-            ?.filter { it.uri !in state.merged }
-            ?.toPage(this, unclaimedRecords) ?: this
+            ?.filter { PreviewKey.final(it.position - startPosition) !in state.merged }
+            ?.toPage(startPosition, this, unclaimedRecords) ?: this
 
     private suspend fun <M : MutablePreviewMap> Sequence<CursorRow>.toPage(
+        startPosition: Int,
         destination: M,
         unclaimedRecords: MutableUnclaimedMap,
     ): M =
         // Restrict parallelism so as to not overload the metadata reader; anecdotally, too
         // many parallel queries causes failures.
-        mapParallel(parallelism = 4) { row -> createPreviewModel(row, unclaimedRecords) }
-            .associateByTo(destination) { it.uri }
+        mapParallel(parallelism = 4) { row ->
+                createPreviewModel(startPosition, row, unclaimedRecords)
+            }
+            .associateByTo(destination) { it.key }
 
     private fun createPreviewModel(
+        startPosition: Int,
         row: CursorRow,
         unclaimedRecords: MutableUnclaimedMap,
     ): PreviewModel =
@@ -298,6 +324,7 @@ constructor(
                     row.previewSize
                         ?: metadata.previewUri?.let { uriMetadataReader.readPreviewSize(it) }
                 PreviewModel(
+                    key = PreviewKey.final(row.position - startPosition),
                     uri = row.uri,
                     previewUri = metadata.previewUri,
                     mimeType = metadata.mimeType,
@@ -308,11 +335,9 @@ constructor(
             .also { updated ->
                 if (unclaimedRecords.remove(row.uri) != null) {
                     // unclaimedRecords contains initially shared (and thus selected) items with
-                    // unknown
-                    // cursor position. Update selection records when any of those items is
-                    // encountered
-                    // in the cursor to maintain proper selection order should other items also be
-                    // selected.
+                    // unknown cursor position. Update selection records when any of those items is
+                    // encountered in the cursor to maintain proper selection order should other
+                    // items also be selected.
                     selectionInteractor.updateSelection(updated)
                 }
             }
@@ -324,7 +349,7 @@ constructor(
         putAllUnclaimedWhere(unclaimed) { it < focusedItemIdx }
 }
 
-private typealias CursorWindow = LoadedWindow<Uri, PreviewModel>
+private typealias CursorWindow = LoadedWindow<PreviewKey, PreviewModel>
 
 /**
  * Values from the initial selection set that have not yet appeared within the Cursor. These values
@@ -336,9 +361,13 @@ private typealias UnclaimedMap = Map<Uri, Pair<Int, PreviewModel>>
 /** Mutable version of [UnclaimedMap]. */
 private typealias MutableUnclaimedMap = MutableMap<Uri, Pair<Int, PreviewModel>>
 
-private typealias MutablePreviewMap = MutableMap<Uri, PreviewModel>
+private typealias UnkeyedMap = Map<Uri, PreviewModel>
+
+private typealias MutableUnkeyedMap = MutableMap<Uri, PreviewModel>
+
+private typealias MutablePreviewMap = MutableMap<PreviewKey, PreviewModel>
 
-private typealias PreviewMap = Map<Uri, PreviewModel>
+private typealias PreviewMap = Map<PreviewKey, PreviewModel>
 
 private fun <M : MutablePreviewMap> M.putAllUnclaimedWhere(
     unclaimedRecords: UnclaimedMap,
@@ -347,7 +376,7 @@ private fun <M : MutablePreviewMap> M.putAllUnclaimedWhere(
     unclaimedRecords
         .asSequence()
         .filter { predicate(it.value.first) }
-        .map { it.key to it.value.second }
+        .map { (_, value) -> value.second.key to value.second }
         .toMap(this)
 
 private fun PagedCursor<CursorRow?>.getPageRows(pageNum: Int): Sequence<CursorRow>? =
diff --git a/java/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/interactor/FetchPreviewsInteractor.kt b/java/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/interactor/FetchPreviewsInteractor.kt
index 50086a23..1fd69351 100644
--- a/java/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/interactor/FetchPreviewsInteractor.kt
+++ b/java/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/interactor/FetchPreviewsInteractor.kt
@@ -22,6 +22,7 @@ import com.android.intentresolver.contentpreview.payloadtoggle.data.repository.P
 import com.android.intentresolver.contentpreview.payloadtoggle.domain.cursor.CursorResolver
 import com.android.intentresolver.contentpreview.payloadtoggle.domain.cursor.PayloadToggle
 import com.android.intentresolver.contentpreview.payloadtoggle.domain.model.CursorRow
+import com.android.intentresolver.contentpreview.payloadtoggle.shared.model.PreviewKey
 import com.android.intentresolver.contentpreview.payloadtoggle.shared.model.PreviewModel
 import com.android.intentresolver.inject.ContentUris
 import com.android.intentresolver.inject.FocusedItemIndex
@@ -64,6 +65,12 @@ constructor(
             .mapParallelIndexed(parallelism = 4) { index, uri ->
                 val metadata = uriMetadataReader.getMetadata(uri)
                 PreviewModel(
+                    key =
+                        if (index == focusedItemIdx) {
+                            PreviewKey.final(0)
+                        } else {
+                            PreviewKey.temp(index)
+                        },
                     uri = uri,
                     previewUri = metadata.previewUri,
                     mimeType = metadata.mimeType,
@@ -71,11 +78,12 @@ constructor(
                         metadata.previewUri?.let {
                             uriMetadataReader.readPreviewSize(it).aspectRatioOrDefault(1f)
                         } ?: 1f,
-                    order = when {
-                        index < focusedItemIdx -> Int.MIN_VALUE + index
-                        index == focusedItemIdx -> 0
-                        else -> Int.MAX_VALUE - selectedItems.size + index + 1
-                    }
+                    order =
+                        when {
+                            index < focusedItemIdx -> Int.MIN_VALUE + index
+                            index == focusedItemIdx -> 0
+                            else -> Int.MAX_VALUE - selectedItems.size + index + 1
+                        },
                 )
             }
 }
diff --git a/java/src/com/android/intentresolver/contentpreview/payloadtoggle/shared/model/PreviewKey.kt b/java/src/com/android/intentresolver/contentpreview/payloadtoggle/shared/model/PreviewKey.kt
new file mode 100644
index 00000000..6b42133e
--- /dev/null
+++ b/java/src/com/android/intentresolver/contentpreview/payloadtoggle/shared/model/PreviewKey.kt
@@ -0,0 +1,49 @@
+/*
+ * Copyright (C) 2024 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.intentresolver.contentpreview.payloadtoggle.shared.model
+
+/** Unique identifier for preview items. */
+sealed interface PreviewKey {
+
+    private data class Temp(override val key: Int, override val isFinal: Boolean = false) :
+        PreviewKey
+
+    private data class Final(override val key: Int, override val isFinal: Boolean = true) :
+        PreviewKey
+
+    /** The identifier, must be unique among like keys types */
+    val key: Int
+    /** Whether this key is final or temporary. */
+    val isFinal: Boolean
+
+    companion object {
+        /**
+         * Creates a temporary key.
+         *
+         * This is used for the initial preview items until final keys can be generated, at which
+         * point it is replaced with a final key.
+         */
+        fun temp(key: Int): PreviewKey = Temp(key)
+
+        /**
+         * Creates a final key.
+         *
+         * This is used for all preview items other than the initial preview items.
+         */
+        fun final(key: Int): PreviewKey = Final(key)
+    }
+}
diff --git a/java/src/com/android/intentresolver/contentpreview/payloadtoggle/shared/model/PreviewModel.kt b/java/src/com/android/intentresolver/contentpreview/payloadtoggle/shared/model/PreviewModel.kt
index 8a479156..d4df8a3a 100644
--- a/java/src/com/android/intentresolver/contentpreview/payloadtoggle/shared/model/PreviewModel.kt
+++ b/java/src/com/android/intentresolver/contentpreview/payloadtoggle/shared/model/PreviewModel.kt
@@ -20,6 +20,8 @@ import android.net.Uri
 
 /** An individual preview presented in Shareousel. */
 data class PreviewModel(
+    /** Unique identifier for this model. */
+    val key: PreviewKey,
     /** Uri for this item; if this preview is selected, this will be shared with the target app. */
     val uri: Uri,
     /** Uri for the preview image. */
@@ -28,7 +30,8 @@ data class PreviewModel(
     val mimeType: String?,
     val aspectRatio: Float = 1f,
     /**
-     * Relative item position in the list that is used to determine items order in the target intent
+     * Relative item position in the list that is used to determine items order in the target
+     * intent.
      */
     val order: Int,
 )
diff --git a/java/src/com/android/intentresolver/contentpreview/payloadtoggle/ui/composable/ShareouselComposable.kt b/java/src/com/android/intentresolver/contentpreview/payloadtoggle/ui/composable/ShareouselComposable.kt
index 4b87d227..eab04aab 100644
--- a/java/src/com/android/intentresolver/contentpreview/payloadtoggle/ui/composable/ShareouselComposable.kt
+++ b/java/src/com/android/intentresolver/contentpreview/payloadtoggle/ui/composable/ShareouselComposable.kt
@@ -96,7 +96,7 @@ private fun Shareousel(viewModel: ShareouselViewModel, keySet: PreviewsModel) {
     Column(
         modifier =
             Modifier.background(MaterialTheme.colorScheme.surfaceContainer)
-                .padding(vertical = 16.dp),
+                .padding(vertical = 16.dp)
     ) {
         PreviewCarousel(keySet, viewModel)
         ActionCarousel(viewModel)
@@ -105,10 +105,7 @@ private fun Shareousel(viewModel: ShareouselViewModel, keySet: PreviewsModel) {
 
 @OptIn(ExperimentalFoundationApi::class)
 @Composable
-private fun PreviewCarousel(
-    previews: PreviewsModel,
-    viewModel: ShareouselViewModel,
-) {
+private fun PreviewCarousel(previews: PreviewsModel, viewModel: ShareouselViewModel) {
     var maxAspectRatio by remember { mutableStateOf(0f) }
     var viewportHeight by remember { mutableStateOf(0) }
     var viewportCenter by remember { mutableStateOf(0) }
@@ -128,7 +125,7 @@ private fun PreviewCarousel(
                             val maxAR =
                                 (maxItemWidth.toFloat() / placeable.height).coerceIn(
                                     0f,
-                                    MAX_ASPECT_RATIO
+                                    MAX_ASPECT_RATIO,
                                 )
                             minItemWidth to maxAR
                         }
@@ -137,7 +134,7 @@ private fun PreviewCarousel(
                     viewportHeight = placeable.height
                     horizontalPadding = ((placeable.width - minItemWidth) / 2).toDp()
                     layout(placeable.width, placeable.height) { placeable.place(0, 0) }
-                },
+                }
     ) {
         if (maxAspectRatio <= 0 && previews.previewModels.isNotEmpty()) {
             // Do not compose the list until we know the viewport size
@@ -148,7 +145,7 @@ private fun PreviewCarousel(
 
         val carouselState =
             rememberLazyListState(
-                prefetchStrategy = remember { ShareouselLazyListPrefetchStrategy() },
+                prefetchStrategy = remember { ShareouselLazyListPrefetchStrategy() }
             )
 
         LazyRow(
@@ -157,7 +154,10 @@ private fun PreviewCarousel(
             contentPadding = PaddingValues(start = horizontalPadding, end = horizontalPadding),
             modifier = Modifier.fillMaxSize().systemGestureExclusion(),
         ) {
-            itemsIndexed(previews.previewModels, key = { _, model -> model.uri }) { index, model ->
+            itemsIndexed(
+                items = previews.previewModels,
+                key = { _, model -> model.key.key to model.key.isFinal },
+            ) { index, model ->
                 val visibleItem by remember {
                     derivedStateOf {
                         carouselState.layoutInfo.visibleItemsInfo.firstOrNull { it.index == index }
@@ -234,7 +234,7 @@ private fun PreviewCarousel(
                         model,
                         viewportHeight,
                         previewIndex,
-                        rememberCoroutineScope()
+                        rememberCoroutineScope(),
                     ),
                     maxAspectRatio,
                 )
@@ -279,7 +279,7 @@ private fun ShareouselCard(viewModel: ShareouselPreviewViewModel, maxAspectRatio
                 .toggleable(
                     value = selected,
                     onValueChange = { scope.launch { viewModel.setSelected(it) } },
-                )
+                ),
     ) { state ->
         val aspectRatio = minOf(maxAspectRatio, maxOf(MIN_ASPECT_RATIO, viewModel.aspectRatio))
         if (state is ValueUpdate.Value) {
@@ -304,7 +304,7 @@ private fun ShareouselCard(viewModel: ShareouselPreviewViewModel, maxAspectRatio
                                 color = borderColor,
                                 shape = RoundedCornerShape(size = 12.dp),
                             )
-                        }
+                        },
                 )
             }
         } else {
@@ -355,7 +355,7 @@ private fun ActionCarousel(viewModel: ShareouselViewModel) {
                             Image(
                                 icon = it,
                                 modifier = Modifier.size(16.dp),
-                                colorFilter = ColorFilter.tint(LocalContentColor.current)
+                                colorFilter = ColorFilter.tint(LocalContentColor.current),
                             )
                         }
                     }
@@ -389,7 +389,7 @@ private fun ShareouselAction(
             AssistChipDefaults.assistChipColors(
                 containerColor = MaterialTheme.colorScheme.surfaceContainerHigh,
                 labelColor = MaterialTheme.colorScheme.onSurface,
-                leadingIconContentColor = MaterialTheme.colorScheme.onSurface
+                leadingIconContentColor = MaterialTheme.colorScheme.onSurface,
             ),
         modifier = modifier,
     )
diff --git a/tests/unit/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/interactor/CursorPreviewsInteractorTest.kt b/tests/unit/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/interactor/CursorPreviewsInteractorTest.kt
index c4ba8105..f43f1467 100644
--- a/tests/unit/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/interactor/CursorPreviewsInteractorTest.kt
+++ b/tests/unit/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/interactor/CursorPreviewsInteractorTest.kt
@@ -34,6 +34,7 @@ import com.android.intentresolver.contentpreview.payloadtoggle.data.repository.p
 import com.android.intentresolver.contentpreview.payloadtoggle.domain.intent.TargetIntentModifier
 import com.android.intentresolver.contentpreview.payloadtoggle.domain.intent.targetIntentModifier
 import com.android.intentresolver.contentpreview.payloadtoggle.domain.model.CursorRow
+import com.android.intentresolver.contentpreview.payloadtoggle.shared.model.PreviewKey
 import com.android.intentresolver.contentpreview.payloadtoggle.shared.model.PreviewModel
 import com.android.intentresolver.contentpreview.readSize
 import com.android.intentresolver.contentpreview.uriMetadataReader
@@ -51,10 +52,10 @@ import org.junit.Test
 class CursorPreviewsInteractorTest {
 
     private fun runTestWithDeps(
-        initialSelection: Iterable<Int> = (1..2),
-        focusedItemIndex: Int = initialSelection.count() / 2,
-        cursor: Iterable<Int> = (0 until 4),
-        cursorStartPosition: Int = cursor.count() / 2,
+        initialSelection: Iterable<Int>,
+        focusedItemIndex: Int,
+        cursor: Iterable<Int>,
+        cursorStartPosition: Int,
         pageSize: Int = 16,
         maxLoadedPages: Int = 3,
         cursorSizes: Map<Int, Size> = emptyMap(),
@@ -81,6 +82,7 @@ class CursorPreviewsInteractorTest {
                 block(
                     TestDeps(
                         initialSelection,
+                        focusedItemIndex,
                         cursor,
                         cursorStartPosition,
                         cursorSizes,
@@ -92,6 +94,7 @@ class CursorPreviewsInteractorTest {
 
     private class TestDeps(
         initialSelectionRange: Iterable<Int>,
+        focusedItemIndex: Int,
         private val cursorRange: Iterable<Int>,
         private val cursorStartPosition: Int,
         private val cursorSizes: Map<Int, Size>,
@@ -117,14 +120,26 @@ class CursorPreviewsInteractorTest {
                     }
                 }
         val initialPreviews: List<PreviewModel> =
-            initialSelectionRange.map { i ->
-                PreviewModel(uri = uri(i), mimeType = "image/bitmap", order = i)
+            initialSelectionRange.mapIndexed { index, i ->
+                PreviewModel(
+                    key =
+                        if (index == focusedItemIndex) {
+                            PreviewKey.final(0)
+                        } else {
+                            PreviewKey.temp(index)
+                        },
+                    uri = uri(i),
+                    mimeType = "image/bitmap",
+                    order = i,
+                )
             }
     }
 
     @Test
     fun initialCursorLoad() =
         runTestWithDeps(
+            initialSelection = (1..2),
+            focusedItemIndex = 1,
             cursor = (0 until 10),
             cursorStartPosition = 2,
             cursorSizes = mapOf(0 to (200 x 100)),
@@ -143,6 +158,7 @@ class CursorPreviewsInteractorTest {
                     .containsExactlyElementsIn(
                         List(6) {
                             PreviewModel(
+                                key = PreviewKey.final((it - 2)),
                                 uri = Uri.fromParts("scheme$it", "ssp$it", "fragment$it"),
                                 mimeType = "image/bitmap",
                                 aspectRatio =
@@ -168,7 +184,9 @@ class CursorPreviewsInteractorTest {
     fun loadMoreLeft_evictRight() =
         runTestWithDeps(
             initialSelection = listOf(24),
+            focusedItemIndex = 0,
             cursor = (0 until 48),
+            cursorStartPosition = 24,
             pageSize = 16,
             maxLoadedPages = 1,
         ) { deps ->
@@ -201,7 +219,9 @@ class CursorPreviewsInteractorTest {
     fun loadMoreRight_evictLeft() =
         runTestWithDeps(
             initialSelection = listOf(24),
+            focusedItemIndex = 0,
             cursor = (0 until 48),
+            cursorStartPosition = 24,
             pageSize = 16,
             maxLoadedPages = 1,
         ) { deps ->
@@ -233,7 +253,9 @@ class CursorPreviewsInteractorTest {
     fun noMoreRight_appendUnclaimedFromInitialSelection() =
         runTestWithDeps(
             initialSelection = listOf(24, 50),
+            focusedItemIndex = 0,
             cursor = listOf(24),
+            cursorStartPosition = 0,
             pageSize = 16,
             maxLoadedPages = 2,
         ) { deps ->
@@ -255,7 +277,9 @@ class CursorPreviewsInteractorTest {
     fun noMoreLeft_appendUnclaimedFromInitialSelection() =
         runTestWithDeps(
             initialSelection = listOf(0, 24),
+            focusedItemIndex = 1,
             cursor = listOf(24),
+            cursorStartPosition = 0,
             pageSize = 16,
             maxLoadedPages = 2,
         ) { deps ->
@@ -283,6 +307,7 @@ class CursorPreviewsInteractorTest {
         ) { deps ->
             previewSelectionsRepository.selections.value =
                 PreviewModel(
+                        key = PreviewKey.final(0),
                         uri = uri(1),
                         mimeType = "image/png",
                         order = 0,
@@ -296,6 +321,7 @@ class CursorPreviewsInteractorTest {
             assertThat(previewSelectionsRepository.selections.value.values)
                 .containsExactly(
                     PreviewModel(
+                        key = PreviewKey.final(0),
                         uri = uri(1),
                         mimeType = "image/bitmap",
                         order = 1,
@@ -307,6 +333,7 @@ class CursorPreviewsInteractorTest {
     fun testReadFailedPages() =
         runTestWithDeps(
             initialSelection = listOf(4),
+            focusedItemIndex = 0,
             cursor = emptyList(),
             cursorStartPosition = 0,
             pageSize = 2,
diff --git a/tests/unit/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/interactor/FetchPreviewsInteractorTest.kt b/tests/unit/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/interactor/FetchPreviewsInteractorTest.kt
index 27c98dc0..09d254f3 100644
--- a/tests/unit/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/interactor/FetchPreviewsInteractorTest.kt
+++ b/tests/unit/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/interactor/FetchPreviewsInteractorTest.kt
@@ -30,6 +30,7 @@ import com.android.intentresolver.contentpreview.payloadtoggle.domain.cursor.pay
 import com.android.intentresolver.contentpreview.payloadtoggle.domain.intent.TargetIntentModifier
 import com.android.intentresolver.contentpreview.payloadtoggle.domain.intent.targetIntentModifier
 import com.android.intentresolver.contentpreview.payloadtoggle.domain.model.CursorRow
+import com.android.intentresolver.contentpreview.payloadtoggle.shared.model.PreviewKey
 import com.android.intentresolver.contentpreview.payloadtoggle.shared.model.PreviewModel
 import com.android.intentresolver.contentpreview.payloadtoggle.shared.model.PreviewsModel
 import com.android.intentresolver.contentpreview.uriMetadataReader
@@ -50,10 +51,10 @@ import org.junit.Test
 class FetchPreviewsInteractorTest {
 
     private fun runTest(
-        initialSelection: Iterable<Int> = (1..2),
-        focusedItemIndex: Int = initialSelection.count() / 2,
-        cursor: Iterable<Int> = (0 until 4),
-        cursorStartPosition: Int = cursor.count() / 2,
+        initialSelection: Iterable<Int>,
+        focusedItemIndex: Int,
+        cursor: Iterable<Int>,
+        cursorStartPosition: Int,
         pageSize: Int = 16,
         maxLoadedPages: Int = 8,
         previewSizes: Map<Int, Size> = emptyMap(),
@@ -110,7 +111,11 @@ class FetchPreviewsInteractorTest {
     fun setsInitialPreviews() =
         runTest(
             initialSelection = (1..3),
-            previewSizes = mapOf(1 to Size(100, 50))) {
+            focusedItemIndex = 1,
+            cursor = (0 until 4),
+            cursorStartPosition = 1,
+            previewSizes = mapOf(1 to Size(100, 50)),
+        ) {
             backgroundScope.launch { fetchPreviewsInteractor.activate() }
             runCurrent()
 
@@ -120,17 +125,20 @@ class FetchPreviewsInteractorTest {
                         previewModels =
                             listOf(
                                 PreviewModel(
+                                    key = PreviewKey.temp(0),
                                     uri = Uri.fromParts("scheme1", "ssp1", "fragment1"),
                                     mimeType = "image/bitmap",
                                     aspectRatio = 2f,
                                     order = Int.MIN_VALUE,
                                 ),
                                 PreviewModel(
+                                    key = PreviewKey.final(0),
                                     uri = Uri.fromParts("scheme2", "ssp2", "fragment2"),
                                     mimeType = "image/bitmap",
                                     order = 0,
                                 ),
                                 PreviewModel(
+                                    key = PreviewKey.temp(2),
                                     uri = Uri.fromParts("scheme3", "ssp3", "fragment3"),
                                     mimeType = "image/bitmap",
                                     order = Int.MAX_VALUE,
@@ -146,48 +154,60 @@ class FetchPreviewsInteractorTest {
         }
 
     @Test
-    fun lookupCursorFromContentResolver() = runTest {
-        backgroundScope.launch { fetchPreviewsInteractor.activate() }
-        fakeCursorResolver.complete()
-        runCurrent()
+    fun lookupCursorFromContentResolver() =
+        runTest(
+            initialSelection = (1..2),
+            focusedItemIndex = 1,
+            cursor = (0 until 4),
+            cursorStartPosition = 2,
+        ) {
+            backgroundScope.launch { fetchPreviewsInteractor.activate() }
+            fakeCursorResolver.complete()
+            runCurrent()
 
-        with(cursorPreviewsRepository) {
-            assertThat(previewsModel.value).isNotNull()
-            assertThat(previewsModel.value!!.startIdx).isEqualTo(0)
-            assertThat(previewsModel.value!!.loadMoreLeft).isNull()
-            assertThat(previewsModel.value!!.loadMoreRight).isNull()
-            assertThat(previewsModel.value!!.previewModels)
-                .containsExactly(
-                    PreviewModel(
-                        uri = Uri.fromParts("scheme0", "ssp0", "fragment0"),
-                        mimeType = "image/bitmap",
-                        order = 0,
-                    ),
-                    PreviewModel(
-                        uri = Uri.fromParts("scheme1", "ssp1", "fragment1"),
-                        mimeType = "image/bitmap",
-                        order = 1,
-                    ),
-                    PreviewModel(
-                        uri = Uri.fromParts("scheme2", "ssp2", "fragment2"),
-                        mimeType = "image/bitmap",
-                        order = 2,
-                    ),
-                    PreviewModel(
-                        uri = Uri.fromParts("scheme3", "ssp3", "fragment3"),
-                        mimeType = "image/bitmap",
-                        order = 3,
-                    ),
-                )
-                .inOrder()
+            with(cursorPreviewsRepository) {
+                assertThat(previewsModel.value).isNotNull()
+                assertThat(previewsModel.value!!.startIdx).isEqualTo(0)
+                assertThat(previewsModel.value!!.loadMoreLeft).isNull()
+                assertThat(previewsModel.value!!.loadMoreRight).isNull()
+                assertThat(previewsModel.value!!.previewModels)
+                    .containsExactly(
+                        PreviewModel(
+                            key = PreviewKey.final(-2),
+                            uri = Uri.fromParts("scheme0", "ssp0", "fragment0"),
+                            mimeType = "image/bitmap",
+                            order = 0,
+                        ),
+                        PreviewModel(
+                            key = PreviewKey.final(-1),
+                            uri = Uri.fromParts("scheme1", "ssp1", "fragment1"),
+                            mimeType = "image/bitmap",
+                            order = 1,
+                        ),
+                        PreviewModel(
+                            key = PreviewKey.final(0),
+                            uri = Uri.fromParts("scheme2", "ssp2", "fragment2"),
+                            mimeType = "image/bitmap",
+                            order = 2,
+                        ),
+                        PreviewModel(
+                            key = PreviewKey.final(1),
+                            uri = Uri.fromParts("scheme3", "ssp3", "fragment3"),
+                            mimeType = "image/bitmap",
+                            order = 3,
+                        ),
+                    )
+                    .inOrder()
+            }
         }
-    }
 
     @Test
     fun loadMoreLeft_evictRight() =
         runTest(
             initialSelection = listOf(24),
+            focusedItemIndex = 0,
             cursor = (0 until 48),
+            cursorStartPosition = 24,
             pageSize = 16,
             maxLoadedPages = 1,
         ) {
@@ -223,7 +243,9 @@ class FetchPreviewsInteractorTest {
     fun loadMoreRight_evictLeft() =
         runTest(
             initialSelection = listOf(24),
+            focusedItemIndex = 0,
             cursor = (0 until 48),
+            cursorStartPosition = 24,
             pageSize = 16,
             maxLoadedPages = 1,
         ) {
@@ -254,7 +276,9 @@ class FetchPreviewsInteractorTest {
     fun noMoreRight_appendUnclaimedFromInitialSelection() =
         runTest(
             initialSelection = listOf(24, 50),
+            focusedItemIndex = 0,
             cursor = listOf(24),
+            cursorStartPosition = 0,
             pageSize = 16,
             maxLoadedPages = 2,
         ) {
@@ -275,7 +299,9 @@ class FetchPreviewsInteractorTest {
     fun noMoreLeft_appendUnclaimedFromInitialSelection() =
         runTest(
             initialSelection = listOf(0, 24),
+            focusedItemIndex = 1,
             cursor = listOf(24),
+            cursorStartPosition = 0,
             pageSize = 16,
             maxLoadedPages = 2,
         ) {
diff --git a/tests/unit/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/interactor/SelectablePreviewInteractorTest.kt b/tests/unit/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/interactor/SelectablePreviewInteractorTest.kt
index 5d9ddbb6..0268a4d5 100644
--- a/tests/unit/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/interactor/SelectablePreviewInteractorTest.kt
+++ b/tests/unit/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/interactor/SelectablePreviewInteractorTest.kt
@@ -24,6 +24,7 @@ import com.android.intentresolver.contentpreview.payloadtoggle.data.repository.p
 import com.android.intentresolver.contentpreview.payloadtoggle.data.repository.previewSelectionsRepository
 import com.android.intentresolver.contentpreview.payloadtoggle.domain.intent.TargetIntentModifier
 import com.android.intentresolver.contentpreview.payloadtoggle.domain.intent.targetIntentModifier
+import com.android.intentresolver.contentpreview.payloadtoggle.shared.model.PreviewKey
 import com.android.intentresolver.contentpreview.payloadtoggle.shared.model.PreviewModel
 import com.android.intentresolver.data.repository.chooserRequestRepository
 import com.android.intentresolver.logging.FakeEventLog
@@ -44,6 +45,7 @@ class SelectablePreviewInteractorTest {
             SelectablePreviewInteractor(
                 key =
                     PreviewModel(
+                        key = PreviewKey.final(1),
                         uri = Uri.fromParts("scheme", "ssp", "fragment"),
                         mimeType = null,
                         order = 0,
@@ -63,6 +65,7 @@ class SelectablePreviewInteractorTest {
             SelectablePreviewInteractor(
                 key =
                     PreviewModel(
+                        key = PreviewKey.final(1),
                         uri = Uri.fromParts("scheme", "ssp", "fragment"),
                         mimeType = "image/bitmap",
                         order = 0,
@@ -75,6 +78,7 @@ class SelectablePreviewInteractorTest {
 
         previewSelectionsRepository.selections.value =
             PreviewModel(
+                    key = PreviewKey.final(1),
                     uri = Uri.fromParts("scheme", "ssp", "fragment"),
                     mimeType = "image/bitmap",
                     order = 0,
@@ -93,6 +97,7 @@ class SelectablePreviewInteractorTest {
             SelectablePreviewInteractor(
                 key =
                     PreviewModel(
+                        key = PreviewKey.final(1),
                         uri = Uri.fromParts("scheme", "ssp", "fragment"),
                         mimeType = "image/bitmap",
                         order = 0,
diff --git a/tests/unit/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/interactor/SelectablePreviewsInteractorTest.kt b/tests/unit/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/interactor/SelectablePreviewsInteractorTest.kt
index c50d2d3f..c90a3091 100644
--- a/tests/unit/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/interactor/SelectablePreviewsInteractorTest.kt
+++ b/tests/unit/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/interactor/SelectablePreviewsInteractorTest.kt
@@ -23,6 +23,7 @@ import com.android.intentresolver.contentpreview.payloadtoggle.data.repository.c
 import com.android.intentresolver.contentpreview.payloadtoggle.data.repository.previewSelectionsRepository
 import com.android.intentresolver.contentpreview.payloadtoggle.domain.intent.TargetIntentModifier
 import com.android.intentresolver.contentpreview.payloadtoggle.domain.intent.targetIntentModifier
+import com.android.intentresolver.contentpreview.payloadtoggle.shared.model.PreviewKey
 import com.android.intentresolver.contentpreview.payloadtoggle.shared.model.PreviewModel
 import com.android.intentresolver.contentpreview.payloadtoggle.shared.model.PreviewsModel
 import com.android.intentresolver.util.runKosmosTest
@@ -41,11 +42,13 @@ class SelectablePreviewsInteractorTest {
                 previewModels =
                     listOf(
                         PreviewModel(
+                            key = PreviewKey.final(1),
                             uri = Uri.fromParts("scheme", "ssp", "fragment"),
                             mimeType = "image/bitmap",
                             order = 0,
                         ),
                         PreviewModel(
+                            key = PreviewKey.final(2),
                             uri = Uri.fromParts("scheme2", "ssp2", "fragment2"),
                             mimeType = "image/bitmap",
                             order = 1,
@@ -59,6 +62,7 @@ class SelectablePreviewsInteractorTest {
             )
         previewSelectionsRepository.selections.value =
             PreviewModel(
+                    key = PreviewKey.final(1),
                     uri = Uri.fromParts("scheme", "ssp", "fragment"),
                     mimeType = null,
                     order = 0,
@@ -72,11 +76,13 @@ class SelectablePreviewsInteractorTest {
         assertThat(keySet.value!!.previewModels)
             .containsExactly(
                 PreviewModel(
+                    key = PreviewKey.final(1),
                     uri = Uri.fromParts("scheme", "ssp", "fragment"),
                     mimeType = "image/bitmap",
                     order = 0,
                 ),
                 PreviewModel(
+                    key = PreviewKey.final(2),
                     uri = Uri.fromParts("scheme2", "ssp2", "fragment2"),
                     mimeType = "image/bitmap",
                     order = 1,
@@ -90,6 +96,7 @@ class SelectablePreviewsInteractorTest {
         val firstModel =
             underTest.preview(
                 PreviewModel(
+                    key = PreviewKey.final(1),
                     uri = Uri.fromParts("scheme", "ssp", "fragment"),
                     mimeType = null,
                     order = 0,
@@ -100,6 +107,7 @@ class SelectablePreviewsInteractorTest {
         val secondModel =
             underTest.preview(
                 PreviewModel(
+                    key = PreviewKey.final(2),
                     uri = Uri.fromParts("scheme2", "ssp2", "fragment2"),
                     mimeType = null,
                     order = 1,
@@ -112,6 +120,7 @@ class SelectablePreviewsInteractorTest {
     fun keySet_reflectsRepositoryUpdate() = runKosmosTest {
         previewSelectionsRepository.selections.value =
             PreviewModel(
+                    key = PreviewKey.final(1),
                     uri = Uri.fromParts("scheme", "ssp", "fragment"),
                     mimeType = null,
                     order = 0,
@@ -124,6 +133,7 @@ class SelectablePreviewsInteractorTest {
         val firstModel =
             underTest.preview(
                 PreviewModel(
+                    key = PreviewKey.final(1),
                     uri = Uri.fromParts("scheme", "ssp", "fragment"),
                     mimeType = null,
                     order = 0,
@@ -140,11 +150,13 @@ class SelectablePreviewsInteractorTest {
                 previewModels =
                     listOf(
                         PreviewModel(
+                            key = PreviewKey.final(1),
                             uri = Uri.fromParts("scheme", "ssp", "fragment"),
                             mimeType = "image/bitmap",
                             order = 0,
                         ),
                         PreviewModel(
+                            key = PreviewKey.final(2),
                             uri = Uri.fromParts("scheme2", "ssp2", "fragment2"),
                             mimeType = "image/bitmap",
                             order = 1,
@@ -163,11 +175,13 @@ class SelectablePreviewsInteractorTest {
         assertThat(previews.value!!.previewModels)
             .containsExactly(
                 PreviewModel(
+                    key = PreviewKey.final(1),
                     uri = Uri.fromParts("scheme", "ssp", "fragment"),
                     mimeType = "image/bitmap",
                     order = 0,
                 ),
                 PreviewModel(
+                    key = PreviewKey.final(2),
                     uri = Uri.fromParts("scheme2", "ssp2", "fragment2"),
                     mimeType = "image/bitmap",
                     order = 1,
diff --git a/tests/unit/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/interactor/SelectionInteractorTest.kt b/tests/unit/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/interactor/SelectionInteractorTest.kt
index c8242333..c24138b8 100644
--- a/tests/unit/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/interactor/SelectionInteractorTest.kt
+++ b/tests/unit/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/interactor/SelectionInteractorTest.kt
@@ -24,6 +24,7 @@ import android.platform.test.flag.junit.SetFlagsRule
 import com.android.intentresolver.Flags
 import com.android.intentresolver.contentpreview.mimetypeClassifier
 import com.android.intentresolver.contentpreview.payloadtoggle.data.repository.previewSelectionsRepository
+import com.android.intentresolver.contentpreview.payloadtoggle.shared.model.PreviewKey
 import com.android.intentresolver.contentpreview.payloadtoggle.shared.model.PreviewModel
 import com.android.intentresolver.util.runKosmosTest
 import com.google.common.truth.Truth.assertThat
@@ -39,9 +40,10 @@ class SelectionInteractorTest {
     fun singleSelection_removalPrevented() = runKosmosTest {
         val initialPreview =
             PreviewModel(
+                key = PreviewKey.final(1),
                 uri = Uri.fromParts("scheme", "ssp", "fragment"),
                 mimeType = null,
-                order = 0
+                order = 0,
             )
         previewSelectionsRepository.selections.value = mapOf(initialPreview.uri to initialPreview)
 
@@ -66,9 +68,10 @@ class SelectionInteractorTest {
     fun singleSelection_itemRemovedNoPendingIntentUpdates() = runKosmosTest {
         val initialPreview =
             PreviewModel(
+                key = PreviewKey.final(1),
                 uri = Uri.fromParts("scheme", "ssp", "fragment"),
                 mimeType = null,
-                order = 0
+                order = 0,
             )
         previewSelectionsRepository.selections.value = mapOf(initialPreview.uri to initialPreview)
 
@@ -92,15 +95,17 @@ class SelectionInteractorTest {
     fun multipleSelections_removalAllowed() = runKosmosTest {
         val first =
             PreviewModel(
+                key = PreviewKey.final(1),
                 uri = Uri.fromParts("scheme", "ssp", "fragment"),
                 mimeType = null,
-                order = 0
+                order = 0,
             )
         val second =
             PreviewModel(
+                key = PreviewKey.final(2),
                 uri = Uri.fromParts("scheme2", "ssp2", "fragment2"),
                 mimeType = null,
-                order = 1
+                order = 1,
             )
         previewSelectionsRepository.selections.value = listOf(first, second).associateBy { it.uri }
 
@@ -109,7 +114,7 @@ class SelectionInteractorTest {
                 previewSelectionsRepository,
                 { Intent() },
                 updateTargetIntentInteractor,
-                mimetypeClassifier
+                mimetypeClassifier,
             )
 
         underTest.unselect(first)
diff --git a/tests/unit/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/interactor/SetCursorPreviewsInteractorTest.kt b/tests/unit/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/interactor/SetCursorPreviewsInteractorTest.kt
index 748459cb..42f1a1b2 100644
--- a/tests/unit/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/interactor/SetCursorPreviewsInteractorTest.kt
+++ b/tests/unit/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/interactor/SetCursorPreviewsInteractorTest.kt
@@ -21,6 +21,7 @@ package com.android.intentresolver.contentpreview.payloadtoggle.domain.interacto
 import android.net.Uri
 import com.android.intentresolver.contentpreview.payloadtoggle.data.repository.cursorPreviewsRepository
 import com.android.intentresolver.contentpreview.payloadtoggle.domain.model.LoadDirection
+import com.android.intentresolver.contentpreview.payloadtoggle.shared.model.PreviewKey
 import com.android.intentresolver.contentpreview.payloadtoggle.shared.model.PreviewModel
 import com.android.intentresolver.util.runKosmosTest
 import com.google.common.truth.Truth.assertThat
@@ -37,6 +38,7 @@ class SetCursorPreviewsInteractorTest {
                 previews =
                     listOf(
                         PreviewModel(
+                            key = PreviewKey.final(1),
                             uri = Uri.fromParts("scheme", "ssp", "fragment"),
                             mimeType = null,
                             order = 0,
@@ -59,9 +61,10 @@ class SetCursorPreviewsInteractorTest {
             assertThat(it.previewModels)
                 .containsExactly(
                     PreviewModel(
+                        key = PreviewKey.final(1),
                         uri = Uri.fromParts("scheme", "ssp", "fragment"),
                         mimeType = null,
-                        order = 0
+                        order = 0,
                     )
                 )
                 .inOrder()
@@ -76,6 +79,7 @@ class SetCursorPreviewsInteractorTest {
                     previews =
                         listOf(
                             PreviewModel(
+                                key = PreviewKey.final(1),
                                 uri = Uri.fromParts("scheme", "ssp", "fragment"),
                                 mimeType = null,
                                 order = 0,
diff --git a/tests/unit/src/com/android/intentresolver/contentpreview/payloadtoggle/ui/viewmodel/ShareouselViewModelTest.kt b/tests/unit/src/com/android/intentresolver/contentpreview/payloadtoggle/ui/viewmodel/ShareouselViewModelTest.kt
index fc7ac751..6dd96040 100644
--- a/tests/unit/src/com/android/intentresolver/contentpreview/payloadtoggle/ui/viewmodel/ShareouselViewModelTest.kt
+++ b/tests/unit/src/com/android/intentresolver/contentpreview/payloadtoggle/ui/viewmodel/ShareouselViewModelTest.kt
@@ -42,6 +42,7 @@ import com.android.intentresolver.contentpreview.payloadtoggle.domain.interactor
 import com.android.intentresolver.contentpreview.payloadtoggle.domain.interactor.selectionInteractor
 import com.android.intentresolver.contentpreview.payloadtoggle.domain.model.ValueUpdate
 import com.android.intentresolver.contentpreview.payloadtoggle.shared.ContentType
+import com.android.intentresolver.contentpreview.payloadtoggle.shared.model.PreviewKey
 import com.android.intentresolver.contentpreview.payloadtoggle.shared.model.PreviewModel
 import com.android.intentresolver.contentpreview.payloadtoggle.shared.model.PreviewsModel
 import com.android.intentresolver.data.model.ChooserRequest
@@ -84,15 +85,17 @@ class ShareouselViewModelTest {
         previewSelectionsRepository.selections.value =
             listOf(
                     PreviewModel(
+                        key = PreviewKey.final(0),
                         uri = Uri.fromParts("scheme", "ssp", "fragment"),
                         mimeType = "image/png",
                         order = 0,
                     ),
                     PreviewModel(
+                        key = PreviewKey.final(1),
                         uri = Uri.fromParts("scheme1", "ssp1", "fragment1"),
                         mimeType = "image/jpeg",
                         order = 1,
-                    )
+                    ),
                 )
                 .associateBy { it.uri }
         runCurrent()
@@ -104,15 +107,17 @@ class ShareouselViewModelTest {
         previewSelectionsRepository.selections.value =
             listOf(
                     PreviewModel(
+                        key = PreviewKey.final(0),
                         uri = Uri.fromParts("scheme", "ssp", "fragment"),
                         mimeType = "video/mpeg",
                         order = 0,
                     ),
                     PreviewModel(
+                        key = PreviewKey.final(1),
                         uri = Uri.fromParts("scheme1", "ssp1", "fragment1"),
                         mimeType = "video/mpeg",
                         order = 1,
-                    )
+                    ),
                 )
                 .associateBy { it.uri }
         runCurrent()
@@ -124,15 +129,17 @@ class ShareouselViewModelTest {
         previewSelectionsRepository.selections.value =
             listOf(
                     PreviewModel(
+                        key = PreviewKey.final(0),
                         uri = Uri.fromParts("scheme", "ssp", "fragment"),
                         mimeType = "image/jpeg",
                         order = 0,
                     ),
                     PreviewModel(
+                        key = PreviewKey.final(1),
                         uri = Uri.fromParts("scheme1", "ssp1", "fragment1"),
                         mimeType = "video/mpeg",
                         order = 1,
-                    )
+                    ),
                 )
                 .associateBy { it.uri }
         runCurrent()
@@ -145,7 +152,7 @@ class ShareouselViewModelTest {
             ChooserRequest(
                 targetIntent = Intent(),
                 launchedFromPackage = "",
-                metadataText = "Hello"
+                metadataText = "Hello",
             )
         chooserRequestRepository.chooserRequest.value = request
 
@@ -162,15 +169,17 @@ class ShareouselViewModelTest {
                     previewModels =
                         listOf(
                             PreviewModel(
+                                key = PreviewKey.final(0),
                                 uri = Uri.fromParts("scheme", "ssp", "fragment"),
                                 mimeType = "image/png",
                                 order = 0,
                             ),
                             PreviewModel(
+                                key = PreviewKey.final(1),
                                 uri = Uri.fromParts("scheme1", "ssp1", "fragment1"),
                                 mimeType = "video/mpeg",
                                 order = 1,
-                            )
+                            ),
                         ),
                     startIdx = 1,
                     loadMoreLeft = null,
@@ -194,6 +203,7 @@ class ShareouselViewModelTest {
             val previewVm =
                 shareouselViewModel.preview.invoke(
                     PreviewModel(
+                        key = PreviewKey.final(1),
                         uri = Uri.fromParts("scheme1", "ssp1", "fragment1"),
                         mimeType = "video/mpeg",
                         order = 0,
@@ -225,15 +235,17 @@ class ShareouselViewModelTest {
                     previewModels =
                         listOf(
                             PreviewModel(
+                                key = PreviewKey.final(0),
                                 uri = Uri.fromParts("scheme", "ssp", "fragment"),
                                 mimeType = "image/png",
                                 order = 0,
                             ),
                             PreviewModel(
+                                key = PreviewKey.final(1),
                                 uri = Uri.fromParts("scheme1", "ssp1", "fragment1"),
                                 mimeType = "video/mpeg",
                                 order = 1,
-                            )
+                            ),
                         ),
                     startIdx = 1,
                     loadMoreLeft = null,
@@ -246,6 +258,7 @@ class ShareouselViewModelTest {
             val previewVm =
                 shareouselViewModel.preview.invoke(
                     PreviewModel(
+                        key = PreviewKey.final(0),
                         uri = Uri.fromParts("scheme", "ssp", "fragment"),
                         mimeType = "video/mpeg",
                         order = 1,
@@ -314,6 +327,7 @@ class ShareouselViewModelTest {
         this.targetIntentModifier = targetIntentModifier
         previewSelectionsRepository.selections.value =
             PreviewModel(
+                    key = PreviewKey.final(1),
                     uri = Uri.fromParts("scheme", "ssp", "fragment"),
                     mimeType = null,
                     order = 0,