From 62ce1c5267f59ed2e15384545b216c2bd696e006 Mon Sep 17 00:00:00 2001
From: nntrivi2001 <vinnt@smartgift.vn>
Date: Wed, 19 Aug 2026 15:01:00 +0700
Subject: [PATCH] fix(ws): drop upgrades arriving after closeWebSocket instead
 of crashing
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

closeWebSocket() nulls both WebSocketServer references but cannot unregister
the `server.on("upgrade")` listeners that initWebSocket/initPtyWebSocket
installed on the shared http.Server. An upgrade landing in that window threw

    TypeError: Cannot read properties of null (reading 'handleUpgrade')

from the listener — unhandled, so it killed the process mid-SIGTERM instead of
letting it exit gracefully, and no server came back up. Observed in
runtime/server.log right after a restart, with the dashboard then showing
"Mất kết nối" and nothing else.

The previous commit's TerminalView reconnect makes this near-certain rather
than rare: every open run console re-attaches to /ws-pty every 1.5s, so a
shutdown almost always has an upgrade in flight. Both listeners now destroy
the socket when their server is gone; clients retry, which is the correct
answer during a shutdown.

Regression test uses a real http.Server and raw upgrade requests — without the
guard the process dies and the test run hangs rather than reporting a failure.
---
 server/__tests__/websocket-shutdown.test.js | 56 +++++++++++++++++++++
 server/websocket.js                         | 14 ++++++
 2 files changed, 70 insertions(+)
 create mode 100644 server/__tests__/websocket-shutdown.test.js

diff --git a/server/__tests__/websocket-shutdown.test.js b/server/__tests__/websocket-shutdown.test.js
new file mode 100644
index 0000000..5b62d20
--- /dev/null
+++ b/server/__tests__/websocket-shutdown.test.js
@@ -0,0 +1,56 @@
+/**
+ * @file websocket-shutdown.test.js
+ * @description Regression test for the shutdown race in server/websocket.js:
+ * `closeWebSocket()` nulls both WebSocketServer references but cannot
+ * unregister the `server.on("upgrade")` listeners it installed, so an upgrade
+ * arriving mid-shutdown used to throw an unhandled TypeError
+ * ("Cannot read properties of null (reading 'handleUpgrade')") and kill the
+ * process. Uses a real http.Server and raw upgrade requests — no ws client.
+ * @author Nguyễn Ngọc Trí Vĩ <vinnt@smartgift.vn>
+ */
+const { describe, it } = require("node:test");
+const assert = require("node:assert/strict");
+const http = require("node:http");
+const net = require("node:net");
+const ws = require("../websocket");
+
+/** Send a raw WS upgrade for `path` and resolve once the socket settles. */
+function rawUpgrade(port, path) {
+  return new Promise((resolve) => {
+    const socket = net.connect(port, "127.0.0.1", () => {
+      socket.write(
+        `GET ${path} HTTP/1.1\r\nHost: 127.0.0.1:${port}\r\n` +
+          `Upgrade: websocket\r\nConnection: Upgrade\r\n` +
+          `Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==\r\nSec-WebSocket-Version: 13\r\n\r\n`
+      );
+    });
+    let body = "";
+    socket.on("data", (d) => {
+      body += d.toString("utf8");
+    });
+    socket.on("close", () => resolve(body));
+    socket.on("error", () => resolve(body));
+  });
+}
+
+describe("websocket shutdown race", () => {
+  it("drops upgrades that land after closeWebSocket instead of crashing", async () => {
+    const server = http.createServer((req, res) => res.end("ok"));
+    ws.initWebSocket(server);
+    ws.initPtyWebSocket(server);
+    await new Promise((r) => server.listen(0, "127.0.0.1", r));
+    const { port } = server.address();
+
+    // The shutdown the real server performs on SIGTERM.
+    ws.closeWebSocket();
+
+    // Both paths, because each has its own upgrade listener and its own null
+    // reference. An unguarded handleUpgrade here takes the process down, so
+    // simply reaching the assertions below is the test passing.
+    await rawUpgrade(port, "/ws");
+    await rawUpgrade(port, "/ws-pty/ccam-lane-1");
+
+    assert.equal(ws.getConnectionCount(), 0);
+    await new Promise((r) => server.close(r));
+  });
+});
diff --git a/server/websocket.js b/server/websocket.js
index b4cffd9..fdf99d8 100644
--- a/server/websocket.js
+++ b/server/websocket.js
@@ -29,6 +29,14 @@ function initWebSocket(server) {
   server.on("upgrade", (req, socket, head) => {
     const url = new URL(req.url, "http://localhost");
     if (url.pathname !== "/ws") return; // not ours — `/ws-pty/*` handles its own path.
+    // closeWebSocket() nulls `wss` but cannot unregister this listener, so an
+    // upgrade arriving mid-shutdown would throw an unhandled TypeError and take
+    // the process down instead of letting it exit gracefully. Clients retry, so
+    // dropping the socket is the correct answer here.
+    if (!wss) {
+      socket.destroy();
+      return;
+    }
     wss.handleUpgrade(req, socket, head, (ws) => {
       wss.emit("connection", ws, req);
     });
@@ -118,6 +126,12 @@ function initPtyWebSocket(server) {
       socket.destroy();
       return;
     }
+    // Same shutdown race as `/ws` above — and far easier to hit here, since
+    // TerminalView re-attaches every 1.5s while a run console is open.
+    if (!ptyWss) {
+      socket.destroy();
+      return;
+    }
     ptyWss.handleUpgrade(req, socket, head, (ws) => {
       ptyWss.emit("connection", ws, match[1]);
     });