"""Connects web clients, sessions and lobbies.""" from __future__ import annotations import asyncio import contextlib import logging import time from pydantic import ValidationError from starlette.websockets import WebSocket, WebSocketDisconnect from wizard_game import GameOptions from . import protocol from .ai_bridge import AIBridge from .config import Settings from .limits import ConnectionLimiter, KeyedRateLimiter, TokenBucket from .load import LoadLevel, LoadMonitor from .lobby import Lobby, LobbyError, LobbyManager from .sessions import Connection, Session, SessionRegistry logger = logging.getLogger(__name__) CLEANUP_INTERVAL = 30.0 class WizardService: def __init__(self, settings: Settings) -> None: self.settings = settings self.sessions = SessionRegistry() self.bridge = AIBridge( settings.ai_decision_timeout, on_change=self._on_ai_availability_changed, max_inflight=settings.ai_max_inflight, breaker_failures=settings.ai_breaker_failures, breaker_cooldown=settings.ai_breaker_cooldown, ) self.lobbies = LobbyManager(self.bridge, settings, on_games_changed=self._broadcast_server_status) self.connections = ConnectionLimiter(settings.max_connections, settings.max_connections_per_ip) self.lobby_creation = KeyedRateLimiter(settings.lobby_create_rate / 60, settings.lobby_create_rate) self.load = LoadMonitor(settings, on_change=lambda _level: self._broadcast_server_status()) self._last_server_status: dict | None = None # --- Capacity ----------------------------------------------------------- @property def accepting_games(self) -> bool: """Whether new lobbies may be created and games started. Running games are never affected; only new work is refused. """ return ( self.load.level is not LoadLevel.OVERLOADED and self.lobbies.active_games() < self.settings.max_active_games ) def server_status(self) -> dict: return { "type": "server_status", "load": self.load.level.value, "accepting_games": self.accepting_games, } def _broadcast_server_status(self) -> None: status = self.server_status() if status == self._last_server_status: return self._last_server_status = status for session in self.sessions.connected(): session.send(status) def _require_capacity(self) -> None: if not self.accepting_games: raise LobbyError( "The server is busy right now, new games can't be started. Please try again in a moment.", protocol.ERROR_BUSY, ) # --- Client connections ------------------------------------------------- async def handle_client(self, websocket: WebSocket) -> None: origin = websocket.headers.get("origin") if self.settings.allowed_origins and origin not in self.settings.allowed_origins: await websocket.close(code=1008) return await websocket.accept() connection = Connection(websocket) connection.start() if not self.connections.acquire(connection.ip): logger.warning("Refusing connection from %s: connection limit reached", connection.ip) connection.send(protocol.error("The server is full right now, please try again later.", protocol.ERROR_BUSY)) await connection.close(code=protocol.CLOSE_TRY_AGAIN_LATER) return bucket = TokenBucket(self.settings.message_rate, self.settings.message_burst) violations = 0 session: Session | None = None try: while True: raw = await websocket.receive_text() if not bucket.allow(): violations += 1 if violations >= self.settings.max_rate_violations: logger.warning("Closing connection from %s: message flood", connection.ip) break if violations == 1 or violations % 10 == 0: connection.send(protocol.error("Too many messages, slow down.", protocol.ERROR_RATE_LIMITED)) continue try: message = protocol.client_message_adapter.validate_json(raw) except ValidationError as exc: connection.send(protocol.error(_validation_message(exc))) continue if isinstance(message, protocol.Hello): session = self._hello(connection, message.token) continue if session is None or session.connection is not connection: connection.send(protocol.error("Send 'hello' first")) continue try: self._dispatch(session, message) except LobbyError as exc: session.send(protocol.error(str(exc), exc.code)) except WebSocketDisconnect: pass finally: self.connections.release(connection.ip) if session is not None and session.connection is connection: session.connection = None session.last_seen = time.monotonic() if lobby := self.lobbies.get(session.lobby_code): lobby.broadcast() await connection.close(code=1008 if violations >= self.settings.max_rate_violations else 1000) def _hello(self, connection: Connection, token: str | None) -> Session: session = self.sessions.get(token) if session is None: session = self.sessions.create() elif session.connection is not None and session.connection is not connection: old = session.connection old.send(protocol.error("This session was opened in another window")) asyncio.create_task(old.close(code=4000)) session.connection = connection session.last_seen = time.monotonic() session.send({ "type": "welcome", "player_id": session.player_id, "token": session.token, "name": session.name, }) session.send(self.server_status()) lobby = self.lobbies.get(session.lobby_code) if lobby is not None and lobby.seat_index(session.player_id) is not None: lobby.add_human(session) # re-binds the seat to this session lobby.broadcast() else: session.lobby_code = None self._send_no_lobby(session) return session def _dispatch(self, session: Session, message: object) -> None: match message: case protocol.CreateLobby(name=name): self._require_capacity() ip = session.connection.ip if session.connection else "unknown" if not self.lobby_creation.allow(ip): raise LobbyError("You are creating lobbies too quickly, please wait a moment.", protocol.ERROR_RATE_LIMITED) self._leave_current_lobby(session) session.name = name lobby = self.lobbies.create() self._join(session, lobby) case protocol.JoinLobby(code=code, name=name): lobby = self.lobbies.get(code) if lobby is None: raise LobbyError(f"Lobby {code} does not exist") if session.lobby_code != lobby.code: self._leave_current_lobby(session) session.name = name self._join(session, lobby) case protocol.LeaveLobby(): self._leave_current_lobby(session) self._send_no_lobby(session) case protocol.AddAI(): lobby = self._lobby_of(session) lobby.add_ai(session) lobby.broadcast() case protocol.RemovePlayer(player_id=player_id): lobby = self._lobby_of(session) seat = lobby.remove_player(session, player_id) if seat.session is not None: seat.session.lobby_code = None seat.session.send(protocol.error("You were removed from the lobby")) self._send_no_lobby(seat.session) lobby.broadcast() case protocol.SetOptions(hidden_bids=hidden_bids): lobby = self._lobby_of(session) lobby.set_options(session, GameOptions(hidden_bids=hidden_bids)) lobby.broadcast() case protocol.StartGame(): lobby = self._lobby_of(session) self._require_capacity() lobby.start(session) lobby.broadcast() case protocol.ReturnToLobby(): lobby = self._lobby_of(session) lobby.return_to_lobby(session) lobby.broadcast() case protocol.GameAction(action=action): self._lobby_of(session).apply_action(session, action) case _: raise LobbyError("Unsupported message") def _join(self, session: Session, lobby: Lobby) -> None: try: lobby.add_human(session) except LobbyError: self.lobbies.discard_if_abandoned(lobby) raise session.lobby_code = lobby.code lobby.broadcast() def _leave_current_lobby(self, session: Session) -> None: lobby = self.lobbies.get(session.lobby_code) session.lobby_code = None if lobby is None: return lobby.leave(session) if not self.lobbies.discard_if_abandoned(lobby): lobby.broadcast() def _lobby_of(self, session: Session) -> Lobby: lobby = self.lobbies.get(session.lobby_code) if lobby is None: raise LobbyError("You are not in a lobby") return lobby @staticmethod def _send_no_lobby(session: Session) -> None: session.send({"type": "state", "lobby": None, "game": None, "events": []}) # --- AI service --------------------------------------------------------- async def handle_ai(self, websocket: WebSocket) -> None: await websocket.accept() try: await self.bridge.serve(websocket) finally: with contextlib.suppress(Exception): await websocket.close() def _on_ai_availability_changed(self) -> None: # Lobbies show whether the AI service is available. for lobby in self.lobbies: lobby.broadcast() # --- Housekeeping ------------------------------------------------------- async def cleanup_loop(self) -> None: while True: await asyncio.sleep(CLEANUP_INTERVAL) try: self.cleanup() except Exception: logger.exception("Cleanup failed") def cleanup(self) -> None: for session in self.sessions.expired(self.settings.session_ttl): self._leave_current_lobby(session) self.sessions.remove(session) self.lobby_creation.prune() def health(self) -> dict: return { "status": "ok", # stays ok under load: an overloaded server must not be restarted **self.load.stats(), "accepting_games": self.accepting_games, "lobbies": len(self.lobbies), "active_games": self.lobbies.active_games(), "sessions": len(self.sessions), "connections": self.connections.total, "ai_available": self.bridge.available, "ai": self.bridge.stats(), } def _validation_message(exc: ValidationError) -> str: errors = exc.errors() if not errors: return "Invalid message" first = errors[0] location = ".".join(str(part) for part in first.get("loc", ()) if part) return f"Invalid message: {location} {first.get('msg', '')}".strip()