jmri_withrottle.session

  1"""``WiThrottleSession`` — one connected client.
  2
  3Holds the per-connection state (client name, acquired throttles,
  4last-heartbeat timestamp) and bridges WiThrottle commands to the
  5active connection's :class:`ThrottleManager`. The TCP server in
  6``server.py`` owns one session per accepted socket.
  7"""
  8from __future__ import annotations
  9
 10import logging
 11from collections.abc import Awaitable, Callable
 12from dataclasses import dataclass, field
 13
 14from jmri_core.beans.consist_manager import ConsistManager
 15from jmri_core.beans.throttle import Throttle
 16from jmri_core.beans.throttle_manager import ThrottleManager
 17from jmri_core.state import ThrottleDirection
 18
 19from jmri_withrottle import protocol
 20
 21_log = logging.getLogger(__name__)
 22
 23# Speed comes in as 0..126 on WiThrottle (1 is the e-stop sentinel,
 24# 0 is stop, 2..126 is motion). Our throttle bean uses 0.0..1.0.
 25WITHROTTLE_SPEED_MAX = 126
 26
 27
 28SendLine = Callable[[str], Awaitable[None]]
 29
 30# Handler for turnout throw/close/toggle commanded from a WiThrottle client.
 31# Called as ``await handler(system_name, action)`` where action is one of
 32# ``"T"`` (throw), ``"C"`` (close), or ``"2"`` (toggle).
 33TurnoutActionHandler = Callable[[str, str], Awaitable[None]]
 34
 35
 36@dataclass(slots=True)
 37class _Acquired:
 38    throttle_char: str          # "T" or "S"
 39    address_token: str          # raw client-supplied token
 40    bean: Throttle              # the underlying core Throttle
 41    # When the acquired address matches a server-side consist, we
 42    # remember it here so throttle actions (V / R / F / X) dispatch
 43    # through the consist manager's fan-out instead of driving only
 44    # the lead's :class:`Throttle`. Software consists on any driver
 45    # then reach every member; CS-managed consists on LocoNet keep
 46    # working via LINK.
 47    consist_address: int | None = None
 48
 49
 50@dataclass(eq=False)   # identity-based eq+hash so set() tracking works
 51class WiThrottleSession:
 52    """One WiThrottle client connection."""
 53
 54    throttle_mgr: ThrottleManager
 55    send_line: SendLine                          # writer to push lines back
 56    consist_mgr: ConsistManager | None = None
 57    fast_clock_provider: Callable[[], str | None] | None = None
 58    turnout_action_handler: TurnoutActionHandler | None = None
 59    client_name: str = ""
 60    hardware_id: str = ""
 61    _acquired: dict[tuple[str, str], _Acquired] = field(default_factory=dict)
 62
 63    # ----- entry point ----------------------------------------------------
 64
 65    async def dispatch(self, line: str) -> bool:
 66        """Process one inbound line. Returns False if the client said Quit."""
 67        if line and _log.isEnabledFor(logging.DEBUG):
 68            _log.debug("WiThrottle IN: %r", line)
 69        msg = protocol.decode_line(line)
 70        if isinstance(msg, protocol.Heartbeat):
 71            # Engine Driver and WiThrottle iOS treat the heartbeat as
 72            # ping/pong — if the server doesn't echo, the client thinks
 73            # the link is dead and tears down the TCP session. Always
 74            # echo back.
 75            await self.send_line("*")
 76            # Piggyback a fast-clock resync on the heartbeat so the
 77            # client's clock stays aligned without needing a separate
 78            # broadcast timer.
 79            if self.fast_clock_provider is not None:
 80                pft = self.fast_clock_provider()
 81                if pft is not None:
 82                    await self.send_line(pft)
 83            return True
 84        if isinstance(msg, protocol.Quit):
 85            await self._release_all()
 86            return False
 87        if isinstance(msg, protocol.HelloName):
 88            self.client_name = msg.name
 89            return True
 90        if isinstance(msg, protocol.HelloHardwareId):
 91            self.hardware_id = msg.hardware_id
 92            return True
 93        if isinstance(msg, protocol.AcquireLoco):
 94            await self._acquire(msg)
 95            return True
 96        if isinstance(msg, protocol.ReleaseLoco):
 97            await self._release(msg.throttle, msg.address_token)
 98            return True
 99        if isinstance(msg, protocol.ThrottleAction):
100            await self._action(msg)
101            return True
102        if isinstance(msg, protocol.TurnoutAction):
103            if self.turnout_action_handler is not None:
104                try:
105                    await self.turnout_action_handler(msg.system_name, msg.action)
106                except Exception as e:  # noqa: BLE001
107                    _log.warning(
108                        "WiThrottle turnout %s action %r failed: %s",
109                        msg.system_name, msg.action, e,
110                    )
111            return True
112        if isinstance(msg, protocol.Unknown):
113            if msg.raw:
114                _log.debug("WiThrottle unknown line: %r", msg.raw)
115            return True
116        return True
117
118    # ----- acquire / release ---------------------------------------------
119
120    async def _acquire(self, msg: protocol.AcquireLoco) -> None:
121        try:
122            addr, is_long = protocol.parse_address_token(msg.address_token)
123        except ValueError:
124            _log.warning("WiThrottle bad address token: %r", msg.address_token)
125            return
126        # Consist first-check: if the address matches a server-side
127        # consist, we still acquire a Throttle for the lead so the
128        # client's initial state (function labels, speed, direction)
129        # renders normally, but subsequent throttle actions fan through
130        # the consist manager to every member. Engine Driver /
131        # WiThrottle iOS have no wire-level notion of "this is a
132        # consist" — MU support is a server-side transparent thing.
133        consist_addr: int | None = None
134        if self.consist_mgr is not None:
135            consist = self.consist_mgr.get(addr)
136            if consist is not None and consist.locos:
137                consist_addr = addr
138                # Use the lead loco's address for the underlying
139                # throttle acquire so the client sees the lead's
140                # function/speed state.
141                lead = consist.locos[0]
142                addr = lead.address
143                is_long = lead.long_address
144        try:
145            bean = await self.throttle_mgr.acquire(addr, long_address=is_long)
146        except Exception as e:  # noqa: BLE001
147            _log.warning("WiThrottle acquire %s failed: %s", msg.address_token, e)
148            return
149        self._acquired[(msg.throttle, msg.address_token)] = _Acquired(
150            throttle_char=msg.throttle,
151            address_token=msg.address_token,
152            bean=bean,
153            consist_address=consist_addr,
154        )
155        await self.send_line(
156            protocol.throttle_acquired(msg.throttle, msg.address_token)
157        )
158        # Broadcast the throttle's initial state so Engine Driver /
159        # WiThrottle iOS know how to render its buttons. Without the
160        # per-function ``F<state><n>`` lines the client doesn't display
161        # the function bar at all; without the initial V / R the slider
162        # sits at an undefined position.
163        await self._broadcast_initial_state(
164            msg.throttle, msg.address_token, bean,
165        )
166
167    async def _broadcast_initial_state(
168        self, throttle: str, address_token: str, bean: Throttle,
169    ) -> None:
170        """Send the state lines Engine Driver expects after an acquire ack.
171
172        Order matches JMRI's ``ThrottleController.notifyControllerAddressFound``.
173
174        The label line matters and must be present with non-empty labels:
175
176        * Roster acquires (``M+L<addr><;>E<id>``) hide the function bar
177          if we skip the labels line — the client expects roster
178          metadata and interprets a missing line as "no functions".
179        * Manual acquires (``M+L<addr><;>L<addr>``) hide the bar if we
180          send *empty* labels — the client interprets those as "loco
181          has no functions". Non-empty defaults (``F0``..``F28``) work
182          for both cases.
183
184        Order sent:
185
186        1. ``M<mux>L<addr><;>]\\[F0]\\[F1]\\[...`` — function labels
187        2. ``F<state><fn>`` — one line per F0..F28 (function states)
188        3. ``V<speed>`` — initial speed (0..126)
189        4. ``R<0|1>`` — initial direction
190        5. ``s<mode>`` — speed step mode (1 = 128-step NMRA DCC)
191        """
192        # Labels first, populated with default F0..F28 so both roster
193        # and manual acquisitions render the function bar.
194        await self.send_line(
195            protocol.throttle_function_labels(throttle, address_token)
196        )
197        # Function states next — must precede speed/direction or the
198        # client hides the function bar.
199        for fn in range(protocol.MAX_FUNCTION + 1):
200            await self.send_line(
201                protocol.throttle_function(
202                    throttle, address_token, fn, bean.functions[fn],
203                )
204            )
205        speed_wire = int(round(WITHROTTLE_SPEED_MAX * max(0.0, bean.speed)))
206        await self.send_line(
207            protocol.throttle_speed(throttle, address_token, speed_wire)
208        )
209        forward = bean.direction == ThrottleDirection.FORWARD
210        await self.send_line(
211            protocol.throttle_direction(throttle, address_token, forward)
212        )
213        # Speed step mode — 1 = NMRA DCC 128-step (the modern default;
214        # NCE, LocoNet, and DCC-EX all speak this).
215        await self.send_line(
216            protocol.throttle_speed_step_mode(throttle, address_token, 1)
217        )
218
219    async def _release(self, throttle: str, address_token: str) -> None:
220        key = (throttle, address_token)
221        a = self._acquired.pop(key, None)
222        if a is None:
223            return
224        try:
225            await self.throttle_mgr.release(a.bean)
226        except Exception as e:  # noqa: BLE001
227            _log.warning("WiThrottle release %s failed: %s", address_token, e)
228        await self.send_line(protocol.throttle_released(throttle, address_token))
229
230    async def _release_all(self) -> None:
231        for (throttle, token) in list(self._acquired):
232            await self._release(throttle, token)
233
234    # ----- per-throttle actions ------------------------------------------
235
236    async def _action(self, msg: protocol.ThrottleAction) -> None:
237        # WiThrottle allows a broadcast address token ``*`` meaning
238        # "every throttle currently acquired on this multiplex
239        # character". Engine Driver uses this for slider drags — so
240        # speed commands arrive as ``M0A*<;>V29`` while direction
241        # arrives as ``M0AL463<;>R0``. Without ``*`` support, direction
242        # works but speed doesn't (the bug we saw on hardware).
243        if msg.address_token == "*":
244            targets = [
245                acq for (t, _tok), acq in self._acquired.items()
246                if t == msg.throttle
247            ]
248        else:
249            target = self._acquired.get((msg.throttle, msg.address_token))
250            targets = [target] if target is not None else []
251        if not targets:
252            return
253        for a in targets:
254            await self._apply_action(a, msg)
255
256    async def _apply_action(
257        self, a: _Acquired, msg: protocol.ThrottleAction,
258    ) -> None:
259        op = msg.action
260        if op == "V":
261            try:
262                wire = int(msg.value)
263            except ValueError:
264                return
265            await self._set_speed(a, wire)
266        elif op == "R":
267            forward = msg.value == "1"
268            direction = (
269                ThrottleDirection.FORWARD if forward
270                else ThrottleDirection.REVERSE
271            )
272            if a.consist_address is not None and self.consist_mgr is not None:
273                await self.consist_mgr.set_direction(a.consist_address, direction)
274            else:
275                await a.bean.set_direction(direction)
276            await self.send_line(
277                protocol.throttle_direction(a.throttle_char, a.address_token, forward)
278            )
279        elif op in ("F", "f"):
280            # Format: state then function number; e.g. "F128" = pressed F28,
281            # "F01" = released F1. "F" is press/release toggle, "f" is
282            # force-set (which we treat the same for alpha).
283            if not msg.value:
284                return
285            state = msg.value[0] == "1"
286            try:
287                fn = int(msg.value[1:])
288            except ValueError:
289                return
290            if not 0 <= fn <= protocol.MAX_FUNCTION:
291                return
292            if a.consist_address is not None and self.consist_mgr is not None:
293                await self.consist_mgr.set_function(a.consist_address, fn, state)
294            else:
295                await a.bean.set_function(fn, state)
296            await self.send_line(
297                protocol.throttle_function(a.throttle_char, a.address_token, fn, state)
298            )
299        elif op == "X":
300            # E-stop: full reverse direction on JMRI is encoded as speed
301            # ``-1`` in our bean; the wire encoding goes via the same
302            # set_speed path. Use the bean's emergency_stop helper.
303            if a.consist_address is not None and self.consist_mgr is not None:
304                await self.consist_mgr.emergency_stop(a.consist_address)
305            else:
306                await a.bean.emergency_stop()
307            await self.send_line(
308                protocol.throttle_speed(a.throttle_char, a.address_token, 1)
309            )
310        elif op == "I":
311            # Idle: speed 0 without changing direction.
312            if a.consist_address is not None and self.consist_mgr is not None:
313                await self.consist_mgr.set_speed(a.consist_address, 0.0)
314            else:
315                await a.bean.set_speed(0.0)
316            await self.send_line(
317                protocol.throttle_speed(a.throttle_char, a.address_token, 0)
318            )
319        # "L" (speed-step mode) accepted silently for alpha.
320
321    async def _set_speed(self, a: _Acquired, wire_speed: int) -> None:
322        if wire_speed < 0:
323            wire_speed = 0
324        if wire_speed > WITHROTTLE_SPEED_MAX:
325            wire_speed = WITHROTTLE_SPEED_MAX
326        # Speed target: fan through the consist manager when this
327        # acquired throttle is bound to a server-side consist.
328        if a.consist_address is not None and self.consist_mgr is not None:
329            if wire_speed == 0:
330                await self.consist_mgr.set_speed(a.consist_address, 0.0)
331            elif wire_speed == 1:
332                await self.consist_mgr.emergency_stop(a.consist_address)
333            else:
334                fraction = (wire_speed - 1) / (WITHROTTLE_SPEED_MAX - 1)
335                await self.consist_mgr.set_speed(a.consist_address, fraction)
336        elif wire_speed <= 1:
337            # 0 = stop, 1 = e-stop sentinel; both map to the same throttle bean call.
338            if wire_speed == 0:
339                await a.bean.set_speed(0.0)
340            else:
341                await a.bean.emergency_stop()
342        else:
343            # Map 2..126 onto 0.0..1.0 linearly.
344            fraction = (wire_speed - 1) / (WITHROTTLE_SPEED_MAX - 1)
345            await a.bean.set_speed(fraction)
346        await self.send_line(
347            protocol.throttle_speed(a.throttle_char, a.address_token, wire_speed)
348        )
349
350    # ----- cleanup -------------------------------------------------------
351
352    async def close(self) -> None:
353        """Release all acquired throttles. Called when the socket goes away."""
354        await self._release_all()