jmri_withrottle

 33@dataclass
 34class WiThrottleServer:
 35    """TCP server speaking the WiThrottle line protocol."""
 36
 37    connection: SystemConnection
 38    # Callable so the server can be told about roster + power state
 39    # without us hard-wiring to jmri-server. Returns a list of dicts
 40    # in the roster JSON shape (id, dccAddress, longAddress).
 41    roster_provider: Callable[[], list[dict]] = field(default_factory=lambda: list)
 42    # Callable returning the current fast-clock ``PFT`` sync line, or
 43    # ``None`` if no fast clock is attached. Called on connect and on
 44    # every heartbeat so Engine Driver / WiThrottle iOS clocks stay in
 45    # sync without needing a separate broadcast path.
 46    fast_clock_provider: Callable[[], str | None] = field(
 47        default_factory=lambda: (lambda: None),
 48    )
 49    # Callable returning the current turnout inventory as a list of
 50    # ``{"system_name", "user_name", "state"}`` dicts. Sent once per
 51    # session (after the roster) so Engine Driver / WiThrottle iOS can
 52    # populate their Turnouts tab.
 53    turnout_provider: Callable[[], list[dict]] = field(
 54        default_factory=lambda: list,
 55    )
 56    # Handler for turnout throw/close/toggle commands from the client.
 57    # Called as ``await handler(system_name, action)`` where action is
 58    # one of ``"T"``, ``"C"``, or ``"2"``. If None, turnout actions from
 59    # the client are silently dropped.
 60    turnout_action_handler: (
 61        Callable[[str, str], Awaitable[None]] | None
 62    ) = None
 63    railroad: str = "Python JMRI"
 64    heartbeat_s: int = protocol.DEFAULT_HEARTBEAT_S
 65    web_port: int = 8000
 66    # How often to push a fresh PFT to every connected client. Engine
 67    # Driver's clock extrapolates from the last sync using the rate;
 68    # 5 s keeps it within a couple of fast-seconds of the server after
 69    # rate/state changes without spamming the wire.
 70    fast_clock_broadcast_s: float = 5.0
 71    _server: asyncio.Server | None = None
 72    _sessions: set[WiThrottleSession] = field(default_factory=set)
 73    _pft_task: asyncio.Task[None] | None = None
 74    # Disconnect callables from ``state_changed`` subscriptions and
 75    # the turnout manager's ``added``/``removed`` signals. Filled on
 76    # ``start()``, drained on ``stop()``.
 77    _turnout_unsubs: list[Callable[[], None]] = field(default_factory=list)
 78
 79    # ----- lifecycle -----------------------------------------------------
 80
 81    async def start(self, host: str = "0.0.0.0", port: int = DEFAULT_PORT) -> None:
 82        if self._server is not None:
 83            return
 84        self._server = await asyncio.start_server(
 85            self._handle, host=host, port=port,
 86        )
 87        self._pft_task = asyncio.create_task(
 88            self._pft_broadcast_loop(), name="withrottle-pft-broadcast",
 89        )
 90        # Wire turnout state broadcasts: every existing turnout gets a
 91        # ``state_changed`` listener that fires a PTA line to every
 92        # connected client. ``added``/``removed`` on the manager keep
 93        # the subscription set in sync so turnouts created later (via
 94        # the JSON API, web UI, or a driver monitor) also broadcast.
 95        # Must run in ``start()`` (post-``conn.open()``), not at
 96        # construction — drivers swap in their own turnout manager
 97        # during ``open()`` and any earlier subscription would leak to
 98        # the base manager.
 99        for t in self.connection.turnouts:
100            self._turnout_unsubs.append(self._subscribe_turnout(t))
101        self._turnout_unsubs.append(
102            self.connection.turnouts.added.connect(self._on_turnout_added)
103        )
104        _log.info("WiThrottle server listening on %s:%d", host, port)
105
106    def _subscribe_turnout(self, t: Turnout) -> Callable[[], None]:
107        """Attach a ``state_changed`` listener to *t* that pushes PTA."""
108        async def _listener(change: TurnoutStateChange) -> None:
109            # Only ``known`` transitions produce a wire-visible state on
110            # the Turnouts tab; commanded transitions still show the
111            # previous known state until the layout confirms.
112            if change.kind != "known":
113                return
114            await self.broadcast_turnout_state(
115                t.system_name, int(change.current),
116            )
117        return t.state_changed.connect(_listener)
118
119    async def _on_turnout_added(self, t: Turnout) -> None:
120        self._turnout_unsubs.append(self._subscribe_turnout(t))
121
122    async def stop(self) -> None:
123        if self._server is None:
124            return
125        if self._pft_task is not None:
126            self._pft_task.cancel()
127            import contextlib  # noqa: PLC0415
128            with contextlib.suppress(asyncio.CancelledError, Exception):
129                await self._pft_task
130            self._pft_task = None
131        for unsub in self._turnout_unsubs:
132            try:
133                unsub()
134            except Exception as e:  # noqa: BLE001
135                _log.debug("turnout unsubscribe failed: %s", e)
136        self._turnout_unsubs.clear()
137        for s in list(self._sessions):
138            await s.close()
139        self._server.close()
140        await self._server.wait_closed()
141        self._server = None
142        _log.info("WiThrottle server stopped")
143
144    async def broadcast_turnout_state(
145        self, system_name: str, state: int,
146    ) -> None:
147        """Push a ``PTA<state><sysName>`` line to every connected session.
148
149        The jmri_server wiring subscribes to each turnout's ``state_changed``
150        signal and calls this so the WiThrottle Turnouts tab stays in
151        sync with layout state — whether the change came from another
152        WiThrottle client, the web UI, or a physical cab.
153        """
154        if not self._sessions:
155            return
156        line = protocol.turnout_state(system_name, state)
157        for sess in list(self._sessions):
158            try:
159                await sess.send_line(line)
160            except Exception as e:  # noqa: BLE001
161                _log.debug("turnout broadcast to session failed: %s", e)
162
163    async def _pft_broadcast_loop(self) -> None:
164        """Periodically push the current PFT to every session.
165
166        Engine Driver clocks extrapolate locally from the last sync line,
167        but if the server-side clock state changes (start, stop, rate
168        change, big time jump) the client can drift until the next sync.
169        A ~15 s broadcast keeps every client aligned within one fast-clock
170        second at 8x rate, with negligible wire cost.
171        """
172        try:
173            while True:
174                await asyncio.sleep(self.fast_clock_broadcast_s)
175                pft = self.fast_clock_provider()
176                if pft is None or not self._sessions:
177                    continue
178                for sess in list(self._sessions):
179                    try:
180                        await sess.send_line(pft)
181                    except Exception as e:  # noqa: BLE001
182                        _log.debug("PFT broadcast to session failed: %s", e)
183        except asyncio.CancelledError:
184            return
185
186    @property
187    def is_running(self) -> bool:
188        return self._server is not None
189
190    @property
191    def sockets(self) -> list:
192        """Bound sockets — useful in tests to discover the assigned port."""
193        if self._server is None:
194            return []
195        return list(self._server.sockets)
196
197    # ----- per-connection handler ----------------------------------------
198
199    async def _handle(
200        self,
201        reader: asyncio.StreamReader,
202        writer: asyncio.StreamWriter,
203    ) -> None:
204        peer = writer.get_extra_info("peername")
205        _log.info("WiThrottle client connected: %s", peer)
206
207        async def send_line(line: str) -> None:
208            try:
209                writer.write((line + "\n").encode("utf-8"))
210                await writer.drain()
211            except (ConnectionResetError, BrokenPipeError):
212                pass
213
214        session = WiThrottleSession(
215            throttle_mgr=self.connection.throttles,
216            send_line=send_line,
217            consist_mgr=getattr(self.connection, "consists", None),
218            fast_clock_provider=self.fast_clock_provider,
219            turnout_action_handler=self.turnout_action_handler,
220        )
221        self._sessions.add(session)
222        try:
223            # Greeting + initial state.
224            for line in protocol.greeting(
225                name=self.railroad,
226                heartbeat_s=self.heartbeat_s,
227                web_port=self.web_port,
228            ):
229                await send_line(line)
230            roster = self.roster_provider()
231            if roster:
232                await send_line(protocol.roster_list(roster))
233            # Push the initial fast-clock reading so the client's clock
234            # doesn't sit at 12:00 PM waiting for the first heartbeat.
235            pft = self.fast_clock_provider()
236            if pft is not None:
237                await send_line(pft)
238            # Turnout state labels + inventory so the Turnouts tab renders.
239            # PTT (titles) must precede PTL (list) so the client already
240            # knows how to render each numeric state code. We skip both
241            # when the connection has no turnouts to avoid confusing
242            # clients that hide the Turnouts tab on an empty PTT.
243            turnouts = self.turnout_provider()
244            if turnouts:
245                await send_line(protocol.turnout_titles())
246                await send_line(protocol.turnout_list(turnouts))
247
248            # Read loop.
249            while True:
250                raw = await reader.readline()
251                if not raw:
252                    break
253                line = raw.decode("utf-8", errors="replace").strip()
254                if not line:
255                    continue
256                keep_going = await session.dispatch(line)
257                if not keep_going:
258                    break
259        except asyncio.CancelledError:
260            raise
261        except Exception as e:  # noqa: BLE001
262            _log.warning("WiThrottle session error from %s: %s", peer, e)
263        finally:
264            await session.close()
265            self._sessions.discard(session)
266            writer.close()
267            with _suppress_close():
268                await writer.wait_closed()
269            _log.info("WiThrottle client disconnected: %s", peer)