diff --git a/ch-13/1-modelling-space/computer/display_arena_from_robot.py b/ch-13/1-modelling-space/computer/display_arena_from_robot.py
index 9b7a6eb..8c36ceb 100644
--- a/ch-13/1-modelling-space/computer/display_arena_from_robot.py
+++ b/ch-13/1-modelling-space/computer/display_arena_from_robot.py
@@ -9,7 +9,7 @@ from robot_ble_connection import BleConnection
 class RobotDisplay:
     def __init__(self):
         self.fig, self.ax = plt.subplots()
-        self.ble_connection = BleConnection(self.handle_data, self.connected)
+        self.ble_connection = BleConnection(self.handle_data)
 
         self.arena = None
         self.display_closed = False
@@ -24,7 +24,7 @@ class RobotDisplay:
         line_part = data.decode("utf-8")
         self.line += line_part
         if not self.line.endswith("\n"):
-          return
+            return
         print(f"Received data: {self.line}")
         data = json.loads(self.line)
         self.line = ""
@@ -35,29 +35,32 @@ class RobotDisplay:
         self.arena = arena
         self.ax.clear()
         for line in arena["arena"]:
-            self.ax.plot([line[0][0], line[1][0]],
-                         [line[0][1], line[1][1]],
-                         color='black')
+            self.ax.plot(
+                [line[0][0], line[1][0]], [line[0][1], line[1][1]], color="black"
+            )
         for line in arena["target_zone"]:
-            self.ax.plot([line[0][0], line[1][0]],
-                         [line[0][1], line[1][1]],
-                         color="red")
+            self.ax.plot(
+                [line[0][0], line[1][0]], [line[0][1], line[1][1]], color="red"
+            )
 
-    def connected(self):
-        request = json.dumps({"command": "arena"}).encode()
-        print(f"Sending request for arena: {request}")
-        self.ble_connection.send_uart_data(request)
-
-    async def main(self): 
+    async def main(self):
         plt.ion()
 
-        asyncio.create_task(self.ble_connection.connect())
-        while not self.display_closed:
-            plt.pause(0.05)
-            plt.draw()
-            await asyncio.sleep(0.01)
-        
-        plt.show()
+        try:
+            await self.ble_connection.connect()
+            request = json.dumps({"type": "arena"}).encode()
+            print(f"Sending request for arena: {request}")
+            self.ble_connection.send_uart_data(request)
+
+            while not self.display_closed:
+                plt.pause(0.05)
+                plt.draw()
+                await asyncio.sleep(0.01)
+
+            plt.show()
+        finally:
+            await self.ble_connection.close()
+
 
 robot_display = RobotDisplay()
 asyncio.run(robot_display.main())
diff --git a/ch-13/1-modelling-space/computer/robot_ble_connection.py b/ch-13/1-modelling-space/computer/robot_ble_connection.py
index 61674b4..9ea514c 100644
--- a/ch-13/1-modelling-space/computer/robot_ble_connection.py
+++ b/ch-13/1-modelling-space/computer/robot_ble_connection.py
@@ -11,11 +11,9 @@ class BleConnection:
     adafruit_tx_uuid = "6E400002-B5A3-F393-E0A9-E50E24DCCA9E"
     ble_name = "Adafruit Bluefruit LE"
 
-    def __init__(self, receive_handler: typing.Callable[[bytes], None], 
-            connected_handler: typing.Callable[[], None]):
+    def __init__(self, receive_handler: typing.Callable[[bytes], None]):
         self.ble_client : bleak.BleakClient = None
         self.receive_handler = receive_handler
-        self.connected_handler = connected_handler
 
     def _uart_handler(self, _, data: bytes):
         self.receive_handler(data)
@@ -27,16 +25,16 @@ class BleConnection:
         print([device.name for device in devices])        
         ble_device_info = [device for device in devices if device.name==self.ble_name][0]
         print("Connecting to {}...".format(ble_device_info.name))
-        async with bleak.BleakClient(ble_device_info.address) as ble_client:
-            self.ble_client = ble_client
-            self.connected_handler()
-            print("Connected to {}".format(ble_device_info.name))
-            asyncio.create_task(
-                self.ble_client.start_notify(self.adafruit_rx_uuid, self._uart_handler)
-            )
-            while True:
-                await asyncio.sleep(1)
+        self.ble_client = bleak.BleakClient(ble_device_info.address)
+        await self.ble_client.connect()
+        print("Connected to {}".format(ble_device_info.name))
+        self.notify_task = asyncio.create_task(
+            self.ble_client.start_notify(self.adafruit_rx_uuid, self._uart_handler)
+        )
 
+    async def close(self):
+        await self.ble_client.disconnect()
+                
     def send_uart_data(self, data: bytes):
         if self.ble_client:
             asyncio.create_task(