Fix photo suggestion stops working after the first try

submodules/TelegramUI/Components/MediaEditorScreen/Sources/MediaEditorScreen.swift

@@ -6764,7 +6764,7 @@ public final class MediaEditorScreenImpl: ViewController, MediaEditorScreen, UID
     var didComplete = false
 
     public var cancelled: (Bool) -> Void = { _ in }
-    public var willComplete: (UIImage?, Bool, @escaping () -> Void) -> Void
+    public var willComplete: (UIImage?, Bool, @escaping () -> Void, @escaping () -> Void) -> Void
     public var completion: ([MediaEditorScreenImpl.Result], @escaping (@escaping () -> Void) -> Void) -> Void
     public var dismissed: () -> Void = { }
     public var willDismiss: () -> Void = { }
@@ -6798,7 +6798,7 @@ public final class MediaEditorScreenImpl: ViewController, MediaEditorScreen, UID
         initialLink: (url: String, name: String?)? = nil,
         transitionIn: TransitionIn?,
         transitionOut: @escaping (Bool, Bool?) -> TransitionOut?,
-        willComplete: @escaping (UIImage?, Bool, @escaping () -> Void) -> Void = { _, _, commit in commit() },
+        willComplete: @escaping (UIImage?, Bool, @escaping () -> Void, @escaping () -> Void) -> Void = { _, _, commit, _ in commit() },
         completion: @escaping ([MediaEditorScreenImpl.Result], @escaping (@escaping () -> Void) -> Void) -> Void
     ) {
         self.context = context

submodules/TelegramUI/Components/MediaEditorScreen/Sources/MediaEditorStoryCompletion.swift

@@ -16,13 +16,9 @@ extension MediaEditorScreenImpl {
         }
 
         self.didComplete = true
-        
+
         self.updateMediaEditorEntities()
 
-        mediaEditor.stop()
-        mediaEditor.invalidate()
-        self.node.entitiesView.invalidate()
-
         if let navigationController = self.navigationController as? NavigationController {
             navigationController.updateRootContainerTransitionOffset(0.0, transition: .immediate)
         }
@@ -463,6 +459,9 @@ extension MediaEditorScreenImpl {
                                 guard let self else {
                                     return
                                 }
+                                self.node.mediaEditor?.stop()
+                                self.node.mediaEditor?.invalidate()
+                                self.node.entitiesView.invalidate()
                                 Logger.shared.log("MediaEditor", "Completed with video \(videoResult)")
                                 self.completion([MediaEditorScreenImpl.Result(media: .video(video: videoResult, coverImage: coverImage, values: values, duration: duration, dimensions: values.resultDimensions), mediaAreas: mediaAreas, caption: caption, coverTimestamp: values.coverImageTimestamp, options: self.state.privacy, stickers: stickers, randomId: randomId)], { [weak self] finished in
                                     self?.node.animateOut(finished: true, saveDraft: false, completion: { [weak self] in
@@ -472,6 +471,8 @@ extension MediaEditorScreenImpl {
                                         }
                                     })
                                 })
+                            }, { [weak self] in
+                                self?.didComplete = false
                             })
                         }
                     })
@@ -505,6 +506,9 @@ extension MediaEditorScreenImpl {
                         guard let self else {
                             return
                         }
+                        self.node.mediaEditor?.stop()
+                        self.node.mediaEditor?.invalidate()
+                        self.node.entitiesView.invalidate()
                         Logger.shared.log("MediaEditor", "Completed with image \(resultImage)")
                         self.completion([MediaEditorScreenImpl.Result(media: .image(image: resultImage, dimensions: PixelDimensions(resultImage.size)), mediaAreas: mediaAreas, caption: caption, coverTimestamp: nil, options: self.state.privacy, stickers: stickers, randomId: randomId)], { [weak self] finished in
                             self?.node.animateOut(finished: true, saveDraft: false, completion: { [weak self] in
@@ -517,6 +521,8 @@ extension MediaEditorScreenImpl {
                         if case let .draft(draft, id) = actualSubject, id == nil {
                             removeStoryDraft(engine: self.context.engine, path: draft.path, delete: true)
                         }
+                    }, { [weak self] in
+                        self?.didComplete = false
                     })
                 }
             })

submodules/TelegramUI/Components/PeerInfo/PeerInfoScreen/Sources/PeerInfoScreenAvatarSetup.swift

@@ -196,10 +196,12 @@ extension PeerInfoScreenImpl {
                         }
                         return nil
                     },
-                    willComplete: { [weak self, weak parentController] image, isVideo, commit in
+                    willComplete: { [weak self, weak parentController] image, isVideo, commit, cancel in
                         if let self, let confirmationAlert, let image {
                             let controller = photoUpdateConfirmationController(context: self.context, peer: peer, image: image, text: isVideo ? confirmationAlert.videoText : confirmationAlert.photoText, doneTitle: confirmationAlert.action, commit: {
                                 commit()
+                            }, onCancel: {
+                                cancel()
                             })
                             parentController?.presentInGlobalOverlay(controller)
                         } else {

submodules/TelegramUI/Components/PeerInfo/PeerInfoScreen/Sources/PhotoUpdateConfirmationController.swift

@@ -21,7 +21,8 @@ func photoUpdateConfirmationController(
     text: String,
     doneTitle: String,
     isDark: Bool = true,
-    commit: @escaping () -> Void
+    commit: @escaping () -> Void,
+    onCancel: (() -> Void)? = nil
 ) -> ViewController {
     let presentationData = context.sharedContext.currentPresentationData.with { $0 }
     let strings = presentationData.strings
@@ -63,16 +64,23 @@ func photoUpdateConfirmationController(
         updatedPresentationData = (presentationData, context.sharedContext.presentationData)
     }
 
+    var didCommit = false
     let alertController = AlertScreen(
         content: content,
         actions: [
             .init(title: strings.Common_Cancel),
             .init(title: doneTitle, type: .default, action: {
+                didCommit = true
                 commit()
             })
         ],
         updatedPresentationData: updatedPresentationData
     )
+    alertController.dismissed = { _ in
+        if !didCommit {
+            onCancel?()
+        }
+    }
     return alertController
 }