Volatile draft to restore the composer after an edit. (#2996)

ElementX/Sources/Mocks/Generated/GeneratedMocks.swift

@@ -4480,6 +4480,47 @@ class ComposerDraftServiceMock: ComposerDraftServiceProtocol {
             return saveDraftReturnValue
         }
     }
+    //MARK: - saveVolatileDraft
+
+    var saveVolatileDraftUnderlyingCallsCount = 0
+    var saveVolatileDraftCallsCount: Int {
+        get {
+            if Thread.isMainThread {
+                return saveVolatileDraftUnderlyingCallsCount
+            } else {
+                var returnValue: Int? = nil
+                DispatchQueue.main.sync {
+                    returnValue = saveVolatileDraftUnderlyingCallsCount
+                }
+
+                return returnValue!
+            }
+        }
+        set {
+            if Thread.isMainThread {
+                saveVolatileDraftUnderlyingCallsCount = newValue
+            } else {
+                DispatchQueue.main.sync {
+                    saveVolatileDraftUnderlyingCallsCount = newValue
+                }
+            }
+        }
+    }
+    var saveVolatileDraftCalled: Bool {
+        return saveVolatileDraftCallsCount > 0
+    }
+    var saveVolatileDraftReceivedDraft: ComposerDraftProxy?
+    var saveVolatileDraftReceivedInvocations: [ComposerDraftProxy] = []
+    var saveVolatileDraftClosure: ((ComposerDraftProxy) -> Void)?
+
+    func saveVolatileDraft(_ draft: ComposerDraftProxy) {
+        saveVolatileDraftCallsCount += 1
+        saveVolatileDraftReceivedDraft = draft
+        DispatchQueue.main.async {
+            self.saveVolatileDraftReceivedInvocations.append(draft)
+        }
+        saveVolatileDraftClosure?(draft)
+    }
     //MARK: - loadDraft
 
     var loadDraftUnderlyingCallsCount = 0
@@ -4544,6 +4585,70 @@ class ComposerDraftServiceMock: ComposerDraftServiceProtocol {
             return loadDraftReturnValue
         }
     }
+    //MARK: - loadVolatileDraft
+
+    var loadVolatileDraftUnderlyingCallsCount = 0
+    var loadVolatileDraftCallsCount: Int {
+        get {
+            if Thread.isMainThread {
+                return loadVolatileDraftUnderlyingCallsCount
+            } else {
+                var returnValue: Int? = nil
+                DispatchQueue.main.sync {
+                    returnValue = loadVolatileDraftUnderlyingCallsCount
+                }
+
+                return returnValue!
+            }
+        }
+        set {
+            if Thread.isMainThread {
+                loadVolatileDraftUnderlyingCallsCount = newValue
+            } else {
+                DispatchQueue.main.sync {
+                    loadVolatileDraftUnderlyingCallsCount = newValue
+                }
+            }
+        }
+    }
+    var loadVolatileDraftCalled: Bool {
+        return loadVolatileDraftCallsCount > 0
+    }
+
+    var loadVolatileDraftUnderlyingReturnValue: ComposerDraftProxy?
+    var loadVolatileDraftReturnValue: ComposerDraftProxy? {
+        get {
+            if Thread.isMainThread {
+                return loadVolatileDraftUnderlyingReturnValue
+            } else {
+                var returnValue: ComposerDraftProxy?? = nil
+                DispatchQueue.main.sync {
+                    returnValue = loadVolatileDraftUnderlyingReturnValue
+                }
+
+                return returnValue!
+            }
+        }
+        set {
+            if Thread.isMainThread {
+                loadVolatileDraftUnderlyingReturnValue = newValue
+            } else {
+                DispatchQueue.main.sync {
+                    loadVolatileDraftUnderlyingReturnValue = newValue
+                }
+            }
+        }
+    }
+    var loadVolatileDraftClosure: (() -> ComposerDraftProxy?)?
+
+    func loadVolatileDraft() -> ComposerDraftProxy? {
+        loadVolatileDraftCallsCount += 1
+        if let loadVolatileDraftClosure = loadVolatileDraftClosure {
+            return loadVolatileDraftClosure()
+        } else {
+            return loadVolatileDraftReturnValue
+        }
+    }
     //MARK: - clearDraft
 
     var clearDraftUnderlyingCallsCount = 0
@@ -4608,6 +4713,41 @@ class ComposerDraftServiceMock: ComposerDraftServiceProtocol {
             return clearDraftReturnValue
         }
     }
+    //MARK: - clearVolatileDraft
+
+    var clearVolatileDraftUnderlyingCallsCount = 0
+    var clearVolatileDraftCallsCount: Int {
+        get {
+            if Thread.isMainThread {
+                return clearVolatileDraftUnderlyingCallsCount
+            } else {
+                var returnValue: Int? = nil
+                DispatchQueue.main.sync {
+                    returnValue = clearVolatileDraftUnderlyingCallsCount
+                }
+
+                return returnValue!
+            }
+        }
+        set {
+            if Thread.isMainThread {
+                clearVolatileDraftUnderlyingCallsCount = newValue
+            } else {
+                DispatchQueue.main.sync {
+                    clearVolatileDraftUnderlyingCallsCount = newValue
+                }
+            }
+        }
+    }
+    var clearVolatileDraftCalled: Bool {
+        return clearVolatileDraftCallsCount > 0
+    }
+    var clearVolatileDraftClosure: (() -> Void)?
+
+    func clearVolatileDraft() {
+        clearVolatileDraftCallsCount += 1
+        clearVolatileDraftClosure?()
+    }
     //MARK: - getReply
 
     var getReplyEventIDUnderlyingCallsCount = 0

ElementX/Sources/Screens/RoomScreen/ComposerToolbar/ComposerToolbarViewModel.swift

@@ -156,8 +156,13 @@ final class ComposerToolbarViewModel: ComposerToolbarViewModelType, ComposerTool
         case .cancelReply:
             set(mode: .default)
         case .cancelEdit:
-            set(mode: .default)
+            if let draft = draftService.loadVolatileDraft() {
-            set(text: "")
+                handleLoadDraft(draft)
+                draftService.clearVolatileDraft()
+            } else {
+                set(text: "")
+                set(mode: .default)
+            }
         case .attach(let attachment):
             state.bindings.composerFocused = false
             actionsSubject.send(.attach(attachment))
@@ -198,6 +203,9 @@ final class ComposerToolbarViewModel: ComposerToolbarViewModelType, ComposerTool
     func process(roomAction: RoomScreenComposerAction) {
         switch roomAction {
         case .setMode(mode: let mode):
+            if state.composerMode.isComposingNewMessage, mode.isEdit {
+                handleSaveDraft(isVolatile: true)
+            }
             set(mode: mode)
         case .setText(let plainText, let htmlText):
             if let htmlText, context.composerFormattingEnabled {
@@ -208,13 +216,22 @@ final class ComposerToolbarViewModel: ComposerToolbarViewModelType, ComposerTool
         case .removeFocus:
             state.bindings.composerFocused = false
         case .clear:
-            set(mode: .default)
+            if let draft = draftService.loadVolatileDraft() {
-            set(text: "")
+                handleLoadDraft(draft)
+                draftService.clearVolatileDraft()
+            } else {
+                set(mode: .default)
+                set(text: "")
+            }
         case .saveDraft:
-            handleSaveDraft()
+            handleSaveDraft(isVolatile: false)
         case .loadDraft:
             Task {
-                await handleLoadDraft()
+                guard case let .success(draft) = await draftService.loadDraft(),
+                      let draft else {
+                    return
+                }
+                handleLoadDraft(draft)
             }
         }
     }
@@ -229,12 +246,7 @@ final class ComposerToolbarViewModel: ComposerToolbarViewModelType, ComposerTool
 
     // MARK: - Private
     
-    private func handleLoadDraft() async {
+    private func handleLoadDraft(_ draft: ComposerDraftProxy) {
-        guard case let .success(draft) = await draftService.loadDraft(),
-              let draft else {
-            return
-        }
-        
         if let html = draft.htmlText {
             context.composerFormattingEnabled = true
             DispatchQueue.main.async {
@@ -269,15 +281,19 @@ final class ComposerToolbarViewModel: ComposerToolbarViewModelType, ComposerTool
         }
     }
     
-    private func handleSaveDraft() {
+    private func handleSaveDraft(isVolatile: Bool) {
         let plainText: String
         let htmlText: String?
         let type: ComposerDraftProxy.ComposerDraftType
         
         if context.composerFormattingEnabled {
             if wysiwygViewModel.isContentEmpty, state.composerMode == .default {
-                Task {
+                if isVolatile {
-                    await draftService.clearDraft()
+                    draftService.clearVolatileDraft()
+                } else {
+                    Task {
+                        await draftService.clearDraft()
+                    }
                 }
                 return
             }
@@ -285,8 +301,12 @@ final class ComposerToolbarViewModel: ComposerToolbarViewModelType, ComposerTool
             htmlText = wysiwygViewModel.content.html
         } else {
             if context.plainComposerText.string.isEmpty, state.composerMode == .default {
-                Task {
+                if isVolatile {
-                    await draftService.clearDraft()
+                    draftService.clearVolatileDraft()
+                } else {
+                    Task {
+                        await draftService.clearDraft()
+                    }
                 }
                 return
             }
@@ -310,15 +330,22 @@ final class ComposerToolbarViewModel: ComposerToolbarViewModelType, ComposerTool
             }
             type = .reply(eventID: eventID)
         default:
-            // Do not save a draft for the other cases
+            if isVolatile {
-            Task {
+                draftService.clearVolatileDraft()
-                await draftService.clearDraft()
+            } else {
+                Task {
+                    await draftService.clearDraft()
+                }
             }
             return
         }
         
-        Task {
+        if isVolatile {
-            await draftService.saveDraft(.init(plainText: plainText, htmlText: htmlText, draftType: type))
+            draftService.saveVolatileDraft(.init(plainText: plainText, htmlText: htmlText, draftType: type))
+        } else {
+            Task {
+                await draftService.saveDraft(.init(plainText: plainText, htmlText: htmlText, draftType: type))
+            }
         }
     }
     

ElementX/Sources/Screens/RoomScreen/RoomScreenInteractionHandler.swift

@@ -276,8 +276,9 @@ class RoomScreenInteractionHandler {
             text = messageTimelineItem.body
         }
         
-        actionsSubject.send(.composer(action: .setText(plainText: text, htmlText: htmlText)))
+        // Always update the mode first and then the text so that the composer has time to save the text draft
         actionsSubject.send(.composer(action: .setMode(mode: .edit(originalItemId: messageTimelineItem.id))))
+        actionsSubject.send(.composer(action: .setText(plainText: text, htmlText: htmlText)))
     }
     
     // MARK: Polls

ElementX/Sources/Screens/RoomScreen/RoomScreenModels.swift

@@ -83,6 +83,15 @@ enum RoomScreenComposerMode: Equatable {
             return nil
         }
     }
+    
+    var isComposingNewMessage: Bool {
+        switch self {
+        case .default, .reply:
+            return true
+        default:
+            return false
+        }
+    }
 }
 
 enum RoomScreenViewPollAction {

ElementX/Sources/Services/ComposerDraft/ComposerDraftService.swift

@@ -21,6 +21,7 @@ import MatrixRustSDK
 final class ComposerDraftService: ComposerDraftServiceProtocol {
     private let roomProxy: RoomProxyProtocol
     private let timelineItemfactory: RoomTimelineItemFactoryProtocol
+    private var volatileDraft: ComposerDraftProxy?
     
     init(roomProxy: RoomProxyProtocol, timelineItemfactory: RoomTimelineItemFactoryProtocol) {
         self.roomProxy = roomProxy
@@ -71,4 +72,16 @@ final class ComposerDraftService: ComposerDraftServiceProtocol {
             return .failure(.failedToClearDraft)
         }
     }
+    
+    func saveVolatileDraft(_ draft: ComposerDraftProxy) {
+        volatileDraft = draft
+    }
+    
+    func loadVolatileDraft() -> ComposerDraftProxy? {
+        volatileDraft
+    }
+    
+    func clearVolatileDraft() {
+        volatileDraft = nil
+    }
 }

ElementX/Sources/Services/ComposerDraft/ComposerDraftServiceProtocol.swift

@@ -74,7 +74,10 @@ enum ComposerDraftServiceError: Error {
 // sourcery: AutoMockable
 protocol ComposerDraftServiceProtocol {
     func saveDraft(_ draft: ComposerDraftProxy) async -> Result<Void, ComposerDraftServiceError>
+    func saveVolatileDraft(_ draft: ComposerDraftProxy)
     func loadDraft() async -> Result<ComposerDraftProxy?, ComposerDraftServiceError>
+    func loadVolatileDraft() -> ComposerDraftProxy?
     func clearDraft() async -> Result<Void, ComposerDraftServiceError>
+    func clearVolatileDraft()
     func getReply(eventID: String) async -> Result<TimelineItemReply, ComposerDraftServiceError>
 }

UnitTests/Sources/ComposerToolbarViewModelTests.swift

@@ -489,6 +489,70 @@ class ComposerToolbarViewModelTests: XCTestCase {
         await fulfillment(of: [loadReplyExpectation], timeout: 10)
         XCTAssertEqual(viewModel.state.composerMode, .default)
     }
+    
+    func testSaveVolatileDraftWhenEditing() {
+        viewModel.context.composerFormattingEnabled = false
+        viewModel.context.plainComposerText = .init(string: "Hello world!")
+        viewModel.process(roomAction: .setMode(mode: .edit(originalItemId: .random)))
+        
+        let draft = draftServiceMock.saveVolatileDraftReceivedDraft
+        XCTAssertNotNil(draft)
+        XCTAssertEqual(draft?.plainText, "Hello world!")
+        XCTAssertNil(draft?.htmlText)
+        XCTAssertEqual(draft?.draftType, .newMessage)
+    }
+    
+    func testRestoreVolatileDraftWhenCancellingEdit() async {
+        let expectation = expectation(description: "Wait for draft to be restored")
+        draftServiceMock.loadVolatileDraftClosure = {
+            defer { expectation.fulfill() }
+            return .init(plainText: "Hello world",
+                         htmlText: nil,
+                         draftType: .newMessage)
+        }
+        
+        viewModel.process(viewAction: .cancelEdit)
+        await fulfillment(of: [expectation])
+        XCTAssertEqual(viewModel.context.plainComposerText, NSAttributedString(string: "Hello world"))
+    }
+    
+    func testRestoreVolatileDraftWhenClearing() async {
+        let expectation1 = expectation(description: "Wait for draft to be restored")
+        draftServiceMock.loadVolatileDraftClosure = {
+            defer { expectation1.fulfill() }
+            return .init(plainText: "Hello world",
+                         htmlText: nil,
+                         draftType: .newMessage)
+        }
+        
+        let expectation2 = expectation(description: "The draft should also be cleared after being loaded")
+        draftServiceMock.clearVolatileDraftClosure = {
+            expectation2.fulfill()
+        }
+        
+        viewModel.process(roomAction: .clear)
+        await fulfillment(of: [expectation1, expectation2])
+        XCTAssertEqual(viewModel.context.plainComposerText, NSAttributedString(string: "Hello world"))
+    }
+    
+    func testRestoreVolatileDraftDoubleClear() async {
+        let expectation1 = expectation(description: "Wait for draft to be restored")
+        draftServiceMock.loadVolatileDraftClosure = {
+            defer { expectation1.fulfill() }
+            return .init(plainText: "Hello world",
+                         htmlText: nil,
+                         draftType: .newMessage)
+        }
+        
+        let expectation2 = expectation(description: "The draft should also be cleared after being loaded")
+        draftServiceMock.clearVolatileDraftClosure = {
+            expectation2.fulfill()
+        }
+        
+        viewModel.process(roomAction: .clear)
+        await fulfillment(of: [expectation1, expectation2])
+        XCTAssertEqual(viewModel.context.plainComposerText, NSAttributedString(string: "Hello world"))
+    }
 }
 
 private extension MentionSuggestionItem {