From d2fd7ee3467a0e822606ef64aaa9c2aacdb15251 Mon Sep 17 00:00:00 2001 From: DuProcess <273172371+DuProcess@users.noreply.github.com> Date: Sat, 27 Jun 2026 18:36:45 -0400 Subject: [PATCH] Make post-submit input gate command agnostic --- src/bin/dosh-client.rs | 130 +++++++++++++---------------------------- 1 file changed, 39 insertions(+), 91 deletions(-) diff --git a/src/bin/dosh-client.rs b/src/bin/dosh-client.rs index ed9edd1..66c0c07 100644 --- a/src/bin/dosh-client.rs +++ b/src/bin/dosh-client.rs @@ -2844,7 +2844,7 @@ async fn run_terminal( let mut pending_user_input: VecDeque> = VecDeque::new(); let mut pending_user_input_bytes = 0usize; let mut startup_input_hold_until: Option = None; - let mut shell_command_line: Vec = 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::>(); @@ -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, Vec)> { Some((now, later)) } -fn update_shell_command_line(line: &mut Vec, 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>, &mut usize), - startup_hold_until: &mut Option, + hold: (&mut Option, &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]