From f9671cb54bdaa40c4844208cae753c415029a84c Mon Sep 17 00:00:00 2001
From: 0xb4lamx <medghaith.sellami@gmail.com>
Date: Tue, 11 Feb 2025 12:00:29 +0100
Subject: [PATCH] fix(server): fast ws disconnect handling during conn setup

---
 packages/server/src/ClientConnection.ts | 21 +++++++++++----------
 packages/server/src/Connection.ts       |  1 +
 2 files changed, 12 insertions(+), 10 deletions(-)

diff --git a/packages/server/src/ClientConnection.ts b/packages/server/src/ClientConnection.ts
index d61a73e39..0c5cab811 100644
--- a/packages/server/src/ClientConnection.ts
+++ b/packages/server/src/ClientConnection.ts
@@ -163,15 +163,6 @@ export class ClientConnection {
       return this.hooks('beforeHandleMessage', beforeHandleMessagePayload)
     })
 
-    // If the WebSocket has already disconnected (wow, that was fast) – then
-    // immediately call close to cleanup the connection and document in memory.
-    if (
-      connection.readyState === WsReadyStates.Closing
-      || connection.readyState === WsReadyStates.Closed
-    ) {
-      instance.close()
-    }
-
     return instance
   }
 
@@ -198,13 +189,23 @@ export class ClientConnection {
 
     this.documentConnections[documentName] = instance
 
+    // If the WebSocket has already disconnected (wow, that was fast) – then
+    // immediately call close to cleanup the connection and document in memory.
+    if (
+      this.websocket.readyState === WsReadyStates.Closing
+      || this.websocket.readyState === WsReadyStates.Closed
+    ) {
+      instance.close()
+      return
+    }
+
     // There’s no need to queue messages anymore.
     // Let’s work through queued messages.
     this.incomingMessageQueue[documentName].forEach(input => {
       this.websocket.emit('message', input)
     })
 
-    this.hooks('connected', {
+    await this.hooks('connected', {
       ...hookPayload,
       documentName,
       context: hookPayload.context,
diff --git a/packages/server/src/Connection.ts b/packages/server/src/Connection.ts
index 592fd9a70..0cc02ec24 100644
--- a/packages/server/src/Connection.ts
+++ b/packages/server/src/Connection.ts
@@ -120,6 +120,7 @@ export class Connection {
       || this.webSocket.readyState === WsReadyStates.Closed
     ) {
       this.close()
+      return
     }
 
     try {