refactor: Update WebRTC connection handling and overlays (#320)

* refactor: Update WebRTC connection handling and overlays

* fix: Update comments for WebRTC connection handling in KvmIdRoute

* chore: Clean up import statements in devices.$id.tsx

ui/src/components/Header.tsx

@@ -36,7 +36,7 @@ export default function DashboardNavbar({
   picture,
   kvmName,
 }: NavbarProps) {
-  const peerConnectionState = useRTCStore(state => state.peerConnectionState);
+  const peerConnection = useRTCStore(state => state.peerConnection);
   const setUser = useUserStore(state => state.setUser);
   const navigate = useNavigate();
   const onLogout = useCallback(async () => {
@@ -82,14 +82,14 @@ export default function DashboardNavbar({
                 <div className="hidden items-center gap-x-2 md:flex">
                   <div className="w-[159px]">
                     <PeerConnectionStatusCard
-                      state={peerConnectionState}
+                      state={peerConnection?.connectionState}
                       title={kvmName}
                     />
                   </div>
                   <div className="hidden w-[159px] md:block">
                     <USBStateStatus
                       state={usbState}
-                      peerConnectionState={peerConnectionState}
+                      peerConnectionState={peerConnection?.connectionState}
                     />
                   </div>
                 </div>

ui/src/components/USBStateStatus.tsx

@@ -30,7 +30,7 @@ export default function USBStateStatus({
   peerConnectionState,
 }: {
   state: USBStates;
-  peerConnectionState: RTCPeerConnectionState | null;
+  peerConnectionState?: RTCPeerConnectionState | null;
 }) {
   const StatusCardProps: StatusProps = {
     configured: {

ui/src/components/VideoOverlay.tsx

@@ -1,6 +1,6 @@
 import React from "react";
 import { ExclamationTriangleIcon } from "@heroicons/react/24/solid";
-import { ArrowRightIcon } from "@heroicons/react/16/solid";
+import { ArrowPathIcon, ArrowRightIcon } from "@heroicons/react/16/solid";
 import { motion, AnimatePresence } from "framer-motion";
 import { LuPlay } from "react-icons/lu";
 
@@ -25,12 +25,12 @@ interface LoadingOverlayProps {
   show: boolean;
 }
 
-export function LoadingOverlay({ show }: LoadingOverlayProps) {
+export function LoadingVideoOverlay({ show }: LoadingOverlayProps) {
   return (
     <AnimatePresence>
       {show && (
         <motion.div
-          className="absolute inset-0 aspect-video h-full w-full"
+          className="aspect-video h-full w-full"
           initial={{ opacity: 0 }}
           animate={{ opacity: 1 }}
           exit={{ opacity: 0 }}
@@ -55,21 +55,59 @@ export function LoadingOverlay({ show }: LoadingOverlayProps) {
   );
 }
 
-interface ConnectionErrorOverlayProps {
+interface LoadingConnectionOverlayProps {
   show: boolean;
+  text: string;
 }
-
-export function ConnectionErrorOverlay({ show }: ConnectionErrorOverlayProps) {
+export function LoadingConnectionOverlay({ show, text }: LoadingConnectionOverlayProps) {
   return (
     <AnimatePresence>
       {show && (
         <motion.div
-          className="absolute inset-0 z-10 aspect-video h-full w-full"
+          className="aspect-video h-full w-full"
           initial={{ opacity: 0 }}
           animate={{ opacity: 1 }}
-          exit={{ opacity: 0 }}
+          exit={{ opacity: 0, transition: { duration: 0 } }}
           transition={{
-            duration: 0.3,
+            duration: 0.4,
+            ease: "easeInOut",
+          }}
+        >
+          <OverlayContent>
+            <div className="flex flex-col items-center justify-center gap-y-1">
+              <div className="animate flex h-12 w-12 items-center justify-center">
+                <LoadingSpinner className="h-8 w-8 text-blue-800 dark:text-blue-200" />
+              </div>
+              <p className="text-center text-sm text-slate-700 dark:text-slate-300">
+                {text}
+              </p>
+            </div>
+          </OverlayContent>
+        </motion.div>
+      )}
+    </AnimatePresence>
+  );
+}
+
+interface ConnectionErrorOverlayProps {
+  show: boolean;
+  setupPeerConnection: () => Promise<void>;
+}
+
+export function ConnectionErrorOverlay({
+  show,
+  setupPeerConnection,
+}: ConnectionErrorOverlayProps) {
+  return (
+    <AnimatePresence>
+      {show && (
+        <motion.div
+          className="aspect-video h-full w-full"
+          initial={{ opacity: 0 }}
+          animate={{ opacity: 1 }}
+          exit={{ opacity: 0, transition: { duration: 0 } }}
+          transition={{
+            duration: 0.4,
             ease: "easeInOut",
           }}
         >
@@ -87,14 +125,21 @@ export function ConnectionErrorOverlay({ show }: ConnectionErrorOverlayProps) {
                       <li>Try restarting both the device and your computer</li>
                     </ul>
                   </div>
-                  <div>
+                  <div className="flex items-center gap-x-2">
                     <LinkButton
                       to={"https://jetkvm.com/docs/getting-started/troubleshooting"}
-                      theme="light"
+                      theme="primary"
                       text="Troubleshooting Guide"
                       TrailingIcon={ArrowRightIcon}
                       size="SM"
                     />
+                    <Button
+                      onClick={() => setupPeerConnection()}
+                      LeadingIcon={ArrowPathIcon}
+                      text="Try again"
+                      size="SM"
+                      theme="light"
+                    />
                   </div>
                 </div>
               </div>

ui/src/components/WebRTCVideo.tsx

@@ -19,9 +19,8 @@ import { useJsonRpc } from "@/hooks/useJsonRpc";
 
 import {
   HDMIErrorOverlay,
+  LoadingVideoOverlay,
   NoAutoplayPermissionsOverlay,
-  ConnectionErrorOverlay,
-  LoadingOverlay,
 } from "./VideoOverlay";
 
 export default function WebRTCVideo() {
@@ -46,15 +45,13 @@ export default function WebRTCVideo() {
 
   // RTC related states
   const peerConnection = useRTCStore(state => state.peerConnection);
-  const peerConnectionState = useRTCStore(state => state.peerConnectionState);
 
   // HDMI and UI states
   const hdmiState = useVideoStore(state => state.hdmiState);
   const hdmiError = ["no_lock", "no_signal", "out_of_range"].includes(hdmiState);
-  const isLoading = !hdmiError && !isPlaying;
-  const isConnectionError = ["error", "failed", "disconnected", "closed"].includes(
-    peerConnectionState || "",
-  );
+  const isVideoLoading = !isPlaying;
+
+  // console.log("peerConnection?.connectionState", peerConnection?.connectionState);
 
   // Keyboard related states
   const { setIsNumLockActive, setIsCapsLockActive, setIsScrollLockActive } =
@@ -379,25 +376,52 @@ export default function WebRTCVideo() {
     }
   }, []);
 
+  const addStreamToVideoElm = useCallback(
+    (mediaStream: MediaStream) => {
+      if (!videoElm.current) return;
+      const videoElmRefValue = videoElm.current;
+      console.log("Adding stream to video element", videoElmRefValue);
+      videoElmRefValue.srcObject = mediaStream;
+      updateVideoSizeStore(videoElmRefValue);
+    },
+    [updateVideoSizeStore],
+  );
+
+  useEffect(
+    function updateVideoStreamOnNewTrack() {
+      if (!peerConnection) return;
+      const abortController = new AbortController();
+      const signal = abortController.signal;
+
+      peerConnection.addEventListener(
+        "track",
+        (e: RTCTrackEvent) => {
+          console.log("Adding stream to video element");
+          addStreamToVideoElm(e.streams[0]);
+        },
+        { signal },
+      );
+
+      return () => {
+        abortController.abort();
+      };
+    },
+    [addStreamToVideoElm, peerConnection],
+  );
+
   useEffect(
     function updateVideoStream() {
       if (!mediaStream) return;
-      if (!videoElm.current) return;
-      if (peerConnection?.iceConnectionState !== "connected") return;
-
-      setTimeout(() => {
-        if (videoElm?.current) {
-          videoElm.current.srcObject = mediaStream;
-        }
-      }, 0);
-      updateVideoSizeStore(videoElm.current);
+      console.log("Updating video stream from mediaStream");
+      // We set the as early as possible
+      addStreamToVideoElm(mediaStream);
     },
     [
       setVideoClientSize,
-      setVideoSize,
       mediaStream,
       updateVideoSizeStore,
-      peerConnection?.iceConnectionState,
+      peerConnection,
+      addStreamToVideoElm,
     ],
   );
 
@@ -474,6 +498,8 @@ export default function WebRTCVideo() {
       const local = resetMousePosition;
       window.addEventListener("blur", local, { signal });
       document.addEventListener("visibilitychange", local, { signal });
+      const preventContextMenu = (e: MouseEvent) => e.preventDefault();
+      videoElmRefValue.addEventListener("contextmenu", preventContextMenu, { signal });
 
       return () => {
         abortController.abort();
@@ -517,17 +543,17 @@ export default function WebRTCVideo() {
   );
 
   const hasNoAutoPlayPermissions = useMemo(() => {
-    if (peerConnectionState !== "connected") return false;
+    if (peerConnection?.connectionState !== "connected") return false;
     if (isPlaying) return false;
     if (hdmiError) return false;
     if (videoHeight === 0 || videoWidth === 0) return false;
     return true;
-  }, [peerConnectionState, isPlaying, hdmiError, videoHeight, videoWidth]);
+  }, [peerConnection?.connectionState, isPlaying, hdmiError, videoHeight, videoWidth]);
 
   return (
     <div className="grid h-full w-full grid-rows-layout">
       <div className="min-h-[39.5px]">
-        <fieldset disabled={peerConnectionState !== "connected"}>
+        <fieldset disabled={peerConnection?.connectionState !== "connected"}>
           <Actionbar
             requestFullscreen={async () =>
               videoElm.current?.requestFullscreen({
@@ -575,28 +601,29 @@ export default function WebRTCVideo() {
                             "cursor-none":
                               settings.mouseMode === "absolute" &&
                               settings.isCursorHidden,
-                            "opacity-0": isLoading || isConnectionError || hdmiError,
+                            "opacity-0": isVideoLoading || hdmiError,
                             "animate-slideUpFade border border-slate-800/30 opacity-0 shadow dark:border-slate-300/20":
                               isPlaying,
                           },
                         )}
                       />
-                      <div
-                        style={{ animationDuration: "500ms" }}
-                        className="pointer-events-none absolute inset-0 flex animate-slideUpFade items-center justify-center opacity-0"
-                      >
-                        <div className="relative h-full max-h-[720px] w-full max-w-[1280px] rounded-md">
-                          <LoadingOverlay show={isLoading} />
-                          <ConnectionErrorOverlay show={isConnectionError} />
-                          <HDMIErrorOverlay show={hdmiError} hdmiState={hdmiState} />
-                          <NoAutoplayPermissionsOverlay
-                            show={hasNoAutoPlayPermissions}
-                            onPlayClick={() => {
-                              videoElm.current?.play();
-                            }}
-                          />
+                      {peerConnection?.connectionState == "connected" && (
+                        <div
+                          style={{ animationDuration: "500ms" }}
+                          className="pointer-events-none absolute inset-0 flex animate-slideUpFade items-center justify-center opacity-0"
+                        >
+                          <div className="relative h-full max-h-[720px] w-full max-w-[1280px] rounded-md">
+                            <LoadingVideoOverlay show={isVideoLoading} />
+                            <HDMIErrorOverlay show={hdmiError} hdmiState={hdmiState} />
+                            <NoAutoplayPermissionsOverlay
+                              show={hasNoAutoPlayPermissions}
+                              onPlayClick={() => {
+                                videoElm.current?.play();
+                              }}
+                            />
+                          </div>
                         </div>
-                      </div>
+                      )}
                     </div>
                   </div>
                   <VirtualKeyboard />