"""
Module de diffusion et multiplexage temps reel des evenements WebSocket.
Fournit une passerelle thread-safe pour envoyer des messages asynchrones depuis des threads de calcul MLX.
"""
from __future__ import annotations
import asyncio
from typing import Any
from ..core.logger import getMirzaLogger
logger = getMirzaLogger("mirza.server.broadcaster")
try:
import orjson
def _fast_dumps(data: dict) -> str:
"""
Serialise un dictionnaire en chaine JSON rapide via orjson.
Parameters
----------
data : dict
Donnees a serialiser.
Returns
-------
str
Chaine JSON formatee.
"""
return orjson.dumps(data).decode("utf-8")
except ImportError:
import json
def _fast_dumps(data: dict) -> str:
"""
Serialise un dictionnaire en chaine JSON standard via le module json.
Parameters
----------
data : dict
Donnees a serialiser.
Returns
-------
str
Chaine JSON formatee.
"""
return json.dumps(data)
[docs]
class EventBroadcaster:
"""
Diffuseur d'evenements temps reel thread-safe pour les clients connectes en WebSocket.
Maintient le registre en memoire de toutes les connexions actives et offre des primitives
permettant aux threads d'inference (LLM, STT, TTS) d'emettre des evenements JSON
vers la boucle d'evenements asynchrone FastAPI sans risque de blocage.
Attributes
----------
activeConnections : set[Any]
Ensemble des instances WebSocket actives connectees au serveur.
eventLoop : asyncio.AbstractEventLoop or None
Reference vers la boucle d'evenements asyncio principale.
"""
[docs]
def __init__(self) -> None:
"""
Initialise le registre vide des connexions WebSocket.
"""
self.activeConnections: set[Any] = set()
self.eventLoop: asyncio.AbstractEventLoop | None = None
[docs]
def setEventLoop(self, loop: asyncio.AbstractEventLoop) -> None:
"""
Associe la boucle d'evenements asynchrone principale de FastAPI au diffuseur.
Parameters
----------
loop : asyncio.AbstractEventLoop
Boucle d'evenements asyncio en cours d'execution.
"""
self.eventLoop = loop
[docs]
def registerConnection(self, websocket: Any) -> None:
"""
Enregistre une nouvelle connexion client WebSocket active dans le registre.
Parameters
----------
websocket : Any
Instance de la socket cliente connectee (FastAPI / Starlette WebSocket).
"""
self.activeConnections.add(websocket)
logger.debug("Connexion WebSocket enregistree (total: %d)", len(self.activeConnections))
[docs]
def unregisterConnection(self, websocket: Any) -> None:
"""
Supprime une connexion client deconnectee du registre.
Parameters
----------
websocket : Any
Instance de la socket cliente a retirer.
"""
self.activeConnections.discard(websocket)
logger.debug("Connexion WebSocket retiree (total: %d)", len(self.activeConnections))
[docs]
def broadcastEvent(self, eventName: str, payload: dict) -> None:
"""
Envoie un message JSON formate a tous les clients widgets et terminaux connectes.
Cette methode est thread-safe : si elle est appelee depuis un thread de travail
(worker thread TTS ou thread d'inference MLX), elle planifie la coroutine d'envoi
sur la boucle principale via `asyncio.run_coroutine_threadsafe`.
Parameters
----------
eventName : str
Nom canonique de l'evenement (ex: 'stateChange', 'thoughtChunk', 'telemetry').
payload : dict
Donnees associees a l'evenement sous forme de dictionnaire serialisable.
"""
if not self.activeConnections or not self.eventLoop:
return
message = _fast_dumps({"event": eventName, "payload": payload})
if self.eventLoop.is_running():
asyncio.run_coroutine_threadsafe(self._sendToAll(message), self.eventLoop)
async def _sendToAll(self, message: str) -> None:
"""
Coroutine interne distribuant le message serialise a chaque connexion du registre.
Parameters
----------
message : str
Chaine JSON complete a envoyer.
"""
disconnected = set()
for ws in list(self.activeConnections):
try:
await ws.send_text(message)
except Exception as exc:
logger.debug("Erreur d'envoi WebSocket, client marque deconnecte: %s", exc)
disconnected.add(ws)
for ws in disconnected:
self.activeConnections.discard(ws)