Recover Vial unlock sessions after reconnect

i18n/en.toml

@@ -259,6 +259,9 @@ return_to_layout_hint = "Right-click or Esc to return to layout"
 
 [unlock]
 highlighted_keys_hint = "Press and hold the highlighted keys"
+start_hint = "Start unlocking, then hold the highlighted keys"
+start = "Start unlocking"
+cancel = "Cancel"
 
 [universal_symbols_setup]
 open_privacy_settings = "Open Privacy Settings"
@@ -405,6 +408,7 @@ refresh_device_data_pending_write = "Finish the pending layer write before refre
 refresh_device_data_missing_device = "Failed to refresh device data: device not found"
 refresh_device_data_missing_info = "Failed to refresh device data: device information unavailable"
 refresh_device_data_delete_failed = "Failed to remove cached device data"
+unlock_recovery_detected = "Active device unlock recovered — hold the highlighted keys"
 
 [dynamic_status]
 device_not_found = "Device not found"

i18n/ru.toml

@@ -259,6 +259,9 @@ return_to_layout_hint = "ПКМ или Esc — вернуться к раскл
 
 [unlock]
 highlighted_keys_hint = "Нажмите и удерживайте подсвеченные клавиши"
+start_hint = "Начните разблокировку, затем удерживайте подсвеченные клавиши"
+start = "Начать разблокировку"
+cancel = "Отмена"
 
 [universal_symbols_setup]
 open_privacy_settings = "Открыть настройки приватности"
@@ -405,6 +408,7 @@ refresh_device_data_pending_write = "Завершите запись слоя п
 refresh_device_data_missing_device = "Не удалось обновить данные устройства: устройство не найдено"
 refresh_device_data_missing_info = "Не удалось обновить данные устройства: информация недоступна"
 refresh_device_data_delete_failed = "Не удалось удалить сохранённые данные устройства"
+unlock_recovery_detected = "Активная разблокировка восстановлена — удерживайте подсвеченные клавиши"
 
 [dynamic_status]
 device_not_found = "Устройство не найдено"

src/app_init.rs

@@ -193,6 +193,7 @@ impl EntropyApp {
             vial_unlock_total: 50,
             vial_unlock_last_poll: None,
             vial_unlock_animation_nonce: 0,
+            vial_unlock_reconnect_after_completion: false,
             #[cfg(not(target_arch = "wasm32"))]
             connect_state: ConnectState::Idle,
             #[cfg(not(target_arch = "wasm32"))]
@@ -230,7 +231,7 @@ impl EntropyApp {
                 .hid_device
                 .as_ref()
                 .and_then(|hid| hid.get_unlock_status().ok())
-                .map(|(unlocked, _)| unlocked)
+                .map(|status| status.unlocked)
                 .map(|unlocked| !unlocked)
                 .unwrap_or(false)
     }
@@ -264,32 +265,28 @@ impl EntropyApp {
 
     #[cfg(not(target_arch = "wasm32"))]
     pub(crate) fn cancel_vial_unlock(&mut self, suppress_macro_auto_unlock: bool) {
-        if let Some(hid) = &self.hid_device {
-            match hid.lock() {
-                Ok(()) => {
-                    self.status_msg = crate::i18n::tr_catalog(
-                        self.app_settings.language,
-                        "status_messages.device_unlock_cancelled",
-                    )
-                    .into();
-                }
-                Err(e) => {
-                    self.status_msg = crate::i18n::tr_catalog_format(
-                        self.app_settings.language,
-                        "status_messages.cancel_unlock_failed",
-                        &[("error", &e.to_string())],
-                    );
-                }
-            }
+        if self.vial_unlock_polling {
+            log::warn!("Ignored Vial unlock cancellation after UNLOCK_START");
+            self.status_msg = crate::i18n::tr_catalog(
+                self.app_settings.language,
+                "status_messages.finish_unlock_before_closing",
+            )
+            .into();
+            return;
         }
+        log::info!("Vial unlock prompt dismissed before UNLOCK_START");
+        self.status_msg = crate::i18n::tr_catalog(
+            self.app_settings.language,
+            "status_messages.device_unlock_cancelled",
+        )
+        .into();
         self.unlock_open = false;
         self.vial_unlock_polling = false;
         self.vial_unlock_last_poll = None;
+        self.vial_unlock_reconnect_after_completion = false;
         self.pending_layout_indicator_open_after_unlock = false;
         self.vial_unlock_counter = 0;
         self.vial_unlock_best = 50;
-        self.matrix_tester_unlock_prompted = false;
-        self.matrix_tester_lock_checked = false;
         if suppress_macro_auto_unlock {
             self.macro_auto_unlock_cancelled = true;
         }

src/app_state.rs

@@ -369,10 +369,26 @@ pub(crate) struct ConnectResult {
     pub(crate) supported_qmk_settings: Vec<u16>,
 }
 
+/// Minimal device state available while Vial firmware blocks normal VIA commands.
+#[cfg(not(target_arch = "wasm32"))]
+pub(crate) struct UnlockRecoveryResult {
+    pub(crate) device_name: String,
+    pub(crate) keyboard_id: u64,
+    pub(crate) hid_device: Option<crate::hid::HidDevice>,
+    pub(crate) layout: KeyboardLayout,
+    pub(crate) unlock_keys: Vec<(u8, u8)>,
+}
+
+#[cfg(not(target_arch = "wasm32"))]
+pub(crate) enum ConnectOutcome {
+    Connected(Box<ConnectResult>),
+    UnlockRecovery(Box<UnlockRecoveryResult>),
+}
+
 #[cfg(not(target_arch = "wasm32"))]
 pub(crate) enum ConnectTaskMessage {
     Progress(String),
-    Done(Box<Result<ConnectResult, String>>),
+    Done(Box<Result<ConnectOutcome, String>>),
 }
 
 #[cfg(not(target_arch = "wasm32"))]
@@ -2980,6 +2996,8 @@ pub struct EntropyApp {
     pub(crate) vial_unlock_total: u8,
     pub(crate) vial_unlock_last_poll: Option<std::time::Instant>,
     pub(crate) vial_unlock_animation_nonce: u64,
+    /// Reload full device state after finishing an unlock recovered during connect.
+    pub(crate) vial_unlock_reconnect_after_completion: bool,
 }
 
 #[cfg(test)]

src/hid.rs

@@ -742,10 +742,7 @@ fn response_matches_command(command: &[u8], resp: &[u8; MSG_LEN]) -> bool {
     };
 
     match cmd {
-        CMD_VIA_GET_PROTOCOL_VERSION => {
-            resp[0] == CMD_VIA_GET_PROTOCOL_VERSION
-                && matches!(u16::from_be_bytes([resp[1], resp[2]]), 9 | 0xFFFF)
-        }
+        CMD_VIA_GET_PROTOCOL_VERSION => response_matches_via_protocol_version(command, resp),
         CMD_VIA_GET_LAYER_COUNT => {
             resp[0] == CMD_VIA_GET_LAYER_COUNT && (1..=32).contains(&resp[1])
         }
@@ -770,6 +767,16 @@ fn response_matches_command(command: &[u8], resp: &[u8; MSG_LEN]) -> bool {
     }
 }
 
+#[cfg(not(target_arch = "wasm32"))]
+pub(crate) fn response_matches_via_protocol_version(command: &[u8], resp: &[u8]) -> bool {
+    command.first() == Some(&CMD_VIA_GET_PROTOCOL_VERSION)
+        && resp.first() == Some(&CMD_VIA_GET_PROTOCOL_VERSION)
+        && resp
+            .get(1..3)
+            .map(|version| matches!(u16::from_be_bytes([version[0], version[1]]), 0 | 9 | 0xFFFF))
+            .unwrap_or(false)
+}
+
 #[cfg(not(target_arch = "wasm32"))]
 fn response_matches_vial_command(command: &[u8], resp: &[u8; MSG_LEN]) -> bool {
     let Some(&subcommand) = command.get(1) else {
@@ -1065,6 +1072,22 @@ mod tests {
         assert!(response_matches_command(&command, &response));
     }
 
+    #[test]
+    fn protocol_version_matcher_accepts_active_unlock_zero_and_rejects_unrelated_reports() {
+        let command = [CMD_VIA_GET_PROTOCOL_VERSION];
+        let active_unlock_response = [CMD_VIA_GET_PROTOCOL_VERSION, 0, 0];
+        let unrelated_response = [CMD_VIA_GET_LAYER_COUNT, 0, 0];
+
+        assert!(response_matches_via_protocol_version(
+            &command,
+            &active_unlock_response
+        ));
+        assert!(!response_matches_via_protocol_version(
+            &command,
+            &unrelated_response
+        ));
+    }
+
     #[test]
     fn qmk_settings_set_rejects_stale_get_response() {
         let mut command = qmk_settings_command(CMD_VIAL_QMK_SETTINGS_SET, 300);

src/hid_parse.rs

@@ -154,8 +154,16 @@ pub(crate) fn parse_vialrgb_supported_effects_payload(
     batch_max
 }
 
-pub(crate) fn parse_unlock_status_response(resp: &[u8]) -> (bool, Vec<(u8, u8)>) {
+#[derive(Clone, Debug, Eq, PartialEq)]
+pub(crate) struct VialUnlockStatus {
+    pub(crate) unlocked: bool,
+    pub(crate) in_progress: bool,
+    pub(crate) keys: Vec<(u8, u8)>,
+}
+
+pub(crate) fn parse_unlock_status_response(resp: &[u8]) -> VialUnlockStatus {
     let unlocked = resp.first().copied() == Some(1);
+    let in_progress = resp.get(1).copied() == Some(1);
     let mut keys = Vec::new();
     let mut i = 2;
     while i + 1 < resp.len() {
@@ -167,13 +175,34 @@ pub(crate) fn parse_unlock_status_response(resp: &[u8]) -> (bool, Vec<(u8, u8)>)
         keys.push((row, col));
         i += 2;
     }
-    (unlocked, keys)
+    VialUnlockStatus {
+        unlocked,
+        in_progress,
+        keys,
+    }
 }
 
 #[cfg(test)]
 mod tests {
     use super::*;
 
+    #[test]
+    fn parses_active_unlock_status_and_keys() {
+        let mut resp = [0xFFu8; 32];
+        resp[0] = 0;
+        resp[1] = 1;
+        resp[2..6].copy_from_slice(&[1, 2, 3, 4]);
+
+        assert_eq!(
+            parse_unlock_status_response(&resp),
+            VialUnlockStatus {
+                unlocked: false,
+                in_progress: true,
+                keys: vec![(1, 2), (3, 4)],
+            }
+        );
+    }
+
     #[test]
     fn parses_and_pads_macro_buffer() {
         let parsed = parse_macro_buffer(b"one\0two", 4);
@@ -300,7 +329,11 @@ mod tests {
         let resp = [1, 0, 3, 4, 5, 6, 0xff, 0xff, 7, 8];
         assert_eq!(
             parse_unlock_status_response(&resp),
-            (true, vec![(3, 4), (5, 6)])
+            VialUnlockStatus {
+                unlocked: true,
+                in_progress: false,
+                keys: vec![(3, 4), (5, 6)],
+            }
         );
     }
 }

src/hid_vial.rs

@@ -1,4 +1,4 @@
-use super::hid_parse::parse_unlock_status_response;
+use super::hid_parse::{parse_unlock_status_response, VialUnlockStatus};
 use super::hid_protocol::*;
 use super::HidDevice;
 use anyhow::{bail, Context, Result};
@@ -81,8 +81,7 @@ impl HidDevice {
     }
 
     /// Check if keyboard is unlocked
-    /// Returns (unlocked, unlock_keys: Vec<(row,col)>)
-    pub fn get_unlock_status(&self) -> Result<(bool, Vec<(u8, u8)>)> {
+    pub fn get_unlock_status(&self) -> Result<VialUnlockStatus> {
         let resp = self
             .usb_send(&[CMD_VIA_VIAL_PREFIX, CMD_VIAL_GET_UNLOCK_STATUS])
             .context("failed to read Vial unlock status")?;

src/ui/device_connect_apply.rs

@@ -341,7 +341,49 @@ impl EntropyApp {
         self.connect_state = ConnectState::Idle;
 
         match result {
-            Ok(r) => {
+            Ok(ConnectOutcome::UnlockRecovery(r)) => {
+                let layer_count = r.layout.layers.len().max(1);
+                log::warn!(
+                    "Applying Vial unlock recovery for {} (keyboard id {:016X}, {} unlock keys)",
+                    r.device_name,
+                    r.keyboard_id,
+                    r.unlock_keys.len()
+                );
+                self.pending_tap_hold_numeric_writes.clear();
+                self.tap_hold_numeric_write_due = None;
+                self.layer_count = layer_count;
+                self.firmware = r.layout.firmware;
+                self.current_device_name = r.device_name.clone();
+                self.current_keyboard_id = Some(r.keyboard_id);
+                self.current_encoder_visibility_id =
+                    encoder_visibility_id(&r.device_name, r.keyboard_id);
+                self.layout = Some(r.layout);
+                self.hid_device = r.hid_device;
+                self.vial_unlock_keys = r.unlock_keys;
+                self.vial_unlock_polling = true;
+                self.vial_unlock_counter = 1;
+                self.vial_unlock_best = 1;
+                self.vial_unlock_total = 1;
+                self.vial_unlock_last_poll = Some(std::time::Instant::now());
+                self.vial_unlock_animation_nonce = self.vial_unlock_animation_nonce.wrapping_add(1);
+                self.vial_unlock_reconnect_after_completion = true;
+                self.unlock_open = true;
+                self.macro_auto_unlock_cancelled = false;
+                self.status_msg = crate::i18n::tr_catalog(
+                    self.app_settings.language,
+                    "status_messages.unlock_recovery_detected",
+                )
+                .into();
+                if let Some(dev) = self
+                    .selected_device
+                    .and_then(|idx| self.device_manager.devices().get(idx))
+                {
+                    self.device_display_names
+                        .insert(dev.display_name_cache_key(), r.device_name);
+                }
+                ctx.request_repaint();
+            }
+            Ok(ConnectOutcome::Connected(r)) => {
                 self.pending_tap_hold_numeric_writes.clear();
                 self.tap_hold_numeric_write_due = None;
                 log::info!(