flynn/src/gateway/ui/lib/ws-client.js

/**
 * Flynn WebSocket RPC Client
 *
 * Promise-based JSON-RPC client with auto-reconnect, event streaming,
 * and connection lifecycle management.
 */
export class FlynnClient {
  constructor(url) {
    this._url = url || `ws://${location.host}`;
    this._ws = null;
    this._requestId = 0;
    this._pending = new Map();     // id -> { resolve, reject }
    this._listeners = new Map();   // id -> { events: Map<event, callback[]> }
    this._reconnectDelay = 1000;
    this._maxReconnectDelay = 30000;
    this._onStatusChange = null;
    this._status = 'disconnected';
    this._autoReconnect = true;
  }

  get status() { return this._status; }

  onStatusChange(callback) {
    this._onStatusChange = callback;
  }

  connect() {
    this._autoReconnect = true;
    this._doConnect();
  }

  disconnect() {
    this._autoReconnect = false;
    if (this._ws) {
      this._ws.close();
      this._ws = null;
    }
    this._setStatus('disconnected');
  }

  _doConnect() {
    this._setStatus('connecting');

    // Build URL with token from URL search params if present
    let wsUrl = this._url;
    const urlParams = new URLSearchParams(window.location.search);
    const token = urlParams.get('token');
    if (token) {
      const sep = wsUrl.includes('?') ? '&' : '?';
      wsUrl = `${wsUrl}${sep}token=${encodeURIComponent(token)}`;
    }

    this._ws = new WebSocket(wsUrl);

    this._ws.onopen = () => {
      this._setStatus('connected');
      this._reconnectDelay = 1000;
    };

    this._ws.onmessage = (event) => {
      this._handleMessage(event.data);
    };

    this._ws.onclose = (event) => {
      this._ws = null;
      // Gateway lock — show specific message and don't auto-reconnect
      if (event.code === 4003) {
        this._setStatus('locked');
        this._autoReconnect = false;
        return;
      }
      this._setStatus('disconnected');
      // Reject all pending requests
      for (const [id, pending] of this._pending) {
        pending.reject(new Error('WebSocket closed'));
      }
      this._pending.clear();
      this._listeners.clear();

      if (this._autoReconnect) {
        setTimeout(() => this._doConnect(), this._reconnectDelay);
        this._reconnectDelay = Math.min(this._reconnectDelay * 2, this._maxReconnectDelay);
      }
    };

    this._ws.onerror = () => {
      // Error is always followed by close
    };
  }

  _setStatus(status) {
    if (this._status !== status) {
      this._status = status;
      this._onStatusChange?.(status);
    }
  }

  _handleMessage(raw) {
    let msg;
    try {
      msg = JSON.parse(raw);
    } catch {
      return;
    }

    // Streamed event (has 'event' field)
    if (msg.event && msg.id != null) {
      const listener = this._listeners.get(msg.id);
      if (listener) {
        const callbacks = listener.events.get(msg.event) || [];
        for (const cb of callbacks) {
          cb(msg.data);
        }
        // Also fire wildcard listeners
        const wildcards = listener.events.get('*') || [];
        for (const cb of wildcards) {
          cb(msg.event, msg.data);
        }
      }
      return;
    }

    // Response or error (matches pending request)
    if (msg.id != null && this._pending.has(msg.id)) {
      const pending = this._pending.get(msg.id);
      this._pending.delete(msg.id);

      if (msg.error) {
        pending.reject(new Error(msg.error.message));
      } else {
        pending.resolve(msg.result);
      }
    }
  }

  /**
   * Send an RPC call and return a promise for the result.
   * For streaming methods (like agent.send), use stream() instead.
   */
  async call(method, params) {
    if (!this._ws || this._ws.readyState !== WebSocket.OPEN) {
      throw new Error('Not connected');
    }

    const id = ++this._requestId;
    return new Promise((resolve, reject) => {
      this._pending.set(id, { resolve, reject });
      this._ws.send(JSON.stringify({ id, method, params }));
    });
  }

  /**
   * Send a streaming RPC call. Returns an object with:
   * - on(event, callback): listen for streaming events
   * - result: promise that resolves when 'done' event fires or rejects on 'error'
   */
  stream(method, params) {
    if (!this._ws || this._ws.readyState !== WebSocket.OPEN) {
      throw new Error('Not connected');
    }

    const id = ++this._requestId;
    const events = new Map();
    this._listeners.set(id, { events });

    const handle = {
      on(event, callback) {
        if (!events.has(event)) {events.set(event, []);}
        events.get(event).push(callback);
        return handle;
      },
      result: new Promise((resolve, reject) => {
        // Auto-wire done/error to resolve/reject the promise
        if (!events.has('done')) {events.set('done', []);}
        events.get('done').push((data) => {
          this._listeners.delete(id);
          resolve(data);
        });
        if (!events.has('error')) {events.set('error', []);}
        events.get('error').push((data) => {
          this._listeners.delete(id);
          reject(new Error(data.message || 'Agent error'));
        });
      }),
    };

    this._ws.send(JSON.stringify({ id, method, params }));
    return handle;
  }
}

// Singleton instance
let _instance = null;

export function getClient() {
  if (!_instance) {
    _instance = new FlynnClient();
    _instance.connect();
  }
  return _instance;
}