From 5358f1f52c26f0946e51914ff312c93653827ece Mon Sep 17 00:00:00 2001
From: DecDuck <declanahofmeyr@gmail.com>
Date: Sat, 16 Nov 2024 17:27:27 +1100
Subject: [PATCH] fix: properly disconnect websockets from task handler

---
 server/api/v1/task/index.get.ts | 2 ++
 server/internal/tasks/index.ts  | 8 ++++++++
 2 files changed, 10 insertions(+)

diff --git a/server/api/v1/task/index.get.ts b/server/api/v1/task/index.get.ts
index a9e0c18..b5da1b8 100644
--- a/server/api/v1/task/index.get.ts
+++ b/server/api/v1/task/index.get.ts
@@ -48,5 +48,7 @@ export default defineWebSocketHandler({
     if (!peer.id) return;
     if (adminSocketSessions[peer.id] === undefined) return;
     delete adminSocketSessions[peer.id];
+
+    taskHandler.disconnectAll(peer.id);
   },
 });
diff --git a/server/internal/tasks/index.ts b/server/internal/tasks/index.ts
index 57bae6c..6c7b197 100644
--- a/server/internal/tasks/index.ts
+++ b/server/internal/tasks/index.ts
@@ -111,6 +111,14 @@ class TaskHandler {
     return true;
   }
 
+  disconnectAll(id: string) {
+    for (const [taskId] of Object.keys(this.taskRegistry)) {
+      delete this.taskRegistry[taskId].clients[id];
+    }
+    
+    delete this.clientRegistry[id];
+  }
+
   disconnect(id: string, taskId: string) {
     if (!this.taskRegistry[taskId]) return false;