49 lines
1.4 KiB
Python
49 lines
1.4 KiB
Python
"""WebSocket progress publisher – decouples task manager from WebSocket."""
|
||
|
||
import json
|
||
from typing import Dict, Set
|
||
|
||
from fastapi import WebSocket
|
||
|
||
# Active WebSocket connections
|
||
active_connections: Dict[str, Set[WebSocket]] = {}
|
||
|
||
|
||
def publish_progress(task_id: str, progress: int, message: str, status: str = "processing"):
|
||
if task_id not in active_connections:
|
||
return
|
||
|
||
data = {
|
||
"type": "progress",
|
||
"data": {
|
||
"task_id": task_id,
|
||
"status": status,
|
||
"progress": progress,
|
||
"message": message
|
||
}
|
||
}
|
||
|
||
to_remove = set()
|
||
for websocket in active_connections.get(task_id, set()):
|
||
try:
|
||
websocket.send_json(data)
|
||
except Exception:
|
||
to_remove.add(websocket)
|
||
|
||
for websocket in to_remove:
|
||
active_connections[task_id].discard(websocket)
|
||
if task_id in active_connections and not active_connections[task_id]:
|
||
del active_connections[task_id]
|
||
|
||
|
||
def register_connection(task_id: str, websocket: WebSocket):
|
||
if task_id not in active_connections:
|
||
active_connections[task_id] = set()
|
||
active_connections[task_id].add(websocket)
|
||
|
||
|
||
def unregister_connection(task_id: str, websocket: WebSocket):
|
||
if task_id in active_connections:
|
||
active_connections[task_id].discard(websocket)
|
||
if not active_connections[task_id]:
|
||
del active_connections[task_id] |