How to Use Zustand with WebSocket for Real-Time Data Updates

You can integrate WebSockets with Zustand by creating the socket connection outside React and calling store.setState() (or the store's action methods) inside the onmessage handler to push real-time server data into your global state.

Zustand is a minimal, unopinionated state management library for React. Because the store created by create() from the pmndrs/zustand repository is a plain JavaScript object with setState and getState methods defined in src/vanilla.ts, you can update it from any asynchronous source—including WebSockets, Server-Sent Events, or GraphQL subscriptions—and have your React components react to those changes instantly.

Why Zustand Works Seamlessly with WebSockets

The core of Zustand lives in src/vanilla.ts, which exports a factory that returns a store object containing setState, getState, getInitialState, and subscribe【/cache/repos/github.com/pmndrs/zustand/main/src/vanilla.ts#L9-L14】. Because this store is not tied to React's lifecycle, you can import it into any module—such as a WebSocket service worker—and mutate state directly.

When you call setState, the implementation merges the partial state (or replaces it if using the function form) and synchronously notifies all listeners【/cache/repos/github.com/pmndrs/zustand/main/src/vanilla.ts#L66-L81】. This synchronous notification is critical for real-time updates, ensuring that your UI reflects new WebSocket data immediately without batching delays.

Setting Up a Zustand Store for Real-Time Data

First, define a store that matches the shape of your incoming WebSocket messages. Create this at the module level so it can be imported by both your React components and your WebSocket handler.

// store.ts
import { create } from 'zustand'

interface ChatState {
  messages: string[]
  addMessage: (msg: string) => void
}

export const useChatStore = create<ChatState>((set) => ({
  messages: [],
  addMessage: (msg) =>
    set((state) => ({ messages: [...state.messages, msg] })),
}))

This creates the vanilla store object that you can access outside of React via useChatStore.getState().

Integrating WebSocket with Zustand

Module-Level WebSocket Connection

Establish the WebSocket connection at the module level to ensure a single shared connection across your application. Import the store and update it directly inside the onmessage handler.

// websocket.ts
import { useChatStore } from './store'

const WS_URL = 'wss://example.com/chat'

const ws = new WebSocket(WS_URL)

ws.addEventListener('open', () => {
  console.log('WebSocket connected')
})

ws.addEventListener('message', (event) => {
  // Parse the incoming payload
  const message = event.data
  
  // Update Zustand store directly
  useChatStore.getState().addMessage(message)
})

ws.addEventListener('error', (error) => {
  console.error('WebSocket error:', error)
})

export const closeConnection = () => ws.close()
export const sendMessage = (msg: string) => ws.send(msg)

Key points:

• The store update happens via useChatStore.getState().addMessage(), which calls setState internally.
• All subscribed React components re-render automatically because Zustand notifies listeners synchronously after the state change.

Handling Incoming Messages with Selectors

In your React components, use selectors to subscribe only to the specific slice of state you need. This prevents unnecessary re-renders when other parts of the store change.

// ChatLog.tsx
import { useChatStore } from './store'

export function ChatLog() {
  // Select only the messages array
  const messages = useChatStore((state) => state.messages)

  return (
    <ul>
      {messages.map((msg, index) => (
        <li key={index}>{msg}</li>
      ))}
    </ul>
  )
}

Using subscribeWithSelector for Granular Updates

For non-React code (such as logging services or other side effects), use the subscribeWithSelector middleware to listen to specific state changes without writing manual selector logic.

According to the source in src/middleware/subscribeWithSelector.ts, this middleware adds an overloaded subscribe method that accepts a selector and an equality function【/cache/repos/github.com/pmndrs/zustand/main/src/middleware/subscribeWithSelector.ts#L28-L39】.

// logger.ts
import { create } from 'zustand'
import { subscribeWithSelector } from 'zustand/middleware'

interface UserState {
  users: string[]
  setUsers: (users: string[]) => void
}

const useUserStore = create(
  subscribeWithSelector<UserState>((set) => ({
    users: [],
    setUsers: (users) => set({ users }),
  }))
)

// Subscribe only to the 'users' array from outside React
useUserStore.subscribe(
  (state) => state.users,
  (newUsers, prevUsers) => {
    console.log('Users updated:', { prev: prevUsers, next: newUsers })
  },
  { fireImmediately: true } // Optional: run callback immediately with current state
)

This approach is ideal when you need to synchronize Zustand state with other real-time systems or analytics trackers.

Managing Connection Lifecycle in React Components

While the WebSocket should generally live at the module level to avoid duplicate connections, you may need to manage the connection lifecycle within a specific component (for example, when the component mounts and unmounts).

Use useEffect to set up the WebSocket and clean it up when the component unmounts. Import the store's action methods to update state.

import { useEffect } from 'react'
import { useChatStore } from './store'

export function ChatComponent() {
  const addMessage = useChatStore((state) => state.addMessage)

  useEffect(() => {
    const ws = new WebSocket('wss://example.com/chat')
    
    const handleMessage = (event: MessageEvent) => {
      addMessage(event.data)
    }
    
    ws.addEventListener('message', handleMessage)
    
    // Cleanup: remove listener and close connection
    return () => {
      ws.removeEventListener('message', handleMessage)
      ws.close()
    }
  }, [addMessage])

  // ... render UI
  return <div>{/* Chat UI */}</div>
}

Important: The store itself does not require cleanup—only the WebSocket connection does. The addMessage function is stable if defined properly in the store, but including it in the dependency array ensures the effect uses the latest reference.

Summary

• Zustand stores are plain JavaScript objects defined in src/vanilla.ts, allowing updates from any asynchronous source outside React.
• Update the store directly from WebSocket onmessage handlers using store.getState().setState() or action methods to trigger synchronous UI updates.
• Use selectors in React components to subscribe only to specific state slices, minimizing re-renders when real-time data changes.
• Leverage subscribeWithSelector from src/middleware/subscribeWithSelector.ts for non-React side effects that need to react to specific state changes.
• Manage WebSocket lifecycle at the module level for shared connections, or inside useEffect for component-scoped connections, ensuring proper cleanup to prevent memory leaks.

Frequently Asked Questions

How do I prevent duplicate WebSocket connections when using Zustand?

Create the WebSocket connection at the module level (outside any React component) and import it where needed. Since Zustand stores are also typically created at the module level in src/vanilla.ts, both the socket and store are instantiated exactly once when the application boots. If you must create the socket inside a component, use useEffect with a cleanup function that closes the connection when the component unmounts.

Can I use Zustand with Server-Sent Events (SSE) instead of WebSockets?

Yes. The same pattern applies to any real-time data source. Create an EventSource connection in a module or effect, and call store.setState() or your store's action methods inside the onmessage event handler. Zustand's vanilla store architecture in src/vanilla.ts is agnostic to the data source, handling updates synchronously regardless of whether they come from WebSockets, SSE, or polling intervals.

What is the difference between subscribe and subscribeWithSelector?

The base subscribe method from src/vanilla.ts notifies listeners whenever any part of the state changes, requiring manual filtering inside the callback. The subscribeWithSelector middleware, defined in src/middleware/subscribeWithSelector.ts, adds an overloaded subscribe method that accepts a selector function as the first argument. This allows you to subscribe only to specific state slices (e.g., (state) => state.messages) and optionally provide an equality function or fireImmediately option, making it ideal for non-React side effects that need precise updates.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}