From 1b9e4586406e867d6f630b59c7a1d15d1ece7538 Mon Sep 17 00:00:00 2001
From: Infi <wingit@geist.ga>
Date: Sat, 3 Jun 2023 18:01:20 +0200
Subject: [PATCH] fix: new messages not coming in when an edit comes in

Signed-off-by: Infi <wingit@geist.ga>
---
 .../chat/views/channel/ChannelScreenViewModel.kt | 16 +++++++++++++---
 1 file changed, 13 insertions(+), 3 deletions(-)

diff --git a/app/src/main/java/chat/revolt/screens/chat/views/channel/ChannelScreenViewModel.kt b/app/src/main/java/chat/revolt/screens/chat/views/channel/ChannelScreenViewModel.kt
index e58778ae..042efa12 100644
--- a/app/src/main/java/chat/revolt/screens/chat/views/channel/ChannelScreenViewModel.kt
+++ b/app/src/main/java/chat/revolt/screens/chat/views/channel/ChannelScreenViewModel.kt
@@ -10,6 +10,7 @@ import androidx.lifecycle.viewModelScope
 import chat.revolt.api.RevoltAPI
 import chat.revolt.api.RevoltJson
 import chat.revolt.api.internals.ULID
+import chat.revolt.api.realtime.RealtimeSocketFrames
 import chat.revolt.api.realtime.frames.receivable.ChannelStartTypingFrame
 import chat.revolt.api.realtime.frames.receivable.ChannelStopTypingFrame
 import chat.revolt.api.realtime.frames.receivable.MessageDeleteFrame
@@ -283,9 +284,12 @@ class ChannelScreenViewModel : ViewModel() {
                             currentMsg.id == it.id
                         } ?: return@onEach // Message not found, ignore.
 
+                        if (messageFrame.author != null)
+                            addUserIfUnknown(messageFrame.author)
+
                         regroupMessages(renderableMessages.map { currentMsg ->
                             if (currentMsg.id == it.id) {
-                                messageFrame
+                                currentMsg.mergeWithPartial(messageFrame)
                             } else {
                                 currentMsg
                             }
@@ -295,9 +299,10 @@ class ChannelScreenViewModel : ViewModel() {
                     is MessageDeleteFrame -> {
                         if (it.channel != channel?.id) return@onEach
 
-                        regroupMessages(renderableMessages.filter { currentMsg ->
+                        val newRenderableMessages = renderableMessages.filter { currentMsg ->
                             currentMsg.id != it.id
-                        })
+                        }
+                        regroupMessages(newRenderableMessages)
                     }
 
                     is ChannelStartTypingFrame -> {
@@ -314,6 +319,11 @@ class ChannelScreenViewModel : ViewModel() {
 
                         _typingUsers.remove(it.user)
                     }
+
+                    is RealtimeSocketFrames.Reconnected -> {
+                        Log.d("ChannelScreen", "Reconnected to WS.")
+                        listenForWsFrame()
+                    }
                 }
             }.catch {
                 Log.e("ChannelScreen", "Failed to receive WS frame", it)