diff --git a/react-compiler.config.js b/react-compiler.config.js index 251878dc7..dcd76e182 100644 --- a/react-compiler.config.js +++ b/react-compiler.config.js @@ -42,6 +42,7 @@ export const REACT_COMPILER_ENABLED_DIRS = [ "src/components/shared/ManageComponent/DeprecatePublishedComponentButton.tsx", "src/components/shared/ManageComponent/PublishComponent.tsx", "src/components/shared/ManageComponent/hooks/useComponentCanvasTasks.ts", + "src/components/shared/TaskDetails/Actions/UnpackSubgraphButton.tsx", // 11-20 useCallback/useMemo // "src/components/ui", // 12 diff --git a/src/components/shared/ReactFlow/FlowCanvas/utils/addTask.ts b/src/components/shared/ReactFlow/FlowCanvas/utils/addTask.ts index 0728d9ee0..f5af8459e 100644 --- a/src/components/shared/ReactFlow/FlowCanvas/utils/addTask.ts +++ b/src/components/shared/ReactFlow/FlowCanvas/utils/addTask.ts @@ -26,7 +26,9 @@ interface AddTaskResult { * Options for creating input/output nodes. * Omits position-related fields (annotations) which are automatically set. */ -type IONodeOptions = Omit, "annotations">; +type IONodeOptions = + | Omit, "annotations"> + | Omit, "annotations">; /** * Creates a task, input, or output node and adds it to the component specification. diff --git a/src/components/shared/ReactFlow/FlowCanvas/utils/removeNode.ts b/src/components/shared/ReactFlow/FlowCanvas/utils/removeNode.ts index a6934cbc1..0a7b8d790 100644 --- a/src/components/shared/ReactFlow/FlowCanvas/utils/removeNode.ts +++ b/src/components/shared/ReactFlow/FlowCanvas/utils/removeNode.ts @@ -82,7 +82,10 @@ export const removeGraphOutput = ( return componentSpec; }; -const removeTask = (taskIdToRemove: string, componentSpec: ComponentSpec) => { +export const removeTask = ( + taskIdToRemove: string, + componentSpec: ComponentSpec, +) => { if (isGraphImplementation(componentSpec.implementation)) { const graphSpec = componentSpec.implementation.graph; diff --git a/src/components/shared/TaskDetails/Actions.tsx b/src/components/shared/TaskDetails/Actions.tsx index 21a952add..4d70d1f47 100644 --- a/src/components/shared/TaskDetails/Actions.tsx +++ b/src/components/shared/TaskDetails/Actions.tsx @@ -12,6 +12,7 @@ import { DownloadYamlButton } from "./Actions/DownloadYamlButton"; import { DuplicateTaskButton } from "./Actions/DuplicateTaskButton"; import { EditComponentButton } from "./Actions/EditComponentButton"; import { NavigateToSubgraphButton } from "./Actions/NavigateToSubgraphButton"; +import { UnpackSubgraphButton } from "./Actions/UnpackSubgraphButton"; import { UpgradeTaskButton } from "./Actions/UpgradeTaskButton"; interface TaskActionsProps { @@ -59,6 +60,9 @@ const TaskActions = ({ const navigateToSubgraph = isSubgraphNode && taskId && !readOnly && ( ); + const unpackSubgraphButton = isSubgraphNode && taskId && !readOnly && ( + + ); const deleteComponent = onDelete && !readOnly && ( ); @@ -72,6 +76,7 @@ const TaskActions = ({ duplicateTask, upgradeTask, navigateToSubgraph, + unpackSubgraphButton, deleteComponent, ].filter(Boolean); diff --git a/src/components/shared/TaskDetails/Actions/UnpackSubgraphButton.tsx b/src/components/shared/TaskDetails/Actions/UnpackSubgraphButton.tsx new file mode 100644 index 000000000..38db3bcdc --- /dev/null +++ b/src/components/shared/TaskDetails/Actions/UnpackSubgraphButton.tsx @@ -0,0 +1,55 @@ +import { useState } from "react"; + +import { useComponentSpec } from "@/providers/ComponentSpecProvider"; +import { unpackSubgraph } from "@/utils/nodes/unpacking/unpackSubgraph"; +import { + getSubgraphDescription, + updateSubgraphSpec, +} from "@/utils/subgraphUtils"; + +import { ActionButton } from "../../Buttons/ActionButton"; +import { removeTask } from "../../ReactFlow/FlowCanvas/utils/removeNode"; + +interface UnpackSubgraphButtonProps { + taskId: string; +} + +export const UnpackSubgraphButton = ({ taskId }: UnpackSubgraphButtonProps) => { + const { + currentGraphSpec, + currentSubgraphSpec, + currentSubgraphPath, + componentSpec, + setComponentSpec, + } = useComponentSpec(); + + const [isLoading, setIsLoading] = useState(false); + + const taskSpec = currentGraphSpec.tasks[taskId]; + const subgraphDescription = taskSpec ? getSubgraphDescription(taskSpec) : ""; + + function handleUnpackSubgraph() { + setIsLoading(true); + const updatedSubgraphSpec = unpackSubgraph(taskId, currentSubgraphSpec); + + const cleanedSubgraphSpec = removeTask(taskId, updatedSubgraphSpec); + + const newRootSpec = updateSubgraphSpec( + componentSpec, + currentSubgraphPath, + cleanedSubgraphSpec, + ); + + setComponentSpec(newRootSpec); + setIsLoading(false); + } + + return ( + + ); +}; diff --git a/src/utils/graphUtils.ts b/src/utils/graphUtils.ts new file mode 100644 index 000000000..ef0210788 --- /dev/null +++ b/src/utils/graphUtils.ts @@ -0,0 +1,65 @@ +import type { XYPosition } from "@xyflow/react"; + +import { + type ComponentSpec, + type GraphSpec, + isGraphImplementation, +} from "./componentSpec"; +import { extractPositionFromAnnotations } from "./nodes/extractPositionFromAnnotations"; + +export const calculateSpecCenter = ( + componentSpec: ComponentSpec, +): XYPosition => { + if (!isGraphImplementation(componentSpec.implementation)) { + return { x: 0, y: 0 }; + } + + const graphSpec: GraphSpec = componentSpec.implementation.graph; + + const allPositions: XYPosition[] = []; + + // Collect positions from tasks + Object.values(graphSpec.tasks).forEach((task) => { + const taskPosition = extractPositionFromAnnotations(task.annotations); + if (taskPosition) { + allPositions.push(taskPosition); + } + }); + + // Collect positions from inputs + componentSpec.inputs?.forEach((input) => { + const inputPosition = extractPositionFromAnnotations(input.annotations); + if (inputPosition) { + allPositions.push(inputPosition); + } + }); + + // Collect positions from outputs + componentSpec.outputs?.forEach((output) => { + const outputPosition = extractPositionFromAnnotations(output.annotations); + if (outputPosition) { + allPositions.push(outputPosition); + } + }); + + if (allPositions.length === 0) { + return { x: 0, y: 0 }; + } + + const sumX = allPositions.reduce((sum, pos) => sum + pos.x, 0); + const sumY = allPositions.reduce((sum, pos) => sum + pos.y, 0); + + return { + x: sumX / allPositions.length, + y: sumY / allPositions.length, + }; +}; + +export const normalizeNodePositionInGroup = ( + nodePosition: XYPosition | undefined, + groupPosition: XYPosition | undefined, + groupCenter: XYPosition, +): XYPosition => ({ + x: (groupPosition?.x || 0) + (nodePosition?.x || 0) - groupCenter.x, + y: (groupPosition?.y || 0) + (nodePosition?.y || 0) - groupCenter.y, +}); diff --git a/src/utils/nodes/unpacking/helpers.ts b/src/utils/nodes/unpacking/helpers.ts new file mode 100644 index 000000000..9ae0692cd --- /dev/null +++ b/src/utils/nodes/unpacking/helpers.ts @@ -0,0 +1,290 @@ +import type { XYPosition } from "@xyflow/react"; + +import addTask from "@/components/shared/ReactFlow/FlowCanvas/utils/addTask"; +import { setGraphOutputValue } from "@/components/shared/ReactFlow/FlowCanvas/utils/setGraphOutputValue"; +import { + type ArgumentType, + type ComponentSpec, + isGraphImplementation, + isGraphInputArgument, + isTaskOutputArgument, + type MetadataSpec, + type TaskOutputArgument, +} from "@/utils/componentSpec"; +import { + calculateSpecCenter, + normalizeNodePositionInGroup, +} from "@/utils/graphUtils"; + +import { extractPositionFromAnnotations } from "../extractPositionFromAnnotations"; + +export const unpackInputs = ( + containerSpec: ComponentSpec, + containerPosition: XYPosition, + componentSpec: ComponentSpec, +): { + spec: ComponentSpec; + inputNameMap: Map; +} => { + let updatedSpec = componentSpec; + const inputNameMap = new Map(); + + const containerCenter = calculateSpecCenter(containerSpec); + + const inputs = containerSpec.inputs; + + inputs?.forEach((input) => { + const position = calculateUnpackedPosition( + input.annotations, + containerPosition, + containerCenter, + ); + + const { spec, ioName } = addTask( + "input", + null, + position, + updatedSpec, + input, + ); + + if (ioName && ioName !== input.name) { + inputNameMap.set(input.name, ioName); + } + + updatedSpec = spec; + }); + + return { spec: updatedSpec, inputNameMap }; +}; + +export const unpackOutputs = ( + containerSpec: ComponentSpec, + containerPosition: XYPosition, + componentSpec: ComponentSpec, +): { + spec: ComponentSpec; + outputNameMap: Map; +} => { + let updatedSpec = componentSpec; + const outputNameMap = new Map(); + + const containerCenter = calculateSpecCenter(containerSpec); + + const outputs = containerSpec.outputs; + + outputs?.forEach((output) => { + const position = calculateUnpackedPosition( + output.annotations, + containerPosition, + containerCenter, + ); + + const { spec, ioName } = addTask( + "output", + null, + position, + updatedSpec, + output, + ); + + if (ioName && ioName !== output.name) { + outputNameMap.set(output.name, ioName); + } + + updatedSpec = spec; + }); + + return { spec: updatedSpec, outputNameMap }; +}; + +export const unpackTasks = ( + containerSpec: ComponentSpec, + containerPosition: XYPosition, + componentSpec: ComponentSpec, + inputNameMap: Map, +): { + spec: ComponentSpec; + taskIdMap: Map; +} => { + let updatedSpec = componentSpec; + const taskIdMap = new Map(); + + if (!isGraphImplementation(containerSpec.implementation)) { + return { spec: updatedSpec, taskIdMap }; + } + + const containerCenter = calculateSpecCenter(containerSpec); + + const tasks = containerSpec.implementation.graph.tasks; + + Object.entries(tasks).forEach(([taskId, task]) => { + const position = calculateUnpackedPosition( + task.annotations, + containerPosition, + containerCenter, + ); + + const { spec, taskId: newTaskId } = addTask( + "task", + task, + position, + updatedSpec, + ); + + if (newTaskId && newTaskId !== taskId) { + taskIdMap.set(taskId, newTaskId); + } + + updatedSpec = spec; + }); + + if (!isGraphImplementation(updatedSpec.implementation)) { + return { spec: updatedSpec, taskIdMap }; + } + + const updatedTasks = { ...updatedSpec.implementation.graph.tasks }; + + Object.entries(tasks).forEach(([oldTaskId, task]) => { + const newTaskId = taskIdMap.get(oldTaskId) || oldTaskId; + const currentTask = updatedTasks[newTaskId]; + + if (!currentTask) { + return; + } + + const remappedArguments = remapTaskArguments( + task.arguments, + taskIdMap, + inputNameMap, + ); + + if (remappedArguments) { + updatedTasks[newTaskId] = { + ...currentTask, + arguments: remappedArguments, + }; + } + }); + + updatedSpec = { + ...updatedSpec, + implementation: { + ...updatedSpec.implementation, + graph: { + ...updatedSpec.implementation.graph, + tasks: updatedTasks, + }, + }, + }; + + return { spec: updatedSpec, taskIdMap }; +}; + +export const copyOutputValues = ( + containerSpec: ComponentSpec, + componentSpec: ComponentSpec, + outputNameMap: Map, + taskIdMap: Map, +): ComponentSpec => { + let updatedSpec = componentSpec; + + if (!isGraphImplementation(containerSpec.implementation)) { + return updatedSpec; + } + + const outputValues = containerSpec.implementation.graph.outputValues || {}; + + Object.entries(outputValues).forEach(([outputName, outputValue]) => { + if (isGraphImplementation(updatedSpec.implementation)) { + const newOutputName = outputNameMap.get(outputName) || outputName; + const remappedTaskOutputArg = remapTaskOutputArgument( + outputValue, + taskIdMap, + ); + + const updatedGraphSpec = setGraphOutputValue( + updatedSpec.implementation.graph, + newOutputName, + remappedTaskOutputArg, + ); + + updatedSpec = { + ...updatedSpec, + implementation: { + ...updatedSpec.implementation, + graph: updatedGraphSpec, + }, + }; + } + }); + + return updatedSpec; +}; + +const remapTaskOutputArgument = ( + taskOutput: TaskOutputArgument, + taskIdMap: Map, +): TaskOutputArgument => { + const newTaskId = + taskIdMap.get(taskOutput.taskOutput.taskId) || taskOutput.taskOutput.taskId; + + return { + taskOutput: { + ...taskOutput.taskOutput, + taskId: newTaskId, + }, + }; +}; + +const remapArgumentValue = ( + arg: ArgumentType, + taskIdMap: Map, + inputNameMap: Map, +) => { + if (isTaskOutputArgument(arg)) { + return remapTaskOutputArgument(arg, taskIdMap); + } + + if (isGraphInputArgument(arg)) { + const newInputName = + inputNameMap.get(arg.graphInput.inputName) || arg.graphInput.inputName; + return { + graphInput: { + ...arg.graphInput, + inputName: newInputName, + }, + }; + } + + return arg; +}; + +const remapTaskArguments = ( + args: Record | undefined, + taskIdMap: Map, + inputNameMap: Map, +) => { + if (!args) return undefined; + + const remappedArgs: Record = {}; + + for (const [key, value] of Object.entries(args)) { + remappedArgs[key] = remapArgumentValue(value, taskIdMap, inputNameMap); + } + + return remappedArgs; +}; + +const calculateUnpackedPosition = ( + nodeAnnotations: MetadataSpec["annotations"], + containerPosition: XYPosition, + containerCenter: XYPosition, +): XYPosition => { + const nodePosition = extractPositionFromAnnotations(nodeAnnotations); + return normalizeNodePositionInGroup( + nodePosition, + containerPosition, + containerCenter, + ); +}; diff --git a/src/utils/nodes/unpacking/unpackSubgraph.ts b/src/utils/nodes/unpacking/unpackSubgraph.ts new file mode 100644 index 000000000..5efb1f106 --- /dev/null +++ b/src/utils/nodes/unpacking/unpackSubgraph.ts @@ -0,0 +1,66 @@ +import type { ComponentSpec } from "@/utils/componentSpec"; +import { isGraphImplementation } from "@/utils/componentSpec"; +import { extractPositionFromAnnotations } from "@/utils/nodes/extractPositionFromAnnotations"; + +import { + copyOutputValues, + unpackInputs, + unpackOutputs, + unpackTasks, +} from "./helpers"; + +export const unpackSubgraph = ( + subgraphTaskId: string, + componentSpec: ComponentSpec, +): ComponentSpec => { + if (!isGraphImplementation(componentSpec.implementation)) { + return componentSpec; + } + + const graphSpec = componentSpec.implementation.graph; + const taskSpec = graphSpec.tasks[subgraphTaskId]; + const subgraphSpec = taskSpec.componentRef.spec; + + if (!subgraphSpec) { + return componentSpec; + } + + const subgraphPosition = extractPositionFromAnnotations(taskSpec.annotations); + + let updatedComponentSpec = componentSpec; + + // Unpack inputs + const { spec: specAfterInputs, inputNameMap } = unpackInputs( + subgraphSpec, + subgraphPosition, + updatedComponentSpec, + ); + updatedComponentSpec = specAfterInputs; + + // Unpack outputs + const { spec: specAfterOutputs, outputNameMap } = unpackOutputs( + subgraphSpec, + subgraphPosition, + updatedComponentSpec, + ); + updatedComponentSpec = specAfterOutputs; + + // Unpack tasks + const { spec: specAfterTasks, taskIdMap } = unpackTasks( + subgraphSpec, + subgraphPosition, + updatedComponentSpec, + inputNameMap, + ); + updatedComponentSpec = specAfterTasks; + + // Copy output values + updatedComponentSpec = copyOutputValues( + subgraphSpec, + updatedComponentSpec, + outputNameMap, + taskIdMap, + ); + + return updatedComponentSpec; +};