supervisor-ui/lib/hooks/useEventSource.ts

import { useEffect, useRef, useState, useCallback } from 'react';
import { createHookLogger } from '@/lib/utils/client-logger';

const logger = createHookLogger('useEventSource');

export type ConnectionStatus = 'connecting' | 'connected' | 'disconnected' | 'error';

export interface EventSourceMessage<T = any> {
  event: string;
  data: T;
}

export interface UseEventSourceOptions {
  enabled?: boolean;
  reconnectInterval?: number;
  maxReconnectAttempts?: number;
  onMessage?: (message: EventSourceMessage) => void;
  onError?: (error: Event) => void;
  onConnect?: () => void;
  onDisconnect?: () => void;
}

export function useEventSource(url: string, options: UseEventSourceOptions = {}) {
  const {
    enabled = true,
    reconnectInterval = 3000,
    maxReconnectAttempts = 10,
    onMessage,
    onError,
    onConnect,
    onDisconnect,
  } = options;

  const [status, setStatus] = useState<ConnectionStatus>('disconnected');
  const [reconnectAttempts, setReconnectAttempts] = useState(0);
  const eventSourceRef = useRef<EventSource | null>(null);
  const reconnectTimeoutRef = useRef<NodeJS.Timeout | null>(null);

  const connect = useCallback(() => {
    if (!enabled || eventSourceRef.current) return;

    logger.info('Connecting to SSE', { url });
    setStatus('connecting');

    try {
      const eventSource = new EventSource(url);
      eventSourceRef.current = eventSource;

      eventSource.addEventListener('connected', () => {
        logger.info('SSE connected successfully', { url });
        setStatus('connected');
        setReconnectAttempts(0);
        onConnect?.();
      });

      eventSource.addEventListener('heartbeat', (event) => {
        // Keep connection alive
        if (status !== 'connected') {
          logger.debug('SSE heartbeat received, updating status to connected');
          setStatus('connected');
        }
      });

      eventSource.addEventListener('process-update', (event) => {
        try {
          const data = JSON.parse(event.data);
          logger.debug('Process update received', {
            processCount: data.processes?.length,
            timestamp: data.timestamp,
          });
          onMessage?.({ event: 'process-update', data });
        } catch (error) {
          logger.error('Failed to parse SSE message', error, { event: event.data });
        }
      });

      eventSource.addEventListener('error', (event) => {
        try {
          const data = JSON.parse((event as MessageEvent).data);
          logger.warn('SSE error event received', { error: data });
          onMessage?.({ event: 'error', data });
        } catch (error) {
          // Not a message error, connection error
        }
      });

      eventSource.onerror = (event) => {
        logger.error('EventSource connection error', event);
        setStatus('error');
        onError?.(event);

        // Close current connection
        eventSource.close();
        eventSourceRef.current = null;

        // Attempt reconnection with exponential backoff
        if (reconnectAttempts < maxReconnectAttempts) {
          const delay = Math.min(reconnectInterval * Math.pow(2, reconnectAttempts), 30000);
          logger.info('Scheduling reconnection', {
            delay,
            attempt: reconnectAttempts + 1,
            maxAttempts: maxReconnectAttempts,
          });

          reconnectTimeoutRef.current = setTimeout(() => {
            setReconnectAttempts((prev) => prev + 1);
            connect();
          }, delay);
        } else {
          logger.warn('Max reconnection attempts reached, disconnecting', {
            maxAttempts: maxReconnectAttempts,
          });
          setStatus('disconnected');
          onDisconnect?.();
        }
      };
    } catch (error) {
      logger.error('Failed to create EventSource', error, { url });
      setStatus('error');
    }
  }, [url, enabled, status, reconnectAttempts, maxReconnectAttempts, reconnectInterval, onMessage, onError, onConnect, onDisconnect]);

  const disconnect = useCallback(() => {
    logger.info('Disconnecting from SSE');

    if (reconnectTimeoutRef.current) {
      clearTimeout(reconnectTimeoutRef.current);
      reconnectTimeoutRef.current = null;
    }

    if (eventSourceRef.current) {
      eventSourceRef.current.close();
      eventSourceRef.current = null;
    }

    setStatus('disconnected');
    setReconnectAttempts(0);
  }, []);

  const reconnect = useCallback(() => {
    logger.info('Manual reconnection requested');
    disconnect();
    setReconnectAttempts(0);
    connect();
  }, [disconnect, connect]);

  useEffect(() => {
    if (enabled) {
      connect();
    } else {
      disconnect();
    }

    return () => {
      disconnect();
    };
  }, [enabled, url]);

  return {
    status,
    reconnectAttempts,
    reconnect,
    disconnect,
  };
}