From f50f168e8f67af9e13072bb80b63b4543d2f8f00 Mon Sep 17 00:00:00 2001
From: smworld01 <32947380+smworld01@users.noreply.github.com>
Date: Thu, 12 Sep 2024 03:28:55 +0900
Subject: [PATCH] [lexical-react] Refactor: Ensure disconnect is called after
 connection is established in useYjsCollaboration (#6619)

Co-authored-by: Bob Ippolito <bob@redivi.com>
---
 .../lexical-react/src/shared/useYjsCollaboration.tsx  | 11 ++++++-----
 1 file changed, 6 insertions(+), 5 deletions(-)

diff --git a/packages/lexical-react/src/shared/useYjsCollaboration.tsx b/packages/lexical-react/src/shared/useYjsCollaboration.tsx
index 7017e96fec3..371a6f4721f 100644
--- a/packages/lexical-react/src/shared/useYjsCollaboration.tsx
+++ b/packages/lexical-react/src/shared/useYjsCollaboration.tsx
@@ -57,9 +57,7 @@ export function useYjsCollaboration(
 ): JSX.Element {
   const isReloadingDoc = useRef(false);
 
-  const connect = useCallback(() => {
-    provider.connect();
-  }, [provider]);
+  const connect = useCallback(() => provider.connect(), [provider]);
 
   const disconnect = useCallback(() => {
     try {
@@ -152,11 +150,14 @@ export function useYjsCollaboration(
         }
       },
     );
-    connect();
+
+    const connectionPromise = connect();
 
     return () => {
       if (isReloadingDoc.current === false) {
-        disconnect();
+        Promise.resolve(connectionPromise).then(() => {
+          disconnect();
+        });
       }
 
       provider.off('sync', onSync);