fix: reassemble split sgr mouse input

refs #1382

2 files changed, 176 insertions, 17 deletions

diff --git a/src/client/input.rs b/src/client/input.rs
index bc7faf35..a49c4316 100644
--- a/src/client/input.rs
+++ b/src/client/input.rs
@@ -133,8 +133,10 @@ fn idle_flush_timeout_ms(
     framer: &crate::raw_input::RawInputByteFramer,
     host_mouse_capture_active: bool,
 ) -> i32 {
-    if host_mouse_capture_active && framer.has_pending_lone_escape() {
-        crate::raw_input::MOUSE_ACTIVE_LONE_ESCAPE_FLUSH_TIMEOUT_MS
+    if host_mouse_capture_active
+        && (framer.has_pending_lone_escape() || framer.has_pending_incomplete_sgr_mouse_sequence())
+    {
+        crate::raw_input::MOUSE_ACTIVE_ESCAPE_SEQUENCE_FLUSH_TIMEOUT_MS
     } else {
         crate::raw_input::RAW_INPUT_IDLE_FLUSH_TIMEOUT_MS
     }
@@ -402,20 +404,33 @@ mod tests {
     }
 
     #[test]
-    fn mouse_active_lone_escape_gets_longer_reassembly_window() {
-        let mut framer = crate::raw_input::RawInputByteFramer::default();
-        assert!(framer.push(b"\x1b").is_empty());
-
+    fn mouse_active_escape_sequences_get_longer_reassembly_window() {
+        let mut escape = crate::raw_input::RawInputByteFramer::default();
+        assert!(escape.push(b"\x1b").is_empty());
+        let mut mouse = crate::raw_input::RawInputByteFramer::default();
+        assert!(mouse.push(b"\x1b[<3").is_empty());
+        let mut unrelated = crate::raw_input::RawInputByteFramer::default();
+        assert!(unrelated.push(b"\x1b[49:33;2:").is_empty());
+
+        for framer in [&escape, &mouse, &unrelated] {
+            assert_eq!(
+                idle_flush_timeout_ms(framer, false),
+                crate::raw_input::RAW_INPUT_IDLE_FLUSH_TIMEOUT_MS
+            );
+        }
+        for framer in [&escape, &mouse] {
+            assert_eq!(
+                idle_flush_timeout_ms(framer, true),
+                crate::raw_input::MOUSE_ACTIVE_ESCAPE_SEQUENCE_FLUSH_TIMEOUT_MS
+            );
+        }
         assert_eq!(
-            idle_flush_timeout_ms(&framer, false),
+            idle_flush_timeout_ms(&unrelated, true),
             crate::raw_input::RAW_INPUT_IDLE_FLUSH_TIMEOUT_MS
         );
-        assert_eq!(
-            idle_flush_timeout_ms(&framer, true),
-            crate::raw_input::MOUSE_ACTIVE_LONE_ESCAPE_FLUSH_TIMEOUT_MS
-        );
+
         let mouse_timeout_ms =
-            std::hint::black_box(crate::raw_input::MOUSE_ACTIVE_LONE_ESCAPE_FLUSH_TIMEOUT_MS);
+            std::hint::black_box(crate::raw_input::MOUSE_ACTIVE_ESCAPE_SEQUENCE_FLUSH_TIMEOUT_MS);
         assert!(mouse_timeout_ms > 100);
     }
 }
diff --git a/src/raw_input.rs b/src/raw_input.rs
index 3bf8cae9..445b4be3 100644
--- a/src/raw_input.rs
+++ b/src/raw_input.rs
@@ -101,8 +101,7 @@ use crate::terminal_theme::{
 
 const ESC: u8 = 0x1b;
 pub(crate) const RAW_INPUT_IDLE_FLUSH_TIMEOUT_MS: i32 = 10;
-#[cfg(any(not(windows), test))]
-pub(crate) const MOUSE_ACTIVE_LONE_ESCAPE_FLUSH_TIMEOUT_MS: i32 = 150;
+pub(crate) const MOUSE_ACTIVE_ESCAPE_SEQUENCE_FLUSH_TIMEOUT_MS: i32 = 150;
 pub(crate) const GHOSTTY_COLOR_SCHEME_DARK_REPORT: &[u8] = b"\x1b[?997;1n";
 pub(crate) const GHOSTTY_COLOR_SCHEME_LIGHT_REPORT: &[u8] = b"\x1b[?997;2n";
 const BRACKETED_PASTE_START: &[u8] = b"\x1b[200~";
@@ -145,6 +144,10 @@ impl RawInputFramer {
         self.byte_framer.has_pending_input()
     }
 
+    pub(crate) fn has_pending_incomplete_sgr_mouse_sequence(&self) -> bool {
+        self.byte_framer.has_pending_incomplete_sgr_mouse_sequence()
+    }
+
     #[cfg(any(windows, test))]
     pub(crate) fn has_pending_bracketed_paste(&self) -> bool {
         self.byte_framer.has_pending_bracketed_paste()
@@ -213,6 +216,10 @@ impl RawInputByteFramer {
         self.buffer.as_slice() == [ESC]
     }
 
+    pub(crate) fn has_pending_incomplete_sgr_mouse_sequence(&self) -> bool {
+        starts_with_incomplete_sgr_mouse_sequence(&self.buffer)
+    }
+
     #[cfg(any(windows, test))]
     pub(crate) fn has_pending_bracketed_paste(&self) -> bool {
         self.buffer.starts_with(BRACKETED_PASTE_START)
@@ -223,6 +230,13 @@ impl RawInputByteFramer {
         let mut chunks = self.drain_available_chunks();
 
         if let Some(family) = self.discard_until {
+            if family == ControlStringFamily::OrphanedSgrMouseTail {
+                self.buffer.clear();
+                self.discard_until = None;
+                self.discarded_tail_bytes = 0;
+                return chunks;
+            }
+
             let keep_split_st = self.buffer.last() == Some(&ESC);
             let keep_discarding = plausible_control_string_tail(family, &self.buffer);
             self.discarded_tail_bytes = self.discarded_tail_bytes.saturating_add(self.buffer.len());
@@ -256,6 +270,18 @@ impl RawInputByteFramer {
             return chunks;
         }
 
+        if starts_with_incomplete_sgr_mouse_sequence(&self.buffer) {
+            tracing::debug!(
+                bytes = ?self.buffer,
+                "discarding incomplete SGR mouse sequence after input timeout"
+            );
+            self.discarded_tail_bytes = self.buffer.len();
+            self.discard_until = (self.discarded_tail_bytes <= MAX_DISCARDED_CONTROL_TAIL_BYTES)
+                .then_some(ControlStringFamily::OrphanedSgrMouseTail);
+            self.buffer.clear();
+            return chunks;
+        }
+
         if self.buffer.starts_with(BRACKETED_PASTE_START)
             && find_subsequence(&self.buffer, BRACKETED_PASTE_END).is_none()
         {
@@ -356,6 +382,18 @@ impl RawInputByteFramer {
             }
 
             if let Some(family) = self.discard_until {
+                if family == ControlStringFamily::OrphanedSgrMouseTail {
+                    if discard_orphaned_sgr_mouse_tail(
+                        &mut self.buffer,
+                        &mut self.discarded_tail_bytes,
+                    ) {
+                        self.discard_until = None;
+                        self.discarded_tail_bytes = 0;
+                        continue;
+                    }
+                    break;
+                }
+
                 let Some(terminator_len) =
                     control_string_terminator_for_family(&self.buffer, family)
                 else {
@@ -438,6 +476,14 @@ pub(crate) fn events_require_host_terminal_theme_query(events: &[RawInputEvent])
         .any(|event| matches!(event, RawInputEvent::HostColorSchemeChanged(_)))
 }
 
+fn input_flush_timeout_ms(framer: &RawInputFramer) -> i32 {
+    if framer.has_pending_incomplete_sgr_mouse_sequence() {
+        MOUSE_ACTIVE_ESCAPE_SEQUENCE_FLUSH_TIMEOUT_MS
+    } else {
+        RAW_INPUT_IDLE_FLUSH_TIMEOUT_MS
+    }
+}
+
 pub fn spawn_input_reader() -> mpsc::Receiver<RawInputEvent> {
     let (tx, rx) = mpsc::channel(256);
 
@@ -455,7 +501,7 @@ pub fn spawn_input_reader() -> mpsc::Receiver<RawInputEvent> {
                 Ok(n) => {
                     send_raw_input_events(framer.push(&scratch[..n]), &tx);
 
-                    if stdin_read_ready(&reader, RAW_INPUT_IDLE_FLUSH_TIMEOUT_MS) == Some(false) {
+                    if stdin_read_ready(&reader, input_flush_timeout_ms(&framer)) == Some(false) {
                         let had_pending = framer.has_pending_input();
                         let events = framer.flush_timeout();
                         let held_escape = had_pending && events.is_empty();
@@ -753,6 +799,13 @@ fn complete_escape_sequence_len(buffer: &[u8]) -> Option<usize> {
     Some(1 + escaped_char_width)
 }
 
+fn starts_with_incomplete_sgr_mouse_sequence(buffer: &[u8]) -> bool {
+    buffer.starts_with(b"\x1b[<")
+        && buffer[3..]
+            .iter()
+            .all(|byte| byte.is_ascii_digit() || *byte == b';')
+}
+
 fn starts_with_incomplete_orphaned_sgr_mouse_tail(buffer: &[u8]) -> bool {
     if buffer.len() > MAX_ORPHANED_SGR_MOUSE_TAIL_BYTES {
         return false;
@@ -792,10 +845,39 @@ fn discard_or_buffer_orphaned_sgr_mouse_tail(
     discarded_tail_bytes: &mut usize,
 ) {
     if !discard_complete_orphaned_sgr_mouse_tail(buffer) {
-        *discard_until = Some(ControlStringFamily::OrphanedSgrMouseTail);
-        *discarded_tail_bytes = 0;
+        *discarded_tail_bytes = buffer.len();
+        *discard_until = (*discarded_tail_bytes <= MAX_DISCARDED_CONTROL_TAIL_BYTES)
+            .then_some(ControlStringFamily::OrphanedSgrMouseTail);
+        buffer.clear();
+    }
+}
+
+fn discard_orphaned_sgr_mouse_tail(buffer: &mut Vec<u8>, discarded_tail_bytes: &mut usize) -> bool {
+    let remaining = MAX_DISCARDED_CONTROL_TAIL_BYTES.saturating_sub(*discarded_tail_bytes);
+    let inspected = buffer.len().min(remaining);
+
+    for index in 0..inspected {
+        match buffer[index] {
+            b'0'..=b'9' | b';' => {}
+            b'M' | b'm' => {
+                buffer.drain(..=index);
+                return true;
+            }
+            _ => {
+                buffer.drain(..index);
+                return true;
+            }
+        }
+    }
+
+    *discarded_tail_bytes = discarded_tail_bytes.saturating_add(inspected);
+    if buffer.len() > inspected {
         buffer.clear();
+        return true;
     }
+
+    buffer.clear();
+    false
 }
 
 fn osc_string_terminator(buffer: &[u8]) -> Option<usize> {
@@ -1491,6 +1573,68 @@ mod tests {
     }
 
     #[test]
+    fn legacy_reader_extends_only_incomplete_sgr_mouse_timeout() {
+        let mut mouse = RawInputFramer::default();
+        assert!(mouse.push(b"\x1b[<3").is_empty());
+        assert_eq!(
+            input_flush_timeout_ms(&mouse),
+            MOUSE_ACTIVE_ESCAPE_SEQUENCE_FLUSH_TIMEOUT_MS
+        );
+
+        let mut escape = RawInputFramer::default();
+        assert!(escape.push(b"\x1b").is_empty());
+        assert_eq!(
+            input_flush_timeout_ms(&escape),
+            RAW_INPUT_IDLE_FLUSH_TIMEOUT_MS
+        );
+    }
+
+    #[test]
+    fn sgr_mouse_sequence_split_after_button_prefix_is_reassembled_before_timeout() {
+        let mut framer = RawInputFramer::default();
+
+        assert!(framer.push(b"\x1b[<3").is_empty());
+        let events = framer.push(b"5;58;30M");
+
+        assert_eq!(events.len(), 1);
+        assert!(matches!(
+            events[0],
+            RawInputEvent::Mouse(MouseEvent {
+                kind: MouseEventKind::Moved,
+                column: 57,
+                row: 29,
+                ..
+            })
+        ));
+    }
+
+    #[test]
+    fn timed_out_split_sgr_mouse_tail_is_discarded_and_following_input_is_preserved() {
+        let mut framer = RawInputFramer::default();
+
+        assert!(framer.push(b"\x1b[<3").is_empty());
+        assert!(framer.flush_timeout().is_empty());
+        let events = framer.push(b"5;58;30Mx");
+
+        assert_eq!(events.len(), 1);
+        assert_raw_key(
+            events.into_iter().next().unwrap(),
+            KeyCode::Char('x'),
+            KeyModifiers::empty(),
+        );
+    }
+
+    #[test]
+    fn timed_out_sgr_mouse_discard_state_clears_at_quiescence() {
+        let mut framer = RawInputByteFramer::default();
+
+        assert!(framer.push(b"\x1b[<3").is_empty());
+        assert!(framer.flush_timeout().is_empty());
+        assert!(framer.flush_timeout().is_empty());
+        assert_eq!(framer.push(b"M"), vec![b"M".to_vec()]);
+    }
+
+    #[test]
     fn sgr_mouse_tail_after_lone_escape_timeout_is_discarded() {
         let mut framer = RawInputFramer::default();