From 4b3a6fea5e53227aaa104996a692cd1fc46cd6df Mon Sep 17 00:00:00 2001
From: David Lechner <david@pybricks.com>
Date: Wed, 10 Mar 2021 10:24:38 -0600
Subject: [PATCH] hub: fix missed notification when connecting

This fixes the first notification being missed due to the notification
handler not being started until after notifications are enabled.

It also fixes the DidConnect action changing the state to Unknown after
a notification was received that set the state to a known value.
---
 src/ble-uart/sagas.ts | 39 +++++++++++++++++++++++++++++----------
 src/hub/reducers.ts   |  2 --
 2 files changed, 29 insertions(+), 12 deletions(-)

diff --git a/src/ble-uart/sagas.ts b/src/ble-uart/sagas.ts
index c703016f..05fd527d 100644
--- a/src/ble-uart/sagas.ts
+++ b/src/ble-uart/sagas.ts
@@ -190,6 +190,16 @@ function* connect(_action: BleDeviceConnectAction): Generator {
             );
     });
 
+    const pybricksControlChannelTask = yield* takeEvery(
+        pybricksControlChannel,
+        handlePybricksControlValueChanged,
+    );
+    const pybricksControlWriteTask = yield* takeEvery(
+        BlePybricksServiceActionType.WriteCommand,
+        writePybricksCommand,
+        pybricksControlChar,
+    );
+
     try {
         // REVISIT: possible Pybricks firmware bug (or chromium bug on Linux)
         // where 'characteristicvaluechanged' is not called after disconnecting
@@ -199,6 +209,8 @@ function* connect(_action: BleDeviceConnectAction): Generator {
         yield* call([pybricksControlChar, 'stopNotifications']);
         yield* call([pybricksControlChar, 'startNotifications']);
     } catch (err) {
+        yield* cancel(pybricksControlChannelTask);
+        yield* cancel(pybricksControlWriteTask);
         pybricksControlChannel.close();
         server.disconnect();
         yield* takeMaybe(disconnectChannel);
@@ -210,6 +222,8 @@ function* connect(_action: BleDeviceConnectAction): Generator {
     try {
         uartService = yield* call([server, 'getPrimaryService'], uartServiceUUID);
     } catch (err) {
+        yield* cancel(pybricksControlChannelTask);
+        yield* cancel(pybricksControlWriteTask);
         pybricksControlChannel.close();
         server.disconnect();
         yield* takeMaybe(disconnectChannel);
@@ -225,6 +239,8 @@ function* connect(_action: BleDeviceConnectAction): Generator {
     try {
         uartRxChar = yield* call([uartService, 'getCharacteristic'], uartRxCharUUID);
     } catch (err) {
+        yield* cancel(pybricksControlChannelTask);
+        yield* cancel(pybricksControlWriteTask);
         pybricksControlChannel.close();
         server.disconnect();
         yield* takeMaybe(disconnectChannel);
@@ -236,6 +252,8 @@ function* connect(_action: BleDeviceConnectAction): Generator {
     try {
         uartTxChar = yield* call([uartService, 'getCharacteristic'], uartTxCharUUID);
     } catch (err) {
+        yield* cancel(pybricksControlChannelTask);
+        yield* cancel(pybricksControlWriteTask);
         pybricksControlChannel.close();
         server.disconnect();
         yield* takeMaybe(disconnectChannel);
@@ -255,6 +273,13 @@ function* connect(_action: BleDeviceConnectAction): Generator {
             uartTxChar.removeEventListener('characteristicvaluechanged', listener);
     });
 
+    const uartTxChannelTask = yield* takeEvery(uartTxChannel, handleUartValueChanged);
+    const uartTxWriteTask = yield* takeEvery(
+        BleUartActionType.Write,
+        writeUart,
+        uartRxChar,
+    );
+
     try {
         // REVISIT: possible Pybricks firmware bug (or chromium bug on Linux)
         // where 'characteristicvaluechanged' is not called after disconnecting
@@ -264,7 +289,11 @@ function* connect(_action: BleDeviceConnectAction): Generator {
         yield* call([uartTxChar, 'stopNotifications']);
         yield* call([uartTxChar, 'startNotifications']);
     } catch (err) {
+        yield* cancel(uartTxWriteTask);
+        yield* cancel(uartTxChannelTask);
         uartTxChannel.close();
+        yield* cancel(pybricksControlChannelTask);
+        yield* cancel(pybricksControlWriteTask);
         pybricksControlChannel.close();
         server.disconnect();
         yield* takeMaybe(disconnectChannel);
@@ -272,16 +301,6 @@ function* connect(_action: BleDeviceConnectAction): Generator {
         return;
     }
 
-    yield* takeEvery(pybricksControlChannel, handlePybricksControlValueChanged);
-    yield* takeEvery(
-        BlePybricksServiceActionType.WriteCommand,
-        writePybricksCommand,
-        pybricksControlChar,
-    );
-
-    yield* takeEvery(uartTxChannel, handleUartValueChanged);
-    yield* takeEvery(BleUartActionType.Write, writeUart, uartRxChar);
-
     yield* put(didConnect());
 
     yield* takeMaybe(disconnectChannel);
diff --git a/src/hub/reducers.ts b/src/hub/reducers.ts
index b0368005..3fe4fb66 100644
--- a/src/hub/reducers.ts
+++ b/src/hub/reducers.ts
@@ -49,8 +49,6 @@ const runtime: Reducer<HubRuntimeState, Action> = (
     switch (action.type) {
         case BleDeviceActionType.DidDisconnect:
             return HubRuntimeState.Disconnected;
-        case BleDeviceActionType.DidConnect:
-            return HubRuntimeState.Unknown;
         case HubMessageActionType.RuntimeStatus:
             switch (action.newStatus) {
                 case HubRuntimeStatusType.Idle: