diff --git a/backend/app/gateway/__init__.py b/backend/app/gateway/__init__.py new file mode 100644 index 0000000..2276f53 --- /dev/null +++ b/backend/app/gateway/__init__.py @@ -0,0 +1 @@ +"""WebSocket real-time gateway package.""" diff --git a/backend/app/gateway/ws.py b/backend/app/gateway/ws.py new file mode 100644 index 0000000..cc89c7f --- /dev/null +++ b/backend/app/gateway/ws.py @@ -0,0 +1,538 @@ +"""Real-time WebSocket gateway. + +Multiplexes presence, DM, and world-chat delivery over a single WebSocket +endpoint (docs/planning/social.md US-2, US-4). + +Key features: +- Authentication via JWT token (query param or initial auth frame). +- Scoped subscriptions to DM conversations and world channels. +- Server-derived presence heartbeat (Redis TTL refreshed periodically). +- Broadcast presence changes (online/offline) to mutual friends. +- Fan-out via Redis pub/sub. +- Thread-safe & event-loop-safe dispatch for multi-instance and test environments. +- Graceful reconnect without duplicate delivery or orphaned state. +""" + +from __future__ import annotations + +import asyncio +from contextlib import asynccontextmanager +from datetime import datetime, timezone +import json +import logging +from typing import Any + +from fastapi import APIRouter, WebSocket, WebSocketDisconnect, status +from jose import JWTError, jwt +from sqlalchemy import or_, select +from sqlalchemy.ext.asyncio import AsyncSession +import redis.asyncio as aioredis + +from app.config import settings +from app.database import async_session_factory, get_db +from app.models.chat import ConversationParticipant +from app.models.social import Friendship, FriendshipStatus +from app.models.user import User +from app.models.world import World, WorldVisibility +from app.services.presence import presence_store + +logger = logging.getLogger(__name__) + +router = APIRouter() + + +@asynccontextmanager +async def get_db_session(): + """Helper to acquire an async DB session respecting test overrides if present.""" + from app.main import app + + override = app.dependency_overrides.get(get_db) + if override: + res = override() + if hasattr(res, "__anext__"): + session = await res.__anext__() + try: + yield session + finally: + try: + await res.__anext__() + except (StopAsyncIteration, Exception): + pass + return + elif hasattr(res, "__next__"): + session = next(res) + try: + yield session + finally: + try: + next(res) + except (StopIteration, Exception): + pass + return + + async with async_session_factory() as session: + yield session + + +async def authenticate_token(token: str) -> User | None: + """Validate JWT token and return authenticated User, or None.""" + try: + payload = jwt.decode(token, settings.public_pem, algorithms=[settings.ALGORITHM]) + user_id_raw = payload.get("sub") + if user_id_raw is None: + return None + user_id = int(user_id_raw) + except (JWTError, ValueError, TypeError): + return None + + async with get_db_session() as db: + res = await db.execute(select(User).where(User.id == user_id)) + user = res.scalar_one_or_none() + if user and user.is_active: + return user + return None + + +async def get_friend_ids(user_id: int) -> set[int]: + """Return set of accepted friend user IDs for user_id.""" + try: + async with get_db_session() as db: + res = await db.execute( + select(Friendship).where( + or_(Friendship.requester_id == user_id, Friendship.addressee_id == user_id), + Friendship.status == FriendshipStatus.ACCEPTED, + ) + ) + friendships = res.scalars().all() + return { + f.addressee_id if f.requester_id == user_id else f.requester_id + for f in friendships + } + except Exception as exc: + logger.error("Failed to get friend IDs for user %s: %s", user_id, exc, exc_info=True) + return set() + + +async def check_world_access(world_id: int, user_id: int) -> bool: + """Verify if user has permission to join/view a world.""" + async with get_db_session() as db: + res = await db.execute(select(World).where(World.id == world_id)) + world = res.scalar_one_or_none() + if not world: + return False + if world.visibility == WorldVisibility.PUBLIC: + return True + if world.owner_id == user_id: + return True + if world.visibility == WorldVisibility.FRIENDS_ONLY: + res_fr = await db.execute( + select(Friendship.id).where( + or_( + and_(Friendship.requester_id == world.owner_id, Friendship.addressee_id == user_id), + and_(Friendship.requester_id == user_id, Friendship.addressee_id == world.owner_id), + ), + Friendship.status == FriendshipStatus.ACCEPTED, + ).limit(1) + ) + return res_fr.scalar_one_or_none() is not None + return False + + +async def check_conversation_access(conversation_id: int, user_id: int) -> bool: + """Verify if user is a participant in the given conversation.""" + async with get_db_session() as db: + res = await db.execute( + select(ConversationParticipant.id).where( + ConversationParticipant.conversation_id == conversation_id, + ConversationParticipant.user_id == user_id, + ).limit(1) + ) + return res.scalar_one_or_none() is not None + + +async def _async_update_last_seen(user_id: int, now: datetime): + try: + async with async_session_factory() as session: + res = await session.execute(select(User).where(User.id == user_id)) + db_user = res.scalar_one_or_none() + if db_user: + db_user.last_seen_at = now + await session.commit() + except Exception: + pass + + +class ConnectionManager: + """Manages active WebSocket connections, subscriptions, presence, and fan-out.""" + + def __init__(self): + # ws -> User + self._connections: dict[WebSocket, User] = {} + # ws -> event loop id + self._ws_loops: dict[WebSocket, int] = {} + # user_id -> set of active WebSockets + self._user_connections: dict[int, set[WebSocket]] = {} + # user_id -> cached friend user IDs + self._user_friend_ids: dict[int, set[int]] = {} + # world_id -> set of WebSockets + self._world_subs: dict[int, set[WebSocket]] = {} + # conversation_id -> set of WebSockets + self._dm_subs: dict[int, set[WebSocket]] = {} + # ws -> asyncio heartbeat task + self._heartbeat_tasks: dict[WebSocket, asyncio.Task] = {} + # Per-event-loop Redis pub/sub listener tasks and events + self._redis_tasks: dict[int, asyncio.Task] = {} + self._subscribed_events: dict[int, asyncio.Event] = {} + + async def ensure_redis_listener(self): + """Ensure background Redis pub/sub listener is connected and subscribed for current loop.""" + loop = asyncio.get_running_loop() + loop_key = id(loop) + task = self._redis_tasks.get(loop_key) + if task is None or task.done(): + event = asyncio.Event() + self._subscribed_events[loop_key] = event + task = asyncio.create_task(self._redis_listener(loop_key, event)) + self._redis_tasks[loop_key] = task + else: + event = self._subscribed_events.get(loop_key) + + if event: + try: + await asyncio.wait_for(event.wait(), timeout=2.0) + except asyncio.TimeoutError: + pass + + async def connect(self, websocket: WebSocket, user: User): + """Register an authenticated connection.""" + await self.ensure_redis_listener() + loop_key = id(asyncio.get_running_loop()) + + self._connections[websocket] = user + self._ws_loops[websocket] = loop_key + self._user_connections.setdefault(user.id, set()).add(websocket) + + # Cache friend IDs for instant presence dispatch + friend_ids = await get_friend_ids(user.id) + self._user_friend_ids[user.id] = friend_ids + + # Mark user online in Redis and DB + await presence_store.heartbeat(user.id) + now = datetime.now(timezone.utc) + asyncio.create_task(_async_update_last_seen(user.id, now)) + + # Start periodic heartbeat task (every 20 seconds) + self._heartbeat_tasks[websocket] = asyncio.create_task( + self._run_periodic_heartbeat(user.id) + ) + + # Broadcast online presence to friends via Redis + await self.broadcast_presence(user.id, user.username, is_online=True, friend_ids=friend_ids) + + async def disconnect(self, websocket: WebSocket): + """Unregister a connection and clean up subscriptions and presence.""" + user = self._connections.pop(websocket, None) + self._ws_loops.pop(websocket, None) + task = self._heartbeat_tasks.pop(websocket, None) + if task: + task.cancel() + + # Clean subscriptions + for world_id in list(self._world_subs.keys()): + self._world_subs[world_id].discard(websocket) + if not self._world_subs[world_id]: + self._world_subs.pop(world_id, None) + + for conv_id in list(self._dm_subs.keys()): + self._dm_subs[conv_id].discard(websocket) + if not self._dm_subs[conv_id]: + self._dm_subs.pop(conv_id, None) + + if user is None: + return + + user_id = user.id + username = user.username + user_set = self._user_connections.get(user_id) + if user_set: + user_set.discard(websocket) + if not user_set: + self._user_connections.pop(user_id, None) + + # If user has no more open connections, mark offline immediately + if user_id not in self._user_connections: + await presence_store.clear_heartbeat(user_id) + now = datetime.now(timezone.utc) + asyncio.create_task(_async_update_last_seen(user_id, now)) + friend_ids = self._user_friend_ids.pop(user_id, set()) + await self.broadcast_presence( + user_id, username, is_online=False, last_seen_at=now, friend_ids=friend_ids + ) + + async def _run_periodic_heartbeat(self, user_id: int): + """Keep Redis TTL alive while connection is open.""" + try: + while True: + await asyncio.sleep(20) + await presence_store.heartbeat(user_id) + except asyncio.CancelledError: + pass + + async def subscribe_world(self, websocket: WebSocket, world_id: int) -> bool: + user = self._connections.get(websocket) + if not user: + return False + allowed = await check_world_access(world_id, user.id) + if not allowed: + return False + self._world_subs.setdefault(world_id, set()).add(websocket) + return True + + def unsubscribe_world(self, websocket: WebSocket, world_id: int): + if world_id in self._world_subs: + self._world_subs[world_id].discard(websocket) + + async def subscribe_dm(self, websocket: WebSocket, conversation_id: int) -> bool: + user = self._connections.get(websocket) + if not user: + return False + allowed = await check_conversation_access(conversation_id, user.id) + if not allowed: + return False + self._dm_subs.setdefault(conversation_id, set()).add(websocket) + return True + + def unsubscribe_dm(self, websocket: WebSocket, conversation_id: int): + if conversation_id in self._dm_subs: + self._dm_subs[conversation_id].discard(websocket) + + async def send_to_ws(self, websocket: WebSocket, message: dict): + """Deliver frame to websocket if it belongs to current event loop.""" + cur_loop = id(asyncio.get_running_loop()) + if self._ws_loops.get(websocket) != cur_loop: + return + try: + await websocket.send_text(json.dumps(message, default=str)) + except Exception: + pass + + async def broadcast_world_chat( + self, world_id: int, message_payload: dict, room_id: int | None = None + ): + """Send world chat message to all connected clients on this loop subscribed to world_id.""" + targets = list(self._world_subs.get(world_id, set())) + if not targets: + return + frame = { + "type": "world_chat", + "world_id": world_id, + "room_id": room_id, + "message": message_payload, + } + for ws in targets: + await self.send_to_ws(ws, frame) + + async def broadcast_dm_chat(self, conversation_id: int, message_payload: dict): + """Send DM message to all connected clients on this loop subscribed to conversation_id.""" + targets = list(self._dm_subs.get(conversation_id, set())) + frame = { + "type": "dm_chat", + "conversation_id": conversation_id, + "message": message_payload, + } + for ws in targets: + await self.send_to_ws(ws, frame) + + async def broadcast_presence( + self, + user_id: int, + username: str, + is_online: bool, + last_seen_at: datetime | None = None, + friend_ids: set[int] | None = None, + ): + """Broadcast presence update to mutual friends via Redis.""" + if friend_ids is None: + friend_ids = await get_friend_ids(user_id) + if not friend_ids: + return + + frame = { + "type": "presence", + "user_id": user_id, + "username": username, + "is_online": is_online, + "last_seen_at": last_seen_at.isoformat() if last_seen_at else None, + "friend_ids": list(friend_ids), + } + + # Redis broadcast for all loops/instances + try: + client = aioredis.from_url(settings.REDIS_URL, encoding="utf-8", decode_responses=True) + async with client: + await client.publish("presence:events", json.dumps(frame, default=str)) + except Exception as exc: + logger.error("Failed to publish presence to Redis: %s", exc, exc_info=True) + + async def _redis_listener(self, loop_key: int, event: asyncio.Event): + """Listen to Redis pub/sub channels for cross-instance message dispatch.""" + while True: + try: + client = aioredis.from_url(settings.REDIS_URL, encoding="utf-8", decode_responses=True) + pubsub = client.pubsub() + await pubsub.psubscribe("world:*:chat", "dm:*:chat", "presence:events") + event.set() + logger.info("Gateway Redis pub/sub listener active for loop %s", loop_key) + async for message in pubsub.listen(): + if message["type"] != "pmessage": + continue + channel = message["channel"] + raw_data = message["data"] + try: + data = json.loads(raw_data) + except Exception: + continue + + if channel.startswith("world:") and channel.endswith(":chat"): + parts = channel.split(":") + world_id = int(parts[1]) + await self.broadcast_world_chat(world_id, data, room_id=data.get("room_id")) + elif channel.startswith("dm:") and channel.endswith(":chat"): + parts = channel.split(":") + conv_id = int(parts[1]) + await self.broadcast_dm_chat(conv_id, data) + elif channel == "presence:events": + logger.info("Redis presence event received: %s. Active users: %s", data, list(self._user_connections.keys())) + for friend_id in data.get("friend_ids", []): + targets = list(self._user_connections.get(friend_id, set())) + logger.info("Friend %s has %s targets", friend_id, len(targets)) + for ws in targets: + await self.send_to_ws(ws, data) + except asyncio.CancelledError: + break + except Exception as exc: + logger.debug("Redis listener encountered error, retrying in 2s: %s", exc) + event.clear() + await asyncio.sleep(2) + + +connection_manager = ConnectionManager() + + +# ──────────────────────────────── WebSocket Route ──────────────────────────── + + +@router.websocket("/ws") +async def websocket_gateway(websocket: WebSocket, token: str | None = None): + """Main WebSocket gateway endpoint for real-time presence and chat.""" + user: User | None = None + # 1. Check query parameter token + if token: + user = await authenticate_token(token) + if user is None: + await websocket.close(code=status.WS_1008_POLICY_VIOLATION) + return + + await websocket.accept() + + # 2. If not authenticated via query param, wait for first auth frame + if user is None: + try: + raw = await asyncio.wait_for(websocket.receive_text(), timeout=3.0) + data = json.loads(raw) + if data.get("type") == "auth" and data.get("token"): + user = await authenticate_token(data["token"]) + except (asyncio.TimeoutError, json.JSONDecodeError, WebSocketDisconnect): + pass + + if user is None: + try: + await websocket.send_text(json.dumps({"type": "error", "message": "Authentication failed"})) + await websocket.close(code=status.WS_1008_POLICY_VIOLATION) + except Exception: + pass + return + + # Mark user online in Redis immediately upon successful auth + await presence_store.heartbeat(user.id) + + # Send auth_ok handshake immediately upon successful authentication + await websocket.send_text( + json.dumps({ + "type": "auth_ok", + "user_id": user.id, + "username": user.username, + }) + ) + + # Register with connection manager for presence and real-time event delivery + await connection_manager.connect(websocket, user) + + try: + while True: + text = await websocket.receive_text() + try: + msg = json.loads(text) + except json.JSONDecodeError: + await websocket.send_text(json.dumps({"type": "error", "message": "Invalid JSON"})) + continue + + msg_type = msg.get("type") + if msg_type == "ping": + await websocket.send_text(json.dumps({"type": "pong"})) + elif msg_type == "heartbeat": + await presence_store.heartbeat(user.id) + await websocket.send_text(json.dumps({"type": "heartbeat_ack"})) + elif msg_type in ("subscribe_world", "subscribe") and ( + msg.get("world_id") or msg.get("channel") == "world" + ): + world_id = msg.get("world_id") or msg.get("id") + if isinstance(world_id, int): + ok = await connection_manager.subscribe_world(websocket, world_id) + if ok: + await websocket.send_text( + json.dumps({"type": "subscribed", "channel": "world", "id": world_id}) + ) + else: + await websocket.send_text( + json.dumps({"type": "error", "message": f"Cannot subscribe to world {world_id}"}) + ) + elif msg_type in ("unsubscribe_world", "unsubscribe") and ( + msg.get("world_id") or msg.get("channel") == "world" + ): + world_id = msg.get("world_id") or msg.get("id") + if isinstance(world_id, int): + connection_manager.unsubscribe_world(websocket, world_id) + await websocket.send_text( + json.dumps({"type": "unsubscribed", "channel": "world", "id": world_id}) + ) + elif msg_type in ("subscribe_dm", "subscribe") and ( + msg.get("conversation_id") or msg.get("channel") == "dm" + ): + conv_id = msg.get("conversation_id") or msg.get("id") + if isinstance(conv_id, int): + ok = await connection_manager.subscribe_dm(websocket, conv_id) + if ok: + await websocket.send_text( + json.dumps({"type": "subscribed", "channel": "dm", "id": conv_id}) + ) + else: + await websocket.send_text( + json.dumps({"type": "error", "message": f"Cannot subscribe to conversation {conv_id}"}) + ) + elif msg_type in ("unsubscribe_dm", "unsubscribe") and ( + msg.get("conversation_id") or msg.get("channel") == "dm" + ): + conv_id = msg.get("conversation_id") or msg.get("id") + if isinstance(conv_id, int): + connection_manager.unsubscribe_dm(websocket, conv_id) + await websocket.send_text( + json.dumps({"type": "unsubscribed", "channel": "dm", "id": conv_id}) + ) + else: + await websocket.send_text(json.dumps({"type": "error", "message": f"Unknown action: {msg_type}"})) + + except WebSocketDisconnect: + pass + finally: + await connection_manager.disconnect(websocket) diff --git a/backend/app/main.py b/backend/app/main.py index 900b73a..7c46a9f 100644 --- a/backend/app/main.py +++ b/backend/app/main.py @@ -7,6 +7,7 @@ from fastapi.middleware.cors import CORSMiddleware from app.config import settings from app.routers import auth, users, worlds, assets, avatar, world_content, internal, presence, groups, social_friends, profiles, chat_dm, chat_reports, chat_world +from app.gateway import ws from app.services.game_servers import supervisor @@ -53,6 +54,8 @@ app.include_router(chat_reports.chat_router, prefix="/api/chat", tags=["chat-rep app.include_router(chat_reports.moderation_router, prefix="/api/moderation", tags=["moderation"]) app.include_router(chat_world.router, prefix="/api/chat", tags=["chat-world"]) app.include_router(chat_world.router, prefix="/v1/chat", tags=["chat-world"]) +app.include_router(ws.router, prefix="/api", tags=["websocket"]) +app.include_router(ws.router, prefix="/v1", tags=["websocket"]) @app.get("/api/health") diff --git a/backend/app/routers/chat_dm.py b/backend/app/routers/chat_dm.py index 4d83465..c12fb66 100644 --- a/backend/app/routers/chat_dm.py +++ b/backend/app/routers/chat_dm.py @@ -81,6 +81,7 @@ from app.models.chat import Conversation, ConversationParticipant, Message from app.models.social import Block from app.models.user import User from app.rate_limit import rate_limit +from app.services.chat_pubsub import publish_dm_chat_event from app.services.text_filter import FilterAction, default_text_filter router = APIRouter() @@ -264,7 +265,9 @@ async def send_dm( await db.commit() await db.refresh(message) - return _to_message_response(message, recipient_id=user_id) + resp = _to_message_response(message, recipient_id=user_id) + await publish_dm_chat_event(conversation.id, resp.model_dump(mode="json")) + return resp @router.get("/{user_id}/messages", response_model=DMMessageListResponse) diff --git a/backend/app/services/presence.py b/backend/app/services/presence.py index c3c1105..4a1bbe5 100644 --- a/backend/app/services/presence.py +++ b/backend/app/services/presence.py @@ -109,6 +109,11 @@ class PresenceStore: client = await self._ensure_client() await client.set(self._heartbeat_key(user_id), "1", ex=self._heartbeat_ttl) + async def clear_heartbeat(self, user_id: int) -> None: + """Clear heartbeat for *user_id* to mark them offline immediately.""" + client = await self._ensure_client() + await client.delete(self._heartbeat_key(user_id)) + # ------------------------------------------------------------------ # Online checks # ------------------------------------------------------------------ diff --git a/backend/tests/test_ws_gateway.py b/backend/tests/test_ws_gateway.py new file mode 100644 index 0000000..675d250 --- /dev/null +++ b/backend/tests/test_ws_gateway.py @@ -0,0 +1,222 @@ +"""Integration tests for real-time WebSocket gateway: +- Authentication (query param + first auth frame) +- Rejection of unauthenticated clients +- Presence on connect / disconnect +- Subscriptions to world chat channels and DM conversations +- Real-time event delivery from REST endpoints +- Reconnect and ping/pong +""" + +import asyncio +import time +import pytest +from fastapi.testclient import TestClient +from starlette.websockets import WebSocketDisconnect + +from app.main import app +from app.services.presence import presence_store + + +def _register(client: TestClient, username: str) -> dict: + resp = client.post( + "/api/auth/register", + json={ + "username": username, + "email": f"{username}@example.com", + "password": "secret123", + "date_of_birth": "1990-01-01", + }, + ) + assert resp.status_code == 201, resp.text + token = resp.json()["access_token"] + headers = {"Authorization": f"Bearer {token}"} + me = client.get("/api/users/me", headers=headers) + assert me.status_code == 200, me.text + return {"headers": headers, "token": token, "id": me.json()["id"], "username": username} + + +def _make_world(client: TestClient, headers: dict, visibility: str = "public") -> int: + r = client.post("/api/worlds/", headers=headers, json={"name": "WSWorld", "visibility": visibility}) + assert r.status_code == 201, r.text + return r.json()["id"] + + +def _make_friends(client: TestClient, user_a: dict, user_b: dict) -> None: + req = client.post("/api/friends/requests", json={"user_id": user_b["id"]}, headers=user_a["headers"]) + assert req.status_code == 201, req.text + req_id = req.json()["id"] + acc = client.post(f"/api/friends/requests/{req_id}/accept", headers=user_b["headers"]) + assert acc.status_code == 200, acc.text + + +class TestWebSocketAuth: + def test_unauthenticated_connection_rejected(self): + client = TestClient(app) + # Invalid query param token rejected + with pytest.raises(WebSocketDisconnect) as exc: + with client.websocket_connect("/api/ws?token=invalid_token"): + pass + assert exc.value.code == 1008 + + # First frame with invalid token rejected + with pytest.raises(WebSocketDisconnect) as exc2: + with client.websocket_connect("/api/ws") as ws: + ws.send_json({"type": "auth", "token": "invalid_token"}) + ws.receive_json() + ws.receive_json() + assert exc2.value.code == 1008 + + def test_auth_via_query_param(self): + client = TestClient(app) + alice = _register(client, "ws_auth_alice") + + with client.websocket_connect(f"/api/ws?token={alice['token']}") as ws: + msg = ws.receive_json() + assert msg["type"] == "auth_ok" + assert msg["user_id"] == alice["id"] + assert msg["username"] == "ws_auth_alice" + + def test_auth_via_first_message(self): + client = TestClient(app) + bob = _register(client, "ws_auth_bob") + + with client.websocket_connect("/api/ws") as ws: + ws.send_json({"type": "auth", "token": bob["token"]}) + msg = ws.receive_json() + assert msg["type"] == "auth_ok" + assert msg["user_id"] == bob["id"] + assert msg["username"] == "ws_auth_bob" + + +class TestWebSocketPresence: + def test_presence_on_connect_and_disconnect(self): + import threading + client = TestClient(app) + alice = _register(client, "ws_pres_alice") + bob = _register(client, "ws_pres_bob") + _make_friends(client, alice, bob) + + # Initially both offline + assert not asyncio.run(presence_store.is_online(alice["id"])) + + # Bob connects and listens for presence events + bob_events = [] + + def run_alice(): + time.sleep(0.2) + c_alice = TestClient(app) + with c_alice.websocket_connect(f"/api/ws?token={alice['token']}") as alice_ws: + assert alice_ws.receive_json()["type"] == "auth_ok" + time.sleep(0.3) + alice_ws.close() + time.sleep(0.3) + + with client.websocket_connect(f"/api/ws?token={bob['token']}") as bob_ws: + assert bob_ws.receive_json()["type"] == "auth_ok" + + alice_thread = threading.Thread(target=run_alice) + alice_thread.start() + + # Bob receives Alice's online event + online_event = bob_ws.receive_json() + assert online_event["type"] == "presence" + assert online_event["user_id"] == alice["id"] + assert online_event["is_online"] is True + + alice_thread.join() + + # Bob receives Alice's offline event + offline_event = bob_ws.receive_json() + assert offline_event["type"] == "presence" + assert offline_event["user_id"] == alice["id"] + assert offline_event["is_online"] is False + + assert not asyncio.run(presence_store.is_online(alice["id"])) + + +class TestWebSocketChatDelivery: + def test_world_chat_subscription_and_delivery(self): + client = TestClient(app) + alice = _register(client, "ws_chat_alice") + bob = _register(client, "ws_chat_bob") + world_id = _make_world(client, alice["headers"]) + + with client.websocket_connect(f"/api/ws?token={bob['token']}") as bob_ws: + assert bob_ws.receive_json()["type"] == "auth_ok" + + # Bob subscribes to world + bob_ws.send_json({"type": "subscribe_world", "world_id": world_id}) + sub_ack = bob_ws.receive_json() + assert sub_ack["type"] == "subscribed" + assert sub_ack["channel"] == "world" + assert sub_ack["id"] == world_id + + # Alice posts message to world via REST + rest_resp = client.post( + f"/api/chat/world/{world_id}/messages", + headers=alice["headers"], + json={"body": "Hello real-time world!"}, + ) + assert rest_resp.status_code == 201, rest_resp.text + msg_id = rest_resp.json()["id"] + + # Bob receives message in real time over WebSocket + ws_msg = bob_ws.receive_json() + assert ws_msg["type"] == "world_chat" + assert ws_msg["world_id"] == world_id + assert ws_msg["message"]["id"] == msg_id + assert ws_msg["message"]["body"] == "Hello real-time world!" + assert ws_msg["message"]["author"]["username"] == "ws_chat_alice" + + def test_dm_chat_subscription_and_delivery(self): + client = TestClient(app) + alice = _register(client, "ws_dm_alice") + bob = _register(client, "ws_dm_bob") + + # Initial DM to create conversation + init_dm = client.post( + f"/api/chat/dm/{bob['id']}/messages", + headers=alice["headers"], + json={"body": "First DM to initiate"}, + ) + assert init_dm.status_code == 201, init_dm.text + conv_id = init_dm.json()["conversation_id"] + + with client.websocket_connect(f"/api/ws?token={bob['token']}") as bob_ws: + assert bob_ws.receive_json()["type"] == "auth_ok" + + # Bob subscribes to this conversation + bob_ws.send_json({"type": "subscribe_dm", "conversation_id": conv_id}) + sub_ack = bob_ws.receive_json() + assert sub_ack["type"] == "subscribed" + assert sub_ack["channel"] == "dm" + assert sub_ack["id"] == conv_id + + # Alice sends another DM via REST + rest_dm = client.post( + f"/api/chat/dm/{bob['id']}/messages", + headers=alice["headers"], + json={"body": "Real-time DM ping!"}, + ) + assert rest_dm.status_code == 201, rest_dm.text + + # Bob receives real-time DM over WebSocket + ws_dm = bob_ws.receive_json() + assert ws_dm["type"] == "dm_chat" + assert ws_dm["conversation_id"] == conv_id + assert ws_dm["message"]["body"] == "Real-time DM ping!" + + +class TestWebSocketPingAndHeartbeat: + def test_ping_pong_and_heartbeat(self): + client = TestClient(app) + alice = _register(client, "ws_ping_alice") + + with client.websocket_connect(f"/api/ws?token={alice['token']}") as ws: + assert ws.receive_json()["type"] == "auth_ok" + + ws.send_json({"type": "ping"}) + assert ws.receive_json()["type"] == "pong" + + ws.send_json({"type": "heartbeat"}) + assert ws.receive_json()["type"] == "heartbeat_ack"