From e056c491a16cfe62ac148cc719104d4190e2649e Mon Sep 17 00:00:00 2001
From: Kerem Yilmaz <kerem@skyvern.com>
Date: Tue, 25 Jun 2024 12:39:25 -0700
Subject: [PATCH] Salih/post streaming cleanup (#513)

---
 .../src/routes/tasks/detail/TaskActions.tsx   |  11 +-
 .../src/routes/tasks/detail/TaskStream.tsx    | 104 ------------------
 2 files changed, 9 insertions(+), 106 deletions(-)
 delete mode 100644 skyvern-frontend/src/routes/tasks/detail/TaskStream.tsx

diff --git a/skyvern-frontend/src/routes/tasks/detail/TaskActions.tsx b/skyvern-frontend/src/routes/tasks/detail/TaskActions.tsx
index 0604d39d..17b70136 100644
--- a/skyvern-frontend/src/routes/tasks/detail/TaskActions.tsx
+++ b/skyvern-frontend/src/routes/tasks/detail/TaskActions.tsx
@@ -14,6 +14,7 @@ import { getClient } from "@/api/AxiosClient";
 import { useCredentialGetter } from "@/hooks/useCredentialGetter";
 import { Skeleton } from "@/components/ui/skeleton";
 import { toast } from "@/components/ui/use-toast";
+import { envCredential } from "@/util/env";
 
 type StreamMessage = {
   task_id: string;
@@ -70,12 +71,18 @@ function TaskActions() {
 
     async function run() {
       // Create WebSocket connection.
-      const credential = await credentialGetter!();
+      let credential = null;
+      if (credentialGetter) {
+        const token = await credentialGetter();
+        credential = `?token=Bearer ${token}`;
+      } else {
+        credential = `?apikey=${envCredential}`;
+      }
       if (socket) {
         socket.close();
       }
       socket = new WebSocket(
-        `${wssBaseUrl}/stream/tasks/${taskId}?token=Bearer ${credential}`,
+        `${wssBaseUrl}/stream/tasks/${taskId}${credential}`,
       );
       // Listen for messages
       socket.addEventListener("message", (event) => {
diff --git a/skyvern-frontend/src/routes/tasks/detail/TaskStream.tsx b/skyvern-frontend/src/routes/tasks/detail/TaskStream.tsx
deleted file mode 100644
index 1341d42a..00000000
--- a/skyvern-frontend/src/routes/tasks/detail/TaskStream.tsx
+++ /dev/null
@@ -1,104 +0,0 @@
-import { Status } from "@/api/types";
-import { useCredentialGetter } from "@/hooks/useCredentialGetter";
-import { useEffect, useState } from "react";
-import { useParams } from "react-router-dom";
-
-type StreamMessage = {
-  task_id: string;
-  status: string;
-  screenshot?: string;
-};
-
-let socket: WebSocket | null = null;
-
-type Props = {
-  status: Status;
-};
-
-const wssBaseUrl = import.meta.env.VITE_WSS_BASE_URL;
-
-function TaskStream({ status }: Props) {
-  const { taskId } = useParams();
-  const credentialGetter = useCredentialGetter();
-  const [imgSrc, setImgSrc] = useState<string>("");
-
-  useEffect(() => {
-    if (!taskId || !credentialGetter) {
-      console.error("TaskStream: Task ID is required");
-      return;
-    }
-
-    async function run() {
-      // Create WebSocket connection.
-      const credential = await credentialGetter!();
-      if (socket) {
-        socket.close();
-      }
-      socket = new WebSocket(
-        `${wssBaseUrl}/stream/tasks/${taskId}?token=Bearer ${credential}`,
-      );
-
-      socket.addEventListener("open", (event) => {
-        console.log("open event", event);
-      });
-
-      // Listen for messages
-      socket.addEventListener("message", (event) => {
-        try {
-          const message: StreamMessage = JSON.parse(event.data);
-          if (message.screenshot) {
-            setImgSrc(message.screenshot);
-          }
-          if (message.status === "completed") {
-            socket?.close();
-          }
-        } catch (e) {
-          console.error("Failed to parse message", e);
-        }
-      });
-
-      socket.addEventListener("close", (event) => {
-        console.log("close event", event);
-        socket = null;
-      });
-    }
-
-    run();
-
-    return () => {
-      if (socket) {
-        socket.close();
-        socket = null;
-      }
-    };
-  }, [credentialGetter, taskId]);
-
-  if (status === Status.Queued) {
-    return (
-      <div className="w-full h-full flex flex-col gap-4 items-center justify-center text-lg bg-slate-900">
-        <span>Your task is queued. Typical queue time is 1-2 minutes.</span>
-        <span>Stream will start when the task is running.</span>
-      </div>
-    );
-  }
-
-  if (status === Status.Running && imgSrc.length === 0) {
-    return (
-      <div className="w-full h-full flex items-center justify-center text-lg bg-slate-900">
-        Starting the stream...
-      </div>
-    );
-  }
-
-  if (status === Status.Running && imgSrc.length > 0) {
-    return (
-      <div className="w-full h-full">
-        <img src={`data:image/png;base64,${imgSrc}`} />
-      </div>
-    );
-  }
-
-  return null;
-}
-
-export { TaskStream };