diff --git a/packages/apollo-react/src/canvas/components/AgentCanvas/nodes/AgentNode.test.tsx b/packages/apollo-react/src/canvas/components/AgentCanvas/nodes/AgentNode.test.tsx index dd18939b22..4b40f25f30 100644 --- a/packages/apollo-react/src/canvas/components/AgentCanvas/nodes/AgentNode.test.tsx +++ b/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(), diff --git a/packages/apollo-react/src/canvas/components/AgentCanvas/nodes/ResourceNode.test.tsx b/packages/apollo-react/src/canvas/components/AgentCanvas/nodes/ResourceNode.test.tsx index 743c12b7b8..d7cefbde2b 100644 --- a/packages/apollo-react/src/canvas/components/AgentCanvas/nodes/ResourceNode.test.tsx +++ b/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(), diff --git a/packages/apollo-react/src/canvas/components/BaseNode/BaseNode.test.tsx b/packages/apollo-react/src/canvas/components/BaseNode/BaseNode.test.tsx index e14095cd99..2d70c778a3 100644 --- a/packages/apollo-react/src/canvas/components/BaseNode/BaseNode.test.tsx +++ b/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()), 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( + <> + + + + + ); + 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(); + await act(async () => {}); + mockUpdateNodeInternals.mockClear(); + + mockHandleConfigs.current = [ + ...makeHandles(Position.Left, 2), + ...makeHandles(Position.Right, 1), + ]; + rerender(); + 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 diff --git a/packages/apollo-react/src/canvas/components/BaseNode/BaseNode.tsx b/packages/apollo-react/src/canvas/components/BaseNode/BaseNode.tsx index 845df726a3..b4c52f4a6e 100644 --- a/packages/apollo-react/src/canvas/components/BaseNode/BaseNode.tsx +++ b/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>) => { const suggestionType = suggestionTypeOverride === null ? undefined : (suggestionTypeOverride ?? data.suggestionType); - const updateNodeInternals = useUpdateNodeInternals(); + const updateNodeInternals = useBatchedUpdateNodeInternals(); const { updateNodeData, updateNode, getNode } = useReactFlow(); const containerRef = useRef(null); const [isHovered, setIsHovered] = useState(false); diff --git a/packages/apollo-react/src/canvas/components/ButtonHandle/SmartHandle.test.tsx b/packages/apollo-react/src/canvas/components/ButtonHandle/SmartHandle.test.tsx index ac03e4b437..70a98b31e6 100644 --- a/packages/apollo-react/src/canvas/components/ButtonHandle/SmartHandle.test.tsx +++ b/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()), + useUpdateNodeInternals: () => mockUpdateNodeInternals, + useStoreApi: () => mockStoreApi, +})); + +const flushFrame = () => + act(() => new Promise((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( + <> + + + + + ); + 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(); + await flushFrame(); + mockUpdateNodeInternals.mockClear(); + + mockUseStore.mockImplementation(connectedStore(100, 400)); + rerender(); + 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(); + await flushFrame(); + mockUpdateNodeInternals.mockClear(); + + rerender(); + await flushFrame(); + + expect(mockUpdateNodeInternals).not.toHaveBeenCalled(); + }); + }); + describe('Edge cases', () => { it('handles missing nodeId gracefully', () => { mockUseNodeId.mockReturnValue(null); diff --git a/packages/apollo-react/src/canvas/components/ButtonHandle/SmartHandle.tsx b/packages/apollo-react/src/canvas/components/ButtonHandle/SmartHandle.tsx index b55e7fb200..8eaac8ceb7 100644 --- a/packages/apollo-react/src/canvas/components/ButtonHandle/SmartHandle.tsx +++ b/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]); diff --git a/packages/apollo-react/src/canvas/components/ContainerNode/ContainerNode.hooks.ts b/packages/apollo-react/src/canvas/components/ContainerNode/ContainerNode.hooks.ts index b3d07dc8e0..d16cb4207f 100644 --- a/packages/apollo-react/src/canvas/components/ContainerNode/ContainerNode.hooks.ts +++ b/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]); } diff --git a/packages/apollo-react/src/canvas/components/ContainerNode/ContainerNode.test.tsx b/packages/apollo-react/src/canvas/components/ContainerNode/ContainerNode.test.tsx index 3c4449a9e9..aae47a9e6f 100644 --- a/packages/apollo-react/src/canvas/components/ContainerNode/ContainerNode.test.tsx +++ b/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() as ReadonlySet }, + 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 | undefined) { beforeEach(() => { mockUpdateNodeData.mockClear(); + mockUpdateNodeInternals.mockClear(); mockExecutionState.current = undefined; mockReadOnlyNodeIds.current = new Set(); mockValidationState.current = undefined; @@ -218,6 +224,37 @@ describe('ContainerNode root', () => { }); }); +describe('ContainerNode internals', () => { + const flushFrame = () => + act(() => new Promise((resolve) => requestAnimationFrame(() => resolve()))); + + it('updates the internals of containers mounted together in one call', async () => { + render( + <> + + + + + ); + 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(); + await flushFrame(); + mockUpdateNodeInternals.mockClear(); + + rerender(); + await flushFrame(); + + expect(mockUpdateNodeInternals).toHaveBeenCalledTimes(1); + expect(mockUpdateNodeInternals).toHaveBeenCalledWith(['container-1']); + }); +}); + describe('ContainerNode header', () => { it('renders no badge by default', () => { renderContainer(); diff --git a/packages/apollo-react/src/canvas/components/LoopNode/LoopNode.test.tsx b/packages/apollo-react/src/canvas/components/LoopNode/LoopNode.test.tsx index 9a9ad9e0c5..084bf1ee97 100644 --- a/packages/apollo-react/src/canvas/components/LoopNode/LoopNode.test.tsx +++ b/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() }), })); diff --git a/packages/apollo-react/src/canvas/hooks/useBatchedUpdateNodeInternals.test.tsx b/packages/apollo-react/src/canvas/hooks/useBatchedUpdateNodeInternals.test.tsx new file mode 100644 index 0000000000..2a30920bab --- /dev/null +++ b/packages/apollo-react/src/canvas/hooks/useBatchedUpdateNodeInternals.test.tsx @@ -0,0 +1,133 @@ +import { act, render } from '@testing-library/react'; +import { ReactFlowProvider, useStoreApi } from '@uipath/apollo-react/canvas/xyflow/react'; +import type { UpdateNodeInternals } from '@uipath/apollo-react/canvas/xyflow/system'; +import { StrictMode, useEffect } from 'react'; +import { afterEach, describe, expect, it, vi } from 'vitest'; +import { useBatchedUpdateNodeInternals } from './useBatchedUpdateNodeInternals'; + +interface InternalsUpdate { + store: unknown; + ids: string[]; +} + +const { updates } = vi.hoisted(() => ({ updates: [] as InternalsUpdate[] })); + +vi.mock('@uipath/apollo-react/canvas/xyflow/react', async (importOriginal) => { + const actual = await importOriginal(); + const { useCallback } = await import('react'); + return { + ...actual, + useUpdateNodeInternals: () => { + const { getState } = actual.useStoreApi(); + return useCallback( + (ids) => { + updates.push({ store: getState, ids: [ids].flat() }); + }, + [getState] + ); + }, + }; +}); + +interface NodeProbeProps { + id: string; +} + +function NodeProbe({ id }: NodeProbeProps) { + const updateNodeInternals = useBatchedUpdateNodeInternals(); + useEffect(() => { + updateNodeInternals(id); + }, [id, updateNodeInternals]); + return null; +} + +interface StoreProbeProps { + stores: unknown[]; +} + +function StoreProbe({ stores }: StoreProbeProps) { + stores.push(useStoreApi().getState); + return null; +} + +const flushMicrotasks = () => act(async () => {}); + +describe('useBatchedUpdateNodeInternals', () => { + afterEach(() => { + updates.length = 0; + }); + + it('sends the ids of nodes mounted together to the store in one update', async () => { + render( + + + + + + ); + await flushMicrotasks(); + + expect(updates).toHaveLength(1); + expect(updates[0]?.ids).toEqual(['a', 'b', 'c']); + }); + + it('sends one update to each store when two canvases mount together', async () => { + const stores: unknown[] = []; + render( + <> + + + + + + + + + + + + ); + await flushMicrotasks(); + + expect(stores[0]).not.toBe(stores[1]); + expect(updates).toEqual([ + { store: stores[0], ids: ['a', 'b'] }, + { store: stores[1], ids: ['a', 'c'] }, + ]); + }); + + it('sends each id once when StrictMode runs the effects twice', async () => { + render( + + + + + + + ); + await flushMicrotasks(); + + expect(updates).toHaveLength(1); + expect(updates[0]?.ids).toEqual(['a', 'b']); + }); + + it('sends a new update when a node calls again after the first update', async () => { + const { rerender } = render( + + + + + ); + await flushMicrotasks(); + + rerender( + + + + + ); + await flushMicrotasks(); + + expect(updates.map((update) => update.ids)).toEqual([['a', 'b'], ['b']]); + }); +}); diff --git a/packages/apollo-react/src/canvas/hooks/useBatchedUpdateNodeInternals.ts b/packages/apollo-react/src/canvas/hooks/useBatchedUpdateNodeInternals.ts new file mode 100644 index 0000000000..0d3010d346 --- /dev/null +++ b/packages/apollo-react/src/canvas/hooks/useBatchedUpdateNodeInternals.ts @@ -0,0 +1,29 @@ +import { useStoreApi, useUpdateNodeInternals } from '@uipath/apollo-react/canvas/xyflow/react'; +import { useCallback } from 'react'; + +type StoreGetState = ReturnType['getState']; + +const pendingIdsByStore = new WeakMap>(); + +/** `useUpdateNodeInternals`, but calls in the same synchronous run reach each React Flow store as one update. */ +export function useBatchedUpdateNodeInternals(): (id: string) => void { + const { getState } = useStoreApi(); + const updateNodeInternals = useUpdateNodeInternals(); + + return useCallback( + (id: string) => { + const pendingIds = pendingIdsByStore.get(getState); + if (pendingIds) { + pendingIds.add(id); + return; + } + const ids = new Set([id]); + pendingIdsByStore.set(getState, ids); + queueMicrotask(() => { + pendingIdsByStore.delete(getState); + updateNodeInternals([...ids]); + }); + }, + [getState, updateNodeInternals] + ); +}