From 6c4e29a01716c9b46c53de4e6fda370e5f731716 Mon Sep 17 00:00:00 2001
From: MartinCupela <32706194+MartinCupela@users.noreply.github.com>
Date: Mon, 3 Jun 2024 12:16:46 +0200
Subject: [PATCH] fix: update quoted message references on message.updated and
 message.deleted events (#1310)

---
 src/channel.ts       |  1 +
 src/channel_state.ts | 49 +++++++++++++++++++++++++++++---------------
 test/unit/channel.js | 39 +++++++++++++++++++++++++++++++++++
 3 files changed, 73 insertions(+), 16 deletions(-)

diff --git a/src/channel.ts b/src/channel.ts
index 5ffb0bcb9..75905bf08 100644
--- a/src/channel.ts
+++ b/src/channel.ts
@@ -1352,6 +1352,7 @@ export class Channel<StreamChatGenerics extends ExtendableGenerics = DefaultGene
         if (event.message) {
           this._extendEventWithOwnReactions(event);
           channelState.addMessageSorted(event.message, false, false);
+          channelState._updateQuotedMessageReferences({ message: event.message });
           if (event.message.pinned) {
             channelState.addPinnedMessage(event.message);
           } else {
diff --git a/src/channel_state.ts b/src/channel_state.ts
index 47d6de817..e63fa6172 100644
--- a/src/channel_state.ts
+++ b/src/channel_state.ts
@@ -1,17 +1,17 @@
 import { Channel } from './channel';
 import {
   ChannelMemberResponse,
-  FormatMessageResponse,
+  DefaultGenerics,
   Event,
   ExtendableGenerics,
-  DefaultGenerics,
-  MessageSetType,
+  FormatMessageResponse,
   MessageResponse,
-  ReactionResponse,
-  UserResponse,
+  MessageSetType,
   PendingMessageResponse,
-  PollVote,
   PollResponse,
+  PollVote,
+  ReactionResponse,
+  UserResponse,
 } from './types';
 import { addToMessageList } from './utils';
 
@@ -242,14 +242,13 @@ export class ChannelState<StreamChatGenerics extends ExtendableGenerics = Defaul
        */
       if (parentID && !initializing) {
         const thread = this.threads[parentID] || [];
-        const threadMessages = this._addToMessageList(
+        this.threads[parentID] = this._addToMessageList(
           thread,
           message,
           timestampChanged,
           'created_at',
           addIfDoesNotExist,
         );
-        this.threads[parentID] = threadMessages;
       }
     }
 
@@ -349,7 +348,13 @@ export class ChannelState<StreamChatGenerics extends ExtendableGenerics = Defaul
     return messageWithReaction;
   }
 
-  removeQuotedMessageReferences(message: MessageResponse<StreamChatGenerics>) {
+  _updateQuotedMessageReferences({
+    message,
+    remove,
+  }: {
+    message: MessageResponse<StreamChatGenerics>;
+    remove?: boolean;
+  }) {
     const parseMessage = (m: ReturnType<ChannelState<StreamChatGenerics>['formatMessage']>) =>
       (({
         ...m,
@@ -358,14 +363,26 @@ export class ChannelState<StreamChatGenerics extends ExtendableGenerics = Defaul
         updated_at: m.updated_at?.toISOString(),
       } as unknown) as MessageResponse<StreamChatGenerics>);
 
-    this.messageSets.forEach((set) => {
-      const updatedMessages = set.messages
-        .filter((msg) => msg.quoted_message_id === message.id)
-        .map(parseMessage)
-        .map((msg) => ({ ...msg, quoted_message: { ...message, attachments: [] } }));
-
+    const update = (messages: FormatMessageResponse<StreamChatGenerics>[]) => {
+      const updatedMessages = messages.reduce<MessageResponse<StreamChatGenerics>[]>((acc, msg) => {
+        if (msg.quoted_message_id === message.id) {
+          acc.push({ ...parseMessage(msg), quoted_message: remove ? { ...message, attachments: [] } : message });
+        }
+        return acc;
+      }, []);
       this.addMessagesSorted(updatedMessages, true);
-    });
+    };
+
+    if (!message.parent_id) {
+      this.messageSets.forEach((set) => update(set.messages));
+    } else if (message.parent_id && this.threads[message.parent_id]) {
+      // prevent going through all the threads even though it is possible to quote a message from another thread
+      update(this.threads[message.parent_id]);
+    }
+  }
+
+  removeQuotedMessageReferences(message: MessageResponse<StreamChatGenerics>) {
+    this._updateQuotedMessageReferences({ message, remove: true });
   }
 
   /**
diff --git a/test/unit/channel.js b/test/unit/channel.js
index 11609a758..ab724086d 100644
--- a/test/unit/channel.js
+++ b/test/unit/channel.js
@@ -542,6 +542,45 @@ describe('Channel _handleChannelEvent', function () {
 		});
 	});
 
+	it('should update quoted_message references on "message.updated" and "message.deleted" event', () => {
+		const originalText = 'XX';
+		const updatedText = 'YY';
+		const parent_id = '0';
+		const parentMesssage = generateMsg({ date: new Date(0).toISOString(), id: parent_id });
+		const quoted_message = generateMsg({
+			date: new Date(2).toISOString(),
+			id: 'quoted-message',
+			text: originalText,
+		});
+		const quotingMessage = generateMsg({
+			date: new Date(3).toISOString(),
+			id: 'quoting-message',
+			quoted_message,
+			quoted_message_id: quoted_message.id,
+		});
+		const updatedQuotedMessage = { ...quoted_message, text: updatedText };
+		const updatedQuotedThreadReply = { ...quoted_message, parent_id, text: updatedText };
+		[
+			[quoted_message, quotingMessage], // channel message
+			[parentMesssage, { ...quoted_message, parent_id }, { ...quotingMessage, parent_id }], // thread message
+		].forEach((messages) => {
+			['message.updated', 'message.deleted'].forEach((eventType) => {
+				channel.state.addMessagesSorted(messages);
+				const isThread = messages.length === 3;
+				const quotingMessage = messages[messages.length - 1];
+				const event = {
+					type: eventType,
+					message: isThread ? updatedQuotedThreadReply : updatedQuotedMessage,
+				};
+				channel._handleChannelEvent(event);
+				expect(
+					channel.state.findMessage(quotingMessage.id, quotingMessage.parent_id).quoted_message.text,
+				).to.equal(updatedQuotedMessage.text);
+				channel.state.clearMessages();
+			});
+		});
+	});
+
 	it('should mark channel visible on channel.visible event', () => {
 		const channelVisibleEvent = {
 			type: 'channel.visible',