cmux-remote/public/js/websocket-client.js

class WebSocketClient {
  constructor() {
    this.ws = null;
    this.handlers = new Map();
    this.reconnectDelay = 1000;
    this.maxReconnectDelay = 30000;
    this.currentDelay = this.reconnectDelay;
    this.shouldReconnect = true;
    this.connected = false;
    this.reconnectAttempt = 0;
    this._heartbeatCheck = null;
    this._lastServerPing = 0;
  }

  connect(authRequired) {
    this.authRequired = authRequired !== false;
    const token = localStorage.getItem('cmux-remote-token');
    if (this.authRequired && !token) {
      window.location.href = '/login.html';
      return;
    }

    const proto = location.protocol === 'https:' ? 'wss:' : 'ws:';
    const url = `${proto}//${location.host}`;

    this.ws = new WebSocket(url);

    this.ws.onopen = () => {
      this.connected = true;
      this.reconnectAttempt = 0;
      this.currentDelay = this.reconnectDelay;
      this._lastServerPing = Date.now();
      this._startHeartbeatCheck();
      this.emit('connected');

      // Authenticate (send token or empty for no-auth mode)
      this.send({ type: 'auth', token: token || '' });
    };

    this.ws.onmessage = (event) => {
      try {
        const msg = JSON.parse(event.data);
        // Any inbound message proves the connection is alive — not just pings.
        // Otherwise a busy stream of screen-diffs with delayed pings could trip
        // the 15s heartbeat and force a needless close/reconnect (screen flash).
        this._lastServerPing = Date.now();
        if (msg.type === 'ping') {
          this.send({ type: 'pong' });
          return;
        }
        this.emit(msg.type, msg);

        // Redirect to login on auth failure (only if auth is required)
        if (msg.type === 'auth-fail' && this.authRequired) {
          localStorage.removeItem('cmux-remote-token');
          window.location.href = '/login.html';
        }
      } catch (err) {
        console.error('[ws] Parse error:', err);
      }
    };

    this.ws.onclose = () => {
      this.connected = false;
      this._stopHeartbeatCheck();
      this.emit('disconnected');

      if (this.shouldReconnect) {
        const attempt = ++this.reconnectAttempt;
        const delay = this.currentDelay;
        this.emit('reconnecting', { attempt, delay });
        setTimeout(() => this.connect(this.authRequired), delay);
        this.currentDelay = Math.min(this.currentDelay * 1.5, this.maxReconnectDelay);
      }
    };

    this.ws.onerror = (err) => {
      console.error('[ws] Error:', err);
    };
  }

  send(msg) {
    if (this.ws && this.ws.readyState === WebSocket.OPEN) {
      this.ws.send(JSON.stringify(msg));
    }
  }

  on(type, handler) {
    if (!this.handlers.has(type)) {
      this.handlers.set(type, []);
    }
    this.handlers.get(type).push(handler);
  }

  off(type, handler) {
    const list = this.handlers.get(type);
    if (list) {
      const idx = list.indexOf(handler);
      if (idx >= 0) list.splice(idx, 1);
    }
  }

  emit(type, data) {
    const list = this.handlers.get(type);
    if (list) {
      for (const handler of list) {
        handler(data);
      }
    }
  }

  subscribe(workspace, surface) {
    this.send({ type: 'subscribe', workspace, surface });
  }

  unsubscribe(surface) {
    this.send({ type: 'unsubscribe', surface });
  }

  sendText(workspace, surface, text) {
    this.send({ type: 'send-text', workspace, surface, text });
  }

  // Send text and have the SERVER press Enter once send_text completes (awaited
  // ordering). Reliable submit without a client-side Enter timing race.
  sendTextSubmit(workspace, surface, text) {
    this.send({ type: 'send-text', workspace, surface, text, submit: true });
  }

  sendKey(workspace, surface, key) {
    this.send({ type: 'send-key', workspace, surface, key });
  }

  listWorkspaces() {
    this.send({ type: 'list-workspaces' });
  }

  requestScroll(workspace, surface, lines) {
    this.send({ type: 'scroll-request', workspace, surface, lines });
  }

  _startHeartbeatCheck() {
    this._stopHeartbeatCheck();
    this._heartbeatCheck = setInterval(() => {
      if (Date.now() - this._lastServerPing > 15000) {
        this.ws?.close();
      }
    }, 5000);
  }

  _stopHeartbeatCheck() {
    if (this._heartbeatCheck) {
      clearInterval(this._heartbeatCheck);
      this._heartbeatCheck = null;
    }
  }

  // Force immediate reconnection (reset backoff, close stale socket)
  reconnectNow() {
    this._stopHeartbeatCheck();
    if (this.ws) {
      // Temporarily disable auto-reconnect so onclose doesn't double-connect
      this.shouldReconnect = false;
      this.ws.close();
    }
    this.shouldReconnect = true;
    this.currentDelay = this.reconnectDelay;
    this.reconnectAttempt = 0;
    this.connect(this.authRequired);
  }

  disconnect() {
    this._stopHeartbeatCheck();
    this.shouldReconnect = false;
    if (this.ws) this.ws.close();
  }

  isConnected() {
    return this.connected;
  }
}

window.WebSocketClient = WebSocketClient;