288 lines
12 KiB
Python
288 lines
12 KiB
Python
"""
|
|
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}") |