From f049888aa7bdca6b735c568d33f8d76517263349 Mon Sep 17 00:00:00 2001
From: dab <dab@dmy.mailworks.org>
Date: Wed, 13 Aug 2025 17:12:05 +0000
Subject: [PATCH] websocket sessions now get a UID + graceful close/cleanup,
 server graceful shutdown

---
 backend/index.js  | 64 +++++++++++++++++++++++++++++++++++++++++++----
 frontend/index.js |  4 +--
 2 files changed, 61 insertions(+), 7 deletions(-)

diff --git a/backend/index.js b/backend/index.js
index 1be5eb0..8a66e5f 100644
--- a/backend/index.js
+++ b/backend/index.js
@@ -1,3 +1,5 @@
+import crypto from 'crypto'
+
 const async_build_js_script = async function (path_js_entry_script, build_options) {
 	const result = await Bun.build({
 		entrypoints: [ path_js_entry_script ],
@@ -25,7 +27,7 @@ const get_siteroot_html = function ({ page_title }) {return `
 </html>
 `.trim()}
 
-const start_server = function ({ config, jsbuild_app_frontend }) {
+const start_server = function ({ context_meta, config, jsbuild_app_frontend }) {
 	const server = Bun.serve({
 		hostname: config.host || '127.0.0.1',
 		port    : config.port || 8800,
@@ -50,8 +52,13 @@ const start_server = function ({ config, jsbuild_app_frontend }) {
 				)
 			}
 			else if (url.pathname === '/ws') {
-				if (server.upgrade(req)) {
-					// do not return a Response if upgrade fails
+				const uid = crypto.randomUUID()
+				if (server.upgrade(req, {
+					data: {
+						uid,
+					},
+				})) {
+					// do not return a Response if upgrade succeeds
 					resp = undefined
 				}
 				resp = new Response('Upgrade failed', { status: 500 })
@@ -62,20 +69,67 @@ const start_server = function ({ config, jsbuild_app_frontend }) {
 			return resp
 		},
 		websocket: {
+			open: function (ws) {
+				const map_obj = {
+					ws,
+					closed_by_server: false,
+				}
+				context_meta.ws_map.set(ws.data.uid, map_obj)
+				return
+			},
+			close: function (ws, code, message) {
+				const map_obj = context_meta.ws_map.get(ws.data.uid)
+				if (map_obj === undefined) {
+					console.warn(['Unexpected: ws_map missing uid', ws.data.uid])
+				}
+				else {
+					if (map_obj.closed_by_server === true) {
+						console.info(['ws closed by server', ws.data.uid])
+						// cleanup will happen in outer context
+					}
+					else {
+						console.info(['ws closed by client', ws.data.uid, { code, message }])
+						// cleanup from this context
+						context_meta.ws_map.delete(ws.data.uid)
+					}
+				}
+				return
+			},
 			message: function (ws, message) {},
 		},
 	})
 	return server
 }
 
+const close_active_ws_sessions = function ({ context_meta }) {
+	const keys_to_delete = []
+	for (const [ uid, map_obj ] of context_meta.ws_map) {
+		const code = 1001
+		const reason = 'Going Away'
+		map_obj.closed_by_server = true
+		map_obj.ws.close(code, reason)
+		keys_to_delete.push(uid)
+	}
+	for (let key of keys_to_delete) {
+		context_meta.ws_map.delete(key)
+	}
+	return
+}
+
 const async_run = async function ({
 	config,
 	jsbuild_app_frontend,
 }) {
-	const server = start_server({ config, jsbuild_app_frontend })
+	const context_meta = {
+		ws_map: new Map(),
+	}
+	const server = start_server({ context_meta, config, jsbuild_app_frontend })
 	console.info(server)
-	process.on('SIGINT', () => {
+	process.on('SIGINT', async () => {
 		console.log('SIGINT intercepted')
+		close_active_ws_sessions({ context_meta })
+		// NOTE: wait for a second for teardown to fully complete
+		await Bun.sleep(2000)
 		process.exit()
 	})
 	return
diff --git a/frontend/index.js b/frontend/index.js
index c3e51a7..292fc26 100644
--- a/frontend/index.js
+++ b/frontend/index.js
@@ -42,9 +42,9 @@ class FramerockUtils {
 			on_open && on_open()
 			return
 		}
-		const func_on_close = () => {
+		const func_on_close = (event) => {
 			console.info('WS CLOSE')
-			on_close && on_close()
+			on_close && on_close(event)
 			return
 		}
 		const func_on_message = (e) => {