fix: allow larger kitty graphics frames

3 files changed, 55 insertions, 38 deletions

diff --git a/src/client/mod.rs b/src/client/mod.rs
index 075eeb55..00e1b138 100644
--- a/src/client/mod.rs
+++ b/src/client/mod.rs
@@ -33,7 +33,7 @@ use tracing::{debug, info, warn};
 use crate::server::headless::client_socket_path;
 use crate::server::protocol::{
     self, ClientMessage, NotifyKind, RenderEncoding, ServerMessage, MAX_FRAME_SIZE,
-    PROTOCOL_VERSION,
+    MAX_GRAPHICS_FRAME_SIZE, PROTOCOL_VERSION,
 };
 
 static RECEIVED_KITTY_GRAPHICS_IDS: OnceLock<Mutex<HashSet<u32>>> = OnceLock::new();
@@ -468,7 +468,17 @@ async fn run_client_loop(
     let server_read_tx = event_tx.clone();
     let read_stream = stream.try_clone().map_err(ClientError::ConnectionFailed)?;
     std::thread::spawn(move || {
-        server_reader_thread(read_stream, server_read_tx, &server_read_quit);
+        let max_frame_size = if kitty_graphics_enabled {
+            MAX_GRAPHICS_FRAME_SIZE
+        } else {
+            MAX_FRAME_SIZE
+        };
+        server_reader_thread(
+            read_stream,
+            server_read_tx,
+            &server_read_quit,
+            max_frame_size,
+        );
     });
 
     // Use the original stream for writing (blocking is fine since we write
@@ -588,6 +598,7 @@ fn server_reader_thread(
     mut stream: UnixStream,
     event_tx: tokio::sync::mpsc::Sender<ClientLoopEvent>,
     should_quit: &Arc<AtomicBool>,
+    max_frame_size: usize,
 ) {
     // Ensure the read stream is in blocking mode to avoid WouldBlock errors
     // from read_exact inside read_message. The stream should already be
@@ -603,7 +614,7 @@ fn server_reader_thread(
             break;
         }
 
-        match protocol::read_message(&mut stream, MAX_FRAME_SIZE) {
+        match protocol::read_message(&mut stream, max_frame_size) {
             Ok(msg) => {
                 if event_tx
                     .blocking_send(ClientLoopEvent::ServerMessage(msg))
@@ -730,30 +741,20 @@ fn forward_clipboard(data: &str) {
 // Frame output
 // ---------------------------------------------------------------------------
 
-const SYNC_OUTPUT_END: &[u8] = b"\x1b[?2026l";
-
 fn write_encoded_frame_with_graphics(
     mut writer: impl io::Write,
     encoded: &[u8],
     graphics: &[u8],
 ) -> io::Result<()> {
+    writer.write_all(encoded)?;
     if graphics.is_empty() {
-        return writer.write_all(encoded);
-    }
-
-    record_received_kitty_graphics(graphics);
-
-    if let Some(sync_end) = rfind_subslice(encoded, SYNC_OUTPUT_END) {
-        writer.write_all(&encoded[..sync_end])?;
-        writer.write_all(graphics)?;
-        writer.write_all(&encoded[sync_end..])?;
         return Ok(());
     }
 
-    writer.write_all(b"\x1b[?2026h")?;
-    writer.write_all(encoded)?;
+    record_received_kitty_graphics(graphics);
+    writer.write_all(b"\x1b7")?;
     writer.write_all(graphics)?;
-    writer.write_all(SYNC_OUTPUT_END)
+    writer.write_all(b"\x1b8")
 }
 
 fn contains_kitty_graphics_bytes(bytes: &[u8]) -> bool {
@@ -827,16 +828,6 @@ fn find_subslice(haystack: &[u8], needle: &[u8]) -> Option<usize> {
         .position(|window| window == needle)
 }
 
-fn rfind_subslice(haystack: &[u8], needle: &[u8]) -> Option<usize> {
-    if needle.is_empty() || needle.len() > haystack.len() {
-        return None;
-    }
-
-    haystack
-        .windows(needle.len())
-        .rposition(|window| window == needle)
-}
-
 // ---------------------------------------------------------------------------
 // Resize polling
 // ---------------------------------------------------------------------------
@@ -920,7 +911,7 @@ mod tests {
     use super::*;
 
     #[test]
-    fn graphics_bytes_are_inserted_before_blit_sync_end() {
+    fn graphics_bytes_are_written_after_blit_with_saved_cursor() {
         let mut output = Vec::new();
         write_encoded_frame_with_graphics(
             &mut output,
@@ -929,15 +920,18 @@ mod tests {
         )
         .unwrap();
 
-        assert_eq!(output, b"\x1b[?2026htextgraphics\x1b[?2026lcursor");
+        assert_eq!(
+            output,
+            b"\x1b[?2026htext\x1b[?2026lcursor\x1b7graphics\x1b8"
+        );
     }
 
     #[test]
-    fn graphics_bytes_fallback_wraps_unsynchronized_blit() {
+    fn empty_graphics_writes_only_blit_frame() {
         let mut output = Vec::new();
-        write_encoded_frame_with_graphics(&mut output, b"text", b"graphics").unwrap();
+        write_encoded_frame_with_graphics(&mut output, b"text", b"").unwrap();
 
-        assert_eq!(output, b"\x1b[?2026htextgraphics\x1b[?2026l");
+        assert_eq!(output, b"text");
     }
 
     #[test]
diff --git a/src/server/headless.rs b/src/server/headless.rs
index 3ff3e587..11c486cb 100644
--- a/src/server/headless.rs
+++ b/src/server/headless.rs
@@ -37,7 +37,9 @@ use crate::detect::AgentState;
 use crate::events::AppEvent;
 use crate::layout::PaneId;
 use crate::server::client_transport::{self, ClientWriter, ServerEvent};
-use crate::server::protocol::{self, FrameData, RenderEncoding, ServerMessage, MAX_FRAME_SIZE};
+use crate::server::protocol::{
+    self, FrameData, RenderEncoding, ServerMessage, MAX_FRAME_SIZE, MAX_GRAPHICS_FRAME_SIZE,
+};
 use crate::server::render_stream::ClientRenderState;
 
 // ---------------------------------------------------------------------------
@@ -1008,13 +1010,21 @@ impl HeadlessServer {
 
     /// Encodes a server message into a length-prefixed frame.
     fn frame_server_message(msg: &ServerMessage) -> Result<Vec<u8>, protocol::FramingError> {
+        Self::frame_server_message_with_max(msg, MAX_FRAME_SIZE)
+    }
+
+    /// Encodes a server message using an explicit payload cap.
+    fn frame_server_message_with_max(
+        msg: &ServerMessage,
+        max_frame_size: usize,
+    ) -> Result<Vec<u8>, protocol::FramingError> {
         let mut framed = Vec::new();
         protocol::write_message(&mut framed, msg)?;
         let payload_len = framed.len().saturating_sub(4);
-        if payload_len > MAX_FRAME_SIZE {
+        if payload_len > max_frame_size {
             return Err(protocol::FramingError::Oversized {
                 claimed: payload_len,
-                max: MAX_FRAME_SIZE,
+                max: max_frame_size,
             });
         }
         Ok(framed)
@@ -1542,11 +1552,11 @@ impl HeadlessServer {
             };
 
             let mut commit_graphics_cache = true;
-            if frame.graphics.len() > MAX_FRAME_SIZE {
+            if frame.graphics.len() > MAX_GRAPHICS_FRAME_SIZE {
                 warn!(
                     client_id,
                     graphics_bytes = frame.graphics.len(),
-                    max = MAX_FRAME_SIZE,
+                    max = MAX_GRAPHICS_FRAME_SIZE,
                     "dropping oversized graphics payload for client frame"
                 );
                 frame.graphics.clear();
@@ -1559,7 +1569,15 @@ impl HeadlessServer {
             };
             let mut frame_to_commit = frame.clone();
 
-            let serialized = match Self::frame_server_message(prepared.message()) {
+            let max_frame_size = if frame.graphics.is_empty() {
+                MAX_FRAME_SIZE
+            } else {
+                MAX_GRAPHICS_FRAME_SIZE
+            };
+            let serialized = match Self::frame_server_message_with_max(
+                prepared.message(),
+                max_frame_size,
+            ) {
                 Ok(framed) => framed,
                 Err(protocol::FramingError::Oversized { claimed, max })
                     if !frame.graphics.is_empty() =>
diff --git a/src/server/protocol.rs b/src/server/protocol.rs
index e0a3c4a0..1a220767 100644
--- a/src/server/protocol.rs
+++ b/src/server/protocol.rs
@@ -19,6 +19,11 @@ pub const PROTOCOL_VERSION: u32 = 5;
 /// rejected to prevent denial-of-service via oversized length prefixes.
 pub const MAX_FRAME_SIZE: usize = 2 * 1024 * 1024;
 
+/// Maximum allowed server-to-client frame payload when Kitty graphics are enabled.
+/// Normal traffic keeps `MAX_FRAME_SIZE`; this larger cap is only for explicit
+/// image payloads that are naturally much larger after base64 encoding.
+pub const MAX_GRAPHICS_FRAME_SIZE: usize = 32 * 1024 * 1024;
+
 /// Length of the u32 little-endian length prefix in bytes.
 const LENGTH_PREFIX_BYTES: usize = 4;