Make post-submit input gate command agnostic

src/bin/dosh-client.rs

@@ -2844,7 +2844,7 @@ async fn run_terminal(
     let mut pending_user_input: VecDeque<Vec<u8>> = VecDeque::new();
     let mut pending_user_input_bytes = 0usize;
     let mut startup_input_hold_until: Option<Instant> = None;
-    let mut shell_command_line: Vec<u8> = Vec::new();
+    let mut hold_all_startup_input = false;
     if let Some(frame) = first_frame {
         if !forward_only {
             render_frame(&frame)?;
@@ -2862,6 +2862,7 @@ async fn run_terminal(
     {
         send_input(&socket, addr, &cred, &mut send_seq, bytes).await?;
         startup_input_hold_until = Some(Instant::now() + STARTUP_INPUT_HOLD);
+        hold_all_startup_input = true;
     }
 
     let (stdin_tx, mut stdin_rx) = mpsc::unbounded_channel::<Vec<u8>>();
@@ -2901,7 +2902,11 @@ async fn run_terminal(
                         }
                         if cred.mode != "view-only" && cred.mode != "forward-only" {
                             if let Some(deadline) = startup_input_hold_until {
-                                if !predictor.alternate_screen && Instant::now() < deadline {
+                                if !predictor.alternate_screen
+                                    && Instant::now() < deadline
+                                    && (hold_all_startup_input
+                                        || should_hold_post_submit_input(&bytes))
+                                {
                                     queue_pending_user_input(
                                         &mut pending_user_input,
                                         &mut pending_user_input_bytes,
@@ -2910,6 +2915,7 @@ async fn run_terminal(
                                     continue;
                                 }
                                 startup_input_hold_until = None;
+                                hold_all_startup_input = false;
                                 flush_pending_user_input(
                                     &socket,
                                     addr,
@@ -2921,25 +2927,32 @@ async fn run_terminal(
                                 )
                                 .await?;
                             }
-                            let hold_after_submit = update_shell_command_line(
-                                &mut shell_command_line,
-                                &bytes,
-                                predictor.alternate_screen,
-                            );
-                            if hold_after_submit
-                                && let Some((send_now, hold_for_later)) =
+                            if !predictor.alternate_screen
+                                && let Some((send_now, mut hold_for_later)) =
                                     split_after_command_submit(&bytes)
                             {
                                 predictor.observe_input(&send_now)?;
                                 send_input(&socket, addr, &cred, &mut send_seq, send_now).await?;
                                 startup_input_hold_until =
                                     Some(Instant::now() + STARTUP_INPUT_HOLD);
-                                if !hold_for_later.is_empty() {
+                                hold_all_startup_input = false;
+                                if should_hold_post_submit_input(&hold_for_later) {
                                     queue_pending_user_input(
                                         &mut pending_user_input,
                                         &mut pending_user_input_bytes,
                                         hold_for_later,
                                     )?;
+                                    continue;
+                                } else if !hold_for_later.is_empty() {
+                                    predictor.observe_input(&hold_for_later)?;
+                                    send_input(
+                                        &socket,
+                                        addr,
+                                        &cred,
+                                        &mut send_seq,
+                                        std::mem::take(&mut hold_for_later),
+                                    )
+                                    .await?;
                                 }
                                 continue;
                             }
@@ -3065,7 +3078,7 @@ async fn run_terminal(
                                     &mut send_seq,
                                     &mut predictor,
                                     (&mut pending_user_input, &mut pending_user_input_bytes),
-                                    &mut startup_input_hold_until,
+                                    (&mut startup_input_hold_until, &mut hold_all_startup_input),
                                 )
                                 .await?;
                             }
@@ -3545,7 +3558,7 @@ async fn run_terminal(
                         &mut send_seq,
                         &mut predictor,
                         (&mut pending_user_input, &mut pending_user_input_bytes),
-                        &mut startup_input_hold_until,
+                        (&mut startup_input_hold_until, &mut hold_all_startup_input),
                     )
                     .await?;
                 }
@@ -3826,66 +3839,8 @@ fn split_after_command_submit(bytes: &[u8]) -> Option<(Vec<u8>, Vec<u8>)> {
     Some((now, later))
 }
 
-fn update_shell_command_line(line: &mut Vec<u8>, bytes: &[u8], alternate_screen: bool) -> bool {
-    if alternate_screen {
-        line.clear();
-        return false;
-    }
-    let mut submitted_tui = false;
-    let mut i = 0;
-    while i < bytes.len() {
-        match bytes[i] {
-            b'\r' | b'\n' => {
-                submitted_tui |= shell_line_starts_tui(line);
-                line.clear();
-                i += 1;
-            }
-            0x08 | 0x7f => {
-                line.pop();
-                i += 1;
-            }
-            0x20..=0x7e => {
-                line.push(bytes[i]);
-                i += 1;
-            }
-            0x1b => {
-                // Cursor editing makes the cheap line model unreliable; drop it
-                // rather than guessing. The command will still run normally.
-                line.clear();
-                i = bytes.len();
-            }
-            _ => {
-                i += 1;
-            }
-        }
-    }
-    submitted_tui
-}
-
-fn shell_line_starts_tui(line: &[u8]) -> bool {
-    let Ok(text) = std::str::from_utf8(line) else {
-        return false;
-    };
-    let command = text
-        .trim_start()
-        .split([' ', '\t', ';', '&', '|'])
-        .next()
-        .unwrap_or("");
-    matches!(
-        command,
-        "tm" | "tmux"
-            | "vim"
-            | "nvim"
-            | "vi"
-            | "less"
-            | "more"
-            | "man"
-            | "top"
-            | "htop"
-            | "btop"
-            | "btm"
-            | "watch"
-    )
+fn should_hold_post_submit_input(bytes: &[u8]) -> bool {
+    matches!(bytes.first(), Some(0x1b | 0x9b))
 }
 
 async fn flush_pending_user_input(
@@ -3912,8 +3867,9 @@ async fn flush_startup_input_if_ready(
     send_seq: &mut u64,
     predictor: &mut Predictor,
     pending: (&mut VecDeque<Vec<u8>>, &mut usize),
-    startup_hold_until: &mut Option<Instant>,
+    hold: (&mut Option<Instant>, &mut bool),
 ) -> Result<()> {
+    let (startup_hold_until, hold_all_startup_input) = hold;
     let Some(deadline) = *startup_hold_until else {
         return Ok(());
     };
@@ -3921,6 +3877,7 @@ async fn flush_startup_input_if_ready(
         return Ok(());
     }
     *startup_hold_until = None;
+    *hold_all_startup_input = false;
     let (pending, pending_bytes) = pending;
     flush_pending_user_input(
         socket,
@@ -5356,9 +5313,9 @@ mod tests {
         queue_pending_user_input, raw_contains_host_table, recv_response_until,
         render_status_clear, render_status_overlay, requested_env, resolved_startup_command,
         rewrite_forward_command, selected_predict_mode, selected_udp_host, server_version_mismatch,
-        split_after_command_submit, ssh_destination_host, ssh_username, ssh_with_user,
-        startup_command, status_ssh_target, toml_bare_key_or_quoted, update_check_requested,
-        update_shell_command_line, valid_forward_host,
+        should_hold_post_submit_input, split_after_command_submit, ssh_destination_host,
+        ssh_username, ssh_with_user, startup_command, status_ssh_target, toml_bare_key_or_quoted,
+        update_check_requested, valid_forward_host,
     };
     use dosh::config::{ClientConfig, CommandExtension, HostConfig};
     use dosh::native::EnvVar;
@@ -6023,21 +5980,12 @@ mod tests {
     }
 
     #[test]
-    fn shell_command_line_detects_tui_submit() {
-        let mut line = Vec::new();
-        assert!(!update_shell_command_line(&mut line, b"echo ok\r", false));
-        assert!(line.is_empty());
-
-        assert!(!update_shell_command_line(&mut line, b"t", false));
-        assert!(!update_shell_command_line(&mut line, b"m", false));
-        assert!(update_shell_command_line(&mut line, b"\r", false));
-        assert!(line.is_empty());
-
-        assert!(update_shell_command_line(
-            &mut line,
-            b"  btop\r\x1b[B",
-            false
-        ));
+    fn post_submit_hold_only_catches_terminal_control_input() {
+        assert!(should_hold_post_submit_input(b"\x1b[B"));
+        assert!(should_hold_post_submit_input(b"\x9bB"));
+        assert!(!should_hold_post_submit_input(b"next command"));
+        assert!(!should_hold_post_submit_input(b"\r"));
+        assert!(!should_hold_post_submit_input(b""));
     }
 
     #[test]