useWebSocket — A Reusable WebSocket Client Hook

Every component that uses a WebSocket needs the same boilerplate: open the connection, handle messages, reconnect on close, clean up on unmount. A useWebSocket hook extracts this lifecycle into one reusable function. The hook exposes a clean interface — the current connection status, a send function, and a lastMessage value — while hiding the reconnection logic, token attachment, and cleanup entirely from the components that use it.

useWebSocket Hook #

// src/hooks/useWebSocket.js
import { useState, useEffect, useRef, useCallback } from "react";
import { useAuthStore } from "@/stores/authStore";
import { config }       from "@/config";

const WS_BASE = config.apiBaseUrl
    .replace(/^http/, "ws")       // http: → ws:, https: → wss:
    .replace(/\/api$/, "");       // strip /api suffix

export const WS_STATUS = {
    CONNECTING:    "connecting",
    OPEN:          "open",
    CLOSING:       "closing",
    CLOSED:        "closed",
    RECONNECTING:  "reconnecting",
};

export function useWebSocket(path, { onMessage, enabled = true } = {}) {
    const [status,      setStatus]      = useState(WS_STATUS.CLOSED);
    const [lastMessage, setLastMessage] = useState(null);

    const wsRef         = useRef(null);
    const reconnectRef  = useRef(null);
    const attemptRef    = useRef(0);
    const mountedRef    = useRef(true);

    const accessToken   = useAuthStore((s) => s.accessToken);

    const connect = useCallback(() => {
        if (!enabled || !accessToken) return;

        // Build URL with auth token as query param
        const url = `${WS_BASE}${path}?token=${encodeURIComponent(accessToken)}`;

        setStatus(WS_STATUS.CONNECTING);
        const ws = new WebSocket(url);
        wsRef.current = ws;

        ws.onopen = () => {
            if (!mountedRef.current) return;
            setStatus(WS_STATUS.OPEN);
            attemptRef.current = 0;   // reset backoff on successful connect
        };

        ws.onmessage = (event) => {
            if (!mountedRef.current) return;
            try {
                const data = JSON.parse(event.data);
                setLastMessage(data);
                onMessage?.(data);
            } catch {
                // Non-JSON frame — ignore or handle as plain text
            }
        };

        ws.onclose = (event) => {
            if (!mountedRef.current) return;
            setStatus(WS_STATUS.CLOSED);

            // Only reconnect on unexpected closes (not code 1000 = intentional)
            if (event.code !== 1000 && enabled) {
                scheduleReconnect();
            }
        };

        ws.onerror = () => {
            // onerror always precedes onclose — let onclose handle reconnect
        };
    }, [path, accessToken, enabled, onMessage]);

    const scheduleReconnect = useCallback(() => {
        // Exponential backoff: 1s, 2s, 4s, 8s … cap at 30s
        const delay = Math.min(1000 * 2 ** attemptRef.current, 30_000);
        attemptRef.current += 1;
        setStatus(WS_STATUS.RECONNECTING);

        reconnectRef.current = setTimeout(() => {
            if (mountedRef.current) connect();
        }, delay);
    }, [connect]);

    const disconnect = useCallback(() => {
        clearTimeout(reconnectRef.current);
        wsRef.current?.close(1000, "Component unmounted");
        wsRef.current = null;
    }, []);

    const send = useCallback((data) => {
        if (wsRef.current?.readyState === WebSocket.OPEN) {
            wsRef.current.send(JSON.stringify(data));
        }
    }, []);

    // Connect on mount / when enabled/token change
    useEffect(() => {
        mountedRef.current = true;
        if (enabled && accessToken) connect();
        return () => {
            mountedRef.current = false;
            disconnect();
        };
    }, [enabled, accessToken]);   // eslint-disable-line react-hooks/exhaustive-deps

    return { status, lastMessage, send };
}

Usage Example #

import { useCallback } from "react";
import { useWebSocket, WS_STATUS } from "@/hooks/useWebSocket";
import { useAuthStore } from "@/stores/authStore";

function NotificationBell() {
    const isLoggedIn = useAuthStore((s) => Boolean(s.user));

    const handleMessage = useCallback((data) => {
        if (data.type === "notification") {
            // handle incoming notification
        }
    }, []);

    const { status, send } = useWebSocket("/ws/notifications", {
        onMessage: handleMessage,
        enabled:   isLoggedIn,
    });

    return (
        <div>
            <span className={`w-2 h-2 rounded-full ${
                status === WS_STATUS.OPEN         ? "bg-green-500"
                : status === WS_STATUS.RECONNECTING ? "bg-yellow-500"
                : "bg-gray-400"
            }`} />
        </div>
    );
}

Common Mistakes #

Mistake 1 — Inline onMessage causes reconnect loop #

❌ Wrong — new function reference on every render:

useWebSocket("/ws/feed", {
    onMessage: (data) => setMessages(m => [...m, data])   // new fn every render!
});

✅ Correct — stable reference with useCallback or a ref.

Mistake 2 — No reconnection on unexpected disconnect #

❌ Wrong — mobile network interruption silently kills the connection:

ws.onclose = () => setStatus("closed");   // no reconnect — user is stuck!

✅ Correct — reconnect on any close code other than 1000 (intentional).