fix(socket): prevent duplicate connections on rapid offline-online toggles (fixes #1390) (#1410)

* fix(socket): prevent duplicate connections on rapid offline-online toggles (fixes #1390)

* fix(chat): resolve upstream merge conflict markers and tts lint warnings

---------

Co-authored-by: unknown <your-email@example.com>

Author: SanjanaTech19

src/__tests__/lib/partySocketReconnect.test.ts

@@ -58,6 +58,83 @@ describe("attachJitteredBackoff", () => {
     attachJitteredBackoff(socket);
     expect(socket._getNextDelay).toBe(first);
   });
+
+  it("guards connection attempts using ConnectionState enum", () => {
+    const mockConnect = jest.fn();
+    const mockDisconnect = jest.fn();
+    const eventListeners: Record<string, Array<() => void>> = {};
+
+    const socket = {
+      _retryCount: 0,
+      _getNextDelay: () => 10,
+      _connect: mockConnect,
+      _disconnect: mockDisconnect,
+      _clearTimeouts: jest.fn(),
+      addEventListener: (event: string, cb: () => void) => {
+        if (!eventListeners[event]) eventListeners[event] = [];
+        eventListeners[event].push(cb);
+      },
+    } as any;
+
+    attachJitteredBackoff(socket);
+
+    // Initial state is CLOSED
+    expect(socket.__worksphereState).toBe("CLOSED");
+
+    // Calling connect transitions state to CONNECTING
+    socket._connect();
+    expect(socket.__worksphereState).toBe("CONNECTING");
+    expect(mockConnect).toHaveBeenCalledTimes(1);
+
+    // Calling connect again while CONNECTING does not trigger mockConnect again
+    socket._connect();
+    expect(mockConnect).toHaveBeenCalledTimes(1);
+
+    // Simulate open event transitions to CONNECTED
+    eventListeners["open"]?.forEach((cb) => cb());
+    expect(socket.__worksphereState).toBe("CONNECTED");
+
+    // Calling connect while CONNECTED does not trigger mockConnect
+    socket._connect();
+    expect(mockConnect).toHaveBeenCalledTimes(1);
+
+    // Calling disconnect transitions state to CLOSED
+    socket._disconnect();
+    expect(socket.__worksphereState).toBe("CLOSED");
+    expect(mockDisconnect).toHaveBeenCalledTimes(1);
+  });
+
+  it("aborts pending reconnect timers on new connection or disconnect", async () => {
+    jest.useFakeTimers();
+
+    const socket = {
+      _retryCount: 1,
+      _getNextDelay: () => 1000,
+      _connect: jest.fn(),
+      _disconnect: jest.fn(),
+      _clearTimeouts: jest.fn(),
+      addEventListener: jest.fn(),
+    } as any;
+
+    attachJitteredBackoff(socket);
+
+    // Start waiting for reconnect
+    const _waitPromise = socket._wait();
+
+    // Reconnect timer is scheduled in setTimeout. Let's call disconnect before it fires.
+    socket._disconnect();
+
+    // Fast-forward time
+    jest.advanceTimersByTime(1000);
+
+    // Wait for any microtasks
+    await Promise.resolve();
+
+    // Verify waitPromise did not resolve yet because the timeout was cleared by _disconnect
+    // Since waitPromise resolves inside setTimeout which was cleared, it remains pending.
+    // If we trigger _connect, it also clears any pending timeouts.
+    jest.useRealTimers();
+  });
 });
 
 import { PartySocketReconnectManager } from "@/lib/partySocketReconnect";

src/lib/partySocketReconnect.ts

@@ -39,20 +39,106 @@ export function jitteredReconnectDelay(
   return Math.round(Math.min(max, Math.max(min, base + jitter)));
 }
 
+export enum ConnectionState {
+  CLOSED = "CLOSED",
+  CONNECTING = "CONNECTING",
+  CONNECTED = "CONNECTED",
+}
+
 type DelaySocket = {
   _retryCount: number;
   _getNextDelay: () => number;
+  _connect?: () => void;
+  _disconnect?: (code?: number, reason?: string) => void;
+  _clearTimeouts?: () => void;
+  _wait?: () => Promise<void>;
+  addEventListener?: (
+    event: string,
+    callback: (...args: any[]) => void,
+  ) => void;
   __worksphereJitter?: boolean;
+  __worksphereState?: ConnectionState;
 };
 
 /** Swap in jittered backoff on a live PartySocket instance (idempotent). */
 export function attachJitteredBackoff<T extends object>(socket: T): T {
   const s = socket as T & DelaySocket;
   if (s.__worksphereJitter) return socket;
 
+  let pendingTimeoutId: any = null;
+  s.__worksphereState = ConnectionState.CLOSED;
+
   s._getNextDelay = function (this: DelaySocket) {
     return jitteredReconnectDelay(this._retryCount);
   };
+
+  s._wait = function (this: any) {
+    if (pendingTimeoutId) {
+      clearTimeout(pendingTimeoutId);
+    }
+    return new Promise<void>((resolve) => {
+      pendingTimeoutId = setTimeout(() => {
+        pendingTimeoutId = null;
+        resolve();
+      }, this._getNextDelay());
+    });
+  };
+
+  const originalClearTimeouts = s._clearTimeouts;
+  s._clearTimeouts = function (this: any) {
+    if (pendingTimeoutId) {
+      clearTimeout(pendingTimeoutId);
+      pendingTimeoutId = null;
+    }
+    if (originalClearTimeouts) {
+      originalClearTimeouts.call(this);
+    }
+  };
+
+  const originalDisconnect = s._disconnect;
+  s._disconnect = function (this: any, code?: number, reason?: string) {
+    s.__worksphereState = ConnectionState.CLOSED;
+    if (pendingTimeoutId) {
+      clearTimeout(pendingTimeoutId);
+      pendingTimeoutId = null;
+    }
+    if (originalDisconnect) {
+      originalDisconnect.call(this, code, reason);
+    }
+  };
+
+  const originalConnect = s._connect;
+  if (originalConnect) {
+    s._connect = function (this: any) {
+      if (
+        s.__worksphereState === ConnectionState.CONNECTING ||
+        s.__worksphereState === ConnectionState.CONNECTED
+      ) {
+        return;
+      }
+      s.__worksphereState = ConnectionState.CONNECTING;
+      if (pendingTimeoutId) {
+        clearTimeout(pendingTimeoutId);
+        pendingTimeoutId = null;
+      }
+      originalConnect.call(this);
+    };
+  }
+
+  if (typeof s.addEventListener === "function") {
+    s.addEventListener("open", () => {
+      s.__worksphereState = ConnectionState.CONNECTED;
+    });
+
+    s.addEventListener("close", () => {
+      s.__worksphereState = ConnectionState.CLOSED;
+    });
+
+    s.addEventListener("error", () => {
+      s.__worksphereState = ConnectionState.CLOSED;
+    });
+  }
+
   s.__worksphereJitter = true;
   return socket;
 }