fix(agent): rearm control capture after device reconnect

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

@@ -11,7 +11,7 @@
 //! (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::binding::Action;
@@ -23,7 +23,7 @@ use openlogi_hid::{
     CaptureChannel, ChannelPool, ChannelRegistry, DIRECT_DEVICE_INDEX, DeviceRoute,
     KEYBOARD_KEY_CIDS,
 };
-use tracing::{info, warn};
+use tracing::{debug, info, warn};
 
 use crate::DpiCycleState;
 use crate::bindings::{bindings_for, gesture_bindings_for, oshook_gestures_for};
@@ -84,6 +84,10 @@ pub struct SharedRuntime {
     /// The keyboard capture session's open channel, reused by Fn-lock writes
     /// (the mouse-oriented [`Self::capture_channel`] points elsewhere).
     pub keyboard_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>,
     /// Receiver access shared by HID++ sessions and pairing. Pairing/host
     /// transitions are exclusive; capture sessions share under read leases.
     pub receiver_access: ReceiverAccess,
@@ -153,6 +157,7 @@ impl Orchestrator {
             channel_pool: ChannelPool::default(),
             keyboard_spec: Arc::new(RwLock::new(None)),
             keyboard_channel: Arc::new(RwLock::new(None)),
+            capture_rearm_generation: Arc::new(AtomicU64::new(0)),
             receiver_access: ReceiverAccess::default(),
             host_switch_links: Arc::new(RwLock::new(Vec::new())),
         };
@@ -339,6 +344,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);
@@ -353,7 +361,11 @@ impl Orchestrator {
                     || a.capabilities != b.capabilities
                     || a.light_capabilities != b.light_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;
@@ -363,11 +375,15 @@ impl Orchestrator {
                 host_switch_links(&self.config, &self.devices),
                 "host_switch_links",
             );
-            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
@@ -828,6 +844,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.
@@ -892,6 +920,7 @@ mod tests {
     use super::{
         AgentDevice, InventoryHealth, Orchestrator, build_devices, configured_wheel_mode,
         host_switch_links, pick_current, plan_reapply, reapply_targets,
+        selected_needs_capture_rearm,
     };
     use openlogi_core::config::{Config, LightSettings, ScrollResolution};
     use openlogi_core::device::{
@@ -1189,6 +1218,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,7 +19,7 @@
 //! 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};
@@ -83,6 +83,7 @@ pub fn spawn(
     dpi_cycle: Arc<RwLock<DpiCycleState>>,
     capture_channel: CaptureChannel,
     thumbwheel_sensitivity: ThumbwheelSensitivity,
+    capture_rearm_generation: Arc<AtomicU64>,
     receiver_access: ReceiverAccess,
 ) {
     spawn_inner(
@@ -91,18 +92,24 @@ pub fn spawn(
         dpi_cycle,
         capture_channel,
         thumbwheel_sensitivity,
+        capture_rearm_generation,
         receiver_access,
         None,
     );
 }
 
 /// Spawn capture in Agent mode, reusing only inventory-published channels.
+#[expect(
+    clippy::too_many_arguments,
+    reason = "agent mode needs maps, dpi, rearm, leases, and registry"
+)]
 pub fn spawn_with_registry(
     hook_maps: SharedHookMaps,
     gesture_bindings: GestureBindings,
     dpi_cycle: Arc<RwLock<DpiCycleState>>,
     capture_channel: CaptureChannel,
     thumbwheel_sensitivity: ThumbwheelSensitivity,
+    capture_rearm_generation: Arc<AtomicU64>,
     receiver_access: ReceiverAccess,
     registry: ChannelRegistry,
 ) {
@@ -112,17 +119,23 @@ pub fn spawn_with_registry(
         dpi_cycle,
         capture_channel,
         thumbwheel_sensitivity,
+        capture_rearm_generation,
         receiver_access,
         Some(registry),
     );
 }
 
+#[expect(
+    clippy::too_many_arguments,
+    reason = "capture manager needs maps, dpi, rearm, leases, and registry"
+)]
 fn spawn_inner(
     hook_maps: SharedHookMaps,
     gesture_bindings: GestureBindings,
     dpi_cycle: Arc<RwLock<DpiCycleState>>,
     capture_channel: CaptureChannel,
     thumbwheel_sensitivity: ThumbwheelSensitivity,
+    capture_rearm_generation: Arc<AtomicU64>,
     receiver_access: ReceiverAccess,
     registry: Option<ChannelRegistry>,
 ) {
@@ -143,6 +156,7 @@ fn spawn_inner(
             dpi_cycle,
             capture_channel,
             thumbwheel_sensitivity,
+            capture_rearm_generation,
             receiver_access,
             registry,
         ));
@@ -186,6 +200,37 @@ pub(crate) fn should_rearm(done_epoch: u64, live_epoch: u64, has_target: 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,
+}
+
+/// Why to stop the current session when the desired target changes.
+fn stop_for_target_change(
+    want: Option<&CaptureTarget>,
+    current: Option<&CaptureTarget>,
+    connection_is_current: bool,
+) -> Option<CaptureStop> {
+    let cur = current?;
+    if want == Some(cur) {
+        return (!connection_is_current).then_some(CaptureStop::Revoked);
+    }
+    // Same route, new generation (reconnect/wake): skip restore — firmware already reset.
+    if want.is_some_and(|next| {
+        next.route == cur.route && next.rearm_generation != cur.rearm_generation
+    }) {
+        return Some(CaptureStop::Revoked);
+    }
+    Some(CaptureStop::Graceful)
+}
+
+#[cfg_attr(
+    not(test),
+    allow(dead_code, reason = "kept for unit tests of stop policy")
+)]
 fn stop_reason<T: PartialEq>(
     want: Option<&T>,
     current: Option<&T>,
@@ -277,18 +322,22 @@ fn spawn_capture_session(launch: CaptureLaunch) -> oneshot::Sender<CaptureStop>
 /// 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.
+#[expect(
+    clippy::too_many_arguments,
+    reason = "capture manager needs maps, dpi, rearm, leases, and registry"
+)]
 async fn manage(
     hook_maps: SharedHookMaps,
     gesture_bindings: GestureBindings,
     dpi_cycle: Arc<RwLock<DpiCycleState>>,
     capture_channel: CaptureChannel,
     thumbwheel_sensitivity: ThumbwheelSensitivity,
+    capture_rearm_generation: Arc<AtomicU64>,
     receiver_access: ReceiverAccess,
     registry: Option<ChannelRegistry>,
 ) {
     let (tx, mut rx) = mpsc::unbounded_channel::<CapturedInput>();
-    // (route, capture_thumbwheel, divert_gesture_button)
-    let mut current: Option<(DeviceRoute, bool, bool)> = None;
+    let mut current: Option<CaptureTarget> = None;
     let mut stop: Option<oneshot::Sender<CaptureStop>> = None;
     let mut stopping_epoch: Option<u64> = None;
     let mut ticker = tokio::time::interval(TARGET_POLL);
@@ -336,12 +385,12 @@ 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 stopping_epoch.is_some() {
@@ -350,7 +399,7 @@ async fn manage(
                 let connection_is_current =
                     capture_connection_is_current(registry.as_ref(), &capture_channel);
                 if let Some(reason) =
-                    stop_reason(want.as_ref(), current.as_ref(), connection_is_current)
+                    stop_for_target_change(want.as_ref(), current.as_ref(), connection_is_current)
                 {
                     if let Some(stop) = stop.take() {
                         let _ = stop.send(reason);
@@ -362,17 +411,17 @@ async fn manage(
                 if want == current {
                     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_session() else {
                         current = None;
                         continue;
                     };
-                    current = Some((route.clone(), capture_thumbwheel, divert_gesture_button));
+                    current = Some(target.clone());
                     epoch = epoch.wrapping_add(1);
                     stop = Some(spawn_capture_session(CaptureLaunch {
-                        route,
-                        capture_thumbwheel,
-                        divert_gesture_button,
+                        route: target.route,
+                        capture_thumbwheel: target.capture_thumbwheel,
+                        divert_gesture_button: target.divert_gesture_button,
                         sink: tx.clone(),
                         channel_slot: Arc::clone(&capture_channel),
                         receiver_lease,

crates/openlogi-agent/src/main.rs

@@ -130,6 +130,7 @@ fn spawn_hidpp_watchers(shared: &SharedRuntime) {
         shared.dpi_cycle.clone(),
         shared.capture_channel.clone(),
         shared.thumbwheel_sensitivity.clone(),
+        shared.capture_rearm_generation.clone(),
         shared.receiver_access.clone(),
         shared.channel_registry.clone(),
     );

crates/openlogi-agent/src/pairing.rs

@@ -277,6 +277,7 @@ mod tests {
             channel_pool: openlogi_hid::ChannelPool::default(),
             keyboard_spec: Arc::new(RwLock::new(None)),
             keyboard_channel: Arc::new(RwLock::new(None)),
+            capture_rearm_generation: Arc::new(0.into()),
             receiver_access: ReceiverAccess::default(),
             host_switch_links: Arc::new(RwLock::new(Vec::new())),
         }