pit-router/app/socketio/handlers.py

"""
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})