fix: detect and recover from silently dropped websocket connections

apps/server/src/adapters/WebsocketAdapter.ts

@@ -38,11 +38,19 @@ import type { IAdapter } from './IAdapter.js';
 type ClientId = string;
 let instance: SocketServer | null = null;
 
+/** Timestamp of the last sign of life received from a client. */
+type LastSeen = number;
+
 class SocketServer implements IAdapter {
   private readonly MAX_PAYLOAD = 1024 * 256; // 256Kb
+  private readonly HEARTBEAT_INTERVAL = 5000;
+  private readonly HEARTBEAT_TIMEOUT = 15000;
 
   private wss: WebSocketServer | null;
   private readonly clients: Map<ClientId, Client>;
+  /** Liveness state is keyed by socket because the heartbeat iterates sockets. */
+  private readonly connections: Map<WebSocket, LastSeen>;
+  private heartbeat: NodeJS.Timeout | null = null;
   private lastConnection: Date | null = null;
   private shouldShowWelcome = true;
 
@@ -54,12 +62,14 @@ class SocketServer implements IAdapter {
     // eslint-disable-next-line @typescript-eslint/no-this-alias -- this logic is used to ensure singleton
     instance = this;
     this.clients = new Map<ClientId, Client>();
+    this.connections = new Map<WebSocket, LastSeen>();
     this.wss = null;
   }
 
   init(server: Server, showWelcome: boolean, prefix?: string) {
     this.shouldShowWelcome = showWelcome;
     this.wss = new WebSocketServer({ path: `${prefix}/ws`, server, maxPayload: this.MAX_PAYLOAD });
+    this.startHeartbeat();
 
     this.wss.on('connection', (ws, req) => {
       // Rejected sockets can emit an error while their close handshake is in progress.
@@ -94,6 +104,7 @@ class SocketServer implements IAdapter {
         origin: '',
         path: '',
       });
+      this.connections.set(ws, Date.now());
 
       this.lastConnection = new Date();
       logger.info(LogOrigin.Client, `${this.clients.size} Connections with new: ${clientName}`);
@@ -105,8 +116,13 @@ class SocketServer implements IAdapter {
       // send store payload on connect
       sendPacket(MessageTag.RuntimeData, eventStore.poll());
 
+      // Browser WebSockets reply to protocol pings automatically.
+      ws.on('pong', () => {
+        this.connections.set(ws, Date.now());
+      });
       ws.on('close', () => {
         this.clients.delete(clientId);
+        this.connections.delete(ws);
         logger.info(LogOrigin.Client, `${this.clients.size} Connections with disconnected: ${clientName}`);
         this.sendClientList();
       });
@@ -180,6 +196,38 @@ class SocketServer implements IAdapter {
     };
   }
 
+  /** Terminates connections that stop answering protocol pings. */
+  private startHeartbeat() {
+    this.stopHeartbeat();
+    this.heartbeat = setInterval(() => {
+      const now = Date.now();
+      this.wss?.clients.forEach((client) => {
+        const lastSeen = this.connections.get(client) ?? 0;
+
+        if (now - lastSeen > this.HEARTBEAT_TIMEOUT) {
+          logger.warning(LogOrigin.Client, 'Terminating unresponsive client');
+          // A non-responsive socket cannot complete a close handshake.
+          client.terminate();
+          return;
+        }
+
+        if (client.readyState === WebSocket.OPEN) {
+          client.ping();
+        }
+      });
+    }, this.HEARTBEAT_INTERVAL);
+
+    // The HTTP server, not this interval, owns process lifetime.
+    this.heartbeat.unref();
+  }
+
+  private stopHeartbeat() {
+    if (this.heartbeat) {
+      clearInterval(this.heartbeat);
+      this.heartbeat = null;
+    }
+  }
+
   private getOrCreateClient(clientId: ClientId): Client {
     if (!this.clients.has(clientId)) {
       this.clients.set(clientId, {
@@ -245,6 +293,7 @@ class SocketServer implements IAdapter {
   }
 
   shutdown(): Promise<void> {
+    this.stopHeartbeat();
     const wss = this.wss;
     if (!wss) {
       return Promise.resolve();
@@ -260,6 +309,7 @@ class SocketServer implements IAdapter {
 
       wss.close(() => {
         this.wss = null;
+        this.connections.clear();
         resolve();
       });
     });

apps/server/src/adapters/__tests__/WebsocketAdapter.test.ts

@@ -1,15 +1,19 @@
 import type { Server } from 'node:http';
 
-import { afterEach, describe, expect, it, vi } from 'vitest';
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
 
 const websocketMocks = vi.hoisted(() => {
   let connectionHandler: ((socket: FakeWebSocket, request: unknown) => void) | undefined;
+  let latestServer: FakeWebSocketServer | undefined;
 
   class FakeWebSocket {
     static readonly OPEN = 1;
 
     readyState = FakeWebSocket.OPEN;
     close = vi.fn();
+    ping = vi.fn();
+    send = vi.fn();
+    terminate = vi.fn();
     handlers = new Map<string, Array<(...args: unknown[]) => void>>();
 
     on(event: string, handler: (...args: unknown[]) => void) {
@@ -31,6 +35,10 @@ const websocketMocks = vi.hoisted(() => {
   class FakeWebSocketServer {
     clients = new Set<FakeWebSocket>();
 
+    constructor() {
+      latestServer = this;
+    }
+
     on(event: string, handler: (socket: FakeWebSocket, request: unknown) => void) {
       if (event === 'connection') {
         connectionHandler = handler;
@@ -47,9 +55,12 @@ const websocketMocks = vi.hoisted(() => {
     FakeWebSocket,
     FakeWebSocketServer,
     getConnectionHandler: () => connectionHandler,
+    getLatestServer: () => latestServer,
   };
 });
 
+const authenticationMocks = vi.hoisted(() => ({ reject: true }));
+
 vi.mock('ws', () => ({
   WebSocket: websocketMocks.FakeWebSocket,
   WebSocketServer: websocketMocks.FakeWebSocketServer,
@@ -57,15 +68,20 @@ vi.mock('ws', () => ({
 
 vi.mock('../../middleware/authenticate.js', () => ({
   authenticateSocket: (_socket: unknown, _request: unknown, next: (error?: Error) => void) => {
-    next(new Error('Unauthorized'));
+    next(authenticationMocks.reject ? new Error('Unauthorized') : undefined);
   },
 }));
 
 import { socket } from '../WebsocketAdapter.js';
 
-describe('WebsocketAdapter authentication', () => {
+describe('WebsocketAdapter', () => {
+  beforeEach(() => {
+    authenticationMocks.reject = true;
+  });
+
   afterEach(async () => {
     await socket.shutdown();
+    vi.useRealTimers();
   });
 
   it('handles an error emitted while rejecting an unauthenticated socket', () => {
@@ -79,4 +95,24 @@ describe('WebsocketAdapter authentication', () => {
     expect(rejectedSocket.close).toHaveBeenCalledWith(1008, 'Unauthorized');
     expect(() => rejectedSocket.emit('error', new Error('socket closed'))).not.toThrow();
   });
+
+  it('terminates a client that sends messages but does not answer protocol pings', () => {
+    vi.useFakeTimers();
+    authenticationMocks.reject = false;
+    socket.init({} as Server, false);
+    const client = new websocketMocks.FakeWebSocket();
+    const server = websocketMocks.getLatestServer();
+    const connectionHandler = websocketMocks.getConnectionHandler();
+
+    expect(server).toBeDefined();
+    expect(connectionHandler).toBeDefined();
+    server?.clients.add(client);
+    connectionHandler?.(client, {});
+
+    vi.advanceTimersByTime(10_000);
+    client.emit('message', Buffer.from(JSON.stringify({ tag: 'ping', payload: null })));
+
+    vi.advanceTimersByTime(10_000);
+    expect(client.terminate).toHaveBeenCalledOnce();
+  });
 });