ai_trade_assistance/src/api/baidu/ws.ts

import { contractConfig } from '../../config/contract'
import type { BaiduWsItem, BaiduWsMessage, BaiduWsOutbound } from './wsTypes'

const WS_URL = 'wss://finance-ws.pae.baidu.com/'
const PING_MS = 6_000
const PATCH_MS = 60_000
const RECONNECT_BASE_MS = 1_000
const RECONNECT_MAX_MS = 30_000

export type BaiduQuoteWsHandlers = {
  onMessage: (msg: BaiduWsMessage) => void
  onError?: (err: unknown) => void
}

function buildItem(): BaiduWsItem {
  return {
    code: contractConfig.code,
    name: contractConfig.name,
    market: 'ab',
    financeType: 'futures',
  }
}

/**
 * Baidu finance quote WebSocket client.
 * Handles subscribe (tick + snapshot), ping every 6s, patch every 60s, and reconnect.
 */
export class BaiduQuoteWs {
  private ws: WebSocket | null = null
  private pingTimer: ReturnType<typeof setInterval> | null = null
  private patchTimer: ReturnType<typeof setInterval> | null = null
  private reconnectTimer: ReturnType<typeof setTimeout> | null = null
  private reconnectAttempt = 0
  private intentionalClose = false
  private handlers: BaiduQuoteWsHandlers

  constructor(handlers: BaiduQuoteWsHandlers) {
    this.handlers = handlers
  }

  connect(): void {
    this.intentionalClose = false
    this.clearReconnect()
    if (
      this.ws &&
      (this.ws.readyState === WebSocket.OPEN || this.ws.readyState === WebSocket.CONNECTING)
    ) {
      return
    }
    this.openSocket()
  }

  disconnect(): void {
    this.intentionalClose = true
    this.clearTimers()
    this.clearReconnect()
    if (this.ws) {
      this.ws.onopen = null
      this.ws.onmessage = null
      this.ws.onerror = null
      this.ws.onclose = null
      try {
        this.ws.close()
      } catch {
        /* ignore */
      }
      this.ws = null
    }
  }

  private openSocket(): void {
    if (this.intentionalClose) return

    const ws = new WebSocket(WS_URL)
    this.ws = ws

    ws.onopen = () => {
      if (this.intentionalClose) {
        try {
          ws.close()
        } catch {
          /* ignore */
        }
        return
      }
      this.reconnectAttempt = 0
      this.sendSubscribe()
      this.startTimers()
    }

    ws.onmessage = (ev) => {
      try {
        const msg = JSON.parse(String(ev.data)) as BaiduWsMessage
        this.handlers.onMessage(msg)
      } catch (e) {
        console.error('[baidu-ws] parse failed', e)
        this.handlers.onError?.(e)
      }
    }

    ws.onerror = (ev) => {
      console.error('[baidu-ws] error', ev)
      this.handlers.onError?.(ev)
    }

    ws.onclose = () => {
      this.clearTimers()
      this.ws = null
      if (!this.intentionalClose) this.scheduleReconnect()
    }
  }

  private send(payload: BaiduWsOutbound): void {
    if (!this.ws || this.ws.readyState !== WebSocket.OPEN) return
    this.ws.send(JSON.stringify(payload))
  }

  private sendSubscribe(): void {
    const item = buildItem()
    this.send({ method: 'subscribe', source: 'pc-web', product: 'tick', items: [item] })
    this.send({ method: 'subscribe', source: 'pc-web', product: 'snapshot', items: [item] })
  }

  private sendPing(): void {
    this.send({ method: 'ping', source: 'pc-web' })
  }

  private sendPatch(): void {
    const item = buildItem()
    this.send({ method: 'patch', source: 'pc-web', product: 'snapshot', items: [item] })
  }

  private startTimers(): void {
    this.clearTimers()
    this.pingTimer = setInterval(() => this.sendPing(), PING_MS)
    this.patchTimer = setInterval(() => this.sendPatch(), PATCH_MS)
  }

  private clearTimers(): void {
    if (this.pingTimer) {
      clearInterval(this.pingTimer)
      this.pingTimer = null
    }
    if (this.patchTimer) {
      clearInterval(this.patchTimer)
      this.patchTimer = null
    }
  }

  private scheduleReconnect(): void {
    this.clearReconnect()
    const delay = Math.min(RECONNECT_BASE_MS * 2 ** this.reconnectAttempt, RECONNECT_MAX_MS)
    this.reconnectAttempt += 1
    this.reconnectTimer = setTimeout(() => {
      this.reconnectTimer = null
      if (this.intentionalClose) return
      this.openSocket()
    }, delay)
  }

  private clearReconnect(): void {
    if (this.reconnectTimer) {
      clearTimeout(this.reconnectTimer)
      this.reconnectTimer = null
    }
  }
}