ui-common: Restore mention state when editing a message

stream-chat-android-ui-common/src/main/kotlin/io/getstream/chat/android/ui/common/feature/messages/composer/MessageComposerController.kt

@@ -489,10 +489,24 @@ public class MessageComposerController(
         setMessageInputInternal(message.text, MessageInput.Source.Edit)
         _editModeMessage.value = fullMessage
         _editModeAttachments.value = attachments
+        restoreSelectedMentions(fullMessage)
         _messageActions.update { it + Edit(fullMessage) }
         syncAttachments()
     }
 
+    private fun restoreSelectedMentions(message: Message) {
+        val mentions = mutableSetOf<Mention>()
+        if (message.mentionedChannel) mentions += Mention.Channel
+        if (message.mentionedHere) mentions += Mention.Here
+        message.mentionedUsers.mapTo(mentions, Mention::User)
+        message.mentionedRoles.mapTo(mentions, Mention::Role)
+        message.mentionedGroups.mapTo(mentions, Mention::Group)
+
+        selectedMentions.clear()
+        selectedMentions += mentions
+        _state.update { it.copy(selectedMentions = selectedMentions.toSet()) }
+    }
+
     /**
      * Called when the input changes and the internal state needs to be updated.
      *
@@ -639,6 +653,7 @@ public class MessageComposerController(
                 setMessageInputInternal(messageAction.message.text, MessageInput.Source.Edit)
                 _editModeMessage.value = messageAction.message
                 _editModeAttachments.value = messageAction.message.attachments
+                restoreSelectedMentions(messageAction.message)
                 _messageActions.update { it + messageAction }
                 syncAttachments()
             }
@@ -655,6 +670,8 @@ public class MessageComposerController(
             setMessageInputInternal("", MessageInput.Source.Default)
             _editModeMessage.value = null
             _editModeAttachments.value = emptyList()
+            selectedMentions.clear()
+            _state.update { it.copy(selectedMentions = emptySet()) }
             syncAttachments()
         }
 

stream-chat-android-ui-common/src/test/kotlin/io/getstream/chat/android/ui/common/feature/messages/composer/MessageComposerControllerTest.kt

@@ -1735,6 +1735,96 @@ internal class MessageComposerControllerTest {
         assertTrue(message.mentionedGroups.isEmpty())
     }
 
+    @Test
+    fun `Given edit mode for message with multiple mentions When one mention is removed Then the others are preserved`() = runTest {
+        // Given
+        val controller = Fixture()
+            .givenAppSettings(mock())
+            .givenAudioPlayer(mock())
+            .givenClientState(randomUser())
+            .givenGlobalState()
+            .givenChannelState()
+            .get()
+        val alice = User(id = "u1", name = "Alice")
+        val bob = User(id = "u2", name = "Bob")
+        val editedMessage = randomMessage(
+            cid = CID,
+            text = "Hey @Alice and @Bob",
+            mentionedUsers = listOf(alice, bob),
+            mentionedChannel = false,
+            mentionedHere = false,
+        )
+        controller.performMessageAction(Edit(editedMessage))
+
+        // When — only @Bob remains in the edited text
+        val message = controller.buildNewMessage("Hey there and @Bob")
+
+        // Then
+        assertEquals(listOf("u2"), message.mentionedUsersIds)
+    }
+
+    @Test
+    fun `Given edit mode for message with mentions of every type When text unchanged Then all mentions are preserved`() = runTest {
+        // Given
+        val controller = Fixture()
+            .givenAppSettings(mock())
+            .givenAudioPlayer(mock())
+            .givenClientState(randomUser())
+            .givenGlobalState()
+            .givenChannelState()
+            .get()
+        val user = User(id = "u1", name = "Alice")
+        val group = UserGroup(id = "g1", name = "platform")
+        val editedMessage = randomMessage(
+            cid = CID,
+            text = "@Alice @channel @here @admin @platform hi",
+            mentionedUsers = listOf(user),
+            mentionedChannel = true,
+            mentionedHere = true,
+            mentionedRoles = listOf("admin"),
+            mentionedGroups = listOf(group),
+        )
+        controller.performMessageAction(Edit(editedMessage))
+
+        // When
+        val message = controller.buildNewMessage("@Alice @channel @here @admin @platform hi")
+
+        // Then
+        assertEquals(listOf("u1"), message.mentionedUsersIds)
+        assertTrue(message.mentionedChannel)
+        assertTrue(message.mentionedHere)
+        assertEquals(listOf("admin"), message.mentionedRoles)
+        assertEquals(listOf(group), message.mentionedGroups)
+    }
+
+    @Test
+    fun `Given edit mode mentions When edit dismissed Then mentions do not leak into next message`() = runTest {
+        // Given
+        val controller = Fixture()
+            .givenAppSettings(mock())
+            .givenAudioPlayer(mock())
+            .givenClientState(randomUser())
+            .givenGlobalState()
+            .givenChannelState()
+            .get()
+        val alice = User(id = "u1", name = "Alice")
+        val editedMessage = randomMessage(
+            cid = CID,
+            text = "Hey @Alice",
+            mentionedUsers = listOf(alice),
+            mentionedChannel = false,
+            mentionedHere = false,
+        )
+        controller.performMessageAction(Edit(editedMessage))
+
+        // When
+        controller.dismissMessageActions()
+        val message = controller.buildNewMessage("Hey @Alice")
+
+        // Then — the dismissed edit must not carry mentions into a fresh message
+        assertTrue(message.mentionedUsersIds.isEmpty())
+    }
+
     @Test
     fun `Given an active command When clearData called Then activeCommand is null`() = runTest {
         // Given