Merge pull request #3132 from element-hq/feature/fga/draft_in_memory_when_editing

Draft : add volatile storage when moving to edit mode.

features/messages/impl/src/main/kotlin/io/element/android/features/messages/impl/draft/ComposerDraftService.kt

@@ -20,6 +20,6 @@ import io.element.android.libraries.matrix.api.core.RoomId
 import io.element.android.libraries.matrix.api.room.draft.ComposerDraft
 
 interface ComposerDraftService {
-    suspend fun loadDraft(roomId: RoomId): ComposerDraft?
+    suspend fun loadDraft(roomId: RoomId, isVolatile: Boolean): ComposerDraft?
-    suspend fun updateDraft(roomId: RoomId, draft: ComposerDraft?)
+    suspend fun updateDraft(roomId: RoomId, draft: ComposerDraft?, isVolatile: Boolean)
 }

features/messages/impl/src/main/kotlin/io/element/android/features/messages/impl/draft/ComposerDraftStore.kt

@@ -0,0 +1,25 @@
+/*
+ * Copyright (c) 2024 New Vector Ltd
+ *
+ * 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
+ *
+ *     https://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 io.element.android.features.messages.impl.draft
+
+import io.element.android.libraries.matrix.api.core.RoomId
+import io.element.android.libraries.matrix.api.room.draft.ComposerDraft
+
+interface ComposerDraftStore {
+    suspend fun loadDraft(roomId: RoomId): ComposerDraft?
+    suspend fun updateDraft(roomId: RoomId, draft: ComposerDraft?)
+}

features/messages/impl/src/main/kotlin/io/element/android/features/messages/impl/draft/DefaultComposerDraftService.kt

@@ -18,44 +18,28 @@ package io.element.android.features.messages.impl.draft
 
 import com.squareup.anvil.annotations.ContributesBinding
 import io.element.android.libraries.di.RoomScope
-import io.element.android.libraries.matrix.api.MatrixClient
 import io.element.android.libraries.matrix.api.core.RoomId
 import io.element.android.libraries.matrix.api.room.draft.ComposerDraft
-import timber.log.Timber
 import javax.inject.Inject
 
 @ContributesBinding(RoomScope::class)
 class DefaultComposerDraftService @Inject constructor(
-    private val client: MatrixClient,
+    private val volatileComposerDraftStore: VolatileComposerDraftStore,
+    private val matrixComposerDraftStore: MatrixComposerDraftStore,
 ) : ComposerDraftService {
-    override suspend fun loadDraft(roomId: RoomId): ComposerDraft? {
+    override suspend fun loadDraft(roomId: RoomId, isVolatile: Boolean): ComposerDraft? {
-        return client.getRoom(roomId)?.use { room ->
+        return getStore(isVolatile).loadDraft(roomId)
-            room.loadComposerDraft()
+    }
-                .onFailure {
+
-                    Timber.e(it, "Failed to load composer draft for room $roomId")
+    override suspend fun updateDraft(roomId: RoomId, draft: ComposerDraft?, isVolatile: Boolean) {
-                }
+        getStore(isVolatile).updateDraft(roomId, draft)
-                .onSuccess { draft ->
-                    room.clearComposerDraft()
-                    Timber.d("Loaded composer draft for room $roomId : $draft")
-                }
-                .getOrNull()
-        }
     }
 
-    override suspend fun updateDraft(roomId: RoomId, draft: ComposerDraft?) {
+    private fun getStore(isVolatile: Boolean): ComposerDraftStore {
-        client.getRoom(roomId)?.use { room ->
+        return if (isVolatile) {
-            val updateDraftResult = if (draft == null) {
+            volatileComposerDraftStore
-                room.clearComposerDraft()
+        } else {
-            } else {
+            matrixComposerDraftStore
-                room.saveComposerDraft(draft)
-            }
-            updateDraftResult
-                .onFailure {
-                    Timber.e(it, "Failed to update composer draft for room $roomId")
-                }
-                .onSuccess {
-                    Timber.d("Updated composer draft for room $roomId")
-                }
         }
     }
 }

features/messages/impl/src/main/kotlin/io/element/android/features/messages/impl/draft/MatrixComposerDraftStore.kt

@@ -0,0 +1,62 @@
+/*
+ * Copyright (c) 2024 New Vector Ltd
+ *
+ * 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
+ *
+ *     https://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 io.element.android.features.messages.impl.draft
+
+import io.element.android.libraries.matrix.api.MatrixClient
+import io.element.android.libraries.matrix.api.core.RoomId
+import io.element.android.libraries.matrix.api.room.draft.ComposerDraft
+import timber.log.Timber
+import javax.inject.Inject
+
+/**
+ * A draft store that persists drafts in the room state.
+ * It can be used to store drafts that should be persisted across app restarts.
+ */
+class MatrixComposerDraftStore @Inject constructor(
+    private val client: MatrixClient,
+) : ComposerDraftStore {
+    override suspend fun loadDraft(roomId: RoomId): ComposerDraft? {
+        return client.getRoom(roomId)?.use { room ->
+            room.loadComposerDraft()
+                .onFailure {
+                    Timber.e(it, "Failed to load composer draft for room $roomId")
+                }
+                .onSuccess { draft ->
+                    room.clearComposerDraft()
+                    Timber.d("Loaded composer draft for room $roomId : $draft")
+                }
+                .getOrNull()
+        }
+    }
+
+    override suspend fun updateDraft(roomId: RoomId, draft: ComposerDraft?) {
+        client.getRoom(roomId)?.use { room ->
+            val updateDraftResult = if (draft == null) {
+                room.clearComposerDraft()
+            } else {
+                room.saveComposerDraft(draft)
+            }
+            updateDraftResult
+                .onFailure {
+                    Timber.e(it, "Failed to update composer draft for room $roomId")
+                }
+                .onSuccess {
+                    Timber.d("Updated composer draft for room $roomId")
+                }
+        }
+    }
+}

features/messages/impl/src/main/kotlin/io/element/android/features/messages/impl/draft/VolatileComposerDraftStore.kt

@@ -0,0 +1,43 @@
+/*
+ * Copyright (c) 2024 New Vector Ltd
+ *
+ * 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
+ *
+ *     https://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 io.element.android.features.messages.impl.draft
+
+import io.element.android.libraries.matrix.api.core.RoomId
+import io.element.android.libraries.matrix.api.room.draft.ComposerDraft
+import javax.inject.Inject
+
+/**
+ * A volatile draft store that keeps drafts in memory only.
+ * It can be used to store drafts that should not be persisted across app restarts.
+ * Currently it's used to store draft message when moving to edit mode.
+ */
+class VolatileComposerDraftStore @Inject constructor() : ComposerDraftStore {
+    private val drafts: MutableMap<RoomId, ComposerDraft> = mutableMapOf()
+
+    override suspend fun loadDraft(roomId: RoomId): ComposerDraft? {
+        // Remove the draft from the map when it is loaded
+        return drafts.remove(roomId)
+    }
+
+    override suspend fun updateDraft(roomId: RoomId, draft: ComposerDraft?) {
+        if (draft == null) {
+            drafts.remove(roomId)
+        } else {
+            drafts[roomId] = draft
+        }
+    }
+}

features/messages/impl/src/main/kotlin/io/element/android/features/messages/impl/messagecomposer/MessageComposerPresenter.kt

@@ -253,7 +253,10 @@ class MessageComposerPresenter @Inject constructor(
         )
 
         LaunchedEffect(Unit) {
-            loadDraft(markdownTextEditorState, richTextEditorState)
+            val draft = draftService.loadDraft(room.roomId, isVolatile = false)
+            if (draft != null) {
+                applyDraft(draft, markdownTextEditorState, richTextEditorState)
+            }
         }
 
         val mentionSpanProvider = LocalMentionSpanProvider.current
@@ -264,26 +267,16 @@ class MessageComposerPresenter @Inject constructor(
                 MessageComposerEvents.CloseSpecialMode -> {
                     if (messageComposerContext.composerMode is MessageComposerMode.Edit) {
                         localCoroutineScope.launch {
-                            textEditorState.reset()
+                            resetComposer(markdownTextEditorState, richTextEditorState, fromEdit = true)
                         }
+                    } else {
+                        messageComposerContext.composerMode = MessageComposerMode.Normal
                     }
-                    messageComposerContext.composerMode = MessageComposerMode.Normal
                 }
                 is MessageComposerEvents.SendMessage -> {
-                    val html = if (showTextFormatting) {
-                        richTextEditorState.messageHtml
-                    } else {
-                        null
-                    }
-                    val markdown = if (showTextFormatting) {
-                        richTextEditorState.messageMarkdown
-                    } else {
-                        markdownTextEditorState.getMessageMarkdown(permalinkBuilder)
-                    }
                     appCoroutineScope.sendMessage(
-                        message = Message(html = html, markdown = markdown),
+                        markdownTextEditorState = markdownTextEditorState,
-                        updateComposerMode = { messageComposerContext.composerMode = it },
+                        richTextEditorState = richTextEditorState,
-                        textEditorState = textEditorState,
                     )
                 }
                 is MessageComposerEvents.SendUri -> appCoroutineScope.sendAttachment(
@@ -386,7 +379,8 @@ class MessageComposerPresenter @Inject constructor(
                     }
                 }
                 MessageComposerEvents.SaveDraft -> {
-                    appCoroutineScope.saveDraft(textEditorState)
+                    val draft = createDraftFromState(markdownTextEditorState, richTextEditorState)
+                    appCoroutineScope.updateDraft(draft, isVolatile = false)
                 }
             }
         }
@@ -407,42 +401,26 @@ class MessageComposerPresenter @Inject constructor(
     }
 
     private fun CoroutineScope.sendMessage(
-        message: Message,
+        markdownTextEditorState: MarkdownTextEditorState,
-        updateComposerMode: (newComposerMode: MessageComposerMode) -> Unit,
+        richTextEditorState: RichTextEditorState,
-        textEditorState: TextEditorState,
     ) = launch {
+        val message = currentComposerMessage(markdownTextEditorState, richTextEditorState, withMentions = true)
         val capturedMode = messageComposerContext.composerMode
-        val mentions = when (textEditorState) {
-            is TextEditorState.Rich -> {
-                textEditorState.richTextEditorState.mentionsState?.let { state ->
-                    buildList {
-                        if (state.hasAtRoomMention) {
-                            add(Mention.AtRoom)
-                        }
-                        for (userId in state.userIds) {
-                            add(Mention.User(UserId(userId)))
-                        }
-                    }
-                }.orEmpty()
-            }
-            is TextEditorState.Markdown -> textEditorState.state.getMentions()
-        }
         // Reset composer right away
-        textEditorState.reset()
+        resetComposer(markdownTextEditorState, richTextEditorState, fromEdit = capturedMode is MessageComposerMode.Edit)
-        updateComposerMode(MessageComposerMode.Normal)
         when (capturedMode) {
-            is MessageComposerMode.Normal -> room.sendMessage(body = message.markdown, htmlBody = message.html, mentions = mentions)
+            is MessageComposerMode.Normal -> room.sendMessage(body = message.markdown, htmlBody = message.html, mentions = message.mentions)
             is MessageComposerMode.Edit -> {
                 val eventId = capturedMode.eventId
                 val transactionId = capturedMode.transactionId
                 timelineController.invokeOnCurrentTimeline {
-                    editMessage(eventId, transactionId, message.markdown, message.html, mentions)
+                    editMessage(eventId, transactionId, message.markdown, message.html, message.mentions)
                 }
             }
 
             is MessageComposerMode.Reply -> {
                 timelineController.invokeOnCurrentTimeline {
-                    replyMessage(capturedMode.eventId, message.markdown, message.html, mentions)
+                    replyMessage(capturedMode.eventId, message.markdown, message.html, message.mentions)
                 }
             }
         }
@@ -537,21 +515,30 @@ class MessageComposerPresenter @Inject constructor(
             }
         }
 
-    private fun CoroutineScope.loadDraft(
+    private fun CoroutineScope.updateDraft(
+        draft: ComposerDraft?,
+        isVolatile: Boolean,
+    ) = launch {
+        draftService.updateDraft(
+            roomId = room.roomId,
+            draft = draft,
+            isVolatile = isVolatile
+        )
+    }
+
+    private suspend fun applyDraft(
+        draft: ComposerDraft,
         markdownTextEditorState: MarkdownTextEditorState,
         richTextEditorState: RichTextEditorState,
-    ) = launch {
+    ) {
-        val draft = draftService.loadDraft(room.roomId) ?: return@launch
         val htmlText = draft.htmlText
         val markdownText = draft.plainText
         if (htmlText != null) {
             showTextFormatting = true
-            richTextEditorState.setHtml(htmlText)
+            setText(htmlText, markdownTextEditorState, richTextEditorState, requestFocus = true)
-            richTextEditorState.requestFocus()
         } else {
             showTextFormatting = false
-            markdownTextEditorState.text.update(markdownText, true)
+            setText(markdownText, markdownTextEditorState, richTextEditorState, requestFocus = true)
-            markdownTextEditorState.requestFocusAction()
         }
         when (val draftType = draft.draftType) {
             ComposerDraftType.NewMessage -> messageComposerContext.composerMode = MessageComposerMode.Normal
@@ -570,11 +557,11 @@ class MessageComposerPresenter @Inject constructor(
         }
     }
 
-    private fun CoroutineScope.saveDraft(
+    private fun createDraftFromState(
-        textEditorState: TextEditorState,
+        markdownTextEditorState: MarkdownTextEditorState,
-    ) = launch {
+        richTextEditorState: RichTextEditorState,
-        val html = textEditorState.messageHtml()
+    ): ComposerDraft? {
-        val markdown = textEditorState.messageMarkdown(permalinkBuilder)
+        val message = currentComposerMessage(markdownTextEditorState, richTextEditorState, withMentions = false)
         val draftType = when (val mode = messageComposerContext.composerMode) {
             is MessageComposerMode.Normal -> ComposerDraftType.NewMessage
             is MessageComposerMode.Edit -> {
@@ -582,22 +569,54 @@ class MessageComposerPresenter @Inject constructor(
             }
             is MessageComposerMode.Reply -> ComposerDraftType.Reply(mode.eventId)
         }
-        val composerDraft = if (draftType == null || markdown.isBlank()) {
+        return if (draftType == null || message.markdown.isBlank()) {
             null
         } else {
             ComposerDraft(
                 draftType = draftType,
-                htmlText = html,
+                htmlText = message.html,
-                plainText = markdown,
+                plainText = message.markdown,
             )
         }
-        draftService.updateDraft(room.roomId, composerDraft)
+    }
+
+    private fun currentComposerMessage(
+        markdownTextEditorState: MarkdownTextEditorState,
+        richTextEditorState: RichTextEditorState,
+        withMentions: Boolean,
+    ): Message {
+        return if (showTextFormatting) {
+            val html = richTextEditorState.messageHtml
+            val markdown = richTextEditorState.messageMarkdown
+            val mentions = richTextEditorState.mentionsState
+                .takeIf { withMentions }
+                ?.let { state ->
+                    buildList {
+                        if (state.hasAtRoomMention) {
+                            add(Mention.AtRoom)
+                        }
+                        for (userId in state.userIds) {
+                            add(Mention.User(UserId(userId)))
+                        }
+                    }
+                }
+                .orEmpty()
+            Message(html = html, markdown = markdown, mentions = mentions)
+        } else {
+            val markdown = markdownTextEditorState.getMessageMarkdown(permalinkBuilder)
+            val mentions = if (withMentions) {
+                markdownTextEditorState.getMentions()
+            } else {
+                emptyList()
+            }
+            Message(html = null, markdown = markdown, mentions = mentions)
+        }
     }
 
     private fun CoroutineScope.toggleTextFormatting(
         enabled: Boolean,
         markdownTextEditorState: MarkdownTextEditorState,
-        richTextEditorState: RichTextEditorState,
+        richTextEditorState: RichTextEditorState
     ) = launch {
         showTextFormatting = enabled
         if (showTextFormatting) {
@@ -615,24 +634,63 @@ class MessageComposerPresenter @Inject constructor(
     }
 
     private fun CoroutineScope.setMode(
-        composerMode: MessageComposerMode,
+        newComposerMode: MessageComposerMode,
         markdownTextEditorState: MarkdownTextEditorState,
-        richTextEditorState: RichTextEditorState
+        richTextEditorState: RichTextEditorState,
     ) = launch {
-        messageComposerContext.composerMode = composerMode
+        val currentComposerMode = messageComposerContext.composerMode
-        when (composerMode) {
+        when (newComposerMode) {
             is MessageComposerMode.Edit -> {
-                setText(composerMode.content, markdownTextEditorState, richTextEditorState)
+                if (currentComposerMode !is MessageComposerMode.Edit) {
+                    val draft = createDraftFromState(markdownTextEditorState, richTextEditorState)
+                    updateDraft(draft, isVolatile = true).join()
+                }
+                setText(newComposerMode.content, markdownTextEditorState, richTextEditorState)
+            }
+            else -> {
+                // When coming from edit, just clear the composer as it'd be weird to reset a volatile draft in this scenario.
+                if (currentComposerMode is MessageComposerMode.Edit) {
+                    setText("", markdownTextEditorState, richTextEditorState)
+                }
             }
-            else -> Unit
         }
+        messageComposerContext.composerMode = newComposerMode
     }
 
-    private suspend fun setText(content: String, markdownTextEditorState: MarkdownTextEditorState, richTextEditorState: RichTextEditorState) {
+    private suspend fun resetComposer(
+        markdownTextEditorState: MarkdownTextEditorState,
+        richTextEditorState: RichTextEditorState,
+        fromEdit: Boolean,
+    ) {
+        // Use the volatile draft only when coming from edit mode otherwise.
+        val draft = draftService.loadDraft(room.roomId, isVolatile = true).takeIf { fromEdit }
+        if (draft != null) {
+            applyDraft(draft, markdownTextEditorState, richTextEditorState)
+        } else {
+            setText("", markdownTextEditorState, richTextEditorState)
+            messageComposerContext.composerMode = MessageComposerMode.Normal
+        }
+    }
+
+    private suspend fun setText(
+        content: String,
+        markdownTextEditorState: MarkdownTextEditorState,
+        richTextEditorState: RichTextEditorState,
+        requestFocus: Boolean = false,
+    ) {
         if (showTextFormatting) {
             richTextEditorState.setHtml(content)
+            if (requestFocus) {
+                richTextEditorState.requestFocus()
+            }
         } else {
+            if (content.isEmpty()) {
+                markdownTextEditorState.selection = IntRange.EMPTY
+            }
             markdownTextEditorState.text.update(content, true)
+            if (requestFocus) {
+                markdownTextEditorState.requestFocusAction()
+            }
         }
     }
 }

features/messages/impl/src/test/kotlin/io/element/android/features/messages/impl/draft/FakeComposerDraftService.kt

@@ -20,9 +20,9 @@ import io.element.android.libraries.matrix.api.core.RoomId
 import io.element.android.libraries.matrix.api.room.draft.ComposerDraft
 
 class FakeComposerDraftService : ComposerDraftService {
-    var loadDraftLambda: (RoomId) -> ComposerDraft? = { null }
+    var loadDraftLambda: (RoomId, Boolean) -> ComposerDraft? = { _, _ -> null }
-    override suspend fun loadDraft(roomId: RoomId) = loadDraftLambda(roomId)
+    override suspend fun loadDraft(roomId: RoomId, isVolatile: Boolean): ComposerDraft? = loadDraftLambda(roomId, isVolatile)
 
-    var saveDraftLambda: (RoomId, ComposerDraft?) -> Unit = { _, _ -> }
+    var saveDraftLambda: (RoomId, ComposerDraft?, Boolean) -> Unit = { _, _, _ -> }
-    override suspend fun updateDraft(roomId: RoomId, draft: ComposerDraft?) = saveDraftLambda(roomId, draft)
+    override suspend fun updateDraft(roomId: RoomId, draft: ComposerDraft?, isVolatile: Boolean) = saveDraftLambda(roomId, draft, isVolatile)
 }

features/messages/impl/src/test/kotlin/io/element/android/features/messages/impl/draft/VolatileComposerDraftStoreTest.kt

@@ -0,0 +1,56 @@
+/*
+ * Copyright (c) 2024 New Vector Ltd
+ *
+ * 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
+ *
+ *     https://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 io.element.android.features.messages.impl.draft
+
+import com.google.common.truth.Truth.assertThat
+import io.element.android.libraries.matrix.api.room.draft.ComposerDraft
+import io.element.android.libraries.matrix.api.room.draft.ComposerDraftType
+import io.element.android.libraries.matrix.test.A_ROOM_ID
+import kotlinx.coroutines.test.runTest
+import org.junit.Test
+
+class VolatileComposerDraftStoreTest {
+    private val roomId = A_ROOM_ID
+    private val sut = VolatileComposerDraftStore()
+    private val draft = ComposerDraft("plainText", "htmlText", ComposerDraftType.NewMessage)
+
+    @Test
+    fun `when storing a non-null draft and then loading it, it's loaded and removed`() = runTest {
+        val initialDraft = sut.loadDraft(roomId)
+        assertThat(initialDraft).isNull()
+
+        sut.updateDraft(roomId, draft)
+
+        val loadedDraft = sut.loadDraft(roomId)
+        assertThat(loadedDraft).isEqualTo(draft)
+
+        val loadedDraftAfter = sut.loadDraft(roomId)
+        assertThat(loadedDraftAfter).isNull()
+    }
+
+    @Test
+    fun `when storing a null draft and then loading it, it's removing the previous one`() = runTest {
+        val initialDraft = sut.loadDraft(roomId)
+        assertThat(initialDraft).isNull()
+
+        sut.updateDraft(roomId, draft)
+        sut.updateDraft(roomId, null)
+
+        val loadedDraft = sut.loadDraft(roomId)
+        assertThat(loadedDraft).isNull()
+    }
+}

features/messages/impl/src/test/kotlin/io/element/android/features/messages/impl/textcomposer/MessageComposerPresenterTest.kt

@@ -173,21 +173,85 @@ class MessageComposerPresenterTest {
 
     @Test
     fun `present - change mode to edit`() = runTest {
-        val presenter = createPresenter(this)
+        val loadDraftLambda = lambdaRecorder { _: RoomId, _: Boolean ->
+            ComposerDraft(A_MESSAGE, A_MESSAGE, ComposerDraftType.NewMessage)
+        }
+        val updateDraftLambda = lambdaRecorder { _: RoomId, _: ComposerDraft?, _: Boolean -> }
+        val draftService = FakeComposerDraftService().apply {
+            this.loadDraftLambda = loadDraftLambda
+            this.saveDraftLambda = updateDraftLambda
+        }
+        val presenter = createPresenter(
+            coroutineScope = this,
+            draftService = draftService,
+        )
         moleculeFlow(RecompositionMode.Immediate) {
             val state = presenter.present()
             remember(state, state.textEditorState.messageHtml()) { state }
         }.test {
             var state = awaitFirstItem()
-            val mode = anEditMode()
+            val mode = anEditMode(message = ANOTHER_MESSAGE)
             state.eventSink.invoke(MessageComposerEvents.SetMode(mode))
             state = awaitItem()
             assertThat(state.mode).isEqualTo(mode)
+            assertThat(state.textEditorState.messageHtml()).isEqualTo(ANOTHER_MESSAGE)
+            state = backToNormalMode(state)
+            // The message that was being edited is cleared and volatile draft is loaded
             assertThat(state.textEditorState.messageHtml()).isEqualTo(A_MESSAGE)
-            state = backToNormalMode(state, skipCount = 1)
 
-            // The message that was being edited is cleared
+            assert(loadDraftLambda)
-            assertThat(state.textEditorState.messageHtml()).isEqualTo("")
+                .isCalledExactly(2)
+                .withSequence(
+                    // Automatic load of draft
+                    listOf(value(A_ROOM_ID), value(false)),
+                    // Load of volatile draft when closing edit mode
+                    listOf(value(A_ROOM_ID), value(true))
+                )
+
+            assert(updateDraftLambda)
+                .isCalledOnce()
+                .with(value(A_ROOM_ID), any(), value(true))
+        }
+    }
+
+    @Test
+    fun `present - change mode to reply after edit`() = runTest {
+        val loadDraftLambda = lambdaRecorder { _: RoomId, _: Boolean ->
+            ComposerDraft(A_MESSAGE, A_MESSAGE, ComposerDraftType.NewMessage)
+        }
+        val updateDraftLambda = lambdaRecorder { _: RoomId, _: ComposerDraft?, _: Boolean -> }
+        val draftService = FakeComposerDraftService().apply {
+            this.loadDraftLambda = loadDraftLambda
+            this.saveDraftLambda = updateDraftLambda
+        }
+        val presenter = createPresenter(
+            coroutineScope = this,
+            draftService = draftService,
+        )
+        moleculeFlow(RecompositionMode.Immediate) {
+            val state = presenter.present()
+            remember(state, state.textEditorState.messageHtml()) { state }
+        }.test {
+            var state = awaitFirstItem()
+            val editMode = anEditMode(message = ANOTHER_MESSAGE)
+            state.eventSink.invoke(MessageComposerEvents.SetMode(editMode))
+            state = awaitItem()
+            assertThat(state.mode).isEqualTo(editMode)
+            assertThat(state.textEditorState.messageHtml()).isEqualTo(ANOTHER_MESSAGE)
+
+            val replyMode = aReplyMode()
+            state.eventSink.invoke(MessageComposerEvents.SetMode(replyMode))
+            state = awaitItem()
+            assertThat(state.mode).isEqualTo(replyMode)
+            assertThat(state.textEditorState.messageHtml()).isEmpty()
+
+            assert(loadDraftLambda)
+                .isCalledOnce()
+                .with(value(A_ROOM_ID), value(false))
+
+            assert(updateDraftLambda)
+                .isCalledOnce()
+                .with(value(A_ROOM_ID), any(), value(true))
         }
     }
 
@@ -317,7 +381,7 @@ class MessageComposerPresenterTest {
 
             assert(editMessageLambda)
                 .isCalledOnce()
-                .with(any(), any(), value(ANOTHER_MESSAGE), value(ANOTHER_MESSAGE), any())
+                .with(value(AN_EVENT_ID), value(null), value(ANOTHER_MESSAGE), value(ANOTHER_MESSAGE), any())
 
             assertThat(analyticsService.capturedEvents).containsExactly(
                 Composer(
@@ -366,7 +430,7 @@ class MessageComposerPresenterTest {
 
             assert(editMessageLambda)
                 .isCalledOnce()
-                .with(any(), any(), value(ANOTHER_MESSAGE), value(ANOTHER_MESSAGE), any())
+                .with(value(null), value(A_TRANSACTION_ID), value(ANOTHER_MESSAGE), value(ANOTHER_MESSAGE), any())
 
             assertThat(analyticsService.capturedEvents).containsExactly(
                 Composer(
@@ -1013,7 +1077,7 @@ class MessageComposerPresenterTest {
 
     @Test
     fun `present - when there is no draft, nothing is restored`() = runTest {
-        val loadDraftLambda = lambdaRecorder<RoomId, ComposerDraft?> { _ -> null }
+        val loadDraftLambda = lambdaRecorder<RoomId, Boolean, ComposerDraft?> { _, _ -> null }
         val composerDraftService = FakeComposerDraftService().apply {
             this.loadDraftLambda = loadDraftLambda
         }
@@ -1024,7 +1088,7 @@ class MessageComposerPresenterTest {
             awaitFirstItem()
             assert(loadDraftLambda)
                 .isCalledOnce()
-                .with(value(A_ROOM_ID))
+                .with(value(A_ROOM_ID), value(false))
 
             ensureAllEventsConsumed()
         }
@@ -1032,7 +1096,7 @@ class MessageComposerPresenterTest {
 
     @Test
     fun `present - when there is a draft for new message with plain text, it is restored`() = runTest {
-        val loadDraftLambda = lambdaRecorder<RoomId, ComposerDraft?> { _ ->
+        val loadDraftLambda = lambdaRecorder<RoomId, Boolean, ComposerDraft?> { _, _ ->
             ComposerDraft(plainText = A_MESSAGE, htmlText = null, draftType = ComposerDraftType.NewMessage)
         }
         val composerDraftService = FakeComposerDraftService().apply {
@@ -1054,7 +1118,7 @@ class MessageComposerPresenterTest {
 
             assert(loadDraftLambda)
                 .isCalledOnce()
-                .with(value(A_ROOM_ID))
+                .with(value(A_ROOM_ID), value(false))
 
             ensureAllEventsConsumed()
         }
@@ -1062,7 +1126,7 @@ class MessageComposerPresenterTest {
 
     @Test
     fun `present - when there is a draft for new message with rich text, it is restored`() = runTest {
-        val loadDraftLambda = lambdaRecorder<RoomId, ComposerDraft?> { _ ->
+        val loadDraftLambda = lambdaRecorder<RoomId, Boolean, ComposerDraft?> { _, _ ->
             ComposerDraft(
                 plainText = A_MESSAGE,
                 htmlText = A_MESSAGE,
@@ -1088,14 +1152,14 @@ class MessageComposerPresenterTest {
             }
             assert(loadDraftLambda)
                 .isCalledOnce()
-                .with(value(A_ROOM_ID))
+                .with(value(A_ROOM_ID), value(false))
             ensureAllEventsConsumed()
         }
     }
 
     @Test
     fun `present - when there is a draft for edit, it is restored`() = runTest {
-        val loadDraftLambda = lambdaRecorder<RoomId, ComposerDraft?> { _ ->
+        val loadDraftLambda = lambdaRecorder<RoomId, Boolean, ComposerDraft?> { _, _ ->
             ComposerDraft(
                 plainText = A_MESSAGE,
                 htmlText = null,
@@ -1122,7 +1186,7 @@ class MessageComposerPresenterTest {
             }
             assert(loadDraftLambda)
                 .isCalledOnce()
-                .with(value(A_ROOM_ID))
+                .with(value(A_ROOM_ID), value(false))
 
             ensureAllEventsConsumed()
         }
@@ -1130,7 +1194,7 @@ class MessageComposerPresenterTest {
 
     @Test
     fun `present - when there is a draft for reply, it is restored`() = runTest {
-        val loadDraftLambda = lambdaRecorder<RoomId, ComposerDraft?> { _ ->
+        val loadDraftLambda = lambdaRecorder<RoomId, Boolean, ComposerDraft?> { _, _ ->
             ComposerDraft(
                 plainText = A_MESSAGE,
                 htmlText = null,
@@ -1165,7 +1229,7 @@ class MessageComposerPresenterTest {
             }
             assert(loadDraftLambda)
                 .isCalledOnce()
-                .with(value(A_ROOM_ID))
+                .with(value(A_ROOM_ID), value(false))
 
             assert(loadReplyDetailsLambda)
                 .isCalledOnce()
@@ -1177,7 +1241,7 @@ class MessageComposerPresenterTest {
 
     @Test
     fun `present - when save draft event is invoked and composer is empty then service is called with null draft`() = runTest {
-        val saveDraftLambda = lambdaRecorder<RoomId, ComposerDraft?, Unit> { _, _ -> }
+        val saveDraftLambda = lambdaRecorder<RoomId, ComposerDraft?, Boolean, Unit> { _, _, _ -> }
         val composerDraftService = FakeComposerDraftService().apply {
             this.saveDraftLambda = saveDraftLambda
         }
@@ -1190,13 +1254,13 @@ class MessageComposerPresenterTest {
             advanceUntilIdle()
             assert(saveDraftLambda)
                 .isCalledOnce()
-                .with(value(A_ROOM_ID), value(null))
+                .with(value(A_ROOM_ID), value(null), value(false))
         }
     }
 
     @Test
     fun `present - when save draft event is invoked and composer is not empty then service is called`() = runTest {
-        val saveDraftLambda = lambdaRecorder<RoomId, ComposerDraft?, Unit> { _, _ -> }
+        val saveDraftLambda = lambdaRecorder<RoomId, ComposerDraft?, Boolean, Unit> { _, _, _ -> }
         val composerDraftService = FakeComposerDraftService().apply {
             this.saveDraftLambda = saveDraftLambda
         }
@@ -1240,17 +1304,34 @@ class MessageComposerPresenterTest {
             advanceUntilIdle()
 
             assert(saveDraftLambda)
-                .isCalledExactly(4)
+                .isCalledExactly(5)
                 .withSequence(
-                    listOf(value(A_ROOM_ID), value(ComposerDraft(plainText = A_MESSAGE, htmlText = null, draftType = ComposerDraftType.NewMessage))),
-                    listOf(value(A_ROOM_ID), value(ComposerDraft(plainText = A_MESSAGE, htmlText = A_MESSAGE, draftType = ComposerDraftType.NewMessage))),
                     listOf(
                         value(A_ROOM_ID),
-                        value(ComposerDraft(plainText = A_MESSAGE, htmlText = A_MESSAGE, draftType = ComposerDraftType.Edit(AN_EVENT_ID)))
+                        value(ComposerDraft(plainText = A_MESSAGE, htmlText = null, draftType = ComposerDraftType.NewMessage)),
+                        value(false)
+                    ),
+                    listOf(
+                        value(A_ROOM_ID),
+                        value(ComposerDraft(plainText = A_MESSAGE, htmlText = A_MESSAGE, draftType = ComposerDraftType.NewMessage)),
+                        value(false)
+                    ),
+                    listOf(
+                        value(A_ROOM_ID),
+                        value(ComposerDraft(plainText = A_MESSAGE, htmlText = A_MESSAGE, draftType = ComposerDraftType.NewMessage)),
+                        // The volatile draft created when switching to edit mode.
+                        value(true)
+                    ),
+                    listOf(
+                        value(A_ROOM_ID),
+                        value(ComposerDraft(plainText = A_MESSAGE, htmlText = A_MESSAGE, draftType = ComposerDraftType.Edit(AN_EVENT_ID))),
+                        value(false)
                     ),
                     listOf(
                         value(A_ROOM_ID),
-                        value(ComposerDraft(plainText = A_MESSAGE, htmlText = A_MESSAGE, draftType = ComposerDraftType.Reply(AN_EVENT_ID)))
+                        // When moving from edit mode, text composer is cleared, so the draft is null
+                        value(null),
+                        value(false)
                     )
                 )
         }

libraries/textcomposer/impl/src/main/kotlin/io/element/android/libraries/textcomposer/model/Message.kt

@@ -16,7 +16,10 @@
 
 package io.element.android.libraries.textcomposer.model
 
+import io.element.android.libraries.matrix.api.room.Mention
+
 data class Message(
     val html: String?,
     val markdown: String,
+    val mentions: List<Mention>,
 )