diff --git a/src/components/shared/ReactFlow/FlowCanvas/FlowCanvas.tsx b/src/components/shared/ReactFlow/FlowCanvas/FlowCanvas.tsx index 34be5f8ab..12d97b8bb 100644 --- a/src/components/shared/ReactFlow/FlowCanvas/FlowCanvas.tsx +++ b/src/components/shared/ReactFlow/FlowCanvas/FlowCanvas.tsx @@ -103,8 +103,13 @@ const FlowCanvas = ({ const { clearContent } = useContextPanel(); const { setReactFlowInstance: setReactFlowInstanceForOverlay } = useNodesOverlay(); - const { componentSpec, setComponentSpec, graphSpec, updateGraphSpec } = - useComponentSpec(); + const { + componentSpec, + setComponentSpec, + graphSpec, + updateGraphSpec, + nodeManager, + } = useComponentSpec(); const { preserveIOSelectionOnSpecChange, resetPrevSpec } = useIOSelectionPersistence(); @@ -292,8 +297,9 @@ const FlowCanvas = ({ connectable: !readOnly && !!nodesConnectable, readOnly: !!readOnly, nodeCallbacks, + nodeManager, }), - [readOnly, nodesConnectable, nodeCallbacks], + [readOnly, nodesConnectable, nodeCallbacks, nodeManager], ); const onConnect = useCallback( @@ -634,6 +640,7 @@ const FlowCanvas = ({ const { updatedComponentSpec, newNodes, updatedNodes } = duplicateNodes( componentSpec, selectedNodes, + nodeManager, { selected: true }, ); @@ -643,7 +650,7 @@ const FlowCanvas = ({ updatedNodes, newNodes, }); - }, [componentSpec, selectedNodes, setComponentSpec, setNodes]); + }, [componentSpec, selectedNodes, nodeManager, setComponentSpec, setNodes]); const onUpgradeNodes = useCallback(async () => { let newGraphSpec = graphSpec; @@ -802,6 +809,7 @@ const FlowCanvas = ({ const { newNodes, updatedComponentSpec } = duplicateNodes( componentSpec, nodesToPaste, + nodeManager, { position: reactFlowCenter, connection: "internal" }, ); @@ -827,6 +835,7 @@ const FlowCanvas = ({ nodes, reactFlowInstance, store, + nodeManager, updateOrAddNodes, setComponentSpec, readOnly, diff --git a/src/components/shared/ReactFlow/FlowCanvas/utils/duplicateNodes.test.ts b/src/components/shared/ReactFlow/FlowCanvas/utils/duplicateNodes.test.ts index d42dddca5..d06a5826c 100644 --- a/src/components/shared/ReactFlow/FlowCanvas/utils/duplicateNodes.test.ts +++ b/src/components/shared/ReactFlow/FlowCanvas/utils/duplicateNodes.test.ts @@ -1,6 +1,7 @@ import type { Node } from "@xyflow/react"; import { describe, expect, it, vi } from "vitest"; +import { NodeManager } from "@/nodeManager"; import type { TaskNodeData } from "@/types/nodes"; import type { ComponentSpec, @@ -17,6 +18,8 @@ import { import { duplicateNodes } from "./duplicateNodes"; +const createMockNodeManager = () => new NodeManager(); + // Mock utility functions const mockTaskSpec: TaskSpec = { componentRef: { name: "test-component" }, @@ -153,8 +156,9 @@ describe("duplicateNodes", () => { }; const nodes: Node[] = []; + const nodeManager = createMockNodeManager(); - expect(() => duplicateNodes(componentSpec, nodes)).toThrow( + expect(() => duplicateNodes(componentSpec, nodes, nodeManager)).toThrow( "ComponentSpec does not contain a graph implementation.", ); }); @@ -178,7 +182,8 @@ describe("duplicateNodes", () => { y: 100, }); - const result = duplicateNodes(componentSpec, [taskNode]); + const nodeManager = createMockNodeManager(); + const result = duplicateNodes(componentSpec, [taskNode], nodeManager); expect(result.newNodes).toHaveLength(1); expect(result.newNodes[0].type).toBe("task"); @@ -208,7 +213,8 @@ describe("duplicateNodes", () => { const inputNode = createMockInputNode("original-input", { x: 50, y: 50 }); - const result = duplicateNodes(componentSpec, [inputNode]); + const nodeManager = createMockNodeManager(); + const result = duplicateNodes(componentSpec, [inputNode], nodeManager); expect(result.newNodes).toHaveLength(1); expect(result.newNodes[0].type).toBe("input"); @@ -243,7 +249,8 @@ describe("duplicateNodes", () => { y: 300, }); - const result = duplicateNodes(componentSpec, [outputNode]); + const nodeManager = createMockNodeManager(); + const result = duplicateNodes(componentSpec, [outputNode], nodeManager); expect(result.newNodes).toHaveLength(1); expect(result.newNodes[0].type).toBe("output"); @@ -274,7 +281,8 @@ describe("duplicateNodes", () => { createMockTaskNode("task2", taskSpec2, { x: 200, y: 200 }), ]; - const result = duplicateNodes(componentSpec, nodes); + const nodeManager = createMockNodeManager(); + const result = duplicateNodes(componentSpec, nodes, nodeManager); expect(result.newNodes).toHaveLength(2); expect(result.newNodes.map((n) => n.id)).toEqual([ @@ -293,7 +301,8 @@ describe("duplicateNodes", () => { const taskNode = createMockTaskNode("original-task", mockTaskSpec); taskNode.selected = true; - const result = duplicateNodes(componentSpec, [taskNode], { + const nodeManager = createMockNodeManager(); + const result = duplicateNodes(componentSpec, [taskNode], nodeManager, { selected: false, }); @@ -317,7 +326,8 @@ describe("duplicateNodes", () => { const taskNode = createMockTaskNode("original-task", taskSpecWithStatus); - const result = duplicateNodes(componentSpec, [taskNode], { + const nodeManager = createMockNodeManager(); + const result = duplicateNodes(componentSpec, [taskNode], nodeManager, { status: false, }); @@ -343,7 +353,8 @@ describe("duplicateNodes", () => { createMockTaskNode("task2", mockTaskSpec, { x: 200, y: 200 }), ]; - const result = duplicateNodes(componentSpec, nodes, { + const nodeManager = createMockNodeManager(); + const result = duplicateNodes(componentSpec, nodes, nodeManager, { position: { x: 500, y: 500 }, }); @@ -399,7 +410,8 @@ describe("duplicateNodes", () => { createMockTaskNode("task2", task2), ]; - const result = duplicateNodes(componentSpec, nodes, { + const nodeManager = createMockNodeManager(); + const result = duplicateNodes(componentSpec, nodes, nodeManager, { connection: "none", }); @@ -422,7 +434,8 @@ describe("duplicateNodes", () => { createMockTaskNode("task2", task2), ]; - const result = duplicateNodes(componentSpec, nodes, { + const nodeManager = createMockNodeManager(); + const result = duplicateNodes(componentSpec, nodes, nodeManager, { connection: "internal", }); @@ -474,7 +487,8 @@ describe("duplicateNodes", () => { createMockTaskNode("task2", task2WithConnections), ]; - const result = duplicateNodes(componentSpec, nodes, { + const nodeManager = createMockNodeManager(); + const result = duplicateNodes(componentSpec, nodes, nodeManager, { connection: "external", }); @@ -507,7 +521,8 @@ describe("duplicateNodes", () => { createMockTaskNode("task2", task2), ]; - const result = duplicateNodes(componentSpec, nodes, { + const nodeManager = createMockNodeManager(); + const result = duplicateNodes(componentSpec, nodes, nodeManager, { connection: "all", }); @@ -549,7 +564,8 @@ describe("duplicateNodes", () => { createMockTaskNode("task1", taskSpec), ]; - const result = duplicateNodes(componentSpec, nodes, { + const nodeManager = createMockNodeManager(); + const result = duplicateNodes(componentSpec, nodes, nodeManager, { connection: "all", }); @@ -586,7 +602,8 @@ describe("duplicateNodes", () => { createMockOutputNode("graph-output"), ]; - const result = duplicateNodes(componentSpec, nodes, { + const nodeManager = createMockNodeManager(); + const result = duplicateNodes(componentSpec, nodes, nodeManager, { connection: "all", }); @@ -609,7 +626,8 @@ describe("duplicateNodes", () => { describe("edge cases", () => { it("should handle empty node array", () => { const componentSpec = createMockComponentSpec(); - const result = duplicateNodes(componentSpec, []); + const nodeManager = createMockNodeManager(); + const result = duplicateNodes(componentSpec, [], nodeManager); expect(result.newNodes).toHaveLength(0); expect(result.nodeIdMap).toEqual({}); @@ -623,7 +641,8 @@ describe("duplicateNodes", () => { const taskNode = createMockTaskNode("original-task", mockTaskSpec); taskNode.measured = { width: 300, height: 200 }; - const result = duplicateNodes(componentSpec, [taskNode]); + const nodeManager = createMockNodeManager(); + const result = duplicateNodes(componentSpec, [taskNode], nodeManager); expect(result.newNodes[0].measured).toEqual({ width: 300, height: 200 }); }); @@ -643,7 +662,8 @@ describe("duplicateNodes", () => { taskSpecWithoutPosition, ); - const result = duplicateNodes(componentSpec, [taskNode]); + const nodeManager = createMockNodeManager(); + const result = duplicateNodes(componentSpec, [taskNode], nodeManager); expect(result.newNodes).toHaveLength(1); expect(result.newNodes[0].position).toEqual({ x: 110, y: 110 }); @@ -658,7 +678,8 @@ describe("duplicateNodes", () => { const taskNode = createMockTaskNode("original-task", mockTaskSpec); - const result = duplicateNodes(componentSpec, [taskNode]); + const nodeManager = createMockNodeManager(); + const result = duplicateNodes(componentSpec, [taskNode], nodeManager); expect(result).toHaveProperty("updatedComponentSpec"); expect(result).toHaveProperty("nodeIdMap"); diff --git a/src/components/shared/ReactFlow/FlowCanvas/utils/duplicateNodes.ts b/src/components/shared/ReactFlow/FlowCanvas/utils/duplicateNodes.ts index fa7e953c0..fcf5796e0 100644 --- a/src/components/shared/ReactFlow/FlowCanvas/utils/duplicateNodes.ts +++ b/src/components/shared/ReactFlow/FlowCanvas/utils/duplicateNodes.ts @@ -1,5 +1,6 @@ import { type Node, type XYPosition } from "@xyflow/react"; +import type { NodeManager } from "@/nodeManager"; import type { IONodeData, NodeData, TaskNodeData } from "@/types/nodes"; import { type ComponentSpec, @@ -44,6 +45,7 @@ type ConnectionMode = "none" | "internal" | "external" | "all"; export const duplicateNodes = ( componentSpec: ComponentSpec, nodesToDuplicate: Node[], + nodeManager: NodeManager, config?: { selected?: boolean; position?: XYPosition; @@ -288,6 +290,7 @@ export const duplicateNodes = ( readOnly: taskData.readOnly, connectable: taskData.connectable, callbacks: convertTaskCallbacksToNodeCallbacks(taskData.callbacks), + nodeManager, }; const newNode = createTaskNode([newTaskId, newTaskSpec], nodeData); @@ -319,6 +322,7 @@ export const duplicateNodes = ( const inputData = originalNode.data as IONodeData; const nodeData: NodeData = { readOnly: inputData.readOnly, + nodeManager, }; const newNode = createInputNode(newInputSpec, nodeData); @@ -350,6 +354,7 @@ export const duplicateNodes = ( const outputData = originalNode.data as IONodeData; const nodeData: NodeData = { readOnly: outputData.readOnly, + nodeManager, }; const newNode = createOutputNode(newOutputSpec, nodeData); diff --git a/src/hooks/useNodeCallbacks.ts b/src/hooks/useNodeCallbacks.ts index 1213988a0..2c8eb0a24 100644 --- a/src/hooks/useNodeCallbacks.ts +++ b/src/hooks/useNodeCallbacks.ts @@ -34,8 +34,13 @@ export const useNodeCallbacks = ({ const notify = useToastNotification(); const reactFlowInstance = useReactFlow(); - const { graphSpec, updateGraphSpec, componentSpec, setComponentSpec } = - useComponentSpec(); + const { + graphSpec, + updateGraphSpec, + componentSpec, + setComponentSpec, + nodeManager, + } = useComponentSpec(); // Workaround for nodes state being stale in task node callbacks const getNodeById = useCallback( @@ -139,6 +144,7 @@ export const useNodeCallbacks = ({ const { updatedComponentSpec, newNodes, updatedNodes } = duplicateNodes( componentSpec, [node], + nodeManager, { selected }, ); @@ -149,7 +155,13 @@ export const useNodeCallbacks = ({ newNodes, }); }, - [componentSpec, getNodeById, setComponentSpec, updateOrAddNodes], + [ + componentSpec, + nodeManager, + getNodeById, + setComponentSpec, + updateOrAddNodes, + ], ); const onUpgrade = useCallback( diff --git a/src/types/nodes.ts b/src/types/nodes.ts index 87cbfcf86..1ee954a52 100644 --- a/src/types/nodes.ts +++ b/src/types/nodes.ts @@ -1,3 +1,4 @@ +import type { NodeManager } from "@/nodeManager"; import type { ArgumentType, ComponentReference, @@ -12,6 +13,7 @@ export interface NodeData extends Record { readOnly: boolean; connectable?: boolean; callbacks?: NodeCallbacks; + nodeManager: NodeManager; } export interface TaskNodeData extends Record { diff --git a/src/utils/nodes/createInputNode.ts b/src/utils/nodes/createInputNode.ts index 6193a664c..639e7dce5 100644 --- a/src/utils/nodes/createInputNode.ts +++ b/src/utils/nodes/createInputNode.ts @@ -8,7 +8,10 @@ import { inputNameToNodeId } from "./nodeIdUtils"; export const createInputNode = (input: InputSpec, nodeData: NodeData) => { const { name, annotations } = input; - const { readOnly } = nodeData; + const { nodeManager, readOnly } = nodeData; + + const newNodeId = nodeManager?.getNodeId(name, "input"); + console.log("Creating input node:", { name, nodeId: newNodeId }); const position = extractPositionFromAnnotations(annotations); const nodeId = inputNameToNodeId(name); diff --git a/src/utils/nodes/createOutputNode.ts b/src/utils/nodes/createOutputNode.ts index 6fee2ee03..627deb924 100644 --- a/src/utils/nodes/createOutputNode.ts +++ b/src/utils/nodes/createOutputNode.ts @@ -8,7 +8,10 @@ import { outputNameToNodeId } from "./nodeIdUtils"; export const createOutputNode = (output: OutputSpec, nodeData: NodeData) => { const { name, annotations } = output; - const { readOnly } = nodeData; + const { nodeManager, readOnly } = nodeData; + + const newNodeId = nodeManager?.getNodeId(name, "output"); + console.log("Creating output node:", { name, nodeId: newNodeId }); const position = extractPositionFromAnnotations(annotations); const nodeId = outputNameToNodeId(name); diff --git a/src/utils/nodes/createTaskNode.ts b/src/utils/nodes/createTaskNode.ts index e58c1d94f..13336279e 100644 --- a/src/utils/nodes/createTaskNode.ts +++ b/src/utils/nodes/createTaskNode.ts @@ -12,7 +12,10 @@ export const createTaskNode = ( nodeData: NodeData, ) => { const [taskId, taskSpec] = task; - const { callbacks, connectable, ...data } = nodeData; + const { nodeManager, callbacks, connectable, ...data } = nodeData; + + const newNodeId = nodeManager?.getNodeId(taskId, "task"); + console.log("Creating task node:", { taskId, nodeId: newNodeId }); const position = extractPositionFromAnnotations(taskSpec.annotations); const nodeId = taskIdToNodeId(taskId);