Make loop node parameter a select (#906)

skyvern-frontend/src/routes/workflows/editor/nodes/LoopNode/LoopNode.tsx

@@ -1,5 +1,11 @@
-import { Input } from "@/components/ui/input";
 import { Label } from "@/components/ui/label";
+import {
+  Select,
+  SelectContent,
+  SelectItem,
+  SelectTrigger,
+  SelectValue,
+} from "@/components/ui/select";
 import { useDeleteNodeCallback } from "@/routes/workflows/hooks/useDeleteNodeCallback";
 import { useNodeLabelChangeHandler } from "@/routes/workflows/hooks/useLabelChangeHandler";
 import { UpdateIcon } from "@radix-ui/react-icons";
@@ -8,11 +14,13 @@ import {
   Handle,
   NodeProps,
   Position,
+  useEdges,
   useNodes,
   useReactFlow,
 } from "@xyflow/react";
-import { useState } from "react";
 import { AppNode } from "..";
+import { useWorkflowParametersState } from "../../useWorkflowParametersState";
+import { getAvailableOutputParameterKeys } from "../../workflowEditorUtils";
 import { EditableNodeTitle } from "../components/EditableNodeTitle";
 import { NodeActionMenu } from "../NodeActionMenu";
 import type { LoopNode } from "./types";
@@ -20,14 +28,19 @@ import type { LoopNode } from "./types";
 function LoopNode({ id, data }: NodeProps<LoopNode>) {
   const { updateNodeData } = useReactFlow();
   const nodes = useNodes<AppNode>();
+  const edges = useEdges();
   const [label, setLabel] = useNodeLabelChangeHandler({
     id,
     initialValue: data.label,
   });
   const deleteNodeCallback = useDeleteNodeCallback();
-  const [inputs, setInputs] = useState({
-    loopValue: data.loopValue,
-  });
+
+  const [workflowParameters] = useWorkflowParametersState();
+  const parameters = workflowParameters.filter(
+    (parameter) => parameter.parameterType !== "credential",
+  );
+  const parameterKeys = parameters.map((parameter) => parameter.key);
+  const outputParameterKeys = getAvailableOutputParameterKeys(nodes, edges, id);
 
   const children = nodes.filter((node) => node.parentId === id);
   const furthestDownChild: Node | null = children.reduce(
@@ -93,19 +106,28 @@ function LoopNode({ id, data }: NodeProps<LoopNode>) {
               />
             </div>
             <div className="space-y-1">
-              <Label className="text-xs text-slate-300">Loop Value</Label>
-              <Input
-                value={inputs.loopValue}
-                onChange={(event) => {
-                  if (!data.editable) {
-                    return;
-                  }
-                  setInputs({ ...inputs, loopValue: event.target.value });
-                  updateNodeData(id, { loopValue: event.target.value });
+              <Label className="text-xs text-slate-300">
+                Loop Value Parameter
+              </Label>
+              <Select
+                value={data.loopValue}
+                onValueChange={(value) => {
+                  updateNodeData(id, { loopValue: value });
                 }}
-                placeholder="What value are you iterating over?"
-                className="nopan"
-              />
+              >
+                <SelectTrigger>
+                  <SelectValue placeholder="Select a parameter" />
+                </SelectTrigger>
+                <SelectContent>
+                  {[...parameterKeys, ...outputParameterKeys].map(
+                    (parameterKey) => (
+                      <SelectItem key={parameterKey} value={parameterKey}>
+                        {parameterKey}
+                      </SelectItem>
+                    ),
+                  )}
+                </SelectContent>
+              </Select>
             </div>
           </div>
         </div>

skyvern-frontend/src/routes/workflows/editor/nodes/TaskNode/TaskNode.tsx

@@ -14,7 +14,6 @@ import { useDeleteNodeCallback } from "@/routes/workflows/hooks/useDeleteNodeCal
 import { useNodeLabelChangeHandler } from "@/routes/workflows/hooks/useLabelChangeHandler";
 import { ListBulletIcon } from "@radix-ui/react-icons";
 import {
-  Edge,
   Handle,
   NodeProps,
   Position,
@@ -24,36 +23,13 @@ import {
 } from "@xyflow/react";
 import { useState } from "react";
 import { AppNode } from "..";
-import { getOutputParameterKey } from "../../workflowEditorUtils";
+import { getAvailableOutputParameterKeys } from "../../workflowEditorUtils";
 import { EditableNodeTitle } from "../components/EditableNodeTitle";
 import { NodeActionMenu } from "../NodeActionMenu";
 import { TaskNodeDisplayModeSwitch } from "./TaskNodeDisplayModeSwitch";
 import { TaskNodeParametersPanel } from "./TaskNodeParametersPanel";
 import type { TaskNode, TaskNodeDisplayMode } from "./types";
 
-function getPreviousNodeIds(
-  nodes: Array<AppNode>,
-  edges: Array<Edge>,
-  target: string,
-): Array<string> {
-  const nodeIds: string[] = [];
-  const node = nodes.find((node) => node.id === target);
-  if (!node) {
-    return nodeIds;
-  }
-  let current = edges.find((edge) => edge.target === target);
-  if (current) {
-    while (current) {
-      nodeIds.push(current.source);
-      current = edges.find((edge) => edge.target === current!.source);
-    }
-  }
-  if (!node.parentId) {
-    return nodeIds;
-  }
-  return [...nodeIds, ...getPreviousNodeIds(nodes, edges, node.parentId)];
-}
-
 function TaskNode({ id, data }: NodeProps<TaskNode>) {
   const { updateNodeData } = useReactFlow();
   const [displayMode, setDisplayMode] = useState<TaskNodeDisplayMode>("basic");
@@ -61,16 +37,7 @@ function TaskNode({ id, data }: NodeProps<TaskNode>) {
   const deleteNodeCallback = useDeleteNodeCallback();
   const nodes = useNodes<AppNode>();
   const edges = useEdges();
-  const previousNodeIds = getPreviousNodeIds(nodes, edges, id);
-  const previousNodes = nodes.filter((node) =>
-    previousNodeIds.includes(node.id),
-  );
-  const labels = previousNodes
-    .filter((node) => node.type !== "nodeAdder")
-    .map((node) => node.data.label);
-  const outputParameterKeys = labels.map((label) =>
-    getOutputParameterKey(label),
-  );
+  const outputParameterKeys = getAvailableOutputParameterKeys(nodes, edges, id);
 
   const [label, setLabel] = useNodeLabelChangeHandler({
     id,

skyvern-frontend/src/routes/workflows/editor/workflowEditorUtils.ts

@@ -677,6 +677,9 @@ function getUpdatedNodesAfterLabelUpdateForParameterKeys(
   }
   const oldLabel = labelUpdatedNode.data.label as string;
   return nodes.map((node) => {
+    if (node.type === "nodeAdder") {
+      return node;
+    }
     if (node.type === "task") {
       return {
         ...node,
@@ -692,6 +695,18 @@ function getUpdatedNodesAfterLabelUpdateForParameterKeys(
         },
       };
     }
+    if (node.type === "loop") {
+      return {
+        ...node,
+        data: {
+          ...node.data,
+          loopValue:
+            node.data.loopValue === getOutputParameterKey(oldLabel)
+              ? getOutputParameterKey(newLabel)
+              : node.data.loopValue,
+        },
+      };
+    }
     return {
       ...node,
       data: {
@@ -814,6 +829,48 @@ function getDefaultValueForParameterType(
   }
 }
 
+function getPreviousNodeIds(
+  nodes: Array<AppNode>,
+  edges: Array<Edge>,
+  target: string,
+): Array<string> {
+  const nodeIds: string[] = [];
+  const node = nodes.find((node) => node.id === target);
+  if (!node) {
+    return nodeIds;
+  }
+  let current = edges.find((edge) => edge.target === target);
+  if (current) {
+    while (current) {
+      nodeIds.push(current.source);
+      current = edges.find((edge) => edge.target === current!.source);
+    }
+  }
+  if (!node.parentId) {
+    return nodeIds;
+  }
+  return [...nodeIds, ...getPreviousNodeIds(nodes, edges, node.parentId)];
+}
+
+function getAvailableOutputParameterKeys(
+  nodes: Array<AppNode>,
+  edges: Array<Edge>,
+  id: string,
+): Array<string> {
+  const previousNodeIds = getPreviousNodeIds(nodes, edges, id);
+  const previousNodes = nodes.filter((node) =>
+    previousNodeIds.includes(node.id),
+  );
+  const labels = previousNodes
+    .filter((node) => node.type !== "nodeAdder")
+    .map((node) => node.data.label);
+  const outputParameterKeys = labels.map((label) =>
+    getOutputParameterKey(label),
+  );
+
+  return outputParameterKeys;
+}
+
 export {
   createNode,
   generateNodeData,
@@ -830,4 +887,6 @@ export {
   getBlockNameOfOutputParameterKey,
   getDefaultValueForParameterType,
   getUpdatedParametersAfterLabelUpdateForSourceParameterKey,
+  getPreviousNodeIds,
+  getAvailableOutputParameterKeys,
 };