Wizard/wizard_server/app/ai_bridge.py

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