From 3906b93c6a5d6ffcb549cc97f9d788e30e687352 Mon Sep 17 00:00:00 2001
From: KoalaTrapLord <265133236+KoalaTrapLord@users.noreply.github.com>
Date: Wed, 25 Mar 2026 19:51:15 -0400
Subject: [PATCH] feat: add WebSocket support for Unity game state streaming
 (#107)

* feat: add ws/ package for WebSocket game state streaming

* feat: add WebSocket ConnectionManager with per-game broadcasting

* feat: add GameStateTracker with deep-diff delta engine

* feat: add WebSocket route with delta push and REST bridge endpoints

* feat: register ws_game_router in lab_api.py
---
 lab_api.py               |   3 +-
 routes/ws_game.py        | 182 +++++++++++++++++++++++++++++++++
 ws/__init__.py           |   1 +
 ws/connection_manager.py | 165 ++++++++++++++++++++++++++++++
 ws/game_state.py         | 211 +++++++++++++++++++++++++++++++++++++++
 5 files changed, 561 insertions(+), 1 deletion(-)
 create mode 100644 routes/ws_game.py
 create mode 100644 ws/__init__.py
 create mode 100644 ws/connection_manager.py
 create mode 100644 ws/game_state.py

diff --git a/lab_api.py b/lab_api.py
index 0db94cc..f75c722 100644
--- a/lab_api.py
+++ b/lab_api.py
@@ -36,11 +36,11 @@ from routes.deepseek import router as deepseek_router
 from routes.deckgen import router as deckgen_router
 from routes.coach import router as coach_router, init_coach_service
 from routes.ml import router as ml_router
+from routes.ws_game import router as ws_game_router
 
 log = logging.getLogger("commander_ai_lab.api")
 
 # ── CORS origins (env-configurable) ────────────────────────────
-_DEFAULT_ORIGINS = "http://localhost:5173,http://localhost:8080"
 _allowed = os.environ.get("ALLOWED_ORIGINS", _DEFAULT_ORIGINS)
 allowed_origins = [o.strip() for o in _allowed.split(",") if o.strip()]
 
@@ -63,6 +63,7 @@ app.include_router(scanner_router)
 app.include_router(deepseek_router)
 app.include_router(deckgen_router)
 app.include_router(coach_router)
+app.include_router(ws_game_router)
 app.include_router(ml_router)
 
 
diff --git a/routes/ws_game.py b/routes/ws_game.py
new file mode 100644
index 0000000..db38ab2
--- /dev/null
+++ b/routes/ws_game.py
@@ -0,0 +1,182 @@
+"""WebSocket route for real-time game state streaming.
+
+Endpoints:
+  WS   /ws/game/{game_id}         — Stream state deltas to Unity
+  POST /api/ws/game/{game_id}/push — Push a state update (called by Forge IPC)
+  POST /api/ws/game/{game_id}/decision — Push AI decision event
+  POST /api/ws/game/{game_id}/end  — Signal game over
+  GET  /api/ws/stats               — Connection statistics
+
+Unity connects to the WebSocket endpoint and receives:
+  1. An initial full snapshot on first state push
+  2. Subsequent delta-only messages as state changes
+  3. Decision events when the AI acts
+  4. A game_over event when the game ends
+"""
+from __future__ import annotations
+
+import logging
+import uuid
+from typing import Optional
+
+from fastapi import APIRouter, WebSocket, WebSocketDisconnect
+from pydantic import BaseModel, Field
+
+from ws.connection_manager import game_manager
+from ws.game_state import get_tracker, remove_tracker
+
+log = logging.getLogger("routes.ws_game")
+
+router = APIRouter(tags=["websocket"])
+
+
+# ────────────────────────────────────────────────────────────
+# WebSocket endpoint  —  Unity connects here
+# ────────────────────────────────────────────────────────────
+
+@router.websocket("/ws/game/{game_id}")
+async def game_ws(websocket: WebSocket, game_id: str):
+    """WebSocket endpoint for Unity game-state streaming.
+
+    Query params:
+        client_id  —  optional unique identifier for this client
+
+    The connection stays open for the duration of the game.
+    The server pushes messages; the client can send:
+        {"type": "ping"}            —  keep-alive
+        {"type": "request_snapshot"} —  request a full re-sync
+    """
+    client_id = (
+        websocket.query_params.get("client_id")
+        or f"unity-{uuid.uuid4().hex[:8]}"
+    )
+
+    client = await game_manager.connect(websocket, game_id, client_id)
+
+    # If the tracker already has a baseline, send the current state
+    tracker = get_tracker(game_id)
+    if tracker.has_baseline:
+        snapshot_msg = tracker.update(tracker._last_snapshot)
+        await websocket.send_json(snapshot_msg)
+
+    try:
+        while True:
+            data = await websocket.receive_json()
+            msg_type = data.get("type", "")
+
+            if msg_type == "ping":
+                await websocket.send_json({"type": "pong"})
+
+            elif msg_type == "request_snapshot":
+                if tracker.has_baseline:
+                    # Re-send current state as full snapshot
+                    import copy
+                    snap = copy.deepcopy(tracker._last_snapshot)
+                    tracker._seq += 1
+                    import time
+                    await websocket.send_json({
+                        "type": "snapshot",
+                        "seq": tracker._seq,
+                        "game_id": game_id,
+                        "timestamp": time.time(),
+                        "state": snap,
+                    })
+                else:
+                    await websocket.send_json({
+                        "type": "error",
+                        "message": "No game state available yet",
+                    })
+
+    except WebSocketDisconnect:
+        pass
+    except Exception as e:
+        log.warning("WS error game=%s client=%s: %s", game_id, client_id, e)
+    finally:
+        await game_manager.disconnect(game_id, client_id)
+
+
+# ────────────────────────────────────────────────────────────
+# REST endpoints  —  called by Forge IPC / policy server
+# ────────────────────────────────────────────────────────────
+
+class PushStateRequest(BaseModel):
+    """Game state snapshot pushed from Forge."""
+    game_id: str = ""
+    turn: int = 0
+    phase: str = ""
+    active_player: int = 0
+    players: list = Field(default_factory=list)
+    stack: list = Field(default_factory=list)
+    legal_actions: list = Field(default_factory=list)
+
+
+class DecisionEvent(BaseModel):
+    """AI decision event pushed after policy inference."""
+    action: str
+    action_index: int = 0
+    confidence: float = 0.0
+    probabilities: dict = Field(default_factory=dict)
+    inference_ms: float = 0.0
+
+
+class GameOverEvent(BaseModel):
+    """Game-over event."""
+    winner_seat: int = -1
+    reason: str = ""
+    turns_played: int = 0
+
+
+@router.post("/api/ws/game/{game_id}/push")
+async def push_game_state(game_id: str, req: PushStateRequest):
+    """Push a game state update to all connected Unity clients.
+
+    Called by the Forge IPC bridge or policy server after each
+    state change. The delta engine diffs against the previous
+    state and broadcasts only what changed.
+    """
+    tracker = get_tracker(game_id)
+    snapshot = req.dict()
+    delta_msg = tracker.update(snapshot)
+    sent = await game_manager.broadcast_to_game(game_id, delta_msg)
+    return {
+        "status": "ok",
+        "type": delta_msg["type"],
+        "seq": delta_msg["seq"],
+        "clients_notified": sent,
+        "changes": len(delta_msg.get("changes", [])),
+    }
+
+
+@router.post("/api/ws/game/{game_id}/decision")
+async def push_decision(game_id: str, event: DecisionEvent):
+    """Push an AI decision event to connected Unity clients."""
+    tracker = get_tracker(game_id)
+    msg = tracker.make_decision_event(
+        action=event.action,
+        action_index=event.action_index,
+        confidence=event.confidence,
+        probabilities=event.probabilities,
+        inference_ms=event.inference_ms,
+    )
+    sent = await game_manager.broadcast_to_game(game_id, msg)
+    return {"status": "ok", "clients_notified": sent}
+
+
+@router.post("/api/ws/game/{game_id}/end")
+async def push_game_over(game_id: str, event: GameOverEvent):
+    """Signal game over to all connected clients and clean up."""
+    tracker = get_tracker(game_id)
+    msg = tracker.make_game_over_event(
+        winner_seat=event.winner_seat,
+        reason=event.reason,
+        turns_played=event.turns_played,
+    )
+    sent = await game_manager.broadcast_to_game(game_id, msg)
+    remove_tracker(game_id)
+    return {"status": "ok", "clients_notified": sent}
+
+
+@router.get("/api/ws/stats")
+async def ws_stats():
+    """Return WebSocket connection statistics."""
+    return game_manager.stats()
diff --git a/ws/__init__.py b/ws/__init__.py
new file mode 100644
index 0000000..221a22c
--- /dev/null
+++ b/ws/__init__.py
@@ -0,0 +1 @@
+"""WebSocket module for real-time game state streaming to Unity."""
diff --git a/ws/connection_manager.py b/ws/connection_manager.py
new file mode 100644
index 0000000..c2bb5ff
--- /dev/null
+++ b/ws/connection_manager.py
@@ -0,0 +1,165 @@
+"""WebSocket connection manager for game-state streaming.
+
+Manages per-game-session WebSocket connections so the server can
+push state deltas to every connected Unity client in real time.
+"""
+from __future__ import annotations
+
+import asyncio
+import logging
+import time
+from dataclasses import dataclass, field
+from typing import Dict, List, Optional, Set
+
+from fastapi import WebSocket, WebSocketDisconnect
+
+log = logging.getLogger("ws.connection_manager")
+
+
+@dataclass
+class ClientInfo:
+    """Metadata for a single connected WebSocket client."""
+    ws: WebSocket
+    client_id: str
+    game_id: str
+    connected_at: float = field(default_factory=time.time)
+    last_ping: float = field(default_factory=time.time)
+    messages_sent: int = 0
+
+
+class ConnectionManager:
+    """Manages WebSocket connections grouped by game_id.
+
+    Usage:
+        manager = ConnectionManager()
+        await manager.connect(websocket, game_id, client_id)
+        await manager.broadcast_to_game(game_id, payload)
+        manager.disconnect(websocket, game_id)
+    """
+
+    def __init__(self, heartbeat_interval: float = 30.0):
+        # game_id -> set of ClientInfo
+        self._games: Dict[str, Dict[str, ClientInfo]] = {}
+        self._heartbeat_interval = heartbeat_interval
+        self._lock = asyncio.Lock()
+
+    # ── Connect / Disconnect ───────────────────────────────
+
+    async def connect(
+        self, ws: WebSocket, game_id: str, client_id: str
+    ) -> ClientInfo:
+        """Accept and register a WebSocket connection."""
+        await ws.accept()
+        client = ClientInfo(ws=ws, client_id=client_id, game_id=game_id)
+        async with self._lock:
+            if game_id not in self._games:
+                self._games[game_id] = {}
+            self._games[game_id][client_id] = client
+        log.info(
+            "WS connected: client=%s game=%s (total=%d)",
+            client_id, game_id, len(self._games[game_id]),
+        )
+        return client
+
+    async def disconnect(self, game_id: str, client_id: str) -> None:
+        """Remove a client from the game session."""
+        async with self._lock:
+            game = self._games.get(game_id)
+            if game and client_id in game:
+                del game[client_id]
+                if not game:
+                    del self._games[game_id]
+        log.info("WS disconnected: client=%s game=%s", client_id, game_id)
+
+    # ── Broadcasting ───────────────────────────────────────
+
+    async def broadcast_to_game(
+        self, game_id: str, payload: dict
+    ) -> int:
+        """Send a JSON payload to every client in a game session.
+
+        Returns the number of clients that received the message.
+        Silently removes clients whose connections are broken.
+        """
+        async with self._lock:
+            game = self._games.get(game_id)
+            if not game:
+                return 0
+            clients = list(game.values())
+
+        sent = 0
+        dead: List[str] = []
+        for client in clients:
+            try:
+                await client.ws.send_json(payload)
+                client.messages_sent += 1
+                sent += 1
+            except Exception:
+                dead.append(client.client_id)
+
+        # Clean up broken connections
+        if dead:
+            async with self._lock:
+                game = self._games.get(game_id)
+                if game:
+                    for cid in dead:
+                        game.pop(cid, None)
+                    if not game:
+                        del self._games[game_id]
+            log.warning(
+                "Removed %d dead connection(s) from game %s",
+                len(dead), game_id,
+            )
+        return sent
+
+    async def send_to_client(
+        self, game_id: str, client_id: str, payload: dict
+    ) -> bool:
+        """Send a JSON payload to a specific client."""
+        async with self._lock:
+            game = self._games.get(game_id)
+            if not game:
+                return False
+            client = game.get(client_id)
+            if not client:
+                return False
+        try:
+            await client.ws.send_json(payload)
+            client.messages_sent += 1
+            return True
+        except Exception:
+            await self.disconnect(game_id, client_id)
+            return False
+
+    # ── Queries ────────────────────────────────────────────
+
+    def get_game_clients(self, game_id: str) -> List[str]:
+        """Return client IDs connected to a game."""
+        game = self._games.get(game_id, {})
+        return list(game.keys())
+
+    def active_games(self) -> List[str]:
+        """Return all game IDs with active connections."""
+        return list(self._games.keys())
+
+    @property
+    def total_connections(self) -> int:
+        return sum(len(g) for g in self._games.values())
+
+    def stats(self) -> dict:
+        """Return connection statistics."""
+        return {
+            "active_games": len(self._games),
+            "total_connections": self.total_connections,
+            "games": {
+                gid: {
+                    "clients": len(clients),
+                    "client_ids": list(clients.keys()),
+                }
+                for gid, clients in self._games.items()
+            },
+        }
+
+
+# Module-level singleton
+game_manager = ConnectionManager()
diff --git a/ws/game_state.py b/ws/game_state.py
new file mode 100644
index 0000000..56b5e9c
--- /dev/null
+++ b/ws/game_state.py
@@ -0,0 +1,211 @@
+"""Game-state delta engine.
+
+Compares successive game-state snapshots and produces minimal
+JSON deltas so that Unity only receives what changed — not
+the full 4-player board every tick.
+
+Delta format sent over WebSocket:
+{
+  "type": "delta" | "snapshot" | "decision" | "game_over",
+  "seq": <monotonic int>,
+  "game_id": "...",
+  "turn": 5,
+  "phase": "main_1",
+  "timestamp": 1711400000.0,
+  "changes": [
+    {"path": "players.0.life", "op": "replace", "value": 37},
+    {"path": "players.1.battlefield", "op": "add", "value": {"name": "Sol Ring", ...}},
+    {"path": "stack", "op": "replace", "value": [...]},
+  ]
+}
+"""
+from __future__ import annotations
+
+import copy
+import logging
+import time
+from typing import Any, Dict, List, Optional, Tuple
+
+log = logging.getLogger("ws.game_state")
+
+
+def _deep_diff(
+    old: Any, new: Any, path: str = ""
+) -> List[Dict[str, Any]]:
+    """Recursively diff two JSON-serialisable structures.
+
+    Returns a list of RFC-6902-style patch operations:
+      {"op": "replace"|"add"|"remove", "path": "...", "value": ...}
+    """
+    changes: List[Dict[str, Any]] = []
+
+    if old is None and new is not None:
+        return [{"op": "add", "path": path, "value": new}]
+    if old is not None and new is None:
+        return [{"op": "remove", "path": path}]
+    if type(old) != type(new):
+        return [{"op": "replace", "path": path, "value": new}]
+
+    if isinstance(new, dict):
+        all_keys = set(list(old.keys()) + list(new.keys()))
+        for key in sorted(all_keys):
+            child_path = f"{path}.{key}" if path else key
+            if key not in old:
+                changes.append({"op": "add", "path": child_path, "value": new[key]})
+            elif key not in new:
+                changes.append({"op": "remove", "path": child_path})
+            else:
+                changes.extend(_deep_diff(old[key], new[key], child_path))
+        return changes
+
+    if isinstance(new, list):
+        # For lists (battlefield, hand, graveyard, stack) we do a
+        # simple length+element comparison. Zone lists change
+        # frequently and a full LCS diff is overkill for game state.
+        if old == new:
+            return []
+        # If the list is short or changed significantly, send whole list
+        if len(old) != len(new) or _list_diff_ratio(old, new) > 0.5:
+            return [{"op": "replace", "path": path, "value": new}]
+        # Element-wise diff for same-length lists with few changes
+        for i, (o, n) in enumerate(zip(old, new)):
+            child_path = f"{path}.{i}"
+            changes.extend(_deep_diff(o, n, child_path))
+        return changes
+
+    # Scalar comparison
+    if old != new:
+        return [{"op": "replace", "path": path, "value": new}]
+    return []
+
+
+def _list_diff_ratio(old: list, new: list) -> float:
+    """Fraction of elements that differ between two same-length lists."""
+    if not old:
+        return 0.0
+    diffs = sum(1 for a, b in zip(old, new) if a != b)
+    return diffs / len(old)
+
+
+class GameStateTracker:
+    """Tracks per-game state and emits deltas on each update.
+
+    Usage:
+        tracker = GameStateTracker(game_id)
+        delta = tracker.update(new_snapshot)
+        # delta is a dict ready to send via WebSocket
+    """
+
+    def __init__(self, game_id: str):
+        self.game_id = game_id
+        self._last_snapshot: Optional[Dict] = None
+        self._seq: int = 0
+        self._created_at: float = time.time()
+
+    def update(self, snapshot: Dict) -> Dict:
+        """Compare snapshot with previous state, return delta message.
+
+        First call returns a full snapshot message (type="snapshot").
+        Subsequent calls return delta messages (type="delta").
+        """
+        self._seq += 1
+        now = time.time()
+
+        if self._last_snapshot is None:
+            # First update — send full snapshot
+            self._last_snapshot = copy.deepcopy(snapshot)
+            return {
+                "type": "snapshot",
+                "seq": self._seq,
+                "game_id": self.game_id,
+                "turn": snapshot.get("turn", 0),
+                "phase": snapshot.get("phase", ""),
+                "timestamp": now,
+                "state": snapshot,
+            }
+
+        # Compute delta
+        changes = _deep_diff(self._last_snapshot, snapshot)
+        self._last_snapshot = copy.deepcopy(snapshot)
+
+        if not changes:
+            return {
+                "type": "heartbeat",
+                "seq": self._seq,
+                "game_id": self.game_id,
+                "timestamp": now,
+            }
+
+        return {
+            "type": "delta",
+            "seq": self._seq,
+            "game_id": self.game_id,
+            "turn": snapshot.get("turn", 0),
+            "phase": snapshot.get("phase", ""),
+            "timestamp": now,
+            "changes": changes,
+        }
+
+    def make_decision_event(
+        self,
+        action: str,
+        action_index: int,
+        confidence: float,
+        probabilities: Optional[Dict[str, float]] = None,
+        inference_ms: float = 0.0,
+    ) -> Dict:
+        """Build a decision event message for the Unity client."""
+        self._seq += 1
+        return {
+            "type": "decision",
+            "seq": self._seq,
+            "game_id": self.game_id,
+            "timestamp": time.time(),
+            "action": action,
+            "action_index": action_index,
+            "confidence": confidence,
+            "probabilities": probabilities or {},
+            "inference_ms": inference_ms,
+        }
+
+    def make_game_over_event(
+        self,
+        winner_seat: int,
+        reason: str = "",
+        turns_played: int = 0,
+    ) -> Dict:
+        """Build a game-over event message."""
+        self._seq += 1
+        return {
+            "type": "game_over",
+            "seq": self._seq,
+            "game_id": self.game_id,
+            "timestamp": time.time(),
+            "winner_seat": winner_seat,
+            "reason": reason,
+            "turns_played": turns_played,
+        }
+
+    @property
+    def sequence(self) -> int:
+        return self._seq
+
+    @property
+    def has_baseline(self) -> bool:
+        return self._last_snapshot is not None
+
+
+# Per-game tracker registry
+_trackers: Dict[str, GameStateTracker] = {}
+
+
+def get_tracker(game_id: str) -> GameStateTracker:
+    """Get or create a state tracker for a game."""
+    if game_id not in _trackers:
+        _trackers[game_id] = GameStateTracker(game_id)
+    return _trackers[game_id]
+
+
+def remove_tracker(game_id: str) -> None:
+    """Remove a game tracker (called on game end)."""
+    _trackers.pop(game_id, None)