From 0e16c0af021c0f35ae735e33717cbd8c90ec4ee1 Mon Sep 17 00:00:00 2001
From: levlam <levlam@telegram.org>
Date: Wed, 31 Jan 2024 02:58:17 +0300
Subject: [PATCH] Hide video and voice message drafts when the message is sent.

---
 td/telegram/DraftMessage.cpp    | 15 +++++++++++++++
 td/telegram/DraftMessage.h      |  3 +++
 td/telegram/MessagesManager.cpp | 17 ++++++++++++++++-
 3 files changed, 34 insertions(+), 1 deletion(-)

diff --git a/td/telegram/DraftMessage.cpp b/td/telegram/DraftMessage.cpp
index 1139c0970..041b3720f 100644
--- a/td/telegram/DraftMessage.cpp
+++ b/td/telegram/DraftMessage.cpp
@@ -362,6 +362,21 @@ DraftMessage::DraftMessage() = default;
 
 DraftMessage::~DraftMessage() = default;
 
+bool DraftMessage::need_clear_local(MessageContentType content_type) const {
+  if (!is_local()) {
+    return false;
+  }
+  switch (local_content_->get_type()) {
+    case DraftMessageContentType::VideoNote:
+      return content_type == MessageContentType::VideoNote;
+    case DraftMessageContentType::VoiceNote:
+      return content_type == MessageContentType::VoiceNote;
+    default:
+      UNREACHABLE();
+      return false;
+  }
+}
+
 bool DraftMessage::need_update_to(const DraftMessage &other, bool from_update) const {
   if (is_local()) {
     return !from_update || other.is_local();
diff --git a/td/telegram/DraftMessage.h b/td/telegram/DraftMessage.h
index 5bbf01a83..9cca70637 100644
--- a/td/telegram/DraftMessage.h
+++ b/td/telegram/DraftMessage.h
@@ -9,6 +9,7 @@
 #include "td/telegram/DialogId.h"
 #include "td/telegram/InputMessageText.h"
 #include "td/telegram/logevent/LogEvent.h"
+#include "td/telegram/MessageContentType.h"
 #include "td/telegram/MessageId.h"
 #include "td/telegram/MessageInputReplyTo.h"
 #include "td/telegram/td_api.h"
@@ -65,6 +66,8 @@ class DraftMessage {
     return local_content_ != nullptr;
   }
 
+  bool need_clear_local(MessageContentType content_type) const;
+
   bool need_update_to(const DraftMessage &other, bool from_update) const;
 
   void add_dependencies(Dependencies &dependencies) const;
diff --git a/td/telegram/MessagesManager.cpp b/td/telegram/MessagesManager.cpp
index 5fb0c0d97..a77b5d6ae 100644
--- a/td/telegram/MessagesManager.cpp
+++ b/td/telegram/MessagesManager.cpp
@@ -12120,6 +12120,7 @@ void MessagesManager::ttl_read_history(Dialog *d, bool is_outgoing, MessageId fr
 
 void MessagesManager::ttl_read_history_impl(DialogId dialog_id, bool is_outgoing, MessageId from_message_id,
                                             MessageId till_message_id, double view_date) {
+  CHECK(dialog_id.get_type() == DialogType::SecretChat);
   CHECK(!from_message_id.is_scheduled());
   CHECK(!till_message_id.is_scheduled());
 
@@ -30180,9 +30181,23 @@ bool MessagesManager::update_dialog_draft_message(Dialog *d, unique_ptr<DraftMes
 }
 
 void MessagesManager::clear_dialog_draft_by_sent_message(Dialog *d, const Message *m, bool need_update_dialog_pos) {
-  if (!m->clear_draft || td_->auth_manager_->is_bot()) {
+  if (td_->auth_manager_->is_bot()) {
     return;
   }
+  if (!m->clear_draft) {
+    const DraftMessage *draft_message = nullptr;
+    if (m->initial_top_thread_message_id.is_valid()) {
+      auto top_m = get_message_force(d, m->initial_top_thread_message_id, "clear_dialog_draft_by_sent_message");
+      if (top_m != nullptr) {
+        draft_message = top_m->thread_draft_message.get();
+      }
+    } else {
+      draft_message = d->draft_message.get();
+    }
+    if (draft_message == nullptr || !draft_message->need_clear_local(m->content->get_type())) {
+      return;
+    }
+  }
   if (m->initial_top_thread_message_id.is_valid()) {
     set_dialog_draft_message(d->dialog_id, m->initial_top_thread_message_id, nullptr).ignore();
   } else {