fix(joint-react): add tests for container key updates on membership changes (#3462)

Author: samuelgja

.changeset/container-membership-notifications.md

@@ -0,0 +1,5 @@
+---
+"@joint/react": patch
+---
+
+fix stale cell keys when a single commit swaps ids without changing the count (membership changes now notify key-list subscribers, and large-graph rendering no longer defers id updates)

packages/joint-react/src/components/graph/__tests__/graph-provider-imperative-churn.test.tsx

@@ -0,0 +1,171 @@
+import { Component, type ReactNode } from 'react';
+import { act, render, waitFor } from '@testing-library/react';
+import { dia } from '@joint/core';
+import { GraphProvider, Paper } from '../..';
+import { useCell } from '../../../hooks/use-cell';
+import { ELEMENT_MODEL_TYPE } from '../../../mvc/element-model';
+import { DEFAULT_CELL_NAMESPACE } from '../../../store/graph-store';
+import type { Computed, ElementRecord } from '../../../types/cell.types';
+
+interface NodeData {
+  readonly label: string;
+}
+
+// Surfaces an error thrown anywhere in its subtree so the test can assert on
+// it instead of the throw tearing down the whole test render.
+class CatchErrorBoundary extends Component<
+  Readonly<{ onCatch: (error: Error) => void; children: ReactNode }>,
+  { hasError: boolean }
+> {
+  state = { hasError: false };
+  static getDerivedStateFromError() {
+    return { hasError: true };
+  }
+  componentDidCatch(error: Error) {
+    this.props.onCatch(error);
+  }
+  render() {
+    return this.state.hasError ? null : this.props.children;
+  }
+}
+
+const caughtState: { error: Error | null } = { error: null };
+function captureError(error: Error) {
+  caughtState.error = error;
+}
+
+// Subscribes to its own cell — the same shape a content-sized node gets from
+// useMeasureElement, so the scenario needs no explicit useCell in app code.
+function SubscribingNode() {
+  const label = useCell((cell: Computed<ElementRecord<NodeData>>) => cell.data.label);
+  return <text>{label}</text>;
+}
+const renderSubscribing = () => <SubscribingNode />;
+
+const PAPER_STYLE = { width: 400, height: 400 } as const;
+
+function elementJSON(id: string, index = 0): dia.Cell.JSON {
+  return {
+    id,
+    type: ELEMENT_MODEL_TYPE,
+    position: { x: (index % 10) * 20, y: Math.floor(index / 10) * 20 },
+    size: { width: 10, height: 10 },
+    data: { label: `label-${id}` },
+  };
+}
+
+function createExternalGraph(): dia.Graph {
+  return new dia.Graph({}, { cellNamespace: DEFAULT_CELL_NAMESPACE });
+}
+
+// Customer scenario: the app owns an external dia.Graph and mutates it
+// imperatively (fromJSON load, stencil drop add→remove→re-add churn,
+// CommandManager undo/redo). React content must follow every membership
+// change of the graph — even when a coalesced commit keeps the cell COUNT
+// unchanged while swapping ids.
+describe('GraphProvider — imperative external-graph churn', () => {
+  afterEach(() => {
+    jest.restoreAllMocks();
+  });
+
+  it('paints a cell swapped in during a same-tick remove+add (stencil churn)', async () => {
+    const graph = createExternalGraph();
+    graph.addCell(elementJSON('a'));
+
+    const { container } = render(
+      <GraphProvider graph={graph}>
+        <Paper style={PAPER_STYLE} id="churn-swap-paper" renderElement={renderSubscribing} />
+      </GraphProvider>
+    );
+
+    await waitFor(() => {
+      expect(container.textContent).toContain('label-a');
+    });
+
+    // Same tick: remove 'a', add 'b'. Both mutations coalesce into a single
+    // container commit whose net size is unchanged.
+    await act(async () => {
+      graph.removeCells([graph.getCell('a')]);
+      graph.addCell(elementJSON('b', 1));
+    });
+
+    await waitFor(() => {
+      expect(container.textContent).toContain('label-b');
+      expect(container.textContent).not.toContain('label-a');
+    });
+  });
+
+  it('paints new cells after graph.fromJSON reload with the same cell count', async () => {
+    const graph = createExternalGraph();
+    graph.addCell(elementJSON('a'));
+    graph.addCell(elementJSON('b', 1));
+
+    const { container } = render(
+      <GraphProvider graph={graph}>
+        <Paper style={PAPER_STYLE} id="churn-fromjson-paper" renderElement={renderSubscribing} />
+      </GraphProvider>
+    );
+
+    await waitFor(() => {
+      expect(container.textContent).toContain('label-a');
+      expect(container.textContent).toContain('label-b');
+    });
+
+    // Imperative reload: same count, entirely different ids.
+    await act(async () => {
+      graph.fromJSON({ cells: [elementJSON('c'), elementJSON('d', 1)] });
+    });
+
+    await waitFor(() => {
+      expect(container.textContent).toContain('label-c');
+      expect(container.textContent).toContain('label-d');
+      expect(container.textContent).not.toContain('label-a');
+    });
+  });
+
+  it('survives delete + undo-style re-add at scale (120 cells)', async () => {
+    jest.spyOn(console, 'error').mockImplementation(() => {});
+    caughtState.error = null;
+
+    const graph = createExternalGraph();
+    const cellCount = 120;
+    for (let index = 0; index < cellCount; index++) {
+      graph.addCell(elementJSON(`el-${index}`, index));
+    }
+
+    const { container } = render(
+      <CatchErrorBoundary onCatch={captureError}>
+        <GraphProvider graph={graph}>
+          <Paper style={PAPER_STYLE} id="churn-undo-paper" renderElement={renderSubscribing} />
+        </GraphProvider>
+      </CatchErrorBoundary>
+    );
+
+    await waitFor(() => {
+      expect(container.textContent).toContain('label-el-119');
+    });
+
+    const removedJSON = graph.getCell('el-115').toJSON();
+
+    // Delete at scale — regression for the historical `useCell(): no cell with
+    // id` crash, where a stale render pass hit the already-removed cell.
+    await act(async () => {
+      graph.removeCells([graph.getCell('el-115')]);
+    });
+
+    await waitFor(() => {
+      expect(container.textContent).not.toContain('label-el-115');
+    });
+    expect(caughtState.error).toBeNull();
+
+    // Undo — CommandManager restores the cell by re-adding the same JSON.
+    await act(async () => {
+      graph.addCell(removedJSON);
+    });
+
+    await waitFor(() => {
+      expect(container.textContent).toContain('label-el-115');
+    });
+    expect(caughtState.error).toBeNull();
+  });
+});

packages/joint-react/src/hooks/__tests__/use-container-keys.test.tsx

@@ -56,6 +56,29 @@ describe('useContainerKeys', () => {
     expect(result.current).toBe(first);
   });
 
+  it('updates keys when one id is swapped for another in a single commit', async () => {
+    const container = createContainer<TestItem>();
+    const readOnly = asReadonlyContainer(container);
+    container.set('a', { id: 'a', value: 1, type: 'item' });
+    container.commitChanges();
+    await flush();
+
+    const { result } = renderHook(() => useContainerKeys(readOnly));
+    expect(result.current).toEqual(['a']);
+
+    // Remove 'a' and add 'b' in ONE commit — the count is unchanged but the
+    // key set is not. This is the shape of a stencil drop (temp cell swap) or
+    // an undo that replaces a cell in the same tick.
+    await act(async () => {
+      container.delete('a');
+      container.set('b', { id: 'b', value: 2, type: 'item' });
+      container.commitChanges();
+      await flush();
+    });
+
+    expect(result.current).toEqual(['b']);
+  });
+
   it('updates the array when a key is added or removed', async () => {
     const container = createContainer<TestItem>();
     const readOnly = asReadonlyContainer(container);

packages/joint-react/src/hooks/use-create-portal-paper.tsx

@@ -2,7 +2,6 @@
 /* eslint-disable @typescript-eslint/no-shadow */
 import { dia, util } from '@joint/core';
 import {
-  useDeferredValue,
   useEffect,
   useLayoutEffect,
   useMemo,
@@ -239,14 +238,8 @@ export function useCreatePortalPaper(
     return { elementIds: elements, linkIds: links };
   }, [allCellIds, graphStore]);
 
-  const deferredElementIdsRaw = useDeferredValue(elementIds);
-  const deferredLinkIdsRaw = useDeferredValue(linkIds);
-  const shouldDefer = elementIds.length > 100 || linkIds.length > 100;
   const featuresContext = useContext(PaperFeaturesContext);
 
-  const deferredElementIds = shouldDefer ? deferredElementIdsRaw : elementIds;
-  const deferredLinkIds = shouldDefer ? deferredLinkIdsRaw : linkIds;
-
   const selectPaperVersion = useMemo(() => createSelectPaperVersion(id), [id]);
 
   // Subscribe to paper version to trigger re-renders on view mount/unmount changes
@@ -463,7 +456,7 @@ export function useCreatePortalPaper(
     if (!hasRenderElement) {
       return null;
     }
-    return deferredElementIds.map((elementId) => {
+    return elementIds.map((elementId) => {
       const elementView = paperStore?.getElementView(elementId);
       if (!elementView?.paper) {
         return null;
@@ -508,7 +501,7 @@ export function useCreatePortalPaper(
     version,
     HTMLRendererContainer,
     areElementsMeasured,
-    deferredElementIds,
+    elementIds,
     hasRenderElement,
     paperStore,
     renderElement,
@@ -520,7 +513,7 @@ export function useCreatePortalPaper(
       return null;
     }
 
-    return deferredLinkIds.map((linkId) => {
+    return linkIds.map((linkId) => {
       const linkView = paperStore?.getLinkView(linkId);
       if (!linkView?.paper) {
         return null;
@@ -542,7 +535,7 @@ export function useCreatePortalPaper(
       );
     });
     // eslint-disable-next-line react-hooks/exhaustive-deps
-  }, [deferredLinkIds, version, hasRenderLink, paperStore, renderLink]);
+  }, [linkIds, version, hasRenderLink, paperStore, renderLink]);
 
   const content = useMemo(
     () => (

packages/joint-react/src/store/__tests__/state-container.test.ts

@@ -433,6 +433,46 @@ describe('createContainer', () => {
       expect(listener).not.toHaveBeenCalled();
     });
 
+    it('notifies when membership changes but the count stays the same (swap in one commit)', async () => {
+      const container = setup();
+      container.set('a', { id: 'a', x: 1, y: 2, type: 'item' });
+      container.commitChanges();
+      await flush();
+
+      const listener = jest.fn();
+      container.subscribeToSize(listener);
+
+      // One coalesced commit: remove 'a' + add 'b' → net size unchanged, but
+      // the key set changed. Subscribers building id lists must be notified.
+      container.delete('a');
+      container.set('b', { id: 'b', x: 3, y: 4, type: 'item' });
+      container.commitChanges();
+      await flush();
+
+      expect(listener).toHaveBeenCalledTimes(1);
+    });
+
+    it('notifies on reset to the same count with different ids', async () => {
+      const container = setup();
+      container.set('a', { id: 'a', x: 1, y: 2, type: 'item' });
+      container.set('b', { id: 'b', x: 3, y: 4, type: 'item' });
+      container.commitChanges();
+      await flush();
+
+      const listener = jest.fn();
+      container.subscribeToSize(listener);
+
+      // Same count, entirely new ids (e.g. graph.fromJSON reload).
+      container.reset([
+        { id: 'c', x: 5, y: 6, type: 'item' },
+        { id: 'd', x: 7, y: 8, type: 'item' },
+      ]);
+      container.commitChanges();
+      await flush();
+
+      expect(listener).toHaveBeenCalledTimes(1);
+    });
+
     it('returns an unsubscribe function', async () => {
       const container = setup();
       const listener = jest.fn();

packages/joint-react/src/store/state-container.ts

@@ -32,6 +32,7 @@ export interface ReadonlyContainer<Cell extends AnyCellRecord> {
   has: (id: CellId) => boolean;
   getSize: () => number;
   subscribe: (id: CellId, listener: () => void) => () => void;
+  /** Notifies on membership changes (ids added / removed), even when the net count is unchanged. */
   subscribeToSize: (listener: () => void) => () => void;
   subscribeToAll: (listener: () => void) => () => void;
 }
@@ -82,7 +83,11 @@ export function createContainer<Cell extends AnyCellRecord>(): Container<Cell> {
   // because `Set.add` is measurably slower than `Array.push` when many
   // unique ids accumulate between commits (each add does a hash lookup).
   let changes: CellId[] = [];
-  let previousSize = 0;
+  // True when the id SET changed since the last commit (add/remove/reset), not
+  // just a value. Gating size-listener notifications on a net count change
+  // would miss same-count membership swaps (remove A + add B in one commit,
+  // same-count reset), leaving key-list subscribers stale.
+  let hasMembershipChanged = false;
   let version = 0;
   return {
     get(id: CellId): Cell | undefined {
@@ -108,6 +113,7 @@ export function createContainer<Cell extends AnyCellRecord>(): Container<Cell> {
       if (index === undefined) {
         indexById.set(id, items.length);
         items.push(value);
+        hasMembershipChanged = true;
       } else {
         items[index] = value;
       }
@@ -130,6 +136,7 @@ export function createContainer<Cell extends AnyCellRecord>(): Container<Cell> {
       items.pop();
       indexById.delete(id);
       changes.push(id);
+      hasMembershipChanged = true;
       version++;
     },
     reset(next: readonly Cell[]) {
@@ -145,6 +152,9 @@ export function createContainer<Cell extends AnyCellRecord>(): Container<Cell> {
         changes.push(item.id as CellId);
         index++;
       }
+      // Reset is a cold path — always flag it. Even when the id set happens to
+      // be identical, key-list subscribers bail out on the stable keys array.
+      hasMembershipChanged = true;
       version++;
     },
     getVersion() {
@@ -169,8 +179,8 @@ export function createContainer<Cell extends AnyCellRecord>(): Container<Cell> {
           listener();
         }
       }
-      if (previousSize !== items.length) {
-        previousSize = items.length;
+      if (hasMembershipChanged) {
+        hasMembershipChanged = false;
         for (const listener of sizeListeners) {
           listener();
         }