evotraders/frontend/src/services/websocket.js

/**
 * WebSocket Client with Dynamic Port Resolution
 * Handles connection, reconnection, and heartbeat
 * Fetches Gateway port from API before connecting
 */

import { WS_URL } from "../config/constants";

// Global port cache
let cachedGatewayPort = null;
let cachedWsUrl = null;

/**
 * Fetch Gateway WebSocket port from API
 */
export async function fetchGatewayPort() {
  try {
    const response = await fetch('/api/runtime/gateway/port');
    if (!response.ok) {
      throw new Error(`HTTP ${response.status}`);
    }
    const data = await response.json();

    if (data.is_running && data.port) {
      cachedGatewayPort = data.port;
      cachedWsUrl = data.ws_url;
      return { port: data.port, wsUrl: data.ws_url };
    }

    return null;
  } catch (error) {
    console.warn('[Gateway] Failed to fetch port:', error);
    return null;
  }
}

/**
 * Get cached or default WebSocket URL
 */
export function getWebSocketUrl() {
  if (cachedWsUrl) {
    return cachedWsUrl;
  }
  return WS_URL;
}

/**
 * Clear cached port (call when Gateway restarts)
 */
export function clearGatewayCache() {
  cachedGatewayPort = null;
  cachedWsUrl = null;
}

export class ReadOnlyClient {
  constructor(onEvent, { wsUrl = null, reconnectDelay = 3000, heartbeatInterval = 5000 } = {}) {
    this.onEvent = onEvent;
    this.wsUrl = wsUrl; // null = auto-resolve from API
    this.baseReconnectDelay = reconnectDelay;
    this.reconnectDelay = reconnectDelay;
    this.maxReconnectDelay = 30000;
    this.heartbeatInterval = heartbeatInterval;
    this.ws = null;
    this.shouldReconnect = false;
    this.reconnectTimer = null;
    this.heartbeatTimer = null;
    this.reconnectAttempts = 0;
    this.lastPongTime = 0;
    this.isConnecting = false;
  }

  async connect() {
    this.shouldReconnect = true;
    this.reconnectAttempts = 0;
    this.reconnectDelay = this.baseReconnectDelay;
    await this._connect();
  }

  async _connect() {
    if (!this.shouldReconnect || this.isConnecting) {
      return;
    }

    this.isConnecting = true;

    // Resolve WebSocket URL if not set
    let targetUrl = this.wsUrl;
    if (!targetUrl) {
      // Try to fetch from API first
      const gatewayInfo = await fetchGatewayPort();
      if (gatewayInfo) {
        targetUrl = gatewayInfo.wsUrl;
        console.log(`[WebSocket] Resolved Gateway port: ${gatewayInfo.port}`);
      } else {
        // Fallback to default
        targetUrl = WS_URL;
        console.log(`[WebSocket] Using default URL: ${targetUrl}`);
      }
    }

    // Clear any existing connection
    if (this.ws) {
      this.ws.onopen = null;
      this.ws.onmessage = null;
      this.ws.onerror = null;
      this.ws.onclose = null;
      if (this.ws.readyState === WebSocket.OPEN || this.ws.readyState === WebSocket.CONNECTING) {
        this.ws.close();
      }
      this.ws = null;
    }

    try {
      this.ws = new WebSocket(targetUrl);

      this.ws.onopen = () => {
        this.reconnectAttempts = 0;
        this.reconnectDelay = this.baseReconnectDelay;
        this.lastPongTime = Date.now();
        this._safeEmit({ type: "system", content: "已连接实时数据服务" });
        console.log("WebSocket connected to", targetUrl);
        this._startHeartbeat();
        this.isConnecting = false;
      };

      this.ws.onmessage = (ev) => {
        try {
          const msg = JSON.parse(ev.data);

          // Update pong time for any message (server is alive)
          this.lastPongTime = Date.now();

          if (msg.type === "pong") {
            return;
          }

          console.log("[WebSocket] Message received:", msg.type || "unknown");
          this._safeEmit(msg);
        } catch (e) {
          console.error("[WebSocket] Parse error:", e);
        }
      };

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

      this.ws.onclose = (event) => {
        const code = event.code || "未知";
        console.log(`[WebSocket] Connection closed: Code=${code}, WasClean=${event.wasClean}`);

        this._stopHeartbeat();
        this.ws = null;
        this.isConnecting = false;

        // Always attempt reconnect if shouldReconnect is true
        if (this.shouldReconnect) {
          this.reconnectAttempts++;
          // Exponential backoff with cap
          this.reconnectDelay = Math.min(
            this.baseReconnectDelay * Math.pow(1.5, this.reconnectAttempts),
            this.maxReconnectDelay
          );

          this._safeEmit({
            type: "system",
            content: "正在尝试连接数据服务..."
          });

          if (this.reconnectTimer) {
            clearTimeout(this.reconnectTimer);
          }

          this.reconnectTimer = setTimeout(() => {
            console.log(`[WebSocket] Reconnect attempt ${this.reconnectAttempts}...`);
            this._connect();
          }, this.reconnectDelay);
        }
      };
    } catch (error) {
      console.error("[WebSocket] Connection error:", error);
      this.isConnecting = false;

      if (this.shouldReconnect) {
        this.reconnectTimer = setTimeout(() => {
          this._connect();
        }, this.reconnectDelay);
      }
    }
  }

  _safeEmit(msg) {
    try {
      this.onEvent(msg);
    } catch (e) {
      console.error("[WebSocket] Error in event handler:", e);
    }
  }

  _startHeartbeat() {
    this._stopHeartbeat();
    this.lastPongTime = Date.now();

    this.heartbeatTimer = setInterval(() => {
      this._sendPing();

      // Check for stale connection (no response in 60s)
      const timeSinceLastPong = Date.now() - this.lastPongTime;
      if (timeSinceLastPong > 60000 && this.ws) {
        console.warn("[WebSocket] Connection appears stale, forcing reconnect");
        this.ws.close();
      }
    }, this.heartbeatInterval);
  }

  _sendPing() {
    if (this.ws && this.ws.readyState === WebSocket.OPEN) {
      try {
        this.ws.send(JSON.stringify({ type: "ping" }));
      } catch (e) {
        console.error("Heartbeat send error:", e);
      }
    }
  }

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

  send(message) {
    if (this.ws && this.ws.readyState === WebSocket.OPEN) {
      try {
        const messageStr = typeof message === "string" ? message : JSON.stringify(message);
        this.ws.send(messageStr);
        return true;
      } catch (e) {
        console.error("Send error:", e);
        return false;
      }
    } else {
      console.warn("WebSocket is not connected, cannot send message");
      return false;
    }
  }

  disconnect() {
    this.shouldReconnect = false;
    this._stopHeartbeat();
    if (this.reconnectTimer) {
      clearTimeout(this.reconnectTimer);
      this.reconnectTimer = null;
    }
    if (this.ws) {
      this.ws.onopen = null;
      this.ws.onmessage = null;
      this.ws.onerror = null;
      this.ws.onclose = null;
      try {
        this.ws.close();
      } catch (e) {
        console.error("Close error:", e);
      }
    }
    this.ws = null;
    this.isConnecting = false;
  }

  /**
   * Reconnect with new port (call after Gateway restart)
   */
  async reconnectWithNewPort() {
    console.log("[WebSocket] Reconnecting with new port...");
    clearGatewayCache();
    this.disconnect();
    this.shouldReconnect = true;
    await this.connect();
  }
}