fix: yjs collaboration plugin in react strict mode

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;

packages/lexical-react/src/LexicalCollaborationPlugin.tsx

@@ -0,0 +1,180 @@
+/**
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ *
+ */
+
+import type {Doc} from 'yjs';
+
+import {
+  type CollaborationContextType,
+  useCollaborationContext,
+} from '@lexical/react/LexicalCollaborationContext';
+import {useLexicalComposerContext} from '@lexical/react/LexicalComposerContext';
+import {
+  Binding,
+  createBinding,
+  ExcludedProperties,
+  Provider,
+} from '@lexical/yjs';
+import {LexicalEditor} from 'lexical';
+import {useEffect, useState} from 'react';
+
+import {InitialEditorStateType} from './LexicalComposer';
+import {
+  CursorsContainerRef,
+  useYjsCollaboration,
+  useYjsFocusTracking,
+  useYjsHistory,
+} from './shared/useYjsCollaboration';
+
+type Props = {
+  id: string;
+  providerFactory: (
+    // eslint-disable-next-line no-shadow
+    id: string,
+    yjsDocMap: Map<string, Doc>,
+  ) => Provider;
+  shouldBootstrap: boolean;
+  username?: string;
+  cursorColor?: string;
+  cursorsContainerRef?: CursorsContainerRef;
+  initialEditorState?: InitialEditorStateType;
+  excludedProperties?: ExcludedProperties;
+  // `awarenessData` parameter allows arbitrary data to be added to the awareness.
+  awarenessData?: object;
+};
+
+export function CollaborationPlugin({
+  id,
+  providerFactory,
+  shouldBootstrap,
+  username,
+  cursorColor,
+  cursorsContainerRef,
+  initialEditorState,
+  excludedProperties,
+  awarenessData,
+}: Props): JSX.Element {
+  const collabContext = useCollaborationContext(username, cursorColor);
+
+  const {yjsDocMap, name, color} = collabContext;
+
+  const [editor] = useLexicalComposerContext();
+
+  useEffect(() => {
+    collabContext.isCollabActive = true;
+
+    return () => {
+      // Reseting flag only when unmount top level editor collab plugin. Nested
+      // editors (e.g. image caption) should unmount without affecting it
+      if (editor._parentEditor == null) {
+        collabContext.isCollabActive = false;
+      }
+    };
+  }, [collabContext, editor]);
+
+  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}
+    />
+  );
+}
+
+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,
+    awarenessData,
+  );
+  collabContext.clientID = binding.clientID;
+  useYjsHistory(editor, binding);
+  useYjsFocusTracking(editor, provider, name, color, awarenessData);
+  return cursors;
+}

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) => {