library-picturebook-activity/lesingle-aicreate-client/demo/CreationProgressTracker.js

/**
 * 创作进度追踪器
 *
 * 三级通信策略：
 *   Level 1: WebSocket (STOMP) — 首选，实时推送，自动重连
 *   Level 2: B2 HTTP 轮询 — WebSocket 断连期间临时补位
 *   Level 3: Redis 进度数据 — 断连期间进度不丢失，重连后立刻获取最新状态
 *
 * 使用示例：
 *   const tracker = new CreationProgressTracker({
 *     serverUrl: 'https://your-server.com',
 *     orgId: 'ORG001',
 *     appKey: 'ORG001',
 *     appSecret: 'your-secret',
 *     workId: '2034987994917769216',
 *     onProgress: (progress, message) => {
 *       console.log(`${progress}% - ${message}`)
 *     },
 *     onComplete: (result) => {
 *       // result.status === 'COMPLETED' => result.pageList 包含完整作品
 *       // result.status === 'FAILED' => result.failReason 包含失败原因
 *     },
 *     onError: (error) => {
 *       console.error('Tracker error:', error)
 *     }
 *   })
 *   tracker.start()
 *
 *   // 用户退出页面时
 *   tracker.stop()
 */

import SockJS from 'sockjs-client/dist/sockjs.min.js'
import { Client } from '@stomp/stompjs'
import { generateSignHeaders } from './HmacSigner.js'

export default class CreationProgressTracker {
  /**
   * @param {Object} options
   * @param {string} options.serverUrl - 服务器地址 (如 https://your-server.com)
   * @param {string} options.orgId - 机构ID
   * @param {string} options.appKey - HMAC 签名 App Key (通常等于 orgId)
   * @param {string} options.appSecret - HMAC 签名密钥
   * @param {string} options.workId - 作品ID
   * @param {Function} options.onProgress - 进度回调 (progress: number, message: string)
   * @param {Function} options.onComplete - 完成回调 (result: CreationResultVO)
   * @param {Function} [options.onError] - 错误回调 (error: Error)
   * @param {number} [options.pollInterval=5000] - B2 轮询间隔 (ms)
   * @param {number} [options.reconnectDelay=3000] - WebSocket 重连延迟 (ms)
   * @param {number} [options.maxPollTimeout=600000] - 最大轮询时间 (ms)，默认10分钟
   */
  constructor(options) {
    this.serverUrl = options.serverUrl
    this.orgId = options.orgId
    this.appKey = options.appKey
    this.appSecret = options.appSecret
    this.workId = options.workId
    this.onProgress = options.onProgress || (() => {})
    this.onComplete = options.onComplete || (() => {})
    this.onError = options.onError || (() => {})
    this.pollInterval = options.pollInterval || 5000
    this.reconnectDelay = options.reconnectDelay || 3000
    this.maxPollTimeout = options.maxPollTimeout || 600000

    this._stompClient = null
    this._pollTimer = null
    this._wsConnected = false
    this._stopped = false
    this._startTime = 0
    this._lastProgress = -1
  }

  /**
   * 启动进度追踪
   */
  start() {
    this._stopped = false
    this._startTime = Date.now()

    // 尝试 WebSocket 连接
    this._connectWebSocket()

    // 安全网：5秒内 WS 未连上则启动轮询
    setTimeout(() => {
      if (!this._wsConnected && !this._stopped) {
        this._startPolling()
      }
    }, 5000)
  }

  /**
   * 停止追踪（用户离开页面时调用）
   */
  stop() {
    this._stopped = true
    this._stopPolling()
    this._disconnectWebSocket()
  }

  // ─── WebSocket (Level 1) ──────────────────────────────────────────

  _connectWebSocket() {
    if (this._stopped) return

    try {
      const wsUrl = `${this.serverUrl}/ws?orgId=${this.orgId}`

      this._stompClient = new Client({
        webSocketFactory: () => new SockJS(wsUrl),
        reconnectDelay: this.reconnectDelay,
        heartbeatIncoming: 10000,
        heartbeatOutgoing: 10000,

        onConnect: () => {
          this._wsConnected = true
          this._stopPolling()  // WS 连上了，停止轮询

          // 订阅进度 topic
          this._stompClient.subscribe(
            `/topic/progress/${this.workId}`,
            (msg) => {
              try {
                const data = JSON.parse(msg.body)
                this._handleProgressData(data)
              } catch (e) {
                // ignore parse errors
              }
            }
          )

          // 重连成功后，立即查一次 B2 获取断连期间可能错过的最终状态
          this._pollOnce()
        },

        onDisconnect: () => {
          this._wsConnected = false
          if (!this._stopped) {
            this._startPolling()  // WS 断了，启动轮询补位
          }
        },

        onStompError: (frame) => {
          this._wsConnected = false
          if (!this._stopped) {
            this._startPolling()
          }
        },

        onWebSocketError: () => {
          this._wsConnected = false
          if (!this._stopped) {
            this._startPolling()
          }
        }
      })

      this._stompClient.activate()
    } catch (e) {
      // WebSocket 不可用（如被防火墙阻断），降级为纯轮询
      this._startPolling()
    }
  }

  _disconnectWebSocket() {
    if (this._stompClient) {
      try {
        this._stompClient.deactivate()
      } catch (e) {
        // ignore
      }
      this._stompClient = null
    }
    this._wsConnected = false
  }

  // ─── B2 HTTP 轮询 (Level 2) ───────────────────────────────────────

  _startPolling() {
    if (this._pollTimer || this._stopped) return
    this._pollTimer = setInterval(() => this._pollOnce(), this.pollInterval)
  }

  _stopPolling() {
    if (this._pollTimer) {
      clearInterval(this._pollTimer)
      this._pollTimer = null
    }
  }

  async _pollOnce() {
    if (this._stopped) return

    // 超时保护
    if (Date.now() - this._startTime > this.maxPollTimeout) {
      this.stop()
      this.onError(new Error('进度查询超时'))
      return
    }

    try {
      const queryParams = { orgId: this.orgId }
      const signHeaders = generateSignHeaders(this.appKey, this.appSecret, queryParams)

      const res = await fetch(
        `${this.serverUrl}/api/v1/query/work/${this.workId}?orgId=${this.orgId}`,
        { headers: { ...signHeaders, 'Content-Type': 'application/json' } }
      )

      if (!res.ok) return

      const json = await res.json()
      if (json.code !== 200 || !json.data) return

      const data = json.data

      // 推送进度
      if (data.progress != null && data.progress !== this._lastProgress) {
        this._handleProgressData({
          progress: data.progress,
          message: data.progressMessage || ''
        })
      }

      // 检查是否终态
      if (data.status === 'COMPLETED' || data.status === 'FAILED') {
        this.stop()
        this.onComplete(data)
      }
    } catch (e) {
      // 网络错误，不中断轮询
    }
  }

  // ─── 进度处理 ─────────────────────────────────────────────────────

  _handleProgressData(data) {
    if (this._stopped) return

    const progress = data.progress != null ? data.progress : 0
    const message = data.message || ''

    // 去重：相同进度不重复回调
    if (progress === this._lastProgress) return
    this._lastProgress = progress

    // 回调
    this.onProgress(progress, message)

    // 终态检测（WebSocket 推送的 progress=100 或 -1）
    if (progress === 100) {
      // 完成，再查一次 B2 获取完整数据（含 pageList）
      setTimeout(() => this._fetchFinalResult(), 1000)
    } else if (progress === -1) {
      // 失败，查一次 B2 获取 failReason
      setTimeout(() => this._fetchFinalResult(), 1000)
    }
  }

  async _fetchFinalResult() {
    try {
      const queryParams = { orgId: this.orgId }
      const signHeaders = generateSignHeaders(this.appKey, this.appSecret, queryParams)

      const res = await fetch(
        `${this.serverUrl}/api/v1/query/work/${this.workId}?orgId=${this.orgId}`,
        { headers: { ...signHeaders, 'Content-Type': 'application/json' } }
      )

      if (!res.ok) return

      const json = await res.json()
      if (json.code === 200 && json.data) {
        this.stop()
        this.onComplete(json.data)
      }
    } catch (e) {
      this.onError(e)
    }
  }
}