Fixing Zustand stale WebSocket subscriptions #

You wired a WebSocket message handler once — inside a useEffect with an empty dependency array, or in a module-level singleton — and now your Zustand store reads the wrong thing. New messages either merge into a snapshot of state from the moment the socket opened, or worse, they call set on a store instance that a remounted component has already torn down. The symptom is maddening: the network tab shows the frame arriving, your handler runs, but the UI shows values that are seconds or whole sessions out of date. This is a stale-closure bug, and it lives at the seam between an event source that outlives React and a WebSocket state sync and optimistic update pipeline that assumes fresh reads.

Root cause #

A WebSocket handler is a long-lived callback. When you register ws.onmessage = handler once, JavaScript closes over every variable the handler references at definition time. If that handler references a Zustand value you destructured from the hook — const { messages } = useStore() — it captures the array reference from the render in which the effect ran. Zustand never mutates that array; every set produces a new object. So the handler keeps appending to the original, frozen-in-time snapshot, and the spread [...messages, incoming] is built on stale data. The store’s actual state has moved on; your closure has not.

The runtime mechanics are pure lexical scoping, not a React quirk. A useEffect(() => { ws.onmessage = ... }, []) runs once, on mount. React’s render closures are immutable per render, so the messages binding inside that effect is permanently bound to the first render’s value. Subsequent renders create new messages bindings, but nothing rebinds the already-registered handler. The same trap catches set calls that target a store slice captured by reference, and it catches StrictMode double-mounts in development, where the second mount can leave a handler from the first, unmounted instance still firing into a store the component no longer renders.

There are three independent failures stacked here, and the fix addresses each:

  1. Stale reads — the handler computes the next state from a captured snapshot instead of live state.
  2. Writes to a torn-down store — a handler registered before unmount keeps calling set after the component is gone, mutating state nobody renders and pinning the closure in memory.
  3. Subscription leaks — registering on every render or never unregistering, a class of bug covered in depth under memory leak prevention.
Stale closure vs live getState in a WebSocket handler A handler registered once captures an old state snapshot, while calling getState inside the handler reads the current store on every message. Buggy: captured snapshot register once closes over state v1 writes onto v1 drops v2, v3... Fixed: getState per message register once no captured state getState() + set reads live v3

Resolution #

The rule is simple: never close over Zustand state values inside a long-lived handler — close over the store object and call getState()/set on it per message. The store object is stable for the lifetime of the module; the state it holds is not. For transient, high-frequency updates that should not re-render the component that owns the socket, use store.subscribe outside React entirely. Register on mount, tear down on unmount.

import { create } from 'zustand';

// 1. State shape. Actions live INSIDE the store and use `set`'s updater form,
//    which always receives the live previous state — no external snapshot.
interface ChatState {
  messages: { id: string; text: string }[];
  status: 'open' | 'closed';
  appendMessage: (m: { id: string; text: string }) => void;
  setStatus: (s: ChatState['status']) => void;
}

export const useChatStore = create<ChatState>((set) => ({
  messages: [],
  status: 'closed',
  // `set((prev) => ...)` reads the CURRENT state at call time, not a closure.
  appendMessage: (m) =>
    set((prev) => ({ messages: [...prev.messages, m] })),
  setStatus: (s) => set({ status: s }),
}));

// 2. Connect OUTSIDE React. The handler closes over `useChatStore` (the stable
//    store object), never over destructured state. So every frame reads live state.
export function connectChatSocket(url: string): () => void {
  const ws = new WebSocket(url);

  const onOpen = () => useChatStore.getState().setStatus('open');

  const onMessage = (event: MessageEvent) => {
    const incoming = JSON.parse(event.data) as { id: string; text: string };
    // getState() resolves the CURRENT store every message — no stale snapshot.
    useChatStore.getState().appendMessage(incoming);
  };

  const onClose = () => useChatStore.getState().setStatus('closed');

  ws.addEventListener('open', onOpen);
  ws.addEventListener('message', onMessage);
  ws.addEventListener('close', onClose);

  // 3. Return a teardown that removes listeners AND closes the socket, so a
  //    remount cannot leave a handler writing into a torn-down store.
  return () => {
    ws.removeEventListener('open', onOpen);
    ws.removeEventListener('message', onMessage);
    ws.removeEventListener('close', onClose);
    if (ws.readyState <= WebSocket.OPEN) ws.close(1000, 'unmount');
  };
}

// 4. Transient updates that must NOT re-render: subscribe to a slice outside
//    React. store.subscribe fires on every change; the selector limits scope.
export function watchUnread(onChange: (count: number) => void): () => void {
  return useChatStore.subscribe(
    // subscribe returns its own unsubscribe — call it on unmount.
    (state) => onChange(state.messages.length),
  );
}

Wire the connection from the component with a useEffect whose cleanup runs the returned teardown. The effect body holds no Zustand values, so nothing it captures can go stale:

import { useEffect } from 'react';

function ChatPanel({ url }: { url: string }) {
  useEffect(() => {
    const disconnect = connectChatSocket(url);
    return disconnect; // runs on unmount and on url change — no leaked handler
  }, [url]);

  // Subscribe to exactly the slice this component renders. This read is fresh
  // because the selector re-runs on each store change, not on each socket frame.
  const messages = useChatStore((s) => s.messages);
  return <ul>{messages.map((m) => <li key={m.id}>{m.text}</li>)}</ul>;
}

Contrast the broken handler — const { messages } = useChatStore() then ws.onmessage = () => useChatStore.setState({ messages: [...messages, incoming] }) — which spreads a frozen array and silently drops every concurrent update. The fixed handler holds no state binding at all; it asks the store for the truth on each frame.

Operational checklist #

  • No WebSocket handler destructures Zustand state — handlers reference the store object and call getState()/setState/set
  • Every set that derives from existing state uses the updater form set((prev) => …)
  • The connect function returns a teardown that removes all listeners and closes the socket, and a useEffect
  • store.subscribe
  • Heap snapshot before/after several mount-unmount cycles shows no growing count of detached WebSocket
  • Rollback plan: if a release regresses, revert to a single module-level connectChatSocket

FAQ #

Why does useStore.getState() fix the stale closure when const { x } = useStore() does not? #

useStore() (with a selector or destructure) returns a value sampled during a render and bound into that render’s closure. A handler registered once keeps that one binding forever. useStore.getState() calls a function on the stable store object at the moment the message arrives, so it returns whatever state is current then. The difference is snapshot-at-registration versus read-at-invocation.

Should the socket live inside a React component or outside it? #

Outside, for anything long-lived. A socket that outlives a component should be owned by a module-level connect function (or a store action) and merely started by a useEffect. The component subscribes to the store slice it renders; it does not own the transport. This keeps the handler’s closure free of React render state.

When do I use store.subscribe instead of the useStore hook? #

Use store.subscribe for transient, high-frequency, or imperative work that must not trigger a React re-render on the owning component — driving a canvas, throttling a counter, or syncing to localStorage. It runs outside React’s render cycle and returns its own unsubscribe. Use the useStore(selector) hook when the value belongs in the rendered tree.

Does StrictMode change anything here? #

Yes. StrictMode mounts, unmounts, and remounts components in development to surface exactly this class of bug. If your cleanup is correct — listeners removed and socket closed on unmount — StrictMode is harmless. If you leak a handler, StrictMode makes it fire twice and corrupt state visibly, which is the early-warning you want.

What changes for Socket.IO instead of raw ws? #

The closure mechanics are identical; only the listener API differs (socket.on('event', fn) / socket.off('event', fn)). Still reference the store object and call getState()/set inside the handler, and still return a teardown that calls socket.off and socket.disconnect() on unmount.

Back to WebSocket State Sync and Optimistic Updates