122 lines
4.0 KiB
Python
122 lines
4.0 KiB
Python
|
"""
|
|||
|
|
WebSocket API
|
|||
|
|
"""
|
|||
|
|
from fastapi import APIRouter, WebSocket, WebSocketDisconnect
|
|||
|
|
from app.websocket.manager import websocket_manager
|
|||
|
|
from app.core.database import SessionLocal
|
|||
|
|
from app.models.execution import Execution
|
|||
|
|
from typing import Optional
|
|||
|
|
import json
|
|||
|
|
import asyncio
|
|||
|
|
|
|||
|
|
|
|||
|
|
router = APIRouter()
|
|||
|
|
|
|||
|
|
|
|||
|
|
@router.websocket("/api/v1/ws/executions/{execution_id}")
|
|||
|
|
async def websocket_execution_status(
|
|||
|
|
websocket: WebSocket,
|
|||
|
|
execution_id: str,
|
|||
|
|
token: Optional[str] = None
|
|||
|
|
):
|
|||
|
|
"""
|
|||
|
|
WebSocket实时推送执行状态
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
websocket: WebSocket连接
|
|||
|
|
execution_id: 执行记录ID
|
|||
|
|
token: JWT Token（可选，通过query参数传递）
|
|||
|
|
"""
|
|||
|
|
# 验证token（可选，如果需要认证）
|
|||
|
|
# user = await get_current_user_optional(token)
|
|||
|
|
|
|||
|
|
# 建立连接
|
|||
|
|
await websocket_manager.connect(websocket, execution_id)
|
|||
|
|
|
|||
|
|
db = SessionLocal()
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
# 发送初始状态
|
|||
|
|
execution = db.query(Execution).filter(Execution.id == execution_id).first()
|
|||
|
|
if execution:
|
|||
|
|
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
|
|||
|
|
|
|||
|
|
# 持续监听并推送状态更新
|
|||
|
|
while True:
|
|||
|
|
try:
|
|||
|
|
# 接收客户端消息（心跳等）
|
|||
|
|
data = await websocket.receive_text()
|
|||
|
|
|
|||
|
|
# 处理客户端消息
|
|||
|
|
try:
|
|||
|
|
message = json.loads(data)
|
|||
|
|
if message.get("type") == "ping":
|
|||
|
|
await websocket_manager.send_personal_message({
|
|||
|
|
"type": "pong"
|
|||
|
|
}, websocket)
|
|||
|
|
except:
|
|||
|
|
pass
|
|||
|
|
|
|||
|
|
except WebSocketDisconnect:
|
|||
|
|
break
|
|||
|
|
|
|||
|
|
# 检查执行状态
|
|||
|
|
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
|
|||
|
|
|
|||
|
|
# 定期发送状态更新（每2秒）
|
|||
|
|
await asyncio.sleep(2)
|
|||
|
|
|
|||
|
|
# 重新查询执行状态
|
|||
|
|
db.refresh(execution)
|
|||
|
|
await websocket_manager.send_personal_message({
|
|||
|
|
"type": "status",
|
|||
|
|
"execution_id": execution_id,
|
|||
|
|
"status": execution.status,
|
|||
|
|
"progress": 50 if execution.status == "running" else 0,
|
|||
|
|
"message": f"执行中..." if execution.status == "running" else "等待执行"
|
|||
|
|
}, websocket)
|
|||
|
|
|
|||
|
|
except WebSocketDisconnect:
|
|||
|
|
pass
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"WebSocket错误: {e}")
|
|||
|
|
try:
|
|||
|
|
await websocket_manager.send_personal_message({
|
|||
|
|
"type": "error",
|
|||
|
|
"message": f"发生错误: {str(e)}"
|
|||
|
|
}, websocket)
|
|||
|
|
except:
|
|||
|
|
pass
|
|||
|
|
finally:
|
|||
|
|
websocket_manager.disconnect(websocket, execution_id)
|
|||
|
|
db.close()
|