48 lines
1.3 KiB
Python
48 lines
1.3 KiB
Python
"""WebSocket endpoint for real-time task progress."""
|
|
|
|
from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Query
|
|
from web.backend.auth.jwt_handler import decode_token
|
|
from jose import JWTError
|
|
|
|
router = APIRouter(tags=["websocket"])
|
|
|
|
|
|
@router.websocket("/ws/task/{task_id}")
|
|
async def task_websocket(
|
|
websocket: WebSocket,
|
|
task_id: str,
|
|
token: str = Query(...),
|
|
):
|
|
"""WebSocket for real-time task progress updates."""
|
|
try:
|
|
payload = decode_token(token)
|
|
username = payload.get("sub")
|
|
if not username:
|
|
await websocket.close(code=4001, reason="Invalid token")
|
|
return
|
|
except (JWTError, Exception):
|
|
await websocket.close(code=4001, reason="Invalid token")
|
|
return
|
|
|
|
await websocket.accept()
|
|
|
|
tm = websocket.app.state.task_manager
|
|
task = tm.get_task(task_id)
|
|
if not task:
|
|
await websocket.send_json({"error": "任务不存在"})
|
|
await websocket.close()
|
|
return
|
|
|
|
tm.subscribe(task_id, websocket)
|
|
await websocket.send_json(task.to_dict())
|
|
|
|
try:
|
|
while True:
|
|
data = await websocket.receive_text()
|
|
if data == "ping":
|
|
await websocket.send_text("pong")
|
|
except WebSocketDisconnect:
|
|
tm.unsubscribe(task_id, websocket)
|
|
except Exception:
|
|
tm.unsubscribe(task_id, websocket)
|