[lexical] Bug Fix: Set activeEditorState when using editor.setEditorState inside of an update

packages/lexical/src/LexicalEditor.ts

@@ -19,12 +19,13 @@ import invariant from 'shared/invariant';
 
 import {$getRoot, $getSelection, TextNode} from '.';
 import {FULL_RECONCILE, NO_DIRTY_NODES} from './LexicalConstants';
-import {cloneEditorState, createEmptyEditorState} from './LexicalEditorState';
+import {createEmptyEditorState} from './LexicalEditorState';
 import {addRootElementEvents, removeRootElementEvents} from './LexicalEvents';
 import {$flushRootMutations, initMutationObserver} from './LexicalMutations';
 import {LexicalNode} from './LexicalNode';
 import {
   $commitPendingUpdates,
+  INTERNAL_$setEditorState,
   internalGetActiveEditor,
   parseEditorState,
   triggerListeners,
@@ -1140,44 +1141,8 @@ export class LexicalEditor {
         "setEditorState: the editor state is empty. Ensure the editor state's root node never becomes empty.",
       );
     }
-
-    // Ensure that we have a writable EditorState so that transforms can run
-    // during a historic operation
-    let writableEditorState = editorState;
-    if (writableEditorState._readOnly) {
-      writableEditorState = cloneEditorState(editorState);
-      writableEditorState._selection = editorState._selection
-        ? editorState._selection.clone()
-        : null;
-    }
-
     $flushRootMutations(this);
-    const pendingEditorState = this._pendingEditorState;
-    const tags = this._updateTags;
-    const tag = options !== undefined ? options.tag : null;
-
-    if (pendingEditorState !== null && !pendingEditorState.isEmpty()) {
-      if (tag != null) {
-        tags.add(tag);
-      }
-      $commitPendingUpdates(this);
-    }
-
-    this._pendingEditorState = writableEditorState;
-    this._dirtyType = FULL_RECONCILE;
-    this._dirtyElements.set('root', false);
-    this._compositionKey = null;
-
-    if (tag != null) {
-      tags.add(tag);
-    }
-
-    // Only commit pending updates if not already in an editor.update
-    // (e.g. dispatchCommand) otherwise this will cause a second commit
-    // with an already read-only state and selection
-    if (!this._updating) {
-      $commitPendingUpdates(this);
-    }
+    INTERNAL_$setEditorState(editorState, this, options);
   }
 
   /**

packages/lexical/src/LexicalUpdates.ts

@@ -11,7 +11,12 @@ import type {LexicalNode, SerializedLexicalNode} from './LexicalNode';
 
 import invariant from 'shared/invariant';
 
-import {$isElementNode, $isTextNode, SELECTION_CHANGE_COMMAND} from '.';
+import {
+  $isElementNode,
+  $isTextNode,
+  EditorSetOptions,
+  SELECTION_CHANGE_COMMAND,
+} from '.';
 import {FULL_RECONCILE, NO_DIRTY_NODES} from './LexicalConstants';
 import {
   CommandPayloadType,
@@ -390,6 +395,62 @@ function $parseSerializedNodeImpl<
   return node;
 }
 
+export function INTERNAL_$setEditorState(
+  editorState: EditorState,
+  editor: LexicalEditor,
+  options?: EditorSetOptions,
+): void {
+  if (editorState.isEmpty()) {
+    invariant(
+      false,
+      "setEditorState: the editor state is empty. Ensure the editor state's root node never becomes empty.",
+    );
+  }
+
+  // Ensure that we have a writable EditorState so that transforms can run
+  // during a historic operation
+  let writableEditorState = editorState;
+  if (writableEditorState._readOnly) {
+    writableEditorState = cloneEditorState(editorState);
+    writableEditorState._selection = editorState._selection
+      ? editorState._selection.clone()
+      : null;
+  }
+
+  const pendingEditorState = editor._pendingEditorState;
+  const tags = editor._updateTags;
+  const tag = options !== undefined ? options.tag : null;
+
+  if (pendingEditorState !== null && !pendingEditorState.isEmpty()) {
+    if (tag != null) {
+      tags.add(tag);
+    }
+    $commitPendingUpdates(editor);
+  }
+
+  editor._pendingEditorState = writableEditorState;
+  editor._dirtyType = FULL_RECONCILE;
+  editor._dirtyElements.set('root', false);
+  editor._compositionKey = null;
+
+  if (tag != null) {
+    tags.add(tag);
+  }
+
+  // Only commit pending updates if not already in an editor.update
+  // (e.g. dispatchCommand) otherwise this will cause a second commit
+  // with an already read-only state and selection
+  if (!editor._updating) {
+    $commitPendingUpdates(editor);
+  } else {
+    invariant(
+      activeEditorState === pendingEditorState,
+      'setEditorState: The previous activeEditorState must be the pendingEditorState',
+    );
+    activeEditorState = writableEditorState;
+  }
+}
+
 export function parseEditorState(
   serializedEditorState: SerializedEditorState,
   editor: LexicalEditor,

packages/lexical/src/__tests__/unit/LexicalEditor.test.tsx

@@ -2568,6 +2568,31 @@ describe('LexicalEditor tests', () => {
     expect(editor._pendingEditorState).toBe(null);
   });
 
+  it('sets the EditorState from a deferred update', async () => {
+    editor = createTestEditor({});
+    const state = editor.parseEditorState(
+      `{"root":{"children":[{"children":[{"detail":0,"format":0,"mode":"normal","style":"","text":"Hello world","type":"text","version":1}],"direction":"ltr","format":"","indent":0,"type":"paragraph","version":1}],"direction":"ltr","format":"","indent":0,"type":"root","version":1}}`,
+    );
+    editor.update(() => {
+      $getRoot().clear().append($createParagraphNode());
+    });
+    editor.update(() => {
+      expect($getRoot().getTextContent()).toBe('');
+      editor.setEditorState(state);
+      // Ensure that the activeEditorState has changed accordingly
+      expect($getRoot().getTextContent()).toBe('Hello world');
+    });
+    await editor.update(() => {
+      // This happens before the update is reconciled
+      expect($getRoot().getTextContent()).toBe('Hello world');
+    });
+    expect(editor._editorState.toJSON()).toEqual(state.toJSON());
+    expect(editor._pendingEditorState).toBe(null);
+    expect(
+      editor.getEditorState().read(() => $getRoot().getTextContent()),
+    ).toBe('Hello world');
+  });
+
   describe('node replacement', () => {
     it('should work correctly', async () => {
       const onError = jest.fn();