Spaces:
Running
Running
| import time | |
| import asyncio | |
| import threading | |
| from fastapi import WebSocket | |
| from datetime import datetime | |
| connections: dict[str, set[WebSocket]] = {} | |
| last_seen_ping: dict[str, float] = {} | |
| def parse_db_timestamp(ts_val) -> int: | |
| """Helper to convert various DB timestamp formats (Unix string or ISO string) to Unix integer.""" | |
| if not ts_val: | |
| return 0 | |
| if isinstance(ts_val, (int, float)): | |
| return int(ts_val) | |
| if isinstance(ts_val, str): | |
| if ts_val.isdigit(): | |
| return int(ts_val) | |
| try: | |
| # Handle ISO format from Supabase | |
| clean_ts = ts_val.replace('Z', '+00:00') | |
| dt = datetime.fromisoformat(clean_ts) | |
| return int(dt.timestamp()) | |
| except Exception: | |
| return 0 | |
| return 0 | |
| # ========================= | |
| # WS MANAGEMENT | |
| # ========================= | |
| async def register_ws(username: str, ws: WebSocket, connections_dict=None): | |
| if connections_dict is None: | |
| connections_dict = connections | |
| uname = username.lower() | |
| if uname not in connections_dict: | |
| connections_dict[uname] = set() | |
| connections_dict[uname].add(ws) | |
| def unregister_ws(username: str, ws: WebSocket, connections_dict=None): | |
| if connections_dict is None: | |
| connections_dict = connections | |
| uname = username.lower() | |
| if uname in connections_dict: | |
| connections_dict[uname].discard(ws) | |
| if not connections_dict[uname]: | |
| del connections_dict[uname] | |
| async def send_to_user(username: str, payload: dict, connections_dict=None): | |
| if connections_dict is None: | |
| connections_dict = connections | |
| uname = username.lower() | |
| if uname in connections_dict: | |
| for ws in list(connections_dict[uname]): | |
| try: | |
| await ws.send_json(payload) | |
| except Exception: | |
| connections_dict[uname].discard(ws) | |
| async def send_follow_notification(target_user, sender, action, connections_dict=None): | |
| if connections_dict is None: | |
| connections_dict = connections | |
| if target_user.lower() in connections_dict: | |
| if action == "follow": | |
| content = f"{sender} started following you" | |
| elif action == "unfollow": | |
| content = f"{sender} unfollowed you" | |
| elif action == "remove_follower": | |
| content = f"{sender} removed you as a follower" | |
| else: | |
| content = f"{sender} updated their connection with you" | |
| payload = { | |
| "type": "notification", | |
| "subtype": action, | |
| "sender": sender, | |
| "content": content, | |
| "timestamp": int(time.time()), | |
| } | |
| for ws in list(connections_dict[target_user.lower()]): | |
| try: | |
| await ws.send_json(payload) | |
| except Exception as e: | |
| connections_dict[target_user.lower()].discard(ws) | |
| # ========================= | |
| # BROADCAST | |
| # ========================= | |
| async def broadcast_like_update(post_id: int, likes: int, timeout: float = 2.0): | |
| message = {"type": "like_update", "post_id": post_id, "likes": likes} | |
| for user_sockets in list(connections.values()): | |
| for ws in list(user_sockets): | |
| asyncio.create_task(_safe_send(ws, message)) | |
| async def broadcast_online(): | |
| now = time.time() | |
| online_users = [ | |
| u | |
| for u, conns in connections.items() | |
| if conns | |
| ] | |
| payload = {"type": "online_list", "users": online_users} | |
| for user_sockets in list(connections.values()): | |
| for ws in list(user_sockets): | |
| asyncio.create_task(_safe_send(ws, payload)) | |
| async def broadcast_avatar_update(username: str, avatar: str): | |
| payload = {"type": "avatar_update", "username": username, "avatar": avatar} | |
| for user_sockets in list(connections.values()): | |
| for ws in list(user_sockets): | |
| asyncio.create_task(_safe_send(ws, payload)) | |
| async def broadcast_hero_update(username: str, hero: str): | |
| payload = {"type": "heroBanner_update", "username": username, "heroBanner": hero} | |
| for user_sockets in list(connections.values()): | |
| for ws in list(user_sockets): | |
| asyncio.create_task(_safe_send(ws, payload)) | |
| async def broadcast_theme_update(username: str, theme: int): | |
| payload = {"type": "theme_update", "username": username, "theme": theme} | |
| uname = username.lower() | |
| if uname in connections: | |
| for ws in list(connections[uname]): | |
| asyncio.create_task(_safe_send(ws, payload)) | |
| async def broadcast_to_all(payload: dict): | |
| """Broadcasts a payload to every active WebSocket connection in parallel.""" | |
| tasks = [] | |
| for user_sockets in list(connections.values()): | |
| for ws in list(user_sockets): | |
| tasks.append(asyncio.create_task(_safe_send(ws, payload))) | |
| if tasks: | |
| # We don't necessarily need to await all of them to finish if we want fire-and-forget, | |
| # but let's at least wait a short time for them to start. | |
| await asyncio.sleep(0) | |
| async def _safe_send(ws: WebSocket, payload: dict, timeout: float = 1.0): | |
| try: | |
| await asyncio.wait_for(ws.send_json(payload), timeout=timeout) | |
| except Exception: | |
| # Socket is likely dead, but we don't unregister here to avoid modification during iteration | |
| # The next ping/pong or cleanup cycle will handle it. | |
| pass | |
| async def evict_user_ws(username: str, connections_dict=None): | |
| """Sends a force-logout signal to all active connections of a user.""" | |
| if connections_dict is None: | |
| connections_dict = connections | |
| uname = username.lower() | |
| if uname in connections_dict: | |
| payload = {"type": "evict_user"} | |
| # Make a copy of the set to avoid modification during iteration | |
| for ws in list(connections_dict[uname]): | |
| try: | |
| await ws.send_json(payload) | |
| await ws.close(code=1000, reason="Account deleted") | |
| except Exception: | |
| pass | |
| # Clean up the connections dictionary | |
| if uname in connections_dict: | |
| del connections_dict[uname] | |
| if uname in last_seen_ping: | |
| del last_seen_ping[uname] | |
| def upload_media_to_cloudinary(data_url: str) -> str: | |
| """Upload a base64 data URL to Cloudinary, return the optimized public URL.""" | |
| import base64, uuid, mimetypes, cloudinary.uploader, os | |
| if not data_url or not data_url.startswith("data:"): | |
| return data_url | |
| try: | |
| import cloudinary | |
| cloudinary.config( | |
| cloud_name=os.environ.get("CLOUDINARY_CLOUD_NAME"), | |
| api_key=os.environ.get("CLOUDINARY_API_KEY"), | |
| api_secret=os.environ.get("CLOUDINARY_API_SECRET") | |
| ) | |
| header, b64 = data_url.split(",", 1) | |
| mime = header.split(";")[0].split(":")[1] | |
| ext = mimetypes.guess_extension(mime) or ".bin" | |
| name = f"{uuid.uuid4().hex}" # filename without extension for public_id | |
| file_bytes = base64.b64decode(b64) | |
| upload_res = cloudinary.uploader.upload( | |
| file_bytes, | |
| public_id=name, | |
| folder="yapstation_media", | |
| resource_type="auto" | |
| ) | |
| # Use 'f_auto,q_auto' for best performance/speed | |
| url = upload_res['secure_url'] | |
| optimized_url = url.replace("/upload/", "/upload/f_auto,q_auto/") | |
| return optimized_url | |
| except Exception as e: | |
| print(f"Cloudinary Upload Error: {e}") | |
| return data_url |