aiagent/backend/app/api/collaboration.py

"""
工作流协作API
支持多人实时协作编辑工作流
"""
from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Depends, HTTPException, status, Query
from sqlalchemy.orm import Session
from app.core.database import get_db, SessionLocal
from app.models.workflow import Workflow
from app.models.workspace import WorkspaceMembership
from app.api.auth import get_current_user
from app.models.user import User
from app.websocket.collaboration_manager import collaboration_manager
from app.core.exceptions import NotFoundError
import json
import logging

logger = logging.getLogger(__name__)

router = APIRouter(prefix="/api/v1/collaboration", tags=["collaboration"])


@router.websocket("/ws/workflows/{workflow_id}")
async def websocket_collaboration(
    websocket: WebSocket,
    workflow_id: str,
    token: str = None
):
    """
    工作流协作WebSocket连接
    
    支持多人实时协作编辑工作流，包括：
    - 节点添加/删除/移动
    - 边添加/删除
    - 节点配置修改
    - 实时同步变更
    
    参数:
    - token: JWT token（通过query参数传递，例如: ?token=xxx）
    """
    await websocket.accept()
    
    # 验证token
    if not token:
        await websocket.close(code=status.WS_1008_POLICY_VIOLATION, reason="缺少token")
        return
    
    db = SessionLocal()
    user = None
    
    try:
        # 验证JWT token
        from app.core.security import decode_access_token
        
        try:
            payload = decode_access_token(token)
            if payload is None:
                await websocket.close(code=status.WS_1008_POLICY_VIOLATION, reason="无效的token")
                return
            
            user_id_from_token = payload.get("sub")
            
            if not user_id_from_token:
                await websocket.close(code=status.WS_1008_POLICY_VIOLATION, reason="无效的token")
                return
            
            # 获取用户信息
            user = db.query(User).filter(User.id == user_id_from_token).first()
            if not user:
                await websocket.close(code=status.WS_1008_POLICY_VIOLATION, reason="用户不存在")
                return
            
        except Exception as e:
            logger.warning(f"Token验证失败: {e}")
            await websocket.close(code=status.WS_1008_POLICY_VIOLATION, reason="token验证失败")
            return
        
        # 获取工作流
        workflow = db.query(Workflow).filter(Workflow.id == workflow_id).first()
        if not workflow:
            await websocket.close(code=status.WS_1008_POLICY_VIOLATION, reason="工作流不存在")
            return
        
        # 检查 workspace 成员资格
        ws_id = payload.get("ws", "")
        if ws_id and user.role != "admin":
            membership = (
                db.query(WorkspaceMembership)
                .filter(
                    WorkspaceMembership.workspace_id == ws_id,
                    WorkspaceMembership.user_id == user.id,
                )
                .first()
            )
            if not membership:
                await websocket.close(code=status.WS_1008_POLICY_VIOLATION, reason="无权访问此工作区")
                return
            # 验证 workflow 属于该 workspace
            if workflow.workspace_id and workflow.workspace_id != ws_id:
                await websocket.close(code=status.WS_1008_POLICY_VIOLATION, reason="无权编辑此工作流")
                return

        # 检查权限（只有工作流所有者可以协作编辑，或者未来可以扩展权限）
        # 暂时只允许所有者编辑
        if workflow.user_id != user.id:
            await websocket.close(code=status.WS_1008_POLICY_VIOLATION, reason="无权限编辑此工作流")
            return
        
        user_id = user.id
        username = user.username
        
        # 建立协作连接
        await collaboration_manager.connect(websocket, workflow_id, user_id, username)
        
        try:
            # 持续监听消息
            while True:
                try:
                    # 接收客户端消息
                    data = await websocket.receive_text()
                    message = json.loads(data)
                    
                    message_type = message.get("type")
                    
                    if message_type == "ping":
                        # 心跳消息
                        await collaboration_manager.send_personal_message({
                            "type": "pong"
                        }, websocket)
                    
                    elif message_type == "operation":
                        # 工作流操作（节点/边的变更）
                        operation = message.get("operation", {})
                        op_type = operation.get("type")
                        
                        # 验证操作类型
                        valid_operations = [
                            "node_add", "node_delete", "node_move", "node_update",
                            "edge_add", "edge_delete", "edge_update",
                            "workflow_update"
                        ]
                        
                        if op_type not in valid_operations:
                            await collaboration_manager.send_personal_message({
                                "type": "error",
                                "message": f"无效的操作类型: {op_type}"
                            }, websocket)
                            continue
                        
                        # 添加操作者信息
                        operation["user_id"] = user_id
                        operation["username"] = username
                        
                        # 广播操作到其他用户
                        await collaboration_manager.broadcast_operation(
                            workflow_id,
                            operation,
                            exclude_websocket=websocket
                        )
                        
                        logger.info(f"用户 {username} 在工作流 {workflow_id} 执行操作: {op_type}")
                    
                    elif message_type == "cursor_move":
                        # 光标移动（可选功能）
                        cursor_info = message.get("cursor", {})
                        await collaboration_manager.broadcast_operation(
                            workflow_id,
                            {
                                "type": "cursor_move",
                                "user_id": user_id,
                                "username": username,
                                "cursor": cursor_info
                            },
                            exclude_websocket=websocket
                        )
                    
                    else:
                        await collaboration_manager.send_personal_message({
                            "type": "error",
                            "message": f"未知的消息类型: {message_type}"
                        }, websocket)
                
                except json.JSONDecodeError:
                    await collaboration_manager.send_personal_message({
                        "type": "error",
                        "message": "无效的JSON格式"
                    }, websocket)
        
        except WebSocketDisconnect:
            pass
        except Exception as e:
            logger.error(f"协作WebSocket错误: {e}", exc_info=True)
            try:
                await collaboration_manager.send_personal_message({
                    "type": "error",
                    "message": f"发生错误: {str(e)}"
                }, websocket)
            except:
                pass
        finally:
            collaboration_manager.disconnect(websocket, workflow_id)
            if db:
                db.close()
    
    except Exception as e:
        logger.error(f"建立协作连接失败: {e}", exc_info=True)
        try:
            await websocket.close()
        except:
            pass


@router.get("/workflows/{workflow_id}/users")
async def get_collaboration_users(
    workflow_id: str,
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user)
):
    """获取当前协作编辑工作流的用户列表"""
    # 验证工作流存在
    workflow = db.query(Workflow).filter(Workflow.id == workflow_id).first()
    if not workflow:
        raise NotFoundError("工作流", workflow_id)
    
    # 获取在线用户
    online_users = collaboration_manager.get_online_users(workflow_id)
    
    return {
        "workflow_id": workflow_id,
        "online_users": online_users,
        "count": len(online_users)
    }