Files
webmail-client/backend/app/ws.py
T
lukasandClaude Opus 4.8 19335116da sync-v2 step 5: activity tiers (hot/warm/cold)
Each owned account is tiered by the coordinator; the runner shapes its work to
match — correctness identical, only latency/socket cost differ.

- hot: persistent INBOX IDLE + secondary poll every POLL_WARM.
- warm: no persistent socket; transient INBOX+special poll every POLL_WARM.
- cold: same transient poll every POLL_COLD.
- engine: poll_all_loop (transient full poll); secondary_poll_loop interval is
  now a parameter; POLL_INTERVAL removed in favor of POLL_WARM/POLL_COLD.
- coordinator: tier eval every 120s from session heat (sess:active:{user}) +
  recent mail (max internal_date within TIER_HOT_ACTIVITY_H), with the
  HOT_BUDGET_PER_WORKER guard shedding least-recently-active session-less hot
  accounts; user.active event promotes a user's accounts to hot at once (≤5s).
- API: login + WS connect/ping publish user.active and refresh sess:active
  (TTL TIER_SESSION_GRACE_MIN). New sync:tier:{id} key, default warm on claim.
- config: POLL_WARM/POLL_COLD/TIER_*/HOT_BUDGET_PER_WORKER.
- tests: 5 tier tests (session/recent-mail promotion, cold, user.active,
  budget demotion) alongside the lease/lifecycle acceptance tests.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
2026-07-19 07:03:54 +02:00

80 lines
2.7 KiB
Python

"""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)
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]