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