Close websocket when finishes

BrunoQuaresma · BrunoQuaresma · commit be252bd46949 · 2024-01-12T19:24:28.000Z
diff --git a/site/src/hooks/useWorkspaceBuildLogs.ts b/site/src/hooks/useWorkspaceBuildLogs.ts
@@ -1,23 +1,26 @@
 import { watchBuildLogsByBuildId } from "api/api";
 import { ProvisionerJobLog } from "api/typesGenerated";
 import { displayError } from "components/GlobalSnackbar/utils";
-import { useState, useEffect } from "react";
+import { useState, useEffect, useRef } from "react";
 
 export const useWorkspaceBuildLogs = (
   // buildId is optional because sometimes the build is not loaded yet
   buildId: string | undefined,
   enabled: boolean = true,
 ) => {
   const [logs, setLogs] = useState<ProvisionerJobLog[]>();
+  const socket = useRef<WebSocket>();
+
   useEffect(() => {
     if (!buildId || !enabled) {
+      socket.current?.close();
       return;
     }
 
     // Every time this hook is called reset the values
     setLogs(undefined);
 
-    const socket = watchBuildLogsByBuildId(buildId, {
+    socket.current = watchBuildLogsByBuildId(buildId, {
       // Retrieve all the logs
       after: -1,
       onMessage: (log) => {
@@ -34,9 +37,9 @@ export const useWorkspaceBuildLogs = (
     });
 
     return () => {
-      socket.close();
+      socket.current?.close();
     };
-  }, [buildId]);
+  }, [buildId, enabled]);
 
   return logs;
 };
