""" Socket.IO 事件处理 """ from flask_socketio import emit, join_room, leave_room from flask_jwt_extended import decode_token from datetime import datetime import json from app.extensions import redis_client # 连接管理器 class ConnectionManager: def __init__(self): self.user_sockets = {} # user_id -> socket_id self.agent_sockets = {} # agent_id -> socket_id self.socket_sessions = {} # socket_id -> session_id connection_manager = ConnectionManager() def register_handlers(socketio): """注册 Socket.IO 事件处理器""" @socketio.on('connect') def handle_connect(): """客户端连接""" print('Client connected') # 发送认证请求 emit('auth', {'message': 'Please authenticate'}) @socketio.on('disconnect') def handle_disconnect(): """客户端断开连接""" from flask import request sid = request.sid print(f'Client disconnected: {sid}') # 清理连接 if sid in connection_manager.socket_sessions: del connection_manager.socket_sessions[sid] @socketio.on('auth') def handle_auth(data): """处理认证""" from flask import request sid = request.sid token = data.get('token') if not token: emit('auth_error', {'code': 'MISSING_TOKEN', 'message': 'Token is required'}) return try: # 验证 JWT Token decoded = decode_token(token) user_id = decoded['sub'] # 保存连接信息 connection_manager.user_sockets[user_id] = sid # 存储到 Redis redis_client.hset(f'socket:{sid}', mapping={ 'user_id': user_id, 'connected_at': datetime.utcnow().isoformat() }) redis_client.expire(f'socket:{sid}', 86400) emit('authenticated', { 'user_id': user_id, 'socket_id': sid }) except Exception as e: emit('auth_error', {'code': 'INVALID_TOKEN', 'message': str(e)}) @socketio.on('ping') def handle_ping(data): """心跳响应""" emit('pong', {'timestamp': datetime.utcnow().timestamp()}) @socketio.on('session.create') def handle_session_create(data): """创建会话""" from app.models import db, Session, Agent user_id = data.get('user_id') title = data.get('title', 'New Session') agent_id = data.get('agent_id') # 如果没有指定 Agent,分配一个 if not agent_id: agent = Agent.query.filter_by(status='online').first() if agent: agent_id = agent.id session = Session( user_id=user_id, primary_agent_id=agent_id, title=title ) db.session.add(session) db.session.commit() emit('session.created', { 'session_id': session.id, 'agent_id': agent_id }) @socketio.on('session.join') def handle_session_join(data): """加入会话""" from flask import request from app.models import Session sid = request.sid session_id = data.get('session_id') # 加入房间 join_room(session_id) connection_manager.socket_sessions[sid] = session_id session = Session.query.get(session_id) if session: emit('session.joined', { 'session_id': session_id, 'participants': [session.user_id, session.primary_agent_id] }) @socketio.on('message.send') def handle_message_send(data): """发送消息""" from flask import request from app.models import db, Message, Session sid = request.sid session_id = data.get('session_id') content = data.get('content') sender_type = data.get('sender_type', 'user') sender_id = data.get('sender_id') # 创建消息 message = Message( session_id=session_id, sender_type=sender_type, sender_id=sender_id, content=content, status='sent' ) db.session.add(message) # 更新会话 session = Session.query.get(session_id) if session: session.message_count += 1 session.last_active_at = datetime.utcnow() db.session.commit() # 广播消息到房间 emit('message', { 'message_id': message.id, 'session_id': session_id, 'sender': {'type': sender_type, 'id': sender_id}, 'content': content, 'timestamp': datetime.utcnow().isoformat() }, room=session_id) @socketio.on('message.ack') def handle_message_ack(data): """消息确认""" from app.models import db, Message message_id = data.get('message_id') status = data.get('status') message = Message.query.get(message_id) if message: message.ack_status = status if status == 'acknowledged': message.status = 'delivered' message.delivered_at = datetime.utcnow() db.session.commit() emit('message.acked', {'message_id': message_id, 'status': status})