fix: refresh host palette on terminal redraw (#4349)

* fix: refresh host palette on terminal redraw

* fix: count overlapping host palette queries

---------

Co-authored-by: Ogulcan Celik <ogulcancelik@gmail.com>

3 files changed, 55 insertions, 3 deletions

diff --git a/src/client/input.rs b/src/client/input.rs
index 23067576..1e308d54 100644
--- a/src/client/input.rs
+++ b/src/client/input.rs
@@ -10,7 +10,7 @@
 //! - We avoid duplicating parsing logic in the client
 //! - Host terminal control replies can be buffered or discarded before they leak
 
-use std::sync::atomic::{AtomicBool, Ordering};
+use std::sync::atomic::{AtomicBool, AtomicU32, Ordering};
 use std::sync::Arc;
 
 #[cfg(unix)]
@@ -39,6 +39,7 @@ pub fn stdin_reader_loop(
     event_tx: mpsc::Sender<ClientLoopEvent>,
     should_quit: &Arc<AtomicBool>,
     host_color_query_sent: bool,
+    host_theme_query_pending: Arc<AtomicU32>,
     host_cell_size_query_sent: bool,
     host_mouse_capture_active: Arc<AtomicBool>,
     host_sgr_pixels_active: Arc<AtomicBool>,
@@ -51,6 +52,7 @@ pub fn stdin_reader_loop(
     {
         let _ = (
             host_color_query_sent,
+            host_theme_query_pending,
             host_cell_size_query_sent,
             host_mouse_capture_active,
             host_sgr_pixels_active,
@@ -64,6 +66,7 @@ pub fn stdin_reader_loop(
         event_tx,
         should_quit,
         host_color_query_sent,
+        host_theme_query_pending,
         host_cell_size_query_sent,
         host_mouse_capture_active,
         host_sgr_pixels_active,
@@ -79,6 +82,7 @@ fn unix_stdin_reader_loop(
     event_tx: mpsc::Sender<ClientLoopEvent>,
     should_quit: &Arc<AtomicBool>,
     host_color_query_sent: bool,
+    host_theme_query_pending: Arc<AtomicU32>,
     host_cell_size_query_sent: bool,
     host_mouse_capture_active: Arc<AtomicBool>,
     host_sgr_pixels_active: Arc<AtomicBool>,
@@ -178,6 +182,11 @@ fn unix_stdin_reader_loop(
         match reader.read(&mut scratch) {
             Ok(0) => break,
             Ok(n) => {
+                // A redraw can issue queries while this thread is blocked in read().
+                // Arm the split-reply guard before framing the returned bytes.
+                for _ in 0..host_theme_query_pending.swap(0, Ordering::AcqRel) {
+                    framer.host_color_query_sent();
+                }
                 let sgr_pixels = *pending_mode
                     .get_or_insert_with(|| host_sgr_pixels_active.load(Ordering::Acquire));
                 if sgr_pixels {
diff --git a/src/client/mod.rs b/src/client/mod.rs
index 1a0f8dd5..2cec6d2d 100644
--- a/src/client/mod.rs
+++ b/src/client/mod.rs
@@ -128,7 +128,7 @@ use terminal_sessions::terminal_control_command_from_json;
 #[cfg(unix)]
 use std::collections::HashMap;
 use std::io::{self, Write as _};
-use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
+use std::sync::atomic::{AtomicBool, AtomicU32, AtomicU64, Ordering};
 use std::sync::Arc;
 #[cfg(unix)]
 use std::sync::Mutex;
@@ -509,6 +509,8 @@ async fn run_client_loop(
     // Spawn the stdin reader thread.
     let will_query_host_terminal_theme =
         state.attach_escape.is_none() && should_query_host_terminal_theme();
+    let host_theme_query_pending = Arc::new(AtomicU32::new(0));
+    let stdin_host_theme_query_pending = host_theme_query_pending.clone();
     // Terminals behind ConPTY report no pixel size through the ioctl, so ask the
     // host terminal directly instead of falling back to an assumed cell size.
     let will_query_host_cell_size = state.attach_escape.is_none()
@@ -530,6 +532,7 @@ async fn run_client_loop(
             stdin_tx,
             &stdin_quit,
             will_query_host_terminal_theme,
+            stdin_host_theme_query_pending,
             will_query_host_cell_size,
             stdin_mouse_capture_active,
             stdin_sgr_pixels_active,
@@ -1181,6 +1184,13 @@ async fn run_client_loop(
                 cell_height_px,
                 pixel_geometry_exact,
             ) => {
+                // Palette updates through OSC do not necessarily produce a color-scheme
+                // notification. Re-query on redraw, including SIGWINCH without a resize,
+                // so desktop theme switchers can refresh the existing panes in place.
+                if will_query_host_terminal_theme {
+                    host_theme_query_pending.fetch_add(1, Ordering::AcqRel);
+                    query_host_terminal_theme();
+                }
                 if !pixel_geometry_exact && host_sgr_pixels_active.load(Ordering::Acquire) {
                     set_mouse_capture(state.mouse_capture_active, false)
                         .map_err(ClientError::ConnectionFailed)?;
diff --git a/src/raw_input.rs b/src/raw_input.rs
index b2c3beda..2694912c 100644
--- a/src/raw_input.rs
+++ b/src/raw_input.rs
@@ -179,8 +179,11 @@ impl RawInputByteFramer {
 
     /// Hold a lone trailing ESC for one idle flush so an OSC 10/11 reply split
     /// at its ESC introducer stitches back together instead of leaking (#549).
+    /// Overlapping queries each add their own replies to the window.
     pub(crate) fn host_color_query_sent(&mut self) {
-        self.host_color_replies_awaited = HOST_COLOR_QUERY_REPLIES;
+        self.host_color_replies_awaited = self
+            .host_color_replies_awaited
+            .saturating_add(HOST_COLOR_QUERY_REPLIES);
         self.held_pending_host_reply_esc = false;
     }
 
@@ -2859,4 +2862,34 @@ mod tests {
         assert!(framer.push(b"\x1b").is_empty());
         assert_eq!(framer.flush_timeout(), vec![b"\x1b".to_vec()]);
     }
+
+    #[test]
+    fn overlapping_host_color_queries_keep_holding_split_reply_escape() {
+        use std::fmt::Write as _;
+
+        let mut framer = RawInputByteFramer::default();
+        framer.host_color_query_sent();
+        framer.host_color_query_sent();
+        let mut first_batch =
+            String::from("\x1b]10;rgb:6565/7b7b/8383\x1b\\\x1b]11;rgb:2424/2727/3a3a\x1b\\");
+        for index in 0..=u8::MAX {
+            let _ = write!(first_batch, "\x1b]4;{index};rgb:1111/2222/3333\x1b\\");
+        }
+        assert_eq!(framer.push(first_batch.as_bytes()).len(), 258);
+
+        // The second batch is still outstanding, so its first reply split at ESC
+        // must not leak into the pane as an Escape key.
+        assert!(framer.push(b"\x1b").is_empty());
+        assert!(framer.flush_timeout().is_empty());
+        let chunks = framer.push(b"]10;rgb:eeee/eeee/eeee\x1b\\");
+        assert_eq!(chunks.len(), 1);
+        let (event, _) = extract_one_event(&chunks[0]).unwrap();
+        assert!(matches!(
+            event,
+            RawInputEvent::HostDefaultColor {
+                kind: DefaultColorKind::Foreground,
+                ..
+            }
+        ));
+    }
 }