feat: (Un)Pack Flex Nodes into Subgraphs

src/components/shared/ReactFlow/FlowCanvas/Subgraphs/create/NodeListItem.tsx

@@ -7,6 +7,7 @@ import { InlineStack } from "@/components/ui/layout";
 import { Paragraph } from "@/components/ui/typography";
 import { cn } from "@/lib/utils";
 
+import { isFlexNode } from "../../types";
 import { getNodeTypeColor } from "./utils";
 
 interface NodeListItemProps {
@@ -34,16 +35,31 @@ export function NodeListItem({
     }
   };
 
+  const isFlexType = isFlexNode(node);
+
+  const displayValue = isFlexType
+    ? `Sticky Note: ${node.data.properties.title.length > 0 ? node.data.properties.title : node.id}`
+    : node.id;
+
   return (
     <li key={node.id} className="flex justify-between items-center">
       <InlineStack
         gap="3"
         className={cn({ "opacity-50": isExcluded })}
         wrap="nowrap"
       >
-        <Badge variant="dot" className={cn(getNodeTypeColor(node.type))} />
+        <Badge
+          variant="dot"
+          className={cn(getNodeTypeColor(node.type))}
+          style={{
+            backgroundColor: isFlexType
+              ? node.data.properties.color
+              : undefined,
+            filter: isFlexType ? "brightness(0.8)" : undefined,
+          }}
+        />
         <Paragraph font="mono" size="xs">
-          {node.id}
+          {displayValue}
         </Paragraph>
         {isOrphaned && (
           <InlineStack gap="1" wrap="nowrap">

src/components/shared/ReactFlow/FlowCanvas/Subgraphs/create/checkForOrphanedNodes.ts

@@ -98,7 +98,7 @@ export const checkForOrphanedNodes = (
     });
 
   const orphanedNodes = selectedNodes.filter(
-    (node) => !connectedNodeIds.has(node.id),
+    (node) => !connectedNodeIds.has(node.id) && node.type !== "flex",
   );
 
   return orphanedNodes;

src/components/shared/ReactFlow/FlowCanvas/Subgraphs/create/utils.ts

@@ -27,6 +27,8 @@ export const canGroupNodes = (nodes: Node[]): GroupingValidation => {
 
 export const getNodeTypeColor = (nodeType: string | undefined): string => {
   switch (nodeType) {
+    case "flex":
+      return "bg-yellow-300";
     case "input":
       return "bg-blue-500";
     case "output":

src/utils/nodes/createSubgraphFromNodes.ts

@@ -1,5 +1,10 @@
 import type { Node, XYPosition } from "@xyflow/react";
 
+import {
+  getFlexNodeAnnotations,
+  serializeFlexNodes,
+} from "@/components/shared/ReactFlow/FlowCanvas/FlexNode/interface";
+import type { FlexNodeData } from "@/components/shared/ReactFlow/FlowCanvas/FlexNode/types";
 import {
   type Bounds,
   calculateNodesCenter,
@@ -23,7 +28,10 @@ import {
   isTaskOutputArgument,
 } from "@/utils/componentSpec";
 
-import { EDITOR_POSITION_ANNOTATION } from "../annotations";
+import {
+  EDITOR_POSITION_ANNOTATION,
+  FLEX_NODES_ANNOTATION,
+} from "../annotations";
 import { extractPositionFromAnnotations } from "../annotations";
 import { generateDigest } from "../componentStore";
 import { getUniqueName, getUniqueTaskName } from "../unique";
@@ -67,12 +75,14 @@ export const createSubgraphFromNodes = async (
   const taskNodes = selectedNodes.filter((node) => node.type === "task");
   const inputNodes = selectedNodes.filter((node) => node.type === "input");
   const outputNodes = selectedNodes.filter((node) => node.type === "output");
+  const flexNodes = selectedNodes.filter((node) => node.type === "flex");
 
   const subgraphTasks: Record<string, TaskSpec> = {};
   const subgraphInputs: InputSpec[] = [];
   const subgraphArguments: Record<string, ArgumentType> = {};
   const subgraphOutputs: OutputSpec[] = [];
   const subgraphOutputValues: Record<string, TaskOutputArgument> = {};
+  const subgraphAnnotations: Record<string, string> = {};
 
   const bounds = getNodesBounds(selectedNodes);
 
@@ -145,6 +155,13 @@ export const createSubgraphFromNodes = async (
     currentGraphSpec,
   );
 
+  processSelectedFlexNodes(
+    flexNodes,
+    bounds,
+    subgraphAnnotations,
+    currentComponentSpec,
+  );
+
   // Create the replacement task that represents the subgraph
   const subgraphPosition = calculateNodesCenter(selectedNodes);
 
@@ -159,6 +176,7 @@ export const createSubgraphFromNodes = async (
     subgraphOutputValues,
     subgraphPosition,
     subgraphArguments,
+    subgraphAnnotations,
   );
 
   const text = await getComponentText(subgraphTask.componentRef);
@@ -179,6 +197,7 @@ const createSubgraphTask = async (
   outputValues: Record<string, TaskOutputArgument>,
   position: XYPosition,
   args: Record<string, ArgumentType>,
+  annotations: Record<string, string> = {},
 ) => {
   let author: string = "Unknown";
   try {
@@ -205,6 +224,7 @@ const createSubgraphTask = async (
         sdk: "https://cloud-pipelines.net/pipeline-editor/",
         "editor.flow-direction": "left-to-right",
         author,
+        ...annotations,
       },
     },
   };
@@ -309,6 +329,38 @@ const processSelectedOutputNodes = (
   });
 };
 
+const processSelectedFlexNodes = (
+  flexNodes: Node[],
+  bounds: Bounds,
+  annotations: Record<string, string>,
+  currentComponentSpec: ComponentSpec,
+): void => {
+  const existingFlexNodes = getFlexNodeAnnotations(currentComponentSpec);
+  const newFlexNodes: FlexNodeData[] = [];
+
+  flexNodes.forEach((node) => {
+    const flexNodeId = node.id;
+
+    const existingFlexNode = existingFlexNodes.find(
+      (flex) => flex.id === flexNodeId,
+    );
+
+    if (existingFlexNode) {
+      const newFlexNode = { ...existingFlexNode };
+      const normalizedPosition = normalizeNodePosition(node, bounds);
+      newFlexNode.position = normalizedPosition;
+
+      newFlexNodes.push(newFlexNode);
+    }
+  });
+
+  if (newFlexNodes.length === 0) {
+    return;
+  }
+
+  annotations[FLEX_NODES_ANNOTATION] = serializeFlexNodes(newFlexNodes);
+};
+
 const processTaskInputConnections = (
   taskSpec: TaskSpec,
   taskPosition: XYPosition,

src/utils/nodes/unpacking/helpers.ts

@@ -1,9 +1,16 @@
 import type { XYPosition } from "@xyflow/react";
 
+import {
+  getFlexNodeAnnotations,
+  serializeFlexNodes,
+} from "@/components/shared/ReactFlow/FlowCanvas/FlexNode/interface";
 import addTask from "@/components/shared/ReactFlow/FlowCanvas/utils/addTask";
 import { setGraphOutputValue } from "@/components/shared/ReactFlow/FlowCanvas/utils/setGraphOutputValue";
 import { setTaskArgument } from "@/components/shared/ReactFlow/FlowCanvas/utils/setTaskArgument";
-import { extractPositionFromAnnotations } from "@/utils/annotations";
+import {
+  extractPositionFromAnnotations,
+  FLEX_NODES_ANNOTATION,
+} from "@/utils/annotations";
 import {
   type ArgumentType,
   type ComponentSpec,
@@ -23,6 +30,51 @@ import {
   normalizeNodePositionInGroup,
 } from "@/utils/graphUtils";
 
+export const unpackFlexNodes = (
+  containerSpec: ComponentSpec,
+  containerPosition: XYPosition,
+  componentSpec: ComponentSpec,
+): ComponentSpec => {
+  const updatedSpec = componentSpec;
+
+  const flexNodes = getFlexNodeAnnotations(containerSpec);
+
+  const containerCenter = calculateSpecCenter(containerSpec);
+
+  const newFlexNodes = flexNodes.map((flexNode) => {
+    const position = normalizeNodePositionInGroup(
+      {
+        x: flexNode.position.x,
+        y: flexNode.position.y,
+      },
+      containerPosition,
+      containerCenter,
+    );
+
+    return {
+      ...flexNode,
+      position,
+    };
+  });
+
+  if (!updatedSpec.metadata) {
+    updatedSpec.metadata = {};
+  }
+
+  if (!updatedSpec.metadata.annotations) {
+    updatedSpec.metadata.annotations = {};
+  }
+
+  const existingFlexNodes = getFlexNodeAnnotations(updatedSpec);
+
+  const mergedFlexNodes = [...existingFlexNodes, ...newFlexNodes];
+
+  updatedSpec.metadata.annotations[FLEX_NODES_ANNOTATION] =
+    serializeFlexNodes(mergedFlexNodes);
+
+  return updatedSpec;
+};
+
 export const unpackInputs = (
   containerSpec: ComponentSpec,
   containerPosition: XYPosition,

src/utils/nodes/unpacking/unpackSubgraph.test.ts

@@ -16,6 +16,7 @@ import {
   reconnectDownstreamOutputs,
   reconnectDownstreamTasks,
   reconnectUpstreamInputsAndTasks,
+  unpackFlexNodes,
   unpackInputs,
   unpackOutputs,
   unpackTasks,
@@ -77,6 +78,7 @@ describe("unpackSubgraph", () => {
     vi.mocked(extractPositionFromAnnotations).mockReturnValue(mockPosition);
     vi.mocked(getOutputNodesConnectedToTask).mockReturnValue({});
     vi.mocked(getDownstreamTaskNodesConnectedToTask).mockReturnValue({});
+    vi.mocked(unpackFlexNodes).mockImplementation((_, __, spec) => spec);
   });
 
   it("should return unchanged spec if implementation is not a graph", () => {

src/utils/nodes/unpacking/unpackSubgraph.ts

@@ -11,6 +11,7 @@ import {
   reconnectDownstreamOutputs,
   reconnectDownstreamTasks,
   reconnectUpstreamInputsAndTasks,
+  unpackFlexNodes,
   unpackInputs,
   unpackOutputs,
   unpackTasks,
@@ -46,6 +47,14 @@ export const unpackSubgraph = (
 
   let updatedComponentSpec = componentSpec;
 
+  // Unpack flex Nodes
+  const specAfterFlexNodes = unpackFlexNodes(
+    subgraphSpec,
+    subgraphPosition,
+    updatedComponentSpec,
+  );
+  updatedComponentSpec = specAfterFlexNodes;
+
   // Unpack inputs
   const { spec: specAfterInputs, inputNameMap } = unpackInputs(
     subgraphSpec,