From 278c0ac0611f901d651f08f8f73aaf1d2cb712f5 Mon Sep 17 00:00:00 2001
From: Rodja Trappe <rodja@zauberzeug.com>
Date: Sat, 28 Sep 2024 06:37:20 +0200
Subject: [PATCH] make on air reconnects more robust by relying more on
 socketio reconnect mechanisms

---
 nicegui/air.py | 48 +++++++++++++++++++++---------------------------
 1 file changed, 21 insertions(+), 27 deletions(-)

diff --git a/nicegui/air.py b/nicegui/air.py
index f7fcf10da..90c01a052 100644
--- a/nicegui/air.py
+++ b/nicegui/air.py
@@ -137,11 +137,16 @@ def _handle_handshake(data: Dict[str, Any]) -> bool:
 
         @self.relay.on('client_disconnect')
         def _handle_disconnect(data: Dict[str, Any]) -> None:
+            self.log.info('disconnected.')
             client_id = data['client_id']
             if client_id not in Client.instances:
                 return
             Client.instances[client_id].handle_disconnect()
 
+        @self.relay.on('connect')
+        async def _handle_connect() -> None:
+            self.log.info('connected.')
+
         @self.relay.on('event')
         def _handle_event(data: Dict[str, Any]) -> None:
             client_id = data['client_id']
@@ -168,10 +173,6 @@ async def _handle_out_of_time() -> None:
             print('Sorry, you have reached the time limit of this NiceGUI On Air preview.', flush=True)
             await self.connect()
 
-        @self.relay.on('reconnect')
-        async def _handle_reconnect(_: Dict[str, Any]) -> None:
-            await self.connect()
-
     async def connect(self) -> None:
         """Connect to the NiceGUI On Air server."""
         if self.connecting:
@@ -181,30 +182,23 @@ async def connect(self) -> None:
             return
         self.log.debug('Going to connect...')
         self.connecting = True
-        backoff_time = 1
         try:
-            while True:
-                try:
-                    if self.relay.connected:
-                        await asyncio.wait_for(self.disconnect(), timeout=5)
-                    self.log.debug('Connecting...')
-                    await self.relay.connect(
-                        f'{RELAY_HOST}?device_token={self.token}',
-                        socketio_path='/on_air/socket.io',
-                        transports=['websocket', 'polling'],  # favor websocket over polling
-                    )
-                    self.log.debug('Connected.')
-                    break
-                except socketio.exceptions.ConnectionError:
-                    self.log.debug('Connection error.', stack_info=True)
-                except ValueError:  # NOTE this sometimes happens when the internal socketio client is not yet ready
-                    self.log.debug('ValueError while connecting.', stack_info=True)
-                except Exception:
-                    log.exception('Could not connect to NiceGUI On Air server.')
-
-                self.log.debug(f'Retrying in {backoff_time} seconds...')
-                await asyncio.sleep(backoff_time)
-                backoff_time = min(backoff_time * 2, 32)
+            if self.relay.connected:
+                await asyncio.wait_for(self.disconnect(), timeout=5)
+            self.log.debug('Connecting...')
+            await self.relay.connect(
+                f'{RELAY_HOST}?device_token={self.token}',
+                socketio_path='/on_air/socket.io',
+                transports=['websocket', 'polling'],  # favor websocket over polling
+            )
+            self.log.debug('Connected.')
+            return
+        except socketio.exceptions.ConnectionError:
+            self.log.debug('Connection error.', stack_info=True)
+        except ValueError:  # NOTE this sometimes happens when the internal socketio client is not yet ready
+            self.log.debug('ValueError while connecting.', stack_info=True)
+        except Exception:
+            log.exception('Could not connect to NiceGUI On Air server.')
         finally:
             self.connecting = False