From a6d6965f4b97efcaafb1fef5d65e5aedb6ed0fbc Mon Sep 17 00:00:00 2001
From: Shuchang Zheng <wintonzheng0325@gmail.com>
Date: Thu, 7 Nov 2024 09:44:48 -0800
Subject: [PATCH] Add proxy location to workflow run screen (#1158)

---
 skyvern-frontend/src/api/types.ts             |  21 ++-
 .../src/routes/workflows/RunWorkflowForm.tsx  | 157 +++++++++++++-----
 2 files changed, 133 insertions(+), 45 deletions(-)

diff --git a/skyvern-frontend/src/api/types.ts b/skyvern-frontend/src/api/types.ts
index ea013479..fd252ad6 100644
--- a/skyvern-frontend/src/api/types.ts
+++ b/skyvern-frontend/src/api/types.ts
@@ -12,6 +12,8 @@ export const ArtifactType = {
   HTMLScrape: "html_scrape",
 } as const;
 
+export type ArtifactType = (typeof ArtifactType)[keyof typeof ArtifactType];
+
 export const Status = {
   Created: "created",
   Running: "running",
@@ -25,7 +27,16 @@ export const Status = {
 
 export type Status = (typeof Status)[keyof typeof Status];
 
-export type ArtifactType = (typeof ArtifactType)[keyof typeof ArtifactType];
+export const ProxyLocation = {
+  Residential: "RESIDENTIAL",
+  ResidentialIE: "RESIDENTIAL_IE",
+  ResidentialES: "RESIDENTIAL_ES",
+  ResidentialIN: "RESIDENTIAL_IN",
+  ResidentialJP: "RESIDENTIAL_JP",
+  None: "NONE",
+} as const;
+
+export type ProxyLocation = (typeof ProxyLocation)[keyof typeof ProxyLocation];
 
 export type ArtifactApiResponse = {
   created_at: string;
@@ -90,7 +101,7 @@ export type CreateTaskRequest = {
   navigation_payload: Record<string, unknown> | string | null;
   extracted_information_schema: Record<string, unknown> | string | null;
   error_code_mapping: Record<string, string> | null;
-  proxy_location: string | null;
+  proxy_location: ProxyLocation | null;
   totp_verification_url: string | null;
   totp_identifier: string | null;
 };
@@ -184,7 +195,7 @@ export type WorkflowApiResponse = {
     parameters: Array<WorkflowParameter>;
     blocks: Array<WorkflowBlock>;
   };
-  proxy_location: string;
+  proxy_location: ProxyLocation | null;
   webhook_callback_url: string;
   created_at: string;
   modified_at: string;
@@ -248,7 +259,7 @@ export type WorkflowRunApiResponse = {
   workflow_run_id: string;
   workflow_id: string;
   status: Status;
-  proxy_location: string;
+  proxy_location: ProxyLocation | null;
   webhook_callback_url: string;
   created_at: string;
   modified_at: string;
@@ -258,7 +269,7 @@ export type WorkflowRunStatusApiResponse = {
   workflow_id: string;
   workflow_run_id: string;
   status: Status;
-  proxy_location: string;
+  proxy_location: ProxyLocation | null;
   webhook_callback_url: string | null;
   created_at: string;
   modified_at: string;
diff --git a/skyvern-frontend/src/routes/workflows/RunWorkflowForm.tsx b/skyvern-frontend/src/routes/workflows/RunWorkflowForm.tsx
index 5f8318f8..e2e6e2a8 100644
--- a/skyvern-frontend/src/routes/workflows/RunWorkflowForm.tsx
+++ b/skyvern-frontend/src/routes/workflows/RunWorkflowForm.tsx
@@ -23,6 +23,14 @@ import { copyText } from "@/util/copyText";
 import { WorkflowParameter } from "./types/workflowTypes";
 import { Input } from "@/components/ui/input";
 import { z } from "zod";
+import { ProxyLocation } from "@/api/types";
+import {
+  Select,
+  SelectContent,
+  SelectItem,
+  SelectTrigger,
+  SelectValue,
+} from "@/components/ui/select";
 
 type Props = {
   workflowParameters: Array<WorkflowParameter>;
@@ -60,38 +68,61 @@ function parseValuesForWorkflowRun(
   );
 }
 
+type RunWorkflowRequestBody = {
+  data: Record<string, unknown>; // workflow parameters and values
+  proxy_location: ProxyLocation | null;
+  webhook_callback_url?: string | null;
+};
+
+function getRunWorkflowRequestBody(
+  values: RunWorkflowFormType,
+  workflowParameters: Array<WorkflowParameter>,
+): RunWorkflowRequestBody {
+  const { webhookCallbackUrl, proxyLocation, ...parameters } = values;
+
+  const parsedParameters = parseValuesForWorkflowRun(
+    parameters,
+    workflowParameters,
+  );
+
+  const body: RunWorkflowRequestBody = {
+    data: parsedParameters,
+    proxy_location: proxyLocation,
+  };
+
+  if (webhookCallbackUrl) {
+    body.webhook_callback_url = webhookCallbackUrl;
+  }
+
+  return body;
+}
+
+type RunWorkflowFormType = Record<string, unknown> & {
+  webhookCallbackUrl: string | null;
+  proxyLocation: ProxyLocation | null;
+};
+
 function RunWorkflowForm({ workflowParameters, initialValues }: Props) {
   const { workflowPermanentId } = useParams();
   const credentialGetter = useCredentialGetter();
   const queryClient = useQueryClient();
-  const form = useForm<Record<string, unknown>>({
-    defaultValues: { ...initialValues, webhookCallbackUrl: null },
+  const form = useForm<RunWorkflowFormType>({
+    defaultValues: {
+      ...initialValues,
+      webhookCallbackUrl: null,
+      proxyLocation: ProxyLocation.Residential,
+    },
   });
   const apiCredential = useApiCredential();
 
   const runWorkflowMutation = useMutation({
-    mutationFn: async (values: Record<string, unknown>) => {
+    mutationFn: async (values: RunWorkflowFormType) => {
       const client = await getClient(credentialGetter);
-
-      const { webhookCallbackUrl, ...parameters } = values;
-
-      const body: {
-        data: Record<string, unknown>;
-        proxy_location: string;
-        webhook_callback_url?: string;
-      } = {
-        data: parameters,
-        proxy_location: "RESIDENTIAL",
-      };
-
-      if (webhookCallbackUrl) {
-        body.webhook_callback_url = webhookCallbackUrl as string;
-      }
-
-      return client.post<unknown, { data: { workflow_run_id: string } }>(
-        `/workflows/${workflowPermanentId}/run`,
-        body,
-      );
+      const body = getRunWorkflowRequestBody(values, workflowParameters);
+      return client.post<
+        RunWorkflowRequestBody,
+        { data: { workflow_run_id: string } }
+      >(`/workflows/${workflowPermanentId}/run`, body);
     },
     onSuccess: (response) => {
       toast({
@@ -123,8 +154,8 @@ function RunWorkflowForm({ workflowParameters, initialValues }: Props) {
     },
   });
 
-  function onSubmit(values: Record<string, unknown>) {
-    const { webhookCallbackUrl, ...parameters } = values;
+  function onSubmit(values: RunWorkflowFormType) {
+    const { webhookCallbackUrl, proxyLocation, ...parameters } = values;
     const parsedParameters = parseValuesForWorkflowRun(
       parameters,
       workflowParameters,
@@ -132,6 +163,7 @@ function RunWorkflowForm({ workflowParameters, initialValues }: Props) {
     runWorkflowMutation.mutate({
       ...parsedParameters,
       webhookCallbackUrl,
+      proxyLocation,
     });
   }
 
@@ -216,7 +248,7 @@ function RunWorkflowForm({ workflowParameters, initialValues }: Props) {
           <FormField
             key="webhookCallbackUrl"
             control={form.control}
-            name={"webhookCallbackUrl"}
+            name="webhookCallbackUrl"
             rules={{
               validate: (value) => {
                 if (value === null) {
@@ -264,6 +296,62 @@ function RunWorkflowForm({ workflowParameters, initialValues }: Props) {
               );
             }}
           />
+          <FormField
+            key="proxyLocation"
+            control={form.control}
+            name="proxyLocation"
+            render={({ field }) => {
+              return (
+                <FormItem>
+                  <div className="flex gap-16">
+                    <FormLabel>
+                      <div className="w-72">
+                        <div className="flex items-center gap-2 text-lg">
+                          Proxy Location
+                        </div>
+                        <h2 className="text-sm text-slate-400">
+                          Route Skyvern through one of our available proxies.
+                        </h2>
+                      </div>
+                    </FormLabel>
+                    <div className="w-full space-y-2">
+                      <FormControl>
+                        <Select
+                          value={field.value ?? ""}
+                          onValueChange={field.onChange}
+                        >
+                          <SelectTrigger className="w-48">
+                            <SelectValue placeholder="Proxy Location" />
+                          </SelectTrigger>
+                          <SelectContent>
+                            <SelectItem value={ProxyLocation.Residential}>
+                              Residential
+                            </SelectItem>
+                            <SelectItem value={ProxyLocation.ResidentialES}>
+                              Residential (Spain)
+                            </SelectItem>
+                            <SelectItem value={ProxyLocation.ResidentialIE}>
+                              Residential (Ireland)
+                            </SelectItem>
+                            <SelectItem value={ProxyLocation.ResidentialIN}>
+                              Residential (India)
+                            </SelectItem>
+                            <SelectItem value={ProxyLocation.ResidentialJP}>
+                              Residential (Japan)
+                            </SelectItem>
+                            <SelectItem value={ProxyLocation.None}>
+                              None
+                            </SelectItem>
+                          </SelectContent>
+                        </Select>
+                      </FormControl>
+                      <FormMessage />
+                    </div>
+                  </div>
+                </FormItem>
+              );
+            }}
+          />
         </div>
 
         <div className="flex justify-end gap-2">
@@ -272,22 +360,10 @@ function RunWorkflowForm({ workflowParameters, initialValues }: Props) {
             variant="secondary"
             onClick={() => {
               const values = form.getValues();
-              const { webhookCallbackUrl, ...parameters } = values;
-              const parsedParameters = parseValuesForWorkflowRun(
-                parameters,
+              const body = getRunWorkflowRequestBody(
+                values,
                 workflowParameters,
               );
-              const body: {
-                data: Record<string, unknown>;
-                proxy_location: string;
-                webhook_callback_url?: string;
-              } = {
-                data: parsedParameters,
-                proxy_location: "RESIDENTIAL",
-              };
-              if (webhookCallbackUrl) {
-                body.webhook_callback_url = webhookCallbackUrl as string;
-              }
 
               const curl = fetchToCurl({
                 method: "POST",
@@ -298,6 +374,7 @@ function RunWorkflowForm({ workflowParameters, initialValues }: Props) {
                   "x-api-key": apiCredential ?? "<your-api-key>",
                 },
               });
+
               copyText(curl).then(() => {
                 toast({
                   variant: "success",