Fix useStableSocket replacing socket on HMR effect re-runs

.changeset/fix-hmr-socket-reconnect.md

@@ -0,0 +1,5 @@
+---
+"partysocket": patch
+---
+
+Fix useStableSocket replacing socket on HMR/StrictMode effect re-runs. When Vite HMR fires, React Fast Refresh re-runs all effects — the old code unconditionally created a new socket, breaking downstream references (event listeners, \_pk identity, etc.). Now detects whether connection options actually changed via referential equality on the memoized options object: if unchanged (HMR), calls `socket.reconnect()` to preserve identity; if changed, creates a new socket as before.

.github/workflows/pkg-pr-new.yml

@@ -17,4 +17,4 @@ jobs:
 
       - run: npm ci
       - run: npm run build
-      - run: npx pkg-pr-new publish
+      - run: npx pkg-pr-new publish './packages/partyserver' './packages/partysocket' './packages/y-partyserver' './packages/partysub' './packages/partyfn' './packages/partysync' './packages/partywhen' './packages/partytracks' './packages/hono-party'

packages/partysocket/src/tests/react-hooks.test.tsx

@@ -3,6 +3,7 @@
  */
 
 import { renderHook, waitFor } from "@testing-library/react";
+import React from "react";
 import { afterAll, beforeAll, describe, expect, test, vitest } from "vitest";
 import { WebSocketServer } from "ws";
 
@@ -1114,3 +1115,174 @@ describe.skipIf(!!process.env.GITHUB_ACTIONS)("useWebSocket", () => {
     result.current.close();
   });
 });
+
+/**
+ * HMR (Hot Module Replacement) resilience tests.
+ *
+ * When Vite HMR fires, React Fast Refresh re-runs all effects without changing
+ * their dependencies. React StrictMode in development does the same thing —
+ * it double-invokes effects (run → cleanup → run). We use StrictMode as a
+ * reliable proxy to test the HMR code path in useStableSocket.
+ *
+ * The critical behavior: when the effect re-runs but socketOptions reference
+ * hasn't changed, we should call socket.reconnect() on the existing instance
+ * instead of creating a new socket. This preserves socket identity so that
+ * downstream code (event listeners, _pk references, etc.) isn't disrupted.
+ */
+describe.skipIf(!!process.env.GITHUB_ACTIONS)(
+  "HMR resilience (useStableSocket)",
+  () => {
+    const strictModeWrapper = ({ children }: { children: React.ReactNode }) =>
+      React.createElement(React.StrictMode, null, children);
+
+    test("usePartySocket preserves socket identity under StrictMode (simulates HMR)", () => {
+      const { result } = renderHook(
+        () =>
+          usePartySocket({
+            host: "example.com",
+            room: "test-room",
+            startClosed: true
+          }),
+        { wrapper: strictModeWrapper }
+      );
+
+      // The socket should still be the same instance — not replaced
+      expect(result.current).toBeDefined();
+      expect(result.current.host).toBe("example.com");
+      expect(result.current.room).toBe("test-room");
+    });
+
+    test("startClosed: true keeps socket closed under StrictMode (no spurious reconnect)", () => {
+      const { result } = renderHook(
+        () =>
+          usePartySocket({
+            host: "example.com",
+            room: "test-room",
+            startClosed: true
+          }),
+        { wrapper: strictModeWrapper }
+      );
+
+      // StrictMode double-invokes effects (run → cleanup → run).
+      // With startClosed: true the HMR branch must NOT call reconnect(),
+      // so the socket should still be CLOSED after the double-invoke.
+      expect(result.current.readyState).toBe(WebSocket.CLOSED);
+    });
+
+    test("usePartySocket preserves socket identity under StrictMode without startClosed", () => {
+      const { result } = renderHook(
+        () =>
+          usePartySocket({
+            host: "example.com",
+            room: "test-room"
+          }),
+        { wrapper: strictModeWrapper }
+      );
+
+      // Without startClosed, the HMR branch calls reconnect() on the
+      // existing socket rather than creating a new instance.
+      // Socket identity must be preserved regardless.
+      expect(result.current).toBeDefined();
+      expect(result.current.host).toBe("example.com");
+      expect(result.current.room).toBe("test-room");
+    });
+
+    test("usePartySocket still creates new socket when options change under StrictMode", async () => {
+      const { result, rerender } = renderHook(
+        ({ room }) =>
+          usePartySocket({
+            host: "example.com",
+            room,
+            startClosed: true
+          }),
+        {
+          initialProps: { room: "room1" },
+          wrapper: strictModeWrapper
+        }
+      );
+
+      const firstSocket = result.current;
+      expect(firstSocket.room).toBe("room1");
+
+      // Change an option — this should create a new socket, not reconnect
+      rerender({ room: "room2" });
+
+      await waitFor(() => {
+        expect(result.current).not.toBe(firstSocket);
+        expect(result.current.room).toBe("room2");
+      });
+    });
+
+    test("useWebSocket preserves socket identity under StrictMode (simulates HMR)", () => {
+      const { result } = renderHook(
+        () =>
+          useWebSocket("ws://example.com", undefined, {
+            startClosed: true
+          }),
+        { wrapper: strictModeWrapper }
+      );
+
+      expect(result.current).toBeDefined();
+      expect(result.current.readyState).toBe(WebSocket.CLOSED);
+    });
+
+    test("useWebSocket still creates new socket when URL changes under StrictMode", async () => {
+      const { result, rerender } = renderHook(
+        ({ url }) =>
+          useWebSocket(url, undefined, {
+            startClosed: true
+          }),
+        {
+          initialProps: { url: "ws://example.com/1" },
+          wrapper: strictModeWrapper
+        }
+      );
+
+      const firstSocket = result.current;
+
+      rerender({ url: "ws://example.com/2" });
+
+      await waitFor(() => {
+        expect(result.current).not.toBe(firstSocket);
+      });
+    });
+
+    test("socket identity is preserved across multiple rerenders with same props", () => {
+      const { result, rerender } = renderHook(
+        () =>
+          usePartySocket({
+            host: "example.com",
+            room: "test-room",
+            startClosed: true
+          }),
+        { wrapper: strictModeWrapper }
+      );
+
+      const originalSocket = result.current;
+
+      // Multiple rerenders with identical props should never replace the socket
+      for (let i = 0; i < 5; i++) {
+        rerender();
+        expect(result.current).toBe(originalSocket);
+      }
+    });
+
+    test("unmount still calls close() under StrictMode", () => {
+      const { result, unmount } = renderHook(
+        () =>
+          usePartySocket({
+            host: "example.com",
+            room: "test-room",
+            startClosed: true
+          }),
+        { wrapper: strictModeWrapper }
+      );
+
+      const closeSpy = vitest.spyOn(result.current, "close");
+
+      unmount();
+
+      expect(closeSpy).toHaveBeenCalled();
+    });
+  }
+);

packages/partysocket/src/use-socket.ts

@@ -65,8 +65,17 @@ export function useStableSocket<
   // track the previous enabled state to detect changes
   const prevEnabledRef = useRef(enabled);
 
+  // track the previous socketOptions reference to distinguish option changes
+  // from HMR/StrictMode effect re-runs. useMemo returns the same reference
+  // when the memo key hasn't changed, so referential equality tells us
+  // whether the connection options actually changed.
+  const prevSocketOptionsRef = useRef(socketOptions);
+
   // finally, initialize the socket
   useEffect(() => {
+    const optionsChanged = prevSocketOptionsRef.current !== socketOptions;
+    prevSocketOptionsRef.current = socketOptions;
+
     // if disabled, close the socket and don't proceed with connection logic
     if (!enabled) {
       socket.close();
@@ -85,16 +94,30 @@ export function useStableSocket<
 
     // we haven't yet restarted the socket
     if (socketInitializedRef.current === socket) {
-      // create new socket
-      const newSocket = createSocketRef.current({
-        ...socketOptions,
-        // when reconnecting because of options change, we always reconnect
-        // (startClosed only applies to initial mount)
-        startClosed: false
-      });
-
-      // update socket reference (this will cause the effect to run again)
-      setSocket(newSocket);
+      if (optionsChanged) {
+        // connection options changed — create new socket with new config
+        const newSocket = createSocketRef.current({
+          ...socketOptions,
+          // when reconnecting because of options change, we always reconnect
+          // (startClosed only applies to initial mount)
+          startClosed: false
+        });
+
+        // update socket reference (this will cause the effect to run again)
+        setSocket(newSocket);
+      } else {
+        // HMR or React Strict Mode effect re-run — reconnect the existing
+        // socket instead of creating a new instance. This preserves the
+        // socket identity (event listeners, _pk, etc.) across Hot Module
+        // Replacement, preventing downstream code from losing its reference
+        // to the live socket.
+        if (socketOptions.startClosed !== true) {
+          socket.reconnect();
+        }
+        return () => {
+          socket.close();
+        };
+      }
     } else {
       // if this is the first time we are running the hook, connect...
       if (!socketInitializedRef.current && socketOptions.startClosed !== true) {