67 lines
2.0 KiB
Python
67 lines
2.0 KiB
Python
"""WebSocket progress publisher – decouples task manager from WebSocket."""
|
||
|
||
import asyncio
|
||
from typing import Dict, Set
|
||
|
||
from fastapi import WebSocket
|
||
|
||
# This will be set by main.py during startup
|
||
MAIN_LOOP = None
|
||
|
||
active_connections: Dict[str, Set[WebSocket]] = {}
|
||
|
||
|
||
def publish_progress(task_id: str, progress: int, message: str, status: str = "processing"):
|
||
"""
|
||
Publish progress update to all connected WebSocket clients for a task.
|
||
"""
|
||
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()
|
||
# Get the loop to use: either the stored one or try to get the current loop
|
||
loop = MAIN_LOOP
|
||
if loop is None:
|
||
try:
|
||
loop = asyncio.get_running_loop()
|
||
except RuntimeError:
|
||
# No running loop, fallback to default event loop
|
||
loop = asyncio.get_event_loop()
|
||
|
||
for websocket in active_connections.get(task_id, set()):
|
||
try:
|
||
asyncio.run_coroutine_threadsafe(websocket.send_json(data), loop)
|
||
except Exception:
|
||
to_remove.add(websocket)
|
||
|
||
# Clean up disconnected clients
|
||
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):
|
||
"""Register a WebSocket connection for a task."""
|
||
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):
|
||
"""Unregister a WebSocket connection for a task."""
|
||
if task_id in active_connections:
|
||
active_connections[task_id].discard(websocket)
|
||
if not active_connections[task_id]:
|
||
del active_connections[task_id]
|