"""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 .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) self.lobbies = LobbyManager(self.bridge, settings) # --- 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() session: Session | None = None try: while True: raw = await websocket.receive_text() 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))) except WebSocketDisconnect: pass finally: 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() 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, }) 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._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) 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) 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()