Spaces:
Sleeping
Sleeping
File size: 2,429 Bytes
b2be963 | 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 | import asyncio
from typing import Dict, List, Any, Optional
import json
from datetime import datetime, timezone
class NotificationManager:
def __init__(self):
# Dictionary of active queues: {user_id: [Queue, Queue, ...]}
# We use a list of queues per user to support multiple tabs/sessions
self.active_connections: Dict[int, List[asyncio.Queue]] = {}
self.lock = asyncio.Lock()
async def subscribe(self, user_id: int) -> asyncio.Queue:
"""Create a new queue for a user's SSE connection."""
queue = asyncio.Queue()
async with self.lock:
if user_id not in self.active_connections:
self.active_connections[user_id] = []
self.active_connections[user_id].append(queue)
return queue
async def unsubscribe(self, user_id: int, queue: asyncio.Queue):
"""Remove a queue when an SSE connection closes."""
async with self.lock:
if user_id in self.active_connections:
if queue in self.active_connections[user_id]:
self.active_connections[user_id].remove(queue)
if not self.active_connections[user_id]:
self.active_connections.pop(user_id, None)
async def broadcast(self, notification_data: Dict[str, Any], target_user_id: Optional[int] = None):
"""
Send a notification to active subscribers.
If target_user_id is None, broadcast to ALL admins (currently all admins see all notifications).
"""
# Prepare the message
# Convert datetime to string if present
if isinstance(notification_data.get("created_at"), datetime):
notification_data["created_at"] = notification_data["created_at"].isoformat()
message = json.dumps(notification_data)
async with self.lock:
if target_user_id:
# Send to specific user
if target_user_id in self.active_connections:
for queue in self.active_connections[target_user_id]:
await queue.put(message)
else:
# Broadcast to everyone (Admins)
for user_queues in self.active_connections.values():
for queue in user_queues:
await queue.put(message)
# Global Manager Instance
notification_manager = NotificationManager()
|