You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

33 lines
900 B

"""
WebSocket进度路由
"""
from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Query
from app.core.progress_manager import progress_manager
import logging
logger = logging.getLogger(__name__)
router = APIRouter()
@router.websocket("/progress/{task_id}")
async def websocket_progress(
websocket: WebSocket,
task_id: str
):
"""WebSocket进度推送"""
await progress_manager.connect(websocket, task_id)
try:
while True:
data = await websocket.receive_text()
if data == "ping":
await websocket.send_text("pong")
elif data == "close":
break
except WebSocketDisconnect:
logger.info(f"WebSocket断开连接: task_id={task_id}")
except Exception as e:
logger.error(f"WebSocket错误: {e}")
finally:
await progress_manager.disconnect(websocket, task_id)