jmri_server.push

  1"""WebSocket push fan-out.
  2
  3When a bean changes state, every connected WebSocket client should receive an
  4unsolicited message with the bean's current data. :class:`PushManager` wires
  5this up by:
  6
  71. Tracking the set of active WebSocket connections.
  82. Installing a listener on every bean's ``state_changed`` Signal — for both
  9   beans that already exist and beans created later (via the manager's
 10   ``added`` signal).
 113. On any change, formatting the bean through the appropriate service and
 12   broadcasting the JMRI envelope to every active client.
 13"""
 14from __future__ import annotations
 15
 16import asyncio
 17import json
 18from collections.abc import Awaitable, Callable
 19from typing import TYPE_CHECKING, Any
 20
 21from jmri_core.bean import AbstractNamedBean
 22
 23if TYPE_CHECKING:
 24    from fastapi import WebSocket
 25
 26    from jmri_server.context import ServerContext
 27    from jmri_server.registry import ServiceRegistry
 28    from jmri_server.services.beans import _BeanService
 29
 30
 31class PushManager:
 32    """Owns the set of WebSocket connections and the listener hookups."""
 33
 34    def __init__(self, ctx: ServerContext, registry: ServiceRegistry) -> None:
 35        self._ctx = ctx
 36        self._registry = registry
 37        self._clients: set[WebSocket] = set()
 38        # disconnect callbacks for everything we've subscribed to:
 39        self._unsubscribes: list[Callable[[], None]] = []
 40        self._wired = False
 41
 42    # ---- WebSocket lifecycle ----------------------------------------------
 43
 44    async def attach(self, ws: WebSocket) -> None:
 45        """Register a newly accepted WebSocket. Wires up listeners on first attach."""
 46        if not self._wired:
 47            self._wire_managers()
 48            self._wired = True
 49        self._clients.add(ws)
 50
 51    async def detach(self, ws: WebSocket) -> None:
 52        self._clients.discard(ws)
 53
 54    async def broadcast(self, message: dict[str, Any]) -> None:
 55        """Send *message* to every active client. Dead sockets are dropped."""
 56        if not self._clients:
 57            return
 58        payload = json.dumps(message)
 59        dead: list[WebSocket] = []
 60        for ws in list(self._clients):
 61            try:
 62                await ws.send_text(payload)
 63            except Exception:  # noqa: BLE001 — socket dead
 64                dead.append(ws)
 65        for ws in dead:
 66            self._clients.discard(ws)
 67
 68    # ---- listener installation --------------------------------------------
 69
 70    def _wire_managers(self) -> None:
 71        """Subscribe to added/removed/state_changed for every bean manager."""
 72        bean_topics = (
 73            "turnout", "sensor", "light", "signalHead", "signalMast",
 74            "memory", "reporter", "block", "route",
 75        )
 76        for topic in bean_topics:
 77            svc = self._registry.get(topic)
 78            if svc is None:
 79                continue
 80            self._wire_bean_service(topic, svc)  # type: ignore[arg-type]
 81
 82        # Power is a singleton, not a Manager.
 83        power = self._registry.get("power")
 84        if power is not None:
 85            self._unsubscribes.append(
 86                self._ctx.connection.power.state_changed.connect(
 87                    self._make_power_listener(power)
 88                )
 89            )
 90
 91        # Command-station power. Drivers with heartbeat detection emit
 92        # ``command_station_power_changed`` when reachability flips;
 93        # broadcast so the /home CS badge updates live without polling.
 94        conn = self._ctx.connection
 95        cs_signal = getattr(conn, "command_station_power_changed", None)
 96        if cs_signal is not None:
 97            async def on_cs_change(on: bool) -> None:
 98                await self.broadcast({
 99                    "type": "commandStationPower",
100                    "data": {"on": bool(on)},
101                })
102            self._unsubscribes.append(cs_signal.connect(on_cs_change))
103
104        # Consist manager exposes ``created`` / ``updated`` / ``deleted``
105        # signals rather than the added/removed/state_changed trio, and
106        # its beans aren't NamedBeans — so wire it here explicitly.
107        # Broadcasts trigger the /consists page's htmx swap and any
108        # scripting client subscribed to the ``consist`` topic.
109        consist_svc = self._registry.get("consist")
110        if consist_svc is not None:
111            self._wire_consist_manager()
112
113        # Throttle manager exposes ``acquired`` / ``released`` and each
114        # ``Throttle`` has its own ``state_changed`` — Throttles are
115        # NOT NamedBeans so they don't fit ``_wire_bean_service``. Wire
116        # them explicitly so that ANY mutation (from the /throttle page,
117        # from the consist fan-out, from the WiThrottle bridge, or
118        # from a script) broadcasts to every open client. Without this,
119        # a consist Forward-button click updates the underlying
120        # throttle beans but the /throttle page for a member loco
121        # never sees the change.
122        throttle_svc = self._registry.get("throttle")
123        if throttle_svc is not None:
124            self._wire_throttle_manager()
125
126    def _wire_throttle_manager(self) -> None:
127        """Broadcast ``{type: "throttle", ...}`` on every throttle mutation.
128
129        Two-tier subscription:
130        * ``throttles.acquired`` fires when a new handle appears — we
131          install a per-throttle listener AND broadcast the initial
132          state so anyone watching sees the fresh throttle immediately.
133        * ``throttles.released`` clears the handle — broadcast a
134          release marker so watchers can drop their local mirror.
135        * Each ``Throttle.state_changed`` fans out any speed / direction
136          / function change through :func:`_project_throttle`.
137        """
138        from jmri_core.beans.throttle import Throttle  # noqa: PLC0415
139        from jmri_server.services.throttle import _project as _project_throttle  # noqa: PLC0415
140
141        throttles = self._ctx.connection.throttles
142
143        def _install_listener(throttle: Throttle) -> None:
144            async def on_change(_payload: Any) -> None:
145                await self.broadcast({
146                    "type": "throttle",
147                    "data": _project_throttle(throttle),
148                })
149            self._unsubscribes.append(throttle.state_changed.connect(on_change))
150
151        # Cover throttles already acquired before the first client
152        # connected (rare — normally acquire is client-driven — but
153        # possible if a script or the consist fan-out ran during startup).
154        for throttle in throttles:
155            _install_listener(throttle)
156
157        async def on_acquired(throttle: Throttle) -> None:
158            _install_listener(throttle)
159            await self.broadcast({
160                "type": "throttle",
161                "data": _project_throttle(throttle),
162            })
163
164        async def on_released(throttle: Throttle) -> None:
165            await self.broadcast({
166                "type": "throttle",
167                "data": _project_throttle(throttle, released=True),
168            })
169
170        self._unsubscribes.append(throttles.acquired.connect(on_acquired))
171        self._unsubscribes.append(throttles.released.connect(on_released))
172
173    def _wire_consist_manager(self) -> None:
174        """Broadcast ``{type: "consist", ...}`` on every create / update /
175        delete so the ``/consists`` page's htmx swap fires and any
176        scripting client subscribed to the ``consist`` topic is
177        notified without polling."""
178        from jmri_core.beans.consist import Consist  # noqa: PLC0415
179        from jmri_server.services.consist import _project  # noqa: PLC0415
180
181        cm = self._ctx.connection.consists
182
183        async def _emit(method: str, consist: Consist) -> None:
184            await self.broadcast({
185                "type": "consist", "method": method,
186                "data": _project(consist, cm),
187            })
188
189        async def on_created(c: Consist) -> None:
190            await _emit("put", c)
191
192        async def on_updated(c: Consist) -> None:
193            await _emit("post", c)
194
195        async def on_deleted(c: Consist) -> None:
196            await _emit("delete", c)
197
198        self._unsubscribes.append(cm.created.connect(on_created))
199        self._unsubscribes.append(cm.updated.connect(on_updated))
200        self._unsubscribes.append(cm.deleted.connect(on_deleted))
201
202    def _wire_bean_service(self, topic: str, svc: _BeanService) -> None:
203        manager = svc.manager()
204
205        # Subscribe state_changed on every bean already in the manager.
206        for bean in manager:
207            self._subscribe_bean(topic, svc, bean)
208
209        # And subscribe to future additions / removals on the manager itself.
210        self._unsubscribes.append(
211            manager.added.connect(self._make_added_listener(topic, svc))
212        )
213        self._unsubscribes.append(
214            manager.removed.connect(self._make_removed_listener(topic))
215        )
216
217    def _subscribe_bean(
218        self, topic: str, svc: _BeanService, bean: AbstractNamedBean
219    ) -> None:
220        listener = self._make_bean_listener(topic, svc, bean)
221        self._unsubscribes.append(bean.state_changed.connect(listener))
222
223    def _make_bean_listener(
224        self, topic: str, svc: _BeanService, bean: AbstractNamedBean
225    ) -> Callable[[Any], Awaitable[None]]:
226        async def listener(_payload: Any) -> None:
227            await self.broadcast({
228                "type": topic,
229                "data": svc.project(bean),
230            })
231        return listener
232
233    def _make_added_listener(
234        self, topic: str, svc: _BeanService
235    ) -> Callable[[AbstractNamedBean], Awaitable[None]]:
236        async def listener(bean: AbstractNamedBean) -> None:
237            self._subscribe_bean(topic, svc, bean)
238            await self.broadcast({
239                "type": topic,
240                "method": "put",
241                "data": svc.project(bean),
242            })
243        return listener
244
245    def _make_removed_listener(
246        self, topic: str
247    ) -> Callable[[AbstractNamedBean], Awaitable[None]]:
248        async def listener(bean: AbstractNamedBean) -> None:
249            await self.broadcast({
250                "type": topic,
251                "method": "delete",
252                "data": {"name": bean.system_name},
253            })
254        return listener
255
256    def _make_power_listener(
257        self, svc: JsonService
258    ) -> Callable[[Any], Awaitable[None]]:
259        async def listener(_payload: Any) -> None:
260            data = await svc.on_get(None)
261            await self.broadcast({"type": "power", "data": data})
262        return listener
263
264    # ---- shutdown ----------------------------------------------------------
265
266    async def aclose(self) -> None:
267        for unsub in self._unsubscribes:
268            unsub()
269        self._unsubscribes.clear()
270        # Best-effort: close any lingering sockets.
271        import contextlib  # noqa: PLC0415
272        for ws in list(self._clients):
273            with contextlib.suppress(Exception):
274                await ws.close()
275        self._clients.clear()
276
277
278# Forward import for type checker on Python 3.12.
279from jmri_server.service import JsonService  # noqa: E402, I001
280
281asyncio  # re-export marker — module name kept available for tests  # noqa: B018