Manage WebSocket reference in screenshot editor context

Author: richiemcilroy

apps/desktop/src/routes/screenshot-editor/context.tsx

@@ -134,6 +134,7 @@ function createScreenshotEditorContext() {
 	});
 
 	const [latestFrame, setLatestFrame] = createLazySignal<FrameData>();
+	const [wsRef, setWsRef] = createSignal<WebSocket | null>(null);
 
 	const [editorInstance] = createResource(async () => {
 		const instance = await commands.createScreenshotEditorInstance();
@@ -176,6 +177,7 @@ function createScreenshotEditorContext() {
 
 		const ws = new WebSocket(instance.framesSocketUrl);
 		ws.binaryType = "arraybuffer";
+		setWsRef(ws);
 		ws.onmessage = async (event) => {
 			const buffer = event.data as ArrayBuffer;
 			if (buffer.byteLength < 24) return;
@@ -239,6 +241,10 @@ function createScreenshotEditorContext() {
 		if (frame?.bitmap) {
 			frame.bitmap.close();
 		}
+		const ws = wsRef();
+		if (ws) {
+			ws.close();
+		}
 	});
 
 	const saveConfig = debounce((config: ProjectConfiguration) => {