Spaces:
Sleeping
Sleeping
| 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() | |