clawtap/src/lib/ws.ts

import { WS } from './ws-types';

export type WsStatus = 'connecting' | 'connected' | 'disconnected' | 'reconnecting';

type MessageHandler = (msg: any) => void;
type StatusHandler = (status: WsStatus) => void;

/** Minimum hidden duration (ms) before forcing reconnect on visibility change. */
const VISIBILITY_RECONNECT_THRESHOLD = 3000;

export class WsClient {
  private ws: WebSocket | null = null;
  private url: string;
  private onMessage: MessageHandler;
  private onStatus: StatusHandler;
  private reconnectDelay = 1000;
  private maxReconnectDelay = 30000;
  private shouldReconnect = true;
  private activeSessionId: string | null = null;
  private activeAdapter: string | null = null;
  private visibilityHandler: (() => void) | null = null;
  private hiddenSince: number | null = null;
  private openedThisAttempt = false;
  private onAuthCheckNeeded: (() => void) | null = null;

  constructor(token: string, onMessage: MessageHandler, onStatus: StatusHandler, onAuthCheckNeeded?: () => void) {
    const proto = location.protocol === 'https:' ? 'wss' : 'ws';
    this.url = `${proto}://${location.host}/ws?token=${encodeURIComponent(token)}`;
    this.onMessage = onMessage;
    this.onStatus = onStatus;
    this.onAuthCheckNeeded = onAuthCheckNeeded || null;
  }

  connect() {
    this.shouldReconnect = true;
    this.openedThisAttempt = false;
    this.onStatus('connecting');

    this.ws = new WebSocket(this.url);

    this.ws.onopen = () => {
      this.openedThisAttempt = true;
      this.reconnectDelay = 1000;
      this.onStatus('connected');

      if (this.activeSessionId) {
        this.send({ type: WS.RECONNECT, sessionId: this.activeSessionId, adapter: this.activeAdapter });
      }
      // Tell server our current visibility state immediately on (re)connect
      this.send({ type: WS.PAGE_VISIBILITY, visible: document.visibilityState === 'visible' });
    };

    this.ws.onmessage = (event) => {
      try {
        const msg = JSON.parse(event.data);
        if (msg.type === WS.SESSION_CREATED) {
          this.activeSessionId = msg.sessionId;
          if (msg.adapter) this.activeAdapter = msg.adapter;
        }
        // Respond to server app-ping only when page is visible.
        // If hidden, no pong → server fires push after 2s timeout.
        if (msg.type === WS.APP_PING) {
          if (document.visibilityState === 'visible') {
            this.send({ type: WS.APP_PONG, sessionId: msg.sessionId });
          }
          return;
        }
        this.onMessage(msg);
      } catch (err) {
        console.error('[ws] Failed to parse message:', err);
      }
    };

    this.ws.onclose = () => {
      this.onStatus('disconnected');
      // The handshake itself failed (e.g. server rejected the auth token) —
      // check whether our token is still valid rather than reconnecting forever.
      if (!this.openedThisAttempt) {
        this.onAuthCheckNeeded?.();
      }
      if (this.shouldReconnect) {
        this.onStatus('reconnecting');
        setTimeout(() => this.connect(), this.reconnectDelay);
        this.reconnectDelay = Math.min(this.reconnectDelay * 2, this.maxReconnectDelay);
      }
    };

    this.ws.onerror = () => {
      this.ws?.close();
    };

    this._startVisibilityWatch();
  }

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

  setActiveSession(sessionId: string | null, adapter?: string | null) {
    this.activeSessionId = sessionId;
    this.activeAdapter = adapter || null;
  }

  /** Force close and reconnect immediately (reset backoff). */
  forceReconnect() {
    if (!this.shouldReconnect) return;
    this.reconnectDelay = 1000;
    this.hiddenSince = null;
    const old = this.ws;
    this.ws = null;
    if (old) {
      old.onclose = null;
      old.onerror = null;
      old.onmessage = null;
      old.close();
    }
    this.connect();
  }

  disconnect() {
    this.shouldReconnect = false;
    this._stopVisibilityWatch();
    this.ws?.close();
    this.ws = null;
  }

  private _startVisibilityWatch() {
    if (this.visibilityHandler) return;
    this.visibilityHandler = () => {
      const visible = document.visibilityState === 'visible';
      this.send({ type: WS.PAGE_VISIBILITY, visible });
      if (!visible) {
        this.hiddenSince = Date.now();
      } else if (this.shouldReconnect) {
        const elapsed = this.hiddenSince ? Date.now() - this.hiddenSince : 0;
        this.hiddenSince = null;
        if (elapsed >= VISIBILITY_RECONNECT_THRESHOLD) {
          this.forceReconnect();
        }
      }
    };
    document.addEventListener('visibilitychange', this.visibilityHandler);
  }

  private _stopVisibilityWatch() {
    if (this.visibilityHandler) {
      document.removeEventListener('visibilitychange', this.visibilityHandler);
      this.visibilityHandler = null;
    }
  }
}