lila/apps/web/src/lib/ws-provider.tsx

import { useCallback, useEffect, useState } from "react";
import type { ReactNode } from "react";
import { WsClient } from "./ws-client.js";
import { WsContext } from "./ws-context.js";

const wsClient = new WsClient();

export const WsProvider = ({ children }: { children: ReactNode }) => {
  const [isConnected, setIsConnected] = useState(false);

  const connect = useCallback(async (url: string): Promise<void> => {
    if (wsClient.isConnected()) return;

    wsClient.onClose = () => setIsConnected(false);
    wsClient.onError = () => setIsConnected(false);
    try {
      await wsClient.connect(url);
      setIsConnected(true);
    } catch (err) {
      setIsConnected(false);
      throw err;
    }
  }, []);

  const disconnect = useCallback((): void => {
    wsClient.disconnect();
    setIsConnected(false);
  }, []);

  useEffect(() => {
    return () => {
      wsClient.disconnect();
      wsClient.clearCallbacks();
      wsClient.onClose = null;
      wsClient.onError = null;
    };
  }, []);

  return (
    <WsContext.Provider
      value={{ client: wsClient, isConnected, connect, disconnect }}
    >
      {children}
    </WsContext.Provider>
  );
};