Support the Kitty Image Protocol (#5428)

* initial implementation - kinda working

* tests and fixes

* fix bounds

* tests and compatibility fixes

* fix: cursor scroll, clears, deletes, key tolerance, and host id range

* fix: keep images alive across resize, relayout and history scroll

* anchor to lines, not pixels

* fix image/text alignments

* config to allow disabling

* fix pinned pane issue

* never let host graphics replies reach the keyboard parser

* fix nested session

* properly handle fullscreening a pane

* rustfmt

* cleanup of ignored test

* add pr

5 files changed, 298 insertions, 12 deletions

diff --git a/zellij-client/src/input_handler.rs b/zellij-client/src/input_handler.rs
index 3c9d269c..8ae35afd 100644
--- a/zellij-client/src/input_handler.rs
+++ b/zellij-client/src/input_handler.rs
@@ -338,6 +338,10 @@ impl InputHandler {
                 self.os_input
                     .send_to_server(ClientToServerMsg::HostTerminalThemeChanged { mode });
             },
+            AnsiStdinInstruction::KittyGraphicsSupport(supported) => {
+                self.os_input
+                    .send_to_server(ClientToServerMsg::KittyGraphicsSupport { supported });
+            },
         }
     }
     fn handle_nested_session_frame_from_host(&mut self, payload_bytes: Vec<u8>) {
diff --git a/zellij-client/src/lib.rs b/zellij-client/src/lib.rs
index 28860dbd..08a7f655 100644
--- a/zellij-client/src/lib.rs
+++ b/zellij-client/src/lib.rs
@@ -928,6 +928,9 @@ pub fn start_client(
         .support_kitty_keyboard_protocol
         .map(|e| !e)
         .unwrap_or(false);
+    let support_kitty_graphics_protocol = config_options
+        .support_kitty_graphics_protocol
+        .unwrap_or(true);
     let should_start_web_server = config_options.web_server.map(|w| w).unwrap_or(false);
     let mut reconnect_to_session = None;
     os_input.unset_raw_mode().unwrap();
@@ -1204,6 +1207,7 @@ pub fn start_client(
                     send_input_instructions,
                     stdin_ansi_parser,
                     explicitly_disable_kitty_keyboard_protocol,
+                    support_kitty_graphics_protocol,
                     Some(resize_sender),
                 )
             }
diff --git a/zellij-client/src/stdin_ansi_parser.rs b/zellij-client/src/stdin_ansi_parser.rs
index c70f9e93..c02b8762 100644
--- a/zellij-client/src/stdin_ansi_parser.rs
+++ b/zellij-client/src/stdin_ansi_parser.rs
@@ -62,6 +62,7 @@ pub enum HostReply {
     /// DSR 997 reply / unsolicited notification reporting the host
     /// terminal's color-palette theme mode (CSI 2031).
     HostTerminalThemeChanged(HostTerminalThemeMode),
+    KittyGraphicsSupport(bool),
 }
 
 /// Retained alias for the pre-refactor type name used by other modules in
@@ -254,6 +255,13 @@ enum SeqStatus {
     Malformed,
 }
 
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+enum StartupKittyProbe {
+    NotSent,
+    AwaitingReply,
+    Resolved,
+}
+
 /// Continuous host-reply parser. Lives for the whole client session.
 pub struct StdinAnsiParser {
     inner: InputParser,
@@ -268,6 +276,8 @@ pub struct StdinAnsiParser {
     partial_paste: Vec<u8>,
     nested_frame_extractor: nested_session::NestedFrameExtractor,
     in_bracketed_paste: bool,
+    startup_kitty_probe: StartupKittyProbe,
+    partial_apc: Vec<u8>,
 }
 
 impl std::fmt::Debug for StdinAnsiParser {
@@ -290,9 +300,15 @@ impl StdinAnsiParser {
             partial_paste: Vec::new(),
             nested_frame_extractor: nested_session::NestedFrameExtractor::new(),
             in_bracketed_paste: false,
+            startup_kitty_probe: StartupKittyProbe::NotSent,
+            partial_apc: Vec::new(),
         }
     }
 
+    pub fn expect_kitty_probe_reply(&mut self) {
+        self.startup_kitty_probe = StartupKittyProbe::AwaitingReply;
+    }
+
     /// Open a forwarding window for `token`. Subsequent reply events that
     /// arrive before the Primary-DA barrier will be accumulated into the
     /// slot's `reply_bytes`, in addition to being dispatched as normal
@@ -356,12 +372,13 @@ impl StdinAnsiParser {
         let (bytes, nested_frames) = self.nested_frame_extractor.extract(bytes);
         let bytes = &bytes[..];
         out.nested_frames = nested_frames;
+        let sanitized = self.extract_kitty_probe_reply(bytes, &mut out.replies);
         // Collect events first (borrow-splits the InputParser across the
         // callback and the post-processing mutations).
         let mut events = Vec::new();
         let mut residue = Vec::new();
         self.inner.parse(
-            bytes,
+            &sanitized,
             |event| {
                 events.push(event);
             },
@@ -398,6 +415,10 @@ impl StdinAnsiParser {
                 } => {
                     match final_byte {
                         b'c' => {
+                            if self.startup_kitty_probe == StartupKittyProbe::AwaitingReply {
+                                self.startup_kitty_probe = StartupKittyProbe::Resolved;
+                                out.replies.push(HostReply::KittyGraphicsSupport(false));
+                            }
                             // Primary-DA — the barrier. Close the slot and
                             // emit the completed forwarded reply if active.
                             if let Some(slot) = self.active_forward.take() {
@@ -434,17 +455,18 @@ impl StdinAnsiParser {
         // other bytes pass through unchanged. The walk is stateful so
         // an OSC/CSI sequence split across `feed()` calls is buffered
         // rather than leaking into residue.
-        residue.extend(self.strip_replies(bytes));
+        residue.extend(self.strip_replies(&sanitized));
         out.residue = residue;
         out.has_partial_state = !self.partial_osc.is_empty()
             || !self.partial_csi.is_empty()
             || !self.partial_paste.is_empty()
-            || !self.nested_frame_extractor.partial_bytes().is_empty();
+            || !self.nested_frame_extractor.partial_bytes().is_empty()
+            || !self.partial_apc.is_empty();
         out
     }
 
     pub fn pending_partial(&self) -> PendingPartial {
-        if !self.partial_paste.is_empty() {
+        if !self.partial_paste.is_empty() || !self.partial_apc.is_empty() {
             PendingPartial::ReplyInProgress
         } else if self.partial_csi.is_empty()
             && self.partial_osc.is_empty()
@@ -487,12 +509,74 @@ impl StdinAnsiParser {
             self.partial_osc.len()
                 + self.partial_csi.len()
                 + self.partial_paste.len()
-                + partial_nested_frame.len(),
+                + partial_nested_frame.len()
+                + self.partial_apc.len(),
         );
         out.append(&mut self.partial_osc);
         out.append(&mut self.partial_csi);
         out.append(&mut self.partial_paste);
         out.append(&mut partial_nested_frame);
+        out.append(&mut self.partial_apc);
+        out
+    }
+
+    fn extract_kitty_probe_reply(&mut self, bytes: &[u8], replies: &mut Vec<HostReply>) -> Vec<u8> {
+        let mut working = std::mem::take(&mut self.partial_apc);
+        if working.is_empty() && self.partial_osc == [0x1b] && bytes.first() == Some(&b'_') {
+            working.append(&mut self.partial_osc);
+        }
+        working.extend_from_slice(bytes);
+        let mut out = Vec::with_capacity(working.len());
+        let mut i = 0;
+        while i < working.len() {
+            let rest = &working[i..];
+            if rest.len() >= 2 && rest[0] == 0x1b && rest[1] == b'_' {
+                match rest.get(2) {
+                    Some(&b'G') => {},
+                    Some(_) => {
+                        out.push(working[i]);
+                        i += 1;
+                        continue;
+                    },
+                    None => {
+                        self.partial_apc = rest.to_vec();
+                        return out;
+                    },
+                }
+                match apc_status(rest) {
+                    SeqStatus::Complete(len) => {
+                        let payload = &rest[2..len - 2];
+                        if self.startup_kitty_probe == StartupKittyProbe::AwaitingReply
+                            && payload.first() == Some(&b'G')
+                            && contains_subslice(payload, b"i=31")
+                        {
+                            replies.push(HostReply::KittyGraphicsSupport(contains_subslice(
+                                payload, b"i=31;OK",
+                            )));
+                            self.startup_kitty_probe = StartupKittyProbe::Resolved;
+                        }
+                        i += len;
+                        continue;
+                    },
+                    SeqStatus::NeedMore => {
+                        let tail = rest.to_vec();
+                        if tail.len() > PARTIAL_BUFFER_CAP_BYTES {
+                            out.extend_from_slice(&tail);
+                        } else {
+                            self.partial_apc = tail;
+                        }
+                        return out;
+                    },
+                    SeqStatus::Malformed => {
+                        out.push(working[i]);
+                        i += 1;
+                        continue;
+                    },
+                }
+            }
+            out.push(working[i]);
+            i += 1;
+        }
         out
     }
 
@@ -649,6 +733,28 @@ fn osc_status(buf: &[u8]) -> SeqStatus {
     SeqStatus::NeedMore
 }
 
+fn apc_status(buf: &[u8]) -> SeqStatus {
+    if buf.get(0) != Some(&0x1b) || buf.get(1) != Some(&b'_') {
+        return SeqStatus::Malformed;
+    }
+    let mut i = 2;
+    while i < buf.len() {
+        match buf[i] {
+            0x1b => match buf.get(i + 1) {
+                Some(&b'\\') => return SeqStatus::Complete(i + 2),
+                Some(_) => return SeqStatus::Malformed,
+                None => return SeqStatus::NeedMore,
+            },
+            _ => i += 1,
+        }
+    }
+    SeqStatus::NeedMore
+}
+
+fn contains_subslice(haystack: &[u8], needle: &[u8]) -> bool {
+    haystack.windows(needle.len()).any(|w| w == needle)
+}
+
 /// Walk a whitelisted CSI report starting at the head of `buf`.
 fn csi_status(buf: &[u8]) -> SeqStatus {
     if buf.get(0) != Some(&0x1b) || buf.get(1) != Some(&b'[') {
diff --git a/zellij-client/src/stdin_ansi_parser_tests.rs b/zellij-client/src/stdin_ansi_parser_tests.rs
index 77d796eb..9b98caa8 100644
--- a/zellij-client/src/stdin_ansi_parser_tests.rs
+++ b/zellij-client/src/stdin_ansi_parser_tests.rs
@@ -1210,6 +1210,139 @@ fn last_resort_force_drain_recovers_a_never_terminated_partial() {
 }
 
 #[test]
+fn kitty_probe_ok_reply_classifies_true() {
+    let mut parser = StdinAnsiParser::new();
+    parser.expect_kitty_probe_reply();
+    let (replies, residue) = feed_once(&mut parser, b"\x1b_Gi=31;OK\x1b\\");
+    assert!(residue.is_empty(), "probe reply must be fully consumed");
+    assert_eq!(replies.len(), 1);
+    match &replies[0] {
+        HostReply::KittyGraphicsSupport(supported) => assert!(*supported),
+        other => panic!("unexpected reply: {:?}", other),
+    }
+}
+
+#[test]
+fn kitty_probe_error_reply_classifies_false() {
+    let mut parser = StdinAnsiParser::new();
+    parser.expect_kitty_probe_reply();
+    let (replies, residue) = feed_once(&mut parser, b"\x1b_Gi=31;ENOTSUPPORTED:query failed\x1b\\");
+    assert!(residue.is_empty());
+    assert_eq!(replies.len(), 1);
+    match &replies[0] {
+        HostReply::KittyGraphicsSupport(supported) => assert!(!*supported),
+        other => panic!("unexpected reply: {:?}", other),
+    }
+}
+
+#[test]
+fn kitty_probe_absence_resolves_false_on_barrier() {
+    let mut parser = StdinAnsiParser::new();
+    parser.expect_kitty_probe_reply();
+    let (replies, residue) = feed_once(&mut parser, b"\x1b[?62;4;52c");
+    assert!(residue.is_empty());
+    assert_eq!(replies.len(), 1);
+    match &replies[0] {
+        HostReply::KittyGraphicsSupport(supported) => assert!(!*supported),
+        other => panic!("unexpected reply: {:?}", other),
+    }
+}
+
+#[test]
+fn kitty_probe_reply_then_barrier_emits_exactly_one_reply() {
+    let mut parser = StdinAnsiParser::new();
+    parser.expect_kitty_probe_reply();
+    let mut chunk = Vec::new();
+    chunk.extend_from_slice(b"\x1b_Gi=31;OK\x1b\\");
+    chunk.extend_from_slice(b"\x1b[?62;4;52c");
+    let (replies, residue) = feed_once(&mut parser, &chunk);
+    assert!(residue.is_empty());
+    let kitty_replies: Vec<bool> = replies
+        .iter()
+        .filter_map(|r| match r {
+            HostReply::KittyGraphicsSupport(supported) => Some(*supported),
+            _ => None,
+        })
+        .collect();
+    assert_eq!(kitty_replies, vec![true]);
+}
+
+#[test]
+fn kitty_probe_reply_fragmented_across_feeds() {
+    let full = b"\x1b_Gi=31;OK\x1b\\";
+    for split in 1..full.len() {
+        let mut parser = StdinAnsiParser::new();
+        parser.expect_kitty_probe_reply();
+        let r1 = parser.feed(&full[..split]);
+        let r2 = parser.feed(&full[split..]);
+        assert!(
+            r1.residue.is_empty(),
+            "split at {}: chunk 1 residue {:?}",
+            split,
+            r1.residue
+        );
+        assert!(
+            r2.residue.is_empty(),
+            "split at {}: chunk 2 residue {:?}",
+            split,
+            r2.residue
+        );
+        assert_eq!(
+            r1.replies.len() + r2.replies.len(),
+            1,
+            "split at {}: exactly one reply across both chunks",
+            split
+        );
+        let reply = r1.replies.into_iter().chain(r2.replies).next().unwrap();
+        match reply {
+            HostReply::KittyGraphicsSupport(supported) => assert!(supported),
+            other => panic!("split at {}: unexpected reply {:?}", split, other),
+        }
+    }
+}
+
+#[test]
+fn apc_outside_probe_window_is_not_classified() {
+    let mut parser = StdinAnsiParser::new();
+    let out = parser.feed(b"\x1b_Gi=31;OK\x1b\\");
+    assert!(out.replies.is_empty());
+}
+
+#[test]
+fn unsolicited_kitty_graphics_reply_never_becomes_keyboard_input() {
+    let mut parser = StdinAnsiParser::new();
+    parser.expect_kitty_probe_reply();
+    let out = parser.feed(b"\x1b_Gi=31;OK\x1b\\");
+    assert_eq!(out.replies.len(), 1);
+
+    let out = parser.feed(b"\x1b_Gi=2000000000;OK\x1b\\");
+    assert!(
+        out.residue.is_empty(),
+        "a host reply to our own image transmission must not be typed into the focused pane, got: {:?}",
+        out.residue
+    );
+    assert!(out.replies.is_empty());
+
+    let mut fragmented = StdinAnsiParser::new();
+    fragmented.expect_kitty_probe_reply();
+    let _ = fragmented.feed(b"\x1b_Gi=31;OK\x1b\\");
+    let first = fragmented.feed(b"\x1b_Gi=2000000");
+    let second = fragmented.feed(b"000;OK\x1b\\");
+    assert!(first.residue.is_empty() && second.residue.is_empty());
+}
+
+#[test]
+fn non_kitty_apc_sequences_are_left_alone() {
+    let mut parser = StdinAnsiParser::new();
+    parser.expect_kitty_probe_reply();
+    let out = parser.feed(b"\x1b_Xsomething\x1b\\");
+    assert!(
+        !out.residue.is_empty(),
+        "only kitty graphics APCs may be swallowed"
+    );
+}
+
+#[test]
 fn unsolicited_theme_notification_classifies_without_outstanding_query() {
     let mut p = StdinAnsiParser::new();
     let out = p.feed(b"\x1b[?997;1n");
diff --git a/zellij-client/src/stdin_handler.rs b/zellij-client/src/stdin_handler.rs
index 681b49b5..60d55ab5 100644
--- a/zellij-client/src/stdin_handler.rs
+++ b/zellij-client/src/stdin_handler.rs
@@ -2,7 +2,7 @@ use crate::keyboard_parser::{KittyKeyboardParser, KittyParseOutcome};
 use crate::os_input_output::ClientOsApi;
 #[cfg(windows)]
 use crate::os_input_output_windows::use_vt_path;
-use crate::stdin_ansi_parser::{PendingPartial, StdinAnsiParser};
+use crate::stdin_ansi_parser::{HostReply, PendingPartial, StdinAnsiParser};
 #[cfg(windows)]
 use crate::stdin_handler_windows::enable_vt_input;
 use crate::InputInstruction;
@@ -21,6 +21,7 @@ pub(crate) fn stdin_loop(
     send_input_instructions: SenderWithContext<InputInstruction>,
     stdin_ansi_parser: Arc<Mutex<StdinAnsiParser>>,
     explicitly_disable_kitty_keyboard_protocol: bool,
+    support_kitty_graphics_protocol: bool,
     resize_sender: Option<std::sync::mpsc::Sender<()>>,
 ) {
     // On Windows we choose between the VT byte path (termwiz/kitty parsing)
@@ -44,11 +45,23 @@ pub(crate) fn stdin_loop(
         let can_query_terminal = true;
 
         if can_query_terminal {
-            let query_string = build_startup_query_string();
+            let query_string = build_startup_query_string(support_kitty_graphics_protocol);
             let _ = os_input
                 .get_stdout_writer()
                 .write(query_string.as_bytes())
                 .unwrap();
+            if support_kitty_graphics_protocol {
+                stdin_ansi_parser.lock().unwrap().expect_kitty_probe_reply();
+            } else {
+                let _ =
+                    send_input_instructions.send(InputInstruction::AnsiStdinInstructions(vec![
+                        HostReply::KittyGraphicsSupport(false),
+                    ]));
+            }
+        } else {
+            let _ = send_input_instructions.send(InputInstruction::AnsiStdinInstructions(vec![
+                HostReply::KittyGraphicsSupport(false),
+            ]));
         }
     }
 
@@ -321,15 +334,24 @@ fn drain_partial_to_keyboard(
 /// Build the fire-and-forget host-query batch sent at client startup.
 /// The host's replies refine `Screen`'s cached state asynchronously as
 /// they arrive; the UI does not block on them.
-fn build_startup_query_string() -> String {
+fn build_startup_query_string(support_kitty_graphics_protocol: bool) -> String {
     // <ESC>[14t => get text area size in pixels,
     // <ESC>[16t => get character cell size in pixels
     // <ESC>]11;?<ESC>\ => get background color
     // <ESC>]10;?<ESC>\ => get foreground color
     // <ESC>[?2026$p => get synchronised output mode
+    // <ESC>_Ga=q,...<ESC>\ => probe kitty graphics support (omitted when the
+    // protocol is disabled), answered by capable terminals only; the trailing
+    // Primary DA is the barrier that resolves the probe negatively when it
+    // goes unanswered
+    let kitty_graphics_probe = if support_kitty_graphics_protocol {
+        "\u{1b}_Ga=q,i=31,s=1,v=1,t=d,f=24;AAAA\u{1b}\u{5c}"
+    } else {
+        ""
+    };
     format!(
-        "{}\u{1b}]11;?\u{1b}\u{5c}\u{1b}]10;?\u{1b}\u{5c}\u{1b}[?2026$p",
-        PIXEL_SIZE_QUERY
+        "{}\u{1b}]11;?\u{1b}\u{5c}\u{1b}]10;?\u{1b}\u{5c}\u{1b}[?2026$p{}\u{1b}[c",
+        PIXEL_SIZE_QUERY, kitty_graphics_probe
     )
 }
 
@@ -346,10 +368,10 @@ mod tests {
 
     #[test]
     fn startup_query_has_no_palette_register_loop() {
-        let query = build_startup_query_string();
+        let query = build_startup_query_string(true);
         assert_eq!(
             query,
-            "\u{1b}[14t\u{1b}[16t\u{1b}]11;?\u{1b}\u{5c}\u{1b}]10;?\u{1b}\u{5c}\u{1b}[?2026$p"
+            "\u{1b}[14t\u{1b}[16t\u{1b}]11;?\u{1b}\u{5c}\u{1b}]10;?\u{1b}\u{5c}\u{1b}[?2026$p\u{1b}_Ga=q,i=31,s=1,v=1,t=d,f=24;AAAA\u{1b}\u{5c}\u{1b}[c"
         );
         assert!(
             !query.contains("\u{1b}]4;"),
@@ -357,4 +379,21 @@ mod tests {
             query
         );
     }
+
+    #[test]
+    fn startup_query_contains_kitty_probe_before_barrier() {
+        let query = build_startup_query_string(true);
+        assert!(query.contains("\u{1b}_Ga=q,i=31,s=1,v=1,t=d,f=24;AAAA\u{1b}\u{5c}"));
+        let probe_pos = query.find("\u{1b}_Ga=q,i=31,").unwrap();
+        let barrier_pos = query.find("\u{1b}[c").unwrap();
+        assert!(probe_pos < barrier_pos);
+    }
+
+    #[test]
+    fn startup_query_omits_kitty_probe_when_the_protocol_is_disabled() {
+        let query = build_startup_query_string(false);
+        assert!(!query.contains("\u{1b}_G"));
+        assert!(query.ends_with("\u{1b}[c"));
+        assert!(query.starts_with("\u{1b}[14t\u{1b}[16t"));
+    }
 }