""" WebSocket endpoints for real-time video processing updates (Task 14.1). Provides live progress updates, transcript streaming, and browser notifications. """ import logging import json from typing import Optional, Dict, Any from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Query from backend.core.websocket_manager import websocket_manager, ProcessingStage, ProgressData logger = logging.getLogger(__name__) router = APIRouter() @router.websocket("/ws/process/{job_id}") async def websocket_processing_endpoint( websocket: WebSocket, job_id: str, user_id: Optional[str] = Query(None) ): """ WebSocket endpoint for real-time video processing updates. Args: websocket: WebSocket connection job_id: Processing job ID to monitor user_id: Optional user ID for authenticated users Message Types: - progress_update: Processing stage and percentage updates - completion_notification: Job completed successfully - error_notification: Processing error occurred - system_message: System-wide announcements - heartbeat: Connection keep-alive Client Message Types: - ping: Connection health check - subscribe_transcript: Enable live transcript streaming - unsubscribe_transcript: Disable transcript streaming - request_status: Get current job status """ try: # Connect the WebSocket for processing updates await websocket_manager.connect(websocket, job_id) # Send initial connection confirmation await websocket.send_json({ "type": "connection_status", "status": "connected", "message": "WebSocket connection established for processing", "job_id": job_id, "user_id": user_id, "supported_messages": [ "progress_update", "completion_notification", "error_notification", "transcript_chunk", "browser_notification", "system_message", "heartbeat" ] }) logger.info(f"Processing WebSocket connected: job={job_id}, user={user_id}") # Keep connection alive and handle incoming messages while True: try: # Wait for messages from the client data = await websocket.receive_json() message_type = data.get("type") if message_type == "ping": # Handle ping/pong for connection health await websocket.send_json({ "type": "pong", "timestamp": logger.handlers[0].formatter.formatTime( logger.makeRecord("", 0, "", 0, "", (), None), None ) if logger.handlers else None }) logger.debug(f"Ping received from job {job_id}") elif message_type == "subscribe_transcript": # Enable live transcript streaming for this connection websocket_manager.enable_transcript_streaming(websocket, job_id) logger.info(f"Enabling transcript streaming for job {job_id}") await websocket.send_json({ "type": "subscription_confirmed", "subscription": "transcript_streaming", "job_id": job_id, "status": "enabled", "message": "You will now receive live transcript chunks as they are processed" }) elif message_type == "unsubscribe_transcript": # Disable transcript streaming for this connection websocket_manager.disable_transcript_streaming(websocket, job_id) logger.info(f"Disabling transcript streaming for job {job_id}") await websocket.send_json({ "type": "subscription_confirmed", "subscription": "transcript_streaming", "job_id": job_id, "status": "disabled", "message": "Transcript streaming has been disabled" }) elif message_type == "request_status": # Send current job status if available stats = websocket_manager.get_stats() job_info = { "job_id": job_id, "connections": stats.get("job_connections", {}).get(job_id, 0), "has_active_processing": job_id in stats.get("active_jobs", []) } await websocket.send_json({ "type": "status_response", "job_id": job_id, "data": job_info }) logger.debug(f"Status request handled for job {job_id}") elif message_type == "cancel_job": # Handle job cancellation request logger.info(f"Job cancellation requested for {job_id}") await websocket.send_json({ "type": "cancellation_acknowledged", "job_id": job_id, "message": "Cancellation request received and forwarded to processing service" }) # Note: Actual job cancellation logic would be handled by the pipeline service # This just acknowledges the request via WebSocket else: logger.warning(f"Unknown message type '{message_type}' from job {job_id}") await websocket.send_json({ "type": "error", "message": f"Unknown message type: {message_type}", "supported_types": ["ping", "subscribe_transcript", "unsubscribe_transcript", "request_status", "cancel_job"] }) except WebSocketDisconnect: logger.info(f"Processing WebSocket disconnected: job={job_id}, user={user_id}") break except json.JSONDecodeError: logger.error(f"Invalid JSON received from job {job_id}") await websocket.send_json({ "type": "error", "message": "Invalid JSON format in message" }) except Exception as e: logger.error(f"Error handling WebSocket message for job {job_id}: {e}") # Send error to client await websocket.send_json({ "type": "error", "message": f"Error processing message: {str(e)}", "error_type": "processing_error" }) except WebSocketDisconnect: logger.info(f"Processing WebSocket disconnected during setup: job={job_id}, user={user_id}") except Exception as e: logger.error(f"Error in processing WebSocket endpoint for job {job_id}: {e}") try: await websocket.send_json({ "type": "error", "message": "WebSocket connection error", "error_details": str(e) }) except: pass # Connection might already be closed finally: # Clean up the connection websocket_manager.disconnect(websocket) logger.info(f"Processing WebSocket cleanup completed: job={job_id}, user={user_id}") @router.websocket("/ws/system") async def websocket_system_endpoint(websocket: WebSocket): """ WebSocket endpoint for system-wide notifications and status. Provides real-time updates about system health, maintenance, etc. """ try: # Connect without job_id for system-wide updates await websocket_manager.connect(websocket) # Send initial system status stats = websocket_manager.get_stats() await websocket.send_json({ "type": "system_status", "status": "connected", "message": "Connected to system notifications", "data": { "total_connections": stats.get("total_connections", 0), "active_jobs": len(stats.get("active_jobs", [])), "server_status": "online" } }) logger.info("System WebSocket connected") # Keep connection alive while True: try: data = await websocket.receive_json() message_type = data.get("type") if message_type == "ping": await websocket.send_json({"type": "pong"}) elif message_type == "get_stats": # Send current system statistics current_stats = websocket_manager.get_stats() await websocket.send_json({ "type": "system_stats", "data": current_stats }) else: logger.warning(f"Unknown system message type: {message_type}") except WebSocketDisconnect: logger.info("System WebSocket disconnected") break except Exception as e: logger.error(f"Error in system WebSocket: {e}") break except Exception as e: logger.error(f"Error in system WebSocket endpoint: {e}") finally: websocket_manager.disconnect(websocket) logger.info("System WebSocket cleanup completed") @router.websocket("/ws/notifications") async def websocket_notifications_endpoint( websocket: WebSocket, user_id: Optional[str] = Query(None) ): """ WebSocket endpoint for browser notifications. Sends notifications for job completions, errors, and system events. """ try: await websocket_manager.connect(websocket) # Send connection confirmation await websocket.send_json({ "type": "notifications_ready", "status": "connected", "user_id": user_id, "message": "Ready to receive browser notifications" }) logger.info(f"Notifications WebSocket connected for user {user_id}") # Keep connection alive for receiving notifications while True: try: data = await websocket.receive_json() message_type = data.get("type") if message_type == "ping": await websocket.send_json({"type": "pong"}) elif message_type == "notification_preferences": # Handle notification preferences from client preferences = data.get("preferences", {}) logger.info(f"Notification preferences updated for user {user_id}: {preferences}") await websocket.send_json({ "type": "preferences_updated", "preferences": preferences, "message": "Notification preferences saved" }) else: logger.warning(f"Unknown notifications message type: {message_type}") except WebSocketDisconnect: logger.info(f"Notifications WebSocket disconnected for user {user_id}") break except Exception as e: logger.error(f"Error in notifications WebSocket: {e}") break except Exception as e: logger.error(f"Error in notifications WebSocket endpoint: {e}") finally: websocket_manager.disconnect(websocket) logger.info(f"Notifications WebSocket cleanup completed for user {user_id}")