fix(agent): rearm control capture after device reconnect

Author: Phecda

crates/openlogi-agent-core/src/orchestrator.rs

@@ -11,13 +11,13 @@
 //! (still valid) values — exactly the GUI's "window never opened" behaviour.
 
 use std::collections::{BTreeMap, HashSet};
-use std::sync::atomic::{AtomicI32, Ordering};
+use std::sync::atomic::{AtomicI32, AtomicU64, Ordering};
 use std::sync::{Arc, RwLock};
 
 use openlogi_core::config::{Config, ScrollResolution};
 use openlogi_core::device::{Capabilities, DeviceInventory};
 use openlogi_hid::{CaptureChannel, DeviceRoute};
-use tracing::warn;
+use tracing::{debug, warn};
 
 use crate::DpiCycleState;
 use crate::bindings::{bindings_for, gesture_bindings_for, oshook_gestures_for};
@@ -58,6 +58,10 @@ pub struct SharedRuntime {
     pub dpi_cycle: Arc<RwLock<DpiCycleState>>,
     pub thumbwheel_sensitivity: Arc<AtomicI32>,
     pub capture_channel: CaptureChannel,
+    /// Incremented when the selected device reconnects or the system wakes, so
+    /// the gesture watcher re-arms volatile HID++ control diversion even when
+    /// the receiver route itself never changed.
+    pub capture_rearm_generation: Arc<AtomicU64>,
     /// Exclusive receiver access shared by HID++ capture and pairing. Capture
     /// and pairing must never open the same receiver HID node concurrently.
     pub receiver_access: ReceiverAccess,
@@ -111,6 +115,7 @@ impl Orchestrator {
                 config.app_settings.thumbwheel_sensitivity,
             )),
             capture_channel: Arc::new(RwLock::new(None)),
+            capture_rearm_generation: Arc::new(AtomicU64::new(0)),
             receiver_access: ReceiverAccess::default(),
         };
         let orch = Self {
@@ -205,6 +210,9 @@ impl Orchestrator {
         // (offline→online), or — via the
         // flag — a system wake where none of those are observable.
         let reapply_all = std::mem::take(&mut self.reapply_all_next_refresh);
+        let next_current = pick_current(&devices, self.config.selected_device());
+        let rearm_capture =
+            selected_needs_capture_rearm(&self.devices, &devices, next_current, reapply_all);
         let followup = std::mem::take(&mut self.reapply_followup);
         let (targets, next_followup) =
             plan_reapply(&self.devices, &devices, &followup, reapply_all);
@@ -218,15 +226,23 @@ impl Orchestrator {
                     || a.route != b.route
                     || a.capabilities != b.capabilities
             });
-        if !changed {
+        if changed {
+            self.devices = devices;
+            self.current = next_current;
+            self.rebuild();
+        } else {
             // Same set and routes — but keep the fresh `online` flags, or a
             // device that woke this tick would read as a transition forever.
             self.devices = devices;
-            return;
         }
-        self.devices = devices;
-        self.current = pick_current(&self.devices, self.config.selected_device());
-        self.rebuild();
+        if rearm_capture {
+            let generation = self
+                .shared
+                .capture_rearm_generation
+                .fetch_add(1, Ordering::Relaxed)
+                .wrapping_add(1);
+            debug!(generation, "selected device requires capture re-arm");
+        }
     }
 
     /// Force a volatile-settings re-apply for every online device on the next
@@ -465,6 +481,18 @@ fn reapply_targets(prev: &[AgentDevice], next: &[AgentDevice], reapply_all: bool
         .collect()
 }
 
+/// Whether this refresh invalidated the selected device's volatile control
+/// diversion. Receiver routes stay connected while a paired mouse sleeps, so
+/// route equality alone cannot tell the capture watcher to re-arm on wake.
+fn selected_needs_capture_rearm(
+    prev: &[AgentDevice],
+    next: &[AgentDevice],
+    selected: usize,
+    reapply_all: bool,
+) -> bool {
+    reapply_targets(prev, next, reapply_all).contains(&selected)
+}
+
 /// Plan this refresh's volatile-settings writes: the [`reapply_targets`] set
 /// plus one confirming re-apply for devices first sighted last refresh, and
 /// the follow-up keys to confirm next refresh.
@@ -519,7 +547,7 @@ fn write_value<T>(lock: &RwLock<T>, value: T, name: &str) {
 mod tests {
     use super::{
         AgentDevice, InventoryHealth, Orchestrator, configured_wheel_mode, plan_reapply,
-        reapply_targets,
+        reapply_targets, selected_needs_capture_rearm,
     };
     use openlogi_core::config::{Config, ScrollResolution};
     use openlogi_core::device::Capabilities;
@@ -633,6 +661,29 @@ mod tests {
         assert_eq!(reapply_targets(&prev, &next, true), vec![0]);
     }
 
+    #[test]
+    fn selected_receiver_reconnect_requests_capture_rearm() {
+        let prev = [dev("selected", 1, false), dev("other", 2, true)];
+        let next = [dev("selected", 1, true), dev("other", 2, true)];
+
+        assert!(selected_needs_capture_rearm(&prev, &next, 0, false));
+        assert!(!selected_needs_capture_rearm(&prev, &next, 1, false));
+    }
+
+    #[test]
+    fn system_wake_requests_capture_rearm_for_selected_online_device() {
+        let devices = [dev("selected", 1, true), dev("other", 2, true)];
+
+        assert!(selected_needs_capture_rearm(&devices, &devices, 0, true));
+    }
+
+    #[test]
+    fn steady_inventory_does_not_cycle_capture() {
+        let devices = [dev("selected", 1, true)];
+
+        assert!(!selected_needs_capture_rearm(&devices, &devices, 0, false));
+    }
+
     #[test]
     fn plan_reapply_confirms_a_first_sighting_once() {
         use std::collections::HashSet;

crates/openlogi-agent-core/src/watchers/gesture.rs

@@ -19,14 +19,14 @@
 //! way regardless.
 
 use std::collections::BTreeMap;
-use std::sync::atomic::{AtomicI32, Ordering};
+use std::sync::atomic::{AtomicI32, AtomicU64, Ordering};
 use std::sync::{Arc, RwLock};
 use std::thread;
 use std::time::{Duration, Instant};
 
 use openlogi_core::binding::{Action, ButtonId, GestureDirection, default_binding};
 use openlogi_core::config::DEFAULT_THUMBWHEEL_SENSITIVITY;
-use openlogi_hid::{CaptureChannel, CapturedInput, DeviceRoute, run_capture_session};
+use openlogi_hid::{CaptureChannel, CaptureStop, CapturedInput, DeviceRoute, run_capture_session};
 use tokio::sync::{mpsc, oneshot};
 use tracing::{debug, warn};
 
@@ -80,6 +80,7 @@ pub fn spawn(
     dpi_cycle: Arc<RwLock<DpiCycleState>>,
     capture_channel: CaptureChannel,
     thumbwheel_sensitivity: ThumbwheelSensitivity,
+    capture_rearm_generation: Arc<AtomicU64>,
     receiver_access: ReceiverAccess,
 ) {
     thread::spawn(move || {
@@ -99,6 +100,7 @@ pub fn spawn(
             dpi_cycle,
             capture_channel,
             thumbwheel_sensitivity,
+            capture_rearm_generation,
             receiver_access,
         ));
     });
@@ -141,6 +143,27 @@ fn should_rearm(done_epoch: u64, live_epoch: u64, has_target: bool) -> bool {
     done_epoch == live_epoch && has_target
 }
 
+#[derive(Debug, Clone, PartialEq, Eq)]
+struct CaptureTarget {
+    route: DeviceRoute,
+    capture_thumbwheel: bool,
+    divert_gesture_button: bool,
+    rearm_generation: u64,
+}
+
+/// A generation-only restart on the same route follows a device reconnect or
+/// system wake. Its old firmware state is already gone, so restoring it would
+/// only delay (or permanently block) the replacement session.
+fn stop_for_transition(current: &CaptureTarget, next: Option<&CaptureTarget>) -> CaptureStop {
+    if next.is_some_and(|next| {
+        next.route == current.route && next.rearm_generation != current.rearm_generation
+    }) {
+        CaptureStop::Abandon
+    } else {
+        CaptureStop::Restore
+    }
+}
+
 /// Keep one capture session alive for the active device, restarting it when the
 /// device or the thumb-wheel arming changes, and dispatch incoming inputs. Runs
 /// for the lifetime of the process.
@@ -150,12 +173,12 @@ async fn manage(
     dpi_cycle: Arc<RwLock<DpiCycleState>>,
     capture_channel: CaptureChannel,
     thumbwheel_sensitivity: ThumbwheelSensitivity,
+    capture_rearm_generation: Arc<AtomicU64>,
     receiver_access: ReceiverAccess,
 ) {
     let (tx, mut rx) = mpsc::unbounded_channel::<CapturedInput>();
-    // (route, capture_thumbwheel, divert_gesture_button)
-    let mut current: Option<(DeviceRoute, bool, bool)> = None;
-    let mut stop: Option<oneshot::Sender<()>> = None;
+    let mut current: Option<CaptureTarget> = None;
+    let mut stop: Option<oneshot::Sender<CaptureStop>> = None;
     let mut ticker = tokio::time::interval(TARGET_POLL);
     let mut accumulators = WheelAccumulators::default();
     // Capture sessions run as detached tasks, so an unexpected exit (a transient
@@ -197,33 +220,39 @@ async fn manage(
                     // thread the full config in. Re-evaluated each tick, so a
                     // ReloadConfig owner change restarts the session accordingly.
                     let divert_gesture = gesture_bindings.read().is_ok_and(|g| !g.is_empty());
-                    target.map(|t| {
-                        (
-                            t,
-                            thumbwheel_armed(&hook_maps, sensitivity),
-                            divert_gesture,
-                        )
+                    let rearm_generation = capture_rearm_generation.load(Ordering::Relaxed);
+                    target.map(|route| CaptureTarget {
+                        route,
+                        capture_thumbwheel: thumbwheel_armed(&hook_maps, sensitivity),
+                        divert_gesture_button: divert_gesture,
+                        rearm_generation,
                     })
                 };
                 if want == current {
                     continue;
                 }
+                debug!(?current, ?want, "capture target state changed");
                 // Target or thumb-wheel arming changed (or first tick): stop the
                 // old session and start one for the new state. Sending on the
                 // oneshot lets the old session restore the diverted controls.
                 if let Some(stop) = stop.take() {
-                    let _ = stop.send(());
+                    let reason = current
+                        .as_ref()
+                        .map_or(CaptureStop::Restore, |current| {
+                            stop_for_transition(current, want.as_ref())
+                        });
+                    let _ = stop.send(reason);
                 }
                 if current.is_some() {
                     current = None;
                     continue;
                 }
-                if let Some((route, capture_thumbwheel, divert_gesture_button)) = want {
+                if let Some(target) = want {
                     let Some(receiver_lease) = receiver_access.try_acquire_for_capture() else {
                         current = None;
                         continue;
                     };
-                    current = Some((route.clone(), capture_thumbwheel, divert_gesture_button));
+                    current = Some(target.clone());
                     let (stop_tx, stop_rx) = oneshot::channel();
                     let sink = tx.clone();
                     let slot = Arc::clone(&capture_channel);
@@ -233,9 +262,9 @@ async fn manage(
                     tokio::spawn(async move {
                         let _receiver_lease = receiver_lease;
                         if let Err(e) = run_capture_session(
-                            route,
-                            capture_thumbwheel,
-                            divert_gesture_button,
+                            target.route,
+                            target.capture_thumbwheel,
+                            target.divert_gesture_button,
                             sink,
                             stop_rx,
                             slot,
@@ -445,8 +474,58 @@ fn advance(
 
 #[cfg(test)]
 mod tests {
+    use openlogi_hid::{CaptureStop, DeviceRoute};
+
     use super::*;
 
+    fn capture_target(route: DeviceRoute, generation: u64) -> CaptureTarget {
+        CaptureTarget {
+            route,
+            capture_thumbwheel: false,
+            divert_gesture_button: true,
+            rearm_generation: generation,
+        }
+    }
+
+    #[test]
+    fn reconnect_restart_abandons_stale_state_on_the_same_route() {
+        let route = DeviceRoute::Unifying {
+            receiver_uid: "receiver".to_string(),
+            slot: 1,
+        };
+        let current = capture_target(route.clone(), 3);
+        let next = capture_target(route, 4);
+
+        assert_eq!(
+            stop_for_transition(&current, Some(&next)),
+            CaptureStop::Abandon
+        );
+    }
+
+    #[test]
+    fn ordinary_target_changes_restore_old_controls() {
+        let current = capture_target(
+            DeviceRoute::Unifying {
+                receiver_uid: "receiver".to_string(),
+                slot: 1,
+            },
+            3,
+        );
+        let next = capture_target(
+            DeviceRoute::Direct {
+                vendor_id: 0x046d,
+                product_id: 0xb023,
+            },
+            4,
+        );
+
+        assert_eq!(
+            stop_for_transition(&current, Some(&next)),
+            CaptureStop::Restore
+        );
+        assert_eq!(stop_for_transition(&current, None), CaptureStop::Restore);
+    }
+
     #[test]
     fn multiplier_is_unity_at_default_sensitivity() {
         assert!((scroll_multiplier(DEFAULT_THUMBWHEEL_SENSITIVITY) - 1.0).abs() < f32::EPSILON);

crates/openlogi-agent/src/main.rs

@@ -161,6 +161,7 @@ async fn run(config: Config) {
         shared.dpi_cycle.clone(),
         shared.capture_channel.clone(),
         shared.thumbwheel_sensitivity.clone(),
+        shared.capture_rearm_generation.clone(),
         shared.receiver_access.clone(),
     );
 

crates/openlogi-agent/src/pairing.rs

@@ -270,6 +270,7 @@ mod tests {
             dpi_cycle: Arc::new(RwLock::new(DpiCycleState::default())),
             thumbwheel_sensitivity: Arc::new(0.into()),
             capture_channel: Arc::new(RwLock::new(None)),
+            capture_rearm_generation: Arc::new(0.into()),
             receiver_access: ReceiverAccess::default(),
         }
     }