diff --git a/bumble/smp.py b/bumble/smp.py
index bf4b257..f22f320 100644
--- a/bumble/smp.py
+++ b/bumble/smp.py
@@ -946,7 +946,9 @@ class Session:
             self.tk = self.passkey.to_bytes(16, byteorder='little')
             logger.debug(f'TK from passkey = {self.tk.hex()}')
 
-        await self.pairing_config.delegate.display_number(self.passkey, digits=6)
+        self.connection.cancel_on_disconnection(
+            self.pairing_config.delegate.display_number(self.passkey, digits=6)
+        )
 
     def input_passkey(self, next_steps: Optional[Callable[[], None]] = None) -> None:
         # Prompt the user for the passkey displayed on the peer
@@ -1569,11 +1571,12 @@ class Session:
         if self.pairing_method == PairingMethod.CTKD_OVER_CLASSIC:
             # Authentication is already done in SMP, so remote shall start keys distribution immediately
             return
-        elif self.sc:
+
+        if self.sc:
+            self.send_public_key_command()
+
             if self.pairing_method == PairingMethod.PASSKEY:
                 self.display_or_input_passkey()
-
-            self.send_public_key_command()
         else:
             if self.pairing_method == PairingMethod.PASSKEY:
                 self.display_or_input_passkey(self.send_pairing_confirm_command)
@@ -1846,10 +1849,10 @@ class Session:
             elif self.pairing_method == PairingMethod.PASSKEY:
                 self.send_pairing_confirm_command()
         else:
+            # Send our public key back to the initiator
+            self.send_public_key_command()
 
             def next_steps() -> None:
-                # Send our public key back to the initiator
-                self.send_public_key_command()
 
                 if self.pairing_method in (
                     PairingMethod.JUST_WORKS,
diff --git a/tests/self_test.py b/tests/self_test.py
index 9098299..4467299 100644
--- a/tests/self_test.py
+++ b/tests/self_test.py
@@ -322,7 +322,9 @@ async def test_self_smp(io_caps, sc, mitm, key_dist):
                 return 0
             else:
                 if (
-                    self.peer_delegate.io_capability
+                    self.io_capability
+                    == PairingDelegate.IoCapability.KEYBOARD_INPUT_ONLY
+                    and self.peer_delegate.io_capability
                     == PairingDelegate.IoCapability.KEYBOARD_INPUT_ONLY
                 ):
                     peer_number = 6789