fix(kitty-graphics): zlib-compress images forwarded to the host terminal (#5608)

* fix(kitty-graphics): zlib-compress images forwarded to the host terminal

`emit_kitty_transmit` sent every image to the host as uncompressed RGBA
base64 (`f=32,t=d`), about 20x the size of a PNG the application sent,
and re-sent it on every tab switch, pane split and resize step. Compress
the cached host payload with zlib and tag the transmit with `o=z`, which
the kitty protocol defines as part of its base feature set.

Measured with a headless client at 429x77 (8x17 px cells): a 121 KB PNG
went from 2.4 MB to 163 KB on the wire, a 943 KB PNG from 19.9 MB to
1.9 MB, with no measurable change in server CPU.

Fixes #5596

* fix: conditional zlib compression depending on outer terminal support

* rustfmt

* add pr

---------

Co-authored-by: Aram Drevekenin <[email protected]>

zellij-client/src/input_handler.rs

@@ -356,6 +356,10 @@ impl InputHandler {
                 self.os_input
                     .send_to_server(ClientToServerMsg::KittyGraphicsSupport { supported });
             },
+            AnsiStdinInstruction::KittyZlibSupport(supported) => {
+                self.os_input
+                    .send_to_server(ClientToServerMsg::KittyZlibSupport { supported });
+            },
             AnsiStdinInstruction::SixelSupport(supported) => {
                 self.os_input
                     .send_to_server(ClientToServerMsg::SixelSupport { supported });

zellij-client/src/stdin_ansi_parser.rs

@@ -65,6 +65,7 @@ pub enum HostReply {
     /// terminal's color-palette theme mode (CSI 2031).
     HostTerminalThemeChanged(HostTerminalThemeMode),
     KittyGraphicsSupport(bool),
+    KittyZlibSupport(bool),
     SixelSupport(bool),
 }
 
@@ -336,6 +337,7 @@ pub struct StdinAnsiParser {
     nested_frame_extractor: nested_session::NestedFrameExtractor,
     in_bracketed_paste: bool,
     startup_kitty_probe: StartupKittyProbe,
+    startup_kitty_zlib_probe: StartupKittyProbe,
     partial_apc: Vec<u8>,
 }
 
@@ -361,6 +363,7 @@ impl StdinAnsiParser {
             nested_frame_extractor: nested_session::NestedFrameExtractor::new(),
             in_bracketed_paste: false,
             startup_kitty_probe: StartupKittyProbe::NotSent,
+            startup_kitty_zlib_probe: StartupKittyProbe::NotSent,
             partial_apc: Vec::new(),
         }
     }
@@ -369,6 +372,10 @@ impl StdinAnsiParser {
         self.startup_kitty_probe = StartupKittyProbe::AwaitingReply;
     }
 
+    pub fn expect_kitty_zlib_probe_reply(&mut self) {
+        self.startup_kitty_zlib_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
@@ -578,6 +585,10 @@ impl StdinAnsiParser {
                                 self.startup_kitty_probe = StartupKittyProbe::Resolved;
                                 out.replies.push(HostReply::KittyGraphicsSupport(false));
                             }
+                            if self.startup_kitty_zlib_probe == StartupKittyProbe::AwaitingReply {
+                                self.startup_kitty_zlib_probe = StartupKittyProbe::Resolved;
+                                out.replies.push(HostReply::KittyZlibSupport(false));
+                            }
                             if let Some(reply) = HostReply::sixel_support_from_primary_da(&raw) {
                                 out.replies.push(reply);
                             }
@@ -758,6 +769,15 @@ impl StdinAnsiParser {
                             )));
                             self.startup_kitty_probe = StartupKittyProbe::Resolved;
                         }
+                        if self.startup_kitty_zlib_probe == StartupKittyProbe::AwaitingReply
+                            && payload.first() == Some(&b'G')
+                            && contains_subslice(payload, b"i=32")
+                        {
+                            replies.push(HostReply::KittyZlibSupport(contains_subslice(
+                                payload, b"i=32;OK",
+                            )));
+                            self.startup_kitty_zlib_probe = StartupKittyProbe::Resolved;
+                        }
                         i += len;
                         continue;
                     },

zellij-client/src/stdin_ansi_parser_tests.rs

@@ -1691,6 +1691,105 @@ fn kitty_probe_reply_fragmented_across_feeds() {
     }
 }
 
+fn zlib_replies(replies: &[HostReply]) -> Vec<bool> {
+    replies
+        .iter()
+        .filter_map(|r| match r {
+            HostReply::KittyZlibSupport(supported) => Some(*supported),
+            _ => None,
+        })
+        .collect()
+}
+
+fn graphics_replies(replies: &[HostReply]) -> Vec<bool> {
+    replies
+        .iter()
+        .filter_map(|r| match r {
+            HostReply::KittyGraphicsSupport(supported) => Some(*supported),
+            _ => None,
+        })
+        .collect()
+}
+
+fn parser_expecting_both_kitty_probes() -> StdinAnsiParser {
+    let mut parser = StdinAnsiParser::new();
+    parser.expect_kitty_probe_reply();
+    parser.expect_kitty_zlib_probe_reply();
+    parser
+}
+
+#[test]
+fn kitty_zlib_probe_ok_reply_classifies_true() {
+    let mut parser = parser_expecting_both_kitty_probes();
+    let (replies, residue) = feed_once(&mut parser, b"\x1b_Gi=32;OK\x1b\\");
+    assert!(residue.is_empty());
+    assert_eq!(zlib_replies(&replies), vec![true]);
+    assert!(graphics_replies(&replies).is_empty());
+}
+
+#[test]
+fn kitty_zlib_probe_error_reply_classifies_false() {
+    let mut parser = parser_expecting_both_kitty_probes();
+    let (replies, residue) = feed_once(
+        &mut parser,
+        b"\x1b_Gi=32;ENODATA:Insufficient image data\x1b\\",
+    );
+    assert!(residue.is_empty());
+    assert_eq!(zlib_replies(&replies), vec![false]);
+}
+
+#[test]
+fn kitty_zlib_probe_absence_resolves_false_on_barrier() {
+    let mut parser = parser_expecting_both_kitty_probes();
+    let (replies, residue) = feed_once(&mut parser, b"\x1b_Gi=31;OK\x1b\\\x1b[?62;22c");
+    assert!(residue.is_empty());
+    assert_eq!(graphics_replies(&replies), vec![true]);
+    assert_eq!(zlib_replies(&replies), vec![false]);
+}
+
+#[test]
+fn both_kitty_probe_replies_in_one_chunk_are_classified_in_order() {
+    let mut parser = parser_expecting_both_kitty_probes();
+    let (replies, residue) = feed_once(
+        &mut parser,
+        b"\x1b_Gi=31;OK\x1b\\\x1b_Gi=32;OK\x1b\\\x1b[?62;22c",
+    );
+    assert!(residue.is_empty());
+    assert_eq!(graphics_replies(&replies), vec![true]);
+    assert_eq!(zlib_replies(&replies), vec![true]);
+    let graphics_position = replies
+        .iter()
+        .position(|r| matches!(r, HostReply::KittyGraphicsSupport(_)))
+        .unwrap();
+    let zlib_position = replies
+        .iter()
+        .position(|r| matches!(r, HostReply::KittyZlibSupport(_)))
+        .unwrap();
+    assert!(graphics_position < zlib_position);
+}
+
+#[test]
+fn kitty_zlib_probe_reply_fragmented_across_feeds() {
+    let full = b"\x1b_Gi=32;OK\x1b\\";
+    for split in 1..full.len() {
+        let mut parser = parser_expecting_both_kitty_probes();
+        let r1 = parser.feed(&full[..split]);
+        let r2 = parser.feed(&full[split..]);
+        assert!(r1.residue.is_empty(), "split at {}", split);
+        assert!(r2.residue.is_empty(), "split at {}", split);
+        let replies: Vec<HostReply> = r1.replies.into_iter().chain(r2.replies).collect();
+        assert_eq!(zlib_replies(&replies), vec![true], "split at {}", split);
+    }
+}
+
+#[test]
+fn kitty_zlib_reply_outside_probe_window_is_not_classified() {
+    let mut parser = StdinAnsiParser::new();
+    let out = parser.feed(b"\x1b_Gi=32;OK\x1b\\");
+    assert!(out.replies.is_empty());
+    assert!(out.residue.is_empty());
+}
+
 #[test]
 fn apc_outside_probe_window_is_not_classified() {
     let mut parser = StdinAnsiParser::new();

zellij-client/src/stdin_handler.rs

@@ -53,6 +53,7 @@ pub(crate) fn stdin_loop(
                 stdin_ansi_parser.open_own_query_batch();
                 if support_kitty_graphics_protocol {
                     stdin_ansi_parser.expect_kitty_probe_reply();
+                    stdin_ansi_parser.expect_kitty_zlib_probe_reply();
                 }
                 let mut stdout = os_input.get_stdout_writer();
                 let _ = stdout.write_all(query_string.as_bytes());
@@ -62,11 +63,13 @@ pub(crate) fn stdin_loop(
                 let _ =
                     send_input_instructions.send(InputInstruction::AnsiStdinInstructions(vec![
                         HostReply::KittyGraphicsSupport(false),
+                        HostReply::KittyZlibSupport(false),
                     ]));
             }
         } else {
             let _ = send_input_instructions.send(InputInstruction::AnsiStdinInstructions(vec![
                 HostReply::KittyGraphicsSupport(false),
+                HostReply::KittyZlibSupport(false),
                 HostReply::SixelSupport(false),
             ]));
         }
@@ -404,7 +407,7 @@ fn build_startup_query_string(
     // 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}"
+        "\u{1b}_Ga=q,i=31,s=1,v=1,t=d,f=24;AAAA\u{1b}\u{5c}\u{1b}_Ga=q,i=32,s=4,v=4,t=d,f=24,o=z;eJxjYCANAAAAMAAB\u{1b}\u{5c}"
     } else {
         ""
     };
@@ -481,7 +484,7 @@ mod tests {
         let query = build_startup_query_string(true, false);
         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}_Ga=q,i=31,s=1,v=1,t=d,f=24;AAAA\u{1b}\u{5c}\u{1b}[c"
+            "\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}_Ga=q,i=32,s=4,v=4,t=d,f=24,o=z;eJxjYCANAAAAMAAB\u{1b}\u{5c}\u{1b}[c"
         );
         assert!(
             !query.contains("\u{1b}]4;"),
@@ -520,6 +523,18 @@ mod tests {
         assert!(probe_pos < barrier_pos);
     }
 
+    #[test]
+    fn startup_query_contains_zlib_probe_between_kitty_probe_and_barrier() {
+        let query = build_startup_query_string(true, false);
+        let probe_pos = query.find("\u{1b}_Ga=q,i=31,").unwrap();
+        let zlib_probe_pos = query
+            .find("\u{1b}_Ga=q,i=32,s=4,v=4,t=d,f=24,o=z;eJxjYCANAAAAMAAB\u{1b}\u{5c}")
+            .unwrap();
+        let barrier_pos = query.find("\u{1b}[c").unwrap();
+        assert!(probe_pos < zlib_probe_pos);
+        assert!(zlib_probe_pos < barrier_pos);
+    }
+
     #[test]
     fn startup_query_omits_kitty_probe_when_the_protocol_is_disabled() {
         let query = build_startup_query_string(false, false);