Building a real-time chat app sounds straightforward until React's rendering model meets WebSockets. In Relay, I built what I thought was a solid real-time integration. Everything seemed to work perfectly—until I started switching between conversations.

What I Built

In Relay, I wanted a clean separation of concerns, so I extracted my WebSocket logic into a custom hook called useWebSocket. It takes an onMessage callback so my Chat UI can decide how to handle incoming messages:

// Chat.tsx
const { sendMessage } = useWebSocket((msg) => {
  // If the message is from the person I'm actively chatting with, add it to the window
  if (selectedUser && msg.senderId === selectedUser.id) {
    addMessage(msg);
    return;
  }

  // Otherwise, show a toast notification
  toast(`${msg.sender.name}: ${msg.content}`);
});

Inside useWebSocket.ts, I wired up the actual WebSocket connection inside a useEffect:

// useWebSocket.ts (The broken version)
export function useWebSocket(onMessage: (msg: Message) => void) {
  const { user } = useAuth();
  useEffect(() => {
    const socket = new WebSocket(WS_URL);

    socket.onmessage = (e) => {
      const msg = JSON.parse(e.data);
      if (msg.type === "receive_message") {
        onMessage(msg.payload);
      }
    };
    return () => socket.close();
  }, [user]); // Only reconnect if the authenticated user changes
}

What Broke

At first glance, this code looks entirely reasonable. I'd open Relay, click on Alice's profile, and we'd start chatting. It worked flawlessly.

Then, I'd switch my active chat to Bob. Bob would send me a message, but instead of the message appearing in Bob's chat window, a toast notification would pop up. Even worse, if Alice sent me a message while I was looking at Bob's chat, her message silently injected itself directly into Bob's chat window!

I was literally watching the messages of one user show up in another user's chat after I switched.

Why It Broke (The Mental Model)

The culprit was a classic React gotcha: the stale closure.

When the useWebSocket hook ran for the first time, the useEffect established the WebSocket connection. The socket.onmessage event listener was created, and it permanently "captured" the onMessage function from that specific initial render.

In that first render, selectedUser was Alice.

When I clicked on Bob in the sidebar, React re-rendered my Chat component. A brand new onMessage function was created where selectedUser was Bob. However, because my useEffect dependency array only included user, the effect didn't re-run. The WebSocket stayed connected—which is what we want.

But this meant the socket.onmessage listener was still holding onto the very first onMessage function it captured—the one where selectedUser is permanently frozen as Alice.

If I simply added onMessage to the useEffect dependency array to fix the stale state, the WebSocket would brutally disconnect and reconnect every single time I clicked a new user, typed a character, or triggered a re-render. That's terrible UX and a massive performance killer.

The useRef Fix

I needed a way for the WebSocket listener to always access the freshest version of the onMessage callback, without actually triggering the useEffect to tear down the socket.

Enter useRef.

Unlike state, updating a ref doesn't trigger a re-render. More importantly, a ref is a mutable object that persists across renders. By continuously storing the latest callback in a ref, the WebSocket listener can just peek into ref.current and always find the most up-to-date function.

Here is the exact code that fixed the bug in Relay:

// useWebSocket.ts (The fixed version)
export function useWebSocket(onMessage: (msg: Message) => void) {
  const { user } = useAuth();

  // 1. Create a ref to hold the callback
  const onMessageRef = useRef(onMessage);
  // 2. Silently update the ref every time the callback changes (on every render)
  useEffect(() => {
    onMessageRef.current = onMessage;
  }, [onMessage]);
  useEffect(() => {
    const socket = new WebSocket(WS_URL);

    socket.onmessage = (e) => {
      const msg = JSON.parse(e.data);
      if (msg.type === "receive_message") {
        // 3. Always call the freshest callback stored in the ref!
        onMessageRef.current(msg.payload);
      }
    };
    return () => socket.close();
  }, [user]); 
}

The Takeaway

Whenever you integrate long-lived event listeners (like WebSockets, window events, or setInterval) with React components, you have to be paranoid about what variables are being captured in their closures. If you need access to constantly changing React state inside a permanent event listener, the useRef escape hatch is your best friend.

I'll be back with more next week. Until then, stay consistent!

Sidenote

Relay is now live: https://relay.nikshrma.dev

It's running on a DigitalOcean droplet and has been a great playground for learning WebSockets, authentication, testing, and deployment.

You can also find my portfolio at https://nikshrma.dev.