Update owned TON amount when sending stake dice to quick replies.

1 files changed, 23 insertions, 8 deletions

diff --git a/td/telegram/QuickReplyManager.cpp b/td/telegram/QuickReplyManager.cpp
index d842e7a50..b6f110502 100644
--- a/td/telegram/QuickReplyManager.cpp
+++ b/td/telegram/QuickReplyManager.cpp
@@ -34,6 +34,7 @@
 #include "td/telegram/ReplyMarkup.hpp"
 #include "td/telegram/SavedMessagesTopicId.h"
 #include "td/telegram/ServerMessageId.h"
+#include "td/telegram/StarManager.h"
 #include "td/telegram/StoryFullId.h"
 #include "td/telegram/Td.h"
 #include "td/telegram/TdDb.h"
@@ -1862,9 +1863,8 @@ void QuickReplyManager::process_send_quick_reply_updates(QuickReplyShortcutId sh
     if (file_upload_id.is_valid()) {
       send_closure_later(G()->file_manager(), &FileManager::delete_partial_remote_location, file_upload_id);
     }
-    on_failed_send_quick_reply_messages(shortcut_id, std::move(random_ids),
-                                        Status::Error(500, "Receive wrong response"));
-    return;
+    return on_failed_send_quick_reply_messages(shortcut_id, std::move(random_ids),
+                                               Status::Error(500, "Receive wrong response"));
   }
 
   auto updates = telegram_api::move_object_as<telegram_api::updates>(updates_ptr);
@@ -1896,9 +1896,8 @@ void QuickReplyManager::process_send_quick_reply_updates(QuickReplyShortcutId sh
     if (!new_shortcut_id.is_server()) {
       LOG(ERROR) << "Failed to find new shortcut identifier for " << shortcut_id;
       reload_quick_reply_shortcuts();
-      on_failed_send_quick_reply_messages(shortcut_id, std::move(random_ids),
-                                          Status::Error(500, "Receive wrong response"));
-      return;
+      return on_failed_send_quick_reply_messages(shortcut_id, std::move(random_ids),
+                                                 Status::Error(500, "Receive wrong response"));
     }
     auto it = get_shortcut_it(shortcut_id);
     if (it != shortcuts_.shortcuts_.end() && (*it)->shortcut_id_ == shortcut_id) {
@@ -1960,6 +1959,13 @@ void QuickReplyManager::process_send_quick_reply_updates(QuickReplyShortcutId sh
                   std::move(static_cast<telegram_api::updateQuickReplyMessage *>(update.get())->message_),
                   "process_send_quick_reply_updates");
               if (message != nullptr && message->shortcut_id == shortcut_id) {
+                auto stake_ton_count = get_message_content_stake_ton_count(message->content.get());
+                if (stake_ton_count > 0) {
+                  auto prize_ton_count = get_message_content_prize_ton_count(message->content.get());
+                  td_->star_manager_->add_pending_owned_ton_count(prize_ton_count, false);
+                  td_->star_manager_->add_pending_owned_ton_count(stake_ton_count - prize_ton_count, true);
+                }
+
                 update_sent_message_content_from_temporary_message(it->get(), message.get(), false);
                 auto old_message_it = get_message_it(s, message->message_id);
                 if (old_message_it == s->messages_.end()) {
@@ -2067,6 +2073,10 @@ void QuickReplyManager::on_failed_send_quick_reply_messages(QuickReplyShortcutId
         update_failed_to_send_message_content(td_, (*it)->content);
         (*it)->file_upload_id = {};
         (*it)->thumbnail_file_upload_id = {};
+        auto stake_ton_count = get_message_content_stake_ton_count((*it)->content.get());
+        if (stake_ton_count > 0) {
+          td_->star_manager_->add_pending_owned_ton_count(stake_ton_count, false);
+        }
         register_message_content(it->get(), "on_failed_send_quick_reply_messages");
 
         break;
@@ -2606,8 +2616,7 @@ void QuickReplyManager::do_send_message_group(QuickReplyShortcutId shortcut_id,
   }
   pending_message_group_sends_.erase(it);
   if (error.is_error()) {
-    on_failed_send_quick_reply_messages(shortcut_id, std::move(random_ids), std::move(error));
-    return;
+    return on_failed_send_quick_reply_messages(shortcut_id, std::move(random_ids), std::move(error));
   }
 
   LOG(INFO) << "Begin to send media group " << media_album_id << " to " << shortcut_id;
@@ -3809,6 +3818,12 @@ void QuickReplyManager::change_message_files(const QuickReplyMessage *m, const v
 void QuickReplyManager::register_new_message(const QuickReplyMessage *m, const char *source) {
   change_message_files(m, {});
   register_message_content(m, source);
+  if (m->message_id.is_yet_unsent()) {
+    auto stake_ton_count = get_message_content_stake_ton_count(m->content.get());
+    if (stake_ton_count > 0) {
+      td_->star_manager_->add_pending_owned_ton_count(-stake_ton_count, false);
+    }
+  }
 }
 
 void QuickReplyManager::register_message_content(const QuickReplyMessage *m, const char *source) const {