From eef5bf9fc9fc54fd1b56c6070fd5f38aca47aec3 Mon Sep 17 00:00:00 2001
From: Minh Vu <vuhoangminh97@gmail.com>
Date: Thu, 10 Sep 2026 23:10:23 +0200
Subject: [PATCH] fix(web): stop reconnecting after explicit disconnect

---
 web/src/api/websocket.test.ts | 108 ++++++++++++++++++++++++++++++++++
 web/src/api/websocket.ts      |  30 +++++++---
 2 files changed, 129 insertions(+), 9 deletions(-)
 create mode 100644 web/src/api/websocket.test.ts

diff --git a/web/src/api/websocket.test.ts b/web/src/api/websocket.test.ts
new file mode 100644
index 00000000..6820495a
--- /dev/null
+++ b/web/src/api/websocket.test.ts
@@ -0,0 +1,108 @@
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
+
+vi.mock('./client', () => ({ getToken: () => null }));
+
+class MockWebSocket {
+  static readonly CONNECTING = 0;
+  static readonly OPEN = 1;
+  static readonly CLOSING = 2;
+  static readonly CLOSED = 3;
+  static instances: MockWebSocket[] = [];
+
+  readonly url: string;
+  readyState = MockWebSocket.CONNECTING;
+  onopen: (() => void) | null = null;
+  onmessage: ((event: { data: string }) => void) | null = null;
+  onclose: (() => void) | null = null;
+  onerror: (() => void) | null = null;
+  readonly send = vi.fn();
+  readonly close = vi.fn(() => {
+    this.readyState = MockWebSocket.CLOSING;
+  });
+
+  constructor(url: string) {
+    this.url = url;
+    MockWebSocket.instances.push(this);
+  }
+
+  open() {
+    this.readyState = MockWebSocket.OPEN;
+    this.onopen?.();
+  }
+
+  closeUnexpectedly() {
+    this.readyState = MockWebSocket.CLOSED;
+    this.onclose?.();
+  }
+}
+
+const { NerveWebSocket } = await import('./websocket');
+
+describe('NerveWebSocket', () => {
+  beforeEach(() => {
+    MockWebSocket.instances = [];
+    vi.stubGlobal('WebSocket', MockWebSocket);
+    vi.useFakeTimers();
+  });
+
+  afterEach(() => {
+    vi.useRealTimers();
+    vi.unstubAllGlobals();
+  });
+
+  it('does not reconnect after an explicit disconnect', () => {
+    const client = new NerveWebSocket();
+    client.connect();
+    const socket = MockWebSocket.instances[0];
+    socket.open();
+
+    client.disconnect();
+    socket.closeUnexpectedly();
+    vi.advanceTimersByTime(3000);
+
+    expect(socket.close).toHaveBeenCalledOnce();
+    expect(MockWebSocket.instances).toHaveLength(1);
+    expect(client.connected).toBe(false);
+  });
+
+  it('drops messages after cancelling a scheduled reconnect', () => {
+    const client = new NerveWebSocket();
+    client.connect();
+    const socket = MockWebSocket.instances[0];
+    socket.open();
+    socket.closeUnexpectedly();
+    expect(vi.getTimerCount()).toBe(1);
+
+    client.disconnect();
+
+    expect(client.send({ type: 'message' })).toBe('dropped');
+    vi.advanceTimersByTime(3000);
+    expect(MockWebSocket.instances).toHaveLength(1);
+  });
+
+  it('allows reconnecting after a later explicit connect', () => {
+    const client = new NerveWebSocket();
+    client.connect();
+    client.disconnect();
+
+    client.connect();
+    const socket = MockWebSocket.instances[1];
+    socket.closeUnexpectedly();
+    vi.advanceTimersByTime(3000);
+
+    expect(MockWebSocket.instances).toHaveLength(3);
+  });
+
+  it('ignores close events from a detached socket', () => {
+    const client = new NerveWebSocket();
+    client.connect();
+    const firstSocket = MockWebSocket.instances[0];
+    client.disconnect();
+
+    client.connect();
+    firstSocket.closeUnexpectedly();
+    vi.advanceTimersByTime(3000);
+
+    expect(MockWebSocket.instances).toHaveLength(2);
+  });
+});
diff --git a/web/src/api/websocket.ts b/web/src/api/websocket.ts
index 489a29a7..5d684b28 100644
--- a/web/src/api/websocket.ts
+++ b/web/src/api/websocket.ts
@@ -73,12 +73,14 @@ export class NerveWebSocket {
   private pingInterval: ReturnType<typeof setInterval> | null = null;
   private _connected = false;
   private _pending: Record<string, unknown>[] = [];
+  private shouldReconnect = true;
 
   get connected() {
     return this._connected;
   }
 
   connect() {
+    this.shouldReconnect = true;
     if (this.ws?.readyState === WebSocket.OPEN) return;
 
     const token = getToken();
@@ -86,15 +88,18 @@ export class NerveWebSocket {
     const host = window.location.host;
     const url = `${protocol}//${host}/ws${token ? `?token=${token}` : ''}`;
 
-    this.ws = new WebSocket(url);
+    const socket = new WebSocket(url);
+    this.ws = socket;
 
-    this.ws.onopen = () => {
+    socket.onopen = () => {
+      if (this.ws !== socket) return;
       this._connected = true;
       this.startPing();
       this.flushPending();
     };
 
-    this.ws.onmessage = (event) => {
+    socket.onmessage = (event) => {
+      if (this.ws !== socket) return;
       try {
         const msg: WSMessage = JSON.parse(event.data);
         this.handlers.forEach((h) => h(msg));
@@ -103,22 +108,29 @@ export class NerveWebSocket {
       }
     };
 
-    this.ws.onclose = () => {
+    socket.onclose = () => {
+      if (this.ws !== socket) return;
       this._connected = false;
       this.stopPing();
-      this.scheduleReconnect();
+      if (this.shouldReconnect) this.scheduleReconnect();
     };
 
-    this.ws.onerror = () => {
+    socket.onerror = () => {
+      if (this.ws !== socket) return;
       this._connected = false;
     };
   }
 
   disconnect() {
-    if (this.reconnectTimer) clearTimeout(this.reconnectTimer);
+    this.shouldReconnect = false;
+    if (this.reconnectTimer !== null) {
+      clearTimeout(this.reconnectTimer);
+      this.reconnectTimer = null;
+    }
     this.stopPing();
-    this.ws?.close();
+    const socket = this.ws;
     this.ws = null;
+    socket?.close();
     this._connected = false;
     this._pending = [];
   }
@@ -206,7 +218,7 @@ export class NerveWebSocket {
   }
 
   private scheduleReconnect() {
-    if (this.reconnectTimer) return;
+    if (!this.shouldReconnect || this.reconnectTimer !== null) return;
     this.reconnectTimer = setTimeout(() => {
       this.reconnectTimer = null;
       this.connect();