"""Module to connect to Automower with websocket.""" import asyncio import contextlib import datetime import logging import zoneinfo from collections.abc import Callable from typing import Any, cast import tzlocal from aiohttp import ClientError, WSMessage, WSMsgType from .auth import AbstractAuth from .commands import AutomowerEndpoint, MowerCommands from .const import REST_POLL_CYCLE, EventTypesV2 from .exceptions import HusqvarnaWSClientError, NoDataAvailableError from .model import Message, MessageData, MowerAttributes, SingleMessageData from .model_input import ( CuttingHeightAttributes, GenericEventData, HeadLightAttributes, MowerDataItem, MowerDataResponse, PositionAttributes, ) from .utils import mower_list_to_dictionary_dataclass _LOGGER = logging.getLogger(__name__) INVALID_MOWER_ID = "0-0" class AutomowerSession: """Automower API to communicate with an Automower. The `AutomowerSession` is the primary API service for this library. It supports operations like getting a status or sending commands. """ __slots__ = ( "_data", "_reconnect_lock", "auth", "commands", "current_mowers", "data", "data_update_cbs", "invalid_mowers", "last_ws_message", "loop", "message_update_cbs", "messages", "mower_tz", "poll", "pong_cbs", "reconnect_task", "rest_task", "single_message", "single_message_cbs", "ws_disconnected_cbs", "ws_ready_cbs", "ws_task", ) def __init__( self, auth: AbstractAuth, mower_tz: zoneinfo.ZoneInfo | None = None, *, poll: bool = False, ) -> None: """Create a session. :param class auth: The AbstractAuth class from aioautomower.auth. :param class mower_tz: The ZoneInfo object for the mower, default is None :param bool poll: Poll data with rest if True. """ self._data: MowerDataResponse | None = None self.auth = auth self.mower_tz = mower_tz or tzlocal.get_localzone() self.data: dict[str, MowerAttributes] = {} self.commands = MowerCommands(self.auth, self.data, self.mower_tz) self.current_mowers: set[str] = set() self.data_update_cbs: list[Callable[[dict[str, MowerAttributes]], None]] = [] self.invalid_mowers: set[str] = set() self.last_ws_message: datetime.datetime self.loop: asyncio.AbstractEventLoop = asyncio.get_running_loop() self.message_update_cbs: list[tuple[str, Callable[[MessageData], None]]] = [] self.single_message_cbs: list[Callable[[SingleMessageData], None]] = [] self.single_message: SingleMessageData | None = None self.messages: dict[str, MessageData] = {} self.poll = poll self.pong_cbs: list[Callable[[datetime.datetime], None]] = [] self.rest_task: asyncio.Task[None] | None = None self.ws_task: asyncio.Task[None] | None = None self.reconnect_task: asyncio.Task[None] | None = None self.ws_disconnected_cbs: list[Callable[[], None]] = [] self.ws_ready_cbs: list[Callable[[], None]] = [] self._reconnect_lock = asyncio.Lock() _LOGGER.debug("self.mower_tz: %s", self.mower_tz) def register_data_callback( self, callback: Callable[[dict[str, MowerAttributes]], None] ) -> None: """Register a data update callback.""" if callback not in self.data_update_cbs: self.data_update_cbs.append(callback) def unregister_data_callback( self, callback: Callable[[dict[str, MowerAttributes]], None] ) -> None: """Unregister a data update callback. :param func callback: Takes one function, which should be unregistered. """ if callback in self.data_update_cbs: self.data_update_cbs.remove(callback) def _schedule_data_callback( self, cb: Callable[[dict[str, MowerAttributes]], None] ) -> None: """Schedule a data callback.""" if self.poll and self.data is None: raise NoDataAvailableError self.loop.call_soon_threadsafe(cb, self.data) def _schedule_data_callbacks(self) -> None: """Schedule a data callbacks.""" for cb in self.data_update_cbs: self._schedule_data_callback(cb) def register_single_message_callback( self, cb: Callable[[SingleMessageData], None] ) -> None: """Register a callback for the latest single message of a specific mower.""" self.single_message_cbs.append(cb) def unregister_single_message_callback( self, cb: Callable[[SingleMessageData], None], ) -> None: """Unregister a single message callback for a specific mower.""" if cb in self.single_message_cbs: self.single_message_cbs.remove(cb) def _schedule_single_message_callbacks(self, message: SingleMessageData) -> None: """Dispatch only the most recent message to registered callbacks.""" for cb in self.single_message_cbs: self.loop.call_soon_threadsafe(cb, message) def register_message_callback( self, callback: Callable[[MessageData], None], mower_id: str, ) -> None: """Register a callback triggered when new messages arrive for specific mower.""" self.message_update_cbs.append((mower_id, callback)) def unregister_message_callback( self, callback: Callable[[MessageData], None], mower_id: str, ) -> None: """Unregister a previously registered message callback.""" self.message_update_cbs.remove((mower_id, callback)) def _schedule_message_callback( self, msg_data: MessageData, cb: Callable[[MessageData], None], ) -> None: """Schedule a single message data callback (thread-safe).""" self.loop.call_soon_threadsafe(cb, msg_data) def _schedule_message_callbacks(self, msg_data: MessageData) -> None: """Schedule all registered message data callbacks for the given mower.""" for mower_id, cb in self.message_update_cbs: if mower_id == msg_data.id: self._schedule_message_callback(msg_data, cb) def register_ws_ready_callback(self, cb: Callable[[], None]) -> None: """Register a callback that is called when WebSocket is ready.""" if cb not in self.ws_ready_cbs: self.ws_ready_cbs.append(cb) def _schedule_ws_ready_callback(self) -> None: """Schedule all ws_ready callbacks (thread-safe).""" if not self.ws_ready_cbs: return for cb in list(self.ws_ready_cbs): try: self.loop.call_soon_threadsafe(cb) except RuntimeError: _LOGGER.exception("Error while scheduling ws_ready callback %s", cb) def unregister_ws_ready_callback(self, cb: Callable[[], None]) -> None: """Unregister a ws_ready callback.""" if cb in self.ws_ready_cbs: self.ws_ready_cbs.remove(cb) def register_ws_disconnected_callback(self, cb: Callable[[], None]) -> None: """Register a callback that is called when WebSocket is disconnected.""" if cb not in self.ws_disconnected_cbs: self.ws_disconnected_cbs.append(cb) def _schedule_ws_disconnected_callback(self) -> None: """Schedule all ws_disconnected callbacks (thread-safe).""" if not self.ws_disconnected_cbs: return for cb in list(self.ws_disconnected_cbs): try: self.loop.call_soon_threadsafe(cb) except RuntimeError: _LOGGER.exception( "Error while scheduling ws_disconnected callback %s", cb ) def unregister_ws_disconnected_callback(self, cb: Callable[[], None]) -> None: """Unregister a ws_disconnected callback.""" if cb in self.ws_disconnected_cbs: self.ws_disconnected_cbs.remove(cb) def register_pong_callback( self, pong_callback: Callable[[datetime.datetime], None] ) -> None: """Register a pong callback. It's not real ping/pong, but a way to check if the websocket is still alive, by receiving an empty message. """ if pong_callback not in self.pong_cbs: self.pong_cbs.append(pong_callback) def _schedule_pong_callback(self, cb: Callable[[datetime.datetime], None]) -> None: """Schedule a pong callback.""" self.loop.call_soon_threadsafe(cb, self.last_ws_message) def _schedule_pong_callbacks(self) -> None: """Schedule pong callbacks.""" for cb in self.pong_cbs: self._schedule_pong_callback(cb) def unregister_pong_callback( self, pong_callback: Callable[[datetime.datetime], None] ) -> None: """Unregister a pong update callback. :param func callback: Takes one function, which should be unregistered. """ if pong_callback in self.pong_cbs: self.pong_cbs.remove(pong_callback) async def connect(self) -> None: """Connect to the API. This method handles the login. Also a REST task will be started, which periodically polls the REST endpoint, when polling is set to true. """ self._schedule_data_callbacks() if self.poll: await self.get_status() for mower_id in self.current_mowers: await self.async_get_messages(mower_id) self.rest_task = asyncio.create_task(self._rest_task()) async def _handle_text_message(self, msg: WSMessage) -> None: """Process a text message to data.""" if not msg.data: self.last_ws_message = datetime.datetime.now(tz=datetime.UTC) _LOGGER.debug("last_ws_message:%s", self.last_ws_message) self._schedule_pong_callbacks() if msg.data: msg_dict: GenericEventData = msg.json() if "type" in msg_dict: if msg_dict["type"] in {event.value for event in EventTypesV2}: _LOGGER.debug("Received websocket message %s", msg_dict) if msg_dict["id"] == INVALID_MOWER_ID: return if msg_dict["id"] not in self.current_mowers: _LOGGER.debug("New mower detected %s", msg_dict["id"]) self.current_mowers.add(msg_dict["id"]) await self.get_status() self._update_data(msg_dict) else: _LOGGER.debug("Received unknown ws type %s", msg_dict["type"]) elif "ready" in msg_dict and "connectionId" in msg_dict: _LOGGER.debug( "Websocket ready=%s (id='%s')", msg_dict["ready"], msg_dict["connectionId"], ) self._schedule_ws_ready_callback() async def _listen(self) -> None: """Start listening to the websocket (and receive initial state).""" ws = self.auth.ws if ws is None: exc = RuntimeError("WebSocket not connected") raise exc while not ws.closed: try: msg = await ws.receive(timeout=300) if msg.type in ( WSMsgType.CLOSE, WSMsgType.CLOSING, WSMsgType.CLOSED, ): break if msg.type == WSMsgType.TEXT: await self._handle_text_message(msg) elif msg.type == WSMsgType.ERROR: continue except TimeoutError: _LOGGER.debug("Timeout receiving from websocket, scheduling reconnect") self.loop.create_task(self.reconnect()) break except ClientError as exc: _LOGGER.debug( "ClientError from websocket, scheduling reconnect: %s", exc ) self.loop.create_task(self.reconnect()) break async def start_listening(self) -> None: """Start listening to the websocket.""" self.ws_task = self.loop.create_task(self._listen()) self.reconnect_task = self.loop.create_task(self._reconnect_scheduler()) async def _reconnect_scheduler(self) -> None: """Reconnect scheduler.""" await asyncio.sleep(7195) _LOGGER.debug("Websocket needs to be reconnected") await self.reconnect() async def reconnect(self) -> None: """Reconnect to the websocket.""" async with self._reconnect_lock: _LOGGER.debug("Attempting to reconnect to WebSocket (make-before-break)") # keep reference to old ws (if any) old_ws = getattr(self.auth, "ws", None) # 1. Make new connection try: await self.auth.websocket_connect() _LOGGER.debug("New WebSocket connection successful") except HusqvarnaWSClientError: _LOGGER.exception("Failed to establish new WebSocket connection: %s") if not self.reconnect_task or self.reconnect_task.done(): self.reconnect_task = asyncio.create_task( self._reconnect_scheduler() ) return # 2. Break old connection and replace tasks _LOGGER.debug("Closing old WebSocket connection and tasks") # cancel old listener task (if running) if self.ws_task and not self.ws_task.done(): self.ws_task.cancel() with contextlib.suppress(asyncio.CancelledError): await asyncio.wait_for(self.ws_task, timeout=2) # cancel old reconnect scheduler (if running) if self.reconnect_task and not self.reconnect_task.done(): self.reconnect_task.cancel() with contextlib.suppress(asyncio.CancelledError): await asyncio.wait_for(self.reconnect_task, timeout=2) # explicitly close previous websocket if it's different and still open new_ws = getattr(self.auth, "ws", None) try: if old_ws is not None and old_ws is not new_ws: # aiohttp ws has .closed property and .close() coroutine with contextlib.suppress(Exception): await old_ws.close() self._schedule_ws_disconnected_callback() _LOGGER.debug("Closed previous websocket connection") except Exception as exc: # noqa: BLE001 _LOGGER.debug("Error closing previous websocket: %s", exc) _LOGGER.debug("Starting new WebSocket listener") await self.start_listening() async def send_empty_message(self, ping_timeout: int = 5) -> bool: """Send a single ping with timeout (in seconds).""" ws = getattr(self.auth, "ws", None) if ws is None: _LOGGER.debug("WebSocket not connected—skipping ping") await self.auth.websocket_connect() return False try: await asyncio.wait_for(ws.send_str(""), timeout=ping_timeout) except TimeoutError: return False except ClientError as err: _LOGGER.debug("Ping failed due to client error: %s", err) await self.reconnect() return False return True async def get_status(self) -> dict[str, MowerAttributes]: """Get mower status via REST.""" mower_list: MowerDataResponse = await self.auth.get_json( AutomowerEndpoint.mowers ) self.invalid_mowers.clear() valid_mowers = [] for mower in mower_list["data"]: if mower["id"] == INVALID_MOWER_ID: self.invalid_mowers.add(mower["attributes"]["system"]["name"]) continue valid_mowers.append(mower) _LOGGER.debug("invalid_mowers: %s", self.invalid_mowers) self._data = { **mower_list, "data": valid_mowers, } self.data = mower_list_to_dictionary_dataclass( self._data, self.mower_tz, ) self.current_mowers = set(self.data.keys()) _LOGGER.debug("current_mowers: %s", self.current_mowers) self.commands = MowerCommands(self.auth, self.data, self.mower_tz) return self.data async def async_get_messages(self, mower_id: str) -> MessageData: """Fetch messages for one mower and merge into self._messages.""" raw_data = await self.auth.get_json( AutomowerEndpoint.messages.format(mower_id=mower_id) ) msg_resp = MessageData.from_dict(raw_data["data"]) msg_resp.id = mower_id self.messages[mower_id] = msg_resp return msg_resp def _update_data(self, new_data: GenericEventData) -> None: """Update internal data with new data from websocket.""" if new_data["type"] == EventTypesV2.MESSAGES: if self.messages: new_msg = Message.from_dict(new_data["attributes"]["message"]) mower_id = new_data["id"] self.messages[mower_id].attributes.messages.insert(0, new_msg) self._schedule_message_callbacks(self.messages[mower_id]) if not self.messages: single_message = SingleMessageData.from_dict(new_data) if single_message != self.single_message: self.single_message = single_message self._schedule_single_message_callbacks(self.single_message) else: _LOGGER.debug( "Received same single message as before, not updating." ) if self._data is None: raise NoDataAvailableError data = self._data["data"] for mower in data: if mower["type"] == "mower" and mower["id"] == new_data["id"]: self._process_event(mower, new_data) break self.data = mower_list_to_dictionary_dataclass(self._data, self.mower_tz) self.commands = MowerCommands(self.auth, self.data, self.mower_tz) self._schedule_data_callbacks() def _process_event(self, mower: MowerDataItem, new_data: GenericEventData) -> None: """Process a specific event type.""" handlers: dict[str, Callable[[MowerDataItem, Any], None]] = { "cuttingHeight": self._handle_cutting_height_event, "headlights": self._handle_headlight_event, "position": self._handle_position_event, } attributes: dict[str, Any] = new_data.get("attributes", {}) for key, handler in handlers.items(): if key in attributes: handler(mower, attributes) # Pass the specific attribute return mower_attributes = mower["attributes"] # General handling for other attributes self._update_nested_dict(cast("dict[str, Any]", mower_attributes), attributes) def _handle_cutting_height_event( self, mower: MowerDataItem, attributes: CuttingHeightAttributes ) -> None: """Handle cuttingHeight-specific updates.""" mower["attributes"]["settings"]["cuttingHeight"] = attributes["cuttingHeight"][ "height" ] def _handle_headlight_event( self, mower: MowerDataItem, attributes: HeadLightAttributes ) -> None: """Handle headLight-specific updates.""" mower["attributes"]["settings"]["headlight"]["mode"] = attributes["headlights"][ "mode" ] def _handle_position_event( self, mower: MowerDataItem, attributes: PositionAttributes ) -> None: mower["attributes"]["positions"].insert(0, attributes["position"]) @staticmethod def _update_nested_dict(original: dict[str, Any], updates: dict[str, Any]) -> None: """Recursively update a nested dictionary with new values.""" for key, value in updates.items(): if ( isinstance(value, dict) and key in original and isinstance(original[key], dict) ): AutomowerSession._update_nested_dict(original[key], value) else: original[key] = value async def _rest_task(self) -> None: """Poll data periodically via Rest.""" while True: await self.get_status() self._schedule_data_callbacks() await asyncio.sleep(REST_POLL_CYCLE) async def close(self) -> None: """Close the session.""" if self.rest_task: if not self.rest_task.cancelled(): self.rest_task.cancel() with contextlib.suppress(asyncio.CancelledError): await asyncio.gather(self.rest_task) if self.ws_task: self.ws_task.cancel() with contextlib.suppress(asyncio.CancelledError): await self.ws_task if self.reconnect_task: self.reconnect_task.cancel() with contextlib.suppress(asyncio.CancelledError): await self.reconnect_task for mower_id, cb in self.message_update_cbs[:]: self.unregister_message_callback(cb, mower_id)