From 160b78a7e9895e67ec85d99b46ddea95f6d6d415 Mon Sep 17 00:00:00 2001
From: evykassirer <evy@zulip.com>
Date: Fri, 15 Mar 2024 16:57:19 -0700
Subject: [PATCH] drafts: Don't restore drafts being actively sent or
 scheduled.

---
 web/src/compose.js        | 21 +++++++++++++++++++--
 web/src/drafts.ts         | 30 +++++++++++++++++++++++++++++-
 web/tests/compose.test.js |  7 ++++++-
 web/tests/drafts.test.js  | 16 ++++++++++++++++
 4 files changed, 70 insertions(+), 4 deletions(-)

diff --git a/web/src/compose.js b/web/src/compose.js
index d6784b741a..abc72f585c 100644
--- a/web/src/compose.js
+++ b/web/src/compose.js
@@ -179,7 +179,11 @@ export function send_message(request = create_message_object()) {
 
     // Silently save / update a draft to ensure the message is not lost in case send fails.
     // We delete the draft on successful send.
-    request.draft_id = drafts.update_draft({no_notify: true, update_count: false});
+    request.draft_id = drafts.update_draft({
+        no_notify: true,
+        update_count: false,
+        is_sending_saving: true,
+    });
 
     let local_id;
     let locally_echoed;
@@ -246,6 +250,10 @@ export function send_message(request = create_message_object()) {
         // We might not have updated the draft count because we assumed the
         // message would send. Ensure that the displayed count is correct.
         drafts.sync_count();
+
+        const draft = drafts.draft_model.getDraft(request.draft_id);
+        draft.is_sending_saving = false;
+        drafts.draft_model.editDraft(request.draft_id, draft);
     }
 
     transmit.send_message(request, success, error);
@@ -404,9 +412,15 @@ function schedule_message_to_custom_date() {
         scheduled_delivery_timestamp,
     };
 
+    const draft_id = drafts.update_draft({
+        no_notify: true,
+        update_count: false,
+        is_sending_saving: true,
+    });
+
     const $banner_container = $("#compose_banners");
     const success = function (data) {
-        drafts.draft_model.deleteDraft($("textarea#compose-textarea").data("draft-id"));
+        drafts.draft_model.deleteDraft(draft_id);
         clear_compose_box();
         const new_row_html = render_success_message_scheduled_banner({
             scheduled_message_id: data.scheduled_message_id,
@@ -425,6 +439,9 @@ function schedule_message_to_custom_date() {
             $banner_container,
             $("textarea#compose-textarea"),
         );
+        const draft = drafts.draft_model.getDraft(draft_id);
+        draft.is_sending_saving = false;
+        drafts.draft_model.editDraft(draft_id, draft);
     };
 
     channel.post({
diff --git a/web/src/drafts.ts b/web/src/drafts.ts
index db765a9655..166950c154 100644
--- a/web/src/drafts.ts
+++ b/web/src/drafts.ts
@@ -35,6 +35,7 @@ const draft_schema = z.intersection(
     z.object({
         content: z.string(),
         updatedAt: z.number(),
+        is_sending_saving: z.boolean().default(false),
     }),
     z.discriminatedUnion("type", [
         z.object({
@@ -60,6 +61,7 @@ const possibly_buggy_draft_schema = z.intersection(
     z.object({
         content: z.string(),
         updatedAt: z.number(),
+        is_sending_saving: z.boolean().default(false),
     }),
     z.discriminatedUnion("type", [
         z.object({
@@ -303,6 +305,7 @@ export function snapshot_message(): LocalStorageDraft | undefined {
             type: "private",
             reply_to: recipient,
             private_message_recipient: recipient,
+            is_sending_saving: false,
         };
     }
     assert(message.type === "stream");
@@ -311,6 +314,7 @@ export function snapshot_message(): LocalStorageDraft | undefined {
         type: "stream",
         stream_id: compose_state.stream_id(),
         topic: compose_state.topic(),
+        is_sending_saving: false,
     };
 }
 
@@ -375,12 +379,15 @@ function maybe_notify(no_notify: boolean): void {
 type UpdateDraftOptions = {
     no_notify?: boolean;
     update_count?: boolean;
+    is_sending_saving?: boolean;
 };
 
 export function update_draft(opts: UpdateDraftOptions = {}): string | undefined {
+    const draft_id = $("textarea#compose-textarea").data("draft-id");
+    const old_draft = draft_model.getDraft(draft_id);
+
     const no_notify = opts.no_notify ?? false;
     const draft = snapshot_message();
-    const draft_id = $("textarea#compose-textarea").data("draft-id");
 
     if (draft === undefined) {
         // The user cleared the compose box, which means
@@ -392,6 +399,12 @@ export function update_draft(opts: UpdateDraftOptions = {}): string | undefined
         return undefined;
     }
 
+    if (opts.is_sending_saving !== undefined) {
+        draft.is_sending_saving = opts.is_sending_saving;
+    } else {
+        draft.is_sending_saving = old_draft ? old_draft.is_sending_saving : false;
+    }
+
     if (draft_id !== undefined) {
         // We don't save multiple drafts of the same message;
         // just update the existing draft.
@@ -507,6 +520,7 @@ export function get_last_draft_based_on_compose_state(): LocalStorageDraftWithId
     );
     return drafts_for_compose_state
         .sort((draft_a, draft_b) => draft_a.updatedAt - draft_b.updatedAt)
+        .filter((draft) => !draft.is_sending_saving)
         .pop();
 }
 
@@ -619,6 +633,20 @@ export function format_draft(draft: LocalStorageDraftWithId): FormattedDraft | u
 export function initialize(): void {
     remove_old_drafts();
 
+    // It's possible that drafts will get still have
+    // `is_sending_saving` set to true if the page was
+    // refreshed in the middle of sending a message. We
+    // reset the field on page reload to ensure that drafts
+    // don't get stuck in that state.
+    const current_drafts = draft_model.get();
+    for (const draft_id of Object.keys(current_drafts)) {
+        const draft = current_drafts[draft_id];
+        if (draft.is_sending_saving) {
+            draft.is_sending_saving = false;
+            draft_model.editDraft(draft_id, draft);
+        }
+    }
+
     window.addEventListener("beforeunload", () => {
         update_draft();
     });
diff --git a/web/tests/compose.test.js b/web/tests/compose.test.js
index adba7d41b2..1fed2a7c16 100644
--- a/web/tests/compose.test.js
+++ b/web/tests/compose.test.js
@@ -239,6 +239,12 @@ test_ui("send_message", ({override, override_rewire, mock_template}) => {
         stub_state.get_events_running_called += 1;
     });
 
+    override_rewire(drafts, "update_draft", () => 100);
+    override(drafts.draft_model, "getDraft", (draft_id) => {
+        assert.equal(draft_id, 100);
+        return {};
+    });
+
     // Tests start here.
     (function test_message_send_success_codepath() {
         stub_state = initialize_state_stub_dict();
@@ -251,7 +257,6 @@ test_ui("send_message", ({override, override_rewire, mock_template}) => {
         override(markdown, "render", noop);
         override(markdown, "get_topic_links", noop);
 
-        override_rewire(drafts, "update_draft", () => 100);
         override_rewire(echo, "try_deliver_locally", (message_request) => {
             const local_id_float = 123.04;
             return echo.insert_local_message(message_request, local_id_float, (messages) =>
diff --git a/web/tests/drafts.test.js b/web/tests/drafts.test.js
index d2f5973b67..f5cb068f97 100644
--- a/web/tests/drafts.test.js
+++ b/web/tests/drafts.test.js
@@ -73,6 +73,7 @@ const draft_1 = {
     type: "stream",
     content: "Test stream message",
     updatedAt: mock_current_timestamp,
+    is_sending_saving: false,
 };
 const draft_2 = {
     private_message_recipient: "aaron@zulip.com",
@@ -80,6 +81,7 @@ const draft_2 = {
     type: "private",
     content: "Test direct message",
     updatedAt: mock_current_timestamp,
+    is_sending_saving: false,
 };
 const short_msg = {
     stream_id,
@@ -87,6 +89,7 @@ const short_msg = {
     type: "stream",
     content: "a",
     updatedAt: mock_current_timestamp,
+    is_sending_saving: false,
 };
 
 function test(label, f) {
@@ -243,6 +246,7 @@ test("initialize", ({override_rewire}) => {
         let called = false;
         override_rewire(drafts, "update_draft", () => {
             called = true;
+            return 100;
         });
         f();
         assert.ok(called);
@@ -262,6 +266,7 @@ test("remove_old_drafts", ({override_rewire}) => {
         type: "stream",
         content: "Test stream message",
         updatedAt: Date.now(),
+        is_sending_saving: false,
     };
     const draft_4 = {
         private_message_recipient: "aaron@zulip.com",
@@ -269,6 +274,7 @@ test("remove_old_drafts", ({override_rewire}) => {
         type: "private",
         content: "Test direct message",
         updatedAt: new Date().setDate(-30),
+        is_sending_saving: false,
     };
     const draft_model = drafts.draft_model;
     const ls = localstorage();
@@ -462,6 +468,7 @@ test("format_drafts", ({override, override_rewire, mock_template}) => {
         content: "Test stream message",
         stream_id: 30,
         updatedAt: feb12().getTime(),
+        is_sending_saving: false,
     };
     const draft_2 = {
         private_message_recipient: "aaron@zulip.com",
@@ -469,6 +476,7 @@ test("format_drafts", ({override, override_rewire, mock_template}) => {
         type: "private",
         content: "Test direct message",
         updatedAt: date(-1),
+        is_sending_saving: false,
     };
     const draft_3 = {
         topic: "topic",
@@ -476,6 +484,7 @@ test("format_drafts", ({override, override_rewire, mock_template}) => {
         stream_id: 40,
         content: "Test stream message 2",
         updatedAt: date(-10),
+        is_sending_saving: false,
     };
     const draft_4 = {
         private_message_recipient: "aaron@zulip.com",
@@ -483,6 +492,7 @@ test("format_drafts", ({override, override_rewire, mock_template}) => {
         type: "private",
         content: "Test direct message 2",
         updatedAt: date(-5),
+        is_sending_saving: false,
     };
     const draft_5 = {
         private_message_recipient: "aaron@zulip.com",
@@ -490,6 +500,7 @@ test("format_drafts", ({override, override_rewire, mock_template}) => {
         type: "private",
         content: "Test direct message 3",
         updatedAt: date(-2),
+        is_sending_saving: false,
     };
 
     const expected = [
@@ -619,6 +630,7 @@ test("filter_drafts", ({override, override_rewire, mock_template}) => {
         content: "Test stream message",
         stream_id: 30,
         updatedAt: feb12().getTime(),
+        is_sending_saving: false,
     };
     const pm_draft_1 = {
         private_message_recipient: "aaron@zulip.com",
@@ -626,6 +638,7 @@ test("filter_drafts", ({override, override_rewire, mock_template}) => {
         type: "private",
         content: "Test direct message",
         updatedAt: date(-1),
+        is_sending_saving: false,
     };
     const stream_draft_2 = {
         topic: "topic",
@@ -633,6 +646,7 @@ test("filter_drafts", ({override, override_rewire, mock_template}) => {
         stream_id: 40,
         content: "Test stream message 2",
         updatedAt: date(-10),
+        is_sending_saving: false,
     };
     const pm_draft_2 = {
         private_message_recipient: "aaron@zulip.com",
@@ -640,6 +654,7 @@ test("filter_drafts", ({override, override_rewire, mock_template}) => {
         type: "private",
         content: "Test direct message 2",
         updatedAt: date(-5),
+        is_sending_saving: false,
     };
     const pm_draft_3 = {
         private_message_recipient: "aaron@zulip.com",
@@ -647,6 +662,7 @@ test("filter_drafts", ({override, override_rewire, mock_template}) => {
         type: "private",
         content: "Test direct message 3",
         updatedAt: date(-2),
+        is_sending_saving: false,
     };
 
     const expected_pm_drafts = [