fix: preserve delayed mouse reports with confirmed keyboard input (#4247)

* fix: preserve delayed mouse reports with confirmed keyboard input

refs #3480

* fix: capture geometry before replaying buffered mouse input

refs #3480

---------

Co-authored-by: akbash-bot <300245827+akbash-bot@users.noreply.github.com>
Co-authored-by: JJ Liebig <jonathan.liebig@gmail.com>

1 files changed, 109 insertions, 0 deletions

diff --git a/src/raw_input.rs b/src/raw_input.rs
index 147d6e39..cfbcba62 100644
--- a/src/raw_input.rs
+++ b/src/raw_input.rs
@@ -150,6 +150,7 @@ pub(crate) struct RawInputByteFramer {
     host_color_scheme_change_tracking: bool,
     host_appearance_query_on_focus: bool,
     split_coalesced_escape: bool,
+    host_escape_disambiguation_active: bool,
 }
 
 const HOST_COLOR_QUERY_REPLIES: u16 = 258;
@@ -218,6 +219,11 @@ impl RawInputByteFramer {
         !self.buffer.is_empty()
     }
 
+    #[cfg(any(unix, test))]
+    pub(crate) fn set_host_escape_disambiguation_active(&mut self, active: bool) {
+        self.host_escape_disambiguation_active = active;
+    }
+
     #[cfg(unix)]
     pub(crate) fn has_pending_lone_escape(&self) -> bool {
         self.buffer.as_slice() == [ESC]
@@ -267,6 +273,16 @@ impl RawInputByteFramer {
             return chunks;
         }
 
+        if self.host_escape_disambiguation_active
+            && starts_with_bounded_incomplete_escape_sequence(&self.buffer)
+        {
+            tracing::trace!(
+                len = self.buffer.len(),
+                "holding incomplete host escape sequence with disambiguation active"
+            );
+            return chunks;
+        }
+
         if self.lone_escape_recently_flushed && self.buffer.starts_with(b"[<") {
             tracing::debug!(
                 len = self.buffer.len(),
@@ -472,6 +488,15 @@ impl RawInputByteFramer {
                 continue;
             }
 
+            if self.host_escape_disambiguation_active
+                && self.buffer.first() == Some(&ESC)
+                && self.buffer.len() > 1
+                && !starts_with_known_escape_introducer(&self.buffer)
+            {
+                self.buffer.drain(..1);
+                continue;
+            }
+
             let Some((event, consumed)) = extract_one_event(&self.buffer) else {
                 break;
             };
@@ -831,6 +856,25 @@ fn starts_with_incomplete_sgr_mouse_sequence(buffer: &[u8]) -> bool {
             .all(|byte| byte.is_ascii_digit() || *byte == b';')
 }
 
+fn starts_with_known_escape_introducer(buffer: &[u8]) -> bool {
+    buffer
+        .get(1)
+        .is_some_and(|byte| matches!(*byte, b'[' | b'O' | b']' | b'P' | b'_' | b'^' | b'X' | ESC))
+}
+
+fn starts_with_bounded_incomplete_escape_sequence(buffer: &[u8]) -> bool {
+    if buffer.len() >= MAX_DISCARDED_CONTROL_TAIL_BYTES {
+        return false;
+    }
+    if buffer == [ESC] || buffer == b"\x1bO" {
+        return true;
+    }
+    let Some(body) = buffer.strip_prefix(b"\x1b[") else {
+        return false;
+    };
+    body.iter().all(|byte| matches!(*byte, 0x20..=0x3f))
+}
+
 #[cfg(any(unix, windows, test))]
 fn starts_with_incomplete_default_mouse_sequence(buffer: &[u8]) -> bool {
     buffer.starts_with(b"\x1b[M") && buffer.len() < 6
@@ -2017,6 +2061,71 @@ mod tests {
     }
 
     #[test]
+    fn confirmed_host_disambiguation_retains_split_sgr_mouse_without_escape() {
+        for (prefix, tail) in [
+            (b"\x1b".as_slice(), b"[<0;5;5M".as_slice()),
+            (b"\x1b[".as_slice(), b"<0;5;5M".as_slice()),
+            (b"\x1b[<0;".as_slice(), b"5;5M".as_slice()),
+        ] {
+            let mut framer = RawInputFramer::default();
+            framer
+                .byte_framer
+                .set_host_escape_disambiguation_active(true);
+
+            assert!(framer.push(prefix).is_empty());
+            assert!(framer.flush_timeout().is_empty());
+            assert!(framer.flush_timeout().is_empty());
+            let events = framer.push(tail);
+
+            assert!(matches!(
+                events.as_slice(),
+                [RawInputEvent::Mouse(MouseEvent {
+                    kind: MouseEventKind::Down(MouseButton::Left),
+                    column: 4,
+                    row: 4,
+                    ..
+                })]
+            ));
+        }
+    }
+
+    #[test]
+    fn confirmed_host_disambiguation_drops_stale_escape_before_plain_input() {
+        let mut framer = RawInputFramer::default();
+        framer
+            .byte_framer
+            .set_host_escape_disambiguation_active(true);
+
+        assert!(framer.push(b"\x1b").is_empty());
+        assert!(framer.flush_timeout().is_empty());
+        let events = framer.push(b"x");
+
+        assert_eq!(events.len(), 1);
+        assert_raw_key(
+            events.into_iter().next().unwrap(),
+            KeyCode::Char('x'),
+            KeyModifiers::empty(),
+        );
+    }
+
+    #[test]
+    fn confirmed_host_disambiguation_keeps_kitty_escape_immediate() {
+        let mut framer = RawInputFramer::default();
+        framer
+            .byte_framer
+            .set_host_escape_disambiguation_active(true);
+
+        let events = framer.push(b"\x1b[27u");
+
+        assert_eq!(events.len(), 1);
+        assert_raw_key(
+            events.into_iter().next().unwrap(),
+            KeyCode::Esc,
+            KeyModifiers::empty(),
+        );
+    }
+
+    #[test]
     fn sgr_mouse_tail_after_lone_escape_timeout_is_discarded() {
         let mut framer = RawInputFramer::default();