From 981e7a4270dec995ff26d74a48b68916df42d11b Mon Sep 17 00:00:00 2001
From: Greyson Parrelli <greyson@signal.org>
Date: Mon, 5 Aug 2024 10:54:53 -0400
Subject: [PATCH] Tweak when decryption drained state is reset.

---
 .../securesms/messages/IncomingMessageObserver.kt           | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/app/src/main/java/org/thoughtcrime/securesms/messages/IncomingMessageObserver.kt b/app/src/main/java/org/thoughtcrime/securesms/messages/IncomingMessageObserver.kt
index 9581411d92..dbef01a279 100644
--- a/app/src/main/java/org/thoughtcrime/securesms/messages/IncomingMessageObserver.kt
+++ b/app/src/main/java/org/thoughtcrime/securesms/messages/IncomingMessageObserver.kt
@@ -376,8 +376,10 @@ class IncomingMessageObserver(private val context: Application) {
         val webSocketDisposable = AppDependencies.webSocketObserver.subscribe { state: WebSocketConnectionState ->
           Log.d(TAG, "WebSocket State: $state")
 
-          // Any state change at all means that we are not drained
-          decryptionDrained = false
+          // Any change to a non-connected state means that we are not drained
+          if (state != WebSocketConnectionState.CONNECTED) {
+            decryptionDrained = false
+          }
 
           if (state == WebSocketConnectionState.CONNECTED) {
             SignalStore.misc.lastWebSocketConnectTime = System.currentTimeMillis()