以下内容节选自我的实战课程《从0到1教你搭建一个基于微信小程序的AI智能体应用平台》,课程包含详细代码和讲解,链接如下: https://edu.csdn.net/course/detail/40753

课程目标

  • 建立 前后端 WebSocket 连接并实现消息收发
  • 统一 连接状态管理与断线重连策略
  • 实现 实时流式对话的基础架构

知识要点

服务端技术栈

  • geventwebsocket:基于 gevent 的 WebSocket 实现,支持高并发
  • 资源注册:/wsbot 端点处理 WebSocket 连接
  • 事件循环:利用 gevent 的协程特性处理大量并发连接
  • 消息路由:根据消息类型分发到不同的处理函数

客户端技术要点

  • 连接管理:建立、维护、关闭 WebSocket 连接
  • 心跳机制:定期发送 ping 消息保持连接活跃
  • 重连策略:指数退避算法,避免频繁重连
  • 消息格式:统一的消息编码和解码规范

安全机制

  • Token 鉴权:连接时验证用户身份
  • 限流控制:防止单个用户过度占用资源
  • 异常处理:优雅处理各种异常情况
  • 数据验证:严格验证消息格式和内容

实战步骤详解

1. 服务端实现

WebSocket 应用类设计

class WSBotApplication:
    def __init__(self):
        self.connections = {}  # 存储活跃连接
        self.user_connections = {}  # 用户到连接的映射
        self.heartbeat_interval = 30  # 心跳间隔(秒)
   
    def on_open(self, ws, user_id, token):
        """连接建立时的处理"""
        # 验证 token 有效性
        if not self.validate_token(token):
            ws.close(code=1008, reason="Invalid token")
            return
       
        # 检查用户是否已有连接
        if user_id in self.user_connections:
            old_ws = self.user_connections[user_id]
            old_ws.close(code=1000, reason="New connection")
       
        # 注册新连接
        self.connections[ws] = {
            'user_id': user_id,
            'token': token,
            'last_heartbeat': time.time(),
            'message_queue': []
        }
        self.user_connections[user_id] = ws
       
        # 发送连接成功消息
        ws.send(json.dumps({
            'type': 'connection_established',
            'message': 'WebSocket connection established'
        }))
   
    def on_message(self, ws, message):
        """处理接收到的消息"""
        try:
            data = json.loads(message)
            msg_type = data.get('type')
           
            if msg_type == 'ping':
                self.handle_heartbeat(ws)
            elif msg_type == 'chat':
                self.handle_chat_message(ws, data)
            elif msg_type == 'image':
                self.handle_image_message(ws, data)
            else:
                ws.send(json.dumps({
                    'type': 'error',
                    'message': 'Unknown message type'
                }))
        except json.JSONDecodeError:
            ws.send(json.dumps({
                'type': 'error',
                'message': 'Invalid JSON format'
            }))
   
    def on_close(self, ws, code, reason):
        """连接关闭时的清理"""
        if ws in self.connections:
            user_id = self.connections[ws]['user_id']
            del self.connections[ws]
            if user_id in self.user_connections:
                del self.user_connections[user_id]

消息处理逻辑

def handle_chat_message(self, ws, data):
    """处理聊天消息"""
    user_id = self.connections[ws]['user_id']
    bot_id = data.get('botId')
    message = data.get('message')
   
    # 验证消息格式
    if not bot_id or not message:
        ws.send(json.dumps({
            'type': 'error',
            'message': 'Missing required fields'
        }))
        return
   
    # 调用 AI 服务生成回复
    try:
        response_stream = self.ai_service.generate_response(
            user_id=user_id,
            bot_id=bot_id,
            message=message
        )
       
        # 流式发送回复
        for chunk in response_stream:
            ws.send(json.dumps({
                'type': 'chunk',
                'content': chunk,
                'botId': bot_id
            }))
       
        # 发送结束标记
        ws.send(json.dumps({
            'type': 'end',
            'botId': bot_id
        }))
       
    except Exception as e:
        ws.send(json.dumps({
            'type': 'error',
            'message': f'AI service error: {str(e)}'
        }))

2. 客户端实现

WebSocket 管理器

class WebSocketManager {
  constructor() {
    this.ws = null;
    this.reconnectAttempts = 0;
    this.maxReconnectAttempts = 5;
    this.reconnectInterval = 1000; // 初始重连间隔
    this.heartbeatInterval = null;
    this.isConnecting = false;
    this.messageQueue = []; // 离线消息队列
  }

  connect(token) {
    if (this.isConnecting || (this.ws && this.ws.readyState === WebSocket.OPEN)) {
      return;
    }

    this.isConnecting = true;
    const wsUrl = `wss://your-domain.com/wsbot?token=${token}`;
   
    try {
      this.ws = new WebSocket(wsUrl);
      this.setupEventHandlers();
    } catch (error) {
      console.error('WebSocket connection failed:', error);
      this.handleConnectionError();
    }
  }

  setupEventHandlers() {
    this.ws.onopen = () => {
      console.log('WebSocket connected');
      this.isConnecting = false;
      this.reconnectAttempts = 0;
      this.startHeartbeat();
      this.processMessageQueue();
    };

    this.ws.onmessage = (event) => {
      try {
        const data = JSON.parse(event.data);
        this.handleMessage(data);
      } catch (error) {
        console.error('Failed to parse message:', error);
      }
    };

    this.ws.onclose = (event) => {
      console.log('WebSocket closed:', event.code, event.reason);
      this.isConnecting = false;
      this.stopHeartbeat();
     
      if (event.code !== 1000) { // 非正常关闭
        this.scheduleReconnect();
      }
    };

    this.ws.onerror = (error) => {
      console.error('WebSocket error:', error);
      this.handleConnectionError();
    };
  }

  sendMessage(message) {
    if (this.ws && this.ws.readyState === WebSocket.OPEN) {
      this.ws.send(JSON.stringify(message));
    } else {
      // 连接未建立,加入队列
      this.messageQueue.push(message);
    }
  }

  scheduleReconnect() {
    if (this.reconnectAttempts >= this.maxReconnectAttempts) {
      console.error('Max reconnection attempts reached');
      return;
    }

    const delay = this.reconnectInterval * Math.pow(2, this.reconnectAttempts);
    this.reconnectAttempts++;
   
    setTimeout(() => {
      this.connect(this.token);
    }, delay);
  }
}

心跳机制实现

startHeartbeat() {
  this.heartbeatInterval = setInterval(() => {
    if (this.ws && this.ws.readyState === WebSocket.OPEN) {
      this.sendMessage({
        type: 'ping',
        timestamp: Date.now()
      });
    }
  }, 30000); // 30秒发送一次心跳
}

stopHeartbeat() {
  if (this.heartbeatInterval) {
    clearInterval(this.heartbeatInterval);
    this.heartbeatInterval = null;
  }
}

3. 消息格式规范

客户端发送消息格式

// 文本消息
{
  "type": "chat",
  "botId": "assistant_001",
  "message": "用户输入内容",
  "sessionId": "session_123",
  "timestamp": 1640995200000
}

// 图片消息
{
  "type": "image",
  "botId": "assistant_001",
  "imageUrl": "https://example.com/image.jpg",
  "message": "图片描述",
  "sessionId": "session_123",
  "timestamp": 1640995200000
}

服务端响应消息格式

// 流式内容块
{
  "type": "chunk",
  "content": "AI回复的部分内容",
  "botId": "assistant_001",
  "sessionId": "session_123"
}

// 消息结束标记
{
  "type": "end",
  "botId": "assistant_001",
  "sessionId": "session_123",
  "totalTokens": 150,
  "cost": 0.02
}

// 错误消息
{
  "type": "error",
  "message": "错误描述",
  "code": "ERROR_CODE",
  "botId": "assistant_001"
}

4. 断线重连与并发控制

单用户单连接策略

def enforce_single_connection(self, user_id, new_ws):
    """确保每个用户只有一个活跃连接"""
    if user_id in self.user_connections:
        old_ws = self.user_connections[user_id]
        # 发送通知给旧连接
        old_ws.send(json.dumps({
            'type': 'connection_replaced',
            'message': 'New connection established, closing old one'
        }))
        old_ws.close(code=1000, reason="Replaced by new connection")
   
    # 注册新连接
    self.user_connections[user_id] = new_ws

消息队列处理

processMessageQueue() {
  while (this.messageQueue.length > 0) {
    const message = this.messageQueue.shift();
    this.sendMessage(message);
  }
}

验收标准

功能验收

  • WebSocket 连接能够正常建立和关闭
  • 消息能够实时双向传输
  • 断线重连机制工作正常
  • 心跳机制保持连接活跃

性能验收

  • 在 3G 网络环境下连接稳定
  • 重连延迟在可接受范围内(< 5秒)
  • 支持至少 1000 个并发连接
  • 消息传输延迟 < 100ms

安全验收

  • 非法 Token 被正确拒绝
  • 连接超时自动断开
  • 消息格式验证严格
  • 防止恶意连接攻击

消息可靠传输实现方案

实现消息的可靠传输,需保证消息不会丢失、重复或乱序,常见方案如下:

1. 消息唯一ID与确认机制
  • 每条消息分配全局唯一ID(如UUID)。
  • 发送方发送消息后,服务端返回ack确认(带消息ID)。
  • 客户端维护未确认消息队列,超时未收到ack则重发。
2. 服务端持久化与幂等处理
  • 服务端收到消息后,先校验消息ID是否已处理,避免重复处理。
  • 可将消息持久化到数据库或队列,确保服务端异常时消息不丢失。
3. 客户端重连补发
  • 客户端断线重连后,将未确认的消息重新发送,服务端根据ID去重。
4. 顺序保证
  • 可为每个会话维护消息序号,服务端按序号递增处理,乱序时缓存等待。
Logo

中国智能体开发者社区,聚焦智能体与大模型开发,提供前沿资讯、实用工具链、开源项目及行业案例。通过技术沙龙、开发者大赛等活动,促进经验交流与协作,助力开发者快速构建创新智能应用。

更多推荐