Support CSI2031 (dark/light theme switching) + remove STDIN cache to avoid stale color queries (#5105)

* initial phase 1 implementation

* various fies

* some fixes

* fix fragmented stdin cases

* some fixes

* renaming

* guard against old clients

* implement csi 2031

* allow toggling theme with actions

* rustfmt

* fix: finalize trailing characters in ansi parser as well

* add pr

1 files changed, 262 insertions, 30 deletions

diff --git a/zellij-server/src/lib.rs b/zellij-server/src/lib.rs
index 1f2b4d14..dea778f4 100644
--- a/zellij-server/src/lib.rs
+++ b/zellij-server/src/lib.rs
@@ -5,6 +5,7 @@ mod os_input_output_unix;
 #[path = "os_input_output_windows.rs"]
 mod os_input_output_windows;
 
+pub mod host_query;
 pub mod os_input_output;
 pub mod output;
 pub mod panes;
@@ -132,6 +133,10 @@ pub enum ServerInstruction {
     WebServerStarted(String), // String -> base_url
     FailedToStartWebServer(String),
     ClearMouseHelpText(ClientId),
+    /// Relay a forwarded-query dispatch from Screen to the server main
+    /// loop. The main loop writes `ServerToClientMsg::ForwardQueryToHost`
+    /// to any connected regular client.
+    ForwardQueryToHost(u32, Vec<u8>),
 }
 
 impl From<&ServerInstruction> for ServerContext {
@@ -180,6 +185,7 @@ impl From<&ServerInstruction> for ServerContext {
                 ServerContext::SendWebClientsForbidden
             },
             ServerInstruction::ClearMouseHelpText(..) => ServerContext::ClearMouseHelpText,
+            ServerInstruction::ForwardQueryToHost(..) => ServerContext::ForwardQueryToHost,
         }
     }
 }
@@ -377,6 +383,28 @@ impl SessionMetaData {
                     ..Default::default()
                 })
             });
+            let host_theme_dark = new_config
+                .options
+                .theme_dark
+                .as_ref()
+                .and_then(|name| new_config.theme_config(Some(name)));
+            let host_theme_light = new_config
+                .options
+                .theme_light
+                .as_ref()
+                .and_then(|name| new_config.theme_config(Some(name)));
+            if new_config.options.theme_dark.is_some() && host_theme_dark.is_none() {
+                log::warn!(
+                    "theme_dark='{}' not found in themes; auto-theme switch disabled for dark.",
+                    new_config.options.theme_dark.as_deref().unwrap_or("?")
+                );
+            }
+            if new_config.options.theme_light.is_some() && host_theme_light.is_none() {
+                log::warn!(
+                    "theme_light='{}' not found in themes; auto-theme switch disabled for light.",
+                    new_config.options.theme_light.as_deref().unwrap_or("?")
+                );
+            }
             self.senders
                 .send_to_screen(ScreenInstruction::Reconfigure {
                     client_id,
@@ -388,6 +416,8 @@ impl SessionMetaData {
                     theme: new_config
                         .theme_config(new_config.options.theme.as_ref())
                         .unwrap_or_else(|| default_palette().into()),
+                    host_theme_dark,
+                    host_theme_light,
                     simplified_ui: new_config.options.simplified_ui.unwrap_or(false),
                     default_shell: new_config.options.default_shell,
                     pane_frames: new_config.options.pane_frames.unwrap_or(true),
@@ -464,10 +494,43 @@ impl Drop for SessionMetaData {
     }
 }
 
+/// Remove a client from session state and synthesize empty
+/// `ForwardedReplyFromHost` instructions for any host-query forwards
+/// that had been dispatched to it. Without this, a client that goes
+/// away while holding an in-flight forward leaves `Screen`'s
+/// `forward_in_flight` flag stuck — the flag is only released by
+/// replies, and no reply will ever come.
+fn remove_client_and_flush_forwards(
+    client_id: ClientId,
+    os_input: &mut Box<dyn ServerOsApi>,
+    session_state: &Arc<RwLock<SessionState>>,
+    session_data: &Arc<RwLock<Option<SessionMetaData>>>,
+) {
+    let _ = os_input.remove_client(client_id);
+    let stuck_tokens = session_state.write().unwrap().remove_client(client_id);
+    if stuck_tokens.is_empty() {
+        return;
+    }
+    if let Some(session) = session_data.read().unwrap().as_ref() {
+        for token in stuck_tokens {
+            let _ = session
+                .senders
+                .send_to_screen(ScreenInstruction::ForwardedReplyFromHost {
+                    token,
+                    reply_bytes: Vec::new(),
+                });
+        }
+    }
+}
+
 macro_rules! remove_client {
-    ($client_id:expr, $os_input:expr, $session_state:expr) => {
-        $os_input.remove_client($client_id).unwrap();
-        $session_state.write().unwrap().remove_client($client_id);
+    ($client_id:expr, $os_input:expr, $session_state:expr, $session_data:expr) => {
+        $crate::remove_client_and_flush_forwards(
+            $client_id,
+            &mut $os_input,
+            &$session_state,
+            &$session_data,
+        );
     };
 }
 
@@ -479,7 +542,7 @@ macro_rules! remove_watcher {
 }
 
 macro_rules! send_to_client {
-    ($client_id:expr, $os_input:expr, $msg:expr, $session_state:expr) => {
+    ($client_id:expr, $os_input:expr, $msg:expr, $session_state:expr, $session_data:expr) => {
         let send_to_client_res = $os_input.send_to_client($client_id, $msg);
         if let Err(e) = send_to_client_res {
             // Try to recover the message
@@ -497,7 +560,7 @@ macro_rules! send_to_client {
             // Log it so it isn't lost
             Err::<(), _>(e).context(context).non_fatal();
             // failed to send to client, remove it
-            remove_client!($client_id, $os_input, $session_state);
+            remove_client!($client_id, $os_input, $session_state, $session_data);
         }
     };
 }
@@ -508,6 +571,13 @@ pub(crate) struct SessionState {
     pipes: HashMap<String, ClientId>,                 // String => pipe_id
     watchers: HashMap<ClientId, bool>, // watcher clients (read-only observers) bool -> is_web_client
     last_active_client: Option<ClientId>, // last client that sent a Key message
+    /// Host-query forward tokens that have been dispatched to a
+    /// specific client and are waiting for a reply. Used to clean up
+    /// when that client disconnects (or when there's no client to
+    /// dispatch to in the first place) — each stuck token gets an
+    /// empty synthetic reply so `Screen`'s `forward_in_flight` slot
+    /// releases and the queued forwards keep moving.
+    forwards_in_flight: HashMap<u32, ClientId>,
 }
 
 impl SessionState {
@@ -517,6 +587,7 @@ impl SessionState {
             pipes: HashMap::new(),
             watchers: HashMap::new(),
             last_active_client: None,
+            forwards_in_flight: HashMap::new(),
         }
     }
     pub fn new_client(&mut self) -> ClientId {
@@ -541,10 +612,25 @@ impl SessionState {
     pub fn associate_pipe_with_client(&mut self, pipe_id: String, client_id: ClientId) {
         self.pipes.insert(pipe_id, client_id);
     }
-    pub fn remove_client(&mut self, client_id: ClientId) {
+    /// Remove a client and return any host-query tokens that had
+    /// been dispatched to this client and were still awaiting a
+    /// reply. Callers must synthesize an empty reply for each token
+    /// (via `ScreenInstruction::ForwardedReplyFromHost`) so
+    /// `Screen`'s in-flight slot releases and any queued forwards
+    /// can proceed.
+    pub fn remove_client(&mut self, client_id: ClientId) -> Vec<u32> {
         self.clients.remove(&client_id);
         self.pipes.retain(|_p_id, c_id| c_id != &client_id);
         self.clear_last_active_client(client_id);
+        let stuck: Vec<u32> = self
+            .forwards_in_flight
+            .iter()
+            .filter_map(|(token, owner)| (*owner == client_id).then_some(*token))
+            .collect();
+        for token in &stuck {
+            self.forwards_in_flight.remove(token);
+        }
+        stuck
     }
     pub fn set_client_size(&mut self, client_id: ClientId, size: Size) {
         self.clients
@@ -648,6 +734,96 @@ impl SessionState {
             self.last_active_client = None;
         }
     }
+    pub fn mark_forward_in_flight(&mut self, token: u32, client_id: ClientId) {
+        self.forwards_in_flight.insert(token, client_id);
+    }
+    pub fn clear_forward_in_flight(&mut self, token: u32) {
+        self.forwards_in_flight.remove(&token);
+    }
+    /// Client to route a host-query forward to. Prefers whichever
+    /// client most recently interacted with the session; falls back
+    /// to any currently-connected non-watcher client. Returns `None`
+    /// only when no regular client is connected.
+    pub fn pick_forward_target(&self) -> Option<ClientId> {
+        if let Some(candidate) = self.last_active_client {
+            if self.clients.contains_key(&candidate) {
+                return Some(candidate);
+            }
+        }
+        self.clients.keys().copied().next()
+    }
+}
+
+#[cfg(test)]
+mod session_state_tests {
+    use super::*;
+
+    fn with_client(id: ClientId) -> SessionState {
+        let mut s = SessionState::new();
+        s.clients.insert(id, None);
+        s
+    }
+
+    #[test]
+    fn pick_forward_target_prefers_last_active_when_still_connected() {
+        let mut s = SessionState::new();
+        s.clients.insert(1, None);
+        s.clients.insert(2, None);
+        s.set_last_active_client(2);
+        assert_eq!(s.pick_forward_target(), Some(2));
+    }
+
+    #[test]
+    fn pick_forward_target_falls_back_when_last_active_disconnected() {
+        let mut s = SessionState::new();
+        s.clients.insert(1, None);
+        s.clients.insert(2, None);
+        // Client 3 was last active but has since disconnected — not in
+        // `clients` map anymore. Must fall through to any connected
+        // client rather than returning None.
+        s.last_active_client = Some(3);
+        let picked = s
+            .pick_forward_target()
+            .expect("some client still connected");
+        assert!(picked == 1 || picked == 2);
+    }
+
+    #[test]
+    fn pick_forward_target_none_when_no_clients() {
+        let s = SessionState::new();
+        assert_eq!(s.pick_forward_target(), None);
+    }
+
+    #[test]
+    fn remove_client_returns_stuck_forward_tokens() {
+        let mut s = with_client(1);
+        s.mark_forward_in_flight(10, 1);
+        s.mark_forward_in_flight(11, 1);
+        // A forward dispatched to a *different* client — must not be
+        // returned when we remove client 1.
+        s.clients.insert(2, None);
+        s.mark_forward_in_flight(12, 2);
+
+        let mut stuck = s.remove_client(1);
+        stuck.sort_unstable();
+        assert_eq!(stuck, vec![10, 11]);
+        assert_eq!(s.forwards_in_flight.get(&12), Some(&2));
+    }
+
+    #[test]
+    fn remove_client_returns_empty_when_no_forwards_in_flight() {
+        let mut s = with_client(1);
+        assert!(s.remove_client(1).is_empty());
+    }
+
+    #[test]
+    fn clear_forward_in_flight_removes_entry() {
+        let mut s = with_client(1);
+        s.mark_forward_in_flight(42, 1);
+        s.clear_forward_in_flight(42);
+        // After clear, removing the client yields no stuck tokens.
+        assert!(s.remove_client(1).is_empty());
+    }
 }
 
 pub fn start_server(mut os_input: Box<dyn ServerOsApi>, socket_path: PathBuf) {
@@ -1048,7 +1224,8 @@ pub fn start_server(mut os_input: Box<dyn ServerOsApi>, socket_path: PathBuf) {
                         client_id,
                         os_input,
                         ServerToClientMsg::UnblockInputThread,
-                        session_state
+                        session_state,
+                        session_data
                     );
                 }
             },
@@ -1062,7 +1239,8 @@ pub fn start_server(mut os_input: Box<dyn ServerOsApi>, socket_path: PathBuf) {
                             ServerToClientMsg::UnblockCliPipeInput {
                                 pipe_name: pipe_name.clone()
                             },
-                            session_state
+                            session_state,
+                            session_data
                         );
                     },
                     None => {
@@ -1075,7 +1253,8 @@ pub fn start_server(mut os_input: Box<dyn ServerOsApi>, socket_path: PathBuf) {
                                 ServerToClientMsg::UnblockCliPipeInput {
                                     pipe_name: pipe_name.clone()
                                 },
-                                session_state
+                                session_state,
+                                session_data
                             );
                         }
                     },
@@ -1092,7 +1271,8 @@ pub fn start_server(mut os_input: Box<dyn ServerOsApi>, socket_path: PathBuf) {
                                 pipe_name: pipe_name.clone(),
                                 output: output.clone()
                             },
-                            session_state
+                            session_state,
+                            session_data
                         );
                     },
                     None => {
@@ -1106,7 +1286,8 @@ pub fn start_server(mut os_input: Box<dyn ServerOsApi>, socket_path: PathBuf) {
                                     pipe_name: pipe_name.clone(),
                                     output: output.clone()
                                 },
-                                session_state
+                                session_state,
+                                session_data
                             );
                         }
                     },
@@ -1136,7 +1317,7 @@ pub fn start_server(mut os_input: Box<dyn ServerOsApi>, socket_path: PathBuf) {
                     os_input.remove_client(client_id).unwrap();
                 } else {
                     // Handle regular client removal
-                    remove_client!(client_id, os_input, session_state);
+                    remove_client!(client_id, os_input, session_state, session_data);
                     drop(completion_tx); // prevent deadlock with route thread
                     if let Some(min_size) = session_state.read().unwrap().min_client_terminal_size()
                     {
@@ -1176,7 +1357,7 @@ pub fn start_server(mut os_input: Box<dyn ServerOsApi>, socket_path: PathBuf) {
                             .collect();
                         // these are just the pipes
                         for client_id in client_ids_to_cleanup {
-                            remove_client!(client_id, os_input, session_state);
+                            remove_client!(client_id, os_input, session_state, session_data);
                         }
 
                         let watcher_client_ids: Vec<ClientId> =
@@ -1211,7 +1392,7 @@ pub fn start_server(mut os_input: Box<dyn ServerOsApi>, socket_path: PathBuf) {
                     os_input.remove_client(client_id).unwrap();
                 } else {
                     // Handle regular client removal
-                    remove_client!(client_id, os_input, session_state);
+                    remove_client!(client_id, os_input, session_state, session_data);
                     if let Some(min_size) = session_state.read().unwrap().min_client_terminal_size()
                     {
                         session_data
@@ -1248,7 +1429,7 @@ pub fn start_server(mut os_input: Box<dyn ServerOsApi>, socket_path: PathBuf) {
                         exit_reason: ExitReason::WebClientsForbidden,
                     },
                 );
-                remove_client!(client_id, os_input, session_state);
+                remove_client!(client_id, os_input, session_state, session_data);
                 if let Some(min_size) = session_state.read().unwrap().min_client_terminal_size() {
                     session_data
                         .write()
@@ -1269,7 +1450,7 @@ pub fn start_server(mut os_input: Box<dyn ServerOsApi>, socket_path: PathBuf) {
                             exit_reason: ExitReason::Normal,
                         },
                     );
-                    remove_client!(client_id, os_input, session_state);
+                    remove_client!(client_id, os_input, session_state, session_data);
                 }
                 break;
             },
@@ -1289,7 +1470,7 @@ pub fn start_server(mut os_input: Box<dyn ServerOsApi>, socket_path: PathBuf) {
                             exit_reason: ExitReason::KickedByHost,
                         },
                     );
-                    remove_client!(client_id, os_input, session_state);
+                    remove_client!(client_id, os_input, session_state, session_data);
                 }
             },
             ServerInstruction::DetachSession(client_ids, completion_tx) => {
@@ -1300,7 +1481,7 @@ pub fn start_server(mut os_input: Box<dyn ServerOsApi>, socket_path: PathBuf) {
                             exit_reason: ExitReason::Normal,
                         },
                     );
-                    remove_client!(*client_id, os_input, session_state);
+                    remove_client!(*client_id, os_input, session_state, session_data);
                 }
                 drop(completion_tx); // we do this here explicitly to signal that the clients have
                                      // already disconnected and to prevent a deadlock below caused
@@ -1348,7 +1529,8 @@ pub fn start_server(mut os_input: Box<dyn ServerOsApi>, socket_path: PathBuf) {
                             ServerToClientMsg::Render {
                                 content: client_render_instruction.clone()
                             },
-                            session_state
+                            session_state,
+                            session_data
                         );
                     }
                 } else {
@@ -1360,7 +1542,7 @@ pub fn start_server(mut os_input: Box<dyn ServerOsApi>, socket_path: PathBuf) {
                                 exit_reason: ExitReason::Normal,
                             },
                         );
-                        remove_client!(client_id, os_input, session_state);
+                        remove_client!(client_id, os_input, session_state, session_data);
                     }
 
                     // Also disconnect all watchers
@@ -1378,7 +1560,7 @@ pub fn start_server(mut os_input: Box<dyn ServerOsApi>, socket_path: PathBuf) {
                                 exit_reason: ExitReason::Normal,
                             },
                         );
-                        remove_client!(watcher_id, os_input, session_state);
+                        remove_client!(watcher_id, os_input, session_state, session_data);
                     }
                     break;
                 }
@@ -1392,13 +1574,13 @@ pub fn start_server(mut os_input: Box<dyn ServerOsApi>, socket_path: PathBuf) {
                             exit_reason: ExitReason::Error(backtrace.clone()),
                         },
                     );
-                    remove_client!(client_id, os_input, session_state);
+                    remove_client!(client_id, os_input, session_state, session_data);
                 }
                 break;
             },
             ServerInstruction::ConnStatus(client_id) => {
                 let _ = os_input.send_to_client(client_id, ServerToClientMsg::Connected);
-                remove_client!(client_id, os_input, session_state);
+                remove_client!(client_id, os_input, session_state, session_data);
             },
             ServerInstruction::Log(
                 lines_to_log,
@@ -1412,7 +1594,8 @@ pub fn start_server(mut os_input: Box<dyn ServerOsApi>, socket_path: PathBuf) {
                     ServerToClientMsg::Log {
                         lines: lines_to_log
                     },
-                    session_state
+                    session_state,
+                    session_data
                 );
             },
             ServerInstruction::LogError(
@@ -1427,7 +1610,8 @@ pub fn start_server(mut os_input: Box<dyn ServerOsApi>, socket_path: PathBuf) {
                     ServerToClientMsg::LogError {
                         lines: lines_to_log
                     },
-                    session_state
+                    session_state,
+                    session_data
                 );
             },
             ServerInstruction::SwitchSession(mut connect_to_session, client_id, completion_tx) => {
@@ -1453,9 +1637,10 @@ pub fn start_server(mut os_input: Box<dyn ServerOsApi>, socket_path: PathBuf) {
                         client_id,
                         os_input,
                         ServerToClientMsg::SwitchSession { connect_to_session },
-                        session_state
+                        session_state,
+                        session_data
                     );
-                    remove_client!(client_id, os_input, session_state);
+                    remove_client!(client_id, os_input, session_state, session_data);
                     drop(completion_tx); // do not deadlock with route thread
 
                     if let Some(min_size) = session_state.read().unwrap().min_client_terminal_size()
@@ -1551,7 +1736,8 @@ pub fn start_server(mut os_input: Box<dyn ServerOsApi>, socket_path: PathBuf) {
                         client_id,
                         os_input,
                         ServerToClientMsg::ConfigFileUpdated,
-                        session_state
+                        session_state,
+                        session_data
                     );
                 }
             },
@@ -1593,7 +1779,8 @@ pub fn start_server(mut os_input: Box<dyn ServerOsApi>, socket_path: PathBuf) {
                         client_id,
                         os_input,
                         ServerToClientMsg::StartWebServer,
-                        session_state
+                        session_state,
+                        session_data
                     );
                 } else {
                     // TODO: test this
@@ -1644,7 +1831,7 @@ pub fn start_server(mut os_input: Box<dyn ServerOsApi>, socket_path: PathBuf) {
                                     exit_reason: ExitReason::WebClientsForbidden,
                                 },
                             );
-                            remove_client!(client_id, os_input, session_state);
+                            remove_client!(client_id, os_input, session_state, session_data);
                         }
                         let web_watcher_client_ids: Vec<ClientId> = session_state
                             .read()
@@ -1707,6 +1894,51 @@ pub fn start_server(mut os_input: Box<dyn ServerOsApi>, socket_path: PathBuf) {
                     .send_to_screen(ScreenInstruction::ClearMouseHelpText(client_id))
                     .unwrap();
             },
+            ServerInstruction::ForwardQueryToHost(token, query_bytes) => {
+                // Pick a regular (non-watcher) client to carry the
+                // forward. Preference is the most recently active
+                // client (whichever last sent input); falls back to
+                // any connected client. When the host terminals of
+                // attached clients differ, the recently-active one
+                // is the best proxy for "what the user is currently
+                // looking at".
+                let target_client_id = {
+                    let mut session = session_state.write().unwrap();
+                    let picked = session.pick_forward_target();
+                    if let Some(cid) = picked {
+                        session.mark_forward_in_flight(token, cid);
+                    }
+                    picked
+                };
+                if let Some(client_id) = target_client_id {
+                    send_to_client!(
+                        client_id,
+                        os_input,
+                        ServerToClientMsg::ForwardQueryToHost { token, query_bytes },
+                        session_state,
+                        session_data
+                    );
+                } else {
+                    // No client to ask — synthesize an empty reply
+                    // so `Screen`'s in-flight slot releases. Without
+                    // this, a detached session (apps still running)
+                    // accumulates forwards in `forward_queue` with
+                    // no way for them to drain until someone
+                    // reattaches.
+                    log::warn!(
+                        "No connected client to forward host query (token={}); returning empty reply",
+                        token
+                    );
+                    if let Some(session) = session_data.read().unwrap().as_ref() {
+                        let _ = session.senders.send_to_screen(
+                            ScreenInstruction::ForwardedReplyFromHost {
+                                token,
+                                reply_bytes: Vec::new(),
+                            },
+                        );
+                    }
+                }
+            },
         }
     }