minor stability changes to the behavior when using the stay-connected arg

Author: shaggysa

pybricksdev/cli/__init__.py

@@ -303,22 +303,12 @@ async def reconnect_hub():
                     # current program, so the menu was canceled and we are now printing
                     # the hub stdout until the user program ends on the hub.
                     try:
-                        await hub._wait_for_user_program_stop(
-                            2.25, raise_error_on_timeout=True
-                        )
+                        await hub._wait_for_power_button_release()
+                        await hub._wait_for_user_program_stop()
 
                     except HubDisconnectError:
                         hub = await reconnect_hub()
 
-                    except asyncio.TimeoutError:
-                        # On windows, it can take significantly longer
-                        # for the device to register that the hub was
-                        # disconnected. If _wait_for_user_program_stop
-                        # throws a timeout error, we can assume that the
-                        # hub was disconnected.
-                        await hub.disconnect()
-                        hub = await reconnect_hub()
-
                 except HubDisconnectError:
                     # let terminal cool off before making a new prompt
                     await asyncio.sleep(0.3)

pybricksdev/connections/pybricks.py

@@ -741,6 +741,46 @@ def handle_power_button_press(status: StatusFlag):
             )
         return awaitable_task.result()
 
+    async def _wait_for_power_button_release(
+        self, program_start_timeout=1, raise_error_on_timeout=False
+    ):
+        power_button_pressed: asyncio.Queue[bool] = asyncio.Queue()
+
+        with self.status_observable.pipe(
+            op.map(lambda s: bool(s & StatusFlag.POWER_BUTTON_PRESSED)),
+            op.distinct_until_changed(),
+        ).subscribe(lambda s: power_button_pressed.put_nowait(s)):
+            # The first item in the queue is the current status. The status
+            # could change before or after the last checksum is received,
+            # so this could be either true or false.
+            is_pressed = await self.race_disconnect(power_button_pressed.get())
+
+            if not is_pressed:
+                # If the button isn't already pressed,
+                # wait a short time for it to become pressed
+                try:
+                    await asyncio.wait_for(
+                        self.race_disconnect(power_button_pressed.get()),
+                        program_start_timeout,
+                    )
+                except asyncio.TimeoutError:
+                    if raise_error_on_timeout:
+                        raise
+                    # If the button never shows as pressed,
+                    # assume that we just missed the status flag
+                    logger.debug(
+                        "timed out waiting for power button press, assuming is was a short press"
+                    )
+                    return
+
+            # At this point, we know the button is pressed, so the
+            # next item in the queue must indicate the button has
+            # been released.
+            is_pressed = await self.race_disconnect(power_button_pressed.get())
+
+            # maybe catch mistake if the code is changed
+            assert not is_pressed
+
     async def _wait_for_user_program_stop(
         self, program_start_timeout=1, raise_error_on_timeout=False
     ):