fix(agent): hand off capture across receiver routes

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

@@ -32,27 +32,27 @@ pub(super) enum CompletionAction {
     Remove { unexpected: bool },
 }
 
-enum SessionPhase {
-    Active(oneshot::Sender<()>),
+enum SessionPhase<Stop> {
+    Active(oneshot::Sender<Stop>),
     Draining,
 }
 
 /// One capture epoch, including its hardware identity, dispatch state and
 /// acknowledged teardown phase.
-pub(super) struct CaptureSession<Target, Dispatch> {
+pub(super) struct CaptureSession<Target, Dispatch, Stop = ()> {
     id: HidppSessionId,
     target: Target,
     dispatch: Dispatch,
-    phase: SessionPhase,
+    phase: SessionPhase<Stop>,
 }
 
-impl<Target, Dispatch> CaptureSession<Target, Dispatch> {
+impl<Target, Dispatch, Stop> CaptureSession<Target, Dispatch, Stop> {
     /// Begin tracking an active capture task.
     pub(super) fn active(
         id: HidppSessionId,
         target: Target,
         dispatch: Dispatch,
-        stop: oneshot::Sender<()>,
+        stop: oneshot::Sender<Stop>,
     ) -> Self {
         Self {
             id,
@@ -108,11 +108,15 @@ impl<Target, Dispatch> CaptureSession<Target, Dispatch> {
     }
 }
 
-impl<Target: PartialEq, Dispatch: Clone + PartialEq> CaptureSession<Target, Dispatch> {
+impl<Target: PartialEq, Dispatch: Clone + PartialEq, Stop> CaptureSession<Target, Dispatch, Stop> {
     /// Reconcile against the latest wanted target and dispatch state. A target
     /// change begins teardown exactly once; dispatch-only changes hot-refresh
     /// the plan while preserving the hardware epoch.
-    pub(super) fn reconcile(&mut self, wanted: Option<(&Target, &Dispatch)>) -> ReconcileAction {
+    pub(super) fn reconcile_with(
+        &mut self,
+        wanted: Option<(&Target, &Dispatch)>,
+        stop_for_change: impl FnOnce(&Target, Option<&Target>) -> Stop,
+    ) -> ReconcileAction {
         if !self.is_active() {
             return ReconcileAction::None;
         }
@@ -125,15 +129,23 @@ impl<Target: PartialEq, Dispatch: Clone + PartialEq> CaptureSession<Target, Disp
             self.dispatch.clone_from(dispatch);
             return ReconcileAction::DispatchChanged;
         }
+        let stop_command = stop_for_change(&self.target, wanted.map(|(target, _)| target));
         let SessionPhase::Active(stop) = std::mem::replace(&mut self.phase, SessionPhase::Draining)
         else {
             return ReconcileAction::None;
         };
-        let _ = stop.send(());
+        let _ = stop.send(stop_command);
         ReconcileAction::Retiring
     }
 }
 
+impl<Target: PartialEq, Dispatch: Clone + PartialEq> CaptureSession<Target, Dispatch> {
+    /// Reconcile a session whose teardown command carries no additional intent.
+    pub(super) fn reconcile(&mut self, wanted: Option<(&Target, &Dispatch)>) -> ReconcileAction {
+        self.reconcile_with(wanted, |_, _| ())
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

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

@@ -30,8 +30,8 @@ use std::time::Duration;
 use openlogi_core::device_order::PhysicalDeviceKey;
 use openlogi_core::scroll::ScrollDelta;
 use openlogi_hid::{
-    CaptureChannel, CaptureSessionOutcome, CapturedInput, DeviceIoGate, PendingCaptureRestore,
-    run_capture_session_with_registry_spec,
+    CaptureChannel, CaptureSessionOutcome, CaptureSessionStop, CapturedInput, DeviceIoGate,
+    PendingCaptureRestore, run_capture_session_with_registry_spec,
 };
 use tokio::sync::{mpsc, oneshot, watch};
 use tokio::time::Instant;
@@ -120,7 +120,7 @@ pub fn spawn(
     });
 }
 
-type RunningSession = CaptureSession<CaptureTarget, DispatchPlan>;
+type RunningSession = CaptureSession<CaptureTarget, DispatchPlan, CaptureSessionStop>;
 
 struct CapturedEvent {
     physical_key: PhysicalDeviceKey,
@@ -223,13 +223,26 @@ fn reconcile_session(
     wanted: Option<(&CaptureTarget, &DispatchPlan)>,
     dispatcher: &mut InputDispatcher,
 ) {
-    if session.reconcile(wanted) == ReconcileAction::DispatchChanged {
+    if session.reconcile_with(wanted, stop_for_target_change) == ReconcileAction::DispatchChanged {
         dispatcher.cancel_session(session.id());
         let config_key = session.dispatch().config_key.clone();
         session.rekey(&config_key);
     }
 }
 
+fn stop_for_target_change(
+    current: &CaptureTarget,
+    wanted: Option<&CaptureTarget>,
+) -> CaptureSessionStop {
+    wanted.map_or(CaptureSessionStop::Shutdown, |next| {
+        if next.route == current.route {
+            CaptureSessionStop::Shutdown
+        } else {
+            CaptureSessionStop::Handoff(next.route.clone())
+        }
+    })
+}
+
 /// Reconcile one tracked slot directly against the latest publication. Input
 /// calls this before dispatch so an event cannot slip between publishing a hot
 /// action update and processing its notification.
@@ -285,6 +298,7 @@ fn acquire_session_lease(
 async fn retry_pending_restores(
     pending_restores: &mut HashMap<PhysicalDeviceKey, PendingRestore>,
     registry: &openlogi_hid::ChannelRegistry,
+    wanted: &[DeviceCapturePlan],
     now: Instant,
 ) {
     let keys: Vec<_> = pending_restores
@@ -296,7 +310,16 @@ async fn retry_pending_restores(
         let Some(pending) = pending_restores.remove(&key) else {
             continue;
         };
-        if let CaptureSessionOutcome::RestorePending(token) = pending.token.retry(registry).await {
+        let outcome = match wanted.iter().find(|plan| plan.target.physical_key == key) {
+            Some(plan) => {
+                pending
+                    .token
+                    .retry_via(plan.target.route.clone(), registry)
+                    .await
+            }
+            None => pending.token.retry(registry).await,
+        };
+        if let CaptureSessionOutcome::RestorePending(token) = outcome {
             pending_restores.insert(
                 key,
                 PendingRestore {
@@ -401,7 +424,8 @@ impl GestureManagerState {
             None
         };
         if restore_lease.is_some() {
-            retry_pending_restores(&mut self.pending_restores, &channels.registry, now).await;
+            retry_pending_restores(&mut self.pending_restores, &channels.registry, wanted, now)
+                .await;
         }
 
         for plan in wanted {

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

@@ -41,7 +41,10 @@ fn live_session_with_epoch(epoch: u64) -> RunningSession {
 
 fn draining_session_with_epoch(epoch: u64) -> RunningSession {
     let mut session = live_session_with_epoch(epoch);
-    assert_eq!(session.reconcile(None), ReconcileAction::Retiring);
+    assert_eq!(
+        session.reconcile_with(None, stop_for_target_change),
+        ReconcileAction::Retiring
+    );
     session
 }
 
@@ -215,7 +218,10 @@ async fn exclusive_request_retires_capture_without_rejecting_owned_input() {
     assert!(wanted_sessions(*requests.borrow(), &plans).is_empty());
 
     let mut session = live_session_with_epoch(7);
-    assert_eq!(session.reconcile(None), ReconcileAction::Retiring);
+    assert_eq!(
+        session.reconcile_with(None, stop_for_target_change),
+        ReconcileAction::Retiring
+    );
     assert!(!session.is_active());
     assert!(
         dispatch_context_for(&session_id(7), Some(&session)).is_some(),
@@ -249,7 +255,10 @@ fn an_active_session_refreshes_bindings_without_rearming_hardware() {
     assert_eq!(session.target(), &new_plan.target);
 
     assert_eq!(
-        session.reconcile(Some((&new_plan.target, &new_plan.dispatch))),
+        session.reconcile_with(
+            Some((&new_plan.target, &new_plan.dispatch)),
+            stop_for_target_change,
+        ),
         ReconcileAction::DispatchChanged,
         "a hot plan refresh must cancel input lifecycles admitted under the old action map"
     );
@@ -288,10 +297,13 @@ fn side_gesture_transition_keeps_the_retiring_plan_until_native_restore() {
         .clear();
     assert_ne!(session.target(), &published_without_hook.target);
     assert_eq!(
-        session.reconcile(Some((
-            &published_without_hook.target,
-            &published_without_hook.dispatch,
-        ))),
+        session.reconcile_with(
+            Some((
+                &published_without_hook.target,
+                &published_without_hook.dispatch,
+            )),
+            stop_for_target_change,
+        ),
         ReconcileAction::Retiring
     );
     assert!(!session.is_active());
@@ -336,6 +348,35 @@ fn capture_target_changes_schedule_the_old_session_for_retirement() {
     );
 }
 
+#[test]
+fn receiver_route_change_requests_a_firmware_handoff() {
+    let old_plan = plan();
+    let mut new_plan = old_plan.clone();
+    let successor_route = DeviceRoute::Bolt {
+        receiver_uid: "receiver-b".to_owned(),
+        slot: 2,
+    };
+    new_plan.target.route.clone_from(&successor_route);
+    let (stop, mut stopped) = oneshot::channel();
+    let mut session =
+        CaptureSession::active(session_id(7), old_plan.target, old_plan.dispatch, stop);
+
+    assert_eq!(
+        session.reconcile_with(
+            Some((&new_plan.target, &new_plan.dispatch)),
+            stop_for_target_change,
+        ),
+        ReconcileAction::Retiring
+    );
+    assert_eq!(
+        stopped
+            .try_recv()
+            .expect("route change should stop the active session"),
+        CaptureSessionStop::Handoff(successor_route),
+        "teardown must restore through the receiver the mouse moved to"
+    );
+}
+
 #[test]
 fn config_key_adoption_hot_refreshes_the_same_physical_capture_slot() {
     let old_plan = plan();
@@ -353,7 +394,10 @@ fn config_key_adoption_hot_refreshes_the_same_physical_capture_slot() {
         .get(&physical_key)
         .map(|plan| (&plan.target, &plan.dispatch));
 
-    assert_eq!(running.reconcile(desired), ReconcileAction::DispatchChanged);
+    assert_eq!(
+        running.reconcile_with(desired, stop_for_target_change),
+        ReconcileAction::DispatchChanged
+    );
     running.rekey(&wanted[&physical_key].dispatch.config_key);
     assert!(running.is_active());
     assert_eq!(running.id().device_key(), "unit:00000001");
@@ -398,7 +442,10 @@ fn active_session_adopts_action_only_plan_changes_without_rearming() {
     );
     assert_eq!(first.target, rebound.target);
     assert_eq!(
-        session.reconcile(Some((&rebound.target, &rebound.dispatch))),
+        session.reconcile_with(
+            Some((&rebound.target, &rebound.dispatch)),
+            stop_for_target_change,
+        ),
         ReconcileAction::DispatchChanged
     );
     assert_eq!(
@@ -439,7 +486,10 @@ fn active_session_adopts_gesture_and_per_app_dispatch_changes() {
     );
     assert_eq!(first.target, gestured.target);
     assert_eq!(
-        session.reconcile(Some((&gestured.target, &gestured.dispatch))),
+        session.reconcile_with(
+            Some((&gestured.target, &gestured.dispatch)),
+            stop_for_target_change,
+        ),
         ReconcileAction::DispatchChanged
     );
     assert_eq!(
@@ -483,7 +533,10 @@ fn active_session_adopts_gesture_and_per_app_dispatch_changes() {
     );
     assert_eq!(base.target, per_app.target);
     assert_eq!(
-        session.reconcile(Some((&per_app.target, &per_app.dispatch))),
+        session.reconcile_with(
+            Some((&per_app.target, &per_app.dispatch)),
+            stop_for_target_change,
+        ),
         ReconcileAction::DispatchChanged
     );
     assert_eq!(
@@ -530,7 +583,10 @@ fn wheel_configuration_changes_refresh_without_rearming_hardware() {
         "both custom bindings require the same HID++ diversion"
     );
     assert_eq!(
-        session.reconcile(Some((&rebound.target, &rebound.dispatch))),
+        session.reconcile_with(
+            Some((&rebound.target, &rebound.dispatch)),
+            stop_for_target_change,
+        ),
         ReconcileAction::DispatchChanged,
         "dispatch-only binding changes must not cycle firmware diversion"
     );
@@ -548,7 +604,10 @@ fn wheel_configuration_changes_refresh_without_rearming_hardware() {
     );
     assert_eq!(rebound.target, rescaled.target);
     assert_eq!(
-        session.reconcile(Some((&rescaled.target, &rescaled.dispatch))),
+        session.reconcile_with(
+            Some((&rescaled.target, &rescaled.dispatch)),
+            stop_for_target_change,
+        ),
         ReconcileAction::DispatchChanged,
         "an already-diverted wheel needs a state reset, not a hardware restart"
     );

crates/openlogi-device/src/lib.rs

@@ -51,8 +51,9 @@ pub use pairing::{
     PasskeyMethod, ReceiverFamily, ReceiverSelector, list_pairing_receivers, run_pairing, unpair,
 };
 pub use session::gesture::{
-    CaptureChannel, CaptureSessionFailure, CaptureSessionOutcome, CapturedInput, GestureError,
-    PendingCaptureRestore, run_capture_session, run_capture_session_with_registry_spec,
+    CaptureChannel, CaptureSessionFailure, CaptureSessionOutcome, CaptureSessionStop,
+    CapturedInput, GestureError, PendingCaptureRestore, run_capture_session,
+    run_capture_session_with_registry_spec,
 };
 pub use session::host_switch::{
     HostSwitchError, HostSwitchStopReason, run_host_switch_session, switch_linked_hosts,

crates/openlogi-device/src/session/capture_restore.rs

@@ -57,12 +57,15 @@ impl ReprogRestore {
     }
 }
 
-#[derive(Clone, Copy)]
+#[derive(Debug, Clone, PartialEq, Eq)]
 pub(crate) enum CaptureStop {
     /// The owner deliberately requested teardown.
     Shutdown,
     /// Inventory removed or replaced the channel that armed capture.
     ChannelChanged,
+    /// The same physical device moved to another route. Restore through the
+    /// successor route instead of timing out against the retired one.
+    Handoff(DeviceRoute),
 }
 
 /// Whether a retry may use the transport on which capture originally ran.
@@ -177,6 +180,30 @@ impl PendingCaptureRestore {
     /// every awaited restore write. A concurrent replacement returns this
     /// token as pending so the new winner is restored on the next attempt.
     pub async fn retry(self, registry: &ChannelRegistry) -> CaptureSessionOutcome {
+        self.retry_current_route(registry).await
+    }
+
+    /// Retry through a newly elected route to the same physical device.
+    ///
+    /// The caller owns physical identity and must only supply a route resolved
+    /// for the device whose firmware this token owns. The route becomes the
+    /// fallback for later retries, so repeated host switches can keep moving
+    /// restoration toward the device's latest live transport.
+    pub async fn retry_via(
+        mut self,
+        route: DeviceRoute,
+        registry: &ChannelRegistry,
+    ) -> CaptureSessionOutcome {
+        self.route = route;
+        // Physical identity elected this route. If a rapid switch returns to
+        // the still-current channel that originally armed capture, restoring
+        // there is now safe; ordinary channel-replacement retries retain the
+        // stricter replacement-only policy.
+        self.retired_policy = RetiredChannelPolicy::CurrentAllowed;
+        self.retry_current_route(registry).await
+    }
+
+    async fn retry_current_route(self, registry: &ChannelRegistry) -> CaptureSessionOutcome {
         let Some(current) = registry.lookup(&self.route) else {
             return CaptureSessionOutcome::RestorePending(self);
         };
@@ -278,6 +305,15 @@ pub(crate) async fn restore_after_stop(
             Some(registry) => pending.retry(registry).await,
             None => CaptureSessionOutcome::RestorePending(pending),
         },
+        CaptureStop::Handoff(route) => {
+            if let Some(registry) = registry {
+                pending.retry_via(route, registry).await
+            } else {
+                let mut pending = pending;
+                pending.route = route;
+                CaptureSessionOutcome::RestorePending(pending)
+            }
+        }
     }
 }