diff --git a/features/messages/src/main/java/io/element/android/x/features/messages/MessageTimelineItemStateMapper.kt b/features/messages/src/main/java/io/element/android/x/features/messages/MessageTimelineItemStateFactory.kt
similarity index 97%
rename from features/messages/src/main/java/io/element/android/x/features/messages/MessageTimelineItemStateMapper.kt
rename to features/messages/src/main/java/io/element/android/x/features/messages/MessageTimelineItemStateFactory.kt
index 3b6ccd2de7..0acaddf336 100644
--- a/features/messages/src/main/java/io/element/android/x/features/messages/MessageTimelineItemStateMapper.kt
+++ b/features/messages/src/main/java/io/element/android/x/features/messages/MessageTimelineItemStateFactory.kt
@@ -20,14 +20,15 @@ import org.matrix.rustcomponents.sdk.MessageFormat
 import org.matrix.rustcomponents.sdk.MessageType
 import org.matrix.rustcomponents.sdk.TimelineKey
 
-class MessageTimelineItemStateMapper(
+class MessageTimelineItemStateFactory(
     private val client: MatrixClient,
     private val room: MatrixRoom,
     private val dispatcher: CoroutineDispatcher,
 ) {
-    var highlightedEventId: String? = null
-
-    suspend fun map(timelineItems: List<MatrixTimelineItem>): List<MessagesTimelineItemState> =
+    suspend fun create(
+        timelineItems: List<MatrixTimelineItem>,
+        highlightedEventId: String? = null,
+    ): List<MessagesTimelineItemState> =
         withContext(dispatcher) {
             val messagesTimelineItemState = ArrayList<MessagesTimelineItemState>()
             for (index in timelineItems.indices.reversed()) {
diff --git a/features/messages/src/main/java/io/element/android/x/features/messages/MessagesViewModel.kt b/features/messages/src/main/java/io/element/android/x/features/messages/MessagesViewModel.kt
index 13bf229d2d..ec92c1e843 100644
--- a/features/messages/src/main/java/io/element/android/x/features/messages/MessagesViewModel.kt
+++ b/features/messages/src/main/java/io/element/android/x/features/messages/MessagesViewModel.kt
@@ -28,7 +28,7 @@ class MessagesViewModel(
     private val client: MatrixClient,
     private val room: MatrixRoom,
     private val timeline: MatrixTimeline,
-    private val messageTimelineItemStateMapper: MessageTimelineItemStateMapper,
+    private val messageTimelineItemStateFactory: MessageTimelineItemStateFactory,
     private val initialState: MessagesViewState
 ) :
     MavericksViewModel<MessagesViewState>(initialState) {
@@ -42,13 +42,13 @@ class MessagesViewModel(
             val matrix = MatrixInstance.getInstance()
             val client = matrix.activeClient()
             val room = client.getRoom(state.roomId) ?: return null
-            val messageTimelineItemStateMapper =
-                MessageTimelineItemStateMapper(client, room, Dispatchers.Default)
+            val messageTimelineItemStateFactory =
+                MessageTimelineItemStateFactory(client, room, Dispatchers.Default)
             return MessagesViewModel(
                 client,
                 room,
                 room.timeline(),
-                messageTimelineItemStateMapper,
+                messageTimelineItemStateFactory,
                 state
             )
         }
@@ -191,20 +191,18 @@ class MessagesViewModel(
 
         combine(
             timeline.timelineItems(),
-            stateFlow
-                .map { it.composerMode }
-                .distinctUntilChanged()
-        ) { timelineItems, messageComposerMode ->
-            // Set the highlightedEventId to messageTimelineItemStateMapper, before the mapping occurs
-            messageTimelineItemStateMapper.highlightedEventId = when (messageComposerMode) {
-                is MessageComposerMode.Normal -> null
-                is MessageComposerMode.Edit -> messageComposerMode.eventId
-                is MessageComposerMode.Quote -> null
-                is MessageComposerMode.Reply -> messageComposerMode.eventId
+            stateFlow.map {
+                when (it.composerMode) {
+                    is MessageComposerMode.Normal -> null
+                    is MessageComposerMode.Edit -> it.composerMode.eventId
+                    is MessageComposerMode.Quote -> null
+                    is MessageComposerMode.Reply -> it.composerMode.eventId
+                }
             }
-            timelineItems
+                .distinctUntilChanged()
+        ) { timelineItems, highlightedEventId ->
+            messageTimelineItemStateFactory.create(timelineItems, highlightedEventId)
         }
-            .map(messageTimelineItemStateMapper::map)
             .execute {
                 copy(timelineItems = it)
             }