"""Bridge between the game server and the AI service. The AI service (``wizard_ai``) connects to ``/internal/ai`` via websocket. Whenever an AI seat has to act, the server sends it the seat's observation and waits for an action: server -> ai: {"type": "decide", "request_id": "...", "timeout": 5.0, "observation": {...}} ai -> server: {"type": "action", "request_id": "...", "action": {...}, "agent": "heuristic"} ai -> server: {"type": "error", "request_id": "...", "message": "..."} ``timeout`` tells the AI service how long the server waits, so it can drop requests that are already too old. Several AI service instances may connect; requests are distributed round robin. A random legal action is played instead (so a game never gets stuck) when no instance is connected, it does not answer in time, it answers with an error or an illegal action, too many requests are already waiting for it (``max_inflight``), or the circuit breaker is open: after ``breaker_failures`` failures in a row the AI is not asked for ``breaker_cooldown`` seconds, so a struggling AI service is not buried under even more requests. """ from __future__ import annotations import asyncio import itertools import logging import random import uuid from collections.abc import Callable from typing import Any from starlette.websockets import WebSocket, WebSocketDisconnect from wizard_game import Action, Observation, action_from_dict logger = logging.getLogger(__name__) class AIConnection: def __init__(self, websocket: WebSocket, name: str) -> None: self.websocket = websocket self.name = name self.pending: dict[str, asyncio.Future[dict[str, Any]]] = {} self._send_lock = asyncio.Lock() async def send(self, message: dict[str, Any]) -> None: async with self._send_lock: await self.websocket.send_json(message) def fail_pending(self) -> None: for future in self.pending.values(): if not future.done(): future.set_exception(ConnectionError("AI service disconnected")) self.pending.clear() class AIBridge: def __init__( self, timeout: float, rng: random.Random | None = None, on_change: Callable[[], None] | None = None, max_inflight: int = 32, breaker_failures: int = 5, breaker_cooldown: float = 30.0, ) -> None: self.timeout = timeout # Called whenever an AI service connects or disconnects. self._on_change = on_change self._connections: list[AIConnection] = [] self._round_robin = itertools.count() self._rng = rng or random.Random() self._max_inflight = max_inflight self._breaker_failures = breaker_failures self._breaker_cooldown = breaker_cooldown self._consecutive_failures = 0 self._breaker_open_until = 0.0 @property def available(self) -> bool: return bool(self._connections) async def decide(self, observation: Observation) -> Action: """Ask the AI service for an action; falls back to a random legal one.""" legal = observation.legal_actions if not legal: raise ValueError("Observation has no legal actions") connection = self._pick_connection() if connection is None: return self._fallback(observation, "no AI service connected", logging.DEBUG) loop = asyncio.get_running_loop() if loop.time() < self._breaker_open_until: return self._fallback(observation, "AI paused after repeated failures", logging.DEBUG) if len(connection.pending) >= self._max_inflight: return self._failure(observation, f"{len(connection.pending)} AI requests already waiting") request_id = uuid.uuid4().hex future: asyncio.Future[dict[str, Any]] = loop.create_future() connection.pending[request_id] = future try: await connection.send({ "type": "decide", "request_id": request_id, "timeout": self.timeout, "observation": observation.to_dict(), }) reply = await asyncio.wait_for(future, self.timeout) except (TimeoutError, ConnectionError, RuntimeError, WebSocketDisconnect) as exc: return self._failure(observation, f"AI request failed: {exc!r}") finally: connection.pending.pop(request_id, None) if reply.get("type") != "action": return self._failure(observation, f"AI error: {reply.get('message')}") try: action = action_from_dict(reply.get("action", {})) except ValueError as exc: return self._failure(observation, f"malformed AI action: {exc}") if action not in legal: return self._failure(observation, f"illegal AI action {action}") self._consecutive_failures = 0 if reply.get("agent"): logger.debug("AI decision by %s agent", reply["agent"]) return action def stats(self) -> dict[str, Any]: loop_time = asyncio.get_running_loop().time() return { "connections": len(self._connections), "pending": sum(len(c.pending) for c in self._connections), "paused": loop_time < self._breaker_open_until, } async def serve(self, websocket: WebSocket) -> None: """Handle one AI service connection until it disconnects.""" connection = AIConnection(websocket, name=websocket.client.host if websocket.client else "?") self._connections.append(connection) logger.info("AI service connected from %s (%d connected)", connection.name, len(self._connections)) self._notify_change() try: while True: message = await websocket.receive_json() if not isinstance(message, dict): continue if message.get("type") == "hello": logger.info("AI service %s: %s", connection.name, message.get("agent")) continue future = connection.pending.get(str(message.get("request_id"))) if future and not future.done(): future.set_result(message) except (WebSocketDisconnect, RuntimeError, ValueError): pass finally: self._connections.remove(connection) connection.fail_pending() logger.info("AI service %s disconnected (%d connected)", connection.name, len(self._connections)) self._notify_change() def _notify_change(self) -> None: if self._on_change is not None: self._on_change() def _pick_connection(self) -> AIConnection | None: if not self._connections: return None return self._connections[next(self._round_robin) % len(self._connections)] def _failure(self, observation: Observation, reason: str) -> Action: """A failed AI request: count it for the circuit breaker, then fall back.""" self._consecutive_failures += 1 if self._consecutive_failures >= self._breaker_failures: self._consecutive_failures = 0 self._breaker_open_until = asyncio.get_running_loop().time() + self._breaker_cooldown logger.warning("AI failed %d times in a row, pausing AI requests for %.0fs", self._breaker_failures, self._breaker_cooldown) return self._fallback(observation, reason) def _fallback(self, observation: Observation, reason: str, level: int = logging.WARNING) -> Action: logger.log(level, "Playing random move for player %d: %s", observation.player, reason) return self._rng.choice(observation.legal_actions)