webmail-client/backend/app/ws.py

"""WebSocket hub (SPEC §9.2).

Auth via session cookie. Relays every §8.3 event belonging to the user as
``{event, payload}`` frames; the client refetches via TanStack Query invalidation.
Client -> server: ``{"type":"ping"}`` only.
"""

from __future__ import annotations

import asyncio
import json
import logging

from fastapi import APIRouter, WebSocket, WebSocketDisconnect

from app.config import get_settings
from app.events import get_event_bus
from app.identity.cookies import SESSION_COOKIE_INSECURE, SESSION_COOKIE_SECURE
from app.identity.sessions import resolve_session
from app.redis_client import get_redis
from app.sync import leases

logger = logging.getLogger(__name__)

router = APIRouter()


@router.websocket("/ws")
async def ws_endpoint(websocket: WebSocket) -> None:
    token = websocket.cookies.get(SESSION_COOKIE_SECURE) or websocket.cookies.get(
        SESSION_COOKIE_INSECURE, ""
    )
    session = await resolve_session(token)
    if session is None:
        await websocket.close(code=1008)  # policy violation
        return

    await websocket.accept()
    user_id = str(session.user_id)
    grace_s = get_settings().tier_session_grace_min * 60
    # Sync v2 §3/§5: an open WS means the user is active -> session heat + promote.
    await leases.mark_user_active(get_redis(), user_id, grace_s)
    await get_event_bus().publish("user.active", {"user_id": user_id})
    pubsub = get_redis().pubsub()
    await pubsub.psubscribe("evt:*")

    async def forward() -> None:
        async for message in pubsub.listen():
            if message.get("type") != "pmessage":
                continue
            try:
                data = json.loads(message["data"])
            except (ValueError, TypeError):
                continue
            payload = data.get("payload", {})
            if str(payload.get("user_id")) == user_id:
                await websocket.send_json(data)
                if data.get("event") == "user.signed_out":
                    # The session this socket was opened with no longer exists.
                    await websocket.close(code=1008)
                    return

    async def receive() -> None:
        while True:
            raw = await websocket.receive_text()
            try:
                if json.loads(raw).get("type") == "ping":
                    # A ping keeps the user's session heat alive (§5).
                    await leases.mark_user_active(get_redis(), user_id, grace_s)
                    await websocket.send_json({"type": "pong"})
            except (ValueError, TypeError):
                continue

    forward_task = asyncio.create_task(forward())
    receive_task = asyncio.create_task(receive())
    try:
        await asyncio.wait({forward_task, receive_task}, return_when=asyncio.FIRST_COMPLETED)
    except WebSocketDisconnect:
        pass
    finally:
        for task in (forward_task, receive_task):
            task.cancel()
        await pubsub.aclose()  # type: ignore[no-untyped-call]