Bug: [lexical-y] A comment comes back over the rest of the document when a collaborator undoes removing it after its author undid the typing and comment together #9689

// @vitest-environment node

import {
  buildEditorFromExtensions,
  getExtensionDependencyFromEditor,
} from '@lexical/extension';
import {
  $addMarkID,
  $getNodeMarkIDs,
  $removeMarkID,
  MarkStateExtension,
} from '@lexical/mark';
import {RichTextExtension} from '@lexical/rich-text';
import * as Y from '@y/y';
import {
  $createParagraphNode,
  $createRangeSelection,
  $createTextNode,
  $getRoot,
  $isElementNode,
  $isTextNode,
  $setSelection,
  configExtension,
  defineExtension,
  type LexicalEditor,
  type RangeSelection,
  UNDO_COMMAND,
} from 'lexical';
import {expect, test} from 'vitest';

import {createYDoc, YCollabExtension} from '../..';
import {seedYState} from '../../headless';

const CONFIG = {extensions: [RichTextExtension, MarkStateExtension], schemaVersion: 1};
const PARAS = ['zero zero zero', 'one one one', 'two two two', 'three three three'];

const $show = () =>
  $getRoot()
    .getChildren()
    .map(b =>
      $isElementNode(b)
        ? b
            .getChildren()
            .map(c => {
              const ids = $isTextNode(c) ? $getNodeMarkIDs(MarkStateExtension, c) : [];
              return `${c.getTextContent()}${ids.length ? `[${ids.join(',')}]` : ''}`;
            })
            .join('|')
        : '?',
    );

function client(seeded: Uint8Array, id: number, gc: boolean) {
  const doc = createYDoc(seeded, {gc});
  doc.clientID = id;
  const editor = buildEditorFromExtensions(
    defineExtension({
      dependencies: [
        RichTextExtension,
        MarkStateExtension,
        configExtension(YCollabExtension, {captureTimeout: 2 ** 30, doc, schemaVersion: 1, synced: true}),
      ],
      name: `repro/${id}`,
    }),
  );
  const {binding} = getExtensionDependencyFromEditor(editor, YCollabExtension).output;
  return {doc, editor, step: () => binding.undoManager?.stopCapturing()};
}

function onRange(editor: LexicalEditor, a: number, b: number, f: (s: RangeSelection) => void) {
  editor.update(
    () => {
      const t = ($getRoot().getChildAtIndex(2) as ReturnType<typeof $createParagraphNode>).getFirstChild();
      if (!$isTextNode(t)) throw new Error('no text');
      const s = $createRangeSelection();
      s.anchor.set(t.getKey(), a, 'text');
      s.focus.set(t.getKey(), b, 'text');
      $setSelection(s);
      f(s);
      $setSelection(null);
    },
    {discrete: true},
  );
}

for (const gc of [false, true]) {
  for (const [idA, idB] of [[1, 2], [2, 1]]) {
    test(`the comment stays off the paragraphs after it (gc ${gc ? 'on' : 'off'}, ids ${idA}, ${idB})`, () => {
      const seeded = seedYState(CONFIG, () => {
        $getRoot().clear().append(...PARAS.map(t => $createParagraphNode().append($createTextNode(t))));
      });
      const A = client(seeded, idA, gc);
      const B = client(seeded, idB, gc);
      const sync = () => {
        Y.applyUpdate(B.doc, Y.encodeStateAsUpdate(A.doc, Y.encodeStateVector(B.doc)));
        Y.applyUpdate(A.doc, Y.encodeStateAsUpdate(B.doc, Y.encodeStateVector(A.doc)));
      };
      onRange(A.editor, 0, 0, s => {
        s.format = 0;
        s.insertText('yy ');
      });
      onRange(A.editor, 0, 3, s => $addMarkID('c1', s));
      A.step();
      sync();
      B.editor.update(() => $removeMarkID('c1'), {discrete: true});
      B.step();
      A.editor.dispatchCommand(UNDO_COMMAND, undefined); // before A sees B's removal
      sync();
      B.editor.dispatchCommand(UNDO_COMMAND, undefined);
      sync();
      expect(A.editor.getEditorState().read($show)).toEqual(PARAS);
      expect(B.editor.getEditorState().read($show)).toEqual(PARAS);
    });
  }
}