WebSocket Command Channel (#2782)

skyvern-frontend/src/routes/workflows/editor/panels/WorkflowParametersPanel.tsx

@@ -141,9 +141,7 @@ function WorkflowParametersPanel() {
                         </span>
                       ) : (
                         <span className="text-sm text-slate-400">
-                          {parameter.parameterType === "onepassword"
-                            ? "credential"
-                            : parameter.parameterType}
+                          {parameter.parameterType}
                         </span>
                       )}
                     </div>

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

@@ -13,24 +13,40 @@ 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 "./workflow-run-stream-vnc.css";
 
 const wssBaseUrl = import.meta.env.VITE_WSS_BASE_URL;
 
+interface CommandTakeControl {
+  kind: "take-control";
+}
+
+interface CommandCedeControl {
+  kind: "cede-control";
+}
+
+type Command = CommandTakeControl | CommandCedeControl;
+
 function WorkflowRunStreamVnc() {
   const { data: workflowRun } = useWorkflowRunQuery();
+
   const { workflowRunId, workflowPermanentId } = useParams<{
     workflowRunId: string;
     workflowPermanentId: string;
   }>();
 
+  const [commandSocket, setCommandSocket] = useState<WebSocket | null>(null);
   const [userIsControlling, setUserIsControlling] = useState<boolean>(false);
   const [vncDisconnectedTrigger, setVncDisconnectedTrigger] = useState(0);
   const prevVncConnectedRef = useRef<boolean>(false);
   const [isVncConnected, setIsVncConnected] = useState<boolean>(false);
+  const [commandDisconnectedTrigger, setCommandDisconnectedTrigger] =
+    useState(0);
+  const prevCommandConnectedRef = useRef<boolean>(false);
+  const [isCommandConnected, setIsCommandConnected] = useState<boolean>(false);
   const showStream = workflowRun && statusIsNotFinalized(workflowRun);
-  const credentialGetter = useCredentialGetter();
   const queryClient = useQueryClient();
   const [canvasContainer, setCanvasContainer] = useState<HTMLDivElement | null>(
     null,
@@ -38,10 +54,42 @@ function WorkflowRunStreamVnc() {
   const setCanvasContainerRef = useCallback((node: HTMLDivElement | null) => {
     setCanvasContainer(node);
   }, []);
-
   const rfbRef = useRef<RFB | null>(null);
+  const clientId = useClientIdStore((state) => state.clientId);
+  const credentialGetter = useCredentialGetter();
 
-  // effect for disconnects only
+  const getWebSocketParams = useCallback(async () => {
+    const clientIdQueryParam = `client_id=${clientId}`;
+    let credentialQueryParam = "";
+
+    if (environment === "local") {
+      credentialQueryParam = `apikey=${envCredential}`;
+    } else {
+      if (credentialGetter) {
+        const token = await credentialGetter();
+        credentialQueryParam = `token=Bearer ${token}`;
+      } else {
+        credentialQueryParam = `apikey=${envCredential}`;
+      }
+    }
+
+    const params = [credentialQueryParam, clientIdQueryParam].join("&");
+
+    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);
@@ -49,6 +97,15 @@ function WorkflowRunStreamVnc() {
     prevVncConnectedRef.current = isVncConnected;
   }, [isVncConnected]);
 
+  // effect for command disconnects only
+  useEffect(() => {
+    if (prevCommandConnectedRef.current && !isCommandConnected) {
+      setCommandDisconnectedTrigger((x) => x + 1);
+    }
+    prevCommandConnectedRef.current = isCommandConnected;
+  }, [isCommandConnected]);
+
+  // vnc socket
   useEffect(
     () => {
       if (!showStream || !canvasContainer || !workflowRunId) {
@@ -61,24 +118,12 @@ function WorkflowRunStreamVnc() {
       }
 
       async function setupVnc() {
-        let credentialQueryParam = "";
-
-        if (environment === "local") {
-          credentialQueryParam = `?apikey=${envCredential}`;
-        } else {
-          if (credentialGetter) {
-            const token = await credentialGetter();
-            credentialQueryParam = `?token=Bearer ${token}`;
-          } else {
-            credentialQueryParam = `?apikey=${envCredential}`;
-          }
-        }
-
         if (rfbRef.current && isVncConnected) {
           return;
         }
 
-        const vncUrl = `${wssBaseUrl}/stream/vnc/workflow_run/${workflowRunId}${credentialQueryParam}`;
+        const wsParams = await getWebSocketParams();
+        const vncUrl = `${wssBaseUrl}/stream/vnc/workflow_run/${workflowRunId}?${wsParams}`;
 
         if (rfbRef.current) {
           rfbRef.current.disconnect();
@@ -102,15 +147,7 @@ function WorkflowRunStreamVnc() {
 
         rfb.addEventListener("disconnect", async (/* e: RfbEvent */) => {
           setIsVncConnected(false);
-
-          queryClient.invalidateQueries({
-            queryKey: ["workflowRun", workflowPermanentId, workflowRunId],
-          });
-          queryClient.invalidateQueries({ queryKey: ["workflowRuns"] });
-          queryClient.invalidateQueries({
-            queryKey: ["workflowTasks", workflowRunId],
-          });
-          queryClient.invalidateQueries({ queryKey: ["runs"] });
+          invalidateQueries();
         });
       }
 
@@ -127,16 +164,75 @@ function WorkflowRunStreamVnc() {
     // cannot include isVncConnected in deps as it will cause infinite loop
     // eslint-disable-next-line react-hooks/exhaustive-deps
     [
-      credentialGetter,
-      workflowRunId,
-      workflowPermanentId,
-      showStream,
-      queryClient,
       canvasContainer,
+      invalidateQueries,
+      showStream,
       vncDisconnectedTrigger, // will re-run on disconnects
+      workflowRunId,
     ],
   );
 
+  // command socket
+  useEffect(() => {
+    let ws: WebSocket | null = null;
+
+    const connect = async () => {
+      const wsParams = await getWebSocketParams();
+      const commandUrl = `${wssBaseUrl}/stream/commands/workflow_run/${workflowRunId}?${wsParams}`;
+      ws = new WebSocket(commandUrl);
+
+      ws.onopen = () => {
+        setIsCommandConnected(true);
+        setCommandSocket(ws);
+      };
+
+      ws.onclose = () => {
+        setIsCommandConnected(false);
+        invalidateQueries();
+        setCommandSocket(null);
+      };
+    };
+
+    connect();
+
+    return () => {
+      try {
+        ws && ws.close();
+      } catch (e) {
+        // pass
+      }
+    };
+  }, [
+    commandDisconnectedTrigger,
+    getWebSocketParams,
+    invalidateQueries,
+    workflowRunId,
+  ]);
+
+  // effect to send a command when the user is controlling, vs not controlling
+  useEffect(() => {
+    if (!isCommandConnected) {
+      return;
+    }
+
+    const sendCommand = (command: Command) => {
+      if (!commandSocket) {
+        console.warn("Cannot send command, as command socket is closed.");
+        console.warn(command);
+        return;
+      }
+
+      commandSocket.send(JSON.stringify(command));
+    };
+
+    if (userIsControlling) {
+      sendCommand({ kind: "take-control" });
+    } else {
+      sendCommand({ kind: "cede-control" });
+    }
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, [userIsControlling, isCommandConnected]);
+
   // Effect to show toast when workflow reaches a final state based on hook updates
   useEffect(() => {
     if (workflowRun) {

skyvern-frontend/src/store/useClientIdStore.ts

@@ -0,0 +1,11 @@
+import { create } from "zustand";
+
+type ClientIdStore = {
+  clientId: string;
+};
+
+const initialClientId = crypto.randomUUID();
+
+export const useClientIdStore = create<ClientIdStore>(() => ({
+  clientId: initialClientId,
+}));