diff --git a/.changeset/fix-hmr-socket-reconnect.md b/.changeset/fix-hmr-socket-reconnect.md new file mode 100644 index 00000000..1de481b0 --- /dev/null +++ b/.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. diff --git a/.github/workflows/pkg-pr-new.yml b/.github/workflows/pkg-pr-new.yml index 430368eb..da26708b 100644 --- a/.github/workflows/pkg-pr-new.yml +++ b/.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' diff --git a/packages/partysocket/src/tests/react-hooks.test.tsx b/packages/partysocket/src/tests/react-hooks.test.tsx index a2c52624..8fc20e0e 100644 --- a/packages/partysocket/src/tests/react-hooks.test.tsx +++ b/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(); + }); + } +); diff --git a/packages/partysocket/src/use-socket.ts b/packages/partysocket/src/use-socket.ts index 065f00e4..dc86fa4d 100644 --- a/packages/partysocket/src/use-socket.ts +++ b/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) {