OBS-overlay/js/websocket-client.js

/**
 * WebSocket client for Jackbox Game Picker API: JWT auth, live session stream, reconnect, heartbeat.
 */

const MAX_RECONNECT_DELAY_MS = 30_000;
const HEARTBEAT_INTERVAL_MS = 30_000;
const INITIAL_RECONNECT_DELAY_MS = 1_000;

/**
 * @typedef {'connecting'|'connected'|'disconnected'|'error'} WsConnectionState
 */

class WebSocketClient {
  constructor(options = {}) {
    const {
      onStatusChange = () => {},
      onEvent = () => {},
      onSessionSubscribed = () => {},
    } = options;

    /** @type {(state: WsConnectionState, message?: string) => void} */
    this._onStatusChange = onStatusChange;
    /** @type {(eventType: string, data: unknown) => void} */
    this._onEvent = onEvent;
    /** @type {(sessionId: string | number) => void} */
    this._onSessionSubscribed = onSessionSubscribed;

    /** @type {WebSocket | null} */
    this._ws = null;
    /** @type {string | null} */
    this._jwtToken = null;
    /** @type {string | null} */
    this._apiUrl = null;
    /** @type {ReturnType<typeof setInterval> | null} */
    this._heartbeatInterval = null;
    /** @type {ReturnType<typeof setTimeout> | null} */
    this._reconnectTimeout = null;
    this._reconnectDelay = INITIAL_RECONNECT_DELAY_MS;
    this._intentionalDisconnect = false;
    /** @type {boolean} */
    this._authenticated = false;
  }

  /**
   * @param {string} apiUrl
   * @param {string} apiKey
   */
  async connect(apiUrl, apiKey) {
    const base = this._normalizeApiUrl(apiUrl);
    const key = String(apiKey).trim();

    if (!base) {
      this._onStatusChange('error', 'API URL is required');
      return;
    }
    if (!key) {
      this._onStatusChange('error', 'API Key is required');
      return;
    }

    this._intentionalDisconnect = false;
    this._apiUrl = base;
    this._onStatusChange('connecting', 'Authenticating...');

    try {
      const response = await fetch(`${base}/api/auth/login`, {
        method: 'POST',
        headers: { 'Content-Type': 'application/json' },
        body: JSON.stringify({ key }),
      });

      if (!response.ok) {
        const errData = await response.json().catch(() => ({}));
        const msg =
          typeof errData.error === 'string'
            ? errData.error
            : response.statusText || String(response.status);
        this._onStatusChange('error', `Auth failed: ${msg}`);
        return;
      }

      const data = await response.json();
      const token = data.token;

      if (!token) {
        this._onStatusChange('error', 'No token in auth response');
        return;
      }

      this._jwtToken = token;
      this._onStatusChange('connecting', 'Connecting WebSocket...');
      this._connectWebSocket(base);
    } catch (err) {
      console.error('[WebSocketClient] Auth error:', err);
      const message = err instanceof Error ? err.message : String(err);
      this._onStatusChange('error', `Auth error: ${message}`);
    }
  }

  disconnect() {
    this._intentionalDisconnect = true;
    if (this._reconnectTimeout) {
      clearTimeout(this._reconnectTimeout);
      this._reconnectTimeout = null;
    }
    this._stopHeartbeat();
    this._authenticated = false;
    this._jwtToken = null;
    if (this._ws) {
      this._ws.close();
      this._ws = null;
    }
    this._onStatusChange('disconnected', 'Disconnected');
  }

  /**
   * @param {string | number} sessionId
   */
  subscribeToSession(sessionId) {
    if (!this._ws || this._ws.readyState !== WebSocket.OPEN) {
      console.warn('[WebSocketClient] subscribeToSession: socket not open');
      return;
    }
    this._ws.send(
      JSON.stringify({ type: 'subscribe', sessionId }),
    );
  }

  get isConnected() {
    return (
      this._authenticated &&
      this._ws !== null &&
      this._ws.readyState === WebSocket.OPEN
    );
  }

  /**
   * @param {string} apiUrl
   */
  _normalizeApiUrl(apiUrl) {
    return String(apiUrl).trim().replace(/\/+$/, '');
  }

  /**
   * @param {string} apiUrl
   */
  _connectWebSocket(apiUrl) {
    const wsUrl = `${apiUrl.replace(/^http/, 'ws')}/api/sessions/live`;

    try {
      this._ws = new WebSocket(wsUrl);
    } catch (err) {
      const message = err instanceof Error ? err.message : String(err);
      this._onStatusChange('error', `WebSocket error: ${message}`);
      return;
    }

    this._ws.addEventListener('open', () => {
      this._authenticated = false;
      this._onStatusChange('connecting', 'Authenticating via WebSocket...');
      this._reconnectDelay = INITIAL_RECONNECT_DELAY_MS;

      if (!this._jwtToken) {
        this._onStatusChange('error', 'Missing JWT for WebSocket auth');
        return;
      }
      this._ws.send(JSON.stringify({ type: 'auth', token: this._jwtToken }));
    });

    this._ws.addEventListener('message', (event) => {
      let message;
      try {
        message = JSON.parse(event.data);
      } catch (e) {
        console.error('[WebSocketClient] Failed to parse message:', event.data);
        return;
      }
      this._handleMessage(message);
    });

    this._ws.addEventListener('close', (event) => {
      console.log('[WebSocketClient] Disconnected, code:', event.code);
      this._stopHeartbeat();
      this._authenticated = false;
      this._ws = null;

      if (!this._intentionalDisconnect) {
        const secs = Math.round(this._reconnectDelay / 1000);
        this._onStatusChange(
          'connecting',
          `Reconnecting in ${secs}s...`,
        );
        this._scheduleReconnect();
      } else {
        this._onStatusChange('disconnected', 'Disconnected');
      }
    });

    this._ws.addEventListener('error', (err) => {
      console.error('[WebSocketClient] Socket error:', err);
    });
  }

  /**
   * @param {Record<string, unknown>} message
   */
  _handleMessage(message) {
    const type = message.type;

    switch (type) {
      case 'auth_success': {
        this._startHeartbeat();
        void this._fetchActiveSessionAndSubscribe();
        this._authenticated = true;
        this._onStatusChange('connected', 'Connected');
        break;
      }

      case 'auth_error': {
        const msg =
          typeof message.message === 'string'
            ? message.message
            : 'WebSocket authentication failed';
        console.error('[WebSocketClient] Auth error:', msg);
        this._onStatusChange('error', `WS auth failed: ${msg}`);
        this._intentionalDisconnect = true;
        if (this._ws) {
          this._ws.close();
        }
        break;
      }

      case 'subscribed': {
        const sessionId = message.sessionId;
        const label =
          sessionId !== undefined && sessionId !== null
            ? `Connected (session ${sessionId})`
            : 'Connected (subscribed)';
        this._onStatusChange('connected', label);
        if (sessionId !== undefined && sessionId !== null) {
          this._onSessionSubscribed(sessionId);
        }
        break;
      }

      case 'pong':
        break;

      case 'error': {
        const serverMsg =
          typeof message.message === 'string'
            ? message.message
            : JSON.stringify(message);
        console.error('[WebSocketClient] Server error:', serverMsg);
        break;
      }

      case 'session.started': {
        const data = message.data;
        /** @type {{ session?: { id?: string | number } } | undefined} */
        let payload;
        if (data && typeof data === 'object' && data !== null) {
          payload = /** @type {{ session?: { id?: string | number } }} */ (
            data
          );
        }
        const id = payload?.session?.id;

        if (
          id !== undefined &&
          id !== null &&
          this._ws &&
          this._ws.readyState === WebSocket.OPEN
        ) {
          this._ws.send(
            JSON.stringify({ type: 'subscribe', sessionId: id }),
          );
        }
        if (data !== undefined) {
          this._onEvent('session.started', data);
        }
        break;
      }

      default: {
        if (
          Object.prototype.hasOwnProperty.call(message, 'data') &&
          message.data !== undefined
        ) {
          if (typeof type === 'string') {
            this._onEvent(type, message.data);
          }
        } else {
          console.log('[WebSocketClient] Unhandled message type:', type);
        }
      }
    }
  }

  async _fetchActiveSessionAndSubscribe() {
    const apiUrl = this._apiUrl;
    const token = this._jwtToken;
    if (!apiUrl || !token) return;

    try {
      const response = await fetch(`${apiUrl}/api/sessions/active`, {
        headers: { Authorization: `Bearer ${token}` },
      });

      if (!response.ok) {
        console.log(
          '[WebSocketClient] Could not fetch active session:',
          response.status,
        );
        return;
      }

      const session = await response.json();
      if (!session || session.id === undefined || session.id === null) {
        console.log(
          '[WebSocketClient] No active session found; waiting for session.started',
        );
        return;
      }

      const sessionId = session.id;
      console.log(
        '[WebSocketClient] Found active session:',
        sessionId,
        '— subscribing',
      );
      this.subscribeToSession(sessionId);

      await this._restoreCurrentGame(sessionId);
    } catch (err) {
      console.error('[WebSocketClient] Error fetching active session:', err);
    }
  }

  /**
   * After subscribing, fetch the session's games to find the currently playing
   * game, then hit status-live for full shard state (players, lobby, etc.).
   * @param {string | number} sessionId
   */
  async _restoreCurrentGame(sessionId) {
    const apiUrl = this._apiUrl;
    const token = this._jwtToken;
    if (!apiUrl || !token) return;

    try {
      const gamesRes = await fetch(`${apiUrl}/api/sessions/${sessionId}/games`, {
        headers: { Authorization: `Bearer ${token}` },
      });

      if (!gamesRes.ok) {
        console.log('[WebSocketClient] Could not fetch session games:', gamesRes.status);
        return;
      }

      const games = await gamesRes.json();
      if (!Array.isArray(games) || games.length === 0) return;

      const playing = games.find((g) => g.status === 'playing');
      if (!playing) return;

      console.log(
        '[WebSocketClient] Restoring in-progress game:',
        playing.title,
        '(room:', playing.room_code, ')',
      );

      this._onEvent('game.added', {
        session: { id: sessionId },
        game: {
          id: playing.game_id,
          title: playing.title,
          pack_name: playing.pack_name,
          min_players: playing.min_players,
          max_players: playing.max_players,
          room_code: playing.room_code,
          manually_added: playing.manually_added,
        },
      });

      const sessionGameId = playing.id;
      const statusRes = await fetch(
        `${apiUrl}/api/sessions/${sessionId}/games/${sessionGameId}/status-live`,
      );

      if (statusRes.ok) {
        const live = await statusRes.json();
        console.log(
          '[WebSocketClient] Restored live status — players:',
          live.playerCount,
          'state:', live.gameState,
        );

        this._onEvent('game.status', {
          sessionId,
          gameId: live.gameId ?? playing.game_id,
          roomCode: live.roomCode ?? playing.room_code,
          appTag: live.appTag,
          maxPlayers: live.maxPlayers ?? playing.max_players,
          playerCount: live.playerCount ?? playing.player_count,
          players: Array.isArray(live.players) ? live.players : [],
          lobbyState: live.lobbyState,
          gameState: live.gameState,
          gameStarted: live.gameStarted,
          gameFinished: live.gameFinished,
          monitoring: live.monitoring,
        });
      } else if (playing.player_count != null && playing.max_players != null) {
        this._onEvent('room.connected', {
          sessionId,
          gameId: playing.game_id,
          roomCode: playing.room_code,
          maxPlayers: playing.max_players,
          playerCount: playing.player_count,
          players: [],
          lobbyState: 'Unknown',
          gameState: 'Unknown',
        });
      }
    } catch (err) {
      console.error('[WebSocketClient] Error restoring current game:', err);
    }
  }

  _startHeartbeat() {
    this._stopHeartbeat();
    this._heartbeatInterval = setInterval(() => {
      if (this._ws && this._ws.readyState === WebSocket.OPEN) {
        this._ws.send(JSON.stringify({ type: 'ping' }));
      }
    }, HEARTBEAT_INTERVAL_MS);
  }

  _stopHeartbeat() {
    if (this._heartbeatInterval) {
      clearInterval(this._heartbeatInterval);
      this._heartbeatInterval = null;
    }
  }

  _scheduleReconnect() {
    if (this._reconnectTimeout) {
      clearTimeout(this._reconnectTimeout);
    }

    const delay = this._reconnectDelay;
    this._reconnectTimeout = setTimeout(() => {
      this._reconnectTimeout = null;
      const apiUrl = this._apiUrl;
      const token = this._jwtToken;

      if (apiUrl && token && !this._intentionalDisconnect) {
        console.log('[WebSocketClient] Attempting reconnect...');
        this._onStatusChange('connecting', 'Reconnecting...');
        this._connectWebSocket(apiUrl);
      }

      this._reconnectDelay = Math.min(
        this._reconnectDelay * 2,
        MAX_RECONNECT_DELAY_MS,
      );
    }, delay);
  }
}

window.OBS = window.OBS || {};
window.OBS.WebSocketClient = WebSocketClient;