aiagent/backend/app/websocket/collaboration_manager.py

"""
工作流协作管理器
管理多人协作编辑工作流的WebSocket连接和消息同步
"""
from typing import Dict, Set, List, Optional
from fastapi import WebSocket
import json
import asyncio
from datetime import datetime
import logging

logger = logging.getLogger(__name__)


class CollaborationManager:
    """工作流协作管理器"""
    
    def __init__(self):
        # workflow_id -> Set[WebSocket] 映射，存储每个工作流的连接
        self.active_connections: Dict[str, Set[WebSocket]] = {}
        # WebSocket -> user_info 映射，存储每个连接的用户信息
        self.connection_users: Dict[WebSocket, Dict] = {}
        # workflow_id -> Dict[user_id, user_info] 映射，存储每个工作流的在线用户
        self.workflow_users: Dict[str, Dict[str, Dict]] = {}
        # 操作锁，用于冲突解决
        self.operation_locks: Dict[str, asyncio.Lock] = {}
    
    async def connect(self, websocket: WebSocket, workflow_id: str, user_id: str, username: str):
        """建立协作连接"""
        await websocket.accept()
        
        if workflow_id not in self.active_connections:
            self.active_connections[workflow_id] = set()
            self.workflow_users[workflow_id] = {}
            self.operation_locks[workflow_id] = asyncio.Lock()
        
        self.active_connections[workflow_id].add(websocket)
        
        user_info = {
            "user_id": user_id,
            "username": username,
            "joined_at": datetime.now().isoformat(),
            "color": self._get_user_color(user_id)  # 为用户分配颜色
        }
        self.connection_users[websocket] = user_info
        self.workflow_users[workflow_id][user_id] = user_info
        
        # 通知其他用户有新用户加入
        await self.broadcast_user_joined(workflow_id, user_info, exclude_websocket=websocket)
        
        # 发送当前在线用户列表给新用户
        await self.send_personal_message({
            "type": "collaboration_init",
            "workflow_id": workflow_id,
            "current_user": user_info,
            "online_users": list(self.workflow_users[workflow_id].values())
        }, websocket)
        
        logger.info(f"用户 {username} ({user_id}) 加入工作流 {workflow_id} 的协作编辑")
    
    def disconnect(self, websocket: WebSocket, workflow_id: str):
        """断开协作连接"""
        if workflow_id in self.active_connections:
            self.active_connections[workflow_id].discard(websocket)
            
            if websocket in self.connection_users:
                user_info = self.connection_users[websocket]
                user_id = user_info["user_id"]
                
                # 从工作流用户列表中移除
                if user_id in self.workflow_users[workflow_id]:
                    del self.workflow_users[workflow_id][user_id]
                
                # 通知其他用户有用户离开
                self.broadcast_user_left(workflow_id, user_id, exclude_websocket=websocket)
                
                del self.connection_users[websocket]
                logger.info(f"用户 {user_info.get('username')} ({user_id}) 离开工作流 {workflow_id} 的协作编辑")
            
            # 如果没有连接了，清理资源
            if not self.active_connections[workflow_id]:
                del self.active_connections[workflow_id]
                del self.workflow_users[workflow_id]
                if workflow_id in self.operation_locks:
                    del self.operation_locks[workflow_id]
    
    async def broadcast_operation(self, workflow_id: str, operation: Dict, exclude_websocket: Optional[WebSocket] = None):
        """广播操作到所有连接的客户端"""
        if workflow_id not in self.active_connections:
            return
        
        message = {
            "type": "operation",
            "workflow_id": workflow_id,
            "operation": operation,
            "timestamp": datetime.now().isoformat()
        }
        
        disconnected = set()
        for websocket in self.active_connections[workflow_id]:
            if websocket == exclude_websocket:
                continue
            
            try:
                await websocket.send_json(message)
            except Exception as e:
                logger.warning(f"发送协作消息失败: {e}")
                disconnected.add(websocket)
        
        # 清理断开的连接
        for ws in disconnected:
            self.disconnect(ws, workflow_id)
    
    async def broadcast_user_joined(self, workflow_id: str, user_info: Dict, exclude_websocket: Optional[WebSocket] = None):
        """广播用户加入消息"""
        if workflow_id not in self.active_connections:
            return
        
        message = {
            "type": "user_joined",
            "workflow_id": workflow_id,
            "user": user_info
        }
        
        disconnected = set()
        for websocket in self.active_connections[workflow_id]:
            if websocket == exclude_websocket:
                continue
            
            try:
                await websocket.send_json(message)
            except Exception as e:
                logger.warning(f"发送用户加入消息失败: {e}")
                disconnected.add(websocket)
        
        for ws in disconnected:
            self.disconnect(ws, workflow_id)
    
    async def broadcast_user_left(self, workflow_id: str, user_id: str, exclude_websocket: Optional[WebSocket] = None):
        """广播用户离开消息"""
        if workflow_id not in self.active_connections:
            return
        
        message = {
            "type": "user_left",
            "workflow_id": workflow_id,
            "user_id": user_id
        }
        
        disconnected = set()
        for websocket in self.active_connections[workflow_id]:
            if websocket == exclude_websocket:
                continue
            
            try:
                await websocket.send_json(message)
            except Exception as e:
                logger.warning(f"发送用户离开消息失败: {e}")
                disconnected.add(websocket)
        
        for ws in disconnected:
            self.disconnect(ws, workflow_id)
    
    async def send_personal_message(self, message: Dict, websocket: WebSocket):
        """发送个人消息"""
        try:
            await websocket.send_json(message)
        except Exception as e:
            logger.warning(f"发送个人消息失败: {e}")
    
    def get_online_users(self, workflow_id: str) -> List[Dict]:
        """获取在线用户列表"""
        if workflow_id not in self.workflow_users:
            return []
        return list(self.workflow_users[workflow_id].values())
    
    def _get_user_color(self, user_id: str) -> str:
        """为用户分配颜色（基于用户ID的哈希）"""
        colors = [
            "#FF6B6B", "#4ECDC4", "#45B7D1", "#FFA07A", "#98D8C8",
            "#F7DC6F", "#BB8FCE", "#85C1E2", "#F8B739", "#52BE80"
        ]
        hash_value = hash(user_id) % len(colors)
        return colors[hash_value]
    
    async def acquire_lock(self, workflow_id: str):
        """获取操作锁（用于冲突解决）"""
        if workflow_id not in self.operation_locks:
            self.operation_locks[workflow_id] = asyncio.Lock()
        return await self.operation_locks[workflow_id].acquire()
    
    def release_lock(self, workflow_id: str):
        """释放操作锁"""
        if workflow_id in self.operation_locks:
            self.operation_locks[workflow_id].release()


# 全局协作管理器实例
collaboration_manager = CollaborationManager()