perf(apollo-react): batch BaseNode's updateNodeInternals calls per React Flow store [ECS-2160]

packages/apollo-react/src/canvas/components/AgentCanvas/nodes/AgentNode.test.tsx

@@ -24,6 +24,7 @@ vi.mock('@uipath/apollo-react/canvas/xyflow/react', () => ({
   useStore: () => ({ edges: [], isConnecting: false }),
   useConnection: () => ({ inProgress: false }),
   useUpdateNodeInternals: () => vi.fn(),
+  useStoreApi: () => ({ getState: vi.fn() }),
   useReactFlow: () => ({
     setNodes: vi.fn(),
     setEdges: vi.fn(),

packages/apollo-react/src/canvas/components/AgentCanvas/nodes/ResourceNode.test.tsx

@@ -12,6 +12,7 @@ vi.mock('@uipath/apollo-react/canvas/xyflow/react', async (importOriginal) => ({
   useStore: () => ({ edges: [], isConnecting: false }),
   useConnection: () => ({ inProgress: false }),
   useUpdateNodeInternals: () => vi.fn(),
+  useStoreApi: () => ({ getState: vi.fn() }),
   useReactFlow: () => ({
     setNodes: vi.fn(),
     setEdges: vi.fn(),

packages/apollo-react/src/canvas/components/BaseNode/BaseNode.test.tsx

@@ -14,6 +14,8 @@ const {
   mockUpdateNode,
   mockUpdateNodeData,
   mockGetNode,
+  mockUpdateNodeInternals,
+  mockStoreApi,
   mockNodeState,
   mockHandleConfigs,
   mockManifest,
@@ -30,6 +32,8 @@ const {
   mockUpdateNode: vi.fn(),
   mockUpdateNodeData: vi.fn(),
   mockGetNode: vi.fn(),
+  mockUpdateNodeInternals: vi.fn(),
+  mockStoreApi: { getState: vi.fn() },
   // Simulates React Flow's stored node: updateNode writes height here, getNode reads it.
   mockNodeState: { current: { height: undefined as number | undefined } },
   mockHandleConfigs: { current: undefined as HandleGroupManifest[] | undefined },
@@ -66,7 +70,8 @@ mockUpdateNode.mockImplementation((_id: string, patch: { height?: number }) => {
 vi.mock('@uipath/apollo-react/canvas/xyflow/react', async (importOriginal) => ({
   ...(await importOriginal<typeof import('@uipath/apollo-react/canvas/xyflow/react')>()),
   useStore: () => mockIsConnecting.current,
-  useUpdateNodeInternals: () => vi.fn(),
+  useUpdateNodeInternals: () => mockUpdateNodeInternals,
+  useStoreApi: () => mockStoreApi,
   useReactFlow: () => ({
     updateNodeData: mockUpdateNodeData,
     updateNode: mockUpdateNode,
@@ -194,6 +199,7 @@ describe('BaseNode', () => {
     mockUpdateNode.mockClear();
     mockUpdateNodeData.mockClear();
     mockGetNode.mockClear();
+    mockUpdateNodeInternals.mockClear();
     mockNodeState.current = { height: undefined };
     mockUseButtonHandles.mockClear();
     mockNodeToolbar.mockClear();
@@ -229,6 +235,39 @@ describe('BaseNode', () => {
     });
   });
 
+  describe('Node internals', () => {
+    it('updates the internals of nodes mounted together in one call', async () => {
+      render(
+        <>
+          <BaseNode {...defaultProps} id="a" />
+          <BaseNode {...defaultProps} id="b" />
+          <BaseNode {...defaultProps} id="c" />
+        </>
+      );
+      await act(async () => {});
+
+      expect(mockUpdateNodeInternals).toHaveBeenCalledTimes(1);
+      expect(mockUpdateNodeInternals).toHaveBeenCalledWith(['a', 'b', 'c']);
+    });
+
+    it('updates the internals again when the handles change after mount', async () => {
+      mockHandleConfigs.current = makeHandles(Position.Left, 2);
+      const { rerender } = render(<BaseNode {...defaultProps} />);
+      await act(async () => {});
+      mockUpdateNodeInternals.mockClear();
+
+      mockHandleConfigs.current = [
+        ...makeHandles(Position.Left, 2),
+        ...makeHandles(Position.Right, 1),
+      ];
+      rerender(<BaseNode {...defaultProps} data={{}} />);
+      await act(async () => {});
+
+      expect(mockUpdateNodeInternals).toHaveBeenCalledTimes(1);
+      expect(mockUpdateNodeInternals).toHaveBeenCalledWith(['test-node']);
+    });
+  });
+
   // The handle count sets `--node-h` (the node height), which is also written to
   // React Flow's node.height so measured height, edges, and handle spacing agree.
   // It is a pure function of handles/footer (never reads the measured height), so the

packages/apollo-react/src/canvas/components/BaseNode/BaseNode.tsx

@@ -1,10 +1,5 @@
 import type { Node, NodeProps } from '@uipath/apollo-react/canvas/xyflow/react';
-import {
-  Position,
-  useReactFlow,
-  useStore,
-  useUpdateNodeInternals,
-} from '@uipath/apollo-react/canvas/xyflow/react';
+import { Position, useReactFlow, useStore } from '@uipath/apollo-react/canvas/xyflow/react';
 import { memo, useCallback, useEffect, useMemo, useRef, useState } from 'react';
 import {
   DEFAULT_NODE_SIZE,
@@ -29,6 +24,7 @@ import {
   useNodeExecutionState,
   useToolbarActionStoreContext,
 } from '../../hooks';
+import { useBatchedUpdateNodeInternals } from '../../hooks/useBatchedUpdateNodeInternals';
 import { isWideNodeShape, type NodeShape } from '../../schema';
 import type { HandleGroupManifest } from '../../schema/node-definition';
 import { resolveAdornments } from '../../utils/adornment-resolver';
@@ -122,7 +118,7 @@ const BaseNodeComponent = (props: NodeProps<Node<BaseNodeData>>) => {
   const suggestionType =
     suggestionTypeOverride === null ? undefined : (suggestionTypeOverride ?? data.suggestionType);
 
-  const updateNodeInternals = useUpdateNodeInternals();
+  const updateNodeInternals = useBatchedUpdateNodeInternals();
   const { updateNodeData, updateNode, getNode } = useReactFlow();
   const containerRef = useRef<HTMLDivElement>(null);
   const [isHovered, setIsHovered] = useState(false);

packages/apollo-react/src/canvas/components/ButtonHandle/SmartHandle.test.tsx

@@ -1,4 +1,4 @@
-import { render, screen } from '@testing-library/react';
+import { act, render, screen } from '@testing-library/react';
 import { Position } from '@uipath/apollo-react/canvas/xyflow/react';
 import { beforeEach, describe, expect, it, vi } from 'vitest';
 import {
@@ -11,6 +11,8 @@ import {
 // Mock @xyflow/react
 const mockUseStore = vi.fn();
 const mockUseNodeId = vi.fn();
+const mockUpdateNodeInternals = vi.fn();
+const mockStoreApi = { getState: vi.fn() };
 
 vi.mock('@xyflow/react', () => ({
   Handle: ({ children, ...props }: any) => {
@@ -42,9 +44,41 @@ vi.mock('@xyflow/react', () => ({
   },
   useStore: (selector: any) => mockUseStore(selector),
   useNodeId: () => mockUseNodeId(),
-  useUpdateNodeInternals: () => vi.fn(),
+  useUpdateNodeInternals: () => mockUpdateNodeInternals,
 }));
 
+vi.mock('@uipath/apollo-react/canvas/xyflow/react', async (importOriginal) => ({
+  ...(await importOriginal<typeof import('@uipath/apollo-react/canvas/xyflow/react')>()),
+  useUpdateNodeInternals: () => mockUpdateNodeInternals,
+  useStoreApi: () => mockStoreApi,
+}));
+
+const flushFrame = () =>
+  act(() => new Promise<void>((resolve) => requestAnimationFrame(() => resolve())));
+
+function connectedStore(targetX: number, targetY: number) {
+  return (selector: (state: unknown) => unknown) =>
+    selector({
+      nodes: [
+        { id: 'test-node', position: { x: 100, y: 100 }, measured: { width: 100, height: 50 } },
+        {
+          id: 'target-node',
+          position: { x: targetX, y: targetY },
+          measured: { width: 100, height: 50 },
+        },
+      ],
+      edges: [
+        {
+          id: 'e1',
+          source: 'test-node',
+          sourceHandle: 'output',
+          target: 'target-node',
+          targetHandle: 'input',
+        },
+      ],
+    });
+}
+
 describe('SmartHandle', () => {
   beforeEach(() => {
     vi.clearAllMocks();
@@ -377,6 +411,50 @@ describe('SmartHandle', () => {
     });
   });
 
+  describe('Node internals', () => {
+    it('updates the internals of nodes mounted together in one call', async () => {
+      mockUseNodeId.mockReturnValueOnce('a').mockReturnValueOnce('b').mockReturnValueOnce('c');
+
+      render(
+        <>
+          <SmartHandle type="source" id="output" />
+          <SmartHandle type="source" id="output" />
+          <SmartHandle type="source" id="output" />
+        </>
+      );
+      await flushFrame();
+
+      expect(mockUpdateNodeInternals).toHaveBeenCalledTimes(1);
+      expect(mockUpdateNodeInternals).toHaveBeenCalledWith(['a', 'b', 'c']);
+    });
+
+    it('updates the internals again when the computed position changes', async () => {
+      mockUseStore.mockImplementation(connectedStore(400, 100));
+      const { rerender } = render(<SmartHandle type="source" id="output" />);
+      await flushFrame();
+      mockUpdateNodeInternals.mockClear();
+
+      mockUseStore.mockImplementation(connectedStore(100, 400));
+      rerender(<SmartHandle type="source" id="output" />);
+      await flushFrame();
+
+      expect(screen.getByTestId('smart-handle')).toHaveAttribute('data-position', Position.Bottom);
+      expect(mockUpdateNodeInternals).toHaveBeenCalledTimes(1);
+      expect(mockUpdateNodeInternals).toHaveBeenCalledWith(['test-node']);
+    });
+
+    it('does not update the internals when a rerender keeps the position', async () => {
+      const { rerender } = render(<SmartHandle type="source" id="output" />);
+      await flushFrame();
+      mockUpdateNodeInternals.mockClear();
+
+      rerender(<SmartHandle type="source" id="output" label="Output" />);
+      await flushFrame();
+
+      expect(mockUpdateNodeInternals).not.toHaveBeenCalled();
+    });
+  });
+
   describe('Edge cases', () => {
     it('handles missing nodeId gracefully', () => {
       mockUseNodeId.mockReturnValue(null);

packages/apollo-react/src/canvas/components/ButtonHandle/SmartHandle.tsx

@@ -7,7 +7,6 @@ import {
   type ReactFlowState,
   useNodeId,
   useStore,
-  useUpdateNodeInternals,
 } from '@xyflow/react';
 import {
   createContext,
@@ -18,6 +17,7 @@ import {
   useRef,
   useState,
 } from 'react';
+import { useBatchedUpdateNodeInternals } from '../../hooks/useBatchedUpdateNodeInternals';
 import { canvasEventBus } from '../../utils/CanvasEventBus';
 import { cx } from '../../utils/CssUtil';
 import { calculateGridAlignedHandlePositions } from '../../utils/handle-positioning';
@@ -564,7 +564,7 @@ export function SmartHandle({
   ...rest
 }: SmartHandleProps) {
   const nodeId = useNodeId();
-  const updateNodeInternals = useUpdateNodeInternals();
+  const updateNodeInternals = useBatchedUpdateNodeInternals();
   const [isHovered, setIsHovered] = useState(false);
   const smartHandleContext = useSmartHandleContext();
 
@@ -618,12 +618,7 @@ export function SmartHandle({
     // Only update if position actually changed
     if (prevPositionRef.current !== computedPosition) {
       prevPositionRef.current = computedPosition;
-
-      // Use requestAnimationFrame to ensure the DOM has updated
-      // before telling React Flow to recalculate
-      requestAnimationFrame(() => {
-        updateNodeInternals(nodeId);
-      });
+      updateNodeInternals(nodeId);
     }
   }, [nodeId, computedPosition, updateNodeInternals]);
 

packages/apollo-react/src/canvas/components/ContainerNode/ContainerNode.hooks.ts

@@ -1,10 +1,7 @@
-import {
-  type ReactFlowState,
-  useStore,
-  useUpdateNodeInternals,
-} from '@uipath/apollo-react/canvas/xyflow/react';
+import { type ReactFlowState, useStore } from '@uipath/apollo-react/canvas/xyflow/react';
 import { useCallback, useEffect } from 'react';
 import { shallow } from 'zustand/shallow';
+import { useBatchedUpdateNodeInternals } from '../../hooks/useBatchedUpdateNodeInternals';
 import type { HandleGroupManifest } from '../../schema/node-definition';
 import { type ContainerResizeMinimums, getContainerResizeMinimums } from '../../utils/container';
 import { DEFAULT_RESIZE_MINIMUMS } from './ContainerNode.constants';
@@ -60,17 +57,11 @@ export function useContainerNodeInternalsRefresh(
   width: number,
   height: number
 ) {
-  const updateNodeInternals = useUpdateNodeInternals();
+  const updateNodeInternals = useBatchedUpdateNodeInternals();
 
   // biome-ignore lint/correctness/useExhaustiveDependencies: dimensions and resolved handle groups are intentional rerun triggers for React Flow internals recalculation
   useEffect(() => {
-    const frameId = requestAnimationFrame(() => {
-      updateNodeInternals(id);
-    });
-
-    return () => {
-      cancelAnimationFrame(frameId);
-    };
+    updateNodeInternals(id);
   }, [id, handleGroups, updateNodeInternals, width, height]);
 }
 

packages/apollo-react/src/canvas/components/ContainerNode/ContainerNode.test.tsx

@@ -1,4 +1,4 @@
-import { fireEvent, render, screen } from '@testing-library/react';
+import { act, fireEvent, render, screen } from '@testing-library/react';
 import type { Node, NodeProps } from '@uipath/apollo-react/canvas/xyflow/react';
 import { beforeEach, describe, expect, it, vi } from 'vitest';
 import type { ContainerNodeData } from './ContainerNode.types';
@@ -11,6 +11,8 @@ const {
   mockManifest,
   mockNodeToolbar,
   mockReadOnlyNodeIds,
+  mockStoreApi,
+  mockUpdateNodeInternals,
   mockValidationState,
 } = vi.hoisted(() => ({
   mockUpdateNodeData: vi.fn(),
@@ -32,6 +34,8 @@ const {
   // biome-ignore lint/suspicious/noExplicitAny: captures the toolbar props for assertions
   mockNodeToolbar: vi.fn() as any,
   mockReadOnlyNodeIds: { current: new Set<string>() as ReadonlySet<string> },
+  mockStoreApi: { getState: vi.fn() },
+  mockUpdateNodeInternals: vi.fn(),
   mockValidationState: { current: undefined as unknown },
 }));
 
@@ -77,7 +81,8 @@ vi.mock('@uipath/apollo-react/canvas/xyflow/react', async (importOriginal) => ({
 
     return typeof selector === 'function' ? selector(state) : false;
   },
-  useUpdateNodeInternals: () => vi.fn(),
+  useUpdateNodeInternals: () => mockUpdateNodeInternals,
+  useStoreApi: () => mockStoreApi,
   useReactFlow: () => ({ updateNodeData: mockUpdateNodeData }),
 }));
 
@@ -180,6 +185,7 @@ function setContainerDisplay(container: Record<string, unknown> | undefined) {
 
 beforeEach(() => {
   mockUpdateNodeData.mockClear();
+  mockUpdateNodeInternals.mockClear();
   mockExecutionState.current = undefined;
   mockReadOnlyNodeIds.current = new Set<string>();
   mockValidationState.current = undefined;
@@ -218,6 +224,37 @@ describe('ContainerNode root', () => {
   });
 });
 
+describe('ContainerNode internals', () => {
+  const flushFrame = () =>
+    act(() => new Promise<void>((resolve) => requestAnimationFrame(() => resolve())));
+
+  it('updates the internals of containers mounted together in one call', async () => {
+    render(
+      <>
+        <ContainerNode {...defaultProps} id="a" />
+        <ContainerNode {...defaultProps} id="b" />
+        <ContainerNode {...defaultProps} id="c" />
+      </>
+    );
+    await flushFrame();
+
+    expect(mockUpdateNodeInternals).toHaveBeenCalledTimes(1);
+    expect(mockUpdateNodeInternals).toHaveBeenCalledWith(['a', 'b', 'c']);
+  });
+
+  it('updates the internals again when the container is resized', async () => {
+    const { rerender } = render(<ContainerNode {...defaultProps} width={400} height={300} />);
+    await flushFrame();
+    mockUpdateNodeInternals.mockClear();
+
+    rerender(<ContainerNode {...defaultProps} width={500} height={300} />);
+    await flushFrame();
+
+    expect(mockUpdateNodeInternals).toHaveBeenCalledTimes(1);
+    expect(mockUpdateNodeInternals).toHaveBeenCalledWith(['container-1']);
+  });
+});
+
 describe('ContainerNode header', () => {
   it('renders no badge by default', () => {
     renderContainer();

packages/apollo-react/src/canvas/components/LoopNode/LoopNode.test.tsx

@@ -77,6 +77,7 @@ vi.mock('@uipath/apollo-react/canvas/xyflow/react', async (importOriginal) => ({
     return typeof selector === 'function' ? selector(state) : false;
   },
   useUpdateNodeInternals: () => vi.fn(),
+  useStoreApi: () => ({ getState: vi.fn() }),
   useReactFlow: () => ({ updateNodeData: vi.fn() }),
 }));