jiggablend/web/src/utils/websocket.js

// Shared WebSocket connection manager
// All components should use this instead of creating their own connections

class WebSocketManager {
  constructor() {
    this.ws = null;
    this.listeners = new Map(); // Map of listener IDs to callback functions
    this.reconnectTimeout = null;
    this.reconnectDelay = 2000;
    this.isConnecting = false;
    this.listenerIdCounter = 0;
    this.verboseLogging = false; // Set to true to enable verbose WebSocket logging
    
    // Track server-side channel subscriptions for re-subscription on reconnect
    this.serverSubscriptions = new Set(); // Channels we want to be subscribed to
    this.confirmedSubscriptions = new Set(); // Channels confirmed by server
    this.pendingSubscriptions = new Set(); // Channels waiting for confirmation
  }

  connect() {
    // If already connected or connecting, don't create a new connection
    if (this.ws && (this.ws.readyState === WebSocket.CONNECTING || this.ws.readyState === WebSocket.OPEN)) {
      return;
    }

    if (this.isConnecting) {
      return;
    }

    this.isConnecting = true;

    try {
      const wsProtocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
      const wsHost = window.location.host;
      const API_BASE = '/api';
      const url = `${wsProtocol}//${wsHost}${API_BASE}/jobs/ws`;
      
      this.ws = new WebSocket(url);

      this.ws.onopen = () => {
        if (this.verboseLogging) {
          console.log('Shared WebSocket connected');
        }
        this.isConnecting = false;
        
        // Re-subscribe to all channels that were previously subscribed
        this.resubscribeToChannels();
        
        this.notifyListeners('open', {});
      };

      this.ws.onmessage = (event) => {
        try {
          const data = JSON.parse(event.data);
          if (this.verboseLogging) {
            console.log('WebSocketManager: Message received:', data.type, data.channel || 'no channel', data);
          }
          this.notifyListeners('message', data);
        } catch (error) {
          console.error('WebSocketManager: Failed to parse message:', error, 'Raw data:', event.data);
        }
      };

      this.ws.onerror = (error) => {
        console.error('Shared WebSocket error:', error);
        this.isConnecting = false;
        this.notifyListeners('error', error);
      };

      this.ws.onclose = (event) => {
        if (this.verboseLogging) {
          console.log('Shared WebSocket closed:', {
            code: event.code,
            reason: event.reason,
            wasClean: event.wasClean
          });
        }
        this.ws = null;
        this.isConnecting = false;
        
        // Clear confirmed/pending but keep serverSubscriptions for re-subscription
        this.confirmedSubscriptions.clear();
        this.pendingSubscriptions.clear();
        
        this.notifyListeners('close', event);
        
        // Always retry connection if we have listeners
        if (this.listeners.size > 0) {
          if (this.reconnectTimeout) {
            clearTimeout(this.reconnectTimeout);
          }
          this.reconnectTimeout = setTimeout(() => {
            if (!this.ws || this.ws.readyState === WebSocket.CLOSED) {
              this.connect();
            }
          }, this.reconnectDelay);
        }
      };
    } catch (error) {
      console.error('Failed to create WebSocket:', error);
      this.isConnecting = false;
      // Retry after delay
      this.reconnectTimeout = setTimeout(() => {
        this.connect();
      }, this.reconnectDelay);
    }
  }

  subscribe(listenerId, callbacks) {
    // Generate ID if not provided
    if (!listenerId) {
      listenerId = `listener_${this.listenerIdCounter++}`;
    }

    if (this.verboseLogging) {
      console.log('WebSocketManager: Subscribing listener:', listenerId, 'WebSocket state:', this.ws ? this.ws.readyState : 'no connection');
    }
    this.listeners.set(listenerId, callbacks);

    // Connect if not already connected
    if (!this.ws || this.ws.readyState === WebSocket.CLOSED) {
      if (this.verboseLogging) {
        console.log('WebSocketManager: WebSocket not connected, connecting...');
      }
      this.connect();
    }

    // If already open, notify immediately
    if (this.ws && this.ws.readyState === WebSocket.OPEN && callbacks.open) {
      if (this.verboseLogging) {
        console.log('WebSocketManager: WebSocket already open, calling open callback for listener:', listenerId);
      }
      // Use setTimeout to ensure this happens after the listener is registered
      setTimeout(() => {
        if (callbacks.open) {
          callbacks.open();
        }
      }, 0);
    }

    return listenerId;
  }

  unsubscribe(listenerId) {
    this.listeners.delete(listenerId);
    
    // If no more listeners, we could close the connection, but let's keep it open
    // in case other components need it
  }

  send(data) {
    if (this.ws && this.ws.readyState === WebSocket.OPEN) {
      if (this.verboseLogging) {
        console.log('WebSocketManager: Sending message:', data);
      }
      this.ws.send(JSON.stringify(data));
    } else {
      console.warn('WebSocketManager: Cannot send message - connection not open. State:', this.ws ? this.ws.readyState : 'no connection', 'Message:', data);
    }
  }

  notifyListeners(eventType, data) {
    this.listeners.forEach((callbacks) => {
      if (callbacks[eventType]) {
        try {
          callbacks[eventType](data);
        } catch (error) {
          console.error('Error in WebSocket listener:', error);
        }
      }
    });
  }

  getReadyState() {
    return this.ws ? this.ws.readyState : WebSocket.CLOSED;
  }

  // Subscribe to a server-side channel (will be re-subscribed on reconnect)
  subscribeToChannel(channel) {
    if (this.serverSubscriptions.has(channel)) {
      // Already subscribed or pending
      return;
    }
    
    this.serverSubscriptions.add(channel);
    
    if (this.ws && this.ws.readyState === WebSocket.OPEN) {
      if (!this.confirmedSubscriptions.has(channel) && !this.pendingSubscriptions.has(channel)) {
        this.pendingSubscriptions.add(channel);
        this.send({ type: 'subscribe', channel });
        if (this.verboseLogging) {
          console.log('WebSocketManager: Subscribing to channel:', channel);
        }
      }
    }
  }

  // Unsubscribe from a server-side channel (won't be re-subscribed on reconnect)
  unsubscribeFromChannel(channel) {
    this.serverSubscriptions.delete(channel);
    this.confirmedSubscriptions.delete(channel);
    this.pendingSubscriptions.delete(channel);
    
    if (this.ws && this.ws.readyState === WebSocket.OPEN) {
      this.send({ type: 'unsubscribe', channel });
      if (this.verboseLogging) {
        console.log('WebSocketManager: Unsubscribing from channel:', channel);
      }
    }
  }

  // Mark a channel subscription as confirmed (call this when server confirms)
  confirmSubscription(channel) {
    this.pendingSubscriptions.delete(channel);
    this.confirmedSubscriptions.add(channel);
    if (this.verboseLogging) {
      console.log('WebSocketManager: Subscription confirmed for channel:', channel);
    }
  }

  // Mark a channel subscription as failed (call this when server rejects)
  failSubscription(channel) {
    this.pendingSubscriptions.delete(channel);
    this.serverSubscriptions.delete(channel);
    if (this.verboseLogging) {
      console.log('WebSocketManager: Subscription failed for channel:', channel);
    }
  }

  // Check if subscribed to a channel
  isSubscribedToChannel(channel) {
    return this.confirmedSubscriptions.has(channel);
  }

  // Re-subscribe to all channels after reconnect
  resubscribeToChannels() {
    if (this.serverSubscriptions.size === 0) {
      return;
    }
    
    if (this.verboseLogging) {
      console.log('WebSocketManager: Re-subscribing to channels:', Array.from(this.serverSubscriptions));
    }
    
    for (const channel of this.serverSubscriptions) {
      if (!this.pendingSubscriptions.has(channel)) {
        this.pendingSubscriptions.add(channel);
        this.send({ type: 'subscribe', channel });
      }
    }
  }

  disconnect() {
    if (this.reconnectTimeout) {
      clearTimeout(this.reconnectTimeout);
      this.reconnectTimeout = null;
    }
    if (this.ws) {
      this.ws.close();
      this.ws = null;
    }
    this.listeners.clear();
    this.serverSubscriptions.clear();
    this.confirmedSubscriptions.clear();
    this.pendingSubscriptions.clear();
  }
}

// Export singleton instance
export const wsManager = new WebSocketManager();