feat: keep connection alive

Author: humbertogontijo

roborock/api.py

@@ -191,7 +191,7 @@ async def _async_response(self, request_id: int, protocol_id: int = 0) -> tuple[
         except (asyncio.TimeoutError, asyncio.CancelledError):
             raise RoborockTimeout(f"id={request_id} Timeout after {QUEUE_TIMEOUT} seconds") from None
         finally:
-            del self._waiting_queue[request_id]
+            self._waiting_queue.pop(request_id, None)
 
     def _get_payload(
         self,

roborock/local_api.py

@@ -30,6 +30,7 @@ def __init__(self, device_data: DeviceData):
         self.remaining = b""
         self.transport: Transport | None = None
         self._mutex = Lock()
+        self.keep_alive_func()
 
     def data_received(self, message):
         if self.remaining:
@@ -45,6 +46,13 @@ def connection_lost(self, exc: Optional[Exception]):
     def is_connected(self):
         return self.transport and self.transport.is_reading()
 
+    def keep_alive_func(self, _=None):
+        keep_alive_task = asyncio.gather(
+            asyncio.sleep(10),
+            self.ping(),
+        )
+        keep_alive_task.add_done_callback(self.keep_alive_func)
+
     async def async_connect(self) -> None:
         async with self._mutex:
             try:
@@ -87,7 +95,18 @@ def build_roborock_message(self, method: RoborockCommand, params: Optional[list
         )
 
     async def ping(self):
-        return await self.send_command(RoborockCommand.APP_WAKEUP_ROBOT)
+        request_id = 2
+        _LOGGER.debug(f"id={request_id} Requesting method ping with None")
+        try:
+            return await self.send_message(RoborockMessage(
+                                protocol=2,
+                                payload=None,
+                                seq=request_id,
+                                version=b'1.0',
+                                random=23
+                                ))
+        except Exception as e:
+            _LOGGER.error(e)
 
     async def send_command(self, method: RoborockCommand, params: Optional[list | dict] = None):
         roborock_message = self.build_roborock_message(method, params)
@@ -96,9 +115,9 @@ async def send_command(self, method: RoborockCommand, params: Optional[list | di
     async def async_local_response(self, roborock_message: RoborockMessage):
         method = roborock_message.get_method()
         request_id: int | None
-        if method and not method.startswith("get"):
+        if not method or not method.startswith("get"):
             request_id = roborock_message.seq
-            response_protocol = 5
+            response_protocol = request_id + 1
         else:
             request_id = roborock_message.get_request_id()
             response_protocol = 4
@@ -107,7 +126,8 @@ async def async_local_response(self, roborock_message: RoborockMessage):
         (response, err) = await self._async_response(request_id, response_protocol)
         if err:
             raise CommandVacuumError("", err) from err
-        _LOGGER.debug(f"id={request_id} Response from {roborock_message.get_method()}: {response}")
+        method = roborock_message.get_method() if roborock_message.protocol != 2 else "ping"
+        _LOGGER.debug(f"id={request_id} Response from method {method}: {response}")
         return response
 
     def _send_msg_raw(self, data: bytes):

roborock/protocol.py

@@ -176,19 +176,21 @@ def _parse(self, stream, context, path):
         subcon1 = Optional(Int16ub)
         length = subcon1.parse_stream(stream, **context)
         if not length:
-            subcon1.parse_stream(stream, **context)  # seek 2
+            if length == 0:
+                subcon1.parse_stream(stream, **context)  # seek 2
             return None
         subcon2 = Bytes(length)
         obj = subcon2.parse_stream(stream, **context)
         return self._decode(obj, context, path)
 
     def _build(self, obj, stream, context, path):
-        obj2 = self._encode(obj, context, path)
-        subcon1 = Int16ub
-        length = len(obj2)
-        subcon1.build_stream(length, stream, **context)
-        subcon2 = Bytes(length)
-        subcon2.build_stream(obj2, stream, **context)
+        if obj is not None:
+            obj2 = self._encode(obj, context, path)
+            subcon1 = Int16ub
+            length = len(obj2)
+            subcon1.build_stream(length, stream, **context)
+            subcon2 = Bytes(length)
+            subcon2.build_stream(obj2, stream, **context)
         return obj
 
     def _encode(self, obj, context, _):

roborock/roborock_message.py

@@ -12,7 +12,7 @@
 @dataclass
 class RoborockMessage:
     protocol: int
-    payload: bytes
+    payload: bytes | None
     seq: int = randint(100000, 999999)
     version: bytes = b"1.0"
     random: int = randint(10000, 99999)