aiagent/backend/app/api/websocket.py

"""
WebSocket API — 实时推送执行状态（workspace 隔离）
"""
from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Query
from app.websocket.manager import websocket_manager
from app.core.database import SessionLocal
from app.core.security import decode_access_token
from app.models.execution import Execution
from app.models.user import User
from app.models.workspace import WorkspaceMembership
from typing import Optional
import json
import asyncio
import logging

logger = logging.getLogger(__name__)

router = APIRouter()


def _get_progress_from_redis(execution_id: str) -> Optional[dict]:
    """从 Redis 读取进度数据。"""
    try:
        from app.core.redis_client import get_redis_client
        redis_client = get_redis_client()
        if redis_client:
            raw = redis_client.get(f"workflow:progress:{execution_id}")
            if raw:
                return json.loads(raw)
    except Exception:
        pass
    return None


def _validate_ws_token(token: Optional[str], execution_id: str) -> dict:
    """验证 WebSocket JWT token，返回 {user_id, workspace_id}。

    失败时返回包含 error/code 的字典，调用方应据此拒绝连接。
    """
    if not token:
        return {"error": "缺少认证令牌", "code": 4001}

    payload = decode_access_token(token)
    if payload is None:
        return {"error": "无效的访问令牌", "code": 4001}

    user_id = payload.get("sub")
    ws_id = payload.get("ws", "")
    if not user_id:
        return {"error": "无效的令牌载荷", "code": 4001}

    # 验证用户存在且未被删除
    db = SessionLocal()
    try:
        user = db.query(User).filter(User.id == user_id).first()
        if not user:
            return {"error": "用户不存在", "code": 4001}
        if user.status == "deleted":
            return {"error": "账号已注销", "code": 4003}

        # 平台管理员绕过 workspace 检查
        if user.role == "admin":
            return {"user_id": user_id, "workspace_id": ws_id, "is_admin": True}

        # 验证 workspace 成员资格
        if ws_id:
            membership = (
                db.query(WorkspaceMembership)
                .filter(
                    WorkspaceMembership.workspace_id == ws_id,
                    WorkspaceMembership.user_id == user_id,
                )
                .first()
            )
            if not membership:
                return {"error": "无权访问此工作区", "code": 4003}

        # 验证 execution 属于该 workspace
        execution = db.query(Execution).filter(Execution.id == execution_id).first()
        if execution and ws_id and execution.workspace_id and execution.workspace_id != ws_id:
            return {"error": "无权查看此执行记录", "code": 4003}

        return {"user_id": user_id, "workspace_id": ws_id}
    finally:
        db.close()


@router.websocket("/api/v1/ws/executions/{execution_id}")
async def websocket_execution_status(
    websocket: WebSocket,
    execution_id: str,
    token: Optional[str] = Query(None),
):
    """
    WebSocket 实时推送执行状态（workspace 隔离）。

    需要 JWT token 通过 query 参数传递: ?token=xxx
    连接建立前验证：JWT 有效性 → workspace 成员资格 → execution workspace 归属。
    """
    # ── 认证 + 鉴权 ──
    auth = _validate_ws_token(token, execution_id)
    if "error" in auth:
        await websocket.accept()
        await websocket.send_json({"type": "error", "message": auth["error"], "code": auth["code"]})
        await websocket.close(code=auth["code"])
        return

    user_id = auth["user_id"]
    ws_id = auth.get("workspace_id", "")

    await websocket_manager.connect(websocket, execution_id, workspace_id=ws_id)

    db = SessionLocal()

    try:
        # 发送初始状态
        execution = db.query(Execution).filter(Execution.id == execution_id).first()
        if execution:
            # 尝试从 Redis 读取进度
            redis_progress = _get_progress_from_redis(execution_id)
            if redis_progress:
                await websocket_manager.send_personal_message({
                    "type": "progress",
                    "execution_id": execution_id,
                    **redis_progress,
                }, websocket)
            else:
                await websocket_manager.send_personal_message({
                    "type": "status",
                    "execution_id": execution_id,
                    "status": execution.status,
                    "progress": 0,
                    "message": "连接已建立"
                }, websocket)
        else:
            await websocket_manager.send_personal_message({
                "type": "error",
                "message": f"执行记录 {execution_id} 不存在"
            }, websocket)
            await websocket.close()
            return

        last_progress = -1

        # 持续监听并推送状态更新
        while True:
            try:
                # 接收客户端消息（心跳等），超时 1 秒以便轮询进度
                try:
                    data = await asyncio.wait_for(websocket.receive_text(), timeout=1.0)
                    try:
                        message = json.loads(data)
                        if message.get("type") == "ping":
                            await websocket_manager.send_personal_message({
                                "type": "pong"
                            }, websocket)
                    except Exception:
                        pass
                except asyncio.TimeoutError:
                    pass  # 超时后轮询进度

            except WebSocketDisconnect:
                break

            # 优先从 Redis 读取进度（推送模式）
            redis_progress = _get_progress_from_redis(execution_id)
            if redis_progress:
                pct = redis_progress.get("progress", -1)
                if pct != last_progress:
                    last_progress = pct
                    await websocket_manager.send_personal_message({
                        "type": "progress",
                        "execution_id": execution_id,
                        **redis_progress,
                    }, websocket)
            else:
                # Redis 不可用时回退到 DB 轮询
                db.refresh(execution)
                pct = 100 if execution.status in ["completed", "failed"] else (50 if execution.status == "running" else 0)
                if pct != last_progress:
                    last_progress = pct
                    await websocket_manager.send_personal_message({
                        "type": "status",
                        "execution_id": execution_id,
                        "status": execution.status,
                        "progress": pct,
                        "message": f"执行中..." if execution.status == "running" else "等待执行"
                    }, websocket)

            # 如果执行完成或失败，发送最终状态并断开
            db.refresh(execution)
            if execution.status in ["completed", "failed"]:
                await websocket_manager.send_personal_message({
                    "type": "status",
                    "execution_id": execution_id,
                    "status": execution.status,
                    "progress": 100,
                    "result": execution.output_data if execution.status == "completed" else None,
                    "error": execution.error_message if execution.status == "failed" else None,
                    "execution_time": execution.execution_time
                }, websocket)

                await asyncio.sleep(1)
                break

    except WebSocketDisconnect:
        pass
    except Exception as e:
        logger.error("WebSocket 错误: %s", e)
        try:
            await websocket_manager.send_personal_message({
                "type": "error",
                "message": f"发生错误: {str(e)}"
            }, websocket)
        except Exception:
            pass
    finally:
        websocket_manager.disconnect(websocket, execution_id)
        db.close()