From eb668baba895ac01621ca054fffe9d9c0d67179c Mon Sep 17 00:00:00 2001
From: Salih Altun <salih@skyvern.com>
Date: Tue, 16 Apr 2024 20:40:59 +0300
Subject: [PATCH] Changes for signed URL artifacts (#197)

---
 skyvern-frontend/src/api/types.ts             |  1 +
 .../src/routes/tasks/detail/JSONArtifact.tsx  | 27 ++++---
 .../src/routes/tasks/detail/StepArtifacts.tsx | 78 +++++++++----------
 .../src/routes/tasks/detail/TaskDetails.tsx   | 12 +--
 .../src/routes/tasks/detail/TextArtifact.tsx  | 27 ++++---
 .../src/routes/tasks/detail/artifactUtils.ts  | 31 ++++++++
 6 files changed, 106 insertions(+), 70 deletions(-)
 create mode 100644 skyvern-frontend/src/routes/tasks/detail/artifactUtils.ts

diff --git a/skyvern-frontend/src/api/types.ts b/skyvern-frontend/src/api/types.ts
index fa96f1a9..1803260b 100644
--- a/skyvern-frontend/src/api/types.ts
+++ b/skyvern-frontend/src/api/types.ts
@@ -31,6 +31,7 @@ export type ArtifactApiResponse = {
   step_id: string;
   artifact_type: ArtifactType;
   uri: string;
+  signed_url?: string | null;
   organization_id: string;
 };
 
diff --git a/skyvern-frontend/src/routes/tasks/detail/JSONArtifact.tsx b/skyvern-frontend/src/routes/tasks/detail/JSONArtifact.tsx
index 128f5fc9..b2d223be 100644
--- a/skyvern-frontend/src/routes/tasks/detail/JSONArtifact.tsx
+++ b/skyvern-frontend/src/routes/tasks/detail/JSONArtifact.tsx
@@ -1,25 +1,32 @@
 import { artifactApiClient } from "@/api/AxiosClient";
+import { ArtifactApiResponse } from "@/api/types";
 import { Skeleton } from "@/components/ui/skeleton";
 import { Textarea } from "@/components/ui/textarea";
 import { useQuery } from "@tanstack/react-query";
+import axios from "axios";
 
 type Props = {
-  uri: string;
+  artifact: ArtifactApiResponse;
 };
 
-function JSONArtifact({ uri }: Props) {
+function JSONArtifact({ artifact }: Props) {
   const { data, isFetching, isError, error } = useQuery<
     Record<string, unknown>
   >({
-    queryKey: ["artifact", uri],
+    queryKey: ["artifact", artifact.artifact_id],
     queryFn: async () => {
-      return artifactApiClient
-        .get(`/artifact/json`, {
-          params: {
-            path: uri.slice(7),
-          },
-        })
-        .then((response) => response.data);
+      if (artifact.uri.startsWith("file://")) {
+        return artifactApiClient
+          .get(`/artifact/json`, {
+            params: {
+              path: artifact.uri.slice(7),
+            },
+          })
+          .then((response) => response.data);
+      }
+      if (artifact.uri.startsWith("s3://") && artifact.signed_url) {
+        return axios.get(artifact.signed_url).then((response) => response.data);
+      }
     },
   });
 
diff --git a/skyvern-frontend/src/routes/tasks/detail/StepArtifacts.tsx b/skyvern-frontend/src/routes/tasks/detail/StepArtifacts.tsx
index 20da68bd..784da2c8 100644
--- a/skyvern-frontend/src/routes/tasks/detail/StepArtifacts.tsx
+++ b/skyvern-frontend/src/routes/tasks/detail/StepArtifacts.tsx
@@ -9,11 +9,11 @@ import { Label } from "@/components/ui/label";
 import { useQuery } from "@tanstack/react-query";
 import { useParams } from "react-router-dom";
 import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs";
-import { artifactApiBaseUrl } from "@/util/env";
 import { ZoomableImage } from "@/components/ZoomableImage";
 import { Skeleton } from "@/components/ui/skeleton";
 import { JSONArtifact } from "./JSONArtifact";
 import { TextArtifact } from "./TextArtifact";
+import { getImageURL } from "./artifactUtils";
 
 type Props = {
   id: string;
@@ -40,46 +40,42 @@ function StepArtifacts({ id, stepProps }: Props) {
     return <div>Error: {error?.message}</div>;
   }
 
-  const llmScreenshotUris = artifacts
-    ?.filter(
-      (artifact) => artifact.artifact_type === ArtifactType.LLMScreenshot,
-    )
-    .map((artifact) => artifact.uri);
+  const llmScreenshots = artifacts?.filter(
+    (artifact) => artifact.artifact_type === ArtifactType.LLMScreenshot,
+  );
 
-  const actionScreenshotUris = artifacts
-    ?.filter(
-      (artifact) => artifact.artifact_type === ArtifactType.ActionScreenshot,
-    )
-    .map((artifact) => artifact.uri);
+  const actionScreenshots = artifacts?.filter(
+    (artifact) => artifact.artifact_type === ArtifactType.ActionScreenshot,
+  );
 
-  const visibleElementsTreeUri = artifacts?.find(
+  const visibleElementsTree = artifacts?.find(
     (artifact) => artifact.artifact_type === ArtifactType.VisibleElementsTree,
-  )?.uri;
+  );
 
-  const visibleElementsTreeTrimmedUri = artifacts?.find(
+  const visibleElementsTreeTrimmed = artifacts?.find(
     (artifact) =>
       artifact.artifact_type === ArtifactType.VisibleElementsTreeTrimmed,
-  )?.uri;
+  );
 
-  const llmPromptUri = artifacts?.find(
+  const llmPrompt = artifacts?.find(
     (artifact) => artifact.artifact_type === ArtifactType.LLMPrompt,
-  )?.uri;
+  );
 
-  const llmRequestUri = artifacts?.find(
+  const llmRequest = artifacts?.find(
     (artifact) => artifact.artifact_type === ArtifactType.LLMRequest,
-  )?.uri;
+  );
 
-  const llmResponseRawUri = artifacts?.find(
+  const llmResponseRaw = artifacts?.find(
     (artifact) => artifact.artifact_type === ArtifactType.LLMResponseRaw,
-  )?.uri;
+  );
 
-  const llmResponseParsedUri = artifacts?.find(
+  const llmResponseParsed = artifacts?.find(
     (artifact) => artifact.artifact_type === ArtifactType.LLMResponseParsed,
-  )?.uri;
+  );
 
-  const htmlRawUri = artifacts?.find(
+  const htmlRaw = artifacts?.find(
     (artifact) => artifact.artifact_type === ArtifactType.HTMLScrape,
-  )?.uri;
+  );
 
   return (
     <Tabs defaultValue="info" className="w-full">
@@ -128,12 +124,12 @@ function StepArtifacts({ id, stepProps }: Props) {
         </div>
       </TabsContent>
       <TabsContent value="screenshot_llm">
-        {llmScreenshotUris && llmScreenshotUris.length > 0 ? (
+        {llmScreenshots && llmScreenshots.length > 0 ? (
           <div className="grid grid-cols-3 gap-4 p-4">
-            {llmScreenshotUris.map((uri, index) => (
+            {llmScreenshots.map((artifact, index) => (
               <ZoomableImage
                 key={index}
-                src={`${artifactApiBaseUrl}/artifact/image?path=${uri.slice(7)}`}
+                src={getImageURL(artifact)}
                 className="object-cover w-full h-full"
                 alt="action-screenshot"
               />
@@ -150,12 +146,12 @@ function StepArtifacts({ id, stepProps }: Props) {
         )}
       </TabsContent>
       <TabsContent value="screenshot_action">
-        {actionScreenshotUris && actionScreenshotUris.length > 0 ? (
+        {actionScreenshots && actionScreenshots.length > 0 ? (
           <div className="grid grid-cols-3 gap-4 p-4">
-            {actionScreenshotUris.map((uri, index) => (
+            {actionScreenshots.map((artifact, index) => (
               <ZoomableImage
                 key={index}
-                src={`${artifactApiBaseUrl}/artifact/image?path=${uri.slice(7)}`}
+                src={getImageURL(artifact)}
                 className="object-cover w-full h-full"
                 alt="action-screenshot"
               />
@@ -172,31 +168,31 @@ function StepArtifacts({ id, stepProps }: Props) {
         )}
       </TabsContent>
       <TabsContent value="element_tree">
-        {visibleElementsTreeUri ? (
-          <JSONArtifact uri={visibleElementsTreeUri} />
+        {visibleElementsTree ? (
+          <JSONArtifact artifact={visibleElementsTree} />
         ) : null}
       </TabsContent>
       <TabsContent value="element_tree_trimmed">
-        {visibleElementsTreeTrimmedUri ? (
-          <JSONArtifact uri={visibleElementsTreeTrimmedUri} />
+        {visibleElementsTreeTrimmed ? (
+          <JSONArtifact artifact={visibleElementsTreeTrimmed} />
         ) : null}
       </TabsContent>
       <TabsContent value="llm_prompt">
-        {llmPromptUri ? <TextArtifact uri={llmPromptUri} /> : null}
+        {llmPrompt ? <TextArtifact artifact={llmPrompt} /> : null}
       </TabsContent>
       <TabsContent value="llm_request">
-        {llmRequestUri ? <JSONArtifact uri={llmRequestUri} /> : null}
+        {llmRequest ? <JSONArtifact artifact={llmRequest} /> : null}
       </TabsContent>
       <TabsContent value="llm_response_raw">
-        {llmResponseRawUri ? <JSONArtifact uri={llmResponseRawUri} /> : null}
+        {llmResponseRaw ? <JSONArtifact artifact={llmResponseRaw} /> : null}
       </TabsContent>
       <TabsContent value="llm_response_parsed">
-        {llmResponseParsedUri ? (
-          <JSONArtifact uri={llmResponseParsedUri} />
+        {llmResponseParsed ? (
+          <JSONArtifact artifact={llmResponseParsed} />
         ) : null}
       </TabsContent>
       <TabsContent value="html_raw">
-        {htmlRawUri ? <TextArtifact uri={htmlRawUri} /> : null}
+        {htmlRaw ? <TextArtifact artifact={htmlRaw} /> : null}
       </TabsContent>
     </Tabs>
   );
diff --git a/skyvern-frontend/src/routes/tasks/detail/TaskDetails.tsx b/skyvern-frontend/src/routes/tasks/detail/TaskDetails.tsx
index b78519eb..feb57f6c 100644
--- a/skyvern-frontend/src/routes/tasks/detail/TaskDetails.tsx
+++ b/skyvern-frontend/src/routes/tasks/detail/TaskDetails.tsx
@@ -11,13 +11,13 @@ import {
   AccordionTrigger,
 } from "@/components/ui/accordion";
 import { StatusBadge } from "@/components/StatusBadge";
-import { artifactApiBaseUrl } from "@/util/env";
 import { Button } from "@/components/ui/button";
 import { ReloadIcon } from "@radix-ui/react-icons";
 import { basicTimeFormat } from "@/util/timeFormat";
 import { StepArtifactsLayout } from "./StepArtifactsLayout";
 import Zoom from "react-medium-image-zoom";
 import { AspectRatio } from "@/components/ui/aspect-ratio";
+import { getRecordingURL, getScreenshotURL } from "./artifactUtils";
 
 function TaskDetails() {
   const { taskId } = useParams();
@@ -63,10 +63,7 @@ function TaskDetails() {
         {task.recording_url ? (
           <div className="flex">
             <Label className="w-32">Recording</Label>
-            <video
-              src={`${artifactApiBaseUrl}/artifact/recording?path=${task.recording_url.slice(7)}`}
-              controls
-            />
+            <video src={getRecordingURL(task)} controls />
           </div>
         ) : null}
         <div className="flex items-center">
@@ -142,10 +139,7 @@ function TaskDetails() {
               {task.screenshot_url ? (
                 <Zoom zoomMargin={16}>
                   <AspectRatio ratio={16 / 9}>
-                    <img
-                      src={`${artifactApiBaseUrl}/artifact/image?path=${task.screenshot_url.slice(7)}`}
-                      alt="screenshot"
-                    />
+                    <img src={getScreenshotURL(task)} alt="screenshot" />
                   </AspectRatio>
                 </Zoom>
               ) : (
diff --git a/skyvern-frontend/src/routes/tasks/detail/TextArtifact.tsx b/skyvern-frontend/src/routes/tasks/detail/TextArtifact.tsx
index f65764df..8906417c 100644
--- a/skyvern-frontend/src/routes/tasks/detail/TextArtifact.tsx
+++ b/skyvern-frontend/src/routes/tasks/detail/TextArtifact.tsx
@@ -1,23 +1,30 @@
 import { artifactApiClient } from "@/api/AxiosClient";
+import { ArtifactApiResponse } from "@/api/types";
 import { Skeleton } from "@/components/ui/skeleton";
 import { Textarea } from "@/components/ui/textarea";
 import { useQuery } from "@tanstack/react-query";
+import axios from "axios";
 
 type Props = {
-  uri: string;
+  artifact: ArtifactApiResponse;
 };
 
-function TextArtifact({ uri }: Props) {
+function TextArtifact({ artifact }: Props) {
   const { data, isFetching, isError, error } = useQuery<string>({
-    queryKey: ["artifact", uri],
+    queryKey: ["artifact", artifact.artifact_id],
     queryFn: async () => {
-      return artifactApiClient
-        .get(`/artifact/text`, {
-          params: {
-            path: uri.slice(7),
-          },
-        })
-        .then((response) => response.data);
+      if (artifact.uri.startsWith("file://")) {
+        return artifactApiClient
+          .get(`/artifact/text`, {
+            params: {
+              path: artifact.uri.slice(7),
+            },
+          })
+          .then((response) => response.data);
+      }
+      if (artifact.uri.startsWith("s3://") && artifact.signed_url) {
+        return axios.get(artifact.signed_url).then((response) => response.data);
+      }
     },
   });
 
diff --git a/skyvern-frontend/src/routes/tasks/detail/artifactUtils.ts b/skyvern-frontend/src/routes/tasks/detail/artifactUtils.ts
new file mode 100644
index 00000000..4e8e9b9d
--- /dev/null
+++ b/skyvern-frontend/src/routes/tasks/detail/artifactUtils.ts
@@ -0,0 +1,31 @@
+import { ArtifactApiResponse, TaskApiResponse } from "@/api/types";
+import { artifactApiBaseUrl } from "@/util/env";
+
+export function getImageURL(artifact: ArtifactApiResponse): string {
+  if (artifact.uri.startsWith("file://")) {
+    return `${artifactApiBaseUrl}/artifact/image?path=${artifact.uri.slice(7)}`;
+  } else if (artifact.uri.startsWith("s3://") && artifact.signed_url) {
+    return artifact.signed_url;
+  }
+  return artifact.uri;
+}
+
+export function getScreenshotURL(task: TaskApiResponse) {
+  if (!task.screenshot_url) {
+    return;
+  }
+  if (task.screenshot_url?.startsWith("file://")) {
+    return `${artifactApiBaseUrl}/artifact/image?path=${task.screenshot_url.slice(7)}`;
+  }
+  return task.screenshot_url;
+}
+
+export function getRecordingURL(task: TaskApiResponse) {
+  if (!task.recording_url) {
+    return;
+  }
+  if (task.recording_url?.startsWith("file://")) {
+    return `${artifactApiBaseUrl}/artifact/recording?path=${task.recording_url.slice(7)}`;
+  }
+  return task.recording_url;
+}