Frontend and backend work together fine, complete workflow is implemented
This commit is contained in:
@@ -1,15 +1,20 @@
|
||||
"""WebSocket progress publisher – decouples task manager from WebSocket."""
|
||||
|
||||
import json
|
||||
import asyncio
|
||||
from typing import Dict, Set
|
||||
|
||||
from fastapi import WebSocket
|
||||
|
||||
# Active WebSocket connections
|
||||
# This will be set by main.py when the app starts
|
||||
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
|
||||
|
||||
@@ -26,10 +31,14 @@ def publish_progress(task_id: str, progress: int, message: str, status: str = "p
|
||||
to_remove = set()
|
||||
for websocket in active_connections.get(task_id, set()):
|
||||
try:
|
||||
websocket.send_json(data)
|
||||
# Use the stored main loop, or fallback to getting the current loop
|
||||
loop = MAIN_LOOP or asyncio.get_running_loop()
|
||||
asyncio.run_coroutine_threadsafe(websocket.send_json(data), loop)
|
||||
except Exception:
|
||||
# Client disconnected or other error
|
||||
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]:
|
||||
@@ -37,13 +46,15 @@ def publish_progress(task_id: str, progress: int, message: str, status: str = "p
|
||||
|
||||
|
||||
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]
|
||||
del active_connections[task_id]
|
||||
|
||||
Reference in New Issue
Block a user