Jon/browser stream component (#2808)

skyvern-frontend/src/components/BrowserStream.tsx

@@ -1,21 +1,22 @@
 import { Status } from "@/api/types";
-import { useWorkflowRunQuery } from "../hooks/useWorkflowRunQuery";
 import { useEffect, useState, useRef, useCallback } from "react";
 import { HandIcon, PlayIcon } from "@radix-ui/react-icons";
 import { Button } from "@/components/ui/button";
 import { Skeleton } from "@/components/ui/skeleton";
 import { statusIsNotFinalized } from "@/routes/tasks/types";
 import { useCredentialGetter } from "@/hooks/useCredentialGetter";
-import { useParams } from "react-router-dom";
 import { envCredential } from "@/util/env";
 import { toast } from "@/components/ui/use-toast";
-import { useQueryClient } from "@tanstack/react-query";
 import RFB from "@novnc/novnc/lib/rfb.js";
 import { environment } from "@/util/env";
 import { cn } from "@/util/utils";
 import { useClientIdStore } from "@/store/useClientIdStore";
+import type {
+  TaskApiResponse,
+  WorkflowRunStatusApiResponse,
+} from "@/api/types";
 
-import "./workflow-run-stream-vnc.css";
+import "./browser-stream.css";
 
 const wssBaseUrl = import.meta.env.VITE_WSS_BASE_URL;
 
@@ -29,13 +30,38 @@ interface CommandCedeControl {
 
 type Command = CommandTakeControl | CommandCedeControl;
 
-function WorkflowRunStreamVnc() {
+type Props = {
-  const { data: workflowRun } = useWorkflowRunQuery();
+  task?: {
+    run: TaskApiResponse;
+  };
+  workflow?: {
+    run: WorkflowRunStatusApiResponse;
+  };
+  // --
+  onClose?: () => void;
+};
 
-  const { workflowRunId, workflowPermanentId } = useParams<{
+function BrowserStream({
-    workflowRunId: string;
+  task = undefined,
-    workflowPermanentId: string;
+  workflow = undefined,
-  }>();
+  // --
+  onClose,
+}: Props) {
+  let showStream: boolean = false;
+  let runId: string;
+  let entity: "task" | "workflow";
+
+  if (task) {
+    runId = task.run.task_id;
+    showStream = statusIsNotFinalized(task.run);
+    entity = "task";
+  } else if (workflow) {
+    runId = workflow.run.workflow_run_id;
+    showStream = statusIsNotFinalized(workflow.run);
+    entity = "workflow";
+  } else {
+    throw new Error("No task or workflow provided");
+  }
 
   const [commandSocket, setCommandSocket] = useState<WebSocket | null>(null);
   const [userIsControlling, setUserIsControlling] = useState<boolean>(false);
@@ -46,8 +72,8 @@ function WorkflowRunStreamVnc() {
     useState(0);
   const prevCommandConnectedRef = useRef<boolean>(false);
   const [isCommandConnected, setIsCommandConnected] = useState<boolean>(false);
-  const showStream = workflowRun && statusIsNotFinalized(workflowRun);
+  // goes up a level
-  const queryClient = useQueryClient();
+  // const queryClient = useQueryClient();
   const [canvasContainer, setCanvasContainer] = useState<HTMLDivElement | null>(
     null,
   );
@@ -78,37 +104,28 @@ function WorkflowRunStreamVnc() {
     return `${params}`;
   }, [clientId, credentialGetter]);
 
-  const invalidateQueries = useCallback(() => {
-    queryClient.invalidateQueries({
-      queryKey: ["workflowRun", workflowPermanentId, workflowRunId],
-    });
-    queryClient.invalidateQueries({ queryKey: ["workflowRuns"] });
-    queryClient.invalidateQueries({
-      queryKey: ["workflowTasks", workflowRunId],
-    });
-    queryClient.invalidateQueries({ queryKey: ["runs"] });
-  }, [queryClient, workflowPermanentId, workflowRunId]);
-
   // effect for vnc disconnects only
   useEffect(() => {
     if (prevVncConnectedRef.current && !isVncConnected) {
       setVncDisconnectedTrigger((x) => x + 1);
+      onClose?.();
     }
     prevVncConnectedRef.current = isVncConnected;
-  }, [isVncConnected]);
+  }, [isVncConnected, onClose]);
 
   // effect for command disconnects only
   useEffect(() => {
     if (prevCommandConnectedRef.current && !isCommandConnected) {
       setCommandDisconnectedTrigger((x) => x + 1);
+      onClose?.();
     }
     prevCommandConnectedRef.current = isCommandConnected;
-  }, [isCommandConnected]);
+  }, [isCommandConnected, onClose]);
 
   // vnc socket
   useEffect(
     () => {
-      if (!showStream || !canvasContainer || !workflowRunId) {
+      if (!showStream || !canvasContainer || !runId) {
         if (rfbRef.current) {
           rfbRef.current.disconnect();
           rfbRef.current = null;
@@ -123,7 +140,16 @@ function WorkflowRunStreamVnc() {
         }
 
         const wsParams = await getWebSocketParams();
-        const vncUrl = `${wssBaseUrl}/stream/vnc/workflow_run/${workflowRunId}?${wsParams}`;
+        const vncUrl =
+          entity === "task"
+            ? `${wssBaseUrl}/stream/vnc/task/${runId}?${wsParams}`
+            : entity === "workflow"
+              ? `${wssBaseUrl}/stream/vnc/workflow_run/${runId}?${wsParams}`
+              : null;
+
+        if (!vncUrl) {
+          throw new Error("No vnc url");
+        }
 
         if (rfbRef.current) {
           rfbRef.current.disconnect();
@@ -147,7 +173,6 @@ function WorkflowRunStreamVnc() {
 
         rfb.addEventListener("disconnect", async (/* e: RfbEvent */) => {
           setIsVncConnected(false);
-          invalidateQueries();
         });
       }
 
@@ -165,20 +190,35 @@ function WorkflowRunStreamVnc() {
     // eslint-disable-next-line react-hooks/exhaustive-deps
     [
       canvasContainer,
-      invalidateQueries,
       showStream,
       vncDisconnectedTrigger, // will re-run on disconnects
-      workflowRunId,
+      runId,
+      entity,
     ],
   );
 
   // command socket
   useEffect(() => {
+    if (!showStream || !canvasContainer || !runId) {
+      return;
+    }
+
     let ws: WebSocket | null = null;
 
     const connect = async () => {
       const wsParams = await getWebSocketParams();
-      const commandUrl = `${wssBaseUrl}/stream/commands/workflow_run/${workflowRunId}?${wsParams}`;
+
+      const commandUrl =
+        entity === "task"
+          ? `${wssBaseUrl}/stream/commands/task/${runId}?${wsParams}`
+          : entity === "workflow"
+            ? `${wssBaseUrl}/stream/commands/workflow_run/${runId}?${wsParams}`
+            : null;
+
+      if (!commandUrl) {
+        throw new Error("No command url");
+      }
+
       ws = new WebSocket(commandUrl);
 
       ws.onopen = () => {
@@ -188,7 +228,6 @@ function WorkflowRunStreamVnc() {
 
       ws.onclose = () => {
         setIsCommandConnected(false);
-        invalidateQueries();
         setCommandSocket(null);
       };
     };
@@ -203,10 +242,12 @@ function WorkflowRunStreamVnc() {
       }
     };
   }, [
+    canvasContainer,
     commandDisconnectedTrigger,
+    entity,
     getWebSocketParams,
-    invalidateQueries,
+    runId,
-    workflowRunId,
+    showStream,
   ]);
 
   // effect to send a command when the user is controlling, vs not controlling
@@ -233,34 +274,41 @@ function WorkflowRunStreamVnc() {
     // eslint-disable-next-line react-hooks/exhaustive-deps
   }, [userIsControlling, isCommandConnected]);
 
-  // Effect to show toast when workflow reaches a final state based on hook updates
+  // Effect to show toast when task or workflow reaches a final state based on hook updates
   useEffect(() => {
-    if (workflowRun) {
+    const run = task ? task.run : workflow ? workflow.run : null;
-      if (
+
-        workflowRun.status === Status.Failed ||
+    if (!run) {
-        workflowRun.status === Status.Terminated
+      return;
-      ) {
-        // Only show toast if VNC is not connected or was never connected,
-        // to avoid double toasting if disconnect handler also triggers similar logic.
-        // However, the disconnect handler now primarily invalidates queries.
-        toast({
-          title: "Run Ended",
-          description: `The workflow run has ${workflowRun.status}.`,
-          variant: "destructive",
-        });
-      } else if (workflowRun.status === Status.Completed) {
-        toast({
-          title: "Run Completed",
-          description: "The workflow run has been completed.",
-          variant: "success",
-        });
-      }
     }
-  }, [workflowRun, workflowRun?.status]);
+
+    const name = task ? "task" : workflow ? "workflow" : null;
+
+    if (!name) {
+      return;
+    }
+
+    if (run.status === Status.Failed || run.status === Status.Terminated) {
+      // Only show toast if VNC is not connected or was never connected,
+      // to avoid double toasting if disconnect handler also triggers similar logic.
+      // However, the disconnect handler now primarily invalidates queries.
+      toast({
+        title: "Run Ended",
+        description: `The ${name} run has ${run.status}.`,
+        variant: "destructive",
+      });
+    } else if (run.status === Status.Completed) {
+      toast({
+        title: "Run Completed",
+        description: `The ${name} run has been completed.`,
+        variant: "success",
+      });
+    }
+  }, [task, workflow]);
 
   return (
     <div
-      className={cn("workflow-run-stream-vnc", {
+      className={cn("browser-stream", {
         "user-is-controlling": userIsControlling,
       })}
       ref={setCanvasContainerRef}
@@ -301,4 +349,4 @@ function WorkflowRunStreamVnc() {
   );
 }
 
-export { WorkflowRunStreamVnc };
+export { BrowserStream };

skyvern-frontend/src/components/browser-stream.css

@@ -1,4 +1,4 @@
-.workflow-run-stream-vnc {
+.browser-stream {
   position: relative;
   width: 100%;
   height: 100%;
@@ -9,11 +9,11 @@
   transition: padding 0.2s ease-in-out;
 }
 
-.workflow-run-stream-vnc.user-is-controlling {
+.browser-stream.user-is-controlling {
   padding: 0rem;
 }
 
-.workflow-run-stream-vnc .overlay-container {
+.browser-stream .overlay-container {
   position: absolute;
   top: 0;
   left: 0;
@@ -24,7 +24,7 @@
   justify-content: center;
 }
 
-.workflow-run-stream-vnc .overlay {
+.browser-stream .overlay {
   position: relative;
   height: auto;
   width: 100%;
@@ -36,15 +36,15 @@
   justify-content: center;
 }
 
-.workflow-run-stream-vnc.user-is-controlling .overlay {
+.browser-stream.user-is-controlling .overlay {
   pointer-events: none;
 }
 
-.workflow-run-stream-vnc.user-is-controlling .overlay-container {
+.browser-stream.user-is-controlling .overlay-container {
   pointer-events: none;
 }
 
-.workflow-run-stream-vnc .take-control {
+.browser-stream .take-control {
   transform: translateY(0);
   transition:
     transform 0.2s ease-in-out,
@@ -52,17 +52,17 @@
   opacity: 0.3;
 }
 
-.workflow-run-stream-vnc .take-control:not(.hide):hover {
+.browser-stream .take-control:not(.hide):hover {
   opacity: 1;
 }
 
-.workflow-run-stream-vnc .take-control.hide {
+.browser-stream .take-control.hide {
   transform: translateY(100%);
   opacity: 0;
   pointer-events: none;
 }
 
-.workflow-run-stream-vnc .relinquish-control {
+.browser-stream .relinquish-control {
   transform: translateY(0);
   transition:
     transform 0.2s ease-in-out,
@@ -71,17 +71,17 @@
   pointer-events: all;
 }
 
-.workflow-run-stream-vnc .relinquish-control:not(.hide):hover {
+.browser-stream .relinquish-control:not(.hide):hover {
   opacity: 1;
 }
 
-.workflow-run-stream-vnc .relinquish-control.hide {
+.browser-stream .relinquish-control.hide {
   transform: translateY(100%);
   opacity: 0;
   pointer-events: none;
 }
 
-.workflow-run-stream-vnc > div > canvas {
+.browser-stream > div > canvas {
   opacity: 0;
   animation: skyvern-anim-fadeIn 1s ease-in forwards;
 }

skyvern-frontend/src/routes/workflows/workflowRun/WorkflowRunOverview.tsx

@@ -1,8 +1,10 @@
 import { ActionsApiResponse } from "@/api/types";
+import { BrowserStream } from "@/components/BrowserStream";
 import { AspectRatio } from "@/components/ui/aspect-ratio";
 import { ActionScreenshot } from "@/routes/tasks/detail/ActionScreenshot";
 import { statusIsFinalized } from "@/routes/tasks/types";
 import { useWorkflowRunQuery } from "../hooks/useWorkflowRunQuery";
+import { useParams } from "react-router-dom";
 import { useWorkflowRunTimelineQuery } from "../hooks/useWorkflowRunTimelineQuery";
 import {
   isAction,
@@ -14,10 +16,11 @@ import {
 import { ObserverThoughtScreenshot } from "./ObserverThoughtScreenshot";
 import { WorkflowRunBlockScreenshot } from "./WorkflowRunBlockScreenshot";
 import { WorkflowRunStream } from "./WorkflowRunStream";
-import { WorkflowRunStreamVnc } from "./WorkflowRunStreamVnc";
 import { useSearchParams } from "react-router-dom";
 import { findActiveItem } from "./workflowTimelineUtils";
 import { Skeleton } from "@/components/ui/skeleton";
+import { useQueryClient } from "@tanstack/react-query";
+import { useCallback } from "react";
 
 export type ActionItem = {
   block: WorkflowRunBlock;
@@ -34,12 +37,33 @@ export type WorkflowRunOverviewActiveElement =
 function WorkflowRunOverview() {
   const [searchParams] = useSearchParams();
   const active = searchParams.get("active");
+  const { workflowPermanentId } = useParams<{
+    workflowPermanentId: string;
+  }>();
+  const queryClient = useQueryClient();
   const { data: workflowRun, isLoading: workflowRunIsLoading } =
     useWorkflowRunQuery();
 
   const { data: workflowRunTimeline, isLoading: workflowRunTimelineIsLoading } =
     useWorkflowRunTimelineQuery();
 
+  const invalidateQueries = useCallback(() => {
+    if (workflowRun) {
+      queryClient.invalidateQueries({
+        queryKey: [
+          "workflowRun",
+          workflowPermanentId,
+          workflowRun.workflow_run_id,
+        ],
+      });
+      queryClient.invalidateQueries({ queryKey: ["workflowRuns"] });
+      queryClient.invalidateQueries({
+        queryKey: ["workflowTasks", workflowRun.workflow_run_id],
+      });
+      queryClient.invalidateQueries({ queryKey: ["runs"] });
+    }
+  }, [queryClient, workflowPermanentId, workflowRun]);
+
   if (workflowRunIsLoading || workflowRunTimelineIsLoading) {
     return (
       <AspectRatio ratio={16 / 9}>
@@ -64,7 +88,10 @@ function WorkflowRunOverview() {
   );
 
   const streamingComponent = workflowRun.browser_session_id ? (
-    <WorkflowRunStreamVnc />
+    <BrowserStream
+      workflow={{ run: workflowRun }}
+      onClose={() => invalidateQueries()}
+    />
   ) : (
     <WorkflowRunStream />
   );