fix(site): Fix websocket connections (#7187)

site/src/api/api.ts

@@ -1106,25 +1106,107 @@ export const watchAgentMetadata = (agentId: string): EventSource => {
   )
 }
 
-export const watchBuildLogs = (
+type WatchBuildLogsByTemplateVersionIdOptions = {
+  after?: number
+  onMessage: (log: TypesGen.ProvisionerJobLog) => void
+  onDone: () => void
+  onError: (error: Error) => void
+}
+export const watchBuildLogsByTemplateVersionId = (
   versionId: string,
-  onMessage: (log: TypesGen.ProvisionerJobLog) => void,
+  {
+    onMessage,
+    onDone,
+    onError,
+    after,
+  }: WatchBuildLogsByTemplateVersionIdOptions,
 ) => {
-  return new Promise<void>((resolve, reject) => {
-    const proto = location.protocol === "https:" ? "wss:" : "ws:"
-    const socket = new WebSocket(
-      `${proto}//${location.host}/api/v2/templateversions/${versionId}/logs?follow=true`,
-    )
-    socket.binaryType = "blob"
-    socket.addEventListener("message", (event) =>
-      onMessage(JSON.parse(event.data) as TypesGen.ProvisionerJobLog),
-    )
-    socket.addEventListener("error", () => {
-      reject(new Error("Connection for logs failed."))
-    })
-    socket.addEventListener("close", () => {
-      // When the socket closes, logs have finished streaming!
-      resolve()
-    })
+  const searchParams = new URLSearchParams({ follow: "true" })
+  if (after !== undefined) {
+    searchParams.append("after", after.toString())
+  }
+  const proto = location.protocol === "https:" ? "wss:" : "ws:"
+  const socket = new WebSocket(
+    `${proto}//${
+      location.host
+    }/api/v2/templateversions/${versionId}/logs?${searchParams.toString()}`,
+  )
+  socket.binaryType = "blob"
+  socket.addEventListener("message", (event) =>
+    onMessage(JSON.parse(event.data) as TypesGen.ProvisionerJobLog),
+  )
+  socket.addEventListener("error", () => {
+    onError(new Error("Connection for logs failed."))
+    socket.close()
   })
+  socket.addEventListener("close", () => {
+    // When the socket closes, logs have finished streaming!
+    onDone()
+  })
+  return socket
+}
+
+type WatchStartupLogsOptions = {
+  after: number
+  onMessage: (logs: TypesGen.WorkspaceAgentStartupLog[]) => void
+  onDone: () => void
+  onError: (error: Error) => void
+}
+
+export const watchStartupLogs = (
+  agentId: string,
+  { after, onMessage, onDone, onError }: WatchStartupLogsOptions,
+) => {
+  const proto = location.protocol === "https:" ? "wss:" : "ws:"
+  const socket = new WebSocket(
+    `${proto}//${location.host}/api/v2/workspaceagents/${agentId}/startup-logs?follow&after=${after}`,
+  )
+  socket.binaryType = "blob"
+  socket.addEventListener("message", (event) => {
+    const logs = JSON.parse(event.data) as TypesGen.WorkspaceAgentStartupLog[]
+    onMessage(logs)
+  })
+  socket.addEventListener("error", () => {
+    onError(new Error("socket errored"))
+  })
+  socket.addEventListener("close", () => {
+    onDone()
+  })
+
+  return socket
+}
+
+type WatchBuildLogsByBuildIdOptions = {
+  after?: number
+  onMessage: (log: TypesGen.ProvisionerJobLog) => void
+  onDone: () => void
+  onError: (error: Error) => void
+}
+export const watchBuildLogsByBuildId = (
+  buildId: string,
+  { onMessage, onDone, onError, after }: WatchBuildLogsByBuildIdOptions,
+) => {
+  const searchParams = new URLSearchParams({ follow: "true" })
+  if (after !== undefined) {
+    searchParams.append("after", after.toString())
+  }
+  const proto = location.protocol === "https:" ? "wss:" : "ws:"
+  const socket = new WebSocket(
+    `${proto}//${
+      location.host
+    }/api/v2/workspacebuilds/${buildId}/logs?${searchParams.toString()}`,
+  )
+  socket.binaryType = "blob"
+  socket.addEventListener("message", (event) =>
+    onMessage(JSON.parse(event.data) as TypesGen.ProvisionerJobLog),
+  )
+  socket.addEventListener("error", () => {
+    onError(new Error("Connection for logs failed."))
+    socket.close()
+  })
+  socket.addEventListener("close", () => {
+    // When the socket closes, logs have finished streaming!
+    onDone()
+  })
+  return socket
 }

site/src/pages/TemplateVersionPage/TemplateVersionEditorPage/TemplateVersionEditorPage.test.tsx

@@ -36,10 +36,13 @@ test("Use custom name and set it as active when publishing", async () => {
   jest
     .spyOn(api, "getTemplateVersion")
     .mockResolvedValue({ ...MockTemplateVersion, id: "new-version-id" })
-  jest.spyOn(api, "watchBuildLogs").mockImplementation((_, onMessage) => {
-    onMessage(MockWorkspaceBuildLogs[0])
-    return Promise.resolve()
-  })
+  jest
+    .spyOn(api, "watchBuildLogsByTemplateVersionId")
+    .mockImplementation((_, options) => {
+      options.onMessage(MockWorkspaceBuildLogs[0])
+      options.onDone()
+      return jest.fn() as never
+    })
   const buildButton = within(topbar).getByRole("button", {
     name: "Build template",
   })
@@ -97,10 +100,13 @@ test("Do not mark as active if promote is not checked", async () => {
   jest
     .spyOn(api, "getTemplateVersion")
     .mockResolvedValue({ ...MockTemplateVersion, id: "new-version-id" })
-  jest.spyOn(api, "watchBuildLogs").mockImplementation((_, onMessage) => {
-    onMessage(MockWorkspaceBuildLogs[0])
-    return Promise.resolve()
-  })
+  jest
+    .spyOn(api, "watchBuildLogsByTemplateVersionId")
+    .mockImplementation((_, options) => {
+      options.onMessage(MockWorkspaceBuildLogs[0])
+      options.onDone()
+      return jest.fn() as never
+    })
   const buildButton = within(topbar).getByRole("button", {
     name: "Build template",
   })
@@ -153,10 +159,13 @@ test("Patch request is not send when the name is not updated", async () => {
   jest
     .spyOn(api, "getTemplateVersion")
     .mockResolvedValue({ ...MockTemplateVersion, id: "new-version-id" })
-  jest.spyOn(api, "watchBuildLogs").mockImplementation((_, onMessage) => {
-    onMessage(MockWorkspaceBuildLogs[0])
-    return Promise.resolve()
-  })
+  jest
+    .spyOn(api, "watchBuildLogsByTemplateVersionId")
+    .mockImplementation((_, options) => {
+      options.onMessage(MockWorkspaceBuildLogs[0])
+      options.onDone()
+      return jest.fn() as never
+    })
   const buildButton = within(topbar).getByRole("button", {
     name: "Build template",
   })

site/src/xServices/templateVersionEditor/templateVersionEditorXService.ts

@@ -49,6 +49,7 @@ export const templateVersionEditorMachine = createMachine(
         | { type: "SET_MISSING_VARIABLE_VALUES"; values: VariableValue[] }
         | { type: "CANCEL_MISSING_VARIABLE_VALUES" }
         | { type: "ADD_BUILD_LOG"; log: ProvisionerJobLog }
+        | { type: "BUILD_DONE" }
         | { type: "PUBLISH" }
         | ({ type: "CONFIRM_PUBLISH" } & PublishVersionData)
         | { type: "CANCEL_PUBLISH" },
@@ -157,14 +158,12 @@ export const templateVersionEditorMachine = createMachine(
         invoke: {
           id: "watchBuildLogs",
           src: "watchBuildLogs",
-          onDone: {
-            target: "fetchingVersion",
-          },
         },
         on: {
           ADD_BUILD_LOG: {
             actions: "addBuildLog",
           },
+          BUILD_DONE: "fetchingVersion",
           CANCEL_VERSION: {
             target: "cancelingBuild",
           },
@@ -360,9 +359,21 @@ export const templateVersionEditorMachine = createMachine(
             throw new Error("version must be set")
           }
 
-          return API.watchBuildLogs(version.id, (log) => {
-            callback({ type: "ADD_BUILD_LOG", log })
+          const socket = API.watchBuildLogsByTemplateVersionId(version.id, {
+            onMessage: (log) => {
+              callback({ type: "ADD_BUILD_LOG", log })
+            },
+            onDone: () => {
+              callback({ type: "BUILD_DONE" })
+            },
+            onError: (error) => {
+              console.error(error)
+            },
           })
+
+          return () => {
+            socket.close()
+          }
         },
       getResources: (ctx) => {
         if (!ctx.version) {

site/src/xServices/workspace/workspaceXService.ts

@@ -117,31 +117,32 @@ export const checks = {
 const permissionsToCheck = (
   workspace: TypesGen.Workspace,
   template: TypesGen.Template,
-) => ({
-  [checks.readWorkspace]: {
-    object: {
-      resource_type: "workspace",
-      resource_id: workspace.id,
-      owner_id: workspace.owner_id,
+) =>
+  ({
+    [checks.readWorkspace]: {
+      object: {
+        resource_type: "workspace",
+        resource_id: workspace.id,
+        owner_id: workspace.owner_id,
+      },
+      action: "read",
     },
-    action: "read",
-  },
-  [checks.updateWorkspace]: {
-    object: {
-      resource_type: "workspace",
-      resource_id: workspace.id,
-      owner_id: workspace.owner_id,
+    [checks.updateWorkspace]: {
+      object: {
+        resource_type: "workspace",
+        resource_id: workspace.id,
+        owner_id: workspace.owner_id,
+      },
+      action: "update",
     },
-    action: "update",
-  },
-  [checks.updateTemplate]: {
-    object: {
-      resource_type: "template",
-      resource_id: template.id,
+    [checks.updateTemplate]: {
+      object: {
+        resource_type: "template",
+        resource_id: template.id,
+      },
+      action: "update",
     },
-    action: "update",
-  },
-})
+  } as const)
 
 export const workspaceMachine = createMachine(
   {
@@ -851,6 +852,10 @@ export const workspaceMachine = createMachine(
         context.eventSource.onerror = () => {
           send({ type: "EVENT_SOURCE_ERROR", error: "sse error" })
         }
+
+        return () => {
+          context.eventSource?.close()
+        }
       },
       getBuilds: async (context) => {
         if (context.workspace) {

site/src/xServices/workspaceAgentLogs/workspaceAgentLogsXService.ts

@@ -22,6 +22,9 @@ export const workspaceAgentLogsMachine = createMachine(
           }
         | {
             type: "FETCH_STARTUP_LOGS"
+          }
+        | {
+            type: "STARTUP_DONE"
           },
       context: {} as {
         agentID: string
@@ -56,16 +59,19 @@ export const workspaceAgentLogsMachine = createMachine(
           id: "streamStartupLogs",
           src: "streamStartupLogs",
         },
+        on: {
+          ADD_STARTUP_LOGS: {
+            actions: "addStartupLogs",
+          },
+          STARTUP_DONE: {
+            target: "loaded",
+          },
+        },
       },
       loaded: {
         type: "final",
       },
     },
-    on: {
-      ADD_STARTUP_LOGS: {
-        actions: "addStartupLogs",
-      },
-    },
   },
   {
     services: {
@@ -79,20 +85,14 @@ export const workspaceAgentLogsMachine = createMachine(
           })),
         ),
       streamStartupLogs: (ctx) => async (callback) => {
-        return new Promise<void>((resolve, reject) => {
-          const proto = location.protocol === "https:" ? "wss:" : "ws:"
-          let after = 0
-          if (ctx.startupLogs && ctx.startupLogs.length > 0) {
-            after = ctx.startupLogs[ctx.startupLogs.length - 1].id
-          }
-          const socket = new WebSocket(
-            `${proto}//${location.host}/api/v2/workspaceagents/${ctx.agentID}/startup-logs?follow&after=${after}`,
-          )
-          socket.binaryType = "blob"
-          socket.addEventListener("message", (event) => {
-            const logs = JSON.parse(
-              event.data,
-            ) as TypesGen.WorkspaceAgentStartupLog[]
+        let after = 0
+        if (ctx.startupLogs && ctx.startupLogs.length > 0) {
+          after = ctx.startupLogs[ctx.startupLogs.length - 1].id
+        }
+
+        const socket = API.watchStartupLogs(ctx.agentID, {
+          after,
+          onMessage: (logs) => {
             callback({
               type: "ADD_STARTUP_LOGS",
               logs: logs.map((log) => ({
@@ -102,14 +102,18 @@ export const workspaceAgentLogsMachine = createMachine(
                 time: log.created_at,
               })),
             })
-          })
-          socket.addEventListener("error", () => {
-            reject(new Error("socket errored"))
-          })
-          socket.addEventListener("open", () => {
-            resolve()
-          })
+          },
+          onDone: () => {
+            callback({ type: "STARTUP_DONE" })
+          },
+          onError: (error) => {
+            console.error(error)
+          },
         })
+
+        return () => {
+          socket.close()
+        }
       },
     },
     actions: {

site/src/xServices/workspaceBuild/workspaceBuildXService.ts

@@ -16,10 +16,14 @@ type LogsContext = {
   logs?: ProvisionerJobLog[]
 }
 
-type LogsEvent = {
-  type: "ADD_LOG"
-  log: ProvisionerJobLog
-}
+type LogsEvent =
+  | {
+      type: "ADD_LOG"
+      log: ProvisionerJobLog
+    }
+  | {
+      type: "BUILD_DONE"
+    }
 
 export const workspaceBuildMachine = createMachine(
   {
@@ -74,16 +78,19 @@ export const workspaceBuildMachine = createMachine(
               id: "streamWorkspaceBuildLogs",
               src: "streamWorkspaceBuildLogs",
             },
+            on: {
+              ADD_LOG: {
+                actions: "addLog",
+              },
+              BUILD_DONE: {
+                target: "loaded",
+              },
+            },
           },
           loaded: {
             type: "final",
           },
         },
-        on: {
-          ADD_LOG: {
-            actions: "addLog",
-          },
-        },
       },
     },
   },
@@ -124,31 +131,26 @@ export const workspaceBuildMachine = createMachine(
       getLogs: async (ctx) =>
         API.getWorkspaceBuildLogs(ctx.buildId, ctx.timeCursor),
       streamWorkspaceBuildLogs: (ctx) => async (callback) => {
-        return new Promise<void>((resolve, reject) => {
-          if (!ctx.logs) {
-            return reject("logs must be set")
-          }
-          const proto = location.protocol === "https:" ? "wss:" : "ws:"
-          const socket = new WebSocket(
-            `${proto}//${location.host}/api/v2/workspacebuilds/${
-              ctx.buildId
-            }/logs?follow=true&after=${ctx.logs[ctx.logs.length - 1].id}`,
-          )
-          socket.binaryType = "blob"
-          socket.addEventListener("message", (event) => {
-            callback({ type: "ADD_LOG", log: JSON.parse(event.data) })
-          })
-          socket.addEventListener("error", () => {
-            reject(new Error("socket errored"))
-          })
-          socket.addEventListener("open", () => {
-            resolve()
-          })
-          socket.addEventListener("close", () => {
-            // When the socket closes, logs have finished streaming!
-            resolve()
-          })
+        if (!ctx.logs) {
+          throw new Error("logs must be set")
+        }
+
+        const after = ctx.logs[ctx.logs.length - 1].id
+        const socket = API.watchBuildLogsByBuildId(ctx.buildId, {
+          after,
+          onMessage: (log) => {
+            callback({ type: "ADD_LOG", log })
+          },
+          onDone: () => {
+            callback({ type: "BUILD_DONE" })
+          },
+          onError: (err) => {
+            console.error(err)
+          },
         })
+        return () => {
+          socket.close()
+        }
       },
     },
   },