fix: yjs collaboration plugin in react strict mode

facebook · Jun 7, 2024 · b700713 · b700713
1 parent da7c439
commit b700713
Show file tree

Hide file tree

Showing 3 changed files with 102 additions and 19 deletions.
diff --git a/packages/lexical-react/src/LexicalCollaborationContext.ts b/packages/lexical-react/src/LexicalCollaborationContext.ts
@@ -10,7 +10,7 @@ import type {Doc} from 'yjs';
 
 import {createContext, useContext} from 'react';
 
-type CollaborationContextType = {
+export type CollaborationContextType = {
   clientID: number;
   color: string;
   isCollabActive: boolean;

diff --git a/packages/lexical-react/src/LexicalCollaborationPlugin.tsx b/packages/lexical-react/src/LexicalCollaborationPlugin.tsx
@@ -8,10 +8,19 @@
 
 import type {Doc} from 'yjs';
 
-import {useCollaborationContext} from '@lexical/react/LexicalCollaborationContext';
+import {
+  type CollaborationContextType,
+  useCollaborationContext,
+} from '@lexical/react/LexicalCollaborationContext';
 import {useLexicalComposerContext} from '@lexical/react/LexicalComposerContext';
-import {ExcludedProperties, Provider} from '@lexical/yjs';
-import {useEffect, useMemo} from 'react';
+import {
+  Binding,
+  createBinding,
+  ExcludedProperties,
+  Provider,
+} from '@lexical/yjs';
+import {LexicalEditor} from 'lexical';
+import {useEffect, useState} from 'react';
 
 import {InitialEditorStateType} from './LexicalComposer';
 import {
@@ -67,22 +76,101 @@ export function CollaborationPlugin({
     };
   }, [collabContext, editor]);
 
-  const provider = useMemo(
-    () => providerFactory(id, yjsDocMap),
-    [id, providerFactory, yjsDocMap],
+  const [provider, setProvider] = useState<Provider>();
+  useEffect(() => {
+    const newProvider = providerFactory(id, yjsDocMap);
+    setProvider(newProvider);
+
+    return () => {
+      newProvider.disconnect();
+    };
+  }, [id, providerFactory, yjsDocMap]);
+
+  const [doc, setDoc] = useState(yjsDocMap.get(id));
+  const [binding, setBinding] = useState<Binding>();
+  useEffect(() => {
+    if (!provider) {
+      return;
+    }
+    const newBinding = createBinding(
+      editor,
+      provider,
+      id,
+      doc ?? yjsDocMap.get(id),
+      yjsDocMap,
+      excludedProperties,
+    );
+    setBinding(newBinding);
+
+    return () => {
+      newBinding.root.destroy(newBinding);
+    };
+  }, [editor, provider, id, yjsDocMap, doc, excludedProperties]);
+
+  if (!provider || !binding) {
+    return <></>;
+  }
+
+  return (
+    <WrapWithProvider
+      awarenessData={awarenessData}
+      binding={binding}
+      collabContext={collabContext}
+      color={color}
+      cursorsContainerRef={cursorsContainerRef}
+      editor={editor}
+      id={id}
+      initialEditorState={initialEditorState}
+      name={name}
+      provider={provider}
+      setDoc={setDoc}
+      shouldBootstrap={shouldBootstrap}
+      yjsDocMap={yjsDocMap}
+    />
   );
+}
 
-  const [cursors, binding] = useYjsCollaboration(
+function WrapWithProvider({
+  editor,
+  id,
+  provider,
+  yjsDocMap,
+  name,
+  color,
+  shouldBootstrap,
+  cursorsContainerRef,
+  initialEditorState,
+  awarenessData,
+  collabContext,
+  binding,
+  setDoc,
+}: {
+  editor: LexicalEditor;
+  id: string;
+  provider: Provider;
+  yjsDocMap: Map<string, Doc>;
+  name: string;
+  color: string;
+  shouldBootstrap: boolean;
+  binding: Binding;
+  setDoc: React.Dispatch<React.SetStateAction<Doc | undefined>>;
+  cursorsContainerRef?: CursorsContainerRef | undefined;
+  initialEditorState?: InitialEditorStateType | undefined;
+  awarenessData?: object;
+  collabContext: CollaborationContextType;
+}) {
+  const [cursors] = useYjsCollaboration(
     editor,
     id,
     provider,
     yjsDocMap,
     name,
     color,
     shouldBootstrap,
+    binding,
+    setDoc,
     cursorsContainerRef,
     initialEditorState,
-    excludedProperties,
     awarenessData,
   );
 

diff --git a/packages/lexical-react/src/shared/useYjsCollaboration.tsx b/packages/lexical-react/src/shared/useYjsCollaboration.tsx
@@ -6,13 +6,12 @@
  *
  */
 
-import type {Binding, ExcludedProperties, Provider} from '@lexical/yjs';
+import type {Binding, Provider} from '@lexical/yjs';
 import type {LexicalEditor} from 'lexical';
 
 import {mergeRegister} from '@lexical/utils';
 import {
   CONNECTED_COMMAND,
-  createBinding,
   createUndoManager,
   initLocalState,
   setLocalStateFocus,
@@ -34,7 +33,7 @@ import {
   UNDO_COMMAND,
 } from 'lexical';
 import * as React from 'react';
-import {useCallback, useEffect, useMemo, useRef, useState} from 'react';
+import {useCallback, useEffect, useMemo, useRef} from 'react';
 import {createPortal} from 'react-dom';
 import {Doc, Transaction, UndoManager, YEvent} from 'yjs';
 
@@ -50,18 +49,13 @@ export function useYjsCollaboration(
   name: string,
   color: string,
   shouldBootstrap: boolean,
+  binding: Binding,
+  setDoc: React.Dispatch<React.SetStateAction<Doc | undefined>>,
   cursorsContainerRef?: CursorsContainerRef,
   initialEditorState?: InitialEditorStateType,
-  excludedProperties?: ExcludedProperties,
   awarenessData?: object,
 ): [JSX.Element, Binding] {
   const isReloadingDoc = useRef(false);
-  const [doc, setDoc] = useState(docMap.get(id));
-
-  const binding = useMemo(
-    () => createBinding(editor, provider, id, doc, docMap, excludedProperties),
-    [editor, provider, id, docMap, doc, excludedProperties],
-  );
 
   const connect = useCallback(() => {
     provider.connect();
@@ -186,6 +180,7 @@ export function useYjsCollaboration(
     provider,
     shouldBootstrap,
     awarenessData,
+    setDoc,
   ]);
   const cursorsContainer = useMemo(() => {
     const ref = (element: null | HTMLElement) => {