Add enabled prop to usePartySocket and useWebSocket hooks

.changeset/bright-sockets-flow.md

@@ -0,0 +1,5 @@
+---
+"partysocket": patch
+---
+
+Add `enabled` prop to `usePartySocket` and `useWebSocket` hooks for conditional connection control

packages/partysocket/src/react.ts

@@ -7,9 +7,11 @@ import {
 
 import type { PartySocketOptions } from "./index";
 import type { EventHandlerOptions } from "./use-handlers";
+import type { SocketOptions } from "./use-socket";
 
 type UsePartySocketOptions = Omit<PartySocketOptions, "host"> &
-  EventHandlerOptions & {
+  EventHandlerOptions &
+  Pick<SocketOptions, "enabled"> & {
     host?: string | undefined;
   };
 

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

@@ -11,7 +11,7 @@ import usePartySocket, { useWebSocket } from "../react";
 const PORT = 50128;
 //  const URL = `ws://localhost:${PORT}`;
 
-describe.skip("usePartySocket", () => {
+describe("usePartySocket", () => {
   let wss: WebSocketServer;
 
   beforeAll(() => {
@@ -653,9 +653,145 @@ describe.skip("usePartySocket", () => {
 
     result.current.close();
   });
+
+  test("does not connect when enabled is false", () => {
+    const { result } = renderHook(() =>
+      usePartySocket({
+        host: `localhost:${PORT}`,
+        room: "test-room",
+        enabled: false
+      })
+    );
+
+    expect(result.current).toBeDefined();
+    expect(result.current.readyState).toBe(WebSocket.CLOSED);
+  });
+
+  test("connects when enabled is true (default)", async () => {
+    wss.once("connection", (ws) => {
+      ws.close();
+    });
+
+    const { result } = renderHook(() =>
+      usePartySocket({
+        host: `localhost:${PORT}`,
+        room: "test-room",
+        enabled: true
+      })
+    );
+
+    await waitFor(
+      () => {
+        expect(result.current.readyState).toBe(WebSocket.OPEN);
+      },
+      { timeout: 3000 }
+    );
+
+    result.current.close();
+  });
+
+  test("disconnects when enabled changes from true to false", async () => {
+    wss.once("connection", (ws) => {
+      // Keep connection open
+    });
+
+    const { result, rerender } = renderHook(
+      ({ enabled }) =>
+        usePartySocket({
+          host: `localhost:${PORT}`,
+          room: "test-room",
+          enabled
+        }),
+      { initialProps: { enabled: true } }
+    );
+
+    await waitFor(
+      () => {
+        expect(result.current.readyState).toBe(WebSocket.OPEN);
+      },
+      { timeout: 3000 }
+    );
+
+    rerender({ enabled: false });
+
+    await waitFor(
+      () => {
+        expect(result.current.readyState).toBe(WebSocket.CLOSED);
+      },
+      { timeout: 3000 }
+    );
+  });
+
+  test("reconnects when enabled changes from false to true", async () => {
+    wss.once("connection", (ws) => {
+      // Keep connection open
+    });
+
+    const { result, rerender } = renderHook(
+      ({ enabled }) =>
+        usePartySocket({
+          host: `localhost:${PORT}`,
+          room: "test-room",
+          enabled
+        }),
+      { initialProps: { enabled: false } }
+    );
+
+    expect(result.current.readyState).toBe(WebSocket.CLOSED);
+
+    rerender({ enabled: true });
+
+    await waitFor(
+      () => {
+        expect(result.current.readyState).toBe(WebSocket.OPEN);
+      },
+      { timeout: 3000 }
+    );
+
+    result.current.close();
+  });
+
+  test("keeps the same socket instance when enabled toggles", async () => {
+    wss.once("connection", () => {
+      // Keep connection open
+    });
+
+    const { result, rerender } = renderHook(
+      ({ enabled }) =>
+        usePartySocket({
+          host: `localhost:${PORT}`,
+          room: "test-room",
+          enabled
+        }),
+      { initialProps: { enabled: true } }
+    );
+
+    await waitFor(
+      () => {
+        expect(result.current.readyState).toBe(WebSocket.OPEN);
+      },
+      { timeout: 3000 }
+    );
+
+    const socketInstance = result.current;
+
+    rerender({ enabled: false });
+
+    await waitFor(
+      () => {
+        expect(result.current.readyState).toBe(WebSocket.CLOSED);
+      },
+      { timeout: 3000 }
+    );
+
+    // Same socket instance should be reused
+    expect(result.current).toBe(socketInstance);
+
+    result.current.close();
+  });
 });
 
-describe.skip("useWebSocket", () => {
+describe("useWebSocket", () => {
   let wss: WebSocketServer;
 
   beforeAll(() => {
@@ -835,4 +971,130 @@ describe.skip("useWebSocket", () => {
 
     result.current.close();
   });
+
+  test("does not connect when enabled is false", () => {
+    const { result } = renderHook(() =>
+      useWebSocket(`ws://localhost:${PORT + 1}`, undefined, {
+        enabled: false
+      })
+    );
+
+    expect(result.current).toBeDefined();
+    expect(result.current.readyState).toBe(WebSocket.CLOSED);
+  });
+
+  test("connects when enabled is true (default)", async () => {
+    wss.once("connection", (ws) => {
+      ws.close();
+    });
+
+    const { result } = renderHook(() =>
+      useWebSocket(`ws://localhost:${PORT + 1}`, undefined, {
+        enabled: true
+      })
+    );
+
+    await waitFor(
+      () => {
+        expect(result.current.readyState).toBe(WebSocket.OPEN);
+      },
+      { timeout: 3000 }
+    );
+
+    result.current.close();
+  });
+
+  test("disconnects when enabled changes from true to false", async () => {
+    wss.once("connection", () => {
+      // Keep connection open
+    });
+
+    const { result, rerender } = renderHook(
+      ({ enabled }) =>
+        useWebSocket(`ws://localhost:${PORT + 1}`, undefined, {
+          enabled
+        }),
+      { initialProps: { enabled: true } }
+    );
+
+    await waitFor(
+      () => {
+        expect(result.current.readyState).toBe(WebSocket.OPEN);
+      },
+      { timeout: 3000 }
+    );
+
+    rerender({ enabled: false });
+
+    await waitFor(
+      () => {
+        expect(result.current.readyState).toBe(WebSocket.CLOSED);
+      },
+      { timeout: 3000 }
+    );
+  });
+
+  test("reconnects when enabled changes from false to true", async () => {
+    wss.once("connection", () => {
+      // Keep connection open
+    });
+
+    const { result, rerender } = renderHook(
+      ({ enabled }) =>
+        useWebSocket(`ws://localhost:${PORT + 1}`, undefined, {
+          enabled
+        }),
+      { initialProps: { enabled: false } }
+    );
+
+    expect(result.current.readyState).toBe(WebSocket.CLOSED);
+
+    rerender({ enabled: true });
+
+    await waitFor(
+      () => {
+        expect(result.current.readyState).toBe(WebSocket.OPEN);
+      },
+      { timeout: 3000 }
+    );
+
+    result.current.close();
+  });
+
+  test("keeps the same socket instance when enabled toggles", async () => {
+    wss.once("connection", () => {
+      // Keep connection open
+    });
+
+    const { result, rerender } = renderHook(
+      ({ enabled }) =>
+        useWebSocket(`ws://localhost:${PORT + 1}`, undefined, {
+          enabled
+        }),
+      { initialProps: { enabled: true } }
+    );
+
+    await waitFor(
+      () => {
+        expect(result.current.readyState).toBe(WebSocket.OPEN);
+      },
+      { timeout: 3000 }
+    );
+
+    const socketInstance = result.current;
+
+    rerender({ enabled: false });
+
+    await waitFor(
+      () => {
+        expect(result.current.readyState).toBe(WebSocket.CLOSED);
+      },
+      { timeout: 3000 }
+    );
+
+    // Same socket instance should be reused
+    expect(result.current).toBe(socketInstance);
+
+    result.current.close();
+  });
 });

packages/partysocket/src/use-socket.ts

@@ -3,10 +3,16 @@ import { useEffect, useMemo, useRef, useState } from "react";
 import type WebSocket from "./ws";
 import type { Options } from "./ws";
 
+export type SocketOptions = Options & {
+  /** Whether the socket should be connected. Defaults to true. */
+  enabled?: boolean;
+};
+
 /** When any of the option values are changed, we should reinitialize the socket */
 export const getOptionsThatShouldCauseRestartWhenChanged = (
-  options: Options
+  options: SocketOptions
 ) => [
+  // Note: enabled is handled separately to avoid creating a new socket on toggle
   options.startClosed,
   options.minUptime,
   options.maxRetries,
@@ -22,7 +28,10 @@ export const getOptionsThatShouldCauseRestartWhenChanged = (
  * Initializes a PartySocket (or WebSocket) and keeps it stable across renders,
  * but reconnects and updates the reference when any of the connection args change.
  */
-export function useStableSocket<T extends WebSocket, TOpts extends Options>({
+export function useStableSocket<
+  T extends WebSocket,
+  TOpts extends SocketOptions
+>({
   options,
   createSocket,
   createSocketMemoKey: createOptionsMemoKey
@@ -31,6 +40,9 @@ export function useStableSocket<T extends WebSocket, TOpts extends Options>({
   createSocket: (options: TOpts) => T;
   createSocketMemoKey: (options: TOpts) => string;
 }) {
+  // extract enabled with default value of true
+  const { enabled = true } = options;
+
   // ensure we only reconnect when necessary
   const shouldReconnect = createOptionsMemoKey(options);
   const socketOptions = useMemo(() => {
@@ -50,8 +62,27 @@ export function useStableSocket<T extends WebSocket, TOpts extends Options>({
   const createSocketRef = useRef(createSocket);
   createSocketRef.current = createSocket;
 
+  // track the previous enabled state to detect changes
+  const prevEnabledRef = useRef(enabled);
+
   // finally, initialize the socket
   useEffect(() => {
+    // if disabled, close the socket and don't proceed with connection logic
+    if (!enabled) {
+      socket.close();
+      prevEnabledRef.current = enabled;
+      return;
+    }
+
+    // if enabled just changed from false to true, reconnect
+    if (!prevEnabledRef.current && enabled) {
+      socket.reconnect();
+      prevEnabledRef.current = enabled;
+      return;
+    }
+
+    prevEnabledRef.current = enabled;
+
     // we haven't yet restarted the socket
     if (socketInitializedRef.current === socket) {
       // create new socket
@@ -76,7 +107,7 @@ export function useStableSocket<T extends WebSocket, TOpts extends Options>({
         socket.close();
       };
     }
-  }, [socket, socketOptions]);
+  }, [socket, socketOptions, enabled]);
 
   return socket;
 }

packages/partysocket/src/use-ws.ts

@@ -6,9 +6,10 @@ import {
 import WebSocket from "./ws";
 
 import type { EventHandlerOptions } from "./use-handlers";
-import type { Options, ProtocolsProvider, UrlProvider } from "./ws";
+import type { SocketOptions } from "./use-socket";
+import type { ProtocolsProvider, UrlProvider } from "./ws";
 
-type UseWebSocketOptions = Options & EventHandlerOptions;
+type UseWebSocketOptions = SocketOptions & EventHandlerOptions;
 
 // A React hook that wraps PartySocket
 export default function useWebSocket(