Implement real-time WebSocket gateway for presence, DM, and world chat (US-4, US-2)

2026-09-15 18:58:19 +00:00
parent 4fc297be43
commit e5bcd61ee4
6 changed files with 773 additions and 1 deletions
@@ -0,0 +1 @@
+"""WebSocket real-time gateway package."""
@@ -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)
@@ -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")
@@ -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)
@@ -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
    # ------------------------------------------------------------------
@@ -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"