Bug: [lexical-y] A link a collaborator removes comes back when another collaborator sets NodeState on it at the same time #9629

/**
 * 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.
 *
 */

// @vitest-environment node

import {buildEditorFromExtensions} from '@lexical/extension';
import {$createLinkNode, $isLinkNode, LinkExtension} from '@lexical/link';
import {RichTextExtension} from '@lexical/rich-text';
import * as Y from '@y/y';
import {
  $createParagraphNode,
  $createTextNode,
  $getRoot,
  $isElementNode,
  $setState,
  configExtension,
  createState,
  defineExtension,
  type ElementNode,
  type LexicalEditor,
} from 'lexical';
import {expect, test} from 'vitest';

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

// NodeState an app keeps on a link (a note, a tracking id).
const note = createState('repro-note', {
  parse: (v: unknown) => (typeof v === 'string' ? v : ''),
});

const app = defineExtension({
  dependencies: [RichTextExtension, LinkExtension],
  name: 'repro-app',
});
const config = {extensions: [app], onWarn: () => {}, schemaVersion: 1};

// "before " + link(https://note.test: "linked words") + " after"
function $build() {
  $getRoot()
    .clear()
    .append(
      $createParagraphNode().append(
        $createTextNode('before '),
        $createLinkNode('https://note.test').append(
          $createTextNode('linked words'),
        ),
        $createTextNode(' after'),
      ),
    );
}

function client(seed: Uint8Array, id: number) {
  const doc = createYDoc(seed);
  doc.clientID = id;
  const editor = buildEditorFromExtensions(
    defineExtension({
      $initialEditorState: null,
      dependencies: [
        app,
        configExtension(YCollabExtension, {doc, schemaVersion: 1}),
      ],
      name: `repro-${id}`,
    }),
  );
  return {doc, editor};
}

function links(editor: LexicalEditor): number {
  return editor.read(() => {
    let n = 0;
    const visit = (node: ElementNode) => {
      for (const child of node.getChildren()) {
        if ($isLinkNode(child)) {
          n++;
        }
        if ($isElementNode(child)) {
          visit(child);
        }
      }
    };
    visit($getRoot());
    return n;
  });
}

for (const ids of [
  [1, 2],
  [2, 1],
]) {
  test(`a removed link stays removed when a collaborator sets NodeState on it (ids ${ids})`, async () => {
    const seed = seedYState(config, $build);
    const a = client(seed, ids[0]);
    const b = client(seed, ids[1]);
    await Promise.resolve();
    // A removes the link (the paragraph's children: text, link, text).
    a.editor.update(
      () => {
        const link = (
          $getRoot().getFirstChildOrThrow() as ElementNode
        ).getChildAtIndex(1);
        if (!$isLinkNode(link)) {
          throw new Error('link expected');
        }
        for (const child of link.getChildren()) {
          link.insertBefore(child);
        }
        link.remove();
      },
      {discrete: true},
    );
    expect(links(a.editor)).toBe(0);
    // B, at the same time, sets NodeState on the link.
    b.editor.update(
      () => {
        const link = (
          $getRoot().getFirstChildOrThrow() as ElementNode
        ).getChildAtIndex(1);
        if (!$isLinkNode(link)) {
          throw new Error('link expected');
        }
        $setState(link, note, 'checked');
      },
      {discrete: true},
    );
    const ua = Y.encodeStateAsUpdate(a.doc);
    const ub = Y.encodeStateAsUpdate(b.doc);
    Y.applyUpdate(a.doc, ub);
    Y.applyUpdate(b.doc, ua);
    await Promise.resolve();
    expect(links(a.editor)).toBe(0);
    expect(links(b.editor)).toBe(0);
  });
}