fix(windows): restore mouse reporting when reattaching (#4080)

2 files changed, 62 insertions, 13 deletions

diff --git a/src/client/terminal_setup.rs b/src/client/terminal_setup.rs
index 9c035346..7441f6db 100644
--- a/src/client/terminal_setup.rs
+++ b/src/client/terminal_setup.rs
@@ -265,16 +265,46 @@ pub(super) fn effective_sgr_pixel_mouse(
     enabled && requested && exact_geometry
 }
 
+#[cfg(any(windows, test))]
+fn set_windows_native_mouse_capture<W: io::Write>(
+    writer: &mut W,
+    enabled: bool,
+    sgr_pixels: bool,
+    set_console_capture: impl FnOnce(bool) -> io::Result<()>,
+) -> io::Result<()> {
+    crate::terminal_modes::clear_host_mouse_reporting(writer)?;
+    set_console_capture(enabled)?;
+    if enabled {
+        crate::terminal_modes::set_windows_mouse_reporting(writer, true, sgr_pixels)?;
+    }
+    Ok(())
+}
+
 pub(super) fn set_mouse_capture(enabled: bool, sgr_pixels: bool) -> io::Result<()> {
-    crate::terminal_modes::clear_host_mouse_reporting(&mut io::stdout())?;
     #[cfg(windows)]
     if is_ssh_session() && windows_vti_input_backend_enabled() {
-        return crate::terminal_modes::set_windows_ssh_mouse_reporting(
+        crate::terminal_modes::clear_host_mouse_reporting(&mut io::stdout())?;
+        return crate::terminal_modes::set_windows_mouse_reporting(
             &mut io::stdout(),
             enabled,
             sgr_pixels,
         );
     }
+    #[cfg(windows)]
+    return set_windows_native_mouse_capture(&mut io::stdout(), enabled, sgr_pixels, |enabled| {
+        if enabled {
+            execute!(io::stdout(), EnableMouseCapture)
+        } else {
+            match execute!(io::stdout(), DisableMouseCapture) {
+                Ok(()) => Ok(()),
+                Err(err) if err.to_string() == "Initial console modes not set" => Ok(()),
+                Err(err) => Err(err),
+            }
+        }
+    });
+    #[cfg(not(windows))]
+    crate::terminal_modes::clear_host_mouse_reporting(&mut io::stdout())?;
+    #[cfg(not(windows))]
     if enabled {
         execute!(io::stdout(), EnableMouseCapture)?;
         if sgr_pixels {
@@ -285,8 +315,6 @@ pub(super) fn set_mouse_capture(enabled: bool, sgr_pixels: bool) -> io::Result<(
     } else {
         match execute!(io::stdout(), DisableMouseCapture) {
             Ok(()) => Ok(()),
-            #[cfg(windows)]
-            Err(err) if err.to_string() == "Initial console modes not set" => Ok(()),
             Err(err) => Err(err),
         }
     }
@@ -443,3 +471,24 @@ impl Drop for TerminalGuard {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn windows_native_mouse_capture_restores_reporting_after_reset() {
+        let mut output = Vec::new();
+
+        set_windows_native_mouse_capture(&mut output, true, false, |enabled| {
+            assert!(enabled);
+            Ok(())
+        })
+        .unwrap();
+
+        assert_eq!(
+            output,
+            b"\x1b[?1006l\x1b[?1016l\x1b[?1015l\x1b[?1005l\x1b[?1003l\x1b[?1002l\x1b[?1000l\x1b[?9l\x1b[?1000h\x1b[?1002h\x1b[?1003h\x1b[?1006h\x1b[?1016l"
+        );
+    }
+}
diff --git a/src/terminal_modes.rs b/src/terminal_modes.rs
index e6b15d36..810fa871 100644
--- a/src/terminal_modes.rs
+++ b/src/terminal_modes.rs
@@ -4,10 +4,10 @@ const DISABLE_HOST_MOUSE_REPORTING_SEQUENCE: &[u8] =
     b"\x1b[?1006l\x1b[?1016l\x1b[?1015l\x1b[?1005l\x1b[?1003l\x1b[?1002l\x1b[?1000l\x1b[?9l";
 
 #[cfg(any(windows, test))]
-const WINDOWS_SSH_MOUSE_REPORTING_ENABLE_SEQUENCE: &[u8] =
+const WINDOWS_MOUSE_REPORTING_ENABLE_SEQUENCE: &[u8] =
     b"\x1b[?1000h\x1b[?1002h\x1b[?1003h\x1b[?1006h";
 #[cfg(any(windows, test))]
-const WINDOWS_SSH_MOUSE_REPORTING_DISABLE_SEQUENCE: &[u8] =
+const WINDOWS_MOUSE_REPORTING_DISABLE_SEQUENCE: &[u8] =
     b"\x1b[?1016l\x1b[?1006l\x1b[?1003l\x1b[?1002l\x1b[?1000l";
 
 pub(crate) fn clear_host_mouse_reporting<W: Write>(writer: &mut W) -> io::Result<()> {
@@ -16,15 +16,15 @@ pub(crate) fn clear_host_mouse_reporting<W: Write>(writer: &mut W) -> io::Result
 }
 
 #[cfg(any(windows, test))]
-pub(crate) fn set_windows_ssh_mouse_reporting<W: Write>(
+pub(crate) fn set_windows_mouse_reporting<W: Write>(
     writer: &mut W,
     enabled: bool,
     sgr_pixels: bool,
 ) -> io::Result<()> {
     writer.write_all(if enabled {
-        WINDOWS_SSH_MOUSE_REPORTING_ENABLE_SEQUENCE
+        WINDOWS_MOUSE_REPORTING_ENABLE_SEQUENCE
     } else {
-        WINDOWS_SSH_MOUSE_REPORTING_DISABLE_SEQUENCE
+        WINDOWS_MOUSE_REPORTING_DISABLE_SEQUENCE
     })?;
     if enabled {
         writer.write_all(if sgr_pixels {
@@ -188,12 +188,12 @@ mod tests {
     }
 
     #[test]
-    fn windows_ssh_mouse_reporting_setup_and_teardown_request_required_modes() {
+    fn windows_mouse_reporting_setup_and_teardown_request_required_modes() {
         let mut output = Vec::new();
 
-        set_windows_ssh_mouse_reporting(&mut output, true, true).unwrap();
-        set_windows_ssh_mouse_reporting(&mut output, true, false).unwrap();
-        set_windows_ssh_mouse_reporting(&mut output, false, false).unwrap();
+        set_windows_mouse_reporting(&mut output, true, true).unwrap();
+        set_windows_mouse_reporting(&mut output, true, false).unwrap();
+        set_windows_mouse_reporting(&mut output, false, false).unwrap();
 
         assert_eq!(
             output,