Add parameterKeys to textPrompt (#1191)

skyvern-frontend/src/routes/workflows/editor/FlowRenderer.tsx

@@ -418,6 +418,19 @@ function FlowRenderer({
           },
         };
       }
+      // TODO: Fix this. When we put these into the same if statement TS fails to recognize that the returned value fits both the task and text prompt node types
+      if (node.type === "textPrompt") {
+        return {
+          ...node,
+          data: {
+            ...node.data,
+            parameterKeys: node.data.parameterKeys.filter(
+              (parameter) =>
+                parameter !== getOutputParameterKey(deletedNodeLabel),
+            ),
+          },
+        };
+      }
       if (node.type === "loop") {
         return {
           ...node,

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

@@ -9,7 +9,7 @@ type Props = {
   onParametersChange: (parameters: Array<string>) => void;
 };
 
-function TaskNodeParametersPanel({
+function ParametersMultiSelect({
   availableOutputParameters,
   parameters,
   onParametersChange,
@@ -42,4 +42,4 @@ function TaskNodeParametersPanel({
   );
 }
 
-export { TaskNodeParametersPanel };
+export { ParametersMultiSelect };

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

@@ -27,7 +27,7 @@ import { AppNode } from "..";
 import { getAvailableOutputParameterKeys } from "../../workflowEditorUtils";
 import { EditableNodeTitle } from "../components/EditableNodeTitle";
 import { NodeActionMenu } from "../NodeActionMenu";
-import { TaskNodeParametersPanel } from "./TaskNodeParametersPanel";
+import { ParametersMultiSelect } from "./ParametersMultiSelect";
 import {
   dataSchemaExampleValue,
   errorMappingExampleValue,
@@ -152,7 +152,7 @@ function TaskNode({ id, data }: NodeProps<TaskNode>) {
                   />
                 </div>
                 <div className="space-y-2">
-                  <TaskNodeParametersPanel
+                  <ParametersMultiSelect
                     availableOutputParameters={outputParameterKeys}
                     parameters={data.parameterKeys}
                     onParametersChange={(parameterKeys) => {

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

@@ -6,12 +6,22 @@ import { CodeEditor } from "@/routes/workflows/components/CodeEditor";
 import { useDeleteNodeCallback } from "@/routes/workflows/hooks/useDeleteNodeCallback";
 import { useNodeLabelChangeHandler } from "@/routes/workflows/hooks/useLabelChangeHandler";
 import { CursorTextIcon } from "@radix-ui/react-icons";
-import { Handle, NodeProps, Position, useReactFlow } from "@xyflow/react";
+import {
+  Handle,
+  NodeProps,
+  Position,
+  useEdges,
+  useNodes,
+  useReactFlow,
+} from "@xyflow/react";
 import { useState } from "react";
 import { EditableNodeTitle } from "../components/EditableNodeTitle";
 import { NodeActionMenu } from "../NodeActionMenu";
 import { helpTooltipContent, type TextPromptNode } from "./types";
 import { HelpTooltip } from "@/components/HelpTooltip";
+import { ParametersMultiSelect } from "../TaskNode/ParametersMultiSelect";
+import { getAvailableOutputParameterKeys } from "../../workflowEditorUtils";
+import { AppNode } from "..";
 
 function TextPromptNode({ id, data }: NodeProps<TextPromptNode>) {
   const { updateNodeData } = useReactFlow();
@@ -22,6 +32,10 @@ function TextPromptNode({ id, data }: NodeProps<TextPromptNode>) {
     jsonSchema: data.jsonSchema,
   });
 
+  const nodes = useNodes<AppNode>();
+  const edges = useEdges();
+  const outputParameterKeys = getAvailableOutputParameterKeys(nodes, edges, id);
+
   const [label, setLabel] = useNodeLabelChangeHandler({
     id,
     initialValue: data.label,
@@ -82,6 +96,15 @@ function TextPromptNode({ id, data }: NodeProps<TextPromptNode>) {
             className="nopan text-xs"
           />
         </div>
+        <div className="space-y-2">
+          <ParametersMultiSelect
+            availableOutputParameters={outputParameterKeys}
+            parameters={data.parameterKeys}
+            onParametersChange={(parameterKeys) => {
+              updateNodeData(id, { parameterKeys });
+            }}
+          />
+        </div>
         <Separator />
         <div className="space-y-2">
           <div className="flex gap-2">

skyvern-frontend/src/routes/workflows/editor/nodes/TextPromptNode/types.ts

@@ -6,6 +6,7 @@ export type TextPromptNodeData = NodeBaseData & {
   jsonSchema: string;
   editable: boolean;
   label: string;
+  parameterKeys: Array<string>;
 };
 
 export type TextPromptNode = Node<TextPromptNodeData, "textPrompt">;
@@ -16,6 +17,7 @@ export const textPromptNodeDefaultData: TextPromptNodeData = {
   prompt: "",
   jsonSchema: "null",
   continueOnFailure: false,
+  parameterKeys: [],
 } as const;
 
 export const helpTooltipContent = {

skyvern-frontend/src/routes/workflows/editor/panels/WorkflowParametersPanel.tsx

@@ -176,7 +176,10 @@ function WorkflowParametersPanel() {
                                 setHasChanges(true);
                                 setNodes((nodes) => {
                                   return nodes.map((node) => {
-                                    if (node.type === "task") {
+                                    if (
+                                      node.type === "task" ||
+                                      node.type === "textPrompt"
+                                    ) {
                                       return {
                                         ...node,
                                         data: {
@@ -270,7 +273,10 @@ function WorkflowParametersPanel() {
                     );
                     setNodes((nodes) => {
                       return nodes.map((node) => {
-                        if (node.type === "task") {
+                        if (
+                          node.type === "task" ||
+                          node.type === "textPrompt"
+                        ) {
                           return {
                             ...node,
                             data: {

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

@@ -201,6 +201,7 @@ function convertToNode(
           editable: true,
           prompt: block.prompt,
           jsonSchema: JSON.stringify(block.json_schema, null, 2),
+          parameterKeys: block.parameters.map((p) => p.key),
         },
       };
     }
@@ -627,6 +628,7 @@ function getWorkflowBlock(node: WorkflowBlockNode): BlockYAML {
         llm_key: "",
         prompt: node.data.prompt,
         json_schema: JSONParseSafe(node.data.jsonSchema),
+        parameter_keys: node.data.parameterKeys,
       };
     }
     default: {
@@ -758,7 +760,7 @@ function getUpdatedNodesAfterLabelUpdateForParameterKeys(
     if (node.type === "nodeAdder" || node.type === "start") {
       return node;
     }
-    if (node.type === "task") {
+    if (node.type === "task" || node.type === "textPrompt") {
       return {
         ...node,
         data: {