File size: 4,869 Bytes
bc1e6dc | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 | from fastapi import APIRouter, WebSocket, Depends, BackgroundTasks, WebSocketDisconnect
from uuid import UUID
from typing import Optional, Dict, Any
import json
from channel.socket_service import SocketService
router = APIRouter(prefix="/ws", tags=["websocket"])
# Singleton instance of SocketService
_socket_service: Optional[SocketService] = None
def get_socket_service() -> SocketService:
global _socket_service
if _socket_service is None:
_socket_service = SocketService()
return _socket_service
@router.websocket("/{user_id}/{device_id}")
async def websocket_endpoint(
websocket: WebSocket,
user_id: UUID,
device_id: str,
background_tasks: BackgroundTasks,
socket_service: SocketService = Depends(get_socket_service)
):
"""
WebSocket endpoint for real-time communication
Args:
websocket: The WebSocket connection
user_id: The ID of the user
device_id: The ID of the device
"""
try:
print(f"New connection request - User: {user_id}, Device: {device_id}")
# Store background tasks reference
socket_service.background_tasks = background_tasks
# Start cleanup task if not already running
background_tasks.add_task(socket_service.cleanup_expired_connections)
# Get initial metadata from client
metadata = None
try:
initial_message = await websocket.receive_json()
print(f"Initial message received: {initial_message}")
if isinstance(initial_message, dict):
metadata = initial_message.get("metadata", {})
except Exception:
# If no initial message or invalid, proceed without metadata
pass
# Connect the websocket
await socket_service.connect(websocket, user_id, device_id, metadata)
# Handle incoming messages
while True:
try:
message = await websocket.receive_json()
# Handle heartbeat responses
if isinstance(message, dict) and message.get("type") == "heartbeat_response":
await socket_service.update_heartbeat(user_id, device_id)
continue
# For other messages, broadcast to all other devices of this user
if isinstance(message, dict):
# Extract the content from the message
if message.get("type") == "message" and "content" in message:
# Create a new message with the content
broadcast_message = {
"type": "message",
"content": message["content"]
}
# Broadcast the content to all other devices of this user
broadcast_count = await socket_service.broadcast_message(
user_id=user_id,
message=broadcast_message,
exclude_device_id=device_id # Exclude the sender's device
)
# Send acknowledgment to the sender
await websocket.send_json({
"type": "message_status",
"status": "delivered",
"broadcast_count": broadcast_count
})
except WebSocketDisconnect:
break
except Exception as e:
# Send error message to client
try:
await websocket.send_json({
"type": "error",
"message": str(e)
})
except Exception:
break
finally:
# Clean up the connection
await socket_service.disconnect(user_id, device_id)
@router.post("/{user_id}/broadcast")
async def broadcast_to_user(
user_id: UUID,
message: Dict[str, Any],
sender_device_id: Optional[str] = None,
socket_service: SocketService = Depends(get_socket_service)
):
"""
Broadcast a message to all active connections of a user, optionally excluding a sender device
Args:
user_id: The ID of the user to broadcast to
message: The message to broadcast
sender_device_id: Optional device ID to exclude from broadcast
Returns:
Dict containing broadcast status and count of successful broadcasts
"""
broadcast_count = await socket_service.broadcast_message(user_id, message, exclude_device_id=sender_device_id)
return {
"status": "message_broadcast",
"broadcast_count": broadcast_count
} |