tools: optional connection timeout

src/infuse_iot/socket_comms.py

@@ -206,16 +206,22 @@ class DataType(enum.IntFlag):
         DATA = 2
         LOGGING = 4
 
-    def __init__(self, infuse_id: int, data_types: DataType):
+    def __init__(self, infuse_id: int, data_types: DataType, timeout_ms: int):
         super().__init__(infuse_id)
         self.data_types = data_types
+        self.timeout_ms = timeout_ms
 
     def to_json(self) -> dict:
-        return {"type": int(self.TYPE), "infuse_id": self.infuse_id, "data_types": self.data_types}
+        return {
+            "type": int(self.TYPE),
+            "infuse_id": self.infuse_id,
+            "data_types": self.data_types,
+            "timeout": self.timeout_ms,
+        }
 
     @classmethod
     def from_json(cls, values: dict) -> Self:
-        return cls(values["infuse_id"], values["data_types"])
+        return cls(values["infuse_id"], values["data_types"], values["timeout"])
 
 
 class GatewayRequestConnectionRelease(GatewayRequestConnection):
@@ -282,11 +288,13 @@ def receive(self) -> ClientNotification | None:
             return None
         return ClientNotification.from_json(json.loads(data.decode("utf-8")))
 
-    def connection_create(self, infuse_id: int, data_types: GatewayRequestConnectionRequest.DataType) -> int:
+    def connection_create(
+        self, infuse_id: int, data_types: GatewayRequestConnectionRequest.DataType, timeout_ms: int
+    ) -> int:
         self._connection_id = infuse_id
 
         # Send the request for the connection
-        req = GatewayRequestConnectionRequest(infuse_id, data_types)
+        req = GatewayRequestConnectionRequest(infuse_id, data_types, timeout_ms)
         self.send(req)
         # Wait for response from the server
         while True:
@@ -306,9 +314,9 @@ def connection_release(self):
         self._connection_id = None
 
     @contextmanager
-    def connection(self, infuse_id: int, data_types: GatewayRequestConnectionRequest.DataType):
+    def connection(self, infuse_id: int, data_types: GatewayRequestConnectionRequest.DataType, timeout_ms: int = 10000):
         try:
-            yield self.connection_create(infuse_id, data_types)
+            yield self.connection_create(infuse_id, data_types, timeout_ms)
         finally:
             self.connection_release()
 

src/infuse_iot/tools/bt_log.py

@@ -29,18 +29,22 @@ def __init__(self, args):
         self._decoder = TDF()
         self._id = args.id
         self._data = args.data
+        self._conn_timeout = args.conn_timeout
 
     @classmethod
     def add_parser(cls, parser):
         parser.add_argument("--id", type=lambda x: int(x, 0), help="Infuse ID to receive logs for")
         parser.add_argument("--data", action="store_true", help="Subscribe to the data characteristic as well")
+        parser.add_argument(
+            "--conn-timeout", type=int, default=10000, help="Timeout to wait for a connection to the device (ms)"
+        )
 
     def run(self):
         try:
             types = GatewayRequestConnectionRequest.DataType.LOGGING
             if self._data:
                 types |= GatewayRequestConnectionRequest.DataType.DATA
-            with self._client.connection(self._id, types) as _:
+            with self._client.connection(self._id, types, self._conn_timeout) as _:
                 while evt := self._client.receive():
                     if evt is None:
                         continue

src/infuse_iot/tools/gateway.py

@@ -333,7 +333,7 @@ def _handle_conn_request(self, req: GatewayRequestConnectionRequest):
 
         connect_args = defs.bt_connect_infuse.request(
             state.bt_addr.to_rpc_struct(),
-            10000,
+            req.timeout_ms,
             subs,
             0,
         )

src/infuse_iot/tools/native_bt.py

@@ -90,7 +90,7 @@ def notification_handler(self, _characteristic: BleakGATTCharacteristic, data: b
 
     async def create_connection(self, request: GatewayRequestConnectionRequest, dev: BLEDevice, queue: asyncio.Queue):
         Console.log_info(f"{dev}: Initiating connection")
-        async with BleakClient(dev) as client:
+        async with BleakClient(dev, timeout=request.timeout_ms / 1000) as client:
             # Modified from bleak example code
             if client._backend.__class__.__name__ == "BleakClientBlueZDBus":
                 await client._backend._acquire_mtu()  # type: ignore

src/infuse_iot/tools/ota_upgrade.py

@@ -36,6 +36,7 @@ class SubCommand(InfuseCommand):
 
     def __init__(self, args):
         self._client = LocalClient(default_multicast_address(), 1.0)
+        self._conn_timeout = args.conn_timeout
         self._min_rssi: int | None = args.rssi
         self._single_id: int | None = args.id
         self._release: ValidRelease = args.release
@@ -70,6 +71,9 @@ def add_parser(cls, parser):
         parser.add_argument("--rssi", type=int, help="Minimum RSSI to attempt upgrade process")
         parser.add_argument("--id", type=lambda x: int(x, 0), help="Single device to upgrade")
         parser.add_argument("--log", type=str, help="File to write upgrade results to")
+        parser.add_argument(
+            "--conn-timeout", type=int, default=10000, help="Timeout to wait for a connection to the device (ms)"
+        )
 
     def progress_table(self):
         table = Table()
@@ -173,7 +177,7 @@ def run(self):
                 self.state_update(live, f"Connecting to {source.infuse_id:016X}")
                 try:
                     with self._client.connection(
-                        source.infuse_id, GatewayRequestConnectionRequest.DataType.COMMAND
+                        source.infuse_id, GatewayRequestConnectionRequest.DataType.COMMAND, self._conn_timeout
                     ) as mtu:
                         self.state_update(live, f"Uploading patch file to {source.infuse_id:016X}")
                         rpc_client = RpcClient(self._client, mtu, source.infuse_id)

src/infuse_iot/tools/rpc.py

@@ -34,6 +34,9 @@ def add_parser(cls, parser):
         addr_group.add_argument("--gateway", action="store_true", help="Run command on local gateway")
         addr_group.add_argument("--id", type=lambda x: int(x, 0), help="Infuse ID to run command on")
         parser.add_argument("--conn-log", action="store_true", help="Request logs from remote device")
+        parser.add_argument(
+            "--conn-timeout", type=int, default=10000, help="Timeout to wait for a connection to the device (ms)"
+        )
         command_list_parser = parser.add_subparsers(title="commands", metavar="<command>", required=True)
 
         for _, name, _ in pkgutil.walk_packages(wrappers.__path__):
@@ -75,7 +78,7 @@ def run(self):
             types = GatewayRequestConnectionRequest.DataType.COMMAND
             if self._args.conn_log:
                 types |= GatewayRequestConnectionRequest.DataType.LOGGING
-            with self._client.connection(self._id, types) as mtu:
+            with self._client.connection(self._id, types, self._args.conn_timeout) as mtu:
                 self._max_payload = mtu
                 rpc_client = RpcClient(self._client, mtu, self._id, self.rx_handler)
                 params = bytes(self._command.request_struct())