From 199a5d0bab5dddfcd625d56593101a46a2c39916 Mon Sep 17 00:00:00 2001
From: Jan Komarek <jan.komarek@satoshilabs.com>
Date: Tue, 30 Apr 2024 14:04:35 +0200
Subject: [PATCH] fix(suite): show correct step for disconnect event

---
 .../suite/src/components/firmware/ReconnectDevicePrompt.tsx | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/packages/suite/src/components/firmware/ReconnectDevicePrompt.tsx b/packages/suite/src/components/firmware/ReconnectDevicePrompt.tsx
index 15f0953da01..0b51e8c06e9 100644
--- a/packages/suite/src/components/firmware/ReconnectDevicePrompt.tsx
+++ b/packages/suite/src/components/firmware/ReconnectDevicePrompt.tsx
@@ -4,7 +4,7 @@ import * as semver from 'semver';
 
 import { pickByDeviceModel, getFirmwareVersion } from '@trezor/device-utils';
 import { H2, Button, ConfirmOnDevice, variables, DeviceAnimation } from '@trezor/components';
-import { DEVICE, Device, DeviceModelInternal } from '@trezor/connect';
+import { DEVICE, Device, DeviceModelInternal, UI } from '@trezor/connect';
 import { Modal, Translation, WebUsbButton } from 'src/components/suite';
 import { DeviceConfirmImage } from 'src/components/suite/DeviceConfirmImage';
 import { useDevice, useFirmware } from 'src/hooks/suite';
@@ -183,8 +183,10 @@ export const ReconnectDevicePrompt = ({ onClose, onSuccess }: ReconnectDevicePro
         ) {
             return 'done';
         }
+        const rebootToBootloaderNotSupported = uiEvent?.type === UI.FIRMWARE_DISCONNECT;
+        const rebootToBootloaderCancelled = device?.connected && device?.mode !== 'bootloader';
 
-        return device?.connected && device?.mode !== 'bootloader'
+        return rebootToBootloaderNotSupported || rebootToBootloaderCancelled
             ? 'waiting-for-reboot'
             : 'disconnected';
     };