AI_Painting_V2.0/src/utils/websocket.js

import { ElNotification } from 'element-plus'
import { h, ref } from 'vue'
import { useDisplayStore } from '@/stores'
import { getToken } from '@/utils/auth'
import { createTask, getTask } from '@/utils/createTask'
import { userError } from '@/utils/tokenError'

export function getChargeType(chargeType) {
  switch (chargeType) {
    case 'painting':
      return 1
    case 'video':
      return 2
    default:
      return 2
  }
}

export function websocketError(code, msg) {
  let message
  switch (code) {
    case 1006:
      message = '用户身份验证失败'
      userError()
      break
    case 4401: // 后端返回常规错误
      message = msg
      break
    case 4402: // 后端返回外部平台提交时的错误
      message = JSON.parse(msg)
      break
    case 4403: // 外部平台的任务结果的错误
      message = msg
      break
    default:
      message = '连接异常，请稍后重试'
  }

  ElNotification({
    title: '生成通知',

    message: h('i', { style: 'color: teal' }, message),
    type: 'error'
  })
}

export function websocketSuccess() {
  // 合并两个通知为一个
  ElNotification({
    title: '生成通知',
    message: h('div', [
      h('div', { style: 'font-weight: bold; color: teal;' }, '生成成功！'),
      h('br'),
      h('div', { style: 'color: orange; margin-top: 5px;' }, '内测状态，请及时下载生成的文件，云端储存与历史记录保留24小时！')
    ]),
    type: 'success',
    duration: 6000 // 增加持续时间以适应更多信息
  })
}

export async function generate(type, data) {
  const progress_text = ref('')
  const message = ref('')
  const useDisplay = useDisplayStore()
  const token = getToken()
  const taskId = crypto.randomUUID()
  let currentTaskId = null
  
  useDisplay.isSubGerenate = true

  const result = await createTask(data, type, taskId, token)
  console.log(result)
  // const token = 'eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9.eyJsb2dpblR5cGUiOiJsb2dpbiIsImxvZ2luSWQiOjY0NDEwODAyMjk1OTgzNzIzMCwicm5TdHIiOiJiWkVwS2JLWFJyZmRIaFFHWXZKTkdzOGdGM0JSRmxQOCJ9.5eQ2GtVdrDntQDe2tnF8vl_DhTfd2uW-KNqzvl1imc0'
  const wsURL = `${import.meta.env.VITE_API_WORKFLOW_WS}/?token=Bearer ${token}`
  const socket = new WebSocket(wsURL)
  console.log('WebSocket连接已建立')

  // 心跳机制相关变量
  let heartbeatInterval = null
  const heartbeatIntervalTime = 20000 // 30秒发送一次心跳

  try {
    // 接收服务器消息
    socket.onmessage = async (event) => {
      // 处理pong响应
      if (event.data === 'pong') {
        console.log('收到心跳响应')
        return
      } else if (event.data === 'please give me taskId') {
        socket.send(`setTaskId:${taskId}`)
        progress_text.value = '信息提交中...'
        return
      } else if (event.data === 'OK! Please continue. ') {
        socket.send(JSON.stringify({
          type: 'generate',
          data: result
        }))
        return
      } else if (event.data === '任务提交成功，正在排队中...') {
        progress_text.value = '视频生成中...'
        currentTaskId = taskId
        
        useDisplay.addGeneratingItem({
          taskId: taskId,
          text: data.prompt || '生成中...',
          name: data.prompt || '生成中...',
          type: type
        })
        setTimeout(() => {
          useDisplay.scrollToBottom()
        }, 100)
        return
      }
      message.value = event.data
    }

    // 处理链接错误
    socket.onerror = (error) => {
      console.error('WebSocket链接出错:', error)

      // 清理心跳定时器
      if (heartbeatInterval) {
        clearInterval(heartbeatInterval)
      }

      // eslint-disable-next-line no-undef
      ElNotification({
        title: '生成通知',
        // eslint-disable-next-line no-undef
        message: h('i', { style: 'color: teal' }, '生成视频失败'),
        type: 'error'
      })
    }

    // 处理链接关闭
    socket.onclose = async (event) => {
      console.log('WebSocket已关闭:', event)
      useDisplay.isSubGerenate = false
      if (heartbeatInterval) {
        clearInterval(heartbeatInterval)
      }
      const res = JSON.parse(message.value)
      if (event.code === 1006) {
        console.error('用户身份验证失败')
        userError()
      } else if (event.code === 1000 && event.reason === 'success') {
        console.log('收到服务器消息:', res)
        const result = await getTask(res)
        if (result.type) {
          
          if (currentTaskId) {
            useDisplay.updateItemToSuccess(currentTaskId, result.urls)
          }
          
          websocketSuccess()
        } else {
          websocketError(4403, res.message)
        }
      } else {
        websocketError(event.code, event.reason)
      }
      if (heartbeatInterval) {
        clearInterval(heartbeatInterval)
      }
    }

    // 等待 WebSocket 连接打开
    socket.onopen = () => {
      console.log('WebSocket连接已建立')

      // 启动心跳机制
      heartbeatInterval = setInterval(() => {
        if (socket.readyState === WebSocket.OPEN) {
          socket.send('ping')
          console.log('发送心跳包')
        }
      }, heartbeatIntervalTime)
    }
  } catch (error) {
    console.log('Error creating AI3D_file:', error)
    // eslint-disable-next-line no-undef
    ElNotification({
      title: '生成通知',
      // eslint-disable-next-line no-undef
      message: h('i', { style: 'color: teal' }, '生成失败，请检查参数后重新提交任务'),
      type: 'error'
    })
  }
}