From b73602dad4151e1e84477b644b067d46a565ca9e Mon Sep 17 00:00:00 2001
From: Shuchang Zheng <wintonzheng0325@gmail.com>
Date: Wed, 6 Nov 2024 12:29:38 -0800
Subject: [PATCH] Add secret parameter to workflow (#1150)

---
 .../routes/workflows/editor/FlowRenderer.tsx  | 32 +++++++--
 .../workflows/editor/WorkflowEditor.tsx       | 14 +++-
 .../panels/WorkflowParameterAddPanel.tsx      | 56 +++++++++++++++-
 .../panels/WorkflowParameterEditPanel.tsx     | 66 ++++++++++++++++++-
 .../editor/panels/WorkflowParametersPanel.tsx | 13 +++-
 .../workflows/editor/workflowEditorUtils.ts   | 25 +------
 6 files changed, 170 insertions(+), 36 deletions(-)

diff --git a/skyvern-frontend/src/routes/workflows/editor/FlowRenderer.tsx b/skyvern-frontend/src/routes/workflows/editor/FlowRenderer.tsx
index 6221f480b..6740e10bb 100644
--- a/skyvern-frontend/src/routes/workflows/editor/FlowRenderer.tsx
+++ b/skyvern-frontend/src/routes/workflows/editor/FlowRenderer.tsx
@@ -16,12 +16,12 @@ import { nanoid } from "nanoid";
 import { useEffect, useState } from "react";
 import {
   AWSSecretParameter,
-  BitwardenSensitiveInformationParameter,
   WorkflowApiResponse,
   WorkflowParameterValueType,
 } from "../types/workflowTypes";
 import {
   BitwardenLoginCredentialParameterYAML,
+  BitwardenSensitiveInformationParameterYAML,
   BlockYAML,
   ContextParameterYAML,
   ParameterYAML,
@@ -74,6 +74,7 @@ function convertToParametersYAML(
   | WorkflowParameterYAML
   | BitwardenLoginCredentialParameterYAML
   | ContextParameterYAML
+  | BitwardenSensitiveInformationParameterYAML
 > {
   return parameters.map((parameter) => {
     if (parameter.parameterType === "workflow") {
@@ -93,6 +94,20 @@ function convertToParametersYAML(
         description: parameter.description || null,
         source_parameter_key: parameter.sourceParameterKey,
       };
+    } else if (parameter.parameterType === "secret") {
+      return {
+        parameter_type: "bitwarden_sensitive_information",
+        key: parameter.key,
+        bitwarden_identity_key: parameter.identityKey,
+        bitwarden_identity_fields: parameter.identityFields,
+        description: parameter.description || null,
+        bitwarden_collection_id: parameter.collectionId,
+        bitwarden_client_id_aws_secret_key: "SKYVERN_BITWARDEN_CLIENT_ID",
+        bitwarden_client_secret_aws_secret_key:
+          "SKYVERN_BITWARDEN_CLIENT_SECRET",
+        bitwarden_master_password_aws_secret_key:
+          "SKYVERN_BITWARDEN_MASTER_PASSWORD",
+      };
     } else {
       return {
         parameter_type: "bitwarden_login_credential",
@@ -131,6 +146,14 @@ export type ParametersState = Array<
       sourceParameterKey: string;
       description?: string | null;
     }
+  | {
+      key: string;
+      parameterType: "secret";
+      identityKey: string;
+      identityFields: Array<string>;
+      collectionId: string;
+      description?: string | null;
+    }
 >;
 
 type Props = {
@@ -246,12 +269,9 @@ function FlowRenderer({
     const parametersInYAMLConvertibleJSON = convertToParametersYAML(parameters);
     const filteredParameters = workflow.workflow_definition.parameters.filter(
       (parameter) => {
-        return (
-          parameter.parameter_type === "aws_secret" ||
-          parameter.parameter_type === "bitwarden_sensitive_information"
-        );
+        return parameter.parameter_type === "aws_secret";
       },
-    ) as Array<AWSSecretParameter | BitwardenSensitiveInformationParameter>;
+    ) as Array<AWSSecretParameter>;
 
     const echoParameters = convertEchoParameters(filteredParameters);
 
diff --git a/skyvern-frontend/src/routes/workflows/editor/WorkflowEditor.tsx b/skyvern-frontend/src/routes/workflows/editor/WorkflowEditor.tsx
index 8545d490a..a6d054527 100644
--- a/skyvern-frontend/src/routes/workflows/editor/WorkflowEditor.tsx
+++ b/skyvern-frontend/src/routes/workflows/editor/WorkflowEditor.tsx
@@ -53,7 +53,8 @@ function WorkflowEditor() {
               (parameter) =>
                 parameter.parameter_type === "workflow" ||
                 parameter.parameter_type === "bitwarden_login_credential" ||
-                parameter.parameter_type === "context",
+                parameter.parameter_type === "context" ||
+                parameter.parameter_type === "bitwarden_sensitive_information",
             )
             .map((parameter) => {
               if (parameter.parameter_type === "workflow") {
@@ -71,6 +72,17 @@ function WorkflowEditor() {
                   sourceParameterKey: parameter.source.key,
                   description: parameter.description,
                 };
+              } else if (
+                parameter.parameter_type === "bitwarden_sensitive_information"
+              ) {
+                return {
+                  key: parameter.key,
+                  parameterType: "secret",
+                  collectionId: parameter.bitwarden_collection_id,
+                  identityKey: parameter.bitwarden_identity_key,
+                  identityFields: parameter.bitwarden_identity_fields,
+                  description: parameter.description,
+                };
               } else {
                 return {
                   key: parameter.key,
diff --git a/skyvern-frontend/src/routes/workflows/editor/panels/WorkflowParameterAddPanel.tsx b/skyvern-frontend/src/routes/workflows/editor/panels/WorkflowParameterAddPanel.tsx
index 35474ecdf..52b984b92 100644
--- a/skyvern-frontend/src/routes/workflows/editor/panels/WorkflowParameterAddPanel.tsx
+++ b/skyvern-frontend/src/routes/workflows/editor/panels/WorkflowParameterAddPanel.tsx
@@ -20,7 +20,7 @@ import { toast } from "@/components/ui/use-toast";
 import { SourceParameterKeySelector } from "../../components/SourceParameterKeySelector";
 
 type Props = {
-  type: "workflow" | "credential" | "context";
+  type: "workflow" | "credential" | "context" | "secret";
   onClose: () => void;
   onSave: (value: ParametersState[number]) => void;
 };
@@ -34,13 +34,16 @@ const workflowParameterTypeOptions = [
   { label: "JSON", value: WorkflowParameterValueType.JSON },
 ];
 
-function header(type: "workflow" | "credential" | "context") {
+function header(type: "workflow" | "credential" | "context" | "secret") {
   if (type === "workflow") {
     return "Add Input Parameter";
   }
   if (type === "credential") {
     return "Add Credential Parameter";
   }
+  if (type === "secret") {
+    return "Add Secret Parameter";
+  }
   return "Add Context Parameter";
 }
 
@@ -62,6 +65,9 @@ function WorkflowParameterAddPanel({ type, onClose, onSave }: Props) {
     string | undefined
   >(undefined);
 
+  const [identityKey, setIdentityKey] = useState("");
+  const [identityFields, setIdentityFields] = useState("");
+
   return (
     <div className="space-y-4">
       <header className="flex items-center justify-between">
@@ -192,6 +198,31 @@ function WorkflowParameterAddPanel({ type, onClose, onSave }: Props) {
           />
         </div>
       )}
+      {type === "secret" && (
+        <>
+          <div className="space-y-1">
+            <Label className="text-xs text-slate-300">Identity Key</Label>
+            <Input
+              value={identityKey}
+              onChange={(e) => setIdentityKey(e.target.value)}
+            />
+          </div>
+          <div className="space-y-1">
+            <Label className="text-xs text-slate-300">Identity Fields</Label>
+            <Input
+              value={identityFields}
+              onChange={(e) => setIdentityFields(e.target.value)}
+            />
+          </div>
+          <div className="space-y-1">
+            <Label className="text-xs text-slate-300">Collection ID</Label>
+            <Input
+              value={collectionId}
+              onChange={(e) => setCollectionId(e.target.value)}
+            />
+          </div>
+        </>
+      )}
       <div className="flex justify-end">
         <Button
           onClick={() => {
@@ -243,6 +274,27 @@ function WorkflowParameterAddPanel({ type, onClose, onSave }: Props) {
                 description,
               });
             }
+            if (type === "secret") {
+              if (!collectionId) {
+                toast({
+                  variant: "destructive",
+                  title: "Failed to add parameter",
+                  description: "Collection ID is required",
+                });
+                return;
+              }
+              onSave({
+                key,
+                parameterType: "secret",
+                collectionId,
+                identityFields: identityFields
+                  .split(",")
+                  .filter((s) => s.length > 0)
+                  .map((field) => field.trim()),
+                identityKey,
+                description,
+              });
+            }
             if (type === "context") {
               if (!sourceParameterKey) {
                 toast({
diff --git a/skyvern-frontend/src/routes/workflows/editor/panels/WorkflowParameterEditPanel.tsx b/skyvern-frontend/src/routes/workflows/editor/panels/WorkflowParameterEditPanel.tsx
index 251ab003a..3cf05c642 100644
--- a/skyvern-frontend/src/routes/workflows/editor/panels/WorkflowParameterEditPanel.tsx
+++ b/skyvern-frontend/src/routes/workflows/editor/panels/WorkflowParameterEditPanel.tsx
@@ -20,7 +20,7 @@ import { toast } from "@/components/ui/use-toast";
 import { SourceParameterKeySelector } from "../../components/SourceParameterKeySelector";
 
 type Props = {
-  type: "workflow" | "credential" | "context";
+  type: "workflow" | "credential" | "context" | "secret";
   onClose: () => void;
   onSave: (value: ParametersState[number]) => void;
   initialValues: ParametersState[number];
@@ -35,13 +35,16 @@ const workflowParameterTypeOptions = [
   { label: "JSON", value: WorkflowParameterValueType.JSON },
 ];
 
-function header(type: "workflow" | "credential" | "context") {
+function header(type: "workflow" | "credential" | "context" | "secret") {
   if (type === "workflow") {
     return "Edit Input Parameter";
   }
   if (type === "credential") {
     return "Edit Credential Parameter";
   }
+  if (type === "secret") {
+    return "Edit Secret Parameter";
+  }
   return "Edit Context Parameter";
 }
 
@@ -61,7 +64,8 @@ function WorkflowParameterEditPanel({
     initialValues.description ?? "",
   );
   const [collectionId, setCollectionId] = useState(
-    initialValues.parameterType === "credential"
+    initialValues.parameterType === "credential" ||
+      initialValues.parameterType === "secret"
       ? initialValues.collectionId
       : "",
   );
@@ -95,6 +99,16 @@ function WorkflowParameterEditPanel({
       : undefined,
   );
 
+  const [identityKey, setIdentityKey] = useState(
+    initialValues.parameterType === "secret" ? initialValues.identityKey : "",
+  );
+
+  const [identityFields, setIdentityFields] = useState(
+    initialValues.parameterType === "secret"
+      ? initialValues.identityFields.join(", ")
+      : "",
+  );
+
   return (
     <div className="space-y-4">
       <header className="flex items-center justify-between">
@@ -225,6 +239,31 @@ function WorkflowParameterEditPanel({
           />
         </div>
       )}
+      {type === "secret" && (
+        <>
+          <div className="space-y-1">
+            <Label className="text-xs text-slate-300">Identity Key</Label>
+            <Input
+              value={identityKey}
+              onChange={(e) => setIdentityKey(e.target.value)}
+            />
+          </div>
+          <div className="space-y-1">
+            <Label className="text-xs text-slate-300">Identity Fields</Label>
+            <Input
+              value={identityFields}
+              onChange={(e) => setIdentityFields(e.target.value)}
+            />
+          </div>
+          <div className="space-y-1">
+            <Label className="text-xs text-slate-300">Collection ID</Label>
+            <Input
+              value={collectionId}
+              onChange={(e) => setCollectionId(e.target.value)}
+            />
+          </div>
+        </>
+      )}
       <div className="flex justify-end">
         <Button
           onClick={() => {
@@ -276,6 +315,27 @@ function WorkflowParameterEditPanel({
                 description,
               });
             }
+            if (type === "secret") {
+              if (!collectionId) {
+                toast({
+                  variant: "destructive",
+                  title: "Failed to add parameter",
+                  description: "Collection ID is required",
+                });
+                return;
+              }
+              onSave({
+                key,
+                parameterType: "secret",
+                collectionId,
+                identityFields: identityFields
+                  .split(",")
+                  .filter((s) => s.length > 0)
+                  .map((field) => field.trim()),
+                identityKey,
+                description,
+              });
+            }
             if (type === "context") {
               if (!sourceParameterKey) {
                 toast({
diff --git a/skyvern-frontend/src/routes/workflows/editor/panels/WorkflowParametersPanel.tsx b/skyvern-frontend/src/routes/workflows/editor/panels/WorkflowParametersPanel.tsx
index 837d443d3..f3c25ab03 100644
--- a/skyvern-frontend/src/routes/workflows/editor/panels/WorkflowParametersPanel.tsx
+++ b/skyvern-frontend/src/routes/workflows/editor/panels/WorkflowParametersPanel.tsx
@@ -41,7 +41,7 @@ function WorkflowParametersPanel() {
     active: boolean;
     operation: "add" | "edit";
     parameter?: ParametersState[number] | null;
-    type: "workflow" | "credential" | "context";
+    type: "workflow" | "credential" | "context" | "secret";
   }>({
     active: false,
     operation: "add",
@@ -103,6 +103,17 @@ function WorkflowParametersPanel() {
             >
               Context Parameter
             </DropdownMenuItem>
+            <DropdownMenuItem
+              onClick={() => {
+                setOperationPanelState({
+                  active: true,
+                  operation: "add",
+                  type: "secret",
+                });
+              }}
+            >
+              Secret Parameter
+            </DropdownMenuItem>
           </DropdownMenuContent>
         </DropdownMenu>
 
diff --git a/skyvern-frontend/src/routes/workflows/editor/workflowEditorUtils.ts b/skyvern-frontend/src/routes/workflows/editor/workflowEditorUtils.ts
index fad6c3fcb..7384c8a43 100644
--- a/skyvern-frontend/src/routes/workflows/editor/workflowEditorUtils.ts
+++ b/skyvern-frontend/src/routes/workflows/editor/workflowEditorUtils.ts
@@ -4,8 +4,6 @@ import { Edge } from "@xyflow/react";
 import { nanoid } from "nanoid";
 import type {
   AWSSecretParameter,
-  BitwardenSensitiveInformationParameter,
-  ContextParameter,
   OutputParameter,
   Parameter,
   WorkflowApiResponse,
@@ -45,11 +43,11 @@ import { fileParserNodeDefaultData } from "./nodes/FileParserNode/types";
 import { LoopNode, loopNodeDefaultData } from "./nodes/LoopNode/types";
 import { NodeAdderNode } from "./nodes/NodeAdderNode/types";
 import { sendEmailNodeDefaultData } from "./nodes/SendEmailNode/types";
+import { StartNode } from "./nodes/StartNode/types";
 import { taskNodeDefaultData } from "./nodes/TaskNode/types";
 import { textPromptNodeDefaultData } from "./nodes/TextPromptNode/types";
 import { NodeBaseData } from "./nodes/types";
 import { uploadNodeDefaultData } from "./nodes/UploadNode/types";
-import { StartNode } from "./nodes/StartNode/types";
 
 export const NEW_NODE_LABEL_PREFIX = "block_";
 
@@ -685,11 +683,7 @@ function generateNodeLabel(existingLabels: Array<string>) {
  * If a parameter is not displayed in the editor, we should echo its value back when saved.
  */
 function convertEchoParameters(
-  parameters: Array<
-    | ContextParameter
-    | BitwardenSensitiveInformationParameter
-    | AWSSecretParameter
-  >,
+  parameters: Array<AWSSecretParameter>,
 ): Array<ParameterYAML> {
   return parameters.map((parameter) => {
     if (parameter.parameter_type === "aws_secret") {
@@ -699,21 +693,6 @@ function convertEchoParameters(
         aws_key: parameter.aws_key,
       };
     }
-    if (parameter.parameter_type === "bitwarden_sensitive_information") {
-      return {
-        key: parameter.key,
-        parameter_type: "bitwarden_sensitive_information",
-        bitwarden_collection_id: parameter.bitwarden_collection_id,
-        bitwarden_identity_key: parameter.bitwarden_identity_key,
-        bitwarden_identity_fields: parameter.bitwarden_identity_fields,
-        bitwarden_client_id_aws_secret_key:
-          parameter.bitwarden_client_id_aws_secret_key,
-        bitwarden_client_secret_aws_secret_key:
-          parameter.bitwarden_client_secret_aws_secret_key,
-        bitwarden_master_password_aws_secret_key:
-          parameter.bitwarden_master_password_aws_secret_key,
-      };
-    }
     throw new Error("Unknown parameter type");
   });
 }