diff --git a/web/backend/api/__init__.py b/web/backend/api/__init__.py index d6cbb4b..c38d5c4 100644 --- a/web/backend/api/__init__.py +++ b/web/backend/api/__init__.py @@ -1 +1,4 @@ -"""API route handlers.""" \ No newline at end of file +"""API route handlers.""" + +from . import upload, split, status, download +# websocket is imported directly in main.py to avoid circular import \ No newline at end of file diff --git a/web/backend/api/websocket.py b/web/backend/api/websocket.py new file mode 100644 index 0000000..5f83519 --- /dev/null +++ b/web/backend/api/websocket.py @@ -0,0 +1,46 @@ +"""WebSocket endpoint for real‑time progress updates.""" + +import json + +from fastapi import APIRouter, WebSocket, WebSocketDisconnect + +from backend.services.task_manager import task_manager +from backend.services.progress_publisher import register_connection, unregister_connection + +router = APIRouter(tags=["websocket"]) + + +@router.websocket("/ws/{task_id}") +async def websocket_endpoint(websocket: WebSocket, task_id: str): + await websocket.accept() + + register_connection(task_id, websocket) + + try: + # Send initial status + status = task_manager.get_status(task_id) + await websocket.send_json({ + "type": "status", + "data": status + }) + + while True: + data = await websocket.receive_text() + try: + message = json.loads(data) + if message.get("type") == "ping": + await websocket.send_json({"type": "pong"}) + elif message.get("type") == "get_status": + status = task_manager.get_status(task_id) + await websocket.send_json({ + "type": "status", + "data": status + }) + except json.JSONDecodeError: + pass + + except WebSocketDisconnect: + unregister_connection(task_id, websocket) + except Exception as e: + unregister_connection(task_id, websocket) + print(f"WebSocket error: {e}") \ No newline at end of file diff --git a/web/backend/main.py b/web/backend/main.py index b38611d..a334ba7 100644 --- a/web/backend/main.py +++ b/web/backend/main.py @@ -5,6 +5,7 @@ from fastapi.middleware.cors import CORSMiddleware from backend.config import settings from backend.api import upload, split, status, download +from backend.api.websocket import router as websocket_router app = FastAPI( title="Audio Splitter Web API", @@ -26,6 +27,7 @@ app.include_router(upload.router) app.include_router(split.router) app.include_router(status.router) app.include_router(download.router) +app.include_router(websocket_router) @app.get("/") diff --git a/web/backend/services/progress_publisher.py b/web/backend/services/progress_publisher.py new file mode 100644 index 0000000..40df706 --- /dev/null +++ b/web/backend/services/progress_publisher.py @@ -0,0 +1,49 @@ +"""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] \ No newline at end of file diff --git a/web/backend/services/splits.zip b/web/backend/services/splits.zip new file mode 100644 index 0000000..8d9f32c Binary files /dev/null and b/web/backend/services/splits.zip differ diff --git a/web/backend/services/splitter.py b/web/backend/services/splitter.py index 826bc72..70eca14 100644 --- a/web/backend/services/splitter.py +++ b/web/backend/services/splitter.py @@ -15,7 +15,9 @@ from backend.models.response import TaskStatus def run_split_task(task_id: str, tracklist: List[TracklistEntry], options: Dict[str, Any]) -> None: try: - task_manager.update_task(task_id, progress=5, message="Initializing...") + task_manager.update_task_with_progress( + task_id, progress=5, message="Initializing..." + ) input_path = FileManager.get_input_path(task_id) if not input_path or not input_path.exists(): @@ -63,24 +65,45 @@ def run_split_task(task_id: str, tracklist: List[TracklistEntry], options: Dict[ from audio_splitter.core import split_audio - task_manager.update_task(task_id, progress=10, message="Starting split...") + # We need to track progress from the core. + # Since the core doesn't have a progress callback, we'll update progress + # based on track list size (approximate). + total_tracks = len(tracks) + progress_base = 10 # Starting progress after init + + task_manager.update_task_with_progress( + task_id, progress=progress_base, message="Starting split..." + ) + + # Run the split + # The core prints progress to stdout, but we can't easily capture it. + # We'll update progress based on track count. + # For now, we'll report progress after the split completes. + # A more advanced implementation would capture stdout or add a callback. split_audio(str(input_path), str(output_dir), tracks, args) + # After split completes, get the output files output_files = FileManager.get_output_files(task_id) - task_manager.update_task( + task_manager.update_task_with_progress( task_id, - status=TaskStatus.DONE, progress=100, message="Split complete", + status=TaskStatus.DONE + ) + + # Add tracks to the task state + task_manager.update_task( + task_id, tracks=output_files ) except Exception as e: - task_manager.update_task( + task_manager.update_task_with_progress( task_id, + progress=0, + message="Split failed", status=TaskStatus.ERROR, - error=str(e), - message="Split failed" + error=str(e) ) \ No newline at end of file diff --git a/web/backend/services/task_manager.py b/web/backend/services/task_manager.py index fb11a4f..aabe245 100644 --- a/web/backend/services/task_manager.py +++ b/web/backend/services/task_manager.py @@ -7,6 +7,7 @@ import time from backend.config import settings from backend.models.response import TaskStatus, TrackInfo +from backend.services.progress_publisher import publish_progress class TaskManager: @@ -33,8 +34,37 @@ class TaskManager: return task_id in self._tasks def get_status(self, task_id: str) -> dict: + """Get task status with datetime objects converted to ISO format.""" with self._lock: - return self._tasks.get(task_id, {}).copy() + if task_id not in self._tasks: + return {} + # Return a copy with datetime objects converted to ISO strings + status = self._tasks[task_id].copy() + return self._prepare_status_for_serialization(status) + + def _prepare_status_for_serialization(self, status: dict) -> dict: + """Convert datetime objects to ISO format strings for JSON serialization.""" + serializable = {} + for key, value in status.items(): + if isinstance(value, datetime): + serializable[key] = value.isoformat() + elif isinstance(value, list): + # Handle lists of objects (e.g., tracks) + serializable[key] = [ + {k: v.isoformat() if isinstance(v, datetime) else v for k, v in item.items()} + if isinstance(item, dict) + else item + for item in value + ] + elif isinstance(value, dict): + # Recursively handle nested dicts + serializable[key] = { + k: v.isoformat() if isinstance(v, datetime) else v + for k, v in value.items() + } + else: + serializable[key] = value + return serializable def update_task(self, task_id: str, **kwargs) -> None: with self._lock: @@ -42,6 +72,22 @@ class TaskManager: self._tasks[task_id].update(kwargs) self._tasks[task_id]["updated_at"] = datetime.now() + def update_task_with_progress(self, task_id: str, progress: int, message: str, + status: str = None, error: str = None) -> None: + update_kwargs = { + "progress": progress, + "message": message + } + if status: + update_kwargs["status"] = status + if error is not None: + update_kwargs["error"] = error + + self.update_task(task_id, **update_kwargs) + + # Publish progress via WebSocket + publish_progress(task_id, progress, message, status or "processing") + def cleanup_old_tasks(self) -> None: now = datetime.now() timeout = timedelta(seconds=settings.cleanup_after_seconds) diff --git a/web/requirements-web.txt b/web/requirements-web.txt index a7351d4..2a4fd0a 100644 --- a/web/requirements-web.txt +++ b/web/requirements-web.txt @@ -4,4 +4,5 @@ python-multipart>=0.0.6 aiofiles>=23.2.0 pydantic>=2.5.0 pydantic-settings>=2.0.0 -python-dotenv>=1.0.0 \ No newline at end of file +python-dotenv>=1.0.0 +websockets>=12.0 \ No newline at end of file