All files / src/hooks useWebSocket.ts

2.06% Statements 2/97
100% Branches 0/0
0% Functions 0/1
2.06% Lines 2/97

Press n or j to go to the next uncovered block, b, p or k for the previous block.

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 1321x                 1x                                                                                                                                                                                                                                                    
import { useEffect, useRef, useCallback, useState } from 'react';
import { useQueryClient } from '@tanstack/react-query';
 
interface WebSocketMessage {
  type: string;
  printer_id?: number;
  data?: Record<string, unknown>;
}
 
export function useWebSocket() {
  const wsRef = useRef<WebSocket | null>(null);
  const reconnectTimeoutRef = useRef<number | null>(null);
  const queryClient = useQueryClient();
  const [isConnected, setIsConnected] = useState(false);
 
  const connect = useCallback(() => {
    if (wsRef.current?.readyState === WebSocket.OPEN) {
      return;
    }
 
    const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
    const wsUrl = `${protocol}//${window.location.host}/api/v1/ws`;
 
    const ws = new WebSocket(wsUrl);
 
    let pingInterval: number | null = null;
 
    ws.onopen = () => {
      console.log('[WebSocket] Connected');
      setIsConnected(true);
      // Start ping interval
      pingInterval = window.setInterval(() => {
        if (ws.readyState === WebSocket.OPEN) {
          ws.send(JSON.stringify({ type: 'ping' }));
        }
      }, 30000);
    };
 
    ws.onmessage = (event) => {
      try {
        const message: WebSocketMessage = JSON.parse(event.data);
        handleMessage(message);
      } catch {
        // Ignore parse errors
      }
    };
 
    ws.onclose = (event) => {
      console.log('[WebSocket] Closed', event.code, event.reason);
      if (pingInterval) {
        clearInterval(pingInterval);
        pingInterval = null;
      }
      setIsConnected(false);
      wsRef.current = null;
 
      // Reconnect after 3 seconds
      reconnectTimeoutRef.current = window.setTimeout(() => {
        connect();
      }, 3000);
    };
 
    ws.onerror = (error) => {
      console.error('[WebSocket] Error', error);
      ws.close();
    };
 
    wsRef.current = ws;
  }, []);
 
  const handleMessage = useCallback((message: WebSocketMessage) => {
    switch (message.type) {
      case 'printer_status':
        // Update the printer status in the query cache
        if (message.printer_id !== undefined) {
          queryClient.setQueryData(
            ['printerStatus', message.printer_id],
            (old: Record<string, unknown> | undefined) => {
              const merged = {
                ...old,
                ...message.data,
              };
              // Preserve last known wifi_signal if new value is null
              if (merged.wifi_signal == null && old?.wifi_signal != null) {
                merged.wifi_signal = old.wifi_signal;
              }
              return merged;
            }
          );
        }
        break;
 
      case 'print_complete':
        // Invalidate archives to refresh the list
        queryClient.invalidateQueries({ queryKey: ['archives'] });
        queryClient.invalidateQueries({ queryKey: ['archiveStats'] });
        break;
 
      case 'archive_created':
        // Invalidate archives to show new archive
        queryClient.invalidateQueries({ queryKey: ['archives'] });
        queryClient.invalidateQueries({ queryKey: ['archiveStats'] });
        break;
 
      case 'pong':
        // Keepalive response, ignore
        break;
    }
  }, [queryClient]);
 
  useEffect(() => {
    connect();
 
    return () => {
      if (reconnectTimeoutRef.current) {
        clearTimeout(reconnectTimeoutRef.current);
      }
      if (wsRef.current) {
        wsRef.current.close();
      }
    };
  }, [connect]);
 
  const sendMessage = useCallback((message: Record<string, unknown>) => {
    if (wsRef.current?.readyState === WebSocket.OPEN) {
      wsRef.current.send(JSON.stringify(message));
    }
  }, []);
 
  return { isConnected, sendMessage };
}