fix(ws): release client state when initial loading fails (#90)

backend/src/ws/server.integration.test.ts

@@ -0,0 +1,45 @@
+import assert from 'node:assert/strict';
+import http from 'node:http';
+import { once } from 'node:events';
+import test from 'node:test';
+import { WebSocket } from 'ws';
+
+// Use an EMPTY disposable database so initial-state reads deliberately fail.
+// Dedicated variable names prevent accidentally using an application database.
+const databaseUrl = process.env['WS_LIFECYCLE_TEST_DATABASE_URL'];
+const redisUrl = process.env['WS_LIFECYCLE_TEST_REDIS_URL'];
+test('failed WebSocket initial state still releases client bookkeeping', {
+  skip: !databaseUrl || !redisUrl,
+  timeout: 15_000,
+}, async () => {
+  process.env['DATABASE_URL'] = databaseUrl;
+  process.env['REDIS_URL'] = redisUrl;
+  process.env['WS_INITIAL_STATE_ENABLED'] = '1';
+  process.env['WARMUP_NETWORKS'] = '';
+  const { initWebSocketServer, closeWebSocketServer } = await import('./server.js');
+  const { closeDb } = await import('../db/index.js');
+  const { websocketClients } = await import('../metrics.js');
+  const server = http.createServer();
+  const wss = initWebSocketServer(server);
+  let client: WebSocket | undefined;
+  try {
+    server.listen(0, '127.0.0.1');
+    await once(server, 'listening');
+    const address = server.address();
+    assert.ok(address && typeof address !== 'string');
+    client = new WebSocket(`ws://127.0.0.1:${address.port}/ws?network=ukmesh`);
+    const [code] = await once(client, 'close');
+    assert.equal(code, 1013);
+    // The client can observe its close frame before the server's close event.
+    for (let attempt = 0; attempt < 50 && wss.clients.size > 0; attempt += 1) {
+      await new Promise((resolve) => setTimeout(resolve, 10));
+    }
+    assert.equal(wss.clients.size, 0);
+    assert.equal((await websocketClients.get()).values[0]?.value, 0);
+  } finally {
+    client?.terminate();
+    await closeWebSocketServer(wss);
+    await new Promise<void>((resolve) => server.close(() => resolve()));
+    await closeDb();
+  }
+});

backend/src/ws/server.ts

@@ -524,6 +524,18 @@ export function initWebSocketServer(httpServer: Server): WebSocketServer {
     };
     clientScopes.set(ws, scope);
 
+    // Install cleanup before initial-state work can reject, close the socket,
+    // or outlive a client that disconnects during the database read.
+    ws.on('close', () => {
+      clientScopes.delete(ws);
+      messageQueue.delete(ws);
+      websocketClients.set(wss.clients.size);
+      console.log('[ws] client disconnected, total:', wss.clients.size);
+    });
+    ws.on('error', (err) => {
+      console.error('[ws] client error', err.message);
+    });
+
     // Synthetic handshake probes explicitly opt out of the expensive initial
     // state. The WebSocket control-frame ping/pong still proves the full HTTP
     // upgrade and bidirectional socket path without touching PostgreSQL.
@@ -588,17 +600,6 @@ export function initWebSocketServer(httpServer: Server): WebSocketServer {
         } satisfies WSMessage));
       }
     }
-
-    ws.on('close', () => {
-      clientScopes.delete(ws);
-      messageQueue.delete(ws);
-      websocketClients.set(wss.clients.size);
-      console.log('[ws] client disconnected, total:', wss.clients.size);
-    });
-
-    ws.on('error', (err) => {
-      console.error('[ws] client error', err.message);
-    });
   });
 
   const heartbeatTimer = setInterval(() => {