aistream-test/测试流式传输uniapp/pages/index/useWebSocket.js

// src/hooks/useWebSocket.js
import { ref, onUnmounted, computed } from 'vue';
import { ProtocolCodec, MessageType } from './ProtocolCodec';

export default function useWebSocket(url, options = {}) {
  // 使用模块级变量存储 SocketTask
  let socketTask = null;
  let reconnectTimer = null;

  // 配置合并
  const _defaultOptions = {
    identity: {},
    protocols: ['binary'],
    reconnectDelay: 3000,
    onMessage: () => {},
    onAuthSuccess: () => {},
    onClose: () => {},
    ...options
  };

  // 响应式状态
  const state = ref({
    isConnected: false,
    isConnecting: false,
    error: null,
    clientId: '',
    context: null
  });

  // 清理连接
  const cleanupSocket = () => {
    if (socketTask) {
      try {
        socketTask.close({ code: 1008, reason: '主动关闭' });
      } catch (e) {
        console.warn('关闭连接时出错:', e);
      }
      socketTask = null;
    }
    
    state.value = {
      isConnected: false,
      isConnecting: false,
      error: null,
      clientId: '',
      context: null
    };
    
    if (reconnectTimer) {
      clearTimeout(reconnectTimer);
      reconnectTimer = null;
    }
  };

  // 发送消息
  const sendMessage = (data, isBinary = true) => {
    return new Promise((resolve, reject) => {
      if (!socketTask || !state.value.isConnected) {
        const err = new Error('WebSocket 未连接');
        state.value.error = err;
        return reject(err);
      }

      const sendData = isBinary ? data : JSON.stringify(data);
      
      socketTask.send({
        data: sendData,
        success: () => resolve(),
        fail: (err) => {
          console.error('消息发送失败:', err);
          reject(new Error(`发送失败: ${err.errMsg || err.message}`));
        }
      });
    });
  };

  // 发送身份校验包
  const sendIdentityPacket = async () => {
    const { user_id, token, name } = _defaultOptions.identity;
    if (!user_id || !token) {
      state.value.error = new Error('身份校验信息缺失（user_id/token）');
      cleanupSocket();
      return;
    }

    const identityPacket = ProtocolCodec.pack(MessageType.IDENTITY, {
      user_id,
      token,
      name: name || `用户${user_id}`
    });

    try {
      await sendMessage(identityPacket);
    } catch (err) {
      console.error('身份校验包发送失败', err);
      cleanupSocket();
      reconnectSocket();
    }
  };

  // 处理服务端消息
  const handleServerMessage = (res) => {
    try {
      const { data } = res;
      let msgType, bodyData;
      
      // 兼容不同数据格式
      if (data instanceof ArrayBuffer) {
        const uint8Array = new Uint8Array(data);
        [msgType, _, bodyData] = ProtocolCodec.unpack(uint8Array);
      } else if (typeof data === 'string') {
        // 如果是字符串，可能是文本消息
        console.warn('收到非二进制消息:', data);
        return;
      } else {
        console.warn('未知的消息数据类型:', typeof data, data);
        return;
      }

      const bodyStr = bodyData.toString('utf-8');
      const parsedBody = bodyStr ? JSON.parse(bodyStr) : {};

      if (msgType === MessageType.IDENTITY_RESP) {
        state.value.context = parsedBody.data;
        state.value.clientId = parsedBody.data?.client_id || '';
        _defaultOptions.onAuthSuccess(parsedBody.data);
      }

      _defaultOptions.onMessage({
        msgType,
        data: parsedBody,
        rawData: data
      });
    } catch (err) {
      console.error('消息处理错误:', err);
      state.value.error = new Error(`消息解析失败：${err.message}`);
      _defaultOptions.onMessage({
        msgType: 'ERROR',
        data: { message: err.message },
        rawData: res.data
      });
    }
  };
  
  // 初始化连接
  const initSocket = () => {
    if (state.value.isConnecting || state.value.isConnected) {
      console.log('连接已存在或正在连接中，跳过重复连接');
      return;
    }
    
    state.value.isConnecting = true;
    state.value.error = null;

    console.log('发起 WebSocket 连接：', url);

    // 创建连接
    socketTask = uni.connectSocket({
      url,
      protocols: _defaultOptions.protocols,
      success: () => {
        console.log('connectSocket API 调用成功');
      },
      fail: (err) => {
        console.error('connectSocket API 调用失败：', err);
        state.value.error = new Error(`连接创建失败：${err.errMsg || err.message}`);
        state.value.isConnecting = false;
        reconnectSocket();
      }
    });
	socketTask.onOpen(() => {
		console.log('xxxxxxxx');
	})
    // 检查实例是否有效
    if (!socketTask) {
      console.error('SocketTask 实例创建失败');
      state.value.error = new Error('SocketTask 实例为空');
      state.value.isConnecting = false;
      reconnectSocket();
      return;
    }

    // 绑定事件处理器
    const openHandler = () => {
      console.log('=== WebSocket 连接成功 ===');
      state.value.isConnecting = false;
      state.value.isConnected = true;
      clearTimeout(reconnectTimer);
      sendIdentityPacket();
    };

    const messageHandler = handleServerMessage;
    
    const closeHandler = (res) => {
      console.log('=== WebSocket 连接关闭 ===', res);
      const closeInfo = {
        code: res.code,
        reason: res.reason,
        isManual: res.code === 1008
      };
      _defaultOptions.onClose(closeInfo);

      state.value.isConnected = false;
      state.value.isConnecting = false;
      
      if (!closeInfo.isManual && res.code !== 1000) {
        reconnectSocket();
      }
    };

    const errorHandler = (err) => {
      console.log('=== WebSocket 连接错误 ===', err);
      state.value.error = new Error(`连接错误：${err.errMsg || err.message}`);
      state.value.isConnecting = false;
      state.value.isConnected = false;
      reconnectSocket();
    };

    // 绑定事件
    socketTask.onOpen(openHandler);
    socketTask.onMessage(messageHandler);
    socketTask.onClose(closeHandler);
    socketTask.onError(errorHandler);

    // 存储事件处理器以便清理
    socketTask._handlers = {
      open: openHandler,
      message: messageHandler,
      close: closeHandler,
      error: errorHandler
    };
  };

  // 重连逻辑
  const reconnectSocket = () => {
    if (reconnectTimer) clearTimeout(reconnectTimer);
    
    reconnectTimer = setTimeout(() => {
      console.log(`尝试重连（延迟${_defaultOptions.reconnectDelay}ms）`);
      initSocket();
    }, _defaultOptions.reconnectDelay);
  };

  // 断开连接并停止重连
  const disconnect = () => {
    if (reconnectTimer) {
      clearTimeout(reconnectTimer);
      reconnectTimer = null;
    }
    cleanupSocket();
  };

  // 对外暴露的方法
  const actions = {
    connect: initSocket,
    disconnect,
    sendBinary: (data) => sendMessage(data, true),
    sendJson: (data) => sendMessage(data, false),
    reconnect: reconnectSocket
  };

  // 组件卸载清理
  onUnmounted(() => {
    disconnect();
  });

  return {
    state: computed(() => ({ ...state.value })),
    ...actions
  };
}